Modifier un script d'entraînement PyTorch - Amazon SageMaker

Modifier un script d'entraînement PyTorch

Voici des exemples de scripts d'entraînement que vous pouvez utiliser pour configurer la bibliothèque de modèles parallèles de SageMaker avec les versions 1.7.1 et 1.6.0 de PyTorch, avec les partitionnements automatique et manuel. Nous vous recommandons de consulter les Importantes considérations et Fonctions de cadre non prises en charge avant de créer un script d'entraînement.

Les modifications que vous devez apporter à votre script d'entraînement pour utiliser la bibliothèque sont répertoriées dans PyTorch.

Si vous optez pour le partitionnement manuel, consultez également Partitionnement manuel avec PyTorch.

Astuce

Pour obtenir des exemples de blocs-notes exécutables de bout en bout illustrant l'utilisation d'un script d'entraînement TensorFlow avec la bibliothèque de modèles parallèles distribués SageMaker, consultez Exemples PyTorch.

Vous noterez que le partitionnement automatique est activé par défaut. Sauf indication contraire, les scripts suivants utilisent le partitionnement automatique.

Importantes considérations

Lorsque vous configurez un script d'entraînement PyTorch à l'aide de la bibliothèque de modèles parallèles distribués de SageMaker, vous devez savoir que :

  • Si vous utilisez une technique d'optimisation reposant sur des normes de gradient globales, par exemple une norme de gradient du modèle tout entier, comme certaines variantes de l'optimiseur LAMB ou de l'écrêtage de gradient global, vous devez rassembler toutes les normes entre toutes les partitions de modèle pour vérifier l'exactitude. Pour ce faire, vous pouvez utiliser les types de données de base de communication de la bibliothèque.

  • Tous les arguments torch.Tensor aux méthodes de transmission des modules nn.Modules dans votre modèle doivent être utilisés dans le calcul de la sortie du module. En d'autres termes, la bibliothèque ne prend pas en charge le cas où il existe un argument torch.Tensor à un module dont la sortie du module ne dépend pas.

  • L'argument à l'appel smp.DistributedModel.backward() doit dépendre de toutes les sorties du modèle. En d'autres termes, il ne peut pas y avoir de sortie de l'appel smp.DistributedModel.forward qui ne soit pas utilisée dans le calcul du tenseur qui est intégré à l'appel smp.DistributedModel.backward.

  • S'il y a des appels torch.cuda.synchronize() dans votre code, vous devrez peut-être appeler torch.cuda.set_device(smp.local_rank()) immédiatement avant l'appel de synchronisation. Sinon, des contextes CUDA inutiles pourraient être créés dans le périphérique 0, ce qui consommerait de la mémoire inutilement.

  • Comme la bibliothèque place nn.Modules sur différents périphériques, les modules du modèle ne doivent pas dépendre d'un état global modifié dans smp.step. Tout état qui reste fixe durant tout l'entraînement, ou qui est modifié en dehors de smp.step d'une manière visible par tous les processus, est autorisé.

  • Lorsque vous utilisez la bibliothèque, vous n'avez pas besoin de déplacer le modèle vers le GPU (par exemple, en utilisant model.to(device)). Si vous essayez de déplacer le modèle vers le GPU avant la partition du modèle (avant le premier appel smp.step), l'appel de déplacement est ignoré. La bibliothèque déplace automatiquement la partie du modèle affectée à un rang, vers son GPU. Une fois que l'entraînement avec la bibliothèque démarre, ne déplacez pas le modèle vers le CPU et ne l'utilisez pas, car il ne contiendra pas des paramètres corrects pour les modules non affectés à la partition maintenue par le processus. Si vous voulez ré-entraîner un modèle ou l'utiliser à des fins d'inférence sans la bibliothèque, après qu'il a été entraîné à l'aide de la bibliothèque de modèles parallèles, nous vous recommandons d'enregistrer le modèle complet à l'aide de notre API de pointage et de le rétro-charger dans un module PyTorch normal.

  • Si vous avez une liste de modules telle que la sortie de l'un en alimente un autre, vous pouvez améliorer la performance de façon significative en remplaçant cette liste par nn.Sequential.

  • La mise à jour du poids (optimizer.step()) doit se produire en dehors de smp.step car c'est à ce moment que la transmission vers l'arrière est entièrement terminée et que les gradients sont prêts. Lorsque vous utilisez un modèle hybride pour le parallélisme des modèles et des données à ce stade, vous êtes également assuré que l'opération AllReduce des gradients va se terminer.

  • Lorsque vous utilisez la bibliothèque conjointement avec le parallélisme des données, assurez-vous que le nombre de lots sur tous les rangs de données parallèles est le même afin que AllReduce ne se bloque pas en attendant un rang qui ne participe pas à l'étape.

  • Si vous lancez une tâche d'entraînement à l'aide d'un type d'instance ml.p4d (tel que ml.p4d.24xlarge), vous devez définir la variable du chargeur de données num_workers=0. Par exemple, vous pouvez définir votre DataLoader de la façon suivante :

    dataloader = torch.utils.data.DataLoader( data, batch_size=batch_size, num_workers=0, pin_memory=True, drop_last=True, shuffle=shuffle, )
  • Les entrées de smp.step doivent être les entrées de modèle générées par le DataLoader. En effet, smp.step divise en interne les tenseurs d'entrée sur toute la dimension du lot et les exécute en pipeline. Transmettre le DataLoader lui-même à la fonction smp.step pour générer les entrées de modèle à l'intérieur ne fonctionne donc pas.

    Par exemple, si vous définissez un DataLoader de la façon suivante :

    train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True)

    Vous devez accéder aux entrées de modèle générées par le train_loader et les transmettre à une fonction décorée smp.step. Ne faites pas transmettre le train_loader directement à smp.step.

    def train(model, device, train_loader, optimizer): model.train() for batch_idx, (data, target) in enumerate(train_loader): ... _, loss_mb = train_step(model, data, target) ... @smp.step def train_step(model, data, target): ... return output, loss
  • Les tenseurs d'entrée à smp.step doivent être déplacés vers le périphérique actuel à l'aide de l'API .to(), et cela après l'appel torch.cuda.set_device(local_rank()).

    Par exemple, vous pouvez définir la fonction train de la façon suivante. Cette fonction ajoute data et target sur le périphérique actuel à l'aide de l'API .to() avant d'utiliser ces tenseurs d'entrée pour appeler train_step.

    def train(model, device, train_loader, optimizer): model.train() for batch_idx, (data, target) in enumerate(train_loader): # smdistributed: Move input tensors to the GPU ID used by the current process, # based on the set_device call. data, target = data.to(device), target.to(device) optimizer.zero_grad() # Return value, loss_mb is a StepOutput object _, loss_mb = train_step(model, data, target) # smdistributed: Average the loss across microbatches. loss = loss_mb.reduce_mean() optimizer.step()

    Dans la fonction train ci-dessus, les tenseurs d'entrée de cette fonction décorée smp.set ont été déplacés vers le périphérique actuel. Le modèle ne doit pas être déplacé vers le périphérique actuel. La bibliothèque déplace automatiquement la partie du modèle affectée à un rang, vers son GPU.

    @smp.step def train_step(model, data, target): output = model(data) loss = F.nll_loss(output, target, reduction="mean") model.backward(loss) return output, loss

