智体通鉴 · AgentGuide

外观

torch.multiprocessing 多进程训练详解 ​

内容整理自学习笔记，仅供面试备考参考；不构成录用、培训或考试承诺。

1. torch.multiprocessing是什么？ ​

torch.multiprocessing 是Python原生 multiprocessing 模块的封装，它注册了自定义的reducer，使用共享内存为不同进程提供相同数据的视图共享。

核心特性：

  • 一旦张量/存储被移动到 shared_memory（调用 share_memory()），就可以将其发送到其他进程而无需其他操作
  • API与原始模块100%兼容，只需将 import multiprocessing 改为 import torch.multiprocessing
  • 所有张量可通过队列发送或通过其他机制共享转移到共享内存中

警告： 如果主进程突然退出（如因传入信号），Python multiprocessing有时无法清理其子进程，可能造成资源泄漏。

2. torch.multiprocessing函数使用 ​

python
# 获取所有共享策略
torch.multiprocessing.get_all_sharing_strategies()

# 获取当前共享策略
torch.multiprocessing.get_sharing_strategy()

# 设置共享策略
torch.multiprocessing.set_sharing_strategy(new_strategy)
函数说明
get_all_sharing_strategies()返回当前系统支持的共享策略集合
get_sharing_strategy()返回当前CPU张量的共享策略
set_sharing_strategy(new_strategy)设置CPU张量的共享策略，参数须为系统支持的策略之一

3. 共享CUDA张量 ​

关键限制： 只支持在Python 3中使用 spawn 或 forkserver 启动方法才支持进程间共享CUDA张量。

Python 2的multiprocessing只能使用fork创建子进程，不支持CUDA运行时。

警告： CUDA API要求导出到其他进程的分配一直保持有效。需确保共享的CUDA张量不会超出必要的作用域。这不影响共享模型参数，但传递其他数据类型需小心。此限制不适用于共享CPU内存。

4. 共享策略详解 ​

共享策略仅适用于CPU张量，CUDA张量始终使用CUDA API。

4.1 文件描述符策略（file_descriptor） ​

这是默认策略（MacOS/OS X除外）。

特点说明
机制使用文件描述符作为共享内存句柄
存储方式从 shm_open 获取的fd被缓存，发送到其他进程时fd也被传送（通过UNIX套接字）
接收方缓存fd并mmap以获得共享视图

缺点： 共享大量张量时会保留大量文件描述符，若系统对打开fd数量有限制且无法提高，应使用 file_system 策略。

4.2 文件系统策略（file_system） ​

特点说明
机制使用 shm_open 的文件名标识共享内存区域
优点不需要缓存文件描述符
缺点容易发生共享内存泄漏——文件创建后不能被删除

内存泄漏问题： 进程崩溃不调用存储析构函数时，文件会保留在系统中直到重启或手动释放。

解决方案： torch.multiprocessing 会产生守护进程 torch_shm_manager，隔离于当前进程组，追踪所有共享内存分配。所有连接进程退出后，遍历并释放残留的共享内存文件。

建议： 如果系统支持 file_descriptor 策略，不建议切换到 file_system。

5. torch.multiprocessing实战：多进程训练模型 ​

python
import torch.multiprocessing as mp

def train(model):
    for data, labels in data_loader:
        optimizer.zero_grad()
        loss_fn(model(data), labels).backward()
        optimizer.step()  # 更新共享参数

# 定义模型
model = nn.Sequential(
    nn.Linear(n_in, n_h1),
    nn.ReLU(),
    nn.Linear(n_h1, n_out)
)

# 将模型移至共享内存（fork方法必须）
model.share_memory()

# 创建多个进程
processes = []
for i in range(4):  # 4个进程
    p = mp.Process(target=train, args=(model,))
    p.start()
    processes.append(p)

for p in processes:
    p.join()

关键步骤：

  1. 模型必须调用 share_memory() 以支持多进程共享参数
  2. 每个进程独立训练，但通过共享内存同步参数更新
  3. 使用 mp.Process 创建子进程，p.join() 等待所有进程完成

6. 共享策略选择指南 ​

策略适用场景优缺点
file_descriptor默认策略，系统fd限制够大时高效、无泄漏风险，但占用大量fd
file_systemfd受限且无法提高时，或MacOS不占fd，但有内存泄漏风险
数据类型共享方式
CPU张量file_descriptor 或 file_system
CUDA张量始终使用CUDA API（Python 3 + spawn/forkserver）