Gewichtsübertragung für RL-Post-Training in unter 2 Sekunden

Wir haben kürzlich 1,3-sekündige maschinenübergreifende Parameteraktualisierungen für Kimi-K2 (1T-Parameter) erreicht, wobei die Gewichte von 256 Trainings-GPUs (BF16) auf 128 Inferenz-GPUs (FP8) übertragen wurden.

Beim asynchronen Reinforcement-Learning-Fine-Tuning laufen Training und Inferenz auf getrennten GPUs. Nach jedem Trainingsschritt müssen neue Gewichte an die Inferenzknoten gesendet werden. Viele bestehende Frameworks benötigen dafür bei Modellen mit Billionen von Parametern mehrere Sekunden oder sogar Minuten.

Durch die Nutzung von RDMA-Punkt-zu-Punkt-Kommunikation können wir die Gewichtsübertragung extrem beschleunigen, ohne die Inferenz-Engine zu ändern, und den Code einfacher schreiben und warten.

Unsere Lösung basiert auf RDMA WRITE, einem einseitigen Primitiv, bei dem die Quelle direkt in den GPU-Speicher des Ziels schreibt.

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

Die Zielseite wird über die Übertragung nicht einmal benachrichtigt. Dies ermöglicht uns latenzarme Übertragungen mit hohem Durchsatz und ohne Kopieraufwand, die von den Trainingsknoten gesteuert werden, ohne jegliche Steuerungslogik auf den Inferenzknoten.

Metadatensammlung – Der Controller erfasst Parametermetadaten von allen Trainings- und Inferenz-GPUs.

Berechnungsplanung – Der Controller berechnet einen statischen Zeitplan für die Gewichtsübertragung und ordnet zu, welche Trainings-GPU welches Parameter an welche Inferenz-GPU sendet und in welcher Reihenfolge.

Zeitplanverteilung – Der Controller sendet den Zeitplan an alle Trainings-GPUs.

Ausführung – Nach jedem Trainingsschritt signalisiert der Controller den Trainings-GPUs, die Übertragungen zu starten.

Nach der Definition des übergeordneten Workflow besteht die größte Herausforderung darin, Gewichtsübertragungen bei einer Größenordnung von Billionen von Parametern effizient auszuführen. Hier beschreiben wir die Details des Ausführungspfads.

Die Parameter beim Training werden gemäß FSDP-Platzierungen verteilt. Mithilfe von full_tensor() können alle GPUs in einem DeviceMesh das vollständige Parameter rekonstruieren, sodass alle als Quelle für die Gewichtsübertragung dienen können.

Mehrere disjunkte DeviceMeshes bilden eine Mesh-Gruppe. Da DeviceMeshes in derselben Gruppe disjunkt sind, interferieren ihre Übertragungen nicht und können vollständig parallel ablaufen. Zwischen Mesh-Gruppen fügen wir eine globale Sperre ein, um die Reihenfolge zu erzwingen.

Die Gewichtsübertragungspipeline überlappt Host-zu-Device-Kopien, Tensor-Zusammenstellung, Fusion und Quantisierung sowie RDMA-Übertragungen über MeshGroups hinweg.

Wir behandeln die Übertragung jedes Parametertensors als Aufgabe. Der Gewichtsübertragungsprozess nutzt verschiedene Arten von Hardware-Ressourcen, daher unterteilen wir eine Gewichtsübertragungsaufgabe in verschiedene Pipeline-Stufen, die sich zeitlich überschneiden:

Host-zu-Gerät-Memcpy — Wenn FSDP Gewicht auf CPU auslagert

Parametervorbereitung — Rekonstruktion des vollen Gewichts mit full_tensor(), Anwendung der Projektionsfusion, bei Bedarf Quantisierung.

RDMA-Übertragung — Zero-Copy-Schreiben in den Speicher der entfernten Inferenz-GPU

Globale Sperre — Nachdem alle full_tensor()-Aufrufe abgeschlossen sind, Synchronisierung über Mesh-Gruppen hinweg mittels GLOO über Ethernet.

In der Implementierung verwalten wir eine FIFO-Warteschlange von Aufgaben für jede Pipeline-Stufe. Sobald der Kopf der Aufgabenwarteschlange die Stufe abschließt, wird er an das Ende der Warteschlange der nächsten Stufe verschoben.

full_tensor() und andere GPU-Operationen führen zu zusätzlicher GPU-Speicherauslastung. Um Out-of-Memory-Fehler zu vermeiden, starten wir die Ausführung einer Aufgabe nur dann, wenn die aktuell laufenden Aufgaben weniger temporären GPU-Speicher beanspruchen als ein konfigurierbarer Schwellenwert.

