FSDP2 + TP/EP

背景：为什么 RL 不再用 Megatron 那一套

在大规模预训练里，Megatron-style 多维并行（DP × TP × PP × EP × CP）几乎是事实标准：模型一旦确定，整个分布式拓扑就是一张静态图，kernel 与切分方式联合优化，能把 MFU 顶到 50% 以上。

但是到了 RL Post-Training（PPO / GRPO / DAPO / On-Policy Distillation 等），训练负载的形态发生了根本性的变化：

• 一个作业里同时跑 多个角色：actor（要训）、reference / reward / critic（推理为主）、rollout（vLLM / SGLang 这类专用推理引擎）。
• 每一步训练之前，actor 的权重要同步到 rollout 引擎——训练侧用一套 parallelism，推理侧用另一套 parallelism，中间要做一次 resharding。
• 显存压力是动态的：rollout 时希望把训练参数 / 优化器状态让给 KV cache；训练时又得把权重抓回 GPU。
• 模型迭代极快（HF 格式 checkpoint 频繁更换），开发速度比峰值 MFU 更重要。

这些差异决定了 RL 系统更倾向于 FSDP 2 为主，EP/TP 按需组合 的栈，而不是 Megatron ND 并行。

FSDP + TP/EP 优势分析

Pretraining vs RL：负载特征对比

Megatron ND 并行在 RL 中的痛点

1. Reshard 代价高昂：Megatron 的权重按 (tp_rank, pp_stage) 物理切片存储，而 vLLM / SGLang 的 TP 排布、EP 排布与 Megatron 并不一一对应。每一步训练完都要做一次 train-layout → infer-layout 的转换，verl 在 Megatron 路径上专门写了 weight_converter 处理 MCore↔HF 转换，链路长而且容易出错（参考 verl/workers/megatron_workers.py 的 _build_rollout）。
2. PP 与 RL 的小 batch 性质冲突：RL 的 batch 通常被切成更小的 mini-batch（PPO epoch、GRPO group），PP 的 bubble 占比迅速放大，气泡损失甚至能吃掉一半的算力。
3. 耦合度高：模型必须改写成 Megatron 的实现（ParallelTransformerLayer 等），遇到新模型（多模态、MoE、Hybrid Attention）要专门移植，节奏跟不上算法迭代。
4. 角色复用差：reference / reward / critic 通常只需要前向，但 Megatron 的拓扑和 actor 是绑死的，难以独立调整 TP/PP 大小。

FSDP 2 + EP/TP 的核心优势

1. 天然 HF 兼容：FSDP2 wrap 之后 state_dict() 的 key 仍然是 HF 原生 fqn，且 value 是 DTensor，做一次 .full_tensor() 就能给 vLLM 的 model.load_weights(...) 直接吃。
2. DTensor 原生 resharding：FSDP2 把每个 sharded 参数表达成 DTensor，重新分发只是一次 redistribute / full_tensor，不需要手写每一种切片到切片的转换。
3. 可组合的 N-D 并行：通过 init_device_mesh 把 FSDP / EP / TP / SP 写成同一张 mesh 的不同维度，“先做 SPMD 再做 FSDP” 就能拼出 FSDP × EP、FSDP × EP × TP、HSDP × EP × TP 等拓扑。
4. 细粒度内存控制：reshard_after_forward=True + CPUOffloadPolicy(pin_memory=True) 让训练阶段的参数 / 优化器状态可以瞬时让出 GPU 给 rollout 的 KV cache，rollout 结束再 load 回来。
5. 没有 PP 气泡：RL 长上下文 + 小 batch 场景下，FSDP 的 zero 3 通信比 PP 气泡更友好；veScale 还把 reduce-scatter / all-gather 跟前后向 overlap 得很好。
6. 角色独立配置：actor / ref / reward / critic 各自的 device mesh 可以独立指定，rollout 那一侧的 (dp, infer_tp, infer_pp) 也是单独 init_device_mesh 的（参考 verl/workers/fsdp_workers.py:572-582）。
7. 代码量小：HF 模型 + 一个 apply_fsdp2 就能跑起来，不需要为每个新模型写 ParallelTransformerLayer。

