実験の追跡と可観測性

Important

この機能は パブリック プレビュー段階です。

AI RuntimeはMLflowと統合され、実験の追跡とオブザーバビリティを提供します。 MLflowは、ノートブックで開発する場合でも、 databricks air CLIでワークロードを提出する場合でも、ジョブのスケジュールを決める場合でも、実行のパラメータ、メトリクス、GPUシステムメトリクス、ログ、アーティファクトを一つの場所で管理できます。

MLflowが提供するもの

MLflow実験は関連するトレーニングランをグループ化し、各ランはトレーニング設定と結果の永続的な記録を提供します。

  • 指標とパラメータ:トレーニング損失、評価指標、学習率、ハイパーパラメータを記録し、それらを実行ごとに比較します。
  • システム メトリック: 管理された実行のGPU、CPU、メモリ利用率をシステム メトリックタブで確認できます。
  • ログ: ジョブバックワークロードのドライバー出力をランの ログ タブで確認できます。
  • アーティファクトとモデル: モデルファイル、構成、その他の出力を実行とともに保存します。 大きなアーティファクトや耐久性のあるアーティファクトにはUnityカタログのボリュームを使いましょう。
  • 共有と協力:実験へのアクセス権を付与し、他のワークスペースユーザーがその実験を閲覧・比較できるようにします。 「MLflow 実験を使用してトレーニング実行を整理する」を参照してください。
  • フレームワーク統合:Hugging Face Transformers、PyTorch Lightning、その他対応ライブラリとのMLflow連携を活用してください。

ジョブに支えられたワークロード、例えばdatabricks air runを使用して送信されたワークロードでは、Azure Databricksジョブ実行がMLflowを補完し、実行状況、コンピュート、試行回数、ドライバー出力を記録します。

MLflowのランがどのように作成されるか

MLflow run の開始または参加方法は、ワークロードの開始方法によります:

ワークロードの実行方法 MLflowランの開始方法 実験の選択方法
Databricks CLI (databricks air run) 各試行ごとに実行が自動的に作成されます。 ワークロードYAMLでexperiment_nameを設定してください。
Databricks AI Runtime SDK (.distributed())を搭載したノートブック 各呼び出しで、自動的に実行が作成されます。 mlflow.set_experiment()の前に.distributed()に電話するか、MLFLOW_EXPERIMENT_NAMEをセットしてください。
コードは直接接続されたノートブック上で実行されます run は自動的に作成されません。 mlflow.start_run() か autologging を使いましょう。 ノートブックで mlflow.set_experiment() を呼び出します。
スケジュール済み ai_runtime_task タスクは実行結果を自動的に作成します。 タスク設定で experiment を設定してください。 詳細は AIランタイムタスクを参照してください。

コードが正しいrunを開始またはそのrunに接続した後、MLflowトラッキングAPIを使ってカスタムトレーニングデータを記録してください。 サポートフレームワークでは、個別のログ呼び出しを追加する代わりに自動ログを使うことができます。 ディープラーニングのパターンについては、 MLflow 3 ディープラーニングワークフローを参照してください。

MLflowトラッキングを設定する際は以下の方法を用いてください:

  • stepをグローバルバッチやエポックのような意味のある値に設定します。 毎バッチではなく、一定の間隔でログを取ること。 MLflowは1ランあたりのメトリックステップ数を制限します。 リソース制限に関するページを参照してください。
  • 分散トレーニングでは、ランク0からのみメトリックをログに記録し、各値を一度だけ記録します。
  • ノートブックのワークフローでは、/Users/<username>/my-experiment や /Workspace/Shared/<team>/my-experiment などの絶対実験パスを使いましょう。 共有したい実験を共有ワークスペースフォルダに入れてください。
  • Unityカタログのボリュームに大きくて耐久性のあるチェックポイントを保存します。 チェックポイントパターンについては、「 AI Runtimeにおけるトレーニングパフォーマンスとレジリエンスの改善」を参照してください。

Note

このページの例ではMLflow 3.7以降を使用します。

Databricks CLIワークロードを追跡

databricks air runでワークロードを提出すると、AIランタイムは実行詳細用のAzure Databricksジョブランと実験データ用のMLflowランを作成します。 元の提出とリトライは1回のジョブ実行を共有し、各試みごとに同じ実験内で別々のMLflow実行が作成され、試行を比較できます。

GPU、CPU、メモリシステムの指標は、管理された各実行ごとに自動的にキャプチャされ、MLflow システムの指標 タブに表示されます。

MLflow実行のシステムメトリックタブにGPU、CPU、メモリの指標が表示されます。

設定し、実行ログを記録する

ワークロードYAMLを使って実験を選択し、実行名とアーティファクトの位置をオプションで設定してください。 以下の例はMLflowフィールドのみを示しています:

experiment_name: my-training
mlflow_run_name: baseline-lr3e5
mlflow_artifact_location: /Volumes/main/default/mlflow-artifacts/my-training

フィールド要件、デフォルト、サポートされる値は Workload YAMLリファレンスを参照してください。

