Entrenamiento resiliente con Rutas

Pathways proporciona beneficios de resiliencia de las siguientes maneras:

  • Suspender y reanudar: Tolerancia ante interrupciones planificadas, como avisos de anulación, sin necesidad de que el usuario escriba ningún código personalizado de control de anulación
  • Entrenamiento elástico: Tolerancia ante fallas de hardware no planificadas sin provocar que el cliente falle, pero que requiere que los usuarios escriban código de recuperación específico del modelo

Antes de comenzar

Asegúrate de tener lo siguiente:

Suspender y reanudar

Por lo general, GKE envía un aviso de anulación a un pod de acelerador antes de que se anule el pod. La tolerancia a la anulación de Pathways está habilitada de forma predeterminada en todas las implementaciones en la nube, y los trabajos de acelerador de Pathways escuchan estos avisos.

Cuando llega un aviso de anulación, Pathways primero determina si la carga de trabajo actual se puede restablecer, es decir, si Pathways 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 reprograme tus trabajos más adelante, Pathways reanudará tu carga de trabajo de AA leyendo su estado persistente.

Si la carga de trabajo no se puede restablecer, Pathways cierra el trabajo de 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 dependen de los componentes 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 definidas con la API de Python colocada de JAX , dependen de los componentes con estado de Pathways, que no se pueden restablecer.

Entrenamiento elástico

El entrenamiento elástico permite que tu trabajo de entrenamiento continúe incluso cuando se producen fallas de hardware. Esto se logra mediante una combinación de capacidades del sistema de Pathways y lógica de recuperación del modelo definida por el usuario:

  • Detección de fallas: Cuando se produce una falla de hardware (por ejemplo, falla un trabajador de TPU ), el sistema de Pathways la detecta y notifica al trabajo de entrenamiento del usuario a través de una excepción la próxima vez que se acceda a los datos ubicados en ese hardware. Esta notificación no falla tu carga de trabajo. Permite que tu código controle la notificación y vuelva a configurar tus recursos para continuar con el procesamiento o salir correctamente.
  • Controlador de elasticidad definido por el usuario: El código del modelo del usuario debe poder controlar esta excepción. Esto es lo que lo convierte en una "recuperación específica del modelo".
    • Creación de instantáneas: El enfoque más común es guardar periódicamente instantáneas del estado de tu modelo. Cuando se produce una falla, puedes cargar desde la instantánea más reciente para reanudar el entrenamiento.
    • Reconfiguración: Es probable que debas volver a configurar tu trabajo de entrenamiento para ajustarlo a la cantidad de porciones disponibles. Por ejemplo, si una porción deja de funcionar, puedes reducir la cantidad de porciones activas en uno hasta que haya un reemplazo disponible. Para obtener más información, consulta Controlador elástico.
    • Actualizaciones de gráficos de datos o de procesamiento: Tu código debe controlar cualquier cambio en la cantidad de dispositivos disponibles para tu procesamiento mediante la recreación del gráfico de procesamiento según sea necesario. Esto puede implicar volver a particionar los datos o volver a compilar tu modelo.
  • Función de Pathways en la recuperación: Pathways proporciona los elementos primitivos para admitir la reconfiguración definida por el usuario:
    • Reemplazo de porciones: Si se reemplaza una porción con errores, se puede informar al cliente una vez que la nueva porción esté disponible. Luego, tu código puede volver a configurarse para usar esta nueva porción.
    • Recuperación transparente: Pathways controla los detalles de nivel inferior de la recuperación, como el restablecimiento de las conexiones a las partes correctas del clúster.
  • Utilidades en pathwaysutils: Un conjunto de utilidades de Pathways definidas en pathways-utils.

Implementa un controlador elástico

La mayor parte del código que tendrás que escribir estará en un controlador elástico definido por el usuario. Este controlador reacciona a los eventos elásticos (como que una porción de TPU deje de estar disponible) recreando la malla y reinicializando el bucle de entrenamiento.

Cada carga de trabajo es única. La complejidad del controlador elástico puede escalarse con la complejidad de la carga de trabajo. Las entradas y salidas del controlador deben ser los argumentos mínimos y los valores que se muestran necesarios para reinicializar el bucle de entrenamiento.

def elastic_handler(elastic_utils, *args, **kwargs):
  mesh = initialize_mesh(**kwargs["mesh_kwargs"])
  initial_state, initial_step, jitted_train_step, other_variables =
      initialize_training_loop(mesh, **kwargs["initialize_training_loop_kwargs"])

  step, snapshot = elastic_utils.get_next_snapshot()
  state = initial_state.replace(**snapshot)

  return state, step, mesh, jitted_train_step, other_variables

Actualiza tu bucle de entrenamiento

Debes realizar los siguientes cambios en tu bucle de entrenamiento:

  1. Crea un administrador elástico.
  2. Une tu bucle de entrenamiento dentro de bloques try-except que controlen jax.errors.JaxRuntimeError.
  3. Dentro de tu controlador jax.errors.JaxRuntimeError, llama a maybe_reshard_down. El administrador elástico volverá a fragmentar si el error está relacionado con un evento elástico o, de lo contrario, lo volverá a generar.
  4. Llama a maybe_snapshot y maybe_reshard_up al final del bucle de entrenamiento.
import pathwaysutils
from pathwaysutils.elastic import manager

pathwaysutils.initialize()

