Vægtoverførsel til RL-post-træning på under 2 sekunder

Vi opnåede for nylig 1,3-sekunders parameterafsendelser på tværs af maskiner for Kimi-K2 (1T parametre) ved at overføre vægte fra 256 trænings-GPU'er (BF16) til 128 inferens-GPU'er (FP8).

Ved asynkron finjustering med forstærkende læring kører træning og inferens på separate GPU'er. Efter hvert træningstrin skal nye vægte skubbes til inferensnoder. Mange eksisterende rammer tager flere sekunder – eller endda minutter – for modeller med billioner af parametre.

Ved at udnytte RDMA-punkt-til-punkt-kommunikation kan vi gøre vægtoverførslen ekstremt hurtig uden at ændre inferensmotoren og gøre koden nemardere at skrive og vedligeholde.

Vores løsning er bygget på RDMA WRITE, et ensidigt primitiv, hvor kilden skriver direkte til destinationens GPU-hukommelse.

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

Destinationssiden får ikke engang besked om overførslen. Dette giver os lav latens, høj gennemnemsnitsevne og nul-kopi-overførsler drevet af træningsnoderne uden kontrollogik på inferensnoderne.

Metadatasamling – Controlleren indsamler parametermetadata fra alle trænings- og inferens-GPU'er.

Planlægningsberegning – Controlleren beregner en statisk vægtoverførselsplan, der knytter hvilken trænings-GPU der sender hvilken parameter til hvilken inferens-GPU, og i hvilken rækkefølge.

Planfordeling – Controlleren sender planlægningen til alle trænings-GPU'er.

Udførelse – Efter hvert træningstrin signalerer controlleren til trænings-GPU'er om at starte overførsler.

Med den overordnede arbejdsproces defineret er den vigtigste udfordring, hvordan man udfører vægtoverførsler effektivt i skala med billioner af parametre. Her beskriver vi detaljerne i udførelsesstien.

Parametre i træning distribueres i henhold til FSDP-placeringer. Ved hjælp af full_tensor() kan alle GPU'er i et DeviceMesh genskabe den fulde parameter, og dermed kan alle fungere som kilde til vægtoverførsel.

Flere disjunkte DeviceMeshes danner en mesh-gruppe. Da DeviceMeshes i samme gruppe er disjunkte, forstyrrer deres overførsler ikke hinanden og kan køre fuldstændigt parallelt. Mellem mesh-grupper indsætter vi en global barriere for at håndhæve rækkefølge.

Vægtoverførsels-pipelinen overlapper kopier fra vært til enhed, tensor-assemblage, fusion og kvantisering samt RDMA-overførsler på tværs af MeshGroups.

Vi behandler overførslen af hver parametertensor som en opgave. Vægtoverførselsprocessen udnytter flere typer hardwarekilder, og derfor opdelte vi en vægtoverførselsopgave i forskellige pipelinestadier, der overlapper i tid:

Vært til enhed memcpy — Hvis FSDP aflæsser vægt til CPU

Parameterforberedelse — Genskab fuld vægt med full_tensor(), anvend projektionsfusion, kvantisér om nødvendigt.

RDMA-overførsel — Nul-kopi-skrivning til ekstern inferens-GPU-hukommelse

Global barriere — Når alle full_tensor()-kald er udført, synkroniseres på tværs af mesh-grupper ved hjælp af GLOO via Ethernet.

Under implementeringen opretholder vi en FIFO-kø af opgaver for hvert pipelinestadie. Hver gang opgaven i køens hoved fuldfører stadiet, flyttes den til halen af den næste stadiekø.

full_tensor() og andre GPU-operationer medfører ekstra GPU-hukommelsesforbrug. For at undgå manglende hukommelsesfejl starter vi udførelsen af en opgave, kun hvis de aktuelle igangværende opgaver optager mindre midlertidig GPU-hukommelse end et konfigurerbart vandmærke.

Flere designvalg gør vores system væsentligt hurtigere at køre og nemmere at vedligeholde end almindelige open source-løsninger.

Vægtoverførsel baseret på Rank0 indsamler vægte, før de distribueres til inferensnoden.

Punkt-til-punkt-kommunikation overfører vægte fra begge træningsnoder til inferensnodens netværksgrænseflader.

Et almindeligt mønster er at kanalisere alle parametre gennem rang-0-GPU'er: indsamle på træningsrang-0, sende til inferensrang-0 og derefter sprede igen. Dette bliver hurtigt et flaskehals, begrænset af en enkelt GPU's PCIe-båndbredde og NIC (f.eks. 400 Gbps ≈ 50 GB/s).

I kontrast gør vores punkt-til-punkt-opsætning det muligt for hver trænings-GPU at sende direkte til hver inferens-GPU, hvorved hele netværksstrukturen mættes i stedet for et enkelt link.

RPC-baseret overførsel bruger en buffer på inferenssiden og worker-kald, mens ensidet overførsel skriver vægte direkte uden deltagelse fra inferens-worker'en.

Nogle systemer er afhængige af at kalde ind i inferensmotorens update_weight()-metode for hver tensor. Det betyder indgribende ændringer i inferenskoden samt overhead fra RPC'er, serialisering og kontrolplanskoordinering.

Med RDMA WRITE-primitivet opdaterer vi vægte lydløst i inferens-GPU-hukommelsen uden ekstra kopier. Ingen kontrolplanmeddelelse og ingen CPU-kontrollogik er involveret. Ingen ændring af inferensmotoren er påkrævet.

Vægtoverførselsprocessen kan udnytte fire typer hardware-ressourcer: (1) Vært-enhed-dataflytning (2) GPU-beregning til projektionsfusion og kvantisering (3) RDMA-netværk til dataplan (4) Ethernet til kontrolplan.

Vores design opdeler vægtoverførselsopgaver i pipelinestadier, hvilket gør det nemt at overlappe på tværs af forskellige hardware-ressourcer.

Nogle implementeringer genberegner en overførselsplan ved hvert træningstrin ved gentagne gange at indsamle metadata og distribuere instruktioner. Dette tilføjer unødvendig kontrolplanslatens.

Vores tidsplan beregnes én gang ved initialisering. Hver træningsiteration afspiller simpelthen planerne: controlleren udsender et "go"-signal, og GPU'er følger deres forudbestemte ruter. Udførelsen er forudsigelig og letvægts.

Det er fristende at indvikle hele vægtopdateringsprocessen i en enkelt monolitisk funktion: indsamle metadata, navnematchning, indsamling inden for noder, projektionsfusion, kvantisering, opsplitning af kommunikationsverden, netværksoverførsel på tværs af noder. Det er svært at programmere korrekt og endnu sværere at optimere.

I vores teknik adskiller vi disse trin som individuelle komponenter. Hver komponent kan enhedstestes, ræsonneres om og optimeres isoleret.

Hurtig, pålidelig vægtoverførsel er en kritisk byggesten til storskala RL-finjustering. Ved at kombinere RDMA WRITE-primitivet, en statisk overførselsplan og pipelined udførelse reducerede vi opdateringer af billioner af parametre til blot 1,3 sekunder på Kimi-K2. Tilgangen er nem at forholde sig til, nem at vedligeholde og undgår flaskehalsene i traditionelle design.

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: ensidede overførsler

Overordnet arbejdsproces

Udførelse af vægtoverførsel

DeviceMesh og Mesh-grupper

Opgavepipeline

Styring af GPU-hukommelsesforbrug

Hvorfor det er hurtigt og enkelt

Punkt-til-punkt-kommunikation

Ensidet dataoverførsel

Pipelining

Statisk tidsplan

Ren adskillelse

Konklusion

Citering