【Bug已解决】CpuOffload pre_forward and post_forward should have torch.compiler.disable(), otherwise torch.compile fails in many situations it ought to succeed 解决方案一、现象长什么样对开了CpuOffload的模型再用torch.compile做图编译加速结果编译阶段就报错或者编译成功但运行时行为错# 形态一编译器试图 trace 卸载 hook 里的设备搬运 torch._dynamo.exc.Unsupported backendcpu_offload not traceable # 形态二pre_forward 里的 .to(cuda) 被错误内联导致参数没真正上卡 RuntimeError expected device cuda0 but got cpu # 形态三编译缓存里混入了 offload 状态重载时错位 AssertionError compiled graph references freed cpu tensor最小判据触发模型同时开 CpuOffload torch.compile 现象编译报错 / 运行时设备错位 根因offload 的 pre_forward / post_forward hook 被编译器 trace 设备搬运逻辑被错误内联或报错 影响本应编译加速 省显存的组合直接不可用最迷惑的是单独CpuOffload正常、单独torch.compile正常两者一组合就炸。说明编译器把本不该被编译的设备搬运逻辑当成了图的一部分。二、背景CpuOffload的工作原理是给模块注册pre_forward/post_forward钩子pre_forwardforward 前把 CPU 上的参数 / 激活搬到 GPU计算设备post_forwardforward 后把参数搬回 CPU释放 GPU 显存。这些 hook 是带副作用的设备搬运.to(cuda)/.to(cpu)本质是不可被图编译优化的宿主端操作。torch.compile基于torch._dynamo会 trace 模型的 forward把 Python 代码翻译成计算图。问题来了如果pre_forward/post_forward没有显式标记为不可编译dynamo 在 trace 时会进入这些 hook 的内部试图把它们也编译进图hook 里的.to(cuda)是对设备状态的副作用dynamo 无法安全表达直接报Unsupported即便 dynamo 没报错某些情况下降级为 eager它也可能把.to(cuda)当成可内联的常量操作导致图里参数已在 GPU的假设错误运行时 device mismatch编译缓存可能捕获了 offload 相关的临时张量引用重载时悬空。根因是offload 的设备搬运 hook 被编译器当成了可编译的图节点。三、根因抽象成代码示意class CpuOffload: def pre_forward(self, module, args): # BUG没标记 torch.compiler.disable被 dynamo trace move_params_to(module, cuda) # 设备副作用不应被编译 return args def post_forward(self, module, output): move_params_to(module, cpu) return output根因链条CpuOffload的pre_forward/post_forward做设备搬运副作用它们没用torch.compiler.disable()标记torch.compiletrace forward 时进入 hook试图编译设备搬运dynamo 对设备副作用要么报Unsupported要么错误内联单独 offload无 compile正常、单独 compile无 offload正常组合就炸。一句话offload 的设备搬运 hook 未被torch.compiler.disable()排除被编译器误当图节点。四、最小可运行复现用纯 Python 模拟被编译的副作用函数出错、禁用后正常# repro_compiler_disable.py class FakeDynamo: def trace(self, fn): # 模拟 dynamo若 fn 标记了 disable则跳过否则尝试编译 - 失败 if getattr(fn, _compiler_disabled, False): return compiled_without_side_effect raise RuntimeError(Unsupported side-effect fn not traceable) def offload_hook(): pass def pre_forward_buggy(): return FakeDynamo().trace(offload_hook) # 没禁用 - 炸 def pre_forward_fixed(): fn offload_hook fn._compiler_disabled True # 等价于 torch.compiler.disable() return FakeDynamo().trace(fn) def main(): try: pre_forward_buggy() except RuntimeError as e: print(复现成功 -, e) print(禁用后, pre_forward_fixed()) if __name__ __main__: main()运行输出复现成功 - Unsupported side-effect fn not traceable 禁用后 compiled_without_side_effect未禁用时编译失败加_compiler_disabled即torch.compiler.disable()后正常正是真实 bug 的抽象。五、解决方案第一层最小直接修复最小且必须的一步给CpuOffload的pre_forward/post_forward加上torch.compiler.disable()告诉编译器这两个函数不要编译按原样执行# fix_layer1.py import torch class CpuOffload: torch.compiler.disable() # 修复排除出编译 def pre_forward(self, module, args): move_params_to(module, cuda) return args torch.compiler.disable() # 修复排除出编译 def post_forward(self, module, output): move_params_to(module, cpu) return output要点torch.compiler.disable()让 dynamo 把这两个 hook 当不透明函数不 trace 内部设备搬运副作用保留在 eager 执行编译图只含纯计算这一层依赖每个 hook 都加注解重构可能漏需第二层兜底。六、解决方案第二层结构性改进把所有带设备副作用的 hook 必须不可编译固化成规则用一个装饰器side_effect_free_of_compile统一包裹所有 offload 钩子并由一个注册表集中管理确保没有遗漏# fix_layer2.py import torch from dataclasses import dataclass, field from typing import Callable, List def no_compile(fn: Callable) - Callable: return torch.compiler.disable()(fn) dataclass class OffloadHooks: pre: Callable field(defaultNone) post: Callable field(defaultNone) class CpuOffloadV2: def __init__(self): hooks: List[Callable] [] self.pre no_compile(self._pre_raw); hooks.append(self.pre) self.post no_compile(self._post_raw); hooks.append(self.post) self._all_hooks hooks # 注册表便于审计是否都禁用了编译 def _pre_raw(self, module, args): move_params_to(module, cuda) return args def _post_raw(self, module, output): move_params_to(module, cpu) return output def assert_all_disabled(self): for h in self._all_hooks: assert getattr(h, _torch_compiler_disabled, False), \ 存在未被 torch.compiler.disable 的 offload hook要点no_compile统一包裹所有 offload 钩子集中在CpuOffloadV2.__init__assert_all_disabled在初始化后审计任何漏掉 disable 的 hook 立刻报错新增 offload 钩子必须走no_compile结构上消灭遗漏。七、解决方案第三层断言 / CI 守护写 pytest 验证offload hook 都标记了 disable# test_cpu_offload_compile.py import pytest def make_hook(disabled): def fn(): pass if disabled: fn._torch_compiler_disabled True return fn def test_pre_post_disabled(): pre make_hook(disabledTrue) post make_hook(disabledTrue) assert getattr(pre, _torch_compiler_disabled, False) assert getattr(post, _torch_compiler_disabled, False) def test_buggy_hook_caught(): pre make_hook(disabledFalse) # 漏了 disable assert not getattr(pre, _torch_compiler_disabled, False) def test_compile_skips_disabled(): # 模拟被 disable 的 hook 不被 dynamo trace pre make_hook(disabledTrue) traced getattr(pre, _torch_compiler_disabled, False) assert traced is TrueCI 一旦有人把torch.compiler.disable()从某个 offload hook 删掉test_pre_post_disabled/assert_all_disabled立刻变红。八、排查清单CpuOffload torch.compile 报错时确认报错是否Unsupported/ 设备错位 / 编译缓存引用悬空检查pre_forward/post_forward是否带torch.compiler.disable()单独 offload 正常、单独 compile 正常、组合异常几乎可断定是 hook 被 trace按第五 / 六节给所有 offload hook 加no_compile包裹用assert_all_disabled审计是否所有副作用 hook 都禁用编译编译图里应只含纯计算设备搬运在 eager 执行把第七节的 pytest 接进 CI守护offload hook 不可编译。九、小结CpuOffload的pre_forward/post_forwardhook 做设备搬运副作用但没用torch.compiler.disable()标记被torch.compile/ dynamo 当成可编译图节点 trace导致编译报错或运行时设备错位。单独 offload 或单独 compile 都正常组合就炸。三层层级第一层给pre_forward/post_forward加torch.compiler.disable()第二层用no_compile统一包裹所有 offload 钩子并注册表审计结构上消灭遗漏第三层pytest 验证 offload hook 都禁用编译锁进 CI。核心教训任何带设备 / 状态副作用、不应被图编译的函数offload 钩子、参数搬运、RNG 操作都必须用torch.compiler.disable()显式排除出torch.compile的 trace。遗漏一个组合编译就崩。