Entrenar modelos de aprendizaje automático a gran escala en GKE con el punto de control multinivel

En esta página se explica cómo usar la creación de puntos de control multinivel para almacenar y gestionar de forma fiable los puntos de control durante el entrenamiento de modelos de aprendizaje automático en GKE. El almacenamiento y la gestión de puntos de control son fundamentales para los trabajos de entrenamiento a gran escala, definidos como aquellos que utilizan más de miles de nodos. Las interrupciones de estas tareas a gran escala son frecuentes (potencialmente cada hora) y la recuperación puede ser lenta.

Ventajas

Usar la creación de puntos de control multinivel ofrece las siguientes ventajas:

  • Gestión de datos de puntos de control totalmente orquestada, incluidas las copias de seguridad, la replicación y la restauración automática de las siguientes cargas de trabajo:
  • Recuperación rápida de los trabajos de entrenamiento a partir de un punto de control almacenado en el nodo local. También puedes hacer una recuperación usando los puntos de control almacenados en otro nodo del clúster de entrenamiento.
  • Restauración rápida de los trabajos de entrenamiento a partir de un punto de control almacenado en una copia de seguridad de Cloud Storage en el peor de los casos, cuando no hay puntos de control en el clúster.

Antes de empezar

Antes de empezar, asegúrate de que has realizado las siguientes tareas:

  • Habilita la API de Google Kubernetes Engine.
    • Habilitar la API de Google Kubernetes Engine
  • Si quieres usar Google Cloud CLI para esta tarea, instálala y, a continuación, inicialízala. Si ya has instalado la gcloud CLI, obtén la versión más reciente ejecutando gcloud components update.

Requisitos

Para usar la creación de puntos de control multinivel, se necesita la versión 1.32.4-gke.1415000 o una posterior del clúster de GKE.

Limitaciones

  • No se admiten clústeres de Autopilot.

Configurar nodos de GKE para usar el almacenamiento de puntos de control multinivel

En esta sección se explica cómo configurar nodos de GKE en clústeres nuevos y ya creados.

Configurar nodos en un clúster nuevo

  1. Crea un clúster con la función de creación de puntos de control multinivel, el controlador CSI de Cloud Storage FUSE y Workload Identity Federation para GKE habilitados. Si usas slices de TPU para tu carga de trabajo de aprendizaje automático, tendrás que ajustar el comando de creación del clúster para incluir la configuración de un grupo de nodos de slice de TPU.

    gcloud container clusters create CLUSTER_NAME \
    --addons=HighScaleCheckpointing,GcsFuseCsiDriver  \
    --node-locations=NODE_LOCATION \
    --workload-pool=PROJECT_ID.svc.id.goog \
    --cluster-version=CLUSTER_VERSION
    --location=CLUSTER_LOCATION \
    --machine-type=MACHINE_TYPE \
    --num-nodes=NUM_NODES

    Sustituye los siguientes valores:

    • CLUSTER_NAME: el nombre del clúster.
    • NODE_LOCATION: la zona de los nodos del clúster. Aquí es donde se encuentra tu capacidad de TPU.
    • PROJECT_ID: tu Google Cloud ID de proyecto.
    • CLUSTER_VERSION: la versión de tu clúster. La versión mínima admitida es la 1.32.4-gke.1415000.
    • CLUSTER_LOCATION: la región en la que quieras crear el clúster.
    • MACHINE_TYPE: el tipo de máquina que se usa en los nodos que ejecutan componentes como el controlador JobSet y el controlador de creación de puntos de control multinivel. Para el entrenamiento a gran escala, recomendamos usar al menos e2-standard-4 máquinas. No usarás este tipo de máquina para entrenar modelos. En su lugar, crearás grupos de nodos independientes para ese fin, que a menudo utilizarán familias de VMs optimizadas para aceleradores.
    • NUM_NODES: el número de nodos que se crearán en cada una de las zonas del clúster.

Configurar nodos en un clúster

Para usar la creación de puntos de control multinivel con un clúster, habilítala junto con el controlador CSI de Cloud Storage FUSE y Workload Identity Federation for GKE con los siguientes comandos. La versión de tu clúster debe ser posterior a la 1.32.3-gke.1170000.

  1. Habilita Workload Identity Federation para GKE:

    gcloud container clusters update CLUSTER_NAME \
    --workload-pool=PROJECT_ID.svc.id.goog \
    --location=CLUSTER_LOCATION

    Sustituye los siguientes valores:

    • CLUSTER_NAME: el nombre del clúster.
    • PROJECT_ID: tu Google Cloud ID de proyecto.
    • CLUSTER_LOCATION: la región del clúster.

  2. Habilita la creación de puntos de control de varios niveles y el controlador de CSI de FUSE de Cloud Storage:

    gcloud container clusters update CLUSTER_NAME \
    --update-addons=HighScaleCheckpointing=ENABLED,GcsFuseCsiDriver=ENABLED \
    --location=CLUSTER_LOCATION