Fonctions de cadre non prises en charge

Les fonctions PyTorch suivantes ne sont pas prises en charge par la bibliothèque de modèles parallèles distribués de SageMaker :

  • Si vous utilisez le parallélisme de données avec le DDP natif PyTorch, le module wrapper torch.nn.parallel.DistributedDataParallel n'est pas pris en charge par la bibliothèque. La bibliothèque gère en interne l'intégration au DDP PyTorch, notamment la diffusion des paramètres et le gradient AllReduce. Lors de l'utilisation de la bibliothèque, les tampons de module ne sont diffusés qu'une seule fois au début de l'entraînement. Si votre modèle possède des tampons de module qui doivent être synchronisés entre des groupes de données parallèles à chaque étape, vous pouvez le faire à l'aide de l'API torch.distributed, en utilisant le groupe de processus qui peut être obtenu via smp.get_dp_process_group().

  • Pour l'entraînement de précision mixte, le module apex.amp n'est pas pris en charge. Nous vous recommandons d'utiliser la bibliothèque avec une précision mixte automatique en utilisant torch.cuda.amp, à la seule exception d'utiliser smp.amp.GradScaler au lieu de la mise en œuvre dans Torch.

  • torch.jit.ScriptModules ou ScriptFunctions ne sont pas pris en charge par smp.DistributedModel.

  • apex : FusedLayerNorm, FusedAdam, FusedLAMB et FusedNovoGrad de apex ne sont pas pris en charge. À la place, vous pouvez utiliser leurs mises en œuvre par la bibliothèque à l'aide des smp.optimizers et des API smp.nn.

