Vektoverføring for RL-ettertrening på under 2 sekunder

Vi oppnådde nylig 1,3 sekunders parameteroppdateringer på tvers av maskiner for Kimi-K2 (1T parametere), og overførte vekter fra 256 trenings-GPU-er (BF16) til 128 inferens-GPU-er (FP8).

Ved asynkron finjustering med forsterkende læring (RL) kjører trening og inferens på separate GPU-er. Etter hvert treningstrinn må nye vekter overføres til inferensnoder. Mange eksisterende rammeverk tar flere sekunder – eller til og med minutter – for modeller med en billion parametere.

Ved å utnytte RDMA-kommunikasjon punkt-til-punkt kan vi gjøre vektoverføringen ekstremt rask, uten å endre inferensmotoren, og gjøre koden enklere å skrive og vedlikeholde.

Løsningen vår er bygget på RDMA WRITE, et ensidig primitiv der kilden skriver direkte til destinasjonens GPU-minne.

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

Destinasjonssiden blir ikke en gang varslet om overføringen. Dette gir oss overføringer med lav latens, høy gjennomstrømming og nullkopiering, drevet av treningsnodene uten kontrollikk på inferensnodene.

Metadatainnsamling – Kontrolleren samler inn parametermetadata fra alle trenings- og inferens-GPU-er.

Beregning av tidsplan – Kontrolleren beregner en statisk tidsplan for vektoverføring som tilordner hvilken trenings-GPU som sender hvilken parameter til hvilken inferens-GPU, og i hvilken rekkefølge.

Distribusjon av tidsplan – Kontrolleren sender tidsplanen til alle trenings-GPU-er.

Utførelse – Etter hvert treningstrinn signaliserer kontrolleren til trenings-GPU-ene om å starte overføringer.

Med den overordnede arbeidsflyten definert er den største utfordringen hvordan man utfører vektoverføringer effektivt i en skala med en billion parametere. Her beskriver vi detaljene i utførelsesbanen.

Parametere i trening distribueres i henhold til FSDP-plasseringer. Ved hjelp av full_tensor() kan alle GPU-er i en DeviceMesh rekonstruere hele parameteren, og dermed fungere som kilde for vektoverføring.

Flere disjunkte DeviceMeshes danner en maskegruppe. Fordi DeviceMeshes i samme gruppe er disjunkte, forstyrrer ikke overføringene deres hverandre og kan kjøre fullt parallelt. Mellom maskegrupper setter vi inn en global barriere for å håndheve rekkefølge.

Vektroverførings-pipelinen overlapper kopieringer fra vert til enhet, tensor-assemblage, fusjon og kvantisering, samt RDMA-overføringer på tvers av MeshGroups.

Vi behandler overføringen av hver parametertensor som en oppgave. Vektoverføringsprosessen bruker flere typer maskinvarekilder, og derfor deler vi opp en vektoverføringsoppgave i ulike rørledningstrinn som overlapper i tid:

Host-til-enhet memcpy — Hvis FSDP laster av vekten til CPU

Parameterklargjøring — Rekonstruer hele vekten med full_tensor(), bruk projeksjonssammenslåing, kvantiser om nødvendig.

RDMA-overføring — Nullkopiskrivning til eksternt inferens-GPU-minne

Global barriere — Etter at alle full_tensor()-kall er fullført, synkroniser på tvers av maskegrupper ved hjelp av GLOO via Ethernet.

I implementeringen opprettholder vi en FIFO-kø med oppgaver for hvert rørledningstrinn. Når oppgaven i starten av køen fullfører trinnet, flyttes den til slutten av køen for neste trinn.

full_tensor() og andre GPU-operasjoner introduserer ekstra GPU-minnebruk. For å unngå minnefeil (out of memory) starter vi utførelsen av en oppgave bare hvis de pågående oppgavene opptar mindre midlertidig GPU-minne enn en konfigurerbar vannmerkeverdi.

