19-FSDP-FSDP2 把大模型切开训练

FSDP 切开的不只是参数，而是一次更新里必须同时活着的参数、梯度、优化器状态和临时激活。

第三组解释了 rollout 如何把 response 生产出来，第四组导读把问题推进到 worker 内部。本文从最常用的 FSDP/FSDP2 后端切入：当 reward、old logprob 和 advantage 已经回到训练 batch 后，verl 怎样让 actor/critic 在多张 GPU 上完成一次大模型更新？

本文的核心判断是：FSDP 在 verl 里不是一个“训练策略”标签，而是一组围绕训练状态切分的工程合同。controller 只看到 update_actor()或 train_mini_batch()；FSDP engine 要负责建模、包裹、分片、micro-batch 执行、梯度裁剪、CPU offload，以及在训练后把权重重新导出给 rollout。理解这条路径，才能看懂为什么同样的 PPO/GRPO loss 会因为 fsdp/fsdp2、fsdp_size、reshard_after_forward、param_offload和 dynamic batch size 呈现不同的显存与同步成本。

先看 FSDP 在一轮 RL update 里的位置。读这张图时注意两条边界：上层 trainer 仍按算法语义组织 batch，真正的模型状态切分从 TrainingWorker进入 FSDPEngine后才发生。

FSDP 在 PPO update 中的位置

这张图对应的源码路径是：RayPPOTrainer._update_actor()把 DataProto 转成 TensorDict，写入 mini_batch_size、epochs、seed等非 tensor 字段，再调用 actor_rollout_wg.update_actor()；ActorRolloutRefWorker.update_actor()转给内部 TrainingWorker.train_mini_batch()；TrainingWorker则通过 EngineRegistry.new(..., backend=self.engine_config.strategy)创建 FSDPEngine（verl/trainer/ppo/ray_trainer.py:1205-1245，verl/workers/engine_workers.py:131-139，verl/workers/engine_workers.py:646-651）。

1. 先切 rank，再切状态

很多人理解 FSDP 时只记住一句话：参数被切到多张 GPU 上。这个说法太粗。verl 里先要决定哪些 rank 组成一个 FSDP group，哪些 rank 属于 data parallel，是否还要引入 Ulysses sequence parallel。没有这层设备网格，后面的参数 shard、loss 聚合和输出收集都没有坐标。

下面这张图把这个坐标系画出来。看图时注意：fsdp_size决定 device mesh 的形态；ulysses_sequence_parallel_size又会把 data parallel rank 和 sequence parallel rank 分开。

device mesh 如何决定 FSDP 切分边界

create_device_mesh()里，如果fsdp_size < 0或大于等于 world size，就创建一维["fsdp"]mesh；否则创建二维["ddp", "fsdp"]mesh（verl/workers/engine/fsdp/utils.py:38-56）。随后get_sharding_strategy()按 mesh 维度选择FULL_SHARD或HYBRID_SHARD（verl/workers/engine/fsdp/utils.py:59-79）。FSDPEngine._init_device_mesh()还会在ulysses_sequence_parallel_size > 1时创建["dp", "sp"]mesh，并把 data parallel size 定义为world_size // ulysses_sequence_parallel_size（verl/workers/engine/fsdp/transformer_impl.py:203-220，587-600）。

这说明 FSDP 的第一层不是“包模型”，而是“给训练状态找坐标”。PPO loss 仍然按 batch 语义计算，但每个 rank 实际拥有的是这个坐标系下的参数 shard、梯度 shard 和 optimizer state shard。

2. FSDP1 和 FSDP2 是两条包装路径

actor.strategy可以选择 fsdp或 fsdp2，但这不是简单开关。两条路径对模型包装、CPU offload 和 reshard 语义都不同。

下面这张图展示两条包装路径。它补的是：FSDP1 直接用 PyTorch FSDP(...)包住模型；FSDP2 则先选中要 fully_shard的模块，再对模块和根模型应用 composable FSDP。

FSDP1 与 FSDP2 的模型包装路径

FSDPEngine._build_model_optimizer()先加载 HF module、可选 LoRA/QAT，再调用 _build_fsdp_module()包装模型，最后才构建 optimizer 和 lr scheduler（verl/workers/engine/fsdp/transformer_impl.py:534-569）。在 FSDP1 分支里，verl 用 FSDP(module, auto_wrap_policy=..., device_mesh=..., sharding_strategy=..., mixed_precision=...)包装模型；forward-only 的 ref policy 会走 CPUOffload(offload_params=True)（verl/workers/engine/fsdp/transformer_impl.py:332-396）。

FSDP2 分支则要求 PyTorch 具备 fully_shardAPI。它构造 MixedPrecisionPolicy和可选 CPUOffloadPolicy，把 reshard_after_forward写入 fsdp_kwargs，保存一份 full state，再调用 apply_fsdp2()和 fsdp2_load_full_state_dict()（verl/workers/engine/fsdp/transformer_impl.py:397-419）。apply_fsdp2()会根据 transformer layer class、embedding 和 lm_head 选择目标模块，逐个 fully_shard()，最后再 shard 根模型；如果配置 forward_prefetch，还会给相邻 FSDPModule 设置 prefetch 目标（verl/utils/fsdp_utils.py:530-595）。

