PyTorch 训练框架内核深度解析

第18章 FSDP：ZeRO 风格的参数分片

作者 杨艺韬 · 14,322 字 · 发布于 · 更新于

"FSDP shards module parameters, gradients, and optimizer states across data parallel workers to save memory at the cost of communication."

—— torch.distributed.fsdp.fully_shard 的 docstring（torch/distributed/fsdp/_fully_shard/_fully_shard.py:101-103）

本章要点

  • DDP 在 70B 模型时崩溃：每张卡需要 280GB（params） + 280GB（grad） + 560GB（Adam state）= 1120GB，单 H100 才 80GB
  • FSDP 灵感来自 DeepSpeed ZeRO：把 params / grads / optimizer state 都切到 N 张卡，每张卡只持有 1/N
  • 核心机制：forward 时 AllGather 把分片重建成完整 param、用完后立即 reshard；backward 时同理 + 最后 ReduceScatter 同步梯度
  • 5 种 ShardingStrategy：NO_SHARD（=DDP）、SHARD_GRAD_OP（ZeRO-2）、FULL_SHARD（ZeRO-3）、HYBRID_SHARD、_HYBRID_SHARD_ZERO2
  • FSDP-1 vs FSDP-2：v2.3 以 torch.distributed._composable.fsdp.fully_shard 引入、v2.6 移到公开路径 torch.distributed.fsdp.fully_shard，模块化 + 与 torch.compile 兼容性好得多
  • prefetch 是性能命脉：在 layer N forward 时 prefetch layer N+1 的 unshard，隐藏通信延迟

18.1 为什么 DDP 不够用

第 17 章 DDP 假设"每张卡放得下完整模型"。70B 模型：

  • params：70B × 4B(fp32) = 280 GB
  • grads：280 GB
  • Adam optimizer state (m + v)：560 GB
  • 加上 activations 几十 GB
  • 合计 1120+ GB

H100 80GB，新一代 Blackwell 单卡也不到 200GB —— 单卡都装不下。这是 DDP 在大模型时代的天花板。

FSDP（Fully Sharded Data Parallel）的解法：把这堆显存按 rank 数 N 均匀切分到 N 张卡。每张卡只持有 1/N 份数据，需要时通过通信临时凑齐完整张量。

8 卡 FSDP:
  每张卡 params: 35 GB
  每张卡 grads:  35 GB
  每张卡 optimizer: 70 GB
  合计 ~140 GB / 卡  ← 80GB 仍然装不下
16 卡:  ~70 GB / 卡  ← 接近但还要 activation
32 卡:  ~35 GB / 卡  ← 舒服了

70B 模型至少 32 张卡才能跑稳。FSDP 的"分片粒度"是这套机制的核心。

18.2 5 种 ShardingStrategy

graph TB
    DDP["DDP / NO_SHARD<br/>每卡完整 params + grads + optimizer"]
    Z2["ZeRO-2 / SHARD_GRAD_OP<br/>fw 后不 reshard（计算期 params 完整）<br/>grads + optimizer 切片"]
    Z3["ZeRO-3 / FULL_SHARD<br/>params + grads + optimizer 全切片<br/>显存最省"]
    HZ["HYBRID_SHARD<br/>同节点内 ZeRO-3，跨节点 DDP<br/>跨节点带宽不足时的折中"]
    DDP -.显存大.-> Z2
    Z2 -.通信增多.-> Z3
    Z3 -.通信占比可能太高.-> HZ
    style Z3 fill:#dcfce7,stroke:#22c55e,stroke-width:2px
    style HZ fill:#fef3c7,stroke:#f59e0b

通信代价：

Strategy params 通信 grads 通信 optimizer 通信 显存 / 卡
NO_SHARD (DDP) 0 AllReduce 0 完整
SHARD_GRAD_OP 1× AllGather（只在 fw；fw 后不 reshard，所以 bw 省掉那次） ReduceScatter 0 计算期 params 完整；grads + opt 为 1/N
FULL_SHARD 2× AllGather (fw + bw) ReduceScatter 0（params 已分片） 全部 1/N
HYBRID_SHARD 节点内 AllGather 节点内 RS + 跨节点 AR - 1/N (节点内)

FULL_SHARD 显存最省、通信最多。按 ZeRO 论文（SC'20）的通信量分析，以参数量 Ψ 计：DDP 的一次 AllReduce 约 2Ψ；ZeRO-2 的 ReduceScatter + AllGather 合计仍是 2Ψ；ZeRO-3 在前向、反向各 AllGather 一次参数、再 ReduceScatter 一次梯度，共约 3Ψ，即 DDP 的 1.5 倍。生产代码里 70B+ 通常用 FULL_SHARD 或 HYBRID_SHARD。

18.3 unshard / reshard 的精确时机

FSDP 把参数按 N 切片后，每张卡平时只持有 1/N。但算子需要完整 param 才能算（Linear 需要完整 weight 矩阵做 GEMM）。所以 forward 前要 AllGather 把分片凑成完整 param，算完立即 reshard（释放完整副本，回到 1/N）：

sequenceDiagram
    autonumber
    participant Layer as Layer N forward
    participant Comm as AllGather
    participant GPU as GPU 显存
    Layer->>Comm: 发起 AllGather (异步)
    Note over GPU: 此时 GPU 上 layer N 的 param 是分片
    Comm->>GPU: AllGather 完成 → 完整 param 在显存
    Layer->>GPU: 跑 layer N 的 GEMM
    GPU->>GPU: reshard: 立即释放完整 param
    Note over GPU: 回到 1/N 状态, 等下次 forward

backward 同理：算每层 grad 前再 AllGather 一次 param，算完 reshard，最后用 ReduceScatter 把这层的 grad 切片同步：

Layer N backward:
  1. AllGather param (因为反向也要用 weight 算 grad_input)
  2. 计算 grad_input + grad_param
  3. ReduceScatter grad_param: 8 卡各自只保留自己的 1/N 切片
  4. 释放完整 param

ReduceScatter 是 AllReduce 的"半成品" —— 每个 rank 只拿到 reduce 结果的 1/N 切片，不再 AllGather。这刚好对应 FSDP 的需求：每个 rank 只更新自己持有的 1/N 参数。

18.4 prefetch：隐藏通信延迟

朴素实现里：

fw layer 1 = AllGather + GEMM + reshard (串行)
fw layer 2 = AllGather + GEMM + reshard
...

每层都等 AllGather 完成才开始算 —— 通信延迟全暴露。prefetch 优化：

fw layer 1: AllGather_1 → GEMM_1 (期间 prefetch AllGather_2)
fw layer 2: GEMM_2 (用已经 prefetch 完的 param) → prefetch AllGather_3
...

每层 GEMM 进行时，下一层的 AllGather 在另一个 stream 上跑。理想情况下通信完全 overlap 到计算里，FSDP 性能接近 DDP（如果带宽足够）。

实现上 FSDP 用 dedicated CUDA stream（_unshard_stream、_pre_unshard_stream 在 torch/distributed/fsdp/_runtime_utils.py:263-269）跑 collective，与训练主 stream 并发。

FSDP-1 的 prefetch 行为由 forward_prefetch=True/False 与 backward 的 backward_prefetch=BACKWARD_PRE/POST 控制。BACKWARD_PRE 在当前层梯度计算之前就 prefetch 下一组参数；BACKWARD_POST 在当前层梯度计算之后才 prefetch（显存更省但 overlap 少，见 torch/distributed/fsdp/api.py:72 的 BackwardPrefetch docstring）。

18.5 FSDP-1 vs FSDP-2

PyTorch 推出了新一代 FSDP：fully_shard（也叫 FSDP-2）。它在 v2.3 以 prototype 形态落在 torch.distributed._composable.fsdp（PR #117776 / #118525），v2.6 移到公开路径 torch.distributed.fsdp（PR #141868）；v2.11 的源码在 torch/distributed/fsdp/_fully_shard/。它解决了 FSDP-1 的几个痛点：

1. 模块化：FSDP-1 把整个模型 wrap 成一个 FullyShardedDataParallel(model)，所有逻辑放一个大类（fully_sharded_data_parallel.py 单文件 2167 行）。FSDP-2 用 fully_shard(submodule) API 给每个 submodule 单独 wrap，更精细控制：

# FSDP-2
from torch.distributed.fsdp import fully_shard

for layer in model.layers:
    fully_shard(layer)            # 每个 transformer 层单独 shard
fully_shard(model)                # root module

2. 与 torch.compile 兼容：Dynamo 会跳过全部 torch.distributed.fsdp 的 FSDP-1 代码、在 FSDP 模块边界 graph break，且只接受 use_orig_params=True（torch/distributed/fsdp/_dynamo_utils.py 的 note、torch/_dynamo/variables/builder.py:2065-2071）。FSDP-2 重新设计了参数管理，hook 可以被 trace 进图（默认仍在图外 eager 执行，见 §18.19）。

3. DTensor 后端：FSDP-2 内部用 DTensor（distributed tensor）抽象，让张量分片成为 first-class concept。

生产代码里 v2.6+ 推荐用 FSDP-2。FSDP-1 仍然保留在 torch.distributed.fsdp.FullyShardedDataParallel，但新功能（如 set_custom_all_gather / set_custom_reduce_scatter 自定义 collective、shard_placement_fn 按参数指定分片维）只在 FSDP-2 上加。

18.6 mixed precision：FSDP 的另一大优化

FSDP 提供专门的 MixedPrecision 配置：

from torch.distributed.fsdp import MixedPrecision

mp_policy = MixedPrecision(
    param_dtype=torch.bfloat16,         # AllGather 时用 bf16 通信
    reduce_dtype=torch.float32,         # ReduceScatter 时用 fp32 (避免梯度精度损失)
    buffer_dtype=torch.bfloat16,
)

model = FSDP(model, mixed_precision=mp_policy)

上面是 FSDP-1 的 MixedPrecision（torch/distributed/fsdp/api.py:112，字段 param_dtype / reduce_dtype / buffer_dtype）。FSDP-2 对应的是 MixedPrecisionPolicy（_fully_shard/_fsdp_api.py:15），字段是 param_dtype / reduce_dtype / output_dtype / cast_forward_inputs —— 没有 buffer_dtype，通过 mp_policy= 传给 fully_shard。

