论文
MegaTrain：单GPU上100B+参数大语言模型的全精度训练
MegaTrain: Full Precision Training of 100B+ Parameter Large Language Models on a Single GPU
摘要
我们推出了 MegaTrain，这是一种以内存为中心的系统，可以在单个 GPU 上以全精度高效训练 100B+ 参数 大语言模型。与传统的以 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 词元上下文的 7B 模型训练。