使用 Pytorch 在同一 GPU 上的进程之间共享 GPU 内存

Question

我正在尝试实现一种在 Pytorch 中进行并发推理的有效方法。

现在，我在我的 GPU 上启动了 2 个进程（我只有 1 个 GPU，两个进程都在同一个设备上）。每个进程都会加载我的 Pytorch 模型并执行推理步骤。

我的问题是我的模型在内存上占用了相当多的空间。我在 GPU 上有 12Gb 的内存，而模型仅占用约 3Gb 的内存（没有数据）。这意味着，我的 2 个进程只为模型占用了 6Gb 的内存。

---

现在我想知道是否可以只加载一次模型，并使用这个模型对 2 个不同的进程进行推理。我想要的是模型只消耗 3Gb 的内存，但仍然有 2 个进程。

---

我遇到了这个提到 IPC 的答案，但据我了解，这意味着进程 #2 将从进程 #1 复制模型，因此我最终仍会为模型分配 6Gb。

我还查看了关于 DataParallel 和 DistributedDataParallel 的 Pytorch 文档，但这似乎是不可能的。

这似乎是我想要的，但我找不到任何关于如何在推理模式下使用 Pytorch 的代码示例。

---

我知道这可能很难为训练做这样的事情，但请注意我只是在谈论推理步骤（模型处于只读模式，无需更新梯度）。有了这个假设，我不确定这是否可能。

Answer 1

GPU 本身有很多线程。在执行数组/张量操作时，它会在数组的一个或多个单元上使用每个线程。这就是为什么看起来一个可以充分利用 GPU 的运算应该在没有多个进程的情况下有效扩展的原因——单个 GPU 内核已经被大规模并行化了。在评论中，您提到在小型基准测试中通过多个进程看到更好的结果。我建议使用更多作业运行基准测试以确保预热，十个内核似乎太小了一个测试。但是，如果您要找到一个彻底的代表性基准来始终更快地运行，那么我会相信良好的基准而不是直觉。我的理解是在默认 CUDA 流上启动的内核会按顺序执行。如果您希望它们并行运行，我认为您需要多个流。查看 PyTorch 代码，我看到getCurrentCUDAStream()内核中的代码，这让我认为 GPU 仍将按顺序运行来自所有进程的任何 PyTorch 代码。这个 NVIDIA 讨论表明这是正确的：https://devtalk.nvidia.com/default/topic/1028054/how-to-launch-cuda-kernel-in-different-processes/较新的 GPU 可能能够并行运行多个内核（使用 MPI？），但似乎这只是通过时间切片实现的，所以我不确定我们是否应该期望更高的总吞吐量：如何使用 Nvidia 多进程服务 (MPS) 运行多个非 MPI CUDA 应用程序？如果您确实需要在两个并行推理调用中共享一个模型的内存，您是否可以只使用多个线程而不是进程，并从两个线程引用相同的模型？要真正让 GPU 并行运行多个内核，您可以在 PyTorch 中使用 nn.Parallel。请参阅此处的讨论： https ://discuss.pytorch.org/t/how-can-l-run-two-blocks-in-parallel/61618/3

Answer 2

通过在数据加载和模型推理中进行并发，您可以在单个进程（只读）推理中获得单个模型的并发性的大部分好处。数据加载与模型运行过程是分开的，这可以手动完成。据我所知，对优化并行数据预加载tensorflow有一些原生支持，您可以查看它作为示例。模型推理在 GPU 上自动并行。您可以通过使用更大的批处理来最大化这种并发性。从架构的角度来看，多个用户也可以通过更高级别的界面与模型对话。