FSDP + TP/EP 实现

整体思路：先做模型并行（EP/TP），再做 FSDP 2

文档开头给出的流程，本质上就是 DeviceMesh 多维拼装 + 顺序 wrap：

1
2
3
4

1. Apply EP: Expert tensors [128, H, I] -> [32, H, I]  (per EP rank)
2. Apply FSDP2 to expert modules: shard along a non-overlapping dim
3. Apply FSDP2 to regular modules: standard dim-0 sharding
4. Result: Expert params [32, H/fsdp_size, I]; regular params standard FSDP2

业界目前主要有两条实现路径：

• 路径 A（通用）：纯 FSDP2 + 智能 shard_placement_fn（代表：verl）。专家参数在训练侧并不显式做 EP，而是依赖 FSDP2 自身把 stacked tensor 切到合适的维度，rollout 侧由 vLLM 的 enable_expert_parallel 自行做 EP。
• 路径 B（显式组合）：先用 SPMD plan 把专家 / embedding 切到独立的 mesh dim 上，再叠一层 FSDP2（代表：mono_rl + veScale）。EP 在训练侧就显式存在，是真正的 FSDP × EP 二维切分。

路径 A：verl 的 apply_fsdp2 + shard_placement_fn

Verl 把每个 transformer block 单独 wrap 一次，再在 root 上 wrap 一次，关键是不再默认沿 dim-0 切，而是让一个回调挑维度：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

def apply_fsdp2(model, fsdp_kwargs, config):
    """model: AutoModelForCausalLM"""
    ...
    for module in modules:                       # 每个 transformer block
        with maybe_patch_fsdp_module(module):
            fully_shard(module, **fsdp_kwargs)
    with maybe_patch_fsdp_module(model):         # 再 wrap 一次 root
        fully_shard(model, **fsdp_kwargs)        # root 不会 reshard_after_forward


def get_shard_placement_fn(fsdp_size):
    """Choose the dimension that can divide fsdp_size to avoid padding"""
    def shard_placement_fn(param):
        shape = list(param.shape)
        for i in range(len(shape)):
            if shape[i] % fsdp_size == 0:
                return Shard(i)
        return Shard(0)
    return shard_placement_fn

verl/utils/fsdp_utils.py:507-548

实际 wire 进 actor 是在 ActorRolloutRefWorker._build_model_optimizer：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

fsdp_kwargs = {
    "mesh": fsdp_mesh,
    "mp_policy": mp_policy,
    "offload_policy": cpu_offload,
    "reshard_after_forward": fsdp_config.reshard_after_forward,
    "shard_placement_fn": get_shard_placement_fn(
        fsdp_size=self.device_mesh.shape[-1]
    ),
}
apply_fsdp2(actor_module, fsdp_kwargs, fsdp_config)
fsdp2_load_full_state_dict(actor_module, full_state, fsdp_mesh, cpu_offload)

verl/workers/fsdp_workers.py:488-510

这套方案的好处：

• 没有维度冲突：shard_placement_fn 拿到每个参数后，只挑能整除 fsdp_size 的第一个维度，stacked expert 张量 [N_expert, H, I] 也能均匀切分而不需要 padding。
• HF 模型零侵入：原生 AutoModelForCausalLM 拿来就能用，不需要改一行 modeling 代码。

代价是：训练侧没有显式 EP，所有 expert 还是被切在 FSDP mesh 上；rollout 侧 vLLM 的 EP 由 expert_parallel_size 单独配置（verl/workers/config/rollout.py:122-185），训练 / 推理 EP 不一致时由 reshard 流程兜底。

路径 B：mono_rl + veScale 的显式 EP × FSDP 2 组合

