RDMA punkt-punkti kommunikatsioon LLM-süsteemidele

LLM-süsteemid on treenimisel ja järeldamisel traditsiooniliselt tuginenud kollektiivsele kommunikatsioonile struktureeritud paralleelsete mustrite jaoks. Kuid esilekerkivad töövoogude mudelid, nagu disagregeeritud järeldamine, eksperdisegu (MoE) ruuting ja asünkroonne tugevdusõppel põhinev peenhäälestus, nõuavad punkt-punkti-kommunikatsiooni paindlikkust.

Meie uusim uurimistöö tutvustab TransferEngine'i, portatiivset RDMA-põhist kommunikatsiooniteeki, mis toimib erinevatel võrguliidese kontrolleritel (NIC), sealhulgas NVIDIA ConnectX-7 ja AWS Elastic Fabric Adapter (EFA). Demonstreerime selle tõhusust kolmes Perplexity tootmissüsteemis:

Disagregeeritud eelpuhastus ja dekodeerimine (Prefill ja Decode): kiht-kihiline KvCache'i edastus leheküljelise kirjutamise ja UVM-jälgijate abil, mis võimaldab eelpuhastuse ja dekodeerimise eksemplaridel sõltumatult skaleeruda.

Kaalude edastus RL-järgseks treeninguks vähem kui 2 sekundiga: välkkiire kaalude sünkroniseerimine treening- ja järeldussõlmede vahel.

Triljoniparameetriliste mudelite toetamine AWS EFA-l: väikese latentsusega MoE edastus- ja kombineerimisoperatsioonid hajutamis- ja tõkkeprimitiivide abil, saavutades tipptasemel latentsuse ConnectX-7-l ja esimese toimiva rakenduse, mis ühildub AWS EFA-ga.

See artikkel keskendub TransferEngine'ile endale – kommunikatsiooniteegile, mis neid süsteeme toidab. Lugege täielikku uurimistööd arXiv-is ja tutvuge koodiga GitHubis.

Kollektiivne kommunikatsioon sobib suurepäraselt struktureeritud paralleelsete mustrite, näiteks tensori paralleelsuse ja andmeparalleelsuse jaoks, kuid see seab piirangud, mis piiravad selle rakendatavust esilekerkivate töövoogude korral:

Fikseeritud liikmelisus: kollektiivid nõuavad initsialiseerimisel globaalse kommunikatsioonikeskkonna loomist koos sünkroniseeritud seadistusega kõigi osalejate vahel. Disagregeeritud järeldamise puhul on see problemaatiline, kuna eelpuhastuse ja dekodeerimise eksemplaride arv peab liiklusmustritele reageerides dünaamiliselt skaleeruma.

Ühtsed tensori kujud: kollektiivid eeldavad, et kõik osalejad vahetavad identse kuju ja andmetüübiga ministreid. MoE ruutingus on igale eksperdile saadetud tokenite arv andmepõhine ja tuvastatakse käitusajal, mis sunnib tihedaid all-to-all kollektiive lisama täiteandmeid halvima juhtumi puhvrisuurusteni.

Operatsioonide järjekord: kollektiivsed API-d jõustavad sünkroniseeritud semantika isegi siis, kui rakendused peavad teadma ainult seda, et teatud hulk operatsioone on lõpule viidud. Näiteks KvCache'i edastamisel võivad kõik lehed eelpuhastuse eksemplarilt dekodeerimise eksemplarile saabuda suvalises järjekorras. Samamoodi saab RL-i kaalude uuendamisel kõiki kaalutensoreid edastada sõltumatult ilma järjekorrapiiranguteta.

RDMA-põhine punkt-punkti-kommunikatsioon lahendab need piirangud loomulikult, kuid selle kasutuselevõttu on takistanud kolm tõket:

Teegi piirangud: kuigi torch.distributed ja NCCL pakuvad punkt-punkti API-sid nagu Send/Recv, pärivad need operatsioonid siiski kollektiivse semantika ja piirangud.

Riistvara mitmekesisus: usaldusväärse ühenduse (RC) transpordiga NVIDIA ConnectX NIC-id on kõige levinum RDMA seadistus, mis tagab tellimuspärase kohaletoimetamise. Pilveteenuse pakkujad, nagu AWS, kasutavad erinevat riistvara – Elastic Fabric Adapter (EFA) skaleeritava usaldusväärse datagrammi (SRD) transpordiga edastab sõnumeid järjekorratult. Samuti nõuab EFA 400 Gbps ribalaiuse saavutamiseks nelja 100 Gbps NIC-i (või kahe 200 Gbps NIC-i) koondamist.

