Ray Train 分布式机器学习训练框架 - 轻松扩展 PyTorch/HuggingFace 到多 GPU 集群

ray-train by orchestra-research/ai-research-skills

129 周安装量

6,300 GitHub Stars

GitHub

安装命令

npx skills add https://github.com/orchestra-research/ai-research-skills --skill ray-train

AI/机器学习 PyTorch 分布式系统

🇨🇳中文介绍

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

# 定义训练函数
def train_func(config):
    # 你的常规 PyTorch 代码
    model = nn.Linear(10, 1)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

    # 为分布式训练做准备（Ray 处理设备放置）
    model = train.torch.prepare_model(model)

    for epoch in range(10):
        # 你的训练循环
        output = model(torch.randn(32, 10))
        loss = output.sum()
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        # 报告指标（自动记录）
        train.report({"loss": loss.item(), "epoch": epoch})

# 运行分布式训练
trainer = TorchTrainer(
    train_func,
    scaling_config=ScalingConfig(
        num_workers=4,  # 4 个 GPU/工作进程
        use_gpu=True
    )
)

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

就这样！ Ray 处理以下事项：

广告位招租

在这里展示您的产品或服务

触达数万 AI 开发者，精准高效

联系我们

工作流 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())

    # 为分布式训练做准备（自动设备放置）
    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()

            # 报告指标
            train.report({"loss": loss.item()})

# 扩展到 8 个 GPU
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):
    # 加载模型和分词器
    model = AutoModelForCausalLM.from_pretrained("gpt2")
    tokenizer = AutoTokenizer.from_pretrained("gpt2")

    # 训练参数（HuggingFace API）
    training_args = TrainingArguments(
        output_dir="./output",
        num_train_epochs=3,
        per_device_train_batch_size=8,
        learning_rate=2e-5,
    )

    # Ray 自动处理分布式训练
    from transformers import Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
    )

    trainer.train()

# 扩展到多节点（2 节点 × 8 GPU = 16 个工作进程）
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):
    # 使用配置中的超参数
    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):
        # 训练循环
        loss = train_epoch(model, optimizer, batch_size)
        train.report({"loss": loss, "epoch": epoch})

# 定义搜索空间
param_space = {
    "lr": tune.loguniform(1e-5, 1e-2),
    "batch_size": tune.choice([16, 32, 64, 128])
}

# 运行 20 次试验，支持早停
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())

    # 尝试从检查点恢复
    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)

        # 每 10 个周期保存一次检查点
        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)
)

# 如果训练失败，自动从检查点恢复
result = trainer.fit()

工作流 5：多节点训练

from ray.train import ScalingConfig

# 连接到 Ray 集群
ray.init(address="auto")  # 或者 ray.init("ray://head-node:10001")

# 跨 4 个节点 × 8 个 GPU = 32 个工作进程进行训练
trainer = TorchTrainer(
    train_func,
    scaling_config=ScalingConfig(
        num_workers=32,
        use_gpu=True,
        resources_per_worker={"GPU": 1, "CPU": 4},
        placement_strategy="SPREAD"  # 跨节点分散
    )
)

result = trainer.fit()

启动 Ray 集群：

# 在头节点上
ray start --head --port=6379

# 在工作节点上
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

# 应显示：
# - 节点数：4
# - GPU 数：32
# - 工作进程：就绪

# 重启头节点
ray stop
ray start --head --port=6379 --dashboard-host=0.0.0.0

# 重启工作节点
ray stop
ray start --address=<head-ip>:6379

问题：内存不足

减少工作进程数量或使用梯度累积：

scaling_config=ScalingConfig(
    num_workers=4,  # 从 8 减少
    use_gpu=True
)

# 在 train_func 中，累积梯度
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
            # 训练...
            start = time.time()
            print(f"数据加载耗时：{data_time:.3f}秒")

如果数据加载慢，增加工作进程数：

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。

单节点：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

🇺🇸English

Ray Train - Distributed Training Orchestration

Quick start

Ray Train scales machine learning training from single GPU to multi-node clusters with minimal code changes.

Installation :

pip install -U "ray[train]"

Basic PyTorch training (single node):

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']}")

That's it! Ray handles:

Distributed coordination
GPU allocation
Fault tolerance
Checkpointing
Metric aggregation

Common workflows

Workflow 1: Scale existing PyTorch code

Original single-GPU code :

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 version (scales to multi-GPU/multi-node):

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()

Benefits : Same code runs on 1 GPU or 1000 GPUs

Workflow 2: HuggingFace Transformers integration

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()

Workflow 3: Hyperparameter tuning with 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}")

Result : Distributed hyperparameter search across cluster

Workflow 4: Checkpointing and fault tolerance

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()

Workflow 5: Multi-node training

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()

Launch Ray cluster :

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

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

When to use vs alternatives

Use Ray Train when :

Training across multiple machines (multi-node)
Need hyperparameter tuning at scale
Want fault tolerance (auto-restart failed workers)
Elastic scaling (add/remove nodes during training)
Unified framework (same code for PyTorch/TF/HF)

Key advantages :

Multi-node orchestration : Easiest multi-node setup
Ray Tune integration : Best-in-class hyperparameter tuning
Fault tolerance : Automatic recovery from failures
Elastic : Add/remove nodes without restarting
Framework agnostic : PyTorch, TensorFlow, HuggingFace, XGBoost

Use alternatives instead :

Accelerate : Single-node multi-GPU, simpler
PyTorch Lightning : High-level abstractions, callbacks
DeepSpeed : Maximum performance, complex setup
Raw DDP : Maximum control, minimal overhead

Common issues

Issue: Ray cluster not connecting

Check ray status:

ray status

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

If not connected:

# 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

Issue: Out of memory

Reduce workers or use gradient accumulation:

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()

Issue: Slow training

Check if data loading is bottleneck:

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")

If data loading is slow, increase workers:

dataloader = DataLoader(dataset, num_workers=8)

Advanced topics

Multi-node setup : See references/multi-node.md for Ray cluster deployment on AWS, GCP, Kubernetes, and SLURM.

Hyperparameter tuning : See references/hyperparameter-tuning.md for Ray Tune integration, search algorithms (Optuna, HyperOpt), and population-based training.

Custom training loops : See references/custom-loops.md for advanced Ray Train usage, custom backends, and integration with other frameworks.

Hardware requirements

Single node : 1+ GPUs (or CPUs)
Multi-node : 2+ machines with network connectivity
Cloud : AWS, GCP, Azure (Ray autoscaling)
On-prem : Kubernetes, SLURM clusters

Supported accelerators :

NVIDIA GPUs (CUDA)
AMD GPUs (ROCm)
TPUs (Google Cloud)
CPUs

Resources

Docs: https://docs.ray.io/en/latest/train/train.html
GitHub: https://github.com/ray-project/ray ⭐ 36,000+
Version: 2.40.0+
Examples: https://docs.ray.io/en/latest/train/examples.html
Slack: https://forms.gle/9TSdDYUgxYs8SA9e8
Used by: OpenAI, Uber, Spotify, Shopify, Instacart

Weekly Installs

Repository

orchestra-resea…h-skills

GitHub Stars

5.6K

First Seen

Feb 7, 2026

Security Audits

Gen Agent Trust HubPass SocketPass SnykWarn

Installed on

codex70

opencode70

gemini-cli68

cursor68

github-copilot67

kimi-cli63

AI 代码实施计划编写技能 | 自动化开发任务分解与 TDD 流程规划工具

50,900 周安装