Preparación distribuida

En esta página se describe cómo ejecutar tareas de entrenamiento distribuidas en Vertex AI.

Requisitos del código

Usa un framework de aprendizaje automático que admita el entrenamiento distribuido. En el 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 una tarea de entrenamiento distribuido con Vertex AI, debes especificar varias máquinas (nodos) en un clúster de entrenamiento. El servicio de entrenamiento asigna los recursos de los tipos de máquinas que especifiques. El trabajo que se está ejecutando en un nodo determinado se denomina 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 un único rol o tarea en el entrenamiento distribuido. Por ejemplo:

  • Réplica principal: se designa exactamente una réplica como réplica principal. Esta tarea gestiona las demás y comunica el estado del trabajo en su conjunto.

  • Trabajadores: se pueden designar una o varias réplicas como trabajadores. Estas réplicas hacen la parte del trabajo que les asignes en la configuración del trabajo.

  • Servidor(es) de parámetros: si tu framework de aprendizaje automático lo admite, se pueden designar una o varias réplicas como servidores de parámetros. Estas réplicas almacenan parámetros del modelo y coordinan el estado del modelo compartido entre los trabajadores.

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

Configurar un trabajo de entrenamiento distribuido

Puedes configurar cualquier trabajo de entrenamiento personalizado como trabajo de entrenamiento distribuido definiendo varios grupos de trabajadores. También puedes ejecutar el entrenamiento distribuido en una canalización de entrenamiento o en una tarea de ajuste de hiperparámetros.

Para configurar una tarea de entrenamiento distribuido, define tu lista de grupos de trabajadores (workerPoolSpecs[]) y asigna un WorkerPoolSpec a cada tipo de tarea:

Posición en workerPoolSpecs[] Tarea realizada en el clúster
Primera (workerPoolSpecs[0]) Principal, jefe, programador o "maestro"
Segundo (workerPoolSpecs[1]) Secundarias, réplicas y trabajadores
Tercero (workerPoolSpecs[2]) Servidores de parámetros y servidores de reducción
Cuarto (workerPoolSpecs[3]) Evaluadores

Debes especificar una réplica principal, que coordina el trabajo realizado por todas las demás réplicas. Usa la primera especificación del grupo de trabajadores solo para tu réplica principal y asigna el valor 1 a su replicaCount:

{
  "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
     {}
   ]
   ...
}

Especificar grupos de trabajadores adicionales

En función de tu framework de aprendizaje automático, 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. Define valores vacíos para los grupos de trabajadores que no quieras usar. De esta forma, podrás omitirlos en la lista workerPoolSpecs[] para especificar los grupos de trabajadores que sí quieras usar. Por ejemplo:

Si quieres especificar un trabajo que solo tenga una réplica principal y un grupo de trabajadores de servidor de parámetros, debes asignar un valor vacío al 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
     {}
   ]
   ...
}

Reducir el tiempo de formación con Reduction Server

Cuando entrenas un modelo de aprendizaje automático grande con varios nodos, la comunicación de gradientes entre nodos puede contribuir a una latencia significativa. Reduction Server es un algoritmo de reducción total que puede aumentar el rendimiento y reducir la latencia del entrenamiento distribuido. Vertex AI pone a tu disposición Reduction Server en una imagen de contenedor Docker que puedes usar en uno de tus grupos de trabajadores durante el entrenamiento distribuido.

Para saber cómo funciona Reduction Server, consulta el artículo Entrenamiento de GPUs distribuidas más rápido con Reduction Server en Vertex AI.

Requisitos previos

Puedes usar Reduction Server si cumples 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 paralelo de datos multihost con GPUs mediante la reducción total NCCL. También puedes usar otros frameworks de aprendizaje automático que usen NCCL.

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

    • Un contenedor de entrenamiento de TensorFlow prediseñado, versión 2.3 o posterior.

    • Un contenedor de entrenamiento de PyTorch prediseñado, versión 1.4 o posterior.

    • Un contenedor personalizado con NCCL 2.7 o una versión posterior y el paquete google-reduction-server instalado. Puedes instalar este paquete en una imagen de contenedor personalizada añadiendo la siguiente línea a tu archivo 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
      

Entrenar con el servidor de reducción

