[Bug已解决] torch.compile(inductor) 静默丢弃前向模式 AD 切向量（返回 None）解决方案

[Bug已解决] torch.compile(inductor) 静默丢弃前向模式 AD 切向量返回 None解决方案一、现象长什么样你用 PyTorch 的前向模式自动微分forward-mode AD——即通过dual张量带切向量 / tangent算前向导数——并把模型用torch.compile(backendinductor)编译后发现编译后的函数对 dual 输入返回了None切向量而不是正确的导数。也就是官方描述的torch.compile with inductor backend silently drops forward-mode AD tangents (returns None)「静默丢弃」是最危险的它不报错只是悄悄给你错误的None / 零导数让你的梯度计算 silently 出错。前向模式 AD 常用于「输入维度远多于输出维度」的雅可比计算如物理仿真、NeRF、某些优化器。本文讲清楚前向模式 AD 是什么、为什么 inductor 会丢切向量、以及如何规避。二、前向模式 AD 与 dual 张量PyTorch 的自动微分有两套反向模式reverse-modeloss.backward()算「输出对参数的梯度」适合参多输出少深度学习主流前向模式forward-mode沿输入方向同时算函数值和「切向量tangent即方向导数」适合输入多输出少的雅可比。前向模式用dual张量承载「原值 切向量」import torch # 开启前向 AD torch.manual_seed(0) x torch.randn(3, requires_gradTrue) # 把 x 包成 dual原值 x切向量是 x 自身方向 tangent torch.randn(3) dual_x torch.zeros_like(x).as_subclass(torch.Tensor) # 占位 # 更标准用法用 torch.autograd.forward_ad from torch.autograd.forward_ad import dual_level, make_dual, unpack_dual def f(x): return x * x * x with dual_level(): d make_dual(x, tangent) # dual (原值 x, 切向量 tangent) out_dual f(d) val, tan unpack_dual(out_dual) print(函数值, val) print(切向量方向导数, tan) # d/dx(x^3) 在 tangent 方向的值数学上f(x)x³其方向导数 3x² · tangent。前向 AD 应该给出正确切向量。但 inductor 编译后可能把它算成None。三、为什么 inductor 会「静默丢弃」切向量torch.compile(inductor)把模型重写成优化 kernel。前向模式 AD 在 PyTorch 内部通过dual张量的元数据和一系列forward_ad规则实现。问题Inductor 生成的 Triton / CUDA kernel可能不识别 / 不支持 dual 张量把它当成普通张量于是切向量信息丢失编译期 FakeTensor 推导对forward_ad规则覆盖不全导致切向量路径被优化掉由于是「返回 None」而非报错用户很难察觉——直到发现导数全错。这是比「报异常」更隐蔽的 bug因为它静默出错。四、可运行对比 eager 与 compile 的前向 AD下面脚本对比「eager 前向 AD」和「compile 后前向 AD」是否一致无 GPU 时纯 CPU 演示import torch from torch.autograd.forward_ad import dual_level, make_dual, unpack_dual def f(x): return x * x * x def forward_ad_eager(x, tangent): with dual_level(): d make_dual(x, tangent) out f(d) return unpack_dual(out) def forward_ad_compiled(x, tangent): compiled_f torch.compile(f) with dual_level(): d make_dual(x, tangent) out compiled_f(d) return unpack_dual(out) if __name__ __main__: x torch.randn(3, requires_gradTrue) t torch.randn(3) v_e, tan_e forward_ad_eager(x, t) print(eager 切向量, tan_e) try: v_c, tan_c forward_ad_compiled(x, t) print(compile 切向量, tan_c) if tan_c is None: print(⚠️ compile 后切向量为 Noneinducitor 丢弃了前向 AD 切向量) else: print(一致, torch.allclose(tan_e, tan_c)) except Exception as e: print(compile 前向 AD 异常, e)如果你看到compile 切向量 None就复现了该 bug。五、解决方案一前向 AD 路径不要用 torch.compile退回 eager最稳需要前向模式 AD 的代码段不进torch.compile退回 eagerimport torch from torch.autograd.forward_ad import dual_level, make_dual, unpack_dual # 前向 AD 部分保持 eager不编译 def forward_ad_only(x, tangent): with dual_level(): d make_dual(x, tangent) out f(d) # f 不编译 return unpack_dual(out) # 只有纯数值前向不含 dual的部分才编译 compiled_f torch.compile(f)或用一个「编译数值部分、AD 在外层」的拆分torch.compiler.disable def f_with_forward_ad(x, tangent): with dual_level(): d make_dual(x, tangent) out f(d) return unpack_dual(out)六、解决方案二用 fullgraphFalse 让 AD 段走 eager如果只有部分路径用前向 AD可以让编译图在 dual 输入处断裂compiled torch.compile(f, fullgraphFalse) # dual 输入可能触发断裂走 eager切向量不丢代价是图被切断性能略降但前向 AD 正确性有保证。七、解决方案三改用反向模式 AD如果适用如果你的场景其实可以用反向模式参多输出少是反向的主场但输入多输出少才是前向的主场且只是想验证可以先用backward算雅可比列import torch # 反向模式近似前向 AD适合验证 x torch.randn(3, requires_gradTrue) y (x * x * x).sum() # 标量输出 y.backward() print(梯度反向, x.grad) # 这与前向 AD 在合适设定下一致但注意反向模式算的是「输出对输入的梯度」的整体和「前向 AD 的方向导数」语义不同。仅当你的雅可比需求能用反向表达时才互换。八、解决方案四升级 PyTorchtorch.compile(inductor) silently drops forward-mode AD tangents是 inductor 对forward_ad支持不完整的 Known Issue。新版本会逐步补全 dual 张量的编译支持。查看并升级import torch print(PyTorch, torch.__version__)九、如何判断你踩的是同一条你用了前向模式 ADdual_level/make_dual/unpack_dual模型被torch.compile(backendinductor)编译编译后unpack_dual的切向量是None或全零而 eager 下正确退回 eager 后恢复。命中即说明踩中该 inductor 静默丢弃前向 AD 切向量的 bug。十、小结torch.compile(inductor)静默丢弃前向模式 AD 切向量返回 None是inducitor 对 dual 张量编译支持不全且报错不声不响。应对前向 AD 路径退回 eager不进 torch.compile第五节最稳fullgraphFalse让 AD 段走 eager第六节若适用改用反向模式 AD 验证第七节升级到补全forward_ad编译支持的 PyTorch第八节。前向模式 AD 是「输入多输出少」场景的雅可比利器但它依赖 PyTorch 内部的 dual 元数据传播。inducitor 在激进优化时可能把这个元数据优化掉——而且不报错。所以凡是带dual的代码编译前务必和 eager 结果比对一次切向量别让None悄悄污染你的导数。