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

[Bug已解决] DDP 模式下 CUDA error: an illegal memory access was encountered 解决方案一、报错长什么样你在用DistributedDataParallelDDP做多卡训练时训练可能跑着跑着崩CUDA error: an illegal memory access was encountered而且特别标注发生在DDP mode下。也就是官方描述的DDP mode: CUDA error: an illegal memory access was encounteredillegal memory access非法内存访问错误码 700意味着某个 CUDA kernel 访问了它不该访问的地址越界、空指针、已释放显存。在 DDP 下出现说明问题往往和分布式特有的逻辑有关——比如梯度桶gradient bucketbuffer、跨 rank 的 all_reduce 通信、或 DDP 包装后的参数布局。本文讲清楚 DDP 下非法访问的常见来源以及如何定位修复。二、DDP 为什么比单机更容易踩非法访问DDP 在单机多卡 / 多机上的工作方式每个 rank 有一份模型副本参数分到一个个梯度桶bucket反向时每个 bucket 的梯度攒满就通过all_reduce在各 rank 间同步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, rankrank, world_sizeworld, timeoutdatetime.timedelta(seconds60)) torch.cuda.set_device(rank) def main(): rank int(os.environ[RANK]) world int(os.environ[WORLD_SIZE]) setup(rank, world) device torch.device(fcuda:{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(), lr0.01) x torch.randn(8, 20, devicedevice) y torch.randn(8, 2, devicedevice) 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(fstep {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_lastTrue动态网络用find_unused_parametersTrue见下。model DDP(model, device_ids[rank], find_unused_parametersTrue)find_unused_parametersTrue让 DDP 在反向时重新遍历图找出真正用到的参数避免「未用参数导致 bucket 错位」。但会略降性能仅在有未用参数时开。五、根因二自定义 CUDA 扩展在 async 通信下越界DDP 的反向通信在后台 stream异步跑。如果你的自定义 autograd 函数里有 CUDA 扩展且它的 kernel 访问了「在通信完成前就被释放」的张量就会非法访问。修复在自定义函数里确保 buffer 生命周期覆盖通信或model.no_sync()控制哪些步做通信或用compute-sanitizer定位越界 kernelcompute-sanitizer --tool memcheck torchrun --nproc_per_node2 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, dtypetorch.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_BLOCKING1让非法访问立刻定位到出事的 kernelexport CUDA_LAUNCH_BLOCKING1 torchrun --nproc_per_node2 train.py配合NCCL_DEBUGINFO能确认是「通信 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_BLOCKING1指向 all_reduce / 梯度 buffer 相关统一 seed、对齐 batch、开find_unused_parameters后恢复。命中即说明是 DDP 特有的非法访问。十、小结DDP 下的illegal memory access比单机更隐蔽因为它和异步梯度通信 bucket 布局耦合。应对所有 rank 用manual_seed 相同结构避免 all_reduce 形状错位第四节动态网络开find_unused_parametersTrue对齐各 rank batch第四节、第五节自定义 CUDA 扩展注意 buffer 生命周期用compute-sanitizer定位越界第五节AMP DDP 用标准 GradScaler 写法第六节CUDA_LAUNCH_BLOCKING1NCCL_DEBUGINFO定位第七节排查硬件 / 驱动 Xid 错误第八节。DDP 把训练「复制」到多卡也把「内存访问」复制成了多份异步的危险。让所有 rank 的模型、数据、通信严格对齐非法访问就失去了滋生土壤。