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
95 changed files with 4240 additions and 3797 deletions
--- 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
@@ -6,7 +6,7 @@
 [![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)

@@ -55,17 +55,23 @@ Compared to ChatGLM's [P-Tuning](https://github.com/THUDM/ChatGLM2-6B/tree/main/

 ## Changelog

-[23/12/01] We supported downloading pre-trained models from the **[ModelScope Hub](https://modelscope.cn/models)** for Chinese mainland users. See [this tutorial](#use-modelscope-models-optional) for usage.
+[24/01/18] We supported **agent tuning** for most models, equipping model with tool using abilities by fine-tuning with `--dataset glaive_toolcall`.

-[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`.
+[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.

@@ -91,19 +97,22 @@ Compared to ChatGLM's [P-Tuning](https://github.com/THUDM/ChatGLM2-6B/tree/main/

 | 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)                   | 1.8B/7B/14B/72B             | 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.
@@ -123,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

@@ -167,6 +176,7 @@ Please refer to [constants.py](src/llmtuner/extras/constants.py) for a full list
 - [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)
@@ -174,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>

@@ -206,13 +217,13 @@ huggingface-cli login

 ### Hardware Requirement

-| Method | Bits |   7B  |  13B  |  30B  |   65B  |
-| ------ | ---- | ----- | ----- | ----- | ------ |
-| Full   |  16  | 140GB | 240GB | 520GB | 1200GB |
-| Freeze |  16  |  20GB |  40GB | 120GB |  240GB |
-| LoRA   |  16  |  16GB |  32GB |  80GB |  160GB |
-| QLoRA  |   8  |  10GB |  16GB |  40GB |   80GB |
-| QLoRA  |   4  |   6GB |  12GB |  24GB |   48GB |
+| 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

@@ -239,9 +250,9 @@ 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 Models (optional)
+### Use ModelScope Hub (optional)

-If you have trouble with downloading models from Hugging Face, you can use LLaMA-Factory together with ModelScope in the following manner.
+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
@@ -255,7 +266,7 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    ... # arguments (same as above)
 ```

-LLaMA Board also supports using the models on the ModelScope Hub.
+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
@@ -271,8 +282,8 @@ CUDA_VISIBLE_DEVICES=0 USE_MODELSCOPE_HUB=1 python src/train_web.py
 ```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 \
@@ -294,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 \
@@ -318,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 \
@@ -343,14 +354,14 @@ 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 \
@@ -374,14 +385,14 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
 ```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 \
@@ -449,7 +460,7 @@ deepspeed --num_gpus 8 --master_port=9901 src/train_bash.py \
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "min_loss_scale": 1
-  },  
+  },
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
@@ -469,20 +480,28 @@ deepspeed --num_gpus 8 --master_port=9901 src/train_bash.py \
 ```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
 ```

 > [!TIP]
@@ -493,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
@@ -503,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
@@ -513,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 \
@@ -528,12 +547,12 @@ 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 \
@@ -553,6 +572,7 @@ 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.
@@ -561,7 +581,7 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \

 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
@@ -6,7 +6,7 @@
 [![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)

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

 ## 更新日志

-[23/12/01] 我们支持了从 **[魔搭社区](https://modelscope.cn/models)** 下载预训练模型。详细用法请参照 [此教程](#使用魔搭社区可跳过)。
+[24/01/18] 我们针对绝大多数模型实现了 **Agent 微调**，微调时指定 `--dataset glaive_toolcall` 即可使模型获得工具调用能力。

-[23/10/21] 我们支持了 **[NEFTune](https://arxiv.org/abs/2310.05914)** 训练技巧。请使用 `--neft_alpha` 参数启用 NEFTune，例如 `--neft_alpha 5`。
+[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)）。

@@ -91,19 +97,22 @@ https://github.com/hiyouga/LLaMA-Factory/assets/16256802/6ba60acc-e2e2-4bec-b846

 | 模型名                                                   | 模型大小                     | 默认模块           | 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)                   | 1.8B/7B/14B/72B             | 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` 参数指定全部模块。
@@ -123,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 训练。

 ## 数据集

@@ -167,6 +176,7 @@ https://github.com/hiyouga/LLaMA-Factory/assets/16256802/6ba60acc-e2e2-4bec-b846
 - [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)
@@ -174,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>

@@ -206,13 +217,13 @@ huggingface-cli login

 ### 硬件依赖

-| 训练方法 | 精度 |   7B  |  13B  |  30B  |   65B  |
-| ------- | ---- | ----- | ----- | ----- | ------ |
-| 全参数   |  16  | 140GB | 240GB | 520GB | 1200GB |
-| 部分参数 |  16  |  20GB |  40GB | 120GB |  240GB |
-| LoRA    |  16  |  16GB |  32GB |  80GB |  160GB |
-| QLoRA   |   8  |  10GB |  16GB |  40GB |   80GB |
-| QLoRA   |   4  |   6GB |  12GB |  24GB |   48GB |
+| 训练方法 | 精度 |   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 |

 ## 如何使用

@@ -241,7 +252,7 @@ pip install https://github.com/jllllll/bitsandbytes-windows-webui/releases/downl

 ### 使用魔搭社区（可跳过）

-如果您在 Hugging Face 模型的下载中遇到了问题，可以通过下述方法使用魔搭社区。
+如果您在 Hugging Face 模型和数据集的下载中遇到了问题，可以通过下述方法使用魔搭社区。

 ```bash
 export USE_MODELSCOPE_HUB=1 # Windows 使用 `set USE_MODELSCOPE_HUB=1`
@@ -255,7 +266,7 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    ... # 参数同上
 ```

-LLaMA Board 同样支持魔搭社区的模型下载。
+LLaMA Board 同样支持魔搭社区的模型和数据集下载。

 ```bash
 CUDA_VISIBLE_DEVICES=0 USE_MODELSCOPE_HUB=1 python src/train_web.py
@@ -271,8 +282,8 @@ CUDA_VISIBLE_DEVICES=0 USE_MODELSCOPE_HUB=1 python src/train_web.py
 ```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 \
@@ -294,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 \
@@ -318,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 \
@@ -343,14 +354,14 @@ 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 \
@@ -374,14 +385,14 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
 ```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 \
@@ -449,7 +460,7 @@ deepspeed --num_gpus 8 --master_port=9901 src/train_bash.py \
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "min_loss_scale": 1
-  },  
+  },
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
@@ -464,25 +475,33 @@ deepspeed --num_gpus 8 --master_port=9901 src/train_bash.py \

 </details>

-### 合并 LoRA 权重并导出完整模型
+### 合并 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
 ```

 > [!TIP]
@@ -493,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
 ```

 ### 浏览器测试
@@ -503,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
 ```

 ### 模型评估
@@ -513,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 \
@@ -528,12 +547,12 @@ 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 \
@@ -553,6 +572,7 @@ 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。
@@ -561,7 +581,7 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \

 本仓库的代码依照 [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/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哈希值（可选，留空不影响训练）",
+  "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/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.3"
+__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,28 +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,
+    ChatCompletionStreamResponse,
+    Finish,
+    Function,
+    FunctionCall,
+    ModelCard,
+    ModelList,
+    Role,
    ScoreEvaluationRequest,
-    ScoreEvaluationResponse
-)
-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
+    ScoreEvaluationResponse,
 )


@@ -40,15 +43,22 @@ if is_uvicorn_available():


@asynccontextmanager
-async def lifespan(app: "FastAPI"): # collects GPU memory
+async def lifespan(app: "FastAPI"):  # collects GPU memory
    yield
    torch_gc()


-def to_json(data: BaseModel) -> str:
-    try: # pydantic v2
+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)


@@ -63,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")
@@ -74,91 +86,119 @@ def create_app(chat_model: "ChatModel") -> "FastAPI":
            raise HTTPException(status_code=status.HTTP_405_METHOD_NOT_ALLOWED, detail="Not allowed")

        if len(request.messages) == 0 or request.messages[-1].role != Role.USER:
-            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid request")
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid length")

-        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
+        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),
-                finish_reason=Finish.STOP if response.finish_reason == "stop" else Finish.LENGTH
-            ))
+            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)
@@ -168,7 +208,12 @@ def create_app(chat_model: "ChatModel") -> "FastAPI":

        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)

@@ -178,4 +223,4 @@ def create_app(chat_model: "ChatModel") -> "FastAPI":
 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,18 +87,18 @@ 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]

@@ -90,7 +110,7 @@ class ScoreEvaluationRequest(BaseModel):


 class ScoreEvaluationResponse(BaseModel):
-    id: Optional[str] = "scoreeval-default"
-    object: Optional[str] = "score.evaluation"
+    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,18 +1,18 @@
-import torch
-import tiktoken
 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
@@ -20,28 +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.can_generate = (finetuning_args.stage == "sft")
+        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)
@@ -56,16 +54,18 @@ class ChatModel:
        max_new_tokens = input_kwargs.pop("max_new_tokens", None)

        generating_args = self.generating_args.to_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"],
-            top_k=top_k or generating_args["top_k"],
-            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
-        ))
+        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"],
+                top_k=top_k or generating_args["top_k"],
+                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,
+            )
+        )

        if isinstance(num_return_sequences, int) and num_return_sequences > 1:
            generating_args["do_sample"] = True
@@ -81,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
@@ -89,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(
@@ -109,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(
-                response_text=response[i],
-                response_length=response_length,
-                prompt_length=prompt_length,
-                finish_reason="stop" if len(eos_index) else "length"
-            ))
+            results.append(
+                Response(
+                    response_text=response[i],
+                    response_length=response_length,
+                    prompt_length=prompt_length,
+                    finish_reason="stop" if len(eos_index) else "length",
+                )
+            )

        return 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

@@ -136,16 +133,7 @@ class ChatModel:
        yield from streamer

    @torch.inference_mode()
-    def get_scores(
-        self,
-        batch_input: List[str],
-        **input_kwargs
-    ) -> List[float]:
-        if isinstance(getattr(self.tokenizer, "tokenizer", None), tiktoken.Encoding): # for tiktoken tokenizer (Qwen)
-            kwargs = dict(allowed_special="all")
-        else:
-            kwargs = dict(add_special_tokens=True)
-
+    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")

@@ -154,9 +142,8 @@ class ChatModel:
            padding=True,
            truncation=True,
            max_length=max_length or getattr(self.model.config, "max_position_embeddings", 1024),
-            pad_to_multiple_of=8,
            return_tensors="pt",
-            **kwargs
+            add_special_tokens=True,
        ).to(device)

        input_ids: torch.Tensor = inputs["input_ids"]
--- 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,135 +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
+    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_dir = dataset_attr.folder

-    for dataset_attr in data_args.dataset_list:
-        logger.info("Loading dataset {}...".format(dataset_attr))
+    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_dir = dataset_attr.folder

-        if dataset_attr.load_from == "hf_hub":
-            data_path = dataset_attr.dataset_name
-            data_name = dataset_attr.subset
-            data_files = None
-        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
-        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))
-                    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)
-            else:
-                raise ValueError("File not found.")
-
-            assert data_path, "File extension must be txt, csv, json or jsonl."
-            checksum(data_files, dataset_attr.dataset_sha1)
+    elif dataset_attr.load_from == "file":
+        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 = 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 NotImplementedError
+            raise ValueError("File not found.")
+
+        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 and (dataset_attr.load_from != "file"))
+            streaming=(data_args.streaming and (dataset_attr.load_from != "file")),
+            **kwargs,
        )

-        if data_args.streaming and (dataset_attr.load_from == "file"):
-            dataset = dataset.to_iterable_dataset() # TODO: add num shards parameter
+    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

-        if max_samples is not None: # truncate dataset
-            dataset = dataset.select(range(min(len(dataset), max_samples)))
+    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))

-        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
+    return align_dataset(dataset, dataset_attr, data_args)

-                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]))

-                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:
@@ -141,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,275 +1,248 @@
-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 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 ...`
+    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]
+
+    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
+    # we drop the small remainder, and if the total_length < block_size, we exclude this batch
+    total_length = (total_length // block_size) * block_size
+    # split by chunks of cutoff_len
+    result = {
+        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+        for k, t in concatenated_examples.items()
+    }
+    return result


-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"],
+def preprocess_supervised_dataset(
+    examples: Dict[str, List[Any]],
    tokenizer: "PreTrainedTokenizer",
+    template: "Template",
    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)
+) -> 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": []}

