[Bug已解决] Inductor all_gather 融合生成「rank 专属缓存」被其他 rank 加载崩溃解决方案一、现象长什么样你在分布式训练里用torch.compile编译一个含all_gather全收集算子的模型并开启了FX 图缓存把编译结果存盘复用见 024。结果某个 rank 编译生成的缓存被另一个 rank 加载时直接崩溃。也就是官方描述的Inductor: allgather fusion creates rank-specific cache that crashes when loaded by other ranks根因Inductor 在融合all_gather时生成的编译产物里混入了「当前 rank 专属」的信息比如 rank id、本 rank 的形状 / 设备号。当这个缓存被别的 rank 反序列化 / 复用rank 专属信息对不上于是崩溃或结果错。本文讲清楚 Inductor 缓存、all_gather 融合为何产生 rank 专属缓存、以及如何规避。二、Inductor 缓存与 all_gather 融合2.1 Inductor 缓存回顾torch.compile把图编译成 kernel并可用 FX 图缓存存盘24 节。缓存的**键cache key**由「模型结构、输入签名、dtype」等决定。理想情况下相同结构应在任何 rank 上生成相同缓存。2.2 all_gather 融合Inductor 会把all_gather和它前后的算子融合成一个更大的 kernel减少通信 / 计算开销。融合时它可能需要知道「这是第几个 rank」「聚合后的总大小」等。如果这些信息被硬编码进编译产物而不是运行时传入缓存就变成「rank 专属」。2.3 为什么被别的 rank 加载会崩假设 rank 0 编译时融合 kernel 里写死了「world_size4, rank0, 输出大小4×N」。rank 1 加载这个缓存运行时它的 rank 是 1但缓存里的 kernel 按 rank 0 的参数执行 → 形状 / 偏移算错 → 崩溃非法访问 / 形状断言。三、可运行诊断「缓存是否 rank 专属」下面脚本演示一个危险模式多个 rank 共享同一个缓存目录编译含 all_gather 的模型。正确做法是每 rank 独立缓存或用无缓存编译。import os import datetime import torch import torch.distributed as dist import torch.nn as nn def setup(rank, world): os.environ[MASTER_ADDR] os.environ.get(MASTER_ADDR, localhost) os.environ[MASTER_PORT] os.environ.get(MASTER_PORT, 29503) dist.init_process_group(nccl, rankrank, world_sizeworld, timeoutdatetime.timedelta(seconds60)) torch.cuda.set_device(rank) def main(): rank int(os.environ.get(RANK, 0)) world int(os.environ.get(WORLD_SIZE, 1)) setup(rank, world) # 含 all_gather 的模型示意每个 rank 本地算一部分再 all_gather local torch.randn(4, 8, devicefcuda:{rank}) gathered [torch.empty_like(local) for _ in range(world)] dist.all_gather(gathered, local) # ❌ 危险所有 rank 共用同一个缓存目录可能加载到别的 rank 的 rank 专属缓存 # os.environ[TORCH_COMPILE_CACHE_DIR] /shared/cache # 不要这样共享 compiled torch.compile(lambda t: t * 2) out compiled(torch.cat(gathered, dim0)) print(f[rank {rank}] 编译输出形状, out.shape) dist.destroy_process_group() if __name__ __main__: main()关键教训分布式场景下不要让各 rank 共享同一个 FX 图缓存目录。四、解决方案一每 rank 独立缓存目录最稳每个 rank 用独立的缓存目录互不加载对方的 rank 专属缓存import os rank int(os.environ.get(RANK, 0)) # 每 rank 独立缓存绝不被别的 rank 加载 os.environ[TORCH_COMPILE_CACHE_DIR] f/tmp/torch_compile_cache_rank{rank}这样每个 rank 自己编译、自己用rank 专属信息只在本 rank 内有效不会跨 rank 污染。五、解决方案二分布式场景直接关掉 FX 图缓存如果共享缓存带来的风险大于「省编译时间」的收益分布式训练里干脆关掉缓存import torch._inductor.config as cfg cfg.fx_graph_cache False # 每次重新编译绝不复用可能 rank 专属的缓存代价每次启动都重新编译几十秒到几分钟。但在分布式场景正确性的优先级高于启动速度。六、解决方案三把 all_gather 移出编译图all_gather融合是产生 rank 专属缓存的源头。如果把它移出torch.compile的图用torch.compiler.disable或fullgraphFalse编译图里就没有 rank 专属信息import torch torch.compiler.disable def my_all_gather(local, world): gathered [torch.empty_like(local) for _ in range(world)] torch.distributed.all_gather(gathered, local) return torch.cat(gathered, dim0) def model_body(x): # 纯计算部分编译all_gather 走 eager return x * 2 1 compiled_body torch.compile(model_body)这样编译图是「rank 无关」的纯计算缓存任何 rank 都能安全加载。all_gather 在外面用distAPI 做稳定可靠。七、解决方案四用一致的输入签名确保缓存键相同即使共享缓存如果能让缓存键不含 rank 信息也能安全复用。做法是编译时用的输入签名在各 rank 完全一致形状、dtype、设备号都用「相对」而非「绝对」。但all_gather融合很容易把 rank 信息渗进去所以最可靠还是第四节 / 第五节。八、解决方案五升级 PyTorchallgather fusion creates rank-specific cache是 Inductor 缓存键计算的 Known Issue——理想情况下缓存键应排除 rank 专属信息。新版本可能已修复让 all_gather 融合的缓存键 rank 无关。查看并升级import torch print(PyTorch, torch.__version__)九、如何判断你踩的是同一条你用torch.compile FX 图缓存24 节模型含all_gather或all_reduce等集合通信多个 rank 共享同一缓存目录现象是「某 rank 编译的缓存被另一 rank 加载时崩溃 / 结果错」每 rank 独立缓存目录后恢复。命中即说明踩中该 rank 专属缓存 bug。十、小结Inductor 的all_gather融合会生成含 rank 专属信息的缓存被其他 rank 加载就崩。应对每 rank 独立缓存目录第四节最稳分布式场景直接关 FX 图缓存第五节正确性优先把all_gather移出编译图第六节让编译图 rank 无关确保编译输入签名各 rank 一致第七节升级到修复缓存键计算的 PyTorch第八节。分布式 编译缓存的组合很微妙缓存本是为了「跨进程复用省时间」但集合通信的融合会把 rank 信息渗进产物让缓存变成「rank 专属」。在分布式场景宁可为每 rank 多花一次编译时间也不要赌共享缓存一定对。