mono_rl 的 fully_shard 是个组装函数：构造一个包含 MP / EP / OE 的命名 mesh，先按 SPMD plan 切一次，再 auto_wrap 进 FSDP 2：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

named_mesh = {ParallelType.MP: default_mp_mesh}
if tp_mesh is not None:
    named_mesh["EP"] = tp_mesh         # EP mesh
if oe_mesh is not None:
    named_mesh["OE"] = oe_mesh         # over-encoded embedding mesh

model = parallelize(model, tp_plan, named_mesh)   # 1) SPMD: EP / OE 先切

for module in model.modules():
    if isinstance(module, block_cls):
        materialize(module)
        modules = auto_wrap(module, SPMDPolicy(),  # 2) 每个 block 套 FSDP2
                            fsdp_kwargs, fuse_fsdp_modules=True, verbose=True)
        ...

auto_wrap(model, SPMDPolicy(), fsdp_kwargs,        # 3) root 再套一次
          fuse_fsdp_modules=False, verbose=True)
apply_spmd_extension(model, set_mesh_attr=True)

mono_rl/worker/engine/fsdp/vescale/fully_shard.py:279-353

每个模型自己提供 EP/OE 的 plan，比如 M 12（Seed 自研多模态 MoE）：

1
2
3
4
5
6
7
8
9

def make_m12_plan(config):
    plan = {
        "OE": { "*over_encoded_embeddings.embedding_list.0.weight": 0 },
        "EP": {
            "*.moe.experts.gate_up_proj": 0,   # 沿 expert 维度切
            "*.moe.experts.down_proj":   0,
        },
    }
    return plan

mono_rl/models/seed_models/modeling_m12.py:82-103

也就是说：128 个 expert 的 stacked weight [128, H, I]，先沿 dim-0（expert 轴）按 EP=4 切成 4 份 [32, H, I]，每个 EP rank 只持有自己那 32 个 expert 的参数。这就是文档开头第 1 步 [128,H,I] -> [32,H,I] 的来源。

关键约束：FSDP 2 的切分维度不能与 EP/TP 重叠

为什么文档说"FSDP2 在 dim-1 (hidden dim) 切分"？因为 PyTorch 的 DeviceMesh 在嵌套 SPMD + FSDP 时存在一个限制——同一个张量维度不能被 SPMD 和 FSDP 同时切。veScale 在 shard_state 里给出了显式策略：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

# if the spmd and fsdp cuts the same dimension,
# torch has bugs in handling this case due to the mismatched order
# in device mesh with parallelism order. More details refer to
# https://github.com/pytorch/pytorch/issues/129206.
# Here we bypass this bug by making sure the shard dimension doesn't
# overlap between fsdp and spmd
spmd_shard_dims = {shard.dim for shard in spec.placements if shard.is_shard()}
remain_dims = [dim for dim in range(global_tensor.ndim)
               if dim not in spmd_shard_dims]
if len(remain_dims) == 0:
    fsdp_placements = [Replicate() for _ in range(mesh_info.mesh.ndim)]
else:
    fsdp_size = mesh_info.shard_mesh_size
    dimlens = [(dim, global_tensor.size(dim)) for dim in remain_dims]
    selected = max(dimlens, key=lambda x: x[1])[0]
    for dim, dimlen in dimlens:
        if dimlen % fsdp_size == 0:
            selected = dim
            break
    fsdp_placements = [Replicate() for _ in range(mesh_info.mesh.ndim)]
    fsdp_placements[mesh_info.shard_mesh_dim] = Shard(selected)

vescale/parallel/fsdp2/extension/spmd.py:265-286

这段逻辑做了三件事：

1. 统计 SPMD（EP/TP）已经切过的维度。
2. 从剩余维度里挑一个：优先挑能整除 fsdp_size 的，否则挑最大的，避免 padding。
3. 生成 FSDP placement：在 fsdp 这一维写 Shard(selected)，其余维度 Replicate()。

回到 expert 张量 [N=128, H, I]：

