PyTorch FSDP训练指南：wrap策略、激活重计算与混合精度配置

在《Ai大模型训练教程：从入门到实战落地的系统课程》里，FSDP（Fully Sharded Data Parallel）通常是“把单机/单卡能跑的模型，稳定扩到多卡并显著省显存”的关键拼图。本篇聚焦三个最影响训练效果与稳定性的实操点：wrap 策略（如何切分分片单元）、激活重计算（checkpointing）、混合精度（AMP / FSDP MixedPrecision）。目标是让你能按步骤把一个 Transformer 类模型在多卡上稳定跑起来，并知道为什么这么配。

默认读者：已经会用 PyTorch DDP/单卡训练，想迁移到 FSDP；训练环境为 torchrun 多进程多 GPU。

一、FSDP你需要先掌握的3个概念

1. 参数分片（Sharding）到底省了什么

FSDP会把模型参数、梯度、优化器状态按 rank 分片存放，从而让每张卡只保存一部分。前向/反向时通过 all-gather / reduce-scatter 临时聚合需要的参数或梯度。

省显存的主要来源：

• 参数：由“每卡一份”变成“每卡 1/N 份”。
• 梯度：同理。
• 优化器状态（如Adam的m/v）：可配合 shard 或 CPU offload 进一步降低峰值。

2. wrap 的粒度决定通信频率与峰值显存

FSDP并不是“包住整个model就完事”，你需要决定“哪些子模块作为一个 FSDP 单元进行分片”。

• wrap 粒度越细：每层都单独分片，显存峰值更低，但 all-gather 次数更多，通信开销更高。
• wrap 粒度越粗：通信次数少，但某些阶段需要聚合的参数块更大，可能导致峰值显存上升甚至 OOM。

3. 混合精度与激活重计算是“省显存/提吞吐”的两大杠杆

• 混合精度：参数/梯度/通信使用 BF16/FP16，通常大幅提升吞吐并降低显存。
• 激活重计算：把前向保存的激活丢掉，在反向时重算，用更多算力换显存。

二、训练骨架：最小可用 FSDP 启动方式

1) torchrun 启动

你需要使用 torchrun 启动多进程：

• 单机 8 卡示例：

命令建议：

• torchrun --nproc_per_node=8 train_fsdp.py ...

关键点：每个进程绑定一张GPU，进程间通过 NCCL 通信。

2) 初始化分布式与设置设备

训练脚本最基础逻辑：

• dist.init_process_group(backend="nccl")
• local_rank = int(os.environ["LOCAL_RANK"])
• torch.cuda.set_device(local_rank)

3) 构建模型后再 wrap

一般流程：

1. build model（尽量在 GPU 上）
2. 应用 wrap 策略生成 FSDP model
3. 构建 optimizer（注意：optimizer 需要拿到 FSDP 包装后的参数）
4. 训练循环

三、wrap策略：怎么选、怎么写、怎么验证

wrap 策略的目标是：在显存、速度、通信之间取得你能接受的平衡。常见有三种：

3.1 全模型一把包（最简单，但不一定最稳）

把整个模型作为一个 FSDP 单元：

• 优点：代码最少，通信次数相对少。
• 缺点：峰值显存可能较高，尤其在大模型/长序列/大 batch 时。

适用场景：

• 模型不算特别大（例如中等规模），主要目的是从 DDP 省掉优化器状态和参数副本。

3.2 按 Transformer Block 包（推荐的默认策略）

对 Transformer 来说，最常用的策略是：每个 Block（或每若干个 Block）一个 FSDP 单元。这样：

• 分片粒度适中
• 通信次数可控
• 峰值显存更低、更稳定

实操做法：使用 auto_wrap_policy

PyTorch 提供了基于模块类型的自动 wrap：

• 你指定一个 transformer_layer_cls（例如 TransformerEncoderLayer、自定义 DecoderLayer 等）
• 自动对这些层进行 wrap

你需要做的事情：

1. 找到你模型中“重复堆叠的基本层”类名（例如 TransformerBlock）
2. 将该类传给 auto wrap policy

如何判断是否 wrap 生效

建议在 rank0 打印模型结构，检查每个 block 外层是否出现 FSDP 包装。
另外也可以统计 FSDP 单元数量：

• 单元太少：显存可能高
• 单元太多：通信过密，性能下降

3.3 基于参数数量的 wrap（按大小切块）

当模型结构不规则（不是标准 Transformer 堆叠）时，可以按“子模块参数量阈值”自动 wrap：

• 大于阈值的模块才 wrap
• 小模块不 wrap

