Svaru pārsūtīšana RL pēcapmācībai mazāk nekā 2 sekundēs

Mēs nesen panācām 1,3 sekunžu starpdatoru parametru atjaunināšanu Kimi-K2 modelim (1 triljons parametru), pārsūtot svarus no 256 apmācības GPU (BF16) uz 128 inferences GPU (FP8).

Asinhronās pastiprināšanas mācīšanās (reinforcement learning) smalknoregulēšanā apmācība un inference darbojas uz atsevišķiem GPU. Pēc katra apmācības solīša jauni svari ir jānosūta uz inferences mezgliem. Daudzām esošajām sistēmām modelim ar triljoniem parametru tas aizņem vairākas sekundes vai pat minūtes.

Izmantojot RDMA punktu-punktam sakarus, mēs varam padarīt svaru pārsūtīšanu neticami ātrāku, nemainot inferences dzinuli, un padarīt kodu vieglāk rakstāmu un uzturamu.

Mūsu risinājums ir balstīts uz RDMA WRITE — vienpusēju primitīvu, kurā avots tieši raksta galamērķa GPU atmiņā.

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

Galamērķa puse par pārsūtīšanu pat saņems paziņojumu. Tas nodrošina zemas aiztures, augstas caurlaidības un nulles kopēšanas (zero-copy) pārsūtīšanu, ko iniciē apmācības mezgli bez jebkādas vadības loģikas inferences mezglos.

Metadatu vākšana — kontrolieris apkopo parametru metadatus no visiem apmācības un inferences GPU.

Grafika aprēķins — kontrolieris aprēķina statisku svaru pārsūtīšanas grafiku, norādot, kurš apmācības GPU nosūta kuru parametru kuram inferences GPU un kādā secībā.

Grafika izplatīšana — kontrolieris nosūta grafiku visiem apmācības GPU.

Izpilde — pēc katra apmācības soļa kontrolieris signalizē apmācības GPU sākt pārsūtīšanu.

Definējot augsta līmeņa darbplūsmu, galvenais izaicinājums ir tas, kā efektīvi izpildīt svaru pārsūtīšanu triljonu parametru mērogā. Šeit mēs aprakstām izpildes ceļa detaļas.

Apmācības parametri tiek sadalīti saskaņā ar FSDP izvietojumu. Izmantojot full_tensor(), visi DeviceMesh esošie GPU var rekonstruēt pilnu parametru, tādējādi visi var kalpot par avotu svaru pārsūtīšanai.

Vairāki nesavienoti DeviceMeshes veido mesh grupu. Tā kā DeviceMeshes vienā grupā ir nesavienoti, to pārsūtīšanas procesi netraucē viens otram un var darboties pilnībā paralēli. Starp mesh grupām mēs ieviešam globālu barjeru secības nodrošināšanai.

Svara pārneses konveijers pārklāj resursdatora–ierīces kopēšanu, tenzoru salikšanu, sapludināšanu un kvantizāciju, kā arī RDMA pārsūtījumus starp MeshGroups.

Mēs katra parametru tenzora pārsūtīšanu uztveram kā uzdevumu. Svaru pārsūtīšanas process izmanto vairāku veidu aparatūras avotus, tāpēc mēs sadalām svaru pārsūtīšanas uzdevumu dažādos konveijera posmos, kas laikā pārklājas:

Resursdatora (Host) uz ierīces memcpy — ja FSDP pārvieto svaru uz CPU

Parametru sagatavošana — rekonstruēt pilnu svaru ar full_tensor(), lietot projicēšanas saplūšanu, nepieciešamības gadījumā veikt kvantēšanu.

RDMA pārsūtīšana — nulles kopēšanas rakstīšana attālā inferences GPU atmiņā

Globālā barjera — kad visi full_tensor() izsaukumi ir pabeigti, sinhronizēt starp mesh grupām, izmantojot GLOO caur Ethernet.

Ieviešanā mēs uzturam FIFO uzdevumu rindu katram konveijera posmam. Kad rindas sākuma uzdevums pabeidz posmu, tas tiek pārvietots uz nākamā posma rindas beigām.

full_tensor() un citas GPU operācijas rada papildu GPU atmiņas patēriņu. Lai izvairītos no atmiņas izsīkuma kļūdas (out of memory), mēs sākam uzdevuma izpildi tikai tad, ja pašreizējie izpildē esošie uzdevumi aizņem mazāk pagaidu GPU atmiņas nekā konfigurējamā robežvērtība (watermark).

