PyTorch : Parallélisme hybride (modèle et données) multi nœuds et multi GPU

La méthodologie présentée dans cette page montre comment adapter les dataloaders, les modèles et la boucle d’entraînement pour profiter de plusieurs GPU pouvant être répartis sur plusieurs nœuds.

Nous nous appuyons sur les pages PyTorch : Parallélisme de données multi-GPU et multi-nœuds et PyTorch : Parallélisme de modèle multi GPU, qu’il convient de consulter préalablement.

L’orchestration entre la distribution de modèle et le parallélisme de données ajoute quelques étapes supplémentaires aux documentations précédentes :

  • configuration de l’environnement Slurm
  • initialisation du code
  • adaptations au niveau du modèle

La création des dataloader est à reprendre sur la page PyTorch : Parallélisme de données multi-GPU et multi-nœuds. La création du modèle et l’implémentation de la boucle d’entraînement sont à copier de PyTorch : Parallélisme de modèle multi GPU.

Le code complet se trouve dans un notebook téléchargeable ou sur Jean-Zay dans le répertoire DSDIR. Pour le copier dans votre espace $WORK :

$ cp $DSDIR/examples_IA/Torch_parallel/demo_model_data_parallelism_pytorch.ipynb $WORK

Ce Notebook Jupyter se lance depuis une frontale, après avoir chargé un module PyTorch. Par exemple :

$ module load pytorch-gpu/py3/1.8.0
$ idrjup --notebook-dir=$WORK

Remarque : Ce code est fonctionnel à partir du module pytorch-gpu/py3/1.8.0, qui exploite la version 2.8.3 de NCCL

Au final nous obtiendrons une configuration similaire à la figure suivante :

 Distribution de modèle avec pipeline et parallélisation des données

Remarque : la numérotation des GPU dans l'illustration ci-dessus trouve son explication dans la section “Initialisation du code”.

Configuration de l’environnement Slurm

Pour rappel :

  • dans le cas de la parallélisation de données (variante Data Distribution Parallel), il faut réserver autant de tâches par nœud que de GPU, c’est torch.distributed.init_process_group qui se charge d’orchestrer les différentes tâches
  • dans le cas de la distribution de modèle, il faut au contraire réserver une seule tâche par modèle et autant de GPU que nécessaires pour déployer une instance du modèle (dans cette documentation, nous parallélisons un modèle sur deux GPU).

Il faut maintenant faire attention à créer un nombre de tâches correspondant au nombre d’entraînements à réaliser en parallèle et non pas au nombre de GPU.

Dans le cas d’une exécution sur trois nœuds comportant chacun 4 GPU (comme la configuration utilisée dans le notebook associé à cette page), on définit :

#SBATCH --gres=gpu:4         # reserve 4 GPU par nœud
#SBATCH --nodes=3            # reserve 3 nœuds
#SBATCH --ntasks-per-node=2  # cree 2 processus par nœud

Cela correspond à 6 processus au total tournant en parallèle et ayant chacun 2 GPU dédiés.

Et dans le cas de deux nœuds comportant chacun 8 GPU, la configuration pour 8 processus ayant chacun 2 GPU dédies sera :

#SBATCH --gres=gpu:8         # reserve 8 GPU par nœud
#SBATCH --nodes=2            # reserve 2 nœuds
#SBATCH --ntasks-per-node=4  # cree 4 processus par nœud

Initialisation du code

Il est important de bien définir les GPU sur lesquels seront envoyées les entrées du modèle ainsi que les couples de GPU sur lesquels le modèle est distribué. Pour cela, on importe le script idr_torch dans le code principal, cela nous permet d'accéder facilement aux variables définies par Slurm (script proposé par l'IDRIS, voir PyTorch : Parallélisme de données multi-GPU et multi-nœuds).

Récupération des variables nécessaires

Grâce à idr_torch, nous avons accès à :

  • hostnames : la liste des noms des nœuds impliqués dans le job.
  • size : équivalent de SLURM_NTASKS (nodes x ntasks_per_node) qui sert lors de la configuration des dataloaders (num_replica dans le jargon pytorch)
  • rank : équivalent de SLURM_PROCID qui permet d’identifier un processus
  • local_rank : équivalent de SLURM_LOCALID qui permet d’identifier un processus sur un nœud, et dans notre cas à affecter les GPU aux bons processus
  • cpus_per_task : équivalent à SLURM_CPUS_PER_TASK qui permet de fixer un nombre adéquat de threads pour lire les données (num_workers dans la déclaration du data_loaders). Attention, la valeur optimale pour num_workers est rarement cpus_per_task qui constitue uniquement une indication de la valeur maximale raisonnable