为什么这个比朴素 model.bfloat16() 好？因为 FSDP 控制每种张量的 dtype 独立：

  • params 在 bf16（省一半 AllGather 带宽）
  • gradients 在 fp32 reduce（保证数值稳定）
  • master weights 仍是 fp32（让 Adam 更新精确）

这套"分别控制"是大模型训练 fp32+bf16 混合精度的标配。第 20 章量化与混合精度会展开。

18.7 DeviceMesh：拓扑的 first-class 表达

FSDP-2 之前，分布式训练靠 ProcessGroup 表达"哪些 rank 参与同一通信"。混合并行（DP + TP + PP）时要手工管理多个 group，容易出错。

torch.distributed.device_mesh（1553 行）引入 DeviceMesh 抽象：把 N 个 rank 排成 K 维网格，每维有自己的 ProcessGroup。

from torch.distributed.device_mesh import init_device_mesh

# 32 卡 = 4 节点 × 8 卡
# 节点内 8 卡 ZeRO-3, 跨节点 4 个 replica DDP
mesh_2d = init_device_mesh(
    device_type="cuda",
    mesh_shape=(4, 8),
    mesh_dim_names=("replica", "shard"),
)

# 取出某一维的 ProcessGroup
shard_pg = mesh_2d["shard"].get_group()    # 节点内 8 卡 group
replica_pg = mesh_2d["replica"].get_group() # 跨节点 4 卡 group

每个 rank 在 mesh 里有 K 个坐标（如 (replica=2, shard=5)）。在某一维上的 collective（如 shard.all_gather）只在那一维的 group 内做，不涉及其他维。

DeviceMesh 让 HSDP / 3D parallel（DP + TP + PP）的实现从"手动管 N 个 group"变成"声明一个 mesh 然后取维度"。这是 FSDP-2 / DTensor 的基础。

18.8 DTensor：分片张量的 first-class 类型

FSDP-2 内部不直接持有 plain Tensor，而是 DTensor（distributed tensor）。一个 DTensor 由三部分组成：

from torch.distributed.tensor import DTensor, Shard, Replicate

dtensor = DTensor.from_local(
    local_tensor,         # 本 rank 持有的切片
    device_mesh,          # 它在哪个 mesh 上
    placements=[Shard(0)] # 在 mesh 维度上怎么分布: Shard / Replicate / Partial
)

三种 Placement：

  • Shard(dim)：按某 dim 切到 mesh 上，每 rank 持有 1/N
  • Replicate()：每 rank 都有完整副本
  • Partial：每 rank 持有"部分和"，下次访问前需 reduce

DTensor 算子（dtensor + dtensor）会自动选最优通信策略：如果两边都是 Shard(0)，相加无需通信；如果一个 Shard(0) 一个 Replicate，会把 Replicate 那边就地切成 Shard(0) 再加，结果仍是 Shard(0)，同样不通信（本机 gloo 两进程实测，CommDebugMode 记录的 collective 数为 0）。需要通信时（如显式 redistribute 到 Replicate()）才发 collective。这套自动 dispatch 让用户不用手写 collective。

FSDP-2 把每个 nn.Parameter 包成 DTensor(local_shard, mesh, [Shard(0)]) 作为常驻的分片表示；但 unshard 并不走 DTensor.redistribute：它把同一 param group 里所有参数的本地分片 copy-in 到一块扁平 buffer，发一次 AllGather（torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py 的 foreach_all_gather），再 copy-out 成普通的 unsharded nn.Parameter 挂回 module（torch/distributed/fsdp/_fully_shard/_fsdp_param.py 的 to_unsharded）。本机 gloo 实测：forward 内部看到的 weight 是 Parameter、形状完整，forward 结束后又变回 DTensor。DTensor 负责表示"分片后是什么"，调度与通信由 FSDP-2 自己完成。

18.9 HSDP 的具体拓扑

HYBRID_SHARD 的关键是用 2D mesh：节点内 shard、跨节点 replicate：

mesh = init_device_mesh("cuda", (2, 8), mesh_dim_names=("inter", "intra"))

# 配 FSDP-2: 节点内 (intra) shard, 跨节点 (inter) replicate
fully_shard(model, mesh=mesh)

通信模式：

  • forward 的 AllGather 只在 intra 维（节点内），走 NVLink（带宽以百 GB/s 计，具体值随卡与代际差别很大）
  • backward 的 ReduceScatter 也在 intra 维
  • backward 完成后额外做一次 AllReduce 在 inter 维（跨节点）同步梯度

这条策略让"高带宽节点内通信"承担分片代价、"低带宽跨节点"只做一次梯度同步。在节点内带宽远大于跨节点的集群上，HSDP 比纯 FULL_SHARD 更快 —— 具体幅度取决于节点内外带宽比与模型规模，须自己 profile。

18.10 _unshard 的 stream 编排

打开 torch/distributed/fsdp/_runtime_utils.py:277 的 _unshard（节选：删去类型注解、docstring 与 rate limiter 外层的 record_function，注释为本专栏所加）：

def _unshard(state, handle, unshard_stream, pre_unshard_stream) -> None:
    if not handle:
        return
    with state._device_handle.stream(pre_unshard_stream):
        # 1. 在 pre-unshard stream 上做准备（mixed precision 的低精度 cast / CPU offload 的 H2D 拷贝）
        ran_pre_unshard = handle.pre_unshard()
    if ran_pre_unshard:
        # 2. wait_stream 让 unshard stream 等 pre-unshard 完成
        unshard_stream.wait_stream(pre_unshard_stream)
    if state.limit_all_gathers:
        event = state._free_event_queue.dequeue_if_needed()
        if event:
            event.synchronize()          # rate limiter: 限制在飞的 AllGather 数量
    with state._device_handle.stream(unshard_stream):
        handle.unshard()                 # 3. 在 unshard stream 上分配 padded buffer 并发起 AllGather
        handle.post_unshard()

两个 stream 的分工：

  • pre_unshard_stream：跑 mixed precision 的低精度 cast、CPU offload 的 H2D 拷贝等"准备工作"
  • unshard_stream：分配 AllGather 的目标 buffer，并跑 AllGather 本身

（分工以 torch/distributed/fsdp/_runtime_utils.py:261-269 建 stream 处的注释为准。）为什么不放在主 stream？因为那样会阻塞下一层 forward；独立 stream 让这些工作与主 stream 上的计算 overlap。

中间那段 limit_all_gathers 分支是 rate limiter：靠 _free_event_queue（_limiter_utils.py）限制同时在飞的 unsharded buffer 数量，避免 prefetch 过深把显存打爆。

wait_stream 是 CUDA 的 stream-to-stream 同步原语：让 unshard stream 等 pre-unshard stream 上之前发的 op 完成。这种"多 stream 协同"是 FSDP 性能的关键 —— 配合 prefetch（§18.4），让通信完全隐藏到计算里。

18.11 activation_checkpoint × FSDP 的协作

第 7 章 §7.5.4 讲过 activation_checkpoint：前向不保存中间激活，反向时重新 forward 一遍取回。这套机制配 FSDP 用时多一层复杂度。

考虑 FSDP+checkpoint 的反向流程：

  1. 正常 forward（FSDP layer N）：
    • AllGather 凑完整 param → forward 计算（不存 activation）→ reshard 释放完整 param
  2. 反向走到 layer N：
    • 需要重新 forward 取 activation → 但 param 已经被 reshard！
    • 必须 再次 AllGather 凑完整 param → 重 forward → 拿到 activation → backward
    • 完成后又 reshard
sequenceDiagram
    autonumber
    participant FW as Forward layer N
    participant Mem as 显存
    participant BW as Backward layer N
    FW->>Mem: AllGather param (1)
    FW->>FW: 跑 forward, 不存 activation
    FW->>Mem: reshard param (释放完整副本)
    Note over BW: 反向到 layer N
    BW->>Mem: AllGather param (2) ← pre-backward hook 里这一次
    BW->>BW: 重 forward 拿到 activation（复用这份完整 param，不再 unshard）
    BW->>BW: backward 算梯度
    BW->>Mem: reshard param
    BW->>Mem: ReduceScatter grad

这里要破掉一个很常见的误解：重 forward 并不会再多发一次 AllGather。_runtime_utils.py 在 pre-forward 与 post-forward 两处都有同一条守卫（:373-380 与 :468-472）——注释写的就是 "For fully_shard + checkpoint, skip pre-forward logic in the recomputed forward"，条件是 handle._training_state == HandleTrainingState.BACKWARD_PRE。也就是说重 forward 既不 unshard 也不 reshard，直接复用 pre-backward hook 刚 AllGather 出来的那份完整 param。

FSDP-2 在 torch/distributed/fsdp/_fully_shard/_fsdp_state.py:246-255 与 :276-279 有同样的守卫（TrainingState.PRE_BACKWARD 时跳过；只有嵌套 FSDP 且参数已被先前的 post_backward reshard 这种边角情形才补一次 unshard）；本机 gloo 两进程实测 apply_activation_checkpointing + fully_shard，3 个 block 开不开 checkpoint 都是 forward 3 次、backward 3 次 AllGather。

所以一个 checkpointed layer 每步仍然是 forward 1 次 + backward 1 次 = 2 次 AllGather，和不开 checkpoint 的 FULL_SHARD 一样。activation checkpointing 换的是"算力换 activation 显存"，不是"通信换显存" —— 它多出来的成本是重跑一遍 forward 的 FLOPs。

torch/distributed/algorithms/_checkpoint/checkpoint_wrapper.py:112 的 CheckpointWrapper 是 PyTorch 给 FSDP 的官方 checkpoint 接口。它的关键设计：

  • apply_activation_checkpointing(model, ...) (:239) 递归给每个匹配 check_fn 的 module 套上 CheckpointWrapper
  • checkpoint_wrapper(module, checkpoint_impl=...)（:196）的默认值是 CheckpointImpl.NO_REENTRANT，即非重入实现：用 saved_tensors_hooks（第 7 章 §7.5.3）实现，与 FSDP / torch.compile 兼容性好
  • 老的重入实现（use_reentrant=True）用 torch.autograd.Function（第 7 章 §7.8），与 FSDP 反向 hook 互动有 corner case，新代码避开
  • 注意底层的 torch.utils.checkpoint.checkpoint 至今没有把默认切成 False：use_reentrant 默认 None，走告警后仍回落到 True（torch/utils/checkpoint.py:480-490）。自己直接调 checkpoint() 时务必显式传 use_reentrant=False

