Запуск моделей с триллионами параметров на AWS EFA

В Perplexity мы используем лучшие модели для нашего продукта, наших API и наших исследовательских команд. Крупные модели с открытым исходным кодом типа Mixture-of-Experts, такие как Kimi-K2, представляют особые сложности, поскольку крупнейшие узлы вывода с 8 графическими процессорами NVIDIA H200 не могут эффективно разместить их, что делает необходимым использование многоузлового развертывания. Мы представляем набор ядер для параллелизма экспертов, которые обеспечивают передовые показатели задержки на ConnectX-7, превосходя производительность DeepEP. Эти же ядра первыми достигли жизнеспособных показателей задержки на AWS Elastic Fabric Adapter (EFA), что позволяет осуществлять развертывание моделей с триллионами параметров.

Попробуйте наши ядра на GitHub и ознакомьтесь с полной научной статьей на arXiv.

Введение

Mixture-of-Experts (MoE) стала стандартной архитектурой для масштабирования до сотен миллиардов или триллионов параметров, обеспечивая при этом приемлемую задержку во время вывода. Модели MoE заменяют плотный слой трансформера набором экспертов и небольшим линейным слоем маршрутизации, который определяет, с какими экспертами нужно умножить токен. Благодаря группировке токенов, связанных с одним и тем же экспертом, задача становится разреженной и легко поддающейся распараллеливанию, так как различные эксперты теперь могут размещаться на разных графических процессорах, распределенных по нескольким узлам.

В отличие от других форм параллелизма, таких как тензорный параллелизм (TP) и параллелизм данных (DP), которые могут быть легко реализованы с использованием библиотек коллективного взаимодействия, таких как torch.distributed и NCCL, маршрутизация MoE включает в себя разреженное одноранговое взаимодействие. Ядро диспетчеризации должно разделить набор входящих токенов и отправить их на узлы, где размещаются эксперты, а его двойственная задача, ядро комбинирования, должна получить их обработанные аналоги и вычислить их взвешенное среднее. Эти передачи не отображаются на коллективные операции и требуют специализированных ядер для достижения низкой задержки.

В Perplexity мы разрабатывали как переносимые межпроцессорные ядра, так и специализированные внутрипроцессорные ядра, представив нашу работу более подробно в предыдущей статье о взаимодействии MoE. Мы также показали в другой публикации в блоге, что многоузловое развертывание может обеспечить меньшую задержку и более высокую пропускную способность для крупных моделей MoE.

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

Вывод через InfiniBand

Тензорный параллелизм и параллелизм экспертов могут эффективно масштабироваться в пределах одного узла, так как NVLink предлагает высокую пропускную способность (900 ГБ/с) при задержке в микросекунды. Хотя системы NVIDIA GB200 и GB300 оснащены достаточным объемом памяти для хранения весов и KV-кэшей моделей с триллионами параметров, развертывание на экземплярах AWS p5en с 8 графическими процессорами H200 ограничено 1120 ГБ HBM, которые должны быть разделены между весами модели и KV-кэшами. Следовательно, определенные модели требуют межпроцессорного развертывания через InfiniBand.

Хотя InfiniBand по-прежнему невероятно быстр, он предлагает пропускную способность лишь в 400 Гбит/с (50 ГБ/с), что добавляет десятки или сотни микросекунд задержки для маршрутизации MoE. К счастью, добавленные задержки можно скрыть за счет использования общих экспертов, микропакетирования и перекрытия вычислений и взаимодействия, чтобы обеспечить жизнеспособную пропускную способность при конкурентоспособной стоимости. DeepSeek впервые применила эти методы с помощью DeepSeek-R1, развернутого на кластерах графических процессоров H800 и адаптеров NVIDIA ConnectX-7 InfiniBand.

В кластерах AWS ситуация иная из-за использования их специализированных адаптеров Elastic Fabric Adapters (EFA). Несмотря на то, что они обеспечивают пиковую пропускную способность 400 Гбит/с при коллективных операциях, обычно используемых в тренировочных нагрузках, они немного уступают ConnectX-7 по размерам сообщений, передаваемых во время диспетчеризации и комбинирования MoE. Кроме того, EFA не поддерживает GPUDirect Async, требуя использования прокси-треда ЦП, связывающего графический процессор с сетевой картой для инициации передачи. Помимо повышенной сложности, дополнительные транзакции PCIe добавляют микросекундные накладные расходы на все передачи. Такие фреймворки, как NVSHMEM, позволяют разработчикам ядер предоставлять реализации, не зависящие от устройства, однако общие API влекут за собой накладные расходы, которые делают развертывание неэффективным как с точки зрения стоимости, так и задержки.

pplx-ядра на базе NVSHMEM

