Svorių perkėlimas RL po-mokymui per mažiau nei 2 sekundes

Neseniai pasiekėme 1,3 sekundės trukmės tarpdalykinio lygio parametrų atnaujinimus „Kimi-K2“ modeliam (1T parametrų), perkeldami svorius iš 256 mokymo GPU (BF16) į 128 inferencijos GPU (FP8).

Asinchroninio sustiprinto mokymo (reinforcement learning) smulkaus derinimo metu mokymas ir inferencija vyksta atskiruose GPU. Po kiekvieno mokymo žingsnio nauji svoriai turi būti perduoti į inferencijos mazgus. Daugeliui esamų sistemų trilijono parametrų modeliams prireikia kelių sekundžių ar net minučių.

Naudodami RDMA taškas-į-tašką ryšį, sugebame padaryti svorių perkėlimą itin greitą, nekeisdami inferencijos variklio, o kodą – lengvesnį rašyti ir prižiūrėti.

Mūsų sprendimas sukurtas remiantis „RDMA WRITE“ – vienpusišku primitivu, kai šaltinis tiesiogiai rašo į paskirties vietos GPU atmintį.

python

def rdma_write(src_ptr, dst_ptr, size, src_mr, dst_mr):
    # Write from local [src_ptr, src_ptr+size) to remote [dst_ptr, dst_ptr+size).
    # src_mr and dst_mr contains Memory Region metadata of both sides.
    ...

Paskirties pusė netgi nebus informuota apie perkėlimą. Tai užtikrina mažo vėlavimo, didelio pralaidumo ir nulinio kopijavimo (zero-copy) perkėlimus, kuriuos inicijuoja mokymo mazgai be jokios valdymo logikos inferencijos mazguose.

Metaduomenų rinkimas – valdiklis surenka parametrų metaduomenis iš visų mokymo ir inferencijos GPU.

Grafiko apskaičiavimas – valdiklis apskaičiuoja statinį svorių perkėlimo grafiką, nustatydamas, kuris mokymo GPU siunčia kurį parametrą kuriam inferencijos GPU ir kokia tvarka.

Grafiko platinimas – valdiklis išsiunčia grafiką visiems mokymo GPU.

Vykdymas – po kiekvieno mokymo žingsnio valdiklis duoda signalą mokymo GPU pradėti perkėlimus.

Apibrėžus aukšto lygio darbo eigą, pagrindinis iššūkis yra tai, kaip efektyviai atlikti svorių perkėlimus trilijono parametrų mastu. Čia aprašome vykdymo kelio detales.

Mokymo metu parametrai paskirstomi pagal FSDP išdėstymus. Naudojant full_tensor(), visi DeviceMesh esantys GPU gali atkurti visą parametrą, todėl visi gali būti naudojami kaip svorio perdavimo šaltinis.

Keli atskiri „DeviceMesh“ formuoja tinklelio grupę. Kadangi tos pačios grupės „DeviceMesh“ yra atskiri, jų perkėlimai netrukdo vienas kitam ir gali vykti visiškai lygiagrečiai. Tarp tinklelio grupių įterpiame pasaulinį barjerą, kad užtikrintume tvarką.

Svorio perdavimo kanalas persidengia su kopijomis iš pagrindinio kompiuterio į įrenginį, tensorių surinkimu, suliejimu ir kvantavimu bei RDMA perdavimais per MeshGroups.

Kiekvieno parametrų tenzoriaus perkėlimą traktuojame kaip užduotį. Svorių perkėlimo procesas naudoja kelių tipų aparatinės įrangos šaltinius, todėl svorio perkėlimo užduotį padalijame į skirtingus konvejerio etapus, kurie persidengia laike:

Perkėlimas iš pagrindinio kompiuterio į įrenginį (Host to device memcpy) – jei FSDP perkelia svorį į CPU

Parametrų paruošimas – atkurti visą svorį naudojant full_tensor(), taikyti projekcijos suliejimą, prireikus kiekuoti.

RDMA perkėlimas – nulinio kopijavimo rašymas į nuotolinę inferencijos GPU atmintį

Pasaulinis barjeras – kai visi full_tensor() iškvietimai baigti, sinchronizuoti tarp tinklelio grupių naudojant GLOO per eternetą.

Diegimo metu kiekvienam konvejerio etapui palaikome FIFO užduočių eilę. Kai eilės pradžioje esanti užduotis užbaigia etapą, ji perkeliama į kito etapo eilės pabaigą.

full_tensor() ir kitos GPU operacijos sukelia papildomą GPU atminties naudojimą. Kad išvengtume atminties trūkumo klaidos, užduoties vykdymą pradedame tik tada, kai dabartinės vykdomos užduotys užima mažiau laikinosios GPU atminties nei konfigūruojama riba.

Keletas dizaino sprendimų daro mūsų sistemą žymiai greitesne ir lengviau prižiūrima nei įprasti atvirojo kodo sprendimai.