• EP 切了 dim-0 → spmd_shard_dims = {0}
• FSDP 必须从 {1, 2} 里挑 → 通常挑 hidden 维 (dim-1)
• 最终 [128, H, I] --EP--> [32, H, I] --FSDP--> [32, H/fsdp, I]

torch native FSDP2 也支持这个语义：shard_placement_fn 可以返回非 Shard(0) 的 placement，但要求该维度能整除 fsdp mesh size：

1
2
3
4
5
6

shard_placement_fn (Optional[Callable[[nn.Parameter], Optional[Shard]]]):
    This callable can be used to override the sharding placement for a
    parameter to shard a parameter on a dimension other than dim-0.
    If sharding on a nonzero dim, we currently require even sharding,
    i.e. the tensor dim size on that dim must be divisible by the FSDP
    Shard mesh size.

vescale/parallel/fsdp2/fully_shard.py:168-177

N-D Mesh 拓扑总结

veScale 在 _fsdp_param.py 里把支持的拓扑总结得很清楚：

1
2
3
4
5
6

# FSDP requires the DP and model parallel TP/EP mesh to have the same parent mesh
...
assert 2 <= self._spmd_mesh.ndim <= 4, (
    "_spmd_mesh.ndim can only be 2 (FSDP+TP/EP), "
    "3 (FSDP+EP+TP, HSDP+TP/EP), or 4 (HSDP+EP+TP)."
)

vescale/parallel/fsdp2/_fsdp_param.py:155-176

也就是说在生产里我们最常见 4 种组合：

训练 ↔ 推理：权重 Resharding

整套机制最关键的就是这一步——也是 RL 框架和预训练框架最不一样的地方。

verl 的 rollout_mode 是个标准范式：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

async def rollout_mode(self):
    if self._is_offload_param:
        load_fsdp_model_to_gpu(self.actor_module_fsdp)

    params = self.actor_module_fsdp.state_dict()       # HF-style fqn
    params = convert_weight_keys(params, ...)

    if self._is_offload_param:
        offload_fsdp_model_to_cpu(self.actor_module_fsdp)

    device = get_device_id()
    per_tensor_param = (
        (name,
         param.to(device, non_blocking=True).full_tensor()  # DTensor → full
            if isinstance(param, DTensor) else param)
        for name, param in params.items()
    )

    if self.config.rollout.free_cache_engine:
        await self.rollout.resume(tags=["weights"])

    await self.rollout.update_weights(
        per_tensor_param, peft_config=peft_config,
        base_sync_done=self.base_sync_done,
    )

verl/workers/fsdp_workers.py:636-707

到 vLLM 那一侧：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

def update_params(self, updated_params, peft_config=None):
    model = self.model_runner.model
    ...
    from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader
    patch_vllm_moe_model_weight_loader(model)         # MoE 融合权重 patch
    device = get_device_id()
    loaded_params = model.load_weights(
        ((name,
          param.to(device, non_blocking=True).full_tensor()
            if isinstance(param, DTensor) else param)
         for name, param in updated_params.items())
    )

verl/workers/sharding_manager/fsdp_vllm.py:283-355

整个流程有几个关键设计点：

1. 生成式（streaming）权重传递：per_tensor_param 是个 generator，按参数粒度逐个 full_tensor() → 喂给 vLLM，避免一次性 all-gather 整个模型。
2. HF fqn 直接对齐 vLLM：vLLM 的 model.load_weights 本来就吃 HF 格式的 (name, tensor) 流，FSDP 2 的 state_dict 天然兼容，无需写"切片到切片"的复杂转换。
3. MoE 融合权重 patch：vLLM 把 gate_up_proj 这种 stacked expert weight 用 fused weight loader 装载，verl 单独打了个 patch 让 DTensor 也走这条路（verl/utils/vllm/patch.py:56-72）。
4. EP 不一致时由 reshard 兜底：训练侧 FSDP（或 FSDP × EP）和 vLLM 侧的 EP 大小可以不一样—— full_tensor() 把张量收齐后由 vLLM 重新按自己的 EP/TP 切回去。
5. 配合 sleep / resume 让出显存：self.rollout.resume(tags=["weights"]) 与 FSDP 的 offload_policy=CPUOffloadPolicy() 配合，rollout 时训练参数被踢到 CPU；rollout 结束再 load 回来。

