[Bug已解决] 无法对 DDP 模型启动 CUDA Graph 捕获报错解决方案

[Bug已解决] 无法对 DDP 模型启动 CUDA Graph 捕获（Unable to launch CUDA Graph with DDP model）解决方案

一、现象长什么样

你用DistributedDataParallel（DDP）做多卡训练，又想用CUDA Graph加速（把一串内核录制成图，重放时省去 CPU 启动开销）。常见的两种尝试：

# 方式 A：用 torch.compile(..., mode="reduce-overhead") 触发 CUDA Graph model = torch.compile(model, mode="reduce-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

本 issue（pytorch/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 梯度的集合通信，这涉及：

1. 多 GPU 间的 NCCL 通信：CUDA Graph 捕获的是「当前流的内核序列」，而 all_reduce 可能跨流 / 涉及 NCCL 内部事件，捕获时不被允许；
2. DDP 的「梯度桶（bucket）」all-reduce 时机：DDP 在 backward 中按 bucket 异步 all_reduce，时机与图捕获的「静态记录」冲突；
3. 参数更新（optimizer.step）：涉及 CPU 侧的参数拷贝 / 优化器状态，不在 GPU 图内，但 DDP 的梯度同步与图边界交错。 于是：对 DDP 模型直接 capture CUDA Graph，会因图内混入了不被允许的集合通信 / 动态通信而失败。

三、为什么 CUDA Graph 与 DDP 冲突

两种典型失败模式：模式 1：捕获期遇到 all_reduceDDP 的 forward 末尾或 backward 中触发 all_reduce，CUDA Graph 捕获到「需要 NCCL 通信」的内核，而 NCCL 在 captured graph 上下文里不被允许（或需特殊 setup），直接报错Unable to launch CUDA Graph。模式 2：图重放时参数已变 / 通信状态不符即使勉强捕获成功，DDP 的 all_reduce 依赖「所有 rank 同步到达」，而图重放是各 rank 独立 replay，若某个 rank 重放快/慢，集合通信的同步语义被破坏 → 死锁或operation not permitted。 另外，torch.compile(mode="reduce-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, device="cuda") 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 Graph（reduce-overhead 改 default）

最实用的绕过：DDP 模型不要用mode="reduce-overhead"（它会强开 CUDA Graph），改用不依赖 graph 的模式：

import torch # 不要：torch.compile(model, mode="reduce-overhead") # DDP 下易触发本 bug # 改用： compiled = torch.compile(model, mode="default") # 无 CUDA Graph，避免冲突 # 或 compiled = torch.compile(model, dynamic=True)

default模式做算子融合但不录 CUDA Graph，DDP 的 all_reduce 正常执行，训练稳定。

七、解决方案三：用 torch.cuda.graph 配合 DDP 的「无通信段」+ 显式同步

如果你有「DDP 后处理」阶段（如只在 rank0 做推理、或单卡 eval），可在不参与 DDP 通信的部分用 CUDA Graph：

import torch import torch.nn as nn # eval 阶段（model.eval() + torch.no_grad），DDP 不插入 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 可安全捕获。训练用普通 DDP，eval 用 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, mode="reduce-overhead")不再Unable to launch CUDA Graph。修复前，用「default 模式 + eval 阶段用 graph」组合。

九、排查清单

1. 报Unable to launch CUDA Graph with DDP model→ 确认是 DDP + graph 冲突（#89546）。
2. 是否用了mode="reduce-overhead"或手动cuda.graph(ddp_model.forward)？是 → 改为mode="default"或只 capture 无通信段。
3. 区分：DDP 的 all_reduce 在 train 的 forward/backward 插入，CUDA Graph 不能整体包裹。
4. 安全用法：CUDA Graph 只用于model.eval()+no_grad的无通信推理段；训练用普通 DDP。
5. 手搓内部块 capture 风险高，非必要不做。
6. 升级：关注 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(mode="reduce-overhead")在 DDP 上触发同样的限制。 应对：

• 训练别用 graph：DDP 训练改用torch.compile(mode="default")，不录 CUDA Graph，all_reduce 正常执行；
• eval 用 graph：model.eval()+no_grad下无 all_reduce，CUDA Graph 可安全捕获用于推理加速；
• 勿整体捕获 DDP：手搓内部块 capture 风险高，非必要不做；
• 等升级：官方持续改进 DDP/FSDP + graph 协同。 记住：CUDA Graph 怕「通信」和「动态」；DDP 的精髓恰恰是 all_reduce 通信，二者天然冲突。训练让 DDP 正常通信，只有无通信的 eval 段才适合 CUDA Graph。