必须用 spawn 启动且 CUDA 操作延后至 DataLoader 创建之后,否则 fork 会复制不可重用的 CUDA 上下文导致死锁;spawn 需在 if name == '__main__': 中调用 torch.multiprocessing.set_start_method('spawn', force=True);DDP 初始化需确保 init_method、world_size 一致,且 torch.cuda.set_device(rank) 在 init_process_group 后、模型构建前执行。

必须在 if __name__ == '__main__': 中显式设 spawn 启动方式,且所有 CUDA 操作(包括 torch.cuda.set_device、model.to('cuda'))必须放在 DataLoader 创建之后——这不是可选优化,而是避免死锁的硬性前提。
为什么 fork 启动在多 GPU 下大概率卡死
Linux 默认用 fork 启动子进程,它会复制父进程已初始化的 CUDA 上下文、NCCL 句柄、OpenCV 线程池等全局状态。这些资源在子进程中不可重用,导致 worker 进程卡在 futex 或 cudaStreamSynchronize 上,表现为 for batch in dataloader: 无响应、CPU 占用低、GPU 显存不动。
- 现象高度随机:可能第 1 个 epoch 就卡,也可能跑完 20 个 epoch 才崩
-
ps aux | grep python可见多个 worker 处于D(不可中断睡眠)状态 - 只要主进程提前调用过
torch.cuda.is_available()、cv2.imread()或matplotlib.use(),风险就已埋下
必须用 spawn,且 force=True
spawn 启动方式让每个子进程从零开始初始化 Python 解释器和 CUDA 环境,从根本上规避 fork 的状态污染问题。但它必须被正确启用:
- 位置必须在
if __name__ == '__main__':块内,且在任何import torch之后、任何DataLoader实例化之前 - 必须加
force=True:torch.multiprocessing.set_start_method('spawn', force=True),否则旧环境或已有 multiprocessing 初始化会报RuntimeError: context has already been set - Jupyter 中无法使用 —— 它的内核已启动 multiprocessing 上下文,必须导出为 .py 脚本后运行
DDP 初始化阶段的三个关键避坑点
即使用了 spawn,torch.distributed.init_process_group 仍极易因配置不一致引发隐性死锁:
图片提示词生成器?不止如此。 马甲系统 —— 把脑海中的画面,翻译成AI能理解的专业表达。 用得越多,它越懂你:首次需要多问几句确认方向,用久了几乎一说就懂。 用得越多,它越快:缓存机制让后续对话越来越省。 RAG进化:成功案例持续入库,越跑越聪明。 输入「新手指南」查看完整功能介绍
立即学习“Python免费学习笔记(深入)”;
-
init_method必须全 rank 严格一致:若用file://,路径需 NFS 共享且文件未被残留进程占用;若用tcp://,rank 0 必须能被其他节点telnet host port连通 -
world_size必须等于实际启动的进程数,不能写死为 4 却只启动 3 个torchrun进程 -
torch.cuda.set_device(rank)必须在init_process_group之后、模型构建之前调用;否则 NCCL 通信时会因torch.cuda.current_device()返回错误设备序号而阻塞在all_gather_object等集体操作上
模型加载与数据预处理的设备隔离原则
GPU 资源不能跨进程共享,任何试图“复用”主进程 CUDA 状态的操作都会失效:
- 主进程加载模型时必须用
map_location='cpu',子进程各自调用.to(rank);传model.to('cuda:0')给Pool.map会导致静默卡死 -
Dataset.__getitem__中禁止出现任何 GPU 操作:torch.randn(..., device='cuda')、model(input)、甚至tokenizer.encode(..., return_tensors='pt')若内部触发 CUDA 张量创建,都会死锁 -
pin_memory=True会强制子进程预分配 CUDA 上下文,加剧冲突;调试期建议显式设为False
真正容易被忽略的是:死锁往往不是某一行代码的问题,而是多个看似独立的配置项(CUDA_VISIBLE_DEVICES、set_start_method、init_process_group 参数、torch.cuda.set_device 顺序)共同构成的脆弱链路。少一个环节对齐,就可能在训练到第 5 个 epoch 时突然卡住,且无任何报错信息。

















