具有共享内存的 Pytorch 多处理导致 matmul 慢 30 倍（只有两个进程）

Posted 2023-02-16

【中文标题】具有共享内存的 Pytorch 多处理导致 matmul 慢 30 倍（只有两个进程）

【英文标题】：Pytorch multiprocessing with shared memory causes matmul to be 30x slower (with only two processes)

【发布时间】：2021-03-11 10:04:47

【问题描述】：

我正在尝试通过使用多处理让多个工作人员同时产生经验来提高我的强化学习算法的速度。每个过程只运行我的神经网络的前向传递，不需要梯度计算。

As I understand it，当跨进程边界传递Tensors 和nn.Modules 时（使用torch.multiprocessing.Queue 或torch.multiprocessing.Pool），张量数据被移动到共享内存，这不应该比非慢- 共享内存。

但是，当我使用 2 个进程（在 8 核机器上）运行我的多进程代码时，我发现我的 pytorch 操作变得慢了 30 倍以上，这抵消了同时运行两个进程带来的加速。

我分析了我的应用程序，以找出具体是哪些操作变慢了。我发现我的大部分时间都花在nn.functional.linear() 上，特别是在Tensor.matmul 通话中的this line：

    output = input.matmul(weight.t())

我只是在这个具体的matmul调用中加了一个定时器，我发现当一个进程运行时，这个操作耗时不到0.3毫秒，但是当两个进程运行时，耗时超过10毫秒。请注意，在两种情况中，权重矩阵已被放入共享内存并跨越进程边界传递给工作进程，唯一的区别是在第二种情况下，有两个工作进程而不是一个。

作为参考，input 和 weight 张量的形状分别为 torch.Size([1, 24, 180]) 和 torch.Size([31, 180])。

什么可能导致这种急剧放缓？使用任何文档中未提及的火炬多处理或共享内存是否有一些微妙之处？我觉得这里肯定有一些隐藏的锁导致了争用，因为这种急剧的减速对我来说毫无意义。

【参考方案1】：

这似乎是由 OpenMP（默认情况下由 pytorch 使用）和多处理的不良交互引起的。这是 pytorch (https://github.com/pytorch/pytorch/issues/17199) 中的一个已知问题，我什至在我用来调试的某些配置中遇到了死锁。使用torch.set_num_threads(1) 关闭 OpenMP 为我解决了这个问题，并立即加快了我在多进程情况下的张量操作，大概是通过绕过 OpenMP 正在做的内部锁定。