Para usar Reduction Server, haz lo siguiente al crear un recurso de entrenamiento personalizado:

  1. Especifica uno de los siguientes URIs en el containerSpec.imageUricampo 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 la ubicación en la que vas a realizar el entrenamiento personalizado puede reducir la latencia.

  2. Al seleccionar el tipo de máquina y el número de nodos del tercer grupo de trabajadores, asegúrate de que el ancho de banda de red total del tercer grupo de trabajadores sea igual o superior al ancho de banda de red total del primer y el segundo grupo de trabajadores.

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

    No se usan GPUs en los nodos del servidor de reducción. Para obtener información sobre el ancho de banda máximo disponible de cada nodo del tercer grupo de trabajadores, consulta las columnas "Ancho de banda de salida máximo (Gbps)" de la familia de máquinas de uso general.

    Por ejemplo, si configuras el primer y el segundo grupo de trabajadores para que usen 5n1-highmem-96 nodos, cada uno con 8 NVIDIA_TESLA_V100 GPUs, cada nodo tendrá un ancho de banda máximo disponible de 100 Gbps, lo que supone un ancho de banda total de 500 Gbps. Para igualar 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, lo que da un ancho de banda total de 512 Gbps.

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

El siguiente comando muestra un ejemplo de cómo crear un recurso CustomJob que usa Reduction Server:

gcloud 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, consulta la guía para crear un CustomJob.

Prácticas recomendadas para entrenar modelos con Reduction Server

Tipo y número de máquinas

En el entrenamiento de Reduction Server, cada trabajador debe conectarse a todos los hosts de reducción. Para minimizar el número de conexiones en el host de trabajador, usa un tipo de máquina con el mayor ancho de banda de red para el host de reductor.

Una buena opción para los hosts de reductores es una máquina virtual N1 o N2 de uso general con al menos 16 vCPUs que proporcione 32 Gbps de ancho de banda de salida, como n1-highcpu-16 y n2-highcpu-16. El ancho de banda de las VMs de nivel 1 para las VMs N1 y N2 aumenta el ancho de banda de salida máximo, que oscila entre 50 y 100 Gbps, por lo que son una buena opción para los nodos de VM de reducción.

El ancho de banda de salida total de los trabajadores y los reductores debe ser el mismo. Por ejemplo, si usas 8 VMs a2-megagpu-16g como trabajadores, deberías usar al menos 25 VMs n1-highcpu-16 como reductores.

`(8 worker VMs * 100 Gbps) / 32 Gbps egress = 25 reducer VMs`.

Agrupar mensajes pequeños

Reduction Server funciona mejor si los mensajes que se van a agregar son lo suficientemente grandes. La mayoría de los frameworks de aprendizaje automático ya ofrecen técnicas con terminología diferente para agrupar tensores de gradiente pequeños antes de realizar la operación de reducción total.

Horovod

Horovod admite Tensor Fusion para agrupar tensores pequeños en lotes para la reducción total. Los tensores se rellenan en un búfer de fusión hasta que el búfer se llena por completo y se ejecuta la operación de reducción total en el búfer. Puedes ajustar el tamaño del búfer de fusión configurando la variable de entorno HOROVOD_FUSION_THRESHOLD.

El valor recomendado para la variable de entorno HOROVOD_FUSION_THRESHOLD es de al menos 128 MB. En este caso, asigna el valor 134217728 (128 * 1024 * 1024) a la variable de entorno HOROVOD_FUSION_THRESHOLD.

PyTorch

PyTorch DistributedDataParallel admite mensajes por lotes como "agrupación de gradientes". Define el parámetro bucket_cap_mb en el constructor DistributedDataParallel para controlar el tamaño de los contenedores de lotes. El tamaño predeterminado es de 25 MB.

PRÁCTICA RECOMENDADA: El valor recomendado de bucket_cap_mb es 64 (64 MB).

Variables de entorno de tu clúster

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

Cuando ejecutas el entrenamiento distribuido con TensorFlow, TF_CONFIG se analiza para crear tf.train.ClusterSpec. Del mismo modo, cuando ejecutes un entrenamiento distribuido con otros frameworks de aprendizaje automático, debes analizar CLUSTER_SPEC para rellenar las variables de entorno o los ajustes que requiera el framework.

El formato de CLUSTER_SPEC

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

Clave Descripción
"cluster"

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 transferir al constructor de tf.train.ClusterSpec.

La descripción del clúster contiene una lista de nombres de réplicas de 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 has especificado al crear el trabajo.
"workerpool2" Este grupo de trabajadores contiene servidores de parámetros, si los has especificado al crear el trabajo.
"workerpool3" Este grupo de trabajadores contiene evaluadores, si los has especificado al crear la tarea.
"environment" La cadena cloud.
"task" Describe la tarea del nodo concreto en el que se está ejecutando el código. Puedes usar esta información para 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 está ejecutando esta tarea. Por ejemplo, "workerpool0" hace referencia a la réplica principal.
"index"

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

