Kaalude edastamine tugevdusõppe järelkoolituseks alla 2 sekundiga

Saavutasime hiljuti Kimi-K2 (1T parameetrit) jaoks 1,3-sekundilised masinatevahelised parameetrite uuendused, edastades kaale 256-lt treeningu GPU-lt (BF16) 128-le järelduse GPU-le (FP8).

Asünkroonses tugevdusõppe häälestuses toimuvad treening ja järeldus eraldi GPU-del. Pärast iga treeningusammu tuleb uued kaalud järeldusesõlmedesse lükata. Paljud olemasolevad raamistikud võtavad triljoni parameetriga mudelite puhul mitu sekundit või isegi minuteid.

Kasutades RDMA punkt-punkt-side, suudame muuta kaalude edastamise ülikiireks ilma järeldusmootorit muutmata ning koodi kirjutamise ja hooldamise lihtsamaks.

Meie lahendus põhineb RDMA WRITE-il, ühepoolsel primitiivil, kus lähtekoht kirjutab otse sihtkoha GPU mällu.

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

Sihtpool ei saa edastusest isegi teavitust. See tagab meile madala latentsusega, suure läbilaskvusega ja nullkopeerimisega ülekanded, mida juhivad treeningsõlmed ilma järeldusesõlmede juhtloogikata.

Metainfo kogumine – Kontroller kogub parameetrite metainfo kõigilt treeningu- ja järelduse GPU-delt.

Plaani arvutamine – Kontroller arvutab staatilise kaalude edastusplaani, kaardistades, milline treeningu GPU saadab millise parameetri millisele järelduse GPU-le ja millises järjekorras.

Plaani jagamine – Kontroller saadab plaani kõigile treeningu GPU-dele.

Täitmine – Pärast iga treeningusammu annab kontroller treeningu GPU-dele signaali ülekannete alustamiseks.

Kõrgetasemelise töövoo määratlemisel on peamiseks väljakutseks see, kuidas teostada kaalude edastamist tõhusalt triljonite parameetrite ulatuses. Siin kirjeldame täitmistee üksikasju.

Treeningu parameetrid jagatakse FSDP paigutuste järgi. Kasutades funktsiooni full_tensor(), saavad kõik DeviceMesh-i GPU-d täieliku parameetri taastada, seega saavad kõik toimida kaalude edastamise allikana.

Mitmed eraldiseisvad DeviceMeshid moodustavad võrgurühma. Kuna sama rühma DeviceMeshid on üksteisest eraldatud, ei sega nende ülekanded üksteist ja need saavad toimuda täielikult paralleelselt. Võrgurühmade vahele lisame järjestuse tagamiseks globaalse tõkke.

Kaalu ülekande torujuhe kattub hostilt seadmele koopiate, tensorite koostamise, liitmise ja kvantimise ning RDMA ülekannetega üle MeshGroupide.

Käsitleme iga parameetri tensori edastamist ülesandena. Kaalude edastamise protsess kasutab mitut tüüpi riistvararessursse, mistõttu jagame kaalude edastamise ülesande erinevateks konveierietappideks, mis ajaliselt kattuvad:

Hostist seadmesse mälu kopeerimine (Host to device memcpy) – kui FSDP laadib kaalu CPU-le

Parameetrite ettevalmistamine – Taasta täielik kaal funktsiooniga full_tensor(), rakenda projektsiooni fusioon, vajadusel kvandigeeri.

RDMA edastus – Nullkopeerimisega kirjutamine kaugjärelduse GPU mällu

Globaalne tõke – Kui kõik full_tensor() kutsed on tehtud, sünkroniseeri võrgurühmade vahel Etherneti kaudu GLOO abil.

Rakendamises säilitame iga konveierietapi jaoks FIFO-ülesannete järjekorra. Kui järjekorra alguses olev ülesanne lõpetab etapi, liigutatakse see järgmise etapi järjekorra lõppu.

Funktsioon full_tensor() ja muud GPU toimingud toovad kaasa täiendava GPU mälukasutuse. Mälutõrke (out of memory) vältimiseks alustame ülesande täitmist ainult siis, kui praegused käimasolevad ülesanded hõivavad vähem ajutist GPU mälu kui konfigureeritav veepiir.

