ノートブックでの分散トレーニング

databricks-sdk-air（@distributed）のデコレーターは、Databricksノートブックから分散トレーニングを実行する最も便利な方法です。 トレーニング関数を装飾して呼び出すと、AI Runtimeがノートブックに接続されているノード上のすべてのGPUで実行します。 同じコードは、クラスタプロビジョニングや分散ランチャーの設定なしに、シングルGPUからマルチGPUへとスケールします。

このページの重要なポイント:

• @distributedデコレーターはノートパソコンの中からノード上のすべてのGPUに対してトレーニング機能を実行します。
• PyTorch DDP、FSDP、DeepSpeedをサポートし、最小限の変更でシングルGPUコードをマルチGPUに移行します。
• ノートパソコンを 8xH100 アクセラレータに接続して、マルチGPUによる本格的なトレーニングを行ってください。 @distributed はデフォルトでそのノード上のすべてのGPUを使用します。

クイック スタート

このページの例は、GPU環境バージョン6以上にプリインストールされているDatabricks AI Runtime SDKを使用しています。 バージョン5以前では、代わりに serverless_gpu パッケージを使いましょう。 詳細は「serverless_gpuからの移行」を参照してください。 ノートブックをアクセラレーターに接続するには、「ノートブックに接続する」をご覧ください。

トレーニング機能を @distributedで飾り、 .distributed()と呼びます。

from databricks.air import distributed

# By default, @distributed uses every GPU on the node your notebook is connected
# to. To use fewer, pass num_accelerators (for example, num_accelerators=2).
@distributed()
def train():
    import os
    import torch
    import torch.distributed as dist

    # Bind this process to its own GPU before training.
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")
    dist.init_process_group("nccl")
    # ... build the model and data on `device`, then run your training loop ...
    dist.destroy_process_group()

train.distributed()

分散トレーニングには 8xH100 アクセラレータが必要です。このアクセラレータでは、8 つの GPU を持つ単一のノードがプロビジョニングされます。 デフォルトでは、 @distributed 8枚すべてのGPUを使用しています。 num_accelerators を少なく使用するように設定しましょう。 accelerator_typeパラメータもオプションで、ノートブックに接続されているアクセラレータから自動的に検出されます。

.distributed() を呼び出すたびに、MLflow ランが作成されます（すでにアクティブなランがある場合は、ネストされた子ランが作成されます）。また、セルの出力にランへのリンクが表示されます。 完全で実行可能なウォークスルーについては、 Full例をご覧ください。

サポートされているフレームワーク

@distributed API は、主要な分散トレーニング ライブラリと統合されます。

• PyTorch 分散データ並列 (DDP): 標準のマルチ GPU データ並列処理。
• 完全シャーディングデータ並列 (FSDP): 大規模モデル向けのメモリ効率に優れたトレーニング。
• DeepSpeed: 大規模なモデル トレーニング用のMicrosoftの最適化ライブラリ。

各ライブラリを使った実際のトレーニングシナリオについては、 分散トレーニングの例を参照してください。

@distributedデコレーターの仕組み

.distributed()で装飾関数を呼び出すと、AI Runtimeは分散ランチャーで手動で設定するような仕組みを処理します。

• シリアライズとファンアウト: この関数はシリアライズされ、選択した各 GPU で起動されます。 すべてのGPUは同じ引数で関数のコピーを実行します。
• 環境同期:Python環境と依存関係はすべてのランクで複製されているため、すべてのプロセスが同じコードを実行します。
• 環境変数のランク付け: LOCAL_RANK などの標準変数が各プロセスごとに入力されます。 関数でそれらを読み、モデルとデータを正しいデバイスに配置してください。
• 結果収集:すべてのランクから返される値は収集され、コールプレイヤーに返されます。
• MLflowのトラッキング:各 .distributed() 呼び出しはMLflowランを生成し、すでにアクティブな場合はネストされた子ランを作成します。これにより、関数から記録された指標が同じ実行に集まります。
• ライフサイクルとタイムアウト:分散実行はノートブックのライフサイクル内で実行されます。 ノートブックを終了するとランは終了します。 デコレーターのデフォルトのタイムアウトは3時間です。 数秒で timeout をパスして変更するか、 timeout=None で無効にしてください。 カスタムタイムアウトにはGPU環境v5以上が必要です。

APIは標準的なPyTorchライブラリである Distributed Data Parallel (DDP)、 Fully Sharded Data Parallel (FSDP)、 DeepSpeedを基盤としています。

