Slurm からの移行
プレビュー
この機能は パブリック プレビュー段階です。
このガイドでは、分散トレーニング ワークロードを Slurm から AI Runtime に移行する方法について説明します。主要な概念の説明、ランの監視と制御の方法の提示、および Slurm バッチ スクリプトの翻訳の解説を行います。
AI Runtime では、YAML ワークロードの構成で、GPU タイプ、総 GPU 数、環境、および各ノードで実行するコマンドを定義します。databricks air run --file train.yaml でその構成を送信すると、AI ランタイム はランのスケジューリング、プロビジョニング、およびリソースのクリーンアップを処理します。
開始する前に、Databricks CLI をインストールし、認証を設定し、クイックスタートに従ってください。
Slurm と AI ランタイム の概念
GPUリソース
Slurmでは、ジョブはパーティションからノードとGPUの割り当てを受け取ります。AI Runtime では、ランごとに GPU の種類と GPU の総数をリクエストします。defaultでは、サービスはオンデマンドでそのコンピュートをプロビジョニングします。
compute.num_accelerators を合計 GPU 数に設定します。compute.accelerator_type によって、各ノードが持つ GPU の数が決まります。たとえば、GPU_8xH100 はノードあたり 8基の H100 GPU を提供するため、16個の GPU を要求すると 2つのノードになります。合計は、ノードあたりの GPU の数の倍数である必要があります。
トレーニング プロセス
Slurm では、sbatch はノードと GPU を要求するジョブを送信し、srun は割り当てられたノードでタスクを起動します。分散 PyTorch トレーニングの一般的なパターンは、srun を使用してノードごとに 1 つの torchrun ランチャーを起動することです。各ランチャーは、GPU ごとに 1 つのトレーニングプロセスを開始します。
AI Runtime は、各ノードで同じ command を 1 回実行します。そのコマンドで torchrun を使用して、トレーニングプロセスを起動します。AI ランタイム は、各ノードのランクと、ノードが連携するために必要な接続情報を提供します。torchrun は、各トレーニングプロセスに RANK、LOCAL_RANK、および WORLD_SIZE を割り当てます。2 つのノードとノードあたり 8 個の GPU を使用する場合、これにより 2 つのランチャーと 16 のトレーニングプロセスが提供されます。
ストレージとリトライ
ランを維持する必要がある共有データセットとチェックポイントには、Unity Catalog ボリューム（/Volumes/<catalog>/<schema>/<volume>/...）を使用します。各ノードのローカルディスクを一時的な作業スペースとして扱います。
timeout_minutes は各試行を制限し、max_retries は失敗したワークロードが再試行される回数を制御します。再試行するたびにコマンドが再度実行されます。チェックポイントからトレーニングを継続するには、トレーニング コードで保存された状態を読み込む必要があります。チェックポイントとリカバリ パターンについては、「Improve training performance and resiliency on AI Runtime」を参照してください。
ランの監視と制御
Slurm | AI Runtime |
|---|---|
|
|
|
|
|
|
|
|
各ワークロードには、Logs と自動的に収集されたシステム メトリクスを持つ MLflow ランがあります。トレーニング コードで MLflow を使用して、パラメーター、トレーニング メトリクス、アーティファクトをログに記録します。MLflow とジョブ実行ページでのランの追跡を参照してください。
例: sbatch スクリプトを翻訳する
それぞれ 8 個の GPU を持つ 2 つのノード（合計 16 個の GPU）の Slurm 起動スクリプト:
#!/bin/bash
#SBATCH --job-name=llama-sft
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=8
#SBATCH --time=02:00:00
srun bash -c '
torchrun \
--nnodes="$SLURM_NNODES" \
--node_rank="$SLURM_NODEID" \
--nproc_per_node=8 \
--master_addr="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1)" \
--master_port=29500 \
train.py
'
対応する AI ランタイム train.yaml：
experiment_name: llama-sft
environment:
version: '4'
dependencies:
- transformers>=4.45
- datasets>=3.0
# 16 GPUs across 2 nodes (GPU_8xH100 = 8 H100 per node).
compute:
num_accelerators: 16
accelerator_type: GPU_8xH100
code_source:
type: snapshot
snapshot:
root_path: .
command: |
cd "$CODE_SOURCE_PATH"
# AI Runtime sets these rendezvous variables on each node.
torchrun \
--nnodes="$NUM_NODES" \
--node_rank="$NODE_RANK" \
--nproc_per_node="${LOCAL_WORLD_SIZE:-8}" \
--master_addr="$MASTER_ADDR" \
--master_port="$MASTER_PORT" \
train.py
timeout_minutes: 120
max_retries: 1
送信してフォローします：
databricks air run --file train.yaml --watch
トレーニング スクリプトを含む完全に実行可能なバージョンについては、FSDP を使用したマルチノード LLM ファインチューニングを参照してください。
移行に関する考慮事項
module loadおよび環境のアクティベーション手順をenvironment構成に置き換えます。管理対象環境のバージョンを選択し、environment.dependencies、インライン、またはrequirements.txtファイルへの-r参照を通じて追加のパッケージを宣言します。カスタム スタックの場合は、カスタム Docker イメージを使用します。code_source.snapshot.root_pathをローカルプロジェクトディレクトリに設定します。ワークロードを送信すると、Databricks CLI がそれをuploadします。$CODE_SOURCE_PATHを使用してuploadされたファイルを参照します。- 環境変数には
env_variablesを使用し、Databricks シークレット参照にはsecretsを使用します。 - 各構成に対して個別のランを送信し、
databricks air run --override key=valueを使用して送信間でフィールドを変更します。