实战决策：开 FSDP 后是否再加 activation_checkpoint？看显存。FSDP 只管 params / grads / optimizer state，activation 那一块它一点都省不到 —— 大 batch、长序列下这部分能吃掉相当可观的显存。activation_checkpoint 按 block 粒度只保留边界处的激活、其余重算，省下的比例取决于 block 内部有多少中间张量。代价是每个 checkpointed block 的 forward 要多跑一遍（通信次数不变，见上）。GPU 时间比显存便宜时这笔账就划算 —— 大模型训练里几乎默认开。

apply_activation_checkpointing 的 check_fn 通常设成"是 transformer block 就 wrap"：

from functools import partial
apply_activation_checkpointing(
    model,
    check_fn=lambda m: isinstance(m, TransformerBlock),
)

这种"按 block 粒度 checkpoint"是 LLM 训练的标准做法。block 内部不 checkpoint —— 否则 attention / mlp 各自重 forward，重算的 FLOPs 与 checkpoint 自身的记账开销都会失控。

18.12 FlatParameter (FSDP-1)：把多个 param 平铺成一个大张量

FSDP-1 内部不直接管理用户的 nn.Parameter，而是把同一 wrap unit 内的多个 param 平铺合并成一个 FlatParameter：

# 用户 module 有 3 个 param
linear.weight  shape [768, 768]    fp32
linear.bias    shape [768]         fp32
norm.weight    shape [768]         fp32

# FSDP-1 合并 (flatten + concat)
flat_param = FlatParameter(torch.cat([
    linear.weight.flatten(),     # 589824 元素
    linear.bias.flatten(),       # 768
    norm.weight.flatten(),       # 768
]))                              # 共 591360 元素

torch/distributed/fsdp/_flat_param.py:202 的 FlatParameter 是 nn.Parameter 子类，元数据里记着每个原始 param 的 shape / stride / numel / FQN 等（dtype 必须统一，torch/distributed/fsdp/_flat_param.py:815 会拒绝 dtype 不一致的参数）。访问 linear.weight 时通过 view 从 FlatParameter 取出对应区段。

为什么要 flatten？

  • AllGather 能一次拉所有 param（vs 分别拉每个 param 的 N 次 collective）
  • shard 时只切一次（按 1/N 切 FlatParameter，不是切几十个 param）
  • memory 连续让 GPU memory bandwidth 利用率高

代价是 view 重建复杂、与 torch.compile 兼容性差。FSDP-2 抛弃 FlatParameter，每个 param 用 DTensor 直接 shard —— 与 compile 兼容性大幅提升；collective 数量并没有变多：同一次 fully_shard 管的参数仍 copy-in 到一块扁平 buffer、发一次 AllGather（本机 gloo 实测 4 个各含 weight + bias 的 unit，forward 共 4 次 AllGather），代价换成了 copy-in / copy-out 两次拷贝。

理解 FlatParameter 让你看到 FSDP-1 用 LOCAL_STATE_DICT 存出的 ckpt 里那些"奇怪 key"（如 0._flat_param，本机 gloo 实测）时不困惑 —— 那是 flat 后的合并张量；FULL_STATE_DICT / SHARDED_STATE_DICT 则还原成原始参数名。

18.13 auto_wrap_policy：决定 sharding 粒度

auto_wrap_policy 让 FSDP-1 自动决定哪些子 module 各自成为一个 wrap unit（FSDP-2 的 fully_shard 没有这个参数，粒度由你对哪些 submodule 调 fully_shard 决定）。粒度选择：

  • 太粗（整个 model 一个 unit）：unshard 一次拉全部参数，参数显存峰值与不分片相同
  • 太细（每个 Linear 一个 unit）：每层 unshard 一次 collective，通信开销爆炸
  • 合适粒度：每个 transformer block 一个 unit，平衡 collective 数量与显存峰值

torch/distributed/fsdp/wrap.py:178 的 ModuleWrapPolicy 是常用 policy：按 module 类型选择 wrap unit。

# FSDP-1：交给 auto_wrap_policy
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import ModuleWrapPolicy

policy = ModuleWrapPolicy({TransformerBlock})    # 每个 TransformerBlock 一个 unit
model = FSDP(model, auto_wrap_policy=policy)

# FSDP-2：手动按同样粒度逐个 fully_shard
for blk in model.layers:
    fully_shard(blk)
fully_shard(model)

其他 policy（都在 torch/distributed/fsdp/wrap.py）：

  • size_based_auto_wrap_policy（:341）：按参数量阈值（默认 min_num_params=1e8）决定
  • transformer_auto_wrap_policy（:298）：按 transformer_layer_cls 集合匹配，语义上等价于 ModuleWrapPolicy
  • 自定义 callable policy

LLM 训练几乎都是"每个 transformer block 一个 unit" —— 与 activation_checkpoint 的 wrap 粒度对齐，让两者协作最优。

18.14 lazy init 与 meta device：训练前不分配显存

70B 模型直接 model = LlamaModel(config) 会立即在 GPU 上分配 280GB params —— 单卡装不下、初始化崩。

torch.device("meta") 是个"假 device"，张量只有 shape / dtype 元信息、没有实际数据。FSDP-2 配 meta device 流程：

with torch.device("meta"):
    model = LlamaModel(config)    # 不分配显存

fully_shard(model, mesh=mesh, ...)   # FSDP 用 meta module 构造 sharded param

model.to_empty(device='cuda')         # 把 meta param 替换成真实显存 (1/N 大小)
init_model_weights(model)             # 用户实现的 init 函数 (按 rank 自己 init 那 1/N)

整个流程全程没有"完整 280GB"在任何 rank 显存里出现。这套 meta 初始化让 70B+ 模型在 80GB 显存的卡上也能启动训练 —— 每卡只物化自己那 1/N（FSDP-1 对应的入口是 param_init_fn 参数）。本机 gloo 两进程实测：meta 上 fully_shard 后 to_empty(device="cpu")，每 rank 的本地分片形状为 [4, 8]（原 weight [8, 8]）。

to_empty() 是个特殊 API：把 meta param 替换成同 shape 的真实张量但不初始化。后续用户调 init_model_weights 在每个 rank 上各自 init 自己持有的那 1/N 数据。这种"分片 init"避免了"先 init 完整 model 再切"的中间显存峰值。

18.15 CPU Offload：参数 / optimizer state 卸到 CPU

显存极紧张时（如单卡训练超大模型），FSDP 提供 cpu_offload：

# FSDP-2
from torch.distributed.fsdp import CPUOffloadPolicy, fully_shard

fully_shard(model, offload_policy=CPUOffloadPolicy(pin_memory=True))

# FSDP-1 的对应写法
from torch.distributed.fsdp import CPUOffload, FullyShardedDataParallel as FSDP
model = FSDP(model, cpu_offload=CPUOffload(offload_params=True))

注意两代 API 名字不同：FSDP-1 是 cpu_offload=CPUOffload(offload_params=...)，FSDP-2 是 offload_policy=CPUOffloadPolicy(pin_memory=...)（_fully_shard/_fsdp_api.py:131 的 OffloadPolicy / :139 的 CPUOffloadPolicy）。

机制：

  • params 平时存在 CPU RAM（host memory）
  • forward 时分片 CPU → GPU 拷贝 + AllGather + GPU 跑 forward；reshard 时直接释放 GPU 上的完整副本（参数没变，不需要拷回）
  • backward 同样先 H2D + AllGather；ReduceScatter 后的梯度分片 GPU → CPU 拷回
  • optimizer step 在 CPU 上跑，optimizer state 也在 CPU（CPUOffloadPolicy docstring：offloads parameters, gradients, and optimizer states to CPU）

代价是每步多出参数 H2D、梯度 D2H 拷贝，走 PCIe，比 NVLink 上的 AllGather 慢一个量级。整个训练吞吐会显著下降（具体幅度取决于 PCIe 代数与模型规模，须自己 profile），但能让"放不下的模型放下"。

optimizer state 不需要单独的开关：开了 offload 后 sharded param 与 grad 本来就在 CPU，optimizer 在这些 CPU 张量上建出的 exp_avg / exp_avg_sq 自然也在 CPU（FSDP-1 的 CPUOffload(offload_params=True) 同理，docstring 写明 optimizer step runs on CPU）。

实战：能用多卡分摊就用多卡（HSDP），cpu_offload 是"硬件不够时的最后兜底"。生产 70B 训练几乎从不用 cpu_offload —— 性能损失太大、不如多租几张卡。

18.16 state_dict_type：完整 vs 分片

FSDP-1 常用的 state_dict 视图有两种（另有保存 FlatParameter 原样分片的 LOCAL_STATE_DICT）：

  • StateDictType.FULL_STATE_DICT（枚举成员在 torch/distributed/fsdp/api.py:270，配套的 FullStateDictConfig 在 torch/distributed/fsdp/api.py:293）：rank 0 收集所有 rank 的分片，组装成完整 state_dict（与未 shard 的 model 视图等价）
  • StateDictType.SHARDED_STATE_DICT（torch/distributed/fsdp/api.py:272，配套 ShardedStateDictConfig 在 torch/distributed/fsdp/api.py:340）：每 rank 输出自己持有的 1/N 切片
from torch.distributed.fsdp import StateDictType

# 完整 state_dict (rank 0 持有完整, 其他 rank 空)
with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT):
    sd = model.state_dict()
    if rank == 0:
        torch.save(sd, "ckpt.pt")

完整 state_dict 简单但 rank 0 要装下完整模型（70B = 280GB），实战不可行。

# 分片 state_dict (每 rank 写自己那 1/N)
with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT):
    sd = model.state_dict()
    dcp.save(sd, checkpoint_id="ckpt_dir")   # import torch.distributed.checkpoint as dcp

