Tensor-Parallelismus ist eine Technik zur Verteilung der Berechnung einer einzelnen Schicht eines neuronalen Netzes auf mehrere Hardware-Geräte. Bei diesem Ansatz werden die Gewichtsmatrizen einer Schicht in Shards aufgeteilt, wobei jedes Gerät einen Teil der Parameter hält und den entsprechenden Teil der Matrixmultiplikation durchführt. Die Teilergebnisse werden anschließend über kollektive Kommunikationsoperationen wie All-Reduce kombiniert, um die Ausgabe der Schicht zu erzeugen. Diese Methode ist eine Form der Modellparallelität, die sich von der Datenparallelität unterscheidet, bei der jedes Gerät eine vollständige Kopie des Modells hält und verschiedene Datenproben verarbeitet. Tensor-Parallelismus ist für das Training und die Inferenz sehr großer Modelle, wie große Sprachmodelle, die die Speicherkapazität eines einzelnen Beschleunigers überschreiten, unerlässlich.
Die Hauptmotivation für Tensor-Parallelismus sind die Speicher- und Rechenanforderungen moderner Deep-Learning-Modelle. Beispielsweise kann ein Modell mit Hunderten von Milliarden Parametern nicht in den Speicher einer einzelnen GPU oder TPU passen, der typischerweise zwischen 16 und 80 Gigabyte liegt. Durch das Sharding der Gewichtsmatrizen ermöglicht Tensor-Parallelismus die Verteilung des Modells auf viele Geräte, was die Verwendung größerer Modelle und schnellere Trainingszeiten ermöglicht. Es reduziert auch den Speicherbedarf pro Gerät und ermöglicht größere Batch-Größen. Allerdings führt Tensor-Parallelismus zu Kommunikationsaufwand, da Geräte Teilergebnisse synchronisieren müssen, was insbesondere bei langsameren Verbindungen zu einem Engpass werden kann.
Tensor-Parallelismus wird häufig in Verbindung mit anderen Parallelisierungsstrategien wie Pipeline-Parallelismus und Datenparallelismus verwendet, um das Training auf Tausende von Geräten zu skalieren. Es ist eine Schlüsselkomponente der Trainingsinfrastruktur für viele hochmoderne Modelle, einschließlich derer von OpenAI, Anthropic und Google DeepMind. Die Technik ist auch für die Inferenz relevant, da sie die Bereitstellung großer Modelle mit geringer Latenz durch die Verteilung der Berechnung auf mehrere Beschleuniger ermöglicht.
Historischer Hintergrund
Das Konzept der Aufteilung von Matrixoperationen auf mehrere Prozessoren reicht bis in die frühe Forschung zum parallelen Rechnen zurück. In den 1980er und 1990er Jahren untersuchten Forscher parallele Implementierungen neuronaler Netze, aber die Hardware- und Softwarewerkzeuge waren begrenzt. Die moderne Form des Tensor-Parallelismus entstand mit dem Aufstieg des GPU-basierten Deep Learnings in den 2010er Jahren. Eine der frühesten bemerkenswerten Anwendungen war das Training großer konvolutionaler Netze, bei dem Forscher an der University of Toronto und anderen Institutionen mit Modellparallelität experimentierten. Es war jedoch das Aufkommen transformerbasierter Modelle, wie der 2017 eingeführten Transformer-Architektur, das Tensor-Parallelismus zu einer Mainstream-Technik machte. Die Aufmerksamkeitsschichten und Feed-Forward-Netze des Transformers bestehen aus großen Matrixmultiplikationen, die natürliche Kandidaten für das Sharding sind.
Im Jahr 2019 veröffentlichten Forscher bei NVIDIA und anderen Organisationen detaillierte Beschreibungen des Tensor-Parallelismus für das Training großer Sprachmodelle. Das von NVIDIA entwickelte Megatron-LM-Framework demonstrierte effizienten Tensor-Parallelismus für Transformermodelle und erreichte nahezu lineare Skalierung auf mehreren GPUs. Diese Arbeit legte das Fundament für nachfolgende groß angelegte Trainingssysteme wie GShard von Google und DeepSpeed von Microsoft, die Tensor-Parallelismus als Kernfunktion integrieren.
Wie Tensor-Parallelismus funktioniert
In einer typischen Schicht eines neuronalen Netzes berechnet der Vorwärtsdurchlauf eine Matrixmultiplikation: Y = XW, wobei X die Eingabeaktivierungen, W die Gewichtsmatrix und Y die Ausgabe ist. Beim Tensor-Parallelismus wird die Gewichtsmatrix W entlang einer oder mehrerer Dimensionen partitioniert. Für eine vollständig verbundene Schicht ist ein üblicher Ansatz, W spaltenweise (d. h. entlang der Ausgabedimension) aufzuteilen. Jedes Gerät hält eine Teilmenge der Spalten von W und berechnet eine Teilausgabe. Die Teilausgaben werden dann mithilfe einer All-Reduce-Operation kombiniert, um das endgültige Y zu erzeugen. Dies wird als spaltenparallele Partitionierung bezeichnet. Alternativ kann W zeilenweise (entlang der Eingabedimension) aufgeteilt werden, was eine All-Gather-Operation für die Eingabe und eine All-Reduce-Operation für die Ausgabe erfordert.
Für Transformermodelle wird Tensor-Parallelismus sowohl auf die Multi-Head-Aufmerksamkeit als auch auf die Feed-Forward-Netze angewendet. Bei der Multi-Head-Aufmerksamkeit werden die Query-, Key- und Value-Gewichtsmatrizen auf Geräte verteilt, und die Aufmerksamkeitsköpfe werden verteilt. Die Ausgabeprojektion wird dann kombiniert. Im Feed-Forward-Netz werden die beiden linearen Schichten typischerweise partitioniert: Die erste Schicht wird spaltenweise aufgeteilt, und die zweite Schicht wird zeilenweise aufgeteilt, sodass die Zwischenaktivierungen verteilt und die Ausgabe reduziert werden. Dieses Design minimiert die Kommunikation, indem nur an bestimmten Punkten All-Reduce-Operationen erforderlich sind.
Das Kommunikationsmuster ist entscheidend. Jede All-Reduce-Operation erfordert den Austausch von Daten zwischen allen Geräten in der Tensor-parallelen Gruppe. Das Kommunikationsvolumen ist proportional zur Größe der Aktivierungen, nicht der Gewichte, was oft kleiner ist. Bei sehr großen Modellen kann die Kommunikation jedoch immer noch erheblich sein. Um dies zu mildern, wird Tensor-Parallelismus typischerweise innerhalb eines einzelnen Knotens mit Hochgeschwindigkeitsverbindungen wie NVIDIAs NVLink oder AMDs Infinity Fabric verwendet, während Datenparallelismus über Knoten hinweg eingesetzt wird.
Vergleich mit anderen Parallelisierungstechniken
Tensor-Parallelismus ist eine von mehreren Modellparallelitätsstrategien. Pipeline-Parallelismus beispielsweise teilt das Modell nach Schichten auf, wobei jedes Gerät für eine zusammenhängende Menge von Schichten verantwortlich ist. Dies reduziert die Kommunikation im Vergleich zum Tensor-Parallelismus, da nur die Aktivierungen an Schichtgrenzen ausgetauscht werden. Pipeline-Parallelismus kann jedoch unter Leerlaufzeiten aufgrund von Pipeline-Blasen leiden. Datenparallelismus hingegen repliziert das Modell auf jedem Gerät und verarbeitet verschiedene Datenproben, was eine Gradientensynchronisierung nach jedem Schritt erfordert. Tensor-Parallelismus ist komplementär zu beiden: Es kann mit Pipeline-Parallelismus kombiniert werden, um den Speicher pro Gerät zu reduzieren, und mit Datenparallelismus, um den Durchsatz zu erhöhen.
Eine weitere verwandte Technik ist Sequenzparallelismus, der die Sequenzdimension in Transformermodellen aufteilt. Dies wird oft in Verbindung mit Tensor-Parallelismus verwendet, um den Speicherverbrauch weiter zu reduzieren. Expertenparallelismus, der in Mixture-of-Experts-Modellen verwendet wird, verteilt die Expertenmodule auf Geräte, was eine Form der Modellparallelität, aber nicht streng genommen Tensor-Parallelismus ist.
Anwendungen und Anwendungsfälle
Tensor-Parallelismus wird häufig beim Training und bei der Inferenz großer Sprachmodelle verwendet. Beispielsweise wurde das von OpenAI entwickelte GPT-3-Modell mit einer Kombination aus Modellparallelität, einschließlich Tensor-Parallelismus, auf Tausenden von NVIDIA-GPUs trainiert. Ebenso verlassen sich die Claude-Modelle von Anthropic und die Gemini-Modelle von Google DeepMind auf Tensor-Parallelismus, um auf Hunderte von Milliarden Parametern zu skalieren. Bei der Inferenz ermöglicht Tensor-Parallelismus die Bereitstellung von Modellen wie GPT-4 mit geringer Latenz, indem die Berechnung auf mehrere GPUs verteilt wird, wie es der NVIDIA Triton Inference Server und die Text Generation Inference von Hugging Face tun.
Über Sprachmodelle hinaus wird Tensor-Parallelismus auch in anderen Bereichen wie Computer Vision und wissenschaftlichem Rechnen eingesetzt, wo große Modelle oder große Batch-Größen verteilte Berechnungen erfordern. Beispielsweise kann das Training großer Vision-Transformer (ViTs) auf hochauflösenden Bildern von Tensor-Parallelismus profitieren.
Herausforderungen und Einschränkungen
Trotz seiner Vorteile hat Tensor-Parallelismus mehrere Herausforderungen. Das Hauptproblem ist der Kommunikationsaufwand. All-Reduce-Operationen erfordern Hochgeschwindigkeitsverbindungen mit geringer Latenz. Auf Systemen mit langsameren Netzwerken wie Ethernet kann Tensor-Parallelismus ineffizient werden. Daher wird es typischerweise innerhalb eines einzelnen Servers oder eines eng gekoppelten Clusters verwendet. Eine weitere Herausforderung ist das Lastenausgleich: Ungleichmäßiges Sharding kann dazu führen, dass einige Geräte unterausgelastet sind. Darüber hinaus erfordert Tensor-Parallelismus eine sorgfältige Speicherverwaltung, da jedes Gerät seinen Shard der Gewichte und die Zwischenaktivierungen speichern muss. Schließlich ist die manuelle Implementierung von Tensor-Parallelismus komplex; es erfordert die Änderung des Modellcodes und die Handhabung von Kommunikationsprimitiven. Glücklicherweise bieten Frameworks wie PyTorch und TensorFlow integrierte Unterstützung für Tensor-Parallelismus, wodurch ein Großteil der Komplexität abstrahiert wird.
Software- und Framework-Unterstützung
Mehrere Deep-Learning-Frameworks und Bibliotheken unterstützen Tensor-Parallelismus. PyTorch bietet das Modul torch.distributed und die API tensor_parallel für das Sharding von Tensoren. TensorFlow bietet tf.distribute mit Unterstützung für Modellparallelität. Spezialisierte Bibliotheken wie Megatron-LM, DeepSpeed und die Transformers von Hugging Face haben Tensor-Parallelismus für Transformermodelle implementiert. Diese Werkzeuge ermöglichen es Forschern und Ingenieuren, Tensor-Parallelismus mit minimalen Codeänderungen anzuwenden, oft durch einfaches Angeben der Anzahl der Geräte und der Sharding-Strategie.
Zukünftige Richtungen
Da Modelle weiter wachsen, wird Tensor-Parallelismus eine entscheidende Technik bleiben. Zukünftige Entwicklungen könnten effizientere Kommunikationsalgorithmen wie hierarchisches All-Reduce und eine bessere Integration mit heterogener Hardware umfassen, einschließlich AMD-GPUs, Intel-Beschleunigern und benutzerdefinierten Chips wie AWS Trainium und den LPUs von Groq. Darüber hinaus ist die automatisierte Parallelitätserkennung, bei der das System die optimale Sharding-Strategie bestimmt, ein aktives Forschungsgebiet. Mit der zunehmenden Nachfrage nach groß angelegter KI wird sich Tensor-Parallelismus weiterentwickeln, um den Herausforderungen von Skalierung und Effizienz gerecht zu werden.