MoonEP: Как Moonshot AI решил проблему дисбаланса в MoE-моделях

В мире больших языковых моделей (LLM) и, в частности, архитектур Mixture-of-Experts (MoE), одной из самых сложных и критичных проблем остается коммуникация между узлами. Когда модель весит триллионы параметров, как, например, Kimi K3 от Moonshot AI, каждый бит данных на шине NVLink или InfiniBand на счету. Компания Moonshot AI недавно сделала важный шаг для всего сообщества разработчиков, открыв исходный код библиотеки MoonEP. Это не просто еще один инструмент оптимизации, а фундаментальное решение проблемы дисбаланса нагрузки при экспертном параллелизме (Expert Parallelism, EP).

Выпуск MoonEP стал частью инициативы Kimi K3 Open Day, где вместе с весами модели Kimi K3 и техническим отчетом были опубликованы три ключевых инфраструктурных проекта: MoonEP, FlashKDA и AgentEnv. В то время как FlashKDA уже был доступен ранее, MoonEP и AgentEnv стали новинками. MoonEP позиционируется как библиотека, которая делает коммуникацию между экспертами более эффективной на масштабе, и, что важно, она распространяется под лицензией MIT, что открывает широкие возможности для интеграции в коммерческие и исследовательские проекты.

Эта библиотека является одним из ключевых инновационных компонентов, обеспечивающих заявленное улучшение масштабируемости на 2,5 раза для Kimi K3 — модели с 2,8 триллионами параметров, нативной поддержкой зрения и контекстным окном в 1 миллион токенов. Но как именно MoonEP достигает таких результатов, и почему существующие решения, такие как DeepEP, не справляются с этой задачей? Давайте разберемся подробно.

01Проблема дисбаланса в экспертном параллелизме

Чтобы понять ценность MoonEP, нужно сначала разобраться в механике экспертного параллелизма. В архитектуре MoE каждый токен входных данных маршрутизируется к «экспертам» — отдельным слоям нейронной сети, которые специализируются на определенных типах данных. Роутер (router) выбирает топ-K наиболее подходящих экспертов для каждого токена. Эти эксперты распределены по разным рангам (узлам) в кластере.

Ключевая проблема заключается в том, что роутеры редко работают идеально сбалансировано. В реальных условиях некоторые эксперты получают значительно больше токенов, чем другие. Это создает так называемый «перекос» (skew). В библиотеке MoonEP этот перекос количественно оценивается метрикой maxvio, которая определяется формулой:

maxvio = max_e (T_e / T̄) - 1

Где T_e — количество токенов, направленных к эксперту e, а T̄ — ожидаемое количество при идеальном балансе. Значение maxvio равное 0 означает идеальный баланс. Однако в реальности это число часто значительно превышает 1, 10 или даже 20.

Последствия этого дисбаланса носят структурный, а не случайный характер. Задержка коллективной операции (collective) определяется самым медленным участником. В контексте MoE это означает, что самый «горячий» ранг (тот, который получил больше всего токенов) определяет время выполнения всей итерации обучения. Кроме того, количество токенов на ранге меняется на каждом шаге. Эти динамические формы активации фрагментируют память GPU и требуют синхронизации на стороне хоста (CPU) для каждого слоя, что создает дополнительные узкие места.

02Основная идея: динамические избыточные эксперты

Главное нововведение MoonEP заключается в жестком инварианте, который библиотека навязывает процессу коммуникации. Независимо от того, насколько сильно перекосилась маршрутизация, каждый ранг получает ровно S × K токенов. Здесь S — количество входных токенов на ранг, а K — количество маршрутизируемых топ-K экспертов для каждого токена.

Как достигается этот баланс? MoonEP использует концепцию динамических избыточных экспертов. Библиотека планирует небольшое количество избыточных экспертов онлайн, непосредственно на основе текущих выходов роутера. Эти дублированные эксперты предварительно загружаются (prefetch) до начала вычислений экспертов. В обратном проходе (backward pass) их градиенты уменьшаются (reduced) обратно к их «домашним» рангам.

Этот дизайн опирается на три ключевых свойства:

1. Идеальный баланс: Гарантия того, что каждый ранг получит ровно S × K токенов, достигается за счет онлайн-планирования избыточных экспертов.
2. Онлайн-планирование: Используется ядро планирования на GPU, близкое к оптимальному, с пренебрежимо малыми накладными расходами. Оно реализовано с использованием DSL CUTLASS CuTe. В файле setup.py жестко зафиксирована версия nvidia-cutlass-dsl==4.4.2, что обеспечивает стабильность и предсказуемость работы.
3. Zero copy и статические формы: Реализованы слитые операции permute/unpermute. Токены записываются напрямую в позиции, сгруппированные по экспертам, на удаленных рангах, а вычисления возвращают представления буферов. Требуется только фиксированный буфер размером S × K. Статически известные формы устраняют необходимость в синхронизации MoE на стороне хоста для каждого слоя.

03Контракт памяти: как MoonEP взаимодействует с фреймворком

MoonEP устанавливает специфический контракт с фреймворками обучения или инференса. Для работы требуется:

• Один непрерывный тензор весов симметричной памяти для каждого проекционного эксперта.
• Массив cu_seqlens, сгенерированный планировщиком.

