Prhub

#56904 [Bugfix] Make GPU sync checks safe under torch.compile

原始 PR 作者 khluu 合并时间 2026-09-16 12:37 文件变更 3 提交数 4 评论 10 代码增减 +32 / -3

执行摘要

修复 GPU 同步检查包装器在 torch.compile 下的崩溃

PR body 明确指出问题根因:"gpu_sync_debug's patched tensor-copy wrappers currently read a ContextVar before checking whether Dynamo is tracing. A full-graph compile therefore fails with an unsupported ContextVar.get() instead of bypassing sync checks as intended."。触发场景是 CI 逃逸回归:Daily main build #88920 首次在 H200 Multimodal Models (Extended Generation 1) job 的 test_voxtral_realtime_cudagraph[FULL_DECODE_ONLY] 中暴露;PR #51167 引入了 Voxtral 编译路径,但该 job 在它的 PR build #88857 和合并后 main build #88899 都被阻塞,导致改动未被真正验证。作者检索 open issue / PR 未发现既有修复,因此需要在本 PR 中一并补上源码修复与测试隔离。

值得精读,尤其是正在做 Dynamo/Inductor 集成或维护 compile 相关测试的工程师:

  1. 关注 _active_check_mode 的短路顺序修正——它揭示了“ContextVar 读取本身就是图断点”这一常被忽视的约束;同一模式可套用到其他在追踪期读取 ContextVar / threading.local 的调试代码。
  2. 关注 _torch_to_in_graph 的设计:allow_in_graph 需要可弱引用对象,因此“Python wrapper 转发原始 C++ 描述符”是通用手法;建议顺手评估 _checked_cuda / _checked_copy_ 是否需要同样处理,避免下一个模型再次踩坑。
  3. 关注测试隔离决策链:先解决显存残留,再解决 fork 与 CUDA 初始化的冲突,最终落到 method="spawn";任何新增 compile-heavy 多模态用例都应默认带上该装饰器。
  4. 评审时无需重复核对模型输出正确性(作者已说明生产改动只改 debug wrapper 控制流),但应确认 cherry-pick 到 v0.30.0 后是否需要重跑等价门禁。
讨论亮点

本 PR 的 review_comments 为 0,人工评审只有 njhill 的 /ci test 和最终 APPROVED;claude[bot] 因 fork 关闭了自动评审。真正的技术交锋全部来自作者 khluu 在 PR 上贴出的“精确 SHA + 定向 lane”验证日志,且暴露了两个非直觉问题:

  1. 直接注册 C++ 描述符不可行的发现:"torch.compiler.allow_in_graph(_torch_to) attempts to weak-reference the saved C++ method_descriptor, so importing gpu_sync_debug.py failed with TypeError: cannot create weak reference to 'method_descriptor' object in both selected H200 lanes." 这直接决定了必须引入 Python wrapper _torch_to_in_graph。

  2. 测试隔离的两段式翻车:先是内存污染——“The following FULL_AND_PIECEWISE parameter started with only 25.73/32.5 GiB free because the prior in-process compiled runner retained about 6.5 GiB, then failed the normal 0.9 startup guard”;改成新进程后又遇到 fork 限制——“Earlier tests had already initialized CUDA in the pytest parent, so both child processes failed immediately with Cannot re-initialize CUDA in forked subprocess”,最终锁定 method="spawn"。

  3. 验证过程本身的可靠性教训:#89010 因手抄的完整 SHA 无效而在 bootstrap 阶段失败、不产生任何测试证据;#89011 的 head 里 spawn 改动未 stage,仍是默认模式。作者据此在正文中明确纠正引用关系,并强调“literal head”门禁的重要性。

  4. 最终的 terminal evidence:Buildkite #89026 在 f7f21d97caec 上四类 job 全绿,两个决定性用例 test_voxtral_realtime_cudagraph[FULL_DECODE_ONLY] 与 [FULL_AND_PIECEWISE] 均通过。

