DeepSpeed, Muon Optimizer 공식 지원으로 대규모 모델 학습 효율 향상

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

핵심 요약

DeepSpeed가 Muon Optimizer를 공식 지원하며 대규모 언어 모델 학습의 효율성을 개선했다. Muon은 신경망의 2D 가중치를 Newton-Schulz 반복법으로 직교화하여 모멘텀 행렬을 최적화하고, Adam 대비 메모리 사용량을 줄이면서 수렴 속도를 높인다. 실험 결과 Moonlight-16B-A3B 모델 파인튜닝 시 AdamW 대비 주요 벤치마크에서 우수한 성능을 보였고, Qwen2.5-3B 모델 학습 시 GPU 메모리 사용량을 9% 절감했다. 이 최적화는 2D 가중치에만 적용되는 하이브리드 방식으로, 임베딩이나 레이어 정규화 등 비2D 파라미터는 기존 AdamW를 유지한다.

배경

DeepSpeed, PyTorch, LLM 학습 및 파인튜닝 경험, Newton-Schulz 반복법 이해

대상 독자

대규모 언어 모델(LLM)을 학습하거나 파인튜닝하는 AI 엔지니어 및 연구자

의미 / 영향

Muon Optimizer의 DeepSpeed 지원은 대규모 모델 학습의 메모리 병목을 완화하여, 더 적은 자원으로도 고성능 모델을 학습할 수 있는 환경을 제공한다. 이는 Kimi-K2나 GLM-5와 같은 초대형 모델 학습에 이미 도입되어 실질적인 성능 향상을 입증하고 있다.

섹션별 상세

Muon은 신경망의 2D 가중치 행렬에 대해 Newton-Schulz 반복법을 적용하여 모멘텀 행렬을 직교화한다. 이 과정은 모멘텀 행렬의 특이값을 균일하게 만들어, AdamW 대비 드물지만 중요한 업데이트 방향을 효과적으로 증폭시킨다. 결과적으로 학습 효율과 수렴 속도가 향상된다.

Adam은 파라미터당 두 개의 모멘텀 버퍼를 사용하지만, Muon은 하나만 사용하여 메모리 상태를 절감한다. Muon은 2D 가중치에만 적용되는 하이브리드 옵티마이저로, 임베딩, 레이어 정규화, 편향 등 비2D 파라미터는 기존 AdamW를 사용한다. 이 구조를 통해 일반적인 트랜스포머 모델에서 옵티마이저 상태 메모리를 약 45% 줄일 수 있다.

DeepSpeed는 기존의 평탄화된 그라디언트 버퍼 구조를 수정하여 Muon 업데이트를 통합했다. Moonlight-16B-A3B 모델 파인튜닝 실험에서 Muon은 MBPP+, MMLU, GSM8K 벤치마크에서 AdamW를 상회하는 성능을 보였다. 또한 Qwen2.5-3B 모델 학습 시 GPU 메모리 사용량을 9% 절감하며 대규모 모델 학습의 실질적인 대안임을 입증했다.

실무 Takeaway

Muon Optimizer는 2D 가중치 행렬의 직교화를 통해 AdamW 대비 메모리 사용량을 줄이고 수렴 속도를 높인다.
DeepSpeed의 Muon 지원은 하이브리드 방식을 채택하여 2D 가중치에는 Muon을, 나머지 파라미터에는 AdamW를 적용해 안정적인 학습을 보장한다.
대규모 모델 학습 시 Muon Optimizer를 적용하면 옵티마이저 상태 메모리를 절감하여 더 큰 모델이나 배치를 처리할 수 있다.

언급된 리소스

GitHubdeepspeed_finetune_demo