Точкові комунікації RDMA для систем LLM

Традиційно системи LLM покладалися на колективні комунікації для структурованих паралельних патернів під час навчання та виведення. Проте нові робочі навантаження, такі як дезагреговане виведення, маршрутизація сукупності експертів (MoE) та асинхронне навчання з підкріпленням шляхом донавчення, вимагають гнучкості точкових комунікацій.

У нашій останній дослідницькій роботі представлено TransferEngine — переносну комунікаційну бібліотеку на базі RDMA, яка працює на різних мережевих адаптерах (NIC), зокрема NVIDIA ConnectX-7 та AWS Elastic Fabric Adapter (EFA). Ми демонструємо її ефективність у трьох продуктивних системах у Perplexity AI:

Дезагреговані префіл та декодування: поетапне передавання KvCache за допомогою посторінкового запису та спостерігачів UVM, що дозволяє екземплярам префілу й декодування масштабуватися незалежно.

Передавання ваг для донавчення за допомогою навчання з підкріпленням менш ніж за 2 секунди: надзвичайно швидка синхронізація ваг між вузлами навчання та виведення.

Запуск моделей із трильйонами параметрів на AWS EFA: диспетчеризація MoE з низькою затримкою та комбіновані операції з використанням примітивів розсіювання та бар'єрів, що забезпечують передову затримку на ConnectX-7 та першу життєздатну реалізацію, сумісну з AWS EFA.

Ця стаття присвячена самій TransferEngine — комунікаційній бібліотеці, яка забезпечує роботу цих систем. Прочитайте повну дослідницьку роботу на arXiv та ознайомтеся з кодом на GitHub.

Колективні комунікації добре підходять для структурованих паралельних патернів, таких як тензорний паралелізм і паралелізм даних, але вони накладають обмеження, які обмежують їхнє застосування для нових робочих навантажень:

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

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

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

Точкові комунікації на базі RDMA природним чином вирішують ці обмеження, проте їхньому впровадженню перешкоджають три бар'єри:

Обмеження бібліотек: хоча torch.distributed та NCCL надають точкові API, такі як Send/Recv, ці операції все ще успадковують семантику та обмеження колективів.

Апаратне різноманіття: мережеві адаптери NVIDIA ConnectX із транспортом Reliable Connection (RC) є найпоширенішим налаштуванням RDMA, що гарантує доставку в порядку відправлення. Хмарні провайдери, такі як AWS, використовують інше обладнання — Elastic Fabric Adapter (EFA) із транспортом Scalable Reliable Datagram (SRD), який доставляє повідомлення не в порядку відправлення. EFA також вимагає агрегації чотирьох мережевих адаптерів по 100 Гбіт/с (або двох по 200 Гбіт/с) для досягнення пропускної здатності 400 Гбіт/с.

Переносимість програмного забезпечення: наявні бібліотеки RDMA страждають від прив'язки до постачальника. DeepEP вимагає RDMA з ініціалізацією на графічному процесорі (IBGDA), доступного лише на ConnectX. NVSHMEM демонструє серйозне зниження продуктивності на EFA. Бібліотеки, такі як Mooncake та NIXL, не мають підтримки EFA або залишаються на ранній стадії розробки. У результаті не існує життєздатного кроспровайдерного рішення для виведення LLM.

TransferEngine подолає ці бар'єри завдяки чотирьом основним можливостям:

Загальні припущення щодо транспорту: замість орієнтації на конкретний адаптер, TransferEngine припускає лише те, що базовий транспорт є надійним, але не впорядкованим, тобто є перетином гарантій, які надають транспорти RC та SRD.

Синхронизація ImmCounter: новий примітив для сповіщення про завершення, який працює без покладання на порядок повідомлень, дозволяючи одному й тому ж коду запускатися як на впорядкованих (RC), так і на невпорядкованих (SRD) транспортах.

Підтримка кількох мережевих адаптерів: прозоре керування кількома мережевими адаптерами на один графічний процесор з автоматичним шардуванням передавань для максимального збільшення пропускної здатності.

