Přenos vah pro RL post-training za méně než 2 sekundy

Nedávno jsme dosáhli aktualizací parametrů mezi stroji za 1,3 sekundy u modelu Kimi-K2 (1T parametrů) přenosem vah z 256 tréninkových GPU (BF16) na 128 inferenčních GPU (FP8).

Při asynchronním jemném doladění pomocí posilovaného učení (RL fine-tuning) běží trénování a inference na samostatných GPU. Po každém tréninkovém kroku je nutné odeslat nové váhy do inferenčních uzlů. Mnoho stávajících frameworků potřebuje u modelů s biliony parametrů několik sekund – nebo dokonce minut.

Využitím bod-bodové komunikace RDMA dokážeme přenos vah zajistit bleskově rychle, aniž bychom měnili inferenční engine, a zároveň zjednodušit psaní a údržbu kódu.

Naše řešení je postaveno na RDMA WRITE, což je jednostranný primitiv, kde zdroj zapisuje přímo do paměti GPU v cílovém zařízení.

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

Cílová strana dokonce ani nedostane o přenosu informaci. To nám poskytuje přenosy s nízkou latencí, vysokou propustností a nulovým kopírováním, které jsou řízeny tréninkovými uzly bez jakékoli řídicí logiky na inferenčních uzlech.

Sběr metadat – Controller shromažďuje metadata parametrů ze všech tréninkových a inferenčních GPU.

Výpočet rozvrhu – Controller vypočítá statický rozvrh přenosu vah, který mapuje, které tréninkové GPU odešle jaký parametr na které inferenční GPU a v jakém pořadí.

Distribuce rozvrhu – Controller odešle rozvrh všem tréninkovým GPU.

Spuštění – Po každém tréninkovém kroku vyšle controller signál tréninkovým GPU k zahájení přenosů.

Po definování vysokoúrovňového pracovního postupu spočívá hlavní výzva v tom, jak efektivně provádět přenosy vah v měřítku bilionů parametrů. Zde popisujeme podrobnosti cesty provedení.

Parametry při trénování jsou distribuovány podle umístění FSDP. Pomocí full_tensor() mohou všechna GPU v DeviceMesh rekonstruovat plný parametr, a proto všechna mohou sloužit jako zdroj pro přenos vah.

Několik disjunktních DeviceMesh tvoří mesh group. Jelikož jsou DeviceMesh ve stejné skupině disjunktní, jejich přenosy se neruší a mohou probíhat plně paralelně. Mezi skupinami mesh groups vkládáme globální bariéru pro vynucení řazení.

Přenosová pipeline vah se překrývá s kopírováním z hostitele na zařízení, sestavováním tenzorů, fúzí a kvantizací a s přenosy RDMA napříč skupinami MeshGroup.

Přenos každého tenzoru parametrů považujeme za úkol. Proces přenosu vah využívá více typů hardwarových zdrojů, proto rozdělujeme úkol přenosu vah do různých fází pipeline, které se v čase překrývají:

Host to device memcpy – pokud FSDP odkládá (offloaduje) váhy do CPU

Příprava parametrů – rekonstrukce úplné váhy pomocí full_tensor(), použití fúze projekce, případně kvantování.

Přenos RDMA – zápis s nulovým kopírováním do paměti vzdáleného inferenčního GPU

Globální bariéra – po dokončení všech volání full_tensor() proběhne synchronizace napříč skupinami mesh groups pomocí GLOO přes Ethernet.

Při implementaci udržujeme FIFO frontu úkolů pro každou fázi pipeline. Kdykoli úkol na začátku fronty dokončí danou fázi, je přesunut na konec fronty další fáze.

full_tensor() a další operace GPU zvyšují dodatečné využití paměti GPU. Abychom se vyhnuli chybě out of memory, zahájíme provádění úkolu pouze tehdy, pokud aktuální probíhající úkoly zabírají méně dočasné paměti GPU, než je konfigurovatelná mezní hodnota (watermark).