"trial" Identificador de la prueba de ajuste de hiperparámetros que se está ejecutando. Cuando configuras el ajuste de hiperparámetros de una tarea, defines un número de pruebas de entrenamiento. Este valor te permite diferenciar en tu código las pruebas que se están ejecutando. El identificador es un valor de cadena que contiene el número de prueba, empezando por 1.
job

El CustomJobSpec que has proporcionado para crear la tarea de entrenamiento actual, representado como un diccionario.

Ejemplo: CLUSTER_SPEC

Este es 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 define la variable de entorno TF_CONFIG en cada réplica de todas las tareas de entrenamiento distribuido. Vertex AI no define TF_CONFIG para las tareas 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 lo que requiere TensorFlow.

La preparación distribuida con TensorFlow funciona de la misma forma cuando se usan contenedores personalizados que cuando se usa un contenedor precompilado.

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

TF_CONFIG campos
cluster

Descripción del clúster de TensorFlow. Un diccionario que asigna uno o varios nombres de tareas (chief, worker, ps o master) a listas de direcciones de red en las que se ejecutan estas tareas. En una tarea de entrenamiento determinada, este diccionario es el mismo en todas las VMs.

Este es un primer argumento válido para el constructor de 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, aunque los uses para tu trabajo.

task

Descripción de la tarea de la VM en la que se ha definido esta variable de entorno. En 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 de un trabajo de entrenamiento distribuido. También puedes usarlo para cambiar el comportamiento de tu código de entrenamiento en diferentes pruebas de una tarea de ajuste de hiperparámetros.

Este diccionario incluye los siguientes pares clave-valor:

task campos
type

El tipo de tarea que realiza esta VM. Este valor se define como worker en los trabajadores, ps en los servidores de parámetros y evaluator en los evaluadores. En el trabajador principal de tu trabajo, el valor se establece en chief o master.

index

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

trial

ID de la prueba de ajuste de hiperparámetros que se está ejecutando en esta VM. Este campo solo se define si la tarea de entrenamiento actual es una tarea de ajuste de hiperparámetros.

En las tareas de ajuste de hiperparámetros, Vertex AI ejecuta el código de entrenamiento repetidamente en muchas pruebas con diferentes hiperparámetros cada vez. Este campo contiene el número de prueba actual, que empieza por 1 en la primera prueba.

cloud

ID que usa Vertex AI internamente. Puedes ignorar este campo.

job

El CustomJobSpec que has proporcionado para crear la tarea de entrenamiento actual, representado como un diccionario.

environment

La cadena cloud.

Ejemplo: TF_CONFIG

El siguiente código de ejemplo 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 una tarea de ajuste de hiperparámetros que se ejecuta en la versión 2.1 o posterior del tiempo de ejecución y usa un trabajador maestro, dos trabajadores y un servidor de parámetros, este código genera el siguiente registro de uno de los trabajadores durante la primera prueba de ajuste de hiperparámetros. En el ejemplo de salida se oculta el campo job para que sea más conciso y se sustituyen algunos IDs 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 define en las tareas de entrenamiento distribuido.

Probablemente no tengas que interactuar directamente con la variable de entorno TF_CONFIG en tu código de entrenamiento. Solo debes acceder a la variable de entorno TF_CONFIG si las estrategias de distribución de TensorFlow y el flujo de trabajo estándar de ajuste de hiperparámetros de Vertex AI, que se describen en las secciones siguientes, no funcionan en tu trabajo.

Preparación distribuida

Vertex AI asigna la variable de entorno TF_CONFIG para ampliar las especificaciones que TensorFlow requiere para el entrenamiento distribuido.

Para llevar a cabo el entrenamiento distribuido con TensorFlow, usa la tf.distribute.Strategy API. En concreto, te recomendamos que uses la API Keras junto con MultiWorkerMirroredStrategy o, si especificas servidores de parámetros para tu trabajo, ParameterServerStrategy. Sin embargo, ten en cuenta que TensorFlow solo ofrece asistencia experimental para estas estrategias.

Estas estrategias de distribución usan la variable de entorno TF_CONFIG para asignar roles a cada máquina virtual de tu tarea de entrenamiento y para facilitar la comunicación entre las máquinas virtuales. No es necesario que accedas directamente a la variable de entorno TF_CONFIG en tu código de entrenamiento, ya que TensorFlow se encarga de ello.

Solo debes analizar directamente la variable de entorno TF_CONFIG si quieres personalizar el comportamiento de las diferentes VMs que ejecutan tu 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 de cada prueba. Tu código de entrenamiento no tiene por qué saber qué prueba se está ejecutando. Además, puedes monitorizar el progreso de los trabajos de ajuste de hiperparámetros en la consola de Google Cloud .

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

Siguientes pasos