Vogaflutningur fyrir eftirþjálfun í RL á undir 2 sekúndum

Við náðum nýlega 1,3 sekúndna breytuuppfærslum milli véla fyrir Kimi-K2 (1T breytur), þar sem vogum var flutt frá 256 þjálfunar-GPU-um (BF16) yfir í 128 ályktunar-GPU-a (FP8).

Í fylgilausri fínstillingu með styrktarnámi fara þjálfun og ályktun fram á aðskildum GPU-um. Eftir hvert þjálfunarskref þarf að ýta nýjum vogum á ályktunarhnúta. Margir núverandi rammar taka nokkrar sekúndur — eða jafnvel mínútur — fyrir líkön með billjón breytum.

Með því að nýta RDMA jafningjasamskipti getum við gert vogaflutninginn gríðarhraðan án þess að breyta ályktunarvélinni, og gert kóðann auðveldari í skrifum og viðhaldi.

Lausn vor er byggð á RDMA WRITE, einhliða grunnaðgerð þar sem uppruninn skrifar beint í minni áfangastaðar-GPU.

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

Áfangastaðurinn fær ekki einu sinni tilkynningu um flutninginn. Þetta gefur okkur lágt seinkunargildi, mikið afköst og núll-afritunarflutning sem þjálfunarhnútarnir knýja áfram án allrar stýrirökfræði á ályktunarhnútunum.

Lýsigagnasöfnun – Stjórnandi safnar lýsigögnum breyta frá öllum þjálfunar- og ályktunar-GPU-um.

Áætlunarútreikningur – Stjórnandi reiknar út kyrrstæða vogaflutningsáætlun sem vörpunar hvaða þjálfunar-GPU sendir hvaða breytu til hvaða ályktunar-GPU, og í hvaða röð.

Áætlunardreifing – Stjórnandi sendir áætlunina til allra þjálfunar-GPU-a.

Framkvæmd – Eftir hvert þjálfunarskref gefur stjórnandinn merki um að þjálfunar-GPU eigi að hefja flutninga.

Með skilgreindu hátt stigs verkflæði er lykiláskorunin hvernig eigi að framkvæma vogaflutninga á skilvirkan hátt í billjóna breytu skala. Hér lýsum við smáatriðum framkvæmdarferilsins.

Breytur í þjálfun eru dreifðar samkvæmt FSDP staðsetningum. Með því að nota full_tensor() geta allir GPU-ar í DeviceMesh endurbyggt allar breyturnar, og þar af leiðandi geta allir þjónað sem uppruni fyrir vogaflutning.

Margir aðskildir DeviceMesh-ar mynda mesh-hóp. Þar sem DeviceMesh-ar í sama hópi eru aðskildir trufla flutningar þeirra hver annan ekki og geta keyrt að fullu samhliða. Milli mesh-hópa setjum við inn alþjóðlega hindrun til að framfylgja röðun.

Pípa fyrir flutning á þyngdum skarast við host-to-device afrit, samsetningu tensa, samruna og magngreiningu, og RDMA-flutninga yfir MeshGroups.

Við meðhöndlum flutning á hverjum breytutensil sem verkefni. Vogaflutningsferlið nýtir margar tegundir vélbúnaðaruppruna, þess vegna skiptum við vogaflutningsverkefni í mismunandi stig leiðslunnar sem skarast í tíma:

Hýsil-til-tækis memcpy — Ef FSDP hleður voga af vinnsluminni yfir á örgjörva

Undirbúningur breytu — Endurbyggja allar vogir með full_tensor(), beita vörpunarsamruna, magnbundna ef þörf krefur.

RDMA-flutningur — Núll-afritunarskrif í minni fjartengts ályktunar-GPU

Alþjóðleg hindrun — Eftir að öllum full_tensor() köllum er lokið, samstilla yfir mesh-hópa með GLOO í gegnum Ethernet.

Við útfærslu viðhöldum við FIFO biðröð verkefna fyrir hvert leiðslustig. Hvenær sem fremsta verkefni biðraðarinnar lýkur stiginu er því fært aftast í biðröð næsta stigs.

full_tensor() og aðrar GPU-aðgerðir kynna til sögunnar aukna notkun á GPU-minni. Til að koma í veg fyrir minnisleysi hefjum við framkvæmd verkefnis að því einu sinni að núverandi verkefni í vinnslu taki minna tímabundið GPU-minni en stillanlegt vatnsmerki.

