Súlyátvitel RL utótanításhoz 2 másodperc alatt

Nemrégiben 1,3 másodperces gépek közötti paraméterfrissítést értünk el a Kimi-K2 (1T paraméter) esetében, átvive a súlyokat 256 tanító GPU-ról (BF16) 128 inferencia GPU-ra (FP8).

Az aszinkron megerősítéses tanulási finomhangolás során a tanítás és az inferencia külön GPU-kon fut. Minden tanítási lépés után az új súlyokat át kell küldeni az inferencia csomópontokra. Számos meglévő keretrendszer esetében ez több másodpercet – vagy akár percet – vesz igénybe a trillió paraméteres modelleknél.

Az RDMA pont-pont kommunikáció kihasználásával képesek a súlyátvitelt villámgyorssá tenni az inferenciamotor megváltoztatása nélkül, és a kódot is könnyebben megírhatóvá és karbantarthatóvá tehetjük.

Megoldásunk az RDMA WRITE-ra épül, amely egy egyoldalas primitív, ahol a forrás közvetlenül ír a célállomás GPU-memóriájába.

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

A céloldal értesítést sem kap az átvitelről. Ez alacsony késleltetésű, nagy sávszélességű, másolatmentes (zero-copy) átviteleket biztosít, amelyeket a tanító csomópontok vezérelnek az inferencia csomópontok vezérlési logikája nélkül.

Metaadatok gyűjtése – A vezérlő összegyűjti a paraméterek metaadatait az összes tanító és inferencia GPU-ról.

Ütemezés számítása – A vezérlő kiszámít egy statikus súlyátviteli ütemezést, amely leképezi, hogy melyik tanító GPU melyik paramétert melyik inferencia GPU-nak küldi, és milyen sorrendben.

Ütemezés elosztása – A vezérlő elküldi az ütemezést az összes tanító GPU-nak.

Végrehajtás – Minden tanítási lépés után a vezérlő jelzést küld a tanító GPU-knak az átvitelek elindításához.

A magas szintű munkafolyamat meghatározásával a kulcsfontosságú kihívás annak eldöntése, hogyan hajtható végre hatékonyan a súlyátvitel trillió paraméteres skálán. Itt ismertetjük a végrehajtási útvonal részleteit.

A tanítás során a paraméterek az FSDP elhelyezések szerint oszlanak meg. A full_tensor() használatával a DeviceMesh összes GPU-ja rekonstruálni tudja a teljes paramétert, így mindegyik szolgálhat forrásként a súlyátvitelhez.

Több diszjunkt DeviceMesh alkot egy mesh-csoportot. Mivel az azonos csoportban lévő DeviceMeshek diszjunktak, az átviteleik nem zavarják egymást, és teljesen párhuzamosan futhatnak. A mesh-csoportok között egy globális korlátot (barrier) iktatunk be a sorrend betartásához.

A súlyátviteli csővezeték átfedésben végzi a hoszt–eszköz másolatokat, a tenzor-összeállítást, az egyesítést és a kvantálást, valamint a MeshGroupokon átívelő RDMA-átvitelt.

Az egyes paramétertenzorok átvitelét feladatként kezeljük. A súlyátviteli folyamat többféle hardverforrást használ, ezért a súlyátviteli feladatot különböző, időben átfedő folyamatszakaszokra bontjuk:

Host-to-device memcpy – Ha az FSDP átteszi a súlyt a CPU-ra

Paraméter-előkészítés – Teljes súly rekonstrukciója a full_tensor() függvénnyel, projekciós fúzió alkalmazása, szükség szerinti kvantálás.

RDMA átvitel – Másolatmentes írás a távoli inferencia GPU memóriájába

Globális korlát – Az összes full_tensor() hívás befejezése után szinkronizálás a mesh-csoportok között a GLOO használatával, Etherneten keresztül.

A megvalósítás során FIFO feladatvárakozó sort (queue) tartunk fenn minden folyamatszakaszhoz. Amikor a sor elején lévő feladat befejezi a szakaszt, átkerül a következő szakasz sorának végére.

A full_tensor() és más GPU-műveletek extra GPU-memóriahasználatot vonnak maga után. A memóriatúllépési hiba (out of memory) elkerülése érdekében egy feladat végrehajtását csak akkor indítjuk el, ha a jelenleg folyamatban lévő feladatok kevesebb ideiglenes GPU-memóriát foglaltak le, mint egy konfigurálható küszöbérték (watermark).

