Amazon EC2 Trn2-Instanzen und Trn2 UltraServer für KI/ML-Training und -Inferenz sind jetzt verfügbar

von Jeff Barr übersetzt durch Tobias Wasle

Die neuen Amazon Elastic Compute Cloud (Amazon EC2) Trn2-Instanzen und Trn2 UltraServer([EN]) sind die leistungsstärksten Optionen von Amazon EC2 für ML-Training und -Inferenz. Auf Basis der zweiten Generation der AWS Trainium-Chips (AWS Trainium2) sind die Trn2-Instanzen 4-mal schneller, bieten 4-mal mehr Speicherbandbreite und 3-mal mehr Speicherkapazität als Trn1-Instanzen der ersten Generation. Trn2-Instanzen bieten ein 30-40% besseres Preis-Leistungs-Verhältnis als die aktuelle Generation der GPU-basierten EC2 P5e- und P5en-Instanzen.

Zusätzlich zu den 16 Trainium2-Chips verfügt jede Trn2-Instanz über 192 vCPUs, 2 TiB Arbeitsspeicher und 3,2 Tbps Elastic Fabric Adapter (EFA) v3 Netzwerkbandbreite, die bis zu 35% geringere Latenz als die vorherige Generation bietet.

Die Trn2 UltraServer, die ein völlig neues Compute-Angebot darstellen, verfügen über 64 Trainium2-Chips, die mit einer breitbandigen, latenzarmen NeuronLink-Verbindung verbunden sind, für maximale Inferenz- und Trainingsleistung bei neuartigen Foundation-Modellen.

Zehntausende von Trainium-Chips treiben bereits Amazon- und AWS-Dienste an. Zum Beispiel unterstützten über 80.000 AWS Inferentia- und Trainium1-Chips den Rufus-Einkaufsassistenten ([EN]) am letzten Prime Day. Die latenzoptimierten Versionen der Llama 3.1 405B([EN], extern) und Claude 3.5 Haiku([EN], extern) Modelle auf Amazon Bedrock laufen bereits auf Trainum2-Chips.

Hoch, weit und raus

Das anhaltende Wachstum in Größe und Komplexität großer Sprach- und Grundmodelle wird durch innovative Formen von Rechenleistung ermöglicht, die in ebenso innovativen Architekturen zusammengesetzt werden. In der Vergangenheit haben wir Skalierung immer in zwei Dimensionen gesehen: Hochskalieren (Verwendung eines größeren Computers) und in die breite Skalieren (Verwendung von mehr Computern). Mit dem Blick auf Trainium2-Chips, die Trn2-Instanz und die noch größeren Rechenangebote die wir noch vorstellen, verschmelzen diese Dimensionen auf unterschiedlichen Ebenen des Gesamtsystems. Um das zu illustrieren, schauen wir in die Bestandteile von Trn2, beginnend beim NeuronCore und skalierend bis zum UltraCluster:

