PyTorch : Parallélisme de modèle multi GPU

La méthodologie présentée dans cette page montre comment adapter concrètement un modèle trop volumineux pour tenir sur un seul GPU. C'est une illustration des concepts présentés dans la page mère : Jean Zay : Distribution multi-GPU et multi-nœuds pour l'apprentissage d'un modèle TensorFlow ou PyTorch.

Le procédé suit plusieurs étapes :

  • adaptation du modèle
  • adaptation de la boucle d'entraînement
  • configuration de l'environnement de calcul Slurm

Les étapes concernant la création de dataloader, la sauvegarde et le chargement de modèles, ne sont pas modifiées par rapport au cas d'un modèle déployé sur un seul GPU.

Ce document présente uniquement les changements à faire pour paralléliser un modèle, une démonstration complète se trouve dans un notebook téléchargeable ici ou récupérable sur Jean-Zay dans le répertoire DSDIR : /examples_IA/Torch_parallel/demo_model_parallelism_pytorch.ipynb. Pour le copier dans votre espace $WORK :

$ cp $DSDIR/examples_IA/Torch_parallel/demo_model_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

La principale source d'information provient de la documentation officielle de PyTorch : Single-Machine Model Parallel Best Practices.

Adaptation du modèle

Pour illustrer la méthodologie, un modèle resnet est distribué sur deux GPU (nous reprenons l’exemple proposé dans la documentation PyTorch). Dans le cas d’autres applications ayant des architectures de modèle différentes des CNN (comme les transformers ou les architectures multi-modèles), la stratégie de découpage du modèle en séquences pourra être différente.

Deux méthodes sont proposées :

  • distribution simple de modèle - dans ce cas, un seul GPU fonctionne à la fois. Cela occasionne une durée plus longue des cycles d’entraînement, due aux communications inter GPU (lors des phases de forward et backward propagation).

 Distribution simple de modèle

  • distribution de modèle, dite pipeline, avec une segmentation sur les batches d'échantillons, ce qui permet aux GPUs de travailler de manière concurrente. Le temps d’exécution est alors réduit par rapport à l’équivalent en mono GPU.

 Distribution de modèle avec pipeline

Ajustements au niveau du modèle

L’idée est de répartir les couches du modèle entre deux GPU (ici, nommés dev0 et dev1).

La fonction forward permet de faire le lien entre les deux GPU.

from torchvision.models.resnet import ResNet, Bottleneck
 
num_classes = 1000
 
class ParallelResnet(ResNet):
    def __init__(self, dev0, dev1, *args, **kwargs):
        super(ParallelResnet, self).__init__(
            Bottleneck, [3, 4, 23, 3], num_classes=num_classes, *args, **kwargs)
        # dev0 and dev1 each point to a GPU device (usually gpu:0 and gpu:1)
        self.dev0 = dev0
        self.dev1 = dev1
 
        # splits the model in two consecutive sequences : seq0 and seq1 
        self.seq0 = nn.Sequential(
            self.conv1,
            self.bn1,
            self.relu,
            self.maxpool,
            self.layer1,
            self.layer2
        ).to(self.dev0)  # sends the first sequence of the model to the first GPU
 
        self.seq1 = nn.Sequential(
            self.layer3,
            self.layer4,
            self.avgpool,
        ).to(self.dev1)  # sends the second sequence of the model to the second GPU
 
        self.fc.to(self.dev1)  # last layer is on the second GPU
 
    def forward(self, x):
        x= self.seq0(x)     # apply first sequence of the model on input x
        x= x.to(self.dev1)  # send the intermediary result to the second GPU
        x = self.seq1(x)    # apply second sequence of the model to x
        return self.fc(x.view(x.size(0), -1))  

La version pipeline part du modèle précédent (version “parallèle”) et va diviser le batch d’échantillons pour que les deux GPU puissent fonctionner de manière quasi concurrente. Le nombre de divisions optimum dépend du contexte (modèle, taille du batch) et doit être estimé au cas par cas.

Le choix de la meilleure division pour resnet50, via un benchmark associé, se trouve sur Single-Machine Model Parallel Best Practices / speed-up-by-pipelining-inputs.

Version implémentant la notion de pipeline :

