Открытая реализация AlphaZero для настольной игры Quoridor на чистом Python обучалась только на процессоре. Одна итерация обучения занимала 783 секунды, полный план в 20 тысяч итераций упирался в месяцы работы. Требовалось ускорить самоигру в 5-20 раз и не тронуть качество игры: обученные веса и результат поиска должны остаться прежними, иначе теряется весь накопленный прогресс обучения.
Профилирование cProfile по одной итерации показало три узких места: сеть считается по одной позиции на ход, проверка достижимости финиша при постановке стены идёт обходом BFS по графу из 81 клетки, и на каждый вызов predict тратится переключение режима сети.
Этап 1, видеокарта. Обучение переведено на CUDA, воркеры самоигры остались на процессоре, устройство выбирается автоматически и задаётся флагом. Смешанный режим здесь важен: 30 процессов самоигры на ядрах обгоняют пачку партий на одной карте, это подтверждено замером на арендованной RTX 4090.
Этап 2, батчевая самоигра. Поиск переписан на генераторы, партии играются одновременно, и все запросы к сети от разных партий собираются в один вызов. В фазе sequential halving симуляции разных кандидатов идут по непересекающимся поддеревьям, поэтому батчинг не меняет результат поиска ни на один ход. Тем же приёмом переписана арена гейтинга, где две сети играют матч.
Этап 3, стены. BFS заменён на битборд-заливку 64-битными масками с предварительным отсечением по барьеру, из predict вынесен eval, маска допустимых ходов переписана на numpy.
Каждый этап отдельным коммитом со своим бенчмарком до и после, замеры на одном железе и одинаковых сидах, длинные прогоны под сторожевым процессом, который перезапускает зависший запуск.
Самоигра на процесс: было 1.36 хода в секунду, стало 12.16, ускорение 8.9 раза. Полная итерация стандартного пресета: было 783 секунды, стало 690 уже на ноутбуке, где батчи не раскрываются на 12 воркерах, а на сервере с 30 процессами выигрыш кратно больше.
Тестов 23 из 23, включая тест на побитовое совпадение нового генераторного пути поиска со старым и тест на учёт всех партий в батчевой арене: качество игры не изменилось, ускорилось только его вычисление.
Работа отдана пулреквестом из 8 коммитов в исходный репозиторий проекта, история покоммитная, каждый этап читается и откатывается отдельно.
Отдельно собран инструмент очной встречи двух чекпоинтов: две модели играют серию партий с чередованием цветов и жадными ходами, как в гейтинге, и владелец проекта видит цифрами, какая из версий сильнее.
