[Bug已解决] torch.ops.symm_mem.one_shot_all_reduce 触发 CUDA_ERROR_ILLEGAL_ADDRESS 解决方案

[Bug已解决] torch.ops.symm_mem.one_shot_all_reduce 触发 CUDA_ERROR_ILLEGAL_ADDRESS 解决方案一、现象长什么样你使用 PyTorch 实验特性_symmetric_memory对称内存里的「一次性全规约one-shot all-reduce」算子时可能遇到CUDA_ERROR_ILLEGAL_ADDRESS也就是官方描述的torch.ops.symm_mem.one_shot_all_reduce causes CUDA_ERROR_ILLEGAL_ADDRESSsymm_mem.one_shot_all_reduce是对称内存体系下的一个高效集合通信原语利用各 rank 的「对称内存窗口」一次操作就把数据规约完one-shot省去传统 NCCL 的多轮握手。但因为它直接操作跨 rank 的内存窗口一旦窗口地址 / 大小 / 对齐有任何不对就会访问非法地址 →CUDA_ERROR_ILLEGAL_ADDRESS错误码 700。 本文讲清楚这个算子的语义、为何会非法访问、以及如何规避。二、one-shot all-reduce 与普通 all_reduce 的区别普通dist.all_reduce走 NCCL多轮通信把各 rank 数据聚合并发回。而symm_mem.one_shot_all_reduce依赖各 rank 预先注册好的对称内存窗口地址对称、可互相直接访问「one-shot」意味着一次操作完成规约类似直接写远端窗口 一次同步极致性能但要求窗口完全正确对齐、大小一致、地址有效。 正因为「直接操作远端内存」它对正确性的要求比 NCCL 严苛得多——NCCL 出错大多报通信错而 one-shot 出错直接是内存非法访问illegal address更底层、更危险。三、为什么会出现 CUDA_ERROR_ILLEGAL_ADDRESS常见成因对称内存窗口未正确注册 / 大小不匹配某个 rank 的窗口比别的 rank 小one-shot 规约时越界访问窗口地址失效注册后参数被改动 / 显存被回收地址无效数据类型 / 对齐不一致各 rank 参与规约的 tensor dtype / 形状不一致导致偏移算错它是实验特性symm_mem带下划线、未稳定one-shot 原语的实现可能在某些 shape / dtype 下有 bug与 CUDA Graph / 异步 stream 冲突one-shot 对执行顺序敏感被图捕获或异步打断会踩内存。四、可运行先确认普通 all_reduce 是否正常排除 NCCL 本身在冒险用 one-shot 前先确认基础集合通信健康import os import datetime import torch import torch.distributed as dist def basic_check(rank, world): os.environ[MASTER_ADDR] os.environ.get(MASTER_ADDR, localhost) os.environ[MASTER_PORT] os.environ.get(MASTER_PORT, 29502) dist.init_process_group(nccl, rankrank, world_sizeworld, timeoutdatetime.timedelta(seconds60)) torch.cuda.set_device(rank) t torch.ones(4, devicefcuda:{rank}) dist.all_reduce(t, opdist.ReduceOp.SUM) torch.cuda.synchronize(rank) print(f[rank {rank}] 普通 all_reduce 结果应为 world{world}{t}) if __name__ __main__: if torch.cuda.is_available(): basic_check(int(os.environ.get(RANK, 0)), int(os.environ.get(WORLD_SIZE, 1)))如果普通all_reduce也失败说明是更底层的 NCCL / 设备问题先修那个再来看 one-shot。五、解决方案一用标准 dist.all_reduce 替代最稳symm_mem.one_shot_all_reduce是实验特性、限制多、易非法访问。绝大多数场景用稳定的dist.all_reduce即可import torch import torch.distributed as dist def safe_all_reduce(tensor): dist.all_reduce(tensor, opdist.ReduceOp.SUM) return tensor # 用法 t torch.randn(4, devicefcuda:{rank}) safe_all_reduce(t)dist.all_reduce经过充分测试不会动「跨 rank 内存窗口」自然没有 illegal address 风险。六、解决方案二确保对称内存窗口大小 / dtype / 对齐一致如果你确实要用symm_mem必须保证所有 rank 的窗口完全一致# 伪代码所有 rank 用相同 shape/dtype 注册对称内存 shape (1024, 1024) dtype torch.float32 # 每个 rank 都按相同 shape/dtype 分配对称内存窗口 buf torch.empty(shape, dtypedtype, devicefcuda:{rank}) # 绝不要某个 rank 用 (1024,1024)、另一个用 (512,512)one-shot all-reduce 对「窗口大小 / dtype」极其敏感任何不一致都会让偏移算错、越界访问。七、解决方案三用 CUDA_LAUNCH_BLOCKING 定位开同步执行确认非法访问就发生在 one-shot 算子export CUDA_LAUNCH_BLOCKING1 torchrun --nproc_per_node2 your_script.py配合compute-sanitizer精确定位越界compute-sanitizer --tool memcheck torchrun --nproc_per_node2 your_script.py八、解决方案四避开与 CUDA Graph / 异步冲突one-shot all-reduce 对执行顺序敏感不要把它塞进 CUDA Graph也不要和异步 stream 混用# ❌ 危险把 one-shot 包进 CUDA Graph # g torch.cuda.CUDAGraph() # with torch.cuda.graph(g): # symm_mem.one_shot_all_reduce(...) # ✅ 安全在普通 stream 上、同步调用 symm_mem.one_shot_all_reduce(tensor) torch.cuda.synchronize(rank)九、解决方案五升级 PyTorchsymm_mem是 PyTorch 较新、快速演进的特性。one-shot all-reduce 的 illegal address bug 很可能在新版本被修。查看并升级import torch print(PyTorch, torch.__version__) print(有 symm_mem, hasattr(torch.ops, symm_mem))十、如何判断你踩的是同一条你调用了torch.ops.symm_mem.one_shot_all_reduce报错是CUDA_ERROR_ILLEGAL_ADDRESS700普通dist.all_reduce正常one-shot 才失败各 rank 窗口可能大小 / dtype 不一致或与 Graph / 异步混用退回dist.all_reduce后恢复。 命中即说明踩中该 one-shot 原语的非法访问 bug。十一、小结torch.ops.symm_mem.one_shot_all_reduce触发CUDA_ERROR_ILLEGAL_ADDRESS是对称内存窗口直接操作远端内存一旦窗口不对就非法访问。应对绝大多数场景用稳定的dist.all_reduce替代第五节最稳真要用symm_mem保证所有 rank 窗口shape / dtype / 对齐完全一致第六节CUDA_LAUNCH_BLOCKING1compute-sanitizer定位越界第七节别把 one-shot 塞进 CUDA Graph / 异步 stream第八节升级到修复该原语 bug 的 PyTorch第九节。 one-shot all-reduce 用「直接写远端内存」换取极致性能但也把「内存安全」的责任完全交给你。在它还带下划线、未稳定时生产里优先用久经考验的dist.all_reduce——性能差一点但绝不会给你非法地址。