Configurar permisos para usar la creación de puntos de control multinivel

En esta sección se explica cómo configurar los permisos para usar la creación de puntos de control de varios niveles.

Conceder acceso a segmentos de Cloud Storage

Los volúmenes efímeros que usa el controlador de CSI de Multi-Tier Checkpointing deben usar segmentos de Cloud Storage.

Para almacenar puntos de control en un segmento de Cloud Storage, la función de creación de puntos de control multinivel necesita acceso al segmento. Asigna el rol de IAM Usuario de objetos de Storage (roles/storage.objectUser) en el bucket a la cuenta de servicio de Kubernetes para la creación de puntos de control multinivel.

gcloud storage buckets add-iam-policy-binding gs://GCS_BUCKET \
    --member "principal://iam.googleapis.com/projects/PROJECT_NUMBER/locations/global/workloadIdentityPools/PROJECT_ID.svc.id.goog/subject/ns/gke-managed-checkpointing/sa/gke-checkpointing-multitier-node" \
    --role "roles/storage.objectUser"

Sustituye los siguientes valores:

(Opcional) Conceder acceso a la cuenta de servicio predeterminada de Compute Engine

Si tus instancias de Compute Engine necesitan acceso de lectura al segmento de Cloud Storage, asigna el rol de IAM de lector de objetos de Storage (roles/storage.objectViewer) a la cuenta de servicio predeterminada de Compute Engine.

gcloud storage buckets add-iam-policy-binding gs://GCS_BUCKET \
    --member serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com \
    --role roles/storage.objectViewer

Despliega el controlador JobSet

El controlador JobSet se encarga de gestionar las tareas por lotes que ejecutan el entrenamiento de tu modelo en GKE, y su asignación de recursos se ajusta para gestionar la carga de trabajo de forma eficiente. Asegúrate de que el lanzador de tareas de entrenamiento implementa y usa JobSet.

Para aumentar la solicitud de memoria a 1 Gi, el límite de memoria a 2 Gi y la solicitud de CPU a 1 para el contenedor de gestor de tu implementación de JobSet, ejecuta el siguiente comando de parche:

kubectl patch -n jobset-system deploy jobset-controller-manager --type json \
    --patch '[{"op": "add", "path": "/spec/template/spec/containers/0/resources", "value": {"limits": {"memory": "2Gi"}, "requests": {"cpu": "1", "memory": "1Gi"}}}]'

Inicializar el controlador de CSI de Multi-Tier Checkpointing

En esta sección se describe cómo inicializar el controlador CSI de Multi-Tier Checkpointing en los nodos en los que se ejecutarán sus cargas de trabajo. El controlador CSI se encarga de gestionar el almacenamiento y la gestión de los puntos de control durante el proceso de entrenamiento del modelo.

Crear un CheckpointConfiguration

CheckpointConfiguration es un recurso personalizado de Kubernetes que especifica las propiedades para desplegar el controlador de CSI de creación de puntos de control multinivel. Este recurso tiene ámbito de clúster.

  1. Crea el siguiente archivo de manifiesto checkpoint.yaml.

    apiVersion: checkpointing.gke.io/v1
kind: CheckpointConfiguration
metadata:
  name: MTC_CONFIG_NAME-configuration
