AWS beschreibt MoE-Reinforcement-Learning auf EKS und meldet 40 % höheren Rollout-Durchsatz
AWS beschreibt eine Architektur für Reinforcement Learning mit Mixture-of-Experts-Modellen (MoE), die Amazon EKS, Elastic Fabric Adapter (EFA) und Amazon S3 verbindet. Dazu nennt das Unternehmen ein Messergebnis: Ein Stack mit DeepEP über EFA erreichte in einem bestimmten Vergleich einen um 40 % höheren aggregierten Durchsatz bei Reinforcement-Learning-Rollouts. Die Zahl ist eine Angabe von AWS für diesen Versuchsaufbau und sollte zusammen mit den Bedingungen des Vergleichs gelesen werden.
Aufteilung der Aufgaben
Im beschriebenen EKS-Cluster übernehmen verschiedene Instanzgruppen verschiedene Arbeiten. GPU-Instanzen erzeugen Rollouts, führen Inferenz mit dem Belohnungsmodell aus und trainieren die Policy. CPU-Instanzen sind für Umgebungen und Vorverarbeitung vorgesehen. Speicheroptimierte Instanzen halten Experience-Puffer und Checkpoint-Caches bereit. Damit ordnet die Architektur der Erzeugung, Zwischenspeicherung und Nutzung von Erfahrungen jeweils einen Platz im Arbeitsablauf zu.
Rollout-Erfahrungen gelangen zunächst in einen speicheroptimierten Puffer. Das Policy-Training verarbeitet daraus Batches und veröffentlicht aktualisierte Checkpoints. Checkpoints und abgeschlossene Artefakte werden in Amazon S3 gespeichert. EKS übernimmt laut AWS die Planung, Skalierung, Wiederherstellung nach Ausfällen und Koordination der Worker-Gruppen. Diese Aufgaben verbinden die laufende Erzeugung von Erfahrungen mit den nachfolgenden Trainingsschritten.
Kommunikation zwischen Experten
Für die Kommunikation innerhalb von GPU-Instanzen nutzt die Architektur NVLink und NVSwitch. Zwischen Instanzen kommt EFA zum Einsatz. DeepEP verwendet einen Dispatch-Kernel, der Tokens an entfernte Experten weiterleitet, und einen Combine-Kernel, der die verarbeiteten Tokens einsammelt. Der beschriebene Pfad zwischen Knoten nutzt libfabric über EFA. AWS gibt an, Kommunikationsprimitive von DeepEP auf libfabric umgestellt und DeepEP v2 mit nativer EFA-Unterstützung ausgestattet zu haben.
Die Bestandteile erfüllen unterschiedliche Funktionen. EKS koordiniert die Worker-Gruppen; der Puffer hält Rollout-Erfahrungen für das Training bereit, und S3 speichert Checkpoints und fertige Artefakte. NVLink und NVSwitch bedienen die Kommunikation innerhalb von GPU-Instanzen, EFA die Kommunikation zwischen Instanzen. Die Dispatch- und Combine-Operationen von DeepEP betreffen den Weg der Tokens zu den Experten und zurück. So lässt sich der gemessene Ablauf einordnen, ohne jeder Komponente einen separat gemessenen Leistungsgewinn zuzuschreiben.
Umfang des Durchsatzvergleichs
Für den von AWS berichteten Vergleich liefen 48 P5en-Instanzen mit einem superdünn besetzten MoE-Modell. Davon waren 16 Instanzen dem Training und 32 der Inferenz zugeordnet. Als Ausgangspunkt diente ein auf Slime basierender Stack ohne durch EFA beschleunigten Expertenparallelismus. Die verbesserte Konfiguration nutzte einen Stack mit DeepEP über EFA. Für diese Gegenüberstellung meldet AWS 40 % mehr aggregierten Reinforcement-Learning-Rollout-Durchsatz.
Gemessen wurde damit der aggregierte Rollout-Durchsatz dieser beiden Konfigurationen. Aus den vorliegenden Angaben ergibt sich kein entsprechender Wert für andere Modelle, Instanztypen oder Bereitstellungen. Ebenso lässt sich der Anteil von DeepEP über EFA an den 40 % nicht isolieren, weil sich die Software-Stacks auch in anderen Komponenten unterschieden. Belastbar ist die Aussage über den beobachteten Unterschied der konkret beschriebenen Setups.
Wiederaufnahme und Versionen
AWS erläutert außerdem, dass unterbrochene Rollout-Aufgaben auf Spot Instances in eine Warteschlange zurückkehren und neu zugewiesen werden können, während andere Rollout-Worker weiter Erfahrungen erzeugen. Das beschreibt einen Weg, unterbrochene Arbeit erneut zu verteilen. Gemessene Kosteneinsparungen durch Spot Instances nennt die Quelle nicht; aus dem Mechanismus folgt daher keine bezifferte Ersparnis.
Die Anleitung führt unter ihren Anforderungen Amazon EKS 1.31 oder neuer, den EFA-Installer 1.49 mit dem AWS-OFI-NCCL-Plugin, DeepEP 2.0.0, NCCL 2.31.2, SGLang 0.5.17 sowie PyTorch 2.12.1 mit CUDA 13.0 auf. Diese Versionsangaben beschreiben den Softwarekontext des dargestellten Vorgehens. Architektur und Messergebnis stehen im technischen Beitrag von AWS.