提出高效训练万亿参数MoE模型的新方案，突破显存与通信瓶颈。
Mixture-of-Parallelisms: Towards Memory-Efficient Training Stack for Mixture-of-Experts Models
- 分层融合多种并行策略，动态优化资源利用
- 支持百万级上下文长度，吞吐比基线高4.7至8.2倍
- 适合大规模模型训练，尤其适用于显存受限场景
本文提出一种面向混合专家（Mixture-of-Experts, MoE）模型的内存高效训练框架。该框架在不同层和训练阶段融合并优化多种现有及新型并行技术，充分适配CPU、CPU内存、GPU HBM内存以及节点间通信带宽的物理限制。其创新性地设计了优化器步骤策略，在保持高吞吐与低内存占用的同时，使从业者可在仅需12个8卡H200节点的条件下，实现无损的万亿参数模型预训练/微调，上下文长度达百万级别。实验表明，相较于经过严格调优的FSDP2基线，该方案在单卡上实现4.7至8.2倍的吞吐提升，且在超长序列（最大100万词元）下仍可稳定训练，而基线在超过64K-128K词元时即因显存不足中断。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。