メインコンテンツまでスキップ
に最終更新

経験の追跡と観察可能性

この記事では、サーバレス GPU コンピュートでMLflowを使用し、GPU の状態を監視し、ログを表示し、モデル チェックポイントを管理する方法について説明します。

MLflowとの連携​

サーバーレス GPU コンピュートは、体験追跡、モデル ロギング、メトリクス視覚化のためにMLflowとネイティブに統合します。

設定に関する推奨事項：

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

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

    Python
    import mlflow
    mlflow.pytorch.autolog()

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

    Python
    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には常に絶対パスを使用してください。

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

  • 以前の実行結果からMLFLOW_RUN_IDを設定することで、以前のトレーニングを再開します。

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

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

GPU 健全性モニタリング​

Databricksノートブック UI を通じて GPU の使用率と状態を直接監視します。 サーバレス GPU コンピュートに接続すると、コンピュート パネルで GPU メトリクス (使用率、メモリ使用量、温度) が表示されます。

閲覧ログ​

  • ノートブックの出力 — 標準出力とトレーニング コードからのエラーがノートブックのセル出力に表示されます。
  • ドライバー ログ — デバッグの起動問題、環境セットアップの問題、ランタイム エラーについては、コンピュート パネルからアクセスできます。
  • MLflowログ — トレーニング メトリクス、問題、およびアーティファクトは、 MLflowエクスペリメント UI で表示できます。

モデルチェックポイント​

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

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

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

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

Python
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 を使用した分散モデルのチェックポイント処理を示しています。

Python
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/Users/<your_email>/のようなユーザー固有のフォルダではなく、 /Workspace/Sharedに保存してください。
  • 現在開発中のコードについては、ユーザー固有のフォルダ/Workspace/Users/<your_email>/内の Git フォルダを使用し、リモート Git リポジトリにプッシュしてください。これにより、複数のユーザーがユーザー固有のクローンとブランチを持つことができ、同時にバージョン管理にはリモートのGitリポジトリを使用できます。 DatabricksでGitを使用する際のベストプラクティスを参照してください。
  • 共同作業者はノートブックを共有したり、コメントしたりできます。

Databricksのグローバル制限​

リソース制限を参照してください。