Mitmed disainivalikud muudavad meie süsteemi käitamise oluliselt kiiremaks ja hooldamise lihtsamaks kui tavalised avatud lähtekoodiga lahendused.

Rank0-põhine kaalu ülekanne kogub kaalud enne nende inference-sõlmele jaotamist.

Punktist punkti side kannab kaalud mõlema treeningsõlme võrguliidestelt inference-sõlmele.

Tavaline muster on suunata kõik parameetrid läbi 0. järgu (rank-0) GPU-de: koguda treeningu 0. järgul, saata järelduse 0. järgule ja seejärel uuesti hajutada. Sellest saab kiiresti pudelikael, mida piirab ühe GPU PCIe ribalaius ja NIC (nt 400 Gbps ≈ 50 GB/s).

Seevastu meie punkt-punkt-häälestus võimaldab igal treeningu GPU-l saata andmeid otse igale järelduse GPU-le, küllastades kogu võrgustiku, mitte ainult ühte ühendust.

RPC-põhine ülekanne kasutab inference-poole bufferit ja töötaja kutseid, samas kui ühepoolne ülekanne kirjutab kaalud otse, ilma inference-töötaja osaluseta.

Mõned süsteemid tuginevad iga tensori puhul järeldusmootori meetodi update_weight() väljakutsumisele. See tähendab sekkuvaid muudatusi järelduskoodis ning lisakoormust RPC-dest, jadastamisest ja juhtploki koordineerimisest.

RDMA WRITE primitiiviga uuendame kaale järelduse GPU mälus vaikselt ja ilma täiendavate koopiateta. Juhtploki sõnumeid ega CPU juhtloogikat pole vaja. Järeldusmootori muutmine ei ole vajalik.

Kaalude edastamise protsess saab kasutada nelja tüüpi riistvararessursse: (1) Host-seade andmeliikumine (2) GPU arvutus projektsiooni fusiooni ja kvandimise jaoks (3) RDMA võrk andmetasandi jaoks (4) Ethernet juhtploki jaoks.

Meie disain jagas kaalude edastamise ülesanded konveierietappideks, võimaldades eri riistvararessursside hõlpsat kattumist.

Mõned teostused arvutavad edastusplaani uuesti igal treeningu sammul, kogudes korduvalt metainfot ja jagades juhiseid. See lisab tarbetut juhtploki latentsust.

Meie plaan arvutatakse üks kord initsialiseerimisel. Iga treeningu iteratsioon lihtsalt kordab plaani: kontroller annab „mine“-signaali ja GPU-d järgivad neile eelnevalt määratud marsruute. Täitmine on prognoositav ja kergekaaluline.

On kiusatus siduda kogu kaalude uuendamise protsess ühte monoliitsesse funktsiooni: metainfo kogumine, nimesobitamine, sõlmesisene kogumine, projektsiooni fusioon, kvantimine, kommunikatsioonimaailma alamtükeldamine, sõlmedevaheline võrguülekkanne. Seda on raske õigesti programmeerida ja veelgi raskem optimeerida.

Meie inseneritöös eraldame need sammud eraldi komponentidena. Iga komponenti saab eraldi testida, analüüsida ja optimeerida.

Kiire ja usaldusväärne kaalude edastamine on suuremahulise tugevdusõppe häälestuse kriitiline ehituskivi. Kombineerides RDMA WRITE primitiivi, staatilise edastusplaani ja konveiertöötluse, vähendasime Kimi-K2 triljoni parameetriga uuendused vaid 1,3 sekundile. Lähenemist on lihtne mõista ja hooldada ning see hoiab ära traditsioonilise disaini pudelikaelad.

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: ühepoolsed edastused

Kõrgetasemeline töövoog

Kaalude edastamise täitmine

DeviceMesh ja võrgurühmad (Mesh Groups)

Ülesannete konveier (Task pipeline)

GPU mälukasutuse juhtimine

Miks see on kiire ja lihtne

Punkt-punkt-side

Ühepoolne andmeedastus

Konveiertöötlus

Staatiline plaan

Puhas eraldatus

Kokkuvõte

Tsitaat