サーバレス GPU コンピュートのベスト プラクティス

この記事では、ノートブックとジョブでサーバレス GPU コンピュートを使用するためのベスト プラクティスの推奨事項を紹介します。

これらの推奨事項に従うことで、Databricks 上のワークロードの生産性、コスト効率、信頼性が向上します。

正しいコンピュートを使用する

• サーバレスGPUコンピュートを使用します。 このオプションには、互換性のために最適化された torch、cuda、torchvision が付属しています。正確なパッケージ バージョンは環境のバージョンによって異なります。

• 環境サイドパネルでアクセラレータを選択します。

  • リモート分散トレーニング ワークロードの場合は、後でリモート H100 にジョブを送信するクライアントとなる A10 GPU を使用します。
  • ノートブック自体で大規模なインタラクティブ ジョブを実行するには、ノートブックを H100 に接続して、1 つのノード (8 つの H100 GPU) を占有します。

• GPU の使用を回避するには、git clone や Spark Dataframe から Mosaic Data Shard (MDS) 形式への変換などの操作のためにノートブックを CPU クラスターに接続することができます。

MLflowの推奨事項

最適な ML 開発サイクルを実現するには、Databricks で MLflow 3 を使用します。次のヒントに従ってください。

• 環境のMLflowバージョン 3.6 以降にアップグレードし、 MLflow 3ディープラーニング ワークフローのMLflowディープラーニング フローに従います。

• MLFlowLoggerのstepパラメーターを適切なバッチ数に設定します。 MLflow 、ログに記録できるメトリクス ステップは 1,000 万件に制限されています。 リソース制限を参照してください。

• PyTorch Lightning をトレーナーとして使用する場合は、 mlflow.pytorch.autolog()を有効にします。

• モデル トレーニング コードをmlflow.start_run() API スコープ内にカプセル化して、MLflow 実行名をカスタマイズします。これにより、実行名を制御でき、以前の実行から再開できるようになります。 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 、システム メトリクスをログに記録するためのMLflowエクスペリメントを起動します。 デフォルトでは、ユーザーが環境変数MLFLOW_EXPERIMENT_NAMEで上書きしない限り、名前/Users/{WORKSPACE_USER}/{get_notebook_name()}が使用されます。

  • MLFLOW_EXPERIMENT_NAME環境変数を設定するときは、絶対パスを使用します。たとえば、 /Users/<username>/my-experiment 。
  • エクスペリメント名には既存のフォルダー名を含めることはできません。 たとえば、 my-experiment既存のフォルダーである場合、上記の例ではエラーが発生します。
  Python

  import os
  from serverless_gpu import distributed
  os.environ['MLFLOW_EXPERIMENT_NAME'] = '/Users/{WORKSPACE_USER}/my_experiment'
  @distributed(gpus=num_gpus, gpu_type=gpu_type, remote=True)
  def run_train():
  # my training code

• 前回の実行からトレーニングを再開するには、次のように前回の実行の MLFLOW_RUN_ID を指定します。

  Python

  import os
  os.environ[‘MLFLOW_RUN_ID’] = <previous_run_id>
  run_train.distributed()

モデルチェックポイント

モデル チェックポイントをUnity Catalog ボリュームに保存します。これにより、他の 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")

複数ユーザーによるコラボレーション

• すべてのユーザーが共有コード（ヘルパーモジュール、environment.yamlなど）にアクセスできるようにするには、/Workspace/Users/<your_email>/のようなユーザー固有のフォルダーの代わりに、 /Workspace/Reposまたは/Workspace/Sharedに git フォルダーを作成します。
• アクティブに開発中のコードの場合は、ユーザー固有のフォルダー/Workspace/Users/<your_email>/内の Git フォルダーを使用し、リモート Git リポジトリにプッシュします。これにより、複数のユーザーがユーザー固有のクローン (およびブランチ) を持つことができ、バージョン管理にはリモート Git リポジトリを引き続き使用できます。Databricks で Git を使用するためのベスト プラクティスを参照してください。
• 共同作業者はノートブックを共有したりコメントしたりできます。

Databricksのグローバル制限

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