Mesh-TensorFlow : Apprentissage profond pour supercalculateurs

Le fractionnement de lots (data-parallelism) est la stratégie dominante pour l'entraînement distribué des Réseaux Neuronaux Profonds (DNN), en raison de sa polyvalence et de sa facilité d'utilisation dans le cadre de la programmation à un seul programme, multiples données (SPMD). Cependant, cette méthode présente plusieurs problèmes, notamment l'impossibilité d'entraîner des modèles très volumineux (en raison des contraintes mémoire), une latence élevée et une inefficacité pour les petits lots. Tous ces problèmes peuvent être résolus par des stratégies de distribution plus générales (model-parallelism). Malheureusement, les algorithmes efficaces de parallélisme du modèle sont souvent complexes à découvrir, à décrire et à mettre en œuvre, particulièrement sur de grands clusters. Nous présentons Mesh-TensorFlow, un langage permettant de spécifier une classe générale de calculs tensoriels distribués. Alors que le parallélisme des données peut être considéré comme le fractionnement des tenseurs et des opérations selon la dimension « batch », dans Mesh-TensorFlow, l'utilisateur peut spécifier n'importe quelle dimension tensorielle à fractionner sur n'importe quelle dimension d'un maillage multidimensionnel de processeurs. Un graphe Mesh-TensorFlow se compile en un programme SPMD composé d'opérations parallèles couplées avec des primitives de communication collective telles que Allreduce. Nous utilisons Mesh-TensorFlow pour implémenter une version efficace du modèle Transformer séquence-à-séquence qui combine le parallélisme des données et celui du modèle. En utilisant des maillages TPU allant jusqu'à 512 cœurs, nous entraînons des modèles Transformer comportant jusqu'à 5 milliards de paramètres, surpassant ainsi les résultats actuels sur la tâche de traduction anglais-français du WMT'14 et sur le benchmark de modélisation linguistique d'un milliard de mots. Mesh-TensorFlow est disponible à l'adresse https://github.com/tensorflow/mesh .