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

Question

我正在实现一个 AlphaZero 版本(AlphaGo 的最新版本)以应用于其他领域。

该算法的关键是状态空间 (CPU) 的蒙特卡洛树搜索与来自评估模式 (GPU) 的神经网络的“直觉”(概率)交错。然后使用 MCTS 结果来训练神经网络。

我已经通过启动多个进程来并行化 CPU 执行，每个进程构建自己的树。这是有效的，现在已导致 GPU 瓶颈! (nvidia-smi 一直显示 GPU 为 100%)

我设计了 2 种策略来并行化 GPU 评估，但是它们都有问题。

• 每个进程仅对来自其自己的树的批处理评估网络。在我最初的天真实现中，这意味着批处理大小为 1。但是，通过重构一些代码并添加 ' virtual loss' 以阻止(但不是完全阻止)同一节点被选择两次，我们可以获得大小为 1-4 的更大的批处理。这里的问题是，在我们评估批处理或准确性受到影响之前，我们不能允许较大的延迟，因此这里的关键是小批量。

• 将批处理发送到中央“神经网络 worker ”线程，该线程组合并评估它们。这可以在 32 或更多的大批量中完成，因此可以使用 GPU非常有效。这里的问题是树 worker 发送 CUDA 张量“往返”，PyTorch 不支持。如果我先克隆它们，它是受支持的，但是所有不断的复制使这种方法比第一种方法慢。

我在想也许我没有看到一个聪明的批处理方案可以使第一种方法起作用。使用多个 GPU 也可以加快第一种方法的速度，但 PyTorch 本身并不支持我想要的那种并行性。也许将所有张量保留在 NN worker 中并且只发送 id 可以改进第二种方法，但是这里的困难在于如何有效地同步以获得大批量而不会使 CPU 线程等待太久。

我在各自的论文中几乎没有发现有关 AlphaZero 或 AlphaGo Zero 如何并行化的信息。然而，我能够在网上找到有限的信息，这让我改进了第一种方法。

如果有任何关于此的建议，我将不胜感激，特别是如果我遗漏了某些要点或方法。

Answer 1

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

预测服务可以读取每个查询主体及其 header (每个查询都不同)，我们可以将这些 header 放入队列中。在等待最多 100 毫秒或当前批处理大于批处理大小时，预测会运行。 GPU给出结果后，我们循环结果，因为顺序与队列中的头相同，我们可以根据每个头通过套接字发送回响应(可以从我们上面保存的字典中查找).

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

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