PyTorch

Les modifications de script d'entraînement suivantes sont requises pour exécuter un modèle PyTorch avec la bibliothèque de modèles parallèles distribués de SageMaker :

  1. Importez et initialisez la bibliothèque avec smp.init().

  2. Enveloppez le modèle avec smp.DistributedModel. N'oubliez pas que tous les tenseurs renvoyés par la méthode forward de l'objet nn.Module sous-jacent seront diffusés sur des périphériques avec parallélisme des modèles. Comme cela induira un surdébit de communication, évitez de renvoyer les tenseurs qui ne sont pas nécessaires en dehors de la méthode d'appel (activations intermédiaires, par exemple).

  3. Enveloppez l'optimiseur avec smp.DistributedOptimizer.

  4. Utilisez l'objet DistributedModel renvoyé au lieu d'un modèle utilisateur.

  5. Mettez la logique en avant et en arrière dans une fonction étape et décorez-la avec smp.step.

  6. Restreignez chaque processus à son propre périphérique via torch.cuda.set_device(smp.local_rank()).

  7. Déplacez les tenseurs d'entrée vers le GPU à l'aide de l'API .to() avant l'appel smp.step (voir l'exemple ci-dessous).

  8. Remplacez torch.Tensor.backward et torch.autograd.backward par DistributedModel.backward.

  9. Effectuez un post-traitement sur les sorties des différents micro-lots à l'aide de méthodes StepOutput telles que reduce_mean.

  10. De façon similaire, s'il y a une étape d'évaluation, placez la logique en avant dans une fonction décorée smp.step et post-traitez les sorties en utilisant l'API StepOutput.

  11. Définissez drop_last=True dans DataLoader. Vous pouvez également ignorer manuellement un lot dans la boucle d'entraînement si la taille du lot n'est pas divisible par le nombre de micro-lots.

Pour en savoir plus sur l'API de la bibliothèque de modèles parallèles distribués de SageMaker, consultez la documentation sur l'API.

import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchnet.dataset import SplitDataset from torchvision import datasets import smdistributed.modelparallel.torch as smp class GroupedNet(nn.Module): def __init__(self): super(GroupedNet, self).__init__() # define layers def forward(self, x): # define forward pass and return model outputs # smdistributed: Define smp.step. Return any tensors needed outside. @smp.step def train_step(model, data, target): output = model(data) loss = F.nll_loss(output, target, reduction="mean") model.backward(loss) return output, loss def train(model, device, train_loader, optimizer): model.train() for batch_idx, (data, target) in enumerate(train_loader): # smdistributed: Move input tensors to the GPU ID used by the current process, # based on the set_device call. data, target = data.to(device), target.to(device) optimizer.zero_grad() # Return value, loss_mb is a StepOutput object _, loss_mb = train_step(model, data, target) # smdistributed: Average the loss across microbatches. loss = loss_mb.reduce_mean() optimizer.step() # smdistributed: initialize the backend smp.init() # smdistributed: Set the device to the GPU ID used by the current process. # Input tensors should be transferred to this device. torch.cuda.set_device(smp.local_rank()) device = torch.device("cuda") # smdistributed: Download only on a single process per instance. # When this is not present, the file is corrupted by multiple processes trying # to download and extract at the same time dataset = datasets.MNIST("../data", train=True, download=False) # smdistributed: Shard the dataset based on data-parallel ranks if smp.dp_size() > 1: partitions_dict = {f"{i}": 1 / smp.dp_size() for i in range(smp.dp_size())} dataset = SplitDataset(dataset, partitions=partitions_dict) dataset.select(f"{smp.dp_rank()}") # smdistributed: Set drop_last=True to ensure that batch size is always divisible # by the number of microbatches train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True) model = GroupedNet() optimizer = optim.Adadelta(model.parameters(), lr=4.0) # smdistributed: Use the DistributedModel container to provide the model # to be partitioned across different ranks. For the rest of the script, # the returned DistributedModel object should be used in place of # the model provided for DistributedModel class instantiation. model = smp.DistributedModel(model) optimizer = smp.DistributedOptimizer(optimizer) train(model, device, train_loader, optimizer)

Partitionnement manuel avec PyTorch

Utilisez les gestionnaires de contexte smp.partition pour placer les modules dans des périphériques spécifiques. Tout module non placé dans un contexte smp.partition est placé dans le default_partition. La default_partition doit être fournie si auto_partition est défini sur False. Les modules qui sont créés dans un contexte smp.partition spécifique sont placés sur la partition correspondante.

Pour en savoir plus sur l'API de la bibliothèque de modèles parallèles distribués de SageMaker, consultez la documentation sur l'API.

import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchnet.dataset import SplitDataset from torchvision import datasets import smdistributed.modelparallel.torch as smp class GroupedNet(nn.Module): def __init__(self): super(GroupedNet, self).__init__() with smp.partition(0): # define child modules on device 0 with smp.partition(1): # define child modules on device 1 def forward(self, x): # define forward pass and return model outputs # smdistributed: Define smp.step. Return any tensors needed outside. @smp.step def train_step(model, data, target): output = model(data) loss = F.nll_loss(output, target, reduction="mean") model.backward(loss) return output, loss def train(model, device, train_loader, optimizer): model.train() for batch_idx, (data, target) in enumerate(train_loader): # smdistributed: Move input tensors to the GPU ID used by the current process, # based on the set_device call. data, target = data.to(device), target.to(device) optimizer.zero_grad() # Return value, loss_mb is a StepOutput object _, loss_mb = train_step(model, data, target) # smdistributed: Average the loss across microbatches. loss = loss_mb.reduce_mean() optimizer.step() # smdistributed: initialize the backend smp.init() # smdistributed: Set the device to the GPU ID used by the current process. # Input tensors should be transferred to this device. torch.cuda.set_device(smp.local_rank()) device = torch.device("cuda") # smdistributed: Download only on a single process per instance. # When this is not present, the file is corrupted by multiple processes trying # to download and extract at the same time dataset = datasets.MNIST("../data", train=True, download=False) # smdistributed: Shard the dataset based on data-parallel ranks if smp.dp_size() > 1: partitions_dict = {f"{i}": 1 / smp.dp_size() for i in range(smp.dp_size())} dataset = SplitDataset(dataset, partitions=partitions_dict) dataset.select(f"{smp.dp_rank()}") # smdistributed: Set drop_last=True to ensure that batch size is always divisible # by the number of microbatches train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True) model = GroupedNet() optimizer = optim.Adadelta(model.parameters(), lr=4.0) # smdistributed: Use the DistributedModel container to provide the model # to be partitioned across different ranks. For the rest of the script, # the returned DistributedModel object should be used in place of # the model provided for DistributedModel class instantiation. model = smp.DistributedModel(model) optimizer = smp.DistributedOptimizer(optimizer) train(model, device, train_loader, optimizer)