Nkolik designových voleb činí náš systém výrazně rychlejším na provoz a jednodušším na údržbu než běžná řešení s otevřeným zdrojovým kódem.

Přenos vah založený na Rank0 nejprve sesbírá váhy a teprve poté je distribuuje na inferenční uzel.

Komunikace point-to-point přenáší váhy z obou trénovacích uzlů do síťových rozhraní inferenčního uzlu.

Běžným vzorem je směřovat všechny parametry přes GPU s rankem 0: shromáždit je na tréninkovém ranku 0, poslat na inferenční rank 0 a poté je zase rozptýlit. To se rychle stává úzkým hrdlem, omezeným propustností PCIe a síťovou kartou (NIC) jednoho GPU (např. 400 Gbps ≈ 50 GB/s).

Naproti tomu naše bod-bodové uspořádání umožňuje každému tréninkovému GPU odesílat data přímo na každé inferenční GPU, čímž vytěžuje celou síťovou infrastrukturu namísto jediného propojení.

Přenos založený na RPC používá buffer na straně inference a volání workerů, zatímco jednosměrný přenos zapisuje váhy přímo bez účasti inference workeru.

Některé systémy spoléhají na volání metody update_weight() inferenčního enginu pro každý tenzor. To znamená invazivní změny v inferenčním kódu a navíc režii spojenou s RPC, serializací a koordinací řídicí roviny.

S primitivním RDMA WRITE aktualizujeme váhy tiše v paměti inferenčního GPU bez dodatečných kopií. Nejsou zapojeny žádné zprávy řídicí roviny ani řídicí logika CPU. Není vyžadována žádná úprava inferenčního enginu.

Proces přenosu vah může využívat čtyři typy hardwarových zdrojů: (1) Přesun dat mezi hostitelem a zařízením (2) Výpočty GPU pro fúzi projekce a kvantování (3) Síť RDMA pro datovou rovinu (4) Ethernet pro řídicí rovinu.

Náš návrh rozdělil úkoly přenosu vah do fází pipeline, což umožňuje snadné překrývání napříč různými hardwarovými zdroji.

Některé implementace přepočítávají rozvrh přenosu v každém tréninkovém kroku, opakovaně shromažďují metadata a distribuují instrukce. To přidává zbytečnou latenci řídicí roviny.

Náš rozvrh se vypočítá jednou při inicializaci. Každá tréninková iterace jednoduše opakuje plán: controller vydá signál „start“ a GPU se řídí svými předem přiřazenými trasami. Provádění je předvídatelné a nenáročné na zdroje.

Je lákavé zaplést celý proces aktualizace vah do jedné monolitické funkce: sběr metadat, párování názvů, shromažďování v rámci uzlu, fúze projekce, kvantování, rozdělení komunikačního světa na podskupiny, síťový přenos mezi uzly. Je těžké to správně naprogramovat a ještě těžší optimalizovat.

V našem inženýrském přístupu oddělujeme tyto kroky jako jednotlivé komponenty. Každá komponenta může být testována, analyzována a optimalizována izolovaně.

Rychlý a spolehlivý přenos vah je klíčovým stavebním kamenem pro velkoškálové doladění pomocí RL. Kombinací primitivu RDMA WRITE, statického rozvrhu přenosu a zpracování v pipeline jsme u modelu Kimi-K2 zkrátili aktualizace s biliony parametrů na pouhé 1,3 sekundy. Přístup je snadno pochopitelný, nenáročný na údržbu a vyhýbá se úzkým hrdlům tradičních návrhů.

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é přenosy

Vysokoúrovňový pracovní postup

Provedení přenosu vah

DeviceMesh a skupiny Mesh Groups

Trubková linie úkolů

Řízení využití paměti GPU

Proč je to rychlé a jednoduché

Bod-bodová komunikace

Jednostranný přenos dat

Vytváření pipeline (Pipelining)

Statický rozvrh

Čisté oddělení

Závěr

Citace