ReaLHF
通过参数重分配动态优化多卡并行策略,显著提升RLHF训练吞吐量的分布式LLM对齐框架
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
通过参数重分配动态优化多卡并行策略,显著提升RLHF训练吞吐量的分布式LLM对齐框架
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
想象一个场景:你有 8 张 A100 显卡,想用 PPO 算法微调一个 70B 参数的大语言模型。理想情况下,训练应该跑得飞快——但实际情况是,GPU 大面积闲置,通信等待时间比真正计算还长。问题出在哪里?
DeepSpeed-Chat 和 OpenRLHF 等现有开源 RLHF 系统存在两大通病:要么对所有模型组件采用相同的并行策略(导致过度并行化,同步开销巨大);要么让不同模型跑在不同 GPU 节点上(导致部分 GPU 因任务依赖而空转)。这就是 OpenPSI 团队在开发 ReaL 时所面对的现实困境。
ReaL(ReaLlocation,参数重分配)诞生于清华大学和 OpenPSI 公司的联合研究,其成果以 ICML 2024 Oral 论文 "Is DPO Superior to PPO for LLM Alignment? A Comprehensive Study" 为契机逐步成型。核心论文 "ReaLHF: Optimized RLHF Training for Large Language Models through Parameter Reallocation" 于 2024 年 6 月发布。
核心创新点:ReaL 打破传统思路,不再将整个 RLHF 流程中所有模型组件(Actor、Critic、Reward、Reference)固定在相同的并行策略上,而是让不同模型调用动态选择最优并行方式——生成阶段用流水线并行(Pipeline Parallelism),训练阶段用张量并行(Tensor Parallelism)——并在执行时让不同模型调用以更小的并行度并发运行,从而最大化 GPU 利用率,同时消除冗余通信。
注意:该项目已于 2025 年归档,代码开发已迁移至新项目 AReaL,但当前仓库仍可正常访问和使用。
可以把 ReaL 想象成一个超级购物中心的智能导览系统。在传统方案里,所有店铺(模型组件)都被安排在相同的楼层(GPU 节点),要么楼层太大导致信息传递缓慢(过度并行),要么某些楼层客流量极少(GPU 空转)。ReaL 的做法是:根据每家店铺的实际需求(计算负载)动态分配楼层(GPU 资源),生成阶段让所有员工参与(高并行度),训练阶段让核心团队精干执行(低并行度),两者同时进行、各取所需。
ReaL 提供了覆盖 RLHF 全流程的功能模块:
1. 算法数据流图(Dataflow Graph)架构
ReaL 引入了一种创新的算法表示方式——将每个 RLHF 算法建模为有向无环图(DAG),图的节点是模型函数调用(MFC),边描述数据依赖或参数版本依赖。例如,PPO 算法被分解为 6 个独立的 MFC:Actor 生成、Critic 推理、Reward 推理、Reference 推理、价值估计(GAE)和策略训练,每个 MFC 分配到不同形状的 GPU 网格上执行。这种设计让用户可以精细控制每个模型调用使用的并行策略和数据分布。
2. 参数重分配引擎
这是 ReaL 最重要的创新。训练过程中,Actor 模型参数会在不同 GPU 之间动态迁移,Scheduler 根据当前负载实时调整各 MFC 的设备网格(Device Mesh)。系统使用 CUDA 自定义 kernels 实现高效的参数读写操作(在 csrc/ 目录下),确保参数迁移本身不成为瓶颈。
3. 高吞吐量生成
集成 vLLM 的 custom all-reduce 和 CUDA Graph 优化,支持大规模并行生成。v0.3.0 版本还引入了 MoE(混合专家)模型的 RLHF 训练支持。
4. 丰富的 RLHF 算法支持
除 PPO 外,还支持 DPO、RAFT、GRPO、ReMax 等多种对齐算法,扩展新算法只需实现对应的数据流图即可接入 ReaL 的高效基础设施。
目录结构:
| 模块 | 说明 |
|---|---|
realhf/api/ | 数据流图、模型 API、配置对象的定义,HuggingFace 模型到 ReaL 的转换 |
realhf/apps/ | 实验启动入口,通过 Hydra 配置驱动,main_start 协调 Scheduler 启动远程进程 |
realhf/impl/ | 数据集、模型接口、模型后端的实现,包含与 Megatron-LM/DeepSpeed 的深度集成 |
realhf/scheduler/ | 类 torchrun 的分布式启动器,支持 Ray/SLURM |
realhf/system/ | Worker 定义(Base Worker、Model Worker、Master Worker) |
realhf/search_engine/ | 自动搜索最优配置的引擎(当前版本未激活) |
csrc/ | CUDA/C++ 扩展:GAE 计算 kernel、参数迁移 kernel |
核心依赖:PyTorch 2.3+、DeepSpeed 0.14.0、Megatron-LM Core 0.6.0、vLLM、HuggingFace Transformers 4.42.3、Ray、Hydra 配置框架。所有依赖均通过 requirements.txt 统一管理。
测试覆盖:项目有 tests/ 目录,包含非 GPU/非分布式的单元测试,每次 PR 均会自动运行。
1. 仓库已归档:ReaL 于 2025 年正式归档,开发已迁移至 AReaL。新用户建议直接使用 AReaL 项目。
2. Python 版本限制:仅支持 Python 3.10~3.11,3.12 及以上版本不兼容,安装时需注意系统 Python 版本。
3. 编译门槛高:若无预构建 Docker 镜像,需手动编译 PyBind11 和 CUDA 扩展,pip install -e . --no-build-isolation 过程中容易因环境差异失败。
4. 硬件要求苛刻:训练 70B 模型至少需要 8 卡 80GB GPU 集群,实验门槛极高,不适合个人开发者或资源有限的研究团队。
ReaL 的参数重分配思想对 RLHF 训练系统设计具有重要启发意义。通过实测,ReaL 在 4×LLaMA-7B 配置下可完成完整 RLHF 流程(SFT/RW/PPO)仅需 30 分钟,在 LLaMA-70B 规模下吞吐量显著优于 DeepSpeed-Chat 和 OpenRLHF。
该项目还推动了学术与工程实践的结合:它不只是一个实验系统,其论文中的方法论(动态并行策略选择、异步模型执行)已被工业界广泛参考。随着 AReaL 的继续开发,ReaL 的核心思想将在更成熟的系统中延续生命力。
图1:ReaL 项目 Logo(来源:GitHub 仓库 docs/source/images/)
图2:PPO 训练 Reward 曲线(来源:GitHub 仓库 docs/source/images/)
图3:SFT 损失曲线(来源:GitHub 仓库 docs/source/images/)