AI21 Labs trainiert ein Sprachmodell mit 178 Milliarden Parametern mithilfe von Amazon-EC2-P4d-Instances und PyTorch

AI21 Labs nutzt Machine Learning, um Sprachmodelle zu entwickeln, die Bedeutungen verstehen. 2021 hat es sich zum Ziel gesetzt, das kürzlich veröffentlichte Jurassic-1 Jumbo zu trainieren, ein autoregressives Sprachmodell mit 178 Milliarden Parametern. Entwickler, die sich für den Betatest registrieren, erhalten Zugriff auf Jurassic-1 Jumbo und können sofort damit beginnen, das Modell an ihren Anwendungsfall anzupassen. Das Software-Startup wollte das Modell effizient trainieren. Deshalb wandte es sich an Amazon Web Services (AWS) und entwickelte eine Lösung mithilfe von Amazon Elastic Compute Cloud (Amazon EC2), einem Webservice, der sichere, skalierbare Rechenkapazität in der Cloud bereitstellt. Die Entscheidung für Amazon EC2 gab dem Unternehmen die Kontrolle über den Trainingsprozess, einschließlich der Knotenzuweisung.

Für leistungsstarke Rechen- und Netzwerkfunktionen entschied sich das Unternehmen für Amazon-EC2-P4d-Instances, die einen hohen Durchsatz und ein Netzwerk mit geringer Latenz für Machine-Learning-Training und Hochleistungs-Computing-Anwendungen in der Cloud bieten. Mithilfe von Amazon-EC2-P4d-Instances konnte AI21 Labs die erforderliche Leistung und den erforderlichen Speicherplatz durch die Verteilung des Modelltrainings auf Hunderte von GPUs erreichen, um über sein Jurassic-1-Jumbo-Modell die Verarbeitung natürlicher Sprache als Service bereitzustellen. Da das Unternehmen nun sein eigenes Großmodell trainiert und steuert, kann es auf die Entwicklung neuer Modelle im gleichen Maßstab hinarbeiten und Innovationen leichter umsetzen.

AI21 Labs wurde 2017 gegründet und verfolgt eine hybride Mission: Forschung zur Verarbeitung natürlicher Sprache und Entwicklung von auf künstlicher Intelligenz basierenden Produkten zum Lesen und Schreiben. Sein Hauptprodukt, Wordtune, ist ein intelligenter Schreib- und Bearbeitungsassistent, der im Oktober 2020 auf den Markt kam und inzwischen fast eine Million Benutzer unterstützt. Sein anderes Hauptprodukt, AI21 Studio, bietet API-Zugriff auf die Jurassic-1-Sprachmodelle des Unternehmens sowie die Entwicklung benutzerdefinierter Modelle. „Wir sind Teil einer kleinen Gruppe von Unternehmen, die Sprachmodelle als Service anbieten und jedem, vom unabhängigen Entwickler bis zum multinationalen Unternehmen, die Möglichkeit geben, Apps und Services auf der Grundlage fortschrittlicher Technologie zur Verarbeitung natürlicher Sprache zu entwickeln“, erklärt Yoav Shoham, Mitbegründer und Co-CEO bei AI21 Labs. „Darüber hinaus verfolgen wir wissenschaftliche Innovationen und bewältigen die Herausforderungen der Softwareentwicklung, die Modelle dieser Größe und Komplexität mit sich bringen.“

Um sein erstes Deep-Learning-Großmodell effizient zu trainieren und die hohen Skalierungs- und Leistungsanforderungen des Modells zu unterstützen, benötigte AI21 Labs leistungsstarke Rechenleistung, effiziente Netzwerkgeschwindigkeit und Zugang zu technischem Support und Beratung. Aus diesen Gründen begann das Unternehmen Anfang 2021 mit der Implementierung einer Lösung auf AWS und entschied sich dafür, das Modell mithilfe von Amazon-EC2-P4d-Instances zu trainieren. Diese Instances werden in Hyperscale-Clustern namens Amazon EC2 UltraClusters bereitgestellt und bieten mehr als 4 000 NVIDIA A100 Grafikprozessoren, eine blockierungsfreie Netzwerkinfrastruktur im Petabit-Maßstab sowie Speicher mit hohem Durchsatz und niedriger Latenz. 

