Databricks分散トレーニング

モデルパラメータが数億〜数千億に達する現代の深層学習では、1台のGPUで訓練を完結させることは現実的ではありません。Databricksは、Sparkクラスタのインフラを活用してPyTorch・TensorFlow・DeepSpeedなどの分散トレーニングをノートブックから直接実行できる統合環境を提供しています。本記事ではTorchDistributor・HorovodRunner・Ray・DeepSpeedの4つのフレームワークを軸に、分散戦略の選択からGPU構成、MLflowとの連携まで実務と試験の両面から解説します。

データ並列 vs モデル並列

分散トレーニングの戦略は大きく2つに分かれます。どちらを採用するかはモデルサイズとGPUメモリの関係で決まります。

データ並列（Data Parallelism）

モデルの完全なコピーを各GPUに配置し、訓練データを分割して各GPUが異なるミニバッチを処理します。各GPUは独立にフォワード・バックワードパスを実行し、勾配をAllReduceで集約してパラメータを同期します。モデルが1台のGPUメモリに収まる場合はデータ並列が第一選択です。

データ並列の概念図:

┌─────────────┐  ┌─────────────┐  ┌─────────────┐
│   GPU 0     │  │   GPU 1     │  │   GPU 2     │
│ Model Copy  │  │ Model Copy  │  │ Model Copy  │
│ Batch 0-99  │  │ Batch 100-199│ │ Batch 200-299│
└──────┬──────┘  └──────┬──────┘  └──────┬──────┘
       │                │                │
       └────────┬───────┴────────┬───────┘
                │  AllReduce     │
                │  (勾配集約)     │
                ▼                ▼
         パラメータ同期 → 次のイテレーション

モデル並列（Model Parallelism）

モデルが1台のGPUメモリに収まらない場合、モデル自体を分割して複数GPUに配置します。テンソル並列（Tensor Parallelism）は1つのレイヤー内の行列演算を分割し、パイプライン並列（Pipeline Parallelism）はレイヤー群を段階的に各GPUに割り当てます。LLMのファインチューニングではこの2つを組み合わせる3D並列（データ＋テンソル＋パイプライン）が主流です。

モデル並列の概念図:

パイプライン並列:
┌──────────┐    ┌──────────┐    ┌──────────┐
│  GPU 0   │───▶│  GPU 1   │───▶│  GPU 2   │
│ Layer 0-7│    │ Layer 8-15│   │Layer 16-23│
└──────────┘    └──────────┘    └──────────┘

テンソル並列 (1レイヤー内):
┌────────────────────────────┐
│       Weight Matrix        │
│  ┌──────┬──────┬──────┐    │
│  │GPU 0 │GPU 1 │GPU 2 │    │
│  │Col0-3│Col4-7│Col8-11│   │
│  └──────┴──────┴──────┘    │
└────────────────────────────┘

TorchDistributor

TorchDistributorはDatabricks Runtime 13.0+で導入された、Sparkクラスタ上でPyTorch DDPを実行するためのAPIです。 Sparkのタスクスケジューリングを利用してワーカーノードにPyTorchプロセスを配置し、 NCCLバックエンドでGPU間通信を行います。 従来のPyTorch DDPではtorchrunコマンドやランチャースクリプトが必要でしたが、 TorchDistributorはノートブックから直接分散訓練を起動できるため、セットアップの手間が大幅に削減されます。

from pyspark.ml.torch.distributor import TorchDistributor
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

