Эффективная и переносимая связь Mixture-of-Experts

Мы представляем высокопроизводительную, переносимую библиотеку с открытым исходным кодом для связи Mixture-of-Experts (MoE), которая обеспечивает в 10 раз более высокую производительность по сравнению со стандартными примитивами связи All-to-All.

Наша реализация включает в себя несколько ключевых технических инноваций, обеспечивающих превосходную эффективность связи MoE:

Связь, инициируемая GPU (IBGDA): поддерживает прямую связь GPU-to-NIC, значительно снижая задержку за счет исключения участия CPU.

Перекрытие связи и вычислений: архитектура разделенного ядра с отдельными этапами отправки и получения позволяет вычислениям выполняться одновременно с передачей данных по сети.

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

Эффективная и переносимая производительность на нескольких узлах: наша реализация достигает скорости до 10 раз выше, чем стандартная связь all-to-all. Несмотря на то, что она примерно в 2 раза медленнее узкоспециализированных реализаций, наш подход предлагает лучшую переносимость между версиями NVSHMEM и сетевыми средами (NVLink, CX-7 и EFA).

Библиотека является полностью открытой и доступна по адресу https://github.com/ppl-ai/pplx-kernels.

В этой статье мы исследуем проблемы экспертного параллелизма в крупномасштабных моделях MoE и описываем наш подход к эффективной отправке и объединению токенов в распределенных средах GPU.

Введение

Модели Mixture-of-Experts (MoE), такие как DeepSeek R1 и Mixtral 8x7B, улучшают плотные модели за счет ограничения количества весов, активируемых для каждого токена. Вместо достижения большого количества параметров за счет увеличения числа слоев или размера линейных слоев в многослойном перцептроне (MLP) каждого декодерного слоя, модели MoE заменяют традиционный MLP несколькими экспертами и маршрутизатором. Например, из 671 млрд параметров DeepSeek R1 только 37 млрд умножаются на данный токен во время инференса. Это значительно сокращает объем вычислений, необходимых для декодирования токена, и улучшает задержку по сравнению с плотной моделью.

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

Чтобы минимизировать задержку в распределенных системах MoE, параллелизм может быть использован на нескольких устройствах, но это создает проблемы со связью. Такие модели, как Mixtral 8x7B или Llama-70B, помещаются на 8 устройствах в пределах одного узла, используя преимущества быстрых и низкозадержных соединений NVLink (до 900 Гбит/с). Однако более крупные модели требуют шардирования экспертов на нескольких узлах с использованием InfiniBand (до 400 Гбит/с), что создает дополнительные проблемы с задержкой.

Развернуть

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

Шардирование и параллелизм

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

Развернуть

Экспертный параллелизм (EP), проиллюстрированный на первом рисунке, распараллеливает только вычисления экспертов. Различные эксперты назначаются на разные устройства, которые хранят их веса. После маршрутизации токены отправляются на соответствующее устройство, а результаты собираются и накапливаются впоследствии. Сложность маршрутизации зависит от степени параллелизма в других частях модели: репликация других слоев может полностью устранить необходимость в маршрутизации, так как каждый ранг может выбирать токены из локально реплицированной таблицы маршрутизации. Однако, если маршрутизацию выполняет только один ранг, требуется широковещательная передача для отправки токенов, индексов и весов соответствующим экспертам. Наконец, операция all-gather или all-to-all broadcast синхронизирует выходные токены с тем рангом, который продолжает выполнение модели. Такая реализация относительно проста, так как torch уже предоставляет необходимые примитивы, хотя некоторые преимущества могут быть получены путем объединения накопления с операцией сбора токенов от экспертов.

Развернуть

Параллелизм только экспертов не масштабируется идеально, так как узлы в кластере могут простаивать, пока модель выполняет слои, не являющиеся экспертными, такие как внимание, нормализация и выборка. Однако вычисления этих слоев, в первую очередь внимания, также могут выиграть от тензорного параллелизма (TP). Большинство моделей полагаются на многоголовое внимание, что означает, что головки внимания и их соответствующие проекции Q, K и V также могут быть шардированы по устройствам, реплицируя или собирая срезы между различными слоями. Если внимание распределено по всем устройствам, операция all-gather может синхронизировать активации, позволяя реплицировать маршрутизацию, что требует примитивов синхронизации, аналогичных случаю экспертного параллелизма, для эффективной реализации. Однако существуют пределы параллелизма на этом уровне, так как уменьшение количества головок внимания ниже определенного порога даст убывающую отдачу.

