我正在实现一个 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 在各自论文中如何并行化的信息。我能够在网上找到有限的信息,但这导致我改进了第一种方法。
如果有任何建议,我将不胜感激,特别是如果我错过了某些观点或方法。