用单块显卡训练千亿参数大模型，突破内存瓶颈。
MegaTrain: Full Precision Training of 100B+ Parameter Large Language Models on a Single GPU
- 参数和优化器存于内存，显卡仅做临时计算
- 120B模型在单张H200上稳定训练，14B模型提速1.84倍
- 适合资源有限却想训大模型的研究者
我们提出MegaTrain，一种以内存为中心的系统，可在单块GPU上以全精度高效训练1000亿参数以上的大语言模型。与传统以GPU为中心的系统不同，MegaTrain将参数和优化器状态存储在主机内存（CPU内存）中，将GPU视为临时计算引擎。对于每一层，我们流式加载参数并输出梯度，最大限度减少设备端持久状态。为应对CPU-GPU带宽瓶颈，采用两项关键优化：1）引入流水线双缓冲执行引擎，通过多个CUDA流重叠参数预取、计算和梯度卸载，实现持续的GPU执行；2）用无状态层模板替代持久化的自动求导图，动态绑定流式加载的权重，消除持久图元数据的同时保持调度灵活性。在配备1.5TB主机内存的单张H200 GPU上，MegaTrain可稳定训练高达120B参数的模型。在训练14B模型时，其训练吞吐量比使用CPU卸载的DeepSpeed ZeRO-3高出1.84倍。此外，MegaTrain还支持在单张GH200上以512k token上下文训练7B模型。
原文摘要 · Abstract (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.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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。