Мінімальна поверхня API: невеликий набір операцій, достатній для побудови складних комунікаційних патернів.

Схема системи TransferEngine, що показує шляхи запитів і зворотних викликів, робочі процесори CPU, буфери UVM, GPU та NIC у двох доменах GPU/NUMA.

TransferEngine створює один робочий тред на кожний графічний процесор. Кожен воркер керує DomainGroup, яка координує від 1 до 4 мережевих адаптерів RDMA залежно від апаратної платформи (один ConnectX-7 забезпечує 400 Гбіт/с, тоді як EFA вимагає чотири адаптери по 100 Гбіт/с або два адаптери по 200 Гбіт/с). Усередині DomainGroup кожний Domain є апаратно-специфічним і керує єдиним мережевим адаптером, обробляючи керування парами черг, надсилання завдань та опитування щодо завершення.

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

TransferEngine пропонує мінімальний API, який абстрагує складність базових інтерфейсів RDMA.

rust

#[serde] struct NetAddr(Bytes);
#[serde] struct MrDesc{ ptr: u64, rkeys: Vec<(NetAddr, u64)> }
struct MrHandle(NonNull<c_void>); type Offset = u64;
struct Pages{ indices: Vec<u32>, stride: u64, offset: Offset }
struct PeerGroupHandle(u64);
struct ScatterDst{ len: u64, src: Offset, dst: (MrDesc,Offset)}
enum OnDone { Callback(fn () -> ()), Flag(Atomic<bool>) }

trait TransferEngine {
  fn main_address() -> NetAddr;
  // Memory Region Management
  fn reg_mr(ptr, len, device) -> (MrHandle, MrDesc);
  // Two-sided Send/Recv
  fn submit_send(addr: NetAddr, msg: &[u8], cb: fn () -> ());
  fn submit_recvs(len: u64, cnt: u64, cb: fn (&[u8]) -> ());
  // One-sided Write
  fn expect_imm_count(imm: u32, count: u32, cb: fn () -> ());
  fn submit_single_write(len: u64, imm: Option<u32>
      src: (MrHandle, Offset), dst: (MrDesc, Offset), OnDone);
  fn submit_paged_writes(page_len: u64, imm: Option<u32>,
      src: (MrHandle, Pages), dst: (MrDesc, Pages), OnDone);
  // One-sided Write to a group of peers
  fn add_peer_group(addrs: Vec<NetAddr>) -> PeerGroupHandle;
  fn submit_scatter(h: Option<PeerGroupHandle>, OnDone,
      imm: Option<u32>, src: MrHandle, dst: Vec<ScatterDst>);
  fn submit_barrier(h: Option<PeerGroupHandle>, OnDone,
      imm: u32, dst: Vec<MrDesc>);
  // Watcher for CPU-GPU synchronization
  fn alloc_uvm_watcher(cb: fn(u64,u64) -> ()) -> NonNull<u64>;
}

API складається з кількох ключових компонентів:

Операції RDMA вимагають реєстрації регіонів пам'яті в мережевому адаптері перед передаванням даних. Функція reg_mr приймає вказівник, довжину та пристрій, повертаючи два дескриптори:

MrHandle: локальний дескриптор, що використовується як джерело для передавань

MrDesc: дескриптор, що серіалізується та яким можна обмінюватися з віддаленими вузлами, щоб надати їм можливість записувати в зареєстровану пам'ять

Структура MrDesc містить список віддалених мережевих адрес та віддалених ключів (по одному на кожен адаптер), які надають віддалений доступ до регіону пам'яті. Такий дизайн природним чином обробляє налаштування з кількома мережевими адаптерами.

Для невеликих повідомлень та повідомлень керуючої площини в стилі RPC, TransferEngine надає submit_send та submit_recvs:

submit_send: надсилає повідомлення на віддалений вузол, ідентифікований за допомогою NetAddr. Повідомлення копіюється всередині функції, тому викликаючий код може одразу повторно використовувати або звільнити буфер.

