Supervisa el procesamiento útil con la biblioteca de ML Goodput Measurement

La biblioteca de ML Goodput Measurement (ml-goodput-measurement) es un paquete de Python que te ayuda a medir la eficiencia de tus cargas de trabajo de entrenamiento de AA que se ejecutan en VMs de Cloud TPU. La biblioteca proporciona métricas que miden el procesamiento útil de la carga de trabajo, que es la proporción del tiempo de uso de la TPU dedicado a realizar un progreso de entrenamiento productivo y preservado. Por el contrario, el procesamiento inútil es la proporción del tiempo total dedicado a actividades no productivas, como la sobrecarga de inicio, las detenciones de E/S y la recuperación de interrupciones.

Puedes visualizar las métricas de procesamiento útil en tiempo real con los paneles de Cloud Monitoring y TensorBoard, lo que te permite identificar cuellos de botella, optimizar el uso de recursos y, en última instancia, reducir los costos de entrenamiento.

Para obtener más información, consulta el repositorio de GitHub de ML Goodput Measurement.

Métricas de procesamiento útil

La biblioteca de ML Goodput Measurement proporciona las siguientes métricas, que también están disponibles para ver en Cloud Monitoring y TensorBoard. Las métricas de esta tabla deben tener el prefijo compute.googleapis.com/workload/. Por ejemplo, el nombre completo de la métrica para goodput_time es compute.googleapis.com/workload/goodput_time.

goodput_time: Es el tiempo total de entrenamiento productivo en segundos. Se puede interpretar como el procesamiento útil acumulativo.
badput_time: Es el tiempo total de entrenamiento no productivo en segundos (inicio, detenciones, recuperación). Se puede interpretar como el procesamiento inútil acumulativo.
total_elapsed_time: Es el tiempo total transcurrido (duración del reloj de pared) de la carga de trabajo en segundos. El tiempo transcurrido se mide desde el momento en que se inicia la aplicación hasta la hora actual o el momento en que se completa el trabajo.
interval_goodput: Es la tasa de procesamiento útil durante un período especificado (por ejemplo, las últimas 24 horas). La métrica proporciona una ventana móvil para el procesamiento útil.
interval_badput: Es la tasa de procesamiento inútil durante un período especificado. La métrica proporciona una ventana móvil para el procesamiento inútil y es útil para identificar problemas transitorios, como picos en las operaciones de E/S.
disruptions: Es el recuento acumulativo de interrupciones del trabajo. Las interrupciones son eventos que hacen que el proceso de entrenamiento se detenga de forma inesperada, lo que requiere un reinicio. Por ejemplo, fallas de hardware y eventos de mantenimiento.
step_time_deviation: Es la cantidad de tiempo no productivo debido a la variación en los tiempos de los pasos de entrenamiento. Esto también se conoce como "fluctuación". La métrica mide la estabilidad sensible a los picos, que es la desviación de los pasos recientes del "peor caso" en comparación con un modelo de referencia histórico en segundos. Un paso es una sola iteración del bucle de entrenamiento.
performance: Es el tiempo estimado más rápido y estable del paso (modelo de referencia) en segundos. El tiempo ideal del paso es el tiempo más rápido que puede tardar un solo paso de entrenamiento en condiciones óptimas, sin ruido ni fluctuaciones transitorios.
max_productive_steps: Es el recuento de pasos más alto alcanzado que se conservó correctamente.

Instalación

Sigue estos pasos para configurar la biblioteca de ML Goodput Measurement con tus cargas de trabajo de TPU:

Habilita la API de Cloud Logging y la API de Cloud Monitoring.
Si realizas la implementación en Google Kubernetes Engine (GKE), configura todos los grupos de nodos con el permiso de acceso cloud-platform.

Instala el paquete ml-goodput-measurement en tu host de entrenamiento y en la máquina de análisis:

pip install ml-goodput-measurement

Supervisa con la biblioteca de Goodput

