训练跑到第两百步突然报 CUDA error: an illegal memory access was encountered,堆栈只指到 loss.backward(),看不出具体算子。

一开始怀疑是「多卡通信时张量形状不一致」,查了一圈,方向不对。

真正麻烦的地方在于,几个原因的现象几乎一样。自定义算子或扩展越界访问显存。少数情况下是索引张量超出范围,gather 时读到了非法地址。

还有一种情况,驱动或框架版本存在已知缺陷。

先打开同步执行让报错停在真正的算子

CUDA_LAUNCH_BLOCKING=1 python train.py

然后用异常检测定位到具体前向算子

python -c "import torch;torch.autograd.set_detect_anomaly(True)"

再往下检查索引是否越界并加断言

assert idx.max().item() < vocab_size

缩小 batch 与序列长度复现,排除显存越界

python train.py --batch-size 1 --max-len 128

恢复的判定其实很直接。开启同步后报错能定位到具体文件和行号,修复越界索引后训练可以稳定跑过原来的步数。

异步执行下报错位置会延后很多步,看到栈顶是 backward 不代表反向有问题,一定要开同步再复现一次。改动前把原配置备份一份,出问题能马上回滚,这比任何技巧都实用。

作者 9823189