Наши предыдущие ядра абстрагировали лежащую в основе реализацию передачи через NVSHMEM. Они полагались на nvshmemx_putmem_signal_nbi_warp для последовательной передачи токенов, используя атомарные счетчики для синхронизации узлов. Хотя это обеспечивало приемлемую производительность на сетевых картах ConnectX-7 с IBGDA, реализация IBRC для ConnectX-7 на основе прокси была значительно медленнее, а реализация EFA не смогла обеспечить задержки менее миллисекунды.

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

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

MoE all-to-all через EFA

Развернуть

Новые ядра следуют гибридной архитектуре ЦП-ГП, где ядра ГП взаимодействуют с моделью на устройстве, а прокси-тред, запущенный на хост-ЦП, обрабатывает взаимодействия с сетевой картой. Новые ядра повторно используют тот же TransferEngine, который мы создали для передач KV-кэша.

Маршрутизация MoE выполняется ядром диспетчеризации, которое получает токены и экспертов для каждого токена для маршрутизации от модели, отправляя их по сети на узлы, которые размещают экспертов. После того как ядра сгруппированных GEMM умножают токены с весами экспертов, ядро комбинирования передает токены обратно на их исходные узлы и вычисляет их взвешенное среднее перед возвращением их в модель.

И ядра диспетчеризации, и ядра комбинирования разделены на отправителя и получателя для облегчения перекрытия и микропакетирования. Отправитель перемешивает токены в буферы отправки, чтобы их можно было записать одноранговым узлам с помощью одной операции записи RDMA. После выполнения отправителя ГП может выполнять другую работу (например, умножение общих экспертов), пока передачи по сети ожидают в фоновом режиме. Ядра получателя блокируются до тех пор, пока не будут получены все передачи. В фоновом режиме рабочий тред на стороне хоста обрабатывает передачи: он получает сигнал от ГП, когда буферы отправки заполнены, после чего отправляет записи RDMA. Как только сетевая карта сигнализирует о завершении, она разблокирует получателя, который возвращает токены в модель.

Ядра тесно связаны с хостом, взаимодействуя друг с другом через унифицированную память и GDRCopy. Унифицированная память CUDA позволяет отображать память устройства или память хоста в адресные пространства обоих, прозрачно обрабатывая связь через PCIe. Пространство (Space) используется как технический термин.

Хотя унифицированная память подходит для массовых передач, опрос с низкой задержкой осуществляется через GDRCopy. GDRCopy использует GPUDirect RDMA для чтения и записи в память ГП с микросекундной задержкой.

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

Буферы отправки и получения

Конструкция буферов отправки для ядра диспетчеризации имеет решающее значение для достижения жизнеспособных задержек, связанных с хост-прокси. Если существует N узлов, каждый из которых отправляет T токенов, каждый из которых направляется к R экспертам из E, узел может получить до N * T * max(R, E / N) токенов в наихудшем возможном случае, если все они будут направлены к нему. Более того, этот минимальный предел может быть достигнут только тогда, когда отправители размещают свои записи непрерывно, поскольку частные буферы получения использовали бы слишком много памяти. Чтобы отправители могли размещать свои записи последовательно без конфликтов, полная информация о маршрутизации должна быть распределена на каждый узел. Следовательно, наши ядра диспетчеризации сначала обмениваются количеством токенов для каждого эксперта, после чего они могут определить, куда уходят их записи и с какого смещения они получают токены от каждого узла.

Простой обмен информацией о маршрутизации привел бы к неоправданно высоким задержкам, так как каждый узел должен был бы ждать прибытия записей перед отправкой собственных записей. Вместо этого мы резервируем небольшое количество частного пространства (Space) отправителя на каждом узле-получателе и отправляем часть токенов вместе с записями информации о маршрутизации. Это обеспечивает полную загрузку пропускной способности на протяжении большей части процесса диспетчеризации. После того как маршруты обменяны, оставшиеся токены передаются через одну запись на каждый удаленный узел. Обмен информацией о маршрутизации и буферы фиксированного размера, а также передачи, заполняющие их, проиллюстрированы на рисунке ниже.

Развернуть

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

Для NVLink буферы отправки и получения внутри узла отображаются в узлы-партнеры, чтобы они могли напрямую отправлять или получать токены от соседних узлов.

Диспетчеризация и комбинирование

На рисунках ниже показаны сроки и взаимодействия между ГП, ЦП и сетевой картой, показывающие одну итерацию диспетчеризации и комбинирования, перемежающуюся с ядрами модели.