分片 state_dict 与第 19 章 §19.6.2 的 DCP 配合，每 rank 并行写自己那份 —— 这是 70B+ 训练 ckpt 的标配。SHARDED_STATE_DICT 写出的不是 ddp model 的视图，而是分片张量字典：FSDP-1 默认每个 entry 是 ShardedTensor（构造 FSDP 时传了 device_mesh 才换成 DTensor，见 torch/distributed/fsdp/_state_dict_utils.py 的 _sharded_post_state_dict_hook；本机 gloo 实测默认得到 ShardedTensor）；FSDP-2 的 model.state_dict() 直接就是 DTensor。DCP 知道怎么序列化这两类分片张量，加载时按当前拓扑重建。

18.17 summon_full_params：临时凑齐完整参数

某些操作（如 model surgery、debug print 完整 weight）需要完整参数。FSDP.summon_full_params context manager 让 FSDP 临时 unshard 整个 group：

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

with FSDP.summon_full_params(model, recurse=True):
    # 此时所有 param 都是完整的 (每 rank 都有完整副本)
    print(model.linear.weight.shape)    # [hidden, hidden] 完整尺寸

实现机制：进入 with 块时触发对所有 wrap unit 的 AllGather → 完整副本临时驻留显存 → 退出时立即 reshard。显存峰值临时翻 N 倍，所以 70B 模型 8 卡 unshard 后单卡瞬间需要 280GB —— 不可行。生产 summon_full_params 用于小 unit / 测试。

writeback=True（默认值）让退出时把完整 param 写回到分片（用于 model surgery 后保存改动）。offload_to_cpu=True 让完整副本驻留 CPU 而非 GPU（省 GPU 但慢）。

18.18 FSDP × clip_grad_norm

第 17 章 §17.35 提过 DDP 下 clip_grad_norm_ 直接用就行（grad 已经全局平均）。FSDP 不行：每 rank 持有 1/N 的 grad，本地 norm 不是全局 norm。

FSDP-1 把它做成 FSDP 包装实例自己的方法——clip_grad_norm_(self, max_norm, norm_type=2.0)，带 @torch.no_grad()（torch/distributed/fsdp/fully_sharded_data_parallel.py:1074）：

# 调用侧用法示例，不是 PyTorch 源码
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = FSDP(model)
model.clip_grad_norm_(max_norm=1.0)

内部实现：

  1. 每 rank 算 local grad norm²（即 sum of squares）
  2. AllReduce sum 得全局 norm²
  3. 算 sqrt 得全局 norm
  4. 每 rank 用全局 norm 缩放自己的 grad 切片

只多一次 AllReduce 标量（极小开销），数学上与 single-rank clip 等价。FSDP-2 没有自己的 clip_grad_norm_：它的 param / grad 都是 DTensor，直接用顶层 torch.nn.utils.clip_grad_norm_，norm 的求和由 DTensor 的算子规则自动变成跨 rank 的 AllReduce。

18.19 FSDP × torch.compile：FSDP-2 的兼容路径

FSDP-1 与 torch.compile 兼容性差（Dynamo 跳过全部 FSDP-1 代码、在每个 FSDP 模块边界 graph break，§18.5）。FSDP-2 重新设计后可以直接套 compile：

fully_shard(model, mesh=mesh)
compiled_model = torch.compile(model)

但先要分清两种模式：

  • 默认模式：torch._dynamo.config.skip_fsdp_hooks 默认 True（torch/_dynamo/config.py:413），FSDP-2 的 hook 被 torch._dynamo.disable 包住（torch/distributed/fsdp/_fully_shard/_fsdp_state.py:60-72 的 disable_if_config_true），在编译图外以 eager 执行 —— AllGather / ReduceScatter 不进图，图里只有计算，hook 处会 graph break（本机 gloo 两进程 + backend="aot_eager" 实测，graph break 原因正是这个 disable）
  • Traceable FSDP2：把 skip_fsdp_hooks 设为 False 并开启 compiled autograd，hook 才会被 trace 进图。FSDP-2 的调度全挂在标准 hook 上：register_forward_pre_hook / register_forward_hook 管 unshard-reshard，输出张量上的 Tensor.register_hook 触发 pre-backward，引擎的 queue_callback 收尾；eager 下负责 post-backward 的 RegisterPostBackwardFunction（torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py:876）在这一模式下不再被调用，改由下一个 param group 的 pre-backward 与 root 的 root_post_backward_callback 补做（见该类的报错说明）；all-gather 的 copy-in / copy-out 是 torch.ops.fsdp.* 自定义算子（torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py），Inductor 另有专门的图 pass（torch/_inductor/comms.py 的 reinplace_fsdp_all_gather、remove_fsdp2_unsharded_param_graph_input_usage）

本机没有 GPU，无法给出实测倍数。能确定的是：两代 FSDP 在默认配置下通信都在图外；FSDP-2 的优势在于参数就是普通 Tensor / DTensor、没有 FlatParameter view，且提供了把 collective 纳入编译图的 Traceable FSDP2 路径。新代码用 FSDP-2，老代码尽快迁移。

18.20 BACKWARD_PRE vs BACKWARD_POST prefetch

§18.4 提过 prefetch 有两种模式：

  • BackwardPrefetch.BACKWARD_PRE（推荐）：layer N+1 的梯度计算开始前就 prefetch layer N。overlap 最多，显存占用也最高（峰值同时持有当前与下一组参数）
  • BackwardPrefetch.BACKWARD_POST：layer N+1 的梯度计算结束后才 prefetch layer N。overlap 少但显存峰值更低

BACKWARD_PRE 是默认，几乎所有场景最优。BACKWARD_POST 仅在显存极紧张时考虑（少 overlap 但少一份 unsharded param 同时驻留）。

forward_prefetch=True 是另一档（FSDP-1 参数，默认 False）：forward 时也 prefetch 下一个 unit 的 AllGather，代价是多一份 unsharded param 同时驻留、显存峰值上升。FSDP-2 没有这个开关：它默认就有隐式 forward 预取（下一组的 copy-in / AllGather 放在独立 stream 上，与当前计算重叠，见 torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py 开头的 implicit prefetching 注释），backward 则默认按 post-forward 的逆序预取；要更激进时用显式的 set_modules_to_forward_prefetch(modules) / set_modules_to_backward_prefetch(modules)（torch/distributed/fsdp/_fully_shard/_fully_shard.py:418 / :438）指定 prefetch 目标 —— docstring 说只传下一个模块时 overlap 与默认相同、只是 CPU 更早发出 AllGather，传两个以上才更激进、也更占显存。

18.21 ZeRO-2 (SHARD_GRAD_OP) 的具体节省机制

§18.2 表里讲 ZeRO-2 是"计算期 params 完整、grads + optimizer 切片"。具体怎么实现？

这里要先纠正一个流传很广的说法：PyTorch 的 SHARD_GRAD_OP 不是"params 根本不切、forward 不用 AllGather"。ShardingStrategy 的 docstring 写得很直白（torch/distributed/fsdp/api.py:43-49）："parameters are sharded outside computation... this strategy unshards before the forward, does not reshard them after the forward, and only reshards them after the backward computation"。所以：

  • forward：params 在计算之外仍是 1/N，forward 前照样要 AllGather 成完整 param —— 与 FULL_SHARD 的 forward 完全一样
  • forward 之后不 reshard：完整 param 一直留在显存里，所以 backward 省掉了 FULL_SHARD 那次 pre-backward AllGather。这才是 ZeRO-2 相对 ZeRO-3 真正省下的通信
  • backward：每 rank 算自己 batch 的本地 grad → ReduceScatter（不是 AllReduce！）让每 rank 拿到全局 grad 的 1/N 切片，然后才 reshard params
  • optimizer step：每 rank 只更新自己持有的 1/N 参数（用 1/N grad + 1/N optimizer state）；下一次 forward 的 AllGather 自然会把更新后的分片拼回来，没有单独一次"step 后 AllGather"

FSDP-2 把同一件事表达成 reshard_after_forward=False，它的 docstring 也是同一句话："keeps the unsharded parameters in memory after forward and avoids the all-gather in backward"（torch/distributed/fsdp/_fully_shard/_fully_shard.py:152-154）。

这套流程让 grads 与 optimizer state 各只占 1/N（与 ZeRO-3 同），而完整 params 从 forward 一直驻留到 backward 结束（vs ZeRO-3 的用完即还）。以 fp32 + Adam 的经典账算（params : grads : optimizer state = 1 : 1 : 2，共 4 份），ZeRO-2 把每卡的 grads 与 optimizer state 各降到 1/N，N 足够大时每卡只剩那 1 份完整 params —— 省掉的是 4 份里的 3 份，同时保留完整 params 的灵活性。

适用场景：模型刚好能装下 params 但 optimizer state（Adam 是 params × 2）装不下。比如 13B 模型 fp32 params 52GB（单张 80GB 卡放得下）但 Adam state 104GB（放不下）—— ZeRO-2 完美。

18.22 NO_SHARD：FSDP 退化成 DDP

ShardingStrategy.NO_SHARD 让 FSDP 不分片任何东西，等价于 DDP：

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, ShardingStrategy

FSDP(model, sharding_strategy=ShardingStrategy.NO_SHARD)
# 等价于
DDP(model)

注意 ShardingStrategy 是 FSDP-1 的枚举（torch/distributed/fsdp/api.py:32，成员 FULL_SHARD / SHARD_GRAD_OP / NO_SHARD / HYBRID_SHARD / _HYBRID_SHARD_ZERO2）。FSDP-2 的 fully_shard 没有 sharding_strategy 参数 —— 对应能力由 mesh 的维数与 reshard_after_forward 表达。

这个看似无意义的选项有真实工程价值：

  • 统一接口：训练框架可以"用 FSDP 一套接口表达 DDP / ZeRO-2 / ZeRO-3 / HSDP"，配置切换简单
  • fallback 路径：HSDP 的跨节点维度本质是 NO_SHARD（跨节点不切，节点内才切）
  • debugging：怀疑 FSDP 引入的 bug 时切到 NO_SHARD 看是否消失

