ColossalChat on Ray：在 Ray 集群上分布式运行 ColossalChat PPO 训练的完整指南

ColossalChat on Ray在 Ray 集群上分布式运行 ColossalChat PPO 训练的完整指南【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI本文围绕仓库中 ColossalChat 社区示例文档 展开讲解如何把 ColossalChat基于 ColossalAI 的 RLHF/ChatGPT 克隆训练流水线的 PPO 训练任务提交到 Ray 集群上运行。读完本文你将掌握 Ray Job 提交、runtime_env 依赖打包、以 Ray Actor 承载多角色 PPO 模型actor / critic / initial / reward的分布式训练框架并能读懂train_prompts_on_ray.py中 torch.distributed 与 ColossalAI 策略DDP / Zero-2 / Gemini的整合方式。一、背景ColossalChat 与 Ray 的结合点ColossalChat 是 ColossalAI 项目下的开源应用目标是提供一条完整的 RLHFReinforcement Learning with Human Feedback流水线用于克隆 ChatGPT 式的对话模型。其核心训练算法是 PPOProximal Policy Optimization一个 PPO 训练过程通常同时需要四类模型角色Actor策略模型负责根据 prompt 采样生成序列并根据经验experience更新参数Critic价值模型估计序列价值与 reward 共同计算 advantageInitial Model训练前的初始策略副本用于计算 baseline 的 action log probabilityKL 惩罚项的参照Reward Model奖励模型对生成序列打分。由于四类模型都要常驻显存并参与计算单机往往难以承载因此社区示例尝试把这套训练放到 Ray 集群上用 Ray Actor 分别托管不同角色的模型分片。原文档开头明确声明了这是一个实验性experimental工作并且由于 ColossalChat 的大版本更新相关内容可能已经过时——这一点在实操前必须牢记详见文末“适用前提与已知限制”。二、四步操作从 Ray 集群到 Dashboard 查看训练1. 搭建支持 GPU 的 Ray 集群按照 Ray 官方的集群搭建文档建立一个带 GPU 支持的集群Head 节点 Worker 节点。搭建完成后记录集群 API Server 的 endpoint其格式形如http://your.head.node.address:8265这个 endpoint 后面提交 Ray Job 时作为参数传入。2. 克隆仓库git clone https://github.com/hpcaitech/ColossalAI.git3. 提交 Ray Job原文档给出的提交命令是python applications/Chat/examples/community/ray/ray_job_script.py http://your.head.node.addrees:8265脚本接收唯一的位置参数——Ray 集群 API Server 的 endpoint。该脚本内部会创建一个JobSubmissionClient并向集群提交一个训练作业。4. 在 Ray Dashboard 上查看作业打开 Ray 集群的 Dashboard即http://your.head.node.address:8265可以实时查看已提交训练任务的运行状态、日志与资源占用。注意原文档中的命令路径applications/Chat/...是旧版目录结构。在当前仓库中该示例实际位于 applications/ColossalChat/examples/community/ray/ 下且 社区示例总览文档 已将 Train prompts on Ray 列为社区贡献示例之一。执行时应使用当前仓库的实际路径python applications/ColossalChat/examples/community/ray/ray_job_script.py http://your.head.node.address:8265三、ray_job_script.py 源码剖析Ray Job 是如何打包和提交的提交端脚本 ray_job_script.py 非常精简其核心逻辑如下from ray.job_submission import JobSubmissionClient def main(api_server_endpointhttp://127.0.0.1:8265): client JobSubmissionClient(api_server_endpoint) client.submit_job( entrypointpython experimental/ray/train_prompts_on_ray.py --strategy colossalai_zero2 --prompt_csv_url ..., runtime_env{ working_dir: applications/Chat, pip: [ torch1.13.1, transformers4.20.1, datasets, loralib, colossalai0.2.4, langchain, tokenizers, fastapi, sse_starlette, wandb, sentencepiece, gpustat, ], }, ) if __name__ __main__: main(sys.argv[1])结合源码可以拆解出三个关键点entrypointRay Job 在集群上执行的入口命令。这里运行的是 train_prompts_on_ray.py并传入--strategy colossalai_zero2使用 ColossalAI 的 Zero-2 并行策略与--prompt_csv_url一个公开的 awesome-chatgpt-prompts 数据集的 CSV 地址作为 PPO 训练的 prompt 来源。runtime_env.working_dir指定作业的工作目录。脚本中写的是旧路径applications/Chat在当前仓库布局下应对应applications/ColossalChat。Ray 会把该目录上传到集群各节点作为作业的工作区。runtime_env.pipRay 的依赖隔离机制——在作业运行环境中按列表安装指定版本依赖。注意其中锁定了torch1.13.1与colossalai0.2.4说明该示例面向的是 ColossalAI 0.2.x 时代的接口例如 Zero-2 是当时的 LowLevelZero stage 2这与其实验性、可能过时的声明相互印证。四、train_prompts_on_ray.py 深度解析Ray Actor 承载的四角色 PPO真正的训练脚本 train_prompts_on_ray.py 是一个约 570 行的自包含实现。它把 PPO 的各角色封装成 Ray Actor并用Actor Group统一调度。4.1 分布式环境搭建DistributedTorchRayActor所有模型角色都继承自BasePPORole - DistributedTorchRayActor。DistributedTorchRayActor.__init__负责为每个 Ray Actor 进程构造标准 PyTorch 分布式环境变量os.environ[MASTER_ADDR] self._master_addr os.environ[MASTER_PORT] str(self._master_port) os.environ[WORLD_SIZE] str(self._world_size) os.environ[RANK] str(self._rank) os.environ[LOCAL_RANK] str(self._local_rank)其中两个细节值得注意MASTER_ADDR 自动发现rank 0 的 master actor 未显式指定地址时通过ray._private.services.get_node_ip_address()获取本机 IP 作为 master 地址其余 worker 通过master_actor.get_master_addr_port()拉取同一地址从而保证跨节点 rendezvous 一致MASTER_PORT 自动选空闲端口通过socket.bind((, 0))让操作系统分配一个未占用的端口避免多作业端口冲突。这意味着每个 Ray Actor 进程内部就是一个标准的 torch.distributed 进程组成员Ray 负责进程与 GPU 的调度torch.distributed 负责进程间集合通信ColossalAI 策略负责参数/梯度切分。4.2 GPU 资源锁定Placement GroupPPORayActorGroup._initiate_actors展示了如何用 Ray Placement Group 为同一角色的所有 worker 锁定资源world_size self._num_nodes * self._num_gpus_per_node if self._num_gpus_per_node 1: bundles [{GPU: self._num_gpus_per_node, CPU: self._num_gpus_per_node} for _ in range(self._num_nodes)] pg placement_group(bundles, strategySTRICT_SPREAD) ray.get(pg.ready())每个 bundle 对应一个节点上的num_gpus_per_node张 GPUSTRICT_SPREAD策略保证 bundle 均匀打散到不同节点避免多 GPU 角色挤在同一台机器上master actor 被调度到 bundle 0后续 worker 按placement_group_bundle_indexrank // num_gpus_per_node依次落到对应节点local_rank rank % num_gpus_per_node这是典型的全局 rank - (节点, 节点内 rank)映射当num_gpus_per_node 1时则退化为简单的options(num_gpus1)调度。每个具体角色如RayPPOActor都装饰了ray.remote(num_gpus1)即一个 Ray Actor 独占一张 GPU。4.3 训练策略与优化器DDP / Zero-2 / GeminiBasePPORole._init_strategy支持三种并行策略--strategy取值底层实现说明ddp默认DDPStrategy朴素数据并行colossalai_zero2LowLevelZeroStrategy(stage2, placement_policycuda)ColossalAI 低层 Zero-2优化器状态 梯度分片colossalai_geminiGeminiStrategy(placement_policycuda, initial_scale2**5)ColossalAI 的 Gemini 显存管理initial_scale32为混合精度初始损失缩放因子优化器选择与之联动_init_optimizer使用 ColossalAI 策略时配套HybridAdamColossalAI 的分布式优化器见 colossalai/nn/optimizer否则使用原生torch.optim.Adam学习率统一为5e-6。模型加载发生在strategy.model_init_context()上下文中配合 ColossalAI 的分片初始化语义。从源码结构看这些策略类从coati.trainer.strategies导入而当前仓库的 coati/trainer 目录 中已不再提供该模块trainer 下现为 ppo/sft/dpo/kto/orpo/grpo/rm 等算法文件说明脚本面向的是旧版 coati API直接在新版环境运行会命中导入错误。4.4 经验制作ExperienceMaker 与 Advantage 计算PPO 的一步训练需要五元组信息序列含 attention mask / action mask、动作 log prob、初始策略 log prob、价值估计、奖励。脚本中ExperienceCompositionRefs把这些都存成ray.ObjectRef零拷贝引用避免数据在 Actor 间复制ExperienceMaker.make_experience才真正ray.get拉取并组装reward compute_reward(r, self.kl_coef, action_log_probs, base_action_log_probs, action_maskaction_mask) advantage reward - value experience Experience(sequences, action_log_probs, value, reward, advantage, attention_mask, action_mask)即奖励中扣除了以kl_coef默认 0.1加权的 KL 惩罚项advantage 采用最直接的reward - value估计。4.5 训练主循环episode / timestep / 经验队列main(args)中的训练流程可以概括为按参数选择模型类--model支持gpt2/bloom/opt分别映射到GPTActor/GPTCritic、BLOOMActor/BLOOMCritic、OPTActor/OPTCritic均带 LoRA 包装见coati.models创建四个 Actor Groupactor / critic / initial / reward各自可配置独立节点数--num_actor_nodes等默认各 1 节点与每节点 GPU 数并行初始化四组模型通过ray.get([...init refs...])一次性并发加载actor、critic 带优化器initial 与 reward 只前向、不建优化器actor 组随后加载 tokenizer、设置采样参数max_length128, do_sampleTrue, temperature1.0, top_k50并从 CSV URL 装载 prompt 采样器训练循环num_episodes×max_timesteps经验排队阶段actor 组采样 prompt 并生成序列同时 initial 组、critic 组、reward 组、actor 组对同一批序列异步计算 baseline log prob、value、reward 与当前 log prob全部以 ObjectRef 形式入队到experience_composition_refs学习阶段当累计步数满足time % update_timesteps 0时把队列中所有五元组分发给空闲 Ray Actor 制作成Experience然后调用 actor 组与 critic 组的async_learn_on_experiences并行反向更新最后清空队列。经验按experience_refs[i::num_actors]轮转分片到各 Actor实现负载均衡。actor 的训练步是标准 PPO clip 策略损失PolicyLoss(eps_clip0.2)critic 是ValueLoss(value_clip0.4)每步均走strategy.backward - optimizer_step - zero_grad从而兼容 DDP 与 Zero/Gemini 策略的通信语义。训练结束后由 rank 0 保存模型 checkpoint可选保存优化器 checkpoint--need_optim_ckpt。4.6 完整命令行参数表脚本通过argparse暴露的全部参数及默认值如下摘自train_prompts_on_ray.py的__main__段参数类型默认值说明--prompt_csv_urlstr无必填prompt 数据集 CSV 的 URL脚本用 pandas 读取其prompt列--strategystrddp可选ddp/colossalai_gemini/colossalai_zero2--modelstrgpt2可选gpt2/bloom/opt--pretrainstrgpt2预训练模型名称或路径--save_pathstractor_checkpoint_prompts.ptactor 模型 checkpoint 保存路径--need_optim_ckptboolFalse是否同时保存优化器 checkpoint--num_episodesint10训练 episode 数--max_timestepsint10每个 episode 的 timestep 上限--update_timestepsint10累计多少 timestep 后触发一次参数更新--train_batch_sizeint8训练批大小参数保留主流程以经验队列为单位--experience_batch_sizeint8每次采样生成的经验批大小--num_actor_nodesint1托管 actor 模型的节点数--num_critic_nodesint1托管 critic 模型的节点数--num_initial_nodesint1托管 initial 模型的节点数--num_reward_nodesint1托管 reward 模型的节点数--num_gpus_per_nodeint1每个 Ray 节点上的 GPU 数注意四个角色各自独立配置节点数这是该示例最灵活的扩展方式可以只在 4 台机器上各部署一个角色也可以把某个角色例如只需前向的 reward 模型单独放到更多节点上资源分配完全由这四个--num_*_nodes参数决定。五、与 coati 内置 Ray 训练器的关系仓库中另有一套更成熟的 Ray 集成位于 coati/ray 包其文档 coati/ray/README.md 介绍了DetachedPPOTrainer / ExperienceMakerHolder的maker-trainer 解耦架构经验制作端maker专职推理与经验生产并远程投递给训练端trainertrainer 周期性把新参数回传给 maker借助 experience buffer 重叠传输与计算。通过.options(name...)命名 Actor 并互相引用可以自定义 2 Maker 1 Trainer、2 Maker 2 Trainer、maker 侧推理量化乃至张量并行等多种拓扑。对比之下examples/community/ray这个社区示例是更早的单进程脚本式探索四类角色各自独立部署、通过ray.ObjectRef队列串起 PPO 主循环没有引入 maker/trainer 解耦层。两者放在一起读可以看到 ColossalChat 的 Ray 化路线从社区实验脚本起步逐步沉淀为 coati 内置的可配置分布式 PPO 训练框架。六、适用前提与已知限制使用本示例前请核对以下事实均来自当前仓库的实际内容官方声明过时风险README.md 与 社区示例总览 开头均有警告——This content may be outdated since the major update of Colossal Chat应把它当作教学性/参考性示例而非可直接照搬的生产脚本。路径已迁移README 中的applications/Chat/examples/community/ray/ray_job_script.py、脚本内working_dirapplications/Chat与 entrypoint 中的experimental/ray/train_prompts_on_ray.py均为旧目录结构当前仓库对应路径是applications/ColossalChat/examples/community/ray/下的 ray_job_script.py 与 train_prompts_on_ray.py。若要在当前仓库布局下运行需相应修正这三处路径引用。依赖版本锁定在旧生态runtime_env 中torch1.13.1、colossalai0.2.4等约束对应旧版 ColossalAIZero-2 等策略 API。从当前仓库源码结构看train_prompts_on_ray.py导入的coati.trainer.strategiesDDPStrategy/GeminiStrategy/LowLevelZeroStrategy在现在的coati/trainer目录中已不存在直接在新版环境执行会失败需要按当前 coati 的 trainer 接口做适配。硬件前提整个流程以 CUDA GPU 集群为前提torch.cuda.current_device()、num_gpus1调度、to(cuda)等调用遍布脚本不支持纯 CPU 集群。七、小结这个社区示例的价值在于它用不到 600 行代码完整演示了Ray 管进程与资源、torch.distributed 管集合通信、ColossalAI 管显存与参数切分三层职责的叠加方式——Placement Group 锁定 GPU 拓扑、环境变量完成 rendezvous、ObjectRef 队列实现多角色流水线、策略抽象统一 DDP 与 Zero/Gemini 的训练步。即便具体 API 随版本演进需要适配这套架构模式对于把任意多角色强化学习训练PPO 四模型扩展到 Ray 集群都具有直接的可借鉴性若追求与当前代码库保持同步的实现应优先参考 coati/ray 中的 maker-trainer 解耦训练器。【免费下载链接】ColossalAIMaking large AI models cheaper, faster and more accessible项目地址: https://gitcode.com/GitHub_Trending/co/ColossalAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考