Укрощение «триллионников»: как Perplexity запустила гигантские ИИ-модели на AWS EFA

Компания Perplexity, использующая в своих продуктах и исследованиях передовые ИИ-модели, столкнулась с известной в индустрии проблемой. Крупные открытые модели класса MoE (Mixture-of-Experts), такие как Kimi-K2, обладают гигантским размером.

Они настолько велики, что их невозможно эффективно разместить на одном стандартном узле, даже на мощных конфигурациях с 8-ю GPU NVIDIA H200. Это вынуждает инженеров использовать многоузловые (multi-node) развертывания, что несет в себе отдельные технические вызовы.

В ответ на эту проблему Perplexity представила новый набор ядер (kernels) для «экспертного параллелизма». Эти ядра не только достигают рекордно низких задержек на оборудовании ConnectX-7, превосходя по производительности DeepEP, но и, что более важно, впервые обеспечивают жизнеспособную производительность на AWS Elastic Fabric Adapter (EFA). По сути, эта разработка открывает возможность для развертывания моделей с триллионом параметров в облачной среде AWS.

Компания опубликовала свои ядра на GitHub, а подробное исследование на arXiv.

MoE: суть архитектуры и сложность коммуникации

Архитектура MoE (Mixture-of-Experts) де-факто стала стандартом для масштабирования моделей до сотен миллиардов или даже триллионов параметров, позволяя сохранить разумную задержку при инференсе (inference time).

Принцип ее работы в том, что вместо одного плотного слоя в трансформере используется набор «экспертов» (отдельных небольших сетей) и специальный «маршрутизатор» (routing layer). Этот маршрутизатор решает, какому именно эксперту отправить каждый отдельный токен. Такая задача отлично распараллеливается, поскольку разные эксперты могут физически находиться на разных GPU, в том числе на разных серверах.

Однако в этом и заключается сложность. В отличие от стандартных видов параллелизма (Tensor Parallelism или Data Parallelism), которые легко реализуются через библиотеки вроде torch.distributed и NCCL, маршрутизация MoE требует «рваных» (sparse) peer-to-peer коммуникаций.

Для этого нужны специализированные ядра:

1. Dispatch-ядро (отправка) - чтобы «разослать» токены по нужным узлам-экспертам.
2. Combine-ядро (сборка) - чтобы «собрать» обработанные результаты обратно и вычислить их средневзвешенное значение.

Эти операции не укладываются в стандартные коллективные вызовы и требуют кастомных низкоуровневых решений для минимизации задержек.

Проблема EFA: почему AWS было сложно

Пока вычисления проходят внутри одного сервера, высокоскоростная шина NVLink (900 ГБ/с) обеспечивает почти мгновенную связь.

Проблемы начинаются, когда модели не помещаются в память одной ноды. В Perplexity отмечают, что даже узел p5en с 8-ю H200 (1120 ГБ HBM) не может вместить триллионную модель, так как память делится между весами модели и KV-кэшем. Это делает меж-узловые (inter-node) развертывания неизбежными.

При выходе за пределы сервера используется InfiniBand (до 400 Гбит/с), который, несмотря на скорость, добавляет десятки и сотни микросекунд задержки к маршрутизации MoE.

На кластерах AWS ситуация исторически была еще сложнее из-за использования проприетарных адаптеров Elastic Fabric Adapters (EFA).

• Во-первых, на размерах пакетов, типичных для MoE, они уступают ConnectX-7.
• Во-вторых (и это главное), EFA не поддерживает GPUDirect Async.

Это означает, что GPU не может напрямую инициировать передачу данных. Вместо этого ему нужен «посредник» - поток на CPU (proxy thread), который связывает GPU и сетевой адаптер (NIC). Эта дополнительная работа и транзакции по шине PCIe «съедают» драгоценные микросекунды.

Предыдущие реализации, например, на базе NVSHMEM, на EFA показывали задержки выше миллисекунды, что делало их неэффективными с точки зрения затрат и производительности.