NeuronCores bilden das Herzstück des Trainium2-Chips. Jeder NeuronCore der dritten Generation umfasst eine skalare Engine (1 Eingang zu 1 Ausgang), eine Vektor-Engine (mehrere Eingänge zu mehreren Ausgängen), eine Tensor-Engine (systolische Array-Multiplikation, Faltung und Transposition) und einen GPSIMD-Kern (General Purpose Single Instruction Multiple Data.

Jeder Trainium2-Chip beherbergt acht dieser NeuronCores und 96 GiB High Bandwidth Memory (HBM) und unterstützt 2,9 TB/Sekunde HBM-Bandbreite. Die Kerne können einzeln adressiert und verwendet werden, oder Paare von physischen Kernen können zu einem einzigen logischen Kern gruppiert werden. Ein einzelner Trainium2-Chip liefert bis zu 1,3 Petaflops dense FP8-Rechenleistung und bis zu 5,2 Petaflops sparse FP8-Rechenleistung und kann dank automatischer Neuordnung der HBM-Warteschlange 95% Auslastung der Speicherbandbreite erreichen.

Jede Trn2-Instanz beherbergt wiederum 16 Trainium2-Chips. Das sind insgesamt 128 NeuronCores, 1,5 TiB HBM und 46 TB/Sekunde HBM-Bandbreite. Insgesamt ergibt das bis zu 20,8 Petaflops dense FP8-Rechenleistung und bis zu 83,2 Petaflops sparse FP8-Rechenleistung. Die Trainium2-Chips sind über NeuronLink in einem 2D-Torus für hohe Bandbreite und geringe Latenz bei der Chip-zu-Chip-Kommunikation mit 1 TB/Sekunde verbunden.

Ein UltraServer beherbergt vier Trn2-Instanzen, die mit latenzarmer, breitbandinger NeuronLink-Verbindung verbunden sind. Das sind 512 NeuronCores, 64 Trainium2-Chips, 6 TiB HBM und 185 TB/Sekunde HBM-Bandbreite. Mathematisch ergibt das bis zu 83 Petaflops dense FP8-Rechenleistung und bis zu 332 Petaflops sparse FP8-Rechenleistung. Zusätzlich zum 2D-Torus, der NeuronCores innerhalb einer Instanz verbindet, sind Kerne an korrespondierenden XY-Positionen in jeder der vier Instanzen in einem Ring verbunden. Für die Inferenz helfen UltraServer, branchenführende Antwortzeiten zu liefern, um die beste Echtzeiterfahrung zu realisieren. In Trainings-Anwendungsgebieten beschleunigen UltraServer das Modelltraining und erhöhen die Traininsgeffizienz durch schnellerer kollektiver Kommunikation im Vergleich zu eigenständigen Instanzen. UltraServer sind darauf ausgelegt, Training und Inferenz auf Billionen-Parameter-Ebene und darüber hinaus zu unterstützen. Sie sind als Vorschau verfügbar und Sie können uns kontaktieren([EN]) , um an der Vorschau teilzunehmen.

Trn2-Instanzen und UltraServer werden in EC2 UltraCluster mit Zugriff aufAmazon FSx for Lustre für High Performance Computing bereitgestellt. Damit wird skalierbares, verteiltes Training über Zehntausende von Trainium-Chips in einem einzigen Petabit-Netzwerk ohne Blockierung zu möglich.

Verwendung von Trn2-Instanzen

Trn2-Instanzen sind ab heute für den Produktionseinsatz in der AWS-Region US East (Ohio) verfügbar und können als Amazon EC2 Kapazitätsblöcke für ML reserviert werden. Sie können bis zu 64 Instanzen für bis zu sechs Monate reservieren, wobei Reservierungen bis zu acht Wochen im Voraus akzeptiert werden, mit sofortigen Startzeiten und der Möglichkeit, Ihre Reservierungen bei Bedarf zu verlängern([EN]) . Um mehr zu erfahren, lesen Sie Ankündigung von Amazon EC2 Capacity Blocks for ML zur Reservierung von GPU-Kapazität für Ihre Machine-Learning-Workloads([EN]) .

Auf der Softwareseite finden Sie mit den AWS Deep Learning AMIs einen guten Einstieg. Diese Images sind vorkonfiguriert mit den Frameworks und Tools, die Sie wahrscheinlich bereits kennen und verwenden: PyTorch([EN]), JAX ([EN]) sowie viele andere.

Wenn Sie das AWS Neuron SDK([EN]) zum Erstellen Ihrer Anwendungen verwenden, können Sie diese übertragen und für die Verwendung auf Trn2-Instanzen neu kompilieren. Das SDK integriert sich nativ mit JAX, PyTorch und wichtigen Bibliotheken wie Hugging Face, PyTorch Lightning und NeMo. Neuron enthält vorgefertigte Optimierungen für verteiltes Training und Inferenz mit den Open-Source-PyTorch-Bibliotheken NxD Training und NxD Inference und bietet gleichzeitig tiefe Einblicke für Profiling und Debugging. Neuron unterstützt auch OpenXLA, einschließlich stabilem HLO und GSPMD, was es PyTorch/XLA- und JAX-Entwicklern ermöglicht die Compiler-Optimierungen von Neuron für Trainium2 zu nutzen.

Über die Autoren

Jeff Barr ist Chief Evangelist bei AWS. Er startete diesen Blog im Jahr 2004 und schreibt seitdem praktisch ununterbrochen Beiträge.