[Bug已解决] DDP 模式下 CUDA error: an illegal memory access was encountered 解决方案

一、报错长什么样

你在用DistributedDataParallel（DDP）做多卡训练时，训练可能跑着跑着崩：

CUDA error: an illegal memory access was encountered

而且特别标注发生在DDP mode下。也就是官方描述的：

DDP mode: CUDA error: an illegal memory access was encountered

illegal memory access（非法内存访问，错误码 700）意味着某个 CUDA kernel 访问了它不该访问的地址（越界、空指针、已释放显存）。在 DDP 下出现，说明问题往往和分布式特有的逻辑有关——比如梯度桶（gradient bucket）buffer、跨 rank 的 all_reduce 通信、或 DDP 包装后的参数布局。本文讲清楚 DDP 下非法访问的常见来源，以及如何定位修复。

二、DDP 为什么比单机更容易踩非法访问

DDP 在单机多卡 / 多机上的工作方式：

1. 每个 rank 有一份模型副本，参数分到一个个梯度桶（bucket）；
2. 反向时，每个 bucket 的梯度攒满，就通过all_reduce在各 rank 间同步；
3. all_reduce 在后台 CUDA stream 上跑，和前向 / 反向的 stream 异步。

这种「多 stream + 异步通信 + 特定 buffer 布局」带来了单机没有的非法访问风险：

• 梯度 buffer 被提前释放 / 复用，而通信 stream 还在读 → 非法访问；
• 某个 rank 的模型结构和别的 rank 不一致，all_reduce 形状对不上，通信踩内存；
• 自定义 autograd 的 CUDA 扩展在 DDP 的 async 通信下越界；
• find_unused_parameters没开，却有不参与损失的参数，导致 buffer 错位。

三、可运行：健壮的 DDP 骨架（避免非法访问）

下面脚本是一个真实可运行的 DDP 训练骨架，内置了能规避多数非法访问的写法：

import os import datetime import torch import torch.distributed as dist import torch.nn as nn import torch.optim as optim from torch.nn.parallel import DistributedDataParallel as DDP def setup(rank, world): os.environ["MASTER_ADDR"] = os.environ.get("MASTER_ADDR", "localhost") os.environ["MASTER_PORT"] = os.environ.get("MASTER_PORT", "29500") dist.init_process_group("nccl", rank=rank, world_size=world, timeout=datetime.timedelta(seconds=60)) torch.cuda.set_device(rank) def main(): rank = int(os.environ["RANK"]) world = int(os.environ["WORLD_SIZE"]) setup(rank, world) device = torch.device(f"cuda:{rank}") # 所有 rank 必须用【完全相同】的结构与初始化，否则 all_reduce 形状错位 torch.manual_seed(0) # 关键：保证各 rank 初始参数一致 model = nn.Sequential(nn.Linear(20, 10), nn.ReLU(), nn.Linear(10, 2)).to(device) model = DDP(model, device_ids=[rank]) opt = optim.SGD(model.parameters(), lr=0.01) x = torch.randn(8, 20, device=device) y = torch.randn(8, 2, device=device) for step in range(10): opt.zero_grad() out = model(x) loss = (out - y).pow(2).mean() loss.backward() opt.step() dist.barrier() # 每步同步，早暴露通信问题 if rank == 0: print(f"step {step} loss={loss.item():.4f}") dist.destroy_process_group() if __name__ == "__main__": main()

关键点是torch.manual_seed(0)保证各 rank 模型结构 / 初始化完全一致，避免 all_reduce 形状错位导致的非法访问。

四、根因一：各 rank 模型结构 / 输入不一致

DDP 的 all_reduce 假设所有 rank 的梯度 buffer 形状相同。如果：

• 某个 rank 因为数据 / 条件分支，模型走了不同的层（比如动态网络）；
• 或batch_size在各 rank 不同，导致某些 bucket 大小不一；

