実験の追跡と可観測性

AI ランタイムは、実験追跡のために MLflow とネイティブに統合され、使用率、メモリ、温度を監視するための組み込みの GPU リソース ウィンドウが含まれています。 MLflow を使用してメトリックと実行をログに記録し、ノートブックと MLflow UI でトレーニング出力を表示し、モデル チェックポイントを Unity カタログ ボリュームに保存し、コードの実行中に GPU の正常性を追跡します。

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 を使用する場合、 .distributed() を呼び出すたびに、MLflow 実験の実行が自動的に作成されます。 アクティブな MLflow 実行内で呼び出された場合は、その代わりに、アクティブな親実行の下にネストされた子実行が作成されます。

  import mlflow
  
  with mlflow.start_run() as outer_run:
      ...
      run_train.distributed()  # creates a nested child run under outer_run
  

• .distributed()で使用される実験をカスタマイズするには、mlflow.set_experiment()を呼び出す前に.distributed()を呼び出すか、MLFLOW_EXPERIMENT_NAME環境変数を設定します。 既定の実験名は /Users/{WORKSPACE_USER}/{notebook-name}です。 常に絶対パスを使用します。

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

  Alternatively:

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

• 前の MLflow 実行を再開するには、 mlflow.start_run(run_id="<previous-run-id>")を使用します。

• .distributed()を使用して以前の MLflow 実行を再開するには、呼び出す前にMLFLOW_RUN_ID設定します。

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

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

ログの表示

• ノートブックの出力: トレーニング コードからの標準出力とエラーがノートブック のセル出力に表示されます。
• MLflow ログ: MLflow 実験 UI には、トレーニング メトリック、パラメーター、成果物が表示されます。

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

分散トレーニングの場合は、モデル のチェックポイントを Unity カタログ ボリュームに非同期的に保存して読み込みます。これは、他の Unity カタログ オブジェクトと同じガバナンスを提供します。 UCVolumeWriter (DCP) API で、UCVolumeReader パッケージの serverless_gpu.data と を使用します。 これらのストレージ バックエンドは、高速なローカル ディレクトリ (サーバーレス GPU ノードで NVMe でサポートされている /tmp) を介してすべての I/O をステージングし、Unity カタログ ボリュームにアップロードまたはダウンロードします。これは、チェックポイント シャードを FUSE マウントに直接書き込むよりも高速です。 メタデータのアトミック性は保持されます。ライターは、データ シャードのアップロードが完了した後にのみ、 .metadata ファイルを発行します。

中断時の作業損失を最小限に抑えられるよう、チェックポイントは十分な頻度で作成してください。ただし、I/O オーバーヘッドでトレーニングが遅くなるほど頻繁に作成してはいけません。 30 分から 1 時間ごとに 1 つのチェックポイントを目指し、ステップの時間とチェックポイント のサイズに基づいて間隔を調整します。

トレーニングの続行中にバックグラウンドでチェックポイントをアップロードするには、UCVolumeWriterにstorage_writerとしてdcp.async_saveを渡します。 非同期保存では、プロセス グループに CPU バックエンドが必要であるため、 torch.distributed.init_process_group(backend="cpu:gloo,cuda:nccl", ...)で初期化します。

import torch.distributed.checkpoint as dcp
from serverless_gpu.data import UCVolumeWriter

checkpoint_path = "/Volumes/my_catalog/my_schema/model/checkpoints"
writer = UCVolumeWriter(checkpoint_path)

future = dcp.async_save(state_dict, storage_writer=writer)
# ...continue training...
future.result()  # blocks until the upload lands on the UC volume

UCVolumeReaderを使用してチェックポイントを読み込みます。

from serverless_gpu.data import UCVolumeReader

reader = UCVolumeReader(checkpoint_path)
dcp.load(state_dict, storage_reader=reader)

データ パイプラインのチェックポイント処理

モデル チェックポイントはモデルとオプティマイザーの状態をキャプチャしますが、データセット内のデータ パイプラインの位置をキャプチャしないため、再開された実行を停止した正確なサンプルに早送りすることはできません。 再開時にはこの点を考慮し、エポック境界から再開するか、処理済みのサンプルまたはシャードを独自の学習状態に記録しておき、再開時にそれらをスキップできるようにします。

GPU リソースの監視

コードが AI ランタイムで実行されている間に GPU の正常性と使用率を監視するには、[ GPU リソース ] ウィンドウを使用します。 このウィンドウでは、単一ノードワークロードとマルチノード ワークロードの両方がサポートされます。

ウィンドウを開くには、ノートブックを AI ランタイムに接続し、[をクリックします。右側のウィンドウの GPU リソース。

このウィンドウには、GPU ごとに次のメトリックが表示されます。

• GPU 使用率
• GPU メモリ使用量
• 気温

このウィンドウでは、10 秒ごとにメトリックがポーリングされ、最大 2 時間の履歴が保持されます。 [最新の値をすぐにフェッチするには、更新します。 非アクティブ状態が 5 分続くと、ウィンドウは一時停止します。を再度開いて監視を再開します。

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

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

Azure Databricksのグローバル制限

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