장문 컨텍스트 MoE 학습, 4가지 메모리 피크 동시 제어 기술 공개

장문 컨텍스트나 대형 배치 크기로 MoE 모델을 학습할 때, 단일 컴포넌트의 메모리 피크가 GPU 용량을 초과하면 학습이 실패하는 문제를 해결하기 위해 4가지 주요 메모리 피크를 모두 제어하는 기술이 제안됐다. 기존 병렬화 방식이 관리하지 못했던 전문가 디스패치, 어휘 투사, 그래디언트 체크포인트 경계, 옵티마이저 상태 등 4가지 피크는 모델, 컨텍스트 길이, 디바이스 수에 따라 다르게 증가한다.

연구팀은 GPU 워킹셋을 시작 시 고정하는 스케줄링을 통해 이 4가지 피크를 모두 제한했다. 구체적으로 PipelinedLLEP은 디스패치 청크에 기여하는 소스별 토큰 수에 상한을 설정하고, Ring-DTP는 어휘 투사 단계에서 활성화 또는 가중치 샤드를 링 구조로 순환시키며 온라인 log-sum-exp로 로짓 블록을 접어 피크를 감소시킨다. Selective Checkpoint Offload(SCO)는 체크포인트 경계의 수명이 긴 단일 텐서를 CPU 메모리로 오프로드하며, OffloadStreamAdamW는 옵티마이저 오프로드 시 직렬 CPU Adam 업데이트를 버킷 파이프라인으로 변환한다. 이 기법들은 계산 및 데이터 이동의 순서와 세분화만 변경하므로 손실 함수와 그래디언트는 정확성을 유지한다.

성능 테스트 결과, MoE 디스패치 피크는 최대 59.3% 감소했고, 어휘 투사 피크는 86.6% 감소했으며, 오프로드된 옵티마이저 스텝은 2.05배 빨라졌다. 120B에서 667B 파라미터의 MoE 모델에 적용했을 때, 1M 컨텍스트 길이에서 학습이 가능하며, 이는 튜닝된 FSDP2 기준선 대비 8~32배의 도달 범위와 최대 10.4배의 처리량을 제공한다.