serverless_gpuから移行

serverless_gpuパッケージは廃止され、新機能も提供されていません。 GPU環境バージョン6からは、Databricks AI Runtime SDK（databricks-sdk-air）がそれを置き換え、databricks.air 名前空間の下で引数を変更した同等のAPIを提供します。 使用するパッケージはGPU環境のバージョンによって異なります:

• バージョン5以下: serverless_gpuのみ含めてください。 これらのバージョンでは引き続き serverless_gpu を使い続けてください。
• バージョン6:両方のパッケージを含みますが、アップデートはAI Runtime SDKのみです。 バージョン6が serverless_gpuを含む最後のGPU環境バージョンです。

AI Runtime SDKはDatabricks SDK for Python(databricks-sdk)とは別のパッケージです。

ノートブックの移行方法:

1. ノートブックをAI Runtime SDKを含むGPU環境バージョン6以上に接続します。
2. インポートを serverless_gpu から databricks.airに変更してください。
3. gpusをnum_acceleratorsに、gpu_typeをaccelerator_typeにリネームし、新しいアクセラレータタイプの値を使いましょう。 古い引数名や GPU タイプの値は databricks.air で引き続き動作しますが、FutureWarning が表示されます。 GPU環境バージョン7でサポートは終了します。

Migrate this notebook from the serverless_gpu SDK to the AI Runtime SDK.

databricks.air.dataを使うDatabricksのCLIワークロードでは、ワークロードのdatabricks-sdk-air[data]にenvironment.dependenciesを追加します。

TorchDistributorより

Sparkクラスターで TorchDistributor を使って分散トレーニングを実行し、ワークロードが1つのノードに収まる場合、@distributed API は AI Runtime で推奨される代替手段です。 DDP、FSDP、DeepSpeedのモデルコードは変わりません。

単一ノードワークロードの移行:

• TorchDistributor(...).run(train_fn, ...)コールをtrain_fnの@distributedデコレーターに置き換え、その後train_fn.distributed(...)で起動します。
• SparkクラスターとGPUワーカーの設定を削除してください。 代わりにノートパソコンを8xH100のアクセラレーターに接続してください。 @distributed デフォルトで8つのGPUすべてを使用します。
• デコレータで修飾された関数内にデータの読み込みを移動してください。 データの 読み込みを参照してください。
• 既存のDDP、FSDP、またはDeepSpeedモデルのコードを保持してください。 デコレーターはこの3つすべてをサポートします。

@distributed 単一のノード上で動作するため( 制限を参照)、すべてのTorchDistributorワークロードを置き換えるわけではありません。 Spark連携に依存するワークロードはTorchDistributorに保管してください。 ローカルマシンや複数のノード間で分散トレーニングを行う場合は、 代わりにパブリックプレビューにあるAI Runtime CLIを使います。 「 Data bricks CLIとAIランタイムを使いましょう」を参照してください。

完全な例

以下の例は、ノートパソコンから8台のH100 GPU上でマルチレイヤーパーセプトロン(MLP)モデルを訓練しています。

1. モデルを設定し、ユーティリティ関数を定義します。

   
   # Define the model
   import os
   import torch
   import torch.distributed as dist
   import torch.nn as nn
   
   def setup():
       torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
       dist.init_process_group("nccl")
   
   def cleanup():
       dist.destroy_process_group()
   
   class SimpleMLP(nn.Module):
       def __init__(self, input_dim=10, hidden_dim=64, output_dim=1):
           super().__init__()
           self.net = nn.Sequential(
               nn.Linear(input_dim, hidden_dim),
               nn.ReLU(),
               nn.Dropout(0.2),
               nn.Linear(hidden_dim, hidden_dim),
               nn.ReLU(),
               nn.Dropout(0.2),
               nn.Linear(hidden_dim, output_dim)
           )
   
       def forward(self, x):
           return self.net(x)
   

2. AI Runtime SDKから distributed デコレーターをインポートします。

   from databricks.air import distributed
   