Vairākas dizaina izvēles padara mūsu sistēmu ievērojami ātrāku un vieglāk uzturamu nekā parastos atvērtā koda risinājumus.

Ar Rank0 balstīta svaru pārnese savāc svarus pirms to izplatīšanas uz inferenču mezglu.

No punkta uz punktu komunikācija pārsūta svarus no abiem apmācības mezgliem uz inferenču mezgla tīkla saskarnēm.

Izplatīts modelis ir virzīt visus parametrus caur 0. ranga GPU: apkopot 0. ranga apmācības mezglā, nosūtīt uz 0. ranga inferences mezglu un pēc tam atkal izkliedēt. Tas ātri kļūst par šauro vietu, ko ierobežo viena GPU PCIe joslas platums un NIC (piem., 400 Gbps ≈ 50 GB/s).

Turpretim mūsu punktu-punktam iestatījums ļauj katram apmācības GPU sūtīt datus tieši uz katru inferences GPU, piesātinot visu tīkla struktūru, nevis vienu savienojumu.

Uz RPC balstīta pārsūtīšana izmanto inferenču puses buferi un darbinieku izsaukumus, savukārt vienpusēja pārsūtīšana raksta svarus tieši bez inferenču darbinieka līdzdalības.

Dažas sistēmas paļaujas uz inferences dzinēja metodes update_weight() izsaukšanu katram tenzoram. Tas nozīmē iejaucošas izmaiņas inferences kodā, kā arī RPC, serializācijas un vadības plaknes koordinācijas radītos papildu izdevumus.

Izmantojot RDMA WRITE primitīvu, mēs nemanāmi atjauninām svarus inferences GPU atmiņā bez papildu kopijām. Netiek iesaistīts neviens vadības plaknes ziņojums un neviena CPU vadības loģika. Nav nepieciešams modificēt inferences dzinuli.

Svaru pārsūtīšanas process var izmantot četri veidu aparatūras resursus: (1) Resursdatora un ierīces datu pārvietošanu (2) GPU skaitļošanu projicēšanas saplūšanai un kvantēšanai (3) RDMA tīklu datu plaknei (4) Ethernet vadības plaknei.

Mūsu dizains sadala svaru pārsūtīšanas uzdevumus konveijera posmos, nodrošinot vieglu pārklāšanos starp dažādiem aparatūras resursiem.

Dažas ieviešanas pārrēķina pārsūtīšanas grafiku katrā apmācības solī, atkārtoti vācot metadatus un izplatot instrukcijas. Tas rada nevajadzīgu vadības plaknes aizturi.

Mūsu grafiks tiek aprēķināts vienu reizi inicializācijas laikā. Katra apmācības iterācija vienkārši atkārto plānu: kontrolieris izsniedz “aiziet” signālu, un GPU seko saviem iepriekš piešķirtajiem maršrutiem. Izpilde ir paredzama un viegla.

Ir vilinoši apvienot visu svaru atjaunināšanas procesu vienā monolītā funkcijā: vākt metadatus, nosaukumu saskaņošanu, mezglu iekšējo apkopošanu, projicēšanas saplūšanu, kvantēšanu, komunikācijas pasaules apakšsadalīšanu, starpmezglu tīkla pārsūtīšanu. To ir grūti pareizi ieprogrammēt un vēl grūtāk optimizēt.

Mūsu inženierijā mēs šos soļus nodalām kā atsevišķus komponentus. Katru komponentu var vienības testēt, analizēt un optimizēt atsevišķi.

Ātra, uzticama svaru pārsūtīšana ir kritisks pamatelements liela mēroga RL smalknoregulēšanai. Apvienojot RDMA WRITE primitīvu, statisku pārsūtīšanas grafiku un konveijera izpildi, mēs samazinājām triljonu parametru atjaunināšanu līdz tikai 1,3 sekundēm Kimi-K2 modelim. Šo pieeju ir viegli saprast, viegli uzturēt, un tā novērš tradicionālo dizainu šaurās vietas.

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: vienpusējas pārsūtīšanas

Augsta līmeņa darbplūsma

Svaru pārsūtīšanas izpilde

DeviceMesh un Mesh grupas

Uzdevumu konveijers

GPU atmiņas izmantošanas kontrole

Kāpēc tas ir ātri un vienkārši

Punktu-punktam sakari

Vienpusēja datu pārsūtīšana

Konveijera veidošana (Pipelining)

Statisks grafiks

Tīra nodalīšana

Secinājums

Citāts