spec:
    cloudStorageBucketName: GCS_BUCKET
    nodeSelector:
        node.kubernetes.io/instance-type: MACHINE_TYPE
    tolerations:
    - key: TOLERATION_KEY
        operator: Exists
        effect: NoSchedule
    inMemoryVolumeSize: IN_MEMORY_VOLUME_SIZE
    gcsFuseMountOptions:
    - implicit-dirs
    - metadata-cache:negative-ttl-secs:0
    - metadata-cache:ttl-secs:-1
    - metadata-cache:stat-cache-max-size-mb:-1
    - metadata-cache:type-cache-max-size-mb:-1
    - file-system:kernel-list-cache-ttl-secs:0
    - read_ahead_kb=1024
    - write:enable-streaming-writes:true

    Haz los cambios siguientes:

    • MTC_CONFIG_NAME: el nombre de tu CheckpointConfiguration. Este nombre es global para el clúster y no es específico de un trabajo.
    • GCS_BUCKET: el nombre del segmento de Cloud Storage en el que almacenarás los datos de puntos de control. Usa el segmento que has configurado en el paso Configurar un segmento de Cloud Storage con permisos.

    • MACHINE_TYPE: el tipo de máquina de los aceleradores correspondientes. El valor puede ser uno de los siguientes:

      Para obtener más información sobre cómo ejecutar cargas de trabajo distribuidas en GPUs con GKE, consulta Ejecutar GPUs con varias instancias. En el caso de las TPU, consulta Crear el grupo de nodos de la porción de TPU.

    • TOLERATION_KEY: este campo permite que el controlador de CSI se programe en nodos con intolerancias coincidentes. Para obtener más información sobre cómo funcionan los taints en los distintos tipos de aceleradores, consulta estas páginas:

    • IN_MEMORY_VOLUME_SIZE: tamaño de la caché de creación de puntos de control en memoria. Especifica la cantidad y la unidad (por ejemplo, 200 Gi).Este valor debe ser:

      • El tamaño del punto de control local de las TPUs multiplicado por 2,2.
      • El tamaño del punto de control local de las GPUs con un solo elemento del mismo nivel multiplicado por 6,6.

  2. Aplica el archivo de manifiesto:

    kubectl apply -f checkpoint.yaml

  3. Comprueba que el controlador CSI se esté ejecutando:

    kubectl get pod -n gke-managed-checkpointing

    La salida debería ser similar a la siguiente. Habrá varias entradas, una por cada nodo acelerado.

    NAME                                                          READY   STATUS    RESTARTS   AGE
multitier-driver-e2b033a7-a4e7-496a-87a3-ffd7fcc2e57b-2d4fz   5/5     Running   0          114s

Desinstalar el controlador de CSI de Multi-Tier Checkpointing

Si quieres dejar de implementar el controlador CSI de Multi-Tier Checkpointing, elimina el recurso CheckpointConfiguration. El controlador de Multi-Tier Checkpointing elimina el controlador de CSI de los nodos. De esta forma, se eliminan los discos de RAM y se libera memoria para otras cargas de trabajo. Por ejemplo:

kubectl delete -f checkpoint.yaml

Gestionar la conservación de datos y la recogida de elementos no utilizados de las copias de seguridad de Cloud Storage

Eres responsable de implementar políticas de retención para las copias de seguridad de los puntos de control de Cloud Storage. La creación de puntos de control de varios niveles solo escribe copias de seguridad de puntos de control en Cloud Storage y nunca las modifica ni las elimina.

Muchas herramientas de código abierto pueden gestionar la conservación y la recogida de elementos no utilizados, entre las que se incluyen las siguientes:

En el siguiente ejemplo se usa backup-warden, donde el directorio backup se monta en una ubicación de copia de seguridad que usa Cloud Storage FUSE:

# Add --delete option to actually delete the backups, as is it only shows what would be deleted (dry-run)
backup-warden -p backup \
    --hourly 24 \
    --daily 7 \
    --weekly 5 \
    --monthly always \
    --yearly always \
    --prefer-recent

Actualizar el manifiesto de JobSet de la carga de trabajo

Actualiza el archivo de manifiesto de JobSet de tu trabajo para incluir el volumen de puntos de control a gran escala. Los detalles dependen de tu carga de trabajo.

Por ejemplo, para ampliar el JobSet de ejemplo de Desplegar multislices de TPU en GKE, sigue estos pasos:

  1. Añade las siguientes líneas al contenedor jax-tpu.

    volumeMounts:
- name: checkpoint
  mountPath: CHECKPOINT_DIR

    Sustituye CHECKPOINT_DIR por la ruta al directorio de puntos de control. Esta es la ubicación donde se genera el replicator.yaml y donde los puntos de control de varios niveles realizan la operación de guardado. Para obtener más información, consulta Integrar la creación de puntos de control de varios niveles en tu aplicación.

  2. Añade las siguientes líneas al campo spec.template.spec de la especificación del trabajo.

    volumes:
- name: checkpoint
  csi:
    driver: multitier-checkpoint.csi.storage.gke.io

Integrar la creación de puntos de control de varios niveles en tu aplicación

Para compartir información sobre las ubicaciones de los puntos de control y la disponibilidad de la réplica, modifica tu aplicación para que use el siguiente protocolo para comunicarse con la creación de puntos de control multinivel.

Startup

En esta sección se describen los pasos iniciales que debe seguir la aplicación para interactuar con la creación de puntos de control multinivel.

El replicador es un componente principal de la creación de puntos de control multinivel que se ejecuta en todos los nodos como parte del controlador CSI. Replicator gestiona la replicación de puntos de control en los niveles de almacenamiento, desde el disco RAM local hasta los nodos de igual y el almacenamiento externo, como Cloud Storage.