def initialize_mesh(**kwargs):
  ...


def initialize_training_loop(**kwargs):
  ...


def train_loop(
    final_step,
    elastic_manager,
    mesh_kwargs,
    initialize_training_loop_kwargs,
):
  mesh = initialize_mesh(**mesh_kwargs)
  initial_state, initial_step, jitted_train_step, other_variables =
      initialize_training_loop(mesh, **initialize_training_loop_kwargs)

  step = initial_step
  while step < final_step:
    try:
      state = jitted_train_step(state)

      elastic_manager.maybe_snapshot(step=step, snapshot=state)
      handler_returns = elastic_manager.maybe_reshard_up(
          step=step,
          snapshot=state,
          elastic_handler=elastic_handler,
          handler_args=(),
          handler_kwargs=dict(
              mesh_kwargs=mesh_kwargs,
              initialize_training_loop_kwargs=initialize_training_loop_kwargs,
          ),
      )
      if handler_returns:
        state, step, mesh, jitted_train_step, other_variables = handler_returns
      step += 1
    except jax.errors.JaxRuntimeError as error:
      handler_returns = elastic_manager.maybe_reshard_down(
          error=error,
          elastic_handler=elastic_handler,
          handler_args=(),
          handler_kwargs=dict(
              mesh_kwargs=mesh_kwargs,
              initialize_training_loop_kwargs=initialize_training_loop_kwargs,
          ),
      )
      if handler_returns:
        state, step, mesh, jitted_train_step, other_variables = handler_returns

  return state


def main():
  elastic_manager = manager.Manager(
      devices=jax.devices(),
      snapshot_period=10,
      snapshot_buffer_size=1,
      reshard_check_period=5,
      max_elastic_down_event_count=10,
      max_reshard_retry_count=3,
  )

  train_loop(100, elastic_manager, {}, {})

Configura el administrador elástico

El administrador elástico se puede configurar de varias maneras. La frecuencia de la creación de instantáneas está determinada por el período de instantáneas. El período de instantáneas afecta la cantidad promedio de pasos perdidos debido a un evento elástico. El período de verificación de la fragmentación determina la frecuencia con la que tu bucle de entrenamiento sondeará la disponibilidad de porciones. El parámetro max_elastic_down_event_count te permite establecer la cantidad de eventos elásticos debido a la pérdida de porciones que admitirá tu bucle de entrenamiento. El parámetro max_reshard_retry_count especifica la cantidad de veces que el administrador elástico debe volver a intentar la fragmentación. El administrador es un objeto singleton y debe crearse solo una vez.

Instantáneas

Según la configuración del administrador elástico, la función puede crear instantáneas de datos en la memoria del host que estará disponible para que tu controlador elástico la use durante un evento elástico.

Reduce la fragmentación

Después de detectar un jax.errors.JaxRuntimeError, Pathways verificará si el error se debe a un evento elástico debido a una porción perdida. Si es así, llamará al controlador elástico en un bucle hasta que se realice correctamente o hasta que se alcance la cantidad máxima de reintentos. Si el error no se debe a un evento elástico, se volverá a generar. Los valores que muestra el controlador elástico se pasan al emisor.

Aumenta la fragmentación

Según la configuración del administrador elástico y si hay porciones no disponibles, Pathways verificará si hay porciones adicionales disponibles. Si es así, guardará de inmediato una instantánea (si aún no se tomó una instantánea preexistente para el paso actual) y llamará al controlador elástico en un bucle hasta que se realice correctamente o hasta que se alcance la cantidad máxima de reintentos. Si se vuelve a fragmentar, los valores que muestra el controlador elástico se pasan al emisor. De lo contrario, se muestra None.

Intercambio en caliente

El intercambio en caliente hace referencia a una función de la API de JobSet de GKE en la que un trabajo de mayor prioridad puede hacerse cargo rápidamente de 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 de JobSet. Si se produce una falla de hardware en una o más porciones, los pods afectados se marcan como fallidos. Cuando vuelvas a programar este Jobset, si elegiste conservar una porción de repuesto en tu clúster de GKE que se podría usar para un trabajo de menor prioridad, el sistema JobSet volverá a asignar la carga de trabajo de la porción fallida del trabajo de mayor prioridad a la porción de repuesto que usa el trabajo de menor prioridad dentro del mismo clúster de GKE. Por lo general, esta reasignación demora menos de un minuto.

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

  1. Modo predeterminado: Si hay porciones de TPU inactivas y de repuesto disponibles en el mismo clúster, el programador de Kubernetes priorizará la programación de los trabajos reiniciados en estas porciones en lugar de esperar a que se reparen las porciones fallidas. Esto proporciona una recuperación más rápida.
  2. Cargas de trabajo heterogéneas: En los clústeres que ejecutan varias cargas de trabajo con una PriorityClass de Kubernetes configurada, 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 anula el trabajo de menor prioridad, lo que permite que el trabajo de mayor prioridad comience de inmediato. Por ejemplo, puedes configurar tus pods de trabajador de Pathways con diferentes prioridades mediante PriorityClass.

Para usar prioridades en tu clúster, define una clase de prioridad, por ejemplo:

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

Luego, adjunta la nueva clase de prioridad a tu trabajo de trabajador de Pathways agregando el siguiente texto al podspec de tu pod pathways-worker.

priorityClassName: high-prior-job

¿Qué sigue?