执行摘要

• 一句话：FSDP2 模型传输改非阻塞，H2D 提速约 3.9 倍
• 推荐动作：值得精读，重点看两点：一是如何在性能优化中处理“依赖未文档化平台行为”的工程姿势——用注释写明约束、用回归测试锁定行为、并严格控制影响面；二是恰当的同步语义文档化，为后续维护者规避异步传输陷阱。对同仓库写 FSDP/rollout 相关代码的开发者，建议在新增模型传输调用前先阅读该 helper 的注释。

功能与动机

PR body 说明问题来自 verl-omni 项目 Qwen-Image 全参 FSDP2 负载的 profiling：old_log_prob 与 update_actor 两个阶段都会把 actor 模型从 CPU 恢复到 GPU 再卸载回 CPU，改动前恢复模型是同步 H2D 拷贝，耗时 1.489 秒，且 D2H 结果是 pageable CPU 内存，使随后的 H2D 无法异步执行。7 月 28 日 weekly report 还测得模型传输 PCIe 利用率仅 10.9%，远低于紧随其后的 optimizer-state 传输（34.6%）。PR 的目标是让整模型 D2H/H2D 都能异步化，并刻意保持整模型传输形态，不拆成逐参数分配与拷贝。

实现拆解

实现分四步完成：

1. 核心逻辑改动（verl/utils/fsdp_utils.py）：offload_fsdp2_model_to_cpu 用 model.to("cpu", non_blocking=True) 替换 model.cpu()，load_fsdp2_model_to_gpu 用 model.to(device, non_blocking=True) 替换 model.to(device)。源码注释详细说明了 pinned 内存行为、host 访问约束、跨 stream 重载的依赖要求，并保留原有整模型传输与 empty_cache() 释放行为不变。这是唯一的 core-logic 变更，其余均为测试与配套调整。

2. CPU 契约测试（tests/utils/test_fsdp2_model_transfer_on_cpu.py，新增）：用 Mock 断言 offload_fsdp2_model_to_cpu 和 load_fsdp2_model_to_gpu 均以 non_blocking=True 调用 model.to()，锁定两个 helper 的 API 契约。

3. 分布式 CUDA 回归（tests/special_distributed/test_fsdp2_pinned_model_transfer.py，新增）：在 2 个 CUDA rank 上用 tiny Qwen2 配置构建 full-parameter FSDP2 模型，验证 D2H offload 后本地分片 is_pinned() 为真，并验证 D2H→H2D 往返数值完全一致（atol/rtol 均为 0）。

4. 测试编排与配套：tests/special_distributed/run_all.sh 追加该回归到分布式测试套件；tests/special_sanity/check_license.py 新增 NVIDIA 2026 版权头常量 license_head_nvidia_26。提交历史还显示 CPU 套件做了可选依赖兼容处理（Megatron/Qwen-VL 缺失时跳过测试、无 FlashAttention 时保留 NPU padding backend、engine graph-release 测试固定 CPU），并在合并 upstream 后回退了强制可选依赖的改动，期间的冲突仅出现在 tests/workers/test_engine_forward_step_detach_on_cpu.py。

关键文件：

• verl/utils/fsdp_utils.py（模块 模型传输；类别 source；类型 core-logic；符号 offload_fsdp2_model_to_cpu, load_fsdp2_model_to_gpu）: 唯一的核心逻辑变更文件，两个 FSDP2 整模型传输 helper 的异步化改造位于此处，是全 PR 的性能收益来源。
• tests/special_distributed/test_fsdp2_pinned_model_transfer.py（模块 分布式测试；类别 test；类型 test-coverage；符号 _local_tensor, _build_full_parameter_fsdp2_model, main）: 新增的 2 rank CUDA 分布式回归，锁定 PyTorch 未文档化的 pinned 内存行为并验证 D2H→H2D 往返正确性，是本次改动可信度的关键支撑。
• tests/utils/test_fsdp2_model_transfer_on_cpu.py（模块 单元测试；类别 test；类型 test-coverage；符号 test_offload_fsdp2_model_to_cpu_uses_non_blocking_copy, test_load_fsdp2_model_to_gpu_uses_non_blocking_copy）: 新增的 CPU mock 契约测试，用最轻量的方式锁住两个 helper 必须以 non_blocking=True 调用 model.to() 的 API 行为，可在任何环境快速执行。
• tests/special_distributed/run_all.sh（模块 测试编排；类别 test；类型 test-coverage）: 把新增的 pinned 传输回归接入分布式测试套件，确保该行为在常规 CI 中被持续验证。
• tests/special_sanity/check_license.py（模块 许可检查；类别 test；类型 config；符号 license_head_nvidia_26）: 新增 NVIDIA 2026 版权头，用于让新测试文件通过 license 检查，属于配套改动。