class PipelinedResnet(ResNet):
    def __init__(self, dev0, dev1, split_size=8, *args, **kwargs):
        super(PipelinedResnet, self).__init__(
            Bottleneck, [3, 4, 23, 3], num_classes=num_classes, *args, **kwargs)
        # dev0 and dev1 each point to a GPU device (usually gpu:0 and gpu:1)
        self.dev0 = dev0
        self.dev1 = dev1
        self.split_size = split_size
 
        # splits the model in two consecutive sequences : seq0 and seq1 
        self.seq0 = nn.Sequential(
            self.conv1,
            self.bn1,
            self.relu,
            self.maxpool,
            self.layer1,
            self.layer2
        ).to(self.dev0)  # sends the first sequence of the model to the first GPU
 
        self.seq1 = nn.Sequential(
            self.layer3,
            self.layer4,
            self.avgpool,
        ).to(self.dev1)  # sends the second sequence of the model to the second GPU
 
        self.fc.to(self.dev1)  # last layer is on the second GPU
 
    def forward(self, x):
        # split setup for x, containing a batch of (image, label) as a tensor
        splits = iter(x.split(self.split_size, dim=0))
        s_next = next(splits)
        # initialisation: 
        # - first mini batch goes through seq0 (on dev0)
        # - the output is sent to dev1
        s_prev = self.seq0(s_next).to(self.dev1)
        ret = []
 
        for s_next in splits:
            # A. s_prev runs on dev1
            s_prev = self.seq1(s_prev)
            ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))
 
            # B. s_next runs on dev0, which can run concurrently with A
            s_prev = self.seq0(s_next).to(self.dev1)
 
        s_prev = self.seq1(s_prev)
        ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))
 
        return torch.cat(ret)

En résumé, la fonction forward est modifiée par l’ajout d’un pipeline, par rapport à la version précédente, à partir du paramètre split_size passé en argument à la création du modèle.

Adaptation de la boucle d’entraînement

Création du modèle

C'est lors de la création du modèle, que celui-ci est chargé en mémoire GPU, il ne faut donc pas ajouter de .to(device) après coup.

Dans le cas d'un job mono-tâche, torch va toujours numéroter les GPU à partir de 0 (même s’il s’agit des cartes graphiques 2 et 3 d’après nvidia-smi). On peut donc se permettre de mettre leur identifiant en dur :

mp_model = PipelinedResnet(dev0='cuda:0', dev1='cuda:1')

Extrait de la boucle d’entraînement

def train(model, optimizer, criterion, train_loader, batch_size, dev0, dev1):
    model.train()
    for batch_counter, (images, labels) in enumerate(train_loader):
        # images are sent to the first GPU
        images = images.to(dev0, non_blocking=True)
        # zero the parameter gradients
        optimizer.zero_grad()
        # forward
        outputs = model(images)
        # labels (ground truth) are sent to the GPU where the outputs of the model
        # reside, which in this case is the second GPU 
        labels = labels.to(outputs.device, non_blocking=True)
        _, preds = torch.max(outputs, 1)
        loss = criterion(outputs, labels)
        # backward + optimize only if in training phase
        loss.backward()
        optimizer.step()

Les inputs (images) sont chargés dans le premier GPU (cuda:0), les outputs (valeurs prédites) se trouvent sur le second GPU (cuda:1, fixé dans la déclaration du modèle). Il est donc nécessaire de charger les labels (valeurs réelles) sur le même GPU que les outputs.

L’option non_blockin=True est à utiliser en association avec l’option pin_memory=True du dataloader. Attention, ces options augmentent généralement la quantité de mémoire RAM nécessaire.

Configuration de l'environnement de calcul Slurm

Une seule tâche doit être déclarée pour un modèle (ou ensemble de modèles) et un nombre adéquat de GPU, qui doivent tous être sur le même nœud. Dans notre exemple, 2 GPU.

#SBATCH --gres=gpu:2
#SBATCH --ntasks-per-node=1

La méthodologie présentée, qui ne dépend que de la bibliothèque PyTorch, se limite à un parallélisme mono-nœud multi-GPU (de 2 GPU, 4 GPU ou 8 GPU) et ne peut s'appliquer à un cas multi-nœuds. Il est vivement conseillé d'associer cette technique avec du parallélisme de données, comme décrit dans la page Parallélisme des données et de modèle avec PyTorch, pour accélérer efficacement les entraînements.

Si votre modèle nécessite un parallèlisme de modèle sur plusieurs nœuds, nous vous invitons à explorer les solutions documentées sur la page Distributed Pipeline Parallelism Using RPC.

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