MegaTrain : Entraînement en pleine précision de modèles de langage de grande taille de 100 milliards de paramètres et plus sur un seul GPU

MegaTrain: Full Precision Training of 100B+ Parameter Large Language Models on a Single GPU

April 6, 2026

Auteurs: Zhengqing Yuan, Hanchi Sun, Lichao Sun, Yanfang Ye

cs.AI

Résumé

Nous présentons MegaTrain, un système centré sur la mémoire qui entraîne efficacement des modèles de langage de plus de 100 milliards de paramètres en pleine précision sur un seul GPU. Contrairement aux systèmes traditionnels centrés sur le GPU, MegaTrain stocke les paramètres et les états de l'optimiseur dans la mémoire hôte (mémoire du CPU) et utilise les GPU comme des moteurs de calcul temporaires. Pour chaque couche, nous transférons les paramètres et calculons les gradients, en minimisant l'état persistant sur le périphérique. Pour lutter contre le goulot d'étranglement de la bande passante CPU-GPU, nous adoptons deux optimisations clés. 1) Nous introduisons un moteur d'exécution pipeline à double tampon qui chevauche le préchargement des paramètres, le calcul et le déchargement des gradients sur plusieurs flux CUDA, permettant une exécution GPU continue. 2) Nous remplaçons les graphes de rétropropagation persistants par des modèles de couche sans état, liant les poids dynamiquement lors de leur transfert, éliminant ainsi les métadonnées de graphe persistantes tout en offrant une flexibilité dans l'ordonnancement. Sur un seul GPU H200 avec 1,5 To de mémoire hôte, MegaTrain entraîne de manière fiable des modèles jusqu'à 120 milliards de paramètres. Il atteint également un débit d'entraînement 1,84 fois supérieur à celui de DeepSpeed ZeRO-3 avec déchargement sur CPU pour l'entraînement de modèles de 14 milliards de paramètres. MegaTrain permet également l'entraînement de modèles de 7 milliards de paramètres avec un contexte de 512 000 tokens sur un seul GH200.

English

We present MegaTrain, a memory-centric system that efficiently trains 100B+ parameter large language models at full precision on a single GPU. Unlike traditional GPU-centric systems, MegaTrain stores parameters and optimizer states in host memory (CPU memory) and treats GPUs as transient compute engines. For each layer, we stream parameters in and compute gradients out, minimizing persistent device state. To battle the CPU-GPU bandwidth bottleneck, we adopt two key optimizations. 1) We introduce a pipelined double-buffered execution engine that overlaps parameter prefetching, computation, and gradient offloading across multiple CUDA streams, enabling continuous GPU execution. 2) We replace persistent autograd graphs with stateless layer templates, binding weights dynamically as they stream in, eliminating persistent graph metadata while providing flexibility in scheduling. On a single H200 GPU with 1.5TB host memory, MegaTrain reliably trains models up to 120B parameters. It also achieves 1.84times the training throughput of DeepSpeed ZeRO-3 with CPU offloading when training 14B models. MegaTrain also enables 7B model training with 512k token context on a single GH200.