Számos tervezési döntés teszi a rendszerünket lényegesen gyorsabban futtathatóvá és könnyebben karbantarthatóvá a gyakori nyílt forráskódú megoldásokhoz képest.

A Rank0-alapú súlyátvitel a súlyokat összegyűjti, mielőtt továbbítaná őket az inferencianódnak.

A ponttól pontig kommunikáció a súlyokat mindkét tréningcsomópont hálózati interfészeiről az inferencianód hálózati interfészeire továbbítja.

Gyakori minta, hogy az összes paramétert a 0. rangú GPU-kon keresztül csatornázzák: gyűjtés a tanítási 0. rangon, küldés az inferencia 0. rangra, majd ismételt szétosztás. Ez gyorsan szűk keresztmetszetté válik, amelyet egyetlen GPU PCIe-sávszélessége és NIC-je (pl. 400 Gbps ≈ 50 GB/s) korlátoz.

Ezzel szemben pont-pont felépítésünk lehetővé teszi, hogy minden tanító GPU közvetlenül küldjön adatot minden inferencia GPU-nak, teljesen telítve a hálózati struktúrát ahelyett, hogy egyetlen kapcsolat korlátozna.

Az RPC-alapú átvitel az inferenciaoldali puffert és a worker-hívásokat használja, míg az egyoldalú átvitel a súlyokat közvetlenül írja, az inferencia-worker részvétele nélkül.

Egyes rendszerek az inferenciamotor update_weight() metódusának meghívására támaszkodnak minden egyes tenzor esetében. Ez tolakodó változtatásokat jelent az inferenciakódban, ráadásul az RPC-k, a szerializáció és a vezérlősík-koordináció is többletterheléssel jár.

Az RDMA WRITE primitívvel a súlyokat észrevétlenül frissítjük az inferencia GPU memóriájában, extra másolatok nélkül. Nincs vezérlősíkon futó üzenet és semmilyen CPU vezérlési logika. Az inferenciamotor módosítására sincs szükség.

A súlyátviteli folyamat négyféle hardvererőforrást tud hasznosítani: (1) Host és eszköz közötti adatmozgatás (2) GPU-számítás a projekciós fúzióhoz és kvantáláshoz (3) RDMA hálózat az adatsíkon (4) Ethernet a vezérlősíkon.

Tervezésünk során a súlyátviteli feladatokat folyamatszakaszokra bontottuk, lehetővé téve a különböző hardvererőforrások közötti egyszerű átfedést.

Egyes megvalósítások minden tanítási lépésben újraszámítják az átviteli ütemezést, ismételten összegyűjtve a metaadatokat és szétosztva az utasításokat. Ez felesleges vezérlősíkon belüli késleltetést okoz.

Ütemezésünket egyszer, az inicializáláskor számítjuk ki. Minden tanítási iteráció egyszerűen megismétli a tervet: a vezérlő kiad egy „indulás” jelzést, és a GPU-k követik az előre hozzárendelt útvonalakat. A végrehajtás kiszámítható és kis erőforrásigényű.

Csábító lehet a teljes súlyfrissítési folyamatot egyetlen monolitikus függvénybe ágyazni: metaadatok gyűjtése, névillesztés, csomóponton belüli gyűjtés, projekciós fúzió, kvantálás, kommunikációs tartomány alosztása, csomópontok közötti hálózati átvitel. Ezt nehéz helyesen megírni, és még nehezebb optimalizálni.

Fejlesztésünk során ezeket a lépéseket különálló komponensekként választjuk szét. Minden komponens egységtesztelhető, elemezhető és külön-külön optimalizálható.

A gyors és megbízható súlyátvitel a nagyszabású RL-finomhangolás kritikus építőköve. Az RDMA WRITE primitív, a statikus átviteli ütemezés és a csővezetékes végrehajtás kombinálásával a trillió paraméteres frissítéseket mindössze 1,3 másodpercre csökkentettük a Kimi-K2-n. A megközelítés könnyen áttekinthető, egyszerűen karbantartható, és elkerüli a hagyományos kialakítások szűk keresztmetszeteit.

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: egyoldalas átvitelek

Magas szintű munkafolyamat

Súlyátvitel végrehajtása

DeviceMesh és Mesh-csoportok

Feladat-folyamat (Task pipeline)

GPU-memóriahasználat szabályozása

Miért gyors és egyszerű?

Pont-pont (point-to-point) kommunikáció

Egyoldalas adóátvitel

Pipelining (csővezetékezés)

Statikus ütemezés

Tiszta szétválasztás

Összegzés

Idézet