Entrenamiento resiliente con Rutas

En esta página, se explica cómo configurar funciones de resiliencia para las cargas de trabajo de Pathways en Google Kubernetes Engine (GKE) de modo que los trabajos de entrenamiento a gran escala puedan recuperarse de interrupciones planificadas y fallas de hardware no planificadas.

Pathways proporciona beneficios de resiliencia de las siguientes maneras:

  • Suspensión y reanudación: Permite que tu carga de trabajo tolere interrupciones planificadas, como avisos de interrupción, sin necesidad de un código personalizado de control de interrupciones.
  • Entrenamiento elástico: Permite que tu carga de trabajo tolere fallas de hardware no planificadas sin que falle el cliente con la lógica de recuperación en tu código de entrenamiento.
  • Puntos de control de varios niveles: Permite que tu carga de trabajo guarde puntos de control locales y replicados por pares de alta frecuencia en los trabajadores TPU con Python ubicado en el mismo lugar, con copias de seguridad asíncronas de Cloud Storage para reducir el tiempo de recuperación.
  • Intercambio en caliente: Permite que una carga de trabajo de JobSet de GKE de mayor prioridad asuma rápidamente el control de las porciones de TPU de repuesto o de menor prioridad cuando se produce una falla de hardware.

Antes de comenzar

Antes de comenzar, completa las siguientes tareas:

Si usas el entrenamiento elástico, instala los siguientes paquetes en la imagen de contenedor de tu carga de trabajo:

  • pathwaysutils
  • Orbax (orbax-checkpoint): >= 0.12.6 (obligatorio para las instantáneas en memoria con ocp.v1.training.pathways.Snapshotter)
  • Grain (grain): Es obligatorio si usas grain.experimental.ElasticIterator.

Si usas la creación de puntos de control de varios niveles, también completa los siguientes requisitos:

  • Sigue las instrucciones de Entrena modelos de aprendizaje automático a gran escala en GKE con Multi-Tier Checkpointing para configurar tu clúster de GKE para Multi-Tier Checkpointing:
    • Habilita los complementos de Multi-Tier Checkpointing y Cloud Storage FUSE (GcsFuseCsiDriver), junto con Workload Identity Federation for GKE, en tu clúster.
    • Crea un bucket de Cloud Storage con el espacio de nombres jerárquico habilitado.
    • Otorga los roles de Identity and Access Management (IAM) roles/storage.objectUser y roles/storage.insightsCollectorService a la cuenta de servicio de Kubernetes de Multi-Tier Checkpointing.
    • Implementa un recurso personalizado CheckpointConfiguration que especifique tu bucket de Cloud Storage, el selector de nodo TPU y el tamaño del volumen en memoria (inMemoryVolumeSize).
  • Compila y envía una imagen de contenedor sidecar de Python coubicada, por ejemplo, con el Dockerfile de sidecar de muestra. En las imágenes de tu carga de trabajo y de sidecar, usa versiones de jax y jaxlib que coincidan y que se alineen con tu versión de Pathways y cumplan con los siguientes requisitos de versión mínima:
    • Versión de Orbax: >= 0.12.4
    • Versión de JAX (incluido Python ubicado en el mismo lugar): >= 0.10.1

Suspensión y reanudación

Por lo general, GKE envía un aviso de interrupción a un Pod acelerador antes de interrumpirlo. Pathways habilita la tolerancia a la interrupción de forma predeterminada en todas las implementaciones de Cloud, y los trabajos del acelerador de Pathways escuchan estos avisos.

Cuando llega un aviso de interrupción, Pathways primero determina si puede guardar y restablecer la carga de trabajo de forma transparente. Si es posible, Pathways intenta suspender de forma transparente tu carga de trabajo de AA escribiendo su estado actual en el almacenamiento persistente, como Cloud Storage, antes de que GKE desaloje tus trabajos de acelerador. Cuando GKE vuelve a programar tus trabajos más adelante, Pathways reanuda tu carga de trabajo de AA leyendo su estado persistente.