关键符号：offload_fsdp2_model_to_cpu, load_fsdp2_model_to_gpu, _build_full_parameter_fsdp2_model, _local_tensor, test_offload_fsdp2_model_to_cpu_uses_non_blocking_copy, test_load_fsdp2_model_to_gpu_uses_non_blocking_copy

关键源码片段

verl/utils/fsdp_utils.py

唯一的核心逻辑变更文件，两个 FSDP2 整模型传输 helper 的异步化改造位于此处，是全 PR 的性能收益来源。

@torch.no_grad()
def offload_fsdp2_model_to_cpu(model, empty_cache: bool = True):
    # PyTorch 当前在 non_blocking=True 的 CUDA 到 CPU 拷贝中分配 pinned 主机内存，
    # 该行为并未在官方文档中明示保证，因此需要分布式回归测试锁定。
    # pinned 参数还保证了随后 CPU 到 GPU 的整模型拷贝能以异步方式下发。
    #
    # 注意：D2H 拷贝完成前，CPU 张量不允许被主机安全访问；empty_cache() 不是同步点。
    # 当前调用方在同一 CUDA stream 上重新加载模型，依赖 stream 顺序保证 D2H 到 H2D 往返正确；
    # 若未来调用方在 host 读取张量，必须先同步；若在另一 stream 重载，必须显式建立 stream 依赖。
    model.to("cpu", non_blocking=True)
    if empty_cache:
        get_torch_device().empty_cache()
​
​
@torch.no_grad()
def load_fsdp2_model_to_gpu(model):
    device = get_device_id()
    # 配合 offload 时生成的 pinned 内存，整模型 H2D 拷贝可异步执行，
    # 避免把整模型传输拆成逐参数分配与拷贝，从而消除拷贝内核间的空隙。
    model.to(device, non_blocking=True)

tests/special_distributed/test_fsdp2_pinned_model_transfer.py

新增的 2 rank CUDA 分布式回归，锁定 PyTorch 未文档化的 pinned 内存行为并验证 D2H→H2D 往返正确性，是本次改动可信度的关键支撑。

def main():
    # 该回归只验证 CUDA 行为，非 CUDA 环境直接跳过
    if get_device_name() != "cuda":
        print("test_fsdp2_pinned_model_transfer skipped: pinned transfer behavior is CUDA-specific")
        return
    assert get_torch_device().device_count() >= 2, "need at least 2 GPUs for test"
    _, rank, world_size = initialize_global_process_group()
    # 构建 2 个 rank 的 DP mesh，覆盖 full-parameter FSDP2 分片场景
    device_mesh = init_device_mesh("cuda", mesh_shape=(world_size,), mesh_dim_names=("dp",))
    model = _build_full_parameter_fsdp2_model(device_mesh)
​
    # 先记录 GPU 上的本地分片作为期望值
    expected_local_params = [_local_tensor(param).detach().clone() for param in model.parameters()]
​
    # 关键断言 1：D2H offload 后，CPU 本地分片必须是 pinned 内存
    offload_fsdp2_model_to_cpu(model, empty_cache=False)
    for param in model.parameters():
        local_param = _local_tensor(param)
        assert local_param.device.type == "cpu"
        assert local_param.is_pinned(), "non-blocking FSDP2 D2H copy must produce pinned CPU parameters"
​
    # 关键断言 2：H2D 加载后数值与 offload 前完全一致（atol 与 rtol 均为 0）
    load_fsdp2_model_to_gpu(model)
    torch.cuda.synchronize()
    for param, expected in zip(model.parameters(), expected_local_params, strict=True):
        local_param = _local_tensor(param)
        assert local_param.device.type == "cuda"
        torch.testing.assert_close(local_param, expected, atol=0.0, rtol=0.0)
​
    torch.distributed.barrier()
    torch.distributed.destroy_process_group()
    if rank == 0:
        print("test_fsdp2_pinned_model_transfer passed")

tests/utils/test_fsdp2_model_transfer_on_cpu.py

新增的 CPU mock 契约测试，用最轻量的方式锁住两个 helper 必须以 non_blocking=True 调用 model.to() 的 API 行为，可在任何环境快速执行。

from unittest.mock import Mock
​
from verl.utils import fsdp_utils
​
​
def test_offload_fsdp2_model_to_cpu_uses_non_blocking_copy():
    model = Mock()
