MegaTrain: Volledige precisie-training van 100B+ parameter grote taalmodelen op een enkele 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

Samenvatting

Wij presenteren MegaTrain, een geheugengericht systeem dat grote taalmodellen van 100B+ parameters efficiënt traint in volledige precisie op een enkele GPU. In tegenstelling tot traditionele GPU-gerichte systemen, slaat MegaTrain parameters en optimizer-statussen op in het hostgeheugen (CPU-geheugen) en behandelt het GPU's als tijdelijke rekenengines. Voor elke laag streamen we parameters naar binnen en berekenen we gradiënten naar buiten, waardoor de persistente toestand op het apparaat wordt geminimaliseerd. Om de CPU-GPU-bandbreedtebeperking te bestrijden, passen we twee belangrijke optimalisaties toe. 1) We introduceren een gepipelinde execution engine met dubbele buffering die parameter-prefetching, berekening en gradient-offloading overlapt over meerdere CUDA-streams, waardoor continue GPU-uitvoering mogelijk wordt. 2) We vervangen persistente autograd-grafieken door stateless laagsjablonen, waarbij we gewichten dynamisch binden tijdens het binnenstromen. Dit elimineert persistente grafiekmetagegevens en biedt tegelijkertijd flexibiliteit in planning. Op een enkele H200-GPU met 1,5 TB hostgeheugen traint MegaTrain betrouwbaar modellen tot 120B parameters. Het behaalt ook een 1,84 keer hogere trainthroughput dan DeepSpeed ZeRO-3 met CPU-offloading bij het trainen van 14B-modellen. MegaTrain maakt tevens training van 7B-modellen met een context van 512k tokens mogelijk op een enkele 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.