병렬 구조 결합: Mixture-of-Experts 모델을 위한 메모리 효율적인 학습 시스템

Mixture-of-Parallelisms: Towards Memory-Efficient Training Stack for Mixture-of-Experts Models

본 논문에서는 Mixture-of-Experts (MoE) 모델을 위한 메모리 효율적인 학습 시스템을 제시합니다. 이 시스템은 기존 및 새로운 병렬 처리 기술들을 MoE 모델의 다양한 계층과 단계에 결합하고 특화하여, CPU, CPU 메모리, GPU HBM 메모리, 그리고 GPU 클러스터 내에서의 CPU-GPU, GPU-GPU, 노드 간 통신 대역폭과 같은 물리적 제약 조건 하에서 최대한의 효율을 달성합니다. 또한, 높은 처리량과 메모리 효율성을 확보하기 위한 새로운 최적화 알고리즘 전략을 포함하고 있으며, 이를 통해 1조 규모의 모델을 거의 12개의 8x H200 GPU 노드로 사용하여, 최첨단 수준의 처리량과 메모리 효율성으로 손실 없는 사전 학습/미세 조정을 수행할 수 있습니다. 실험 결과, MoP는 정교하게 조정된 FSDP2 기준보다 GPU당 4.7배에서 8.2배 더 높은 처리량을 제공하며 (규모가 커질수록 그 격차가 벌어짐), 기준 시스템이 64~128K 토큰을 넘어서 메모리 부족에 직면하는 반면, 최대 1M 토큰까지 학습을 지속할 수 있습니다.

This paper showcases a memory-efficient training stack for Mixture-of-Experts (MoE) models. It is a training paradigm that combines and specializes various existing and novel parallelism techniques at different layers and stages of the Mixture-of-Experts (MoE) model training pipeline. It leverages these techniques to achieve maximal efficiency given the physical constraints of CPU, CPU memory, GPU HBM memory, and the CPU-GPU, GPU-GPU, and node-node communication bandwidth of the GPU cluster. It also contains a novel strategy for the optimizer step to achieve high throughput and memory efficiency, enabling practitioners to conduct lossless pre-training/fine-tuning of trillion-parameter scale models, at a million context length, with just under 12 8x H200 GPU nodes, with state-of-the-art throughput and memory efficiency. In our experiments, MoP delivers 4.7x--8.2x higher per-GPU throughput than a strongly-tuned FSDP2 baseline (with the gap widening at larger scale) and sustains training at context lengths up to 1M tokens, where the baseline runs out of memory beyond 64--128K.