MegaTrain: 在单块 GPU 上以全精度训练 100B+ 参数大型语言模型
计算与语言
2026-04-08 v1 分布式、并行与集群计算
操作系统
摘要
我们提出 MegaTrain，一个以内存为中心的系统，能够在单块 GPU 上高效地以全精度训练参数规模超过 100B 的大型语言模型。与传统的以 GPU 为中心的系统不同，MegaTrain 将参数和优化器状态存储在主机内存（CPU 内存）中，并将 GPU 视为瞬态计算引擎。对于每一层，我们流式传输参数并计算梯度输出，最小化持久设备状态。为应对 CPU-GPU 带宽瓶颈，我们采用两项关键优化。1) 我们引入了一种流水线双缓冲执行引擎，在多个 CUDA 流之间重叠参数预取、计算和梯度卸载，实现连续的 GPU 执行。2) 我们用无状态层模板替代持久自动微分图，在参数流式传输时动态绑定权重，消除持久图元数据，同时提供调度灵活性。在拥有 1.5TB 主机内存的单块 H200 GPU 上，MegaTrain 可可靠地训练多达 120B 参数的模型。在训练 14B 模型时，其训练吞吐量达到 DeepSpeed ZeRO-3 带 CPU 卸载的 1.84 倍。MegaTrain 还可在单块 GH200 上实现 512k token 上下文长度的 7B 模型训练。
引用
@article{arxiv.2604.05091,
title = {MegaTrain: Full Precision Training of 100B+ Parameter Large Language Models on a Single GPU},
author = {Zhengqing Yuan and Hanchi Sun and Lichao Sun and Yanfang Ye},
journal= {arXiv preprint arXiv:2604.05091},
year = {2026}
}