NVIDIA: Reshard kann die Refitting-Zeit bei Reinforcement Learning verkürzen

Warum Refitting bei Reinforcement Learning zum Engpass wird

Beim Post-Training mit Reinforcement Learning erzeugen Systeme Trainingsverläufe, indem sie mit dem aktuellen Policy-Modell Inferenz ausführen. Training und Inferenz können parallel auf getrennten GPU-Ressourcen laufen. Die Inferenz-GPUs müssen jedoch regelmäßig die neuesten Gewichte erhalten, die von den Trainings-GPUs erzeugt wurden. NVIDIA bezeichnet diesen Vorgang als Refit.

Die beiden Seiten verteilen Modellgewichte häufig unterschiedlich, weil sie verschiedene Rechenanforderungen haben. Ihre Konfigurationen können unterschiedliche Kombinationen aus Daten-, Kontext-, Tensor-, Experten- und Pipeline-Parallelismus verwenden. Beim Refitting muss außerdem jedes Modellgewicht auf den Inferenz-GPUs aktualisiert werden, sodass das Kommunikationsvolumen mit der Gesamtgröße des Modells wächst.

Bestehende Implementierungen nutzen häufig Gather und Broadcast. Die Trainings-GPUs führen zunächst ein All-Gather ihrer lokalen Shards durch, um einen vollständigen Parameter zu rekonstruieren, und übertragen diesen vollständigen Tensor anschließend per Broadcast an die Inferenz-GPUs. Jede Inferenz-GPU behält nur den benötigten Teil. Dadurch werden auch Daten übertragen, die am Ziel nicht benötigt werden.

NVIDIA nennt Kimi K3 und Qwen 3.8 2.4T A95B als Beispiele. Bei BF16-Gewichten mit zwei Byte pro Parameter schätzt das Unternehmen die Größe der Gewichte auf etwa 5.6 TB beziehungsweise 4.8 TB. Unter Einbeziehung von All-Gather und Broadcast ergeben sich geschätzte logische Nutzdaten von etwa 11.2 TB beziehungsweise 9.6 TB.

Reshard überträgt sich überschneidende Bereiche direkt

NVIDIAs Ansatz formuliert das Refitting als DTensor-zu-DTensor-Transfer. Die verteilte Tensor-Abstraktion von PyTorch, kurz DTensor, beschreibt die globale Form eines Tensors, die logische Anordnung der GPUs sowie die Art, wie der Tensor über diese Anordnung verteilt oder repliziert wird.

Mit dieser Beschreibung wird das Trainings-Layout zur Quelle und das Inferenz-Layout zum Ziel. Das Reshard-Kollektiv berechnet, welche Teile der Quell-Shards die Ziel-Shards überlappen. Dadurch können die benötigten Bereiche direkt übertragen werden, statt zunächst einen vollständigen Zwischentensor zu rekonstruieren.

NVIDIA veranschaulicht das Konzept anhand einer Gewichtsmatrix mit 8 x 8 Elementen. Zwei Trainings-GPUs halten Zeilen-Shards mit den Abmessungen 4 x 8, während vier Inferenz-GPUs Spalten-Shards mit den Abmessungen 8 x 2 benötigen. Jedes Ziel kann von jeder Quelle einen Bereich von 4 x 2 Elementen empfangen und so genau den benötigten Shard zusammensetzen.

Das Unternehmen beschreibt Reshard als Aufruf für kollektive Kommunikation bei beliebigen DTensor-zu-DTensor-Transfers. Wie andere Kollektivoperationen ermöglicht es Anwendungen, das Kommunikationsmuster als eine Operation aufzurufen, während die zugrunde liegende Bibliothek die Transferlogik verwaltet.

Open-Source-Implementierung unterstützt unterschiedliche GPU-Layouts

NCCL M2N Reshard ist NVIDIAs quelloffene Referenzimplementierung des Kollektivs und über das GitHub-Repository NVIDIA/nccl-extensions verfügbar. Die auf NCCL basierende Lösung überträgt einen auf M Quell-GPUs verteilten Tensor in das von N Ziel-GPUs benötigte Layout. Beim Refitting mit Reinforcement Learning bilden die Trainings-GPUs die Quellgruppe und die Inferenz-GPUs die Zielgruppe.

Laut der KI-generierten Zusammenfassung von NVIDIA zu dem Beitrag zeigten Benchmarks in NeMo RL auf NVIDIA-DGX-H100- und DGX-GB200-Systemen beim Refitting mit Reshard einen Geschwindigkeitsvorteil von 2.7x bis 5.7x gegenüber herkömmlichem Refitting mit Gather und Broadcast. Getestet wurde dies unter anderem mit Qwen 3 235B und DeepSeek v3 671B.

Dieselbe Zusammenfassung besagt, dass ein hybrider Ansatz in NeMo RL Reshard für Feed-Forward-Projektionsgewichte verwendet, die den Großteil der Nutzdaten ausmachen, während für kleinere Gewichte Gather und Broadcast beibehalten werden. NVIDIA verweist Entwickler für unterstützte Backends, Layouts und Präzisionsmodi auf den Refitting-Leitfaden von NeMo RL sowie für Implementierungsdetails auf die Konfigurationsanforderungen von NCCL Reshard.

Für Nutzer, die verteilte Reinforcement-Learning-Systeme entwickeln, bietet die Ankündigung eine Möglichkeit, unnötigen Datenverkehr beim Transfer von Gewichten zu reduzieren, während Modelle und GPU-Installationen wachsen. Der praktische Einsatzbereich hängt von den Layouts und Implementierungsoptionen ab, die von den jeweiligen NeMo-RL- und NCCL-Konfigurationen unterstützt werden.