논문

Flattening Every Memory Peak in Long-Context Mixture-of-Experts Training

Hugging Face Daily Papers

#MoE#Distributed Training#Memory Optimization#LLM#RAG

MoE 모델의 대규모 컨텍스트 및 대형 배치 학습 시 발생하는 4가지 핵심 메모리 피크를 제어하여 학습 효율을 극대화하는 기술

논문이 다루는 내용

MoE 모델을 긴 컨텍스트나 큰 배치 사이즈로 학습할 때, 평균 메모리 사용량이 아닌 특정 시점의 피크 메모리 초과로 인해 학습이 실패하는 문제가 발생합니다. 연구진은 전문가 분산(Expert Dispatch), 어휘 투영(Vocabulary Projection), 그래디언트 체크포인트, 옵티마이저 상태라는 네 가지 주요 메모리 병목을 식별했습니다. 이를 해결하기 위해 PipelinedLLEP, Ring-DTP, SCO, OffloadStreamAdamW라는 네 가지 스케줄링 기법을 제안하여 GPU 워킹 셋을 고정했습니다. 이 방법들은 계산 순서와 데이터 이동의 입도(granularity)만 변경하므로 손실 함수와 그래디언트의 정확도를 유지합니다. 실험 결과, 120B~667B 규모의 MoE 모델에서 기존 FSDP2 대비 최대 1M 컨텍스트 길이를 지원하며 훨씬 높은 처리량을 달성했습니다.

핵심 결과

4가지 핵심 메모리 피크(Expert Dispatch, Vocab Projection, Checkpoint, Optimizer)를 제어하는 통합 프레임워크 제안
계산 순서와 데이터 이동 최적화를 통해 수학적 정확도(Loss/Gradient)를 유지하면서 메모리 사용량 제어
대규모 MoE 모델(최대 667B)에서 초장거리 컨텍스트(1M) 학습 가능성 입증

실무에서 볼 만한 점

메모리 부족(OOM) 문제로 인해 컨텍스트 길이를 늘리지 못하는 상황에서, 하드웨어 자원을 효율적으로 분배하여 모델 크기와 컨텍스트 길이를 동시에 확장할 수 있는 실질적인 가이드를 제공합니다.

읽을 때 확인할 점

기존 MoE 학습 워크로드에서 메모리 프로파일링을 통해 어떤 피크가 가장 먼저 발생하는지 확인
제안된 4가지 기법 중 현재 인프라(CPU-GPU 대역폭 등)에 가장 적합한 기법 우선 적용 테스트
FSDP2와 같은 기존 병렬화 전략과 결합 시의 성능 변화 측정

Hugging Face에서 보기 ↗PDF 보기 ↗