Масштабируемые MoE‑модели: как Perplexity решает проблемы AWS EFA

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

Архитектура MoE и её траектория коммуникации

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

Почему традиционный EFA был проблемой

Внутри одного узла связь NVLink (≈900 ГБ/с) почти мгновенна, но при межузловой работе InfiniBand повышает задержку до десятков микросекунд. AWS EFA, используемая в кластерах, не поддерживает GPUDirect Async, вынуждая GPU инициировать передачу через CPU‑прокси. Это приводит к потере драгоценных микросекунд и делает MoE‑коммуникации дорогими.

Гибридная CPU‑GPU архитектура и конфликт «прокси‑проблемы»

Ключевая идея Perplexity ─ разделить работу ядра в две части: sender (передача токенов) и receiver (приём). Эти части синхронизируются через TransferEngine, а промежуточный слой CPU‑прокси обрабатывает сетевые вызовы.

• Dispatch‑ядро формирует пакеты токенов и передаёт их в прокси.
• Combine‑ядро собирает результаты и вычисляет среднее значение.
• Объединение вычислений и коммуникаций значительно снижает суммарную задержку.

Оптимизация буферов и TransferEngine

Сначала ядра обмениваются только счётчиками токенов, позволяя каждому узлу заранее вычислить смещения в буфере. После этого происходит двухэтапная передача: небольшая «первичная» часть и большой RDMA‑write. TransferEngine получил новые функции scatter и barrier, а также пред‑регистрируемые WR‑шаблоны, которые ускоряют передачу мелких пакетов.

Оценка производительности

Показатели, полученные на H200 с ConnectX‑7 и EFA:

• Общая задержка (dispatch + combine): 459 µs для 16 экспертов, 582 µs для 32 экспертов и 692 µs для 64 экспертов.
• На ConnectX‑7 решение превзошло в среднем DeepEP и традиционный IBGDA благодаря быстрому combine‑ядру.
• Разрыв между EFA и ConnectX‑7 оказался существенно ниже, чем ожидалось.

Пользовательские кейсы и масштабирование

В тестах на DeepSeek‑V3 (671 млрд параметров) многопоточное развертывание показало пропускную способность, сравнимую с оптимизированной одно‑узловой конфигурацией. Для Kimi‑K2 (1 трис) кастомные ядра впервые сделали возможным запуск модели на AWS с «жизнеспособной» задержкой.

Будущие направления развития

Perplexity продолжает работу с инженерами AWS над улучшением libfabric. В планах — экспериментировать с EFA‑Direct, чтобы ещё более сократить задержки, а также исследовать варианты использования более агрессивных архитектур NVLink/PCIe для внутренних соединений.