@che_shr_cat: 1/ 我们一直误解了GPU内存的使用方式。如果GPU根本不需要存储你的模型呢？MegaTrain 实现了…

1/ 一直以来，我们对待 GPU 内存的方式都是错的。

如果 GPU 根本不需要存储你的模型呢？

MegaTrain 能够在一块 GPU 上实现 100B+ 参数规模 LLM 的全精度训练，它将显存变成了一个临时的、无状态的缓存。

秘诀何在？反转内存层次结构。

2/ 传统的卸载方案（如 ZeRO-Offload）将 CPU 内存视为慢速溢出缓冲区，导致严重的 PCIe 瓶颈。

MegaTrain 彻底颠覆了这一思路。CPU 主机持有权威的 12 字节/参数状态；GPU 仅动态持有当前激活层的参数。

3/ MegaTrain 协调三条并行的 CUDA 流： • S_comp（在缓冲区 0 上执行） • S_H2D（预取第 i+1 层的权重到缓冲区 1） • S_D2H（将梯度流式传回 CPU）

这个流水线将 CPU 到 GPU 的数据传输延迟隐藏在了活跃的数学计算背后。

4/ 它是如何绕过 PyTorch 沉重的自动求导计算图的？

MegaTrain 不再使用持久的全局有向无环图（DAG），而是采用无状态的层模板，不包含任何静态权重指针。

一个专门的 Bind 原语将来自活跃流式缓冲区的视图动态映射到执行槽位。

5/ 在反向传播过程中，系统通过一个手动的 LocalBackward 原语绕过全局 DAG。

它在受限的块内执行局部的正向重计算，计算单个层的梯度，并立即将其卸载。显存占用保持严格平坦。

6/ 为了充分利用 PCIe Gen4/Gen5 的带宽，MegaTrain 采用了层连续内存分片（Layer-Contiguous Memory Tiling）。

某一层的所有参数、梯度和状态都被打包成页对齐的连续块，从而能够在物理线路极限（Gen4 约 26 GB/s）上实现单次突发 DMA 传输。

7/ 缺点？性能严重受限于物理互连。

此外，宽度扩展比深度扩展更糟糕，因为层的参数量呈二次方增长。

最后，如果层传输时间 > 层计算时间（例如，高度稀疏的 MoE），流水线重叠就会失效。

8/ 回报是巨大的： • 在单块配备 1.5TB DDR5 内存的 H200 上完整训练了一个 120B 模型。 • 保持了精确的数值精度（在 MetaMathQA 14B 上达到 92.5%）。 • 在 GH200 上实现了 284 TFLOPS（是 ZeRO-3 Offload 的 1.8 倍）。

主机内存线性扩展，GPU 显存占用平坦。

9/ 这将大语言模型训练的瓶颈从极其昂贵的 GPU 显存容量转移到了廉价、通用的 CPU 内存上。

它让 100B+ 参数规模的微调和对齐从大规模分布式集群转移到了单台工作站节点上。

10/ 在此阅读我对 MegaTrain 的完整拆解： https://arxiviq.substack.com/p/megatrain-full-precision-training…

查阅论文： https://arxiv.org/abs/2604.05091

我们是否正走向一个廉价单节点训练的未来？欢迎讨论。关注我，获取每日技术论文深度解读。

11/ 我还草拟了双缓冲流式机制在底层的工作方式。下面请看图示说明。