RDMA pont-pont kommunikáció LLM-rendszerekhez

Az LLM-rendszerek eddig hagyományosan a kollektív kommunikációra támaszkodtak a strukturált párhuzamos mintákhoz a betanítás és a következtetés során. Az újabb munkaterhelések – mint például a szétválasztott következtetés, a szakértő-keverék (MoE) útválasztás és az aszinkron megerősítéses tanulással történő finomhangolás – azonban igénylik a pont-pont közötti kommunikáció rugalmasságát.

Legfrissebb kutatási cikkünk bemutatja a TransferEngine-t, egy hordozható, RDMA-alapú kommunikációs könyvtárat, amely különböző hálózati illesztővezérlők (NIC-ek) – többek között az NVIDIA ConnectX-7 és az AWS Elastic Fabric Adapter (EFA) – között működik. Hatékonyságát a Perplexity AI három éles rendszerében mutatjuk be:

Szétválasztott előzetes kitöltés és dekódolás: Rétegről rétegre történő KvCache-átvitel lapozott írások és UVM-figyelők használatával, ami lehetővé teszi az előzetes kitöltési és dekódolási példányok független skálázását.

Súlyátvitel megerősítéses tanulás utáni betanításhoz 2 másodperc alatt: Villámgyors súlyszinkronizálás a betanítási és a következtetési csomópontok között.

Trillió paraméteres modellek engedélyezése AWS EFA-n: Alacsony késleltetésű MoE-diszpécser és -kombináló műveletek szórási és korlát (barrier) primitívek használatával, a ConnectX-7-en elért csúcskategóriás késleltetéssel és az első életképes, az AWS EFA-val kompatibilis megvalósítással.

Ez a cikk magára a TransferEngine-re, az ezeket a rendszereket hajtó kommunikációs könyvtárra fókuszál. Olvassa el a teljes kutatási cikket az arXiv-on, és fedezze fel a kódot a GitHub-on.

A kollektív kommunikáció kiválóan alkalmas a strukturált párhuzamos mintákhoz, mint például a tenzorpárhuzamosság és az adatszintű párhuzamosság, de olyan korlátozásokat is ró, amelyek korlátozzák az alkalmazhatóságát az újabb munkaterheléseknél:

Rögzített tagság: A kollektívák globális kommunikációs világ létrehozását igénylik az inicializálás időpontjában, az összes résztvevő közötti szinkronizált beállítással. A szétválasztott következtetésnél ez problémás, mivel az előzetes kitöltési és dekódolási példányok számának dinamikusan kell skálázódnia a forgalmi mintákra válaszul.

Egységes tenzorformák: A kollektívák feltételezik, hogy minden résztvevő azonos formájú és adattípusú tenzorokat cserél. Az MoE-útválasztásban az egyes szakértőknek küldött tokenek száma adatfüggő, és futásidőben derül ki, ami arra kényszeríti a sűrű all-to-all kollektívákat, hogy a legrosszabb esetre vonatkozó pufferméretekre töltsék ki (pad) azokat.

Műveleti sorrend: A kollektív API-k szinkronizált szemantikát kényszerítenek ki, még akkor is, amikor az alkalmazásoknak csak azt kell tudniuk, hogy egy műveletsorozat befejeződött. Például a KvCache-átvitel során az előzetes kitöltési példány összes oldala bármilyen sorrendben megérkezhet a dekódolási példányba. Hasonlóképpen, az RL-súlyfrissítésekben az összes súlytenzor átvihető függetlenül, sorrendi korlátozások nélkül.

Az RDMA-alapú pont-pont közötti kommunikáció természetes módon kezeli ezeket a korlátokat, de három akadály gátolta az elterjedését:

Könyvtári korlátozások: Bár a torch.distributed és a NCCL biztosít pont-pont közötti API-kat, mint például a Send/Recv, ezek a műveletek még mindig öröklik a kollektív szemantikát és korlátozásokat.

Hardveres sokféleség: Az NVIDIA ConnectX NIC-ek megbízható kapcsolat (RC) átvitellel a leggyakoribb RDMA-beállítások, amelyek garantálják a sorrendhelyes kézbesítést. A felhőszolgáltatók, például az AWS, más hardvert használnak: az Elastic Fabric Adapter (EFA) méretezhető megbízható datagram (SRD) átvitellel sorrenden kívül szállítja üzeneteket. Az EFA ráadásul négy 100 Gbps-os NIC (vagy két 200 Gbps-os NIC) aggregálását igényli a 400 Gbps-os sávszélesség eléréséhez.

