ray-train

Ray Train：在 Ray 集群上编排分布式训练、容错与 GPU 调度

补充说明

Ray Train 是 Ray 生态里的分布式训练编排层，负责把训练函数/类分发到多个 worker，申请 GPU、初始化进程组、处理失败重试和 checkpoint。它不替代 PyTorch DDP/FSDP/DeepSpeed，而是站在它们之上做集群调度和生命周期管理；单机多卡可以用 TorchTrainer，多机则依赖 Ray 集群。

Ray Train 适合“训练任务多、需要排队和弹性资源”的场景，也适合把 Hugging Face Trainer、XGBoost、LightGBM 等接入统一调度。它引入了 Ray 的 actor/object store 模型，调试时要比原生 torchrun 多一层。

常用示例

import ray
from ray.train import ScalingConfig, RunConfig
from ray.train.torch import TorchTrainer

def train_func(config):
    # 这里写单进程训练逻辑；Ray 会为每个 worker 设置分布式环境
    import torch
    model = MyModel().cuda()
    model = ray.train.torch.prepare_model(model)
    dataloader = ray.train.torch.prepare_data_loader(loader)
    for batch in dataloader:
        loss = model(batch).sum()
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

trainer = TorchTrainer(
    train_func,
    scaling_config=ScalingConfig(num_workers=8, use_gpu=True),
    run_config=RunConfig(name="my-training", storage_path="/tmp/ray-results"),
)
result = trainer.fit()
print(result.metrics)

# 本地启动 Ray 集群
ray start --head --num-gpus=8

# 提交作业
ray job submit --address http://127.0.0.1:8265 \
  -- python train_ray.py

# 查看集群资源
ray status

# 恢复训练：保存并加载 checkpoint
from ray.train import Checkpoint
checkpoint = Checkpoint.from_directory("/tmp/checkpoint")
trainer = TorchTrainer(..., resume_from_checkpoint=checkpoint)

易错点与排错

• 把 Ray Train 当 DDP：它负责调度和 worker 生命周期，进程组/模型包装仍要用 ray.train.torch.prepare_model 等接口。
• 训练函数里重复初始化分布式：Ray 已经设置好环境，不要再手动 init_process_group，否则可能冲突。
• GPU 分配不匹配：ScalingConfig(use_gpu=True) 只表示每个 worker 用 GPU；num_workers 要和集群可用 GPU 数匹配。
• Object store 内存打满：大对象/数据集通过 Ray object store 传递会占用内存；用共享存储或流式读取，避免反复 put/get。
• checkpoint 路径不可共享：多机 worker 写本地路径，主进程读不到；用共享存储或 Ray 的 Checkpoint 抽象。
• 失败重试掩盖问题：Ray 会自动重试 worker，但数据/代码 bug 会反复失败；看 RunConfig 的失败策略和日志。
• 版本兼容性：Ray、PyTorch、CUDA、Transformers 的组合要验证；Ray 集群各节点版本必须一致。
• 调试困难：Ray 多了一层 actor 和日志；先在单 worker 跑通，再扩到多 worker，保留每个 worker 的日志。

相关命令

• ray — Ray 核心运行时与对象存储。
• ray-serve — Ray 上的在线推理服务。
• torchrun — 原生 PyTorch 分布式启动。
• fsdp — 分片训练。
• slurm — 另一套集群作业调度。

参考链接

• Ray Train 文档
• Ray GitHub
• Ray TorchTrainer

读完自测

合上页面，先回答，再展开对照。

1. Ray Train 和 PyTorch DDP/FSDP 的关系是什么？
2. ScalingConfig(num_workers=8, use_gpu=True) 表示什么？
3. 多机训练时 checkpoint 为什么要用共享存储或 Ray Checkpoint？