El archivo replicator.yaml actúa como un plano de control dinámico entre tu trabajo de entrenamiento de aprendizaje automático (código del framework) y el componente Replicator. Tu aplicación de aprendizaje automático genera este archivo de forma programática en el volumen local (RAMDisk), al que pueden acceder tanto el trabajo de entrenamiento como el servicio Replicator. Este manifiesto permite que el framework de AA proporcione instrucciones de configuración en tiempo de ejecución y de gestión del ciclo de vida al replicador, que son distintas de los parámetros de infraestructura estáticos (por ejemplo, la frecuencia de subida de Cloud Storage) definidos durante la configuración del backend.

Para ver un ejemplo concreto de esta interacción, consulta lo siguiente:

Tu aplicación debe seguir estos pasos durante el inicio:

  1. Espera hasta que el archivo replicator.yaml no esté presente, lo que indica que Replicator está listo para que tu aplicación lo configure. El archivo replicator.yaml se genera en la ubicación CHECKPOINT_DIR que has configurado en la sección Actualizar el manifiesto de JobSet de la carga de trabajo.

    Cuando se crea el trabajo de entrenamiento del modelo, el archivo replicator.yaml no existe y tu aplicación puede continuar inmediatamente. Sin embargo, si el trabajo se ha reiniciado (por ejemplo, debido a un fallo o a una intervención manual), es posible que el sistema siga procesando la instancia de trabajo anterior y que el replicator.yaml de esa instancia siga presente en el volumen local.

  2. Tu aplicación o trabajo de aprendizaje automático crea el archivo replicator.yaml con una configuración similar a la siguiente.

    Orbax

    job-name: orbax
framework: orbax
assume-data-parallelism: 3
node-rank: 0
nodes: 32
peer-ranks: [1, 16] or peers-per-node: 2
backup-interval-minutes: 30

    PyTorch

    job-name: nemo
framework: pytorch.distributed
node-rank: 0
nodes: 32
peer-ranks: [1, 16] or peers-per-node: 2
backup-interval-minutes: 30

    Esta configuración de ejemplo tiene los siguientes campos:

    • name: el nombre de la tarea de entrenamiento.
    • framework: el framework de aprendizaje automático que usa la tarea de entrenamiento.
    • node-rank: identificador único del nodo actual en el trabajo de entrenamiento distribuido. Representa el rango del nodo que crea este archivo. Cada nodo que participe en la ejecución tendrá su propio rango.
    • nodes: el número total de nodos que participan en el trabajo de entrenamiento distribuido. Este valor procede de los metadatos del pod. El trabajo de entrenamiento de AA también puede ver este valor.
    • peer-ranks o peers-per-node: dos formas alternativas de especificar la topología de replicación. Solo debe incluirse uno de estos dos parámetros.
      • peer-ranks: rangos explícitos de los nodos del mismo nivel en los que se deben replicar los datos del punto de control del nodo actual. De esta forma, se puede controlar con precisión qué nodos específicos actúan como partners de replicación.
      • peers-per-node: número de nodos de igual a igual por nodo que el replicador debe seleccionar automáticamente para la replicación.
    • backup-interval-minutes: la frecuencia, en minutos, con la que se crean copias de seguridad de los puntos de control en Cloud Storage. Te recomendamos que asignes a este valor 30 minutos o más.

  3. Espera a que el sistema elimine el archivo replicator.yaml. Esto indica que el replicador se ha reiniciado y ha realizado una limpieza. Este paso te permite evitar archivos obsoletos o temporales en el volumen local cuando tu aplicación realice los pasos de la siguiente sección.

Restaurar desde el último punto de control correcto

  1. Una vez inicializado el replicador, la creación de puntos de control multinivel crea un enlace simbólico por trabajador de TPU o GPU. Estos enlaces simbólicos se crean en el mismo volumen local montado que el archivo replicator.yaml, donde la tarea guarda los puntos de control.

    Los enlaces simbólicos tienen el formato <job-name>-s{step}-n<node-rank>-w<worker-index>.restore.

  2. Restaura cada trabajador desde el archivo .restore correspondiente. Por ejemplo, consulta el ejemplo de gestor de puntos de control replicados de Orbax en la siguiente sección.

Guarda el punto de control.

Tu aplicación realiza estos pasos varias veces mientras avanza el trabajo de entrenamiento. Las operaciones de guardado se realizan en la ubicación CHECKPOINT_DIR que has configurado en Actualizar el manifiesto de JobSet de la carga de trabajo.