所以，FSDP1/FSDP2 的差异不是名字新旧，而是状态管理 API 的差异。对文章主线来说，关键不是背 PyTorch 术语，而是记住：verl 把两者都压到同一个 BaseEngine合同下面，让上层 trainer 不关心具体包装方式。

3. micro-batch 才是真正执行训练的单位

从算法看，PPO update 用的是 mini-batch 和 epoch；从引擎看，真正进入 forward/backward 的是 micro-batch。第 22 篇会专门拆 batch 名词，这里先看 FSDP 为什么必须在这一层工作：大模型不能把完整 PPO mini-batch 一次塞进 GPU，只能按样本数或 token 数切成更小的执行单元。

下面这张图把 train_mini_batch()到 forward_backward_batch()的路径展开。看图时注意：controller 传来的 batch 先被 worker 按 DP rank 切分，engine 再继续按 micro-batch 执行。

FSDP micro-batch 训练执行路径

TrainingWorker.train_mini_batch()会读取 mini_batch_size、epochs、seed和 dataloader 参数，把全局 mini-batch 换算成每个 DP rank 的 mini_batch_size_per_gpu，再迭代调用 self.train_batch()（verl/workers/engine_workers.py:238-326）。train_batch()会补入 use_dynamic_bsz、max_token_len_per_gpu、micro_batch_size_per_gpu等工程参数，然后调用 self.engine.train_batch()（verl/workers/engine_workers.py:328-382）。

进入 BaseEngine.train_batch()后，统一顺序是 optimizer_zero_grad()、forward_backward_batch()、optimizer_step()（verl/workers/engine/base.py:112-131）。FSDP 具体实现里，forward_backward_batch()先把 loss_mask在 data parallel group 上 all-reduce 得到全局 token 数，再调用 prepare_micro_batches()。如果启用 dynamic batch size，prepare_micro_batches()按 max_token_len_per_gpu * sp_size重排样本；否则按固定 micro_batch_size_per_gpu切 TensorDict（verl/workers/engine/fsdp/transformer_impl.py:608-645，verl/workers/engine/utils.py:58-96）。

每个 micro-batch 进入 forward_step()：数据移动到设备，准备 model inputs，调用 self.module(..., use_cache=False)，再由 loss function 计算 actor loss 或 critic loss。非 forward-only 时，FSDP engine 对每个 micro-batch 执行 loss.backward()，最后由 optimizer_step()统一裁剪梯度并更新参数（verl/workers/engine/fsdp/transformer_impl.py:1210-1251，656-702）。

这里的工程含义是：算法层的 batch 控制“用哪些样本更新几轮”，FSDP 层的 micro-batch 控制“每张 GPU 一次承受多少 token 和激活”。把两者混在一起，就很难解释显存、吞吐和 loss 归一化之间的关系。

4. offload 和 reshard 是显存账本的一部分

FSDP 训练大模型时，显存压力不只来自参数。一次 update 里还会短暂出现 activations、gradients、optimizer state、logits、entropy 计算和 all-gather 结果。verl 因此把 offload 和 reshard 变成 engine 合同的一部分，而不是外部脚本里的手动优化。

FSDPEngineConfig里和这件事直接相关的字段包括 param_offload、optimizer_offload、offload_policy、reshard_after_forward、forward_prefetch、fsdp_size和 ulysses_sequence_parallel_size；配置文件也明确区分了 FSDP1 的 param_offload/optimizer_offload与 FSDP2 的 offload_policy（verl/workers/config/engine.py:250-266，verl/trainer/config/engine/fsdp.yaml:10-63）。

FSDPEngine.initialize()在模型、optimizer 和 checkpoint manager 建好后，会根据 param_offload、optimizer_offload把模型、optimizer 和 grad offload 到 CPU（verl/workers/engine/fsdp/transformer_impl.py:175-201）。真正训练时，EngineTrainModeCtx通过 BaseEngineCtx在进入 train mode 时把需要的 model/optimizer/grad 搬回设备，退出时再按配置搬回 CPU；FSDP 的 to()方法内部会调用 load_fsdp_model_to_gpu()、load_fsdp_optimizer()、offload_fsdp_model_to_cpu()和 offload_fsdp_optimizer()（verl/workers/engine/base.py:250-264，verl/workers/engine/fsdp/transformer_impl.py:712-738，verl/utils/fsdp_utils.py:166-247）。

这类机制的取舍很直接：offload 能降低常驻显存，但会增加 CPU/GPU 搬运；reshard_after_forward能更积极地释放参数副本，但可能增加后续 all-gather 频率；dynamic batch size 能缓解长短样本混排带来的 token 峰值，但会增加重排和恢复输出顺序的复杂度。FSDP 文章不需要给出唯一最优配置，但要让读者知道这些开关对应的是同一张显存账本。

5. 训练权重还要导回 rollout