经验阈值：

• 先从 1e7（1000万参数）左右试起
• 若显存仍高就降低阈值；若通信过重就提高阈值

3.4 wrap的经验法则（能直接用）

1. 优先按 Transformer Block wrap：最稳、最可解释。

2. embedding、lm_head 是否单独 wrap：

   • 词表很大（如 50k-250k）时，embedding/lm_head 参数巨大，建议单独 wrap 或与相邻层拆开，降低峰值。

3. 不要过度细粒度到每个 Linear：

   • 会导致 all-gather 次数爆炸，吞吐明显变差。

4. 先跑通再调优：

   • 第一步目标是“不 OOM + loss 正常下降”，第二步才是“吞吐/利用率”。

四、激活重计算（Activation Checkpointing）：放哪儿、怎么放、放多少

激活重计算是 FSDP 训练大模型时最常见的“救命配置”。它对显存的影响通常比你微调 wrap 更直接。

4.1 原理与代价

• 正常训练：前向会保存各层激活用于反向，显存占用大。
• 重计算：前向不保存（或少保存）激活，反向时重新做一次前向计算来恢复激活。

代价：

• 训练速度变慢（常见 10%~40% 不等，取决于层结构与重算比例）
• 但能显著降低峰值显存

4.2 最推荐的放置点：每个 Transformer Block

经验上，把 checkpoint 放在每个 block 上是最常见、收益最大的做法。

两种常见策略：

1. 全量重算：所有 block 都 checkpoint

   • 显存最低
   • 速度损失最大

2. 间隔重算：每 2~4 个 block checkpoint 一次

   • 显存与速度折中

4.3 PyTorch中常用的实现方式

实现路径一般有两类：

• 使用 torch.utils.checkpoint.checkpoint 包裹 block 的 forward
• 或使用 FSDP 相关的 checkpoint wrapper（例如对模块进行 checkpoint wrapping）

实操建议：

• 优先使用“模块级 checkpoint wrapper”，更干净、可维护；
• 如果你只想快速验证，直接在 block forward 里用 checkpoint(fn, *args) 也可行。

4.4 checkpoint 与 FSDP 的组合注意事项

1. 确保输入输出是 Tensor 且 requires_grad 逻辑正确：checkpoint 对非 Tensor 输入支持有限。
2. 使用 checkpoint 时，某些 inplace 操作（如 inplace dropout、inplace residual）可能导致 autograd 问题；建议避免或改成非 inplace。
3. 若你开启了混合精度，checkpoint 内部重算的 dtype 要与外部一致（通常由 autocast 控制）。

4.5 何时需要开 checkpoint（判断标准）

• 你的 OOM 出现在 forward 或 backward 中期，且 batch/seq_len 无法再降：优先开 checkpoint。
• 你已经 FSDP 分片了参数但还是 OOM：说明主要瓶颈可能在激活，而非参数副本。

五、混合精度配置：FSDP MixedPrecision 怎么选 BF16/FP16

5.1 BF16 vs FP16 的选择建议

• 优先 BF16（如果你的 GPU 支持，如 A100/H100/部分 RTX 新卡）：

  • 数值范围更大，更不容易溢出
  • 通常不需要 loss scaling 或只需要很轻量的处理

• FP16：

  • 也能提速省显存，但更容易出现 overflow/NaN
  • 常需要 GradScaler

结论：能用 BF16 就用 BF16，除非你有明确原因必须 FP16。

5.2 FSDP MixedPrecision 的核心点

FSDP 混合精度不是只有 autocast 一种方式。更常见的组合是：

• torch.autocast 控制算子执行精度（前向/部分反向）

• FSDP 的 MixedPrecision 控制：

  • 参数 dtype（param_dtype）
  • 梯度 reduce/scatter 的 dtype（reduce_dtype）
  • buffer dtype（buffer_dtype，比如 LayerNorm 的 buffer）

实操推荐组合：

• BF16 训练：

  • param_dtype = bf16
  • reduce_dtype = bf16
  • buffer_dtype = bf16 或 fp32（视稳定性）

• FP16 训练：

  • param_dtype = fp16
  • reduce_dtype = fp16
  • buffer_dtype = fp16
  • 额外配合 GradScaler

5.3 LayerNorm / RMSNorm 要不要 FP32？

在一些模型里，norm 层用 fp32 能提升稳定性，代价是略增显存/耗时。

经验建议：

• 训练不稳定（loss 抖动大、出现 NaN）：尝试让 norm 计算保持 fp32。
• 训练稳定：全 bf16 通常更快。