Гибридная архитектура: как обошли «прокси-проблему»

Новые ядра Perplexity используют гибридную CPU-GPU архитектуру. GPU-ядро взаимодействует с моделью, в то время как выделенный поток-прокси на CPU обрабатывает взаимодействие с сетевым адаптером. В решении используется TransferEngine, ранее созданный компанией для передачи KV-кэша.

Ключевая идея - разделение ядер (dispatch и combine) на две половины: sender (отправитель) и receiver (получатель).

Это позволяет реализовать перекрытие вычислений и коммуникаций (computation-communication overlapping). Отправитель (sender) подготавливает токены и передает их CPU-прокси для отправки по сети. Пока данные «летят» по сети, GPU не простаивает, а может выполнять другую работу, например, вычисления общего (shared) эксперта. Получатель (receiver) в это время ожидает сигнала о завершении всех сетевых передач.

Для связи между GPU и CPU используются как Unified Memory (для объемных передач), так и GDRCopy (для опроса с низкой задержкой). При этом внутри одного узла (intra-node) ядра активно используют NVLink для прямого обмена токенами, что дополнительно разгружает сеть (до 1/8 трафика для EP64).

Эффективный обмен: буферы и TransferEngine

Серьезной задачей была организация буферов. Если каждый ранг будет ждать полной информации о маршрутах, возникнет простой. Если использовать приватные буферы для каждого, не хватит памяти.

Решение нашли в двухэтапной передаче:

1. Сначала ядра обмениваются только счетчиками токенов для каждого эксперта. Это позволяет всем рангам заранее рассчитать точные смещения (offsets) в буферах.
2. Чтобы не терять время, вместе с этой маршрутной информацией отправляется первая порция токенов, которая помещается в небольшой зарезервированный буфер.
3. После обмена маршрутами оставшиеся токены передаются одним большим пакетом (single RDMA write) в уже известные ячейки памяти.

TransferEngine также пришлось доработать. Изначально он был оптимизирован для других задач. Для MoE, где нужно отправлять много мелких пакетов всем узлам, были добавлены две операции: scatter (рассылка) и barrier (сигнализация). Также были внедрены оптимизации, такие как WR templating (пред-регистрация запросов).

Интересно, что на EFA, где используются две NIC по 200 Гбит/с, инженеры Perplexity решили шардировать (распределять) не байты каждого пакета, а группы узлов-получателей (peer groups) между двумя картами.

Оценка: что получилось в итоге

Тестирование проводилось на H200 с адаптерами ConnectX-7 и EFA.

• Задержки (Latency): Ядра показали комбинированную задержку (dispatch + combine) 459 мкс для EP16, 582 мкс для EP32 и 692 мкс для EP64. Это значительно превосходит другие EFA-решения, такие как UCCL-EP.
• Сравнение с ConnectX-7: На ConnectX-7 новое решение, даже используя прокси-поток на CPU, обошло по скорости реализации без прокси (IBGDA), включая DeepEP, в основном за счет более быстрого combine-ядра. Это доказывает эффективность специализированного прокси.
• EFA vs ConnectX-7: Разрыв в производительности между EFA и ConnectX-7 в задачах MoE оказался «менее существенным», чем можно было ожидать.

Пропускная способность (End-to-End Throughput):

Самое главное - как это работает в реальной модели.

• DeepSeek-V3 (671B): Многоузловое развертывание показало производительность на уровне или даже выше, чем оптимизированная одно-узловая конфигурация, доказывая эффективность масштабирования.
• Kimi-K2 (1T): Для этой модели, которая физически не помещается на одном узле, новые ядра вообще сделали возможным ее запуск. Модель теперь работает на EFA с «жизнеспособной задержкой» (viable latency).

Будущая работа

В Perplexity отмечают, что продолжают тесное сотрудничество с инженерами AWS для улучшения libfabric, а также планируют экспериментировать с efa-direct для дальнейшего снижения задержек.