Я реализую версию AlphaZero (самое последнее воплощение AlphaGo) для применения к некоторому другому домену.
Суть алгоритма - поиск по дереву Монте-Карло с чередованием пространства состояний (CPU)с «интуицией» (вероятностями) из нейронной сети в режиме eval (GPU). Затем результат MCTS используется для обучения нейронной сети.
Я уже распараллелил выполнение ЦП, запустив несколько процессов, каждый из которых создает свое собственное дерево. Это эффективно и теперь привело к узкому месту GPU ! (nvidia-smi все время показывает GPU на 100%)
Я разработал 2 стратегии для распараллеливания оценок GPU, однако у них обеих есть проблемы.
Каждый процесс оценивает сеть только по пакетам из своего собственного дерева. В моей первоначальной наивной реализации это означало, что размер пакета равен 1. Тем не менее, путем рефакторинга некоторого кода и добавления «виртуальных потерь», чтобы препятствовать (но неполностью заблокировать) один и тот же узел, если его выбрать дважды, мы можем получить большие партии размером 1-4. Проблема здесь в том, что мы не можем допустить больших задержек, пока мы не оценим партию или точность не пострадает, поэтому здесь важен небольшой размер партии.
Отправка партий в центральную станцию "«Нейросетевой работник», который объединяет и оценивает их. Это может быть сделано в большой партии из 32 или более, поэтому графический процессор может использоваться очень эффективно. Проблема здесь в том, что работники дерева посылают тензоры CUDA «туда-обратно», что не поддерживается PyTorch. Это поддерживается, если я их сначала клонирую, но все это постоянное копирование делает этот подход медленнее, чем первый.
Я думал, может быть, умная схема пакетирования, которую я не вижу, можетзаставить первый подход работать. Использование нескольких графических процессоров также может ускорить первый подход, но тот тип параллелизма, который мне нужен, изначально не поддерживается PyTorch. Может быть, сохранение всех тензоров в NN-работнике и только отправка идентификаторов может улучшить второй подход, однако сложность здесь заключается в том, как эффективно синхронизировать, чтобы получить большой пакет, не заставляя потоки ЦП ждать слишком долго.
Я обнаружилрядом нет информации о том, как AlphaZero или AlphaGo Zero были распараллелены в соответствующих документах. Однако я смог найти ограниченную информацию в Интернете, что позволило мне улучшить первый подход.
Буду признателен за любые советы по этому вопросу, особенно если есть какой-то пункт или подход, который я пропустил.