FSDP2 Qwen3 Decoder Layer Walkthrough

本页用一个简化的 Qwen3 decoder layer + 4 GPU FSDP2 设定，逐算子展示每一步的数据分布、本地计算、以及 DP 通信。目标是让读者能看到 FSDP2 在真实现代 Transformer 上到底在搬哪些张量、搬到哪里、搬多少。

简化的 Qwen3 结构

完整的 Qwen3 模型是：

class Qwen3Model(nn.Module):
    embed_tokens: nn.Embedding               # [vocab_size, hidden_size]
    layers: nn.ModuleList[Qwen3DecoderLayer] # N 层 decoder
    norm: Qwen3RMSNorm                        # 最终 norm
    # lm_head 通常与 embed_tokens 权重共享

本 walkthrough 聚焦于一层 Qwen3DecoderLayer——完整模型不过是把这一层重复 $N_L$ 次，每层都是独立的 FSDP 单元，机制完全相同。

配置（精简版，便于数值追踪）

参数	值
hidden_size ($d$)	512
num_attention_heads ($h_q$)	8
num_key_value_heads ($h_{kv}$)	4（GQA 2:1）
head_dim ($d_h$)	64
intermediate_size ($d_{\text{ffn}}$)	1024
DP world size ($N$)	4
每 rank batch ($B$)	2
序列长度 ($L$)	8

Decoder Layer 的结构

简化后的 Qwen3DecoderLayer（去掉 bias、position_ids、cache 等）：

class Qwen3DecoderLayer(nn.Module):
    def __init__(self, config):
        self.input_layernorm = Qwen3RMSNorm(hidden_size)
        self.self_attn = Qwen3Attention(config)
        self.post_attention_layernorm = Qwen3RMSNorm(hidden_size)
        self.mlp = Qwen3MLP(config)

    def forward(self, x):
        # Pre-norm attention + residual
        residual = x
        x = self.input_layernorm(x)
        x = self.self_attn(x)
        x = residual + x
        # Pre-norm MLP + residual
        residual = x
        x = self.post_attention_layernorm(x)
        x = self.mlp(x)
        x = residual + x
        return x

class Qwen3Attention(nn.Module):
    def __init__(self, config):
        self.q_proj = nn.Linear(d, h_q * d_h, bias=False)      # [512, 512]
        self.k_proj = nn.Linear(d, h_kv * d_h, bias=False)     # [256, 512]
        self.v_proj = nn.Linear(d, h_kv * d_h, bias=False)     # [256, 512]
        self.o_proj = nn.Linear(h_q * d_h, d, bias=False)      # [512, 512]
        self.q_norm = Qwen3RMSNorm(d_h)   # Qwen3 特有：head-dim RMSNorm
        self.k_norm = Qwen3RMSNorm(d_h)

    def forward(self, x):
        B, L, _ = x.shape
        q = self.q_proj(x).view(B, L, h_q,  d_h)
        k = self.k_proj(x).view(B, L, h_kv, d_h)
        v = self.v_proj(x).view(B, L, h_kv, d_h)
        q = self.q_norm(q)                       # Qwen3 独有
        k = self.k_norm(k)                       # Qwen3 独有
        q, k = apply_rotary_pos_emb(q, k)        # RoPE
        attn_out = sdpa(q, k, v, is_causal=True) # GQA 内部 repeat k/v
        attn_out = attn_out.reshape(B, L, h_q * d_h)
        return self.o_proj(attn_out)

class Qwen3MLP(nn.Module):
    def __init__(self, config):
        self.gate_proj = nn.Linear(d, d_ffn, bias=False)  # [1024, 512]
        self.up_proj   = nn.Linear(d, d_ffn, bias=False)  # [1024, 512]
        self.down_proj = nn.Linear(d_ffn, d, bias=False)  # [512, 1024]

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))

和经典 Transformer 的差异：Pre-Norm + RMSNorm + GQA + SwiGLU + RoPE，再加上 Qwen3 特有的 Q-norm / K-norm（head-dim 上的 RMSNorm，训练稳定性优化）。

参数列表（单层）

