実験の追跡と可観測性

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 の開始または参加方法は、ワークロードの開始方法によります:

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

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

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

Databricks CLIワークロードを追跡

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

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

実行の構成とログ出力

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

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

フィールドの要件、デフォルト値、サポートされている値については、ワークロード 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 リソース をクリックします。

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

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

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

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

Azure Databricksのグローバル制限

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