MaxText で ML 診断を使用する

MaxText は、モデル トレーニング用に TPU と GPU で実行するように構築されたオープンソースの大規模言語モデル（LLM）ライブラリです。ML Diagnostics SDK は MaxText に事前に統合されています。managed-mldiagnostics フラグを使用すると、MaxText で ML Diagnostics を有効にできます。

有効にすると、次の操作を行えるようになります。

• すべての MaxText 構成を使用して、マネージド MachineLearning 実行を作成します。
• profiler="xplane" フラグでプロファイリングが有効になっている場合は、プロファイリング トレースをアップロードします。
• log_period フラグを使用して、定義された間隔でトレーニング指標をアップロードします。

MaxText で ML 診断を使用する場合は、次のフラグを使用します。

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 フラグを設定すると、すべてのホストからマルチホスト プロファイルがキャプチャされます。