Проблемы train-inference mismatch в MoE моделях при RL

Еще немного интересных наблюдений из RL для MoE. В нем есть одна проблема, которая делает обучение гораздо более склонным к нестабильностям. В любом современном RL есть train-inference mismatch, обычно вызванный двумя вещами: 1) из-за асинхронности инференс отстает от трейнера и 2) из-за разной имплементации кернелов/разного железа/неассоциативности сложения (у одной версии dense чекпоинта pi_train(tok) / pi_inference(tok) отличается обычно в четвертом знаке). Однако, в MoE появляется еще один очень неприятный источник такого расхождения: routed experts. Проблема в том, что во время инференса могут активироваться одни эксперты, а во время forward_pass на стороне трейнера – другие, и незначительное изменение в скорах роутера (но достаточное для того, чтобы topK изменился) могут значительно изменить итоговый результат. Основные решения такие: - Просто фризим роутер и уменьшаем асинк ⇒ роутер реже пересекает границу, при которой меняется topK - R2/R3, они же Routing Replays: сохраняем экспертов, активируемых версией модели для инференса и форсируем forward pass через них в трейнере. В статье https://arxiv.org/abs/2512.01374 можно почитать про разницу R2 vs R3 и всякие доп детали. Но проблема в том, что экспертов-то мы форсируем, но скоры роутера меняются, и это может быть проблемой (можно порассуждать на тему, что будет, если инференс версия дала i-ому эксперту высокий скор, а трейнер – низкий) И на самом деле у всех результаты смешанные: у кого-то работает одно, у других – другое. Вчера прочитал https://kiddyboots216.github.io/mismatch/#conclusion-router-replay и увидел еще один способ: Total Router Recall – сохраняем не только экспертов, но и их скоры, а в трейнере фризим роутер. Интересно, что можно предложить способ, при котором мы будем роутер учить, используя скоры и экспертов с инференса, правда работать это будет плохо. В общем, для всех интересующихся RL-ем рекомендую полистать блогпост, там много чего интересно помимо реплеев