Prenos váh pre dodatočný tréning RL za menej ako 2 sekundy

Nedávno sa nám podarilo dosiahnuť aktualizáciu parametrov medzi strojmi za 1,3 sekundy pre model Kimi-K2 (1T parametrov), pričom sme prenášali váhy z 256 tréningových GPU (BF16) na 128 inferenčných GPU (FP8).

Pri asynchrónnom dolaďovaní pomocou posilňovania (reinforcement learning fine-tuning) prebieha trénovanie a inferencia na samostatných GPU. Po každom tréningovom kroku je nutné poslať nové váhy do inferenčných uzlov. Mnoho existujúcich rámcov potrebuje pri modeloch s biliónmi parametrov niekoľko sekúnd – alebo dokonca minút.

Využitím bod-bodovej komunikácie RDMA dokážeme zabezpečiť bleskový prenos váh bez nutnosti zmeny inferenčného enginu a zároveň zjednodušiť písanie a údržbu kódu.

Naše riešenie je postavené na príkaze RDMA WRITE, čo je jednostranný primitív, pri ktorom zdroj zapisuje priamo do pamäte GPU cieľového zariadenia.

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

Cieľová strana dokonca ani nedostane upozornenie o prenose. To nám prináša nízku latenciu, vysokú priepustnosť a prenosy bez kopírovania (zero-copy), ktoré sú riadené tréningovými uzlami bez akejkoľvek riadiacej logiky na inferenčných uzloch.

Zber metadát – Kontrolér zhromažďuje metadáta parametrov zo všetkých tréningových a inferenčných GPU.

Výpočet plánu – Kontrolér vypočíta statický plán prenosu váh, pričom určuje, ktorá tréningová GPU posiela aký parameter ktorej inferenčnej GPU a v akom poradí.

Distribúcia plánu – Kontrolér odošle plán do všetkých tréningových GPU.

Vykonanie – Po každom tréningovom kroku dá kontrolér signál tréningovým GPU na spustenie prenosov.

Po definovaní vysokoúrovňového pracovného postupu je kľúčovou výzvou efektívne vykonanie prenosu váh v mierke biliónov parametrov. Tu popisujeme podrobnosti cesty vykonávania.

Parametre pri trénovaní sú distribuované podľa umiestnení FSDP. Pomocou funkcie full_tensor() môžu všetky GPU v štruktúre DeviceMesh zrekonštruovať celý parameter, takže všetky môžu slúžiť ako zdroj pre prenos váh.

Viacero disjunktných štruktúr DeviceMesh tvoria skupinu mesh. Keďže štruktúry DeviceMesh v rovnakej skupine sú disjunktné, ich prenosy sa navzájom nerušia a môžu prebiehať plne paralelne. Medzi skupinami mesh vkladáme globálnu bariéru na vynútenie poradia.

Pipeline prenosu váh sa prekrýva s kopírovaním z hostiteľa do zariadenia, zostavovaním tensorov, fúziou a kvantizáciou a s prenosmi RDMA naprieč skupinami MeshGroup.

Prenos každého tenzora parametrov považujeme za úlohu. Proces prenosu váh využíva viacero typov hardvérových zdrojov, preto rozdeľujeme úlohu prenosu váh do rôznych štádií rúry, ktoré sa časovo prekrývajú:

Kopírovanie z hostiteľa do zariadenia (Host to device memcpy) — Ak FSDP presúva váhu do CPU

Príprava parametrov — Rekonštrukcia úplnej váhy pomocou full_tensor(), použitie fúzie projekcií, v prípade potreby kvantizácia.

Prenos cez RDMA — Zápis bez kopírovania do vzdialenej pamäte inferenčnej GPU

Globálna bariéra — Po dokončení všetkých volaní full_tensor() prebehne synchronizácia naprieč skupinami mesh pomocou GLOO cez Ethernet.

Pri implementácii udržiavame FIFO rad úloh pre každé štádium rúry. Keď úloha na začiatku radu dokončí dané štádium, presunie sa na koniec radu nasledujúceho štádia.

Funkcia full_tensor() a ďalšie operácie GPU spôsobujú dodatočné využitie pamäte GPU. Aby sme predišli chybe nedostatku pamäte (out of memory), spúšťame vykonávanie úlohy iba vtedy, ak aktuálne prebiehajúce úlohy zaberajú menej dočasnej pamäte GPU, než je konfigurovateľná prahová hodnota (watermark).