Работа начинается с отправки диспетчеризации, при этом ядро агрегирует информацию о маршрутизации и вычисляет количество токенов для каждого эксперта через атомарное накопление в общей памяти внутри одного блока. Как только количество токенов становится доступным, оно копируется в память хоста, и прокси-тред немедленно уведомляется о начале отправки информации о маршрутизации всем остальным узлам-партнерам. Параллельно диспетчеризация продолжает вычислять кумулятивную сумму количества токенов, определяя смещения, необходимые для упаковки токенов в непрерывные буферы отправки для каждого партнера. Как только смещения становятся доступными, все токены, подлежащие отправке через EFA, упаковываются в буферы отправки. Прокси-тред уведомляется о необходимости разместить записи в буферах отправки межпроцессорных узлов. Внутри узла ядра также немедленно отправляют токены в частные буферы внутри того же узла через NVLink.

Развернуть

Развернуть

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

На стороне получателя ядро сначала извлекает токены из того же узла, которые выходят за пределы емкости частных буферов, через NVLink. Далее оно блокируется до тех пор, пока хост-прокси не просигнализирует о завершении всех передач. Как только сигнал получен, токены перемещаются из буферов получения в выходной тензор, переданный в групповой GEMM. Поскольку буферы повторно используются между диспетчеризацией и комбинированием, барьеры NVLink и RDMA гарантируют, что буферы получателя были прочитаны, прежде чем позволить следующему ядру перезаписать их.

Ядра комбинирования следуют примерно аналогичному шаблону, перемешивая токены, сигнализируя прокси-треду и ожидая завершения. Однако метаданные для размещения их записей легко доступны, и все удаленные смещения известны, поэтому все токены, передаваемые внутри узла, отправляются через NVLink к получателям, а межпроцессорные токены передаются через одну запись. Перед возвратом результатов ядро аккумулирует токены из буферов получателя, записывая их в выходные тензоры.

TransferEngine для маршрутизации MoE

Для прототипирования нашей реализации мы изначально повторно использовали Transfer Engine, созданный для дезагрегированного префилла и декодирования, чтобы выпускать и отслеживать записи RDMA. Однако, в отличие от дезагрегированного декодирования, где каждый узел размещает большое количество записей на один и тот же узел, диспетчеризация MoE размещает небольшие передачи всем узлам за одну операцию. «Из коробки» производительность не была оптимальной, и наша реализация отставала от DeepEP на ConnectX-7 примерно на 20 мкс. Следовательно, мы специализировали TransferEngine для поддержки двух операций, scatter и barrier, чтобы оптимизировать взаимодействия с большим количеством узлов.

Обе операции, scatter и barrier, предоставляют группы узлов-партнеров, которые могут быть предварительно зарегистрированы для предварительного заполнения структур данных, используемых базовой реализацией транспорта (шаблонизация рабочих запросов, WR). Scatter копирует различные срезы из одного и того же исходного буфера в разные смещения на всех членах группы узлов. Barrier сигнализирует всем партнерам через мгновенное сообщение. Постановка в очередь и обработка этих операций в массовом режиме экономит ценные микросекунды на критическом пути. На EFA, где две сетевые карты по 200 Гбит/с агрегируются для обеспечения пропускной способности 400 Гбит/с, мы сегментируем узлы группы по двум сетевым картам вместо сегментации байтов каждой передачи.

Хотя первоначально мы разработали TransferEngine исключительно на EFA с использованием libfabric, ядра MoE работают на ConnectX-7 без изменения ни одной строки кода после добавления поддержки ConnectX-7 в TransferEngine с использованием libibverbs. По сравнению с EFA на основе Scalable Reliable Datagram (SRD), ConnectX-7 требует дополнительной работы по настройке соединения и управлению партнерами. Мы используем пару очередей Unreliable Datagram (UD) для загрузки надежного соединения (RC). Мы создаем две пары очередей для каждого удаленного узла-партнера, одну для двусторонней передачи SEND/RECV и одну для односторонней WRITE_IMM, поскольку и RECV, и WRITE_IMM потребляют операции RECV в порядке, в то время как две операции требуют разных размеров буфера. Кроме того, мы ослабляем порядок транзакций PCIe между ConnectX-7 и памятью ГП. Подобно EFA, мы внедряем оптимизацию шаблонизации WR. Кроме того, при отправке нескольких небольших сообщений одному и тому же узлу мы используем цепочку WR для повышения скорости сообщений.

Поскольку работа на стороне ГП идентична для обоих адаптеров, мы довели базовую реализацию транспорта до предела, чтобы превзойти производительность DeepEP, подчеркнув, что разрыв между маршрутизацией MoE на EFA и ConnectX-7 определяется возможностями базового оборудования.

Оценка

Задержка диспетчеризации/комбинирования

