Viktöverföring för RL-postträning på under 2 sekunder

Vi uppnådde nyligen 1,3 sekunders parameteruppdateringar mellan maskiner för Kimi-K2 (1T-parametrar) genom att överföra vikter från 256 tränings-GPU:er (BF16) till 128 inferens-GPU:er (FP8).

Vid finjustering med asynkron förstärkningsinlärning körs träning och inferens på separata GPU:er. Efter varje träningssteg måste nya vikter skickas till inferensnoder. Många befintliga ramverk tar flera sekunder – eller till och med minuter – för modeller med biljoner parametrar.

Genom att utnyttja RDMA-punkt-till-punkt-kommunikation kan vi göra viktöverföringen blixtsnabb utan att ändra inferensmotorn, samt göra koden enklare att skriva och underhålla.

Vår lösning är byggd på RDMA WRITE, en ensidig primitiv där källan direkt skriver till destinationens 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.
    ...

Destinationssidan får inte ens en meddelandenotis om överföringen. Detta ger oss överföringar med låg latens, hög genomströmning och noll kopiering som styrs av träningsnoderna utan kontrolllogik på inferensnoderna.

Metadatainsamling – Kontrollern samlar in parametermetadata från alla tränings- och inferens-GPU:er.

Schemaberäkning – Kontrollern beräknar ett statiskt schema för viktöverföring som mappar vilken tränings-GPU som skickar vilken parameter till vilken inferens-GPU, och i vilken ordning.

Schemadistribution – Kontrollern skickar schemat till alla tränings-GPU:er.

Utförande – Efter varje träningssteg signalerar kontrollern till tränings-GPU:erna att starta överföringarna.

Med det övergripande arbetsflödet definierat är den stora utmaningen hur man utför viktöverföringar effektivt i en skala med biljoner parametrar. Här beskriver vi detaljerna för exekveringsvägen.

Parametrar i träning distribueras enligt FSDP-placeringar. Genom att använda full_tensor() kan alla GPU:er i ett DeviceMesh återskapa hela parametern, och därmed kan alla fungera som källa för viktöverföring.

Flera disjunkta DeviceMesh bildar en mesh-grupp. Eftersom DeviceMesh i samma grupp är disjunkta stör inte deras överföringar varandra och kan köras helt parallellt. Mellan mesh-grupper infogar vi en global barriär för att säkerställa ordningsföljd.

Viktöverföringspipelinen överlappar kopior från värd till enhet, tensoruppbyggnad, fusion och kvantisering samt RDMA-överföringar över MeshGroups.

Vi behandlar överföringen av varje parametertensor som en uppgift. Viktöverföringsprocessen använder flera typer av hårdvarukällor, vilket är anledningen till att vi delar upp en viktöverföringsuppgift i olika pipelinesteg som överlappar i tid:

Host-till-enhet-memcpy – Om FSDP offloadar vikt till CPU

Parameterförberedelse – Återskapa full vikt med full_tensor(), tillämpa projektionsfusion, kvantisera vid behov.

RDMA-överföring – Nollkopieringsskrivning till fjärrminnet för inferens-GPU

Global barriär – När alla full_tensor()-anrop är klara, synkronisera över mesh-grupper med GLOO via Ethernet.

Vid implementering underhåller vi en FIFO-kö med uppgifter för varje pipelinesteg. När den främsta uppgiften i kön slutför steget flyttas den till slutet av nästa stegs kö.

full_tensor() och andra GPU-operationer introducerar extra GPU-minnesanvändning. För att undvika slut-på-minne-fel (OOM) startar vi utförandet av en uppgift endast om de aktuella pågående uppgifterna upptar mindre tillfälligt GPU-minne än en konfigurerbar tröskel.

Flera designval gör vårt system betydligt snabbare att köra och enklare att underhålla än vanliga open source-lösningar.

Rank0-baserad viktöverföring samlar vikter innan de distribueras till inferensnoden.

Punkt-till-punkt-kommunikation överför vikter från båda träningsnoderna till inferensnodens nätverksgränssnitt.

Ett vanligt mönster är att kanalisera alla parametrar genom rang-0-GPU:er: samla in på tränings-rang-0, skicka till inferens-rang-0 och sedan sprida ut igen. Detta blir snabbt en flaskhals, begränsad av en enda GPU:s PCIe-bandbredd och NIC (t.ex. 400 Gbps ≈ 50 GB/s).

Vår punkt-till-punkt-konfiguration tillåter däremot varje tränings-GPU att skicka direkt till varje inferens-GPU, vilket mättar hela nätverksstrukturen snarare än en enskild länk.

RPC-baserad överföring använder en buffert på inferenssidan och anrop från arbetare, medan envägsöverföring skriver vikter direkt utan medverkan från inferensarbetaren.

Vissa system förlitar sig på att anropa inferensmotorns metod update_weight() för varje tensor. Det innebär påträngande ändringar i inferenskoden samt omkostnader från RPC:er, serialisering och kontrollplanskoordinering.

Med RDMA WRITE-primitiven uppdaterar vi vikter tyst i inferens-GPU-minnet utan extra kopior. Inget kontrollplansmeddelande och ingen CPU-kontrolllogik är involverad. Ingen modifiering av inferensmotorn krävs.

Viktöverföringsprocessen kan utnyttja fyra typer av hårdvaruresurser: (1) Dataflytt mellan värd och enhet (2) GPU-beräkning för projektionsfusion och kvantisering (3) RDMA-nätverk för dataplan (4) Ethernet för kontrollplan.

Vår design delar upp viktöverföringsuppgifter i pipelinesteg, vilket möjliggör enkel överlappning mellan olika hårdvaruresurser.

Vissa implementeringar beräknar om ett överföringsschema vid varje träningssteg genom att upprepade gånger samla in metadata och distribuera instruktioner. Detta lägger till onödig latens i kontrollplanet.

Vårt schema beräknas en gång vid initiering. Varje träningsiteration spelar helt enkelt upp planerna igen: kontrollern utfärdar en ”kör”-signal och GPU:erna följer sina förutbestämda vägar. Utförandet är förutsägbart och lättviktsorienterat.

Det är frestande att fläta in hela viktuppdateringsprocessen i en enda monolitisk funktion: samla in metadata, namnmatchning, samling inom noden, projektionsfusion, kvantisering, uppdelning av kommunikationsvärlden och nätverksöverföring mellan noder. Det är svårt att programmera korrekt och ännu svårare att optimera.

I vår konstruktion separerar vi dessa steg som individuella komponenter. Varje komponent kan enhetstestas, analyseras och optimeras isolerat.

Snabb och tillförlitlig viktöverföring är en kritisk byggsten för storskalig RL-finjustering. Genom att kombinera RDMA WRITE-primitiven, ett statiskt överföringsschema och pipelined exekvering reducerade vi uppdateringar av biljoner parametrar till endast 1,3 sekunder på Kimi-K2. Tillvägagångssättet är enkelt att resonera kring, lätt att underhålla och undviker flaskhalsarna i traditionella designer.

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: ensidiga överföringar

Övergripande arbetsflöde

Utförande av viktöverföring

DeviceMesh och Mesh-grupper

Uppgiftspipeline

Kontroll av GPU-minnesanvändning

Varför det är snabbt och enkelt

Punkt-till-punkt-kommunikation

Ensidig dataöverföring

Pipelining

Statiskt schema

Ren separation

Slutsats

Citat