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：性能最高，但配置复杂
• 原生 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)

高级主题

多节点配置：参见 references/multi-node.md，了解如何在 AWS、GCP、Kubernetes 和 SLURM 上部署 Ray 集群。

超参数调优：参见 references/hyperparameter-tuning.md，了解 Ray Tune 集成、搜索算法（Optuna、HyperOpt）和基于种群的训练。

自定义训练循环：参见 references/custom-loops.md，了解 Ray Train 的高级用法、自定义后端以及与其他框架的集成。

硬件要求

• 单节点：1 个或更多 GPU（或 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

Agent 操作流程

1. 检查环境。 检查 GPU 类型、内存、驱动和 CUDA 版本（nvidia-smi）、框架版本，以及数据集的位置和大小。
2. 明确输入。 根据数据或用户提供的信息确认格式、标识符和参数。任何会影响结果的值都应询问确认，不要猜测。
3. 先运行小规模版本。 执行冒烟测试：使用小型模型或数据子集，运行少量步骤，并确认损失下降且检查点已保存。
4. 按照上述说明和参考资料执行完整任务。
5. 验证结果。 在留出数据上跟踪指标，与基线进行比较，并记录随机种子、配置和硬件信息。
6. 报告。 说明运行了什么（版本、命令、参数）、检查了什么，以及还有哪些不确定之处。

如果出现这种情况	请这样做
CUDA 显存不足	减小批次大小，启用梯度累积、梯度检查点或混合精度，或者对模型进行分片。
损失为 NaN 或发散	降低学习率，检查数据中是否存在无效值，并启用梯度裁剪。
这些说明中提到的函数、标志或端点在已安装版本中不存在	查阅已安装版本自身的文档（help()、--help、官方文档），进行相应调整，并告知用户。绝不捏造 API。
必需的输入、标识符或参数含糊不清	询问用户，或在运行前明确说明所作的假设。

诚信规则

• 绝不捏造结果、参数、标识符、引用或统计数据。如果某项操作无法运行或验证，请如实说明。
• 没有日志时，绝不声称取得了训练结果；启动大型任务前，先估算计算成本。
• 此处与版本相关的细节可能已经过时：请根据已安装版本的官方文档进行确认。
• 对于会产生费用、占用共享 GPU 或云配额、涉及个人或患者数据，或无法撤销的操作，必须事先征求同意。

相关技能

• ray-data：用于机器学习工作负载的可扩展数据处理。
• pytorch-lightning-distributed：高层级 PyTorch 框架，提供 Trainer 类、自动分布式训练（DDP/FSDP/DeepSpeed）、回调系统，以及精简的样板代码……
• weights-and-biases：使用 Weights & Biases（W&B、wandb）记录和跟踪机器学习实验：指标、超参数、检查点、参数搜索、制品……