기술 요약: DMuon – AdamW 수준의 오버헤드로 구현하는 효율적인 분산형 Muon 학습

1. 문제 정의

행렬 인식 최적화 도구(Matrix-aware optimizers), 특히 Muon은 가중치 행렬의 모멘텀 누적 그레디언트에 뉴턴-슐츠(Newton–Schulz, NS) 반복법을 적용함으로써 요소별(element-wise) 최적화 도구인 AdamW보다 우수한 수렴 특성을 입증했습니다. 이러한 접근 방식은 대규모 모델(예: LLM 및 Embodied Foundation Models)의 사전 학습(pretraining)에서 상당한 연산 효율성 이득을 제공합니다.

그러나 현대적인 분산 학습 환경에 Muon을 배포하는 데는 심각한 시스템 병목 현상이 존재합니다. 현대의 분산 학습 스택(예: ZeRO, FSDP)은 파라미터 샤드(shard)마다 독립적으로 업데이트를 계산하는 **요소별 최적화(element-wise optimization)**라는 암묵적인 계약 하에 작동합니다. Muon은 필요한 행렬 곱($XX^\top$)을 계산하기 위해 단일 장치에 전체 가중치 행렬이 존재해야 한다는 점에서 이 계약을 위반합니다.

나이브한(naive) 분산 구현("Gather-then-Compute")에서는 모든 랭크(rank)가 다음 과정을 거쳐야 합니다:

1. 모든 샤드로부터 축소된 전체 그레디언트 행렬을 All-gather로 수집합니다.
2. 로컬에서 NS 반복법을 수행합니다.
3. 자신이 소유하지 않은 샤드에 대한 결과는 폐기합니다.

이 방식은 두 가지 주요 비용을 발생시킵니다:

• 통신 오버헤드: 매 단계마다 전체 행렬을 재구성하는 것은 모델 크기와 분산 너비(distributed width)에 비례하는 집합 통신(collective communication)을 유발합니다.
• 중복 연산: 모든 랭크가 동일하게 값비싼 직교화(orthogonalization) 과정을 수행하므로, 연산 비용이 랭크 수($D$)만큼 배가됩니다.

실제로 이러한 비용으로 인해 분산 환경에서의 Muon 옵티마이저 스텝 시간은 AdamW의 포워드 및 백워드 패스를 합친 시간을 상회하거나 이에 필적하게 되며, 이는 Muon의 알고리즘적 이점을 퇴색시킵니다. 이러한 페널티는 특히 표준 LLM 사전 학습에 비해 옵티마이저 스텝의 상대적 비중이 큰 Embodied Model 학습(예: VLA)에서 더욱 심각하게 나타납니다.

2. 방법론: DMuon 시스템 설계

DMuon은 Muon과 AdamW 사이의 성능 격차를 줄이기 위해 설계된 오픈 소스 분산 구현체입니다. 이는 기존 학습 프레임워크(예: PyTorch/FSDP2)를 수정하지 않고도 통신 최적화, 부하 분산, 커널 수준의 효율성을 통합하는 공동 설계된(co-designed) 시스템 아키텍처를 통해 이를 달enc합니다.

2.1 소유자 중심 실행 전략 (Owner-Centric Execution Strategy)

모든 랭크에서 연산을 복제하는 대신, DMuon은 각 행렬 파라미터에 단 하나의 **소유자 랭크(owner rank)**를 할당합니다.

• 소유자 역할: 소유자는 권위 있는 옵티마이저 상태를 유지하며, 각 행렬에 대해 정확히 한 번만 Muon 업데이트(Newton–Schulz 반복)를 수행합니다.
• 비소유자 역할: 다른 랭크들은 포워드/백워드 계산에 필요한 경우에만 파라미터를 실체화하며, 직교화 과정을 수행하지 않습니다.
• 통신 흐름:
  • 포워드(Forward): 소유자는 업데이트된 파라미터를 비동기적으로 비소유자들에게 게시하며, 이는 레이어 실행과 파이프라이닝됩니다.
  • 백워드(Backward): 그레디언트는 패킹된 버퍼(packed buffers)에 누적되어 소유자에게 리듀스(reduce)됩니다.
  • 오버랩(Overlap): DMuon은 미세한 가중치 레이아웃을 가진 2단계 계층 구조(노드 내 및 노드 간)를 채택하여 통신을 계산과 오버랩시킴으로써 집합 통신 경합을 방지합니다.

2.2 효율적인 Gram Newton–Schulz 반복

소유자 측의 업데이트 연산 비용을 줄이기 위해, DMuon은 Gram Newton–Schulz 정식화를 채택합니다.

• Gram 공간 재귀: 원래의 $m \times n$ 공간에서 연산하는 대신, $m \times m$ Gram 공간($G = XX^\top$)에서 반복을 수행합니다. 이는 $m < n$일 때 복잡도를 $O(m^2n)$에서 $O(m^3)$으로 줄여줍니다.
• 대칭성 활용: $G$는 대칭 행렬이므로, DMuon은 주된 행렬 곱 연산의 산술 작업량을 절반으로 줄이는 SYRK 스타일의 실행 경로를 사용합니다.
• 배칭(Batching): 작은 행렬에서의 GPU 점유율 문제를 해결하기 위해, DMuon은 유사한 형태의 행렬들을 배치하고 Gram NS 재귀를 공동으로 실행하여 커널 디스패치 오버헤드를 분담(amortize)합니다.

2.3 연산 인지형 부하 분산 (Computation-Aware Load Balancing)

