MoE 모델 훈련을 위한 초고속 희소 체크포인팅 시스템

초록

MoEvement은 전문가(Expert) 파라미터를 부분적으로, 반복마다 순차적으로 저장하는 ‘희소 체크포인팅’ 기법과, 저장된 희소 스냅샷을 점진적으로 밀집 형태로 복원하는 변환 메커니즘, 그리고 파이프라인 단계 경계에서 활성값·그래디언트를 로그하는 ‘업스트림 로깅’ 기법을 결합해 MoE 모델 훈련 시 체크포인팅 오버헤드를 최대 4배, 복구 오버헤드를 최대 31배 감소시킨다.

상세 분석

MoEvement은 기존 체크포인팅이 밀집 모델에 최적화돼 있어, 파라미터 수가 수십 배로 늘어나는 MoE 모델에 적용하면 I/O와 메모리 병목이 급증한다는 근본적인 문제를 해결한다. 첫 번째 핵심 아이디어인 희소 체크포인팅은 전체 전문가 집합을 여러 서브셋으로 나누고, 각 서브셋을 서로 다른 반복에 저장한다. 이렇게 하면 매 반복마다 기록해야 할 파라미터 양이 전체의 1/N(여기서 N은 전문가 수) 수준으로 감소해 GPU‑CPU 간 데이터 전송 비용이 크게 낮아진다. 저자는 전문가 활성화 빈도에 기반해 우선순위를 매겨 자주 사용되는 전문가를 더 자주 스냅샷하도록 설계했으며, 이는 전체 학습 흐름에 거의 영향을 주지 않으면서도 최신 파라미터를 유지한다는 장점을 제공한다.

두 번째 핵심은 희소‑밀집 변환 메커니즘이다. 희소 체크포인트는 선택된 전문가에 대해서는 FP32, 나머지는 FP16 형태로 저장한다. 복구 시에는 먼저 FP16 파라미터를 로드하고, 필요에 따라 최신 FP32 스냅샷이 존재하는 전문가만 재계산한다. 이렇게 단계적으로 밀집 체크포인트를 재구성하면, 전통적인 전역 롤백 방식에서 발생하는 대규모 재연산을 피하면서도 동기식 학습 시퀀스와 정확도를 보장한다.

세 번째 혁신은 업스트림 로깅이다. 파이프라인 병렬화가 적용된 MoE 훈련에서는 각 단계의 출력(활성값)과 역전파 시의 그래디언트가 다음 단계로 전달된다. MoEvement은 이러한 경계 데이터를 로컬 메모리에 로그해 두고, 장애 발생 시 해당 데이터 파라미터가 포함된 데이터‑패럴렐 그룹만 최근 희소 체크포인트까지 롤백한다. 결과적으로 전역 롤백에 비해 복구 시간은 평균 1/31 수준으로 감소하고, 다른 워커는 중단 없이 학습을 지속할 수 있다.

실험 결과는 설득력 있다. 64개의 전문가를 가진 DeepSeek‑16.4B 모델을 96개의 A100 GPU에서 훈련했을 때, Gemini 기반 체크포인팅은 체크포인트당 2.5배의 오버헤드를 초래했지만 MoEvement은 동일 환경에서 오버헤드를 4배까지 감소시켰다. 또한 MTBF가 10분 수준으로 짧아도 ETTR(Effective Training Time Ratio)이 0.94 이상 유지돼, 실제 클라우드 비용 절감 효과가 크게 기대된다. 정확도 측면에서도 FP32‑FP16 혼합 정밀도 유지와 전문가 일관성 보장을 통해 기존 모델과 동일한 수렴 곡선을 보였다.

요약하면, MoEvement은 (1) 체크포인트 크기를 희소화해 I/O와 메모리 압박을 완화하고, (2) 단계적 밀집 복원을 통해 정확도와 동기화 요구를 만족시키며, (3) 로그 기반 국소 복구로 전역 롤백 비용을 최소화한다. 이 세 가지 설계는 대규모 MoE 훈련에서 발생하는 빈번한 장애와 거대한 파라미터 스토어 문제를 동시에 해결하는 통합 솔루션으로 평가할 수 있다.