Para usar la biblioteca de AA Goodput Measurement, inicializa una instancia de GoodputRecorder, instrumenta tu código de entrenamiento envolviendo las secciones clave con el administrador de contexto record_event y supervisa las métricas en tiempo real con GoodputMonitor. El supervisor ejecuta un proceso en segundo plano para calcular periódicamente las métricas de procesamiento útil a partir de los eventos registrados y las sube a Cloud Monitoring y TensorBoard para el análisis y la visualización en tiempo real.

Inicializa la grabadora de Goodput

Inicializa GoodputRecorder, que es el componente principal de la biblioteca de ML Goodput Measurement.

import datetime
import jax
from ml_goodput_measurement import measurement

# Define a unique logger name for this specific run
logger_name = f'goodput_{config.run_name}'

# Instantiate the recorder
goodput_recorder = measurement.GoodputRecorder(
    job_name=config.run_name,
    logger_name=logger_name,
    logging_enabled=(jax.process_index() == 0)
)

Cómo grabar eventos

Usa el administrador de contexto record_event para envolver tu código de entrenamiento.

def train_loop(config):

  # 1. Wrap the entire Job (Start/End)
  with goodput_recorder.record_event(measurement.Event.JOB):

    # 2. Record Hardware Initialization
    with goodput_recorder.record_event(measurement.Event.ACCELERATOR_INIT):
      # ... perform device mesh setup ...
      initialize_tpu(config)

    # 3. Record Training Prep
    with goodput_recorder.record_event(measurement.Event.TRAINING_PREP):
      # ... create checkpoint managers, setup model ...
      model = training_prep(config)

    # 4. Main Training Loop
    for step in range(config.steps):

      # Record Data Loading
      with goodput_recorder.record_event(measurement.Event.DATA_LOADING):
        batch = get_next_batch()

      # Record Step Start (CRITICAL: Pass the step number!)
      with goodput_recorder.record_event(measurement.Event.STEP, step):
        output = train_step(model, batch)

      # 5. Record Custom Events (e.g., Evaluation)
      if step % eval_interval == 0:
        with goodput_recorder.record_event(measurement.Event.CUSTOM, "eval_step"):
          run_evaluation()

Si deseas ver un ejemplo de cómo se puede integrar el código de entrenamiento con la biblioteca de Goodput mediante MaxText, consulta goodput.py.

Supervisar eventos

Supervisa las métricas con GoodputMonitor, que inicia un proceso en segundo plano para calcular y subir métricas mientras se ejecuta el trabajo. Envuelve la lógica en un administrador de contexto para asegurarte de que el proceso de carga se inicie y se detenga según lo previsto.

Define un administrador de contexto auxiliar para controlar la configuración y la administración del ciclo de vida.

import contextlib
from ml_goodput_measurement import monitoring

@contextlib.contextmanager
def maybe_monitor_goodput(config):
  """Monitor goodput if enabled and on the main process."""
  if not config.monitor_goodput or jax.process_index() != 0:
    yield
    return

  goodput_monitor = None
  try:
    # Configure GCPOptions for Cloud Monitoring
    gcp_options = monitoring.GCPOptions(
      enable_gcp_goodput_metrics=config.enable_gcp_goodput_metrics
    )

    # Instantiate the monitor
    goodput_monitor = monitoring.GoodputMonitor(
      job_name=config.run_name,
      logger_name=f"goodput_{config.run_name}",
      tensorboard_dir=config.tensorboard_dir,
      upload_interval=config.goodput_upload_interval_seconds,
      monitoring_enabled=True,
      pathway_enabled=config.enable_pathways_goodput,
      include_badput_breakdown=True,
      gcp_options=gcp_options,
    )

    # Start the background upload process
    goodput_monitor.start_goodput_uploader()
    print("Started Goodput upload to Tensorboard & GCM in the background!")
    yield

  finally:
    # Ensure clean shutdown of the background process
    if goodput_monitor:
      goodput_monitor.stop_goodput_uploader()
      print("Flushed final metrics and safe exited from Goodput monitoring.")

Para medir el rendimiento de la ventana móvil, configura el supervisor para que use start_rolling_window_goodput_uploader y stop_rolling_window_goodput_uploader.

try:
    self._rolling_window_monitor.start_rolling_window_goodput_uploader(
        self.config.rolling_window_size
    )