Развернуть

Для оптимального использования всех устройств и поддержки очень высокой степени экспертного параллелизма, до 128 или 256 GPU, требуется параллелизм данных (DP). Согласно этой схеме, устройства, по которым разделена модель, группируются для одновременной обработки запросов, вычисляя внимание и поддерживая KV-кэши, шардированные по их локальной группе. Несколько экземпляров этих параллельных групп совместно работают над оценкой экспертов, причем каждая размещает другой поднабор весов экспертов. Основываясь на количестве головок внимания, группа может обычно охватывать размер целого узла, так как внутриузловая связь быстрее. Например, на рисунке выше один ранг DP независимо обслуживает два запроса, A и B, обрабатывая внимание, нормализацию и любые другие задачи для запросов. Другой ранг DP обрабатывает отдельный запрос C. Однако первый узел размещает половину экспертов, а другой узел — другую половину, поэтому после маршрутизации токены от A и B могут быть отправлены на второй узел и наоборот. Это приводит к проблеме разреженной связи: каждое устройство может отправлять разное количество токенов на любой другой целевой ранг. Существующие примитивы из torch, в основном all_to_all, не особенно хорошо подходят, так как они могут потребовать некоторой формы заполнения (padding) или синхронизации GPU-CPU и широковещательной передачи метаданных. Для эффективного осуществления связи требуются пользовательские ядра для обработки информации о маршрутизации и инициирования связи с устройств без вмешательства CPU. После маршрутизации ядро отправки (dispatch) должно отправлять токены на ранги, на которые они были маршрутизированы, в то время как на стороне объединения (combine) активации, принадлежащие запросам в текущей группе DP, должны быть собраны. Кроме того, работа должна быть сбалансирована внутри группы DP, которая может иметь несколько устройств, чтобы минимизировать широковещательную передачу и правильно реплицировать токены на всех рангах для последующих слоев.

NVSHMEM

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

Операции NVSHMEM построены вокруг концепции симметричной памяти: они работают с буферами, которые были выделены на всех устройствах, участвующих в связи между устройствами. Локальное устройство может использовать адрес, полученный в рамках собственного локального буфера, вместе с индексом целевого ранга, чтобы указать место назначения операции. Рисунок ниже иллюстрирует эту концепцию: оба GPU выделяют симметричные буферы одинакового размера, сохраняя на них указатели src. Первый ранг хочет отправить 3 целых числа второму, помещая их в начало буфера: nvshmem_int_put_nbi определяет начальный адрес из локального буфера, указывая целевое устройство. Второй ранг получает смещение из своего собственного буфера, отправляя один элемент на первое устройство со смещением на один. Хотя адреса назначения всегда должны быть симметричными буферами, выделенными с помощью nvshmem_alloc, исходные буферы могут быть произвольными областями памяти устройства, при условии, что они предварительно зарегистрированы в NVSHMEM.

Хотя NVSHMEM предоставляет широкий спектр примитивов, наши ядра полагаются только на 3 функции, выстраивая на них всю синхронизацию и ограждение (fencing).

Развернуть

nvshmemx_putmem_signal_nbi_warp: передает блок данных с одного устройства на другое, одновременно устанавливая флаг на удаленном устройстве. Операция либо устанавливает (NVSHMEM_SIGNAL_SET), либо увеличивает (NVSHMEM_SIGNAL_ADD) 64-битное местоположение. Флаг обновляется после передачи всего блока памяти. Если удаленное устройство заметило изменение флага, оно может безопасно получить доступ к буферу в своей собственной памяти. Эта функция полезна для объединения передачи данных и синхронизации.

nvshmemx_signal_op работает с одной областью памяти, обычно 64-битным флагом, атомарно устанавливая или увеличивая его. Это полезно при отправке метаданных и синхронизации устройств.

nvshmem_uint64_wait_until используется на принимающей стороне сигнала для опроса флага, пока удаленное устройство не обновит его.

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

