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

Modifier un script d'entraînement PyTorch

Dans la bibliothèque parallèle de données SageMaker v1.4.0 et ultérieure, la bibliothèque est disponible en tant qu'option de backend pour le package distribué PyTorch. Il vous suffit d'importer la bibliothèque une fois au début de votre script d'entraînement et de la définir comme le backend distribué de PyTorch pendant l'initialisation. Avec la seule ligne de spécification du backend, vous pouvez garder votre script d'entraînement PyTorch inchangé et utiliser directement les modules distribués de PyTorch. Pour trouver la dernière documentation sur les API de la bibliothèque, consultez la section SageMaker distributed data parallel APIs for PyTorch (API parallèles de données distribuées de SageMaker pour PyTorch) dans la documentation du kit SDK Python de SageMaker. Pour en savoir plus sur le package distribué PyTorch et les options de backend, consultez la section Distributed communication package - torch.distributed (Package de communication distribué — torch.distributed).

Important

Étant donné que la bibliothèque de parallélisme de données distribuées SageMaker v1.4.0 et ultérieures fonctionne en tant que backend distribué de PyTorch, les API smdistributed suivantes pour le package distribué de PyTorch sont rendues obsolètes.

Si vous devez utiliser les versions antérieures de la bibliothèque (v1.3.0 ou antérieures), consultez la documentation archivée de la bibliothèque parallèle de données distribuées SageMaker dans la documentation du kit SDK Python de SageMaker.

Utiliser la bibliothèque parallèle de données distribuées SageMaker comme backend de torch.distributed

Pour utiliser la bibliothèque de données parallèles distribuées SageMaker, la seule chose à faire est d'importer le client PyTorch de la bibliothèque de données parallèles distribuées SageMaker (smdistributed.dataparallel.torch.torch_smddp). Le client enregistre smddp en tant que backend pour PyTorch. Lorsque vous initialisez le groupe de processus distribués PyTorch en utilisant l'API torch.distributed.init_process_group, assurez-vous de spécifier 'smddp' pour l'argument backend.

import smdistributed.dataparallel.torch.torch_smddp import torch.distributed as dist dist.init_process_group(backend='smddp')
Note

Le backend smddp ne prend actuellement pas en charge la création de groupes de sous-processus avec l'API torch.distributed.new_group(). Vous ne pouvez pas utiliser le backend smddp en même temps que d'autres backends de groupes de processus tels que NCCL et Gloo.

Si vous avez déjà un script PyTorch qui fonctionne et que vous avez seulement besoin d'ajouter les spécifications du backend, vous pouvez passer à Utilisation des estimateurs de cadre SageMaker pour PyTorch et TensorFlow dans la rubrique Étape 2 : lancer une tâche d'entraînement distribuée SageMaker à l'aide du kit SDK Python SageMaker.

Si vous devez encore modifier votre script d'entraînement pour utiliser correctement le package distribué par PyTorch, suivez le reste des procédures de cette page.

Préparation d'un script d'entraînement PyTorch pour l'entraînement distribué

Les étapes suivantes fournissent des conseils supplémentaires sur la façon de préparer votre script d'entraînement pour exécuter avec succès une tâche d'entraînement distribuée à l'aide de PyTorch.

Note

Dans la version 1.4.0, la bibliothèque parallèle de données distribuées de SageMaker prend en charge les types de données primitives collectives suivants de l'interface torch.distributed : all_reduce, broadcast, reduce, all_gather et barrier.

  1. Importez les modules distribués de PyTorch.

    import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP
  2. Après avoir analysé les arguments et défini un paramètre de taille de lot (batch_size=args.batch_size, par exemple), ajoutez deux lignes de code pour redimensionner la taille du lot par employé (GPU). L'opération DataLoader de PyTorch ne gère pas automatiquement le redimensionnement du lot pour l'entraînement distribué.

    batch_size //= dist.get_world_size() batch_size = max(batch_size, 1)
  3. Épinglez chaque GPU à un seul processus de bibliothèque de données parallèles SageMaker avec local_rank : cela fait référence au rang relatif du processus au sein d'un nœud donné.

    Vous pouvez récupérer le rang du processus à partir de la variable d'environnement LOCAL_RANK.

    import os local_rank = os.environ["LOCAL_RANK"] torch.cuda.set_device(local_rank)
  4. Après avoir défini un modèle, encapsulez-le avec l'API DistributedDataParallel de PyTorch.

    model = ... # Wrap the model with the PyTorch DistributedDataParallel API model = DDP(model)
  5. Lorsque vous appelez l'API torch.utils.data.distributed.DistributedSampler, indiquez le nombre total de processus (GPU) participant à l'entraînement sur tous les nœuds du cluster. C'est ce que l'on appelle world_size et vous pouvez récupérer le numéro à partir de l'API torch.distributed.get_world_size(). Précisez également le rang de chaque processus parmi tous les processus utilisant l'API torch.distributed.get_rank().

    from torch.utils.data.distributed import DistributedSampler train_sampler = DistributedSampler( train_dataset, num_replicas = dist.get_world_size(), rank = dist.get_rank() )
  6. Modifiez votre script pour enregistrer les points de contrôle uniquement sur le processus leader (rang 0). Le processus leader possède un modèle synchronisé. Cela évite également que d'autres processus écrasent les points de contrôle et les corrompent éventuellement.

    if dist.get_rank() == 0: torch.save(...)

L'exemple de code suivant montre la structure d'un script d'entraînement PyTorch avec smddp comme backend.

import os import torch # SageMaker data parallel: Import the library PyTorch API import smdistributed.dataparallel.torch.torch_smddp # SageMaker data parallel: Import PyTorch's distributed API import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP # SageMaker data parallel: Initialize the process group dist.init_process_group(backend='smddp') class Net(nn.Module): ... # Define model def train(...): ... # Model training def test(...): ... # Model evaluation def main(): # SageMaker data parallel: Scale batch size by world size batch_size //= dist.get_world_size() batch_size = max(batch_size, 1) # Prepare dataset train_dataset = torchvision.datasets.MNIST(...) # SageMaker data parallel: Set num_replicas and rank in DistributedSampler train_sampler = torch.utils.data.distributed.DistributedSampler( train_dataset, num_replicas=dist.get_world_size(), rank=dist.get_rank()) train_loader = torch.utils.data.DataLoader(..) # SageMaker data parallel: Wrap the PyTorch model with the library's DDP model = DDP(Net().to(device)) # SageMaker data parallel: Pin each GPU to a single library process. local_rank = os.environ["LOCAL_RANK"] torch.cuda.set_device(local_rank) model.cuda(local_rank) # Train optimizer = optim.Adadelta(...) scheduler = StepLR(...) for epoch in range(1, args.epochs + 1): train(...) if rank == 0: test(...) scheduler.step() # SageMaker data parallel: Save model on master node. if dist.get_rank() == 0: torch.save(...) if __name__ == '__main__': main()

Une fois que vous avez terminé d'adapter votre scénario d'entraînement, passez à Étape 2 : lancer une tâche d'entraînement distribuée SageMaker à l'aide du kit SDK Python SageMaker.