Release v0.1.8

Former-commit-id: d9666411375964d334d0a93ec162b27e05f70d49
update flashattn, fix ppo save model
2023-09-11 17:31:34 +08:00 · 2023-09-11 17:25:36 +08:00 · 2023-09-10 22:23:23 +08:00 · 2023-09-10 21:04:20 +08:00 · 2023-09-10 21:01:20 +08:00 · 2023-09-10 20:52:21 +08:00
45 changed files with 1133 additions and 726 deletions
--- a/README.md
+++ b/README.md
@@ -12,57 +12,48 @@

 ## Changelog

+[23/09/10] Now we support using **[FlashAttention](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 (experimental feature).
+
 [23/08/18] Now we support **resuming training**, upgrade `transformers` to `4.31.0` to enjoy this feature.

 [23/08/12] Now we support **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.

-[23/08/11] Now we support **[DPO training](https://arxiv.org/abs/2305.18290)** for instruction-tuned models. See [this example](#dpo-training) to train your models (experimental feature).
-
-[23/08/03] Now we support training the **Qwen-7B** model in this repo. Try `--model_name_or_path Qwen/Qwen-7B-Chat` and `--lora_target c_attn` arguments to train the Qwen-7B model. Remember to use `--template chatml` argument when you are using the Qwen-7B-Chat model.
+[23/08/11] Now we support **[DPO training](https://arxiv.org/abs/2305.18290)** for instruction-tuned models. See [this example](#dpo-training) to train your models.

 [23/07/31] Now we support **dataset streaming**. Try `--streaming` and `--max_steps 10000` arguments to load your dataset in streaming mode.

-[23/07/29] We release two instruction-tuned 13B models at Hugging Face. See these Hugging Face Repos ([LLaMA-2](https://huggingface.co/hiyouga/Llama-2-Chinese-13b-chat) / [Baichuan](https://huggingface.co/hiyouga/baichuan-13b-sft)) for details.
-
-[23/07/19] Now we support training the **LLaMA-2** models in this repo. Try `--model_name_or_path meta-llama/Llama-2-7b-hf` argument to use the LLaMA-2 model. Remember to use `--template llama2` argument when you are using the LLaMA-2-chat model.
+[23/07/29] We release two instruction-tuned 13B models at Hugging Face. See these Hugging Face Repos ([LLaMA-2](https://huggingface.co/hiyouga/Llama-2-Chinese-13b-chat) / [Baichuan](https://huggingface.co/hiyouga/Baichuan-13B-sft)) for details.

 [23/07/18] Now we develop an **all-in-one Web UI** for training, evaluation and inference. Try `train_web.py` to fine-tune models in your Web browser. Thank [@KanadeSiina](https://github.com/KanadeSiina) and [@codemayq](https://github.com/codemayq) for their efforts in the development.

-[23/07/11] Now we support training the **Baichuan-13B** model in this repo. Try `--model_name_or_path baichuan-inc/Baichuan-13B-Base` and `--lora_target W_pack` arguments to train the Baichuan-13B model. Remember to use `--template baichuan` argument when you are using the Baichuan-13B-Chat model.
-
 [23/07/09] Now we release **[FastEdit](https://github.com/hiyouga/FastEdit)** ⚡🩹, an easy-to-use package for editing the factual knowledge of large language models efficiently. Please follow [FastEdit](https://github.com/hiyouga/FastEdit) if you are interested.

-[23/07/07] Now we support training the **InternLM-7B** model in this repo. Try `--model_name_or_path internlm/internlm-7b` argument to use the InternLM model. Remember to use `--template intern` argument when you are using the InternLM-chat model.
-
-[23/07/05] Now we support training the **Falcon-7B/40B** models in this repo. Try `--model_name_or_path tiiuae/falcon-7b` and `--lora_target query_key_value` arguments to use the Falcon model.
-
-[23/06/29] We provide a **reproducible example** of training a chat model using instruction-following datasets, see this [Hugging Face Repo](https://huggingface.co/hiyouga/baichuan-7b-sft) for details.
+[23/06/29] We provide a **reproducible example** of training a chat model using instruction-following datasets, see [Baichuan-7B-sft](https://huggingface.co/hiyouga/Baichuan-7B-sft) for details.

 [23/06/22] Now we align the [demo API](src/api_demo.py) with the [OpenAI's](https://platform.openai.com/docs/api-reference/chat) format where you can insert the fine-tuned model in **arbitrary ChatGPT-based applications**.

-[23/06/15] Now we support training the **Baichuan-7B** model in this repo. Try `--model_name_or_path baichuan-inc/Baichuan-7B` and `--lora_target W_pack` arguments to use the Baichuan-7B model.
-
 [23/06/03] Now we support quantized training and inference (aka **[QLoRA](https://github.com/artidoro/qlora)**). Try `--quantization_bit 4/8` argument to work with quantized models.

-[23/05/31] Now we support training the **BLOOM & BLOOMZ** models in this repo. Try `--model_name_or_path bigscience/bloomz-7b1-mt` and `--lora_target query_key_value` arguments to use the BLOOMZ model.
-
 ## Supported Models

-| Model                                                    | Model size                  | Default module    | Template |
-| -------------------------------------------------------- | --------------------------- | ----------------- |----------|
-| [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   |
-| [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   | -        |
-| [Falcon](https://huggingface.co/tiiuae/falcon-7b)        | 7B/40B                      | query_key_value   | -        |
-| [Baichuan](https://github.com/baichuan-inc/baichuan-13B) | 7B/13B                      | W_pack            | baichuan |
-| [InternLM](https://github.com/InternLM/InternLM)         | 7B                          | q_proj,v_proj     | intern   |
-| [Qwen](https://github.com/QwenLM/Qwen-7B)                | 7B                          | c_attn            | chatml   |
-| [XVERSE](https://github.com/xverse-ai/XVERSE-13B)        | 13B                         | q_proj,v_proj     | -        |
-| [ChatGLM2](https://github.com/THUDM/ChatGLM2-6B)         | 6B                          | query_key_value   | chatglm2 |
+| Model                                                    | Model size                  | Default module    | Template  |
+| -------------------------------------------------------- | --------------------------- | ----------------- | --------- |
+| [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    |
+| [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   | -         |
+| [Falcon](https://huggingface.co/tiiuae/falcon-7b)        | 7B/40B                      | query_key_value   | -         |
+| [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 |
+| [InternLM](https://github.com/InternLM/InternLM)         | 7B                          | q_proj,v_proj     | intern    |
+| [Qwen](https://github.com/QwenLM/Qwen-7B)                | 7B                          | c_attn            | chatml    |
+| [XVERSE](https://github.com/xverse-ai/XVERSE-13B)        | 13B                         | q_proj,v_proj     | xverse    |
+| [ChatGLM2](https://github.com/THUDM/ChatGLM2-6B)         | 6B                          | query_key_value   | chatglm2  |

- **Default module** is used for the `--lora_target` argument. Please use `python src/train_bash.py -h` to see all available options.
- For the "base" models, the `--template` argument can be chosen from `default`, `alpaca`, `vicuna` etc. But make sure to use the corresponding template for the "chat" models.
+> [!NOTE]
+> **Default module** is used for the `--lora_target` argument, you can use `--lora_target all` to specify all the available modules.
+>
+> For the "base" models, the `--template` argument can be chosen from `default`, `alpaca`, `vicuna` etc. But make sure to use the corresponding template for the "chat" models.

 ## Supported Training Approaches

@@ -74,7 +65,8 @@
 | PPO Training           |                    |                    | :white_check_mark: | :white_check_mark: |
 | DPO Training           | :white_check_mark: |                    | :white_check_mark: | :white_check_mark: |

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

 ## Provided Datasets

@@ -105,6 +97,7 @@
  - [Web QA (zh)](https://huggingface.co/datasets/suolyer/webqa)
  - [UltraChat (en)](https://github.com/thunlp/UltraChat)
  - [WebNovel (zh)](https://huggingface.co/datasets/zxbsmk/webnovel_cn)
+  - [Ad Gen (zh)](https://huggingface.co/datasets/HasturOfficial/adgen)
 - For reward modeling or DPO training:
  - [HH-RLHF (en)](https://huggingface.co/datasets/Anthropic/hh-rlhf)
  - [Open Assistant (multilingual)](https://huggingface.co/datasets/OpenAssistant/oasst1)
@@ -123,7 +116,7 @@ huggingface-cli login

 - Python 3.8+ and PyTorch 1.13.1+
 - 🤗Transformers, Datasets, Accelerate, PEFT and TRL
- sentencepiece and tiktoken
+- sentencepiece, protobuf and tiktoken
 - jieba, rouge-chinese and nltk (used at evaluation)
 - gradio and matplotlib (used in web_demo.py)
 - uvicorn, fastapi and sse-starlette (used in api_demo.py)
@@ -136,7 +129,8 @@ And **powerful GPUs**!

 Please refer to `data/example_dataset` for checking the details about the format of dataset files. You can either use a single `.json` file or a [dataset loading script](https://huggingface.co/docs/datasets/dataset_script) with multiple files to create a custom dataset.

-Note: please update `data/dataset_info.json` to use your custom dataset. About the format of this file, please refer to `data/README.md`.
+> [!NOTE]
+> Please update `data/dataset_info.json` to use your custom dataset. About the format of this file, please refer to `data/README.md`.

 ### Dependence Installation (optional)

@@ -162,10 +156,14 @@ CUDA_VISIBLE_DEVICES=0 python src/train_web.py

 We strongly recommend using the all-in-one Web UI for newcomers since it can also generate training scripts **automatically**.

-Currently the web UI only supports training on **a single GPU**.
+> [!WARNING]
+> Currently the web UI only supports training on **a single GPU**.

 ### Train on a single GPU

+> [!IMPORTANT]
+> If you want to train models on multiple GPUs, please refer to [Distributed Training](#distributed-training).
+
 #### Pre-Training

 ```bash
@@ -174,7 +172,6 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --model_name_or_path path_to_llama_model \
    --do_train \
    --dataset wiki_demo \
-    --template default \
    --finetuning_type lora \
    --lora_target q_proj,v_proj \
    --output_dir path_to_pt_checkpoint \
@@ -299,19 +296,13 @@ accelerate config # configure the environment
 accelerate launch src/train_bash.py # arguments (same as above)
 ```

-<details><summary>Example config.yaml for training with DeepSpeed ZeRO-2</summary>
+<details><summary>Example config for LoRA training</summary>

 ```yaml
 compute_environment: LOCAL_MACHINE
-deepspeed_config:
-  gradient_accumulation_steps: 4
-  gradient_clipping: 0.5
-  offload_optimizer_device: none
-  offload_param_device: none
-  zero3_init_flag: false
-  zero_stage: 2
-distributed_type: DEEPSPEED
+distributed_type: MULTI_GPU
 downcast_bf16: 'no'
+gpu_ids: all
 machine_rank: 0
 main_training_function: main
 mixed_precision: fp16
@@ -335,10 +326,11 @@ deepspeed --num_gpus 8 --master_port=9901 src/train_bash.py \
    ... # arguments (same as above)
 ```

-<details><summary>Example ds_config.json for training with DeepSpeed ZeRO-2</summary>
+<details><summary>Example config for full-parameter training with DeepSpeed ZeRO-2</summary>

 ```json
 {
+  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
@@ -386,7 +378,8 @@ python src/api_demo.py \
    --checkpoint_dir path_to_checkpoint
 ```

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

 ### CLI Demo

@@ -425,7 +418,8 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --predict_with_generate
 ```

-We recommend using `--per_device_eval_batch_size=1` and `--max_target_length 128` at 4/8-bit evaluation.
+> [!NOTE]
+> We recommend using `--per_device_eval_batch_size=1` and `--max_target_length 128` at 4/8-bit evaluation.

 ### Predict

@@ -444,12 +438,6 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --predict_with_generate
 ```

-## TODO
-
- [ ] Supporting flash attention ([torch](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html) / [xformers](https://github.com/facebookresearch/xformers) / [flashattn](https://github.com/Dao-AILab/flash-attention)).
- [ ] Implementing multi-query attention for faster inference.
- [ ] Supporting full-parameter RLHF training.
-
 ## License

 This repository is licensed under the [Apache-2.0 License](LICENSE).
@@ -461,6 +449,7 @@ Please follow the model licenses to use the corresponding model weights:
 - [BLOOM](https://huggingface.co/spaces/bigscience/license)
 - [Falcon](LICENSE)
 - [Baichuan](https://huggingface.co/baichuan-inc/baichuan-7B/resolve/main/baichuan-7B%20%E6%A8%A1%E5%9E%8B%E8%AE%B8%E5%8F%AF%E5%8D%8F%E8%AE%AE.pdf)
+- [Baichuan2](https://huggingface.co/baichuan-inc/Baichuan2-7B-Base/resolve/main/Baichuan%202%E6%A8%A1%E5%9E%8B%E7%A4%BE%E5%8C%BA%E8%AE%B8%E5%8F%AF%E5%8D%8F%E8%AE%AE.pdf)
 - [InternLM](https://github.com/InternLM/InternLM#open-source-license)
 - [Qwen](https://huggingface.co/Qwen/Qwen-7B-Chat/blob/main/LICENSE)
 - [XVERSE](https://github.com/xverse-ai/XVERSE-13B/blob/main/MODEL_LICENSE.pdf)
@@ -481,7 +470,7 @@ If this work is helpful, please kindly cite as:

 ## Acknowledgement

-This repo is a sibling of [ChatGLM-Efficient-Tuning](https://github.com/hiyouga/ChatGLM-Efficient-Tuning). They share a similar code structure of efficient tuning on large language models.
+This repo benefits from [PEFT](https://github.com/huggingface/peft), [QLoRA](https://github.com/artidoro/qlora) and [OpenChatKit](https://github.com/togethercomputer/OpenChatKit). Thanks for their wonderful works.

 ## Star History

--- a/README_zh.md
+++ b/README_zh.md
@@ -12,57 +12,48 @@

 ## 更新日志

+[23/09/10] 现在我们支持了 LLaMA 模型的 **[FlashAttention](https://github.com/Dao-AILab/flash-attention)**。如果您使用的是 RTX4090、A100 或 H100 GPU，请使用 `--flash_attn` 参数以启用 FlashAttention-2（实验性功能）。
+
 [23/08/18] 现在我们支持了**训练状态恢复**，请将 `transformers` 升级至 `4.31.0` 以启用此功能。

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

-[23/08/11] 现在我们支持了指令模型的 **[DPO 训练](https://arxiv.org/abs/2305.18290)**。详情请参阅[此示例](#dpo-训练)（实验性功能）。
-
-[23/08/03] 现在我们支持了 **Qwen-7B** 模型的训练。请尝试使用 `--model_name_or_path Qwen/Qwen-7B-Chat` 和 `--lora_target c_attn` 参数。使用 Qwen-7B-Chat 模型时请添加 `--template chatml` 参数。
+[23/08/11] 现在我们支持了指令模型的 **[DPO 训练](https://arxiv.org/abs/2305.18290)**。详情请参阅[此示例](#dpo-训练)。

 [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)）。
-
-[23/07/19] 现在我们支持了 **LLaMA-2** 模型的训练。请尝试使用 `--model_name_or_path meta-llama/Llama-2-7b-hf` 参数。使用 LLaMA-2-chat 模型时请添加 `--template llama2` 参数。
+[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)）。

 [23/07/18] 我们开发了支持训练和测试的**浏览器一体化界面**。请尝试使用 `train_web.py` 在您的浏览器中微调模型。感谢 [@KanadeSiina](https://github.com/KanadeSiina) 和 [@codemayq](https://github.com/codemayq) 在该功能开发中付出的努力。

-[23/07/11] 现在我们支持了 **Baichuan-13B** 模型的训练。请尝试使用 `--model_name_or_path baichuan-inc/Baichuan-13B-Base` 和 `--lora_target W_pack` 参数。使用 Baichuan-13B-Chat 模型时请添加 `--template baichuan` 参数。
-
 [23/07/09] 我们开源了 **[FastEdit](https://github.com/hiyouga/FastEdit)** ⚡🩹，一个简单易用的、能迅速编辑大模型事实记忆的工具包。如果您感兴趣请关注我们的 [FastEdit](https://github.com/hiyouga/FastEdit) 项目。

-[23/07/07] 现在我们支持了 **InternLM-7B** 模型的训练。请尝试使用 `--model_name_or_path internlm/internlm-7b` 参数。使用 InternLM-chat 模型时请添加 `--template intern` 参数。
-
-[23/07/05] 现在我们支持了 **Falcon-7B/40B** 模型的训练。请尝试使用 `--model_name_or_path tiiuae/falcon-7b` 和 `--lora_target query_key_value` 参数。
-
-[23/06/29] 我们提供了一个**可复现的**指令模型微调示例，详细内容请查阅 [Hugging Face 项目](https://huggingface.co/hiyouga/baichuan-7b-sft)。
+[23/06/29] 我们提供了一个**可复现的**指令模型微调示例，详细内容请查阅 [Baichuan-7B-sft](https://huggingface.co/hiyouga/Baichuan-7B-sft)。

 [23/06/22] 我们对齐了[示例 API](src/api_demo.py) 与 [OpenAI API](https://platform.openai.com/docs/api-reference/chat) 的格式，您可以将微调模型接入**任意基于 ChatGPT 的应用**中。

-[23/06/15] 现在我们支持了 **Baichuan-7B** 模型的训练。请尝试使用 `--model_name_or_path baichuan-inc/Baichuan-7B` 和 `--lora_target W_pack` 参数。
-
 [23/06/03] 现在我们实现了 4 比特的 LoRA 训练（也称 **[QLoRA](https://github.com/artidoro/qlora)**）。请尝试使用 `--quantization_bit 4` 参数进行 4 比特量化微调。

-[23/05/31] 现在我们支持了 **BLOOM & BLOOMZ** 模型的训练。请尝试使用 `--model_name_or_path bigscience/bloomz-7b1-mt` 和 `--lora_target query_key_value` 参数。
-
 ## 模型

-| 模型名                                                   | 模型大小                     | 默认模块           | Template |
-| -------------------------------------------------------- | --------------------------- | ----------------- |----------|
-| [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   |
-| [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   | -        |
-| [Falcon](https://huggingface.co/tiiuae/falcon-7b)        | 7B/40B                      | query_key_value   | -        |
-| [Baichuan](https://github.com/baichuan-inc/baichuan-13B) | 7B/13B                      | W_pack            | baichuan |
-| [InternLM](https://github.com/InternLM/InternLM)         | 7B                          | q_proj,v_proj     | intern   |
-| [Qwen](https://github.com/QwenLM/Qwen-7B)                | 7B                          | c_attn            | chatml   |
-| [XVERSE](https://github.com/xverse-ai/XVERSE-13B)        | 13B                         | q_proj,v_proj     | -        |
-| [ChatGLM2](https://github.com/THUDM/ChatGLM2-6B)         | 6B                          | query_key_value   | chatglm2 |
+| 模型名                                                   | 模型大小                     | 默认模块           | Template  |
+| -------------------------------------------------------- | --------------------------- | ----------------- | --------- |
+| [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    |
+| [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   | -         |
+| [Falcon](https://huggingface.co/tiiuae/falcon-7b)        | 7B/40B                      | query_key_value   | -         |
+| [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 |
+| [InternLM](https://github.com/InternLM/InternLM)         | 7B                          | q_proj,v_proj     | intern    |
+| [Qwen](https://github.com/QwenLM/Qwen-7B)                | 7B                          | c_attn            | chatml    |
+| [XVERSE](https://github.com/xverse-ai/XVERSE-13B)        | 13B                         | q_proj,v_proj     | xverse    |
+| [ChatGLM2](https://github.com/THUDM/ChatGLM2-6B)         | 6B                          | query_key_value   | chatglm2  |

- **默认模块**是 `--lora_target` 参数的部分可选项。请使用 `python src/train_bash.py -h` 查看全部可选项。
- 对于所有“基座”（Base）模型，`--template` 参数可以是 `default`, `alpaca`, `vicuna` 等任意值。但“对话”（Chat）模型请务必使用对应的模板。
+> [!NOTE]
+> **默认模块**应作为 `--lora_target` 参数的默认值，可使用 `--lora_target all` 参数指定全部模块。
+>
+> 对于所有“基座”（Base）模型，`--template` 参数可以是 `default`, `alpaca`, `vicuna` 等任意值。但“对话”（Chat）模型请务必使用对应的模板。

 ## 训练方法

@@ -74,7 +65,8 @@
 | PPO 训练               |                    |                    | :white_check_mark: | :white_check_mark: |
 | DPO 训练               | :white_check_mark: |                    | :white_check_mark: | :white_check_mark: |

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

 ## 数据集

@@ -105,7 +97,8 @@
  - [Web QA (zh)](https://huggingface.co/datasets/suolyer/webqa)
  - [UltraChat (en)](https://github.com/thunlp/UltraChat)
  - [WebNovel (zh)](https://huggingface.co/datasets/zxbsmk/webnovel_cn)
- 用于奖励模型或 DPO 训练：
+  - [Ad Gen (zh)](https://huggingface.co/datasets/HasturOfficial/adgen)
+- 用于训练奖励模型或 DPO 训练：
  - [HH-RLHF (en)](https://huggingface.co/datasets/Anthropic/hh-rlhf)
  - [Open Assistant (multilingual)](https://huggingface.co/datasets/OpenAssistant/oasst1)
  - [GPT-4 Generated Data (en&zh)](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)
@@ -123,7 +116,7 @@ huggingface-cli login

 - Python 3.8+ 和 PyTorch 1.13.1+
 - 🤗Transformers, Datasets, Accelerate, PEFT 和 TRL
- sentencepiece 和 tiktoken
+- sentencepiece, protobuf 和 tiktoken
 - jieba, rouge-chinese 和 nltk (用于评估)
 - gradio 和 matplotlib (用于网页端交互)
 - uvicorn, fastapi 和 sse-starlette (用于 API)
@@ -136,7 +129,8 @@ huggingface-cli login

 关于数据集文件的格式，请参考 `data/example_dataset` 文件夹的内容。构建自定义数据集时，既可以使用单个 `.json` 文件，也可以使用一个[数据加载脚本](https://huggingface.co/docs/datasets/dataset_script)和多个文件。

-注意：使用自定义数据集时，请更新 `data/dataset_info.json` 文件，该文件的格式请参考 `data/README.md`。
+> [!NOTE]
+> 使用自定义数据集时，请更新 `data/dataset_info.json` 文件，该文件的格式请参考 `data/README.md`。

 ### 环境搭建（可跳过）

@@ -162,10 +156,14 @@ CUDA_VISIBLE_DEVICES=0 python src/train_web.py

 我们极力推荐新手使用浏览器一体化界面，因为它还可以**自动**生成运行所需的命令行脚本。

-目前网页 UI 仅支持**单卡训练**。
+> [!WARNING]
+> 目前网页 UI 仅支持**单卡训练**。

 ### 单 GPU 训练

+> [!IMPORTANT]
+> 如果您使用多张 GPU 训练模型，请移步[多 GPU 分布式训练](#多-gpu-分布式训练)部分。
+
 #### 预训练

 ```bash
@@ -174,7 +172,6 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --model_name_or_path path_to_llama_model \
    --do_train \
    --dataset wiki_demo \
-    --template default \
    --finetuning_type lora \
    --lora_target q_proj,v_proj \
    --output_dir path_to_pt_checkpoint \
@@ -298,19 +295,13 @@ accelerate config # 首先配置分布式环境
 accelerate launch src/train_bash.py # 参数同上
 ```

-<details><summary>使用 DeepSpeed ZeRO-2 进行全参数微调的 Accelerate 配置示例</summary>
+<details><summary>LoRA 训练的 Accelerate 配置示例</summary>

 ```yaml
 compute_environment: LOCAL_MACHINE
-deepspeed_config:
-  gradient_accumulation_steps: 4
-  gradient_clipping: 0.5
-  offload_optimizer_device: none
-  offload_param_device: none
-  zero3_init_flag: false
-  zero_stage: 2
-distributed_type: DEEPSPEED
+distributed_type: MULTI_GPU
 downcast_bf16: 'no'
+gpu_ids: all
 machine_rank: 0
 main_training_function: main
 mixed_precision: fp16
@@ -334,10 +325,11 @@ deepspeed --num_gpus 8 --master_port=9901 src/train_bash.py \
    ... # 参数同上
 ```

-<details><summary>使用 DeepSpeed ZeRO-2 进行全参数微调的 DeepSpeed 配置示例</summary>
+<details><summary>使用 DeepSpeed ZeRO-2 进行全参数训练的 DeepSpeed 配置示例</summary>

 ```json
 {
+  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
@@ -385,7 +377,8 @@ python src/api_demo.py \
    --checkpoint_dir path_to_checkpoint
 ```

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

 ### 命令行测试

@@ -424,7 +417,8 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --predict_with_generate
 ```

-我们建议在量化模型的评估中使用 `--per_device_eval_batch_size=1` 和 `--max_target_length 128`。
+> [!NOTE]
+> 我们建议在量化模型的评估中使用 `--per_device_eval_batch_size=1` 和 `--max_target_length 128`。

 ### 模型预测

@@ -443,12 +437,6 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --predict_with_generate
 ```

-## TODO
-
- [ ] 实现 flash attention ([torch](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html) / [xformers](https://github.com/facebookresearch/xformers) / [flashattn](https://github.com/Dao-AILab/flash-attention))。
- [ ] 在推理阶段使用 Multi-query attention 进行加速。
- [ ] 支持 RLHF 的全参数微调。
-
 ## 协议

 本仓库的代码依照 [Apache-2.0](LICENSE) 协议开源。
@@ -460,6 +448,7 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
 - [BLOOM](https://huggingface.co/spaces/bigscience/license)
 - [Falcon](LICENSE)
 - [Baichuan](https://huggingface.co/baichuan-inc/baichuan-7B/resolve/main/baichuan-7B%20%E6%A8%A1%E5%9E%8B%E8%AE%B8%E5%8F%AF%E5%8D%8F%E8%AE%AE.pdf)
+- [Baichuan2](https://huggingface.co/baichuan-inc/Baichuan2-7B-Base/resolve/main/Baichuan%202%E6%A8%A1%E5%9E%8B%E7%A4%BE%E5%8C%BA%E8%AE%B8%E5%8F%AF%E5%8D%8F%E8%AE%AE.pdf)
 - [InternLM](https://github.com/InternLM/InternLM#open-source-license)
 - [Qwen](https://huggingface.co/Qwen/Qwen-7B-Chat/blob/main/LICENSE)
 - [XVERSE](https://github.com/xverse-ai/XVERSE-13B/blob/main/MODEL_LICENSE.pdf)
@@ -480,7 +469,7 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \

 ## 致谢

-本项目是 [ChatGLM-Efficient-Tuning](https://github.com/hiyouga/ChatGLM-Efficient-Tuning) 的同类项目。采用了类似的代码结构和训练方法。
+本项目受益于 [PEFT](https://github.com/huggingface/peft)、[QLoRA](https://github.com/artidoro/qlora) 和 [OpenChatKit](https://github.com/togethercomputer/OpenChatKit)，感谢以上诸位作者的付出。

 ## Star History

--- a/data/README.md
+++ b/data/README.md
@@ -6,6 +6,7 @@ If you are using a custom dataset, please provide your dataset definition in the
  "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)",
  "file_sha1": "the SHA-1 hash value of the dataset file. (optional)",
+  "ranking": "whether the examples contains ranked responses or not. (default: false)",
  "columns": {
    "prompt": "the name of the column in the datasets containing the prompts. (default: instruction)",
    "query": "the name of the column in the datasets containing the queries. (default: input)",
@@ -16,3 +17,16 @@ If you are using a custom dataset, please provide your dataset definition in the
 ```

 where the `prompt` and `response` columns should contain non-empty values. The `query` column will be concatenated with the `prompt` column and used as input for the model. The `history` column should contain a list where each element is a string tuple representing a query-response pair.
+
+For datasets used in reward modeling or DPO training, the `response` column should be a string list, with the preferred answers appearing first, for example:
+
+```json
+{
+  "instruction": "Question",
+  "input": "",
+  "output": [
+    "Chosen answer",
+    "Rejected answer"
+  ]
+}
+```
--- a/data/README_zh.md
+++ b/data/README_zh.md
@@ -6,6 +6,7 @@
  "script_url": "包含数据加载脚本的本地文件夹名称（若指定，则忽略下列两个参数）",
  "file_name": "该目录下数据集文件的名称（若上述参数未指定，则此项必需）",
  "file_sha1": "数据集文件的SHA-1哈希值（可选）",
+  "ranking": "数据集是否包含排序后的回答（默认：false）",
  "columns": {
    "prompt": "数据集代表提示词的表头名称（默认：instruction）",
    "query": "数据集代表请求的表头名称（默认：input）",
@@ -16,3 +17,16 @@
 ```

 其中 `prompt` 和 `response` 列应当是非空的字符串。`query` 列的内容将会和 `prompt` 列拼接作为模型输入。`history` 列应当是一个列表，其中每个元素是一个字符串二元组，分别代表用户请求和模型答复。
+
+对于训练奖励模型或 DPO 训练的数据集，`response` 列应当是一个字符串列表，排在前面的代表更优的答案，例如：
+
+```json
+{
+  "instruction": "Question",
+  "input": "",
+  "output": [
+    "Chosen answer",
+    "Rejected answer"
+  ]
+}
+```
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,11 +1,12 @@
 torch>=1.13.1
-transformers>=4.29.1
+transformers>=4.30.0
 datasets>=2.12.0
 accelerate>=0.21.0
-peft>=0.4.0
-trl>=0.5.0
+peft==0.4.0
+trl>=0.7.1
 scipy
 sentencepiece
+protobuf
 tiktoken
 jieba
 rouge-chinese
--- a/src/llmtuner/init.py
+++ b/src/llmtuner/init.py
@@ -6,4 +6,4 @@ from llmtuner.tuner import export_model, run_exp
 from llmtuner.webui import create_ui, create_web_demo


-__version__ = "0.1.7"
+__version__ = "0.1.8"
--- a/src/llmtuner/chat/stream_chat.py
+++ b/src/llmtuner/chat/stream_chat.py
@@ -1,7 +1,7 @@
 import torch
 from typing import Any, Dict, Generator, List, Optional, Tuple
 from threading import Thread
-from transformers import TextIteratorStreamer
+from transformers import GenerationConfig, TextIteratorStreamer

 from llmtuner.extras.misc import dispatch_model, get_logits_processor
 from llmtuner.extras.template import get_template_and_fix_tokenizer
@@ -14,7 +14,6 @@ class ChatModel:
        model_args, data_args, finetuning_args, self.generating_args = get_infer_args(args)
        self.model, self.tokenizer = load_model_and_tokenizer(model_args, finetuning_args)
        self.model = dispatch_model(self.model)
-        self.model = self.model.eval() # enable evaluation mode
        self.template = get_template_and_fix_tokenizer(data_args.template, self.tokenizer)
        self.system_prompt = data_args.system_prompt

@@ -41,26 +40,30 @@ class ChatModel:
        max_length = input_kwargs.pop("max_length", None)
        max_new_tokens = input_kwargs.pop("max_new_tokens", None)

-        gen_kwargs = self.generating_args.to_dict()
-        gen_kwargs.update(dict(
-            input_ids=input_ids,
-            do_sample=do_sample if do_sample is not None else gen_kwargs["do_sample"],
-            temperature=temperature or gen_kwargs["temperature"],
-            top_p=top_p or gen_kwargs["top_p"],
-            top_k=top_k or gen_kwargs["top_k"],
-            repetition_penalty=repetition_penalty or gen_kwargs["repetition_penalty"],
-            eos_token_id=list(set([self.tokenizer.eos_token_id] + self.tokenizer.additional_special_tokens_ids)),
-            pad_token_id=self.tokenizer.pad_token_id,
-            logits_processor=get_logits_processor()
+        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"],
+            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 max_length:
-            gen_kwargs.pop("max_new_tokens", None)
-            gen_kwargs["max_length"] = max_length
+            generating_args.pop("max_new_tokens", None)
+            generating_args["max_length"] = max_length

        if max_new_tokens:
-            gen_kwargs.pop("max_length", None)
-            gen_kwargs["max_new_tokens"] = max_new_tokens
+            generating_args.pop("max_length", None)
+            generating_args["max_new_tokens"] = max_new_tokens
+
+        gen_kwargs = dict(
+            inputs=input_ids,
+            generation_config=GenerationConfig(**generating_args),
+            logits_processor=get_logits_processor()
+        )

        return gen_kwargs, prompt_length

--- a/src/llmtuner/dsets/preprocess.py
+++ b/src/llmtuner/dsets/preprocess.py
@@ -31,11 +31,15 @@ def preprocess_dataset(
            yield query, response, history, system

    def preprocess_pretrain_dataset(examples: Dict[str, List[Any]]) -> Dict[str, Any]:
-        # build grouped texts with format `X1 X2 X3 ...` (without <eos>)
-        if isinstance(getattr(tokenizer, "tokenizer", None), tiktoken.Encoding): # for tiktoken tokenizer (Qwen)
-            kwargs = dict(allowed_special="all")
+        # build grouped texts with format `X1 X2 X3 ...`
+        if isinstance(getattr(tokenizer, "tokenizer", None), tiktoken.Encoding):
+            kwargs = dict(allowed_special="all") # for tiktoken tokenizer (Qwen)
        else:
-            kwargs = dict(add_special_tokens=False)
+            kwargs = dict(add_special_tokens=True)
+
+        if hasattr(tokenizer, "add_bos_token") and hasattr(tokenizer, "add_eos_token"):
+            setattr(tokenizer, "add_bos_token", True) # for LLaMA tokenizer
+            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()}
@@ -59,7 +63,9 @@ def preprocess_dataset(
        for query, response, history, system in construct_example(examples):
            input_ids, labels = [], []

-            for source_ids, target_ids in template.encode_multiturn(tokenizer, query, response, history, system):
+            for turn_idx, (source_ids, target_ids) in enumerate(template.encode_multiturn(
+                tokenizer, query, response, history, system
+            )):
                if len(source_ids) > data_args.max_source_length:
                    source_ids = source_ids[:data_args.max_source_length]
                if len(target_ids) > data_args.max_target_length:
@@ -68,8 +74,17 @@ def preprocess_dataset(
                if len(input_ids) + len(source_ids) + len(target_ids) > max_length:
                    break

+                if 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 += [IGNORE_INDEX] * len(source_ids) + target_ids
+                labels += source_mask + target_ids
+
+            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))
@@ -89,6 +104,9 @@ def preprocess_dataset(
            if len(target_ids) > data_args.max_target_length:
                target_ids = target_ids[:data_args.max_target_length]

+            if template.efficient_eos:
+                target_ids += [tokenizer.eos_token_id]
+
            model_inputs["input_ids"].append(source_ids)
            model_inputs["attention_mask"].append([1] * len(source_ids))
            model_inputs["labels"].append(target_ids)
@@ -109,6 +127,10 @@ def preprocess_dataset(
            if len(rejected_ids) > data_args.max_target_length:
                rejected_ids = rejected_ids[:data_args.max_target_length]

+            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)
--- a/src/llmtuner/extras/callbacks.py
+++ b/src/llmtuner/extras/callbacks.py
@@ -5,7 +5,9 @@ from typing import TYPE_CHECKING
 from datetime import timedelta

 from transformers import TrainerCallback
-from transformers.trainer_utils import has_length
+from transformers.trainer_callback import TrainerControl, TrainerState
+from transformers.trainer_utils import has_length, PREFIX_CHECKPOINT_DIR
+from transformers.training_args import TrainingArguments

 from llmtuner.extras.constants import LOG_FILE_NAME
 from llmtuner.extras.logging import get_logger
@@ -17,6 +19,24 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)


+class SavePeftModelCallback(TrainerCallback):
+
+    def on_save(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+        r"""
+        Event called after a checkpoint save.
+        """
+        if args.should_save:
+            output_dir = os.path.join(args.output_dir, "{}-{}".format(PREFIX_CHECKPOINT_DIR, state.global_step))
+            getattr(kwargs.get("model"), "pretrained_model").save_pretrained(output_dir)
+
+    def on_train_end(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+        r"""
+        Event called at the end of training.
+        """
+        if args.should_save:
+            getattr(kwargs.get("model"), "pretrained_model").save_pretrained(args.output_dir)
+
+
 class LogCallback(TrainerCallback):

    def __init__(self, runner=None):
--- a/src/llmtuner/extras/constants.py
+++ b/src/llmtuner/extras/constants.py
@@ -2,21 +2,17 @@ IGNORE_INDEX = -100

 LOG_FILE_NAME = "trainer_log.jsonl"

-VALUE_HEAD_FILE_NAME = "value_head.bin"
-
-FINETUNING_ARGS_NAME = "finetuning_args.json"
-
 LAYERNORM_NAMES = ["norm", "ln_f", "ln_attn", "ln_mlp"]

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

-STAGES = [
-    "SFT",
-    "Reward Modeling",
-    "PPO",
-    "DPO",
-    "Pre-Training"
-]
+TRAINING_STAGES = {
+    "Supervised Fine-Tuning": "sft",
+    "Reward Modeling": "rm",
+    "PPO": "ppo",
+    "DPO": "dpo",
+    "Pre-Training": "pt"
+}

 SUPPORTED_MODELS = {
    "LLaMA-7B": "huggyllama/llama-7b",
@@ -46,11 +42,16 @@ SUPPORTED_MODELS = {
    "Baichuan-7B": "baichuan-inc/Baichuan-7B",
    "Baichuan-13B": "baichuan-inc/Baichuan-13B-Base",
    "Baichuan-13B-Chat": "baichuan-inc/Baichuan-13B-Chat",
+    "Baichuan2-7B": "baichuan-inc/Baichuan2-7B-Base",
+    "Baichuan2-13B": "baichuan-inc/Baichuan2-13B-Base",
+    "Baichuan2-7B-Chat": "baichuan-inc/Baichuan2-7B-Chat",
+    "Baichuan2-13B-Chat": "baichuan-inc/Baichuan2-13B-Chat",
    "InternLM-7B": "internlm/internlm-7b",
    "InternLM-7B-Chat": "internlm/internlm-chat-7b",
    "Qwen-7B": "Qwen/Qwen-7B",
    "Qwen-7B-Chat": "Qwen/Qwen-7B-Chat",
    "XVERSE-13B": "xverse/XVERSE-13B",
+    "XVERSE-13B-Chat": "xverse/XVERSE-13B-Chat",
    "ChatGLM2-6B-Chat": "THUDM/chatglm2-6b"
 }

@@ -62,6 +63,7 @@ DEFAULT_MODULE = {
    "BLOOMZ": "query_key_value",
    "Falcon": "query_key_value",
    "Baichuan": "W_pack",
+    "Baichuan2": "W_pack",
    "InternLM": "q_proj,v_proj",
    "Qwen": "c_attn",
    "XVERSE": "q_proj,v_proj",
@@ -72,7 +74,9 @@ DEFAULT_TEMPLATE = {
    "LLaMA2": "llama2",
    "ChineseLLaMA2": "llama2_zh",
    "Baichuan": "baichuan",
+    "Baichuan2": "baichuan2",
    "InternLM": "intern",
    "Qwen": "chatml",
+    "XVERSE": "xverse",
    "ChatGLM2": "chatglm2"
 }
--- a/src/llmtuner/extras/misc.py
+++ b/src/llmtuner/extras/misc.py
@@ -1,9 +1,8 @@
+import gc
 import torch
-from typing import TYPE_CHECKING, List, Optional, Tuple
+from typing import TYPE_CHECKING, Tuple
 from transformers import InfNanRemoveLogitsProcessor, LogitsProcessorList

-from llmtuner.extras.constants import LAYERNORM_NAMES
-
 if TYPE_CHECKING:
    from transformers.modeling_utils import PreTrainedModel

@@ -28,12 +27,6 @@ class AverageMeter:
        self.avg = self.sum / self.count


-def get_logits_processor() -> LogitsProcessorList:
-    logits_processor = LogitsProcessorList()
-    logits_processor.append(InfNanRemoveLogitsProcessor())
-    return logits_processor
-
-
 def count_parameters(model: torch.nn.Module) -> Tuple[int, int]:
    r"""
    Returns the number of trainable parameters and number of all parameters in the model.
@@ -56,48 +49,17 @@ def count_parameters(model: torch.nn.Module) -> Tuple[int, int]:
    return trainable_params, all_param


-# 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/c0209c35abbf88c63aa267800d98a8e212ed0a42/src/peft/utils/other.py#L35
-def prepare_model_for_training(
-    model: "PreTrainedModel",
-    finetuning_type: str,
-    output_layer_name: Optional[str] = "lm_head",
-    use_gradient_checkpointing: Optional[bool] = True,
-    layer_norm_names: Optional[List[str]] = LAYERNORM_NAMES
-) -> "PreTrainedModel":
-    for name, param in model.named_parameters():
-        if param.ndim == 1 and any(layer_norm_name in name for layer_norm_name in layer_norm_names):
-            param.data = param.data.to(torch.float32)
-
-    if use_gradient_checkpointing:
-        if hasattr(model, "enable_input_require_grads"):
-            model.enable_input_require_grads()
-        else:
-            def make_inputs_require_grad(module, input, output):
-                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
-
-    if finetuning_type != "full" and hasattr(model, output_layer_name):
-        output_layer: torch.nn.Linear = getattr(model, output_layer_name)
-        input_dtype = output_layer.weight.dtype
-
-        class CastOutputToFloat(torch.nn.Sequential):
-
-            def forward(self, x: torch.Tensor) -> torch.Tensor:
-                return super().forward(x.to(input_dtype)).to(torch.float32)
-
-        setattr(model, output_layer_name, CastOutputToFloat(output_layer))
-
-    return model
+def get_logits_processor() -> LogitsProcessorList:
+    logits_processor = LogitsProcessorList()
+    logits_processor.append(InfNanRemoveLogitsProcessor())
+    return logits_processor


 def torch_gc() -> None:
    r"""
    Collects GPU memory.
    """
+    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()
--- a/src/llmtuner/extras/models/init.py
+++ b/src/llmtuner/extras/models/init.py
--- a/src/llmtuner/extras/models/flash_llama.py
+++ b/src/llmtuner/extras/models/flash_llama.py
@@ -0,0 +1,305 @@
+# coding=utf-8
+# Modified from:
+# [1] https://huggingface.co/Birchlabs/flash_llama/blob/main/modeling_flash_llama.py
+# [2] https://github.com/lm-sys/FastChat/blob/main/fastchat/train/llama2_flash_attn_monkey_patch.py
+# [3] https://huggingface.co/togethercomputer/LLaMA-2-7B-32K/blob/main/modeling_flash_llama.py
+# [4] https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
+# With fix from Alex Birch: https://huggingface.co/togethercomputer/LLaMA-2-7B-32K/discussions/17
+
+import torch
+from typing import Optional, Tuple
+from transformers.utils import logging
+from transformers.models.llama.configuration_llama import LlamaConfig
+
+
+try:
+    from flash_attn.flash_attn_interface import (
+        flash_attn_kvpacked_func,
+        flash_attn_varlen_kvpacked_func,
+    )
+    from flash_attn.bert_padding import unpad_input, pad_input
+    flash_attn_v2_installed = True
+    print('>>>> Flash Attention installed')
+except ImportError:
+    flash_attn_v2_installed = False
+    raise ImportError('Please install Flash Attention: `pip install flash-attn --no-build-isolation`')
+
+try:
+    from flash_attn.layers.rotary import apply_rotary_emb_func
+    flash_rope_installed = True
+    print('>>>> Flash RoPE installed')
+except ImportError:
+    flash_rope_installed = False
+    raise ImportError('Please install RoPE kernels: `pip install git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/rotary`')
+
+
+logger = logging.get_logger(__name__)
+
+
+class LlamaRMSNorm(torch.nn.Module):
+
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return (self.weight * hidden_states).to(input_dtype) # for fp32 weight
+
+
+class FlashRotaryEmbedding(torch.nn.Module):
+
+    def __init__(
+        self,
+        dim: int,
+        base=10000.0,
+        interleaved=False,
+        scale_base=None,
+        scaling_factor=1.0,
+        pos_idx_in_fp32=True,
+        device=None
+    ):
+        super().__init__()
+        self.dim = dim
+        self.base = float(base)
+        self.pos_idx_in_fp32 = pos_idx_in_fp32
+        # Generate and save the inverse frequency buffer (non trainable)
+        inv_freq = self._compute_inv_freq(device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.interleaved = interleaved
+        self.scale_base = scale_base
+        self.scaling_factor = scaling_factor
+        scale = (
+            (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim)
+            if scale_base is not None else None
+        )
+        self.register_buffer("scale", scale)
+
+        self._seq_len_cached = 0
+        self._cos_cached = None
+        self._sin_cached = None
+        self._cos_k_cached = None
+        self._sin_k_cached = None
+
+    def _compute_inv_freq(self, device=None):
+        return 1 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))
+
+    def _update_cos_sin_cache(self, seqlen, device=None, dtype=None):
+        if (
+            seqlen > self._seq_len_cached or self._cos_cached.device != device
+            or self._cos_cached.dtype != dtype
+            or (self.training and self._cos_cached.is_inference())
+        ):
+            self._seq_len_cached = seqlen
+            if self.pos_idx_in_fp32:
+                t = torch.arange(seqlen, device=device, dtype=torch.float32)
+                t /= self.scaling_factor
+                if self.inv_freq.dtype != torch.float32:
+                    inv_freq = self.inv_freq.to(torch.float32)
+                else:
+                    inv_freq = self.inv_freq
+            else:
+                t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
+                t /= self.scaling_factor
+                inv_freq = self.inv_freq
+            freqs = torch.outer(t, inv_freq)
+            if self.scale is None:
+                self._cos_cached = torch.cos(freqs).to(dtype)
+                self._sin_cached = torch.sin(freqs).to(dtype)
+            else:
+                power = (
+                    (torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - seqlen // 2) / self.scale_base
+                )
+                scale = self.scale.to(device=power.device) ** power.unsqueeze(-1)
+                # We want the multiplication by scale to happen in fp32
+                self._cos_cached = (torch.cos(freqs) * scale).to(dtype)
+                self._sin_cached = (torch.sin(freqs) * scale).to(dtype)
+                self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)
+                self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor, seqlen_offset: int = 0) -> Tuple[torch.Tensor, torch.Tensor]:
+        r"""
+        q: (batch, seqlen, nheads, headdim)
+        k: (batch, seqlen, nheads, headdim)
+        seqlen_offset: can be used in generation where the qkv being passed in is only the last
+        token in the batch.
+        """
+        self._update_cos_sin_cache(q.shape[1] + seqlen_offset, device=q.device, dtype=q.dtype)
+        if self.scale is None:
+            return apply_rotary_emb_func(
+                q, self._cos_cached[seqlen_offset:], self._sin_cached[seqlen_offset:],
+                self.interleaved, True # inplace=True
+            ), apply_rotary_emb_func(
+                k, self._cos_cached[seqlen_offset:], self._sin_cached[seqlen_offset:],
+                self.interleaved, True # inplace=True
+            )
+        else:
+            assert False
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    r"""
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, slen, _, num_key_value_heads, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, :, :, None, :].expand(batch, slen, 2, num_key_value_heads, n_rep, head_dim)
+    return hidden_states.reshape(batch, slen, 2, num_key_value_heads * n_rep, head_dim)
+
+
+class LlamaAttention(torch.nn.Module):
+
+    def __init__(self, config: "LlamaConfig"):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+
+        if (self.head_dim * self.num_heads) != self.hidden_size:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {self.num_heads})."
+            )
+
+        self.q_proj = torch.nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = torch.nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = torch.nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = torch.nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+
+        self.register_buffer(
+            "norm_factor",
+            torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)).to(torch.get_default_dtype()),
+            persistent=False,
+        )
+
+        if self.config.rope_scaling is None:
+            scaling_factor = 1
+        else:
+            scaling_type = self.config.rope_scaling["type"]
+            scaling_factor = self.config.rope_scaling["factor"]
+            assert scaling_type == "linear"
+
+        self.rotary_emb = FlashRotaryEmbedding(
+            self.head_dim, base=10000, interleaved=False, scaling_factor=scaling_factor
+        )
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+    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
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, h_size = hidden_states.size()
+
+        has_layer_past = past_key_value is not None
+
+        if has_layer_past:
+            past_kv = past_key_value[0]
+            past_len = past_key_value[1]
+        else:
+            past_len = 0
+
+        q = self.q_proj(hidden_states)
+        k = self.k_proj(hidden_states)
+        v = self.v_proj(hidden_states)
+
+        q = q.view(bsz, q_len, self.num_heads, self.head_dim)
+        k = k.view(bsz, q_len, self.num_key_value_heads, self.head_dim)
+        v = v.view(bsz, q_len, self.num_key_value_heads, self.head_dim)
+
+        q, k = self.rotary_emb(q, k, past_len)
+
+        kv = torch.stack([k, v], 2)
+        kv = repeat_kv(kv, self.num_key_value_groups)
+
+        # Cache QKV values
+        if has_layer_past:
+            new_len = past_len+q.size(1)
+            if new_len > past_kv.size(1):
+                past_kv = torch.cat(
+                    [past_kv, torch.empty(bsz, 256, 2, kv.size(3), kv.size(4), dtype=kv.dtype, device=kv.device)], 1
+                )
+            past_kv[:, past_len:new_len] = kv
+            kv = past_kv[:, :new_len]
+        else:
+            past_kv = kv
+
+        past_key_value = (past_kv, past_len + q.size(1)) if use_cache else None
+
+        if attention_mask is not None:
+            # varlen, ignore padding tokens, efficient for large batch with many paddings
+            logger.warning_once("padded sequences is less efficient")
+
+            unpadded_kv, indices_k, cu_seqlens_k, max_seqlen_k = unpad_input(kv, attention_mask)
+            unpadded_q, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(q, attention_mask[:, -q.size(1):])
+            attn_outputs = flash_attn_varlen_kvpacked_func(
+                unpadded_q, unpadded_kv, cu_seqlens_q, cu_seqlens_k,
+                max_seqlen_q, max_seqlen_k,
+                dropout_p=0.0, softmax_scale=1.0/self.norm_factor,
+                causal=(not has_layer_past), return_attn_probs=output_attentions
+            )
+
+            attn_output = attn_outputs[0] if output_attentions else attn_outputs
+            attn_output = pad_input(
+                attn_output, indices_q, bsz, q_len
+            ).reshape(bsz, q_len, h_size)
+            attn_weights = attn_outputs[2] if output_attentions else None
+
+        else:
+            # no padding tokens, more efficient
+            attn_outputs = flash_attn_kvpacked_func(
+                q, kv, dropout_p=0.0, softmax_scale=1.0/self.norm_factor,
+                causal=(not has_layer_past), return_attn_probs=output_attentions
+            )
+            attn_output = attn_outputs[0] if output_attentions else attn_outputs
+            attn_output = attn_output.reshape(bsz, q_len, h_size)
+            attn_weights = attn_outputs[2] if output_attentions else None
+
+        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 key_padding_mask. Fills in the past kv length for use in forward.
+def _prepare_decoder_attention_mask(
+    self, attention_mask, input_shape, inputs_embeds, past_key_values_length
+):
+    # [bsz, seq_len]
+    if past_key_values_length > 0 and attention_mask is not None:
+        attention_mask = torch.cat(
+            (
+                torch.full(
+                    (input_shape[0], past_key_values_length),
+                    True,
+                    dtype=attention_mask.dtype,
+                    device=attention_mask.device
+                ),
+                attention_mask
+            ),
+            dim=-1
+        )
+
+    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
--- a/src/llmtuner/extras/save_and_load.py
+++ b/src/llmtuner/extras/save_and_load.py
@@ -1,49 +1,21 @@
 import os
 import torch
-from typing import Dict
+from transformers.trainer import WEIGHTS_NAME

-from transformers.trainer import WEIGHTS_NAME, WEIGHTS_INDEX_NAME
-from transformers.modeling_utils import load_sharded_checkpoint
-
-from llmtuner.extras.constants import VALUE_HEAD_FILE_NAME
 from llmtuner.extras.logging import get_logger


 logger = get_logger(__name__)


-def get_state_dict(model: torch.nn.Module) -> Dict[str, torch.Tensor]:
-    state_dict: Dict[str, torch.Tensor] = model.state_dict()
-    filtered_state_dict = {}
-
-    for k, v in model.named_parameters():
-        if v.requires_grad:
-            filtered_state_dict[k] = state_dict[k].cpu().clone().detach()
-
-    return filtered_state_dict
-
-
-def load_trainable_params(model: torch.nn.Module, checkpoint_dir: os.PathLike) -> bool:
-    weights_file = os.path.join(checkpoint_dir, WEIGHTS_NAME)
-    if os.path.exists(weights_file):
-        model_state_dict = torch.load(weights_file, map_location="cpu")
-        model.load_state_dict(model_state_dict, strict=False) # skip missing keys
-    elif os.path.exists(os.path.join(checkpoint_dir, WEIGHTS_INDEX_NAME)):
-        load_sharded_checkpoint(model, checkpoint_dir, strict=False)
-    else:
-        logger.warning("Provided path ({}) does not contain pre-trained weights.".format(checkpoint_dir))
-        return False
-    return True
-
-
 def load_valuehead_params(model: torch.nn.Module, checkpoint_dir: os.PathLike) -> bool:
-    valuehead_file = os.path.join(checkpoint_dir, VALUE_HEAD_FILE_NAME)
-    if not os.path.exists(valuehead_file):
+    vhead_file = os.path.join(checkpoint_dir, WEIGHTS_NAME)
+    if not os.path.exists(vhead_file):
        logger.warning("Provided path ({}) does not contain valuehead weights.".format(checkpoint_dir))
        return False
-    valuehead_state_dict = torch.load(valuehead_file, map_location="cpu")
-    model.register_buffer("reward_head_weight", valuehead_state_dict["summary.weight"])
-    model.register_buffer("reward_head_bias", valuehead_state_dict["summary.bias"])
-    model.register_buffer("default_head_weight", torch.zeros_like(valuehead_state_dict["summary.weight"]))
-    model.register_buffer("default_head_bias", torch.zeros_like(valuehead_state_dict["summary.bias"]))
+    vhead_params = torch.load(vhead_file, map_location="cpu")
+    model.register_buffer("reward_head_weight", vhead_params["v_head.summary.weight"], persistent=False)
+    model.register_buffer("reward_head_bias", vhead_params["v_head.summary.bias"], persistent=False)
+    model.register_buffer("default_head_weight", torch.zeros_like(vhead_params["v_head.summary.weight"]), persistent=False)
+    model.register_buffer("default_head_bias", torch.zeros_like(vhead_params["v_head.summary.bias"]), persistent=False)
    return True
--- a/src/llmtuner/extras/template.py
+++ b/src/llmtuner/extras/template.py
@@ -20,6 +20,7 @@ class Template:
    sep: List[Union[str, Dict[str, str]]]
    stop_words: List[str]
    use_history: bool
+    efficient_eos: bool

    def encode_oneturn(
        self,
@@ -74,19 +75,19 @@ class Template:
        self,
        tokenizer: "PreTrainedTokenizer"
    ) -> Tuple[List[int], List[int]]:
-        if (
-            tokenizer.bos_token_id is not None
-            and getattr(tokenizer, "add_bos_token", True)
-        ): # baichuan-13b has no bos token
+        if tokenizer.bos_token_id is not None and getattr(tokenizer, "add_bos_token", True):
            bos_ids = [tokenizer.bos_token_id]
-        else:
-            bos_ids = [] # bos token is optional
+        else: # baichuan, qwen and gpt2 models have no bos token
+            bos_ids = []

-        if tokenizer.eos_token_id is not None:
-            eos_ids = [tokenizer.eos_token_id]
-        else:
+        if tokenizer.eos_token_id is None:
            raise ValueError("EOS token is required.")

+        if self.efficient_eos: # used in baichuan, qwen, chatglm, etc.
+            eos_ids = []
+        else:
+            eos_ids = [tokenizer.eos_token_id]
+
        return bos_ids, eos_ids

    def _encode(
@@ -137,6 +138,8 @@ class Template:
        token_ids = []
        for elem in context:
            if isinstance(elem, str):
+                if len(elem) == 0:
+                    continue
                elem = elem.replace("{{system}}", system, 1) if system is not None else elem
                elem = elem.replace("{{query}}", query, 1) if query is not None else elem
                elem = elem.replace("{{idx}}", idx, 1) if idx is not None else elem
@@ -184,7 +187,8 @@ def register_template(
    system: str,
    sep: List[Union[str, Dict[str, str]]],
    stop_words: Optional[List[str]] = [],
-    use_history: Optional[bool] = True
+    use_history: Optional[bool] = True,
+    efficient_eos: Optional[bool] = False
 ) -> None:
    template_class = Llama2Template if "llama2" in name else Template
    templates[name] = template_class(
@@ -193,7 +197,8 @@ def register_template(
        system=system,
        sep=sep,
        stop_words=stop_words,
-        use_history=use_history
+        use_history=use_history,
+        efficient_eos=efficient_eos
    )


@@ -201,30 +206,23 @@ def get_template_and_fix_tokenizer(
    name: str,
    tokenizer: "PreTrainedTokenizer"
 ) -> Template:
-    template = templates.get(name, None)
-    assert template is not None, "Template {} does not exist.".format(name)
-
-    additional_special_tokens = template.stop_words
-    if len(template.stop_words): # inplace method
-        if tokenizer.eos_token_id is not None:
-            additional_special_tokens.append(tokenizer.eos_token)
-
-        tokenizer.eos_token = additional_special_tokens[0] # use the first stop word as eos token
-        additional_special_tokens.pop(0)
-        logger.info("Replace eos token: {}".format(tokenizer.eos_token))
-
    if tokenizer.eos_token_id is None:
        tokenizer.eos_token = "<|endoftext|>"
        logger.info("Add eos token: {}".format(tokenizer.eos_token))

    if tokenizer.pad_token_id is None:
-        if tokenizer.unk_token_id is not None:
-            tokenizer.pad_token = tokenizer.unk_token
-        else:
-            tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token = tokenizer.eos_token
        logger.info("Add pad token: {}".format(tokenizer.pad_token))

-    tokenizer.add_special_tokens(dict(additional_special_tokens=additional_special_tokens))
+    if name is None:
+        return None
+
+    template = templates.get(name, None)
+    assert template is not None, "Template {} does not exist.".format(name)
+    tokenizer.add_special_tokens(
+        dict(additional_special_tokens=template.stop_words),
+        replace_additional_special_tokens=False
+    )
    return template


@@ -282,7 +280,7 @@ register_template(
        "Always answer as helpfully as possible, while being safe.  "
        "Your answers should not include any harmful, unethical, "
        "racist, sexist, toxic, dangerous, or illegal content. "
-        "Please ensure that your responses are socially unbiased and positive in nature.\n"
+        "Please ensure that your responses are socially unbiased and positive in nature.\n\n"
        "If a question does not make any sense, or is not factually coherent, "
        "explain why instead of answering something not correct. "
        "If you don't know the answer to a question, please don't share false information."
@@ -461,12 +459,13 @@ register_template(
    ],
    system="",
    sep=[
+        {"token": "<eoa>"},
        "\n"
    ],
    stop_words=[
-        "</s>", # internlm cannot replace eos token
        "<eoa>"
-    ]
+    ],
+    efficient_eos=True
 )


@@ -476,18 +475,36 @@ Supports: https://huggingface.co/baichuan-inc/Baichuan-13B-Chat
 register_template(
    name="baichuan",
    prefix=[
-        "{{system}}",
-        {"token": "<reserved_102>"} # user token
+        "{{system}}"
    ],
    prompt=[
+        {"token": "<reserved_102>"}, # user token
        "{{query}}",
-        {"token": "<reserved_103>"} # assistant token
+        {"token": "<reserved_103>"}  # assistant token
    ],
    system="",
    sep=[],
-    stop_words=[
-        "<reserved_102>" # user token
-    ]
+    efficient_eos=True
+)
+
+
+r"""
+Supports: https://huggingface.co/baichuan-inc/Baichuan2-7B-Chat
+          https://huggingface.co/baichuan-inc/Baichuan2-13B-Chat
+"""
+register_template(
+    name="baichuan2",
+    prefix=[
+        "{{system}}"
+    ],
+    prompt=[
+        {"token": "<reserved_106>"}, # user token
+        "{{query}}",
+        {"token": "<reserved_107>"}  # assistant token
+    ],
+    system="",
+    sep=[],
+    efficient_eos=True
 )


@@ -500,7 +517,6 @@ register_template(
    prefix=[
        {"token": "<|system|>"},
        "\n{{system}}",
-        {"token": "<|end|>"}
    ],
    prompt=[
        {"token": "<|user|>"},
@@ -511,11 +527,13 @@ register_template(
    ],
    system="",
    sep=[
+        {"token": "<|end|>"},
        "\n"
    ],
    stop_words=[
        "<|end|>"
-    ]
+    ],
+    efficient_eos=True
 )


@@ -526,8 +544,7 @@ register_template(
    name="chatml",
    prefix=[
        {"token": "<|im_start|>"},
-        "system\n{{system}}",
-        {"token": "<|im_end|>"}
+        "system\n{{system}}"
    ],
    prompt=[
        {"token": "<|im_start|>"},
@@ -539,11 +556,13 @@ register_template(
    ],
    system="You are a helpful assistant.",
    sep=[
+        {"token": "<|im_end|>"},
        "\n"
    ],
    stop_words=[
        "<|im_end|>"
-    ]
+    ],
+    efficient_eos=True
 )


@@ -563,5 +582,22 @@ register_template(
    system="",
    sep=[
        "\n\n"
-    ]
+    ],
+    efficient_eos=True
+)
+
+
+r"""
+Supports: https://huggingface.co/xverse/XVERSE-13B-Chat
+"""
+register_template(
+    name="xverse",
+    prefix=[
+        "{{system}}"
+    ],
+    prompt=[
+        "Human: {{query}}\n\nAssistant: "
+    ],
+    system="",
+    sep=[]
 )
--- a/src/llmtuner/hparams/data_args.py
+++ b/src/llmtuner/hparams/data_args.py
@@ -11,23 +11,23 @@ class DatasetAttr:
    dataset_name: Optional[str] = None
    dataset_sha1: Optional[str] = None
    system_prompt: Optional[str] = None
+    ranking: Optional[bool] = False
+    prompt: Optional[str] = "instruction"
+    query: Optional[str] = "input"
+    response: Optional[str] = "output"
+    history: Optional[str] = None

    def __repr__(self) -> str:
        return self.dataset_name

-    def __post_init__(self):
-        self.prompt = "instruction"
-        self.query = "input"
-        self.response = "output"
-        self.history = None
-

@dataclass
 class DataArguments:
    r"""
    Arguments pertaining to what data we are going to input our model for training and evaluation.
    """
-    template: str = field(
+    template: Optional[str] = field(
+        default=None,
        metadata={"help": "Which template to use for constructing prompts in training and inference."}
    )
    dataset: Optional[str] = field(
@@ -47,7 +47,7 @@ class DataArguments:
        metadata={"help": "Enable streaming mode."}
    )
    buffer_size: Optional[int] = field(
-        default=16384,
+        default=1024,
        metadata={"help": "Size of the buffer to randomly sample examples from in streaming mode."}
    )
    mix_strategy: Optional[Literal["concat", "interleave_under", "interleave_over"]] = field(
@@ -129,5 +129,6 @@ class DataArguments:
                dataset_attr.response = dataset_info[name]["columns"].get("response", None)
                dataset_attr.history = dataset_info[name]["columns"].get("history", None)

+            dataset_attr.ranking = dataset_info[name].get("ranking", False)
            dataset_attr.system_prompt = prompt_list[i]
            self.dataset_list.append(dataset_attr)
--- a/src/llmtuner/hparams/finetuning_args.py
+++ b/src/llmtuner/hparams/finetuning_args.py
@@ -21,7 +21,8 @@ class FinetuningArguments:
                  Falcon choices: [\"32\", \"60\"], \
                  Baichuan choices: [\"32\", \"40\"] \
                  Qwen choices: [\"32\"], \
-                  XVERSE choices: [\"40\"]"}
+                  XVERSE choices: [\"40\"], \
+                  ChatGLM2 choices: [\"28\"]"}
    )
    num_layer_trainable: Optional[int] = field(
        default=3,
@@ -31,7 +32,7 @@ class FinetuningArguments:
        default="mlp",
        metadata={"help": "Name of trainable modules for partial-parameter (freeze) fine-tuning. \
                  LLaMA choices: [\"mlp\", \"self_attn\"], \
-                  BLOOM & Falcon choices: [\"mlp\", \"self_attention\"], \
+                  BLOOM & Falcon & ChatGLM2 choices: [\"mlp\", \"self_attention\"], \
                  Baichuan choices: [\"mlp\", \"self_attn\"], \
                  Qwen choices: [\"mlp\", \"attn\"], \
                  LLaMA-2, InternLM, XVERSE choices: the same as LLaMA."}
@@ -52,7 +53,7 @@ class FinetuningArguments:
        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 choices: [\"query_key_value\", \"self_attention.dense\", \"mlp.dense\"], \
+                  BLOOM & Falcon & ChatGLM2 choices: [\"query_key_value\", \"self_attention.dense\", \"mlp.dense\"], \
                  Baichuan choices: [\"W_pack\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"], \
                  Qwen choices: [\"c_attn\", \"attn.c_proj\", \"w1\", \"w2\", \"mlp.c_proj\"], \
                  LLaMA-2, InternLM, XVERSE choices: the same as LLaMA."}
--- a/src/llmtuner/hparams/model_args.py
+++ b/src/llmtuner/hparams/model_args.py
@@ -16,7 +16,7 @@ class ModelArguments:
        metadata={"help": "Where to store the pretrained models downloaded from huggingface.co."}
    )
    use_fast_tokenizer: Optional[bool] = field(
-        default=False,
+        default=True,
        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."}
    )
    use_auth_token: Optional[bool] = field(
@@ -27,10 +27,6 @@ class ModelArguments:
        default="main",
        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."}
    )
-    padding_side: Optional[Literal["left", "right"]] = field(
-        default="left",
-        metadata={"help": "The side on which the model should have padding applied."}
-    )
    quantization_bit: Optional[int] = field(
        default=None,
        metadata={"help": "The number of bits to quantize the model."}
@@ -47,6 +43,10 @@ class ModelArguments:
        default=None,
        metadata={"help": "Adopt scaled rotary positional embeddings."}
    )
+    flash_attn: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Enable flash attention for faster training."}
+    )
    checkpoint_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Path to the directory(s) containing the delta model checkpoints as well as the configurations."}
--- a/src/llmtuner/tuner/core/adapter.py
+++ b/src/llmtuner/tuner/core/adapter.py
@@ -11,7 +11,7 @@ from peft import (
 from peft.utils import CONFIG_NAME, WEIGHTS_NAME

 from llmtuner.extras.logging import get_logger
-from llmtuner.extras.save_and_load import load_trainable_params
+from llmtuner.tuner.core.utils import find_all_linear_modules

 if TYPE_CHECKING:
    from transformers.modeling_utils import PreTrainedModel
@@ -52,9 +52,6 @@ def init_adapter(
            else:
                param.data = param.data.to(torch.float32)

-        if model_args.checkpoint_dir is not None:
-            assert load_trainable_params(model, model_args.checkpoint_dir[0]), "Model checkpoint is not correctly loaded."
-
    if finetuning_args.finetuning_type == "lora":
        logger.info("Fine-tuning method: LoRA")
        latest_checkpoint = None
@@ -81,13 +78,18 @@ def init_adapter(
                model = PeftModel.from_pretrained(model, latest_checkpoint, is_trainable=is_trainable)

        if is_trainable and latest_checkpoint is None: # create new lora weights while training
+            if len(finetuning_args.lora_target) == 1 and finetuning_args.lora_target[0] == "all":
+                target_modules = find_all_linear_modules(model, model_args.quantization_bit)
+            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=finetuning_args.lora_target
+                target_modules=target_modules
            )
            model = get_peft_model(model, lora_config)

--- a/src/llmtuner/tuner/core/loader.py
+++ b/src/llmtuner/tuner/core/loader.py
@@ -4,6 +4,7 @@ import torch
 from types import MethodType
 from typing import TYPE_CHECKING, Literal, Optional, Tuple

+import transformers
 from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
@@ -15,14 +16,19 @@ from transformers import (
 )
 from transformers.utils import check_min_version
 from transformers.utils.versions import require_version
-from transformers.deepspeed import is_deepspeed_zero3_enabled
 from trl import AutoModelForCausalLMWithValueHead

+try:
+    from transformers.deepspeed import is_deepspeed_zero3_enabled
+except ImportError:
+    from transformers.integrations import is_deepspeed_zero3_enabled
+
 from llmtuner.extras.logging import reset_logging, get_logger
-from llmtuner.extras.misc import count_parameters, prepare_model_for_training
+from llmtuner.extras.misc import count_parameters
 from llmtuner.extras.save_and_load import load_valuehead_params
 from llmtuner.hparams import FinetuningArguments
 from llmtuner.tuner.core.adapter import init_adapter
+from llmtuner.tuner.core.utils import prepare_model_for_training

 if TYPE_CHECKING:
    from transformers import PreTrainedTokenizer
@@ -32,11 +38,11 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)


-check_min_version("4.29.1")
+check_min_version("4.30.0")
 require_version("datasets>=2.12.0", "To fix: pip install datasets>=2.12.0")
 require_version("accelerate>=0.21.0", "To fix: pip install accelerate>=0.21.0")
-require_version("peft>=0.4.0", "To fix: pip install peft>=0.4.0")
-require_version("trl>=0.5.0", "To fix: pip install trl>=0.5.0")
+require_version("peft==0.4.0", "To fix: pip install peft==0.4.0")
+require_version("trl>=0.7.1", "To fix: pip install trl>=0.7.1")


 def load_model_and_tokenizer(
@@ -64,18 +70,23 @@ def load_model_and_tokenizer(
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.model_name_or_path,
        use_fast=model_args.use_fast_tokenizer,
-        padding_side=model_args.padding_side,
+        padding_side="right", # training with left-padded tensors in fp16 precision may cause overflow
        **config_kwargs
    )

-    if finetuning_args.finetuning_type == "full" and model_args.checkpoint_dir is not None:
+    # Fix tokenizer (for ChatGLM2)
+    if "PreTrainedTokenizerBase" not in str(tokenizer._pad.__func__):
+        tokenizer._pad = MethodType(PreTrainedTokenizerBase._pad, tokenizer)
+
+    if finetuning_args.finetuning_type != "lora" and model_args.checkpoint_dir is not None:
        model_to_load = model_args.checkpoint_dir[0]
    else:
        model_to_load = model_args.model_name_or_path

    config = AutoConfig.from_pretrained(model_to_load, **config_kwargs)

-    if hasattr(config, "fp16") and hasattr(config, "bf16"): # fix Qwen config
+    # Fix config (for Qwen)
+    if is_trainable and hasattr(config, "fp16") and hasattr(config, "bf16"):
        if model_args.compute_dtype == torch.bfloat16:
            setattr(config, "bf16", True)
        else:
@@ -91,11 +102,12 @@ def load_model_and_tokenizer(
                setattr(config, "use_logn_attn", True)
                logger.info("Using dynamic NTK scaling.")

-        elif hasattr(config, "rope_scaling"): # for LLaMA models
+        elif hasattr(config, "rope_scaling"): # for LLaMA and Falcon models
            require_version("transformers>=4.31.0", "RoPE scaling requires transformers>=4.31.0")

            if is_trainable:
                if model_args.rope_scaling == "dynamic":
+                    assert not model_args.flash_attn, "Flash attention does not support dynamic rope scaling."
                    logger.warning(
                        "Dynamic NTK may not work well with fine-tuning. "
                        "See: https://github.com/huggingface/transformers/pull/24653"
@@ -118,6 +130,18 @@ def load_model_and_tokenizer(
        else:
            logger.warning("Current model does not support RoPE scaling.")

+    # Set flash attention
+    if model_args.flash_attn and getattr(config, "model_type", None) == "llama":
+        import transformers.models.llama.modeling_llama as LlamaModule
+        from llmtuner.extras.models.flash_llama import LlamaRMSNorm, LlamaAttention, _prepare_decoder_attention_mask
+        LlamaModule.LlamaRMSNorm = LlamaRMSNorm
+        LlamaModule.LlamaAttention = LlamaAttention
+        LlamaModule.LlamaModel._prepare_decoder_attention_mask = _prepare_decoder_attention_mask
+        if not hasattr(config, "num_key_value_heads"):
+            setattr(config, "num_key_value_heads", getattr(config, "num_attention_heads"))
+        if getattr(config, "pretraining_tp", 1) != 1:
+            setattr(config, "pretraining_tp", 1)
+
    # Quantization configurations (using bitsandbytes library).
    is_mergeable = True
    if model_args.quantization_bit is not None:
@@ -171,10 +195,12 @@ def load_model_and_tokenizer(
    # Initialize adapters
    model = prepare_model_for_training(model, finetuning_args.finetuning_type) if is_trainable else model
    model = init_adapter(model, model_args, finetuning_args, is_trainable, is_mergeable)
+    model = model.train() if is_trainable else model.eval()

    # Prepare model with valuehead for RLHF
    if stage == "rm" or stage == "ppo":
        model: AutoModelForCausalLMWithValueHead = AutoModelForCausalLMWithValueHead.from_pretrained(model)
+        model._keys_to_ignore_on_save = None
        reset_logging()
        if stage == "rm" and model_args.checkpoint_dir is not None: # load valuehead weights to evaluate reward model
            logger.warning("Only the last checkpoint containing valuehead will be loaded as the valuehead.")
--- a/src/llmtuner/tuner/core/parser.py
+++ b/src/llmtuner/tuner/core/parser.py
@@ -5,6 +5,7 @@ import datasets
 import transformers
 from typing import Any, Dict, Optional, Tuple
 from transformers import HfArgumentParser, Seq2SeqTrainingArguments
+from transformers.utils.versions import require_version
 from transformers.trainer_utils import get_last_checkpoint

 from llmtuner.extras.logging import get_logger
@@ -95,6 +96,9 @@ def get_train_args(
    # Check arguments (do not check finetuning_args since it may be loaded from checkpoints)
    data_args.init_for_training()

+    if general_args.stage != "pt" and data_args.template is None:
+        raise ValueError("Please specify which `template` to use.")
+
    if general_args.stage != "sft" and training_args.predict_with_generate:
        raise ValueError("`predict_with_generate` cannot be set as True except SFT.")

@@ -110,9 +114,17 @@ def get_train_args(
    if general_args.stage in ["ppo", "dpo"] and not training_args.do_train:
        raise ValueError("PPO and DPO stages can only be performed at training.")

+    if general_args.stage in ["rm", "dpo"]:
+        for dataset_attr in data_args.dataset_list:
+            if not dataset_attr.ranking:
+                raise ValueError("Please use ranked datasets for reward modeling or DPO training.")
+
    if general_args.stage == "ppo" and model_args.reward_model is None:
        raise ValueError("Reward model is necessary for PPO training.")

+    if general_args.stage == "ppo" and training_args.deepspeed is not None:
+        raise ValueError("PPO training is incompatible with DeepSpeed, use Accelerate instead.")
+
    if general_args.stage == "ppo" and data_args.streaming:
        raise ValueError("Streaming mode does not suppport PPO training currently.")

@@ -156,10 +168,9 @@ def get_train_args(
        and finetuning_args.finetuning_type == "lora"
    ):
        logger.warning("`ddp_find_unused_parameters` needs to be set as False for LoRA in DDP training.")
-        training_args.ddp_find_unused_parameters = False
-
-    if training_args.optim == "adamw_hf":
-        training_args.optim = "adamw_torch" # suppress warning
+        training_args_dict = training_args.to_dict()
+        training_args_dict.update(dict(ddp_find_unused_parameters=False))
+        training_args = Seq2SeqTrainingArguments(**training_args_dict)

    if (
        training_args.resume_from_checkpoint is None
@@ -167,12 +178,15 @@ def get_train_args(
        and os.path.isdir(training_args.output_dir)
        and not training_args.overwrite_output_dir
    ):
+        require_version("transformers>=4.31.0", "Resuming training requires transformers>=4.31.0.")
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError("Output directory already exists and is not empty. Use `overwrite_output_dir`.")

        if last_checkpoint is not None:
-            training_args.resume_from_checkpoint = last_checkpoint
+            training_args_dict = training_args.to_dict()
+            training_args_dict.update(dict(resume_from_checkpoint=last_checkpoint))
+            training_args = Seq2SeqTrainingArguments(**training_args_dict)
            logger.info(
                "Resuming from checkpoint. Change `output_dir` or use `overwrite_output_dir` to avoid."
            )
@@ -210,6 +224,9 @@ def get_infer_args(
 ]:
    model_args, data_args, finetuning_args, generating_args = parse_infer_args(args)

+    if data_args.template is None:
+        raise ValueError("Please specify which `template` to use.")
+
    if model_args.quantization_bit is not None and finetuning_args.finetuning_type != "lora":
        raise ValueError("Quantization is only compatible with the LoRA method.")

--- a/src/llmtuner/tuner/core/trainer.py
+++ b/src/llmtuner/tuner/core/trainer.py
@@ -1,107 +0,0 @@
-import os
-import torch
-from typing import TYPE_CHECKING, Dict, Optional
-
-from transformers import Seq2SeqTrainer
-from transformers.trainer import TRAINING_ARGS_NAME, WEIGHTS_NAME
-from transformers.modeling_utils import PreTrainedModel, unwrap_model
-from peft import PeftModel
-from trl import PreTrainedModelWrapper
-
-from llmtuner.extras.constants import FINETUNING_ARGS_NAME, VALUE_HEAD_FILE_NAME
-from llmtuner.extras.logging import get_logger
-from llmtuner.extras.save_and_load import get_state_dict, load_trainable_params
-
-if TYPE_CHECKING:
-    from transformers import PreTrainedTokenizer, Seq2SeqTrainingArguments, TrainerState
-    from llmtuner.hparams import FinetuningArguments
-
-
-logger = get_logger(__name__)
-
-
-class PeftModelMixin:
-    r"""
-    Patches the save and load methods in Hugging Face Trainer for PeftModel and ModelWithValueHead.
-    """
-
-    def __init__(self) -> None: # for type checking
-        self.model: PreTrainedModel = None
-        self.tokenizer: "PreTrainedTokenizer" = None
-        self.args: "Seq2SeqTrainingArguments" = None
-        self.finetuning_args: "FinetuningArguments" = None
-        self.state: "TrainerState" = None
-        raise AssertionError("Mixin should not be initialized.")
-
-    def _save(self, output_dir: Optional[str] = None, state_dict: Optional[Dict[str, torch.Tensor]] = None) -> None:
-        r"""
-        Saves trainable parameters as model checkpoint.
-
-        This function will only be executed at the process zero.
-
-        Subclass and override to inject custom behavior. It should not be directly used by external scripts.
-        """
-        output_dir = output_dir if output_dir is not None else self.args.output_dir
-        os.makedirs(output_dir, exist_ok=True)
-        logger.info(f"Saving model checkpoint to {output_dir}")
-
-        model = unwrap_model(self.model)
-        if isinstance(model, PreTrainedModelWrapper):
-            # Custom state dict: https://github.com/lvwerra/trl/blob/v0.4.7/trl/models/modeling_value_head.py#L200
-            model_state_dict = state_dict or model.state_dict()
-            v_head_state_dict = {
-                name.replace("v_head.", ""): model_state_dict[name].cpu().clone().detach()
-                for name in model_state_dict.keys() if name.startswith("v_head.")
-            }
-
-            torch.save(v_head_state_dict, os.path.join(output_dir, VALUE_HEAD_FILE_NAME))
-            model = model.pretrained_model
-
-        state_dict = state_dict or get_state_dict(model)
-        if isinstance(model, (PeftModel, PreTrainedModel)):
-            model.config.use_cache = True
-            model.save_pretrained(output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors)
-            model.config.use_cache = False
-        else:
-            torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
-
-        if self.finetuning_args.finetuning_type == "full" and self.tokenizer is not None:
-            try:
-                self.tokenizer.save_pretrained(output_dir)
-            except:
-                logger.warning("Cannot save tokenizer, copy the files manually.")
-
-        with open(os.path.join(output_dir, TRAINING_ARGS_NAME), "w", encoding="utf-8") as f:
-            f.write(self.args.to_json_string() + "\n")
-
-        self.finetuning_args.save_to_json(os.path.join(output_dir, FINETUNING_ARGS_NAME))
-
-    def _load_best_model(self):
-        r"""
-        Loads trainable parameters from model checkpoint.
-
-        Subclass and override to inject custom behavior. It should not be directly used by external scripts.
-        """
-        logger.info(f"Loading best model from {self.state.best_model_checkpoint} (score: {self.state.best_metric}).")
-        model = unwrap_model(self.model)
-
-        if isinstance(model, PreTrainedModelWrapper):
-            model.v_head.load_state_dict(torch.load(
-                os.path.join(self.state.best_model_checkpoint, VALUE_HEAD_FILE_NAME), map_location="cpu"
-            ))
-            model = model.pretrained_model
-
-        if isinstance(model, PeftModel):
-            model.load_adapter(self.state.best_model_checkpoint, model.active_adapter)
-        else: # freeze/full-tuning
-            load_trainable_params(model, self.state.best_model_checkpoint)
-
-
-class PeftTrainer(PeftModelMixin, Seq2SeqTrainer):
-    r"""
-    Inherits Seq2SeqTrainer to support parameter-efficient checkpoints.
-    """
-
-    def __init__(self, finetuning_args: "FinetuningArguments", **kwargs):
-        Seq2SeqTrainer.__init__(self, **kwargs)
-        self.finetuning_args = finetuning_args
--- a/src/llmtuner/tuner/core/utils.py
+++ b/src/llmtuner/tuner/core/utils.py
@@ -0,0 +1,72 @@
+import torch
+from typing import TYPE_CHECKING, List, Optional
+
+from llmtuner.extras.constants import LAYERNORM_NAMES
+
+if TYPE_CHECKING:
+    from transformers.modeling_utils import PreTrainedModel
+
+
+def find_all_linear_modules(
+    model: "PreTrainedModel",
+    quantization_bit: Optional[int] = None,
+    output_layer_name: Optional[str] = "lm_head"
+) -> List[str]:
+    if quantization_bit is not None:
+        import bitsandbytes as bnb
+        linear_cls = bnb.nn.Linear4bit if quantization_bit == 4 else bnb.nn.Linear8bitLt
+    else:
+        linear_cls = torch.nn.Linear
+
+    module_names = set()
+    for name, module in model.named_modules():
+        if output_layer_name not in name and isinstance(module, linear_cls):
+            module_names.add(name.split(".")[-1])
+
+    if output_layer_name in module_names:
+        module_names.pop(output_layer_name)
+
+    return list(module_names)
+
+
+def prepare_model_for_training(
+    model: "PreTrainedModel",
+    finetuning_type: str,
+    output_layer_name: Optional[str] = "lm_head",
+    use_gradient_checkpointing: Optional[bool] = True,
+    layer_norm_names: Optional[List[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
+    """
+    for name, param in model.named_parameters():
+        if param.ndim == 1 and any(layer_norm_name in name for layer_norm_name in layer_norm_names):
+            param.data = param.data.to(torch.float32)
+
+    if use_gradient_checkpointing:
+        if hasattr(model, "enable_input_require_grads"):
+            model.enable_input_require_grads()
+        else:
+            def make_inputs_require_grad(module, input, output):
+                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
+
+    if finetuning_type != "full" and hasattr(model, output_layer_name):
+        output_layer: torch.nn.Linear = getattr(model, output_layer_name)
+        input_dtype = output_layer.weight.dtype
+
+        class CastOutputToFloat(torch.nn.Sequential):
+
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return super().forward(x.to(input_dtype)).to(torch.float32)
+
+        setattr(model, output_layer_name, CastOutputToFloat(output_layer))
+
+    return model
--- a/src/llmtuner/tuner/dpo/collator.py
+++ b/src/llmtuner/tuner/dpo/collator.py
@@ -16,7 +16,7 @@ class DPODataCollatorWithPadding(DataCollatorForSeq2Seq):
            if self.tokenizer.padding_side == "left":
                start, end = feature.size(0) - answer_len, feature.size(0)
            else:
-                start, end = prompt_len, answer_len
+                start, end = prompt_len, prompt_len + answer_len
            padded_tensor = self.label_pad_token_id * torch.ones_like(feature)
            padded_tensor[start:end] = feature[start:end]
            padded_labels.append(padded_tensor)
--- a/src/llmtuner/tuner/dpo/trainer.py
+++ b/src/llmtuner/tuner/dpo/trainer.py
@@ -1,42 +1,49 @@
 import torch
 from collections import defaultdict
-from peft import PeftModel
 from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
 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 llmtuner.tuner.core.trainer import PeftModelMixin

 if TYPE_CHECKING:
    from transformers import PreTrainedModel
-    from llmtuner.hparams import FinetuningArguments, GeneratingArguments


-class DPOPeftTrainer(PeftModelMixin, DPOTrainer):
+class CustomDPOTrainer(DPOTrainer):

    def __init__(
        self,
-        finetuning_args: "FinetuningArguments",
-        generating_args: "GeneratingArguments",
+        beta: float,
+        model: Union["PreTrainedModel", torch.nn.Module],
        ref_model: Optional[Union["PreTrainedModel", torch.nn.Module]] = None,
+        disable_dropout: Optional[bool] = True,
        **kwargs
    ):
-        self.finetuning_args = finetuning_args
-        self.generating_args = generating_args
+        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.label_pad_token_id = IGNORE_INDEX
        self.padding_value = 0
-        self.beta = finetuning_args.dpo_beta
+        self.beta = beta
        self._stored_metrics = defaultdict(lambda: defaultdict(list))

-        Trainer.__init__(self, **kwargs)
+        Trainer.__init__(self, model=model, **kwargs)
        if not hasattr(self, "accelerator"):
            raise AttributeError("Please update `transformers`.")

        if ref_model is not None:
-            self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
+            if self.is_deepspeed_enabled:
+                self.ref_model, = self.accelerator._prepare_deepspeed(self.ref_model)
+                self.ref_model.eval()
+            else:
+                self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)

    def concatenated_forward(
        self,
@@ -44,27 +51,12 @@ class DPOPeftTrainer(PeftModelMixin, DPOTrainer):
        batch: Optional[Dict[str, torch.Tensor]] = None
    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
        batch_copied = BatchEncoding({k: v.detach().clone() for k, v in batch.items()}) # avoid error
-        unwrapped_model: "PreTrainedModel" = self.accelerator.unwrap_model(self.model)

-        if not torch.is_grad_enabled():
-            unwrapped_model.gradient_checkpointing_disable()
-
-        if model is None and isinstance(unwrapped_model, PeftModel): # peft model has no ref_model
-            with unwrapped_model.disable_adapter():
-                all_logits = self.model(
-                    input_ids=batch_copied["input_ids"],
-                    attention_mask=batch_copied["attention_mask"],
-                    return_dict=True
-                ).logits.to(torch.float32)
-        else:
-            all_logits = model(
-                input_ids=batch_copied["input_ids"],
-                attention_mask=batch_copied["attention_mask"],
-                return_dict=True
-            ).logits.to(torch.float32)
-
-        if not torch.is_grad_enabled():
-            unwrapped_model.gradient_checkpointing_enable()
+        all_logits = model(
+            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_logits,
--- a/src/llmtuner/tuner/dpo/workflow.py
+++ b/src/llmtuner/tuner/dpo/workflow.py
@@ -3,17 +3,18 @@
 from copy import deepcopy
 from peft import PeftModel
 from typing import TYPE_CHECKING, Optional, List
+from transformers import Seq2SeqTrainingArguments

 from llmtuner.dsets import get_dataset, preprocess_dataset, split_dataset
 from llmtuner.extras.constants import IGNORE_INDEX
 from llmtuner.extras.ploting import plot_loss
 from llmtuner.tuner.core import load_model_and_tokenizer
 from llmtuner.tuner.dpo.collator import DPODataCollatorWithPadding
-from llmtuner.tuner.dpo.trainer import DPOPeftTrainer
+from llmtuner.tuner.dpo.trainer import CustomDPOTrainer

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


 def run_dpo(
@@ -21,7 +22,6 @@ def run_dpo(
    data_args: "DataArguments",
    training_args: "Seq2SeqTrainingArguments",
    finetuning_args: "FinetuningArguments",
-    generating_args: "GeneratingArguments",
    callbacks: Optional[List["TrainerCallback"]] = None
 ):
    dataset = get_dataset(model_args, data_args)
@@ -32,15 +32,15 @@ def run_dpo(
        label_pad_token_id=IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
    )

-    training_args.remove_unused_columns = False # important for pairwise dataset
-    ref_model = deepcopy(model) if not isinstance(model, PeftModel) else None
+    training_args_dict = training_args.to_dict()
+    training_args_dict.update(dict(remove_unused_columns=False)) # important for pairwise dataset
+    training_args = Seq2SeqTrainingArguments(**training_args_dict)

    # Initialize our Trainer
-    trainer = DPOPeftTrainer(
-        finetuning_args=finetuning_args,
-        generating_args=generating_args,
-        ref_model=ref_model,
+    trainer = CustomDPOTrainer(
+        beta=finetuning_args.dpo_beta,
        model=model,
+        ref_model=deepcopy(model) if not isinstance(model, PeftModel) else None,
        args=training_args,
        tokenizer=tokenizer,
        data_collator=data_collator,
--- a/src/llmtuner/tuner/ppo/trainer.py
+++ b/src/llmtuner/tuner/ppo/trainer.py
@@ -2,29 +2,28 @@ import os
 import math
 import torch
 from tqdm import tqdm
-from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple

-from transformers import TrainerState, TrainerControl
+from transformers import GenerationConfig, Trainer, TrainerState, TrainerControl
+from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR

 from trl import PPOTrainer
 from trl.core import LengthSampler, PPODecorators, logprobs_from_logits

 from llmtuner.extras.logging import get_logger
 from llmtuner.extras.misc import AverageMeter, count_parameters, get_logits_processor
-from llmtuner.tuner.core.trainer import PeftTrainer
-from llmtuner.tuner.ppo.utils import replace_model
+from llmtuner.tuner.ppo.utils import cast_layernorm_dtype, replace_model

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


 logger = get_logger(__name__)


-class PPOPeftTrainer(PPOTrainer, PeftTrainer):
+class CustomPPOTrainer(PPOTrainer, Trainer):
    r"""
    Inherits PPOTrainer.
    """
@@ -32,17 +31,18 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
    def __init__(
        self,
        training_args: "Seq2SeqTrainingArguments",
-        finetuning_args: "FinetuningArguments",
        generating_args: "GeneratingArguments",
-        callbacks: List["LogCallback"],
+        callbacks: List["TrainerCallback"],
        compute_dtype: torch.dtype,
        **kwargs
    ):
        PPOTrainer.__init__(self, **kwargs)
+        if getattr(self.accelerator.state, "deepspeed_plugin", None) is not None:
+            raise ValueError("PPOTrainer is incompatible with DeepSpeed.")
+
        self.args = training_args
-        self.finetuning_args = finetuning_args
        self.generating_args = generating_args
-        self.log_callback = callbacks[0]
+        self.log_callback, self.save_callback = callbacks[0], callbacks[1]
        self.compute_dtype = compute_dtype
        self.state = TrainerState()
        self.control = TrainerControl()
@@ -75,10 +75,11 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
            logger.info(f"  Number of trainable parameters = {count_parameters(self.model)[0]}")

        # Keyword arguments for `model.generate`
-        gen_kwargs = self.generating_args.to_dict()
-        gen_kwargs["eos_token_id"] = list(set([self.tokenizer.eos_token_id] + self.tokenizer.additional_special_tokens_ids))
-        gen_kwargs["pad_token_id"] = self.tokenizer.pad_token_id
-        gen_kwargs["logits_processor"] = get_logits_processor()
+        generating_args = self.generating_args.to_dict()
+        generating_args.update(dict(
+            eos_token_id=[self.tokenizer.eos_token_id] + self.tokenizer.additional_special_tokens_ids,
+            pad_token_id=self.tokenizer.pad_token_id
+        ))

        length_sampler = LengthSampler(max_target_length // 2, max_target_length)
        unwrapped_model: "AutoModelForCausalLMWithValueHead" = self.accelerator.unwrap_model(self.model)
@@ -96,17 +97,21 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
            # Cast to inference mode
            unwrapped_model.gradient_checkpointing_disable()
            unwrapped_model.config.use_cache = True
+            self.model.eval()

            # Get inputs
-            queries, responses = self.get_inputs(batch, length_sampler, **gen_kwargs)
+            queries, responses = self.get_inputs(batch, length_sampler, generating_args)
+            self.tokenizer.padding_side = "right" # change padding side
            rewards = self.get_rewards(queries, responses, unwrapped_model)

            # Cast to training mode
            unwrapped_model.gradient_checkpointing_enable()
            unwrapped_model.config.use_cache = False
+            self.model.train()

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

@@ -128,7 +133,12 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
                reward_meter.reset()

            if (step+1) % self.args.save_steps == 0: # save checkpoint
-                self.save_model(os.path.join(self.args.output_dir, f"checkpoint-{step+1}"))
+                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)
+                )

            if self.control.should_epoch_stop or self.control.should_training_stop:
                break
@@ -137,34 +147,46 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
                dataiter = iter(self.dataloader)
                steps_trained = 0

-        self.log_callback.on_train_end(self.args, self.state, self.control)
+        self.log_callback.on_train_end(
+            self.args, self.state, self.control, model=self.accelerator.unwrap_model(self.model)
+        )

    @torch.no_grad()
    def get_inputs(
        self,
        batch: Dict[str, torch.Tensor],
-        length_sampler: Optional[Callable] = None,
-        **generation_kwargs
+        length_sampler: Callable,
+        generating_args: Dict[str, Any]
    ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
        r"""
        Generates model's responses given queries.
        """
-        if length_sampler is not None:
-            generation_kwargs["max_new_tokens"] = length_sampler()
+        generating_args["max_new_tokens"] = length_sampler()
+        gen_kwargs = dict(
+            generation_config=GenerationConfig(**generating_args),
+            logits_processor=get_logits_processor(),
+            **batch
+        )

+        input_ids = batch["input_ids"]
+        self.model, layer_norm_params = cast_layernorm_dtype(self.model, self.compute_dtype)
        unwrapped_model: "AutoModelForCausalLMWithValueHead" = self.accelerator.unwrap_model(self.model)
-        response: torch.Tensor = unwrapped_model.generate(**batch, **generation_kwargs)
-
-        # Temporary hack to ensure the generation config is not initialized for each iteration of the evaluation loop
-        # Inspired by: https://github.com/huggingface/transformers/blob/v4.28.1/src/transformers/trainer_seq2seq.py#L273
-        if unwrapped_model.pretrained_model.generation_config._from_model_config:
-            unwrapped_model.pretrained_model.generation_config._from_model_config = False
+        response: torch.Tensor = unwrapped_model.generate(**gen_kwargs)
+        self.model, _ = cast_layernorm_dtype(self.model, self.compute_dtype, layer_norm_params)
+        query, response = input_ids.detach().cpu(), response[:, input_ids.size(-1):].detach().cpu()

        queries, responses = [], []
-        query, response = batch["input_ids"].detach().cpu(), response[:, batch["input_ids"].size(-1):].detach().cpu()
        for i in range(len(query)):
            query_length = (query[i] != self.tokenizer.pad_token_id).nonzero()[0]
-            response_length = (response[i] != self.tokenizer.pad_token_id).nonzero()[-1] + 1
+            response_index = (response[i] != self.tokenizer.pad_token_id).nonzero()
+
+            if len(response_index) == 0:
+                response_length = 1 # allow empty response
+            elif self.tokenizer.pad_token_id == self.tokenizer.eos_token_id:
+                response_length = response_index[-1] + 2 # save the EOS token
+            else:
+                response_length = response_index[-1] + 1
+
            queries.append(query[i, query_length:]) # remove padding from left
            responses.append(response[i, :response_length]) # remove padding from right

@@ -189,7 +211,11 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
        if values.size(0) != batch["input_ids"].size(0): # adapt to chatglm2
            values = torch.transpose(values, 0, 1)

-        rewards = [reward for reward in values[:, -1].float().detach().cpu()] # use fp32 type
+        rewards = []
+        for i in range(values.size(0)):
+            end_index = batch["attention_mask"][i].nonzero()[-1] # use the score on the EOS token
+            rewards.append(values[i, end_index].float().detach().cpu()) # use fp32 type
+
        replace_model(unwrapped_model, target="default")
        return rewards

@@ -200,7 +226,8 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
        queries: torch.Tensor,
        responses: torch.Tensor,
        model_inputs: dict,
-        return_logits: Optional[bool] = False
+        return_logits: Optional[bool] = False,
+        response_masks: Optional[torch.Tensor] = None
    ):
        r"""
        Calculates model outputs in multiple batches.
@@ -218,6 +245,8 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
            input_kwargs = {key: value[i * fbs : (i + 1) * fbs] for key, value in model_inputs.items()}
            query_batch = queries[i * fbs : (i + 1) * fbs]
            response_batch = responses[i * fbs : (i + 1) * fbs]
+            if response_masks is not None:
+                response_masks_batch = response_masks[i * fbs : (i + 1) * fbs]
            input_ids = input_kwargs["input_ids"]
            attention_mask = input_kwargs["attention_mask"]

@@ -233,12 +262,19 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):

            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]
                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:]
+
                masks[j, :start] = 0
                masks[j, end:] = 0
+                if response_masks is not None:
+                    masks[j, start:end] = masks[j, start:end] * response_masks_batch[j][start:end]

            if return_logits:
                all_logits.append(logits)
--- a/src/llmtuner/tuner/ppo/utils.py
+++ b/src/llmtuner/tuner/ppo/utils.py
@@ -1,4 +1,7 @@
-from typing import TYPE_CHECKING, Literal
+import torch
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple
+
+from llmtuner.extras.constants import LAYERNORM_NAMES

 if TYPE_CHECKING:
    from trl import AutoModelForCausalLMWithValueHead
@@ -15,3 +18,23 @@ def replace_model(model: "AutoModelForCausalLMWithValueHead", target: Literal["d
        "summary.weight": getattr(model, "{}_head_weight".format(target)),
        "summary.bias": getattr(model, "{}_head_bias".format(target))
    })
+
+
+def cast_layernorm_dtype(
+    model: "AutoModelForCausalLMWithValueHead",
+    compute_dtype: torch.dtype,
+    layer_norm_params: Optional[Dict[str, torch.Tensor]] = None,
+    layer_norm_names: Optional[List[str]] = LAYERNORM_NAMES
+) -> Tuple["AutoModelForCausalLMWithValueHead", Dict[str, torch.Tensor]]:
+
+    layer_norm_state_dict = {}
+
+    for name, param in model.named_parameters():
+        if param.ndim == 1 and any(layer_norm_name in name for layer_norm_name in layer_norm_names):
+            if layer_norm_params is None:
+                layer_norm_state_dict[name] = param.data.detach().clone() # store float32 weights for stability
+                param.data = param.data.to(compute_dtype)
+            else:
+                param.data = layer_norm_params[name] # restore float32 weights
+
+    return model, layer_norm_state_dict
--- a/src/llmtuner/tuner/ppo/workflow.py
+++ b/src/llmtuner/tuner/ppo/workflow.py
@@ -4,14 +4,14 @@ import math
 from trl import PPOConfig
 from torch.optim import AdamW
 from typing import TYPE_CHECKING, Optional, List
-from transformers import DataCollatorForSeq2Seq
+from transformers import DataCollatorWithPadding
 from transformers.optimization import get_scheduler
-from transformers.utils.versions import require_version

 from llmtuner.dsets import get_dataset, preprocess_dataset
+from llmtuner.extras.callbacks import SavePeftModelCallback
 from llmtuner.extras.ploting import plot_loss
 from llmtuner.tuner.core import load_model_and_tokenizer
-from llmtuner.tuner.ppo.trainer import PPOPeftTrainer
+from llmtuner.tuner.ppo.trainer import CustomPPOTrainer

 if TYPE_CHECKING:
    from transformers import Seq2SeqTrainingArguments, TrainerCallback
@@ -29,7 +29,9 @@ def run_ppo(
    dataset = get_dataset(model_args, data_args)
    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="ppo")
    dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="ppo")
-    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, label_pad_token_id=tokenizer.pad_token_id)
+
+    tokenizer.padding_side = "left" # use left-padding in generation while using right-padding in training
+    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

    ppo_config = PPOConfig(
        model_name=model_args.model_name_or_path,
@@ -44,7 +46,6 @@ def run_ppo(
    )

    if finetuning_args.ppo_score_norm:
-        require_version("trl>=0.5.1.dev0", "To fix: pip install git+https://github.com/huggingface/trl.git")
        ppo_config.use_score_scaling = True
        ppo_config.use_score_norm = True

@@ -61,11 +62,10 @@ def run_ppo(
    )

    # Initialize our Trainer
-    ppo_trainer = PPOPeftTrainer(
+    ppo_trainer = CustomPPOTrainer(
        training_args=training_args,
-        finetuning_args=finetuning_args,
        generating_args=generating_args,
-        callbacks=callbacks,
+        callbacks=callbacks + [SavePeftModelCallback()],
        compute_dtype=model_args.compute_dtype,
        config=ppo_config,
        model=model,
--- a/src/llmtuner/tuner/pt/workflow.py
+++ b/src/llmtuner/tuner/pt/workflow.py
@@ -2,12 +2,11 @@

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

 from llmtuner.dsets import get_dataset, preprocess_dataset, split_dataset
 from llmtuner.extras.ploting import plot_loss
 from llmtuner.tuner.core import load_model_and_tokenizer
-from llmtuner.tuner.core.trainer import PeftTrainer

 if TYPE_CHECKING:
    from transformers import Seq2SeqTrainingArguments, TrainerCallback
@@ -27,8 +26,7 @@ def run_pt(
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    # Initialize our Trainer
-    trainer = PeftTrainer(
-        finetuning_args=finetuning_args,
+    trainer = Trainer(
        model=model,
        args=training_args,
        tokenizer=tokenizer,
--- a/src/llmtuner/tuner/rm/trainer.py
+++ b/src/llmtuner/tuner/rm/trainer.py
@@ -2,9 +2,9 @@ import os
 import json
 import torch
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+from transformers import Trainer

 from llmtuner.extras.logging import get_logger
-from llmtuner.tuner.core.trainer import PeftTrainer

 if TYPE_CHECKING:
    from transformers.trainer import PredictionOutput
@@ -14,7 +14,7 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)


-class PairwisePeftTrainer(PeftTrainer):
+class PairwiseTrainer(Trainer):
    r"""
    Inherits PeftTrainer to compute pairwise loss.
    """
@@ -32,21 +32,54 @@ class PairwisePeftTrainer(PeftTrainer):
        r"""
        Computes pairwise loss. The first n examples are chosen and the last n examples are rejected.

-        We use score on the EOS token to represent reward of the whole sentence.
-
-        Subclass and override to inject custom behavior. It should not be directly used by external scripts.
-
-        Note that the first element will be removed from the output tuple.
+        Subclass and override to inject custom behavior.

+        Note that the first element will be removed from the output tuple. 
        See: https://github.com/huggingface/transformers/blob/v4.30.2/src/transformers/trainer.py#L3509
        """
-        batch_size = inputs["input_ids"].size(0) // 2
+        # Compute rewards
        _, _, values = model(**inputs, output_hidden_states=True, return_dict=True)
        if values.size(0) != inputs["input_ids"].size(0): # adapt to chatglm2
            values = torch.transpose(values, 0, 1)
-        r_accept, r_reject = values[:, -1].split(batch_size, dim=0)
-        loss = -torch.log(torch.sigmoid(r_accept - r_reject)).mean()
-        return (loss, [loss, r_accept, r_reject]) if return_outputs else loss
+
+        # Split the inputs and rewards into two parts, chosen and rejected
+        batch_size = inputs["input_ids"].size(0) // 2
+        chosen_input_ids, rejected_input_ids = inputs["input_ids"][:batch_size], inputs["input_ids"][batch_size:]
+        chosen_attn_mask, rejected_attn_mask = (
+            inputs["attention_mask"][:batch_size], inputs["attention_mask"][batch_size:]
+        )
+        chosen_rewards, rejected_rewards = values[:batch_size], values[batch_size:]
+        chosen_scores, rejected_scores = [], []
+
+        # Compute pairwise loss. Only backprop on the different tokens before padding
+        # Inspired by: https://github.com/CarperAI/trlx/blob/main/examples/summarize_rlhf/reward_model/reward_model.py
+        loss = 0
+        for i in range(batch_size):
+            chosen_length = chosen_attn_mask[i].nonzero()[-1] + 1
+            rejected_length = rejected_attn_mask[i].nonzero()[-1] + 1
+            check_divergence = (chosen_input_ids[i] != rejected_input_ids[i]).nonzero()
+
+            if len(check_divergence) == 0:
+                end_index = chosen_length
+                div_index = end_index - 1
+            else:
+                end_index = max(chosen_length, rejected_length)
+                div_index = check_divergence[0]
+
+            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 EOS 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
+        if return_outputs:
+            chosen_scores, rejected_scores = torch.stack(chosen_scores), torch.stack(rejected_scores)
+            return loss, [loss, chosen_scores, rejected_scores]
+
+        return loss

    def save_predictions(
        self,
@@ -63,10 +96,10 @@ class PairwisePeftTrainer(PeftTrainer):
        output_prediction_file = os.path.join(self.args.output_dir, "generated_predictions.jsonl")
        logger.info(f"Saving prediction results to {output_prediction_file}")

-        acc_scores, rej_scores = predict_results.predictions
+        chosen_scores, rejected_scores = predict_results.predictions

        with open(output_prediction_file, "w", encoding="utf-8") as writer:
            res: List[str] = []
-            for acc_score, rej_score in zip(acc_scores, rej_scores):
-                res.append(json.dumps({"accept": round(float(acc_score), 2), "reject": round(float(rej_score), 2)}))
+            for c_score, r_score in zip(chosen_scores, rejected_scores):
+                res.append(json.dumps({"chosen": round(float(c_score), 2), "rejected": round(float(r_score), 2)}))
            writer.write("\n".join(res))
--- a/src/llmtuner/tuner/rm/workflow.py
+++ b/src/llmtuner/tuner/rm/workflow.py
@@ -1,18 +1,19 @@
 # Inspired by:
-# https://github.com/lvwerra/trl/blob/main/examples/summarization/scripts/reward_summarization.py
 # 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 transformers import Seq2SeqTrainingArguments

 from llmtuner.dsets import get_dataset, preprocess_dataset, split_dataset
+from llmtuner.extras.callbacks import SavePeftModelCallback
 from llmtuner.extras.ploting import plot_loss
 from llmtuner.tuner.core import load_model_and_tokenizer
 from llmtuner.tuner.rm.metric import compute_accuracy
 from llmtuner.tuner.rm.collator import PairwiseDataCollatorWithPadding
-from llmtuner.tuner.rm.trainer import PairwisePeftTrainer
+from llmtuner.tuner.rm.trainer import PairwiseTrainer

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


@@ -28,16 +29,17 @@ def run_rm(
    dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="rm")
    data_collator = PairwiseDataCollatorWithPadding(tokenizer)

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

    # Initialize our Trainer
-    trainer = PairwisePeftTrainer(
-        finetuning_args=finetuning_args,
+    trainer = PairwiseTrainer(
        model=model,
        args=training_args,
        tokenizer=tokenizer,
        data_collator=data_collator,
-        callbacks=callbacks,
+        callbacks=callbacks + [SavePeftModelCallback()],
        compute_metrics=compute_accuracy,
        **split_dataset(dataset, data_args, training_args)
    )
--- a/src/llmtuner/tuner/sft/trainer.py
+++ b/src/llmtuner/tuner/sft/trainer.py
@@ -4,10 +4,10 @@ import torch
 import numpy as np
 import torch.nn as nn
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+from transformers import Seq2SeqTrainer

 from llmtuner.extras.constants import IGNORE_INDEX
 from llmtuner.extras.logging import get_logger
-from llmtuner.tuner.core.trainer import PeftTrainer

 if TYPE_CHECKING:
    from transformers.trainer import PredictionOutput
@@ -16,7 +16,7 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)


-class Seq2SeqPeftTrainer(PeftTrainer):
+class CustomSeq2SeqTrainer(Seq2SeqTrainer):
    r"""
    Inherits PeftTrainer to compute generative metrics such as BLEU and ROUGE.
    """
@@ -33,27 +33,29 @@ class Seq2SeqPeftTrainer(PeftTrainer):

        Subclass and override to inject custom behavior.
        """
-        prompt_len, label_len = inputs["input_ids"].size(-1), inputs["labels"].size(-1)
-        if prompt_len > label_len:
-            inputs["labels"] = self._pad_tensors_to_target_len(inputs["labels"], inputs["input_ids"])
-        if label_len > prompt_len:
-            inputs["input_ids"] = self._pad_tensors_to_target_len(inputs["input_ids"], inputs["labels"])
-            if "attention_mask" in inputs:
-                inputs["attention_mask"] = self._pad_tensors_to_target_len(
-                    inputs["attention_mask"], inputs["labels"], pad_token_id=0
-                )
-            if "position_ids" in inputs:
-                inputs["position_ids"] = self._pad_tensors_to_target_len(
-                    inputs["position_ids"], inputs["labels"], pad_token_id=0
-                )
+        if self.args.predict_with_generate:
+            assert self.tokenizer.padding_side == "left", "This method only accepts left-padded tensor."
+            assert self.tokenizer.pad_token_id is not None, "Pad token is required."
+            prompt_len, label_len = inputs["input_ids"].size(-1), inputs["labels"].size(-1)
+            if prompt_len > label_len:
+                inputs["labels"] = self._pad_tensors_to_target_len(inputs["labels"], inputs["input_ids"])
+            if label_len > prompt_len:
+                inputs["input_ids"] = self._pad_tensors_to_target_len(inputs["input_ids"], inputs["labels"])
+                if "attention_mask" in inputs:
+                    inputs["attention_mask"] = self._pad_tensors_to_target_len(
+                        inputs["attention_mask"], inputs["labels"], pad_token_id=0
+                    )
+                if "position_ids" in inputs:
+                    inputs["position_ids"] = self._pad_tensors_to_target_len(
+                        inputs["position_ids"], inputs["labels"], pad_token_id=0
+                    )

        loss, generated_tokens, labels = super().prediction_step(
            model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys
        )
-        if generated_tokens is not None:
-            generated_tokens[:, :max(prompt_len, label_len)] = (
-                self.tokenizer.pad_token_id * torch.ones_like(generated_tokens[:, :max(prompt_len, label_len)])
-            )
+        if generated_tokens is not None and self.args.predict_with_generate:
+            generated_tokens[:, :max(prompt_len, label_len)] = self.tokenizer.pad_token_id
+            generated_tokens = generated_tokens.contiguous()

        return loss, generated_tokens, labels

@@ -65,16 +67,8 @@ class Seq2SeqPeftTrainer(PeftTrainer):
    ) -> torch.Tensor:
        r"""
        Pads the tensor to the same length as the target tensor.
-
-        Should only be called when predict_with_generate=True.
        """
-        if pad_token_id is None:
-            if self.tokenizer is not None and hasattr(self.tokenizer, "pad_token_id"):
-                assert self.tokenizer.padding_side == "left", "This method only accepts left-padded tensor."
-                pad_token_id = self.tokenizer.pad_token_id
-            else:
-                raise ValueError("PAD token is required.")
-
+        pad_token_id = pad_token_id if pad_token_id is not None else self.tokenizer.pad_token_id
        padded_tensor = 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
--- a/src/llmtuner/tuner/sft/workflow.py
+++ b/src/llmtuner/tuner/sft/workflow.py
@@ -1,7 +1,7 @@
 # Inspired by: https://github.com/huggingface/transformers/blob/v4.29.2/examples/pytorch/summarization/run_summarization.py

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

 from llmtuner.dsets import get_dataset, preprocess_dataset, split_dataset
 from llmtuner.extras.constants import IGNORE_INDEX
@@ -9,10 +9,10 @@ from llmtuner.extras.misc import get_logits_processor
 from llmtuner.extras.ploting import plot_loss
 from llmtuner.tuner.core import load_model_and_tokenizer
 from llmtuner.tuner.sft.metric import ComputeMetrics
-from llmtuner.tuner.sft.trainer import Seq2SeqPeftTrainer
+from llmtuner.tuner.sft.trainer import CustomSeq2SeqTrainer

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


@@ -27,20 +27,25 @@ def run_sft(
    dataset = get_dataset(model_args, data_args)
    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="sft")
    dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="sft")
+
+    if training_args.predict_with_generate:
+        tokenizer.padding_side = "left" # use left-padding in generation
+
    data_collator = DataCollatorForSeq2Seq(
        tokenizer=tokenizer,
        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.generation_max_length = training_args.generation_max_length if \
-                training_args.generation_max_length is not None else data_args.max_target_length
-    training_args.generation_num_beams = data_args.eval_num_beams if \
-                data_args.eval_num_beams is not None else training_args.generation_num_beams
+    training_args_dict = training_args.to_dict()
+    training_args_dict.update(dict(
+        generation_max_length=training_args.generation_max_length or data_args.max_target_length,
+        generation_num_beams=data_args.eval_num_beams or training_args.generation_num_beams
+    ))
+    training_args = Seq2SeqTrainingArguments(**training_args_dict)

    # Initialize our Trainer
-    trainer = Seq2SeqPeftTrainer(
-        finetuning_args=finetuning_args,
+    trainer = CustomSeq2SeqTrainer(
        model=model,
        args=training_args,
        tokenizer=tokenizer,
@@ -52,7 +57,7 @@ def run_sft(

    # Keyword arguments for `model.generate`
    gen_kwargs = generating_args.to_dict()
-    gen_kwargs["eos_token_id"] = list(set([tokenizer.eos_token_id] + tokenizer.additional_special_tokens_ids))
+    gen_kwargs["eos_token_id"] = [tokenizer.eos_token_id] + tokenizer.additional_special_tokens_ids
    gen_kwargs["pad_token_id"] = tokenizer.pad_token_id
    gen_kwargs["logits_processor"] = get_logits_processor()

--- a/src/llmtuner/webui/chat.py
+++ b/src/llmtuner/webui/chat.py
@@ -53,7 +53,7 @@ class WebChatModel(ChatModel):
            model_name_or_path=model_name_or_path,
            checkpoint_dir=checkpoint_dir,
            finetuning_type=finetuning_type,
-            quantization_bit=int(quantization_bit) if quantization_bit != "None" else None,
+            quantization_bit=int(quantization_bit) if quantization_bit and quantization_bit != "None" else None,
            template=template,
            system_prompt=system_prompt
        )
--- a/src/llmtuner/webui/common.py
+++ b/src/llmtuner/webui/common.py
@@ -6,7 +6,7 @@ import gradio as gr
 from peft.utils import WEIGHTS_NAME as PEFT_WEIGHTS_NAME
 from transformers.trainer import WEIGHTS_NAME, WEIGHTS_INDEX_NAME

-from llmtuner.extras.constants import DEFAULT_TEMPLATE, SUPPORTED_MODELS
+from llmtuner.extras.constants import DEFAULT_TEMPLATE, SUPPORTED_MODELS, TRAINING_STAGES


 DEFAULT_CACHE_DIR = "cache"
@@ -16,8 +16,8 @@ USER_CONFIG = "user.config"
 DATA_CONFIG = "dataset_info.json"


-def get_save_dir(model_name: str) -> str:
-    return os.path.join(DEFAULT_SAVE_DIR, os.path.split(model_name)[-1])
+def get_save_dir(*args) -> os.PathLike:
+    return os.path.join(DEFAULT_SAVE_DIR, *args)


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


 def save_config(lang: str, model_name: str, model_path: str) -> None:
@@ -56,7 +56,7 @@ def get_template(model_name: str) -> str:

 def list_checkpoint(model_name: str, finetuning_type: str) -> Dict[str, Any]:
    checkpoints = []
-    save_dir = os.path.join(get_save_dir(model_name), finetuning_type)
+    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 (
@@ -78,6 +78,10 @@ def load_dataset_info(dataset_dir: str) -> Dict[str, Any]:
        return {}


-def list_dataset(dataset_dir: Optional[str] = None) -> Dict[str, Any]:
+def list_dataset(
+    dataset_dir: Optional[str] = None, training_stage: Optional[str] = list(TRAINING_STAGES.keys())[0]
+) -> Dict[str, Any]:
    dataset_info = load_dataset_info(dataset_dir if dataset_dir is not None else DEFAULT_DATA_DIR)
-    return gr.update(value=[], choices=list(dataset_info.keys()))
+    ranking = TRAINING_STAGES[training_stage] in ["rm", "dpo"]
+    datasets = [k for k, v in dataset_info.items() if v.get("ranking", False) == ranking]
+    return gr.update(value=[], choices=datasets)
--- a/src/llmtuner/webui/components/infer.py
+++ b/src/llmtuner/webui/components/infer.py
@@ -16,7 +16,7 @@ def create_infer_tab(top_elems: Dict[str, "Component"]) -> Dict[str, "Component"

    info_box = gr.Textbox(show_label=False, interactive=False)

-    chat_model = WebChatModel()
+    chat_model = WebChatModel(lazy_init=True)
    chat_box, chatbot, history, chat_elems = create_chat_box(chat_model)

    load_btn.click(
--- a/src/llmtuner/webui/components/train.py
+++ b/src/llmtuner/webui/components/train.py
@@ -3,7 +3,7 @@ from transformers.trainer_utils import SchedulerType

 import gradio as gr

-from llmtuner.extras.constants import STAGES
+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 can_preview, get_preview, gen_plot
@@ -15,14 +15,17 @@ if TYPE_CHECKING:

 def create_train_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dict[str, "Component"]:
    with gr.Row():
-        training_stage = gr.Dropdown(choices=STAGES, value=STAGES[0], scale=2)
+        training_stage = gr.Dropdown(
+            choices=list(TRAINING_STAGES.keys()), value=list(TRAINING_STAGES.keys())[0], scale=2
+        )
        dataset_dir = gr.Textbox(value=DEFAULT_DATA_DIR, scale=2)
        dataset = gr.Dropdown(multiselect=True, scale=4)
        data_preview_btn = gr.Button(interactive=False, scale=1)

    preview_box, preview_count, preview_samples, close_btn = create_preview_box()

-    dataset_dir.change(list_dataset, [dataset_dir], [dataset])
+    training_stage.change(list_dataset, [dataset_dir, training_stage], [dataset])
+    dataset_dir.change(list_dataset, [dataset_dir, training_stage], [dataset])
    dataset.change(can_preview, [dataset_dir, dataset], [data_preview_btn])
    data_preview_btn.click(
        get_preview,
@@ -53,7 +56,6 @@ def create_train_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dic
            save_steps = gr.Slider(value=100, minimum=10, maximum=5000, step=10)
            warmup_steps = gr.Slider(value=0, minimum=0, maximum=5000, step=1)
            compute_type = gr.Radio(choices=["fp16", "bf16"], value="fp16")
-            padding_side = gr.Radio(choices=["left", "right"], value="left")

    with gr.Accordion(label="LoRA config", open=False) as lora_tab:
        with gr.Row():
@@ -119,7 +121,6 @@ def create_train_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dic
        save_steps,
        warmup_steps,
        compute_type,
-        padding_side,
        lora_rank,
        lora_dropout,
        lora_target,
@@ -165,7 +166,6 @@ def create_train_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dic
        save_steps=save_steps,
        warmup_steps=warmup_steps,
        compute_type=compute_type,
-        padding_side=padding_side,
        lora_tab=lora_tab,
        lora_rank=lora_rank,
        lora_dropout=lora_dropout,
--- a/src/llmtuner/webui/interface.py
+++ b/src/llmtuner/webui/interface.py
@@ -59,7 +59,7 @@ def create_web_demo() -> gr.Blocks:
    chat_model = WebChatModel(lazy_init=False)

    with gr.Blocks(title="Web Demo", css=CSS) as demo:
-        lang = gr.Dropdown(choices=["en", "zh"])
+        lang = gr.Dropdown(choices=["en", "zh"], value="en")

        _, _, _, chat_elems = create_chat_box(chat_model, visible=True)

--- a/src/llmtuner/webui/locales.py
+++ b/src/llmtuner/webui/locales.py
@@ -287,16 +287,6 @@ LOCALES = {
            "info": "是否启用 FP16 或 BF16 混合精度训练。"
        }
    },
-    "padding_side": {
-        "en": {
-            "label": "Padding side",
-            "info": "The side on which the model should have padding applied."
-        },
-        "zh": {
-            "label": "填充位置",
-            "info": "使用左填充或右填充。"
-        }
-    },
    "lora_tab": {
        "en": {
            "label": "LoRA configurations"
--- a/src/llmtuner/webui/manager.py
+++ b/src/llmtuner/webui/manager.py
@@ -19,12 +19,15 @@ class Manager:
        }

        user_config = load_config()
-        if lang:
-            refresh_dict["lang"] = {"value": lang}
-        else:
-            refresh_dict["lang"] = {"value": user_config["lang"] if user_config["lang"] else "en"}
+        if not lang:
+            if user_config.get("lang", None):
+                lang = user_config["lang"]
+            else:
+                lang = "en"

-        if user_config["last_model"]:
+        refresh_dict["lang"] = {"value": lang}
+
+        if user_config.get("last_model", None):
            refresh_dict["model_name"] = {"value": user_config["last_model"]}
            refresh_dict["model_path"] = {"value": get_model_path(user_config["last_model"])}

--- a/src/llmtuner/webui/runner.py
+++ b/src/llmtuner/webui/runner.py
@@ -8,11 +8,11 @@ from transformers.trainer import TRAINING_ARGS_NAME
 from typing import Any, Dict, Generator, List, Tuple

 from llmtuner.extras.callbacks import LogCallback
-from llmtuner.extras.constants import DEFAULT_MODULE
+from llmtuner.extras.constants import DEFAULT_MODULE, TRAINING_STAGES
 from llmtuner.extras.logging import LoggerHandler
 from llmtuner.extras.misc import torch_gc
 from llmtuner.tuner import run_exp
-from llmtuner.webui.common import get_model_path, get_save_dir
+from llmtuner.webui.common import get_model_path, get_save_dir, load_config
 from llmtuner.webui.locales import ALERTS
 from llmtuner.webui.utils import gen_cmd, get_eval_results, update_process_bar

@@ -87,7 +87,6 @@ class Runner:
        save_steps: int,
        warmup_steps: int,
        compute_type: str,
-        padding_side: str,
        lora_rank: int,
        lora_dropout: float,
        lora_target: str,
@@ -98,21 +97,25 @@ class Runner:
    ) -> Tuple[str, str, List[str], str, Dict[str, Any]]:
        if checkpoints:
            checkpoint_dir = ",".join(
-                [os.path.join(get_save_dir(model_name), finetuning_type, ckpt) for ckpt in checkpoints]
+                [get_save_dir(model_name, finetuning_type, ckpt) for ckpt in checkpoints]
            )
        else:
            checkpoint_dir = None

-        output_dir = os.path.join(get_save_dir(model_name), finetuning_type, output_dir)
+        output_dir = get_save_dir(model_name, finetuning_type, output_dir)
+
+        user_config = load_config()
+        cache_dir = user_config.get("cache_dir", None)

        args = dict(
-            stage="sft",
+            stage=TRAINING_STAGES[training_stage],
            model_name_or_path=get_model_path(model_name),
            do_train=True,
-            overwrite_cache=True,
+            overwrite_cache=False,
+            cache_dir=cache_dir,
            checkpoint_dir=checkpoint_dir,
            finetuning_type=finetuning_type,
-            quantization_bit=int(quantization_bit) if quantization_bit != "None" else None,
+            quantization_bit=int(quantization_bit) if quantization_bit in ["8", "4"] else None,
            template=template,
            system_prompt=system_prompt,
            dataset_dir=dataset_dir,
@@ -129,30 +132,22 @@ class Runner:
            logging_steps=logging_steps,
            save_steps=save_steps,
            warmup_steps=warmup_steps,
-            padding_side=padding_side,
            lora_rank=lora_rank,
            lora_dropout=lora_dropout,
            lora_target=lora_target or DEFAULT_MODULE.get(model_name.split("-")[0], "q_proj,v_proj"),
-            resume_lora_training=resume_lora_training,
+            resume_lora_training=(
+                False if TRAINING_STAGES[training_stage] in ["rm", "ppo", "dpo"] else resume_lora_training
+            ),
            output_dir=output_dir
        )
        args[compute_type] = True

-        if training_stage == "Reward Modeling":
-            args["stage"] = "rm"
-            args["resume_lora_training"] = False
-        elif training_stage == "PPO":
-            args["stage"] = "ppo"
-            args["resume_lora_training"] = False
+        if args["stage"] == "ppo":
            args["reward_model"] = reward_model
-            args["padding_side"] = "left"
            val_size = 0
-        elif training_stage == "DPO":
-            args["stage"] = "dpo"
-            args["resume_lora_training"] = False
+
+        if args["stage"] == "dpo":
            args["dpo_beta"] = dpo_beta
-        elif training_stage == "Pre-Training":
-            args["stage"] = "pt"

        if val_size > 1e-6:
            args["val_size"] = val_size
@@ -181,22 +176,26 @@ class Runner:
    ) -> Tuple[str, str, List[str], str, Dict[str, Any]]:
        if checkpoints:
            checkpoint_dir = ",".join(
-                [os.path.join(get_save_dir(model_name), finetuning_type, checkpoint) for checkpoint in checkpoints]
+                [get_save_dir(model_name, finetuning_type, ckpt) for ckpt in checkpoints]
            )
-            output_dir = os.path.join(get_save_dir(model_name), finetuning_type, "eval_" + "_".join(checkpoints))
+            output_dir = get_save_dir(model_name, finetuning_type, "eval_" + "_".join(checkpoints))
        else:
            checkpoint_dir = None
-            output_dir = os.path.join(get_save_dir(model_name), finetuning_type, "eval_base")
+            output_dir = get_save_dir(model_name, finetuning_type, "eval_base")
+
+        user_config = load_config()
+        cache_dir = user_config.get("cache_dir", None)

        args = dict(
            stage="sft",
            model_name_or_path=get_model_path(model_name),
            do_eval=True,
-            overwrite_cache=True,
+            overwrite_cache=False,
            predict_with_generate=True,
+            cache_dir=cache_dir,
            checkpoint_dir=checkpoint_dir,
            finetuning_type=finetuning_type,
-            quantization_bit=int(quantization_bit) if quantization_bit != "None" else None,
+            quantization_bit=int(quantization_bit) if quantization_bit in ["8", "4"] else None,
            template=template,
            system_prompt=system_prompt,
            dataset_dir=dataset_dir,
--- a/src/llmtuner/webui/utils.py
+++ b/src/llmtuner/webui/utils.py
@@ -3,7 +3,7 @@ import json
 import gradio as gr
 import matplotlib.figure
 import matplotlib.pyplot as plt
-from typing import TYPE_CHECKING, Any, Dict, Generator, List, Tuple
+from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional, Tuple
 from datetime import datetime

 from llmtuner.extras.ploting import smooth
@@ -36,6 +36,7 @@ def get_time() -> str:
 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)
+
    if (
        len(dataset) > 0
        and "file_name" in dataset_info[dataset[0]]
@@ -46,18 +47,26 @@ def can_preview(dataset_dir: str, dataset: list) -> Dict[str, Any]:
        return gr.update(interactive=False)


-def get_preview(dataset_dir: str, dataset: list) -> Tuple[int, list, Dict[str, Any]]:
+def get_preview(
+    dataset_dir: str, dataset: list, start: Optional[int] = 0, end: Optional[int] = 2
+) -> Tuple[int, list, Dict[str, Any]]:
    with open(os.path.join(dataset_dir, DATA_CONFIG), "r", encoding="utf-8") as f:
        dataset_info = json.load(f)
-    data_file = dataset_info[dataset[0]]["file_name"]
+
+    data_file: str = dataset_info[dataset[0]]["file_name"]
    with open(os.path.join(dataset_dir, data_file), "r", encoding="utf-8") as f:
-        data = json.load(f)
-    return len(data), data[:2], gr.update(visible=True)
+        if data_file.endswith(".json"):
+            data = json.load(f)
+        elif data_file.endswith(".jsonl"):
+            data = [json.loads(line) for line in f]
+        else:
+            data = [line for line in f]
+    return len(data), data[start:end], gr.update(visible=True)


 def can_quantize(finetuning_type: str) -> Dict[str, Any]:
    if finetuning_type != "lora":
-        return gr.update(value="", interactive=False)
+        return gr.update(value="None", interactive=False)
    else:
        return gr.update(interactive=True)

@@ -65,7 +74,7 @@ def can_quantize(finetuning_type: str) -> Dict[str, Any]:
 def gen_cmd(args: Dict[str, Any]) -> str:
    if args.get("do_train", None):
        args["plot_loss"] = True
-    cmd_lines = ["CUDA_VISIBLE_DEVICES=0 python "]
+    cmd_lines = ["CUDA_VISIBLE_DEVICES=0 python src/train_bash.py "]
    for k, v in args.items():
        if v is not None and v != "":
            cmd_lines.append("    --{} {} ".format(k, str(v)))
@@ -81,7 +90,7 @@ def get_eval_results(path: os.PathLike) -> str:


 def gen_plot(base_model: str, finetuning_type: str, output_dir: str) -> matplotlib.figure.Figure:
-    log_file = os.path.join(get_save_dir(base_model), finetuning_type, output_dir, "trainer_log.jsonl")
+    log_file = get_save_dir(base_model, finetuning_type, output_dir, "trainer_log.jsonl")
    if not os.path.isfile(log_file):
        return None

@@ -130,7 +139,7 @@ def save_model(
        return

    checkpoint_dir = ",".join(
-            [os.path.join(get_save_dir(model_name), finetuning_type, checkpoint) for checkpoint in checkpoints]
+            [get_save_dir(model_name, finetuning_type, ckpt) for ckpt in checkpoints]
        )

    if not save_dir:
--- a/tests/llamafy_baichuan2.py
+++ b/tests/llamafy_baichuan2.py
@@ -0,0 +1,65 @@
+# coding=utf-8
+# Converts the Baichuan2-7B model in the same format as LLaMA2-7B.
+# Usage: python llamafy_baichuan2.py --llama2_json llama2.index.json --input_dir input --output_dir output
+# 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
+from collections import OrderedDict
+
+
+SHARD_A = "pytorch_model-00001-of-00002.bin"
+SHARD_B = "pytorch_model-00002-of-00002.bin"
+
+
+def llamafy_baichuan2(
+    llama2_json: str,
+    input_dir: str,
+    output_dir: str
+):
+    baichuan2_state_dict = OrderedDict()
+    for filepath in os.listdir(input_dir):
+        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 = OrderedDict()
+    total_size = 0
+    for key, value in baichuan2_state_dict.items():
+        total_size += 2 * value.numel() # half precision
+        if "W_pack" in key:
+            llama2_state_dict[key.replace("W_pack", "q_proj")] = value[:4096, :]
+            llama2_state_dict[key.replace("W_pack", "k_proj")] = value[4096:2*4096, :]
+            llama2_state_dict[key.replace("W_pack", "v_proj")] = value[2*4096:, :]
+        elif "lm_head" in key:
+            llama2_state_dict[key] = torch.nn.functional.normalize(value)
+        else:
+            llama2_state_dict[key] = value
+
+    with open(os.path.join(input_dir, llama2_json), "r", encoding="utf-8") as f:
+        llama2_index = json.load(f)
+
+    merged_index = OrderedDict()
+    merged_index["metadata"] = {"total_size": total_size}
+    merged_index["weight_map"] = llama2_index["weight_map"]
+
+    state_dict_a, state_dict_b = OrderedDict(), OrderedDict()
+    for key, value in llama2_state_dict.items():
+        if merged_index["weight_map"][key] == SHARD_A:
+            state_dict_a[key] = value
+        else:
+            state_dict_b[key] = value
+
+    os.makedirs(output_dir, exist_ok=True)
+    torch.save(state_dict_a, os.path.join(output_dir, SHARD_A))
+    torch.save(state_dict_b, os.path.join(output_dir, SHARD_B))
+    with open(os.path.join(output_dir, "pytorch_model.bin.index.json"), "w", encoding="utf-8") as f:
+        json.dump(merged_index, f, indent=2)
+    print("Completed!")
+
+
+if __name__ == "__main__":
+    fire.Fire(llamafy_baichuan2)
--- a/tests/modeling_baichuan.py
+++ b/tests/modeling_baichuan.py
@@ -1,4 +1,6 @@
 # Copyright (c) 2023, Baichuan Intelligent Technology. All rights reserved.
+# Modified by hiyouga, to support attention mask, the alibi implementation is largely borrowed from
+# https://github.com/huggingface/transformers/blob/main/src/transformers/models/bloom/modeling_bloom.py

 import math
 from typing import List, Optional, Tuple, Union
@@ -12,7 +14,6 @@ from transformers import PreTrainedModel
 from transformers.activations import ACT2FN
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from transformers.utils import logging
-from transformers.generation.utils import GenerationConfig

 from .configuration_baichuan import BaichuanConfig

@@ -128,7 +129,7 @@ class MLP(nn.Module):

 class BaichuanAttention(nn.Module):

-    def __init__(self, config: BaichuanConfig):
+    def __init__(self, config: "BaichuanConfig"):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
@@ -223,7 +224,7 @@ class BaichuanAttention(nn.Module):

 class BaichuanLayer(nn.Module):

-    def __init__(self, config: BaichuanConfig):
+    def __init__(self, config: "BaichuanConfig"):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.self_attn = BaichuanAttention(config=config)
@@ -342,7 +343,7 @@ class BaichuanPreTrainedModel(PreTrainedModel):

 class BaichuanModel(BaichuanPreTrainedModel):

-    def __init__(self, config: BaichuanConfig):
+    def __init__(self, config: "BaichuanConfig"):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size
@@ -651,93 +652,3 @@ class BaichuanForCausalLM(BaichuanPreTrainedModel):
            for layer_past in standardized_past
        )
        return self._convert_to_baichuan_cache(reordered_past)
-
-    def quantize(self, bits: int):
-        try:
-            from .quantizer import QLinear
-        except ImportError:
-            raise ImportError(
-                f"Needs QLinear to run quantize."
-            )
-
-        for layer in self.model.layers:
-            layer.self_attn.W_pack = QLinear(
-                bits=bits,
-                weight=layer.self_attn.W_pack.weight,
-                bias = None,
-            )
-            layer.self_attn.o_proj = QLinear(
-                bits=bits,
-                weight=layer.self_attn.o_proj.weight,
-                bias = None,
-            )
-            layer.mlp.gate_proj = QLinear(
-                bits=bits,
-                weight=layer.mlp.gate_proj.weight,
-                bias = None,
-            )
-            layer.mlp.down_proj = QLinear(
-                bits=bits,
-                weight=layer.mlp.down_proj.weight,
-                bias = None,
-            )
-            layer.mlp.up_proj = QLinear(
-                bits=bits,
-                weight=layer.mlp.up_proj.weight,
-                bias = None,
-            )
-        return self
-
-    def _build_chat_input(self, tokenizer, messages: List[dict], max_new_tokens: int=0):
-        max_new_tokens = max_new_tokens or self.generation_config.max_new_tokens
-        max_input_tokens = self.config.model_max_length - max_new_tokens
-        max_input_tokens = max(self.config.model_max_length // 2, max_input_tokens)
-        total_input, round_input = [], []
-        for i, message in enumerate(messages[::-1]):
-            content_tokens = tokenizer.encode(message['content'])
-            if message['role'] == 'user':
-                round_input = [self.generation_config.user_token_id] + content_tokens + round_input
-                if total_input and len(total_input) + len(round_input) > max_input_tokens:
-                    break
-                else:
-                    total_input = round_input + total_input
-                    if len(total_input) >= max_input_tokens:
-                        break
-                    else:
-                        round_input = []
-            elif message['role'] == 'assistant':
-                round_input = [
-                    self.generation_config.assistant_token_id
-                ] + content_tokens + [
-                    self.generation_config.eos_token_id
-                ] + round_input
-            else:
-                raise ValueError(f"message role not supported yet: {message['role']}")
-        total_input = total_input[-max_input_tokens:]  # truncate left
-        total_input.append(self.generation_config.assistant_token_id)
-        total_input = torch.LongTensor([total_input]).to(self.device)
-        return total_input
-
-    @torch.no_grad()
-    def chat(self, tokenizer, messages: List[dict], stream=False,
-             generation_config: Optional[GenerationConfig]=None):
-        generation_config = generation_config or self.generation_config
-        input_ids = self._build_chat_input(tokenizer, messages, generation_config.max_new_tokens)
-        if stream:
-            from transformers_stream_generator.main import NewGenerationMixin, StreamGenerationConfig
-            self.__class__.generate = NewGenerationMixin.generate
-            self.__class__.sample_stream = NewGenerationMixin.sample_stream
-            stream_config = StreamGenerationConfig(**generation_config.to_dict(), do_stream=True)
-
-            def stream_generator():
-                outputs = []
-                for token in self.generate(input_ids, generation_config=stream_config):
-                    outputs.append(token.item())
-                    yield tokenizer.decode(outputs, skip_special_tokens=True)
-
-            return stream_generator()
-        else:
-            self.__class__.generate = PreTrainedModel.generate  # disable stream
-            outputs = self.generate(input_ids, generation_config=generation_config)
-            response = tokenizer.decode(outputs[0][len(input_ids[0]):], skip_special_tokens=True)
-            return response
Author	SHA1	Message	Date
hiyouga	f3e638ac6a	Release v0.1.8 Former-commit-id: d9666411375964d334d0a93ec162b27e05f70d49	2023-09-11 17:31:34 +08:00
hiyouga	42e0b30476	update flashattn, fix ppo save model Former-commit-id: 0b08bc3dac246d4aa3f89afb7172529dcad9c39f	2023-09-11 17:25:36 +08:00
hiyouga	a09a7b650d	remove PeftTrainer Former-commit-id: cc0cff3e991f194732d278e627648e528118a719	2023-09-10 22:23:23 +08:00
hiyouga	332d7bbd56	truncate readme Former-commit-id: fed5d0cc87e4a5a023f2edae622f2820bded1509	2023-09-10 21:04:20 +08:00
hiyouga	d3b6fece71	update readme Former-commit-id: c42fe77fec2918fe8811d48ec88e9a7c1e6f07ab	2023-09-10 21:01:20 +08:00
hiyouga	9d963b82de	update readme Former-commit-id: b4109cfe548e091cd20fa84815dce5ff3974a090	2023-09-10 20:52:21 +08:00
hiyouga	a402161631	support FlashAttention2 Former-commit-id: 23e56c5554b948d4f08ad87849b261eafd2c7890	2023-09-10 20:43:56 +08:00
hiyouga	b481ad58e6	fix #850 Former-commit-id: e5975c4c6b8bd47ec506b0d4a4703bee05495436	2023-09-10 14:22:03 +08:00
hiyouga	f91c5f2638	fix lora target Former-commit-id: d822e41e7ac7e310ee49e347fc45754284ce30b8	2023-09-09 17:04:45 +08:00
hiyouga	7143c551ab	support lora target auto find Former-commit-id: bce9984733d88bf013847eed523d1c75fdf0995e	2023-09-09 15:38:37 +08:00
hiyouga	50e93392dd	fix chatglm2 tokenizer Former-commit-id: 1ab60b4a93fa1be5dfe6ffbd4deb64c0f9d9b431	2023-09-09 13:50:29 +08:00
hiyouga	9f83e93839	add baichuan2 convert script Former-commit-id: 4d676e0ea9e59c1be13ecb47734917ba78938ac8	2023-09-08 22:59:41 +08:00
hiyouga	692b132dbf	fix bug in DPO data collator Former-commit-id: 4fc262cdf1347691e253bdfbd96568db5a49c086	2023-09-08 20:45:07 +08:00
hiyouga	e70b3e8947	fix #761 Former-commit-id: be76f6cbe5143f781b6b39603b80392253b3080a	2023-09-08 20:22:18 +08:00
hiyouga	612d97db6f	change to right-padding, update reward score #803 Former-commit-id: baa90415bc8f5ebd423d001378b51c3a3a6c2ec7	2023-09-08 20:04:31 +08:00
hiyouga	bb1b67c076	fix chatglm template Former-commit-id: 69a824628b4d6a56a680a7e713b217877c6c15c5	2023-09-08 14:45:58 +08:00
hiyouga	5a75c31caa	update requirements Former-commit-id: d796a4a5709c390629bafbeb7c91fccf6a9076d0	2023-09-07 19:26:25 +08:00
hiyouga	8b9210286b	fix #818 Former-commit-id: e81fd458c279ed2f3cee780e517482b425c8886d	2023-09-07 19:19:53 +08:00
hiyouga	b5acec34f7	add deepspeed check in PPO training Former-commit-id: e203ec7f71f504ccbaa89c27d20b8a0d9fa53f7e	2023-09-07 19:12:40 +08:00
hiyouga	86d835878c	fix #809 Former-commit-id: 2783ca75365d7c373cefba039788a48f0b8f35fc	2023-09-07 19:04:32 +08:00
hiyouga	eae7b331d3	fix baichuan templates Former-commit-id: f48a49e835b32f3991cfad8874c7b9c78953809f	2023-09-07 18:54:14 +08:00
hiyouga	ed89e29bcc	update baichuan2 template Former-commit-id: 16d9f8ba176443c5b397233da621600d6e1e1eec	2023-09-06 21:43:06 +08:00
hiyouga	c2b1886aff	add Baichuan2 models Former-commit-id: 90b3f02c44c0b8cc1b59f37af3a1ec28874a8a61	2023-09-06 18:40:11 +08:00
hiyouga	218f36bca5	add Baichuan2 models Former-commit-id: 36960025e9274b574f57e7a7bf453cd96956e922	2023-09-06 18:36:04 +08:00
hoshi-hiyouga	b91fc1f5b3	Merge pull request #786 from kinghuin/patch-1 fix utils.py bug Former-commit-id: 26aad616340748e1594a60119ca9434908bf7465	2023-09-05 10:49:34 +08:00
Q	2a22bf9c15	fix utils.py bug Former-commit-id: dc490117d50c3cbc070b804bac89400f4290272f	2023-09-05 10:38:01 +08:00
hiyouga	62e2037125	fix #763 Former-commit-id: e424b928a35097b783af879a2290f59b2158801d	2023-09-01 23:13:05 +08:00
hiyouga	e5b72c6a77	refactor dataset_attr, add eos in pt, fix #757 Former-commit-id: 0feec9a830b917b36686b61938a66e842eccf930	2023-09-01 19:00:45 +08:00
codingma	93be211f80	Merge pull request #741 from hiyouga/feature-addDatasetCheck Feature add dataset check Former-commit-id: 4b6dabe73d2c7edc94cd495390577c8bcf88428b	2023-08-31 20:57:36 +08:00
codemayq	9ae3fb4ced	update llama2 template Former-commit-id: 01de1d51d9fa5a22a338b6ed18ffad4d0ad5e3e8	2023-08-30 16:23:56 +08:00
codemayq	f641075789	add dataset stage check Former-commit-id: 5c719a7ce988339d034a653456da9742dc2cec7c	2023-08-30 16:23:08 +08:00
codingma	f7658db1b6	Merge pull request #651 from hiyouga/feature-dataset_stage add dataset stage Former-commit-id: 3b0ef57405cbc22ff8ce4eef2cfcb73872519db5	2023-08-28 16:03:45 +08:00
codemayq	b869bc1a20	add ad gen dataset Former-commit-id: fcd0788aa4dda0cecc1420d369d371032a207810	2023-08-27 20:35:32 +08:00
codemayq	a72d756d77	add text format dataset preview in webui Former-commit-id: cd30871aadb40cd3d598a6d0b415946744d2d550	2023-08-24 19:45:36 +08:00
codemayq	d3fd8f89b8	add stage in DatasetAttr Former-commit-id: 9c55200d8de0623640f529dbf39b8b0f169636d3	2023-08-23 20:54:53 +08:00
hiyouga	180a05a446	fix import error Former-commit-id: b3207a974a45038591b8cbbcf20d1ca1142d6679	2023-08-23 20:45:03 +08:00
hiyouga	eb9ac9ee1f	fix #649 Former-commit-id: e6120a937ddb4f3c0b9bcb2466742f5cf4f77f8c	2023-08-23 20:21:15 +08:00
codemayq	a6662b73f5	add readme for dataset Former-commit-id: bdcb0ea40e726e4c5752f938b379ed9a18e7e1d0	2023-08-23 19:55:45 +08:00
codemayq	cbc7db3478	add dataset stage and filter dataset when stage chosen in webui Former-commit-id: 26e4136449a4df6028d834fd16a0f4a7c532759d	2023-08-23 18:54:23 +08:00
hiyouga	4606340f0f	fix webui Former-commit-id: 95304b6822d9fe04bcddc1ee246a56389bd5f96a	2023-08-23 11:03:35 +08:00
hoshi-hiyouga	d4b4ccd597	Merge pull request #644 from hiyouga/fix-quantization_bit fix quantization bit is "" Former-commit-id: e1a8eca182e532b48e472919b4474656a726b40c	2023-08-23 10:45:45 +08:00
codemayq	9c3f4e3a37	fix quantization bit is "" Former-commit-id: 0dcab66f8843e2887f9f7ca66334122fef35c5b7	2023-08-23 10:08:17 +08:00
codemayq	440e00d8f9	fix quantization is "" Former-commit-id: 2469cc16d1dd3f5ee822edc18b2d7021ff7cba03	2023-08-23 10:04:03 +08:00
hiyouga	6310613699	update template Former-commit-id: a95f3a4d62de1073a78125401cf4289ec0523156	2023-08-22 19:46:09 +08:00
hoshi-hiyouga	f55907dbea	Merge pull request #629 from panpan0000/main add rm dataset explanation Former-commit-id: c2b4571d0ffb6298d6e07212982d9c13efd65adf	2023-08-22 13:41:44 +08:00
Peter Pan	5cac87d317	add rm dataset explanation Signed-off-by: Peter Pan <Peter.Pan@daocloud.io> Former-commit-id: 1efb95025be6501f1b30b20e7c711d3590b5d1ee	2023-08-22 01:33:59 -04:00
hoshi-hiyouga	9c0622de13	Merge pull request #619 from hiyouga/feature-templateTest add template encode test Former-commit-id: 8a1587ae49fff3968e0182f4fcc9a65dfdb260fc	2023-08-21 20:56:34 +08:00
codemayq	37b93c8b71	add template encode test Former-commit-id: c15e0d6847cbc055d8376b3c43ac4fbd17b5877a	2023-08-21 20:51:24 +08:00
hiyouga	d6be98cda6	fix #617 Former-commit-id: a7bdaf1c92c7d798caf8438dc42a8972632ec584	2023-08-21 18:16:11 +08:00
hiyouga	4d128acc17	fix #608 Former-commit-id: c02a6809124fcfd06628c49c95d419ec2d8cc8ef	2023-08-21 17:49:36 +08:00
hiyouga	516df9ecce	fix baichuan template for training #597 #616 Former-commit-id: 6530c1d972301eac9ef058b3235618bb09833f15	2023-08-21 17:41:51 +08:00
hiyouga	8eec1d50e1	fix #595 Former-commit-id: a360ccf9aa0484ce783eaa5857cf698b3ac2051e	2023-08-20 16:40:00 +08:00
hoshi-hiyouga	cfb096d43a	Merge pull request #596 from beat4ocean/beat fix KeyError: 'lang' bug Former-commit-id: dd22541cdf1b832d20bb894d78c034afce841bfb	2023-08-20 16:37:40 +08:00
beat4ocean	713fa28804	fix KeyError: 'lang' bug Former-commit-id: 4d4d9172b1f362cb4876315f1f5739e417055065	2023-08-20 15:32:36 +08:00
hiyouga	5549f35939	fix ppo trainer #551 Former-commit-id: 050a5447c191b8c50a0826a0f03bae499bff8b48	2023-08-20 14:07:11 +08:00