3. モデルトレーニングコードを関数でラップし、 @distributed デコレーターで関数を装飾します。 装飾関数は分散実行のエントリーポイントなので、すべてのトレーニングロジック、データ読み込み、モデル初期化を定義してください。

   @distributed()
   def run_train(num_epochs: int, batch_size: int) -> None:
       import mlflow
       import torch.optim as optim
       from torch.nn.parallel import DistributedDataParallel as DDP
       from torch.utils.data import DataLoader, DistributedSampler, TensorDataset
   
       # 1. Set up multi-GPU environment
       setup()
       device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
   
       # 2. Apply the Torch distributed data parallel (DDP) library for data-parallel training.
       model = SimpleMLP().to(device)
       model = DDP(model, device_ids=[device])
   
       # 3. Create and load dataset.
       x = torch.randn(5000, 10)
       y = torch.randn(5000, 1)
   
       dataset = TensorDataset(x, y)
       sampler = DistributedSampler(dataset)
       dataloader = DataLoader(dataset, sampler=sampler, batch_size=batch_size)
   
       # 4. Define the training loop.
       optimizer = optim.Adam(model.parameters(), lr=0.001)
       loss_fn = nn.MSELoss()
   
       for epoch in range(num_epochs):
           sampler.set_epoch(epoch)
           model.train()
           total_loss = 0.0
           for step, (xb, yb) in enumerate(dataloader):
               xb, yb = xb.to(device), yb.to(device)
               optimizer.zero_grad()
               loss = loss_fn(model(xb), yb)
               # Log loss to MLflow metric
               mlflow.log_metric("loss", loss.item(), step=step)
   
               loss.backward()
               optimizer.step()
               total_loss += loss.item() * xb.size(0)
   
           mlflow.log_metric("total_loss", total_loss)
           print(f"Total loss for epoch {epoch}: {total_loss}")
   
       cleanup()
   

4. 分散トレーニングはユーザー定義の引数で分散関数を呼び出して実行します。

   run_train.distributed(num_epochs=3, batch_size=1)
   

5. 実行すると、ノートブック のセル出力に MLflow 実行リンクが生成されます。 MLflow の実行リンクをクリックするか、 実験 パネルで見つけて実行結果を表示します。 実験名のカスタマイズ、メトリックの追跡、実行の再開の詳細については、「 実験の追跡と可観測性」を参照してください。

データ読み込み

@distributed関数内にデータロードコードを配置します。 データセットは pickleで許容される最大サイズを超えることがあるため、デコレーター内で生成または読み込むことでシリアライズエラーを避けられます:

from databricks.air import distributed

# Don't do this: loading the dataset outside the function means it's captured
# and pickled with the function, which can raise a pickle error.
# dataset = get_dataset(file_path)

@distributed()
def run_train():
    # Load the dataset inside the decorated function instead.
    dataset = get_dataset(file_path)
    ...

Unity カタログ ボリュームに格納されているファイル ベースのデータの場合は、UCVolumeDatasetからのdatabricks.air.dataを使用します。このは、ローカル キャッシュを使用してファイルをストリーミングし、それらをランクとワーカー間で自動的にパーティション分割します。 分散トレーニングのチェックポイントをボリュームに保存するには、UCVolumeWriter と UCVolumeReader を使用します。 AI Runtime でのデータのロードおよびDistributed Checkpoint (DCP) によるチェックポイントを参照してください。

Limitations

• @distributed 1つのノードのみをサポートしている:ノートブックが接続されているノードのGPUをまたぐ動作です。 複数のノードをまたいで訓練するには、AI Runtime Databricks CLIをご利用ください。 「 Data bricks CLIとAIランタイムを使いましょう」を参照してください。
• accelerator_typeを設定する場合、ノートブックが接続されているアクセラレーター（例えば、"GPU_8xH100"）と一致しなければなりません。 不一致が作業負荷の失敗を引き起こします。 パラメータは任意で、省略すると自動検出されます。
• AIランタイムはGPU環境v4以上をサポートしています。 カスタムタイムアウト( timeout パラメータ)はGPU環境v5以上が必要です。
• デコレーターはデフォルトで3時間後にタイムアウトします。 数秒で timeout をパスして変更するか、 timeout=None で無効にしてください。
• 実行はノートブックのライフサイクル内で行われます。 ノートブックを終了するとランは終了します。

その他のリソース

• @distributedデコレーター、AcceleratorType、Ray APIについては、Databricks AI Runtime SDK APIリファレンスを参照してください。
• トレーニングパイプラインをより効率的かつ回復力のあるものにするパターンについては、 パフォーマンスとレジリエンスのガイドをご覧ください。
• エンドツーエンドのトレーニングシナリオについては、 分散トレーニングの例を参照してください。