Table des matières

Jean Zay : Distribution multi-GPU et multi-nœuds pour l'apprentissage d'un modèle TensorFlow ou PyTorch

L'apprentissage d'un modèle de réseau de neurones est confronté à deux problématiques majeures :

  • Le temps d'apprentissage peut être très long (plusieurs semaines).
  • Le modèle peut être trop volumineux pour la mémoire du GPU (80 GB, 40 GB, 32 GB ou 16 GB selon la partition utilisée).

La distribution permet de répondre à ces deux problématiques. Nous distinguons deux techniques :

  • Le parallélisme de données (Data Parallelism) qui permet d'accélérer l'apprentissage du modèle en distribuant chaque batch de données et donc chaque itération d'upgrade du modèle. Dans ce cas, le modèle est répliqué sur chaque GPU utilisé pour traiter un sous-ensemble des données. Il est adapté aux apprentissages impliquant des données volumineuses ou des batches de grande taille. Il s'agit du procédé de distribution le plus documenté et le plus facile à mettre en œuvre.
  • Le parallélisme de modèle (Model Parallelism) qui est adapté aux modèles trop volumineux pour être répliqués et manipulés avec la taille de batch visé dans chaque GPU. Il est plus difficile à mettre en œuvre, particulièrement en multi-nœuds. Il peut être optimisé avec des techniques de pipeline, mais il ne permet pas d'accélérer l'apprentissage aussi significativement que le parallélisme de données.

Les deux techniques peuvent être utilisées en même temps (parallélisme hybride) pour résoudre les deux problèmes à la fois en groupant les GPU en fonction de la tâche à effectuer.

Fonctionnalités intégrées à TensorFlow et PyTorch

Parallélisme de données

Le parallélisme de données consiste à répliquer le modèle sur l'ensemble des GPU et à diviser le batch de données en mini-batches. Chaque GPU prend alors en charge un mini-batch pour un apprentissage en parallèle sur l'ensemble des mini-batches.

Sous-batch en multi GPU

Schéma de distribution d'un batch de données sur 256 GPU

Dans ce cas, le processus d'entraînement est le suivant :

  1. Le script d'apprentissage est exécuté en mode distribué sur plusieurs workers (i.e. plusieurs GPU ou plusieurs groupes de GPU). Ainsi, chaque GPU :
    1. lit une partie des données (mini-batch) ;
    2. puis entraîne le modèle à partir de son mini-batch ;
    3. et calcule la mise à jour des paramètres du modèle local (gradients locaux).
  2. Les gradients globaux sont calculés en moyennant les gradients locaux retournés par l'ensemble des GPU. Il est nécessaire à cette étape de synchroniser les différents workers.
  3. Le modèle est mis à jour sur l'ensemble des GPU.
  4. L'algorithme est répété à partir de l'étape 1.a. jusqu'à la fin de l'apprentissage.

Data Parallelism

Représentation d'un apprentissage distribué sur 3 GPU selon la méthode de parallélisme de données

Mise en pratique

Parallélisme de modèle

Mise en garde: Il n'est pas évident à mettre en place car il oblige à minima à définir une hyperclasse pour le modèle ou à redéfinir complètement le modèle de façon séquentielle. Il sera préférable d'utiliser d'autres optimisations pour traiter les problèmes d'occupation de la mémoire (Gradient Checkpointing, ZeRO de Deepspeed, …). Seuls les très gros modèles (par exemples, pour les Transformers les plus performants) utiliseront le parallélisme de modèle par absolue nécessité.

Le parallélisme de modèle consiste à segmenter un modèle sur plusieurs GPU afin de réduire l'empreinte mémoire du modèle sur chaque GPU (nécessaire pour des modèles trop gros ne pouvant pas tenir dans la mémoire d'un seul GPU). En théorie, l'occupation mémoire par GPU est ainsi divisée par le nombre de GPU.

Le modèle est divisé en groupe de couches successives (nommées séquences) et chaque GPU prend en charge uniquement une de ces séquences, ce qui limite le nombre de paramètres à traiter localement.

Contrairement au parallélisme de données, il n'est pas nécessaire de synchroniser les gradients : les GPU collaborent de manière séquentielle lors de chaque itération d'apprentissage.

La figure suivante illustre une parallélisation du modèle sur 2 GPU, divisant ainsi par 2 la taille mémoire occupée dans un GPU.

Model parallelism

Illustration d'un entraînement distribué sur 2 GPU selon la méthode de parallélisme de modèle

Notez que le parallélisme de modèle n'accélère pas l'apprentissage car un GPU doit attendre que les couches antérieures du modèle aient été traitées par un autre GPU avant de pouvoir exécuter ses propres couches. L'apprentissage est même un peu plus long à cause des temps de transfert des données entre GPU.

Néanmoins, il est possible et conseillé d'optimiser le procédé pour gagner du temps en utilisant une technique de data pipeline : la segmentation de batch (on parle alors de Pipeline Parallelism). Cela consiste à diviser le batch en sous-batches afin de limiter les temps d'attente et de faire travailler les GPU quasi-simultanément. Cependant il y aura toujours un petit temps d'attente (la bulle sur le schéma).

Le schéma ci-dessous illustre la quasi-simultanéité et le gain obtenu grâce à la technique de segmentation lorsque le batch de données est divisé en 4 sous-batches, pour un entraînement distribué sur 4 GPU.

Pipeline et parallélisme de modèle

Comparaison entre parallèlisme de modèle basique et en pipeline avec 4 sous-batches. (Source:GPipe)

Le nombre optimal de sous-batches dépend du modèle considéré.

Mise en pratique

Pour aller plus loin

  • Des scripts exemples, adaptés à Jean-Zay, pour l'utilisation des librairies PEFT, Deepspeed et accelerate sont disponibles sur le dépôt plm4all. Ces scripts permettent une utilisation pour de l'inférence ou du finetunning de modèles.Ils ont pour vocation à servir de base pour l'implémentation des différentes librairies.

Documentation et sources