Configura el entrenamiento distribuido para PyTorch

En este documento, se explica cómo crear un trabajo de entrenamiento distribuido de PyTorch. Cuando creas un trabajo de entrenamiento distribuido, AI Platform Training ejecuta tu código en un clúster de instancias de máquina virtual (VM), también conocido como nodos, con variables de entorno que admiten el entrenamiento distribuido de PyTorch. Esto puede ayudar a tu trabajo de entrenamiento a escalar para manejar una gran cantidad de datos.

En esta guía, se supone que usas contenedores de PyTorch compilados con anterioridad para el entrenamiento, como se describe en Comienza a usar PyTorch. La adaptación de tu código de PyTorch para el entrenamiento distribuido requiere cambios mínimos.

Especifica la estructura del clúster de entrenamiento

Para el entrenamiento distribuido de PyTorch, configura tu trabajo a fin de que use un nodo de trabajador principal y uno o más nodos de trabajador. Estas funciones tienen los siguientes comportamientos:

  • Trabajador principal: la VM con rango 0. Este nodo establece conexiones entre los nodos del clúster.
  • Trabajador: Los nodos restantes en el clúster. Cada nodo realiza una parte del entrenamiento, según lo especifica el código de tu aplicación de entrenamiento.

A fin de obtener más información sobre cómo especificar el nodo de trabajador principal y los nodos de trabajador para tu clúster de entrenamiento, consulta Especifica los tipos de máquina o los niveles de escala.

Especifica imágenes de contenedor

Cuando creas un trabajo de entrenamiento, especifica la imagen de un contenedor de Docker para que el trabajador principal la use en el campo trainingInput.masterConfig.imageUri y especifica la imagen de un contenedor de Docker para usar en el Campo trainingInput.workerConfig.imageUri. Consulta la lista de contenedores prediseñados de PyTorch.

Si usas el comando gcloud ai-platform jobs submit training para crear tu trabajo de entrenamiento, puedes especificar estos campos con las marcas --master-image-uri y --worker-image-uri.

Sin embargo, si no especificas el campo trainingInput.workerConfig.imageUri, su valor predeterminado será trainingInput.masterConfig.imageUri. Generalmente, tiene sentido usar el mismo contenedor de PyTorch compilado previamente en cada nodo.

Actualiza tu código de entrenamiento

En tu aplicación de entrenamiento, agrega el siguiente código para inicializar el clúster de entrenamiento:

import torch

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

Reemplaza BACKEND con uno de los backends de entrenamiento distribuidos compatibles que se describen en la siguiente sección. El argumento de palabra clave init_method='env://' le indica a PyTorch que use variables de entorno para inicializar la comunicación en el clúster. Obtén más información en la sección de variables de entorno de esta guía.

Además, actualiza tu código de entrenamiento para usar la clase torch.nn.parallel.DistributedDataParallel. Por ejemplo, si creaste un módulo PyTorch llamado model en tu código, agrega la línea siguiente:

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

Para obtener más información sobre la configuración del entrenamiento distribuido, consulta la guía de la documentación de PyTorch en el entrenamiento distribuido.

Backends de entrenamiento distribuido

AI Platform Training admite los siguientes backends para el entrenamiento distribuido de PyTorch:

  • gloo: recomendado para trabajos de entrenamiento de CPU
  • nccl: recomendada para los trabajos de entrenamiento de GPU

Lee sobre las diferencias entre los backends.

Variables del entorno

Cuando creas un trabajo de entrenamiento de PyTorch distribuido, AI Platform Training establece las siguientes variables de entorno en cada nodo:

  • WORLD_SIZE: La cantidad total de nodos del clúster. Esta variable tiene el mismo valor en todos los nodos.
  • RANK: un identificador único para cada nodo. En el trabajador principal, esto se configura como 0. En cada trabajador, se establece en un valor diferente de 1 a WORLD_SIZE - 1.
  • MASTER_ADDR: el nombre de host del nodo trabajador principal. Esta variable tiene el mismo valor en cada nodo.
  • MASTER_PORT: El puerto al que se comunica el nodo trabajador principal. Esta variable tiene el mismo valor en cada nodo.

PyTorch usa estas variables de entorno para inicializar el clúster.

¿Qué sigue?