参数名	形状	元素数
input_layernorm.weight	[512]	512
self_attn.q_proj.weight	[512, 512]	262,144
self_attn.k_proj.weight	[256, 512]	131,072
self_attn.v_proj.weight	[256, 512]	131,072
self_attn.q_norm.weight	[64]	64
self_attn.k_norm.weight	[64]	64
self_attn.o_proj.weight	[512, 512]	262,144
post_attention_layernorm.weight	[512]	512
mlp.gate_proj.weight	[1024, 512]	524,288
mlp.up_proj.weight	[1024, 512]	524,288
mlp.down_proj.weight	[512, 1024]	524,288
单层合计		$\approx 2.36\text{M}$

fp16 下每层约 4.72 MB。

FSDP2 设置

用 4 GPU 做 FSDP2，每个 decoder layer 作为一个 FSDP 单元：

from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import fully_shard

mesh = init_device_mesh("cuda", (4,), mesh_dim_names=("dp",))

model = Qwen3Model(config)
for layer in model.layers:
    fully_shard(layer, mesh=mesh)   # 每层一个 FSDP 单元
fully_shard(model, mesh=mesh)        # 根模型

包装后每个参数被转换为 DTensor(placements=(Shard(0),))，在 dim-0 上均分到 4 个 GPU 上。FSDP2 的 pre-/post-forward hook 会挂在每个 FSDP 单元（即每一层）上。

初始状态（一层 forward 开始之前）

每个 rank 上的参数分片

参数	完整形状	每 GPU 分片 [row/4, col]	每 GPU 元素数
input_layernorm.weight	[512]	[128]	128
q_proj.weight	[512, 512]	[128, 512]	65,536
k_proj.weight	[256, 512]	[64, 512]	32,768
v_proj.weight	[256, 512]	[64, 512]	32,768
q_norm.weight	[64]	[16]	16
k_norm.weight	[64]	[16]	16
o_proj.weight	[512, 512]	[128, 512]	65,536
post_attention_layernorm.weight	[512]	[128]	128
gate_proj.weight	[1024, 512]	[256, 512]	131,072
up_proj.weight	[1024, 512]	[256, 512]	131,072
down_proj.weight	[512, 1024]	[128, 1024]	131,072
每 GPU 合计			$\approx 590\text{K}$

每 rank 单层参数显存 $\approx 1.18$ MB（fp16），正好是完整层的 $1/4$。

每个 rank 上的激活值（输入）

global batch size = 8，$N=4$ 个 DP rank，每 rank 处理 2 个样本：

Rank	持有的输入 hidden_states	来自
0	[2, 8, 512]	batch 样本 0, 1
1	[2, 8, 512]	batch 样本 2, 3
2	[2, 8, 512]	batch 样本 4, 5
3	[2, 8, 512]	batch 样本 6, 7

激活值在 FSDP 中从不分片——每个 rank 在自己的数据子集上做完整前向，产生各自完整的激活张量。这块显存与 FSDP 策略无关，只取决于 local batch 和序列长度。

Forward 逐算子 Walkthrough

Step 0：Pre-Forward Hook → All-Gather 层参数

FSDP2 在进入本层 forward 前，触发注册在该 FSDP 单元上的 pre-forward hook，做一次 All-Gather：

对每个参数：
  DTensor(Shard(0))   →   plain torch.Tensor（完整形状）

所有 rank 并行执行：每个 rank 把自己的 $1/4$ 分片广播给其他 3 个 rank，同时收齐其他 3 个 rank 的分片。底层是一次 NCCL ncclAllGather（实际上 FSDP2 会在一个单元内对多个参数批量调度，尽量合并成少数几个大 AllGather）。

通信量：每 rank 发送 $\sim \frac{N-1}{N} \times 4.72\text{MB} = 3.54$ MB。

每 rank 状态：

内容	状态
11 个参数	✅ 完整（plain Tensor）
hidden_states	[2, 8, 512]（本 rank 的数据）

此时每 rank 显存里临时有完整的 4.72 MB 层参数。

---

Step 1：input_layernorm(x)

输入（每 rank 相同结构，数据不同）：

• x: [2, 8, 512]
• input_layernorm.weight: [512]

计算：

$\text{RMSNorm}(x)_i = \frac{x_i}{\sqrt{\frac{1}{d}\sum_j x_j^2 + \epsilon}} \cdot w_i$

沿 hidden_size 维度做归一化，乘以可学习的 scale。

输出：x_norm: [2, 8, 512]

通信：无。每 rank 有完整的 weight，本地独立完成。

---

Step 2：Q / K / V 投影

输入：

• x_norm: [2, 8, 512]
• q_proj.weight: [512, 512]
• k_proj.weight: [256, 512]
• v_proj.weight: [256, 512]

计算：

q = x_norm @ q_proj.weight.T    → [2, 8, 512]
k = x_norm @ k_proj.weight.T    → [2, 8, 256]
v = x_norm @ v_proj.weight.T    → [2, 8, 256]

注意 GQA：$h_{kv} = h_q / 2 = 4$，所以 K、V 的输出维度是 Q 的一半。

Reshape 为多头：

q: [2, 8, 8, 64]   # (B, L, h_q,  d_h)
k: [2, 8, 4, 64]   # (B, L, h_kv, d_h)
v: [2, 8, 4, 64]

通信：无。三个独立的矩阵乘，每 rank 在自己的 [2, 8, 512] 上和完整权重相乘。

---

Step 3：Q-Norm / K-Norm（Qwen3 特有）

Qwen3 在算 attention 分数之前对 Q、K 的每个 head 单独做一次 RMSNorm（沿 head_dim 维度），用来稳定训练，尤其是 long-context 场景。

输入：

• q: [2, 8, 8, 64]，q_norm.weight: [64]
• k: [2, 8, 4, 64]，k_norm.weight: [64]

计算：

q = q_norm(q)   # RMSNorm 沿最后一维，形状不变
k = k_norm(k)

通信：无。

---

Step 4：RoPE

输入：q, k，加上 RoPE 的正余弦缓存（通常预计算，不算模型参数）

计算：对 q 和 k 的 head_dim 维度做旋转位置编码，形状不变。详见 rope。

通信：无。

---

Step 5：Scaled Dot-Product Attention

输入：

• q: [2, 8, 8, 64]
• k: [2, 8, 4, 64]
• v: [2, 8, 4, 64]

计算：

1. GQA repeat：K、V 沿 kv_head 维度 repeat 以匹配 Q 的 head 数（每个 KV head 被 2 个 Q head 共享）

   k: [2, 8, 4, 64] → repeat_interleave(dim=2, repeats=2) → [2, 8, 8, 64]
   v: [2, 8, 4, 64] → repeat_interleave(dim=2, repeats=2) → [2, 8, 8, 64]

   （FlashAttention 内部可以原地做 GQA，不必显式 materialize repeat 后的 tensor）

2. Attention：

   $\text{Attn}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_h}} + \text{causal\_mask}\right) V$

   attn_out: [2, 8, 8, 64]

3. Flatten heads：[2, 8, 8, 64] → [2, 8, 512]

通信：无。注意力计算是纯粹的 per-token 运算，每 rank 的样本之间互不影响。

---

Step 6：o_proj

输入：attn_out: [2, 8, 512]，o_proj.weight: [512, 512]

计算：attn_out @ o_proj.weight.T → [2, 8, 512]

通信：无。

---

Step 7：残差连接

hidden_states = residual + attn_output    # [2, 8, 512]

residual 是 Step 1 之前保存的原始 x。

通信：无。

---

Step 8：post_attention_layernorm

和 Step 1 结构完全一致，对 [2, 8, 512] 的 hidden 做 RMSNorm，输出 [2, 8, 512]。

通信：无。

---

Step 9：MLP（SwiGLU）

输入：

• h: [2, 8, 512]
• gate_proj.weight: [1024, 512]
• up_proj.weight: [1024, 512]
• down_proj.weight: [512, 1024]

计算：

gate  = h @ gate_proj.weight.T       → [2, 8, 1024]
up    = h @ up_proj.weight.T         → [2, 8, 1024]
fused = F.silu(gate) * up            → [2, 8, 1024]   # element-wise
out   = fused @ down_proj.weight.T   → [2, 8, 512]

SwiGLU 相比经典 FFN 多一个 gate_proj：gate_proj 过 SiLU 激活后与 up_proj 做点乘（element-wise），再经过 down_proj 降回 hidden_size。

通信：无。三个独立的矩阵乘 + 一个 element-wise 点乘。

---

Step 10：残差连接

hidden_states = residual + mlp_output    # [2, 8, 512]

这就是本层最终的输出 hidden_states，也是下一层 decoder layer 的输入。

---

Step 11：Post-Forward Hook → Reshard

本层 forward 结束后，FSDP2 触发 post-forward hook：

对每个参数：
  plain torch.Tensor（完整）   →   DTensor(Shard(0))