Flere designvalg gjør systemet vårt betydelig raskere å kjøre og enklere å vedlikeholde enn vanlige åpen kildekode-løsninger.

Rank0-basert vektroverføring samler vekter før de distribueres til inferensnoden.

Punkt-til-punkt-kommunikasjon overfører vekter fra begge treningsnodene til inferensnodens nettverksgrensesnitt.

Et vanlig mønster er å kanalisere alle parametere gjennom rang-0-GPU-er: samle inn på treningsrang-0, sende til inferensrang-0, og deretter spre igjen. Dette blir raskt en flaskehals, begrenset av en enkelt GPU-s PCIe-båndbredde og NIC (f.eks. 400 Gbps ≈ 50 GB/s).

I kontrast lar oppsettet vårt for punkt-til-punkt hver trenings-GPU sende direkte til hver inferens-GPU, noe som metter hele nettverksstrukturen i stedet for en enkelt kobling.

RPC-basert overføring bruker en buffer på inferenssiden og kall fra arbeidere, mens ensidig overføring skriver vekter direkte uten deltakelse fra inferensarbeideren.

Noen systemer er avhengige av å kalle inferensmotorens update_weight()-metode for hver tensor. Det betyr inngripende endringer i inferenskoden, pluss overhead fra RPC-er, serialisering og kontrollplanskoordinering.

Med RDMA WRITE-primitivet oppdaterer vi vekter i det stille på inferens-GPU-minnet, uten ekstra kopier. Ingen kontrollplanmelding og ingen CPU-kontrolllogikk er involvert. Ingen modifikasjon av inferensmotoren er nødvendig.

Vektoverføringsprosessen kan utnytte fire typer maskinvareessurser: (1) Dataforflytning mellom host og enhet (2) GPU-beregning for projeksjonssammenslåing og kvantisering (3) RDMA-nettverk for dataplan (4) Ethernet for kontrollplan.

Designet vårt deler opp vektoverføringsoppgaver i rørledningstrinn, noe som gir enkel overlapp på tvers av ulike maskinvareessurser.

Noen implementeringer beregner en overføringstidsplan på nytt ved hvert treningstrinn, og samler gjentatte ganger inn metadata og distribuerer instruksjoner. Dette gir unødvendig latens i kontrollplanet.

Tidsplanen vår beregnes én gang ved initialisering. Hver treningsiterasjon gjentar ganske enkelt planene: kontrolleren utsteder et «gå»-signal, og GPU-er følger sine forhåndsijesernte ruter. Utførelsen er forutsigbar og lettvekts.

Det er fristende å flette hele vektoppdateringsprosessen inn i én monolitisk funksjon: samle metadata, navnematching, innsamling innenfor node, projeksjonssammenslåing, kvantisering, oppdeling av kommunikasjonsverden, nettverksoverføring mellom noder. Det er vanskelig å programmere riktig, og enda vanskeligere å optimalisere.

I vår ingeniørkunst skiller vi disse trinnene som individuelle komponenter. Hver komponent kan enhetstestes, resonneres om og optimaliseres isolert.

Rask og pålitelig vektoverføring er en kritisk byggestein for storskala RL-finjustering. Ved å kombinere RDMA WRITE-primitivet, en statisk overføringstidsplan og rørledningsbasert utførelse reduserte vi oppdateringer med en billion parametere til bare 1,3 sekunder på Kimi-K2. Tilnærmingen er enkel å resonnere rundt, lett å vedlikeholde og unngår flaskehalsene i tradisjonelle 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: ensidige overføringer

Overordnet arbeidsflyt

Utførelse av vektoverføring

DeviceMesh og maskegrupper

Oppgaverørledningslinje

Kontroll av GPU-minnebruk

Hvorfor det er raskt og enkelt

Kommunikasjon punkt-til-punkt

Ensidig dataoverføring

Rørledningsføring

Statisk tidsplan

Ren separasjon

Konklusjon

Sitat