RDMA taško į tašką ryšys LLM sistemoms

Tradicindiškai LLM sistemos mokymo ir išvedimo metu rėmėsi kolektyviniu ryšiu, skirtu struktūrizuotiems lygiagretumio šablonams. Tačiau atsirandantys darbo krūviai, pavyzdžiui, disperguotas išvedimas, „Mixture-of-Experts“ (MoE) maršrutizavimas ir asinchroninis sustiprinimo mokymosi (angl. reinforcement learning) derinimas, reikalauja taško į tašką ryšio lankstumo.

Mūsų naujausiame tyrimo straipsnyje pristatoma „TransferEngine“ – nešiojama RDMA pagrindu sukurta ryšio biblioteka, veikianti įvairiuose tinklo sąsajos valdikliuose (NIC), įskaitant „NVIDIA ConnectX-7“ ir „AWS Elastic Fabric Adapter“ (EFA). Jos veiksmingumą demonstruojame trijose „Perplexity“ gamybinėse sistemose:

Disperguotas pirminis užpildymas ir dekodavimas (angl. Prefill and Decode): sluoksnis po sluoksnio vykdomas „KvCache“ perkėlimas naudojant puslapiuojamus įrašymus ir UVM stebėtojus, leidžiančius pirminio užpildymo ir dekodavimo egzemplioriams mastelį keisti nepriklausomai.

Svorių perkėlimas RL po-mokymui per mažiau nei 2 sekundes: itin greitas svorių sinchronizavimas tarp mokymo ir išvedimo mazgų.

Trilijono parametrų modelių įgalinimas „AWS EFA“ pagrindu: mažo vėlavimo MoE paskirstymo ir sujungimo operacijos, naudojant išsklaidymo ir barjero primityvus, užtikrinančios pažangiausią vėlavimą „ConnectX-7“ ir pirmąją gyvybingą realizaciją, suderinamą su „AWS EFA“.

Šiame straipsnyje daugiausia dėmesio skiriama pačiai „TransferEngine“ – ryšio bibliotekai, kuri valdo šias sistemas. Skaitykite visą tyrimo straipsnį „arXiv“ ir tyrinėkite kodą GitHub.

Kolektyvinis ryšys puikiai tinka struktūrizuotiems lygiagretumo šablonams, pavyzdžiui, tenzorių lygiagretumui ir duomenų lygiagretumui, tačiau jis nustato apribojimus, kurie riboja jo pritaikomumą naujiems darbo krūviams:

Fiksuota narystė: kolektyviniams ryšiams reikia inicijuojant sukurti pasaulinę ryšio aplinką su sinchronizuota sąranka tarp visų dalyvių. Disperguotam išvedimui tai kelia problemų, nes pirminio užpildymo ir dekodavimo egzempliorių skaičius turi dinamiškai keistis reaguodamas į srauto šablonus.

Vienodos tenzorių formos: kolektyviniai ryšiai daro prielaidą, kad visi dalyviai keičiasi identiškų formų ir duomenų tipų tenzoriais. MoE maršrutizavimo metu kiekvienam ekspertui siunčiamų žetonų skaičius priklauso nuo duomenų ir nustatomas vykdymo metu, todėl tankūs „all-to-all“ kolektyviniai ryšiai priversti taikyti užpildą iki blogiausio atvejo buferio dydžių.

Operacijų eiliškumas: kolektyvinės API užtikrina sinchronizuotą semantiką net tada, kai programoms reikia tik žinoti, kad operacijų rinkinys baigtas. Pavyzdžiui, vykdant „KvCache“ perkėlimą, visi puslapiai iš pirminio užpildymo egzempliorio į dekodavimo egzempliorių gali atvykti bet kokia tvarka. Panašiai ir atnaujinant RL svorius, visus svorio tenzorius galima perkelti nepriklausomai be eiliškumo apribojimų.

RDMA pagrindu sukurtas taško į tašką ryšys natūraliai pašalina šiuos apribojimus, tačiau jo pritaikymui truko tris barjerus:

