将 ML 诊断与 MaxText 搭配使用

MaxText 是一个开源的大语言模型 (LLM) 库，旨在在 TPU 和 GPU 上运行，用于模型训练。ML Diagnostics SDK 已预先与 MaxText 集成。您可以使用 managed-mldiagnostics 标志通过 MaxText 启用 ML 诊断。

启用后，您将能够执行以下操作：

• 创建包含所有 MaxText 配置的受管理 MachineLearning 运行。
• 如果通过 profiler="xplane" 标志启用了性能剖析，则上传性能剖析轨迹。
• 使用 log_period 标志以定义的间隔上传训练指标。

将 ML 诊断与 MaxText 搭配使用时，请使用以下标志：

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

如需在 MaxText 中启用机器学习诊断，您可以更改运行的配置文件，也可以从命令行传递标志。例如，通过 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 标志可捕获来自所有主机的多主机配置文件。