update readme

Former-commit-id: 561481a8008fde5a3273558460193864a09866ed

.github/ISSUE_TEMPLATE/bug-report.yml

@@ -0,0 +1,58 @@
+name: "\U0001F41B Bug / Help"
+description: Create a report to help us improve the LLaMA Factory
+body:
+  - type: checkboxes
+    id: reminder
+    attributes:
+      label: Reminder
+      description: |
+        Please ensure you have read the README carefully and searched the existing issues.
+        请确保您已经认真阅读了 README 并且搜索过现有的 Issue。
+
+      options:
+        - label: I have read the README and searched the existing issues.
+          required: true
+
+  - type: textarea
+    id: reproduction
+    validations:
+      required: true
+    attributes:
+      label: Reproduction
+      description: |
+        Please provide code snippets, error messages and stack traces that reproduces the problem.
+        请提供运行参数，错误信息以及异常堆栈以便于我们复现该问题。
+        Remember to use Markdown tags to correctly format your code.
+        请合理使用 Markdown 标签来格式化您的文本。
+
+      placeholder: |
+        python src/train_bash.py ...
+
+  - type: textarea
+    id: expected-behavior
+    validations:
+      required: false
+    attributes:
+      label: Expected behavior
+      description: |
+        Please provide a clear and concise description of what you would expect to happen.
+        请提供您原本的目的，即这段代码的期望行为。
+
+  - type: textarea
+    id: system-info
+    validations:
+      required: false
+    attributes:
+      label: System Info
+      description: |
+        Please share your system info with us. You can run the command **transformers-cli env** and copy-paste its output below.
+        请提供您的系统信息。您可以在命令行运行 **transformers-cli env** 并将其输出复制到该文本框中。
+
+      placeholder: transformers version, platform, python version, ...
+
+  - type: textarea
+    id: others
+    validations:
+      required: false
+    attributes:
+      label: Others

CODE_OF_CONDUCT.md

