MoE-оптимизация: как снизить трафик в 2.6 раза без потери качества

Проблема: All-to-All как узкое горлышко

В архитектурах Mixture-of-Experts (MoE) каждый токен маршрутизируется к $k$ из $E$ экспертов. При экспертной параллелизации (Expert Parallelism, EP) эксперты распределены по GPU, что требует выполнения операций all-to-all collectives на каждом шаге обучения для отправки токенов и сбора результатов. На кластерах с 8 GPU AMD Instinct MI300X эти коммуникации занимают 45% времени шага при маршрутизации top-2 и до 60% при top-6. Это критически замедляет pretraining крупных моделей.

Решение: Корреляция и перемешивание

Авторы (Radha Gulhane, Quentin Anthony, Beren Millidge) обнаружили, что ранние этапы обучения формируют устойчивые корреляции в выборе экспертов. Токены часто выбирают одни и те же пары экспертов, а выбор в одном слое предсказывает выбор в следующем. На основе этого предложены два метода:

• Correlated Expert Placement: эксперты, часто выбираемые вместе, размещаются на одном GPU. Это позволяет диспетчеру отправлять каждый токен на каждый GPU только один раз, устраняя до 58% строк диспетчеризации.
• Token Shuffling: при sequence parallelism токен перемещается на GPU, где, по прогнозу, будут эксперты следующего слоя, во время операции reduce-scatter после attention-блока.

Результаты: Ускорение без компромиссов

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

Значение для индустрии

Результаты показывают, что оптимизация коммуникаций через анализ паттернов маршрутизации эффективнее, чем просто добавление мощности. Это позволяет масштабировать MoE-модели (такие как Mixtral или Grok) на тысячи GPU с меньшими затратами на сеть, что критично для следующего поколения LLM.

Источник: arXiv cs.CL ↗