Prhub

#51440 [CI Test] Add specific unit test for mrv2 offloading

原始 PR 作者 yewentao256 合并时间 2026-08-08 06:49 文件变更 2 提交数 3 评论 5 代码增减 +58 / -10

执行摘要

为 MRv2 offload 新增定向单测并回退 CI 变更

PR body 明确说明:'A following up PR for https://github.com/vllm-project/vllm/pull/51413 We add a specific unit test for MRv2 offloading, and revert the non-meaningful ci update from previous PR'。即 #51413 引入 MRv2 weight offloading 功能后,缺乏针对该路径的定向回归测试,且前一个 PR 中混入了一些无意义的 CI 更新需要回退。

值得快速阅读,重点看 tests/basic_correctness/test_cpu_offload.py 中的 test_mrv2_weight_offloading:它展示了如何穿透 engine_core 到 GPUModelRunner,并通过 get_offloader() 直接校验 offloader 内部状态,这种测试思路对验证 MRv2 内部机制很有参考价值。对于负责 offload 或 MRv2 功能的工程师,建议精读并留意该测试对内部实现的耦合度。

讨论亮点

本 PR 没有实质性 review 评论或代码讨论线程。唯一值得一提的交互是 mergify[bot] 提示 pre-commit 检查失败并要求运行 pre-commit run --all-files,作者随后通过合并 main 分支(commit 7fdfb1a)修复;维护者 njhill 在最终 review 中直接批准,回复 'Thanks @yewentao256'。

实现拆解

  1. 在 tests/basic_correctness/test_cpu_offload.py 中新增 test_mrv2_weight_offloading 测试:通过 offload_kwargs + offloader_type 参数化覆盖 UVAOffloader(cpu_offload_gb=1)与 PrefetchOffloader(group_size/num_in_group/prefetch_step 全为 1)两条路径;monkeypatch 关闭 VLLM_ENABLE_V1_MULTIPROCESSING 并调用 envs.disable_envs_cache() 规避环境变量缓存,启动 tiny-random-Llama 完成 eager 推理后,穿透 engine_core 断言 model_runner 是 GPUModelRunner,再通过 get_offloader() 断言全局 offloader 类型及内部状态(cpu_offload_bytes > 0,或 total_offloaded_bytes > 0 且 buffer_pool 非空),finally 中恢复原始 offloader。
  2. 简化原有 test_cpu_offload:移除 use_v2_model_runner 参数化及 env1 中的 VLLM_USE_V2_MODEL_RUNNER 设置,因为 MRv2 已成为默认路径,避免重复覆盖。
  3. 清理 tests/basic_correctness/test_prefetch_offload.py:删除 pytest 导入、use_v2_model_runner 参数化以及 env1/env2 的环境变量传递,回退前 PR 中的非必要 CI 更新。
  4. 验证方式:PR body 给出了 pytest tests/basic_correctness/test_cpu_offload.py::test_mrv2_weight_offloading 在有/无 MRv2 变更时的通过/失败对比,证明该测试能有效区分新旧行为。本次为纯测试变更,不涉及源码、配置或部署配套调整。
文件 模块 状态 重要度
tests/basic_correctness/test_cpu_offload.py 卸载测试 modified 5.89
tests/basic_correctness/test_prefetch_offload.py 预取卸载 modified 4.16

关键符号

test_mrv2_weight_offloading test_cpu_offload test_prefetch_offload_llama

关键源码片段

tests/basic_correctness/test_cpu_offload.py test-coverage

本 PR 的核心变更文件:新增 test_mrv2_weight_offloading 定向测试,覆盖 UVA 与 Prefetch 两条 offload 路径;同时简化既有 test_cpu_offload,移除已冗余的 VLLM_USE_V2_MODEL_RUNNER 参数化。

# tests/basic_correctness/test_cpu_offload.py
# 新增的 MRv2 权重 offload 专用测试,直接校验 offloader 内部状态
@pytest.mark.parametrize(
    ("offload_kwargs", "offloader_type"),
    [
        ({"cpu_offload_gb": 1}, UVAOffloader),
        (
            {
                "offload_group_size": 1,
                "offload_num_in_group": 1,
                "offload_prefetch_step": 1,
            },
            PrefetchOffloader,
        ),
    ],
)
def test_mrv2_weight_offloading(vllm_runner, monkeypatch, offload_kwargs, offloader_type):
    # 关闭 v1 多进程,确保 driver worker 上能直接访问 offloader 单例
    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
    envs.disable_envs_cache() # 让本次环境变量变更对后续读取立即生效
    original_offloader = get_offloader()
​
    try:
        with vllm_runner(
            "hmellor/tiny-random-LlamaForCausalLM",
            enforce_eager=True,
            gpu_memory_utilization=0.02,
            max_model_len=128,
            max_num_seqs=1,
            **offload_kwargs,
        ) as vllm_model:
            # 穿透到 GPUModelRunner,确认走的是 v2 model runner 路径
            engine_core = vllm_model.llm.llm_engine.engine_core.engine_core
            model_runner = engine_core.model_executor.driver_worker.worker.model_runner
            assert isinstance(model_runner, GPUModelRunner)
​
            # 分别校验 UVA 与 Prefetch 两条 offload 路径的内部状态
            offloader = get_offloader()
            assert isinstance(offloader, offloader_type)
            if isinstance(offloader, UVAOffloader):
                assert offloader.cpu_offload_bytes > 0
            else:
                assert offloader.total_offloaded_bytes > 0
                assert offloader.buffer_pool is not None
    finally:
        # 恢复全局 offloader,避免影响其他测试
        set_offloader(original_offloader)
        envs.disable_envs_cache()

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

风险主要集中在测试本身:

1) 测试依赖 vllm.model_executor.offloader 与 vllm.v1.worker.gpu.model_runner 的内部结构(offloader 类名、属性如 cpu_offload_bytes / total_offloaded_bytes / buffer_pool),后续 MRv2 重构若调整这些内部符号,该测试需要同步更新;
2) 测试通过 vllm_runner 启动真实推理进程,单用例耗时约 76 秒,会显著增加 CI 时长;
3) 原 test_cpu_offload 与 test_prefetch_offload_llama 删除了 use_v2_model_runner 参数化,虽然 v2 已是默认路径,但显式覆盖 v1 旧路径的能力被移除,若后续需要重新支持 v1 runner 需恢复该参数化;
4) 测试中 monkeypatch 与 envs.disable_envs_cache() 的配合依赖全局环境缓存行为,并行执行测试时可能产生干扰,不过 finally 中恢复 offloader 降低了串扰风险。

对用户无任何运行时影响,因为未改动生产代码。对测试与 CI 体系影响较大:新增一个针对 MRv2 weight offloading 的定向回归测试,能有效保护 #51413 引入的功能,防止后续重构导致 offload 行为回退;同时清理了旧测试中已冗余的 VLLM_USE_V2_MODEL_RUNNER 参数化,减少重复用例数量。该测试模式(借助 get_offloader()/set_offloader() 对全局 offloader 状态做断言)也为后续 offload 相关测试提供了可借鉴的范式。

依赖 offloader 内部实现 CI 时长增加 减少旧参数化覆盖

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论