Entrenamiento distribuido

En esta página, se describe cómo ejecutar trabajos de entrenamiento distribuidos en Vertex AI.

Requisitos de código

Usar un marco de trabajo de AA que admita el entrenamiento distribuido En tu código de entrenamiento, puedes usar las variables de entorno CLUSTER_SPEC o TF_CONFIG para hacer referencia a partes específicas de tu clúster de entrenamiento.

Estructura del clúster de entrenamiento

Si ejecutas un trabajo de entrenamiento distribuido con IA de Vertex, debes especificar varias máquinas (nodos) en un clúster de entrenamiento. El servicio de entrenamiento asigna los recursos para los tipos de máquina que especifiques. El trabajo en ejecución en un nodo dado se llama réplica. Un grupo de réplicas con la misma configuración se denomina grupo de trabajadores.

A cada réplica del clúster de entrenamiento se le asigna una función o tarea únicas en el entrenamiento distribuido. Por ejemplo:

  • Réplica principal: Se designa exactamente una réplica como réplica principal. Esta tarea administra las demás y también informa el estado del trabajo en su conjunto.

  • Trabajadores: Se pueden designar una o más réplicas como trabajadores. Estas réplicas hacen su parte del trabajo según lo designes en tu configuración del trabajo.

  • Servidores de parámetros: Si el framework de AA lo admite, una o más réplicas se pueden designar como servidores de parámetros. Estas réplicas almacenan los parámetros del modelo y coordinan el estado del modelo compartido entre los trabajadores.

  • Evaluadores: Si tu framework de AA lo admite, se pueden designar una o más réplicas como evaluadores. Estas réplicas se pueden usar para evaluar el modelo. Si usas TensorFlow, ten en cuenta que, por lo general, TensorFlow espera que no uses más de un evaluador.

Configura un trabajo de entrenamiento distribuido

Puedes configurar cualquier trabajo de entrenamiento personalizado como un trabajo de entrenamiento distribuido si defines varios grupos de trabajadores. También puedes ejecutar entrenamiento distribuido dentro de una canalización de entrenamiento o un trabajo de ajuste de hiperparámetros.

A fin de configurar un trabajo de entrenamiento distribuido, define tu lista de grupos de trabajadores (workerPoolSpecs[]) y designa una WorkerPoolSpec para cada tipo de tarea:

Posición en workerPoolSpecs[] Tarea realizada en el clúster
Primera (workerPoolSpecs[0]) Principal o programadora
Segunda (workerPoolSpecs[1]) Secundaria, réplicas o trabajadores
Tercera (workerPoolSpecs[2]) Servidores de parámetros, servidor de reducción
Cuarta (workerPoolSpecs[3]) Evaluadores

Debes especificar una réplica principal que coordine el trabajo realizado por todas las demás réplicas. Usa la primera especificación de grupo de trabajadores solo para tu réplica principal y establece replicaCount en 1:

{
  "workerPoolSpecs": [
     // `WorkerPoolSpec` for worker pool 0, primary replica, required
     {
       "machineSpec": {...},
       "replicaCount": 1,
       "diskSpec": {...},
       ...
     },
     // `WorkerPoolSpec` for worker pool 1, optional
     {},
     // `WorkerPoolSpec` for worker pool 2, optional
     {},
     // `WorkerPoolSpec` for worker pool 3, optional
     {}
   ]
   ...
}

Especifica grupos de trabajadores adicionales

Según tu framework de AA, puedes especificar grupos de trabajadores adicionales para otros fines. Por ejemplo, si usas TensorFlow, puedes especificar grupos de trabajadores para configurar réplicas de trabajadores, réplicas de servidores de parámetros y réplicas de evaluadores.

El orden de los grupos de trabajadores que especifiques en la lista workerPoolSpecs[] determina el tipo de grupo de trabajadores. Establece valores vacíos para los grupos de trabajadores que no deseas usar, de modo que puedas omitirlos en la lista workerPoolSpecs[] para especificar los grupos de trabajadores que deseas usar. Por ejemplo:

Si deseas especificar un trabajo que solo tenga una réplica principal y un grupo de trabajadores del servidor de parámetros, debes configurar un valor vacío para el grupo de trabajadores 1:

{
  "workerPoolSpecs": [
     // `WorkerPoolSpec` for worker pool 0, required
     {
       "machineSpec": {...},
       "replicaCount": 1,
       "diskSpec": {...},
       ...
     },
     // `WorkerPoolSpec` for worker pool 1, optional
     {},
     // `WorkerPoolSpec` for worker pool 2, optional
     {
       "machineSpec": {...},
       "replicaCount": 1,
       "diskSpec": {...},
       ...
     },
     // `WorkerPoolSpec` for worker pool 3, optional
     {}
   ]
   ...
}

Reduce el tiempo de entrenamiento con el servidor de reducción

Cuando entrenas un modelo de AA grande con varios nodos, la comunicación de gradientes entre nodos puede contribuir con una latencia significativa. El servidor de reducción es un algoritmo de reducción total que puede aumentar la capacidad de procesamiento y reducir la latencia para el entrenamiento distribuido. Vertex AI pone a disposición el servidor de reducción en una imagen de contenedor de Docker que puedes usar para uno de tus grupos de trabajadores durante el entrenamiento distribuido.

Puedes usar el servidor de reducción si cumples con los siguientes requisitos:

  • Estás realizando un entrenamiento distribuido con trabajadores de GPU.

  • Tu código de entrenamiento usa TensorFlow o PyTorch y está configurado para el entrenamiento de datos paralelos de varios hosts con GPU mediante el uso de NCCL para reducir todo. (También puedes usar otros marcos de trabajo de AA que usen NCCL).

  • Los contenedores que se ejecutan en tu nodo principal (workerPoolSpecs[0]) y trabajadores (workerPoolSpecs[1]) son compatibles con el servidor de reducción. En particular, cada contenedor es uno de los siguientes:

    • Un contenedor de entrenamiento de TensorFlow compilado previamente, versión 2.3 o posterior

    • Un contenedor de entrenamiento de Pytorch, versión 1.4 o posterior.

    • Un contenedor personalizado con NCCL 2.7 o posterior y el paquete google-reduction-server instalado. Puedes instalar este paquete en una imagen de contenedor personalizada si agregas la siguiente línea a tu Dockerfile:

      RUN echo "deb https://packages.cloud.google.com/apt google-fast-socket main" | tee /etc/apt/sources.list.d/google-fast-socket.list && \
          curl -s -L https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - && \
          apt update && apt install -y google-reduction-server
      

Para usar el servidor de reducción, haz lo siguiente cuando crees un recurso de entrenamiento personalizado:

  1. Especifica uno de los siguientes URI en el campo containerSpec.imageUri del tercer grupo de trabajadores (workerPoolSpecs[2]):

    • us-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest
    • europe-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest
    • asia-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest

    Elegir la multirregión más cercana a donde realizas el entrenamiento personalizado puede reducir la latencia.

  2. Cuando selecciones el tipo de máquina y la cantidad de nodos del tercer grupo de trabajadores, asegúrate de que el ancho de banda total de la red del tercer grupo de trabajadores coincida o supere el ancho de banda total de la red del primer y segundo grupo de trabajadores

    Para obtener información sobre el ancho de banda máximo disponible de cada nodo en el segundo grupo de trabajadores, consulta Ancho de banda de red y GPU.

    No usas GPU para los nodos del servidor de reducción. Para conocer el ancho de banda máximo disponible de cada nodo en el tercer grupo de trabajadores, consulta las columnas “Ancho de banda de salida máximo (Gbps)” en Familia de máquinas de uso general.

    Por ejemplo, si configuras el primer y el segundo grupo de trabajadores para que usen 5 nodos n1-highmem-96, cada uno con 8 GPU NVIDIA_TESLA_V100, cada nodo tiene un ancho de banda máximo disponible de 100 Gbps, para un total de 500 Gbps. Para hacer coincidir este ancho de banda en el tercer grupo de trabajadores, puedes usar 16 nodos n1-highcpu-16, cada uno con un ancho de banda máximo de 32 Gbps, para un ancho de banda total de 512 Gbps.

    Te recomendamos que uses el tipo de máquina n1-highcpu-16 para los nodos del servidor de reducción, ya que este tipo de máquina ofrece un ancho de banda relativamente alto para sus recursos.