实现拆解

  1. 短路顺序修正(vllm/utils/gpu_sync_debug.py::_active_check_mode):把 torch.compiler.is_compiling() 提到 _allow_depth.get() 之前求值(原顺序为 not _sync_check_enabled or _allow_depth.get() or torch.compiler.is_compiling())。_sync_check_enabled 只是普通模块全局,先判它无副作用;而 _allow_depth / _checking 两个 ContextVar 一旦在 Dynamo 追踪期间被读取就会报 unsupported ContextVar.get()。新顺序保证编译期在触碰 ContextVar 前返回 None,_checking.get() 也不会被求值,_checked_to / _checked_cuda / _checked_copy_ 的检查体在编译期整体退化为 no-op 直通。

  2. 为 Tensor.to 建立编译期逃逸通道(新增 _torch_to_in_graph):_install_copy_checkers 会替换 torch.Tensor.to 等属性,Dynamo 之后只能通过模块全局 _torch_to 看到导入期捕获的 C++ 方法描述符。直接用 torch.compiler.allow_in_graph(_torch_to) 会在 import 阶段失败,因为 allow_in_graph 内部需要弱引用该对象。实现改为注册一个可弱引用的 Python wrapper,内部仍调用原始描述符,从而既保持无递归,又让 Dynamo 把这次调用当作不透明算子,AOT/Inductor 仍能 lower 到底层 C++ 调用。

  3. 切换调用点(vllm/utils/gpu_sync_debug.py::_checked_to):把末尾的 return _torch_to(self, *args, **kwargs) 改为 return _torch_to_in_graph(self, *args, **kwargs)。这是唯一的行为入口切换;由于 _active_check_mode() 在编译期已返回 None,_parse_to 与 implicit sync 判定都不会在图中执行,包装器只承担“把原始 to 交给编译器”的职责。

  4. 测试配套:focused 回归用例(tests/utils_/test_gpu_sync_debug.py):新增 test_copy_checker_is_traceable_while_compiling,先用 monkeypatch 打开 _sync_check_enabled 并把模式设为 error,显式调用 gsd._install_copy_checkers() 复现包装后的运行态;内部函数 convert_like(x, like) 使用 x.to(device=like.device, dtype=like.dtype),严格复刻 Voxtral 中 device/dtype 来自另一个张量的动态调用;用 torch.compile(convert_like, fullgraph=True)(默认 backend,覆盖 AOT/Inductor)验证 dtype 是否被正确转换。相比早期草稿,这里放弃了 backend="eager",避免掩盖真实编译路径问题。

  5. 测试配套:Voxtral CUDA graph 用例进程隔离(tests/models/multimodal/generation/test_voxtral_realtime.py):给 test_voxtral_realtime_cudagraph 加上 @create_new_process_for_each_test(method="spawn"),并补充 create_new_process_for_each_test 的导入。两个动因在讨论中被逐一验证:其一是 FULL_DECODE_ONLY 的进程内编译/CUDA graph 捕获会残留约 6.5 GiB 显存,使后续 FULL_AND_PIECEWISE 参数只剩 25.73/32.5 GiB 可用、撞上 0.9 启动守卫并级联失败;其二是默认隔离辅助函数在父 pytest 进程已初始化 CUDA 时选择 fork,子进程直接报 Cannot re-initialize CUDA in forked subprocess,必须显式指定 spawn。

  6. 无 schema / 配置 / 部署改动:唯一的环境变量仍是既有的 gpu_sync_debug 开关;PR 被标记到 v0.30.0 cherry picks milestone,意味着验证证据需要在 cherry-pick 后重跑。

文件 模块 状态 重要度
vllm/utils/gpu_sync_debug.py 同步调试 modified 6.82
tests/utils_/test_gpu_sync_debug.py 同步调试 modified 5.69
tests/models/multimodal/generation/test_voxtral_realtime.py 实时语音 modified 4.13

关键符号

_active_check_mode _torch_to_in_graph _checked_to test_copy_checker_is_traceable_while_compiling convert_like test_voxtral_realtime_cudagraph

关键源码片段

vllm/utils/gpu_sync_debug.py core-logic