Bibliotekos apribojimai: nors „torch.distributed“ ir „NCCL“ teikia taško į tašką API, tokias kaip „Send/Recv“, šios operacijos vis tiek paveldi kolektyvinę semantiką ir apribojimus.

Aparatinės įrangos įvairovė: „NVIDIA ConnectX“ NIC su patikimo ryšio (RC) transportu yra dažniausia RDMA sąranka, užtikrinanti pristatymą iš eilės. Debesų paslaugų teikėjai, tokie kaip „AWS“, naudoja kitokią aparatinę įrangą – „Elastic Fabric Adapter“ (EFA) su keičiamo masto patikimo datagramos (SRD) transportu, kuris pranešimus pristato ne iš eilės. EFA taip pat reikalauja apjungti keturis 100 Gbps NIC (arba du 200 Gbps NIC), kad būtų pasiektas 400 Gbps pralaidumas.

Programinės įrangos perkeliamumas: esamos RDMA bibliotekos kenčia nuo pririšimo prie teikėjo. „DeepEP“ reikalauja GPU inicijuojamo RDMA (IBGDA), kuris pasiekiamas tik „ConnectX“. „NVSHMEM“ rodo didelį našumo pablogėjimą naudojant EFA. Bibliotekos, tokios kaip „Mooncake“ ir „NIXL“, neturi EFA palaikymo arba yra ankstyvoje kūrimo stadijoje. Dėl to LLM išvedimui nėra jokio perspektyvaus sprendimo, veikiančio tarp skirtingų teikėjų.

„TransferEngine“ pašalina šiuos barjerus per keturias pagrindines galimybes:

Bendros transporto prielaidos: vietoj to, kad taikytų į konkretų NIC, „TransferEngine“ daro prielaidą tik tuo, kad pagrindinis transportas yra patikimas, bet nesureguliuotas, t. y. tai RC ir SRD transportų teikiamų garantijų sankirta.

„ImmCounter“ sinchronizavimas: naujas baigimo pranešimo primityvas, veikiantis nepasikliaujant pranešimų eiliškumu, leidžiantis tam pačiam kodui veikti tiek tvarkinguose (RC), tiek netvarkinguose (SRD) transportuose.

Kelių NIC palaikymas: skaidrus kelių NIC valdymas vienam GPU, automatiškai dalijantis perkėlimus dalimis (angl. sharding), kad būtų maksimaliai padidintas pralaidumas.

Minimalus API paviršius: nedidelis operacijų rinkinys, kurio pakanka sudėtingiems ryšio šablonams kurti.

TransferEngine sistemos dizainas, rodantis užklausų ir atgalinių iškvietimų kelius, CPU darbuotojus, UVM buferius, GPU ir NIC tinklo plokštes dviejuose GPU/NUMA domenuose.

„TransferEngine“ sukuria vieną darbinę giją vienam GPU. Kiekvienas darbuotojas valdo „DomainGroup“, kuri koordinuoja nuo 1 iki 4 RDMA NIC, priklausomai nuo aparatinės įrangos platformos (vienas „ConnectX-7“ suteikia 400 Gbps, o EFA reikalauja keturių 100 Gbps adapterių arba dviejų 200 Gbps adapterių). „DomainGroup“ viduje kiekvienas „Domain“ yra specifinis aparatiniai įrangai ir valdo vieną NIC, apdorodamas eilių porų valdymą, užduočių pateikimą ir užbaigimo apklausą.

Kiekvienas „TransferEngine“ egzempliorius pateikia vieną pagrindinį adresą (angl. main address), skirtą lygiaverčių mazgų identifikavimui ir atradimui. Nuotoliniai lygiaverčiai mazgai keičiasi „NetAddr“ struktūromis, kurios užkoduoja tinklo adresus ir gali būti serializuojamos juostos išoriniam ryšiui. Visi sistemos lygiaverčiai mazgai turi naudoti vienodą NIC skaičių vienam GPU. Ši vienoda konfigūracija leidžia „TransferEngine“ intelektualiai paskirstyti perkėlimus: pradedant perkėlimą, tiek šaltinis, tiek paskirtis tiksliai žino, kurie NIC yra priasiejami, todėl įgalinamas automatinis apkrovos paskirstymas tarp visų adapterių.

