マルチGPUワークロード

備考

ベータ版

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

単一ノード上の複数の GPU に分散ワークロードを起動するには、GPU Python API使用します。 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	TorchDistributor
インフラストラクチャー	完全なサーバレス、クラスター管理なし	GPU ワーカーを備えたSparkクラスターが必要
設定	単一のデコレーター、最小限の構成	SparkクラスターとTorchDistributorセットアップが必要です
フレームワークのサポート	PyTorch DDP、FSDP、DeepSpeed	主にPyTorch DDP
データ読み込み	デコレーター内部では、 Unity Catalogボリュームを使用します。	Sparkまたはファイルシステム経由

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

クイックスタート

分散トレーニング用のサーバレス GPU API 、 Databricksノートブックおよびジョブ内でサーバレス GPU に接続するときにプレインストールされます。 GPU環境はバージョン4以上を推奨します。分散トレーニングに使用するには、 distributedトレーニング関数を配布するためのデコレータ。

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

警告

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

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

Python
# 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) モデルをトレーニングする完全な例です。ノートブック:

モデルをセットアップし、ユーティリティ関数を定義します。

Python

# 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 ライブラリと分散モジュールをインポートします。
Python
```
import serverless_gpu
from serverless_gpu import distributed
```

モデルのトレーニングコードを関数でラップし、その関数に@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()

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

分散実行の詳細

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

コンピュートマネージャー: リソースの割り当てと管理を処理します。
Runtime環境： Python環境と依存関係を管理します
ランチャー: ジョブの実行とモニタリングを調整します。

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

この関数はシリアル化され、指定された数のGPUに分散されます。
各 GPU は同じ問題を持つ関数のコピーを実行します
環境はすべてのGPU間で同期されています
結果はすべてのGPUから収集され、返されます。

API は、分散データ並列(DDP) などの一般的な並列トレーニングライブラリをサポートしています。完全シャーディングデータ並列処理（FSDP）ディープスピード。

ノートブックの例には、さまざまなライブラリを使用したより実践的な分散トレーニングシナリオが掲載されています。

よくある質問

データ読み込みコードはどこに配置すべきですか？

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

Python
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ドキュメントを参照してください。

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

serverless_gpu APIとTorchDistributorの比較​

クイックスタート​

分散実行の詳細​

よくある質問​

データ読み込みコードはどこに配置すべきですか？​

もっと詳しく知る​