Autres variables indispensables :

  • NTASKS_PER_NODE, variable d’environnement Slurm, accessible uniquement si l’on utilise l’instruction #SBATCH --ntasks-per-node=x ou calculable avec les variables fournis par idr_torch.
  • torch.cuda.device_count(), nombre de cartes GPU visibles du processus (soit toutes les cartes du nœud !).

Ces deux variables servent à créer des listes de GPU pour la distribution de modèle.

Définition des couples de GPU par processus

L’objectif est d’affecter un couple de GPU (dans le cas présent) à chaque processus, et à les donner en paramètres lors de la création du modèle.

Si l’on utilise le backend NCCL (conseillé sur Jean Zay car plus performant), il ne faut pas que les GPU pour un modèle aient des identifiants consécutifs. Par exemple, si l’on définit [0,1] et [2,3] pour une exécution mono nœud et 4 GPU, il y aura un deadlock lors de la communication avec le processus maître (voir Model parallel with DDP get 'Socket Timeout' error when using NCCL, while GLOO works fine #25767 pour une description du bug).

Cette remarque est valable pour le module pytorch-gpu-1.8.0 (avril 2021). Ce qui mène au code suivant, qui permet de respecter ces consignes :

NTASKS_PER_NODE = os.environ['SLURM_NTASKS_PER_NODE']
 
dev0 = idr_torch.local_rank % torch.cuda.device_count()
dev1 = dev0 + NTASKS_PER_NODE  
torch.cuda.set_device(dev0)

Mise en place du mode de communication

backend= 'nccl'   # alternate option is 'gloo'
dist.init_process_group(backend=backend, init_method='env://', 
                        world_size=idr_torch.size, rank=idr_torch.rank)

Adaptations au niveau du modèle

Encapsulage du modèle

Après la création du modèle distribué, avec les variables dev0 et dev1 définies correctement, on rajoute l’instruction suivante, pour le rendre compatible avec le parallélisme de données :

ddp_mp_model = DistributedDataParallel(mp_model)

Cela permet à PyTorch d’envoyer les fractions de batches aux différentes instances du modèle et d’assurer la synchronisation lors de sa mise à jour.

Si l’on regarde le profil de la fonction DistributedDataParallel(...), on y trouve les options suivantes :

  • device_ids= [...]
  • output_device= i

Dans le cas d’un modèle en parallélisme hybride, il est primordial de ne pas fixer ces variables.

Sauvegarde et chargement du modèle

La sauvegarde des paramètres du modèle se fait uniquement sur le processus maître :

if idr_torch.rank == 0
    torch.save(model.module.state_dict(), name)

Le fait d’ajouter le parallélisme de données sur un modèle réparti sur deux GPU (ou plus), ajoute une “indirection” au dictionnaire de paramètres (model.module au lieu de model), que l’on peut éviter dans le fichier de sauvegarde en récupérant directement le dictionnaire au niveau de model.module plutôt que model (comportement par défaut).

Le chargement doit être réalisé par tous les processus. A noter que nous n’avons testé que des configurations “statiques” : chargement d’un modèle bi-gpu sur une configuration aussi bi-gpu. Pour plus d’informations, se référer au tutoriel Saving and Loading Models.

mp_model = PipelinedResnet(dev0, dev1)
if load_model:
    mp_model.load_state_dict(torch.load(load_model))
ddp_mp_model = DistributedDataParallel(mp_model)

Compléments d’information

Quelques remarques pour finir :

  • l’optimisation d’un code distribué, modèle et data, est complexe.
  • d’autres techniques permettent de limiter la mémoire GPU nécessaire pour un modèle, elles peuvent être plus simples à mettre en œuvre et tout aussi efficaces. Par exemple, la Mixed-Precision (ou AMP), le Gradient Checkpointing, l'optimisation ZeRO.

⇐ Revenir à la page mère sur l'apprentissage distribué