https://github.com/volcengine/verl/blob/main/verl/trainer/main_ppo.py#L62
这段代码实现了 PPO 分布式训练的主流程:首先检查 Ray 是否已初始化,若未初始化则根据配置启动 Ray 集群并设置相关环境变量;随后根据配置判断是否需要性能分析(nsight),创建远程 TaskRunner 实例并在集群中异步执行训练任务,最后可选地保存 Ray 的 timeline 性能分析文件。整体实现了分布式训练的启动、调度和性能追踪。
# Define a function to run the PPO-like training process
def run_ppo(config) -> None:
...
if OmegaConf.select(config.trainer, "profile_steps") is not None and len(OmegaConf.select(config.trainer, "profile_steps")) > 0:
nsight_options = OmegaConf.to_container(config.trainer.controller_nsight_options)
runner = TaskRunner.options(runtime_env={"nsight": nsight_options}).remote()
else:
runner = TaskRunner.remote()
ray.get(runner.run.remote(config))
...
主要就是有一个经典的Ray启动任务流程,如下:
runner = TaskRunner.remote()
ray.get(runner.run.remote(config))
https://github.com/volcengine/verl/blob/main/verl/trainer/main_ppo.py#L91
这段代码定义了一个基于 Ray 的远程任务类 TaskRunner,用于在分布式环境下启动和管理 PPO 强化学习训练流程。其核心流程如下:
环境准备与模型加载
首先打印当前节点信息,并解析配置。随后通过 copy_to_local 自动将模型权重从 HDFS 或本地路径加载到本地磁盘(或共享内存),并据此初始化分词器(tokenizer)和处理器(processor),支持多模态和远程代码信任选项。
策略与资源适配
根据配置选择不同的分布式策略(如 FSDP、Megatron),动态导入对应的 worker 类和 worker group 类,并为每种角色(如 Actor、Critic、RewardModel、RefPolicy)注册 Ray 远程 worker,实现灵活的分布式资源调度。
这部分主要关注不同策略有不同的worker,我们就先关注FSDP的worker即可:
https://github.com/volcengine/verl/blob/76f63cffa5081564d8fea93a1cb3ce8bd5bdcc39/verl/workers/fsdp_workers.py#L95
选取worker之后,用ray.remote()创建一个ray的远程actor。
后续有注册奖励函数,自定义奖励函数在如下文件中自定义实现:
奖励函数与数据集构建
根据配置加载多源奖励函数(支持规则、模型、代码沙箱等多种奖励方式),并构建训练和验证数据集及采样器,适配 RLHF 等场景。
训练器初始化与启动
最后实例化 RayPPOTrainer,传入所有必要组件和资源池,初始化所有分布式 worker,并正式启动训练流程。
整体上,这段代码实现了一个高度模块化、可扩展的分布式 PPO 训练主控逻辑,适用于大规模 RLHF 或自定义奖励的强化学习任务。