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 (32 Go ou 16 Go selon la partition utilisée).

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

  • Le paralé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 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.

Notes :

  • 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.
  • La taille du modèle peut également être optimisée en réduisant la précision des poids du modèle. L'état de l'art consiste à entraîner des modèles en float 32, mais ils peuvent aussi être entraînés en float 16 ou en int 8 au prix d'une baisse de la précision souvent non critique.

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

Nous avons également rédigé une documentation de mise en pratique pour pouvoir faire du parallélisme de données sur Jean Zay avec les stratégies de distribution intégrées à TensorFlow. Cependant, nous déconseillons son utilisation au vu des mauvaises performances constatées en multi-GPU. Nous conseillons la solution de distribution Horovod.

Parallélisme de modèle

Le parallélisme de modèle consiste à segmenter un modèle sur plusieurs GPU afin :

  • de limiter le volume de calcul à effectuer par un GPU lors de l'apprentissage
  • et 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 et chaque GPU prend en charge uniquement un de ces groupes, ce qui limite le nombre de paramètres à traiter localement.

Notez que, dans ce cas, il n'est pas nécessaire de synchroniser les gradients : les GPU collaborent de manière séquentielle pour réaliser l'apprentissage sur chaque donnée.

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 ait é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.

Cependant, il est possible d'optimiser le procédé pour gagner du temps en utilisant une technique de data pipeline : la segmentation de batch. Cela consiste à diviser le batch en sous-batches afin de limiter les temps d'attente et de faire travailler les GPU simultanément lors de l'étape de Forward.

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

Pipeline et parallélisme de modèle

Illustration de la technique de segmentation de batch avec 3 sous-batches.

Le nombre optimal de sous-batches dépend du modèle considéré. Le graphe ci-dessous illustre le temps d'exécution d'un apprentissage ResNet-50 (réseau de neurones profond convolutif composé de 50 couches) sur un batch de données de 120 images. L'entraînement est distribué sur 2 GPU selon la méthode de parallélisme de modèle et on utilise la technique de segmentation de batch en faisant varier le nombre de sous-batches.

Optimisation du pipeline

Evolution du temps d'entraînement de ResNet-50 en fonction du nombre de sous-batches

Vous noterez que pour cet exemple, le nombre optimal de sous-batches est environ de 12.

Notes :

  • La segmentation de batch permet aussi de limiter la mémoire nécessaire par GPU. Si le problème n'est pas tant la taille du modèle mais plutôt la taille des entrées (images médicales de très haute résolution, images 3D, vidéos,…), la technique du data pipeline même sur un seul GPU peut alors résoudre le problème de mémoire.
  • Il est possible aussi de combler les temps d'attente des GPU avec des techniques de data pipeline propres au chargement des données (prefetch, preprocessing via les GPU, Data Augmentation, …).

Pour aller plus loin

Documentation et sources