5.4 一个可落地的“稳定优先”配置清单

当你第一次把模型迁到 FSDP，建议按以下顺序启用：

1. 先用 BF16（若支持）
2. 开启 activation checkpoint（至少按 block）
3. wrap 按 block

4. 若仍不稳定：

   • norm fp32
   • 降低学习率或启用 gradient clipping

六、把三者组合起来：一套可复制的配置思路

下面给出一个“从可跑到更快”的调参路线，你可以在实际训练中按阶段推进。

6.1 阶段A：先跑通（稳定第一）

目标：loss 正常下降、不 OOM。

• wrap：按 Transformer Block
• 激活重计算：所有 Block checkpoint
• 混合精度：BF16（或 FP16 + GradScaler）
• batch：先用较小 micro-batch，配合 gradient accumulation 达到目标 global batch
• 额外建议：开启 gradient clipping（如 1.0）避免偶发爆炸

6.2 阶段B：降低通信开销（提吞吐）

目标：提高 tokens/s。

• wrap：尝试“每 2 个 block wrap 一次”（粒度稍变粗）或调大按参数量阈值
• checkpoint：改为间隔重算（每 2~4 层一次）
• 混合精度：保持 BF16；检查是否可全 bf16 buffer

6.3 阶段C：进一步压显存（把模型做大或序列变长）

目标：在现有硬件上把 seq_len 或模型规模推上去。

• wrap：更细粒度（每个 block 独立 wrap，embedding/lm_head 单独 wrap）
• checkpoint：全量重算
• 如仍不够：考虑 CPU offload（参数或优化器状态），但吞吐会明显下降

七、常见问题排查（实战中最常见的坑）

7.1 训练变慢很多：是 wrap 太细还是 checkpoint 太多？

排查顺序：

1. 先关掉 checkpoint，看吞吐恢复多少；
2. 再把 wrap 粒度变粗（例如从每层 wrap -> 每个 block wrap）；
3. 观察 NCCL 通信占比（可用 profiler）；如果通信时间占比很高，说明 wrap 太细或 all-gather 太频繁。

7.2 出现 NaN / Inf：先查混合精度与学习率

常用处理：

• BF16：通常先检查学习率是否过大、是否需要 gradient clipping。
• FP16：启用 GradScaler；必要时降低 loss scale 或改 BF16。
• 对 norm 层保持 fp32（尤其是深层 Transformer）。

7.3 OOM 仍然发生：优先怀疑激活而非参数

行动建议：

1. 开启/加大 activation checkpoint 覆盖范围
2. 降低 micro-batch 或 seq_len
3. 把 embedding/lm_head 单独 wrap
4. 检查是否有额外缓存（如保存 logits、保存所有层 hidden states）

7.4 优化器与参数绑定错误

症状：loss 不下降或梯度不更新。
要点：

• optimizer 一定要基于 FSDP 包装后的 model.parameters() 创建。
• 不要在 wrap 前创建 optimizer。

八、一个“可直接照着做”的落地示例（配置层面）

这里给出一个不依赖你具体模型代码的配置清单，你可以把它翻译成自己的训练脚本：

8.1 推荐默认（多数大模型微调/预训练都适用）

• wrap：Transformer Block 级别 auto wrap
• activation checkpoint：每个 block
• mixed precision：BF16（param/reduce/buffer 尽量 bf16）
• 梯度：gradient clipping=1.0
• batch：micro-batch 尽量小 + accumulation

8.2 性能优先（显存还有余量）

• wrap：每 2~4 个 block wrap
• checkpoint：间隔重算
• mixed precision：BF16

8.3 显存极限（单卡显存吃紧）

• wrap：每个 block wrap + embedding/lm_head 单独 wrap
• checkpoint：全量重算
• mixed precision：BF16 + 必要时 norm fp32
• 必要时：optimizer state offload（代价是速度）

九、小结：三件事的最短决策路径

• wrap 策略：默认按 Transformer Block；显存不够就更细，速度太慢就更粗。
• 激活重计算：OOM 先开它；通常放在每个 Block 最有效。
• 混合精度：优先 BF16；不稳定再考虑 norm fp32、clip、降 LR；FP16 记得 GradScaler。

这三项组合好，FSDP 才真正从“能用”变成“可落地、可扩展、可复现”的大模型训练方案。下一步在系列《Ai大模型训练教程》中，你可以继续把 attention 实现（如 FlashAttention）、数据管线（tokenizer/packing）、以及 profiler 结合进来，把 tokens/s 与显存利用率进一步拉满。