ChatGLM-LoRA-RLHF-PyTorch
在消费级显卡(RTX 2080 Ti 12GB)上完成 ChatGLM-6B 的 LoRA + RL
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
在消费级显卡(RTX 2080 Ti 12GB)上完成 ChatGLM-6B 的 LoRA + RL
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
2023 年初,ChatGPT 和 GPT-4 的横空出世让整个 AI 圈为之震撼。然而这些强大模型背后依赖的 RLHF(人类反馈强化学习)训练方法,对于普通开发者和小型研究团队来说,几乎是一个"不可能完成的任务"。RLHF 需要三阶段的复杂流程——有监督微调(SFT)、奖励模型训练(RM)和强化学习微调(PPO)——每一步都需要海量 GPU 算力支撑,动辄需要 A100 这样价值数万元的专业计算卡。
在这样的背景下,一个来自个人开发者的开源项目悄然出现,试图用最接地气的方式解决这个问题。jackaduma/ChatGLM-LoRA-RLHF-PyTorch 的目标非常明确:让普通研究者只需要一块消费级 RTX 2080 Ti(12GB 显存)就能跑完 ChatGLM-6B 的 RLHF 训练全流程。
这个项目从 2023 年 4 月发布至今,持续活跃维护,README 最近一次更新为 2026 年 4 月,吸引了不少对大模型微调感兴趣的中文开发者社区用户。
全参数微调 60 亿参数的大模型需要约 120GB 显存(FP16),这远超消费级显卡的上限。LoRA(Low-Rank Adaptation) 的核心思想是"不要微调整个模型权重矩阵,而是只训练少量附加的低秩矩阵"。
具体来说,对于原始权重矩阵 W₀ ∈ ℝ^{d×k},LoRA 冻结 W₀,然后新增两个低秩矩阵 A ∈ ℝ^{r×k} 和 B ∈ ℝ^{d×r}(其中 r ≪ min(d, k)),最终输出为:
W = W₀ + BA
在 ChatGLM-6B 中,LoRA 配置为 rank=8、alpha=16,仅更新约 0.1% 的参数,显存占用从 120GB 降至约 12-14GB,使得 2080 Ti 成为可能。
该项目提供了 RLHF 的完整三阶段流程:
第一阶段:有监督微调(SFT) supervised_finetune.py 实现了基于 HuggingFace Trainer 的标准微调流程。代码中关键的显存优化手段包括:
load_in_8bit=True(INT8 量化)—— 权重加载时即做量化gradient_checkpointing_enable() —— 用时间换空间,前向传播时不保存中间激活值gradient_accumulation_steps=32 —— 每 32 步才更新一次参数,等效 batch size=32第二阶段:奖励模型训练(RM)
train_reward_model.py 和 reward_model.py 实现了自定义奖励模型。由于 ChatGLM 官方暂不支持 HuggingFace 的 AutoModelForSequenceClassification 接口,项目作者自行实现了 RewardModel 类。
RewardModel 的核心是 PairWiseLoss(成对损失):
probs = torch.sigmoid(chosen_reward - reject_reward)
log_probs = torch.log(probs)
loss = -log_probs.mean()
即对于同一个 prompt,让"好回答"的奖励值高于"差回答",通过 sigmoid 概率化后最大化正确排序的对数似然。
第三阶段:PPO 强化学习(README 中标注为 TODO)
从代码中可以看到多个针对 ChatGLM 架构的特殊处理:
项目结构清晰,所有脚本均围绕 RLHF 三阶段组织:
| 文件 | 作用 | 复杂度 |
|---|---|---|
| supervised_finetune.py | SFT 训练 | ★★★ |
| train_reward_model.py | 奖励模型训练 | ★★★★ |
| reward_model.py | 奖励模型核心类 | ★★★ |
| merge_peft_adapter.py | LoRA 适配器合并到基座模型 | ★★ |
| tokenize_dataset_rows.py | 数据预处理(JSONL → tokenized dataset) | ★★ |
| cover_alpaca2jsonl.py | Alpaca 格式转换 | ★ |
| data/ | 训练数据目录 | — |
| models/ | 本地 ChatGLM 模型代码 | — |
| misc/ | 打赏二维码(作者接地气) | — |
优点:
不足:
| 资源 | 最低要求 | 推荐配置 |
|---|---|---|
| GPU | RTX 2080 Ti 12GB | RTX 3090/4090 24GB |
| RAM | 32GB | 64GB |
| 显存 | 12GB+ | 24GB+ |
| 磁盘 | 30GB | 100GB+ |
| 训练时间 | 10-20小时/epoch | 大幅缩短 |
与开箱即用的推理工具不同,这个项目需要手动完成以下步骤:
这是一个面向有深度学习训练经验的开发者项目,不适合新手直接上手。
jackaduma 的工作填补了中文开源社区在 ChatGLM RLHF 训练方案上的空白。尽管 ChatGLM 官方后来推出了更完善的微调工具链(如 ChatGLM-Finetuning),但该项目代码逻辑清晰、注释详细,对于希望深入理解 RLHF 实现细节的学习者来说,是一个难得的低门槛入门素材。
更重要的是,项目将 LoRA 量化技术与 RLHF 训练结合的思路,为后来众多 LoRA+RLHF 项目(如 Stanford AlpacaFarm、RL4LMs 等)提供了中文社区的参考实现。
ChatGLM-LoRA-RLHF-PyTorch 是一个专注于"让 RLHF 训练在消费级硬件上落地"的技术实践项目。它以 ChatGLM-6B 为基座模型,通过 LoRA 量化大幅降低显存需求,提供了 SFT 和 RM 两个完整训练阶段的开源代码。对于 AI 研究者和开发者来说,这个项目的价值在于提供了一个透明、可读的 RLHF 训练实现参考;而对于想快速微调自己数据集的用户,建议关注 ChatGLM 官方的最新微调方案或 Axolotl、TRL 等更成熟的训练框架。