-    if data_args.train_on_prompt and template.efficient_eos:
-        raise ValueError("Current template does not support `train_on_prompt`.")
+    for i in range(len(examples["prompt"])):
+        if len(examples["prompt"][i]) == 0 or len(examples["response"][i]) != 1:
+            continue

-    def preprocess_pretrain_dataset(examples: Dict[str, List[Any]]) -> Dict[str, List[List[int]]]:
-        # build grouped texts with format `X1 X2 X3 ...`
-        if isinstance(getattr(tokenizer, "tokenizer", None), tiktoken.Encoding): # for tiktoken tokenizer (Qwen)
-            kwargs = dict(allowed_special="all")
-        else:
-            kwargs = dict(add_special_tokens=True)
-
-        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
-        # we drop the small remainder, and if the total_length < block_size, we exclude this batch
-        total_length = (total_length // block_size) * block_size
-        # split by chunks of cutoff_len
-        result = {
-            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]]]:
-        # 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 != ""):
-                continue
-
-            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]
-
-                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
-
-            if template.efficient_eos:
-                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]]]:
-        # 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": []}
+        messages = examples["prompt"][i] + examples["response"][i]
        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 != ""):
-                continue
+        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:
+                source_mask = [tokenizer.eos_token_id] + [IGNORE_INDEX] * (len(source_ids) - 1)
+            else:
+                source_mask = [IGNORE_INDEX] * len(source_ids)

-            for turn_idx, (source_ids, target_ids) in enumerate(template.encode_multiturn(
-                tokenizer, query, response, history, system
-            )):
-                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
+            input_ids += source_ids + target_ids
+            labels += source_mask + target_ids

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

-        total_length = len(input_ids)
-        block_size = data_args.cutoff_len
-        # we drop the small remainder, and if the total_length < block_size, we exclude this batch
-        total_length = (total_length // block_size) * block_size
-        # split by chunks of cutoff_len
-        for i in range(0, total_length, block_size):
-            model_inputs["input_ids"].append(input_ids[i: i + block_size])
-            model_inputs["attention_mask"].append([1] * block_size)
-            model_inputs["labels"].append(labels[i: i + block_size])
+        model_inputs["input_ids"].append(input_ids)
+        model_inputs["attention_mask"].append([1] * len(input_ids))
+        model_inputs["labels"].append(labels)

-        return model_inputs
+    return model_inputs

-    def preprocess_unsupervised_dataset(examples: Dict[str, List[Any]]) -> 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 != ""):
-                continue
+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 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]
+        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)

-            if template.efficient_eos:
-                labels += [tokenizer.eos_token_id]
+            input_ids += source_ids + target_ids
+            labels += source_mask + target_ids

-            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]
+    if template.efficient_eos:
+        input_ids += [tokenizer.eos_token_id]
+        labels += [tokenizer.eos_token_id]

-            model_inputs["input_ids"].append(input_ids)
-            model_inputs["attention_mask"].append([1] * len(input_ids))
-            model_inputs["labels"].append(labels)
+    total_length = len(input_ids)
+    block_size = data_args.cutoff_len
+    # we drop the small remainder, and if the total_length < block_size, we exclude this batch
+    total_length = (total_length // block_size) * block_size
+    # split by chunks of cutoff_len
+    for i in range(0, total_length, block_size):
+        model_inputs["input_ids"].append(input_ids[i : i + block_size])
+        model_inputs["attention_mask"].append([1] * block_size)
+        model_inputs["labels"].append(labels[i : i + block_size])

-        return model_inputs
+    return model_inputs

-    def preprocess_pairwise_dataset(examples: Dict[str, List[Any]]) -> 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):
-                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)
+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": []}

-            if template.efficient_eos:
-                chosen_ids += [tokenizer.eos_token_id]
-                rejected_ids += [tokenizer.eos_token_id]
+    for i in range(len(examples["prompt"])):
+        if len(examples["prompt"][i]) == 0 or len(examples["response"][i]) != 1:
+            continue

-            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:
-        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(
-            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:
-        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"]))
-        print("chosen:\n{}".format(tokenizer.decode(example["chosen_ids"], skip_special_tokens=False)))
-        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:
-        print("input_ids:\n{}".format(example["input_ids"]))
-        print("inputs:\n{}".format(tokenizer.decode(example["input_ids"], skip_special_tokens=False)))
-
-    if stage == "pt":
-        preprocess_func = preprocess_pretrain_dataset
-        print_function = print_unsupervised_dataset_example
-    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
-    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"
-            )
-
-        dataset = dataset.map(
-            preprocess_func,
-            batched=True,
-            remove_columns=column_names,
-            **kwargs
+        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 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 template.efficient_eos:
+            labels += [tokenizer.eos_token_id]

-        if training_args.should_log:
-            try:
-                print_function(next(iter(dataset)))
-            except StopIteration:
-                raise RuntimeError("Empty dataset!")
+        model_inputs["input_ids"].append(input_ids)
+        model_inputs["attention_mask"].append([1] * len(input_ids))
+        model_inputs["labels"].append(labels)

-        return dataset
+    return model_inputs
+
+
+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 i in range(len(examples["prompt"])):
+        if len(examples["prompt"][i]) == 0 or len(examples["response"][i]) < 2:
+            continue
+
+        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]
+
+        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]], 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(
+            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]], 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"]))
+    print("chosen:\n{}".format(tokenizer.decode(example["chosen_ids"], skip_special_tokens=False)))
+    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]], 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 = 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:
+        if data_args.sft_packing:
+            preprocess_func = partial(
+                preprocess_packed_supervised_dataset, tokenizer=tokenizer, template=template, data_args=data_args
+            )
+        else:
+            preprocess_func = partial(
+                preprocess_supervised_dataset, tokenizer=tokenizer, template=template, data_args=data_args
+            )
+
+        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)
+
+    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,13 +38,18 @@ 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
+        if data_args.val_size > 1e-6:  # Split the dataset
            if data_args.streaming:
                val_set = dataset.take(int(data_args.val_size))
                train_set = dataset.skip(int(data_args.val_size))
@@ -57,5 +63,5 @@ def split_dataset(
            if data_args.streaming:
                dataset = dataset.shuffle(buffer_size=data_args.buffer_size, seed=training_args.seed)
            return {"train_dataset": dataset}
-    else: # do_eval or do_predict
+    else:  # do_eval or do_predict
        return {"eval_dataset": 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,41 +1,34 @@
 # Inspired by: https://github.com/hendrycks/test/blob/master/evaluate_flan.py

-import os
-import json
-import torch
 import inspect
-import tiktoken
-import numpy as np
-from tqdm import tqdm, trange
+import json
+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)
-        self.tokenizer.padding_side = "right" # avoid overflow issue in batched inference for llama2
+        self.tokenizer.padding_side = "right"  # avoid overflow issue in batched inference for llama2
        self.model = dispatch_model(self.model)
        self.template = get_template_and_fix_tokenizer(self.data_args.template, self.tokenizer)
        self.eval_template = get_eval_template(self.eval_args.lang)
-        self.choice_inputs = self._encode_choices()
-
-    def _encode_choices(self) -> List[int]:
-        if isinstance(getattr(self.tokenizer, "tokenizer", None), tiktoken.Encoding): # for tiktoken tokenizer (Qwen)
-            kwargs = dict(allowed_special="all")
-        else:
-            kwargs = dict(add_special_tokens=False)
-
-        return [self.tokenizer.encode(self.eval_template.prefix + ch, **kwargs)[-1] for ch in CHOICES]
+        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]:
@@ -46,16 +39,11 @@ class Evaluator:
        return [chr(ord("A") + offset.item()) for offset in torch.argmax(choice_probs, dim=-1)]

    def eval(self) -> None:
-        if "token" in inspect.signature(cached_file).parameters:
-            kwargs = {"token": self.model_args.hf_hub_token}
-        elif "use_auth_token" in inspect.signature(cached_file).parameters: # for transformers==4.31.0
-            kwargs = {"use_auth_token": self.model_args.hf_hub_token}
-
        mapping = cached_file(
-            path_or_repo_id = os.path.join(self.eval_args.task_dir, self.eval_args.task),
+            path_or_repo_id=os.path.join(self.eval_args.task_dir, self.eval_args.task),
            filename="mapping.json",
            cache_dir=self.model_args.cache_dir,
-            **kwargs
+            token=self.model_args.hf_hub_token,
        )

        with open(mapping, "r", encoding="utf-8") as f:
@@ -65,37 +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,
                cache_dir=self.model_args.cache_dir,
                download_mode=self.eval_args.download_mode,
-                token=self.model_args.hf_hub_token
+                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)
@@ -105,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([
-            "{:>15}: {:.2f}".format(category_name, 100 * np.mean(category_correct))
-            for category_name, category_correct in category_corrects.items() if len(category_correct)
-        ])
+        score_info = "\n".join(
+            [
+                "{:>15}: {:.2f}".format(category_name, 100 * np.mean(category_correct))
+                for category_name, category_correct in category_corrects.items()
+                if len(category_correct)
+            ]
+        )
        print(score_info)
        if self.eval_args.save_dir is not None:
            os.makedirs(self.eval_args.save_dir, exist_ok=False)
--- 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,56 +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.modeling_utils import custom_object_save, unwrap_model
-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__)


