Hochleistungs-GPU-Speichertransfer auf AWS Sagemaker Hyperpod

Moderne Deep-Learning-Infrastruktur erfordert häufig die Übertragung großer Datenmengen zwischen GPUs über Maschinen hinweg. Bei Perplexity stießen wir auf eine einzigartige technische Herausforderung: die effiziente Übertragung nicht zusammenhängender GPU-Speicherbereiche zwischen Maschinen mit maximal möglicher Geschwindigkeit. Unsere Zielplattform, AWS p5-Instanzen, bietet eine beeindruckende Netzwerkbandbreite von 3200 Gbit/s über 32 Netzwerkkarten. Dieser Artikel beschreibt unseren Weg zur Entwicklung einer maßgeschneiderten Hochleistungs-Netzwerklösung, die 97,1 % dieser theoretischen Bandbreite erreicht.

Die technische Herausforderung

Unser Anwendungsfall stellte mehrere wichtige technische Anforderungen:

• Übertragung nicht zusammenhängender Speicherblöcke zwischen entfernten GPUs mit hoher Bandbreite
• Fähigkeit, Knoten dynamisch zu Kubernetes-Bereitstellungen hinzuzufügen oder daraus zu entfernen, ohne laufende Vorgänge zu unterbrechen
• Unterstützung für Peer-to-Peer-Kommunikationsmuster

Obwohl die NCCL-Bibliothek von NVIDIA der De-facto-Standard für verteiltes Deep Learning ist, war sie für unseren Anwendungsfall nicht ideal:

• NCCL zeichnet sich durch kollektive Kommunikation aus, erfordert jedoch die Einrichtung einer statischen „Welt“, was bei der Anpassung der beteiligten Knoten einen Neustart des gesamten Clusters notwendig macht.
• Das synchrone Kommunikationsmodell von NCCL erhöht die Komplexität für unsere asynchrone Arbeitslast
• Wir wollten zur Optimierung eine direkte Kontrolle über unsere Speichertransfermuster
• Die Entwicklung unserer eigenen Lösung bot wertvolle Lernmöglichkeiten

Moderne Hochleistungsnetzwerke

Um unsere Lösung zu verstehen, untersuchen wir zunächst, wie sich moderne Hochleistungsnetzwerke von herkömmlichen Netzwerken unterscheiden.

Die meisten Netzwerke, die wir täglich nutzen, basieren auf TCP/IP-Protokollen, bei denen Anwendungen über Sockets mit der Netzwerkkarte durch den Betriebssystem-Kernel kommunizieren. Hochleistungsnetzwerke verwenden jedoch RDMA (Remote Direct Memory Access) – einen völlig anderen Hardware- und Software-Stack, der den direkten Speicherzugriff zwischen Maschinen ohne Einbeziehung der CPU ermöglicht.

AWS bietet den Elastic Fabric Adapter (EFA) an, eine benutzerdefinierte Netzwerkschnittstelle, die das kundenspezifische Amazon-Protokoll namens Scalable Reliable Datagram (SRD) implementiert. Im Gegensatz zur herkömmlichen TCP/IP-Vernetzung, bei der Daten mehrfach zwischen User Space, Kernel Space und Netzwerkpuffern kopiert werden müssen, ermöglicht EFA mit RDMA den direkten Datentransfer zwischen GPU-Speicher und Netzwerkkarte unter vollständiger Umgehung der CPU.

Philosophie des Hochleistungs-Netzwerkdesigns

Der Aufbau von Hochleistungs-Netzwerksystemen erfordert das Überdenken grundlegender Annahmen:

1. Pufferbesitz: Im Gegensatz zu herkömmlichen Sockets, bei denen der Kernel Netzwerkpuffer verwaltet und ein Kopieren zwischen User Space und Kernel Space erforderlich ist, erfordert RDMA, dass Anwendungen ihre Puffer selbst verwalten. Wenn eine Anwendung einen Netzwerkvorgang einleitet, überträgt sie den Pufferbesitz an die Netzwerkkarte, bis der Vorgang abgeschlossen ist, wodurch das Kopieren von Daten überflüssig wird.
2. Speicherregistrierung: Anwendungen müssen Speicherbereiche beim Betriebssystem-Kernel registrieren. Der Kernel richtet virtuelle Adresszuordnungen ein, die es CPU, GPUs und Netzwerkkarten ermöglichen, dieselben virtuellen Adressen zu verstehen. Diese Registrierung ist ein einmaliger Vorgang, der nachfolgende Zero-Copy-Datentransfers ermöglicht.
3. Steuerungsebene vs. Datenebene: Hochleistungsnetzwerke unterteilen Vorgänge in zwei Kategorien:
   1. Steuerungsebenen-Operationen (wie Verbindungseinrichtung und Speicherregistrierung) durchlaufen den Kernel, um die Sicherheit zu gewährleisten
   2. Datenebenen-Operationen (eigentlicher Datentransfer) umgehen für maximale Leistung den Kernel
4. Empfang vor Übertragung: Ohne Kernel-verwaltete Puffer müssen Anwendungen Empfangsvorgänge vorab posten und angeben, wo eingehende Daten platziert werden sollen. Dies ist ein grundlegender Wandel gegenüber dem Socket-Modell, bei dem Anwendungen Daten jederzeit empfangen können.
5. Polling-basierter Abschluss: Anstatt auf Kernel-Benachrichtigungen durch Mechanismen wie epoll zu warten, fragen Anwendungen direkt Hardware-Completion-Queues ab. Dies eliminiert den Overhead von Systemaufrufen und ermöglicht eine sofortige Reaktion auf abgeschlossene Vorgänge.
6. Hardware-Topologie-Bewusstsein: Das Verständnis und die Optimierung der Hardwaretopologie sind entscheidend für die Erzielung maximaler Leistung.

Verständnis der Hardwaretopologie

AWS p5-Instanzen verfügen über eine ausgeklügelte interne Architektur. Wie unten dargestellt, enthält jede Instanz zwei CPU-Sockel, die zwei NUMA-Knoten bilden, wobei jeder NUMA-Knoten eine Verbindung zu vier PCIe-Switches herstellt:

Vergrößern

Unter jedem PCIe-Switch finden wir:

• Vier 100 Gbit/s EFA-Netzwerkkarten
• Eine NVIDIA H100 GPU
• Eine NVMe-SSD

Die Datenpfade für TCP/IP- und RDMA-Übertragungen zeigen grundlegende Unterschiede in ihren Ansätzen:

Vergrößern

Bei TCP/IP (linke Seite) müssen Daten mehrfach kopiert werden:

1. Vom GPU-Speicher zum Anwendungspuffer im User Space
2. Vom Puffer im User Space zum Kernel-Puffer
3. Vom Kernel-Puffer zur Netzwerkkarte
4. Der umgekehrte Prozess findet auf der Empfängerseite statt

Jeder Kopiervorgang verbraucht CPU-Zyklen und Speicherbandbreite. Die Anwendung muss außerdem für jeden Netzwerkvorgang zwischen User Space und Kernel Space kontextwechseln.

Im Gegensatz dazu bietet RDMA (rechte Seite) einen echten Zero-Copy-Datentransfer:

1. Die Netzwerkkarte liest direkt aus dem GPU-Speicher
2. Daten gelangen direkt zur entfernten Netzwerkkarte
3. Die entfernte Netzwerkkarte schreibt direkt in den Ziel-GPU-Speicher
4. Die Anwendung muss lediglich eine Completion Queue (CQ) im User Space überprüfen, um festzustellen, wann der Transfer abgeschlossen ist

Mit RDMA und der richtigen Hardware-Kopplung erfordert die Datenübertragung zwischen zwei GPUs nur das Durchlaufen des lokalen PCIe-Switches und des Netzwerks:

Vergrößern

Im Gegensatz dazu müssen bei TCP/IP-Übertragungen Daten mehrfach durch den Hauptspeicher kopiert werden, was zu erheblichen Überlastungen des PCIe-Busses führt:

Vergrößern

Entwicklung mit libfabric

Wir haben libfabric verwendet – ein Framework, das eine generische Schnittstelle für Fabric-Dienste bietet. Unsere Implementierung verwendet zwei Arten von RDMA-Vorgängen:

• Zweiseitiges RDMA (SEND/RECV) für Steuernachrichten, die Metadaten über Speicherbereiche enthalten (z. B. Offsets und Größen).
• Einseitiges RDMA-WRITE für den eigentlichen Datentransfer, wobei jeder WRITE-Vorgang einen zusammenhängenden Speicherblock verarbeitet.

Weg zur Spitzenleistung

Unsere Entwicklung durchlief mehrere Phasen:

1. Implementierung einer einfachen unidirektionalen Nachrichtenübertragung mittels SEND/RECV
2. Erweiterung auf bidirektionale Kommunikation
3. Hinzufügen von GPUDirect RDMA-WRITE für direkten GPU-zu-GPU-Transfer
4. Erweiterung zur Bewältigung mehrerer gleichzeitiger Transfers
5. Einführung von Vorgangs-Queuing für Robustheit
6. Erreichung von 97,4 % Bandbreitenauslastung auf einer einzelnen Netzwerkkarte

Bei der Skalierung auf 32 Netzwerkkarten haben wir mehrere entscheidende Optimierungen implementiert:

1. Vorgangs-Queuing: Anstatt Vorgänge direkt an Netzwerkkarten zu übermitteln, verwalten wir eine Queue auf Anwendungsebene. Dies sorgt für Robustheit bei Netzwerküberlastung und vereinfacht das Programmiermodell.
2. Netzwerk-Warmup: Das Vorab-Herstellen von Verbindungen verbesserte die Startleistung.
3. Multithreading: Dedizierte Threads für die Netzwerkoperationen jeder GPU.
4. CPU-Core-Pinning: Binden von Threads an spezifische CPU-Kerne, um NUMA-Effekte und Cache-Fehler zu vermeiden.
5. State Sharding: Reduzierung von Konflikten zwischen Threads.
6. Vorgangs-Batching: Gemeinsames Übermitteln mehrerer Vorgänge.
7. Lazy Operation Posting: Vorgänge werden zuerst in der Anwendung in eine Warteschlange gestellt. Nach dem Abfragen der Completion Queue versuchen wir, ausstehende Vorgänge an die Netzwerkkarte zu übermitteln, um eine effiziente Nutzung der Netzwerkressourcen sicherzustellen.
8. NUMA-bewusste Ressourcenzuweisung: Zuweisung von libfabric-Ressourcen wie Completion Queues auf dem korrekten NUMA-Knoten, um die Latenz beim Speicherzugriff zu minimieren.

Durch diese Optimierungen erreichten wir eine endgültige Leistung von 3108 Gbit/s – 97,1 % der theoretischen maximalen Bandbreite.

Das Video unten zeigt unser Befehlszeilenprogramm in Aktion. Es überträgt nicht zusammenhängende GPU-Speicherblöcke an einen entfernten Knoten und erreicht eine Übertragungsgeschwindigkeit von 3108,283 Gbit/s – was die nahezu theoretische Bandbreitenauslastung der Netzwerkinfrastruktur demonstriert:

Fazit

Der Aufbau eines Hochleistungs-Netzwerksystems erfordert das Verständnis sowohl der Hardwarearchitektur als auch der Systemdesign-Prinzipien. Während Bibliotheken wie NCCL hervorragende Lösungen für allgemeine Muster bieten, sind manchmal benutzerdefinierte Lösungen für spezifische Anforderungen erforderlich.

Unsere Reise zeigt, dass das Erreichen einer nahezu theoretischen Netzwerkleistung durch sorgfältige Beachtung der Systemarchitektur, der Hardwaretopologie und verschiedener Optimierungstechniken möglich ist. Der Schlüssel liegt nicht nur im Verständnis einzelner Komponenten, sondern darin, wie sie interagieren, um ein vollständiges System zu bilden.

Die vollständige technische Analyse dieser Reise, einschließlich Implementierungsdetails und Codebeispielen, ist in unserem Open-Source-Repository verfügbar.

Zitat

@online{perplexity2025HighPerformanceGpuMemoryTransferOnAwsSagemakerHyperpod,
  author = {{Perplexity} and Lequn Chen and Phi Hoang},
  title = {{High-performance GPU memory transfer on AWS Sagemaker Hyperpod}},
  date = {2025-02-10},
  url = {https://www.perplexity.ai/hub/blog/high-performance-gpu-memory-transfer-on-aws-sagemaker-hyperpod}
}