由Zhengqing Yuan等研究人员开发的全新训练系统MegaTrain，近日在arXiv发布论文，展示了在单块GPU上运行超大规模语言模型的突破性能力。该技术通过将内存管理重心从GPU转移至主机内存，大幅降低了硬件门槛。
内存中心化架构的突破
传统的大模型训练依赖于昂贵的GPU集群，因为模型参数和优化器状态通常需要常驻显存。MegaTrain采取了截然不同的策略：它将GPU仅视为瞬时的计算引擎，而将参数和优化器状态存储在主机（CPU）内存中。在训练过程中，系统逐层将参数流式传输至GPU进行计算，随后将梯度回传至内存。
为了克服CPU与GPU之间的数据传输带宽瓶颈，研究团队引入了双缓冲执行引擎。该引擎能够通过多个CUDA流，实现参数预取、计算与梯度卸载的流水线式重叠，确保GPU始终处于满载的计算状态。此外，该系统摒弃了传统的静态自动求导图，转而使用无状态的层模板，能够根据流式传输的权重动态绑定计算任务，从而消除了持久化图元数据带来的调度限制。
实验数据显示，在配备1.5TB主机内存的单张H200 GPU上，MegaTrain能够稳定训练高达1200亿参数的模型。与目前主流的DeepSpeed ZeRO-3 CPU卸载方案相比，该系统在训练140亿参数模型时，吞吐量提升了1.84倍。
此外，MegaTrain在单张GH200 GPU上实现了70亿参数模型在512k上下文长度下的训练任务。这一进展为开发者提供了更具成本效益的超大规模模型训练路径，无需依赖大规模并行计算集群即可进行高精度模型开发。