-def _save_model_with_valuehead(model: "AutoModelForCausalLMWithValueHead", output_dir: str) -> None:
-    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)
-    elif getattr(model.pretrained_model, "_auto_class", None): # must not a peft model
-        custom_object_save(model.pretrained_model, output_dir, config=model.pretrained_model.config)
-
-
-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:
-            _save_model_with_valuehead(
-                model=unwrap_model(kwargs.pop("model")),
-                output_dir=os.path.join(args.output_dir, "{}-{}".format(PREFIX_CHECKPOINT_DIR, state.global_step))
+            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,
            )

-    def on_train_end(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
-        r"""
-        Event called at the end of training.
-        """
-        if args.should_save:
-            _save_model_with_valuehead(model=unwrap_model(kwargs.pop("model")), output_dir=args.output_dir)
-

 class LogCallback(TrainerCallback):
-
    def __init__(self, runner=None):
        self.runner = runner
        self.in_training = False
@@ -116,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.
        """
@@ -142,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(
-                logs["loss"] or 0, logs["learning_rate"] or 0, logs["epoch"] or 0
-            ))
+            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,14 +1,18 @@
+from collections import OrderedDict, defaultdict
 from enum import Enum
-from collections import defaultdict, OrderedDict
 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"}
@@ -17,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()
@@ -26,18 +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, Dict[DownloadSource, 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():
@@ -56,19 +65,19 @@ register_model_group(
    models={
        "Baichuan-7B-Base": {
            DownloadSource.DEFAULT: "baichuan-inc/Baichuan-7B",
-            DownloadSource.MODELSCOPE: "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"
+            DownloadSource.MODELSCOPE: "baichuan-inc/Baichuan-13B-Base",
        },
        "Baichuan-13B-Chat": {
            DownloadSource.DEFAULT: "baichuan-inc/Baichuan-13B-Chat",
-            DownloadSource.MODELSCOPE: "baichuan-inc/Baichuan-13B-Chat"
-        }
+            DownloadSource.MODELSCOPE: "baichuan-inc/Baichuan-13B-Chat",
+        },
    },
    module="W_pack",
-    template="baichuan"
+    template="baichuan",
 )


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


@@ -100,18 +109,18 @@ register_model_group(
    models={
        "BLOOM-560M": {
            DownloadSource.DEFAULT: "bigscience/bloom-560m",
-            DownloadSource.MODELSCOPE: "AI-ModelScope/bloom-560m"
+            DownloadSource.MODELSCOPE: "AI-ModelScope/bloom-560m",
        },
        "BLOOM-3B": {
            DownloadSource.DEFAULT: "bigscience/bloom-3b",
-            DownloadSource.MODELSCOPE: "AI-ModelScope/bloom-3b"
+            DownloadSource.MODELSCOPE: "AI-ModelScope/bloom-3b",
        },
        "BLOOM-7B1": {
            DownloadSource.DEFAULT: "bigscience/bloom-7b1",
-            DownloadSource.MODELSCOPE: "AI-ModelScope/bloom-7b1"
-        }
+            DownloadSource.MODELSCOPE: "AI-ModelScope/bloom-7b1",
+        },
    },
-    module="query_key_value"
+    module="query_key_value",
 )


@@ -119,18 +128,18 @@ register_model_group(
    models={
        "BLOOMZ-560M": {
            DownloadSource.DEFAULT: "bigscience/bloomz-560m",
-            DownloadSource.MODELSCOPE: "AI-ModelScope/bloomz-560m"
+            DownloadSource.MODELSCOPE: "AI-ModelScope/bloomz-560m",
        },
        "BLOOMZ-3B": {
            DownloadSource.DEFAULT: "bigscience/bloomz-3b",
-            DownloadSource.MODELSCOPE: "AI-ModelScope/bloomz-3b"
+            DownloadSource.MODELSCOPE: "AI-ModelScope/bloomz-3b",
        },
        "BLOOMZ-7B1-mt": {
            DownloadSource.DEFAULT: "bigscience/bloomz-7b1-mt",
-            DownloadSource.MODELSCOPE: "AI-ModelScope/bloomz-7b1-mt"
-        }
+            DownloadSource.MODELSCOPE: "AI-ModelScope/bloomz-7b1-mt",
+        },
    },
-    module="query_key_value"
+    module="query_key_value",
 )


@@ -138,14 +147,14 @@ register_model_group(
    models={
        "BlueLM-7B-Base": {
            DownloadSource.DEFAULT: "vivo-ai/BlueLM-7B-Base",
-            DownloadSource.MODELSCOPE: "vivo-ai/BlueLM-7B-Base"
+            DownloadSource.MODELSCOPE: "vivo-ai/BlueLM-7B-Base",
        },
        "BlueLM-7B-Chat": {
            DownloadSource.DEFAULT: "vivo-ai/BlueLM-7B-Chat",
-            DownloadSource.MODELSCOPE: "vivo-ai/BlueLM-7B-Chat"
-        }
+            DownloadSource.MODELSCOPE: "vivo-ai/BlueLM-7B-Chat",
+        },
    },
-    template="bluelm"
+    template="bluelm",
 )


@@ -153,11 +162,11 @@ register_model_group(
    models={
        "ChatGLM2-6B-Chat": {
            DownloadSource.DEFAULT: "THUDM/chatglm2-6b",
-            DownloadSource.MODELSCOPE: "ZhipuAI/chatglm2-6b"
+            DownloadSource.MODELSCOPE: "ZhipuAI/chatglm2-6b",
        }
    },
    module="query_key_value",
-    template="chatglm2"
+    template="chatglm2",
 )


@@ -165,15 +174,15 @@ register_model_group(
    models={
        "ChatGLM3-6B-Base": {
            DownloadSource.DEFAULT: "THUDM/chatglm3-6b-base",
-            DownloadSource.MODELSCOPE: "ZhipuAI/chatglm3-6b-base"
+            DownloadSource.MODELSCOPE: "ZhipuAI/chatglm3-6b-base",
        },
        "ChatGLM3-6B-Chat": {
            DownloadSource.DEFAULT: "THUDM/chatglm3-6b",
-            DownloadSource.MODELSCOPE: "ZhipuAI/chatglm3-6b"
-        }
+            DownloadSource.MODELSCOPE: "ZhipuAI/chatglm3-6b",
+        },
    },
    module="query_key_value",
-    template="chatglm3"
+    template="chatglm3",
 )


@@ -181,76 +190,91 @@ register_model_group(
    models={
        "ChineseLLaMA2-1.3B": {
            DownloadSource.DEFAULT: "hfl/chinese-llama-2-1.3b",
-            DownloadSource.MODELSCOPE: "AI-ModelScope/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"
+            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"
+            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"
+            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"
+            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"
-        }
+            DownloadSource.MODELSCOPE: "AI-ModelScope/chinese-alpaca-2-13b",
+        },
    },
-    template="llama2_zh"
+    template="llama2_zh",
 )


 register_model_group(
    models={
-        "DeepseekLLM-7B-Base": {
+        "DeepSeekLLM-7B-Base": {
            DownloadSource.DEFAULT: "deepseek-ai/deepseek-llm-7b-base",
-            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-llm-7b-base"
+            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-llm-7b-base",
        },
-        "DeepseekLLM-67B-Base": {
+        "DeepSeekLLM-67B-Base": {
            DownloadSource.DEFAULT: "deepseek-ai/deepseek-llm-67b-base",
-            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-llm-67b-base"
+            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-llm-67b-base",
        },
-        "DeepseekLLM-7B-Chat": {
+        "DeepSeekLLM-7B-Chat": {
            DownloadSource.DEFAULT: "deepseek-ai/deepseek-llm-7b-chat",
-            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-llm-7b-chat"
+            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-llm-7b-chat",
        },
-        "DeepseekLLM-67B-Chat": {
+        "DeepSeekLLM-67B-Chat": {
            DownloadSource.DEFAULT: "deepseek-ai/deepseek-llm-67b-chat",
-            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-llm-67b-chat"
-        }
+            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-llm-67b-chat",
+        },
    },
-    template="deepseek"
+    template="deepseek",
 )


 register_model_group(
    models={
-        "DeepseekCoder-6.7B-Base": {
+        "DeepSeekCoder-6.7B-Base": {
            DownloadSource.DEFAULT: "deepseek-ai/deepseek-coder-6.7b-base",
-            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-coder-6.7b-base"
+            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-coder-6.7b-base",
        },
-        "DeepseekCoder-33B-Base": {
+        "DeepSeekCoder-33B-Base": {
            DownloadSource.DEFAULT: "deepseek-ai/deepseek-coder-33b-base",
-            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-coder-33b-base"
+            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-coder-33b-base",
        },
-        "DeepseekCoder-6.7B-Chat": {
+        "DeepSeekCoder-6.7B-Chat": {
            DownloadSource.DEFAULT: "deepseek-ai/deepseek-coder-6.7b-instruct",
-            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-coder-6.7b-instruct"
+            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-coder-6.7b-instruct",
        },
-        "DeepseekCoder-33B-Chat": {
+        "DeepSeekCoder-33B-Chat": {
            DownloadSource.DEFAULT: "deepseek-ai/deepseek-coder-33b-instruct",
-            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-coder-33b-instruct"
-        }
+            DownloadSource.MODELSCOPE: "deepseek-ai/deepseek-coder-33b-instruct",
+        },
    },
-    template="deepseekcoder"
+    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",
 )


@@ -258,31 +282,31 @@ register_model_group(
    models={
        "Falcon-7B": {
            DownloadSource.DEFAULT: "tiiuae/falcon-7b",
-            DownloadSource.MODELSCOPE: "AI-ModelScope/falcon-7b"
+            DownloadSource.MODELSCOPE: "AI-ModelScope/falcon-7b",
        },
        "Falcon-40B": {
            DownloadSource.DEFAULT: "tiiuae/falcon-40b",
-            DownloadSource.MODELSCOPE: "AI-ModelScope/falcon-40b"
+            DownloadSource.MODELSCOPE: "AI-ModelScope/falcon-40b",
        },
        "Falcon-180B": {
            DownloadSource.DEFAULT: "tiiuae/falcon-180b",
-            DownloadSource.MODELSCOPE: "modelscope/falcon-180B"
+            DownloadSource.MODELSCOPE: "modelscope/falcon-180B",
        },
        "Falcon-7B-Chat": {
            DownloadSource.DEFAULT: "tiiuae/falcon-7b-instruct",
-            DownloadSource.MODELSCOPE: "AI-ModelScope/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"
+            DownloadSource.MODELSCOPE: "AI-ModelScope/falcon-40b-instruct",
        },
        "Falcon-180B-Chat": {
            DownloadSource.DEFAULT: "tiiuae/falcon-180b-chat",
-            DownloadSource.MODELSCOPE: "modelscope/falcon-180B-chat"
-        }
+            DownloadSource.MODELSCOPE: "modelscope/falcon-180B-chat",
+        },
    },
    module="query_key_value",
-    template="falcon"
+    template="falcon",
 )


@@ -290,22 +314,46 @@ register_model_group(
    models={
        "InternLM-7B": {
            DownloadSource.DEFAULT: "internlm/internlm-7b",
-            DownloadSource.MODELSCOPE: "Shanghai_AI_Laboratory/internlm-7b"
+            DownloadSource.MODELSCOPE: "Shanghai_AI_Laboratory/internlm-7b",
        },
        "InternLM-20B": {
            DownloadSource.DEFAULT: "internlm/internlm-20b",
-            DownloadSource.MODELSCOPE: "Shanghai_AI_Laboratory/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"
+            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"
-        }
+            DownloadSource.MODELSCOPE: "Shanghai_AI_Laboratory/internlm-chat-20b",
+        },
    },
-    template="intern"
+    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",
 )


@@ -313,31 +361,28 @@ register_model_group(
    models={
        "LingoWhale-8B": {
            DownloadSource.DEFAULT: "deeplang-ai/LingoWhale-8B",
-            DownloadSource.MODELSCOPE: "DeepLang/LingoWhale-8B"
+            DownloadSource.MODELSCOPE: "DeepLang/LingoWhale-8B",
        }
    },
-    module="qkv_proj"
+    module="qkv_proj",
 )


 register_model_group(
    models={
-        "LLaMA-7B": {
-            DownloadSource.DEFAULT: "huggyllama/llama-7b",
-            DownloadSource.MODELSCOPE: "skyline2006/llama-7b"
-        },
+        "LLaMA-7B": {DownloadSource.DEFAULT: "huggyllama/llama-7b", DownloadSource.MODELSCOPE: "skyline2006/llama-7b"},
        "LLaMA-13B": {
            DownloadSource.DEFAULT: "huggyllama/llama-13b",
-            DownloadSource.MODELSCOPE: "skyline2006/llama-13b"
+            DownloadSource.MODELSCOPE: "skyline2006/llama-13b",
        },
        "LLaMA-30B": {
            DownloadSource.DEFAULT: "huggyllama/llama-30b",
-            DownloadSource.MODELSCOPE: "skyline2006/llama-30b"
+            DownloadSource.MODELSCOPE: "skyline2006/llama-30b",
        },
        "LLaMA-65B": {
            DownloadSource.DEFAULT: "huggyllama/llama-65b",
-            DownloadSource.MODELSCOPE: "skyline2006/llama-65b"
-        }
+            DownloadSource.MODELSCOPE: "skyline2006/llama-65b",
+        },
    }
 )

@@ -346,30 +391,30 @@ register_model_group(
    models={
        "LLaMA2-7B": {
            DownloadSource.DEFAULT: "meta-llama/Llama-2-7b-hf",
-            DownloadSource.MODELSCOPE: "modelscope/Llama-2-7b-ms"
+            DownloadSource.MODELSCOPE: "modelscope/Llama-2-7b-ms",
        },
        "LLaMA2-13B": {
            DownloadSource.DEFAULT: "meta-llama/Llama-2-13b-hf",
-            DownloadSource.MODELSCOPE: "modelscope/Llama-2-13b-ms"
+            DownloadSource.MODELSCOPE: "modelscope/Llama-2-13b-ms",
        },
        "LLaMA2-70B": {
            DownloadSource.DEFAULT: "meta-llama/Llama-2-70b-hf",
-            DownloadSource.MODELSCOPE: "modelscope/Llama-2-70b-ms"
+            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"
+            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"
+            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"
-        }
+            DownloadSource.MODELSCOPE: "modelscope/Llama-2-70b-chat-ms",
+        },
    },
-    template="llama2"
+    template="llama2",
 )


@@ -377,14 +422,33 @@ register_model_group(
    models={
        "Mistral-7B": {
            DownloadSource.DEFAULT: "mistralai/Mistral-7B-v0.1",
-            DownloadSource.MODELSCOPE: "AI-ModelScope/Mistral-7B-v0.1"
+            DownloadSource.MODELSCOPE: "AI-ModelScope/Mistral-7B-v0.1",
        },
        "Mistral-7B-Chat": {
            DownloadSource.DEFAULT: "mistralai/Mistral-7B-Instruct-v0.1",
-            DownloadSource.MODELSCOPE: "AI-ModelScope/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"
+    template="mistral",
+)
+
+
+register_model_group(
+    models={
+        "Mixtral-8x7B": {
+            DownloadSource.DEFAULT: "mistralai/Mixtral-8x7B-v0.1",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/Mixtral-8x7B-v0.1",
+        },
+        "Mixtral-8x7B-Chat": {
+            DownloadSource.DEFAULT: "mistralai/Mixtral-8x7B-Instruct-v0.1",
+            DownloadSource.MODELSCOPE: "AI-ModelScope/Mixtral-8x7B-Instruct-v0.1",
+        },
+    },
+    template="mistral",
 )


@@ -392,93 +456,87 @@ register_model_group(
    models={
        "OpenChat3.5-7B-Chat": {
            DownloadSource.DEFAULT: "openchat/openchat_3.5",
-            DownloadSource.MODELSCOPE: "myxiongmodel/openchat_3.5"
+            DownloadSource.MODELSCOPE: "myxiongmodel/openchat_3.5",
        }
    },
-    template="openchat"
+    template="openchat",
 )


 register_model_group(
    models={
-        "Phi1.5-1.3B": {
-            DownloadSource.DEFAULT: "microsoft/phi-1_5",
-            DownloadSource.MODELSCOPE: "allspace/PHI_1-5"
-        }
-    },
-    module="Wqkv"
+        "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": {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"
+            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"
+            DownloadSource.MODELSCOPE: "qwen/Qwen-14B-Chat",
        },
        "Qwen-72B-Chat": {
            DownloadSource.DEFAULT: "Qwen/Qwen-72B-Chat",
-            DownloadSource.MODELSCOPE: "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"
+            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"
+            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"
+            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"
+            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"
+            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"
+            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"
+            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"
-        }
+            DownloadSource.MODELSCOPE: "qwen/Qwen-72B-Chat-Int4",
+        },
    },
    module="c_attn",
-    template="qwen"
+    template="qwen",
+)
+
+
+register_model_group(
+    models={
+        "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",
 )


@@ -486,7 +544,7 @@ register_model_group(
    models={
        "Skywork-13B-Base": {
            DownloadSource.DEFAULT: "Skywork/Skywork-13B-base",
-            DownloadSource.MODELSCOPE: "skywork/Skywork-13B-base"
+            DownloadSource.MODELSCOPE: "skywork/Skywork-13B-base",
        }
    }
 )
@@ -496,60 +554,51 @@ register_model_group(
    models={
        "Vicuna1.5-7B-Chat": {
            DownloadSource.DEFAULT: "lmsys/vicuna-7b-v1.5",
-            DownloadSource.MODELSCOPE: "Xorbits/vicuna-7b-v1.5"
+            DownloadSource.MODELSCOPE: "Xorbits/vicuna-7b-v1.5",
        },
        "Vicuna1.5-13B-Chat": {
            DownloadSource.DEFAULT: "lmsys/vicuna-13b-v1.5",
-            DownloadSource.MODELSCOPE: "Xorbits/vicuna-13b-v1.5"
-        }
+            DownloadSource.MODELSCOPE: "Xorbits/vicuna-13b-v1.5",
+        },
    },
-    template="vicuna"
+    template="vicuna",
 )


 register_model_group(
    models={
-        "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"
-        }
+        "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="xuanyuan"
+    template="xuanyuan",
 )


 register_model_group(
    models={
-        "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-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",
        },
        "XVERSE-7B-Chat": {
            DownloadSource.DEFAULT: "xverse/XVERSE-7B-Chat",
-            DownloadSource.MODELSCOPE: "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"
-        }
+            DownloadSource.MODELSCOPE: "xverse/XVERSE-13B-Chat",
+        },
+        "XVERSE-65B-Chat": {
+            DownloadSource.DEFAULT: "xverse/XVERSE-65B-Chat",
+            DownloadSource.MODELSCOPE: "xverse/XVERSE-65B-Chat",
+        },
    },
-    template="xverse"
+    template="xverse",
 )


@@ -557,37 +606,52 @@ register_model_group(
    models={
        "Yayi-7B": {
            DownloadSource.DEFAULT: "wenge-research/yayi-7b-llama2",
-            DownloadSource.MODELSCOPE: "AI-ModelScope/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"
-        }
+            DownloadSource.MODELSCOPE: "AI-ModelScope/yayi-13b-llama2",
+        },
    },
-    template="yayi"
+    template="yayi",
 )


 register_model_group(
    models={
-        "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-34B-Chat": {
-            DownloadSource.DEFAULT: "01-ai/Yi-34B-Chat",
-            DownloadSource.MODELSCOPE: "01ai/Yi-34B-Chat"
+        "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",
        },
        "Yi-34B-int8-Chat": {
            DownloadSource.DEFAULT: "01-ai/Yi-34B-Chat-8bits",
-            DownloadSource.MODELSCOPE: "01ai/Yi-34B-Chat-8bits"
-        }
+            DownloadSource.MODELSCOPE: "01ai/Yi-34B-Chat-8bits",
+        },
    },
-    template="yi"
+    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",
 )


@@ -595,12 +659,12 @@ register_model_group(
    models={
        "Zephyr-7B-Alpha-Chat": {
            DownloadSource.DEFAULT: "HuggingFaceH4/zephyr-7b-alpha",
-            DownloadSource.MODELSCOPE: "AI-ModelScope/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"
-        }
+            DownloadSource.MODELSCOPE: "modelscope/zephyr-7b-beta",
+        },
    },
-    template="zephyr"
+    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,35 +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

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

@@ -68,6 +77,74 @@ def count_parameters(model: torch.nn.Module) -> Tuple[int, int]:
    return trainable_params, all_param


+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:
+        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":
    r"""
    Gets logits processor that removes NaN and Inf logits.
@@ -89,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.
@@ -115,12 +181,11 @@ def try_download_model_from_ms(model_args: "ModelArguments") -> None:
        return

    try:
-        from modelscope import snapshot_download # type: ignore
+        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
+            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`")
--- a/src/llmtuner/extras/packages.py
+++ b/src/llmtuner/extras/packages.py
@@ -9,52 +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")
-_requests_available = is_package_available("requests")
-_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 _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,224 +1,197 @@
 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 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["Cache"] = None,
+    output_attentions: bool = False,
+    **kwargs,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+    bsz, q_len, _ = hidden_states.size()

-    def forward(
-        self,
-        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
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-        bsz, q_len, _ = hidden_states.size()
+    query_states = self.q_proj(hidden_states)
+    key_states = self.k_proj(hidden_states)
+    value_states = self.v_proj(hidden_states)

-        query_states = self.q_proj(hidden_states)
-        key_states = self.k_proj(hidden_states)
-        value_states = self.v_proj(hidden_states)
+    query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+    value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

-        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+    kv_seq_len = key_states.shape[-2]
+    if past_key_value is not None:
+        kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)

-        kv_seq_len = key_states.shape[-2]
-        if past_key_value is not None:
-            kv_seq_len += past_key_value[0].shape[-2]
+    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)

-        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:
+        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)

-        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)
+    key_states = repeat_kv(key_states, self.num_key_value_groups)
+    value_states = repeat_kv(value_states, self.num_key_value_groups)

-        past_key_value = (key_states, value_states) if use_cache else None
+    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

-        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)
-
-        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 = 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)
-                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)
-            if attention_mask is not None:
-                attention_mask = attention_mask[:, :, :groupsz, :groupsz].repeat(num_groups, 1, 1, 1)
-
-        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
-
-        if attention_mask is not None:
-            attn_weights = attn_weights + attention_mask
-
-        # upcast attention to fp32
-        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
-        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 = attn_output.reshape(bsz, q_len, self.hidden_size)
-        attn_output = self.o_proj(attn_output)
-
-        if not output_attentions:
-            attn_weights = None
-
-        return attn_output, attn_weights, past_key_value
-
-
-class LlamaFlashAttention2(LlamaAttention):
-
-    def forward(
-        self,
-        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
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-        # LlamaFlashAttention2 attention does not support output_attentions
-        output_attentions = False
-
-        bsz, q_len, _ = hidden_states.size()
-
-        query_states = self.q_proj(hidden_states)
-        key_states = self.k_proj(hidden_states)
-        value_states = self.v_proj(hidden_states)
-
-        # FlashAttention requires the input to have the shape (bsz, seq_len, n_heads, head_dim)
-        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-
-        kv_seq_len = key_states.shape[-2]
-        if past_key_value is not None:
-            kv_seq_len += past_key_value[0].shape[-2]
-
-        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)
-
-        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)
-
-        query_states = query_states.transpose(1, 2) # (bsz, seq_len, n_heads, head_dim)
-        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)
-
-        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)
-                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)
-
-        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,
+        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,
            )
-            attn_output = pad_input(attn_output_unpad, indices_q, bsz, q_len)
+            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)
+        if attention_mask is not None:
+            attention_mask = attention_mask[:, :, :groupsz, :groupsz].repeat(num_groups, 1, 1, 1)
+
+    attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+
+    # 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 = attn_output.reshape(bsz, q_len, self.hidden_size)
+    attn_output = self.o_proj(attn_output)
+
+    if not output_attentions:
+        attn_weights = None
+
+    return attn_output, attn_weights, past_key_value
+
+
+# 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,
+    **kwargs,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+    # LlamaFlashAttention2 attention does not support output_attentions
+    output_attentions = False
+
+    bsz, q_len, _ = hidden_states.size()
+
+    query_states = self.q_proj(hidden_states)
+    key_states = self.k_proj(hidden_states)
+    value_states = self.v_proj(hidden_states)
+
+    # FlashAttention requires the input to have the shape (bsz, seq_len, n_heads, head_dim)
+    query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+    value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+    kv_seq_len = key_states.shape[-2]
+    if past_key_value is not None:
+        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:
+        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)
+
+    key_states = repeat_kv(key_states, self.num_key_value_groups)
+    value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+    query_states = query_states.transpose(1, 2)  # (bsz, seq_len, n_heads, head_dim)
+    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:
-            attn_output = flash_attn_func(
-                query_states, key_states, value_states, 0.0, softmax_scale=None, causal=True
+            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,
            )
+            return state.reshape(bsz * num_groups, groupsz, self.num_heads, self.head_dim)

-        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)
-            ))
+        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[:, :, :groupsz, :groupsz].repeat(num_groups, 1, 1, 1)

-        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
-        attn_output = self.o_proj(attn_output)
+    attn_output: torch.Tensor = self._flash_attention_forward(
+        query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
+    )

-        if not output_attentions:
-            attn_weights = None
+    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),
+            )
+        )

-        return attn_output, attn_weights, past_key_value
+    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
+    attn_output = self.o_proj(attn_output)
+
+    if not output_attentions:
+        attn_weights = None
+
+    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
@@ -20,7 +22,7 @@ def smooth(scalars: List[float]) -> List[float]:
    """
    last = scalars[0]
    smoothed = list()
-    weight = 1.8 * (1 / (1 + math.exp(-0.05 * len(scalars))) - 0.5) # a sigmoid function
+    weight = 1.8 * (1 / (1 + math.exp(-0.05 * len(scalars))) - 0.5)  # a sigmoid function
    for next_val in scalars:
        smoothed_val = last * weight + (1 - weight) * next_val
        smoothed.append(smoothed_val)
@@ -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,33 +1,5 @@
-import os
-import json
-from typing import List, Literal, Optional
 from dataclasses import dataclass, field
-
-
-DATA_CONFIG = "dataset_info.json"
-
-
-@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
@@ -36,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):
@@ -125,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, DATA_CONFIG), "r") as f:
-                dataset_info = json.load(f)
-        except Exception as err:
-            if self.dataset is not None:
-                raise ValueError("Cannot open {} due to {}.".format(os.path.join(self.dataset_dir, DATA_CONFIG), str(err)))
-            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 {}.".format(name, DATA_CONFIG))
-
-            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
@@ -10,17 +10,18 @@ class FreezeArguments:
    """
    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."}
+        default=3, metadata={"help": "The number of trainable layers for partial-parameter (freeze) fine-tuning."}
    )


@@ -31,33 +32,32 @@ class LoraArguments:
    """
    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."
+        },
    )
-    lora_alpha: Optional[float] = field(
-        default=None,
-        metadata={"help": "The scale factor for LoRA fine-tuning (default: lora_rank * 2.0)."}
-    )
-    lora_dropout: Optional[float] = field(
-        default=0.1,
-        metadata={"help": "Dropout rate for the LoRA fine-tuning."}
-    )
-    lora_rank: Optional[int] = field(
-        default=8,
-        metadata={"help": "The intrinsic dimension for LoRA fine-tuning."}
+    lora_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-1.5 choices: [\"Wqkv\", \"out_proj\", \"fc1\", \"fc2\"], \
-                  Others choices: the same as LLaMA."}
+        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.'
+        },
    )
-    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_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,61 +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."}
+        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."}
+        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", "api"]] = field(
        default="lora",
-        metadata={"help": "The type of the reward model in PPO training. Lora model only supports lora training."}
+        metadata={"help": "The type of the reward model in PPO training. Lora model only supports lora training."},
    )


@@ -130,32 +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."}
-    )
-    export_size: Optional[int] = field(
-        default=1,
-        metadata={"help": "The file shard size (in GB) of 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):
@@ -165,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."
--- 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,56 +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 or modelscope.cn/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):
@@ -67,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,89 +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 _verify_model_args(model_args: "ModelArguments", finetuning_args: "FinetuningArguments") -> None:
-    if model_args.quantization_bit is not None and finetuning_args.finetuning_type != "lora":
-        raise ValueError("Quantization is only compatible with the LoRA method.")
+def _parse_args(parser: "HfArgumentParser", args: Optional[Dict[str, Any]] = None) -> Tuple[Any]:
+    if args is not None:
+        return parser.parse_dict(args)

-    if (
-        model_args.checkpoint_dir is not None
-        and len(model_args.checkpoint_dir) != 1
-        and finetuning_args.finetuning_type != "lora"
-    ):
-        raise ValueError("Multiple checkpoints are only available for LoRA tuning.")
+    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_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:
-        # 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.")

@@ -99,12 +109,12 @@ def get_train_args(args: Optional[Dict[str, Any]] = None) -> _TRAIN_CLS:
    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"] and (not all([data_attr.ranking for data_attr in data_args.dataset_list])):
-        raise ValueError("Please use ranked datasets for reward modeling or DPO training.")
-
    if finetuning_args.stage == "ppo" and model_args.shift_attn:
        raise ValueError("PPO training is incompatible with S^2-Attn.")

+    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.")

@@ -116,7 +126,7 @@ def get_train_args(args: Optional[Dict[str, Any]] = None) -> _TRAIN_CLS:

    _verify_model_args(model_args, finetuning_args)

-    if training_args.do_train and model_args.quantization_bit is not None and (not finetuning_args.upcast_layernorm):
+    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):
@@ -139,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:
@@ -153,14 +170,22 @@ def get_train_args(args: Optional[Dict[str, Any]] = None) -> _TRAIN_CLS:
            training_args_dict = training_args.to_dict()
            training_args_dict.update(dict(resume_from_checkpoint=last_checkpoint))
            training_args = Seq2SeqTrainingArguments(**training_args_dict)
-            logger.info("Resuming training from {}. Change `output_dir` or use `overwrite_output_dir` to avoid.".format(
-                training_args.resume_from_checkpoint
-            ))
+            logger.info(
+                "Resuming training from {}. Change `output_dir` or use `overwrite_output_dir` to avoid.".format(
+                    training_args.resume_from_checkpoint
+                )
+            )

-    if finetuning_args.stage in ["rm", "ppo"] and training_args.resume_from_checkpoint is not None:
-        logger.warning("Add {} to `checkpoint_dir` to resume training from checkpoint.".format(
-            training_args.resume_from_checkpoint
-        ))
+    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
    model_args.compute_dtype = (
@@ -169,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.
@@ -182,7 +212,8 @@ 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.")
@@ -193,7 +224,8 @@ def get_infer_args(args: Optional[Dict[str, Any]] = None) -> _INFER_CLS:


 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.")
--- 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, get_modelcard_args, 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
+
+        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)]
+        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)]  # noqa: C416

        trainable_layers = []
        for module_name in finetuning_args.name_module_trainable:
@@ -62,47 +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 model_args.adapter_name_or_path is not None:
            is_mergeable = True
-            if getattr(model, "quantization_method", None) == "gptq":
-                assert len(model_args.checkpoint_dir) == 1, "GPTQ quantized model only accepts a single checkpoint."
+            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_trainable and finetuning_args.resume_lora_training) or (not is_mergeable):
-                checkpoints_to_merge, checkpoint_to_resume = model_args.checkpoint_dir[:-1], model_args.checkpoint_dir[-1]
-            else:
-                checkpoints_to_merge = model_args.checkpoint_dir
+            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

-            for checkpoint in checkpoints_to_merge:
-                model = PeftModel.from_pretrained(model, checkpoint)
+            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:
+                adapter_to_merge = model_args.adapter_name_or_path
+
+            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)
            else:
                target_modules = finetuning_args.lora_target

-            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
-            )
-            model = get_peft_model(model, lora_config)
+            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.checkpoint_dir is not None:
-        logger.info("Loaded fine-tuned model from checkpoint(s): {}".format(",".join(model_args.checkpoint_dir)))
+            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,
+                    modules_to_save=finetuning_args.additional_target,
+                    **peft_kwargs,
+                )
+                model = get_peft_model(model, lora_config)
+
+        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,56 +1,39 @@
-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 get_logger
-from llmtuner.extras.misc import count_parameters, infer_optim_dtype, try_download_model_from_ms
-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,
-    add_valuehead: Optional[bool] = False
-) -> Tuple[PreTrainedModel, "PreTrainedTokenizer"]:
+    add_valuehead: Optional[bool] = False,
+) -> Tuple["PreTrainedModel", "PreTrainedTokenizer"]:
    r"""
    Loads pretrained model and tokenizer.

@@ -63,174 +46,88 @@ def load_model_and_tokenizer(
        "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.")
-            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.")
+            model, _ = FastLlamaModel.from_pretrained(**unsloth_kwargs)
+        elif getattr(config, "model_type", None) == "mistral":
+            model, _ = FastMistralModel.from_pretrained(**unsloth_kwargs)
        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 gptq or awq)
-    if getattr(config, "quantization_config", None):
-        if model_args.quantization_bit is not None: # remove bnb quantization
-            model_args.quantization_bit = None
-        config_kwargs["device_map"] = {"": int(os.environ.get("LOCAL_RANK", "0"))}
-        quantization_config = getattr(config, "quantization_config", None)
-        logger.info("Loading {}-bit quantized model.".format(quantization_config.get("bits", -1)))
+    if model is None:
+        model = AutoModelForCausalLM.from_pretrained(
+            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,
+        )

-    # 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.")
+    patch_model(model, tokenizer, model_args, is_trainable)
+    register_autoclass(config, model, tokenizer)

-        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)
-
-        if 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"] = {"": int(os.environ.get("LOCAL_RANK", "0"))}
-        logger.info("Quantizing model to {} bit.".format(model_args.quantization_bit))
-
-    # Load pre-trained models (without valuehead)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_to_load,
-        config=config,
-        torch_dtype=model_args.compute_dtype,
-        low_cpu_mem_usage=(not is_deepspeed_zero3_enabled()),
-        **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)
-
-    # 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)

-    # Prepare model with valuehead for RLHF
    if add_valuehead:
        model: "AutoModelForCausalLMWithValueHead" = AutoModelForCausalLMWithValueHead.from_pretrained(model)
-        setattr(model, "_keys_to_ignore_on_save", [name for name, _ in model.named_parameters() if "pretrained_model" in name])
-        setattr(model, "tie_weights", MethodType(lambda _: None, model)) # use empty method
-        vhead_path = (
-            model_args.checkpoint_dir[-1] if model_args.checkpoint_dir is not None else model_args.model_name_or_path
-        )
+        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(
-        trainable_params, all_param, 100 * trainable_params / all_param
-    ))
+    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,17 +1,19 @@
-import torch
 import inspect
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple
+from typing import TYPE_CHECKING, Any, Dict, List

+import torch
+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__)
@@ -19,27 +21,32 @@ 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, "quantization_method", None): # already set on current device
+    if getattr(model, "quantization_method", None):  # already set on current device
        return model

-    if torch.cuda.device_count() > 1 and getattr(model.config, "model_type", None) != "chatglm":
+    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") -> List[str]:
@@ -51,6 +58,7 @@ def find_all_linear_modules(model: "PreTrainedModel") -> List[str]:
        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))
@@ -61,10 +69,7 @@ def find_all_linear_modules(model: "PreTrainedModel") -> List[str]:

    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)))
@@ -72,112 +77,49 @@ def find_all_linear_modules(model: "PreTrainedModel") -> List[str]:


 def get_modelcard_args(
-    model_args: "ModelArguments",
-    data_args: "DataArguments",
-    finetuning_args: "FinetuningArguments"
+    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
-    }
-
-    if "token" in inspect.signature(cached_file).parameters:
-        kwargs["token"] = model_args.hf_hub_token
-    elif "use_auth_token" in inspect.signature(cached_file).parameters: # for transformers==4.31.0
-        kwargs["use_auth_token"] = model_args.hf_hub_token
-    else:
-        logger.warning("Ignore `hf_hub_token` since matched parameter is not found.")
-
-    try:
-        vhead_file = cached_file(filename=WEIGHTS_NAME, **kwargs)
-        return torch.load(vhead_file, map_location="cpu")
-    except Exception as err:
-        logger.info("Failed to load {}: {}".format(WEIGHTS_NAME, str(err)))
+    kwargs = {"path_or_repo_id": path_or_repo_id, "cache_dir": model_args.cache_dir, "token": model_args.hf_hub_token}

    try:
        from safetensors import safe_open
-        vhead_file = cached_file(filename=SAFE_WEIGHTS_NAME, **kwargs)
-        with safe_open(vhead_file, framework="pt", device="cpu") as f:
-            return {
-                "v_head.summary.weight": f.get_tensor("v_head.summary.weight"),
-                "v_head.summary.bias": f.get_tensor("v_head.summary.bias")
-            }
-    except Exception as err:
-        logger.info("Failed to load {}: {}".format(SAFE_WEIGHTS_NAME, str(err)))

-    logger.warning("Provided path ({}) does not contain valuehead weights.".format(path_or_repo_id))
+        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=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


-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 and getattr(model, "supports_gradient_checkpointing", False):
-        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


@@ -20,7 +21,7 @@ class DPODataCollatorWithPadding(DataCollatorForSeq2Seq):
            padded_tensor = self.label_pad_token_id * torch.ones_like(feature)
            padded_tensor[start:end] = feature[start:end]
            padded_labels.append(padded_tensor)
-        return torch.stack(padded_labels, dim=0).contiguous() # in contiguous memory
+        return torch.stack(padded_labels, dim=0).contiguous()  # in contiguous memory

    def __call__(self, features: Sequence[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
        r"""
@@ -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({
-                    "input_ids": feature["prompt_ids"] + feature[key],
-                    "attention_mask": [1] * (prompt_len + answer_len)
-                })
+                concatenated_features.append(
+                    {
+                        "input_ids": feature["prompt_ids"] + feature[key],
+                        "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.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)
-                ): # quantized models are already set on the correct device
+                    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
+        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,20 +1,23 @@
 # Inspired by: https://github.com/huggingface/trl/blob/main/examples/research_projects/stack_llama_2/scripts/dpo_llama2.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.constants import IGNORE_INDEX
-from llmtuner.extras.ploting import plot_loss
-from llmtuner.hparams import ModelArguments
-from llmtuner.model import load_model_and_tokenizer
-from llmtuner.train.dpo.collator import DPODataCollatorWithPadding
-from llmtuner.train.dpo.trainer import CustomDPOTrainer
-from llmtuner.train.utils import create_modelcard_and_push, create_ref_model
+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(
@@ -22,38 +25,39 @@ 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)
-    dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="rm")
+    dataset = get_dataset(tokenizer, model_args, data_args, training_args, stage="rm")
    data_collator = DPODataCollatorWithPadding(
        tokenizer=tokenizer,
        pad_to_multiple_of=8,
-        label_pad_token_id=IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
+        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
+    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)

    # Update arguments
    training_args_dict = training_args.to_dict()
-    training_args_dict.update(dict(remove_unused_columns=False)) # important for pairwise dataset
+    training_args_dict.update(dict(remove_unused_columns=False))  # important for pairwise dataset
    training_args = Seq2SeqTrainingArguments(**training_args_dict)

    # 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
@@ -69,7 +73,7 @@ def run_dpo(
    # Evaluation
    if training_args.do_eval:
        metrics = trainer.evaluate(metric_key_prefix="eval")
-        if id(model) == id(ref_model): # unable to compute rewards without a reference model
+        if id(model) == id(ref_model):  # unable to compute rewards without a reference model
            remove_keys = [key for key in metrics.keys() if "rewards" in key]
            for key in remove_keys:
                metrics.pop(key)
--- 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,27 +1,28 @@
+import math
 import os
 import sys
-import math
-import torch
-from tqdm import tqdm
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple

-from transformers import GenerationConfig, Trainer, TrainerState, TrainerControl
-from transformers.utils import WEIGHTS_NAME, SAFE_WEIGHTS_NAME
-from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
+import torch
+from tqdm import tqdm
+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, get_rewards_from_server, 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__)
@@ -40,7 +41,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
        generating_args: "GeneratingArguments",
        callbacks: List["TrainerCallback"],
        reward_model: "AutoModelForCausalLMWithValueHead",
-        **kwargs
+        **kwargs,
    ):
        PPOTrainer.__init__(self, **kwargs)

@@ -52,7 +53,7 @@ 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()
@@ -61,7 +62,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
            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")
@@ -71,7 +72,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
                if not (
                    getattr(reward_model.pretrained_model, "is_loaded_in_8bit", False)
                    or getattr(reward_model.pretrained_model, "is_loaded_in_4bit", False)
-                ): # quantized models are already set on the correct device
+                ):  # quantized models are already set on the correct device
                    self.reward_model = self._prepare_deepspeed(self.reward_model)
            else:
                self.reward_model = self.accelerator.prepare_model(self.reward_model, evaluation_mode=True)
@@ -111,9 +112,11 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
            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(
+                "  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))
@@ -138,10 +141,12 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
            self.model.eval()

            # Get inputs
-            self.tokenizer.padding_side = "right" # change padding side
+            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)
@@ -154,7 +159,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer):

            # Run PPO step
            stats = self.step(queries, responses, rewards)
-            self.tokenizer.padding_side = "left" # restore padding side
+            self.tokenizer.padding_side = "left"  # restore padding side
            loss_meter.update(float(stats["ppo/loss/total"]), n=len(rewards))
            reward_meter.update(torch.stack(rewards).mean().item(), n=len(rewards))

@@ -163,18 +168,18 @@ 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
            self.log_callback.on_step_end(self.args, self.state, self.control)

-            if self.is_local_process_zero() and (step+1) % self.args.logging_steps == 0:
+            if self.is_local_process_zero() and (step + 1) % self.args.logging_steps == 0:
                logs = dict(
                    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
@@ -183,10 +188,10 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
                loss_meter.reset()
                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)
-                ))
+            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_callback.on_save(
                    self.args, self.state, self.control, model=self.accelerator.unwrap_model(self.model)
                )
@@ -204,33 +209,36 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
        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)
        generate_output: torch.Tensor = unwrapped_model.generate(
-            generation_config=self.generation_config,
-            logits_processor=get_logits_processor(),
-            **batch
+            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 = batch["input_ids"].detach().cpu()
-        response = generate_output[:, batch["input_ids"].size(-1):].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:
-                response_length = 1 # allow empty response
+                response_length = 1  # allow empty response
            else:
                response_length = response_index[-1].item() + 1

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

        return queries, responses

@@ -239,7 +247,7 @@ 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.
@@ -259,17 +267,17 @@ class CustomPPOTrainer(PPOTrainer, Trainer):

        batch = self.prepare_model_inputs(queries, responses)

-        with torch.cuda.amp.autocast(dtype=self.model_args.compute_dtype): # support bf16
+        with torch.cuda.amp.autocast(dtype=self.model_args.compute_dtype):  # support bf16
            _, _, values = reward_model(**batch, output_hidden_states=True, return_dict=True)

-        if getattr(unwrapped_model.config, "model_type", None) == "chatglm": # assume same architecture
+        if getattr(unwrapped_model.config, "model_type", None) == "chatglm":  # assume same architecture
            values = torch.transpose(values, 0, 1)

        rewards = []
        for i in range(values.size(0)):
            end_indexes = (batch["input_ids"][i] != self.tokenizer.pad_token_id).nonzero()
            end_index = end_indexes[-1].item() if len(end_indexes) else 0
-            rewards.append(values[i, end_index].float().detach().cpu()) # use fp32 type
+            rewards.append(values[i, end_index].float().detach().cpu())  # use fp32 type

        if self.finetuning_args.reward_model_type == "lora":
            replace_model(unwrapped_model, target="default")
@@ -284,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.
@@ -307,7 +315,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
            input_ids = input_kwargs["input_ids"]
            attention_mask = input_kwargs["attention_mask"]

-            with torch.cuda.amp.autocast(dtype=self.model_args.compute_dtype): # support bf16
+            with torch.cuda.amp.autocast(dtype=self.model_args.compute_dtype):  # support bf16
                logits, _, values = model(**input_kwargs)

            unwrapped_model: "AutoModelForCausalLMWithValueHead" = self.accelerator.unwrap_model(self.model)
@@ -320,14 +328,12 @@ class CustomPPOTrainer(PPOTrainer, Trainer):

            for j in range(len(query_batch)):
                start = len(query_batch[j]) - 1
-                if attention_mask[j, 0] == 0: # offset left padding
+                if attention_mask[j, 0] == 0:  # offset left padding
                    start += attention_mask[j, :].nonzero()[0].item()
                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
@@ -361,9 +367,9 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
                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"
+                    " 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(self.args.should_save, output_dir, [WEIGHTS_NAME, SAFE_WEIGHTS_NAME])
-                self.model.save_checkpoint(output_dir) # wrapped model
+                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,8 +1,10 @@
 import json
-import torch
 from typing import TYPE_CHECKING, Dict, List, Literal, Optional

-from llmtuner.extras.packages import is_requests_available
+import torch
+
+from ...extras.packages import is_requests_available
+

 if TYPE_CHECKING:
    from transformers import PreTrainedModel
@@ -21,16 +23,18 @@ def get_rewards_from_server(server_url: str, messages: List[str]) -> List[torch.


 def replace_model(model: "AutoModelForCausalLMWithValueHead", target: Literal["default", "reward"]) -> None:
-    if target == "reward": # save default head temporarily
+    if target == "reward":  # save default head temporarily
        valuehead_state_dict: Dict[str, torch.Tensor] = model.v_head.state_dict()
        setattr(model, "default_head_weight", valuehead_state_dict["summary.weight"].detach().clone())
        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({
-        "summary.weight": model.get_buffer("{}_head_weight".format(target)).detach().clone(),
-        "summary.bias": model.get_buffer("{}_head_bias".format(target)).detach().clone()
-    })
+    model.pretrained_model.set_adapter(target)  # set the LoRA adapter to be active
+    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(),
+        }
+    )


 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,13 +29,14 @@ 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, add_valuehead=True)
-    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
+    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
@@ -55,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
@@ -70,7 +75,7 @@ def run_ppo(
        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
@@ -79,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,
@@ -88,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(resume_from_checkpoint=training_args.resume_from_checkpoint)
        ppo_trainer.save_model()
-        ppo_trainer.save_state() # must be called after save_model to have a folder
+        if 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,17 +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 load_model_and_tokenizer
-from llmtuner.train.utils import create_modelcard_and_push
+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(
@@ -19,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)
-    dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="pt")
+    dataset = get_dataset(tokenizer, model_args, data_args, training_args, stage="pt")
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    # Initialize our Trainer
@@ -33,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
--- 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__)
@@ -21,13 +23,10 @@ class PairwiseTrainer(Trainer):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
-        self.can_return_loss = True # override property to return eval_loss
+        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.
@@ -68,9 +67,9 @@ class PairwiseTrainer(Trainer):
            assert div_index > 0
            chosen_trunc_rewards = chosen_rewards[i, div_index:end_index]
            rejected_trunc_rewards = rejected_rewards[i, div_index:end_index]
-            if return_outputs: # use the score on the last token except pad token for inference
-                chosen_scores.append(chosen_rewards[i, chosen_length-1])
-                rejected_scores.append(rejected_rewards[i, rejected_length-1])
+            if return_outputs:  # use the score on the last token except pad token for inference
+                chosen_scores.append(chosen_rewards[i, chosen_length - 1])
+                rejected_scores.append(rejected_rewards[i, rejected_length - 1])
            loss += -torch.nn.functional.logsigmoid(chosen_trunc_rewards - rejected_trunc_rewards).mean()

        loss = loss / batch_size
@@ -80,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,20 +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 load_model_and_tokenizer
-from llmtuner.train.rm.collator import PairwiseDataCollatorWithPadding
-from llmtuner.train.rm.metric import compute_accuracy
-from llmtuner.train.rm.trainer import PairwiseTrainer
-from llmtuner.train.utils import create_modelcard_and_push
+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(
@@ -22,16 +26,17 @@ 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, add_valuehead=True)
-    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, 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()
-    training_args_dict.update(dict(remove_unused_columns=False)) # important for pairwise dataset
+    training_args_dict.update(dict(remove_unused_columns=False))  # important for pairwise dataset
    training_args = Seq2SeqTrainingArguments(**training_args_dict)

    # Initialize our Trainer
@@ -40,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(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()
--- 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
@@ -33,16 +35,16 @@ class CustomSeq2SeqTrainer(Seq2SeqTrainer):

        Subclass and override to inject custom behavior.
        """
-        labels = inputs["labels"].detach().clone() if "labels" in inputs else None # backup labels
+        labels = inputs["labels"].detach().clone() if "labels" in inputs else None  # backup labels
        if self.args.predict_with_generate:
            assert self.tokenizer.padding_side == "left", "This method only accepts left-padded tensor."
            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: # truncate the labels instead of padding the inputs (llama2 fp16 compatibility)
+            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( # ignore the returned labels (may be truncated)
+        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,23 +53,16 @@ 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.
        """
        assert self.tokenizer.pad_token_id is not None, "Pad token is required."
        padded_tensor = self.tokenizer.pad_token_id * torch.ones_like(tgt_tensor)
-        padded_tensor[:, -src_tensor.shape[-1]:] = src_tensor # adopt left-padding
-        return padded_tensor.contiguous() # in contiguous memory
+        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,15 +74,23 @@ 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}")

-        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)
+        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
+                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_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)

        with open(output_prediction_file, "w", encoding="utf-8") as writer:
--- a/src/llmtuner/train/sft/workflow.py
+++ b/src/llmtuner/train/sft/workflow.py
@@ -1,20 +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 load_model_and_tokenizer
-from llmtuner.train.sft.metric import ComputeMetrics
-from llmtuner.train.sft.trainer import CustomSeq2SeqTrainer
-from llmtuner.train.utils import create_modelcard_and_push
+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(
@@ -23,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)
-    dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="sft")
+    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
+        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=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
+        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(
-        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
-    ))
+    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,
+        )
+    )
    training_args = Seq2SeqTrainingArguments(**training_args_dict)

    # Initialize our Trainer
@@ -54,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`
@@ -76,7 +83,7 @@ def run_sft(
    # Evaluation
    if training_args.do_eval:
        metrics = trainer.evaluate(metric_key_prefix="eval", **gen_kwargs)
-        if training_args.predict_with_generate: # eval_loss will be wrong if predict_with_generate is enabled
+        if training_args.predict_with_generate:  # eval_loss will be wrong if predict_with_generate is enabled
            metrics.pop("eval_loss", None)
        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)
@@ -84,7 +91,7 @@ def run_sft(
    # Predict
    if training_args.do_predict:
        predict_results = trainer.predict(dataset, metric_key_prefix="predict", **gen_kwargs)
-        if training_args.predict_with_generate: # predict_loss will be wrong if predict_with_generate is enabled
+        if training_args.predict_with_generate:  # predict_loss will be wrong if predict_with_generate is enabled
            predict_results.metrics.pop("predict_loss", None)
        trainer.log_metrics("predict", predict_results.metrics)
        trainer.save_metrics("predict", predict_results.metrics)
--- 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
@@ -36,19 +41,48 @@ def run_exp(args: Optional[Dict[str, Any]] = None, callbacks: Optional[List["Tra

 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)

-    if getattr(model, "quantization_method", None) in ["gptq", "awq"]:
-        raise ValueError("Cannot export a GPTQ or AWQ quantized model.")
+    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.")

-    model.config.use_cache = True
-    model.save_pretrained(finetuning_args.export_dir, max_shard_size="{}GB".format(finetuning_args.export_size))
+    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.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,15 +1,18 @@
-import torch
 from typing import TYPE_CHECKING, Optional, Union

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

 if TYPE_CHECKING:
    from transformers import Seq2SeqTrainingArguments, Trainer
    from transformers.modeling_utils import PreTrainedModel
    from trl import AutoModelForCausalLMWithValueHead
-    from llmtuner.hparams import DataArguments
+
+    from ..hparams import DataArguments


 logger = get_logger(__name__)
@@ -20,7 +23,7 @@ def create_modelcard_and_push(
    model_args: "ModelArguments",
    data_args: "DataArguments",
    training_args: "Seq2SeqTrainingArguments",
-    finetuning_args: "FinetuningArguments"
+    finetuning_args: "FinetuningArguments",
 ) -> None:
    if training_args.do_train:
        if training_args.push_to_hub:
@@ -33,9 +36,7 @@ def create_modelcard_and_push(


 def create_ref_model(
-    model_args: "ModelArguments",
-    finetuning_args: "FinetuningArguments",
-    add_valuehead: Optional[bool] = False
+    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.
@@ -44,11 +45,13 @@ 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(
-            model_name_or_path=finetuning_args.ref_model,
-            checkpoint_dir=finetuning_args.ref_model_checkpoint,
-            quantization_bit=finetuning_args.ref_model_quantization_bit
-        ))
+        ref_model_args_dict.update(
+            dict(
+                model_name_or_path=finetuning_args.ref_model,
+                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(
@@ -68,9 +71,7 @@ def create_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.
@@ -81,24 +82,30 @@ def create_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
+        for name, param in model.named_parameters():  # https://github.com/huggingface/peft/issues/1090
            if "default" in name:
-                param.data = param.data.to(torch.float32) # trainable params should in fp32
+                param.data = param.data.to(torch.float32)  # trainable params should in fp32
        vhead_params = load_valuehead_params(finetuning_args.reward_model, model_args)
        assert vhead_params is not None, "Reward model is not correctly loaded."
        model.register_buffer("reward_head_weight", vhead_params["v_head.summary.weight"], persistent=False)
        model.register_buffer("reward_head_bias", vhead_params["v_head.summary.bias"], persistent=False)
-        model.register_buffer("default_head_weight", torch.zeros_like(vhead_params["v_head.summary.weight"]), persistent=False)
-        model.register_buffer("default_head_bias", torch.zeros_like(vhead_params["v_head.summary.bias"]), persistent=False)
+        model.register_buffer(
+            "default_head_weight", torch.zeros_like(vhead_params["v_head.summary.weight"]), persistent=False
+        )
+        model.register_buffer(
+            "default_head_bias", torch.zeros_like(vhead_params["v_head.summary.bias"]), persistent=False
+        )
        logger.info("Loaded adapter weights of reward model from {}".format(finetuning_args.reward_model))
        return None
    else:
        reward_model_args_dict = model_args.to_dict()
-        reward_model_args_dict.update(dict(
-            model_name_or_path=finetuning_args.reward_model,
-            checkpoint_dir=finetuning_args.reward_model_checkpoint,
-            quantization_bit=finetuning_args.reward_model_quantization_bit
-        ))
+        reward_model_args_dict.update(
+            dict(
+                model_name_or_path=finetuning_args.reward_model,
+                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(
--- 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 gradio as gr
-from gradio.components import Component # cannot use TYPE_CHECKING here
-from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional, Tuple
+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 ..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

-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

 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
@@ -26,11 +26,12 @@ class WebChatModel(ChatModel):
        self.tokenizer = None
        self.generating_args = GeneratingArguments()

-        if not lazy_init: # read arguments from command line
+        if not lazy_init:  # read arguments from command line
            super().__init__()

-        if demo_mode: # load demo_config.json if exists
+        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)
@@ -38,7 +39,7 @@ class WebChatModel(ChatModel):
                super().__init__(args)
            except AssertionError:
                print("Please provided model name and template in `demo_config.json`.")
-            except:
+            except Exception:
                print("Cannot find `demo_config.json` at current directory.")

    @property
@@ -63,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)

@@ -104,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,39 +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
-)

-from llmtuner.extras.constants import (
+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
+    DownloadSource,
 )
-from llmtuner.extras.misc import use_modelscope
-from llmtuner.hparams.data_args import DATA_CONFIG
+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"
-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:
@@ -48,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}


@@ -65,13 +54,13 @@ 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()
-    path_dict: Dict[DownloadSource, str] = SUPPORTED_MODELS.get(model_name, [])
-    model_path = user_config["path_dict"].get(model_name, None) or path_dict.get(DownloadSource.DEFAULT, "")
+    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
+    ):  # replace path
        model_path = path_dict.get(DownloadSource.MODELSCOPE)
    return model_path

@@ -90,18 +79,20 @@ 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, Dict[str, Any]]:
--- 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(
-        system=system,
-        query=query,
-        submit_btn=submit_btn,
-        clear_btn=clear_btn,
-        max_new_tokens=max_new_tokens,
-        top_p=top_p,
-        temperature=temperature
+    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,
+        ),
    )
--- 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]:
-    with open(os.path.join(dataset_dir, DATA_CONFIG), "r", encoding="utf-8") as f:
-        dataset_info = json.load(f)
+    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,9 +33,7 @@ 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)
@@ -41,9 +42,7 @@ def create_eval_tab(engine: "Engine") -> Dict[str, "Component"]:
        output_dir = gr.Textbox()

    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
-    ))
+    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()
@@ -58,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,47 +1,66 @@
-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_size=max_shard_size
+        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]
@@ -51,9 +70,11 @@ def save_model(

 def create_export_tab(engine: "Engine") -> Dict[str, "Component"]:
    with gr.Row():
-        export_dir = gr.Textbox()
        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)

