Gewichtsoverdracht voor RL Post-Training in minder dan 2 sec

We hebben onlangs 1,3 seconden behaald voor parameter-updates tussen machines voor Kimi-K2 (1T parameters), waarbij gewichten werden overgedragen van 256 trainings-GPU's (BF16) naar 128 inferentie-GPU's (FP8).

Bij asynchrone fine-tuning van reinforcement learning draaien training en inferentie op afzonderlijke GPU's. Na elke trainingsstap moeten nieuwe gewichten naar inferentie-knooppunten worden gepusht. Veel bestaande frameworks doen hier enkele seconden — of zelfs minuten — over voor modellen met biljoenen parameters.

Door gebruik te maken van RDMA point-to-point-communicatie kunnen we de gewichtsoverdracht razendsnel maken, zonder de inferentie-engine te wijzigen, en de code gemakkelijker schrijven en onderhouden.

Onze oplossing is gebouwd op RDMA WRITE, een eenzijdig primitief waarbij de bron direct in het GPU-geheugen van de bestemming schrijft.

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

De ontvangende kant krijgt niet eens een melding van de overdracht. Dit levert ons low-latency, high-throughput, zero-copy overdrachten op, aangestuurd door de trainingsknooppunten zonder enige besturingslogica op de inferentie-knooppunten.

Verzameling van metadata – De controller verzamelt parametermetadata van alle trainings- en inferentie-GPU's.

Berekening van planning – De controller berekent een statische planning voor gewichtsoverdracht, waarbij wordt toegewezen welke trainings-GPU welke parameter naar welke inferentie-GPU stuurt, en in welke volgorde.

Distributie van planning – De controller stuurt de planning naar alle trainings-GPU's.

Uitvoering – Na elke trainingsstap geeft de controller een signaal aan de trainings-GPU's om de overdrachten te starten.

Met de gedefinieerde workflow op hoog niveau is de grootste uitdaging hoe gewichtsoverdrachten efficiënt kunnen worden uitgevoerd op een schaal van biljoenen parameters. Hier beschrijven we de details van het uitvoeringstraject.

Parameters tijdens de training worden gedistribueerd volgens FSDP-plaatsingen. Met behulp van full_tensor() kunnen alle GPU's in een DeviceMesh de volledige parameter reconstrueren, waardoor ze allemaal als bron voor gewichtsoverdracht kunnen dienen.

Meerdere afzonderlijke DeviceMeshes vormen een mesh-groep. Omdat DeviceMeshes in dezelfde groep gescheiden zijn, interfereren hun overdrachten niet en kunnen ze volledig parallel draaien. Tussen mesh-groepen voegen we een globale barrière in om de volgorde af te dwingen.

De gewichtsoverdracht-pipeline overlapt host-naar-device-kopieën, tensorassemblage, fusie en kwantisering, en RDMA-overdrachten over MeshGroups.

We behandelen de overdracht van elke parametertensor als een taak. Het proces van gewichtsoverdracht maakt gebruik van meerdere soorten hardwarebronnen, daarom splitsen we een taak voor gewichtsoverdracht op in verschillende pijplijnfasen die elkaar in tijd overlappen:

Host-naar-apparaat memcpy — Indien FSDP gewicht naar CPU offloadt

Parameter voorbereiding — Reconstrueer volledig gewicht met full_tensor(), pas projectiefusie toe, kwantiseer indien nodig.

RDMA-overdracht — Zero-copy schrijven naar extern inferentie-GPU-geheugen

Globale barrière — Nadat alle full_tensor()-aanroepen zijn voltooid, synchroniseren over mesh-groepen heen met GLOO via Ethernet.

In de implementatie onderhouden we een FIFO-wachtrij met taken voor elke pijplijnfase. Wanneer de taak aan het begin van de wachtrij de fase voltooit, wordt deze verplaatst naar het einde van de volgende fasewachtrij.

full_tensor() en andere GPU-bewerkingen introduceren extra GPU-geheugengebruik. Om een out-of-memory fout te voorkomen, starten we de uitvoering van een taak alleen als de huidige lopende taken minder tijdelijk GPU-geheugen in beslag nemen dan een configureerbaar watermerk.

