ray-train

Ray Train - 分散トレーニングのオーケストレーション

クイックスタート

Ray Train を使うと、最小限のコード変更で、単一 GPU からマルチノードクラスターまで機械学習のトレーニングをスケールできます。

インストール：

pip install -U "ray[train]"

基本的な PyTorch トレーニング（単一ノード）：

import ray
from ray import train
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer
import torch
import torch.nn as nn

# Define training function
def train_func(config):
    # Your normal PyTorch code
    model = nn.Linear(10, 1)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

    # Prepare for distributed (Ray handles device placement)
    model = train.torch.prepare_model(model)

    for epoch in range(10):
        # Your training loop
        output = model(torch.randn(32, 10))
        loss = output.sum()
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        # Report metrics (logged automatically)
        train.report({"loss": loss.item(), "epoch": epoch})

# Run distributed training
trainer = TorchTrainer(
    train_func,
    scaling_config=ScalingConfig(
        num_workers=4,  # 4 GPUs/workers
        use_gpu=True
    )
)

result = trainer.fit()
print(f"Final loss: {result.metrics['loss']}")

これだけです！ Ray が以下を処理します：

• 分散処理の調整
• GPU の割り当て
• フォールトトレランス
• チェックポイントの保存
• メトリクスの集約

よくあるワークフロー

ワークフロー 1：既存の PyTorch コードをスケールする

元の単一 GPU コード：

model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters())

for epoch in range(epochs):
    for batch in dataloader:
        loss = model(batch)
        loss.backward()
        optimizer.step()

Ray Train 版（マルチ GPU／マルチノードにスケール）：

from ray.train.torch import TorchTrainer
from ray import train

def train_func(config):
    model = MyModel()
    optimizer = torch.optim.Adam(model.parameters())

    # Prepare for distributed (automatic device placement)
    model = train.torch.prepare_model(model)
    dataloader = train.torch.prepare_data_loader(dataloader)

    for epoch in range(epochs):
        for batch in dataloader:
            loss = model(batch)
            loss.backward()
            optimizer.step()

            # Report metrics
            train.report({"loss": loss.item()})

# Scale to 8 GPUs
trainer = TorchTrainer(
    train_func,
    scaling_config=ScalingConfig(num_workers=8, use_gpu=True)
)
trainer.fit()

メリット：同じコードを 1 GPU でも 1000 GPU でも実行できます

ワークフロー 2：HuggingFace Transformers との統合

from ray.train.huggingface import TransformersTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments

def train_func(config):
    # Load model and tokenizer
    model = AutoModelForCausalLM.from_pretrained("gpt2")
    tokenizer = AutoTokenizer.from_pretrained("gpt2")

    # Training arguments (HuggingFace API)
    training_args = TrainingArguments(
        output_dir="./output",
        num_train_epochs=3,
        per_device_train_batch_size=8,
        learning_rate=2e-5,
    )

    # Ray automatically handles distributed training
    from transformers import Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
    )

    trainer.train()

# Scale to multi-node (2 nodes × 8 GPUs = 16 workers)
trainer = TransformersTrainer(
    train_func,
    scaling_config=ScalingConfig(
        num_workers=16,
        use_gpu=True,
        resources_per_worker={"GPU": 1}
    )
)

result = trainer.fit()

ワークフロー 3：Ray Tune によるハイパーパラメーター調整

from ray import tune
from ray.train.torch import TorchTrainer
from ray.tune.schedulers import ASHAScheduler

def train_func(config):
    # Use hyperparameters from config
    lr = config["lr"]
    batch_size = config["batch_size"]

    model = MyModel()
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    model = train.torch.prepare_model(model)

    for epoch in range(10):
        # Training loop
        loss = train_epoch(model, optimizer, batch_size)
        train.report({"loss": loss, "epoch": epoch})

# Define search space
param_space = {
    "lr": tune.loguniform(1e-5, 1e-2),
    "batch_size": tune.choice([16, 32, 64, 128])
}

# Run 20 trials with early stopping
tuner = tune.Tuner(
    TorchTrainer(
        train_func,
        scaling_config=ScalingConfig(num_workers=4, use_gpu=True)
    ),
    param_space=param_space,
    tune_config=tune.TuneConfig(
        num_samples=20,
        scheduler=ASHAScheduler(metric="loss", mode="min")
    )
)

