pytorch多进程推理不能直接用torch.multiprocessing.spawn共享模型,因cuda上下文不隔离会导致冲突;须每个进程独立构建加载模型、绑定独立gpu设备、用distributedsampler切分数据、以all_gather合并输出。

PyTorch多进程推理为什么不能直接用torch.multiprocessing.spawn加载同一个模型?
因为模型对象(尤其是含nn.Module和cuda()状态的)不是进程安全的:多个子进程共享模型引用时,会触发CUDA上下文冲突或梯度图残留,常见报错是RuntimeError: unable to open shared memory object或CUDA error: initialization error。
正确做法是每个进程独立构建、加载模型——不是“复制”主进程的模型实例,而是各自调用model = MyModel().load_state_dict(...)。
- 模型加载必须在子进程内部完成,不能在主进程里加载再传入
- 如果用
torch.load(..., map_location='cpu')预加载再分发,依然会因CUDA上下文未隔离而失败 - 务必为每个进程设置独立的CUDA设备号,通过
torch.cuda.set_device(rank)绑定
如何用torch.distributed做数据切分而非模型并行?
分布式推理不等于模型并行;多数场景只需要把输入数据按batch切片,让每个GPU跑一个子集,最后合并结果。这时候用torch.distributed.init_process_group(backend='nccl')只是为了同步和切分,不需要DistributedDataParallel。
关键函数是torch.utils.data.distributed.DistributedSampler,它会在每个进程里自动跳过不属于本rank的数据索引。
-
DistributedSampler必须配合DataLoader(sampler=...)使用,且shuffle=False时行为可预测 - 若自己手写切分逻辑(比如用
data[start:end]),要确保start = rank * chunk_size,且chunk_size = len(data) // world_size,余数部分需显式处理,否则丢样本 - 注意
num_workers > 0时,子进程内再fork会触发CUDA初始化错误,建议设为0或改用spawn启动方式
torch.multiprocessing.spawn vs mp.spawn:启动方式选哪个?
二者本质相同,但torch.multiprocessing.spawn默认启用fork方式,在Linux上可能继承主进程的CUDA上下文,导致子进程崩溃;而mp.spawn(即torch.multiprocessing.spawn的别名)支持显式指定start_method='spawn',这才是安全选择。
- 必须加参数
start_method='spawn',尤其在CUDA环境下 - 入口函数(如
main_worker)必须接收rank和args,并在开头调用init_process_group - 不要在入口函数外定义模型或数据,避免被序列化时带入不可pickle对象(如CUDA tensor)
推理输出怎么合并?要注意哪些隐式同步点?
每个进程输出自己的torch.Tensor结果后,需要用torch.distributed.all_gather收集,而不是靠Python list拼接——后者无法跨进程通信,list里全是空或本地副本。
容易忽略的是:所有gather操作前,必须保证各进程到达同一逻辑点,否则死锁。最稳妥的方式是所有进程都执行一次all_gather,哪怕某rank没数据也要传一个占位tensor。
- 用
torch.empty预分配output tensor,大小统一,避免shape不一致报错 - 如果某rank无数据,仍要调用
all_gather,否则其他rank卡在等待 - CPU tensor也能用
all_gather,但需确保device一致;混合CPU/GPU时建议全转到CPU再gather
事情说清了就结束。真正麻烦的不是写几行spawn或all_gather,而是每个进程里模型加载路径、设备绑定、数据切片边界、tensor device一致性这四点稍有错位,就会静默丢数据或卡死——最好先用2个进程+小数据集跑通全流程,再扩规模。
Python免费学习笔记(深入):立即使用
在学习笔记中,你将探索 Python 的核心概念和高级技巧!











