2、深度学习训练基础回顾：张量、自动微分、优化器（SGD/Adam）与分布式训练的基本概念（进程、通信、同步）

在深入大模型混合精度与数据并行训练之前，必须夯实底层基础。大模型训练的本质是对海量张量执行自动微分，并借助优化器迭代更新参数；而分布式训练则是在此基础上引入多进程协作、通信原语与同步机制。本章将系统梳理这些核心概念，并建立它们与后续混合精度、数据并行策略之间的映射关系。

深度学习训练基础与分布式扩展全景 单机训练核心 张量 (Tensor) 多维数组 + 数据类型 + 设备 自动微分 (Autograd) 计算图 + 反向传播 + 梯度 优化器 (SGD/Adam) 参数更新 + 动量 + 自适应学习率 迭代 分布式扩展 进程 (Process) 独立训练实例 + 独立显存 通信 (Communication) AllReduce / Broadcast / P2P 同步 (Synchronization) 屏障 + 全局梯度聚合 + 一致性 数据并行 / 模型并行 混合精度训练：FP16/BF16 计算 + FP32 主权重

上图展示了本章的知识框架：左侧为单机训练的三块基石（张量、自动微分、优化器），右侧为分布式训练的三个核心概念（进程、通信、同步），底部则引出混合精度训练作为两者的交汇点。下面逐一展开。

2.1 张量：深度学习的基本操作单元

核心原理

张量（Tensor）是深度学习框架中的核心数据结构，本质上是多维数组，但相比传统数组，它额外携带了三个关键属性：数据类型（dtype）、设备（device）和梯度（grad）。在大模型训练中，张量的维度通常非常高，例如一个 7B 参数模型的权重矩阵形状为 [num_layers, hidden_size, hidden_size]，而中间激活张量则可能达到 [batch_size, seq_len, hidden_size] 的规模。

数据类型直接决定了内存占用和计算精度。FP32（32位浮点）是基准精度，但大模型训练中我们常使用 FP16（半精度）或 BF16（脑浮点）来减少显存占用并加速计算。设备属性则决定了张量存储在 CPU 还是 GPU 上，以及具体在哪块 GPU 上——这是分布式训练中张量切分的基础。

判断方法：张量生命周期

理解张量在训练中的流转路径至关重要：

  1. 前向传播：输入数据经过各层计算，产生中间激活张量。
  2. 损失计算：模型输出与标签计算损失标量张量。
  3. 反向传播：损失对每个参数张量求梯度，梯度张量被暂存。
  4. 优化器更新：根据梯度张量更新参数张量。

在混合精度训练中，还需要额外维护一份 FP32 的“主权重”张量副本，用于参数更新时的数值稳定性。

案例拆解：显存占用估算

假设训练一个 1.5B 参数的模型，使用 Adam 优化器，FP32 训练时每个参数需要：权重（4字节）+ 梯度（4字节）+ 一阶动量（4字节）+ 二阶动量（4字节）= 16 字节。总显存需求约为 1.5B × 16 = 24GB。若改用混合精度（FP16 权重+FP32 主权重），则每个参数需要：FP16 权重（2字节）+ FP32 主权重（4字节）+ FP32 梯度（4字节）+ FP32 动量（8字节）= 18 字节，看似没有减少，但实际计算时 FP16 权重参与前向/反向，显存占用可降低约 30%-40%，因为激活张量也以 FP16 存储。

经验提醒：在 PyTorch 中，通过 tensor.dtype 和 tensor.device 可以快速检查张量属性。调试显存溢出（OOM）时，优先检查是否有意外创建的 FP32 大张量，以及是否所有张量都正确放置在目标 GPU 上。

2.2 自动微分：反向传播的工程实现

核心原理

自动微分（Autograd）是深度学习框架的“魔法引擎”。它通过构建动态计算图，在前向传播时记录每个张量的操作历史，然后在反向传播时利用链式法则自动计算梯度。计算图的节点是张量，边是操作（如矩阵乘法、ReLU、LayerNorm）。

在大模型训练中，自动微分面临两个关键挑战：计算图内存和梯度累积精度。对于 100 层以上的 Transformer 模型，计算图可能包含数万个节点，如果全部保存中间激活，显存会迅速耗尽。因此，框架实现了激活检查点（Activation Checkpointing）技术，只保存部分层的激活，反向传播时重新计算。

操作流程：梯度计算与累积

# PyTorch 自动微分示例
import torch
x = torch.randn(64, 1024, requires_grad=True)
w = torch.randn(1024, 4096, requires_grad=True)
b = torch.randn(4096, requires_grad=True)

# 前向
y = torch.matmul(x, w) + b
loss = y.square().mean()

# 反向
loss.backward()

# 此时 w.grad, b.grad, x.grad 已填充
print(w.grad.shape)  # torch.Size([1024, 4096])

在分布式训练中，每个进程计算的是本地梯度（基于本地数据分片的梯度），需要通过通信操作将各进程的梯度求和（或求平均），得到全局梯度后再更新参数。

常见误区：梯度覆盖与累积

一个常见错误是在循环中多次调用 backward() 而不清零梯度，导致梯度累加。PyTorch 默认是累加行为，因此每个 step 需要显式调用 optimizer.zero_grad()。但在梯度累积（Gradient Accumulation）场景下，我们反而需要利用这种累加特性，通过多次前向/反向累积梯度，再执行一次参数更新，以模拟更大的 batch size。

风险提示：在混合精度训练中，梯度可能会下溢为 0（FP16 最小正数为 6e-8）。如果梯度值小于该阈值，反向传播后梯度会丢失。解决方案是使用损失缩放（Loss Scaling）或改用 BF16（其指数位与 FP32 相同，不会下溢）。

2.3 优化器：SGD 与 Adam 的数学原理与选择

核心原理

优化器负责根据梯度更新模型参数。最基础的 SGD 更新公式为：

θ_{t+1} = θ_t - η * g_t

其中 η 是学习率，g_t 是 t 时刻的梯度。SGD 简单高效，但收敛速度慢且对学习率敏感。Adam 则通过维护一阶动量（m）和二阶动量（v）来自适应调整每个参数的学习率：

m_t = β1 * m_{t-1} + (1-β1) * g_t
v_t = β2 * v_{t-1} + (1-β2) * g_t^2
θ_{t+1} = θ_t - η * m_t / (√v_t + ε)

Adam 在大模型训练中几乎成为默认选择，因为它对学习率的鲁棒性更强，且能处理稀疏梯度。但 Adam 需要额外存储两倍于参数大小的动量张量，显存开销显著增加。

对比分析：SGD vs Adam

维度 SGD + Momentum Adam
收敛速度 较慢，需要精细调参 快，自适应学习率
显存开销 1 份动量（与参数同大小） 2 份动量（一阶+二阶）
泛化性能 通常更好（尤其在 CV 任务） 可能略差，但大模型场景差异不大
混合精度适配 简单，只需 FP32 主权重 需要额外处理动量精度，通常保持 FP32
大模型推荐 较少使用 主流选择（AdamW 变体）

实战场景：AdamW 与权重衰减

在大模型训练中，实际使用的是 AdamW（Adam with Weight Decay），它将权重衰减从梯度更新中解耦，直接在参数上执行 L2 正则化。PyTorch 中的实现为：

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)

在混合精度训练中，优化器状态（动量）必须保持 FP32 精度，否则会因精度不足导致训练不稳定。这也是混合精度训练中“主权重”和“优化器状态”始终以 FP32 存储的原因。

2.4 分布式训练基本概念：进程、通信与同步

核心原理

分布式训练的核心思想是将一个大任务拆分为多个子任务，由多个进程并行执行。每个进程拥有独立的模型副本（数据并行）或模型分片（模型并行），并持有自己的数据分片。进程之间通过通信交换梯度或激活信息，并通过同步机制保证训练的一致性。

