【Bug已解决】重入式 autograd Function 在 CPU：CUDA 上报错（DISABLED test）解决方案

[Bug已解决] 重入式 autograd Function 在 CPU/CUDA 上报错DISABLED test解决方案一、现象长什么样PyTorch CI 里有一条DISABLED test_reentrant_parent_error_on_cpu_cuda (__main__.TestAutogradDeviceTypeCUDA)它测试的是重入式reentrantautograd.Function在反向传播时遇到「父图错误」时的行为在 CPU 和 CUDA 上。被DISABLED说明当一个 autograd.Function 在反向里「重入」调用了另一个需要梯度的函数且在 CPU/CUDA 下错误处理不一致或崩溃官方暂时禁用了测试。「重入式 autograd」指你在自定义autograd.Function的backward里又调用了另一个会建计算图、需要梯度的操作。这是一种高级用法但极易踩坑。本文讲清楚重入式 autograd 是什么、为什么在 CPU/CUDA 下会出错、以及正确写法。二、autograd.Function 与「重入」自定义autograd.Functionimport torch class MyFn(torch.autograd.Function): staticmethod def forward(ctx, x): ctx.save_for_backward(x) return x * 2 staticmethod def backward(ctx, grad_output): (x,) ctx.saved_tensors return grad_output * 2「重入」是指在backward里又去调用一个会建图的函数class ReentrantFn(torch.autograd.Function): staticmethod def forward(ctx, x): ctx.save_for_backward(x) return x * 2 staticmethod def backward(ctx, grad_output): (x,) ctx.saved_tensors # 重入在 backward 里又建了一个需要梯度的图 y (x * x) # y 需要梯度 g torch.autograd.grad(y, x, grad_outputstorch.ones_like(x))[0] return grad_output * 2 g这种「backward 里再反向」的模式如果不按规定来会让 autograd 引擎的图栈状态混乱导致错误尤其和「父图」的交互。三、为什么在 CPU/CUDA 下错误处理会不一致autograd 引擎在不同设备上的实现细节略有差异CPU 和 CUDA 的 backward 调度、流同步。当重入式 Function 触发「父图错误」比如父图已经在反向、子图又重入导致图栈错乱错误在 CPU 和 CUDA 上表现不同一边可能清晰报「reentrant call not allowed」另一边可能崩溃 / 报含糊的 device 错误或两边报错信息不一致CI 的断言失败 →DISABLED。这正是官方描述的「on cpu_cuda」CPU 和 CUDA 都要覆盖但行为不一致。四、可运行复现重入式 backward下面脚本演示一个危险的重入式 backward纯 CPU 可跑CUDA 同理import torch class ReentrantBad(torch.autograd.Function): staticmethod def forward(ctx, x): ctx.save_for_backward(x) return x 1 staticmethod def backward(ctx, grad_output): (x,) ctx.saved_tensors # 危险backward 里再建需要梯度的图且直接反向 y x * x gy torch.autograd.grad(y.sum(), x)[0] return grad_output gy def demo(): x torch.randn(3, requires_gradTrue) try: out ReentrantBad.apply(x) out.sum().backward() print(重入 backward 成功梯度, x.grad) except Exception as e: print(重入 backward 报错, type(e).__name__, -, str(e)[:160]) if __name__ __main__: demo()如果你看到报错或结果明显不对就复现了重入式 autograd 的坑。五、解决方案一不要在 backward 里重入建图用 ctx 保存所需量绝大多数重入需求其实可以用ctx.save_for_backward规避——把 backward 需要的量在 forward 里存好backward 里只用这些量做纯数值计算不再建图import torch class CleanFn(torch.autograd.Function): staticmethod def forward(ctx, x): # 把 backward 需要的量提前存好 ctx.save_for_backward(x) return x 1 staticmethod def backward(ctx, grad_output): (x,) ctx.saved_tensors # 不重入直接用保存的 x 做数值计算 # 假设我们要 dy/dx 2x对应 yx^2 的梯度 g 2 * x return grad_output g # 使用 x torch.randn(3, requires_gradTrue) out CleanFn.apply(x) out.sum().backward() print(干净 backward 梯度, x.grad)这样 backward 是「纯数值」不建新图不会触发重入错误CPU/CUDA 行为一致。六、解决方案二如果必须重入用 torch.autograd.grad 的正确姿势确实需要在 backward 里对子图求导时确保子图是局部的不影响外层图栈用torch.autograd.grad(..., create_graphFalse)默认得到梯度不把子图接回父图或用torch.no_grad()包住不需要梯度的子计算。class ControlledReentrant(torch.autograd.Function): staticmethod def forward(ctx, x): ctx.save_for_backward(x) return x 1 staticmethod def backward(ctx, grad_output): (x,) ctx.saved_tensors with torch.no_grad(): # 明确不建图 y (x * x).sum() gy torch.autograd.grad(y, x, retain_graphFalse)[0] return grad_output gytorch.no_grad()让子计算不参与 autograd避免重入污染父图。七、解决方案三用 Function 的嵌套而非 backward 内重入如果逻辑复杂把「子图求导」拆成独立的autograd.Function在主函数的 forward 里就完成而不是塞进 backwardclass SubFn(torch.autograd.Function): staticmethod def forward(ctx, x): y x * x ctx.save_for_backward(x) return y staticmethod def backward(ctx, g): (x,) ctx.saved_tensors return g * 2 * x class MainFn(torch.autograd.Function): staticmethod def forward(ctx, x): ctx.save_for_backward(x) z SubFn.apply(x) # 子图在 forward 里完成backward 自动链式 return z 1 staticmethod def backward(ctx, g): (x,) ctx.saved_tensors return g # 子图梯度由 autograd 引擎自动传播无需手动重入把重入变成「嵌套 Function由引擎自动链式反向」最稳。八、解决方案四升级 PyTorchtest_reentrant_parent_error_on_cpu_cuda是 autograd 引擎对重入错误处理的 Known IssueCPU/CUDA 行为应一致。新版本可能已统一错误处理。查看并升级import torch print(PyTorch, torch.__version__)九、如何判断你踩的是同一条你写了自定义autograd.Function且在backward里又建了需要梯度的图现象是「反向时崩溃 / 含糊设备错误 / CPU 与 CUDA 报错不一致」把 backward 改成纯数值ctx 取值后恢复。命中即说明踩中重入式 autograd 的坑。十、小结test_reentrant_parent_error_on_cpu_cuda揭示重入式 autograd.Function 的错误处理在 CPU/CUDA 下不一致。应对backward 里不重入建图用ctx.save_for_backward存好量backward 只做纯数值第五节最稳必须重入时用torch.no_grad()包住子计算不污染父图第六节把子图求导拆成嵌套 Function由引擎自动链式反向第七节升级到统一错误处理的 PyTorch第八节。重入式 autograd 是高级技巧但「backward 里再反向」最容易让图栈状态乱掉尤其在 CPU/CUDA 实现差异下表现不一。最省心的原则是backward 只做数值建图留给 forward。让 autograd 引擎去管链式传播而不是你手动重入。