Komunikacja punkt-punkt RDMA dla systemów LLM

Systemy LLM tradycyjnie polegały na komunikacji zbiorowej w przypadku ustrukturyzowanych wzorców równoległych podczas trenowania i wnioskowania. Jednakże pojawiające się obciążenia robocze, takie jak wnioskowanie zdezagregowane, routing typu Mixture-of-Experts (MoE) oraz asynchroniczne dostrajanie w ramach uczenia ze wzmocnieniem, wymagają elastyczności komunikacji punkt-punkt.

Nasz najnowszy artykuł badawczy przedstawia TransferEngine, przenośną bibliotekę komunikacyjną opartą na RDMA, która działa na różnych kontrolerach interfejsu sieciowego (NIC), w tym NVIDIA ConnectX-7 i AWS Elastic Fabric Adapter (EFA). Prezentujemy jej skuteczność w trzech systemach produkcyjnych w Perplexity AI:

Zdezagregowane wypełnianie (prefill) i dekodowanie: Transfer KvCache warstwa po warstwie przy użyciu zapisów stronicowanych i obserwatorów UVM, co pozwala na niezależne skalowanie instancji wypełniania i dekodowania.

Transfer wag do uczenia post-treningowego RL w czasie poniżej 2 sekund: Błyskawiczna synchronizacja wag między węzłami trenowania i wnioskowania.

Umożliwienie działania modeli z bilionami parametrów w AWS EFA: Operacje wysyłania i łączenia MoE o niskim opóźnieniu z wykorzystaniem prymitywów scatter i barrier, osiągające najnowocześniejsze opóźnienia w ConnectX-7 oraz pierwszą realną implementację kompatybilną z AWS EFA.

Ten artykuł skupia się na samym TransferEngine, bibliotece komunikacyjnej, która zasila te systemy. Przeczytaj pełny artykuł badawczy na arXiv i poznaj kod na GitHub.

Komunikacja zbiorowa sprawdza się w przypadku ustrukturyzowanych wzorców równoległych, takich jak równoległość tensorowa i równoległość danych, ale nakłada ograniczenia, które ograniczają jej zastosowanie w nowych obciążeniach roboczych:

Stałe członkostwo: Operacje zbiorowe wymagają ustanowienia globalnego świata komunikacji w czasie inicjalizacji, przy zsynchronizowanej konfiguracji wszystkich uczestników. W przypadku zdezagregowanego wnioskowania jest to problematyczne, ponieważ liczba instancji wypełniania i dekodowania musi być skalowana dynamicznie w odpowiedzi na wzorce ruchu.

Ujednolicone kształty tensorów: Operacje zbiorowe zakładają, że wszyscy uczestnicy wymieniają tensory o identycznych kształtach i typach danych. W routingu MoE liczba tokenów wysyłanych do każdego eksperta jest zależna od danych i odkrywana w czasie wykonywania, co wymusza na gęstych operacjach all-to-all stosowanie wypełniania (padding) do najgorszych rozmiarów bufora.

Kolejność operacji: Interfejsy API komunikacji zbiorowej wymuszają zsynchronizowaną semantykę, nawet gdy aplikacje muszą jedynie wiedzieć, że zestaw operacji został zakończony. Na przykład w transferze KvCache wszystkie strony z instancji wypełniania do instancji dekodowania mogą docierać w dowolnej kolejności. Podobnie, w aktualizacjach wag RL, wszystkie tensory wag mogą być przesyłane niezależnie bez ograniczeń dotyczących kolejności.

Komunikacja punkt-punkt oparta na RDMA naturalnie eliminuje te ograniczenia, ale trzy bariery uniemożliwiły jej przyjęcie:

Ograniczenia bibliotek: Chociaż torch.distributed i NCCL udostępniają interfejsy API punkt-punkt, takie jak Send/Recv, operacje te nadal dziedziczą semantykę i ograniczenia komunikacji zbiorowej.

Różnorodność sprzętowa: Karty sieciowe NVIDIA ConnectX z transportem Reliable Connection (RC) są najpopularniejszą konfiguracją RDMA, gwarantującą dostarczanie w kolejności. Dostawcy chmury, tacy jak AWS, używają innego sprzętu; Elastic Fabric Adapter (EFA) z transportem Scalable Reliable Datagram (SRD) dostarcza komunikaty bez zachowania kolejności. EFA wymaga również agregacji czterech kart sieciowych 100 Gbps (lub dwóch 200 Gbps), aby osiągnąć przepustowość 400 Gbps.

