Configurar 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 el código en un clúster de instancias de máquina virtual (VM), también conocidas como nodos, con variables de entorno que admiten el entrenamiento distribuido de PyTorch. Esto puede ayudar a que tu trabajo de entrenamiento se escale para controlar una gran cantidad de datos.

En esta guía, se da por sentado que usas contenedores de PyTorch compilados con anterioridad para el entrenamiento, como se describe en Comienza a usar PyTorch. Adaptar 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 el trabajo a fin de que use un nodo trabajador principal y uno o más nodos trabajadores. Estas funciones tienen los siguientes comportamientos:

  • Trabajador principal: La VM con rango 0. Este nodo configura las conexiones entre los nodos del clúster.
  • Trabajador: Los nodos restantes del clúster. Cada nodo realiza una parte del entrenamiento, como 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, debes especificar la imagen de un contenedor de Docker para que el trabajador principal la use en el campo trainingInput.masterConfig.imageUri y debes especificar la imagen de un Contenedor de Docker para cada trabajador la use en el campo trainingInput.workerConfig.imageUri. Consulta la lista de contenedores de PyTorch compilados con anterioridad.

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 es el de trainingInput.masterConfig.imageUri. A menudo, tiene sentido usar el mismo contenedor de PyTorch compilado con anterioridad en todos los nodos.

Actualiza tu código de entrenamiento

En la 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 por uno de los backends de entrenamiento distribuido 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 Variables de entorno de esta guía.

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

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

Si deseas obtener más información sobre cómo configurar el entrenamiento distribuido, lee la guía de documentación de PyTorch para el entrenamiento distribuido.

Backends de entrenamiento distribuido

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

  • gloo: Se recomienda para los trabajos de entrenamiento de CPU.
  • nccl: Se recomienda para los trabajos de entrenamiento de GPU

Lee sobre las diferencias entre los backends.

Variables del entorno

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

  • WORLD_SIZE: La cantidad total de nodos en el clúster. Esta variable tiene el mismo valor en cada nodo.
  • RANK: Un identificador único para cada nodo. En el trabajador principal, se establece en 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 en el 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?