[Bug已解决] torch.compile FX 图缓存：bfloat16 + 动态形状预热加载失效（DISABLED test）解决方案

[Bug已解决] torch.compile FX 图缓存bfloat16 动态形状预热加载失效DISABLED test解决方案一、现象长什么样PyTorch CI 里有一条DISABLED test_cache_hot_load_caching_precompile_dynamic_True_device_cuda_bfloat16 (__main__.TestFxGraphCache)它测试的是 **FX 图缓存FX Graph Cache**功能在「CUDA bfloat16 动态形状dynamicTrue」组合下能否正确「预热加载hot load已缓存的预编译图」。被DISABLED说明这个组合下图缓存的加载 / 复用有 bug官方暂时禁用了测试。FX 图缓存是torch.compile的加速神器——把编译好的计算图缓存到磁盘下次启动直接加载省去重新编译的几十秒甚至几分钟。但如果「缓存加载」在 bfloat16 动态形状下失效你就白白等了编译、还拿了错误结果。本文讲清楚图缓存是什么、为什么 bfloat16 dynamic 容易出问题以及如何安全地用它。二、torch.compile 的 FX 图缓存是什么torch.compile(model)默认会把模型 lowering 成 FX 图再交给 Inductor 编译成 kernel。编译可能很慢首次几十秒。FX 图缓存把这个编译结果序列化的 FX 图 / 编译产物存到磁盘下次用相同模型 相同输入签名时直接「热加载」跳过编译。开启方式环境变量具体键名以版本文档为准export TORCH_COMPILE_CACHE_DIR/tmp/torch_compile_cache # 或 export TORCHINDUCTOR_FX_GRAPH_CACHE1或在代码里import torch._inductor.config as cfg cfg.fx_graph_cache True这样第二次运行同一程序编译时间从「几十秒」降到「不到一秒」。三、为什么 bfloat16 动态形状的组合会踩坑3.1 动态形状dynamicTruetorch.compile(f, dynamicTrue)让编译出的图能接受变化形状的输入而不必为每个形状重新编译。动态形状的图里含有「符号维度symint」序列化 / 反序列化比静态形状复杂。3.2 bfloat16bfloat16 是「脑浮点」格式指数位和 float32 一样、尾数位少专为深度学习设计。但它的 kernel 特化、dtype 标记在缓存的图里需要正确记录。3.3 两者叠加 → 缓存键 / 加载出错当图缓存同时涉及「动态符号维度」和「bfloat16 dtype」时缓存系统可能缓存键cache key算错导致「看起来命中其实没命中」白等反序列化时 dtype / 符号维度恢复错误加载出错误的图甚至崩溃预热加载hot load路径在 bfloat16 下走到未支持的分支。于是 CI 用「加载后结果一致性」断言守护发现不对就DISABLED。四、可运行验证图缓存是否真的命中下面脚本演示如何开图缓存并粗略判断「第二次运行是否明显更快」无 GPU 时演示逻辑import time import torch def demo_cache(): if not torch.cuda.is_available(): print(无 GPU仅说明用法。) return model torch.nn.Sequential( torch.nn.Linear(64, 64), torch.nn.ReLU(), ).cuda().to(torch.bfloat16) x torch.randn(8, 64, devicecuda, dtypetorch.bfloat16) compiled torch.compile(model, dynamicTrue) for run in range(2): t0 time.time() with torch.no_grad(): out compiled(x) torch.cuda.synchronize() dt time.time() - t0 print(f第 {run1} 次运行耗时 {dt:.3f}s f若第2次明显更快说明缓存命中) print(输出形状, out.shape) if __name__ __main__: demo_cache()如果你发现第 2 次仍然很慢、或结果不一致就印证了图缓存在该组合下有问题。五、解决方案一动态形状场景用 static关掉 dynamic绕过如果动态形状是触发 bug 的面之一而你实际输入形状比较固定可以不用 dynamic改用静态编译。静态图的缓存序列化更简单、更稳# 静态编译输入形状固定时更稳缓存更可靠 compiled torch.compile(model, dynamicFalse)用dynamicFalse图缓存只在「完全匹配的形状」下命中但命中时正确率高避开 dynamic 的缓存 bug。六、解决方案二bfloat16 场景先 fp32 编译再 cast如果 bfloat16 是触发 bug 的面可以先用 float32 编译缓存最稳推理/训练时再转 bfloat16model model.float().cuda() compiled torch.compile(model) # float32 编译缓存可靠 # 使用时再转 bfloat16 路径AMP x x.bfloat16() with torch.autocast(cuda, dtypetorch.bfloat16): out compiled(x)这样「编译 缓存」走最成熟的 float32 路径bfloat16 只出现在运行时由 AMP 管理避开「bf16 缓存」的 bug。七、解决方案三不要跨版本 / 跨机器共享缓存图缓存是强版本相关的PyTorch 版本、CUDA 版本、甚至源码哈希变了缓存就可能失效或加载错误。务必缓存目录按 PyTorch 版本隔离不要把一个机器编译的缓存目录直接拷到另一台机器用尤其架构不同升级 PyTorch 后清空旧缓存。# 升级后清缓存 rm -rf /tmp/torch_compile_cache export TORCH_COMPILE_CACHE_DIR/tmp/torch_compile_cache_$(python -c import torch; print(torch.__version__))八、解决方案四用 checksum 验证缓存一致性如果你依赖缓存加速且担心加载错误可以在加载后做一次轻量「正确性校验」import torch def verify_compiled(compiled_fn, x, reference_eager_fn): with torch.no_grad(): y_c compiled_fn(x) y_e reference_eager_fn(x) ok torch.allclose(y_c.float(), y_e.float(), atol1e-2, rtol1e-2) print(缓存加载后与 eager 一致, ok) if not ok: print(⚠️ 缓存可能加载错误建议清空缓存重编)把这段加进启动逻辑能在「缓存加载出错」时立刻告警而不是默默用错误结果训练。九、解决方案五升级 PyTorchTestFxGraphCache的 disabled 测试是持续在修的。新版本可能已重新启用 bfloat16 dynamic 的缓存。查看并升级import torch print(PyTorch, torch.__version__)十、如何判断你踩的是同一条你开了 FX 图缓存fx_graph_cacheTrue或对应环境变量数据类型是 bfloat16、且用了dynamicTrue现象是「第二次启动仍很慢」或「加载后结果不对 / 崩溃」清空缓存、关 dynamic、或转 fp32 后恢复正常。命中即说明踩中该 disabled 测试覆盖的问题。十一、小结DISABLED test_cache_hot_load_caching_precompile_dynamic_True_device_cuda_bfloat16揭示FX 图缓存在 bfloat16 动态形状下加载不可靠。应对形状固定就关 dynamicdynamicFalse静态图缓存更稳bfloat16 场景先 fp32 编译缓存再 AMP 转 bf16运行缓存按版本隔离、升级后清空别跨机共享加载后做正确性校验异常立刻告警升级到修复该缓存 bug 的 PyTorch 版本。图缓存是「省编译时间」的利器但「缓存键算错 / 加载出错」会让你既浪费了等待、又拿到错误结果。在它还不稳的组合bf16 dynamic上宁可多等一次编译也别赌缓存一定对。