Przenośność oprogramowania: Istniejące biblioteki RDMA cierpią z powodu uzależnienia od dostawcy (vendor lock-in). DeepEP wymaga RDMA inicjowanego przez GPU (IBGDA), dostępnego tylko w ConnectX. NVSHMEM wykazuje poważny spadek wydajności w EFA. Biblioteki takie jak Mooncake i NIXL nie mają wsparcia dla EFA lub znajdują się we wczesnej fazie rozwoju. W rezultacie nie istnieje żadne realne, międzyplatformowe rozwiązanie dla wnioskowania LLM.

TransferEngine rozwiązuje te bariery dzięki czterem kluczowym możliwościom:

Wspólne założenia transportowe: Zamiast celować w konkretną kartę sieciową, TransferEngine zakłada jedynie, że bazowy transport jest niezawodny, ale nieuporządkowany, tj. stanowi część wspólną gwarancji zapewnianych przez transporty RC i SRD.

Synchronizacja ImmCounter: Nowatorski prymityw do powiadamiania o zakończeniu pracy, który działa bez polegania na kolejności komunikatów, umożliwiając uruchomienie tego samego kodu zarówno w transportach uporządkowanych (RC), jak i nieuporządkowanych (SRD).

Obsługa wielu kart sieciowych: Przezroczyste zarządzanie wieloma kartami sieciowymi na GPU, automatycznie dzielące transfery w celu maksymalizacji przepustowości.

Minimalna powierzchnia API: Mały zestaw operacji wystarczający do zbudowania złożonych wzorców komunikacyjnych.

Schemat projektowy systemu TransferEngine, przedstawiający ścieżki żądań i wywołań zwrotnych, procesory CPU, bufory UVM, GPU i NIC w dwóch domenach GPU/NUMA.

TransferEngine uruchamia jeden wątek na GPU. Każdy wątek zarządza DomainGroup, która koordynuje od 1 do 4 kart sieciowych RDMA, w zależności od platformy sprzętowej (pojedynczy ConnectX-7 zapewnia 400 Gbps, podczas gdy EFA wymaga czterech adapterów 100 Gbps lub dwóch 200 Gbps). W ramach DomainGroup każda Domain jest specyficzna dla sprzętu i zarządza pojedynczą kartą sieciową, obsługując zarządzanie parami kolejek, przesyłanie pracy i odpytywanie o zakończenie.

Każda instancja TransferEngine udostępnia jeden główny adres (main address) w celu identyfikacji i wykrywania partnerów. Zdalni partnerzy wymieniają struktury NetAddr, które kodują adresy sieciowe i mogą być serializowane do komunikacji poza pasmem (out-of-band). Wszyscy partnerzy w systemie muszą używać tej samej liczby kart sieciowych na GPU. Ta jednolita konfiguracja pozwala TransferEngine inteligentnie dzielić transfery: podczas inicjowania transferu zarówno źródło, jak i miejsce docelowe wiedzą dokładnie, które karty sieciowe są dostępne, umożliwiając automatyczne równoważenie obciążenia między wszystkimi adapterami.

TransferEngine udostępnia minimalne API, które abstrahuje od złożoności bazowych interfejsów 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 składa się z kilku kluczowych komponentów:

Operacje RDMA wymagają, aby obszary pamięci zostały zarejestrowane w karcie sieciowej przed transferem danych. Funkcja reg_mr przyjmuje wskaźnik, długość i urządzenie, zwracając dwa uchwyty:

MrHandle: Lokalny uchwyt używany jako źródło transferów

MrDesc: Serializowalny deskryptor, który można wymienić ze zdalnymi partnerami, aby umożliwić im zapis do zarejestrowanej pamięci

Struktura MrDesc zawiera listę zdalnych adresów sieciowych i zdalnych kluczy (jeden na kartę sieciową), które dają zdalny dostęp do obszaru pamięci. Ten projekt naturalnie obsługuje konfiguracje z wieloma kartami sieciowymi.

W przypadku małych komunikatów i komunikatów płaszczyzny sterowania typu RPC, TransferEngine udostępnia submit_send i submit_recvs:

submit_send: Wysyła komunikat do zdalnego partnera zidentyfikowanego przez NetAddr. Komunikat jest kopiowany wewnętrznie, dzięki czemu wywołujący może natychmiast ponownie użyć lub zwolnić bufor.

submit_recvs: Umieszcza pulę buforów odbiorczych. Gdy nadejdzie komunikat, wywoływane jest wywołanie zwrotne (callback) z odebranymi danymi. Bufor jest automatycznie zwracany do puli po zakończeniu wywołania zwrotnego.

W przypadku transferów dużych danych TransferEngine zapewnia jednostronne operacje zapisu typu zero-copy:

submit_single_write: Przesyła ciągły obszar ze źródła do miejsca docelowego

submit_paged_writes: Przesyła nieciągłe strony określone przez indeksy, kroki i przesunięcia

Obie operacje akceptują opcjonalną 32-bitową wartość natychmiastową (immediate value), która jest dostarczana do odbiorcy atomowo po zakończeniu.

Silnik automatycznie dzieli duże transfery między wiele kart sieciowych, aby zmaksymalizować przepustowość.

Dla operacji skierowanych do wielu partnerów (częstych w routingu MoE), TransferEngine zapewnia zoptymalizowane prymitywy:

add_peer_group: Wstępnie rejestruje zestaw adresów partnerów w celu szybkiego wyszukiwania

submit_scatter: Wysyła różne fragmenty danych do każdego partnera w grupie

submit_barrier: Wysyła powiadomienie tylko o wartości natychmiastowej do wszystkich partnerów (transfer zerobajtowy)

Operacje te są zoptymalizowane pod kątem niskiego opóźnienia poprzez buforowanie wewnętrznych struktur danych i tworzenie szablonów żądań pracy RDMA, unikając powtarzalnego narzutu konfiguracji.

Kluczowym wyzwaniem w systemach LLM jest koordynacja transferów RDMA z postępem GPU. Idealnie, transfer danych powinien rozpocząć się, gdy tylko GPU wygeneruje wynik, ale odpytywanie GPU z procesora CPU jest nieefektywne.

TransferEngine rozwiązuje to za pomocą obserwatorów UVM: zunifikowanych lokalizacji pamięci wirtualnej, które mogą być aktualizowane przez jądra GPU (nawet w ramach grafów CUDA) i wydajnie odpytywane przez CPU za pomocą GDRCopy. Gdy obserwator UVM zostaje przydzielony, rejestrowane jest wywołanie zwrotne, które jest wyzwalane, gdy GPU aktualizuje wartość.

Umożliwia to wzorce takie jak transfer KvCache warstwa po warstwie: po zakończeniu każdej warstwy transformatora GPU zwiększa wartość obserwatora, wyzwalając transfer RDMA KvCache tej warstwy, podczas gdy kolejne warstwy kontynuują obliczenia.

TransferEngine zapewnia mechanizmy powiadamiania o zakończeniu zarówno dla strony inicjującej (osoba wywołująca submit_write), jak i strony odbiorczej (cel operacji zapisu).

Powiadomienie po stronie inicjatora: Po przesłaniu operacji zapisu osoba wywołująca może wybrać powiadomienie o zakończeniu za pomocą funkcji wywołania zwrotnego lub flagi atomowej. To powiadomienie potwierdza, że dane zostały pomyślnie przesłane z bufora źródłowego.

Powiadomienie po stronie odbiorcy: Kluczową innowacją jest ImmCounter, komponent, który śledzi liczniki indeksowane przez 32-bitowe numery natychmiastowe. Gdy operacja zapisu niesie wartość natychmiastową, ImmCounter odbiorcy jest zwiększany atomowo po dostarczeniu danych. Funkcja expect_imm_count pozwala odbiorcy zarejestrować wywołanie zwrotne, które uruchamia się, gdy konkretna wartość natychmiastowa osiągnie docelową liczbę. Liczniki mogą być przezrozyście synchronizowane z GPU przez GDRCopy, bezpośrednio odpytywane lub monitorowane przez wywołania zwrotne.

Ten mechanizm po stronie odbiorcy zapewnia niezawodną synchronizację bez wymogu zachowania kolejności komunikatów. Działa równie dobrze na kartach sieciowych ConnectX (które gwarantują dostarczanie w kolejności przez transport RC) i kartach sieciowych EFA (które dostarczają komunikaty bez zachowania kolejności przez transport SRD).

TransferEngine jest zaimplementowany w języku Rust, ze szczególną dbałością o optymalizację alokacji, wątków i synchronizacji w celu zminimalizowania opóźnień i osiągnięcia wysokiej przepustowości. Silnik uruchamia jeden wątek na każdą DomainGroup, z których każdy jest przypięty do rdzenia procesora na węźle NUMA, do którego podłączone są urządzenia, co minimalizuje opóźnienia związane z planowaniem i dostępem do pamięci. Struktury danych specyficzne dla domeny są przydzielane po przypięciu, aby zapewnić rezerwację pamięci w odpowiednim węźle NUMA. Jeden wątek obsługuje do 4 Domain (z których każda zarządza pojedynczą kartą sieciową), podczas gdy inny dedykowany wątek odpytuje GPU w celu aktualizacji obserwatorów UVM. Komunikacja między wątkami odbywa się za pomocą kolejek bez blokad.

