如何在 GPU 上有效地并行化 AlphaZero？

Question

我正在实现一个 AlphaZero 版本（AlphaGo 的最新版本），以应用于其他领域。该算法的关键是状态空间 (CPU) 的蒙特卡洛树搜索与来自 eval 模式 (GPU) 的神经网络的“直觉”（概率）交错。然后使用 MCTS 结果来训练神经网络。我已经通过启动多个进程来并行化 CPU 执行，每个进程都建立自己的树。这是有效的，现在已经导致GPU 瓶颈！（nvidia-smi 始终以 100% 显示 GPU）我设计了 2 种策略来并行化 GPU 评估，但是它们都有问题。每个进程仅在其自己的树中的批次上评估网络。在我最初的幼稚实现中，这意味着批量大小为 1。但是，通过重构一些代码并添加“虚拟损失”来阻止（但不完全阻止）同一节点被选中两次，我们可以获得更大的批量大小 1 -4。这里的问题是，在我们评估批次或准确性受到影响之前，我们不能允许大的延迟，所以小批量是这里的关键。将批次发送到一个中央“神经网络工作者”线程，该线程组合和评估它们。这可以在 32 个或更多的大批量中完成，因此可以非常有效地使用 GPU。这里的问题是树工作者发送 CUDA 张量“往返”，这不受 PyTorch 支持。如果我先克隆它们是受支持的，但是所有不断的复制使这种方法比第一种方法慢。我在想也许一个我没有看到的聪明的批处理方案可以使第一种方法起作用。使用多个 GPU 也可以加速第一种方法，但是 PyTorch 本身并不支持我想要的那种并行性。也许将所有张量保留在 NN worker 中并且只发送 id 可以改进第二种方法，但是这里的困难是如何有效地同步以获得大批量而不使 CPU 线程等待太久。我几乎没有发现有关 AlphaZero 或 AlphaGo Zero 在各自论文中如何并行化的信息。我能够在网上找到有限的信息，但这导致我改进了第一种方法。如果有任何建议，我将不胜感激，特别是如果我错过了某些观点或方法。

Answer 1

以 tensorflow 服务为例，预测服务可以运行在不同的进程中，运行一个服务来接收来自 worker 的请求（运行一个 MCTS 进程并将预测请求发送到预测服务）。我们可以保留一个从套接字地址到套接字本身的字典。

预测服务可以读取每个查询正文及其标题（每个查询都不同），我们可以将这些标题放入队列中。在等待最多 100 毫秒或当前批次大于批次大小时，预测运行。在 GPU 给出结果后，我们循环遍历结果，并且由于顺序与队列中的标头相同，我们可以根据每个标头通过套接字发回响应（可以从我们上面保存的 dict 中查找） ）。

由于每个查询都带有不同的标头，因此您不能错过请求、响应和套接字。虽然您可以在运行多个工作程序的同时运行使用 GPU 卡的 tensorflow 服务，以保持批处理大小足够大以获得更大的吞吐量。

我也在这个 repo 中找到了一个批处理机制： https ://github.com/richemslie/galvanise_zero