マルチ GPU ワークロードとマルチノードワークロード

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

簡単スタート

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

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

# Import the distributed decorator
from serverless_gpu import distributed

# Decorate your training function with @distributed and specify the number of GPUs, the GPU type,
# and whether or not the GPUs are remote
@distributed(gpus=8, gpu_type='A10', remote=True)
def run_train():
    ...

ノートブックから 8 つの A10 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='A10', remote=True)
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 node 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 から返されます

remoteが True に設定されている場合、ワークロードはリモート GPU に分散されます。 remoteが False に設定されている場合、ワークロードは、現在のノートブックに接続されている単一の GPU ノードで実行されます。ノードに複数の GPU チップがある場合は、それらのすべてが利用されます。

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

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

Ray を使用した起動

また、サーバーレス gpu API では、@ray_launchの上に階層化された @distributed デコレーターを使用した Ray を使用した分散トレーニングの起動もサポートされています。各 ray_launch タスクは、まず、Ray ヘッドワーカーを決定し、IP を収集するために、聖火分散型ランデブーをブートストラップします。ランク 0 は ray start --head (メトリックエクスポートが有効な場合) を開始し、 RAY_ADDRESSを設定し、装飾された関数を Ray ドライバーとして実行します。他のノードは ray start --address 経由で参加し、ドライバーが完了マーカーを書き込むまで待機します。

追加の構成の詳細:

各ノードで Ray システムメトリックの収集を有効にするには、RayMetricsMonitorでremote=Trueを使用します。
標準の Ray API を使用して、装飾された関数内で Ray ランタイムオプション (アクター、データセット、配置グループ、スケジュール) を定義します。
デコレーターの引数またはノートブック環境で、関数 の外部 にあるクラスター全体のコントロール (GPU の数と種類、リモートモードとローカルモード、非同期動作、Databricks プール環境変数) を管理します。

次の例は、 @ray_launchの使用方法を示しています。

from serverless_gpu.ray import ray_launch
@ray_launch(gpus=16, remote=True, gpu_type='A10')
def foo():
    import os
    import ray
    print(ray.state.available_resources_per_node())
    return 1
foo.distributed()

完全な例については、このノートブックを参照してください。Ray を起動して、複数の A10 GPU で Resnet18 ニューラルネットワークをトレーニングします。

FAQs

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

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

from serverless_gpu import distributed

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

予約済みの GPU プールを使用できますか?

ワークスペースで予約済みの GPU プールが使用可能な場合（管理者に確認してみてください）、remote、True、そして@distributedのデコレーターを指定すると、ワークロードは既定で予約済みの GPU プールで起動されます。オンデマンド GPU プールを使用する場合は、次に示すように、分散関数を呼び出す前に、環境変数 DATABRICKS_USE_RESERVED_GPU_POOL を False に設定してください。

import os
os.environ['DATABRICKS_USE_RESERVED_GPU_POOL'] = 'False'
@distributed(gpus=8, remote=True)
def run_train():
    ...

詳細情報

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

フィードバック

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

Last updated on 2025-12-19