Build on Trainium

Bei „Build on Trainium“ handelt es sich um ein Investitionsprogramm in Höhe von 110 Mio. USD, bei dem die KI-Forschung und die Hochschulbildung im Mittelpunkt stehen. Es soll die nächste Generation von Innovationen und die Entwicklung mit AWS Trainium unterstützen. AWS Trainium ist ein systolischer KI-Array-Chip, der eigens für die Weiterentwicklung modernster KI-Ideen und -Anwendungen entwickelt wurde. „Build on Trainium“ finanziert neuartige KI-Forschung mit Trainium und investiert in führende akademische Teams, um Innovationen in wichtigen Bereichen zu entwickeln, darunter u. a. neue Modellarchitekturen, ML-Bibliotheken, Optimierungen und groß angelegte verteilte Systeme. Diese auf mehrere Jahre angelegte Initiative legt den Grundstein für die Zukunft der KI, indem sie die akademische Gemeinschaft dazu anregt, die Open-Source-Community rund um Trainium zu nutzen, in sie zu investieren und dazu beizutragen. Durch die Kombination dieser Vorteile mit dem Neuron Software Development Kit (SDK) und der kürzlich erfolgten Einführung des Neuron Kernel Interface (NKI) können Trainium-Kunden nun in großem Maßstab in der Cloud Innovationen entwickeln.

Wir haben einen speziellen Trainium-Forschungscluster mit bis zu 40 000 Trainium-Chips eingerichtet, der über Amazon-EC2-Trn1-Instances verfügbar sein wird, die über ein einzelnes, nicht blockierendes Netzwerk im Peta-Bit-Maßstab mit Amazon EC2 UltraClusters verbunden sind. Forschungsteams und Studierende können über selbstverwaltete Reservierungen von Kapazitätsblöcken mithilfe von Amazon EC2 Capacity Blocks für ML auf diese Chips zugreifen.

Wir führen mehrere Runden der Aufforderung zur Einreichung von Vorschlägen (CFP) der Amazon Research Awards (ARA) für die breite Forschungsgemeinschaft durch. Ausgewählte Vorschläge erhalten AWS-Trainium-Guthaben und erhalten Zugang zum Trainium-Forschungscluster. Die Initiative „Build on Trainium“ begrüßt Forschungsvorschläge, die populäre Open-Source-ML-Bibliotheken und -Frameworks nutzen und einen Beitrag zu Open-Source-Projekten leisten, um die Ressourcen für die ML-Entwicklergemeinschaft zu verbessern.

Neuron Kernel Interface (NKI) ist eine neue Programmierschnittstelle für die AWS-KI-Chips Trainium und Inferentia. NKI bietet direkten Zugriff auf Hardware-Primitive und Anweisungen, die für AWS Trainium und Inferentia verfügbar sind. So können Forschende Compute-Kernel erstellen und optimieren, um die bestmögliche Leistung zu erzielen. Es handelt sich dabei um eine Python-basierte Programmierumgebung, die häufig verwendete Triton-ähnliche Syntax und Semantik auf Kachelebene nutzt. Mit NKI können Forschende Deep-Learning-Modelle um neue Funktionen, Optimierungen und wissenschaftliche Innovationen erweitern. Weitere Informationen finden Sie auf der Seite mit der NKI-Dokumentation.

Im Device Realization Lab des MIT verwenden wir AWS Trainium, um die Grenzen der medizinischen KI-Forschung zu erweitern. Unsere Modelle zur 3D-Ultraschallsegmentierung und Schallgeschwindigkeitsschätzung trainieren schneller und effizienter als je zuvor, wodurch die Versuchszeit um mehr als die Hälfte reduziert wird und gleichzeitig eine hochmoderne Genauigkeit erreicht wird. AWS Trainium hat es uns ermöglicht, unsere Forschung auf eine Weise zu skalieren, die mit herkömmlichen GPU-Systemen nicht machbar war. Durch das Training unserer vollständig konvolutionellen 3D-Neuralnetzwerke auf AWS Trainium (trn.32xlarge) erreichten wir eine Leistung auf dem neuesten Stand der Technik mit einem um 50 % höheren Durchsatz und niedrigeren Kosten im Vergleich zu NVIDIA A100-Instances. Mithilfe eines Trainium-Clusters mit 32 Knoten führten wir effizient über 180 Ablationsexperimente durch, wodurch die gesamte Trainingszeit von Monaten auf Wochen reduziert und die medizinische KI-Innovation am MIT beschleunigt wurde. In Zukunft planen wir, Trainium zu verwenden, um KI-Agentenmodelle zu trainieren, die den digitalen Ultraschall-Workflow bedienen und automatisieren können, was dem Arzt viel Zeit spart und die Patienten besser versorgen kann.“