具体做的事：释放 Step 0 all-gathered 参数中”不属于本 rank”的那 $3/4$，只保留 $1/4$ 本地分片。这一步不需要通信——所有 rank 都只保留自己原本的分片。

每 rank 状态：

内容	状态
11 个参数	✅ 回到 DTensor(Shard(0))，各持 $1/4$
hidden_states	[2, 8, 512]（本层输出，进入下一层）

显存从 4.72 MB 回落到 1.18 MB。中间激活值仍然保留在每 rank 上——它们要留着给反向计算用。

---

层间过渡

下一层 decoder layer 的 forward 又会从 Step 0 开始：pre-forward hook 触发新一次 All-Gather（只 gather 下一层的参数，和本层无关），然后重复 Step 1 ~ Step 11。

层与层之间没有 DP 通信——hidden_states 在 rank 本地从上一层直接流到下一层。

Backward 简述

反向传播的流程与前向对称，关键区别是多一次 Reduce-Scatter：

Step	操作	通信
0	Pre-backward hook：All-Gather 层参数	All-Gather（$\sim 3.54$ MB/rank）
1–10	逆序执行每个算子的反向（local 计算）	无
11	Post-backward hook：Reduce-Scatter 梯度 + Reshard 参数	Reduce-Scatter（$\sim 3.54$ MB/rank）

Reduce-Scatter 的作用：

• 反向计算产出的本地梯度 ∇W（每 rank 一份完整的 [1024, 512] 等）
• Reduce-Scatter 把 4 个 rank 的完整梯度按 dim-0 分 4 片相加求平均
• 每个 rank 最终只保留自己负责的那 $1/4$ 梯度（已经是跨 rank 的平均值）

优化器更新：每个 rank 用自己的 $1/4$ 梯度 + 自己的 $1/4$ 优化器状态（momentum、variance）更新自己的 $1/4$ 参数。无通信。

实际训练中，FSDP2 的 Backward Prefetching（BACKWARD_PRE）会让 layer $i$ 的反向计算与 layer $i-1$ 的 All-Gather 重叠，所以大部分通信时间被隐藏在计算时间内。详见 fsdp 的「降低通信开销」节。

单层单步通信量汇总

阶段	操作	每 rank 通信量（fp16）
前向	All-Gather 层参数	$\frac{3}{4} \times 4.72 \approx 3.54$ MB
反向	All-Gather 层参数	$\sim 3.54$ MB
反向	Reduce-Scatter 层梯度	$\sim 3.54$ MB
单层单步合计		$\sim 10.6$ MB/rank

对于 $N_L$ 层的 Qwen3 模型，每训练步总通信量 $\approx N_L \times 10.6$ MB/rank。但由于 compute-comm overlap，实际 wall-clock 远小于这个量。

关键要点

1. 参数”用完即弃”：完整参数只在本层 forward / backward 执行期间短暂存在，其它时候保持分片状态
2. 激活值从不分片：每个 rank 始终持有自己数据子集的完整激活值，这部分显存与 FSDP 无关
3. 所有算子都是 local 计算：FSDP 只负责”参数 gather / reshard”和”梯度 reduce-scatter”；中间的 MatMul、Softmax、RoPE、LayerNorm、SwiGLU 全部在 rank 本地独立执行
4. Qwen3 的特殊算子（Q-norm、K-norm、GQA、SwiGLU、RoPE）对 FSDP 透明：它们都是 per-sample / per-token 的 local 运算，不需要跨 DP 通信。FSDP 看到的只是”一个 nn.Module 的 forward/backward”
5. 层与层之间完全无 DP 通信：hidden_states 在本 rank 上从上一层流到下一层，不需要同步
6. 一个 FSDP 单元 = 一次 All-Gather 的批量：把整层所有参数（11 个）合并成少数几次 NCCL AllGather，比对每个参数单独 gather 更高效

相关页面

• FSDP 概念详解见 fsdp
• FSDP2 DeviceMesh 配置（Shard / HYBRID_SHARD / Replicate）见 fsdp 的「分片策略」节
• 更简单的 2 GPU × 2 层模型 walkthrough 见 fsdp-training-step-walkthrough
• 4 GPU × 简单 TransformerBlock 的分片示例见 fsdp-transformer-block-sharding
• Ring AllReduce 数值示例见 ring-allreduce-walkthrough
• Transformer 架构与现代 LLM 对比见 transformer
• RoPE 位置编码见 rope