在 PyTorch 中，分布式训练基于 torch.distributed 包实现，底层依赖 NCCL（NVIDIA 集体通信库）或 Gloo 后端。关键概念包括：

  • World Size：参与训练的总进程数。
  • Rank：每个进程的唯一编号（0 到 world_size-1）。
  • Local Rank：进程在单机内的编号，用于指定 GPU 设备。
  • Group：进程的子集，用于分组通信。

通信原语：AllReduce 详解

数据并行训练中最核心的通信操作是 AllReduce，它实现所有进程的梯度求和并广播结果。其底层实现有 Ring-AllReduce 和 Tree-AllReduce 两种算法。Ring-AllReduce 将进程组成环状，分两步完成：

  1. Reduce-Scatter：每个进程将梯度分块，沿环传递并累加，最终每个进程持有部分梯度的全局和。
  2. All-Gather：每个进程将持有的全局梯度分块广播给所有其他进程，最终所有进程获得完整梯度。

这种算法的通信复杂度为 O(P)（P 为进程数），且带宽利用率高，是大规模训练的首选。

同步机制：同步训练 vs 异步训练

同步训练要求所有进程完成本地反向传播后，统一执行 AllReduce 聚合梯度，然后同时更新参数。这种方式保证所有进程的模型始终一致，但受限于最慢的进程（掉队者问题）。异步训练则允许进程独立更新，通信开销低，但可能导致模型不一致，收敛不稳定。大模型训练几乎全部采用同步训练。

关键理解：在数据并行中，每个进程的 batch size 是全局 batch size 的 1/world_size。例如，4 个进程，每个进程 batch size=8，则全局 batch size=32。梯度 AllReduce 后，需要将梯度除以 world_size（或使用求平均模式）以保持学习率语义一致。

案例拆解：初始化与通信验证

以下是一个标准的 PyTorch 分布式初始化流程：

import torch.distributed as dist
import torch.multiprocessing as mp

def worker(rank, world_size):
    dist.init_process_group('nccl', rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)
    
    # 创建张量并执行 AllReduce
    tensor = torch.ones(4).cuda() * rank
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    print(f'Rank {rank}: {tensor.tolist()}')
    
    dist.destroy_process_group()

if __name__ == '__main__':
    world_size = 4
    mp.spawn(worker, args=(world_size,), nprocs=world_size)

输出结果应为每个进程都得到 [6, 6, 6, 6]（0+1+2+3=6）。这个简单实验验证了通信是否正常。在实际训练中，还需要设置 MASTER_ADDR 和 MASTER_PORT 环境变量，并处理多机通信的网卡选择。

常见误区与复盘

  • 误区一：忘记设置随机种子。分布式训练中，每个进程的模型初始化必须一致（数据并行场景），否则各进程模型不同步，AllReduce 后梯度混乱。应在所有进程中使用相同的种子初始化模型。
  • 误区二：AllReduce 的除法位置。有些实现将梯度求和后除以 world_size，有些则在优化器中处理。必须统一约定，否则学习率语义不一致。
  • 误区三：忽略 NCCL 超时设置。大规模训练中，单卡故障会导致 AllReduce 挂起。应设置 timeout 参数并实现故障检测。
风险提示：在混合精度 + 分布式训练的组合场景中，梯度 AllReduce 的精度选择至关重要。如果使用 FP16 梯度进行通信，会引入精度损失；推荐使用 FP32 梯度通信，或在通信前将梯度转换为 FP32。NCCL 支持 bf16 通信模式，但需确认硬件支持。

本章内容为后续混合精度与数据并行训练奠定了坚实基础。张量是操作对象，自动微分是梯度来源，优化器是更新策略，而进程、通信、同步则是将这一切扩展到多卡/多机环境的骨架。理解这些概念之间的相互作用，是掌握大模型训练工程化的前提。

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