„Die Forschungsgruppe CMU Catalyst arbeitet an der Optimierung von ML-Systemen. Unser Projekt zielt darauf ab, die Optimierung über verschiedene ML-Systeme hinweg zu vereinfachen. Trainium ist einzigartig, da es sowohl eine Steuerung auf niedriger Ebene als auch eine zugängliche Programmierschnittstelle über das Neuron Kernel Interface (NKI) bietet.

Mit Unterstützung von AWS im Rahmen des ‚Build on Trainium‘-Programms war unser Forscher in der Lage, fortschrittliche Optimierungen für einen wichtigen Kernel – FlashAttention – zu untersuchen. Was uns am meisten überraschte, war die Geschwindigkeit, mit der wir iterieren konnten: In nur einer Woche erreichten wir mithilfe der öffentlich zugänglichen NKI-, Neuron Profiler- und Architekturdokumentation bedeutende Verbesserungen zusätzlich zum bisherigen Stand der Technik. Die Kombination aus leistungsstarken Tools und klaren Hardware-Einblicken machte unserem Team eine ausgeklügelte Optimierung auf niedriger Ebene zugänglich.

AWS Trainium und Neuron Kernel Interface (NKI) ermöglicht es Forschern wie uns, schneller innovativ zu sein und Hindernisse zu beseitigen, die normalerweise hardwarespezifische Optimierungsarbeiten verlangsamen.“

 

„Durch das ‚Build on Trainium‘-Programm hat sein Team vollen Zugriff auf den neuen NKI-Open-Source-Compiler-Stack von AWS Neuron erhalten – einschließlich direkter Einblicke in die Trainium ISA und APIs für eine präzise Planung und Speicherzuweisung. Dieses Maß an Sichtbarkeit und Kontrolle ermöglicht es seinen Studierenden, die Optimierungsmöglichkeiten einfacher zu analysieren und leistungsfähige Implementierungen effektiver zu entdecken.“

Christopher Fletcher, Associate Professor of Computer Science, University of California, Berkeley
 

„Der Zugang zu AWS Trainium und Inferentia hat maßgeblich dazu beigetragen, unsere Forschung und Ausbildung im Bereich groß angelegter, effizienter KI-Systeme voranzutreiben. Wir verwenden diese Plattformen für Schulungen und Inferenzoptimierungen von Mixture-of-Experts-Modellen, indem wir Prototypen für neue Laufzeit- und Planungstechniken entwickeln, die die Skalierbarkeit, Effizienz und Portabilität auf neuen Beschleunigerarchitekturen verbessern. Mithilfe des Neuron Developer-Stacks entwickeln UIUC-Forscher neue Laufzeit- und Planungstechniken, die die Effizienz und Portabilität von KI-Workloads verbessern. Das Team ist besonders beeindruckt von der Offenheit des Neuron Developer-Stacks, die diese Plattformen für die Laufzeitforschung wertvoll macht und Innovationen in Bezug auf Sparsität, Speicherhierarchien und Kommunikationseffizienz ermöglicht, die über herkömmliche GPU-Architekturen hinausgehen.“

„Durch die Nutzung von AWS Trainium und dem ‚Build on Trainium‘-Programm konnten meine Studenten und ich unsere Quantenschaltungssimulationen erheblich beschleunigen. Das Projekt brachte eine starke Gruppe von Studenten zusammen, die gemeinsam einen Hochleistungssimulator bauten, der tiefere Experimente und praktisches Lernen in einem Umfang ermöglichte, der zuvor einfach nicht möglich war.“

„Unser Forschungsteam bei UTS untersucht die Integration von Tree-Ring-Wasserzeichenalgorithmen, indem es benutzerdefinierte Neuron-NKI-Kernel entwickelt. Der Zugriff auf den Open-Source-Neuron-Stack über das ‚Build on Trainium‘-Programm war transformativ. Es gibt uns einen beispiellosen Einblick in die Trainium-Architektur und die Möglichkeit, direkt auf Hardwareebene zu arbeiten. Der Zugang zu Trainium hat es unserem Team ermöglicht, unsere Wasserzeichen-Workloads erheblich zu beschleunigen, die Iterationszyklen zu reduzieren und komplexere Modelle und Techniken zu untersuchen. Diese Zugriffstiefe ermöglicht es unseren Forschern, neue Ideen zu prototypisieren, mit Optimierungen auf niedriger Ebene zu experimentieren und die Grenzen dessen zu erweitern, was Wasserzeichensysteme mit modernen KI-Beschleunigern erreichen können.“