Tarkvara portatiivsus: olemasolevad RDMA teegid kannatavad tarnijapõhise lukustuse all. DeepEP nõuab GPU-algatatud RDMA-d (IBGDA), mis on saadaval ainult ConnectX-il. NVSHMEM näitab EFA-l tõsist jõudluse halvenemist. Teekidel nagu Mooncake ja NIXL puudub EFA tugi või on nad varases arendusjärgus. Seetõttu puudub LLM-i järeldamiseks igasugune toimiv tarnijateülene lahendus.

TransferEngine lahendab need tõkked nelja põhivõime kaudu:

Ühised transpordieeldused: konkreetse NIC-i sihtimise asemel eeldab TransferEngine vaid seda, et aluseks olev transport on usaldusväärne, kuid järjestamata, st RC- ja SRD-transpordi pakutavate tagatiste ristumiskoht.

ImmCounteri sünkroniseerimine: uudne primitiiv lõpetamise teavitamiseks, mis toimib sõnumite järjekorrast sõltumata, võimaldades sama koodi käitamist nii järjestatud (RC) kui ka järjestamata (SRD) transpordil.

Mitme NIC-i tugi: mitme NIC-i läbipaistev haldamine GPU kohta, jagades edastused automaatselt ribalaiuse maksimeerimiseks.

Minimaalne API pind: väike hulk operatsioone, millest piisab keeruliste kommunikatsioonimustrite ülesehitamiseks.

TransferEngine'i süsteemidisain, mis näitab päringu- ja tagasihelistamisteid, CPU tööüksuseid, UVM-puhvreid, GPU-sid ja NIC-e üle kahe GPU/NUMA domeeni.

TransferEngine loob ühe töölõime GPU kohta. Iga tööline haldab DomainGroup-i, mis koordineerib 1 kuni 4 RDMA NIC-i, sõltuvalt riistvaraplatvormist (üksainus ConnectX-7 pakub 400 Gbps, samas kui EFA nõuab nelja 100 Gbps adapterit või kahte 200 Gbps adapterit). DomainGroup-i sees on iga Domain riistvaraspetsiifiline ja haldab ühte NIC-i, tegeledes järjekorrapaaride halduse, tööde esitamise ja lõpetamise küsitlemisega.

Iga TransferEngine'i eksemplar avalikustab ühe põhiaadressi vastastikuseks tuvastamiseks ja avastamiseks. Kaugpartnerid vahetavad NetAddr struktuure, mis kodeerivad võrguaadresse ja mida saab ribaväliseks kommunikatsiooniks seriaalida. Kõik süsteemi partnerid peavad kasutama sama palju NIC-e GPU kohta. See ühtne konfiguratsioon võimaldab TransferEngine'il edastusi nutikalt tükeldada: edastuse algatamisel teavad nii saatja kui ka sihtkoht täpselt, millised NIC-id on saadaval, võimaldades automaatset koormuse tasakaalustamist kõigi adapterite vahel.

TransferEngine pakub minimaalset API-t, mis abstraheerib aluseks olevate RDMA liideste keerukuse.

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 koosneb mitmest põhikomponendist:

RDMA operatsioonid nõuavad, et mälu piirkonnad registreeritaks enne andmete edastamist NIC-is. Funktsioon reg_mr võtab vastu viite, pikkuse ja seadme ning tagastab kaks käepidet:

MrHandle: kohalik käepide, mida kasutatakse edastuste allikana

MrDesc: seriaalitav descriptor, mida saab kaugpartneritega vahetada, et võimaldada neil registreeritud mällu kirjutada

Struktuur MrDesc sisaldab loendit kaugvõrgu aadressidest ja kaugturbeklahvidest (üks NIC-i kohta), mis annavad kaugjuurdepääsu mälupiirkonnale. See disain käsitleb mitme NIC-i seadistusi loomulikult.

Väikeste sõnumite ja RPC-stiilis kontrolltasandi sõnumite jaoks pakub TransferEngine funktsioone submit_send ja submit_recvs:

submit_send: saadab sõnumi NetAddr abil tuvastatud kaugpartnerile. Sõnum kopeeritakse sisemiselt, nii et väljakutsuja saab puhvri kohe uuesti kasutusse võtta või vabastada.

submit_recvs: paigutab vastuvõtupuhvrite kogumi. Sõnumi saabumisel kutsutakse tagasiandmefunktsioon saadud andmetega. Puhver tagastatakse automaatselt kogumisse pärast tagasiandmefunktsiooni lõpetamist.

Mahukate andmeedastuste jaoks pakub TransferEngine nullkopeerimisega ühesuunalisi kirjutamisoperatsioone:

submit_single_write: edastab külgneva piirkonna allikast sihtkohta

submit_paged_writes: edastab indeksi, sammude ja nihkete abil määratud mitte-külgnevad lehed

Mõlemad operatsioonid aktsepteerivad valikulist 32-bitist vahetut väärtust, mis toimetatakse vastuvõtjale lõpetamisel aatomiliselt.

Mootor jagab suured edastused automaatselt mitme NIC-i vahel, et maksimeerida ribalaiust.

Mitut partnerit sihtivate operatsioonide jaoks (tavaline MoE ruutingus) pakub TransferEngine optimeeritud primitiive:

add_peer_group: registreerib ette partnereadresside komplekti kiireks otsinguks

submit_scatter: saadab rühma igale partnerile erinevad andmelõigud

submit_barrier: saadab kõigile partneritele ainult vahetu teate (nullbaidine edastus)

Need operatsioonid on madala latentsuse jaoks optimeeritud sisemiste andmestruktuuride vahemällu salvestamise ja RDMA töötaotluste mallipõhise loomise kaudu, vältides korduvaid seadistuskulusid.

LLM-süsteemide peamine väljakutse on RDMA edastuste koordineerimine GPU edusammudega. Ideaalis peaks andmeedastus algama kohe, kui GPU toodab väljundit, kuid GPU küsitlemine CPU poolt on ebatõhus.

TransferEngine lahendab selle UVM-jälgijate (UVM Watchers) abil: ühendatud virtuaalmälu asukohad, mida GPU kernelid saavad värskendada (isegi CUDA graafikute raames) ja mida CPU saab GDRCopy abil tõhusalt küsitleda. UVM-jälgija eraldamisel registreeritakse tagasiandmefunktsioon, mis käivitub alati, kui GPU väärtust värskendab.

See võimaldab selliseid mustreid nagu kiht-kihiline KvCache'i edastus: pärast iga transformaatorikihi valmimist suurendab GPU jälgijat, käivitades selle kihi KvCache'i RDMA edastuse, samal ajal kui järgmised kihid jätkavad arvutamist.

TransferEngine pakub lõpetamise teavitamise mehhanisme nii algataja poolele (funktsiooni submit_write väljakutsuja) kui ka vastuvõtja poolele (kirjutamisoperatsioonide sihtmärk).

Algatajapoolne teatis: kirjutamisoperatsiooni esitamisel saab väljakutsuja valida, kas teda teavitatakse lõpetamisest tagasiandmefunktsiooni või aatomlipu kaudu. See teatis kinnitab, et andmed on allikapuuvrist edukalt edastatud.

Vastuvõtjapoolne teatis: peamine uuendus on ImmCounter, komponent, mis jälgib 32-bitiste vahetute numbrite järgi indekseeritud loendureid. Kui kirjutamisoperatsioon kannab vahetut väärtust, suurendatakse vastuvõtja ImmCounterit andmete kohaletoimetamisel aatomiliselt. Funktsioon expect_imm_count võimaldab vastuvõtjal registreerida tagasiandmefunktsiooni, mis käivitub, kui konkreetne vahetu väärtus jõuab sihtarvuni. Loendureid saab GDRCopy kaudu GPU-ga läbipaistvalt sünkroniseerida, otse küsitleda või tagasiandmefunktsioonide kaudu jälgida.

See vastuvõtjapoolne mehhanism tagab usaldusväärse sünkroniseerimise ilma sõnumite järjekorda nõudmata. See toimib sama hästi ConnectX NIC-idel (mis tagavad tellimuspärase kohaletoimetamise RC-transpordi kaudu) ja EFA NIC-idel (mis edastavad sõnumeid järjekorratult SRD-transpordi kaudu).

TransferEngine on rakendatud Rustis, optimeerides hoolikalt mälueraldisteid, lõimimist ja sünkroniseerimist, et minimeerida latentsust ja saavutada kõrge läbilaskevõime. Mootor loob ühe töölõime iga DomainGroup kohta, joondades igaühe NUMA sõlme CPU tuumale, millega seadmed on ühendatud, minimeerides nii ajastuse kui ka mälu juurdepääsu latentsust. Domeenispetsiifilised andmestruktuurid eraldatakse pärast joondamist, et tagada mälu reserveerimine õiges NUMA sõlmes. Üks töölõim haldab kuni 4 Domain-i (millest igaüks haldab ühte NIC-i), samas kui teine spetsiaalne lõim küsitleb GPU-d UVM-jälgijate värskendamiseks. Lõimedevaheline kommunikatsioon toimub lukuvabade järjekordade kaudu.

