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
56 changed files with 2229 additions and 1220 deletions
--- a/.gitignore
+++ b/.gitignore
@@ -0,0 +1,160 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
--- a/README.md
+++ b/README.md
@@ -12,63 +12,61 @@

 ## Changelog

-[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/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/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/18] Now we support **resuming training**, upgrade `transformers` to `4.31.0` to enjoy this feature.

-[23/07/31] Now we support dataset streaming. Try `--streaming` and `--max_steps 100` arguments to stream your dataset.
+[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/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/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/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/31] Now we support **dataset streaming**. Try `--streaming` and `--max_steps 10000` arguments to load your dataset in streaming mode.

-[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/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/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/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/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/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

-| Approach               | Full-parameter | Partial-parameter | LoRA | QLoRA |
-| ---------------------- | -------------- | ----------------- | ---- | ----- |
-| Pre-Training           | ✅            | ✅                | ✅   | ✅   |
-| Supervised Fine-Tuning | ✅            | ✅                | ✅   | ✅   |
-| Reward Model Training  |                |                   | ✅   | ✅   |
-| PPO Training           |                |                   | ✅   | ✅   |
-| DPO Training           | ✅            |                    | ✅   | ✅   |
+| Approach               |   Full-parameter   | Partial-parameter  |       LoRA         |       QLoRA        |
+| ---------------------- | ------------------ | ------------------ | ------------------ | ------------------ |
+| Pre-Training           | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: |
+| Supervised Fine-Tuning | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: |
+| Reward Modeling        |                    |                    | :white_check_mark: | :white_check_mark: |
+| PPO Training           |                    |                    | :white_check_mark: | :white_check_mark: |
+| DPO Training           | :white_check_mark: |                    | :white_check_mark: | :white_check_mark: |
+
+> [!NOTE]
+> Use `--quantization_bit 4/8` argument to enable QLoRA.

 ## Provided Datasets

@@ -99,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)
- For reward modelling or DPO training:
+  - [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)
  - [GPT-4 Generated Data (en&zh)](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)
@@ -117,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)
@@ -130,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)

@@ -154,18 +154,26 @@ pip install https://github.com/jllllll/bitsandbytes-windows-webui/releases/downl
 CUDA_VISIBLE_DEVICES=0 python src/train_web.py
 ```

-Currently the web UI only supports training on **a single GPU**.
+We strongly recommend using the all-in-one Web UI for newcomers since it can also generate training scripts **automatically**.

-### Pre-Training
+> [!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
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage pt \
-    --model_name_or_path path_to_your_model \
+    --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 \
    --overwrite_cache \
    --per_device_train_batch_size 4 \
@@ -179,16 +187,17 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --fp16
 ```

-### Supervised Fine-Tuning
+#### Supervised Fine-Tuning

 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage sft \
-    --model_name_or_path path_to_your_model \
+    --model_name_or_path path_to_llama_model \
    --do_train \
    --dataset alpaca_gpt4_en \
    --template default \
    --finetuning_type lora \
+    --lora_target q_proj,v_proj \
    --output_dir path_to_sft_checkpoint \
    --overwrite_cache \
    --per_device_train_batch_size 4 \
@@ -202,18 +211,17 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --fp16
 ```

-Remember to specify `--lora_target W_pack` if you are using Baichuan models.
-
-### Reward Model Training
+#### Reward Modeling

 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage rm \
-    --model_name_or_path path_to_your_model \
+    --model_name_or_path path_to_llama_model \
    --do_train \
    --dataset comparison_gpt4_en \
    --template default \
    --finetuning_type lora \
+    --lora_target q_proj,v_proj \
    --resume_lora_training False \
    --checkpoint_dir path_to_sft_checkpoint \
    --output_dir path_to_rm_checkpoint \
@@ -222,22 +230,23 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --lr_scheduler_type cosine \
    --logging_steps 10 \
    --save_steps 1000 \
-    --learning_rate 1e-5 \
+    --learning_rate 1e-6 \
    --num_train_epochs 1.0 \
    --plot_loss \
    --fp16
 ```

-### PPO Training
+#### PPO Training

 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage ppo \
-    --model_name_or_path path_to_your_model \
+    --model_name_or_path path_to_llama_model \
    --do_train \
    --dataset alpaca_gpt4_en \
    --template default \
    --finetuning_type lora \
+    --lora_target q_proj,v_proj \
    --resume_lora_training False \
    --checkpoint_dir path_to_sft_checkpoint \
    --reward_model path_to_rm_checkpoint \
@@ -249,19 +258,21 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --save_steps 1000 \
    --learning_rate 1e-5 \
    --num_train_epochs 1.0 \
-    --plot_loss
+    --plot_loss \
+    --fp16
 ```

-### DPO Training
+#### DPO Training

 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage dpo \
-    --model_name_or_path path_to_your_model \
+    --model_name_or_path path_to_llama_model \
    --do_train \
    --dataset comparison_gpt4_en \
    --template default \
    --finetuning_type lora \
+    --lora_target q_proj,v_proj \
    --resume_lora_training False \
    --checkpoint_dir path_to_sft_checkpoint \
    --output_dir path_to_dpo_checkpoint \
@@ -278,24 +289,20 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \

 ### Distributed Training

+#### Use Huggingface Accelerate
+
 ```bash
 accelerate config # configure the environment
 accelerate launch src/train_bash.py # arguments (same as above)
 ```

-<details><summary>Example configuration for full-tuning 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
@@ -311,12 +318,95 @@ use_cpu: false

 </details>

+#### Use DeepSpeed
+
+```bash
+deepspeed --num_gpus 8 --master_port=9901 src/train_bash.py \
+    --deepspeed ds_config.json \
+    ... # arguments (same as above)
+```
+
+<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",
+  "zero_allow_untested_optimizer": true,
+  "fp16": {
+    "enabled": "auto",
+    "loss_scale": 0,
+    "initial_scale_power": 16,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  },  
+  "zero_optimization": {
+    "stage": 2,
+    "allgather_partitions": true,
+    "allgather_bucket_size": 5e8,
+    "reduce_scatter": true,
+    "reduce_bucket_size": 5e8,
+    "overlap_comm": false,
+    "contiguous_gradients": true
+  }
+}
+```
+
+</details>
+
+### Export model
+
+```bash
+python src/export_model.py \
+    --model_name_or_path path_to_llama_model \
+    --template default \
+    --finetuning_type lora \
+    --checkpoint_dir path_to_checkpoint \
+    --output_dir path_to_export
+```
+
+### API Demo
+
+```bash
+python src/api_demo.py \
+    --model_name_or_path path_to_llama_model \
+    --template default \
+    --finetuning_type lora \
+    --checkpoint_dir path_to_checkpoint
+```
+
+> [!NOTE]
+> Visit `http://localhost:8000/docs` for API documentation.
+
+### CLI Demo
+
+```bash
+python src/cli_demo.py \
+    --model_name_or_path path_to_llama_model \
+    --template default \
+    --finetuning_type lora \
+    --checkpoint_dir path_to_checkpoint
+```
+
+### Web Demo
+
+```bash
+python src/web_demo.py \
+    --model_name_or_path path_to_llama_model \
+    --template default \
+    --finetuning_type lora \
+    --checkpoint_dir path_to_checkpoint
+```
+
 ### Evaluation (BLEU and ROUGE_CHINESE)

 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage sft \
-    --model_name_or_path path_to_your_model \
+    --model_name_or_path path_to_llama_model \
    --do_eval \
    --dataset alpaca_gpt4_en \
    --template default \
@@ -328,14 +418,15 @@ 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

 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage sft \
-    --model_name_or_path path_to_your_model \
+    --model_name_or_path path_to_llama_model \
    --do_predict \
    --dataset alpaca_gpt4_en \
    --template default \
@@ -347,55 +438,6 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --predict_with_generate
 ```

-### API Demo
-
-```bash
-python src/api_demo.py \
-    --model_name_or_path path_to_your_model \
-    --template default \
-    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint
-```
-
-Visit `http://localhost:8000/docs` for API documentation.
-
-### CLI Demo
-
-```bash
-python src/cli_demo.py \
-    --model_name_or_path path_to_your_model \
-    --template default \
-    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint
-```
-
-### Web Demo
-
-```bash
-python src/web_demo.py \
-    --model_name_or_path path_to_your_model \
-    --template default \
-    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint
-```
-
-### Export model
-
-```bash
-python src/export_model.py \
-    --model_name_or_path path_to_your_model \
-    --template default \
-    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint \
-    --output_dir path_to_export
-```
-
-## 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).
@@ -407,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)
@@ -427,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,62 +12,61 @@

 ## 更新日志

-[23/08/11] 现在我们支持了指令模型的 **[DPO 训练](https://arxiv.org/abs/2305.18290)**。详情请参阅[此示例](#dpo-训练)（实验性功能）。
+[23/09/10] 现在我们支持了 LLaMA 模型的 **[FlashAttention](https://github.com/Dao-AILab/flash-attention)**。如果您使用的是 RTX4090、A100 或 H100 GPU，请使用 `--flash_attn` 参数以启用 FlashAttention-2（实验性功能）。

-[23/08/03] 现在我们支持了 **Qwen-7B** 模型的训练。请尝试使用 `--model_name_or_path Qwen/Qwen-7B-Chat` 和 `--lora_target c_attn` 参数。使用 Qwen-7B-Chat 模型请添加 `--template chatml` 参数。
+[23/08/18] 现在我们支持了**训练状态恢复**，请将 `transformers` 升级至 `4.31.0` 以启用此功能。

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

-[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/08/11] 现在我们支持了指令模型的 **[DPO 训练](https://arxiv.org/abs/2305.18290)**。详情请参阅[此示例](#dpo-训练)。

-[23/07/19] 现在我们支持了 **LLaMA-2** 模型的训练。请尝试使用 `--model_name_or_path meta-llama/Llama-2-7b-hf` 参数。请注意使用 LLaMA-2-chat 模型需要添加 `--template llama2` 参数。
+[23/07/31] 现在我们支持了**数据流式加载**。请尝试使用 `--streaming` 和 `--max_steps 10000` 参数来流式加载数据集。

-[23/07/18] 我们开发了支持训练和测试的浏览器一键微调界面。请尝试使用 `train_web.py` 在您的浏览器中微调模型。感谢 [@KanadeSiina](https://github.com/KanadeSiina) 和 [@codemayq](https://github.com/codemayq) 在该功能开发中付出的努力。
+[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/11] 现在我们支持了 **Baichuan-13B** 模型的训练。请尝试使用 `--model_name_or_path baichuan-inc/Baichuan-13B-Base` 和 `--lora_target W_pack` 参数。请注意使用 Baichuan-13B-Chat 模型需要添加 `--template baichuan` 参数。
+[23/07/18] 我们开发了支持训练和测试的**浏览器一体化界面**。请尝试使用 `train_web.py` 在您的浏览器中微调模型。感谢 [@KanadeSiina](https://github.com/KanadeSiina) 和 [@codemayq](https://github.com/codemayq) 在该功能开发中付出的努力。

-[23/07/09] 我们开源了 [FastEdit](https://github.com/hiyouga/FastEdit)⚡🩹，一个简单易用的、能迅速编辑大模型事实记忆的工具包。如果您感兴趣请关注我们的 [FastEdit](https://github.com/hiyouga/FastEdit) 项目。
+[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/06/29] 我们提供了一个**可复现的**指令模型微调示例，详细内容请查阅 [Baichuan-7B-sft](https://huggingface.co/hiyouga/Baichuan-7B-sft)。

-[23/07/05] 现在我们支持了 **Falcon-7B/40B** 模型的训练。请尝试使用 `--model_name_or_path tiiuae/falcon-7b` 和 `--lora_target query_key_value` 参数。
+[23/06/22] 我们对齐了[示例 API](src/api_demo.py) 与 [OpenAI API](https://platform.openai.com/docs/api-reference/chat) 的格式，您可以将微调模型接入**任意基于 ChatGPT 的应用**中。

-[23/06/29] 我们提供了一个**可复现的**指令模型微调示例，详细内容请查阅 [Hugging Face 项目](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` 参数。
+[23/06/03] 现在我们实现了 4 比特的 LoRA 训练（也称 **[QLoRA](https://github.com/artidoro/qlora)**）。请尝试使用 `--quantization_bit 4` 参数进行 4 比特量化微调。

 ## 模型

-| 模型名                                                   | 模型大小                     | 默认模块           | 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     | -        |
+| 模型名                                                   | 模型大小                     | 默认模块           | 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）模型请务必使用对应的模板。

 ## 训练方法

-| 方法        | 全参数训练 | 部分参数训练 | LoRA | QLoRA |
-| ---------- | ---------- | ----------- | ---- | ----- |
-| 预训练      | ✅        | ✅         | ✅   | ✅   |
-| 指令监督微调 | ✅        | ✅         | ✅   | ✅   |
-| 奖励模型训练 |           |             | ✅   | ✅   |
-| PPO 训练    |           |             | ✅   | ✅   |
-| DPO 训练    | ✅        |             | ✅   | ✅   |
+| 方法                   |     全参数训练      |    部分参数训练     |       LoRA         |       QLoRA        |
+| ---------------------- | ------------------ | ------------------ | ------------------ | ------------------ |
+| 预训练                 | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: |
+| 指令监督微调            | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: |
+| 奖励模型训练            |                    |                    | :white_check_mark: | :white_check_mark: |
+| PPO 训练               |                    |                    | :white_check_mark: | :white_check_mark: |
+| DPO 训练               | :white_check_mark: |                    | :white_check_mark: | :white_check_mark: |
+
+> [!NOTE]
+> 请使用 `--quantization_bit 4/8` 参数来启用 QLoRA 训练。

 ## 数据集

@@ -98,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)
@@ -116,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)
@@ -129,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`。

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

@@ -147,24 +148,32 @@ pip install -r requirements.txt
 pip install https://github.com/jllllll/bitsandbytes-windows-webui/releases/download/wheels/bitsandbytes-0.39.1-py3-none-win_amd64.whl
 ```

-### 浏览器一键微调/测试
+### 浏览器一体化界面

 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_web.py
 ```

-目前网页 UI 仅支持**单卡训练**。
+我们极力推荐新手使用浏览器一体化界面，因为它还可以**自动**生成运行所需的命令行脚本。

-### 预训练
+> [!WARNING]
+> 目前网页 UI 仅支持**单卡训练**。
+
+### 单 GPU 训练
+
+> [!IMPORTANT]
+> 如果您使用多张 GPU 训练模型，请移步[多 GPU 分布式训练](#多-gpu-分布式训练)部分。
+
+#### 预训练

 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage pt \
-    --model_name_or_path path_to_your_model \
+    --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 \
    --overwrite_cache \
    --per_device_train_batch_size 4 \
@@ -178,16 +187,17 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --fp16
 ```

-### 指令监督微调
+#### 指令监督微调

 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage sft \
-    --model_name_or_path path_to_your_model \
+    --model_name_or_path path_to_llama_model \
    --do_train \
    --dataset alpaca_gpt4_zh \
    --template default \
    --finetuning_type lora \
+    --lora_target q_proj,v_proj \
    --output_dir path_to_sft_checkpoint \
    --overwrite_cache \
    --per_device_train_batch_size 4 \
@@ -201,18 +211,17 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --fp16
 ```

-使用 Baichuan 模型时请指定 `--lora_target W_pack` 参数。
-
-### 奖励模型训练
+#### 奖励模型训练

 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage rm \
-    --model_name_or_path path_to_your_model \
+    --model_name_or_path path_to_llama_model \
    --do_train \
    --dataset comparison_gpt4_zh \
    --template default \
    --finetuning_type lora \
+    --lora_target q_proj,v_proj \
    --resume_lora_training False \
    --checkpoint_dir path_to_sft_checkpoint \
    --output_dir path_to_rm_checkpoint \
@@ -221,22 +230,23 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --lr_scheduler_type cosine \
    --logging_steps 10 \
    --save_steps 1000 \
-    --learning_rate 1e-5 \
+    --learning_rate 1e-6 \
    --num_train_epochs 1.0 \
    --plot_loss \
    --fp16
 ```

-### PPO 训练
+#### PPO 训练

 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage ppo \
-    --model_name_or_path path_to_your_model \
+    --model_name_or_path path_to_llama_model \
    --do_train \
    --dataset alpaca_gpt4_zh \
    --template default \
    --finetuning_type lora \
+    --lora_target q_proj,v_proj \
    --resume_lora_training False \
    --checkpoint_dir path_to_sft_checkpoint \
    --reward_model path_to_rm_checkpoint \
@@ -251,16 +261,17 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --plot_loss
 ```

-### DPO 训练
+#### DPO 训练

 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage dpo \
-    --model_name_or_path path_to_your_model \
+    --model_name_or_path path_to_llama_model \
    --do_train \
    --dataset comparison_gpt4_zh \
    --template default \
    --finetuning_type lora \
+    --lora_target q_proj,v_proj \
    --resume_lora_training False \
    --checkpoint_dir path_to_sft_checkpoint \
    --output_dir path_to_dpo_checkpoint \
@@ -277,24 +288,20 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \

 ### 多 GPU 分布式训练

+#### 使用 Huggingface Accelerate
+
 ```bash
 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
@@ -310,12 +317,95 @@ use_cpu: false

 </details>

-### 指标评估（BLEU分数和汉语ROUGE分数）
+#### 使用 DeepSpeed
+
+```bash
+deepspeed --num_gpus 8 --master_port=9901 src/train_bash.py \
+    --deepspeed ds_config.json \
+    ... # 参数同上
+```
+
+<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",
+  "zero_allow_untested_optimizer": true,
+  "fp16": {
+    "enabled": "auto",
+    "loss_scale": 0,
+    "initial_scale_power": 16,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  },  
+  "zero_optimization": {
+    "stage": 2,
+    "allgather_partitions": true,
+    "allgather_bucket_size": 5e8,
+    "reduce_scatter": true,
+    "reduce_bucket_size": 5e8,
+    "overlap_comm": false,
+    "contiguous_gradients": true
+  }
+}
+```
+
+</details>
+
+### 导出微调后的模型
+
+```bash
+python src/export_model.py \
+    --model_name_or_path path_to_llama_model \
+    --template default \
+    --finetuning_type lora \
+    --checkpoint_dir path_to_checkpoint \
+    --output_dir path_to_export
+```
+
+### API 服务
+
+```bash
+python src/api_demo.py \
+    --model_name_or_path path_to_llama_model \
+    --template default \
+    --finetuning_type lora \
+    --checkpoint_dir path_to_checkpoint
+```
+
+> [!NOTE]
+> 关于 API 文档请见 `http://localhost:8000/docs`。
+
+### 命令行测试
+
+```bash
+python src/cli_demo.py \
+    --model_name_or_path path_to_llama_model \
+    --template default \
+    --finetuning_type lora \
+    --checkpoint_dir path_to_checkpoint
+```
+
+### 浏览器测试
+
+```bash
+python src/web_demo.py \
+    --model_name_or_path path_to_llama_model \
+    --template default \
+    --finetuning_type lora \
+    --checkpoint_dir path_to_checkpoint
+```
+
+### 指标评估（BLEU 分数和汉语 ROUGE 分数）

 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage sft \
-    --model_name_or_path path_to_your_model \
+    --model_name_or_path path_to_llama_model \
    --do_eval \
    --dataset alpaca_gpt4_zh \
    --template default \
@@ -327,14 +417,15 @@ 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`。

 ### 模型预测

 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --stage sft \
-    --model_name_or_path path_to_your_model \
+    --model_name_or_path path_to_llama_model \
    --do_predict \
    --dataset alpaca_gpt4_zh \
    --template default \
@@ -346,55 +437,6 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
    --predict_with_generate
 ```

-### API 服务
-
-```bash
-python src/api_demo.py \
-    --model_name_or_path path_to_your_model \
-    --template default \
-    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint
-```
-
-关于 API 文档请见 `http://localhost:8000/docs`。
-
-### 命令行测试
-
-```bash
-python src/cli_demo.py \
-    --model_name_or_path path_to_your_model \
-    --template default \
-    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint
-```
-
-### 浏览器测试
-
-```bash
-python src/web_demo.py \
-    --model_name_or_path path_to_your_model \
-    --template default \
-    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint
-```
-
-### 导出微调模型
-
-```bash
-python src/export_model.py \
-    --model_name_or_path path_to_your_model \
-    --template default \
-    --finetuning_type lora \
-    --checkpoint_dir path_to_checkpoint \
-    --output_dir path_to_export
-```
-
-## 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) 协议开源。
@@ -406,8 +448,11 @@ python src/export_model.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)
+- [ChatGLM2](https://github.com/THUDM/ChatGLM2-6B/blob/main/MODEL_LICENSE)

 ## 引用

@@ -424,7 +469,7 @@ python src/export_model.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.6"
+__version__ = "0.1.8"
--- a/src/llmtuner/api/app.py
+++ b/src/llmtuner/api/app.py
@@ -47,15 +47,15 @@ def create_app(chat_model: ChatModel) -> FastAPI:

    @app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
    async def create_chat_completion(request: ChatCompletionRequest):
-        if request.messages[-1].role != Role.USER:
+        if len(request.messages) < 1 or request.messages[-1].role != Role.USER:
            raise HTTPException(status_code=400, detail="Invalid request")

        query = request.messages[-1].content
        prev_messages = request.messages[:-1]
        if len(prev_messages) > 0 and prev_messages[0].role == Role.SYSTEM:
-            prefix = prev_messages.pop(0).content
+            system = prev_messages.pop(0).content
        else:
-            prefix = None
+            system = None

        history = []
        if len(prev_messages) % 2 == 0:
@@ -64,11 +64,11 @@ def create_app(chat_model: ChatModel) -> FastAPI:
                    history.append([prev_messages[i].content, prev_messages[i+1].content])

        if request.stream:
-            generate = predict(query, history, prefix, request)
+            generate = predict(query, history, system, request)
            return EventSourceResponse(generate, media_type="text/event-stream")

        response, (prompt_length, response_length) = chat_model.chat(
-            query, history, prefix, temperature=request.temperature, top_p=request.top_p, max_new_tokens=request.max_tokens
+            query, history, system, temperature=request.temperature, top_p=request.top_p, max_new_tokens=request.max_tokens
        )

        usage = ChatCompletionResponseUsage(
@@ -85,7 +85,7 @@ def create_app(chat_model: ChatModel) -> FastAPI:

        return ChatCompletionResponse(model=request.model, choices=[choice_data], usage=usage)

-    async def predict(query: str, history: List[Tuple[str, str]], prefix: str, request: ChatCompletionRequest):
+    async def predict(query: str, history: List[Tuple[str, str]], system: str, request: ChatCompletionRequest):
        choice_data = ChatCompletionResponseStreamChoice(
            index=0,
            delta=DeltaMessage(role=Role.ASSISTANT),
@@ -95,7 +95,7 @@ def create_app(chat_model: ChatModel) -> FastAPI:
        yield chunk.json(exclude_unset=True, ensure_ascii=False)

        for new_text in chat_model.stream_chat(
-            query, history, prefix, temperature=request.temperature, top_p=request.top_p, max_new_tokens=request.max_tokens
+            query, history, system, temperature=request.temperature, top_p=request.top_p, max_new_tokens=request.max_tokens
        ):
            if len(new_text) == 0:
                continue
--- a/src/llmtuner/chat/stream_chat.py
+++ b/src/llmtuner/chat/stream_chat.py
@@ -1,10 +1,9 @@
 import torch
-from types import MethodType
 from typing import Any, Dict, Generator, List, Optional, Tuple
 from threading import Thread
-from transformers import PreTrainedModel, TextIteratorStreamer
+from transformers import GenerationConfig, TextIteratorStreamer

-from llmtuner.extras.misc import dispatch_model, get_logits_processor, get_stopping_criteria
+from llmtuner.extras.misc import dispatch_model, get_logits_processor
 from llmtuner.extras.template import get_template_and_fix_tokenizer
 from llmtuner.tuner.core import get_infer_args, load_model_and_tokenizer

@@ -15,22 +14,20 @@ 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() # change to eval mode
        self.template = get_template_and_fix_tokenizer(data_args.template, self.tokenizer)
-        self.source_prefix = data_args.source_prefix
-        self.model.generate = MethodType(PreTrainedModel.generate, self.model) # disable custom method (for Qwen)
+        self.system_prompt = data_args.system_prompt

    def process_args(
        self,
        query: str,
        history: Optional[List[Tuple[str, str]]] = None,
-        prefix: Optional[str] = None,
+        system: Optional[str] = None,
        **input_kwargs
    ) -> Tuple[Dict[str, Any], int]:
-        prefix = prefix or self.source_prefix
+        system = system or self.system_prompt

        prompt, _ = self.template.encode_oneturn(
-            tokenizer=self.tokenizer, query=query, resp="", history=history, prefix=prefix
+            tokenizer=self.tokenizer, query=query, resp="", history=history, system=system
        )
        input_ids = torch.tensor([prompt], device=self.model.device)
        prompt_length = len(input_ids[0])
@@ -43,25 +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"],
-            logits_processor=get_logits_processor(),
-            stopping_criteria=get_stopping_criteria(self.tokenizer.additional_special_tokens_ids)
+        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

@@ -70,10 +72,10 @@ class ChatModel:
        self,
        query: str,
        history: Optional[List[Tuple[str, str]]] = None,
-        prefix: Optional[str] = None,
+        system: Optional[str] = None,
        **input_kwargs
    ) -> Tuple[str, Tuple[int, int]]:
-        gen_kwargs, prompt_length = self.process_args(query, history, prefix, **input_kwargs)
+        gen_kwargs, prompt_length = self.process_args(query, history, system, **input_kwargs)
        generation_output = self.model.generate(**gen_kwargs)
        outputs = generation_output.tolist()[0][prompt_length:]
        response = self.tokenizer.decode(outputs, skip_special_tokens=True)
@@ -85,10 +87,10 @@ class ChatModel:
        self,
        query: str,
        history: Optional[List[Tuple[str, str]]] = None,
-        prefix: Optional[str] = None,
+        system: Optional[str] = None,
        **input_kwargs
    ) -> Generator[str, None, None]:
-        gen_kwargs, _ = self.process_args(query, history, prefix, **input_kwargs)
+        gen_kwargs, _ = self.process_args(query, history, system, **input_kwargs)
        streamer = TextIteratorStreamer(self.tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
        gen_kwargs["streamer"] = streamer

--- a/src/llmtuner/dsets/loader.py
+++ b/src/llmtuner/dsets/loader.py
@@ -1,48 +1,25 @@
 import os
-import hashlib
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, List, Union

-from datasets import Value, concatenate_datasets, interleave_datasets, load_dataset
+from datasets import concatenate_datasets, interleave_datasets, load_dataset

+from llmtuner.dsets.utils import checksum, EXT2TYPE
 from llmtuner.extras.logging import get_logger

 if TYPE_CHECKING:
-    from datasets import Dataset
+    from datasets import Dataset, IterableDataset
    from llmtuner.hparams import ModelArguments, DataArguments


 logger = get_logger(__name__)


-EXT2TYPE = {
-    "csv": "csv",
-    "json": "json",
-    "jsonl": "json",
-    "txt": "text"
-}
-
-
-def checksum(data_files: List[str], file_sha1: Optional[str] = None) -> None:
-    if file_sha1 is None:
-        logger.warning("Checksum failed: missing SHA-1 hash value in dataset_info.json.")
-        return
-
-    if len(data_files) != 1:
-        logger.warning("Checksum failed: too many files.")
-        return
-
-    with open(data_files[0], "rb") as f:
-        sha1 = hashlib.sha1(f.read()).hexdigest()
-        if sha1 != file_sha1:
-            logger.warning("Checksum failed: mismatched SHA-1 hash value at {}.".format(data_files[0]))
-
-
 def get_dataset(
    model_args: "ModelArguments",
    data_args: "DataArguments"
-) -> "Dataset":
+) -> Union["Dataset", "IterableDataset"]:
    max_samples = data_args.max_samples
-    all_datasets: List["Dataset"] = [] # support multiple datasets
+    all_datasets: List[Union["Dataset", "IterableDataset"]] = [] # support multiple datasets

    for dataset_attr in data_args.dataset_list:
        logger.info("Loading dataset {}...".format(dataset_attr))
@@ -92,14 +69,11 @@ def get_dataset(
            if getattr(dataset_attr, column_name) and getattr(dataset_attr, column_name) != column_name:
                dataset = dataset.rename_column(getattr(dataset_attr, column_name), column_name)

-        if dataset_attr.source_prefix: # add prefix
+        if dataset_attr.system_prompt: # add system prompt
            if data_args.streaming:
-                features = dataset.features
-                features["prefix"] = Value(dtype="string", id=None)
-                dataset = dataset.map(lambda _: {"prefix": dataset_attr.source_prefix}, features=features)
+                dataset = dataset.map(lambda _: {"system": dataset_attr.system_prompt})
            else:
-                prefix_data = [dataset_attr.source_prefix] * len(dataset)
-                dataset = dataset.add_column("prefix", prefix_data)
+                dataset = dataset.add_column("system", [dataset_attr.system_prompt] * len(dataset))

        all_datasets.append(dataset)

--- a/src/llmtuner/dsets/preprocess.py
+++ b/src/llmtuner/dsets/preprocess.py
@@ -1,25 +1,25 @@
 import tiktoken
-from typing import TYPE_CHECKING, Any, Dict, Generator, List, Literal
+from typing import TYPE_CHECKING, Any, Dict, Generator, List, Literal, Union
 from itertools import chain

 from llmtuner.extras.constants import IGNORE_INDEX
 from llmtuner.extras.template import get_template_and_fix_tokenizer

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


 def preprocess_dataset(
-    dataset: "Dataset",
+    dataset: Union["Dataset", "IterableDataset"],
    tokenizer: "PreTrainedTokenizer",
    data_args: "DataArguments",
    training_args: "Seq2SeqTrainingArguments",
    stage: Literal["pt", "sft", "rm", "ppo"]
-) -> "Dataset":
-    column_names = list(dataset.column_names)
+) -> Union["Dataset", "IterableDataset"]:
+    column_names = list(next(iter(dataset)).keys())
    template = get_template_and_fix_tokenizer(data_args.template, tokenizer)

    def construct_example(examples: Dict[str, List[Any]]) -> Generator[Any, None, None]:
@@ -27,15 +27,19 @@ def preprocess_dataset(
            query, response = examples["prompt"][i], examples["response"][i]
            query = query + "\n" + examples["query"][i] if "query" in examples and examples["query"][i] else query
            history = examples["history"][i] if "history" in examples else None
-            prefix = examples["prefix"][i] if "prefix" in examples else None
-            yield query, response, history, prefix
+            system = examples["system"][i] if "system" in examples else None
+            yield query, response, history, system

    def preprocess_pretrain_dataset(examples: Dict[str, List[Any]]) -> Dict[str, 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()}
@@ -56,10 +60,12 @@ def preprocess_dataset(
        model_inputs = {"input_ids": [], "attention_mask": [], "labels": []}
        max_length = data_args.max_source_length + data_args.max_target_length

-        for query, response, history, prefix in construct_example(examples):
+        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, prefix):
+            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))
@@ -78,17 +93,20 @@ def preprocess_dataset(
        return model_inputs

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

-        for query, response, history, prefix in construct_example(examples):
-            source_ids, target_ids = template.encode_oneturn(tokenizer, query, response, history, prefix)
+        for query, response, history, system in construct_example(examples):
+            source_ids, target_ids = template.encode_oneturn(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:
                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)
@@ -98,9 +116,9 @@ def preprocess_dataset(
    def preprocess_pairwise_dataset(examples):
        # build input pairs with format `<bos> X`, `Y1 <eos>` and `Y2 <eos>`
        model_inputs = {"prompt_ids": [], "chosen_ids": [], "rejected_ids": []}
-        for query, response, history, prefix in construct_example(examples):
-            prompt_ids, chosen_ids = template.encode_oneturn(tokenizer, query, response[0], history, prefix)
-            _, rejected_ids = template.encode_oneturn(tokenizer, query, response[1], history, prefix)
+        for query, response, history, system in construct_example(examples):
+            prompt_ids, chosen_ids = template.encode_oneturn(tokenizer, query, response[0], history, system)
+            _, rejected_ids = template.encode_oneturn(tokenizer, query, response[1], history, system)

            if len(prompt_ids) > data_args.max_source_length:
                prompt_ids = prompt_ids[:data_args.max_source_length]
@@ -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/dsets/utils.py
+++ b/src/llmtuner/dsets/utils.py
@@ -1,4 +1,7 @@
-from typing import TYPE_CHECKING, Dict, Union
+import hashlib
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
+
+from llmtuner.extras.logging import get_logger

 if TYPE_CHECKING:
    from datasets import Dataset, IterableDataset
@@ -6,6 +9,32 @@ if TYPE_CHECKING:
    from llmtuner.hparams import DataArguments


+logger = get_logger(__name__)
+
+
+EXT2TYPE = {
+    "csv": "csv",
+    "json": "json",
+    "jsonl": "json",
+    "txt": "text"
+}
+
+
+def checksum(data_files: List[str], file_sha1: Optional[str] = None) -> None:
+    if file_sha1 is None:
+        logger.warning("Checksum failed: missing SHA-1 hash value in dataset_info.json.")
+        return
+
+    if len(data_files) != 1:
+        logger.warning("Checksum failed: too many files.")
+        return
+
+    with open(data_files[0], "rb") as f:
+        sha1 = hashlib.sha1(f.read()).hexdigest()
+        if sha1 != file_sha1:
+            logger.warning("Checksum failed: mismatched SHA-1 hash value at {}.".format(data_files[0]))
+
+
 def split_dataset(
    dataset: Union["Dataset", "IterableDataset"],
    data_args: "DataArguments",
--- 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,14 +2,18 @@ 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"]

+TRAINING_STAGES = {
+    "Supervised Fine-Tuning": "sft",
+    "Reward Modeling": "rm",
+    "PPO": "ppo",
+    "DPO": "dpo",
+    "Pre-Training": "pt"
+}
+
 SUPPORTED_MODELS = {
    "LLaMA-7B": "huggyllama/llama-7b",
    "LLaMA-13B": "huggyllama/llama-13b",
@@ -21,6 +25,10 @@ SUPPORTED_MODELS = {
    "LLaMA2-7B-Chat": "meta-llama/Llama-2-7b-chat-hf",
    "LLaMA2-13B-Chat": "meta-llama/Llama-2-13b-chat-hf",
    "LLaMA2-70B-Chat": "meta-llama/Llama-2-70b-chat-hf",
+    "ChineseLLaMA2-7B": "ziqingyang/chinese-llama-2-7b",
+    "ChineseLLaMA2-13B": "ziqingyang/chinese-llama-2-13b",
+    "ChineseLLaMA2-7B-Chat": "ziqingyang/chinese-alpaca-2-7b",
+    "ChineseLLaMA2-13B-Chat": "ziqingyang/chinese-alpaca-2-13b",
    "BLOOM-560M": "bigscience/bloom-560m",
    "BLOOM-3B": "bigscience/bloom-3b",
    "BLOOM-7B1": "bigscience/bloom-7b1",
@@ -34,23 +42,41 @@ 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",
-    "ChatGLM2-6B": "THUDM/chatglm2-6b"
+    "XVERSE-13B-Chat": "xverse/XVERSE-13B-Chat",
+    "ChatGLM2-6B-Chat": "THUDM/chatglm2-6b"
 }

 DEFAULT_MODULE = {
    "LLaMA": "q_proj,v_proj",
    "LLaMA2": "q_proj,v_proj",
+    "ChineseLLaMA2": "q_proj,v_proj",
    "BLOOM": "query_key_value",
    "BLOOMZ": "query_key_value",
    "Falcon": "query_key_value",
    "Baichuan": "W_pack",
+    "Baichuan2": "W_pack",
    "InternLM": "q_proj,v_proj",
    "Qwen": "c_attn",
    "XVERSE": "q_proj,v_proj",
    "ChatGLM2": "query_key_value"
 }
+
+DEFAULT_TEMPLATE = {
+    "LLaMA2": "llama2",
+    "ChineseLLaMA2": "llama2_zh",
+    "Baichuan": "baichuan",
+    "Baichuan2": "baichuan2",
+    "InternLM": "intern",
+    "Qwen": "chatml",
+    "XVERSE": "xverse",
+    "ChatGLM2": "chatglm2"
+}
--- a/src/llmtuner/extras/logging.py
+++ b/src/llmtuner/extras/logging.py
@@ -8,6 +8,9 @@ class LoggerHandler(logging.Handler):
        super().__init__()
        self.log = ""

+    def reset(self):
+        self.log = ""
+
    def emit(self, record):
        if record.name == "httpx":
            return
--- a/src/llmtuner/extras/misc.py
+++ b/src/llmtuner/extras/misc.py
@@ -1,13 +1,7 @@
+import gc
 import torch
-from typing import TYPE_CHECKING, List, Optional, Tuple
-from transformers import (
-    LogitsProcessor,
-    LogitsProcessorList,
-    StoppingCriteria,
-    StoppingCriteriaList
-)
-
-from llmtuner.extras.constants import LAYERNORM_NAMES
+from typing import TYPE_CHECKING, Tuple
+from transformers import InfNanRemoveLogitsProcessor, LogitsProcessorList

 if TYPE_CHECKING:
    from transformers.modeling_utils import PreTrainedModel
@@ -33,37 +27,6 @@ class AverageMeter:
        self.avg = self.sum / self.count


-class InvalidScoreLogitsProcessor(LogitsProcessor):
-
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
-        if torch.isnan(scores).any() or torch.isinf(scores).any():
-            scores.zero_()
-            scores[..., 0] = 1.0
-        return scores
-
-
-def get_logits_processor() -> LogitsProcessorList:
-    logits_processor = LogitsProcessorList()
-    logits_processor.append(InvalidScoreLogitsProcessor())
-    return logits_processor
-
-
-class StopWordsCriteria(StoppingCriteria):
-
-    def __init__(self, stop_ids: List[int]) -> None:
-        super().__init__()
-        self.stop_ids = stop_ids
-
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-        return any([stop_id in input_ids[:, -1] for stop_id in self.stop_ids])
-
-
-def get_stopping_criteria(stop_ids: List[int]) -> StoppingCriteriaList:
-    stopping_criteria = StoppingCriteriaList()
-    stopping_criteria.append(StopWordsCriteria(stop_ids))
-    return stopping_criteria
-
-
 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.
@@ -86,52 +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):
-        if hasattr(model, "config") and hasattr(model.config, "pretraining_tp"):
-            model.config.pretraining_tp = 1 # disable TP for LoRA (https://github.com/huggingface/peft/pull/728)
-
-        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()
@@ -142,6 +70,9 @@ def dispatch_model(model: "PreTrainedModel") -> "PreTrainedModel":
    Dispatches a pre-trained model to GPUs with balanced memory.
    Borrowed from: https://github.com/huggingface/transformers/blob/v4.31.0/src/transformers/modeling_utils.py#L2803
    """
+    if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False): # do nothing
+        return model
+
    if torch.cuda.device_count() > 1:
        from accelerate import dispatch_model
        from accelerate.utils import infer_auto_device_map, get_balanced_memory
--- 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
@@ -16,9 +16,11 @@ class Template:

    prefix: List[Union[str, Dict[str, str]]]
    prompt: List[Union[str, Dict[str, str]]]
+    system: str
    sep: List[Union[str, Dict[str, str]]]
    stop_words: List[str]
    use_history: bool
+    efficient_eos: bool

    def encode_oneturn(
        self,
@@ -26,18 +28,18 @@ class Template:
        query: str,
        resp: str,
        history: Optional[List[Tuple[str, str]]] = None,
-        prefix: Optional[str] = None
+        system: Optional[str] = None
    ) -> Tuple[List[int], List[int]]:
        r"""
        Returns a single pair of token ids representing prompt and response respectively.
        """
-        prefix, history = self._format(query, resp, history, prefix)
-        encoded_pairs = self._encode(tokenizer, prefix, history)
+        system, history = self._format(query, resp, history, system)
+        encoded_pairs = self._encode(tokenizer, system, history)
        prompt_ids = []
        for query_ids, resp_ids in encoded_pairs[:-1]:
            prompt_ids = prompt_ids + query_ids + resp_ids
-        prompt_ids = prompt_ids + encoded_pairs[-1][0]
-        return prompt_ids, encoded_pairs[-1][1]
+        prompt_ids, answer_ids = prompt_ids + encoded_pairs[-1][0], encoded_pairs[-1][1]
+        return prompt_ids, answer_ids

    def encode_multiturn(
        self,
@@ -45,13 +47,13 @@ class Template:
        query: str,
        resp: str,
        history: Optional[List[Tuple[str, str]]] = None,
-        prefix: Optional[str] = None
+        system: Optional[str] = None
    ) -> List[Tuple[List[int], List[int]]]:
        r"""
        Returns multiple pairs of token ids representing prompts and responses respectively.
        """
-        prefix, history = self._format(query, resp, history, prefix)
-        encoded_pairs = self._encode(tokenizer, prefix, history)
+        system, history = self._format(query, resp, history, system)
+        encoded_pairs = self._encode(tokenizer, system, history)
        return encoded_pairs

    def _format(
@@ -59,36 +61,39 @@ class Template:
        query: str,
        resp: str,
        history: Optional[List[Tuple[str, str]]] = None,
-        prefix: Optional[str] = None
-    ) -> Tuple[List[Union[str, Dict[str, str]]], List[Tuple[str, str]]]:
+        system: Optional[str] = None
+    ) -> Tuple[str, List[Tuple[str, str]]]:
        r"""
        Aligns inputs to the standard format.
        """
-        prefix = [prefix] if prefix else self.prefix # use prefix if provided
+        system = system or self.system # use system if provided
        history = history if (history and self.use_history) else []
        history = history + [(query, resp)]
-        return prefix, history
+        return system, history

    def _get_special_ids(
        self,
        tokenizer: "PreTrainedTokenizer"
    ) -> Tuple[List[int], List[int]]:
-        if tokenizer.bos_token_id:
+        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:
-            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(
        self,
        tokenizer: "PreTrainedTokenizer",
-        prefix: List[Union[str, Dict[str, str]]],
+        system: str,
        history: List[Tuple[str, str]]
    ) -> List[Tuple[List[int], List[int]]]:
        r"""
@@ -101,8 +106,9 @@ class Template:
        encoded_pairs = []
        for turn_idx, (query, resp) in enumerate(history):
            if turn_idx == 0:
-                if prefix: # has prefix
-                    prefix_ids = bos_ids + self._convert_inputs_to_ids(tokenizer, context=prefix) + sep_ids
+                prefix_ids = self._convert_inputs_to_ids(tokenizer, context=self.prefix, system=system)
+                if len(prefix_ids) != 0: # has prefix
+                    prefix_ids = bos_ids + prefix_ids + sep_ids
                else:
                    prefix_ids = bos_ids
            else:
@@ -117,8 +123,9 @@ class Template:
        self,
        tokenizer: "PreTrainedTokenizer",
        context: List[Union[str, Dict[str, str]]],
-        query: Optional[str] = "",
-        idx: Optional[str] = ""
+        system: Optional[str] = None,
+        query: Optional[str] = None,
+        idx: Optional[str] = None
    ) -> List[int]:
        r"""
        Converts context to token ids.
@@ -131,13 +138,17 @@ class Template:
        token_ids = []
        for elem in context:
            if isinstance(elem, str):
-                elem = elem.replace("{{query}}", query, 1)
-                elem = elem.replace("{{idx}}", idx, 1)
+                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
                token_ids = token_ids + tokenizer.encode(elem, **kwargs)
            elif isinstance(elem, dict):
                token_ids = token_ids + [tokenizer.convert_tokens_to_ids(elem.get("token"))]
            else:
                raise NotImplementedError
+
        return token_ids


@@ -147,7 +158,7 @@ class Llama2Template(Template):
    def _encode(
        self,
        tokenizer: "PreTrainedTokenizer",
-        prefix: List[Union[str, Dict[str, str]]],
+        system: str,
        history: List[Tuple[str, str]]
    ) -> List[Tuple[List[int], List[int]]]:
        r"""
@@ -157,10 +168,9 @@ class Llama2Template(Template):
        """
        bos_ids, eos_ids = self._get_special_ids(tokenizer)
        encoded_pairs = []
-        assert isinstance(prefix[0], str), "LLaMA-2 template only accepts list containing a single string."
        for turn_idx, (query, resp) in enumerate(history):
-            if turn_idx == 0: # llama2 template has not sep_ids
-                query = prefix[0] + query
+            if turn_idx == 0: # llama2 template has no sep_ids
+                query = self.prefix[0].replace("{{system}}", system) + query
            query_ids = self._convert_inputs_to_ids(tokenizer, context=self.prompt, query=query)
            resp_ids = self._convert_inputs_to_ids(tokenizer, context=[resp])
            encoded_pairs.append((bos_ids + query_ids, resp_ids + eos_ids))
@@ -174,17 +184,21 @@ def register_template(
    name: str,
    prefix: List[Union[str, Dict[str, str]]],
    prompt: List[Union[str, Dict[str, str]]],
+    system: str,
    sep: List[Union[str, Dict[str, str]]],
-    stop_words: List[str],
-    use_history: bool
+    stop_words: Optional[List[str]] = [],
+    use_history: Optional[bool] = True,
+    efficient_eos: Optional[bool] = False
 ) -> None:
    template_class = Llama2Template if "llama2" in name else Template
    templates[name] = template_class(
        prefix=prefix,
        prompt=prompt,
+        system=system,
        sep=sep,
        stop_words=stop_words,
-        use_history=use_history
+        use_history=use_history,
+        efficient_eos=efficient_eos
    )


@@ -192,13 +206,6 @@ 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)
-
-    if len(template.stop_words): # inplace method
-        tokenizer.eos_token = template.stop_words[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))
@@ -207,7 +214,15 @@ def get_template_and_fix_tokenizer(
        tokenizer.pad_token = tokenizer.eos_token
        logger.info("Add pad token: {}".format(tokenizer.pad_token))

-    tokenizer.add_special_tokens(dict(additional_special_tokens=template.stop_words))
+    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


@@ -220,8 +235,8 @@ register_template(
    prompt=[
        "{{query}}"
    ],
+    system="",
    sep=[],
-    stop_words=[],
    use_history=False
 )

@@ -232,17 +247,18 @@ Default template.
 register_template(
    name="default",
    prefix=[
-        "A chat between a curious user and an artificial intelligence assistant. "
-        "The assistant gives helpful, detailed, and polite answers to the user's questions."
+        "{{system}}"
    ],
    prompt=[
        "Human: {{query}}\nAssistant: "
    ],
+    system=(
+        "A chat between a curious user and an artificial intelligence assistant. "
+        "The assistant gives helpful, detailed, and polite answers to the user's questions."
+    ),
    sep=[
        "\n"
-    ],
-    stop_words=[],
-    use_history=True
+    ]
 )


@@ -254,38 +270,39 @@ Supports: https://huggingface.co/meta-llama/Llama-2-7b-chat-hf
 register_template(
    name="llama2",
    prefix=[
-        "<<SYS>>\nYou are a helpful, respectful and honest assistant. "
-        "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"
-        "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.\n<</SYS>>\n\n"
+        "<<SYS>>\n{{system}}\n<</SYS>>\n\n"
    ],
    prompt=[
        "[INST] {{query}} [/INST] "
    ],
-    sep=[],
-    stop_words=[],
-    use_history=True
+    system=(
+        "You are a helpful, respectful and honest assistant. "
+        "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\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."
+    ),
+    sep=[]
 )


 r"""
 Supports: https://github.com/ymcui/Chinese-LLaMA-Alpaca-2
+          https://huggingface.co/ziqingyang/chinese-alpaca-2-7b
 """
 register_template(
    name="llama2_zh",
    prefix=[
-        "<<SYS>>\nYou are a helpful assistant. 你是一个乐于助人的助手。\n<</SYS>>\n\n"
+        "<<SYS>>\n{{system}}\n<</SYS>>\n\n"
    ],
    prompt=[
        "[INST] {{query}} [/INST] "
    ],
-    sep=[],
-    stop_words=[],
-    use_history=True
+    system="You are a helpful assistant. 你是一个乐于助人的助手。",
+    sep=[]
 )


@@ -296,17 +313,18 @@ Supports: https://huggingface.co/tatsu-lab/alpaca-7b-wdiff
 register_template(
    name="alpaca",
    prefix=[
-        "Below is an instruction that describes a task. "
-        "Write a response that appropriately completes the request."
+        "{{system}}"
    ],
    prompt=[
        "### Instruction:\n{{query}}\n\n### Response:\n"
    ],
+    system=(
+        "Below is an instruction that describes a task. "
+        "Write a response that appropriately completes the request."
+    ),
    sep=[
        "\n\n"
-    ],
-    stop_words=[],
-    use_history=True
+    ]
 )


@@ -317,15 +335,16 @@ Supports: https://huggingface.co/lmsys/vicuna-7b-delta-v1.1
 register_template(
    name="vicuna",
    prefix=[
-        "A chat between a curious user and an artificial intelligence assistant. "
-        "The assistant gives helpful, detailed, and polite answers to the user's questions."
+        "{{system}}"
    ],
    prompt=[
        "USER: {{query}} ASSISTANT: "
    ],
-    sep=[],
-    stop_words=[],
-    use_history=True
+    system=(
+        "A chat between a curious user and an artificial intelligence assistant. "
+        "The assistant gives helpful, detailed, and polite answers to the user's questions."
+    ),
+    sep=[]
 )


@@ -334,15 +353,16 @@ Supports: https://huggingface.co/BelleGroup/BELLE-LLaMA-EXT-13B
 """
 register_template(
    name="belle",
-    prefix=[],
+    prefix=[
+        "{{system}}"
+    ],
    prompt=[
        "Human: {{query}}\n\nBelle: "
    ],
+    system="",
    sep=[
        "\n\n"
-    ],
-    stop_words=[],
-    use_history=True
+    ]
 )


@@ -351,15 +371,16 @@ Supports: https://github.com/CVI-SZU/Linly
 """
 register_template(
    name="linly",
-    prefix=[],
+    prefix=[
+        "{{system}}"
+    ],
    prompt=[
        "User: {{query}}\nBot: "
    ],
+    system="",
    sep=[
        "\n"
-    ],
-    stop_words=[],
-    use_history=True
+    ]
 )


@@ -368,15 +389,16 @@ Supports: https://github.com/Neutralzz/BiLLa
 """
 register_template(
    name="billa",
-    prefix=[],
+    prefix=[
+        "{{system}}"
+    ],
    prompt=[
        "Human: {{query}}\nAssistant: "
    ],
+    system="",
    sep=[
        "\n"
-    ],
-    stop_words=[],
-    use_history=True
+    ]
 )


@@ -385,18 +407,19 @@ Supports: https://huggingface.co/IDEA-CCNL/Ziya-LLaMA-13B-v1
 """
 register_template(
    name="ziya",
-    prefix=[],
+    prefix=[
+        "{{system}}"
+    ],
    prompt=[
        {"token": "<human>"},
        ":{{query}}\n",
        {"token": "<bot>"},
        ":"
    ],
+    system="",
    sep=[
        "\n"
-    ],
-    stop_words=[],
-    use_history=True
+    ]
 )


@@ -406,17 +429,18 @@ Supports: https://huggingface.co/qhduan/aquilachat-7b
 register_template(
    name="aquila",
    prefix=[
-        "A chat between a curious human and an artificial intelligence assistant. "
-        "The assistant gives helpful, detailed, and polite answers to the human's questions."
+        "{{system}}"
    ],
    prompt=[
        "Human: {{query}}###Assistant: "
    ],
+    system=(
+        "A chat between a curious human and an artificial intelligence assistant. "
+        "The assistant gives helpful, detailed, and polite answers to the human's questions."
+    ),
    sep=[
        "###"
-    ],
-    stop_words=[],
-    use_history=True
+    ]
 )


@@ -425,19 +449,23 @@ Supports: https://huggingface.co/internlm/internlm-chat-7b
 """
 register_template(
    name="intern",
-    prefix=[],
+    prefix=[
+        "{{system}}"
+    ],
    prompt=[
        "<|User|>:{{query}}",
        {"token": "<eoh>"},
        "\n<|Bot|>:"
    ],
+    system="",
    sep=[
+        {"token": "<eoa>"},
        "\n"
    ],
    stop_words=[
        "<eoa>"
    ],
-    use_history=True
+    efficient_eos=True
 )


@@ -447,17 +475,36 @@ Supports: https://huggingface.co/baichuan-inc/Baichuan-13B-Chat
 register_template(
    name="baichuan",
    prefix=[
-        {"token": "<reserved_102>"} # user token (a little difference in the first turn)
+        "{{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}}"
    ],
-    use_history=True
+    prompt=[
+        {"token": "<reserved_106>"}, # user token
+        "{{query}}",
+        {"token": "<reserved_107>"}  # assistant token
+    ],
+    system="",
+    sep=[],
+    efficient_eos=True
 )


@@ -469,8 +516,7 @@ register_template(
    name="starchat",
    prefix=[
        {"token": "<|system|>"},
-        "\n",
-        {"token": "<|end|>"}
+        "\n{{system}}",
    ],
    prompt=[
        {"token": "<|user|>"},
@@ -479,13 +525,15 @@ register_template(
        "\n",
        {"token": "<|assistant|>"}
    ],
+    system="",
    sep=[
+        {"token": "<|end|>"},
        "\n"
    ],
    stop_words=[
        "<|end|>"
    ],
-    use_history=True
+    efficient_eos=True
 )


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


@@ -524,14 +573,31 @@ register_template(
    name="chatglm2",
    prefix=[
        {"token": "[gMASK]"},
-        {"token": "sop"}
+        {"token": "sop"},
+        "{{system}}"
    ],
    prompt=[
        "[Round {{idx}}]\n\n问：{{query}}\n\n答："
    ],
+    system="",
    sep=[
        "\n\n"
    ],
-    stop_words=[],
-    use_history=True
+    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
@@ -10,24 +10,24 @@ class DatasetAttr:
    load_from: str
    dataset_name: Optional[str] = None
    dataset_sha1: Optional[str] = None
-    source_prefix: 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(
@@ -86,9 +86,9 @@ class DataArguments:
        default=True,
        metadata={"help": "Whether to ignore the tokens corresponding to padded labels in the loss computation or not."}
    )
-    source_prefix: Optional[str] = field(
+    system_prompt: Optional[str] = field(
        default=None,
-        metadata={"help": "A prefix to add before every source text. Use `|` to separate multiple prefixes in training."}
+        metadata={"help": "System prompt to add before the user query. Use `|` to separate multiple prompts in training."}
    )
    val_size: Optional[float] = field(
        default=0,
@@ -100,12 +100,9 @@ class DataArguments:
        with open(os.path.join(self.dataset_dir, "dataset_info.json"), "r") as f:
            dataset_info = json.load(f)

-        if self.source_prefix is not None:
-            prefix_list = self.source_prefix.split("|")
-            prefix_list = prefix_list * len(dataset_names) if len(prefix_list) == 1 else prefix_list
-            assert len(prefix_list) == len(dataset_names), "The number of prefixes should be either identical with datasets or 1."
-        else:
-            prefix_list = [None] * len(dataset_names)
+        prompt_list = self.system_prompt.split("|") if self.system_prompt else [None]
+        prompt_list = prompt_list * (len(dataset_names) // len(prompt_list))
+        assert len(prompt_list) == len(dataset_names), "Number of system prompts should be equal to datasets or 1."

        if self.interleave_probs is not None:
            self.interleave_probs = [float(prob.strip()) for prob in self.interleave_probs.split(",")]
@@ -126,12 +123,12 @@ class DataArguments:
                    dataset_sha1=dataset_info[name].get("file_sha1", None)
                )

-            dataset_attr.source_prefix = prefix_list[i]
-
            if "columns" in dataset_info[name]:
                dataset_attr.prompt = dataset_info[name]["columns"].get("prompt", None)
                dataset_attr.query = dataset_info[name]["columns"].get("query", None)
                dataset_attr.response = dataset_info[name]["columns"].get("response", None)
                dataset_attr.history = dataset_info[name]["columns"].get("history", None)

+            dataset_attr.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
@@ -8,7 +8,7 @@ class FinetuningArguments:
    r"""
    Arguments pertaining to which techniques we are going to fine-tuning with.
    """
-    finetuning_type: Optional[Literal["none", "freeze", "lora", "full"]] = field(
+    finetuning_type: Optional[Literal["lora", "freeze", "full", "none"]] = field(
        default="lora",
        metadata={"help": "Which fine-tuning method to use."}
    )
@@ -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."}
@@ -49,10 +50,10 @@ class FinetuningArguments:
        metadata={"help": "Dropout rate for the LoRA fine-tuning."}
    )
    lora_target: Optional[str] = field(
-        default="q_proj,v_proj",
+        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."}
@@ -61,6 +62,10 @@ class FinetuningArguments:
        default=True,
        metadata={"help": "Whether to resume training from the last LoRA weights or create new weights after merging them."}
    )
+    ppo_score_norm: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Use score normalization in PPO Training."}
+    )
    dpo_beta: Optional[float] = field(
        default=0.1,
        metadata={"help": "The beta parameter for the DPO loss."}
@@ -77,7 +82,7 @@ class FinetuningArguments:

        self.trainable_layers = ["{:d}.{}".format(idx, self.name_module_trainable) for idx in trainable_layer_ids]

-        assert self.finetuning_type in ["none", "freeze", "lora", "full"], "Invalid fine-tuning method."
+        assert self.finetuning_type in ["lora", "freeze", "full", "none"], "Invalid fine-tuning method."

    def save_to_json(self, json_path: str):
        r"""Saves the content of this instance in JSON format inside `json_path`."""
--- a/src/llmtuner/hparams/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."}
@@ -43,9 +39,13 @@ class ModelArguments:
        default=True,
        metadata={"help": "Whether to use double quantization in int4 training or not."}
    )
-    compute_dtype: Optional[torch.dtype] = field(
+    rope_scaling: Optional[Literal["linear", "dynamic"]] = field(
        default=None,
-        metadata={"help": "Used in quantization configs. Do not specify this argument manually."}
+        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,
@@ -63,8 +63,19 @@ class ModelArguments:
        default=None,
        metadata={"help": "Auth token to log in with Hugging Face Hub."}
    )
+    compute_dtype: Optional[torch.dtype] = field(
+        default=None,
+        metadata={"help": "Used in quantization configs. Do not specify this argument manually."}
+    )
+    model_max_length: Optional[int] = field(
+        default=None,
+        metadata={"help": "Used in rope scaling. Do not specify this argument manually."}
+    )

    def __post_init__(self):
+        if self.compute_dtype is not None or self.model_max_length is not None:
+            raise ValueError("These arguments cannot be specified.")
+
        if self.checkpoint_dir is not None: # support merging multiple lora weights
            self.checkpoint_dir = [cd.strip() for cd in self.checkpoint_dir.split(",")]

--- 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
@@ -1,7 +1,10 @@
 import os
+import math
 import torch
+from types import MethodType
 from typing import TYPE_CHECKING, Literal, Optional, Tuple

+import transformers
 from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
@@ -13,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
@@ -30,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(
@@ -62,20 +70,84 @@ 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 model_args.checkpoint_dir is not None and finetuning_args.finetuning_type == "full":
+    # 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)
-    is_mergeable = True
+
+    # 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:
+            setattr(config, "fp16", True)
+
+    # Set RoPE scaling
+    if model_args.rope_scaling is not None:
+        if hasattr(config, "use_dynamic_ntk"): # for Qwen models
+            if is_trainable:
+                logger.warning("Qwen model does not support RoPE scaling in training.")
+            else:
+                setattr(config, "use_dynamic_ntk", True)
+                setattr(config, "use_logn_attn", True)
+                logger.info("Using dynamic NTK scaling.")
+
+        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"
+                    )
+
+                current_max_length = getattr(config, "max_position_embeddings", None)
+                if current_max_length and model_args.model_max_length > current_max_length:
+                    scaling_factor = float(math.ceil(model_args.model_max_length / current_max_length))
+                else:
+                    logger.warning("Input length is smaller than max length. Consider increase input length.")
+                    scaling_factor = 1.0
+            else:
+                scaling_factor = 2.0
+
+            setattr(config, "rope_scaling", {"type": model_args.rope_scaling, "factor": scaling_factor})
+            logger.info("Using {} scaling strategy and setting scaling factor to {}".format(
+                model_args.rope_scaling, scaling_factor
+            ))
+
+        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:
+        if is_deepspeed_zero3_enabled():
+            raise ValueError("DeepSpeed ZeRO-3 is incompatible with quantization.")
+
        if model_args.quantization_bit == 8:
            require_version("bitsandbytes>=0.37.0", "To fix: pip install bitsandbytes>=0.37.0")
            config_kwargs["load_in_8bit"] = True
@@ -92,18 +164,26 @@ def load_model_and_tokenizer(
            )

        is_mergeable = False
-        config_kwargs["device_map"] = {"": int(os.environ.get("LOCAL_RANK", "0"))}
+        config_kwargs["device_map"] = {"": int(os.environ.get("LOCAL_RANK", "0"))} if is_trainable else "auto"
        logger.info("Quantizing model to {} bit.".format(model_args.quantization_bit))

-    # Load and prepare pretrained models (without valuehead).
+    # Load and prepare pre-trained models (without valuehead).
    model = AutoModelForCausalLM.from_pretrained(
        model_to_load,
        config=config,
-        torch_dtype=torch.bfloat16 if model_args.compute_dtype == torch.bfloat16 else torch.float16,
+        torch_dtype=model_args.compute_dtype,
        low_cpu_mem_usage=(not is_deepspeed_zero3_enabled()),
        **config_kwargs
    )

+    # Disable custom generate method (for Qwen)
+    if "GenerationMixin" not in str(model.generate.__func__):
+        model.generate = MethodType(PreTrainedModel.generate, model)
+
+    # Fix LM head (for ChatGLM2)
+    if not hasattr(model, "lm_head") and hasattr(model, "transformer"):
+        setattr(model, "lm_head", model.transformer.output_layer)
+
    # Register auto class to save the custom code files.
    if isinstance(config, PretrainedConfig) and "AutoConfig" in getattr(config, "auto_map", {}):
        config.__class__.register_for_auto_class()
@@ -115,11 +195,13 @@ 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()

-    if stage == "rm" or stage == "ppo": # add value head
-        model: "AutoModelForCausalLMWithValueHead" = AutoModelForCausalLMWithValueHead.from_pretrained(model)
+    # 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.")
            if load_valuehead_params(model, model_args.checkpoint_dir[-1]):
@@ -133,9 +215,11 @@ def load_model_and_tokenizer(
            model.pretrained_model.load_adapter(model_args.reward_model, "reward", is_trainable=False)
            assert load_valuehead_params(model, model_args.reward_model), "Reward model is not correctly loaded."

+    # Prepare model for inference
    if not is_trainable:
        model.requires_grad_(False) # fix all model params
-        model = model.half() if model_args.quantization_bit is None else model # cast from fp32 to fp16
+        infer_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 # detect cuda capability
+        model = model.to(infer_dtype) if model_args.quantization_bit is None else model

    trainable_params, all_param = count_parameters(model)
    logger.info("trainable params: {:d} || all params: {:d} || trainable%: {:.4f}".format(
--- a/src/llmtuner/tuner/core/parser.py
+++ b/src/llmtuner/tuner/core/parser.py
@@ -5,6 +5,8 @@ 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
 from llmtuner.hparams import (
@@ -94,33 +96,50 @@ 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.")

-    if training_args.do_train and training_args.predict_with_generate:
-        raise ValueError("`predict_with_generate` cannot be set as True while training.")
-
    if general_args.stage == "sft" and training_args.do_predict and not training_args.predict_with_generate:
        raise ValueError("Please enable `predict_with_generate` to save model predictions.")

    if general_args.stage in ["rm", "ppo"] and finetuning_args.finetuning_type != "lora":
-        raise ValueError("RM and PPO training can only be performed with the LoRA method.")
+        raise ValueError("RM and PPO stages can only be performed with the LoRA method.")
+
+    if general_args.stage in ["rm", "ppo"] and training_args.resume_from_checkpoint is not None:
+        raise ValueError("RM and PPO stages do not support `resume_from_checkpoint`.")

    if general_args.stage in ["ppo", "dpo"] and not training_args.do_train:
-        raise ValueError("PPO and DPO stage can only be performed at training.")
+        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 training_args.max_steps == -1 and data_args.streaming:
-        raise ValueError("Please specify `max_steps` in streaming mode.")
+    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.")

+    if training_args.max_steps == -1 and data_args.streaming:
+        raise ValueError("Please specify `max_steps` in streaming mode.")
+
    if data_args.val_size > 1e-6 and data_args.val_size < 1 and data_args.streaming:
        raise ValueError("Streaming mode should have an integer val size.")

+    if training_args.do_train and training_args.predict_with_generate:
+        raise ValueError("`predict_with_generate` cannot be set as True while training.")
+
+    if training_args.do_train and finetuning_args.finetuning_type == "lora" and finetuning_args.lora_target is None:
+        raise ValueError("Please specify `lora_target` in LoRA training.")
+
    if model_args.quantization_bit is not None and finetuning_args.finetuning_type != "lora":
        raise ValueError("Quantization is only compatible with the LoRA method.")

@@ -134,35 +153,58 @@ def get_train_args(
    if model_args.quantization_bit is not None and (not training_args.do_train):
        logger.warning("Evaluating model in 4/8-bit mode may cause lower scores.")

-    if training_args.do_train and (not training_args.fp16):
-        logger.warning("We recommend enable fp16 mixed precision training.")
+    if training_args.do_train and (not training_args.fp16) and (not training_args.bf16):
+        logger.warning("We recommend enable mixed precision training.")

+    # postprocess data_args
+    if data_args.max_samples is not None and data_args.streaming:
+        logger.warning("`max_samples` is incompatible with `streaming`. Disabling max_samples.")
+        data_args.max_samples = None
+
+    # postprocess training_args
    if (
        training_args.local_rank != -1
        and training_args.ddp_find_unused_parameters is None
        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
+        training_args_dict = training_args.to_dict()
+        training_args_dict.update(dict(ddp_find_unused_parameters=False))
+        training_args = Seq2SeqTrainingArguments(**training_args_dict)

-    if data_args.max_samples is not None and data_args.streaming:
-        logger.warning("`max_samples` is incompatible with `streaming`. Disabling max_samples.")
-        data_args.max_samples = None
+    if (
+        training_args.resume_from_checkpoint is None
+        and training_args.do_train
+        and os.path.isdir(training_args.output_dir)
+        and not training_args.overwrite_output_dir
+    ):
+        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`.")

-    training_args.optim = "adamw_torch" if training_args.optim == "adamw_hf" else training_args.optim # suppress warning
+        if last_checkpoint is not None:
+            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."
+            )

-    if model_args.quantization_bit is not None:
-        if training_args.fp16:
-            model_args.compute_dtype = torch.float16
-        elif training_args.bf16:
-            model_args.compute_dtype = torch.bfloat16
-        else:
-            model_args.compute_dtype = torch.float32
+    # postprocess model_args
+    if training_args.bf16:
+        if not torch.cuda.is_bf16_supported():
+            raise ValueError("Current device does not support bf16 training.")
+        model_args.compute_dtype = torch.bfloat16
+    else:
+        model_args.compute_dtype = torch.float16
+
+    model_args.model_max_length = data_args.max_source_length + data_args.max_target_length

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

@@ -182,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,68 +1,62 @@
 import torch
 from collections import defaultdict
-from peft import PeftModel
 from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
-from transformers import Trainer
+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:
-            if hasattr(self, "accelerator"):
-                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:
-                raise AttributeError("Please update `transformers`.")
+                self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)

    def concatenated_forward(
        self,
        model: Optional[torch.nn.Module] = None,
        batch: Optional[Dict[str, torch.Tensor]] = None
    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
-        unwrapped_model: "PreTrainedModel" = self.accelerator.unwrap_model(self.model)
-        if not torch.is_grad_enabled():
-            unwrapped_model.gradient_checkpointing_disable()
+        batch_copied = BatchEncoding({k: v.detach().clone() for k, v in batch.items()}) # avoid error

-        if model is None and isinstance(unwrapped_model, PeftModel): # peft model has no ref_model
-            with unwrapped_model.disable_adapter():
-                all_logits: torch.Tensor = self.model(
-                    batch["input_ids"],
-                    attention_mask=batch["attention_mask"],
-                    return_dict=True
-                ).logits.to(torch.float32)
-        else:
-            all_logits: torch.Tensor = model(
-                batch["input_ids"],
-                attention_mask=batch["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,
@@ -50,7 +50,7 @@ def run_dpo(

    # Training
    if training_args.do_train:
-        train_result = trainer.train()
+        train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
        trainer.log_metrics("train", train_result.metrics)
        trainer.save_metrics("train", train_result.metrics)
        trainer.save_state()
--- 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
+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, get_stopping_criteria
-from llmtuner.tuner.core.trainer import PeftTrainer
+from llmtuner.extras.misc import AverageMeter, count_parameters, get_logits_processor
 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,16 +31,19 @@ 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()

@@ -73,9 +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["logits_processor"] = get_logits_processor()
-        gen_kwargs["stopping_criteria"] = get_stopping_criteria(self.tokenizer.additional_special_tokens_ids)
+        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)
@@ -93,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))

@@ -125,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
@@ -134,36 +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
+        )

-        self.model, layer_norm_params = cast_layernorm_dtype(self.model)
+        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)
-        self.model, _ = cast_layernorm_dtype(self.model, layer_norm_params)
-
-        # 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

@@ -181,11 +204,94 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
        """
        replace_model(unwrapped_model, target="reward")
        batch = self.prepare_model_inputs(queries, responses)
-        _, _, values = self.model(**batch, output_hidden_states=True, return_dict=True)
-        rewards = [reward for reward in values[:, -1].float().detach().cpu()] # use fp32 type
+
+        with torch.cuda.amp.autocast(dtype=self.compute_dtype): # support bf16
+            _, _, values = self.model(**batch, output_hidden_states=True, return_dict=True)
+
+        if values.size(0) != batch["input_ids"].size(0): # adapt to chatglm2
+            values = torch.transpose(values, 0, 1)
+
+        rewards = []
+        for i in range(values.size(0)):
+            end_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

+    @PPODecorators.empty_cuda_cache()
+    def batched_forward_pass(
+        self,
+        model: "AutoModelForCausalLMWithValueHead",
+        queries: torch.Tensor,
+        responses: torch.Tensor,
+        model_inputs: dict,
+        return_logits: Optional[bool] = False,
+        response_masks: Optional[torch.Tensor] = None
+    ):
+        r"""
+        Calculates model outputs in multiple batches.
+
+        Subclass and override to inject custom behavior.
+        """
+        bs = len(queries)
+        fbs = self.config.mini_batch_size
+        all_logprobs = []
+        all_logits = []
+        all_masks = []
+        all_values = []
+
+        for i in range(math.ceil(bs / fbs)):
+            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"]
+
+            with torch.cuda.amp.autocast(dtype=self.compute_dtype): # support bf16
+                logits, _, values = model(**input_kwargs)
+
+            if values.size(0) != input_ids.size(0): # adapt to chatglm2
+                values = torch.transpose(values, 0, 1)
+
+            logprobs = logprobs_from_logits(logits[:, :-1, :], input_ids[:, 1:])
+            masks = torch.zeros_like(attention_mask)
+            masks[:, :-1] = attention_mask[:, 1:]
+
+            for j in range(len(query_batch)):
+                start = len(query_batch[j]) - 1
+                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)
+            else:
+                del logits
+
+            all_values.append(values)
+            all_logprobs.append(logprobs)
+            all_masks.append(masks)
+
+        return (
+            torch.cat(all_logprobs),
+            torch.cat(all_logits)[:, :-1] if return_logits else None,
+            torch.cat(all_values)[:, :-1],
+            torch.cat(all_masks)[:, :-1],
+        )
+
    def save_model(self, output_dir: Optional[str] = None) -> None:
        r"""
        Saves model checkpoint.
--- a/src/llmtuner/tuner/ppo/utils.py
+++ b/src/llmtuner/tuner/ppo/utils.py
@@ -22,18 +22,19 @@ def replace_model(model: "AutoModelForCausalLMWithValueHead", target: Literal["d

 def cast_layernorm_dtype(
    model: "AutoModelForCausalLMWithValueHead",
-    layer_norm_names: List[str] = LAYERNORM_NAMES,
-    layer_norm_params: Optional[Dict[str, torch.Tensor]] = None
+    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 not None:
-                param.data = layer_norm_params[name] # restore float32 weights
-            else:
+            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(torch.float16)
+                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,13 +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 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
@@ -28,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,
@@ -37,9 +40,15 @@ def run_ppo(
        batch_size=training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps,
        gradient_accumulation_steps=training_args.gradient_accumulation_steps,
        ppo_epochs=1,
-        max_grad_norm=training_args.max_grad_norm
+        max_grad_norm=training_args.max_grad_norm,
+        seed=training_args.seed,
+        optimize_cuda_cache=True
    )

+    if finetuning_args.ppo_score_norm:
+        ppo_config.use_score_scaling = True
+        ppo_config.use_score_norm = True
+
    optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=training_args.learning_rate)
    total_train_batch_size = (
        training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps * training_args.world_size
@@ -53,11 +62,11 @@ 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,
        ref_model=None,
--- 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,
@@ -39,7 +37,7 @@ def run_pt(

    # Training
    if training_args.do_train:
-        train_result = trainer.train()
+        train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
        trainer.log_metrics("train", train_result.metrics)
        trainer.save_metrics("train", train_result.metrics)
        trainer.save_state()
@@ -56,6 +54,5 @@ def run_pt(
            perplexity = float("inf")

        metrics["perplexity"] = perplexity
-
        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)
--- 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,19 +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)
-        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
+        if values.size(0) != inputs["input_ids"].size(0): # adapt to chatglm2
+            values = torch.transpose(values, 0, 1)
+
+        # 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,
@@ -61,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/metric.py
+++ b/src/llmtuner/tuner/sft/metric.py
@@ -25,7 +25,7 @@ class ComputeMetrics:
        Uses the model predictions to compute metrics.
        """
        preds, labels = eval_preds
-        score_dict = {"accuracy": [], "rouge-1": [], "rouge-2": [], "rouge-l": [], "bleu-4": []}
+        score_dict = {"rouge-1": [], "rouge-2": [], "rouge-l": [], "bleu-4": []}

        preds = np.where(preds != IGNORE_INDEX, preds, self.tokenizer.pad_token_id)
        labels = np.where(labels != IGNORE_INDEX, labels, self.tokenizer.pad_token_id)
@@ -49,6 +49,5 @@ class ComputeMetrics:

            bleu_score = sentence_bleu([list(label)], list(pred), smoothing_function=SmoothingFunction().method3)
            score_dict["bleu-4"].append(round(bleu_score * 100, 4))
-            score_dict["accuracy"].append(float(len(label) != 0 and pred[:len(label)] == label))

        return {k: float(np.mean(v)) for k, v in score_dict.items()}
--- 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,28 +33,31 @@ 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
        )
-        generated_tokens = (
-            generated_tokens[:, max(prompt_len, label_len):] if generated_tokens is not None else None
-        )
+        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)
+        return loss, generated_tokens, labels

    def _pad_tensors_to_target_len(
        self,
@@ -64,19 +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:
-                if self.model.config.pad_token_id is not None:
-                    pad_token_id = self.model.config.pad_token_id
-                else:
-                    raise ValueError("Pad_token_id must be set in the configuration of the model.")
-
+        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,18 +1,18 @@
 # 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
-from llmtuner.extras.misc import get_logits_processor, get_stopping_criteria
+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,12 +57,13 @@ def run_sft(

    # Keyword arguments for `model.generate`
    gen_kwargs = generating_args.to_dict()
+    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()
-    gen_kwargs["stopping_criteria"] = get_stopping_criteria(tokenizer.additional_special_tokens_ids)

    # Training
    if training_args.do_train:
-        train_result = trainer.train()
+        train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
        trainer.log_metrics("train", train_result.metrics)
        trainer.save_metrics("train", train_result.metrics)
        trainer.save_state()
--- a/src/llmtuner/tuner/tune.py
+++ b/src/llmtuner/tuner/tune.py
@@ -35,7 +35,7 @@ def run_exp(args: Optional[Dict[str, Any]] = None, callbacks: Optional[List["Tra


 def export_model(args: Optional[Dict[str, Any]] = None, max_shard_size: Optional[str] = "10GB"):
-    model_args, _, training_args, finetuning_args, _ = get_train_args(args)
+    model_args, _, training_args, finetuning_args, _, _ = get_train_args(args)
    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args)
    model.save_pretrained(training_args.output_dir, max_shard_size=max_shard_size)
    try:
--- a/src/llmtuner/webui/chat.py
+++ b/src/llmtuner/webui/chat.py
@@ -26,7 +26,7 @@ class WebChatModel(ChatModel):
        finetuning_type: str,
        quantization_bit: str,
        template: str,
-        source_prefix: str
+        system_prompt: str
    ):
        if self.model is not None:
            yield ALERTS["err_exists"][lang]
@@ -53,9 +53,9 @@ 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 else None,
+            quantization_bit=int(quantization_bit) if quantization_bit and quantization_bit != "None" else None,
            template=template,
-            source_prefix=source_prefix
+            system_prompt=system_prompt
        )
        super().__init__(args)

@@ -73,7 +73,7 @@ class WebChatModel(ChatModel):
        chatbot: List[Tuple[str, str]],
        query: str,
        history: List[Tuple[str, str]],
-        prefix: str,
+        system: str,
        max_new_tokens: int,
        top_p: float,
        temperature: float
@@ -81,7 +81,7 @@ class WebChatModel(ChatModel):
        chatbot.append([query, ""])
        response = ""
        for new_text in self.stream_chat(
-            query, history, prefix, max_new_tokens=max_new_tokens, top_p=top_p, temperature=temperature
+            query, history, system, max_new_tokens=max_new_tokens, top_p=top_p, temperature=temperature
        ):
            response += new_text
            response = self.postprocess(response)
--- 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 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,14 +29,16 @@ def load_config() -> Dict[str, Any]:
        with open(get_config_path(), "r", encoding="utf-8") as f:
            return json.load(f)
    except:
-        return {"last_model": "", "path_dict": {}}
+        return {"lang": None, "last_model": None, "path_dict": {}, "cache_dir": None}


-def save_config(model_name: str, model_path: str) -> None:
+def save_config(lang: str, model_name: str, model_path: str) -> None:
    os.makedirs(DEFAULT_CACHE_DIR, exist_ok=True)
    user_config = load_config()
-    user_config["last_model"] = model_name
-    user_config["path_dict"][model_name] = model_path
+    user_config["lang"] = lang or user_config["lang"]
+    if model_name:
+        user_config["last_model"] = model_name
+        user_config["path_dict"][model_name] = model_path
    with open(get_config_path(), "w", encoding="utf-8") as f:
        json.dump(user_config, f, indent=2, ensure_ascii=False)

@@ -46,9 +48,15 @@ def get_model_path(model_name: str) -> str:
    return user_config["path_dict"].get(model_name, SUPPORTED_MODELS.get(model_name, ""))


+def get_template(model_name: str) -> str:
+    if model_name.endswith("Chat") and model_name.split("-")[0] in DEFAULT_TEMPLATE:
+        return DEFAULT_TEMPLATE[model_name.split("-")[0]]
+    return "default"
+
+
 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 (
@@ -70,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/init.py
+++ b/src/llmtuner/webui/components/init.py
@@ -1,5 +1,5 @@
 from llmtuner.webui.components.top import create_top
-from llmtuner.webui.components.sft import create_sft_tab
+from llmtuner.webui.components.train import create_train_tab
 from llmtuner.webui.components.eval import create_eval_tab
 from llmtuner.webui.components.infer import create_infer_tab
 from llmtuner.webui.components.export import create_export_tab
--- a/src/llmtuner/webui/components/chatbot.py
+++ b/src/llmtuner/webui/components/chatbot.py
@@ -17,7 +17,7 @@ def create_chat_box(

        with gr.Row():
            with gr.Column(scale=4):
-                prefix = gr.Textbox(show_label=False)
+                system = gr.Textbox(show_label=False)
                query = gr.Textbox(show_label=False, lines=8)
                submit_btn = gr.Button(variant="primary")

@@ -31,7 +31,7 @@ def create_chat_box(

    submit_btn.click(
        chat_model.predict,
-        [chatbot, query, history, prefix, max_new_tokens, top_p, temperature],
+        [chatbot, query, history, system, max_new_tokens, top_p, temperature],
        [chatbot, history],
        show_progress=True
    ).then(
@@ -41,7 +41,7 @@ def create_chat_box(
    clear_btn.click(lambda: ([], []), outputs=[chatbot, history], show_progress=True)

    return chat_box, chatbot, history, dict(
-        prefix=prefix,
+        system=system,
        query=query,
        submit_btn=submit_btn,
        clear_btn=clear_btn,
--- a/src/llmtuner/webui/components/eval.py
+++ b/src/llmtuner/webui/components/eval.py
@@ -14,13 +14,13 @@ def create_eval_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dict
    with gr.Row():
        dataset_dir = gr.Textbox(value=DEFAULT_DATA_DIR, scale=2)
        dataset = gr.Dropdown(multiselect=True, scale=4)
-        preview_btn = gr.Button(interactive=False, scale=1)
+        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])
-    dataset.change(can_preview, [dataset_dir, dataset], [preview_btn])
-    preview_btn.click(
+    dataset.change(can_preview, [dataset_dir, dataset], [data_preview_btn])
+    data_preview_btn.click(
        get_preview,
        [dataset_dir, dataset],
        [preview_count, preview_samples, preview_box],
@@ -35,6 +35,7 @@ def create_eval_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dict
        predict = gr.Checkbox(value=True)

    with gr.Row():
+        cmd_preview_btn = gr.Button()
        start_btn = gr.Button()
        stop_btn = gr.Button()

@@ -44,35 +45,36 @@ def create_eval_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dict
    with gr.Box():
        output_box = gr.Markdown()

-    start_btn.click(
-        runner.run_eval,
-        [
-            top_elems["lang"],
-            top_elems["model_name"],
-            top_elems["checkpoints"],
-            top_elems["finetuning_type"],
-            top_elems["quantization_bit"],
-            top_elems["template"],
-            top_elems["source_prefix"],
-            dataset_dir,
-            dataset,
-            max_source_length,
-            max_target_length,
-            max_samples,
-            batch_size,
-            predict
-        ],
-        [
-            output_box,
-            process_bar
-        ]
-    )
+    input_components = [
+        top_elems["lang"],
+        top_elems["model_name"],
+        top_elems["checkpoints"],
+        top_elems["finetuning_type"],
+        top_elems["quantization_bit"],
+        top_elems["template"],
+        top_elems["system_prompt"],
+        dataset_dir,
+        dataset,
+        max_source_length,
+        max_target_length,
+        max_samples,
+        batch_size,
+        predict
+    ]
+
+    output_components = [
+        output_box,
+        process_bar
+    ]
+
+    cmd_preview_btn.click(runner.preview_eval, input_components, output_components)
+    start_btn.click(runner.run_eval, input_components, output_components)
    stop_btn.click(runner.set_abort, queue=False)

    return dict(
        dataset_dir=dataset_dir,
        dataset=dataset,
-        preview_btn=preview_btn,
+        data_preview_btn=data_preview_btn,
        preview_count=preview_count,
        preview_samples=preview_samples,
        close_btn=close_btn,
@@ -81,6 +83,7 @@ def create_eval_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dict
        max_samples=max_samples,
        batch_size=batch_size,
        predict=predict,
+        cmd_preview_btn=cmd_preview_btn,
        start_btn=start_btn,
        stop_btn=stop_btn,
        output_box=output_box
--- 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(
@@ -28,7 +28,7 @@ def create_infer_tab(top_elems: Dict[str, "Component"]) -> Dict[str, "Component"
            top_elems["finetuning_type"],
            top_elems["quantization_bit"],
            top_elems["template"],
-            top_elems["source_prefix"]
+            top_elems["system_prompt"]
        ],
        [info_box]
    ).then(
--- a/src/llmtuner/webui/components/top.py
+++ b/src/llmtuner/webui/components/top.py
@@ -4,7 +4,7 @@ import gradio as gr

 from llmtuner.extras.constants import METHODS, SUPPORTED_MODELS
 from llmtuner.extras.template import templates
-from llmtuner.webui.common import list_checkpoint, get_model_path, save_config
+from llmtuner.webui.common import list_checkpoint, get_model_path, get_template, save_config
 from llmtuner.webui.utils import can_quantize

 if TYPE_CHECKING:
@@ -15,27 +15,32 @@ def create_top() -> Dict[str, "Component"]:
    available_models = list(SUPPORTED_MODELS.keys()) + ["Custom"]

    with gr.Row():
-        lang = gr.Dropdown(choices=["en", "zh"], value="en", scale=1)
+        lang = gr.Dropdown(choices=["en", "zh"], scale=1)
        model_name = gr.Dropdown(choices=available_models, scale=3)
        model_path = gr.Textbox(scale=3)

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

    with gr.Accordion(label="Advanced config", open=False) as advanced_tab:
        with gr.Row():
-            quantization_bit = gr.Dropdown([8, 4], scale=1)
-            template = gr.Dropdown(value="default", choices=list(templates.keys()), scale=1)
-            source_prefix = gr.Textbox(scale=2)
+            quantization_bit = gr.Dropdown(choices=["None", "8", "4"], value="None", scale=1)
+            template = gr.Dropdown(choices=list(templates.keys()), value="default", scale=1)
+            system_prompt = gr.Textbox(scale=2)
+
+    lang.change(save_config, [lang, model_name, model_path])

    model_name.change(
        list_checkpoint, [model_name, finetuning_type], [checkpoints]
    ).then(
        get_model_path, [model_name], [model_path]
+    ).then(
+        get_template, [model_name], [template]
    ) # do not save config since the below line will save
-    model_path.change(save_config, [model_name, model_path])
+
+    model_path.change(save_config, [lang, model_name, model_path])

    finetuning_type.change(
        list_checkpoint, [model_name, finetuning_type], [checkpoints]
@@ -43,7 +48,9 @@ def create_top() -> Dict[str, "Component"]:
        can_quantize, [finetuning_type], [quantization_bit]
    )

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

    return dict(
        lang=lang,
@@ -55,5 +62,5 @@ def create_top() -> Dict[str, "Component"]:
        advanced_tab=advanced_tab,
        quantization_bit=quantization_bit,
        template=template,
-        source_prefix=source_prefix
+        system_prompt=system_prompt
    )
--- a/src/llmtuner/webui/components/train.py
+++ b/src/llmtuner/webui/components/train.py
@@ -3,7 +3,8 @@ from transformers.trainer_utils import SchedulerType

 import gradio as gr

-from llmtuner.webui.common import list_dataset, DEFAULT_DATA_DIR
+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

@@ -12,17 +13,21 @@ if TYPE_CHECKING:
    from llmtuner.webui.runner import Runner


-def create_sft_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dict[str, "Component"]:
+def create_train_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dict[str, "Component"]:
    with gr.Row():
+        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)
-        preview_btn = gr.Button(interactive=False, scale=1)
+        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])
-    dataset.change(can_preview, [dataset_dir, dataset], [preview_btn])
-    preview_btn.click(
+    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,
        [dataset_dir, dataset],
        [preview_count, preview_samples, preview_box],
@@ -40,7 +45,7 @@ def create_sft_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dict[
        batch_size = gr.Slider(value=4, minimum=1, maximum=512, step=1)
        gradient_accumulation_steps = gr.Slider(value=4, minimum=1, maximum=512, step=1)
        lr_scheduler_type = gr.Dropdown(
-            value="cosine", choices=[scheduler.value for scheduler in SchedulerType]
+            choices=[scheduler.value for scheduler in SchedulerType], value="cosine"
        )
        max_grad_norm = gr.Textbox(value="1.0")
        val_size = gr.Slider(value=0, minimum=0, maximum=1, step=0.001)
@@ -51,7 +56,6 @@ def create_sft_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dict[
            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():
@@ -60,7 +64,21 @@ def create_sft_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dict[
            lora_target = gr.Textbox(scale=2)
            resume_lora_training = gr.Checkbox(value=True, scale=1)

+    with gr.Accordion(label="RLHF config", open=False) as rlhf_tab:
+        with gr.Row():
+            dpo_beta = gr.Slider(value=0.1, minimum=0, maximum=1, step=0.01, scale=2)
+            reward_model = gr.Dropdown(scale=2)
+            refresh_btn = gr.Button(scale=1)
+
+    refresh_btn.click(
+        list_checkpoint,
+        [top_elems["model_name"], top_elems["finetuning_type"]],
+        [reward_model],
+        queue=False
+    )
+
    with gr.Row():
+        cmd_preview_btn = gr.Button()
        start_btn = gr.Button()
        stop_btn = gr.Button()

@@ -78,44 +96,47 @@ def create_sft_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dict[
        with gr.Column(scale=1):
            loss_viewer = gr.Plot()

-    start_btn.click(
-        runner.run_train,
-        [
-            top_elems["lang"],
-            top_elems["model_name"],
-            top_elems["checkpoints"],
-            top_elems["finetuning_type"],
-            top_elems["quantization_bit"],
-            top_elems["template"],
-            top_elems["source_prefix"],
-            dataset_dir,
-            dataset,
-            max_source_length,
-            max_target_length,
-            learning_rate,
-            num_train_epochs,
-            max_samples,
-            batch_size,
-            gradient_accumulation_steps,
-            lr_scheduler_type,
-            max_grad_norm,
-            val_size,
-            logging_steps,
-            save_steps,
-            warmup_steps,
-            compute_type,
-            padding_side,
-            lora_rank,
-            lora_dropout,
-            lora_target,
-            resume_lora_training,
-            output_dir
-        ],
-        [
-            output_box,
-            process_bar
-        ]
-    )
+    input_components = [
+        top_elems["lang"],
+        top_elems["model_name"],
+        top_elems["checkpoints"],
+        top_elems["finetuning_type"],
+        top_elems["quantization_bit"],
+        top_elems["template"],
+        top_elems["system_prompt"],
+        training_stage,
+        dataset_dir,
+        dataset,
+        max_source_length,
+        max_target_length,
+        learning_rate,
+        num_train_epochs,
+        max_samples,
+        batch_size,
+        gradient_accumulation_steps,
+        lr_scheduler_type,
+        max_grad_norm,
+        val_size,
+        logging_steps,
+        save_steps,
+        warmup_steps,
+        compute_type,
+        lora_rank,
+        lora_dropout,
+        lora_target,
+        resume_lora_training,
+        dpo_beta,
+        reward_model,
+        output_dir
+    ]
+
+    output_components = [
+        output_box,
+        process_bar
+    ]
+
+    cmd_preview_btn.click(runner.preview_train, input_components, output_components)
+    start_btn.click(runner.run_train, input_components, output_components)
    stop_btn.click(runner.set_abort, queue=False)

    process_bar.change(
@@ -123,9 +144,10 @@ def create_sft_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dict[
    )

    return dict(
+        training_stage=training_stage,
        dataset_dir=dataset_dir,
        dataset=dataset,
-        preview_btn=preview_btn,
+        data_preview_btn=data_preview_btn,
        preview_count=preview_count,
        preview_samples=preview_samples,
        close_btn=close_btn,
@@ -144,12 +166,16 @@ def create_sft_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dict[
        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,
        lora_target=lora_target,
        resume_lora_training=resume_lora_training,
+        rlhf_tab=rlhf_tab,
+        dpo_beta=dpo_beta,
+        reward_model=reward_model,
+        refresh_btn=refresh_btn,
+        cmd_preview_btn=cmd_preview_btn,
        start_btn=start_btn,
        stop_btn=stop_btn,
        output_dir=output_dir,
--- a/src/llmtuner/webui/interface.py
+++ b/src/llmtuner/webui/interface.py
@@ -3,7 +3,7 @@ from transformers.utils.versions import require_version

 from llmtuner.webui.components import (
    create_top,
-    create_sft_tab,
+    create_train_tab,
    create_eval_tab,
    create_infer_tab,
    create_export_tab,
@@ -24,8 +24,8 @@ def create_ui() -> gr.Blocks:
    with gr.Blocks(title="Web Tuner", css=CSS) as demo:
        top_elems = create_top()

-        with gr.Tab("SFT"):
-            sft_elems = create_sft_tab(top_elems, runner)
+        with gr.Tab("Train"):
+            train_elems = create_train_tab(top_elems, runner)

        with gr.Tab("Evaluate"):
            eval_elems = create_eval_tab(top_elems, runner)
@@ -36,7 +36,7 @@ def create_ui() -> gr.Blocks:
        with gr.Tab("Export"):
            export_elems = create_export_tab(top_elems)

-        elem_list = [top_elems, sft_elems, eval_elems, infer_elems, export_elems]
+        elem_list = [top_elems, train_elems, eval_elems, infer_elems, export_elems]
        manager = Manager(elem_list)

        demo.load(
@@ -67,7 +67,7 @@ def create_web_demo() -> gr.Blocks:

        demo.load(manager.gen_label, [lang], [lang] + list(chat_elems.values()))

-        lang.change(manager.gen_label, [lang], [lang] + list(chat_elems.values()), queue=False)
+        lang.select(manager.gen_label, [lang], [lang] + list(chat_elems.values()), queue=False)

    return demo

--- a/src/llmtuner/webui/locales.py
+++ b/src/llmtuner/webui/locales.py
@@ -77,7 +77,7 @@ LOCALES = {
            "info": "构建提示词时使用的模板"
        }
    },
-    "source_prefix": {
+    "system_prompt": {
        "en": {
            "label": "System prompt (optional)",
            "info": "A sequence used as the default system prompt."
@@ -87,6 +87,16 @@ LOCALES = {
            "info": "默认使用的系统提示词"
        }
    },
+    "training_stage": {
+        "en": {
+            "label": "Stage",
+            "info": "The stage to perform in training."
+        },
+        "zh": {
+            "label": "训练阶段",
+            "info": "目前采用的训练方式。"
+        }
+    },
    "dataset_dir": {
        "en": {
            "label": "Data dir",
@@ -105,12 +115,12 @@ LOCALES = {
            "label": "数据集"
        }
    },
-    "preview_btn": {
+    "data_preview_btn": {
        "en": {
-            "value": "Preview"
+            "value": "Preview dataset"
        },
        "zh": {
-            "value": "预览"
+            "value": "预览数据集"
        }
    },
    "preview_count": {
@@ -277,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"
@@ -335,6 +335,42 @@ LOCALES = {
            "info": "接着上次的 LoRA 权重训练或创建一个新的 LoRA 权重。"
        }
    },
+    "rlhf_tab": {
+        "en": {
+            "label": "RLHF configurations"
+        },
+        "zh": {
+            "label": "RLHF 参数设置"
+        }
+    },
+    "dpo_beta": {
+        "en": {
+            "label": "DPO beta",
+            "info": "Value of the beta parameter in the DPO loss."
+        },
+        "zh": {
+            "label": "DPO beta 参数",
+            "info": "DPO 损失函数中 beta 超参数大小。"
+        }
+    },
+    "reward_model": {
+        "en": {
+            "label": "Reward model",
+            "info": "Checkpoint of the reward model for PPO training."
+        },
+        "zh": {
+            "label": "奖励模型",
+            "info": "PPO 训练中奖励模型的断点路径。"
+        }
+    },
+    "cmd_preview_btn": {
+        "en": {
+            "value": "Preview command"
+        },
+        "zh": {
+            "value": "预览命令"
+        }
+    },
    "start_btn": {
        "en": {
            "value": "Start"
@@ -409,7 +445,7 @@ LOCALES = {
            "value": "模型未加载，请先加载模型。"
        }
    },
-    "prefix": {
+    "system": {
        "en": {
            "placeholder": "System prompt (optional)"
        },
--- a/src/llmtuner/webui/manager.py
+++ b/src/llmtuner/webui/manager.py
@@ -12,13 +12,22 @@ class Manager:
    def __init__(self, elem_list: List[Dict[str, Component]]):
        self.elem_list = elem_list

-    def gen_refresh(self) -> Dict[str, Any]:
+    def gen_refresh(self, lang: str) -> Dict[str, Any]:
        refresh_dict = {
            "dataset": {"choices": list_dataset()["choices"]},
            "output_dir": {"value": get_time()}
        }
+
        user_config = load_config()
-        if user_config["last_model"]:
+        if not lang:
+            if user_config.get("lang", None):
+                lang = user_config["lang"]
+            else:
+                lang = "en"
+
+        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"])}

@@ -26,10 +35,12 @@ class Manager:

    def gen_label(self, lang: str) -> Dict[Component, Dict[str, Any]]: # cannot use TYPE_CHECKING
        update_dict = {}
-        refresh_dict = self.gen_refresh()
+        refresh_dict = self.gen_refresh(lang)

        for elems in self.elem_list:
            for name, component in elems.items():
-                update_dict[component] = gr.update(**LOCALES[name][lang], **refresh_dict.get(name, {}))
+                update_dict[component] = gr.update(
+                    **LOCALES[name][refresh_dict["lang"]["value"]], **refresh_dict.get(name, {})
+                )

        return update_dict
--- a/src/llmtuner/webui/runner.py
+++ b/src/llmtuner/webui/runner.py
@@ -5,16 +5,16 @@ import threading
 import time
 import transformers
 from transformers.trainer import TRAINING_ARGS_NAME
-from typing import Generator, List, Tuple
+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 get_eval_results, update_process_bar
+from llmtuner.webui.utils import gen_cmd, get_eval_results, update_process_bar


 class Runner:
@@ -22,39 +22,36 @@ class Runner:
    def __init__(self):
        self.aborted = False
        self.running = False
+        self.logger_handler = LoggerHandler()
+        self.logger_handler.setLevel(logging.INFO)
+        logging.root.addHandler(self.logger_handler)
+        transformers.logging.add_handler(self.logger_handler)

    def set_abort(self):
        self.aborted = True
        self.running = False

-    def initialize(
+    def _initialize(
        self, lang: str, model_name: str, dataset: List[str]
-    ) -> Tuple[str, str, LoggerHandler, LogCallback]:
+    ) -> str:
        if self.running:
-            return None, ALERTS["err_conflict"][lang], None, None
+            return ALERTS["err_conflict"][lang]

        if not model_name:
-            return None, ALERTS["err_no_model"][lang], None, None
+            return ALERTS["err_no_model"][lang]

-        model_name_or_path = get_model_path(model_name)
-        if not model_name_or_path:
-            return None, ALERTS["err_no_path"][lang], None, None
+        if not get_model_path(model_name):
+            return ALERTS["err_no_path"][lang]

        if len(dataset) == 0:
-            return None, ALERTS["err_no_dataset"][lang], None, None
+            return ALERTS["err_no_dataset"][lang]

        self.aborted = False
-        self.running = True
+        self.logger_handler.reset()
+        self.trainer_callback = LogCallback(self)
+        return ""

-        logger_handler = LoggerHandler()
-        logger_handler.setLevel(logging.INFO)
-        logging.root.addHandler(logger_handler)
-        transformers.logging.add_handler(logger_handler)
-        trainer_callback = LogCallback(self)
-
-        return model_name_or_path, "", logger_handler, trainer_callback
-
-    def finalize(
+    def _finalize(
        self, lang: str, finish_info: str
    ) -> str:
        self.running = False
@@ -64,7 +61,7 @@ class Runner:
        else:
            return finish_info

-    def run_train(
+    def _parse_train_args(
        self,
        lang: str,
        model_name: str,
@@ -72,7 +69,8 @@ class Runner:
        finetuning_type: str,
        quantization_bit: str,
        template: str,
-        source_prefix: str,
+        system_prompt: str,
+        training_stage: str,
        dataset_dir: str,
        dataset: List[str],
        max_source_length: int,
@@ -89,37 +87,37 @@ class Runner:
        save_steps: int,
        warmup_steps: int,
        compute_type: str,
-        padding_side: str,
        lora_rank: int,
        lora_dropout: float,
        lora_target: str,
        resume_lora_training: bool,
+        dpo_beta: float,
+        reward_model: str,
        output_dir: str
-    ) -> Generator[str, None, None]:
-        model_name_or_path, error, logger_handler, trainer_callback = self.initialize(lang, model_name, dataset)
-        if error:
-            yield error, gr.update(visible=False)
-            return
-
+    ) -> 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]
            )
        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",
-            model_name_or_path=model_name_or_path,
+            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 else None,
+            quantization_bit=int(quantization_bit) if quantization_bit in ["8", "4"] else None,
            template=template,
-            source_prefix=source_prefix,
+            system_prompt=system_prompt,
            dataset_dir=dataset_dir,
            dataset=",".join(dataset),
            max_source_length=max_source_length,
@@ -134,15 +132,22 @@ class Runner:
            logging_steps=logging_steps,
            save_steps=save_steps,
            warmup_steps=warmup_steps,
-            fp16=(compute_type == "fp16"),
-            bf16=(compute_type == "bf16"),
-            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 args["stage"] == "ppo":
+            args["reward_model"] = reward_model
+            val_size = 0
+
+        if args["stage"] == "dpo":
+            args["dpo_beta"] = dpo_beta

        if val_size > 1e-6:
            args["val_size"] = val_size
@@ -150,25 +155,9 @@ class Runner:
            args["eval_steps"] = save_steps
            args["load_best_model_at_end"] = True

-        run_kwargs = dict(args=args, callbacks=[trainer_callback])
-        thread = threading.Thread(target=run_exp, kwargs=run_kwargs)
-        thread.start()
+        return lang, model_name, dataset, output_dir, args

-        while thread.is_alive():
-            time.sleep(2)
-            if self.aborted:
-                yield ALERTS["info_aborting"][lang], gr.update(visible=False)
-            else:
-                yield logger_handler.log, update_process_bar(trainer_callback)
-
-        if os.path.exists(os.path.join(output_dir, TRAINING_ARGS_NAME)):
-            finish_info = ALERTS["info_finished"][lang]
-        else:
-            finish_info = ALERTS["err_failed"][lang]
-
-        yield self.finalize(lang, finish_info), gr.update(visible=False)
-
-    def run_eval(
+    def _parse_eval_args(
        self,
        lang: str,
        model_name: str,
@@ -176,7 +165,7 @@ class Runner:
        finetuning_type: str,
        quantization_bit: str,
        template: str,
-        source_prefix: str,
+        system_prompt: str,
        dataset_dir: str,
        dataset: List[str],
        max_source_length: int,
@@ -184,32 +173,31 @@ class Runner:
        max_samples: str,
        batch_size: int,
        predict: bool
-    ) -> Generator[str, None, None]:
-        model_name_or_path, error, logger_handler, trainer_callback = self.initialize(lang, model_name, dataset)
-        if error:
-            yield error, gr.update(visible=False)
-            return
-
+    ) -> 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=model_name_or_path,
+            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 else None,
+            quantization_bit=int(quantization_bit) if quantization_bit in ["8", "4"] else None,
            template=template,
-            source_prefix=source_prefix,
+            system_prompt=system_prompt,
            dataset_dir=dataset_dir,
            dataset=",".join(dataset),
            max_source_length=max_source_length,
@@ -223,7 +211,33 @@ class Runner:
            args.pop("do_eval", None)
            args["do_predict"] = True

-        run_kwargs = dict(args=args, callbacks=[trainer_callback])
+        return lang, model_name, dataset, output_dir, args
+
+    def preview_train(self, *args) -> Generator[Tuple[str, Dict[str, Any]], None, None]:
+        lang, model_name, dataset, _, args = self._parse_train_args(*args)
+        error = self._initialize(lang, model_name, dataset)
+        if error:
+            yield error, gr.update(visible=False)
+        else:
+            yield gen_cmd(args), gr.update(visible=False)
+
+    def preview_eval(self, *args) -> Generator[Tuple[str, Dict[str, Any]], None, None]:
+        lang, model_name, dataset, _, args = self._parse_eval_args(*args)
+        error = self._initialize(lang, model_name, dataset)
+        if error:
+            yield error, gr.update(visible=False)
+        else:
+            yield gen_cmd(args), gr.update(visible=False)
+
+    def run_train(self, *args) -> Generator[Tuple[str, Dict[str, Any]], None, None]:
+        lang, model_name, dataset, output_dir, args = self._parse_train_args(*args)
+        error = self._initialize(lang, model_name, dataset)
+        if error:
+            yield error, gr.update(visible=False)
+            return
+
+        self.running = True
+        run_kwargs = dict(args=args, callbacks=[self.trainer_callback])
        thread = threading.Thread(target=run_exp, kwargs=run_kwargs)
        thread.start()

@@ -232,11 +246,37 @@ class Runner:
            if self.aborted:
                yield ALERTS["info_aborting"][lang], gr.update(visible=False)
            else:
-                yield logger_handler.log, update_process_bar(trainer_callback)
+                yield self.logger_handler.log, update_process_bar(self.trainer_callback)
+
+        if os.path.exists(os.path.join(output_dir, TRAINING_ARGS_NAME)):
+            finish_info = ALERTS["info_finished"][lang]
+        else:
+            finish_info = ALERTS["err_failed"][lang]
+
+        yield self._finalize(lang, finish_info), gr.update(visible=False)
+
+    def run_eval(self, *args) -> Generator[str, None, None]:
+        lang, model_name, dataset, output_dir, args = self._parse_eval_args(*args)
+        error = self._initialize(lang, model_name, dataset)
+        if error:
+            yield error, gr.update(visible=False)
+            return
+
+        self.running = True
+        run_kwargs = dict(args=args, callbacks=[self.trainer_callback])
+        thread = threading.Thread(target=run_exp, kwargs=run_kwargs)
+        thread.start()
+
+        while thread.is_alive():
+            time.sleep(2)
+            if self.aborted:
+                yield ALERTS["info_aborting"][lang], gr.update(visible=False)
+            else:
+                yield self.logger_handler.log, update_process_bar(self.trainer_callback)

        if os.path.exists(os.path.join(output_dir, "all_results.json")):
            finish_info = get_eval_results(os.path.join(output_dir, "all_results.json"))
        else:
            finish_info = ALERTS["err_failed"][lang]

-        yield self.finalize(lang, finish_info), gr.update(visible=False)
+        yield self._finalize(lang, finish_info), gr.update(visible=False)
--- 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,22 +47,42 @@ 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)


+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 src/train_bash.py "]
+    for k, v in args.items():
+        if v is not None and v != "":
+            cmd_lines.append("    --{} {} ".format(k, str(v)))
+    cmd_text = "\\\n".join(cmd_lines)
+    cmd_text = "```bash\n{}\n```".format(cmd_text)
+    return cmd_text
+
+
 def get_eval_results(path: os.PathLike) -> str:
    with open(path, "r", encoding="utf-8") as f:
        result = json.dumps(json.load(f), indent=4)
@@ -69,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

@@ -118,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
hiyouga	6eed1db36c	Release v0.1.7 Former-commit-id: 81abe8d6cabaa1ebe74dc32a5dc143389e4c9f31	2023-08-18 17:21:27 +08:00
hiyouga	948124f55e	tiny fix Former-commit-id: 0ee159654ac6339c162745b004e2152ba6fe3c81	2023-08-18 13:07:35 +08:00
hiyouga	2b191ca776	support ppo score norm (trl 0.5.1.dev required) Former-commit-id: 2b25db6d260ec1532281a592e873579346c7d21c	2023-08-18 12:02:42 +08:00
hiyouga	be4d2822ea	fix PPO trainer #551 , update readme Former-commit-id: faead74849470cebae9e37cde5fab2a71b32aa43	2023-08-18 11:43:10 +08:00
hiyouga	736ddd0319	update readme Former-commit-id: beaf2fb737dbe64d35334d88b42935c89ef09eee	2023-08-18 01:51:55 +08:00
hiyouga	dfa289aa72	Update .gitignore Former-commit-id: a1772a4dfef8dfaf7c2c321fad0a70ccf95fe6a0	2023-08-18 01:43:42 +08:00
hiyouga	c2644f939a	update training resuming Former-commit-id: 2ec75c31f609e65116ac3b621eeb7d8ccbf69135	2023-08-18 01:41:17 +08:00
hoshi-hiyouga	f11c1ae562	Merge pull request #434 from niuba/main add last_checkpoint support Former-commit-id: b78d461f2826c194c332ead37825704c2cb8b910	2023-08-18 01:38:31 +08:00
hoshi-hiyouga	3126164aa6	Merge branch 'main' into main Former-commit-id: 870d2c7bf74d0da5a927bef4b8b01d15cc66a3e9	2023-08-18 01:37:23 +08:00
hiyouga	ed10486cad	support bf16 ppo #551 Former-commit-id: 092088967de7409a2d51847cfc7afc83a8887320	2023-08-18 00:40:32 +08:00
hiyouga	04fa430c6c	fix ChatGLM2 ppo #527 #528 Former-commit-id: 60d6ad64d7c9f6445b0df8de0153c3a311974198	2023-08-18 00:34:59 +08:00
hiyouga	fa1893b59c	fix generation bug #532 Former-commit-id: c071121e67374e5f09798db57cfc8668617a36ae	2023-08-17 22:21:34 +08:00
hiyouga	e993e717a5	fix streaming in pt stage #548 #549 Former-commit-id: 050e992bee2a9293cc7399b578de807b5bf9bddc	2023-08-17 17:59:26 +08:00
hiyouga	c80e56423a	update readme Former-commit-id: b74af3c9cf29e1690ae4d5acb27599b1abd152e2	2023-08-17 11:00:22 +08:00
hiyouga	ffa09a01d6	fix baichuan and intern template Former-commit-id: e1fd18fa6ef1009f978aca5210a259251a0b19a6	2023-08-17 01:27:20 +08:00
hiyouga	7d04f8567b	fix generation Former-commit-id: 66a0300d312ef91c24fcf80667fa3b0bb8e1a342	2023-08-16 22:39:54 +08:00
hiyouga	baa709674f	fix system prompt Former-commit-id: 411e775aa939bdd154a3f1e92921ede90d989f18	2023-08-16 01:35:52 +08:00
hiyouga	ca9a494d0c	fix baichuan template #481 Former-commit-id: 7608c6c25877d97ef26a1c209c4073c9c42f4535	2023-08-15 11:38:21 +08:00
hoshi-hiyouga	37eb8c05cc	Merge pull request #516 from liuyanyi/add_gitignore [Enhance] Add .gitignore file Former-commit-id: 12cfe5482f5ef95d8c386d0af0de381e72eab0f9	2023-08-15 11:25:40 +08:00
hiyouga	7c046edb7b	fix ChatGLM RLHF Former-commit-id: 4e43e887e432ceb7e9287b4e309b63af3c3ba1bf	2023-08-15 11:19:20 +08:00
Yanyi Liu	22cea38b20	Add .gitignore Former-commit-id: a2ebdeef81706596617da4409fc5da71739bccdc	2023-08-15 11:13:45 +08:00
hiyouga	ef2ca0a827	alert pad_token source Former-commit-id: f26a84e0d927d2554890daf431a93652e18f4235	2023-08-15 00:07:56 +08:00
hiyouga	7f0b908de2	update webui Former-commit-id: da30d0fb4abdb825f3383ddd106bb06a84695b7a	2023-08-14 22:45:26 +08:00
hoshi-hiyouga	5fc5e776ff	Merge pull request #511 from hiyouga/feature-autoTemplate add template match and stage in webui Former-commit-id: 413752ecba845cddaff5fb48db7d3d24b960eec1	2023-08-14 22:44:04 +08:00
codemayq	93b281c016	auto match template when change model_name Former-commit-id: ab2d7ab0572765ce33a52ac71641062d5d904db4	2023-08-14 20:56:05 +08:00
codemayq	9585699918	add template match and stage in webui Former-commit-id: d6283e7f041f08f76d18350cb5f6a6c58ca80e92	2023-08-14 20:42:59 +08:00
hiyouga	bceaba551d	fix ChatGLM lm_head #494 Former-commit-id: bf0048abdaeb2b9592d38ac991704ad014370b47	2023-08-14 14:14:48 +08:00
hiyouga	0bfeed3a7e	fix bug in webui Former-commit-id: c95f0f687689934379b6c24abf872ffcde06073b	2023-08-14 11:38:42 +08:00
hiyouga	70a780c3c0	fix webui cache Former-commit-id: 9aba5c197fbc8abaab77f454374f8b497f0310d0	2023-08-14 11:37:01 +08:00
hiyouga	d74ab5306c	update readme_zh Former-commit-id: bdfe7e0285fdeb3a2728669dbdabf70c9652735c	2023-08-14 11:13:25 +08:00
hiyouga	688e8601ab	web UI integrating RLHF Former-commit-id: 137fd146b90f89a1164b56e6d507b30b1f5c2437	2023-08-14 10:48:47 +08:00
hiyouga	4933ab5956	fix #480 Former-commit-id: ec15ca8fffacba2c34e1849c5ce90ca9989d66a2	2023-08-14 00:23:56 +08:00
hiyouga	6c7225a5d4	fix webui Former-commit-id: 2c8b7414be9b43e20cc1d0575cc4dc1c7545fd86	2023-08-12 23:52:07 +08:00
hiyouga	a22982f2fa	tiny fix Former-commit-id: 50a34c043de6d9e1410291e1d8c1ea9d53754e9e	2023-08-12 22:02:43 +08:00
hiyouga	c95479dddb	fix rope scaling Former-commit-id: 2e0dd36700ec5e8294581c1db4b9431f755fc5f8	2023-08-12 22:00:01 +08:00
hiyouga	fc48bd8da0	update readme Former-commit-id: 94ac570cb62aa9cd5dba105f0bb4c4da43eca042	2023-08-12 21:29:06 +08:00
hiyouga	d5323bfa3f	update readme Former-commit-id: ecfe87f34b383901f8e97ffb90af459cd55419b1	2023-08-12 21:25:19 +08:00
hiyouga	e9d4a2b507	update readme Former-commit-id: eadbe9b7a0b6c8897e7a763b519cc5b7e00f3b2c	2023-08-12 21:23:05 +08:00
hiyouga	37bcbe8046	update readme Former-commit-id: 6fa381400c21fa249cebcdff8c3afd72f8de20b3	2023-08-12 21:00:11 +08:00
hiyouga	fdfb644f0a	support rope scaling, fix #475 #476 #478 Former-commit-id: 337d5f68b72230e545e7a94ca789187c7a2b7187	2023-08-12 20:46:27 +08:00
hoshi-hiyouga	cde9f3db57	Merge pull request #479 from hiyouga/feature-addCmdExport add sft script preview in webui Former-commit-id: 060225e57d13d8164beb6920410c181fbb28b77a	2023-08-12 20:41:52 +08:00
codemayq	8bf5a98815	add sft script preview in webui Former-commit-id: 2b72649b404750226aa418b61ef5a6c9ac03938f	2023-08-12 13:53:55 +08:00
hiyouga	be566a15a5	fix unusual output of 8bit models #278 #391 Former-commit-id: 337ce5272b81f5561162beb08814b0e5abf23703	2023-08-12 00:25:29 +08:00
niuba	458955d0fb	add last_checkpoint support Former-commit-id: 9f1977e4de00b14a9d1b555c25bcaf12998d5046	2023-08-09 16:39:27 +08:00