@@ -63,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,55 +22,40 @@ 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
-    ) # do not save config since the below line will save
+    )  # do not save config since the below line will save

    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,17 +1,17 @@
-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
@@ -26,10 +26,7 @@ class Engine:
        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"]}
@@ -49,13 +46,17 @@ class Engine:
                else:
                    yield self._form_dict({"eval.resume_btn": {"value": True}})
            else:
-                yield self._form_dict({
-                    "train.output_dir": {"value": "train_" + get_time()},
-                    "eval.output_dir": {"value": "eval_" + 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 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": "样例"
-        }
+        "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": "学习率预热采用的步数。"
-        }
+        "en": {"label": "Warmup steps", "info": "Number of steps used for warmup."},
+        "zh": {"label": "预热步数", "info": "学习率预热采用的步数。"},
    },
-    "neft_alpha": {
-        "en": {
-            "label": "NEFTune Alpha",
-            "info": "Magnitude of noise adding to embedding vectors."
-        },
-        "zh": {
-            "label": "NEFTune 噪声参数",
-            "info": "嵌入向量所添加的噪声大小。"
-        }
+    "neftune_alpha": {
+        "en": {"label": "NEFTune Alpha", "info": "Magnitude of noise adding to embedding vectors."},
+        "zh": {"label": "NEFTune 噪声参数", "info": "嵌入向量所添加的噪声大小。"},
    },
-    "train_on_prompt": {
+    "sft_packing": {
        "en": {
-            "label": "Train on prompt",
-            "info": "Compute loss on the prompt tokens in supervised fine-tuning."
+            "label": "Pack sequences",
+            "info": "Pack sequences into samples of fixed length 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."
-        },
-        "zh": {
-            "label": "DPO beta 参数",
-            "info": "DPO 损失函数中 beta 超参数大小。"
-        }
+        "en": {"label": "DPO beta", "info": "Value of the beta parameter in the DPO loss."},
+        "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": "Output dir",
-            "info": "Directory for saving results."
-        },
-        "zh": {
-            "label": "输出目录",
-            "info": "保存结果的路径。"
-        }
+        "en": {"label": "Output dir", "info": "Directory for saving results."},
+        "zh": {"label": "输出目录", "info": "保存结果的路径。"},
    },
-    "output_box": {
-        "en": {
-            "value": "Ready."
-        },
-        "zh": {
-            "value": "准备就绪。"
-        }
+    "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": "单个模型文件的最大大小。"},
    },
