100 Billionen Parameter KI-Trainingsmodelle

Recommender-KI-Systeme sind heute ein wichtiger Bestandteil von Internetdiensten: Unternehmen mit Milliardenumsätzen wie Amazon und Netflix werden direkt von Empfehlungsdiensten angetrieben.

KI-Recommender werden besser, wenn sie größer werden. In letzter Zeit wurden mehrere Modelle mit Milliarden Parametern bis hin zu Billionen herausgebracht. Jeder Sprung in der Modellkapazität hat zu einer deutlichen Qualitätsverbesserung geführt. Die Ära der 100 Billionen Parameter steht vor der Tür.

Das komplizierte, dichte neuronale Netzwerk mit dichtem Rest ist mit mehr als 100 TFLOPs in jeder Trainingsiteration zunehmend rechenintensiv. Daher ist es wichtig, einen ausgeklügelten Mechanismus zu haben, um einen Cluster mit heterogenen Ressourcen für solche Trainingsaufgaben zu verwalten.

Kürzlich haben das Kwai Seattle AI Lab und das DS3 Lab der ETH Zürich zusammengearbeitet, um ein neuartiges System namens „Persia“ vorzuschlagen, um dieses Problem durch sorgfältiges Co-Design sowohl des Trainingsalgorithmus als auch des Trainingssystems anzugehen. Auf der Algorithmusebene verwendet Persia einen hybriden Trainingsalgorithmus, um die Einbettungsschicht und dichte neuronale Netzwerkmodule unterschiedlich zu handhaben. Die Einbettungsschicht wird asynchron trainiert, um den Durchsatz von Trainingsbeispielen zu verbessern, während das restliche neuronale Netzwerk synchron trainiert wird, um die statistische Effizienz zu bewahren. Auf Systemebene wurde eine breite Palette von Systemoptimierungen für die Speicherverwaltung und die Reduzierung der Kommunikation implementiert, um das volle Potenzial des Hybridalgorithmus auszuschöpfen.

Cloud-Ressourcen für 100 Billionen Parameter-KI-Modelle

Persiens 100-Billionen-Parameter-KI-Workload wird auf den folgenden heterogenen Ressourcen ausgeführt:

3,000 Kerne rechenintensiver virtueller Maschinen
8 virtuelle A2-Maschinen, die insgesamt 64 A100-Nvidia-GPUs hinzufügen
30 virtuelle Maschinen mit hohem Arbeitsspeicher, jede mit 12 TB RAM, insgesamt 360 TB
Orchestrierung mit Kubernetes
Alle Ressourcen mussten gleichzeitig in derselben Zone gestartet werden, um die Netzwerklatenz zu minimieren. Google Cloud konnte die erforderliche Kapazität innerhalb kürzester Zeit bereitstellen.

KI-Training benötigt Ressourcen in Schüben.

Google Kubernetes Engine (GKE) wurde verwendet, um die Bereitstellung der 138 VMs und Softwarecontainer zu orchestrieren. Die Containerisierung der Workload ermöglicht auch die Portierung und Wiederholbarkeit des Trainings.

Ergebnisse und Schlussfolgerungen
Mit Unterstützung der Google Cloud-Infrastruktur demonstrierte das Team die Skalierbarkeit von Persia auf bis zu 100 Billionen Parameter. Der hybride verteilte Trainingsalgorithmus führte ausgeklügelte Systemlockerungen für eine effiziente Nutzung heterogener Cluster ein, während er so schnell wie Vanilla SGD konvergierte. Google Cloud war unerlässlich, um die Einschränkungen der lokalen Hardware zu überwinden, und erwies sich als optimale Computerumgebung für verteiltes Machine Learning-Training in großem Umfang.

Persia wurde als Open-Source-Projekt auf GitHub mit Setup-Anweisungen für Google Cloud veröffentlicht – jeder aus Wissenschaft und Industrie würde es leicht finden, Deep-Learning-Recommender-Modelle mit einer Skala von 100 Billionen Parametern zu trainieren.