TensorX
返回文献探索

Paper · arXiv 2604.05091

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

Zhengqing Yuan, Hanchi Sun, Lichao Sun, Yanfang Ye

48 upvotesApril 6, 2026arXiv 预印本
AI 摘要

MegaTrain enables efficient training of large language models with over 100 billion parameters on a single GPU by utilizing host memory storage and optimized data streaming techniques.

large language modelsfull precisionhost memoryoptimizer statestransient compute enginesparameter streaminggradient offloadingpipelined double-buffered execution engineCUDA streamsstateless layer templatespersistent autograd graphsDeepSpeed ZeRO-3CPU-GPU bandwidth bottleneck

Abstract

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.

北京市昌平区探索星信息技术及软件开发工作室

京ICP备2026059466号