Si la carga de trabajo no se puede restablecer, Pathways cierra el trabajo del acelerador y reenvía la falla a tu carga de trabajo si se configura el entrenamiento elástico. Si no se configura el entrenamiento elástico, GKE reinicia toda la carga de trabajo según la política de reinicio de JobSet.

Las cargas de trabajo de AA típicas definidas con JAX dependen de componentes de XLA de Pathways sin estado, que se pueden restablecer con una instantánea de memoria de alto ancho de banda (HBM). Ciertas cargas de trabajo de AA, como las que se definen con la API de Python de JAX colocada, dependen de componentes de Pathways con estado y no se pueden restablecer.

Entrenamiento elástico

El entrenamiento elástico permite que tu trabajo de entrenamiento continúe cuando se producen fallas de hardware o interrupciones de porciones, sin reiniciar toda la carga de trabajo. En un trabajo de entrenamiento de varias porciones, Pathways mantiene activos el proceso del cliente de un solo controlador y las porciones de TPU en buen estado cuando una porción individual deja de estar disponible.

El entrenamiento elástico combina las capacidades del sistema de Pathways con la lógica de recuperación en tu código de entrenamiento:

  • Detección de fallas: Cuando se produce una falla de hardware o una interrupción en una porción de TPU, Pathways anula la asignación de la porción afectada y genera un jax.errors.JaxRuntimeError la próxima vez que tu carga de trabajo acceda a los datos de esa porción. Debido a que el proceso del cliente permanece activo, tu código puede detectar este error, volver a configurar tus recursos y continuar con el entrenamiento.
  • Modos de recuperación: La clase pathwaysutils.elastic.manager.Manager y su decorador elastic_retry admiten dos estrategias de recuperación:
    • Pausa y reanudación: Si configuras minimum_slice_count como None o como el recuento total de porciones, el entrenamiento se detiene cuando falla una porción y se reanuda después de que todas las porciones vuelven a estar en buen estado. Debido a que el recuento de porciones no cambia entre los intentos, tu carga de trabajo puede reutilizar la malla de dispositivos existente y los ejecutables compilados.
    • Cambio de tamaño de la réplica: Si configuras minimum_slice_count como un valor menor que el recuento total de porciones, el entrenamiento continúa en las porciones activas restantes cuando una porción falla, ya que se vuelve a compilar la malla del dispositivo y se vuelve a compilar para la topología más pequeña. Un subproceso de supervisión en segundo plano verifica periódicamente los segmentos inactivos y registra los segmentos que vuelven a estar en buen estado en elastic_manager.available_inactive_slices. Luego, tu bucle de entrenamiento puede reconstruir la malla de dispositivos con los sectores de reemplazo, reinicializar para la topología expandida y colocar el estado actual en la malla nueva.
  • Recuperación de estado: Para reanudar el entrenamiento después de un evento elástico, tu carga de trabajo restablece los parámetros del modelo, el estado del optimizador y el estado del iterador de datos con instantáneas en la memoria. La clase Snapshotter en orbax.checkpoint copia de forma asíncrona el estado replicado de la HBM de la TPU a la memoria de la CPU de pinned_host en cada porción. Durante la recuperación, Snapshotter reconstruye el estado a partir de las réplicas de host sobrevivientes y lo vuelve a fragmentar directamente en la malla de TPU activa sin leer desde el almacenamiento remoto.

Habilita la elasticidad en el servidor proxy de Pathways

Para habilitar los conjuntos de dispositivos elásticos en Pathways en GKE, agrega la marca --num_elastic_instances al campo args del contenedor pathways-proxy en tu manifiesto de JobSet y especifica la cantidad de porciones de TPU elásticas en tu carga de trabajo:

- name: pathways-proxy
  image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/proxy_server:latest
  args:
  - --resource_manager_address=$(PATHWAYS_HEAD):29001
  - --server_port=29000
  - --gcs_scratch_location=gs://BUCKET_NAME
  - --num_elastic_instances=NUM_SLICES

Reemplaza lo siguiente:

  • BUCKET_NAME: Es el bucket de Cloud Storage que se usa para los archivos temporales de Pathways.
  • NUM_SLICES: Es la cantidad de porciones de TPU en tu carga de trabajo.

Filtra los dispositivos activos y compila la malla de dispositivos

En una carga de trabajo de Pathways, jax.devices() devuelve todos los dispositivos de la asignación inicial durante el ciclo de vida del proceso del cliente. Cuando uses el entrenamiento elástico, usa elastic_manager.active_devices para que la malla de dispositivos y el tamaño del lote solo usen porciones en buen estado y activas:

import jax
import numpy as np
from pathwaysutils.elastic import manager

elastic_manager = manager.Manager()


def create_mesh(
    elastic_manager: manager.Manager, devices_per_slice: int
) -> jax.sharding.Mesh:
  device_array = np.asarray(elastic_manager.active_devices).reshape(
      (elastic_manager.active_slice_count, devices_per_slice)
  )
  return jax.sharding.Mesh(device_array, axis_names=("replica", "model"))

Cuando cambia la cantidad de porciones activas en el modo de cambio de tamaño de la réplica, usa un cargador de datos elástico, como grain.experimental.ElasticIterator. Los iteradores de datos estándar por host guardan el estado por host y no se pueden restablecer en una cantidad diferente de hosts. En cambio, ElasticIterator hace un seguimiento de una sola posición del conjunto de datos global que puedes restablecer en cualquier cantidad de porciones activas.

Bucle de entrenamiento elástico

Un bucle de entrenamiento elástico coordina la supervisión de porciones, la recuperación ante fallas y la creación de instantáneas en la memoria, como se muestra en el siguiente ejemplo de código:

  • Tu función train crea una sola instancia de manager.Manager() y un Snapshotter de Orbax fuera del wrapper de reintentos para que el seguimiento de porciones y las instantáneas de la memoria del host persistan entre los intentos.
  • Aplicar @elastic_manager.elastic_retry a run_training_loop configura la función para que espere porciones en buen estado, establece jax.default_device en una porción activa, mantiene activo el subproceso del monitor en segundo plano y vuelve a intentarlo en un jax.errors.JaxRuntimeError.
  • Al comienzo de un intento de recuperación, tu código debe volver a compilar la malla del dispositivo desde elastic_manager.active_devices y llamar a snapshotter.load para restablecer la instantánea más reciente en las fragmentaciones de malla actualizadas. Permite que los búferes del dispositivo del intento fallido salgan del alcance para que JAX pueda liberar su memoria.
  • Durante cada paso del bucle, tu código llama periódicamente a snapshotter.save para crear una copia de seguridad del estado del entrenamiento y verifica elastic_manager.available_inactive_slices para detectar porciones de reemplazo en buen estado. Cuando hay porciones de reemplazo disponibles, tu código las agrega a elastic_manager.active_slice_indices, vuelve a compilar la malla, reinicializa el bucle de entrenamiento y coloca el estado en la malla expandida.

En el siguiente ejemplo, se muestra cómo implementar un bucle de entrenamiento elástico combinando manager.Manager, elastic_retry y Snapshotter en memoria de Orbax:

import jax
import orbax.checkpoint as ocp
import pathwaysutils
from pathwaysutils.elastic import manager

pathwaysutils.initialize()


