Modifier un script PyTorch d'entraînement - Amazon SageMaker
Fractionnement automatique avec PyTorchDécoupage manuel avec PyTorchConsidérationsFonctionnalités de framework non prises en charge

Les traductions sont fournies par des outils de traduction automatique. En cas de conflit entre le contenu d'une traduction et celui de la version originale en anglais, la version anglaise prévaudra.

Modifier un script PyTorch d'entraînement

Dans cette section, vous apprendrez à modifier les scripts d' PyTorch apprentissage afin de configurer la bibliothèque de parallélisme des SageMaker modèles pour le partitionnement automatique et le partitionnement manuel.

Note

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

Astuce

Pour obtenir des exemples de end-to-end blocs-notes illustrant l'utilisation d'un script de PyTorch formation avec la bibliothèque de parallélisme des SageMaker modèles, reportez-vous à. PyTorchExemples

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

Fractionnement automatique avec PyTorch

Les modifications de script d'entraînement suivantes sont nécessaires pour exécuter un script d' PyTorch entraînement avec SageMaker la bibliothèque de parallélisme des modèles :

  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 l'entraînement 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 l'entraînement FP16, vous devez configurer la mise à l'échelle des pertes statiques ou dynamiques. 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 SageMaker de la bibliothèque de parallélisme des modèles, consultez la documentation de 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)

Découpage 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 SageMaker de la bibliothèque de parallélisme des modèles, consultez la documentation de 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)

Considérations

Lorsque vous configurez un script d' PyTorch entraînement à l'aide SageMaker de la bibliothèque de parallélisme des modèles, vous devez tenir compte des points suivants :

  • 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 souhaitez réentraîner un modèle ou l'utiliser à des fins d'inférence sans la bibliothèque après l'avoir entraîné à l'aide de la bibliothèque de parallélisme des modèles, la méthode recommandée est d'enregistrer le modèle complet à l'aide de notre API de point de contrôle et de le charger à nouveau dans un module normal. PyTorch

  • 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. Lors de l'utilisation d'un modèle hybride avec parallélisme des modèles et des données, à ce stade, AllReduce la fin des dégradés est également garantie.

  • Lorsque vous utilisez la bibliothèque en combinaison avec le parallélisme des données, assurez-vous que le nombre de lots sur tous les classements data parallel est le même afin de AllReduce ne pas attendre 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

Fonctionnalités de framework non prises en charge

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

  • Si vous utilisez le parallélisme des données avec le PyTorch DDP natif, le module torch.nn.parallel.DistributedDataParallelwrapper n'est pas pris en charge par la bibliothèque. La bibliothèque gère en interne l'intégration au PyTorch DDP, y compris 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.