执行摘要

• 一句话：修复 FSDP 训练中 rollout 缓冲区不同步
• 推荐动作：对于从事 FSDP 与 vLLM 集成开发的工程师，此 PR 值得精读，特别是其通过拆分权重更新来安全同步缓冲区的设计模式。测试代码中使用的桩模块技巧也可作为在无硬件依赖下进行单元测试的参考。

功能与动机

Fix FSDP-to-rollout weight sync for models whose registered buffers are updated during training. The existing rollout update path applies parameter updates through load_weights(...), but registered buffers from the FSDP state dict can be missed. This may cause rollout/training mismatch for models that rely on mutable buffers such as e_score_correction_bias.

实现拆解

1. 新增 verl/workers/rollout/vllm_rollout/weight_update_utils.py，提供 split_buffer_updates 和 apply_buffer_updates 函数。
2. 在 verl/workers/rollout/vllm_rollout/utils.py 中导入上述函数，在 _update_weights 中拆分权重，参数继续通过 load_weights 加载，缓冲区通过 apply_buffer_updates 更新。新增 _apply_buffer_updates_all_models 确保主模型和 MTP drafter 同步。
3. 调整 FP8 和标准分支的权重加载逻辑，仅传递 param_updates。
4. 新增 CPU 测试文件，通过桩模块模拟 vLLM 依赖，验证缓冲区路由和集成行为。

关键文件：

• verl/workers/rollout/vllm_rollout/weight_update_utils.py（模块 采样；类别 source；类型 core-logic；符号 split_buffer_updates, apply_buffer_updates, WeightUpdate）: 新文件，提供核心的缓冲区拆分和应用函数。
• verl/workers/rollout/vllm_rollout/utils.py（模块 采样；类别 source；类型 core-logic；符号 _apply_buffer_updates_all_models, _update_weights）: 修改了权重同步的核心路径，整合缓冲区更新。
• tests/workers/rollout/test_vllm_weight_update_utils_on_cpu.py（模块 采样；类别 test；类型 test-coverage；符号 _load_weight_update_utils, _load_vllm_rollout_utils, _ToyBlock, _ToyModel）: 新增 CPU 测试，覆盖缓冲区拆分、应用及集成逻辑。

关键符号：split_buffer_updates, apply_buffer_updates, _apply_buffer_updates_all_models

关键源码片段

verl/workers/rollout/vllm_rollout/utils.py

修改了权重同步的核心路径，整合缓冲区更新。

def _apply_buffer_updates_all_models(self, buffer_updates, main_named_buffers):
    """将缓冲区更新应用到主模型和同步的 MTP drafter 模型。
​
    The main model (yielded first) reuses the prebuilt ``main_named_buffers`` map;
    the drafter builds its own from its own model.  Returns the number of buffers
    applied to the main model.
    """
    models = list(self._iter_all_models())
    loaded = apply_buffer_updates(models[0], buffer_updates, named_buffers=main_named_buffers)
    for model in models[1:]:
        apply_buffer_updates(model, buffer_updates)
    return loaded
​
# 在 _update_weights 中（摘录核心分支）：
else:
    param_updates, buffer_updates, named_buffers = split_buffer_updates(
        self.model_runner.model, weights
    )
    # 处理参数更新（FP8 / 标准）...
    # 应用缓冲区更新到所有模型
    loaded_buffers = self._apply_buffer_updates_all_models(buffer_updates, named_buffers)

评论区精华

Review 中 gemini-code-assist 建议 split_buffer_updates 返回 named_buffers 字典以避免重复遍历模型缓冲区，作者采纳。Copilot 指出在 FP8 分支中缓冲区更新未同步到 MTP drafter 模型，最终实现通过 _apply_buffer_updates_all_models 循环所有模型解决。Copilot 建议测试应恢复 sys.modules 的修改，最终代码使用 try/finally 保证恢复。关于 non_blocking 参数的建议未采纳。

• 性能：重复遍历模型缓冲区 (performance): 作者通过 split_buffer_updates 返回 named_buffers，apply_buffer_updates 接受可选的 named_buffers 参数，避免了重复遍历。
• FP8 分支中 drafter 模型缓冲区未同步 (correctness): 最终实现中 _apply_buffer_updates_all_models 循环所有模型，包括 drafter，因此已修复。
• 测试中 sys.modules 泄露 (testing): 最终代码使用 try/finally 恢复 sys.modules 的原始状态。
• 使用 non_blocking=True 提高性能 (performance): 作者未采纳，可能出于安全考虑或因为当前非阻塞不需要。

风险与影响

• 风险：核心路径变更：_update_weights 是 rollout 权重同步的关键函数，拆分逻辑可能引入回归。测试依赖桩模块：测试使用桩模块模拟 vLLM 依赖，可能不完全反映真实行为。兼容性：若模型缓冲区未通过 register_buffer 注册，将不会被同步。性能：缓冲区数量通常较少，性能影响可忽略。
• 影响：用户影响：修复了依赖可变缓冲区（如 e_score_correction_bias）的模型在 FSDP 训练中的不一致。对于无此类缓冲区的模型无影响。系统影响：新增一个工具模块，仅 rollout 路径使用。团队影响：为 FSDP 训练提供了正确的缓冲区同步机制。
• 风险标记：核心路径变更, 测试依赖桩模块, FP8 分支兼容性, MTP drafter 同步

关联脉络

• PR #6648 [megatron] fix: MTP compatible with latest mcore: 同样修改了 verl/workers/rollout/vllm_rollout/utils.py，涉及 MTP rollout 配置，可能与本 PR 有冲突或需要协调合并。
• PR #6620 [vllm] fix: use data-parallel rank in vLLM ZMQ handles: 也修改了 verl/workers/rollout/vllm_rollout/utils.py，修复 ZMQ socket 冲突，与本 PR 处于相同文件区域。