-    "loss_viewer": {
-        "en": {
-            "label": "Loss"
-        },
-        "zh": {
-            "label": "损失"
-        }
+    "export_quantization_bit": {
+        "en": {"label": "Export quantization bit.", "info": "Quantizing the exported model."},
+        "zh": {"label": "导出量化等级", "info": "量化导出模型。"},
    },
-    "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."
-        },
-        "zh": {
-            "label": "导出目录",
-            "info": "保存导出模型的文件夹路径。"
-        }
+        "en": {"label": "Export dir", "info": "Directory to save exported model."},
+        "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,45 +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")
-            ])
+        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="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=get_save_dir(get("top.model_name"), get("top.finetuning_type"), get("eval.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)
@@ -235,9 +243,11 @@ 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)
--- 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,13 +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)
    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,11 +3,12 @@
 # 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
+from deepspeed.accelerator import get_accelerator  # type: ignore
+from deepspeed.profiling.flops_profiler import get_model_profile  # type: ignore

 from llmtuner import ChatModel

@@ -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,21 +3,23 @@
 # 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
-BASE_BS = 4_000_000 # from llama paper
+BASE_LR = 3e-4  # 1.5e-4 for 30B-70B models
+BASE_BS = 4_000_000  # from llama paper


 def calculate_lr(
@@ -25,21 +27,22 @@ def calculate_lr(
    dataset: str,
    cutoff_len: int,  # i.e. maximum input length during training
    batch_size: int,  # total batch size, namely (batch size * gradient accumulation * world size)
-    is_mistral: bool, # mistral model uses a smaller learning rate,
-    dataset_dir: Optional[str] = "data"
+    is_mistral: bool,  # mistral model uses a smaller learning rate,
+    dataset_dir: Optional[str] = "data",
 ):
-    model_args, data_args, training_args, finetuning_args, _ = get_train_args(dict(
-        stage="sft",
-        model_name_or_path=model_name_or_path,
-        dataset=dataset,
-        dataset_dir=dataset_dir,
-        template="default",
-        cutoff_len=cutoff_len,
-        output_dir="dummy_dir"
-    ))
-    trainset = get_dataset(model_args, data_args)
-    _, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, is_trainable=False, stage="sft")
-    trainset = preprocess_dataset(trainset, tokenizer, data_args, training_args, stage="sft")
+    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",
+        )
+    )
+    _, 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
@@ -49,14 +52,16 @@ def calculate_lr(
        valid_tokens += torch.sum(batch["labels"] != IGNORE_INDEX).item()
        total_tokens += torch.numel(batch["labels"])

-    batch_max_len = cutoff_len * batch_size # max tokens in a batch
+    batch_max_len = cutoff_len * batch_size  # max tokens in a batch
    valid_ratio = valid_tokens / total_tokens
    batch_valid_len = batch_max_len * valid_ratio
-    lr = BASE_LR * math.sqrt(batch_valid_len / BASE_BS) # lr ~ sqrt(batch_size)
+    lr = BASE_LR * math.sqrt(batch_valid_len / BASE_BS)  # lr ~ sqrt(batch_size)
    lr = lr / 6.0 if is_mistral else lr
-    print("Optimal learning rate is {:.2e} for valid ratio% {:.2f} and effective batch size {:.2f}".format(
-        lr, valid_ratio * 100, batch_valid_len
-    ))
+    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,57 +4,65 @@
 # 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, :]
-            llama2_state_dict[key.replace("W_pack", "k_proj")] = value[proj_size:2*proj_size, :]
-            llama2_state_dict[key.replace("W_pack", "v_proj")] = value[2*proj_size:, :]
+            llama2_state_dict[key.replace("W_pack", "k_proj")] = value[proj_size : 2 * proj_size, :]
+            llama2_state_dict[key.replace("W_pack", "v_proj")] = value[2 * proj_size :, :]
        elif "lm_head" in key:
            llama2_state_dict[key] = torch.nn.functional.normalize(value)
        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():
-        torch.save(shard, os.path.join(output_dir, shard_file))
-    
+    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,18 +76,14 @@ 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_config(input_dir, output_dir)    
+    save_weight(input_dir, output_dir, shard_size, save_safetensors)
+    save_config(input_dir, output_dir)


 if __name__ == "__main__":
--- 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.v_proj")] = value[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():
-        torch.save(shard, os.path.join(output_dir, shard_file))
+    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,18 +128,14 @@ 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)
-    save_config(input_dir, output_dir, torch_dtype)    
+    torch_dtype = save_weight(input_dir, output_dir, shard_size, save_safetensors)
+    save_config(input_dir, output_dir, torch_dtype)


 if __name__ == "__main__":
--- 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)