NO_SHARD 让 FSDP-1 的接口涵盖 DDP 全部功能。FSDP-2 换了一套表达：不分片就是不对该 module 调 fully_shard，ZeRO-2 语义则由 reshard_after_forward=False 表达。

18.23 use_orig_params=True：兼容老 optimizer

FSDP-1 用 FlatParameter（§18.12）会让用户原本的 nn.Parameter 引用失效 —— optimizer 创建时拿到的是 user param，FSDP wrap 后这些 param 引用指向已经"被合并"的位置。

use_orig_params=True 让 FSDP-1 保留原始 param 引用：用户构造 optimizer 时传 model.parameters()，FSDP 内部建立 orig_param ↔ FlatParameter 切片的映射。optimizer step 时按 orig_param 更新，FSDP 自动写回 FlatParameter。

这条选项让"现有 optimizer 代码不改"就能上 FSDP。代价是内部多一层映射开销（小）。生产代码强烈建议开启。FSDP-2 默认行为就这样（每个 param 独立 DTensor，根本不需要这个选项）。

18.24 FSDP × 3D parallel (DP + TP + PP)

70B+ 训练经常需要 3D parallel：

  • Data Parallel (DP)：FSDP 在 DP 维度上分片
  • Tensor Parallel (TP)：把单层 weight 切到多卡（如 attention QKV 切 4 路）
  • Pipeline Parallel (PP)：把多层切到多卡，流水线执行

DeviceMesh（§18.7）让 3D 配置变简单：

mesh = init_device_mesh("cuda", (PP, DP, TP), mesh_dim_names=("pp", "dp", "tp"))

# 先 TP：在 tp 维度上切单层 weight
# tp_plan 形如 {"attn.q_proj": ColwiseParallel(), "attn.o_proj": RowwiseParallel(), ...}
parallelize_module(model, mesh["tp"], tp_plan)

# 再 FSDP：在 dp 维度上 shard（此时参数已是 TP 的 DTensor，FSDP-2 在其上再叠一层 dp 分片）
fully_shard(model, mesh=mesh["dp"])

# PP 用 mesh["pp"] 配 PipelineStage

3D parallel（再加上序列维就是 4D）是 405B 级别训练的基础 —— Llama 3 405B 论文（arXiv:2407.21783）报告用了最多 16K 张 H100。每个维度的 group 用 functional collectives（第 16 章 §16.12）通信，互不干扰。

实战参数选择：8 卡节点内 TP=8（NVLink 高带宽）+ 4 节点 DP（FSDP）+ 4 PP（跨更多节点）= 128 卡训练。具体取决于模型与硬件拓扑。

18.25 FSDP × LoRA：高效微调

LoRA（Low-Rank Adaptation）只训练注入的低秩矩阵 A、B，冻结原始 weight。FSDP 与 LoRA 的协作：

# 冻结 base model 参数
for p in base_model.parameters():
    p.requires_grad = False

# 注入 LoRA adapter (有 grad)
inject_lora(base_model, rank=8)

# FSDP wrap (只 shard 有 grad 的部分?)
fully_shard(base_model, ...)

关键问题：FSDP 默认 shard 所有 params，包括 frozen 的 base weights。但 frozen weights 不需要 grad / optimizer state，shard 浪费空间吗？

实际不浪费 —— frozen weights 仍要 unshard：forward 要用，backward 里算 grad_input 也要用（梯度得穿过它们流到前面层的 LoRA adapter）；但它们没有 grad 要 ReduceScatter、也没有 optimizer state，显存账上反而占更少。FSDP 在内部识别 requires_grad=False 的 param、跳过它们的 grad 处理。

LoRA + FSDP 是 fine-tuning 70B 模型的标配 —— base model 冻结后只训几十兆 LoRA 参数，FSDP 让 base model 分片到多卡能装下。

18.26 FSDP × gradient accumulation

显存不够调 batch size 时常用 gradient accumulation：跑 N 个 micro batch 累积梯度、再一次 step。

FSDP 下的微妙点：每 micro batch 默认会触发 ReduceScatter 同步梯度。N 次 micro batch = N 次 collective —— 浪费。

FSDP-1 用 with model.no_sync():（torch/distributed/fsdp/fully_sharded_data_parallel.py:1034）；FSDP-2 没有 no_sync，改成显式开关 set_requires_gradient_sync(bool)（_fully_shard/_fully_shard.py:319，docstring 里直接写着 This is the equivalence of no_sync in FSDP1）：

# FSDP-2
model.set_requires_gradient_sync(False)
for i in range(accumulation_steps - 1):
    loss = model(batch[i])
    loss.backward()        # 不触发 ReduceScatter / AllReduce, 只本地累积

model.set_requires_gradient_sync(True)
loss = model(batch[-1])
loss.backward()            # 这次正常触发 ReduceScatter
optimizer.step()

它同时关掉 reduce_grads 与 all_reduce_grads（HSDP 下两维一起关）。N-1 次 micro batch 跳过 collective，只在最后一次同步累积总梯度：通信次数从 N 降到 1。

DDP 那边仍是 no_sync 上下文管理器，思想一致。

18.27 FSDP × 量化训练 / 推理

第 20 章会讲量化。FSDP 与量化协作有几个工程点：

训练时 fp8：把 AllGather 的通信 dtype 降到 fp8（8 bit）能比 bf16 再省一半带宽。但 fp8 数值范围窄、需要 per-tensor scale，MixedPrecisionPolicy.param_dtype 本身只管 cast、不管 scale —— 要正确处理 scale 得靠 torchao 的 float8 或 NVIDIA TransformerEngine 这类库把量化逻辑塞进 param 的表示里，再让 FSDP-2 去 AllGather 那个已量化的表示。

推理时 INT8：FSDP shard 的 model 直接量化遇到障碍 —— 量化要看完整张量做 calibration。一般做法：训练用 FSDP / 量化前把 model 收到单卡 / 量化后再加载（不分片或者用 PT2E 静态图）。

LLM 推理几乎不用 FSDP（vLLM / SGLang 用 TP + KV cache 管理替代）。FSDP 主要是训练时的工具，推理时通常切换到其他并行策略。

18.28 FSDP root module 的特殊性

FSDP 中的"root module"（最外层 fully_shard 的 module）有几个特殊职责：

  • 管理整个 FSDP unit 树：root 持有所有子 unit 的引用，调度它们的 unshard / reshard 时机
  • 管理 stream：root 在 lazy init 里建好通信 stream，子 unit 共用（FSDP-2 是 FSDPCommContext 的 all_gather_copy_in_stream / all_gather_stream / reduce_scatter_stream / all_reduce_stream；FSDP-1 是 _unshard_stream / _pre_unshard_stream 等）
  • 触发 root_pre_forward：第一次 forward 时 root 做整体初始化（首次 AllGather、CUDA stream sync 等）
  • forward 完成后调 reshard：root 决定 root unit 自身的 reshard 时机（最后一个 reshard 在整 forward 完成后）

root 在首次 forward 的 lazy init 里确定（torch/distributed/fsdp/_fully_shard/_fsdp_state.py 的 _lazy_init）：最先跑 forward 的 fully_shard 模块就是 root，通常就是最外层 model。如果用 fully_shard(layer) 给每层都 wrap，最外层依然是 root，每层是 sub-unit。root 的 lifecycle 决定整体训练流程，理解这点能解释为什么"在 root 之外的代码看到的 param 是分片的、root forward 内部看到完整的"。

18.29 FSDP × torch.export 与部署

FSDP 是训练工具，部署时通常不带 FSDP 直接 export。部署流程：

  1. 训练完成后用 FSDP.summon_full_params(model) 或者 state_dict_type=FULL_STATE_DICT 收集完整权重
  2. 在单卡（或推理用的并行配置）重建 model（不带 FSDP wrap）
  3. 加载完整权重
  4. torch.export.export(model, ...) + AOTI（第 15 章 §15.10.1）

这条流程让训练时的 FSDP 与部署时的 AOTI 完全解耦。FSDP 不存在于部署 binary 里 —— 它是纯粹的训练时工具。

实战：HuggingFace Transformers 的 FSDP 训练流程结尾通常调 unwrap_model(model) 取出原始 nn.Module、再保存 state_dict。这是与 FSDP 解耦的标准做法。

18.30 FSDP-2 与 DTensor 的层次关系

FSDP-2 内部把每个 nn.Parameter 替换成 DTensor（§18.8）。这意味着：

  • 用户视角：forward 之外，model.linear.weight 是 DTensor（local_shard 是 1/N），optimizer、clip_grad_norm_、state_dict 看到的都是它
  • forward 视角：FSDP-2 的 hook 自己发 AllGather，把 unsharded 的普通 nn.Parameter 换上去，计算算子看到的是普通 Tensor，而不是靠 DTensor 算子去触发 placement 转换（§18.8）
  • backward 视角：FSDP-2 自己发 ReduceScatter，把结果包成 DTensor 梯度（本机 gloo 实测 weight.grad 是 Shard(0) 的 DTensor）

也就是说，FSDP-2 把分片的表示下沉到 DTensor，调度与通信仍由自己负责（按 param group 批量 AllGather / ReduceScatter）。

顺带纠正一个常见印象：FSDP-2 并不比 FSDP-1「代码更少」。v2.11 里 torch/distributed/fsdp/_fully_shard/ 共 4552 行，而 FSDP-1 的两个主文件 fully_sharded_data_parallel.py（2167 行）+ _runtime_utils.py（1654 行）合计 3821 行。FSDP-2 的优势不在体量，而在抽象层次：手写的 flat-buffer 簿记换成了类型化的 placement。

DTensor 思想的胜利：把分片表示为类型而非协议。Tensor 的 placement 是 Tensor 类型的一部分，编译器 / 运行时都能利用这个信息做优化。这是 PyTorch 分布式训练演进的下一代方向。

18.31 完整训练 step 时间分解

70B Llama 训练单 step 时间（H100，32 卡 4 节点 HSDP，bf16）：

forward
  ├─ AllGather params        ← 理想情况完全 overlap 进上一个 unit 的计算
  └─ forward 计算
backward                     ← 通常是 forward 的两倍量级 (要算 grad_input + grad_weight)
  ├─ AllGather params        ← overlap
  ├─ backward 计算
  └─ ReduceScatter grads     ← overlap
