次の方法で共有

実験の追跡と可観測性

Important

単一ノード タスクの AI ランタイムは パブリック プレビュー段階です。 マルチ GPU ワークロード用の分散トレーニング API は ベータ版のままです。

この記事では、MLflow の使用、GPU の正常性の監視、ログの表示、AI ランタイムでのモデル チェックポイントの管理を行う方法について説明します。

MLflow の統合

AI ランタイムは、実験の追跡、モデルのログ記録、メトリックの視覚化のために MLflow とネイティブに統合されます。

セットアップの推奨事項:

  • MLflow をバージョン 3.7 以降にアップグレードし、 ディープ ラーニング ワークフロー パターンに従います。

  • PyTorch Lightning の自動ログ記録を有効にする:

    import mlflow
mlflow.pytorch.autolog()

  • mlflow.start_run() API スコープ内にモデル トレーニング コードをカプセル化して、MLflow の実行名をカスタマイズします。 これにより、実行名を制御でき、前の実行から再開できます。run_nameまたは MLflow をサポートするサード パーティ製ライブラリ (Hugging Face Transformers など) では、mlflow.start_run(run_name="your-custom-name") パラメーターを使用して実行名をカスタマイズできます。 それ以外の場合、既定の実行名は jobTaskRun-xxxxx。

    from transformers import TrainingArguments
args = TrainingArguments(
    report_to="mlflow",
    run_name="llama7b-sft-lr3e5",  # <-- MLflow run name
    logging_steps=50,
)

  • サーバーレス GPU API は、既定の名前 /Users/{WORKSPACE_USER}/{get_notebook_name()}を使用して MLflow 実験を自動的に起動します。 ユーザーは環境変数 MLFLOW_EXPERIMENT_NAMEで上書きできます。 MLFLOW_EXPERIMENT_NAME 環境変数には常に絶対パスを使用します。

    import os
os.environ["MLFLOW_EXPERIMENT_NAME"] = "/Users/<username>/my-experiment"

  • 前の実行から MLFLOW_RUN_ID を設定して、前のトレーニングを再開します。

    mlflow.start_run(run_id="<previous-run-id>")

  • stepの MLFlowLogger パラメーターを適切なバッチ番号に設定します。 MLflow には 1,000 万のメトリック ステップの制限があります。大規模なトレーニング実行で 1 つのバッチごとにログを記録すると、この制限に達する可能性があります。 リソース制限に関するページを参照してください。

ログの表示

  • ノートブックの出力 - トレーニング コードからの標準出力とエラーがノートブック のセル出力に表示されます。
  • ドライバー ログ - スタートアップの問題、環境のセットアップの問題、ランタイム エラーをデバッグするために、コンピューティング パネルからアクセスできます。
  • MLflow ログ - トレーニング メトリック、パラメーター、成果物は、MLflow 実験 UI で表示できます。

モデルのチェックポイント処理

モデル チェックポイントを Unity カタログ ボリュームに保存します。これは、他の Unity カタログ オブジェクトと同じガバナンスを提供します。 Databricks ノートブックからボリューム内のファイルを参照するには、次のパス形式を使用します。

/Volumes/<catalog>/<schema>/<volume>/<path>/<file-name>

チェックポイントをローカル ストレージに保存するのと同じ方法で、ボリュームに保存します。

次の例は、PyTorch チェックポイントを Unity カタログ ボリュームに書き込む方法を示しています。

import torch

checkpoint = {
    "epoch": epoch,  # last finished epoch
    "model_state_dict": model.state_dict(),  # weights & buffers
    "optimizer_state_dict": optimizer.state_dict(),  # optimizer state
    "loss": loss,  # optional current loss
    "metrics": {"val_acc": val_acc},  # optional metrics
    # Add scheduler state, RNG state, and other metadata as needed.
}
checkpoint_path = "/Volumes/my_catalog/my_schema/model/checkpoints/ckpt-0001.pt"
torch.save(checkpoint, checkpoint_path)

この方法は、分散チェックポイントでも機能します。 次の例は、Torch Distributed Checkpoint API を使用した分散モデルのチェックポイント処理を示しています。

import torch.distributed.checkpoint as dcp

def save_checkpoint(self, checkpoint_path):
    state_dict = self.get_state_dict(model, optimizer)
    dcp.save(state_dict, checkpoint_id=checkpoint_path)

trainer.save_checkpoint("/Volumes/my_catalog/my_schema/model/checkpoints")

マルチユーザー コラボレーション

  • すべてのユーザーが共有コード (ヘルパー モジュールや環境 YAML ファイルなど) にアクセスできるようにするには、/Workspace/Sharedなどのユーザー固有のフォルダーではなく、/Workspace/Users/<your_email>/に格納します。
  • 開発中のコードの場合は、ユーザー固有のフォルダー内の Git フォルダー /Workspace/Users/<your_email>/ 使用し、リモート Git リポジトリにプッシュします。 これにより、複数のユーザーがユーザー固有のクローンとブランチを持つことができますが、バージョン管理にはリモート Git リポジトリを引き続き使用できます。 Databricks で Git を使用するための ベスト プラクティス を参照してください。
  • コラボレーターはノートブック を共有し、コメント することができます。

Databricks のグローバル制限

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

  • Last updated on