Prijenos težina za RL post-trening za manje od 2 sekunde

Nedavno smo postigli ažuriranje parametara između strojeva za 1,3 sekunde za model Kimi-K2 (1T parametara), prenoseći težine s 256 GPU-a za treniranje (BF16) na 128 GPU-a za inferenciju (FP8).

Kod fine-tuninga uz potpomognuto učenje putem asinkrone povratne sprege, treniranje i inferencija pokreću se na zasebnim GPU-ovima. Nakon svakog koraka treniranja, nove se težine moraju poslati na čvorove za inferenciju. Mnogim postojećim okvirima potrebno je nekoliko sekundi – ili čak minuta – za modele s bilijunima parametara.

Korištenjem RDMA komunikacije točka‑do‑točka možemo prijenos težine učiniti iznimno brzim, bez promjene inference enginea, te olakšati pisanje i održavanje koda.

Naše je rješenje izgrađeno na RDMA WRITE, jednostranom primitku gdje izvor izravno piše u memoriju GPU-a odredišta.

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

Odredišna strana neće čak ni primiti obavijest o prijenosu. To nam omogućuje prijenose niske latencije, visoke propusnosti i nula kopiranja koje pokreću čvorovi za treniranje bez ikakve upravljačke logike na čvorovima za inferenciju.

Prikupljanje metapodataka – Upravljač prikuplja metadatske parametre sa svih GPU-ova za treniranje i inferenciju.

Izračun rasporeda – Upravljač izračunava statički raspored prijenosa težina, mapirajući koji GPU za treniranje šalje koji parametar kojem GPU-u za inferenciju i kojim redoslijedom.

Distribucija rasporeda – Upravljač šalje raspored svim GPU-ovima za treniranje.

Izvođenje – Nakon svakog koraka treniranja, upravljač signalizira GPU-ovima za treniranje da započnu prijenose.

Nakon definiranja tijeka rada na visokoj razini, ključni je izazov kako učinkovito izvršiti prijenose težina na razini bilijuna parametara. Ovdje opisujemo detalje staze izvođenja.

Parametri u treniranju distribuirani su prema FSDP smještajima. Korištenjem funkcije full_tensor(), svi GPU-ovi u DeviceMesh mogu rekonstruirati cijeli parametar, stoga svi mogu služiti kao izvor za prijenos težina.

Više disjunktnih DeviceMesh struktura čini mesh grupu. Budući da su DeviceMesh strukture u istoj grupi disjunktne, njihovi prijenosi ne ometaju jedni druge i mogu se izvoditi u potpunosti paralelno. Između mesh grupa umetamo globalnu barijeru kako bismo nametnuli redoslijed.

Cjevovod prijenosa težina preklapa kopije s hosta na uređaj, sastavljanje tenzora, fuziju i kvantizaciju te RDMA prijenose kroz MeshGroups.

Prijenos svakog tenzora parametara tretiramo kao zadatak. Proces prijenosa težina koristi više vrsta hardverskih izvora, stoga zadatak prijenosa težina dijelimo na različite faze cjevovoda koje se vremenski preklapaju:

Memcpy s domaćina na uređaj — Ako FSDP prebacuje težinu na CPU

Priprema parametara — Rekonstruirajte potpunu težinu pomoću full_tensor(), primijenite fuziju projekcije, kvantizirajte ako je potrebno.

RDMA prijenos — Pisanje bez kopiranja u udaljenu memoriju GPU-a za inferenciju

Globalna barijera — Nakon što su svi pozivi full_tensor() dovršeni, sinkronizirajte se preko mesh grupa koristeći GLOO putem Etherneta.

U implementaciji održavamo FIFO red zadataka za svaku fazu cjevovoda. Kad god zadatak na čelu reda završi fazu, premješta se na začelje reda sljedeće faze.

Funkcija full_tensor() i druge operacije GPU-a uvode dodatnu upotrebu GPU memorije. Kako bismo izbjegli pogrešku nedostatka memorije (out of memory), pokrećemo izvršavanje zadatka samo ako trenutni zadaci u tijeku zauzimaju manje privremene GPU memorije od konfigurirane granice.