Muon 업데이트 비용은 행렬 형태에 따라 크게 달라지므로, 단순한 할당(예: 라운드 로빈)은 스트래글러(straggler) 랭크를 발생시킵니다. DMuon은 소유자 할당 문제를 **혼합 정수 선형 계획법(MILP)**으로 정식화합니다.

• 프로파일링: 초기화 단계에서 시스템은 각 고유한 행렬 형태에 대해 다양한 배치 크기 및 커널 스케줄의 실행 시간을 벤치마킹합니다.
• 최적화: MILP 솔버는 예측된 Makespan(최대 부하)을 최소화하도록 하여, 가장 무거운 행렬들이 워크로드를 균형 있게 분산하도록 보장합니다.
• 폴백(Fallback): 탐색 공간이 극도로 큰 경우, 초기화 시간을 제한하기 위해 그리디(greedy) 휴리스틱을 사용합니다.

2.4 형태 적응형 커널 자동 튜닝 (Shape-Adaptive Kernel Autotuning)

DMuon은 고성능 커널을 생성하기 위해 DSL 기반 자동 튜닝 워크플로우(TileLang 및 CUTE DSL 사용)를 활용합니다.

• 타겟 하드웨어와 행렬 형태에 특화된 타일 형태(tile shapes), 블록 크기, 소프트웨어 파이프라인 구성을 탐색합니다.
• 최적의 구성은 캐싱되어 재사용되며, 이를 통해 커널 스케줄이 트랜스포머 학습에서 발생하는 특정 차원 문제에 대해 항상 최적에 가깝게 유지되도록 합니다.

3. 주요 기여

1. 미세 조정된 통신 최적화: 계층적 스케줄링과 미세한 가중치 레이아웃을 통해 통신 오버헤드를 최소화하고 정확한 옵티마이저 의미론을 보존하는 소유자 중심 런타임.
2. 형태 적응형 실행 스택: 이질적인 행렬 워크로드를 효율적으로 처리하기 위해 배치형 Gram Newton–Schulz, 대칭 인식 커널 및 자동 튜닝을 결합한 고성능 엔진.
3. 연산 인지형 부하 분산: 이론적 FLOP 수가 아닌 실제 런타임 프로파일링에 기반하여 행렬 소유자를 할당하는 측정된 비용 최적화(MILP) 방식.
4. 드롭인 통합(Drop-in Integration): 프레임워크 수준의 수정 없이 최소한의 노력으로 기존 학습 파이프라인(예: FSDP2)에 통합 가능한 모듈.
5. 프로덕션 규모 검증: 로보틱스(Embodied) 및 언어 모델 학습 워크로드 모두에서 성공적인 배포.

4. 실험 결과

저자들은 네 가지 프로덕션 워크로드(Wall-OSS-0.5, Pi0, Wall-WM, Qwen2.5-7B)에 대해 8개에서 256개의 A800 GPU 규모로 DMuon을 평가했습니다.

• 옵티마이저 스텝 시간: DMuon은 나이브한 "Gather-then-Compute" 방식의 분산 Muon 구현과 비교하여 옵티마이저 스텝 시간에서 6.85배에서 163.00배의 속도 향상을 달성했습니다.
• 엔드 투 엔드(End-to-End) 스텝 시간: DMuon은 총 스텝 시간에서 1.48배에서 3.01배의 속도 향상을 달성했습니다.
• AdamW 대비 오버헤드: DMuon의 스텝당 지연 시간은 AdamW 수준에 근접하게 낮추었으며, 평가된 워크로드 전반에 걸쳐 평균 **+2%**의 오 오버헤드만을 기록했습니다.
• 확장성(Scaling): Wall-OSS에서 DMuon은 GPU 수가 증가함에 따라 AdamW에 근접한 처리량을 유지하는 반면, 나이브한 Muon은 심각한 통신 병목 현상을 겪게 됩니다.

속도 향상 세부 분석:

• 대칭 Gram 커널: 48% (주요 곱셈에서의 산술 연산 절반 감소).
• 소유자 스케줄링 및 부하 분산: 32% ($D$배의 중복 연산 제거).
• 자동 튜닝 및 배칭: 16% (작은 행렬에 대한 커널 런치 오버헤드 감소).

5. 의의 및 주장

본 논문은 DMuon이 행렬 인식 옵티마이저와 분산 학습 인프라 사이의 시스템 불일치를 성공적으로 해결했다고 주장합니다. 중복 연산을 제거하고 통신을 최적화함으로써, DMuion은 Muon의 수렴 이점(높은 연산 효율성 및 더 나은 최적화 기하학 등)이 실제 벽시계(wall-clock) 학습 속도 향상으로 이어지게 합니다.

저자들은 DMuon을 새로운 알고리즘이 아니라, Muon을 대규모 프로덕션 학습에서 사용 가능한 AdamW의 드롭인 대체제로 만드는 필수적인 시스템 조력자로 규정합니다. 이 연구는 행렬 인식 최적화의 "페널티"를 무시할 수 있는 수준으로 낮춤으로써, 대규모 언어 모델(LLM)과 Embodied Foundation Models(VLA) 모두를 위한 효율적인 스케일링을 가능하게 함을 보여줍니다.

논문은 DMuon이 Muon의 업데이트 규칙이나 수렴 특성을 변경하지 않으며, 오직 시스템 오버헤드만을 다룬다는 점을 명시합니다. 따라서 분산 환경에서 이점이 가장 두드러지며, 단일 GPU 학습에서는 제거할 분산 중복성이 없기 때문에 보다 완만한 개선(약 2배의 옵티마이저 속도 향상)을 보입니다.