Surveiller le débit utile avec la bibliothèque ML Goodput Measurement

La bibliothèque ML Goodput Measurement (ml-goodput-measurement) est un package Python qui vous aide à mesurer l'efficacité de vos charges de travail d'entraînement de ML exécutées sur des VM Cloud TPU. La bibliothèque fournit des métriques qui mesurent le débit utile de la charge de travail, qui correspond à la proportion du temps d'utilisation du TPU consacrée à des progrès d'entraînement productifs et préservés. À l'inverse, le débit inutile correspond à la proportion du temps total consacré à des activités non productives telles que la surcharge de démarrage, les blocages d'E/S et la récupération après une interruption.

Vous pouvez visualiser les métriques de débit utile en temps réel avec Cloud Monitoring et TensorBoard, ce qui vous permet d'identifier les goulots d'étranglement, d'optimiser l'utilisation des ressources et, en fin de compte, de réduire les coûts d'entraînement.

Pour en savoir plus, consultez le dépôt GitHub ML Goodput Measurement.

Métriques de débit utile

La bibliothèque ML Goodput Measurement fournit les métriques suivantes, qui sont également disponibles dans Cloud Monitoring et TensorBoard. Les métriques de ce tableau doivent être précédées de compute.googleapis.com/workload/. Par exemple, le nom complet de la métrique goodput_time est compute.googleapis.com/workload/goodput_time.

  • goodput_time : temps total d'entraînement productif en secondes. Cela peut être interprété comme le débit utile cumulé.

  • badput_time: temps total d'entraînement non productif en secondes (démarrage, blocages, récupération). Cela peut être interprété comme le débit inutile cumulé.

  • total_elapsed_time: temps total écoulé (durée réelle) de la charge de travail en secondes. Le temps écoulé est mesuré à partir du démarrage de l'application jusqu'à l'heure actuelle ou à l'heure de fin de la tâche.

  • interval_goodput: taux de débit utile sur une période spécifiée (par exemple, les dernières 24 heures). La métrique fournit une fenêtre glissante pour le débit utile.

  • interval_badput : taux de débit inutile sur une période spécifiée. La métrique fournit une fenêtre glissante pour le débit inutile et est utile pour identifier les problèmes temporaires tels que les pics d'opérations d'E/S.

  • disruptions : nombre cumulé d'interruptions de tâches. Les interruptions sont des événements qui entraînent l'arrêt inattendu du processus d'entraînement, ce qui nécessite un redémarrage. Par exemple, les défaillances matérielles et les événements de maintenance.

  • step_time_deviation: quantité de temps non productif due à la variation des temps d'étape d'entraînement. On parle également de "gigue". La métrique mesure la stabilité sensible aux pics, qui correspond à l'écart des étapes "les plus défavorables" récentes par rapport à une référence historique en secondes. Une étape correspond à une seule itération de la boucle d'entraînement.

  • performance : temps d'étape stable le plus rapide estimé (référence) en secondes. Le temps d'étape idéal est le temps le plus rapide qu'une seule étape d'entraînement peut prendre dans des conditions optimales, sans bruit ni gigue transitoires.

  • max_productive_steps: nombre de pas le plus élevé atteint et préservé.

Installation

Pour configurer la bibliothèque ML Goodput Measurement avec vos charges de travail TPU, procédez comme suit :

  1. Activez l'API Cloud Logging et l'API Cloud Monitoring.
  2. Si vous effectuez un déploiement sur Google Kubernetes Engine (GKE), configurez tous les pools de nœuds avec le niveau d'accès cloud-platform.

Installez le package ml-goodput-measurement sur votre hôte d'entraînement et votre machine d'analyse :

pip install ml-goodput-measurement

Surveiller avec la bibliothèque Goodput

Pour utiliser la bibliothèque ML Goodput Measurement, initialisez une instance GoodputRecorder, instrumentez votre code d'entraînement en encapsulant les sections clés avec le gestionnaire de contexte record_event, et surveillez les métriques en temps réel avec GoodputMonitor. Le moniteur exécute un processus en arrière-plan pour calculer périodiquement les métriques de débit utile à partir des événements enregistrés, puis les importe dans Cloud Monitoring et TensorBoard pour l'analyse et la visualisation en temps réel.

Initialiser l'enregistreur Goodput

Initialisez GoodputRecorder, qui est le composant principal de la bibliothèque 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)
)

Enregistrer des événements

Utilisez le gestionnaire de contexte record_event pour encapsuler votre code d'entraînement.

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()

Pour obtenir un exemple d'intégration du code d'entraînement à la bibliothèque Goodput à l'aide de MaxText, consultez goodput.py.

Surveiller les événements

Surveillez les métriques avec GoodputMonitor, qui démarre un processus en arrière-plan pour calculer et importer les métriques pendant l'exécution de la tâche. Encapsulez la logique dans un gestionnaire de contexte pour vous assurer que le processus d'importation démarre et s'arrête comme prévu.

Définissez un gestionnaire de contexte d'assistance pour gérer la configuration et la gestion du cycle de vie.

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.")

Pour mesurer les performances de la fenêtre glissante, configurez le moniteur pour qu'il utilise start_rolling_window_goodput_uploader et 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()

Encapsulez votre point d'entrée d'entraînement principal avec le gestionnaire de contexte. Cela garantit que la surveillance démarre avant le début de l'entraînement et que les métriques finales sont effacées à la fin de l'entraînement.