⏶43
MegaTrain：在单张 GPU 上进行千亿级参数大语言模型的全精度训练
发表
由 
Zhengqing Yuan 提交
Zhengqing Yuan 提交作者: Zhengqing Yuan, Hanchi Sun, Lichao Sun, Yanfang Ye
摘要
AI 生成总结
MegaTrain 通过利用主机内存存储和优化的数据流技术，实现了在单个 GPU 上高效训练参数量超过 1000 亿的大语言模型。我们展示了 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 token 上下文的 7B 模型训练。
我们推出了 MegaTrain，一个以内存为中心的系统，能够高效地在单张 GPU 上以全精度训练参数量超过 100B 的大语言模型。与传统的以 GPU 为中心的系统不同，MegaTrain 将参数和优化器状态存储在主机内存（CPU 内存）中，并将 GPU 视为瞬时计算引擎。对于每一层，我们采用流式参数输入和梯度输出的方式，最大限度地减少持久化的设备状态。为了应对 CPU-GPU 带宽瓶颈，我们采用了两项关键优化。1) 我们引入了流水线双缓冲执行引擎，可在多个 CUDA 流中重叠参数预取、计算和梯度卸载，从而实现 GPU 的持续执行。2) 我们用无状态层模板取代了持久化的 autograd 图，随着权重的流入动态绑定权重，消除了持久化的图元数据，同时提供了调度灵活性。在配备 1.5TB 主机内存的单张 H200 GPU 上，MegaTrain 能够可靠地训练高达 120B 参数的模型。在训练 14B 模型时，其训练吞吐量达到了带有 CPU 卸载的 DeepSpeed ZeRO-3 的 1.84 倍。MegaTrain 还能在单张 GH200 上实现支持 512k token 上下文的 7B 模型训练。