[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_reduce DDP 的 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。