GPUs werden seit langem verwendet, um Berechnungen zu beschleunigen, und in den letzten Jahren sind sie mit der Unterstützung spezifischer Datenformate, die für Algorithmen des maschinellen Lernens typisch sind, überwachsen, da sie praktisch ihre eigenen Grafikblöcke verloren haben. Aber in naher Zukunft könnten sie von spezialisierten KI-Prozessoren in vielerlei Hinsicht übertroffen werden, einschließlich der neuen AWS-Entwicklung, dem Trainium-Chip.
Auf der AWS Re: Invent stellte das Unternehmen mit seinen P3dn (Nvidia V100) und P4 (Nvidia A100) Instanzen seine Fortschritte beim maschinellen Lernen vor. Die erste Variante debütierte 2018, als das BERT-Large-Modell ein Beispiel für Komplexität war und mit 256 GB Speicher und 100 GbE-Netzwerk beeindruckende Ergebnisse lieferte. Allerdings wächst die Komplexität von Machine-Learning-Modellen jedes Jahr um fast eine Größenordnung, und das Wachstum der Fähigkeiten von KI-Beschleunigern hinkt diesen Raten deutlich hinterher.
Als die P4d-Variante letztes Jahr eingeführt wurde, vervierfachte sie ihre Rechenleistung und ihren Speicher, während das berühmte GPT-3-Modell die Komplexität von BERT-Large um das 500-fache übertraf. Und jetzt sind 175 Milliarden Parameter der letzteren nichts im Vergleich zu 10 Billionen in neuen Modellen. Wir müssen auch den lokalen Speicher erhöhen (Trainium hat 512 GB HBM mit einer Gesamtbandbreite von 13,1 TB / s) und verteiltes Lernen aktiver nutzen.
Bei letzterem Ansatz wurde das Netzwerk-Subsystem zum Engpass, und bei der Entwicklung des Elastic Fabric Adapter (EFA)-Stacks hat das Unternehmen dies berücksichtigt, indem es neuen Trn1-Instanzen eine Verbindung mit 800 Gb/s (doppelt so viel wie P4d) und mit ultra-niedrige Latenz und eine optimiertere Version von Trn1n, die die doppelte Bandbreite hat und 1,6 Tbit / s erreicht. Für die Kommunikation zwischen den Chips selbst innerhalb der Instanz wird der NeuroLink-Interconnect mit einer Geschwindigkeit von 768 GB/s verwendet.
Aber es geht nicht nur darum, GPT-3 in weniger als zwei Wochen trainieren zu können, auch die Menge der eingesetzten Ressourcen ist wichtig. Im Fall von P3d würde dies 600 gleichzeitig laufende Instanzen erfordern, und selbst der Wechsel zur Ampere-Architektur würde diese Zahl auf 200 reduzieren. Aber das Training auf Basis von Trainium-Chips erfordert nur 130 Trn1-Instanzen. Dank Optimierungen betragen die Kommunikationskosten für neue Instanzen nur 7 % gegenüber 14 % für Ampere und satte 49 % für Volta.
Trainium basiert auf einem systolischen Array (Google hat den gleichen Ansatz für seine TPUs verwendet), d.h. besteht aus vielen sehr eng verwandten Recheneinheiten, die von Nachbarn empfangene Daten selbstständig verarbeiten und das Ergebnis an den nächsten Nachbarn weitergeben. Insbesondere dieser Ansatz eliminiert zahlreiche Register und Speicherzugriffe, was für "klassische" GPUs typisch ist, aber solchen Beschleunigern die Flexibilität nimmt.
In Trainium bleibt laut AWS die Flexibilität erhalten – der Beschleuniger verfügt über 16 voll programmierbare (in C/C++) Handler. Es hat auch andere Optimierungen. Zum Beispiel die Hardwarebeschleunigung der stochastischen Rundung, die bei ultragroßen Modellen durch Overhead zu "teuer" wird, obwohl sie die Effizienz des Trainings mit gemischter Präzision verbessern kann. All dies ermöglicht es Ihnen, bis zu 3,4 Pflops in Berechnungen mit niedriger Genauigkeit und bis zu 840 Teraflops in FP32-Berechnungen zu erreichen.
AWS hat versucht, den Übergang zu Trainium für Entwickler so reibungslos wie möglich zu gestalten, da das AWS Neuron SDK gängige Frameworks für maschinelles Lernen unterstützt. Das Unternehmen wird Kunden jedoch nicht zwangsweise auf Trn1-Instanzen drängen und wird weiterhin eine Auswahl anderer Beschleuniger anbieten, da der Übergang beispielsweise vom CUDA-Ökosystem schwierig sein könnte. In Sachen Machine Learning für den eigenen Bedarf ist Amazon jedoch mittlerweile komplett unabhängig – es verfügt sowohl über eine moderne CPU Graviton3 als auch über einen Inferencing-Beschleuniger Inferentia.
2021-12-06 22:44:05
Autor: Vitalii Babkin