Usa ML Diagnostics con MaxText

MaxText es una biblioteca de modelos de lenguaje grandes (LLM) de código abierto creada para ejecutarse en TPU y GPU para el entrenamiento de modelos. El SDK de ML Diagnostics viene preintegrado con MaxText. Puedes habilitar ML Diagnostics con MaxText usando la marca managed-mldiagnostics.

Cuando esté habilitada, podrás realizar las siguientes acciones:

• Crea una ejecución de MachineLearning administrada con todas las configuraciones de MaxText.
• Sube los registros de seguimiento de la creación de perfiles si la marca profiler="xplane" habilita la creación de perfiles.
• Sube las métricas de entrenamiento en un intervalo definido con la marca log_period.

Usa las siguientes marcas cuando uses ML Diagnostics con MaxText:

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

Para habilitar ML Diagnostics en MaxText, puedes cambiar el archivo de configuración de tu ejecución o pasar las marcas desde la línea de comandos. Por ejemplo, pasa las siguientes marcas con el comando 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"

En este ejemplo, establecer la marca upload_all_profiler_results=True captura los perfiles de varios hosts de todos los hosts.