- 人工智能
- 分布式训练
- 强化学习
- 任务调度
- 模型推理服务
【免费下载链接】ray
Ray is an AI compute engine. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.
导读
Local Mode(本地模式)是 Ray Train V2 提供的一种轻量级运行方式:通过ScalingConfig(num_workers=0)让训练函数直接在当前进程内执行,完全不启动 Ray 训练 Worker Actor,从而获得接近原生 Python 的调试与迭代体验。本指南以 doc/source/train/user-guides/local_mode.rst 为骨架,结合 python/ray/train/v2 目录下的控制器实现与测试用例,完整讲解 Local Mode 的两种执行形态(单进程、torchrun多进程)、启用方式、适用场景、与 Ray Data 的配合、从本地到分布式的平滑迁移,以及需要绕开的限制与 API 差异。读完本文,你将能够把任意ray.train训练代码以本地模式快速跑通、用 pytest 为训练逻辑编写单测,并用torchrun在单机多卡乃至多机多卡上调试分布式训练逻辑。
重要说明:本指南针对Ray Train V2。若你的代码基于 V1 API,可参考官方 Train V2 迁移讨论(issue #49454)。当前仓库中 V2 是新版训练栈,
Ray Train V2 必须通过环境变量RAY_TRAIN_V2_ENABLED=1显式启用,这一约束同样体现在 data_parallel_trainer.py 的_validate_configs校验中。
什么是 Local Mode:从"分布式编排"回到"当前进程"
常规的 Ray Train 训练流程是:Driver 进程启动一个 Controller,由 Controller 拉起一组 Ray Actor(Worker)作为训练进程,再通过 actor-to-actor 的通信完成同步、指标上报与检查点持久化。Local Mode 则完全不同——它跳过 Worker 的创建与调度,直接在当前进程中调用你的训练函数。
这一点在源码中有非常直接的体现:data_parallel_trainer.py 中通过一行判断切换执行路径:
self.running_in_local_mode = self.scaling_config.num_workers == 0在fit()中(data_parallel_trainer.py):
if self.running_in_local_mode: return self._initialize_and_run_local_controller(train_fn)本地控制器(utils.py)的核心逻辑非常朴素——设置训练工具上下文、调用训练函数、把最后上报的指标与检查点包装成Result返回:
def run(self, train_func): set_train_fn_utils(LocalTrainFnUtils(...)) result = train_func() train_fn_utils = get_train_fn_utils() return Result( metrics=train_fn_utils._get_last_metrics(), checkpoint=train_fn_utils.get_checkpoint(), path=None, # 没有持久化存储路径 error=None, return_value=result, )Local Mode 支持两种执行形态:
- 单进程模式(Single-process mode):训练函数在单个进程内运行,适合快速迭代与调试;
- 多进程模式(
torchrun):借助torchrun启动多个进程做多卡训练,适合用熟悉的工具调试分布式训练逻辑。
两种形态下,你使用的ray.trainAPI 完全一致,训练代码无需任何额外改动。
如何启用 Local Mode
只需在ScalingConfig中设置num_workers=0:
from ray.train import ScalingConfig from ray.train.torch import TorchTrainer def train_func(config): # Your training logic pass trainer = TorchTrainer( train_loop_per_worker=train_func, scaling_config=ScalingConfig(num_workers=0), ) result = trainer.fit()ScalingConfig的官方语义在 v2/api/config.py 中有明确记载:num_workers表示要启动的 Worker(Ray Actor)数量,当其为 0 时,训练函数将以 Local Mode 在当前进程运行。其他常用参数如use_gpu(默认False,为 True 时每个 Worker 保留 1 块 GPU)、resources_per_worker(以"CPU"、"GPU"为键覆盖每个 Worker 的资源预留)在 Local Mode 下同样可以传入,use_gpu=True在多进程模式下用于指示进程绑定 GPU 设备。
值得强调的是:Local Mode 提供与分布式训练完全相同的ray.trainAPI 表面,因此你可以先用本地模式验证训练逻辑,再无缝扩展到分布式训练,期间训练函数本身零改动。
何时使用 Local Mode
单进程 Local Mode 适合:
- 快速开发迭代:本地改动后立即运行,验证训练函数;
- 编写单元测试:在简化环境中验证训练逻辑的正确性(下文给出 pytest 示例);
- 调试训练逻辑:使用标准 Python 调试工具(如
pdb、IDE 断点)单步跟踪训练代码、定位问题。
多进程torchrunLocal Mode 适合:
- 验证多卡逻辑:用熟悉的
torchrun命令在多个 GPU 上验证分布式训练代码; - 迁移既有代码:把基于
torchrun的既有训练脚本接入 Ray Train,同时保留原有开发工作流; - 调试分布式行为:借助
torchrun的进程管理能力隔离并定位分布式训练逻辑中的问题。
提示:在 Local Mode 下 Ray Train 不启动 Worker Actor,但你的训练代码仍然可以使用其他 Ray 能力——例如在单进程模式中使用 Ray Data,或按需启动 Ray Actor(详见下文"使用 Ray Data"与"局限性与 API 差异")。
单进程 Local Mode 实战
PyTorch 训练示例
下面用TorchTrainer演示单进程 Local Mode 的完整用法,训练逻辑是一个简单的线性层 + SGD:
import torch from torch import nn import ray from ray.train import ScalingConfig from ray.train.torch import TorchTrainer def train_func(config): model = nn.Linear(10, 1) optimizer = torch.optim.SGD(model.parameters(), lr=config["lr"]) for epoch in range(config["epochs"]): # Training loop loss = model(torch.randn(32, 10)).sum() loss.backward() optimizer.step() # Report metrics ray.train.report({"loss": loss.item()}) trainer = TorchTrainer( train_loop_per_worker=train_func, train_loop_config={"lr": 0.01, "epochs": 3}, scaling_config=ScalingConfig(num_workers=0), ) result = trainer.fit() print(f"Final loss: {result.metrics['loss']}")运行这段代码,result.metrics中即可拿到最后一次ray.train.report的指标。从实现上看,这是因为LocalTrainFnUtils.report()(train_fn_utils.py)在内存中记录了最新的 metrics 与 checkpoint,LocalController随后把它们装配成Result返回。仓库自带的端到端测试 test_local_mode.py 验证了同样的行为:result.error is None、result.metrics["loss"]非空、result.checkpoint存在。
Local Mode 适用于所有 Ray Train 框架集成,包括 PyTorch Lightning、Hugging Face Transformers、LightGBM、XGBoost、TensorFlow 等。测试目录 test_local_mode.py 中分别有
test_lightning_trainer_local_mode、test_e2e_hf_local_mode、test_lightgbm_trainer_local_mode、test_xgboost_trainer_local_mode、test_tensorflow_linear_local_mode以及test_jax_trainer_local_mode作为佐证。
用 Local Mode 编写单元测试
Local Mode 最常见的用途之一是为训练逻辑编写快速单测——不需要启动集群,也不需要 mock 分布式通信:
import pytest import ray from ray.train import ScalingConfig from ray.train.torch import TorchTrainer def test_training_runs(): def train_func(config): # Report minimal training result ray.train.report({"loss": 0.5}) trainer = TorchTrainer( train_loop_per_worker=train_func, scaling_config=ScalingConfig(num_workers=0), ) result = trainer.fit() assert result.error is None assert result.metrics["loss"] == 0.5test_local_mode.py中的test_data_parallel_trainer_local_mode是仓库内最精简的同类测试:训练函数内部ray.train.report(metrics={"test": 1}, checkpoint=...),随后断言result.metrics == {"test": 1}且result.checkpoint存在(test_local_mode.py)。这个模式可以直接复用到你自己的训练管线测试中。
在 Local Mode 中使用 Ray Data
单进程 Local Mode 可以无缝配合 Ray Data:Ray Data 负责数据的加载与预处理,并把数据交给当前进程内的训练函数消费:
import ray from ray.train import ScalingConfig from ray.train.torch import TorchTrainer def train_func(config): # Get the dataset shard train_dataset = ray.train.get_dataset_shard("train") # Iterate over batches for batch in train_dataset.iter_batches(batch_size=32): # Training logic pass # Create a Ray Dataset dataset = ray.data.read_csv("s3://bucket/data.csv") trainer = TorchTrainer( train_loop_per_worker=train_func, scaling_config=ScalingConfig(num_workers=0), datasets={"train": dataset}, ) result = trainer.fit()在 Local Mode 下,get_dataset_shard直接返回传入的 dataset 本身(见 train_fn_utils.py 的实现,以及 utils.py 中对 callable 数据集参数的实例化逻辑)。test_e2e_hf_local_mode测试还展示了 Local Mode + Ray Data + Hugging Face Transformers 的组合用法(test_local_mode.py)。
警告:当使用
torchrun进行多进程训练时,Local Mode不支持 Ray Data。多进程训练请使用标准 PyTorch 数据加载机制,例如DataLoader配合DistributedSampler。这一限制在控制器实现中有硬性断言(torch.py):当 world size 大于 1 且传入了 datasets 时会直接报错 "Ray Data is not supported in local mode with multiple workers."。
多进程 Local Mode:用 torchrun 调试多卡/多机训练
Local Mode 支持通过torchrun启动多进程,从而在本地对多 GPU 训练进行与生产一致的过程编排调试。启动后,Ray Train 会自动检测torchrun设置的环境变量并据此配置分布式训练。
检测逻辑位于 torch.py 的has_torchrun_env(),它要求以下 6 个环境变量同时存在:
| 环境变量 | 含义 |
|---|---|
RANK | 当前进程在全局作业中的排名 |
LOCAL_RANK | 当前进程在节点内的排名 |
WORLD_SIZE | 参与作业的进程总数 |
LOCAL_WORLD_SIZE | 当前节点上参与作业的进程数 |
MASTER_ADDR | 主节点(rank 0)的 IP 或主机名 |
MASTER_PORT | 主节点上用于通信的空闲端口 |
检测到这些变量后,LocalTorchController._set_train_fn_utils()会依次执行:
- 调用
torch.distributed.init_process_group(backend="nccl" if torch.cuda.is_available() else "gloo")初始化进程组(torch.py); - 从分布式环境读取
world_size、global_rank,从环境变量读取local_rank; - GPU 可用时执行
torch.cuda.set_device(local_rank)绑定当前进程到对应设备; - 计算
nproc_per_node = LOCAL_WORLD_SIZE、node_rank = global_rank // nproc_per_node; - 把这些值写入
LocalTrainFnUtils,供ray.train.get_context()查询(torch.py)。
仓库测试 test_local_mode.py 的test_torch_distributed_variables_local_train_fn_utils完整验证了这套解析逻辑,例如在RANK=2 / WORLD_SIZE=4 / LOCAL_RANK=1 / LOCAL_WORLD_SIZE=2时,get_world_size()==4、get_world_rank()==2、get_local_rank()==1、get_node_rank()==1,且 CPU 场景调用init_process_group(backend="gloo")、GPU 场景调用backend="nccl"并执行set_device(1)。
单节点多 GPU 训练
以下示例展示如何用torchrun在单节点多卡上运行 Local Mode 训练。它使用标准 PyTorchDataLoader加载数据(FashionMNIST),便于直接迁移既有的 PyTorch 训练代码。
首先创建训练脚本train_script.py:
import os import tempfile import torch import torch.distributed as dist from torch import nn from torch.utils.data import DataLoader from torchvision.datasets import FashionMNIST from torchvision.transforms import ToTensor, Normalize, Compose from filelock import FileLock import ray from ray.train import Checkpoint, ScalingConfig, get_context from ray.train.torch import TorchTrainer def train_func(config): # Load dataset with file locking to avoid multiple downloads transform = Compose([ToTensor(), Normalize((0.5,), (0.5,))]) data_dir = "./data" # Only local rank 0 downloads the dataset local_rank = get_context().get_local_rank() if local_rank == 0: with FileLock(os.path.join(data_dir, "fashionmnist.lock")): train_dataset = FashionMNIST( root=data_dir, train=True, download=True, transform=transform ) # Wait for rank 0 to finish downloading dist.barrier() # Now all ranks can safely load the dataset train_dataset = FashionMNIST( root=data_dir, train=True, download=False, transform=transform ) train_loader = DataLoader( train_dataset, batch_size=config["batch_size"], shuffle=True ) # Prepare dataloader for distributed training train_loader = ray.train.torch.prepare_data_loader(train_loader) # Prepare model for distributed training model = nn.Sequential( nn.Flatten(), nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 10) ) model = ray.train.torch.prepare_model(model) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=config["lr"]) # Training loop for epoch in range(config["epochs"]): # Set epoch for distributed sampler if ray.train.get_context().get_world_size() > 1: train_loader.sampler.set_epoch(epoch) epoch_loss = 0.0 for batch_idx, (images, labels) in enumerate(train_loader): outputs = model(images) loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() epoch_loss += loss.item() avg_loss = epoch_loss / len(train_loader) # Report metrics and checkpoint with tempfile.TemporaryDirectory() as temp_dir: torch.save(model.state_dict(), os.path.join(temp_dir, "model.pt")) ray.train.report( {"loss": avg_loss, "epoch": epoch}, checkpoint=Checkpoint.from_directory(temp_dir) ) # Configure trainer for local mode trainer = TorchTrainer( train_loop_per_worker=train_func, train_loop_config={"lr": 0.001, "epochs": 10, "batch_size": 32}, scaling_config=ScalingConfig(num_workers=0, use_gpu=True), ) result = trainer.fit()然后通过torchrun启动训练:
# Train on 4 GPUs on a single node torchrun --nproc-per-node=4 train_script.py训练过程中,你可以随时通过ray.train.get_context()读取分布式信息:
from ray.train import get_context context = get_context() print(f"World size: {context.get_world_size()}") print(f"World rank: {context.get_world_rank()}") print(f"Local rank: {context.get_local_rank()}")从源码角度,这些值的来源是LocalTrainFnUtils构造时传入的分布式元数据,而TorchTrainer会重写_get_local_controller()返回LocalTorchController(torch_trainer.py),这正是TorchTrainer能支持torchrun多进程、而基础DataParallelTrainer只支持单进程的原因。
多节点多 GPU 训练
torchrun同样支持跨节点启动 Local Mode 训练。下面的示例在 2 个节点、每节点 4 块 GPU 上运行。
主节点(192.168.1.1)上执行:
RAY_TRAIN_V2_ENABLED=1 torchrun \ --nnodes=2 \ --nproc-per-node=4 \ --node_rank=0 \ --rdzv_backend=c10d \ --rdzv_endpoint=192.168.1.1:29500 \ --rdzv_id=job_id \ train_script.pyWorker 节点上执行:
RAY_TRAIN_V2_ENABLED=1 torchrun \ --nnodes=2 \ --nproc-per-node=4 \ --node_rank=1 \ --rdzv_backend=c10d \ --rdzv_endpoint=192.168.1.1:29500 \ --rdzv_id=job_id \ train_script.py其中--rdzv_backend=c10d使用 PyTorch 内置的 c10d rendezvous 服务完成进程发现,--rdzv_endpoint指向主节点的 IP 与端口,--rdzv_id标识本次作业。Ray Train 通过MASTER_ADDR/MASTER_PORT等环境变量自动感知多节点拓扑,并正确计算node_rank。
警告:与单节点多卡一样,多进程(
torchrun)Local Mode 不支持 Ray Data,请使用DataLoader+DistributedSampler等标准 PyTorch 数据加载方式。
从 Local Mode 平滑过渡到分布式训练
当本地调试完成后,向分布式训练迁移只需要把num_workers改成大于 0 的值,训练函数代码无需任何改动:
trainer = TorchTrainer( train_loop_per_worker=train_func, train_loop_config=config, - scaling_config=ScalingConfig(num_workers=0), + scaling_config=ScalingConfig(num_workers=4, use_gpu=True), )切换之后,Ray Train 会自动接管分布式协调:拉起 4 个 Worker Actor、完成进程组初始化、指标聚合与检查点持久化——这些正是 Local Mode 所省略的编排逻辑。得益于 data_parallel_trainer.py 中本地与分布式两条执行路径的清晰划分,两种模式下的用户侧 API 完全对齐。
局限性与 API 行为差异
Local Mode 提供的是 Ray Train API 的简化实现,目的是在无分布式编排的条件下快速调试,因此部分能力不可用或行为不同。
Local Mode 中不可用的功能
- Worker 级容错:Ray Train 的自动容错(如失败后自动重启 Worker)在 Local Mode 下不生效。即使配置了
ray.train.FailureConfig,其设置也不会被应用。 - 回调(Callbacks):在
ray.train.RunConfig中指定的用户自定义回调不会被触发。 - Ray Data + 多进程训练:使用
torchrun的 Local Mode 多进程训练不支持 Ray Data,请改用标准 PyTorch 数据加载机制。
API 行为差异一览
下表总结了ray.trainAPI 在 Local Mode 下的行为差异,右侧同时给出仓库源码中的对应实现位置,方便对照:
| API | Local Mode 下的行为 | 源码依据 |
|---|---|---|
ray.train.report | 检查点仅保存在内存中(不持久化到存储);忽略checkpoint_upload_mode、checkpoint_upload_fn、validation、delete_local_checkpoint_after_upload等参数;指标仅本地打印日志,不经过上报管线;不触发跨 Worker 同步屏障 | train_fn_utils.py:LocalTrainFnUtils.report仅记录_last_metrics/_last_checkpoint并logger.info |
ray.train.get_checkpoint | 返回内存中最后一个检查点;不从持久化存储加载 | train_fn_utils.py |
ray.train.get_all_reported_checkpoints | 始终返回空列表;不追踪检查点历史 | train_fn_utils.py:直接返回[] |
ray.train.collective.barrier | 空操作(No-op) | train_fn_utils.py:pass |
ray.train.collective.broadcast_from_rank_zero | 原样返回数据,不做广播 | train_fn_utils.py:return data |
ray.train.get_context().get_storage() | 抛出NotImplementedError | context.py |
此外,LocalTrainFnUtils.is_distributed()恒为False、get_preemption_info()恒为None(单进程内没有抢占监控),Result.path为None(无持久化存储路径),这些差异从实现上解释了"本地模式为何如此轻量"。
常见陷阱与排查建议
综合文档与源码,使用 Local Mode 时有几个值得注意的点:
RAY_TRAIN_V2_ENABLED=1必须显式设置:V2 Trainer API 在未启用 V2 时会直接抛出ValueError(data_parallel_trainer.py)。多节点torchrun示例中已包含该环境变量,单机场景也不要遗漏。- 多进程 + Ray Data 会硬报错:
torchrun模式下若同时传入datasets,控制器会断言失败并提示改用标准DataLoader(torch.py)。单进程模式才支持 Ray Data。 - 下载型数据集注意多进程竞争:多个进程同时下载数据集可能互相干扰,参考示例中的
FileLock+ 仅 local rank 0 下载 +dist.barrier()模式。 - Local Mode 不等于"免费分布式":
report不做跨进程同步、barrier是空操作,涉及集体通信语义的逻辑仍需在真实分布式环境下回归验证。
总结
Local Mode 是 Ray Train V2 中衔接"本地开发"与"分布式生产"的桥梁:num_workers=0一行配置即可获得无 Actor 开销的快速调试环境,配合 pytest 可以做训练逻辑的单测验证,配合torchrun可以在单机多卡/多机多卡上以熟悉的方式调试分布式训练;而迁移到分布式时,只需把num_workers调回正值。理解其"内存化上报、无持久化、简化集体通信"的实现边界(可对照 train_fn_utils.py 与 torch.py),能帮助你在调试效率与行为保真之间做出正确取舍。
- 人工智能
- 分布式训练
- 强化学习
- 任务调度
- 模型推理服务
【免费下载链接】ray
Ray is an AI compute engine. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.
相关推荐
Ray Train V2 完整 API 参考指南:从 Trainer 到分布式训练工具箱
Ray Train V2 完整 API 参考指南:从 Trainer 到分布式训练工具箱 导读 本文是 Ray 项目 doc/source/train/api/
人工智能分布式训练强化学习任务调度模型推理服务Ray Train 分布式 PyTorch 训练配置实战:ScalingConfig、RunConfig 与 TorchTrainer 完整指南
Ray Train 分布式 PyTorch 训练配置实战:ScalingConfig、RunConfig 与 TorchTrainer 完整指南 本文以 Ray
人工智能分布式训练强化学习任务调度模型推理服务Ray Train 实战:将 PyTorch Lightning 训练脚本改造为分布式训练
Ray Train 实战:将 PyTorch Lightning 训练脚本改造为分布式训练 本指南以 Ray 开源仓库中 getting started pyt
人工智能分布式训练强化学习任务调度模型推理服务
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考