Megatron-LM : Formation de modèles linguistiques à plusieurs milliards de paramètres en utilisant le parallélisme de modèle

Des travaux récents en modélisation linguistique montrent que l'entraînement de grands modèles de transformateurs fait progresser l'état de l'art dans les applications de traitement du langage naturel (Natural Language Processing, NLP). Cependant, les modèles très volumineux peuvent être difficiles à entraîner en raison des contraintes mémoire. Dans ce travail, nous présentons nos techniques pour entraîner des modèles de transformateurs très volumineux et mettons en œuvre une approche simple et efficace de parallélisme inter-couche qui permet d'entraîner des modèles de transformateurs avec des milliards de paramètres. Notre approche ne nécessite pas un nouveau compilateur ou des modifications de bibliothèque, est orthogonale et complémentaire au parallélisme inter-couche, et peut être entièrement mise en œuvre par l'insertion de quelques opérations de communication dans le PyTorch natif. Nous illustrons cette approche en faisant converger des modèles basés sur des transformateurs jusqu'à 8,3 milliards de paramètres en utilisant 512 GPU. Nous maintenons 15,1 PetaFLOPs sur toute l'application avec une efficacité d'évolutivité de 76 % par rapport à une base solide sur un seul GPU qui maintient 39 TeraFLOPs, soit 30 % des FLOPs maximums. Pour démontrer que les grands modèles linguistiques peuvent encore faire progresser l'état de l'art (State of the Art, SOTA), nous entraînons un modèle linguistique de transformateur à 8,3 milliards de paramètres similaire à GPT-2 et un modèle à 3,9 milliards de paramètres similaire à BERT. Nous montrons que la prise en compte attentive du placement de la normalisation inter-couche dans les modèles similaires à BERT est cruciale pour atteindre une meilleure performance lorsque la taille du modèle augmente. En utilisant le modèle GPT-2, nous obtenons des résultats SOTA sur les jeux de données WikiText103 (10,8 contre une perplexité SOTA de 15,8) et LAMBADA (66,5 % contre une précision SOTA de 63,2 %). Notre modèle BERT obtient des résultats SOTA sur le jeu de données RACE (90,9 % contre une précision SOTA de 89,4 %).