def train(
    final_step: int,
    snapshot_interval: int,  # in steps
    devices_per_slice: int,
):
  elastic_manager = manager.Manager()
  snapshotter = ocp.v1.training.pathways.Snapshotter(replica_axis_index=0)

  @elastic_manager.elastic_retry(
      minimum_slice_count=1,
      poll_interval=10,  # in seconds
      timeout=300,  # in seconds
      retry_policy=manager.ElasticRetryLimit(max_attempts=10),
  )
  def run_training_loop():
    mesh = create_mesh(elastic_manager, devices_per_slice)
    state, abstract_state, jitted_train_step, data_iterator = (
        initialize_training_loop(mesh)
    )
    step = 0

    # Restore from the latest in-memory snapshot after a slice failure.
    if snapshotter.latest is not None:
      snapshotter.join()
      restored = snapshotter.load({
          "state": abstract_state,
          "step": step,
          "iterator_state": data_iterator.get_state(),
      })
      state = restored["state"]
      step = int(restored["step"]) + 1
      data_iterator.set_state(restored["iterator_state"])

    while step < final_step:
      batch = next(data_iterator)
      state = jitted_train_step(state, batch)

      if step % snapshot_interval == 0:
        snapshotter.save(
            step,
            {
                "state": state,
                "step": step,
                "iterator_state": data_iterator.get_state(),
            },
        )

      # If replacement slices are healthy, reinitialize and place state onto the new mesh.
      if elastic_manager.available_inactive_slices:
        elastic_manager.active_slice_indices |= (
            elastic_manager.available_inactive_slices
        )
        mesh = create_mesh(elastic_manager, devices_per_slice)
        iterator_state = data_iterator.get_state()
        _, abstract_state, jitted_train_step, data_iterator = (
            initialize_training_loop(mesh)
        )
        data_iterator.set_state(iterator_state)
        state_shardings = jax.tree.map(lambda x: x.sharding, abstract_state)
        state = jax.device_put(state, state_shardings)

      step += 1

    return state

  state = run_training_loop()
  snapshotter.join()
  return state

Configura el administrador elástico y el comportamiento de reintento

Crea una instancia de manager.Manager(devices=None) una vez por proceso. Si omites devices, el administrador usa jax.devices() y agrupa los dispositivos por su atributo slice_index.

La instancia Manager expone las siguientes propiedades para inspeccionar la topología del clúster durante el entrenamiento:

  • slice_to_devices: Es una asignación de cada slice_index a la secuencia de objetos jax.Device que pertenecen a ese segmento.
  • all_slice_indices: Es un frozenset de todos los índices de porciones de la asignación.
  • active_slice_indices: Es un frozenset de índices de porciones que participan en la malla de procesamiento activa.
  • active_devices: Es una lista de objetos jax.Device que pertenecen a los segmentos activos, ordenados por slice_index.
  • inactive_slice_indices: Es un frozenset de índices de porciones en la asignación que no participan en la malla de procesamiento.
  • available_inactive_slices: Es un frozenset de índices de porciones inactivas que el subproceso del monitor en segundo plano detectó como en buen estado y listo para volver a unirse.
  • total_slice_count y active_slice_count: Es la cantidad total de porciones asignadas y la cantidad de porciones activas.
  • default_device: Es el dispositivo 0 de una porción activa, que elastic_retry establece automáticamente como el dispositivo JAX predeterminado para cada intento.

Puedes personalizar el comportamiento de reintento y escalamiento pasando los siguientes argumentos a @elastic_manager.elastic_retry:

  • minimum_slice_count: Es la cantidad mínima de porciones en buen estado que se requieren antes de comenzar la ejecución inicial o un intento de reintento. Si se establece en None (el valor predeterminado), elastic_retry espera a que se completen todas las porciones de la asignación (total_slice_count) y opera en modo de pausa y reanudación. Si se establece en un número entero más pequeño (como 1), elastic_retry opera en el modo de cambio de tamaño de la réplica.
  • poll_interval: Es el intervalo en segundos entre las verificaciones de estado de las porciones que realiza el administrador de entrenamiento elástico (Manager) durante el entrenamiento. La configuración predeterminada es 10.
  • timeout: Es la cantidad máxima de segundos que se debe esperar para que las porciones de minimum_slice_count se activen antes de generar un TimeoutError. La configuración predeterminada es None (sin tiempo de espera).
  • pre_callback: Es una devolución de llamada opcional que se ejecuta al comienzo de cada intento después de que el administrador actualiza active_slice_indices y antes de que se ejecute la función de bucle de entrenamiento unida.
  • on_elastic_event_callback: Es una devolución de llamada opcional que se ejecuta después de que elastic_retry detecta un jax.errors.JaxRuntimeError de reducción de porciones, antes de que comience el siguiente intento de reintento. Puedes usar esta devolución de llamada para limpiar los puntos de control incompletos o restablecer los iteradores de datos.
  • retry_policy: Es un (attempt: int, error: Exception) -> bool que se puede llamar y que determina si se debe seguir reintentando después de un evento elástico. Usa manager.ElasticRetryLimit(max_attempts) para limitar la cantidad total de intentos. Si se omite, elastic_retry reintenta de forma indefinida.