跨节点 AllReduce (HSDP only)  ← 只在需要同步梯度的那一次 backward 末尾
optimizer.step               ← 每 rank 只更新自己那 1/N, 占比很小

本机无 GPU，给不出具体毫秒数 —— 这类数字必须在你自己的集群上用 profiler（第 21 章）实测。能确定的是相对关系：FSDP 比 DDP 慢，换来的是显存从「单卡装不下」降到「每卡 1/N + 一个 unit 的 unsharded peak」。这就是 FSDP 的工程哲学：用通信换显存。

如果显存够（中等模型），用 DDP 性能更好；如果装不下，FSDP 是唯一选择。混合 HSDP 是中间值 —— 跨节点不分（DDP 风格）、节点内分（FSDP 风格），平衡通信与显存。

18.32 内存账：FSDP vs DDP 的精确对比

7B 模型 fp32 训练详细内存账（每卡）：

| 项                    | DDP    | FSDP-2 ZeRO-3 (8 卡) |
| params                | 28 GB  |  3.5 GB              |
| grads                 | 28 GB  |  3.5 GB              |
| optimizer (Adam m+v)  | 56 GB  |  7 GB                |
| activations           |  8 GB  |  8 GB (相同)         |
| 临时 buffer            |  4 GB  | 28 GB (unsharded peak)|
| 单卡总               |124 GB  | 50 GB                |

注意 FSDP-2 的"临时 buffer"列：这 28GB 是把整个 model 当成一个 wrap unit 的退化情形——一次 unshard 就要凑齐全部 params。真实配置按 transformer block 逐个 fully_shard（§18.13），同时在飞的 unsharded buffer 只有「当前 unit + prefetch 中的下一个 unit」那么大，通常是几百 MB 到几 GB 量级。unsharded peak 是 per-unit 的，不是整模型的——这正是 §18.13 强调 unit 粒度的原因，也是为什么 FSDP-2 的实际显存峰值不是简单的 1/N。

70B 模型同样表（32 卡 ZeRO-3）：

DDP:  params 280 + grads 280 + opt 560 + activ 几十 + buffer = 1100+ GB → 装不下
FSDP-2: params 8.75 + grads 8.75 + opt 17.5 + activ 几十 + buffer 280 (unsharded) = ~330 GB

上面那个 280 同样是"整模型一个 unit"的极端值。按 block 粒度 wrap 之后，unsharded buffer 只剩「当前 unit + 预取的下一个 unit」，70B 才真正能在 32 卡上跑起来 —— 这告诉我们 unit 粒度不能太大（每 unit unshard 后的临时显存是关键约束）。

18.33 实测带宽 vs 计算的临界点

FSDP 性能取决于"通信能不能 overlap 进计算"。临界点：

  • 节点内 NVLink：带宽以百 GB/s 计，一个 transformer block 的 params AllGather 通常远快于该 block 的计算 → 能完全 overlap
  • 节点间 InfiniBand：按每卡计带宽明显低于 NVLink（常差到接近一个数量级），但配合 prefetch 多数情况下还能 overlap
  • 节点间普通以太网：带宽比 NVLink 低一到两个数量级，同样大小的 AllGather 会比计算慢得多，基本无法 overlap —— 这就是为什么以太网集群用 HSDP 而非纯 FULL_SHARD

具体数字随代际（NVLink 4/5、IB HDR/NDR、100/400 GbE）差异极大，别抄别人的表，用 torch.distributed 自带的 benchmark 或 NCCL tests 在自己集群上测。

判断你的硬件是否适合纯 FULL_SHARD：

  • 算每 backward step 的 GPU 计算时间（用 profiler 看）
  • 算每个 wrap unit 的 AllGather 通信时间（unit_param_size / 带宽）
  • 比较两者，AllGather > 计算时间就要切 HSDP 或减小 unit 粒度

把 unit 的 param_bytes / 实测带宽 与 profiler 里该 unit 的 kernel 时间放一起比：前者远小于后者就是健康的 overlap，两者接近就说明要么减小 unit 粒度、要么换 HSDP。

18.34 FSDP × torch.compile 的协作细节

FSDP-2 + compile 的具体内部协作：

fully_shard(model, mesh=mesh)
compiled = torch.compile(model)

out = compiled(x)

发生的事：

以下 1–4 步描述的是 Traceable FSDP2（torch._dynamo.config.skip_fsdp_hooks = False + compiled autograd）；默认配置下 FSDP-2 的 hook 在图外 eager 执行，collective 不进图（§18.19）。

  1. 第一次调用：Dynamo trace model(x)，FSDP-2 的 pre/post-forward hook 被一起 trace 进来
  2. AOTAutograd functionalize：unsharded param 的 resize_ / torch.ops.fsdp.copy_ 这类 mutation 被 functionalize 进图
  3. Inductor 图 pass：torch/_inductor/fx_passes/post_grad.py 在 skip_fsdp_hooks=False 时额外跑 remove_fsdp2_unsharded_param_graph_input_usage 与 reinplace_fsdp_all_gather（均在 torch/_inductor/comms.py），清理这些 resize / copy、让 all-gather 原地写
  4. Inductor codegen：collective 仍以 extern call 的形式留在生成的 wrapper 里（Inductor 不会把 NCCL 塞进 Triton kernel），但它进入了同一张图的调度，能和前后的 compute kernel 一起排序、一起决定 buffer 复用

最终编译产物里 collective 与 compute 处在同一张图、由同一份 wrapper 顺序发射。这与 FSDP-1 + compile（在每个 FSDP 模块边界 graph break）相比有本质提升。

本机无法实测加速比。定性上，compile 在这条路径上省掉的是 eager 的 Python 调度开销与 kernel 启动开销，模型越大、单 kernel 越大，相对收益越小；Traceable FSDP2 额外的收益来自 collective 与 compute 被放进同一个调度决策里。

18.35 FSDP × Pipeline Parallel

Pipeline Parallel (PP) 把模型按层切到多卡，与 FSDP 在不同维度上分。3D parallel 中两者协同：

mesh = init_device_mesh("cuda", (PP, DP), mesh_dim_names=("pp", "dp"))

# 每 pipeline stage 内用 FSDP
for stage_id in range(PP):
    if local_pp_rank == stage_id:
        layers = model.layers[stage_id*L:(stage_id+1)*L]   # nn.ModuleList
        for layer in layers:        # fully_shard 不接受没有 forward 的 ModuleList 容器
            fully_shard(layer, mesh=mesh["dp"])

# Pipeline schedule 处理跨 stage 通信；stage_module 是把这段 layers 包起来、有 forward 的 nn.Module
pipe = PipelineStage(stage_module, ...)

每个 PP stage 内部用 FSDP shard 自己那部分 layers。stage 之间用 P2P send/recv 传 activation。这种"stage 内 FSDP、stage 间 PP"是 405B 训练的标准配置。

实战调优：PP 的"bubble"（pipeline 启动 / 结束的空闲时间）与 FSDP 的 collective 时间互动复杂，需要 profiler（第 21 章）实测每个 stage 的 timeline 优化。

18.36 FSDP × evaluation：临时切到完整模型

训练循环里偶尔要做 eval（计算 val loss / metrics）。FSDP 模型默认是分片状态，eval 时需要完整 forward。两条路：

A. 保持 FSDP 状态做 eval（推荐）：

model.eval()
with torch.no_grad():
    for batch in val_loader:
        out = model(batch)    # FSDP 仍 unshard / reshard, 与训练相同路径

eval 期间 FSDP 仍触发 AllGather + reshard，但因为 no_grad，没有 backward、也不需要 ReduceScatter grads。开销是单纯的 forward unshard。

B. summon_full_params 一次性 unshard（FSDP-1 接口，fully_shard 没有对应的 summon_full_params）：

with FSDP.summon_full_params(model, recurse=True):
    model.eval()
    with torch.no_grad():
        for batch in val_loader:
            out = model(batch)

完整模型驻留显存，eval 期间不触发 collective —— 但单卡显存压力翻 N 倍。只在 eval batch 多 + 模型小时用，70B 训练绝对不能用（OOM）。

实战推荐 A 路线：FSDP eval 速度足够，省心。

18.37 FSDP 错误诊断速查表

症状 可能原因 诊断
反向卡在 AllGather 某 rank 进度不一致 / NCCL hang 开 TORCH_NCCL_DESYNC_DEBUG=1，watchdog 超时后报告哪个 rank 没跟上（TORCH_NCCL_ASYNC_ERROR_HANDLING 在 v2.11 默认已是 3，见 torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:910-911）
OOM 在 forward 第一层 wrap 粒度太粗（整 model 一个 unit） 按 transformer block 粒度逐个 fully_shard（FSDP-1 则用 ModuleWrapPolicy）
比 DDP 慢很多 跨节点带宽不够，AllGather 无法 overlap 改 HYBRID_SHARD / _HYBRID_SHARD_ZERO2，或减小 unit 粒度
fully_shard 之后替换子模块 / 参数，前向报类型或形状错误 model surgery 发生在 wrap 之后 检查是否在 fully_shard 之后改了 model，改动挪到 wrap 之前
ckpt save 时 rank 0 OOM 用了 FULL_STATE_DICT 切到 SHARDED_STATE_DICT + DCP
加载 HF safetensors 后训练异常 weight 没正确切片到 DTensor 用 torch.distributed.checkpoint.state_dict.set_model_state_dict（torch/distributed/checkpoint/state_dict.py:1369）加载，别手写 load_state_dict
forward 与 eager 数值不一致 precision 配置错 检查 MixedPrecisionPolicy 的 reduce_dtype（应为 fp32）

这几类症状覆盖了 FSDP 日常调试里的绝大多数情况，遇到问题先按表对号入座能省掉大量瞎猜。

18.38 FSDP 的 unshard 时机决策

FSDP 的 unshard 调度可以示意为（简化的概念图，不对应某个具体函数）：