Vďaka niekoľkým návrhovým rozhodnutiam je náš systém výrazne rýchlejší na spúšťanie a jednoduchší na údržbu než bežné open-source riešenia.

Prenos váh založený na Rank0 zhromažďuje váhy pred ich distribúciou do inferenčného uzla.

Komunikácia typu point-to-point prenáša váhy z oboch tréningových uzlov do sieťových rozhraní inferenčného uzla.

Bežným vzorom je smerovanie všetkých parametrov cez GPU s hodnosťou 0 (rank-0): zhromaždenie na tréningovej GPU rank-0, odoslanie na inferenčnú GPU rank-0 a následné opätovné rozptýlenie. To sa rýchlo stáva úzkym hrdlom, ktoré je obmedzené šírkou pásma PCIe a sieťovej karty jednej GPU (napr. 400 Gbps ≈ 50 GB/s).

Naša bod-bodová štruktúra naopak umožňuje každému tréningovému GPU posielať dáta priamo každej inferenčnej GPU, čím sa vyťaží celá sieťová štruktúra namiesto jediného prepojenia.

Prenos založený na RPC používa buffer na strane inferencie a volania workerov, zatiaľ čo jednosmerný prenos zapisuje váhy priamo bez účasti inferenčného workeru.

Niektoré systémy sa spoliehajú na volanie metódy update_weight() inferenčného enginu pre každý tenzor. To znamená invazívne zmeny v inferenčnom kóde, dodatočnú réžiu z RPC, serializáciu a koordináciu riadiacej roviny.

S primitívom RDMA WRITE aktualizujeme váhy v pamäti inferenčnej GPU ticho a bez dodatočných kópií. Nezapočítava sa žiadna správa riadiacej roviny ani riadiaca logika CPU. Nevyžadujú sa žiadne úpravy inferenčného enginu.

Proces prenosu váh môže využívať štyri typy hardvérových zdrojov: (1) presun dát medzi hostiteľom a zariadením, (2) výpočty GPU pre fúziu a kvantizáciu projekcií, (3) sieť RDMA pre dátovú rovinu, (4) Ethernet pre riadiacu rovinu.

Náš návrh rozdeľuje úlohy prenosu váh do štádií rúry, čo umožňuje jednoduché prekrývanie naprieč rôznymi hardvérovými zdrojmi.

Niektoré implementácie prepočítavajú plán prenosu pri každom tréningovom kroku, pričom opakovane zbierajú metadáta a distribuujú inštrukcie. To pridáva zbytočnú latenciu riadiacej roviny.

Náš plán sa vypočíta raz pri inicializácii. Každá tréningová iterácia jednoducho opakuje plán: kontrolér vydá signál „go“ a GPU nasledujú svoje vopred priradené trasy. Vykonávanie je predvídateľné a nenáročné na zdroje.

Je lákavé spojiť celý proces aktualizácie váh do jednej monolitickej funkcie: zber metadát, párovanie názvov, zhromažďovanie v rámci uzla, fúzia projekcií, kvantizácia, rozdelenie komunikačného sveta a sieťový prenos medzi uzlami. Je ťažké naprogramovať to správne a ešte ťažšie optimalizovať.

V našom inžinierstve oddeľujeme tieto kroky ako samostatné komponenty. Každý komponent možno unit testovať, analyzovať a optimalizovať izolovane.

Rýchly a spoľahlivý prenos váh je kľúčovým stavebným kameňom pre rozsiahle dolaďovanie pomocou posilňovania. Kombináciou primitívu RDMA WRITE, statického plánu prenosu a rúrového vykonávania sme znížili aktualizácie pre bilióny parametrov na iba 1,3 sekundy na modeli Kimi-K2. Prístup je jednoduchý na pochopenie, ľahký na údržbu a vyhýba sa úzkym hrdlám tradičných návrhov.

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: jednostranné prenosy

Vysokoúrovňový pracovný postup

Vykonanie prenosu váh

DeviceMesh a skupiny mesh (Mesh Groups)

Rúra úloh (Task pipeline)

Riadenie využitia pamäte GPU

Prečo je to rýchle a jednoduché

Bod-bodová komunikácia (Point-to-point communication)

Jednostranný prenos dát

Rúrové spracovanie (Pipelining)

Statický plán

Čisté oddelenie

Záver

Citácia