关于 ray-train
ray-train 是用于机器学习分布式训练编排的技能资料，面向需要从单机扩展到多节点集群的开发者。它涵盖 Ray Train 与 Ray Tune 工作流、容错和检查点管理。
主要功能
- 01扩展 PyTorch 训练
- 02集成 HuggingFace Transformers
- 03与 Ray Tune 联合调参
- 04故障恢复与检查点
- 05多节点训练编排
适用场景
- 01将单 GPU 训练扩展到多机
- 02在集群中搜索超参数
- 03训练失败后恢复检查点
安装
复制下面的命令，在终端中运行。
$
npx skills add https:/安装后在 Claude Code 中输入/ray-train确认已加载。
附带文件references/multi-node.md
详细介绍
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 与其他方案
适合使用 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
评论
支持 Markdown