Групповое умножение матриц (GEMM) потребляет единый тензор весов размером [E+B, H, H'], где:

• E — общее количество маршрутизируемых экспертов,
• B — количество слотов предварительной выборки (prefetch) на ранг,
• H — скрытый размер (hidden size),
• H' — промежуточный размер экспертного FFN.

Массив cu_seqlens[E+B], возвращаемый функцией dispatch, указывает, какие строки эксперта активны. Непрерывность (contiguity) является жестким требованием, поскольку групповое GEMM обращается к экспертам исключительно по индексу строки. Структура данных разбивается следующим образом:

• Строки [0, E) содержат локальных экспертов всех рангов, по E/R строк на ранг. Каждый чанк физически является памятью параметров домашнего ранга, отображенной везде через симметричную память.
• Строки [E, E+B) — это локальные слоты предварительной выборки, заполняемые функцией buffer.prefetch_weight.

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

04Настройка параметров для обучения и инференса

Параметр B (количество слотов предварительной выборки) зависит от рабочей нагрузки:

• Обучение (Training): Должно использовать B = E/R. Планировщик дублирует экспертов максимум из одной удаленной домашней группы на ранг. Эта граница гарантирует, что каждый эксперт, к которому обращается групповое GEMM, является локальным.
• Инференс (Inference): Позволяет использовать B < E/R. README рекомендует B = 3–4. Если рангу требуется больше различных удаленных экспертов, чем B, групповое GEMM читает переполненные веса напрямую с домашнего ранга через симметричное отображение. Это немного медленнее, но не влияет на корректность вычислений.

В обратном проходе (backward pass) при обучении структура весов в fp32 зеркально отражается в буфере градиентов размером [E+B, H, H'] для каждой проекции. Критически важно, что строки [E, E+B)Backinged (поддерживаются) отдельным буфером уменьшения (reduce buffer), а не градиентами параметров. Градиенты дублированных экспертов являются временными и должны оставаться невидимыми для собственного reduce-механизма фреймворка. Каждый ранг отображает все буферы уменьшения как единое представление [R, B, H, H'], затем reduce_grad читает слоты собственных экспертов с каждого ранга по NVLink, накапливает их в локальный градиент параметров и обнуляет потребленные слоты.

05Бенчмарки: MoonEP против DeepEP v2

Опубликованные бенчмарки сравнивают MoonEP с DeepEP v2. Оба теста проводились на GPU NVIDIA H20 с экспертным параллелизмом (EP) равным 8, с изменением дисбаланса роутера. Скрипт сравнения benchmarks/bench_vs_deepep.py использует следующие параметры по умолчанию:

• S=8192
• E=384
• H=7168
• K=8
• H'=2048
• 32 SM (Streaming Multiprocessors)
• Целевые значения maxvio: 0.2, 1, 10 и 20.

Обе библиотеки получают идентичную матрицу маршрутизации из общего семени. Результаты показали три ключевых вывода:

1. Zero copy ускоряет коммуникацию: Отсутствие копирования из буфера коммуникации в буфер пользователя, которое доминирует в эпилоге, делает сырую коммуникацию MoonEP быстрее. Время коммуникации MoonEP стабильно ниже, чем у DeepEP v2, на всех уровнях дисбаланса.
2. Идеальный баланс защищает от перекоса: Время коммуникации MoonEP остается почти плоским (неизменным) по мере роста maxvio, в то время как DeepEP v2, чья задержка определяется самым горячим рангом, постепенно деградирует.
3. Стабильность памяти: MoonEP практически невосприимчив к дисбалансу, что делает его предсказуемым выбором для долгосрочного обучения.

06Что это значит на практике

Открытие MoonEP под лицензией MIT знаменует собой важный момент для экосистемы открытого ИИ. Для исследователей и инженеров, работающих с MoE-моделями, это означает возможность достичь более высокой эффективности использования GPU без необходимости писать сложные кастомные ядра коммуникации с нуля. Особенно актуально это для пользователей из РФ, где доступ к некоторым облачным сервисам или закрытым библиотекам может быть ограничен, а локальный запуск на имеющемся оборудовании (включая NVIDIA H20, которые активно поставляются в Россию) требует максимальной оптимизации.

Интеграция MoonEP позволяет:

• Снизить задержку при обучении крупных MoE-моделей за счет устранения узких мест в коммуникации.
• Упростить управление памятью благодаря статическим формам буферов и глобальному пулу для prefetch.
• Использовать проверенные бенчмарки для оценки производительности своих кластеров.

В сочетании с другими открытыми проектами Moonshot, такими как FlashKDA и AgentEnv, MoonEP формирует полный стек инструментов для работы с современными агентными и мультимодальными моделями. Если вы планируете обучать или развертывать модели типа Kimi K3 или аналогичные MoE-архитектуры, изучение исходного кода MoonEP на GitHub является обязательным шагом для понимания передовых практик в области распределенного обучения.

Исходный код MoonEP доступен по адресу MoonshotAI/MoonEP на GitHub. Официальное объявление также можно найти в аккаунте @Kimi_Moonshot.

Источник: MarkTechPost ↗