核心修复文件:调整 `_active_check_mode` 短路顺序避免编译期读取 ContextVar,并新增 `_torch_to_in_graph` 让被替换的 C++ `Tensor.to` 描述符可被 Dynamo 接纳,`_checked_to` 改走该入口。

# 导入期捕获原始 C++ 方法描述符;_install_copy_checkers 之后会替换这些属性。
_torch_to = torch.Tensor.to
_torch_cuda = torch.Tensor.cuda
_torch_copy_ = torch.Tensor.copy_
​
​
# Tensor.to 被替换后,Dynamo 只能通过这个模块全局看到保存的 C++ 描述符,
# 无法再把它当作可识别的 tensor 方法。把描述符包进一个可弱引用的 Python
# 可调用对象,allow_in_graph 才能接纳它,下游编译器仍可 lower 原始算子。
# 注意:不能直接 allow_in_graph(_torch_to),因为 method descriptor 无法被弱引用。
@torch.compiler.allow_in_graph
def _torch_to_in_graph(self, *args, **kwargs):
    return _torch_to(self, *args, **kwargs)
​
​
# torch 自己的 Tensor.to 重载签名解析器,返回 (device, dtype, non_blocking, memory_format)
_parse_to = getattr(torch._C._nn, "_parse_to", None)
​
​
def _active_check_mode() -> str | None:
    """当前调用线程所处的检查模式;没有则返回 None。"""
    # 关键顺序:先判 is_compiling(),再读 _allow_depth / _checking 两个 ContextVar。
    # Dynamo 追踪期间读取 ContextVar 会直接报 unsupported ContextVar.get(),
    # 因此编译期必须在触碰任何 ContextVar 之前短路返回 None。
    if not _sync_check_enabled or torch.compiler.is_compiling() or _allow_depth.get():
        return None
    return _checking.get()
​
​
def _checked_to(self, *args, **kwargs):
    # 只检查 H2D:D2H 的 to 会在支持的 torch 版本里分配 pinned 目标,本身是异步的。
    if (mode := _active_check_mode()) is not None and _parse_to is not None:
        try:
            device, _, non_blocking, _ = _parse_to(*args, **kwargs)
        except Exception:
            pass # 参数非法:交给原始调用抛出原有异常
        else:
            if non_blocking and (reason := _cpu_copy_stall_reason(self, device)):
                _report_implicit_sync("to", "H2D", reason, mode)
    # 改走 allow_in_graph 包装后的入口,避免 Dynamo 直接面对 C++ method descriptor
    return _torch_to_in_graph(self, *args, **kwargs)
tests/utils_/test_gpu_sync_debug.py test-coverage

新增 focused 回归用例 `test_copy_checker_is_traceable_while_compiling`,用动态 device/dtype 的 `convert_like` 与 `fullgraph=True` 精确复现 Voxtral 的失败场景。

@create_new_process_for_each_test()
def test_copy_checker_is_traceable_while_compiling(monkeypatch):
    """Dynamo 追踪期间,拷贝包装器必须绕过 ContextVar。"""
    monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error")
    monkeypatch.setattr(gsd, "_sync_check_enabled", True)
    # 显式安装被替换的 Tensor.to / cuda / copy_ 包装器,复现真实运行态
    gsd._install_copy_checkers()
​
    def convert_like(x, like):
        # 关键:device 与 dtype 都来自另一个张量,完全对齐 Voxtral 全图编译时的动态参数
        return x.to(device=like.device, dtype=like.dtype)
​
    # fullgraph=True 让任何 graph break 直接失败;用默认 backend 才能真正覆盖 AOT/Inductor,
    # 而不是被 backend="eager" 掩盖问题。
    compiled = torch.compile(convert_like, fullgraph=True)
    result = compiled(
        torch.ones(4, device="cuda", dtype=torch.float64),
        torch.ones(4, device="cuda", dtype=torch.float32),
    )
    # 断言动态 dtype 转换确实生效,证明原始 C++ Tensor.to 被正确 lower
    assert result.dtype == torch.float32
tests/models/multimodal/generation/test_voxtral_realtime.py test-coverage