„Rank0“ pagrindu veikiantis svorio perdavimas surenka svorius prieš paskirstydamas juos į inferencijos mazgą.

Taškas–taškas ryšys perduoda svorius iš abiejų mokymo mazgų į inferencijos mazgo tinklo sąsajas.

Įprastas šablonas yra nukreipti visus parametrus per 0-ojo rango GPU: suvesti mokymo 0-ajame ranke, siųsti į inferencijos 0-ąjį ranką, tada vėl išsklaidyti. Tai greitai tampa kliūtimi, ribojama vieno GPU PCIe pralaidumo ir tinklo plokštės (pvz., 400 Gbps ≈ 50 GB/s).

Priešingai, mūsų taškas-į-tašką sąranka leidžia kiekvienam mokymo GPU siųsti duomenis tiesiogiai į kiekvieną inferencijos GPU, taip užpildant visą tinklo infrastruktūrą, o ne tik vieną jungtį.

RPC pagrindu veikiantis perdavimas naudoja inferencijos pusės buferį ir darbuotojų iškvietimus, o vienpusis perdavimas įrašo svorius tiesiogiai, nedalyvaujant inferencijos darbuotojui.

Kai kurios sistemos remiasi inferencijos variklio update_weight() metodo iškvietimu kiekvienam tenzoriui. Tai reiškia invazinius inferencijos kodo pakeitimus bei papildomas išlaidas iš RPC, serializavimo ir valdymo plokštumos koordinavimo.

Naudodami „RDMA WRITE“ primitivą, svorius atnaujiname tyliai inferencijos GPU atmintyje, be papildomų kopijų. Nėra jokių valdymo plokštumos pranešimų ir jokios CPU valdymo logikos. Nereikia jokių inferencijos variklio pakeitimų.

Svorių perkėlimo procesas gali išnaudoti keturių tipų aparatinės įrangos išteklius: (1) Pagrindinio kompiuterio ir įrenginio duomenų judėjimą (2) GPU skaičiavimus projekcijų suliejimui ir kiekuotei (3) RDMA tinklą duomenų plokštumai (4) Eternetą valdymo plokštumai.

Mūsų dizainas padalijo svorio perkėlimo užduotis į konvejerio etapus, leidžiančius lengvai persidengti tarp skirtingų aparatinės įrangos išteklių.

Kai kurie diegimai iš naujo apskaičiuoja perkėlimo grafiką kiekvienu mokymo žingsniu, pakartotinai renkdami metaduomenis ir dalindami instrukcijas. Tai sukelia nereikalingą valdymo plokštumos vėlavimą.

Mūsų grafikas apskaičiuojamas vieną kartą inicijuojant. Kiekviena mokymo iteracija tiesiog atkartoja planą: valdiklis išduoda „pirmyn“ signalą, o GPU vadovaujasi iš anksto priskirtais maršrutais. Vykdymas yra nuspėjamas ir lengvas.

Gundoma visą svorių atnaujinimo procesą įausti į vieną monolitine funkciją: surinkti metaduomenis, vardų atitikimą, mazgo vidinį surinkimą, projekcijų suliejimą, kiekuotę, ryšio pasaulio dalijimą, tarpmazginį tinklo perkėlimą. Tai sunku teisingai užkoduoti ir dar sunkiau optimizuoti.

Mūsų inžinerijoje šiuos žingsnius atskiriame kaip individualius komponentus. Kiekvienas komponentas gali būti testuojamas atskirai, analizuojamas ir optimizuojamas izoliuotai.

Greitas ir patikimas svorio perkėlimas yra kritinis statybinis blokas didelio masto RL smulkiam derinimui. Apjungę RDMA WRITE primitivą, statinį perkėlimo grafiką ir konvejerinį vykdymą, sumažinome trilijono parametrų atnaujinimus iki vos 1,3 sekundės „Kimi-K2“ modelyje. Šis metodas yra lengvai suprantamas, paprastai prižiūrimas ir leidžia išvengti tradicinių dizainų kliūčių.

BibTeX

@online{perplexity2025WeightTransferForRlPostTrainingInUnder2Seconds,
  author = {{Perplexity} and Lequn Chen and Erin McKnight and Vivian Wang},
  title = {{Weight transfer for RL post-training in under 2 seconds}},
  date = {2025-09-24},
  url = {https://www.perplexity.ai/hub/blog/weight-transfer-for-rl-post-training-in-under-2-seconds}
}

RDMA WRITE: vienpusiški perkėlimai

Aukšto lygio darbo eiga

Svorių perkėlimo vykdymas

„DeviceMesh“ ir tinklelio grupės

Uužduočių konvejeris

GPU atminties naudojimo valdymas

Kodėl tai greita ir paprasta

Taškas-į-tašką ryšys

Vienpusiškas duomenų perkėlimas

Konvejerinis apdorojimas

Statinis grafikas

Švarus atskyrimas

Išvada

Citata