我们解决了哪些问题

把上面散落的点收拢一下，FSDP + EP/TP 这套栈在 RL 系统里实际解决了：

1. 训练-推理 resharding 难题：DTensor + HF state_dict 让"FSDP layout → vLLM layout"变成一个 streaming full_tensor() 调用，不再需要为每对 (train_tp, infer_tp, ep) 写转换。
2. MoE 切分与 stacked expert weight：通过 SPMD plan 把 expert 维度（dim-0）独立成 EP mesh，再让 FSDP 在剩余维度（hidden dim）上切，避免与 EP 冲突，也避开 PyTorch 的 mesh 排序 bug（issue #129206 ）。
3. Rollout / Train 显存切换：reshard_after_forward=True + CPUOffloadPolicy + vLLM 的 sleep / resume，让同一批 GPU 在 rollout 阶段把空间让给 KV cache，在 train 阶段把权重抓回来。
4. HF 生态零成本接入：HuggingFace 出新模型 → AutoModel.from_pretrained → apply_fsdp2，几乎零移植成本；Megatron 路径上要花数周写 ParallelTransformerLayer 和 weight_converter。
5. 多角色异构拓扑：actor / ref / reward / critic / rollout 各自独立 init_device_mesh，可以分别选 FSDP size、EP size、TP size，而不像 Megatron 那样全部绑死在同一个 ND grid 上。
6. 长上下文 + 小 batch 友好：RL 的 batch 在 PPO / GRPO epoch 切分后非常小，PP 气泡占比直线上升；FSDP zero 3 通信能与前后向 overlap，反而比 PP 更适合 RL。
7. 观测与可调试性：每个 DTensor 都有显式 placement，权重值 / 梯度可以直接 .full_tensor() 拿出来对比，不像 Megatron 切片后必须借助 dist debugger 才能拼回原貌。

值得提一句的是：不是说 Megatron 错了——大模型 pretraining 仍然是 Megatron + MCore 的天下。但是 RL Post-Training 的负载特征（多角色、resharding 重、迭代快）确实更适合 FSDP 2 + 按需 EP/TP 这套栈，这也解释了为什么 verl / OpenRLHF / mono_rl / TRL 等 RL 框架都不约而同地选择了这条路线。

参考资料

• FSDP2 Under the Hood — A Deep Dive：FSDP2 的 FSDPState / FSDPParamGroup / FSDPParam 内部实现拆解。
• verl PyTorch FSDP Backend：verl 官方文档里 FSDP 1/FSDP 2 与 vLLM 的对接说明。
• verl PR #1026 - support fsdp2 training and inference in fsdp_workers：FSDP 2 在 verl 的接入 PR，讨论了 8% / 27% 的显存收益。
• PyTorch issue #129206 - DeviceMesh ordering with FSDP + SPMD：FSDP/SPMD 切分维度不能重叠的根本原因。
• torchtitan — fsdp.md：torchtitan 关于 FSDP 2 的设计说明。
• 代码参考：
  • verl/utils/fsdp_utils.py、verl/workers/fsdp_workers.py、verl/workers/sharding_manager/fsdp_vllm.py
  • mono_rl/worker/engine/fsdp/vescale/fully_shard.py、mono_rl/models/seed_models/modeling_m12.py
  • vescale/parallel/fsdp2/fully_shard.py、vescale/parallel/fsdp2/extension/spmd.py、vescale/parallel/fsdp2/_fsdp_param.py
• https://zhuanlan.zhihu.com/p/1990790333823481023