Entrenamiento resiliente con Rutas

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 para controlar las 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.

Antes de comenzar

Antes de comenzar, completa las siguientes tareas:

Suspensión y reanudación

Por lo general, Google Kubernetes Engine (GKE) envía un aviso de interrupción a un Pod acelerador antes de que se interrumpa. La tolerancia a la interrupción de Pathways está habilitada 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 así, intenta suspender de forma transparente tu carga de trabajo de AA escribiendo su estado actual en un almacenamiento persistente, como Cloud Storage, antes de que GKE expulse 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 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 se basan en componentes de XLA de Pathways sin estado, que se pueden restablecer con una instantánea de memoria de ancho de banda alto (HBM). Ciertas cargas de trabajo de AA, como las que se definen con la API de Python de JAX ubicada en el mismo lugar, dependen de componentes de Pathways con estado. Estas cargas de trabajo 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: Cuando minimum_slice_count es None o igual a la cantidad total de porciones en el trabajo, el entrenamiento se detiene cuando se pierde una porción y espera hasta que todas las porciones vuelvan a estar en buen estado antes de reanudarse. Debido a que el recuento de porciones no cambia entre los intentos, tu carga de trabajo puede volver a usar la malla de dispositivos existente y los ejecutables compilados.
    • Cambio de tamaño de la réplica: Cuando minimum_slice_count es menor que la cantidad total de porciones, el entrenamiento continúa de inmediato en las porciones activas restantes reconstruyendo la malla del dispositivo y volviendo 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 volver a compilar 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, mientras que ElasticIterator hace un seguimiento de una sola posición del conjunto de datos global que se puede 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.
  • Decorar la función run_training_loop interna con @elastic_manager.elastic_retry espera porciones en buen estado, establece jax.default_device en una porción activa, mantiene activo el subproceso del monitor en segundo plano en todos los pasos y vuelve a intentar el bucle cuando se produce una jax.errors.JaxRuntimeError.
  • Cuando se recupera de una falla de porción al comienzo de un intento, 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, mientras que los búferes del dispositivo del intento fallido quedan fuera del alcance.
  • 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 se 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 se detecta un jax.errors.JaxRuntimeError de reducción, 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). El decorador 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, lo que sondea el estado de las porciones hasta que al menos minimum_slice_count porciones estén activas, 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.

Cambio en caliente

El intercambio en caliente se refiere a una función de la API de JobSet de GKE en la que un trabajo de mayor prioridad puede tomar rápidamente los recursos de un trabajo de menor prioridad, lo que minimiza el tiempo de inactividad y garantiza una recuperación más rápida.

Cuando se crea 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 segmentos, los Pods afectados se marcan como fallidos. Si tu clúster de GKE tiene una porción de repuesto que ejecuta un trabajo de menor prioridad, el sistema 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?