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

我们提出了 MegaTrain，一种以内存为中心的系统，能够在单块 GPU 上以全精度高效训练参数量超过 1000 亿（100B+）的大语言模型。与传统的以 GPU 为中心的系统不同，MegaTrain 将模型参数和优化器状态存储在主机内存（即 CPU 内存）中，并将 GPU 视为临时性的计算引擎。对于每一层网络，系统按需流式加载参数并即时计算梯度，从而最大限度地减少 GPU 上需长期驻留的状态数据。为应对 CPU 与 GPU 之间的带宽瓶颈，我们采用了两项关键优化： 1）设计了一种流水线式的双缓冲执行引擎，利用多个 CUDA 流重叠完成参数预取、前向/反向计算以及梯度卸载操作，确保 GPU 始终处于连续计算状态； 2）摒弃了传统持久化自动微分计算图（autograd graph），转而采用无状态的层模板（stateless layer templates），在参数流式加载过程中动态绑定权重，既消除了计算图元数据的持久化开销，又保持了调度策略的高度灵活性。 在配备 1.5TB 主机内存的单块 H200 GPU 上，MegaTrain 可稳定训练参数量高达 1200 亿（120B）的模型；在训练 140 亿（14B）参数模型时，其训练吞吐量达到 DeepSpeed ZeRO-3 配合 CPU 卸载方案的 1.84 倍；此外，MegaTrain 还支持在单块 GH200 GPU 上训练具有 512K 令牌上下文长度的 70 亿（7B）参数模型。

如何在单个GPU上高效训练百亿参数级大语言模型（100B+），突破GPU显存容量瓶颈，避免传统分布式训练或混合精度/压缩带来的精度损失与系统复杂性。这是一个极具现实挑战的新问题——此前业界普遍认为单卡全精度训练100B+模型在工程上不可行。

颠覆GPU-centric范式，提出memory-centric架构：将全部模型参数和优化器状态持久存储于大容量主机内存（CPU RAM），GPU仅作为无状态、瞬时计算单元；通过双缓冲流水线（重叠参数加载-计算-梯度卸载）和动态绑定的stateless层模板（取代静态autograd图）消除GPU端持久状态与图开销，直击CPU-GPU带宽瓶颈。

在单块H200 GPU（配1.5TB CPU内存）上实现120B参数LLM全精度稳定训练；相较DeepSpeed ZeRO-3+CPU offloading提速1.84×（14B模型）；首次支持单GH200卡训练7B模型+512k长上下文；未提及其开源代码，但实验设计严谨（跨模型规模、硬件平台、基线对比）；值得深入的方向包括：通用化stateless执行引擎、支持异构内存层级（CXL）、与FlashAttention等算子深度协同、面向推理微调的轻量化迁移。

ZeRO-3 (DeepSpeed, 2021); OffloadEngine (Colossal-AI, 2022); SwapAdvisor (OSDI'23); vLLM with PagedAttention (2023); TensorRT-LLM CPU Offload Extensions (2024)