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

Question

我正在实现AlphaZero的一个版本（AlphaGo的最新版本）以应用于其他领域。你知道吗

该算法的关键是对状态空间（CPU）的蒙特卡罗树搜索，它与来自评估模式（GPU）的神经网络的“直觉”（概率）交织在一起。然后利用MCTS结果训练神经网络。你知道吗

我已经通过启动多个进程来并行化CPU执行，每个进程建立自己的树。这是有效的，现在已经导致了GPU瓶颈！（nvidia smi始终100%显示GPU）

我设计了两种策略来并行GPU评估，但是它们都有问题。你知道吗

• 每个进程只对自己树中的批进行网络评估。在我最初的天真实现中，这意味着批大小为1。但是，通过重构一些代码并添加一个“虚拟丢失”来阻止（但不是完全阻止）同一个节点被选中两次，我们可以得到大小为1-4的更大批。这里的问题是，我们不能允许大的延迟，直到我们评估批或准确性受到影响，所以小批量是关键。

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

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

我几乎没有发现AlphaZero和AlphaGo-Zero在各自的论文中是如何并行的。我能够在网上找到有限的信息，但是这使我改进了第一种方法。你知道吗

我会很感激任何关于这方面的建议，特别是如果有一些要点或方法我错过了。你知道吗