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

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

Спробуйте наші ядра на GitHub і прочитайте повний дослідницький документ на arXiv.

Архітектура «суміш експертів» (Mixture-of-Experts, MoE) стала стандартною архітектурою для масштабування до сотень мільярдів або трильйонів параметрів, водночас забезпечуючи прийнятну затримку під час виведення. Моделі MoE замінюють щільний шар (dense layer) трансформера набором експертів і невеликим лінійним шаром маршрутизації, який визначає, з яким експертом слід помножити токен. Завдяки групуванню токенів, пов'язаних з тим самим експертом, завдання стає розрідженим і легко розпаралелюваним, оскільки різні експерти тепер можуть розміщуватися на різних графічних процесорах, розподілених між багатьма вузлами.

На відміну від інших форм паралелізму, таких як тензорний паралелизм (Tensor Parallelism, TP) і паралелізм даних (Data Parallelism, DP), які можна легко реалізувати за допомогою бібліотек колективного зв'язку, таких як torch.distributed і NCCL, маршрутизація MoE передбачає розріджений одноранговий зв'язок (peer-to-peer communication). Ядро диспетчеризації (dispatch kernel) повинно розділити набір вхідних токенов і направити їх до рангів, що розміщують експертів, тоді як його двоїста задача, ядро об'єднання (combine kernel), повинна отримати їхні оброблені аналоги та обчислити їх зважене середнє. Ці передачі не зводяться до колективних операцій і вимагають спеціалізованих ядер для досягнення низької затримки.

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

У цій статті ми детально розглядаємо реалізацію наших нових міжвузлових ядер, які значно покращують продуктивність нашої попередньої роботи на базі NVSHMEM. Окрім демонстрації найсучасніших результатів на ConnectX-K7, ми також показуємо, що в межах нашого середовища виконання виведення ці ядра можна використовувати для ефективного обслуговування найбільших моделей з відкритим кодом.

Тензорний паралелізм і паралелізм експертів можуть ефективно масштабуватися в межах одного вузла, оскільки 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, що вимагає наявності проксі-треду ЦП, який з'єднує графічний процесор із мережевою платою (NIC) для ініціювання передач. Окрім додаткової складності, додаткові транзакції PCIe додають накладні витрати мікросекундного масштабу до всіх передач. Такі фреймворки, як NVSHMEM, дозволяють розробникам ядер надавати незалежні від пристрою реалізації, проте загальні API створюють накладні витрати, через які розгортання стає неефективним як з погляду вартості, так і затримки.

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

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

З огляду на ці спостереження, ми розробили наші ядра на базі EFA, що надають той самий інтерфейс, оптимізувавши взаємодію ядро пристрою — хост-проксі «під капотом» задля досягнення пікової продуктивності.

Гібридний зв’язок RDMA та NVLink із проксі хоста, що показує ранги ініціювання й отримання, працівників проксі та NIC.

Нові ядра дотримуються гібридної архітектури ЦП-ГП (CPU-GPU), де ядра графічного процесора взаємодіють з моделлю на пристрої, а проксі-тред, що працює на хост-ЦП, керує взаємодією з мережевою платою (NIC). Нові ядра повторно використовують той самий TransferEngine, який ми створили для передачі кешу KV.

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

І ядра диспетчеризації, і ядра об'єднання розділені на половину відправника та половину одержувача для полегшення перекриття та мікропакетування. Відправник перетасовує токени в буфери відправки, щоб їх можна було записати на однорангові вузли за допомогою однієї операції запису RDMA. Після виконання відправника графічний процесор може виконувати інші завдання (наприклад, множення спільного експерта), поки передачі через мережу очікують у фоновому режимі. Ядра одержувача блокуються доти, доки не будуть отримані всі передачі. У фоновому режимі робочий тред на стороні хоста обробляє передачі: графічний процесор сигналізує йому, коли буфери відправки заповнені, після чого він надсилає запити RDMA write. Щойно мережева плата сигналізує про завершення, вона розблоковує одержувача, який повертає токени моделі.

Ядри тісно пов'язані з хостом, спілкуючись одне з одним як через уніфіковану пам'ять (unified memory), так і через GDRCopy. Уніфікована пам'ять CUDA дозволяє відображати пам'ять пристрою або пам'ять хоста в адресні простори обох, прозоро керуючи зв'язком через PCIe.

Хоча уніфікована пам'ять підходить для масових передач, опитування з низькою затримкою здійснюється за допомогою GDRCopy. GDRCopy використовує GPUDirect RDMA для читання та запису в пам'ять графічного процесора із затримкою в мікросекунди.

На додаток до RDMA, у межах вузла ядра також використовують NVLink для обміну токенами. Це може розвантажити значну частину трафіку — до чверті для EP32 і до однієї восьмої для EP64 — досягаючи загалом нижчих затримок.

