release v0.5.0 (real)

Former-commit-id: 2146e1d9195c179fa8f92144ec2b7034e1a9f942
finish agent
2024-01-21 01:54:49 +08:00 · 2024-01-21 01:47:33 +08:00 · 2024-01-21 00:03:09 +08:00 · 2024-01-20 23:33:50 +08:00 · 2024-01-20 23:27:10 +08:00 · 2024-01-20 23:22:09 +08:00
100 changed files with 6219 additions and 3725 deletions
--- a/.github/ISSUE_TEMPLATE/bug-report.yml
+++ b/.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
--- a/.gitignore
+++ b/.gitignore
@@ -157,4 +157,9 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+.idea/
+
+# custom .gitignore
+user.config
+saves/
+cache/
--- a/11
+++ b/11
@@ -0,0 +1,11 @@
+.PHONY: quality style
+
+check_dirs := src tests
+
+quality:
+	black --check $(check_dirs)
+	ruff $(check_dirs)
+
+style:
+	black $(check_dirs)
+	ruff $(check_dirs) --fix
--- a/README.md
+++ b/README.md
@@ -1,4 +1,4 @@
-# LLaMA Factory: Training and Evaluating Large Language Models with Minimal Effort
+![# LLaMA Factory](assets/logo.png)

 [![GitHub Repo stars](https://img.shields.io/github/stars/hiyouga/LLaMA-Factory?style=social)](https://github.com/hiyouga/LLaMA-Factory/stargazers)
 [![GitHub Code License](https://img.shields.io/github/license/hiyouga/LLaMA-Factory)](LICENSE)
@@ -6,8 +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/c2EPEt5NU?compact=true&style=flat)](https://discord.gg/c2EPEt5NU)
+[![Discord](https://dcbadge.vercel.app/api/server/rKfvV9r9FK?compact=true&style=flat)](https://discord.gg/rKfvV9r9FK)
 [![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).

@@ -15,23 +16,62 @@

 ## LLaMA Board: A One-stop Web UI for Getting Started with LLaMA Factory

-Preview LLaMA Board at **[🤗 Spaces](https://huggingface.co/spaces/hiyouga/LLaMA-Board)**.
+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)
+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)
+
+<details><summary>Definitions</summary>
+
+- **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.
+
+</details>
+
 ## 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`.
+[24/01/18] We supported **agent tuning** for most models, equipping model with tool using abilities by fine-tuning with `--dataset glaive_toolcall`.
+
+[23/12/23] We supported **[unsloth](https://github.com/unslothai/unsloth)**'s implementation to boost LoRA tuning for the LLaMA, Mistral and Yi models. Try `--use_unsloth` argument to activate unsloth patch. It achieves 1.7x speed in our benchmark, check [this page](https://github.com/hiyouga/LLaMA-Factory/wiki/Performance-comparison) for details.
+
+[23/12/12] We supported fine-tuning the latest MoE model **[Mixtral 8x7B](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1)** in our framework. See hardware requirement [here](#hardware-requirement).
+
+<details><summary>Full Changelog</summary>
+
+[23/12/01] We supported downloading pre-trained models and datasets from the **[ModelScope Hub](https://modelscope.cn/models)** for Chinese mainland users. See [this tutorial](#use-modelscope-hub-optional) for usage.
+
+[23/10/21] We supported **[NEFTune](https://arxiv.org/abs/2310.05914)** trick for fine-tuning. Try `--neftune_noise_alpha` argument to activate NEFTune, e.g., `--neftune_noise_alpha 5`.

 [23/09/27] We supported **$S^2$-Attn** proposed by [LongLoRA](https://github.com/dvlab-research/LongLoRA) for the LLaMA models. Try `--shift_attn` argument to enable shift short attention.

 [23/09/23] We integrated MMLU, C-Eval and CMMLU benchmarks in this repo. See [this example](#evaluation) to evaluate your models.

-[23/09/10] We supported using **[FlashAttention-2](https://github.com/Dao-AILab/flash-attention)** for the LLaMA models. Try `--flash_attn` argument to enable FlashAttention-2 if you are using RTX4090, A100 or H100 GPUs.
+[23/09/10] We supported **[FlashAttention-2](https://github.com/Dao-AILab/flash-attention)**. Try `--flash_attn` argument to enable FlashAttention-2 if you are using RTX4090, A100 or H100 GPUs.

 [23/08/12] We supported **RoPE scaling** to extend the context length of the LLaMA models. Try `--rope_scaling linear` argument in training and `--rope_scaling dynamic` argument at inference to extrapolate the position embeddings.

@@ -51,23 +91,28 @@ https://github.com/hiyouga/LLaMA-Factory/assets/16256802/6ba60acc-e2e2-4bec-b846

 [23/06/03] We supported quantized training and inference (aka **[QLoRA](https://github.com/artidoro/qlora)**). Try `--quantization_bit 4/8` argument to work with quantized models.

+</details>
+
 ## Supported Models

 | Model                                                    | Model size                  | Default module    | Template  |
 | -------------------------------------------------------- | --------------------------- | ----------------- | --------- |
-| [Baichuan](https://github.com/baichuan-inc/Baichuan-13B) | 7B/13B                      | W_pack            | baichuan  |
-| [Baichuan2](https://github.com/baichuan-inc/Baichuan2)   | 7B/13B                      | W_pack            | baichuan2 |
+| [Baichuan2](https://huggingface.co/baichuan-inc)         | 7B/13B                      | W_pack            | baichuan2 |
 | [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    |
-| [InternLM](https://github.com/InternLM/InternLM)         | 7B/20B                      | q_proj,v_proj     | intern    |
+| [ChatGLM3](https://huggingface.co/THUDM/chatglm3-6b)     | 6B                          | query_key_value   | chatglm3  |
+| [DeepSeek (MoE)](https://huggingface.co/deepseek-ai)     | 7B/16B/67B                  | q_proj,v_proj     | deepseek  |
+| [Falcon](https://huggingface.co/tiiuae)                  | 7B/40B/180B                 | query_key_value   | falcon    |
+| [InternLM2](https://huggingface.co/internlm)             | 7B/20B                      | wqkv              | intern2   |
 | [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    |
 | [Mistral](https://huggingface.co/mistralai)              | 7B                          | q_proj,v_proj     | mistral   |
-| [Phi-1.5](https://huggingface.co/microsoft/phi-1_5)      | 1.3B                        | Wqkv              | -         |
-| [Qwen](https://github.com/QwenLM/Qwen)                   | 7B/14B                      | c_attn            | qwen      |
-| [XVERSE](https://github.com/xverse-ai)                   | 7B/13B/65B                  | q_proj,v_proj     | xverse    |
+| [Mixtral](https://huggingface.co/mistralai)              | 8x7B                        | q_proj,v_proj     | mistral   |
+| [Phi-1.5/2](https://huggingface.co/microsoft)            | 1.3B/2.7B                   | q_proj,v_proj     | -         |
+| [Qwen](https://huggingface.co/Qwen)                      | 1.8B/7B/14B/72B             | c_attn            | qwen      |
+| [XVERSE](https://huggingface.co/xverse)                  | 7B/13B/65B                  | q_proj,v_proj     | xverse    |
+| [Yi](https://huggingface.co/01-ai)                       | 6B/34B                      | q_proj,v_proj     | yi        |
+| [Yuan](https://huggingface.co/IEITYuan)                  | 2B/51B/102B                 | q_proj,v_proj     | yuan      |

 > [!NOTE]
 > **Default module** is used for the `--lora_target` argument, you can use `--lora_target all` to specify all the available modules.
@@ -87,7 +132,7 @@ Please refer to [constants.py](src/llmtuner/extras/constants.py) for a full list
 | DPO Training           | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: |

 > [!NOTE]
-> Use `--quantization_bit 4/8` argument to enable QLoRA.
+> Use `--quantization_bit 4` argument to enable QLoRA.

 ## Provided Datasets

@@ -130,6 +175,8 @@ Please refer to [constants.py](src/llmtuner/extras/constants.py) for a full list
 - [Firefly 1.1M (zh)](https://huggingface.co/datasets/YeungNLP/firefly-train-1.1M)
 - [Web QA (zh)](https://huggingface.co/datasets/suolyer/webqa)
 - [WebNovel (zh)](https://huggingface.co/datasets/zxbsmk/webnovel_cn)
+- [Nectar (en)](https://huggingface.co/datasets/berkeley-nest/Nectar)
+- [deepctrl (en&zh)](https://www.modelscope.cn/datasets/deepctrl/deepctrl-sft-data)
 - [Ad Gen (zh)](https://huggingface.co/datasets/HasturOfficial/adgen)
 - [ShareGPT Hyperfiltered (en)](https://huggingface.co/datasets/totally-not-an-llm/sharegpt-hyperfiltered-3k)
 - [ShareGPT4 (en&zh)](https://huggingface.co/datasets/shibing624/sharegpt_gpt4)
@@ -137,6 +184,7 @@ Please refer to [constants.py](src/llmtuner/extras/constants.py) for a full list
 - [AgentInstruct (en)](https://huggingface.co/datasets/THUDM/AgentInstruct)
 - [LMSYS Chat 1M (en)](https://huggingface.co/datasets/lmsys/lmsys-chat-1m)
 - [Evol Instruct V2 (en)](https://huggingface.co/datasets/WizardLM/WizardLM_evol_instruct_V2_196k)
+- [Glaive Function Calling V2 (en)](https://huggingface.co/datasets/glaiveai/glaive-function-calling-v2)

 </details>

@@ -145,6 +193,7 @@ Please refer to [constants.py](src/llmtuner/extras/constants.py) for a full list
 - [HH-RLHF (en)](https://huggingface.co/datasets/Anthropic/hh-rlhf)
 - [Open Assistant (multilingual)](https://huggingface.co/datasets/OpenAssistant/oasst1)
 - [GPT-4 Generated Data (en&zh)](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)
+- [Nectar (en)](https://huggingface.co/datasets/berkeley-nest/Nectar)

 </details>

@@ -166,7 +215,15 @@ huggingface-cli login
 - gradio and matplotlib (used in web UI)
 - uvicorn, fastapi and sse-starlette (used in API)

-And **powerful GPUs**!
+### Hardware Requirement
+
+| Method | Bits |   7B  |  13B  |  30B  |   65B  |   8x7B |
+| ------ | ---- | ----- | ----- | ----- | ------ | ------ |
+| Full   |  16  | 160GB | 320GB | 600GB | 1200GB |  900GB |
+| Freeze |  16  |  20GB |  40GB | 120GB |  240GB |  200GB |
+| LoRA   |  16  |  16GB |  32GB |  80GB |  160GB |  120GB |
+| QLoRA  |   8  |  10GB |  16GB |  40GB |   80GB |   80GB |
+| QLoRA  |   4  |   6GB |  12GB |  24GB |   48GB |   32GB |

 ## Getting Started

@@ -193,6 +250,28 @@ If you want to enable the quantized LoRA (QLoRA) on the Windows platform, you wi
 pip install https://github.com/jllllll/bitsandbytes-windows-webui/releases/download/wheels/bitsandbytes-0.39.1-py3-none-win_amd64.whl
 ```

+### Use ModelScope Hub (optional)
+
+If you have trouble with downloading models and datasets from Hugging Face, you can use LLaMA-Factory together with ModelScope in the following manner.
+
+```bash
+export USE_MODELSCOPE_HUB=1 # `set USE_MODELSCOPE_HUB=1` for Windows
+```
+
+Then you can train the corresponding model by specifying a model ID of the ModelScope Hub. (find a full list of model IDs at [ModelScope Hub](https://modelscope.cn/models))
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
+    --model_name_or_path modelscope/Llama-2-7b-ms \
+    ... # arguments (same as above)
+```
+
+LLaMA Board also supports using the models and datasets on the ModelScope Hub.
+
+```bash
+CUDA_VISIBLE_DEVICES=0 USE_MODELSCOPE_HUB=1 python src/train_web.py
+```
+
 ### Train on a single GPU

 > [!IMPORTANT]
@@ -203,8 +282,8 @@ pip install https://github.com/jllllll/bitsandbytes-windows-webui/releases/downl
 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage pt \
-    --model_name_or_path path_to_llama_model \
    --do_train \
+    --model_name_or_path path_to_llama_model \
    --dataset wiki_demo \
    --finetuning_type lora \
    --lora_target q_proj,v_proj \
@@ -226,8 +305,8 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage sft \
-    --model_name_or_path path_to_llama_model \
    --do_train \
+    --model_name_or_path path_to_llama_model \
    --dataset alpaca_gpt4_en \
    --template default \
    --finetuning_type lora \
@@ -250,14 +329,14 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage rm \
-    --model_name_or_path path_to_llama_model \
    --do_train \
+    --model_name_or_path path_to_llama_model \
+    --adapter_name_or_path path_to_sft_checkpoint \
+    --create_new_adapter \
    --dataset comparison_gpt4_en \
    --template default \
    --finetuning_type lora \
    --lora_target q_proj,v_proj \
-    --resume_lora_training False \
-    --checkpoint_dir path_to_sft_checkpoint \
    --output_dir path_to_rm_checkpoint \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 4 \
@@ -275,19 +354,21 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage ppo \
-    --model_name_or_path path_to_llama_model \
    --do_train \
+    --model_name_or_path path_to_llama_model \
+    --adapter_name_or_path path_to_sft_checkpoint \
+    --create_new_adapter \
    --dataset alpaca_gpt4_en \
    --template default \
    --finetuning_type lora \
    --lora_target q_proj,v_proj \
-    --resume_lora_training False \
-    --checkpoint_dir path_to_sft_checkpoint \
    --reward_model path_to_rm_checkpoint \
    --output_dir path_to_ppo_checkpoint \
    --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 \
@@ -296,19 +377,22 @@ 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
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage dpo \
-    --model_name_or_path path_to_llama_model \
    --do_train \
+    --model_name_or_path path_to_llama_model \
+    --adapter_name_or_path path_to_sft_checkpoint \
+    --create_new_adapter \
    --dataset comparison_gpt4_en \
    --template default \
    --finetuning_type lora \
    --lora_target q_proj,v_proj \
-    --resume_lora_training False \
-    --checkpoint_dir path_to_sft_checkpoint \
    --output_dir path_to_dpo_checkpoint \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 4 \
@@ -391,28 +475,36 @@ 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 \
    --model_name_or_path path_to_llama_model \
+    --adapter_name_or_path path_to_checkpoint \
    --template default \
    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint \
-    --export_dir path_to_export
+    --export_dir path_to_export \
+    --export_size 2 \
+    --export_legacy_format False
 ```

+> [!WARNING]
+> Merging LoRA weights into a quantized model is not supported.
+
+> [!TIP]
+> Use `--export_quantization_bit 4` and `--export_quantization_dataset data/c4_demo.json` to quantize the model after merging the LoRA weights.
+
 ### API Demo

 ```bash
 python src/api_demo.py \
    --model_name_or_path path_to_llama_model \
+    --adapter_name_or_path path_to_checkpoint \
    --template default \
-    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint
+    --finetuning_type lora
 ```

-> [!NOTE]
+> [!TIP]
 > Visit `http://localhost:8000/docs` for API documentation.

 ### CLI Demo
@@ -420,9 +512,9 @@ python src/api_demo.py \
 ```bash
 python src/cli_demo.py \
    --model_name_or_path path_to_llama_model \
+    --adapter_name_or_path path_to_checkpoint \
    --template default \
-    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint
+    --finetuning_type lora
 ```

 ### Web Demo
@@ -430,9 +522,9 @@ python src/cli_demo.py \
 ```bash
 python src/web_demo.py \
    --model_name_or_path path_to_llama_model \
+    --adapter_name_or_path path_to_checkpoint \
    --template default \
-    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint
+    --finetuning_type lora
 ```

 ### Evaluation
@@ -440,9 +532,9 @@ python src/web_demo.py \
 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/evaluate.py \
    --model_name_or_path path_to_llama_model \
-    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint \
+    --adapter_name_or_path path_to_checkpoint \
    --template vanilla \
+    --finetuning_type lora \
    --task mmlu \
    --split test \
    --lang en \
@@ -455,19 +547,23 @@ CUDA_VISIBLE_DEVICES=0 python src/evaluate.py \
 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage sft \
-    --model_name_or_path path_to_llama_model \
    --do_predict \
+    --model_name_or_path path_to_llama_model \
+    --adapter_name_or_path path_to_checkpoint \
    --dataset alpaca_gpt4_en \
    --template default \
    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint \
    --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
@@ -476,12 +572,16 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
 - **[DISC-LawLLM](https://github.com/FudanDISC/DISC-LawLLM)**: A large language model specialized in Chinese legal domain, based on Baichuan-13B, is capable of retrieving and reasoning on legal knowledge.
 - **[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.
+- **[MachineMindset](https://github.com/PKU-YuanGroup/Machine-Mindset/)**: A series of MBTI Personality large language models, capable of giving any LLM 16 different personality types based on different datasets and training methods.
+
+> [!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).

-Please follow the model licenses to use the corresponding model weights: [Baichuan](https://huggingface.co/baichuan-inc/Baichuan-13B-Base/resolve/main/Community%20License%20for%20Baichuan-13B%20Model.pdf) / [Baichuan2](https://huggingface.co/baichuan-inc/Baichuan2-13B-Chat/resolve/main/Community%20License%20for%20Baichuan2%20Model.pdf) / [BLOOM](https://huggingface.co/spaces/bigscience/license) / [ChatGLM3](https://github.com/THUDM/ChatGLM3/blob/main/MODEL_LICENSE) / [Falcon](https://huggingface.co/tiiuae/falcon-180B/blob/main/LICENSE.txt) / [InternLM](https://github.com/InternLM/InternLM#license) / [LLaMA](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) / [LLaMA-2](https://ai.meta.com/llama/license/) / [Mistral](LICENSE) / [Phi-1.5](https://huggingface.co/microsoft/phi-1_5/resolve/main/Research%20License.docx) / [Qwen](https://github.com/QwenLM/Qwen/blob/main/LICENSE) / [XVERSE](https://github.com/xverse-ai/XVERSE-13B/blob/main/MODEL_LICENSE.pdf)
+Please follow the model licenses to use the corresponding model weights: [Baichuan2](https://huggingface.co/baichuan-inc/Baichuan2-7B-Base/blob/main/Community%20License%20for%20Baichuan%202%20Model.pdf) / [BLOOM](https://huggingface.co/spaces/bigscience/license) / [ChatGLM3](https://github.com/THUDM/ChatGLM3/blob/main/MODEL_LICENSE) / [DeepSeek](https://github.com/deepseek-ai/DeepSeek-LLM/blob/main/LICENSE-MODEL) / [Falcon](https://huggingface.co/tiiuae/falcon-180B/blob/main/LICENSE.txt) / [InternLM2](https://github.com/InternLM/InternLM#license) / [LLaMA](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) / [LLaMA-2](https://ai.meta.com/llama/license/) / [Mistral](LICENSE) / [Phi-1.5/2](https://huggingface.co/microsoft/phi-1_5/resolve/main/Research%20License.docx) / [Qwen](https://github.com/QwenLM/Qwen/blob/main/Tongyi%20Qianwen%20LICENSE%20AGREEMENT) / [XVERSE](https://github.com/xverse-ai/XVERSE-13B/blob/main/MODEL_LICENSE.pdf) / [Yi](https://huggingface.co/01-ai/Yi-6B/blob/main/LICENSE) / [Yuan](https://github.com/IEIT-Yuan/Yuan-2.0/blob/main/LICENSE-Yuan)

 ## Citation

--- a/README_zh.md
+++ b/README_zh.md
@@ -1,4 +1,4 @@
-# LLaMA Factory: 轻松的大模型训练与评估
+![# LLaMA Factory](assets/logo.png)

 [![GitHub Repo stars](https://img.shields.io/github/stars/hiyouga/LLaMA-Factory?style=social)](https://github.com/hiyouga/LLaMA-Factory/stargazers)
 [![GitHub Code License](https://img.shields.io/github/license/hiyouga/LLaMA-Factory)](LICENSE)
@@ -6,8 +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/c2EPEt5NU?compact=true&style=flat)](https://discord.gg/c2EPEt5NU)
+[![Discord](https://dcbadge.vercel.app/api/server/rKfvV9r9FK?compact=true&style=flat)](https://discord.gg/rKfvV9r9FK)
 [![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)。

@@ -15,7 +16,7 @@

 ## LLaMA Board: 通过一站式网页界面快速上手 LLaMA Factory

-通过 **[🤗 Spaces](https://huggingface.co/spaces/hiyouga/LLaMA-Board)** 预览 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。（该模式目前仅支持单卡训练）

@@ -23,21 +24,60 @@

 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)
+
+<details><summary>变量定义</summary>
+
+- **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`。
+
+</details>
+
 ## 更新日志

-[23/10/21] 我们支持了 **[NEFTune](https://arxiv.org/abs/2310.05914)** 训练技巧。请使用 `--neft_alpha` 参数启用 NEFTune，例如 `--neft_alpha 5`。
+[24/01/18] 我们针对绝大多数模型实现了 **Agent 微调**，微调时指定 `--dataset glaive_toolcall` 即可使模型获得工具调用能力。
+
+[23/12/23] 我们针对 LLaMA, Mistral 和 Yi 模型支持了 **[unsloth](https://github.com/unslothai/unsloth)** 的 LoRA 训练加速。请使用 `--use_unsloth` 参数启用 unsloth 优化。该方法可提供 1.7 倍的训练速度，详情请查阅[此页面](https://github.com/hiyouga/LLaMA-Factory/wiki/Performance-comparison)。
+
+[23/12/12] 我们支持了微调最新的混合专家模型 **[Mixtral 8x7B](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1)**。硬件需求请查阅[此处](#硬件依赖)。
+
+<details><summary>展开日志</summary>
+
+[23/12/01] 我们支持了从 **[魔搭社区](https://modelscope.cn/models)** 下载预训练模型和数据集。详细用法请参照 [此教程](#使用魔搭社区可跳过)。
+
+[23/10/21] 我们支持了 **[NEFTune](https://arxiv.org/abs/2310.05914)** 训练技巧。请使用 `--neftune_noise_alpha` 参数启用 NEFTune，例如 `--neftune_noise_alpha 5`。

 [23/09/27] 我们针对 LLaMA 模型支持了 [LongLoRA](https://github.com/dvlab-research/LongLoRA) 提出的 **$S^2$-Attn**。请使用 `--shift_attn` 参数以启用该功能。

 [23/09/23] 我们在项目中集成了 MMLU、C-Eval 和 CMMLU 评估集。使用方法请参阅[此示例](#模型评估)。

-[23/09/10] 我们针对 LLaMA 模型支持了 **[FlashAttention-2](https://github.com/Dao-AILab/flash-attention)**。如果您使用的是 RTX4090、A100 或 H100 GPU，请使用 `--flash_attn` 参数以启用 FlashAttention-2。
+[23/09/10] 我们支持了 **[FlashAttention-2](https://github.com/Dao-AILab/flash-attention)**。如果您使用的是 RTX4090、A100 或 H100 GPU，请使用 `--flash_attn` 参数以启用 FlashAttention-2。

 [23/08/12] 我们支持了 **RoPE 插值**来扩展 LLaMA 模型的上下文长度。请使用 `--rope_scaling linear` 参数训练模型或使用 `--rope_scaling dynamic` 参数评估模型。

 [23/08/11] 我们支持了指令模型的 **[DPO 训练](https://arxiv.org/abs/2305.18290)**。使用方法请参阅[此示例](#dpo-训练)。

-[23/07/31] 我们支持了**数据流式加载**。请尝试使用 `--streaming` 和 `--max_steps 10000` 参数来流式加载数据集。
+[23/07/31] 我们支持了**数据流式加载**。请使用 `--streaming` 和 `--max_steps 10000` 参数来流式加载数据集。

 [23/07/29] 我们在 Hugging Face 发布了两个 13B 指令微调模型。详细内容请查阅我们的 Hugging Face 项目（[LLaMA-2](https://huggingface.co/hiyouga/Llama-2-Chinese-13b-chat) / [Baichuan](https://huggingface.co/hiyouga/Baichuan-13B-sft)）。

@@ -51,23 +91,28 @@ https://github.com/hiyouga/LLaMA-Factory/assets/16256802/6ba60acc-e2e2-4bec-b846

 [23/06/03] 我们实现了 4 比特的 LoRA 训练（也称 **[QLoRA](https://github.com/artidoro/qlora)**）。请使用 `--quantization_bit 4` 参数进行 4 比特量化微调。

+</details>
+
 ## 模型

 | 模型名                                                   | 模型大小                     | 默认模块           | Template  |
 | -------------------------------------------------------- | --------------------------- | ----------------- | --------- |
-| [Baichuan](https://github.com/baichuan-inc/Baichuan-13B) | 7B/13B                      | W_pack            | baichuan  |
-| [Baichuan2](https://github.com/baichuan-inc/Baichuan2)   | 7B/13B                      | W_pack            | baichuan2 |
+| [Baichuan2](https://huggingface.co/baichuan-inc)         | 7B/13B                      | W_pack            | baichuan2 |
 | [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    |
-| [InternLM](https://github.com/InternLM/InternLM)         | 7B/20B                      | q_proj,v_proj     | intern    |
+| [ChatGLM3](https://huggingface.co/THUDM/chatglm3-6b)     | 6B                          | query_key_value   | chatglm3  |
+| [DeepSeek (MoE)](https://huggingface.co/deepseek-ai)     | 7B/16B/67B                  | q_proj,v_proj     | deepseek  |
+| [Falcon](https://huggingface.co/tiiuae)                  | 7B/40B/180B                 | query_key_value   | falcon    |
+| [InternLM2](https://huggingface.co/internlm)             | 7B/20B                      | wqkv              | intern2   |
 | [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    |
 | [Mistral](https://huggingface.co/mistralai)              | 7B                          | q_proj,v_proj     | mistral   |
-| [Phi-1.5](https://huggingface.co/microsoft/phi-1_5)      | 1.3B                        | Wqkv              | -         |
-| [Qwen](https://github.com/QwenLM/Qwen)                   | 7B/14B                      | c_attn            | qwen      |
-| [XVERSE](https://github.com/xverse-ai)                   | 7B/13B/65B                  | q_proj,v_proj     | xverse    |
+| [Mixtral](https://huggingface.co/mistralai)              | 8x7B                        | q_proj,v_proj     | mistral   |
+| [Phi-1.5/2](https://huggingface.co/microsoft)            | 1.3B/2.7B                   | q_proj,v_proj     | -         |
+| [Qwen](https://huggingface.co/Qwen)                      | 1.8B/7B/14B/72B             | c_attn            | qwen      |
+| [XVERSE](https://huggingface.co/xverse)                  | 7B/13B/65B                  | q_proj,v_proj     | xverse    |
+| [Yi](https://huggingface.co/01-ai)                       | 6B/34B                      | q_proj,v_proj     | yi        |
+| [Yuan](https://huggingface.co/IEITYuan)                  | 2B/51B/102B                 | q_proj,v_proj     | yuan      |

 > [!NOTE]
 > **默认模块**应作为 `--lora_target` 参数的默认值，可使用 `--lora_target all` 参数指定全部模块。
@@ -87,7 +132,7 @@ https://github.com/hiyouga/LLaMA-Factory/assets/16256802/6ba60acc-e2e2-4bec-b846
 | DPO 训练               | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: |

 > [!NOTE]
-> 请使用 `--quantization_bit 4/8` 参数来启用 QLoRA 训练。
+> 请使用 `--quantization_bit 4` 参数来启用 QLoRA 训练。

 ## 数据集

@@ -130,6 +175,8 @@ https://github.com/hiyouga/LLaMA-Factory/assets/16256802/6ba60acc-e2e2-4bec-b846
 - [Firefly 1.1M (zh)](https://huggingface.co/datasets/YeungNLP/firefly-train-1.1M)
 - [Web QA (zh)](https://huggingface.co/datasets/suolyer/webqa)
 - [WebNovel (zh)](https://huggingface.co/datasets/zxbsmk/webnovel_cn)
+- [Nectar (en)](https://huggingface.co/datasets/berkeley-nest/Nectar)
+- [deepctrl (en&zh)](https://www.modelscope.cn/datasets/deepctrl/deepctrl-sft-data)
 - [Ad Gen (zh)](https://huggingface.co/datasets/HasturOfficial/adgen)
 - [ShareGPT Hyperfiltered (en)](https://huggingface.co/datasets/totally-not-an-llm/sharegpt-hyperfiltered-3k)
 - [ShareGPT4 (en&zh)](https://huggingface.co/datasets/shibing624/sharegpt_gpt4)
@@ -137,6 +184,7 @@ https://github.com/hiyouga/LLaMA-Factory/assets/16256802/6ba60acc-e2e2-4bec-b846
 - [AgentInstruct (en)](https://huggingface.co/datasets/THUDM/AgentInstruct)
 - [LMSYS Chat 1M (en)](https://huggingface.co/datasets/lmsys/lmsys-chat-1m)
 - [Evol Instruct V2 (en)](https://huggingface.co/datasets/WizardLM/WizardLM_evol_instruct_V2_196k)
+- [Glaive Function Calling V2 (en)](https://huggingface.co/datasets/glaiveai/glaive-function-calling-v2)

 </details>

@@ -145,6 +193,7 @@ https://github.com/hiyouga/LLaMA-Factory/assets/16256802/6ba60acc-e2e2-4bec-b846
 - [HH-RLHF (en)](https://huggingface.co/datasets/Anthropic/hh-rlhf)
 - [Open Assistant (multilingual)](https://huggingface.co/datasets/OpenAssistant/oasst1)
 - [GPT-4 Generated Data (en&zh)](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)
+- [Nectar (en)](https://huggingface.co/datasets/berkeley-nest/Nectar)

 </details>

@@ -157,7 +206,7 @@ pip install --upgrade huggingface_hub
 huggingface-cli login
 ```

-## 软件依赖
+## 软硬件依赖

 - Python 3.8+ 和 PyTorch 1.13.1+
 - 🤗Transformers, Datasets, Accelerate, PEFT 和 TRL
@@ -166,7 +215,15 @@ huggingface-cli login
 - gradio 和 matplotlib (用于网页端交互)
 - uvicorn, fastapi 和 sse-starlette (用于 API)

-以及 **强而有力的 GPU**！
+### 硬件依赖
+
+| 训练方法 | 精度 |   7B  |  13B  |  30B  |   65B  |   8x7B |
+| ------- | ---- | ----- | ----- | ----- | ------ | ------ |
+| 全参数   |  16  | 160GB | 320GB | 600GB | 1200GB |  900GB |
+| 部分参数 |  16  |  20GB |  40GB | 120GB |  240GB |  200GB |
+| LoRA    |  16  |  16GB |  32GB |  80GB |  160GB |  120GB |
+| QLoRA   |   8  |  10GB |  16GB |  40GB |   80GB |   80GB |
+| QLoRA   |   4  |   6GB |  12GB |  24GB |   48GB |   32GB |

 ## 如何使用

@@ -193,6 +250,28 @@ pip install -r requirements.txt
 pip install https://github.com/jllllll/bitsandbytes-windows-webui/releases/download/wheels/bitsandbytes-0.39.1-py3-none-win_amd64.whl
 ```

+### 使用魔搭社区（可跳过）
+
+如果您在 Hugging Face 模型和数据集的下载中遇到了问题，可以通过下述方法使用魔搭社区。
+
+```bash
+export USE_MODELSCOPE_HUB=1 # Windows 使用 `set USE_MODELSCOPE_HUB=1`
+```
+
+接着即可通过指定模型名称来训练对应的模型。（在[魔搭社区](https://modelscope.cn/models)查看所有可用的模型）
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
+    --model_name_or_path modelscope/Llama-2-7b-ms \
+    ... # 参数同上
+```
+
+LLaMA Board 同样支持魔搭社区的模型和数据集下载。
+
+```bash
+CUDA_VISIBLE_DEVICES=0 USE_MODELSCOPE_HUB=1 python src/train_web.py
+```
+
 ### 单 GPU 训练

 > [!IMPORTANT]
@@ -203,8 +282,8 @@ pip install https://github.com/jllllll/bitsandbytes-windows-webui/releases/downl
 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage pt \
-    --model_name_or_path path_to_llama_model \
    --do_train \
+    --model_name_or_path path_to_llama_model \
    --dataset wiki_demo \
    --finetuning_type lora \
    --lora_target q_proj,v_proj \
@@ -226,8 +305,8 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage sft \
-    --model_name_or_path path_to_llama_model \
    --do_train \
+    --model_name_or_path path_to_llama_model \
    --dataset alpaca_gpt4_zh \
    --template default \
    --finetuning_type lora \
@@ -250,14 +329,14 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage rm \
-    --model_name_or_path path_to_llama_model \
    --do_train \
+    --model_name_or_path path_to_llama_model \
+    --adapter_name_or_path path_to_sft_checkpoint \
+    --create_new_adapter \
    --dataset comparison_gpt4_zh \
    --template default \
    --finetuning_type lora \
    --lora_target q_proj,v_proj \
-    --resume_lora_training False \
-    --checkpoint_dir path_to_sft_checkpoint \
    --output_dir path_to_rm_checkpoint \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 4 \
@@ -275,39 +354,45 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage ppo \
-    --model_name_or_path path_to_llama_model \
    --do_train \
+    --model_name_or_path path_to_llama_model \
+    --adapter_name_or_path path_to_sft_checkpoint \
+    --create_new_adapter \
    --dataset alpaca_gpt4_zh \
    --template default \
    --finetuning_type lora \
    --lora_target q_proj,v_proj \
-    --resume_lora_training False \
-    --checkpoint_dir path_to_sft_checkpoint \
    --reward_model path_to_rm_checkpoint \
    --output_dir path_to_ppo_checkpoint \
    --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
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage dpo \
-    --model_name_or_path path_to_llama_model \
    --do_train \
+    --model_name_or_path path_to_llama_model \
+    --adapter_name_or_path path_to_sft_checkpoint \
+    --create_new_adapter \
    --dataset comparison_gpt4_zh \
    --template default \
    --finetuning_type lora \
    --lora_target q_proj,v_proj \
-    --resume_lora_training False \
-    --checkpoint_dir path_to_sft_checkpoint \
    --output_dir path_to_dpo_checkpoint \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 4 \
@@ -390,28 +475,36 @@ deepspeed --num_gpus 8 --master_port=9901 src/train_bash.py \

 </details>

-### 导出微调后的完整模型
+### 合并 LoRA 权重并导出模型

 ```bash
 python src/export_model.py \
    --model_name_or_path path_to_llama_model \
+    --adapter_name_or_path path_to_checkpoint \
    --template default \
    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint \
-    --export_dir path_to_export
+    --export_dir path_to_export \
+    --export_size 2 \
+    --export_legacy_format False
 ```

+> [!WARNING]
+> 尚不支持量化模型的 LoRA 权重合并及导出。
+
+> [!TIP]
+> 合并 LoRA 权重之后可再次使用 `--export_quantization_bit 4` 和 `--export_quantization_dataset data/c4_demo.json` 量化模型。
+
 ### API 服务

 ```bash
 python src/api_demo.py \
    --model_name_or_path path_to_llama_model \
+    --adapter_name_or_path path_to_checkpoint \
    --template default \
-    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint
+    --finetuning_type lora
 ```

-> [!NOTE]
+> [!TIP]
 > 关于 API 文档请见 `http://localhost:8000/docs`。

 ### 命令行测试
@@ -419,9 +512,9 @@ python src/api_demo.py \
 ```bash
 python src/cli_demo.py \
    --model_name_or_path path_to_llama_model \
+    --adapter_name_or_path path_to_checkpoint \
    --template default \
-    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint
+    --finetuning_type lora
 ```

 ### 浏览器测试
@@ -429,9 +522,9 @@ python src/cli_demo.py \
 ```bash
 python src/web_demo.py \
    --model_name_or_path path_to_llama_model \
+    --adapter_name_or_path path_to_checkpoint \
    --template default \
-    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint
+    --finetuning_type lora
 ```

 ### 模型评估
@@ -439,9 +532,9 @@ python src/web_demo.py \
 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/evaluate.py \
    --model_name_or_path path_to_llama_model \
-    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint \
+    --adapter_name_or_path path_to_checkpoint \
    --template vanilla \
+    --finetuning_type lora \
    --task ceval \
    --split validation \
    --lang zh \
@@ -454,19 +547,23 @@ CUDA_VISIBLE_DEVICES=0 python src/evaluate.py \
 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage sft \
-    --model_name_or_path path_to_llama_model \
    --do_predict \
+    --model_name_or_path path_to_llama_model \
+    --adapter_name_or_path path_to_checkpoint \
    --dataset alpaca_gpt4_zh \
    --template default \
    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint \
    --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 的项目
@@ -475,12 +572,16 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
 - **[DISC-LawLLM](https://github.com/FudanDISC/DISC-LawLLM)**: 中文法律领域大模型 DISC-LawLLM，基于 Baichuan-13B 微调而得，具有法律推理和知识检索能力。
 - **[Sunsimiao](https://github.com/thomas-yanxin/Sunsimiao)**: 孙思邈中文医疗大模型 Sumsimiao，基于 Baichuan-7B 和 ChatGLM-6B 在中文医疗数据上微调而得。
 - **[CareGPT](https://github.com/WangRongsheng/CareGPT)**: 医疗大模型项目 CareGPT，基于 LLaMA2-7B 和 Baichuan-13B 在中文医疗数据上微调而得。
+- **[MachineMindset](https://github.com/PKU-YuanGroup/Machine-Mindset/)**：MBTI性格大模型项目，根据数据集与训练方式让任意 LLM 拥有 16 个不同的性格类型。
+
+> [!TIP]
+> 如果您有项目希望添加至上述列表，请通过邮件联系或者创建一个 PR。

 ## 协议

 本仓库的代码依照 [Apache-2.0](LICENSE) 协议开源。

-使用模型权重时，请遵循对应的模型协议：[Baichuan](https://huggingface.co/baichuan-inc/Baichuan-13B-Base/resolve/main/Community%20License%20for%20Baichuan-13B%20Model.pdf) / [Baichuan2](https://huggingface.co/baichuan-inc/Baichuan2-13B-Chat/resolve/main/Community%20License%20for%20Baichuan2%20Model.pdf) / [BLOOM](https://huggingface.co/spaces/bigscience/license) / [ChatGLM3](https://github.com/THUDM/ChatGLM3/blob/main/MODEL_LICENSE) / [Falcon](https://huggingface.co/tiiuae/falcon-180B/blob/main/LICENSE.txt) / [InternLM](https://github.com/InternLM/InternLM#license) / [LLaMA](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) / [LLaMA-2](https://ai.meta.com/llama/license/) / [Mistral](LICENSE) / [Phi-1.5](https://huggingface.co/microsoft/phi-1_5/resolve/main/Research%20License.docx) / [Qwen](https://github.com/QwenLM/Qwen/blob/main/LICENSE) / [XVERSE](https://github.com/xverse-ai/XVERSE-13B/blob/main/MODEL_LICENSE.pdf)
+使用模型权重时，请遵循对应的模型协议：[Baichuan2](https://huggingface.co/baichuan-inc/Baichuan2-7B-Base/blob/main/Community%20License%20for%20Baichuan%202%20Model.pdf) / [BLOOM](https://huggingface.co/spaces/bigscience/license) / [ChatGLM3](https://github.com/THUDM/ChatGLM3/blob/main/MODEL_LICENSE) / [DeepSeek](https://github.com/deepseek-ai/DeepSeek-LLM/blob/main/LICENSE-MODEL) / [Falcon](https://huggingface.co/tiiuae/falcon-180B/blob/main/LICENSE.txt) / [InternLM2](https://github.com/InternLM/InternLM#license) / [LLaMA](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) / [LLaMA-2](https://ai.meta.com/llama/license/) / [Mistral](LICENSE) / [Phi-1.5/2](https://huggingface.co/microsoft/phi-1_5/resolve/main/Research%20License.docx) / [Qwen](https://github.com/QwenLM/Qwen/blob/main/Tongyi%20Qianwen%20LICENSE%20AGREEMENT) / [XVERSE](https://github.com/xverse-ai/XVERSE-13B/blob/main/MODEL_LICENSE.pdf) / [Yi](https://huggingface.co/01-ai/Yi-6B/blob/main/LICENSE) / [Yuan](https://github.com/IEIT-Yuan/Yuan-2.0/blob/main/LICENSE-Yuan)

 ## 引用

--- a/assets/benchmark.svg
+++ b/assets/benchmark.svg
--- a/data/README.md
+++ b/data/README.md
@@ -2,11 +2,13 @@ If you are using a custom dataset, please provide your dataset definition in the

 ```json
 "dataset_name": {
-  "hf_hub_url": "the name of the dataset repository on the Hugging Face hub. (if specified, ignore below 3 arguments)",
-  "script_url": "the name of the directory containing a dataset loading script. (if specified, ignore below 2 arguments)",
-  "file_name": "the name of the dataset file in the this directory. (required if above are not specified)",
+  "hf_hub_url": "the name of the dataset repository on the Hugging Face hub. (if specified, ignore script_url and file_name)",
+  "ms_hub_url": "the name of the dataset repository on the ModelScope hub. (if specified, ignore script_url and file_name)",
+  "script_url": "the name of the directory containing a dataset loading script. (if specified, ignore file_name)",
+  "file_name": "the name of the dataset file in this directory. (required if above are not specified)",
  "file_sha1": "the SHA-1 hash value of the dataset file. (optional, does not affect training)",
  "subset": "the name of the subset. (optional, default: None)",
+  "folder": "the name of the folder of the dataset repository on the Hugging Face hub. (optional, default: None)",
  "ranking": "whether the dataset is a preference dataset or not. (default: false)",
  "formatting": "the format of the dataset. (optional, default: alpaca, can be chosen from {alpaca, sharegpt})",
  "columns": {
@@ -16,7 +18,8 @@ If you are using a custom dataset, please provide your dataset definition in the
    "history": "the column name in the dataset containing the histories. (default: None, for alpaca)",
    "messages": "the column name in the dataset containing the messages. (default: conversations, for sharegpt)",
    "role": "the key in the message represents the identity. (default: from, for sharegpt)",
-    "content": "the key in the message represents the content. (default: value, for sharegpt)"
+    "content": "the key in the message represents the content. (default: value, for sharegpt)",
+    "system": "the column name in the dataset containing the system prompts. (default: None, for both)"
  }
 }
 ```
@@ -31,6 +34,7 @@ Currently we support dataset in **alpaca** or **sharegpt** format, the dataset i
    "instruction": "user instruction (required)",
    "input": "user input (optional)",
    "output": "model response (required)",
+    "system": "system prompt (optional)",
    "history": [
      ["user instruction in the first round (optional)", "model response in the first round (optional)"],
      ["user instruction in the second round (optional)", "model response in the second round (optional)"]
@@ -47,6 +51,7 @@ Regarding the above dataset, the `columns` in `dataset_info.json` should be:
    "prompt": "instruction",
    "query": "input",
    "response": "output",
+    "system": "system",
    "history": "history"
  }
 }
@@ -54,7 +59,7 @@ Regarding the above dataset, the `columns` in `dataset_info.json` should be:

 where the `prompt` and `response` columns should contain non-empty values, represent instruction and response respectively. The `query` column will be concatenated with the `prompt` column and used as input for the model.

-The `history` column is a list consisting string tuples representing query-response pairs in history. Note that the responses **in each round will be used for training**.
+The `system` column will be used as the system prompt in the template. The `history` column is a list consisting string tuples representing query-response pairs in history. Note that the responses **in each round will be used for training**.

 For the pre-training datasets, only the `prompt` column will be used for training.

@@ -85,7 +90,8 @@ The dataset in sharegpt format should follow the below format:
        "from": "gpt",
        "value": "model response"
      }
-    ]
+    ],
+    "system": "system prompt (optional)"
  }
 ]
 ```
@@ -97,7 +103,8 @@ Regarding the above dataset, the `columns` in `dataset_info.json` should be:
  "columns": {
    "messages": "conversations",
    "role": "from",
-    "content": "value"
+    "content": "value",
+    "system": "system"
  }
 }
 ```
--- a/data/README_zh.md
+++ b/data/README_zh.md
@@ -2,11 +2,13 @@

 ```json
 "数据集名称": {
-  "hf_hub_url": "Hugging Face 上的项目地址（若指定，则忽略下列三个参数）",
-  "script_url": "包含数据加载脚本的本地文件夹名称（若指定，则忽略下列两个参数）",
+  "hf_hub_url": "Hugging Face 的数据集仓库地址（若指定，则忽略 script_url 和 file_name）",
+  "ms_hub_url": "ModelScope 的数据集仓库地址（若指定，则忽略 script_url 和 file_name）",
+  "script_url": "包含数据加载脚本的本地文件夹名称（若指定，则忽略 file_name）",
  "file_name": "该目录下数据集文件的名称（若上述参数未指定，则此项必需）",
  "file_sha1": "数据集文件的 SHA-1 哈希值（可选，留空不影响训练）",
  "subset": "数据集子集的名称（可选，默认：None）",
+  "folder": "Hugging Face 仓库的文件夹名称（可选，默认：None）",
  "ranking": "是否为偏好数据集（可选，默认：False）",
  "formatting": "数据集格式（可选，默认：alpaca，可以为 alpaca 或 sharegpt）",
  "columns": {
@@ -16,7 +18,8 @@
    "history": "数据集代表历史对话的表头名称（默认：None，用于 alpaca 格式）",
    "messages": "数据集代表消息列表的表头名称（默认：conversations，用于 sharegpt 格式）",
    "role": "消息中代表发送者身份的键名（默认：from，用于 sharegpt 格式）",
-    "content": "消息中代表文本内容的键名（默认：value，用于 sharegpt 格式）"
+    "content": "消息中代表文本内容的键名（默认：value，用于 sharegpt 格式）",
+    "system": "数据集代表系统提示的表头名称（默认：None，用于两种格式）"
  }
 }
 ```
@@ -31,6 +34,7 @@
    "instruction": "用户指令（必填）",
    "input": "用户输入（选填）",
    "output": "模型回答（必填）",
+    "system": "系统提示词（选填）",
    "history": [
      ["第一轮指令（选填）", "第一轮回答（选填）"],
      ["第二轮指令（选填）", "第二轮回答（选填）"]
@@ -47,6 +51,7 @@
    "prompt": "instruction",
    "query": "input",
    "response": "output",
+    "system": "system",
    "history": "history"
  }
 }
@@ -54,7 +59,7 @@

 其中 `prompt` 和 `response` 列应当是非空的字符串，分别代表用户指令和模型回答。`query` 列的内容将会和 `prompt` 列拼接作为模型输入。

-`history` 列是由多个字符串二元组构成的列表，分别代表历史消息中每轮的指令和回答。注意每轮的模型回答**均会被用于训练**。
+`system` 为模板中的系统提示词。`history` 列是由多个字符串二元组构成的列表，分别代表历史消息中每轮的指令和回答。注意每轮的模型回答**均会被用于训练**。

 对于预训练数据集，仅 `prompt` 列中的内容会用于模型训练。

@@ -85,7 +90,8 @@
        "from": "gpt",
        "value": "模型回答"
      }
-    ]
+    ],
+    "system": "系统提示词（选填）"
  }
 ]
 ```
@@ -97,7 +103,8 @@
  "columns": {
    "messages": "conversations",
    "role": "from",
-    "content": "value"
+    "content": "value",
+    "system": "system"
  }
 }
 ```
--- a/data/glaive_toolcall_10k.json.REMOVED.git-id
+++ b/data/glaive_toolcall_10k.json.REMOVED.git-id
@@ -0,0 +1 @@
+4748dff00d1dc42768a5b6cc772143c313017812
--- a/data/sharegpt_zh_27k.json.REMOVED.git-id
+++ b/data/sharegpt_zh_27k.json.REMOVED.git-id
@@ -1 +0,0 @@
-38c89869c6aeca2a3af9ea1e09afe460f9b46810
--- a/data/wiki_demo.txt
+++ b/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.
--- a/data/wiki_demo.txt.REMOVED.git-id
+++ b/data/wiki_demo.txt.REMOVED.git-id
@@ -0,0 +1 @@
+c9cf509b7fdac5490cfd6dae72c2d7b8a60af6cb
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,3 +1,37 @@
 [build-system]
 requires = ["setuptools>=61.0"]
 build-backend = "setuptools.build_meta"
+
+[tool.black]
+line-length = 119
+target-version = ["py38"]
+
+[tool.ruff]
+ignore = ["C408", "C901", "E501", "E731", "E741", "W605"]
+select = ["C", "E", "F", "I", "W"]
+line-length = 119
+
+[tool.ruff.isort]
+lines-after-imports = 2
+known-first-party = ["llmtuner"]
+
+[isort]
+default_section = "FIRSTPARTY"
+known_first_party = "llmtuner"
+known_third_party = [
+    "accelerate",
+    "datasets",
+    "gradio",
+    "numpy",
+    "peft",
+    "torch",
+    "transformers",
+    "trl"
+]
+line_length = 119
+lines_after_imports = 2
+multi_line_output = 3
+include_trailing_comma = true
+force_grid_wrap = 0
+use_parentheses = true
+ensure_newline_before_comments = true
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,14 +1,14 @@
 torch>=1.13.1
-transformers>=4.31.0,<4.35.0
-datasets>=2.14.0
+transformers>=4.36.2
+datasets>=2.14.3
 accelerate>=0.21.0
-peft>=0.6.0
-trl>=0.7.4
+peft>=0.7.0
+trl>=0.7.6
 gradio>=3.38.0,<4.0.0
 scipy
+einops
 sentencepiece
 protobuf
-tiktoken
 jieba
 rouge-chinese
 nltk
--- a/src/api_demo.py
+++ b/src/api_demo.py
@@ -1,3 +1,5 @@
+import os
+
 import uvicorn

 from llmtuner import ChatModel, create_app
@@ -6,8 +8,8 @@ from llmtuner import ChatModel, create_app
 def main():
    chat_model = ChatModel()
    app = create_app(chat_model)
-    print("Visit http://localhost:8000/docs for API document.")
-    uvicorn.run(app, host="0.0.0.0", port=8000, workers=1)
+    print("Visit http://localhost:{}/docs for API document.".format(os.environ.get("API_PORT", 8000)))
+    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("API_PORT", 8000)), workers=1)


 if __name__ == "__main__":
--- a/src/cli_demo.py
+++ b/src/cli_demo.py
@@ -1,17 +1,19 @@
 from llmtuner import ChatModel
 from llmtuner.extras.misc import torch_gc

+
 try:
    import platform
+
    if platform.system() != "Windows":
-        import readline
+        import readline  # noqa: F401
 except ImportError:
    print("Install `readline` for a better experience.")


 def main():
    chat_model = ChatModel()
-    history = []
+    messages = []
    print("Welcome to the CLI application, use `clear` to remove the history, use `exit` to exit the application.")

    while True:
@@ -27,20 +29,20 @@ def main():
            break

        if query.strip() == "clear":
-            history = []
+            messages = []
            torch_gc()
            print("History has been removed.")
            continue

+        messages.append({"role": "user", "content": query})
        print("Assistant: ", end="", flush=True)

        response = ""
-        for new_text in chat_model.stream_chat(query, history):
+        for new_text in chat_model.stream_chat(messages):
            print(new_text, end="", flush=True)
            response += new_text
        print()
-
-        history = history + [(query, response)]
+        messages.append({"role": "assistant", "content": response})


 if __name__ == "__main__":
--- a/src/llmtuner/init.py
+++ b/src/llmtuner/init.py
@@ -1,10 +1,11 @@
 # Level: api, webui > chat, eval, train > data, model > extras, hparams

-from llmtuner.api import create_app
-from llmtuner.chat import ChatModel
-from llmtuner.eval import Evaluator
-from llmtuner.train import export_model, run_exp
-from llmtuner.webui import create_ui, create_web_demo
+from .api import create_app
+from .chat import ChatModel
+from .eval import Evaluator
+from .train import export_model, run_exp
+from .webui import create_ui, create_web_demo


-__version__ = "0.3.0"
+__version__ = "0.5.0"
+__all__ = ["create_app", "ChatModel", "Evaluator", "export_model", "run_exp", "create_ui", "create_web_demo"]
--- a/src/llmtuner/api/init.py
+++ b/src/llmtuner/api/init.py
@@ -1 +1,4 @@
-from llmtuner.api.app import create_app
+from .app import create_app
+
+
+__all__ = ["create_app"]
--- a/src/llmtuner/api/app.py
+++ b/src/llmtuner/api/app.py
@@ -1,26 +1,31 @@
+import asyncio
 import json
-from typing import List, Tuple
-from pydantic import BaseModel
+import os
 from contextlib import asynccontextmanager
+from typing import Any, Dict, Sequence

-from llmtuner.api.protocol import (
-    Role,
-    Finish,
-    ModelCard,
-    ModelList,
-    ChatMessage,
-    DeltaMessage,
+from pydantic import BaseModel
+
+from ..chat import ChatModel
+from ..data import Role as DataRole
+from ..extras.misc import torch_gc
+from ..extras.packages import is_fastapi_availble, is_starlette_available, is_uvicorn_available
+from .protocol import (
+    ChatCompletionMessage,
    ChatCompletionRequest,
    ChatCompletionResponse,
-    ChatCompletionStreamResponse,
    ChatCompletionResponseChoice,
    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
+    ChatCompletionResponseUsage,
+    ChatCompletionStreamResponse,
+    Finish,
+    Function,
+    FunctionCall,
+    ModelCard,
+    ModelList,
+    Role,
+    ScoreEvaluationRequest,
+    ScoreEvaluationResponse,
 )


@@ -43,10 +48,17 @@ async def lifespan(app: "FastAPI"): # collects GPU memory
    torch_gc()


-def to_json(data: BaseModel) -> str:
+def dictify(data: "BaseModel") -> Dict[str, Any]:
+    try:  # pydantic v2
+        return data.model_dump(exclude_unset=True)
+    except AttributeError:  # pydantic v1
+        return data.dict(exclude_unset=True)
+
+
+def jsonify(data: "BaseModel") -> str:
    try:  # pydantic v2
        return json.dumps(data.model_dump(exclude_unset=True), ensure_ascii=False)
-    except: # pydantic v1
+    except AttributeError:  # pydantic v1
        return data.json(exclude_unset=True, ensure_ascii=False)


@@ -61,6 +73,8 @@ def create_app(chat_model: "ChatModel") -> "FastAPI":
        allow_headers=["*"],
    )

+    semaphore = asyncio.Semaphore(int(os.environ.get("MAX_CONCURRENT", 1)))
+
    @app.get("/v1/models", response_model=ModelList)
    async def list_models():
        model_card = ModelCard(id="gpt-3.5-turbo")
@@ -68,98 +82,145 @@ 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) == 0 or request.messages[-1].role != Role.USER:
-            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid request")
+        if not chat_model.can_generate:
+            raise HTTPException(status_code=status.HTTP_405_METHOD_NOT_ALLOWED, detail="Not allowed")

-        query = request.messages[-1].content
-        prev_messages = request.messages[:-1]
-        if len(prev_messages) and prev_messages[0].role == Role.SYSTEM:
-            system = prev_messages.pop(0).content
+        if len(request.messages) == 0 or request.messages[-1].role != Role.USER:
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid length")
+
+        messages = [dictify(message) for message in request.messages]
+        if len(messages) and messages[0]["role"] == Role.SYSTEM:
+            system = messages.pop(0)["content"]
        else:
            system = None

-        history = []
-        if len(prev_messages) % 2 == 0:
-            for i in range(0, len(prev_messages), 2):
-                if prev_messages[i].role == Role.USER and prev_messages[i+1].role == Role.ASSISTANT:
-                    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:
+        if len(messages) % 2 == 0:
            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Only supports u/a/u/a/u...")

+        for i in range(len(messages)):
+            if i % 2 == 0 and messages[i]["role"] not in [Role.USER, Role.TOOL]:
+                raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid role")
+            elif i % 2 == 1 and messages[i]["role"] not in [Role.ASSISTANT, Role.FUNCTION]:
+                raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid role")
+            elif messages[i]["role"] == Role.TOOL:
+                messages[i]["role"] = DataRole.OBSERVATION
+
+        tool_list = request.tools
+        if len(tool_list):
+            try:
+                tools = json.dumps([tool_list[0]["function"]], ensure_ascii=False)
+            except Exception:
+                raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid tools")
+        else:
+            tools = ""
+
+        async with semaphore:
+            loop = asyncio.get_running_loop()
+            return await loop.run_in_executor(None, chat_completion, messages, system, tools, request)
+
+    def chat_completion(messages: Sequence[Dict[str, str]], system: str, tools: str, request: ChatCompletionRequest):
        if request.stream:
-            generate = predict(query, history, system, request)
+            generate = stream_chat_completion(messages, system, tools, request)
            return EventSourceResponse(generate, media_type="text/event-stream")

        responses = chat_model.chat(
-            query, history, system,
+            messages,
+            system,
+            tools,
            do_sample=request.do_sample,
            temperature=request.temperature,
            top_p=request.top_p,
            max_new_tokens=request.max_tokens,
-            num_return_sequences=request.n
+            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),
+            if tools:
+                result = chat_model.template.format_tools.extract(response.response_text)
+            else:
+                result = response.response_text
+
+            if isinstance(result, tuple):
+                name, arguments = result
+                function = Function(name=name, arguments=arguments)
+                response_message = ChatCompletionMessage(
+                    role=Role.ASSISTANT, tool_calls=[FunctionCall(function=function)]
+                )
+                finish_reason = Finish.TOOL
+            else:
+                response_message = ChatCompletionMessage(role=Role.ASSISTANT, content=result)
                finish_reason = Finish.STOP if response.finish_reason == "stop" else Finish.LENGTH
-            ))
+
+            choices.append(
+                ChatCompletionResponseChoice(index=i, message=response_message, finish_reason=finish_reason)
+            )
            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
+            total_tokens=prompt_length + response_length,
        )

        return ChatCompletionResponse(model=request.model, choices=choices, usage=usage)

-    async def predict(query: str, history: List[Tuple[str, str]], system: str, request: ChatCompletionRequest):
+    def stream_chat_completion(
+        messages: Sequence[Dict[str, str]], system: str, tools: str, request: ChatCompletionRequest
+    ):
        choice_data = ChatCompletionResponseStreamChoice(
-            index=0,
-            delta=DeltaMessage(role=Role.ASSISTANT),
-            finish_reason=None
+            index=0, delta=ChatCompletionMessage(role=Role.ASSISTANT, content=""), finish_reason=None
        )
        chunk = ChatCompletionStreamResponse(model=request.model, choices=[choice_data])
-        yield to_json(chunk)
+        yield jsonify(chunk)

        for new_text in chat_model.stream_chat(
-            query, history, system,
+            messages,
+            system,
+            tools,
            do_sample=request.do_sample,
            temperature=request.temperature,
            top_p=request.top_p,
-            max_new_tokens=request.max_tokens
+            max_new_tokens=request.max_tokens,
        ):
            if len(new_text) == 0:
                continue

            choice_data = ChatCompletionResponseStreamChoice(
-                index=0,
-                delta=DeltaMessage(content=new_text),
-                finish_reason=None
+                index=0, delta=ChatCompletionMessage(content=new_text), finish_reason=None
            )
            chunk = ChatCompletionStreamResponse(model=request.model, choices=[choice_data])
-            yield to_json(chunk)
+            yield jsonify(chunk)

        choice_data = ChatCompletionResponseStreamChoice(
-            index=0,
-            delta=DeltaMessage(),
-            finish_reason=Finish.STOP
+            index=0, delta=ChatCompletionMessage(), finish_reason=Finish.STOP
        )
        chunk = ChatCompletionStreamResponse(model=request.model, choices=[choice_data])
-        yield to_json(chunk)
+        yield jsonify(chunk)
        yield "[DONE]"

+    @app.post("/v1/score/evaluation", response_model=ScoreEvaluationResponse, status_code=status.HTTP_200_OK)
+    async def create_score_evaluation(request: ScoreEvaluationRequest):
+        if chat_model.can_generate:
+            raise HTTPException(status_code=status.HTTP_405_METHOD_NOT_ALLOWED, detail="Not allowed")
+
+        if len(request.messages) == 0:
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid request")
+
+        async with semaphore:
+            loop = asyncio.get_running_loop()
+            return await loop.run_in_executor(None, get_score, request)
+
+    def get_score(request: ScoreEvaluationRequest):
+        scores = chat_model.get_scores(request.messages, max_length=request.max_length)
+        return ScoreEvaluationResponse(model=request.model, scores=scores)
+
    return app


 if __name__ == "__main__":
    chat_model = ChatModel()
    app = create_app(chat_model)
-    uvicorn.run(app, host="0.0.0.0", port=8000, workers=1)
+    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("API_PORT", 8000)), workers=1)
--- a/src/llmtuner/api/protocol.py
+++ b/src/llmtuner/api/protocol.py
@@ -1,30 +1,48 @@
 import time
-from enum import Enum
-from pydantic import BaseModel, Field
+from enum import Enum, unique
 from typing import List, Optional

+from pydantic import BaseModel, Field
+from typing_extensions import Literal

+
+@unique
 class Role(str, Enum):
    USER = "user"
    ASSISTANT = "assistant"
    SYSTEM = "system"
+    FUNCTION = "function"
+    TOOL = "tool"


+@unique
 class Finish(str, Enum):
    STOP = "stop"
    LENGTH = "length"
+    TOOL = "tool_calls"


 class ModelCard(BaseModel):
    id: str
-    object: Optional[str] = "model"
-    created: Optional[int] = Field(default_factory=lambda: int(time.time()))
-    owned_by: Optional[str] = "owner"
+    object: Literal["model"] = "model"
+    created: int = Field(default_factory=lambda: int(time.time()))
+    owned_by: Literal["owner"] = "owner"


 class ModelList(BaseModel):
-    object: Optional[str] = "list"
-    data: Optional[List[ModelCard]] = []
+    object: Literal["list"] = "list"
+    data: List[ModelCard] = []
+
+
+class Function(BaseModel):
+    name: str
+    arguments: str
+
+
+class FunctionCall(BaseModel):
+    id: Literal["call_default"] = "call_default"
+    type: Literal["function"] = "function"
+    function: Function


 class ChatMessage(BaseModel):
@@ -32,31 +50,33 @@ class ChatMessage(BaseModel):
    content: str


-class DeltaMessage(BaseModel):
+class ChatCompletionMessage(BaseModel):
    role: Optional[Role] = None
    content: Optional[str] = None
+    tool_calls: Optional[List[FunctionCall]] = None


 class ChatCompletionRequest(BaseModel):
    model: str
    messages: List[ChatMessage]
-    do_sample: Optional[bool] = True
+    tools: Optional[list] = []
+    do_sample: bool = True
    temperature: Optional[float] = None
    top_p: Optional[float] = None
-    n: Optional[int] = 1
+    n: int = 1
    max_tokens: Optional[int] = None
-    stream: Optional[bool] = False
+    stream: bool = False


 class ChatCompletionResponseChoice(BaseModel):
    index: int
-    message: ChatMessage
+    message: ChatCompletionMessage
    finish_reason: Finish


 class ChatCompletionResponseStreamChoice(BaseModel):
    index: int
-    delta: DeltaMessage
+    delta: ChatCompletionMessage
    finish_reason: Optional[Finish] = None


@@ -67,17 +87,30 @@ class ChatCompletionResponseUsage(BaseModel):


 class ChatCompletionResponse(BaseModel):
-    id: Optional[str] = "chatcmpl-default"
-    object: Optional[str] = "chat.completion"
-    created: Optional[int] = Field(default_factory=lambda: int(time.time()))
+    id: Literal["chatcmpl-default"] = "chatcmpl-default"
+    object: Literal["chat.completion"] = "chat.completion"
+    created: int = Field(default_factory=lambda: int(time.time()))
    model: str
    choices: List[ChatCompletionResponseChoice]
    usage: ChatCompletionResponseUsage


 class ChatCompletionStreamResponse(BaseModel):
-    id: Optional[str] = "chatcmpl-default"
-    object: Optional[str] = "chat.completion.chunk"
-    created: Optional[int] = Field(default_factory=lambda: int(time.time()))
+    id: Literal["chatcmpl-default"] = "chatcmpl-default"
+    object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
+    created: int = Field(default_factory=lambda: int(time.time()))
    model: str
    choices: List[ChatCompletionResponseStreamChoice]
+
+
+class ScoreEvaluationRequest(BaseModel):
+    model: str
+    messages: List[str]
+    max_length: Optional[int] = None
+
+
+class ScoreEvaluationResponse(BaseModel):
+    id: Literal["scoreeval-default"] = "scoreeval-default"
+    object: Literal["score.evaluation"] = "score.evaluation"
+    model: str
+    scores: List[float]
--- a/src/llmtuner/chat/init.py
+++ b/src/llmtuner/chat/init.py
@@ -1 +1,4 @@
-from llmtuner.chat.chat_model import ChatModel
+from .chat_model import ChatModel
+
+
+__all__ = ["ChatModel"]
--- a/src/llmtuner/chat/chat_model.py
+++ b/src/llmtuner/chat/chat_model.py
@@ -1,17 +1,18 @@
-import torch
 from dataclasses import dataclass
-from typing import Any, Dict, Generator, List, Literal, Optional, Tuple
 from threading import Thread
+from typing import Any, Dict, Generator, List, Literal, Optional, Sequence, Tuple
+
+import torch
 from transformers import GenerationConfig, TextIteratorStreamer

-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
+from ..data import get_template_and_fix_tokenizer
+from ..extras.misc import get_logits_processor
+from ..hparams import get_infer_args
+from ..model import dispatch_model, load_model_and_tokenizer


@dataclass
 class Response:
-
    response_text: str
    response_length: int
    prompt_length: int
@@ -19,25 +20,26 @@ class Response:


 class ChatModel:
-
    def __init__(self, args: Optional[Dict[str, Any]] = None) -> None:
        model_args, data_args, finetuning_args, self.generating_args = get_infer_args(args)
-        self.model, self.tokenizer = load_model_and_tokenizer(model_args, finetuning_args)
-        self.tokenizer.padding_side = "left"
+        self.can_generate = finetuning_args.stage == "sft"
+        self.model, self.tokenizer = load_model_and_tokenizer(
+            model_args, finetuning_args, is_trainable=False, add_valuehead=(not self.can_generate)
+        )
+        self.tokenizer.padding_side = "left" if self.can_generate else "right"
        self.model = dispatch_model(self.model)
        self.template = get_template_and_fix_tokenizer(data_args.template, self.tokenizer)
-        self.system_prompt = data_args.system_prompt

    def _process_args(
        self,
-        query: str,
-        history: Optional[List[Tuple[str, str]]] = None,
+        messages: Sequence[Dict[str, str]],
        system: Optional[str] = None,
-        **input_kwargs
+        tools: Optional[str] = None,
+        **input_kwargs,
    ) -> Tuple[Dict[str, Any], int]:
-        system = system or self.system_prompt
+        paired_messages = messages + [{"role": "assistant", "content": ""}]
        prompt, _ = self.template.encode_oneturn(
-            tokenizer=self.tokenizer, query=query, resp="", history=history, system=system
+            tokenizer=self.tokenizer, messages=paired_messages, system=system, tools=tools
        )
        prompt_length = len(prompt)
        input_ids = torch.tensor([prompt], device=self.model.device)
@@ -52,7 +54,8 @@ class ChatModel:
        max_new_tokens = input_kwargs.pop("max_new_tokens", None)

        generating_args = self.generating_args.to_dict()
-        generating_args.update(dict(
+        generating_args.update(
+            dict(
                do_sample=do_sample if do_sample is not None else generating_args["do_sample"],
                temperature=temperature or generating_args["temperature"],
                top_p=top_p or generating_args["top_p"],
@@ -60,8 +63,9 @@ class ChatModel:
                num_return_sequences=num_return_sequences or 1,
                repetition_penalty=repetition_penalty or generating_args["repetition_penalty"],
                eos_token_id=[self.tokenizer.eos_token_id] + self.tokenizer.additional_special_tokens_ids,
-            pad_token_id=self.tokenizer.pad_token_id
-        ))
+                pad_token_id=self.tokenizer.pad_token_id,
+            )
+        )

        if isinstance(num_return_sequences, int) and num_return_sequences > 1:
            generating_args["do_sample"] = True
@@ -77,7 +81,7 @@ class ChatModel:
        gen_kwargs = dict(
            inputs=input_ids,
            generation_config=GenerationConfig(**generating_args),
-            logits_processor=get_logits_processor()
+            logits_processor=get_logits_processor(),
        )

        return gen_kwargs, prompt_length
@@ -85,17 +89,12 @@ class ChatModel:
    @torch.inference_mode()
    def chat(
        self,
-        query: str,
-        history: Optional[List[Tuple[str, str]]] = None,
+        messages: Sequence[Dict[str, str]],
        system: Optional[str] = None,
-        **input_kwargs
+        tools: Optional[str] = None,
+        **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)
+        gen_kwargs, prompt_length = self._process_args(messages, system, tools, **input_kwargs)
        generate_output = self.model.generate(**gen_kwargs)
        response_ids = generate_output[:, prompt_length:]
        response = self.tokenizer.batch_decode(
@@ -105,24 +104,26 @@ class ChatModel:
        for i in range(len(response)):
            eos_index = (response_ids[i] == self.tokenizer.eos_token_id).nonzero()
            response_length = (eos_index[0].item() + 1) if len(eos_index) else len(response_ids[i])
-            results.append(Response(
+            results.append(
+                Response(
                    response_text=response[i],
                    response_length=response_length,
                    prompt_length=prompt_length,
-                finish_reason="stop" if len(eos_index) else "length"
-            ))
+                    finish_reason="stop" if len(eos_index) else "length",
+                )
+            )

        return results

    @torch.inference_mode()
    def stream_chat(
        self,
-        query: str,
-        history: Optional[List[Tuple[str, str]]] = None,
+        messages: Sequence[Dict[str, str]],
        system: Optional[str] = None,
-        **input_kwargs
+        tools: Optional[str] = None,
+        **input_kwargs,
    ) -> Generator[str, None, None]:
-        gen_kwargs, _ = self._process_args(query, history, system, **input_kwargs)
+        gen_kwargs, _ = self._process_args(messages, system, tools, **input_kwargs)
        streamer = TextIteratorStreamer(self.tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
        gen_kwargs["streamer"] = streamer

@@ -130,3 +131,31 @@ class ChatModel:
        thread.start()

        yield from streamer
+
+    @torch.inference_mode()
+    def get_scores(self, batch_input: List[str], **input_kwargs) -> List[float]:
+        max_length = input_kwargs.pop("max_length", None)
+        device = getattr(self.model.pretrained_model, "device", "cuda")
+
+        inputs = self.tokenizer(
+            batch_input,
+            padding=True,
+            truncation=True,
+            max_length=max_length or getattr(self.model.config, "max_position_embeddings", 1024),
+            return_tensors="pt",
+            add_special_tokens=True,
+        ).to(device)
+
+        input_ids: torch.Tensor = inputs["input_ids"]
+        _, _, values = self.model(**inputs, output_hidden_states=True, return_dict=True)
+
+        if getattr(self.model.config, "model_type", None) == "chatglm":
+            values = torch.transpose(values, 0, 1)
+
+        scores = []
+        for i in range(input_ids.size(0)):
+            end_indexes = (input_ids[i] != self.tokenizer.pad_token_id).nonzero()
+            end_index = end_indexes[-1].item() if len(end_indexes) else 0
+            scores.append(values[i, end_index].nan_to_num().item())
+
+        return scores
--- a/src/llmtuner/data/init.py
+++ b/src/llmtuner/data/init.py
@@ -1,4 +1,6 @@
-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
+from .loader import get_dataset
+from .template import get_template_and_fix_tokenizer, templates
+from .utils import Role, split_dataset
+
+
+__all__ = ["get_dataset", "get_template_and_fix_tokenizer", "templates", "Role", "split_dataset"]
--- a/src/llmtuner/data/aligner.py
+++ b/src/llmtuner/data/aligner.py
@@ -0,0 +1,108 @@
+from functools import partial
+from typing import TYPE_CHECKING, Any, Dict, List, Union
+
+from .utils import Role
+
+
+if TYPE_CHECKING:
+    from datasets import Dataset, IterableDataset
+
+    from ..hparams import DataArguments
+    from .parser import DatasetAttr
+
+
+def convert_alpaca(examples: Dict[str, List[Any]], dataset_attr: "DatasetAttr") -> Dict[str, List[Any]]:
+    outputs = {"prompt": [], "response": [], "system": [], "tools": []}
+    for i in range(len(examples[dataset_attr.prompt])):
+        prompt = []
+        if dataset_attr.history:
+            for old_prompt, old_response in examples[dataset_attr.history][i]:
+                prompt.append({"role": Role.USER, "content": old_prompt})
+                prompt.append({"role": Role.ASSISTANT, "content": old_response})
+
+        instruction = examples[dataset_attr.prompt][i]
+        if dataset_attr.query and examples[dataset_attr.query][i]:
+            instruction += "\n" + examples[dataset_attr.query][i]
+        prompt.append({"role": Role.USER, "content": instruction})
+
+        if dataset_attr.response:
+            if isinstance(examples[dataset_attr.response][i], list):
+                response = [
+                    {"role": Role.ASSISTANT, "content": content} for content in examples[dataset_attr.response][i]
+                ]
+            else:
+                response = [{"role": Role.ASSISTANT, "content": examples[dataset_attr.response][i]}]
+        else:
+            response = []
+
+        outputs["prompt"].append(prompt)
+        outputs["response"].append(response)
+        outputs["system"].append(examples[dataset_attr.system][i] if dataset_attr.system else "")
+        outputs["tools"].append("")
+
+    return outputs
+
+
+def convert_sharegpt(examples: Dict[str, List[Any]], dataset_attr: "DatasetAttr") -> Dict[str, List[Any]]:
+    outputs = {"prompt": [], "response": [], "system": [], "tools": []}
+    tag_mapping = {
+        dataset_attr.user_tag: Role.USER,
+        dataset_attr.assistant_tag: Role.ASSISTANT,
+        dataset_attr.observation_tag: Role.OBSERVATION,
+        dataset_attr.function_tag: Role.FUNCTION,
+    }
+    for i, messages in enumerate(examples[dataset_attr.messages]):
+        messages = messages[: len(messages) // 2 * 2]  # should be multiples of 2
+        if len(messages) == 0:
+            continue
+
+        prompt = []
+        response = []
+        for turn_idx, message in enumerate(messages):
+            if turn_idx % 2 == 0:
+                accept_tags = [dataset_attr.user_tag, dataset_attr.observation_tag]
+            else:
+                accept_tags = [dataset_attr.assistant_tag, dataset_attr.function_tag]
+
+            if message[dataset_attr.role_tag] not in accept_tags:
+                raise ValueError("Invalid role tag in {}.".format(messages))
+
+            prompt.append(
+                {"role": tag_mapping[message[dataset_attr.role_tag]], "content": message[dataset_attr.content_tag]}
+            )
+
+        last_message = prompt.pop(-1)
+        response.append(last_message)
+        outputs["prompt"].append(prompt)
+        outputs["response"].append(response)
+        outputs["system"].append(examples[dataset_attr.system][i] if dataset_attr.system else "")
+        outputs["tools"].append(examples[dataset_attr.tools][i] if dataset_attr.tools else "")
+
+    return outputs
+
+
+def align_dataset(
+    dataset: Union["Dataset", "IterableDataset"], dataset_attr: "DatasetAttr", data_args: "DataArguments"
+) -> Union["Dataset", "IterableDataset"]:
+    r"""
+    Aligned dataset:
+        prompt: [{"role": "user", "content": "..."}]
+        response: [{"role": "assistant", "content": "..."}]
+        system: "..."
+        tools: "..."
+    """
+    if dataset_attr.formatting == "alpaca":
+        convert_func = partial(convert_alpaca, dataset_attr=dataset_attr)
+    else:
+        convert_func = partial(convert_sharegpt, dataset_attr=dataset_attr)
+
+    column_names = list(next(iter(dataset)).keys())
+    kwargs = {}
+    if not data_args.streaming:
+        kwargs = dict(
+            num_proc=data_args.preprocessing_num_workers,
+            load_from_cache_file=(not data_args.overwrite_cache),
+            desc="Converting format of dataset",
+        )
+
+    return dataset.map(convert_func, batched=True, remove_columns=column_names, **kwargs)
--- a/src/llmtuner/data/formatter.py
+++ b/src/llmtuner/data/formatter.py
@@ -0,0 +1,148 @@
+import json
+import re
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Literal, Sequence, Set, Tuple, Union
+
+
+SLOTS = Sequence[Union[str, Set[str], Dict[str, str]]]
+
+
+JSON_FORMAT_PROMPT = (
+    """, in a JSON format representing the kwargs (e.g. ```{"input": "hello world", "num_beams": 5}```)"""
+)
+
+
+TOOL_SYSTEM_PROMPT = (
+    "You have access to the following tools:\n{tool_text}"
+    "Use the following format to answer the question:\n"
+    "```\n"
+    "Action: the action to take, should be one of [{tool_names}] if using a tool.\n"
+    "Action Input: the input to the action{format_prompt}.\n"
+    "```"
+)
+
+
+def default_tool_formatter(tools: List[Dict[str, Any]]) -> str:
+    tool_text = ""
+    tool_names = []
+    for tool in tools:
+        param_text = ""
+        for name, param in tool["parameters"]["properties"].items():
+            required = ", required" if name in tool["parameters"].get("required", []) else ""
+            enum = ", should be one of [{}]".format(", ".join(param["enum"])) if param.get("enum", None) else ""
+            param_text += "  - {name} ({type}{required}): {desc}{enum}\n".format(
+                name=name,
+                type=param.get("type", ""),
+                required=required,
+                desc=param.get("description", ""),
+                enum=enum,
+            )
+
+        tool_text += "> Tool Name: {name}\nTool Description: {desc}\nTool Args:\n{args}\n".format(
+            name=tool["name"], desc=tool.get("description", ""), args=param_text
+        )
+        tool_names.append(tool["name"])
+
+    return TOOL_SYSTEM_PROMPT.format(
+        tool_text=tool_text, tool_names=", ".join(tool_names), format_prompt=JSON_FORMAT_PROMPT
+    )
+
+
+def default_tool_extractor(content: str) -> Union[str, Tuple[str, str]]:
+    regex = re.compile(r"Action:\s*([a-zA-Z0-9_]+).*?Action Input:\s*(.*)", re.DOTALL)
+    action_match = re.search(regex, content)
+    if not action_match:
+        return content
+
+    tool_name = action_match.group(1).strip()
+    tool_input = action_match.group(2).strip().strip('"').strip("```")
+    try:
+        arguments = json.loads(tool_input)
+    except json.JSONDecodeError:
+        return content
+
+    return tool_name, json.dumps(arguments, ensure_ascii=False)
+
+
+@dataclass
+class Formatter(ABC):
+    slots: SLOTS = field(default_factory=list)
+    tool_format: Literal["default"] = "default"
+
+    @abstractmethod
+    def apply(self, **kwargs) -> SLOTS:
+        ...
+
+    def extract(self, content: str) -> Union[str, Tuple[str, str]]:
+        raise NotImplementedError
+
+
+@dataclass
+class EmptyFormatter(Formatter):
+    def apply(self, **kwargs) -> SLOTS:
+        return self.slots
+
+
+@dataclass
+class StringFormatter(Formatter):
+    def apply(self, **kwargs) -> SLOTS:
+        elements = []
+        for slot in self.slots:
+            if isinstance(slot, str):
+                for name, value in kwargs.items():
+                    slot = slot.replace("{{" + name + "}}", value, 1)
+                elements.append(slot)
+            elif isinstance(slot, (dict, set)):
+                elements.append(slot)
+            else:
+                raise ValueError("Input must be string, set[str] or dict[str, str], got {}".format(type(slot)))
+
+        return elements
+
+
+@dataclass
+class FunctionFormatter(Formatter):
+    def apply(self, **kwargs) -> SLOTS:
+        content = kwargs.pop("content")
+        try:
+            function = json.loads(content)
+            name = function["name"]
+            arguments = json.dumps(function["arguments"], ensure_ascii=False)
+        except Exception:
+            name, arguments = "", ""
+
+        elements = []
+        for slot in self.slots:
+            if isinstance(slot, str):
+                slot = slot.replace("{{name}}", name).replace("{{arguments}}", arguments)
+                elements.append(slot)
+            elif isinstance(slot, (dict, set)):
+                elements.append(slot)
+            else:
+                raise ValueError("Input must be string, set[str] or dict[str, str], got {}".format(type(slot)))
+
+        return elements
+
+
+@dataclass
+class ToolFormatter(Formatter):
+    def apply(self, **kwargs) -> SLOTS:
+        content = kwargs.pop("content")
+        try:
+            tools = json.loads(content)
+            if not len(tools):
+                return [""]
+
+            if self.tool_format == "default":
+                return [default_tool_formatter(tools)]
+            else:
+                raise NotImplementedError
+        except Exception:
+            return [""]
+
+    def extract(self, content: str) -> Union[str, Tuple[str, str]]:
+        if self.tool_format == "default":
+            return default_tool_extractor(content)
+        else:
+            raise NotImplementedError
--- a/src/llmtuner/data/loader.py
+++ b/src/llmtuner/data/loader.py
@@ -1,132 +1,121 @@
+import inspect
 import os
-from typing import TYPE_CHECKING, Any, Dict, List, Union
+from typing import TYPE_CHECKING, List, Literal, Union

-from datasets import concatenate_datasets, interleave_datasets, load_dataset
+from datasets import concatenate_datasets, interleave_datasets, load_dataset, load_from_disk
+
+from ..extras.constants import FILEEXT2TYPE
+from ..extras.logging import get_logger
+from .aligner import align_dataset
+from .parser import get_dataset_list
+from .preprocess import get_preprocess_and_print_func
+from .template import get_template_and_fix_tokenizer
+from .utils import checksum

-from llmtuner.data.utils import checksum, EXT2TYPE
-from llmtuner.extras.logging import get_logger

 if TYPE_CHECKING:
    from datasets import Dataset, IterableDataset
-    from llmtuner.hparams import ModelArguments, DataArguments
+    from transformers import Seq2SeqTrainingArguments
+    from transformers.tokenization_utils import PreTrainedTokenizer
+
+    from ..hparams import DataArguments, ModelArguments
+    from .parser import DatasetAttr


 logger = get_logger(__name__)


-def get_dataset(
+def load_single_dataset(
+    dataset_attr: "DatasetAttr",
    model_args: "ModelArguments",
-    data_args: "DataArguments"
-) -> Union["Dataset", "IterableDataset"]:
-    max_samples = data_args.max_samples
-    all_datasets: List[Union["Dataset", "IterableDataset"]] = [] # support multiple datasets
-
-    for dataset_attr in data_args.dataset_list:
-        logger.info("Loading dataset {}...".format(dataset_attr))
-
-        if dataset_attr.load_from == "hf_hub":
+    data_args: "DataArguments",
+):
+    data_path, data_name, data_dir, data_files = None, None, None, None
+    if dataset_attr.load_from in ["hf_hub", "ms_hub"]:
        data_path = dataset_attr.dataset_name
        data_name = dataset_attr.subset
-            data_files = None
+        data_dir = dataset_attr.folder
+
    elif dataset_attr.load_from == "script":
        data_path = os.path.join(data_args.dataset_dir, dataset_attr.dataset_name)
        data_name = dataset_attr.subset
-            data_files = None
+        data_dir = dataset_attr.folder
+
    elif dataset_attr.load_from == "file":
-            data_path, data_name = None, None
-            data_files: List[str] = []
-            if os.path.isdir(os.path.join(data_args.dataset_dir, dataset_attr.dataset_name)): # is directory
-                for file_name in os.listdir(os.path.join(data_args.dataset_dir, dataset_attr.dataset_name)):
-                    data_files.append(os.path.join(data_args.dataset_dir, dataset_attr.dataset_name, file_name))
+        data_files = []
+        local_path: str = os.path.join(data_args.dataset_dir, dataset_attr.dataset_name)
+        if os.path.isdir(local_path):  # is directory
+            for file_name in os.listdir(local_path):
+                data_files.append(os.path.join(local_path, file_name))
                if data_path is None:
-                        data_path = EXT2TYPE.get(file_name.split(".")[-1], None)
-                    else:
-                        assert data_path == EXT2TYPE.get(file_name.split(".")[-1], None), "file types are not identical."
-            elif os.path.isfile(os.path.join(data_args.dataset_dir, dataset_attr.dataset_name)): # is file
-                data_files.append(os.path.join(data_args.dataset_dir, dataset_attr.dataset_name))
-                data_path = EXT2TYPE.get(dataset_attr.dataset_name.split(".")[-1], None)
+                    data_path = FILEEXT2TYPE.get(file_name.split(".")[-1], None)
+                elif data_path != FILEEXT2TYPE.get(file_name.split(".")[-1], None):
+                    raise ValueError("File types should be identical.")
+        elif os.path.isfile(local_path):  # is file
+            data_files.append(local_path)
+            data_path = FILEEXT2TYPE.get(local_path.split(".")[-1], None)
        else:
            raise ValueError("File not found.")

-            assert data_path, "File extension must be txt, csv, json or jsonl."
+        if data_path is None:
+            raise ValueError("File extension must be txt, csv, json or jsonl.")
+
        checksum(data_files, dataset_attr.dataset_sha1)
    else:
        raise NotImplementedError

+    if dataset_attr.load_from == "ms_hub":
+        try:
+            from modelscope import MsDataset
+            from modelscope.utils.config_ds import MS_DATASETS_CACHE
+
+            cache_dir = model_args.cache_dir or MS_DATASETS_CACHE
+            dataset = MsDataset.load(
+                dataset_name=data_path,
+                subset_name=data_name,
+                data_dir=data_dir,
+                data_files=data_files,
+                split=data_args.split,
+                cache_dir=cache_dir,
+                token=model_args.ms_hub_token,
+                use_streaming=(data_args.streaming and (dataset_attr.load_from != "file")),
+            ).to_hf_dataset()
+        except ImportError:
+            raise ImportError("Please install modelscope via `pip install modelscope -U`")
+    else:
+        if "trust_remote_code" in inspect.signature(load_dataset).parameters:  # for datasets==2.16.0
+            kwargs = {"trust_remote_code": True}
+        else:
+            kwargs = {}
+
        dataset = load_dataset(
            path=data_path,
            name=data_name,
+            data_dir=data_dir,
            data_files=data_files,
            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")),
+            **kwargs,
        )

-        if max_samples is not None: # truncate dataset
-            dataset = dataset.select(range(min(len(dataset), max_samples)))
+    if data_args.streaming and (dataset_attr.load_from == "file"):  # faster than specifying streaming=True
+        dataset = dataset.to_iterable_dataset()  # TODO: add num shards parameter

-        def convert_format(examples: Dict[str, List[Any]]) -> Dict[str, List[Any]]:
-            # convert dataset from sharegpt format to alpaca format
-            outputs = {"prompt": [], "query": [], "response": [], "history": []}
-            for msg_list in examples[dataset_attr.messages]:
-                msg_list = msg_list[:len(msg_list) // 2 * 2] # should be multiples of 2
-                if len(msg_list) == 0:
-                    continue
+    if data_args.max_samples is not None:  # truncate dataset
+        num_samples = min(data_args.max_samples, len(dataset))
+        dataset = dataset.select(range(num_samples))

-                msg_pairs = []
-                user_role, assistant_role = None, None
-                for idx in range(0, len(msg_list), 2):
-                    if user_role is None and assistant_role is None:
-                        user_role = msg_list[idx][dataset_attr.role]
-                        assistant_role = msg_list[idx + 1][dataset_attr.role]
-                    else:
-                        if (
-                            msg_list[idx][dataset_attr.role] != user_role
-                            or msg_list[idx+1][dataset_attr.role] != assistant_role
-                        ):
-                            raise ValueError("Only accepts conversation in u/a/u/a/u/a order.")
-                    msg_pairs.append((msg_list[idx][dataset_attr.content], msg_list[idx + 1][dataset_attr.content]))
+    return align_dataset(dataset, dataset_attr, data_args)

-                if len(msg_pairs) != 0:
-                    outputs["prompt"].append(msg_pairs[-1][0])
-                    outputs["query"].append("")
-                    outputs["response"].append(msg_pairs[-1][1])
-                    outputs["history"].append(msg_pairs[:-1])

-            return outputs
-
-        if dataset_attr.formatting == "sharegpt": # convert format
-            column_names = list(next(iter(dataset)).keys())
-            kwargs = {}
-            if not data_args.streaming:
-                kwargs = dict(
-                    num_proc=data_args.preprocessing_num_workers,
-                    load_from_cache_file=(not data_args.overwrite_cache),
-                    desc="Converting format of dataset"
-                )
-
-            dataset = dataset.map(
-                convert_format,
-                batched=True,
-                remove_columns=column_names,
-                **kwargs
-            )
-        else:
-            for column_name in ["prompt", "query", "response", "history"]: # align dataset
-                if getattr(dataset_attr, column_name) and getattr(dataset_attr, column_name) != column_name:
-                    dataset = dataset.rename_column(getattr(dataset_attr, column_name), column_name)
-
-        if dataset_attr.system_prompt: # add system prompt
-            system_prompt = dataset_attr.system_prompt
-            if data_args.streaming:
-                dataset = dataset.map(lambda _: {"system": system_prompt})
-            else:
-                dataset = dataset.add_column("system", [system_prompt] * len(dataset))
-
-        all_datasets.append(dataset)
-
-    if len(data_args.dataset_list) == 1:
+def merge_dataset(
+    all_datasets: List[Union["Dataset", "IterableDataset"]],
+    data_args: "DataArguments",
+    training_args: "Seq2SeqTrainingArguments",
+) -> Union["Dataset", "IterableDataset"]:
+    if len(all_datasets) == 1:
        return all_datasets[0]
    elif data_args.mix_strategy == "concat":
        if data_args.streaming:
@@ -138,8 +127,67 @@ def get_dataset(
        return interleave_datasets(
            datasets=all_datasets,
            probabilities=data_args.interleave_probs,
-            seed=data_args.seed,
-            stopping_strategy="first_exhausted" if data_args.mix_strategy.endswith("under") else "all_exhausted"
+            seed=training_args.seed,
+            stopping_strategy="first_exhausted" if data_args.mix_strategy.endswith("under") else "all_exhausted",
        )
    else:
        raise ValueError("Unknown mixing strategy.")
+
+
+def get_dataset(
+    tokenizer: "PreTrainedTokenizer",
+    model_args: "ModelArguments",
+    data_args: "DataArguments",
+    training_args: "Seq2SeqTrainingArguments",
+    stage: Literal["pt", "sft", "rm", "ppo"],
+    # split: Optional[str] = "train", # TODO: add split
+) -> Union["Dataset", "IterableDataset"]:
+    template = get_template_and_fix_tokenizer(data_args.template, tokenizer)
+    if data_args.train_on_prompt and template.efficient_eos:
+        raise ValueError("Current template does not support `train_on_prompt`.")
+
+    # Load from cache
+    if data_args.cache_path is not None:
+        if os.path.exists(data_args.cache_path):
+            logger.warning("Loading dataset from disk will ignore other data arguments.")
+            dataset = load_from_disk(data_args.cache_path)
+            if data_args.streaming:
+                dataset = dataset.to_iterable_dataset()
+            return dataset
+
+        if data_args.streaming:
+            raise ValueError("Turn off dataset streaming to save cache files.")
+
+    with training_args.main_process_first(desc="load dataset"):
+        all_datasets = []
+        for dataset_attr in get_dataset_list(data_args):  # TODO: add split
+            all_datasets.append(load_single_dataset(dataset_attr, model_args, data_args))
+        dataset = merge_dataset(all_datasets, data_args, training_args)
+
+    with training_args.main_process_first(desc="pre-process dataset"):
+        preprocess_func, print_function = get_preprocess_and_print_func(
+            tokenizer, template, data_args, training_args, stage
+        )
+        column_names = list(next(iter(dataset)).keys())
+        kwargs = {}
+        if not data_args.streaming:
+            kwargs = dict(
+                num_proc=data_args.preprocessing_num_workers,
+                load_from_cache_file=(not data_args.overwrite_cache),
+                desc="Running tokenizer on dataset",
+            )
+
+        dataset = dataset.map(preprocess_func, batched=True, remove_columns=column_names, **kwargs)
+
+        if data_args.cache_path is not None and not os.path.exists(data_args.cache_path):
+            if training_args.should_save:
+                dataset.save_to_disk(data_args.cache_path)
+                logger.info("Dataset cache saved at {}.".format(data_args.cache_path))
+
+        if training_args.should_log:
+            try:
+                print_function(next(iter(dataset)))
+            except StopIteration:
+                raise RuntimeError("Empty dataset!")
+
+        return dataset
--- a/src/llmtuner/data/parser.py
+++ b/src/llmtuner/data/parser.py
@@ -0,0 +1,103 @@
+import json
+import os
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, List, Literal, Optional
+
+from ..extras.constants import DATA_CONFIG
+from ..extras.misc import use_modelscope
+
+
+if TYPE_CHECKING:
+    from ..hparams import DataArguments
+
+
+@dataclass
+class DatasetAttr:
+    load_from: Literal["hf_hub", "ms_hub", "script", "file"]
+    dataset_name: Optional[str] = None
+    dataset_sha1: Optional[str] = None
+    subset: Optional[str] = None
+    folder: Optional[str] = None
+    ranking: Optional[bool] = False
+    formatting: Optional[Literal["alpaca", "sharegpt"]] = "alpaca"
+
+    system: Optional[str] = None
+
+    prompt: Optional[str] = "instruction"
+    query: Optional[str] = "input"
+    response: Optional[str] = "output"
+    history: Optional[str] = None
+
+    messages: Optional[str] = "conversations"
+    tools: Optional[str] = None
+
+    role_tag: Optional[str] = "from"
+    content_tag: Optional[str] = "value"
+    user_tag: Optional[str] = "human"
+    assistant_tag: Optional[str] = "gpt"
+    observation_tag: Optional[str] = "observation"
+    function_tag: Optional[str] = "function_call"
+
+    def __repr__(self) -> str:
+        return self.dataset_name
+
+
+def get_dataset_list(data_args: "DataArguments") -> List["DatasetAttr"]:
+    dataset_names = [ds.strip() for ds in data_args.dataset.split(",")] if data_args.dataset is not None else []
+    try:
+        with open(os.path.join(data_args.dataset_dir, DATA_CONFIG), "r") as f:
+            dataset_info = json.load(f)
+    except Exception as err:
+        if data_args.dataset is not None:
+            raise ValueError(
+                "Cannot open {} due to {}.".format(os.path.join(data_args.dataset_dir, DATA_CONFIG), str(err))
+            )
+        dataset_info = None
+
+    if data_args.interleave_probs is not None:
+        data_args.interleave_probs = [float(prob.strip()) for prob in data_args.interleave_probs.split(",")]
+
+    dataset_list: List[DatasetAttr] = []
+    for name in dataset_names:
+        if name not in dataset_info:
+            raise ValueError("Undefined dataset {} in {}.".format(name, DATA_CONFIG))
+
+        has_hf_url = "hf_hub_url" in dataset_info[name]
+        has_ms_url = "ms_hub_url" in dataset_info[name]
+
+        if has_hf_url or has_ms_url:
+            if (use_modelscope() and has_ms_url) or (not has_hf_url):
+                dataset_attr = DatasetAttr("ms_hub", dataset_name=dataset_info[name]["ms_hub_url"])
+            else:
+                dataset_attr = DatasetAttr("hf_hub", dataset_name=dataset_info[name]["hf_hub_url"])
+        elif "script_url" in dataset_info[name]:
+            dataset_attr = DatasetAttr("script", dataset_name=dataset_info[name]["script_url"])
+        else:
+            dataset_attr = DatasetAttr(
+                "file",
+                dataset_name=dataset_info[name]["file_name"],
+                dataset_sha1=dataset_info[name].get("file_sha1", None),
+            )
+
+        dataset_attr.subset = dataset_info[name].get("subset", None)
+        dataset_attr.folder = dataset_info[name].get("folder", None)
+        dataset_attr.ranking = dataset_info[name].get("ranking", False)
+        dataset_attr.formatting = dataset_info[name].get("formatting", "alpaca")
+
+        if "columns" in dataset_info[name]:
+            if dataset_attr.formatting == "alpaca":
+                column_names = ["prompt", "query", "response", "history"]
+            else:
+                column_names = ["messages", "tools"]
+
+            column_names += ["system"]
+            for column_name in column_names:
+                setattr(dataset_attr, column_name, dataset_info[name]["columns"].get(column_name, None))
+
+        if dataset_attr.formatting == "sharegpt" and "tags" in dataset_info[name]:
+            for tag in ["role_tag", "content_tag", "user_tag", "assistant_tag", "observation_tag", "function_tag"]:
+                setattr(dataset_attr, tag, dataset_info[name]["tags"].get(tag, None))
+
+        dataset_list.append(dataset_attr)
+
+    return dataset_list
--- a/src/llmtuner/data/preprocess.py
+++ b/src/llmtuner/data/preprocess.py
@@ -1,64 +1,32 @@
-import os
-import tiktoken
+from functools import partial
 from itertools import chain
-from typing import TYPE_CHECKING, Any, Dict, Generator, List, Literal, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Tuple

-from datasets import load_from_disk
+from ..extras.constants import IGNORE_INDEX
+from ..extras.logging import get_logger

-from llmtuner.data.template import get_template_and_fix_tokenizer
-from llmtuner.extras.constants import IGNORE_INDEX
-from llmtuner.extras.logging import get_logger

 if TYPE_CHECKING:
-    from datasets import Dataset, IterableDataset
    from transformers import Seq2SeqTrainingArguments
    from transformers.tokenization_utils import PreTrainedTokenizer
-    from llmtuner.hparams import DataArguments
+
+    from ..hparams import DataArguments
+    from .template import Template


 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",
-    data_args: "DataArguments",
-    training_args: "Seq2SeqTrainingArguments",
-    stage: Literal["pt", "sft", "rm", "ppo"]
-) -> Union["Dataset", "IterableDataset"]:
-    template = get_template_and_fix_tokenizer(data_args.template, tokenizer)
-
-    if data_args.train_on_prompt and template.efficient_eos:
-        raise ValueError("Current template does not support `train_on_prompt`.")
-
-    def preprocess_pretrain_dataset(examples: Dict[str, List[Any]]) -> Dict[str, List[List[int]]]:
+def preprocess_pretrain_dataset(
+    examples: Dict[str, List[Any]], tokenizer: "PreTrainedTokenizer", data_args: "DataArguments"
+) -> 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)
-            kwargs = dict(allowed_special="all")
-        else:
-            kwargs = dict(add_special_tokens=True)
+    text_examples = [examples["prompt"][i][0]["content"] for i in range(len(examples["prompt"]))]
+    tokenized_examples = tokenizer(text_examples, add_special_tokens=False)
+    for i in range(len(tokenized_examples["input_ids"])):
+        tokenized_examples["input_ids"][i] += [tokenizer.eos_token_id]
+        tokenized_examples["attention_mask"][i] += [1]

-        if hasattr(tokenizer, "add_eos_token"): # for LLaMA tokenizer
-            add_eos_token_flag = getattr(tokenizer, "add_eos_token")
-            setattr(tokenizer, "add_eos_token", True)
-
-        tokenized_examples = tokenizer(examples["prompt"], **kwargs)
    concatenated_examples = {k: list(chain(*tokenized_examples[k])) for k in tokenized_examples.keys()}
    total_length = len(concatenated_examples[list(concatenated_examples.keys())[0]])
    block_size = data_args.cutoff_len
@@ -69,31 +37,30 @@ def preprocess_dataset(
        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
        for k, t in concatenated_examples.items()
    }
-        # make sure the saved tokenizer is the same as the original one
-        if hasattr(tokenizer, "add_eos_token"):
-            setattr(tokenizer, "add_eos_token", add_eos_token_flag)
    return result

-    def preprocess_supervised_dataset(examples: Dict[str, List[Any]]) -> Dict[str, List[List[int]]]:
+
+def preprocess_supervised_dataset(
+    examples: Dict[str, List[Any]],
+    tokenizer: "PreTrainedTokenizer",
+    template: "Template",
+    data_args: "DataArguments",
+) -> Dict[str, List[List[int]]]:
    # build inputs with format `<bos> X Y <eos>` and labels with format `<ignore> ... <ignore> Y <eos>`
    # for multiturn examples, we only mask the prompt part in each prompt-response pair.
    model_inputs = {"input_ids": [], "attention_mask": [], "labels": []}

-        for query, response, history, system in construct_example(examples):
-            if not (isinstance(query, str) and isinstance(response, str) and query != "" and response != ""):
+    for i in range(len(examples["prompt"])):
+        if len(examples["prompt"][i]) == 0 or len(examples["response"][i]) != 1:
            continue

+        messages = examples["prompt"][i] + examples["response"][i]
        input_ids, labels = [], []
-            for turn_idx, (source_ids, target_ids) in enumerate(template.encode_multiturn(
-                tokenizer, query, response, history, system
-            )):
-                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 target_len > max_target_len:
-                    target_ids = target_ids[:max_target_len]
-
+        for turn_idx, (source_ids, target_ids) in enumerate(
+            template.encode_multiturn(
+                tokenizer, messages, examples["system"][i], examples["tools"][i], data_args.cutoff_len
+            )
+        ):
            if data_args.train_on_prompt:
                source_mask = source_ids
            elif turn_idx != 0 and template.efficient_eos:
@@ -108,34 +75,38 @@ def preprocess_dataset(
            input_ids += [tokenizer.eos_token_id]
            labels += [tokenizer.eos_token_id]

-            if len(input_ids) > data_args.cutoff_len:
-                input_ids = input_ids[:data_args.cutoff_len]
-                labels = labels[:data_args.cutoff_len]
-
        model_inputs["input_ids"].append(input_ids)
        model_inputs["attention_mask"].append([1] * len(input_ids))
        model_inputs["labels"].append(labels)

    return model_inputs

-    def preprocess_packed_supervised_dataset(examples: Dict[str, List[Any]]) -> Dict[str, List[List[int]]]:
+
+def preprocess_packed_supervised_dataset(
+    examples: Dict[str, List[Any]],
+    tokenizer: "PreTrainedTokenizer",
+    template: "Template",
+    data_args: "DataArguments",
+) -> Dict[str, List[List[int]]]:
    # build inputs with format `<bos> X1 Y1 <eos> <bos> X2 Y2 <eos>`
    # and labels with format `<ignore> ... <ignore> Y1 <eos> <ignore> ... <ignore> Y2 <eos>`
    model_inputs = {"input_ids": [], "attention_mask": [], "labels": []}
    input_ids, labels = [], []
-        for query, response, history, system in construct_example(examples):
-            if not (isinstance(query, str) and isinstance(response, str) and query != "" and response != ""):
+    for i in range(len(examples["prompt"])):
+        if len(examples["prompt"][i]) == 0 or len(examples["response"][i]) != 1:
            continue

-            for turn_idx, (source_ids, target_ids) in enumerate(template.encode_multiturn(
-                tokenizer, query, response, history, system
-            )):
+        messages = examples["prompt"][i] + examples["response"][i]
+        for turn_idx, (source_ids, target_ids) in enumerate(
+            template.encode_multiturn(tokenizer, messages, examples["system"][i], examples["tools"][i])
+        ):
            if data_args.train_on_prompt:
                source_mask = source_ids
            elif turn_idx != 0 and template.efficient_eos:
                source_mask = [tokenizer.eos_token_id] + [IGNORE_INDEX] * (len(source_ids) - 1)
            else:
                source_mask = [IGNORE_INDEX] * len(source_ids)
+
            input_ids += source_ids + target_ids
            labels += source_mask + target_ids

@@ -155,67 +126,80 @@ def preprocess_dataset(

    return model_inputs

-    def preprocess_unsupervised_dataset(examples: Dict[str, List[Any]]) -> Dict[str, List[List[int]]]:
+
+def preprocess_unsupervised_dataset(
+    examples: Dict[str, List[Any]],
+    tokenizer: "PreTrainedTokenizer",
+    template: "Template",
+    data_args: "DataArguments",
+) -> Dict[str, List[List[int]]]:
    # build inputs with format `<bos> X` and labels with format `Y <eos>`
    model_inputs = {"input_ids": [], "attention_mask": [], "labels": []}

-        for query, response, history, system in construct_example(examples):
-            if not (isinstance(query, str) and query != ""):
+    for i in range(len(examples["prompt"])):
+        if len(examples["prompt"][i]) == 0 or len(examples["response"][i]) != 1:
            continue

-            input_ids, labels = template.encode_oneturn(tokenizer, query, response, history, system)
+        messages = examples["prompt"][i] + examples["response"][i]
+        input_ids, labels = template.encode_oneturn(
+            tokenizer, messages, examples["system"][i], examples["tools"][i], data_args.cutoff_len
+        )

        if template.efficient_eos:
            labels += [tokenizer.eos_token_id]

-            if len(input_ids) > data_args.cutoff_len:
-                input_ids = input_ids[:data_args.cutoff_len]
-            if len(labels) > data_args.cutoff_len:
-                labels = labels[:data_args.cutoff_len]
-
        model_inputs["input_ids"].append(input_ids)
        model_inputs["attention_mask"].append([1] * len(input_ids))
        model_inputs["labels"].append(labels)

    return model_inputs

-    def preprocess_pairwise_dataset(examples: Dict[str, List[Any]]) -> Dict[str, List[List[int]]]:
+
+def preprocess_pairwise_dataset(
+    examples: Dict[str, List[Any]],
+    tokenizer: "PreTrainedTokenizer",
+    template: "Template",
+    data_args: "DataArguments",
+) -> Dict[str, List[List[int]]]:
    # build input pairs with format `<bos> X`, `Y1 <eos>` and `Y2 <eos>`
    model_inputs = {"prompt_ids": [], "chosen_ids": [], "rejected_ids": []}
-        for query, response, history, system in construct_example(examples):
-            if not (isinstance(query, str) and isinstance(response, list) and query != "" and len(response) > 1):
+    for i in range(len(examples["prompt"])):
+        if len(examples["prompt"][i]) == 0 or len(examples["response"][i]) < 2:
            continue

-            prompt_ids, chosen_ids = template.encode_oneturn(tokenizer, query, response[0], history, system)
-            _, rejected_ids = template.encode_oneturn(tokenizer, query, response[1], history, system)
+        chosen_messages = examples["prompt"][i] + [examples["response"][i][0]]
+        rejected_messages = examples["prompt"][i] + [examples["response"][i][1]]
+
+        prompt_ids, chosen_ids = template.encode_oneturn(
+            tokenizer, chosen_messages, examples["system"][i], examples["tools"][i], data_args.cutoff_len
+        )
+        _, rejected_ids = template.encode_oneturn(
+            tokenizer, rejected_messages, examples["system"][i], examples["tools"][i], data_args.cutoff_len
+        )

        if template.efficient_eos:
            chosen_ids += [tokenizer.eos_token_id]
            rejected_ids += [tokenizer.eos_token_id]

-            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 target_len > max_target_len:
-                chosen_ids = chosen_ids[:max_target_len]
-                rejected_ids = rejected_ids[:max_target_len]
-
        model_inputs["prompt_ids"].append(prompt_ids)
        model_inputs["chosen_ids"].append(chosen_ids)
        model_inputs["rejected_ids"].append(rejected_ids)

    return model_inputs

-    def print_supervised_dataset_example(example: Dict[str, List[int]]) -> None:
+
+def print_supervised_dataset_example(example: Dict[str, List[int]], tokenizer: "PreTrainedTokenizer") -> None:
    print("input_ids:\n{}".format(example["input_ids"]))
    print("inputs:\n{}".format(tokenizer.decode(example["input_ids"], skip_special_tokens=False)))
    print("label_ids:\n{}".format(example["labels"]))
-        print("labels:\n{}".format(
+    print(
+        "labels:\n{}".format(
            tokenizer.decode(list(filter(lambda x: x != IGNORE_INDEX, example["labels"])), skip_special_tokens=False)
-        ))
+        )
+    )

-    def print_pairwise_dataset_example(example: Dict[str, List[int]]) -> None:
+
+def print_pairwise_dataset_example(example: Dict[str, List[int]], tokenizer: "PreTrainedTokenizer") -> None:
    print("prompt_ids:\n{}".format(example["prompt_ids"]))
    print("prompt:\n{}".format(tokenizer.decode(example["prompt_ids"], skip_special_tokens=False)))
    print("chosen_ids:\n{}".format(example["chosen_ids"]))
@@ -223,53 +207,42 @@ def preprocess_dataset(
    print("rejected_ids:\n{}".format(example["rejected_ids"]))
    print("rejected:\n{}".format(tokenizer.decode(example["rejected_ids"], skip_special_tokens=False)))

-    def print_unsupervised_dataset_example(example: Dict[str, List[int]]) -> None:
+
+def print_unsupervised_dataset_example(example: Dict[str, List[int]], tokenizer: "PreTrainedTokenizer") -> None:
    print("input_ids:\n{}".format(example["input_ids"]))
    print("inputs:\n{}".format(tokenizer.decode(example["input_ids"], skip_special_tokens=False)))

+
+def get_preprocess_and_print_func(
+    tokenizer: "PreTrainedTokenizer",
+    template: "Template",
+    data_args: "DataArguments",
+    training_args: "Seq2SeqTrainingArguments",
+    stage: Literal["pt", "sft", "rm", "ppo"],
+) -> Tuple[Callable, Callable]:
    if stage == "pt":
-        preprocess_func = preprocess_pretrain_dataset
-        print_function = print_unsupervised_dataset_example
+        preprocess_func = partial(preprocess_pretrain_dataset, tokenizer=tokenizer, data_args=data_args)
+        print_function = partial(print_unsupervised_dataset_example, tokenizer=tokenizer)
    elif stage == "sft" and not training_args.predict_with_generate:
-        preprocess_func = preprocess_packed_supervised_dataset if data_args.sft_packing else preprocess_supervised_dataset
-        print_function = print_supervised_dataset_example
-    elif stage == "rm":
-        preprocess_func = preprocess_pairwise_dataset
-        print_function = print_pairwise_dataset_example
+        if data_args.sft_packing:
+            preprocess_func = partial(
+                preprocess_packed_supervised_dataset, tokenizer=tokenizer, template=template, data_args=data_args
+            )
        else:
-        preprocess_func = preprocess_unsupervised_dataset
-        print_function = print_unsupervised_dataset_example
-
-    if data_args.cache_path is not None and os.path.exists(data_args.cache_path):
-        logger.warning("Loading dataset from disk will ignore other data arguments.")
-        return load_from_disk(data_args.cache_path)
-
-    with training_args.main_process_first(desc="dataset map pre-processing"):
-        column_names = list(next(iter(dataset)).keys())
-        kwargs = {}
-        if not data_args.streaming:
-            kwargs = dict(
-                num_proc=data_args.preprocessing_num_workers,
-                load_from_cache_file=(not data_args.overwrite_cache),
-                desc="Running tokenizer on dataset"
+            preprocess_func = partial(
+                preprocess_supervised_dataset, tokenizer=tokenizer, template=template, data_args=data_args
            )

-        dataset = dataset.map(
-            preprocess_func,
-            batched=True,
-            remove_columns=column_names,
-            **kwargs
+        print_function = partial(print_supervised_dataset_example, tokenizer=tokenizer)
+    elif stage == "rm":
+        preprocess_func = partial(
+            preprocess_pairwise_dataset, tokenizer=tokenizer, template=template, data_args=data_args
        )
+        print_function = partial(print_pairwise_dataset_example, tokenizer=tokenizer)
+    else:
+        preprocess_func = partial(
+            preprocess_unsupervised_dataset, tokenizer=tokenizer, template=template, data_args=data_args
+        )
+        print_function = partial(print_unsupervised_dataset_example, tokenizer=tokenizer)

-        if data_args.cache_path is not None and not os.path.exists(data_args.cache_path):
-            if training_args.should_save:
-                dataset.save_to_disk(data_args.cache_path)
-            raise SystemExit("Dataset saved, rerun this script with the same `--cache_path`.")
-
-        if training_args.should_log:
-            try:
-                print_function(next(iter(dataset)))
-            except StopIteration:
-                raise RuntimeError("Empty dataset!")
-
-        return dataset
+    return preprocess_func, print_function
--- a/src/llmtuner/data/template.py
+++ b/src/llmtuner/data/template.py
--- a/src/llmtuner/data/utils.py
+++ b/src/llmtuner/data/utils.py
@@ -1,25 +1,26 @@
 import hashlib
-from typing import TYPE_CHECKING, Dict, List, Optional, Union
+from enum import Enum, unique
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+
+from ..extras.logging import get_logger

-from llmtuner.extras.logging import get_logger

 if TYPE_CHECKING:
    from datasets import Dataset, IterableDataset
    from transformers import TrainingArguments
+
    from llmtuner.hparams import DataArguments


 logger = get_logger(__name__)


-EXT2TYPE = {
-    "arrow": "arrow",
-    "csv": "csv",
-    "json": "json",
-    "jsonl": "json",
-    "parquet": "parquet",
-    "txt": "text"
-}
+@unique
+class Role(str, Enum):
+    USER = "user"
+    ASSISTANT = "assistant"
+    OBSERVATION = "observation"
+    FUNCTION = "function"


 def checksum(data_files: List[str], file_sha1: Optional[str] = None) -> None:
@@ -37,10 +38,15 @@ def checksum(data_files: List[str], file_sha1: Optional[str] = None) -> None:
            logger.warning("Checksum failed: mismatched SHA-1 hash value at {}.".format(data_files[0]))


+def infer_max_len(source_len: int, target_len: int, max_len: int, reserved_label_len: int) -> Tuple[int, int]:
+    max_target_len = int(max_len * (target_len / (source_len + target_len)))
+    max_target_len = max(max_target_len, reserved_label_len)
+    max_source_len = max_len - max_target_len
+    return max_source_len, max_target_len
+
+
 def split_dataset(
-    dataset: Union["Dataset", "IterableDataset"],
-    data_args: "DataArguments",
-    training_args: "TrainingArguments"
+    dataset: Union["Dataset", "IterableDataset"], data_args: "DataArguments", training_args: "TrainingArguments"
 ) -> Dict[str, "Dataset"]:
    if training_args.do_train:
        if data_args.val_size > 1e-6:  # Split the dataset
--- a/src/llmtuner/eval/init.py
+++ b/src/llmtuner/eval/init.py
@@ -1 +1,4 @@
-from llmtuner.eval.evaluator import Evaluator
+from .evaluator import Evaluator
+
+
+__all__ = ["Evaluator"]
--- a/src/llmtuner/eval/evaluator.py
+++ b/src/llmtuner/eval/evaluator.py
@@ -1,24 +1,24 @@
 # Inspired by: https://github.com/hendrycks/test/blob/master/evaluate_flan.py

-import os
+import inspect
 import json
-import torch
-import tiktoken
-import numpy as np
-from tqdm import tqdm, trange
+import os
 from typing import Any, Dict, List, Optional

+import numpy as np
+import torch
 from datasets import load_dataset
+from tqdm import tqdm, trange
 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
+from ..data import get_template_and_fix_tokenizer
+from ..extras.constants import CHOICES, SUBJECTS
+from ..hparams import get_eval_args
+from ..model import dispatch_model, load_model_and_tokenizer
+from .template import get_eval_template


 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)
@@ -26,15 +26,9 @@ class Evaluator:
        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]
+        self.choice_inputs = [
+            self.tokenizer.encode(self.eval_template.prefix + ch, add_special_tokens=False)[-1] for ch in CHOICES
+        ]

    @torch.inference_mode()
    def batch_inference(self, batch_input: Dict[str, torch.Tensor]) -> List[str]:
@@ -50,8 +44,8 @@ class Evaluator:
            filename="mapping.json",
            cache_dir=self.model_args.cache_dir,
            token=self.model_args.hf_hub_token,
-            revision=self.model_args.model_revision
        )
+
        with open(mapping, "r", encoding="utf-8") as f:
            categorys: Dict[str, Dict[str, str]] = json.load(f)

@@ -59,35 +53,45 @@ class Evaluator:
        pbar = tqdm(categorys.keys(), desc="Processing subjects", position=0)
        results = {}
        for subject in pbar:
+            if "trust_remote_code" in inspect.signature(load_dataset).parameters:  # for datasets==2.16.0
+                kwargs = {"trust_remote_code": True}
+            else:
+                kwargs = {}
+
            dataset = load_dataset(
                path=os.path.join(self.eval_args.task_dir, self.eval_args.task),
                name=subject,
-                download_mode="force_redownload"
+                cache_dir=self.model_args.cache_dir,
+                download_mode=self.eval_args.download_mode,
+                token=self.model_args.hf_hub_token,
+                **kwargs,
            )
            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(
+                support_set = (
+                    dataset["train"].shuffle().select(range(min(self.eval_args.n_shot, len(dataset["train"]))))
+                )
+                messages = 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):
+                input_ids, _ = self.template.encode_oneturn(tokenizer=self.tokenizer, messages=messages)
+                inputs.append({"input_ids": input_ids, "attention_mask": [1] * len(input_ids)})
+                labels.append(messages[-1]["content"])
+
+            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))
+            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)
@@ -97,10 +101,13 @@ class Evaluator:
        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([
+        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)
-        ])
+                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)
--- a/src/llmtuner/eval/template.py
+++ b/src/llmtuner/eval/template.py
@@ -1,7 +1,9 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Dict, List, Tuple

-from llmtuner.extras.constants import CHOICES
+from ..data import Role
+from ..extras.constants import CHOICES
+

 if TYPE_CHECKING:
    from datasets import Dataset
@@ -9,60 +11,39 @@ if TYPE_CHECKING:

@dataclass
 class EvalTemplate:
-
    system: str
    choice: str
    answer: str
    prefix: str

-    def parse_example(
-        self,
-        example: Dict[str, str]
-    ) -> Tuple[str, 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))]
+        self, target_data: Dict[str, str], support_set: "Dataset", subject_name: str
+    ) -> List[Dict[str, str]]:
+        messages = []
+        for k in range(len(support_set)):
+            prompt, response = self.parse_example(support_set[k])
+            messages.append({"role": Role.USER, "content": prompt})
+            messages.append({"role": Role.ASSISTANT, "content": response})

-        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
+        prompt, response = self.parse_example(target_data)
+        messages.append({"role": Role.USER, "content": prompt})
+        messages.append({"role": Role.ASSISTANT, "content": response})
+        messages[0]["content"] = self.system.format(subject=subject_name) + messages[0]["content"]
+        return messages


-eval_templates: Dict[str, EvalTemplate] = {}
+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 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:
+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
@@ -73,7 +54,7 @@ register_eval_template(
    system="The following are multiple choice questions (with answers) about {subject}.\n\n",
    choice="\n{choice}. {content}",
    answer="\nAnswer: ",
-    prefix=" "
+    prefix=" ",
 )


@@ -82,5 +63,5 @@ register_eval_template(
    system="以下是中国关于{subject}考试的单项选择题，请选出其中的正确答案。\n\n",
    choice="\n{choice}. {content}",
    answer="\n答案：",
-    prefix="\n"
+    prefix="\n",
 )
--- a/src/llmtuner/extras/callbacks.py
+++ b/src/llmtuner/extras/callbacks.py
@@ -1,53 +1,38 @@
-import os
 import json
+import os
 import time
-from typing import TYPE_CHECKING
 from datetime import timedelta
+from typing import TYPE_CHECKING

 from transformers import TrainerCallback
-from transformers.trainer_utils import has_length, PREFIX_CHECKPOINT_DIR
+from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR, has_length
+
+from .constants import LOG_FILE_NAME
+from .logging import get_logger
+from .misc import fix_valuehead_checkpoint

-from llmtuner.extras.constants import LOG_FILE_NAME
-from llmtuner.extras.logging import get_logger

 if TYPE_CHECKING:
-    from transformers import TrainingArguments, TrainerState, TrainerControl
-    from trl import AutoModelForCausalLMWithValueHead
+    from transformers import TrainerControl, TrainerState, TrainingArguments


 logger = get_logger(__name__)


-class SavePeftModelCallback(TrainerCallback):
-
+class FixValueHeadModelCallback(TrainerCallback):
    def on_save(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
        r"""
        Event called after a checkpoint save.
        """
        if args.should_save:
-            output_dir = os.path.join(args.output_dir, "{}-{}".format(PREFIX_CHECKPOINT_DIR, state.global_step))
-            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):
-                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: "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):
-                model.pretrained_model.save_pretrained(args.output_dir)
+            fix_valuehead_checkpoint(
+                model=kwargs.pop("model"),
+                output_dir=os.path.join(args.output_dir, "{}-{}".format(PREFIX_CHECKPOINT_DIR, state.global_step)),
+                safe_serialization=args.save_safetensors,
+            )


 class LogCallback(TrainerCallback):
-
    def __init__(self, runner=None):
        self.runner = runner
        self.in_training = False
@@ -113,7 +98,9 @@ class LogCallback(TrainerCallback):
            self.cur_steps = 0
            self.max_steps = 0

-    def on_predict(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", *other, **kwargs):
+    def on_predict(
+        self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", *other, **kwargs
+    ):
        r"""
        Event called after a successful prediction.
        """
@@ -139,18 +126,22 @@ class LogCallback(TrainerCallback):
            epoch=state.log_history[-1].get("epoch", None),
            percentage=round(self.cur_steps / self.max_steps * 100, 2) if self.max_steps != 0 else 100,
            elapsed_time=self.elapsed_time,
-            remaining_time=self.remaining_time
+            remaining_time=self.remaining_time,
        )
        if self.runner is not None:
-            logger.info("{{'loss': {:.4f}, 'learning_rate': {:2.4e}, 'epoch': {:.2f}}}".format(
+            logger.info(
+                "{{'loss': {:.4f}, 'learning_rate': {:2.4e}, 'epoch': {:.2f}}}".format(
                    logs["loss"] or 0, logs["learning_rate"] or 0, logs["epoch"] or 0
-            ))
+                )
+            )

        os.makedirs(args.output_dir, exist_ok=True)
        with open(os.path.join(args.output_dir, "trainer_log.jsonl"), "a", encoding="utf-8") as f:
            f.write(json.dumps(logs) + "\n")

-    def on_prediction_step(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+    def on_prediction_step(
+        self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs
+    ):
        r"""
        Event called after a prediction step.
        """
--- a/src/llmtuner/extras/constants.py
+++ b/src/llmtuner/extras/constants.py
@@ -1,13 +1,18 @@
-from collections import defaultdict, OrderedDict
+from collections import OrderedDict, defaultdict
+from enum import Enum
 from typing import Dict, Optional


 CHOICES = ["A", "B", "C", "D"]

+DATA_CONFIG = "dataset_info.json"
+
 DEFAULT_MODULE = defaultdict(str)

 DEFAULT_TEMPLATE = defaultdict(str)

+FILEEXT2TYPE = {"arrow": "arrow", "csv": "csv", "json": "json", "jsonl": "json", "parquet": "parquet", "txt": "text"}
+
 IGNORE_INDEX = -100

 LAYERNORM_NAMES = {"norm", "ln"}
@@ -16,6 +21,8 @@ LOG_FILE_NAME = "trainer_log.jsonl"

 METHODS = ["full", "freeze", "lora"]

+PEFT_METHODS = ["lora"]
+
 SUBJECTS = ["Average", "STEM", "Social Sciences", "Humanities", "Other"]

 SUPPORTED_MODELS = OrderedDict()
@@ -25,14 +32,21 @@ TRAINING_STAGES = {
    "Reward Modeling": "rm",
    "PPO": "ppo",
    "DPO": "dpo",
-    "Pre-Training": "pt"
+    "Pre-Training": "pt",
 }

+V_HEAD_WEIGHTS_NAME = "value_head.bin"
+
+V_HEAD_SAFE_WEIGHTS_NAME = "value_head.safetensors"
+
+
+class DownloadSource(str, Enum):
+    DEFAULT = "hf"
+    MODELSCOPE = "ms"
+

 def register_model_group(
-    models: Dict[str, str],
-    module: Optional[str] = None,
-    template: Optional[str] = None
+    models: Dict[str, Dict[DownloadSource, str]], module: Optional[str] = None, template: Optional[str] = None
 ) -> None:
    prefix = None
    for name, path in models.items():
@@ -49,230 +63,608 @@ def register_model_group(

 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"
+        "Baichuan-7B-Base": {
+            DownloadSource.DEFAULT: "baichuan-inc/Baichuan-7B",
+            DownloadSource.MODELSCOPE: "baichuan-inc/baichuan-7B",
+        },
+        "Baichuan-13B-Base": {
+            DownloadSource.DEFAULT: "baichuan-inc/Baichuan-13B-Base",
+            DownloadSource.MODELSCOPE: "baichuan-inc/Baichuan-13B-Base",
+        },
+        "Baichuan-13B-Chat": {
+            DownloadSource.DEFAULT: "baichuan-inc/Baichuan-13B-Chat",
+            DownloadSource.MODELSCOPE: "baichuan-inc/Baichuan-13B-Chat",
+        },
    },
    module="W_pack",
-    template="baichuan"
+    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"
+        "Baichuan2-7B-Base": {
+            DownloadSource.DEFAULT: "baichuan-inc/Baichuan2-7B-Base",
+            DownloadSource.MODELSCOPE: "baichuan-inc/Baichuan2-7B-Base",
+        },
+        "Baichuan2-13B-Base": {
+            DownloadSource.DEFAULT: "baichuan-inc/Baichuan2-13B-Base",
+            DownloadSource.MODELSCOPE: "baichuan-inc/Baichuan2-13B-Base",
+        },
+        "Baichuan2-7B-Chat": {
+            DownloadSource.DEFAULT: "baichuan-inc/Baichuan2-7B-Chat",
+            DownloadSource.MODELSCOPE: "baichuan-inc/Baichuan2-7B-Chat",
+        },
+        "Baichuan2-13B-Chat": {
+            DownloadSource.DEFAULT: "baichuan-inc/Baichuan2-13B-Chat",
+            DownloadSource.MODELSCOPE: "baichuan-inc/Baichuan2-13B-Chat",
+        },
    },
    module="W_pack",
-    template="baichuan2"
+    template="baichuan2",
 )


 register_model_group(
    models={
-        "BLOOM-560M": "bigscience/bloom-560m",
-        "BLOOM-3B": "bigscience/bloom-3b",
-        "BLOOM-7B1": "bigscience/bloom-7b1"
+        "BLOOM-560M": {
+            DownloadSource.DEFAULT: "bigscience/bloom-560m",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/bloom-560m",
        },
-    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"
+        "BLOOM-3B": {
+            DownloadSource.DEFAULT: "bigscience/bloom-3b",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/bloom-3b",
        },
-    module="query_key_value"
-)
-
-
-register_model_group(
-    models={
-        "BlueLM-7B-Base": "vivo-ai/BlueLM-7B-Base",
-        "BlueLM-7B-Chat": "vivo-ai/BlueLM-7B-Chat"
+        "BLOOM-7B1": {
+            DownloadSource.DEFAULT: "bigscience/bloom-7b1",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/bloom-7b1",
        },
-    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"
+        "BLOOMZ-560M": {
+            DownloadSource.DEFAULT: "bigscience/bloomz-560m",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/bloomz-560m",
+        },
+        "BLOOMZ-3B": {
+            DownloadSource.DEFAULT: "bigscience/bloomz-3b",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/bloomz-3b",
+        },
+        "BLOOMZ-7B1-mt": {
+            DownloadSource.DEFAULT: "bigscience/bloomz-7b1-mt",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/bloomz-7b1-mt",
+        },
    },
    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"
+        "BlueLM-7B-Base": {
+            DownloadSource.DEFAULT: "vivo-ai/BlueLM-7B-Base",
+            DownloadSource.MODELSCOPE: "vivo-ai/BlueLM-7B-Base",
        },
-    template="llama2_zh"
+        "BlueLM-7B-Chat": {
+            DownloadSource.DEFAULT: "vivo-ai/BlueLM-7B-Chat",
+            DownloadSource.MODELSCOPE: "vivo-ai/BlueLM-7B-Chat",
+        },
+    },
+    template="bluelm",
 )


 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"
+        "ChatGLM2-6B-Chat": {
+            DownloadSource.DEFAULT: "THUDM/chatglm2-6b",
+            DownloadSource.MODELSCOPE: "ZhipuAI/chatglm2-6b",
+        }
    },
    module="query_key_value",
-    template="falcon"
+    template="chatglm2",
 )


 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"
+        "ChatGLM3-6B-Base": {
+            DownloadSource.DEFAULT: "THUDM/chatglm3-6b-base",
+            DownloadSource.MODELSCOPE: "ZhipuAI/chatglm3-6b-base",
        },
-    template="intern"
-)
-
-
-register_model_group(
-    models={
-        "LingoWhale-8B": "deeplang-ai/LingoWhale-8B"
+        "ChatGLM3-6B-Chat": {
+            DownloadSource.DEFAULT: "THUDM/chatglm3-6b",
+            DownloadSource.MODELSCOPE: "ZhipuAI/chatglm3-6b",
        },
-    module="qkv_proj"
+    },
+    module="query_key_value",
+    template="chatglm3",
 )


 register_model_group(
    models={
-        "LLaMA-7B": "huggyllama/llama-7b",
-        "LLaMA-13B": "huggyllama/llama-13b",
-        "LLaMA-30B": "huggyllama/llama-30b",
-        "LLaMA-65B": "huggyllama/llama-65b"
+        "ChineseLLaMA2-1.3B": {
+            DownloadSource.DEFAULT: "hfl/chinese-llama-2-1.3b",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/chinese-llama-2-1.3b",
+        },
+        "ChineseLLaMA2-7B": {
+            DownloadSource.DEFAULT: "hfl/chinese-llama-2-7b",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/chinese-llama-2-7b",
+        },
+        "ChineseLLaMA2-13B": {
+            DownloadSource.DEFAULT: "hfl/chinese-llama-2-13b",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/chinese-llama-2-13b",
+        },
+        "ChineseLLaMA2-1.3B-Chat": {
+            DownloadSource.DEFAULT: "hfl/chinese-alpaca-2-1.3b",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/chinese-alpaca-2-1.3b",
+        },
+        "ChineseLLaMA2-7B-Chat": {
+            DownloadSource.DEFAULT: "hfl/chinese-alpaca-2-7b",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/chinese-alpaca-2-7b",
+        },
+        "ChineseLLaMA2-13B-Chat": {
+            DownloadSource.DEFAULT: "hfl/chinese-alpaca-2-13b",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/chinese-alpaca-2-13b",
+        },
+    },
+    template="llama2_zh",
+)
+
+
+register_model_group(
+    models={
+        "DeepSeekLLM-7B-Base": {
+            DownloadSource.DEFAULT: "deepseek-ai/deepseek-llm-7b-base",
+            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-llm-7b-base",
+        },
+        "DeepSeekLLM-67B-Base": {
+            DownloadSource.DEFAULT: "deepseek-ai/deepseek-llm-67b-base",
+            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-llm-67b-base",
+        },
+        "DeepSeekLLM-7B-Chat": {
+            DownloadSource.DEFAULT: "deepseek-ai/deepseek-llm-7b-chat",
+            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-llm-7b-chat",
+        },
+        "DeepSeekLLM-67B-Chat": {
+            DownloadSource.DEFAULT: "deepseek-ai/deepseek-llm-67b-chat",
+            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-llm-67b-chat",
+        },
+    },
+    template="deepseek",
+)
+
+
+register_model_group(
+    models={
+        "DeepSeekCoder-6.7B-Base": {
+            DownloadSource.DEFAULT: "deepseek-ai/deepseek-coder-6.7b-base",
+            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-coder-6.7b-base",
+        },
+        "DeepSeekCoder-33B-Base": {
+            DownloadSource.DEFAULT: "deepseek-ai/deepseek-coder-33b-base",
+            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-coder-33b-base",
+        },
+        "DeepSeekCoder-6.7B-Chat": {
+            DownloadSource.DEFAULT: "deepseek-ai/deepseek-coder-6.7b-instruct",
+            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-coder-6.7b-instruct",
+        },
+        "DeepSeekCoder-33B-Chat": {
+            DownloadSource.DEFAULT: "deepseek-ai/deepseek-coder-33b-instruct",
+            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-coder-33b-instruct",
+        },
+    },
+    template="deepseekcoder",
+)
+
+
+register_model_group(
+    models={
+        "DeepSeekMoE-16B-Base": {
+            DownloadSource.DEFAULT: "deepseek-ai/deepseek-moe-16b-base",
+            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-moe-16b-base",
+        },
+        "DeepSeekMoE-16B-Chat": {
+            DownloadSource.DEFAULT: "deepseek-ai/deepseek-moe-16b-chat",
+            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-moe-16b-chat",
+        },
+    },
+    template="deepseek",
+)
+
+
+register_model_group(
+    models={
+        "Falcon-7B": {
+            DownloadSource.DEFAULT: "tiiuae/falcon-7b",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/falcon-7b",
+        },
+        "Falcon-40B": {
+            DownloadSource.DEFAULT: "tiiuae/falcon-40b",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/falcon-40b",
+        },
+        "Falcon-180B": {
+            DownloadSource.DEFAULT: "tiiuae/falcon-180b",
+            DownloadSource.MODELSCOPE: "modelscope/falcon-180B",
+        },
+        "Falcon-7B-Chat": {
+            DownloadSource.DEFAULT: "tiiuae/falcon-7b-instruct",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/falcon-7b-instruct",
+        },
+        "Falcon-40B-Chat": {
+            DownloadSource.DEFAULT: "tiiuae/falcon-40b-instruct",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/falcon-40b-instruct",
+        },
+        "Falcon-180B-Chat": {
+            DownloadSource.DEFAULT: "tiiuae/falcon-180b-chat",
+            DownloadSource.MODELSCOPE: "modelscope/falcon-180B-chat",
+        },
+    },
+    module="query_key_value",
+    template="falcon",
+)
+
+
+register_model_group(
+    models={
+        "InternLM-7B": {
+            DownloadSource.DEFAULT: "internlm/internlm-7b",
+            DownloadSource.MODELSCOPE: "Shanghai_AI_Laboratory/internlm-7b",
+        },
+        "InternLM-20B": {
+            DownloadSource.DEFAULT: "internlm/internlm-20b",
+            DownloadSource.MODELSCOPE: "Shanghai_AI_Laboratory/internlm-20b",
+        },
+        "InternLM-7B-Chat": {
+            DownloadSource.DEFAULT: "internlm/internlm-chat-7b",
+            DownloadSource.MODELSCOPE: "Shanghai_AI_Laboratory/internlm-chat-7b",
+        },
+        "InternLM-20B-Chat": {
+            DownloadSource.DEFAULT: "internlm/internlm-chat-20b",
+            DownloadSource.MODELSCOPE: "Shanghai_AI_Laboratory/internlm-chat-20b",
+        },
+    },
+    template="intern",
+)
+
+
+register_model_group(
+    models={
+        "InternLM2-7B": {
+            DownloadSource.DEFAULT: "internlm/internlm2-7b",
+            DownloadSource.MODELSCOPE: "Shanghai_AI_Laboratory/internlm2-7b",
+        },
+        "InternLM2-20B": {
+            DownloadSource.DEFAULT: "internlm/internlm2-20b",
+            DownloadSource.MODELSCOPE: "Shanghai_AI_Laboratory/internlm2-20b",
+        },
+        "InternLM2-7B-Chat": {
+            DownloadSource.DEFAULT: "internlm/internlm2-chat-7b",
+            DownloadSource.MODELSCOPE: "Shanghai_AI_Laboratory/internlm2-chat-7b",
+        },
+        "InternLM2-20B-Chat": {
+            DownloadSource.DEFAULT: "internlm/internlm2-chat-20b",
+            DownloadSource.MODELSCOPE: "Shanghai_AI_Laboratory/internlm2-chat-20b",
+        },
+    },
+    module="wqkv",
+    template="intern2",
+)
+
+
+register_model_group(
+    models={
+        "LingoWhale-8B": {
+            DownloadSource.DEFAULT: "deeplang-ai/LingoWhale-8B",
+            DownloadSource.MODELSCOPE: "DeepLang/LingoWhale-8B",
+        }
+    },
+    module="qkv_proj",
+)
+
+
+register_model_group(
+    models={
+        "LLaMA-7B": {DownloadSource.DEFAULT: "huggyllama/llama-7b", DownloadSource.MODELSCOPE: "skyline2006/llama-7b"},
+        "LLaMA-13B": {
+            DownloadSource.DEFAULT: "huggyllama/llama-13b",
+            DownloadSource.MODELSCOPE: "skyline2006/llama-13b",
+        },
+        "LLaMA-30B": {
+            DownloadSource.DEFAULT: "huggyllama/llama-30b",
+            DownloadSource.MODELSCOPE: "skyline2006/llama-30b",
+        },
+        "LLaMA-65B": {
+            DownloadSource.DEFAULT: "huggyllama/llama-65b",
+            DownloadSource.MODELSCOPE: "skyline2006/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"
+        "LLaMA2-7B": {
+            DownloadSource.DEFAULT: "meta-llama/Llama-2-7b-hf",
+            DownloadSource.MODELSCOPE: "modelscope/Llama-2-7b-ms",
        },
-    template="llama2"
+        "LLaMA2-13B": {
+            DownloadSource.DEFAULT: "meta-llama/Llama-2-13b-hf",
+            DownloadSource.MODELSCOPE: "modelscope/Llama-2-13b-ms",
+        },
+        "LLaMA2-70B": {
+            DownloadSource.DEFAULT: "meta-llama/Llama-2-70b-hf",
+            DownloadSource.MODELSCOPE: "modelscope/Llama-2-70b-ms",
+        },
+        "LLaMA2-7B-Chat": {
+            DownloadSource.DEFAULT: "meta-llama/Llama-2-7b-chat-hf",
+            DownloadSource.MODELSCOPE: "modelscope/Llama-2-7b-chat-ms",
+        },
+        "LLaMA2-13B-Chat": {
+            DownloadSource.DEFAULT: "meta-llama/Llama-2-13b-chat-hf",
+            DownloadSource.MODELSCOPE: "modelscope/Llama-2-13b-chat-ms",
+        },
+        "LLaMA2-70B-Chat": {
+            DownloadSource.DEFAULT: "meta-llama/Llama-2-70b-chat-hf",
+            DownloadSource.MODELSCOPE: "modelscope/Llama-2-70b-chat-ms",
+        },
+    },
+    template="llama2",
 )


 register_model_group(
    models={
-        "Mistral-7B": "mistralai/Mistral-7B-v0.1",
-        "Mistral-7B-Chat": "mistralai/Mistral-7B-Instruct-v0.1"
+        "Mistral-7B": {
+            DownloadSource.DEFAULT: "mistralai/Mistral-7B-v0.1",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/Mistral-7B-v0.1",
        },
-    template="mistral"
+        "Mistral-7B-Chat": {
+            DownloadSource.DEFAULT: "mistralai/Mistral-7B-Instruct-v0.1",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/Mistral-7B-Instruct-v0.1",
+        },
+        "Mistral-7B-v0.2-Chat": {
+            DownloadSource.DEFAULT: "mistralai/Mistral-7B-Instruct-v0.2",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/Mistral-7B-Instruct-v0.2",
+        },
+    },
+    template="mistral",
 )


 register_model_group(
    models={
-        "OpenChat3.5-7B-Chat": "openchat/openchat_3.5"
+        "Mixtral-8x7B": {
+            DownloadSource.DEFAULT: "mistralai/Mixtral-8x7B-v0.1",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/Mixtral-8x7B-v0.1",
        },
-    template="openchat"
+        "Mixtral-8x7B-Chat": {
+            DownloadSource.DEFAULT: "mistralai/Mixtral-8x7B-Instruct-v0.1",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/Mixtral-8x7B-Instruct-v0.1",
+        },
+    },
+    template="mistral",
 )


 register_model_group(
    models={
-        "Phi1.5-1.3B": "microsoft/phi-1_5"
+        "OpenChat3.5-7B-Chat": {
+            DownloadSource.DEFAULT: "openchat/openchat_3.5",
+            DownloadSource.MODELSCOPE: "myxiongmodel/openchat_3.5",
+        }
    },
-    module="Wqkv"
+    template="openchat",
 )


 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"
+        "Phi-1.5-1.3B": {DownloadSource.DEFAULT: "microsoft/phi-1_5", DownloadSource.MODELSCOPE: "allspace/PHI_1-5"},
+        "Phi-2-2.7B": {DownloadSource.DEFAULT: "microsoft/phi-2", DownloadSource.MODELSCOPE: "AI-ModelScope/phi-2"},
+    }
+)
+
+
+register_model_group(
+    models={
+        "Qwen-1.8B": {DownloadSource.DEFAULT: "Qwen/Qwen-1_8B", DownloadSource.MODELSCOPE: "qwen/Qwen-1_8B"},
+        "Qwen-7B": {DownloadSource.DEFAULT: "Qwen/Qwen-7B", DownloadSource.MODELSCOPE: "qwen/Qwen-7B"},
+        "Qwen-14B": {DownloadSource.DEFAULT: "Qwen/Qwen-14B", DownloadSource.MODELSCOPE: "qwen/Qwen-14B"},
+        "Qwen-72B": {DownloadSource.DEFAULT: "Qwen/Qwen-72B", DownloadSource.MODELSCOPE: "qwen/Qwen-72B"},
+        "Qwen-1.8B-Chat": {
+            DownloadSource.DEFAULT: "Qwen/Qwen-1_8B-Chat",
+            DownloadSource.MODELSCOPE: "qwen/Qwen-1_8B-Chat",
+        },
+        "Qwen-7B-Chat": {DownloadSource.DEFAULT: "Qwen/Qwen-7B-Chat", DownloadSource.MODELSCOPE: "qwen/Qwen-7B-Chat"},
+        "Qwen-14B-Chat": {
+            DownloadSource.DEFAULT: "Qwen/Qwen-14B-Chat",
+            DownloadSource.MODELSCOPE: "qwen/Qwen-14B-Chat",
+        },
+        "Qwen-72B-Chat": {
+            DownloadSource.DEFAULT: "Qwen/Qwen-72B-Chat",
+            DownloadSource.MODELSCOPE: "qwen/Qwen-72B-Chat",
+        },
+        "Qwen-1.8B-int8-Chat": {
+            DownloadSource.DEFAULT: "Qwen/Qwen-1_8B-Chat-Int8",
+            DownloadSource.MODELSCOPE: "qwen/Qwen-1_8B-Chat-Int8",
+        },
+        "Qwen-1.8B-int4-Chat": {
+            DownloadSource.DEFAULT: "Qwen/Qwen-1_8B-Chat-Int4",
+            DownloadSource.MODELSCOPE: "qwen/Qwen-1_8B-Chat-Int4",
+        },
+        "Qwen-7B-int8-Chat": {
+            DownloadSource.DEFAULT: "Qwen/Qwen-7B-Chat-Int8",
+            DownloadSource.MODELSCOPE: "qwen/Qwen-7B-Chat-Int8",
+        },
+        "Qwen-7B-int4-Chat": {
+            DownloadSource.DEFAULT: "Qwen/Qwen-7B-Chat-Int4",
+            DownloadSource.MODELSCOPE: "qwen/Qwen-7B-Chat-Int4",
+        },
+        "Qwen-14B-int8-Chat": {
+            DownloadSource.DEFAULT: "Qwen/Qwen-14B-Chat-Int8",
+            DownloadSource.MODELSCOPE: "qwen/Qwen-14B-Chat-Int8",
+        },
+        "Qwen-14B-int4-Chat": {
+            DownloadSource.DEFAULT: "Qwen/Qwen-14B-Chat-Int4",
+            DownloadSource.MODELSCOPE: "qwen/Qwen-14B-Chat-Int4",
+        },
+        "Qwen-72B-int8-Chat": {
+            DownloadSource.DEFAULT: "Qwen/Qwen-72B-Chat-Int8",
+            DownloadSource.MODELSCOPE: "qwen/Qwen-72B-Chat-Int8",
+        },
+        "Qwen-72B-int4-Chat": {
+            DownloadSource.DEFAULT: "Qwen/Qwen-72B-Chat-Int4",
+            DownloadSource.MODELSCOPE: "qwen/Qwen-72B-Chat-Int4",
+        },
    },
    module="c_attn",
-    template="qwen"
+    template="qwen",
 )


 register_model_group(
    models={
-        "Skywork-13B-Base": "Skywork/Skywork-13B-base"
+        "SOLAR-10.7B": {DownloadSource.DEFAULT: "upstage/SOLAR-10.7B-v1.0"},
+        "SOLAR-10.7B-Chat": {
+            DownloadSource.DEFAULT: "upstage/SOLAR-10.7B-Instruct-v1.0",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/SOLAR-10.7B-Instruct-v1.0",
+        },
+    },
+    template="solar",
+)
+
+
+register_model_group(
+    models={
+        "Skywork-13B-Base": {
+            DownloadSource.DEFAULT: "Skywork/Skywork-13B-base",
+            DownloadSource.MODELSCOPE: "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"
+        "Vicuna1.5-7B-Chat": {
+            DownloadSource.DEFAULT: "lmsys/vicuna-7b-v1.5",
+            DownloadSource.MODELSCOPE: "Xorbits/vicuna-7b-v1.5",
        },
-    template="vicuna"
+        "Vicuna1.5-13B-Chat": {
+            DownloadSource.DEFAULT: "lmsys/vicuna-13b-v1.5",
+            DownloadSource.MODELSCOPE: "Xorbits/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"
+        "XuanYuan-70B": {DownloadSource.DEFAULT: "Duxiaoman-DI/XuanYuan-70B"},
+        "XuanYuan-70B-Chat": {DownloadSource.DEFAULT: "Duxiaoman-DI/XuanYuan-70B-Chat"},
+        "XuanYuan-70B-int8-Chat": {DownloadSource.DEFAULT: "Duxiaoman-DI/XuanYuan-70B-Chat-8bit"},
+        "XuanYuan-70B-int4-Chat": {DownloadSource.DEFAULT: "Duxiaoman-DI/XuanYuan-70B-Chat-4bit"},
    },
-    template="xverse"
+    template="xuanyuan",
 )


 register_model_group(
    models={
-        "Yayi-7B": "wenge-research/yayi-7b-llama2",
-        "Yayi-13B": "wenge-research/yayi-13b-llama2"
+        "XVERSE-7B": {DownloadSource.DEFAULT: "xverse/XVERSE-7B", DownloadSource.MODELSCOPE: "xverse/XVERSE-7B"},
+        "XVERSE-13B": {DownloadSource.DEFAULT: "xverse/XVERSE-13B", DownloadSource.MODELSCOPE: "xverse/XVERSE-13B"},
+        "XVERSE-65B": {DownloadSource.DEFAULT: "xverse/XVERSE-65B", DownloadSource.MODELSCOPE: "xverse/XVERSE-65B"},
+        "XVERSE-65B-2": {
+            DownloadSource.DEFAULT: "xverse/XVERSE-65B-2",
+            DownloadSource.MODELSCOPE: "xverse/XVERSE-65B-2",
        },
-    template="yayi"
+        "XVERSE-7B-Chat": {
+            DownloadSource.DEFAULT: "xverse/XVERSE-7B-Chat",
+            DownloadSource.MODELSCOPE: "xverse/XVERSE-7B-Chat",
+        },
+        "XVERSE-13B-Chat": {
+            DownloadSource.DEFAULT: "xverse/XVERSE-13B-Chat",
+            DownloadSource.MODELSCOPE: "xverse/XVERSE-13B-Chat",
+        },
+        "XVERSE-65B-Chat": {
+            DownloadSource.DEFAULT: "xverse/XVERSE-65B-Chat",
+            DownloadSource.MODELSCOPE: "xverse/XVERSE-65B-Chat",
+        },
+    },
+    template="xverse",
 )


 register_model_group(
    models={
-        "Yi-6B": "01-ai/Yi-6B",
-        "Yi-34B": "01-ai/Yi-34B"
-    }
+        "Yayi-7B": {
+            DownloadSource.DEFAULT: "wenge-research/yayi-7b-llama2",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/yayi-7b-llama2",
+        },
+        "Yayi-13B": {
+            DownloadSource.DEFAULT: "wenge-research/yayi-13b-llama2",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/yayi-13b-llama2",
+        },
+    },
+    template="yayi",
 )


 register_model_group(
    models={
-        "Zephyr-7B-Alpha-Chat": "HuggingFaceH4/zephyr-7b-alpha",
-        "Zephyr-7B-Beta-Chat": "HuggingFaceH4/zephyr-7b-beta"
+        "Yi-6B": {DownloadSource.DEFAULT: "01-ai/Yi-6B", DownloadSource.MODELSCOPE: "01ai/Yi-6B"},
+        "Yi-34B": {DownloadSource.DEFAULT: "01-ai/Yi-34B", DownloadSource.MODELSCOPE: "01ai/Yi-34B"},
+        "Yi-6B-Chat": {DownloadSource.DEFAULT: "01-ai/Yi-6B-Chat", DownloadSource.MODELSCOPE: "01ai/Yi-6B-Chat"},
+        "Yi-34B-Chat": {DownloadSource.DEFAULT: "01-ai/Yi-34B-Chat", DownloadSource.MODELSCOPE: "01ai/Yi-34B-Chat"},
+        "Yi-6B-int8-Chat": {
+            DownloadSource.DEFAULT: "01-ai/Yi-6B-Chat-8bits",
+            DownloadSource.MODELSCOPE: "01ai/Yi-6B-Chat-8bits",
        },
-    template="zephyr"
+        "Yi-34B-int8-Chat": {
+            DownloadSource.DEFAULT: "01-ai/Yi-34B-Chat-8bits",
+            DownloadSource.MODELSCOPE: "01ai/Yi-34B-Chat-8bits",
+        },
+    },
+    template="yi",
+)
+
+
+register_model_group(
+    models={
+        "Yuan2-2B-Chat": {
+            DownloadSource.DEFAULT: "IEITYuan/Yuan2-2B-hf",
+            DownloadSource.MODELSCOPE: "YuanLLM/Yuan2.0-2B-hf",
+        },
+        "Yuan2-51B-Chat": {
+            DownloadSource.DEFAULT: "IEITYuan/Yuan2-51B-hf",
+            DownloadSource.MODELSCOPE: "YuanLLM/Yuan2.0-51B-hf",
+        },
+        "Yuan2-102B-Chat": {
+            DownloadSource.DEFAULT: "IEITYuan/Yuan2-102B-hf",
+            DownloadSource.MODELSCOPE: "YuanLLM/Yuan2.0-102B-hf",
+        },
+    },
+    template="yuan",
+)
+
+
+register_model_group(
+    models={
+        "Zephyr-7B-Alpha-Chat": {
+            DownloadSource.DEFAULT: "HuggingFaceH4/zephyr-7b-alpha",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/zephyr-7b-alpha",
+        },
+        "Zephyr-7B-Beta-Chat": {
+            DownloadSource.DEFAULT: "HuggingFaceH4/zephyr-7b-beta",
+            DownloadSource.MODELSCOPE: "modelscope/zephyr-7b-beta",
+        },
+    },
+    template="zephyr",
 )
--- a/src/llmtuner/extras/logging.py
+++ b/src/llmtuner/extras/logging.py
@@ -1,5 +1,5 @@
-import sys
 import logging
+import sys


 class LoggerHandler(logging.Handler):
@@ -27,8 +27,7 @@ def get_logger(name: str) -> logging.Logger:
    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"
+        fmt="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S"
    )
    handler = logging.StreamHandler(sys.stdout)
    handler.setFormatter(formatter)
--- a/src/llmtuner/extras/misc.py
+++ b/src/llmtuner/extras/misc.py
@@ -1,31 +1,44 @@
 import gc
 import os
-import sys
-import torch
-from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
-from transformers import InfNanRemoveLogitsProcessor, LogitsProcessorList
+from typing import TYPE_CHECKING, Dict, Tuple

-try:
+import torch
+from peft import PeftModel
+from transformers import InfNanRemoveLogitsProcessor, LogitsProcessorList, PreTrainedModel
 from transformers.utils import (
-        is_torch_bf16_cpu_available,
+    SAFE_WEIGHTS_NAME,
+    WEIGHTS_NAME,
    is_torch_bf16_gpu_available,
    is_torch_cuda_available,
-        is_torch_npu_available
+    is_torch_npu_available,
+    is_torch_xpu_available,
 )
+
+from .constants import V_HEAD_SAFE_WEIGHTS_NAME, V_HEAD_WEIGHTS_NAME
+from .logging import get_logger
+
+
 _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()
-except ImportError:
-    _is_fp16_available = torch.cuda.is_available()
-    _is_bf16_available = torch.cuda.is_bf16_supported()
+try:
+    _is_bf16_available = is_torch_bf16_gpu_available()
+except Exception:
+    _is_bf16_available = False
+

 if TYPE_CHECKING:
-    from transformers import HfArgumentParser
+    from trl import AutoModelForCausalLMWithValueHead
+
+    from llmtuner.hparams import ModelArguments
+
+
+logger = get_logger(__name__)


 class AverageMeter:
    r"""
    Computes and stores the average and current value.
    """
+
    def __init__(self):
        self.reset()

@@ -64,14 +77,72 @@ def count_parameters(model: torch.nn.Module) -> Tuple[int, int]:
    return trainable_params, all_param


-def get_current_device() -> str:
-    import accelerate
-    from accelerate import Accelerator
-    dummy_accelerator = Accelerator()
-    if accelerate.utils.is_xpu_available():
-        return "xpu:{}".format(dummy_accelerator.local_process_index)
+def fix_valuehead_checkpoint(
+    model: "AutoModelForCausalLMWithValueHead", output_dir: str, safe_serialization: bool
+) -> None:
+    r"""
+    The model is already unwrapped.
+
+    There are three cases:
+    1. full tuning without ds_zero3: state_dict = {"model.layers.*": ..., "v_head.summary.*": ...}
+    2. lora tuning without ds_zero3: state_dict = {"v_head.summary.*": ...}
+    3. under deepspeed zero3: state_dict = {"pretrained_model.model.layers.*": ..., "v_head.summary.*": ...}
+
+    We assume `stage3_gather_16bit_weights_on_model_save=true`.
+    """
+    if not isinstance(model.pretrained_model, (PreTrainedModel, PeftModel)):
+        return
+
+    if safe_serialization:
+        from safetensors import safe_open
+        from safetensors.torch import save_file
+
+        path_to_checkpoint = os.path.join(output_dir, SAFE_WEIGHTS_NAME)
+        with safe_open(path_to_checkpoint, framework="pt", device="cpu") as f:
+            state_dict: Dict[str, torch.Tensor] = {key: f.get_tensor(key) for key in f.keys()}
    else:
-        return dummy_accelerator.local_process_index if torch.cuda.is_available() else "cpu"
+        path_to_checkpoint = os.path.join(output_dir, WEIGHTS_NAME)
+        state_dict: Dict[str, torch.Tensor] = torch.load(path_to_checkpoint, map_location="cpu")
+
+    decoder_state_dict = {}
+    v_head_state_dict = {}
+    for name, param in state_dict.items():
+        if name.startswith("v_head."):
+            v_head_state_dict[name] = param
+        else:
+            decoder_state_dict[name.replace("pretrained_model.", "")] = param
+
+    os.remove(path_to_checkpoint)
+    model.pretrained_model.save_pretrained(
+        output_dir, state_dict=decoder_state_dict or None, safe_serialization=safe_serialization
+    )
+
+    if safe_serialization:
+        save_file(v_head_state_dict, os.path.join(output_dir, V_HEAD_SAFE_WEIGHTS_NAME), metadata={"format": "pt"})
+    else:
+        torch.save(v_head_state_dict, os.path.join(output_dir, V_HEAD_WEIGHTS_NAME))
+
+    logger.info("Value head model saved at: {}".format(output_dir))
+
+
+def get_current_device() -> torch.device:
+    r"""
+    Gets the current available device.
+    """
+    if is_torch_xpu_available():
+        device = "xpu:{}".format(os.environ.get("LOCAL_RANK", "0"))
+    elif is_torch_npu_available():
+        device = "npu:{}".format(os.environ.get("LOCAL_RANK", "0"))
+    elif is_torch_cuda_available():
+        device = "cuda:{}".format(os.environ.get("LOCAL_RANK", "0"))
+    else:
+        device = "cpu"
+
+    return torch.device(device)
+
+
+def get_device_count() -> int:
+    return torch.cuda.device_count()


 def get_logits_processor() -> "LogitsProcessorList":
@@ -95,17 +166,6 @@ def infer_optim_dtype(model_dtype: torch.dtype) -> torch.dtype:
        return torch.float32


-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:
    r"""
    Collects GPU memory.
@@ -114,3 +174,22 @@ def torch_gc() -> None:
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()
+
+
+def try_download_model_from_ms(model_args: "ModelArguments") -> None:
+    if not use_modelscope() or os.path.exists(model_args.model_name_or_path):
+        return
+
+    try:
+        from modelscope import snapshot_download
+
+        revision = "master" if model_args.model_revision == "main" else model_args.model_revision
+        model_args.model_name_or_path = snapshot_download(
+            model_args.model_name_or_path, revision=revision, cache_dir=model_args.cache_dir
+        )
+    except ImportError:
+        raise ImportError("Please install modelscope via `pip install modelscope -U`")
+
+
+def use_modelscope() -> bool:
+    return bool(int(os.environ.get("USE_MODELSCOPE_HUB", "0")))
--- a/src/llmtuner/extras/packages.py
+++ b/src/llmtuner/extras/packages.py
@@ -9,47 +9,41 @@ def is_package_available(name: str) -> bool:
 def get_package_version(name: str) -> str:
    try:
        return importlib.metadata.version(name)
-    except:
+    except Exception:
        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
+    return is_package_available("fastapi")


 def is_flash_attn2_available():
-    return _flash_attn2_available
+    return is_package_available("flash_attn") and get_package_version("flash_attn").startswith("2")


 def is_jieba_available():
-    return _jieba_available
+    return is_package_available("jieba")


 def is_matplotlib_available():
-    return _matplotlib_available
+    return is_package_available("matplotlib")


 def is_nltk_available():
-    return _nltk_available
+    return is_package_available("nltk")
+
+
+def is_requests_available():
+    return is_package_available("requests")


 def is_rouge_available():
-    return _rouge_available
+    return is_package_available("rouge_chinese")


 def is_starlette_available():
-    return _starlette_available
+    return is_package_available("sse_starlette")


 def is_uvicorn_available():
-    return _uvicorn_available
+    return is_package_available("uvicorn")
--- a/src/llmtuner/extras/patches/llama_patch.py
+++ b/src/llmtuner/extras/patches/llama_patch.py
@@ -1,38 +1,30 @@
 import math
+from typing import Optional, Tuple
+
 import torch
 import torch.nn as nn
-from typing import Optional, Tuple
+from transformers.models.llama.modeling_llama import (
+    Cache,
+    LlamaAttention,
+    LlamaFlashAttention2,
+    apply_rotary_pos_emb,
+    repeat_kv,
+)
 from transformers.utils import logging
-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


 logger = logging.get_logger(__name__)


 # Modified from: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
-class LlamaShiftShortAttention(LlamaAttention):
-
-    def forward(
-        self,
+def llama_torch_attn_forward(
+    self: "LlamaAttention",
    hidden_states: torch.Tensor,
    attention_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.LongTensor] = None,
-        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+    past_key_value: Optional["Cache"] = None,
    output_attentions: bool = False,
-        use_cache: bool = False,
-        **kwargs
+    **kwargs,
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
    bsz, q_len, _ = hidden_states.size()

@@ -46,18 +38,15 @@ class LlamaShiftShortAttention(LlamaAttention):

    kv_seq_len = key_states.shape[-2]
    if past_key_value is not None:
-            kv_seq_len += past_key_value[0].shape[-2]
+        kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)

    cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)

-        if past_key_value is not None: # reuse k, v, self_attention
-            key_states = torch.cat([past_key_value[0], key_states], dim=2)
-            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+    if past_key_value is not None:
+        cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
+        key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)

-        past_key_value = (key_states, value_states) if use_cache else None
-
-        if getattr(self, "num_key_value_groups"):
    key_states = repeat_kv(key_states, self.num_key_value_groups)
    value_states = repeat_kv(value_states, self.num_key_value_groups)

@@ -65,11 +54,13 @@ class LlamaShiftShortAttention(LlamaAttention):
        groupsz = int(q_len * getattr(self.config, "group_size_ratio"))
        assert q_len % groupsz == 0, "q_len {} should be divisible by group size {}.".format(q_len, groupsz)
        num_groups = q_len // groupsz
+
        def shift(state: torch.Tensor) -> torch.Tensor:
            state = state.transpose(1, 2)  # output: (bsz, seq_len, n_heads, head_dim)
-                state = torch.cat((
-                    state[:, :, :self.num_heads//2], state[:, :, self.num_heads//2:].roll(-groupsz//2, dims=1)
-                ), dim=2)
+            state = torch.cat(
+                (state[:, :, : self.num_heads // 2], state[:, :, self.num_heads // 2 :].roll(-groupsz // 2, dims=1)),
+                dim=2,
+            )
            return state.reshape(bsz * num_groups, groupsz, self.num_heads, self.head_dim).transpose(1, 2)

        query_states, key_states, value_states = shift(query_states), shift(key_states), shift(value_states)
@@ -83,14 +74,18 @@ class LlamaShiftShortAttention(LlamaAttention):

    # upcast attention to fp32
    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
    attn_output = torch.matmul(attn_weights, value_states)  # (bsz, :, seq_len, :) or (bsz*n_group, :, groupsz, :)
    attn_output = attn_output.transpose(1, 2).contiguous()

    if getattr(self.config, "group_size_ratio", None) and self.training:  # shift back
        attn_output.reshape(bsz, q_len, self.num_heads, self.head_dim)
-            attn_output = torch.cat((
-                attn_output[:, :, :self.num_heads//2], attn_output[:, :, self.num_heads//2:].roll(groupsz//2, dims=1)
-            ))
+        attn_output = torch.cat(
+            (
+                attn_output[:, :, : self.num_heads // 2],
+                attn_output[:, :, self.num_heads // 2 :].roll(groupsz // 2, dims=1),
+            )
+        )

    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
    attn_output = self.o_proj(attn_output)
@@ -101,17 +96,15 @@ class LlamaShiftShortAttention(LlamaAttention):
    return attn_output, attn_weights, past_key_value


-class LlamaFlashAttention2(LlamaAttention):
-
-    def forward(
-        self,
+# Modified from: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
+def llama_flash_attn_forward(
+    self: "LlamaFlashAttention2",
    hidden_states: torch.Tensor,
    attention_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.LongTensor] = None,
    past_key_value: Optional[Tuple[torch.Tensor]] = None,
    output_attentions: bool = False,
-        use_cache: bool = False,
-        **kwargs
+    **kwargs,
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
    # LlamaFlashAttention2 attention does not support output_attentions
    output_attentions = False
@@ -129,26 +122,15 @@ class LlamaFlashAttention2(LlamaAttention):

    kv_seq_len = key_states.shape[-2]
    if past_key_value is not None:
-            kv_seq_len += past_key_value[0].shape[-2]
+        kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)

    cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)

-        if past_key_value is not None: # reuse k, v, self_attention
-            key_states = torch.cat([past_key_value[0], key_states], dim=2)
-            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+    if past_key_value is not None:
+        cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
+        key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)

-        past_key_value = (key_states, value_states) if use_cache else None
-
-        # cast to half precision
-        input_dtype = query_states.dtype
-        if input_dtype == torch.float32:
-            logger.warning_once("The input hidden states seems to be silently casted in float32.")
-            query_states = query_states.to(self.config.torch_dtype)
-            key_states = key_states.to(self.config.torch_dtype)
-            value_states = value_states.to(self.config.torch_dtype)
-
-        if getattr(self, "num_key_value_groups", None):
    key_states = repeat_kv(key_states, self.num_key_value_groups)
    value_states = repeat_kv(value_states, self.num_key_value_groups)

@@ -156,49 +138,50 @@ class LlamaFlashAttention2(LlamaAttention):
    key_states = key_states.transpose(1, 2)  # (bsz, seq_len, n_heads, head_dim)
    value_states = value_states.transpose(1, 2)  # (bsz, seq_len, n_heads, head_dim)

+    dropout_rate = self.attention_dropout if self.training else 0.0
+
+    input_dtype = query_states.dtype
+    if input_dtype == torch.float32:
+        if torch.is_autocast_enabled():
+            target_dtype = torch.get_autocast_gpu_dtype()
+        elif hasattr(self.config, "_pre_quantization_dtype"):
+            target_dtype = self.config._pre_quantization_dtype
+        else:
+            target_dtype = self.q_proj.weight.dtype
+
+        logger.warning_once("The input hidden states seems to be silently casted in float32.")
+        query_states = query_states.to(target_dtype)
+        key_states = key_states.to(target_dtype)
+        value_states = value_states.to(target_dtype)
+
    if getattr(self.config, "group_size_ratio", None) and self.training:  # shift
        groupsz = int(q_len * getattr(self.config, "group_size_ratio"))
        assert q_len % groupsz == 0, "q_len {} should be divisible by group size {}.".format(q_len, groupsz)
        num_groups = q_len // groupsz
+
        def shift(state: torch.Tensor) -> torch.Tensor:
-                state = torch.cat((
-                    state[:, :, :self.num_heads//2], state[:, :, self.num_heads//2:].roll(-groupsz//2, dims=1)
-                ), dim=2)
+            state = torch.cat(
+                (state[:, :, : self.num_heads // 2], state[:, :, self.num_heads // 2 :].roll(-groupsz // 2, dims=1)),
+                dim=2,
+            )
            return state.reshape(bsz * num_groups, groupsz, self.num_heads, self.head_dim)

        query_states, key_states, value_states = shift(query_states), shift(key_states), shift(value_states)
        if attention_mask is not None:
-                attention_mask = attention_mask.reshape(bsz * num_groups, groupsz)
+            attention_mask = attention_mask[:, :, :groupsz, :groupsz].repeat(num_groups, 1, 1, 1)

-        if attention_mask is not None:
-            logger.warning_once("Padded sequences are less efficient in FlashAttention.")
-            # -q_len: assumes left padding when q_len != kv_len
-            unpadded_q, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(query_states, attention_mask[:, -q_len:])
-            unpadded_k, _, cu_seqlens_k, max_seqlen_k = unpad_input(key_states, attention_mask)
-            unpadded_v, _, _, _ = unpad_input(value_states, attention_mask)
-            attn_output_unpad = flash_attn_varlen_func(
-                unpadded_q,
-                unpadded_k,
-                unpadded_v,
-                cu_seqlens_q=cu_seqlens_q,
-                cu_seqlens_k=cu_seqlens_k,
-                max_seqlen_q=max_seqlen_q,
-                max_seqlen_k=max_seqlen_k,
-                dropout_p=0.0,
-                softmax_scale=None,
-                causal=True,
-            )
-            attn_output = pad_input(attn_output_unpad, indices_q, bsz, q_len)
-        else:
-            attn_output = flash_attn_func(
-                query_states, key_states, value_states, 0.0, softmax_scale=None, causal=True
+    attn_output: torch.Tensor = self._flash_attention_forward(
+        query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
    )

    if getattr(self.config, "group_size_ratio", None) and self.training:  # shift back
        attn_output.reshape(bsz, q_len, self.num_heads, self.head_dim)
-            attn_output = torch.cat((
-                attn_output[:, :, :self.num_heads//2], attn_output[:, :, self.num_heads//2:].roll(groupsz//2, dims=1)
-            ))
+        attn_output = torch.cat(
+            (
+                attn_output[:, :, : self.num_heads // 2],
+                attn_output[:, :, self.num_heads // 2 :].roll(groupsz // 2, dims=1),
+            )
+        )

    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
    attn_output = self.o_proj(attn_output)
@@ -209,16 +192,6 @@ class LlamaFlashAttention2(LlamaAttention):
    return attn_output, attn_weights, past_key_value


-# Disable the transformation of the attention mask in LlamaModel as flash attention
-# takes a boolean padding_mask. Fills in the past kv length for use in forward.
-def _prepare_decoder_attention_mask(
-    self,
-    attention_mask: torch.Tensor,
-    input_shape: torch.Tensor,
-    inputs_embeds: torch.Tensor,
-    past_key_values_length: int
-) -> torch.Tensor:
-    if attention_mask is not None and torch.all(attention_mask):
-        return None  # This uses the faster call when training with full samples
-
-    return attention_mask
+def apply_llama_patch() -> None:
+    LlamaAttention.forward = llama_torch_attn_forward
+    LlamaFlashAttention2.forward = llama_flash_attn_forward
--- a/src/llmtuner/extras/ploting.py
+++ b/src/llmtuner/extras/ploting.py
@@ -1,11 +1,13 @@
-import os
-import math
 import json
+import math
+import os
 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
+from .logging import get_logger
+from .packages import is_matplotlib_available
+

 if is_matplotlib_available():
    import matplotlib.pyplot as plt
@@ -29,7 +31,6 @@ def smooth(scalars: List[float]) -> List[float]:


 def plot_loss(save_dictionary: os.PathLike, keys: Optional[List[str]] = ["loss"]) -> None:
-
    with open(os.path.join(save_dictionary, TRAINER_STATE_NAME), "r", encoding="utf-8") as f:
        data = json.load(f)

--- a/src/llmtuner/hparams/init.py
+++ b/src/llmtuner/hparams/init.py
@@ -3,3 +3,16 @@ from .evaluation_args import EvaluationArguments
 from .finetuning_args import FinetuningArguments
 from .generating_args import GeneratingArguments
 from .model_args import ModelArguments
+from .parser import get_eval_args, get_infer_args, get_train_args
+
+
+__all__ = [
+    "DataArguments",
+    "EvaluationArguments",
+    "FinetuningArguments",
+    "GeneratingArguments",
+    "ModelArguments",
+    "get_eval_args",
+    "get_infer_args",
+    "get_train_args",
+]
--- a/src/llmtuner/hparams/data_args.py
+++ b/src/llmtuner/hparams/data_args.py
@@ -1,30 +1,5 @@
-import os
-import json
-from typing import List, Literal, Optional
 from dataclasses import dataclass, field
-
-
-@dataclass
-class DatasetAttr:
-
-    load_from: str
-    dataset_name: Optional[str] = None
-    dataset_sha1: Optional[str] = None
-    system_prompt: Optional[str] = None
-    subset: Optional[str] = None
-    ranking: Optional[bool] = False
-    formatting: Optional[Literal["alpaca", "sharegpt"]] = "alpaca"
-
-    prompt: Optional[str] = "instruction"
-    query: Optional[str] = "input"
-    response: Optional[str] = "output"
-    history: Optional[str] = None
-    messages: Optional[str] = "conversations"
-    role: Optional[str] = "from"
-    content: Optional[str] = "value"
-
-    def __repr__(self) -> str:
-        return self.dataset_name
+from typing import Literal, Optional


@dataclass
@@ -33,84 +8,66 @@ class DataArguments:
    Arguments pertaining to what data we are going to input our model for training and evaluation.
    """
    template: Optional[str] = field(
-        default=None,
-        metadata={"help": "Which template to use for constructing prompts in training and inference."}
+        default=None, metadata={"help": "Which template to use for constructing prompts in training and inference."}
    )
    dataset: Optional[str] = field(
        default=None,
-        metadata={"help": "The name of provided dataset(s) to use. Use commas to separate multiple datasets."}
+        metadata={"help": "The name of provided dataset(s) to use. Use commas to separate multiple datasets."},
    )
    dataset_dir: Optional[str] = field(
-        default="data",
-        metadata={"help": "Path to the folder containing the datasets."}
+        default="data", metadata={"help": "Path to the folder containing the datasets."}
    )
    split: Optional[str] = field(
-        default="train",
-        metadata={"help": "Which dataset split to use for training and evaluation."}
+        default="train", metadata={"help": "Which dataset split to use for training and evaluation."}
    )
    cutoff_len: Optional[int] = field(
-        default=1024,
-        metadata={"help": "The maximum length of the model inputs after tokenization."}
+        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."}
+        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."}
-    )
-    streaming: Optional[bool] = field(
-        default=False,
-        metadata={"help": "Enable dataset streaming."}
+        default=False, metadata={"help": "Whether to disable the mask on the prompt or not."}
    )
+    streaming: Optional[bool] = field(default=False, metadata={"help": "Enable dataset streaming."})
    buffer_size: Optional[int] = field(
-        default=16384,
-        metadata={"help": "Size of the buffer to randomly sample examples from in dataset streaming."}
+        default=16384, metadata={"help": "Size of the buffer to randomly sample examples from in dataset streaming."}
    )
    mix_strategy: Optional[Literal["concat", "interleave_under", "interleave_over"]] = field(
        default="concat",
-        metadata={"help": "Strategy to use in dataset mixing (concat/interleave) (undersampling/oversampling)."}
+        metadata={"help": "Strategy to use in dataset mixing (concat/interleave) (undersampling/oversampling)."},
    )
    interleave_probs: Optional[str] = field(
        default=None,
-        metadata={"help": "Probabilities to sample data from datasets. Use commas to separate multiple datasets."}
+        metadata={"help": "Probabilities to sample data from datasets. Use commas to separate multiple datasets."},
    )
    overwrite_cache: Optional[bool] = field(
-        default=False,
-        metadata={"help": "Overwrite the cached training and evaluation sets."}
+        default=False, metadata={"help": "Overwrite the cached training and evaluation sets."}
    )
    preprocessing_num_workers: Optional[int] = field(
-        default=None,
-        metadata={"help": "The number of processes to use for the preprocessing."}
+        default=None, metadata={"help": "The number of processes to use for the preprocessing."}
    )
    max_samples: Optional[int] = field(
-        default=None,
-        metadata={"help": "For debugging purposes, truncate the number of examples for each dataset."}
+        default=None, metadata={"help": "For debugging purposes, truncate the number of examples for each dataset."}
    )
    eval_num_beams: Optional[int] = field(
        default=None,
-        metadata={"help": "Number of beams to use for evaluation. This argument will be passed to `model.generate`"}
+        metadata={"help": "Number of beams to use for evaluation. This argument will be passed to `model.generate`"},
    )
    ignore_pad_token_for_loss: Optional[bool] = field(
        default=True,
-        metadata={"help": "Whether to ignore the tokens corresponding to padded labels in the loss computation or not."}
-    )
-    system_prompt: Optional[str] = field(
-        default=None,
-        metadata={"help": "System prompt to add before the user query. Use `|` to separate multiple prompts in training."}
+        metadata={
+            "help": "Whether to ignore the tokens corresponding to padded labels in the loss computation or not."
+        },
    )
    val_size: Optional[float] = field(
-        default=0,
-        metadata={"help": "Size of the development set, should be an integer or a float in range `[0,1)`."}
+        default=0, metadata={"help": "Size of the development set, should be an integer or a float in range `[0,1)`."}
    )
    sft_packing: Optional[bool] = field(
-        default=False,
-        metadata={"help": "Packing the questions and answers in the supervised fine-tuning stage."}
+        default=False, metadata={"help": "Packing the questions and answers in the supervised fine-tuning stage."}
    )
    cache_path: Optional[str] = field(
-        default=None,
-        metadata={"help": "Path to save or load the preprocessed datasets."}
+        default=None, metadata={"help": "Path to save or load the preprocessed datasets."}
    )

    def __post_init__(self):
@@ -122,55 +79,3 @@ class DataArguments:

        if self.streaming and self.max_samples is not None:
            raise ValueError("`max_samples` is incompatible with `streaming`.")
-
-        if self.streaming and self.cache_path:
-            raise ValueError("`cache_path` is incompatible with `streaming`.")
-
-    def init_for_training(self, seed: int): # support mixing multiple datasets
-        self.seed = seed
-        dataset_names = [ds.strip() for ds in self.dataset.split(",")] if self.dataset is not None else []
-        try:
-            with open(os.path.join(self.dataset_dir, "dataset_info.json"), "r") as f:
-                dataset_info = json.load(f)
-        except Exception:
-            if self.dataset is not None:
-                raise ValueError("Cannot find dataset_info.json in `dataset_dir`.")
-            dataset_info = None
-
-        prompt_list = self.system_prompt.split("|") if self.system_prompt else [None]
-        prompt_list = prompt_list * (len(dataset_names) // len(prompt_list))
-        assert len(prompt_list) == len(dataset_names), "Number of system prompts should be equal to datasets or 1."
-
-        if self.interleave_probs is not None:
-            self.interleave_probs = [float(prob.strip()) for prob in self.interleave_probs.split(",")]
-
-        self.dataset_list: List[DatasetAttr] = []
-        for i, name in enumerate(dataset_names):
-            if name not in dataset_info:
-                raise ValueError("Undefined dataset {} in dataset_info.json.".format(name))
-
-            if "hf_hub_url" in dataset_info[name]:
-                dataset_attr = DatasetAttr("hf_hub", dataset_name=dataset_info[name]["hf_hub_url"])
-            elif "script_url" in dataset_info[name]:
-                dataset_attr = DatasetAttr("script", dataset_name=dataset_info[name]["script_url"])
-            else:
-                dataset_attr = DatasetAttr(
-                    "file",
-                    dataset_name=dataset_info[name]["file_name"],
-                    dataset_sha1=dataset_info[name].get("file_sha1", None)
-                )
-
-            if "columns" in dataset_info[name]:
-                dataset_attr.prompt = dataset_info[name]["columns"].get("prompt", None)
-                dataset_attr.query = dataset_info[name]["columns"].get("query", None)
-                dataset_attr.response = dataset_info[name]["columns"].get("response", None)
-                dataset_attr.history = dataset_info[name]["columns"].get("history", None)
-                dataset_attr.messages = dataset_info[name]["columns"].get("messages", None)
-                dataset_attr.role = dataset_info[name]["columns"].get("role", None)
-                dataset_attr.content = dataset_info[name]["columns"].get("content", None)
-
-            dataset_attr.subset = dataset_info[name].get("subset", None)
-            dataset_attr.ranking = dataset_info[name].get("ranking", False)
-            dataset_attr.formatting = dataset_info[name].get("formatting", "alpaca")
-            dataset_attr.system_prompt = prompt_list[i]
-            self.dataset_list.append(dataset_attr)
--- a/src/llmtuner/hparams/evaluation_args.py
+++ b/src/llmtuner/hparams/evaluation_args.py
@@ -1,6 +1,6 @@
 import os
-from typing import Literal, Optional
 from dataclasses import dataclass, field
+from typing import Literal, Optional

 from datasets import DownloadMode

@@ -10,46 +10,20 @@ class EvaluationArguments:
    r"""
    Arguments pertaining to specify the evaluation parameters.
    """
-    task: str = field(
-        metadata={"help": "Name of the evaluation task."}
-    )
+    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."}
+        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."}
+        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.")
--- a/src/llmtuner/hparams/finetuning_args.py
+++ b/src/llmtuner/hparams/finetuning_args.py
@@ -1,6 +1,6 @@
 import json
-from typing import Literal, Optional
 from dataclasses import asdict, dataclass, field
+from typing import Literal, Optional


@dataclass
@@ -8,19 +8,20 @@ class FreezeArguments:
    r"""
    Arguments pertaining to the freeze (partial-parameter) training.
    """
-    num_layer_trainable: Optional[int] = field(
-        default=3,
-        metadata={"help": "Number of trainable layers for partial-parameter (freeze) fine-tuning."}
-    )
    name_module_trainable: Optional[str] = field(
        default="mlp",
-        metadata={"help": "Name of trainable modules for partial-parameter (freeze) fine-tuning. \
+        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."}
+                  LLaMA choices: ["mlp", "self_attn"], \
+                  BLOOM & Falcon & ChatGLM choices: ["mlp", "self_attention"], \
+                  Qwen choices: ["mlp", "attn"], \
+                  Phi 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."}
    )


@@ -29,35 +30,34 @@ class LoraArguments:
    r"""
    Arguments pertaining to the LoRA training.
    """
-    lora_rank: Optional[int] = field(
-        default=8,
-        metadata={"help": "The intrinsic dimension for LoRA fine-tuning."}
-    )
-    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_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."}
-    )
    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."}
+        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."}
+    lora_alpha: Optional[int] = field(
+        default=None, metadata={"help": "The scale factor for LoRA fine-tuning (default: lora_rank * 2)."}
+    )
+    lora_dropout: Optional[float] = field(default=0.0, 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 choices: ["Wqkv", "out_proj", "fc1", "fc2"], \
+                  Others choices: the same as LLaMA.'
+        },
+    )
+    lora_bf16_mode: Optional[bool] = field(
+        default=False, metadata={"help": "Whether or not to train lora adapters in bf16 precision."}
+    )
+    create_new_adapter: Optional[bool] = field(
+        default=False, metadata={"help": "Whether or not to create a new adapter with randomly initialized weight."}
    )


@@ -66,53 +66,53 @@ 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."}
+    dpo_beta: Optional[float] = field(default=0.1, metadata={"help": "The beta parameter for the DPO loss."})
+    dpo_loss: Optional[Literal["sigmoid", "hinge", "ipo", "kto"]] = field(
+        default="sigmoid", metadata={"help": "The type of DPO loss to use."}
+    )
+    dpo_ftx: Optional[float] = field(
+        default=0, metadata={"help": "The supervised fine-tuning loss coefficient in DPO training."}
+    )
+    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."}
+        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."}
+        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."}
+        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."}
+        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."}
+        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_adapters: Optional[str] = field(
+        default=None, metadata={"help": "Path to the adapters of the reference model."}
    )
    ref_model_quantization_bit: Optional[int] = field(
-        default=None,
-        metadata={"help": "The number of bits to quantize the reference model."}
+        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."}
+        default=None, metadata={"help": "Path to the reward model used for the PPO training."}
    )
-    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_adapters: Optional[str] = field(
+        default=None, metadata={"help": "Path to the adapters of the reward model."}
    )
    reward_model_quantization_bit: Optional[int] = field(
-        default=None,
-        metadata={"help": "The number of bits to quantize the reward model."}
+        default=None, metadata={"help": "The number of bits to quantize the reward model."}
    )
-    reward_model_type: Optional[Literal["lora", "full"]] = field(
+    reward_model_type: Optional[Literal["lora", "full", "api"]] = field(
        default="lora",
-        metadata={"help": "The checkpoint type of the reward model. The lora type only supports lora training."}
+        metadata={"help": "The type of the reward model in PPO training. Lora model only supports lora training."},
    )


@@ -122,28 +122,13 @@ class FinetuningArguments(FreezeArguments, LoraArguments, RLHFArguments):
    Arguments pertaining to which techniques we are going to fine-tuning with.
    """
    stage: Optional[Literal["pt", "sft", "rm", "ppo", "dpo"]] = field(
-        default="sft",
-        metadata={"help": "Which stage will be performed in training."}
+        default="sft", metadata={"help": "Which stage will be performed in training."}
    )
    finetuning_type: Optional[Literal["lora", "freeze", "full"]] = field(
-        default="lora",
-        metadata={"help": "Which fine-tuning method to use."}
-    )
-    upcast_layernorm: Optional[bool] = field(
-        default=False,
-        metadata={"help": "Whether to upcast the layernorm weights in fp32."}
-    )
-    neft_alpha: Optional[float] = field(
-        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."}
+        default="lora", metadata={"help": "Which fine-tuning method to use."}
    )
    plot_loss: Optional[bool] = field(
-        default=False,
-        metadata={"help": "Whether to plot the training loss after fine-tuning or not."}
+        default=False, metadata={"help": "Whether or not to save the training loss curves."}
    )

    def __post_init__(self):
@@ -153,11 +138,9 @@ class FinetuningArguments(FreezeArguments, LoraArguments, RLHFArguments):
            return arg

        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_alpha = self.lora_alpha or self.lora_rank * 2
        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."
@@ -167,7 +150,7 @@ class FinetuningArguments(FreezeArguments, LoraArguments, RLHFArguments):
            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.")
+            raise ValueError("Freeze/Full PPO training needs `reward_model_type=full`.")

    def save_to_json(self, json_path: str):
        r"""Saves the content of this instance in JSON format inside `json_path`."""
--- a/src/llmtuner/hparams/generating_args.py
+++ b/src/llmtuner/hparams/generating_args.py
@@ -1,5 +1,5 @@
-from typing import Any, Dict, Optional
 from dataclasses import asdict, dataclass, field
+from typing import Any, Dict, Optional


@dataclass
@@ -8,40 +8,37 @@ class GeneratingArguments:
    Arguments pertaining to specify the decoding parameters.
    """
    do_sample: Optional[bool] = field(
-        default=True,
-        metadata={"help": "Whether or not to use sampling, use greedy decoding otherwise."}
+        default=True, metadata={"help": "Whether or not to use sampling, use greedy decoding otherwise."}
    )
    temperature: Optional[float] = field(
-        default=0.95,
-        metadata={"help": "The value used to modulate the next token probabilities."}
+        default=0.95, metadata={"help": "The value used to modulate the next token probabilities."}
    )
    top_p: Optional[float] = field(
        default=0.7,
-        metadata={"help": "The smallest set of most probable tokens with probabilities that add up to top_p or higher are kept."}
+        metadata={
+            "help": "The smallest set of most probable tokens with probabilities that add up to top_p or higher are kept."
+        },
    )
    top_k: Optional[int] = field(
        default=50,
-        metadata={"help": "The number of highest probability vocabulary tokens to keep for top-k filtering."}
+        metadata={"help": "The number of highest probability vocabulary tokens to keep for top-k filtering."},
    )
    num_beams: Optional[int] = field(
-        default=1,
-        metadata={"help": "Number of beams for beam search. 1 means no beam search."}
+        default=1, metadata={"help": "Number of beams for beam search. 1 means no beam search."}
    )
    max_length: Optional[int] = field(
        default=512,
-        metadata={"help": "The maximum length the generated tokens can have. It can be overridden by max_new_tokens."}
+        metadata={"help": "The maximum length the generated tokens can have. It can be overridden by max_new_tokens."},
    )
    max_new_tokens: Optional[int] = field(
        default=512,
-        metadata={"help": "The maximum numbers of tokens to generate, ignoring the number of tokens in the prompt."}
+        metadata={"help": "The maximum numbers of tokens to generate, ignoring the number of tokens in the prompt."},
    )
    repetition_penalty: Optional[float] = field(
-        default=1.0,
-        metadata={"help": "The parameter for repetition penalty. 1.0 means no penalty."}
+        default=1.0, metadata={"help": "The parameter for repetition penalty. 1.0 means no penalty."}
    )
    length_penalty: Optional[float] = field(
-        default=1.0,
-        metadata={"help": "Exponential penalty to the length that is used with beam-based generation."}
+        default=1.0, metadata={"help": "Exponential penalty to the length that is used with beam-based generation."}
    )

    def to_dict(self) -> Dict[str, Any]:
--- a/src/llmtuner/hparams/model_args.py
+++ b/src/llmtuner/hparams/model_args.py
@@ -1,5 +1,5 @@
-from typing import Any, Dict, Literal, Optional
 from dataclasses import asdict, dataclass, field
+from typing import Any, Dict, Literal, Optional


@dataclass
@@ -8,55 +8,85 @@ class ModelArguments:
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune.
    """
    model_name_or_path: str = field(
-        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models."}
+        metadata={"help": "Path to the model weight or identifier from huggingface.co/models or modelscope.cn/models."}
+    )
+    adapter_name_or_path: Optional[str] = field(
+        default=None, metadata={"help": "Path to the adapter weight or identifier from huggingface.co/models."}
    )
    cache_dir: Optional[str] = field(
        default=None,
-        metadata={"help": "Where to store the pretrained models downloaded from huggingface.co."}
+        metadata={"help": "Where to store the pre-trained models downloaded from huggingface.co or modelscope.cn."},
    )
    use_fast_tokenizer: Optional[bool] = field(
-        default=True,
-        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."}
+        default=False,
+        metadata={"help": "Whether or not to use one of the fast tokenizer (backed by the tokenizers library)."},
+    )
+    resize_vocab: Optional[bool] = field(
+        default=False, metadata={"help": "Whether or not to resize the tokenizer vocab and the embedding layers."}
    )
    split_special_tokens: Optional[bool] = field(
        default=False,
-        metadata={"help": "Whether or not the special tokens should be split during the tokenization process."}
+        metadata={"help": "Whether or not the special tokens should be split during the tokenization process."},
    )
    model_revision: Optional[str] = field(
        default="main",
-        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."}
+        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
    )
    quantization_bit: Optional[int] = field(
-        default=None,
-        metadata={"help": "The number of bits to quantize the model."}
+        default=None, metadata={"help": "The number of bits to quantize the model."}
    )
    quantization_type: Optional[Literal["fp4", "nf4"]] = field(
-        default="nf4",
-        metadata={"help": "Quantization data type to use in int4 training."}
+        default="nf4", metadata={"help": "Quantization data type to use in int4 training."}
    )
    double_quantization: Optional[bool] = field(
-        default=True,
-        metadata={"help": "Whether to use double quantization in int4 training or not."}
+        default=True, metadata={"help": "Whether or not to use double quantization in int4 training."}
    )
    rope_scaling: Optional[Literal["linear", "dynamic"]] = field(
-        default=None,
-        metadata={"help": "Adopt scaled rotary positional embeddings."}
-    )
-    checkpoint_dir: Optional[str] = field(
-        default=None,
-        metadata={"help": "Path to the directory(s) containing the model checkpoints as well as the configurations."}
+        default=None, metadata={"help": "Which scaling strategy should be adopted for the RoPE embeddings."}
    )
    flash_attn: Optional[bool] = field(
-        default=False,
-        metadata={"help": "Enable FlashAttention-2 for faster training."}
+        default=False, metadata={"help": "Enable FlashAttention-2 for faster training."}
    )
    shift_attn: Optional[bool] = field(
-        default=False,
-        metadata={"help": "Enable shift short attention (S^2-Attn) proposed by LongLoRA."}
+        default=False, metadata={"help": "Enable shift short attention (S^2-Attn) proposed by LongLoRA."}
    )
-    hf_hub_token: Optional[str] = field(
-        default=None,
-        metadata={"help": "Auth token to log in with Hugging Face Hub."}
+    use_unsloth: Optional[bool] = field(
+        default=False, metadata={"help": "Whether or not to use unsloth's optimization for the LoRA training."}
+    )
+    disable_gradient_checkpointing: Optional[bool] = field(
+        default=False, metadata={"help": "Whether or not to disable gradient checkpointing."}
+    )
+    upcast_layernorm: Optional[bool] = field(
+        default=False, metadata={"help": "Whether or not to upcast the layernorm weights in fp32."}
+    )
+    upcast_lmhead_output: Optional[bool] = field(
+        default=False, metadata={"help": "Whether or not to upcast the output of lm_head in fp32."}
+    )
+    hf_hub_token: Optional[str] = field(default=None, metadata={"help": "Auth token to log in with Hugging Face Hub."})
+    ms_hub_token: Optional[str] = field(default=None, metadata={"help": "Auth token to log in with ModelScope Hub."})
+    export_dir: Optional[str] = field(
+        default=None, metadata={"help": "Path to the directory to save the exported model."}
+    )
+    export_size: Optional[int] = field(
+        default=1, metadata={"help": "The file shard size (in GB) of the exported model."}
+    )
+    export_quantization_bit: Optional[int] = field(
+        default=None, metadata={"help": "The number of bits to quantize the exported model."}
+    )
+    export_quantization_dataset: Optional[str] = field(
+        default=None, metadata={"help": "Path to the dataset or dataset name to use in quantizing the exported model."}
+    )
+    export_quantization_nsamples: Optional[int] = field(
+        default=128, metadata={"help": "The number of samples used for quantization."}
+    )
+    export_quantization_maxlen: Optional[int] = field(
+        default=1024, metadata={"help": "The maximum length of the model inputs used for quantization."}
+    )
+    export_legacy_format: Optional[bool] = field(
+        default=False, metadata={"help": "Whether or not to save the `.bin` files instead of `.safetensors`."}
+    )
+    export_hub_model_id: Optional[str] = field(
+        default=None, metadata={"help": "The name of the repository if push the model to the Hugging Face hub."}
    )

    def __post_init__(self):
@@ -66,10 +96,14 @@ class ModelArguments:
        if self.split_special_tokens and self.use_fast_tokenizer:
            raise ValueError("`split_special_tokens` is only supported for slow tokenizers.")

-        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.adapter_name_or_path is not None:  # support merging multiple lora weights
+            self.adapter_name_or_path = [path.strip() for path in self.adapter_name_or_path.split(",")]

        assert self.quantization_bit in [None, 8, 4], "We only accept 4-bit or 8-bit quantization."
+        assert self.export_quantization_bit in [None, 8, 4, 3, 2], "We only accept 2/3/4/8-bit quantization."
+
+        if self.export_quantization_bit is not None and self.export_quantization_dataset is None:
+            raise ValueError("Quantization dataset is necessary for exporting.")

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)
--- a/src/llmtuner/hparams/parser.py
+++ b/src/llmtuner/hparams/parser.py
@@ -1,77 +1,99 @@
+import logging
 import os
-import torch
-import datasets
-import transformers
+import sys
 from typing import Any, Dict, Optional, Tuple
+
+import datasets
+import torch
+import transformers
 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
-)
+from ..extras.logging import get_logger
+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


 logger = get_logger(__name__)


-_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
-]
+_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) -> _TRAIN_CLS:
-    parser = HfArgumentParser(_TRAIN_ARGS)
-    return parse_args(parser, args)
+def _parse_args(parser: "HfArgumentParser", args: Optional[Dict[str, Any]] = None) -> Tuple[Any]:
+    if args is not None:
+        return parser.parse_dict(args)
+
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"):
+        return parser.parse_yaml_file(os.path.abspath(sys.argv[1]))
+
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        return parser.parse_json_file(os.path.abspath(sys.argv[1]))
+
+    (*parsed_args, unknown_args) = parser.parse_args_into_dataclasses(return_remaining_strings=True)
+
+    if unknown_args:
+        print(parser.format_help())
+        print("Got unknown args, potentially deprecated arguments: {}".format(unknown_args))
+        raise ValueError("Some specified arguments are not used by the HfArgumentParser: {}".format(unknown_args))
+
+    return (*parsed_args,)


-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
-    if training_args.should_log:
-        # The default of training_args.log_level is passive, so we set log level at info here to have that default.
-        transformers.utils.logging.set_verbosity_info()
-
-    log_level = training_args.get_process_log_level()
+def _set_transformers_logging(log_level: Optional[int] = logging.INFO) -> None:
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

-    # Check arguments
-    data_args.init_for_training(training_args.seed)

+def _verify_model_args(model_args: "ModelArguments", finetuning_args: "FinetuningArguments") -> None:
+    if model_args.quantization_bit is not None:
+        if finetuning_args.finetuning_type != "lora":
+            raise ValueError("Quantization is only compatible with the LoRA method.")
+
+        if finetuning_args.create_new_adapter:
+            raise ValueError("Cannot create new adapter upon a quantized model.")
+
+    if model_args.adapter_name_or_path is not None and len(model_args.adapter_name_or_path) != 1:
+        if finetuning_args.finetuning_type != "lora":
+            raise ValueError("Multiple adapters are only available for LoRA tuning.")
+
+        if model_args.quantization_bit is not None:
+            raise ValueError("Quantized model only accepts a single adapter. Merge them first.")
+
+
+def _parse_train_args(args: Optional[Dict[str, Any]] = None) -> _TRAIN_CLS:
+    parser = HfArgumentParser(_TRAIN_ARGS)
+    return _parse_args(parser, args)
+
+
+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
+    if training_args.should_log:
+        _set_transformers_logging()
+
+    # Check arguments
    if finetuning_args.stage != "pt" and data_args.template is None:
        raise ValueError("Please specify which `template` to use.")

@@ -81,23 +103,18 @@ def get_train_args(args: Optional[Dict[str, Any]] = None) -> _TRAIN_CLS:
    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 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:
+    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, 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.shift_attn:
        raise ValueError("PPO training is incompatible with S^2-Attn.")

+    if finetuning_args.stage == "ppo" and finetuning_args.reward_model_type == "lora" and model_args.use_unsloth:
+        raise ValueError("Unsloth does not support lora reward model.")
+
    if training_args.max_steps == -1 and data_args.streaming:
        raise ValueError("Please specify `max_steps` in streaming mode.")

@@ -107,17 +124,9 @@ def get_train_args(args: Optional[Dict[str, Any]] = None) -> _TRAIN_CLS:
    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.")
+    _verify_model_args(model_args, finetuning_args)

-    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.")
-
-    if training_args.do_train and model_args.quantization_bit is not None and (not finetuning_args.upcast_layernorm):
+    if training_args.do_train and model_args.quantization_bit is not None and (not model_args.upcast_layernorm):
        logger.warning("We recommend enable `upcast_layernorm` in quantized training.")

    if training_args.do_train and (not training_args.fp16) and (not training_args.bf16):
@@ -140,11 +149,18 @@ def get_train_args(args: Optional[Dict[str, Any]] = None) -> _TRAIN_CLS:
        training_args_dict.update(dict(ddp_find_unused_parameters=False))
        training_args = Seq2SeqTrainingArguments(**training_args_dict)

+    if finetuning_args.stage in ["rm", "ppo"] and finetuning_args.finetuning_type in ["full", "freeze"]:
+        can_resume_from_checkpoint = False
+        training_args.resume_from_checkpoint = None
+    else:
+        can_resume_from_checkpoint = True
+
    if (
        training_args.resume_from_checkpoint is None
        and training_args.do_train
        and os.path.isdir(training_args.output_dir)
        and not training_args.overwrite_output_dir
+        and can_resume_from_checkpoint
    ):
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
@@ -155,7 +171,20 @@ def get_train_args(args: Optional[Dict[str, Any]] = None) -> _TRAIN_CLS:
            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."
+                "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 finetuning_args.finetuning_type == "lora"
+        and training_args.resume_from_checkpoint is not None
+    ):
+        logger.warning(
+            "Add {} to `adapter_name_or_path` to resume training from checkpoint.".format(
+                training_args.resume_from_checkpoint
+            )
        )

    # postprocess model_args
@@ -165,10 +194,15 @@ def get_train_args(args: Optional[Dict[str, Any]] = None) -> _TRAIN_CLS:
    model_args.model_max_length = data_args.cutoff_len

    # Log on each process the small summary:
-    logger.info("Process rank: {}, device: {}, n_gpu: {}\n  distributed training: {}, compute dtype: {}".format(
-        training_args.local_rank, training_args.device, training_args.n_gpu,
-        bool(training_args.local_rank != -1), str(model_args.compute_dtype)
-    ))
+    logger.info(
+        "Process rank: {}, device: {}, n_gpu: {}\n  distributed training: {}, compute dtype: {}".format(
+            training_args.local_rank,
+            training_args.device,
+            training_args.n_gpu,
+            bool(training_args.local_rank != -1),
+            str(model_args.compute_dtype),
+        )
+    )
    logger.info(f"Training/evaluation parameters {training_args}")

    # Set seed before initializing model.
@@ -178,32 +212,25 @@ def get_train_args(args: Optional[Dict[str, Any]] = None) -> _TRAIN_CLS:


 def get_infer_args(args: Optional[Dict[str, Any]] = None) -> _INFER_CLS:
-    model_args, data_args, finetuning_args, generating_args = parse_infer_args(args)
+    model_args, data_args, finetuning_args, generating_args = _parse_infer_args(args)
+    _set_transformers_logging()

    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)
+    model_args, data_args, eval_args, finetuning_args = _parse_eval_args(args)
+    _set_transformers_logging()

    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.")
+    _verify_model_args(model_args, finetuning_args)

    transformers.set_seed(eval_args.seed)

--- a/src/llmtuner/model/init.py
+++ b/src/llmtuner/model/init.py
@@ -1,5 +1,5 @@
-# Level: loader > adapter > parser, utils
+from .loader import load_model_and_tokenizer
+from .utils import dispatch_model, get_modelcard_args, load_valuehead_params

-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, generate_model_card, load_valuehead_params
+
+__all__ = ["load_model_and_tokenizer", "dispatch_model", "get_modelcard_args", "load_valuehead_params"]
--- a/src/llmtuner/model/adapter.py
+++ b/src/llmtuner/model/adapter.py
@@ -1,23 +1,25 @@
-import torch
+import inspect
 from typing import TYPE_CHECKING
-from peft import PeftModel, TaskType, LoraConfig, get_peft_model

-from llmtuner.extras.logging import get_logger
-from llmtuner.model.utils import find_all_linear_modules
+import torch
+from peft import LoraConfig, PeftModel, TaskType, get_peft_model
+from transformers.integrations import is_deepspeed_zero3_enabled
+
+from ..extras.logging import get_logger
+from .utils import find_all_linear_modules
+

 if TYPE_CHECKING:
    from transformers.modeling_utils import PreTrainedModel
-    from llmtuner.hparams import ModelArguments, FinetuningArguments
+
+    from ..hparams import FinetuningArguments, ModelArguments


 logger = get_logger(__name__)


 def init_adapter(
-    model: "PreTrainedModel",
-    model_args: "ModelArguments",
-    finetuning_args: "FinetuningArguments",
-    is_trainable: bool
+    model: "PreTrainedModel", model_args: "ModelArguments", finetuning_args: "FinetuningArguments", is_trainable: bool
 ) -> "PreTrainedModel":
    r"""
    Initializes the adapters.
@@ -27,8 +29,8 @@ def init_adapter(
    Note that the trainable parameters must be cast to float32.
    """

-    if (not is_trainable) and model_args.checkpoint_dir is None:
-        logger.info("Checkpoint is not found at evaluation, load the original model.")
+    if (not is_trainable) and model_args.adapter_name_or_path is None:
+        logger.info("Adapter is not found at evaluation, load the base model.")
        return model

    if finetuning_args.finetuning_type == "full" and is_trainable:
@@ -44,10 +46,11 @@ def init_adapter(
        )
        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_layer_ids = [k for k in range(-finetuning_args.num_layer_trainable)]  # noqa: C416

        trainable_layers = []
        for module_name in finetuning_args.name_module_trainable:
@@ -62,42 +65,74 @@ def init_adapter(

    if finetuning_args.finetuning_type == "lora":
        logger.info("Fine-tuning method: LoRA")
-        checkpoint_to_resume = None
+        adapter_to_resume = None

-        if model_args.checkpoint_dir is not None:
-            if is_trainable and finetuning_args.resume_lora_training:
-                checkpoints_to_merge, checkpoint_to_resume = model_args.checkpoint_dir[:-1], model_args.checkpoint_dir[-1]
+        if model_args.adapter_name_or_path is not None:
+            is_mergeable = True
+            if getattr(model, "quantization_method", None):  # merge lora in quantized model is unstable
+                assert len(model_args.adapter_name_or_path) == 1, "Quantized model only accepts a single adapter."
+                is_mergeable = False
+
+            if is_deepspeed_zero3_enabled():
+                assert len(model_args.adapter_name_or_path) == 1, "Cannot use multiple adapters in DeepSpeed ZeRO-3."
+                is_mergeable = False
+
+            if (is_trainable and not finetuning_args.create_new_adapter) or (not is_mergeable):
+                adapter_to_merge = model_args.adapter_name_or_path[:-1]
+                adapter_to_resume = model_args.adapter_name_or_path[-1]
            else:
-                checkpoints_to_merge = model_args.checkpoint_dir
+                adapter_to_merge = model_args.adapter_name_or_path

-            for checkpoint in checkpoints_to_merge:
-                model = PeftModel.from_pretrained(model, checkpoint)
+            for adapter in adapter_to_merge:
+                model = PeftModel.from_pretrained(model, adapter)
                model = model.merge_and_unload()

-            if len(checkpoints_to_merge) > 0:
-                logger.info("Merged {} model checkpoint(s).".format(len(checkpoints_to_merge)))
+            if len(adapter_to_merge) > 0:
+                logger.info("Merged {} adapter(s).".format(len(adapter_to_merge)))

-            if checkpoint_to_resume is not None: # resume lora training
-                model = PeftModel.from_pretrained(model, checkpoint_to_resume, is_trainable=is_trainable)
+            if adapter_to_resume is not None:  # resume lora training
+                model = PeftModel.from_pretrained(model, adapter_to_resume, is_trainable=is_trainable)

-        if is_trainable and checkpoint_to_resume is None: # create new lora weights while training
+        if is_trainable and adapter_to_resume is None:  # create new lora weights while training
            if len(finetuning_args.lora_target) == 1 and finetuning_args.lora_target[0] == "all":
-                target_modules = find_all_linear_modules(model, model_args.quantization_bit)
+                target_modules = find_all_linear_modules(model)
            else:
                target_modules = finetuning_args.lora_target

+            peft_kwargs = {
+                "r": finetuning_args.lora_rank,
+                "target_modules": target_modules,
+                "lora_alpha": finetuning_args.lora_alpha,
+                "lora_dropout": finetuning_args.lora_dropout,
+            }
+
+            if model_args.use_unsloth:
+                from unsloth import FastLlamaModel, FastMistralModel  # type: ignore
+
+                unsloth_peft_kwargs = {"model": model, "max_seq_length": model_args.model_max_length}
+                if "loftq_config" in inspect.signature(FastLlamaModel.get_peft_model).parameters:
+                    unsloth_peft_kwargs["loftq_config"] = {}
+
+                if getattr(model.config, "model_type", None) == "llama":
+                    model = FastLlamaModel.get_peft_model(**peft_kwargs, **unsloth_peft_kwargs)
+                elif getattr(model.config, "model_type", None) == "mistral":
+                    model = FastMistralModel.get_peft_model(**peft_kwargs, **unsloth_peft_kwargs)
+                else:
+                    raise NotImplementedError
+
+            else:
                lora_config = LoraConfig(
                    task_type=TaskType.CAUSAL_LM,
                    inference_mode=False,
-                r=finetuning_args.lora_rank,
-                lora_alpha=finetuning_args.lora_alpha,
-                lora_dropout=finetuning_args.lora_dropout,
-                target_modules=target_modules,
-                modules_to_save=finetuning_args.additional_target
+                    modules_to_save=finetuning_args.additional_target,
+                    **peft_kwargs,
                )
                model = get_peft_model(model, lora_config)

-    if model_args.checkpoint_dir is not None:
-        logger.info("Loaded fine-tuned model from checkpoint(s): {}".format(",".join(model_args.checkpoint_dir)))
+        for param in filter(lambda p: p.requires_grad, model.parameters()):
+            param.data = param.data.to(torch.bfloat16 if finetuning_args.lora_bf16_mode else torch.float32)
+
+    if model_args.adapter_name_or_path is not None:
+        logger.info("Loaded adapter(s): {}".format(",".join(model_args.adapter_name_or_path)))

    return model
--- a/src/llmtuner/model/loader.py
+++ b/src/llmtuner/model/loader.py
@@ -1,224 +1,133 @@
-import os
-import math
-import torch
-from types import MethodType
-from typing import TYPE_CHECKING, Literal, Optional, Tuple
+from typing import TYPE_CHECKING, Optional, Tuple

-from transformers import (
-    AutoConfig,
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    BitsAndBytesConfig,
-    PretrainedConfig,
-    PreTrainedModel,
-    PreTrainedTokenizerBase
-)
-from transformers.models.llama import modeling_llama as LlamaModule
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from transformers.integrations import is_deepspeed_zero3_enabled
 from transformers.utils.versions import require_version
 from trl import AutoModelForCausalLMWithValueHead

-try:
-    from transformers.integrations import is_deepspeed_zero3_enabled
-except ImportError: # https://github.com/huggingface/transformers/releases/tag/v4.33.1
-    from transformers.deepspeed import is_deepspeed_zero3_enabled
+from ..extras.logging import get_logger
+from ..extras.misc import count_parameters, get_current_device, try_download_model_from_ms
+from .adapter import init_adapter
+from .patcher import patch_config, patch_model, patch_tokenizer, patch_valuehead_model
+from .utils import load_valuehead_params, register_autoclass

-from llmtuner.extras.logging import reset_logging, get_logger
-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.model.adapter import init_adapter
-from llmtuner.model.utils import load_valuehead_params, prepare_model_for_training

 if TYPE_CHECKING:
-    from transformers import PreTrainedTokenizer
-    from llmtuner.hparams import ModelArguments
+    from transformers import PreTrainedModel, PreTrainedTokenizer
+
+    from ..hparams import FinetuningArguments, ModelArguments


 logger = get_logger(__name__)


-require_version("transformers>=4.31.0,<4.35.0", "To fix: pip install \"transformers>=4.31.0,<4.35.0\"")
-require_version("datasets>=2.14.0", "To fix: pip install datasets>=2.14.0")
+require_version("transformers>=4.36.2", "To fix: pip install transformers>=4.36.2")
+require_version("datasets>=2.14.3", "To fix: pip install datasets>=2.14.3")
 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.4", "To fix: pip install trl>=0.7.4")
+require_version("peft>=0.7.0", "To fix: pip install peft>=0.7.0")
+require_version("trl>=0.7.6", "To fix: pip install trl>=0.7.6")


 def load_model_and_tokenizer(
    model_args: "ModelArguments",
    finetuning_args: "FinetuningArguments",
    is_trainable: Optional[bool] = False,
-    stage: Optional[Literal["pt", "sft", "rm", "ppo"]] = "sft"
-) -> Tuple[PreTrainedModel, "PreTrainedTokenizer"]:
+    add_valuehead: Optional[bool] = False,
+) -> Tuple["PreTrainedModel", "PreTrainedTokenizer"]:
    r"""
    Loads pretrained model and tokenizer.

    Support both training and inference.
    """

+    try_download_model_from_ms(model_args)
+
    config_kwargs = {
        "trust_remote_code": True,
        "cache_dir": model_args.cache_dir,
        "revision": model_args.model_revision,
-        "token": model_args.hf_hub_token
+        "token": model_args.hf_hub_token,
    }

    tokenizer = AutoTokenizer.from_pretrained(
        model_args.model_name_or_path,
        use_fast=model_args.use_fast_tokenizer,
        split_special_tokens=model_args.split_special_tokens,
-        padding_side="right", # training with left-padded tensors in fp16 precision may cause overflow
-        **config_kwargs
+        padding_side="right",
+        **config_kwargs,
    )
+    patch_tokenizer(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
+    config = AutoConfig.from_pretrained(model_args.model_name_or_path, **config_kwargs)
+    patch_config(config, tokenizer, model_args, config_kwargs, is_trainable)

-    config = AutoConfig.from_pretrained(model_to_load, **config_kwargs)
+    model = None
+    if is_trainable and model_args.use_unsloth:
+        require_version("unsloth", "Follow the instructions at: https://github.com/unslothai/unsloth")
+        from unsloth import FastLlamaModel, FastMistralModel  # type: ignore

-    # Fix tokenizer (for ChatGLM2 and ChatGLM3)
-    if getattr(config, "model_type", None) == "chatglm":
-        tokenizer._pad = MethodType(PreTrainedTokenizerBase._pad, tokenizer)
-
-    # Set model dtype
-    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":
-        for dtype_name, dtype in [("fp16", torch.float16), ("bf16", torch.bfloat16), ("fp32", torch.float32)]:
-            setattr(config, dtype_name, getattr(config, "torch_dtype", None) == dtype)
-
-    # Set RoPE scaling
-    if model_args.rope_scaling is not None:
-        if not hasattr(config, "rope_scaling"):
-            logger.warning("Current model does not support RoPE scaling.")
-        else:
-            if is_trainable:
-                if model_args.rope_scaling == "dynamic":
-                    logger.warning(
-                        "Dynamic NTK may not work well with fine-tuning. "
-                        "See: https://github.com/huggingface/transformers/pull/24653"
-                    )
-
-                current_max_length = getattr(config, "max_position_embeddings", None)
-                if current_max_length and model_args.model_max_length > current_max_length:
-                    scaling_factor = float(math.ceil(model_args.model_max_length / current_max_length))
-                else:
-                    logger.warning("Input length is smaller than max length. Consider increase input length.")
-                    scaling_factor = 1.0
-            else:
-                scaling_factor = 2.0
-
-            setattr(config, "rope_scaling", {"type": model_args.rope_scaling, "factor": scaling_factor})
-            logger.info("Using {} scaling strategy and setting scaling factor to {}".format(
-                model_args.rope_scaling, scaling_factor
-            ))
-
-    # Set FlashAttention-2
-    if model_args.flash_attn:
+        unsloth_kwargs = {
+            "model_name": model_args.model_name_or_path,
+            "max_seq_length": model_args.model_max_length,
+            "dtype": model_args.compute_dtype,
+            "load_in_4bit": model_args.quantization_bit == 4,
+            "token": model_args.hf_hub_token,
+            "device_map": get_current_device(),
+            "rope_scaling": getattr(config, "rope_scaling", None),
+        }
        if getattr(config, "model_type", None) == "llama":
-            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.")
+            model, _ = FastLlamaModel.from_pretrained(**unsloth_kwargs)
+        elif getattr(config, "model_type", None) == "mistral":
+            model, _ = FastMistralModel.from_pretrained(**unsloth_kwargs)
        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.")
-    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.")
+            logger.warning("Unsloth does not support model type {}.".format(getattr(config, "model_type", None)))
+            model_args.use_unsloth = False

-    # Set shift short attention (S^2-Attn)
-    if is_trainable and model_args.shift_attn:
-        if getattr(config, "model_type", None) == "llama":
-            setattr(config, "group_size_ratio", 0.25)
-            logger.info("Using shift short attention with group_size_ratio=1/4.")
-        else:
-            logger.warning("Current model does not support shift short attention.")
+        if model_args.adapter_name_or_path:
+            model_args.adapter_name_or_path = None
+            logger.warning("Unsloth does not support loading adapters.")

-    # 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.")
-
-        if model_args.quantization_bit == 8:
-            require_version("bitsandbytes>=0.37.0", "To fix: pip install bitsandbytes>=0.37.0")
-            config_kwargs["load_in_8bit"] = True
-            config_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
-
-        if model_args.quantization_bit == 4:
-            require_version("bitsandbytes>=0.39.0", "To fix: pip install bitsandbytes>=0.39.0")
-            config_kwargs["load_in_4bit"] = True
-            config_kwargs["quantization_config"] = BitsAndBytesConfig(
-                load_in_4bit=True,
-                bnb_4bit_compute_dtype=model_args.compute_dtype,
-                bnb_4bit_use_double_quant=model_args.double_quantization,
-                bnb_4bit_quant_type=model_args.quantization_type
-            )
-
-        config_kwargs["device_map"] = {"": get_current_device()}
-        logger.info("Quantizing model to {} bit.".format(model_args.quantization_bit))
-
-    # Load pre-trained models (without valuehead)
+    if model is None:
        model = AutoModelForCausalLM.from_pretrained(
-        model_to_load,
+            model_args.model_name_or_path,
            config=config,
            torch_dtype=model_args.compute_dtype,
            low_cpu_mem_usage=(not is_deepspeed_zero3_enabled()),
-        **config_kwargs
+            **config_kwargs,
        )

-    # Disable custom generate method (for Qwen and Baichuan2)
-    if isinstance(model, PreTrainedModel) and "GenerationMixin" not in str(model.generate.__func__):
-        model.generate = MethodType(PreTrainedModel.generate, model)
+    patch_model(model, tokenizer, model_args, is_trainable)
+    register_autoclass(config, model, tokenizer)

-    # Fix LM head (for ChatGLM2 and ChatGLM3)
-    if getattr(config, "model_type", None) == "chatglm":
-        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
-    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", {}):
-        model.__class__.register_for_auto_class()
-    if isinstance(tokenizer, PreTrainedTokenizerBase) and "AutoTokenizer" in tokenizer.init_kwargs.get("auto_map", {}):
-        tokenizer.__class__.register_for_auto_class()
-
-    # Initialize adapters
-    model = prepare_model_for_training(model=model, finetuning_args=finetuning_args) if is_trainable else model
    model = init_adapter(model, model_args, finetuning_args, is_trainable)
-    model = model.train() if is_trainable else model.eval()

-    # Prepare model with valuehead for RLHF
-    if stage in ["rm", "ppo"]:
+    if add_valuehead:
        model: "AutoModelForCausalLMWithValueHead" = AutoModelForCausalLMWithValueHead.from_pretrained(model)
-        vhead_path = (
-            model_args.checkpoint_dir[-1] if model_args.checkpoint_dir is not None else model_args.model_name_or_path
-        )
+        patch_valuehead_model(model)
+
+        if model_args.adapter_name_or_path is not None:
+            vhead_path = model_args.adapter_name_or_path[-1]
+        else:
+            vhead_path = 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:
-        model.requires_grad_(False) # fix all model params
-        model = model.to(model_args.compute_dtype) if model_args.quantization_bit is None else model
+        model.requires_grad_(False)
+        model = model.to(model_args.compute_dtype) if not getattr(model, "quantization_method", None) else model
+        model.eval()
+    else:
+        model.train()

    trainable_params, all_param = count_parameters(model)
-    logger.info("trainable params: {:d} || all params: {:d} || trainable%: {:.4f}".format(
+    logger.info(
+        "trainable params: {:d} || all params: {:d} || trainable%: {:.4f}".format(
            trainable_params, all_param, 100 * trainable_params / all_param
-    ))
+        )
+    )

    if not is_trainable:
        logger.info("This IS expected that the trainable params is 0 if you are using model for inference only.")
--- a/src/llmtuner/model/patcher.py
+++ b/src/llmtuner/model/patcher.py
@@ -0,0 +1,299 @@
+import math
+import os
+import random
+from contextlib import nullcontext
+from types import MethodType
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+
+import torch
+from datasets import load_dataset
+from transformers import BitsAndBytesConfig, GPTQConfig, PreTrainedModel, PreTrainedTokenizerBase
+from transformers.integrations import is_deepspeed_zero3_enabled
+from transformers.utils.versions import require_version
+
+from ..extras.constants import FILEEXT2TYPE, LAYERNORM_NAMES
+from ..extras.logging import get_logger
+from ..extras.misc import get_current_device, infer_optim_dtype
+from ..extras.packages import is_flash_attn2_available
+from ..extras.patches.llama_patch import apply_llama_patch
+
+
+if TYPE_CHECKING:
+    from transformers import PretrainedConfig, PreTrainedTokenizer
+    from trl import AutoModelForCausalLMWithValueHead
+
+    from ..hparams import ModelArguments
+
+
+logger = get_logger(__name__)
+SUPPORTED_CLASS_FOR_S2ATTN = ["llama"]
+
+
+def _noisy_mean_initialization(embed_weight: torch.Tensor, num_new_tokens: int):
+    embedding_dim = embed_weight.size(1)
+    avg_weight = embed_weight[:-num_new_tokens].mean(dim=0, keepdim=True)
+    noise_weight = torch.empty_like(embed_weight[-num_new_tokens:])
+    noise_weight.normal_(mean=0, std=(1.0 / math.sqrt(embedding_dim)))
+    embed_weight[-num_new_tokens:] = avg_weight + noise_weight
+
+
+def _resize_embedding_layer(model: "PreTrainedModel", tokenizer: "PreTrainedTokenizer") -> None:
+    r"""
+    Resize token embeddings.
+    """
+    if is_deepspeed_zero3_enabled():
+        import deepspeed  # type: ignore
+
+        params = [model.get_input_embeddings().weight]
+        if model.get_output_embeddings() is not None and not model.config.tie_word_embeddings:
+            params.append(model.get_output_embeddings().weight)
+
+        context_maybe_zero3 = deepspeed.zero.GatheredParameters(params, modifier_rank=0)
+    else:
+        context_maybe_zero3 = nullcontext()
+
+    with context_maybe_zero3:
+        current_embedding_size = model.get_input_embeddings().weight.size(0)
+
+    if len(tokenizer) > current_embedding_size:
+        if not isinstance(model.get_output_embeddings(), torch.nn.Linear):
+            logger.warning("Current model does not support resizing token embeddings.")
+            return
+
+        model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=64)
+        with context_maybe_zero3:
+            new_embedding_size = model.get_input_embeddings().weight.size(0)
+            num_new_tokens = new_embedding_size - current_embedding_size
+            _noisy_mean_initialization(model.get_input_embeddings().weight.data, num_new_tokens)
+            _noisy_mean_initialization(model.get_output_embeddings().weight.data, num_new_tokens)
+
+        logger.info("Resized token embeddings from {} to {}.".format(current_embedding_size, new_embedding_size))
+
+
+def _get_quantization_dataset(tokenizer: "PreTrainedTokenizer", model_args: "ModelArguments") -> List[str]:
+    r"""
+    Inspired by: https://github.com/huggingface/optimum/blob/v1.16.0/optimum/gptq/data.py#L133
+    TODO: remove tokenizer.decode() https://github.com/huggingface/optimum/pull/1600
+    """
+    if os.path.isfile(model_args.export_quantization_dataset):
+        data_path = FILEEXT2TYPE.get(model_args.export_quantization_dataset.split(".")[-1], None)
+        data_files = model_args.export_quantization_dataset
+    else:
+        data_path = model_args.export_quantization_dataset
+        data_files = None
+
+    dataset = load_dataset(path=data_path, data_files=data_files, split="train", cache_dir=model_args.cache_dir)
+    maxlen = model_args.export_quantization_maxlen
+
+    samples = []
+    for _ in range(model_args.export_quantization_nsamples):
+        while True:
+            sample_idx = random.randint(0, len(dataset) - 1)
+            sample: Dict[str, torch.Tensor] = tokenizer(dataset[sample_idx]["text"], return_tensors="pt")
+            if sample["input_ids"].size(1) >= maxlen:
+                break  # TODO: fix large maxlen
+
+        word_idx = random.randint(0, sample["input_ids"].size(1) - maxlen - 1)
+        input_ids = sample["input_ids"][:, word_idx : word_idx + maxlen]
+        samples.append(tokenizer.decode(input_ids[0].tolist(), skip_special_tokens=True))
+
+    return samples
+
+
+def _configure_rope(config: "PretrainedConfig", model_args: "ModelArguments", is_trainable: bool) -> None:
+    if not hasattr(config, "rope_scaling"):
+        logger.warning("Current model does not support RoPE scaling.")
+        return
+
+    if is_trainable:
+        if model_args.rope_scaling == "dynamic":
+            logger.warning(
+                "Dynamic NTK scaling may not work well with fine-tuning. "
+                "See: https://github.com/huggingface/transformers/pull/24653"
+            )
+
+        current_max_length = getattr(config, "max_position_embeddings", None)
+        if current_max_length and model_args.model_max_length > current_max_length:
+            scaling_factor = float(math.ceil(model_args.model_max_length / current_max_length))
+        else:
+            logger.warning("Input length is smaller than max length. Consider increase input length.")
+            scaling_factor = 1.0
+    else:
+        scaling_factor = 2.0
+
+    setattr(config, "rope_scaling", {"type": model_args.rope_scaling, "factor": scaling_factor})
+    logger.info(
+        "Using {} scaling strategy and setting scaling factor to {}".format(model_args.rope_scaling, scaling_factor)
+    )
+
+
+def _configure_flashattn(config_kwargs: Dict[str, Any]) -> None:
+    if not is_flash_attn2_available():
+        logger.warning("FlashAttention2 is not installed.")
+        return
+
+    config_kwargs["use_flash_attention_2"] = True
+    logger.info("Using FlashAttention-2 for faster training and inference.")
+
+
+def _configure_longlora(config: "PretrainedConfig") -> None:
+    if getattr(config, "model_type", None) in SUPPORTED_CLASS_FOR_S2ATTN:
+        setattr(config, "group_size_ratio", 0.25)
+        apply_llama_patch()
+        logger.info("Using shift short attention with group_size_ratio=1/4.")
+    else:
+        logger.warning("Current model does not support shift short attention.")
+
+
+def _configure_quantization(
+    config: "PretrainedConfig",
+    tokenizer: "PreTrainedTokenizer",
+    model_args: "ModelArguments",
+    config_kwargs: Dict[str, Any],
+) -> None:
+    r"""
+    Priority: GPTQ-quantized (training) > AutoGPTQ (export) > Bitsandbytes (training)
+    """
+    if getattr(config, "quantization_config", None):  # gptq
+        if is_deepspeed_zero3_enabled():
+            raise ValueError("DeepSpeed ZeRO-3 is incompatible with quantization.")
+
+        config_kwargs["device_map"] = {"": get_current_device()}
+        quantization_config: Dict[str, Any] = getattr(config, "quantization_config", None)
+        if quantization_config.get("quant_method", None) == "gptq" and quantization_config.get("bits", -1) == 4:
+            quantization_config["use_exllama"] = False  # disable exllama
+        logger.info("Loading {}-bit GPTQ-quantized model.".format(quantization_config.get("bits", -1)))
+
+    elif model_args.export_quantization_bit is not None:  # auto-gptq
+        require_version("optimum>=1.16.0", "To fix: pip install optimum>=1.16.0")
+        require_version("auto_gptq>=0.5.0", "To fix: pip install auto_gptq>=0.5.0")
+        from accelerate.utils import get_max_memory
+
+        if getattr(config, "model_type", None) == "chatglm":
+            raise ValueError("ChatGLM model is not supported.")
+
+        config_kwargs["quantization_config"] = GPTQConfig(
+            bits=model_args.export_quantization_bit,
+            tokenizer=tokenizer,
+            dataset=_get_quantization_dataset(tokenizer, model_args),
+        )
+        config_kwargs["device_map"] = "auto"
+        config_kwargs["max_memory"] = get_max_memory()
+        logger.info("Quantizing model to {} bit.".format(model_args.export_quantization_bit))
+
+    elif model_args.quantization_bit is not None:  # bnb
+        if is_deepspeed_zero3_enabled():
+            raise ValueError("DeepSpeed ZeRO-3 is incompatible with quantization.")
+
+        if model_args.quantization_bit == 8:
+            require_version("bitsandbytes>=0.37.0", "To fix: pip install bitsandbytes>=0.37.0")
+            config_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
+
+        elif model_args.quantization_bit == 4:
+            require_version("bitsandbytes>=0.39.0", "To fix: pip install bitsandbytes>=0.39.0")
+            config_kwargs["quantization_config"] = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=model_args.compute_dtype,
+                bnb_4bit_use_double_quant=model_args.double_quantization,
+                bnb_4bit_quant_type=model_args.quantization_type,
+            )
+
+        config_kwargs["device_map"] = {"": get_current_device()}
+        logger.info("Quantizing model to {} bit.".format(model_args.quantization_bit))
+
+
+def _prepare_model_for_training(
+    model: "PreTrainedModel", model_args: "ModelArguments", output_layer_name: Optional[str] = "lm_head"
+) -> None:
+    r"""
+    Includes:
+        (1) cast the layernorm in fp32
+        (2) make output embedding layer require grads
+        (3) add the upcasting of the lm_head in fp32
+    Inspired by: https://github.com/huggingface/peft/blob/v0.7.1/src/peft/utils/other.py#L72
+    """
+    if model_args.upcast_layernorm:
+        for name, param in model.named_parameters():
+            if param.ndim == 1 and any(ln_name in name for ln_name in LAYERNORM_NAMES):
+                param.data = param.data.to(torch.float32)
+        logger.info("Upcasting layernorm weights in float32.")
+
+    if not model_args.disable_gradient_checkpointing:
+        if not getattr(model, "supports_gradient_checkpointing", False):
+            logger.warning("Current model does not support gradient checkpointing.")
+        else:
+            model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+            model.config.use_cache = False  # turn off when gradient checkpointing is enabled
+            logger.info("Gradient checkpointing enabled.")
+
+    if hasattr(model, output_layer_name) and model_args.upcast_lmhead_output:
+
+        def fp32_forward_post_hook(module: torch.nn.Module, args: Tuple[torch.Tensor], output: torch.Tensor):
+            return output.to(torch.float32)
+
+        output_layer = getattr(model, output_layer_name)
+        if isinstance(output_layer, torch.nn.Linear) and output_layer.weight.dtype != torch.float32:
+            output_layer.register_forward_hook(fp32_forward_post_hook)
+
+
+def patch_tokenizer(tokenizer: "PreTrainedTokenizer") -> None:
+    if "PreTrainedTokenizerBase" not in str(tokenizer._pad.__func__):
+        tokenizer._pad = MethodType(PreTrainedTokenizerBase._pad, tokenizer)
+
+
+def patch_config(
+    config: "PretrainedConfig",
+    tokenizer: "PreTrainedTokenizer",
+    model_args: "ModelArguments",
+    config_kwargs: Dict[str, Any],
+    is_trainable: bool,
+) -> None:
+    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))
+
+    if getattr(config, "model_type", None) == "qwen":
+        for dtype_name, dtype in [("fp16", torch.float16), ("bf16", torch.bfloat16), ("fp32", torch.float32)]:
+            setattr(config, dtype_name, model_args.compute_dtype == dtype)
+
+    if model_args.rope_scaling is not None:
+        _configure_rope(config, model_args, is_trainable)
+
+    if model_args.flash_attn:
+        _configure_flashattn(config_kwargs)
+
+    if is_trainable and model_args.shift_attn:
+        _configure_longlora(config)
+
+    _configure_quantization(config, tokenizer, model_args, config_kwargs)
+
+
+def patch_model(
+    model: "PreTrainedModel", tokenizer: "PreTrainedTokenizer", model_args: "ModelArguments", is_trainable: bool
+) -> None:
+    if "GenerationMixin" not in str(model.generate.__func__):
+        model.generate = MethodType(PreTrainedModel.generate, model)
+
+    if getattr(model.config, "model_type", None) == "chatglm":
+        setattr(model, "lm_head", model.transformer.output_layer)
+        setattr(model, "_keys_to_ignore_on_save", ["lm_head.weight"])
+
+    if model_args.resize_vocab:
+        _resize_embedding_layer(model, tokenizer)
+
+    if is_trainable:
+        _prepare_model_for_training(model, model_args)
+
+
+def patch_valuehead_model(model: "AutoModelForCausalLMWithValueHead") -> None:
+    def tie_weights(self: "AutoModelForCausalLMWithValueHead") -> None:
+        if isinstance(self.pretrained_model, PreTrainedModel):
+            self.pretrained_model.tie_weights()
+
+    def get_input_embeddings(self: "AutoModelForCausalLMWithValueHead") -> torch.nn.Module:
+        if isinstance(self.pretrained_model, PreTrainedModel):
+            return self.pretrained_model.get_input_embeddings()
+
+    ignore_modules = [name for name, _ in model.named_parameters() if "pretrained_model" in name]
+    setattr(model, "_keys_to_ignore_on_save", ignore_modules)
+    setattr(model, "tie_weights", MethodType(tie_weights, model))
+    setattr(model, "get_input_embeddings", MethodType(get_input_embeddings, model))
--- a/src/llmtuner/model/utils.py
+++ b/src/llmtuner/model/utils.py
@@ -1,16 +1,19 @@
+import inspect
+from typing import TYPE_CHECKING, Any, Dict, List
+
 import torch
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple
-
+from transformers import PreTrainedModel
 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
+from ..extras.constants import V_HEAD_SAFE_WEIGHTS_NAME, V_HEAD_WEIGHTS_NAME
+from ..extras.logging import get_logger
+from ..extras.misc import get_current_device
+

 if TYPE_CHECKING:
-    from transformers.modeling_utils import PreTrainedModel
-    from llmtuner.hparams import DataArguments
+    from transformers import PretrainedConfig, PreTrainedTokenizer
+
+    from ..hparams import DataArguments, FinetuningArguments, ModelArguments


 logger = get_logger(__name__)
@@ -18,41 +21,47 @@ 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
+    Dispatches a pre-trained model to GPUs with balanced memory when the GPU is available.
+    Borrowed from: https://github.com/huggingface/transformers/blob/v4.36.2/src/transformers/modeling_utils.py#L3570
    """
-    if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False): # do nothing
+    if getattr(model, "quantization_method", None):  # already set on current device
        return model

-    if torch.cuda.device_count() > 1:
+    if (
+        torch.cuda.device_count() > 1
+        and isinstance(model, PreTrainedModel)
+        and model._no_split_modules is not None
+        and model.config.model_type != "chatglm"
+    ):
        from accelerate import dispatch_model
-        from accelerate.utils import infer_auto_device_map, get_balanced_memory
+        from accelerate.utils import get_balanced_memory, infer_auto_device_map

-        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}
+        kwargs = {"dtype": model.dtype, "no_split_module_classes": model._get_no_split_modules("auto")}
        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)
+        device_map_kwargs = {"device_map": device_map}
+        if "skip_keys" in inspect.signature(dispatch_model).parameters:
+            device_map_kwargs["skip_keys"] = model._skip_keys_device_placement
+        return dispatch_model(model, **device_map_kwargs)
    else:
-        return model.cuda()
+        return model.to(device=get_current_device())


-def find_all_linear_modules(
-    model: "PreTrainedModel",
-    quantization_bit: Optional[int] = None
-) -> List[str]:
+def find_all_linear_modules(model: "PreTrainedModel") -> 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
-    else:
+    quantization_method = getattr(model, "quantization_method", None)
+    if quantization_method is None:
        linear_cls = torch.nn.Linear
+    elif quantization_method == "bitsandbytes":
+        import bitsandbytes as bnb
+
+        linear_cls = bnb.nn.Linear4bit if getattr(model, "is_loaded_in_4bit", False) else bnb.nn.Linear8bitLt
+    else:
+        raise ValueError("Finding linear modules for {} models is not supported.".format(quantization_method))

    output_layer_names = ["lm_head"]
    if model.config.model_type == "chatglm":
@@ -60,106 +69,57 @@ def find_all_linear_modules(

    module_names = set()
    for name, module in model.named_modules():
-        if (
-            isinstance(module, linear_cls)
-            and not any([output_layer in name for output_layer in output_layer_names])
-        ):
+        if isinstance(module, linear_cls) and not any(output_layer in name for output_layer in output_layer_names):
            module_names.add(name.split(".")[-1])

    logger.info("Found linear modules: {}".format(",".join(module_names)))
    return list(module_names)


-def generate_model_card(
-    model_args: "ModelArguments",
-    data_args: "DataArguments",
-    finetuning_args: "FinetuningArguments"
+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 [])
+        "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]:
+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,
-        "token": model_args.hf_hub_token
-    }
+    kwargs = {"path_or_repo_id": path_or_repo_id, "cache_dir": model_args.cache_dir, "token": model_args.hf_hub_token}
+
    try:
-        vhead_file = cached_file(filename=WEIGHTS_NAME, **kwargs)
-    except:
+        from safetensors import safe_open
+
+        vhead_file = cached_file(filename=V_HEAD_SAFE_WEIGHTS_NAME, **kwargs)
+        with safe_open(vhead_file, framework="pt", device="cpu") as f:
+            return {key: f.get_tensor(key) for key in f.keys()}
+    except Exception as err:
+        logger.info("Failed to load {}: {}".format(V_HEAD_SAFE_WEIGHTS_NAME, str(err)))
+
    try:
-            vhead_file = cached_file(filename=SAFE_WEIGHTS_NAME, **kwargs)
-        except:
-            logger.warning("Provided path ({}) does not contain valuehead weights.".format(path_or_repo_id))
+        vhead_file = cached_file(filename=V_HEAD_WEIGHTS_NAME, **kwargs)
+        return torch.load(vhead_file, map_location="cpu")
+    except Exception as err:
+        logger.info("Failed to load {}: {}".format(V_HEAD_WEIGHTS_NAME, str(err)))
+
+    logger.info("Provided path ({}) does not contain value head weights.".format(path_or_repo_id))
+    logger.info("Ignore these messages if you are not resuming the training of a value head model.")
    return None

-    return torch.load(vhead_file, map_location="cpu")

-
-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[Set[str]] = LAYERNORM_NAMES
-) -> "PreTrainedModel":
-    r"""
-    Includes:
-        (1) cast the layernorm in fp32
-        (2) make output embedding layer require grads
-        (3) upcast the lm_head to fp32
-    Inspired by: https://github.com/huggingface/peft/blob/v0.2.0/src/peft/utils/other.py#L33
-    """
-    if finetuning_args.upcast_layernorm:
-        for name, param in model.named_parameters():
-            if param.ndim == 1 and any(ln_name in name for ln_name in layernorm_names):
-                param.data = param.data.to(torch.float32)
-        logger.info("Upcasting weights in layernorm in float32.")
-
-    if finetuning_args.neft_alpha > 1e-6:
-        def neftune_forward_hook(module: torch.nn.Module, args: Tuple[torch.Tensor], output: torch.Tensor):
-            if module.training:
-                dims = torch.tensor(output.size(1) * output.size(2))
-                mag_norm = finetuning_args.neft_alpha / torch.sqrt(dims)
-                output = output + torch.zeros_like(output).uniform_(-mag_norm, mag_norm)
-            return output
-
-        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 hasattr(model, "enable_input_require_grads"):
-            model.enable_input_require_grads()
-        else:
-            def make_inputs_require_grad(module: torch.nn.Module, args: Tuple[torch.Tensor], output: torch.Tensor):
-                output.requires_grad_(True)
-            model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
-
-        model.gradient_checkpointing_enable()
-        model.config.use_cache = False # turn off when gradient checkpointing is enabled
-        logger.info("Gradient checkpointing enabled.")
-
-    if finetuning_args.finetuning_type != "full" and hasattr(model, output_layer_name):
-        output_layer = getattr(model, output_layer_name)
-        if isinstance(output_layer, torch.nn.Linear):
-            def fp32_forward_pre_hook(module: torch.nn.Module, args: Tuple[torch.Tensor]):
-                return args[0].to(output_layer.weight.dtype)
-            def fp32_forward_post_hook(module: torch.nn.Module, args: Tuple[torch.Tensor], output: torch.Tensor):
-                return output.to(torch.float32)
-            output_layer.register_forward_pre_hook(fp32_forward_pre_hook)
-            output_layer.register_forward_hook(fp32_forward_post_hook)
-
-    return model
+def register_autoclass(config: "PretrainedConfig", model: "PreTrainedModel", tokenizer: "PreTrainedTokenizer"):
+    if "AutoConfig" in getattr(config, "auto_map", {}):
+        config.__class__.register_for_auto_class()
+    if "AutoModelForCausalLM" in getattr(config, "auto_map", {}):
+        model.__class__.register_for_auto_class()
+    if "AutoTokenizer" in tokenizer.init_kwargs.get("auto_map", {}):
+        tokenizer.__class__.register_for_auto_class()
--- a/src/llmtuner/train/init.py
+++ b/src/llmtuner/train/init.py
@@ -1 +1,4 @@
-from llmtuner.train.tuner import export_model, run_exp
+from .tuner import export_model, run_exp
+
+
+__all__ = ["export_model", "run_exp"]
--- a/src/llmtuner/train/dpo/init.py
+++ b/src/llmtuner/train/dpo/init.py
@@ -1 +1,4 @@
-from llmtuner.train.dpo.workflow import run_dpo
+from .workflow import run_dpo
+
+
+__all__ = ["run_dpo"]
--- a/src/llmtuner/train/dpo/collator.py
+++ b/src/llmtuner/train/dpo/collator.py
@@ -1,6 +1,7 @@
-import torch
 from dataclasses import dataclass
 from typing import Any, Dict, List, Sequence, Tuple
+
+import torch
 from transformers import DataCollatorForSeq2Seq


@@ -34,10 +35,12 @@ class DPODataCollatorWithPadding(DataCollatorForSeq2Seq):
        for key in ("chosen_ids", "rejected_ids"):
            for feature in features:
                prompt_len, answer_len = len(feature["prompt_ids"]), len(feature[key])
-                concatenated_features.append({
+                concatenated_features.append(
+                    {
                        "input_ids": feature["prompt_ids"] + feature[key],
-                    "attention_mask": [1] * (prompt_len + answer_len)
-                })
+                        "attention_mask": [1] * (prompt_len + answer_len),
+                    }
+                )
                label_positions.append((prompt_len, answer_len))

        batch = self.tokenizer.pad(
--- a/src/llmtuner/train/dpo/trainer.py
+++ b/src/llmtuner/train/dpo/trainer.py
@@ -1,40 +1,50 @@
-import torch
 from collections import defaultdict
+from contextlib import nullcontext
 from typing import TYPE_CHECKING, Dict, Literal, Optional, Tuple, Union
+
+import torch
 from transformers import BatchEncoding, Trainer
 from trl import DPOTrainer
 from trl.trainer.utils import disable_dropout_in_model

-from llmtuner.extras.constants import IGNORE_INDEX
+from ...extras.constants import IGNORE_INDEX
+

 if TYPE_CHECKING:
    from transformers import PreTrainedModel


 class CustomDPOTrainer(DPOTrainer):
-
    def __init__(
        self,
        beta: float,
+        loss_type: Literal["sigmoid", "hinge", "ipo", "kto"],
+        ftx_gamma: float,
        model: Union["PreTrainedModel", torch.nn.Module],
        ref_model: Optional[Union["PreTrainedModel", torch.nn.Module]] = None,
        disable_dropout: Optional[bool] = True,
-        loss_type: Optional[Literal["sigmoid", "hinge"]] = "sigmoid",
-        **kwargs
+        **kwargs,
    ):
        if disable_dropout:
            disable_dropout_in_model(model)
            if ref_model is not None:
                disable_dropout_in_model(ref_model)

-        self.is_encoder_decoder = model.config.is_encoder_decoder
-        self.ref_model = ref_model
        self.use_dpo_data_collator = True  # hack to avoid warning
        self.generate_during_eval = False  # disable at evaluation
        self.label_pad_token_id = IGNORE_INDEX
        self.padding_value = 0
+        self.is_encoder_decoder = model.config.is_encoder_decoder
+        self.precompute_ref_log_probs = False
+        self._precomputed_train_ref_log_probs = False
+        self._precomputed_eval_ref_log_probs = False
+        self._peft_has_been_casted_to_bf16 = False
+
+        self.ref_model = ref_model
        self.beta = beta
+        self.label_smoothing = 0
        self.loss_type = loss_type
+        self.ftx_gamma = ftx_gamma
        self._stored_metrics = defaultdict(lambda: defaultdict(list))

        Trainer.__init__(self, model=model, **kwargs)
@@ -44,32 +54,95 @@ class CustomDPOTrainer(DPOTrainer):
        if ref_model is not None:
            if self.is_deepspeed_enabled:
                if not (
-                    getattr(ref_model, "is_loaded_in_8bit", False)
-                    or getattr(ref_model, "is_loaded_in_4bit", False)
+                    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 sft_loss(self, chosen_logits: torch.FloatTensor, chosen_labels: torch.LongTensor) -> torch.Tensor:
+        r"""
+        Computes supervised cross-entropy loss of given labels under the given logits.
+
+        Returns:
+            A tensor of shape (batch_size,) containing the cross-entropy loss of each samples.
+        """
+        all_logps = self.get_batch_logps(chosen_logits, chosen_labels, average_log_prob=True)
+        return -all_logps
+
    def concatenated_forward(
-        self,
-        model: Optional[torch.nn.Module] = None,
-        batch: Optional[Dict[str, torch.Tensor]] = None
+        self, model: "PreTrainedModel", batch: Dict[str, torch.Tensor]
    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
        batch_copied = BatchEncoding({k: v.detach().clone() for k, v in batch.items()})  # avoid error

        all_logits = model(
-            input_ids=batch_copied["input_ids"],
-            attention_mask=batch_copied["attention_mask"],
-            return_dict=True
+            input_ids=batch_copied["input_ids"], attention_mask=batch_copied["attention_mask"], return_dict=True
        ).logits.to(torch.float32)

-        all_logps = self._get_batch_logps(
+        all_logps = self.get_batch_logps(
            all_logits,
            batch["labels"],
-            average_log_prob=False
+            average_log_prob=False,
+            label_pad_token_id=self.label_pad_token_id,
        )
        batch_size = batch["input_ids"].size(0) // 2
        chosen_logps, rejected_logps = all_logps.split(batch_size, dim=0)
        chosen_logits, rejected_logits = all_logits.split(batch_size, dim=0)
        return chosen_logps, rejected_logps, chosen_logits, rejected_logits
+
+    def get_batch_loss_metrics(
+        self,
+        model: "PreTrainedModel",
+        batch: Dict[str, torch.Tensor],
+        train_eval: Optional[Literal["train", "eval"]] = "train",
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        r"""
+        Computes the DPO loss and other metrics for the given batch of inputs for train or test.
+        """
+        metrics = {}
+        (
+            policy_chosen_logps,
+            policy_rejected_logps,
+            policy_chosen_logits,
+            policy_rejected_logits,
+        ) = self.concatenated_forward(model, batch)
+        with torch.no_grad():
+            if self.ref_model is None:
+                ref_model = self.model
+                ref_context = self.accelerator.unwrap_model(self.model).disable_adapter()
+            else:
+                ref_model = self.ref_model
+                ref_context = nullcontext()
+
+            with ref_context:
+                (
+                    reference_chosen_logps,
+                    reference_rejected_logps,
+                    _,
+                    _,
+                ) = self.concatenated_forward(ref_model, batch)
+
+        losses, chosen_rewards, rejected_rewards = self.dpo_loss(
+            policy_chosen_logps,
+            policy_rejected_logps,
+            reference_chosen_logps,
+            reference_rejected_logps,
+        )
+        if self.ftx_gamma > 1e-6:
+            batch_size = batch["input_ids"].size(0) // 2
+            chosen_labels, _ = batch["labels"].split(batch_size, dim=0)
+            losses += self.ftx_gamma * self.sft_loss(policy_chosen_logits, chosen_labels)
+
+        reward_accuracies = (chosen_rewards > rejected_rewards).float()
+
+        prefix = "eval_" if train_eval == "eval" else ""
+        metrics[f"{prefix}rewards/chosen"] = chosen_rewards.cpu().mean()
+        metrics[f"{prefix}rewards/rejected"] = rejected_rewards.cpu().mean()
+        metrics[f"{prefix}rewards/accuracies"] = reward_accuracies.cpu().mean()
+        metrics[f"{prefix}rewards/margins"] = (chosen_rewards - rejected_rewards).cpu().mean()
+        metrics[f"{prefix}logps/rejected"] = policy_rejected_logps.detach().cpu().mean()
+        metrics[f"{prefix}logps/chosen"] = policy_chosen_logps.detach().cpu().mean()
+        metrics[f"{prefix}logits/rejected"] = policy_rejected_logits.detach().cpu().mean()
+        metrics[f"{prefix}logits/chosen"] = policy_chosen_logits.detach().cpu().mean()
+
+        return losses.mean(), metrics
--- a/src/llmtuner/train/dpo/workflow.py
+++ b/src/llmtuner/train/dpo/workflow.py
@@ -1,21 +1,23 @@
 # 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 typing import TYPE_CHECKING, List, Optional
+
 from transformers import Seq2SeqTrainingArguments

-from llmtuner.data import get_dataset, preprocess_dataset, split_dataset
-from llmtuner.extras.constants import IGNORE_INDEX
-from llmtuner.extras.ploting import plot_loss
-from llmtuner.hparams import ModelArguments
-from llmtuner.model import generate_model_card, load_model_and_tokenizer
-from llmtuner.train.utils import create_ref_model
-from llmtuner.train.dpo.collator import DPODataCollatorWithPadding
-from llmtuner.train.dpo.trainer import CustomDPOTrainer
+from ...data import get_dataset, split_dataset
+from ...extras.constants import IGNORE_INDEX
+from ...extras.ploting import plot_loss
+from ...hparams import ModelArguments
+from ...model import load_model_and_tokenizer
+from ...train.dpo.collator import DPODataCollatorWithPadding
+from ...train.dpo.trainer import CustomDPOTrainer
+from ...train.utils import create_modelcard_and_push, create_ref_model
+

 if TYPE_CHECKING:
    from transformers import TrainerCallback
-    from llmtuner.hparams import DataArguments, FinetuningArguments
+
+    from ...hparams import DataArguments, FinetuningArguments


 def run_dpo(
@@ -23,22 +25,21 @@ def run_dpo(
    data_args: "DataArguments",
    training_args: "Seq2SeqTrainingArguments",
    finetuning_args: "FinetuningArguments",
-    callbacks: Optional[List["TrainerCallback"]] = None
+    callbacks: Optional[List["TrainerCallback"]] = None,
 ):
-    dataset = get_dataset(model_args, data_args)
-    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="sft")
-    dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="rm")
+    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train)
+    dataset = get_dataset(tokenizer, model_args, data_args, training_args, stage="rm")
    data_collator = DPODataCollatorWithPadding(
        tokenizer=tokenizer,
-        pad_to_multiple_of=4,
-        label_pad_token_id=IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
+        pad_to_multiple_of=8,
+        label_pad_token_id=IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id,
    )

    # Create reference model
    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")
+        ref_model = create_ref_model(model_args, finetuning_args)

    # Update arguments
    training_args_dict = training_args.to_dict()
@@ -48,13 +49,15 @@ def run_dpo(
    # Initialize our Trainer
    trainer = CustomDPOTrainer(
        beta=finetuning_args.dpo_beta,
+        loss_type=finetuning_args.dpo_loss,
+        ftx_gamma=finetuning_args.dpo_ftx,
        model=model,
        ref_model=ref_model,
        args=training_args,
        tokenizer=tokenizer,
        data_collator=data_collator,
        callbacks=callbacks,
-        **split_dataset(dataset, data_args, training_args)
+        **split_dataset(dataset, data_args, training_args),
    )

    # Training
@@ -78,8 +81,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)
--- a/src/llmtuner/train/ppo/init.py
+++ b/src/llmtuner/train/ppo/init.py
@@ -1 +1,4 @@
-from llmtuner.train.ppo.workflow import run_ppo
+from .workflow import run_ppo
+
+
+__all__ = ["run_ppo"]
--- a/src/llmtuner/train/ppo/trainer.py
+++ b/src/llmtuner/train/ppo/trainer.py
@@ -1,25 +1,28 @@
+import math
 import os
 import sys
-import math
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
+
 import torch
 from tqdm import tqdm
-from typing import TYPE_CHECKING, List, Optional, Tuple
-
-from transformers import BatchEncoding, GenerationConfig, Trainer, TrainerState, TrainerControl
+from transformers import GenerationConfig, Trainer, TrainerControl, TrainerState
+from transformers.trainer_pt_utils import remove_dummy_checkpoint
 from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
-
+from transformers.utils import SAFE_WEIGHTS_NAME, WEIGHTS_NAME
 from trl import PPOTrainer
 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.train.ppo.utils import dump_layernorm, restore_layernorm, replace_model
+from ...extras.callbacks import FixValueHeadModelCallback, LogCallback
+from ...extras.logging import get_logger
+from ...extras.misc import AverageMeter, count_parameters, get_logits_processor
+from .utils import dump_layernorm, get_rewards_from_server, replace_model, restore_layernorm
+

 if TYPE_CHECKING:
    from transformers import Seq2SeqTrainingArguments, TrainerCallback
    from trl import AutoModelForCausalLMWithValueHead
-    from llmtuner.hparams import ModelArguments, FinetuningArguments, GeneratingArguments
+
+    from ...hparams import FinetuningArguments, GeneratingArguments, ModelArguments


 logger = get_logger(__name__)
@@ -38,7 +41,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
        generating_args: "GeneratingArguments",
        callbacks: List["TrainerCallback"],
        reward_model: "AutoModelForCausalLMWithValueHead",
-        **kwargs
+        **kwargs,
    ):
        PPOTrainer.__init__(self, **kwargs)

@@ -50,22 +53,22 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
        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()
+            **generating_args.to_dict(),
        )

        self.state = TrainerState()
        self.control = TrainerControl()
+        self.is_deepspeed_enabled = self.accelerator.distributed_type == "DEEPSPEED" and hasattr(
+            self.accelerator.state, "deepspeed_plugin"
+        )
        self.log_callback, self.save_callback = callbacks[0], callbacks[1]
-        assert isinstance(self.log_callback, LogCallback) and isinstance(self.save_callback, SavePeftModelCallback)
+        assert isinstance(self.log_callback, LogCallback) and isinstance(self.save_callback, FixValueHeadModelCallback)

        if self.args.max_steps > 0:
            logger.info("max_steps is given, it will override any value given in num_train_epochs")

-        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 finetuning_args.reward_model_type == "full":
+            if self.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)
@@ -74,18 +77,24 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
            else:
                self.reward_model = self.accelerator.prepare_model(self.reward_model, evaluation_mode=True)

-    def ppo_train(self) -> None:
+    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)
@@ -100,13 +109,18 @@ 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)
@@ -130,7 +144,9 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
            self.tokenizer.padding_side = "right"  # change padding side
            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_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)
@@ -152,7 +168,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
                    batch["query"] = self.tokenizer.batch_decode(queries, skip_special_tokens=True)
                    batch["response"] = self.tokenizer.batch_decode(responses, skip_special_tokens=True)
                    self.log_stats(stats, batch, rewards)
-                except:
+                except Exception:
                    logger.warning("Failed to save stats due to unknown errors.")

            self.state.global_step += 1
@@ -163,7 +179,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
                    loss=round(loss_meter.avg, 4),
                    reward=round(reward_meter.avg, 4),
                    learning_rate=stats["ppo/learning_rate"],
-                    epoch=round(step / steps_in_epoch, 2)
+                    epoch=round(step / steps_in_epoch, 2),
                )
                tqdm.write(str(logs))
                logs["step"] = step
@@ -173,9 +189,9 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
                reward_meter.reset()

            if (step + 1) % self.args.save_steps == 0:  # save checkpoint
-                self.save_model(os.path.join(
-                    self.args.output_dir, "{}-{}".format(PREFIX_CHECKPOINT_DIR, self.state.global_step)
-                ))
+                self.save_model(
+                    os.path.join(self.args.output_dir, "{}-{}".format(PREFIX_CHECKPOINT_DIR, self.state.global_step))
+                )
                self.save_callback.on_save(
                    self.args, self.state, self.control, model=self.accelerator.unwrap_model(self.model)
                )
@@ -189,27 +205,31 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
        )

    @torch.no_grad()
-    def get_inputs(self, batch: BatchEncoding) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
+    def get_inputs(self, batch: Dict[str, torch.Tensor]) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
        r"""
        Generates model's responses given queries.
        """
-        if self.finetuning_args.upcast_layernorm:
+        if self.model_args.upcast_layernorm:
            layernorm_params = dump_layernorm(self.model)

+        if batch["input_ids"].size(0) == 1:  # handle llama2 ppo with gradient accumulation > 1
+            start_index = (batch["input_ids"][0] != self.tokenizer.pad_token_id).nonzero()[0].item()
+            for k, v in batch.items():
+                batch[k] = v[:, start_index:]
+
        unwrapped_model: "AutoModelForCausalLMWithValueHead" = self.accelerator.unwrap_model(self.model)
-        response: torch.Tensor = unwrapped_model.generate(
-            generation_config=self.generation_config,
-            logits_processor=get_logits_processor(),
-            **batch
+        generate_output: torch.Tensor = unwrapped_model.generate(
+            generation_config=self.generation_config, logits_processor=get_logits_processor(), **batch
        )

-        if self.finetuning_args.upcast_layernorm:
+        if self.model_args.upcast_layernorm:
            restore_layernorm(self.model, layernorm_params)

-        query, response = batch["input_ids"].detach().cpu(), response[:, batch["input_ids"].size(-1):].detach().cpu()
+        query = batch["input_ids"].detach().cpu()
+        response = generate_output[:, batch["input_ids"].size(-1) :].detach().cpu()
        queries, responses = [], []
        for i in range(len(query)):
-            query_length = (query[i] != self.tokenizer.pad_token_id).nonzero()[0].item()
+            query_start_index = (query[i] != self.tokenizer.pad_token_id).nonzero()[0].item()
            response_index = (response[i] != self.tokenizer.pad_token_id).nonzero()

            if len(response_index) == 0:
@@ -217,7 +237,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
            else:
                response_length = response_index[-1].item() + 1

-            queries.append(query[i, query_length:]) # remove padding from left
+            queries.append(query[i, query_start_index:])  # remove padding from left
            responses.append(response[i, :response_length])  # remove padding from right

        return queries, responses
@@ -227,21 +247,30 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
        self,
        queries: List[torch.Tensor],
        responses: List[torch.Tensor],
-        unwrapped_model: "AutoModelForCausalLMWithValueHead"
+        unwrapped_model: "AutoModelForCausalLMWithValueHead",
    ) -> List[torch.Tensor]:
        r"""
        Computes scores using given reward model.
+
+        Both inputs and outputs are put on CPU.
        """
-        if self.reward_model is None:
+        if self.finetuning_args.reward_model_type == "api":
+            token_ids = [torch.cat((q, r), dim=-1).tolist() for q, r in zip(queries, responses)]
+            messages = self.tokenizer.batch_decode(token_ids, skip_special_tokens=True)
+            return get_rewards_from_server(self.reward_model, messages)
+
+        if self.finetuning_args.reward_model_type == "lora":
            replace_model(unwrapped_model, target="reward")
+            reward_model = self.model
+        else:
+            reward_model = self.reward_model

        batch = self.prepare_model_inputs(queries, responses)

        with torch.cuda.amp.autocast(dtype=self.model_args.compute_dtype):  # support bf16
-            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
+        if getattr(unwrapped_model.config, "model_type", None) == "chatglm":  # assume same architecture
            values = torch.transpose(values, 0, 1)

        rewards = []
@@ -250,7 +279,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
            end_index = end_indexes[-1].item() if len(end_indexes) else 0
            rewards.append(values[i, end_index].float().detach().cpu())  # use fp32 type

-        if self.reward_model is None:
+        if self.finetuning_args.reward_model_type == "lora":
            replace_model(unwrapped_model, target="default")

        return rewards
@@ -263,7 +292,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
        responses: torch.Tensor,
        model_inputs: dict,
        return_logits: Optional[bool] = False,
-        response_masks: Optional[torch.Tensor] = None
+        response_masks: Optional[torch.Tensor] = None,
    ):
        r"""
        Calculates model outputs in multiple batches.
@@ -289,7 +318,8 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
            with torch.cuda.amp.autocast(dtype=self.model_args.compute_dtype):  # support bf16
                logits, _, values = model(**input_kwargs)

-            if values.size(0) != input_ids.size(0): # adapt to chatglm2
+            unwrapped_model: "AutoModelForCausalLMWithValueHead" = self.accelerator.unwrap_model(self.model)
+            if getattr(unwrapped_model.config, "model_type", None) == "chatglm":
                values = torch.transpose(values, 0, 1)

            logprobs = logprobs_from_logits(logits[:, :-1, :], input_ids[:, 1:])
@@ -303,9 +333,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
                end = start + len(response_batch[j])

                if response_masks is not None:
-                    response_masks_batch = torch.cat(
-                        (torch.zeros_like(query_batch[j]), response_masks_batch[j])
-                    )[1:]
+                    response_masks_batch = torch.cat((torch.zeros_like(query_batch[j]), response_masks_batch[j]))[1:]

                masks[j, :start] = 0
                masks[j, end:] = 0
@@ -335,4 +363,13 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
        Subclass and override to inject custom behavior.
        """
        if self.args.should_save:
-            self._save(output_dir)
+            try:
+                self._save(output_dir, state_dict=self.accelerator.get_state_dict(self.model))
+            except ValueError:
+                logger.warning(
+                    " stage3_gather_16bit_weights_on_model_save=false. Saving the full checkpoint instead,"
+                    " use zero_to_fp32.py to recover weights"
+                )
+                self._save(output_dir, state_dict={})
+                remove_dummy_checkpoint(True, output_dir, [WEIGHTS_NAME, SAFE_WEIGHTS_NAME])
+                self.model.save_checkpoint(output_dir)
--- a/src/llmtuner/train/ppo/utils.py
+++ b/src/llmtuner/train/ppo/utils.py
@@ -1,10 +1,26 @@
+import json
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional
+
 import torch
-from typing import TYPE_CHECKING, Dict, Literal, Optional
+
+from ...extras.packages import is_requests_available
+

 if TYPE_CHECKING:
    from transformers import PreTrainedModel
    from trl import AutoModelForCausalLMWithValueHead

+if is_requests_available():
+    import requests
+
+
+def get_rewards_from_server(server_url: str, messages: List[str]) -> List[torch.Tensor]:
+    headers = {"Content-Type": "application/json"}
+    payload = {"model": "model", "messages": messages}
+    response = requests.post(server_url, json=payload, headers=headers)
+    rewards = json.loads(response.text)["scores"]
+    return torch.Tensor(rewards)
+

 def replace_model(model: "AutoModelForCausalLMWithValueHead", target: Literal["default", "reward"]) -> None:
    if target == "reward":  # save default head temporarily
@@ -13,10 +29,12 @@ def replace_model(model: "AutoModelForCausalLMWithValueHead", target: Literal["d
        setattr(model, "default_head_bias", valuehead_state_dict["summary.bias"].detach().clone())

    model.pretrained_model.set_adapter(target)  # set the LoRA adapter to be active
-    model.v_head.load_state_dict({
+    model.v_head.load_state_dict(
+        {
            "summary.weight": model.get_buffer("{}_head_weight".format(target)).detach().clone(),
-        "summary.bias": model.get_buffer("{}_head_bias".format(target)).detach().clone()
-    })
+            "summary.bias": model.get_buffer("{}_head_bias".format(target)).detach().clone(),
+        }
+    )


 def dump_layernorm(model: "PreTrainedModel") -> Dict[str, torch.Tensor]:
--- a/src/llmtuner/train/ppo/workflow.py
+++ b/src/llmtuner/train/ppo/workflow.py
@@ -1,22 +1,26 @@
 # Inspired by: https://github.com/lvwerra/trl/blob/main/examples/research_projects/stack_llama/scripts/rl_training.py

 import math
-from trl import PPOConfig
+from typing import TYPE_CHECKING, List, Optional
+
 from torch.optim import AdamW
-from typing import TYPE_CHECKING, Optional, List
 from transformers import DataCollatorWithPadding
 from transformers.optimization import get_scheduler
+from trl import PPOConfig
+
+from ...data import get_dataset
+from ...extras.callbacks import FixValueHeadModelCallback
+from ...extras.misc import fix_valuehead_checkpoint
+from ...extras.ploting import plot_loss
+from ...model import load_model_and_tokenizer
+from ...train.ppo.trainer import CustomPPOTrainer
+from ...train.utils import create_ref_model, create_reward_model

-from llmtuner.data import get_dataset, preprocess_dataset
-from llmtuner.extras.callbacks import SavePeftModelCallback
-from llmtuner.extras.ploting import plot_loss
-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
-    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments
+
+    from ...hparams import DataArguments, FinetuningArguments, GeneratingArguments, ModelArguments


 def run_ppo(
@@ -25,27 +29,29 @@ def run_ppo(
    training_args: "Seq2SeqTrainingArguments",
    finetuning_args: "FinetuningArguments",
    generating_args: "GeneratingArguments",
-    callbacks: Optional[List["TrainerCallback"]] = None
+    callbacks: Optional[List["TrainerCallback"]] = None,
 ):
-    dataset = get_dataset(model_args, data_args)
-    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="ppo")
-    dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="ppo")
+    model, tokenizer = load_model_and_tokenizer(
+        model_args, finetuning_args, training_args.do_train, add_valuehead=True
+    )
+    dataset = get_dataset(tokenizer, model_args, data_args, training_args, stage="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")
+    ref_model = create_ref_model(model_args, finetuning_args, add_valuehead=True)
    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_device_cache=True,
@@ -54,7 +60,7 @@ def run_ppo(
        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}
+        accelerator_kwargs={"step_scheduler_with_optimizer": False},
    )

    # Create optimizer and scheduler
@@ -62,16 +68,14 @@ def run_ppo(
    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(
        training_args.lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=training_args.get_warmup_steps(num_training_steps),
-        num_training_steps=num_training_steps
+        num_training_steps=num_training_steps,
    )

    # Initialize our Trainer
@@ -80,7 +84,7 @@ def run_ppo(
        training_args=training_args,
        finetuning_args=finetuning_args,
        generating_args=generating_args,
-        callbacks=callbacks + [SavePeftModelCallback()],
+        callbacks=callbacks + [FixValueHeadModelCallback()],
        reward_model=reward_model,
        config=ppo_config,
        model=model,
@@ -89,13 +93,15 @@ def run_ppo(
        dataset=dataset,
        data_collator=data_collator,
        optimizer=optimizer,
-        lr_scheduler=lr_scheduler
+        lr_scheduler=lr_scheduler,
    )

    # 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()
+        if training_args.should_save:
+            fix_valuehead_checkpoint(model, training_args.output_dir, training_args.save_safetensors)
        ppo_trainer.save_state()  # must be called after save_model to have a folder
        if ppo_trainer.is_world_process_zero() and finetuning_args.plot_loss:
            plot_loss(training_args.output_dir, keys=["loss", "reward"])
--- a/src/llmtuner/train/pt/init.py
+++ b/src/llmtuner/train/pt/init.py
@@ -1 +1,4 @@
-from llmtuner.train.pt.workflow import run_pt
+from .workflow import run_pt
+
+
+__all__ = ["run_pt"]
--- a/src/llmtuner/train/pt/workflow.py
+++ b/src/llmtuner/train/pt/workflow.py
@@ -1,16 +1,20 @@
 # Inspired by: https://github.com/huggingface/transformers/blob/v4.34.1/examples/pytorch/language-modeling/run_clm.py

 import math
-from typing import TYPE_CHECKING, Optional, List
+from typing import TYPE_CHECKING, List, Optional
+
 from transformers import DataCollatorForLanguageModeling, Trainer

-from llmtuner.data import get_dataset, preprocess_dataset, split_dataset
-from llmtuner.extras.ploting import plot_loss
-from llmtuner.model import generate_model_card, load_model_and_tokenizer
+from ...data import get_dataset, split_dataset
+from ...extras.ploting import plot_loss
+from ...model import load_model_and_tokenizer
+from ...train.utils import create_modelcard_and_push
+

 if TYPE_CHECKING:
    from transformers import Seq2SeqTrainingArguments, TrainerCallback
-    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
+
+    from ...hparams import DataArguments, FinetuningArguments, ModelArguments


 def run_pt(
@@ -18,11 +22,10 @@ def run_pt(
    data_args: "DataArguments",
    training_args: "Seq2SeqTrainingArguments",
    finetuning_args: "FinetuningArguments",
-    callbacks: Optional[List["TrainerCallback"]] = None
+    callbacks: Optional[List["TrainerCallback"]] = None,
 ):
-    dataset = get_dataset(model_args, data_args)
-    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="pt")
-    dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="pt")
+    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train)
+    dataset = get_dataset(tokenizer, model_args, data_args, training_args, stage="pt")
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    # Initialize our Trainer
@@ -32,7 +35,7 @@ def run_pt(
        tokenizer=tokenizer,
        data_collator=data_collator,
        callbacks=callbacks,
-        **split_dataset(dataset, data_args, training_args)
+        **split_dataset(dataset, data_args, training_args),
    )

    # Training
@@ -58,8 +61,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)
--- a/src/llmtuner/train/rm/init.py
+++ b/src/llmtuner/train/rm/init.py
@@ -1 +1,4 @@
-from llmtuner.train.rm.workflow import run_rm
+from .workflow import run_rm
+
+
+__all__ = ["run_rm"]
--- a/src/llmtuner/train/rm/collator.py
+++ b/src/llmtuner/train/rm/collator.py
@@ -1,6 +1,7 @@
-import torch
 from dataclasses import dataclass
 from typing import Any, Dict, Sequence
+
+import torch
 from transformers import DataCollatorWithPadding


@@ -20,8 +21,9 @@ class PairwiseDataCollatorWithPadding(DataCollatorWithPadding):
        features = [
            {
                "input_ids": feature["prompt_ids"] + feature[key],
-                "attention_mask": [1] * (len(feature["prompt_ids"]) + len(feature[key]))
+                "attention_mask": [1] * (len(feature["prompt_ids"]) + len(feature[key])),
            }
-            for key in ("chosen_ids", "rejected_ids") for feature in features
+            for key in ("chosen_ids", "rejected_ids")
+            for feature in features
        ]
        return super().__call__(features)
--- a/src/llmtuner/train/rm/metric.py
+++ b/src/llmtuner/train/rm/metric.py
@@ -1,6 +1,7 @@
-import numpy as np
 from typing import Dict, Sequence, Tuple, Union

+import numpy as np
+

 def compute_accuracy(eval_preds: Sequence[Union[np.ndarray, Tuple[np.ndarray]]]) -> Dict[str, float]:
    preds, _ = eval_preds
--- a/src/llmtuner/train/rm/trainer.py
+++ b/src/llmtuner/train/rm/trainer.py
@@ -1,14 +1,16 @@
-import os
 import json
-import torch
+import os
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+
+import torch
 from transformers import Trainer

-from llmtuner.extras.logging import get_logger
+from ...extras.logging import get_logger
+

 if TYPE_CHECKING:
-    from transformers.trainer import PredictionOutput
    from transformers.modeling_utils import PreTrainedModel
+    from transformers.trainer import PredictionOutput


 logger = get_logger(__name__)
@@ -24,10 +26,7 @@ class PairwiseTrainer(Trainer):
        self.can_return_loss = True  # override property to return eval_loss

    def compute_loss(
-        self,
-        model: "PreTrainedModel",
-        inputs: Dict[str, torch.Tensor],
-        return_outputs: Optional[bool] = False
+        self, model: "PreTrainedModel", inputs: Dict[str, torch.Tensor], return_outputs: Optional[bool] = False
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[torch.Tensor]]]:
        r"""
        Computes pairwise loss. The first n examples are chosen and the last n examples are rejected.
@@ -39,7 +38,9 @@ class PairwiseTrainer(Trainer):
        """
        # Compute rewards
        _, _, values = model(**inputs, output_hidden_states=True, return_dict=True)
-        if values.size(0) != inputs["input_ids"].size(0): # adapt to chatglm2
+
+        unwrapped_model: "PreTrainedModel" = self.accelerator.unwrap_model(self.model)
+        if getattr(unwrapped_model.config, "model_type", None) == "chatglm":
            values = torch.transpose(values, 0, 1)

        # Split the inputs and rewards into two parts, chosen and rejected
@@ -78,10 +79,7 @@ class PairwiseTrainer(Trainer):

        return loss

-    def save_predictions(
-        self,
-        predict_results: "PredictionOutput"
-    ) -> None:
+    def save_predictions(self, predict_results: "PredictionOutput") -> None:
        r"""
        Saves model predictions to `output_dir`.

--- a/src/llmtuner/train/rm/workflow.py
+++ b/src/llmtuner/train/rm/workflow.py
@@ -1,19 +1,24 @@
 # Inspired by: https://github.com/CarperAI/trlx/blob/main/examples/summarize_rlhf/reward_model/train_reward_model_gptj.py

-from typing import TYPE_CHECKING, Optional, List
+from typing import TYPE_CHECKING, List, Optional
+
 from transformers import Seq2SeqTrainingArguments

-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.model import generate_model_card, 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 ...data import get_dataset, split_dataset
+from ...extras.callbacks import FixValueHeadModelCallback
+from ...extras.misc import fix_valuehead_checkpoint
+from ...extras.ploting import plot_loss
+from ...model import load_model_and_tokenizer
+from ...train.rm.collator import PairwiseDataCollatorWithPadding
+from ...train.rm.metric import compute_accuracy
+from ...train.rm.trainer import PairwiseTrainer
+from ...train.utils import create_modelcard_and_push
+

 if TYPE_CHECKING:
    from transformers import TrainerCallback
-    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
+
+    from ...hparams import DataArguments, FinetuningArguments, ModelArguments


 def run_rm(
@@ -21,12 +26,13 @@ def run_rm(
    data_args: "DataArguments",
    training_args: "Seq2SeqTrainingArguments",
    finetuning_args: "FinetuningArguments",
-    callbacks: Optional[List["TrainerCallback"]] = None
+    callbacks: Optional[List["TrainerCallback"]] = None,
 ):
-    dataset = get_dataset(model_args, data_args)
-    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="rm")
-    dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="rm")
-    data_collator = PairwiseDataCollatorWithPadding(tokenizer, pad_to_multiple_of=4)
+    model, tokenizer = load_model_and_tokenizer(
+        model_args, finetuning_args, training_args.do_train, add_valuehead=True
+    )
+    dataset = get_dataset(tokenizer, model_args, data_args, training_args, stage="rm")
+    data_collator = PairwiseDataCollatorWithPadding(tokenizer, pad_to_multiple_of=8)

    # Update arguments
    training_args_dict = training_args.to_dict()
@@ -39,15 +45,17 @@ def run_rm(
        args=training_args,
        tokenizer=tokenizer,
        data_collator=data_collator,
-        callbacks=callbacks + [SavePeftModelCallback()],
+        callbacks=callbacks + [FixValueHeadModelCallback()],
        compute_metrics=compute_accuracy,
-        **split_dataset(dataset, data_args, training_args)
+        **split_dataset(dataset, data_args, training_args),
    )

    # 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()
+        if training_args.should_save:
+            fix_valuehead_checkpoint(model, training_args.output_dir, training_args.save_safetensors)
        trainer.log_metrics("train", train_result.metrics)
        trainer.save_metrics("train", train_result.metrics)
        trainer.save_state()
@@ -68,8 +76,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)
--- a/src/llmtuner/train/sft/init.py
+++ b/src/llmtuner/train/sft/init.py
@@ -1 +1,4 @@
-from llmtuner.train.sft.workflow import run_sft
+from .workflow import run_sft
+
+
+__all__ = ["run_sft"]
--- a/src/llmtuner/train/sft/metric.py
+++ b/src/llmtuner/train/sft/metric.py
@@ -1,11 +1,11 @@
-import numpy as np
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Dict, Sequence, Tuple, Union

-from llmtuner.extras.constants import IGNORE_INDEX
-from llmtuner.extras.packages import (
-    is_jieba_available, is_nltk_available, is_rouge_available
-)
+import numpy as np
+
+from ...extras.constants import IGNORE_INDEX
+from ...extras.packages import is_jieba_available, is_nltk_available, is_rouge_available
+

 if TYPE_CHECKING:
    from transformers.tokenization_utils import PreTrainedTokenizer
@@ -14,7 +14,7 @@ if is_jieba_available():
    import jieba

 if is_nltk_available():
-    from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
+    from nltk.translate.bleu_score import SmoothingFunction, sentence_bleu

 if is_rouge_available():
    from rouge_chinese import Rouge
--- a/src/llmtuner/train/sft/trainer.py
+++ b/src/llmtuner/train/sft/trainer.py
@@ -1,13 +1,15 @@
-import os
 import json
-import torch
-import numpy as np
-import torch.nn as nn
+import os
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
 from transformers import Seq2SeqTrainer

-from llmtuner.extras.constants import IGNORE_INDEX
-from llmtuner.extras.logging import get_logger
+from ...extras.constants import IGNORE_INDEX
+from ...extras.logging import get_logger
+

 if TYPE_CHECKING:
    from transformers.trainer import PredictionOutput
@@ -39,10 +41,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:
@@ -51,11 +53,7 @@ class CustomSeq2SeqTrainer(Seq2SeqTrainer):

        return loss, generated_tokens, labels

-    def _pad_tensors_to_target_len(
-        self,
-        src_tensor: torch.Tensor,
-        tgt_tensor: torch.Tensor
-    ) -> torch.Tensor:
+    def _pad_tensors_to_target_len(self, src_tensor: torch.Tensor, tgt_tensor: torch.Tensor) -> torch.Tensor:
        r"""
        Pads the tensor to the same length as the target tensor.
        """
@@ -64,10 +62,7 @@ class CustomSeq2SeqTrainer(Seq2SeqTrainer):
        padded_tensor[:, -src_tensor.shape[-1] :] = src_tensor  # adopt left-padding
        return padded_tensor.contiguous()  # in contiguous memory

-    def save_predictions(
-        self,
-        predict_results: "PredictionOutput"
-    ) -> None:
+    def save_predictions(self, predict_results: "PredictionOutput") -> None:
        r"""
        Saves model predictions to `output_dir`.

@@ -79,14 +74,27 @@ 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)
+        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))
--- a/src/llmtuner/train/sft/workflow.py
+++ b/src/llmtuner/train/sft/workflow.py
@@ -1,19 +1,23 @@
 # Inspired by: https://github.com/huggingface/transformers/blob/v4.34.1/examples/pytorch/summarization/run_summarization.py

-from typing import TYPE_CHECKING, Optional, List
+from typing import TYPE_CHECKING, List, Optional
+
 from transformers import DataCollatorForSeq2Seq, Seq2SeqTrainingArguments

-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.model import generate_model_card, load_model_and_tokenizer
-from llmtuner.train.sft.metric import ComputeMetrics
-from llmtuner.train.sft.trainer import CustomSeq2SeqTrainer
+from ...data import get_dataset, split_dataset
+from ...extras.constants import IGNORE_INDEX
+from ...extras.misc import get_logits_processor
+from ...extras.ploting import plot_loss
+from ...model import load_model_and_tokenizer
+from ...train.sft.metric import ComputeMetrics
+from ...train.sft.trainer import CustomSeq2SeqTrainer
+from ...train.utils import create_modelcard_and_push
+

 if TYPE_CHECKING:
    from transformers import TrainerCallback
-    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments
+
+    from ...hparams import DataArguments, FinetuningArguments, GeneratingArguments, ModelArguments


 def run_sft(
@@ -22,27 +26,31 @@ def run_sft(
    training_args: "Seq2SeqTrainingArguments",
    finetuning_args: "FinetuningArguments",
    generating_args: "GeneratingArguments",
-    callbacks: Optional[List["TrainerCallback"]] = None
+    callbacks: Optional[List["TrainerCallback"]] = None,
 ):
-    dataset = get_dataset(model_args, data_args)
-    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="sft")
-    dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="sft")
+    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train)
+    dataset = get_dataset(tokenizer, model_args, data_args, training_args, stage="sft")

    if training_args.predict_with_generate:
        tokenizer.padding_side = "left"  # use left-padding in generation

+    if getattr(model, "is_quantized", False) and not training_args.do_train:
+        setattr(model, "_hf_peft_config_loaded", True)  # hack here: make model compatible with prediction
+
    data_collator = DataCollatorForSeq2Seq(
        tokenizer=tokenizer,
-        pad_to_multiple_of=4 if tokenizer.padding_side == "right" else None, # for shift short attention
-        label_pad_token_id=IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
+        pad_to_multiple_of=8 if tokenizer.padding_side == "right" else None,  # for shift short attention
+        label_pad_token_id=IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id,
    )

    # Override the decoding parameters of Seq2SeqTrainer
    training_args_dict = training_args.to_dict()
-    training_args_dict.update(dict(
+    training_args_dict.update(
+        dict(
            generation_max_length=training_args.generation_max_length or data_args.cutoff_len,
-        generation_num_beams=data_args.eval_num_beams or training_args.generation_num_beams
-    ))
+            generation_num_beams=data_args.eval_num_beams or training_args.generation_num_beams,
+        )
+    )
    training_args = Seq2SeqTrainingArguments(**training_args_dict)

    # Initialize our Trainer
@@ -53,7 +61,7 @@ def run_sft(
        data_collator=data_collator,
        callbacks=callbacks,
        compute_metrics=ComputeMetrics(tokenizer) if training_args.predict_with_generate else None,
-        **split_dataset(dataset, data_args, training_args)
+        **split_dataset(dataset, data_args, training_args),
    )

    # Keyword arguments for `model.generate`
@@ -90,8 +98,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)
--- a/src/llmtuner/train/tuner.py
+++ b/src/llmtuner/train/tuner.py
@@ -1,13 +1,18 @@
 from typing import TYPE_CHECKING, Any, Dict, List, Optional

-from llmtuner.extras.callbacks import LogCallback
-from llmtuner.extras.logging import get_logger
-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
+import torch
+from transformers import PreTrainedModel
+
+from ..extras.callbacks import LogCallback
+from ..extras.logging import get_logger
+from ..hparams import get_infer_args, get_train_args
+from ..model import load_model_and_tokenizer
+from .dpo import run_dpo
+from .ppo import run_ppo
+from .pt import run_pt
+from .rm import run_rm
+from .sft import run_sft
+

 if TYPE_CHECKING:
    from transformers import TrainerCallback
@@ -34,16 +39,50 @@ def run_exp(args: Optional[Dict[str, Any]] = None, callbacks: Optional[List["Tra
        raise ValueError("Unknown task.")


-def export_model(args: Optional[Dict[str, Any]] = None, max_shard_size: Optional[str] = "10GB"):
+def export_model(args: Optional[Dict[str, Any]] = None):
    model_args, _, finetuning_args, _ = get_infer_args(args)
+
+    if model_args.export_dir is None:
+        raise ValueError("Please specify `export_dir`.")
+
+    if model_args.adapter_name_or_path is not None and model_args.export_quantization_bit is not None:
+        raise ValueError("Please merge adapters before quantizing the model.")
+
    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args)
-    model.config.use_cache = True
-    model.save_pretrained(finetuning_args.export_dir, max_shard_size=max_shard_size)
+
+    if getattr(model, "quantization_method", None) and model_args.adapter_name_or_path is not None:
+        raise ValueError("Cannot merge adapters to a quantized model.")
+
+    if not isinstance(model, PreTrainedModel):
+        raise ValueError("The model is not a `PreTrainedModel`, export aborted.")
+
+    setattr(model.config, "use_cache", True)
+    if getattr(model.config, "torch_dtype", None) == "bfloat16":
+        model = model.to(torch.bfloat16).to("cpu")
+    else:
+        model = model.to(torch.float16).to("cpu")
+        setattr(model.config, "torch_dtype", "float16")
+
+    model.save_pretrained(
+        save_directory=model_args.export_dir,
+        max_shard_size="{}GB".format(model_args.export_size),
+        safe_serialization=(not model_args.export_legacy_format),
+    )
+    if model_args.export_hub_model_id is not None:
+        model.push_to_hub(
+            model_args.export_hub_model_id,
+            token=model_args.hf_hub_token,
+            max_shard_size="{}GB".format(model_args.export_size),
+            safe_serialization=(not model_args.export_legacy_format),
+        )
+
    try:
        tokenizer.padding_side = "left"  # restore padding side
        tokenizer.init_kwargs["padding_side"] = "left"
-        tokenizer.save_pretrained(finetuning_args.export_dir)
-    except:
+        tokenizer.save_pretrained(model_args.export_dir)
+        if model_args.export_hub_model_id is not None:
+            tokenizer.push_to_hub(model_args.export_hub_model_id, token=model_args.hf_hub_token)
+    except Exception:
        logger.warning("Cannot save tokenizer, please copy the files manually.")


--- a/src/llmtuner/train/utils.py
+++ b/src/llmtuner/train/utils.py
@@ -1,22 +1,42 @@
-import torch
-from typing import TYPE_CHECKING, Literal, Union
+from typing import TYPE_CHECKING, Optional, Union
+
+import torch
+
+from ..extras.logging import get_logger
+from ..hparams import FinetuningArguments, ModelArguments
+from ..model import get_modelcard_args, load_model_and_tokenizer, load_valuehead_params

-from llmtuner.extras.logging import get_logger
-from llmtuner.hparams import ModelArguments, FinetuningArguments
-from llmtuner.model import 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 ..hparams import DataArguments
+

 logger = get_logger(__name__)


-def create_ref_model(
+def create_modelcard_and_push(
+    trainer: "Trainer",
    model_args: "ModelArguments",
+    data_args: "DataArguments",
+    training_args: "Seq2SeqTrainingArguments",
    finetuning_args: "FinetuningArguments",
-    stage: Literal["ppo", "dpo"]
+) -> 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", add_valuehead: Optional[bool] = False
 ) -> Union["PreTrainedModel", "AutoModelForCausalLMWithValueHead"]:
    r"""
    Creates reference model for PPO/DPO training. Evaluation mode is not supported.
@@ -25,34 +45,42 @@ def create_ref_model(
    """
    if finetuning_args.ref_model is not None:
        ref_model_args_dict = model_args.to_dict()
-        ref_model_args_dict.update(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
-        ))
+                adapter_name_or_path=finetuning_args.ref_model_adapters,
+                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)
+        ref_model, _ = load_model_and_tokenizer(
+            ref_model_args, ref_finetuning_args, is_trainable=False, add_valuehead=add_valuehead
+        )
        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)
+            ref_model, _ = load_model_and_tokenizer(
+                model_args, finetuning_args, is_trainable=False, add_valuehead=add_valuehead
+            )
            logger.info("Created reference model from the model itself.")

    return ref_model


 def create_reward_model(
-    model: "AutoModelForCausalLMWithValueHead",
-    model_args: "ModelArguments",
-    finetuning_args: "FinetuningArguments"
+    model: "AutoModelForCausalLMWithValueHead", model_args: "ModelArguments", finetuning_args: "FinetuningArguments"
 ) -> "AutoModelForCausalLMWithValueHead":
    r"""
    Creates reward model for PPO training.
    """
-    if finetuning_args.reward_model_type == "lora":
+    if finetuning_args.reward_model_type == "api":
+        assert finetuning_args.reward_model.startswith("http"), "Please provide full url."
+        logger.info("Use reward server {}".format(finetuning_args.reward_model))
+        return finetuning_args.reward_model
+    elif 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:
@@ -61,20 +89,28 @@ def create_reward_model(
        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)
+        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(
+        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
-        ))
+                adapter_name_or_path=finetuning_args.reward_model_adapters,
+                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))
+        reward_model, _ = load_model_and_tokenizer(
+            reward_model_args, reward_finetuning_args, is_trainable=False, add_valuehead=True
+        )
+        logger.info("Loaded 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
--- a/src/llmtuner/webui/init.py
+++ b/src/llmtuner/webui/init.py
@@ -1 +1,4 @@
-from llmtuner.webui.interface import create_ui, create_web_demo
+from .interface import create_ui, create_web_demo
+
+
+__all__ = ["create_ui", "create_web_demo"]
--- a/src/llmtuner/webui/chatter.py
+++ b/src/llmtuner/webui/chatter.py
@@ -1,24 +1,24 @@
+import json
+from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional, Sequence, Tuple
+
 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 import ChatModel
-from llmtuner.extras.misc import torch_gc
-from llmtuner.hparams import GeneratingArguments
-from llmtuner.webui.common import get_save_dir
-from llmtuner.webui.locales import ALERTS
+from ..chat import ChatModel
+from ..data import Role
+from ..extras.misc import torch_gc
+from ..hparams import GeneratingArguments
+from .common import get_save_dir
+from .locales import ALERTS
+

 if TYPE_CHECKING:
-    from llmtuner.webui.manager import Manager
+    from .manager import Manager


 class WebChatModel(ChatModel):
-
    def __init__(
-        self,
-        manager: "Manager",
-        demo_mode: Optional[bool] = False,
-        lazy_init: Optional[bool] = True
+        self, manager: "Manager", demo_mode: Optional[bool] = False, lazy_init: Optional[bool] = True
    ) -> None:
        self.manager = manager
        self.demo_mode = demo_mode
@@ -29,8 +29,18 @@ class WebChatModel(ChatModel):
        if not lazy_init:  # read arguments from command line
            super().__init__()

-        if demo_mode: # load openchat 3.5 by default
-            super().__init__(dict(model_name_or_path="openchat/openchat_3.5", template="openchat"))
+        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 Exception:
+                print("Cannot find `demo_config.json` at current directory.")

    @property
    def loaded(self) -> bool:
@@ -54,24 +64,26 @@ class WebChatModel(ChatModel):
            yield error
            return

-        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")
-            ])
+        if get("top.adapter_path"):
+            adapter_name_or_path = ",".join(
+                [
+                    get_save_dir(get("top.model_name"), get("top.finetuning_type"), adapter)
+                    for adapter in get("top.adapter_path")
+                ]
+            )
        else:
-            checkpoint_dir = None
+            adapter_name_or_path = None

        yield ALERTS["info_loading"][lang]
        args = dict(
            model_name_or_path=get("top.model_path"),
-            checkpoint_dir=checkpoint_dir,
+            adapter_name_or_path=adapter_name_or_path,
            finetuning_type=get("top.finetuning_type"),
            quantization_bit=int(get("top.quantization_bit")) if get("top.quantization_bit") in ["8", "4"] else None,
            template=get("top.template"),
-            system_prompt=get("top.system_prompt"),
-            flash_attn=get("top.flash_attn"),
-            shift_attn=get("top.shift_attn"),
-            rope_scaling=get("top.rope_scaling") if get("top.rope_scaling") in ["linear", "dynamic"] else None
+            flash_attn=(get("top.booster") == "flash_attn"),
+            use_unsloth=(get("top.booster") == "unsloth"),
+            rope_scaling=get("top.rope_scaling") if get("top.rope_scaling") in ["linear", "dynamic"] else None,
        )
        super().__init__(args)

@@ -81,6 +93,7 @@ class WebChatModel(ChatModel):
        lang = data[self.manager.get_elem_by_name("top.lang")]

        if self.demo_mode:
+            gr.Warning(ALERTS["err_demo"][lang])
            yield ALERTS["err_demo"][lang]
            return

@@ -94,21 +107,37 @@ class WebChatModel(ChatModel):
        self,
        chatbot: List[Tuple[str, str]],
        query: str,
-        history: List[Tuple[str, str]],
+        messages: Sequence[Tuple[str, str]],
        system: str,
+        tools: str,
        max_new_tokens: int,
        top_p: float,
-        temperature: float
-    ) -> Generator[Tuple[List[Tuple[str, str]], List[Tuple[str, str]]], None, None]:
+        temperature: float,
+    ) -> Generator[Tuple[Sequence[Tuple[str, str]], Sequence[Tuple[str, str]]], None, None]:
        chatbot.append([query, ""])
+        query_messages = messages + [{"role": Role.USER, "content": query}]
        response = ""
        for new_text in self.stream_chat(
-            query, history, system, max_new_tokens=max_new_tokens, top_p=top_p, temperature=temperature
+            query_messages, system, tools, max_new_tokens=max_new_tokens, top_p=top_p, temperature=temperature
        ):
            response += new_text
-            new_history = history + [(query, response)]
-            chatbot[-1] = [query, self.postprocess(response)]
-            yield chatbot, new_history
+            if tools:
+                result = self.template.format_tools.extract(response)
+            else:
+                result = response
+
+            if isinstance(result, tuple):
+                name, arguments = result
+                arguments = json.loads(arguments)
+                tool_call = json.dumps({"name": name, "arguments": arguments}, ensure_ascii=False)
+                output_messages = query_messages + [{"role": Role.FUNCTION, "content": tool_call}]
+                bot_text = "```json\n" + tool_call + "\n```"
+            else:
+                output_messages = query_messages + [{"role": Role.ASSISTANT, "content": result}]
+                bot_text = result
+
+            chatbot[-1] = [query, self.postprocess(bot_text)]
+            yield chatbot, output_messages

    def postprocess(self, response: str) -> str:
        blocks = response.split("```")
--- a/src/llmtuner/webui/common.py
+++ b/src/llmtuner/webui/common.py
@@ -1,32 +1,28 @@
-import os
 import json
-import gradio as gr
+import os
+from collections import defaultdict
 from typing import Any, Dict, Optional
-from transformers.utils import (
-    WEIGHTS_NAME,
-    WEIGHTS_INDEX_NAME,
-    SAFE_WEIGHTS_NAME,
-    SAFE_WEIGHTS_INDEX_NAME,
-    ADAPTER_WEIGHTS_NAME,
-    ADAPTER_SAFE_WEIGHTS_NAME
+
+import gradio as gr
+from peft.utils import SAFETENSORS_WEIGHTS_NAME, WEIGHTS_NAME
+
+from ..extras.constants import (
+    DATA_CONFIG,
+    DEFAULT_MODULE,
+    DEFAULT_TEMPLATE,
+    PEFT_METHODS,
+    SUPPORTED_MODELS,
+    TRAINING_STAGES,
+    DownloadSource,
 )
-
-from llmtuner.extras.constants import DEFAULT_MODULE, DEFAULT_TEMPLATE, SUPPORTED_MODELS, TRAINING_STAGES
+from ..extras.misc import use_modelscope


+ADAPTER_NAMES = {WEIGHTS_NAME, SAFETENSORS_WEIGHTS_NAME}
 DEFAULT_CACHE_DIR = "cache"
 DEFAULT_DATA_DIR = "data"
 DEFAULT_SAVE_DIR = "saves"
 USER_CONFIG = "user.config"
-DATA_CONFIG = "dataset_info.json"
-CKPT_NAMES = [
-    WEIGHTS_NAME,
-    WEIGHTS_INDEX_NAME,
-    SAFE_WEIGHTS_NAME,
-    SAFE_WEIGHTS_INDEX_NAME,
-    ADAPTER_WEIGHTS_NAME,
-    ADAPTER_SAFE_WEIGHTS_NAME
-]


 def get_save_dir(*args) -> os.PathLike:
@@ -41,7 +37,7 @@ def load_config() -> Dict[str, Any]:
    try:
        with open(get_config_path(), "r", encoding="utf-8") as f:
            return json.load(f)
-    except:
+    except Exception:
        return {"lang": None, "last_model": None, "path_dict": {}, "cache_dir": None}


@@ -58,7 +54,15 @@ def save_config(lang: str, model_name: Optional[str] = None, model_path: Optiona

 def get_model_path(model_name: str) -> str:
    user_config = load_config()
-    return user_config["path_dict"].get(model_name, None) or SUPPORTED_MODELS.get(model_name, "")
+    path_dict: Dict[DownloadSource, str] = SUPPORTED_MODELS.get(model_name, defaultdict(str))
+    model_path = user_config["path_dict"].get(model_name, None) or path_dict.get(DownloadSource.DEFAULT, None)
+    if (
+        use_modelscope()
+        and path_dict.get(DownloadSource.MODELSCOPE)
+        and model_path == path_dict.get(DownloadSource.DEFAULT)
+    ):  # replace path
+        model_path = path_dict.get(DownloadSource.MODELSCOPE)
+    return model_path


 def get_prefix(model_name: str) -> str:
@@ -75,26 +79,28 @@ def get_template(model_name: str) -> str:
    return "default"


-def list_checkpoint(model_name: str, finetuning_type: str) -> Dict[str, Any]:
-    checkpoints = []
-    if model_name:
+def list_adapters(model_name: str, finetuning_type: str) -> Dict[str, Any]:
+    if finetuning_type not in PEFT_METHODS:
+        return gr.update(value=[], choices=[], interactive=False)
+
+    adapters = []
+    if model_name and finetuning_type == "lora":
        save_dir = get_save_dir(model_name, finetuning_type)
        if save_dir and os.path.isdir(save_dir):
-            for checkpoint in os.listdir(save_dir):
-                if (
-                    os.path.isdir(os.path.join(save_dir, checkpoint))
-                    and any([os.path.isfile(os.path.join(save_dir, checkpoint, name)) for name in CKPT_NAMES])
+            for adapter in os.listdir(save_dir):
+                if os.path.isdir(os.path.join(save_dir, adapter)) and any(
+                    os.path.isfile(os.path.join(save_dir, adapter, name)) for name in ADAPTER_NAMES
                ):
-                    checkpoints.append(checkpoint)
-    return gr.update(value=[], choices=checkpoints)
+                    adapters.append(adapter)
+    return gr.update(value=[], choices=adapters, interactive=True)


-def load_dataset_info(dataset_dir: str) -> Dict[str, Any]:
+def load_dataset_info(dataset_dir: str) -> Dict[str, Dict[str, Any]]:
    try:
        with open(os.path.join(dataset_dir, DATA_CONFIG), "r", encoding="utf-8") as f:
            return json.load(f)
-    except:
-        print("Cannot find {} in {}.".format(DATA_CONFIG, dataset_dir))
+    except Exception as err:
+        print("Cannot open {} due to {}.".format(os.path.join(dataset_dir, DATA_CONFIG), str(err)))
        return {}


--- a/src/llmtuner/webui/components/init.py
+++ b/src/llmtuner/webui/components/init.py
@@ -1,6 +1,16 @@
-from llmtuner.webui.components.top import create_top
-from llmtuner.webui.components.train import create_train_tab
-from llmtuner.webui.components.eval import create_eval_tab
-from llmtuner.webui.components.infer import create_infer_tab
-from llmtuner.webui.components.export import create_export_tab
-from llmtuner.webui.components.chatbot import create_chat_box
+from .chatbot import create_chat_box
+from .eval import create_eval_tab
+from .export import create_export_tab
+from .infer import create_infer_tab
+from .top import create_top
+from .train import create_train_tab
+
+
+__all__ = [
+    "create_chat_box",
+    "create_eval_tab",
+    "create_export_tab",
+    "create_infer_tab",
+    "create_top",
+    "create_train_tab",
+]
--- a/src/llmtuner/webui/components/chatbot.py
+++ b/src/llmtuner/webui/components/chatbot.py
@@ -1,22 +1,27 @@
-import gradio as gr
 from typing import TYPE_CHECKING, Dict, Optional, Tuple

+import gradio as gr
+
+from ..utils import check_json_schema
+
+
 if TYPE_CHECKING:
    from gradio.blocks import Block
    from gradio.components import Component
-    from llmtuner.webui.engine import Engine
+
+    from ..engine import Engine


 def create_chat_box(
-    engine: "Engine",
-    visible: Optional[bool] = False
+    engine: "Engine", visible: Optional[bool] = False
 ) -> Tuple["Block", "Component", "Component", Dict[str, "Component"]]:
    with gr.Box(visible=visible) as chat_box:
        chatbot = gr.Chatbot()
-        history = gr.State([])
+        messages = gr.State([])
        with gr.Row():
            with gr.Column(scale=4):
                system = gr.Textbox(show_label=False)
+                tools = gr.Textbox(show_label=False, lines=2)
                query = gr.Textbox(show_label=False, lines=8)
                submit_btn = gr.Button(variant="primary")

@@ -27,23 +32,29 @@ def create_chat_box(
                top_p = gr.Slider(0.01, 1, value=gen_kwargs.top_p, step=0.01)
                temperature = gr.Slider(0.01, 1.5, value=gen_kwargs.temperature, step=0.01)

+    tools.input(check_json_schema, [tools, engine.manager.get_elem_by_name("top.lang")])
+
    submit_btn.click(
        engine.chatter.predict,
-        [chatbot, query, history, system, max_new_tokens, top_p, temperature],
-        [chatbot, history],
-        show_progress=True
-    ).then(
-        lambda: gr.update(value=""), outputs=[query]
-    )
+        [chatbot, query, messages, system, tools, max_new_tokens, top_p, temperature],
+        [chatbot, messages],
+        show_progress=True,
+    ).then(lambda: gr.update(value=""), outputs=[query])

-    clear_btn.click(lambda: ([], []), outputs=[chatbot, history], show_progress=True)
+    clear_btn.click(lambda: ([], []), outputs=[chatbot, messages], show_progress=True)

-    return chat_box, chatbot, history, dict(
+    return (
+        chat_box,
+        chatbot,
+        messages,
+        dict(
            system=system,
+            tools=tools,
            query=query,
            submit_btn=submit_btn,
            clear_btn=clear_btn,
            max_new_tokens=max_new_tokens,
            top_p=top_p,
-        temperature=temperature
+            temperature=temperature,
+        ),
    )
--- a/src/llmtuner/webui/components/data.py
+++ b/src/llmtuner/webui/components/data.py
@@ -1,9 +1,11 @@
-import os
 import json
-import gradio as gr
+import os
 from typing import TYPE_CHECKING, Any, Dict, Tuple

-from llmtuner.webui.common import DATA_CONFIG
+import gradio as gr
+
+from ...extras.constants import DATA_CONFIG
+

 if TYPE_CHECKING:
    from gradio.components import Component
@@ -21,8 +23,11 @@ def next_page(page_index: int, total_num: int) -> int:


 def can_preview(dataset_dir: str, dataset: list) -> Dict[str, Any]:
+    try:
        with open(os.path.join(dataset_dir, DATA_CONFIG), "r", encoding="utf-8") as f:
            dataset_info = json.load(f)
+    except Exception:
+        return gr.update(interactive=False)

    if (
        len(dataset) > 0
@@ -45,7 +50,7 @@ def get_preview(dataset_dir: str, dataset: list, page_index: int) -> Tuple[int,
        elif data_file.endswith(".jsonl"):
            data = [json.loads(line) for line in f]
        else:
-            data = [line for line in f]
+            data = [line for line in f]  # noqa: C416
    return len(data), data[PAGE_SIZE * page_index : PAGE_SIZE * (page_index + 1)], gr.update(visible=True)


@@ -64,32 +69,17 @@ def create_preview_box(dataset_dir: "gr.Textbox", dataset: "gr.Dropdown") -> Dic
        with gr.Row():
            preview_samples = gr.JSON(interactive=False)

-    dataset.change(
-        can_preview, [dataset_dir, dataset], [data_preview_btn], queue=False
-    ).then(
+    dataset.change(can_preview, [dataset_dir, dataset], [data_preview_btn], queue=False).then(
        lambda: 0, outputs=[page_index], queue=False
    )
    data_preview_btn.click(
-        get_preview,
-        [dataset_dir, dataset, page_index],
-        [preview_count, preview_samples, preview_box],
-        queue=False
+        get_preview, [dataset_dir, dataset, page_index], [preview_count, preview_samples, preview_box], queue=False
    )
-    prev_btn.click(
-        prev_page, [page_index], [page_index], queue=False
-    ).then(
-        get_preview,
-        [dataset_dir, dataset, page_index],
-        [preview_count, preview_samples, preview_box],
-        queue=False
+    prev_btn.click(prev_page, [page_index], [page_index], queue=False).then(
+        get_preview, [dataset_dir, dataset, page_index], [preview_count, preview_samples, preview_box], queue=False
    )
-    next_btn.click(
-        next_page, [page_index, preview_count], [page_index], queue=False
-    ).then(
-        get_preview,
-        [dataset_dir, dataset, page_index],
-        [preview_count, preview_samples, preview_box],
-        queue=False
+    next_btn.click(next_page, [page_index, preview_count], [page_index], queue=False).then(
+        get_preview, [dataset_dir, dataset, page_index], [preview_count, preview_samples, preview_box], queue=False
    )
    close_btn.click(lambda: gr.update(visible=False), outputs=[preview_box], queue=False)
    return dict(
@@ -99,5 +89,5 @@ def create_preview_box(dataset_dir: "gr.Textbox", dataset: "gr.Dropdown") -> Dic
        prev_btn=prev_btn,
        next_btn=next_btn,
        close_btn=close_btn,
-        preview_samples=preview_samples
+        preview_samples=preview_samples,
    )
--- a/src/llmtuner/webui/components/eval.py
+++ b/src/llmtuner/webui/components/eval.py
@@ -1,12 +1,15 @@
-import gradio as gr
 from typing import TYPE_CHECKING, Dict

-from llmtuner.webui.common import list_dataset, DEFAULT_DATA_DIR
-from llmtuner.webui.components.data import create_preview_box
+import gradio as gr
+
+from ..common import DEFAULT_DATA_DIR, list_dataset
+from .data import create_preview_box
+

 if TYPE_CHECKING:
    from gradio.components import Component
-    from llmtuner.webui.engine import Engine
+
+    from ..engine import Engine


 def create_eval_tab(engine: "Engine") -> Dict[str, "Component"]:
@@ -30,19 +33,16 @@ def create_eval_tab(engine: "Engine") -> Dict[str, "Component"]:
        predict = gr.Checkbox(value=True)

    input_elems.update({cutoff_len, max_samples, batch_size, predict})
-    elem_dict.update(dict(
-        cutoff_len=cutoff_len, max_samples=max_samples, batch_size=batch_size, predict=predict
-    ))
+    elem_dict.update(dict(cutoff_len=cutoff_len, max_samples=max_samples, batch_size=batch_size, predict=predict))

    with gr.Row():
        max_new_tokens = gr.Slider(10, 2048, value=128, step=1)
        top_p = gr.Slider(0.01, 1, value=0.7, step=0.01)
        temperature = gr.Slider(0.01, 1.5, value=0.95, step=0.01)
+        output_dir = gr.Textbox()

-    input_elems.update({max_new_tokens, top_p, temperature})
-    elem_dict.update(dict(
-        max_new_tokens=max_new_tokens, top_p=top_p, temperature=temperature
-    ))
+    input_elems.update({max_new_tokens, top_p, temperature, output_dir})
+    elem_dict.update(dict(max_new_tokens=max_new_tokens, top_p=top_p, temperature=temperature, output_dir=output_dir))

    with gr.Row():
        cmd_preview_btn = gr.Button()
@@ -57,10 +57,16 @@ def create_eval_tab(engine: "Engine") -> Dict[str, "Component"]:
        output_box = gr.Markdown()

    output_elems = [output_box, process_bar]
-    elem_dict.update(dict(
-        cmd_preview_btn=cmd_preview_btn, start_btn=start_btn, stop_btn=stop_btn,
-        resume_btn=resume_btn, process_bar=process_bar, output_box=output_box
-    ))
+    elem_dict.update(
+        dict(
+            cmd_preview_btn=cmd_preview_btn,
+            start_btn=start_btn,
+            stop_btn=stop_btn,
+            resume_btn=resume_btn,
+            process_bar=process_bar,
+            output_box=output_box,
+        )
+    )

    cmd_preview_btn.click(engine.runner.preview_eval, input_elems, output_elems)
    start_btn.click(engine.runner.run_eval, input_elems, output_elems)
--- a/src/llmtuner/webui/components/export.py
+++ b/src/llmtuner/webui/components/export.py
@@ -1,58 +1,80 @@
-import gradio as gr
 from typing import TYPE_CHECKING, Dict, Generator, List

-from llmtuner.train import export_model
-from llmtuner.webui.common import get_save_dir
-from llmtuner.webui.locales import ALERTS
+import gradio as gr
+
+from ...train import export_model
+from ..common import get_save_dir
+from ..locales import ALERTS
+

 if TYPE_CHECKING:
    from gradio.components import Component
-    from llmtuner.webui.engine import Engine
+
+    from ..engine import Engine
+
+
+GPTQ_BITS = ["8", "4", "3", "2"]


 def save_model(
    lang: str,
    model_name: str,
    model_path: str,
-    checkpoints: List[str],
+    adapter_path: List[str],
    finetuning_type: str,
    template: str,
    max_shard_size: int,
-    export_dir: str
+    export_quantization_bit: int,
+    export_quantization_dataset: str,
+    export_dir: str,
 ) -> Generator[str, None, None]:
    error = ""
    if not model_name:
        error = ALERTS["err_no_model"][lang]
    elif not model_path:
        error = ALERTS["err_no_path"][lang]
-    elif not checkpoints:
-        error = ALERTS["err_no_checkpoint"][lang]
    elif not export_dir:
        error = ALERTS["err_no_export_dir"][lang]
+    elif export_quantization_bit in GPTQ_BITS and not export_quantization_dataset:
+        error = ALERTS["err_no_dataset"][lang]
+    elif export_quantization_bit not in GPTQ_BITS and not adapter_path:
+        error = ALERTS["err_no_adapter"][lang]

    if error:
        gr.Warning(error)
        yield error
        return

+    if adapter_path:
+        adapter_name_or_path = ",".join(
+            [get_save_dir(model_name, finetuning_type, adapter) for adapter in adapter_path]
+        )
+    else:
+        adapter_name_or_path = None
+
    args = dict(
        model_name_or_path=model_path,
-        checkpoint_dir=",".join([get_save_dir(model_name, finetuning_type, ckpt) for ckpt in checkpoints]),
+        adapter_name_or_path=adapter_name_or_path,
        finetuning_type=finetuning_type,
        template=template,
-        export_dir=export_dir
+        export_dir=export_dir,
+        export_size=max_shard_size,
+        export_quantization_bit=int(export_quantization_bit) if export_quantization_bit in GPTQ_BITS else None,
+        export_quantization_dataset=export_quantization_dataset,
    )

    yield ALERTS["info_exporting"][lang]
-    export_model(args, max_shard_size="{}GB".format(max_shard_size))
+    export_model(args)
    yield ALERTS["info_exported"][lang]


 def create_export_tab(engine: "Engine") -> Dict[str, "Component"]:
    with gr.Row():
-        export_dir = gr.Textbox()
-        max_shard_size = gr.Slider(value=10, minimum=1, maximum=100)
+        max_shard_size = gr.Slider(value=1, minimum=1, maximum=100)
+        export_quantization_bit = gr.Dropdown(choices=["none", "8", "4", "3", "2"], value="none")
+        export_quantization_dataset = gr.Textbox(value="data/c4_demo.json")

+    export_dir = gr.Textbox()
    export_btn = gr.Button()
    info_box = gr.Textbox(show_label=False, interactive=False)

@@ -62,18 +84,22 @@ def create_export_tab(engine: "Engine") -> Dict[str, "Component"]:
            engine.manager.get_elem_by_name("top.lang"),
            engine.manager.get_elem_by_name("top.model_name"),
            engine.manager.get_elem_by_name("top.model_path"),
-            engine.manager.get_elem_by_name("top.checkpoints"),
+            engine.manager.get_elem_by_name("top.adapter_path"),
            engine.manager.get_elem_by_name("top.finetuning_type"),
            engine.manager.get_elem_by_name("top.template"),
            max_shard_size,
-            export_dir
+            export_quantization_bit,
+            export_quantization_dataset,
+            export_dir,
        ],
-        [info_box]
+        [info_box],
    )

    return dict(
-        export_dir=export_dir,
        max_shard_size=max_shard_size,
+        export_quantization_bit=export_quantization_bit,
+        export_quantization_dataset=export_quantization_dataset,
+        export_dir=export_dir,
        export_btn=export_btn,
-        info_box=info_box
+        info_box=info_box,
    )
--- a/src/llmtuner/webui/components/infer.py
+++ b/src/llmtuner/webui/components/infer.py
@@ -1,11 +1,14 @@
-import gradio as gr
 from typing import TYPE_CHECKING, Dict

-from llmtuner.webui.components.chatbot import create_chat_box
+import gradio as gr
+
+from .chatbot import create_chat_box
+

 if TYPE_CHECKING:
    from gradio.components import Component
-    from llmtuner.webui.engine import Engine
+
+    from ..engine import Engine


 def create_infer_tab(engine: "Engine") -> Dict[str, "Component"]:
@@ -22,18 +25,12 @@ def create_infer_tab(engine: "Engine") -> Dict[str, "Component"]:
    chat_box, chatbot, history, chat_elems = create_chat_box(engine, visible=False)
    elem_dict.update(dict(chat_box=chat_box, **chat_elems))

-    load_btn.click(
-        engine.chatter.load_model, input_elems, [info_box]
-    ).then(
+    load_btn.click(engine.chatter.load_model, input_elems, [info_box]).then(
        lambda: gr.update(visible=engine.chatter.loaded), outputs=[chat_box]
    )

-    unload_btn.click(
-        engine.chatter.unload_model, input_elems, [info_box]
-    ).then(
+    unload_btn.click(engine.chatter.unload_model, input_elems, [info_box]).then(
        lambda: ([], []), outputs=[chatbot, history]
-    ).then(
-        lambda: gr.update(visible=engine.chatter.loaded), outputs=[chat_box]
-    )
+    ).then(lambda: gr.update(visible=engine.chatter.loaded), outputs=[chat_box])

    return elem_dict
--- a/src/llmtuner/webui/components/top.py
+++ b/src/llmtuner/webui/components/top.py
@@ -1,10 +1,12 @@
-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.webui.common import get_model_path, get_template, list_checkpoint, save_config
-from llmtuner.webui.utils import can_quantize
+import gradio as gr
+
+from ...data import templates
+from ...extras.constants import METHODS, SUPPORTED_MODELS
+from ..common import get_model_path, get_template, list_adapters, save_config
+from ..utils import can_quantize
+

 if TYPE_CHECKING:
    from gradio.components import Component
@@ -20,25 +22,17 @@ def create_top() -> Dict[str, "Component"]:

    with gr.Row():
        finetuning_type = gr.Dropdown(choices=METHODS, value="lora", scale=1)
-        checkpoints = gr.Dropdown(multiselect=True, scale=5)
+        adapter_path = gr.Dropdown(multiselect=True, scale=5, allow_custom_value=True)
        refresh_btn = gr.Button(scale=1)

    with gr.Accordion(label="Advanced config", open=False) as advanced_tab:
        with gr.Row():
-            quantization_bit = gr.Dropdown(choices=["none", "8", "4"], value="none", scale=1)
-            template = gr.Dropdown(choices=list(templates.keys()), value="default", scale=1)
-            system_prompt = gr.Textbox(scale=2)
-
-    with gr.Accordion(label="Model config (LLaMA only)", open=False) as llama_tab:
-        with gr.Row():
-            with gr.Column():
-                flash_attn = gr.Checkbox(value=False)
-                shift_attn = gr.Checkbox(value=False)
+            quantization_bit = gr.Dropdown(choices=["none", "8", "4"], value="none")
+            template = gr.Dropdown(choices=list(templates.keys()), value="default")
            rope_scaling = gr.Radio(choices=["none", "linear", "dynamic"], value="none")
+            booster = gr.Radio(choices=["none", "flash_attn", "unsloth"], value="none")

-    model_name.change(
-        list_checkpoint, [model_name, finetuning_type], [checkpoints], queue=False
-    ).then(
+    model_name.change(list_adapters, [model_name, finetuning_type], [adapter_path], queue=False).then(
        get_model_path, [model_name], [model_path], queue=False
    ).then(
        get_template, [model_name], [template], queue=False
@@ -46,29 +40,22 @@ def create_top() -> Dict[str, "Component"]:

    model_path.change(save_config, inputs=[lang, model_name, model_path], queue=False)

-    finetuning_type.change(
-        list_checkpoint, [model_name, finetuning_type], [checkpoints], queue=False
-    ).then(
+    finetuning_type.change(list_adapters, [model_name, finetuning_type], [adapter_path], queue=False).then(
        can_quantize, [finetuning_type], [quantization_bit], queue=False
    )

-    refresh_btn.click(
-        list_checkpoint, [model_name, finetuning_type], [checkpoints], queue=False
-    )
+    refresh_btn.click(list_adapters, [model_name, finetuning_type], [adapter_path], queue=False)

    return dict(
        lang=lang,
        model_name=model_name,
        model_path=model_path,
        finetuning_type=finetuning_type,
-        checkpoints=checkpoints,
+        adapter_path=adapter_path,
        refresh_btn=refresh_btn,
        advanced_tab=advanced_tab,
        quantization_bit=quantization_bit,
        template=template,
-        system_prompt=system_prompt,
-        llama_tab=llama_tab,
-        flash_attn=flash_attn,
-        shift_attn=shift_attn,
-        rope_scaling=rope_scaling
+        rope_scaling=rope_scaling,
+        booster=booster,
    )
--- a/src/llmtuner/webui/components/train.py
+++ b/src/llmtuner/webui/components/train.py
@@ -1,15 +1,18 @@
-import gradio as gr
 from typing import TYPE_CHECKING, Dict
+
+import gradio as gr
 from transformers.trainer_utils import SchedulerType

-from llmtuner.extras.constants import TRAINING_STAGES
-from llmtuner.webui.common import list_checkpoint, list_dataset, DEFAULT_DATA_DIR
-from llmtuner.webui.components.data import create_preview_box
-from llmtuner.webui.utils import gen_plot
+from ...extras.constants import TRAINING_STAGES
+from ..common import DEFAULT_DATA_DIR, list_adapters, list_dataset
+from ..components.data import create_preview_box
+from ..utils import gen_plot
+

 if TYPE_CHECKING:
    from gradio.components import Component
-    from llmtuner.webui.engine import Engine
+
+    from ..engine import Engine


 def create_train_tab(engine: "Engine") -> Dict[str, "Component"]:
@@ -28,54 +31,67 @@ def create_train_tab(engine: "Engine") -> Dict[str, "Component"]:
    dataset_dir.change(list_dataset, [dataset_dir, training_stage], [dataset], queue=False)

    input_elems.update({training_stage, dataset_dir, dataset})
-    elem_dict.update(dict(
-        training_stage=training_stage, dataset_dir=dataset_dir, dataset=dataset, **preview_elems
-    ))
+    elem_dict.update(dict(training_stage=training_stage, dataset_dir=dataset_dir, dataset=dataset, **preview_elems))

    with gr.Row():
        cutoff_len = gr.Slider(value=1024, minimum=4, maximum=8192, step=1)
        learning_rate = gr.Textbox(value="5e-5")
        num_train_epochs = gr.Textbox(value="3.0")
        max_samples = gr.Textbox(value="100000")
-        compute_type = gr.Radio(choices=["fp16", "bf16"], value="fp16")
+        compute_type = gr.Radio(choices=["fp16", "bf16", "fp32"], value="fp16")

    input_elems.update({cutoff_len, learning_rate, num_train_epochs, max_samples, compute_type})
-    elem_dict.update(dict(
-        cutoff_len=cutoff_len, learning_rate=learning_rate, num_train_epochs=num_train_epochs,
-        max_samples=max_samples, compute_type=compute_type
-    ))
+    elem_dict.update(
+        dict(
+            cutoff_len=cutoff_len,
+            learning_rate=learning_rate,
+            num_train_epochs=num_train_epochs,
+            max_samples=max_samples,
+            compute_type=compute_type,
+        )
+    )

    with gr.Row():
        batch_size = gr.Slider(value=4, minimum=1, maximum=512, step=1)
        gradient_accumulation_steps = gr.Slider(value=4, minimum=1, maximum=512, step=1)
-        lr_scheduler_type = gr.Dropdown(
-            choices=[scheduler.value for scheduler in SchedulerType], value="cosine"
-        )
+        lr_scheduler_type = gr.Dropdown(choices=[scheduler.value for scheduler in SchedulerType], value="cosine")
        max_grad_norm = gr.Textbox(value="1.0")
        val_size = gr.Slider(value=0, minimum=0, maximum=1, step=0.001)

    input_elems.update({batch_size, gradient_accumulation_steps, lr_scheduler_type, max_grad_norm, val_size})
-    elem_dict.update(dict(
-        batch_size=batch_size, gradient_accumulation_steps=gradient_accumulation_steps,
-        lr_scheduler_type=lr_scheduler_type, max_grad_norm=max_grad_norm, val_size=val_size
-    ))
+    elem_dict.update(
+        dict(
+            batch_size=batch_size,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            lr_scheduler_type=lr_scheduler_type,
+            max_grad_norm=max_grad_norm,
+            val_size=val_size,
+        )
+    )

-    with gr.Accordion(label="Advanced config", open=False) as advanced_tab:
+    with gr.Accordion(label="Extra config", open=False) as extra_tab:
        with gr.Row():
            logging_steps = gr.Slider(value=5, minimum=5, maximum=1000, step=5)
            save_steps = gr.Slider(value=100, minimum=10, maximum=5000, step=10)
            warmup_steps = gr.Slider(value=0, minimum=0, maximum=5000, step=1)
-            neft_alpha = gr.Slider(value=0, minimum=0, maximum=10, step=0.1)
+            neftune_alpha = gr.Slider(value=0, minimum=0, maximum=10, step=0.1)

            with gr.Column():
-                train_on_prompt = gr.Checkbox(value=False)
+                sft_packing = gr.Checkbox(value=False)
                upcast_layernorm = gr.Checkbox(value=False)

-    input_elems.update({logging_steps, save_steps, warmup_steps, neft_alpha, train_on_prompt, upcast_layernorm})
-    elem_dict.update(dict(
-        advanced_tab=advanced_tab, logging_steps=logging_steps, save_steps=save_steps, warmup_steps=warmup_steps,
-        neft_alpha=neft_alpha, train_on_prompt=train_on_prompt, upcast_layernorm=upcast_layernorm
-    ))
+    input_elems.update({logging_steps, save_steps, warmup_steps, neftune_alpha, sft_packing, upcast_layernorm})
+    elem_dict.update(
+        dict(
+            extra_tab=extra_tab,
+            logging_steps=logging_steps,
+            save_steps=save_steps,
+            warmup_steps=warmup_steps,
+            neftune_alpha=neftune_alpha,
+            sft_packing=sft_packing,
+            upcast_layernorm=upcast_layernorm,
+        )
+    )

    with gr.Accordion(label="LoRA config", open=False) as lora_tab:
        with gr.Row():
@@ -83,29 +99,38 @@ def create_train_tab(engine: "Engine") -> Dict[str, "Component"]:
            lora_dropout = gr.Slider(value=0.1, minimum=0, maximum=1, step=0.01, scale=1)
            lora_target = gr.Textbox(scale=1)
            additional_target = gr.Textbox(scale=1)
-            resume_lora_training = gr.Checkbox(value=True, scale=1)
+            create_new_adapter = gr.Checkbox(scale=1)

-    input_elems.update({lora_rank, lora_dropout, lora_target, additional_target, resume_lora_training})
-    elem_dict.update(dict(
-        lora_tab=lora_tab, lora_rank=lora_rank, lora_dropout=lora_dropout, lora_target=lora_target,
-        additional_target=additional_target, resume_lora_training=resume_lora_training,
-    ))
+    input_elems.update({lora_rank, lora_dropout, lora_target, additional_target, create_new_adapter})
+    elem_dict.update(
+        dict(
+            lora_tab=lora_tab,
+            lora_rank=lora_rank,
+            lora_dropout=lora_dropout,
+            lora_target=lora_target,
+            additional_target=additional_target,
+            create_new_adapter=create_new_adapter,
+        )
+    )

    with gr.Accordion(label="RLHF config", open=False) as rlhf_tab:
        with gr.Row():
            dpo_beta = gr.Slider(value=0.1, minimum=0, maximum=1, step=0.01, scale=1)
-            reward_model = gr.Dropdown(scale=3)
+            dpo_ftx = gr.Slider(value=0, minimum=0, maximum=10, step=0.01, scale=1)
+            reward_model = gr.Dropdown(scale=2, allow_custom_value=True)
            refresh_btn = gr.Button(scale=1)

    refresh_btn.click(
-        list_checkpoint,
+        list_adapters,
        [engine.manager.get_elem_by_name("top.model_name"), engine.manager.get_elem_by_name("top.finetuning_type")],
        [reward_model],
-        queue=False
+        queue=False,
    )

-    input_elems.update({dpo_beta, reward_model})
-    elem_dict.update(dict(rlhf_tab=rlhf_tab, dpo_beta=dpo_beta, reward_model=reward_model, refresh_btn=refresh_btn))
+    input_elems.update({dpo_beta, dpo_ftx, reward_model})
+    elem_dict.update(
+        dict(rlhf_tab=rlhf_tab, dpo_beta=dpo_beta, dpo_ftx=dpo_ftx, reward_model=reward_model, refresh_btn=refresh_btn)
+    )

    with gr.Row():
        cmd_preview_btn = gr.Button()
@@ -135,20 +160,28 @@ def create_train_tab(engine: "Engine") -> Dict[str, "Component"]:
    stop_btn.click(engine.runner.set_abort, queue=False)
    resume_btn.change(engine.runner.monitor, outputs=output_elems)

-    elem_dict.update(dict(
-        cmd_preview_btn=cmd_preview_btn, start_btn=start_btn, stop_btn=stop_btn, output_dir=output_dir,
-        resume_btn=resume_btn, process_bar=process_bar, output_box=output_box, loss_viewer=loss_viewer
-    ))
+    elem_dict.update(
+        dict(
+            cmd_preview_btn=cmd_preview_btn,
+            start_btn=start_btn,
+            stop_btn=stop_btn,
+            output_dir=output_dir,
+            resume_btn=resume_btn,
+            process_bar=process_bar,
+            output_box=output_box,
+            loss_viewer=loss_viewer,
+        )
+    )

    output_box.change(
        gen_plot,
        [
            engine.manager.get_elem_by_name("top.model_name"),
            engine.manager.get_elem_by_name("top.finetuning_type"),
-            output_dir
+            output_dir,
        ],
        loss_viewer,
-        queue=False
+        queue=False,
    )

    return elem_dict
--- a/src/llmtuner/webui/engine.py
+++ b/src/llmtuner/webui/engine.py
@@ -1,18 +1,19 @@
-import gradio as gr
-from gradio.components import Component # cannot use TYPE_CHECKING here
 from typing import Any, Dict, Generator, Optional

-from llmtuner.webui.chatter import WebChatModel
-from llmtuner.webui.common import get_model_path, list_dataset, load_config
-from llmtuner.webui.locales import LOCALES
-from llmtuner.webui.manager import Manager
-from llmtuner.webui.runner import Runner
-from llmtuner.webui.utils import get_time
+import gradio as gr
+from gradio.components import Component  # cannot use TYPE_CHECKING here
+
+from .chatter import WebChatModel
+from .common import get_model_path, list_dataset, load_config
+from .locales import LOCALES
+from .manager import Manager
+from .runner import Runner
+from .utils import get_time


 class Engine:
-
    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()
        self.runner = Runner(self.manager, demo_mode=demo_mode)
@@ -22,13 +23,10 @@ class Engine:
        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 = {
-            "top.lang": {"value": lang},
-            "infer.chat_box": {"visible": self.chatter.loaded}
-        }
+        init_dict = {"top.lang": {"value": lang}, "infer.chat_box": {"visible": self.chatter.loaded}}

        if not self.pure_chat:
            init_dict["train.dataset"] = {"choices": list_dataset()["choices"]}
@@ -48,10 +46,17 @@ class Engine:
                else:
                    yield self._form_dict({"eval.resume_btn": {"value": True}})
            else:
-                yield self._form_dict({"train.output_dir": {"value": get_time()}})
+                yield self._form_dict(
+                    {
+                        "train.output_dir": {"value": "train_" + get_time()},
+                        "eval.output_dir": {"value": "eval_" + get_time()},
+                    }
+                )

    def change_lang(self, lang: str) -> Dict[Component, Dict[str, Any]]:
        return {
            component: gr.update(**LOCALES[name][lang])
-            for elems in self.manager.all_elems.values() for name, component in elems.items() if name in LOCALES
+            for elems in self.manager.all_elems.values()
+            for name, component in elems.items()
+            if name in LOCALES
        }
--- a/src/llmtuner/webui/interface.py
+++ b/src/llmtuner/webui/interface.py
@@ -1,21 +1,22 @@
-import gradio as gr
 from typing import Optional
+
+import gradio as gr
 from transformers.utils.versions import require_version

-from llmtuner.webui.components import (
+from .common import save_config
+from .components import (
+    create_chat_box,
+    create_eval_tab,
+    create_export_tab,
+    create_infer_tab,
    create_top,
    create_train_tab,
-    create_eval_tab,
-    create_infer_tab,
-    create_export_tab,
-    create_chat_box
 )
-from llmtuner.webui.common import save_config
-from llmtuner.webui.css import CSS
-from llmtuner.webui.engine import Engine
+from .css import CSS
+from .engine import Engine


-require_version("gradio>=3.38.0,<4.0.0", "To fix: pip install \"gradio>=3.38.0,<4.0.0\"")
+require_version("gradio>=3.38.0,<4.0.0", 'To fix: pip install "gradio>=3.38.0,<4.0.0"')


 def create_ui(demo_mode: Optional[bool] = False) -> gr.Blocks:
@@ -23,11 +24,9 @@ def create_ui(demo_mode: Optional[bool] = False) -> gr.Blocks:

    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(
-                "<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\">"
+                '<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")
@@ -75,4 +74,4 @@ def create_web_demo() -> gr.Blocks:
 if __name__ == "__main__":
    demo = create_ui()
    demo.queue()
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=False, inbrowser=True)
+    demo.launch(server_name="0.0.0.0", share=False, inbrowser=True)
--- a/src/llmtuner/webui/locales.py
+++ b/src/llmtuner/webui/locales.py
@@ -1,702 +1,222 @@
 LOCALES = {
-    "lang": {
-        "en": {
-            "label": "Lang"
-        },
-        "zh": {
-            "label": "语言"
-        }
-    },
-    "model_name": {
-        "en": {
-            "label": "Model name"
-        },
-        "zh": {
-            "label": "模型名称"
-        }
-    },
+    "lang": {"en": {"label": "Lang"}, "zh": {"label": "语言"}},
+    "model_name": {"en": {"label": "Model name"}, "zh": {"label": "模型名称"}},
    "model_path": {
-        "en": {
-            "label": "Model path",
-            "info": "Path to pretrained model or model identifier from Hugging Face."
-        },
-        "zh": {
-            "label": "模型路径",
-            "info": "本地模型的文件路径或 Hugging Face 的模型标识符。"
-        }
-    },
-    "finetuning_type": {
-        "en": {
-            "label": "Finetuning method"
-        },
-        "zh": {
-            "label": "微调方法"
-        }
-    },
-    "checkpoints": {
-        "en": {
-            "label": "Checkpoints"
-        },
-        "zh": {
-            "label": "模型断点"
-        }
-    },
-    "refresh_btn": {
-        "en": {
-            "value": "Refresh checkpoints"
-        },
-        "zh": {
-            "value": "刷新断点"
-        }
-    },
-    "advanced_tab": {
-        "en": {
-            "label": "Advanced configurations"
-        },
-        "zh": {
-            "label": "高级设置"
-        }
+        "en": {"label": "Model path", "info": "Path to pretrained model or model identifier from Hugging Face."},
+        "zh": {"label": "模型路径", "info": "本地模型的文件路径或 Hugging Face 的模型标识符。"},
    },
+    "finetuning_type": {"en": {"label": "Finetuning method"}, "zh": {"label": "微调方法"}},
+    "adapter_path": {"en": {"label": "Adapter path"}, "zh": {"label": "适配器路径"}},
+    "refresh_btn": {"en": {"value": "Refresh adapters"}, "zh": {"value": "刷新适配器"}},
+    "advanced_tab": {"en": {"label": "Advanced configurations"}, "zh": {"label": "高级设置"}},
    "quantization_bit": {
-        "en": {
-            "label": "Quantization bit",
-            "info": "Enable 4/8-bit model quantization (QLoRA)."
-        },
-        "zh": {
-            "label": "量化等级",
-            "info": "启用 4/8 比特模型量化（QLoRA）。"
-        }
+        "en": {"label": "Quantization bit", "info": "Enable 4/8-bit model quantization (QLoRA)."},
+        "zh": {"label": "量化等级", "info": "启用 4/8 比特模型量化（QLoRA）。"},
    },
    "template": {
-        "en": {
-            "label": "Prompt template",
-            "info": "The template used in constructing prompts."
-        },
-        "zh": {
-            "label": "提示模板",
-            "info": "构建提示词时使用的模板"
-        }
-    },
-    "system_prompt": {
-        "en": {
-            "label": "System prompt (optional)",
-            "info": "A sequence used as the default system prompt."
-        },
-        "zh": {
-            "label": "系统提示词（非必填）",
-            "info": "默认使用的系统提示词"
-        }
-    },
-    "llama_tab": {
-        "en": {
-            "label": "Model configurations (LLaMA only)"
-        },
-        "zh": {
-            "label": "模型设置（仅LLaMA）"
-        }
-    },
-    "flash_attn": {
-        "en": {
-            "label": "Use FlashAttention-2"
-        },
-        "zh": {
-            "label": "使用 FlashAttention-2"
-        }
-    },
-    "shift_attn": {
-        "en": {
-            "label": "Use shift short attention (S^2-Attn)"
-        },
-        "zh": {
-            "label": "使用 shift short attention (S^2-Attn)"
-        }
-    },
-    "rope_scaling": {
-        "en": {
-            "label": "RoPE scaling"
-        },
-        "zh": {
-            "label": "RoPE 插值方法"
-        }
+        "en": {"label": "Prompt template", "info": "The template used in constructing prompts."},
+        "zh": {"label": "提示模板", "info": "构建提示词时使用的模板"},
    },
+    "rope_scaling": {"en": {"label": "RoPE scaling"}, "zh": {"label": "RoPE 插值方法"}},
+    "booster": {"en": {"label": "Booster"}, "zh": {"label": "加速方式"}},
    "training_stage": {
-        "en": {
-            "label": "Stage",
-            "info": "The stage to perform in training."
-        },
-        "zh": {
-            "label": "训练阶段",
-            "info": "目前采用的训练方式。"
-        }
+        "en": {"label": "Stage", "info": "The stage to perform in training."},
+        "zh": {"label": "训练阶段", "info": "目前采用的训练方式。"},
    },
    "dataset_dir": {
-        "en": {
-            "label": "Data dir",
-            "info": "Path of the data directory."
-        },
-        "zh": {
-            "label": "数据路径",
-            "info": "数据文件夹的路径。"
-        }
-    },
-    "dataset": {
-        "en": {
-            "label": "Dataset"
-        },
-        "zh": {
-            "label": "数据集"
-        }
-    },
-    "data_preview_btn": {
-        "en": {
-            "value": "Preview dataset"
-        },
-        "zh": {
-            "value": "预览数据集"
-        }
-    },
-    "preview_count": {
-        "en": {
-            "label": "Count"
-        },
-        "zh": {
-            "label": "数量"
-        }
-    },
-    "page_index": {
-        "en": {
-            "label": "Page"
-        },
-        "zh": {
-            "label": "页数"
-        }
-    },
-    "prev_btn": {
-        "en": {
-            "value": "Prev"
-        },
-        "zh": {
-            "value": "上一页"
-        }
-    },
-    "next_btn": {
-        "en": {
-            "value": "Next"
-        },
-        "zh": {
-            "value": "下一页"
-        }
-    },
-    "close_btn": {
-        "en": {
-            "value": "Close"
-        },
-        "zh": {
-            "value": "关闭"
-        }
-    },
-    "preview_samples": {
-        "en": {
-            "label": "Samples"
-        },
-        "zh": {
-            "label": "样例"
-        }
+        "en": {"label": "Data dir", "info": "Path to the data directory."},
+        "zh": {"label": "数据路径", "info": "数据文件夹的路径。"},
    },
+    "dataset": {"en": {"label": "Dataset"}, "zh": {"label": "数据集"}},
+    "data_preview_btn": {"en": {"value": "Preview dataset"}, "zh": {"value": "预览数据集"}},
+    "preview_count": {"en": {"label": "Count"}, "zh": {"label": "数量"}},
+    "page_index": {"en": {"label": "Page"}, "zh": {"label": "页数"}},
+    "prev_btn": {"en": {"value": "Prev"}, "zh": {"value": "上一页"}},
+    "next_btn": {"en": {"value": "Next"}, "zh": {"value": "下一页"}},
+    "close_btn": {"en": {"value": "Close"}, "zh": {"value": "关闭"}},
+    "preview_samples": {"en": {"label": "Samples"}, "zh": {"label": "样例"}},
    "cutoff_len": {
-        "en": {
-            "label": "Cutoff length",
-            "info": "Max tokens in input sequence."
-        },
-        "zh": {
-            "label": "截断长度",
-            "info": "输入序列分词后的最大长度。"
-        }
+        "en": {"label": "Cutoff length", "info": "Max tokens in input sequence."},
+        "zh": {"label": "截断长度", "info": "输入序列分词后的最大长度。"},
    },
    "learning_rate": {
-        "en": {
-            "label": "Learning rate",
-            "info": "Initial learning rate for AdamW."
-        },
-        "zh": {
-            "label": "学习率",
-            "info": "AdamW 优化器的初始学习率。"
-        }
+        "en": {"label": "Learning rate", "info": "Initial learning rate for AdamW."},
+        "zh": {"label": "学习率", "info": "AdamW 优化器的初始学习率。"},
    },
    "num_train_epochs": {
-        "en": {
-            "label": "Epochs",
-            "info": "Total number of training epochs to perform."
-        },
-        "zh": {
-            "label": "训练轮数",
-            "info": "需要执行的训练总轮数。"
-        }
+        "en": {"label": "Epochs", "info": "Total number of training epochs to perform."},
+        "zh": {"label": "训练轮数", "info": "需要执行的训练总轮数。"},
    },
    "max_samples": {
-        "en": {
-            "label": "Max samples",
-            "info": "Maximum samples per dataset."
-        },
-        "zh": {
-            "label": "最大样本数",
-            "info": "每个数据集最多使用的样本数。"
-        }
+        "en": {"label": "Max samples", "info": "Maximum samples per dataset."},
+        "zh": {"label": "最大样本数", "info": "每个数据集最多使用的样本数。"},
    },
    "compute_type": {
-        "en": {
-            "label": "Compute type",
-            "info": "Whether to use fp16 or bf16 mixed precision training."
-        },
-        "zh": {
-            "label": "计算类型",
-            "info": "是否启用 FP16 或 BF16 混合精度训练。"
-        }
+        "en": {"label": "Compute type", "info": "Whether to use fp16 or bf16 mixed precision training."},
+        "zh": {"label": "计算类型", "info": "是否启用 FP16 或 BF16 混合精度训练。"},
    },
    "batch_size": {
-        "en": {
-            "label": "Batch size",
-            "info": "Number of samples to process per GPU."
-        },
-        "zh":{
-            "label": "批处理大小",
-            "info": "每块 GPU 上处理的样本数量。"
-        }
+        "en": {"label": "Batch size", "info": "Number of samples to process per GPU."},
+        "zh": {"label": "批处理大小", "info": "每块 GPU 上处理的样本数量。"},
    },
    "gradient_accumulation_steps": {
-        "en": {
-            "label": "Gradient accumulation",
-            "info": "Number of gradient accumulation steps."
-        },
-        "zh": {
-            "label": "梯度累积",
-            "info": "梯度累积的步数。"
-        }
+        "en": {"label": "Gradient accumulation", "info": "Number of gradient accumulation steps."},
+        "zh": {"label": "梯度累积", "info": "梯度累积的步数。"},
    },
    "lr_scheduler_type": {
        "en": {
            "label": "LR Scheduler",
            "info": "Name of learning rate scheduler.",
        },
-        "zh": {
-            "label": "学习率调节器",
-            "info": "采用的学习率调节器名称。"
-        }
+        "zh": {"label": "学习率调节器", "info": "采用的学习率调节器名称。"},
    },
    "max_grad_norm": {
-        "en": {
-            "label": "Maximum gradient norm",
-            "info": "Norm for gradient clipping.."
-        },
-        "zh": {
-            "label": "最大梯度范数",
-            "info": "用于梯度裁剪的范数。"
-        }
+        "en": {"label": "Maximum gradient norm", "info": "Norm for gradient clipping.."},
+        "zh": {"label": "最大梯度范数", "info": "用于梯度裁剪的范数。"},
    },
    "val_size": {
-        "en": {
-            "label": "Val size",
-            "info": "Proportion of data in the dev set."
-        },
-        "zh": {
-            "label": "验证集比例",
-            "info": "验证集占全部样本的百分比。"
-        }
+        "en": {"label": "Val size", "info": "Proportion of data in the dev set."},
+        "zh": {"label": "验证集比例", "info": "验证集占全部样本的百分比。"},
    },
+    "extra_tab": {"en": {"label": "Extra configurations"}, "zh": {"label": "其它参数设置"}},
    "logging_steps": {
-        "en": {
-            "label": "Logging steps",
-            "info": "Number of steps between two logs."
-        },
-        "zh": {
-            "label": "日志间隔",
-            "info": "每两次日志输出间的更新步数。"
-        }
+        "en": {"label": "Logging steps", "info": "Number of steps between two logs."},
+        "zh": {"label": "日志间隔", "info": "每两次日志输出间的更新步数。"},
    },
    "save_steps": {
-        "en": {
-            "label": "Save steps",
-            "info": "Number of steps between two checkpoints."
-        },
-        "zh": {
-            "label": "保存间隔",
-            "info": "每两次断点保存间的更新步数。"
-        }
+        "en": {"label": "Save steps", "info": "Number of steps between two checkpoints."},
+        "zh": {"label": "保存间隔", "info": "每两次断点保存间的更新步数。"},
    },
    "warmup_steps": {
+        "en": {"label": "Warmup steps", "info": "Number of steps used for warmup."},
+        "zh": {"label": "预热步数", "info": "学习率预热采用的步数。"},
+    },
+    "neftune_alpha": {
+        "en": {"label": "NEFTune Alpha", "info": "Magnitude of noise adding to embedding vectors."},
+        "zh": {"label": "NEFTune 噪声参数", "info": "嵌入向量所添加的噪声大小。"},
+    },
+    "sft_packing": {
        "en": {
-            "label": "Warmup steps",
-            "info": "Number of steps used for warmup."
+            "label": "Pack sequences",
+            "info": "Pack sequences into samples of fixed length in supervised fine-tuning.",
        },
-        "zh": {
-            "label": "预热步数",
-            "info": "学习率预热采用的步数。"
-        }
-    },
-    "neft_alpha": {
-        "en": {
-            "label": "NEFTune Alpha",
-            "info": "Magnitude of noise adding to embedding vectors."
-        },
-        "zh": {
-            "label": "NEFTune 噪声参数",
-            "info": "嵌入向量所添加的噪声大小。"
-        }
-    },
-    "train_on_prompt": {
-        "en": {
-            "label": "Train on prompt",
-            "info": "Compute loss on the prompt tokens in supervised fine-tuning."
-        },
-        "zh": {
-            "label": "计算输入损失",
-            "info": "在监督微调时候计算输入序列的损失。"
-        }
+        "zh": {"label": "序列打包", "info": "在有监督微调阶段将序列打包为相同长度的样本。"},
    },
    "upcast_layernorm": {
-        "en": {
-            "label": "Upcast LayerNorm",
-            "info": "Upcast weights of layernorm in float32."
-        },
-        "zh": {
-            "label": "缩放归一化层",
-            "info": "将归一化层权重缩放至 32 位浮点数。"
-        }
-    },
-    "lora_tab": {
-        "en": {
-            "label": "LoRA configurations"
-        },
-        "zh": {
-            "label": "LoRA 参数设置"
-        }
+        "en": {"label": "Upcast LayerNorm", "info": "Upcast weights of layernorm in float32."},
+        "zh": {"label": "缩放归一化层", "info": "将归一化层权重缩放至 32 位精度。"},
    },
+    "lora_tab": {"en": {"label": "LoRA configurations"}, "zh": {"label": "LoRA 参数设置"}},
    "lora_rank": {
-        "en": {
-            "label": "LoRA rank",
-            "info": "The rank of LoRA matrices."
-        },
-        "zh": {
-            "label": "LoRA 秩",
-            "info": "LoRA 矩阵的秩。"
-        }
+        "en": {"label": "LoRA rank", "info": "The rank of LoRA matrices."},
+        "zh": {"label": "LoRA 秩", "info": "LoRA 矩阵的秩。"},
    },
    "lora_dropout": {
-        "en": {
-            "label": "LoRA Dropout",
-            "info": "Dropout ratio of LoRA weights."
-        },
-        "zh": {
-            "label": "LoRA 随机丢弃",
-            "info": "LoRA 权重随机丢弃的概率。"
-        }
+        "en": {"label": "LoRA Dropout", "info": "Dropout ratio of LoRA weights."},
+        "zh": {"label": "LoRA 随机丢弃", "info": "LoRA 权重随机丢弃的概率。"},
    },
    "lora_target": {
        "en": {
            "label": "LoRA modules (optional)",
-            "info": "Name(s) of target modules to apply LoRA. Use commas to separate multiple modules."
+            "info": "Name(s) of target modules to apply LoRA. Use commas to separate multiple modules.",
        },
-        "zh": {
-            "label": "LoRA 作用模块（非必填）",
-            "info": "应用 LoRA 的目标模块名称。使用英文逗号分隔多个名称。"
-        }
+        "zh": {"label": "LoRA 作用模块（非必填）", "info": "应用 LoRA 的目标模块名称。使用英文逗号分隔多个名称。"},
    },
    "additional_target": {
        "en": {
            "label": "Additional modules (optional)",
-            "info": "Name(s) of modules apart from LoRA layers to be set as trainable. Use commas to separate multiple modules."
+            "info": "Name(s) of modules apart from LoRA layers to be set as trainable. Use commas to separate multiple modules.",
        },
-        "zh": {
-            "label": "附加模块（非必填）",
-            "info": "除 LoRA 层以外的可训练模块名称。使用英文逗号分隔多个名称。"
-        }
+        "zh": {"label": "附加模块（非必填）", "info": "除 LoRA 层以外的可训练模块名称。使用英文逗号分隔多个名称。"},
    },
-    "resume_lora_training": {
+    "create_new_adapter": {
        "en": {
-            "label": "Resume LoRA training",
-            "info": "Whether to resume training from the last LoRA weights or create new lora weights."
+            "label": "Create new adapter",
+            "info": "Whether to create a new adapter with randomly initialized weight or not.",
        },
-        "zh": {
-            "label": "继续上次的训练",
-            "info": "接着上次的 LoRA 权重训练或创建一个新的 LoRA 权重。"
-        }
-    },
-    "rlhf_tab": {
-        "en": {
-            "label": "RLHF configurations"
-        },
-        "zh": {
-            "label": "RLHF 参数设置"
-        }
+        "zh": {"label": "新建适配器", "info": "是否创建一个经过随机初始化的新适配器。"},
    },
+    "rlhf_tab": {"en": {"label": "RLHF configurations"}, "zh": {"label": "RLHF 参数设置"}},
    "dpo_beta": {
-        "en": {
-            "label": "DPO beta",
-            "info": "Value of the beta parameter in the DPO loss."
+        "en": {"label": "DPO beta", "info": "Value of the beta parameter in the DPO loss."},
+        "zh": {"label": "DPO beta 参数", "info": "DPO 损失函数中 beta 超参数大小。"},
    },
-        "zh": {
-            "label": "DPO beta 参数",
-            "info": "DPO 损失函数中 beta 超参数大小。"
-        }
+    "dpo_ftx": {
+        "en": {"label": "DPO-ftx weight", "info": "The weight of SFT loss in the DPO-ftx."},
+        "zh": {"label": "DPO-ftx 权重", "info": "DPO-ftx 中 SFT 损失的权重大小。"},
    },
    "reward_model": {
        "en": {
            "label": "Reward model",
-            "info": "Checkpoint of the reward model for PPO training. (Needs to refresh checkpoints)"
+            "info": "Adapter of the reward model for PPO training. (Needs to refresh adapters)",
        },
-        "zh": {
-            "label": "奖励模型",
-            "info": "PPO 训练中奖励模型的断点路径。（需要刷新断点）"
-        }
-    },
-    "cmd_preview_btn": {
-        "en": {
-            "value": "Preview command"
-        },
-        "zh": {
-            "value": "预览命令"
-        }
-    },
-    "start_btn": {
-        "en": {
-            "value": "Start"
-        },
-        "zh": {
-            "value": "开始"
-        }
-    },
-    "stop_btn": {
-        "en": {
-            "value": "Abort"
-        },
-        "zh": {
-            "value": "中断"
-        }
+        "zh": {"label": "奖励模型", "info": "PPO 训练中奖励模型的适配器路径。（需要刷新适配器）"},
    },
+    "cmd_preview_btn": {"en": {"value": "Preview command"}, "zh": {"value": "预览命令"}},
+    "start_btn": {"en": {"value": "Start"}, "zh": {"value": "开始"}},
+    "stop_btn": {"en": {"value": "Abort"}, "zh": {"value": "中断"}},
    "output_dir": {
-        "en": {
-            "label": "Checkpoint name",
-            "info": "Directory to save checkpoint."
+        "en": {"label": "Output dir", "info": "Directory for saving results."},
+        "zh": {"label": "输出目录", "info": "保存结果的路径。"},
    },
-        "zh": {
-            "label": "断点名称",
-            "info": "保存模型断点的文件夹名称。"
-        }
+    "output_box": {"en": {"value": "Ready."}, "zh": {"value": "准备就绪。"}},
+    "loss_viewer": {"en": {"label": "Loss"}, "zh": {"label": "损失"}},
+    "predict": {"en": {"label": "Save predictions"}, "zh": {"label": "保存预测结果"}},
+    "load_btn": {"en": {"value": "Load model"}, "zh": {"value": "加载模型"}},
+    "unload_btn": {"en": {"value": "Unload model"}, "zh": {"value": "卸载模型"}},
+    "info_box": {"en": {"value": "Model unloaded, please load a model first."}, "zh": {"value": "模型未加载，请先加载模型。"}},
+    "system": {"en": {"placeholder": "System prompt (optional)"}, "zh": {"placeholder": "系统提示词（非必填）"}},
+    "tools": {"en": {"placeholder": "Tools (optional)"}, "zh": {"placeholder": "工具列表（非必填）"}},
+    "query": {"en": {"placeholder": "Input..."}, "zh": {"placeholder": "输入..."}},
+    "submit_btn": {"en": {"value": "Submit"}, "zh": {"value": "提交"}},
+    "clear_btn": {"en": {"value": "Clear history"}, "zh": {"value": "清空历史"}},
+    "max_length": {"en": {"label": "Maximum length"}, "zh": {"label": "最大长度"}},
+    "max_new_tokens": {"en": {"label": "Maximum new tokens"}, "zh": {"label": "最大生成长度"}},
+    "top_p": {"en": {"label": "Top-p"}, "zh": {"label": "Top-p 采样值"}},
+    "temperature": {"en": {"label": "Temperature"}, "zh": {"label": "温度系数"}},
+    "max_shard_size": {
+        "en": {"label": "Max shard size (GB)", "info": "The maximum size for a model file."},
+        "zh": {"label": "最大分块大小（GB）", "info": "单个模型文件的最大大小。"},
    },
-    "output_box": {
-        "en": {
-            "value": "Ready."
+    "export_quantization_bit": {
+        "en": {"label": "Export quantization bit.", "info": "Quantizing the exported model."},
+        "zh": {"label": "导出量化等级", "info": "量化导出模型。"},
    },
-        "zh": {
-            "value": "准备就绪。"
-        }
-    },
-    "loss_viewer": {
-        "en": {
-            "label": "Loss"
-        },
-        "zh": {
-            "label": "损失"
-        }
-    },
-    "predict": {
-        "en": {
-            "label": "Save predictions"
-        },
-        "zh": {
-            "label": "保存预测结果"
-        }
-    },
-    "load_btn": {
-        "en": {
-            "value": "Load model"
-        },
-        "zh": {
-            "value": "加载模型"
-        }
-    },
-    "unload_btn": {
-        "en": {
-            "value": "Unload model"
-        },
-        "zh": {
-            "value": "卸载模型"
-        }
-    },
-    "info_box": {
-        "en": {
-            "value": "Model unloaded, please load a model first."
-        },
-        "zh": {
-            "value": "模型未加载，请先加载模型。"
-        }
-    },
-    "system": {
-        "en": {
-            "placeholder": "System prompt (optional)"
-        },
-        "zh": {
-            "placeholder": "系统提示词（非必填）"
-        }
-    },
-    "query": {
-        "en": {
-            "placeholder": "Input..."
-        },
-        "zh": {
-            "placeholder": "输入..."
-        }
-    },
-    "submit_btn": {
-        "en": {
-            "value": "Submit"
-        },
-        "zh": {
-            "value": "提交"
-        }
-    },
-    "clear_btn": {
-        "en": {
-            "value": "Clear history"
-        },
-        "zh": {
-            "value": "清空历史"
-        }
-    },
-    "max_length": {
-        "en": {
-            "label": "Maximum length"
-        },
-        "zh": {
-            "label": "最大长度"
-        }
-    },
-    "max_new_tokens": {
-        "en": {
-            "label": "Maximum new tokens"
-        },
-        "zh": {
-            "label": "最大生成长度"
-        }
-    },
-    "top_p": {
-        "en": {
-            "label": "Top-p"
-        },
-        "zh": {
-            "label": "Top-p 采样值"
-        }
-    },
-    "temperature": {
-        "en": {
-            "label": "Temperature"
-        },
-        "zh": {
-            "label": "温度系数"
-        }
+    "export_quantization_dataset": {
+        "en": {"label": "Export quantization dataset.", "info": "The calibration dataset used for quantization."},
+        "zh": {"label": "导出量化数据集", "info": "量化过程中使用的校准数据集。"},
    },
    "export_dir": {
-        "en": {
-            "label": "Export dir",
-            "info": "Directory to save exported model."
+        "en": {"label": "Export dir", "info": "Directory to save exported model."},
+        "zh": {"label": "导出目录", "info": "保存导出模型的文件夹路径。"},
    },
-        "zh": {
-            "label": "导出目录",
-            "info": "保存导出模型的文件夹路径。"
-        }
-    },
-    "max_shard_size": {
-        "en": {
-            "label": "Max shard size (GB)",
-            "info": "The maximum size for a model file."
-        },
-        "zh": {
-            "label": "最大分块大小（GB）",
-            "info": "模型文件的最大大小。"
-        }
-    },
-    "export_btn": {
-        "en": {
-            "value": "Export"
-        },
-        "zh": {
-            "value": "开始导出"
-        }
-    }
+    "export_btn": {"en": {"value": "Export"}, "zh": {"value": "开始导出"}},
 }


 ALERTS = {
-    "err_conflict": {
-        "en": "A process is in running, please abort it firstly.",
-        "zh": "任务已存在，请先中断训练。"
-    },
-    "err_exists": {
-        "en": "You have loaded a model, please unload it first.",
-        "zh": "模型已存在，请先卸载模型。"
-    },
-    "err_no_model": {
-        "en": "Please select a model.",
-        "zh": "请选择模型。"
-    },
-    "err_no_path": {
-        "en": "Model not found.",
-        "zh": "模型未找到。"
-    },
-    "err_no_dataset": {
-        "en": "Please choose a dataset.",
-        "zh": "请选择数据集。"
-    },
-    "err_no_checkpoint": {
-        "en": "Please select a checkpoint.",
-        "zh": "请选择断点。"
-    },
-    "err_no_export_dir": {
-        "en": "Please provide export dir.",
-        "zh": "请填写导出目录"
-    },
-    "err_failed": {
-        "en": "Failed.",
-        "zh": "训练出错。"
-    },
+    "err_conflict": {"en": "A process is in running, please abort it firstly.", "zh": "任务已存在，请先中断训练。"},
+    "err_exists": {"en": "You have loaded a model, please unload it first.", "zh": "模型已存在，请先卸载模型。"},
+    "err_no_model": {"en": "Please select a model.", "zh": "请选择模型。"},
+    "err_no_path": {"en": "Model not found.", "zh": "模型未找到。"},
+    "err_no_dataset": {"en": "Please choose a dataset.", "zh": "请选择数据集。"},
+    "err_no_adapter": {"en": "Please select an adapter.", "zh": "请选择一个适配器。"},
+    "err_no_export_dir": {"en": "Please provide export dir.", "zh": "请填写导出目录"},
+    "err_failed": {"en": "Failed.", "zh": "训练出错。"},
    "err_demo": {
        "en": "Training is unavailable in demo mode, duplicate the space to a private one first.",
-        "zh": "展示模式不支持训练，请先复制到私人空间。"
+        "zh": "展示模式不支持训练，请先复制到私人空间。",
    },
-    "info_aborting": {
-        "en": "Aborted, wait for terminating...",
-        "zh": "训练中断，正在等待线程结束……"
-    },
-    "info_aborted": {
-        "en": "Ready.",
-        "zh": "准备就绪。"
-    },
-    "info_finished": {
-        "en": "Finished.",
-        "zh": "训练完毕。"
-    },
-    "info_loading": {
-        "en": "Loading model...",
-        "zh": "加载中……"
-    },
-    "info_unloading": {
-        "en": "Unloading model...",
-        "zh": "卸载中……"
-    },
-    "info_loaded": {
-        "en": "Model loaded, now you can chat with your model!",
-        "zh": "模型已加载，可以开始聊天了！"
-    },
-    "info_unloaded": {
-        "en": "Model unloaded.",
-        "zh": "模型已卸载。"
-    },
-    "info_exporting": {
-        "en": "Exporting model...",
-        "zh": "正在导出模型……"
-    },
-    "info_exported": {
-        "en": "Model exported.",
-        "zh": "模型导出完成。"
-    }
+    "err_device_count": {"en": "Multiple GPUs are not supported yet.", "zh": "尚不支持多 GPU 训练。"},
+    "err_tool_name": {"en": "Tool name not found.", "zh": "工具名称未找到。"},
+    "err_json_schema": {"en": "Invalid JSON schema.", "zh": "Json 格式错误。"},
+    "info_aborting": {"en": "Aborted, wait for terminating...", "zh": "训练中断，正在等待线程结束……"},
+    "info_aborted": {"en": "Ready.", "zh": "准备就绪。"},
+    "info_finished": {"en": "Finished.", "zh": "训练完毕。"},
+    "info_loading": {"en": "Loading model...", "zh": "加载中……"},
+    "info_unloading": {"en": "Unloading model...", "zh": "卸载中……"},
+    "info_loaded": {"en": "Model loaded, now you can chat with your model!", "zh": "模型已加载，可以开始聊天了！"},
+    "info_unloaded": {"en": "Model unloaded.", "zh": "模型已卸载。"},
+    "info_exporting": {"en": "Exporting model...", "zh": "正在导出模型……"},
+    "info_exported": {"en": "Model exported.", "zh": "模型导出完成。"},
 }
--- a/src/llmtuner/webui/manager.py
+++ b/src/llmtuner/webui/manager.py
@@ -1,11 +1,11 @@
 from typing import TYPE_CHECKING, Dict, List, Set

+
 if TYPE_CHECKING:
    from gradio.components import Component


 class Manager:
-
    def __init__(self) -> None:
        self.all_elems: Dict[str, Dict[str, "Component"]] = {}

@@ -21,14 +21,12 @@ class Manager:
            self.all_elems["top"]["lang"],
            self.all_elems["top"]["model_name"],
            self.all_elems["top"]["model_path"],
-            self.all_elems["top"]["checkpoints"],
+            self.all_elems["top"]["adapter_path"],
            self.all_elems["top"]["finetuning_type"],
            self.all_elems["top"]["quantization_bit"],
            self.all_elems["top"]["template"],
-            self.all_elems["top"]["system_prompt"],
-            self.all_elems["top"]["flash_attn"],
-            self.all_elems["top"]["shift_attn"],
-            self.all_elems["top"]["rope_scaling"]
+            self.all_elems["top"]["rope_scaling"],
+            self.all_elems["top"]["booster"],
        }

    def list_elems(self) -> List["Component"]:
--- a/src/llmtuner/webui/runner.py
+++ b/src/llmtuner/webui/runner.py
@@ -1,29 +1,29 @@
+import logging
 import os
 import time
-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, Optional, Tuple

+import gradio as gr
 import transformers
+from gradio.components import Component  # cannot use TYPE_CHECKING here
 from transformers.trainer import TRAINING_ARGS_NAME

-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.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
+from ..extras.callbacks import LogCallback
+from ..extras.constants import TRAINING_STAGES
+from ..extras.logging import LoggerHandler
+from ..extras.misc import get_device_count, torch_gc
+from ..train import run_exp
+from .common import get_module, get_save_dir, load_config
+from .locales import ALERTS
+from .utils import gen_cmd, get_eval_results, update_process_bar
+

 if TYPE_CHECKING:
-    from llmtuner.webui.manager import Manager
+    from .manager import Manager


 class Runner:
-
    def __init__(self, manager: "Manager", demo_mode: Optional[bool] = False) -> None:
        self.manager = manager
        self.demo_mode = demo_mode
@@ -67,6 +67,9 @@ class Runner:
        if self.demo_mode and (not from_preview):
            return ALERTS["err_demo"][lang]

+        if not from_preview and get_device_count() > 1:
+            return ALERTS["err_device_count"][lang]
+
        self.aborted = False
        self.logger_handler.reset()
        self.trainer_callback = LogCallback(self)
@@ -86,26 +89,28 @@ class Runner:
        get = lambda name: data[self.manager.get_elem_by_name(name)]
        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")])
+        if get("top.adapter_path"):
+            adapter_name_or_path = ",".join(
+                [
+                    get_save_dir(get("top.model_name"), get("top.finetuning_type"), adapter)
+                    for adapter in get("top.adapter_path")
+                ]
+            )
        else:
-            checkpoint_dir = None
+            adapter_name_or_path = None

        args = dict(
            stage=TRAINING_STAGES[get("train.training_stage")],
-            model_name_or_path=get("top.model_path"),
            do_train=True,
+            model_name_or_path=get("top.model_path"),
+            adapter_name_or_path=adapter_name_or_path,
            cache_dir=user_config.get("cache_dir", None),
-            checkpoint_dir=checkpoint_dir,
            finetuning_type=get("top.finetuning_type"),
            quantization_bit=int(get("top.quantization_bit")) if get("top.quantization_bit") in ["8", "4"] else None,
            template=get("top.template"),
-            system_prompt=get("top.system_prompt"),
-            flash_attn=get("top.flash_attn"),
-            shift_attn=get("top.shift_attn"),
            rope_scaling=get("top.rope_scaling") if get("top.rope_scaling") in ["linear", "dynamic"] else None,
+            flash_attn=(get("top.booster") == "flash_attn"),
+            use_unsloth=(get("top.booster") == "unsloth"),
            dataset_dir=get("train.dataset_dir"),
            dataset=",".join(get("train.dataset")),
            cutoff_len=get("train.cutoff_len"),
@@ -119,24 +124,22 @@ class Runner:
            logging_steps=get("train.logging_steps"),
            save_steps=get("train.save_steps"),
            warmup_steps=get("train.warmup_steps"),
-            neft_alpha=get("train.neft_alpha"),
-            train_on_prompt=get("train.train_on_prompt"),
+            neftune_noise_alpha=get("train.neftune_alpha") or None,
+            sft_packing=get("train.sft_packing"),
            upcast_layernorm=get("train.upcast_layernorm"),
            lora_rank=get("train.lora_rank"),
            lora_dropout=get("train.lora_dropout"),
            lora_target=get("train.lora_target") or get_module(get("top.model_name")),
-            additional_target=get("train.additional_target") if get("train.additional_target") else None,
-            resume_lora_training=get("train.resume_lora_training"),
-            output_dir=get_save_dir(get("top.model_name"), get("top.finetuning_type"), get("train.output_dir"))
+            additional_target=get("train.additional_target") or None,
+            create_new_adapter=get("train.create_new_adapter"),
+            output_dir=get_save_dir(get("top.model_name"), get("top.finetuning_type"), get("train.output_dir")),
+            fp16=(get("train.compute_type") == "fp16"),
+            bf16=(get("train.compute_type") == "bf16"),
        )
-        args[get("train.compute_type")] = True
        args["disable_tqdm"] = True

        if TRAINING_STAGES[get("train.training_stage")] in ["rm", "ppo", "dpo"]:
-            args["resume_lora_training"] = (args["quantization_bit"] is not None)
-
-        if args["quantization_bit"] is not None:
-            args["upcast_layernorm"] = True
+            args["create_new_adapter"] = args["quantization_bit"] is None

        if args["stage"] == "ppo":
            args["reward_model"] = get_save_dir(
@@ -146,6 +149,7 @@ class Runner:

        if args["stage"] == "dpo":
            args["dpo_beta"] = get("train.dpo_beta")
+            args["dpo_ftx"] = get("train.dpo_ftx")

        if get("train.val_size") > 1e-6 and args["stage"] != "ppo":
            args["val_size"] = get("train.val_size")
@@ -159,49 +163,49 @@ class Runner:
        get = lambda name: data[self.manager.get_elem_by_name(name)]
        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")])
-            output_dir = get_save_dir(
-                get("top.model_name"), get("top.finetuning_type"), "eval_" + "_".join(get("top.checkpoints"))
+        if get("top.adapter_path"):
+            adapter_name_or_path = ",".join(
+                [
+                    get_save_dir(get("top.model_name"), get("top.finetuning_type"), adapter)
+                    for adapter in get("top.adapter_path")
+                ]
            )
        else:
-            checkpoint_dir = None
-            output_dir = get_save_dir(get("top.model_name"), get("top.finetuning_type"), "eval_base")
+            adapter_name_or_path = None

        args = dict(
            stage="sft",
            model_name_or_path=get("top.model_path"),
-            do_eval=True,
-            predict_with_generate=True,
+            adapter_name_or_path=adapter_name_or_path,
            cache_dir=user_config.get("cache_dir", None),
-            checkpoint_dir=checkpoint_dir,
            finetuning_type=get("top.finetuning_type"),
            quantization_bit=int(get("top.quantization_bit")) if get("top.quantization_bit") in ["8", "4"] else None,
            template=get("top.template"),
-            system_prompt=get("top.system_prompt"),
-            flash_attn=get("top.flash_attn"),
-            shift_attn=get("top.shift_attn"),
            rope_scaling=get("top.rope_scaling") if get("top.rope_scaling") in ["linear", "dynamic"] else None,
+            flash_attn=(get("top.booster") == "flash_attn"),
+            use_unsloth=(get("top.booster") == "unsloth"),
            dataset_dir=get("eval.dataset_dir"),
            dataset=",".join(get("eval.dataset")),
            cutoff_len=get("eval.cutoff_len"),
            max_samples=int(get("eval.max_samples")),
            per_device_eval_batch_size=get("eval.batch_size"),
+            predict_with_generate=True,
            max_new_tokens=get("eval.max_new_tokens"),
            top_p=get("eval.top_p"),
            temperature=get("eval.temperature"),
-            output_dir=output_dir
+            output_dir=get_save_dir(get("top.model_name"), get("top.finetuning_type"), get("eval.output_dir")),
        )

        if get("eval.predict"):
-            args.pop("do_eval", None)
            args["do_predict"] = True
+        else:
+            args["do_eval"] = True

        return args

-    def _preview(self, data: Dict[Component, Any], do_train: bool) -> Generator[Tuple[str, Dict[str, Any]], None, None]:
+    def _preview(
+        self, data: Dict[Component, Any], do_train: bool
+    ) -> Generator[Tuple[str, Dict[str, Any]], None, None]:
        error = self._initialize(data, do_train, from_preview=True)
        if error:
            gr.Warning(error)
@@ -239,9 +243,12 @@ class Runner:
        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")
-        ))
+        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:
--- a/src/llmtuner/webui/utils.py
+++ b/src/llmtuner/webui/utils.py
@@ -1,15 +1,18 @@
-import os
 import json
-import gradio as gr
-from typing import TYPE_CHECKING, Any, Dict
+import os
 from datetime import datetime
+from typing import TYPE_CHECKING, Any, Dict
+
+import gradio as gr
+
+from ..extras.packages import is_matplotlib_available
+from ..extras.ploting import smooth
+from .common import get_save_dir
+from .locales import ALERTS

-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
+    from ..extras.callbacks import LogCallback

 if is_matplotlib_available():
    import matplotlib.figure
@@ -22,16 +25,13 @@ def update_process_bar(callback: "LogCallback") -> Dict[str, Any]:

    percentage = round(100 * callback.cur_steps / callback.max_steps, 0) if callback.max_steps != 0 else 100.0
    label = "Running {:d}/{:d}: {} < {}".format(
-        callback.cur_steps,
-        callback.max_steps,
-        callback.elapsed_time,
-        callback.remaining_time
+        callback.cur_steps, callback.max_steps, callback.elapsed_time, callback.remaining_time
    )
    return gr.update(label=label, value=percentage, visible=True)


 def get_time() -> str:
-    return datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
+    return datetime.now().strftime("%Y-%m-%d-%H-%M-%S")


 def can_quantize(finetuning_type: str) -> Dict[str, Any]:
@@ -41,12 +41,24 @@ def can_quantize(finetuning_type: str) -> Dict[str, Any]:
        return gr.update(interactive=True)


+def check_json_schema(text: str, lang: str) -> None:
+    try:
+        tools = json.loads(text)
+        for tool in tools:
+            assert "name" in tool
+    except AssertionError:
+        gr.Warning(ALERTS["err_tool_name"][lang])
+    except json.JSONDecodeError:
+        gr.Warning(ALERTS["err_json_schema"][lang])
+
+
 def gen_cmd(args: Dict[str, Any]) -> str:
    args.pop("disable_tqdm", None)
    args["plot_loss"] = args.get("do_train", None)
-    cmd_lines = ["CUDA_VISIBLE_DEVICES=0 python src/train_bash.py "]
+    current_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "0")
+    cmd_lines = ["CUDA_VISIBLE_DEVICES={} python src/train_bash.py ".format(current_devices)]
    for k, v in args.items():
-        if v is not None and v != "":
+        if v is not None and v is not False and v != "":
            cmd_lines.append("    --{} {} ".format(k, str(v)))
    cmd_text = "\\\n".join(cmd_lines)
    cmd_text = "```bash\n{}\n```".format(cmd_text)
--- a/src/train_web.py
+++ b/src/train_web.py
@@ -4,7 +4,7 @@ from llmtuner import create_ui
 def main():
    demo = create_ui()
    demo.queue()
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=False, inbrowser=True)
+    demo.launch(server_name="0.0.0.0", share=False, inbrowser=True)


 if __name__ == "__main__":
--- a/src/web_demo.py
+++ b/src/web_demo.py
@@ -4,7 +4,7 @@ from llmtuner import create_web_demo
 def main():
    demo = create_web_demo()
    demo.queue()
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=False, inbrowser=True)
+    demo.launch(server_name="0.0.0.0", share=False, inbrowser=True)


 if __name__ == "__main__":
--- a/tests/cal_flops.py
+++ b/tests/cal_flops.py
@@ -3,9 +3,10 @@
 # Usage: python cal_flops.py --model_name_or_path path_to_model --batch_size 1 --seq_length 512
 # Inspired by: https://www.deepspeed.ai/tutorials/flops-profiler/

+from typing import Optional
+
 import fire
 import torch
-from typing import Optional
 from deepspeed.accelerator import get_accelerator  # type: ignore
 from deepspeed.profiling.flops_profiler import get_model_profile  # type: ignore

@@ -16,25 +17,13 @@ def calculate_flops(
    model_name_or_path: str,
    batch_size: Optional[int] = 1,
    seq_length: Optional[int] = 256,
-    flash_attn: Optional[bool] = False
+    flash_attn: Optional[bool] = False,
 ):
    with get_accelerator().device(0):
-        chat_model = ChatModel(dict(
-            model_name_or_path=model_name_or_path,
-            template="vanilla",
-            flash_attn=flash_attn
-        ))
+        chat_model = ChatModel(dict(model_name_or_path=model_name_or_path, template="vanilla", flash_attn=flash_attn))
        fake_input = torch.ones((batch_size, seq_length), dtype=torch.long, device=chat_model.model.device)
-        input_dict = {
-            "input_ids": fake_input,
-            "labels": fake_input.clone()
-        }
-        flops, macs, params = get_model_profile(
-            chat_model.model,
-            kwargs=input_dict,
-            print_profile=True,
-            detailed=True
-        )
+        input_dict = {"input_ids": fake_input, "labels": fake_input.clone()}
+        flops, macs, params = get_model_profile(chat_model.model, kwargs=input_dict, print_profile=True, detailed=True)
        print("FLOPs:", flops)
        print("MACs:", macs)
        print("Params:", params)
--- a/tests/cal_lr.py
+++ b/tests/cal_lr.py
@@ -3,17 +3,19 @@
 # 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
+
+import fire
+import torch
 from torch.utils.data import DataLoader
+from tqdm import tqdm
 from transformers import DataCollatorForSeq2Seq

-from llmtuner.data import get_dataset, preprocess_dataset
+from llmtuner.data import get_dataset
 from llmtuner.extras.constants import IGNORE_INDEX
-from llmtuner.model import get_train_args, load_model_and_tokenizer
+from llmtuner.hparams import get_train_args
+from llmtuner.model import load_model_and_tokenizer


 BASE_LR = 3e-4  # 1.5e-4 for 30B-70B models
@@ -26,20 +28,21 @@ def calculate_lr(
    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"
+    dataset_dir: Optional[str] = "data",
 ):
-    model_args, data_args, training_args, finetuning_args, _ = get_train_args(dict(
+    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")
+            output_dir="dummy_dir",
+        )
+    )
+    _, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, is_trainable=False, add_valuehead=False)
+    trainset = get_dataset(tokenizer, model_args, 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
@@ -54,9 +57,11 @@ def calculate_lr(
    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(
+    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__":
--- a/tests/llamafy_baichuan2.py
+++ b/tests/llamafy_baichuan2.py
@@ -4,31 +4,36 @@
 # Inspired by: https://huggingface.co/fireballoon/baichuan-llama-7b/blob/main/convert_baichuan_to_llama.py
 # Converted model: https://huggingface.co/hiyouga/Baichuan2-7B-Base-LLaMAfied

-import os
-import fire
 import json
-import torch
+import os
 from collections import OrderedDict
-from transformers.modeling_utils import shard_checkpoint, WEIGHTS_NAME, WEIGHTS_INDEX_NAME
-from typing import Any, Dict
+from typing import Any, Dict, Optional
+
+import fire
+import torch
+from safetensors.torch import save_file
+from tqdm import tqdm
+from transformers.modeling_utils import (
+    SAFE_WEIGHTS_INDEX_NAME,
+    SAFE_WEIGHTS_NAME,
+    WEIGHTS_INDEX_NAME,
+    WEIGHTS_NAME,
+    shard_checkpoint,
+)


 CONFIG_NAME = "config.json"


-def save_weight(
-    input_dir: str,
-    output_dir: str,
-    shard_size: str
-):
+def save_weight(input_dir: str, output_dir: str, shard_size: str, save_safetensors: bool):
    baichuan2_state_dict: Dict[str, torch.Tensor] = OrderedDict()
-    for filepath in os.listdir(input_dir):
+    for filepath in tqdm(os.listdir(input_dir), desc="Load weights"):
        if os.path.isfile(os.path.join(input_dir, filepath)) and filepath.endswith(".bin"):
            shard_weight = torch.load(os.path.join(input_dir, filepath), map_location="cpu")
            baichuan2_state_dict.update(shard_weight)

    llama2_state_dict: Dict[str, torch.Tensor] = OrderedDict()
-    for key, value in baichuan2_state_dict.items():
+    for key, value in tqdm(baichuan2_state_dict.items(), desc="Convert format"):
        if "W_pack" in key:
            proj_size = value.size(0) // 3
            llama2_state_dict[key.replace("W_pack", "q_proj")] = value[:proj_size, :]
@@ -39,22 +44,25 @@ def save_weight(
        else:
            llama2_state_dict[key] = value

-    shards, index = shard_checkpoint(llama2_state_dict, max_shard_size=shard_size, weights_name=WEIGHTS_NAME)
-    for shard_file, shard in shards.items():
+    weights_name = SAFE_WEIGHTS_NAME if save_safetensors else WEIGHTS_NAME
+    shards, index = shard_checkpoint(llama2_state_dict, max_shard_size=shard_size, weights_name=weights_name)
+
+    for shard_file, shard in tqdm(shards.items(), desc="Save weights"):
+        if save_safetensors:
+            save_file(shard, os.path.join(output_dir, shard_file), metadata={"format": "pt"})
+        else:
            torch.save(shard, os.path.join(output_dir, shard_file))

    if index is None:
        print("Model weights saved in {}".format(os.path.join(output_dir, WEIGHTS_NAME)))
    else:
-        with open(os.path.join(output_dir, WEIGHTS_INDEX_NAME), "w", encoding="utf-8") as f:
+        index_name = SAFE_WEIGHTS_INDEX_NAME if save_safetensors else WEIGHTS_INDEX_NAME
+        with open(os.path.join(output_dir, index_name), "w", encoding="utf-8") as f:
            json.dump(index, f, indent=2, sort_keys=True)
        print("Model weights saved in {}".format(output_dir))


-def save_config(
-    input_dir: str,
-    output_dir: str
-):
+def save_config(input_dir: str, output_dir: str):
    with open(os.path.join(input_dir, CONFIG_NAME), "r", encoding="utf-8") as f:
        llama2_config_dict: Dict[str, Any] = json.load(f)

@@ -68,17 +76,13 @@ def save_config(
    print("Model config saved in {}".format(os.path.join(output_dir, CONFIG_NAME)))


-def llamafy_baichuan2(
-    input_dir: str,
-    output_dir: str,
-    shard_size: str
-):
+def llamafy_baichuan2(input_dir: str, output_dir: str, shard_size: str, save_safetensors: Optional[bool] = False):
    try:
        os.makedirs(output_dir, exist_ok=False)
    except Exception as e:
        raise print("Output dir already exists", e)

-    save_weight(input_dir, output_dir, shard_size)
+    save_weight(input_dir, output_dir, shard_size, save_safetensors)
    save_config(input_dir, output_dir)


--- a/tests/llamafy_internlm2.py
+++ b/tests/llamafy_internlm2.py
@@ -0,0 +1,112 @@
+# coding=utf-8
+# Converts the InternLM2 model in the same format as LLaMA2.
+# Usage: python llamafy_internlm2.py --input_dir input --output_dir output --shard_size 10GB
+# Warning: We have found that the converted model cannot infer correctly. It will be fixed later.
+
+import json
+import os
+from collections import OrderedDict
+from typing import Any, Dict, Optional
+
+import fire
+import torch
+from safetensors.torch import save_file
+from tqdm import tqdm
+from transformers.modeling_utils import (
+    SAFE_WEIGHTS_INDEX_NAME,
+    SAFE_WEIGHTS_NAME,
+    WEIGHTS_INDEX_NAME,
+    WEIGHTS_NAME,
+    shard_checkpoint,
+)
+
+
+CONFIG_NAME = "config.json"
+
+
+def save_weight(input_dir: str, output_dir: str, shard_size: str, save_safetensors: bool):
+    with open(os.path.join(input_dir, CONFIG_NAME), "r", encoding="utf-8") as f:
+        internlm2_config_dict: Dict[str, Any] = json.load(f)
+
+    internlm2_state_dict: Dict[str, torch.Tensor] = OrderedDict()
+    for filepath in tqdm(os.listdir(input_dir), desc="Load weights"):
+        if os.path.isfile(os.path.join(input_dir, filepath)) and filepath.endswith(".bin"):
+            shard_weight = torch.load(os.path.join(input_dir, filepath), map_location="cpu")
+            internlm2_state_dict.update(shard_weight)
+
+    llama2_state_dict: Dict[str, torch.Tensor] = OrderedDict()
+    for key, value in tqdm(internlm2_state_dict.items(), desc="Convert format"):
+        if "output" in key:
+            llama2_state_dict[key.replace("output", "lm_head")] = value
+        elif "tok_embeddings" in key:
+            llama2_state_dict[key.replace("tok_embeddings", "embed_tokens")] = value
+        elif "wqkv" in key:
+            num_q_heads = internlm2_config_dict["num_attention_heads"]
+            num_kv_heads = internlm2_config_dict["num_key_value_heads"]
+            q_size = value.size(0) // (num_q_heads + 2 * num_kv_heads) * num_q_heads
+            kv_size = value.size(0) // (num_q_heads + 2 * num_kv_heads) * num_kv_heads
+            llama2_state_dict[key.replace("attention.wqkv", "self_attn.q_proj")] = value[:q_size, ...]
+            llama2_state_dict[key.replace("attention.wqkv", "self_attn.k_proj")] = value[
+                q_size : q_size + kv_size, ...
+            ]
+            llama2_state_dict[key.replace("attention.wqkv", "self_attn.v_proj")] = value[q_size + kv_size :, ...]
+        elif "wo" in key:
+            llama2_state_dict[key.replace("attention.wo", "self_attn.o_proj")] = value
+        elif "attention_norm" in key:
+            llama2_state_dict[key.replace("attention_norm", "input_layernorm")] = value
+        elif "ffn_norm" in key:
+            llama2_state_dict[key.replace("ffn_norm", "post_attention_layernorm")] = value
+        elif "w1" in key:
+            llama2_state_dict[key.replace("feed_forward.w1", "mlp.gate_proj")] = value
+        elif "w2" in key:
+            llama2_state_dict[key.replace("feed_forward.w2", "mlp.down_proj")] = value
+        elif "w3" in key:
+            llama2_state_dict[key.replace("feed_forward.w3", "mlp.up_proj")] = value
+        else:
+            llama2_state_dict[key] = value
+
+    weights_name = SAFE_WEIGHTS_NAME if save_safetensors else WEIGHTS_NAME
+    shards, index = shard_checkpoint(llama2_state_dict, max_shard_size=shard_size, weights_name=weights_name)
+
+    for shard_file, shard in tqdm(shards.items(), desc="Save weights"):
+        if save_safetensors:
+            save_file(shard, os.path.join(output_dir, shard_file), metadata={"format": "pt"})
+        else:
+            torch.save(shard, os.path.join(output_dir, shard_file))
+
+    if index is None:
+        print("Model weights saved in {}".format(os.path.join(output_dir, WEIGHTS_NAME)))
+    else:
+        index_name = SAFE_WEIGHTS_INDEX_NAME if save_safetensors else WEIGHTS_INDEX_NAME
+        with open(os.path.join(output_dir, index_name), "w", encoding="utf-8") as f:
+            json.dump(index, f, indent=2, sort_keys=True)
+        print("Model weights saved in {}".format(output_dir))
+
+
+def save_config(input_dir: str, output_dir: str):
+    with open(os.path.join(input_dir, CONFIG_NAME), "r", encoding="utf-8") as f:
+        llama2_config_dict: Dict[str, Any] = json.load(f)
+
+    llama2_config_dict["architectures"] = ["LlamaForCausalLM"]
+    llama2_config_dict.pop("auto_map", None)
+    llama2_config_dict.pop("bias", None)
+    llama2_config_dict.pop("rope_scaling", None)
+    llama2_config_dict["model_type"] = "llama"
+
+    with open(os.path.join(output_dir, CONFIG_NAME), "w", encoding="utf-8") as f:
+        json.dump(llama2_config_dict, f, indent=2)
+    print("Model config saved in {}".format(os.path.join(output_dir, CONFIG_NAME)))
+
+
+def llamafy_internlm2(input_dir: str, output_dir: str, shard_size: str, save_safetensors: Optional[bool] = False):
+    try:
+        os.makedirs(output_dir, exist_ok=False)
+    except Exception as e:
+        raise print("Output dir already exists", e)
+
+    save_weight(input_dir, output_dir, shard_size, save_safetensors)
+    save_config(input_dir, output_dir)
+
+
+if __name__ == "__main__":
+    fire.Fire(llamafy_internlm2)
--- a/tests/llamafy_qwen.py
+++ b/tests/llamafy_qwen.py
@@ -1,33 +1,40 @@
 # coding=utf-8
 # Converts the Qwen models in the same format as LLaMA2.
 # Usage: python llamafy_qwen.py --input_dir input --output_dir output --shard_size 10GB
+# Converted model: https://huggingface.co/hiyouga/Qwen-14B-Chat-LLaMAfied

-import os
-import fire
 import json
-import torch
+import os
 from collections import OrderedDict
+from typing import Any, Dict, Optional
+
+import fire
+import torch
 from safetensors import safe_open
-from transformers.modeling_utils import shard_checkpoint, WEIGHTS_NAME, WEIGHTS_INDEX_NAME
+from safetensors.torch import save_file
+from tqdm import tqdm
+from transformers.modeling_utils import (
+    SAFE_WEIGHTS_INDEX_NAME,
+    SAFE_WEIGHTS_NAME,
+    WEIGHTS_INDEX_NAME,
+    WEIGHTS_NAME,
+    shard_checkpoint,
+)
 from transformers.utils import check_min_version
-from typing import Any, Dict
+

 try:
    check_min_version("4.34.0")
-except:
+except Exception:
    raise ValueError("Please upgrade `transformers` to 4.34.0")


 CONFIG_NAME = "config.json"


-def save_weight(
-    input_dir: str,
-    output_dir: str,
-    shard_size: str
-) -> str:
+def save_weight(input_dir: str, output_dir: str, shard_size: str, save_safetensors: bool) -> str:
    qwen_state_dict: Dict[str, torch.Tensor] = OrderedDict()
-    for filepath in os.listdir(input_dir):
+    for filepath in tqdm(os.listdir(input_dir), desc="Load weights"):
        if os.path.isfile(os.path.join(input_dir, filepath)) and filepath.endswith(".safetensors"):
            with safe_open(os.path.join(input_dir, filepath), framework="pt", device="cpu") as f:
                for key in f.keys():
@@ -35,7 +42,7 @@ def save_weight(

    llama2_state_dict: Dict[str, torch.Tensor] = OrderedDict()
    torch_dtype = None
-    for key, value in qwen_state_dict.items():
+    for key, value in tqdm(qwen_state_dict.items(), desc="Convert format"):
        if torch_dtype is None:
            torch_dtype = value.dtype
        if "wte" in key:
@@ -47,13 +54,15 @@ def save_weight(
            if "attn.c_attn" in key:
                proj_size = value.size(0) // 3
                llama2_state_dict[key.replace("attn.c_attn", "self_attn.q_proj")] = value[:proj_size, ...]
-                llama2_state_dict[key.replace("attn.c_attn", "self_attn.k_proj")] = value[proj_size:2*proj_size, ...]
+                llama2_state_dict[key.replace("attn.c_attn", "self_attn.k_proj")] = value[
+                    proj_size : 2 * proj_size, ...
+                ]
                llama2_state_dict[key.replace("attn.c_attn", "self_attn.v_proj")] = value[2 * proj_size :, ...]
            elif "attn.c_proj" in key:
                llama2_state_dict[key.replace("attn.c_proj", "self_attn.o_proj")] = value
-                llama2_state_dict[key.replace("attn.c_proj.weight", "self_attn.o_proj.bias")] = (
-                    torch.zeros_like(value[:, 0]).squeeze()
-                )
+                llama2_state_dict[key.replace("attn.c_proj.weight", "self_attn.o_proj.bias")] = torch.zeros_like(
+                    value[:, 0]
+                ).squeeze()
            elif "ln_1" in key:
                llama2_state_dict[key.replace("ln_1", "input_layernorm")] = value
            elif "ln_2" in key:
@@ -69,25 +78,27 @@ def save_weight(
            else:
                raise KeyError("Unable to process key {}".format(key))

-    shards, index = shard_checkpoint(llama2_state_dict, max_shard_size=shard_size, weights_name=WEIGHTS_NAME)
-    for shard_file, shard in shards.items():
+    weights_name = SAFE_WEIGHTS_NAME if save_safetensors else WEIGHTS_NAME
+    shards, index = shard_checkpoint(llama2_state_dict, max_shard_size=shard_size, weights_name=weights_name)
+
+    for shard_file, shard in tqdm(shards.items(), desc="Save weights"):
+        if save_safetensors:
+            save_file(shard, os.path.join(output_dir, shard_file), metadata={"format": "pt"})
+        else:
            torch.save(shard, os.path.join(output_dir, shard_file))

    if index is None:
-        print("Model weights saved in {}".format(os.path.join(output_dir, WEIGHTS_NAME)))
+        print("Model weights saved in {}".format(os.path.join(output_dir, weights_name)))
    else:
-        with open(os.path.join(output_dir, WEIGHTS_INDEX_NAME), "w", encoding="utf-8") as f:
+        index_name = SAFE_WEIGHTS_INDEX_NAME if save_safetensors else WEIGHTS_INDEX_NAME
+        with open(os.path.join(output_dir, index_name), "w", encoding="utf-8") as f:
            json.dump(index, f, indent=2, sort_keys=True)
        print("Model weights saved in {}".format(output_dir))

    return str(torch_dtype).replace("torch.", "")


-def save_config(
-    input_dir: str,
-    output_dir: str,
-    torch_dtype: str
-):
+def save_config(input_dir: str, output_dir: str, torch_dtype: str):
    with open(os.path.join(input_dir, CONFIG_NAME), "r", encoding="utf-8") as f:
        qwen_config_dict: Dict[str, Any] = json.load(f)

@@ -117,17 +128,13 @@ def save_config(
    print("Model config saved in {}".format(os.path.join(output_dir, CONFIG_NAME)))


-def llamafy_qwen(
-    input_dir: str,
-    output_dir: str,
-    shard_size: str
-):
+def llamafy_qwen(input_dir: str, output_dir: str, shard_size: str, save_safetensors: Optional[bool] = False):
    try:
        os.makedirs(output_dir, exist_ok=False)
    except Exception as e:
        raise print("Output dir already exists", e)

-    torch_dtype = save_weight(input_dir, output_dir, shard_size)
+    torch_dtype = save_weight(input_dir, output_dir, shard_size, save_safetensors)
    save_config(input_dir, output_dir, torch_dtype)


--- a/tests/loftq_init.py
+++ b/tests/loftq_init.py
@@ -0,0 +1,82 @@
+# coding=utf-8
+# Initializes LoRA weights with LoRA-fine-tuning-aware Quantization (LoftQ)
+# Usage: python loftq_init.py --model_name_or_path path_to_model --save_dir output_dir
+# Inspired by: https://github.com/huggingface/peft/blob/main/examples/loftq_finetuning/quantize_save_load.py
+
+import os
+from typing import TYPE_CHECKING, Optional
+
+import fire
+import torch
+import torch.nn as nn
+from peft import LoftQConfig, LoraConfig, TaskType, get_peft_model
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel
+
+
+class Shell(nn.Module):
+    def __init__(self, weight: torch.Tensor, bias: Optional[torch.Tensor] = None):
+        super().__init__()
+        self.weight = nn.Parameter(weight, requires_grad=False)
+        if bias is not None:
+            self.bias = nn.Parameter(bias, requires_grad=False)
+
+
+def unwrap_model(model: nn.Module, pattern=".base_layer") -> None:
+    for name in set([k.split(pattern)[0] for k, _ in model.named_modules() if pattern in k]):  # noqa: C403
+        parent_name = ".".join(name.split(".")[:-1])
+        child_name = name.split(".")[-1]
+        parent_module = model.get_submodule(parent_name)
+        child_module = getattr(parent_module, child_name)
+        base_layer = getattr(child_module, "base_layer")
+        weight = getattr(base_layer, "weight", None)
+        bias = getattr(base_layer, "bias", None)
+        setattr(parent_module, child_name, Shell(weight, bias))
+
+    print("Model unwrapped.")
+
+
+def quantize_loftq(
+    model_name_or_path: str,
+    save_dir: str,
+    loftq_bits: Optional[int] = 4,
+    loftq_iter: Optional[int] = 1,
+    lora_alpha: Optional[int] = None,
+    lora_rank: Optional[int] = 16,
+    lora_target: Optional[str] = "q_proj,v_proj",
+    save_safetensors: Optional[bool] = False,
+):
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(model_name_or_path, trust_remote_code=True, torch_dtype="auto")
+    loftq_config = LoftQConfig(loftq_bits=loftq_bits, loftq_iter=loftq_iter)
+    lora_config = LoraConfig(
+        task_type=TaskType.CAUSAL_LM,
+        inference_mode=True,
+        r=lora_rank,
+        lora_alpha=lora_alpha if lora_alpha is not None else lora_rank * 2,
+        lora_dropout=0.1,
+        target_modules=[name.strip() for name in lora_target.split(",")],
+        init_lora_weights="loftq",
+        loftq_config=loftq_config,
+    )
+
+    # Init LoftQ model
+    lora_model = get_peft_model(model, lora_config)
+    base_model: "PreTrainedModel" = lora_model.get_base_model()
+
+    # Save LoftQ model
+    setattr(lora_model.base_model.peft_config["default"], "base_model_name_or_path", save_dir)
+    setattr(lora_model.base_model.peft_config["default"], "init_lora_weights", True)
+    lora_model.save_pretrained(os.path.join(save_dir, "adapters"), safe_serialization=save_safetensors)
+
+    # Save base model
+    unwrap_model(base_model)
+    base_model.save_pretrained(save_dir, safe_serialization=save_safetensors)
+    tokenizer.save_pretrained(save_dir)
+
+
+if __name__ == "__main__":
+    fire.Fire(quantize_loftq)
--- a/tests/quantize.py
+++ b/tests/quantize.py
@@ -1,49 +0,0 @@
-# coding=utf-8
-# Quantizes models with AutoGPTQ (https://github.com/PanQiWei/AutoGPTQ).
-# Usage: python quantize.py --input_dir path_to_llama_model --output_dir path_to_quant_model --data_file alpaca.json
-#                           --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
-from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
-
-
-def quantize(input_dir: str, output_dir: str, data_file: str, max_length: int, max_samples: int):
-    tokenizer = AutoTokenizer.from_pretrained(input_dir, use_fast=False, padding_side="left")
-
-    def format_example(examples):
-        prefix=("A chat between a curious user and an artificial intelligence assistant. "
-                "The assistant gives helpful, detailed, and polite answers to the user's questions.")
-        texts = []
-        for i in range(len(examples["instruction"])):
-            prompt = prefix + "\n"
-            if "history" in examples:
-                for user_query, bot_resp in examples["history"][i]:
-                    prompt += "Human: {}\nAssistant: {}\n".format(user_query, bot_resp)
-            prompt += "Human: {}\nAssistant: {}".format(
-                examples["instruction"][i] + "\n" + examples["input"][i], examples["output"][i]
-            )
-            texts.append(prompt)
-        return tokenizer(texts, truncation=True, max_length=max_length)
-
-    dataset = load_dataset("json", data_files=data_file)["train"]
-    column_names = list(dataset.column_names)
-    dataset = dataset.select(range(min(len(dataset), max_samples)))
-    dataset = dataset.map(format_example, batched=True, remove_columns=column_names)
-    dataset = dataset.shuffle()
-
-    quantize_config = BaseQuantizeConfig(
-        bits=4,
-        group_size=128,
-        desc_act=False
-    )
-
-    model = AutoGPTQForCausalLM.from_pretrained(input_dir, quantize_config, trust_remote_code=True)
-    model.quantize(dataset)
-    model.save_quantized(output_dir)
-
-
-if __name__ == "__main__":
-    fire.Fire(quantize)