NVSHMEM использует GPUDirect RDMA, чтобы NIC мог напрямую читать и записывать память GPU для обмена данными без участия CPU. Кроме того, на NIC ConnectX, NVSHMEM поддерживает связь, инициируемую GPU (также известную как Infiniband GPUDirect Async или IBGDA), что позволяет GPU напрямую отправлять операции RDMA на NIC и напрямую опрашивать очередь завершения NIC. На платформах, которые не поддерживают связь, инициируемую GPU, NVSHMEM порождает «прокси»-тред на хост-CPU для инициирования связи и опроса завершения от имени GPU. Программа NVSHMEM переносима независимо от того, поддерживается связь, инициируемая GPU, или нет. Однако связь, инициируемая GPU, значительно сокращает задержку, поскольку полностью обходит обходной путь через CPU.

Переносимые ядра

Мы реализуем связь MoE через пару ядер отправки (dispatch) и объединения (combine). Ядра dispatch отвечают за чтение токенов и информации о маршрутизации на каждом ранге, отправляя их соответствующим экспертам. Ядра combine собирают активации, произведенные экспертами, и отправляют их обратно на их исходные ранги, а также вычисляют взвешенное среднее от выбранных экспертов на основе весов, вычисленных маршрутизатором. Ядра дополнительно разделены на компонент отправки и получения, чтобы позволить передаче данных перекрываться с вычислениями. Ядра отправки все являются неблокирующими и несинхронизирующими: они просто отправляют все записи на удаленные устройства. На другом конце ядра combine только читают из памяти, ожидая, пока все необходимые данные не будут переданы. После отправки работы на NIC, пока данные передаются асинхронно по проводам, GPU могут выполнять другую полезную работу локально, например, применяя общих экспертов или вычисляя внимание.

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

Развернуть

Единственная форма глобальной синхронизации, как проиллюстрировано на рисунке выше, реализована в ядре combine-receive. Как только данные отправлены из send, ранг возвращается и продолжает выполнение, когда получает все токены, которые ему требуются от других рангов, и входит в ядра combine send, как только они обработаны. Барьер в combine-receive гарантирует, что ни один ранг не сможет опередить и запустить dispatch-send, пока любой другой ранг все еще ожидает получения данных, что потенциально может вызвать разрушительное перекрытие. Синхронизация выполняется в ядре combine для простоты, в качестве оптимизации тот же барьер можно было бы реализовать между ядрами dispatch-receive и combine-send, еще больше уменьшив задержку, позволив NIC больше времени для завершения передачи, пока GPU выполняет другие ядра.