Дизайн буферів відправлення для ядра диспетчеризації має вирішальне значення для досягнення життєздатних затримок із залученням хост-проксі. Якщо є N рангів, кожен з яких диспетчеризує T токенів, кожен з яких маршрутизується до R експертів із E, ранг може отримати до N * T * max(R, E / N) токенів у найгіршому можливому випадку, якщо всі вони будуть спрямовані на нього. Крім того, цю мінімальну межу можна досягти лише тоді, коли відправники розміщують свої записи суміжно, оскільки приватні буфери отримання використали б занадто багато пам'яті. Щоб відправники могли послідовно розміщувати свої записи без конфліктів, повна інформація про маршрутизацію має бути розповсюджена на кожен ранг. Отже, наші ядра диспетчеризації спочатку обмінюються кількістю токенів на експерта, після чого кожне з них може визначити, куди йдуть їхні записи та з якого зміщення вони отримують токени від кожного рангу.

Простий обмін інформацією про маршрутизацію призвів би до невиправдано високих затримок, оскільки кожен ранг мав би чекати на прибуття записів перед відправкою власних. Натомість ми резервуємо невелику кількість приватного простору на кожному ранзі-одержувачі та надсилаємо деякі токени разом із записами інформації про маршрутизацію. Це забезпечує повне використання смуги пропускання протягом більшої частини процесу диспетчеризації. Після обміну маршрутами решта токенів передається за допомогою одного запису на віддалений одноранговий вузл. Обмін інформацією про маршрутизацію та буферами фіксованого розміру разом із передачами, що їх заповнюють, проілюстровано на рисунку нижче.

Розташування буфера для MoE-диспетчеризації, що показує, як розміщуються та обмінюються даними токенів і метаданими маршрутизації між двома рангами.

Ядра об'єднання є набагато простішими, оскільки інформація про маршрутизацію вже доступна їм із етапу диспетчеризації. Вони можуть копіювати токени за один запис на одноранговий вузл, після чого зчитують дані з буфера отримання для обчислення зваженого середнього токенів у вихідних тензорах.

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

На рисунках нижче детально описано час і взаємодію між графічним процесором, ЦП та мережевою платою (NIC), показуючи одну ітерацію диспетчеризації та об'єднання, що чергується з ядрами моделі.

Робота починається у відправці диспетчеризації (dispatch send): ядро агрегує інформацію про маршрутизацію та обчислює кількість токенів на експерта за допомогою атомарного накопичення у спільній пам'яті в межах одного блоку. Щойно кількість токенів стає доступною, вона копіюється в пам'ять хоста, і проксі негайно сповіщається про необхідність розпочати диспетчеризацію інформації про маршрутизацію на всі інші однорангові ранги. Паралельно диспетчеризація продовжує обчислювати суму кількості токенів, що накопичується, визначаючи зміщення, необхідні для пакування токенів у суміжні буфери відправки для кожного рангу. Щойно зміщення стають доступними, усі токени, які потрібно надіслати через EFA, пакуються в буфери відправки. Проксі отримує сповіщення про розміщення записів у буферах відправки міжвузлових однорангових вузлів. Усередині вузла ядра також негайно надсилають токени до приватних буферів у межах того самого вузла через NVLink.

Часування та координація для MoE-диспетчеризації, що показує поетапні операції GPU, CPU та NIC і їхню синхронізацію.

Часування та координація для MoE-комбінування, що показує операції GPU, CPU та NIC і їхню синхронізацію.

Коли ранг отримує інформацію про маршрутизацію від усіх однорангових вузлів, він централізує її та обчислює різні зміщення, необхідні для ідентифікації токенів для диспетчеризації та об'єднання. Ця робота виконується на ЦП, поки передачі вже перебувають в очікуванні, тому вона не є критичною до затримок. Згодом, коли всі зміщення стають відомі, розміщуються записи для решти токенів, які перевищують розмір приватного буфера, і фаза відправки диспетчеризації завершується.

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

Ядра об'єднання слідують приблизно схожому шаблону: перетасовують токени, подають сигнал проксі та чекають на завершення. Проте метадані для розміщення їхніх записів є легкодоступними, і всі віддалені зміщення відомі, тому всі токени, передані всередині вузла, надсилаються одержувачам через NVLink, а міжвузлові токени передаються за допомогою єдиного запису. Перед поверненням результатів ядро накопичує токени з буферів одержувача, записуючи їх у вихідні тензори.

Для створення прототипу нашої реалізації ми спочатку повторно використали TransferEngine, створений для дезагрегованого заповнення (prefill) та декодування (decode), щоб ініціювати та відстежувати записи RDMA. Проте, на відміну від дезагрегованого декодування, де кожен вузол розміщує велику кількість записів на той самий одноранговий вузл, диспетчеризація MoE розміщує невеликі передачі для всіх однорангових вузлів в одній операції. "З коробки" продуктивність не була оптимальною, і наша реалізація відставала від DeepEP на ConnectX-7 приблизно на 20 мкс. Отже, ми спеціалізували TransferEngine для підтримки двох операцій — розсіювання (scatter) і бар'єру (barrier), — щоб оптимізувати взаємодію з великою кількістю однорангових вузлів.

