2、深度学习训练基础：模型训练流程、显存瓶颈、数据并行与模型并行概念
在正式进入 DeepSpeed 的零冗余优化之前，必须先回到深度学习训练本身：一次训练迭代到底发生了什么？显存被谁吃掉了？当单卡放不下模型时，我们有哪些拆分思路？只有把这些问题看清楚，后面理解 ZeRO 的“分片”才有落脚点。
一、一次训练迭代的完整流程
无论模型多大，单步训练都可以拆成四个阶段：前向传播、损失计算、反向传播、参数更新。前向传播把输入 x 逐层变换得到预测 y_hat；损失函数比较预测与标签得到标量 loss；反向传播从 loss 出发，按链式法则求出每个参数对应的梯度；优化器再根据梯度更新参数。
for x, y in dataloader:
y_hat = model(x) # 前向传播
loss = criterion(y_hat, y) # 损失计算
loss.backward() # 反向传播，产生 .grad
optimizer.step() # 参数更新
optimizer.zero_grad() # 梯度清零关键点在于：反向传播需要用到前向传播时保存的中间激活值。也就是说，前向阶段不仅要算出结果，还要把每层的输入、输出等张量留在显存里，供反向阶段计算梯度使用。这直接决定了显存占用的规模。
二、显存到底被什么占满了
训练显存大致由四部分组成，理解它们的比例是排查 OOM 的基础：
- 模型参数：权重本身，FP16 下每个参数 2 字节。
- 梯度：与参数量同规模，通常也是 2 字节。
- 优化器状态：Adam 会为每个参数保存一阶动量与二阶动量，混合精度训练还常保留 FP32 主权重，这部分往往是参数量的数倍。
- 激活值：前向保存的中间张量，与 batch size、序列长度强相关，是波动最大的一块。
| 组成 | 典型精度 | 相对参数量倍数 | 是否随 batch 变化 |
|---|---|---|---|
| 模型参数 | FP16 | 2 字节/参数 | 否 |
| 梯度 | FP16 | 2 字节/参数 | 否 |
| 优化器状态 | FP32 | 约 12 字节/参数 | 否 |
| 激活值 | FP16 | 与层数、序列长度相关 | 是 |
三、数据并行与模型并行
当单卡放不下时，第一反应通常是加卡。但“加卡”有两种完全不同的切法。
数据并行（Data Parallelism，DP）：每张卡都持有一份完整模型副本，把一个大 batch 切成若干小 batch 分给各卡。各卡独立前向、反向，得到各自梯度，再通过 AllReduce 求平均，最后每张卡用相同梯度更新，从而保持参数一致。它的优点是实现简单、扩展性好；缺点是每张卡都要存完整模型、梯度、优化器状态和激活，显存并没有被分摊。
模型并行（Model Parallelism，MP）：把模型本身切开。张量并行把单个矩阵乘按行或列拆到多卡，流水线并行把不同层分到不同卡。它能让单卡放不下的超大模型跑起来，但通信频繁、对切分策略和调度要求高，且容易出现卡间负载不均。
四、一个真实场景：为什么 8 卡 DP 仍然 OOM
假设有一个 7B 参数模型，用 8 张 40GB 卡做数据并行，混合精度 + Adam：
参数 FP16: 7B × 2 = 14 GB
梯度 FP16: 7B × 2 = 14 GB
优化器状态 FP32: 7B × 12 = 84 GB
激活值: 随 batch 变化，约 10~30 GB
合计: 约 122~142 GB / 每卡可以看到，即使有 8 张卡，每张卡依然要装下全部 122GB 以上的状态，40GB 显存远远不够。数据并行只分摊了数据，没有分摊模型状态。这正是 ZeRO 要解决的问题：把参数、梯度、优化器状态也切开，分给各卡保存，需要时再通信取回。
五、小结与关键概念对照
- 训练迭代 = 前向 + 损失 + 反向 + 更新，反向依赖前向保存的激活。
- 显存四大块：参数、梯度、优化器状态、激活值，其中优化器状态占比最高。
- 数据并行复制模型、切分数据，省时间不省显存。
- 模型并行切分模型，能装下超大模型，但通信与调度复杂。
- ZeRO 属于“数据并行 + 状态分片”，在保留 DP 易用性的同时分摊显存。
理解了这些基础，接下来看 DeepSpeed 如何用 ZeRO 把优化器状态、梯度、参数逐级分片，就能明白它为什么能在不改变模型代码的前提下大幅降低单卡显存。