„TransferEngine“ pateikia minimalią API, kuri abstrahuoja pagrindinių RDMA sąsajų sudėtingumą.

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 susideda iš kelių pagrindinių komponentų:

RDMA operacijos reikalauja, kad atminties sritys būtų užregistruotos NIC prieš perduodant duomenis. Funkcija „reg_mr“ priima rodyklę, ilgį ir įrenginį, grąžindama du rankenėles (angl. handles):

„MrHandle“: vietinė rankenėlė, naudojama kaip perkėlimų šaltinis

„MrDesc“: serializuojamas aprašas, kuriuo galima keistis su nuotoliniais lygiaverčiais mazgais, kad jie galėtų rašyti į užregistruotą atmintį

„MrDesc“ struktūroje yra nuotolinių tinklo adresų sąrašas ir nuotolinis raktas (po vieną NIC), suteikiantys nuotolinę prieigą prie atminties srities. Šis dizainas natūraliai tvarko kelių NIC sąrankas.

Mažoms žinutėms ir RPC stiliaus valdymo plokštumos žinutėms „TransferEngine“ teikia „submit_send“ ir „submit_recvs“:

„submit_send“: išsiunčia žinutę nuotoliniam lygiaverčiui mazgui, identifikuojamam pagal „NetAddr“. Žinutė nukopijuojama viduje, todėl kviečiantysis gali nedelsdamas pakartotinai naudoti arba atlaisvinti buferį.

„submit_recvs“: paskelbia gavimo buferiaraštį. Atvykus žinutei, iškviečiamas atgalinis skambinys su gautais duomenimis. Buferis po atgalinio skambučio įvykdymo automatiškai grąžinamas į telkinį.

Didelės apimties duomenų perkėlimams „TransferEngine“ teikia nulinio kopijavimo vienpusio rašymo operacijas:

„submit_single_write“: perkelia gretutinę sritį iš šaltinio į paskirtį

„submit_paged_writes“: perkelia negretutinius puslapius, nurodytus indeksais, žingsniais ir poslinkiais

Abi operacijos priima pasirenkamą 32 bitų tiesioginę reikšmę, kuri po užbaigimo atomiškai pristatoma gavėjui.

Variklis automatiškai padalija didelius perkėlimus per kelis NIC, kad maksimaliai padidintų pralaidumą.

Operacijoms, skirtoms keliems lygiaverčiams mazgams (dažniausiai MoE maršrutizavime), „TransferEngine“ teikia optimizuotus primityvus:

„add_peer_group“: iš anksto užregistruoja lygiaverčių mazgų adresų rinkinį greitai paieškai

„submit_scatter“: išsiunčia skirtingus duomenų pjūvius kiekvienam grupės lygiaverčiui mazgui

„submit_barrier“: išsiunčia tik tiesioginį pranešimą visiems lygiaverčiams mazgams (nulinio baito perkėlimas)

Šios operacijos yra optimizuotos žemam vėlavimui, talpinant vidines duomenų struktūras ir kuriant RDMA užduočių užklausų šablonus, taip išvengiant pasikartojančių sąrankos sąnaudų.

Pagrindinis LLM sistemų iššūkis yra RDMA perkėlimų koordinavimas su GPU eiga. Idealiu atveju duomenų perkėlimas turėtų prasidėti vos tik GPU sugeneruoja išvestį, tačiau GPU apklausa iš CPU yra neefektyvi.

„TransferEngine“ tai sprendžia naudodamas UVM stebėtojus (angl. UVM Watchers): vieningos virtualios atminties vietas, kurias gali atnaujinti GPU branduoliai (netgi CUDA grafikose) ir kurias CPU efektyviai apklausia naudodamas „GDRCopy“. Kai paskiriamas UVM stebėtojas, užregistruojamas atgalinis skambinys, kuris iškviečiamas kiekvieną kartą, kai GPU atnaujina reikšmę.

