[add] merge from jingyaogong
screnwei
parent
8495189bbd
commit
408babf307
|
|
@ -1,6 +1,7 @@
|
|||
/model/__pycache__
|
||||
/dataset
|
||||
/out
|
||||
.idea
|
||||
.DS_Store
|
||||
MiniMind2/
|
||||
.idea/
|
||||
MiniMind2/
|
||||
__pycache__/
|
||||
130
README.md
130
README.md
|
|
@ -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,26 @@
|
|||
### 👉**更新日志**
|
||||
|
||||
<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等热门三方生态。
|
||||
- 规范代码和目录结构。
|
||||
- 改动词表`<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 +236,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 +256,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 +302,8 @@ print(torch.cuda.is_available())
|
|||
|
||||
### 3.开始训练
|
||||
|
||||
目录位于`trainer`
|
||||
|
||||
**3.1 预训练（学知识）**
|
||||
|
||||
```bash
|
||||
|
|
@ -478,12 +509,12 @@ quality（当然也还不算high，提升数据质量无止尽）。
|
|||
|
||||
---
|
||||
|
||||
## Ⅷ 数据集下载
|
||||
## Ⅷ MiniMind训练数据集
|
||||
|
||||
> [!NOTE]
|
||||
> 2025-02-05后，开源MiniMind最终训练所用的所有数据集，因此无需再自行预处理大规模数据集，避免重复性的数据处理工作。
|
||||
|
||||
MiniMind训练数据集 ([ModelScope](https://www.modelscope.cn/datasets/gongjy/minimind_dataset/files) | [HuggingFace](https://huggingface.co/datasets/jingyaogong/minimind_dataset/tree/main))
|
||||
MiniMind训练数据集下载地址： [ModelScope](https://www.modelscope.cn/datasets/gongjy/minimind_dataset/files) | [HuggingFace](https://huggingface.co/datasets/jingyaogong/minimind_dataset/tree/main)
|
||||
|
||||
> 无需全部clone，可单独下载所需的文件
|
||||
|
||||
|
|
@ -640,6 +671,8 @@ Zero模型权重保存为 `full_sft_512_zero.pth`（见下文MiniMind模型文
|
|||
|
||||
## Ⅱ 主要训练步骤
|
||||
|
||||
> 所有训练脚本均 `cd ./trainer` 目录执行
|
||||
|
||||
### **1. 预训练(Pretrain)**:
|
||||
|
||||
LLM首先要学习的并非直接与人交流，而是让网络参数中充满知识的墨水，“墨水” 理论上喝的越饱越好，产生大量的对世界的知识积累。
|
||||
|
|
@ -677,6 +710,8 @@ python train_full_sft.py
|
|||
|
||||
## Ⅲ 其它训练步骤
|
||||
|
||||
> 所有训练脚本均 `cd ./trainer` 目录执行
|
||||
|
||||
### **3. 人类反馈强化学习(Reinforcement Learning from Human Feedback, RLHF)**
|
||||
|
||||
在前面的训练步骤中，模型已经具备了基本的对话能力，但是这样的能力完全基于单词接龙，缺少正反样例的激励。
|
||||
|
|
@ -1214,16 +1249,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 +1297,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]
|
||||
|
|
|
|||
141
README_en.md
141
README_en.md
|
|
@ -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,36 @@ 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.
|
||||
|
||||
• 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 +264,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 +330,8 @@ needs and GPU resources.
|
|||
|
||||
### 3. Start Training
|
||||
|
||||
The directory is located at `trainer`
|
||||
|
||||
**3.1 Pretraining (Learning Knowledge)**
|
||||
|
||||
```bash
|
||||
|
|
@ -517,15 +557,14 @@ Big respect!
|
|||
|
||||
---
|
||||
|
||||
## Ⅷ Dataset Download
|
||||
## Ⅷ MiniMind Dataset Download
|
||||
|
||||
> [!NOTE]
|
||||
> After `2025-02-05`, MiniMind’s open-source datasets for final training are provided, so there is no need for
|
||||
> you to preprocess large datasets by yourself anymore. This helps avoid redundant work.
|
||||
|
||||
MiniMind Training Datasets are available for download from:
|
||||
|
||||
Dataset ([ModelScope](https://www.modelscope.cn/datasets/gongjy/minimind_dataset/files) | [HuggingFace](https://huggingface.co/datasets/jingyaogong/minimind_dataset/tree/main))
|
||||
Available for download from:
|
||||
[ModelScope](https://www.modelscope.cn/datasets/gongjy/minimind_dataset/files) | [HuggingFace](https://huggingface.co/datasets/jingyaogong/minimind_dataset/tree/main)
|
||||
|
||||
> You don’t need to clone everything, just download the necessary files.
|
||||
|
||||
|
|
@ -696,6 +735,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 +785,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 +1404,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 +1455,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]
|
||||
|
|
|
|||
|
|
@ -35,9 +35,8 @@ class PretrainDataset(Dataset):
|
|||
sample = self.samples[index]
|
||||
|
||||
# 构建输入文本
|
||||
text = f"{self.tokenizer.bos_token}{str(sample['text'])}{self.tokenizer.eos_token}"
|
||||
encoding = self.tokenizer(
|
||||
text,
|
||||
str(sample['text']),
|
||||
max_length=self.max_length,
|
||||
padding='max_length',
|
||||
truncation=True,
|
||||
|
|
@ -58,8 +57,8 @@ class SFTDataset(Dataset):
|
|||
self.tokenizer = tokenizer
|
||||
self.max_length = max_length
|
||||
self.samples = self.load_data(jsonl_path)
|
||||
self.bos_id = tokenizer('<s>assistant', add_special_tokens=False).input_ids
|
||||
self.eos_id = tokenizer('</s>', add_special_tokens=False).input_ids
|
||||
self.bos_id = tokenizer('<|im_start|>assistant', add_special_tokens=False).input_ids
|
||||
self.eos_id = tokenizer('<|im_end|>', add_special_tokens=False).input_ids
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
|
|
@ -126,8 +125,8 @@ class DPODataset(Dataset):
|
|||
self.tokenizer = tokenizer
|
||||
self.max_length = max_length
|
||||
self.padding = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
|
||||
self.bos_id = tokenizer('<s>assistant', add_special_tokens=False).input_ids
|
||||
self.eos_id = tokenizer('</s>', add_special_tokens=False).input_ids
|
||||
self.bos_id = tokenizer('<|im_start|>assistant', add_special_tokens=False).input_ids
|
||||
self.eos_id = tokenizer('<|im_end|>', add_special_tokens=False).input_ids
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
self.data = []
|
||||
for line in f:
|
||||
|
|
@ -202,8 +201,8 @@ class RLAIFDataset(Dataset):
|
|||
self.tokenizer = tokenizer
|
||||
self.max_length = max_length
|
||||
self.samples = self.load_data(jsonl_path)
|
||||
self.bos_id = tokenizer('<s>assistant', add_special_tokens=False).input_ids
|
||||
self.eos_id = tokenizer('</s>', add_special_tokens=False).input_ids
|
||||
self.bos_id = tokenizer('<|im_start|>assistant', add_special_tokens=False).input_ids
|
||||
self.eos_id = tokenizer('<|im_end|>', add_special_tokens=False).input_ids
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
|
|
@ -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=512, 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)
|
||||
# 携带历史对话上下文条数
|
||||
# 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] == '<EFBFBD>') 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
File diff suppressed because one or more lines are too long
385
model/model.py
385
model/model.py
|
|
@ -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
|
||||
|
|
@ -18,7 +18,7 @@ class LoRA(nn.Module):
|
|||
return self.B(self.A(x))
|
||||
|
||||
|
||||
def apply_lora(model, rank=16):
|
||||
def apply_lora(model, rank=8):
|
||||
for name, module in model.named_modules():
|
||||
if isinstance(module, nn.Linear) and module.weight.shape[0] == module.weight.shape[1]:
|
||||
lora = LoRA(module.weight.shape[0], module.weight.shape[1], rank=rank).to(model.device)
|
||||
|
|
|
|||
|
|
@ -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 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
|
||||
File diff suppressed because it is too large
Load Diff
|
|
@ -4,7 +4,7 @@
|
|||
"add_prefix_space": false,
|
||||
"added_tokens_decoder": {
|
||||
"0": {
|
||||
"content": "<unk>",
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
|
|
@ -12,7 +12,7 @@
|
|||
"special": true
|
||||
},
|
||||
"1": {
|
||||
"content": "<s>",
|
||||
"content": "<|im_start|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
|
|
@ -20,7 +20,7 @@
|
|||
"special": true
|
||||
},
|
||||
"2": {
|
||||
"content": "</s>",
|
||||
"content": "<|im_end|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
|
|
@ -29,15 +29,15 @@
|
|||
}
|
||||
},
|
||||
"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 %}"
|
||||
}
|
||||
|
|
@ -1,19 +1,19 @@
|
|||
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 = []
|
||||
conversation_history = conversation_history_origin.copy()
|
||||
history_messages_num = 2 # 设置为偶数（Q+A），为0则每次不携带历史对话进行独立QA
|
||||
while True:
|
||||
conversation_history = conversation_history_origin.copy()
|
||||
query = input('[Q]: ')
|
||||
conversation_history.append({"role": "user", "content": query})
|
||||
response = client.chat.completions.create(
|
||||
model="minimind",
|
||||
messages=conversation_history,
|
||||
messages=conversation_history[-history_messages_num:],
|
||||
stream=stream
|
||||
)
|
||||
if not stream:
|
||||
|
|
|
|||
|
|
@ -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=False)
|
||||
|
||||
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'
|
||||
|
||||
# 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)
|
||||
|
|
|
|||
|
|
@ -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] == '<EFBFBD>') 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)
|
||||
|
|
|
|||
|
|
@ -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": "你是一个优秀的聊天机器人，总是给我正确的回应！"},
|
||||
|
|
|
|||
|
|
@ -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] == '<EFBFBD>') 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,37 @@ 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}/pretrain_{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} 百万')
|
||||
|
||||
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 +120,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 +136,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 +178,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,18 +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 random
|
||||
import time
|
||||
import math
|
||||
import warnings
|
||||
import torch
|
||||
from torch import optim, nn
|
||||
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_lora import *
|
||||
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
||||
from dataset.lm_dataset import SFTDataset
|
||||
from model.model_lora import load_lora, save_lora, apply_lora
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
|
@ -68,27 +72,29 @@ 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()
|
||||
lora_save_path = f'{args.save_dir}/lora/{args.lora_name}_{lm_config.hidden_size}.pth'
|
||||
os.makedirs(os.path.dirname(lora_save_path), exist_ok=True)
|
||||
# 【区别1】只保存lora权重即可
|
||||
save_lora(model, f'{args.save_dir}/lora/{args.lora_name}_{lm_config.dim}.pth')
|
||||
save_lora(model, lora_save_path)
|
||||
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}/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)
|
||||
return model.to(args.device), tokenizer
|
||||
|
|
@ -108,10 +114,10 @@ def init_distributed_mode():
|
|||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="MiniMind SFT with LoRA")
|
||||
parser.add_argument("--out_dir", type=str, default="out")
|
||||
parser.add_argument("--epochs", type=int, default=50)
|
||||
parser.add_argument("--batch_size", type=int, default=16)
|
||||
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=10)
|
||||
parser.add_argument("--batch_size", type=int, default=32)
|
||||
parser.add_argument("--learning_rate", type=float, default=1e-4)
|
||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16")
|
||||
parser.add_argument("--use_wandb", action="store_true")
|
||||
|
|
@ -122,21 +128,22 @@ if __name__ == "__main__":
|
|||
parser.add_argument("--grad_clip", type=float, default=1.0)
|
||||
parser.add_argument("--warmup_iters", type=int, default=0)
|
||||
parser.add_argument("--log_interval", type=int, default=100)
|
||||
parser.add_argument("--save_interval", type=int, default=1)
|
||||
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/lora_identity.jsonl")
|
||||
parser.add_argument("--lora_name", type=str, default="lora_identity", help="根据任务保存成lora_(英文/医学/心理...)")
|
||||
parser.add_argument("--data_path", type=str, default="../dataset/lora_medical.jsonl")
|
||||
parser.add_argument("--lora_name", type=str, default="lora_medical", 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 +189,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,
|
||||
Loading…
Reference in New Issue