[Bug已解决] 无法对 DDP 模型启动 CUDA Graph 捕获Unable to launch CUDA Graph with DDP model解决方案一、现象长什么样你用DistributedDataParallelDDP做多卡训练又想用CUDA Graph加速把一串内核录制成图重放时省去 CPU 启动开销。常见的两种尝试# 方式 A用 torch.compile(..., modereduce-overhead) 触发 CUDA Graph model torch.compile(model, modereduce-overhead) # 方式 B手动 cuda graph 捕获 g torch.cuda.CUDAGraph() with torch.cuda.graph(g): out model(input)结果报错Unable to launch CUDA Graph with DDP model ... RuntimeError: CUDA graph capture failed with DDP ... CUDA error: operation not permitted on captured graph本 issuepytorch/pytorch#89546说的就是对 DDP 包装的模型做 CUDA Graph 捕获常常失败——因为 DDP 的前向/反向里嵌了集合通信all_reduce 梯度而 CUDA Graph 捕获对「图里有跨流 / 通信 / 动态控制流」非常敏感。 本文聚焦CUDA Graph 为什么和 DDP 冲突、两种失败模式、怎么在 DDP 下安全地用 CUDA Graph 或退而求其次。二、背景CUDA Graph 与 DDP 的冲突点CUDA Graph把一系列 GPU 内核及它们的依赖录制成静态图之后「重放replay」整图避免每次都从 CPU 逐个 launch 内核的开销。要求图内操作确定、无动态分支、可在同一流上重放不能有「图捕获期间不被允许」的操作如某些同步、跨设备、动态内存分配。DDP 的问题DDP 在 forward 的最后或 backward 中会插入all_reduce 梯度的集合通信这涉及多 GPU 间的 NCCL 通信CUDA Graph 捕获的是「当前流的内核序列」而 all_reduce 可能跨流 / 涉及 NCCL 内部事件捕获时不被允许DDP 的「梯度桶bucket」all-reduce 时机DDP 在 backward 中按 bucket 异步 all_reduce时机与图捕获的「静态记录」冲突参数更新optimizer.step涉及 CPU 侧的参数拷贝 / 优化器状态不在 GPU 图内但 DDP 的梯度同步与图边界交错。 于是对 DDP 模型直接 capture CUDA Graph会因图内混入了不被允许的集合通信 / 动态通信而失败。三、为什么 CUDA Graph 与 DDP 冲突两种典型失败模式模式 1捕获期遇到 all_reduceDDP 的 forward 末尾或 backward 中触发 all_reduceCUDA Graph 捕获到「需要 NCCL 通信」的内核而 NCCL 在 captured graph 上下文里不被允许或需特殊 setup直接报错Unable to launch CUDA Graph。模式 2图重放时参数已变 / 通信状态不符即使勉强捕获成功DDP 的 all_reduce 依赖「所有 rank 同步到达」而图重放是各 rank 独立 replay若某个 rank 重放快/慢集合通信的同步语义被破坏 → 死锁或operation not permitted。 另外torch.compile(modereduce-overhead)在 DDP 模型上也会尝试用 CUDA Graph 包装于是同样触发该限制——这也是很多人「开了 reduce-overhead 的 DDP 训练崩」的原因。四、最小可运行演示守卫 概念下面演示「CUDA Graph 捕获普通模型」与「DDP 下不应直接捕获」的区别多卡需分布式环境import torch import torch.nn as nn def demo_cuda_graph_single_gpu(): if not torch.cuda.is_available(): print([skip] 无 GPU仅说明 CUDA Graph 用法) return model nn.Linear(16, 16).cuda() x torch.randn(4, 16, devicecuda) g torch.cuda.CUDAGraph() with torch.cuda.graph(g): out model(x) # 普通模型可捕获 g.replay() print(普通模型 CUDA Graph 重放形状:, out.shape) def demo_ddp_warning(): # DDP 模型不要直接 capture 整个 forward会因 all_reduce 失败 print(DDP 模型应只 capture『不含 all_reduce 的内部计算』或用 reduce-overhead 谨慎开启) print(不要with torch.cuda.graph(g): ddp_model(x) # 会因 all_reduce 失败) if __name__ __main__: demo_cuda_graph_single_gpu() demo_ddp_warning()要点普通模型能捕获DDP 模型「整个 forward 含 all_reduce」不能整体捕获。五、解决方案一只 capture DDP 内部的纯计算部分若一定要用 CUDA Graph只捕获「不含集合通信」的计算段如单个 Transformer 块的前向而不是整个ddp_model(x)import torch import torch.nn as nn # 把 DDP 内部模块拆出来单独 capture示例结构 class Block(nn.Module): def forward(self, x): return torch.relu(self.fc(x)) if hasattr(self, fc) else x def capture_inner(block, x): if not torch.cuda.is_available(): return None g torch.cuda.CUDAGraph() with torch.cuda.graph(g): out block(x) # 仅纯计算无 all_reduce return g, out # 使用先正常跑 DDP 的 all_reduce再对内部块重放需小心梯度/通信边界但注意这非常 tricky——DDP 的梯度同步和 backward 的 all_reduce 仍在外面capture 内部块要保证「重放时输入来自真实流、且不与 DDP 钩子冲突」。一般不推荐手搓容易引入更难查的 bug。六、解决方案二用 torch.compile 但关掉 CUDA Graphreduce-overhead 改 default最实用的绕过DDP 模型不要用modereduce-overhead它会强开 CUDA Graph改用不依赖 graph 的模式import torch # 不要torch.compile(model, modereduce-overhead) # DDP 下易触发本 bug # 改用 compiled torch.compile(model, modedefault) # 无 CUDA Graph避免冲突 # 或 compiled torch.compile(model, dynamicTrue)default模式做算子融合但不录 CUDA GraphDDP 的 all_reduce 正常执行训练稳定。七、解决方案三用 torch.cuda.graph 配合 DDP 的「无通信段」 显式同步如果你有「DDP 后处理」阶段如只在 rank0 做推理、或单卡 eval可在不参与 DDP 通信的部分用 CUDA Graphimport torch import torch.nn as nn # eval 阶段model.eval() torch.no_gradDDP 不插入 all_reduce torch.no_grad() def eval_with_graph(model, x): if not torch.cuda.is_available(): return model(x) g torch.cuda.CUDAGraph() with torch.cuda.graph(g): out model(x) # eval 无 all_reduce可捕获 g.replay() return out # 训练阶段仍用普通 DDP不捕获CUDA Graph 只用于 eval 推理原理DDP 的 all_reduce 只在model.train()的 forward/backward 插入model.eval()no_grad下没有通信CUDA Graph 可安全捕获。训练用普通 DDPeval 用 graph 加速。八、解决方案四升级并关注 FSDP/DDP CUDA Graph 协同#89546 是长期已知限制PyTorch 在持续改进如torch.compile对 DDP 的 graph 处理、FSDP 的 graph 支持。升级pip install --upgrade torch --index-url https://download.pytorch.org/whl/cu124判断改善同模型torch.compile(ddp_model, modereduce-overhead)不再Unable to launch CUDA Graph。修复前用「default 模式 eval 阶段用 graph」组合。九、排查清单报Unable to launch CUDA Graph with DDP model→ 确认是 DDP graph 冲突#89546。是否用了modereduce-overhead或手动cuda.graph(ddp_model.forward)是 → 改为modedefault或只 capture 无通信段。区分DDP 的 all_reduce 在 train 的 forward/backward 插入CUDA Graph 不能整体包裹。安全用法CUDA Graph 只用于model.eval()no_grad的无通信推理段训练用普通 DDP。手搓内部块 capture 风险高非必要不做。升级关注 DDP/FSDP CUDA Graph 协同的改进。十、小结Unable to launch CUDA Graph with DDP model#89546的本质是DDP 在 forward/backward 中插入了 all_reduce 等集合通信而 CUDA Graph 捕获要求图内是「静态、无跨流通信、可重放」的内核序列把含 all_reduce 的 DDP 前向整体录成 CUDA Graph会因通信操作不被允许而失败。torch.compile(modereduce-overhead)在 DDP 上触发同样的限制。 应对训练别用 graphDDP 训练改用torch.compile(modedefault)不录 CUDA Graphall_reduce 正常执行eval 用 graphmodel.eval()no_grad下无 all_reduceCUDA Graph 可安全捕获用于推理加速勿整体捕获 DDP手搓内部块 capture 风险高非必要不做等升级官方持续改进 DDP/FSDP graph 协同。 记住CUDA Graph 怕「通信」和「动态」DDP 的精髓恰恰是 all_reduce 通信二者天然冲突。训练让 DDP 正常通信只有无通信的 eval 段才适合 CUDA Graph。