Tai leidžia taikyti tokius šablonus kaip sluoksnis po sluoksnio vykdomas „KvCache“ perkėlimas: po kiekvieno „transformer“ sluoksnio užbaigimo GPU padidina stebėtojo reikšmę, inicijuodamas to sluoksnio „KvCache“ RDMA perkėlimą, kol vėlesni sluoksniai tęsia skaičiavimus.

„TransferEngine“ teikia baigimo pranešimo mechanizmus tiek iniciatoriaus pusei („submit_write“ kviečiančiajam), tiek gavėjo pusei (rašymo operacijų tikslui).

Iniciatoriaus pusės pranešimas: pateikus rašymo operaciją, kviečiantis asmuo gali pasirinkti būti informuotas apie užbaigimą per atgalinio skambučio funkciją arba atominę vėliavėlę. Šis pranešimas patvirtina, kad duomenys buvo sėkmingai perkelti iš šaltinio buferio.

Gavėjo pusės pranešimas: pagrindinė inovacija yra „ImmCounter“ – komponentas, kuris seka skaitiklius, indeksuotus 32 bitų tiesioginėmis reikšmėmis. Kai rašymo operacija perneša tiesioginę reikšmę, gavėjo „ImmCounter“ atomiškai padidinamas pristatant duomenis. Funkcija „expect_imm_count“ leidžia gavėjui užregistruoti atgalinį skambinį, kuris suveikia, kai konkrečios tiesioginės reikšmės skaitiklis pasiekia tikslinį skaičių. Skaitikliai gali būti skaidriai sinchronizuojami su GPU per „GDRCopy“, apklausiami tiesiogiai arba stebimi per atgalinius skambučius.

Šis gavėjo pusės mechanizmas užtikrina patikimą sinchronizavimą nereikalaudamas pranešimų eiliškumo. Jis vienodai gerai veikia „ConnectX“ NIC (kurie garantuoja pristatymą iš eilės per RC transportą) ir EFA NIC (kurie pristato pranešimus ne iš eilės per SRD transportą).

„TransferEngine“ įdiegta „Rust“ kalba, kruopščiai optimizuojant paskirstymus, gijas ir sinchronizavimą, kad būtų sumažintas vėlavimas ir pasiektas didelis pralaidumas. Variklis sukuria vieną darbinę giją vienai „DomainGroup“, kuri kiekviena pririšama prie CPU branduolio NUMA mazge, prie kurio prijungti įrenginiai, taip sumažinant tiek planavimo, tiek prieigos prie atminties vėlavimą. Domenui būdingos duomenų struktūros paskirstomos po pririšimo, siekiant užtikrinti, kad atmintis būtų rezervuota teisingame NUMA mazge. Viena darbinė gija aptarnauja iki 4 domenų (kiekvienas valdo vieną NIC), o kita skirta gija apklausia GPU, kad atnaujintų UVM stebėtojus. Ryšys tarp gijų vyksta per užrakto neturinčias eiles.

API perduoda užklausas „DomainGroup“, aptarnaujančiai atitinkamą įrenginį, o pirmoji iš jų taip pat aptarnauja pagrindinį kompiuterį (host). Glaudžiame cikle domeno darbininkas apklausia naujas užduotis, teikdamas pirmenybę naujų užklausų pateikimui. Priklausomai nuo aparatinės įrangos ir konfigūracijos, užklausos yra padalijamos ir jų apkrova paskirstoma tarp galimų domenų. Pirmasis sudėtinės užklausos rašymas nedelsiant paskelbiamas NIC siuntimo eilėje. Išsekus naujoms užklausoms, darbininkas pereina prie laukiančių operacijų vykdymo, paskelbdamas rašymus aparatinės įrangos konvejeriui užpildyti. Vėliau užbaigimo eilės apklausiamos, siekiant užklausti apie baigtus perkėlimus ir tiesioginio skaitiklio padidėjimus. Įvykiai agreguojami, kad būtų pateikti pranešimai apie kiekvieną perkėlimą, perduodant perkėlimą dedikuotai atgalinio skambučio gijai, kurią bendrina visos grupės.