results = tuner.fit()
best = results.get_best_result(metric="loss", mode="min")
print(f"Best hyperparameters: {best.config}")

結果：クラスター全体で分散ハイパーパラメーター検索を実行できます

ワークフロー 4：チェックポイントとフォールトトレランス

from ray import train
from ray.train import Checkpoint

def train_func(config):
    model = MyModel()
    optimizer = torch.optim.Adam(model.parameters())

    # Try to resume from checkpoint
    checkpoint = train.get_checkpoint()
    if checkpoint:
        with checkpoint.as_directory() as checkpoint_dir:
            state = torch.load(f"{checkpoint_dir}/model.pt")
            model.load_state_dict(state["model"])
            optimizer.load_state_dict(state["optimizer"])
            start_epoch = state["epoch"]
    else:
        start_epoch = 0

    model = train.torch.prepare_model(model)

    for epoch in range(start_epoch, 100):
        loss = train_epoch(model, optimizer)

        # Save checkpoint every 10 epochs
        if epoch % 10 == 0:
            checkpoint = Checkpoint.from_directory(
                train.get_context().get_trial_dir()
            )
            torch.save({
                "model": model.state_dict(),
                "optimizer": optimizer.state_dict(),
                "epoch": epoch
            }, checkpoint.path / "model.pt")

            train.report({"loss": loss}, checkpoint=checkpoint)

trainer = TorchTrainer(
    train_func,
    scaling_config=ScalingConfig(num_workers=8, use_gpu=True)
)

# Automatically resumes from checkpoint if training fails
result = trainer.fit()

ワークフロー 5：マルチノードトレーニング

from ray.train import ScalingConfig

# Connect to Ray cluster
ray.init(address="auto")  # Or ray.init("ray://head-node:10001")

# Train across 4 nodes × 8 GPUs = 32 workers
trainer = TorchTrainer(
    train_func,
    scaling_config=ScalingConfig(
        num_workers=32,
        use_gpu=True,
        resources_per_worker={"GPU": 1, "CPU": 4},
        placement_strategy="SPREAD"  # Spread across nodes
    )
)

result = trainer.fit()

Ray クラスターを起動：

# On head node
ray start --head --port=6379

# On worker nodes
ray start --address=<head-node-ip>:6379

代替手段と比べた使いどころ

Ray Train を使うケース：

• 複数のマシンにまたがってトレーニングする（マルチノード）
• 大規模なハイパーパラメーター調整が必要
• フォールトトレランスが必要（失敗したワーカーを自動再起動）
• エラスティックスケーリング（トレーニング中にノードを追加／削除）
• 統一されたフレームワークを使いたい（PyTorch／TF／HF で同じコード）

主な利点：

• マルチノードオーケストレーション：最も簡単なマルチノード構成
• Ray Tune との統合：クラス最高水準のハイパーパラメーター調整
• フォールトトレランス：障害から自動的に復旧
• エラスティック：再起動せずにノードを追加／削除
• フレームワーク非依存：PyTorch、TensorFlow、HuggingFace、XGBoost

代わりに以下を使うケース：

• Accelerate：単一ノードのマルチ GPU、よりシンプル
• PyTorch Lightning：高レベルの抽象化、コールバック
• DeepSpeed：最大限のパフォーマンス、複雑なセットアップ
• Raw DDP：最大限の制御、最小限のオーバーヘッド

よくある問題

問題：Ray クラスターに接続できない

Ray のステータスを確認します：

ray status

# Should show:
# - Nodes: 4
# - GPUs: 32
# - Workers: Ready

接続されていない場合：

# Restart head node
ray stop
ray start --head --port=6379 --dashboard-host=0.0.0.0

# Restart worker nodes
ray stop
ray start --address=<head-ip>:6379

問題：メモリ不足

ワーカー数を減らすか、勾配蓄積を使用します：

scaling_config=ScalingConfig(
    num_workers=4,  # Reduce from 8
    use_gpu=True
)

# In train_func, accumulate gradients
for i, batch in enumerate(dataloader):
    loss = model(batch) / accumulation_steps
    loss.backward()

    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

問題：トレーニングが遅い

データの読み込みがボトルネックになっていないか確認します：

import time

def train_func(config):
    for epoch in range(epochs):
        start = time.time()
        for batch in dataloader:
            data_time = time.time() - start
            # Train...
            start = time.time()
            print(f"Data loading: {data_time:.3f}s")

