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

L'entraînement de modèles Mixture-of-Experts (MoE) sur des contextes longs échoue dès qu'un composant dépasse la mémoire de l'appareil. L'article identifie quatre pics de mémoire non bornés : la répartition des experts, la projection du vocabulaire, les frontières de points de contrôle des gradients et l'état de l'optimiseur. Les auteurs introduisent quatre planifications — PipelinedLLEP, Ring-DTP, Selective Checkpoint Offload (SCO) et OffloadStreamAdamW — pour borner ces pics sans modifier la perte ni les gradients. Testées sur des modèles MoE de 120B à 667B paramètres, ces méthodes réduisent le pic de répartition MoE de 59,3 %, le pic de projection du vocabulaire de 86.6 %, et accélèrent l'étape de l'optimiseur de 2.05fois, permettant un entraînement à une longueur de contexte de 1M.

In matched component tests, they cut the MoE dispatch peak by up to 59.3% without losing throughput, the vocabulary projection peak by 86.6%, and the offloaded optimizer step by 2.05times faster.

Composed on MoE models from 120B to 667B parameters, they train at 1M context length, 8--32times the reach of a tuned FSDP2 baseline, and up to 10.4times its throughput.