LRMoE урезает обратный проход MoE и ускоряет обучение на 21–45%

Команда авторов из академической среды предложила LRMoE — систему для обучения MoE-моделей на узлах, где памяти GPU не хватает. MoE (Mixture-of-Experts, «смесь экспертов») — архитектура, в которой слой состоит из десятков или сотен небольших подсетей-экспертов, а для каждого токена маршрутизатор включает только несколько из них. Один такой выбор «токен — эксперт» авторы называют парой. Главная цифра работы: на моделях Qwen3-30B-A3B и GLM4.5-106B-A12B накопленное время обучения MoE-слоёв до заранее заданного значения функции потерь (loss) сократилось на 21,2–44,8% по сравнению с SYMI. Это предыдущая система, которая разделяет хранение параметров модели и состояний оптимизатора, чтобы гибче размещать экспертов.

Почему перекос нагрузки не лечится копиями экспертов

Эксперты разложены по GPU, и токены летят к нужным картам через all-to-all — обмен «каждый с каждым». Популярные эксперты получают больше токенов, их GPU становятся узким местом, остальные ждут. Синхронный шаг идёт со скоростью самого загруженного участника.

Обычные ответы на это — переставить экспертов или сделать копии самых популярных. Так работают FasterMoE, FlexMoE, SmartMoE, LAER-MoE, UltraEP и TopoEP. Но у копий есть цена:

• реплика занимает память, которая нужна под буферы токенов и сохранённые активации;
• копии надо синхронизировать, а градиенты суммировать;
• общий объём работы не уменьшается, он лишь переезжает на другие карты.

Авторы показывают, что при жёстком лимите на GPU даже идеально ровная раскладка может оставить суммарную нагрузку выше того, что память и сеть способны обслужить. Неравномерность и перегрузка в целом — разные проблемы, и выравнивание решает только первую.

Что делает LRMoE: полный прямой проход, обрезанный обратный

Идея в том, чтобы считать нагрузку обратного прохода не неизбежным следствием маршрутизации, а регулируемой величиной. Работа идёт в два этапа.

• Прямой проход остаётся полным. Все пары токен–эксперт считаются как обычно, потому что на этом этапе градиентов ещё нет и оценить ценность пар нечем.
• Квоты до вычислений. Контроллер видит, сколько токенов пришло к каждому эксперту, и с учётом памяти и сети решает, сколько пар каждого эксперта оставить в обратном проходе. Он же планирует размещение реплик. Квота фиксируется на весь вызов слоя.
• Отбор по метрике LRM. Когда градиенты доходят до слоя, для каждой пары оценивается вклад в обновление весов эксперта. Вклад пары — внешнее произведение входящего градиента и сохранённого входа проекции. Поэтому квадрат его нормы раскладывается на произведение трёх множителей: квадрата веса маршрутизации, квадрата нормы градиента и квадрата нормы входа. Полные матрицы градиентов для этого считать не нужно. В пределах квоты остаются пары с наибольшим счётом.
• Перераспределение. Оставшиеся пары целочисленным алгоритмом раскладываются по допустимым копиям своего эксперта. Цель — минимизировать пиковую нагрузку на одну GPU. Новых реплик в обратном проходе не появляется.
• Физическое уплотнение. Оставшиеся строки градиентов и активаций собираются в плотные блоки по получателю и эксперту. Обмен all-to-all становится компактнее, а умножения матриц идут как grouped GEMM переменного размера без дополнения нулями. Результат затем возвращается в исходную раскладку.

Важная деталь: просто замаскировать лишние строки в полноразмерных тензорах ничего не даст, ни времени, ни памяти. Выигрыш появляется только когда уменьшаются реальные размеры пересылок и матриц.

Что показали эксперименты и где подвох

Помимо основной цифры в 21,2–44,8%, авторы отдельно прогнали воспроизведение на GPU и получили снижение на 57,8% суммарной памяти под размещение экспертов, реплики и буферы токенов по сравнению с фиксированным владением экспертов. Это замер памяти, а не времени обучения.

Подвох в том, что обратный проход становится приближённым: часть вкладов в градиент просто выбрасывается. Метрика LRM — оценка на уровне одной проекции, и гарантий сходимости она не даёт. Поэтому авторы сравнивают системы по времени до фиксированного loss, а не по скорости шага: ускорение засчитывается, только если модель дошла до того же качества. Накладные расходы на подсчёт счётов, упаковку и маршрутизацию включены в измерения.

Для инженеров вывод такой. Если обучение MoE упирается в память GPU и сеть, а добавить реплик уже некуда, имеет смысл смотреть на сокращение самой работы обратного прохода, а не только на её перекладывание. Оговорки тоже важны. Результаты получены на двух моделях и сравнены с одной базовой системой, поэтому перенос на другие архитектуры и длинные запуски надо проверять отдельно. Цена приблизительных градиентов может проявиться на более жёстких целях по качеству.