Prhub

#41235 [Bugfix][Compile] Fix gc.collect/empty_cache patch arity in CUDAGraphWrapper

原始 PR 作者 roikoren755 合并时间 2026-04-30 05:59 文件变更 1 提交数 1 评论 0 代码增减 +7 / -2

执行摘要

修复 CUDAGraph 捕获时 gc.collect 参数不匹配

CUDAGraphWrapper 在分段 CUDA 图捕获期间,使用无参 lambda 修补 gc.collect 和 torch.accelerator.empty_cache 以抑制它们。但 gc.collect 有时会被 torch._dynamo.convert_frame._compile 以参数调用(如 gc.collect(generation)),导致 TypeError: () takes 0 positional arguments but 1 was given,使得 worker 进程崩溃。该问题在 GB200 节点上稳定复现。

建议立刻合入。该 PR 修复了一个明确的崩溃问题,修改极小且安全。对于使用嵌套 torch.compile 的模型(如 GB200 节点上的场景)至关重要。

讨论亮点

无 review 讨论。提交仅由 ProExpertProg 批准,无提交评论或讨论。

实现拆解

  1. 修改 gc.collect 的 patch lambda:在文件 vllm/compilation/cuda_graph.py 第 293-294 行,将 lambda: None 改为 lambda *args, **kwargs: None
  2. 修改 torch.accelerator.empty_cache 的 patch lambda:同样在第 296-301 行,将 lambda: None 改为 lambda *args, **kwargs: None
  3. 保持其他逻辑不变:条件判断 self.cudagraph_options.gc_disableExitStack 的使用方式未受影响。
  4. 无额外测试或配置变更:本 PR 仅涉及两处 lambda 参数的修改,无测试文件、配置文件或文档的变更。
文件 模块 状态 重要度
vllm/compilation/cuda_graph.py 编译模块 modified 5.14

关键符号

CUDAGraphWrapper.__call__

关键源码片段

vllm/compilation/cuda_graph.py core-logic

唯一变更文件,核心逻辑修复所在。

# 在 CUDAGraphWrapper.__call__ 方法中,当 gc_disable 为 True 时,
# 使用 patch 上下文管理器抑制 gc.collect 和 torch.accelerator.empty_cache。
# 原 lambda 不接受参数,但 gc.collect 在某些场景下会被传入参数调用,
# 导致 TypeError。修改后接受任意参数,兼容所有调用方式。with ExitStack() as stack:
    if self.cudagraph_options.gc_disable:
        # 使用 *args, **kwargs 使得 patch 能兼容任何调用签名
        stack.enter_context(
            patch("gc.collect", lambda *args, **kwargs: None)
        )
        stack.enter_context(
            patch(
                "torch.accelerator.empty_cache",
                lambda *args, **kwargs: None,
            )
        )

评论区精华

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

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

风险与影响

风险极低。变更仅将两个 lambda 从无参数改为接受任意参数,不会改变抑制 gc.collect 和 empty_cache 的语义。潜在风险在于如果将来有其他代码依赖于这些函数未被调用时的副作用,但当前逻辑就是完全抑制它们,因此不会引入回归。

影响范围小,仅影响启用了 CUDAGraphOptions.gc_disable 的分段 CUDA 图捕获路径。修复了一个在特定场景下(嵌套 @torch.compile 调用 + CUDA 图捕获)的崩溃问题,提高了系统稳定性。

修改极小 无测试覆盖

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论