λLLM Wiki
5 min read#fsdp#distributed-training#allreduce#transformer#all-gather#reduce-scatter

FSDP Transformer Block Sharding Example

M
memset0
Apr 9, 2026

本页用一个具体的 TransformerBlock + 4 GPU 设定，完整展示 FSDP 的分片布局、All-Gather 机制、包装粒度，以及 FSDP1 / FSDP2 的代码对比。

模型定义

一个 12 层的 Transformer，每个 TransformerBlock 包含：

TransformerBlock (hidden_dim=512, num_heads=8, ffn_dim=2048)
├── LayerNorm  (norm1)       — 512 个 scale 参数
├── Self-Attention
│   ├── W_Q  : [512, 512]   — 262,144 个参数
│   ├── W_K  : [512, 512]   — 262,144
│   ├── W_V  : [512, 512]   — 262,144
│   └── W_O  : [512, 512]   — 262,144
├── LayerNorm  (norm2)       — 512 个 scale 参数
└── FFN (MLP)
    ├── W_up  : [512, 2048]  — 1,048,576
    └── W_down: [2048, 512]  — 1,048,576

单个 TransformerBlock 参数总量 ≈ 3.15M，fp16 下约 6.3 MB。12 层合计 ≈ 37.8M 参数。

FSDP 包装粒度

FSDP 允许选择以什么为单位做分片和通信。这个单位称为 FSDP 单元(FSDP unit)。包装粒度直接决定了 All-Gather 的范围和峰值显存：

包装方式FSDP 单元All-Gather 范围前向峰值显存评价
整个模型一个单元整个 Transformer（12 层）一次性聚合全部 37.8M 参数完整模型大小❌ 失去分片意义
每个 TransformerBlock 一个单元一个 TransformerBlock（3.15M 参数）每次只聚合一层的参数一层完整参数 ≈ 6.3 MB✅ 推荐：显存低 + 通信批量化
每个 Linear 层一个单元一个 Linear（如 W_Q: 262K 参数）每次只聚合一个权重矩阵一个权重矩阵⚠️ 通信次数过多

推荐做法是以 TransformerBlock 为 FSDP 单元——这是显存节省和通信开销之间的最佳平衡点。

4 GPU 分片布局

使用 4 个 GPU，FSDP 在 dim-0 上均分每个参数。以一个 TransformerBlock 为例：

FSDP2（DTensor per-parameter 分片）

每个参数独立在 dim-0 上切分为 4 片：

参数原始形状每个 GPU 持有每 GPU 参数量
WQW_Q[512, 512][128, 512]65,536
WKW_K[512, 512][128, 512]65,536
WVW_V[512, 512][128, 512]65,536
WOW_O[512, 512][128, 512]65,536
WupW_{up}[512, 2048][128, 2048]262,144
WdownW_{down}[2048, 512][512, 512]262,144
合计786,432

每个 GPU 存储该层约 786K 参数 ≈ 1.57 MB (fp16)，是完整层的 1/41/4。

具体示意——以 WQW_Q [512, 512] 为例：

原始 W_Q (512 行 × 512 列):
┌─────────────────────────┐
│  行 0~127   (128 行)    │ ← GPU 0 持有
├─────────────────────────┤
│  行 128~255 (128 行)    │ ← GPU 1 持有
├─────────────────────────┤
│  行 256~383 (128 行)    │ ← GPU 2 持有
├─────────────────────────┤
│  行 384~511 (128 行)    │ ← GPU 3 持有
└─────────────────────────┘

FSDP1（FlatParameter 模块级分片）

FSDP1 不按参数切分，而是把整个模块的所有参数压平拼接成一个 1D 张量 FlatParameter，再均分：

FlatParameter = concat(flatten(W_Q), flatten(W_K), ..., flatten(W_down))
             = 一个长度 3,145,728 的 1D 张量

