Can Model Merging Improve Aggregation in DiLoCo?

기술 요약: DiLoCo에서의 집합을 위한 모델 병합 (Model Merging)

문제 정의

Distributed Low-Communication (DiLoCo) 훈련은 글로벌 동기화를 최소화함으로써 통신 연결이 좋지 않은 컴퓨팅 클러스터("아일랜드") 전반에서 대규모 언어 모델을 훈련하는 방법으로 부상했습니다. DiLoCo에서 워커(worker)들은 글로벌 옵티마이저를 통해 파라미터 업데이트(의사 그래디언트, pseudo-gradients)를 집합하기 전에 여러 번의 로컬 최적화 단계(내부 단계, inner steps)를 수행합니다. 이 방식은 통신 비용을 줄이는 데는 효과적이지만, 워커의 수($R$)나 로컬 단계의 수($H$)가 증가함에 따라 성능 저하가 발생합니다. 이러한 저하는 로컬 복제본들이 크게 발산하여, 평균화된 의사 그래디언트가 긴밀하게 동기화된 데이터 병렬 훈련에서 생성되는 업데이트와 정렬되지 않기 때문에 발생합니다.

동시에, 모델 병합(Model Merging) 기술은 독립적으로 미세 조정된 모델들을 하나의 모델로 결합하기 위해 개발되었습니다. 병합의 핵심 과제는 "태스크 간섭(task interference)"으로, 이는 서로 다른 태스크로부터의 파라미터 업데이트가 충돌하여, 특히 병합할 모델의 수가 많아질 때 성능 저하를 유발하는 현상입니다.

저자들은 두 분야 사이에 구조적 유사성이 있음을 식려냈습니다: DiLoCo의 의사 그래디언트 집합 단계(SGD 외곽 옵티마이저 사용 시)는 수학적으로 태스크 산술(Task Arithmetic) 모델 병합과 동일합니다. 이 비유에서:

• 외곽 단계 시작 시의 공유 글로벌 모델은 사전 훈련된 초기화 모델에 해당합니다.
• 로컬 워커 복제본은 미세 조정된 "전문가(expert)" 모델들에 해당합니다.
• 로컬 업데이트(의사 그래디언트)는 "태스크 벡터(task vectors)"에 해당합니다.
• 워커의 수는 태스크의 수에 해당합니다.

핵심 문제는, 태스크 간섭을 완화하기 위해 특별히 설계된 고급 모델 병합 기술을 DiLoCo의 집합 단계에 재용용하여, 워커 수나 로컬 단계가 증가함에 따라 발생하는 성능 저하를 개선할 수 있는지 여부입니다.

방법론

이론적 가교 (Theoretical Bridge)

본 논문은 SGD를 사용하는 DiLoCo의 외곽 업데이트가 태스크 산술의 직접적인 적용임을 입증합니다. 의사 그래디언트의 집합을 모델 병합 문제로 간주함으로써, 저자들은 표준 평균화 연산을 부정적 간섭을 줄이도록 설계된 정교한 병합 알고리즘으로 대체할 것을 제안합니다.

병합 방법의 실험적 평가

저자들은 DiLoCo의 집합 단계에 드롭인 교체(drop-in replacement)로 사용할 수 있도록 여러 최첨단 병합 방법들을 체계적으로 벤치마킹했습니다. 평가된 방법들은 다음과 같습니다:

• 희소성 기반 (Sparsity-based): TIES (TrIm, Elect Sign) 및 DARE (Drop And REscale).
• 직교화 기반 (Orthogonalization-based): TSV (Task Singular Vector), Iso-C (Isotropic merging in Common subspace), 그리고 Iso-CTS.

병합 방법에 대한 조사 결과:

• 희소성 방법 (TIES, DARE): 매우 낮은 희소성 수준에서만 표준 DiLoCo SGD와 유사한 성능을 보였습니다. 희소성이 증가함에 따라 성능이 저하되었으며, 이는 이 방법들이 이 문맥에서 단순 평균화보다 뚜렷한 이점을 제공하지 못했음을 시사합니다.
• TSV: 성능이 낮았으며(높은 검증 손실), 단계당 여러 번의 특이값 분해(SVD)를 필요로 하여 계산 비용이 과도했습니다.
• Iso-CTS: 베이스라인보다 우수한 성능을 보였으나 계산 비용이 많이 들었으며, 로컬 태스크가 짧고 유사한 DiLoCo 환경에서는 Iso-C에 비해 큰 이점이 없었습니다.
• Iso-C: 가장 유망한 후보로 떠올랐습니다. 이는 집합된 의사 그래디언트의 특이값 스펙트럼을 그 평균값으로 대체하여, 고에너지 방향을 평탄화하고 균형 잡힌 공통 부분 공간(common subspace)을 생성함으로써 성능을 개선합니다. 이는 파라미터 행렬당 단 한 번의 SVD만을 요구하므로 유리한 비용 대비 성능 트레이드오프를 제공합니다.

제안된 알고리즘: IsoLoCo

Iso-C의 성공을 바탕으로, 저자들은 IsoLoCo (Isotropic Distributed Low-Communication)를 제안합니다. IsoLoCo는 다음과 같이 DiLoCo를 위해 Iso-C를 적응시킵니다:

