Mesh-TensorFlow: Tiefes Lernen für Supercomputer

Batch-Splitting (Datenparallellität) ist die vorherrschende verteilte Trainierungsstrategie für Tiefneuronale Netze (DNNs), aufgrund ihrer universellen Anwendbarkeit und ihrer Eignung für die Programmierung nach dem Single-Program-Multiple-Data (SPMD)-Prinzip. Dennoch leidet Batch-Splitting unter Problemen wie der Unfähigkeit, sehr große Modelle aufgrund von Speicherbeschränkungen zu trainieren, hoher Latenz und Ineffizienz bei kleinen Batch-Größen. All diese Probleme können durch allgemeinere Verteilungsstrategien (Modellparallellität) gelöst werden. Leider neigen effiziente modellparallele Algorithmen dazu, kompliziert zu sein, sowohl in ihrer Entdeckung als auch in ihrer Beschreibung und Implementierung, insbesondere in großen Clustern. Wir stellen Mesh-TensorFlow vor, eine Sprache zur Spezifizierung einer allgemeinen Klasse verteilter Tensorberechnungen. Während Datenparallellität als Aufteilung von Tensoren und Operationen entlang der "Batch"-Dimension betrachtet werden kann, kann der Benutzer in Mesh-TensorFlow beliebige Tensor-Dimensionen über beliebige Dimensionen eines mehrdimensionalen Prozessorsystems aufteilen. Ein Mesh-TensorFlow-Graph wird in ein SPMD-Programm kompiliert, das aus parallelen Operationen besteht, die mit kollektiven Kommunikationsprimitiven wie Allreduce gekoppelt sind. Wir verwenden Mesh-TensorFlow, um eine effiziente datenparallele und modellparallele Version des Transformer sequenz-zu-sequenz-Modells zu implementieren. Mit TPU-Gittern von bis zu 512 Kernen trainieren wir Transformer-Modelle mit bis zu fünf Milliarden Parametern und übertreffen damit den Stand der Technik bei der Übersetzungs-Aufgabe vom WMT'14 Englisch-Französisch-Korpus sowie beim Benchmark für Sprachmodelle mit einer Milliarde Wörtern. Mesh-TensorFlow ist unter https://github.com/tensorflow/mesh verfügbar.