Utilizzare ML Diagnostics con MaxText

MaxText è una libreria open source di modelli linguistici di grandi dimensioni (LLM) creata per essere eseguita su TPU e GPU per l'addestramento dei modelli. L'SDK ML Diagnostics è preintegrato con MaxText. Puoi attivare la diagnostica ML con MaxText utilizzando il flag managed-mldiagnostics.

Se l'opzione è abilitata, potrai eseguire le seguenti operazioni:

  • Crea un'esecuzione di MachineLearning gestita con tutte le configurazioni MaxText.
  • Carica le tracce di profilazione, se la profilazione è abilitata dal flag profiler="xplane".
  • Carica le metriche di addestramento a un intervallo definito utilizzando il flag log_period.

Utilizza i seguenti flag quando utilizzi ML Diagnostics con MaxText:

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

Per attivare ML Diagnostics in MaxText, puoi modificare il file di configurazione dell'esecuzione o passare i flag dalla riga di comando. Ad esempio, passa i seguenti flag con il 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"

In questo esempio, l'impostazione del flag upload_all_profiler_results=True acquisisce i profili multihost da tutti gli host.