2、深度学习训练基础：模型训练流程、计算图、反向传播、优化器、学习率调度、混合精度训练

在进入 DeepSpeed 的各类优化技术之前，必须先建立对深度学习训练“标准流程”的清晰认知。DeepSpeed 并不是替代训练框架的魔法，而是在这套标准流程的关键环节上做并行化、显存优化和通信压缩。理解每一步在做什么、消耗什么资源，才能判断 DeepSpeed 的优化点究竟落在哪里。

一、模型训练的标准流程

一次完整的训练迭代（iteration）通常包含以下步骤：

  1. 前向传播（Forward）：将一批数据输入模型，逐层计算，得到预测输出。
  2. 损失计算（Loss）：用损失函数比较预测值与真实标签，得到一个标量损失。
  3. 反向传播（Backward）：从损失出发，沿计算图反向求每个参数的梯度。
  4. 优化器更新（Optimizer Step）：根据梯度更新参数，如 SGD、Adam、AdamW。
  5. 梯度清零（Zero Grad）：清空上一轮梯度，准备下一轮。

这五步循环往复，直到损失收敛或达到预设的训练步数。看似简单，但每一步都涉及大量张量运算和显存占用，尤其是反向传播和优化器状态。

训练迭代的核心公式可以概括为：
前向：y = f(x; θ)
损失：L = loss(y, y_true)
反向：g = ∂L/∂θ
更新：θ = θ - lr · update(g)

二、计算图与反向传播

计算图（Computational Graph）是理解反向传播的关键。它把前向计算中的每个操作记录为节点，张量作为边，形成一个有向无环图。PyTorch 采用动态图机制，每次前向都会重新构建计算图，这为调试和灵活控制提供了便利，但也带来额外的构建开销。

反向传播本质上是链式法则在计算图上的应用：从损失节点出发，沿图反向传播梯度，每个节点根据局部导数计算上游梯度。对于包含数亿参数的模型，反向传播的计算量约为前向的两倍，显存占用也显著增加，因为需要保存中间激活值用于梯度计算。

计算图与反向传播示意 输入 x 线性层 Wx+b 激活函数 输出 ŷ 损失 L 反向传播：梯度沿计算图回传 前向传播：激活值逐层计算 优化器更新参数 θ θ ← θ − lr · update(g)

三、优化器与学习率调度

优化器决定如何利用梯度更新参数。常见优化器包括：

  • SGD：最基础，仅按梯度方向更新，收敛稳定但速度慢。
  • Momentum：引入动量项，加速收敛并抑制震荡。
  • Adam / AdamW：自适应学习率，为每个参数维护一阶和二阶矩估计，适合稀疏梯度和大规模模型。

Adam 类优化器需要为每个参数保存动量（m）和方差（v），因此优化器状态显存约为参数量的 2 倍。对于 10 亿参数的模型，仅优化器状态就可能占用数 GB 显存，这正是 DeepSpeed ZeRO 要解决的核心问题之一。

学习率调度（LR Scheduler）控制训练过程中学习率的变化。常见策略包括：

调度策略 特点 适用场景
Step Decay 每隔固定步数衰减为原来的 γ 倍 传统 CNN 训练
Cosine Annealing 按余弦曲线平滑衰减到最小值 Transformer 预训练
Warmup + Decay 先线性升温再衰减，避免早期震荡 大模型训练标配
Linear Decay 线性递减，简单可控 微调阶段
大模型训练通常采用 Warmup + Cosine 或 Warmup + Linear 组合。Warmup 步数一般为总步数的 1%～5%，学习率峰值需根据 batch size 调整，batch 越大峰值可适当提高。

四、混合精度训练

混合精度训练（Mixed Precision Training）是 DeepSpeed 和现代训练框架的默认配置。其核心思想是：前向和反向使用 FP16/BF16 加速计算并减少显存，而参数更新和优化器状态保留 FP32 以保证数值稳定性。

典型流程如下：

  1. 维护一份 FP32 主权重（master weights）。
  2. 每次迭代将主权重转换为 FP16/BF16 副本用于前向和反向。
  3. 反向得到 FP16 梯度后，转换为 FP32 并做梯度缩放（loss scaling）防止下溢。
  4. 用 FP32 梯度更新 FP32 主权重。
# 混合精度训练核心逻辑（伪代码）
master_weights = fp32(params)
for x, y in dataloader:
    with autocast(dtype=torch.float16):
        y_hat = model(x)
        loss = criterion(y_hat, y)
    scaled_loss = loss * scaler.get_scale()
    scaled_loss.backward()
    scaler.unscale_(optimizer)
    scaler.step(optimizer)
    scaler.update()

BF16 相比 FP16 有更大的动态范围，通常不需要 loss scaling，在 A100/H100 等硬件上更受青睐。DeepSpeed 支持 FP16 和 BF16 两种模式，并可与 ZeRO 优化器状态分区结合使用。

常见误区：认为混合精度只是把 dtype 改成 float16。实际上，若不保留 FP32 主权重、不做 loss scaling，训练很容易出现梯度下溢或参数更新失效，导致 loss 不下降甚至发散。

五、复盘：训练流程中的资源瓶颈

把上述流程串起来看，显存主要消耗在四块：模型参数、梯度、优化器状态、中间激活值。以 Adam + FP16 混合精度为例，一个 10 亿参数模型的显存占用大致如下：

  • FP16 参数副本：约 2 GB
  • FP32 主权重：约 4 GB
  • FP32 梯度：约 4 GB
  • Adam 动量与方差：约 8 GB
  • 激活值：随 batch size 和序列长度变化，可能数 GB 到数十 GB

可以看到，优化器状态和梯度占据了大部分显存，而它们并不直接参与前向计算。DeepSpeed 的 ZeRO 技术正是通过分区（partition）这些状态，让每张卡只保存一部分，从而在不改变模型代码的前提下大幅降低单卡显存需求。理解这一点，就理解了 DeepSpeed 优化技术的出发点。

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