Мы оценили задержки диспетчеризации и комбинирования на системах H200, оснащенных адаптерами ConnectX-7 и EFA, в стандартной конфигурации DeepSeek-V3 и Kimi-K2 со скрытой размерностью 7168, с блочным масштабированием fp8 для диспетчеризации и bf16 для комбинирования. Мы рассматриваем размер пакета 128 для диспетчеризации и 4096 для комбинирования. На ConnectX-7 мы сравниваем с DeepEP и нашими предыдущими pplx-ядрами на базе NVSHMEM, используя как GPUDirect Async (IBGDA), так и хост-прокси (IBRC). Мы достигаем комбинированных задержек 459 мкс, 582 мкс и 692 мкс на EP16, EP32 и EP64 соответственно. Мы превосходим другие ядра на базе EFA, полагающиеся на общие прокси-треды, такие как UCCL-EP, которые достигают 519 мкс, 966 мкс и 1159 мкс.

Развернуть

При декодировании, несмотря на то, что наши ядра на несколько микросекунд медленнее, чем DeepEP, мы достигаем в целом более низких задержек, так как наши ядра комбинирования значительно быстрее. Хотя мы используем хост-прокси, благодаря нашим массовым передачам мы превосходим производительность всех реализаций IBGDA без прокси. Мы также показываем, что реализация прокси-треда имеет значение, так как наш узкоспециализированный прокси на порядок быстрее, чем общая реализация NVSHMEM. Несмотря на тот факт, что EFA в два раза медленнее, чем ConnectX-7 при передаче пакетов размером 256 КБ, которые близки к типичному размеру пакета, передаваемого при маршрутизации MoE, разрыв между двумя адаптерами в целом менее существенен.

Развернуть

Несмотря на то, что наши оптимизированные для декодирования ядра могут масштабироваться до размеров пакетов префилла без каких-либо дополнительных правок, производительность отстает от DeepEP. Это связано главным образом с тем, что мы аккумулируем токены только на стороне получателя при комбинировании, тогда как DeepEP аккумулирует на стороне отправителя, чтобы уменьшить объем передаваемых данных. Однако эта оптимизация может иметь последствия для точности. Кроме того, поскольку ядра передают данные за одну запись, мы используем больше памяти для буферов отправки и получения.

Пропускная способность Kimi-K2 / DeepSeek-V3

Чтобы доказать жизнеспособность наших ядер, мы проводим бенчмаркинг сквозной производительности нашего внутреннего движка вывода, обслуживающего DeepSeek-V3 (671 млрд параметров) и Kimi-K2 (1 трлн параметров) на нескольких экземплярах p5en, оснащенных адаптерами Elastic Fabric Adapters (EFA) и графическими процессорами H200. Мы запускаем модель в режиме проверки черновика и сообщаем о задержке, предполагая уровень принятия 80% с использованием одного слоя MTP.

В DeepSeek-V3 мы сравниваем с внутриузловой базовой линией EP=8, использующей эффективные ядра диспетчеризации и комбинирования только через NVLink. Межпроцессорные развертывания в целом либо соответствуют, либо превосходят производительность одноузлового развертывания, особенно на пакетах среднего размера, что демонстрирует, что масштабирование по узлам действительно приносит улучшения в общую пропускную способность.

На Kimi-K2 ядра обеспечивают обслуживание. Из-за своего размера модель не может поместиться в пределах одного узла при использовании параллелизма данных (DP) на узле H200, поэтому для обслуживания модели требуется межпроцессорное развертывание. Поскольку модель имеет меньше голов внимания, мы достигаем лучших задержек, чем в DeepSeek-V3, и можем обслуживать модель на EFA с жизнеспособной задержкой.

Развернуть

Развернуть

Мы также рассматривали микропакетирование, хотя оно улучшало задержки только при очень больших размерах пакетов или при использовании более одного чернового токена. Чтобы получить выгоду от разделения задач сгруппированных GEMM на две, размеры пакетов токенов должны быть достаточно плотными. Основная часть времени все еще тратится на диспетчеризацию и комбинирование MoE, что указывает на то, что любые дальнейшие улучшения ядер впоследствии улучшат задержки и облегчат микропакетирование.

Будущие работы

Мы тесно сотрудничаем с инженерами AWS для улучшения производительности на EFA. Со своей стороны, мы будем следить за обновлениями libfabric, которые уменьшат накладные расходы на плоскость данных. Кроме того, мы поэкспериментируем с efa-direct, чтобы еще больше сократить накладные расходы пространства пользователя и напрямую взаимодействовать с оборудованием.

Ссылки

Технический отчет DeepSeek R1

Технический отчет Kimi K2

DeepEP

UCCL-EP

Улучшение сетевой производительности HPC-систем с использованием NVIDIA Magnum IO NVSHMEM и GPUDirect Async

Цитата

@online{perplexity2025EnablingTrillionParameterModelsOnAwsEfa,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{Enabling trillion-parameter models on AWS EFA}},
  date = {2025-11-04},
  url = {https://www.perplexity.ai/hub/blog/enabling-trillion-parameter-models-on-aws-efa}
}