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 crees un trabajo de entrenamiento, especifica la imagen de un contenedor de Docker para el
trabajador principal para usarlo en la clase trainingInput.masterConfig.imageUri
y especifica la imagen de una
contenedor de Docker para cada trabajador que se usará en el
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 en0
. En cada trabajador, se establece en un valor diferente de1
aWORLD_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?
- Para obtener más información sobre el entrenamiento con PyTorch en AI Platform Training, sigue el instructivo Comienza a usar PyTorch.
- A fin de obtener más información sobre el entrenamiento distribuido de PyTorch en general, lee la guía para el entrenamiento distribuido de la documentación de PyTorch.