submit_recvs: розміщує пул буферів приймання. Коли надходить повідомлення, викликається зворотний виклик із отриманими даними. Буфер автоматично повертається до пулу після завершення зворотного виклику.

Для передавання великих обсягів даних TransferEngine надає операції одностороннього запису без копіювання:

submit_single_write: передає нерозривний регіон від джерела до призначення

submit_paged_writes: передає незв'язні сторінки, визначені індексами, кроками та зміщеннями

Обидві операції приймають необов'язкове 32-бітне негайне значення, яке доставляється отримувачу атомарно після завершення.

Рушій автоматично шардує великі передавання між кількома мережевими адаптерами для максимізації пропускної здатності.

Для операцій, спрямованих на кілька вузлів (поширених у маршрутизації MoE), TransferEngine надає оптимізовані примітиви:

add_peer_group: попередньо реєструє набір адрес вузлів для швидкого пошуку

submit_scatter: надсилає різні зрізи даних кожному вузлу в групі

submit_barrier: надсилає сповіщення, що містить лише негайне значення, усім вузлам (передавання нульового банта)

Ці операції оптимізовані для низької затримки завдяки кешуванню внутрішніх структур даних і створенню шаблонів робочих запитів RDMA, що дозволяє уникнути накладних витрат на повторюване налаштування.

Ключовою проблемою в системах LLM є координація передавань RDMA з прогресом роботи графічного процесора. В ідеалі передавання даних має починатися щойно графічний процесор створює вихідні дані, але опитування графічного процесора з центрального є неефективним.

TransferEngine вирішує цю проблему за допомогою спостерігачів UVM (UVM Watchers): розташувань єдиної віртуальної пам'яті, які можуть оновлюватися ядрами графічного процесора (навіть усередині графів CUDA) та ефективно опитуватися центральним процесором за допомогою GDRCopy. Коли спостерігач UVM виділяється, реєструється зворотний виклик, який викликається щоразу, коли графічний процесор оновлює значення.

Це дозволяє реалізувати такі патерни, як поетапне передавання KvCache: після завершення кожного шару трансформера графічний процесор збільшує значення спостерігача, ініціюючи передавання RDMA для KvCache цього шару, тоді як наступні шари продовжують обчислення.

TransferEngine надає механізми сповіщення про завершення як для боку ініціатора (викликача submit_write), так і для боку отримувача (цілі операцій запису).

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

Сповіщення на боці отримувача: ключовою інновацією є ImmCounter — компонент, який відстежує лічильники, індексовані 32-бітними негайними числами. Коли операція запису містить негайне значення, ImmCounter отримувача атомарно збільшується під час доставки даних. Функція expect_imm_count дозволяє отримувачу зареєструвати зворотний виклик, який спрацьовує, коли певне негайне значення досягає цільового лічильника. Лічильники можуть прозоро синхронізуватися з графічним процесором через GDRCopy, опитуватися безпосередньо або контролюватися за допомогою зворотних викликів.

Цей механізм на боці отримувача забезпечує надійну синхронизацію без вимагання порядку повідомлення. Він однаково добре працює на адаптерах ConnectX (які гарантують доставку в порядку відправлення через транспорт RC) та адаптерах EFA (які доставляють повідомлення не в порядку відправлення через транспорт SRD).

TransferEngine реалізовано мовою Rust, ретельно оптимізуючи виділення пам'яті, тредінг і синхронизацію для мінімізації затримки та досягнення високої пропускної здатності. Рушій створює один робочий тред на кожну DomainGroup, прикріплений до ядра центрального процесора на вузлі NUMA, до якого підключено пристрої, що мінімізує затримку планування та доступу до пам'яті. Структури даних, специфічні для домену, виділяються після прив'язки, щоб гарантувати резервування пам'яті в правильному вузлі NUMA. Один робочий тред обробляє до 4 Domain (кожен керує одним мережевим адаптером), тоді як інший виділений тред опитує графічний процесор для оновлення спостерігачів UVM. Міжтредовий зв'язок здійснюється за допомогою черг без блокувань.

