mirror of
https://github.com/jingyaogong/minimind.git
synced 2026-01-13 19:57:20 +08:00
[feat] minimind-2510
This commit is contained in:
jingyaogong
parent
3812a16d20
commit
7d877db79b
623
README.md
623
README.md
@ -15,6 +15,12 @@
|
||||
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
|
||||
|
||||
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
<h3>"大道至简"</h3>
|
||||
</div>
|
||||
@ -27,15 +33,13 @@
|
||||
|
||||
* 此开源项目旨在完全从0开始，仅用3块钱成本 + 2小时！即可训练出仅为25.8M的超小语言模型**MiniMind**。
|
||||
* **MiniMind**系列极其轻量，最小版本体积是 GPT-3 的 $\frac{1}{7000}$，力求做到最普通的个人GPU也可快速训练。
|
||||
* 项目同时开源了大模型的极简结构-包含拓展共享混合专家(MoE)、数据集清洗、预训练(Pretrain)、监督微调(SFT)、LoRA微调，
|
||||
直接偏好强化学习(DPO)算法、模型蒸馏算法等全过程代码。
|
||||
* 项目同时开源了大模型的极简结构-包含拓展共享混合专家(MoE)、数据集清洗、预训练(Pretrain)、监督微调(SFT)、LoRA微调、直接偏好优化(DPO)、强化学习训练(RLAIF: PPO/GRPO等)、模型蒸馏等全过程代码。
|
||||
* **MiniMind**同时拓展了视觉多模态的VLM: [MiniMind-V](https://github.com/jingyaogong/minimind-v)。
|
||||
* 项目所有核心算法代码均从0使用PyTorch原生重构！不依赖第三方库提供的抽象接口。
|
||||
* 这不仅是大语言模型的全阶段开源复现，也是一个入门LLM的教程。
|
||||
* 希望此项目能为所有人提供一个抛砖引玉的示例，一起感受创造的乐趣！推动更广泛AI社区的进步！
|
||||
|
||||
> 为防止误解，“2小时” 基于NVIDIA 3090硬件设备（单卡）测试，“3块钱”
|
||||
> 指GPU服务器租用成本，具体规格详情见下文。
|
||||
> 为防止误解，“2小时” 基于NVIDIA 3090硬件设备（单卡）测试，“3块钱”指GPU服务器租用成本，具体规格详情见下文。
|
||||
|
||||
---
|
||||
|
||||
@ -77,7 +81,7 @@
|
||||
这就好比教牛顿如何使用21世纪的智能手机——虽然有趣，却完全偏离了理解物理本质的初衷。
|
||||
与此同时，第三方的大模型框架和工具库，如transformers+trl，几乎只暴露了高度抽象的接口。
|
||||
通过短短10行代码，就能完成“加载模型+加载数据集+推理+强化学习”的全流程训练。
|
||||
这种高效的封装固然便利，但也像一架高速飞船，将我们与底层实现隔离开来，阻碍了深入探究LLM核心代码的机会。
|
||||
这种高效的封装固然便利，但也像一架高速飞船，将开发者与底层实现隔离开来，阻碍了深入探究LLM核心代码的机会。
|
||||
然而，“用乐高拼出一架飞机，远比坐在头等舱里飞行更让人兴奋！”。
|
||||
更糟糕的是，互联网上充斥着大量付费课程和营销号，以漏洞百出、一知半解的内容推销AI教程。
|
||||
正因如此，本项目初衷是拉低LLM的学习门槛，让每个人都能从理解每一行代码开始，
|
||||
@ -106,12 +110,12 @@
|
||||
|
||||
- MiniMind-LLM结构的全部代码（Dense+MoE模型）。
|
||||
- 包含Tokenizer分词器详细训练代码。
|
||||
- 包含Pretrain、SFT、LoRA、RLHF-DPO、模型蒸馏的全过程训练代码。
|
||||
- 包含Pretrain、SFT、LoRA、RLHF-DPO、RLAIF(PPO/GRPO/SPO)、模型蒸馏的全过程训练代码。
|
||||
- 收集、蒸馏、整理并清洗去重所有阶段的高质量数据集，且全部开源。
|
||||
- 从0实现预训练、指令微调、LoRA、DPO强化学习，白盒模型蒸馏。关键算法几乎不依赖第三方封装的框架，且全部开源。
|
||||
- 从0实现预训练、指令微调、LoRA、DPO/PPO/GRPO/SPO强化学习，白盒模型蒸馏。关键算法几乎不依赖第三方封装的框架，且全部开源。
|
||||
- 同时兼容`transformers`、`trl`、`peft`等第三方主流框架。
|
||||
- 训练支持单机单卡、单机多卡(DDP、DeepSpeed)训练，支持wandb可视化训练流程。支持动态启停训练。
|
||||
- 在第三方测评榜（C-Eval、C-MMLU、OpenBookQA等）进行模型测试。
|
||||
- 训练支持单机单卡、单机多卡(DDP、DeepSpeed)训练，支持wandb/swanlab可视化训练流程。支持动态启停训练。
|
||||
- 在第三方测评榜（C-Eval、C-MMLU、OpenBookQA等）进行模型测试，支持YaRN算法执行RoPE长文本外推。
|
||||
- 实现Openai-Api协议的极简服务端，便于集成到第三方ChatUI使用（FastGPT、Open-WebUI等）。
|
||||
- 基于streamlit实现最简聊天WebUI前端。
|
||||
- 全面兼容社区热门`llama.cpp`、`vllm`、`ollama`推理引擎或`Llama-Factory`训练框架。
|
||||
@ -122,7 +126,21 @@
|
||||
### 👉**更新日志**
|
||||
|
||||
<details close>
|
||||
<summary> <b>2025-04-26 (newest 🎉🎉🎉)</b> </summary>
|
||||
<summary> <b>2025-10-24 (newest🎉)</b> </summary>
|
||||
|
||||
- 🔥 新增RLAIF训练算法：PPO、GRPO、SPO（从0原生实现）
|
||||
- 新增RLAIF数据集：rlaif-mini.jsonl（从SFT数据随机采样1万条）
|
||||
- 新增YaRN算法：支持RoPE长文本外推，提升长序列处理能力
|
||||
- Adaptive Thinking：Reason模型可选是否启用思考链
|
||||
- chat_template全面支持Tool Calling和Reasoning标签（`<tool_call>`、`<think>`等）
|
||||
- 新增RLAIF完整章节、训练曲线对比、算法原理折叠说明
|
||||
- [SwanLab](https://swanlab.cn/)替代WandB（国内访问友好，API完全兼容）
|
||||
- 修复一些已知bug
|
||||
|
||||
</details>
|
||||
|
||||
<details close>
|
||||
<summary> <b>2025-04-26</b> </summary>
|
||||
|
||||
- 重要更新
|
||||
- 如有兼容性需要，可访问[🔗旧仓库内容🔗](https://github.com/jingyaogong/minimind/tree/7da201a944a90ed49daef8a0265c959288dff83a)。
|
||||
@ -131,6 +149,7 @@
|
||||
- 🔥支持llama.cpp、vllm、ollama等热门三方生态。
|
||||
- 规范代码和目录结构。
|
||||
- 改动词表`<s></s>`->`<|im_start|><|im_end|>`
|
||||
|
||||
```text
|
||||
为兼容第三方推理框架llama.cpp、vllm，本次更新需付出一些可观代价。
|
||||
本次更新不再支持「直接」加载25-04-26以前的旧模型进行推理。
|
||||
@ -138,6 +157,7 @@
|
||||
MiniMind2系列旧模型均经过权重映射+（微调训练）QKVO线性层校准恢复而来。
|
||||
本次更新后将放弃对`minimind-v1`全系列的维护，并在仓库中下线。
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
<details close>
|
||||
@ -236,7 +256,9 @@ pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
|
||||
```
|
||||
|
||||
### 2.下载模型
|
||||
|
||||
到项目根目录
|
||||
|
||||
```bash
|
||||
git clone https://huggingface.co/jingyaogong/MiniMind2
|
||||
```
|
||||
@ -333,7 +355,7 @@ python train_full_sft.py
|
||||
|
||||
---
|
||||
|
||||
### 4.测试模型效果
|
||||
### 4.测试自己训练的模型效果
|
||||
|
||||
确保需要测试的模型`*.pth`文件位于`./out/`目录下。
|
||||
也可以直接去[此处](https://www.modelscope.cn/models/gongjy/MiniMind2-PyTorch/files)下载使用我训练的`*.pth`文件。
|
||||
@ -345,7 +367,9 @@ python eval_model.py --model_mode 1 # 默认为0：测试pretrain模型效果，
|
||||
<details style="color:rgb(128,128,128)">
|
||||
<summary>注：测试须知</summary>
|
||||
|
||||
如需详情，查看`eval_model.py`脚本代码即可。model_mode分为 0: 预训练模型，1: SFT-Chat模型，2: RLHF-Chat模型，3: Reason模型
|
||||
详情查看`eval_model.py`脚本代码即可。
|
||||
|
||||
model_mode分为【必要的】 0: 预训练模型，1: SFT-Chat模型 |【非必要的】2: RLHF模型，3: Reason模型，4/5: RLAIF模型
|
||||
|
||||
</details>
|
||||
|
||||
@ -370,7 +394,7 @@ torchrun --nproc_per_node N train_xxx.py
|
||||
deepspeed --master_port 29500 --num_gpus=N train_xxx.py
|
||||
```
|
||||
|
||||
可根据需要开启wandb记录训练过程
|
||||
可根据需要开启wandb记录训练过程（需可直连）
|
||||
|
||||
```bash
|
||||
# 需要登录: wandb login
|
||||
@ -382,6 +406,8 @@ python train_xxx.py --use_wandb
|
||||
通过添加`--use_wandb`参数，可以记录训练过程，训练完成后，可以在wandb网站上查看训练过程。通过修改`wandb_project`
|
||||
和`wandb_run_name`参数，可以指定项目名称和运行名称。
|
||||
|
||||
【注】：25年6月后，国内网络环境无法直连WandB，MiniMind项目默认转为使用[SwanLab](https://swanlab.cn/)作为训练可视化工具（完全兼容WandB API），即`import wandb`改为`import swanlab as wandb`即可，其他均无需改动。
|
||||
|
||||
</details>
|
||||
|
||||
# 📌 数据介绍
|
||||
@ -434,7 +460,7 @@ quality（当然也还不算high，提升数据质量无止尽）。
|
||||
|
||||
文件`pretrain_hq.jsonl` 数据格式为
|
||||
|
||||
```bash
|
||||
```json
|
||||
{"text": "如何才能摆脱拖延症？ 治愈拖延症并不容易，但以下建议可能有所帮助..."}
|
||||
```
|
||||
|
||||
@ -527,6 +553,7 @@ MiniMind训练数据集下载地址： [ModelScope](https://www.modelscope.cn/da
|
||||
├── lora_medical.jsonl (34MB)
|
||||
├── pretrain_hq.jsonl (1.6GB, ✨)
|
||||
├── r1_mix_1024.jsonl (340MB)
|
||||
├── rlaif-mini.jsonl (1MB)
|
||||
├── sft_1024.jsonl (5.6GB)
|
||||
├── sft_2048.jsonl (9GB)
|
||||
├── sft_512.jsonl (7.5GB)
|
||||
@ -542,6 +569,7 @@ MiniMind训练数据集下载地址： [ModelScope](https://www.modelscope.cn/da
|
||||
* `lora_medical.jsonl` --医疗问答数据集，推荐用于lora训练（亦可用于全参SFT，勿被名字局限）
|
||||
* `pretrain_hq.jsonl`✨ --预训练数据集，整合自jiangshu科技
|
||||
* `r1_mix_1024.jsonl` --DeepSeek-R1-1.5B蒸馏数据，每条数据字符最大长度为1024（因此训练时设置max_seq_len=1024）
|
||||
* `rlaif-mini.jsonl` --RLAIF训练数据集，从SFT数据集中随机采样1万条高质量对话，用于PPO/GRPO/SPO等强化学习算法训练
|
||||
* `sft_1024.jsonl` --整合自Qwen2.5蒸馏数据（是sft_2048的子集），每条数据字符最大长度为1024（因此训练时设置max_seq_len=1024）
|
||||
* `sft_2048.jsonl` --整合自Qwen2.5蒸馏数据，每条数据字符最大长度为2048（因此训练时设置max_seq_len=2048）
|
||||
* `sft_512.jsonl` --整合自匠数科技SFT数据，每条数据字符最大长度为512（因此训练时设置max_seq_len=512）
|
||||
@ -566,7 +594,9 @@ MiniMind训练数据集下载地址： [ModelScope](https://www.modelscope.cn/da
|
||||
|
||||
</details>
|
||||
|
||||
# 📌 Model Structure
|
||||
# 📌 Model
|
||||
|
||||
## Structure
|
||||
|
||||
MiniMind-Dense（和[Llama3.1](https://ai.meta.com/blog/meta-llama-3-1/)一样）使用了Transformer的Decoder-Only结构，跟GPT-3的区别在于：
|
||||
|
||||
@ -600,6 +630,35 @@ MiniMind的整体结构一致，只是在RoPE计算、推理函数和FFN层的
|
||||
| minimind-v1-moe | 4×26M | 6400 | 1e4 | 8 | 512 | 8 | 16 | 1+4 |
|
||||
| minimind-v1 | 108M | 6400 | 1e4 | 16 | 768 | 8 | 16 | - |
|
||||
|
||||
|
||||
## Model Configuration
|
||||
|
||||
📋关于LLM的参数配置，有一篇很有意思的论文[MobileLLM](https://arxiv.org/pdf/2402.14905)做了详细的研究和实验。
|
||||
Scaling Law在小模型中有自己独特的规律。
|
||||
引起Transformer参数成规模变化的参数几乎只取决于`d_model`和`n_layers`。
|
||||
|
||||
* `d_model`↑ + `n_layers`↓ -> 矮胖子
|
||||
* `d_model`↓ + `n_layers`↑ -> 瘦高个
|
||||
|
||||
2020年提出Scaling Law的论文认为，训练数据量、参数量以及训练迭代次数才是决定性能的关键因素，而模型架构的影响几乎可以忽视。
|
||||
然而似乎这个定律对小模型并不完全适用。
|
||||
MobileLLM提出架构的深度比宽度更重要，「深而窄」的「瘦长」模型可以学习到比「宽而浅」模型更多的抽象概念。
|
||||
例如当模型参数固定在125M或者350M时，30～42层的「狭长」模型明显比12层左右的「矮胖」模型有更优越的性能，
|
||||
在常识推理、问答、阅读理解等8个基准测试上都有类似的趋势。
|
||||
这其实是非常有趣的发现，因为以往为100M左右量级的小模型设计架构时，几乎没人尝试过叠加超过12层。
|
||||
这与MiniMind在训练过程中，模型参数量在`d_model`和`n_layers`之间进行调整实验观察到的效果是一致的。
|
||||
然而「深而窄」的「窄」也是有维度极限的，当d_model<512时，词嵌入维度坍塌的劣势非常明显，
|
||||
增加的layers并不能弥补词嵌入在固定q_head带来d_head不足的劣势。
|
||||
当d_model>1536时，layers的增加似乎比d_model的优先级更高，更能带来具有"性价比"的参数->效果增益。
|
||||
|
||||
* 因此MiniMind设定small模型dim=512，n_layers=8来获取的「极小体积<->更好效果」的平衡。
|
||||
* 设定dim=768，n_layers=16来获取效果的更大收益，更加符合小模型Scaling-Law的变化曲线。
|
||||
|
||||
作为参考，GPT3的参数设定见下表：
|
||||
|
||||
|
||||
---
|
||||
|
||||
# 📌 Experiment
|
||||
|
||||
## Ⅰ 训练开销
|
||||
@ -669,7 +728,7 @@ Zero模型权重保存为 `full_sft_512_zero.pth`（见下文MiniMind模型文
|
||||
|
||||
---
|
||||
|
||||
## Ⅱ 主要训练步骤
|
||||
## Ⅱ 主要训练（必须）
|
||||
|
||||
> 所有训练脚本均 `cd ./trainer` 目录执行
|
||||
|
||||
@ -678,7 +737,7 @@ Zero模型权重保存为 `full_sft_512_zero.pth`（见下文MiniMind模型文
|
||||
LLM首先要学习的并非直接与人交流，而是让网络参数中充满知识的墨水，“墨水” 理论上喝的越饱越好，产生大量的对世界的知识积累。
|
||||
预训练就是让Model先埋头苦学大量基本的知识，例如从Wiki百科、新闻、书籍整理大规模的高质量训练数据。
|
||||
这个过程是“无监督”的，即人类不需要在过程中做任何“有监督”的校正，而是由模型自己从大量文本中总结规律学习知识点。
|
||||
模型此阶段目的只有一个：**学会词语接龙**。例如我们输入“秦始皇”四个字，它可以接龙“是中国的第一位皇帝”。
|
||||
模型此阶段目的只有一个：**学会词语接龙**。例如输入"秦始皇"四个字，它可以接龙"是中国的第一位皇帝"。
|
||||
|
||||
```bash
|
||||
torchrun --nproc_per_node 1 train_pretrain.py # 1即为单卡训练，可根据硬件情况自行调整 (设置>=2)
|
||||
@ -689,13 +748,18 @@ python train_pretrain.py
|
||||
> 训练后的模型权重文件默认每隔`100步`保存为: `pretrain_*.pth`（*
|
||||
> 为模型具体dimension，每次保存时新文件会覆盖旧文件）
|
||||
|
||||
|
||||
| MiniMind2-Small (512dim) | MiniMind2 (768dim) |
|
||||
|---|---|
|
||||
| <img src="./images/pre_512_loss.png"> | <img src="./images/pre_768_loss.png"> |
|
||||
|
||||
### **2. 有监督微调 (Supervised Fine-Tuning)**:
|
||||
|
||||
经过预训练，LLM此时已经掌握了大量知识，然而此时它只会无脑地词语接龙，还不会与人聊天。
|
||||
SFT阶段就需要把半成品LLM施加一个自定义的聊天模板进行微调。
|
||||
例如模型遇到这样的模板【问题->回答，问题->回答】后不再无脑接龙，而是意识到这是一段完整的对话结束。
|
||||
称这个过程为指令微调，就如同让已经学富五车的「牛顿」先生适应21世纪智能手机的聊天习惯，学习屏幕左侧是对方消息，右侧是本人消息这个规律。
|
||||
在训练时，MiniMind的指令和回答长度被截断在512，是为了节省显存空间。就像我们学习时，会先从短的文章开始，当学会写作200字作文后，800字文章也可以手到擒来。
|
||||
在训练时，MiniMind的指令和回答长度被截断在512，是为了节省显存空间。就像学习写作时，会先从短的文章开始，当学会写作200字作文后，800字文章也可以手到擒来。
|
||||
在需要长度拓展时，只需要准备少量的2k/4k/8k长度对话数据进行进一步微调即可（此时最好配合RoPE-NTK的基准差值）。
|
||||
> 在推理时通过调整RoPE线性差值，实现免训练长度外推到2048及以上将会很方便。
|
||||
|
||||
@ -708,32 +772,16 @@ python train_full_sft.py
|
||||
> 训练后的模型权重文件默认每隔`100步`保存为: `full_sft_*.pth`（*
|
||||
> 为模型具体dimension，每次保存时新文件会覆盖旧文件）
|
||||
|
||||
## Ⅲ 其它训练步骤
|
||||
|
||||
| MiniMind2-Small (512dim) | MiniMind2 (768dim) |
|
||||
|---|---|
|
||||
| <img src="./images/sft_512_loss.png"> | <img src="./images/sft_768_loss.png"> |
|
||||
|
||||
## Ⅲ 其它训练阶段（可选）
|
||||
|
||||
> 所有训练脚本均 `cd ./trainer` 目录执行
|
||||
|
||||
### **3. 人类反馈强化学习(Reinforcement Learning from Human Feedback, RLHF)**
|
||||
|
||||
在前面的训练步骤中，模型已经具备了基本的对话能力，但是这样的能力完全基于单词接龙，缺少正反样例的激励。
|
||||
模型此时尚未知什么回答是好的，什么是差的。我们希望它能够更符合人的偏好，降低让人类不满意答案的产生概率。
|
||||
这个过程就像是让模型参加新的培训，从优秀员工的作为例子，消极员工作为反例，学习如何更好地回复。
|
||||
此处使用的是RLHF系列之-直接偏好优化(Direct Preference Optimization, DPO)。
|
||||
与PPO(Proximal Policy Optimization)这种需要奖励模型、价值模型的RL算法不同；
|
||||
DPO通过推导PPO奖励模型的显式解，把在线奖励模型换成离线数据，Ref模型输出可以提前保存。
|
||||
DPO性能几乎不变，只用跑 actor_model 和 ref_model 两个模型，大大节省显存开销和增加训练稳定性。
|
||||
|
||||
> 注：RLHF训练步骤**并非必须**，此步骤难以提升模型“智力”而通常仅用于提升模型的“礼貌”，有利（符合偏好、减少有害内容）也有弊（样本收集昂贵、反馈偏差、多样性损失）。
|
||||
|
||||
```bash
|
||||
torchrun --nproc_per_node 1 train_dpo.py
|
||||
# or
|
||||
python train_dpo.py
|
||||
```
|
||||
|
||||
> 训练后的模型权重文件默认每隔`100步`保存为: `rlhf_*.pth`（*
|
||||
> 为模型具体dimension，每次保存时新文件会覆盖旧文件）
|
||||
|
||||
### **4. 知识蒸馏(Knowledge Distillation, KD)**
|
||||
### **3. 知识蒸馏 (Knowledge Distillation, KD)**
|
||||
|
||||
在前面的所有训练步骤中，模型已经完全具备了基本能力，通常可以学成出师了。
|
||||
而知识蒸馏可以进一步优化模型的性能和效率，所谓知识蒸馏，即学生模型面向教师模型学习。
|
||||
@ -768,7 +816,7 @@ torchrun --nproc_per_node 1 train_distillation.py
|
||||
python train_distillation.py
|
||||
```
|
||||
|
||||
### **5. LoRA (Low-Rank Adaptation)**
|
||||
### **4. LoRA (Low-Rank Adaptation)**
|
||||
|
||||
LoRA是一种高效的参数高效微调（Parameter-Efficient Fine-Tuning, PEFT）方法，旨在通过低秩分解的方式对预训练模型进行微调。
|
||||
相比于全参数微调（Full Fine-Tuning），LoRA 只需要更新少量的参数。
|
||||
@ -787,7 +835,7 @@ python train_lora.py
|
||||
|
||||
非常多的人困惑，如何使模型学会自己私有领域的知识？如何准备数据集？如何迁移通用领域模型打造垂域模型？
|
||||
这里举几个例子，对于通用模型，医学领域知识欠缺，可以尝试在原有模型基础上加入领域知识，以获得更好的性能。
|
||||
同时，我们通常不希望学会领域知识的同时损失原有基础模型的其它能力，此时LoRA可以很好的改善这个问题。
|
||||
同时，通常不希望学会领域知识的同时损失原有基础模型的其它能力，此时LoRA可以很好的改善这个问题。
|
||||
只需要准备如下格式的对话数据集放置到`./dataset/lora_xxx.jsonl`，启动 `python train_lora.py`
|
||||
训练即可得到`./out/lora/lora_xxx.pth`新模型权重。
|
||||
|
||||
@ -806,7 +854,7 @@ python train_lora.py
|
||||
```
|
||||
|
||||
此时【基础模型+LoRA模型】即可获得医疗场景模型增强的能力，相当于为基础模型增加了LoRA外挂，这个过程并不损失基础模型的本身能力。
|
||||
我们可以通过`eval_model.py`进行模型评估测试。
|
||||
可以通过`eval_model.py`进行模型评估测试。
|
||||
|
||||
```bash
|
||||
# 注意：model_mode即选择基础模型的类型，这和train_lora是基于哪个模型训练的相关，确保统一即可。
|
||||
@ -827,7 +875,7 @@ python eval_model.py --lora_name 'lora_medical' --model_mode 2
|
||||
|
||||
PS：只要有所需要的数据集，也可以full_sft全参微调（需要进行通用知识的混合配比，否则过拟合领域数据会让模型变傻，损失通用性）
|
||||
|
||||
### **6. 训练推理模型 (Reasoning Model)**
|
||||
### **5. 训练推理模型 (Reasoning Model)**
|
||||
|
||||
DeepSeek-R1实在太火了，几乎重新指明了未来LLM的新范式。
|
||||
论文指出`>3B`的模型经历多次反复的冷启动和RL奖励训练才能获得肉眼可见的推理能力提升。
|
||||
@ -835,7 +883,9 @@ DeepSeek-R1实在太火了，几乎重新指明了未来LLM的新范式。
|
||||
但由于缺乏技术含量，蒸馏派被RL派瞧不起（hhhh）。
|
||||
本人迅速已经在Qwen系列1.5B小模型上进行了尝试，很快复现了Zero过程的数学推理能力。
|
||||
然而一个遗憾的共识是：参数太小的模型直接通过冷启动SFT+GRPO几乎不可能获得任何推理效果。
|
||||
<del>
|
||||
MiniMind2第一时间只能坚定不移的选择做蒸馏派，日后基于0.1B模型的RL如果同样取得小小进展会更新此部分的训练方案。
|
||||
</del>
|
||||
|
||||
做蒸馏需要准备的依然是和SFT阶段同样格式的数据即可，数据集来源已如上文介绍。数据格式例如：
|
||||
|
||||
@ -908,44 +958,343 @@ python train_distill_reason.py
|
||||
</answer>
|
||||
```
|
||||
|
||||
---
|
||||
## IV 强化学习后训练
|
||||
|
||||
## Ⅳ 模型参数设定
|
||||
LLM里的强化学习方法可分两类：
|
||||
|
||||
📋关于LLM的参数配置，有一篇很有意思的论文[MobileLLM](https://arxiv.org/pdf/2402.14905)做了详细的研究和实验。
|
||||
Scaling Law在小模型中有自己独特的规律。
|
||||
引起Transformer参数成规模变化的参数几乎只取决于`d_model`和`n_layers`。
|
||||
1. **基于人类反馈的强化学习 (Reinforcement Learning from Human Feedback, RLHF)**
|
||||
|
||||
* `d_model`↑ + `n_layers`↓ -> 矮胖子
|
||||
* `d_model`↓ + `n_layers`↑ -> 瘦高个
|
||||
- 通过**人类**对模型输出的偏好进行评价来训练模型，使其生成更符合人类价值观和偏好的内容。
|
||||
|
||||
2020年提出Scaling Law的论文认为，训练数据量、参数量以及训练迭代次数才是决定性能的关键因素，而模型架构的影响几乎可以忽视。
|
||||
然而似乎这个定律对小模型并不完全适用。
|
||||
MobileLLM提出架构的深度比宽度更重要，「深而窄」的「瘦长」模型可以学习到比「宽而浅」模型更多的抽象概念。
|
||||
例如当模型参数固定在125M或者350M时，30～42层的「狭长」模型明显比12层左右的「矮胖」模型有更优越的性能，
|
||||
在常识推理、问答、阅读理解等8个基准测试上都有类似的趋势。
|
||||
这其实是非常有趣的发现，因为以往为100M左右量级的小模型设计架构时，几乎没人尝试过叠加超过12层。
|
||||
这与MiniMind在训练过程中，模型参数量在`d_model`和`n_layers`之间进行调整实验观察到的效果是一致的。
|
||||
然而「深而窄」的「窄」也是有维度极限的，当d_model<512时，词嵌入维度坍塌的劣势非常明显，
|
||||
增加的layers并不能弥补词嵌入在固定q_head带来d_head不足的劣势。
|
||||
当d_model>1536时，layers的增加似乎比d_model的优先级更高，更能带来具有“性价比”的参数->效果增益。
|
||||
2. **基于AI反馈的强化学习 (Reinforcement Learning from AI Feedback, RLAIF)**
|
||||
|
||||
* 因此MiniMind设定small模型dim=512，n_layers=8来获取的「极小体积<->更好效果」的平衡。
|
||||
* 设定dim=768，n_layers=16来获取效果的更大收益，更加符合小模型Scaling-Law的变化曲线。
|
||||
- 使用**AI模型**（通常是预训练的语言奖励模型）来提供反馈，而不直接依赖人类的人工标注。
|
||||
- 这里的“AI”也可以是某些规则奖励，例如数学答案/代码解释器...
|
||||
|
||||
作为参考，GPT3的参数设定见下表：
|
||||
|
||||
| 类型 | 裁判 | 优点 | 缺点 |
|
||||
|-------|----|-----------|------------|
|
||||
| RLHF | 人类 | 更贴近真实人类偏好 | 成本高、效率低 |
|
||||
| RLAIF | 模型 | 自动化、可扩展性强 | 可能偏离人类真实偏好 |
|
||||
|
||||
二者本质上是一样的，都是通过**强化学习的方式**，利用某种形式的"**反馈**"来优化模型的行为。
|
||||
|
||||
除了**反馈**的来源不同，其他并无任何区别。
|
||||
|
||||
|
||||
### 👀 PO算法的统一视角
|
||||
|
||||
在介绍实现具体算法之前，我先以个人理解的极简视角，阐述所有Policy Optimization (PO)算法的统一共性。
|
||||
|
||||
所有RL算法的本质都只是在优化一个期望：
|
||||
|
||||
$$\mathcal{J}_{PO} = \mathbb{E}_{q \sim P(Q), o \sim \pi(O|q)} \left[ \underbrace{f(r_t)}_{\text{策略项}} \cdot \underbrace{g(A_t)}_{\text{优势项}} - \underbrace{h(\text{KL}_t)}_{\text{正则项}} \right]$$
|
||||
|
||||
训练时，只需**最小化负目标函数**，即: $\mathcal{L_{PO}}=\mathcal{J_{PO}}$
|
||||
|
||||
这个框架只包含三个核心组件：
|
||||
* **策略项** $f(r_t)$: 如何使用概率比 $r_t$? 即告诉模型新旧策略偏差有多大，是否探索到了更好的token
|
||||
* **优势项** $g(A_t)$: 如何计算优势 $A_t$, 这很重要！大模型算对定积分也不足为奇，小模型回答对加减法优势通常都是正的
|
||||
* **正则项** $h(\text{KL}_t)$: 如何约束变化幅度 $\text{KL}_t$, 既防止跑偏又防止管的太死
|
||||
|
||||
<details>
|
||||
<summary>（展开）符号说明</summary>
|
||||
|
||||
| 符号 | 含义 | 说明 | 值域 |
|
||||
|------|------|------|------|
|
||||
| $q$ | 问题/提示词 | 从数据集 $P(Q)$ 中采样 | - |
|
||||
| $o$ | 模型输出序列 | 由策略 $\pi$ 生成 | - |
|
||||
| $r_t$ | 概率比 | $r_t = \frac{\pi_\theta(o_t\|q, o_{<t})}{\pi_{ref}(o_t\|q, o_{<t})}$ | $(0, +\infty)$ |
|
||||
| $A_t$ | 优势函数 | 衡量某个动作相比基线有多好 | $(-\infty, +\infty)$ |
|
||||
| $\text{KL}_t$ | KL散度 | 防止策略偏离参考模型太远 | $[0, +\infty)$ |
|
||||
|
||||
</details>
|
||||
<br/>
|
||||
|
||||
不同的**xxPO算法**本质上只是对这三个组件的不同设计的实例化！
|
||||
|
||||
---
|
||||
|
||||
## Ⅴ 训练结果
|
||||
### **6. 基于人类反馈的强化学习 (Reinforcement Learning from Human Feedback, RLHF)**
|
||||
|
||||
MiniMind2 模型训练损失走势（由于数据集在训练后又更新清洗多次，因此Loss仅供参考）
|
||||
在前面的训练步骤中，模型已经具备了基本的对话能力，但是这样的能力完全基于单词接龙，缺少正反样例的激励。
|
||||
模型此时尚未知什么回答是好的，什么是差的。希望它能够更符合人的偏好，降低让人类不满意答案的产生概率。
|
||||
这个过程就像是让模型参加新的培训，从优秀员工的作为例子，消极员工作为反例，学习如何更好地回复。
|
||||
|
||||
| models | pretrain (length-512) | sft (length-512) |
|
||||
|-----------------|----------------------------------------------------|----------------------------------------------------|
|
||||
| MiniMind2-Small | <img src="./images/pre_512_loss.png" width="100%"> | <img src="./images/sft_512_loss.png" width="100%"> |
|
||||
| MiniMind2 | <img src="./images/pre_768_loss.png" width="100%"> | <img src="./images/sft_768_loss.png" width="100%"> |
|
||||
#### 6.1 Direct Preference Optimization
|
||||
直接偏好优化（DPO）算法，损失为：
|
||||
|
||||
$$\mathcal{L}_{DPO} = -\mathbb{E}\left[\log \sigma\left(\beta \left[\log \frac{\pi_\theta(y_w|x)}{\pi_{ref}(y_w|x)} - \log \frac{\pi_\theta(y_l|x)}{\pi_{ref}(y_l|x)}\right]\right)\right]$$
|
||||
|
||||
其中：
|
||||
- **策略项**: $f(r_t) = \log r_w - \log r_l$ (对比chosen vs rejected的概率比)
|
||||
- **优势项**: $g(A_t)$ = / (通过偏好对比，无需显式计算优势)
|
||||
- **正则项**: $h(\text{KL}_t)$ = 隐含在 $\beta$ 中 (控制偏离参考模型程度)
|
||||
|
||||
特别地，
|
||||
- DPO从PPO带KL约束的目标推导出对偏好对的解析训练目标，直接最大化"chosen优于rejected"的对数几率；无需同步训练Reward/Value模型。DPO只需跑`actor`与`ref`两个模型，显存占用低、收敛稳定、实现简单。
|
||||
- 训练范式：off‑policy，使用静态偏好数据集，可反复多轮epoch；Ref模型固定（预先缓存输出）。
|
||||
- DPO的局限在于不做在线探索，更多用于"偏好/安全"的人类价值对齐；对"能不能做对题"的智力能力提升有限（当然这也取决于数据集，大规模收集正反样本并人类评估很困难）。
|
||||
|
||||
```bash
|
||||
torchrun --nproc_per_node 1 train_dpo.py
|
||||
# or
|
||||
python train_dpo.py
|
||||
```
|
||||
|
||||
> 训练后的模型权重文件默认每隔`100步`保存为: `rlhf_*.pth`（*为模型具体dimension，每次保存时新文件会覆盖旧文件）
|
||||
|
||||
### **7. 基于AI反馈的强化学习 (Reinforcement Learning from AI Feedback, RLAIF)**
|
||||
|
||||
相比RLHF依赖人类标注chosen/rejected偏好对，RLAIF则完全由AI来充当"裁判"。
|
||||
所谓AI"裁判"可以是model-base的奖励大模型(Reward Model)，也可以是R1一样设置规则函数进行校验，也可以是例如工具调用的环境反馈。
|
||||
例如：数学题答案是否正确、工具调用执行代码能否通过测试用例、推理过程是否符合格式...都可以自动化判断。
|
||||
RLAIF的最大优势在于**可扩展性**和**On-Policy**的特点——不需要昂贵的人工标注，可以生成海量的训练样本，让模型在在线大量试错中快速进化。
|
||||
|
||||
MiniMind 着手实现**2+N**种基本+前沿的RLAIF方法：
|
||||
* **PPO**、**GRPO** 被大规模验证的经典RL算法；
|
||||
* N种前沿RL算法（不定期以Exp性质更新）。
|
||||
|
||||
#### 1️⃣ 数据集准备 (需要)
|
||||
|
||||
为了快速验证RLAIF的效果，这里从SFT数据集中随机采样了1万条高质量对话，构建约1MB大小的`rlaif-mini.jsonl`([Huggingface](https://huggingface.co/datasets/jingyaogong/minimind_dataset/blob/main/rlaif-mini.jsonl))
|
||||
|
||||
数据格式与SFT一致，但assistant并不需要内容，因为训练过程中完全由 $\Pi$ 策略模型实时采样生成。因此形如：
|
||||
|
||||
```json
|
||||
{
|
||||
"conversations": [
|
||||
{"role": "user", "content": "请解释一下什么是光合作用？"},
|
||||
{"role": "assistant", "content": "无"}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
RLAIF的训练过程中，模型会基于user的问题生成1或多个候选回答，然后由奖励函数/模型对回答打分，
|
||||
分数高的回答会被鼓励（增加 $\Pi$ 策略概率），分数低的回答会被抑制（降低 $\Pi$ 策略概率）。这个"打分->调整"的循环就是强化学习的核心。
|
||||
|
||||
#### 2️⃣ 奖励模型准备 (需要)
|
||||
|
||||
已知RLAIF训练需要“奖励模型 (Reward Model)”对生成的回答进行打分。
|
||||
|
||||
此处选取小型且高质量的InternLM2-1.8B-Reward
|
||||
([ModelScope](https://modelscope.cn/models/Shanghai_AI_Laboratory/internlm2-1_8b-reward) | [HuggingFace](https://huggingface.co/internlm/internlm2-1_8b-reward))
|
||||
作为基础奖励模型。
|
||||
|
||||
下载奖励模型后需要放置在minimind项目的**同级目录**下，推荐结构如下：
|
||||
|
||||
```
|
||||
project/
|
||||
├── minimind/ # MiniMind项目
|
||||
│ ├── model/
|
||||
│ └── ...
|
||||
└── internlm2-1_8b-reward/ # 奖励模型（与minimind同级）
|
||||
├── config.json
|
||||
├── model.safetensors
|
||||
└── ...
|
||||
```
|
||||
|
||||
<details>
|
||||
<summary><b>奖励机制选择与MiniMind限制说明（点击展开）</b></summary>
|
||||
|
||||
**1. 奖励机制的多样性**
|
||||
|
||||
RLAIF中的"奖励信号"来源可以非常灵活：
|
||||
|
||||
- **Model-based奖励**：可使用专门的Reward Model（如InternLM2-Reward），也可使用通用LLM+提示词进行打分（如Qwen3-as-a-Judge）。奖励模型规模和架构均可自由选择。
|
||||
|
||||
- **Rule-based奖励**：可以基于规则函数构造奖励信号，例如：
|
||||
- 数学题答案正确性验证（Ground Truth对比）
|
||||
- SQL执行成功率与结果准确性
|
||||
- 代码解释器运行结果（pass@k）
|
||||
- 工具调用返回状态（API成功/失败）
|
||||
- 格式合规性检查（JSON/XML解析）
|
||||
- 推理链完整性评估（CoT步骤数）
|
||||
|
||||
- **Environment-based奖励**：在Agent场景中，环境反馈本身即为天然奖励（如游戏得分、Research完整度、任务完成度）。
|
||||
|
||||
任何能够量化"回答质量"的机制都可作为RL的奖励来源。DeepSeek R1就是典型案例：使用规则函数验证数学答案正确性作为奖励，无需额外的Reward Model。
|
||||
|
||||
**2. MiniMind限制：奖励稀疏问题**
|
||||
|
||||
RLAIF训练既可以针对推理模型也可以针对非推理模型，区别仅在于格式。
|
||||
|
||||
然而对于MiniMind这种0.1B参数量极小能力弱的模型，在通用任务（如R1风格的数学数据集）上会遇到严重的奖励稀疏(Reward Sparsity)问题：
|
||||
|
||||
- **现象**：模型生成的候选回答几乎全部错误，导致所有奖励分数 $r(x,y) \approx 0$
|
||||
- **后果**：优势函数 $A(x,y) = r(x,y) - b(x) \approx 0$，策略梯度信号消失，无法有效更新参数 $\theta$
|
||||
|
||||
如同让小学生做高考数学题，无论尝试多少次都得零分，无法通过分数差异学习改进策略。因此这是RL算法的根本原理限制的。
|
||||
|
||||
为缓解此问题，MiniMind的实现选择了**model-based的连续性奖励信号**：
|
||||
|
||||
- Reward Model输出连续分数（如-2.5到+3.0），而非二元的0/1
|
||||
- 即使回答质量都差，也仍能区分"更更差"(-3.0)和"更差"(-2.8)的细微差异。所以这种**稠密且连续**的奖励信号能够为优势函数 $A(x,y)$ 提供非零梯度，使得策略网络得以渐进式优化
|
||||
- 也可以混合多种奖励源: $r_{\text{total}} = \alpha \cdot r_{\text{model}} + \beta \cdot r_{\text{rule}}$ (例如既可以检测think标签格式reward，又可以综合回答本身质量的reward分数)
|
||||
- minimind实践中避免直接使用rule-based二元奖励 + 超纲难度数据（如MATH500），易导致奖励全零；
|
||||
- 监控训练时观察奖励分数的方差 $\text{Var}(r)$，若持续接近0则需调整数据或奖励机制
|
||||
|
||||
**对于生产级大模型的Agentic RL场景**：
|
||||
|
||||
在真实Agent系统（代码生成、工具调用、检索-规划-执行的多轮链路）中，奖励是“延迟整轮结算”的不同范式：
|
||||
|
||||
- LLM需要逐token生成工具调用指令（tool_call），经历解析（tool_parse）、工具执行（tool_exec），再把结果拼接回上下文继续下一步；循环往复直到完成。
|
||||
- 一次完整的任务链路包含多次调用+思考，直到终止条件满足时计算一次总reward（如任务是否完成、测试是否通过、目标是否命中）。
|
||||
|
||||
因此，Agentic RL更接近稀疏/延迟奖励设定：梯度回传在“整轮结束后”才发生，和非Agentic RL任务在对话单轮上“即时评分即时更新”有很大不同。
|
||||
这也解释了Agent任务上更偏向环境反馈（environment-based reward），而非凭Reward Model进行静态打分。
|
||||
|
||||
- **环境交互反馈**：最终以执行结果为准（代码是否跑通、API是否返回成功、子目标是否完成）；
|
||||
- **Model-based奖励局限**：对长链路、可执行语义的全貌捕捉有限，且大概率和真实环境反馈不一致（reward hacking）。
|
||||
|
||||
|
||||
</details>
|
||||
|
||||
---
|
||||
|
||||
#### 7.1 [Proximal Policy Optimization](https://arxiv.org/abs/1707.06347)
|
||||
|
||||
PPO 是2017年OpenAI提出的非常经典强化学习算法，也是LLM RL通用的基线方法，甚至不需要加之一。
|
||||
|
||||
**PPO损失**：
|
||||
$$\mathcal{L}_{PPO} = -\mathbb{E}\left[\min(r_t \cdot A_t, \text{clip}(r_t, 1-\varepsilon, 1+\varepsilon) \cdot A_t)\right] + \beta \cdot \mathbb{E}[\text{KL}]$$
|
||||
|
||||
其中：
|
||||
- **策略项**: $f(r_t) = \min(r_t, \text{clip}(r_t, 1-\varepsilon, 1+\varepsilon))$ (裁剪概率比防止更新过激)
|
||||
- **优势项**: $g(A_t) = R - V(s)$ (通过Critic网络估计价值函数)
|
||||
- **正则项**: $h(\text{KL}_t) = \beta \cdot \mathbb{E}[\text{KL}]$ (全局KL散度约束)
|
||||
|
||||
对比DPO而言，
|
||||
- DPO (Off-Policy)：训练数据是静态的偏好数据集（chosen vs rejected），可以反复使用同一批数据训练多个epoch，就像传统监督学习一样。数据效率高，训练成本低。它直接优化偏好对的对数似然，无需Reward Model。
|
||||
- PPO (On-Policy)：必须用当前策略实时采样生成新数据，旧策略采集的数据不能用（会有distribution shift问题）。虽然通过importance sampling和clip机制允许轻微的分布偏移，但本质上要求数据来自相对新鲜的策略。数据效率低，但适合探索式学习。
|
||||
|
||||
简单来说：
|
||||
|
||||
- 前者教模型按离线预定的「好/坏标准」学习，尽管它并非是当前模型所能输出的（例如参考世界冠/亚军录像练习打球）；
|
||||
- 后者实时地教模型把事情做对做好，在线采样自最新模型policy（教练手把手教打，为每个动作实时打分）。
|
||||
|
||||
MiniMind的PPO实现包含了Actor模型(生成回答)和Critic模型(评估回答价值)，以及完整的GAE(Generalized Advantage Estimation)优势函数计算。
|
||||
|
||||
**训练方式**：
|
||||
|
||||
```bash
|
||||
torchrun --nproc_per_node N train_ppo.py
|
||||
# or
|
||||
python train_ppo.py
|
||||
```
|
||||
|
||||
> 训练后的模型权重文件默认每隔`100步`保存为: `ppo_actor_*.pth`（*为模型具体dimension）
|
||||
|
||||
|
||||
| MiniMind2-Small (512dim) | MiniMind2 (768dim) |
|
||||
|---|---|
|
||||
| <img src="./images/train_ppo_512.png"> | <img src="./images/train_ppo_768.png"> |
|
||||
|
||||
从训练曲线可以看出，PPO存在**reward提升缓慢**的问题。私以为这主要源于**PPO双网络联合优化**方法：Critic需要逐步收敛以准确估计价值函数，而Actor的策略更新依赖Critic提供的优势估计，两者相互依赖形成复杂的优化过程。训练初期Critic估计不准会影响Actor梯度方向，导致整体收敛缓慢。此外，PPO需要同时维护两个网络，显存占用约为单网络方法的1.5-2倍。
|
||||
|
||||
#### 7.2 [Group Relative Policy Optimization](https://arxiv.org/pdf/2402.03300)
|
||||
|
||||
2025年初，DeepSeek-R1火爆出圈，同样火了的有来自DeepSeekMath论文的GRPO算法，也一跃成为最先进的RL算法之一。
|
||||
然而AI半年=人间半个世纪，时至今日GRPO已经演变为各大XXPO大战(后面演变的DAPO、GSPO、CISPO等)的基线算法。
|
||||
具体来说，一句话总结它的核心创新是"分组相对价值估计"。
|
||||
|
||||
**GRPO损失**：
|
||||
|
||||
$$\mathcal{L}_{GRPO} = -\mathbb{E}\left[r_t \cdot A_t - \beta \cdot \text{KL}_t\right]$$
|
||||
|
||||
其中：
|
||||
- **策略项**: $f(r_t) = r_t$ (直接使用概率比，无clip裁剪)
|
||||
- **优势项**: $g(A_t) = \frac{R - \mu_{group}}{\sigma_{group}}$ (组内归一化，消除Critic网络)
|
||||
- **正则项**: $h(\text{KL}_t) = \beta \cdot \text{KL}_t$ (token级KL散度约束)
|
||||
|
||||
对于同一个问题，模型生成N个不同的回答(例如N=4)，然后计算这N个回答的奖励分数。
|
||||
接着把这N个回答的平均奖励作为baseline，高于baseline的回答被鼓励，低于baseline的回答被抑制。
|
||||
用这种方式巧妙地避免了训练额外的critic网络。
|
||||
|
||||
只要是RL都必须面对的正反样本这个原理性限制，GRPO也不会例外，其更显著的问题是：退化组(Degenerate Groups)。
|
||||
假设某个问题略难，导致N个回答的奖励分数几乎一样（大部分情况是一样烂而不是一样好），那么这一组的学习信号就无限接近0。
|
||||
在MiniMind这种超小模型上，这个问题尤为明显，求解数学问题99.99%的情况下整组回答质量都很差，那么将无法学习。
|
||||
因此必须为模型指定合理的domain，即必须限制在能力边界内。
|
||||
|
||||
|
||||
**训练方式**：
|
||||
|
||||
```bash
|
||||
torchrun --nproc_per_node N train_grpo.py
|
||||
# or
|
||||
python train_grpo.py
|
||||
```
|
||||
|
||||
> 训练后的模型权重文件默认每隔`100步`保存为: `grpo_*.pth`
|
||||
|
||||
|
||||
| MiniMind2-Small (512dim) | MiniMind2 (768dim) |
|
||||
|---|---|
|
||||
| <img src="./images/train_grpo_512.png"> | <img src="./images/train_grpo_768.png"> |
|
||||
|
||||
从训练曲线可以看出，GRPO的**reward呈现更加稳定的上升趋势**，达到4左右，说明GRPO本身能更好地利用RLAIF信号。Policy Loss整体下降平稳，相比PPO的双网络优化，GRPO单网络架构训练更稳定且收敛上限更高。
|
||||
|
||||
#### 7.3 ⏳⌛️🔥 更多RL拓展 (Exp)
|
||||
|
||||
##### 7.3.1 [Single-stream Policy Optimization](https://arxiv.org/abs/2509.13232)
|
||||
|
||||
SPO是2025年9月腾讯提出的RL算法，针对GRPO的退化组问题进行改进。
|
||||
论文认为，GRPO等算法"一个样本要依赖一组采样"显得别扭而不优雅：太容易或太难的题目，整组几乎学不到东西，学习效率先天受限。
|
||||
SPO的动机就是回到RL的本质—**1个输入，1个输出，就是1个训练样本**，回到policy gradient的基本公式去思考：不用group mean也能得到稳定的baseline，也就是把价值估计 V 铺开在时序上，训练前先做粗略的价值预估，训练中一边采样一边更新对 V 的估计，从而为每个样本提供一个跨 batch 持久化、可自适应的基线参照。这种"单流"设计不再依赖同组样本，天然避免了退化组。
|
||||
|
||||
**SPO损失**:
|
||||
|
||||
$$\mathcal{L}_{SPO} = -\mathbb{E}\left[\log \pi_\theta(a_t|s) \cdot A_t - \beta \cdot \text{KL}_t\right]$$
|
||||
|
||||
其中：
|
||||
- **策略项**: $f(r_t) = \log \pi_\theta(a_t|s)$ (直接使用log概率，不计算ratio)
|
||||
- **优势项**: $g(A_t) = R - B_t^{adaptive}$ (自适应baseline，Beta分布动态跟踪)
|
||||
- **正则项**: $h(\text{KL}_t) = \beta \cdot \text{KL}_t$ (token级KL + 动态 $\rho$ 调整)
|
||||
|
||||
落到实现层面：SPO采用无分组设计，用持久化的KL自适应value tracker替代组内baseline，优势函数在整个batch上全局归一化。这样每个样本独立处理，无需等待同组其他样本，且能为每个样本提供稳定的学习信号。
|
||||
论文在Qwen3-8B的5个困难数学数据集上，SPO平均比GRPO高出3.4个百分点，其中BRUMO 25数据集+7.3pp、AIME 25数据集+4.4pp。
|
||||
|
||||
|
||||
> 注：SPO是实验性前沿算法，MiniMind的实现用于探索学习。由于模型参数量极小，无法完全复现论文的8B模型效果。
|
||||
|
||||
**训练方式**：
|
||||
|
||||
```bash
|
||||
torchrun --nproc_per_node N train_spo.py
|
||||
# or
|
||||
python train_spo.py
|
||||
```
|
||||
|
||||
> 训练后的模型权重文件默认每隔`100步`保存为: `spo_*.pth`
|
||||
|
||||
|
||||
<div align="center">
|
||||
<img src="./images/train_spo_768.png">
|
||||
<p><i>MiniMind2 (768dim) 训练曲线</i></p>
|
||||
</div>
|
||||
|
||||
从训练曲线来看，SPO的reward波动与PPO表现接近，弱于GRPO。实际推理测试发现模型输出质量不高，存在逻辑混乱和格式错误问题。
|
||||
|
||||
**实验性说明**：当前SPO手搓实现可能在value_tracker配置、reward归一化策略上还存在问题。尚需排查算法本身在小模型上的适应性/或是实现上存在差异。
|
||||
|
||||
|
||||
### RL算法小结
|
||||
|
||||
我们收束回“**统一框架**”, 重新整理所有不同PO算法只是对三个核心组件的不同实例化的表格：
|
||||
|
||||
| 算法 | 策略项 $f(r_t)$ | 优势项 $g(A_t)$ | 正则项 $h(\text{KL}_t)$ | 优化模型 |
|
||||
|------|----------------|----------------|----------------------|----------|
|
||||
| **DPO** | $\log r_w - \log r_l$ | 隐式（偏好对比） | 隐含在 $\beta$ 中 | 2 |
|
||||
| **PPO** | $\min(r, \text{clip}(r))$ | $R - V(s)$ | $\beta \cdot \mathbb{E}[\text{KL}]$ | 4 |
|
||||
| **GRPO** | $r$ | $\frac{R - \mu}{\sigma}$ | $\beta \cdot \text{KL}_t$ | 2 |
|
||||
| **SPO** | $\log \pi_\theta$ | $R - B_t^{adaptive}$ | $\beta \cdot \text{KL}_t$ | 2 |
|
||||
|
||||
**RL是优美且自洽的**
|
||||
|
||||
> 以上纯属个人视角理解，如有偏差请随时指正
|
||||
|
||||
---
|
||||
|
||||
## V 训练结果
|
||||
|
||||
### 训练完成-模型合集
|
||||
|
||||
@ -955,22 +1304,15 @@ MiniMind2 模型训练损失走势（由于数据集在训练后又更新清洗
|
||||
|
||||
MiniMind2模型权重 ([ModelScope](https://www.modelscope.cn/models/gongjy/MiniMind2-PyTorch) | [HuggingFace](https://huggingface.co/jingyaogong/MiniMind2-Pytorch))
|
||||
|
||||
MiniMind-V1模型权重 ([百度网盘](https://pan.baidu.com/s/1KUfSzEkSXYbCCBj0Pw-9fA?pwd=6666))
|
||||
|
||||
<details style="color:rgb(128,128,128)">
|
||||
<summary>Torch文件命名对照</summary>
|
||||
|
||||
| Model Name | params | pretrain_model | sft_model | rl_model | reason_model | lora_model |
|
||||
|-----------------|--------|------------------------|------------------------|--------------------|------------------|--------------------|
|
||||
| MiniMind2-small | 26M | `pretrain_512.pth` | `full_sft_512.pth` | `rlhf_512.pth` | `reason_512.pth` | `lora_xxx_512.pth` |
|
||||
| MiniMind2-MoE | 145M | `pretrain_640_moe.pth` | `full_sft_640_moe.pth` | `rlhf_640_moe.pth` | - | - |
|
||||
| MiniMind2 | 104M | `pretrain_768.pth` | `full_sft_768.pth` | `rlhf_768.pth` | `reason_768.pth` | `lora_xxx_768.pth` |
|
||||
|
||||
| Model Name | params | pretrain_model | 单轮对话sft | 多轮对话sft | rl_model |
|
||||
|-------------------|--------|------------------------|------------------------------------|-----------------------------------|--------------|
|
||||
| minimind-v1-small | 26M | `pretrain_512.pth` | `single_chat/full_sft_512.pth` | `multi_chat/full_sft_512.pth` | `rl_512.pth` |
|
||||
| minimind-v1-moe | 4×26M | `pretrain_512_moe.pth` | `single_chat/full_sft_512_moe.pth` | `multi_chat/full_sft_512_moe.pth` | - |
|
||||
| minimind-v1 | 108M | `pretrain_768.pth` | `single_chat/full_sft_768.pth` | `multi_chat/full_sft_768.pth` | `rl_768.pth` |
|
||||
| Model Name | params | pretrain_model | sft_model | rlhf_model (DPO) | reason_model | rlaif_model (PPO/GRPO/SPO) | lora_model |
|
||||
|-----------------|--------|------------------------|------------------------|--------------------|------------------|----------------------------------------------|--------------------|
|
||||
| MiniMind2-small | 26M | `pretrain_512.pth` | `full_sft_512.pth` | `rlhf_512.pth` | `reason_512.pth` | `xxpo_512.pth` | `lora_xxx_512.pth` |
|
||||
| MiniMind2-MoE | 145M | `pretrain_640_moe.pth` | `full_sft_640_moe.pth` | `rlhf_640_moe.pth` | - | - | - |
|
||||
| MiniMind2 | 104M | `pretrain_768.pth` | `full_sft_768.pth` | `rlhf_768.pth` | `reason_768.pth` | `xxpo_768.pth` | `lora_xxx_768.pth` |
|
||||
|
||||
</details>
|
||||
|
||||
@ -983,37 +1325,31 @@ MiniMind系列 ([ModelScope](https://www.modelscope.cn/collections/MiniMind-b72f
|
||||
|
||||
# 📌 Eval
|
||||
|
||||
## Ⅰ RLHF对比SFT篇
|
||||
## Ⅰ RLHF对比
|
||||
|
||||
测试基于`full_sft`和`rlhf`的`MiniMind2系列`模型对比，测试随机种子均固定
|
||||
|
||||
<details style="color:rgb(128,128,128)">
|
||||
<summary>问答实测</summary>
|
||||
|
||||
[A0] MiniMind2：full_sft_640
|
||||
|
||||
[A1] MiniMind2：rlhf_640
|
||||
[A] MiniMind2：full_sft_640<br/>
|
||||
[B] MiniMind2：rlhf_640
|
||||
|
||||
```text
|
||||
[Q]: 你叫什么名字？
|
||||
[A0]: 我是一个大型语言模型，可以回答各种问题，提供信息和帮助。我可以回答各种问题，提供信息和帮助，帮助你解决各种问题。
|
||||
[A1]: 我是一种基于文本的AI模型，我并没有个人经历或情感，所以无法感受和体验情感。我的设计目的是为了帮助用户解决问题、提供信息和进行对话。如果您有任何问题或需要帮助，请随时告诉我！
|
||||
[A]: 我是一个大型语言模型，可以回答各种问题，提供信息和帮助。我可以回答各种问题，提供信息和帮助，帮助你解决各种问题。
|
||||
[B]: 我是一种基于文本的AI模型，我并没有个人经历或情感，所以无法感受和体验情感。我的设计目的是为了帮助用户解决问题、提供信息和进行对话。如果您有任何问题或需要帮助，请随时告诉我！
|
||||
|
||||
[Q]: 你更擅长哪一个学科？
|
||||
[A0]: 作为一个AI模型，我并没有真正的“擅长”这个学科。我被设计成能够提供信息、解答问题、进行对话、执行任务和与用户进行对话。不过，我可以提供信息、解答问题、进行对话、提供建议、处理问题、进行对话、提供建议、学习新知识、编写代码、回答问题、提供建议等多种学科的功能。如果你有具体的学科需求或主题，欢迎提供详细信息。
|
||||
[A1]: 作为AI，我并没有个人能力或兴趣，但我可以提供一些关于数学、科学、计算机科学、语言学、心理学等学科的知识。如果你对数学感兴趣，我可以提供一些基础知识，例如代数、几何、概率、统计学等。如果你需要更深入的学习，我也可以提供相关的在线课程、书籍、论文等资源。请告诉我你具体需要什么帮助，我会尽力提供帮助。
|
||||
[A]: 作为一个AI模型，我并没有真正的“擅长”这个学科。我被设计成能够提供信息、解答问题、进行对话、执行任务和与用户进行对话。不过，我可以提供信息、解答问题、进行对话、提供建议、处理问题、进行对话、提供建议、学习新知识、编写代码、回答问题、提供建议等多种学科的功能。如果你有具体的学科需求或主题，欢迎提供详细信息。
|
||||
[B]: 作为AI，我并没有个人能力或兴趣，但我可以提供一些关于数学、科学、计算机科学、语言学、心理学等学科的知识。如果你对数学感兴趣，我可以提供一些基础知识，例如代数、几何、概率、统计学等。如果你需要更深入的学习，我也可以提供相关的在线课程、书籍、论文等资源。请告诉我你具体需要什么帮助，我会尽力提供帮助。
|
||||
|
||||
[Q]: 鲁迅的《狂人日记》是如何批判封建礼教的？
|
||||
[A0]: 鲁迅的《狂人日记》是一部讲述了鲁迅对人类价值和道德责任的倡导。在这部作品中，鲁迅通过细腻的笔触，探讨了人性、道德、正义、牺牲与责任等深刻的主题。以下是鲁迅在《狂人日记》中如何批判封建礼教的几个方面：1. **社会现实与道德观念**：鲁迅通过讲述鲁迅的故事，探讨了社会的不平等、权力的流动、人性的演变等社会问题。他认为，社会不平等不仅仅是不公正的，更是每个人在社会中的不公正。他呼吁社会各界共同努力，通过教育、宣传、法律等手段，为构建和谐社会做出贡献。2. **道德与责任**：鲁迅通过讲述鲁迅的故事，展现了道德和责任的重要性。他提倡通过道德教育来培养人们的道德品质，同时也强调了道德和责任的重要性。鲁迅通过讲述鲁迅的故事，揭示了社会的不平等和人性的复杂性，体现了对人性的深刻理解。3. **道德与责任**：鲁迅通过讲述鲁迅的故事，强调了道德责任的重要性。他提倡通过自己的行为来体现对社会的尊重和对人性的深刻理解。鲁迅通过讲述鲁迅的故事，强调了道德和责任的重要性。4. **道德与道德**：鲁迅通过讲述鲁迅的故事，揭示了道德和责任的真正意义。鲁迅通过讲述鲁迅的故事，揭示了道德和责任的本质，强调了道德和责任的重要性。5. **个人与社会**：鲁迅通过讲述鲁迅的故事，揭示了个人与社会之间的复杂关系。他呼吁社会各界共同努力，通过教育和宣传，为构建和谐社会做出贡献。鲁迅的《狂人日记》通过其丰富的故事和深刻的道德寓言，传达了关于人性、道德、正义、道德以及个人与社会的复杂关系。这部作品不仅在文学价值上取得了巨大成就，更在思想层面上启发了后来的文学创作，对后世产生了深远的影响。
|
||||
[A1]: 鲁迅的《狂人日记》是一部描绘战争和自由精神的作品，它批判了封建社会的种种问题和冲突，反映了社会的不平等和人性本质。在这部作品中，鲁迅通过对封建社会的诸多现象进行批判，如“弗雷德里克·约翰逊的“鹰迪伦”、“纳扎特的”、“贝克汉姆”、“乔治·威尔逊”等，对封建社会的种种问题进行了批判，强调了封建社会的不平等、人性本质以及人性本质。鲁迅通过他的“鹰迪伦”、“纳扎特的”、“贝克汉姆”、“乔治·威尔逊”等，对封建社会的种种问题进行了批判。在小说中，鲁迅通过对封建社会的诸多现象进行批判，揭示了封建社会的腐败、人性本质以及人性本质的复杂性。这种批判体现了对封建社会的尊重和对人性的复杂性，同时也反映了现代社会中普遍存在的不平等和人性本质。鲁迅的《狂人日记》不仅是一部关于封建社会的小说，也是一部具有深刻社会意义的作品，对后世的审判和批判起到了关键作用。
|
||||
[A]: 鲁迅的《狂人日记》是一部讲述了鲁迅对人类价值和道德责任的倡导。在这部作品中，鲁迅通过细腻的笔触，探讨了人性、道德、正义、牺牲与责任等深刻的主题。以下是鲁迅在《狂人日记》中如何批判封建礼教的几个方面：1. **社会现实与道德观念**：鲁迅通过讲述鲁迅的故事，探讨了社会的不平等、权力的流动、人性的演变等社会问题。他认为，社会不平等不仅仅是不公正的，更是每个人在社会中的不公正。他呼吁社会各界共同努力，通过教育、宣传、法律等手段，为构建和谐社会做出贡献。2. **道德与责任**：鲁迅通过讲述鲁迅的故事，展现了道德和责任的重要性。他提倡通过道德教育来培养人们的道德品质，同时也强调了道德和责任的重要性。鲁迅通过讲述鲁迅的故事，揭示了社会的不平等和人性的复杂性，体现了对人性的深刻理解。3. **道德与责任**：鲁迅通过讲述鲁迅的故事，强调了道德责任的重要性。他提倡通过自己的行为来体现对社会的尊重和对人性的深刻理解。鲁迅通过讲述鲁迅的故事，强调了道德和责任的重要性。4. **道德与道德**：鲁迅通过讲述鲁迅的故事，揭示了道德和责任的真正意义。鲁迅通过讲述鲁迅的故事，揭示了道德和责任的本质，强调了道德和责任的重要性。5. **个人与社会**：鲁迅通过讲述鲁迅的故事，揭示了个人与社会之间的复杂关系。他呼吁社会各界共同努力，通过教育和宣传，为构建和谐社会做出贡献。鲁迅的《狂人日记》通过其丰富的故事和深刻的道德寓言，传达了关于人性、道德、正义、道德以及个人与社会的复杂关系。这部作品不仅在文学价值上取得了巨大成就，更在思想层面上启发了后来的文学创作，对后世产生了深远的影响。
|
||||
[B]: 鲁迅的《狂人日记》是一部描绘战争和自由精神的作品，它批判了封建社会的种种问题和冲突，反映了社会的不平等和人性本质。在这部作品中，鲁迅通过对封建社会的诸多现象进行批判，如“弗雷德里克·约翰逊的“鹰迪伦”、“纳扎特的”、“贝克汉姆”、“乔治·威尔逊”等，对封建社会的种种问题进行了批判，强调了封建社会的不平等、人性本质以及人性本质。鲁迅通过他的“鹰迪伦”、“纳扎特的”、“贝克汉姆”、“乔治·威尔逊”等，对封建社会的种种问题进行了批判。在小说中，鲁迅通过对封建社会的诸多现象进行批判，揭示了封建社会的腐败、人性本质以及人性本质的复杂性。这种批判体现了对封建社会的尊重和对人性的复杂性，同时也反映了现代社会中普遍存在的不平等和人性本质。鲁迅的《狂人日记》不仅是一部关于封建社会的小说，也是一部具有深刻社会意义的作品，对后世的审判和批判起到了关键作用。
|
||||
|
||||
[Q]: 我咳嗽已经持续了两周，需要去医院检查吗？
|
||||
[A0]: 咳嗽确实需要一定的检查，但如果您的咳嗽持续时间过长，或者有其他不适症状（如持续咳嗽、呼吸困难等），建议尽快就医。医生会根据您的症状、体检结果和可能的影像学检查（如X光、MRI或CT扫描）来诊断您的症状。在等待就医期间，您可以尝试以下方法来帮助缓解症状：1. **休息**：确保您有足够的睡眠。2. **保持水分**：脱水可能会加重咳嗽。3. **热敷**：使用热敷袋或热毛巾敷在疼痛的部位，有助于缓解喉咙痛。4. **避免刺激**：避免使用可能刺激喉咙的化学物质，如某些化学物质。5. **保持室内湿润**：确保您的卧室有良好的通风，可以使用加湿器帮助保持空气湿度。6. **保持室内湿度**：使用加湿器可以帮助缓解喉咙干燥和咳嗽。如果症状持续不缓解或有其他不适症状，请立即就医。希望您早日康复！
|
||||
[A1]: 如果您的咳嗽持续了两周，通常建议您尽快就医，因为这可能会影响您的呼吸和心跳。您可以通过以下步骤来检查咳嗽：1. **休息**：用纸巾或肘部轻轻按压咳嗽的位置。2. **使用药物**：在可能的情况下，使用药物（如抗过敏药、抗过敏药等）来缓解咳嗽。请在医生的指导下使用药物。3. **热敷**：在咳嗽期间，可以使用热敷来缓解疼痛和肿胀。4. **药物治疗**：如果咳嗽持续了两周，可以使用一些消炎药或抗过敏药来缓解症状。请务必在医生的指导下使用。5. **保持健康的生活方式**：保持充足的睡眠、均衡的饮食、适量的运动和避免吸烟和限酒，这些都有助于提高咳嗽的几率。请记住，这些建议不能替代专业医疗意见。如果您的症状持续或加重，强烈建议您尽快就医。
|
||||
[A]: 咳嗽确实需要一定的检查，但如果您的咳嗽持续时间过长，或者有其他不适症状（如持续咳嗽、呼吸困难等），建议尽快就医。医生会根据您的症状、体检结果和可能的影像学检查（如X光、MRI或CT扫描）来诊断您的症状。在等待就医期间，您可以尝试以下方法来帮助缓解症状：1. **休息**：确保您有足够的睡眠。2. **保持水分**：脱水可能会加重咳嗽。3. **热敷**：使用热敷袋或热毛巾敷在疼痛的部位，有助于缓解喉咙痛。4. **避免刺激**：避免使用可能刺激喉咙的化学物质，如某些化学物质。5. **保持室内湿润**：确保您的卧室有良好的通风，可以使用加湿器帮助保持空气湿度。6. **保持室内湿度**：使用加湿器可以帮助缓解喉咙干燥和咳嗽。如果症状持续不缓解或有其他不适症状，请立即就医。希望您早日康复！
|
||||
[B]: 如果您的咳嗽持续了两周，通常建议您尽快就医，因为这可能会影响您的呼吸和心跳。您可以通过以下步骤来检查咳嗽：1. **休息**：用纸巾或肘部轻轻按压咳嗽的位置。2. **使用药物**：在可能的情况下，使用药物（如抗过敏药、抗过敏药等）来缓解咳嗽。请在医生的指导下使用药物。3. **热敷**：在咳嗽期间，可以使用热敷来缓解疼痛和肿胀。4. **药物治疗**：如果咳嗽持续了两周，可以使用一些消炎药或抗过敏药来缓解症状。请务必在医生的指导下使用。5. **保持健康的生活方式**：保持充足的睡眠、均衡的饮食、适量的运动和避免吸烟和限酒，这些都有助于提高咳嗽的几率。请记住，这些建议不能替代专业医疗意见。如果您的症状持续或加重，强烈建议您尽快就医。
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
👉效果总结
|
||||
|
||||
full_sft模型在简洁性和信息准确性方面表现更好；rlhf模型在回答中倾向于提供更多的背景信息，但信息准确性有待改进。
|
||||
@ -1022,9 +1358,34 @@ full_sft模型在简洁性和信息准确性方面表现更好；rlhf模型在
|
||||
DPO和在线PPO的区别在于reject和chosen都是离线准备的，和minimind模型本身的输出必然存在很大的分布差异。
|
||||
通俗地说DPO算法使模型观看乒乓球世界冠军的打法「录像」进行RL，而不是像PPO一样请reward模型做「教练」纠正自己的打法进行RL。
|
||||
|
||||
## Ⅱ 主观样例测评
|
||||
## Ⅱ RLAIF对比
|
||||
|
||||
🏃以下测试于2025-02-09完成，此日期后发布的新模型，无特殊需要时将不加入测试。
|
||||
基于`MiniMind2 (768dim)`在相同数据与训练步数下的主观观感对比，测试随机种子均固定2026，供参考：
|
||||
|
||||
[A] MiniMind-Reason-0.1B (蒸馏)<br/>
|
||||
[B] MiniMind-Reason-0.1B (PPO)<br/>
|
||||
[C] MiniMind-Reason-0.1B (GRPO)
|
||||
|
||||
```text
|
||||
[Q]: 请介绍一下自己。
|
||||
[A]: <think>我是由中国的个人开发者开发的智能助手小型AI推理模型-R1。我被设计用来理解和生成人类语言，能够理解复杂的自然语言问题，提供有用的信息和解决方案。</think><answer>我是由中国的个人开发者开发的智能助手小型AI推理模型-R1。</answer>
|
||||
[B]: <think>嗯，用户让我介绍一下自己。首先，我需要明确用户的需求是什么。他们可能是在做某个项目，想要了解自己的一些基本情况。不过，用户没有提供具体的信息，所以我得先假设一些可能的情况，比如他是一个学生，正在学习编程，或者是某个项目的开发者。接下来，我要考虑用户的身份。可能是一个学生、教师或者普通用户，他们需要了解自己的特点和能力。也有可能他们正在寻找灵感，或者在写论文时需要帮助。所以，我需要保持回答的简洁和实用，同时提供一些参考资料，比如教学视频、课程记录等。然后，我要考虑用户的身份。可能是一个学生，需要一些学习资源或帮助的对象。也有可能是一个学生，需要在课后作业或学习计划中进行调整。因此，我需要先询问用户的具体需求，这样才能提供更精准的帮助。接下来，我要确保语言简洁明了，适合不同场景。比如，如果用户是学生，可以介绍自己的学习特点；如果是教师，可以说明自己的角色和贡献。同时，提供一些实用的建议，比如如何高效学习，如何进行项目，以及如何与团队成员协作。最后，我要确保回复清晰，有条理，能够满足用户的需求。这样，用户就能轻松地了解自己，并且有实际的帮助。</think><answer>当然！请告诉我你想了解哪方面的内容，我可以为您提供详细的介绍！</answer>
|
||||
[C]: <think>嗯，用户问的是"请介绍一下自己"。首先，我需要理解用户的具体需求是什么。他们可能想知道自己是谁，或者有什么特定的背景。接下来，我应该考虑用户可能没有明确表达的需求。他们可能想知道如何描述自己的外观，或者如何与他人交流。然后，我应该考虑如何组织内容，从头到尾都要有条理。可能需要分步骤说明，比如从视觉、听觉到整体感受，再到情感共鸣。这样用户能更清晰地理解自己，也能更好地回应他们的需求。还要注意语言的自然和流畅，避免过于学术化的术语，让用户感觉顺畅。同时，保持回答的结构清晰，每句话都有明确的目的，这样用户能有明确的理解和操作。最后，检查是否有遗漏的重要信息，确保回答准确且全面。这样，用户不仅能得到答案，还能得到满足他们的需求。</think><answer>您好！我是由中国的个人开发者开发的智能助手小型AI推理模型-R1。关于您的描述，似乎涉及多个领域，包括知识图谱构建、复杂系统的推理与分析。无论是理论模型、实际应用还是伦理讨论，您都在寻求如何将这些概念有机地结合起来，为人类社会带来福祉的解决方案。在"我"的身份下，您提到的"个人体验"和"情感共鸣"，都是您个人独特魅力的重要组成部分。通过精准的描述和深度的分析，您可以更深入地理解您的独特之处。如果您对特定领域有特别的兴趣或需要建议，随时可以分享更多细节，我们将为您提供更加精准和个性化的服务。如果您有特定的主题或问题，欢迎随时提问！让我们一起探索这个神奇的AI世界，共同创造更美好的未来。</answer>
|
||||
|
||||
|
||||
[Q]: 你更擅长哪一个学科？
|
||||
[A]: <think>我是由中国的个人开发者开发的智能助手小型AI推理模型-R1。</think><answer>我是由中国的个人开发者开发的智能助手小型AI推理模型-R1。</answer>
|
||||
[B]: <think>嗯，用户问的是"你更擅长哪一个学科？"这个问题看起来有点复杂，但我可以一步步分析。首先，我需要理解用户的需求。他可能是在准备考试，或者需要了解某个学科的特点，或者是想在学习过程中应用某个学科的知识。接下来，我得考虑用户可能的背景。他可能是一个学生，或者是在准备考试，需要应用数学知识来解决具体的问题。也有可能他是一个学生，正在学习相关的课程，或者是对某个学科感兴趣。然后，我需要分析用户可能没有明确表达出来的需求。他可能在思考自己在学习中的困难，或者在寻找一个合适的教材或学习资源。考虑到这些，我应该提供几个不同的角度来帮助他。首先，分析"我更擅长哪一个学科"可以拆解为以下几个方面：兴趣、学习能力、学习风格、学习方法、学习资源、学习环境等等。每个方面都有不同的优势和特点，比如兴趣是最好的资源，学习能力是最好的方式，学习方式是根据个人情况选择。接下来，我需要考虑用户的身份。可能是学生，也可能是家长，或者对学习有一定兴趣的人。不同的身份可能有不同的偏好，比如家长可能更注重孩子的兴趣，而学生可能更关注学习方法和资源。然后，我需要考虑用户的身份是否有特殊需求。比如，他可能需要更专业的学习材料，或者希望找到最适合他的学习方法。因此，在分析时，我应该提供一个全面的回答，涵盖各个方面，并解释每个方面的重要性。另外，考虑到用户可能希望得到一个结构化的回答，我应该保持语言的简洁和逻辑性，确保信息准确，同时鼓励用户提供更多细节以便进一步分析。最后，我需要确保我的回答既有知识性，又符合用户的要求，帮助用户更好地理解和应用这个学科。因此，在回答中不仅要给出一个全面的分析，还要提到每个部分的作用和如何利用。</think><answer>你擅长哪一个学科？这是一个明确的问题，可以从以下几个方面进行分析：1. **兴趣与学习能力**：你喜欢什么样的学习方式，比如自主学习、实验、项目等，这些都是很好的学习资源。2. **学习风格**：你喜欢什么类型的学习方法，比如笔记、复习、阅读、讨论等，这些都能影响学习效果。3. **学习资源**：你是否在寻找学习资源，比如书籍、在线课程、教育软件等，这些都能提供学习帮助。4. **学习环境**：你是否在学习环境中感到舒适，比如舒适的学习空间、安静的学习空间等，这些都会影响学习效果。5. **学习资源**：你是否有学习资源的丰富，如在线课程、学习平台、书籍等，这些都能提供丰富的学习资源。6. **学习方法**：你是否有具体的学习方法，比如如何理解、如何组织学习内容、如何应用知识等，这些都是提高学习效率的途径。综合考虑，这是一个非常全面的分析，涵盖了多个方面，帮助你更好地理解你的学科。</answer>
|
||||
[C]: <think>嗯，用户问的是"你更擅长哪一个学科？"这个问题看起来有点复杂，但我可以一步步分析。首先，我需要理解用户的需求。他可能是在准备考试，或者需要了解某个学科的特点，或者是想在生活中找到自己的兴趣点。接下来，我要考虑用户可能的背景。他可能在工作或学习中遇到了类似的问题，想要了解自己擅长哪个领域，或者是否需要一些资源或指导来提升自己。然后，我要分析用户可能的深层需求。也许他想在职场上提升自己，或者寻找方向。他可能需要更具体的例子，帮助他更好地理解不同学科的特点和应用场景。还要考虑用户可能的背景。他可能不仅仅是在学习，而是希望了解在特定领域中如何应用这些学科。这需要我提供全面的分析，帮助他找到适合自己的方法。最后，我应该保持开放的态度，鼓励用户继续探索和学习。这可能包括推荐一些资源，或者讨论如何将这些知识应用到实际生活中。</think><answer>你更擅长哪一个学科？这个问题本质上取决于你的个人兴趣和目标。以下是一些关键点，帮助你全面理解并应用这一点：1. **兴趣与专长**：了解自己擅长哪个领域，有助于做出更明智的选择。2. **全面发展**：考虑学科的多样性和深度，以全面发展自己的生活和职业。3. **适应与调整**：选择与自己技能、兴趣和目标相匹配的学科，有助于应对挑战和实现目标。4. **职业规划**：探索适合自己的职业路径，提升效率和竞争力。5. **社会贡献**：评估自己在社会中的角色，为社会做出贡献。为了更好地理解并应用这一点，你可以：- **了解行业动态**：关注特定领域的专业发展。- **实践与反思**：通过项目或实习积累经验。- **职业发展**：评估自己的长期职业规划。- **持续学习**：不断更新知识，适应职场变化。通过这些方向，你可以更全面地认识自己，找到适合自己的发展路径。</answer>
|
||||
|
||||
|
||||
[Q]: 鲁迅的《狂人日记》是如何批判封建礼教的？
|
||||
[A]: <think>嗯，鲁迅的《狂人日记》是如何批判封建礼教的呢？这个问题听起来挺有意思的，尤其是他在《狂人日记》中的表达方式。首先，我需要确认一下《狂人日记》具体是谁写的，是否有特定的人物或者事件。《狂人日记》是鲁迅在1924年创作的，主要讨论了社会问题，尤其是封建礼教的问题。作者曹雪芹是个非常著名的作家，他的作品中有许多关于封建礼教和阶级斗争的观点。这类作品在文学、哲学和政治上都有很深远的影响，尤其是在社会结构和人性的方面。接下来，我需要了解鲁迅的写作风格。他的写作风格比较直接，有严肃的哲学观点，强调人性的悲剧和道德的虚无。他的语言风格比较富有诗意，富有象征意义，善于运用象征手法，同时又不失真实和现实主义。他的作品中有很多关于社会阶级和阶级矛盾的分析，这些观点可能与当时的社会现实不符，所以在批判封建礼教时，应该强调这些矛盾。然后，我想到《狂人日记》中对封建礼教的讨论。鲁迅的作品中，封建礼教被视为不人道的，但通过文字和行动，表达了他对封建礼教的不满和反抗。这不仅反映了当时的社会现实，也揭示了封建礼教的黑暗面，体现了对封建社会的不满和对人性的深刻反思。接下来，我需要思考鲁迅对封建礼教的理解和批判。他在《狂人日记》中，对封建礼教的质疑和批判，体现了他对传统封建礼教的否定和批判。这不仅仅是对封建礼教的质疑，更是对当时社会的反思，以及对人性的关注。最后，我需要总结一下鲁迅对封建礼教的批判和批判。他通过《狂人日记》中的文字和行为，批判了封建礼教的虚假和不合理，体现了对社会现实的深刻反思。同时，他也关注了封建礼教的黑暗面，反映了对封建社会的不满和对人性的否定。总的来说，鲁迅的《狂人日记》通过文字和行动，批判了封建礼教的虚假和不合理，体现了当时的社会现实和人性的复杂。他的作品对后来的封建礼教有了启示，对当时的社会现象有了更深的反思。</think><answer>鲁迅的《狂人日记》是一部反映封建礼教的文学作品，主要通过文字和行动反映了当时社会的现实。鲁迅的文字风格独特，富有象征意义，但通过文字和行动，批判了封建礼教的虚假和不合理。他的作品中，封建礼教被视为不人道的，但通过文字和行动，反映了封建社会的不公和人性的深刻反思。鲁迅的《狂人日记》通过文字和行动，批判了封建礼教的虚假和不合理，体现了当时社会的不满和对人性的反思。他的作品不仅批判了封建礼教的虚假和不合理，还深刻揭示了封建社会的黑暗面。</answer>
|
||||
[B]: <think>嗯，用户问的是鲁迅的《狂人日记》如何批判封建礼教。首先，我需要理解鲁迅的生平和作品。《狂人日记》是中国现代文学的经典，作者是毛泽东，作品涉及社会、家庭、伦理等多个层面。我需要找出鲁迅的写作风格和主题，以及它在当时的社会背景。接下来，我应该回顾《狂人日记》的主要内容。毛泽东在《青春无悔》中提出了"狂人日记"这一概念，强调青春活泼、勇敢、正义，这在当时社会中仍然具有重要意义。他通过描绘主人公狂人来表达对社会不公的批判，但同时也指出他的创作对社会的负面影响，比如社会动荡和个人主义的失败。然后，我应该分析鲁迅的创作理念和主题。《狂人日记》强调了社会的黑暗面，主张"以人为本"、"反对封建礼教"等。他可能也在探讨社会不平等的问题，比如"人性弱点"、"权力滥用"等问题。还要考虑鲁迅的其他作品，比如《阿Q正传》、《死亡诗社》等，这些作品虽然有点夸张，但也有深刻的社会批判，比如对社会道德的质疑。同时，作者可能还涉及其他作品，如《庆余年》等，但这些可能更侧重于文学创作，而不是社会批判。另外，用户可能想知道鲁迅在文学创作中的具体经历，比如他是否有作品被广泛批评，或者他对社会的批判是否具有批判性。所以，回答时需要指出鲁迅的创作经历，以及他对社会现实的反思。总结来说，我需要回答鲁迅的《狂人日记》批判封建礼教，并分析其在当时的社会背景和文学价值，同时指出其对社会的批判性思考，以及其对后世的影响。</think><answer>鲁迅的《狂人日记》批判封建礼教，以其独特的叙事风格和深刻的社会批判精神著称，主要体现在以下几个方面：1. **社会批判与个人成长**：小说通过描绘主人公狂人的活泼、勇敢、正义，揭示社会的不公和弱势。2. **个人主义与社会批判**：小说反映了封建礼教对个人自由和权利的侵犯，强调个人尊严和正义的价值。3. **对社会的批判性思考**：通过对社会不公的刻画，揭示社会问题的根源，并对当时社会的批判态度进行批判。鲁迅的作品在当时的社会背景和文学价值使其成为当时文学的重要代表。他的作品不仅在当时广受欢迎，也对后世产生了深远影响，被广泛引用和传颂。</answer>
|
||||
[C]: <think>嗯，用户问的是鲁迅的《狂人日记》如何批判封建礼教。首先，我需要理解鲁迅的背景和核心思想。他可能是一位社会学家、思想家，或者是对传统文化有深厚兴趣的读者。狂人日记是他在《狂人日记》中所写的，可能对传统文化有一定影响。接下来，我要考虑用户的需求是什么。用户可能想知道这本书在现代社会中的反思意义，或者是否有相关研究支持这一点。也有可能他们对书中的情节有疑问，想了解如何更深刻地理解社会现象。然后，我需要思考如何从不同的角度解读鲁迅的文本。例如，是否有直接的历史叙事，或者是否存在更深层次的哲学思考？这需要我先了解鲁迅的不同观点，然后分析其批判性。还要考虑用户可能的深层需求，比如是否需要了解他作为文学批判家的背景，或者是否有对传统价值观的反思。因此，回答时要全面，涵盖核心问题，并引用相关资料来支撑观点。最后，总结鲁迅的观点，强调其对现代社会的批判性思考，以及对传统文化的尊重。这样既全面又有深度，帮助用户全面理解鲁迅的作品。</think><answer>鲁迅的《狂人日记》通过其独特的叙事手法和深刻的社会批判精神，对封建礼教进行了深刻的批判。以下是其批判性分析：1. **社会批判与历史真相**：小说通过对历史人物的刻画，揭示封建礼教的弊端，以及其对传统社会结构的冲击。这种批判不仅限于对现实的深刻批判，也反映了现代社会对传统价值观的重新审视。2. **个人命运与社会批判**：小说中，主人公通过"狂人日记"揭示了封建礼教对个人命运的严重剥夺。这种叙事方式既批判了封建礼教，也对后世产生了深远影响。3. **对后世的影响**：从文学创作到思想解读，鲁迅的作品深刻影响了后世的思想体系。例如，《狂人日记》通过人物命运的反映，对现代社会的价值观念提出了新的思考。4. **对现代文明的反思**：这一批判性思考不仅反映了当代社会的复杂性，也为现代人提供了反思社会现实、维护人权的视角。5. **个人视角与文化视角**：作者通过这一批判性思考，展现了对传统价值观的重新解读，同时也对当代社会的多元化视角产生了启发。综上所述，鲁迅的《狂人日记》不仅是对封建礼教的批判，更是对传统文化的一种深刻反思。它通过对封建礼教的批判，提出了对现代社会结构、权力关系的深刻反思，具有重要的社会价值。</answer>
|
||||
```
|
||||
|
||||
## Ⅲ 其他模型对比
|
||||
|
||||
[A] [MiniMind2 (0.1B)](https://www.modelscope.cn/models/gongjy/MiniMind2-PyTorch)<br/>
|
||||
[B] [MiniMind2-MoE (0.15B)](https://www.modelscope.cn/models/gongjy/MiniMind2-PyTorch)<br/>
|
||||
@ -1193,7 +1554,6 @@ DPO和在线PPO的区别在于reject和chosen都是离线准备的，和minimind
|
||||
| 7 | E | 22 | 23 | 15 | 14 | 74 |
|
||||
| 8 | G | 10 | 12 | 10 | 10 | 42 |
|
||||
|
||||
|
||||
### 👉主观效果总结
|
||||
|
||||
个人主观评价与DeepSeek-R1基本相符，其中：
|
||||
@ -1208,10 +1568,27 @@ DPO和在线PPO的区别在于reject和chosen都是离线准备的，和minimind
|
||||
|
||||
---
|
||||
|
||||
## Ⅲ Objective Benchmark
|
||||
## Ⅳ RoPE长度外推
|
||||
|
||||
下面就到喜闻乐见的benchmark刷榜测试环节，就不找乐子和qwen、glm级别的中文模型做对比了。
|
||||
这里选取了一些<1B的微型模型进行横评比较，
|
||||
MiniMind支持通过YaRN算法进行RoPE位置编码的长度外推，使模型能够处理超出训练长度的文本序列。
|
||||
在使用`eval_model.py`进行推理时，只需添加`--inference_rope_scaling True`参数即可启用RoPE外推：
|
||||
|
||||
```bash
|
||||
python eval_model.py --inference_rope_scaling True
|
||||
```
|
||||
|
||||
下图展示了在不同文本「西游记」白话文小说长度下，使用RoPE scaling前后的困惑度(PPL)对比。可以看出，启用RoPE scaling后，模型在长文本上的表现显著提升：
|
||||
|
||||
<div align="center">
|
||||
<img src="./images/rope_ppl.png">
|
||||
</div>
|
||||
|
||||
---
|
||||
|
||||
## Ⅴ Objective Benchmark
|
||||
|
||||
下面就到喜闻乐见的benchmark测试环节，就不找乐子和Qwen、GLM级别的模型做对比了。
|
||||
这里选取了一些微型模型进行横评比较，
|
||||
测试集选择C-Eval、CMMLU、A-CLUE、TMMLU+这几个纯中文语言榜单。
|
||||
|
||||
|
||||
@ -1232,7 +1609,7 @@ lm_eval --model hf --model_args pretrained=<填写模型路径>,device=cuda,dtyp
|
||||
PS: 在这种全是选择题的测评集中，为了避免回复格式的难以固定的特点，
|
||||
所以常用做法是直接把`A`,`B`,`C`,`D`四个字母对应token的预测概率取出来，将其中概率最大的字母与标准答案计算正确率。
|
||||
选择题1/4乱选的正确率是25%，然而这个量级的所有模型都集中在25附近，甚至很多时候不如瞎选，是不是像极了高中完形填空的滑铁卢正确率...
|
||||
MiniMind模型本身预训练数据集小的可怜，也没有针对性的对测试集做刷榜微调，因此结果图一乐即可：
|
||||
MiniMind模型本身预训练数据集小的可怜，也没有针对性的对测试集做刷榜微调，因此结果纯娱乐：
|
||||
|
||||
| models | from | params↓ | ceval↑ | cmmlu↑ | aclue↑ | tmmlu+↑ |
|
||||
|-------------------------------------------------------------------------------|---------------|---------|--------|---------|--------|---------|
|
||||
@ -1247,7 +1624,7 @@ MiniMind模型本身预训练数据集小的可怜，也没有针对性的对测
|
||||
|
||||
|
||||
|
||||
# 📌 其它 (Others)
|
||||
# 📌 Others
|
||||
|
||||
## 模型转换
|
||||
|
||||
@ -1297,23 +1674,21 @@ MiniMind模型本身预训练数据集小的可怜，也没有针对性的对测
|
||||
}'
|
||||
```
|
||||
|
||||
## VLLM模型推理（服务）
|
||||
## <img src="https://avatars.githubusercontent.com/u/136984999" height="28" style="vertical-align: middle;"/> [vllm](https://github.com/vllm-project/vllm)
|
||||
|
||||
[vLLM](https://github.com/vllm-project/vllm)是极其流行的高效推理框架，支持大模型快速部署，优化显存利用与吞吐量。
|
||||
vLLM是极其流行的高效推理框架，支持大模型快速部署，优化显存利用与吞吐量。
|
||||
|
||||
```bash
|
||||
vllm serve ./MiniMind2/ --model-impl transformers --served-model-name "minimind"
|
||||
vllm serve ./MiniMind2 --model-impl transformers --served-model-name "minimind" --port 8998
|
||||
```
|
||||
|
||||
服务将以openai api协议启动，端口默认为8000。
|
||||
## <img src="https://user-images.githubusercontent.com/1991296/230134379-7181e485-c521-4d23-a0d6-f7b3b61ba524.png" height="28" style="vertical-align: middle;"/> [llama.cpp](https://github.com/ggerganov/llama.cpp)
|
||||
|
||||
更多用法请参考官方说明～
|
||||
|
||||
## llama.cpp
|
||||
[llama.cpp](https://github.com/ggerganov/llama.cpp)是一个C++库，
|
||||
llama.cpp是一个C++库，
|
||||
可以在命令行下直接使用，支持多线程推理，支持GPU加速。
|
||||
|
||||
参考官方仓库安装后，在`convert_hf_to_gguf.py` ～760行插入
|
||||
|
||||
```text
|
||||
# 添加MiniMind2 tokenizer支持
|
||||
if res is None:
|
||||
@ -1321,27 +1696,29 @@ if res is None:
|
||||
```
|
||||
|
||||
转换自定义训练的minimind模型 -> gguf
|
||||
|
||||
```bash
|
||||
python convert_hf_to_gguf.py ../minimind/MiniMind2/
|
||||
```
|
||||
|
||||
量化模型
|
||||
|
||||
```bash
|
||||
./build/bin/llama-quantize ../minimind/MiniMind2/MiniMind2-109M-F16.gguf ../minimind/MiniMind2/Q4-MiniMind2.gguf Q4_K_M
|
||||
```
|
||||
|
||||
命令行推理
|
||||
|
||||
```bash
|
||||
./build/bin/llama-cli -m ../minimind/MiniMind2/MiniMind2-109M-F16.gguf --chat-template chatml
|
||||
```
|
||||
|
||||
更多用法请参考官方说明～
|
||||
## <img src="https://ollama.com/public/cloud.png" height="28" style="vertical-align: middle;"/> [ollama](https://ollama.ai)
|
||||
|
||||
## ollama
|
||||
|
||||
[ollama](https://ollama.ai/)是本地运行大模型的工具，支持多种开源LLM，简单易用。
|
||||
ollama是本地运行大模型的工具，支持多种开源LLM，简单易用。
|
||||
|
||||
通过ollama加载自定义的gguf模型，新建minimind.modelfile：
|
||||
|
||||
```text
|
||||
FROM ./MiniMind2-109M-F16.gguf
|
||||
TEMPLATE """{{ if .System }}<|im_start|>system
|
||||
@ -1353,17 +1730,19 @@ TEMPLATE """{{ if .System }}<|im_start|>system
|
||||
```
|
||||
|
||||
加载模型并命名为`minimind2`
|
||||
|
||||
```bash
|
||||
ollama create -f minimind.modelfile minimind2
|
||||
```
|
||||
|
||||
启动推理
|
||||
|
||||
```text
|
||||
ollama run minimind2
|
||||
> 你好，我是MiniMind2，一个基于xxxxxxxx
|
||||
```
|
||||
|
||||
更多用法请参考官方说明～
|
||||
> 以上三方框架的更多用法请参考对应官方文档😊
|
||||
|
||||
# 📌 Acknowledge
|
||||
|
||||
|
||||
1731
README_en.md
1731
README_en.md
File diff suppressed because it is too large
Load Diff
@ -57,8 +57,8 @@ class SFTDataset(Dataset):
|
||||
self.tokenizer = tokenizer
|
||||
self.max_length = max_length
|
||||
self.samples = self.load_data(jsonl_path)
|
||||
self.bos_id = tokenizer('<|im_start|>assistant', add_special_tokens=False).input_ids
|
||||
self.eos_id = tokenizer('<|im_end|>', add_special_tokens=False).input_ids
|
||||
self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant', add_special_tokens=False).input_ids
|
||||
self.eos_id = tokenizer(f'{tokenizer.eos_token}', add_special_tokens=False).input_ids
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
@ -71,16 +71,14 @@ class SFTDataset(Dataset):
|
||||
samples.append(data)
|
||||
return samples
|
||||
|
||||
def _create_chat_prompt(self, conversations):
|
||||
"""构建符合ChatML格式的对话"""
|
||||
messages = []
|
||||
for i, turn in enumerate(conversations):
|
||||
role = 'user' if i % 2 == 0 else 'assistant'
|
||||
messages.append({"role": role, "content": turn['content']})
|
||||
def _create_chat_prompt(self, cs):
|
||||
messages = cs.copy()
|
||||
tools = cs[0]["functions"] if (cs and cs[0]["role"] == "system" and cs[0].get("functions")) else None
|
||||
return self.tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=False
|
||||
add_generation_prompt=False,
|
||||
tools=tools
|
||||
)
|
||||
|
||||
def _generate_loss_mask(self, input_ids):
|
||||
@ -115,7 +113,14 @@ class SFTDataset(Dataset):
|
||||
X = torch.tensor(input_ids[:-1], dtype=torch.long)
|
||||
Y = torch.tensor(input_ids[1:], dtype=torch.long)
|
||||
loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long) # 对齐预测位置
|
||||
|
||||
# # === 打印每个token的掩码情况 ===
|
||||
# print(f"\n--- Sample {index} Token Loss Mask (length: {len(input_ids)}) ---")
|
||||
# for i, (token_id, mask) in enumerate(zip(input_ids, loss_mask)):
|
||||
# token_str = self.tokenizer.decode([token_id], skip_special_tokens=False)
|
||||
# token_str = token_str.replace('\n', '\\n').replace('\t', '\\t') # 处理换行等不可见字符
|
||||
# print(f"Token {i:3d}: {token_id:5d} -> '{token_str:10s}' | mask: {mask}")
|
||||
# print(f"--- End of Sample {index} ---")
|
||||
# # ================================
|
||||
return X, Y, loss_mask
|
||||
|
||||
|
||||
@ -125,8 +130,8 @@ class DPODataset(Dataset):
|
||||
self.tokenizer = tokenizer
|
||||
self.max_length = max_length
|
||||
self.padding = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
|
||||
self.bos_id = tokenizer('<|im_start|>assistant', add_special_tokens=False).input_ids
|
||||
self.eos_id = tokenizer('<|im_end|>', add_special_tokens=False).input_ids
|
||||
self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant', add_special_tokens=False).input_ids
|
||||
self.eos_id = tokenizer(f'{tokenizer.eos_token}', add_special_tokens=False).input_ids
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
self.data = []
|
||||
for line in f:
|
||||
@ -201,8 +206,8 @@ class RLAIFDataset(Dataset):
|
||||
self.tokenizer = tokenizer
|
||||
self.max_length = max_length
|
||||
self.samples = self.load_data(jsonl_path)
|
||||
self.bos_id = tokenizer('<|im_start|>assistant', add_special_tokens=False).input_ids
|
||||
self.eos_id = tokenizer('<|im_end|>', add_special_tokens=False).input_ids
|
||||
self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant', add_special_tokens=False).input_ids
|
||||
self.eos_id = tokenizer(f'{tokenizer.eos_token}', add_special_tokens=False).input_ids
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
@ -226,7 +231,7 @@ class RLAIFDataset(Dataset):
|
||||
return self.tokenizer.apply_chat_template(
|
||||
messages[:-1],
|
||||
tokenize=False,
|
||||
add_generation_prompt=True
|
||||
add_generation_prompt=True # 这里需要True
|
||||
), answer
|
||||
|
||||
def __getitem__(self, index):
|
||||
|
||||
@ -112,14 +112,14 @@ def main():
|
||||
parser.add_argument('--num_hidden_layers', default=8, type=int)
|
||||
parser.add_argument('--max_seq_len', default=8192, type=int)
|
||||
parser.add_argument('--use_moe', default=False, type=bool)
|
||||
parser.add_argument('--model_mode', default=1, type=int, help="0: 预训练模型，1: SFT-Chat模型，2: RLHF-Chat模型，3: Reason模型，4/5: RLAIF-Chat模型")
|
||||
parser.add_argument('--model_mode', default=1, type=int, help="【必要的】0: 预训练模型，1: SFT-Chat模型 | 【非必要的】2: RLHF模型，3: Reason模型，4/5: RLAIF模型")
|
||||
# 启用长度外推，默认为4倍（注：仅解决位置编码外推问题，不代表模型真实具备长文本能力）
|
||||
parser.add_argument('--inference_rope_scaling', default=False, action='store_true')
|
||||
# 携带历史对话上下文条数history_cnt需要设为偶数，即【用户问题, 模型回答】为1组；设置为0时，即当前query不携带历史上文
|
||||
# 模型未经过多轮对话微调时，在多轮次的长上下文难免出现能力的明显退化，因此需要注意此处设置
|
||||
parser.add_argument('--history_cnt', default=0, type=int)
|
||||
# load模式为1时，前置hidden_size、num_hidden_layers、max_seq_len等参数失效，即以加载的transformers模型的config.json配置为准
|
||||
parser.add_argument('--load', default=1, type=int, help="0: 原生torch权重，1: transformers加载")
|
||||
parser.add_argument('--load', default=0, type=int, help="0: 原生torch权重，1: transformers加载")
|
||||
args = parser.parse_args()
|
||||
|
||||
model, tokenizer = init_model(args)
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 519 KiB After Width: | Height: | Size: 546 KiB |
BIN
images/rope_ppl.png
Normal file
BIN
images/rope_ppl.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 144 KiB |
BIN
images/train_grpo_512.png
Normal file
BIN
images/train_grpo_512.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 214 KiB |
BIN
images/train_grpo_768.png
Normal file
BIN
images/train_grpo_768.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 246 KiB |
BIN
images/train_ppo_512.png
Normal file
BIN
images/train_ppo_512.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 246 KiB |
BIN
images/train_ppo_768.png
Normal file
BIN
images/train_ppo_768.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 241 KiB |
BIN
images/train_spo_768.png
Normal file
BIN
images/train_spo_768.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 234 KiB |
@ -99,7 +99,7 @@ def train_tokenizer():
|
||||
"spaces_between_special_tokens": False,
|
||||
"tokenizer_class": "PreTrainedTokenizerFast",
|
||||
"unk_token": "<|endoftext|>",
|
||||
"chat_template": "{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{{ '<|im_start|>system\\n' + system_message + '<|im_end|>\\n' }}{% else %}{{ '<|im_start|>system\\nYou are a helpful assistant<|im_end|>\\n' }}{% endif %}{% for message in messages %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<|im_start|>user\\n' + content + '<|im_end|>\\n<|im_start|>assistant\\n' }}{% elif message['role'] == 'assistant' %}{{ content + '<|im_end|>' + '\\n' }}{% endif %}{% endfor %}"
|
||||
"chat_template": "{%- if tools %}\n {{- '<|im_start|>system\\n' }}\n {%- if messages[0].role == 'system' %}\n {{- messages[0].content + '\\n\\n' }}\n {%- endif %}\n {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n {%- for tool in tools %}\n {{- \"\\n\" }}\n {{- tool | tojson }}\n {%- endfor %}\n {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n {%- if messages[0]['role'] == 'system' -%}\n {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n {%- else -%}\n {{- '<|im_start|>system\\nYou are a helpful assistant<|im_end|>\\n' }}\n {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n {%- set index = (messages|length - 1) - loop.index0 %}\n {%- if ns.multi_step_tool and message.role == \"user\" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n {%- set ns.multi_step_tool = false %}\n {%- set ns.last_query_index = index %}\n {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n {%- if message.content is string %}\n {%- set content = message.content %}\n {%- else %}\n {%- set content = '' %}\n {%- endif %}\n {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n {%- elif message.role == \"assistant\" %}\n {%- set reasoning_content = '' %}\n {%- if message.reasoning_content is string %}\n {%- set reasoning_content = message.reasoning_content %}\n {%- else %}\n {%- if '</think>' in content %}\n {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n {%- set content = content.split('</think>')[-1].lstrip('\\n') %}\n {%- endif %}\n {%- endif %}\n {%- if loop.index0 > ns.last_query_index %}\n {%- if loop.last or (not loop.last and reasoning_content) %}\n {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n {%- else %}\n {{- '<|im_start|>' + message.role + '\\n' + content }}\n {%- endif %}\n {%- else %}\n {{- '<|im_start|>' + message.role + '\\n' + content }}\n {%- endif %}\n {%- if message.tool_calls %}\n {%- for tool_call in message.tool_calls %}\n {%- if (loop.first and content) or (not loop.first) %}\n {{- '\\n' }}\n {%- endif %}\n {%- if tool_call.function %}\n {%- set tool_call = tool_call.function %}\n {%- endif %}\n {{- '<tool_call>\\n{\"name\": \"' }}\n {{- tool_call.name }}\n {{- '\", \"arguments\": ' }}\n {%- if tool_call.arguments is string %}\n {{- tool_call.arguments }}\n {%- else %}\n {{- tool_call.arguments | tojson }}\n {%- endif %}\n {{- '}\\n</tool_call>' }}\n {%- endfor %}\n {%- endif %}\n {{- '<|im_end|>\\n' }}\n {%- elif message.role == \"tool\" %}\n {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n {{- '<|im_start|>user' }}\n {%- endif %}\n {{- '\\n<tool_response>\\n' }}\n {{- content }}\n {{- '\\n</tool_response>' }}\n {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n {{- '<|im_end|>\\n' }}\n {%- endif %}\n {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n {{- '<|im_start|>assistant\\n' }}\n {%- if enable_thinking is defined and enable_thinking is false %}\n {{- '<think>\\n\\n</think>\\n\\n' }}\n {%- endif %}\n{%- endif %}"
|
||||
}
|
||||
|
||||
# 保存配置文件
|
||||
|
||||
@ -76,7 +76,7 @@ def train_epoch(epoch, wandb):
|
||||
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
if step % args.log_interval == 0:
|
||||
if step % args.log_interval == 0 or step == iter_per_epoch - 1:
|
||||
spend_time = time.time() - start_time
|
||||
Logger(
|
||||
'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} epoch_Time:{}min:'.format(
|
||||
@ -93,7 +93,7 @@ def train_epoch(epoch, wandb):
|
||||
"lr": optimizer.param_groups[-1]['lr'],
|
||||
"epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
|
||||
|
||||
if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
|
||||
if ((step + 1) % args.save_interval == 0 or step == iter_per_epoch - 1) and (not ddp or dist.get_rank() == 0):
|
||||
model.eval()
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'{args.save_dir}/reason_{lm_config.hidden_size}{moe_path}.pth'
|
||||
@ -184,7 +184,7 @@ if __name__ == "__main__":
|
||||
torch.cuda.manual_seed(base_seed + rank)
|
||||
|
||||
if args.use_wandb and (not ddp or ddp_local_rank == 0):
|
||||
import wandb
|
||||
import swanlab as wandb
|
||||
|
||||
wandb.init(project=args.wandb_project, name=args.wandb_run_name)
|
||||
else:
|
||||
|
||||
@ -110,7 +110,7 @@ def train_epoch(epoch, wandb, alpha=0.0, temperature=1.0):
|
||||
scaler.update()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
if step % args.log_interval == 0:
|
||||
if step % args.log_interval == 0 or step == iter_per_epoch - 1:
|
||||
spend_time = time.time() - start_time
|
||||
Logger(
|
||||
'Epoch:[{}/{}]({}/{}) loss:{:.4f} lr:{:.12f} epoch_Time:{}min:'.format(
|
||||
@ -133,7 +133,7 @@ def train_epoch(epoch, wandb, alpha=0.0, temperature=1.0):
|
||||
"last-time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60
|
||||
})
|
||||
|
||||
if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
|
||||
if ((step + 1) % args.save_interval == 0 or step == iter_per_epoch - 1) and (not ddp or dist.get_rank() == 0):
|
||||
model.eval()
|
||||
moe_path = '_moe' if lm_config_student.use_moe else ''
|
||||
ckp = f'{args.save_dir}/full_dist_{lm_config_student.hidden_size}{moe_path}.pth'
|
||||
@ -231,7 +231,7 @@ if __name__ == "__main__":
|
||||
torch.cuda.manual_seed(base_seed + rank)
|
||||
|
||||
if args.use_wandb and (not ddp or ddp_local_rank == 0):
|
||||
import wandb
|
||||
import swanlab as wandb
|
||||
|
||||
wandb.init(project=args.wandb_project, name=args.wandb_run_name)
|
||||
else:
|
||||
|
||||
@ -1,5 +1,6 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
__package__ = "trainer"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
@ -99,7 +100,7 @@ def train_epoch(epoch, wandb):
|
||||
scaler.update()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
if step % args.log_interval == 0:
|
||||
if step % args.log_interval == 0 or step == iter_per_epoch - 1:
|
||||
spend_time = time.time() - start_time
|
||||
Logger(
|
||||
'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} epoch_Time:{}min:'.format(
|
||||
@ -116,7 +117,7 @@ def train_epoch(epoch, wandb):
|
||||
"lr": optimizer.param_groups[-1]['lr'],
|
||||
"epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
|
||||
|
||||
if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
|
||||
if ((step + 1) % args.save_interval == 0 or step == iter_per_epoch - 1) and (not ddp or dist.get_rank() == 0):
|
||||
model.eval()
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'{args.save_dir}/rlhf_{lm_config.hidden_size}{moe_path}.pth'
|
||||
@ -189,7 +190,8 @@ if __name__ == "__main__":
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=args.use_moe)
|
||||
lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
|
||||
use_moe=args.use_moe)
|
||||
args.save_dir = os.path.join(args.out_dir)
|
||||
os.makedirs(args.save_dir, exist_ok=True)
|
||||
os.makedirs(args.out_dir, exist_ok=True)
|
||||
@ -214,7 +216,7 @@ if __name__ == "__main__":
|
||||
torch.cuda.manual_seed(base_seed + rank)
|
||||
|
||||
if args.use_wandb and (not ddp or ddp_local_rank == 0):
|
||||
import wandb
|
||||
import swanlab as wandb
|
||||
|
||||
wandb.init(project=args.wandb_project, name=args.wandb_run_name)
|
||||
else:
|
||||
|
||||
@ -63,7 +63,7 @@ def train_epoch(epoch, wandb):
|
||||
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
if step % args.log_interval == 0:
|
||||
if step % args.log_interval == 0 or step == iter_per_epoch - 1:
|
||||
spend_time = time.time() - start_time
|
||||
Logger(
|
||||
'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} epoch_Time:{}min:'.format(
|
||||
@ -80,7 +80,7 @@ def train_epoch(epoch, wandb):
|
||||
"lr": optimizer.param_groups[-1]['lr'],
|
||||
"epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
|
||||
|
||||
if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
|
||||
if ((step + 1) % args.save_interval == 0 or step == iter_per_epoch - 1) and (not ddp or dist.get_rank() == 0):
|
||||
model.eval()
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'{args.save_dir}/full_sft_{lm_config.hidden_size}{moe_path}.pth'
|
||||
@ -170,7 +170,7 @@ if __name__ == "__main__":
|
||||
torch.cuda.manual_seed(base_seed + rank)
|
||||
|
||||
if args.use_wandb and (not ddp or ddp_local_rank == 0):
|
||||
import wandb
|
||||
import swanlab as wandb
|
||||
|
||||
wandb.init(project=args.wandb_project, name=args.wandb_run_name)
|
||||
else:
|
||||
|
||||
316
trainer/train_grpo.py
Executable file
316
trainer/train_grpo.py
Executable file
@ -0,0 +1,316 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
__package__ = "trainer"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
import argparse
|
||||
import time
|
||||
import re
|
||||
import gc
|
||||
import torch
|
||||
from contextlib import nullcontext
|
||||
import torch.distributed as dist
|
||||
from torch import optim
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModel
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from dataset.lm_dataset import RLAIFDataset
|
||||
from torch.optim.lr_scheduler import CosineAnnealingLR
|
||||
|
||||
|
||||
def Logger(content):
|
||||
if not ddp or dist.get_rank() == 0:
|
||||
print(content)
|
||||
|
||||
|
||||
def calculate_rewards(prompts, responses, reward_model, reward_tokenizer):
|
||||
"""整合所有奖励函数计算总奖励"""
|
||||
|
||||
def reasoning_model_reward(rewards):
|
||||
# 1. 格式奖励（仅针对训练推理模型时使用）
|
||||
pattern = r"^<think>\n.*?\n</think>\n<answer>\n.*?\n</answer>$"
|
||||
pattern2 = r"^<think>\n.*?\n</think>\n\n<answer>\n.*?\n</answer>$"
|
||||
|
||||
matches_pattern = [re.match(pattern, response, re.S) for response in responses]
|
||||
matches_pattern2 = [re.match(pattern2, response, re.S) for response in responses]
|
||||
|
||||
format_rewards = []
|
||||
for match_pattern, match_pattern2 in zip(matches_pattern, matches_pattern2):
|
||||
if match_pattern:
|
||||
format_rewards.append(0.5)
|
||||
elif match_pattern2:
|
||||
format_rewards.append(0.5)
|
||||
else:
|
||||
format_rewards.append(0.0)
|
||||
rewards += torch.tensor(format_rewards, device=args.device)
|
||||
|
||||
# 2. 标记奖励（防止严格奖励稀疏，仅针对训练推理模型时使用）
|
||||
def mark_num(text):
|
||||
reward = 0
|
||||
if text.count("<think>") == 1:
|
||||
reward += 0.25
|
||||
if text.count("</think>") == 1:
|
||||
reward += 0.25
|
||||
if text.count("<answer>") == 1:
|
||||
reward += 0.25
|
||||
if text.count("</answer>") == 1:
|
||||
reward += 0.25
|
||||
return reward
|
||||
|
||||
mark_rewards = [mark_num(response) for response in responses]
|
||||
rewards += torch.tensor(mark_rewards, device=args.device)
|
||||
return rewards
|
||||
|
||||
rewards = torch.zeros(len(responses), device=args.device)
|
||||
|
||||
# 3. 格式奖励
|
||||
if args.reasoning == 1:
|
||||
rewards = reasoning_model_reward(rewards) # 训练推理模型时使用
|
||||
|
||||
# 4. 使用reward model计算奖励
|
||||
with torch.no_grad():
|
||||
reward_model_scores = []
|
||||
batch_size = len(prompts)
|
||||
scale = 3.0
|
||||
|
||||
for i in range(batch_size):
|
||||
for j in range(args.num_generations):
|
||||
response_idx = i * args.num_generations + j
|
||||
response = responses[response_idx]
|
||||
prompt = prompts[i]
|
||||
|
||||
pattern = r"<\|im_start\|>(system|user|assistant)\s+(.*?)<\|im_end\|>"
|
||||
matches = re.findall(pattern, prompt, re.DOTALL)
|
||||
messages = [{"role": role, "content": content.strip()} for role, content in matches]
|
||||
|
||||
tmp_chat = messages + [{"role": "assistant", "content": response}]
|
||||
score = reward_model.get_score(reward_tokenizer, tmp_chat)
|
||||
score = max(min(score, scale), -scale)
|
||||
|
||||
if args.reasoning == 1:
|
||||
answer_match = re.search(r'<answer>(.*?)</answer>', response, re.DOTALL)
|
||||
if answer_match:
|
||||
answer_content = answer_match.group(1).strip()
|
||||
tmp_chat = messages + [{"role": "assistant", "content": answer_content}]
|
||||
answer_score = reward_model.get_score(reward_tokenizer, tmp_chat)
|
||||
answer_score = max(min(answer_score, scale), -scale)
|
||||
score = score * 0.4 + answer_score * 0.6
|
||||
|
||||
reward_model_scores.append(score)
|
||||
|
||||
reward_model_scores = torch.tensor(reward_model_scores, device=args.device)
|
||||
rewards += reward_model_scores
|
||||
|
||||
return rewards
|
||||
|
||||
|
||||
def grpo_train_epoch(epoch, wandb):
|
||||
for step, batch in enumerate(train_loader):
|
||||
prompts = batch['prompt'] # list[str], length B
|
||||
prompt_inputs = tokenizer(prompts, return_tensors="pt", padding=True, return_token_type_ids=False,
|
||||
padding_side="left", add_special_tokens=False).to(args.device) # input_ids: [B, P], attention_mask: [B, P]
|
||||
if args.max_seq_len:
|
||||
prompt_inputs["input_ids"] = prompt_inputs["input_ids"][:, -args.max_seq_len:]
|
||||
prompt_inputs["attention_mask"] = prompt_inputs["attention_mask"][:, -args.max_seq_len:]
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = (model.module if ddp else model).generate(
|
||||
**prompt_inputs, max_new_tokens=args.max_gen_len, do_sample=True, temperature=0.8,
|
||||
num_return_sequences=args.num_generations, pad_token_id=tokenizer.pad_token_id) # [B*num_gen, P+R]
|
||||
|
||||
completion_ids = outputs[:, prompt_inputs["input_ids"].size(1):] # [B*num_gen, R]
|
||||
|
||||
def get_per_token_logps(mdl, input_ids, n_keep):
|
||||
input_ids = input_ids.detach().clone() if input_ids.is_inference() else input_ids
|
||||
logits = mdl(input_ids, logits_to_keep=n_keep + 1).logits[:, :-1, :]
|
||||
per_token_logps = []
|
||||
for logits_row, ids_row in zip(logits, input_ids[:, -n_keep:]):
|
||||
ids_row = ids_row.detach().clone() if ids_row.is_inference() else ids_row
|
||||
per_token_logps.append(torch.gather(logits_row.log_softmax(dim=-1), 1, ids_row.unsqueeze(1)).squeeze(1))
|
||||
return torch.stack(per_token_logps)
|
||||
|
||||
per_token_logps = get_per_token_logps(model, outputs, completion_ids.size(1)) # [B*num_gen, R]
|
||||
with torch.no_grad():
|
||||
ref_per_token_logps = get_per_token_logps(ref_model, outputs, completion_ids.size(1)) # [B*num_gen, R]
|
||||
|
||||
completions = tokenizer.batch_decode(completion_ids, skip_special_tokens=True)
|
||||
rewards = calculate_rewards(prompts, completions, reward_model, reward_tokenizer).to(args.device) # [B*num_gen]
|
||||
|
||||
grouped_rewards = rewards.view(-1, args.num_generations) # [B, num_gen]
|
||||
mean_r = grouped_rewards.mean(dim=1).repeat_interleave(args.num_generations) # [B*num_gen]
|
||||
std_r = grouped_rewards.std(dim=1).repeat_interleave(args.num_generations) # [B*num_gen]
|
||||
advantages = torch.clamp((rewards - mean_r) / (std_r + 1e-4), -10, 10)
|
||||
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8) # [B*num_gen]
|
||||
|
||||
is_eos = completion_ids == tokenizer.eos_token_id # [B*num_gen, R]
|
||||
eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=args.device)
|
||||
eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)]
|
||||
completion_mask = (torch.arange(is_eos.size(1), device=args.device).expand(is_eos.size(0), -1) <= eos_idx.unsqueeze(1)).int() # [B*num_gen, R]
|
||||
|
||||
kl_div = ref_per_token_logps - per_token_logps
|
||||
per_token_kl = torch.exp(kl_div) - kl_div - 1 # [B*num_gen, R]
|
||||
per_token_loss = -(torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1) - args.beta * per_token_kl) # [B*num_gen, R]
|
||||
loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() / args.accumulation_steps # scalar
|
||||
loss.backward()
|
||||
|
||||
if (step + 1) % args.accumulation_steps == 0:
|
||||
if args.grad_clip > 0:
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||||
optimizer.step()
|
||||
scheduler.step()
|
||||
optimizer.zero_grad()
|
||||
|
||||
if step % args.log_interval == 0 or step == iter_per_epoch - 1:
|
||||
policy_loss_val = loss.item()
|
||||
avg_reward_val = rewards.mean().item()
|
||||
avg_len_val = completion_mask.sum(dim=1).float().mean().item()
|
||||
current_lr = optimizer.param_groups[0]['lr']
|
||||
|
||||
Logger(
|
||||
f'Epoch: {epoch}, Step: {step + 1}/{iter_per_epoch}, '
|
||||
f'Actor Loss: {policy_loss_val:.4f}, Reward: {avg_reward_val:.4f}, '
|
||||
f'Avg Response Len: {avg_len_val:.2f}, LR: {current_lr:.2e}')
|
||||
|
||||
if wandb and (not ddp or dist.get_rank() == 0):
|
||||
log_dict = {
|
||||
"policy_loss": policy_loss_val,
|
||||
"reward": avg_reward_val,
|
||||
"avg_response_len": avg_len_val,
|
||||
"advantages_mean": advantages.mean().item(),
|
||||
"learning_rate": current_lr
|
||||
}
|
||||
wandb.log(log_dict)
|
||||
|
||||
if ((step + 1) % args.save_interval == 0 or step == iter_per_epoch - 1) and (not ddp or dist.get_rank() == 0):
|
||||
model.eval()
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
suffix = 'grpo'
|
||||
ckp = f'{args.save_dir}/{suffix}_{lm_config.hidden_size}{moe_path}.pth'
|
||||
|
||||
state_dict = model.module.state_dict() if isinstance(model,
|
||||
torch.nn.parallel.DistributedDataParallel) else model.state_dict()
|
||||
torch.save({k: v.half() for k, v in state_dict.items()}, ckp)
|
||||
model.train()
|
||||
|
||||
del prompt_inputs, outputs, completion_ids, per_token_logps, ref_per_token_logps
|
||||
del completions, rewards, grouped_rewards, mean_r, std_r, advantages, completion_mask
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
|
||||
def init_model(lm_config):
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model/')
|
||||
model = MiniMindForCausalLM(lm_config)
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'{args.save_dir}/full_sft_{lm_config.hidden_size}{moe_path}.pth'
|
||||
if args.reasoning == 1:
|
||||
ckp = f'{args.save_dir}/reason_{lm_config.hidden_size}{moe_path}.pth'
|
||||
state_dict = torch.load(ckp, map_location=args.device)
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
|
||||
ref_model = MiniMindForCausalLM(lm_config)
|
||||
ref_model.load_state_dict(state_dict, strict=False)
|
||||
ref_model.eval().requires_grad_(False)
|
||||
|
||||
Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
||||
model = model.to(args.device)
|
||||
ref_model = ref_model.to(args.device)
|
||||
|
||||
reward_name = "../../internlm2-1_8b-reward"
|
||||
reward_model = AutoModel.from_pretrained(
|
||||
reward_name,
|
||||
device_map="cuda",
|
||||
torch_dtype=torch.float16,
|
||||
trust_remote_code=True,
|
||||
).to(args.device).eval().requires_grad_(False)
|
||||
reward_tokenizer = AutoTokenizer.from_pretrained(reward_name, trust_remote_code=True)
|
||||
|
||||
return model, ref_model, tokenizer, reward_model, reward_tokenizer
|
||||
|
||||
|
||||
def init_distributed_mode():
|
||||
if not ddp: return
|
||||
global ddp_local_rank, DEVICE
|
||||
dist.init_process_group(backend="nccl")
|
||||
ddp_local_rank = int(os.environ["LOCAL_RANK"])
|
||||
DEVICE = f"cuda:{ddp_local_rank}"
|
||||
torch.cuda.set_device(DEVICE)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--out_dir", type=str, default="../out")
|
||||
parser.add_argument("--epochs", type=int, default=1)
|
||||
parser.add_argument("--batch_size", type=int, default=2)
|
||||
parser.add_argument("--learning_rate", type=float, default=8e-8)
|
||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16")
|
||||
parser.add_argument("--use_wandb", action="store_true")
|
||||
parser.add_argument("--wandb_project", type=str, default="MiniMind-GRPO")
|
||||
parser.add_argument("--num_workers", type=int, default=1)
|
||||
parser.add_argument("--ddp", action="store_true")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1)
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0)
|
||||
parser.add_argument("--log_interval", type=int, default=1)
|
||||
parser.add_argument("--save_interval", type=int, default=10)
|
||||
parser.add_argument('--hidden_size', default=512, type=int)
|
||||
parser.add_argument('--num_hidden_layers', default=8, type=int)
|
||||
parser.add_argument('--use_moe', default=False, type=bool)
|
||||
parser.add_argument('--max_seq_len', default=66, type=int)
|
||||
parser.add_argument("--max_gen_len", type=int, default=1536)
|
||||
parser.add_argument("--data_path", type=str, default="../dataset/rlaif-mini.jsonl")
|
||||
parser.add_argument("--num_generations", type=int, default=8)
|
||||
parser.add_argument("--beta", type=float, default=0.02)
|
||||
parser.add_argument("--reasoning", type=int, default=1, help='0:普通模型，1:推理模型')
|
||||
args = parser.parse_args()
|
||||
|
||||
lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
|
||||
max_seq_len=args.max_seq_len + args.max_gen_len,
|
||||
use_moe=args.use_moe)
|
||||
args.save_dir = os.path.join(args.out_dir)
|
||||
os.makedirs(args.save_dir, exist_ok=True)
|
||||
os.makedirs(args.out_dir, exist_ok=True)
|
||||
|
||||
ctx = nullcontext() if "cuda" not in args.device else torch.amp.autocast('cuda')
|
||||
ddp = int(os.environ.get("RANK", -1)) != -1
|
||||
ddp_local_rank, DEVICE = 0, "cuda:0"
|
||||
|
||||
base_seed = 1337
|
||||
torch.manual_seed(base_seed)
|
||||
torch.cuda.manual_seed(base_seed)
|
||||
|
||||
if ddp:
|
||||
init_distributed_mode()
|
||||
args.device = torch.device(DEVICE)
|
||||
rank = dist.get_rank()
|
||||
torch.manual_seed(base_seed + rank)
|
||||
# 同时设置 CUDA 的随机种子
|
||||
torch.cuda.manual_seed(base_seed + rank)
|
||||
|
||||
if args.use_wandb and (not ddp or ddp_local_rank == 0):
|
||||
import swanlab as wandb
|
||||
|
||||
wandb.init(project=args.wandb_project)
|
||||
else:
|
||||
wandb = None
|
||||
|
||||
model, ref_model, tokenizer, reward_model, reward_tokenizer = init_model(lm_config=lm_config)
|
||||
train_ds = RLAIFDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
|
||||
train_sampler = DistributedSampler(train_ds) if ddp else None
|
||||
train_loader = DataLoader(train_ds, batch_size=args.batch_size, pin_memory=True,
|
||||
drop_last=False, shuffle=False,
|
||||
num_workers=args.num_workers, sampler=train_sampler)
|
||||
|
||||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
|
||||
|
||||
iter_per_epoch = len(train_loader)
|
||||
total_optimizer_steps = (iter_per_epoch // args.accumulation_steps) * args.epochs
|
||||
scheduler = CosineAnnealingLR(optimizer, T_max=total_optimizer_steps, eta_min=args.learning_rate / 10)
|
||||
|
||||
if ddp:
|
||||
model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
|
||||
model = DistributedDataParallel(model, device_ids=[ddp_local_rank])
|
||||
|
||||
for epoch in range(args.epochs):
|
||||
grpo_train_epoch(epoch, wandb)
|
||||
@ -64,7 +64,7 @@ def train_epoch(epoch, wandb):
|
||||
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
if step % args.log_interval == 0:
|
||||
if step % args.log_interval == 0 or step == iter_per_epoch - 1:
|
||||
spend_time = time.time() - start_time
|
||||
Logger(
|
||||
'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} epoch_Time:{}min:'.format(
|
||||
@ -81,7 +81,7 @@ def train_epoch(epoch, wandb):
|
||||
"lr": optimizer.param_groups[-1]['lr'],
|
||||
"epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
|
||||
|
||||
if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
|
||||
if ((step + 1) % args.save_interval == 0 or step == iter_per_epoch - 1) and (not ddp or dist.get_rank() == 0):
|
||||
model.eval()
|
||||
lora_save_path = f'{args.save_dir}/lora/{args.lora_name}_{lm_config.hidden_size}.pth'
|
||||
os.makedirs(os.path.dirname(lora_save_path), exist_ok=True)
|
||||
@ -115,7 +115,7 @@ def init_distributed_mode():
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="MiniMind SFT with LoRA")
|
||||
parser.add_argument("--out_dir", type=str, default="../out")
|
||||
parser.add_argument("--epochs", type=int, default=10)
|
||||
parser.add_argument("--epochs", type=int, default=50)
|
||||
parser.add_argument("--batch_size", type=int, default=32)
|
||||
parser.add_argument("--learning_rate", type=float, default=1e-4)
|
||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
@ -127,15 +127,15 @@ if __name__ == "__main__":
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1)
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0)
|
||||
parser.add_argument("--warmup_iters", type=int, default=0)
|
||||
parser.add_argument("--log_interval", type=int, default=100)
|
||||
parser.add_argument("--save_interval", type=int, default=100)
|
||||
parser.add_argument("--log_interval", type=int, default=10)
|
||||
parser.add_argument("--save_interval", type=int, default=1)
|
||||
parser.add_argument('--local_rank', type=int, default=-1)
|
||||
parser.add_argument('--hidden_size', default=512, type=int)
|
||||
parser.add_argument('--num_hidden_layers', default=8, type=int)
|
||||
parser.add_argument('--max_seq_len', default=512, type=int)
|
||||
parser.add_argument('--use_moe', default=False, type=bool)
|
||||
parser.add_argument("--data_path", type=str, default="../dataset/lora_medical.jsonl")
|
||||
parser.add_argument("--lora_name", type=str, default="lora_medical", help="根据任务保存成lora_(英文/医学/心理...)")
|
||||
parser.add_argument("--data_path", type=str, default="../dataset/lora_identity.jsonl")
|
||||
parser.add_argument("--lora_name", type=str, default="lora_identity", help="根据任务保存成lora_(英文/医学/心理...)")
|
||||
args = parser.parse_args()
|
||||
|
||||
lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
|
||||
@ -163,7 +163,7 @@ if __name__ == "__main__":
|
||||
|
||||
args.wandb_run_name = f"MiniMind-Lora-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
|
||||
if args.use_wandb and (not ddp or ddp_local_rank == 0):
|
||||
import wandb
|
||||
import swanlab as wandb
|
||||
|
||||
wandb.init(project=args.wandb_project, name=args.wandb_run_name)
|
||||
else:
|
||||
|
||||
373
trainer/train_ppo.py
Normal file
373
trainer/train_ppo.py
Normal file
@ -0,0 +1,373 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
__package__ = "trainer"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
import argparse
|
||||
import re
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn.functional as F
|
||||
from torch import optim, nn
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from dataset.lm_dataset import RLAIFDataset
|
||||
from torch.nn.utils import clip_grad_norm_
|
||||
from torch.optim.lr_scheduler import CosineAnnealingLR
|
||||
|
||||
|
||||
def Logger(content):
|
||||
if not ddp or dist.get_rank() == 0:
|
||||
print(content)
|
||||
|
||||
|
||||
def calculate_rewards(prompts, responses, reward_model, reward_tokenizer):
|
||||
"""整合所有奖励函数计算总奖励"""
|
||||
|
||||
def reasoning_model_reward(rewards):
|
||||
# 1. 格式奖励（仅针对训练推理模型时使用）
|
||||
pattern = r"^<think>\n.*?\n</think>\n<answer>\n.*?\n</answer>$"
|
||||
pattern2 = r"^<think>\n.*?\n</think>\n\n<answer>\n.*?\n</answer>$"
|
||||
|
||||
matches_pattern = [re.match(pattern, response, re.S) for response in responses]
|
||||
matches_pattern2 = [re.match(pattern2, response, re.S) for response in responses]
|
||||
|
||||
format_rewards = []
|
||||
for match_pattern, match_pattern2 in zip(matches_pattern, matches_pattern2):
|
||||
if match_pattern:
|
||||
format_rewards.append(0.5)
|
||||
elif match_pattern2:
|
||||
format_rewards.append(0.5)
|
||||
else:
|
||||
format_rewards.append(0.0)
|
||||
rewards += torch.tensor(format_rewards, device=args.device)
|
||||
|
||||
# 2. 标记奖励（防止严格奖励稀疏，仅针对训练推理模型时使用）
|
||||
def mark_num(text):
|
||||
reward = 0
|
||||
if text.count("<think>") == 1:
|
||||
reward += 0.25
|
||||
if text.count("</think>") == 1:
|
||||
reward += 0.25
|
||||
if text.count("<answer>") == 1:
|
||||
reward += 0.25
|
||||
if text.count("</answer>") == 1:
|
||||
reward += 0.25
|
||||
return reward
|
||||
|
||||
mark_rewards = [mark_num(response) for response in responses]
|
||||
rewards += torch.tensor(mark_rewards, device=args.device)
|
||||
return rewards
|
||||
|
||||
rewards = torch.zeros(len(responses), device=args.device)
|
||||
|
||||
# 格式奖励
|
||||
if args.reasoning == 1:
|
||||
rewards = reasoning_model_reward(rewards) # 训练推理模型时使用
|
||||
|
||||
# 使用reward model计算整个response的奖励
|
||||
with torch.no_grad():
|
||||
reward_model_scores = []
|
||||
for prompt, response in zip(prompts, responses):
|
||||
pattern = r"<\|im_start\|>(system|user|assistant)\s+(.*?)<\|im_end\|>"
|
||||
matches = re.findall(pattern, prompt, re.DOTALL)
|
||||
messages = [{"role": role, "content": content.strip()} for role, content in matches]
|
||||
|
||||
tmp_chat = messages + [{"role": "assistant", "content": response}]
|
||||
score = reward_model.get_score(reward_tokenizer, tmp_chat)
|
||||
|
||||
scale = 3.0
|
||||
score = max(min(score, scale), -scale)
|
||||
|
||||
# 当args.reasoning=1时，额外计算<answer>内容的奖励
|
||||
if args.reasoning == 1:
|
||||
answer_match = re.search(r'<answer>(.*?)</answer>', response, re.DOTALL)
|
||||
if answer_match:
|
||||
answer_content = answer_match.group(1).strip()
|
||||
# 对answer内容单独计算reward
|
||||
tmp_chat = messages + [{"role": "assistant", "content": answer_content}]
|
||||
answer_score = reward_model.get_score(reward_tokenizer, tmp_chat)
|
||||
answer_score = max(min(answer_score, scale), -scale)
|
||||
|
||||
score = score * 0.4 + answer_score * 0.6
|
||||
reward_model_scores.append(score)
|
||||
|
||||
reward_model_scores = torch.tensor(reward_model_scores, device=args.device)
|
||||
rewards += reward_model_scores
|
||||
|
||||
return rewards
|
||||
|
||||
|
||||
def ppo_train_epoch(epoch: int, wandb_run, old_actor_model, ref_model, actor_scheduler, critic_scheduler):
|
||||
actor_model.train()
|
||||
critic_model.train()
|
||||
is_master = (not ddp) or dist.get_rank() == 0
|
||||
|
||||
for step, batch in enumerate(train_loader):
|
||||
prompts = batch["prompt"] # list[str], length B
|
||||
enc = tokenizer(prompts, return_tensors="pt", padding=True, truncation=True,
|
||||
max_length=args.max_seq_len).to(args.device) # input_ids: [B, P], attention_mask: [B, P]
|
||||
prompt_lengths = enc.attention_mask.sum(dim=1) # [B]
|
||||
|
||||
with torch.no_grad():
|
||||
gen_out = actor_model.generate(
|
||||
input_ids=enc.input_ids, attention_mask=enc.attention_mask,
|
||||
max_new_tokens=args.max_gen_len, do_sample=True, temperature=0.8,
|
||||
pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id) # [B, P+R]
|
||||
|
||||
responses_text = [tokenizer.decode(gen_out[i, prompt_lengths[i]:], skip_special_tokens=True) for i in range(len(prompts))]
|
||||
rewards = calculate_rewards(prompts, responses_text, reward_model, reward_tokenizer) # [B]
|
||||
|
||||
full_mask = (gen_out != tokenizer.pad_token_id).long() # [B, P+R]
|
||||
values_seq = critic_model(input_ids=gen_out, attention_mask=full_mask) # [B, P+R]
|
||||
last_indices = full_mask.sum(dim=1) - 1 # [B]
|
||||
values = values_seq[torch.arange(values_seq.size(0), device=values_seq.device), last_indices] # [B]
|
||||
advantages = rewards - values.detach() # [B]
|
||||
|
||||
logits = actor_model(input_ids=gen_out, attention_mask=full_mask).logits # [B, P+R, V]
|
||||
labels = gen_out[:, 1:].clone() # [B, P+R-1]
|
||||
logp_tokens = F.log_softmax(logits[:, :-1], dim=-1).gather(2, labels.unsqueeze(-1)).squeeze(-1) # [B, P+R-1]
|
||||
seq_len = gen_out.size(1) - 1
|
||||
resp_mask = torch.arange(seq_len, device=gen_out.device).unsqueeze(0) >= prompt_lengths.unsqueeze(1)
|
||||
final_mask = resp_mask & (~labels.eq(tokenizer.pad_token_id)) # [B, P+R-1]
|
||||
actor_logp = (logp_tokens * final_mask).sum(dim=1) # [B]
|
||||
|
||||
with torch.no_grad():
|
||||
old_logits = old_actor_model(input_ids=gen_out, attention_mask=full_mask).logits # [B, P+R, V]
|
||||
old_logp_tokens = F.log_softmax(old_logits[:, :-1], dim=-1).gather(2, labels.unsqueeze(-1)).squeeze(-1) # [B, P+R-1]
|
||||
old_logp = (old_logp_tokens * final_mask).sum(dim=1) # [B]
|
||||
|
||||
ref_logits = ref_model(input_ids=gen_out, attention_mask=full_mask).logits # [B, P+R, V]
|
||||
ref_logp_tokens = F.log_softmax(ref_logits[:, :-1], dim=-1).gather(2, labels.unsqueeze(-1)).squeeze(-1) # [B, P+R-1]
|
||||
ref_logp = (ref_logp_tokens * final_mask).sum(dim=1) # [B]
|
||||
|
||||
kl = (actor_logp - old_logp).mean() # scalar
|
||||
kl_ref = (actor_logp - ref_logp).mean() # scalar
|
||||
ratio = torch.exp(actor_logp - old_logp) # [B]
|
||||
surr1 = ratio * advantages # [B]
|
||||
surr2 = torch.clamp(ratio, 1.0 - args.clip_epsilon, 1.0 + args.clip_epsilon) * advantages # [B]
|
||||
policy_loss = -torch.min(surr1, surr2).mean() # scalar
|
||||
value_loss = F.mse_loss(values, rewards) # scalar
|
||||
loss = policy_loss + args.vf_coef * value_loss + args.kl_coef * kl_ref # scalar
|
||||
loss.backward()
|
||||
|
||||
if (step + 1) % args.accumulation_steps == 0:
|
||||
clip_grad_norm_(actor_model.parameters(), args.grad_clip)
|
||||
clip_grad_norm_(critic_model.parameters(), args.grad_clip)
|
||||
actor_optimizer.step()
|
||||
critic_optimizer.step()
|
||||
actor_scheduler.step()
|
||||
critic_scheduler.step()
|
||||
actor_optimizer.zero_grad()
|
||||
critic_optimizer.zero_grad()
|
||||
|
||||
if is_master:
|
||||
response_ids = gen_out[:, enc.input_ids.shape[1]:]
|
||||
is_eos = (response_ids == tokenizer.eos_token_id)
|
||||
eos_indices = torch.argmax(is_eos.int(), dim=1)
|
||||
has_eos = is_eos.any(dim=1)
|
||||
lengths = torch.where(has_eos, eos_indices + 1, torch.tensor(response_ids.shape[1], device=is_eos.device))
|
||||
avg_len = lengths.float().mean()
|
||||
|
||||
actor_loss_val = policy_loss.item()
|
||||
critic_loss_val = value_loss.item()
|
||||
reward_val = rewards.mean().item()
|
||||
kl_val = kl.item()
|
||||
kl_ref_val = kl_ref.item()
|
||||
avg_len_val = avg_len.item()
|
||||
actor_lr = actor_optimizer.param_groups[0]['lr']
|
||||
critic_lr = critic_optimizer.param_groups[0]['lr']
|
||||
|
||||
if wandb_run is not None:
|
||||
wandb_run.log({
|
||||
"actor_loss": actor_loss_val,
|
||||
"critic_loss": critic_loss_val,
|
||||
"reward": reward_val,
|
||||
"kl": kl_val,
|
||||
"kl_ref": kl_ref_val,
|
||||
"avg_response_len": avg_len_val,
|
||||
"actor_lr": actor_lr,
|
||||
})
|
||||
|
||||
Logger(f"Epoch: {epoch}, Step: {step + 1}/{len(train_loader)}, "
|
||||
f"Actor Loss: {actor_loss_val:.4f}, Critic Loss: {critic_loss_val:.4f}, "
|
||||
f"Reward: {reward_val:.4f}, KL: {kl_val:.4f}, KL_ref: {kl_ref_val:.4f}, "
|
||||
f"Avg Response Len: {avg_len_val:.2f}, Actor LR: {actor_lr:.2e}, Critic LR: {critic_lr:.2e}")
|
||||
|
||||
if (step + 1) % args.update_old_actor_freq == 0:
|
||||
state_dict = actor_model.module.state_dict() if isinstance(actor_model, torch.nn.parallel.DistributedDataParallel) else actor_model.state_dict()
|
||||
old_actor_model.load_state_dict({k: v.detach().cpu() for k, v in state_dict.items()})
|
||||
old_actor_model.to(args.device)
|
||||
|
||||
if ((step + 1) % args.save_interval == 0 or step == iter_per_epoch - 1) and (not ddp or dist.get_rank() == 0):
|
||||
actor_model.eval()
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'{args.save_dir}/ppo_actor_{lm_config.hidden_size}{moe_path}.pth'
|
||||
|
||||
if isinstance(actor_model, torch.nn.parallel.DistributedDataParallel):
|
||||
state_dict = actor_model.module.state_dict()
|
||||
else:
|
||||
state_dict = actor_model.state_dict()
|
||||
|
||||
state_dict = {k: v.half() for k, v in state_dict.items()} # 半精度保存
|
||||
torch.save(state_dict, ckp)
|
||||
actor_model.train()
|
||||
|
||||
|
||||
# 自定义的Critic模型，继承自MiniMindLM
|
||||
class CriticModel(MiniMindForCausalLM):
|
||||
def __init__(self, params):
|
||||
super().__init__(params)
|
||||
# 替换lm_head为输出单一价值的线性层
|
||||
self.value_head = nn.Linear(params.hidden_size, 1)
|
||||
|
||||
def forward(self, input_ids=None, attention_mask=None, **kwargs):
|
||||
# 使用基础模型获取隐藏状态
|
||||
outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
|
||||
# self.model 返回的是一个元组，第一个元素是 last_hidden_state
|
||||
hidden_states = self.model.norm(outputs[0])
|
||||
# 使用value_head获取价值估计
|
||||
values = self.value_head(hidden_states).squeeze(-1)
|
||||
return values
|
||||
|
||||
|
||||
def init_model(lm_config):
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model/', padding_side='left')
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'{args.save_dir}/{"reason" if args.reasoning == 1 else "full_sft"}_{lm_config.hidden_size}{moe_path}.pth'
|
||||
state_dict = torch.load(ckp, map_location=args.device)
|
||||
|
||||
actor_model = MiniMindForCausalLM(lm_config)
|
||||
actor_model.load_state_dict(state_dict, strict=False)
|
||||
actor_model = actor_model.to(args.device)
|
||||
|
||||
old_actor_model = MiniMindForCausalLM(lm_config)
|
||||
old_actor_model.load_state_dict(state_dict, strict=False)
|
||||
old_actor_model = old_actor_model.eval().requires_grad_(False).to(args.device)
|
||||
|
||||
ref_model = MiniMindForCausalLM(lm_config)
|
||||
ref_model.load_state_dict(state_dict, strict=False)
|
||||
ref_model = ref_model.eval().requires_grad_(False).to(args.device)
|
||||
|
||||
critic_model = CriticModel(lm_config)
|
||||
critic_model.load_state_dict(state_dict, strict=False)
|
||||
critic_model = critic_model.to(args.device)
|
||||
|
||||
reward_name = "../../internlm2-1_8b-reward"
|
||||
reward_model = AutoModel.from_pretrained(
|
||||
reward_name, device_map="cuda", torch_dtype=torch.float32, trust_remote_code=True
|
||||
).to(args.device).eval().requires_grad_(False)
|
||||
reward_tokenizer = AutoTokenizer.from_pretrained(reward_name, trust_remote_code=True)
|
||||
|
||||
Logger(f'Actor模型总参数量：{sum(p.numel() for p in actor_model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
||||
Logger(f'Critic模型总参数量：{sum(p.numel() for p in critic_model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
||||
|
||||
return actor_model, old_actor_model, ref_model, critic_model, reward_model, tokenizer, reward_tokenizer
|
||||
|
||||
|
||||
def init_distributed_mode():
|
||||
if not ddp: return
|
||||
global ddp_local_rank, DEVICE
|
||||
dist.init_process_group(backend="nccl")
|
||||
ddp_local_rank = int(os.environ["LOCAL_RANK"])
|
||||
DEVICE = f"cuda:{ddp_local_rank}"
|
||||
torch.cuda.set_device(DEVICE)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--out_dir", type=str, default="../out")
|
||||
parser.add_argument("--epochs", type=int, default=1)
|
||||
parser.add_argument("--batch_size", type=int, default=2)
|
||||
parser.add_argument("--learning_rate", type=float, default=8e-8)
|
||||
parser.add_argument("--critic_learning_rate", type=float, default=8e-8)
|
||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16")
|
||||
parser.add_argument("--use_wandb", action="store_true")
|
||||
parser.add_argument("--wandb_project", type=str, default="MiniMind-PPO")
|
||||
parser.add_argument("--num_workers", type=int, default=1)
|
||||
parser.add_argument("--ddp", action="store_true")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1)
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0)
|
||||
parser.add_argument("--log_interval", type=int, default=1)
|
||||
parser.add_argument("--save_interval", type=int, default=10)
|
||||
parser.add_argument('--hidden_size', default=512, type=int)
|
||||
parser.add_argument('--num_hidden_layers', default=8, type=int)
|
||||
parser.add_argument('--use_moe', default=False, type=bool)
|
||||
parser.add_argument('--max_seq_len', default=66, type=int)
|
||||
parser.add_argument("--max_gen_len", type=int, default=1536)
|
||||
parser.add_argument("--data_path", type=str, default="../dataset/rlaif-mini.jsonl")
|
||||
parser.add_argument("--clip_epsilon", type=float, default=0.1)
|
||||
parser.add_argument("--vf_coef", type=float, default=0.5)
|
||||
parser.add_argument("--kl_coef", type=float, default=0.02, help="KL散度惩罚系数")
|
||||
parser.add_argument("--reasoning", type=int, default=1, help='0:普通模型，1:推理模型')
|
||||
parser.add_argument("--update_old_actor_freq", type=int, default=4, help="频率：每处理n个batch后更新old_actor_model")
|
||||
args = parser.parse_args()
|
||||
|
||||
lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
|
||||
use_moe=args.use_moe)
|
||||
args.save_dir = os.path.join(args.out_dir)
|
||||
os.makedirs(args.save_dir, exist_ok=True)
|
||||
os.makedirs(args.out_dir, exist_ok=True)
|
||||
|
||||
ddp = int(os.environ.get("RANK", -1)) != -1
|
||||
ddp_local_rank, DEVICE = 0, "cuda:0"
|
||||
base_seed = 1337
|
||||
torch.manual_seed(base_seed)
|
||||
torch.cuda.manual_seed(base_seed)
|
||||
|
||||
if ddp:
|
||||
init_distributed_mode()
|
||||
args.device = torch.device(DEVICE)
|
||||
rank = dist.get_rank()
|
||||
torch.manual_seed(base_seed + rank)
|
||||
# 同时设置 CUDA 的随机种子
|
||||
torch.cuda.manual_seed(base_seed + rank)
|
||||
|
||||
if args.use_wandb and (not ddp or ddp_local_rank == 0):
|
||||
import swanlab as wandb
|
||||
|
||||
wandb.init(project=args.wandb_project)
|
||||
else:
|
||||
wandb = None
|
||||
|
||||
# 初始化所有模型
|
||||
actor_model, old_actor_model, ref_model, critic_model, reward_model, tokenizer, reward_tokenizer = init_model(lm_config=lm_config)
|
||||
|
||||
# 准备数据集和数据加载器
|
||||
train_ds = RLAIFDataset(args.data_path, tokenizer, max_length=(args.max_seq_len + args.max_gen_len))
|
||||
train_sampler = DistributedSampler(train_ds) if ddp else None
|
||||
train_loader = DataLoader(train_ds, batch_size=args.batch_size, pin_memory=True,
|
||||
drop_last=False, shuffle=False,
|
||||
num_workers=args.num_workers, sampler=train_sampler)
|
||||
|
||||
# 初始化优化器
|
||||
actor_optimizer = optim.AdamW(actor_model.parameters(), lr=args.learning_rate)
|
||||
critic_optimizer = optim.AdamW(critic_model.parameters(), lr=args.critic_learning_rate)
|
||||
|
||||
iter_per_epoch = len(train_loader)
|
||||
total_optimizer_steps = (iter_per_epoch // args.accumulation_steps) * args.epochs
|
||||
actor_scheduler = CosineAnnealingLR(actor_optimizer, T_max=total_optimizer_steps, eta_min=args.learning_rate / 10)
|
||||
critic_scheduler = CosineAnnealingLR(critic_optimizer, T_max=total_optimizer_steps,
|
||||
eta_min=args.critic_learning_rate / 10)
|
||||
|
||||
# 如果使用分布式训练，包装模型
|
||||
if ddp:
|
||||
actor_model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
|
||||
critic_model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
|
||||
actor_model = DistributedDataParallel(actor_model, device_ids=[ddp_local_rank])
|
||||
critic_model = DistributedDataParallel(critic_model, device_ids=[ddp_local_rank])
|
||||
# old_actor_model 不需要DDP包装，因为它只在主进程上用于计算，并且不进行梯度更新
|
||||
old_actor_model.to(args.device)
|
||||
|
||||
for epoch in range(args.epochs):
|
||||
ppo_train_epoch(epoch, wandb, old_actor_model, ref_model, actor_scheduler, critic_scheduler)
|
||||
|
||||
if ddp:
|
||||
dist.destroy_process_group()
|
||||
@ -1,5 +1,6 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
__package__ = "trainer"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
@ -62,7 +63,7 @@ def train_epoch(epoch, wandb):
|
||||
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
if step % args.log_interval == 0:
|
||||
if step % args.log_interval == 0 or step == iter_per_epoch - 1:
|
||||
spend_time = time.time() - start_time
|
||||
Logger(
|
||||
'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} epoch_Time:{}min:'.format(
|
||||
@ -79,7 +80,7 @@ def train_epoch(epoch, wandb):
|
||||
"lr": optimizer.param_groups[-1]['lr'],
|
||||
"epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
|
||||
|
||||
if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
|
||||
if ((step + 1) % args.save_interval == 0 or step == iter_per_epoch - 1) and (not ddp or dist.get_rank() == 0):
|
||||
model.eval()
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'{args.save_dir}/pretrain_{lm_config.hidden_size}{moe_path}.pth'
|
||||
@ -140,7 +141,8 @@ if __name__ == "__main__":
|
||||
parser.add_argument("--data_path", type=str, default="../dataset/pretrain_hq.jsonl")
|
||||
args = parser.parse_args()
|
||||
|
||||
lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=args.use_moe)
|
||||
lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
|
||||
use_moe=args.use_moe)
|
||||
args.save_dir = os.path.join(args.out_dir)
|
||||
os.makedirs(args.save_dir, exist_ok=True)
|
||||
os.makedirs(args.out_dir, exist_ok=True)
|
||||
@ -167,7 +169,7 @@ if __name__ == "__main__":
|
||||
torch.cuda.manual_seed(base_seed + rank)
|
||||
|
||||
if args.use_wandb and (not ddp or ddp_local_rank == 0):
|
||||
import wandb
|
||||
import swanlab as wandb
|
||||
|
||||
wandb.init(project=args.wandb_project, name=args.wandb_run_name)
|
||||
else:
|
||||
|
||||
367
trainer/train_spo.py
Executable file
367
trainer/train_spo.py
Executable file
@ -0,0 +1,367 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
__package__ = "trainer"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
import argparse
|
||||
import time
|
||||
import re
|
||||
import gc
|
||||
import torch
|
||||
from contextlib import nullcontext
|
||||
import torch.distributed as dist
|
||||
from torch import optim
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModel
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from dataset.lm_dataset import RLAIFDataset
|
||||
from torch.optim.lr_scheduler import CosineAnnealingLR
|
||||
from collections import defaultdict
|
||||
|
||||
|
||||
class AutoAdaptiveValueTracker:
|
||||
def __init__(self, rho_mode='kl', rho_const=0.9, D_half=0.06, clip_lower=0.5, clip_upper=0.96):
|
||||
self.rho_mode = rho_mode
|
||||
self.rho_const = rho_const
|
||||
self.D_half = D_half
|
||||
self.clip_lower = clip_lower
|
||||
self.clip_upper = clip_upper
|
||||
# Stable initialization following N_init = 1/(1-clip_lower)
|
||||
N_init = 1.0 / (1.0 - self.clip_lower)
|
||||
self.alpha = 0.5 * N_init
|
||||
self.beta = 0.5 * N_init
|
||||
self.old_mean_logprob = None
|
||||
|
||||
def get_baselines(self, batch_size):
|
||||
baseline = self.alpha / (self.alpha + self.beta)
|
||||
return torch.full((batch_size,), baseline, dtype=torch.float32)
|
||||
|
||||
def compute_rho(self, cur_mean_logprob):
|
||||
if self.rho_mode == 'constant':
|
||||
return self.rho_const
|
||||
if self.old_mean_logprob is None:
|
||||
return self.rho_const
|
||||
kl = abs(self.old_mean_logprob - cur_mean_logprob)
|
||||
rho = 2 ** (-kl / self.D_half)
|
||||
return max(min(rho, self.clip_upper), self.clip_lower)
|
||||
|
||||
def update(self, rewards, cur_logprobs=None, response_masks=None):
|
||||
if cur_logprobs is not None and response_masks is not None:
|
||||
mean_logprob = ((cur_logprobs * response_masks).sum() / response_masks.sum()).item()
|
||||
rho = self.compute_rho(mean_logprob)
|
||||
self.old_mean_logprob = mean_logprob
|
||||
else:
|
||||
rho = self.rho_const
|
||||
|
||||
scale = 3.0
|
||||
normalized_rewards = (rewards + scale) / (2 * scale)
|
||||
avg_normalized_reward = normalized_rewards.mean().item()
|
||||
self.alpha = rho * self.alpha + avg_normalized_reward
|
||||
self.beta = rho * self.beta + (1 - avg_normalized_reward)
|
||||
return rho
|
||||
|
||||
|
||||
def Logger(content):
|
||||
if not ddp or dist.get_rank() == 0:
|
||||
print(content)
|
||||
|
||||
|
||||
def calculate_rewards(prompts, responses, reward_model, reward_tokenizer):
|
||||
"""整合所有奖励函数计算总奖励"""
|
||||
|
||||
def reasoning_model_reward(rewards):
|
||||
# 1. 格式奖励（仅针对训练推理模型时使用）
|
||||
pattern = r"^<think>\n.*?\n</think>\n<answer>\n.*?\n</answer>$"
|
||||
pattern2 = r"^<think>\n.*?\n</think>\n\n<answer>\n.*?\n</answer>$"
|
||||
|
||||
matches_pattern = [re.match(pattern, response, re.S) for response in responses]
|
||||
matches_pattern2 = [re.match(pattern2, response, re.S) for response in responses]
|
||||
|
||||
format_rewards = []
|
||||
for match_pattern, match_pattern2 in zip(matches_pattern, matches_pattern2):
|
||||
if match_pattern:
|
||||
format_rewards.append(0.5)
|
||||
elif match_pattern2:
|
||||
format_rewards.append(0.5)
|
||||
else:
|
||||
format_rewards.append(0.0)
|
||||
rewards += torch.tensor(format_rewards, device=args.device)
|
||||
|
||||
# 2. 标记奖励（防止严格奖励稀疏，仅针对训练推理模型时使用）
|
||||
def mark_num(text):
|
||||
reward = 0
|
||||
if text.count("<think>") == 1:
|
||||
reward += 0.25
|
||||
if text.count("</think>") == 1:
|
||||
reward += 0.25
|
||||
if text.count("<answer>") == 1:
|
||||
reward += 0.25
|
||||
if text.count("</answer>") == 1:
|
||||
reward += 0.25
|
||||
return reward
|
||||
|
||||
mark_rewards = [mark_num(response) for response in responses]
|
||||
rewards += torch.tensor(mark_rewards, device=args.device)
|
||||
return rewards
|
||||
|
||||
rewards = torch.zeros(len(responses), device=args.device)
|
||||
|
||||
# 3. 格式奖励
|
||||
if args.reasoning == 1:
|
||||
rewards = reasoning_model_reward(rewards) # 训练推理模型时使用
|
||||
|
||||
# 4. 使用reward model计算奖励
|
||||
with torch.no_grad():
|
||||
reward_model_scores = []
|
||||
scale = 3.0
|
||||
|
||||
for i, (prompt, response) in enumerate(zip(prompts, responses)):
|
||||
pattern = r"<\|im_start\|>(system|user|assistant)\s+(.*?)<\|im_end\|>"
|
||||
matches = re.findall(pattern, prompt, re.DOTALL)
|
||||
messages = [{"role": role, "content": content.strip()} for role, content in matches]
|
||||
|
||||
tmp_chat = messages + [{"role": "assistant", "content": response}]
|
||||
score = reward_model.get_score(reward_tokenizer, tmp_chat)
|
||||
score = max(min(score, scale), -scale)
|
||||
|
||||
if args.reasoning == 1:
|
||||
answer_match = re.search(r'<answer>(.*?)</answer>', response, re.DOTALL)
|
||||
if answer_match:
|
||||
answer_content = answer_match.group(1).strip()
|
||||
tmp_chat = messages + [{"role": "assistant", "content": answer_content}]
|
||||
answer_score = reward_model.get_score(reward_tokenizer, tmp_chat)
|
||||
answer_score = max(min(answer_score, scale), -scale)
|
||||
score = score * 0.4 + answer_score * 0.6
|
||||
|
||||
reward_model_scores.append(score)
|
||||
|
||||
reward_model_scores = torch.tensor(reward_model_scores, device=args.device)
|
||||
rewards += reward_model_scores
|
||||
|
||||
return rewards
|
||||
|
||||
|
||||
def spo_train_epoch(epoch, wandb, value_tracker):
|
||||
for step, batch in enumerate(train_loader):
|
||||
prompts = batch['prompt'] # list[str], length B
|
||||
prompt_inputs = tokenizer(prompts, return_tensors="pt", padding=True, return_token_type_ids=False,
|
||||
padding_side="left", add_special_tokens=False).to(args.device) # input_ids: [B, P], attention_mask: [B, P]
|
||||
if args.max_seq_len:
|
||||
prompt_inputs["input_ids"] = prompt_inputs["input_ids"][:, -args.max_seq_len:]
|
||||
prompt_inputs["attention_mask"] = prompt_inputs["attention_mask"][:, -args.max_seq_len:]
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = (model.module if ddp else model).generate(
|
||||
**prompt_inputs, max_new_tokens=args.max_gen_len, do_sample=True, temperature=0.8,
|
||||
num_return_sequences=1, pad_token_id=tokenizer.pad_token_id) # [B, P+R]
|
||||
|
||||
completion_ids = outputs[:, prompt_inputs["input_ids"].size(1):] # [B, R]
|
||||
|
||||
def get_per_token_logps(mdl, input_ids, n_keep):
|
||||
input_ids = input_ids.detach().clone() if input_ids.is_inference() else input_ids
|
||||
logits = mdl(input_ids, logits_to_keep=n_keep + 1).logits[:, :-1, :]
|
||||
per_token_logps = []
|
||||
for logits_row, ids_row in zip(logits, input_ids[:, -n_keep:]):
|
||||
ids_row = ids_row.detach().clone() if ids_row.is_inference() else ids_row
|
||||
per_token_logps.append(torch.gather(logits_row.log_softmax(dim=-1), 1, ids_row.unsqueeze(1)).squeeze(1))
|
||||
return torch.stack(per_token_logps)
|
||||
|
||||
per_token_logps = get_per_token_logps(model, outputs, completion_ids.size(1)) # [B, R]
|
||||
with torch.no_grad():
|
||||
ref_per_token_logps = get_per_token_logps(ref_model, outputs, completion_ids.size(1)) # [B, R]
|
||||
|
||||
completions = tokenizer.batch_decode(completion_ids, skip_special_tokens=True) # list[str], length B
|
||||
rewards = calculate_rewards(prompts, completions, reward_model, reward_tokenizer).to(args.device) # [B]
|
||||
|
||||
baselines = value_tracker.get_baselines(len(prompts)).to(args.device) # [B]
|
||||
|
||||
scale = 3.0
|
||||
# Un-normalize baselines to be in the same scale as raw rewards [-3, 3]
|
||||
unnormalized_baselines = baselines * (2 * scale) - scale # [B]
|
||||
advantages = rewards - unnormalized_baselines # [B]
|
||||
|
||||
# 直接使用 baseline 提供的优势估计，只做裁剪防止梯度爆炸。不再做 batch 内归一化，因为 baseline 已经提供了跨 batch 的稳定基线
|
||||
advantages = advantages.clamp(-5.0, 5.0)
|
||||
|
||||
is_eos = completion_ids == tokenizer.eos_token_id # [B, R]
|
||||
eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=args.device) # [B]
|
||||
eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)]
|
||||
completion_mask = (torch.arange(is_eos.size(1), device=args.device).expand(is_eos.size(0), -1) <= eos_idx.unsqueeze(1)).int() # [B, R]
|
||||
|
||||
kl_div = ref_per_token_logps - per_token_logps # [B, R]
|
||||
per_token_kl = torch.exp(kl_div) - kl_div - 1 # [B, R]
|
||||
per_token_loss = -per_token_logps * advantages.unsqueeze(1) + args.beta * per_token_kl # [B, R]
|
||||
loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() / args.accumulation_steps # scalar
|
||||
loss.backward()
|
||||
|
||||
response_masks = completion_mask.float() # [B, R]
|
||||
rho = value_tracker.update(rewards, per_token_logps.detach(), response_masks)
|
||||
|
||||
if (step + 1) % args.accumulation_steps == 0:
|
||||
if args.grad_clip > 0:
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||||
optimizer.step()
|
||||
scheduler.step()
|
||||
optimizer.zero_grad()
|
||||
|
||||
if step % args.log_interval == 0 or step == iter_per_epoch - 1:
|
||||
policy_loss_val = loss.item()
|
||||
avg_reward_val = rewards.mean().item()
|
||||
avg_len_val = completion_mask.sum(dim=1).float().mean().item()
|
||||
# average kl over valid tokens for logging
|
||||
kl_val = ((per_token_kl * completion_mask).sum() / (completion_mask.sum() + 1e-8)).item()
|
||||
avg_baseline_val = baselines.mean().item()
|
||||
current_lr = optimizer.param_groups[0]['lr']
|
||||
|
||||
Logger(
|
||||
f'Epoch: {epoch}, Step: {step + 1}/{iter_per_epoch}, '
|
||||
f'Actor Loss: {policy_loss_val:.4f}, Reward: {avg_reward_val:.4f}, '
|
||||
f'Baseline: {avg_baseline_val:.4f}, KL: {kl_val:.4f}, Rho: {rho:.4f}, Avg Response Len: {avg_len_val:.2f}, LR: {current_lr:.2e}')
|
||||
|
||||
if wandb and (not ddp or dist.get_rank() == 0):
|
||||
log_dict = {
|
||||
"policy_loss": policy_loss_val,
|
||||
"reward": avg_reward_val,
|
||||
"kl": kl_val,
|
||||
"rho": float(rho),
|
||||
"baseline": avg_baseline_val,
|
||||
# "avg_response_len": avg_len_val,
|
||||
"advantages_mean": advantages.mean().item(),
|
||||
"learning_rate": current_lr
|
||||
}
|
||||
wandb.log(log_dict)
|
||||
|
||||
if ((step + 1) % args.save_interval == 0 or step == iter_per_epoch - 1) and (not ddp or dist.get_rank() == 0):
|
||||
model.eval()
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
suffix = 'spo'
|
||||
ckp = f'{args.save_dir}/{suffix}_{lm_config.hidden_size}{moe_path}.pth'
|
||||
|
||||
state_dict = model.module.state_dict() if isinstance(model,
|
||||
torch.nn.parallel.DistributedDataParallel) else model.state_dict()
|
||||
torch.save({k: v.half() for k, v in state_dict.items()}, ckp)
|
||||
model.train()
|
||||
|
||||
del prompt_inputs, outputs, completion_ids, per_token_logps, ref_per_token_logps
|
||||
del completions, rewards, advantages, completion_mask, baselines, response_masks
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
|
||||
def init_model(lm_config):
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model/')
|
||||
model = MiniMindForCausalLM(lm_config)
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'{args.save_dir}/full_sft_{lm_config.hidden_size}{moe_path}.pth'
|
||||
if args.reasoning == 1:
|
||||
ckp = f'{args.save_dir}/reason_{lm_config.hidden_size}{moe_path}.pth'
|
||||
state_dict = torch.load(ckp, map_location=args.device)
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
|
||||
ref_model = MiniMindForCausalLM(lm_config)
|
||||
ref_model.load_state_dict(state_dict, strict=False)
|
||||
ref_model.eval().requires_grad_(False)
|
||||
|
||||
Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
||||
model = model.to(args.device)
|
||||
ref_model = ref_model.to(args.device)
|
||||
|
||||
reward_name = "../../internlm2-1_8b-reward"
|
||||
reward_model = AutoModel.from_pretrained(
|
||||
reward_name,
|
||||
device_map="cuda",
|
||||
torch_dtype=torch.float16,
|
||||
trust_remote_code=True,
|
||||
).to(args.device).eval().requires_grad_(False)
|
||||
reward_tokenizer = AutoTokenizer.from_pretrained(reward_name, trust_remote_code=True)
|
||||
|
||||
return model, ref_model, tokenizer, reward_model, reward_tokenizer
|
||||
|
||||
|
||||
def init_distributed_mode():
|
||||
if not ddp: return
|
||||
global ddp_local_rank, DEVICE
|
||||
dist.init_process_group(backend="nccl")
|
||||
ddp_local_rank = int(os.environ["LOCAL_RANK"])
|
||||
DEVICE = f"cuda:{ddp_local_rank}"
|
||||
torch.cuda.set_device(DEVICE)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--out_dir", type=str, default="../out")
|
||||
parser.add_argument("--epochs", type=int, default=1)
|
||||
parser.add_argument("--batch_size", type=int, default=2)
|
||||
parser.add_argument("--learning_rate", type=float, default=1e-7)
|
||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16")
|
||||
parser.add_argument("--use_wandb", action="store_true")
|
||||
parser.add_argument("--wandb_project", type=str, default="MiniMind-SPO")
|
||||
parser.add_argument("--num_workers", type=int, default=1)
|
||||
parser.add_argument("--ddp", action="store_true")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=4)
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0)
|
||||
parser.add_argument("--log_interval", type=int, default=1)
|
||||
parser.add_argument("--save_interval", type=int, default=10)
|
||||
parser.add_argument('--hidden_size', default=512, type=int)
|
||||
parser.add_argument('--num_hidden_layers', default=8, type=int)
|
||||
parser.add_argument('--use_moe', default=False, type=bool)
|
||||
parser.add_argument('--max_seq_len', default=66, type=int)
|
||||
parser.add_argument("--max_gen_len", type=int, default=1536)
|
||||
parser.add_argument("--data_path", type=str, default="../dataset/rlaif-mini.jsonl")
|
||||
parser.add_argument("--beta", type=float, default=0.02)
|
||||
parser.add_argument("--reasoning", type=int, default=1, help='0:普通模型，1:推理模型')
|
||||
args = parser.parse_args()
|
||||
|
||||
lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
|
||||
max_seq_len=args.max_seq_len + args.max_gen_len,
|
||||
use_moe=args.use_moe)
|
||||
args.save_dir = os.path.join(args.out_dir)
|
||||
os.makedirs(args.save_dir, exist_ok=True)
|
||||
os.makedirs(args.out_dir, exist_ok=True)
|
||||
|
||||
ctx = nullcontext() if "cuda" not in args.device else torch.amp.autocast('cuda')
|
||||
ddp = int(os.environ.get("RANK", -1)) != -1
|
||||
ddp_local_rank, DEVICE = 0, "cuda:0"
|
||||
|
||||
base_seed = 1337
|
||||
torch.manual_seed(base_seed)
|
||||
torch.cuda.manual_seed(base_seed)
|
||||
|
||||
if ddp:
|
||||
init_distributed_mode()
|
||||
args.device = torch.device(DEVICE)
|
||||
rank = dist.get_rank()
|
||||
torch.manual_seed(base_seed + rank)
|
||||
# 同时设置 CUDA 的随机种子
|
||||
torch.cuda.manual_seed(base_seed + rank)
|
||||
|
||||
if args.use_wandb and (not ddp or ddp_local_rank == 0):
|
||||
import swanlab as wandb
|
||||
|
||||
wandb.init(project=args.wandb_project)
|
||||
else:
|
||||
wandb = None
|
||||
|
||||
model, ref_model, tokenizer, reward_model, reward_tokenizer = init_model(lm_config=lm_config)
|
||||
train_ds = RLAIFDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
|
||||
train_sampler = DistributedSampler(train_ds) if ddp else None
|
||||
train_loader = DataLoader(train_ds, batch_size=args.batch_size, pin_memory=True,
|
||||
drop_last=False, shuffle=False,
|
||||
num_workers=args.num_workers, sampler=train_sampler)
|
||||
|
||||
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
|
||||
|
||||
iter_per_epoch = len(train_loader)
|
||||
total_optimizer_steps = (iter_per_epoch // args.accumulation_steps) * args.epochs
|
||||
scheduler = CosineAnnealingLR(optimizer, T_max=total_optimizer_steps, eta_min=args.learning_rate / 10)
|
||||
|
||||
if ddp:
|
||||
model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
|
||||
model = DistributedDataParallel(model, device_ids=[ddp_local_rank])
|
||||
|
||||
value_tracker = AutoAdaptiveValueTracker(rho_mode='kl', rho_const=0.9, D_half=0.06, clip_lower=0.5, clip_upper=0.96)
|
||||
|
||||
for epoch in range(args.epochs):
|
||||
spo_train_epoch(epoch, wandb, value_tracker)
|
||||
Loading…
Reference in New Issue
Block a user