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.84$\times$ 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.

翻译：我们提出MegaTrain——一种以内存为中心的系统，能够在单GPU上以全精度高效训练超过1000亿参数的大语言模型。不同于传统以GPU为中心的系统，MegaTrain将参数和优化器状态存储于主机内存（CPU内存），并将GPU视为瞬时计算引擎。针对每一层，我们流式加载参数并回传梯度，最大限度减少持久化设备状态。为攻克CPU-GPU带宽瓶颈，我们采用两项关键优化：1) 引入流水线双缓冲执行引擎，通过多CUDA流重叠参数预取、计算和梯度卸载，实现GPU持续执行；2) 用无状态层模板替代持久化自动求导图，在参数流式输入时动态绑定权重，既消除持久化图元数据，又提供调度灵活性。在搭载1.5TB主机内存的单块H200 GPU上，MegaTrain可稳定训练最高1200亿参数的模型；训练140亿参数模型时，其吞吐量达到使用CPU卸载的DeepSpeed ZeRO-3的1.84倍。此外，MegaTrain可在单块GH200上实现512K token上下文的70亿参数模型训练。