API suunab päringud edasi sobivat seadet teenindavale DomainGroup-ile, kusjuures esimene teenindab ka hosti. Tihedas tsüklis küsitleb domeeni tööline uute tööde olemasolu, seades uute päringute esitamise esikohale. Sõltuvalt riistvarast ja konfiguratsioonist jagatakse päringud tükkideks ja tasakaalustatakse saadaolevate Domain-ide vahel. Liitpäringu esimene kirjutamistoiming postitatakse viivitamatult NIC-i saatmisjärjekorda. Kui uued päringud on ammendatud, asub tööline menetlema ootel olevaid operatsioone, postitades kirjutusi riistvaratoru täitmiseks. Seejärel küsitletakse lõpetamisjärjekordi, et pärida lõpetatud edastusi ja vahetuid loenduri suurendamisi. Sündmused koondatakse edastuspõhiste teadete edastamiseks, andes edastuse üle kõigi rühmade vahel jagatud spetsiaalsele tagasiandmefunktsiooni lõimele.

Tükeldamine DomainGroup-is on paindlik. Edastused võivad sihtida konkreetseid NIC-e indeksi järgi. Ühe kirjutamise saab jagada. Leheküljepõhised edastused, hajutamis- ja tõkkeoperatsioonid, mis kõik teisendatakse mitmeks kirjutamiseks, saavad jaguneda kõigi NIC-ide vahel.

TransferEngine'i Domain-id on spetsialiseerunud ja optimeeritud nende kontrolli all oleva riistvara jaoks:

AWS EFA: EFA teostus kasutab libfabricut, hallates iga NIC-i kohta fabric-domain’i DomainGroup-is. Peamised teostusotsused ja optimeerimised hõlmavad järgmist:

Nullsuurusega kirjutused: EFA kaldub RDMA spetsifikatsioonist kõrvale, nõudes kehtivaid sihtkoha deskriptoreid isegi nullsuurusega kirjutuste jaoks. Seetõttu nõuab TransferEngine kehtivaid deskriptoreid kõigi edastuste jaoks.

Töötaotluste mallipõhine loomine: eeltäidab ja säilitab libfabric-u struktuuride ühised väljad enne postitamist, vähendades massedastuste ja partnerirühma operatsioonide toimingupõhiseid kulusid

NVIDIA ConnectX-7: ConnectX-7 rakendus kasutab libibverbs-i RC-transpordiga. Peamised rakendusotsused ja optimeerimised hõlmavad järgmist:

UD järjekorrapaarid käepigistuste jaoks: iga Domain kasutab RC-ühenduse käepigistuste vahetamiseks usaldamatu datagrammi (UD) järjekorrapaari.

Eraldi RC järjekorrapaarid: loob 2 RC järjekorrapaari partneri kohta, ühe kahesuunalise Send/Recv operatsioonide jaoks ja teise ühesuunaliste Write- ja WriteImm-operatsioonide jaoks. See eraldamine on vajalik, kuna nii Recv- kui ka WriteImm-lõpetamised tarbivad töötaotlusi postitamise järjekorras, võimaldades kõrgetasemelist Recv-semantikat, toetades samal ajal WriteImm-i ilma häireteta.

Töötaotluste mallipõhine loomine: eeltäidab massedastuste ja partnerirühmade töötaotluste struktuuride ühiseid välju.

Töötaotluste aheldamine: lingib kuni 4 töötaotlust ibv_send_wr viida next kaudu, vähendades NIC-i saadetavate uksekellahelide arvu.

Lõdvem järjestus: võimaldab IBV_ACCESS_RELAXED_ORDERING seadistust, et lubada NIC-i ja GPU mälu vahelisi järjestuseta PCIe tehinguid.

Sihtides erinevate RDMA-transporditüüpide tagatiste ristumiskohta ja tuues sisse järjestusest sõltumatu sünkroniseerimise ImmCounteri kaudu, pakub TransferEngine portatiivse aluse punkt-punkti-kommunikatsiooniks LLM-süsteemides. Teek on avatud lähtekoodiga ja saadaval GitHubis, ning kutsume kogukonda üles seda uurima ja sellele uusi lahendusi ehitama, et toetada suuremahuliste masinõppesüsteemide esilekerkivaid kommunikatsioonimustreid.

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

Motivatsioon

TransferEngine

API disain

Mälu registreerimine

Kahesuunaline Send/Recv

Ühesuunaline kirjutamine

Partnerite rühmad ja massoperatsioonid

UVM-jälgija: GPU-st CPU-le signaalimine

Lõpetamise teatis

Rakendamine

Riistvaraspetsiifilised optimeerimised

Kokkuvõte

Viide