Moonshot AI открыла MoonEP: обучение Mixture of Experts без простоя GPU в 2,5 раза эффективнее

Библиотека вышла под лицензией MIT (полностью свободная для коммерческого и личного использования) вместе с моделью Kimi K3 и ещё двумя кодовыми базами: FlashKDA и AgentEnv. По данным репозитория Moonshot AI на GitHub, MoonEP стала одной из технологий, обеспечивших 2,5-кратное улучшение эффективности масштабирования Kimi K3, MoE-модели с 2,8 триллиона параметров, встроенным зрением и контекстным окном в миллион токенов (токен, минимальная единица текста, которую обрабатывает модель, примерно три четверти слова).

Что понадобится

• Кластер с GPU NVIDIA (бенчмарки опубликованы для H20, 8 узлов с expert parallelism, EP = 8)
• Python 3.x, PyTorch, пакет nvidia-cutlass-dsl==4.4.2 (устанавливается автоматически через setup.py)
• Доступ к репозиторию MoonshotAI/MoonEP на GitHub (MIT-лицензия, клонирование бесплатно)
• Базовое понимание распределённого обучения и архитектуры MoE
• Время на интеграцию: от нескольких часов на тестовый запуск до нескольких дней на встраивание в рабочий пайплайн

Какую проблему решает MoonEP и как именно?

При обучении MoE-моделей маршрутизатор (router) отправляет каждый токен к нескольким лучшим экспертам, но эксперты живут на разных GPU. Маршрутизатор редко распределяет токены равномерно: одни эксперты получают намного больше работы, другие простаивают.

Репозиторий MoonEP измеряет перекос метрикой maxvio: если maxvio = 0, баланс идеальный, если 10 или 20, один эксперт перегружен в разы. Проблема структурная: скорость всего шага определяется самым медленным узлом. Вдобавок объём данных на каждом узле меняется от шага к шагу, что фрагментирует память GPU и вынуждает синхронизироваться с процессором на каждом слое.

Пошаговая инструкция

1. Клонируйте репозиторий и установите зависимости.

git clone https://github.com/MoonshotAI/MoonEP.git
cd MoonEP
pip install -e .

Файл setup.py автоматически подтянет nvidia-cutlass-dsl==4.4.2.

1. Разберитесь с ключевым контрактом библиотеки. MoonEP гарантирует жёсткий инвариант: каждый узел (rank) получает ровно S × K токенов, где S означает количество входных токенов на узел, а K означает число экспертов, к которым маршрутизатор отправляет каждый токен. Перекос маршрутизатора больше не влияет на распределение.

2. Настройте размер пула предзагрузки (параметр B). Для обучения задайте B = E/R, где E означает общее число экспертов, а R означает количество узлов. Это жёсткое требование: планировщик дублирует экспертов максимум из одной удалённой группы на узел, и такой размер B гарантирует, что каждый эксперт в групповом GEMM окажется локальным.

# Пример: 384 эксперта, 8 узлов
B = E // R  # B = 384 // 8 = 48 для обучения

Для инференса (инференс, этап, когда модель уже обучена и просто выдаёт ответы) допустимо B = 3–4. Если узлу потребуется больше удалённых экспертов, чем помещается в B слотов, библиотека прочитает веса напрямую с «домашнего» узла через симметричную память. Это чуть медленнее, но на корректность не влияет.

1. Подготовьте тензоры весов. MoonEP требует один непрерывный тензор весов формы [E+B, H, H'] на каждую проекцию эксперта, где H означает размер скрытого слоя, H' означает промежуточный размер FFN эксперта. Непрерывность обязательна: групповой GEMM адресует экспертов чисто по индексу строки.

2. Вызовите dispatch для маршрутизации. Функция dispatch возвращает cu_seqlens[E+B], массив указателей, которые определяют, какие строки экспертов активны. Токены записываются напрямую в свои позиции на удалённых узлах без промежуточного копирования (zero copy).

3. Запустите бенчмарк для проверки.

cd benchmarks
python bench_vs_deepep.py

По умолчанию скрипт использует S=8192, E=384, H=7168, K=8, H'=2048, 32 потоковых мультипроцессора (SM) и тестирует maxvio на уровнях 0.2, 1, 10 и 20.

1. Для обучения настройте буфер градиентов. Библиотека зеркально дублирует тензор весов буфером градиентов формы [E+B, H, H'] в fp32. Градиенты дублированных экспертов хранятся отдельно от параметрических градиентов: функция reduce_grad считывает слоты своих экспертов со всех узлов через NVLink, накапливает их в локальный градиент и обнуляет использованные слоты.

Что это даёт вам на практике?

Разработчику MoE-моделей в России. MoonEP решает проблему, из-за которой mixture of experts обучение на больших кластерах теряет эффективность: динамический дисбаланс нагрузки. Библиотека открыта под MIT, можно встраивать в коммерческий пайплайн. Если вы уже используете DeepEP, скрипт bench_vs_deepep.py позволит сравнить производительность на вашем железе напрямую.

Автору Дзена, который пишет про ИИ. Архитектура mixture of experts обучение и инференс делает доступнее, потому что активируется только часть параметров. Понимание того, как решается проблема баланса, даёт вам материал для объяснения читателям, почему модели вроде Mixtral или Kimi K3 работают быстрее, чем «плотные» модели с таким же числом параметров.

Предпринимателю. Пул предзагрузки (B экспертных весов) расходуется на весь пайплайн, а не на каждый слой отдельно. Это означает, что дополнительная память при масштабировании растёт предсказуемо, а не линейно с глубиной модели. Для планирования бюджета на GPU это конкретный аргумент.

Если вы обучаете MoE-модели на нескольких GPU и теряете время на дисбаланс экспертов, MoonEP стоит протестировать на вашей конфигурации: клонируете репозиторий, запускаете бенчмарк и за час получаете ответ, даёт ли библиотека выигрыш именно на вашем железе.