给逃逸失败的原始用例加 `@create_new_process_for_each_test(method="spawn")`:避免前一参数的编译 /CUDA graph 显存残留影响后一参数,并绕开 fork 子进程的 CUDA 重初始化失败。

from ....utils import ROCM_ENGINE_KWARGS, create_new_process_for_each_test
​
​
@pytest.mark.skipif(
    not current_platform.is_cuda(), reason="Full cudagraphs require CUDA"
)
@pytest.mark.parametrize(
    "cudagraph_mode",
    [CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_AND_PIECEWISE],
)
# 每个参数跑在全新进程里,原因有两条:
# 1) 前一个参数在进程内编译并捕获 CUDA graph 后残留约 6.5 GiB 显存,
# 会让后一个参数只剩 25.73/32.5 GiB 可用、撞上 0.9 启动守卫并级联失败;
# 2) 父 pytest 进程可能已初始化 CUDA,默认的 fork 子进程会直接报
# "Cannot re-initialize CUDA in forked subprocess",必须显式指定 spawn。
@create_new_process_for_each_test(method="spawn")
def test_voxtral_realtime_cudagraph(
    audio_assets, tokenizer, vllm_runner, monkeypatch, cudagraph_mode
):
    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
​
    vllm_kwargs = {**ENGINE_CONFIG}
    vllm_kwargs["model_name"] = vllm_kwargs.pop("model")
    vllm_kwargs["enforce_eager"] = False
    vllm_kwargs["compilation_config"] = {
        "cudagraph_mode": cudagraph_mode,
        # max_num_seqs 是 4,因此这是唯一值得捕获的 decode size
        "cudagraph_capture_sizes": [4],
    }
​
    with vllm_runner(**vllm_kwargs) as vllm_model:
        inputs, sampling_params = _build_inputs(audio_assets, tokenizer)
        outputs = vllm_model.llm.generate(inputs, sampling_params=sampling_params)
        _assert_expected_text(outputs)

评论区精华

编译期必须在读取 ContextVar 之前短路 正确性

PR body 指出包装器“read a ContextVar before checking whether Dynamo is tracing”,导致 full-graph 编译报 unsupported ContextVar.get()。讨论将 `torch.compiler.is_compiling()` 前移到 `_allow_depth.get()` 之前;#88941 验证该改动通过 H200 Async/Inputs/Utils/Worker,但随即暴露第二个问题——Dynamo 拒绝保存的 C++ `TensorBase.to` 描述符(真实 Voxtral 路径用另一个张量提供 device/dtype)。

结论:先修短路顺序,再补描述符逃逸通道;两步缺一不可,#88969 起 Voxtral FULL_DECODE_ONLY 能走通 Dynamo、Inductor、capture 与推理。 · 已解决

allow_in_graph 不能直接注册 C++ method descriptor 设计

revision 版本尝试 `torch.compiler.allow_in_graph(_torch_to)`,Buildkite #88962 在测试收集前就失败:"importing `gpu_sync_debug.py` failed with `TypeError: cannot create weak reference to 'method_descriptor' object`",因为 allow_in_graph 内部需要弱引用被注册对象。

结论:改为注册可弱引用的 Python wrapper `_torch_to_in_graph`,内部调用原描述符,保持无递归;Dynamo 视其为不透明调用,AOT/Inductor 仍能 lower 原始 op。 · 已解决

Voxtral CUDA graph 用例的进程隔离:显存残留与 fork-CUDA 冲突 测试

#88969 中 FULL_DECODE_ONLY 通过后,FULL_AND_PIECEWISE 因前一个进程内 compiled runner 残留约 6.5 GiB 显存、只剩 25.73/32.5 GiB 而撞上 0.9 启动守卫;#89000 / #89011 改用新进程后又因默认 fork 与父进程已初始化的 CUDA 冲突,报 `Cannot re-initialize CUDA in forked subprocess`(且 #89011 的 head 实际未 stage 该改动)。

结论:显式使用 `@create_new_process_for_each_test(method="spawn")`;#89026 在 f7f21d97caec 上两个参数(FULL_DECODE_ONLY、FULL_AND_PIECEWISE)均通过,lane 结果为 32 passed / 26 skipped / 64 deselected。 · 已解决

