MegaTrain: Addestramento in Precisione Completa di Modelli Linguistici di Grandi Dimensioni con oltre 100B di Parametri su una Singola GPU
MegaTrain: Full Precision Training of 100B+ Parameter Large Language Models on a Single GPU
April 6, 2026
Autori: Zhengqing Yuan, Hanchi Sun, Lichao Sun, Yanfang Ye
cs.AI
Abstract
Presentiamo MegaTrain, un sistema di tipo memory-centric che addestra efficientemente modelli linguistici di grandi dimensioni da 100+ miliardi di parametri in precisione completa su una singola GPU. A differenza dei sistemi tradizionali di tipo GPU-centric, MegaTrain memorizza i parametri e gli stati dell'ottimizzatore nella memoria host (memoria della CPU) e tratta le GPU come motori di calcolo transitori. Per ogni livello, trasferiamo i parametri in entrata e calcoliamo i gradienti in uscita, minimizzando lo stato persistente sul dispositivo. Per contrastare il collo di bottiglia della larghezza di banda CPU-GPU, adottiamo due ottimizzazioni chiave. 1) Introduciamo un motore di esecuzione pipeline con doppio buffer che sovrappone il precaricamento dei parametri, il calcolo e lo scaricamento dei gradienti su più stream CUDA, consentendo un'esecuzione GPU continua. 2) Sostituiamo i persistenti grafi di autograd con template di livello stateless, associando i pesi dinamicamente man mano che vengono trasferiti, eliminando così i metadati persistenti del grafo e fornendo al contempo flessibilità nella schedulazione. Su una singola GPU H200 con 1,5 TB di memoria host, MegaTrain addestra in modo affidabile modelli fino a 120 miliardi di parametri. Raggiunge anche una velocità di addestramento 1,84 volte superiore a quella di DeepSpeed ZeRO-3 con offload su CPU quando addestra modelli da 14B. MegaTrain consente inoltre l'addestramento di modelli da 7B con un contesto di 512k token su una singola 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.