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

備考

ベータ版

この機能はベータ版です。ワークスペース管理者は、 プレビュー ページからこの機能へのアクセスを制御できます。「Databricks プレビューの管理」を参照してください。

Serverless GPU Python APIの @distributed デコレータは、Databricksノートブックから分散トレーニングを実行するための最も便利な方法です。トレーニング関数をデコレータで装飾して呼び出すと、AIランタイムが、ノートブックが接続されているノード上のすべてのGPUでそれを実行します。クラスターのプロビジョニングや分散ランチャーの構成を行うことなく、同じコードを単一GPUからマルチGPUへとスケールさせることができます。

ヒント

主なポイント

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

クイックスタート

ノートブックがServerless GPU に接続されている場合、serverless_gpu パッケージがプレインストールされます。トレーニング関数を @distributed でデコレートし、.distributed() で呼び出します。

Python

from serverless_gpu import distributed

# gpus is the number of GPUs on the node. gpu_type is optional and
# auto-detected from the accelerator your notebook is connected to.
@distributed(gpus=8, gpu_type="H100")
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()

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

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

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

• PyTorch 分散データ並列 (DDP) : 標準的なマルチ GPU データ並列処理です。
• Fully Sharded Data Parallel (FSDP) ：大規模モデル向けのメモリ効率の高いトレーニング。
• DeepSpeed : Microsoft の大規模モデルトレーニング用最適化ライブラリ。

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

@distributed デコレータの仕組み

.distributed()AI ランタイム でデコレータ付き関数を呼び出すと、分散ランチャーを使用して手動で構成するメカニズムが AI ランタイム によって処理されます。

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

この API は、標準の PyTorch ライブラリ（Distributed Data Parallel (DDP)、Fully Sharded Data Parallel (FSDP)、DeepSpeed）を基盤として構築されています。

TorchDistributor からの移行

現在 Spark 上で TorchDistributor を使用して分散 PyTorch を実行しており、ワークロードが単一のノードに収まる場合、新しいディープラーニングのワークロードには serverless_gpu @distributed API を使用することをお勧めします。Spark クラスターを削除し、単一 GPU からマルチ GPU まで同じコードパスを提供します。

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

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

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

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

@distributed は単一ノードで実行されるため（制限事項を参照）、すべての TorchDistributor ワークロードの代わりになるわけではありません。Spark 統合に依存するワークロードは、TorchDistributor 上に維持します。ローカルマシンまたは複数のノードにわたって分散トレーニングを実行するには、代わりに パブリックプレビュー 段階にある AI Runtime CLI を使用します。AI ランタイム 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. serverless_gpu ライブラリと distributed モジュールをインポートします。

   Python

   import serverless_gpu
   from serverless_gpu import distributed

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

   Python

   @distributed(gpus=8, gpu_type='H100')
   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-parellel 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 実行リンクが生成されます。MLflow 実行リンクをクリックするか、「 エクスペリメント 」パネルで探して実行結果を表示します。エクスペリメント名のカスタマイズ、メトリクスの追跡、実行の再開の詳細については、エクスペリメントの追跡とオブザーバビリティを参照してください。

データ読み込み

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

Python

from serverless_gpu import distributed

# This may cause a pickle error because the dataset is captured by the function.
dataset = get_dataset(file_path)

@distributed(gpus=8, gpu_type='H100')
def run_train():
    # Load the dataset inside the decorated function instead.
    dataset = get_dataset(file_path)
    ...

Unity Catalog ボリュームに保存されているファイルベースのデータには、ローカルキャッシングを使用してファイルをストリームし、ランクとワーカーに自動的にパーティション分割する serverless_gpu.data からの UCVolumeDataset を使用します。分散トレーニングをボリュームにチェックポイントするには、UCVolumeWriter と UCVolumeReader を使用します。AI Runtimeへのデータの読み込みおよびモデルのチェックポイント処理を参照してください。

制限事項

• 分散トレーニングは、ノートブックが接続されている単一ノード上のGPU全体で実行されます。完全なマルチGPUトレーニングを行うには、8個のGPUを持つ1つのノードをプロビジョニングする 8xH100 アクセラレータに接続し、gpus=8 を設定します。
• アクセラレータタイプが一致している必要があります。@distributed で gpu_type を設定する場合、ノートブックが接続されているアクセラレータ（"H100" または "A10"）と一致している必要があります。不一致が発生すると、ワークロードが失敗します。このパラメーターはオプションであり、省略した場合は自動検出されます。
• AI ランタイムでは、GPU 環境 v4 以上を推奨しています。カスタムタイムアウト（timeout パラメーター）には、GPU 環境 v5 以上が必要です。
• このデコレータは、defaultでは 3 時間後にタイムアウトします。変更するには timeout を秒単位で渡し、無効にするには timeout=None を渡します。
• 実行は、ノートブックのライフサイクル内で実行されます。ノートブックを終了すると、ランが終了します。

もっと詳しく知る

• @distributed デコレーター、GPUType、および Ray API の詳細については、Serverless GPU Python API のリファレンスドキュメントを参照してください。
• トレーニングパイプラインの効率性と回復力を高めるパターンについては、パフォーマンスと回復力に関するガイドを参照してください。
• エンドツーエンドのトレーニングシナリオについては、ノートブックの例を参照してください。