En un evento elástico, JAX libera los búferes del dispositivo del intento fallido una vez que se desenrolla el marco de pila y ninguna variable hace referencia a esos arrays. Evita mantener referencias a arrays de dispositivos de un intento fallido dentro de variables globales o devoluciones de llamada de larga duración para que se libere memoria antes de que comience el siguiente intento.

Instantáneas en memoria con Orbax

La clase ocp.v1.training.pathways.Snapshotter (disponible en la versión 0.12.6 de orbax-checkpoint y versiones posteriores) mantiene una copia de seguridad en la memoria de tu estado de entrenamiento en la memoria de pinned_host en todas las porciones:

  • Snapshotter(replica_axis_index=0): Inicializa el creador de instantáneas y especifica qué eje de malla representa las réplicas de porciones paralelas de datos. Todos los arrays de JAX en el estado guardado se deben replicar a lo largo de este eje de malla para que las porciones sobrevivientes conserven una copia completa del estado del modelo y del optimizador cuando falla una porción. Si tu estado contiene arrays de un solo dispositivo, replícalos en la malla antes de llamar a save.
  • save(step, state): Transfiere de forma asíncrona state a la memoria de pinned_host y lo divide en árboles por réplica a lo largo de replica_axis_index. Si una transferencia de instantánea anterior todavía está en curso, save omite la solicitud nueva para evitar que se detenga el bucle de entrenamiento.
  • load(abstract_state, reset_snapshot_state=True): Sondea las réplicas de host guardadas para identificar los sectores sobrevivientes, concatena las réplicas en buen estado y vuelve a fragmentar el estado en las fragmentaciones de destino especificadas por abstract_state.
  • join(): Bloquea hasta que se complete cualquier operación de instantánea en cola.
  • latest: Devuelve los metadatos (incluido step) de la instantánea completada más recientemente, o None si no se completó ninguna instantánea.

Reduce la fragmentación en caso de falla de la porción

Cuando una porción de TPU falla o se interrumpe durante la ejecución, Pathways genera una jax.errors.JaxRuntimeError (como DATA_LOSS, DEADLINE_EXCEEDED, NOT_FOUND o INTERNAL). En una función decorada con @elastic_manager.elastic_retry, elastic_retry detecta esta excepción y verifica elastic.is_error_due_to_slice_down(error):

  • Si el error no se debe a un evento de reducción, elastic_retry vuelve a generar la excepción de inmediato.
  • Si el error se debe a un evento de reducción de porciones, elastic_retry detiene cualquier seguimiento activo del generador de perfiles de JAX, ejecuta on_elastic_event_callback y evalúa retry_policy(attempt, error).
  • Antes de comenzar el siguiente intento, elastic_retry llama a elastic.wait_for_slices una vez para sondear el estado de los segmentos hasta que al menos minimum_slice_count segmentos estén activos. Luego, actualiza elastic_manager.active_slice_indices y ejecuta la función de bucle de entrenamiento en las porciones restantes.

Aumenta la fragmentación cuando se vuelven a unir los segmentos