flowchart TB
    Start[forward 开始]
    Start --> Stage{当前 unit?}
    Stage -->|root unit| RootUnshard[unshard root + 立即 prefetch unit 1]
    Stage -->|sub unit| SubCheck{已被 prefetch?}
    SubCheck -->|是| Use[直接用]
    SubCheck -->|否| LazyUnshard[紧急 unshard]
    Use --> Compute[执行 forward]
    Compute --> Reshard[reshard 当前 unit]
    Reshard --> Prefetch[启动 prefetch 下下个 unit]
    Prefetch --> Stage
    style RootUnshard fill:#fef3c7
    style LazyUnshard fill:#fee2e2

核心规则：当前 unit 计算时，下一个 unit 的 AllGather 应该已经在飞。FSDP-2 默认靠独立 stream 做隐式预取，需要更激进时用 set_modules_to_forward_prefetch；FSDP-1 则要开 forward_prefetch=True（§18.20）。

如果实际跑出来 prefetch 没赶上（latency-bound 场景），会触发 lazy unshard —— 当前 unit 等 AllGather 完成，整个 timeline 出现 gap。这时 profiler 的 distributed view 能直接看到。优化：调整 wrap policy 让 unit 更小（每个 unit AllGather 时间更短、prefetch 更容易及时）。

18.39 grad accumulation 的 FSDP 显存账

§18.26 提了 no_sync 跳过中间 ReduceScatter。但还有显存账值得展开：

正常 FSDP grad accumulation 流程：

micro_step 1: forward (unshard / reshard) + backward (有 ReduceScatter)
              → grad 保留在 1/N 切片
micro_step 2: 同上, grad 累积到 1/N 切片
...
micro_step N: 同上 + optimizer.step

每 micro step 都触发 ReduceScatter，累积 N 次通信。no_sync 模式：

set_requires_gradient_sync(False)    # 内部不触发 ReduceScatter
    micro_step 1: forward + backward (grad 保留为 unsharded, 与完整 params 同尺寸)
    micro_step 2: grad 就地累积, 仍是 unsharded
    ...
    micro_step N-1: 同上

set_requires_gradient_sync(True)
micro_step N: forward + backward (这次触发 ReduceScatter, grad 切到 1/N)
optimizer.step()

关键差异：跳过 reduce 期间累积的 grad 是 unsharded 的（每 rank 持有与完整 params 同尺寸的一份），而不是 1/N。注意它不随 micro batch 数放大 —— 是就地累加，不是攒 N 份；但相对分片状态仍然是 N 倍的 grad 显存。

所以 FSDP 下要不要跳过 reduce 是一道显存 / 通信的取舍题：显存宽裕就跳（省 N-1 次 collective），显存紧张就每个 micro batch 都 reduce（grad 常驻 1/N）。DDP 那边没有这个矛盾（grad 本来就是完整的），这是 FSDP 与 DDP 工程取舍的不同。

18.40 FSDP 与其他并行框架对比

框架 核心思想 对比
PyTorch FSDP-2 DTensor + DeviceMesh 与 PyTorch 生态深度集成、与 compile 兼容
DeepSpeed ZeRO Shard params/grads/state 更早实现，是 FSDP 的灵感来源；与 PyTorch 集成靠 wrapper
Megatron-LM TP + PP 为主 NVIDIA 主推，专攻 TP/PP；DP 侧用自带的 DDP，可配 distributed optimizer 分片 optimizer state
ColossalAI 多策略统一抽象 国产框架，支持多种并行策略
OneFlow 全局视角自动并行 算子级自动决定并行策略

PyTorch FSDP-2 的工程优势是与 torch.compile / torch.export / DTensor 共生 —— 不是独立框架而是 PyTorch 生态的一等公民。其他框架要么是 wrapper（依赖 PyTorch 但侵入式扩展），要么是独立运行时（迁移成本高）。

实战选择：

  • 用 PyTorch 训练 → 用 FSDP-2
  • 已有 DeepSpeed 代码 → 继续用 DeepSpeed（迁移成本高、收益不一定值得）
  • 极致 TP/PP 优化 → Megatron-LM
  • 国产芯片支持 → ColossalAI / 厂商定制

18.41 FSDP × HF Transformers / Lightning 实战

HuggingFace Trainer 与 PyTorch Lightning 都支持 FSDP，但配置方式不同：

HuggingFace TrainingArguments：

training_args = TrainingArguments(
    fsdp="full_shard auto_wrap",
    fsdp_config={
        "transformer_layer_cls_to_wrap": ["LlamaDecoderLayer"],
        "min_num_params": 0,
    },
)

PyTorch Lightning Strategy：

trainer = pl.Trainer(strategy="fsdp", devices=8, ...)
# 或者更细粒度
trainer = pl.Trainer(strategy=FSDPStrategy(
    auto_wrap_policy={LlamaDecoderLayer},
    sharding_strategy="FULL_SHARD",
))

这两段配置用的都是 FSDP-1 的语汇（sharding_strategy、auto_wrap_policy / transformer_layer_cls_to_wrap）—— 框架替用户把 wrap policy 写好了。但理解原始 API 让你能在 Trainer / Lightning 配置出错时定位是哪一层包装的问题。

国内 Llama-Factory、Firefly 等微调框架也内置 FSDP 支持，配置类似 HF Trainer。

18.42 FSDP 与 elastic 训练的协作

第 17 章 §17.29 / §17.40 讲过 elastic（torchrun + max_restarts）。FSDP 与 elastic 协作时有几个细节：

  • 重启时 mesh 配置必须一致：如果重启时用了不同 world_size，必须用 DCP 加载（自动 reshard 到新 mesh）
  • ckpt 频率要够高：失败重启等于丢失从最近 ckpt 到现在的所有计算。间隔要在「存 ckpt 的开销」与「故障时丢掉的计算」之间权衡
  • DCP 的写入要 robust：跨节点写文件可能因为某 rank 网络抖动失败，DCP 要支持"部分 rank 重写"

实战：torchrun + DCP + 定期 ckpt 是 LLM 训练的常见组合。失败时整 job 重启、加载最近 ckpt 继续训。大规模集群上长时间训练期间故障不可避免（Llama 3 405B 论文报告 54 天预训练里有 419 次意外中断），没有 elastic 自动恢复就要全人工介入。

18.43 实战训练流程的完整推荐

70B 模型从零训练的标准 FSDP-2 流程：

# 1. 设置环境
import os
import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    apply_activation_checkpointing,
)

# 2. 初始化 process group
dist.init_process_group(backend='nccl')
local_rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)

# 3. 创建 DeviceMesh (32 卡 4 节点 HSDP)
mesh = init_device_mesh("cuda", (4, 8), mesh_dim_names=("inter", "intra"))

# 4. meta device 构造模型
with torch.device("meta"):
    model = LlamaModel(config)

# 5. activation_checkpoint 每层
apply_activation_checkpointing(
    model,
    check_fn=lambda m: isinstance(m, LlamaDecoderLayer),
)

# 6. FSDP-2 wrap (HSDP)
mp_policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32,
)
for layer in model.layers:
    fully_shard(layer, mesh=mesh, mp_policy=mp_policy)
fully_shard(model, mesh=mesh, mp_policy=mp_policy)

# 7. lazy init weights
model.to_empty(device='cuda')
init_model_weights(model)

# 8. compile
model = torch.compile(model)

# 9. optimizer 与 scheduler
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, fused=True)
scheduler = ...

# 10. 训练循环
for batch in loader:
    optimizer.zero_grad()
    loss = model(batch).loss
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)   # param 是 DTensor, norm 求和自动跨 rank
    optimizer.step()
    scheduler.step()

整套流程把全专栏前面章节的内容串起来：DeviceMesh（§18.7）+ FSDP-2（§18.5）+ 混合精度（§18.6 + 第 20 章）+ activation_checkpoint（§18.11）+ HSDP（§18.9）+ torch.compile（第 12-15 章）+ fused optimizer（第 10 章）+ DTensor grad clip（§18.18）。

这是 70B 级 LLM 训练的一套常见配置。理解整章后回到这段代码，每行配置都对应某个具体的工程决策。

18.44 FSDP 调试：常用日志与工具

调 FSDP 训练的几个常用工具：

TORCH_DISTRIBUTED_DEBUG=DETAIL：开启后进程组被包进 ProcessGroupWrapper（torch/distributed/distributed_c10d.py 的 _create_process_group_wrapper），每个 collective 发出前先借一个 gloo 辅助组在各 rank 间核对操作类型与张量形状，不一致时报出各 rank 正在跑的 collective，能定位"哪个 rank 的 collective 对不上"。

TORCH_LOGS=+fsdp：FSDP-2 的日志都是 DEBUG 级，要带 + 才会打出来（本机 gloo 实测 TORCH_LOGS=fsdp 无输出）；打印每个 unit 的 FSDP::pre_forward / post_forward / pre_backward / post_backward 事件，看调度顺序是否符合预期。

torch.profiler.record_function 标记：FSDP 的关键路径自带 profiler 标记（FSDP-2 如 torch/distributed/fsdp/_fully_shard/_fsdp_state.py:316 的 FSDP::root_post_backward_callback；FSDP-1 如 torch/distributed/fsdp/_runtime_utils.py 里的 FullyShardedDataParallel.rate_limiter），在 trace 里能直接按名字定位。

chrome trace + distributed view（第 21 章 §21.9）：终极调试工具。能看到每 rank 的 collective timeline + 各 rank 等待时间。

实战调试流程：先开 TORCH_LOGS=+fsdp 看调度顺序对不对、再用 chrome trace 看实际 timeline。多数 FSDP 性能问题（如 unshard 没及时 prefetch）能在这两层定位。

18.45 FSDP 设计上的几个隐形约束

FSDP 不是万能的。几个不能突破的工程约束：

1. 每个 unit 的 unshard 必须能装下完整 params：FSDP 的最大单 unit param 大小不能超过单 rank GPU 显存。70B 单层 1B 参数 → 4GB（fp32）单 rank 装得下。但 1T 模型单层 100B → 400GB 装不下，FSDP 必须配合 TP 才能跑。

2. 同一 unit 的所有 param 必须同 dtype：FlatParameter（FSDP-1）合并要求 dtype 一致；FSDP-2 用 DTensor 没这个限制但一个 unit 同 dtype 仍是最优。