@@ -0,0 +1,128 @@
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, religion, or sexual identity
+and orientation.
+
+We pledge to act and interact in ways that contribute to an open, welcoming,
+diverse, inclusive, and healthy community.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment for our
+community include:
+
+* Demonstrating empathy and kindness toward other people
+* Being respectful of differing opinions, viewpoints, and experiences
+* Giving and gracefully accepting constructive feedback
+* Accepting responsibility and apologizing to those affected by our mistakes,
+  and learning from the experience
+* Focusing on what is best not just for us as individuals, but for the
+  overall community
+
+Examples of unacceptable behavior include:
+
+* The use of sexualized language or imagery, and sexual attention or
+  advances of any kind
+* Trolling, insulting or derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or email
+  address, without their explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Enforcement Responsibilities
+
+Community leaders are responsible for clarifying and enforcing our standards of
+acceptable behavior and will take appropriate and fair corrective action in
+response to any behavior that they deem inappropriate, threatening, offensive,
+or harmful.
+
+Community leaders have the right and responsibility to remove, edit, or reject
+comments, commits, code, wiki edits, issues, and other contributions that are
+not aligned to this Code of Conduct, and will communicate reasons for moderation
+decisions when appropriate.
+
+## Scope
+
+This Code of Conduct applies within all community spaces, and also applies when
+an individual is officially representing the community in public spaces.
+Examples of representing our community include using an official e-mail address,
+posting via an official social media account, or acting as an appointed
+representative at an online or offline event.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the community leaders responsible for enforcement at
+`hoshihiyouga AT gmail DOT com`.
+All complaints will be reviewed and investigated promptly and fairly.
+
+All community leaders are obligated to respect the privacy and security of the
+reporter of any incident.
+
+## Enforcement Guidelines
+
+Community leaders will follow these Community Impact Guidelines in determining
+the consequences for any action they deem in violation of this Code of Conduct:
+
+### 1. Correction
+
+**Community Impact**: Use of inappropriate language or other behavior deemed
+unprofessional or unwelcome in the community.
+
+**Consequence**: A private, written warning from community leaders, providing
+clarity around the nature of the violation and an explanation of why the
+behavior was inappropriate. A public apology may be requested.
+
+### 2. Warning
+
+**Community Impact**: A violation through a single incident or series
+of actions.
+
+**Consequence**: A warning with consequences for continued behavior. No
+interaction with the people involved, including unsolicited interaction with
+those enforcing the Code of Conduct, for a specified period of time. This
+includes avoiding interactions in community spaces as well as external channels
+like social media. Violating these terms may lead to a temporary or
+permanent ban.
+
+### 3. Temporary Ban
+
+**Community Impact**: A serious violation of community standards, including
+sustained inappropriate behavior.
+
+**Consequence**: A temporary ban from any sort of interaction or public
+communication with the community for a specified period of time. No public or
+private interaction with the people involved, including unsolicited interaction
+with those enforcing the Code of Conduct, is allowed during this period.
+Violating these terms may lead to a permanent ban.
+
+### 4. Permanent Ban
+
+**Community Impact**: Demonstrating a pattern of violation of community
+standards, including sustained inappropriate behavior,  harassment of an
+individual, or aggression toward or disparagement of classes of individuals.
+
+**Consequence**: A permanent ban from any sort of public interaction within
+the community.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 2.0, available at
+https://www.contributor-covenant.org/version/2/0/code_of_conduct.html.
+
+Community Impact Guidelines were inspired by [Mozilla's code of conduct
+enforcement ladder](https://github.com/mozilla/diversity).
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see the FAQ at
+https://www.contributor-covenant.org/faq. Translations are available at
+https://www.contributor-covenant.org/translations.

README.md

@@ -6,7 +6,9 @@
 [![PyPI](https://img.shields.io/pypi/v/llmtuner)](https://pypi.org/project/llmtuner/)
 [![Downloads](https://static.pepy.tech/badge/llmtuner)](https://pypi.org/project/llmtuner/)
 [![GitHub pull request](https://img.shields.io/badge/PRs-welcome-blue)](https://github.com/hiyouga/LLaMA-Factory/pulls)
-[![Discord](https://dcbadge.vercel.app/api/server/e73gccsSd?compact=true&style=flat)](https://discord.gg/e73gccsSd)
+[![Discord](https://dcbadge.vercel.app/api/server/c2EPEt5NU?compact=true&style=flat)](https://discord.gg/c2EPEt5NU)
+[![Spaces](https://img.shields.io/badge/🤗-Open%20In%20Spaces-blue)](https://huggingface.co/spaces/hiyouga/LLaMA-Board)
+[![Studios](https://img.shields.io/badge/ModelScope-Open%20In%20Studios-blue)](https://modelscope.cn/studios/hiyouga/LLaMA-Board)
 
 👋 Join our [WeChat](assets/wechat.jpg).
 
@@ -14,12 +16,39 @@
 
 ## LLaMA Board: A One-stop Web UI for Getting Started with LLaMA Factory
 
-Launch **LLaMA Board** via `CUDA_VISIBLE_DEVICES=0 python src/train_web.py`. (multiple GPUs are not supported yet)
+Preview LLaMA Board at **[🤗 Spaces](https://huggingface.co/spaces/hiyouga/LLaMA-Board)** or **[ModelScope](https://modelscope.cn/studios/hiyouga/LLaMA-Board)**.
+
+Launch LLaMA Board via `CUDA_VISIBLE_DEVICES=0 python src/train_web.py`. (multiple GPUs are not supported yet in this mode)
 
 Here is an example of altering the self-cognition of an instruction-tuned language model within 10 minutes on a single GPU.
 
 https://github.com/hiyouga/LLaMA-Factory/assets/16256802/6ba60acc-e2e2-4bec-b846-2d88920d5ba1
 
+## Table of Contents
+
+- [Benchmark](#benchmark)
+- [Changelog](#changelog)
+- [Supported Models](#supported-models)
+- [Supported Training Approaches](#supported-training-approaches)
+- [Provided Datasets](#provided-datasets)
+- [Requirement](#requirement)
+- [Getting Started](#getting-started)
+- [Projects using LLaMA Factory](#projects-using-llama-factory)
+- [License](#license)
+- [Citation](#citation)
+- [Acknowledgement](#acknowledgement)
+
+## Benchmark
+
+Compared to ChatGLM's [P-Tuning](https://github.com/THUDM/ChatGLM2-6B/tree/main/ptuning), LLaMA-Factory's LoRA tuning offers up to **3.7 times faster** training speed with a better Rouge score on the advertising text generation task. By leveraging 4-bit quantization technique, LLaMA-Factory's QLoRA further improves the efficiency regarding the GPU memory.
+
+![benchmark](assets/benchmark.svg)
+
+- **Training Speed**: the number of training samples processed per second during the training. (bs=4, cutoff_len=1024)
+- **Rouge Score**: Rouge-2 score on the development set of the [advertising text generation](https://aclanthology.org/D19-1321.pdf) task. (bs=4, cutoff_len=1024)
+- **GPU Memory**: Peak GPU memory usage in 4-bit quantized training. (bs=1, cutoff_len=1024)
+- We adopt `pre_seq_len=128` for ChatGLM's P-Tuning and `lora_rank=32` for LLaMA-Factory's LoRA tuning.
+
 ## Changelog
 
 [23/10/21] We supported **[NEFTune](https://arxiv.org/abs/2310.05914)** trick for fine-tuning. Try `--neft_alpha` argument to activate NEFTune, e.g., `--neft_alpha 5`.
@@ -57,7 +86,7 @@ https://github.com/hiyouga/LLaMA-Factory/assets/16256802/6ba60acc-e2e2-4bec-b846
 | [BLOOM](https://huggingface.co/bigscience/bloom)         | 560M/1.1B/1.7B/3B/7.1B/176B | query_key_value   | -         |
 | [BLOOMZ](https://huggingface.co/bigscience/bloomz)       | 560M/1.1B/1.7B/3B/7.1B/176B | query_key_value   | -         |
 | [ChatGLM3](https://github.com/THUDM/ChatGLM3)            | 6B                          | query_key_value   | chatglm3  |
-| [Falcon](https://huggingface.co/tiiuae/falcon-7b)        | 7B/40B/180B                 | query_key_value   | -         |
+| [Falcon](https://huggingface.co/tiiuae/falcon-7b)        | 7B/40B/180B                 | query_key_value   | falcon    |
 | [InternLM](https://github.com/InternLM/InternLM)         | 7B/20B                      | q_proj,v_proj     | intern    |
 | [LLaMA](https://github.com/facebookresearch/llama)       | 7B/13B/33B/65B              | q_proj,v_proj     | -         |
 | [LLaMA-2](https://huggingface.co/meta-llama)             | 7B/13B/70B                  | q_proj,v_proj     | llama2    |
@@ -71,7 +100,7 @@ https://github.com/hiyouga/LLaMA-Factory/assets/16256802/6ba60acc-e2e2-4bec-b846
 >
 > For the "base" models, the `--template` argument can be chosen from `default`, `alpaca`, `vicuna` etc. But make sure to use the **corresponding template** for the "chat" models.
 
-Please refer to [template.py](src/llmtuner/extras/template.py) for a full list of models we supported.
+Please refer to [constants.py](src/llmtuner/extras/constants.py) for a full list of models we supported.
 
 ## Supported Training Approaches
 
@@ -79,9 +108,9 @@ Please refer to [template.py](src/llmtuner/extras/template.py) for a full list o
 | ---------------------- | ------------------ | ------------------ | ------------------ | ------------------ |
 | Pre-Training           | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: |
 | Supervised Fine-Tuning | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: |
-| Reward Modeling        |                    |                    | :white_check_mark: | :white_check_mark: |
-| PPO Training           |                    |                    | :white_check_mark: | :white_check_mark: |
-| DPO Training           | :white_check_mark: |                    | :white_check_mark: | :white_check_mark: |
+| Reward Modeling        | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: |
+| PPO Training           | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: |
+| DPO Training           | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: |
 
 > [!NOTE]
 > Use `--quantization_bit 4/8` argument to enable QLoRA.
@@ -122,6 +151,7 @@ Please refer to [template.py](src/llmtuner/extras/template.py) for a full list o
 - [OpenPlatypus (en)](https://huggingface.co/datasets/garage-bAInd/Open-Platypus)
 - [CodeAlpaca 20k (en)](https://huggingface.co/datasets/sahil2801/CodeAlpaca-20k)
 - [Alpaca CoT (multilingual)](https://huggingface.co/datasets/QingyiSi/Alpaca-CoT)
+- [OpenOrca (en)](https://huggingface.co/datasets/Open-Orca/OpenOrca)
 - [MathInstruct (en)](https://huggingface.co/datasets/TIGER-Lab/MathInstruct)
 - [Firefly 1.1M (zh)](https://huggingface.co/datasets/YeungNLP/firefly-train-1.1M)
 - [Web QA (zh)](https://huggingface.co/datasets/suolyer/webqa)
@@ -158,7 +188,7 @@ huggingface-cli login
 - Python 3.8+ and PyTorch 1.13.1+
 - 🤗Transformers, Datasets, Accelerate, PEFT and TRL
 - sentencepiece, protobuf and tiktoken
-- fire, jieba, rouge-chinese and nltk (used at evaluation and predict)
+- jieba, rouge-chinese and nltk (used at evaluation and predict)
 - gradio and matplotlib (used in web UI)
 - uvicorn, fastapi and sse-starlette (used in API)
 
@@ -284,6 +314,8 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
     --per_device_train_batch_size 2 \
     --gradient_accumulation_steps 4 \
     --lr_scheduler_type cosine \
+    --top_k 0 \
+    --top_p 0.9 \
     --logging_steps 10 \
     --save_steps 1000 \
     --learning_rate 1e-5 \
@@ -292,6 +324,9 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
     --fp16
 ```
 
+> [!WARNING]
+> Use `--per_device_train_batch_size=1` for LLaMA-2 models in fp16 PPO training.
+
 #### DPO Training
 
 ```bash
@@ -387,7 +422,7 @@ deepspeed --num_gpus 8 --master_port=9901 src/train_bash.py \
 
 </details>
 
-### Export model
+### Merge LoRA weights and export model
 
 ```bash
 python src/export_model.py \
@@ -408,7 +443,7 @@ python src/api_demo.py \
     --checkpoint_dir path_to_checkpoint
 ```
 
-> [!NOTE]
+> [!TIP]
 > Visit `http://localhost:8000/docs` for API documentation.
 
 ### CLI Demo
@@ -460,10 +495,14 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
     --output_dir path_to_predict_result \
     --per_device_eval_batch_size 8 \
     --max_samples 100 \
-    --predict_with_generate
+    --predict_with_generate \
+    --fp16
 ```
 
-> [!NOTE]
+> [!WARNING]
+> Use `--per_device_train_batch_size=1` for LLaMA-2 models in fp16 predict.
+
+> [!TIP]
 > We recommend using `--per_device_eval_batch_size=1` and `--max_target_length 128` at 4/8-bit predict.
 
 ## Projects using LLaMA Factory
@@ -473,6 +512,9 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
 - **[Sunsimiao](https://github.com/thomas-yanxin/Sunsimiao)**: A large language model specialized in Chinese medical domain, based on Baichuan-7B and ChatGLM-6B.
 - **[CareGPT](https://github.com/WangRongsheng/CareGPT)**: A series of large language models for Chinese medical domain, based on LLaMA2-7B and Baichuan-13B.
 
+> [!TIP]
+> If you have a project that should be incorporated, please contact via email or create a pull request.
+
 ## License
 
 This repository is licensed under the [Apache-2.0 License](LICENSE).

README_zh.md

@@ -6,7 +6,9 @@
 [![PyPI](https://img.shields.io/pypi/v/llmtuner)](https://pypi.org/project/llmtuner/)
 [![Downloads](https://static.pepy.tech/badge/llmtuner)](https://pypi.org/project/llmtuner/)
 [![GitHub pull request](https://img.shields.io/badge/PRs-welcome-blue)](https://github.com/hiyouga/LLaMA-Factory/pulls)
-[![Discord](https://dcbadge.vercel.app/api/server/e73gccsSd?compact=true&style=flat)](https://discord.gg/e73gccsSd)
+[![Discord](https://dcbadge.vercel.app/api/server/c2EPEt5NU?compact=true&style=flat)](https://discord.gg/c2EPEt5NU)
+[![Spaces](https://img.shields.io/badge/🤗-Open%20In%20Spaces-blue)](https://huggingface.co/spaces/hiyouga/LLaMA-Board)
+[![Studios](https://img.shields.io/badge/ModelScope-Open%20In%20Studios-blue)](https://modelscope.cn/studios/hiyouga/LLaMA-Board)
 
 👋 加入我们的[微信群](assets/wechat.jpg)。
 
@@ -14,12 +16,39 @@
 
 ## LLaMA Board: 通过一站式网页界面快速上手 LLaMA Factory
 
-使用 `CUDA_VISIBLE_DEVICES=0 python src/train_web.py` 启动 **LLaMA Board**。（该界面目前仅支持单卡训练）
+通过 **[🤗 Spaces](https://huggingface.co/spaces/hiyouga/LLaMA-Board)** 或 **[ModelScope](https://modelscope.cn/studios/hiyouga/LLaMA-Board)** 预览 LLaMA Board。
+
+使用 `CUDA_VISIBLE_DEVICES=0 python src/train_web.py` 启动 LLaMA Board。（该模式目前仅支持单卡训练）
 
 下面是使用单张 GPU 在 10 分钟内更改对话式大型语言模型自我认知的示例。
 
 https://github.com/hiyouga/LLaMA-Factory/assets/16256802/6ba60acc-e2e2-4bec-b846-2d88920d5ba1
 
+## 目录
+
+- [性能指标](#性能指标)
+- [更新日志](#更新日志)
+- [模型](#模型)
+- [训练方法](#训练方法)
+- [数据集](#数据集)
+- [软件依赖](#软件依赖)
+- [如何使用](#如何使用)
+- [使用了 LLaMA Factory 的项目](#使用了-llama-factory-的项目)
+- [协议](#协议)
+- [引用](#引用)
+- [致谢](#致谢)
+
+## 性能指标
+
+与 ChatGLM 官方的 [P-Tuning](https://github.com/THUDM/ChatGLM2-6B/tree/main/ptuning) 微调相比，LLaMA-Factory 的 LoRA 微调提供了 **3.7 倍**的加速比，同时在广告文案生成任务上取得了更高的 Rouge 分数。结合 4 比特量化技术，LLaMA-Factory 的 QLoRA 微调进一步降低了 GPU 显存消耗。
+
+![benchmark](assets/benchmark.svg)
+
+- **Training Speed**: 训练阶段每秒处理的样本数量。（批处理大小=4，截断长度=1024）
+- **Rouge Score**: [广告文案生成](https://aclanthology.org/D19-1321.pdf)任务验证集上的 Rouge-2 分数。（批处理大小=4，截断长度=1024）
+- **GPU Memory**: 4 比特量化训练的 GPU 显存峰值。（批处理大小=1，截断长度=1024）
+- 我们在 ChatGLM 的 P-Tuning 中采用 `pre_seq_len=128`，在 LLaMA-Factory 的 LoRA 微调中采用 `lora_rank=32`。
+
 ## 更新日志
 
 [23/10/21] 我们支持了 **[NEFTune](https://arxiv.org/abs/2310.05914)** 训练技巧。请使用 `--neft_alpha` 参数启用 NEFTune，例如 `--neft_alpha 5`。
@@ -57,7 +86,7 @@ https://github.com/hiyouga/LLaMA-Factory/assets/16256802/6ba60acc-e2e2-4bec-b846
 | [BLOOM](https://huggingface.co/bigscience/bloom)         | 560M/1.1B/1.7B/3B/7.1B/176B | query_key_value   | -         |
 | [BLOOMZ](https://huggingface.co/bigscience/bloomz)       | 560M/1.1B/1.7B/3B/7.1B/176B | query_key_value   | -         |
 | [ChatGLM3](https://github.com/THUDM/ChatGLM3)            | 6B                          | query_key_value   | chatglm3  |
-| [Falcon](https://huggingface.co/tiiuae/falcon-7b)        | 7B/40B/180B                 | query_key_value   | -         |
+| [Falcon](https://huggingface.co/tiiuae/falcon-7b)        | 7B/40B/180B                 | query_key_value   | falcon    |
 | [InternLM](https://github.com/InternLM/InternLM)         | 7B/20B                      | q_proj,v_proj     | intern    |
 | [LLaMA](https://github.com/facebookresearch/llama)       | 7B/13B/33B/65B              | q_proj,v_proj     | -         |
 | [LLaMA-2](https://huggingface.co/meta-llama)             | 7B/13B/70B                  | q_proj,v_proj     | llama2    |
@@ -71,7 +100,7 @@ https://github.com/hiyouga/LLaMA-Factory/assets/16256802/6ba60acc-e2e2-4bec-b846
 >
 > 对于所有“基座”（Base）模型，`--template` 参数可以是 `default`, `alpaca`, `vicuna` 等任意值。但“对话”（Chat）模型请务必使用**对应的模板**。
 
-项目所支持模型的完整列表请参阅 [template.py](src/llmtuner/extras/template.py)。
+项目所支持模型的完整列表请参阅 [constants.py](src/llmtuner/extras/constants.py)。
 
 ## 训练方法
 
@@ -79,9 +108,9 @@ https://github.com/hiyouga/LLaMA-Factory/assets/16256802/6ba60acc-e2e2-4bec-b846
 | ---------------------- | ------------------ | ------------------ | ------------------ | ------------------ |
 | 预训练                 | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: |
 | 指令监督微调            | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: |
-| 奖励模型训练            |                    |                    | :white_check_mark: | :white_check_mark: |
-| PPO 训练               |                    |                    | :white_check_mark: | :white_check_mark: |
-| DPO 训练               | :white_check_mark: |                    | :white_check_mark: | :white_check_mark: |
+| 奖励模型训练            | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: |
+| PPO 训练               | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: |
+| DPO 训练               | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: |
 
 > [!NOTE]
 > 请使用 `--quantization_bit 4/8` 参数来启用 QLoRA 训练。
@@ -122,6 +151,7 @@ https://github.com/hiyouga/LLaMA-Factory/assets/16256802/6ba60acc-e2e2-4bec-b846
 - [OpenPlatypus (en)](https://huggingface.co/datasets/garage-bAInd/Open-Platypus)
 - [CodeAlpaca 20k (en)](https://huggingface.co/datasets/sahil2801/CodeAlpaca-20k)
 - [Alpaca CoT (multilingual)](https://huggingface.co/datasets/QingyiSi/Alpaca-CoT)
+- [OpenOrca (en)](https://huggingface.co/datasets/Open-Orca/OpenOrca)
 - [MathInstruct (en)](https://huggingface.co/datasets/TIGER-Lab/MathInstruct)
 - [Firefly 1.1M (zh)](https://huggingface.co/datasets/YeungNLP/firefly-train-1.1M)
 - [Web QA (zh)](https://huggingface.co/datasets/suolyer/webqa)
@@ -158,7 +188,7 @@ huggingface-cli login
 - Python 3.8+ 和 PyTorch 1.13.1+
 - 🤗Transformers, Datasets, Accelerate, PEFT 和 TRL
 - sentencepiece, protobuf 和 tiktoken
-- fire, jieba, rouge-chinese 和 nltk (用于评估及预测)
+- jieba, rouge-chinese 和 nltk (用于评估及预测)
 - gradio 和 matplotlib (用于网页端交互)
 - uvicorn, fastapi 和 sse-starlette (用于 API)
 
@@ -284,13 +314,19 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
     --per_device_train_batch_size 2 \
     --gradient_accumulation_steps 4 \
     --lr_scheduler_type cosine \
+    --top_k 0 \
+    --top_p 0.9 \
     --logging_steps 10 \
     --save_steps 1000 \
     --learning_rate 1e-5 \
     --num_train_epochs 1.0 \
-    --plot_loss
+    --plot_loss \
+    --fp16
 ```
 
+> [!WARNING]
+> 如果使用 fp16 精度进行 LLaMA-2 模型的 PPO 训练，请使用 `--per_device_train_batch_size=1`。
+
 #### DPO 训练
 
 ```bash
@@ -386,7 +422,7 @@ deepspeed --num_gpus 8 --master_port=9901 src/train_bash.py \
 
 </details>
 
-### 导出微调后的完整模型
+### 合并 LoRA 权重并导出完整模型
 
 ```bash
 python src/export_model.py \
@@ -407,7 +443,7 @@ python src/api_demo.py \
     --checkpoint_dir path_to_checkpoint
 ```
 
-> [!NOTE]
+> [!TIP]
 > 关于 API 文档请见 `http://localhost:8000/docs`。
 
 ### 命令行测试
@@ -459,10 +495,14 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
     --output_dir path_to_predict_result \
     --per_device_eval_batch_size 8 \
     --max_samples 100 \
-    --predict_with_generate
+    --predict_with_generate \
+    --fp16
 ```
 
-> [!NOTE]
+> [!WARNING]
+> 如果使用 fp16 精度进行 LLaMA-2 模型的预测，请使用 `--per_device_eval_batch_size=1`。
+
+> [!TIP]
 > 我们建议在量化模型的预测中使用 `--per_device_eval_batch_size=1` 和 `--max_target_length 128`。
 
 ## 使用了 LLaMA Factory 的项目
@@ -472,6 +512,9 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
 - **[Sunsimiao](https://github.com/thomas-yanxin/Sunsimiao)**: 孙思邈中文医疗大模型 Sumsimiao，基于 Baichuan-7B 和 ChatGLM-6B 在中文医疗数据上微调而得。
 - **[CareGPT](https://github.com/WangRongsheng/CareGPT)**: 医疗大模型项目 CareGPT，基于 LLaMA2-7B 和 Baichuan-13B 在中文医疗数据上微调而得。
 
+> [!TIP]
+> 如果您有项目希望添加至上述列表，请通过邮件联系或者创建一个 PR。
+
 ## 协议
 
 本仓库的代码依照 [Apache-2.0](LICENSE) 协议开源。

data/belle_multiturn/belle_multiturn.py

@@ -24,9 +24,7 @@ class BelleMultiturn(datasets.GeneratorBasedBuilder):
 
     def _info(self):
         features = datasets.Features({
-            "instruction": datasets.Value("string"),
-            "output": datasets.Value("string"),
-            "history": datasets.Sequence(datasets.Sequence(datasets.Value("string")))
+            "conversations": [{"from": datasets.Value("string"), "value": datasets.Value("string")}]
         })
         return datasets.DatasetInfo(
             description=_DESCRIPTION,
@@ -51,6 +49,7 @@ class BelleMultiturn(datasets.GeneratorBasedBuilder):
         with open(filepath, "r", encoding="utf-8") as f:
             for key, row in enumerate(f):
                 data = json.loads(row)
+                conversations = []
                 prompt = data["instruction"].strip()
                 response = data["output"].strip()
 
@@ -58,7 +57,8 @@ class BelleMultiturn(datasets.GeneratorBasedBuilder):
                 human_idx = prompt.rfind("Human:")
                 query = prompt[human_idx+6:assist_idx].strip()
                 prompt = prompt[:human_idx].strip()
-                history = []
+                conversations.insert(0, {"from": "gpt", "value": response})
+                conversations.insert(0, {"from": "human", "value": query})
 
                 while prompt.rfind("Assistant:") != -1:
                     assist_idx = prompt.rfind("Assistant:")
@@ -66,13 +66,10 @@ class BelleMultiturn(datasets.GeneratorBasedBuilder):
                     if human_idx != -1:
                         old_query = prompt[human_idx+6:assist_idx].strip()
                         old_resp = prompt[assist_idx+10:].strip()
-                        history.insert(0, (old_query, old_resp))
+                        conversations.insert(0, {"from": "gpt", "value": old_resp})
+                        conversations.insert(0, {"from": "human", "value": old_query})
                     else:
                         break
                     prompt = prompt[:human_idx].strip()
 
-                yield key, {
-                    "instruction": query,
-                    "output": response,
-                    "history": history
-                }
+                yield key, {"conversations": conversations}

data/ultra_chat/ultra_chat.py

@@ -66,6 +66,4 @@ class UltraChat(datasets.GeneratorBasedBuilder):
                         "from": "human" if i % 2 == 0 else "gpt",
                         "value": content[i]
                     } for i in range(len(content))]
-                    yield key, {
-                        "conversations": conversations
-                    }
+                    yield key, {"conversations": conversations}

data/wiki_demo.txt

@@ -1,50 +0,0 @@
-Machine learning (ML) is a field devoted to understanding and building methods that let machines "learn" – that is, methods that leverage data to improve computer performance on some set of tasks.
-Machine learning algorithms build a model based on sample data, known as training data, in order to make predictions or decisions without being explicitly programmed to do so. Machine learning algorithms are used in a wide variety of applications, such as in medicine, email filtering, speech recognition, agriculture, and computer vision, where it is difficult or unfeasible to develop conventional algorithms to perform the needed tasks.
-A subset of machine learning is closely related to computational statistics, which focuses on making predictions using computers, but not all machine learning is statistical learning. The study of mathematical optimization delivers methods, theory and application domains to the field of machine learning. Data mining is a related field of study, focusing on exploratory data analysis through unsupervised learning.
-Some implementations of machine learning use data and neural networks in a way that mimics the working of a biological brain.
-In its application across business problems, machine learning is also referred to as predictive analytics.
-Learning algorithms work on the basis that strategies, algorithms, and inferences that worked well in the past are likely to continue working well in the future. These inferences can sometimes be obvious, such as "since the sun rose every morning for the last 10,000 days, it will probably rise tomorrow morning as well". Other times, they can be more nuanced, such as "X% of families have geographically separate species with color variants, so there is a Y% chance that undiscovered black swans exist".
-Machine learning programs can perform tasks without being explicitly programmed to do so. It involves computers learning from data provided so that they carry out certain tasks. For simple tasks assigned to computers, it is possible to program algorithms telling the machine how to execute all steps required to solve the problem at hand; on the computer's part, no learning is needed. For more advanced tasks, it can be challenging for a human to manually create the needed algorithms. In practice, it can turn out to be more effective to help the machine develop its own algorithm, rather than having human programmers specify every needed step.
-The discipline of machine learning employs various approaches to teach computers to accomplish tasks where no fully satisfactory algorithm is available. In cases where vast numbers of potential answers exist, one approach is to label some of the correct answers as valid. This can then be used as training data for the computer to improve the algorithm(s) it uses to determine correct answers. For example, to train a system for the task of digital character recognition, the MNIST dataset of handwritten digits has often been used.
-The term machine learning was coined in 1959 by Arthur Samuel, an IBM employee and pioneer in the field of computer gaming and artificial intelligence. The synonym self-teaching computers was also used in this time period.
-By the early 1960s an experimental "learning machine" with punched tape memory, called Cybertron, had been developed by Raytheon Company to analyze sonar signals, electrocardiograms, and speech patterns using rudimentary reinforcement learning. It was repetitively "trained" by a human operator/teacher to recognize patterns and equipped with a "goof" button to cause it to re-evaluate incorrect decisions. A representative book on research into machine learning during the 1960s was Nilsson's book on Learning Machines, dealing mostly with machine learning for pattern classification. Interest related to pattern recognition continued into the 1970s, as described by Duda and Hart in 1973. In 1981 a report was given on using teaching strategies so that a neural network learns to recognize 40 characters (26 letters, 10 digits, and 4 special symbols) from a computer terminal.
-Tom M. Mitchell provided a widely quoted, more formal definition of the algorithms studied in the machine learning field: "A computer program is said to learn from experience E with respect to some class of tasks T and performance measure P if its performance at tasks in T, as measured by P, improves with experience E." This definition of the tasks in which machine learning is concerned offers a fundamentally operational definition rather than defining the field in cognitive terms. This follows Alan Turing's proposal in his paper "Computing Machinery and Intelligence", in which the question "Can machines think?" is replaced with the question "Can machines do what we (as thinking entities) can do?".
-Modern-day machine learning has two objectives, one is to classify data based on models which have been developed, the other purpose is to make predictions for future outcomes based on these models. A hypothetical algorithm specific to classifying data may use computer vision of moles coupled with supervised learning in order to train it to classify the cancerous moles. A machine learning algorithm for stock trading may inform the trader of future potential predictions.
-As a scientific endeavor, machine learning grew out of the quest for artificial intelligence (AI). In the early days of AI as an academic discipline, some researchers were interested in having machines learn from data. They attempted to approach the problem with various symbolic methods, as well as what were then termed "neural networks"; these were mostly perceptrons and other models that were later found to be reinventions of the generalized linear models of statistics. Probabilistic reasoning was also employed, especially in automated medical diagnosis.: 488 
-However, an increasing emphasis on the logical, knowledge-based approach caused a rift between AI and machine learning. Probabilistic systems were plagued by theoretical and practical problems of data acquisition and representation.: 488  By 1980, expert systems had come to dominate AI, and statistics was out of favor. Work on symbolic/knowledge-based learning did continue within AI, leading to inductive logic programming, but the more statistical line of research was now outside the field of AI proper, in pattern recognition and information retrieval.: 708–710, 755  Neural networks research had been abandoned by AI and computer science around the same time. This line, too, was continued outside the AI/CS field, as "connectionism", by researchers from other disciplines including Hopfield, Rumelhart, and Hinton. Their main success came in the mid-1980s with the reinvention of backpropagation.: 25 
-Machine learning (ML), reorganized and recognized as its own field, started to flourish in the 1990s. The field changed its goal from achieving artificial intelligence to tackling solvable problems of a practical nature. It shifted focus away from the symbolic approaches it had inherited from AI, and toward methods and models borrowed from statistics, fuzzy logic, and probability theory.
-Machine learning and data mining often employ the same methods and overlap significantly, but while machine learning focuses on prediction, based on known properties learned from the training data, data mining focuses on the discovery of (previously) unknown properties in the data (this is the analysis step of knowledge discovery in databases). Data mining uses many machine learning methods, but with different goals; on the other hand, machine learning also employs data mining methods as "unsupervised learning" or as a preprocessing step to improve learner accuracy. Much of the confusion between these two research communities (which do often have separate conferences and separate journals, ECML PKDD being a major exception) comes from the basic assumptions they work with: in machine learning, performance is usually evaluated with respect to the ability to reproduce known knowledge, while in knowledge discovery and data mining (KDD) the key task is the discovery of previously unknown knowledge. Evaluated with respect to known knowledge, an uninformed (unsupervised) method will easily be outperformed by other supervised methods, while in a typical KDD task, supervised methods cannot be used due to the unavailability of training data.
-Machine learning also has intimate ties to optimization: many learning problems are formulated as minimization of some loss function on a training set of examples. Loss functions express the discrepancy between the predictions of the model being trained and the actual problem instances (for example, in classification, one wants to assign a label to instances, and models are trained to correctly predict the pre-assigned labels of a set of examples).
-The difference between optimization and machine learning arises from the goal of generalization: while optimization algorithms can minimize the loss on a training set, machine learning is concerned with minimizing the loss on unseen samples. Characterizing the generalization of various learning algorithms is an active topic of current research, especially for deep learning algorithms.
-Machine learning and statistics are closely related fields in terms of methods, but distinct in their principal goal: statistics draws population inferences from a sample, while machine learning finds generalizable predictive patterns. According to Michael I. Jordan, the ideas of machine learning, from methodological principles to theoretical tools, have had a long pre-history in statistics. He also suggested the term data science as a placeholder to call the overall field.
-Leo Breiman distinguished two statistical modeling paradigms: data model and algorithmic model, wherein "algorithmic model" means more or less the machine learning algorithms like Random Forest.
-Some statisticians have adopted methods from machine learning, leading to a combined field that they call statistical learning.
-Analytical and computational techniques derived from deep-rooted physics of disordered systems can be extended to large-scale problems, including machine learning, e.g., to analyze the weight space of deep neural networks. Statistical physics is thus finding applications in the area of medical diagnostics.
-A core objective of a learner is to generalize from its experience. Generalization in this context is the ability of a learning machine to perform accurately on new, unseen examples/tasks after having experienced a learning data set. The training examples come from some generally unknown probability distribution (considered representative of the space of occurrences) and the learner has to build a general model about this space that enables it to produce sufficiently accurate predictions in new cases.
-The computational analysis of machine learning algorithms and their performance is a branch of theoretical computer science known as computational learning theory via the Probably Approximately Correct Learning (PAC) model. Because training sets are finite and the future is uncertain, learning theory usually does not yield guarantees of the performance of algorithms. Instead, probabilistic bounds on the performance are quite common. The bias–variance decomposition is one way to quantify generalization error.
-For the best performance in the context of generalization, the complexity of the hypothesis should match the complexity of the function underlying the data. If the hypothesis is less complex than the function, then the model has under fitted the data. If the complexity of the model is increased in response, then the training error decreases. But if the hypothesis is too complex, then the model is subject to overfitting and generalization will be poorer.
-In addition to performance bounds, learning theorists study the time complexity and feasibility of learning. In computational learning theory, a computation is considered feasible if it can be done in polynomial time. There are two kinds of time complexity results: Positive results show that a certain class of functions can be learned in polynomial time. Negative results show that certain classes cannot be learned in polynomial time.
-Machine learning approaches are traditionally divided into three broad categories, which correspond to learning paradigms, depending on the nature of the "signal" or "feedback" available to the learning system:
-Supervised learning: The computer is presented with example inputs and their desired outputs, given by a "teacher", and the goal is to learn a general rule that maps inputs to outputs.
-Unsupervised learning: No labels are given to the learning algorithm, leaving it on its own to find structure in its input. Unsupervised learning can be a goal in itself (discovering hidden patterns in data) or a means towards an end (feature learning).
-Reinforcement learning: A computer program interacts with a dynamic environment in which it must perform a certain goal (such as driving a vehicle or playing a game against an opponent). As it navigates its problem space, the program is provided feedback that's analogous to rewards, which it tries to maximize. Although each algorithm has advantages and limitations, no single algorithm works for all problems.
-Supervised learning algorithms build a mathematical model of a set of data that contains both the inputs and the desired outputs. The data is known as training data, and consists of a set of training examples. Each training example has one or more inputs and the desired output, also known as a supervisory signal. In the mathematical model, each training example is represented by an array or vector, sometimes called a feature vector, and the training data is represented by a matrix. Through iterative optimization of an objective function, supervised learning algorithms learn a function that can be used to predict the output associated with new inputs. An optimal function will allow the algorithm to correctly determine the output for inputs that were not a part of the training data. An algorithm that improves the accuracy of its outputs or predictions over time is said to have learned to perform that task.
-Types of supervised-learning algorithms include active learning, classification and regression. Classification algorithms are used when the outputs are restricted to a limited set of values, and regression algorithms are used when the outputs may have any numerical value within a range. As an example, for a classification algorithm that filters emails, the input would be an incoming email, and the output would be the name of the folder in which to file the email.
-Similarity learning is an area of supervised machine learning closely related to regression and classification, but the goal is to learn from examples using a similarity function that measures how similar or related two objects are. It has applications in ranking, recommendation systems, visual identity tracking, face verification, and speaker verification.
-Unsupervised learning algorithms take a set of data that contains only inputs, and find structure in the data, like grouping or clustering of data points. The algorithms, therefore, learn from test data that has not been labeled, classified or categorized. Instead of responding to feedback, unsupervised learning algorithms identify commonalities in the data and react based on the presence or absence of such commonalities in each new piece of data. A central application of unsupervised learning is in the field of density estimation in statistics, such as finding the probability density function. Though unsupervised learning encompasses other domains involving summarizing and explaining data features. Unsupervised learning algorithms streamlined the process of survey and graph large indel based haplotypes of a gene of interest from pan-genome. 
-Cluster analysis is the assignment of a set of observations into subsets (called clusters) so that observations within the same cluster are similar according to one or more predesignated criteria, while observations drawn from different clusters are dissimilar. Different clustering techniques make different assumptions on the structure of the data, often defined by some similarity metric and evaluated, for example, by internal compactness, or the similarity between members of the same cluster, and separation, the difference between clusters. Other methods are based on estimated density and graph connectivity.
-Semi-supervised learning falls between unsupervised learning (without any labeled training data) and supervised learning (with completely labeled training data). Some of the training examples are missing training labels, yet many machine-learning researchers have found that unlabeled data, when used in conjunction with a small amount of labeled data, can produce a considerable improvement in learning accuracy.
-In weakly supervised learning, the training labels are noisy, limited, or imprecise; however, these labels are often cheaper to obtain, resulting in larger effective training sets.
-Reinforcement learning is an area of machine learning concerned with how software agents ought to take actions in an environment so as to maximize some notion of cumulative reward. Due to its generality, the field is studied in many other disciplines, such as game theory, control theory, operations research, information theory, simulation-based optimization, multi-agent systems, swarm intelligence, statistics and genetic algorithms. In machine learning, the environment is typically represented as a Markov decision process (MDP). Many reinforcements learning algorithms use dynamic programming techniques. Reinforcement learning algorithms do not assume knowledge of an exact mathematical model of the MDP and are used when exact models are infeasible. Reinforcement learning algorithms are used in autonomous vehicles or in learning to play a game against a human opponent.
-Dimensionality reduction is a process of reducing the number of random variables under consideration by obtaining a set of principal variables. In other words, it is a process of reducing the dimension of the feature set, also called the "number of features". Most of the dimensionality reduction techniques can be considered as either feature elimination or extraction. One of the popular methods of dimensionality reduction is principal component analysis (PCA). PCA involves changing higher-dimensional data (e.g., 3D) to a smaller space (e.g., 2D). This results in a smaller dimension of data (2D instead of 3D), while keeping all original variables in the model without changing the data. The manifold hypothesis proposes that high-dimensional data sets lie along low-dimensional manifolds, and many dimensionality reduction techniques make this assumption, leading to the area of manifold learning and manifold regularization.
-Although machine learning has been transformative in some fields, machine-learning programs often fail to deliver expected results. Reasons for this are numerous: lack of (suitable) data, lack of access to the data, data bias, privacy problems, badly chosen tasks and algorithms, wrong tools and people, lack of resources, and evaluation problems.
-In 2018, a self-driving car from Uber failed to detect a pedestrian, who was killed after a collision. Attempts to use machine learning in healthcare with the IBM Watson system failed to deliver even after years of time and billions of dollars invested.
-Machine learning has been used as a strategy to update the evidence related to a systematic review and increased reviewer burden related to the growth of biomedical literature. While it has improved with training sets, it has not yet developed sufficiently to reduce the workload burden without limiting the necessary sensitivity for the findings research themselves.
-Machine learning approaches in particular can suffer from different data biases. A machine learning system trained specifically on current customers may not be able to predict the needs of new customer groups that are not represented in the training data. When trained on human-made data, machine learning is likely to pick up the constitutional and unconscious biases already present in society. Language models learned from data have been shown to contain human-like biases. Machine learning systems used for criminal risk assessment have been found to be biased against black people. In 2015, Google photos would often tag black people as gorillas, and in 2018 this still was not well resolved, but Google reportedly was still using the workaround to remove all gorillas from the training data, and thus was not able to recognize real gorillas at all. Similar issues with recognizing non-white people have been found in many other systems. In 2016, Microsoft tested a chatbot that learned from Twitter, and it quickly picked up racist and sexist language. Because of such challenges, the effective use of machine learning may take longer to be adopted in other domains. Concern for fairness in machine learning, that is, reducing bias in machine learning and propelling its use for human good is increasingly expressed by artificial intelligence scientists, including Fei-Fei Li, who reminds engineers that "There's nothing artificial about AI...It's inspired by people, it's created by people, and—most importantly—it impacts people. It is a powerful tool we are only just beginning to understand, and that is a profound responsibility."
-Learners can also disappoint by "learning the wrong lesson". A toy example is that an image classifier trained only on pictures of brown horses and black cats might conclude that all brown patches are likely to be horses. A real-world example is that, unlike humans, current image classifiers often do not primarily make judgments from the spatial relationship between components of the picture, and they learn relationships between pixels that humans are oblivious to, but that still correlate with images of certain types of real objects. Modifying these patterns on a legitimate image can result in "adversarial" images that the system misclassifies.
-Adversarial vulnerabilities can also result in nonlinear systems, or from non-pattern perturbations. Some systems are so brittle that changing a single adversarial pixel predictably induces misclassification.[citation needed] Machine learning models are often vulnerable to manipulation and/or evasion via adversarial machine learning.
-Researchers have demonstrated how backdoors can be placed undetectably into classifying (e.g., for categories "spam" and well-visible "not spam" of posts) machine learning models which are often developed and/or trained by third parties. Parties can change the classification of any input, including in cases for which a type of data/software transparency is provided, possibly including white-box access.
-Machine learning poses a host of ethical questions. Systems that are trained on datasets collected with biases may exhibit these biases upon use (algorithmic bias), thus digitizing cultural prejudices. For example, in 1988, the UK's Commission for Racial Equality found that St. George's Medical School had been using a computer program trained from data of previous admissions staff and this program had denied nearly 60 candidates who were found to be either women or had non-European sounding names. Using job hiring data from a firm with racist hiring policies may lead to a machine learning system duplicating the bias by scoring job applicants by similarity to previous successful applicants. Responsible collection of data and documentation of algorithmic rules used by a system thus is a critical part of machine learning.
-AI can be well-equipped to make decisions in technical fields, which rely heavily on data and historical information. These decisions rely on the objectivity and logical reasoning. Because human languages contain biases, machines trained on language corpora will necessarily also learn these biases.
-Other forms of ethical challenges, not related to personal biases, are seen in health care. There are concerns among health care professionals that these systems might not be designed in the public's interest but as income-generating machines. This is especially true in the United States where there is a long-standing ethical dilemma of improving health care, but also increase profits. For example, the algorithms could be designed to provide patients with unnecessary tests or medication in which the algorithm's proprietary owners hold stakes. There is potential for machine learning in health care to provide professionals an additional tool to diagnose, medicate, and plan recovery paths for patients, but this requires these biases to be mitigated.
-Since the 2010s, advances in both machine learning algorithms and computer hardware have led to more efficient methods for training deep neural networks (a particular narrow subdomain of machine learning) that contain many layers of non-linear hidden units. By 2019, graphic processing units (GPUs), often with AI-specific enhancements, had displaced CPUs as the dominant method of training large-scale commercial cloud AI. OpenAI estimated the hardware computing used in the largest deep learning projects from AlexNet (2012) to AlphaZero (2017), and found a 300,000-fold increase in the amount of compute required, with a doubling-time trendline of 3.4 months.

data/wiki_demo.txt.REMOVED.git-id

@@ -0,0 +1 @@
+c9cf509b7fdac5490cfd6dae72c2d7b8a60af6cb

requirements.txt

@@ -3,13 +3,12 @@ transformers>=4.31.0,<4.35.0
 datasets>=2.14.0
 accelerate>=0.21.0
 peft>=0.6.0
-trl==0.7.2
+trl>=0.7.4
 gradio>=3.38.0,<4.0.0
 scipy
 sentencepiece
 protobuf
 tiktoken
-fire
 jieba
 rouge-chinese
 nltk

src/cli_demo.py

@@ -1,5 +1,12 @@
-import readline
 from llmtuner import ChatModel
+from llmtuner.extras.misc import torch_gc
+
+try:
+    import platform
+    if platform.system() != "Windows":
+        import readline
+except ImportError:
+    print("Install `readline` for a better experience.")
 
 
 def main():
@@ -21,6 +28,7 @@ def main():
 
         if query.strip() == "clear":
             history = []
+            torch_gc()
             print("History has been removed.")
             continue
 

src/evaluate.py

@@ -1,190 +1,10 @@
-# coding=utf-8
-# Evaluates the performance of pre-trained models.
-# Usage: python evaluate.py --model_name_or_path path_to_model --checkpoint_dir path_to_ckpt --template vanilla
-#                           --task ceval --split validation --lang zh --n_shot 5 --batch_size 4 --save_name result
-# Inspired by: https://github.com/hendrycks/test/blob/master/evaluate_flan.py
-
-import os
-import fire
-import json
-import torch
-import numpy as np
-import transformers
-from collections import Counter
-from datasets import load_dataset
-from dataclasses import dataclass
-from tqdm import tqdm, trange
-from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple
-
-from llmtuner import ChatModel
-
-if TYPE_CHECKING:
-    from datasets import Dataset
+from llmtuner import Evaluator
 
 
-choices = ["A", "B", "C", "D"]
-
-
-@dataclass
-class EvalTemplate:
-
-    system: str
-    choice: str
-    answer: str
-    prefix: str
-
-    def parse_example(
-        self,
-        example: Dict[str, str]
-    ) -> Tuple[str, str]:
-        candidates = [self.choice.format(choice=ch, content=example[ch]) for ch in choices if ch in example]
-        return "".join([example["question"]] + candidates + [self.answer]), example["answer"]
-
-    def format_example(
-        self,
-        target_data: Dict[str, str],
-        support_set: "Dataset",
-        subject_name: str,
-        use_history: bool
-    ) -> Tuple[str, str, List[Tuple[str, str]]]:
-        query, resp = self.parse_example(target_data)
-        history = [self.parse_example(support_set[k]) for k in range(len(support_set))]
-
-        if len(history):
-            temp = history.pop(0)
-            history.insert(0, (self.system.format(subject=subject_name) + temp[0], temp[1]))
-        else:
-            query = self.system.format(subject=subject_name) + query
-
-        if not use_history:
-            query = "\n\n".join(["".join(item) for item in history] + [query])
-            history = []
-        return query.strip(), resp, history
-
-
-eval_templates = {
-    "en": EvalTemplate(
-        system="The following are multiple choice questions (with answers) about {subject}.\n\n",
-        choice="\n{choice}. {content}",
-        answer="\nAnswer: ",
-        prefix=" "
-    ),
-    "zh": EvalTemplate(
-        system="以下是中国关于{subject}考试的单项选择题，请选出其中的正确答案。\n\n",
-        choice="\n{choice}. {content}",
-        answer="\n答案：",
-        prefix="\n"
-    )
-}
-
-
-@torch.inference_mode()
-def batch_inference(
-    chat_model: ChatModel,
-    batch_input: Dict[str, torch.Tensor],
-    prefix_char: str
-) -> List[str]:
-    logits = chat_model.model(**batch_input).logits
-    lengths = torch.sum(batch_input["attention_mask"], dim=-1)
-    nextword_logits = torch.stack([logits[i, lengths[i] - 1] for i in range(len(lengths))], dim=0)
-    probs = torch.nn.functional.softmax(
-        torch.stack(
-            [
-                nextword_logits[:, chat_model.tokenizer.encode(prefix_char + choice, add_special_tokens=False)[-1]]
-                for choice in choices
-            ],
-            dim=-1
-        ),
-        dim=-1
-    ).detach()
-    return [chr(ord("A") + offset.item()) for offset in torch.argmax(probs, dim=-1)]
-
-
-def evaluate(
-    model_name_or_path: str,
-    finetuning_type: Optional[str] = "lora",
-    checkpoint_dir: Optional[str] = None,
-    template: Optional[str] = "vanilla",
-    task: Optional[str] = "ceval",
-    dataset_dir: Optional[str] = "evaluation",
-    split: Optional[Literal["validation", "test"]] = "validation",
-    lang: Optional[Literal["zh", "en"]] = "zh",
-    n_shot: Optional[int] = 5,
-    n_avg: Optional[int] = 1,
-    batch_size: Optional[int] = 4,
-    save_name: Optional[str] = None,
-    seed: Optional[int] = 42
-):
-    with open(os.path.join(dataset_dir, task, "mapping.json"), "r", encoding="utf-8") as f:
-        categorys: Dict[str, Dict[str, str]] = json.load(f)
-
-    transformers.set_seed(seed)
-    chat_model = ChatModel(dict(
-        model_name_or_path=model_name_or_path,
-        finetuning_type=finetuning_type,
-        checkpoint_dir=checkpoint_dir,
-        template=template
-    ))
-    chat_model.tokenizer.padding_side = "right" # avoid overflow issue in batched inference for llama2
-    eval_template = eval_templates[lang]
-
-    category_corrects: Dict[str, np.ndarray] = {
-        subj: np.array([], dtype="bool") for subj in ["Average", "STEM", "Social Sciences", "Humanities", "Other"]
-    }
-    pbar = tqdm(categorys.keys(), desc="Processing subjects", position=0)
-    results = {}
-    for subject in pbar:
-        dataset = load_dataset(os.path.join(dataset_dir, task), subject)
-        labels, answers, all_outputs = [], [], []
-        for epoch in range(n_avg):
-            pbar.set_postfix_str("{} Trial: {}".format(categorys[subject]["name"], epoch))
-            inputs, outputs = [], []
-            for i in trange(len(dataset[split]), desc="Formatting batches", position=1, leave=False):
-                support_set = dataset["train"].shuffle().select(range(min(n_shot, len(dataset["train"]))))
-                query, resp, history = eval_template.format_example(
-                    target_data=dataset[split][i],
-                    support_set=support_set,
-                    subject_name=categorys[subject]["name"],
-                    use_history=chat_model.template.use_history
-                )
-                input_ids, _ = chat_model.template.encode_oneturn(
-                    tokenizer=chat_model.tokenizer, query=query, resp=resp, history=history
-                )
-                inputs.append({"input_ids": input_ids, "attention_mask": [1] * len(input_ids)})
-                if epoch == 0:
-                    labels.append(resp)
-
-            for i in trange(0, len(inputs), batch_size, desc="Predicting batches", position=1, leave=False):
-                batch_input = chat_model.tokenizer.pad(
-                    inputs[i : i + batch_size], return_attention_mask=True, return_tensors="pt"
-                ).to(chat_model.model.device)
-                preds = batch_inference(chat_model, batch_input, eval_template.prefix)
-                outputs += preds
-            all_outputs.append(outputs)
-
-        for i in range(len(all_outputs[0])):
-            count = Counter([all_outputs[epoch][i] for epoch in range(n_avg)])
-            answers.append(count.most_common(1)[0][0])
-
-        corrects = (np.array(answers) == np.array(labels))
-        category_name = categorys[subject]["category"]
-        category_corrects[category_name] = np.concatenate([category_corrects[category_name], corrects], axis=0)
-        category_corrects["Average"] = np.concatenate([category_corrects["Average"], corrects], axis=0)
-        results[subject] = {str(i): answers[i] for i in range(len(answers))}
-
-    score_info = "\n".join([
-        "{:>15}: {:.2f}".format(category_name, 100 * np.mean(category_correct))
-        for category_name, category_correct in category_corrects.items() if len(category_correct)
-    ])
-
-    print(score_info)
-    if save_name is not None:
-        with open(save_name + ".json", "w", encoding="utf-8", newline="\n") as f:
-            json.dump(results, f, indent=2)
-
-        with open(save_name + ".log", "w", encoding="utf-8", newline="\n") as f:
-            f.write(score_info)
+def main():
+    evaluator = Evaluator()
+    evaluator.eval()
 
 
 if __name__ == "__main__":
-    fire.Fire(evaluate)
+    main()

src/llmtuner/__init__.py

@@ -1,9 +1,10 @@
-# Level: api, webui > chat > tuner > dsets > extras, hparams
+# Level: api, webui > chat, eval, train > data, model > extras, hparams
 
 from llmtuner.api import create_app
 from llmtuner.chat import ChatModel
-from llmtuner.tuner import export_model, run_exp
+from llmtuner.eval import Evaluator
+from llmtuner.train import export_model, run_exp
 from llmtuner.webui import create_ui, create_web_demo
 
 
-__version__ = "0.2.1"
+__version__ = "0.3.2"

src/llmtuner/api/app.py

@@ -1,14 +1,8 @@
 import json
-import uvicorn
-from fastapi import FastAPI, HTTPException, status
-from fastapi.middleware.cors import CORSMiddleware
-from contextlib import asynccontextmanager
-from sse_starlette import EventSourceResponse
 from typing import List, Tuple
 from pydantic import BaseModel
+from contextlib import asynccontextmanager
 
-from llmtuner.extras.misc import torch_gc
-from llmtuner.chat import ChatModel
 from llmtuner.api.protocol import (
     Role,
     Finish,
@@ -23,10 +17,28 @@ from llmtuner.api.protocol import (
     ChatCompletionResponseStreamChoice,
     ChatCompletionResponseUsage
 )
+from llmtuner.chat import ChatModel
+from llmtuner.extras.misc import torch_gc
+from llmtuner.extras.packages import (
+    is_fastapi_availble, is_starlette_available, is_uvicorn_available
+)
+
+
+if is_fastapi_availble():
+    from fastapi import FastAPI, HTTPException, status
+    from fastapi.middleware.cors import CORSMiddleware
+
+
+if is_starlette_available():
+    from sse_starlette import EventSourceResponse
+
+
+if is_uvicorn_available():
+    import uvicorn
 
 
 @asynccontextmanager
-async def lifespan(app: FastAPI): # collects GPU memory
+async def lifespan(app: "FastAPI"): # collects GPU memory
     yield
     torch_gc()
 
@@ -38,7 +50,7 @@ def to_json(data: BaseModel) -> str:
         return data.json(exclude_unset=True, ensure_ascii=False)
 
 
-def create_app(chat_model: ChatModel) -> FastAPI:
+def create_app(chat_model: "ChatModel") -> "FastAPI":
     app = FastAPI(lifespan=lifespan)
 
     app.add_middleware(
@@ -56,12 +68,12 @@ def create_app(chat_model: ChatModel) -> FastAPI:
 
     @app.post("/v1/chat/completions", response_model=ChatCompletionResponse, status_code=status.HTTP_200_OK)
     async def create_chat_completion(request: ChatCompletionRequest):
-        if len(request.messages) < 1 or request.messages[-1].role != Role.USER:
+        if len(request.messages) == 0 or request.messages[-1].role != Role.USER:
             raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid request")
 
         query = request.messages[-1].content
         prev_messages = request.messages[:-1]
-        if len(prev_messages) > 0 and prev_messages[0].role == Role.SYSTEM:
+        if len(prev_messages) and prev_messages[0].role == Role.SYSTEM:
             system = prev_messages.pop(0).content
         else:
             system = None
@@ -73,12 +85,14 @@ def create_app(chat_model: ChatModel) -> FastAPI:
                     history.append([prev_messages[i].content, prev_messages[i+1].content])
                 else:
                     raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Only supports u/a/u/a/u...")
+        else:
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Only supports u/a/u/a/u...")
 
         if request.stream:
             generate = predict(query, history, system, request)
             return EventSourceResponse(generate, media_type="text/event-stream")
 
-        response, (prompt_length, response_length) = chat_model.chat(
+        responses = chat_model.chat(
             query, history, system,
             do_sample=request.do_sample,
             temperature=request.temperature,
@@ -87,18 +101,23 @@ def create_app(chat_model: ChatModel) -> FastAPI:
             num_return_sequences=request.n
         )
 
+        prompt_length, response_length = 0, 0
+        choices = []
+        for i, response in enumerate(responses):
+            choices.append(ChatCompletionResponseChoice(
+                index=i,
+                message=ChatMessage(role=Role.ASSISTANT, content=response.response_text),
+                finish_reason=Finish.STOP if response.finish_reason == "stop" else Finish.LENGTH
+            ))
+            prompt_length = response.prompt_length
+            response_length += response.response_length
+
         usage = ChatCompletionResponseUsage(
             prompt_tokens=prompt_length,
             completion_tokens=response_length,
             total_tokens=prompt_length+response_length
         )
 
-        choices = [ChatCompletionResponseChoice(
-            index=i,
-            message=ChatMessage(role=Role.ASSISTANT, content=choice),
-            finish_reason=Finish.STOP
-        ) for i, choice in enumerate(response)]
-
         return ChatCompletionResponse(model=request.model, choices=choices, usage=usage)
 
     async def predict(query: str, history: List[Tuple[str, str]], system: str, request: ChatCompletionRequest):

src/llmtuner/chat/__init__.py

@@ -1 +1 @@
-from llmtuner.chat.stream_chat import ChatModel
+from llmtuner.chat.chat_model import ChatModel

src/llmtuner/chat/chat_model.py

@@ -1,11 +1,21 @@
 import torch
-from typing import Any, Dict, Generator, List, Optional, Tuple
+from dataclasses import dataclass
+from typing import Any, Dict, Generator, List, Literal, Optional, Tuple
 from threading import Thread
 from transformers import GenerationConfig, TextIteratorStreamer
 
-from llmtuner.extras.misc import dispatch_model, get_logits_processor
-from llmtuner.extras.template import get_template_and_fix_tokenizer
-from llmtuner.tuner.core import get_infer_args, load_model_and_tokenizer
+from llmtuner.data.template import get_template_and_fix_tokenizer
+from llmtuner.extras.misc import get_logits_processor
+from llmtuner.model import dispatch_model, get_infer_args, load_model_and_tokenizer
+
+
+@dataclass
+class Response:
+
+    response_text: str
+    response_length: int
+    prompt_length: int
+    finish_reason: Literal["stop", "length"]
 
 
 class ChatModel:
@@ -18,7 +28,7 @@ class ChatModel:
         self.template = get_template_and_fix_tokenizer(data_args.template, self.tokenizer)
         self.system_prompt = data_args.system_prompt
 
-    def process_args(
+    def _process_args(
         self,
         query: str,
         history: Optional[List[Tuple[str, str]]] = None,
@@ -79,17 +89,30 @@ class ChatModel:
         history: Optional[List[Tuple[str, str]]] = None,
         system: Optional[str] = None,
         **input_kwargs
-    ) -> Tuple[List[str], Tuple[int, int]]:
-        gen_kwargs, prompt_length = self.process_args(query, history, system, **input_kwargs)
+    ) -> List[Response]:
+        r"""
+        Args: query, history, system, **input_kwargs
+
+        Returns: [(response_text, prompt_length, response_length)] * n (default n=1)
+        """
+        gen_kwargs, prompt_length = self._process_args(query, history, system, **input_kwargs)
         generate_output = self.model.generate(**gen_kwargs)
         response_ids = generate_output[:, prompt_length:]
-        response = self.tokenizer.batch_decode(response_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
-        response_length = 0
-        for i in range(len(response_ids)):
+        response = self.tokenizer.batch_decode(
+            response_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
+        )
+        results = []
+        for i in range(len(response)):
             eos_index = (response_ids[i] == self.tokenizer.eos_token_id).nonzero()
-            response_length += eos_index[0].item() if len(eos_index) else len(response_ids[i])
+            response_length = (eos_index[0].item() + 1) if len(eos_index) else len(response_ids[i])
+            results.append(Response(
+                response_text=response[i],
+                response_length=response_length,
+                prompt_length=prompt_length,
+                finish_reason="stop" if len(eos_index) else "length"
+            ))
 
-        return response, (prompt_length, response_length)
+        return results
 
     @torch.inference_mode()
     def stream_chat(
@@ -99,7 +122,7 @@ class ChatModel:
         system: Optional[str] = None,
         **input_kwargs
     ) -> Generator[str, None, None]:
-        gen_kwargs, _ = self.process_args(query, history, system, **input_kwargs)
+        gen_kwargs, _ = self._process_args(query, history, system, **input_kwargs)
         streamer = TextIteratorStreamer(self.tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
         gen_kwargs["streamer"] = streamer
 

src/llmtuner/data/__init__.py

@@ -0,0 +1,4 @@
+from llmtuner.data.loader import get_dataset
+from llmtuner.data.preprocess import preprocess_dataset
+from llmtuner.data.template import get_template_and_fix_tokenizer
+from llmtuner.data.utils import split_dataset

src/llmtuner/data/loader.py

@@ -3,7 +3,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Union
 
 from datasets import concatenate_datasets, interleave_datasets, load_dataset
 
-from llmtuner.dsets.utils import checksum, EXT2TYPE
+from llmtuner.data.utils import checksum, EXT2TYPE
 from llmtuner.extras.logging import get_logger
 
 if TYPE_CHECKING:
@@ -60,9 +60,12 @@ def get_dataset(
             split=data_args.split,
             cache_dir=model_args.cache_dir,
             token=model_args.hf_hub_token,
-            streaming=data_args.streaming
+            streaming=(data_args.streaming and (dataset_attr.load_from != "file"))
         )
 
+        if data_args.streaming and (dataset_attr.load_from == "file"):
+            dataset = dataset.to_iterable_dataset() # TODO: add num shards parameter
+
         if max_samples is not None: # truncate dataset
             dataset = dataset.select(range(min(len(dataset), max_samples)))
 

src/llmtuner/data/preprocess.py

@@ -1,13 +1,13 @@
 import os
 import tiktoken
 from itertools import chain
-from typing import TYPE_CHECKING, Any, Dict, Generator, List, Literal, Union
+from typing import TYPE_CHECKING, Any, Dict, Generator, List, Literal, Tuple, Union
 
 from datasets import load_from_disk
 
+from llmtuner.data.template import get_template_and_fix_tokenizer
 from llmtuner.extras.constants import IGNORE_INDEX
 from llmtuner.extras.logging import get_logger
-from llmtuner.extras.template import get_template_and_fix_tokenizer
 
 if TYPE_CHECKING:
     from datasets import Dataset, IterableDataset
@@ -19,6 +19,22 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)
 
 
+def construct_example(examples: Dict[str, List[Any]]) -> Generator[Any, None, None]:
+    for i in range(len(examples["prompt"])):
+        query, response = examples["prompt"][i], examples["response"][i]
+        query = query + "\n" + examples["query"][i] if "query" in examples and examples["query"][i] else query
+        history = examples["history"][i] if "history" in examples else None
+        system = examples["system"][i] if "system" in examples else None
+        yield query, response, history, system
+
+
+def infer_max_len(source_len: int, target_len: int, data_args: "DataArguments") -> Tuple[int, int]:
+    max_target_len = int(data_args.cutoff_len * (target_len / (source_len + target_len)))
+    max_target_len = max(max_target_len, data_args.reserved_label_len)
+    max_source_len = data_args.cutoff_len - max_target_len
+    return max_source_len, max_target_len
+
+
 def preprocess_dataset(
     dataset: Union["Dataset", "IterableDataset"],
     tokenizer: "PreTrainedTokenizer",
@@ -31,14 +47,6 @@ def preprocess_dataset(
     if data_args.train_on_prompt and template.efficient_eos:
         raise ValueError("Current template does not support `train_on_prompt`.")
 
-    def construct_example(examples: Dict[str, List[Any]]) -> Generator[Any, None, None]:
-        for i in range(len(examples["prompt"])):
-            query, response = examples["prompt"][i], examples["response"][i]
-            query = query + "\n" + examples["query"][i] if "query" in examples and examples["query"][i] else query
-            history = examples["history"][i] if "history" in examples else None
-            system = examples["system"][i] if "system" in examples else None
-            yield query, response, history, system
-
     def preprocess_pretrain_dataset(examples: Dict[str, List[Any]]) -> Dict[str, List[List[int]]]:
         # build grouped texts with format `X1 X2 X3 ...`
         if isinstance(getattr(tokenizer, "tokenizer", None), tiktoken.Encoding): # for tiktoken tokenizer (Qwen)
@@ -79,13 +87,11 @@ def preprocess_dataset(
             for turn_idx, (source_ids, target_ids) in enumerate(template.encode_multiturn(
                 tokenizer, query, response, history, system
             )):
-                total_len = len(source_ids) + len(target_ids)
-                max_source_len = int(data_args.cutoff_len * (len(source_ids) / total_len))
-                max_target_len = int(data_args.cutoff_len * (len(target_ids) / total_len))
-
-                if len(source_ids) > max_source_len:
+                source_len, target_len = len(source_ids), len(target_ids)
+                max_source_len, max_target_len = infer_max_len(source_len, target_len, data_args)
+                if source_len > max_source_len:
                     source_ids = source_ids[:max_source_len]
-                if len(target_ids) > max_target_len:
+                if target_len > max_target_len:
                     target_ids = target_ids[:max_target_len]
 
                 if data_args.train_on_prompt:
@@ -187,15 +193,12 @@ def preprocess_dataset(
                 chosen_ids += [tokenizer.eos_token_id]
                 rejected_ids += [tokenizer.eos_token_id]
 
-            total_len = len(prompt_ids) + max(len(chosen_ids), len(rejected_ids))
-            max_source_len = int(data_args.cutoff_len * (len(prompt_ids) / total_len))
-            max_target_len = int(data_args.cutoff_len * (max(len(chosen_ids), len(rejected_ids)) / total_len))
-
-            if len(prompt_ids) > max_source_len:
+            source_len, target_len = len(prompt_ids), max(len(chosen_ids), len(rejected_ids))
+            max_source_len, max_target_len = infer_max_len(source_len, target_len, data_args)
+            if source_len > max_source_len:
                 prompt_ids = prompt_ids[:max_source_len]
-            if len(chosen_ids) > max_target_len:
+            if target_len > max_target_len:
                 chosen_ids = chosen_ids[:max_target_len]
-            if len(rejected_ids) > max_target_len:
                 rejected_ids = rejected_ids[:max_target_len]
 
             model_inputs["prompt_ids"].append(prompt_ids)

src/llmtuner/data/template.py

@@ -114,7 +114,7 @@ class Template:
             else:
                 prefix_ids = sep_ids + bos_ids
 
-            query_ids = self._convert_inputs_to_ids(tokenizer, context=self.prompt, query=query, idx=str(turn_idx))
+            query_ids = self._convert_inputs_to_ids(tokenizer, context=self.prompt, query=query, idx=str(turn_idx+1))
             resp_ids = self._convert_inputs_to_ids(tokenizer, context=[resp])
             encoded_pairs.append((prefix_ids + query_ids, resp_ids + eos_ids))
         return encoded_pairs
@@ -225,9 +225,6 @@ def get_template_and_fix_tokenizer(
     return template
 
 
-r"""
-Supports: https://huggingface.co/tatsu-lab/alpaca-7b-wdiff
-"""
 register_template(
     name="alpaca",
     prefix=[
@@ -246,11 +243,6 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/BAAI/AquilaChat-7B
-          https://huggingface.co/BAAI/AquilaChat2-7B
-          https://huggingface.co/BAAI/AquilaChat2-34B
-"""
 register_template(
     name="aquila",
     prefix=[
@@ -273,9 +265,6 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/baichuan-inc/Baichuan-13B-Chat
-"""
 register_template(
     name="baichuan",
     prefix=[
@@ -292,10 +281,6 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/baichuan-inc/Baichuan2-7B-Chat
-          https://huggingface.co/baichuan-inc/Baichuan2-13B-Chat
-"""
 register_template(
     name="baichuan2",
     prefix=[
@@ -312,9 +297,6 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/BelleGroup/BELLE-LLaMA-EXT-13B
-"""
 register_template(
     name="belle",
     prefix=[
@@ -330,9 +312,6 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/vivo-ai/BlueLM-7B-Chat
-"""
 register_template(
     name="bluelm",
     prefix=[
@@ -348,9 +327,6 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/THUDM/chatglm2-6b
-"""
 register_template(
     name="chatglm2",
     prefix=[
@@ -369,23 +345,54 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/THUDM/chatglm3-6b
-"""
 register_template(
     name="chatglm3",
     prefix=[
         {"token": "[gMASK]"},
         {"token": "sop"},
+        {"token": "<|system|>"},
+        "\n",
         "{{system}}"
     ],
     prompt=[
         {"token": "<|user|>"},
         "\n",
         "{{query}}",
-        {"token": "<|assistant|>"}
+        {"token": "<|assistant|>"},
+        "\n" # add an extra newline to avoid error in ChatGLM's process_response method
     ],
-    system="",
+    system=(
+        "You are ChatGLM3, a large language model trained by Zhipu.AI. "
+        "Follow the user's instructions carefully. Respond using markdown."
+    ),
+    sep=[],
+    stop_words=[
+        "<|user|>",
+        "<|observation|>"
+    ],
+    efficient_eos=True
+)
+
+
+register_template(
+    name="chatglm3_raw", # the raw template for tool tuning
+    prefix=[
+        {"token": "[gMASK]"},
+        {"token": "sop"},
+        {"token": "<|system|>"},
+        "\n",
+        "{{system}}"
+    ],
+    prompt=[
+        {"token": "<|user|>"},
+        "\n",
+        "{{query}}",
+        {"token": "<|assistant|>"}
+    ],
+    system=(
+        "You are ChatGLM3, a large language model trained by Zhipu.AI. "
+        "Follow the user's instructions carefully. Respond using markdown."
+    ),
     sep=[],
     stop_words=[
         "<|user|>",
@@ -395,11 +402,6 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/deepseek-ai/deepseek-coder-1.3b-instruct
-          https://huggingface.co/deepseek-ai/deepseek-coder-6.7b-instruct
-          https://huggingface.co/deepseek-ai/deepseek-coder-33b-instruct
-"""
 register_template(
     name="deepseek",
     prefix=[
@@ -426,9 +428,6 @@ register_template(
 )
 
 
-r"""
-Default template.
-"""
 register_template(
     name="default",
     prefix=[
@@ -447,10 +446,22 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/internlm/internlm-chat-7b
-          https://huggingface.co/internlm/internlm-chat-20b
-"""
+register_template(
+    name="falcon",
+    prefix=[
+        "{{system}}"
+    ],
+    prompt=[
+        "User: {{query}}\nFalcon:"
+    ],
+    system="",
+    sep=[
+        "\n"
+    ],
+    efficient_eos=True
+)
+
+
 register_template(
     name="intern",
     prefix=[
@@ -473,11 +484,6 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/meta-llama/Llama-2-7b-chat-hf
-          https://huggingface.co/meta-llama/Llama-2-13b-chat-hf
-          https://huggingface.co/meta-llama/Llama-2-70b-chat-hf
-"""
 register_template(
     name="llama2",
     prefix=[
@@ -500,10 +506,6 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/ziqingyang/chinese-alpaca-2-7b
-          https://huggingface.co/ziqingyang/chinese-alpaca-2-13b
-"""
 register_template(
     name="llama2_zh",
     prefix=[
@@ -517,9 +519,6 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1
-"""
 register_template(
     name="mistral",
     prefix=[
@@ -529,13 +528,12 @@ register_template(
         "[INST] {{query}} [/INST]"
     ],
     system="",
-    sep=[]
+    sep=[
+        " "
+    ]
 )
 
 
-r"""
-Supports: https://huggingface.co/openchat/openchat_3.5
-"""
 register_template(
     name="openchat",
     prefix=[
@@ -557,10 +555,6 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/Qwen/Qwen-7B-Chat
-          https://huggingface.co/Qwen/Qwen-14B-Chat
-"""
 register_template(
     name="qwen",
     prefix=[
@@ -587,10 +581,6 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/HuggingFaceH4/starchat-alpha
-          https://huggingface.co/HuggingFaceH4/starchat-beta
-"""
 register_template(
     name="starchat",
     prefix=[
@@ -631,10 +621,6 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/lmsys/vicuna-7b-v1.5
-          https://huggingface.co/lmsys/vicuna-13b-v1.5
-"""
 register_template(
     name="vicuna",
     prefix=[
@@ -651,10 +637,6 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/xverse/XVERSE-7B-Chat
-          https://huggingface.co/xverse/XVERSE-13B-Chat
-"""
 register_template(
     name="xverse",
     prefix=[
@@ -668,11 +650,6 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/wenge-research/yayi-7b
-          https://huggingface.co/wenge-research/yayi-7b-llama2
-          https://huggingface.co/wenge-research/yayi-13b-llama2
-"""
 register_template(
     name="yayi",
     prefix=[
@@ -705,10 +682,6 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/HuggingFaceH4/zephyr-7b-alpha
-          https://huggingface.co/HuggingFaceH4/zephyr-7b-beta
-"""
 register_template(
     name="zephyr",
     prefix=[
@@ -727,11 +700,6 @@ register_template(
 )
 
 
-r"""
-Supports: https://huggingface.co/IDEA-CCNL/Ziya-LLaMA-13B-v1
-          https://huggingface.co/IDEA-CCNL/Ziya-LLaMA-13B-v1.1
-          https://huggingface.co/IDEA-CCNL/Ziya2-13B-Chat
-"""
 register_template(
     name="ziya",
     prefix=[

src/llmtuner/dsets/__init__.py

@@ -1,3 +0,0 @@
-from llmtuner.dsets.loader import get_dataset
-from llmtuner.dsets.preprocess import preprocess_dataset
-from llmtuner.dsets.utils import split_dataset

src/llmtuner/eval/__init__.py

@@ -0,0 +1 @@
+from llmtuner.eval.evaluator import Evaluator

src/llmtuner/eval/evaluator.py

@@ -0,0 +1,124 @@
+# Inspired by: https://github.com/hendrycks/test/blob/master/evaluate_flan.py
+
+import os
+import json
+import torch
+import inspect
+import tiktoken
+import numpy as np
+from tqdm import tqdm, trange
+from typing import Any, Dict, List, Optional
+
+from datasets import load_dataset
+from transformers.utils import cached_file
+
+from llmtuner.data.template import get_template_and_fix_tokenizer
+from llmtuner.eval.template import get_eval_template
+from llmtuner.extras.constants import CHOICES, SUBJECTS
+from llmtuner.model import dispatch_model, get_eval_args, load_model_and_tokenizer
+
+
+class Evaluator:
+
+    def __init__(self, args: Optional[Dict[str, Any]] = None) -> None:
+        self.model_args, self.data_args, self.eval_args, finetuning_args = get_eval_args(args)
+        self.model, self.tokenizer = load_model_and_tokenizer(self.model_args, finetuning_args)
+        self.tokenizer.padding_side = "right" # avoid overflow issue in batched inference for llama2
+        self.model = dispatch_model(self.model)
+        self.template = get_template_and_fix_tokenizer(self.data_args.template, self.tokenizer)
+        self.eval_template = get_eval_template(self.eval_args.lang)
+        self.choice_inputs = self._encode_choices()
+
+    def _encode_choices(self) -> List[int]:
+        if isinstance(getattr(self.tokenizer, "tokenizer", None), tiktoken.Encoding): # for tiktoken tokenizer (Qwen)
+            kwargs = dict(allowed_special="all")
+        else:
+            kwargs = dict(add_special_tokens=False)
+
+        return [self.tokenizer.encode(self.eval_template.prefix + ch, **kwargs)[-1] for ch in CHOICES]
+
+    @torch.inference_mode()
+    def batch_inference(self, batch_input: Dict[str, torch.Tensor]) -> List[str]:
+        logits = self.model(**batch_input).logits
+        lengths = torch.sum(batch_input["attention_mask"], dim=-1)
+        word_probs = torch.stack([logits[i, lengths[i] - 1] for i in range(len(lengths))], dim=0)
+        choice_probs = torch.nn.functional.softmax(word_probs[:, self.choice_inputs], dim=-1).detach()
+        return [chr(ord("A") + offset.item()) for offset in torch.argmax(choice_probs, dim=-1)]
+
+    def eval(self) -> None:
+        if "token" in inspect.signature(cached_file).parameters:
+            kwargs = {"token": self.model_args.hf_hub_token}
+        elif "use_auth_token" in inspect.signature(cached_file).parameters: # for transformers==4.31.0
+            kwargs = {"use_auth_token": self.model_args.hf_hub_token}
+
+        mapping = cached_file(
+            path_or_repo_id = os.path.join(self.eval_args.task_dir, self.eval_args.task),
+            filename="mapping.json",
+            cache_dir=self.model_args.cache_dir,
+            **kwargs
+        )
+
+        with open(mapping, "r", encoding="utf-8") as f:
+            categorys: Dict[str, Dict[str, str]] = json.load(f)
+
+        category_corrects = {subj: np.array([], dtype="bool") for subj in SUBJECTS}
+        pbar = tqdm(categorys.keys(), desc="Processing subjects", position=0)
+        results = {}
+        for subject in pbar:
+            dataset = load_dataset(
+                path=os.path.join(self.eval_args.task_dir, self.eval_args.task),
+                name=subject,
+                cache_dir=self.model_args.cache_dir,
+                download_mode=self.eval_args.download_mode,
+                token=self.model_args.hf_hub_token
+            )
+            pbar.set_postfix_str(categorys[subject]["name"])
+            inputs, outputs, labels = [], [], []
+            for i in trange(len(dataset[self.data_args.split]), desc="Formatting batches", position=1, leave=False):
+                support_set = dataset["train"].shuffle().select(range(min(self.eval_args.n_shot, len(dataset["train"]))))
+                query, resp, history = self.eval_template.format_example(
+                    target_data=dataset[self.data_args.split][i],
+                    support_set=support_set,
+                    subject_name=categorys[subject]["name"],
+                    use_history=self.template.use_history
+                )
+                input_ids, _ = self.template.encode_oneturn(
+                    tokenizer=self.tokenizer, query=query, resp=resp, history=history
+                )
+                inputs.append({"input_ids": input_ids, "attention_mask": [1] * len(input_ids)})
+                labels.append(resp)
+
+            for i in trange(0, len(inputs), self.eval_args.batch_size, desc="Predicting batches", position=1, leave=False):
+                batch_input = self.tokenizer.pad(
+                    inputs[i : i + self.eval_args.batch_size], return_attention_mask=True, return_tensors="pt"
+                ).to(self.model.device)
+                preds = self.batch_inference(batch_input)
+                outputs += preds
+
+            corrects = (np.array(outputs) == np.array(labels))
+            category_name = categorys[subject]["category"]
+            category_corrects[category_name] = np.concatenate([category_corrects[category_name], corrects], axis=0)
+            category_corrects["Average"] = np.concatenate([category_corrects["Average"], corrects], axis=0)
+            results[subject] = {str(i): outputs[i] for i in range(len(outputs))}
+
+        pbar.close()
+        self._save_results(category_corrects, results)
+
+    def _save_results(self, category_corrects: Dict[str, np.ndarray], results: Dict[str, Dict[int, str]]) -> None:
+        score_info = "\n".join([
+            "{:>15}: {:.2f}".format(category_name, 100 * np.mean(category_correct))
+            for category_name, category_correct in category_corrects.items() if len(category_correct)
+        ])
+        print(score_info)
+        if self.eval_args.save_dir is not None:
+            os.makedirs(self.eval_args.save_dir, exist_ok=False)
+            with open(os.path.join(self.eval_args.save_dir, "results.json"), "w", encoding="utf-8", newline="\n") as f:
+                json.dump(results, f, indent=2)
+
+            with open(os.path.join(self.eval_args.save_dir, "results.log"), "w", encoding="utf-8", newline="\n") as f:
+                f.write(score_info)
+
+
+if __name__ == "__main__":
+    evaluator = Evaluator()
+    evaluator.eval()

src/llmtuner/eval/template.py

@@ -0,0 +1,86 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict, List, Tuple
+
+from llmtuner.extras.constants import CHOICES
+
+if TYPE_CHECKING:
+    from datasets import Dataset
+
+
+@dataclass
+class EvalTemplate:
+
+    system: str
+    choice: str
+    answer: str
+    prefix: str
+
+    def parse_example(
+        self,
+        example: Dict[str, str]
+    ) -> Tuple[str, str]:
+        candidates = [self.choice.format(choice=ch, content=example[ch]) for ch in CHOICES if ch in example]
+        return "".join([example["question"]] + candidates + [self.answer]), example["answer"]
+
+    def format_example(
+        self,
+        target_data: Dict[str, str],
+        support_set: "Dataset",
+        subject_name: str,
+        use_history: bool
+    ) -> Tuple[str, str, List[Tuple[str, str]]]:
+        query, resp = self.parse_example(target_data)
+        history = [self.parse_example(support_set[k]) for k in range(len(support_set))]
+
+        if len(history):
+            temp = history.pop(0)
+            history.insert(0, (self.system.format(subject=subject_name) + temp[0], temp[1]))
+        else:
+            query = self.system.format(subject=subject_name) + query
+
+        if not use_history:
+            query = "\n\n".join(["".join(item) for item in history] + [query])
+            history = []
+        return query.strip(), resp, history
+
+
+eval_templates: Dict[str, EvalTemplate] = {}
+
+
+def register_eval_template(
+    name: str,
+    system: str,
+    choice: str,
+    answer: str,
+    prefix: str
+) -> None:
+    eval_templates[name] = EvalTemplate(
+        system=system,
+        choice=choice,
+        answer=answer,
+        prefix=prefix
+    )
+
+
+def get_eval_template(name: str) -> EvalTemplate:
+    eval_template = eval_templates.get(name, None)
+    assert eval_template is not None, "Template {} does not exist.".format(name)
+    return eval_template
+
+
+register_eval_template(
+    name="en",
+    system="The following are multiple choice questions (with answers) about {subject}.\n\n",
+    choice="\n{choice}. {content}",
+    answer="\nAnswer: ",
+    prefix=" "
+)
+
+
+register_eval_template(
+    name="zh",
+    system="以下是中国关于{subject}考试的单项选择题，请选出其中的正确答案。\n\n",
+    choice="\n{choice}. {content}",
+    answer="\n答案：",
+    prefix="\n"
+)

src/llmtuner/extras/callbacks.py

@@ -12,6 +12,7 @@ from llmtuner.extras.logging import get_logger
 
 if TYPE_CHECKING:
     from transformers import TrainingArguments, TrainerState, TrainerControl
+    from trl import AutoModelForCausalLMWithValueHead
 
 
 logger = get_logger(__name__)
@@ -25,18 +26,24 @@ class SavePeftModelCallback(TrainerCallback):
         """
         if args.should_save:
             output_dir = os.path.join(args.output_dir, "{}-{}".format(PREFIX_CHECKPOINT_DIR, state.global_step))
-            model = kwargs.pop("model")
+            model: "AutoModelForCausalLMWithValueHead" = kwargs.pop("model")
+            model.pretrained_model.config.save_pretrained(output_dir)
+            if model.pretrained_model.can_generate():
+                model.pretrained_model.generation_config.save_pretrained(output_dir)
             if getattr(model, "is_peft_model", False):
-                getattr(model, "pretrained_model").save_pretrained(output_dir)
+                model.pretrained_model.save_pretrained(output_dir)
 
     def on_train_end(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
         r"""
         Event called at the end of training.
         """
         if args.should_save:
-            model = kwargs.pop("model")
+            model: "AutoModelForCausalLMWithValueHead" = kwargs.pop("model")
+            model.pretrained_model.config.save_pretrained(args.output_dir)
+            if model.pretrained_model.can_generate():
+                model.pretrained_model.generation_config.save_pretrained(args.output_dir)
             if getattr(model, "is_peft_model", False):
-                getattr(model, "pretrained_model").save_pretrained(args.output_dir)
+                model.pretrained_model.save_pretrained(args.output_dir)
 
 
 class LogCallback(TrainerCallback):

src/llmtuner/extras/constants.py

@@ -1,11 +1,25 @@
+from collections import defaultdict, OrderedDict
+from typing import Dict, Optional
+
+
+CHOICES = ["A", "B", "C", "D"]
+
+DEFAULT_MODULE = defaultdict(str)
+
+DEFAULT_TEMPLATE = defaultdict(str)
+
 IGNORE_INDEX = -100
 
+LAYERNORM_NAMES = {"norm", "ln"}
+
 LOG_FILE_NAME = "trainer_log.jsonl"
 
-LAYERNORM_NAMES = ["norm", "ln_f", "ln_attn", "ln_mlp", "ln_1", "ln_2", "ln1", "ln2"]
-
 METHODS = ["full", "freeze", "lora"]
 
+SUBJECTS = ["Average", "STEM", "Social Sciences", "Humanities", "Other"]
+
+SUPPORTED_MODELS = OrderedDict()
+
 TRAINING_STAGES = {
     "Supervised Fine-Tuning": "sft",
     "Reward Modeling": "rm",
@@ -14,79 +28,251 @@ TRAINING_STAGES = {
     "Pre-Training": "pt"
 }
 
-SUPPORTED_MODELS = {
-    "LLaMA-7B": "huggyllama/llama-7b",
-    "LLaMA-13B": "huggyllama/llama-13b",
-    "LLaMA-30B": "huggyllama/llama-30b",
-    "LLaMA-65B": "huggyllama/llama-65b",
-    "LLaMA2-7B": "meta-llama/Llama-2-7b-hf",
-    "LLaMA2-13B": "meta-llama/Llama-2-13b-hf",
-    "LLaMA2-70B": "meta-llama/Llama-2-70b-hf",
-    "LLaMA2-7B-Chat": "meta-llama/Llama-2-7b-chat-hf",
-    "LLaMA2-13B-Chat": "meta-llama/Llama-2-13b-chat-hf",
-    "LLaMA2-70B-Chat": "meta-llama/Llama-2-70b-chat-hf",
-    "ChineseLLaMA2-7B": "ziqingyang/chinese-llama-2-7b",
-    "ChineseLLaMA2-13B": "ziqingyang/chinese-llama-2-13b",
-    "ChineseLLaMA2-7B-Chat": "ziqingyang/chinese-alpaca-2-7b",
-    "ChineseLLaMA2-13B-Chat": "ziqingyang/chinese-alpaca-2-13b",
-    "BLOOM-560M": "bigscience/bloom-560m",
-    "BLOOM-3B": "bigscience/bloom-3b",
-    "BLOOM-7B1": "bigscience/bloom-7b1",
-    "BLOOMZ-560M": "bigscience/bloomz-560m",
-    "BLOOMZ-3B": "bigscience/bloomz-3b",
-    "BLOOMZ-7B1-mt": "bigscience/bloomz-7b1-mt",
-    "Falcon-7B": "tiiuae/falcon-7b",
-    "Falcon-40B": "tiiuae/falcon-40b",
-    "Falcon-7B-Chat": "tiiuae/falcon-7b-instruct",
-    "Falcon-40B-Chat": "tiiuae/falcon-40b-instruct",
-    "Baichuan-7B": "baichuan-inc/Baichuan-7B",
-    "Baichuan-13B": "baichuan-inc/Baichuan-13B-Base",
-    "Baichuan-13B-Chat": "baichuan-inc/Baichuan-13B-Chat",
-    "Baichuan2-7B": "baichuan-inc/Baichuan2-7B-Base",
-    "Baichuan2-13B": "baichuan-inc/Baichuan2-13B-Base",
-    "Baichuan2-7B-Chat": "baichuan-inc/Baichuan2-7B-Chat",
-    "Baichuan2-13B-Chat": "baichuan-inc/Baichuan2-13B-Chat",
-    "InternLM-7B": "internlm/internlm-7b",
-    "InternLM-20B": "internlm/internlm-20b",
-    "InternLM-7B-Chat": "internlm/internlm-chat-7b",
-    "InternLM-20B-Chat": "internlm/internlm-chat-20b",
-    "Qwen-7B": "Qwen/Qwen-7B",
-    "Qwen-14B": "Qwen/Qwen-14B",
-    "Qwen-7B-Chat": "Qwen/Qwen-7B-Chat",
-    "Qwen-14B-Chat": "Qwen/Qwen-14B-Chat",
-    "XVERSE-13B": "xverse/XVERSE-13B",
-    "XVERSE-13B-Chat": "xverse/XVERSE-13B-Chat",
-    "ChatGLM2-6B-Chat": "THUDM/chatglm2-6b",
-    "ChatGLM3-6B-Base": "THUDM/chatglm3-6b-base",
-    "ChatGLM3-6B-Chat": "THUDM/chatglm3-6b",
-    "Phi1.5-1.3B": "microsoft/phi-1_5"
-}
 
-DEFAULT_MODULE = {
-    "LLaMA": "q_proj,v_proj",
-    "LLaMA2": "q_proj,v_proj",
-    "ChineseLLaMA2": "q_proj,v_proj",
-    "BLOOM": "query_key_value",
-    "BLOOMZ": "query_key_value",
-    "Falcon": "query_key_value",
-    "Baichuan": "W_pack",
-    "Baichuan2": "W_pack",
-    "InternLM": "q_proj,v_proj",
-    "Qwen": "c_attn",
-    "XVERSE": "q_proj,v_proj",
-    "ChatGLM2": "query_key_value",
-    "ChatGLM3": "query_key_value",
-    "Phi1.5": "Wqkv"
-}
+def register_model_group(
+    models: Dict[str, str],
+    module: Optional[str] = None,
+    template: Optional[str] = None
+) -> None:
+    prefix = None
+    for name, path in models.items():
+        if prefix is None:
+            prefix = name.split("-")[0]
+        else:
+            assert prefix == name.split("-")[0], "prefix should be identical."
+        SUPPORTED_MODELS[name] = path
+    if module is not None:
+        DEFAULT_MODULE[prefix] = module
+    if template is not None:
+        DEFAULT_TEMPLATE[prefix] = template
 
-DEFAULT_TEMPLATE = {
-    "LLaMA2": "llama2",
-    "ChineseLLaMA2": "llama2_zh",
-    "Baichuan": "baichuan",
-    "Baichuan2": "baichuan2",
-    "InternLM": "intern",
-    "Qwen": "chatml",
-    "XVERSE": "xverse",
-    "ChatGLM2": "chatglm2",
-    "ChatGLM3": "chatglm3"
-}
+
+register_model_group(
+    models={
+        "Baichuan-7B-Base": "baichuan-inc/Baichuan-7B",
+        "Baichuan-13B-Base": "baichuan-inc/Baichuan-13B-Base",
+        "Baichuan-13B-Chat": "baichuan-inc/Baichuan-13B-Chat"
+    },
+    module="W_pack",
+    template="baichuan"
+)
+
+
+register_model_group(
+    models={
+        "Baichuan2-7B-Base": "baichuan-inc/Baichuan2-7B-Base",
+        "Baichuan2-13B-Base": "baichuan-inc/Baichuan2-13B-Base",
+        "Baichuan2-7B-Chat": "baichuan-inc/Baichuan2-7B-Chat",
+        "Baichuan2-13B-Chat": "baichuan-inc/Baichuan2-13B-Chat"
+    },
+    module="W_pack",
+    template="baichuan2"
+)
+
+
+register_model_group(
+    models={
+        "BLOOM-560M": "bigscience/bloom-560m",
+        "BLOOM-3B": "bigscience/bloom-3b",
+        "BLOOM-7B1": "bigscience/bloom-7b1"
+    },
+    module="query_key_value"
+)
+
+
+register_model_group(
+    models={
+        "BLOOMZ-560M": "bigscience/bloomz-560m",
+        "BLOOMZ-3B": "bigscience/bloomz-3b",
+        "BLOOMZ-7B1-mt": "bigscience/bloomz-7b1-mt"
+    },
+    module="query_key_value"
+)
+
+
+register_model_group(
+    models={
+        "BlueLM-7B-Base": "vivo-ai/BlueLM-7B-Base",
+        "BlueLM-7B-Chat": "vivo-ai/BlueLM-7B-Chat"
+    },
+    template="bluelm"
+)
+
+
+register_model_group(
+    models={
+        "ChatGLM2-6B-Chat": "THUDM/chatglm2-6b"
+    },
+    module="query_key_value",
+    template="chatglm2"
+)
+
+
+register_model_group(
+    models={
+        "ChatGLM3-6B-Base": "THUDM/chatglm3-6b-base",
+        "ChatGLM3-6B-Chat": "THUDM/chatglm3-6b"
+    },
+    module="query_key_value",
+    template="chatglm3"
+)
+
+
+register_model_group(
+    models={
+        "ChineseLLaMA2-1.3B": "hfl/chinese-llama-2-1.3b",
+        "ChineseLLaMA2-7B": "hfl/chinese-llama-2-7b",
+        "ChineseLLaMA2-13B": "hfl/chinese-llama-2-13b",
+        "ChineseLLaMA2-1.3B-Chat": "hfl/chinese-alpaca-2-1.3b",
+        "ChineseLLaMA2-7B-Chat": "hfl/chinese-alpaca-2-7b",
+        "ChineseLLaMA2-13B-Chat": "hfl/chinese-alpaca-2-13b"
+    },
+    template="llama2_zh"
+)
+
+
+register_model_group(
+    models={
+        "Falcon-7B": "tiiuae/falcon-7b",
+        "Falcon-40B": "tiiuae/falcon-40b",
+        "Falcon-180B": "tiiuae/falcon-180B",
+        "Falcon-7B-Chat": "tiiuae/falcon-7b-instruct",
+        "Falcon-40B-Chat": "tiiuae/falcon-40b-instruct",
+        "Falcon-180B-Chat": "tiiuae/falcon-180B-chat"
+    },
+    module="query_key_value",
+    template="falcon"
+)
+
+
+register_model_group(
+    models={
+        "InternLM-7B": "internlm/internlm-7b",
+        "InternLM-20B": "internlm/internlm-20b",
+        "InternLM-7B-Chat": "internlm/internlm-chat-7b",
+        "InternLM-20B-Chat": "internlm/internlm-chat-20b"
+    },
+    template="intern"
+)
+
+
+register_model_group(
+    models={
+        "LingoWhale-8B": "deeplang-ai/LingoWhale-8B"
+    },
+    module="qkv_proj"
+)
+
+
+register_model_group(
+    models={
+        "LLaMA-7B": "huggyllama/llama-7b",
+        "LLaMA-13B": "huggyllama/llama-13b",
+        "LLaMA-30B": "huggyllama/llama-30b",
+        "LLaMA-65B": "huggyllama/llama-65b"
+    }
+)
+
+
+register_model_group(
+    models={
+        "LLaMA2-7B": "meta-llama/Llama-2-7b-hf",
+        "LLaMA2-13B": "meta-llama/Llama-2-13b-hf",
+        "LLaMA2-70B": "meta-llama/Llama-2-70b-hf",
+        "LLaMA2-7B-Chat": "meta-llama/Llama-2-7b-chat-hf",
+        "LLaMA2-13B-Chat": "meta-llama/Llama-2-13b-chat-hf",
+        "LLaMA2-70B-Chat": "meta-llama/Llama-2-70b-chat-hf"
+    },
+    template="llama2"
+)
+
+
+register_model_group(
+    models={
+        "Mistral-7B": "mistralai/Mistral-7B-v0.1",
+        "Mistral-7B-Chat": "mistralai/Mistral-7B-Instruct-v0.1"
+    },
+    template="mistral"
+)
+
+
+register_model_group(
+    models={
+        "OpenChat3.5-7B-Chat": "openchat/openchat_3.5"
+    },
+    template="openchat"
+)
+
+
+register_model_group(
+    models={
+        "Phi1.5-1.3B": "microsoft/phi-1_5"
+    },
+    module="Wqkv"
+)
+
+
+register_model_group(
+    models={
+        "Qwen-7B": "Qwen/Qwen-7B",
+        "Qwen-14B": "Qwen/Qwen-14B",
+        "Qwen-7B-Chat": "Qwen/Qwen-7B-Chat",
+        "Qwen-14B-Chat": "Qwen/Qwen-14B-Chat"
+    },
+    module="c_attn",
+    template="qwen"
+)
+
+
+register_model_group(
+    models={
+        "Skywork-13B-Base": "Skywork/Skywork-13B-base"
+    }
+)
+
+
+register_model_group(
+    models={
+        "Vicuna1.5-7B-Chat": "lmsys/vicuna-7b-v1.5",
+        "Vicuna1.5-13B-Chat": "lmsys/vicuna-13b-v1.5"
+    },
+    template="vicuna"
+)
+
+
+register_model_group(
+    models={
+        "XVERSE-7B": "xverse/XVERSE-7B",
+        "XVERSE-13B": "xverse/XVERSE-13B",
+        "XVERSE-65B": "xverse/XVERSE-65B",
+        "XVERSE-7B-Chat": "xverse/XVERSE-7B-Chat",
+        "XVERSE-13B-Chat": "xverse/XVERSE-13B-Chat"
+    },
+    template="xverse"
+)
+
+
+register_model_group(
+    models={
+        "Yayi-7B": "wenge-research/yayi-7b-llama2",
+        "Yayi-13B": "wenge-research/yayi-13b-llama2"
+    },
+    template="yayi"
+)
+
+
+register_model_group(
+    models={
+        "Yi-6B": "01-ai/Yi-6B",
+        "Yi-34B": "01-ai/Yi-34B"
+    }
+)
+
+
+register_model_group(
+    models={
+        "Zephyr-7B-Alpha-Chat": "HuggingFaceH4/zephyr-7b-alpha",
+        "Zephyr-7B-Beta-Chat": "HuggingFaceH4/zephyr-7b-beta"
+    },
+    template="zephyr"
+)

src/llmtuner/extras/logging.py

@@ -3,6 +3,9 @@ import logging
 
 
 class LoggerHandler(logging.Handler):
+    r"""
+    Logger handler used in Web UI.
+    """
 
     def __init__(self):
         super().__init__()
@@ -19,16 +22,10 @@ class LoggerHandler(logging.Handler):
         self.log += "\n\n"
 
 
-def reset_logging():
-    r"""
-    Removes basic config of root logger
-    """
-    root = logging.getLogger()
-    list(map(root.removeHandler, root.handlers))
-    list(map(root.removeFilter, root.filters))
-
-
 def get_logger(name: str) -> logging.Logger:
+    r"""
+    Gets a standard logger with a stream hander to stdout.
+    """
     formatter = logging.Formatter(
         fmt="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
         datefmt="%m/%d/%Y %H:%M:%S"
@@ -41,3 +38,12 @@ def get_logger(name: str) -> logging.Logger:
     logger.addHandler(handler)
 
     return logger
+
+
+def reset_logging() -> None:
+    r"""
+    Removes basic config of root logger. (unused in script)
+    """
+    root = logging.getLogger()
+    list(map(root.removeHandler, root.handlers))
+    list(map(root.removeFilter, root.filters))

src/llmtuner/extras/misc.py

@@ -1,6 +1,8 @@
 import gc
+import os
+import sys
 import torch
-from typing import TYPE_CHECKING, Tuple
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
 from transformers import InfNanRemoveLogitsProcessor, LogitsProcessorList
 
 try:
@@ -11,13 +13,16 @@ try:
         is_torch_npu_available
     )
     _is_fp16_available = is_torch_npu_available() or is_torch_cuda_available()
-    _is_bf16_available = is_torch_bf16_gpu_available() or is_torch_bf16_cpu_available
+    _is_bf16_available = is_torch_bf16_gpu_available() or is_torch_bf16_cpu_available()
 except ImportError:
     _is_fp16_available = torch.cuda.is_available()
-    _is_bf16_available = torch.cuda.is_bf16_supported()
+    try:
+        _is_bf16_available = torch.cuda.is_bf16_supported()
+    except:
+        _is_bf16_available = False
 
 if TYPE_CHECKING:
-    from transformers.modeling_utils import PreTrainedModel
+    from transformers import HfArgumentParser
 
 
 class AverageMeter:
@@ -62,6 +67,24 @@ def count_parameters(model: torch.nn.Module) -> Tuple[int, int]:
     return trainable_params, all_param
 
 
+def get_current_device() -> str:
+    import accelerate
+    dummy_accelerator = accelerate.Accelerator()
+    if accelerate.utils.is_xpu_available():
+        return "xpu:{}".format(dummy_accelerator.local_process_index)
+    else:
+        return dummy_accelerator.local_process_index if torch.cuda.is_available() else "cpu"
+
+
+def get_logits_processor() -> "LogitsProcessorList":
+    r"""
+    Gets logits processor that removes NaN and Inf logits.
+    """
+    logits_processor = LogitsProcessorList()
+    logits_processor.append(InfNanRemoveLogitsProcessor())
+    return logits_processor
+
+
 def infer_optim_dtype(model_dtype: torch.dtype) -> torch.dtype:
     r"""
     Infers the optimal dtype according to the model_dtype and device compatibility.
@@ -74,13 +97,15 @@ def infer_optim_dtype(model_dtype: torch.dtype) -> torch.dtype:
         return torch.float32
 
 
-def get_logits_processor() -> LogitsProcessorList:
-    r"""
-    Gets logits processor that removes NaN and Inf logits.
-    """
-    logits_processor = LogitsProcessorList()
-    logits_processor.append(InfNanRemoveLogitsProcessor())
-    return logits_processor
+def parse_args(parser: "HfArgumentParser", args: Optional[Dict[str, Any]] = None) -> Tuple[Any]:
+    if args is not None:
+        return parser.parse_dict(args)
+    elif len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"):
+        return parser.parse_yaml_file(os.path.abspath(sys.argv[1]))
+    elif len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        return parser.parse_json_file(os.path.abspath(sys.argv[1]))
+    else:
+        return parser.parse_args_into_dataclasses()
 
 
 def torch_gc() -> None:
@@ -91,28 +116,3 @@ def torch_gc() -> None:
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
         torch.cuda.ipc_collect()
-
-
-def dispatch_model(model: "PreTrainedModel") -> "PreTrainedModel":
-    r"""
-    Dispatches a pre-trained model to GPUs with balanced memory.
-    Borrowed from: https://github.com/huggingface/transformers/blob/v4.31.0/src/transformers/modeling_utils.py#L2803
-    """
-    if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False): # do nothing
-        return model
-
-    if torch.cuda.device_count() > 1:
-        from accelerate import dispatch_model
-        from accelerate.utils import infer_auto_device_map, get_balanced_memory
-
-        if model._no_split_modules is None:
-            raise ValueError("The model class needs to implement the `_no_split_modules` attribute.")
-
-        kwargs = {"dtype": model.dtype, "no_split_module_classes": model._no_split_modules}
-        max_memory = get_balanced_memory(model, **kwargs)
-        # Make sure tied weights are tied before creating the device map.
-        model.tie_weights()
-        device_map = infer_auto_device_map(model, max_memory=max_memory, **kwargs)
-        return dispatch_model(model, device_map)
-    else:
-        return model.cuda()

src/llmtuner/extras/packages.py

@@ -0,0 +1,55 @@
+import importlib.metadata
+import importlib.util
+
+
+def is_package_available(name: str) -> bool:
+    return importlib.util.find_spec(name) is not None
+
+
+def get_package_version(name: str) -> str:
+    try:
+        return importlib.metadata.version(name)
+    except:
+        return "0.0.0"
+
+
+_fastapi_available = is_package_available("fastapi")
+_flash_attn2_available = is_package_available("flash_attn") and get_package_version("flash_attn").startswith("2")
+_jieba_available = is_package_available("jieba")
+_matplotlib_available = is_package_available("matplotlib")
+_nltk_available = is_package_available("nltk")
+_rouge_available = is_package_available("rouge_chinese")
+_starlette_available = is_package_available("sse_starlette")
+_uvicorn_available = is_package_available("uvicorn")
+
+
+def is_fastapi_availble():
+    return _fastapi_available
+
+
+def is_flash_attn2_available():
+    return _flash_attn2_available
+
+
+def is_jieba_available():
+    return _jieba_available
+
+
+def is_matplotlib_available():
+    return _matplotlib_available
+
+
+def is_nltk_available():
+    return _nltk_available
+
+
+def is_rouge_available():
+    return _rouge_available
+
+
+def is_starlette_available():
+    return _starlette_available
+
+
+def is_uvicorn_available():
+    return _uvicorn_available

src/llmtuner/extras/patches/llama_patch.py

@@ -3,13 +3,19 @@ import torch
 import torch.nn as nn
 from typing import Optional, Tuple
 from transformers.utils import logging
-from transformers.models.llama.modeling_llama import LlamaAttention, apply_rotary_pos_emb, repeat_kv
+from transformers.models.llama.modeling_llama import LlamaAttention, apply_rotary_pos_emb
 
 try:
+    from transformers.models.llama.modeling_llama import repeat_kv
+except ImportError:
+    print("Please upgrade `transformers`.")
+
+from llmtuner.extras.packages import is_flash_attn2_available
+
+
+if is_flash_attn2_available():
     from flash_attn import flash_attn_func, flash_attn_varlen_func # type: ignore
     from flash_attn.bert_padding import pad_input, unpad_input # type: ignore
-except ImportError:
-    print("FlashAttention-2 is not installed, ignore this if you are not using FlashAttention.")
 
 
 logger = logging.get_logger(__name__)

src/llmtuner/extras/ploting.py

@@ -1,11 +1,14 @@
 import os
 import math
 import json
-import matplotlib.pyplot as plt
 from typing import List, Optional
 from transformers.trainer import TRAINER_STATE_NAME
 
 from llmtuner.extras.logging import get_logger
+from llmtuner.extras.packages import is_matplotlib_available
+
+if is_matplotlib_available():
+    import matplotlib.pyplot as plt
 
 
 logger = get_logger(__name__)

src/llmtuner/hparams/__init__.py

@@ -1,4 +1,5 @@
 from .data_args import DataArguments
+from .evaluation_args import EvaluationArguments
 from .finetuning_args import FinetuningArguments
 from .generating_args import GeneratingArguments
 from .model_args import ModelArguments

src/llmtuner/hparams/data_args.py

@@ -42,7 +42,7 @@ class DataArguments:
     )
     dataset_dir: Optional[str] = field(
         default="data",
-        metadata={"help": "The name of the folder containing datasets."}
+        metadata={"help": "Path to the folder containing the datasets."}
     )
     split: Optional[str] = field(
         default="train",
@@ -52,6 +52,10 @@ class DataArguments:
         default=1024,
         metadata={"help": "The maximum length of the model inputs after tokenization."}
     )
+    reserved_label_len: Optional[int] = field(
+        default=1,
+        metadata={"help": "The maximum length reserved for label after tokenization."}
+    )
     train_on_prompt: Optional[bool] = field(
         default=False,
         metadata={"help": "Whether to disable the mask on the prompt or not."}
@@ -110,6 +114,9 @@ class DataArguments:
     )
 
     def __post_init__(self):
+        if self.reserved_label_len >= self.cutoff_len:
+            raise ValueError("`reserved_label_len` must be smaller than `cutoff_len`.")
+
         if self.streaming and self.val_size > 1e-6 and self.val_size < 1:
             raise ValueError("Streaming mode should have an integer val size.")
 

src/llmtuner/hparams/evaluation_args.py

@@ -0,0 +1,55 @@
+import os
+from typing import Literal, Optional
+from dataclasses import dataclass, field
+
+from datasets import DownloadMode
+
+
+@dataclass
+class EvaluationArguments:
+    r"""
+    Arguments pertaining to specify the evaluation parameters.
+    """
+    task: str = field(
+        metadata={"help": "Name of the evaluation task."}
+    )
+    task_dir: Optional[str] = field(
+        default="evaluation",
+        metadata={"help": "Path to the folder containing the evaluation datasets."}
+    )
+    batch_size: Optional[int] = field(
+        default=4,
+        metadata={"help": "The batch size per GPU for evaluation."}
+    )
+    seed: Optional[int] = field(
+        default=42,
+        metadata={"help": "Random seed to be used with data loaders."}
+    )
+    lang: Optional[Literal["en", "zh"]] = field(
+        default="en",
+        metadata={"help": "Language used at evaluation."}
+    )
+    n_shot: Optional[int] = field(
+        default=5,
+        metadata={"help": "Number of examplars for few-shot learning."}
+    )
+    save_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Path to save the evaluation results."}
+    )
+    download_mode: Optional[DownloadMode] = field(
+        default=DownloadMode.REUSE_DATASET_IF_EXISTS,
+        metadata={"help": "Download mode used for the evaluation datasets."}
+    )
+
+    def __post_init__(self):
+        task_available = []
+        for folder in os.listdir(self.task_dir):
+            if os.path.isdir(os.path.join(self.task_dir, folder)):
+                task_available.append(folder)
+
+        if self.task not in task_available:
+            raise ValueError("Task {} not found in {}.".format(self.task, self.task_dir))
+
+        if self.save_dir is not None and os.path.exists(self.save_dir):
+            raise ValueError("`save_dir` already exists, use another one.")

src/llmtuner/hparams/finetuning_args.py

@@ -4,7 +4,128 @@ from dataclasses import asdict, dataclass, field
 
 
 @dataclass
-class FinetuningArguments:
+class FreezeArguments:
+    r"""
+    Arguments pertaining to the freeze (partial-parameter) training.
+    """
+    name_module_trainable: Optional[str] = field(
+        default="mlp",
+        metadata={"help": "Name of trainable modules for partial-parameter (freeze) fine-tuning. \
+                  Use commas to separate multiple modules. \
+                  LLaMA choices: [\"mlp\", \"self_attn\"], \
+                  BLOOM & Falcon & ChatGLM choices: [\"mlp\", \"self_attention\"], \
+                  Qwen choices: [\"mlp\", \"attn\"], \
+                  Phi-1.5 choices: [\"mlp\", \"mixer\"], \
+                  Others choices: the same as LLaMA."}
+    )
+    num_layer_trainable: Optional[int] = field(
+        default=3,
+        metadata={"help": "The number of trainable layers for partial-parameter (freeze) fine-tuning."}
+    )
+
+
+@dataclass
+class LoraArguments:
+    r"""
+    Arguments pertaining to the LoRA training.
+    """
+    additional_target: Optional[str] = field(
+        default=None,
+        metadata={"help": "Name(s) of modules apart from LoRA layers to be set as trainable and saved in the final checkpoint."}
+    )
+    lora_alpha: Optional[float] = field(
+        default=None,
+        metadata={"help": "The scale factor for LoRA fine-tuning (default: lora_rank * 2.0)."}
+    )
+    lora_dropout: Optional[float] = field(
+        default=0.1,
+        metadata={"help": "Dropout rate for the LoRA fine-tuning."}
+    )
+    lora_rank: Optional[int] = field(
+        default=8,
+        metadata={"help": "The intrinsic dimension for LoRA fine-tuning."}
+    )
+    lora_target: Optional[str] = field(
+        default=None,
+        metadata={"help": "Name(s) of target modules to apply LoRA. Use commas to separate multiple modules. \
+                  LLaMA choices: [\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"], \
+                  BLOOM & Falcon & ChatGLM choices: [\"query_key_value\", \"dense\", \"dense_h_to_4h\", \"dense_4h_to_h\"], \
+                  Baichuan choices: [\"W_pack\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"], \
+                  Qwen choices: [\"c_attn\", \"attn.c_proj\", \"w1\", \"w2\", \"mlp.c_proj\"], \
+                  Phi-1.5 choices: [\"Wqkv\", \"out_proj\", \"fc1\", \"fc2\"], \
+                  Others choices: the same as LLaMA."}
+    )
+    resume_lora_training: Optional[bool] = field(
+        default=True,
+        metadata={"help": "Whether to resume training from the last LoRA weights or create new weights after merging them."}
+    )
+
+
+@dataclass
+class RLHFArguments:
+    r"""
+    Arguments pertaining to the PPO and DPO training.
+    """
+    dpo_beta: Optional[float] = field(
+        default=0.1,
+        metadata={"help": "The beta parameter for the DPO loss."}
+    )
+    ppo_buffer_size: Optional[int] = field(
+        default=1,
+        metadata={"help": "The number of mini-batches to make experience buffer in a PPO optimization step."}
+    )
+    ppo_epochs: Optional[int] = field(
+        default=4,
+        metadata={"help": "The number of epochs to perform in a PPO optimization step."}
+    )
+    ppo_logger: Optional[str] = field(
+        default=None,
+        metadata={"help": "Log with either \"wandb\" or \"tensorboard\" in PPO training."}
+    )
+    ppo_score_norm: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Use score normalization in PPO training."}
+    )
+    ppo_target: Optional[float] = field(
+        default=6.0,
+        metadata={"help": "Target KL value for adaptive KL control in PPO training."}
+    )
+    ppo_whiten_rewards: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Whiten the rewards before compute advantages in PPO training."}
+    )
+    ref_model: Optional[str] = field(
+        default=None,
+        metadata={"help": "Path to the reference model used for the PPO or DPO training."}
+    )
+    ref_model_checkpoint: Optional[str] = field(
+        default=None,
+        metadata={"help": "Path to the directory(s) containing the model checkpoints of the reference model."}
+    )
+    ref_model_quantization_bit: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of bits to quantize the reference model."}
+    )
+    reward_model: Optional[str] = field(
+        default=None,
+        metadata={"help": "Path to the directory containing the checkpoints of the reward model."}
+    )
+    reward_model_checkpoint: Optional[str] = field(
+        default=None,
+        metadata={"help": "Path to the directory(s) containing the model checkpoints of the reward model."}
+    )
+    reward_model_quantization_bit: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of bits to quantize the reward model."}
+    )
+    reward_model_type: Optional[Literal["lora", "full"]] = field(
+        default="lora",
+        metadata={"help": "The checkpoint type of the reward model. The lora type only supports lora training."}
+    )
+
+
+@dataclass
+class FinetuningArguments(FreezeArguments, LoraArguments, RLHFArguments):
     r"""
     Arguments pertaining to which techniques we are going to fine-tuning with.
     """
@@ -12,77 +133,10 @@ class FinetuningArguments:
         default="sft",
         metadata={"help": "Which stage will be performed in training."}
     )
-    finetuning_type: Optional[Literal["lora", "freeze", "full", "none"]] = field(
+    finetuning_type: Optional[Literal["lora", "freeze", "full"]] = field(
         default="lora",
         metadata={"help": "Which fine-tuning method to use."}
     )
-    num_layer_trainable: Optional[int] = field(
-        default=3,
-        metadata={"help": "Number of trainable layers for partial-parameter (freeze) fine-tuning."}
-    )
-    name_module_trainable: Optional[Literal["mlp", "self_attn", "self_attention"]] = field(
-        default="mlp",
-        metadata={"help": "Name of trainable modules for partial-parameter (freeze) fine-tuning. \
-                  LLaMA choices: [\"mlp\", \"self_attn\"], \
-                  BLOOM & Falcon & ChatGLM choices: [\"mlp\", \"self_attention\"], \
-                  Qwen choices: [\"mlp\", \"attn\"], \
-                  Phi-1.5 choices: [\"mlp\", \"mixer\"], \
-                  LLaMA-2, BlueLM, Baichuan, InternLM, Mistral, Skywork, XVERSE, Yi choices: the same as LLaMA."}
-    )
-    lora_rank: Optional[int] = field(
-        default=8,
-        metadata={"help": "The intrinsic dimension for LoRA fine-tuning."}
-    )
-    lora_alpha: Optional[float] = field(
-        default=32.0,
-        metadata={"help": "The scale factor for LoRA fine-tuning (similar with the learning rate)."}
-    )
-    lora_dropout: Optional[float] = field(
-        default=0.1,
-        metadata={"help": "Dropout rate for the LoRA fine-tuning."}
-    )
-    lora_target: Optional[str] = field(
-        default=None,
-        metadata={"help": "Name(s) of target modules to apply LoRA. Use commas to separate multiple modules. \
-                  LLaMA choices: [\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"], \
-                  BLOOM & Falcon & ChatGLM choices: [\"query_key_value\", \"self_attention.dense\", \"mlp.dense\"], \
-                  Baichuan choices: [\"W_pack\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"], \
-                  Qwen choices: [\"c_attn\", \"attn.c_proj\", \"w1\", \"w2\", \"mlp.c_proj\"], \
-                  Phi-1.5 choices: [\"Wqkv\", \"out_proj\", \"fc1\", \"fc2\"], \
-                  LLaMA-2, BlueLM, InternLM, Mistral, Skywork, XVERSE, Yi choices: the same as LLaMA."}
-    )
-    additional_target: Optional[str] = field(
-        default=None,
-        metadata={"help": "Name(s) of modules apart from LoRA layers to be set as trainable and saved in the final checkpoint."}
-    )
-    resume_lora_training: Optional[bool] = field(
-        default=True,
-        metadata={"help": "Whether to resume training from the last LoRA weights or create new weights after merging them."}
-    )
-    ppo_score_norm: Optional[bool] = field(
-        default=False,
-        metadata={"help": "Use score normalization in PPO training."}
-    )
-    ppo_logger: Optional[str] = field(
-        default=None,
-        metadata={"help": "Log with either 'wandb' or 'tensorboard' in PPO training."}
-    )
-    ppo_target: Optional[float] = field(
-        default=6.0,
-        metadata={"help": "Target KL value for adaptive KL control in PPO training."}
-    )
-    dpo_beta: Optional[float] = field(
-        default=0.1,
-        metadata={"help": "The beta parameter for the DPO loss."}
-    )
-    dpo_ref_model: Optional[str] = field(
-        default=None,
-        metadata={"help": "Path to the reference model used for the DPO training."}
-    )
-    dpo_ref_model_checkpoint: Optional[str] = field(
-        default=None,
-        metadata={"help": "Path to the directory(s) containing the model checkpoints of the reference model."}
-    )
     upcast_layernorm: Optional[bool] = field(
         default=False,
         metadata={"help": "Whether to upcast the layernorm weights in fp32."}
@@ -91,15 +145,37 @@ class FinetuningArguments:
         default=0,
         metadata={"help": "The alpha parameter to control the noise magnitude in NEFTune."}
     )
+    export_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Path to the directory to save the exported model."}
+    )
+    plot_loss: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Whether to plot the training loss after fine-tuning or not."}
+    )
 
     def __post_init__(self):
-        if isinstance(self.lora_target, str): # support custom target modules/layers of LoRA
-            self.lora_target = [target.strip() for target in self.lora_target.split(",")]
+        def split_arg(arg):
+            if isinstance(arg, str):
+                return [item.strip() for item in arg.split(",")]
+            return arg
 
-        if isinstance(self.additional_target, str):
-            self.additional_target = [target.strip() for target in self.additional_target.split(",")]
+        self.name_module_trainable = split_arg(self.name_module_trainable)
+        self.lora_alpha = self.lora_alpha or float(self.lora_rank * 2.0)
+        self.lora_target = split_arg(self.lora_target)
+        self.additional_target = split_arg(self.additional_target)
+        self.ref_model_checkpoint = split_arg(self.ref_model_checkpoint)
+        self.reward_model_checkpoint = split_arg(self.reward_model_checkpoint)
 
         assert self.finetuning_type in ["lora", "freeze", "full"], "Invalid fine-tuning method."
+        assert self.ref_model_quantization_bit in [None, 8, 4], "We only accept 4-bit or 8-bit quantization."
+        assert self.reward_model_quantization_bit in [None, 8, 4], "We only accept 4-bit or 8-bit quantization."
+
+        if self.stage == "ppo" and self.reward_model is None:
+            raise ValueError("Reward model is necessary for PPO training.")
+
+        if self.stage == "ppo" and self.reward_model_type == "lora" and self.finetuning_type != "lora":
+            raise ValueError("Lora reward model only supports lora training.")
 
     def save_to_json(self, json_path: str):
         r"""Saves the content of this instance in JSON format inside `json_path`."""
@@ -112,4 +188,5 @@ class FinetuningArguments:
         r"""Creates an instance from the content of `json_path`."""
         with open(json_path, "r", encoding="utf-8") as f:
             text = f.read()
+
         return cls(**json.loads(text))

src/llmtuner/hparams/model_args.py

@@ -54,22 +54,10 @@ class ModelArguments:
         default=False,
         metadata={"help": "Enable shift short attention (S^2-Attn) proposed by LongLoRA."}
     )
-    reward_model: Optional[str] = field(
-        default=None,
-        metadata={"help": "Path to the directory containing the checkpoints of the reward model."}
-    )
-    plot_loss: Optional[bool] = field(
-        default=False,
-        metadata={"help": "Whether to plot the training loss after fine-tuning or not."}
-    )
     hf_hub_token: Optional[str] = field(
         default=None,
         metadata={"help": "Auth token to log in with Hugging Face Hub."}
     )
-    export_dir: Optional[str] = field(
-        default=None,
-        metadata={"help": "Path to the directory to save the exported model."}
-    )
 
     def __post_init__(self):
         self.compute_dtype = None
@@ -81,8 +69,7 @@ class ModelArguments:
         if self.checkpoint_dir is not None: # support merging multiple lora weights
             self.checkpoint_dir = [cd.strip() for cd in self.checkpoint_dir.split(",")]
 
-        if self.quantization_bit is not None:
-            assert self.quantization_bit in [4, 8], "We only accept 4-bit or 8-bit quantization."
+        assert self.quantization_bit in [None, 8, 4], "We only accept 4-bit or 8-bit quantization."
 
     def to_dict(self) -> Dict[str, Any]:
         return asdict(self)

src/llmtuner/model/__init__.py

@@ -0,0 +1,5 @@
+# Level: loader > adapter > parser, utils
+
+from llmtuner.model.loader import load_model_and_tokenizer
+from llmtuner.model.parser import get_train_args, get_infer_args, get_eval_args
+from llmtuner.model.utils import dispatch_model, get_modelcard_args, load_valuehead_params

src/llmtuner/model/adapter.py

@@ -1,18 +1,9 @@
-import os
 import torch
 from typing import TYPE_CHECKING
-
-from transformers.utils import cached_file
-from transformers.trainer import WEIGHTS_NAME, SAFE_WEIGHTS_NAME
-from peft import (
-    PeftModel,
-    TaskType,
-    LoraConfig,
-    get_peft_model
-)
+from peft import PeftModel, TaskType, LoraConfig, get_peft_model
 
 from llmtuner.extras.logging import get_logger
-from llmtuner.tuner.core.utils import find_all_linear_modules
+from llmtuner.model.utils import find_all_linear_modules
 
 if TYPE_CHECKING:
     from transformers.modeling_utils import PreTrainedModel
@@ -38,20 +29,31 @@ def init_adapter(
 
     if (not is_trainable) and model_args.checkpoint_dir is None:
         logger.info("Checkpoint is not found at evaluation, load the original model.")
+        return model
 
     if finetuning_args.finetuning_type == "full" and is_trainable:
         logger.info("Fine-tuning method: Full")
         model = model.float()
 
-    if finetuning_args.finetuning_type == "freeze":
+    if finetuning_args.finetuning_type == "freeze" and is_trainable:
         logger.info("Fine-tuning method: Freeze")
-        num_layers = getattr(model.config, "num_layers")
+        num_layers = (
+            getattr(model.config, "num_hidden_layers", None)
+            or getattr(model.config, "num_layers", None)
+            or getattr(model.config, "n_layer", None)
+        )
+        if not num_layers:
+            raise ValueError("Current model does not support freeze tuning.")
         if finetuning_args.num_layer_trainable > 0: # fine-tuning the last n layers if num_layer_trainable > 0
             trainable_layer_ids = [num_layers - k - 1 for k in range(finetuning_args.num_layer_trainable)]
         else: # fine-tuning the first n layers if num_layer_trainable < 0
             trainable_layer_ids = [k for k in range(-finetuning_args.num_layer_trainable)]
 
-        trainable_layers = ["{:d}.{}".format(idx, finetuning_args.name_module_trainable) for idx in trainable_layer_ids]
+        trainable_layers = []
+        for module_name in finetuning_args.name_module_trainable:
+            for idx in trainable_layer_ids:
+                trainable_layers.append("{:d}.{}".format(idx, module_name))
+
         for name, param in model.named_parameters():
             if not any(trainable_layer in name for trainable_layer in trainable_layers):
                 param.requires_grad_(False)
@@ -63,7 +65,12 @@ def init_adapter(
         checkpoint_to_resume = None
 
         if model_args.checkpoint_dir is not None:
-            if is_trainable and finetuning_args.resume_lora_training:
+            is_mergeable = True
+            if getattr(model, "quantization_method", None) == "gptq":
+                assert len(model_args.checkpoint_dir) == 1, "GPTQ quantized model only accepts a single checkpoint."
+                is_mergeable = False
+
+            if (is_trainable and finetuning_args.resume_lora_training) or (not is_mergeable):
                 checkpoints_to_merge, checkpoint_to_resume = model_args.checkpoint_dir[:-1], model_args.checkpoint_dir[-1]
             else:
                 checkpoints_to_merge = model_args.checkpoint_dir
@@ -99,30 +106,3 @@ def init_adapter(
         logger.info("Loaded fine-tuned model from checkpoint(s): {}".format(",".join(model_args.checkpoint_dir)))
 
     return model
-
-
-def load_valuehead_params(
-    model: "PreTrainedModel",
-    model_args: "ModelArguments"
-) -> bool:
-    kwargs = {
-        "path_or_repo_id": model_args.reward_model,
-        "cache_dir": model_args.cache_dir,
-        "token": model_args.hf_hub_token,
-        "revision": model_args.model_revision
-    }
-    try:
-        vhead_file = cached_file(filename=WEIGHTS_NAME, **kwargs)
-    except:
-        try:
-            vhead_file = cached_file(filename=SAFE_WEIGHTS_NAME, **kwargs)
-        except:
-            logger.warning("Provided path ({}) does not contain valuehead weights.".format(model_args.reward_model))
-            return False
-
-    vhead_params = torch.load(vhead_file, map_location="cpu")
-    model.register_buffer("reward_head_weight", vhead_params["v_head.summary.weight"], persistent=False)
-    model.register_buffer("reward_head_bias", vhead_params["v_head.summary.bias"], persistent=False)
-    model.register_buffer("default_head_weight", torch.zeros_like(vhead_params["v_head.summary.weight"]), persistent=False)
-    model.register_buffer("default_head_bias", torch.zeros_like(vhead_params["v_head.summary.bias"]), persistent=False)
-    return True

src/llmtuner/model/loader.py

@@ -1,4 +1,3 @@
-import os
 import math
 import torch
 from types import MethodType
@@ -15,7 +14,6 @@ from transformers import (
 )
 from transformers.models.llama import modeling_llama as LlamaModule
 from transformers.utils.versions import require_version
-from peft import PeftModel
 from trl import AutoModelForCausalLMWithValueHead
 
 try:
@@ -24,11 +22,12 @@ except ImportError: # https://github.com/huggingface/transformers/releases/tag/v
     from transformers.deepspeed import is_deepspeed_zero3_enabled
 
 from llmtuner.extras.logging import reset_logging, get_logger
-from llmtuner.extras.misc import count_parameters, infer_optim_dtype
+from llmtuner.extras.misc import count_parameters, get_current_device, infer_optim_dtype
+from llmtuner.extras.packages import is_flash_attn2_available
 from llmtuner.extras.patches import llama_patch as LlamaPatches
 from llmtuner.hparams import FinetuningArguments
-from llmtuner.tuner.core.adapter import init_adapter, load_valuehead_params
-from llmtuner.tuner.core.utils import prepare_model_for_training
+from llmtuner.model.adapter import init_adapter
+from llmtuner.model.utils import load_valuehead_params, prepare_model_for_training
 
 if TYPE_CHECKING:
     from transformers import PreTrainedTokenizer
@@ -42,7 +41,7 @@ require_version("transformers>=4.31.0,<4.35.0", "To fix: pip install \"transform
 require_version("datasets>=2.14.0", "To fix: pip install datasets>=2.14.0")
 require_version("accelerate>=0.21.0", "To fix: pip install accelerate>=0.21.0")
 require_version("peft>=0.6.0", "To fix: pip install peft>=0.6.0")
-require_version("trl==0.7.2", "To fix: pip install trl==0.7.2")
+require_version("trl>=0.7.4", "To fix: pip install trl>=0.7.4")
 
 
 def load_model_and_tokenizer(
@@ -73,6 +72,7 @@ def load_model_and_tokenizer(
     )
 
     if finetuning_args.finetuning_type != "lora" and model_args.checkpoint_dir is not None:
+        logger.info("Use `model_name_or_path` to specify the model trained with full/freeze method.")
         model_to_load = model_args.checkpoint_dir[0]
     else:
         model_to_load = model_args.model_name_or_path
@@ -84,10 +84,9 @@ def load_model_and_tokenizer(
         tokenizer._pad = MethodType(PreTrainedTokenizerBase._pad, tokenizer)
 
     # Set model dtype
-    if model_args.compute_dtype is not None: # for training
-        setattr(config, "torch_dtype", model_args.compute_dtype)
-    else: # for evaluation, priority: bf16 > fp16 > fp32
+    if model_args.compute_dtype is None: # priority: bf16 > fp16 > fp32
         model_args.compute_dtype = infer_optim_dtype(model_dtype=getattr(config, "torch_dtype", None))
+    setattr(config, "torch_dtype", model_args.compute_dtype)
 
     # Fix config (for Qwen)
     if getattr(config, "model_type", None) == "qwen":
@@ -123,13 +122,16 @@ def load_model_and_tokenizer(
     # Set FlashAttention-2
     if model_args.flash_attn:
         if getattr(config, "model_type", None) == "llama":
-            LlamaModule.LlamaAttention = LlamaPatches.LlamaFlashAttention2
-            LlamaModule.LlamaModel._prepare_decoder_attention_mask = LlamaPatches._prepare_decoder_attention_mask
-            logger.info("Using FlashAttention-2 for faster training and inference.")
+            if is_flash_attn2_available():
+                LlamaModule.LlamaAttention = LlamaPatches.LlamaFlashAttention2
+                LlamaModule.LlamaModel._prepare_decoder_attention_mask = LlamaPatches._prepare_decoder_attention_mask
+                logger.info("Using FlashAttention-2 for faster training and inference.")
+            else:
+                logger.warning("FlashAttention-2 is not installed.")
         elif getattr(config, "model_type", None) in ["qwen", "Yi"]:
             logger.info("Current model automatically enables FlashAttention if installed.")
         else:
-            logger.warning("Current model does not support FlashAttention-2.")
+            logger.warning("Current model does not support FlashAttention.")
     elif is_trainable and model_args.shift_attn and getattr(config, "model_type", None) == "llama":
         LlamaModule.LlamaAttention = LlamaPatches.LlamaShiftShortAttention
         logger.warning("Using `--flash_attn` for faster training in large context length.")
@@ -142,7 +144,7 @@ def load_model_and_tokenizer(
         else:
             logger.warning("Current model does not support shift short attention.")
 
-    # Quantization configurations (using bitsandbytes library).
+    # Quantization configurations (using bitsandbytes library)
     if model_args.quantization_bit is not None:
         if is_deepspeed_zero3_enabled():
             raise ValueError("DeepSpeed ZeRO-3 is incompatible with quantization.")
@@ -162,10 +164,10 @@ def load_model_and_tokenizer(
                 bnb_4bit_quant_type=model_args.quantization_type
             )
 
-        config_kwargs["device_map"] = {"": int(os.environ.get("LOCAL_RANK", "0"))} if is_trainable else "auto"
+        config_kwargs["device_map"] = {"": get_current_device()}
         logger.info("Quantizing model to {} bit.".format(model_args.quantization_bit))
 
-    # Load and prepare pre-trained models (without valuehead).
+    # Load pre-trained models (without valuehead)
     model = AutoModelForCausalLM.from_pretrained(
         model_to_load,
         config=config,
@@ -183,7 +185,7 @@ def load_model_and_tokenizer(
         setattr(model, "lm_head", model.transformer.output_layer)
         setattr(model, "_keys_to_ignore_on_save", ["lm_head.weight"])
 
-    # Register auto class to save the custom code files.
+    # Register auto class to save the custom code files
     if isinstance(config, PretrainedConfig) and "AutoConfig" in getattr(config, "auto_map", {}):
         config.__class__.register_for_auto_class()
     if isinstance(model, PreTrainedModel) and "AutoModelForCausalLM" in getattr(config, "auto_map", {}):
@@ -197,25 +199,17 @@ def load_model_and_tokenizer(
     model = model.train() if is_trainable else model.eval()
 
     # Prepare model with valuehead for RLHF
-    if stage == "rm" or stage == "ppo":
+    if stage in ["rm", "ppo"]:
         model: "AutoModelForCausalLMWithValueHead" = AutoModelForCausalLMWithValueHead.from_pretrained(model)
-        reset_logging()
-        if stage == "rm" and model_args.checkpoint_dir is not None: # load valuehead weights to evaluate reward model
-            logger.warning("Only the last checkpoint containing valuehead will be loaded.")
-            if load_valuehead_params(model, model_args):
-                model.v_head.load_state_dict({
-                    "summary.weight": getattr(model, "reward_head_weight"),
-                    "summary.bias": getattr(model, "reward_head_bias")
-                })
-
-        if stage == "ppo": # load reward model
-            logger.info("Load reward model from {}".format(model_args.reward_model))
-            if isinstance(model.pretrained_model, PeftModel):
-                model.pretrained_model.load_adapter(model_args.reward_model, "reward")
-            for name, param in model.named_parameters(): # https://github.com/huggingface/peft/issues/1090
-                if "default" in name:
-                    param.data = param.data.to(torch.float32) # trainable params should in fp32
-            assert load_valuehead_params(model, model_args), "Reward model is not correctly loaded."
+        setattr(model, "_keys_to_ignore_on_save", [name for name, _ in model.named_parameters() if "pretrained_model" in name])
+        setattr(model, "tie_weights", MethodType(lambda _: None, model)) # use empty method
+        vhead_path = (
+            model_args.checkpoint_dir[-1] if model_args.checkpoint_dir is not None else model_args.model_name_or_path
+        )
+        vhead_params = load_valuehead_params(vhead_path, model_args)
+        if vhead_params is not None:
+            model.load_state_dict(vhead_params, strict=False)
+            logger.info("Loaded valuehead from checkpoint: {}".format(vhead_path))
 
     # Prepare model for inference
     if not is_trainable:

src/llmtuner/model/parser.py

@@ -1,5 +1,4 @@
 import os
-import sys
 import torch
 import datasets
 import transformers
@@ -8,9 +7,11 @@ from transformers import HfArgumentParser, Seq2SeqTrainingArguments
 from transformers.trainer_utils import get_last_checkpoint
 
 from llmtuner.extras.logging import get_logger
+from llmtuner.extras.misc import parse_args
 from llmtuner.hparams import (
     ModelArguments,
     DataArguments,
+    EvaluationArguments,
     FinetuningArguments,
     GeneratingArguments
 )
@@ -19,62 +20,54 @@ from llmtuner.hparams import (
 logger = get_logger(__name__)
 
 
-def _parse_args(parser: HfArgumentParser, args: Optional[Dict[str, Any]] = None) -> Tuple[Any]:
-    if args is not None:
-        return parser.parse_dict(args)
-    elif len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"):
-        return parser.parse_yaml_file(os.path.abspath(sys.argv[1]))
-    elif len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
-        return parser.parse_json_file(os.path.abspath(sys.argv[1]))
-    else:
-        return parser.parse_args_into_dataclasses()
+_TRAIN_ARGS = [
+    ModelArguments, DataArguments, Seq2SeqTrainingArguments, FinetuningArguments, GeneratingArguments
+]
+_TRAIN_CLS = Tuple[
+    ModelArguments, DataArguments, Seq2SeqTrainingArguments, FinetuningArguments, GeneratingArguments
+]
+_INFER_ARGS = [
+    ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments
+]
+_INFER_CLS = Tuple[
+    ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments
+]
+_EVAL_ARGS = [
+    ModelArguments, DataArguments, EvaluationArguments, FinetuningArguments
+]
+_EVAL_CLS = Tuple[
+    ModelArguments, DataArguments, EvaluationArguments, FinetuningArguments
+]
 
 
-def parse_train_args(
-    args: Optional[Dict[str, Any]] = None
-) -> Tuple[
-    ModelArguments,
-    DataArguments,
-    Seq2SeqTrainingArguments,
-    FinetuningArguments,
-    GeneratingArguments
-]:
-    parser = HfArgumentParser((
-        ModelArguments,
-        DataArguments,
-        Seq2SeqTrainingArguments,
-        FinetuningArguments,
-        GeneratingArguments
-    ))
-    return _parse_args(parser, args)
+def _verify_model_args(model_args: "ModelArguments", finetuning_args: "FinetuningArguments") -> None:
+    if model_args.quantization_bit is not None and finetuning_args.finetuning_type != "lora":
+        raise ValueError("Quantization is only compatible with the LoRA method.")
+
+    if (
+        model_args.checkpoint_dir is not None
+        and len(model_args.checkpoint_dir) != 1
+        and finetuning_args.finetuning_type != "lora"
+    ):
+        raise ValueError("Multiple checkpoints are only available for LoRA tuning.")
 
 
-def parse_infer_args(
-    args: Optional[Dict[str, Any]] = None
-) -> Tuple[
-    ModelArguments,
-    DataArguments,
-    FinetuningArguments,
-    GeneratingArguments
-]:
-    parser = HfArgumentParser((
-        ModelArguments,
-        DataArguments,
-        FinetuningArguments,
-        GeneratingArguments
-    ))
-    return _parse_args(parser, args)
+def parse_train_args(args: Optional[Dict[str, Any]] = None) -> _TRAIN_CLS:
+    parser = HfArgumentParser(_TRAIN_ARGS)
+    return parse_args(parser, args)
 
 
-def get_train_args(
-    args: Optional[Dict[str, Any]] = None
-) -> Tuple[
-    ModelArguments,
-    DataArguments,
-    Seq2SeqTrainingArguments,
-    FinetuningArguments,
-    GeneratingArguments
-]:
+def parse_infer_args(args: Optional[Dict[str, Any]] = None) -> _INFER_CLS:
+    parser = HfArgumentParser(_INFER_ARGS)
+    return parse_args(parser, args)
+
+
+def parse_eval_args(args: Optional[Dict[str, Any]] = None) -> _EVAL_CLS:
+    parser = HfArgumentParser(_EVAL_ARGS)
+    return parse_args(parser, args)
+
+
+def get_train_args(args: Optional[Dict[str, Any]] = None) -> _TRAIN_CLS:
     model_args, data_args, training_args, finetuning_args, generating_args = parse_train_args(args)
 
     # Setup logging
@@ -100,24 +93,14 @@ def get_train_args(
     if finetuning_args.stage == "sft" and training_args.do_predict and not training_args.predict_with_generate:
         raise ValueError("Please enable `predict_with_generate` to save model predictions.")
 
-    if finetuning_args.stage in ["rm", "ppo"]:
-        if finetuning_args.finetuning_type != "lora":
-            raise ValueError("RM and PPO stages can only be performed with the LoRA method.")
-        if training_args.resume_from_checkpoint is not None:
-            raise ValueError("RM and PPO stages do not support `resume_from_checkpoint`.")
-        if training_args.load_best_model_at_end:
-            raise ValueError("RM and PPO stages do not support `load_best_model_at_end`.")
+    if finetuning_args.stage in ["rm", "ppo"] and training_args.load_best_model_at_end:
+        raise ValueError("RM and PPO stages do not support `load_best_model_at_end`.")
 
     if finetuning_args.stage == "ppo" and not training_args.do_train:
-        raise ValueError("PPO training does not support evaluation.")
+        raise ValueError("PPO training does not support evaluation, use the SFT stage to evaluate models.")
 
-    if finetuning_args.stage in ["rm", "dpo"]:
-        for dataset_attr in data_args.dataset_list:
-            if not dataset_attr.ranking:
-                raise ValueError("Please use ranked datasets for reward modeling or DPO training.")
-
-    if finetuning_args.stage == "ppo" and model_args.reward_model is None:
-        raise ValueError("Reward model is necessary for PPO training.")
+    if finetuning_args.stage in ["rm", "dpo"] and (not all([data_attr.ranking for data_attr in data_args.dataset_list])):
+        raise ValueError("Please use ranked datasets for reward modeling or DPO training.")
 
     if finetuning_args.stage == "ppo" and model_args.shift_attn:
         raise ValueError("PPO training is incompatible with S^2-Attn.")
@@ -131,15 +114,7 @@ def get_train_args(
     if training_args.do_train and finetuning_args.finetuning_type == "lora" and finetuning_args.lora_target is None:
         raise ValueError("Please specify `lora_target` in LoRA training.")
 
-    if model_args.quantization_bit is not None and finetuning_args.finetuning_type != "lora":
-        raise ValueError("Quantization is only compatible with the LoRA method.")
-
-    if (
-        model_args.checkpoint_dir is not None
-        and len(model_args.checkpoint_dir) != 1
-        and finetuning_args.finetuning_type != "lora"
-    ):
-        raise ValueError("Only LoRA tuning accepts multiple checkpoints.")
+    _verify_model_args(model_args, finetuning_args)
 
     if training_args.do_train and model_args.quantization_bit is not None and (not finetuning_args.upcast_layernorm):
         logger.warning("We recommend enable `upcast_layernorm` in quantized training.")
@@ -150,6 +125,9 @@ def get_train_args(
     if (not training_args.do_train) and model_args.quantization_bit is not None:
         logger.warning("Evaluating model in 4/8-bit mode may cause lower scores.")
 
+    if (not training_args.do_train) and finetuning_args.stage == "dpo" and finetuning_args.ref_model is None:
+        logger.warning("Specify `ref_model` for computing rewards at evaluation.")
+
     # postprocess training_args
     if (
         training_args.local_rank != -1
@@ -175,9 +153,14 @@ def get_train_args(
             training_args_dict = training_args.to_dict()
             training_args_dict.update(dict(resume_from_checkpoint=last_checkpoint))
             training_args = Seq2SeqTrainingArguments(**training_args_dict)
-            logger.info(
-                "Resuming from checkpoint. Change `output_dir` or use `overwrite_output_dir` to avoid."
-            )
+            logger.info("Resuming training from {}. Change `output_dir` or use `overwrite_output_dir` to avoid.".format(
+                training_args.resume_from_checkpoint
+            ))
+
+    if finetuning_args.stage in ["rm", "ppo"] and training_args.resume_from_checkpoint is not None:
+        logger.warning("Add {} to `checkpoint_dir` to resume training from checkpoint.".format(
+            training_args.resume_from_checkpoint
+        ))
 
     # postprocess model_args
     model_args.compute_dtype = (
@@ -198,27 +181,25 @@ def get_train_args(
     return model_args, data_args, training_args, finetuning_args, generating_args
 
 
-def get_infer_args(
-    args: Optional[Dict[str, Any]] = None
-) -> Tuple[
-    ModelArguments,
-    DataArguments,
-    FinetuningArguments,
-    GeneratingArguments
-]:
+def get_infer_args(args: Optional[Dict[str, Any]] = None) -> _INFER_CLS:
     model_args, data_args, finetuning_args, generating_args = parse_infer_args(args)
 
     if data_args.template is None:
         raise ValueError("Please specify which `template` to use.")
 
-    if model_args.quantization_bit is not None and finetuning_args.finetuning_type != "lora":
-        raise ValueError("Quantization is only compatible with the LoRA method.")
-
-    if (
-        model_args.checkpoint_dir is not None
-        and len(model_args.checkpoint_dir) != 1
-        and finetuning_args.finetuning_type != "lora"
-    ):
-        raise ValueError("Only LoRA tuning accepts multiple checkpoints.")
+    _verify_model_args(model_args, finetuning_args)
 
     return model_args, data_args, finetuning_args, generating_args
+
+
+def get_eval_args(args: Optional[Dict[str, Any]] = None) -> _EVAL_CLS:
+    model_args, data_args, eval_args, finetuning_args = parse_eval_args(args)
+
+    if data_args.template is None:
+        raise ValueError("Please specify which `template` to use.")
+
+    _verify_model_args(model_args, finetuning_args)
+
+    transformers.set_seed(eval_args.seed)
+
+    return model_args, data_args, eval_args, finetuning_args

src/llmtuner/model/utils.py

@@ -1,21 +1,54 @@
 import torch
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+import inspect
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple
+
+from transformers.utils import cached_file
+from transformers.trainer import WEIGHTS_NAME, SAFE_WEIGHTS_NAME
 
 from llmtuner.extras.constants import LAYERNORM_NAMES
 from llmtuner.extras.logging import get_logger
+from llmtuner.hparams import ModelArguments, FinetuningArguments
 
 if TYPE_CHECKING:
     from transformers.modeling_utils import PreTrainedModel
-    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
+    from llmtuner.hparams import DataArguments
 
 
 logger = get_logger(__name__)
 
 
+def dispatch_model(model: "PreTrainedModel") -> "PreTrainedModel":
+    r"""
+    Dispatches a pre-trained model to GPUs with balanced memory.
+    Borrowed from: https://github.com/huggingface/transformers/blob/v4.31.0/src/transformers/modeling_utils.py#L2803
+    """
+    if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False): # do nothing
+        return model
+
+    if torch.cuda.device_count() > 1:
+        from accelerate import dispatch_model
+        from accelerate.utils import infer_auto_device_map, get_balanced_memory
+
+        if model._no_split_modules is None:
+            raise ValueError("The model class needs to implement the `_no_split_modules` attribute.")
+
+        kwargs = {"dtype": model.dtype, "no_split_module_classes": model._no_split_modules}
+        max_memory = get_balanced_memory(model, **kwargs)
+        # Make sure tied weights are tied before creating the device map.
+        model.tie_weights()
+        device_map = infer_auto_device_map(model, max_memory=max_memory, **kwargs)
+        return dispatch_model(model, device_map)
+    else:
+        return model.cuda()
+
+
 def find_all_linear_modules(
     model: "PreTrainedModel",
     quantization_bit: Optional[int] = None
 ) -> List[str]:
+    r"""
+    Finds all available modules to apply lora.
+    """
     if quantization_bit is not None:
         import bitsandbytes as bnb
         linear_cls = bnb.nn.Linear4bit if quantization_bit == 4 else bnb.nn.Linear8bitLt
@@ -38,25 +71,68 @@ def find_all_linear_modules(
     return list(module_names)
 
 
-def generate_model_card(
+def get_modelcard_args(
     model_args: "ModelArguments",
     data_args: "DataArguments",
     finetuning_args: "FinetuningArguments"
 ) -> Dict[str, Any]:
     return {
         "tasks": "text-generation",
+        "license": "other",
         "finetuned_from": model_args.model_name_or_path,
         "dataset": [dataset.strip() for dataset in data_args.dataset.split(",")],
         "tags": ["llama-factory"] + (["lora"] if finetuning_args.finetuning_type == "lora" else [])
     }
 
 
+def load_valuehead_params(
+    path_or_repo_id: str,
+    model_args: "ModelArguments"
+) -> Dict[str, torch.Tensor]:
+    r"""
+    Loads value head parameters from Hugging Face Hub or local disk.
+
+    Returns: dict with keys `v_head.summary.weight` and `v_head.summary.bias`.
+    """
+    kwargs = {
+        "path_or_repo_id": path_or_repo_id,
+        "cache_dir": model_args.cache_dir
+    }
+
+    if "token" in inspect.signature(cached_file).parameters:
+        kwargs["token"] = model_args.hf_hub_token
+    elif "use_auth_token" in inspect.signature(cached_file).parameters: # for transformers==4.31.0
+        kwargs["use_auth_token"] = model_args.hf_hub_token
+    else:
+        logger.warning("Ignore `hf_hub_token` since matched parameter is not found.")
+
+    try:
+        vhead_file = cached_file(filename=WEIGHTS_NAME, **kwargs)
+        return torch.load(vhead_file, map_location="cpu")
+    except Exception as err:
+        logger.info("Failed to load {}: {}".format(WEIGHTS_NAME, str(err)))
+
+    try:
+        from safetensors import safe_open
+        vhead_file = cached_file(filename=SAFE_WEIGHTS_NAME, **kwargs)
+        with safe_open(vhead_file, framework="pt", device="cpu") as f:
+            return {
+                "v_head.summary.weight": f.get_tensor("v_head.summary.weight"),
+                "v_head.summary.bias": f.get_tensor("v_head.summary.bias")
+            }
+    except Exception as err:
+        logger.info("Failed to load {}: {}".format(SAFE_WEIGHTS_NAME, str(err)))
+
+    logger.warning("Provided path ({}) does not contain valuehead weights.".format(path_or_repo_id))
+    return None
+
+
 def prepare_model_for_training(
     model: "PreTrainedModel",
     finetuning_args: "FinetuningArguments",
     output_layer_name: Optional[str] = "lm_head",
     use_gradient_checkpointing: Optional[bool] = True,
-    layernorm_names: Optional[List[str]] = LAYERNORM_NAMES
+    layernorm_names: Optional[Set[str]] = LAYERNORM_NAMES
 ) -> "PreTrainedModel":
     r"""
     Includes:
@@ -82,7 +158,7 @@ def prepare_model_for_training(
         model.get_input_embeddings().register_forward_hook(neftune_forward_hook)
         logger.info("Using noisy embedding with alpha={:.2f}".format(finetuning_args.neft_alpha))
 
-    if use_gradient_checkpointing:
+    if use_gradient_checkpointing and getattr(model, "supports_gradient_checkpointing", False):
         if hasattr(model, "enable_input_require_grads"):
             model.enable_input_require_grads()
         else:

src/llmtuner/train/__init__.py

@@ -0,0 +1 @@
+from llmtuner.train.tuner import export_model, run_exp

src/llmtuner/train/dpo/__init__.py

@@ -0,0 +1 @@
+from llmtuner.train.dpo.workflow import run_dpo

src/llmtuner/train/dpo/trainer.py

@@ -1,6 +1,4 @@
 import torch
-import deepspeed # type: ignore
-from copy import deepcopy
 from collections import defaultdict
 from typing import TYPE_CHECKING, Dict, Literal, Optional, Tuple, Union
 from transformers import BatchEncoding, Trainer
@@ -11,7 +9,6 @@ from llmtuner.extras.constants import IGNORE_INDEX
 
 if TYPE_CHECKING:
     from transformers import PreTrainedModel
-    from trl import PreTrainedModelWrapper
 
 
 class CustomDPOTrainer(DPOTrainer):
@@ -46,40 +43,14 @@ class CustomDPOTrainer(DPOTrainer):
 
         if ref_model is not None:
             if self.is_deepspeed_enabled:
-                self.ref_model = self._prepare_deepspeed(self.ref_model)
+                if not (
+                    getattr(ref_model, "is_loaded_in_8bit", False)
+                    or getattr(ref_model, "is_loaded_in_4bit", False)
+                ): # quantized models are already set on the correct device
+                    self.ref_model = self._prepare_deepspeed(self.ref_model)
             else:
                 self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
 
-    def _prepare_deepspeed(self, model: "PreTrainedModelWrapper"):
-        # Adapted from accelerate: https://github.com/huggingface/accelerate/blob/739b135f8367becb67ffaada12fe76e3aa60fefd/src/accelerate/accelerator.py#L1473
-        deepspeed_plugin = self.accelerator.state.deepspeed_plugin
-        config_kwargs = deepcopy(deepspeed_plugin.deepspeed_config)
-        if model is not None:
-            if hasattr(model, "config"):
-                hidden_size = (
-                    max(model.config.hidden_sizes)
-                    if getattr(model.config, "hidden_sizes", None)
-                    else getattr(model.config, "hidden_size", None)
-                )
-                if hidden_size is not None and config_kwargs["zero_optimization"]["stage"] == 3:
-                    # Note that `stage3_prefetch_bucket_size` can produce DeepSpeed messages like: `Invalidate trace cache @ step 0: expected module 1, but got module 0`
-                    # This is expected and is not an error, see: https://github.com/microsoft/DeepSpeed/discussions/4081
-                    config_kwargs.update(
-                        {
-                            "zero_optimization.reduce_bucket_size": hidden_size * hidden_size,
-                            "zero_optimization.stage3_param_persistence_threshold": 10 * hidden_size,
-                            "zero_optimization.stage3_prefetch_bucket_size": 0.9 * hidden_size * hidden_size,
-                        }
-                    )
-
-        # If ZeRO-3 is used, we shard both the active and reference model.
-        # Otherwise, we assume the reference model fits in memory and is initialized on each device with ZeRO disabled (stage 0)
-        if config_kwargs["zero_optimization"]["stage"] != 3:
-            config_kwargs["zero_optimization"]["stage"] = 0
-        model, *_ = deepspeed.initialize(model=model, config=config_kwargs)
-        model.eval()
-        return model
-
     def concatenated_forward(
         self,
         model: Optional[torch.nn.Module] = None,

src/llmtuner/train/dpo/workflow.py

@@ -1,26 +1,22 @@
 # Inspired by: https://github.com/huggingface/trl/blob/main/examples/research_projects/stack_llama_2/scripts/dpo_llama2.py
 
-from peft import PeftModel
 from typing import TYPE_CHECKING, Optional, List
 from transformers import Seq2SeqTrainingArguments
 
-from llmtuner.dsets import get_dataset, preprocess_dataset, split_dataset
+from llmtuner.data import get_dataset, preprocess_dataset, split_dataset
 from llmtuner.extras.constants import IGNORE_INDEX
-from llmtuner.extras.logging import get_logger
 from llmtuner.extras.ploting import plot_loss
 from llmtuner.hparams import ModelArguments
-from llmtuner.tuner.core import generate_model_card, load_model_and_tokenizer
-from llmtuner.tuner.dpo.collator import DPODataCollatorWithPadding
-from llmtuner.tuner.dpo.trainer import CustomDPOTrainer
+from llmtuner.model import load_model_and_tokenizer
+from llmtuner.train.dpo.collator import DPODataCollatorWithPadding
+from llmtuner.train.dpo.trainer import CustomDPOTrainer
+from llmtuner.train.utils import create_modelcard_and_push, create_ref_model
 
 if TYPE_CHECKING:
     from transformers import TrainerCallback
     from llmtuner.hparams import DataArguments, FinetuningArguments
 
 
-logger = get_logger(__name__)
-
-
 def run_dpo(
     model_args: "ModelArguments",
     data_args: "DataArguments",
@@ -38,23 +34,10 @@ def run_dpo(
     )
 
     # Create reference model
-    if finetuning_args.dpo_ref_model is not None:
-        ref_model_args_dict = model_args.to_dict()
-        ref_model_args_dict.update(dict(
-            model_name_or_path=finetuning_args.dpo_ref_model,
-            checkpoint_dir=finetuning_args.dpo_ref_model_checkpoint
-        ))
-        ref_model_args = ModelArguments(**ref_model_args_dict)
-        ref_model, _ = load_model_and_tokenizer(ref_model_args, finetuning_args, is_trainable=False, stage="sft")
-        logger.info("Created reference model from {}".format(finetuning_args.dpo_ref_model))
-    elif training_args.do_train:
-        if isinstance(model, PeftModel):
-            ref_model = None
-        else:
-            ref_model, _ = load_model_and_tokenizer(model_args, finetuning_args, is_trainable=False, stage="sft")
-            logger.info("Created reference model from the model itself.")
-    else:
+    if finetuning_args.ref_model is None and (not training_args.do_train): # use the model itself
         ref_model = model
+    else:
+        ref_model = create_ref_model(model_args, finetuning_args, stage="dpo")
 
     # Update arguments
     training_args_dict = training_args.to_dict()
@@ -80,14 +63,13 @@ def run_dpo(
         trainer.log_metrics("train", train_result.metrics)
         trainer.save_metrics("train", train_result.metrics)
         trainer.save_state()
-        if trainer.is_world_process_zero() and model_args.plot_loss:
+        if trainer.is_world_process_zero() and finetuning_args.plot_loss:
             plot_loss(training_args.output_dir, keys=["loss", "eval_loss"])
 
     # Evaluation
     if training_args.do_eval:
         metrics = trainer.evaluate(metric_key_prefix="eval")
         if id(model) == id(ref_model): # unable to compute rewards without a reference model
-            logger.warning("Pass `dpo_ref_model` for computing rewards at evaluation.")
             remove_keys = [key for key in metrics.keys() if "rewards" in key]
             for key in remove_keys:
                 metrics.pop(key)
@@ -95,8 +77,4 @@ def run_dpo(
         trainer.save_metrics("eval", metrics)
 
     # Create model card
-    if training_args.do_train:
-        if training_args.push_to_hub:
-            trainer.push_to_hub(**generate_model_card(model_args, data_args, finetuning_args))
-        else:
-            trainer.create_model_card(**generate_model_card(model_args, data_args, finetuning_args))
+    create_modelcard_and_push(trainer, model_args, data_args, training_args, finetuning_args)

src/llmtuner/train/ppo/__init__.py

@@ -0,0 +1 @@
+from llmtuner.train.ppo.workflow import run_ppo

src/llmtuner/train/ppo/trainer.py

@@ -3,9 +3,9 @@ import sys
 import math
 import torch
 from tqdm import tqdm
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, List, Optional, Tuple
 
-from transformers import GenerationConfig, Trainer, TrainerState, TrainerControl
+from transformers import BatchEncoding, GenerationConfig, Trainer, TrainerState, TrainerControl
 from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
 
 from trl import PPOTrainer
@@ -14,7 +14,7 @@ from trl.core import PPODecorators, logprobs_from_logits
 from llmtuner.extras.callbacks import LogCallback, SavePeftModelCallback
 from llmtuner.extras.logging import get_logger
 from llmtuner.extras.misc import AverageMeter, count_parameters, get_logits_processor
-from llmtuner.tuner.ppo.utils import dump_layernorm, restore_layernorm, replace_model
+from llmtuner.train.ppo.utils import dump_layernorm, restore_layernorm, replace_model
 
 if TYPE_CHECKING:
     from transformers import Seq2SeqTrainingArguments, TrainerCallback
@@ -37,36 +37,61 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
         finetuning_args: "FinetuningArguments",
         generating_args: "GeneratingArguments",
         callbacks: List["TrainerCallback"],
+        reward_model: "AutoModelForCausalLMWithValueHead",
         **kwargs
     ):
         PPOTrainer.__init__(self, **kwargs)
+
         self.args = training_args
         self.model_args = model_args
         self.finetuning_args = finetuning_args
+        self.reward_model = reward_model
+
         self.generation_config = GenerationConfig(
             pad_token_id=self.tokenizer.pad_token_id,
             eos_token_id=[self.tokenizer.eos_token_id] + self.tokenizer.additional_special_tokens_ids,
             **generating_args.to_dict()
         )
+
         self.state = TrainerState()
         self.control = TrainerControl()
         self.log_callback, self.save_callback = callbacks[0], callbacks[1]
         assert isinstance(self.log_callback, LogCallback) and isinstance(self.save_callback, SavePeftModelCallback)
+
         if self.args.max_steps > 0:
             logger.info("max_steps is given, it will override any value given in num_train_epochs")
 
-    def ppo_train(self) -> None:
+        if reward_model is not None:
+            is_deepspeed_enabled = self.accelerator.distributed_type == "DEEPSPEED" and hasattr(
+                self.accelerator.state, "deepspeed_plugin"
+            )
+            if is_deepspeed_enabled:
+                if not (
+                    getattr(reward_model.pretrained_model, "is_loaded_in_8bit", False)
+                    or getattr(reward_model.pretrained_model, "is_loaded_in_4bit", False)
+                ): # quantized models are already set on the correct device
+                    self.reward_model = self._prepare_deepspeed(self.reward_model)
+            else:
+                self.reward_model = self.accelerator.prepare_model(self.reward_model, evaluation_mode=True)
+
+    def ppo_train(self, resume_from_checkpoint: Optional[str] = None) -> None:
         r"""
         Implements training loop for the PPO stage, like _inner_training_loop() in Huggingface's Trainer.
         """
+        if resume_from_checkpoint is not None:
+            raise ValueError("`resume_from_checkpoint` will be supported in the future version.")
+
         total_train_batch_size = (
-            self.args.per_device_train_batch_size * self.args.gradient_accumulation_steps * self.args.world_size
+            self.args.per_device_train_batch_size
+            * self.args.gradient_accumulation_steps
+            * self.finetuning_args.ppo_buffer_size
+            * self.args.world_size
         )
         if self.args.max_steps > 0:
             num_examples = total_train_batch_size * self.args.max_steps
             num_train_epochs = sys.maxsize
             max_steps = self.args.max_steps
-            steps_in_epoch = self.args.max_steps * self.args.gradient_accumulation_steps
+            steps_in_epoch = self.args.max_steps
         else:
             len_dataloader = len(self.dataloader)
             num_examples = len(self.dataset)
@@ -81,13 +106,16 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
 
         if self.is_world_process_zero():
             logger.info("***** Running training *****")
-            logger.info(f"  Num examples = {num_examples}")
-            logger.info(f"  Num Epochs = {num_train_epochs}")
-            logger.info(f"  Instantaneous batch size per device = {self.args.per_device_train_batch_size}")
-            logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size}")
-            logger.info(f"  Gradient Accumulation steps = {self.args.gradient_accumulation_steps}")
-            logger.info(f"  Total optimization steps = {max_steps}")
-            logger.info(f"  Number of trainable parameters = {count_parameters(self.model)[0]}")
+            logger.info("  Num examples = {}".format(num_examples))
+            logger.info("  Num Epochs = {}".format(num_train_epochs))
+            logger.info("  Instantaneous batch size per device = {}".format(self.args.per_device_train_batch_size))
+            logger.info("  Total train batch size (w. parallel, buffer, distributed & accumulation) = {}".format(
+                total_train_batch_size
+            ))
+            logger.info("  Gradient Accumulation steps = {}".format(self.args.gradient_accumulation_steps))
+            logger.info("  Num optimization epochs per batch = {}".format(self.finetuning_args.ppo_epochs))
+            logger.info("  Total training steps = {}".format(max_steps))
+            logger.info("  Number of trainable parameters = {}".format(count_parameters(self.model)[0]))
 
         unwrapped_model: "AutoModelForCausalLMWithValueHead" = self.accelerator.unwrap_model(self.model)
         dataiter = iter(self.dataloader)
@@ -108,9 +136,14 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
             self.model.eval()
 
             # Get inputs
-            queries, responses = self.get_inputs(batch)
             self.tokenizer.padding_side = "right" # change padding side
-            rewards = self.get_rewards(queries, responses, unwrapped_model)
+            queries, responses, rewards = [], [], []
+            for idx in range(0, self.config.batch_size, self.config.mini_batch_size):
+                mini_batch_queries, mini_batch_responses = self.get_inputs(batch[idx:idx+self.config.mini_batch_size])
+                mini_batch_rewards = self.get_rewards(mini_batch_queries, mini_batch_responses, unwrapped_model)
+                queries.extend(mini_batch_queries)
+                responses.extend(mini_batch_responses)
+                rewards.extend(mini_batch_rewards)
 
             # Cast to training mode
             unwrapped_model.gradient_checkpointing_enable()
@@ -165,7 +198,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
         )
 
     @torch.no_grad()
-    def get_inputs(self, batch: Dict[str, torch.Tensor]) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
+    def get_inputs(self, batch: BatchEncoding) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
         r"""
         Generates model's responses given queries.
         """
@@ -208,25 +241,30 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
         r"""
         Computes scores using given reward model.
         """
-        replace_model(unwrapped_model, target="reward")
+        if self.reward_model is None:
+            replace_model(unwrapped_model, target="reward")
+
         batch = self.prepare_model_inputs(queries, responses)
 
         with torch.cuda.amp.autocast(dtype=self.model_args.compute_dtype): # support bf16
-            _, _, values = self.model(**batch, output_hidden_states=True, return_dict=True)
+            reward_model = self.reward_model if self.reward_model is not None else self.model
+            _, _, values = reward_model(**batch, output_hidden_states=True, return_dict=True)
 
         if values.size(0) != batch["input_ids"].size(0): # adapt to chatglm2
             values = torch.transpose(values, 0, 1)
 
         rewards = []
         for i in range(values.size(0)):
-            end_indexes = (batch["input_ids"][i] != self.tokenizer.eos_token_id).nonzero()
+            end_indexes = (batch["input_ids"][i] != self.tokenizer.pad_token_id).nonzero()
             end_index = end_indexes[-1].item() if len(end_indexes) else 0
             rewards.append(values[i, end_index].float().detach().cpu()) # use fp32 type
 
-        replace_model(unwrapped_model, target="default")
+        if self.reward_model is None:
+            replace_model(unwrapped_model, target="default")
+
         return rewards
 
-    @PPODecorators.empty_cuda_cache()
+    @PPODecorators.empty_device_cache()
     def batched_forward_pass(
         self,
         model: "AutoModelForCausalLMWithValueHead",

src/llmtuner/train/ppo/workflow.py

@@ -7,11 +7,12 @@ from typing import TYPE_CHECKING, Optional, List
 from transformers import DataCollatorWithPadding
 from transformers.optimization import get_scheduler
 
-from llmtuner.dsets import get_dataset, preprocess_dataset
+from llmtuner.data import get_dataset, preprocess_dataset
 from llmtuner.extras.callbacks import SavePeftModelCallback
 from llmtuner.extras.ploting import plot_loss
-from llmtuner.tuner.core import load_model_and_tokenizer
-from llmtuner.tuner.ppo.trainer import CustomPPOTrainer
+from llmtuner.model import load_model_and_tokenizer
+from llmtuner.train.utils import create_ref_model, create_reward_model
+from llmtuner.train.ppo.trainer import CustomPPOTrainer
 
 if TYPE_CHECKING:
     from transformers import Seq2SeqTrainingArguments, TrainerCallback
@@ -33,30 +34,36 @@ def run_ppo(
     tokenizer.padding_side = "left" # use left-padding in generation while using right-padding in training
     data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
 
+    # Create reference model and reward model
+    ref_model = create_ref_model(model_args, finetuning_args, stage="ppo")
+    reward_model = create_reward_model(model, model_args, finetuning_args)
+
+    # Create ppo config
+    backward_batch_size = training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps
     ppo_config = PPOConfig(
         model_name=model_args.model_name_or_path,
         learning_rate=training_args.learning_rate,
         mini_batch_size=training_args.per_device_train_batch_size,
-        batch_size=training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps,
+        batch_size=backward_batch_size * finetuning_args.ppo_buffer_size,
         gradient_accumulation_steps=training_args.gradient_accumulation_steps,
-        ppo_epochs=1,
+        ppo_epochs=finetuning_args.ppo_epochs,
         max_grad_norm=training_args.max_grad_norm,
         seed=training_args.seed,
-        optimize_cuda_cache=True,
+        optimize_device_cache=True,
         target=finetuning_args.ppo_target,
         log_with=finetuning_args.ppo_logger,
         use_score_scaling=finetuning_args.ppo_score_norm,
         use_score_norm=finetuning_args.ppo_score_norm,
+        whiten_rewards=finetuning_args.ppo_whiten_rewards,
         accelerator_kwargs={"step_scheduler_with_optimizer": False}
     )
 
+    # Create optimizer and scheduler
     optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=training_args.learning_rate)
     if training_args.max_steps > 0:
         num_training_steps = training_args.max_steps
     else:
-        total_train_batch_size = (
-            training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps * training_args.world_size
-        )
+        total_train_batch_size = backward_batch_size * finetuning_args.ppo_buffer_size * training_args.world_size
         num_training_steps = training_args.num_train_epochs * math.ceil(len(dataset) / total_train_batch_size)
 
     lr_scheduler = get_scheduler(
@@ -73,9 +80,10 @@ def run_ppo(
         finetuning_args=finetuning_args,
         generating_args=generating_args,
         callbacks=callbacks + [SavePeftModelCallback()],
+        reward_model=reward_model,
         config=ppo_config,
         model=model,
-        ref_model=None,
+        ref_model=ref_model,
         tokenizer=tokenizer,
         dataset=dataset,
         data_collator=data_collator,
@@ -85,8 +93,8 @@ def run_ppo(
 
     # Training
     if training_args.do_train:
-        ppo_trainer.ppo_train()
+        ppo_trainer.ppo_train(resume_from_checkpoint=training_args.resume_from_checkpoint)
         ppo_trainer.save_model()
         ppo_trainer.save_state() # must be called after save_model to have a folder
-        if ppo_trainer.is_world_process_zero() and model_args.plot_loss:
+        if ppo_trainer.is_world_process_zero() and finetuning_args.plot_loss:
             plot_loss(training_args.output_dir, keys=["loss", "reward"])

src/llmtuner/train/pt/__init__.py

@@ -0,0 +1 @@
+from llmtuner.train.pt.workflow import run_pt

src/llmtuner/train/pt/workflow.py

@@ -4,9 +4,10 @@ import math
 from typing import TYPE_CHECKING, Optional, List
 from transformers import DataCollatorForLanguageModeling, Trainer
 
-from llmtuner.dsets import get_dataset, preprocess_dataset, split_dataset
+from llmtuner.data import get_dataset, preprocess_dataset, split_dataset
 from llmtuner.extras.ploting import plot_loss
-from llmtuner.tuner.core import generate_model_card, load_model_and_tokenizer
+from llmtuner.model import load_model_and_tokenizer
+from llmtuner.train.utils import create_modelcard_and_push
 
 if TYPE_CHECKING:
     from transformers import Seq2SeqTrainingArguments, TrainerCallback
@@ -42,7 +43,7 @@ def run_pt(
         trainer.log_metrics("train", train_result.metrics)
         trainer.save_metrics("train", train_result.metrics)
         trainer.save_state()
-        if trainer.is_world_process_zero() and model_args.plot_loss:
+        if trainer.is_world_process_zero() and finetuning_args.plot_loss:
             plot_loss(training_args.output_dir, keys=["loss", "eval_loss"])
 
     # Evaluation
@@ -58,8 +59,4 @@ def run_pt(
         trainer.save_metrics("eval", metrics)
 
     # Create model card
-    if training_args.do_train:
-        if training_args.push_to_hub:
-            trainer.push_to_hub(**generate_model_card(model_args, data_args, finetuning_args))
-        else:
-            trainer.create_model_card(**generate_model_card(model_args, data_args, finetuning_args))
+    create_modelcard_and_push(trainer, model_args, data_args, training_args, finetuning_args)

src/llmtuner/train/rm/__init__.py

@@ -0,0 +1 @@
+from llmtuner.train.rm.workflow import run_rm

src/llmtuner/train/rm/workflow.py

@@ -3,13 +3,14 @@
 from typing import TYPE_CHECKING, Optional, List
 from transformers import Seq2SeqTrainingArguments
 
-from llmtuner.dsets import get_dataset, preprocess_dataset, split_dataset
+from llmtuner.data import get_dataset, preprocess_dataset, split_dataset
 from llmtuner.extras.callbacks import SavePeftModelCallback
 from llmtuner.extras.ploting import plot_loss
-from llmtuner.tuner.core import generate_model_card, load_model_and_tokenizer
-from llmtuner.tuner.rm.metric import compute_accuracy
-from llmtuner.tuner.rm.collator import PairwiseDataCollatorWithPadding
-from llmtuner.tuner.rm.trainer import PairwiseTrainer
+from llmtuner.model import load_model_and_tokenizer
+from llmtuner.train.rm.collator import PairwiseDataCollatorWithPadding
+from llmtuner.train.rm.metric import compute_accuracy
+from llmtuner.train.rm.trainer import PairwiseTrainer
+from llmtuner.train.utils import create_modelcard_and_push
 
 if TYPE_CHECKING:
     from transformers import TrainerCallback
@@ -46,12 +47,12 @@ def run_rm(
 
     # Training
     if training_args.do_train:
-        train_result = trainer.train()
+        train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
         trainer.save_model()
         trainer.log_metrics("train", train_result.metrics)
         trainer.save_metrics("train", train_result.metrics)
         trainer.save_state()
-        if trainer.is_world_process_zero() and model_args.plot_loss:
+        if trainer.is_world_process_zero() and finetuning_args.plot_loss:
             plot_loss(training_args.output_dir, keys=["loss", "eval_loss"])
 
     # Evaluation
@@ -68,8 +69,4 @@ def run_rm(
         trainer.save_predictions(predict_results)
 
     # Create model card
-    if training_args.do_train:
-        if training_args.push_to_hub:
-            trainer.push_to_hub(**generate_model_card(model_args, data_args, finetuning_args))
-        else:
-            trainer.create_model_card(**generate_model_card(model_args, data_args, finetuning_args))
+    create_modelcard_and_push(trainer, model_args, data_args, training_args, finetuning_args)

src/llmtuner/train/sft/__init__.py

@@ -0,0 +1 @@
+from llmtuner.train.sft.workflow import run_sft

src/llmtuner/train/sft/metric.py

@@ -2,15 +2,23 @@ import numpy as np
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Dict, Sequence, Tuple, Union
 
-import jieba
-from rouge_chinese import Rouge
-from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
-
 from llmtuner.extras.constants import IGNORE_INDEX
+from llmtuner.extras.packages import (
+    is_jieba_available, is_nltk_available, is_rouge_available
+)
 
 if TYPE_CHECKING:
     from transformers.tokenization_utils import PreTrainedTokenizer
 
+if is_jieba_available():
+    import jieba
+
+if is_nltk_available():
+    from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
+
+if is_rouge_available():
+    from rouge_chinese import Rouge
+
 
 @dataclass
 class ComputeMetrics:

src/llmtuner/train/sft/trainer.py

@@ -39,10 +39,10 @@ class CustomSeq2SeqTrainer(Seq2SeqTrainer):
             prompt_len, label_len = inputs["input_ids"].size(-1), inputs["labels"].size(-1)
             if prompt_len > label_len:
                 inputs["labels"] = self._pad_tensors_to_target_len(inputs["labels"], inputs["input_ids"])
-            if label_len > prompt_len:
-                inputs["labels"] = inputs["labels"][:, :prompt_len] # truncate the labels instead of padding the inputs
+            if label_len > prompt_len: # truncate the labels instead of padding the inputs (llama2 fp16 compatibility)
+                inputs["labels"] = inputs["labels"][:, :prompt_len]
 
-        loss, generated_tokens, _ = super().prediction_step(
+        loss, generated_tokens, _ = super().prediction_step( # ignore the returned labels (may be truncated)
             model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys
         )
         if generated_tokens is not None and self.args.predict_with_generate:
@@ -79,14 +79,19 @@ class CustomSeq2SeqTrainer(Seq2SeqTrainer):
         output_prediction_file = os.path.join(self.args.output_dir, "generated_predictions.jsonl")
         logger.info(f"Saving prediction results to {output_prediction_file}")
 
-        preds = np.where(predict_results.predictions != IGNORE_INDEX, predict_results.predictions, self.tokenizer.pad_token_id)
         labels = np.where(predict_results.label_ids != IGNORE_INDEX, predict_results.label_ids, self.tokenizer.pad_token_id)
+        preds = np.where(predict_results.predictions != IGNORE_INDEX, predict_results.predictions, self.tokenizer.pad_token_id)
 
+        for i in range(len(preds)):
+            pad_len = np.nonzero(preds[i] != self.tokenizer.pad_token_id)[0]
+            if len(pad_len):
+                preds[i] = np.concatenate((preds[i][pad_len[0]:], preds[i][:pad_len[0]]), axis=-1) # move pad token to last
+
+        decoded_labels = self.tokenizer.batch_decode(labels, skip_special_tokens=True, clean_up_tokenization_spaces=False)
         decoded_preds = self.tokenizer.batch_decode(preds, skip_special_tokens=True, clean_up_tokenization_spaces=True)
-        decoded_labels = self.tokenizer.batch_decode(labels, skip_special_tokens=True, clean_up_tokenization_spaces=True)
 
         with open(output_prediction_file, "w", encoding="utf-8") as writer:
             res: List[str] = []
-            for pred, label in zip(decoded_preds, decoded_labels):
+            for label, pred in zip(decoded_labels, decoded_preds):
                 res.append(json.dumps({"label": label, "predict": pred}, ensure_ascii=False))
             writer.write("\n".join(res))

src/llmtuner/train/sft/workflow.py

@@ -3,13 +3,14 @@
 from typing import TYPE_CHECKING, Optional, List
 from transformers import DataCollatorForSeq2Seq, Seq2SeqTrainingArguments
 
-from llmtuner.dsets import get_dataset, preprocess_dataset, split_dataset
+from llmtuner.data import get_dataset, preprocess_dataset, split_dataset
 from llmtuner.extras.constants import IGNORE_INDEX
 from llmtuner.extras.misc import get_logits_processor
 from llmtuner.extras.ploting import plot_loss
-from llmtuner.tuner.core import generate_model_card, load_model_and_tokenizer
-from llmtuner.tuner.sft.metric import ComputeMetrics
-from llmtuner.tuner.sft.trainer import CustomSeq2SeqTrainer
+from llmtuner.model import load_model_and_tokenizer
+from llmtuner.train.sft.metric import ComputeMetrics
+from llmtuner.train.sft.trainer import CustomSeq2SeqTrainer
+from llmtuner.train.utils import create_modelcard_and_push
 
 if TYPE_CHECKING:
     from transformers import TrainerCallback
@@ -69,7 +70,7 @@ def run_sft(
         trainer.log_metrics("train", train_result.metrics)
         trainer.save_metrics("train", train_result.metrics)
         trainer.save_state()
-        if trainer.is_world_process_zero() and model_args.plot_loss:
+        if trainer.is_world_process_zero() and finetuning_args.plot_loss:
             plot_loss(training_args.output_dir, keys=["loss", "eval_loss"])
 
     # Evaluation
@@ -90,8 +91,4 @@ def run_sft(
         trainer.save_predictions(predict_results)
 
     # Create model card
-    if training_args.do_train:
-        if training_args.push_to_hub:
-            trainer.push_to_hub(**generate_model_card(model_args, data_args, finetuning_args))
-        else:
-            trainer.create_model_card(**generate_model_card(model_args, data_args, finetuning_args))
+    create_modelcard_and_push(trainer, model_args, data_args, training_args, finetuning_args)

src/llmtuner/train/tuner.py

@@ -2,12 +2,12 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 from llmtuner.extras.callbacks import LogCallback
 from llmtuner.extras.logging import get_logger
-from llmtuner.tuner.core import get_train_args, get_infer_args, load_model_and_tokenizer
-from llmtuner.tuner.pt import run_pt
-from llmtuner.tuner.sft import run_sft
-from llmtuner.tuner.rm import run_rm
-from llmtuner.tuner.ppo import run_ppo
-from llmtuner.tuner.dpo import run_dpo
+from llmtuner.model import get_train_args, get_infer_args, load_model_and_tokenizer
+from llmtuner.train.pt import run_pt
+from llmtuner.train.sft import run_sft
+from llmtuner.train.rm import run_rm
+from llmtuner.train.ppo import run_ppo
+from llmtuner.train.dpo import run_dpo
 
 if TYPE_CHECKING:
     from transformers import TrainerCallback
@@ -37,12 +37,17 @@ def run_exp(args: Optional[Dict[str, Any]] = None, callbacks: Optional[List["Tra
 def export_model(args: Optional[Dict[str, Any]] = None, max_shard_size: Optional[str] = "10GB"):
     model_args, _, finetuning_args, _ = get_infer_args(args)
     model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args)
+
+    if getattr(model, "quantization_method", None) == "gptq":
+        raise ValueError("Cannot export a GPTQ quantized model.")
+
     model.config.use_cache = True
-    model.save_pretrained(model_args.export_dir, max_shard_size=max_shard_size)
+    model.save_pretrained(finetuning_args.export_dir, max_shard_size=max_shard_size)
+
     try:
         tokenizer.padding_side = "left" # restore padding side
         tokenizer.init_kwargs["padding_side"] = "left"
-        tokenizer.save_pretrained(model_args.export_dir)
+        tokenizer.save_pretrained(finetuning_args.export_dir)
     except:
         logger.warning("Cannot save tokenizer, please copy the files manually.")
 

src/llmtuner/train/utils.py

@@ -0,0 +1,99 @@
+import torch
+from typing import TYPE_CHECKING, Literal, Union
+
+from llmtuner.extras.logging import get_logger
+from llmtuner.hparams import ModelArguments, FinetuningArguments
+from llmtuner.model import get_modelcard_args, load_model_and_tokenizer, load_valuehead_params
+
+if TYPE_CHECKING:
+    from transformers import Seq2SeqTrainingArguments, Trainer
+    from transformers.modeling_utils import PreTrainedModel
+    from trl import AutoModelForCausalLMWithValueHead
+    from llmtuner.hparams import DataArguments
+
+
+logger = get_logger(__name__)
+
+
+def create_modelcard_and_push(
+    trainer: "Trainer",
+    model_args: "ModelArguments",
+    data_args: "DataArguments",
+    training_args: "Seq2SeqTrainingArguments",
+    finetuning_args: "FinetuningArguments"
+) -> None:
+    if training_args.do_train:
+        if training_args.push_to_hub:
+            trainer.push_to_hub(**get_modelcard_args(model_args, data_args, finetuning_args))
+            return
+        try:
+            trainer.create_model_card(**get_modelcard_args(model_args, data_args, finetuning_args))
+        except Exception as err:
+            logger.warning("Failed to create model card: {}".format(str(err)))
+
+
+def create_ref_model(
+    model_args: "ModelArguments",
+    finetuning_args: "FinetuningArguments",
+    stage: Literal["ppo", "dpo"]
+) -> Union["PreTrainedModel", "AutoModelForCausalLMWithValueHead"]:
+    r"""
+    Creates reference model for PPO/DPO training. Evaluation mode is not supported.
+
+    The valuehead parameter is randomly initialized since it is useless for PPO training.
+    """
+    if finetuning_args.ref_model is not None:
+        ref_model_args_dict = model_args.to_dict()
+        ref_model_args_dict.update(dict(
+            model_name_or_path=finetuning_args.ref_model,
+            checkpoint_dir=finetuning_args.ref_model_checkpoint,
+            quantization_bit=finetuning_args.ref_model_quantization_bit
+        ))
+        ref_model_args = ModelArguments(**ref_model_args_dict)
+        ref_finetuning_args = FinetuningArguments(finetuning_type="lora")
+        ref_model, _ = load_model_and_tokenizer(ref_model_args, ref_finetuning_args, is_trainable=False, stage=stage)
+        logger.info("Created reference model from {}".format(finetuning_args.ref_model))
+    else:
+        if finetuning_args.finetuning_type == "lora":
+            ref_model = None
+        else:
+            ref_model, _ = load_model_and_tokenizer(model_args, finetuning_args, is_trainable=False, stage=stage)
+            logger.info("Created reference model from the model itself.")
+
+    return ref_model
+
+
+def create_reward_model(
+    model: "AutoModelForCausalLMWithValueHead",
+    model_args: "ModelArguments",
+    finetuning_args: "FinetuningArguments"
+) -> "AutoModelForCausalLMWithValueHead":
+    r"""
+    Creates reward model for PPO training.
+    """
+    if finetuning_args.reward_model_type == "lora":
+        model.pretrained_model.load_adapter(finetuning_args.reward_model, "reward")
+        for name, param in model.named_parameters(): # https://github.com/huggingface/peft/issues/1090
+            if "default" in name:
+                param.data = param.data.to(torch.float32) # trainable params should in fp32
+        vhead_params = load_valuehead_params(finetuning_args.reward_model, model_args)
+        assert vhead_params is not None, "Reward model is not correctly loaded."
+        model.register_buffer("reward_head_weight", vhead_params["v_head.summary.weight"], persistent=False)
+        model.register_buffer("reward_head_bias", vhead_params["v_head.summary.bias"], persistent=False)
+        model.register_buffer("default_head_weight", torch.zeros_like(vhead_params["v_head.summary.weight"]), persistent=False)
+        model.register_buffer("default_head_bias", torch.zeros_like(vhead_params["v_head.summary.bias"]), persistent=False)
+        logger.info("Loaded adapter weights of reward model from {}".format(finetuning_args.reward_model))
+        return None
+    else:
+        reward_model_args_dict = model_args.to_dict()
+        reward_model_args_dict.update(dict(
+            model_name_or_path=finetuning_args.reward_model,
+            checkpoint_dir=finetuning_args.reward_model_checkpoint,
+            quantization_bit=finetuning_args.reward_model_quantization_bit
+        ))
+        reward_model_args = ModelArguments(**reward_model_args_dict)
+        reward_finetuning_args = FinetuningArguments(finetuning_type="lora")
+        reward_model, _ = load_model_and_tokenizer(reward_model_args, reward_finetuning_args, is_trainable=False, stage="ppo")
+        logger.info("Load full weights of reward model from {}".format(finetuning_args.reward_model))
+        logger.warning("Please ensure the ppo model and reward model share SAME tokenizer and vocabulary.")
+        return reward_model

src/llmtuner/tuner/__init__.py

@@ -1 +0,0 @@
-from llmtuner.tuner.tune import export_model, run_exp

src/llmtuner/tuner/core/__init__.py

@@ -1,3 +0,0 @@
-from llmtuner.tuner.core.parser import get_train_args, get_infer_args
-from llmtuner.tuner.core.loader import load_model_and_tokenizer
-from llmtuner.tuner.core.utils import generate_model_card

src/llmtuner/tuner/dpo/__init__.py

@@ -1 +0,0 @@
-from llmtuner.tuner.dpo.workflow import run_dpo

src/llmtuner/tuner/ppo/__init__.py

@@ -1 +0,0 @@
-from llmtuner.tuner.ppo.workflow import run_ppo

src/llmtuner/tuner/pt/__init__.py

@@ -1 +0,0 @@
-from llmtuner.tuner.pt.workflow import run_pt

src/llmtuner/tuner/rm/__init__.py

@@ -1 +0,0 @@
-from llmtuner.tuner.rm.workflow import run_rm

src/llmtuner/tuner/sft/__init__.py

@@ -1 +0,0 @@
-from llmtuner.tuner.sft.workflow import run_sft

src/llmtuner/webui/chatter.py

@@ -2,7 +2,7 @@ import gradio as gr
 from gradio.components import Component # cannot use TYPE_CHECKING here
 from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional, Tuple
 
-from llmtuner.chat.stream_chat import ChatModel
+from llmtuner.chat import ChatModel
 from llmtuner.extras.misc import torch_gc
 from llmtuner.hparams import GeneratingArguments
 from llmtuner.webui.common import get_save_dir
@@ -14,14 +14,33 @@ if TYPE_CHECKING:
 
 class WebChatModel(ChatModel):
 
-    def __init__(self, manager: "Manager", lazy_init: Optional[bool] = True) -> None:
+    def __init__(
+        self,
+        manager: "Manager",
+        demo_mode: Optional[bool] = False,
+        lazy_init: Optional[bool] = True
+    ) -> None:
         self.manager = manager
+        self.demo_mode = demo_mode
         self.model = None
         self.tokenizer = None
         self.generating_args = GeneratingArguments()
-        if not lazy_init:
+
+        if not lazy_init: # read arguments from command line
             super().__init__()
 
+        if demo_mode: # load demo_config.json if exists
+            import json
+            try:
+                with open("demo_config.json", "r", encoding="utf-8") as f:
+                    args = json.load(f)
+                assert args.get("model_name_or_path", None) and args.get("template", None)
+                super().__init__(args)
+            except AssertionError:
+                print("Please provided model name and template in `demo_config.json`.")
+            except:
+                print("Cannot find `demo_config.json` at current directory.")
+
     @property
     def loaded(self) -> bool:
         return self.model is not None
@@ -36,6 +55,8 @@ class WebChatModel(ChatModel):
             error = ALERTS["err_no_model"][lang]
         elif not get("top.model_path"):
             error = ALERTS["err_no_path"][lang]
+        elif self.demo_mode:
+            error = ALERTS["err_demo"][lang]
 
         if error:
             gr.Warning(error)
@@ -67,6 +88,11 @@ class WebChatModel(ChatModel):
 
     def unload_model(self, data: Dict[Component, Any]) -> Generator[str, None, None]:
         lang = data[self.manager.get_elem_by_name("top.lang")]
+
+        if self.demo_mode:
+            yield ALERTS["err_demo"][lang]
+            return
+
         yield ALERTS["info_unloading"][lang]
         self.model = None
         self.tokenizer = None

src/llmtuner/webui/common.py

@@ -61,13 +61,17 @@ def get_model_path(model_name: str) -> str:
     return user_config["path_dict"].get(model_name, None) or SUPPORTED_MODELS.get(model_name, "")
 
 
+def get_prefix(model_name: str) -> str:
+    return model_name.split("-")[0]
+
+
 def get_module(model_name: str) -> str:
-    return DEFAULT_MODULE.get(model_name.split("-")[0], "q_proj,v_proj")
+    return DEFAULT_MODULE.get(get_prefix(model_name), "q_proj,v_proj")
 
 
 def get_template(model_name: str) -> str:
-    if model_name.endswith("Chat") and model_name.split("-")[0] in DEFAULT_TEMPLATE:
-        return DEFAULT_TEMPLATE[model_name.split("-")[0]]
+    if model_name and model_name.endswith("Chat") and get_prefix(model_name) in DEFAULT_TEMPLATE:
+        return DEFAULT_TEMPLATE[get_prefix(model_name)]
     return "default"
 
 

src/llmtuner/webui/components/export.py

@@ -1,7 +1,7 @@
 import gradio as gr
 from typing import TYPE_CHECKING, Dict, Generator, List
 
-from llmtuner.tuner import export_model
+from llmtuner.train import export_model
 from llmtuner.webui.common import get_save_dir
 from llmtuner.webui.locales import ALERTS
 

src/llmtuner/webui/components/top.py

@@ -1,8 +1,8 @@
 import gradio as gr
 from typing import TYPE_CHECKING, Dict
 
+from llmtuner.data.template import templates
 from llmtuner.extras.constants import METHODS, SUPPORTED_MODELS
-from llmtuner.extras.template import templates
 from llmtuner.webui.common import get_model_path, get_template, list_checkpoint, save_config
 from llmtuner.webui.utils import can_quantize
 

src/llmtuner/webui/css.py

@@ -1,4 +1,11 @@
 CSS = r"""
+.duplicate-button {
+  margin: auto !important;
+  color: white !important;
+  background: black !important;
+  border-radius: 100vh !important;
+}
+
 .modal-box {
   position: fixed !important;
   top: 50%;

src/llmtuner/webui/engine.py

@@ -12,17 +12,18 @@ from llmtuner.webui.utils import get_time
 
 class Engine:
 
-    def __init__(self, pure_chat: Optional[bool] = False) -> None:
+    def __init__(self, demo_mode: Optional[bool] = False, pure_chat: Optional[bool] = False) -> None:
+        self.demo_mode = demo_mode
         self.pure_chat = pure_chat
-        self.manager: "Manager" = Manager()
-        self.runner: "Runner" = Runner(self.manager)
-        self.chatter: "WebChatModel" = WebChatModel(manager=self.manager, lazy_init=(not pure_chat))
+        self.manager = Manager()
+        self.runner = Runner(self.manager, demo_mode=demo_mode)
+        self.chatter = WebChatModel(manager=self.manager, demo_mode=demo_mode, lazy_init=(not pure_chat))
 
     def _form_dict(self, resume_dict: Dict[str, Dict[str, Any]]):
         return {self.manager.get_elem_by_name(k): gr.update(**v) for k, v in resume_dict.items()}
 
     def resume(self) -> Generator[Dict[Component, Dict[str, Any]], None, None]:
-        user_config = load_config()
+        user_config = load_config() if not self.demo_mode else {}
         lang = user_config.get("lang", None) or "en"
 
         init_dict = {

src/llmtuner/webui/interface.py

@@ -1,4 +1,5 @@
 import gradio as gr
+from typing import Optional
 from transformers.utils.versions import require_version
 
 from llmtuner.webui.components import (
@@ -17,24 +18,35 @@ from llmtuner.webui.engine import Engine
 require_version("gradio>=3.38.0,<4.0.0", "To fix: pip install \"gradio>=3.38.0,<4.0.0\"")
 
 
-def create_ui() -> gr.Blocks:
-    engine = Engine(pure_chat=False)
+def create_ui(demo_mode: Optional[bool] = False) -> gr.Blocks:
+    engine = Engine(demo_mode=demo_mode, pure_chat=False)
 
     with gr.Blocks(title="LLaMA Board", css=CSS) as demo:
+        if demo_mode:
+            gr.HTML(
+                "<h1><center>LLaMA Board: A One-stop Web UI for Getting Started with LLaMA Factory</center></h1>"
+            )
+            gr.HTML(
+                "<h3><center>Visit <a href=\"https://github.com/hiyouga/LLaMA-Factory\" target=\"_blank\">"
+                "LLaMA Factory</a> for details.</center></h3>"
+            )
+            gr.DuplicateButton(value="Duplicate Space for private use", elem_classes="duplicate-button")
+
         engine.manager.all_elems["top"] = create_top()
         lang: "gr.Dropdown" = engine.manager.get_elem_by_name("top.lang")
 
         with gr.Tab("Train"):
             engine.manager.all_elems["train"] = create_train_tab(engine)
 
-        with gr.Tab("Evaluate"):
+        with gr.Tab("Evaluate & Predict"):
             engine.manager.all_elems["eval"] = create_eval_tab(engine)
 
         with gr.Tab("Chat"):
             engine.manager.all_elems["infer"] = create_infer_tab(engine)
 
-        with gr.Tab("Export"):
-            engine.manager.all_elems["export"] = create_export_tab(engine)
+        if not demo_mode:
+            with gr.Tab("Export"):
+                engine.manager.all_elems["export"] = create_export_tab(engine)
 
         demo.load(engine.resume, outputs=engine.manager.list_elems())
         lang.change(engine.change_lang, [lang], engine.manager.list_elems(), queue=False)

src/llmtuner/webui/locales.py

@@ -659,6 +659,10 @@ ALERTS = {
         "en": "Failed.",
         "zh": "训练出错。"
     },
+    "err_demo": {
+        "en": "Training is unavailable in demo mode, duplicate the space to a private one first.",
+        "zh": "展示模式不支持训练，请先复制到私人空间。"
+    },
     "info_aborting": {
         "en": "Aborted, wait for terminating...",
         "zh": "训练中断，正在等待线程结束……"

src/llmtuner/webui/runner.py

@@ -4,7 +4,7 @@ import logging
 import gradio as gr
 from threading import Thread
 from gradio.components import Component # cannot use TYPE_CHECKING here
-from typing import TYPE_CHECKING, Any, Dict, Generator, List, Tuple
+from typing import TYPE_CHECKING, Any, Dict, Generator, Optional, Tuple
 
 import transformers
 from transformers.trainer import TRAINING_ARGS_NAME
@@ -13,7 +13,7 @@ from llmtuner.extras.callbacks import LogCallback
 from llmtuner.extras.constants import TRAINING_STAGES
 from llmtuner.extras.logging import LoggerHandler
 from llmtuner.extras.misc import torch_gc
-from llmtuner.tuner import run_exp
+from llmtuner.train import run_exp
 from llmtuner.webui.common import get_module, get_save_dir, load_config
 from llmtuner.webui.locales import ALERTS
 from llmtuner.webui.utils import gen_cmd, get_eval_results, update_process_bar
@@ -24,13 +24,13 @@ if TYPE_CHECKING:
 
 class Runner:
 
-    def __init__(self, manager: "Manager") -> None:
+    def __init__(self, manager: "Manager", demo_mode: Optional[bool] = False) -> None:
         self.manager = manager
+        self.demo_mode = demo_mode
         """ Resume """
         self.thread: "Thread" = None
         self.do_train = True
         self.running_data: Dict["Component", Any] = None
-        self.monitor_inputs: Dict[str, str] = None
         """ State """
         self.aborted = False
         self.running = False
@@ -46,9 +46,8 @@ class Runner:
 
     def set_abort(self) -> None:
         self.aborted = True
-        self.running = False
 
-    def _initialize(self, data: Dict[Component, Any], do_train: bool) -> str:
+    def _initialize(self, data: Dict[Component, Any], do_train: bool, from_preview: bool) -> str:
         get = lambda name: data[self.manager.get_elem_by_name(name)]
         lang, model_name, model_path = get("top.lang"), get("top.model_name"), get("top.model_path")
         dataset = get("train.dataset") if do_train else get("eval.dataset")
@@ -65,6 +64,9 @@ class Runner:
         if len(dataset) == 0:
             return ALERTS["err_no_dataset"][lang]
 
+        if self.demo_mode and (not from_preview):
+            return ALERTS["err_demo"][lang]
+
         self.aborted = False
         self.logger_handler.reset()
         self.trainer_callback = LogCallback(self)
@@ -72,6 +74,7 @@ class Runner:
 
     def _finalize(self, lang: str, finish_info: str) -> str:
         self.thread = None
+        self.running_data = None
         self.running = False
         torch_gc()
         if self.aborted:
@@ -84,9 +87,9 @@ class Runner:
         user_config = load_config()
 
         if get("top.checkpoints"):
-            checkpoint_dir = ",".join([
-                get_save_dir(get("top.model_name"), get("top.finetuning_type"), ckpt) for ckpt in get("top.checkpoints")
-            ])
+            checkpoint_dir = ",".join([get_save_dir(
+                get("top.model_name"), get("top.finetuning_type"), ckpt
+            ) for ckpt in get("top.checkpoints")])
         else:
             checkpoint_dir = None
 
@@ -136,7 +139,10 @@ class Runner:
             args["upcast_layernorm"] = True
 
         if args["stage"] == "ppo":
-            args["reward_model"] = get("train.reward_model")
+            args["reward_model"] = get_save_dir(
+                get("top.model_name"), get("top.finetuning_type"), get("train.reward_model")
+            )
+            args["reward_model_type"] = "lora" if get("top.finetuning_type") == "lora" else "full"
 
         if args["stage"] == "dpo":
             args["dpo_beta"] = get("train.dpo_beta")
@@ -154,9 +160,9 @@ class Runner:
         user_config = load_config()
 
         if get("top.checkpoints"):
-            checkpoint_dir = ",".join([
-                get_save_dir(get("top.model_name"), get("top.finetuning_type"), ckpt) for ckpt in get("top.checkpoints")
-            ])
+            checkpoint_dir = ",".join([get_save_dir(
+                get("top.model_name"), get("top.finetuning_type"), ckpt
+            ) for ckpt in get("top.checkpoints")])
             output_dir = get_save_dir(
                 get("top.model_name"), get("top.finetuning_type"), "eval_" + "_".join(get("top.checkpoints"))
             )
@@ -196,7 +202,7 @@ class Runner:
         return args
 
     def _preview(self, data: Dict[Component, Any], do_train: bool) -> Generator[Tuple[str, Dict[str, Any]], None, None]:
-        error = self._initialize(data, do_train)
+        error = self._initialize(data, do_train, from_preview=True)
         if error:
             gr.Warning(error)
             yield error, gr.update(visible=False)
@@ -205,16 +211,14 @@ class Runner:
             yield gen_cmd(args), gr.update(visible=False)
 
     def _launch(self, data: Dict[Component, Any], do_train: bool) -> Generator[Tuple[str, Dict[str, Any]], None, None]:
-        error = self._initialize(data, do_train)
+        error = self._initialize(data, do_train, from_preview=False)
         if error:
             gr.Warning(error)
             yield error, gr.update(visible=False)
         else:
             args = self._parse_train_args(data) if do_train else self._parse_eval_args(data)
             run_kwargs = dict(args=args, callbacks=[self.trainer_callback])
-            self.running = True
             self.do_train, self.running_data = do_train, data
-            self.monitor_inputs = dict(lang=data[self.manager.get_elem_by_name("top.lang")], output_dir=args["output_dir"])
             self.thread = Thread(target=run_exp, kwargs=run_kwargs)
             self.thread.start()
             yield from self.monitor()
@@ -232,7 +236,12 @@ class Runner:
         yield from self._launch(data, do_train=False)
 
     def monitor(self) -> Generator[Tuple[str, Dict[str, Any]], None, None]:
-        lang, output_dir = self.monitor_inputs["lang"], self.monitor_inputs["output_dir"]
+        get = lambda name: self.running_data[self.manager.get_elem_by_name(name)]
+        self.running = True
+        lang = get("top.lang")
+        output_dir = get_save_dir(get("top.model_name"), get("top.finetuning_type"), get(
+            "{}.output_dir".format("train" if self.do_train else "eval")
+        ))
         while self.thread.is_alive():
             time.sleep(2)
             if self.aborted:

src/llmtuner/webui/utils.py

@@ -1,17 +1,20 @@
 import os
 import json
 import gradio as gr
-import matplotlib.figure
-import matplotlib.pyplot as plt
 from typing import TYPE_CHECKING, Any, Dict
 from datetime import datetime
 
+from llmtuner.extras.packages import is_matplotlib_available
 from llmtuner.extras.ploting import smooth
 from llmtuner.webui.common import get_save_dir
 
 if TYPE_CHECKING:
     from llmtuner.extras.callbacks import LogCallback
 
+if is_matplotlib_available():
+    import matplotlib.figure
+    import matplotlib.pyplot as plt
+
 
 def update_process_bar(callback: "LogCallback") -> Dict[str, Any]:
     if not callback.max_steps:
@@ -56,7 +59,7 @@ def get_eval_results(path: os.PathLike) -> str:
     return "```json\n{}\n```\n".format(result)
 
 
-def gen_plot(base_model: str, finetuning_type: str, output_dir: str) -> matplotlib.figure.Figure:
+def gen_plot(base_model: str, finetuning_type: str, output_dir: str) -> "matplotlib.figure.Figure":
     if not base_model:
         return
     log_file = get_save_dir(base_model, finetuning_type, output_dir, "trainer_log.jsonl")

tests/cal_flops.py

@@ -12,7 +12,7 @@ from deepspeed.profiling.flops_profiler import get_model_profile # type: ignore
 from llmtuner import ChatModel
 
 
-def calculate(
+def calculate_flops(
     model_name_or_path: str,
     batch_size: Optional[int] = 1,
     seq_length: Optional[int] = 256,
@@ -41,4 +41,4 @@ def calculate(
 
 
 if __name__ == "__main__":
-    fire.Fire(calculate)
+    fire.Fire(calculate_flops)

tests/cal_lr.py

@@ -0,0 +1,63 @@
+# coding=utf-8
+# Calculates the optimal learning rate for 7B/13B models using LLaMA's hyper-parameters.
+# Usage: python cal_lr.py --model_name_or_path path_to_model --dataset alpaca_en --cutoff_len 1024 --batch_size 16
+# Inspired by: https://github.com/imoneoi/openchat/blob/master/ochat/training_deepspeed/train.py
+
+import fire
+import math
+import torch
+from tqdm import tqdm
+from typing import Optional
+from torch.utils.data import DataLoader
+from transformers import DataCollatorForSeq2Seq
+
+from llmtuner.data import get_dataset, preprocess_dataset
+from llmtuner.extras.constants import IGNORE_INDEX
+from llmtuner.model import get_train_args, load_model_and_tokenizer
+
+
+BASE_LR = 3e-4      # 1.5e-4 for 30B-70B models
+BASE_BS = 4_000_000 # from llama paper
+
+
+def calculate_lr(
+    model_name_or_path: str,
+    dataset: str,
+    cutoff_len: int,  # i.e. maximum input length during training
+    batch_size: int,  # total batch size, namely (batch size * gradient accumulation * world size)
+    is_mistral: bool, # mistral model uses a smaller learning rate,
+    dataset_dir: Optional[str] = "data"
+):
+    model_args, data_args, training_args, finetuning_args, _ = get_train_args(dict(
+        stage="sft",
+        model_name_or_path=model_name_or_path,
+        dataset=dataset,
+        dataset_dir=dataset_dir,
+        template="default",
+        cutoff_len=cutoff_len,
+        output_dir="dummy_dir"
+    ))
+    trainset = get_dataset(model_args, data_args)
+    _, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, is_trainable=False, stage="sft")
+    trainset = preprocess_dataset(trainset, tokenizer, data_args, training_args, stage="sft")
+    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, label_pad_token_id=IGNORE_INDEX)
+    dataloader = DataLoader(
+        dataset=trainset, batch_size=batch_size, shuffle=True, collate_fn=data_collator, pin_memory=True
+    )
+    valid_tokens, total_tokens = 0, 0
+    for batch in tqdm(dataloader):
+        valid_tokens += torch.sum(batch["labels"] != IGNORE_INDEX).item()
+        total_tokens += torch.numel(batch["labels"])
+
+    batch_max_len = cutoff_len * batch_size # max tokens in a batch
+    valid_ratio = valid_tokens / total_tokens
+    batch_valid_len = batch_max_len * valid_ratio
+    lr = BASE_LR * math.sqrt(batch_valid_len / BASE_BS) # lr ~ sqrt(batch_size)
+    lr = lr / 6.0 if is_mistral else lr
+    print("Optimal learning rate is {:.2e} for valid ratio% {:.2f} and effective batch size {:.2f}".format(
+        lr, valid_ratio * 100, batch_valid_len
+    ))
+
+
+if __name__ == "__main__":
+    fire.Fire(calculate_lr)

tests/quantize.py

@@ -4,7 +4,6 @@
 #                           --max_length 1024 --max_samples 1024
 # dataset format: instruction (string), input (string), output (string), history (List[string])
 
-
 import fire
 from datasets import load_dataset
 from transformers import AutoTokenizer