在Pytorch中，如何在同一个GPU上的不同进程之间共享GPU内存？

Question

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

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

我的问题是，我的模型占用了相当多的内存空间。我在GPU上有12GB的内存，而模型本身就占用了大约3GB的内存（不包括数据）。这意味着，我的两个进程共占用了6GB的内存，仅用于模型。

---

我在想是否有可能只加载一次模型，然后在两个不同的进程中使用该模型进行推断。我的目标是仅使用3GB的内存来消耗模型，但仍然具有2个进程。

---

我在这个答案中看到了IPC的提及，但据我理解，它意味着进程#2将从进程#1复制模型，因此最终仍会分配6GB的模型空间。

我还查阅了Pytorch文档，关于DataParallel和DistributedDataParallel，但似乎不可行。 这篇文章似乎是我想要的，但我找不到任何关于如何在Pytorch的推断模式下使用的代码示例。

---

我知道这可能很难作为训练的一部分来完成，但请注意我只是在谈论推理步骤（模型处于只读模式，不需要更新梯度）。在这种情况下，我不确定是否可能。

- Astariul

---

我不明白为什么你不能只是使用相同的（只读）模型进行推断。您可以将不同的数据批次传递到同一模型中，数据加载和推断可以并行进行。多个用户也可以通过更高级别的接口与模型进行交互。是什么瓶颈导致您使用两个进程呢？ - THN

我会使用一个进程来加载一个模型并进行推理。这对大多数目的都是可行的。您想要实现什么样的目标？ - THN

通过在单个进程中进行数据加载并行处理（与模型运行过程分离，可以手动完成；tensorflow具有最佳并行数据预加载的本地支持，您可以查看它以获取示例），您可以获得大部分并发性的好处。同时使用单一模型即可。 - THN

如果您运行这样的进程，然后将其fork成两个不同的进程，每个进程都充当服务器并开始侦听不同的套接字，那会怎么样呢？问题在于通过共享GPU内存，您必须同步这两个进程，以便它们不会同时使用（相同的）GPU内存。 - Raz Rotenberg

@THN 我不知道“你可以在单个进程的单个模型上获得多数并发性的好处”。 我认为，如果内存允许，加载两个进程是更有效率的，这样它们就可以并行运行。请发表答案！ - Astariul

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

- nairbv

---

非常感谢您提供如此详细的答案。我一定要阅读所有这些资源。 - Astariul

Answer 2

您可以通过在数据加载和模型推断过程中进行并发，使用单个模型在单个进程中获得大部分并发的好处（只读推断）。数据加载与模型运行过程分离，这可以手动完成。据我所知，tensorflow对于最佳并行数据预加载有一些本地支持，您可以查看示例。模型推断在GPU上自动并行化。您可以通过使用更大的批次来最大化此并发性。从架构角度来看，多个用户也可以通过更高级别的接口与模型交互。

- THN

---

你应该关注作业调度问题，这在操作系统中已经有很多研究并且有多种算法。实际上，作业不会同时到达，因此你可以在加载另一个作业时处理当前的作业。如果必要，你可以将作业批量处理，或者如果等待时间可以忽略，则按顺序处理，或者如果作业太大，则将其分成若干部分进行处理。 - THN

我为我的特定情况进行了一些基准测试：如果有10个客户端请求预测，使用同一GPU上的2个进程为所有客户端提供服务需要0.96秒。仅使用单个进程进行相同实验需要1.42秒。 - Astariul

很好，你实际测试了一下，但请注意每个结果都是一个轶事。如果所有请求同时到达，并且它们仅占用GPU的一小部分，并且您单独处理每个请求，则使用2个或更多进程肯定会更快。但有些情况下，一个进程已经足够好，比如请求随机到来；或者一个进程更好，比如模型庞大且请求可以批处理在一起。毕竟，您需要查看自己的典型用例，找出瓶颈，并决定在哪里进行优化。 - THN

使用多个CPU进程来读取请求、加载数据并将它们批处理，然后在一个GPU进程上运行，与您最初关于在GPU上共享内存（实际上是模型参数）的问题相同。您仍然需要为此而努力。 - THN

顺便提一下，PyTorch也有并行数据加载器。 - nairbv