PyTorch中DDP结合CUDA Graph报错求助：捕获阶段CUDA操作失败

在4张GPU上运行PyTorch代码，尝试结合DDP（分布式数据并行）与CUDA Graph进行训练时，出现CUDA捕获阶段错误，报错信息及原始代码如下：

原始代码 #

setup(rank, gpus)

dataset = RandomDataset(input_shape, 80*batch_size, rank)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
data_iter = iter(dataloader)

model = model(pretrained=True).to(rank)
optimizer = optim.SGD(model.parameters(), lr=0.0001)
criterion = torch.nn.CrossEntropyLoss()

s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())

with torch.cuda.stream(s):
    print("[MAKING DDP Model]")
    model = DDP(model)
    print("[MODEL CREATED]")

    for i in range(11):
        optimizer.zero_grad(set_to_none=True)
        inputs, labels = next(data_iter)
        output = model(inputs)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()

capture_input = torch.empty((batch_size, 3, input_shape, input_shape)).to(rank)
capture_target = torch.argmax(torch.from_numpy(np.eye(1000)[np.random.choice(1000, batch_size)]), axis=1).to(rank)

g = torch.cuda.CUDAGraph()

optimizer.zero_grad(set_to_none=True)
with torch.cuda.graph(g):
    capture_y_pred = model(capture_input)
    capture_loss = criterion(capture_y_pred, capture_target)
    capture_loss.backward()
optimizer.step()

print("RECORDED")

for i in range(20):
    inputs, label = next(data_iter)
    capture_input.copy_(inputs)
    capture_target.copy_(label)
    g.replay()
    optimizer.step()

print("DATASET DONE")

报错信息 #

RuntimeError: CUDA error: operation failed due to a previous error during capture
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

核心问题点 #

1. 自定义CUDA流同步异常：将DDP初始化和预热步骤放在自定义CUDA流中执行，后续Graph捕获时，自定义流内的异步操作可能未完成，触发捕获阶段的CUDA错误。
2. DDP通信与CUDA Graph兼容性问题：DDP的梯度同步（如allreduce）属于异步操作，默认依赖默认CUDA流，直接在Graph中捕获DDP前向/反向传播会导致动态通信操作无法被Graph正确记录。
3. 未同步CUDA流：预热步骤完成后未同步所有GPU的CUDA流，导致Graph捕获时仍有未完成的异步任务。
4. 优化器更新逻辑混乱：原代码中Graph捕获阶段仅执行反向传播，优化器更新放在捕获外；重放Graph后又重复执行更新，导致梯度更新逻辑错误。

修复后的代码 #

setup(rank, gpus)

dataset = RandomDataset(input_shape, 80*batch_size, rank)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
data_iter = iter(dataloader)

# 初始化模型与DDP，必须在默认流中执行
model = model(pretrained=True).to(rank)
# 设置find_unused_parameters=False，避免动态参数导致Graph捕获异常
model = DDP(model, find_unused_parameters=False)
optimizer = optim.SGD(model.parameters(), lr=0.0001)
criterion = torch.nn.CrossEntropyLoss()

# 预热步骤：在默认流中执行，稳定DDP状态
for i in range(11):
    optimizer.zero_grad(set_to_none=True)
    inputs, labels = next(data_iter)
    output = model(inputs)
    loss = criterion(output, labels)
    loss.backward()
    optimizer.step()

# 同步所有GPU的CUDA流，确保预热阶段的异步操作全部完成
torch.cuda.synchronize(rank)

# 使用真实样本初始化捕获张量，避免未定义内存问题
inputs_sample, labels_sample = next(data_iter)
capture_input = inputs_sample.to(rank, non_blocking=True)
capture_target = labels_sample.to(rank, non_blocking=True)

g = torch.cuda.CUDAGraph()

# 将完整训练步纳入Graph捕获范围
with torch.cuda.graph(g):
    optimizer.zero_grad(set_to_none=True)
    capture_y_pred = model(capture_input)
    capture_loss = criterion(capture_y_pred, capture_target)
    capture_loss.backward()
    optimizer.step()

print("RECORDED")

# 重放Graph进行训练
for i in range(20):
    inputs, label = next(data_iter)
    # 非阻塞式复制数据到捕获张量
    capture_input.copy_(inputs, non_blocking=True)
    capture_target.copy_(label, non_blocking=True)
    # 重放Graph，执行完整训练步
    g.replay()

print("DATASET DONE")

关键修复说明 #

• 移除自定义CUDA流：DDP初始化和预热必须在默认CUDA流中执行，避免流间同步异常。
• 强制CUDA同步：预热后执行torch.cuda.synchronize(rank)，确保所有GPU上的异步操作（包括DDP梯度同步）完成后再启动Graph捕获。
• DDP参数配置：添加find_unused_parameters=False，确保所有模型参数都参与DDP通信，消除Graph捕获时的动态参数不确定性。
• 调整Graph捕获范围：将优化器更新步骤纳入Graph，确保每次重放都执行完整的训练流程（清零梯度→前向→反向→参数更新）。
• 合理初始化捕获张量：使用真实样本初始化捕获张量，避免未定义内存引发的错误，同时通过non_blocking=True提升数据复制效率。

内容的提问来源于stack exchange，提问作者davbleca