API перенаправляє запити до DomainGroup, що обслуговує відповідний пристрій, причому перший із них також обслуговує хост. У тісному циклі доменний воркер опитує нові завдання, надаючи пріоритет надсиланню нових запитів. Залежно від апаратного забезпечення та конфігурації, запити шардуються та збалансовуються за навантаженням між доступними Domain. Перший запис складеного запису негайно розміщується в черзі надсилання мережевого адаптера. Щойно нові запити вичерпуються, воркер переходить до виконання незавершених операцій, розміщуючи записи для заповнення апаратного конвеєру. Згодом черги завершення опитуються для запиту завершених передавань і збільшень негайних лічильників. Події агрегуються для надання сповіщень по кожному передаванню, передаючи процес спеціалізованому треду зворотного виклику, спільному для всіх груп.

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

Domain у складі TransferEngine спеціалізовані та оптимізовані під апаратне забезпечення, яким вони керують:

AWS EFA: реалізація EFA використовує libfabric, керуючи доменною fabric-областю на кожний мережевий адаптер усередині DomainGroup. Ключові рішення щодо реалізації та оптимізації включають:

Записи нульового розміру: EFA відходить від специфікації RDMA, вимагаючи дійсних дескрипторів цілі навіть для записів нульового розміру. Тому TransferEngine забезпечує наявність дійсних дескрипторів для всіх передавань.

Створення шаблонів робочих запитів: попередньо заповнює та зберігає загальні поля структур libfabric перед розміщенням, зменшуючи накладні витрати на одну операцію для масових передавань і групових операцій вузлів

NVIDIA ConnectX-7: реалізація ConnectX-7 використовує libibverbs із транспортом RC. Ключові рішення щодо реалізації та оптимізації включають:

Пари черг UD для рукостискань: кожний Domain використовує пару черг ненадійних датаграм (UD) для обміну рукостисканнями з'єднання RC.

Окремі пари черг RC: створює 2 пари черг RC на кожен вузол: одну для двосторонніх операцій Send/Recv, а іншу для односторонніх операцій Write і WriteImm. Це розділення є необхідним, оскільки завершення як Recv, так і WriteImm споживають робочі запити в порядку їх розміщення, дозволяючи використовувати семантику Recv високого рівня та підтримуючи WriteImm без створення перешкод.

Створення шаблонів робочих запитів: попередньо заповнює загальні поля структур робочих запитів для масових передавань і груп вузлів.

Ланцюжкове зв'язування робочих запитів: зв'язує до 4 робочих запитів через вказівник next структури ibv_send_wr, зменшуючи кількість сигналів дзвінка (doorbell rings) до мережевого адаптера.

Слабкий порядок: вмикає IBV_ACCESS_RELAXED_ORDERING, щоб дозволити транзакції PCIe не в порядку відправлення між мережевим адаптером і пам'яттю графічного процесора.

Орієнтуючись на перетин гарантій між різними транспортами RDMA та впроваджуючи незалежну від порядку синхронизацію за допомогою ImmCounter, TransferEngine надає переносний фундамент для точкових комунікацій у системах LLM. Бібліотека є відкритою та доступною на GitHub, і ми закликаємо спільноту досліджувати її та створювати на її основі рішення для нових комунікаційних патернів у масштабних системах машинного навчання.

BibTeX

@online{perplexity2025RdmaPointToPointCommunicationForLlmSystems,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{RDMA point-to-point communication for LLM systems}},
  date = {2025-11-05},
  url = {https://www.perplexity.ai/hub/blog/rdma-point-to-point-communication-for-llm-systems}
}

Мотивація

TransferEngine

Дизайн API

Реєстрація пам'яті

Двосторонні Send/Recv

Односторонній запис

Групи вузлів та масові операції

Спостерігач UVM: сигналізація від графічного процесора до центрального процесора

Сповіщення про завершення

Реалізація

Апаратно-специфічні оптимізації

Висновок

Цитування