Con el siguiente comando, se proporciona un ejemplo de cómo crear un recurso CustomJob que use el servidor de reducción:

gcloud beta ai custom-jobs create \
  --region=LOCATION \
  --display-name=JOB_NAME \
  --worker-pool-spec=machine-type=n1-highmem-96,replica-count=1,accelerator-type=NVIDIA_TESLA_V100,accelerator-count=8,container-image-uri=CUSTOM_CONTAINER_IMAGE_URI \
  --worker-pool-spec=machine-type=n1-highmem-96,replica-count=4,accelerator-type=NVIDIA_TESLA_V100,accelerator-count=8,container-image-uri=CUSTOM_CONTAINER_IMAGE_URI \
  --worker-pool-spec=machine-type=n1-highcpu-16,replica-count=16,container-image-uri=us-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest

Para obtener más información, lee la guía para crear un CustomJob.

Variables de entorno para tu clúster

Vertex AI propaga una variable de entorno, CLUSTER_SPEC, en cada réplica para describir cómo se configura el clúster general. Al igual que TF_CONFIG de TensorFlow, CLUSTER_SPEC describe cada réplica del clúster, incluido el índice y la función (réplica principal, trabajador, servidor de parámetros o evaluador).

Cuando ejecutas un entrenamiento distribuido con TensorFlow, se analiza TF_CONFIG para compilar tf.train.ClusterSpec. De manera similar, cuando ejecutas el entrenamiento distribuido con otros marcos de trabajo de AA, debes analizar CLUSTER_SPEC para propagar las variables de entorno o la configuración requerida por el marco.

El formato de CLUSTER_SPEC

La variable de entorno CLUSTER_SPEC es una string JSON con el siguiente formato:

Clave Descripción
"cluster"

La descripción del clúster de tu contenedor personalizado. Al igual que con TF_CONFIG, este objeto tiene el formato de una especificación de clúster de TensorFlow, y se puede pasar al constructor de tf.train.ClusterSpec.

La descripción del clúster contiene una lista de nombres de réplicas para cada grupo de trabajadores que especifiques.

"workerpool0" Todos los trabajos de entrenamiento distribuido tienen una réplica principal en el primer grupo de trabajadores.
"workerpool1" Este grupo de trabajadores contiene réplicas de trabajadores, si las especificaste cuando creaste el trabajo.
"workerpool2" Este grupo de trabajadores contiene servidores de parámetros, si los especificaste cuando creaste tu trabajo.
"workerpool3" Este grupo de trabajadores contiene evaluadores, si los especificaste cuando creaste el trabajo.
"environment" La string cloud.
"task" Describe la tarea del nodo particular en el que se está ejecutando el código. Puedes usar esta información a fin de escribir código para trabajadores específicos en un trabajo distribuido. Esta entrada es un diccionario con las siguientes claves:
"type" El tipo de grupo de trabajadores en el que se ejecuta esta tarea. Por ejemplo, "workerpool0" hace referencia a la réplica principal.
"index"

El índice basado en cero de la tarea. Por ejemplo, si el trabajo de entrenamiento incluye dos trabajadores, este valor se establece en 0 en uno de ellos y en 1 en el otro.

"trial" El identificador de la prueba de ajuste de hiperparámetros que se está ejecutando actualmente. Cuando configuras el ajuste de hiperparámetros de tu trabajo, estableces una cantidad de pruebas para entrenar. Este valor te proporciona una forma de diferenciar en tu código las pruebas que se están ejecutando. El identificador es un valor de string que contiene el número de prueba, que comienza en 1.
job

La CustomJobSpec que proporcionaste para crear el trabajo de entrenamiento actual, representada como un diccionario.

Ejemplo CLUSTER_SPEC

A continuación, se muestra un valor de ejemplo:


{
   "cluster":{
      "workerpool0":[
         "cmle-training-workerpool0-ab-0:2222"
      ],
      "workerpool1":[
         "cmle-training-workerpool1-ab-0:2222",
         "cmle-training-workerpool1-ab-1:2222"
      ],
      "workerpool2":[
         "cmle-training-workerpool2-ab-0:2222",
         "cmle-training-workerpool2-ab-1:2222"
      ],
      "workerpool3":[
         "cmle-training-workerpool3-ab-0:2222",
         "cmle-training-workerpool3-ab-1:2222",
         "cmle-training-workerpool3-ab-2:2222"
      ]
   },
   "environment":"cloud",
   "task":{
      "type":"workerpool0",
      "index":0,
      "trial":"TRIAL_ID"
   },
   "job": {
      ...
   }
}

El formato de TF_CONFIG

Además de CLUSTER_SPEC, Vertex AI establece la variable de entorno TF_CONFIG en cada réplica de todos los trabajos de entrenamiento distribuidos. Vertex AI no configura TF_CONFIG para trabajos de entrenamiento de una sola réplica.

CLUSTER_SPEC y TF_CONFIG comparten algunos valores, pero tienen formatos diferentes. Ambas variables de entorno incluyen campos adicionales más allá de los que requiere TensorFlow.

El entrenamiento distribuido con TensorFlow funciona de la misma manera con contenedores personalizados que con un contenedor compilado previamente.

La variable de entorno TF_CONFIG es una string JSON con el siguiente formato:

Campos TF_CONFIG
cluster

La descripción del clúster de TensorFlow. Un diccionario que asigna uno o más nombres de tareas (chief, worker, ps o master) a las listas de direcciones de red en las que se ejecutan estas tareas. Para un trabajo de entrenamiento determinado, este diccionario es el mismo en todas las VM.

Este es un primer argumento válido para el constructor tf.train.ClusterSpec. Ten en cuenta que este diccionario nunca contiene evaluator como clave, ya que los evaluadores no se consideran parte del clúster de entrenamiento, incluso si los usas para el trabajo.

task

La descripción de la tarea de la VM en la que se configura esta variable de entorno. Para un trabajo de entrenamiento determinado, este diccionario es diferente en cada VM. Puedes usar esta información para personalizar el código que se ejecuta en cada VM en un trabajo de entrenamiento distribuido. También puedes usarla para cambiar el comportamiento del código de entrenamiento en diferentes pruebas de un trabajo de ajuste de hiperparámetros.

En este diccionario, se incluyen los siguientes pares clave-valor:

Campos task
type

El tipo de tarea que realiza esta VM. Este valor se establece en worker en los trabajadores, ps en los servidores de parámetros y evaluator en los evaluadores. En el trabajador principal del trabajo, el valor se establece en chief o master. Obtén más información sobre la diferencia entre los dos en la sección chief en comparación con master de este documento.

index

El índice basado en cero de la tarea. Por ejemplo, si el trabajo de entrenamiento incluye dos trabajadores, este valor se establece en 0 en uno de ellos y en 1 en el otro.

trial

El ID de la prueba de ajuste de hiperparámetros que se está ejecutando en el momento en esta VM. Este campo solo se configura si el trabajo de entrenamiento actual es un trabajo de ajuste de hiperparámetros.

En los trabajos de ajuste de hiperparámetros, Vertex AI ejecuta tu código de entrenamiento de forma repetida en muchas pruebas con diferentes hiperparámetros cada vez. Este campo contiene el número de prueba actual, que comienza en 1 para la primera prueba.

cloud

Un ID que utiliza Vertex AI de forma interna. Puedes ignorar este campo.

job

La CustomJobSpec que proporcionaste para crear el trabajo de entrenamiento actual, representada como un diccionario.

environment

La string cloud.

Ejemplo TF_CONFIG

Con el siguiente código de ejemplo, se imprime la variable de entorno TF_CONFIG en los registros de entrenamiento:

import json
import os

tf_config_str = os.environ.get('TF_CONFIG')
tf_config_dict  = json.loads(tf_config_str)

# Convert back to string just for pretty printing
print(json.dumps(tf_config_dict, indent=2))

En un trabajo de ajuste de hiperparámetros que se ejecuta en la versión 2.1 o posterior del entorno de ejecución y usa un trabajador principal, dos trabajadores y un servidor de parámetros, este código produce el siguiente registro para uno de los trabajadores durante la primera prueba de ajuste de hiperparámetros. El resultado de ejemplo oculta el campo job para mayor concisión y reemplaza algunos ID por valores genéricos.

{
  "cluster": {
    "chief": [
      "training-workerpool0-[ID_STRING_1]-0:2222"
    ],
    "ps": [
      "training-workerpool2-[ID_STRING_1]-0:2222"
    ],
    "worker": [
      "training-workerpool1-[ID_STRING_1]-0:2222",
      "training-workerpool1-[ID_STRING_1]-1:2222"
    ]
  },
  "environment": "cloud",
  "job": {
    ...
  },
  "task": {
    "cloud": "[ID_STRING_2]",
    "index": 0,
    "trial": "1",
    "type": "worker"
  }
}

Cuándo usar TF_CONFIG

TF_CONFIG solo se configura para los trabajos de entrenamiento distribuido.

Es probable que no necesite interactuar con la variable de entorno TF_CONFIG directamente en tu código de entrenamiento. Solo accede a la variable de entorno TF_CONFIG si las estrategias de distribución de TensorFlow y el flujo de trabajo de ajuste de hiperparámetros estándar de Vertex AI, descritos en las siguientes secciones, no funcionan en tu trabajo.

Entrenamiento distribuido

Vertex AI establece la variable de entorno TF_CONFIG a fin de extender las especificaciones que TensorFlow requiere para el entrenamiento distribuido.

Para realizar un entrenamiento distribuido con TensorFlow, usa la API tf.distribute.Strategy. En particular, te recomendamos que uses la API de Keras junto con la MultiWorkerMirroredStrategy o, si especificas los servidores de parámetros para el trabajo, la ParameterServerStrategy. Sin embargo, ten en cuenta que, en la actualidad, TensorFlow solo proporciona asistencia experimental para estas estrategias.

Estas estrategias de distribución usan la variable de entorno TF_CONFIG para asignar funciones a cada VM en el trabajo de entrenamiento y facilitar la comunicación entre las VM. No necesitas acceder a la variable de entorno TF_CONFIG directamente en el código de entrenamiento, ya que TensorFlow la controla por ti.

Solo analiza directamente la variable de entorno TF_CONFIG si deseas personalizar el comportamiento de las diferentes VM que ejecutan el trabajo de entrenamiento.

Ajuste de hiperparámetros

Cuando ejecutas un trabajo de ajuste de hiperparámetros, Vertex AI proporciona diferentes argumentos al código de entrenamiento para cada prueba. No es necesario que el código de entrenamiento tenga en cuenta la prueba que se encuentra en ejecución. Además, puedes supervisar el progreso de los trabajos de ajuste de hiperparámetros en Cloud Console.

Si es necesario, el código puede leer el número de prueba actual del campo trial del campo task de la variable de entorno TF_CONFIG.

¿Qué sigue?