Szoftveres hordozhatóság: A meglévő RDMA-könyvtárakat a gyártói bezárás sújtja. A DeepEP GPU-kezdeményezett RDMA-t (IBGDA) igényel, amely csak a ConnectX-en érhető el. A NVSHMEM súlyos teljesítményromlást mutat EFA-n. Az olyan könyvtárak, mint a Mooncake és a NIXL, nem támogatják az EFA-t, vagy korai fejlesztési stádiumban vannak. Ennek eredményeként nem létezik életképes, szolgáltatókon átívelő megoldás az LLM-következtetéshez.

A TransferEngine négy alapvető képességgel hidalja át ezeket az akadályokat:

Közös átviteli feltételezések: Ahelyett, hogy egy adott NIC-et célozna meg, a TransferEngine csupán azt feltételezi, hogy az alapprobléma szerinti átvitel megbízható, de nem rendezett – azaz az RC- és az SRD-átvitelek által biztosított garanciák metszete.

ImmCounter szinkronizálás: Újszerű primitív a befejezési értesítéshez, amely üzenetsorrend-függetlenül működik, lehetővé téve, hogy ugyanaz a kód fusson rendezett (RC) és rendezetlen (SRD) átviteleken egyaránt.

Több NIC támogatása: GPU-nként több NIC átlátszó kezelése, automatikusan többszörözve (sharding) az átviteleket a sávszélesség maximalizálása érdekében.

Minimális API-felület: Műveletek kis készlete, amely elegendő összetett kommunikációs minták felépítéséhez.

TransferEngine rendszerterv, amely bemutatja a kérés- és visszahívási útvonalakat, a CPU-munkásokat, az UVM puffereket, a GPU-kat és a NIC-eket két GPU/NUMA tartományon át.

A TransferEngine GPU-nként egy munkaszálat (worker thread) indít. Minden munkaszál egy DomainGroup-ot kezel, amely a hardverplatformtól függően 1–4 RDMA NIC-et koordinál (egyetlen ConnectX-7 400 Gbps-ot biztosít, míg az EFA négy 100 Gbps-os vagy két 200 Gbps-os adaptert igényel). Egy DomainGroup-on belül minden Domain hardverspecifikus, és egyetlen NIC-et kezel, ellátva a sorpár-kezelést (queue pair), a munkák beküldését és a befejezés lekérdezését.

Minden TransferEngine-példány egyetlen fő címet (main address) tesz közzé a társkereséshez és felderítéshez. A távoli partnerek olyan NetAddr struktúrákat cserélnek, amelyek kódolják a hálózati címeket, és sorosíthatók a sávon kívüli (out-of-band) kommunikációhoz. A rendszer minden partnerének GPU-nként azonos számú NIC-et kell használnia. Ez az egységes konfiguráció lehetővé teszi a TransferEngine számára az átvitelek intelligens darabolását: az átvitel indításakor a forrás és a célhely is pontosan tudja, hogy mely NIC-ek érhetők el, így automatikus terheléselosztás valósul meg az összes adapter között.

A TransferEngine minimális API-t tesz közzé, amely elrejti az alacsony szintű RDMA-interfészek összetettségét.

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

Az API számos kulcsfontosságú komponenst tartalmaz:

Az RDMA-műveletek megkövetelik, hogy a memóriaterületek regisztrálva legyenek a NIC-nél az adatátvitel előtt. A reg_mr függvény egy mutatót (pointer), egy hosszt és egy eszközt fogad, és két leírót (handle) ad vissza:

MrHandle: Helyi leíró, amelyet az átvitelek forrásaként használnak

MrDesc: Sorosítható leíró, amely megosztható a távoli partnerekkel, lehetővé téve számukra, hogy írjanak a regisztrált memóriába

A MrDesc struktúra távoli hálózati címek és távoli kulcsok listáját tartalmazza (NIC-enként egyet), amelyek távoli hozzáférést biztosítanak a memóriaterülethez. Ez a kialakítás természetes módon kezeli a több-NIC-es beállításokat.

Kis üzenetekhez és RPC-stílusú vezérlősík-üzenetekhez a TransferEngine a submit_send és a submit_recvs függvényeket biztosítja:

submit_send: Elküldi az üzenetet egy NetAddr által azonosított távoli partnernek. Az üzenet belsőleg másolásra kerül, így a hívó azonnal újrafelhasználhatja vagy felszabadíthatja a puffert.

submit_recvs: Fogadópufferek készletét teszi közzé. Üzenet érkezésekor egy visszahívás (callback) fut le a kapott adatokkal. A puffer a visszahívás befejeződése után automatikusan visszakerül a készletbe.

Tömeges adatátvitelekhez a TransferEngine másolásmentes (zero-copy) egyoldalas írási műveleteket biztosít:

submit_single_write: Folytonos tartományt visz át a forrástól a célig

