Prenos težina za RL post-obuku za manje od 2 sekunde

Nedavno smo ostvarili ažuriranje parametara između mašina za 1,3 sekunde za model Kimi-K2 (1T parametara), prenoseći težine sa 256 GPU-ova za obuku (BF16) na 128 GPU-ova za inferenciju (FP8).

U asinhronom finjunesingu potkrepljenom učenje potkrepljenjem (reinforcement learning fine-tuning), obuka i inferencija se izvršavaju na odvojenim GPU-ovima. Nakon svakog koraka obuke, nove težine se moraju proslediti čvorovima za inferenciju. Mnogim postojećim okruženjima je potrebno nekoliko sekundi — ili čak minuta — za modele sa bilionima parametara.

Korišćenjem RDMA tačka-u-tačku komunikacije, u mogućnosti smo da prenos težina učinimo izuzetno brzim, bez izmene mašine za inferenciju, čime kod postaje lakši za pisanje i održavanje.

Naše rešenje je zasnovano na RDMA WRITE-u, jednostranom primitivu gde izvor direktno upisuje podatke 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 dobiti obaveštenje o prenosu. Ovo nam omogućava nisku latenciju, visok protok i prenose bez kopiranja koje pokreću čvorovi za obuku bez ikakve kontrolne logike na čvorovima za inferenciju.

Prikupljanje metapodataka – Kontrolor prikuplja metapodatke o parametrima sa svih GPU-ova za obuku i inferenciju.

Izračunavanje rasporeda – Kontrolor izračunava statički raspored prenosa težina, mapirajući koji GPU za obuku šalje koji parametar kom GPU-u za inferenciju i po kom redosledu.

Distribucija rasporeda – Kontrolor šalje raspored svim GPU-ovima za obuku.

Izvršavanje – Nakon svakog koraka obuke, kontrolor signalizira GPU-ovima za obuku da započnu prenose.

Sa definisanim tokom posla na visokom nivou, glavni izazov je kako efikasno izvršiti prenose težina na skali od bilion parametara. Ovde opisujemo detalje putanje izvršavanja.

Parametri u obuci su distribuirani u skladu sa FSDP rasporedima. Korišćenjem full_tensor() funkcije, svi GPU-ovi u okviru DeviceMesh mogu da rekonstruišu pun parametar, čime svi mogu služiti kao izvor za prenos težina.

Više disjunktnih DeviceMesh-ova formira mrežnu grupu (mesh group). Pošto su DeviceMesh-ovi u istoj grupi disjunktni, njihovi prenosi se ne mešaju i mogu se izvršavati u potpunosti paralelno. Između mrežnih grupa uvodimo globalnu barijeru radi obezbeđivanja redosleda.

Cevovod za prenos težine preklapa kopiranje sa hosta na uređaj, sklapanje tenzora, fuziju i kvantizaciju, kao i RDMA prenose preko MeshGroups.

Prenos svakog tenzora parametara tretiramo kao zadatak. Proces prenosa težina koristi više vrsta hardverskih izvora, stoga zadatak prenosa težina delimo na različite faze cevovoda koje se vremenski preklapaju:

Kopiranje memorije sa domaćina na uređaj (Host to device memcpy) — Ako FSDP rasterećuje težinu na CPU

Priprema parametara — Rekonstrukcija pune težine pomoću full_tensor(), primena fuzije projekcije, kvantizacija ako je potrebna.

RDMA prenos — Upisivanje bez kopiranja u udaljenu memoriju GPU-a za inferenciju

Globalna barijera — Nakon što se završe svi pozivi funkcije full_tensor(), vrši se sinhronizacija između mrežnih grupa korišćenjem GLOO-a preko Etherneta.

U implementaciji održavamo FIFO red zadataka za svaku fazu cevovoda. Kada zadatak na čelu reda završi fazu, premešta se na začelje reda sledeće faze.

full_tensor() i druge GPU operacije uvode dodatnu potrošnju GPU memorije. Da bismo izbegli grešku nedostatka memorije (out of memory), izvršavanje zadatka započinjemo samo ako trenutni aktivni zadaci zauzimaju manje privremene GPU memorije od podesivog praga (watermark).

