当前位置：首页 > news >正文

机器学习数学基础专题【左扬精讲】—— 混合精度与分布式训练：从单机单卡到 DDP、FSDP、DeepSpeed

news 2026/7/13 19:46:43

机器学习数学基础专题【左扬精讲】—— 混合精度与分布式训练：从单机单卡到 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 加速能力。

How — PyTorch AMP 的两种使用方式

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 并减小缩放因子。

How — GradScaler 的关键操作

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。

How — autocast 与 bfloat16

在 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 训练无法加速，训练时间线性增长。
无法利用多机多卡集群资源。

How — 分布式训练的基本通信原语

分布式训练的核心是梯度同步。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 等优化无法使用。

How — DDP 的标准启动方式

# 启动脚本（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。

How — 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、混合精度、梯度累积等生产级功能。

How — DeepSpeed 零代码改造

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 分片、通信和混合精度。用户代码几乎不需要改动。

混合精度与分布式训练速记