submit_paged_writes: Indexek, lépésközök (strides) és eltolások által megadott, nem folytonos oldalakat visz át

Mindkét művelet elfogad egy opcionális 32 bites azonnali értéket, amelyet a fogadó atomi módon kap meg a befejezéskor.

A motor automatikusan feldarabolja a nagy átviteleket több NIC között a sávszélesség maximalizálása érdekében.

Több partnert célzó műveletekhez (gyakori az MoE-útválasztásban) a TransferEngine optimalizált primitíveket biztosít:

add_peer_group: Előre regisztrál egy partnercím-készletet a gyors keresés érdekében

submit_scatter: Különböző adatszeleteket küld a csoport minden egyes partnerének

submit_barrier: Csak azonnali értesítést küld az összes partnernek (nulla bájtos átvitel)

Ezeket a műveleteket alacsony késleltetésre optimalizálták a belső adatszerkezetek gyorsítótárazásával és az RDMA-munkakérések sablonosításával, elkerülve az ismétlődő beállítási overheadet.

Az LLM-rendszerek egyik kulcsfontosságú kihívása az RDMA-átvitelek összehangolása a GPU előrehaladásával. Ideális esetben az adatátvitelnek amint lehet el kellene indulnia, amint a GPU kimenetet állít elő, de a GPU lekérdezése a CPU-ból nem hatékony.

A TransferEngine ezt UVM Watcherekkel (egységes virtuális memóriájú figyelőkkel) oldja meg: egységes virtuális memóriabeli helyekkel, amelyeket a GPU-kernelek frissíthetnek (akár CUDA-gráfokon belül is), és a CPU a GDRCopy használatával hatékonyan lekérdezhet. Amikor egy UVM-figyelőt lefoglalnak, egy visszahívás regisztrálódik, amely minden alkalommal lefut, amikor a GPU frissíti az értéket.

Ez olyan mintákat tesz lehetővé, mint a rétegről rétegre történő KvCache-átvitel: minden transzformátorréteg befejeződése után a GPU növeli a figyelő értékét, ami kiváltja az adott réteg KvCache-ének RDMA-átvitelét, miközben a további rétegek számítása folytatódik.

A TransferEngine befejezési értesítési mechanizmusokat biztosít mind a kezdeményező oldalon (a submit_write hívójánál), mind a fogadó oldalon (az írási műveletek célpontjánál).

Kezdeményező oldali értesítés: Írási művelet beküldésekor a hívó dönthet úgy, hogy visszahívási függvényen vagy atomi elemen keresztül értesítést kap a befejezésről. Ez az értesítés megerősíti, hogy az adatok sikeresen átkerültek a forráspufferből.

Fogadó oldali értesítés: A kulcsfontosságú innováció az ImmCounter, egy olyan komponens, amely a 32 bites azonnali számokkal indexelt számlálókat követi nyomon. Ha egy írási művelet azonnali értéket hordoz, a fogadó ImmCounter számlálója atomi módon növekszik az adat kézbesítésekor. Az expect_imm_count függvény lehetővé teszi a fogadó számára, hogy olyan visszahívást regisztráljon, amely akkor lép működésbe, amikor egy adott azonnali érték eléri a célszámot. A számlálók átláthatóan szinkronizálhatók a GPU-val a GDRCopyn keresztül, közvetlenül lekérdezhetők, vagy visszahívásokon keresztül felügyelhetők.

Ez a fogadó oldali mechanizmus megbízható szinkronizációt biztosít anélkül, hogy üzenetsorrendre lenne szükség. Ugyanolyan jól működik ConnectX NIC-eken (amelyek garantálják a sorrendhelyes kézbesítést az RC-átvitel felett) és EFA NIC-eken (amelyek sorrenden kívül kézbesítik az üzeneteket az SRD-átvitel felett).

A TransferEngine Rust nyelven íródott, gondosan optimalizálva a memóriafoglalásokat, a szálkezelést és a szinkronizációt a késleltetés minimalizálása és a magas áteresztőképesség elérése érdekében. A motor DomainGroup-onként egy munkaszálat indít, amelyek mindegyike a NUMA-csomópont egy-egy CPU-magjához van rögzítve (pinning), amelyhez az eszközök csatlakoznak, minimalizálva ezzel mind a ütemezési, mind a memóriaelérési késleltetést. A egy domainhez tartozó adatszerkezetek a rögzítés után kerülnek lefoglalásra annak biztosítása érdekében, hogy a memória a megfelelő NUMA-csomópontban legyen fenntartva. Egy munkaszál legfeljebb 4 Domain-t kezel (amelyek mindegyike egyetlen NIC-et irányít), míg egy másik dedikált szál a GPU-t kérdezi le az UVM-figyelők frissítéséhez. A szálak közötti kommunikáció zárolásmentes sorokon (lock-free queues) keresztül történik.