Cuando tu carga de trabajo se ejecuta en una cantidad reducida de porciones en el modo de cambio de tamaño de la réplica, el subproceso del monitor en segundo plano ejecuta verificaciones de estado periódicas en elastic_manager.inactive_slice_indices. Cuando las porciones reparadas o reprogramadas aprueban la verificación de estado, el subproceso de supervisión agrega sus índices de porción a elastic_manager.available_inactive_slices.

Cuando elastic_manager.available_inactive_slices no está vacío al final de un paso de entrenamiento, la carga de trabajo vuelve a escalar verticalmente en su lugar sin perder el progreso del entrenamiento:

  • Agregar los índices de porciones en elastic_manager.available_inactive_slices a elastic_manager.active_slice_indices incorpora las porciones unidas en elastic_manager.active_devices y las borra de available_inactive_slices.
  • Reconstruir la malla de dispositivos desde elastic_manager.active_devices y reinicializar el paso de entrenamiento y el iterador de datos adapta el cálculo a la topología expandida.
  • Cuando se llama a jax.device_put, se coloca el estado de entrenamiento actual en las nuevas fragmentaciones de malla sin restablecer desde una instantánea.

Puntos de control de varios niveles

Los trabajos de entrenamiento a gran escala experimentan interrupciones frecuentes, y escribir o restablecer puntos de control directamente desde y hacia Cloud Storage puede generar cuellos de botella en el progreso del entrenamiento. La creación de puntos de control de varios niveles reduce la sobrecarga de la creación de puntos de control y el tiempo de recuperación coordinando tres niveles de almacenamiento en GKE:

  • Disco de RAM local: Los trabajadores TPU escriben puntos de control de alta frecuencia en un volumen de CSI en memoria (multitier-checkpoint.csi.storage.gke.io) activado en cada nodo trabajador. Este volumen persiste entre los reinicios del Pod en el mismo nodo.
  • Replicación de pares en el clúster: El servicio de replicador de GKE replica fragmentos de puntos de control locales en los trabajadores TPU pares del clúster.
  • Copia de seguridad de Cloud Storage: El servicio de replicador crea copias de seguridad de forma asíncrona de los puntos de control locales en un bucket de Cloud Storage con el espacio de nombres jerárquico habilitado, sin bloquear los pasos de entrenamiento de TPU.

A diferencia de JAX con varios controladores, en el que cada host de TPU ejecuta su propia secuencia de comandos de entrenamiento, Pathways ejecuta un solo controlador en el Pod de pathways-head y envía el trabajo a los Pods de pathways-worker. Para leer y escribir fragmentos de puntos de control directamente en el disco RAM local de cada host de TPU sin enrutar los datos de los puntos de control a través del controlador único, Pathways usa Python ubicado en el mismo lugar. Cada Pod de worker ejecuta un contenedor de colocated-python-sidecar junto con el contenedor de pathways-worker, y ambos contenedores activan el volumen de CSI de puntos de control de varios niveles y un directorio de memoria compartido (/tmp/sidecar). Durante la inicialización y el entrenamiento, Orbax envía la configuración, los guardados y los restablecimientos de los puntos de control desde el controlador de Pathways a los sidecars de los trabajadores en cada host de TPU. La creación de puntos de control de varios niveles en Pathways también admite el entrenamiento elástico cuando se reduce o aumenta la cantidad de porciones de TPU activas.

Pathways admite la creación de puntos de control de varios niveles con MaxText. Para obtener instrucciones de configuración, consulta la guía de puntos de control de varios niveles de MaxText.

Configura tu carga de trabajo de Pathways

Puedes configurar una carga de trabajo de Pathways para la creación de puntos de control de varios niveles con Cluster Toolkit (gcluster job submit) o con un manifiesto de JobSet de Pathways personalizado.

Usa Cluster Toolkit

Cuando envíes un trabajo con gcluster job submit --pathways, pasa las marcas --gke-mtc-enabled, --gke-mtc-ramdisk-dir y --pathways-colocated-python-sidecar-image:

gcluster job submit \
    --pathways \
    --gke-mtc-enabled \
    --gke-mtc-ramdisk-dir="/tmp/mtc_checkpoints" \
    --pathways-colocated-python-sidecar-image="SIDECAR_IMAGE" \
    ...

Reemplaza SIDECAR_IMAGE por el URI de la imagen del contenedor sidecar de Python ubicado en el mismo lugar.

Cuando configuras estas marcas, gcluster establece automáticamente lo siguiente:

  • Pasa --sidecar_name=external al contenedor pathways-proxy.
  • Ejecuta el contenedor colocated-python-sidecar en cada Pod worker y activa un directorio /tmp/sidecar compartido (--cloud_pathways_sidecar_shm_directory=/tmp/sidecar) en pathways-worker y colocated-python-sidecar.
  • Activa el volumen de CSI de multitier-checkpoint.csi.storage.gke.io en la ruta de acceso especificada por --gke-mtc-ramdisk-dir en pathways-worker y colocated-python-sidecar en los Pods de worker, pero no en el Pod de pathways-head.

Usa un manifiesto de JobSet de Pathways personalizado

Si administras tu propio manifiesto de JobSet de Pathways en lugar de usar gcluster, pasa --sidecar_name=external a pathways-proxy. En la especificación del Pod de worker, activa el volumen de memoria de archivo adicional compartido y el volumen de CSI de multitier-checkpoint.csi.storage.gke.io en los contenedores pathways-worker y colocated-python-sidecar:

containers:
- name: pathways-worker
  args:
  - --cloud_pathways_sidecar_shm_directory=/tmp/sidecar
  volumeMounts:
  - name: shared-tmp
    mountPath: /tmp/sidecar
  - name: checkpoint-mount
    mountPath: /tmp/mtc_checkpoints
- name: colocated-python-sidecar
  volumeMounts:
  - name: shared-tmp
    mountPath: /tmp/sidecar
  - name: checkpoint-mount
    mountPath: /tmp/mtc_checkpoints
volumes:
- name: shared-tmp
  emptyDir:
    medium: Memory
- name: checkpoint-mount
  csi:
    driver: multitier-checkpoint.csi.storage.gke.io

Usa los puntos de control de varios niveles de Orbax en una carga de trabajo de JAX personalizada

En una secuencia de comandos de entrenamiento de JAX personalizada que se ejecuta en Pathways, usa el paquete orbax.checkpoint.experimental.emergency.multi_tier_checkpointing y establece use_colocated_python=True en initialize_multi_tier_checkpointing y ReplicatorCheckpointManagerOptions:

from etils import epath
import jax
import orbax.checkpoint as ocp
from orbax.checkpoint.experimental.emergency.multi_tier_checkpointing import (
    initialization,
    replicator_checkpoint_manager,
)
import pathwaysutils

pathwaysutils.initialize()

local_checkpoint_dir = epath.Path("/tmp/mtc_checkpoints")

# Initialize multi-tier checkpointing across worker sidecars.
# When using elastic training, pass the active slice devices to `devices`
# on initial setup and whenever the active slice count changes.
initialization.initialize_multi_tier_checkpointing(
    local_checkpoint_directory=local_checkpoint_dir,
    backup_interval_minutes=10,
    num_slices=num_slices,
    run_name=run_name,
    data_parallelism=data_parallelism,
    use_colocated_python=True,
    devices=jax.devices(),
)

options = replicator_checkpoint_manager.ReplicatorCheckpointManagerOptions(
    save_interval_steps=20,
    use_colocated_python=True,
)
checkpoint_manager = (
    replicator_checkpoint_manager.ReplicatorCheckpointManager(
        local_checkpoint_dir,
        options=options,
        global_mesh=mesh,
    )
)

