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.

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)

Post-traitement et analyse

Vous pouvez calculer les métriques de débit utile pour une tâche terminée à partir de n'importe quelle machine, comme une VM de processeur standard ou votre ordinateur portable. Vous n'avez pas besoin d'utiliser un TPU pour le post-traitement et l'analyse.

Le code suivant génère le débit utile total de la tâche :

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]}%")

Pour en savoir plus sur les types de débit inutile (BadputType), consultez Détails de la répartition du débit inutile.

Pour spécifier une période donnée, utilisez get_job_goodput_interval :

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

L'heure de début (start_time_utc) et l'heure de fin (end_time_utc) sont datetime objets. Pour en savoir plus, consultez goodput.py.

Surveiller avec les tableaux de bord Goodput

Pour vous aider à surveiller et à visualiser vos charges de travail d'entraînement de machine learning, Google Cloud propose deux ensembles de tableaux de bord Goodput : le tableau de bord GKE JobSet Goodput et le tableau de bord Cloud ML Goodput. Utilisez le tableau de bord GKE JobSet pour diagnostiquer les problèmes d'infrastructure ou de planification, et les tableaux de bord Cloud ML Goodput pour identifier les goulots d'étranglement dans le code d'entraînement.

Google Cloud Tableau de bord ML Goodput

Le tableau de bord ML Goodput mesure l'efficacité au niveau de l'application de votre script d'entraînement. Il fournit des insights sur le temps consacré à l'entraînement productif et sur les sources de débit inutile telles que le chargement des données, l'initialisation ou la récupération après des interruptions.

Pour afficher le tableau de bord ML Goodput dans Cloud Monitoring, procédez comme suit :

  1. Dans la Google Cloud console, accédez à la page Cloud Monitoring.
    Accéder à la console Monitoring
  2. Dans le volet de navigation, cliquez sur Tableaux de bord.
  3. Dans le champ Recherche de filtres, saisissez "Tableau de bord Cloud ML Goodput".

Le tableau de bord ML Goodput inclut les métriques suivantes (précédées de workload/) :

  • goodput_time : temps cumulé consacré aux étapes d'entraînement productives.
  • badput_time: temps cumulé consacré aux activités non productives, ainsi que la source de débit inutile :
    • ACCELERATOR_INITIALIZATION : temps nécessaire à la configuration des TPU.
    • TRAINING_PREP : chargement des points de contrôle, création du modèle et de l'optimiseur.
    • PROGRAM_STARTUP : compilation JIT, traçage de graphiques.
    • DATA_LOADING_SYNC : temps passé bloqué sur l'entrée de données.
    • CHECKPOINT_SAVE : temps nécessaire à l'enregistrement de l'état du modèle.
    • CHECKPOINT_RESTORE : temps nécessaire à la restauration de l'état du modèle.
    • WASTED_PROGRESS: temps productif perdu en raison d'interruptions survenues avant un point de contrôle.
    • INFRASTRUCTURE_RECOVERY : temps d'arrêt lors des redémarrages de tâches.
    • CUSTOM_BADPUT_EVENTS: événements synchrones définis par l'utilisateur, tels que des évaluations.
  • step_time_deviation: mesure la gigue et l'instabilité des temps d'étape d'entraînement.
  • interval_goodput: débit utile calculé sur des fenêtres glissantes (par exemple, la dernière heure).

Tableau de bord GKE JobSet

Le tableau de bord GKE JobSet Goodput se concentre sur l'efficacité de la couche d'orchestration et vous aide à déterminer si GKE planifie rapidement les pods et les maintient en cours d'exécution.

Pour afficher des informations complètes sur l'état et les performances des JobSets, accédez au tableau de bord Surveillance des JobSets dans la Google Cloud console :

Accéder au tableau de bord de surveillance des JobSets

Le tableau de bord de surveillance des JobSets inclut les métriques suivantes :

  • kubernetes.io/jobset/scheduling_goodput: fraction du temps pendant lequel toutes les ressources requises (pods) pour le JobSet sont prêtes et disponibles pour effectuer des tâches, par rapport à l'heure de création du JobSet. Les valeurs faibles indiquent des retards dans la planification des pods, l'extraction d'images ou l'allocation de ressources dans GKE.

  • kubernetes.io/jobset/proxy_runtime_goodput: fraction estimée du temps pendant lequel les TPU sont activement utilisés, en fonction de signaux au niveau du système tels que le cycle de service de l'accélérateur. Cela fournit une vue de haut niveau de la productivité d'exécution sans instrumentation d'application.

  • Métriques du pool de nœuds: Informations sur l'état, la disponibilité et les interruptions affectant les pools de nœuds hébergeant le JobSet. Cela permet de corréler les baisses de débit utile avec les problèmes de nœuds sous-jacents.

Pour en savoir plus sur le tableau de bord GKE JobSet, consultez Surveiller le débitutiledesJobSets.

Guide de dépannage

Cette section fournit des informations de dépannage pour vous aider à identifier et à résoudre les problèmes courants que vous pouvez rencontrer lors de l'utilisation de la bibliothèque ML Goodput Measurement.

Métriques Cloud Monitoring manquantes

Si des métriques Cloud Monitoring sont manquantes :

  • Vérifiez les autorisations IAM : le compte de service associé à votre charge de travail doit disposer des autorisations monitoring.timeSeries.create et monitoring.metricDescriptors.create.
  • Vérifiez la configuration : les métriques de débit utile doivent être explicitement activées dans votre configuration. Par exemple, vérifiez que enable_gcp_goodput_metrics=True est transmis dans votre GCPOptions.

Métriques générales manquantes

Si des métriques générales sont manquantes :

  • Vérifiez l'état du moniteur : la surveillance doit être activée pour l'instanciation de votre moniteur. Transmettez monitoring_enabled=True au moniteur.

Métriques corrompues

Si vos métriques sont corrompues ou affichent des valeurs étranges :

  • Assurez-vous que les noms d'exécution sont uniques : toutes les valeurs run_name doivent être uniques pour chaque expérience. La réutilisation d'un nom d'exécution mélange les journaux des anciennes et des nouvelles exécutions.

Journaux manquants

Si des journaux Cloud Logging sont manquants :

  • Activez l'API Cloud Logging : l'API Cloud Logging doit être activée dans votre Google Cloud projet.
  • Vérifiez les paramètres de l'enregistreur : l'option logging_enabled=True doit être transmise à l'enregistreur Goodput.
  • Vérifiez le processus principal : le processus principal (où jax.process_index() == 0) doit signaler activement les journaux d'application.

Points de contrôle de débit inutile manquants

Si des points de contrôle de débit inutile sont manquants :

  • Vérifiez votre bibliothèque de points de contrôle : la bibliothèque Goodput suit automatiquement le temps de point de contrôle pour Orbax. Toutefois, si vous utilisez d'autres bibliothèques (telles que le point de contrôle PyTorch), vous devez encapsuler manuellement vos appels d'enregistrement et de restauration à l'aide de record_custom_badput_event_start_time et record_custom_badput_event_end_time.