机器学习数学基础专题【左扬精讲】—— 混合精度与分布式训练：从单机单卡到 DDP、FSDP、DeepSpeed
当模型参数量突破数十亿乃至数千亿时，单卡显存和单机算力成为瓶颈。混合精度训练（AMP）用半精度浮点减少显存占用、加速计算；分布式训练把计算和参数分散到多卡多机上，实现横向扩展。
本文从底层原理出发，系统讲解 PyTorch 生态中混合精度和分布式训练的核心机制、代码实现与选型指南。
本文所有代码片段均为通用 PyTorch 2.x 写法，可直接在 PyTorch 2.0+ 环境中运行验证 。
本文涉及的核心 PyTorch 模块（PyTorch 2.x）：
torch.cuda.amp ← 自动混合精度（AMP）
torch.cuda.amp.GradScaler ← 梯度缩放，防止 fp16 下溢
torch.nn.parallel.DistributedDataParallel ← DDP：数据并行
torch.distributed ← 分布式通信后端（NCCL / Gloo）
torch.distributed.fsdp ← Fully Sharded Data Parallel：参数分片混合精度 AMP GradScaler 分布式训练 DDP FSDP DeepSpeed 数据并行 参数分片 NCCL ZeRO
学习重点提示
★ 必须掌握
- 必须掌握：混合精度的核心 — 用 fp16 计算、fp32 保存权重，在速度和精度之间取平衡
- 必须掌握：GradScaler 解决 fp16 梯度下溢问题
- 必须掌握：DDP 把模型复制到每张卡，数据拆分，梯度做 All-Reduce
- 必须掌握：FSDP 在 DDP 基础上进一步分片参数和优化器状态
★ 建议掌握
- 建议掌握：ZeRO Stage 1/2/3 的渐进式优化思路
- 建议掌握：NCCL vs Gloo 通信后端的选型
- 建议掌握：DeepSpeed 如何封装 FSDP 和 ZeRO
本文目录
- 一、混合精度训练 — What & Why
- 二、GradScaler：防止 fp16 梯度下溢
- 三、自动混合精度（AMP）的 PyTorch 实现
- 四、分布式训练概述 — 从单机到多机
- 五、DDP（DistributedDataParallel）：数据并行
- 六、FSDP（Fully Sharded Data Parallel）：参数分片
- 七、DeepSpeed：封装 ZeRO 的高层框架
- 八、常见问题 FAQ（20 组）
- 九、Roadmap 后续预告
一、混合精度训练 — What & Why
What — 混合精度训练是什么？
混合精度训练（Automatic Mixed Precision, AMP）是让部分计算使用半精度（float16 / bfloat16），部分计算保持单精度（float32）的技术。前向传播和反向传播用 fp16 加速计算、节省显存，但权重参数保持 fp32 以保证数值稳定性。
显存对比（7B 参数模型）：
fp32: 7B × 4 bytes = 28 GB（参数）+ 优化器状态 ≈ 84 GB
fp16: 7B × 2 bytes = 14 GB（参数）+ 优化器状态 ≈ 42 GB
混合精度可节省约 30%~50% 显存
Why — 为什么不用全 fp16？
问题一：fp16 的动态范围太窄
fp16 的指数位只有 5 位（vs fp32 的 8 位），最大值约 65504。大模型的激活值可能超过这个范围，导致溢出（inf）。
问题二：fp16 的精度不足
fp16 的尾数只有 10 位，对极小的梯度值（如 1e-7）无法精确表示，可能导致下溢（变为 0）。
问题三：权重更新必须在 fp32 下进行
梯度累加、动量更新等操作需要 fp32 精度。如果把权重也存为 fp16，累积误差会迅速放大，训练崩溃。
没有混合精度会发生什么？
- 显存占用翻倍，batch size 被迫减半，吞吐量大幅下降。
- 训练速度降低 30%~50%，大模型训练时间从几周变成几个月。
- 无法充分利用 Tensor Core 的 fp16 加速能力。
PyTorch 2.x 推荐使用 torch.cuda.amp.autocast 上下文管理器：
# 方式 1: 手动 autocast + GradScaler
scaler = torch.cuda.amp.GradScaler()for batch in dataloader:optimizer.zero_grad()with torch.cuda.amp.autocast():outputs = model(inputs)loss = loss_fn(outputs, targets)scaler.scale(loss).backward()scaler.unscale_(optimizer)torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)scaler.step(optimizer)scaler.update()注意：scaler.scale(loss).backward() 将损失放大以匹配 fp16 的动态范围；scaler.unscale_() 在梯度裁剪前把梯度恢复为 fp32 可表示的范围。
二、GradScaler：防止 fp16 梯度下溢
What — GradScaler 在做什么？
torch.cuda.amp.GradScaler 在反向传播前将损失乘以一个缩放因子（scale），放大梯度使其超出 fp16 的下溢阈值。在优化器更新前，再把梯度缩放回来。
Why — 为什么需要动态缩放？
不同 batch 的梯度量级不同。如果用固定缩放因子，大梯度的 batch 会溢出，小梯度的 batch 仍可能下溢。GradScaler 内部维护了一个动态缩放因子：如果某次反向传播没有 inf/nan，就逐渐增大缩放因子；如果有 inf/nan，就跳过该 batch 并减小缩放因子。
scaler = torch.cuda.amp.GradScaler(init_scale=2.**16,growth_interval=2000,growth_factor=2.0,backoff_factor=0.5,enabled=True,
)# 训练步
scaler.scale(loss).backward() # 1. 放大 loss → 反传放大梯度
scaler.unscale_(optimizer) # 2. 还原梯度到 fp32
torch.nn.utils.clip_grad_norm_(...) # 3. 在 fp32 下裁剪
scaler.step(optimizer) # 4. 检查 inf/nan，安全则更新
scaler.update() # 5. 调整缩放因子scaler.step() 会自动检查梯度中是否有 inf/nan。如果有，就跳过该 batch 的更新，并减小缩放因子。这是 AMP 安全性的关键保障 。
三、自动混合精度（AMP）的 PyTorch 实现
What — autocast 做了什么？
torch.cuda.amp.autocast 是一个上下文管理器。进入上下文后，PyTorch 自动选择每个运算的最合适精度：矩阵乘法、卷积等计算密集型操作用 fp16； reductions（如 softmax、layer_norm）等对精度敏感的操作保持 fp32。
Why — 为什么不是所有操作都用 fp16？
reductions 操作的输出值通常很小（如 softmax 输出 < 1），fp16 的精度不足会导致结果失真。PyTorch 的 autocast 有一个"白名单"，只有经过验证安全的运算才会用 fp16。
在 Ampere 架构（A100/3090）及以上 GPU 上，可以使用 bfloat16 替代 fp16。bfloat16 有和 fp32 相同的指数位（8 位），动态范围相同，只是精度降低。这意味着 不需要 GradScaler：
# Ampere+ GPU 可以用 bfloat16，无需 GradScaler
with torch.cuda.amp.autocast(dtype=torch.bfloat16):outputs = model(inputs)loss = loss_fn(outputs, targets)loss.backward() # 直接反传，不需要 scaler
optimizer.step()bfloat16 的 trade-off 是：动态范围与 fp32 相同（不会溢出），但尾数只有 7 位（精度较低）。对于大 batch 训练，精度损失通常可以接受 。
四、分布式训练概述 — 从单机到多机
What — 分布式训练解决什么问题？
分布式训练把数据和计算分散到多个 GPU / 多台机器上，实现：① 更大的有效 batch size，② 更快的训练速度，③ 训练超出单卡显存容量的大模型。
PyTorch 生态中主要有三种并行策略：
- DDP（DistributedDataParallel）：数据并行，每张卡有完整模型副本
- FSDP（Fully Sharded Data Parallel）：数据并行 + 参数/梯度/优化器状态分片
- Pipeline Parallelism：模型并行，每张卡负责模型的一部分层
Why — 为什么数据并行是默认选择？
数据并行的实现最简单：每张卡复制一份模型，输入不同 batch，做 All-Reduce 同步梯度。对于 Transformer 类模型（单层计算量均匀），数据并行的负载均衡最好。模型并行（流水线）需要手动切分层，实现复杂，且存在流水线气泡问题。
没有分布式训练会发生什么？
- 7B 模型需要约 28GB 仅存参数，加上优化器状态约 84GB，远超单卡（A100 80GB）能力。
- 大 batch 训练无法加速，训练时间线性增长。
- 无法利用多机多卡集群资源。
分布式训练的核心是梯度同步。PyTorch 通过 torch.distributed 提供通信原语：
import torch.distributed as dist# All-Reduce: 所有进程求和梯度后分发
dist.all_reduce(grad, op=dist.ReduceOp.SUM)
grad /= world_size # 取平均# 常用通信后端
# - NCCL: NVIDIA GPU 间通信，速度最快
# - Gloo: CPU 通信 / 跨平台All-Reduce 是 DDP 的核心操作。每张卡计算完自己的梯度后，通过 All-Reduce 把梯度做全局求和再除以卡数，得到全局平均梯度。
五、DDP（DistributedDataParallel）：数据并行
What — DDP 的工作原理？
DDP 把模型副本复制到每张卡上，每个进程处理不同的 batch。前向/反向传播在本地完成，然后通过 All-Reduce 同步梯度。同步完成后，所有卡的参数一致，下一轮继续。
Why — 为什么 DDP 比 DataParallel（DP）更好？
问题一：DP 的主线程瓶颈
PyTorch 的旧版 DataParallel 把所有 batch 先拉到主卡，再逐个分发。主卡成为瓶颈，无法充分利用多卡。
问题二：DDP 的每卡独立进程
DDP 每张卡一个独立进程，各自拉取数据、前向传播、反向传播，只同步梯度。通信和计算可以部分重叠，效率更高。
没有 DDP 会发生什么？
- DP 的主卡瓶颈导致加速比远低于卡数（4 卡可能只有 1.5x 加速）。
- 无法扩展到多机场景。
- Gradient Bucketing 等优化无法使用。
# 启动脚本（train.py）
import os
import torch
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDPdef setup():dist.init_process_group("nccl")local_rank = int(os.environ["LOCAL_RANK"])torch.cuda.set_device(local_rank)return local_rankdef main():local_rank = setup()model = MyModel().cuda()model = DDP(model, device_ids=[local_rank])optimizer = torch.optim.AdamW(model.parameters(), lr=3e-5)for batch in dataloader:loss = model(batch)loss.backward()optimizer.step()if __name__ == "__main__":main()# 命令行启动（4 卡）
# torchrun --nproc_per_node=4 train.py注意：DDP 要求每个进程有独立的 DataLoader 子集（通过 DistributedSampler 实现），否则不同卡的数据重复，等效 batch size 翻倍。
六、FSDP（Fully Sharded Data Parallel）：参数分片
What — FSDP 是什么？
FSDP 在 DDP 的基础上，将模型的参数、梯度和优化器状态在多个 GPU 之间分片（shard）。每张卡只持有参数的一个分片，而非完整副本。
对于 7B 参数模型，4 卡 FSDP 下每张卡只需存储 7B/4 = 1.75B 的参数 + 优化器状态，约 21 GB，远低于 DDP 的 84 GB。
Why — 为什么 DDP 不够用？
问题一：DDP 每个卡存完整参数
7B 模型的 fp32 参数占 28GB，AdamW 优化器状态（动量 + 二阶矩）再加 56GB，每卡总计约 84GB。A100-80GB 几乎放不下一个训练步。
问题二：大模型无法放到单卡上
175B 参数的模型（如 GPT-3）光参数就占 700GB（fp32），DDP 完全无法处理。
没有 FSDP 会发生什么？
- 模型无法装入单卡显存，训练无法开始。
- 必须减少 batch size，导致训练不稳定。
- 大模型训练只能在超大规模集群上使用 DeepSpeed/FSDP。
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecisionmp = MixedPrecision(param_dtype=torch.float16, reduce_dtype=torch.float32)model = MyModel().cuda()
model = FSDP(model, mixed_precision=mp)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-5)# 前向传播
with torch.cuda.amp.autocast():loss = model(batch)loss.backward()
optimizer.step()
# FSDP 自动处理参数 gather/scatter 和梯度分片FSDP 的核心优势在于参数只在需要时才 gather 到各卡，用完即释放。这大幅降低了峰值显存 。
七、DeepSpeed：封装 ZeRO 的高层框架
What — DeepSpeed 是什么？
DeepSpeed 是微软开发的高性能分布式训练框架，核心是 ZeRO（Zero Redundancy Optimizer）。ZeRO 把 DDP 中冗余的优化器状态、梯度和参数在多个 GPU 之间分片，消除了 DDP 的冗余存储。
ZeRO 有三个渐进式 Stage：
- ZeRO Stage 1：分片优化器状态（节省 4x 显存）
- ZeRO Stage 2：分片优化器状态 + 梯度（节省 8x 显存）
- ZeRO Stage 3：分片优化器状态 + 梯度 + 参数（节省 >16x 显存）
Why — 为什么 DeepSpeed 比纯 FSDP 更流行？
DeepSpeed 封装了 ZeRO 的所有 Stage，提供了更高层的 API（deepspeed.initialize），只需几行代码即可将普通训练脚本转换为分布式训练。它还内置了 CPU offload、NVMe offload、混合精度、梯度累积等生产级功能。
import deepspeedmodel, optimizer, _, _ = deepspeed.initialize(args=args,model=model,optimizer=optimizer,config="ds_config.json",
)# 之后的训练循环和单卡完全一样
for batch in dataloader:loss = model(batch)model.backward(loss)model.step()DeepSpeed 的 deepspeed.initialize 会自动包装模型和优化器，处理 ZeRO 分片、通信和混合精度。用户代码几乎不需要改动 。
混合精度与分布式训练速记
- 混合精度：fp16 计算 + fp32 权重 + GradScaler = 2~3x 加速 + 50% 显存节省
- bfloat16：无需 GradScaler，适合 Ampere+ GPU
- DDP：每卡完整模型副本，All-Reduce 同步梯度
- FSDP：参数分片，多卡共享参数量
- DeepSpeed ZeRO：三层渐进式分片，从 Stage 1 到 Stage 3
八、常见问题 FAQ（20 组）
FAQ 常见问题解答
Q1. 混合精度训练能提速多少？
通常 1.5x ~ 3x，取决于模型和 GPU。Ampere 架构（A100/3090）的 Tensor Core 对 fp16 有原生硬件加速，速度提升最明显。
Q2. fp16 和 bfloat16 有什么区别？
bfloat16 有相同的指数位（8 位）但更少的尾数位（7 vs 10）。bfloat16 的动态范围和 fp32 一致（不会溢出），但精度更低。bfloat16 不需要 GradScaler，但可能在某些任务上精度略低于 fp16。
Q3. GradScaler 什么时候会跳过更新？
当梯度中出现 inf 或 nan 时。GradScaler 会减小缩放因子并跳过该 batch。如果连续多次跳过，缩放因子会持续缩小，可能导致训练变慢。
Q4. DDP 和 DataParallel 有什么区别？
DDP 每卡一个独立进程，DP 单进程多线程。DDP 避免了主卡瓶颈，支持 Gradient Bucketing，性能更好，是多卡训练的事实标准。
Q5. DDP 的 effective batch size 怎么算？
有效 batch size = 单卡 batch size × GPU 数量。如果要保持相同的有效 batch size，卡数增加时单卡 batch size 应该等比减小。
Q6. DistributedSampler 是必须的吗？
是的，在 DDP 中必须使用。DistributedSampler 确保每个进程拿到不同的数据子集。如果不使用，所有卡会处理相同的 batch，相当于把 batch size 乘以卡数但梯度没有同步优化。
Q7. NCCL 和 Gloo 怎么选？
NCCL 用于 NVIDIA GPU，Gloo 用于 CPU 或非 NVIDIA GPU。NCCL 在 GPU 间通信速度远快于 Gloo。如果环境中没有 NCCL（如 AMD GPU），自动降级到 Gloo。
Q8. FSDP 和 DeepSpeed ZeRO Stage 3 有什么区别？
设计目标相同，实现不同。FSDP 是 PyTorch 原生实现，集成度更高；DeepSpeed ZeRO Stage 3 功能更丰富（CPU/NVMe offload），但需要额外的 DeepSpeed 依赖。
Q9. 什么是 ZeRO Offload？
把优化器状态放到 CPU 内存或 NVMe 上。CPU offload 用系统内存换显存，NVMe offload 进一步用磁盘换显存。代价是通信延迟增加，训练速度下降。
Q10. 梯度累积和分布式训练可以一起用吗？
可以。梯度累积是在多次前向/反向传播后才做一次参数更新，等效于增大 batch size。在 DDP 中，梯度累积的步数需要在所有卡上同步。
Q11. DDP 中的 find_unused_parameters 是什么？
当模型有部分参数未参与前向传播时，需要设为 True。如条件计算、动态图结构。设为 True 会增加通信开销，应尽量避免。
Q12. 为什么 AMP 有时会降低精度？
fp16 的精度只有 10 位尾数。对于数值范围大或梯度极小的场景，fp16 可能丢失精度。GradScaler 可以缓解但不能完全消除。
Q13. FSDP 的 wrap() 策略是什么？
决定哪些层被分片。可以按层（auto_wrap_policy）自动决定分片粒度，也可以手动用 fully_shard 包装特定子模块。
Q14. 分布式训练的 deadlock 怎么避免？
确保所有进程的调用顺序一致。使用 torchrun 启动确保进程同步。在训练循环中，所有进程必须执行相同数量的前后向传播步数。
Q15. torchrun 和 python -m torch.distributed.launch 有什么区别？
torchrun 是推荐的新方式，支持 elastic 启动。旧版 torch.distributed.launch 已废弃。torchrun 支持节点故障后自动恢复（elastic training）。
Q16. 分布式训练的 batch size 怎么调？
通常按卡数线性增大，同时按平方根比例增大学习率。这是线性缩放规则（Linear Scaling Rule）：batch size 翻倍，学习率翻倍。
Q17. 什么是流水线并行（Pipeline Parallelism）？
把模型的不同层放到不同 GPU 上。如 GPU 0 跑 Transformer 前 20 层，GPU 1 跑后 20 层。适合模型大到单卡连参数都放不下的场景，但存在流水线气泡开销。
Q18. 混合精度 + 分布式需要特殊配置吗？
需要，但 DeepSpeed 自动处理。手动配置时，AMP 的 autocast 在 DDP 内可以正常工作，只需确保 GradScaler 只在 rank 0 上更新。
Q19. 什么是梯度 bucketing？
DDP 将多个小梯度合并为一个大张量，一次性做 All-Reduce。这减少了通信次数，提高了带宽利用率。PyTorch DDP 默认开启 bucket_cap_mb=25。
Q20. 如何选择 DDP、FSDP 还是 DeepSpeed？
模型能装下选 DDP，装不下选 FSDP 或 DeepSpeed ZeRO Stage 2/3。DeepSpeed 提供最多功能（offload、梯度累积、混合精度一体化），但引入外部依赖。FSDP 是 PyTorch 原生方案，集成更简单。
全篇总纲
- 1 个加速公式：混合精度 = 2~3x 加速 + 50% 显存节省
- 1 个安全网：GradScaler 防止 fp16 下溢
- 1 个并行基准：DDP 是数据并行的最小可用方案
- 1 个容量突破：FSDP / DeepSpeed ZeRO 让超大模型可训练
- 2 个选型原则：能 DDP 就不 FSDP，能原生就不第三方
九、Roadmap 后续预告
后续预告
- 第 10 篇：大模型推理与部署 — 从 PyTorch 到 ONNX / TensorRT / vLLM