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

Modifier un script d'entraînement PyTorch

Dans cette section, vous apprendrez à modifier les scripts d'entraînement PyTorch pour configurer la bibliothèque parallèle de modèles distribués SageMaker, le partitionnement automatique et le partitionnement manuel.

Note

Pour connaître les versions de PyTorch prises en charge par la bibliothèque, consultez Cadres pris en et Régions AWS.

Astuce

Pour des exemples de bloc-notes de bout en bout qui montrent comment utiliser un script d'entraînement PyTorch avec la bibliothèque parallèle de modèles distribués SageMaker, Exemples PyTorch.

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

PyTorch

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

  1. Importez et initialisez la bibliothèque avec smdistributed.modelparallel.torch.init().

  2. Enveloppez le modèle avec smdistributed.modelparallel.torch.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).

    Note

    Pour les entraînements FP16, vous devez utiliser le gestionnaire de contexte smdistributed.modelparallel.torch.model_creation () pour encapsuler le modèle. Pour plus d'informations, consultez Entraînement FP16 avec parallélisme des modèles.

  3. Enveloppez l'optimiseur avec smdistributed.modelparallel.torch.DistributedOptimizer.

    Note

    Pour la formation FP16, vous devez configurer la mise à l'échelle des pertes statique ou dynamique. Pour plus d'informations, consultez Entraînement FP16 avec parallélisme des modèles.

  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 smdistributed.modelparallel.torch.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)

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 num_workers=0 du chargeur de données. 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.