1. 등방성 교정 (Isotropic Correction): 2D 파라미터 행렬(가중치 레이어)에 대해, 평균된 의사 그래디언트의 SVD를 계산하고, 특이값을 그 평균값으로 교체한 뒤 업데이트를 재구성합니다. 이는 부분 공간 구조를 보존하면서 스펙트럼을 정규화합니다.
2. 네스테로프 모멘텀 (Nesterov Momentum): 표준 Iso-C와 달리, IsoLoCo는 외곽 옵티마이저에 네스테로프 모멘텀을 통합합니다. 이는 FedOpt/DiLoCo 환경에서 모멘텀의 이점을 보여준 선행 연구에 근거합니다. 교정된 의사 그래디언트는 최종 업데이트를 생성하기 위해 모멘텀 버퍼로 전달됩니다.

구현 최적화

SVD의 계산 비용을 해결하기 위해, 저자들은 **뉴턴-슐츠 반복법(Newton–Schulz iterations)**을 사용하여 IsoLoCo의 빠른 구현체를 개발했습니다. 이는 표준 SVD를 반복적인 직교화로 대체하고 RMS 프록시를 사용하여 특이값 평균을 구합니다. 이 최적화는 성능 저하를 거의 일으키지 않으면서 직교화 속도를 최대 2개 차수(100배)까지 높여줍니다.

주요 결과

워커 수에 따른 성능

Llama 스타일 모델(178M, 512M, 1B 파라미터)을 DCLM 데이터셋으로 실험했습니다.

• 워커 수에 따른 스케일링: IsoLoCo는 워커 수가 증가함에 따라 표준 DiLoCo를 크게 능가합니다.
  • 178M 규모에서 $R=128$일 때, IsoLoCo는 DiLoCo 대비 데이터 병렬 베이스라인과의 검증 손실 격차를 5.8% 줄였습니다.
  • $R=128$에서 IsoLoCo는 데이터 병렬 훈련 대비 9.7%의 저하를 보인 반면, DiLoCo는 16.5%의 저하를 보였습니다.
  • IsoLoCo는 낮은 워커 수($R \le 4$)에서도 데이터 병렬 베이스라인보다 우수한 성능을 보였습니다.
• 모델 크기: IsoLoCo의 DiLoCo 대비 이점은 모델 크기(178M ~ 1B) 전반에 걸쳐 지속되지만, 모델이 커질수록 상대적인 개선 폭은 약간 작아집니다 (1B에서 1.2% vs 178M에서 약 2.4%). 이는 큰 모델이 워커 수에 따라 자연스럽게 더 잘 확장되기 때문으로 보입니다.

내부 단계(Inner Steps)의 영향

IsoLoCo는 동기화 간격($H = 30$ ~ $240$)이 증가하는 환경에서 평가되었습니다. IsoLoCo는 모든 간격에서 일관되게 DiLoCo보다 우수한 성능을 보였으며, 이는 더 긴 로컬 발산 기간에 대한 견고성을 입증합니다.

모멘텀 기반 DiLoCo와의 비교

특히, Iso-C 집합(모멘텀 없이)을 사용하는 DiLoCo SGD가 표준 DiLoCo(네스테로프 모멘텀 사용)보다 이미 우수했습니다. 여기에 네스테로프 모멘텀을 추가하여 IsoLoCo를 만든 결과, 성능 격차가 더욱 벌어졌으며, 이는 병합 기반 집합과 모멘텀 메커니즘이 상호 보완적임을 나타냅니다.

Muon과의 호환성

본 논문은 IsoLoCo가 Muon 내부 옵티마이저(MuLoCo)와 결합될 수 있음을 언급합니다. 이 구성에서 IsoLoCo는 $R \ge 32$인 경우 모든 방법 중 가장 낮은 검증 손실을 달이트었으며, AdamW 내부 옵티마이저를 사용하는 IsoLoCo는 높은 워커 수에서 MuLoCo와 대등하거나 이를 능가하면서도 훨씬 낮은 직교화 비용을 보였습니다.

의의 및 주장

본 논문은 저통신 분산 훈련을 모델 병합의 관점에서 바라보는 것이 성능 향상을 위한 강력한 프레임워크를 제공한다고 주장합니다. 주요 기여는 다음과 같습니다:

1. 이론적 통합: DiLoCo의 의사 그래디언트 집합과 태스크 산술 모델 병합 사이의 직접적인 유사성을 확립했습니다.
2. 경험적 검증: 모델 병합 기술, 특히 Iso-C가 단순 평균화 및 모멘텀 기반 DiLoCo보다 우수한 집합 규칙으로서 역할을 할 수 있음을 입증했습니다.
3. IsoLoCo: 많은 워커 수나 긴 동기화 간격이 있는 환경에서 데이터 병렬 골드 스탠다드와의 성능 격차를 크게 좁히는 실용적인 모멘텀 강화 알고리즘을 제안했습니다.
4. 효율성: SVD 집약적인 병합 과정을 반복적인 분산 훈련에서 계산 가능하게 만드는 Newton-Schulz 기반 구현을 제공했습니다.

저자들은 병합에서 영감을 얻은 집합 방식이 모델 품질을 희생하지 않으면서 지리적으로 분산되거나 연결이 불량한 클러스터 전반에서 훈련 규모를 확장할 수 있는 효과적인 전략이라고 결론짓습니다.