Der Ansatz des Unternehmens wurde durch GPUDirectRDMA mit niedriger Latenz und hoher Bandbreite sowie Elastic Fabric Adapter (EFA), einer Netzwerkschnittstelle für Amazon-EC2-Instances, weiter optimiert. Damit können Kunden Anwendungen, die ein hohes Maß an Kommunikation zwischen Knoten erfordern, in großem Maßstab auf AWS ausführen. Aufgrund der Größe des Modells musste das Team parallele Verarbeitung verwenden, um eine effiziente Trainingszeit zu erreichen. Deshalb nutzte es die Netzwerkfunktionen auf AWS, um das verteilte Training und die Modellparallelität zu unterstützen. „Amazon-EC2-P4d-Instances bieten Hochleistungsnetzwerke mit 400 Gbit/s auf EFA“, sagt Opher Lieber, technischer Leiter von Jurassic bei AI21 Labs. „Die GPU-zu-GPU-Netzwerkgeschwindigkeit wirkt sich direkt auf die Fähigkeit aus, effizient skalieren zu können, und bleibt selbst bei der Skalierung auf Hunderte von GPUs kosteneffektiv.“

AI21 Labs begann damit, seine Codebasis auf für EFA aktivierten Amazon-EC2-P4d-Instances zu erstellen. Anschließend testete und verifizierte es die Leistung und effiziente Skalierung seines Multiknoten-Trainingsansatzes. Als Nächstes führte das Team ein schnelles Training des Modells in voller Größe durch, das Hunderte von GPUs verwendet, um Funktion und Leistung zu überprüfen. Von dort aus konnte das Unternehmen mit dem Training seines Jurassic-1-Jumbo-Modells auf AWS beginnen. Für die Orchestrierung entschied sich das Unternehmen für eine interne Lösung, die Instances mithilfe eines AWS Softwareentwicklungskits – AWS SDK for Python (Boto3) – zuweist. Diese Lösung vereinfacht die Integration von Python-Anwendungen, -Bibliotheken oder -Skripten des Kunden in verschiedene AWS-Services.

Für die Speicherung entschied sich AI21 Labs für Amazon Simple Storage Service (Amazon S3), der branchenführende Skalierbarkeit, Datenverfügbarkeit, Sicherheit und Leistung bietet. „Wir konnten mithilfe des AWS-Teams eine sehr gute Leistung auf Amazon S3 erzielen – daher fiel uns die Wahl sowohl hinsichtlich der Leistung als auch des Preises leicht“, sagt Lieber. Das Team verwendet Amazon-S3-Buckets, um Prüfpunkte effizient und verteilt zu speichern und zu laden. Zur Protokollierung von Trainingsfortschritten und Ereignissen verwendet das Team Amazon CloudWatch, einen Service zur Überprüfung und Überwachung. 

AI21 Labs nutzte bei der Implementierung seiner Lösung die Unterstützung von AWS. Das Team konsultierte AWS-Spezialisten, die bei Fragen und Anliegen in Bezug auf Service-Level, Architektur und Hardware beratend zur Seite standen. Darüber hinaus verbesserte das Unternehmen die Leistung von Jurassic-1 Jumbo mithilfe von PyTorch in AWS, einem Open-Source-Framework für Deep Learning, das die Entwicklung von Machine-Learning-Modellen und deren Bereitstellung in der Produktion vereinfacht. 

AI21 Labs schloss das mehrmonatige Training im Juni 2021 ab. Das neue Großmodell, ein autoregressives Sprachmodell, verfügt über 178 Milliarden Parameter und ist damit vergleichbar mit dem Angebot des Konkurrenzunternehmens. Es bietet auch ein differenziertes Vokabular mit 256 000 Einträgen, das erweiterte Funktionen zur Textdarstellung sowie Unterstützung für benannte Entitäten bietet. Das Unternehmen bietet jetzt Jurassic-1 Jumbo (zusammen mit seinem Gegenstück Jurassic-1 Large mit 7 Milliarden Parametern) in der offenen Beta-Version über das AI21-Studio-Angebot des Unternehmens an. Mithilfe des Service kann eine Vielzahl von Entwicklern Produkte nach dem Jurassic-1-Jumbo-Modell entwickeln. AI21 Labs hat bereits in vielen Branchen Anklang gefunden, darunter Marketing, Content-Erstellung, Gaming, medizinische Forschung, Automobilbau, Telekommunikation und Finanzen.

Da AI21 Labs sein Modell besitzt und direkten Zugriff darauf hat, kann es Anpassungen und Innovationen vornehmen, ohne auf Dritte angewiesen zu sein. Außerdem kann es fortlaufende Innovationsziele verfolgen, die ein wesentlicher Bestandteil seiner Mission sind. AI21 Labs entwickelt derzeit Prototypen für weitere Modelle, die ebenfalls im großen Maßstab trainiert werden sollen. „Das Training und der Besitz unserer eigenen Großmodelle werden weiterhin ein entscheidender Differenzierungsfaktor sowohl für unsere Wordtune- als auch für unsere AI21-Studio-Angebote sein“, sagt Shoham.