マルチ GPU ワークロード

Important

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

サーバーレス GPU Python API を使用して、1 つのノード上の複数の GPU 間で分散ワークロードを起動できます。この API には、GPU プロビジョニング、環境のセットアップ、ワークロードの分散の詳細を抽象化する、シンプルで統一されたインターフェイスが用意されています。コードの変更を最小限に抑えて、単一 GPU トレーニングから、同じノートブックからマルチ GPU 分散実行にシームレスに移行できます。

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

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

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

serverless_gpu API と TorchDistributor

次の表は、 serverless_gpu@distributed API と TorchDistributor を比較しています。

特徴	`serverless_gpu` `@distributed` API	トーチ分配器
インフラストラクチャ	完全にサーバーレス、クラスター管理なし	GPU ワーカーを使用する Spark クラスターが必要
セットアップ	単一デコレータ、最小構成	Spark クラスターと TorchDistributor のセットアップが必要
フレームワークサポート	PyTorch DDP、FSDP、DeepSpeed	主に PyTorch DDP
データの読み込み	デコレーター内では、Unity カタログボリュームを使用します	Spark またはファイルシステム経由

serverless_gpu API は、Databricks の新しいディープラーニングワークロードに推奨されるアプローチです。 TorchDistributor は、Spark クラスターと密接に結合されたワークロードで引き続き使用できます。

簡単スタート

分散トレーニング用のサーバーレス GPU API は、Databricks ノートブックとジョブ内のサーバーレス GPU に接続されている場合にプレインストールされます。 GPU 環境 4 以降をお勧めします。分散トレーニングに使用するには、 distributed デコレーターをインポートして使用し、トレーニング関数を配布します。

モデルトレーニングコードを関数でラップし、 @distributed デコレーターで関数を装飾します。修飾された関数は、分散実行のエントリポイントになります。すべてのトレーニングロジック、データ読み込み、モデルの初期化は、この関数内で定義する必要があります。

Warnung

gpu_typeの @distributed パラメーターは、ノートブックが接続されているアクセラレータの種類と一致している必要があります。たとえば、 @distributed(gpus=8, gpu_type='H100') ノートブックが H100 アクセラレータに接続されている必要があります。アクセラレータの種類が一致しない (H100 の指定中に A10 に接続するなど) を使用すると、ワークロードが失敗します。

次のコードスニペットは、 @distributedの基本的な使用方法を示しています。

# Import the distributed decorator
from serverless_gpu import distributed

# Decorate your training function with @distributed and specify the number of GPUs and GPU type
@distributed(gpus=8, gpu_type='H100')
def run_train():
    ...

ノートブックから 8 H100 GPU で多層パーセプトロン (MLP) モデルをトレーニングする完全な例を次に示します。

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


# Define the model
import os
import torch
import torch.distributed as dist
import torch.nn as nn

def setup():
    dist.init_process_group("nccl")
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))

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)

serverless_gpu ライブラリと分散モジュールをインポートします。
```
import serverless_gpu
from serverless_gpu import distributed
```

モデルトレーニングコードを関数でラップし、 @distributed デコレーターで関数を装飾します。

@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()

ユーザー定義の引数を使用して分散関数を呼び出して、分散トレーニングを実行します。
```
run_train.distributed(num_epochs=3, batch_size=1)
```
実行すると、ノートブックのセル出力に MLflow 実行リンクが生成されます。 MLflow の実行リンクをクリックするか、実験パネルで見つけて実行結果を表示します。

分散実行の詳細

サーバーレス GPU API は、いくつかの主要なコンポーネントで構成されています。

コンピューティングマネージャー: リソースの割り当てと管理を処理する
ランタイム環境: Python 環境と依存関係を管理します
ランチャー: ジョブの実行と監視を調整する

分散モードで実行する場合:

関数はシリアル化され、指定された数の GPU に分散されます
各 GPU は、同じパラメーターを使用して関数のコピーを実行します
環境はすべての GPU で同期されます
結果が収集され、すべての GPU から返されます

この API は、分散データ並列 (DDP)、フルシャードデータ並列 (FSDP)、 DeepSpeed などの一般的な並列トレーニングライブラリをサポートしています。

ノートブックの例のさまざまなライブラリを使用して、より実際の分散トレーニングシナリオを見つけることができます。

FAQs

データ読み込みコードはどこに配置する必要がありますか?

分散トレーニングにサーバーレス GPU API を使用する場合は、 @distributed デコレーター内でデータ読み込みコードを移動します。データセットのサイズは pickle で許可されている最大サイズを超える可能性があるため、次に示すように、デコレーター内でデータセットを生成することをお勧めします。

from serverless_gpu import distributed

# this may cause pickle error
dataset = get_dataset(file_path)
@distributed(gpus=8, gpu_type='H100')
def run_train():
  # good practice
  dataset = get_dataset(file_path)
  ....

詳細情報

API リファレンスについては、サーバーレス GPU Python API のドキュメントを参照してください。

フィードバック

このページはお役に立ちましたか?

Last updated on 2026-03-21