搭配 MaxText 使用 ML 診斷

MaxText 是開放原始碼的大型語言模型 (LLM) 程式庫，專為在 TPU 和 GPU 上執行模型訓練而建構。ML Diagnostics SDK 已預先與 MaxText 整合。您可以使用 managed-mldiagnostics 旗標，透過 MaxText 啟用 ML 診斷。

啟用後，您將可以執行下列操作：

• 使用所有 MaxText 設定建立受管理 MachineLearning 執行作業。
• 如果透過 profiler="xplane" 旗標啟用剖析功能，請上傳剖析追蹤記錄。
• 使用 log_period 標記，按照定義的間隔上傳訓練指標。

搭配 MaxText 使用 ML Diagnostics 時，請使用下列標記：

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

如要在 MaxText 中啟用 ML 診斷功能，您可以變更執行作業的設定檔，也可以從指令列傳遞旗標。舉例來說，請使用 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"

在本範例中，設定 upload_all_profiler_results=True 標記會從所有主機擷取多主機設定檔。