finally:
    if self._rolling_window_monitor:
        self._rolling_window_monitor.stop_rolling_window_goodput_uploader()

Envuelve tu punto de entrada de entrenamiento principal con el administrador de contexto. Esto garantiza que la supervisión comience antes de que comience el entrenamiento y que se borren las métricas finales cuando finalice el entrenamiento.

def main():
  # ... Load configuration ...

  # Wrap the entire execution
  with maybe_monitor_goodput(config):
    # Run the training loop (which contains the GoodputRecorder events)
    train_loop(config)

Análisis y procesamiento posterior

Puedes calcular las métricas de procesamiento útil para un trabajo completado desde cualquier máquina, como una VM de CPU estándar o tu laptop. No necesitas usar una TPU para el análisis y el procesamiento posterior.

El siguiente código muestra el procesamiento útil total del trabajo:

from ml_goodput_measurement import goodput

calculator = goodput.GoodputCalculator(
    job_name="my-run-name",
    logger_name="goodput_my-run-name"
)

goodput, badput, last_step = calculator.get_job_goodput(
    include_badput_breakdown=True
)

print(f"Goodput: {goodput}%")
print(f"Badput (Infra Recovery): {badput[goodput.BadputType.INFRASTRUCTURE_RECOVERY_FROM_DISRUPTION]}%")

Para obtener información sobre los tipos de procesamiento inútil (BadputType), consulta Detalles del desglose del procesamiento inútil.

Para especificar un período específico, usa get_job_goodput_interval:

goodput_pct, badput, _, _, _ = calculator.get_job_goodput_interval(
    start_time_utc,
    end_time_utc
)

La hora de inicio (start_time_utc) y la hora de finalización (end_time_utc) son objetos datetime. Para obtener más información, consulta goodput.py.

Supervisa con los paneles de Goodput

Para ayudarte a supervisar y visualizar tus cargas de trabajo de entrenamiento de aprendizaje automático, Google Cloud ofrece dos conjuntos de paneles de Goodput: el panel de GKE JobSet Goodput y el panel de Cloud ML Goodput. Usa el panel de GKE JobSet para diagnosticar problemas de infraestructura o programación, y los paneles de Cloud ML Goodput para identificar cuellos de botella dentro del código de entrenamiento.

Google Cloud Panel de ML Goodput

El panel de ML Goodput mide la eficiencia a nivel de la aplicación de tu secuencia de comandos de entrenamiento. Proporciona estadísticas sobre el tiempo dedicado al entrenamiento productivo y las fuentes de procesamiento inútil, como la carga de datos, la inicialización o la recuperación de interrupciones.

Para ver el panel de ML Goodput en Cloud Monitoring, haz lo siguiente:

En la Google Cloud consola de, dirígete a la página Cloud Monitoring.
Ir a la consola de Monitoring
En el panel de navegación, haz clic en Paneles.
En el campo Buscar filtro, ingresa "Panel de Cloud ML Goodput".

El panel de ML Goodput incluye las siguientes métricas (con el prefijo workload/):

goodput_time: Es el tiempo acumulativo dedicado a los pasos de entrenamiento productivos.
badput_time: Es el tiempo acumulativo dedicado a actividades no productivas, junto con la fuente de procesamiento inútil:
- ACCELERATOR_INITIALIZATION: Es el tiempo para configurar las TPU.
- TRAINING_PREP: Es la carga de puntos de control, la creación de modelos y optimizadores.
- PROGRAM_STARTUP: Es la compilación JIT, el seguimiento de gráficos.
- DATA_LOADING_SYNC: Es el tiempo dedicado al bloqueo de la entrada de datos.
- CHECKPOINT_SAVE: Es el tiempo para guardar el estado del modelo.
- CHECKPOINT_RESTORE: Es el tiempo para restablecer el estado del modelo.
- WASTED_PROGRESS: Es el tiempo productivo perdido debido a las interrupciones que ocurren antes de un punto de control.
- INFRASTRUCTURE_RECOVERY: Es el tiempo de inactividad durante los reinicios del trabajo.
- CUSTOM_BADPUT_EVENTS: Son eventos síncronos definidos por el usuario, como las evaluaciones.
step_time_deviation: Mide la fluctuación y la inestabilidad en los tiempos de los pasos de entrenamiento.
interval_goodput: Es el procesamiento útil calculado en ventanas móviles (p.ej., la última hora).

