MegaTrain: Treinamento em Precisão Total de Modelos de Linguagem Grandes com 100B+ de Parâmetros em uma Única GPU

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

April 6, 2026

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

cs.AI

Resumo

Apresentamos o MegaTrain, um sistema centrado na memória que treina eficientemente modelos de linguagem grandes com 100 bilhões ou mais de parâmetros, em precisão total, em uma única GPU. Diferente dos sistemas tradicionais centrados na GPU, o MegaTrain armazena parâmetros e estados do otimizador na memória hospedeira (memória da CPU) e trata as GPUs como unidades de computação transitórias. Para cada camada, transmitimos os parâmetros para dentro e calculamos os gradientes para fora, minimizando o estado persistente no dispositivo. Para combater o gargalo de largura de banda entre CPU e GPU, adotamos duas otimizações principais. 1) Introduzimos um motor de execução pipeline com double-buffering que sobrepõe a pré-busca de parâmetros, a computação e o descarregamento de gradientes através de múltiplos fluxos CUDA, permitindo execução contínua na GPU. 2) Substituímos grafos de autograd persistentes por modelos de camada stateless, vinculando pesos dinamicamente à medida que são transmitidos, eliminando metadados persistentes do grafo enquanto oferece flexibilidade no agendamento. Em uma única GPU H200 com 1,5 TB de memória hospedeira, o MegaTrain treina com confiabilidade modelos de até 120 bilhões de parâmetros. Ele também alcança um throughput de treinamento 1,84 vezes maior que o DeepSpeed ZeRO-3 com descarregamento para CPU ao treinar modelos de 14B. O MegaTrain também permite o treinamento de modelos de 7B com contexto de 512 mil tokens em um único 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.