Nekoliko dizajnerskih odluka čini naš sistem znatno bržim za pokretanje i lakšim za održavanje u poređenju sa uobičajenim rešenjima otvorenog koda.

Prenos težine zasnovan na Rank0 prvo prikuplja težine, a zatim ih distribuira inferentnom čvoru.

Komunikacija od tačke do tačke prenosi težine sa oba trening čvora na mrežne interfejse inferentnog čvora.

Uobičajeni šablon je usmeravanje svih parametara kroz GPU-ove sa rangom 0: prikupljanje na rangu 0 za obuku, slanje na rang 0 za inferenciju, a zatim ponovno rasipanje. Ovo brzo postaje usko grlo, ograničeno propusnim opsegom PCIe magistrale i mrežne kartice (NIC) jednog GPU-a (npr. 400 Gbps ≈ 50 GB/s).

Nasuprot tome, naša postavka tačka-u-tačku omogućava svakom GPU-u za obuku da šalje podatke direktno svakom GPU-u za inferenciju, čime se zasićuje celokupna mrežna struktura umesto samo jedne veze.

Prenos zasnovan na RPC-u koristi bafer na strani inferencije i pozive radnika, dok jednostrani prenos upisuje težine direktno, bez učešća inferentnog radnika.

Neki sistemi se oslanjaju na pozivanje metode update_weight() mašine za inferenciju za svaki tenzor. To podrazumeva invazivne izmene u kodu za inferenciju, kao i opterećenje usled RPC poziva, serijalizacije i koordinacije kontrolne ravni.

Sa RDMA WRITE primitivom, težine ažuriramo neprimetno u memoriji GPU-a za inferenciju, bez dodatnih kopiranja. Nisu uključene poruke kontrolne ravni niti kontrolna logika CPU-a. Nisu potrebne nikakve izmene na mašini za inferenciju.

Proces prenosa težina može da koristi četiri vrste hardverskih resursa: (1) Kretanje podataka između domaćina i uređaja (2) GPU izračunavanja za fuziju projekcije i kvantizaciju (3) RDMA mrežu za ravan podataka (4) Ethernet za kontrolnu ravan.

Naš dizajn deli zadatke prenosa težina u faze cevovoda, omogućavajući lako preklapanje između različitih hardverskih resursa.

Neke implementacije iznova izračunavaju raspored prenosa u svakom koraku obuke, uzastopno prikupljajući metapodatke i distribuirajući instrukcije. Ovo dodaje nepotrebnu latenciju kontrolne ravni.

Naš raspored se izračunava jednom, prilikom inicijalizacije. Svaka iteracija obuke jednostavno ponovo pokreće plan: kontrolor izdaje signal „kreni“, a GPU-ovi prate svoje unapred dodeljene rute. Izvršavanje je predvidljivo i lako.

Privlačno je celokupan proces ažuriranja težina upakovati u jednu monolitnu funkciju: prikupljanje metapodataka, uparivanje imena, prikupljanje unutar čvora, fuzija projekcije, kvantizacija, deljenje komunikacionog sveta, prenos mrežom između čvorova. Ovakav pristup je težak za ispravno programiranje, a još teži za optimizaciju.

U našem inženjerskom pristupu, ove korake odvajamo kao pojedinačne komponente. Svaka komponenta se može pojedinačno testirati, analizirati i optimizovati izolovano.

Brz i pouzdan prenos težina predstavlja ključnu gradivnu jedinicu za finjunesing RL-a velikih razmera. Kombinovanjem RDMA WRITE primitiva, statičkog rasporeda prenosa i izvršavanja u cevovodu, sveli smo ažuriranja modela sa bilionima parametara na samo 1,3 sekunde na Kimi-K2 modelu. Pristup je jednostavan za razumevanje, lako se održava i izbegava 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: jednostrani prenosi

Tok posla na visokom nivou

Izvršavanje prenosa težina

DeviceMesh i mrežne grupe (Mesh Groups)

Cevovod zadataka (Task pipeline)

Kontrola korišćenja GPU memorije

Zašto je ovo rešenje brzo i jednostavno

Komunikacija tačka-u-tačku

Jednostrani prenos podataka

Uvođenje cevovoda (Pipelining)

Statički raspored

Jasno odvajanje

Zaključak

Citiranje