def train_fn():
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    device = torch.device(f"cuda:{rank % torch.cuda.device_count()}")

    model = nn.Sequential(
        nn.Linear(784, 512), nn.ReLU(),
        nn.Linear(512, 256), nn.ReLU(),
        nn.Linear(256, 10)
    ).to(device)
    model = DDP(model, device_ids=[device])

    dataset = load_training_data()  # 訓練データの読み込み
    sampler = DistributedSampler(dataset)
    loader = DataLoader(dataset, batch_size=64, sampler=sampler)

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = nn.CrossEntropyLoss()

    for epoch in range(10):
        sampler.set_epoch(epoch)
        for batch_x, batch_y in loader:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)
            loss = loss_fn(model(batch_x), batch_y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

    dist.destroy_process_group()
    return model

# Sparkクラスタ上で4GPU分散訓練を実行
distributor = TorchDistributor(
    num_processes=4,
    local_mode=False,    # True: ドライバノードのみ、False: ワーカーノードに分散
    use_gpu=True
)
trained_model = distributor.run(train_fn)

local_mode=Trueはドライバノード上のマルチGPU実行（デバッグ用途）、local_mode=Falseはワーカーノードへの分散実行です。num_processesは使用するGPUの総数を指定します。 クラスタのワーカーが各2GPUで2ノードならnum_processes=4です。

HorovodRunner

HorovodはUber社が開発したフレームワーク非依存の分散トレーニングライブラリです。PyTorch・TensorFlow・Kerasに対応し、MPIベースのAllReduce通信で勾配を集約します。DatabricksではHorovodRunnerというラッパーAPIが提供されており、Sparkクラスタ上でHorovodプロセスをバリアモードで起動します。

import horovod.torch as hvd
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, DistributedSampler

def train_hvd():
    hvd.init()
    rank = hvd.rank()
    torch.cuda.set_device(hvd.local_rank())
    device = torch.device("cuda")

    model = nn.Sequential(
        nn.Linear(784, 512), nn.ReLU(),
        nn.Linear(512, 256), nn.ReLU(),
        nn.Linear(256, 10)
    ).to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3 * hvd.size())
    optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())
    hvd.broadcast_parameters(model.state_dict(), root_rank=0)

    dataset = load_training_data()
    sampler = DistributedSampler(dataset, num_replicas=hvd.size(), rank=hvd.rank())
    loader = DataLoader(dataset, batch_size=64, sampler=sampler)

    loss_fn = nn.CrossEntropyLoss()
    for epoch in range(10):
        sampler.set_epoch(epoch)
        for batch_x, batch_y in loader:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)
            loss = loss_fn(model(batch_x), batch_y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

    return model

from sparkdl import HorovodRunner
hr = HorovodRunner(np=4, driver_log_verbosity="all")
trained_model = hr.run(train_hvd)

HorovodRunnerのnpパラメータはプロセス数です。np=-1を指定するとクラスタの全GPUを自動検出して使用します。 TorchDistributorとの主な違いは、HorovodがDistributedOptimizerでoptimizerをラップし、broadcast_parametersで初期パラメータを明示的に同期する必要がある点です。 TorchDistributorではPyTorch標準のDDPラッパーをそのまま使うため、既存のDDPコードを最小限の変更で移行できます。

Ray on Databricks

Ray 2.x以降、DatabricksクラスタはRayクラスタとしても動作できるようになりました。ray.util.spark.setup_ray_cluster()でSparkワーカー上にRayワーカーを起動し、 Ray TrainやRay TuneでPyTorch・TensorFlowの分散訓練とハイパーパラメータチューニングを統合実行できます。

Rayの強みは訓練とチューニングの統合です。Ray Tuneの探索アルゴリズム（ASHA、PBTなど）とRay Trainの分散訓練を組み合わせることで、複数のハイパーパラメータ設定を並列に分散訓練できます。単純な分散訓練だけならTorchDistributorの方がセットアップが簡単ですが、大規模なハイパーパラメータサーチを伴う場合はRayが有利です。

DeepSpeed統合

MicrosoftのDeepSpeedは、ZeROオプティマイザによるメモリ効率化が最大の特徴です。ZeRO Stage 1はオプティマイザ状態を分割、Stage 2は勾配も分割、Stage 3はモデルパラメータも分割し、GPUメモリ消費を劇的に削減します。Databricks上ではTorchDistributorの訓練関数内でDeepSpeedを初期化して使用します。

70Bパラメータ規模のLLMをファインチューニングする場合、ZeRO Stage 3 + Offload（CPUメモリへのパラメータ退避）を組み合わせることで、A100 80GB×8台程度のクラスタでも実行可能になります。DeepSpeedはJSON設定ファイルで各種最適化を制御するため、コード変更なしにメモリ最適化レベルやバッチサイズ調整が可能です。

4つの分散フレームワーク比較

シングルノード vs マルチノードの判断基準

分散トレーニングにはノード間通信のオーバーヘッドが伴います。闇雲にマルチノードにすると、通信コストがボトルネックになり逆に遅くなるケースがあります。以下の基準で判断します。

GPU選定とクラスタ設定

DatabricksでGPUクラスタを構成する際、ワークロードに応じて適切なGPUインスタンスを選択します。

A10Gはコストパフォーマンスに優れ、7B以下のモデルのファインチューニングや推論に適しています。A100はLLMの本格的な訓練に必要なメモリ帯域とTFLOPSを備え、分散訓練のスタンダードです。H100はFP8対応のTransformer Engineを搭載し、A100比で約3倍の訓練スループットを実現しますが、コストも相応に高くなります。

クラスタ設定では、Databricks Runtime ML（MLR）を選択し、GPUドライバ・CUDA・cuDNN・NCCL・PyTorch・TensorFlowが事前インストールされた環境を使います。マルチノード分散ではワーカー間のネットワーク帯域が重要で、AWSならEFAネットワーク、AzureならInfiniBand対応のインスタンスを選択すると通信ボトルネックを軽減できます。

MLflow連携（分散実行時のトラッキング）

分散トレーニングでは複数のプロセスが同時にMLflowにログを書き込むと、 メトリクスの重複やランの競合が発生します。 ベストプラクティスはrank 0のプロセスのみがMLflowにログを記録するパターンです。

import mlflow
import torch.distributed as dist

def train_fn():
    dist.init_process_group("nccl")
    rank = dist.get_rank()

    # ... モデル定義・訓練ループ ...

    # rank 0 のみが MLflow にログを記録
    if rank == 0:
        with mlflow.start_run():
            mlflow.log_param("num_gpus", dist.get_world_size())
            mlflow.log_param("batch_size_per_gpu", 64)
            mlflow.log_param("effective_batch_size", 64 * dist.get_world_size())
            mlflow.log_metric("final_loss", avg_loss)
            mlflow.pytorch.log_model(model.module, "model")

    dist.destroy_process_group()

model.moduleはDDPラッパーを外した元のモデルです。 DDPラッパーのまま保存すると推論時にDDP初期化が必要になるため、 必ず.moduleでアンラップしてからMLflowに記録します。 effective_batch_sizeはGPU数×per-GPU batch sizeで、学習率のスケーリングの根拠としてログに残します。

試験で問われるポイント

分散トレーニングはML Professional試験で深く問われるトピックです。以下のポイントを整理しておくと得点源になります。

• TorchDistributor vs HorovodRunner: TorchDistributorはPyTorch専用でSpark統合がネイティブ、HorovodRunnerはフレームワーク非依存でMPIベース。新規PyTorchプロジェクトならTorchDistributor、TensorFlow分散や既存Horovodコードの移行ならHorovodRunnerを選択
• データ並列 vs モデル並列: モデルが1GPUに収まるならデータ並列、収まらないならモデル並列。試験では「70Bパラメータの LLMを4台のA100で訓練したい」→ モデル並列（ZeRO Stage 3）が必要、のような判断を求められる
• local_modeの意味: TorchDistributorのlocal_mode=Trueはドライバノード上のローカル実行（デバッグ用）、Falseはワーカーノードへの分散実行
• MLflowログのベストプラクティス：rank 0のみがログを書き、DDPラッパーを外してモデルを保存する
• np=-1の意味：HorovodRunnerのnp=-1はクラスタの全GPUを自動使用するオプション

問題で確認

よくある質問