3. unit 边界不能跨越复杂的 module 控制流：if-branch / loop 内的 module 不适合自己成为 wrap unit，因为不是每次 forward 都被调到，prefetch 决策困难。

4. FSDP-2 的 mesh 一旦决定就不能改：训练中途换 mesh 配置（如从 8 卡变 16 卡）需要先 ckpt 落盘 → 用新 mesh 重 init → DCP load 时自动 reshard。不能"在线变"。

理解这些约束让你设计训练架构时心里有数。"一切都用 FSDP" 不是答案 —— 极限场景下还要 TP / PP / 自定义并行。

18.46 几条工程经验

1. ShardingStrategy 选择：完整 params 单卡放得下、只是 optimizer state 放不下时用 SHARD_GRAD_OP（ZeRO-2，grads 与 optimizer state 降到 1/N，见 §18.21）；连 params 都放不下用 FULL_SHARD；跨节点带宽不足用 HYBRID_SHARD

2. 用 v2.6+ 公开路径下的 FSDP-2 (torch.distributed.fsdp.fully_shard)：除非你已经有大量 FSDP-1 代码，否则新代码直接用 FSDP-2

3. backward_prefetch=BACKWARD_PRE 是默认，forward_prefetch 默认是 False（torch/distributed/fsdp/fully_sharded_data_parallel.py:404 / :410）：反向 prefetch 开箱即用，forward prefetch 要显式开、且换显存峰值（§18.20）

4. wrapping policy 设到合适粒度：每个 transformer block 单独 shard 通常最优。粒度太细（每个 Linear 都 shard）通信开销过大；太粗（整个模型一个 shard）overlap 不好

5. use_orig_params=True：让 FSDP 不破坏 param 引用，原 optimizer 能直接复用。FSDP-1 的兼容选项，FSDP-2 默认就这样

6. cpu_offload=CPUOffload(offload_params=True)：参数卸到 CPU，省更多显存，代价是 H2D 拷贝拖慢训练。仅在显存极紧张时用

7. checkpoint 与 FSDP：FSDP 的 ckpt 用 torch.distributed.checkpoint（DCP）存分布式格式。每 rank 只存自己那 1/N，避免 rank 0 写 1TB 文件的瓶颈

8. activation_checkpoint 与 FSDP 叠加用：FSDP 省 params/grads/optimizer 显存，checkpoint 省 activation 显存。两者正交、可叠加

18.47 实战决策路径

flowchart TD
    Start[要训练大模型]
    Start --> Size{模型 / GPU?}
    Size -->|装得下| DDP[用 DDP]
    Size -->|装不下| Bandwidth{跨节点带宽?}
    Bandwidth -->|InfiniBand 高带宽| Full[FULL_SHARD]
    Bandwidth -->|普通 ethernet| Hybrid[HYBRID_SHARD<br/>节点内 shard 跨节点 DDP]
    Full --> Mp[配 MixedPrecision]
    Hybrid --> Mp
    Mp --> Ckpt[启用 activation_checkpoint]
    Ckpt --> Compile[用 torch.compile 提速]
    style Full fill:#dcfce7,stroke:#22c55e
    style Hybrid fill:#fef3c7,stroke:#f59e0b

70B Llama 训练的典型配置：32 张 H100、HYBRID_SHARD（同 8 卡节点 ZeRO-3、跨 4 节点 DDP）、bf16 通信 + fp32 reduce、每个 transformer block 单独 fully_shard 并做 activation_checkpoint（与 §18.43 一致）。

18.48 FSDP 与 PyTorch 演进路线

FSDP 在 PyTorch 演进路线上的位置：

几个能从仓库 git 历史里查证的节点：

  • v1.11：torch/distributed/fsdp 目录公开（PR #72084 "make fsdp folder to be public"），FSDP-1 登场，思想源自 DeepSpeed ZeRO
  • v2.3：FSDP-2 的 fully_shard 前端首次落地，路径是私有的 torch.distributed._composable.fsdp（PR #117776，reland #118525）
  • v2.6：fully_shard 移到公开的 torch.distributed.fsdp（PR #141868）
  • v2.11（本专栏基准）：两代并存 —— torch/distributed/fsdp/__init__.py 的 __all__ 里明确分成 # FSDP1 与 # FSDP2 两块；FSDP-2 侧还多出 share_comm_ctx、register_fsdp_forward_method 等导出，以及 FSDPModule.set_custom_all_gather / set_custom_reduce_scatter 自定义 collective 接口

v2.11 的源码里 FSDP-1 整体尚未被标记 deprecated（只有 FSDP.state_dict_type / set_state_dict_type、full_optim_state_dict 这类个别老接口带 deprecation warning，见 torch/distributed/fsdp/fully_sharded_data_parallel.py:674-680 与 :1246-1253），但新功能只往 FSDP-2 加。新代码用 FSDP-2。

FSDP / ZeRO 这套"分片 + 临时凑齐"的思想已经被各家训练框架广泛借鉴。理解 FSDP 不只是理解 PyTorch 一个工具，是理解整个大模型训练时代的工程基础。

18.49 完整决策树：用 DDP 还是 FSDP-2？

flowchart TB
    Start[要训练大模型]
    Start --> SizeCheck{单 rank 装得下完整 params + grads + optimizer state?}
    SizeCheck -->|是| DDP[DDP - 性能最优]
    SizeCheck -->|否, 装不下 optimizer state| ZERO2[FSDP SHARD_GRAD_OP - ZeRO-2]
    SizeCheck -->|否, 连 params 也装不下| ZERO3[FSDP FULL_SHARD - ZeRO-3]
    ZERO2 --> Bandwidth{跨节点带宽足够?}
    ZERO3 --> Bandwidth
    Bandwidth -->|是, NVLink/IB| Flat[平面 mesh: 单维度 shard]
    Bandwidth -->|否, ethernet| Hybrid[HSDP 2D mesh: 节点内 shard 跨节点 replicate]
    Flat --> Compile{要 torch.compile?}
    Hybrid --> Compile
    Compile -->|是| FSDP2[FSDP-2 fully_shard API]
    Compile -->|否, 用老代码| FSDP1["FSDP-1 (维护模式)"]
    style DDP fill:#dcfce7
    style FSDP2 fill:#fef3c7,stroke:#f59e0b,stroke-width:2px
    style FSDP1 fill:#fee2e2

这条决策树覆盖常见的大模型训练场景；极端场景（如 1T+ 模型必须配 TP/PP）需要专门设计。

18.50 跨专栏关联

  • 第 16 章 ProcessGroup：FSDP 的 AllGather / ReduceScatter 都通过 ProcessGroupNCCL
  • 第 4 章 Caching Allocator：FSDP 的 unshard 触发临时分配大量显存，对 allocator 压力大，常需设 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
  • 第 7 章 §7.5.3 saved_tensors_hooks：FSDP 的 activation checkpoint 与这套 hook 配合
  • 第 13 章 AOTAutograd：开启 Traceable FSDP2 时，FSDP-2 的 collective 也走 AOTAutograd，让通信 op 进入 Inductor 的编译图（默认配置下在图外，见 §18.19）

18.51 FSDP × Engine：与第 8 章的协作

第 17 章 §17.47 讲过 DDP × Engine。FSDP 的协作更复杂：

  • forward 时插入 unshard hook：FSDP 通过 register_forward_pre_hook（第 9 章 §9.8）在 forward 前触发 unshard
  • forward 后插入 reshard hook：register_forward_hook 触发 reshard
  • backward 前通过输出张量上的 Tensor.register_hook 触发 unshard（torch/distributed/fsdp/_fully_shard/_fsdp_state.py:356 的 _register_pre_backward_hook）
  • backward 完成的 ReduceScatter 由 RegisterPostBackwardFunction（一个 torch.autograd.Function，torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py:876）+ 引擎的 final callback（torch/distributed/fsdp/_fully_shard/_fsdp_state.py:313 的 _root_post_backward_final_callback）驱动（与 DDP §17.3 的思路类似，但用 ReduceScatter 而非 AllReduce）

整套机制让 FSDP 把分片调度寄生在 PyTorch 的 hook 体系上。Engine 完全不知道 FSDP 存在，只是按 DAG 调度反向 —— hook 触发的 collective 是 Engine 视角的"普通副作用"。

这种"分布式策略寄生在框架 hook 之上"的设计是 PyTorch 分布式训练能持续演进的根本（DDP / FSDP / 用户自定义并行都共享 hook 接口）。

18.52 整章信息密度的小结

读完 Ch 18 你应该能：

  • 决策：DDP / ZeRO-2 / ZeRO-3 / HSDP 各自适用场景（§18.49 决策树）
  • 配置：选 wrap 粒度、配 MixedPrecisionPolicy、决定是否开 CPU offload
  • 理解：unshard / reshard 时机、prefetch 怎么 overlap、activation_checkpoint 怎么与 FSDP 配合
  • 调试：错误诊断速查表 + chrome trace 看 timeline
  • 接生态：HF Trainer / Lightning / DCP / torch.compile 怎么协作
  • 预判演进：FSDP-1 → FSDP-2 → 未来与 DTensor 进一步融合

70B Llama 训练的实战配置（§18.43）把整章串起来 —— 每行配置都对应某节的工程决策。这类配置在大模型训练里很常见，理解它就理解了"为什么 LLM 训练长这样、不能简化成更朴素的形式"。

18.53 设计启示

FSDP 的核心思想：

第一：显存与计算 / 通信可以互换：FSDP 把显存压力转成通信压力，让"训练大模型"从"硬件极限"变成"调优问题"

第二：分片 + 临时凑齐是分布式数据结构的通用模式：DTensor / 分布式哈希表 / 分片数据库都用这套思路。每节点常态只持有部分数据，需要时凑齐

第三：prefetch 是 overlap 通信的标配：任何"计算 - 通信 - 计算"链路都要考虑 prefetch，让通信延迟隐藏到计算时间里

第四：模块级 wrap 比模型级 wrap 灵活：FSDP-2 的 fully_shard(submodule) 思想可以借鉴到所有"框架增强 module"的场景

下一章拆序列化：torch.save / torch.load + safetensors + Distributed Checkpoint，看大模型训练的 ckpt 怎么管理。