GPU 0: FlatParameter[0 : 786,432]
GPU 1: FlatParameter[786,432 : 1,572,864]
GPU 2: FlatParameter[1,572,864 : 2,359,296]
GPU 3: FlatParameter[2,359,296 : 3,145,728]

每个 GPU 持有的是一段连续的数字，不对应某个具体参数——一个 GPU 的分片可能横跨 WQW_Q 的后半部分和 WKW_K 的前半部分。

All-Gather 机制详解

当一个 TransformerBlock 的前向计算即将开始时，FSDP 需要把分散在 4 个 GPU 上的参数分片聚合为完整参数。这就是 All-Gather 操作。

完整流程（以 FSDP2 为例）

Step 1: All-Gather 参数

以 WQW_Q 为例：

GPU 0 持有 W_Q[0:128, :]    ─┐
GPU 1 持有 W_Q[128:256, :]   │  All-Gather
GPU 2 持有 W_Q[256:384, :]   │  ──────────→  每个 GPU 都获得完整的
GPU 3 持有 W_Q[384:512, :]  ─┘               W_Q [512, 512]

同时对 WK,WV,WO,Wup,WdownW_K, W_V, W_O, W_{up}, W_{down} 也做 All-Gather。All-Gather 完成后，每个 GPU 在显存中临时持有该层的完整参数（6.3 MB）。

每个 GPU 的 All-Gather 发送量 = N−1N×层参数量=34×6.3≈4.7\frac{N-1}{N} \times \text{层参数量} = \frac{3}{4} \times 6.3 \approx 4.7 MB。

Step 2: 前向计算

每个 GPU 使用完整的 WQ,WK,WV,WOW_Q, W_K, W_V, W_O 对自己的数据子集做 Self-Attention，再用完整的 Wup,WdownW_{up}, W_{down} 做 FFN。

注意：4 个 GPU 的输入数据各不相同（数据并行），因此计算出的激活值也不同，且不需要通信。

Step 3: 释放非本地参数

计算完成后，每个 GPU 立即丢弃从其他 GPU 收集来的参数分片：

GPU 0:
  计算中：持有完整 W_Q [512, 512] = 6.3 MB   ← 峰值
  计算后：只保留 W_Q[0:128, :] 等本地分片 = 1.57 MB  ← 回到分片状态

这就是 FSDP 的核心 trade-off：用通信换显存。完整参数只在计算时短暂存在。

FSDP1 的 All-Gather

FSDP1 的 All-Gather 作用在 FlatParameter 上：

GPU 0 持有 FlatParam[0:786K]    ─┐
GPU 1 持有 FlatParam[786K:1.57M] │  All-Gather
GPU 2 持有 FlatParam[1.57M:2.36M]│  ──────────→  每个 GPU 都获得
GPU 3 持有 FlatParam[2.36M:3.15M]┘               完整 FlatParam[0:3.15M]

然后 FSDP1 把这个 1D 张量反压平(unflatten) 还原为 WQ,WK,…W_Q, W_K, \ldots 的原始形状。

关键区别：即使当前计算只需要 WQW_Q（262K 参数），FSDP1 也必须 All-Gather 整个 FlatParameter（3.15M 参数）。FSDP2 原则上允许更细粒度的通信，但当前实现中仍以 FSDP 单元为粒度批量 All-Gather。

反向阶段的通信

反向传播时，每层需要两次通信：

┌────────────────────────────────────────────────┐
│ Backward - Layer i                              │
│                                                 │
│ 1. All-Gather 参数 (和前向一样，因为已释放)       │
│ 2. 反向计算 → 各 GPU 得到本地梯度                │
│ 3. 释放非本地参数                                │
│ 4. Reduce-Scatter 梯度:                         │
│    GPU 0 的 ∇W = [g0,g1,g2,g3] ─┐              │
│    GPU 1 的 ∇W = [g0,g1,g2,g3]  │ Reduce       │
│    GPU 2 的 ∇W = [g0,g1,g2,g3]  │ Scatter      │
│    GPU 3 的 ∇W = [g0,g1,g2,g3] ─┘              │
│    → GPU 0 得到 avg(g0)                          │
│    → GPU 1 得到 avg(g1)                          │
│    → GPU 2 得到 avg(g2)                          │
│    → GPU 3 得到 avg(g3)                          │
└────────────────────────────────────────────────┘

