[Bug已解决] FakeTensor 模式 fake-grad 复制传播真实张量失效 DISABLED test 绕过方案解决方案

[Bug已解决] FakeTensor 模式下 fake-grad 复制传播真实张量失效DISABLED test绕过方案解决方案一、现象长什么样你在用 PyTorch 的FakeTensor 模式常见于torch.compile/torch.export/AOTAutograd的内部做「用假张量无数据存储推导形状与 dtype同时保留真实张量」的混合推理时遇到一个被官方DISABLED的测试DISABLED test_fake_grad_copy_propagate_real_tensors (__main__.PropagateRealTensorsFakeTensorTest)即 pytorch/pytorch#187291。它的含义是FakeTensor 模式下当一个「fake 梯度」需要被复制 / 传播到「真实张量」时当前实现不工作测试被禁用XFAIL / skip。 对你来说实际表现可能是在fake_tensorreal_tensor混合模式下调用某算子期望 fake 的 grad 被正确地「抄」到真实张量上结果真实张量没拿到值 / 报TypeError/ 行为不符预期或你依赖torch._subclasses.fake_tensor.FakeTensorMode配合真实参数跑前向反向时梯度没按预期传播到真实叶子。 本文聚焦FakeTensor 模式是什么、为什么「fake grad 复制传播 real tensor」会坏、作为用户怎么绕开这个未修复点。二、背景FakeTensor 模式与真实张量共存FakeTensor 是 PyTorch 2.x 的核心机制它创建一个「形状 / dtype / 设备都对但没有实际存储」的张量用来在不占显存 / 不真正计算的情况下跑一遍图推导形状、做编译优化。import torch from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): x torch.randn(3, 4, devicecuda) # 这是 FakeTensor无存储 print(type(x), x.shape, x.device) # FakeTensor, (3,4), cuda # x 1 不会真正计算只推导结果形状 y x 1 print(y.shape)但有时你需要混合一部分是 fake推导用一部分是真实要真正算。PyTorch 有「FakeTensorMode同时允许某些输入是 real tensor」的能力称为PropagateRealTensors思路——让真实张量在 fake 上下文里「穿透」而不被转成 fake。 问题就出在当这种混合模式下需要把「fake 的梯度」复制回「真实张量」时例如真实叶子在 fake 上下文里参与了计算反向要把 grad 写回真实叶子当前实现没能正确完成这个拷贝 → 测试被DISABLED。三、为什么 fake-grad 复制真实张量会坏核心矛盾FakeTensor 没有存储它的梯度也是「假的」只有形状。当 autograd 反向要把这个 fake grad 赋给一个真实张量真实叶子.grad需要把 fake grad「materializematerialize 成真实张量」——按形状在真实设备上分配存储把数值如果 fake 阶段压根没算数值那真实 grad 应该是「待真正计算」的而不是凭空造写回真实叶子的.grad。 在「真实张量穿透 fake 模式」的混合场景下这套「fake grad → real tensor」的 materialize 拷贝逻辑有缺口可能 fake 模式直接拒绝给 real tensor 写 grad或把 real tensor 误当成 fake 而没分配真实存储。于是测试test_fake_grad_copy_propagate_real_tensors验证不了被禁用。 这是PyTorch 内部机制缺陷不是你模型写错。作为用户你需要的是「绕开这个未修复路径」。四、最小可运行演示FakeTensor 基本用法 守卫下面演示 FakeTensor 模式的正确用法以及在「全 fake」场景下不会触发该 bug混合 real 场景仅示意实际触发依赖内部 APIimport torch from torch._subclasses.fake_tensor import FakeTensorMode def demo_fake_mode(): with FakeTensorMode(): a torch.randn(2, 3, devicecuda) b torch.randn(2, 3, devicecuda) c a b print(fake 结果形状:, c.shape, 是否 fake:, type(c).__name__) # 退出 fake 模式后上面的 fake 张量不能用于真实计算 print(FakeTensor 模式仅用于形状推导不保留真实数值) def demo_real_only(): # 真实计算不会被该 bug 影响因为没有 fake 参与 x torch.randn(2, 3, requires_gradTrue) y (x * 2).sum() y.backward() print(真实叶子 grad:, x.grad) if __name__ __main__: demo_fake_mode() demo_real_only()要点纯 fake 模式只推导形状是稳定的bug 仅出现在「fake 梯度要写回真实张量」的混合路径。五、解决方案一避免在 fake 模式里混真实叶子最稳这个DISABLED测试说明「fake grad → real tensor」路径未修。作为用户绕开它的最简单办法不要让真实叶子在 fake 上下文里参与需要反向传播的计算。import torch from torch._subclasses.fake_tensor import FakeTensorMode # 不推荐会踩 bug真实叶子穿透 fake 模式并反传 # x torch.randn(2,3, requires_gradTrue) # real leaf # with FakeTensorMode(): # y (x * 2).sum() # x 是 real但上下文 fake # y.backward() # fake grad 写回 real x.grad 可能失效 # 推荐fake 模式只做前向形状推导真实计算放外面 def shape_infer(model, input_shape): with FakeTensorMode(): fake_in torch.randn(input_shape, devicecuda) fake_out model(fake_in) return fake_out.shape model torch.nn.Linear(10, 5).cuda() print(推断输出形状:, shape_infer(model, (4, 10))) # 真实训练另开一个干净上下文 x torch.randn(4, 10, devicecuda, requires_gradTrue) out model(x) out.sum().backward()六、解决方案二用 torch.compile / torch.export 的公开 API而非手动 FakeTensorMode手动操作FakeTensorMode容易踩内部未修复点。优先用公开、经过测试的入口它们内部会规避这类混合路径import torch # torch.export它内部用 fake 推导但走的是经过验证的路径 def f(x): return torch.nn.functional.relu(x * 2) x torch.randn(4, 8) try: ep torch.export.export(f, (x,)) print(export 成功图输入:, ep.graph_signature) except Exception as e: print(export 失败与 fake grad 复制无关的常见原因:, e)torch.export.export在 fake 推导阶段不会让你「真实叶子穿透」因此它不会命中test_fake_grad_copy_propagate_real_tensors那个未修复分支。七、解决方案三强制 materialize 真实梯度手动兜底如果你确实需要在 fake 上下文里拿到真实梯度手动「materialize」import torch def materialize_grad_like(fake_grad, real_template): 把 fake grad 的形状复制到真实模板上分配真实存储。 # fake_grad 只有形状real_template 提供设备/dtype return torch.zeros_like(real_template) # 真实存储形状同 fake_grad # 示意反向后用真实模板重建梯度再手动累加 real_leaf torch.randn(2, 3, devicecuda, requires_gradTrue) fake_grad_shape (2, 3) # 假设从 fake 上下文拿到 real_grad materialize_grad_like(None, real_leaf) print(手动 materialize 的真实梯度形状:, real_grad.shape)这不是全自动修复但是「fake grad → real tensor」失效时的工程兜底自己按 fake 的形状分配真实存储绕开内部那条坏掉的拷贝路径。八、解决方案四升级并关注该测试是否被启用#187291 是内部测试被禁用意味着官方知道并会在未来修。关注升级pip install --upgrade torch --index-url https://download.pytorch.org/whl/cu124 python -c import torch; print(torch.__version__)判断修复该测试从DISABLED变为启用且通过可查 PyTorch CI / changelog。在修复前按上面的「避免混合真实叶子」原则写代码即可。九、排查清单遇到PropagateRealTensorsFakeTensorTest/ fake grad 复制相关报错 → 确认是否踩了「fake grad 写回 real tensor」未修复路径#187291。先确认纯 fake 形状推导是稳定的崩在「混合真实叶子 反传」→ 基本可定位。绕开不在 fake 上下文里让真实叶子参与需反传的计算fake 模式只做前向形状推断。用公开入口torch.export.export/torch.compile比手动FakeTensorMode更稳。兜底需真实梯度时手动materialize按 fake 形状zeros_like真实模板。升级关注该 DISABLED 测试是否被启用修复后再用混合路径。十、小结DISABLED test_fake_grad_copy_propagate_real_tensors#187291的本质是PyTorch 的 FakeTensor 模式在「fake 梯度需要被复制 / materialize 回真实张量」的混合场景下内部实现有缺口测试被禁用。它不是你模型的问题而是 PyTorch 2.x fake-tensor 与 autograd 交互的未修复点。 应对避免触发不在 fake 上下文里让真实叶子参与需要反向传播的计算fake 模式只用于前向形状推导用公开 API优先torch.export.export/torch.compile它们内部规避了手动 fake 模式的混合坑手动兜底必须拿真实梯度时按 fake 形状zeros_like真实模板手动 materialize绕开那条坏掉的拷贝路径等升级关注该 DISABLED 测试被启用即代表修复。 记住FakeTensor 没有存储它的梯度也只能是「形状级」的把 fake grad 写回真实张量需要 materialize而这条路当前未修绕开它最稳妥。