RL 后训练和普通 SFT 最大的不同之一，是 actor update 完成后，新权重必须回到 rollout 引擎。FSDP 内部保存的是分片状态，但 rollout backend 要接收的是逐 tensor 的模型权重更新。get_per_tensor_param()就是这条边界上的转换器。

下面这张图展示权重导出的路径。看图时注意：FSDP engine 先把模型加载回设备并拿到 state_dict，再处理 LoRA、key 转换和 DTensor full tensor，最后交给 rollout 或 checkpoint engine。

FSDP 权重导出到 rollout 的边界

源码里，FSDPEngine.get_per_tensor_param()会先 load_fsdp_model_to_gpu()，然后按 LoRA 是否存在、是否 merge 选择 collect_lora_params()或 self.module.state_dict()；随后 convert_weight_keys()把权重 key 转成 rollout 可以识别的形式。对 DTensor，它会调用 full_tensor()并转成 bf16，再以 (name, tensor)生成器形式返回（verl/workers/engine/fsdp/transformer_impl.py:785-858）。

这条导出路径在 ActorRolloutRefWorker.update_weights()里被消费。非 naive 模式下，worker 直接拿 self.actor.engine.get_per_tensor_param()交给 checkpoint engine 的 send_weights()；naive 模式下，它先恢复 rollout weights 内存，再通过 rollout backend 的 update_weights()同步，最后在需要时把 actor 参数重新 offload（verl/workers/engine_workers.py:663-743）。更外层的 CheckpointEngineManager.update_weights()会在非 naive 路径里 abort 未完成请求、释放 KV cache、建立 trainer/rollout 通信组、同步权重，再恢复 KV cache 和 generation（verl/checkpoint_engine/base.py:448-492）。

这也解释了为什么第 19 篇不能只停在 backward：FSDP 把大模型切开训练，但 RL 系统还必须把切开的权重重新拼成 rollout 能消费的更新流。第 24 篇会把这条同步链路展开，这里先把接口位置钉住。

小结：FSDP 是 worker 内部的状态切分合同

本文从 RayPPOTrainer._update_actor()追到 FSDPEngine，看到 FSDP/FSDP2 在 verl 里的职责链：

PPO/GRPO batch
-> TrainingWorker.train_mini_batch()
-> EngineRegistry selects fsdp/fsdp2
-> device mesh + FSDP wrapping
-> micro-batch forward/backward
-> gradient clip + optimizer step
-> get_per_tensor_param()
-> rollout weight sync

放回系列地图，第 19 篇补上的是 workers/resources -> training engine/distributed core里的第一块：大模型训练状态如何被 FSDP 切开，又如何在一次 RL update 后重新变成可同步权重。下一篇会进入 Megatron/MCore：当模型规模和 MoE 结构继续放大，系统会把 tensor、pipeline、expert 和 context parallel 都纳入同一个训练 engine 合同。

本文源码索引

• verl/trainer/ppo/ray_trainer.py:1205-1245：actor update 如何把 DataProto 转成 TensorDict，并写入 PPO mini-batch、epoch、shuffle 等字段。
• verl/trainer/ppo/ray_trainer.py:1247-1272：critic update 如何复用相似路径进入训练 worker。
• verl/workers/engine_workers.py:131-139：TrainingWorker如何通过 EngineRegistry.new()选择 fsdp/fsdp2后端。
• verl/workers/engine_workers.py:238-326：train_mini_batch()如何按 DP rank 迭代 mini-batch。
• verl/workers/engine_workers.py:328-382：worker 如何把 dynamic batch、micro-batch 等工程参数注入 train_batch()。
• verl/workers/engine/base.py:112-131：BaseEngine.train_batch()的统一训练步顺序。
• verl/workers/config/engine.py:250-266、verl/trainer/config/engine/fsdp.yaml:10-63：FSDP/FSDP2 相关配置字段。
• verl/workers/engine/fsdp/utils.py:38-79：device mesh 与 FSDP sharding strategy 的选择。
• verl/workers/engine/fsdp/transformer_impl.py:203-220、587-600：FSDP 与 Ulysses sequence parallel 的 rank 坐标。
• verl/workers/engine/fsdp/transformer_impl.py:332-440：FSDP1/FSDP2 的模型包装路径。
• verl/utils/fsdp_utils.py:530-595：FSDP2 如何选择模块并调用 fully_shard()。
• verl/workers/engine/fsdp/transformer_impl.py:608-645、1210-1251：micro-batch forward/backward 和 loss 计算路径。
• verl/workers/engine/fsdp/transformer_impl.py:656-702：FSDP/FSDP2 梯度裁剪和 optimizer step。
• verl/workers/engine/fsdp/transformer_impl.py:712-738、verl/utils/fsdp_utils.py:166-247：模型和 optimizer 在 CPU/GPU 之间的 offload/load。
• verl/workers/engine/fsdp/transformer_impl.py:785-858：get_per_tensor_param()如何导出逐 tensor 权重。
• verl/workers/engine_workers.py:663-743、verl/checkpoint_engine/base.py:448-492：训练权重如何进入 rollout 或 checkpoint engine 同步路径。