# Restore from the latest checkpoint if one exists.
start_step = 1
latest_step = checkpoint_manager.latest_step()
if latest_step is not None:
  restore_args = jax.tree.map(
      lambda x: ocp.type_handlers.ArrayRestoreArgs(sharding=x.sharding),
      state,
  )
  restored = checkpoint_manager.restore(
      latest_step,
      args=ocp.args.Composite(
          state=ocp.args.PyTreeRestore(
              item=state,
              restore_args=restore_args,
          )
      ),
  )
  state = restored.state
  start_step = latest_step + 1

# Step 0 is reserved by the replicator service to indicate "no checkpoint".
for step in range(start_step, final_step + 1):
  state = jitted_train_step(state)
  checkpoint_manager.save(
      step,
      args=ocp.args.Composite(
          state=ocp.args.PyTreeSave(item=state)
      ),
  )

checkpoint_manager.wait_until_finished()
checkpoint_manager.close()

Cambio en caliente

El intercambio en caliente es una función de la API de JobSet de GKE que permite que un trabajo de mayor prioridad se haga cargo rápidamente de los recursos de un trabajo de menor prioridad, lo que minimiza el tiempo de inactividad y acelera la recuperación.

Cuando creas un JobSet, GKE programa la carga de trabajo en varias porciones, como se especifica en la configuración del JobSet. Si se produce una falla de hardware en uno o más sectores, GKE marca los Pods afectados como con errores. Si tu clúster de GKE tiene una porción de repuesto que ejecuta un trabajo de menor prioridad, el controlador de JobSet reasigna la carga de trabajo de la porción con errores del trabajo de mayor prioridad a esa porción de repuesto. Esta reasignación suele tardar menos de un minuto.

Cuando se reinicia un JobSet, el intercambio en caliente puede ocurrir en las siguientes situaciones:

  • Modo predeterminado: Si hay porciones de TPU inactivas y de repuesto disponibles dentro del mismo clúster, el programador de Kubernetes prioriza la programación de los trabajos reiniciados en estas porciones en lugar de esperar a que se reparen las porciones con errores. Este comportamiento de programación proporciona una recuperación más rápida.
  • Cargas de trabajo heterogéneas: En los clústeres que ejecutan varias cargas de trabajo con un recurso PriorityClass de Kubernetes configurado, un JobSet reiniciado puede activar un intercambio en caliente. Si la afinidad del trabajo reiniciado coincide con los recursos de un trabajo de menor prioridad, Kubernetes interrumpe el trabajo de menor prioridad para que el trabajo de mayor prioridad pueda iniciarse sin esperar a que se repare la porción. Por ejemplo, puedes configurar tus Pods de trabajador de Pathways con diferentes prioridades usando un recurso PriorityClass.

Para usar prioridades en tu clúster, define un manifiesto PriorityClass y guárdalo como high-prior-job.yaml:

apiVersion: scheduling.k8s.io/v1
kind: PriorityClass
metadata:
  name: high-prior-job
value: 2000
globalDefault: false
description: "This priority class should be used for high priority job."

Aplica este YAML a tu clúster de GKE:

kubectl apply -f high-prior-job.yaml

A continuación, adjunta la nueva clase de prioridad a tu carga de trabajo de Pathways. Si envías tu carga de trabajo con Cluster Toolkit, pasa --priority=high-prior-job al comando gcluster job submit, que establece priorityClassName en los Pods pathways-head y worker. Si escribes tu propio manifiesto de JobSet, configura priorityClassName en la plantilla de Pod del trabajo replicado worker, que ejecuta el contenedor pathways-worker:

replicatedJobs:
- name: worker
  template:
    spec:
      template:
        spec:
          priorityClassName: high-prior-job

Si creaste tu clúster con un plano de TPU de Cluster Toolkit con Pathways habilitado, el clúster ya incluye las clases de prioridad very-low, low, medium, high y very-high, por lo que puedes pasar una de ellas a --priority en lugar de crear la tuya.

¿Qué sigue?