【Bug已解决】torch.export.export 对 CUDA：MPS 示例输入报 Unhandled FakeTensor Device Propagation 解决方案

[Bug已解决] torch.export.export 对 CUDA/MPS 示例输入报 Unhandled FakeTensor Device Propagation 解决方案一、报错长什么样你用torch.export.export(model, example_inputs)把模型导出成稳定计算图时如果example_inputs是CUDA 或 MPS 设备上的张量可能遇到RuntimeError: Unhandled FakeTensor Device Propagation而不是一个「清晰的、面向用户的错误」比如「export 暂不支持 MPS 设备」。也就是官方描述的torch.export.export raises RuntimeError: Unhandled FakeTensor Device Propagation for CUDA/MPS example inputs instead of a clear user-facing errortorch.export在导出时内部用 FakeTensor 模式推导整张图的形状 / 设备。当示例输入在 CUDA / MPS 上时FakeTensor 的「设备传播」逻辑遇到某些情况没处理好就抛出这个内部错误而不是友好提示。本文讲清楚 torch.export 的流程、为何 CUDA/MPS 示例输入会触发、以及如何规避。二、torch.export 是什么torch.export.export(model, example_inputs)把模型「捕获」成一个ExportedProgram稳定计算图可用于 AOTInductor 部署见 26 节、移动端、边缘设备。import torch model torch.nn.Linear(8, 4) x torch.randn(2, 8) ep torch.export.export(model, (x,)) print(ep.graph_module) # 稳定的计算图导出过程用 FakeTensor 模式「跑一遍」模型不碰真实数据记录所有算子的形状 / 设备 / 控制流生成图。三、为什么 CUDA/MPS 示例输入会触发 FakeTensor 设备传播错误torch.export的 FakeTensor 推导主要假设示例输入在CPU上CPU FakeTensor 支持最完整。当示例输入在 CUDA / MPS 上时FakeTensor 需要把「真实 CUDA 张量」转成「假的 CUDA 张量」来推导但某些算子在 FakeTensor 模式下对「CUDA / MPS 设备」的传播逻辑没实现好遇到Unhandled于是抛出内部错误而不是「export 暂不支持在 CUDA 上直接导出请先用 CPU 示例输入」这种清晰提示。简单说export 期望你用 CPU 示例输入做推导但你给了 CUDA/MPS 输入它没优雅处理。四、可运行复现与对比下面脚本演示「CPU 示例输入正常、CUDA 示例输入可能报错」无 GPU 时只演示 CPUimport torch def demo_export(device): model torch.nn.Linear(8, 4) x torch.randn(2, 8, devicedevice) try: ep torch.export.export(model, (x,)) print(f[{device}] 导出成功) except Exception as e: print(f[{device}] 导出失败{type(e).__name__}: {str(e)[:120]}) if __name__ __main__: demo_export(cpu) # 通常成功 if torch.cuda.is_available(): demo_export(cuda) # 可能报 Unhandled FakeTensor Device Propagation elif torch.backends.mps.is_available(): demo_export(mps)如果你看到cuda/mps那行报Unhandled FakeTensor Device Propagation就复现了该问题。五、解决方案一用 CPU 示例输入做 export最稳torch.export本质是「形状 / 设备推导」用 CPU 示例输入就能完成。导出后再在实际设备CUDA/MPS上跑import torch model torch.nn.Linear(8, 4) # 用 CPU 示例输入导出FakeTensor 在 CPU 上支持最完整 x_cpu torch.randn(2, 8) ep torch.export.export(model, (x_cpu,)) # 之后在实际设备上运行用 AOTInductor 或直接调用 if torch.cuda.is_available(): model_cuda ep.module().cuda() x_cuda torch.randn(2, 8, devicecuda) out model_cuda(x_cuda) print(CUDA 上运行导出模型输出, out.shape)这是官方推荐的 export 流程导出用 CPU 示例部署用目标设备。六、解决方案二导出后用 .to(device) 移到目标设备ExportedProgram可以先在 CPU 上导出再整体.to()到 CUDA / MPSimport torch model torch.nn.Sequential(torch.nn.Linear(8, 4), torch.nn.ReLU()) ep torch.export.export(model, (torch.randn(2, 8),)) # 移到目标设备 device cuda if torch.cuda.is_available() else cpu ep_cuda ep.to(device) x torch.randn(2, 8, devicedevice) out ep_cuda.module()(x) print(设备, device, 输出, out.shape)ep.to(device)把导出图里的参数 / 缓冲区搬到目标设备避免导出阶段就用 CUDA 示例输入踩 FakeTensor 坑。七、解决方案三用 dynamic_shapes 时也要 CPU 示例如果你的模型用动态形状dynamic_shapes示例输入同样用 CPUimport torch from torch.export import Dim model torch.nn.Linear(8, 4) batch Dim(batch, min1, max32) ep torch.export.export( model, (torch.randn(2, 8),), # CPU 示例 dynamic_shapes{x: {0: batch}}, ) print(带动态形状的 CPU 导出成功)示例输入保持 CPU动态约束用Dim声明FakeTensor 推导最稳。八、解决方案四MPS 用户特别注意MPSApple Silicon上的 export 支持比 CUDA 更滞后见 034 的 MPS dtype 缺口。如果你在 MPS 上 export 报错几乎肯定要退回 CPU 示例输入# MPS 用户export 一律用 CPU 示例运行再 .to(mps) x_cpu torch.randn(2, 8) # CPU ep torch.export.export(model, (x_cpu,)) ep_mps ep.to(mps)九、解决方案五升级 PyTorch错误信息改善Unhandled FakeTensor Device Propagation是「内部错误未转成用户友好提示」的 bug。新版本可能已让它「清晰报错」或直接支持 CUDA/MPS 示例输入。查看并升级import torch print(PyTorch, torch.__version__)十、如何判断你踩的是同一条你调用torch.export.export示例输入在 CUDA 或 MPS 上报错是Unhandled FakeTensor Device Propagation内部错误不清晰改用 CPU 示例输入后导出成功。命中即说明踩中该 export 的设备传播 bug。十一、小结torch.export.export对 CUDA/MPS 示例输入报Unhandled FakeTensor Device Propagation是导出期的 FakeTensor 设备推导不支持这些设备且没转成清晰错误。应对用 CPU 示例输入做 export第五节最稳官方推荐导出后.to(device)移到目标设备运行第六节动态形状也用 CPU 示例 Dim声明第七节MPS 用户必须退回 CPU 示例第八节升级到改善错误 / 支持该设备的 PyTorch第九节。torch.export的 FakeTensor 推导在 CPU 上最成熟。把示例输入留在 CPU、把设备搬运留到导出之后——这个「先推导、后部署」的顺序能避开绝大多数 export 期的设备传播错误。