ML Diagnostics mit MaxText verwenden

MaxText ist eine Open-Source-Bibliothek für Large Language Models (LLMs), die für die Ausführung auf TPUs und GPUs für das Modelltraining entwickelt wurde. Das ML Diagnostics SDK ist bereits in MaxText integriert. Sie können ML Diagnostics mit MaxText aktivieren, indem Sie das Flag managed-mldiagnostics verwenden.

Wenn die Option aktiviert ist, haben Sie folgende Möglichkeiten:

  • Erstellen Sie einen verwalteten MachineLearning-Lauf mit allen MaxText-Konfigurationen.
  • Laden Sie Profiling-Traces hoch, wenn das Profiling durch das Flag profiler="xplane" aktiviert ist.
  • Mit dem Flag log_period können Sie Trainingsmesswerte in einem definierten Intervall hochladen.

Verwenden Sie die folgenden Flags, wenn Sie ML Diagnostics mit MaxText verwenden:

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

Wenn Sie ML Diagnostics in MaxText aktivieren möchten, können Sie entweder die Konfigurationsdatei Ihres Laufs ändern oder die Flags über die Befehlszeile übergeben. Übergeben Sie beispielsweise die folgenden Flags mit dem Befehl 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"

Wenn Sie in diesem Beispiel das Flag upload_all_profiler_results=True festlegen, werden Multi-Host-Profile von allen Hosts erfasst.