​
    fsdp_utils.offload_fsdp2_model_to_cpu(model, empty_cache=False)
​
    # 契约锁定：offload 必须使用 non_blocking=True，才能让后续 H2D 异步执行
    model.to.assert_called_once_with("cpu", non_blocking=True)
​
​
def test_load_fsdp2_model_to_gpu_uses_non_blocking_copy(monkeypatch):
    model = Mock()
    device = object()
    monkeypatch.setattr(fsdp_utils, "get_device_id", lambda: device)
​
    fsdp_utils.load_fsdp2_model_to_gpu(model)
​
    # 契约锁定：load 必须使用 non_blocking=True，配合 pinned 内存发起异步 H2D
    model.to.assert_called_once_with(device, non_blocking=True)

评论区精华

该 PR 的 review 讨论非常精简：SamitHuang 在 verl/utils/fsdp_utils.py 留下唯一一条评论“Looks clean to me”，随后 SamitHuang 与 wuxibin89 双双 APPROVED，无任何修改要求或未解决的疑虑。实质性的设计论证（为何接受 PyTorch 未文档化的 pinned 内存行为、为何保持整模型传输、异步传输的同步约束）主要写在 PR body 与源码注释中，属于“论证前置”式写法。另外 PR 声明草案与实现由 OpenAI Codex 协助生成，提交者逐行复核并对变更负责，主 commit 带有 Co-authored-by: OpenAI Codex。

• 整体改动评价（Looks clean to me） (other): 两位 reviewer 均批准，无修改要求；设计论证主要在 PR body 与源码注释中前置完成。

风险与影响

• 风险：
  1. 依赖未文档化行为：实现依赖 PyTorch 当前“非阻塞 CUDA→CPU 拷贝产生 pinned 内存”这一未在官方文档中明示保证的行为。若上游版本改变该行为，H2D 将退回同步拷贝，性能收益消失；新增分布式回归中的 is_pinned() 断言可在 CI 中捕获该回归。
  2. 异步同步语义：D2H 拷贝未完成时 CPU 张量不允许被 host 安全访问，且 empty_cache() 不是同步点。当前调用方（old_log_prob、update_actor）在同一 CUDA stream 上重载模型，靠 stream 顺序保证往返正确，但这是约定而非强制，未来新调用方若跨 stream 重载或直接在 host 读取需显式同步。
  3. 内存驻留开销：pinned 内存不可换出，全参模型 offload 后整份参数量驻留物理内存，多卡共机时主机内存压力上升。
  4. 影响面控制良好：核心功能改动仅两个 helper 共 2 行行为变化，FSDP1、optimizer offload、Megatron、fully-async 路径均未触碰，整体回归风险低。
     - 影响：对用户：colocated FSDP2 全参训练（如 Qwen-Image、verl-omni 负载）每个训练步的 old_log_prob 与 update_actor 阶段的模型重载时间从约 1.49 秒降至约 0.38 秒，减少 PCIe 空闲，直接提升训练吞吐。对系统：无 API 与配置变更，行为透明，仅改变拷贝异步性与 CPU 内存类型（pageable → pinned）。对团队：改动小而聚焦，测试完备（CPU mock 测试 + 2 rank CUDA 回归 + 纳入分布式测试套件），维护成本低。
     - 风险标记：依赖未文档化 PyTorch 行为, 异步传输同步语义易误用, 核心训练热路径, pinned 内存驻留开销

关联脉络

• PR #5081 [rollout] experimental fully-async sharded old-log-prob save/load: PR body 明确列为非重复工作：它改 fully-async 实验路径的 old-log-prob 存取，本 PR 改当前 FSDP engine 的整模型传输 helper。
• PR #5651 [megatron] pinned-copy offload for optimizer FP32 master weights: PR body 明确列为非重复工作：它用显式 pinned 拷贝卸载 Megatron optimizer 权重，不修改 FSDP2 actor 模型传输，但与本 PR 共享 pinned 内存提速思路。
• PR #7182 [fsdp] CPU offload control for forward-only ref/reward models: PR body 明确列为非重复工作：它只管 forward-only 模型的 offload 开关，不覆盖 actor 在 old_log_prob / update_actor 的手动传输。
• PR #7362 add fsdpturbo backend engine support: 历史 PR 中同属 FSDP 训练后端能力扩展，fsdp_turbo 复用 fsdp_utils 基础设施，本 PR 建立的 pinned 传输语义是其共享底座的一部分。