PyTorch 多进程编程权威指南：torch.multiprocessing 共享内存、CUDA 张量传递与进程管理
【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch
导读
本文以 PyTorch 官方文档 docs/source/multiprocessing.md 为主体，结合仓库内 torch/multiprocessing 的实际实现与底层 C++ 共享内存管理器，系统讲解torch.multiprocessing的完整能力：从共享内存的底层机制（file_descriptor与file_system两种策略）、CUDA 张量跨进程传递的最佳实践，到spawn多进程管理与错误传播。读完本文，你将掌握如何让 Tensor 在进程间零拷贝共享、如何规避 CUDA 共享时的内存泄漏与段错误陷阱，以及如何用torch.multiprocessing.spawn可靠地管理多个子进程的生命周期。
一、torch.multiprocessing 是什么
torch.multiprocessing是 Python 标准库multiprocessing的封装，它注册了自定义的 reducer，利用共享内存让不同进程中的 Tensor 拥有同一份底层数据的共享视图。核心流程是：一旦 Tensor/Storage 被移动到共享内存（调用 torch.Tensor.share_memory_），就可以在不做任何拷贝的情况下把它发送给其他进程。
从 torch/multiprocessing/init.py 的模块文档可以看到其设计哲学：
该模块 API 与原版 100% 兼容——只需把
import multiprocessing改为import torch.multiprocessing，所有通过队列发送或通过其他机制共享的 Tensor 都会自动移动到共享内存。
由于 API 高度相似，PyTorch 并未重复文档化大部分内容，而是推荐直接参考 Python 标准库的multiprocessing文档。这意味着你已有的多进程代码几乎可以零改动迁移到 PyTorch 版本，从而获得 Tensor 的共享内存传递能力。
1.1 内部初始化链路
在 torch/multiprocessing/init.py 中：
from multiprocessing import * # 继承标准库全部 API torch._C._multiprocessing_init() # 初始化 C++ 侧共享内存支持随后调用init_reductions()（定义于 torch/multiprocessing/reductions.py），向ForkingPickler注册 Tensor/Storage 的自定义序列化函数。这些 reducer 维护一个SharedCache（reductions.py），它是一个"multiprocessing 句柄 → Storage 弱引用"的字典，带锁保护并在 fork 后重置锁，避免死锁。
二、策略管理（Strategy management）
共享策略决定 CPU Tensor 底层数据如何在不同进程间传递。torch.multiprocessing提供三个顶层管理函数（定义于 torch/multiprocessing/init.py）：
| 函数 | 作用 |
|---|---|
get_all_sharing_strategies() | 返回当前系统支持的全部共享策略集合 |
get_sharing_strategy() | 返回当前生效的共享策略名称 |
set_sharing_strategy(new_strategy) | 设置共享策略；传入不在支持集合中的值时抛出AssertionError |
2.1 平台默认值
策略的默认选择逻辑写在 torch/multiprocessing/init.py：
if sys.platform == "darwin" or sys.platform == "win32": _sharing_strategy = "file_system" _all_sharing_strategies = {"file_system"} else: _sharing_strategy = "file_descriptor" _all_sharing_strategies = {"file_descriptor", "file_system"}- Linux 等平台：默认
file_descriptor，支持{"file_descriptor", "file_system"}； - macOS / Windows：只支持
file_system（file_descriptor依赖 UNIX 域套接字传递文件描述符，在这些平台不可用）。
仓库测试 test/distributed/rpc/test_share_memory.py 演示了如何在测试中临时切换到file_system并在结束后恢复原策略，可作为日常使用的参考写法。
三、共享 CUDA 张量：约束与三条最佳实践
CUDA 张量跨进程共享仅支持 Python 3，且必须使用spawn或forkserver启动方式（fork无法安全地继承 CUDA 上下文）。这与 CPU 张量有本质区别：发送进程必须一直持有原始张量，直到接收进程不再需要它为止。引用计数在底层已经实现，但需要用户遵守以下最佳实践。
3.1 实践一：在消费端尽快释放内存
消费进程拿到张量后应立即使用并删除，否则发送进程必须持续为其保留内存：
# 好：及时释放 x = queue.get() # 用 x 做点什么 del x# 差：持有到进程结束，生产者被迫一直保留 x 的内存 x = queue.get() # 用 x 做点什么 # 之后还做其他事（生产者不得不一直把 x 留在内存里）3.2 实践二：让生产者进程运行到所有消费者退出
这可以防止生产者提前释放消费者仍在使用中的内存：
# 生产者 # 发送张量、做其他事 event.wait() # 等待消费者发出完成信号# 消费者 # 接收张量并使用 event.set() # 通知生产者可以结束3.3 实践三：不要转发接收到的张量
CUDA 张量不能在进程间"接力"传递，否则无法正常工作甚至触发段错误：
# 这样不行 x = queue.get() queue_2.put(x)# 正确做法：先创建进程本地副本再转发 x = queue.get() x_clone = x.clone() queue_2.put(x_clone)# 在同一进程内对同一个队列先 put 再 get，极可能导致段错误 queue.put(tensor) x = queue.get()3.4 崩溃时的内存风险
原文档同时给出两条重要警告：
- 若主进程突然退出（例如收到信号），Python 的
multiprocessing有时无法清理子进程，中断解释器后可能出现资源泄漏； - 若消费进程因致命信号异常死亡，只要发送进程还在运行，共享张量就可能被永久保留在内存中。
CUDA 事件（Event）也可以通过 IPC 句柄在进程间共享，相关实现见 torch/multiprocessing/reductions.py 的_reduce_event/_rebuild_event。
四、共享策略详解：file_descriptor 与 file_system
本节仅适用于 CPU Tensor——CUDA 张量始终使用 CUDA 的 IPC API（这是它们唯一可行的共享方式）。
4.1 文件描述符策略file_descriptor
这是默认策略（macOS 与 Windows 除外）。工作方式：存储一旦移入共享内存，就从shm_open获得一个文件描述符并缓存到对象上；发送时通过 UNIX 套接字等机制把文件描述符传递给接收进程；接收进程同样缓存该描述符并执行mmap，从而获得存储数据的共享视图。
该策略的代价是：共享大量张量时，系统会长期保持大量打开的文件描述符。若你的系统对打开文件数有较低上限且无法提升，应改用file_system策略。
4.2 文件系统策略file_system
该策略使用shm_open生成的文件名来标识共享内存区域。好处是不需要在实现中缓存文件描述符；坏处是容易产生共享内存泄漏——文件创建后不能立即删除（其他进程还需要通过它打开自己的视图），一旦进程崩溃或被杀死而没有调用存储析构函数，这些文件就会残留在系统中，持续占用内存直到系统重启或手动释放。
为缓解这个问题，torch.multiprocessing会派生一个名为torch_shm_manager的守护进程：它把自己隔离出当前进程组，跟踪进程组内所有的共享内存分配；当所有相关进程退出后，它会等待片刻确保不会有新连接，然后遍历该组分配的全部共享内存文件，发现仍存在的便予以释放。这一机制在 C++ 侧由 torch/lib/libshm/core.cpp 的start_manager()实现——通过fork+pipe派生守护进程并交换握手信息（socket handle），随后以execl启动名为torch_shm_manager的可执行文件。
原文档的结论是：该清理方法经过测试，对各种故障足够健壮；但只要系统限制允许且file_descriptor可用，官方仍不推荐切换到file_system。
五、派生子进程：torch.multiprocessing.spawn
手动创建多个Process实例并逐个join的做法，在单进程场景下没有问题，但面对多个进程时存在隐患：串行join意味着进程按顺序终止，若第一个进程不退出，后续进程的终止将无人察觉；同时标准库也没有原生的错误传播机制。
torch.multiprocessing.spawn正是为解决这些问题而生，它提供错误传播、乱序终止检测，并在某个进程出错时主动终止其余进程。
5.1 spawn 函数签名
spawn定义于 torch/multiprocessing/spawn.py：
spawn(fn, args=(), nprocs=1, join=True, daemon=False, start_method="spawn")| 参数 | 说明 |
|---|---|
fn | 子进程入口函数，必须定义在模块顶层（multiprocessing 要求可 pickle）；以fn(i, *args)调用，i为进程索引 |
args | 传给fn的参数元组 |
nprocs | 派生的进程数 |
join | 是否阻塞等待所有进程结束；为True返回None，为False返回ProcessContext |
daemon | 子进程的 daemon 标志 |
start_method | 已废弃参数——该方法总是使用spawn；如需其他启动方式请用start_processes() |
5.2 错误传播与清理机制
从 torch/multiprocessing/spawn.py 可以看到，每个子进程由_wrap包装执行：发生异常时，traceback 会被 pickle 写入一个独立的临时错误文件（pytorch-errorfile-*.pickle，使用文件而非队列是为了避免死锁），随后以退出码 1 退出。
父进程侧（ProcessContext.join，spawn.py）通过multiprocessing.connection.wait监听所有进程的 sentinel：一旦发现某个进程退出码非零，就执行分级清理——先给出grace_period让其他进程自行退出，然后发送SIGTERM，若再等待一个宽限期仍未退出则升级为SIGKILL。随后根据错误文件是否存在抛出两种异常：
ProcessRaisedException：子进程抛出了异常，父进程会收到包含原始 traceback 的完整报错；ProcessExitedException：子进程因信号或非零退出码终止，异常携带exit_code与signal_name信息。
两者均继承自ProcessException（spawn.py），记录error_index与error_pid以便定位出错的进程。
5.3 ProcessContext 与 SpawnContext
调用spawn(..., join=False)会返回ProcessContext，其join(timeout=None, grace_period=None)方法可多次调用：
- 返回
True表示所有进程已全部成功结束； - 返回
False表示还有进程未结束； - 任一进程失败时抛异常，并清理其余进程。
SpawnContext是ProcessContext的历史名称，自 1.4 版本起已重命名，仍保留并发出弃用警告（spawn.py）。
5.4 start_processes：更通用的底层 API
若需要fork、forkserver等其他启动方式，应使用start_processes(fn, args, nprocs, join, daemon, start_method)（spawn.py）。它在forkserver且设置环境变量TORCH_MP_PARALLEL_START=1时通过ThreadPoolExecutor并行启动进程以提升性能（见 ENV_VAR_PARALLEL_START）。该 API 当前未写入官方文档，但在 IPython Notebook 等环境中fork往往比spawn表现更好，值得了解。
六、配套模块：pool、queue 与 reductions
原文档为这三个模块预留了文档位（见 docs/source/multiprocessing.md），它们在仓库中均有完整实现：
6.1 torch.multiprocessing.pool
torch/multiprocessing/pool.py 提供了Pool类，继承自multiprocessing.pool.Pool但改用自己的SimpleQueue，从而让 Tensor 以共享内存形式在池的进程间传递而非序列化底层数据。其clean_worker在标准 worker 结束后显式调用gc.collect()，确保所有析构函数（包括共享内存释放）被真正执行。
6.2 torch.multiprocessing.queue
torch/multiprocessing/queue.py 中的Queue与SimpleQueue通过ConnectionWrapper在发送时使用ForkingPickler序列化对象——这正是 Tensor 走自定义共享内存 reducer 的关键路径，也是set_sharing_strategy等设置能生效的底层原因。
6.3 torch.multiprocessing.reductions
torch/multiprocessing/reductions.py 是共享机制的核心，包含：
StorageWeakRef：对 Storage 的弱引用，配合SharedCache跟踪句柄与 Storage 的对应关系；SharedCache：带锁与 fork 重置机制的句柄缓存，长度超限时自动清理已失效引用；rebuild_tensor/_reduce_event等重建函数：负责在接收端根据句柄与元数据（offset、size、stride、requires_grad）恢复张量。
七、实战要点总结
- 迁移成本极低：把
import multiprocessing换成import torch.multiprocessing即可让队列中的 Tensor 自动共享内存，无需拷贝。 - CPU 张量：Linux 默认
file_descriptor策略高效但占用 fd 较多；受限于系统 fd 上限时切换set_sharing_strategy("file_system")，并依赖torch_shm_manager守护进程清理泄漏文件。 - CUDA 张量：务必使用
spawn/forkserver；在消费端及时del、让生产者运行到消费者全部退出、绝不转发收到的张量，避免内存滞留与段错误。 - 多进程管理：优先使用
torch.multiprocessing.spawn而非手写循环join，它自带错误传播（ProcessRaisedException/ProcessExitedException）、乱序终止检测与SIGTERM→SIGKILL分级清理，显著提升健壮性。 - 守护进程说明：
torch_shm_manager由 PyTorch 在需要时自动派生（实现见 torch/lib/libshm/core.cpp），用户无需手动启动，只需了解其内存清理职责即可。
【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