每个 GPU 只保留自己负责的 1/41/4 梯度。

FSDP1 vs FSDP2 代码对比

模型定义（两者共用）

import torch
import torch.nn as nn

class TransformerBlock(nn.Module):
    def __init__(self, hidden_dim=512, num_heads=8, ffn_dim=2048):
        super().__init__()
        self.norm1 = nn.LayerNorm(hidden_dim)
        self.self_attn = nn.MultiheadAttention(hidden_dim, num_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(hidden_dim)
        self.ffn = nn.Sequential(
            nn.Linear(hidden_dim, ffn_dim),
            nn.GELU(),
            nn.Linear(ffn_dim, hidden_dim),
        )

    def forward(self, x):
        x = x + self.self_attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
        x = x + self.ffn(self.norm2(x))
        return x

class Transformer(nn.Module):
    def __init__(self, num_layers=12, hidden_dim=512, num_heads=8, ffn_dim=2048):
        super().__init__()
        self.layers = nn.ModuleList([
            TransformerBlock(hidden_dim, num_heads, ffn_dim)
            for _ in range(num_layers)
        ])

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

FSDP1：包装器 API

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, ShardingStrategy
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from functools import partial

model = Transformer()

# 自动包装策略：以 TransformerBlock 为单位
auto_wrap_policy = partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={TransformerBlock},
)

# FSDP(module) 返回一个新的包装器对象
model = FSDP(
    model,
    auto_wrap_policy=auto_wrap_policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD,  # ZeRO Stage 3
    use_orig_params=True,  # 保留原始参数名（否则只有 FlatParameter）
)

# 包装后：
# model 是 FSDP 对象
# model.layers[0] 也是 FSDP 对象（被自动包装）
# 每个 TransformerBlock 的所有参数被压平为一个 FlatParameter

FSDP2：就地转换 API

from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy

model = Transformer()

# 混合精度策略
mp_policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,   # 前向/反向用 bf16
    reduce_dtype=torch.float32,   # 梯度归约用 fp32
)

# fully_shard(module) 就地修改，不返回新对象
# 先包装子模块，再包装根模型
for layer in model.layers:
    fully_shard(layer, mp_policy=mp_policy)  # 每层独立分片
fully_shard(model, mp_policy=mp_policy)       # 根模型

# 包装后：
# model 仍然是原始的 Transformer 类型（不是包装器）
# model.layers[0].self_attn.in_proj_weight 是 DTensor（Shard(0)）
# 每个参数都是独立的 DTensor，保留原始名称和 requires_grad

关键区别总结

# FSDP1: 查看参数
for name, param in model.named_parameters():
    print(name, param.shape)
# → _fsdp_wrapped_module._fpw_module.layers.0._fsdp_wrapped_module.flat_param
#   torch.Size([786432])  ← 一个压平的大张量

# FSDP2: 查看参数
for name, param in model.named_parameters():
    print(name, type(param))
# → layers.0.self_attn.in_proj_weight  DTensor(Shard(0))  ← 保留原始结构
# → layers.0.self_attn.out_proj.weight DTensor(Shard(0))
# → layers.0.ffn.0.weight              DTensor(Shard(0))
# → ...

FSDP1 包装后参数名变成 flat_param，无法按名字访问单个权重。FSDP2 保留原始参数结构，每个参数是一个独立的 DTensor。

这个区别直接影响了部分参数冻结的可行性——FSDP1 的 FlatParameter 要求所有参数的 requires_grad 一致，而 FSDP2 的 DTensor 允许混合冻结/可训练参数（详见 raw/note/fsdp-partial-parameter-freezing.md）。

相关页面

Comments