AWS Lambda上PyTorch用multiprocessing.Pipe并行无加速问题排查
问题背景
我尝试使用multiprocessing.Pipe在AWS Lambda上并行化基于PyTorch的代码，本地机器增加进程数能获得明显性能提升，但Lambda上却无法实现并行加速。为测试该问题，我编写代码将n次矩阵乘法分配到m个进程中执行，使用10GB内存的Lambda实例得到'torch'模式测试结果如下：
| 测试场景 | Lambda耗时 | 本地耗时 |
|---|---|---|
| 150次矩阵乘法（1进程） | 9.3s | 21.7s |
| 150次矩阵乘法（2进程） | 9.9s | 15.3s |
| 150次矩阵乘法（3进程） | 9.9s | 18.1s |
| 350次矩阵乘法（1进程） | 15.6s | 41.9s |
| 350次矩阵乘法（2进程） | 16.0s | 25.4s |
| 350次矩阵乘法（3进程） | 16.2s | 25.5s |
| 700次矩阵乘法（1进程） | 27.8s | 79.8s |
| 700次矩阵乘法（2进程） | 28.4s | 45.7s |
| 700次矩阵乘法（3进程） | 28.8s | 41.7s |
Lambda配置
- 内存规格：10GB
- Docker镜像：
public.ecr.aws/lambda/python:3.9 - Torch安装命令：
pip3 install torch==1.12.0+cpu torchvision==0.13.0+cpu -f https://download.pytorch.org/whl/torch_stable.html
补充说明：代码支持三种模式——torch（PyTorch实现）、numpy（NumPy实现）、lists（Python列表实现），其中lists和numpy模式下Lambda与本地的并行性能表现一致，仅PyTorch模式出现异常。
原因分析
PyTorch默认多线程与Lambda vCPU限制冲突
AWS Lambda的vCPU分配规则为每GB内存对应0.2vCPU，10GB内存实例仅提供2个vCPU核心。而PyTorch默认通过OpenMP启用多线程优化，单进程即可占满全部2个核心。此时新增进程会引发CPU资源竞争，不仅无法并行计算，还会增加进程调度开销，导致耗时小幅上升。进程间张量传递的序列化开销
使用multiprocessing.Pipe传递PyTorch张量时，会触发完整的序列化与反序列化操作（PyTorch张量默认不支持fork-safe的共享内存传递）。Lambda环境中进程间内存拷贝的开销比本地更显著，进一步抵消了并行计算的潜在收益。与NumPy/Python列表的差异
- NumPy的矩阵乘法默认线程数较少，或可通过环境变量灵活控制，多进程可利用剩余CPU资源；
- Python列表实现的矩阵乘法为纯单线程计算，多进程可直接将任务拆分到不同核心，因此Lambda上并行有效。
解决方法
- 限制PyTorch单进程线程数
通过设置环境变量OMP_NUM_THREADS=1，强制PyTorch单进程仅使用1个线程，让多进程能分别占用独立vCPU核心，避免资源竞争。在Lambda代码开头添加：
import os os.environ["OMP_NUM_THREADS"] = "1"
使用PyTorch专用多进程工具
替换multiprocessing.Pipe为torch.multiprocessing模块，该模块支持张量的共享内存传递，大幅减少序列化开销。例如使用torch.multiprocessing.Queue，或采用spawn启动方式传递共享张量。调整Lambda实例规格
若需更高并行度，可选择内存更大的Lambda实例（如17GB内存对应4vCPU），在限制单进程线程数后，多进程能充分利用更多核心资源。
内容的提问来源于stack exchange，提问作者Rgkpdx
