Modifier un script d'entraînement PyTorch

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 parallélisme de modè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 de 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 de données, assurez-vous que le nombre de lots sur tous les rangs de parallélisme de données 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

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.