all_reduce 就会在读 / 写不匹配的 buffer 时非法访问。

修复：

• 所有 rank 用相同随机种子初始化模型（第四节manual_seed）；
• 各 rank 的 batch 大小保持一致（DataLoader 用drop_last=True）；
• 动态网络用find_unused_parameters=True（见下）。

model = DDP(model, device_ids=[rank], find_unused_parameters=True)

find_unused_parameters=True让 DDP 在反向时重新遍历图，找出真正用到的参数，避免「未用参数导致 bucket 错位」。但会略降性能，仅在有未用参数时开。

五、根因二：自定义 CUDA 扩展在 async 通信下越界

DDP 的反向通信在后台 stream异步跑。如果你的自定义 autograd 函数里有 CUDA 扩展，且它的 kernel 访问了「在通信完成前就被释放」的张量，就会非法访问。

修复：

• 在自定义函数里确保 buffer 生命周期覆盖通信；
• 或model.no_sync()控制哪些步做通信；
• 或用compute-sanitizer定位越界 kernel：

compute-sanitizer --tool memcheck torchrun --nproc_per_node=2 train.py

六、根因三：梯度 buffer 被提前释放（混合精度 / AMP）

用torch.amp时，梯度是 float16，而 DDP 的 bucket 通常是 float32。如果 AMP 的 GradScaler / 类型转换和 DDP bucket 交互出错，可能踩内存。

修复：用标准 AMP + DDP 写法，确保backward在autocast上下文正确：

scaler = torch.cuda.amp.GradScaler() with torch.autocast("cuda", dtype=torch.float16): out = model(x) loss = out.sum() scaler.scale(loss).backward() scaler.step(opt) scaler.update()

注意：scaler.step(opt)必须在loss.backward()之后、且在 DDP 的 all_reduce 完成前不要动参数。

七、根因四：CUDA_LAUNCH_BLOCKING 定位

和前面一样，CUDA_LAUNCH_BLOCKING=1让非法访问立刻定位到出事的 kernel：

export CUDA_LAUNCH_BLOCKING=1 torchrun --nproc_per_node=2 train.py

配合NCCL_DEBUG=INFO，能确认是「通信 stream 上的非法访问」还是「前向 / 反向 kernel 越界」。

八、根因五：GPU 硬件 / 驱动不稳定

如果所有代码写法都对，且只在某张特定卡上出现，可能是硬件 / 驱动（见 005 的 Xid 检查）：

dmesg | grep -i nvidia dmesg | grep -i xid

有Xid 79/Xid 48之类，就是硬件问题，需检修。

九、如何判断你踩的是 DDP 专属非法访问

• 报错是illegal memory access，且标注 DDP mode；
• 单机 / 单卡时不报，上 DDP 才报；
• 各 rank 可能结构 / batch 不一致；
• CUDA_LAUNCH_BLOCKING=1指向 all_reduce / 梯度 buffer 相关；
• 统一 seed、对齐 batch、开find_unused_parameters后恢复。

命中即说明是 DDP 特有的非法访问。

十、小结

DDP 下的illegal memory access比单机更隐蔽，因为它和异步梯度通信 + bucket 布局耦合。应对：

1. 所有 rank 用manual_seed+ 相同结构，避免 all_reduce 形状错位（第四节）；
2. 动态网络开find_unused_parameters=True，对齐各 rank batch（第四节、第五节）；
3. 自定义 CUDA 扩展注意 buffer 生命周期，用compute-sanitizer定位越界（第五节）；
4. AMP + DDP 用标准 GradScaler 写法（第六节）；
5. CUDA_LAUNCH_BLOCKING=1+NCCL_DEBUG=INFO定位（第七节）；
6. 排查硬件 / 驱动 Xid 错误（第八节）。

DDP 把训练「复制」到多卡，也把「内存访问」复制成了多份异步的危险。让所有 rank 的模型、数据、通信严格对齐，非法访问就失去了滋生土壤。