API przekazuje żądania do DomainGroup obsługującej odpowiednie urządzenie, przy czym pierwsza z nich obsługuje również hosta. W ciasnej pętli wątek domeny odpytuje o nową pracę, priorytetyzując przesyłanie nowych żądań. W zależności od sprzętu i konfiguracji żądania są dzielone i równoważone między dostępne Domain. Pierwszy zapis żądania złożonego jest natychmiast umieszczany w kolejce wysyłania karty sieciowej. Po wyczerpaniu nowych żądań wątek przechodzi do przetwarzania oczekujących operacji, wysyłając zapisy w celu wypełnienia potoku sprzętowego. Następnie odpytywane są kolejki zakończeń w celu sprawdzenia zakończonych transferów i przyrostów liczników natychmiastowych. Zdarzenia są agregowane w celu dostarczania powiadomień na transfer, przekazując transfer do dedykowanego wątku wywołań zwrotnych współdzielonego przez wszystkie grupy.

Podział (sharding) wewnątrz DomainGroup jest elastyczny. Transfery mogą być kierowane do konkretnych kart sieciowych według indeksu. Pojedynczy zapis może zostać podzielony. Transfery stronicowane, scatter i operacje barierowe, które wszystkie tłumaczą się na wiele zapisów, mogą być dzielone między wszystkie karty sieciowe.

Domain w ramach TransferEngine są wyspecjalizowane i zoptymalizowane pod kątem sprzętu, którym zarządzają:

AWS EFA: Implementacja EFA wykorzystuje libfabric, zarządzając domeną fabric na kartę sieciową wewnątrz DomainGroup. Kluczowe decyzje implementacyjne i optymalizacje obejmują:

Zapisy o zerowym rozmiarze: EFA odbiega od specyfikacji RDMA, wymagając ważnych deskryptorów docelowych nawet dla zapisów o zerowym rozmiarze. Dlatego TransferEngine wymusza poprawne deskryptory dla wszystkich transferów.

Szablonowanie żądań pracy: Wstępnie wypełnia i zachowuje wspólne pola struktur libfabric przed wysłaniem, redukując narzut operacji na masowe transfery i operacje w grupach partnerów

NVIDIA ConnectX-7: Implementacja ConnectX-7 wykorzystuje libibverbs z transportem RC. Kluczowe decyzje implementacyjne i optymalizacje obejmują:

Pary kolejek UD dla uzgadniania: Każda Domain używa pary kolejek Unreliable Datagram (UD) do wymiany uzgadniania połączeń RC.

Oddzielne pary kolejek RC: Tworzy 2 pary kolejek RC na partnera, jedną dla dwustronnych operacji Send/Recv i drugą dla jednostronnych operacji Write i WriteImm. To rozdzielenie jest konieczne, ponieważ zakończenia Recv i WriteImm zużywają żądania pracy w kolejności przesyłania, umożliwiając wysokopoziomową semantykę Recv przy jednoczesnej obsłudze WriteImm bez zakłóceń.

Szablonowanie żądań pracy: Wstępnie wypełnia wspólne pola struktur żądań pracy dla masowych transferów i grup partnerów.

Łańcuchowanie żądań pracy: Łączy do 4 żądań pracy za pomocą wskaźnika next struktury ibv_send_wr, redukując liczbę wywołań doorbell do karty sieciowej.

Rozluźniona kolejność: Włącza IBV_ACCESS_RELAXED_ORDERING, aby zezwolić na transakcje PCIe bez zachowania kolejności między kartą sieciową a pamięcią GPU.

Celując w część wspólną gwarancji między różnymi transportami RDMA i wprowadzając niezależną od kolejności synchronizację za pomocą ImmCounter, TransferEngine zapewnia przenośną podstawę dla komunikacji punkt-punkt w systemach LLM. Biblioteka jest oprogramowaniem open source i jest dostępna na GitHub. Zachęcamy społeczność do odkrywania jej i budowania na jej bazie rozwiązań dla nowych wzorców komunikacyjnych w wielkoskalowych systemach uczenia maszynowego.

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

Motywacja

TransferEngine

Projekt API

Rejestracja pamięci

Dwustronne Send/Recv

Jednostronny zapis (One-Sided Write)

Grupy partnerów i operacje masowe

Obserwator UVM: Sygnalizacja GPU-do-CPU

Powiadomienie o zakończeniu

Implementacja

Optymalizacje specyficzne dla sprzętu

Wnioski

Cytat