Оба ядра разделены по всем доступным SM устройств: в то время как ядра dispatch-send и combine-receive должны распараллеливаться по максимальному количеству токенов на ранг (max_m), ядра dispatch-receive и combine-send должны обрабатывать столько токенов от каждого ранга и локального эксперта, максимум для max_m * (num_experts // EP) * (EP // TP).

Отправка (Dispatch)

Ядро dispatch отвечает за отправку токенов и коэффициентов масштабирования блоков коллегам, которые размещают экспертов, к которым были маршрутизированы токены. Сторона отправителя опирается на специализацию варпов для распараллеливания двух задач: агрегирования информации о маршрутизации для определения того, сколько токенов отправляется с каждого ранга, и упаковки токенов в сообщения для отправки на удаленное устройство. Сторона получателя сначала ждет получения всех счетчиков токенов, затем считывает и распаковывает полезные нагрузки из общей памяти в тензоры, которые будут переданы в нижестоящие ядра Grouped GEMM. Оно также сохраняет информацию в буферах, общих с ядрами combine, чтобы указать, куда каждый токен должен быть отправлен обратно. Эта информация необходима, так как ядро receive перемешивает токены в непрерывных буферах в недетерминированном порядке.

Развернуть

В части отправителя один варп из 32 тредов на каждом блоке отвечает за чтение информации о маршрутизации и подсчет количества токенов, которое, как ожидается, должен получить каждый целевой эксперт с этого ранга, отправляя счетчик плюс один с использованием nvshmem_signal_op. Счетчик увеличивается на единицу, так как переход от нуля к ненулевому значению на удаленной стороне означает получение счетчиков от ранга-источника. Параллельно остальные варпы сотрудничают для копирования токенов в симметричную память по всем блокам параллельно, упаковывая активации, коэффициенты масштабирования и их индекс на локальном ранге в непрерывный блок. Индекс требуется отправителю combine для определения адреса, по которому будет записан токен. Затем, убедившись, что все данные были скопированы через барьер по группам варпов, варпы снова работают независимо, каждый отправляя один и тот же буфер разным экспертам параллельно. Токены отправляются с использованием nvshmemx_putmem_signal_nbi_warp, которая также атомарно увеличивает количество отправленных токенов с локального ранга на удаленном устройстве. Внутри группы DP, поскольку каждый ранг владеет репликой токена для отправки, dispatch сбалансирован равномерно, при этом каждое устройство отправляет поднабор токенов.

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

Объединение (Combine)

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

Развернуть

На стороне отправителя ядра параллельно обходят список токенов, назначенных всем локальным экспертам, записывая их в буфер на ранге назначения. Целевой ранг, индекс эксперта и индекс токена считываются из буферов для каждого токена, заполненных ядрами scater при получении токенов. У каждого отправителя есть своя собственная область личной памяти для каждого эксперта, в которую нужно писать, как показано на рисунке выше, что позволяет избежать необходимости синхронизации. Подобно dispatch, combine атомарно увеличивает счетчики для каждого токена на ранге назначения, чтобы указать на получение данных: когда счетчик совпадает с количеством экспертов, которым был отправлен токен, содержимое токена может быть доступно.

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

Тесты (Benchmarks)

Мы оцениваем наши ядра на кластере из 128 GPU, распределенных по 16 узлам, соединенных через InfiniBand и NIC CX-7. Мы пробуем как GPUDirect Async (IBGDA), так и Reliable Connection (RC) с прокси-сервером на CPU. Мы сравниваем их с высокооптимизированной реализацией DeepSeek, а также с плотными примитивами, предоставляемыми PyTorch (через NCCL) или NVSHMEM.

Развернуть

Хотя в плотной операции NVSHMEM all-to-all производительность IBRC и IBGDA схожа (6378 мкс против 6180 мкс), IBGDA значительно быстрее с разреженными ядрами. Сложив Dispatch и Combine, IBGDA использует 902 мкс, тогда как IBRC требует 3223 мкс — улучшение задержки в 3,6 раза. В то время как all-to-all ограничено пропускной способностью, разреженные ядра транслируют на порядки меньше данных, будучи ограниченными задержкой. Инициируя передачу по сети напрямую с GPU, без необходимости прокси-сервера на CPU для координации GPU и NIC, сквозная задержка значительно сокращается.

Хотя наши переносимые ядра примерно в 2 раза медленнее, чем высокооптимизированные ядра DeepSeek, они улучшают задержку в 10 раз по сравнению с плотными ядрами (902 мкс против 9944 мкс).

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

Развернуть

На одном узле (EP8) NVSHMEM использует NVLINK для транспортировки, обеспечивая меньшую задержку и более высокую пропускную способность, чем межсетевое взаимодействие. Наши переносимые ядра примерно в 2,5 раза быстрее, чем DeepEP на одном узле (186 мкс против 481 мкс).

Дальнейшие улучшения

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

Заключение

Мы представили высокопроизводительную переносимую библиотеку для связи MoE, которая обеспечивает в 10 раз более быструю производительность по сравнению со стандартной связью all-to-all, сохраняя при этом совместимость с разнообразными конфигурациями оборудования. На развертываниях с одним узлом с NVLink наше решение демонстрирует задержку в 2,5 раза ниже, чем у предыдущих реализаций.

Наш подход балансирует производительность и переносимость за счет ключевых инноваций, включая поддержку связи, инициируемой GPU, архитектуру разделенных ядер, обеспечивающую перекрытие вычислений и связи, и эффективную отправку токенов с использованием минимальных примитивов NVSHMEM. Хотя наша библиотека примерно в 2 раза медленнее узкоспециализированных реализаций на многоузловых установках, она предлагает превосходную гибкость в различных сетевых средах (NVLink, CX-7 и EFA).

По мере масштабирования моделей MoE, эффективные стратегии связи, подобные нашей, станут все более важными для практического развертывания. Наша полностью открытая реализация доступна по адресу https://github.com/ppl-ai/pplx-kernels.

Ссылки

DeepEP

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

NVIDIA NVSHMEM

Цитирование

@online{perplexity2025EfficientAndPortableMixtureOfExpertsCommunication,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Phi Hoang},
  title = {{Efficient and portable mixture-of-experts communication}},
  date = {2025-04-02},
  url = {https://perplexity.coolkeypoint.com/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}