Moonshot AI открыла MoonEP: библиотека для идеального баланса экспертов в MoE-обучении

Moonshot AI представила MoonEP — библиотеку с открытым исходным кодом (MIT) для экспертного параллелизма (EP) в распределённых моделях Mixture-of-Experts. Релиз приурочен к Kimi K3 Open Day, где также выложили веса модели K3, технический отчёт и ещё два инструмента: FlashKDA и AgentEnv. MoonEP — одна из ключевых технологий, позволившая достичь 2.5-кратного улучшения эффективности масштабирования при обучении Kimi K3 — MoE-модели на 2,8 трлн параметров с нативным зрением и окном контекста до 1 млн токенов.

Проблема, которую решает MoonEP, — дисбаланс токенов между экспертами при экспертом параллелизме. Роутер отправляет каждый токен топ-K экспертам, которые живут на разных рангах, но распределение почти всегда неравномерно. Из-за этого задержка коллектива определяется самым загруженным рангом, а динамические активации фрагментируют память GPU и требуют синхронизации на каждом слое.

Решение MoonEP — жёсткая гарантия, что каждый ранг получает ровно S × K токенов, независимо от перекоса роутинга. Это достигается за счёт планирования небольшого числа дублирующихся экспертов в реальном времени на основе текущих выходов роутера. Копии экспертов предзагружаются до вычисления, а в обратном проходе их градиенты возвращаются на домашние ранги. Ключевые свойства: идеальный баланс, онлайн-планирование с минимальными накладными расходами (реализовано на CUTLASS CuTe DSL) и статические формы буферов без копирования и без синхронизации хоста на каждом слое.

Библиотека уже доступна на GitHub. Для её сборки требуется nvidia-cutlass-dsl==4.4.2. Это не просто инкрементальный патч, а архитектурное решение, которое может стать стандартом для обучения крупных MoE-моделей.