AI RuntimeはMLFLOW_RUN_IDを通じてマネージド実行 IDを公開します。 そのランに自分のパラメータ、指標、アーティファクトをログに記録するには、 MLflowトラッキングAPIで再度開いてください:

import os

import mlflow

# Log from rank 0 only because all nodes share the same MLFLOW_RUN_ID.
if os.environ.get("RANK", "0") == "0":
    with mlflow.start_run(run_id=os.environ["MLFLOW_RUN_ID"]):
        mlflow.log_param("learning_rate", 3e-4)
        for step, loss in enumerate(training_losses):
            mlflow.log_metric("train_loss", loss, step=step)

実行とログの検索

ログはDatabricksのCLIとMLflowの実行ページの両方から利用可能です。 以下のコマンドを使ってワークロードの検索、ワークロードの検査、ログのストリーミングやダウンロードを行います。

databricks air list --all-status
databricks air get <job-run-id>
databricks air logs <job-run-id>
databricks air logs <job-run-id> --node 2
databricks air logs <job-run-id> --download-to ./logs/

ノートブックのワークロードを追跡

ノートブックからは、添付されたセッションで直接コードを実行したり、Databricks AI Runtime SDKの @distributed デコレーターを通じて分散トレーニングを実行したりできます。

ノートブック上で直接コードを実行

付属ノートブック上で直接実行されるコードは、自動的にMLflow runを作成するわけではありません。 実験を選択し、実行を開始するか、自動ログを有効にする:

import mlflow

mlflow.set_experiment("/Users/<username>/my-experiment")

with mlflow.start_run(run_name="baseline"):
    train()

PyTorch Lightningのオートログを有効にするには、以下のパターンを使用します。

mlflow.pytorch.autolog()

with mlflow.start_run(run_name="lightning-baseline"):
    trainer.fit(model, datamodule=datamodule)

Hugging Face Transformers の場合は、以下のように TrainingArguments を設定します:

from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="/Volumes/<catalog>/<schema>/<volume>/checkpoints",
    report_to="mlflow",
    run_name="llama7b-sft-lr3e5",
    logging_steps=50,
)

Databricks AI Runtime SDKを使いましょう

ノートブックからの分散トレーニングには、Databricks AI Runtime SDK（@distributed）のdatabricks-sdk-airデコレーターを使いましょう。 SDKは、GPU環境バージョン6以降では databricks.air 名前空間で、GPU環境バージョン5以下では serverless_gpu 名で利用可能です。

from databricks.air import distributed

@distributed()
def train():
    # Training code
    ...

train.distributed()

各 .distributed() コールは自動的にMLflow runを生成します。 デフォルトの実験は /Users/{WORKSPACE_USER}/{notebook-name} です。 別の実験を選択するには、 .distributed()を呼ぶ前に設定してください:

mlflow.set_experiment("/Users/<username>/my-experiment")
train.distributed()

アクティブな実行内で .distributed() を呼び出すと、AI Runtimeはネストされた子実行を生成します:

with mlflow.start_run() as parent_run:
    train.distributed()

前回の実行を再開するには、MLFLOW_RUN_ID の前に .distributed() を設定してください:

import os

os.environ["MLFLOW_RUN_ID"] = "<previous-run-id>"
train.distributed()

セル出力の run link またはノートブックの Experiment パネルを使ってMLflowランを開いてください。

ログを見る

  • ノートブック出力: 標準出力とトレーニングコードのエラーはノートブックセルの出力に表示されます。
  • MLflowログ: .distributed(remote=True) 呼び出しなどのジョブに裏付けられたノートブック ワークロードでは、MLflow ログ タブに関連するジョブ実行の出力が表示されます。

ライブGPUリソースの監視

ノートブックの GPUリソース ペインでは、シングルノードおよびマルチノードのワークロードに対するリアルタイムのGPU利用率、メモリ、温度を表示します。 実行に紐づく永続レコードが必要なときは、MLflow System metrics タブを使いましょう。

ペインを開くには、ノートブックをAI Runtimeに接続し、右側のペインでチップ アイコン。GPU リソースをクリックしてください。

各 GPU の使用率、メモリ、温度のメトリックを示す [GPU リソース] ウィンドウ。

このウィンドウでは、10 秒ごとにメトリックがポーリングされ、最大 2 時間の履歴が保持されます。 更新アイコンをクリックします。 最新の値を取得するには、更新します。 5分間操作がないと、ペインが一時停止します。 再度開いて監視を再開してください。

ログへのアクセスのトラブルシューティング

MLflowの実行ページのLogsタブは、関連するAzure Databricksジョブ実行からの出力をストリーミングします。 MLflow実験へのアクセスがジョブへのアクセスを意味するわけではありません。

タブに「これらのログにアクセスできません」と表示された場合は、「管理可能」権限を持つユーザーまたはワークスペース管理者に、ジョブに対して少なくとも「閲覧可能」権限を付与してもらうよう依頼してください。 「ジョブへのアクセスを制御する」を参照してください。

Azure Databricksのグローバル制限

リソース制限に関するページを参照してください。

  • 最終更新: