news 2026/9/20 10:55:20

Ray Train V2 Local Mode 完整实战指南:不启动 Worker 的分布式训练本地调试方案

作者头像

张小明

前端开发工程师

1.2k 24
文章封面图
Ray Train V2 Local Mode 完整实战指南:不启动 Worker 的分布式训练本地调试方案
  • 人工智能
  • 分布式训练
  • 强化学习
  • 任务调度
  • 模型推理服务

【免费下载链接】ray

Ray is an AI compute engine. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.

项目地址:https://gitcode.com/gh_mirrors/ra/ray
点击查看免费下载

导读

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 Noneresult.metrics["loss"]非空、result.checkpoint存在。

Local Mode 适用于所有 Ray Train 框架集成,包括 PyTorch Lightning、Hugging Face Transformers、LightGBM、XGBoost、TensorFlow 等。测试目录 test_local_mode.py 中分别有test_lightning_trainer_local_modetest_e2e_hf_local_modetest_lightgbm_trainer_local_modetest_xgboost_trainer_local_modetest_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.5

test_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()会依次执行:

  1. 调用torch.distributed.init_process_group(backend="nccl" if torch.cuda.is_available() else "gloo")初始化进程组(torch.py);
  2. 从分布式环境读取world_sizeglobal_rank,从环境变量读取local_rank
  3. GPU 可用时执行torch.cuda.set_device(local_rank)绑定当前进程到对应设备;
  4. 计算nproc_per_node = LOCAL_WORLD_SIZEnode_rank = global_rank // nproc_per_node
  5. 把这些值写入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()==4get_world_rank()==2get_local_rank()==1get_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.py

Worker 节点上执行:

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 下的行为差异,右侧同时给出仓库源码中的对应实现位置,方便对照:

APILocal Mode 下的行为源码依据
ray.train.report检查点仅保存在内存中(不持久化到存储);忽略checkpoint_upload_modecheckpoint_upload_fnvalidationdelete_local_checkpoint_after_upload等参数;指标仅本地打印日志,不经过上报管线;不触发跨 Worker 同步屏障train_fn_utils.py:LocalTrainFnUtils.report仅记录_last_metrics/_last_checkpointlogger.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()抛出NotImplementedErrorcontext.py

此外,LocalTrainFnUtils.is_distributed()恒为Falseget_preemption_info()恒为None(单进程内没有抢占监控),Result.pathNone(无持久化存储路径),这些差异从实现上解释了"本地模式为何如此轻量"。

常见陷阱与排查建议

综合文档与源码,使用 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.

项目地址:https://gitcode.com/gh_mirrors/ra/ray
点击查看免费下载

相关推荐

上一篇:PROJ核心功能解析:深度理解坐标参考系统(CRS)转换
下一篇:ESP-IDF 中使用 esp-iot-solution GProf 组件做 GNU Profiler 性能分析

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

版权声明: 本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若内容造成侵权/违法违规/事实不符,请联系邮箱:809451989@qq.com进行投诉反馈,一经查实,立即删除!
网站建设 2026/9/20 10:54:11

前端转战AI应用开发:手把手打造内部知识库问答助手

前端这个圈子这些年有个很有趣的现象:一到技术转型节点,跳得最欢的往往不是后端,而是天天跟页面打交道的前端。前两篇我们聊了本地大模型部署和对话网页怎么搭,今天这篇我打算换个节奏,从一个真实需求出发,…

作者头像 李华
网站建设 2026/9/20 10:54:05

eNSP安装避坑指南:从VirtualBox版本选择到高频报错排查

/* MD / 富文本中的 .toc(含博客园搬家等嵌套结构);.toc-box 在侧栏,不受影响 */#content_views .toc,/* 编辑器常在目录前后插入空 p(:empty 仍占 20px),一并去掉避免顶空隙 */#content_views.markdown_views > p:empty:has(+ .toc),#content_views.markdown_views …

作者头像 李华
网站建设 2026/9/20 10:53:34

OpenResearch本地优先研究工作流:原理、验证与AI可替换实践

1. 项目概述:一个被误读的开源研究协作范式“OpenResearch”这个词最近在开发者社区里频繁出现,但很多人一看到就下意识联想到某个具体工具、某个CLI命令,甚至直接去搜“orx install”或者“autoresearch setup”。其实这恰恰暴露了一个普遍存…

作者头像 李华
网站建设 2026/9/20 10:52:28

GAC水平集图像分割:PDE驱动的自演化边界模型

简介:本资源是面向图像处理初学者与计算机视觉学习者的Matlab实践项目,聚焦Geodesic Active Contours(GAC)水平集图像分割算法的完整实现,解决边界模糊、光照不均等典型图像分割难题,适用于医学影像分析、目…

作者头像 李华
网站建设 2026/9/20 10:52:25

第四次工业革命:从技术堆叠到系统重构的产业逻辑

1. 从“技术堆叠”到“系统重构”:产业逻辑到底在变什么很多技术人聊到第四次工业革命,第一反应是“AI、物联网、大数据、云计算”,然后把这些词像贴标签一样往PPT上堆。但如果你真的在制造业、物流、能源或者医疗行业待过,就会发…

作者头像 李华
网站建设 2026/9/20 10:51:53

GitHub热榜观察与国内开发者实用指南:从访问加速到项目评估

最近几个月,我每天打开 GitHub 热榜的次数,比打开朋友圈还勤。这个习惯从 2023 年 AI 应用爆发那阵子养成的,一直保留到现在。GitHub 热榜基本就是全球开发者用代码投票出来的“流行风向标”,你不需要读论文、刷资讯,只…

作者头像 李华