Utiliser les diagnostics de ML avec MaxText

MaxText est une bibliothèque de grands modèles de langage (LLM) Open Source conçue pour s'exécuter sur des TPU et des GPU pour l'entraînement de modèles. Le SDK ML Diagnostics est préintégré à MaxText. Vous pouvez activer les diagnostics ML avec MaxText à l'aide du flag managed-mldiagnostics.

Si cette fonctionnalité est activée, vous pourrez effectuer les actions suivantes :

• Créez une exécution MachineLearning gérée avec toutes les configurations MaxText.
• Importez les traces de profilage si le profilage est activé par l'indicateur profiler="xplane".
• Importez les métriques d'entraînement à un intervalle défini à l'aide de l'indicateur log_period.

Lorsque vous utilisez ML Diagnostics avec MaxText, utilisez les indicateurs suivants :

managed_mldiagnostics: True  # Enable the managed diagnostics
managed_mldiagnostics_run_group: GROUP_NAME  # Used to group multiple runs. (Optional)

Pour activer les diagnostics ML dans MaxText, vous pouvez modifier le fichier de configuration de votre exécution ou transmettre les indicateurs à partir de la ligne de commande. Par exemple, transmettez les indicateurs suivants avec la commande MaxText.train :

python3 -m MaxText.train \
  src/MaxText/configs/base.yml \
  run_name="demo-mldiagnostics-run-2" \
  model_name="<your_chosen_model>" \
  base_output_directory=gs://<your_gcs_folder>/ \
  dataset_type=synthetic \
  steps=100 \
  log_period=10 \
  profiler=xplane \
  upload_all_profiler_results=True \
  managed_mldiagnostics=True \
  managed_mldiagnostics_run_group="demo-mldiagnostics-group"

Dans cet exemple, la définition de l'indicateur upload_all_profiler_results=True capture les profils multi-hôtes de tous les hôtes.