І операції розсіювання, і бар'єрні операції надають однорангові групи, які можна попередньо зареєструвати для попереднього заповнення структур даних, що використовуються базовою реалізацією транспорту (шаблонізація Work Request (WR)). Розсіювання копіює різні зрізи з одного вихідного буфера в різні зміщення на всіх учасниках однорангової групи. Бар'єр сигналізує всім одноранговим вузлам через immediate. Ставлення в чергу та обробка цих операцій оптом економить цінні мікросекунди на критичному шляху. На EFA, де дві мережеві плати (NIC) по 200 Гбіт/с об'єднуються для забезпечення пропускної здатності 400 Гбіт/с, ми розділяємо однорангові вузли групи між двома мережевими платами замість шардування байтів кожної передачі.

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

Оскільки робота на стороні графічного процесора ідентична на обох адаптерах, ми довели базову реалізацію транспорту до її меж, щоб перевершити продуктивність DeepEP і підкреслити, що розрив між маршрутизацією EFA та ConnectX-7 MoE визначається можливостями базового апаратного забезпечення.

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

Затримка диспетчеризації та комбінування для декоду MoE EP64, із порівнянням показаних реалізацій і шляхів зв’язку.

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

Затримка диспетчеризації та комбінування для префілінгу MoE EP64, із порівнянням шляхів зв’язку pplx і DeepEP.

Попри те, що наші оптимізовані для декодування ядра можуть масштабуватися до розмірів пакетів заповнення (prefill) без будь-яких додаткових налаштувань, продуктивність відстає від DeepEP. Це передусім пов'язано з тим, що ми накопичуємо токени лише на стороні одержувача об'єднання, тоді як DeepEP накопичує їх у відправнику, щоб зменшити обсяг переданих даних. Однак ця оптимізація може мати наслідки для точності. Крім того, оскільки ядра передають дані за один запис, ми використовуємо більше пам'яті для буферів відправлення та отримання.

Щоб довести життєздатність наших ядер, ми проводимо бенчмаркінг продуктивності наскрізного (end-to-end) виконання нашого власного механізму виведення (inference engine), що обслуговує DeepSeek-V3 (671 млрд параметрів) та Kimi-K2 (1 трлн параметрів) на кількох інстансах p5en, оснащених адаптерами Elastic Fabric Adapters (EFA) та графічними процесорами H200. Ми запускаємо модель у режимі верифікації чернетки (draft verification mode) та звітуємо про затримку, припускаючи рівень прийняття 80% з використанням одного шару MTP.

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

На Kimi-K2 ядра уможливлюють обслуговування. Через свій розмір модель не може поміститися в межах одного вузла з використанням паралелізму даних (DP) на вузлі H200, тому для обслуговування моделі необхідне міжвузлове розгортання. Оскільки модель має менше голів уваги (attention heads), ми досягаємо кращих затримок, ніж у DeepSeek-V3, і можемо обслуговувати модель на EFA з прийнятною затримкою.

Пропускна здатність DeepSeek-V3 на GPU залежно від швидкості виводу для конфігурацій EP32, EP16 і EP8.

Пропускна здатність Kimi-K2 на GPU залежно від швидкості виводу для конфігурацій EP32 і EP16.

Ми також розглянули мікропакетування, хоча воно покращило затримки лише при дуже великих розмірах пакетів або при використанні більш ніж одного чернового токена. Щоб отримати вигоду від розділення завдань згрупованої GEMM на дві частини, розміри пакетів токенів мають бути достатньо щільними. Основна частина часу все ще витрачається на диспетчеризацію та об'єднання MoE, що вказує на те, що будь-які подальші вдосконалення ядер відповідно покращать затримки та сприятимуть мікропакетуванню.

Ми тісно співпрацюємо з інженерами AWS над покращенням продуктивності на EFA. Зі свого боку, ми будемо стежити за оновленнями libfabric, які зменшать накладні витрати на площину даних (data plane overhead). Крім того, ми будемо експериментувати з efa-direct, щоб ще більше зменшити накладні витрати простору користувача та напряму взаємодіяти з апаратним забезпеченням.

Технічний звіт про DeepSeek R1

Технічний звіт про Kimi K2

DeepEP

UCCL-EP

Покращення продуктивності мережі систем HPC за допомогою NVIDIA Magnum IO NVSHMEM та GPUDirect Async

BibTeX

@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}
}

Вступ

Виведення через InfiniBand

pplx-kernels на базі NVSHMEM

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

Буфери відправлення та отримання

Диспетчеризація та об'єднання

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

Оцінка

Затримка диспетчеризації/об'єднання

Пропускна здатність Kimi-K2 / DeepSeek-V3

Майбутня робота

Посилання

Цитата