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

備考

この機能はベータ版です。

@distributed デコレータは、Databricks AI ランタイム SDK (databricks-sdk-air) の Databricks ノートブックから分散トレーニングを実行するための最も便利な方法です。トレーニングの関数にデコレータを追加し、呼び出します。AI ランタイムはノートブックが接続されているノードのすべてのGPUにわたってそれを実行します。プロビジョニングするクラスターや構成する分散ランチャーなしに、同じコードが単一GPUからマルチGPUへとスケールします。

このページのポイント：

• @distributedデコレータは、ノートブック内からノードのすべてのGPUでトレーニング関数を実行します。
• これはPyTorch DDP、FSDP、DeepSpeedをサポートしており、最小限の変更で単一GPUコードをマルチGPUに移行します。
• ノートブックを 8xH100 アクセラレータに接続し、完全なマルチ GPU トレーニングを行います。defaultでは、 @distributedはノード上のすべてのGPUを使用します。

注記

このページでは、Databricks AI ランタイム SDKを使用したDatabricksノートブックからの分散トレーニングについて説明します。ローカルマシンから分散トレーニングワークロードを送信するには、AI ランタイムのDatabricks CLIコマンドを使用します。AI ランタイムでDatabricks CLIを使用するを参照してください。

クイックスタート

このページの例では、GPU 環境バージョン 6 以上にプリインストールされている Databricks AI ランタイム SDK を使用しています。バージョン 5 以下の場合は、代わりに serverless_gpu パッケージを使用してください。serverless_gpuからの移行を参照してください。ノートブックをアクセラレータに接続する方法については、ステップ1: ノートブックをGPUにアタッチするを参照してください。

トレーニング関数に@distributedをデコレートし、.distributed()で呼び出します。

Python

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または8xB300アクセラレータが必要であり、これにより8個のGPUを持つ単一ノードがプロビジョニングされます。defaultでは、@distributed は 8 つの GPU すべてを使用します。使用する数を減らすには num_accelerators を設定します。accelerator_typeパラメーターもオプションであり、ノートブックが接続されているアクセラレータから自動的に検出されます。

各.distributed()呼び出しは、MLflowラン（またはすでにアクティブな場合はネストされた子ラン）を作成し、セル出力にランLinkを出力します。完全に実行可能なチュートリアルについては、「 完全な例」を参照してください。

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

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

• PyTorch 分散データ並列 （DDP）：標準的なマルチGPUデータ並列処理。
• 完全シャーディングデータパラレル (FSDP) : 大規模モデルのメモリ効率の高いトレーニング。
• DeepSpeed : 大規模モデルトレーニングのためのMicrosoftの最適化ライブラリ。

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

@distributed デコレータの仕組み

装飾された関数を .distributed() で呼び出すと、AI ランタイムは分散ランチャーで手動で設定するメカニズムを処理します:

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

APIは、標準のPyTorchライブラリ：分散データ並列（DDP）完全シャーディングデータパラレル（FSDP）、DeepSpeedに基づいています。

からの移行 serverless_gpu

serverless_gpu パッケージは非推奨であり、新機能は追加されません。GPU 環境バージョン 6 以降では、Databricks AI ランタイム SDK (databricks-sdk-air) に置き換わり、名前が変更された引数を持つ同等の APIs が databricks.air 名前空間の下で提供されます。使用するパッケージは GPU 環境バージョンによって異なります。

• バージョン 5 以下: serverless_gpu のみを含めます。これらのバージョンで serverless_gpu の使用を続けます。
• バージョン6: 両方のパッケージが含まれますが、AI ランタイム SDKのみが更新を受け取ります。バージョン6はserverless_gpuが含まれる最後のGPU環境バージョンです。

AIランタイムSDKは、Databricks SDK for Python (databricks-sdk) とは別のパッケージです。

	serverless_gpu （非推奨）	AI ランタイム SDK (databricks.air)
GPU環境バージョン	6以下	6以上
インポート	from serverless_gpu import distributed, ray_init	from databricks.air import distributed, ray_init
GPU 数引数	gpus	num_accelerators
GPUタイプの引数	gpu_type (GPUType)	accelerator_type (AcceleratorType)
GPUタイプ値	"A10"、"H100"、 "B300"	"GPU_1xA10"、"GPU_1xH100"、"GPU_8xH100"、 "GPU_8xB300"
データの読み込みとチェックポイント モジュール	serverless_gpu.data	databricks.air.data。AI ランタイムへのデータの読み込みを参照してください。

	serverless_gpu （非推奨）	AI ランタイム SDK (databricks.air)
GPU環境バージョン	6以下	6以上
インポート	from serverless_gpu import distributed, ray_init	from databricks.air import distributed, ray_init
GPU 数引数	gpus	num_accelerators
GPUタイプの引数	gpu_type (GPUType)	accelerator_type (AcceleratorType)
GPUタイプ値	"A10"、"H100"、 "B300"	"GPU_1xA10"、"GPU_1xH100"、"GPU_8xH100"、 "GPU_8xB300"
データの読み込みとチェックポイント モジュール	serverless_gpu.data	databricks.air.data。AI ランタイムへのデータの読み込みを参照してください。