Nokkrar hönnunarákvarðanir gera kerfið okkar verulega hraðvirkara í keyrslu og auðveldara í viðhaldi en algengar opnar lausnir.

Þyngdaflutningur byggður á Rank0 safnar þyngdum áður en þeim er dreift til ályktunarnóðsins.

Point-to-point samskipti flytja þyngdir frá báðum þjálfunarnóðum til netviðmóta ályktunarnóðsins.

Algett mynstur er að beina öllum breytum í gegnum röð-0 GPU: safna saman á þjálfunarröð-0, senda í ályktunarröð-0 og dreifa svo aftur. Þetta verður fljótt flöskuháls, takmarkað af PCIe-bandbreidd eins GPU og netkorti (t.d. 400 Gbps ≈ 50 GB/s).

Hins vegar gerir jafningjauppsetning okkar öllum þjálfunar-GPU-um kleift að senda beint til allra ályktunar-GPU-a, og nýta þannig allan netvefinn í stað eins tengils.

RPC-byggður flutningur notar biðminni á ályktunarhliðinni og vinnuköll, en einhliða flutningur skrifar þyngdir beint án þátttöku ályktunarvinnslu.

Sum kerfi reiða sig á að kalla á update_weight() aðferð ályktunarvélarinnar fyrir hvern tensil. Það þýðir innrásargjarnar breytingar á ályktunarkóðanum, auk yfirbyrðis frá fjarforritaköllum (RPC), raðgreiningu og stjórnborðssamhæfingu.

Með RDMA WRITE grunnaðgerðinni uppfærum við vogir hljóðlega í minni ályktunar-GPU-a, án aukaafrita. Engin stjórnborðsskilaboð og engin örgjörvastýrirökfræði kemur við sögu. Engar breytingar á ályktunarvélinni eru nauðsynlegar.

Vogaflutningsferlið getur nýtt fjórar tegundir vélbúnaðarauðlinda: (1) Gagnatilfærsla á hýsil og tæki (2) GPU-útreikningar fyrir vörpunarsamruna og magnbundningu (3) RDMA-netkerfi fyrir gagnaplan (4) Ethernet fyrir stjórnplan.

Hönnun okkar skiptir vogaflutningsverkefnum í leiðslustig, sem gerir kleift að láta mismunandi vélbúnaðarauðlindir skarast auðveldlega.

Sumar útfærslur reikna flutningsáætlun upp á nýtt í hverju þjálfunarskrefi, safna lýsigögnum ítrekað og dreifa leiðbeiningum. Þetta bætir við óþarfa seinkun á stjórnplani.

Áætlun okkar er reiknuð út einu sinni við frumstillingu. Hver þjálfunarítrun endurspilar einfaldlega áætluninni: stjórnandinn gefur „áfram“-merki og GPU-arnir fylgja fyrirfram úthlutuðum leiðum. Framkvæmdin er fyrirsjáanleg og létt.

Það er freistandi að flækja allt vogauppfærsluferlið í einu einliða falli: safna lýsigögnum, nafnsamsvörun, söfnun innan hnúta, vörpunarsamruni, magnbundning, skipting samskiptaheims, netflutningur milli hnúta. Það er erfitt að forrita rétt og enn erfiðara að fínstilla.

Í verkfræði okkar skiljum við þessi skref að sem stökum hlutum. Hægt er að einingarprófa hvern hlut, rökræða um hann og fínstilla hann sjálfstætt.

Hraður og áreiðanlegur vogaflutningur er mikilvæg byggingareining fyrir stórtæka fínstillingu með styrktarnámi. Með því að sameina RDMA WRITE grunnaðgerðina, kyrrstæða flutningsáætlun og leiðslukeyrslu færðum við uppfærslur á billjón breytum niður í aðeins 1,3 sekúndur á Kimi-K2. Nálgunin er auðveld í röksemdafærslu, einföld í viðhaldi og kemur í veg fyrir flöskuhálsa hefðbundinnar hönnunar.

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: einhliða flutningar

Hátt stigs verkflæði

Framkvæmd vogaflutnings

DeviceMesh og Mesh-hópar

Verkefnaleiðsla

Stjórnun á notkun GPU-minnis

Af hverju þetta er hratt og einfalt

Jafningjasamskipti

Einhliða gagnaflutningur

Leiðsla

Kyrrstæð áætlun

Hreinn aðskilnaður

Niðurstaða

Tilvitnun