如何在 GPU 上有效并行化 AlphaZero?

How do I effectively parallelize AlphaZero on the GPU?

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

算法的关键是 Monte Carlo 树搜索状态 space (CPU) 与 eval 中神经网络的 'intuition'(概率)交错模式(GPU)。然后使用 MCTS 结果来训练神经网络。

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

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

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

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

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

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

预测服务可以读取每个查询 body 及其 header(每个查询都不同),我们可以将 header 放在 queue 中.在等待最多 100 毫秒或当前批次大于批次大小时,预测 运行s。 GPU给出结果后,我们将结果循环一遍,由于顺序与queue中的header相同,我们可以根据每个[=20]通过套接字发送回响应=](可以从我们上面保留的字典中查找)。

由于每个查询都带有不同的header,您不能错过请求、响应和套接字。虽然您可以 运行 使用 GPU 卡服务的张量流,同时 运行 多个工作人员以保持批处理大小足够大以获得更大的吞吐量。

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