Dalinimas viduje „DomainGroup“ yra lankstus. Perkėlimai gali būti nukreipti į konkrečius NIC pagal indeksą. Vieno rašymo operacija gali būti padalinta. Puslapiuojami perkėlimai, išsklaidymo ir barjero operacijos, kurios visos virsta keliais rašymais, gali būti padalintos per visus NIC.

„TransferEngine“ esantys domenai yra specializuoti ir optimizuoti jų valdomai aparatinėi įrangai:

AWS EFA: EFA įgyvendinimas naudoja libfabric, valdydamas kiekvienam NIC skirtą fabric domeną DomainGroup viduje. Pagrindiniai įgyvendinimo sprendimai ir optimizacijos apima:

Nulinio dydžio rašymai: EFA nukrypsta nuo RDMA specifikacijos, reikalaudama galiojančių tikslinių aprašų net ir nulinio dydžio rašymams. Todėl „TransferEngine“ reikalauja galiojančių aprašų visiems perkėlimams.

Užduoties užklausos šablonų kūrimas: prieš skelbiant iš anksto užpildo ir išsaugo bendrus „libfabric“ struktūrų laukus, sumažindamas operacijų sąnaudas masiniams perkėlimams ir lygiaverčių mazgų grupės operacijoms

„NVIDIA ConnectX-7“: „ConnectX-7“ realizacija naudoja „libibverbs“ su RC transportu. Pagrindiniai įgyvendinimo sprendimai ir optimizavimai apima:

UD eilių poros pasisveikinimams (handshakes): kiekvienas domenas naudoja nepatikimo datagramo (UD) eilių porą RC ryšio pasisveikinimams keistis.

Atskiros RC eilių poros: sukuria 2 RC eilių poras vienam lygiaverčiui mazgui: vieną dvipusėms Send/Recv operacijoms, o kitą – vienpusėms Write ir WriteImm operacijoms. Šis atskyrimas būtinas, nes tiek Recv, tiek WriteImm užbaigimai sunaudoja užduočių užklausas paskelbimo tvarka, leidžiant aukšto lygio Recv semantiką ir kartu palaikant WriteImm be trukdžių.

Užduoties užklausos šablonų kūrimas: iš anksto užpildo bendrus užduoties užklausos struktūrų laukus masiniams perkėlimams ir lygiaverčių mazgų grupėms.

Užduoties užklausos grandininis sujungimas: susieja iki 4 užduočių užklausų per „ibv_send_wr“ rodyklę „next“, sumažindamas durų skambučių (doorbell rings) į NIC skaičių.

Laisvesnis eiliškumas: įgalina „IBV_ACCESS_RELAXED_ORDERING“, kad būtų leidžiamos ne iš eilės atliekamos PCIe transakcijos tarp NIC ir GPU atminties.

Nukreiptas į garantijų sankirtą tarp skirtingų RDMA transportų ir įvesdamas nuo eiliškumo nepriklausomą sinchronizavimą per „ImmCounter“, „TransferEngine“ suteikia nešiojamą pagrindą taško į tašką ryšiui LLM sistemose. Biblioteka yra atvirojo kodo ir pasiekiama GitHub, ir mes skatiname bendruomenę ją tyrinėti bei kurti jos pagrindu naujiems ryšio šablonams didelio masto mašininio mokymosi sistemose.

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

Motyvacija

„TransferEngine“

API dizainas

Atminties registracija

Dvipusis siuntimas / gavimas (angl. Two-Sided Send/Recv)

Vienpusis įrašymas (angl. One-Sided Write)

Lygiaverčių mazgų grupės ir masinės operacijos

UVM stebėtojas: GPU į CPU signalizavimas

Baigimo pranešimas

Diegimas

Specifiniai aparatinės įrangos optimizavimai

Išvada

Citata