データの読み込みが遅い場合は、ワーカー数を増やします：

dataloader = DataLoader(dataset, num_workers=8)

高度なトピック

マルチノード構成：AWS、GCP、Kubernetes、SLURM での Ray クラスターのデプロイについては、references/multi-node.md を参照してください。

ハイパーパラメーター調整：Ray Tune との統合、検索アルゴリズム（Optuna、HyperOpt）、集団ベースのトレーニングについては、references/hyperparameter-tuning.md を参照してください。

カスタムトレーニングループ：Ray Train の高度な使用方法、カスタムバックエンド、他のフレームワークとの統合については、references/custom-loops.md を参照してください。

ハードウェア要件

• 単一ノード：GPU 1 基以上（または CPU）
• マルチノード：ネットワーク接続されたマシン 2 台以上
• クラウド：AWS、GCP、Azure（Ray のオートスケーリング）
• オンプレミス：Kubernetes、SLURM クラスター

対応アクセラレーター：

• NVIDIA GPU（CUDA）
• AMD GPU（ROCm）
• TPU（Google Cloud）
• CPU

リソース

• ドキュメント：https://docs.ray.io/en/latest/train/train.html
• GitHub：https://github.com/ray-project/ray ⭐ 36,000+
• バージョン：2.40.0+
• サンプル：https://docs.ray.io/en/latest/train/examples.html
• Slack：https://forms.gle/9TSdDYUgxYs8SA9e8
• 利用企業：OpenAI、Uber、Spotify、Shopify、Instacart

エージェントの作業手順

1. 環境を確認します。 GPU の種類、メモリ、ドライバーと CUDA のバージョン（nvidia-smi）、フレームワークのバージョン、データセットの場所とサイズを確認します。
2. 入力を明確にします。 データまたはユーザーに確認して、形式、識別子、パラメーターを確定します。結果に影響する値は推測せず、質問してください。
3. まず小規模な実行を行います。 スモークテストとして、小さなモデルまたはサブセットで少ないステップ数だけ実行し、損失が減少することとチェックポイントが保存されることを確認します。
4. 上記の手順と参考情報を使用して、タスク全体を実行します。
5. 結果を検証します。 ホールドアウトデータで指標を追跡し、ベースラインと比較して、シード、設定、ハードウェアを記録します。
6. 報告します。 実行内容（バージョン、コマンド、パラメーター）、確認した内容、未解決の不確実な点を明示します。

この状況が発生した場合	対応
CUDA のメモリ不足	バッチサイズを小さくする、勾配累積／チェックポイントを有効にする、混合精度を使用する、またはモデルをシャーディングします。
損失が NaN になる、または発散する	学習率を下げ、データに無効な値がないか確認し、勾配クリッピングを有効にします。
これらの手順に記載された関数、フラグ、またはエンドポイントがインストール済みバージョンにない	インストール済みバージョンのドキュメント（help()、--help、公式ドキュメント）を確認し、適宜対応してユーザーに伝えます。API を捏造してはいけません。
必須の入力、識別子、またはパラメーターが曖昧	ユーザーに質問するか、実行前に仮定を明示します。

誠実性に関するルール

• 結果、パラメーター、識別子、引用、統計情報を捏造してはいけません。実行または検証できない場合は、その旨を明確に伝えます。
• ログなしに学習結果を主張してはいけません。大規模なジョブを開始する前に、計算コストを見積もります。
• ここに記載されたバージョン固有の情報は古い可能性があるものとして扱い、インストール済みバージョンの公式ドキュメントで確認します。
• 費用が発生する操作、共有 GPU やクラウドのクォータを消費する操作、個人データまたは患者データを扱う操作、取り消せない操作を行う前には、必ず確認を求めます。

関連スキル

• ray-data：ML ワークロード向けのスケーラブルなデータ処理。
• pytorch-lightning-distributed：Trainer クラス、自動分散学習（DDP/FSDP/DeepSpeed）、コールバックシステム、最小限のボイラープレートを備えた高水準の PyTorch フレームワーク…
• weights-and-biases：Weights & Biases（W&B、wandb）を使用して、機械学習の実験を記録・追跡します。指標、ハイパーパラメーター、チェックポイント、スイープ、アーティファクトなど…