Az API továbbítja a kéréseket a megfelelő eszközt kiszolgáló DomainGroup-nak, és az első egyben a hosztot is kiszolgálja. Egy szoros ciklusban (tight loop) a domain-dolgozó új munkák után kutat, prioritásként kezelve az új kérések beküldését. A hardvertől és a konfigurációtól függően a kérések megoszlanak és terheléselosztáson esnek át az elérhető Domain-ek között. Egy összetett kérés első írása azonnal bekerül a NIC küldési sorába. Miután az új kérések kifogytak, a dolgozó a függőben lévő műveletek előrehaladásával folytatja, írásokat küldve a hardveres folyamat feltöltésére. Ezt követően a befejezési sorok lekérdezésre kerülnek a befejezett átvitelek és a növekményes azonnali számlálók lekérdezése céljából. Az események összesítésre kerülnek az átvitelenkénti értesítések kézbesítéséhez, átadva az átvitelt az összes csoport által megosztott, dedikált visszahívási szálnak.

A DomainGroup belüli darabolás rugalmas. Az átvitelek index alapján célozhatnak meg specifikus NIC-eket. Egyetlen írás is felosztható. Az oldalonkénti átvitelek, a szórási és a korlátműveletek – amelyek mindegyike több írásra konvertálódik – az összes NIC között megoszthatók.

A TransferEngine-en belüli Domain-ek specializáltak és az általuk felügyelt hardverre vannak optimalizálva:

AWS EFA: Az EFA megvalósítás a libfabric könyvtárat használja, NIC-enként egy fabric domaint kezelve a DomainGroup-on belül. A kulcsfontosságú megvalósítási döntések és optimalizációk a következők:

Nulla méretű írások: Az EFA eltér az RDMA-specifikációtól azzal, hogy még a nulla méretű írásokhoz is érvényes célleírókat követel meg. Emiatt a TransferEngine minden átvitelhez érvényes leírókat kényszerít ki.

Munkakérés-sablonosítás: Előre kitölti és megtartja a libfabric struktúrák közös mezőit a beküldés előtt, csökkentve a műveletenkénti overheadet a tömeges átviteleknél és a partnercsoport-műveleteknél

NVIDIA ConnectX-7: A ConnectX-7 megvalósítás a libibverbs könyvtárat használja RC-átvitellel. A kulcsfontosságú megvalósítási döntések és optimalizációk a következők:

UD sorpárok kézfogásokhoz: Minden egyes Domain megbízhatatlan datagram (UD) sorpárt használ az RC-kapcsolati kézfogások cseréjéhez.

Különálló RC sorpárok: Partnerenként 2 RC sorpárt hoz létre: az egyiket a kétirányú Send/Recv műveletekhez, a másikat pedig az egyoldalas Write és WriteImm műveletekhez. Ez a szétválasztás azért szükséges, mert mind a Recv, mind a WriteImm befejezések a beküldési sorrendben fogyasztják a munkakéréseket, így lehetővé téve a magas szintű Recv szemantikát, miközben a WriteImm is támogatott interferencia nélkül.

Munkakérés-sablonosítás: Előre kitölti a munkakérés-struktúrák közös mezőit a tömeges átvitelekhez és a partnercsoportokhoz.

Munkakérés láncolása: Legfeljebb 4 munkakérést láncol össze az ibv_send_wr next mutatóján keresztül, csökkentve a NIC felé irányuló csengőhívások (doorbell rings) számát.

Laza sorrendiség: Engedélyezi az IBV_ACCESS_RELAXED_ORDERING használatát, hogy lehetővé téve a sorrenden kívüli PCIe-tranzakciókat a NIC és a GPU memóriája között.

Azáltal, hogy a különböző RDMA-átvitelek garanciáinak metszetét célozza meg, és az ImmCounter révén bevezeti a sorrendfüggetlen szinkronizációt, a TransferEngine hordozható alapot biztosít a pont-pont közötti kommunikációhoz az LLM-rendszerekben. A könyvtár nyílt forráskódú és elérhető a GitHub-on; arra buzdítjuk a közösséget, hogy fedezze fel és építsen rá a nagyszabású gépi tanulási rendszerek új kommunikációs mintáihoz.

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

Motiváció

TransferEngine

API-tervezés

Memóriaregisztráció

Kétirányú Send/Recv

Egyoldalas írás (One-Sided Write)

Partnercsoportok és tömeges műveletek

UVM Watcher: GPU–CPU jelzés

Befejezési értesítés

Megvalósítás

Hardverspecifikus optimalizációk

Összegzés

Hivatkozás