ノートブックを移行するには：

1. ノートブックを、AIランタイムSDKが含まれているGPU環境バージョン6以降に接続します。
2. インポートをserverless_gpuからdatabricks.airに変更します。
3. gpusをnum_acceleratorsに、gpu_typeをaccelerator_typeに名前変更し、新しいアクセラレータタイプの値を使用してください。古い引数名とGPUタイプの値はdatabricks.airで引き続き機能しますが、FutureWarningを出力します。サポートは GPU 環境バージョン 7 で終了します。

プロンプト

これを行うように Genie Code に伝えてください：

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

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

TorchDistributorから

SparkクラスターでTorchDistributorを使用して分散トレーニングを実行し、ワークロードが単一ノードに収まる場合は、AIランタイムで@distributed APIを使用することをお勧めします。DDP、FSDP、またはDeepSpeedモデルコードは変更されません。

機能	@distributed API	TorchDistributor
インフラストラクチャ	完全にServerless、クラスター管理不要	GPU ワーカーを持つ Spark クラスターが必要です
設定	単一のデコレータ、最小限の構成	SparkクラスターとTorchDistributorのセットアップが必要です
フレームワークのサポート	PyTorch DDP、FSDP、DeepSpeed	主に PyTorch DDP
データの読み込み	デコレーター内で、Unityカタログボリューム (UCVolumeDatasetストリーミングファイルデータの場合) を使用します	Spark またはファイルシステム経由

機能	@distributed API	TorchDistributor
インフラストラクチャ	完全にServerless、クラスター管理不要	GPU ワーカーを持つ Spark クラスターが必要です
設定	単一のデコレータ、最小限の構成	SparkクラスターとTorchDistributorのセットアップが必要です
フレームワークのサポート	PyTorch DDP、FSDP、DeepSpeed	主に PyTorch DDP
データの読み込み	デコレーター内で、Unityカタログボリューム (UCVolumeDatasetストリーミングファイルデータの場合) を使用します	Spark またはファイルシステム経由

単一ノードのワークロードを移行するには：

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

@distributed デコレーターは ベータ版 であり、シングル ノードで実行されるため (制限 を参照)、すべての TorchDistributor ワークロードを置き換えるわけではありません。Spark 統合に依存するワークロードは TorchDistributor に保持します。ローカル マシンから、または複数のノード間で分散トレーニングを実行するには、AI ランタイム の Databricks CLI コマンドを使用します。AI ランタイム で Databricks CLI を使用するを参照してください。

完全な例

次の例では、ノートブックで8 つの H100 GPU で多層パーセプトロン (MLP) モデルをトレーニングします。

1. モデルのセットアップとユーティリティ関数の定義を行います。

   Python

   
   # 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 ランタイム SDK から distributed デコレータをインポートします。

   Python

   from databricks.air import distributed

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

   Python

   @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. ユーザー定義引数で分散関数を呼び出すことで、分散トレーニングをランします。

   Python

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

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

データ読み込み

データロードコードは@distributed関数の中に配置してください。データセットがpickleで許可されている最大サイズを超える可能性があるため、デコレーター内で生成またはロードすることでシリアル化エラーを回避します。

Python

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 Catalog ボリュームに保存されたファイルベースのデータの場合、 databricks.air.data の UCVolumeDataset を使用します。これは、ローカルキャッシュを使用してファイルをストリームし、ランクおよびワーカー全体に自動的にパーティション分割します。分散トレーニングのチェックポイントをボリュームに保存するには、UCVolumeWriter と UCVolumeReader を使用します。AI ランタイムへのデータの読み込みおよびDistributed Checkpoint (DCP)を使用したチェックポイントを参照してください。

制限

• @distributed 単一ノードのみをサポートします：ノートブックが接続されているノード上のGPUで実行されます。複数ノードでトレーニングするには、AIランタイムDatabricks CLIを使用します。AIランタイムでDatabricks CLIを使用するを参照してください。
• accelerator_type を設定する場合、ノートブックが接続されているアクセラレーターと一致している必要があります（例: "GPU_8xH100"）。不一致が発生すると、ワークロードに失敗します。このパラメーターはオプションで、省略すると自動的に検出されます。
• AIランタイムは、GPU環境v4以降をサポートしています。カスタムタイムアウト (timeoutパラメーター) には、GPU環境v5以降が必要です。
• デコレーターは、defaultで 3 時間後にタイムアウトします。timeoutを秒単位で渡して変更するか、 timeout=Noneを渡して無効にします。
• 実行はノートブックのライフサイクル内で実行されます。ノートブックを終了すると、ランも終了します。

その他のリソース

• @distributed デコレータ、 AcceleratorType、および Ray APIの詳細については、 Databricks AI ランタイム SDK API リファレンスを参照してください。
• トレーニングパイプラインの効率と回復力を向上させるパターンについては、「パフォーマンスと回復力のガイド」を参照してください。
• エンドツーエンドのトレーニングシナリオについては、 分散トレーニングの例を参照してください。