Orbax

Crea puntos de control de Orbax. El directorio se denomina con el número del paso. El replicador detecta el directorio de puntos de control recién creado, realiza la replicación o la copia de seguridad según sea necesario y se limpia automáticamente.

Para obtener más información sobre cómo usar el gestor de puntos de control de réplicas de Orbax, consulta el archivo MaxtTest checkpointing. Para ver un ejemplo de interacción del servicio de replicación, consulta el archivo MaxText max_utils.

PyTorch

Usa InClusterLocalCheckpointIO como pytorch_lightning.CheckpointIO personalizado para habilitar la creación de puntos de control distribuidos correcta con el almacenamiento local. El siguiente comando de ejemplo habilita la creación de puntos de control multinivel mediante una implementación de referencia basada en el framework NVIDIA NeMo:

torchrun train.py <other_train_flags> \
    --local-ckpt-dir=CHECKPOINT_DIR \
    --local-ckpt-interval=20 \
    --job-name=JOB_NAME \
    --enable-high-scale-ckpt

Haz los cambios siguientes:

  • CHECKPOINT_DIR: la ruta al directorio de puntos de control.
  • JOB_NAME: el nombre de la carga de trabajo de entrenamiento.

Actualizaciones del clúster

En el caso de las actualizaciones de clústeres, puedes eliminar y volver a crear tu objeto CheckpointConfiguration antes o después de la actualización. Esta acción es obligatoria porque el daemonset del controlador de puntos de control de nodos, implementado dinámicamente por este objeto, no se actualizará automáticamente.

Si quieres mantener la especificación de DaemonSet, no tienes que hacer nada.

Solucionar problemas

En esta sección se ofrecen instrucciones para solucionar problemas con la creación de puntos de control de varios niveles. Para obtener información general sobre cómo solucionar problemas de almacenamiento, consulta Solucionar problemas de Cloud Storage en GKE.

No se ha habilitado la creación de puntos de control de varios niveles

El siguiente error indica que la creación de puntos de control multinivel no está habilitada en tu clúster:

error: unable to recognize "checkpoint.yaml": no matches for kind "CheckpointConfiguration" in version "checkpointing.gke.io/v1"

Es posible que se produzca este error después de ejecutar kubectl apply -f checkpoint.yaml en el paso Crear un CheckpointConfiguration.

Para solucionar este problema, comprueba si has habilitado la creación de puntos de control multinivel en tu clúster con el siguiente comando:

gcloud container clusters describe CLUSTER_NAME \
    --project PROJECT_ID
    --location CLUSTER_LOCATION

Si la creación de puntos de control de varios niveles está habilitada, la salida debería ser similar a la siguiente:

addonsConfig:
  gcePersistentDiskCsiDriverConfig:
    enabled: true
  gcsFuseCsiDriverConfig:
    enabled: true
  highScaleCheckpointingConfig:
    enabled: true
  kubernetesDashboard:
    disabled: true
  networkPolicyConfig:
    disabled: true

Si la creación de puntos de control de varios niveles no está habilitada, actualiza tu clúster para habilitarla.

El controlador de CSI de comprobación de puntos de control multinivel no puede montar volúmenes

Puede que te encuentres con este problema si el controlador CSI no puede montar el volumen de Cloud Storage. Puede haber varias líneas similares a esta.

kubectl get pod -n gke-managed-checkpointing
NAME                                                          READY   STATUS     RESTARTS   AGE
multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9   0/5     Init:0/1   0          6m32s

Para solucionar este problema, comprueba los eventos de Pod del controlador CSI, como se muestra en el siguiente ejemplo:

kubectl describe pod multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9 -n gke-managed-checkpointing

Events:
  Type     Reason       Age                 From               Message
  ----     ------       ----                ----               -------
  Normal   Scheduled    17m                 default-scheduler  Successfully assigned gke-managed-checkpointing/multitier-driver-14694e4d-774f-4104-8bba-f0bd82fd7557-5vxr9 to gke-my-cluster-default-pool-353c773f-6d8q
  Warning  FailedMount  82s (x16 over 17m)  kubelet            MountVolume.SetUp failed for volume "gcs" : rpc error: code = PermissionDenied desc = failed to get GCS bucket "checkpointing-test-bucket": googleapi: Error 403: Caller does not have storage.objects.list access to the Google Cloud Storage bucket. Permission 'storage.objects.list' denied on resource (or it may not exist)., forbidden

Si el problema se debe a un error PermissionDenied del bucket de Cloud Storage, como se muestra en el ejemplo, puede solucionarlo configurando correctamente los permisos.

Siguientes pasos