Verschillende ontwerpkeuzes maken ons systeem aanzienlijk sneller in uitvoering en gemakkelijker te onderhouden dan algemene open-sourceoplossingen.

Gewichtsoverdracht op basis van Rank0 verzamelt gewichten voordat ze naar de inferentienode worden gedistribueerd.

Point-to-pointcommunicatie draagt gewichten over van beide trainingsnodes naar de netwerkinterfaces van de inferentienode.

Een veelgebruikt patroon is om alle parameters via rang-0 GPU's te kanaliseren: verzamelen op trainings-rang-0, sturen naar inferentie-rang-0, en vervolgens weer verspreiden. Dit wordt al snel een knelpunt, beperkt door de PCIe-bandbreedte van één GPU en de NIC (bijv. 400 Gbps ≈ 50 GB/s).

Onze point-to-point-opstelling daarentegen stelt elke trainings-GPU in staat om rechtstreeks naar elke inferentie-GPU te sturen, waardoor het volledige netwerkfabric wordt verzadigd in plaats van slechts één link.

RPC-gebaseerde overdracht gebruikt een buffer aan de inferentiezijde en worker-aanroepen, terwijl eenzijdige overdracht gewichten rechtstreeks wegschrijft zonder deelname van de inference-worker.

Sommige systemen vertrouwen op het aanroepen van de update_weight()-methode van de inferentie-engine voor elke tensor. Dat betekent ingrijpende wijzigingen in de inferentiecode, plus overhead door RPC's, serialisatie en coördinatie van het besturingsvlak.

Met het RDMA WRITE-primitief updaten we gewichten geruisloos in het inferentie-GPU-geheugen, zonder extra kopieën. Er zijn geen berichten voor het besturingsvlak en geen CPU-besturingslogica bij betrokken. Er is geen aanpassing aan de inferentie-engine vereist.

Het proces van gewichtsoverdracht kan gebruikmaken van vier soorten hardwarebronnen: (1) Gegevensverplaatsing tussen host en apparaat (2) GPU-berekening voor projectiefusie en kwantisering (3) RDMA-netwerk voor gegevensvlak (4) Ethernet voor besturingsvlak.

Ons ontwerp splitst taken voor gewichtsoverdracht op in pijplijnfasen, waardoor eenvoudige overlapping tussen verschillende hardwarebronnen mogelijk is.

Sommige implementaties herberekenen een overdrachtsplanning bij elke trainingsstap, waarbij herhaaldelijk metadata wordt verzameld en instructies worden gedistribueerd. Dit voegt onnodige latentie toe aan het besturingsvlak.

Onze planning wordt eenmalig berekend bij initialisatie. Elke trainingsiteratie herhaalt simpelweg het plan: de controller geeft een "go"-signaal en de GPU's volgen hun vooraf toegewezen routes. De uitvoering is voorspelbaar en lichtgewicht.

Het is verleidelijk om het hele proces van gewichtsupdates te verweven in één monolithische functie: metadata verzamelen, naamovereenkomst, verzameling binnen knooppunten, projectiefusie, kwantisering, subslicing van de communicatiewereld, netwerkoverdracht tussen knooppunten. Het is moeilijk om dit correct te programmeren, en nog moeilijker om te optimaliseren.

In onze engineering scheiden we deze stappen als individuele componenten. Elke component kan afzonderlijk worden getest, beredeneerd en geoptimaliseerd.

Snelle, betrouwbare gewichtsoverdracht is een cruciaal onderdeel voor RL-fine-tuning op grote schaal. Door het RDMA WRITE-primitief, een statische overdrachtsplanning en gepijplijnde uitvoering te combineren, hebben we updates voor biljoenen parameters teruggebracht tot slechts 1,3 seconden op Kimi-K2. De aanpak is eenvoudig te begrijpen, gemakkelijk te onderhouden en vermijdt de knelpunten van traditionele ontwerpen.

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: eenzijdige overdrachten

Workflow op hoog niveau

Uitvoering van gewichtsoverdracht

DeviceMesh en Mesh-groepen

Taakpijplijn

Controle van GPU-geheugengebruik

Waarom het snel en eenvoudig is

Point-to-point-communicatie

Eenzijdige gegevensoverdracht

Pipelining

Statische planning

Schone scheiding

Conclusie

Citatie