Mehrere Designentscheidungen machen unser System deutlich schneller im Betrieb und einfacher zu warten als gängige Open-Source-Lösungen.

Die auf Rank0 basierende Gewichtsübertragung sammelt Gewichte, bevor sie an den Inferenzknoten verteilt werden.

Die Punkt-zu-Punkt-Kommunikation überträgt Gewichte von beiden Trainingsknoten an die Netzwerkschnittstellen des Inferenzknotens.

Ein gängiges Muster ist das Durchleiten aller Parameter über Rang-0-GPUs: Sammeln auf Trainings-Rang-0, Senden an Inferenz-Rang-0, dann erneutes Verteilen. Dies wird schnell zu einem Engpass, begrenzt durch die PCIe-Bandbreite und die Netzwerkkarte (NIC) einer einzelnen GPU (z. B. 400 Gbit/s ≈ 50 GB/s).

Im Gegensatz dazu ermöglicht unser Punkt-zu-Punkt-Setup jeder Trainings-GPU, direkt an jede Inferenz-GPU zu senden, wodurch das gesamte Netzwerk-Fabric gesättigt wird, anstatt nur eine einzelne Verbindung.

Die RPC-basierte Übertragung verwendet einen Buffer auf der Inferenzseite und Worker-Aufrufe, während die einseitige Übertragung Gewichte direkt schreibt, ohne Beteiligung des Inferenz-Workers.

Einige Systeme verlassen sich darauf, für jeden Tensor die Methode update_weight() der Inferenz-Engine aufzurufen. Das bedeutet invasive Änderungen am Inferenz-Code sowie Overhead durch RPCs, Serialisierung und Koordinierung der Steuerungsebene.

Mit dem RDMA WRITE-Primitiv aktualisieren wir Gewichte im Inferenz-GPU-Speicher lautlos, ohne zusätzliche Kopien. Es sind keine Nachrichten der Steuerungsebene und keine CPU-Steuerungslogik beteiligt. Es ist keine Modifikation der Inferenz-Engine erforderlich.

Der Gewichtsübertragungsprozess kann vier Arten von Hardware-Ressourcen nutzen: (1) Host-Gerät-Datenbewegung (2) GPU-Berechnung für Projektionsfusion und Quantisierung (3) RDMA-Netzwerk für die Datenebene (4) Ethernet für die Steuerungsebene.

Unser Design unterteilt Gewichtsübertragungsaufgaben in Pipeline-Stufen, was eine einfache Überlappung über verschiedene Hardware-Ressourcen hinweg ermöglicht.

Einige Implementierungen berechnen bei jedem Trainingsschritt einen Übertragungszeitplan neu und sammeln wiederholt Metadaten und verteilen Anweisungen. Dies fügt unnötige Latenz auf der Steuerungsebene hinzu.

Unser Zeitplan wird bei der Initialisierung einmal berechnet. Jede Trainingsiteration spielt den Plan einfach erneut ab: Der Controller sendet ein „Go“-Signal und die GPUs folgen ihren vordefinierten Routen. Die Ausführung ist vorhersehbar und leichtgewichtig.

Es ist verlockend, den gesamten Gewichtsaktualisierungsprozess in einer monolithischen Funktion zu verflechten: Metadaten sammeln, Namensabgleich, Intra-Node-Sammeln, Projektionsfusion, Quantisierung, Sub-Slicing der Kommunikationswelt, Inter-Node-Netzwerkübertragung. Dies ist schwer korrekt zu programmieren und noch schwerer zu optimieren.

In unserer Entwicklung trennen wir diese Schritte als individuelle Komponenten. Jede Komponente kann isoliert getestet, bewertet und optimiert werden.

Schnelle, zuverlässige Gewichtsübertragung ist ein kritischer Baustein für das RL-Fine-Tuning im großen Maßstab. Durch die Kombination des RDMA WRITE-Primitivs, eines statischen Übertragungszeitplans und pipelined Ausführung haben wir Aktualisierungen bei Billionen von Parametern auf Kimi-K2 auf nur 1,3 Sekunden reduziert. Der Ansatz ist logisch einfach, leicht zu warten und vermeidet die Engpässe herkömmlicher Designs.

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: einseitige Übertragungen

Übergeordneter Workflow

Ausführung der Gewichtsübertragung

DeviceMesh und Mesh-Gruppen

Aufgaben-Pipeline

Steuerung der GPU-Speicherauslastung

Warum es schnell und einfach ist

Punkt-zu-Punkt-Kommunikation

Einseitige Datenübertragung

Pipelining

Statischer Zeitplan

Saubere Trennung

Fazit

Zitat