# 把训练账本放进多张设备

课堂先看到“复制完整状态”和“分片保存、临用时取齐”的差别。本篇把这个动作展开成完整显存账本、前向/反向通信与不同并行切分，帮助判断省掉的重复在哪里，以及新代价在哪里。

一个模型文件能装进 GPU，不等于同一设备能训练它。训练还要保留更新所需的状态，并给中间计算留出空间。本主题采用与课堂一致的 8B 全参数训练构造，把“放什么”“分给谁”“怎样通信”分开解释；数字用于复算机制，不对应某个商用模型的实测峰值。

<a id="state-ledger"></a>
## 从持久状态开始记账

N44先沿工作权重→预测→loss→梯度走完一次更新，再把各对象放入账本。梯度描述当前损失对参数的局部变化；FP32主权重用于以较高精度累积更新，工作权重用于低精度计算，Adam两份历史统计帮助组织后续步长。激活和算子临时量与这些跨步保留的状态要分开。

在本课混合精度 Adam 假设下，80 亿个参数各有 2 字节工作权重、2 字节梯度、4 字节 FP32 主权重，以及各 4 字节的一阶和二阶动量。合计每参数 16 字节，对应 128 GB。这里 GB=10⁹ bytes；激活、算子临时空间、通信缓冲和运行环境均未计入。实现可能使用不同的梯度精度、主权重策略或优化器，因此应先统一口径再比数字。[ZeRO 原论文](https://arxiv.org/abs/1910.02054v3)讨论了模型状态冗余；本课重新计算自己的四卡账本，不照抄附件中不一致的内存口径或公式。

| 每卡持久状态 | 工作权重 | 梯度 | 主权重与两份动量 | 合计 |
|---|---:|---:|---:|---:|
| DDP | 16 GB | 16 GB | 96 GB | 128 GB |
| ZeRO-1 | 16 GB | 16 GB | 24 GB | 56 GB |
| ZeRO-2 | 16 GB | 4 GB | 24 GB | 44 GB |
| ZeRO-3 | 4 GB | 4 GB | 24 GB | 32 GB |

这张表中的 ZeRO-3 每卡 32 GB 不是训练峰值。当前层收集后的参数、预取下一层时同时驻留的数据，以及尚未释放的激活，都可能把峰值抬高。降低持久状态只是其中一项工作。

<a id="data-and-shards"></a>
## 数据分工与保存分工

分布式数据并行 DDP 给不同设备相同模型和不同小批数据。各卡独立前向和反向，再同步梯度。在各卡样本数相同、损失归一化一致的条件下，汇总并平均本地梯度对应联合小批次的更新。样本数、序列有效 token 数不同而仍盲目等权平均，未必得到想要的训练目标。

DDP 复制了模型状态，所以单卡容量没有自动缩小。ZeRO 则按阶段分担保存责任：先分优化器相关状态，再分梯度，最后分工作参数。它保留数据并行的基本职责：每个数据 rank 仍把自己的数据送过全部层。不能把 ZeRO-3 画成“第一张卡只跑第一层”。[PyTorch DDP 说明](https://docs.pytorch.org/tutorials/intermediate/ddp_tutorial.html)和[FSDP2 说明](https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html)提供具体工程入口；分片思想相近并不意味着所有实现的通信调度完全相同。

<a id="zero3"></a>
## ZeRO-3 的一次完整计算

沿课堂采用的“层算完就释放完整参数”的教学调度：每卡带自己的数据进入本层；各卡通过 all-gather 交换分片，取得本层完整参数；本地前向得到下一层激活；释放临时完整参数，保留所属分片。全网形成损失后，反向按相反层序返回，在需要时再次 all-gather 参数，计算本地梯度。随后 reduce-scatter 合并各卡数据贡献，并让每卡只保留负责片段的合并梯度；最后各卡更新所属主权重、动量和工作参数分片。下一次前向取齐的是这些已经共同更新的参数分片，因此仍在训练同一个模型；不能把各卡拥有不同保存责任理解为各自训练一个互不一致的模型。

all-gather 是各处的分片在每个参与者处组成完整视图；reduce-scatter 是先逐元素合并贡献，再把结果分片交给参与者。all-reduce 则使各参与者取得完整的合并结果，从结果关系上可以理解为 reduce-scatter 后接 all-gather。这里描述的是数据语义，不表示库一定以这两个独立调用实现。NCCL 是提供这些 GPU 集合通信的工具之一，见[集合通信文档](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html)。

为了隐藏通信等待，真实系统会把梯度组织成桶，在某些层仍反向时同步已就绪的梯度，也可能提前收集下一层参数。较积极的预取有机会减少停顿，却会增加同时驻留的数据。数据已经量化或分片，也不意味着它移动没有代价；带宽、延迟和设备拓扑决定重叠是否足够。某些教材把全分片通信概括为比普通数据并行多一轮参数传输，这需要固定精度、设备数、释放策略和计数口径，不能当作所有部署不变的倍数。

<a id="tp-pp"></a>
## 切分计算本身：TP 与 PP

张量并行 TP 在层内切矩阵。若按输入维度拆分，x=[x₀,x₁]，W 由 W₀、W₁ 两个行块组成，则 xW=x₀W₀+x₁W₁；各卡的局部结果需要求和。若按输出维度切列，则得到不同输出分段，通信可安排到后续层。某个注意力头内部计算互不依赖，不代表整个 Transformer 层都不需要通信。实际并行方案必须把前后矩阵、残差和梯度一起考虑。

流水线并行 PP 按连续层分配阶段。一个微批次依次穿过阶段，后续微批次填入前级，从而使多阶段同时忙碌。刚开始尚未填满、结束时逐步排空，以及阶段耗时不均衡，都形成空泡。完整训练还要安排反向和激活：GPipe 风格先做一组前向再做反向，1F1B 将前向与反向交错，其内存与时间取舍不同。课堂仅展示两份微批次的前向，不能据此声称已经画出最优训练调度。

DP、TP、PP 和状态分片可以组合，但需要明确设备组：哪些设备处理不同数据，哪些共同计算一层，哪些承接相邻层。Megatron 对应大模型并行组织；NCCL 对应集合通信；二者不是可互相替代的同层工具。对自己的实验，先问限制是模型状态、激活、单层尺寸还是通信，再选最简单能解除该限制的组合。核心 CPU Lab 可以复算账本和分片语义，但不能测量真实多卡吞吐。

课堂依次在 [N44更新](../classroom/N44.html?step=b0)、[N45账本](../classroom/N45.html?step=b0)、[N46数据副本](../classroom/N46.html?step=b0)、[N47临时取齐](../classroom/N47.html?step=b0)、[N48共同更新](../classroom/N48.html?step=b0) 展开。原论文、课程材料的具体版本及采用位置见 [来源清单](../sources.html)。教学计算与硬件性能实测分别报告。