mirror of
https://github.com/jingyaogong/minimind.git
synced 2026-08-03 12:17:47 +00:00
250426
This commit is contained in:
jingyaogong
@@ -93,9 +93,9 @@
|
||||
|
||||
| 模型 (大小) | 推理占用 (约) | Release |
|
||||
|-------------------------|----------|------------|
|
||||
| MiniMind2-small (26M) | 0.5 GB | 2025.02.06 |
|
||||
| MiniMind2-MoE (145M) | 1.0 GB | 2025.02.06 |
|
||||
| MiniMind2 (104M) | 1.0 GB | 2025.02.06 |
|
||||
| MiniMind2-small (26M) | 0.5 GB | 2025.04.26 |
|
||||
| MiniMind2-MoE (145M) | 1.0 GB | 2025.04.26 |
|
||||
| MiniMind2 (104M) | 1.0 GB | 2025.04.26 |
|
||||
| minimind-v1-small (26M) | 0.5 GB | 2024.08.28 |
|
||||
| minimind-v1-moe (4×26M) | 1.0 GB | 2024.09.17 |
|
||||
| minimind-v1 (108M) | 1.0 GB | 2024.09.01 |
|
||||
@@ -114,6 +114,7 @@
|
||||
- 在第三方测评榜（C-Eval、C-MMLU、OpenBookQA等）进行模型测试。
|
||||
- 实现Openai-Api协议的极简服务端，便于集成到第三方ChatUI使用（FastGPT、Open-WebUI等）。
|
||||
- 基于streamlit实现最简聊天WebUI前端。
|
||||
- 全面兼容社区热门`llama.cpp`、`vllm`、`ollama`推理引擎或`Llama-Factory`训练框架。
|
||||
- 复现(蒸馏/RL)大型推理模型DeepSeek-R1的MiniMind-Reason模型，**数据+模型**全部开源！
|
||||
|
||||
希望此开源项目可以帮助LLM初学者快速入门！
|
||||
@@ -121,7 +122,27 @@
|
||||
### 👉**更新日志**
|
||||
|
||||
<details close>
|
||||
<summary> <b>2025-02-09 (newest 🎉🎉🎉)</b> </summary>
|
||||
<summary> <b>2025-04-26 (newest 🎉🎉🎉)</b> </summary>
|
||||
|
||||
- 重要更新
|
||||
- 如有兼容性需要，可访问[🔗旧仓库内容🔗](https://github.com/jingyaogong/minimind/tree/7da201a944a90ed49daef8a0265c959288dff83a)。
|
||||
- MiniMind模型参数完全改名，对齐Transformers库模型（统一命名）。
|
||||
- generate方式重构，继承自GenerationMixin类。
|
||||
- 🔥支持llama.cpp、vllm、ollama等热门三方生态。
|
||||
- 规范代码和目录结构。
|
||||
- 🔥更新：从0实现PPO、GRPO的训练代码。
|
||||
- 改动词表`<s></s>`->`<|im_start|><|im_end|>`
|
||||
```text
|
||||
为兼容第三方推理框架llama.cpp、vllm，本次更新需付出一些可观代价。
|
||||
本次更新不再支持「直接」加载25-04-26以前的旧模型进行推理。
|
||||
由于Llama位置编码方式与minimind存在区别，导致映射Llama模型后QK值存在差异
|
||||
MiniMind2系列旧模型均经过权重映射+（微调训练）QKVO线性层校准恢复而来。
|
||||
本次更新后将放弃对`minimind-v1`全系列的维护，并在仓库中下线。
|
||||
```
|
||||
</details>
|
||||
|
||||
<details close>
|
||||
<summary> <b>2025-02-09</b> </summary>
|
||||
|
||||
- 迎来发布以来重大更新，Release MiniMind2 Series。
|
||||
- 代码几乎全部重构，使用更简洁明了的统一结构。
|
||||
@@ -216,19 +237,19 @@ pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
|
||||
```
|
||||
|
||||
### 2.下载模型
|
||||
|
||||
到项目根目录
|
||||
```bash
|
||||
git clone https://huggingface.co/jingyaogong/MiniMind2
|
||||
```
|
||||
|
||||
### 3.命令行问答
|
||||
### （可选）命令行问答
|
||||
|
||||
```bash
|
||||
# load=0: load from pytorch model, load=1: load from transformers-hf model
|
||||
python eval_model.py --load 1 --model_mode 2
|
||||
```
|
||||
|
||||
### 4.或启动WebUI
|
||||
### （可选）启动WebUI
|
||||
|
||||
```bash
|
||||
# 可能需要`python>=3.10` 安装 `pip install streamlit`
|
||||
@@ -236,6 +257,15 @@ python eval_model.py --load 1 --model_mode 2
|
||||
streamlit run web_demo.py
|
||||
```
|
||||
|
||||
### （可选）第三方推理框架
|
||||
|
||||
```bash
|
||||
# ollama
|
||||
ollama run jingyaogong/minimind2
|
||||
# vllm
|
||||
vllm serve ./MiniMind2/ --served-model-name "minimind"
|
||||
```
|
||||
|
||||
## Ⅱ 从0开始自己训练
|
||||
|
||||
### 1.环境准备
|
||||
@@ -273,6 +303,8 @@ print(torch.cuda.is_available())
|
||||
|
||||
### 3.开始训练
|
||||
|
||||
目录位于`trainer`
|
||||
|
||||
**3.1 预训练（学知识）**
|
||||
|
||||
```bash
|
||||
@@ -640,6 +672,8 @@ Zero模型权重保存为 `full_sft_512_zero.pth`（见下文MiniMind模型文
|
||||
|
||||
## Ⅱ 主要训练步骤
|
||||
|
||||
> 所有训练脚本均 `cd ./trainer` 目录执行
|
||||
|
||||
### **1. 预训练(Pretrain)**:
|
||||
|
||||
LLM首先要学习的并非直接与人交流，而是让网络参数中充满知识的墨水，“墨水” 理论上喝的越饱越好，产生大量的对世界的知识积累。
|
||||
@@ -677,6 +711,8 @@ python train_full_sft.py
|
||||
|
||||
## Ⅲ 其它训练步骤
|
||||
|
||||
> 所有训练脚本均 `cd ./trainer` 目录执行
|
||||
|
||||
### **3. 人类反馈强化学习(Reinforcement Learning from Human Feedback, RLHF)**
|
||||
|
||||
在前面的训练步骤中，模型已经具备了基本的对话能力，但是这样的能力完全基于单词接龙，缺少正反样例的激励。
|
||||
@@ -1214,16 +1250,13 @@ MiniMind模型本身预训练数据集小的可怜，也没有针对性的对测
|
||||
|
||||
# 📌 其它 (Others)
|
||||
|
||||
### 推理与导出
|
||||
## 模型转换
|
||||
|
||||
* [./scripts/convert_model.py](./scripts/convert_model.py)可以将torch/transformers模型互相转换。
|
||||
|
||||
* MiniMind的HuggingFace集合地址：
|
||||
[MiniMind](https://huggingface.co/collections/jingyaogong/minimind-66caf8d999f5c7fa64f399e5)
|
||||
* [./scripts/convert_model.py](./scripts/convert_model.py)可以实现`torch模型/transformers`模型之间的转换
|
||||
|
||||
---
|
||||
|
||||
### 基于MiniMind-API服务接口
|
||||
## 基于MiniMind-API服务接口
|
||||
|
||||
* [./scripts/serve_openai_api.py](./scripts/serve_openai_api.py)完成了兼容openai-api的最简聊天接口，方便将自己的模型接入第三方UI
|
||||
例如FastGPT、OpenWebUI、Dify等等。
|
||||
@@ -1265,6 +1298,74 @@ MiniMind模型本身预训练数据集小的可怜，也没有针对性的对测
|
||||
}'
|
||||
```
|
||||
|
||||
## VLLM模型推理（服务）
|
||||
|
||||
[vLLM](https://github.com/vllm-project/vllm)是极其流行的高效推理框架，支持大模型快速部署，优化显存利用与吞吐量。
|
||||
|
||||
```bash
|
||||
vllm serve ./MiniMind2/ --model-impl transformers --served-model-name "minimind"
|
||||
```
|
||||
|
||||
服务将以openai api协议启动，端口默认为8000。
|
||||
|
||||
更多用法请参考官方说明～
|
||||
|
||||
## llama.cpp
|
||||
[llama.cpp](https://github.com/ggerganov/llama.cpp)是一个C++库，
|
||||
可以在命令行下直接使用，支持多线程推理，支持GPU加速。
|
||||
|
||||
参考官方仓库安装后，在`convert_hf_to_gguf.py` ～760行插入
|
||||
```text
|
||||
# 添加MiniMind2 tokenizer支持
|
||||
if res is None:
|
||||
res = "smollm"
|
||||
```
|
||||
|
||||
转换自定义训练的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
|
||||
```
|
||||
|
||||
更多用法请参考官方说明～
|
||||
|
||||
## ollama
|
||||
|
||||
[ollama](https://ollama.ai/)是本地运行大模型的工具，支持多种开源LLM，简单易用。
|
||||
|
||||
通过ollama加载自定义的gguf模型，新建minimind.modelfile：
|
||||
```text
|
||||
FROM ./MiniMind2-109M-F16.gguf
|
||||
TEMPLATE """{{ if .System }}<|im_start|>system
|
||||
{{ .System }}<|im_end|>
|
||||
{{ end }}{{ if .Prompt }}<|im_start|>user
|
||||
{{ .Prompt }}<|im_end|>
|
||||
{{ end }}<|im_start|>assistant
|
||||
"""
|
||||
```
|
||||
|
||||
加载模型并命名为`minimind2`
|
||||
```bash
|
||||
ollama create -f minimind.modelfile minimind2
|
||||
```
|
||||
|
||||
启动推理
|
||||
```text
|
||||
ollama run minimind2
|
||||
> 你好，我是MiniMind2，一个基于xxxxxxxx
|
||||
```
|
||||
|
||||
更多用法请参考官方说明～
|
||||
|
||||
# 📌 Acknowledge
|
||||
|
||||
> [!NOTE]
|
||||
|
||||
+123
-13
@@ -100,9 +100,9 @@ the entire process of building a language model from 0 to 1. Let's enjoy the fun
|
||||
|
||||
| Model (Size) | Inference Usage (Approx.) | Release |
|
||||
|-------------------------|---------------------------|------------|
|
||||
| MiniMind2-small (26M) | 0.5 GB | 2025.02.06 |
|
||||
| MiniMind2-MoE (145M) | 1.0 GB | 2025.02.06 |
|
||||
| MiniMind2 (104M) | 1.0 GB | 2025.02.06 |
|
||||
| MiniMind2-small (26M) | 0.5 GB | 2025.04.26 |
|
||||
| MiniMind2-MoE (145M) | 1.0 GB | 2025.04.26 |
|
||||
| MiniMind2 (104M) | 1.0 GB | 2025.04.26 |
|
||||
| minimind-v1-small (26M) | 0.5 GB | 2024.08.28 |
|
||||
| minimind-v1-moe (4×26M) | 1.0 GB | 2024.09.17 |
|
||||
| minimind-v1 (108M) | 1.0 GB | 2024.09.01 |
|
||||
@@ -123,6 +123,7 @@ the entire process of building a language model from 0 to 1. Let's enjoy the fun
|
||||
- Model testing on third-party evaluation benchmarks (C-Eval, C-MMLU, OpenBookQA, etc.).
|
||||
- A minimal server implementing the Openai-Api protocol, easy to integrate into third-party ChatUI applications (
|
||||
FastGPT, Open-WebUI, etc.).
|
||||
- Fully compatible with popular community inference engines like llama.cpp, vllm, ollama, or training frameworks such as Llama-Factory.
|
||||
- A simple chat WebUI front-end implemented using streamlit.
|
||||
- Reproduction (distillation/RL) of the large inference model DeepSeek-R1 as the MiniMind-Reason model, **data + model**
|
||||
all open-source!
|
||||
@@ -131,8 +132,38 @@ We hope this open-source project can help LLM beginners quickly get started!
|
||||
|
||||
### 👉**Update log**
|
||||
|
||||
<details close>
|
||||
<summary> <b>2025-04-26 (newest 🎉🎉🎉)</b> </summary>
|
||||
|
||||
• Major Updates
|
||||
|
||||
• For compatibility needs, visit [🔗Legacy Repository Content🔗](https://github.com/jingyaogong/minimind/tree/7da201a944a90ed49daef8a0265c959288dff83a).
|
||||
|
||||
• MiniMind model parameters have been fully renamed to align with Transformers library models (unified naming).
|
||||
|
||||
• The `generate` method has been refactored, now inheriting from the `GenerationMixin` class.
|
||||
|
||||
• 🔥 Support for popular third-party ecosystems like llama.cpp, vllm, and ollama.
|
||||
|
||||
• Standardized code and directory structure.
|
||||
|
||||
• 🔥 New: Added training code for PPO and GRPO from scratch.
|
||||
|
||||
• Updated vocabulary tokens: `<s></s>` → `<|im_start|><|im_end|>`.
|
||||
|
||||
|
||||
```text
|
||||
To ensure compatibility with third-party inference frameworks (llama.cpp, vllm), this update comes at a non-trivial cost.
|
||||
Models saved before 2025-04-26 can no longer be **directly** loaded for inference.
|
||||
Due to differences in positional encoding between Llama and MiniMind, QK values diverge after weight mapping.
|
||||
MiniMind2 legacy models have been restored via weight mapping + (fine-tuning) QKVO linear layer calibration.
|
||||
After this update, maintenance for the entire `minimind-v1` series will be discontinued, and the models will be removed from the repository.
|
||||
```
|
||||
</details>
|
||||
|
||||
|
||||
<details close>
|
||||
<summary> <b>2025-02-09 (newest 🎉🎉🎉)</b> </summary>
|
||||
<summary> <b>2025-02-09</b> </summary>
|
||||
|
||||
- Major update since the release, with the release of MiniMind2 Series.
|
||||
- Almost all code has been refactored, using a more streamlined and unified structure.
|
||||
@@ -235,21 +266,30 @@ pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
|
||||
git clone https://huggingface.co/jingyaogong/MiniMind2
|
||||
```
|
||||
|
||||
### 3. Command-line Q&A
|
||||
### (Optional) Command-line Q&A
|
||||
|
||||
```bash
|
||||
# load=0: load from pytorch model, load=1: load from transformers-hf model
|
||||
python eval_model.py --load 1 --model_mode 2
|
||||
```
|
||||
|
||||
### 4. Or Start WebUI
|
||||
### (Optional) Launch WebUI
|
||||
|
||||
```bash
|
||||
# You may need `python>=3.10` and install `pip install streamlit`.
|
||||
# May require `python>=3.10`, install with `pip install streamlit`
|
||||
# cd scripts
|
||||
streamlit run web_demo.py
|
||||
```
|
||||
|
||||
### (Optional) Third-party inference frameworks
|
||||
|
||||
```bash
|
||||
# ollama
|
||||
ollama run jingyaogong/minimind2
|
||||
# vllm
|
||||
vllm serve ./MiniMind2/ --served-model-name "minimind"
|
||||
```
|
||||
|
||||
## Ⅱ Training from Scratch
|
||||
|
||||
### 1. Environment Setup
|
||||
@@ -292,6 +332,8 @@ needs and GPU resources.
|
||||
|
||||
### 3. Start Training
|
||||
|
||||
The directory is located at `trainer`
|
||||
|
||||
**3.1 Pretraining (Learning Knowledge)**
|
||||
|
||||
```bash
|
||||
@@ -696,6 +738,8 @@ download and test the model's performance.
|
||||
|
||||
## Ⅱ Main Training Steps
|
||||
|
||||
> All training scripts are executed in the `cd ./trainer` directory.
|
||||
|
||||
### **1. Pretraining**:
|
||||
|
||||
The first task for LLM is not to interact directly with humans, but to fill the network parameters with knowledge. The "
|
||||
@@ -744,6 +788,8 @@ python train_full_sft.py
|
||||
|
||||
## Ⅲ Other Training Steps
|
||||
|
||||
> All training scripts are executed in the `cd ./trainer` directory.
|
||||
|
||||
### **3. Reinforcement Learning from Human Feedback (RLHF)**
|
||||
|
||||
In the previous training steps, the model has acquired basic conversational abilities, but these are entirely based on
|
||||
@@ -1361,16 +1407,13 @@ is mainly for fun, so take the results lightly:
|
||||
|
||||
# 📌 Others
|
||||
|
||||
### Inference and Export
|
||||
## Model Conversion
|
||||
|
||||
* [./scripts/convert_model.py](./scripts/convert_model.py) can convert models between torch/transformers.
|
||||
|
||||
* MiniMind's HuggingFace collection link:
|
||||
[MiniMind](https://huggingface.co/collections/jingyaogong/minimind-66caf8d999f5c7fa64f399e5)
|
||||
* [./scripts/convert_model.py](./scripts/convert_model.py) can be used to convert between `torch models` and `transformers` models.
|
||||
|
||||
---
|
||||
|
||||
### Based on MiniMind-API Service Interface
|
||||
## Based on MiniMind-API Service Interface
|
||||
|
||||
* [./scripts/serve_openai_api.py](./scripts/serve_openai_api.py) provides the simplest chat interface compatible with
|
||||
the OpenAI API,
|
||||
@@ -1415,6 +1458,73 @@ is mainly for fun, so take the results lightly:
|
||||
}'
|
||||
```
|
||||
|
||||
## VLLM Model Inference (Service)
|
||||
|
||||
[vLLM](https://github.com/vllm-project/vllm) is an extremely popular and efficient inference framework that supports fast deployment of large models, optimizing memory utilization and throughput.
|
||||
|
||||
```bash
|
||||
vllm serve ./MiniMind2/ --model-impl transformers --served-model-name "minimind"
|
||||
```
|
||||
|
||||
The service will start using the OpenAI API protocol, with the default port being 8000.
|
||||
|
||||
For more usage, please refer to the official documentation.
|
||||
|
||||
## llama.cpp
|
||||
[llama.cpp](https://github.com/ggerganov/llama.cpp) is a C++ library that can be used directly from the command line, supporting multi-threaded inference and GPU acceleration.
|
||||
|
||||
After installation (refer to the official repository), insert the following code at line 760 of `convert_hf_to_gguf.py`:
|
||||
```text
|
||||
# Add MiniMind2 tokenizer support
|
||||
if res is None:
|
||||
res = "smollm"
|
||||
```
|
||||
|
||||
Convert a custom-trained MiniMind model to gguf:
|
||||
```bash
|
||||
python convert_hf_to_gguf.py ../minimind/MiniMind2/
|
||||
```
|
||||
|
||||
Quantize the model:
|
||||
```bash
|
||||
./build/bin/llama-quantize ../minimind/MiniMind2/MiniMind2-109M-F16.gguf ../minimind/MiniMind2/Q4-MiniMind2.gguf Q4_K_M
|
||||
```
|
||||
|
||||
Command line inference:
|
||||
```bash
|
||||
./build/bin/llama-cli -m ../minimind/MiniMind2/MiniMind2-109M-F16.gguf --chat-template chatml
|
||||
```
|
||||
|
||||
For more usage, please refer to the official documentation.
|
||||
|
||||
## ollama
|
||||
|
||||
[ollama](https://ollama.ai/) is a tool for running large models locally, supporting multiple open-source LLMs, and is easy to use.
|
||||
|
||||
To load a custom gguf model with ollama, create a new file `minimind.modelfile`:
|
||||
```text
|
||||
FROM ./MiniMind2-109M-F16.gguf
|
||||
TEMPLATE """{{ if .System }}<|im_start|>system
|
||||
{{ .System }}<|im_end|>
|
||||
{{ end }}{{ if .Prompt }}<|im_start|>user
|
||||
{{ .Prompt }}<|im_end|>
|
||||
{{ end }}<|im_start|>assistant
|
||||
"""
|
||||
```
|
||||
|
||||
Load the model and name it `minimind2`:
|
||||
```bash
|
||||
ollama create -f minimind.modelfile minimind2
|
||||
```
|
||||
|
||||
Start inference:
|
||||
```text
|
||||
ollama run minimind2
|
||||
> Hello, I am MiniMind2, based on xxxxxxxx
|
||||
```
|
||||
|
||||
For more usage, please refer to the official documentation.
|
||||
|
||||
# 📌 Acknowledge
|
||||
|
||||
> [!NOTE]
|
||||
|
||||
+41
-52
@@ -1,37 +1,32 @@
|
||||
import argparse
|
||||
import random
|
||||
import time
|
||||
import numpy as np
|
||||
import torch
|
||||
import warnings
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
from model.model import MiniMindLM
|
||||
from model.LMConfig import LMConfig
|
||||
import numpy as np
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from model.model_lora import *
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
||||
def init_model(args):
|
||||
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
|
||||
tokenizer = AutoTokenizer.from_pretrained('./model/')
|
||||
if args.load == 0:
|
||||
moe_path = '_moe' if args.use_moe else ''
|
||||
modes = {0: 'pretrain', 1: 'full_sft', 2: 'rlhf', 3: 'reason', 4: 'grpo'}
|
||||
ckp = f'./{args.out_dir}/{modes[args.model_mode]}_{args.dim}{moe_path}.pth'
|
||||
ckp = f'./{args.out_dir}/{modes[args.model_mode]}_{args.hidden_size}{moe_path}.pth'
|
||||
|
||||
model = MiniMindLM(LMConfig(
|
||||
dim=args.dim,
|
||||
n_layers=args.n_layers,
|
||||
max_seq_len=args.max_seq_len,
|
||||
model = MiniMindForCausalLM(MiniMindConfig(
|
||||
hidden_size=args.hidden_size,
|
||||
num_hidden_layers=args.num_hidden_layers,
|
||||
use_moe=args.use_moe
|
||||
))
|
||||
|
||||
state_dict = torch.load(ckp, map_location=args.device)
|
||||
model.load_state_dict({k: v for k, v in state_dict.items() if 'mask' not in k}, strict=True)
|
||||
model.load_state_dict(torch.load(ckp, map_location=args.device), strict=True)
|
||||
|
||||
if args.lora_name != 'None':
|
||||
apply_lora(model)
|
||||
load_lora(model, f'./{args.out_dir}/lora/{args.lora_name}_{args.dim}.pth')
|
||||
load_lora(model, f'./{args.out_dir}/lora/{args.lora_name}_{args.hidden_size}.pth')
|
||||
else:
|
||||
transformers_model_path = './MiniMind2'
|
||||
tokenizer = AutoTokenizer.from_pretrained(transformers_model_path)
|
||||
@@ -108,19 +103,18 @@ def main():
|
||||
parser.add_argument('--temperature', default=0.85, type=float)
|
||||
parser.add_argument('--top_p', default=0.85, type=float)
|
||||
parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str)
|
||||
# 此处max_seq_len（最大允许输入长度）并不意味模型具有对应的长文本的性能，仅防止QA出现被截断的问题
|
||||
# MiniMind2-moe (145M)：(dim=640, n_layers=8, use_moe=True)
|
||||
# MiniMind2-Small (26M)：(dim=512, n_layers=8)
|
||||
# MiniMind2 (104M)：(dim=768, n_layers=16)
|
||||
parser.add_argument('--dim', default=512, type=int)
|
||||
parser.add_argument('--n_layers', default=8, type=int)
|
||||
# 此处max_seq_len（最大输出长度）并不意味模型具有对应的长文本的性能，仅防止QA出现被截断的问题
|
||||
# MiniMind2-moe (145M)：(hidden_size=640, num_hidden_layers=8, use_moe=True)
|
||||
# MiniMind2-Small (26M)：(hidden_size=512, num_hidden_layers=8)
|
||||
# MiniMind2 (104M)：(hidden_size=768, num_hidden_layers=16)
|
||||
parser.add_argument('--hidden_size', default=640, type=int)
|
||||
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('--use_moe', default=True, type=bool)
|
||||
# 携带历史对话上下文条数
|
||||
# history_cnt需要设为偶数，即【用户问题, 模型回答】为1组；设置为0时，即当前query不携带历史上文
|
||||
# 模型未经过外推微调时，在更长的上下文的chat_template时难免出现性能的明显退化，因此需要注意此处设置
|
||||
parser.add_argument('--history_cnt', default=0, type=int)
|
||||
parser.add_argument('--stream', default=True, type=bool)
|
||||
parser.add_argument('--load', default=0, type=int, help="0: 原生torch权重，1: transformers加载")
|
||||
parser.add_argument('--model_mode', default=1, type=int,
|
||||
help="0: 预训练模型，1: SFT-Chat模型，2: RLHF-Chat模型，3: Reason模型，4: RLAIF-Chat模型")
|
||||
@@ -130,6 +124,8 @@ def main():
|
||||
|
||||
prompts = get_prompt_datas(args)
|
||||
test_mode = int(input('[0] 自动测试\n[1] 手动输入\n'))
|
||||
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
|
||||
|
||||
messages = []
|
||||
for idx, prompt in enumerate(prompts if test_mode == 0 else iter(lambda: input('👶: '), '')):
|
||||
setup_seed(random.randint(0, 2048))
|
||||
@@ -143,38 +139,31 @@ def main():
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True
|
||||
)[-args.max_seq_len - 1:] if args.model_mode != 0 else (tokenizer.bos_token + prompt)
|
||||
) if args.model_mode != 0 else (tokenizer.bos_token + prompt)
|
||||
|
||||
answer = new_prompt
|
||||
with torch.no_grad():
|
||||
x = torch.tensor(tokenizer(new_prompt)['input_ids'], device=args.device).unsqueeze(0)
|
||||
outputs = model.generate(
|
||||
x,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
max_new_tokens=args.max_seq_len,
|
||||
temperature=args.temperature,
|
||||
top_p=args.top_p,
|
||||
stream=args.stream,
|
||||
pad_token_id=tokenizer.pad_token_id
|
||||
)
|
||||
inputs = tokenizer(
|
||||
new_prompt,
|
||||
return_tensors="pt",
|
||||
truncation=True
|
||||
).to(args.device)
|
||||
|
||||
print('🤖️: ', end='')
|
||||
try:
|
||||
if not args.stream:
|
||||
print(tokenizer.decode(outputs.squeeze()[x.shape[1]:].tolist(), skip_special_tokens=True), end='')
|
||||
else:
|
||||
history_idx = 0
|
||||
for y in outputs:
|
||||
answer = tokenizer.decode(y[0].tolist(), skip_special_tokens=True)
|
||||
if (answer and answer[-1] == '�') or not answer:
|
||||
continue
|
||||
print(answer[history_idx:], end='', flush=True)
|
||||
history_idx = len(answer)
|
||||
except StopIteration:
|
||||
print("No answer")
|
||||
print('\n')
|
||||
print('🤖️: ', end='')
|
||||
generated_ids = model.generate(
|
||||
inputs["input_ids"],
|
||||
max_new_tokens=args.max_seq_len,
|
||||
num_return_sequences=1,
|
||||
do_sample=True,
|
||||
attention_mask=inputs["attention_mask"],
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
streamer=streamer,
|
||||
top_p=args.top_p,
|
||||
temperature=args.temperature
|
||||
)
|
||||
|
||||
messages.append({"role": "assistant", "content": answer})
|
||||
response = tokenizer.decode(generated_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
|
||||
messages.append({"role": "assistant", "content": response})
|
||||
print('\n\n')
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -1,61 +0,0 @@
|
||||
from transformers import PretrainedConfig
|
||||
from typing import List
|
||||
|
||||
|
||||
class LMConfig(PretrainedConfig):
|
||||
model_type = "minimind"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int = 512,
|
||||
n_layers: int = 8,
|
||||
n_heads: int = 8,
|
||||
n_kv_heads: int = 2,
|
||||
vocab_size: int = 6400,
|
||||
hidden_dim: int = None,
|
||||
multiple_of: int = 64,
|
||||
norm_eps: float = 1e-5,
|
||||
max_seq_len: int = 8192,
|
||||
rope_theta: int = 1e6,
|
||||
dropout: float = 0.0,
|
||||
flash_attn: bool = True,
|
||||
####################################################
|
||||
# Here are the specific configurations of MOE
|
||||
# When use_moe is false, the following is invalid
|
||||
####################################################
|
||||
use_moe: bool = False,
|
||||
####################################################
|
||||
num_experts_per_tok: int = 2,
|
||||
n_routed_experts: int = 4,
|
||||
n_shared_experts: bool = True,
|
||||
scoring_func: str = 'softmax',
|
||||
aux_loss_alpha: float = 0.1,
|
||||
seq_aux: bool = True,
|
||||
norm_topk_prob: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
self.dim = dim
|
||||
self.n_layers = n_layers
|
||||
self.n_heads = n_heads
|
||||
self.n_kv_heads = n_kv_heads
|
||||
self.vocab_size = vocab_size
|
||||
self.hidden_dim = hidden_dim
|
||||
self.multiple_of = multiple_of
|
||||
self.norm_eps = norm_eps
|
||||
self.max_seq_len = max_seq_len
|
||||
self.rope_theta = rope_theta
|
||||
self.dropout = dropout
|
||||
self.flash_attn = flash_attn
|
||||
####################################################
|
||||
# Here are the specific configurations of MOE
|
||||
# When use_moe is false, the following is invalid
|
||||
####################################################
|
||||
self.use_moe = use_moe
|
||||
self.num_experts_per_tok = num_experts_per_tok # 每个token选择的专家数量
|
||||
self.n_routed_experts = n_routed_experts # 总的专家数量
|
||||
self.n_shared_experts = n_shared_experts # 共享专家
|
||||
self.scoring_func = scoring_func # 评分函数，默认为'softmax'
|
||||
self.aux_loss_alpha = aux_loss_alpha # 辅助损失的alpha参数
|
||||
self.seq_aux = seq_aux # 是否在序列级别上计算辅助损失
|
||||
self.norm_topk_prob = norm_topk_prob # 是否标准化top-k概率
|
||||
super().__init__(**kwargs)
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -1,43 +0,0 @@
|
||||
{
|
||||
"add_bos_token": false,
|
||||
"add_eos_token": false,
|
||||
"add_prefix_space": false,
|
||||
"added_tokens_decoder": {
|
||||
"0": {
|
||||
"content": "<unk>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"1": {
|
||||
"content": "<s>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"2": {
|
||||
"content": "</s>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
}
|
||||
},
|
||||
"additional_special_tokens": [],
|
||||
"bos_token": "<s>",
|
||||
"clean_up_tokenization_spaces": false,
|
||||
"eos_token": "</s>",
|
||||
"legacy": true,
|
||||
"model_max_length": 32768,
|
||||
"pad_token": "<unk>",
|
||||
"sp_model_kwargs": {},
|
||||
"spaces_between_special_tokens": false,
|
||||
"tokenizer_class": "PreTrainedTokenizerFast",
|
||||
"unk_token": "<unk>",
|
||||
"chat_template": "{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{{ '<s>system\\n' + system_message + '</s>\\n' }}{% else %}{{ '<s>system\\n你是 MiniMind，是一个有用的人工智能助手。</s>\\n' }}{% endif %}{% for message in messages %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<s>user\\n' + content + '</s>\\n<s>assistant\\n' }}{% elif message['role'] == 'assistant' %}{{ content + '</s>' + '\\n' }}{% endif %}{% endfor %}"
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
-385
@@ -1,385 +0,0 @@
|
||||
import math
|
||||
import struct
|
||||
import inspect
|
||||
import time
|
||||
|
||||
from .LMConfig import LMConfig
|
||||
from typing import Any, Optional, Tuple, List, Union
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
from transformers import PreTrainedModel
|
||||
from transformers.modeling_outputs import CausalLMOutputWithPast
|
||||
|
||||
|
||||
class RMSNorm(torch.nn.Module):
|
||||
def __init__(self, dim: int, eps: float = 1e-6):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
|
||||
def _norm(self, x):
|
||||
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
||||
|
||||
def forward(self, x):
|
||||
return self.weight * self._norm(x.float()).type_as(x)
|
||||
|
||||
|
||||
def precompute_pos_cis(dim: int, end: int = int(32 * 1024), theta: float = 1e6):
|
||||
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
|
||||
t = torch.arange(end, device=freqs.device) # type: ignore
|
||||
freqs = torch.outer(t, freqs).float() # type: ignore
|
||||
pos_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64
|
||||
return pos_cis
|
||||
|
||||
|
||||
def apply_rotary_emb(xq, xk, pos_cis):
|
||||
def unite_shape(pos_cis, x):
|
||||
ndim = x.ndim
|
||||
assert 0 <= 1 < ndim
|
||||
assert pos_cis.shape == (x.shape[1], x.shape[-1])
|
||||
shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
|
||||
return pos_cis.view(*shape)
|
||||
|
||||
xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
|
||||
xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
|
||||
pos_cis = unite_shape(pos_cis, xq_)
|
||||
xq_out = torch.view_as_real(xq_ * pos_cis).flatten(3)
|
||||
xk_out = torch.view_as_real(xk_ * pos_cis).flatten(3)
|
||||
return xq_out.type_as(xq), xk_out.type_as(xk)
|
||||
|
||||
|
||||
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
|
||||
"""torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
|
||||
bs, slen, n_kv_heads, head_dim = x.shape
|
||||
if n_rep == 1:
|
||||
return x
|
||||
return (
|
||||
x[:, :, :, None, :]
|
||||
.expand(bs, slen, n_kv_heads, n_rep, head_dim)
|
||||
.reshape(bs, slen, n_kv_heads * n_rep, head_dim)
|
||||
)
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: LMConfig):
|
||||
super().__init__()
|
||||
self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
|
||||
assert args.n_heads % self.n_kv_heads == 0
|
||||
self.n_local_heads = args.n_heads
|
||||
self.n_local_kv_heads = self.n_kv_heads
|
||||
self.n_rep = self.n_local_heads // self.n_local_kv_heads
|
||||
self.head_dim = args.dim // args.n_heads
|
||||
self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
|
||||
self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
|
||||
self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
|
||||
self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
|
||||
self.attn_dropout = nn.Dropout(args.dropout)
|
||||
self.resid_dropout = nn.Dropout(args.dropout)
|
||||
self.dropout = args.dropout
|
||||
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn
|
||||
# print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
|
||||
mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
|
||||
mask = torch.triu(mask, diagonal=1)
|
||||
self.register_buffer("mask", mask, persistent=False)
|
||||
|
||||
def forward(self,
|
||||
x: torch.Tensor,
|
||||
pos_cis: torch.Tensor,
|
||||
past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
||||
use_cache=False):
|
||||
bsz, seq_len, _ = x.shape
|
||||
xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
|
||||
xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
|
||||
xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
|
||||
xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
|
||||
|
||||
xq, xk = apply_rotary_emb(xq, xk, pos_cis)
|
||||
# kv_cache实现
|
||||
if past_key_value is not None:
|
||||
xk = torch.cat([past_key_value[0], xk], dim=1)
|
||||
xv = torch.cat([past_key_value[1], xv], dim=1)
|
||||
past_kv = (xk, xv) if use_cache else None
|
||||
|
||||
xq, xk, xv = (
|
||||
xq.transpose(1, 2),
|
||||
repeat_kv(xk, self.n_rep).transpose(1, 2),
|
||||
repeat_kv(xv, self.n_rep).transpose(1, 2)
|
||||
)
|
||||
if self.flash and seq_len != 1:
|
||||
dropout_p = self.dropout if self.training else 0.0
|
||||
output = F.scaled_dot_product_attention(
|
||||
xq, xk, xv,
|
||||
attn_mask=None,
|
||||
dropout_p=dropout_p,
|
||||
is_causal=True
|
||||
)
|
||||
else:
|
||||
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
|
||||
scores += self.mask[:, :, :seq_len, :seq_len]
|
||||
scores = F.softmax(scores.float(), dim=-1).type_as(xq)
|
||||
scores = self.attn_dropout(scores)
|
||||
output = scores @ xv
|
||||
|
||||
output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
|
||||
output = self.resid_dropout(self.wo(output))
|
||||
return output, past_kv
|
||||
|
||||
|
||||
class FeedForward(nn.Module):
|
||||
def __init__(self, config: LMConfig):
|
||||
super().__init__()
|
||||
if config.hidden_dim is None:
|
||||
hidden_dim = 4 * config.dim
|
||||
hidden_dim = int(2 * hidden_dim / 3)
|
||||
config.hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of)
|
||||
self.w1 = nn.Linear(config.dim, config.hidden_dim, bias=False)
|
||||
self.w2 = nn.Linear(config.hidden_dim, config.dim, bias=False)
|
||||
self.w3 = nn.Linear(config.dim, config.hidden_dim, bias=False)
|
||||
self.dropout = nn.Dropout(config.dropout)
|
||||
|
||||
def forward(self, x):
|
||||
return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
|
||||
|
||||
|
||||
class MoEGate(nn.Module):
|
||||
def __init__(self, config: LMConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.top_k = config.num_experts_per_tok
|
||||
self.n_routed_experts = config.n_routed_experts
|
||||
|
||||
self.scoring_func = config.scoring_func
|
||||
self.alpha = config.aux_loss_alpha
|
||||
self.seq_aux = config.seq_aux
|
||||
|
||||
self.norm_topk_prob = config.norm_topk_prob
|
||||
self.gating_dim = config.dim
|
||||
self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
import torch.nn.init as init
|
||||
init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
||||
|
||||
def forward(self, hidden_states):
|
||||
bsz, seq_len, h = hidden_states.shape
|
||||
hidden_states = hidden_states.view(-1, h)
|
||||
logits = F.linear(hidden_states, self.weight, None)
|
||||
if self.scoring_func == 'softmax':
|
||||
scores = logits.softmax(dim=-1)
|
||||
else:
|
||||
raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}')
|
||||
|
||||
topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
|
||||
|
||||
if self.top_k > 1 and self.norm_topk_prob:
|
||||
denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
|
||||
topk_weight = topk_weight / denominator
|
||||
|
||||
if self.training and self.alpha > 0.0:
|
||||
scores_for_aux = scores
|
||||
aux_topk = self.top_k
|
||||
topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
|
||||
if self.seq_aux:
|
||||
scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
|
||||
ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)
|
||||
ce.scatter_add_(1, topk_idx_for_aux_loss,
|
||||
torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device)).div_(
|
||||
seq_len * aux_topk / self.n_routed_experts)
|
||||
aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha
|
||||
else:
|
||||
mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
|
||||
ce = mask_ce.float().mean(0)
|
||||
Pi = scores_for_aux.mean(0)
|
||||
fi = ce * self.n_routed_experts
|
||||
aux_loss = (Pi * fi).sum() * self.alpha
|
||||
else:
|
||||
aux_loss = 0
|
||||
return topk_idx, topk_weight, aux_loss
|
||||
|
||||
|
||||
class MOEFeedForward(nn.Module):
|
||||
def __init__(self, config: LMConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.experts = nn.ModuleList([
|
||||
FeedForward(config)
|
||||
for _ in range(config.n_routed_experts)
|
||||
])
|
||||
self.gate = MoEGate(config)
|
||||
if config.n_shared_experts is not None:
|
||||
self.shared_experts = FeedForward(config)
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
orig_shape = x.shape
|
||||
bsz, seq_len, _ = x.shape
|
||||
# 使用门控机制选择专家
|
||||
topk_idx, topk_weight, aux_loss = self.gate(x)
|
||||
x = x.view(-1, x.shape[-1])
|
||||
flat_topk_idx = topk_idx.view(-1)
|
||||
if self.training:
|
||||
x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)
|
||||
y = torch.empty_like(x, dtype=torch.float16)
|
||||
for i, expert in enumerate(self.experts):
|
||||
y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(y.dtype) # 确保类型一致
|
||||
y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
|
||||
y = y.view(*orig_shape)
|
||||
else:
|
||||
y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)
|
||||
if self.config.n_shared_experts is not None:
|
||||
y = y + self.shared_experts(identity)
|
||||
self.aux_loss = aux_loss
|
||||
return y
|
||||
|
||||
@torch.no_grad()
|
||||
def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
|
||||
expert_cache = torch.zeros_like(x)
|
||||
idxs = flat_expert_indices.argsort()
|
||||
tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
|
||||
token_idxs = idxs // self.config.num_experts_per_tok
|
||||
# 当tokens_per_expert = [6, 15, 20, 26]，tokens_per_expert.shape[0]即为专家数量（此时为4）
|
||||
# 且token_idxs = [3, 7, 19, 21, 24, 25, 4, 5, 6, 10, 11, 12...] 时
|
||||
# 意味token_idxs[:6] -> [3, 7, 19, 21, 24, 25]这6个位置属于专家0处理的token（每个token有可能被多个专家处理，这取决于num_experts_per_tok）
|
||||
# 接下来9个位置token_idxs[6:15] -> [4, 5, 6, 10, 11, 12...]属于专家1处理的token...依此类推
|
||||
for i, end_idx in enumerate(tokens_per_expert):
|
||||
start_idx = 0 if i == 0 else tokens_per_expert[i - 1]
|
||||
if start_idx == end_idx:
|
||||
continue
|
||||
expert = self.experts[i]
|
||||
exp_token_idx = token_idxs[start_idx:end_idx]
|
||||
expert_tokens = x[exp_token_idx]
|
||||
expert_out = expert(expert_tokens).to(expert_cache.dtype)
|
||||
expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
|
||||
expert_cache.scatter_add_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out)
|
||||
|
||||
return expert_cache
|
||||
|
||||
|
||||
class MiniMindBlock(nn.Module):
|
||||
def __init__(self, layer_id: int, config: LMConfig):
|
||||
super().__init__()
|
||||
self.n_heads = config.n_heads
|
||||
self.dim = config.dim
|
||||
self.head_dim = config.dim // config.n_heads
|
||||
self.attention = Attention(config)
|
||||
|
||||
self.layer_id = layer_id
|
||||
self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)
|
||||
self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)
|
||||
self.feed_forward = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
|
||||
|
||||
def forward(self, x, pos_cis, past_key_value=None, use_cache=False):
|
||||
h_attn, past_kv = self.attention(
|
||||
self.attention_norm(x),
|
||||
pos_cis,
|
||||
past_key_value=past_key_value,
|
||||
use_cache=use_cache
|
||||
)
|
||||
h = x + h_attn
|
||||
out = h + self.feed_forward(self.ffn_norm(h))
|
||||
return out, past_kv
|
||||
|
||||
|
||||
class MiniMindLM(PreTrainedModel):
|
||||
config_class = LMConfig
|
||||
|
||||
def __init__(self, params: LMConfig = None):
|
||||
self.params = params or LMConfig()
|
||||
super().__init__(self.params)
|
||||
self.vocab_size, self.n_layers = params.vocab_size, params.n_layers
|
||||
self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)
|
||||
self.dropout = nn.Dropout(params.dropout)
|
||||
self.layers = nn.ModuleList([MiniMindBlock(l, params) for l in range(self.n_layers)])
|
||||
self.norm = RMSNorm(params.dim, eps=params.norm_eps)
|
||||
self.output = nn.Linear(params.dim, params.vocab_size, bias=False)
|
||||
self.tok_embeddings.weight = self.output.weight
|
||||
self.register_buffer("pos_cis",
|
||||
precompute_pos_cis(dim=params.dim // params.n_heads, theta=params.rope_theta),
|
||||
persistent=False)
|
||||
self.OUT = CausalLMOutputWithPast()
|
||||
|
||||
def forward(self,
|
||||
input_ids: Optional[torch.Tensor] = None,
|
||||
past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
|
||||
use_cache: bool = False,
|
||||
logits_to_keep: Union[int, torch.Tensor] = 0,
|
||||
**args):
|
||||
past_key_values = past_key_values or [None] * len(self.layers)
|
||||
start_pos = args.get('start_pos', 0)
|
||||
h = self.dropout(self.tok_embeddings(input_ids))
|
||||
pos_cis = self.pos_cis[start_pos:start_pos + input_ids.size(1)]
|
||||
past_kvs = []
|
||||
for l, layer in enumerate(self.layers):
|
||||
h, past_kv = layer(
|
||||
h, pos_cis,
|
||||
past_key_value=past_key_values[l],
|
||||
use_cache=use_cache
|
||||
)
|
||||
past_kvs.append(past_kv)
|
||||
|
||||
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
||||
logits = self.output(self.norm(h)[:, slice_indices, :])
|
||||
aux_loss = sum(l.feed_forward.aux_loss for l in self.layers if isinstance(l.feed_forward, MOEFeedForward))
|
||||
self.OUT.__setitem__('last_hidden_state', h)
|
||||
self.OUT.__setitem__('logits', logits)
|
||||
self.OUT.__setitem__('aux_loss', aux_loss)
|
||||
self.OUT.__setitem__('past_key_values', past_kvs)
|
||||
return self.OUT
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate(self, input_ids, eos_token_id=2, max_new_tokens=1024, temperature=0.75, top_p=0.90,
|
||||
stream=False, rp=1., use_cache=True, pad_token_id=0, num_return_sequences=1, **args):
|
||||
# 流式生成
|
||||
if stream:
|
||||
return self._stream(input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args)
|
||||
|
||||
# 直接生成
|
||||
generated = []
|
||||
for i in range(input_ids.size(0)):
|
||||
non_pad = input_ids[i][input_ids[i] != pad_token_id].unsqueeze(0)
|
||||
for _ in range(num_return_sequences):
|
||||
out = self._stream(non_pad, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args)
|
||||
tokens_list = [tokens[:, -1:] for tokens in out]
|
||||
gen = torch.cat(tokens_list, dim=-1) if tokens_list else non_pad
|
||||
full_sequence = torch.cat([non_pad, gen], dim=-1)
|
||||
generated.append(full_sequence)
|
||||
|
||||
max_length = max(seq.size(1) for seq in generated)
|
||||
generated = [
|
||||
torch.cat(
|
||||
[seq, torch.full((1, max_length - seq.size(1)), pad_token_id, dtype=seq.dtype, device=seq.device)],
|
||||
dim=-1)
|
||||
for seq in generated
|
||||
]
|
||||
output = torch.cat(generated, dim=0)
|
||||
res = output.view(input_ids.size(0) * num_return_sequences, -1)
|
||||
return res
|
||||
|
||||
def _stream(self, input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args):
|
||||
start, first_seq, past_kvs = input_ids.shape[1], True, None
|
||||
while input_ids.shape[1] < max_new_tokens - 1:
|
||||
if first_seq or not use_cache:
|
||||
out, first_seq = self(input_ids, past_key_values=past_kvs, use_cache=use_cache, **args), False
|
||||
else:
|
||||
out = self(input_ids[:, -1:], past_key_values=past_kvs, use_cache=use_cache,
|
||||
start_pos=input_ids.shape[1] - 1, **args)
|
||||
logits, past_kvs = out.logits[:, -1, :], out.past_key_values
|
||||
logits[:, list(set(input_ids.tolist()[0]))] /= rp
|
||||
logits /= (temperature + 1e-9)
|
||||
if top_p is not None and top_p < 1.0:
|
||||
sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
|
||||
sorted_probs = F.softmax(sorted_logits, dim=-1)
|
||||
cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
|
||||
sorted_indices_to_remove = cumulative_probs > top_p
|
||||
sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
|
||||
sorted_indices_to_remove[:, 0] = False
|
||||
indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
|
||||
logits[indices_to_remove] = -float('Inf')
|
||||
input_ids_next = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
|
||||
input_ids = torch.cat((input_ids, input_ids_next), dim=1)
|
||||
yield input_ids[:, start:]
|
||||
if input_ids_next.item() == eos_token_id:
|
||||
break
|
||||
@@ -0,0 +1,446 @@
|
||||
# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
|
||||
# MiniMind Config
|
||||
# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
|
||||
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
|
||||
class MiniMindConfig(PretrainedConfig):
|
||||
model_type = "minimind"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dropout: float = 0.0,
|
||||
bos_token_id: int = 1,
|
||||
eos_token_id: int = 2,
|
||||
hidden_act: str = 'silu',
|
||||
hidden_size: int = 512,
|
||||
intermediate_size: int = None,
|
||||
max_position_embeddings: int = 32768,
|
||||
num_attention_heads: int = 8,
|
||||
num_hidden_layers: int = 8,
|
||||
num_key_value_heads: int = 2,
|
||||
vocab_size: int = 6400,
|
||||
rms_norm_eps: float = 1e-05,
|
||||
rope_theta: int = 1000000.0,
|
||||
flash_attn: bool = True,
|
||||
####################################################
|
||||
# Here are the specific configurations of MOE
|
||||
# When use_moe is false, the following is invalid
|
||||
####################################################
|
||||
use_moe: bool = False,
|
||||
num_experts_per_tok: int = 2,
|
||||
n_routed_experts: int = 4,
|
||||
n_shared_experts: int = 1,
|
||||
scoring_func: str = 'softmax',
|
||||
aux_loss_alpha: float = 0.1,
|
||||
seq_aux: bool = True,
|
||||
norm_topk_prob: bool = True,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
self.dropout = dropout
|
||||
self.bos_token_id = bos_token_id
|
||||
self.eos_token_id = eos_token_id
|
||||
self.hidden_act = hidden_act
|
||||
self.hidden_size = hidden_size
|
||||
self.intermediate_size = intermediate_size
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.num_hidden_layers = num_hidden_layers
|
||||
self.num_key_value_heads = num_key_value_heads
|
||||
self.vocab_size = vocab_size
|
||||
self.rms_norm_eps = rms_norm_eps
|
||||
self.rope_theta = rope_theta
|
||||
self.flash_attn = flash_attn
|
||||
####################################################
|
||||
# Here are the specific configurations of MOE
|
||||
# When use_moe is false, the following is invalid
|
||||
####################################################
|
||||
self.use_moe = use_moe
|
||||
self.num_experts_per_tok = num_experts_per_tok # 每个token选择的专家数量
|
||||
self.n_routed_experts = n_routed_experts # 总的专家数量
|
||||
self.n_shared_experts = n_shared_experts # 共享专家
|
||||
self.scoring_func = scoring_func # 评分函数，默认为'softmax'
|
||||
self.aux_loss_alpha = aux_loss_alpha # 辅助损失的alpha参数
|
||||
self.seq_aux = seq_aux # 是否在序列级别上计算辅助损失
|
||||
self.norm_topk_prob = norm_topk_prob # 是否标准化top-k概率
|
||||
|
||||
|
||||
# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
|
||||
# MiniMind Model
|
||||
# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
|
||||
|
||||
import math
|
||||
import torch
|
||||
from torch import nn
|
||||
from transformers.activations import ACT2FN
|
||||
from typing import Optional, Tuple, List, Union
|
||||
import torch.nn.functional as F
|
||||
from transformers import PreTrainedModel, GenerationMixin, PretrainedConfig
|
||||
from transformers.modeling_outputs import CausalLMOutputWithPast
|
||||
|
||||
|
||||
class RMSNorm(torch.nn.Module):
|
||||
def __init__(self, dim: int, eps: float = 1e-5):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
|
||||
def _norm(self, x):
|
||||
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
||||
|
||||
def forward(self, x):
|
||||
return self.weight * self._norm(x.float()).type_as(x)
|
||||
|
||||
|
||||
def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), theta: float = 1e6):
|
||||
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
|
||||
t = torch.arange(end, device=freqs.device)
|
||||
freqs = torch.outer(t, freqs).float()
|
||||
freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1)
|
||||
freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1)
|
||||
return freqs_cos, freqs_sin
|
||||
|
||||
|
||||
def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
|
||||
def rotate_half(x):
|
||||
return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)
|
||||
|
||||
q_embed = (q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))
|
||||
k_embed = (k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))
|
||||
return q_embed, k_embed
|
||||
|
||||
|
||||
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
|
||||
"""torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
|
||||
bs, slen, num_key_value_heads, head_dim = x.shape
|
||||
if n_rep == 1:
|
||||
return x
|
||||
return (
|
||||
x[:, :, :, None, :]
|
||||
.expand(bs, slen, num_key_value_heads, n_rep, head_dim)
|
||||
.reshape(bs, slen, num_key_value_heads * n_rep, head_dim)
|
||||
)
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: MiniMindConfig):
|
||||
super().__init__()
|
||||
self.num_key_value_heads = args.num_attention_heads if args.num_key_value_heads is None else args.num_key_value_heads
|
||||
assert args.num_attention_heads % self.num_key_value_heads == 0
|
||||
self.n_local_heads = args.num_attention_heads
|
||||
self.n_local_kv_heads = self.num_key_value_heads
|
||||
self.n_rep = self.n_local_heads // self.n_local_kv_heads
|
||||
self.head_dim = args.hidden_size // args.num_attention_heads
|
||||
self.q_proj = nn.Linear(args.hidden_size, args.num_attention_heads * self.head_dim, bias=False)
|
||||
self.k_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
|
||||
self.v_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
|
||||
self.o_proj = nn.Linear(args.num_attention_heads * self.head_dim, args.hidden_size, bias=False)
|
||||
self.attn_dropout = nn.Dropout(args.dropout)
|
||||
self.resid_dropout = nn.Dropout(args.dropout)
|
||||
self.dropout = args.dropout
|
||||
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn
|
||||
# print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
|
||||
|
||||
def forward(self,
|
||||
x: torch.Tensor,
|
||||
position_embeddings: Tuple[torch.Tensor, torch.Tensor], # 修改为接收cos和sin
|
||||
past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
||||
use_cache=False,
|
||||
attention_mask: Optional[torch.Tensor] = None):
|
||||
bsz, seq_len, _ = x.shape
|
||||
xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
|
||||
xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
|
||||
xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
|
||||
|
||||
cos, sin = position_embeddings
|
||||
xq, xk = apply_rotary_pos_emb(xq, xk, cos[:seq_len], sin[:seq_len])
|
||||
|
||||
# kv_cache实现
|
||||
if past_key_value is not None:
|
||||
xk = torch.cat([past_key_value[0], xk], dim=1)
|
||||
xv = torch.cat([past_key_value[1], xv], dim=1)
|
||||
past_kv = (xk, xv) if use_cache else None
|
||||
|
||||
xq, xk, xv = (
|
||||
xq.transpose(1, 2),
|
||||
repeat_kv(xk, self.n_rep).transpose(1, 2),
|
||||
repeat_kv(xv, self.n_rep).transpose(1, 2)
|
||||
)
|
||||
|
||||
if False and self.flash and seq_len != 1:
|
||||
dropout_p = self.dropout if self.training else 0.0
|
||||
attn_mask = None
|
||||
if attention_mask is not None:
|
||||
attn_mask = attention_mask.view(bsz, 1, 1, -1).expand(bsz, self.n_local_heads, seq_len, -1)
|
||||
attn_mask = attn_mask.bool() if attention_mask is not None else None
|
||||
|
||||
output = F.scaled_dot_product_attention(xq, xk, xv, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=True)
|
||||
else:
|
||||
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
|
||||
scores = scores + torch.triu(
|
||||
torch.full((seq_len, seq_len), float("-inf"), device=scores.device),
|
||||
diagonal=1
|
||||
).unsqueeze(0).unsqueeze(0) # scores+mask
|
||||
|
||||
if attention_mask is not None:
|
||||
extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
|
||||
extended_attention_mask = (1.0 - extended_attention_mask) * -1e9
|
||||
scores = scores + extended_attention_mask
|
||||
|
||||
scores = F.softmax(scores.float(), dim=-1).type_as(xq)
|
||||
scores = self.attn_dropout(scores)
|
||||
output = scores @ xv
|
||||
|
||||
output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
|
||||
output = self.resid_dropout(self.o_proj(output))
|
||||
return output, past_kv
|
||||
|
||||
|
||||
class FeedForward(nn.Module):
|
||||
def __init__(self, config: MiniMindConfig):
|
||||
super().__init__()
|
||||
if config.intermediate_size is None:
|
||||
intermediate_size = int(config.hidden_size * 8 / 3)
|
||||
config.intermediate_size = 64 * ((intermediate_size + 64 - 1) // 64)
|
||||
self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
||||
self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
|
||||
self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
||||
self.dropout = nn.Dropout(config.dropout)
|
||||
self.act_fn = ACT2FN[config.hidden_act]
|
||||
|
||||
def forward(self, x):
|
||||
return self.dropout(self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)))
|
||||
|
||||
|
||||
class MoEGate(nn.Module):
|
||||
def __init__(self, config: MiniMindConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.top_k = config.num_experts_per_tok
|
||||
self.n_routed_experts = config.n_routed_experts
|
||||
|
||||
self.scoring_func = config.scoring_func
|
||||
self.alpha = config.aux_loss_alpha
|
||||
self.seq_aux = config.seq_aux
|
||||
|
||||
self.norm_topk_prob = config.norm_topk_prob
|
||||
self.gating_dim = config.hidden_size
|
||||
self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
import torch.nn.init as init
|
||||
init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
||||
|
||||
def forward(self, hidden_states):
|
||||
bsz, seq_len, h = hidden_states.shape
|
||||
hidden_states = hidden_states.view(-1, h)
|
||||
logits = F.linear(hidden_states, self.weight, None)
|
||||
if self.scoring_func == 'softmax':
|
||||
scores = logits.softmax(dim=-1)
|
||||
else:
|
||||
raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}')
|
||||
|
||||
topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
|
||||
|
||||
if self.top_k > 1 and self.norm_topk_prob:
|
||||
denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
|
||||
topk_weight = topk_weight / denominator
|
||||
|
||||
if self.training and self.alpha > 0.0:
|
||||
scores_for_aux = scores
|
||||
aux_topk = self.top_k
|
||||
topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
|
||||
if self.seq_aux:
|
||||
scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
|
||||
ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)
|
||||
ce.scatter_add_(1, topk_idx_for_aux_loss,
|
||||
torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device)).div_(
|
||||
seq_len * aux_topk / self.n_routed_experts)
|
||||
aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha
|
||||
else:
|
||||
mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
|
||||
ce = mask_ce.float().mean(0)
|
||||
Pi = scores_for_aux.mean(0)
|
||||
fi = ce * self.n_routed_experts
|
||||
aux_loss = (Pi * fi).sum() * self.alpha
|
||||
else:
|
||||
aux_loss = 0
|
||||
return topk_idx, topk_weight, aux_loss
|
||||
|
||||
|
||||
class MOEFeedForward(nn.Module):
|
||||
def __init__(self, config: MiniMindConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.experts = nn.ModuleList([
|
||||
FeedForward(config)
|
||||
for _ in range(config.n_routed_experts)
|
||||
])
|
||||
self.gate = MoEGate(config)
|
||||
if config.n_shared_experts > 0:
|
||||
self.shared_experts = nn.ModuleList([
|
||||
FeedForward(config)
|
||||
for _ in range(config.n_shared_experts)
|
||||
])
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
orig_shape = x.shape
|
||||
bsz, seq_len, _ = x.shape
|
||||
# 使用门控机制选择专家
|
||||
topk_idx, topk_weight, aux_loss = self.gate(x)
|
||||
x = x.view(-1, x.shape[-1])
|
||||
flat_topk_idx = topk_idx.view(-1)
|
||||
if self.training:
|
||||
x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)
|
||||
y = torch.empty_like(x, dtype=torch.float16)
|
||||
for i, expert in enumerate(self.experts):
|
||||
y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(y.dtype) # 确保类型一致
|
||||
y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
|
||||
y = y.view(*orig_shape)
|
||||
else:
|
||||
y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)
|
||||
if self.config.n_shared_experts > 0:
|
||||
for expert in self.shared_experts:
|
||||
y = y + expert(identity)
|
||||
self.aux_loss = aux_loss
|
||||
return y
|
||||
|
||||
@torch.no_grad()
|
||||
def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
|
||||
expert_cache = torch.zeros_like(x)
|
||||
idxs = flat_expert_indices.argsort()
|
||||
tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
|
||||
token_idxs = idxs // self.config.num_experts_per_tok
|
||||
# 当tokens_per_expert = [6, 15, 20, 26]，tokens_per_expert.shape[0]即为专家数量（此时为4）
|
||||
# 且token_idxs = [3, 7, 19, 21, 24, 25, 4, 5, 6, 10, 11, 12...] 时
|
||||
# 意味token_idxs[:6] -> [3, 7, 19, 21, 24, 25]这6个位置属于专家0处理的token（每个token有可能被多个专家处理，这取决于num_experts_per_tok）
|
||||
# 接下来9个位置token_idxs[6:15] -> [4, 5, 6, 10, 11, 12...]属于专家1处理的token...依此类推
|
||||
for i, end_idx in enumerate(tokens_per_expert):
|
||||
start_idx = 0 if i == 0 else tokens_per_expert[i - 1]
|
||||
if start_idx == end_idx:
|
||||
continue
|
||||
expert = self.experts[i]
|
||||
exp_token_idx = token_idxs[start_idx:end_idx]
|
||||
expert_tokens = x[exp_token_idx]
|
||||
expert_out = expert(expert_tokens).to(expert_cache.dtype)
|
||||
expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
|
||||
expert_cache.scatter_add_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out)
|
||||
|
||||
return expert_cache
|
||||
|
||||
|
||||
class MiniMindBlock(nn.Module):
|
||||
def __init__(self, layer_id: int, config: MiniMindConfig):
|
||||
super().__init__()
|
||||
self.num_attention_heads = config.num_attention_heads
|
||||
self.hidden_size = config.hidden_size
|
||||
self.head_dim = config.hidden_size // config.num_attention_heads
|
||||
self.self_attn = Attention(config)
|
||||
|
||||
self.layer_id = layer_id
|
||||
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
|
||||
|
||||
def forward(self, hidden_states, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
|
||||
residual = hidden_states
|
||||
hidden_states, present_key_value = self.self_attn(
|
||||
self.input_layernorm(hidden_states), position_embeddings,
|
||||
past_key_value, use_cache, attention_mask
|
||||
)
|
||||
hidden_states += residual
|
||||
hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
|
||||
return hidden_states, present_key_value
|
||||
|
||||
|
||||
class MiniMindModel(nn.Module):
|
||||
def __init__(self, config: MiniMindConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.vocab_size, self.num_hidden_layers = config.vocab_size, config.num_hidden_layers
|
||||
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
|
||||
self.dropout = nn.Dropout(config.dropout)
|
||||
self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(self.num_hidden_layers)])
|
||||
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.hidden_size // config.num_attention_heads,
|
||||
end=config.max_position_embeddings, theta=config.rope_theta)
|
||||
self.register_buffer("freqs_cos", freqs_cos, persistent=False)
|
||||
self.register_buffer("freqs_sin", freqs_sin, persistent=False)
|
||||
|
||||
def forward(self,
|
||||
input_ids: Optional[torch.Tensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
|
||||
use_cache: bool = False,
|
||||
**kwargs):
|
||||
batch_size, seq_length = input_ids.shape
|
||||
past_key_values = past_key_values or [None] * len(self.layers)
|
||||
start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
|
||||
|
||||
hidden_states = self.dropout(self.embed_tokens(input_ids))
|
||||
|
||||
position_embeddings = (
|
||||
self.freqs_cos[start_pos:start_pos + seq_length],
|
||||
self.freqs_sin[start_pos:start_pos + seq_length]
|
||||
)
|
||||
|
||||
presents = []
|
||||
for layer_idx, (layer, past_key_value) in enumerate(zip(self.layers, past_key_values)):
|
||||
hidden_states, present = layer(
|
||||
hidden_states,
|
||||
position_embeddings,
|
||||
past_key_value=past_key_value,
|
||||
use_cache=use_cache,
|
||||
attention_mask=attention_mask
|
||||
)
|
||||
presents.append(present)
|
||||
|
||||
hidden_states = self.norm(hidden_states)
|
||||
|
||||
aux_loss = sum(
|
||||
layer.mlp.aux_loss
|
||||
for layer in self.layers
|
||||
if isinstance(layer.mlp, MOEFeedForward)
|
||||
)
|
||||
|
||||
return hidden_states, presents, aux_loss
|
||||
|
||||
|
||||
class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
|
||||
config_class = MiniMindConfig
|
||||
|
||||
def __init__(self, config: MiniMindConfig = None):
|
||||
self.config = config or MiniMindConfig()
|
||||
super().__init__(self.config)
|
||||
self.model = MiniMindModel(self.config)
|
||||
self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
|
||||
self.model.embed_tokens.weight = self.lm_head.weight
|
||||
self.OUT = CausalLMOutputWithPast()
|
||||
|
||||
def forward(self,
|
||||
input_ids: Optional[torch.Tensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
|
||||
use_cache: bool = False,
|
||||
logits_to_keep: Union[int, torch.Tensor] = 0,
|
||||
**args):
|
||||
h, past_kvs, aux_loss = self.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
past_key_values=past_key_values,
|
||||
use_cache=use_cache,
|
||||
**args
|
||||
)
|
||||
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
||||
logits = self.lm_head(h[:, slice_indices, :])
|
||||
self.OUT.__setitem__('last_hidden_state', h)
|
||||
self.OUT.__setitem__('logits', logits)
|
||||
self.OUT.__setitem__('aux_loss', aux_loss)
|
||||
self.OUT.__setitem__('past_key_values', past_kvs)
|
||||
return self.OUT
|
||||
@@ -1,8 +1,8 @@
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(
|
||||
api_key="none",
|
||||
base_url="http://localhost:8998/v1"
|
||||
api_key="ollama",
|
||||
base_url="http://127.0.0.1:8998/v1"
|
||||
)
|
||||
stream = True
|
||||
conversation_history_origin = []
|
||||
|
||||
+46
-34
@@ -1,33 +1,57 @@
|
||||
import torch
|
||||
import warnings
|
||||
import sys
|
||||
import os
|
||||
import sys
|
||||
|
||||
__package__ = "scripts"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
from model.LMConfig import LMConfig
|
||||
from model.model import MiniMindLM
|
||||
import torch
|
||||
import warnings
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaConfig, LlamaForCausalLM
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
|
||||
warnings.filterwarnings('ignore', category=UserWarning)
|
||||
|
||||
|
||||
def convert_torch2transformers(torch_path, transformers_path):
|
||||
def export_tokenizer(transformers_path):
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model/minimind_tokenizer')
|
||||
tokenizer.save_pretrained(transformers_path)
|
||||
|
||||
LMConfig.register_for_auto_class()
|
||||
MiniMindLM.register_for_auto_class("AutoModelForCausalLM")
|
||||
lm_model = MiniMindLM(lm_config)
|
||||
# MoE模型需使用此函数转换
|
||||
def convert_torch2transformers_minimind(torch_path, transformers_path, dtype=torch.bfloat16):
|
||||
MiniMindConfig.register_for_auto_class()
|
||||
MiniMindForCausalLM.register_for_auto_class("AutoModelForCausalLM")
|
||||
lm_model = MiniMindForCausalLM(lm_config)
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
state_dict = torch.load(torch_path, map_location=device)
|
||||
lm_model.load_state_dict(state_dict, strict=False)
|
||||
lm_model = lm_model.to(dtype) # 转换模型权重精度
|
||||
model_params = sum(p.numel() for p in lm_model.parameters() if p.requires_grad)
|
||||
print(f'模型参数: {model_params / 1e6} 百万 = {model_params / 1e9} B (Billion)')
|
||||
lm_model.save_pretrained(transformers_path, safe_serialization=False)
|
||||
export_tokenizer(transformers_path)
|
||||
print(f"模型已保存为 Transformers 格式: {transformers_path}")
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model/')
|
||||
tokenizer.save_pretrained(transformers_path)
|
||||
print(f"模型已保存为 Transformers-MiniMind 格式: {transformers_path}")
|
||||
|
||||
|
||||
# LlamaForCausalLM结构兼容第三方生态
|
||||
def convert_torch2transformers_llama(torch_path, transformers_path, dtype=torch.bfloat16):
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
state_dict = torch.load(torch_path, map_location=device)
|
||||
llama_config = LlamaConfig(
|
||||
vocab_size=lm_config.vocab_size,
|
||||
hidden_size=lm_config.hidden_size,
|
||||
intermediate_size=64 * ((int(lm_config.hidden_size * 8 / 3) + 64 - 1) // 64),
|
||||
num_hidden_layers=lm_config.num_hidden_layers,
|
||||
num_attention_heads=lm_config.num_attention_heads,
|
||||
num_key_value_heads=lm_config.num_key_value_heads,
|
||||
max_position_embeddings=lm_config.max_seq_len,
|
||||
rms_norm_eps=lm_config.rms_norm_eps,
|
||||
rope_theta=lm_config.rope_theta,
|
||||
)
|
||||
llama_model = LlamaForCausalLM(llama_config)
|
||||
llama_model.load_state_dict(state_dict, strict=False)
|
||||
llama_model = llama_model.to(dtype) # 转换模型权重精度
|
||||
llama_model.save_pretrained(transformers_path)
|
||||
model_params = sum(p.numel() for p in llama_model.parameters() if p.requires_grad)
|
||||
print(f'模型参数: {model_params / 1e6} 百万 = {model_params / 1e9} B (Billion)')
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model/')
|
||||
tokenizer.save_pretrained(transformers_path)
|
||||
print(f"模型已保存为 Transformers-Llama 格式: {transformers_path}")
|
||||
|
||||
|
||||
def convert_transformers2torch(transformers_path, torch_path):
|
||||
@@ -36,27 +60,15 @@ def convert_transformers2torch(transformers_path, torch_path):
|
||||
print(f"模型已保存为 PyTorch 格式: {torch_path}")
|
||||
|
||||
|
||||
# don't need to use
|
||||
def push_to_hf(export_model_path):
|
||||
def init_model():
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model/minimind_tokenizer')
|
||||
model = AutoModelForCausalLM.from_pretrained(export_model_path, trust_remote_code=True)
|
||||
return model, tokenizer
|
||||
|
||||
model, tokenizer = init_model()
|
||||
# model.push_to_hub(model_path)
|
||||
# tokenizer.push_to_hub(model_path, safe_serialization=False)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
lm_config = LMConfig(dim=512, n_layers=8, max_seq_len=8192, use_moe=False)
|
||||
lm_config = MiniMindConfig(hidden_size=768, num_hidden_layers=16, max_seq_len=8192, use_moe=True)
|
||||
|
||||
torch_path = f"../out/rlhf_{lm_config.dim}{'_moe' if lm_config.use_moe else ''}.pth"
|
||||
torch_path = f"../out/full_sft_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth"
|
||||
|
||||
transformers_path = '../MiniMind2-Small'
|
||||
transformers_path = '../MiniMind2-MoE'
|
||||
|
||||
# convert torch to transformers model
|
||||
convert_torch2transformers(torch_path, transformers_path)
|
||||
convert_torch2transformers_minimind(torch_path, transformers_path)
|
||||
|
||||
# # convert transformers to torch model
|
||||
# convert_transformers2torch(transformers_path, torch_path)
|
||||
# # # convert transformers to torch model
|
||||
# # convert_transformers2torch(transformers_path, torch_path)
|
||||
|
||||
+74
-61
@@ -9,12 +9,14 @@ import time
|
||||
import torch
|
||||
import warnings
|
||||
import uvicorn
|
||||
|
||||
from threading import Thread
|
||||
from queue import Queue
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
from model.LMConfig import LMConfig
|
||||
from model.model import MiniMindLM
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from model.model_lora import apply_lora, load_lora
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
@@ -23,30 +25,25 @@ app = FastAPI()
|
||||
|
||||
|
||||
def init_model(args):
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model/minimind_tokenizer')
|
||||
if args.load == 0:
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model/')
|
||||
moe_path = '_moe' if args.use_moe else ''
|
||||
modes = {0: 'pretrain', 1: 'full_sft', 2: 'rlhf', 3: 'reason'}
|
||||
ckp = f'../{args.out_dir}/{modes[args.model_mode]}_{args.dim}{moe_path}.pth'
|
||||
|
||||
model = MiniMindLM(LMConfig(
|
||||
dim=args.dim,
|
||||
n_layers=args.n_layers,
|
||||
ckp = f'../{args.out_dir}/{modes[args.model_mode]}_{args.hidden_size}{moe_path}.pth'
|
||||
model = MiniMindForCausalLM(MiniMindConfig(
|
||||
hidden_size=args.hidden_size,
|
||||
num_hidden_layers=args.num_hidden_layers,
|
||||
max_seq_len=args.max_seq_len,
|
||||
use_moe=args.use_moe
|
||||
))
|
||||
|
||||
state_dict = torch.load(ckp, map_location=device)
|
||||
model.load_state_dict({k: v for k, v in state_dict.items() if 'mask' not in k}, strict=True)
|
||||
|
||||
model.load_state_dict(torch.load(ckp, map_location=device), strict=True)
|
||||
if args.lora_name != 'None':
|
||||
apply_lora(model)
|
||||
load_lora(model, f'../{args.out_dir}/{args.lora_name}_{args.dim}.pth')
|
||||
load_lora(model, f'../{args.out_dir}/{args.lora_name}_{args.hidden_size}.pth')
|
||||
else:
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
'./MiniMind2',
|
||||
trust_remote_code=True
|
||||
)
|
||||
model_path = '../MiniMind2'
|
||||
model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
||||
print(f'MiniMind模型参数量: {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.2f}M(illion)')
|
||||
return model.eval().to(device), tokenizer
|
||||
|
||||
@@ -58,42 +55,61 @@ class ChatRequest(BaseModel):
|
||||
top_p: float = 0.92
|
||||
max_tokens: int = 8192
|
||||
stream: bool = False
|
||||
tools: list = []
|
||||
|
||||
|
||||
class CustomStreamer(TextStreamer):
|
||||
def __init__(self, tokenizer, queue):
|
||||
super().__init__(tokenizer, skip_prompt=True, skip_special_tokens=True)
|
||||
self.queue = queue
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def on_finalized_text(self, text: str, stream_end: bool = False):
|
||||
self.queue.put(text)
|
||||
if stream_end:
|
||||
self.queue.put(None)
|
||||
|
||||
|
||||
def generate_stream_response(messages, temperature, top_p, max_tokens):
|
||||
try:
|
||||
new_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)[-max_tokens:]
|
||||
x = tokenizer(new_prompt).data['input_ids']
|
||||
x = (torch.tensor(x, dtype=torch.long, device=device)[None, ...])
|
||||
with torch.no_grad():
|
||||
res_y = model.generate(
|
||||
x,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
inputs = tokenizer(new_prompt, return_tensors="pt", truncation=True).to(device)
|
||||
|
||||
queue = Queue()
|
||||
streamer = CustomStreamer(tokenizer, queue)
|
||||
|
||||
def _generate():
|
||||
model.generate(
|
||||
inputs.input_ids,
|
||||
max_new_tokens=max_tokens,
|
||||
do_sample=True,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
stream=True,
|
||||
rp=1.,
|
||||
pad_token_id=tokenizer.pad_token_id
|
||||
attention_mask=inputs.attention_mask,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
streamer=streamer
|
||||
)
|
||||
history_idx = 0
|
||||
for y in res_y:
|
||||
answer = tokenizer.decode(y[0].tolist(), skip_special_tokens=True)
|
||||
if (answer and answer[-1] == '�') or not answer:
|
||||
continue
|
||||
delta = answer[history_idx:]
|
||||
history_idx = len(answer)
|
||||
json_data = {
|
||||
'id': f'chatcmpl-{int(time.time())}',
|
||||
'object': 'chat.completion.chunk',
|
||||
'created': int(time.time()),
|
||||
'model': 'minimind',
|
||||
'choices': [{'index': 0, 'delta': {'content': delta}, 'finish_reason': None}]
|
||||
}
|
||||
yield f"data: {json.dumps(json_data)}\n\n"
|
||||
|
||||
Thread(target=_generate).start()
|
||||
|
||||
while True:
|
||||
text = queue.get()
|
||||
if text is None:
|
||||
yield json.dumps({
|
||||
"choices": [{
|
||||
"delta": {},
|
||||
"finish_reason": "stop"
|
||||
}]
|
||||
}, ensure_ascii=False)
|
||||
break
|
||||
|
||||
yield json.dumps({
|
||||
"choices": [{"delta": {"content": text}}]
|
||||
}, ensure_ascii=False)
|
||||
|
||||
except Exception as e:
|
||||
yield f"data: {json.dumps({'error': str(e)})}\n\n"
|
||||
yield json.dumps({"error": str(e)})
|
||||
|
||||
|
||||
@app.post("/v1/chat/completions")
|
||||
@@ -101,12 +117,12 @@ async def chat_completions(request: ChatRequest):
|
||||
try:
|
||||
if request.stream:
|
||||
return StreamingResponse(
|
||||
generate_stream_response(
|
||||
(f"data: {chunk}\n\n" for chunk in generate_stream_response(
|
||||
messages=request.messages,
|
||||
temperature=request.temperature,
|
||||
top_p=request.top_p,
|
||||
max_tokens=request.max_tokens
|
||||
),
|
||||
)),
|
||||
media_type="text/event-stream"
|
||||
)
|
||||
else:
|
||||
@@ -115,20 +131,19 @@ async def chat_completions(request: ChatRequest):
|
||||
tokenize=False,
|
||||
add_generation_prompt=True
|
||||
)[-request.max_tokens:]
|
||||
x = tokenizer(new_prompt).data['input_ids']
|
||||
x = (torch.tensor(x, dtype=torch.long, device=device)[None, ...])
|
||||
inputs = tokenizer(new_prompt, return_tensors="pt", truncation=True).to(device)
|
||||
with torch.no_grad():
|
||||
res_y = model.generate(
|
||||
x,
|
||||
generated_ids = model.generate(
|
||||
inputs["input_ids"],
|
||||
max_length=inputs["input_ids"].shape[1] + request.max_tokens,
|
||||
do_sample=True,
|
||||
attention_mask=inputs["attention_mask"],
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
max_new_tokens=request.max_tokens,
|
||||
temperature=request.temperature,
|
||||
top_p=request.top_p,
|
||||
stream=False,
|
||||
rp=1.,
|
||||
pad_token_id=tokenizer.pad_token_id
|
||||
temperature=request.temperature
|
||||
)
|
||||
answer = tokenizer.decode(res_y.squeeze()[x.shape[1]:].tolist(), skip_special_tokens=True)
|
||||
answer = tokenizer.decode(generated_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
|
||||
return {
|
||||
"id": f"chatcmpl-{int(time.time())}",
|
||||
"object": "chat.completion",
|
||||
@@ -142,7 +157,6 @@ async def chat_completions(request: ChatRequest):
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@@ -151,14 +165,13 @@ if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Server for MiniMind")
|
||||
parser.add_argument('--out_dir', default='out', type=str)
|
||||
parser.add_argument('--lora_name', default='None', type=str)
|
||||
parser.add_argument('--dim', default=512, type=int)
|
||||
parser.add_argument('--n_layers', default=8, type=int)
|
||||
parser.add_argument('--hidden_size', default=768, type=int)
|
||||
parser.add_argument('--num_hidden_layers', default=16, type=int)
|
||||
parser.add_argument('--max_seq_len', default=8192, type=int)
|
||||
parser.add_argument('--use_moe', default=False, type=bool)
|
||||
parser.add_argument('--load', default=0, type=int, help="0: 从原生torch权重，1: 利用transformers加载")
|
||||
parser.add_argument('--model_mode', default=1, type=int, help="0: 预训练模型，1: SFT-Chat模型，2: RLHF-Chat模型，3: Reason模型")
|
||||
|
||||
parser.add_argument('--model_mode', default=1, type=int,
|
||||
help="0: 预训练模型，1: SFT-Chat模型，2: RLHF-Chat模型，3: Reason模型")
|
||||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||
model, tokenizer = init_model(parser.parse_args())
|
||||
|
||||
uvicorn.run(app, host="0.0.0.0", port=8998)
|
||||
|
||||
+15
-20
@@ -1,14 +1,9 @@
|
||||
import random
|
||||
from tqdm import tqdm
|
||||
from transformers import AutoTokenizer
|
||||
import json
|
||||
from datasets import load_dataset
|
||||
from tokenizers import (
|
||||
decoders,
|
||||
models,
|
||||
normalizers,
|
||||
pre_tokenizers,
|
||||
processors,
|
||||
trainers,
|
||||
Tokenizer,
|
||||
)
|
||||
@@ -32,7 +27,7 @@ def train_tokenizer():
|
||||
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
|
||||
|
||||
# 定义特殊token
|
||||
special_tokens = ["<unk>", "<s>", "</s>"]
|
||||
special_tokens = ["<|endoftext|>", "<|im_start|>", "<|im_end|>"]
|
||||
|
||||
# 设置训练器并添加特殊token
|
||||
trainer = trainers.BpeTrainer(
|
||||
@@ -52,15 +47,15 @@ def train_tokenizer():
|
||||
tokenizer.decoder = decoders.ByteLevel()
|
||||
|
||||
# 检查特殊token的索引
|
||||
assert tokenizer.token_to_id("<unk>") == 0
|
||||
assert tokenizer.token_to_id("<s>") == 1
|
||||
assert tokenizer.token_to_id("</s>") == 2
|
||||
assert tokenizer.token_to_id("<|endoftext|>") == 0
|
||||
assert tokenizer.token_to_id("<|im_start|>") == 1
|
||||
assert tokenizer.token_to_id("<|im_end|>") == 2
|
||||
|
||||
# 保存tokenizer
|
||||
tokenizer_dir = "../model/minimind_tokenizer"
|
||||
tokenizer_dir = "../model/"
|
||||
os.makedirs(tokenizer_dir, exist_ok=True)
|
||||
tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json"))
|
||||
tokenizer.model.save("../model/minimind_tokenizer")
|
||||
tokenizer.model.save("../model/")
|
||||
|
||||
# 手动创建配置文件
|
||||
config = {
|
||||
@@ -69,7 +64,7 @@ def train_tokenizer():
|
||||
"add_prefix_space": False,
|
||||
"added_tokens_decoder": {
|
||||
"0": {
|
||||
"content": "<unk>",
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": False,
|
||||
"normalized": False,
|
||||
"rstrip": False,
|
||||
@@ -77,7 +72,7 @@ def train_tokenizer():
|
||||
"special": True
|
||||
},
|
||||
"1": {
|
||||
"content": "<s>",
|
||||
"content": "<|im_start|>",
|
||||
"lstrip": False,
|
||||
"normalized": False,
|
||||
"rstrip": False,
|
||||
@@ -85,7 +80,7 @@ def train_tokenizer():
|
||||
"special": True
|
||||
},
|
||||
"2": {
|
||||
"content": "</s>",
|
||||
"content": "<|im_end|>",
|
||||
"lstrip": False,
|
||||
"normalized": False,
|
||||
"rstrip": False,
|
||||
@@ -94,17 +89,17 @@ def train_tokenizer():
|
||||
}
|
||||
},
|
||||
"additional_special_tokens": [],
|
||||
"bos_token": "<s>",
|
||||
"bos_token": "<|im_start|>",
|
||||
"clean_up_tokenization_spaces": False,
|
||||
"eos_token": "</s>",
|
||||
"eos_token": "<|im_end|>",
|
||||
"legacy": True,
|
||||
"model_max_length": 32768,
|
||||
"pad_token": "<unk>",
|
||||
"pad_token": "<|endoftext|>",
|
||||
"sp_model_kwargs": {},
|
||||
"spaces_between_special_tokens": False,
|
||||
"tokenizer_class": "PreTrainedTokenizerFast",
|
||||
"unk_token": "<unk>",
|
||||
"chat_template": "{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{{ '<s>system\\n' + system_message + '</s>\\n' }}{% else %}{{ '<s>system\\n你是 MiniMind，是一个有用的人工智能助手。</s>\\n' }}{% endif %}{% for message in messages %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<s>user\\n' + content + '</s>\\n<s>assistant\\n' }}{% elif message['role'] == 'assistant' %}{{ content + '</s>' + '\\n' }}{% endif %}{% endfor %}"
|
||||
"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 %}"
|
||||
}
|
||||
|
||||
# 保存配置文件
|
||||
@@ -118,7 +113,7 @@ def eval_tokenizer():
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
# 加载预训练的tokenizer
|
||||
tokenizer = AutoTokenizer.from_pretrained("../model/minimind_tokenizer")
|
||||
tokenizer = AutoTokenizer.from_pretrained("../model/")
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": "你是一个优秀的聊天机器人，总是给我正确的回应！"},
|
||||
|
||||
+100
-65
@@ -1,14 +1,13 @@
|
||||
import random
|
||||
import re
|
||||
import time
|
||||
from threading import Thread
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
import streamlit as st
|
||||
import torch
|
||||
|
||||
st.set_page_config(page_title="MiniMind", initial_sidebar_state="collapsed")
|
||||
|
||||
# 在文件开头的 CSS 样式中修改按钮样式
|
||||
st.markdown("""
|
||||
<style>
|
||||
/* 添加操作按钮样式 */
|
||||
@@ -70,7 +69,9 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
|
||||
def process_assistant_content(content):
|
||||
if 'R1' not in MODEL_PATHS[selected_model][1]:
|
||||
if model_source == "API" and 'R1' not in api_model_name:
|
||||
return content
|
||||
if model_source != "API" and 'R1' not in MODEL_PATHS[selected_model][1]:
|
||||
return content
|
||||
|
||||
if '<think>' in content and '</think>' in content:
|
||||
@@ -119,7 +120,6 @@ def init_chat_messages():
|
||||
if message["role"] == "assistant":
|
||||
with st.chat_message("assistant", avatar=image_url):
|
||||
st.markdown(process_assistant_content(message["content"]), unsafe_allow_html=True)
|
||||
# 在消息内容下方添加按钮
|
||||
if st.button("🗑", key=f"delete_{i}"):
|
||||
st.session_state.messages.pop(i)
|
||||
st.session_state.messages.pop(i - 1)
|
||||
@@ -137,8 +137,6 @@ def init_chat_messages():
|
||||
|
||||
return st.session_state.messages
|
||||
|
||||
|
||||
# 添加这两个辅助函数
|
||||
def regenerate_answer(index):
|
||||
st.session_state.messages.pop()
|
||||
st.session_state.chat_messages.pop()
|
||||
@@ -153,32 +151,34 @@ def delete_conversation(index):
|
||||
st.rerun()
|
||||
|
||||
|
||||
# 侧边栏模型选择
|
||||
st.sidebar.title("模型设定调整")
|
||||
|
||||
st.sidebar.text("【注】训练数据偏差，增加上下文记忆时\n多轮对话（较单轮）容易出现能力衰减")
|
||||
# st.sidebar.text("训练数据偏差，增加上下文记忆时\n多轮对话（较单轮）容易出现能力衰减")
|
||||
st.session_state.history_chat_num = st.sidebar.slider("Number of Historical Dialogues", 0, 6, 0, step=2)
|
||||
# st.session_state.history_chat_num = 0
|
||||
st.session_state.max_new_tokens = st.sidebar.slider("Max Sequence Length", 256, 8192, 8192, step=1)
|
||||
st.session_state.top_p = st.sidebar.slider("Top-P", 0.8, 0.99, 0.85, step=0.01)
|
||||
st.session_state.temperature = st.sidebar.slider("Temperature", 0.6, 1.2, 0.85, step=0.01)
|
||||
|
||||
# 模型路径映射
|
||||
MODEL_PATHS = {
|
||||
"MiniMind2-R1 (0.1B)": ["../MiniMind2-R1", "MiniMind2-R1"],
|
||||
"MiniMind2-Small-R1 (0.02B)": ["../MiniMind2-Small-R1", "MiniMind2-Small-R1"],
|
||||
"MiniMind2 (0.1B)": ["../MiniMind2", "MiniMind2"],
|
||||
"MiniMind2-MoE (0.15B)": ["../MiniMind2-MoE", "MiniMind2-MoE"],
|
||||
"MiniMind2-Small (0.02B)": ["../MiniMind2-Small", "MiniMind2-Small"],
|
||||
"MiniMind-V1 (0.1B)": ["../minimind-v1", "MiniMind-V1"],
|
||||
"MiniMind-V1-MoE (0.1B)": ["../minimind-v1-moe", "MiniMind-V1-MoE"],
|
||||
"MiniMind-V1-Small (0.02B)": ["../minimind-v1-small", "MiniMind-V1-Small"],
|
||||
}
|
||||
model_source = st.sidebar.radio("选择模型来源", ["本地模型", "API"], index=0)
|
||||
|
||||
selected_model = st.sidebar.selectbox('Models', list(MODEL_PATHS.keys()), index=2) # 默认选择 MiniMind2
|
||||
model_path = MODEL_PATHS[selected_model][0]
|
||||
if model_source == "API":
|
||||
api_url = st.sidebar.text_input("API URL", value="http://127.0.0.1:8000/v1")
|
||||
api_model_id = st.sidebar.text_input("Model ID", value="minimind")
|
||||
api_model_name = st.sidebar.text_input("Model Name", value="MiniMind2")
|
||||
api_key = st.sidebar.text_input("API Key", value="none", type="password")
|
||||
slogan = f"Hi, I'm {api_model_name}"
|
||||
else:
|
||||
MODEL_PATHS = {
|
||||
"MiniMind2-R1 (0.1B)": ["../MiniMind2-R1", "MiniMind2-R1"],
|
||||
"MiniMind2-Small-R1 (0.02B)": ["../MiniMind2-Small-R1", "MiniMind2-Small-R1"],
|
||||
"MiniMind2 (0.1B)": ["../MiniMind2", "MiniMind2"],
|
||||
"MiniMind2-MoE (0.15B)": ["../MiniMind2-MoE", "MiniMind2-MoE"],
|
||||
"MiniMind2-Small (0.02B)": ["../MiniMind2-Small", "MiniMind2-Small"]
|
||||
}
|
||||
|
||||
slogan = f"Hi, I'm {MODEL_PATHS[selected_model][1]}"
|
||||
selected_model = st.sidebar.selectbox('Models', list(MODEL_PATHS.keys()), index=2) # 默认选择 MiniMind2
|
||||
model_path = MODEL_PATHS[selected_model][0]
|
||||
slogan = f"Hi, I'm {MODEL_PATHS[selected_model][1]}"
|
||||
|
||||
image_url = "https://www.modelscope.cn/api/v1/studio/gongjy/MiniMind/repo?Revision=master&FilePath=images%2Flogo2.png&View=true"
|
||||
|
||||
@@ -205,23 +205,22 @@ def setup_seed(seed):
|
||||
|
||||
|
||||
def main():
|
||||
model, tokenizer = load_model_tokenizer(model_path)
|
||||
if model_source == "本地模型":
|
||||
model, tokenizer = load_model_tokenizer(model_path)
|
||||
else:
|
||||
model, tokenizer = None, None
|
||||
|
||||
# 初始化消息列表
|
||||
if "messages" not in st.session_state:
|
||||
st.session_state.messages = []
|
||||
st.session_state.chat_messages = []
|
||||
|
||||
# Use session state messages
|
||||
messages = st.session_state.messages
|
||||
|
||||
# 在显示历史消息的循环中
|
||||
for i, message in enumerate(messages):
|
||||
if message["role"] == "assistant":
|
||||
with st.chat_message("assistant", avatar=image_url):
|
||||
st.markdown(process_assistant_content(message["content"]), unsafe_allow_html=True)
|
||||
if st.button("×", key=f"delete_{i}"):
|
||||
# 删除当前消息及其之后的所有消息
|
||||
st.session_state.messages = st.session_state.messages[:i - 1]
|
||||
st.session_state.chat_messages = st.session_state.chat_messages[:i - 1]
|
||||
st.rerun()
|
||||
@@ -230,14 +229,11 @@ def main():
|
||||
f'<div style="display: flex; justify-content: flex-end;"><div style="display: inline-block; margin: 10px 0; padding: 8px 12px 8px 12px; background-color: gray; border-radius: 10px; color:white; ">{message["content"]}</div></div>',
|
||||
unsafe_allow_html=True)
|
||||
|
||||
# 处理新的输入或重新生成
|
||||
prompt = st.chat_input(key="input", placeholder="给 MiniMind 发送消息")
|
||||
|
||||
# 检查是否需要重新生成
|
||||
if hasattr(st.session_state, 'regenerate') and st.session_state.regenerate:
|
||||
prompt = st.session_state.last_user_message
|
||||
regenerate_index = st.session_state.regenerate_index # 获取重新生成的位置
|
||||
# 清除所有重新生成相关的状态
|
||||
regenerate_index = st.session_state.regenerate_index
|
||||
delattr(st.session_state, 'regenerate')
|
||||
delattr(st.session_state, 'last_user_message')
|
||||
delattr(st.session_state, 'regenerate_index')
|
||||
@@ -246,48 +242,87 @@ def main():
|
||||
st.markdown(
|
||||
f'<div style="display: flex; justify-content: flex-end;"><div style="display: inline-block; margin: 10px 0; padding: 8px 12px 8px 12px; background-color: gray; border-radius: 10px; color:white; ">{prompt}</div></div>',
|
||||
unsafe_allow_html=True)
|
||||
messages.append({"role": "user", "content": prompt})
|
||||
st.session_state.chat_messages.append({"role": "user", "content": prompt})
|
||||
messages.append({"role": "user", "content": prompt[-st.session_state.max_new_tokens:]})
|
||||
st.session_state.chat_messages.append({"role": "user", "content": prompt[-st.session_state.max_new_tokens:]})
|
||||
|
||||
with st.chat_message("assistant", avatar=image_url):
|
||||
placeholder = st.empty()
|
||||
random_seed = random.randint(0, 2 ** 32 - 1)
|
||||
setup_seed(random_seed)
|
||||
|
||||
st.session_state.chat_messages = system_prompt + st.session_state.chat_messages[
|
||||
-(st.session_state.history_chat_num + 1):]
|
||||
new_prompt = tokenizer.apply_chat_template(
|
||||
st.session_state.chat_messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True
|
||||
)[-(st.session_state.max_new_tokens - 1):]
|
||||
|
||||
x = torch.tensor(tokenizer(new_prompt)['input_ids'], device=device).unsqueeze(0)
|
||||
with torch.no_grad():
|
||||
res_y = model.generate(x, tokenizer.eos_token_id, max_new_tokens=st.session_state.max_new_tokens,
|
||||
temperature=st.session_state.temperature,
|
||||
top_p=st.session_state.top_p, stream=True)
|
||||
if model_source == "API":
|
||||
try:
|
||||
for y in res_y:
|
||||
answer = tokenizer.decode(y[0].tolist(), skip_special_tokens=True)
|
||||
if (answer and answer[-1] == '�') or not answer:
|
||||
continue
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(
|
||||
api_key=api_key,
|
||||
base_url=api_url
|
||||
)
|
||||
history_num = st.session_state.history_chat_num + 1 # +1 是为了包含当前的用户消息
|
||||
conversation_history = system_prompt + st.session_state.chat_messages[-history_num:]
|
||||
answer = ""
|
||||
response = client.chat.completions.create(
|
||||
model=api_model_id,
|
||||
messages=conversation_history,
|
||||
stream=True,
|
||||
temperature=st.session_state.temperature
|
||||
)
|
||||
|
||||
for chunk in response:
|
||||
content = chunk.choices[0].delta.content or ""
|
||||
answer += content
|
||||
placeholder.markdown(process_assistant_content(answer), unsafe_allow_html=True)
|
||||
except StopIteration:
|
||||
print("No answer")
|
||||
|
||||
assistant_answer = answer.replace(new_prompt, "")
|
||||
messages.append({"role": "assistant", "content": assistant_answer})
|
||||
st.session_state.chat_messages.append({"role": "assistant", "content": assistant_answer})
|
||||
except Exception as e:
|
||||
answer = f"API调用出错: {str(e)}"
|
||||
placeholder.markdown(answer, unsafe_allow_html=True)
|
||||
else:
|
||||
random_seed = random.randint(0, 2 ** 32 - 1)
|
||||
setup_seed(random_seed)
|
||||
|
||||
with st.empty():
|
||||
if st.button("×", key=f"delete_{len(messages) - 1}"):
|
||||
st.session_state.messages = st.session_state.messages[:-2]
|
||||
st.session_state.chat_messages = st.session_state.chat_messages[:-2]
|
||||
st.rerun()
|
||||
st.session_state.chat_messages = system_prompt + st.session_state.chat_messages[
|
||||
-(st.session_state.history_chat_num + 1):]
|
||||
new_prompt = tokenizer.apply_chat_template(
|
||||
st.session_state.chat_messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True
|
||||
)
|
||||
|
||||
inputs = tokenizer(
|
||||
new_prompt,
|
||||
return_tensors="pt",
|
||||
truncation=True
|
||||
).to(device)
|
||||
|
||||
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
|
||||
generation_kwargs = {
|
||||
"input_ids": inputs.input_ids,
|
||||
"max_length": inputs.input_ids.shape[1] + st.session_state.max_new_tokens,
|
||||
"num_return_sequences": 1,
|
||||
"do_sample": True,
|
||||
"attention_mask": inputs.attention_mask,
|
||||
"pad_token_id": tokenizer.pad_token_id,
|
||||
"eos_token_id": tokenizer.eos_token_id,
|
||||
"temperature": st.session_state.temperature,
|
||||
"top_p": 0.85,
|
||||
"streamer": streamer,
|
||||
}
|
||||
|
||||
Thread(target=model.generate, kwargs=generation_kwargs).start()
|
||||
|
||||
answer = ""
|
||||
for new_text in streamer:
|
||||
answer += new_text
|
||||
placeholder.markdown(process_assistant_content(answer), unsafe_allow_html=True)
|
||||
|
||||
messages.append({"role": "assistant", "content": answer})
|
||||
st.session_state.chat_messages.append({"role": "assistant", "content": answer})
|
||||
with st.empty():
|
||||
if st.button("×", key=f"delete_{len(messages) - 1}"):
|
||||
st.session_state.messages = st.session_state.messages[:-2]
|
||||
st.session_state.chat_messages = st.session_state.chat_messages[:-2]
|
||||
st.rerun()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
|
||||
|
||||
main()
|
||||
|
||||
@@ -1,23 +1,22 @@
|
||||
import os
|
||||
import platform
|
||||
import sys
|
||||
|
||||
__package__ = "trainer"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
import argparse
|
||||
import time
|
||||
import math
|
||||
import warnings
|
||||
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
|
||||
from torch import optim, nn
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
from model.model import MiniMindLM
|
||||
from model.LMConfig import LMConfig
|
||||
from model.dataset import SFTDataset
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from dataset.lm_dataset import SFTDataset
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
@@ -85,34 +84,35 @@ def train_epoch(epoch, wandb):
|
||||
args.epochs,
|
||||
step,
|
||||
iter_per_epoch,
|
||||
loss.item(),
|
||||
loss.item() * args.accumulation_steps,
|
||||
optimizer.param_groups[-1]['lr'],
|
||||
spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
|
||||
|
||||
if (wandb is not None) and (not ddp or dist.get_rank() == 0):
|
||||
wandb.log({"loss": loss,
|
||||
wandb.log({"loss": loss * args.accumulation_steps,
|
||||
"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):
|
||||
model.eval()
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'{args.save_dir}/reason_{lm_config.dim}{moe_path}.pth'
|
||||
ckp = f'{args.save_dir}/reason_{lm_config.hidden_size}{moe_path}.pth'
|
||||
|
||||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||||
state_dict = model.module.state_dict()
|
||||
else:
|
||||
state_dict = model.state_dict()
|
||||
|
||||
state_dict = {k: v.half() for k, v in state_dict.items()} # 半精度保存
|
||||
torch.save(state_dict, ckp)
|
||||
model.train()
|
||||
|
||||
|
||||
def init_model(lm_config):
|
||||
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
|
||||
model = MiniMindLM(lm_config)
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model')
|
||||
model = MiniMindForCausalLM(lm_config)
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'./out/rlhf_{lm_config.dim}{moe_path}.pth'
|
||||
ckp = f'{args.save_dir}/rlhf_{lm_config.hidden_size}{moe_path}.pth'
|
||||
state_dict = torch.load(ckp, map_location=args.device)
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
||||
@@ -134,7 +134,7 @@ def init_distributed_mode():
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="MiniMind Distill Reasoning")
|
||||
parser.add_argument("--out_dir", type=str, default="out")
|
||||
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=8)
|
||||
parser.add_argument("--learning_rate", type=float, default=1e-6)
|
||||
@@ -150,19 +150,20 @@ if __name__ == "__main__":
|
||||
parser.add_argument("--log_interval", type=int, default=1)
|
||||
parser.add_argument("--save_interval", type=int, default=50)
|
||||
parser.add_argument('--local_rank', type=int, default=-1)
|
||||
parser.add_argument('--dim', default=512, type=int)
|
||||
parser.add_argument('--n_layers', default=8, type=int)
|
||||
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=1024, type=int)
|
||||
parser.add_argument('--use_moe', default=False, type=bool)
|
||||
parser.add_argument("--data_path", type=str, default="./dataset/r1_mix_1024.jsonl")
|
||||
parser.add_argument("--data_path", type=str, default="../dataset/r1_mix_1024.jsonl")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
lm_config = LMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, 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)
|
||||
tokens_per_iter = args.batch_size * lm_config.max_seq_len
|
||||
tokens_per_iter = args.batch_size * args.max_seq_len
|
||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||
|
||||
args.wandb_run_name = f"MiniMind-Distill-Reasoning-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
|
||||
@@ -191,7 +192,7 @@ if __name__ == "__main__":
|
||||
|
||||
model, tokenizer = init_model(lm_config)
|
||||
|
||||
train_ds = SFTDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
|
||||
train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
|
||||
train_sampler = DistributedSampler(train_ds) if ddp else None
|
||||
train_loader = DataLoader(
|
||||
train_ds,
|
||||
@@ -1,22 +1,23 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
__package__ = "trainer"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
import argparse
|
||||
import time
|
||||
import math
|
||||
import warnings
|
||||
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
|
||||
from torch import optim, nn
|
||||
from torch import optim
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
from model.model import MiniMindLM
|
||||
from model.LMConfig import LMConfig
|
||||
from model.dataset import SFTDataset
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from dataset.lm_dataset import SFTDataset
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
@@ -32,9 +33,9 @@ def get_lr(current_step, total_steps, lr):
|
||||
|
||||
def distillation_loss_fn(student_logits, teacher_logits, temperature=1.0, reduction='batchmean'):
|
||||
with torch.no_grad():
|
||||
teacher_probs = F.softmax(teacher_logits / temperature, dim=-1).detach()
|
||||
teacher_probs = F.softmax(teacher_logits / temperature, hidden_size=-1).detach()
|
||||
|
||||
student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)
|
||||
student_log_probs = F.log_softmax(student_logits / temperature, hidden_size=-1)
|
||||
|
||||
kl = F.kl_div(
|
||||
student_log_probs,
|
||||
@@ -98,7 +99,7 @@ def train_epoch(epoch, wandb, alpha=0.0, temperature=1.0):
|
||||
distill_loss = torch.tensor(0.0, device=args.device)
|
||||
|
||||
# 3) 总损失 = alpha * CE + (1-alpha) * Distill
|
||||
loss = alpha * ce_loss + (1 - alpha) * distill_loss
|
||||
loss = (alpha * ce_loss + (1 - alpha) * distill_loss) / args.accumulation_steps
|
||||
|
||||
scaler.scale(loss).backward()
|
||||
|
||||
@@ -135,20 +136,21 @@ def train_epoch(epoch, wandb, alpha=0.0, temperature=1.0):
|
||||
if (step + 1) % args.save_interval == 0 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.dim}{moe_path}.pth'
|
||||
ckp = f'{args.save_dir}/full_dist_{lm_config_student.hidden_size}{moe_path}.pth'
|
||||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||||
state_dict = model.module.state_dict()
|
||||
else:
|
||||
state_dict = model.state_dict()
|
||||
state_dict = {k: v.half() for k, v in state_dict.items()} # 半精度保存
|
||||
torch.save(state_dict, ckp)
|
||||
model.train()
|
||||
|
||||
|
||||
def init_student_model(lm_config):
|
||||
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
|
||||
model = MiniMindLM(lm_config)
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model/')
|
||||
model = MiniMindForCausalLM(lm_config)
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'./out/full_sft_{lm_config.dim}{moe_path}.pth'
|
||||
ckp = f'{args.save_dir}/full_sft_{lm_config.hidden_size}{moe_path}.pth'
|
||||
state_dict = torch.load(ckp, map_location=args.device)
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
Logger(f'学生模型(LLM)总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
||||
@@ -158,9 +160,9 @@ def init_student_model(lm_config):
|
||||
|
||||
|
||||
def init_teacher_model(lm_config):
|
||||
model = MiniMindLM(lm_config)
|
||||
model = MiniMindForCausalLM(lm_config)
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'./out/full_sft_{lm_config.dim}{moe_path}.pth'
|
||||
ckp = f'{args.save_dir}/full_sft_{lm_config.hidden_size}{moe_path}.pth'
|
||||
state_dict = torch.load(ckp, map_location=args.device)
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
Logger(f'教师模型(LLM)总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
||||
@@ -182,7 +184,7 @@ def init_distributed_mode():
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="MiniMind Full SFT")
|
||||
parser.add_argument("--out_dir", type=str, default="out")
|
||||
parser.add_argument("--out_dir", type=str, default="../out")
|
||||
parser.add_argument("--epochs", type=int, default=6)
|
||||
parser.add_argument("--batch_size", type=int, default=32)
|
||||
parser.add_argument("--learning_rate", type=float, default=5e-6)
|
||||
@@ -197,18 +199,18 @@ if __name__ == "__main__":
|
||||
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("--max_seq_len", type=int, default=512)
|
||||
parser.add_argument('--local_rank', type=int, default=-1)
|
||||
parser.add_argument("--data_path", type=str, default="./dataset/sft_data.jsonl")
|
||||
parser.add_argument("--data_path", type=str, default="../dataset/sft_xxx.jsonl")
|
||||
|
||||
args = parser.parse_args()
|
||||
# 定义学生模型和教师模型
|
||||
lm_config_student = LMConfig(dim=512, n_layers=8, max_seq_len=512)
|
||||
lm_config_teacher = LMConfig(dim=768, n_layers=16, max_seq_len=512)
|
||||
max_seq_len = lm_config_student.max_seq_len
|
||||
lm_config_student = MiniMindConfig(hidden_size=512, num_hidden_layers=8)
|
||||
lm_config_teacher = MiniMindConfig(hidden_size=768, num_hidden_layers=16)
|
||||
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)
|
||||
tokens_per_iter = args.batch_size * max_seq_len
|
||||
tokens_per_iter = args.batch_size * args.max_seq_len
|
||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||
|
||||
args.wandb_run_name = f"MiniMind-Dist-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
|
||||
@@ -239,7 +241,7 @@ if __name__ == "__main__":
|
||||
model, tokenizer = init_student_model(lm_config_student)
|
||||
teacher_model = init_teacher_model(lm_config_teacher)
|
||||
|
||||
train_ds = SFTDataset(args.data_path, tokenizer, max_length=max_seq_len)
|
||||
train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
|
||||
train_sampler = DistributedSampler(train_ds) if ddp else None
|
||||
train_loader = DataLoader(
|
||||
train_ds,
|
||||
@@ -1,23 +1,22 @@
|
||||
import os
|
||||
import platform
|
||||
import sys
|
||||
__package__ = "trainer"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
import argparse
|
||||
import time
|
||||
import math
|
||||
import warnings
|
||||
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
|
||||
from torch import optim, nn
|
||||
from torch import optim
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
from model.model import MiniMindLM
|
||||
from model.LMConfig import LMConfig
|
||||
from model.dataset import DPODataset
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from dataset.lm_dataset import DPODataset
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
@@ -35,17 +34,17 @@ def logits_to_probs(logits, labels):
|
||||
# logits shape: (batch_size, seq_len, vocab_size)
|
||||
# labels shape: (batch_size, seq_len)
|
||||
# probs shape: (batch_size, seq_len)
|
||||
log_probs = F.log_softmax(logits, dim=2)
|
||||
probs = torch.gather(log_probs, dim=2, index=labels.unsqueeze(2)).squeeze(-1)
|
||||
log_probs = F.log_softmax(logits, hidden_size=2)
|
||||
probs = torch.gather(log_probs, hidden_size=2, index=labels.unsqueeze(2)).squeeze(-1)
|
||||
return probs
|
||||
|
||||
|
||||
def dpo_loss(ref_probs, probs, mask, beta):
|
||||
# ref_probs 和 probs 都是 shape: (batch_size, seq_len)
|
||||
# https://github.com/jingyaogong/minimind/issues/298
|
||||
seq_lengths = mask.sum(dim=1, keepdim=True) # (batch_size, 1)
|
||||
ref_probs = (ref_probs * mask).sum(dim=1) / seq_lengths.squeeze()
|
||||
probs = (probs * mask).sum(dim=1) / seq_lengths.squeeze()
|
||||
seq_lengths = mask.sum(hidden_size=1, keephidden_size=True) # (batch_size, 1)
|
||||
ref_probs = (ref_probs * mask).sum(hidden_size=1) / seq_lengths.squeeze()
|
||||
probs = (probs * mask).sum(hidden_size=1) / seq_lengths.squeeze()
|
||||
|
||||
# 将 chosen 和 rejected 数据分开
|
||||
batch_size = ref_probs.shape[0]
|
||||
@@ -70,9 +69,9 @@ def train_epoch(epoch, wandb):
|
||||
y_rejected = batch['y_rejected'].to(args.device)
|
||||
mask_chosen = batch['mask_chosen'].to(args.device)
|
||||
mask_rejected = batch['mask_rejected'].to(args.device)
|
||||
x = torch.cat([x_chosen, x_rejected], dim=0)
|
||||
y = torch.cat([y_chosen, y_rejected], dim=0)
|
||||
mask = torch.cat([mask_chosen, mask_rejected], dim=0)
|
||||
x = torch.cat([x_chosen, x_rejected], hidden_size=0)
|
||||
y = torch.cat([y_chosen, y_rejected], hidden_size=0)
|
||||
mask = torch.cat([mask_chosen, mask_rejected], hidden_size=0)
|
||||
|
||||
lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate)
|
||||
for param_group in optimizer.param_groups:
|
||||
@@ -108,38 +107,38 @@ def train_epoch(epoch, wandb):
|
||||
args.epochs,
|
||||
step,
|
||||
iter_per_epoch,
|
||||
loss.item(),
|
||||
loss.item() * args.accumulation_steps,
|
||||
optimizer.param_groups[-1]['lr'],
|
||||
spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
|
||||
|
||||
if (wandb is not None) and (not ddp or dist.get_rank() == 0):
|
||||
wandb.log({"loss": loss,
|
||||
wandb.log({"loss": loss * args.accumulation_steps,
|
||||
"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):
|
||||
model.eval()
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'{args.save_dir}/rlhf_{lm_config.dim}{moe_path}.pth'
|
||||
ckp = f'{args.save_dir}/rlhf_{lm_config.hidden_size}{moe_path}.pth'
|
||||
|
||||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||||
state_dict = model.module.state_dict()
|
||||
else:
|
||||
state_dict = model.state_dict()
|
||||
|
||||
state_dict = {k: v.half() for k, v in state_dict.items()} # 半精度保存
|
||||
torch.save(state_dict, ckp)
|
||||
model.train()
|
||||
|
||||
|
||||
def init_model(lm_config):
|
||||
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
|
||||
model = MiniMindLM(lm_config)
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model/')
|
||||
model = MiniMindForCausalLM(lm_config)
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'./out/full_sft_{lm_config.dim}{moe_path}.pth'
|
||||
ckp = f'{args.save_dir}/full_sft_{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 = MiniMindLM(lm_config)
|
||||
ref_model = MiniMindForCausalLM(lm_config)
|
||||
ref_model.load_state_dict(state_dict, strict=False)
|
||||
ref_model.eval()
|
||||
ref_model.requires_grad_(False)
|
||||
@@ -165,7 +164,7 @@ def init_distributed_mode():
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="MiniMind RLHF")
|
||||
parser.add_argument("--out_dir", type=str, default="out")
|
||||
parser.add_argument("--out_dir", type=str, default="../out")
|
||||
parser.add_argument("--epochs", type=int, default=2)
|
||||
parser.add_argument("--batch_size", type=int, default=8)
|
||||
# sft阶段学习率为 「5e-6」->「5e-7」长度512，建议离线正负样本「概率」偏好对齐阶段lr <=「1e-8」长度3000，否则很容易遗忘训坏
|
||||
@@ -182,19 +181,19 @@ if __name__ == "__main__":
|
||||
parser.add_argument("--log_interval", type=int, default=100)
|
||||
parser.add_argument("--save_interval", type=int, default=100)
|
||||
parser.add_argument('--local_rank', type=int, default=-1)
|
||||
parser.add_argument('--dim', default=512, type=int)
|
||||
parser.add_argument('--n_layers', default=8, type=int)
|
||||
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=1024, type=int)
|
||||
parser.add_argument('--use_moe', default=False, type=bool)
|
||||
parser.add_argument("--data_path", type=str, default="./dataset/dpo.jsonl")
|
||||
parser.add_argument("--data_path", type=str, default="../dataset/dpo.jsonl")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
lm_config = LMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, 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)
|
||||
tokens_per_iter = args.batch_size * lm_config.max_seq_len
|
||||
tokens_per_iter = args.batch_size * args.max_seq_len
|
||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||
|
||||
args.wandb_run_name = f"MiniMind-Full-DPO-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
|
||||
@@ -223,7 +222,7 @@ if __name__ == "__main__":
|
||||
|
||||
model, ref_model, tokenizer = init_model(lm_config)
|
||||
|
||||
train_ds = DPODataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
|
||||
train_ds = DPODataset(args.data_path, tokenizer, max_length=args.max_seq_len)
|
||||
train_sampler = DistributedSampler(train_ds) if ddp else None
|
||||
train_loader = DataLoader(
|
||||
train_ds,
|
||||
@@ -1,23 +1,22 @@
|
||||
import os
|
||||
import platform
|
||||
import sys
|
||||
|
||||
__package__ = "trainer"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
import argparse
|
||||
import time
|
||||
import math
|
||||
import warnings
|
||||
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
|
||||
from torch import optim, nn
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
from model.model import MiniMindLM
|
||||
from model.LMConfig import LMConfig
|
||||
from model.dataset import SFTDataset
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from dataset.lm_dataset import SFTDataset
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
@@ -72,37 +71,46 @@ def train_epoch(epoch, wandb):
|
||||
args.epochs,
|
||||
step,
|
||||
iter_per_epoch,
|
||||
loss.item(),
|
||||
loss.item() * args.accumulation_steps,
|
||||
optimizer.param_groups[-1]['lr'],
|
||||
spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
|
||||
|
||||
if (wandb is not None) and (not ddp or dist.get_rank() == 0):
|
||||
wandb.log({"loss": loss,
|
||||
wandb.log({"loss": loss * args.accumulation_steps,
|
||||
"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):
|
||||
model.eval()
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'{args.save_dir}/full_sft_{lm_config.dim}{moe_path}.pth'
|
||||
|
||||
ckp = f'{args.save_dir}/full_sft_{lm_config.hidden_size}{moe_path}.pth'
|
||||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||||
state_dict = model.module.state_dict()
|
||||
else:
|
||||
state_dict = model.state_dict()
|
||||
|
||||
state_dict = {k: v.half() for k, v in state_dict.items()} # 半精度保存
|
||||
torch.save(state_dict, ckp)
|
||||
model.train()
|
||||
|
||||
|
||||
def init_model(lm_config):
|
||||
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
|
||||
model = MiniMindLM(lm_config)
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model')
|
||||
model = MiniMindForCausalLM(lm_config)
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'./out/pretrain_{lm_config.dim}{moe_path}.pth'
|
||||
ckp = f'{args.save_dir}/full_sft_{lm_config.hidden_size}{moe_path}.pth'
|
||||
state_dict = torch.load(ckp, map_location=args.device)
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
||||
|
||||
# 冻结所有参数
|
||||
for param in model.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
# 只解冻注意力机制中的投影层参数
|
||||
for name, param in model.named_parameters():
|
||||
if any(proj in name for proj in ['q_proj', 'k_proj', 'v_proj', 'o_proj']):
|
||||
param.requires_grad = True
|
||||
|
||||
Logger(f'LLM可训练总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
||||
model = model.to(args.device)
|
||||
return model, tokenizer
|
||||
|
||||
@@ -121,10 +129,10 @@ def init_distributed_mode():
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="MiniMind Full SFT")
|
||||
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=32)
|
||||
parser.add_argument("--learning_rate", type=float, default=5e-5)
|
||||
parser.add_argument("--out_dir", type=str, default="../out")
|
||||
parser.add_argument("--epochs", type=int, default=2)
|
||||
parser.add_argument("--batch_size", type=int, default=16)
|
||||
parser.add_argument("--learning_rate", type=float, default=5e-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")
|
||||
@@ -137,19 +145,20 @@ if __name__ == "__main__":
|
||||
parser.add_argument("--log_interval", type=int, default=100)
|
||||
parser.add_argument("--save_interval", type=int, default=100)
|
||||
parser.add_argument('--local_rank', type=int, default=-1)
|
||||
parser.add_argument('--dim', default=512, type=int)
|
||||
parser.add_argument('--n_layers', default=8, type=int)
|
||||
parser.add_argument('--hidden_size', default=768, type=int)
|
||||
parser.add_argument('--num_hidden_layers', default=16, 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/sft_mini_512.jsonl")
|
||||
parser.add_argument("--data_path", type=str, default="../dataset/sft_512.jsonl")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
lm_config = LMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, 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)
|
||||
tokens_per_iter = args.batch_size * lm_config.max_seq_len
|
||||
tokens_per_iter = args.batch_size * args.max_seq_len
|
||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||
|
||||
args.wandb_run_name = f"MiniMind-Full-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
|
||||
@@ -178,7 +187,7 @@ if __name__ == "__main__":
|
||||
|
||||
model, tokenizer = init_model(lm_config)
|
||||
|
||||
train_ds = SFTDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
|
||||
train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
|
||||
train_sampler = DistributedSampler(train_ds) if ddp else None
|
||||
train_loader = DataLoader(
|
||||
train_ds,
|
||||
@@ -1,7 +1,9 @@
|
||||
import os
|
||||
import platform
|
||||
import sys
|
||||
__package__ = "trainer"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
import argparse
|
||||
import random
|
||||
import time
|
||||
import math
|
||||
import warnings
|
||||
@@ -9,9 +11,8 @@ import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
from model.model import MiniMindLM
|
||||
from model.LMConfig import LMConfig
|
||||
from model.dataset import SFTDataset
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from dataset.lm_dataset import SFTDataset
|
||||
from model.model_lora import *
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
@@ -68,27 +69,27 @@ def train_epoch(epoch, wandb):
|
||||
args.epochs,
|
||||
step,
|
||||
iter_per_epoch,
|
||||
loss.item(),
|
||||
loss.item() * args.accumulation_steps,
|
||||
optimizer.param_groups[-1]['lr'],
|
||||
spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
|
||||
|
||||
if (wandb is not None) and (not ddp or dist.get_rank() == 0):
|
||||
wandb.log({"loss": loss,
|
||||
wandb.log({"loss": loss * args.accumulation_steps,
|
||||
"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):
|
||||
model.eval()
|
||||
# 【区别1】只保存lora权重即可
|
||||
save_lora(model, f'{args.save_dir}/lora/{args.lora_name}_{lm_config.dim}.pth')
|
||||
save_lora(model, f'{args.save_dir}/lora/{args.lora_name}_{lm_config.hidden_size}.pth')
|
||||
model.train()
|
||||
|
||||
|
||||
def init_model(lm_config):
|
||||
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
|
||||
model = MiniMindLM(lm_config)
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model/')
|
||||
model = MiniMindForCausalLM(lm_config)
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'./out/rlhf_{lm_config.dim}{moe_path}.pth'
|
||||
ckp = f'{args.save_dir}/rlhf_{lm_config.hidden_size}{moe_path}.pth'
|
||||
state_dict = torch.load(ckp, map_location=args.device)
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
return model.to(args.device), tokenizer
|
||||
@@ -108,7 +109,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("--out_dir", type=str, default="../out")
|
||||
parser.add_argument("--epochs", type=int, default=50)
|
||||
parser.add_argument("--batch_size", type=int, default=16)
|
||||
parser.add_argument("--learning_rate", type=float, default=5e-5)
|
||||
@@ -124,19 +125,19 @@ if __name__ == "__main__":
|
||||
parser.add_argument("--log_interval", type=int, default=100)
|
||||
parser.add_argument("--save_interval", type=int, default=1)
|
||||
parser.add_argument('--local_rank', type=int, default=-1)
|
||||
parser.add_argument('--dim', default=512, type=int)
|
||||
parser.add_argument('--n_layers', default=8, type=int)
|
||||
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_identity.jsonl")
|
||||
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 = LMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, 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)
|
||||
tokens_per_iter = args.batch_size * lm_config.max_seq_len
|
||||
tokens_per_iter = args.batch_size * args.max_seq_len
|
||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||
|
||||
ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
|
||||
@@ -182,7 +183,7 @@ if __name__ == "__main__":
|
||||
|
||||
# 只对 LoRA 参数进行优化
|
||||
optimizer = optim.AdamW(lora_params, lr=args.learning_rate)
|
||||
train_ds = SFTDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
|
||||
train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
|
||||
train_sampler = DistributedSampler(train_ds) if ddp else None
|
||||
train_loader = DataLoader(
|
||||
train_ds,
|
||||
@@ -1,23 +1,21 @@
|
||||
import os
|
||||
import platform
|
||||
import sys
|
||||
__package__ = "trainer"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
import argparse
|
||||
import time
|
||||
import math
|
||||
import warnings
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch import optim, nn
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.optim.lr_scheduler import CosineAnnealingLR
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from contextlib import nullcontext
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from model.model import MiniMindLM
|
||||
from model.LMConfig import LMConfig
|
||||
from model.dataset import PretrainDataset
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from dataset.lm_dataset import PretrainDataset
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
@@ -84,21 +82,22 @@ def train_epoch(epoch, wandb):
|
||||
if (step + 1) % args.save_interval == 0 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.dim}{moe_path}.pth'
|
||||
ckp = f'{args.save_dir}/pretrain_{lm_config.hidden_size}{moe_path}.pth'
|
||||
|
||||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||||
state_dict = model.module.state_dict()
|
||||
else:
|
||||
state_dict = model.state_dict()
|
||||
|
||||
state_dict = {k: v.half() for k, v in state_dict.items()} # 半精度保存
|
||||
torch.save(state_dict, ckp)
|
||||
model.train()
|
||||
|
||||
|
||||
def init_model(lm_config):
|
||||
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
|
||||
model = MiniMindLM(lm_config).to(args.device)
|
||||
Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model/')
|
||||
model = MiniMindForCausalLM(lm_config).to(args.device)
|
||||
Logger(f'LLM可训练总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
||||
return model, tokenizer
|
||||
|
||||
|
||||
@@ -117,7 +116,7 @@ def init_distributed_mode():
|
||||
# torchrun --nproc_per_node 2 1-pretrain.py
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="MiniMind Pretraining")
|
||||
parser.add_argument("--out_dir", type=str, default="out")
|
||||
parser.add_argument("--out_dir", type=str, default="../out")
|
||||
# 若要以最快速度实现zero则epochs设置为1轮；否则应当利用有限的数据训练2~6个epochs。
|
||||
parser.add_argument("--epochs", type=int, default=1)
|
||||
parser.add_argument("--batch_size", type=int, default=32)
|
||||
@@ -134,18 +133,18 @@ if __name__ == "__main__":
|
||||
parser.add_argument("--log_interval", type=int, default=100)
|
||||
parser.add_argument("--save_interval", type=int, default=100)
|
||||
parser.add_argument('--local_rank', type=int, default=-1)
|
||||
parser.add_argument('--dim', default=512, type=int)
|
||||
parser.add_argument('--n_layers', default=8, type=int)
|
||||
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/pretrain_hq.jsonl")
|
||||
parser.add_argument("--data_path", type=str, default="../dataset/pretrain_hq.jsonl")
|
||||
args = parser.parse_args()
|
||||
|
||||
lm_config = LMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, 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)
|
||||
tokens_per_iter = args.batch_size * lm_config.max_seq_len
|
||||
tokens_per_iter = args.batch_size * args.max_seq_len
|
||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||
|
||||
args.wandb_run_name = f"MiniMind-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
|
||||
@@ -175,7 +174,7 @@ if __name__ == "__main__":
|
||||
wandb = None
|
||||
|
||||
model, tokenizer = init_model(lm_config)
|
||||
train_ds = PretrainDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
|
||||
train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
|
||||
train_sampler = DistributedSampler(train_ds) if ddp else None
|
||||
train_loader = DataLoader(
|
||||
train_ds,
|
||||
Reference in New Issue
Block a user