Nekoliko dizajnerskih odluka čini naš sustav znatno bržim za rad i lakšim za održavanje od uobičajenih rješenja otvorenog koda.

Prijenos težina temeljen na Rank0 prikuplja težine prije njihove distribucije na inferencijski čvor.

Komunikacija točka‑točka prenosi težine s oba čvora za treniranje na mrežna sučelja inferencijskog čvora.

Uobičajeni uzorak je usmjeravanje svih parametara kroz GPU-ove ranga 0: prikupljanje na rangu 0 za treniranje, slanje na rang 0 za inferenciju, a zatim ponovno raspršivanje. To brzo postaje usko grlo, ograničeno PCIe propusnošću jednog GPU-a i mrežnom karticom (npr. 400 Gbps ≈ 50 GB/s).

Nasuprot tome, naša postavka točka-do-točke omogućuje svakom GPU-u za treniranje da šalje izravno svakom GPU-u za inferenciju, zasićujući cjelokupnu mrežnu strukturu umjesto jedne veze.

Prijenos temeljen na RPC-u koristi međuspremnik na strani inferencije i pozive radnika, dok jednostrani prijenos zapisuje težine izravno bez sudjelovanja radnika za inferenciju.

Neki sustavi se oslanjaju na pozivanje metode update_weight() pogona za inferenciju za svaki tenzor. To znači invazivne promjene u kodu za inferenciju, plus opterećenje od RPC-ova, serijalizacije i koordinacije kontrolne ravnine.

S RDMA WRITE primitkom, tiho ažuriramo težine u memoriji GPU-a za inferenciju, bez dodatnih kopiranja. Nema poruke kontrolne ravnine i nije uključena nikakva upravljačka logika CPU-a. Nije potrebna nikakva modifikacija pogona za inferenciju.

Proces prijenosa težina može iskoristiti četiri vrste hardverskih resursa: (1) Kretanje podataka između domaćina i uređaja (2) Računanje GPU-a za fuziju projekcije i kvantizaciju (3) RDMA mrežu za ravninu podataka (4) Ethernet za upravljačku ravninu.

Naš dizajn dijeli zadatke prijenosa težina u faze cjevovoda, omogućujući jednostavno preklapanje preko različitih hardverskih resursa.

Neke implementacije ponovno izračunavaju raspored prijenosa u svakom koraku treniranja, opetovano prikupljajući metadatke i distribuirajući upute. To dodaje nepotrebnu latenciju kontrolne ravnine.

Naš se raspored izračunava jednom prilikom inicijalizacije. Svaka iteracija treniranja jednostavno ponavlja plan: upravljač izdaje signal "kreni", a GPU-ovi slijede svoje unaprijed dodijeljene rute. Izvršavanje je predvidljivo i lagano.

Mamljivo je zapetljati cijeli proces ažuriranja težina u jednu monolitnu funkciju: prikupljanje metadata, usklađivanje imena, prikupljanje unutar čvora, fuzija projekcije, kvantizacija, podjela komunikacijskog svijeta, mrežni prijenos između čvorova. Teško je ispravno programirati, a još teže optimizirati.

U našem inženjeringu odvajamo ove korake kao pojedinačne komponente. Svaka se komponenta može zasebno testirati, analizirati i optimizirati izolovano.

Brz i pouzdan prijenos težina ključni je građevni blok za fine-tuning uz potpomognuto učenje velikih razmjera. Kombinacijom RDMA WRITE primitka, statičkog rasporeda prijenosa i izvršavanja u cjevovodu, sveli smo ažuriranja modela s bilijunima parametara na samo 1,3 sekunde na Kimi-K2 modelu. Pristup je jednostavan za razumijevanje, lako se održava i izbjegava uska grla tradicionalnih dizajna.

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: jednolični prijenosi

Tijek rada na visokoj razini

Izvođenje prijenosa težina

DeviceMesh i Mesh grupe

Cjevovod zadataka

Upravljanje upotrebom memorije GPU-a

Zašto je brzo i jednostavno

Komunikacija točka-do-točke

Jednostrani prijenos podataka

Uvođenje cjevovoda (Pipelining)

Statički raspored

Čisto odvajanje

Zaključak

Citiranje