Panel de GKE JobSet

El panel de GKE JobSet Goodput se centra en la eficiencia de la capa de organización y te ayuda a comprender si GKE programa pods rápidamente y los mantiene en ejecución.

Para ver información completa sobre el estado y el rendimiento de JobSets, ve al panel de supervisión de JobSet en la Google Cloud consola de:

Ir al panel de supervisión de JobSet

El panel de supervisión de JobSet incluye las siguientes métricas:

kubernetes.io/jobset/scheduling_goodput: Es la fracción de tiempo en la que todos los recursos requeridos (pods) para el JobSet están en un estado listo y disponibles para realizar el trabajo, en relación con la hora de creación del JobSet. Los valores bajos indican retrasos en la programación de pods, la extracción de imágenes o la asignación de recursos dentro de GKE.
kubernetes.io/jobset/proxy_runtime_goodput: Es una fracción estimada de tiempo en la que las TPU se usan de forma activa, según los indicadores a nivel del sistema, como el ciclo de trabajo del acelerador. Esto proporciona una vista de alto nivel de la productividad del entorno de ejecución sin la instrumentación de la aplicación.
Métricas del grupo de nodos: Información sobre el estado, la disponibilidad y las interrupciones que afectan a los grupos de nodos que alojan el JobSet. Esto ayuda a correlacionar las caídas de procesamiento útil con los problemas subyacentes de los nodos.

Para obtener más información sobre el panel de GKE JobSet, consulta Supervisa el procesamiento útil de JobSet.

Guía de solución de problemas

En esta sección, se proporciona información para solucionar problemas que te ayudará a identificar y resolver problemas comunes que podrías encontrar mientras usas la biblioteca de ML Goodput Measurement.

Faltan métricas de Cloud Monitoring

Si te faltan métricas de Cloud Monitoring, haz lo siguiente:

Verifica los permisos de IAM: La cuenta de servicio adjunta a tu carga de trabajo debe tener los permisos monitoring.timeSeries.create y monitoring.metricDescriptors.create.
Verifica la configuración: Las métricas de procesamiento útil deben habilitarse de forma explícita en tu configuración. Por ejemplo, verifica que enable_gcp_goodput_metrics=True se pase dentro de tu GCPOptions.

Faltan métricas generales

Si te faltan métricas generales, haz lo siguiente:

Verifica el estado del supervisor: La supervisión debe estar habilitada para la creación de instancias del supervisor. Pasa monitoring_enabled=True al supervisor.

Métricas dañadas

Si tus métricas están dañadas o muestran valores extraños, haz lo siguiente:

Asegúrate de que los nombres de ejecución sean únicos: Todos los valores de run_name deben ser únicos para cada experimento. Si vuelves a usar un nombre de ejecución, se mezclarán los registros de las ejecuciones antiguas y nuevas.

Faltan registros

Si te faltan registros de Cloud Logging, haz lo siguiente:

Habilita la API de Cloud Logging: La API de Cloud Logging debe estar habilitada en tu Google Cloud proyecto.
Verifica la configuración de la grabadora: La marca logging_enabled=True debe pasarse a la grabadora de Goodput.
Verifica el proceso principal: El proceso principal (en el que jax.process_index() == 0) debe informar de forma activa los registros de la aplicación.

Faltan puntos de control de procesamiento inútil

Si te faltan puntos de control de procesamiento inútil, haz lo siguiente:

Verifica tu biblioteca de puntos de control: La biblioteca de Goodput realiza un seguimiento automático del tiempo de creación de puntos de control para Orbax. Sin embargo, si usas otras bibliotecas (como la creación de puntos de control de PyTorch), debes envolver manualmente tus llamadas de guardado y restablecimiento con record_custom_badput_event_start_time y record_custom_badput_event_end_time.