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

Question

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

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

我的问题是我的模型占用了相当多的内存空间。我的 GPU 上有 12Gb 内存，模型仅占用约 3Gb 内存(不包含数据)。这意味着我的 2 个进程总共需要 6Gb 内存用于模型。

<小时/>

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

<小时/>

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

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

This似乎是我想要的，但我找不到任何关于如何在推理模式下使用 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？)，但无论如何，这似乎只是通过时间切片实现的，所以我不确定我们是否应该期望更高的总吞吐量:

How do I use Nvidia Multi-process Service (MPS) to run multiple non-MPI CUDA applications?

如果您确实需要在两个并行推理调用之间共享一个模型的内存，是否可以仅使用多个线程而不是进程，并从两个线程引用同一模型？

要真正让 GPU 并行运行多个内核，您可以在 PyTorch 中使用 nn.Parallel。请参阅此处的讨论: https://discuss.pytorch.org/t/how-can-l-run-two-blocks-in-parallel/61618/3