FSDP Training Step WalkthroughVerified

本页以一个极简模型为例，逐步展示 FSDP 一个完整训练步骤中每个 rank 的存储状态和通信操作。

设定

• 2 个 GPU：Rank 0, Rank 1
• 模型：2 层，每层一个权重向量
  • Layer 1 权重：$W_1 = [w_1, w_2, w_3, w_4]$
  • Layer 2 权重：$W_2 = [w_5, w_6, w_7, w_8]$
• FSDP 分片：每个参数在 dim-0 上均分为 2 片
• 数据：Rank 0 处理 $x_0$，Rank 1 处理 $x_1$（不同的 mini-batch 子集）

初始状态

	$W_1$ 分片	$W_2$ 分片	优化器状态	激活值
Rank 0	$[w_1, w_2]$	$[w_5, w_6]$	$[w_1, w_2, w_5, w_6]$ 的 momentum/variance	—
Rank 1	$[w_3, w_4]$	$[w_7, w_8]$	$[w_3, w_4, w_7, w_8]$ 的 momentum/variance	—

每个 rank 显存 = 原始的 $1/2$。

Step 1: Forward — Layer 1

1.1 All-Gather $W_1$

Rank 0 广播 [w1, w2]  ──→  两个 rank 都获得
Rank 1 广播 [w3, w4]  ──→  完整 W1 = [w1, w2, w3, w4]

📡 通信：All-Gather，每个 rank 发送 2 个参数

1.2 前向计算

每个 rank 用完整的 $W_1$ 对自己的数据做前向传播：

• Rank 0：$a_1^{(0)} = W_1 \cdot x_0$
• Rank 1：$a_1^{(1)} = W_1 \cdot x_1$

激活值 $a_1$ 在各 rank 上不同（因为数据不同），且不参与通信。

1.3 释放非本地参数

Rank 0 丢弃 [w3, w4]，只保留自己的 [w1, w2]
Rank 1 丢弃 [w1, w2]，只保留自己的 [w3, w4]

💾 释放后显存回到分片状态。这是 FSDP 节省显存的关键——完整参数只在计算时短暂存在。

Step 2: Forward — Layer 2

2.1 All-Gather $W_2$

两个 rank 交换分片 → 各自获得完整 $W_2 = [w_5, w_6, w_7, w_8]$

2.2 前向计算

• Rank 0：$a_2^{(0)} = W_2 \cdot a_1^{(0)}$
• Rank 1：$a_2^{(1)} = W_2 \cdot a_1^{(1)}$

2.3 释放非本地参数

各 rank 只保留自己的 $W_2$ 分片。

Step 3: 计算 Loss

各 rank 独立计算 loss。由于数据不同，loss 值不同：

• Rank 0：$\mathcal{L}_0 = \text{loss}(a_2^{(0)}, y_0)$
• Rank 1：$\mathcal{L}_1 = \text{loss}(a_2^{(1)}, y_1)$

Step 4: Backward — Layer 2

反向传播从最后一层开始。

4.1 All-Gather $W_2$

完整参数前向后已释放，需要重新收集才能做反向计算。

📡 通信：All-Gather（和前向阶段相同的操作）

4.2 反向计算

每个 rank 用完整 $W_2$ 计算本地梯度：

• Rank 0：$g_2^{(0)} = \nabla_{W_2} \mathcal{L}_0$
• Rank 1：$g_2^{(1)} = \nabla_{W_2} \mathcal{L}_1$

两个 rank 的梯度不同（因为数据不同），需要求平均以保证参数更新一致。

4.3 释放非本地参数

4.4 Reduce-Scatter 梯度

                     求和 + 切分
g2(Rank 0) = [g5₀, g6₀, g7₀, g8₀]  ──┐
                                        ├──→ Rank 0 得到 [(g5₀+g5₁)/2, (g6₀+g6₁)/2]
g2(Rank 1) = [g5₁, g6₁, g7₁, g8₁]  ──┘     Rank 1 得到 [(g7₀+g7₁)/2, (g8₀+g8₁)/2]

📡 通信：Reduce-Scatter。每个 rank 只保留自己负责的那 $1/N$ 梯度（已完成求和平均）。

Step 5: Backward — Layer 1

与 Step 4 完全相同的流程：

1. All-Gather $W_1$ → 聚合完整参数
2. 反向计算 → 各 rank 得到本地梯度
3. 释放非本地参数
4. Reduce-Scatter 梯度 → 各 rank 得到自己负责的 $1/N$ 平均梯度

Step 6: Optimizer Update

每个 rank 独立更新自己的参数分片，无需任何通信：

	用到的梯度	更新的参数	用到的优化器状态
Rank 0	$\nabla[w_1, w_2, w_5, w_6]$	$[w_1, w_2, w_5, w_6]$	对应的 momentum/variance
Rank 1	$\nabla[w_3, w_4, w_7, w_8]$	$[w_3, w_4, w_7, w_8]$	对应的 momentum/variance

更新完成后，回到与初始状态相同的分片布局，准备下一个训练步骤。

单步通信汇总

阶段	操作	通信量（每 rank）
Forward Layer 1	All-Gather $W_1$	2 个参数
Forward Layer 2	All-Gather $W_2$	2 个参数
Backward Layer 2	All-Gather $W_2$	2 个参数
Backward Layer 2	Reduce-Scatter $\nabla W_2$	2 个梯度
Backward Layer 1	All-Gather $W_1$	2 个参数
Backward Layer 1	Reduce-Scatter $\nabla W_1$	2 个梯度
Optimizer Step	—	0（无通信）
总计		12 个元素

对比 DDP（只需 AllReduce 梯度）：

• DDP 通信：AllReduce 8 个梯度 ≈ 发送 $\frac{2(N-1)}{N} \times 8 = 8$ 个元素
• FSDP 通信：12 个元素（多 50%，与理论一致）

关键 takeaway

1. 参数用完即弃：完整参数只在计算时短暂存在，计算后立刻释放
2. 梯度归约即切分：Reduce-Scatter 一步完成梯度求和 + 分片分发
3. 优化器零通信：每个 rank 独立更新，互不干扰
4. 激活值不参与通信：每个 rank 的激活值只属于自己的数据子集
5. 前向 All-Gather 是 FSDP 相比 DDP 额外付出的代价：DDP 始终持有完整参数所以不需要这一步

相关页面

• FSDP 概念详解见 fsdp
• AllReduce 算法见 allreduce
• Ring AllReduce 数值示例见 ring-allreduce-walkthrough