第2章 深度学习训练基础：模型训练流程、显存瓶颈、计算瓶颈、分布式训练基本概念

在真正使用 DeepSpeed 之前，必须先理解深度学习训练到底在做什么、瓶颈通常出现在哪里，以及分布式训练为什么能解决这些问题。很多初学者一上来就配置 ZeRO、Offload、混合精度，却说不清每一步在优化什么，结果调参全靠猜。本章把训练流程、显存与计算瓶颈、分布式训练的基本概念串成一条清晰的链路，为后续 DeepSpeed 的各个特性打好地基。

一、模型训练的标准流程

一次完整的训练迭代（iteration）可以拆成四个核心阶段，它们循环往复直到收敛：

  1. 前向传播（Forward）：输入一个 batch 的数据，逐层计算，得到模型输出（logits）。
  2. 损失计算（Loss）：用损失函数比较预测值与真实标签，得到一个标量 loss。
  3. 反向传播（Backward）：从 loss 出发，沿计算图反向求每个参数的梯度。
  4. 参数更新（Optimizer Step）：优化器根据梯度更新权重，例如 SGD 或 Adam。

用伪代码表示，训练循环大致如下：

for batch in dataloader:
    outputs = model(batch.inputs)      # 前向
    loss = criterion(outputs, batch.labels)  # 损失
    loss.backward()                    # 反向，梯度写入 param.grad
    optimizer.step()                   # 更新参数
    optimizer.zero_grad()              # 清空梯度

看似简单，但每一步都伴随着大量的张量运算和显存占用。理解这些占用从何而来，是判断瓶颈类型的关键。

二、显存瓶颈：训练时显存被谁吃掉了

训练显存主要由四部分组成：模型参数、梯度、优化器状态、激活值。以 Adam 为例，每个参数除了本身，还要保存一份梯度、一份一阶动量、一份二阶动量，相当于参数量的 4 倍左右（FP32 下）。激活值则是前向过程中每层输出，反向传播时需要复用，batch 越大、序列越长，激活值越庞大。

下面这张图展示了单卡训练时显存的主要构成与它们随 batch 增长的变化关系：

单卡训练显存构成与瓶颈示意 模型参数 梯度 优化器状态 激活值 固定开销 随 batch 增长 参数量 × 4 字节（FP32） 与参数量同规模 Adam 动量：参数量的 2 倍 与 batch、序列长度成正比 显存瓶颈典型表现 • CUDA out of memory • batch 无法继续增大 • 只能缩小模型或序列长度 缓解方向 • 混合精度（FP16/BF16） • 梯度累积 / 梯度检查点 • ZeRO 分片、CPU/NVMe Offload

当显存不足时，训练会直接抛出 CUDA out of memory。此时要么减小 batch，要么引入显存优化技术。DeepSpeed 的 ZeRO 正是针对参数、梯度、优化器状态做分片，把单卡显存压力分摊到多卡上。

三、计算瓶颈：GPU 到底在等什么

显存够用不代表训练就快。计算瓶颈通常表现为 GPU 利用率低、step 时间波动大。常见原因有三类：

  • 数据加载慢：CPU 预处理、磁盘 IO 跟不上 GPU，GPU 空转等数据。
  • 通信开销大：多卡训练时梯度同步、参数广播占用大量时间。
  • 算子效率低：小 batch、小矩阵乘无法打满 Tensor Core。

判断瓶颈类型的一个实用方法是看 GPU 利用率和 step 时间：如果利用率长期低于 50%，多半是数据或通信瓶颈；如果利用率高但 step 时间仍长，则可能是计算本身或显存带宽受限。

提示：显存瓶颈和计算瓶颈经常同时出现。先解决显存问题让训练能跑起来，再通过 profiling 定位计算瓶颈，是更稳妥的顺序。

四、分布式训练基本概念

分布式训练的核心思想是把计算和存储分摊到多张卡上。需要先分清几个高频术语：

术语 含义 典型作用
数据并行（DP） 每张卡持有完整模型，输入不同数据 提升吞吐，梯度需同步
模型并行（MP） 模型切分到多卡，各卡算一部分层 放下超大模型
流水线并行（PP） 按层分段，micro-batch 流水执行 减少气泡，提升利用率
张量并行（TP） 单层内矩阵切分到多卡 降低单卡参数与激活
All-Reduce 多卡梯度求平均并广播 数据并行的同步核心
通信组（Process Group） 参与通信的进程集合 决定通信范围与拓扑

数据并行是最常用的起点：每张卡跑相同模型、不同数据，反向结束后通过 All-Reduce 把梯度求平均，保证各卡参数一致。它的代价是通信量与参数量成正比，模型越大通信越重。DeepSpeed 的 ZeRO 在数据并行基础上进一步把优化器状态、梯度、参数分片，从而在保持数据并行易用性的同时大幅降低单卡显存。

常见误区：把“多卡”等同于“一定更快”。如果 batch 太小、通信占比过高，多卡反而可能比单卡慢。分布式训练要先确认瓶颈类型，再选择并行策略。

五、案例复盘：一次显存溢出的排查

假设你在单张 24GB 卡上训练一个 1.3B 参数的模型，batch size 设为 8 时出现 OOM。按下面的顺序排查：

  1. 先算固定开销：1.3B 参数 FP32 约 5.2GB，梯度约 5.2GB，Adam 状态约 10.4GB，合计已超 20GB。
  2. 再算激活值：batch=8、序列长度 512 时激活值可能再占数 GB，必然溢出。
  3. 缓解手段：改用混合精度可把参数与梯度降到约一半；开启梯度检查点可显著降低激活值；若仍不够，则引入 ZeRO-2 或 ZeRO-3 做分片。

这个例子说明：显存瓶颈往往不是单一原因，而是参数、梯度、优化器状态、激活值共同叠加的结果。理解每一部分的量级，才能有针对性地选择优化手段，而不是盲目调小 batch 或换卡。

← 上一章 📖 返回目录 下一章 →