AI 辅助改动与验证纪律(精确 SHA / 定向 lane) other

PR body 声明 "prepared with AI assistance. Human review is required";claude[bot] 因来自 fork 关闭自动评审,仅提示可由 maintainer 手动触发。验证过程中出现两类过程失误:#89010 使用手抄的无效完整 SHA 导致 bootstrap 失败(不含任何测试证据);#89011 的 head 未包含预期的 spawn 改动。作者据此在正文中逐条更正引用,强调 literal signed head 门禁。

结论:njhill 发出 `/ci test` 并最终 APPROVED;人工 line review 由维护者在合并前完成,未留下未解决问题。 · 已解决

风险与影响

  1. 修复只覆盖 Tensor.to(潜在遗漏面):_checked_cuda / _checked_copy_ 仍然直接调用模块全局 _torch_cuda / _torch_copy_,它们同样是导入期捕获的 C++ 方法描述符。若某模型在 full-graph 编译中调用 x.cuda() 或 y.copy_(x),理论上可能复现同类 Dynamo 拒绝,而新增测试只覆盖了 Tensor.to 路径。这是基于源码结构的推断,本次 CI 未给出证据,但属于明确未覆盖面。

  2. 编译期同步检查能力被有意放弃:_active_check_mode() 在 is_compiling() 为真时返回 None,意味着一旦启用 torch.compile,gpu_sync_debug 的 implicit sync 检测在编译图中完全失效。这是设计取舍(编译图里本就无法抛 Python 异常),但使用者在“编译 + debug”组合下会失去覆盖率,需要在文档或使用说明中明确。

  3. 条件重排的语义风险低但需复核:重排只改变求值顺序与短路时机,布尔结果不变;_sync_check_enabled 仍是首个求值项,确认它是普通全局而非 ContextVar。

  4. 测试隔离策略带来的 CI 成本:test_voxtral_realtime_cudagraph 两个参数各自 spawn 全新进程并重建 4B 模型,H200 上的显存与时间开销上升;不过该用例仅覆盖 CUDA graph 场景,forward / generator 用例未加,隔离范围控制得较克制。spawn 模式下 fixture(audio_assets、tokenizer、vllm_runner)必须可跨进程重建,这也是后续其他 compile-heavy 测试复用时需要注意的前提。

  5. 验证依赖 CI 且与 cherry-pick 绑定:开发机无 CUDA/PyTorch 运行时,全部结论来自 Buildkite 精确 SHA 门禁;PR 被标记进 v0.30.0 cherry picks,cherry-pick 后的分叉 head 需要重新获取等价证据。

  6. 回归风险面:源码改动仅影响调试包装器的控制流,生产默认路径(未开启 gpu_sync_debug)不受影响,因此主路径回归风险低。

  • 用户侧:开启 gpu_sync_debug 同时使用 full-graph torch.compile 的组合不再直接崩溃;Voxtral realtime 在 enforce_eager=False + FULL_DECODE_ONLY / FULL_AND_PIECEWISE 下的可用性恢复,Voxtral 使用者受益。
  • 系统侧:H200 Multimodal Models (Extended Generation 1) 每日 lane 恢复绿色,消除了一个从 PR #51167 逃逸、只在 daily build #88920 才暴露的长期红灯。
  • 团队侧:确立了两个可复用范式——C++ 方法描述符需经可弱引用 Python wrapper 才能进 allow_in_graph;含 CUDA graph 的多模态用例必须显式 spawn 隔离。同时示范了“精确 SHA + 定向 lane”的验证纪律(含 bootstrapping 失败与未 stage 改动的反面案例)。
  • 影响范围有限:32 行改动集中在调试工具与两个测试文件,不触及调度、KV cache、内核等核心推理链路,属于“低风险、高可读性”的修复 PR。
调试包装器控制流变更 编译期不再做 implicit sync 检查 修复仅覆盖 Tensor.to 描述符 测试改 spawn 提升 CI 开销 验证依赖 CI 精确 SHA 且需随 cherry-pick 重跑

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论