Configuration d'un entraînement distribué pour PyTorch

Ce document explique comment créer un job d'entraînement PyTorch distribué. Lorsque vous créez une tâche d'entraînement distribué, AI Platform Training exécute votre code sur un cluster d'instances de machines virtuelles (VM), également appelées nœuds, avec des variables d'environnement qui acceptent l'entraînement PyTorch distribué. Cela peut aider votre tâche d'entraînement à s'adapter à une grande quantité de données.

Dans ce guide, nous partons du principe que vous utilisez des conteneurs PyTorch prédéfinis pour l'entraînement, comme décrit sur la page Premiers pas avec PyTorch. L'adaptation de votre code PyTorch pour l'entraînement distribué nécessite des modifications minimales.

Spécifier la structure du cluster d'entraînement

Pour l'entraînement PyTorch distribué, configurez votre tâche de sorte qu'elle utilise un nœud de calcul maître et un ou plusieurs nœuds de calcul. Ces rôles se comportent de la façon suivante :

  • Nœud de calcul maître : VM avec le rang 0. Ce nœud configure les connexions entre les nœuds du cluster.
  • Nœud de calcul : nœuds restants du cluster. Chaque nœud se charge d'une partie de l'entraînement, comme spécifié par le code de votre application d'entraînement.

Pour savoir comment spécifier le nœud de calcul maître et les nœuds de calcul pour votre cluster d'entraînement, consultez la page Spécifier des types de machines ou des niveaux d'évolutivité.

Spécifier des images de conteneurs

Lorsque vous créez un job d'entraînement, spécifiez l'image d'un conteneur Docker pour le nœud de calcul maître à utiliser dans trainingInput.masterConfig.imageUri , puis spécifiez l'image d'un conteneur Docker pour chaque nœud de calcul à utiliser dans la trainingInput.workerConfig.imageUri . Consultez la liste des conteneurs PyTorch prédéfinis.

Si vous utilisez la commande gcloud ai-platform jobs submit training pour créer votre tâche d'entraînement, vous pouvez spécifier ces champs avec les options --master-image-uri et --worker-image-uri.

Toutefois, si vous ne spécifiez pas le champ trainingInput.workerConfig.imageUri, sa valeur par défaut est trainingInput.masterConfig.imageUri. Il est souvent judicieux d'utiliser le même conteneur PyTorch prédéfini sur chaque nœud.

Mettre à jour votre code d'entraînement

Dans votre application d'entraînement, ajoutez le code suivant pour initialiser le cluster d'entraînement :

import torch

torch.distributed.init_process_group(
    backend='BACKEND',
    init_method='env://'
)

Remplacez BACKEND par l'un des backends d'entraînement distribué compatibles décrits dans la section suivante. L'argument de mot clé init_method='env://' indique à PyTorch d'utiliser des variables d'environnement pour initialiser la communication dans le cluster. Pour en savoir plus, consultez la section Variables d'environnement de ce guide.

En outre, mettez à jour votre code d'entraînement de sorte qu'il utilise la classe torch.nn.parallel.DistributedDataParallel. Par exemple, si vous avez créé un module PyTorch appelé model dans votre code, ajoutez la ligne suivante :

model = torch.nn.parallel.DistributedDataParallel(model)

Pour en savoir plus sur la configuration de l'entraînement distribué, consultez le guide d'entraînement distribué de PyTorch.

Backends d'entraînement distribué

AI Platform Training est compatible avec les backends suivants pour l'entraînement PyTorch distribué :

  • gloo : recommandé pour les tâches d'entraînement sur processeur
  • nccl : recommandé pour les tâches d'entraînement sur GPU

Découvrez les différences entre les backends.

Variables d'environnement

Lorsque vous créez une tâche d'entraînement distribué PyTorch, AI Platform Training définit les variables d'environnement suivantes sur chaque nœud :

  • WORLD_SIZE : nombre total de nœuds dans le cluster. Chaque nœud prend la même valeur pour cette variable.
  • RANK : identifiant unique de chaque nœud. Pour le nœud de calcul maître, ce paramètre est défini sur 0. Pour chaque nœud de calcul, sa valeur varie de 1 à WORLD_SIZE - 1.
  • MASTER_ADDR : nom d'hôte du nœud de calcul maître. Chaque nœud prend la même valeur pour cette variable.
  • MASTER_PORT : port sur lequel le nœud de calcul maître communique. Chaque nœud prend la même valeur pour cette variable.

PyTorch utilise ces variables d'environnement pour initialiser le cluster.

Étape suivante