Release v0.1.8

Former-commit-id: d9666411375964d334d0a93ec162b27e05f70d49

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

README.md

@@ -8,83 +8,104 @@
 
 👋 Join our [WeChat](assets/wechat.jpg).
 
+\[ English | [中文](README_zh.md) \]
+
 ## Changelog
 
-[23/07/11] Now we support training the **Baichuan-13B** model in this repo. Try `--model_name_or_path baichuan-inc/Baichuan-13B-Base`, `--padding_side right` and `--lora_target W_pack` arguments to train the Baichuan-13B model. Remember to use `--prompt_template baichuan` argument when you are using the Baichuan-13B-Chat model.
+[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/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/08/18] Now we support **resuming training**, upgrade `transformers` to `4.31.0` to enjoy this feature.
 
-[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 `--prompt_template intern` argument when you are using the InternLM-chat model.
+[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/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/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/06/29] We provide a **reproducible example** of training a chat model using instruction-following datasets, see this [HuggingFace Repo](https://huggingface.co/hiyouga/baichuan-7b-sft) for details.
+[23/07/31] Now we support **dataset streaming**. Try `--streaming` and `--max_steps 10000` arguments to load your dataset in streaming mode.
+
+[23/07/29] We release two instruction-tuned 13B models at Hugging Face. See these Hugging Face Repos ([LLaMA-2](https://huggingface.co/hiyouga/Llama-2-Chinese-13b-chat) / [Baichuan](https://huggingface.co/hiyouga/Baichuan-13B-sft)) for details.
+
+[23/07/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/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. If you want to train with RTX3090, use `git checkout baichuan-7b-rtx3090` to switch to the `baichuan-7b-rtx3090` branch and try the `--baichuan_rtx_gpu true` argument. (Other RTX series GPUs can also be tried)
-
-[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 model. (experimental feature)
-
-[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.
+[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.
 
 ## Supported Models
 
-- [LLaMA](https://github.com/facebookresearch/llama) (7B/13B/33B/65B)
-- [BLOOM](https://huggingface.co/bigscience/bloom) & [BLOOMZ](https://huggingface.co/bigscience/bloomz) (560M/1.1B/1.7B/3B/7.1B/176B)
-- [Falcon](https://huggingface.co/tiiuae/falcon-7b) (7B/40B)
-- [Baichuan](https://huggingface.co/baichuan-inc/baichuan-7B) (7B/13B)
-- [InternLM](https://github.com/InternLM/InternLM) (7B)
+| 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  |
+
+> [!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
 
-- [(Continually) pre-training](https://s3-us-west-2.amazonaws.com/openai-assets/research-covers/language-unsupervised/language_understanding_paper.pdf)
-  - Full-parameter tuning
-  - Partial-parameter tuning
-  - [LoRA](https://arxiv.org/abs/2106.09685)
-  - [QLoRA](https://arxiv.org/abs/2305.14314)
-- [Supervised fine-tuning](https://arxiv.org/abs/2109.01652)
-  - Full-parameter tuning
-  - Partial-parameter tuning
-  - [LoRA](https://arxiv.org/abs/2106.09685)
-  - [QLoRA](https://arxiv.org/abs/2305.14314)
-- [RLHF](https://arxiv.org/abs/2203.02155)
-  - [LoRA](https://arxiv.org/abs/2106.09685)
-  - [QLoRA](https://arxiv.org/abs/2305.14314)
+| 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
 
 - For pre-training:
-  - [Wiki Demo](data/wiki_demo.txt)
+  - [Wiki Demo (en)](data/wiki_demo.txt)
+  - [RefinedWeb (en)](https://huggingface.co/datasets/tiiuae/falcon-refinedweb)
+  - [StarCoder (en)](https://huggingface.co/datasets/bigcode/starcoderdata)
+  - [Wikipedia (en)](https://huggingface.co/datasets/olm/olm-wikipedia-20221220)
+  - [Wikipedia (zh)](https://huggingface.co/datasets/pleisto/wikipedia-cn-20230720-filtered)
 - For supervised fine-tuning:
-  - [Stanford Alpaca](https://github.com/tatsu-lab/stanford_alpaca)
-  - [Stanford Alpaca (Chinese)](https://github.com/ymcui/Chinese-LLaMA-Alpaca)
-  - [GPT-4 Generated Data](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)
-  - [BELLE 2M](https://huggingface.co/datasets/BelleGroup/train_2M_CN)
-  - [BELLE 1M](https://huggingface.co/datasets/BelleGroup/train_1M_CN)
-  - [BELLE 0.5M](https://huggingface.co/datasets/BelleGroup/train_0.5M_CN)
-  - [BELLE Dialogue 0.4M](https://huggingface.co/datasets/BelleGroup/generated_chat_0.4M)
-  - [BELLE School Math 0.25M](https://huggingface.co/datasets/BelleGroup/school_math_0.25M)
-  - [BELLE Multiturn Chat 0.8M](https://huggingface.co/datasets/BelleGroup/multiturn_chat_0.8M)
-  - [Guanaco Dataset](https://huggingface.co/datasets/JosephusCheung/GuanacoDataset)
-  - [Firefly 1.1M](https://huggingface.co/datasets/YeungNLP/firefly-train-1.1M)
-  - [CodeAlpaca 20k](https://huggingface.co/datasets/sahil2801/CodeAlpaca-20k)
-  - [Alpaca CoT](https://huggingface.co/datasets/QingyiSi/Alpaca-CoT)
-  - [Web QA (Chinese)](https://huggingface.co/datasets/suolyer/webqa)
-  - [UltraChat](https://github.com/thunlp/UltraChat)
-  - [Open Assistant](https://huggingface.co/datasets/OpenAssistant/oasst1)
-  - [Open Assistant (Chinese)](https://huggingface.co/datasets/OpenAssistant/oasst1)
-  - [WebNovel (Chinese)](https://huggingface.co/datasets/zxbsmk/webnovel_cn)
-- For reward model training:
-  - [HH-RLHF](https://huggingface.co/datasets/Anthropic/hh-rlhf)
-  - [Open Assistant](https://huggingface.co/datasets/OpenAssistant/oasst1)
-  - [Open Assistant (Chinese)](https://huggingface.co/datasets/OpenAssistant/oasst1)
-  - [GPT-4 Generated Data](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)
-  - [GPT-4 Generated Data (Chinese)](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)
+  - [Stanford Alpaca (en)](https://github.com/tatsu-lab/stanford_alpaca)
+  - [Stanford Alpaca (zh)](https://github.com/ymcui/Chinese-LLaMA-Alpaca)
+  - [GPT-4 Generated Data (en&zh)](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)
+  - [Open Assistant (multilingual)](https://huggingface.co/datasets/OpenAssistant/oasst1)
+  - [Self-cognition (zh)](data/self_cognition.json)
+  - [ShareGPT (zh)](https://huggingface.co/datasets/QingyiSi/Alpaca-CoT/tree/main/Chinese-instruction-collection)
+  - [Guanaco Dataset (multilingual)](https://huggingface.co/datasets/JosephusCheung/GuanacoDataset)
+  - [BELLE 2M (zh)](https://huggingface.co/datasets/BelleGroup/train_2M_CN)
+  - [BELLE 1M (zh)](https://huggingface.co/datasets/BelleGroup/train_1M_CN)
+  - [BELLE 0.5M (zh)](https://huggingface.co/datasets/BelleGroup/train_0.5M_CN)
+  - [BELLE Dialogue 0.4M (zh)](https://huggingface.co/datasets/BelleGroup/generated_chat_0.4M)
+  - [BELLE School Math 0.25M (zh)](https://huggingface.co/datasets/BelleGroup/school_math_0.25M)
+  - [BELLE Multiturn Chat 0.8M (zh)](https://huggingface.co/datasets/BelleGroup/multiturn_chat_0.8M)
+  - [Firefly 1.1M (zh)](https://huggingface.co/datasets/YeungNLP/firefly-train-1.1M)
+  - [LIMA (en)](https://huggingface.co/datasets/GAIR/lima)
+  - [CodeAlpaca 20k (en)](https://huggingface.co/datasets/sahil2801/CodeAlpaca-20k)
+  - [Alpaca CoT (multilingual)](https://huggingface.co/datasets/QingyiSi/Alpaca-CoT)
+  - [Web QA (zh)](https://huggingface.co/datasets/suolyer/webqa)
+  - [UltraChat (en)](https://github.com/thunlp/UltraChat)
+  - [WebNovel (zh)](https://huggingface.co/datasets/zxbsmk/webnovel_cn)
+  - [Ad Gen (zh)](https://huggingface.co/datasets/HasturOfficial/adgen)
+- For reward modeling or DPO training:
+  - [HH-RLHF (en)](https://huggingface.co/datasets/Anthropic/hh-rlhf)
+  - [Open Assistant (multilingual)](https://huggingface.co/datasets/OpenAssistant/oasst1)
+  - [GPT-4 Generated Data (en&zh)](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)
 
 Please refer to [data/README.md](data/README.md) for details.
 
-Some datasets require confirmation before using them, so we recommend logging in with your HuggingFace account using these commands.
+Some datasets require confirmation before using them, so we recommend logging in with your Hugging Face account using these commands.
 
 ```bash
 pip install --upgrade huggingface_hub
@@ -95,25 +116,21 @@ huggingface-cli login
 
 - Python 3.8+ and PyTorch 1.13.1+
 - 🤗Transformers, Datasets, Accelerate, PEFT and TRL
+- 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)
 
 And **powerful GPUs**!
 
-If you want to enable quantized LoRA (QLoRA) on the Windows platform, you should install a pre-built version of `bitsandbytes` library, which supports CUDA 11.1 to 12.1.
-
-```bash
-pip install https://github.com/jllllll/bitsandbytes-windows-webui/releases/download/wheels/bitsandbytes-0.39.1-py3-none-win_amd64.whl
-```
-
 ## Getting Started
 
 ### Data Preparation (optional)
 
 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)
 
@@ -125,25 +142,38 @@ cd LLaMA-Efficient-Tuning
 pip install -r requirements.txt
 ```
 
-### LLaMA Weights Preparation (optional)
-
-1. Download the weights of the LLaMA models.
-2. Convert them to HF format using the following command.
+If you want to enable the quantized LoRA (QLoRA) on the Windows platform, you will be required to install a pre-built version of `bitsandbytes` library, which supports CUDA 11.1 to 12.1.
 
 ```bash
-python -m transformers.models.llama.convert_llama_weights_to_hf \
-    --input_dir path_to_llama_weights --model_size 7B --output_dir path_to_llama_model
+pip install https://github.com/jllllll/bitsandbytes-windows-webui/releases/download/wheels/bitsandbytes-0.39.1-py3-none-win_amd64.whl
 ```
 
-### (Continually) Pre-Training
+### All-in-one Web UI
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python src/train_web.py
+```
+
+We strongly recommend using the all-in-one Web UI for newcomers since it can also generate training scripts **automatically**.
+
+> [!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 \
     --finetuning_type lora \
+    --lora_target q_proj,v_proj \
     --output_dir path_to_pt_checkpoint \
     --overwrite_cache \
     --per_device_train_batch_size 4 \
@@ -157,15 +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 \
@@ -179,36 +211,43 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
     --fp16
 ```
 
-### 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 \
-    --per_device_train_batch_size 4 \
+    --per_device_train_batch_size 2 \
     --gradient_accumulation_steps 4 \
     --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 (RLHF)
+#### 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 \
     --output_dir path_to_ppo_checkpoint \
@@ -219,30 +258,51 @@ CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
     --save_steps 1000 \
     --learning_rate 1e-5 \
     --num_train_epochs 1.0 \
+    --plot_loss \
+    --fp16
+```
+
+#### DPO Training
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
+    --stage dpo \
+    --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 \
-    --plot_loss
+    --checkpoint_dir path_to_sft_checkpoint \
+    --output_dir path_to_dpo_checkpoint \
+    --per_device_train_batch_size 2 \
+    --gradient_accumulation_steps 4 \
+    --lr_scheduler_type cosine \
+    --logging_steps 10 \
+    --save_steps 1000 \
+    --learning_rate 1e-5 \
+    --num_train_epochs 1.0 \
+    --plot_loss \
+    --fp16
 ```
 
 ### 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
@@ -258,40 +318,126 @@ use_cpu: false
 
 </details>
 
-### Evaluation (BLEU and ROUGE_CHINESE)
+#### Use DeepSpeed
 
 ```bash
-CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
-    --stage pt \
-    --model_name_or_path path_to_your_model \
-    --do_eval \
-    --dataset alpaca_gpt4_en \
-    --checkpoint_dir path_to_checkpoint \
-    --output_dir path_to_eval_result \
-    --per_device_eval_batch_size 8 \
-    --max_samples 50 \
-    --predict_with_generate
+deepspeed --num_gpus 8 --master_port=9901 src/train_bash.py \
+    --deepspeed ds_config.json \
+    ... # arguments (same as above)
 ```
 
-We recommend using `--per_device_eval_batch_size=1` and `--max_target_length 128` at 4/8-bit evaluation.
+<details><summary>Example config for full-parameter training with DeepSpeed ZeRO-2</summary>
 
-### API / CLI / Web Demo
-
-```bash
-python src/xxx_demo.py \
-    --model_name_or_path path_to_your_model \
-    --checkpoint_dir path_to_checkpoint
+```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_your_model \
+    --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_llama_model \
+    --do_eval \
+    --dataset alpaca_gpt4_en \
+    --template default \
+    --finetuning_type lora \
+    --checkpoint_dir path_to_checkpoint \
+    --output_dir path_to_eval_result \
+    --per_device_eval_batch_size 8 \
+    --max_samples 100 \
+    --predict_with_generate
+```
+
+> [!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_llama_model \
+    --do_predict \
+    --dataset alpaca_gpt4_en \
+    --template default \
+    --finetuning_type lora \
+    --checkpoint_dir path_to_checkpoint \
+    --output_dir path_to_predict_result \
+    --per_device_eval_batch_size 8 \
+    --max_samples 100 \
+    --predict_with_generate
+```
+
 ## License
 
 This repository is licensed under the [Apache-2.0 License](LICENSE).
@@ -299,10 +445,15 @@ This repository is licensed under the [Apache-2.0 License](LICENSE).
 Please follow the model licenses to use the corresponding model weights:
 
 - [LLaMA](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md)
+- [LLaMA-2](https://ai.meta.com/llama/license/)
 - [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)
+- [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)
 
 ## Citation
 
@@ -319,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
 

README_zh.md

@@ -0,0 +1,476 @@
+# LLaMA Efficient Tuning
+
+[![GitHub Repo stars](https://img.shields.io/github/stars/hiyouga/LLaMA-Efficient-Tuning?style=social)](https://github.com/hiyouga/LLaMA-Efficient-Tuning/stargazers)
+[![GitHub Code License](https://img.shields.io/github/license/hiyouga/LLaMA-Efficient-Tuning)](LICENSE)
+[![GitHub last commit](https://img.shields.io/github/last-commit/hiyouga/LLaMA-Efficient-Tuning)](https://github.com/hiyouga/LLaMA-Efficient-Tuning/commits/main)
+[![PyPI](https://img.shields.io/pypi/v/llmtuner)](https://pypi.org/project/llmtuner/)
+[![GitHub pull request](https://img.shields.io/badge/PRs-welcome-blue)](https://github.com/hiyouga/LLaMA-Efficient-Tuning/pulls)
+
+👋 加入我们的[微信群](assets/wechat.jpg)。
+
+\[ [English](README.md) | 中文 \]
+
+## 更新日志
+
+[23/09/10] 现在我们支持了 LLaMA 模型的 **[FlashAttention](https://github.com/Dao-AILab/flash-attention)**。如果您使用的是 RTX4090、A100 或 H100 GPU，请使用 `--flash_attn` 参数以启用 FlashAttention-2（实验性功能）。
+
+[23/08/18] 现在我们支持了**训练状态恢复**，请将 `transformers` 升级至 `4.31.0` 以启用此功能。
+
+[23/08/12] 现在我们支持了 **RoPE 插值**来扩展 LLaMA 模型的上下文长度。请使用 `--rope_scaling linear` 参数训练模型或使用 `--rope_scaling dynamic` 参数评估模型。
+
+[23/08/11] 现在我们支持了指令模型的 **[DPO 训练](https://arxiv.org/abs/2305.18290)**。详情请参阅[此示例](#dpo-训练)。
+
+[23/07/31] 现在我们支持了**数据流式加载**。请尝试使用 `--streaming` 和 `--max_steps 10000` 参数来流式加载数据集。
+
+[23/07/29] 我们在 Hugging Face 发布了两个 13B 指令微调模型。详细内容请查阅我们的 Hugging Face 项目（[LLaMA-2](https://huggingface.co/hiyouga/Llama-2-Chinese-13b-chat) / [Baichuan](https://huggingface.co/hiyouga/Baichuan-13B-sft)）。
+
+[23/07/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/06/29] 我们提供了一个**可复现的**指令模型微调示例，详细内容请查阅 [Baichuan-7B-sft](https://huggingface.co/hiyouga/Baichuan-7B-sft)。
+
+[23/06/22] 我们对齐了[示例 API](src/api_demo.py) 与 [OpenAI API](https://platform.openai.com/docs/api-reference/chat) 的格式，您可以将微调模型接入**任意基于 ChatGPT 的应用**中。
+
+[23/06/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  |
+| [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  |
+
+> [!NOTE]
+> **默认模块**应作为 `--lora_target` 参数的默认值，可使用 `--lora_target all` 参数指定全部模块。
+>
+> 对于所有“基座”（Base）模型，`--template` 参数可以是 `default`, `alpaca`, `vicuna` 等任意值。但“对话”（Chat）模型请务必使用对应的模板。
+
+## 训练方法
+
+| 方法                   |     全参数训练      |    部分参数训练     |       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 训练。
+
+## 数据集
+
+- 用于预训练：
+  - [Wiki Demo (en)](data/wiki_demo.txt)
+  - [RefinedWeb (en)](https://huggingface.co/datasets/tiiuae/falcon-refinedweb)
+  - [StarCoder (en)](https://huggingface.co/datasets/bigcode/starcoderdata)
+  - [Wikipedia (en)](https://huggingface.co/datasets/olm/olm-wikipedia-20221220)
+  - [Wikipedia (zh)](https://huggingface.co/datasets/pleisto/wikipedia-cn-20230720-filtered)
+- 用于指令监督微调：
+  - [Stanford Alpaca (en)](https://github.com/tatsu-lab/stanford_alpaca)
+  - [Stanford Alpaca (zh)](https://github.com/ymcui/Chinese-LLaMA-Alpaca)
+  - [GPT-4 Generated Data (en&zh)](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)
+  - [Open Assistant (multilingual)](https://huggingface.co/datasets/OpenAssistant/oasst1)
+  - [Self-cognition (zh)](data/self_cognition.json)
+  - [ShareGPT (zh)](https://huggingface.co/datasets/QingyiSi/Alpaca-CoT/tree/main/Chinese-instruction-collection)
+  - [Guanaco Dataset (multilingual)](https://huggingface.co/datasets/JosephusCheung/GuanacoDataset)
+  - [BELLE 2M (zh)](https://huggingface.co/datasets/BelleGroup/train_2M_CN)
+  - [BELLE 1M (zh)](https://huggingface.co/datasets/BelleGroup/train_1M_CN)
+  - [BELLE 0.5M (zh)](https://huggingface.co/datasets/BelleGroup/train_0.5M_CN)
+  - [BELLE Dialogue 0.4M (zh)](https://huggingface.co/datasets/BelleGroup/generated_chat_0.4M)
+  - [BELLE School Math 0.25M (zh)](https://huggingface.co/datasets/BelleGroup/school_math_0.25M)
+  - [BELLE Multiturn Chat 0.8M (zh)](https://huggingface.co/datasets/BelleGroup/multiturn_chat_0.8M)
+  - [Firefly 1.1M (zh)](https://huggingface.co/datasets/YeungNLP/firefly-train-1.1M)
+  - [LIMA (en)](https://huggingface.co/datasets/GAIR/lima)
+  - [CodeAlpaca 20k (en)](https://huggingface.co/datasets/sahil2801/CodeAlpaca-20k)
+  - [Alpaca CoT (multilingual)](https://huggingface.co/datasets/QingyiSi/Alpaca-CoT)
+  - [Web QA (zh)](https://huggingface.co/datasets/suolyer/webqa)
+  - [UltraChat (en)](https://github.com/thunlp/UltraChat)
+  - [WebNovel (zh)](https://huggingface.co/datasets/zxbsmk/webnovel_cn)
+  - [Ad Gen (zh)](https://huggingface.co/datasets/HasturOfficial/adgen)
+- 用于训练奖励模型或 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)
+
+使用方法请参考 [data/README.md](data/README_zh.md) 文件。
+
+部分数据集的使用需要确认，我们推荐使用下述命令登录您的 Hugging Face 账户。
+
+```bash
+pip install --upgrade huggingface_hub
+huggingface-cli login
+```
+
+## 软件依赖
+
+- Python 3.8+ 和 PyTorch 1.13.1+
+- 🤗Transformers, Datasets, Accelerate, PEFT 和 TRL
+- sentencepiece, protobuf 和 tiktoken
+- jieba, rouge-chinese 和 nltk (用于评估)
+- gradio 和 matplotlib (用于网页端交互)
+- uvicorn, fastapi 和 sse-starlette (用于 API)
+
+以及 **强而有力的 GPU**！
+
+## 如何使用
+
+### 数据准备（可跳过）
+
+关于数据集文件的格式，请参考 `data/example_dataset` 文件夹的内容。构建自定义数据集时，既可以使用单个 `.json` 文件，也可以使用一个[数据加载脚本](https://huggingface.co/docs/datasets/dataset_script)和多个文件。
+
+> [!NOTE]
+> 使用自定义数据集时，请更新 `data/dataset_info.json` 文件，该文件的格式请参考 `data/README.md`。
+
+### 环境搭建（可跳过）
+
+```bash
+git clone https://github.com/hiyouga/LLaMA-Efficient-Tuning.git
+conda create -n llama_etuning python=3.10
+conda activate llama_etuning
+cd LLaMA-Efficient-Tuning
+pip install -r requirements.txt
+```
+
+如果要在 Windows 平台上开启量化 LoRA（QLoRA），需要安装预编译的 `bitsandbytes` 库, 支持 CUDA 11.1 到 12.1.
+
+```bash
+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
+```
+
+我们极力推荐新手使用浏览器一体化界面，因为它还可以**自动**生成运行所需的命令行脚本。
+
+> [!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_llama_model \
+    --do_train \
+    --dataset wiki_demo \
+    --finetuning_type lora \
+    --lora_target q_proj,v_proj \
+    --output_dir path_to_pt_checkpoint \
+    --overwrite_cache \
+    --per_device_train_batch_size 4 \
+    --gradient_accumulation_steps 4 \
+    --lr_scheduler_type cosine \
+    --logging_steps 10 \
+    --save_steps 1000 \
+    --learning_rate 5e-5 \
+    --num_train_epochs 3.0 \
+    --plot_loss \
+    --fp16
+```
+
+#### 指令监督微调
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
+    --stage sft \
+    --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 \
+    --gradient_accumulation_steps 4 \
+    --lr_scheduler_type cosine \
+    --logging_steps 10 \
+    --save_steps 1000 \
+    --learning_rate 5e-5 \
+    --num_train_epochs 3.0 \
+    --plot_loss \
+    --fp16
+```
+
+#### 奖励模型训练
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
+    --stage rm \
+    --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 \
+    --per_device_train_batch_size 2 \
+    --gradient_accumulation_steps 4 \
+    --lr_scheduler_type cosine \
+    --logging_steps 10 \
+    --save_steps 1000 \
+    --learning_rate 1e-6 \
+    --num_train_epochs 1.0 \
+    --plot_loss \
+    --fp16
+```
+
+#### PPO 训练
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
+    --stage ppo \
+    --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 \
+    --output_dir path_to_ppo_checkpoint \
+    --per_device_train_batch_size 2 \
+    --gradient_accumulation_steps 4 \
+    --lr_scheduler_type cosine \
+    --logging_steps 10 \
+    --save_steps 1000 \
+    --learning_rate 1e-5 \
+    --num_train_epochs 1.0 \
+    --plot_loss
+```
+
+#### DPO 训练
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
+    --stage dpo \
+    --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 \
+    --per_device_train_batch_size 2 \
+    --gradient_accumulation_steps 4 \
+    --lr_scheduler_type cosine \
+    --logging_steps 10 \
+    --save_steps 1000 \
+    --learning_rate 1e-5 \
+    --num_train_epochs 1.0 \
+    --plot_loss \
+    --fp16
+```
+
+### 多 GPU 分布式训练
+
+#### 使用 Huggingface Accelerate
+
+```bash
+accelerate config # 首先配置分布式环境
+accelerate launch src/train_bash.py # 参数同上
+```
+
+<details><summary>LoRA 训练的 Accelerate 配置示例</summary>
+
+```yaml
+compute_environment: LOCAL_MACHINE
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+gpu_ids: all
+machine_rank: 0
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false
+```
+
+</details>
+
+#### 使用 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_llama_model \
+    --do_eval \
+    --dataset alpaca_gpt4_zh \
+    --template default \
+    --finetuning_type lora \
+    --checkpoint_dir path_to_checkpoint \
+    --output_dir path_to_eval_result \
+    --per_device_eval_batch_size 8 \
+    --max_samples 100 \
+    --predict_with_generate
+```
+
+> [!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_llama_model \
+    --do_predict \
+    --dataset alpaca_gpt4_zh \
+    --template default \
+    --finetuning_type lora \
+    --checkpoint_dir path_to_checkpoint \
+    --output_dir path_to_predict_result \
+    --per_device_eval_batch_size 8 \
+    --max_samples 100 \
+    --predict_with_generate
+```
+
+## 协议
+
+本仓库的代码依照 [Apache-2.0](LICENSE) 协议开源。
+
+使用模型权重时，请遵循对应的模型协议：
+
+- [LLaMA](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md)
+- [LLaMA-2](https://ai.meta.com/llama/license/)
+- [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)
+
+## 引用
+
+如果您觉得此项目有帮助，请考虑以下列格式引用
+
+```bibtex
+@Misc{llama-efficient-tuning,
+  title = {LLaMA Efficient Tuning},
+  author = {hiyouga},
+  howpublished = {\url{https://github.com/hiyouga/LLaMA-Efficient-Tuning}},
+  year = {2023}
+}
+```
+
+## 致谢
+
+本项目受益于 [PEFT](https://github.com/huggingface/peft)、[QLoRA](https://github.com/artidoro/qlora) 和 [OpenChatKit](https://github.com/togethercomputer/OpenChatKit)，感谢以上诸位作者的付出。
+
+## Star History
+
+![Star History Chart](https://api.star-history.com/svg?repos=hiyouga/LLaMA-Efficient-Tuning&type=Date)

data/README.md

@@ -1,53 +1,32 @@
-Data format in `dataset_info.json`:
+If you are using a custom dataset, please provide your dataset definition in the following format in `dataset_info.json`.
+
 ```json
 "dataset_name": {
-    "hf_hub_url": "the name of the dataset repository on the HuggingFace hub. (if specified, ignore below 3 arguments)",
-    "script_url": "the name of the directory containing a dataset loading script. (if specified, ignore below 2 arguments)",
-    "file_name": "the name of the dataset file in the this directory. (required if above are not specified)",
-    "file_sha1": "the SHA-1 hash value of the dataset file. (optional)",
-    "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)",
-        "response": "the name of the column in the datasets containing the responses. (default: output)",
-        "history": "the name of the column in the datasets containing the history of chat. (default: None)"
-    }
+  "hf_hub_url": "the name of the dataset repository on the HuggingFace hub. (if specified, ignore below 3 arguments)",
+  "script_url": "the name of the directory containing a dataset loading script. (if specified, ignore below 2 arguments)",
+  "file_name": "the name of the dataset file in the this directory. (required if above are not specified)",
+  "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)",
+    "response": "the name of the column in the datasets containing the responses. (default: output)",
+    "history": "the name of the column in the datasets containing the history of chat. (default: None)"
+  }
 }
 ```
 
-`dataset_info.json` 中的数据集定义格式：
+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
-"数据集名称": {
-    "hf_hub_url": "HuggingFace上的项目地址（若指定，则忽略下列三个参数）",
-    "script_url": "包含数据加载脚本的本地文件夹名称（若指定，则忽略下列两个参数）",
-    "file_name": "该目录下数据集文件的名称（若上述参数未指定，则此项必需）",
-    "file_sha1": "数据集文件的SHA-1哈希值（可选）",
-    "columns": {
-        "prompt": "数据集代表提示词的表头名称（默认：instruction）",
-        "query": "数据集代表请求的表头名称（默认：input）",
-        "response": "数据集代表回答的表头名称（默认：output）",
-        "history": "数据集代表历史对话的表头名称（默认：None）"
-    }
+{
+  "instruction": "Question",
+  "input": "",
+  "output": [
+    "Chosen answer",
+    "Rejected answer"
+  ]
 }
 ```
-
-部分预置数据集简介：
-
-| 数据集名称 | 规模 | 描述 |
-| --- | --- | --- |
-| [Stanford Alpaca](https://github.com/tatsu-lab/stanford_alpaca) | 52k | 斯坦福大学开源的 Alpaca 数据集，训练了 Alpaca 这类早期基于 LLaMA 的模型 |
-| [Stanford Alpaca (Chinese)](https://github.com/ymcui/Chinese-LLaMA-Alpaca) | 51k | 使用 ChatGPT 翻译的 Alpaca 数据集 |
-| [GPT-4 Generated Data](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM) | 100k+ | 基于 GPT-4 的 self-instruction 数据集 |
-| [BELLE 2M](https://huggingface.co/datasets/BelleGroup/train_2M_CN) | 2m | 包含约 200 万条由 [BELLE](https://github.com/LianjiaTech/BELLE) 项目生成的中文指令数据 |
-| [BELLE 1M](https://huggingface.co/datasets/BelleGroup/train_1M_CN) | 1m | 包含约 100 万条由 [BELLE](https://github.com/LianjiaTech/BELLE) 项目生成的中文指令数据 |
-| [BELLE 0.5M](https://huggingface.co/datasets/BelleGroup/train_0.5M_CN) | 500k  | 包含约 50 万条由 [BELLE](https://github.com/LianjiaTech/BELLE) 项目生成的中文指令数据 |
-| [BELLE Dialogue 0.4M](https://huggingface.co/datasets/BelleGroup/generated_chat_0.4M) | 400k | 包含约 40 万条由 [BELLE](https://github.com/LianjiaTech/BELLE) 项目生成的个性化角色对话数据，包含角色介绍 |
-| [BELLE School Math 0.25M](https://huggingface.co/datasets/BelleGroup/school_math_0.25M) | 250k  | 包含约 25 万条由 [BELLE](https://github.com/LianjiaTech/BELLE) 项目生成的中文数学题数据，包含解题过程 |
-| [BELLE Multiturn Chat 0.8M](https://huggingface.co/datasets/BelleGroup/multiturn_chat_0.8M) | 800k | 包含约 80 万条由 [BELLE](https://github.com/LianjiaTech/BELLE) 项目生成的用户与助手的多轮对话 |
-| [Guanaco Dataset](https://huggingface.co/datasets/JosephusCheung/GuanacoDataset) | 100k+ | 包含日文、简繁体中文、英文等多类数据，数据集原用于 Guanaco 模型训练 |
-| [Firefly 1.1M](https://huggingface.co/datasets/YeungNLP/firefly-train-1.1M) | 1.1M  | 中文对话大模型 firefly（流萤）的中文数据集，包含多个 NLP 任务 |
-| [CodeAlpaca 20k](https://huggingface.co/datasets/sahil2801/CodeAlpaca-20k) | 20k | 英文代码生成任务数据集 |
-| [Alpaca CoT](https://huggingface.co/datasets/QingyiSi/Alpaca-CoT) | 6M | 用于微调的指令数据集集合 |
-| [Web QA](https://huggingface.co/datasets/suolyer/webqa) | 36k | 百度知道汇集的中文问答数据集 |
-| [UltraChat](https://github.com/thunlp/UltraChat) | 1.57M | 清华 NLP 发布的大规模多轮对话数据集 |
-
-注：BELLE 数据集是由 ChatGPT 产生的数据集，不保证数据准确性，所有类 GPT 模型产生的 self-instruction 数据集均不能保证其准确性。

data/README_zh.md

@@ -0,0 +1,32 @@
+如果您使用自定义数据集，请务必在 `dataset_info.json` 文件中以如下格式提供您的数据集定义。
+
+```json
+"数据集名称": {
+  "hf_hub_url": "HuggingFace上的项目地址（若指定，则忽略下列三个参数）",
+  "script_url": "包含数据加载脚本的本地文件夹名称（若指定，则忽略下列两个参数）",
+  "file_name": "该目录下数据集文件的名称（若上述参数未指定，则此项必需）",
+  "file_sha1": "数据集文件的SHA-1哈希值（可选）",
+  "ranking": "数据集是否包含排序后的回答（默认：false）",
+  "columns": {
+    "prompt": "数据集代表提示词的表头名称（默认：instruction）",
+    "query": "数据集代表请求的表头名称（默认：input）",
+    "response": "数据集代表回答的表头名称（默认：output）",
+    "history": "数据集代表历史对话的表头名称（默认：None）"
+  }
+}
+```
+
+其中 `prompt` 和 `response` 列应当是非空的字符串。`query` 列的内容将会和 `prompt` 列拼接作为模型输入。`history` 列应当是一个列表，其中每个元素是一个字符串二元组，分别代表用户请求和模型答复。
+
+对于训练奖励模型或 DPO 训练的数据集，`response` 列应当是一个字符串列表，排在前面的代表更优的答案，例如：
+
+```json
+{
+  "instruction": "Question",
+  "input": "",
+  "output": [
+    "Chosen answer",
+    "Rejected answer"
+  ]
+}
+```

data/oaast_sft.json.REMOVED.git-id

@@ -1 +1 @@
-0a57fbc1d8cb08a8cd71c5eb8425cf59206ffed6
+57fd080be5bffe4153fe3ee26a175e3d56da30f3

data/pretrain_data/wiki_long.jsonl

@@ -1,2 +0,0 @@
-{"id": 0,"title": "大卫·亨利","content": "大卫·亨利\n\n大卫·克莱顿·亨利（David Clayton Henrie，），美国演员。近来在迪士尼频道原创电视影集《少年魔法师》（Wizards of Waverly Place）当中演出贾斯汀·鲁索（Justin Russo）一角。\n\n大卫·亨利出生在加州Mission Viejo，在凤凰城长大。他的胞弟劳伦斯·亨利（Lorenzo Henrie）也是演员。大卫·亨利就读夏安传统学校。家中是信奉罗马天主教。 \n\n大卫在2007年拍摄少年魔法师期间认识女演员露西·海尔（Lucy Hale），之后与其交往，于2009年分手。\n\n10岁时，大卫·亨利和SAG在凤凰城签订了合约，并开始走出去试镜。 9岁的时候，在沙加缅度进行商业拍摄，SAG董事建议大卫·亨利搬到洛杉矶。在10岁那年夏天，他和他的家人搬到了好莱坞。他预定他的前2支商业试镜，扮演主要角色为汉堡王和桂格燕麦。他初演电视节目为Providence。 \n\n到了13岁，大卫有了他的第一次重大突破，在福克斯公司的喜剧The Pitts饰演 Petey Pitt一角。大卫下出作品为的Hallmark movie为Monster Maker，和琳达布莱儿、乔治甘迺迪共同演出，并要求回来Hallmark movie公司。 \n\n在18岁时，大卫得到了迪士尼频道原创系列演出机会，该节目2007年10月12日首播。大卫2008年参加了迪士尼频道的游戏节目。他是绿色团队的队长，隔年，为旋风队队长。他在迪士尼原创电影《少年魔法师》之后在《酷爸的疯狂假期》中有饰演一角。\n"}
-{"id": 1,"title": "大卫·亨利","content": "大卫·亨利\n\n大卫·克莱顿·亨利（David Clayton Henrie，），美国演员。近来在迪士尼频道原创电视影集《少年魔法师》（Wizards of Waverly Place）当中演出贾斯汀·鲁索（Justin Russo）一角。\n\n大卫·亨利出生在加州Mission Viejo，在凤凰城长大。他的胞弟劳伦斯·亨利（Lorenzo Henrie）也是演员。大卫·亨利就读夏安传统学校。家中是信奉罗马天主教。 \n\n大卫在2007年拍摄少年魔法师期间认识女演员露西·海尔（Lucy Hale），之后与其交往，于2009年分手。\n\n10岁时，大卫·亨利和SAG在凤凰城签订了合约，并开始走出去试镜。 9岁的时候，在沙加缅度进行商业拍摄，SAG董事建议大卫·亨利搬到洛杉矶。在10岁那年夏天，他和他的家人搬到了好莱坞。他预定他的前2支商业试镜，扮演主要角色为汉堡王和桂格燕麦。他初演电视节目为Providence。 \n\n到了13岁，大卫有了他的第一次重大突破，在福克斯公司的喜剧The Pitts饰演 Petey Pitt一角。大卫下出作品为的Hallmark movie为Monster Maker，和琳达布莱儿、乔治甘迺迪共同演出，并要求回来Hallmark movie公司。 \n\n在18岁时，大卫得到了迪士尼频道原创系列演出机会，该节目2007年10月12日首播。大卫2008年参加了迪士尼频道的游戏节目。他是绿色团队的队长，隔年，为旋风队队长。他在迪士尼原创电影《少年魔法师》之后在《酷爸的疯狂假期》中有饰演一角。\n"}

data/sharegpt_zh_27k.json.REMOVED.git-id

@@ -0,0 +1 @@
+38c89869c6aeca2a3af9ea1e09afe460f9b46810

requirements.txt

@@ -1,16 +1,19 @@
 torch>=1.13.1
-transformers>=4.29.1
+transformers>=4.30.0
 datasets>=2.12.0
-accelerate>=0.19.0
-peft>=0.3.0
-trl>=0.4.4
+accelerate>=0.21.0
+peft==0.4.0
+trl>=0.7.1
+scipy
 sentencepiece
+protobuf
+tiktoken
 jieba
 rouge-chinese
 nltk
 gradio>=3.36.0
 uvicorn
-pydantic==1.10.7
-fastapi
+pydantic==1.10.11
+fastapi==0.95.1
 sse-starlette
 matplotlib

src/api_demo.py

@@ -1,13 +1,14 @@
-# coding=utf-8
-# Implements API for fine-tuned models in OpenAI's format. (https://platform.openai.com/docs/api-reference/chat)
-# Usage: python api_demo.py --model_name_or_path path_to_model --checkpoint_dir path_to_checkpoint
-# Visit http://localhost:8000/docs for document.
-
 import uvicorn
 
-from llmtuner import create_app
+from llmtuner import ChatModel, create_app
+
+
+def main():
+    chat_model = ChatModel()
+    app = create_app(chat_model)
+    uvicorn.run(app, host="0.0.0.0", port=8000, workers=1)
+    print("Visit http://localhost:8000/docs for API document.")
 
 
 if __name__ == "__main__":
-    app = create_app()
-    uvicorn.run(app, host="0.0.0.0", port=8000, workers=1)
+    main()

src/cli_demo.py

@@ -1,12 +1,8 @@
-# coding=utf-8
-# Implements stream chat in command line for fine-tuned models.
-# Usage: python cli_demo.py --model_name_or_path path_to_model --checkpoint_dir path_to_checkpoint
-
-from llmtuner import ChatModel, get_infer_args
+from llmtuner import ChatModel
 
 
 def main():
-    chat_model = ChatModel(*get_infer_args())
+    chat_model = ChatModel()
     history = []
     print("Welcome to the CLI application, use `clear` to remove the history, use `exit` to exit the application.")
 

src/export_model.py

@@ -1,16 +1,8 @@
-# coding=utf-8
-# Exports the fine-tuned model.
-# Usage: python export_model.py --checkpoint_dir path_to_checkpoint --output_dir path_to_save_model
-
-from llmtuner import get_train_args, load_model_and_tokenizer
+from llmtuner import export_model
 
 
 def main():
-    model_args, _, training_args, finetuning_args, _ = get_train_args()
-    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args)
-    model.save_pretrained(training_args.output_dir, max_shard_size="10GB")
-    tokenizer.save_pretrained(training_args.output_dir)
-    print("model and tokenizer have been saved at:", training_args.output_dir)
+    export_model()
 
 
 if __name__ == "__main__":

src/llmtuner/__init__.py

@@ -1,6 +1,9 @@
+# Level: api, webui > chat > tuner > dsets > extras, hparams
+
 from llmtuner.api import create_app
 from llmtuner.chat import ChatModel
-from llmtuner.tuner import get_train_args, get_infer_args, load_model_and_tokenizer, run_pt, run_sft, run_rm, run_ppo
+from llmtuner.tuner import export_model, run_exp
+from llmtuner.webui import create_ui, create_web_demo
 
 
-__version__ = "0.0.9"
+__version__ = "0.1.8"

src/llmtuner/api/app.py

@@ -5,10 +5,11 @@ from contextlib import asynccontextmanager
 from sse_starlette import EventSourceResponse
 from typing import List, Tuple
 
-from llmtuner.tuner import get_infer_args
 from llmtuner.extras.misc import torch_gc
-from llmtuner.chat.stream_chat import ChatModel
+from llmtuner.chat import ChatModel
 from llmtuner.api.protocol import (
+    Role,
+    Finish,
     ModelCard,
     ModelList,
     ChatMessage,
@@ -28,9 +29,7 @@ async def lifespan(app: FastAPI): # collects GPU memory
     torch_gc()
 
 
-def create_app():
-    chat_model = ChatModel(*get_infer_args())
-
+def create_app(chat_model: ChatModel) -> FastAPI:
     app = FastAPI(lifespan=lifespan)
 
     app.add_middleware(
@@ -48,28 +47,28 @@ def create_app():
 
     @app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
     async def create_chat_completion(request: ChatCompletionRequest):
-        if request.messages[-1].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
 
+        query = request.messages[-1].content
         prev_messages = request.messages[:-1]
-        if len(prev_messages) > 0 and prev_messages[0].role == "system":
-            prefix = prev_messages.pop(0).content
+        if len(prev_messages) > 0 and prev_messages[0].role == Role.SYSTEM:
+            system = prev_messages.pop(0).content
         else:
-            prefix = None
+            system = None
 
         history = []
         if len(prev_messages) % 2 == 0:
             for i in range(0, len(prev_messages), 2):
-                if prev_messages[i].role == "user" and prev_messages[i+1].role == "assistant":
+                if prev_messages[i].role == Role.USER and prev_messages[i+1].role == Role.ASSISTANT:
                     history.append([prev_messages[i].content, prev_messages[i+1].content])
 
         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(
@@ -80,23 +79,23 @@ def create_app():
 
         choice_data = ChatCompletionResponseChoice(
             index=0,
-            message=ChatMessage(role="assistant", content=response),
-            finish_reason="stop"
+            message=ChatMessage(role=Role.ASSISTANT, content=response),
+            finish_reason=Finish.STOP
         )
 
-        return ChatCompletionResponse(model=request.model, choices=[choice_data], usage=usage, object="chat.completion")
+        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="assistant"),
+            delta=DeltaMessage(role=Role.ASSISTANT),
             finish_reason=None
         )
-        chunk = ChatCompletionStreamResponse(model=request.model, choices=[choice_data], object="chat.completion.chunk")
+        chunk = ChatCompletionStreamResponse(model=request.model, choices=[choice_data])
         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
@@ -106,15 +105,15 @@ def create_app():
                 delta=DeltaMessage(content=new_text),
                 finish_reason=None
             )
-            chunk = ChatCompletionStreamResponse(model=request.model, choices=[choice_data], object="chat.completion.chunk")
+            chunk = ChatCompletionStreamResponse(model=request.model, choices=[choice_data])
             yield chunk.json(exclude_unset=True, ensure_ascii=False)
 
         choice_data = ChatCompletionResponseStreamChoice(
             index=0,
             delta=DeltaMessage(),
-            finish_reason="stop"
+            finish_reason=Finish.STOP
         )
-        chunk = ChatCompletionStreamResponse(model=request.model, choices=[choice_data], object="chat.completion.chunk")
+        chunk = ChatCompletionStreamResponse(model=request.model, choices=[choice_data])
         yield chunk.json(exclude_unset=True, ensure_ascii=False)
         yield "[DONE]"
 
@@ -122,5 +121,6 @@ def create_app():
 
 
 if __name__ == "__main__":
-    app = create_app()
+    chat_model = ChatModel()
+    app = create_app(chat_model)
     uvicorn.run(app, host="0.0.0.0", port=8000, workers=1)

src/llmtuner/api/protocol.py

@@ -1,6 +1,18 @@
 import time
+from enum import Enum
 from pydantic import BaseModel, Field
-from typing import List, Literal, Optional
+from typing import List, Optional
+
+
+class Role(str, Enum):
+    USER = "user"
+    ASSISTANT = "assistant"
+    SYSTEM = "system"
+
+
+class Finish(str, Enum):
+    STOP = "stop"
+    LENGTH = "length"
 
 
 class ModelCard(BaseModel):
@@ -19,12 +31,12 @@ class ModelList(BaseModel):
 
 
 class ChatMessage(BaseModel):
-    role: Literal["user", "assistant", "system"]
+    role: Role
     content: str
 
 
 class DeltaMessage(BaseModel):
-    role: Optional[Literal["user", "assistant", "system"]] = None
+    role: Optional[Role] = None
     content: Optional[str] = None
 
 
@@ -41,13 +53,13 @@ class ChatCompletionRequest(BaseModel):
 class ChatCompletionResponseChoice(BaseModel):
     index: int
     message: ChatMessage
-    finish_reason: Literal["stop", "length"]
+    finish_reason: Finish
 
 
 class ChatCompletionResponseStreamChoice(BaseModel):
     index: int
     delta: DeltaMessage
-    finish_reason: Optional[Literal["stop", "length"]] = None
+    finish_reason: Optional[Finish] = None
 
 
 class ChatCompletionResponseUsage(BaseModel):
@@ -58,7 +70,7 @@ class ChatCompletionResponseUsage(BaseModel):
 
 class ChatCompletionResponse(BaseModel):
     id: Optional[str] = "chatcmpl-default"
-    object: Literal["chat.completion"]
+    object: Optional[str] = "chat.completion"
     created: Optional[int] = Field(default_factory=lambda: int(time.time()))
     model: str
     choices: List[ChatCompletionResponseChoice]
@@ -67,7 +79,7 @@ class ChatCompletionResponse(BaseModel):
 
 class ChatCompletionStreamResponse(BaseModel):
     id: Optional[str] = "chatcmpl-default"
-    object: Literal["chat.completion.chunk"]
+    object: Optional[str] = "chat.completion.chunk"
     created: Optional[int] = Field(default_factory=lambda: int(time.time()))
     model: str
     choices: List[ChatCompletionResponseStreamChoice]

src/llmtuner/chat/stream_chat.py

@@ -1,36 +1,38 @@
+import torch
 from typing import Any, Dict, Generator, List, Optional, Tuple
 from threading import Thread
-from transformers import TextIteratorStreamer
+from transformers import GenerationConfig, TextIteratorStreamer
 
-from llmtuner.extras.misc import get_logits_processor
-from llmtuner.extras.template import Template
-from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments
-from llmtuner.tuner import load_model_and_tokenizer
+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
 
 
 class ChatModel:
 
-    def __init__(
-        self,
-        model_args: ModelArguments,
-        data_args: DataArguments,
-        finetuning_args: FinetuningArguments,
-        generating_args: GeneratingArguments
-    ) -> None:
+    def __init__(self, args: Optional[Dict[str, Any]] = None) -> None:
+        model_args, data_args, finetuning_args, self.generating_args = get_infer_args(args)
         self.model, self.tokenizer = load_model_and_tokenizer(model_args, finetuning_args)
-        self.template = Template(data_args.prompt_template)
-        self.source_prefix = data_args.source_prefix if data_args.source_prefix else ""
-        self.generating_args = generating_args
+        self.model = dispatch_model(self.model)
+        self.template = get_template_and_fix_tokenizer(data_args.template, self.tokenizer)
+        self.system_prompt = data_args.system_prompt
 
     def process_args(
-        self, query: str, history: Optional[List[Tuple[str, str]]] = None, prefix: Optional[str] = None, **input_kwargs
+        self,
+        query: str,
+        history: Optional[List[Tuple[str, str]]] = None,
+        system: Optional[str] = None,
+        **input_kwargs
     ) -> Tuple[Dict[str, Any], int]:
-        prefix = prefix if prefix else self.source_prefix
+        system = system or self.system_prompt
 
-        inputs = self.tokenizer([self.template.get_prompt(query, history, prefix)], return_tensors="pt")
-        inputs = inputs.to(self.model.device)
-        prompt_length = len(inputs["input_ids"][0])
+        prompt, _ = self.template.encode_oneturn(
+            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])
 
+        do_sample = input_kwargs.pop("do_sample", None)
         temperature = input_kwargs.pop("temperature", None)
         top_p = input_kwargs.pop("top_p", None)
         top_k = input_kwargs.pop("top_k", None)
@@ -38,45 +40,61 @@ 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=inputs["input_ids"],
-            temperature=temperature if temperature else gen_kwargs["temperature"],
-            top_p=top_p if top_p else gen_kwargs["top_p"],
-            top_k=top_k if top_k else gen_kwargs["top_k"],
-            repetition_penalty=repetition_penalty if repetition_penalty else gen_kwargs["repetition_penalty"],
-            logits_processor=get_logits_processor()
+        generating_args = self.generating_args.to_dict()
+        generating_args.update(dict(
+            do_sample=do_sample if do_sample is not None else generating_args["do_sample"],
+            temperature=temperature or generating_args["temperature"],
+            top_p=top_p or generating_args["top_p"],
+            top_k=top_k or generating_args["top_k"],
+            repetition_penalty=repetition_penalty or generating_args["repetition_penalty"],
+            eos_token_id=[self.tokenizer.eos_token_id] + self.tokenizer.additional_special_tokens_ids,
+            pad_token_id=self.tokenizer.pad_token_id
         ))
 
         if max_length:
-            gen_kwargs.pop("max_new_tokens", None)
-            gen_kwargs["max_length"] = max_length
+            generating_args.pop("max_new_tokens", None)
+            generating_args["max_length"] = max_length
 
         if max_new_tokens:
-            gen_kwargs.pop("max_length", None)
-            gen_kwargs["max_new_tokens"] = max_new_tokens
+            generating_args.pop("max_length", None)
+            generating_args["max_new_tokens"] = max_new_tokens
+
+        gen_kwargs = dict(
+            inputs=input_ids,
+            generation_config=GenerationConfig(**generating_args),
+            logits_processor=get_logits_processor()
+        )
 
         return gen_kwargs, prompt_length
 
+    @torch.inference_mode()
     def chat(
-        self, query: str, history: Optional[List[Tuple[str, str]]] = None, prefix: Optional[str] = None, **input_kwargs
+        self,
+        query: str,
+        history: Optional[List[Tuple[str, 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)
         response_length = len(outputs)
         return response, (prompt_length, response_length)
 
+    @torch.inference_mode()
     def stream_chat(
-        self, query: str, history: Optional[List[Tuple[str, str]]] = None, prefix: Optional[str] = None, **input_kwargs
+        self,
+        query: str,
+        history: Optional[List[Tuple[str, 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
 
         thread = Thread(target=self.model.generate, kwargs=gen_kwargs)
         thread.start()
 
-        for new_text in streamer:
-            yield new_text
+        yield from streamer

src/llmtuner/dsets/__init__.py

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

src/llmtuner/dsets/callbacks.py

@@ -1,63 +0,0 @@
-import os
-import json
-import time
-from datetime import timedelta
-
-from transformers import (
-    TrainerCallback,
-    TrainerControl,
-    TrainerState,
-    TrainingArguments
-)
-
-
-class LogCallback(TrainerCallback):
-
-    def __init__(self, runner=None):
-        self.runner = runner
-        self.start_time = time.time()
-        self.tracker = {}
-
-    def on_step_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
-        r"""
-        Event called at the beginning of a training step. If using gradient accumulation, one training step
-        might take several inputs.
-        """
-        if self.runner is not None and self.runner.aborted:
-            control.should_epoch_stop = True
-            control.should_training_stop = True
-
-    def on_substep_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
-        r"""
-        Event called at the end of an substep during gradient accumulation.
-        """
-        if self.runner is not None and self.runner.aborted:
-            control.should_epoch_stop = True
-            control.should_training_stop = True
-
-    def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs) -> None:
-        r"""
-        Event called after logging the last logs.
-        """
-        if "loss" not in state.log_history[-1]:
-            return
-        cur_time = time.time()
-        cur_steps = state.log_history[-1].get("step")
-        elapsed_time = cur_time - self.start_time
-        avg_time_per_step = elapsed_time / cur_steps if cur_steps != 0 else 0
-        remaining_steps = state.max_steps - cur_steps
-        remaining_time = remaining_steps * avg_time_per_step
-        self.tracker = {
-            "current_steps": cur_steps,
-            "total_steps": state.max_steps,
-            "loss": state.log_history[-1].get("loss", None),
-            "reward": state.log_history[-1].get("reward", None),
-            "learning_rate": state.log_history[-1].get("learning_rate", None),
-            "epoch": state.log_history[-1].get("epoch", None),
-            "percentage": round(cur_steps / state.max_steps * 100, 2) if state.max_steps != 0 else 100,
-            "elapsed_time": str(timedelta(seconds=int(elapsed_time))),
-            "remaining_time": str(timedelta(seconds=int(remaining_time)))
-        }
-        os.makedirs(args.output_dir, exist_ok=True)
-        with open(os.path.join(args.output_dir, "trainer_log.jsonl"), "a", encoding="utf-8") as f:
-            f.write(json.dumps(self.tracker) + "\n")

src/llmtuner/dsets/loader.py

@@ -1,40 +1,27 @@
 import os
-import hashlib
-from typing import List
+from typing import TYPE_CHECKING, List, Union
 
-from datasets import Dataset, concatenate_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
-from llmtuner.hparams import ModelArguments, DataArguments
+
+if TYPE_CHECKING:
+    from datasets import Dataset, IterableDataset
+    from llmtuner.hparams import ModelArguments, DataArguments
 
 
 logger = get_logger(__name__)
 
 
 def get_dataset(
-    model_args: ModelArguments,
-    data_args: DataArguments
-) -> Dataset:
-
-    def checksum(file_path, hash):
-        with open(file_path, "rb") as datafile:
-            binary_data = datafile.read()
-        sha1 = hashlib.sha1(binary_data).hexdigest()
-        if sha1 != hash:
-            logger.warning("Checksum failed for {}. It may vary depending on the platform.".format(file_path))
-
-    ext2type = {
-        "csv": "csv",
-        "json": "json",
-        "jsonl": "json",
-        "txt": "text"
-    }
-
+    model_args: "ModelArguments",
+    data_args: "DataArguments"
+) -> 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))
 
         if dataset_attr.load_from == "hf_hub":
@@ -47,60 +34,59 @@ def get_dataset(
             data_path = None
             data_files: List[str] = []
 
-            if os.path.isdir(os.path.join(data_args.dataset_dir, dataset_attr.dataset_name)):
+            if os.path.isdir(os.path.join(data_args.dataset_dir, dataset_attr.dataset_name)): # directory
                 for file_name in os.listdir(os.path.join(data_args.dataset_dir, dataset_attr.dataset_name)):
                     data_files.append(os.path.join(data_args.dataset_dir, dataset_attr.dataset_name, file_name))
-
                     if data_path is None:
-                        data_path = ext2type.get(data_files[0].split(".")[-1], None)
+                        data_path = EXT2TYPE.get(file_name.split(".")[-1], None)
                     else:
-                        assert data_path == ext2type.get(data_files[-1].split(".")[-1], None), "file type does not match."
-            elif os.path.isfile(os.path.join(data_args.dataset_dir, dataset_attr.dataset_name)):
+                        assert data_path == EXT2TYPE.get(file_name.split(".")[-1], None), "file type does not match."
+            elif os.path.isfile(os.path.join(data_args.dataset_dir, dataset_attr.dataset_name)): # single file
                 data_files.append(os.path.join(data_args.dataset_dir, dataset_attr.dataset_name))
-                data_path = ext2type.get(data_files[0].split(".")[-1], None)
+                data_path = EXT2TYPE.get(dataset_attr.dataset_name.split(".")[-1], None)
             else:
                 raise ValueError("File not found.")
 
             assert data_path, "File extension must be txt, csv, json or jsonl."
-
-            if len(data_files) == 1 and dataset_attr.dataset_sha1 is not None:
-                checksum(data_files[0], dataset_attr.dataset_sha1)
-            else:
-                logger.warning("Checksum failed: missing SHA-1 hash value in dataset_info.json or too many files.")
+            checksum(data_files, dataset_attr.dataset_sha1)
         else:
             raise NotImplementedError
 
-        raw_datasets = load_dataset(
+        dataset = load_dataset(
             data_path,
             data_files=data_files,
+            split=data_args.split,
             cache_dir=model_args.cache_dir,
+            streaming=data_args.streaming,
             use_auth_token=True if model_args.use_auth_token else None
         )
-        dataset = raw_datasets[data_args.split]
 
         if max_samples is not None:
             max_samples_temp = min(len(dataset), max_samples)
             dataset = dataset.select(range(max_samples_temp))
 
-        dummy_data = [None] * len(dataset)
-        prefix_data = [dataset_attr.source_prefix] * len(dataset)
-        for column_name, target_name in [
-            ("prompt_column", "prompt"),
-            ("query_column", "query"),
-            ("response_column", "response"),
-            ("history_column", "history")
-        ]: # every dataset will have 4 columns same as each other
-            if getattr(dataset_attr, column_name) != target_name:
-                if getattr(dataset_attr, column_name):
-                    dataset = dataset.rename_column(getattr(dataset_attr, column_name), target_name)
-                else: # None or empty string
-                    dataset = dataset.add_column(target_name, dummy_data)
-        dataset = dataset.add_column("prefix", prefix_data)
+        for column_name in ["prompt", "query", "response", "history"]: # align datasets
+            if getattr(dataset_attr, column_name) and getattr(dataset_attr, column_name) != column_name:
+                dataset = dataset.rename_column(getattr(dataset_attr, column_name), column_name)
+
+        if dataset_attr.system_prompt: # add system prompt
+            if data_args.streaming:
+                dataset = dataset.map(lambda _: {"system": dataset_attr.system_prompt})
+            else:
+                dataset = dataset.add_column("system", [dataset_attr.system_prompt] * len(dataset))
+
         all_datasets.append(dataset)
 
     if len(data_args.dataset_list) == 1:
-        all_datasets = all_datasets[0]
+        return all_datasets[0]
+    elif data_args.mix_strategy == "concat":
+        if data_args.streaming:
+            logger.warning("The samples between different datasets will not be mixed in streaming mode.")
+        return concatenate_datasets(all_datasets)
+    elif data_args.mix_strategy.startswith("interleave"):
+        if not data_args.streaming:
+            logger.warning("We recommend using `mix_strategy=concat` in non-streaming mode.")
+        stopping_strategy = "first_exhausted" if data_args.mix_strategy.endswith("under") else "all_exhausted"
+        return interleave_datasets(all_datasets, data_args.interleave_probs, stopping_strategy=stopping_strategy)
     else:
-        all_datasets = concatenate_datasets(all_datasets)
-
-    return all_datasets
+        raise ValueError("Unknown mixing strategy.")

src/llmtuner/dsets/preprocess.py

@@ -1,172 +1,193 @@
-from typing import Literal
+import tiktoken
+from typing import TYPE_CHECKING, Any, Dict, Generator, List, Literal, Union
 from itertools import chain
-from transformers import Seq2SeqTrainingArguments
-from transformers.tokenization_utils import PreTrainedTokenizer
-
-from datasets import Dataset
 
 from llmtuner.extras.constants import IGNORE_INDEX
-from llmtuner.extras.template import Template
-from llmtuner.hparams import DataArguments
+from llmtuner.extras.template import get_template_and_fix_tokenizer
+
+if TYPE_CHECKING:
+    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,
-    tokenizer: PreTrainedTokenizer,
-    data_args: DataArguments,
-    training_args: Seq2SeqTrainingArguments,
+    dataset: Union["Dataset", "IterableDataset"],
+    tokenizer: "PreTrainedTokenizer",
+    data_args: "DataArguments",
+    training_args: "Seq2SeqTrainingArguments",
     stage: Literal["pt", "sft", "rm", "ppo"]
-) -> Dataset:
+) -> Union["Dataset", "IterableDataset"]:
+    column_names = list(next(iter(dataset)).keys())
+    template = get_template_and_fix_tokenizer(data_args.template, tokenizer)
 
-    column_names = list(dataset.column_names)
-    prompt_template = Template(data_args.prompt_template)
-
-    # support question with a single answer or multiple answers
-    def get_dialog(examples):
+    def construct_example(examples: Dict[str, List[Any]]) -> Generator[Any, None, None]:
         for i in range(len(examples["prompt"])):
-            if examples["prompt"][i] and examples["response"][i]:
-                query, answer = examples["prompt"][i], examples["response"][i]
-                query = query + "\n" + examples["query"][i] if examples["query"][i] else query
-                prefix = examples["prefix"][i] if examples["prefix"][i] else ""
-                dialog = prompt_template.get_dialog(query, answer, examples["history"][i], prefix)
-                yield dialog
+            query, response = examples["prompt"][i], examples["response"][i]
+            query = query + "\n" + examples["query"][i] if "query" in examples and examples["query"][i] else query
+            history = examples["history"][i] if "history" in examples else None
+            system = examples["system"][i] if "system" in examples else None
+            yield query, response, history, system
 
-    def preprocess_pretrain_dataset(examples):
-        # build grouped texts with format `<bos> X1 X2 X3 ...` (without <eos>)
-        text_ids = tokenizer(examples["prompt"], add_special_tokens=False)["input_ids"]
-        concatenated_ids = list(chain(*text_ids))
-        total_length = len(concatenated_ids)
-        block_size = data_args.max_source_length - 1
+    def preprocess_pretrain_dataset(examples: Dict[str, List[Any]]) -> Dict[str, Any]:
+        # 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=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()}
+        total_length = len(concatenated_examples[list(concatenated_examples.keys())[0]])
+        block_size = data_args.max_source_length
         # we drop the small remainder, and if the total_length < block_size, we exclude this batch
         total_length = (total_length // block_size) * block_size
         # split by chunks of max_source_length
-        result = [[tokenizer.bos_token_id] + concatenated_ids[i: i + block_size]
-                  for i in range(0, total_length, block_size)]
-        return {
-            "input_ids": result,
-            "labels": result.copy()
+        result = {
+            k: [t[i: i + block_size] for i in range(0, total_length, block_size)]
+            for k, t in concatenated_examples.items()
         }
+        return result
 
-    def preprocess_supervised_dataset(examples):
+    def preprocess_supervised_dataset(examples: Dict[str, List[Any]]) -> Dict[str, Any]:
         # build inputs with format `<bos> X Y <eos>` and labels with format `<ignore> ... <ignore> Y <eos>`
-        # for input with history, we build multiple input-label pairs just like:
-        # https://github.com/lm-sys/FastChat/blob/f17c092f64840fa6354ed52789dccb2daa793d0b/fastchat/train/train.py#L112
-        model_inputs = {"input_ids": [], "labels": []}
+        # for multiturn examples, we only mask the prompt part in each prompt-response pair.
+        model_inputs = {"input_ids": [], "attention_mask": [], "labels": []}
         max_length = data_args.max_source_length + data_args.max_target_length
 
-        for dialog in get_dialog(examples):
+        for query, response, history, system in construct_example(examples):
             input_ids, labels = [], []
 
-            for i in range(len(dialog) // 2):
-                source_ids = tokenizer.encode(text=dialog[2*i], add_special_tokens=(i == 0))
-                target_ids = tokenizer.encode(text=dialog[2*i+1], add_special_tokens=False)
-
+            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 - 1: # eos token
-                    target_ids = target_ids[:data_args.max_target_length - 1]
+                if len(target_ids) > data_args.max_target_length:
+                    target_ids = target_ids[:data_args.max_target_length]
 
-                if len(input_ids) + len(source_ids) + len(target_ids) + 1 > max_length:
+                if len(input_ids) + len(source_ids) + len(target_ids) > max_length:
                     break
 
-                input_ids += source_ids + target_ids + [tokenizer.eos_token_id]
-                labels += [IGNORE_INDEX] * len(source_ids) + target_ids + [tokenizer.eos_token_id]
+                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 += 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))
             model_inputs["labels"].append(labels)
 
         return model_inputs
 
-    def preprocess_unsupervised_dataset(examples):
-        # build inputs with format `<bos> X` and labels with format `<bos> Y`
-        model_inputs = {"input_ids": [], "labels": []}
+    def preprocess_unsupervised_dataset(examples: Dict[str, List[Any]]) -> Dict[str, Any]:
+        # build inputs with format `<bos> X` and labels with format `Y <eos>`
+        model_inputs = {"input_ids": [], "attention_mask": [], "labels": []}
 
-        for dialog in get_dialog(examples):
-            prompt, answer = "".join(dialog[:-1]), dialog[-1]
-
-            source_ids = tokenizer.encode(text=prompt, add_special_tokens=True)
-            target_ids = tokenizer.encode(text=answer, add_special_tokens=True)
+        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)
 
         return model_inputs
 
     def preprocess_pairwise_dataset(examples):
-        # build input pairs with format `<bos> X Y1 <eos>` and `<bos> X Y2 <eos>`
-        model_inputs = {"accept_ids": [], "reject_ids": []}
-        for dialog in get_dialog(examples):
-            prompt, answer = "".join(dialog[:-1]), dialog[-1]
+        # build input pairs with format `<bos> X`, `Y1 <eos>` and `Y2 <eos>`
+        model_inputs = {"prompt_ids": [], "chosen_ids": [], "rejected_ids": []}
+        for query, response, history, system in construct_example(examples):
+            prompt_ids, chosen_ids = template.encode_oneturn(tokenizer, query, response[0], history, system)
+            _, rejected_ids = template.encode_oneturn(tokenizer, query, response[1], history, system)
 
-            source_ids = tokenizer.encode(text=prompt, add_special_tokens=True)
-            accept_ids = tokenizer.encode(text=answer[0], add_special_tokens=False)
-            reject_ids = tokenizer.encode(text=answer[1], add_special_tokens=False)
+            if len(prompt_ids) > data_args.max_source_length:
+                prompt_ids = prompt_ids[:data_args.max_source_length]
+            if len(chosen_ids) > data_args.max_target_length:
+                chosen_ids = chosen_ids[:data_args.max_target_length]
+            if len(rejected_ids) > data_args.max_target_length:
+                rejected_ids = rejected_ids[:data_args.max_target_length]
 
-            if len(source_ids) > data_args.max_source_length:
-                source_ids = source_ids[:data_args.max_source_length]
-            if len(accept_ids) > data_args.max_target_length - 1: # eos token
-                accept_ids = accept_ids[:data_args.max_target_length - 1]
-            if len(reject_ids) > data_args.max_target_length - 1: # eos token
-                reject_ids = reject_ids[:data_args.max_target_length - 1]
+            if template.efficient_eos:
+                chosen_ids += [tokenizer.eos_token_id]
+                rejected_ids += [tokenizer.eos_token_id]
 
-            accept_ids = source_ids + accept_ids + [tokenizer.eos_token_id]
-            reject_ids = source_ids + reject_ids + [tokenizer.eos_token_id]
-
-            model_inputs["accept_ids"].append(accept_ids)
-            model_inputs["reject_ids"].append(reject_ids)
+            model_inputs["prompt_ids"].append(prompt_ids)
+            model_inputs["chosen_ids"].append(chosen_ids)
+            model_inputs["rejected_ids"].append(rejected_ids)
         return model_inputs
 
     def print_supervised_dataset_example(example):
         print("input_ids:\n{}".format(example["input_ids"]))
         print("inputs:\n{}".format(tokenizer.decode(example["input_ids"], skip_special_tokens=False)))
         print("label_ids:\n{}".format(example["labels"]))
-        print("labels:\n{}".format(
-            tokenizer.decode([d if d != IGNORE_INDEX else tokenizer.pad_token_id for d in example["labels"]],
-                             skip_special_tokens=False)
-        ))
+        print("labels:\n{}".format(tokenizer.decode([
+            token_id if token_id != IGNORE_INDEX else tokenizer.pad_token_id for token_id in example["labels"]
+        ], skip_special_tokens=False)))
 
     def print_pairwise_dataset_example(example):
-        print("accept_ids:\n{}".format(example["accept_ids"]))
-        print("accepts:\n{}".format(tokenizer.decode(example["accept_ids"], skip_special_tokens=False)))
-        print("reject_ids:\n{}".format(example["reject_ids"]))
-        print("rejects:\n{}".format(tokenizer.decode(example["reject_ids"], skip_special_tokens=False)))
+        print("prompt_ids:\n{}".format(example["prompt_ids"]))
+        print("prompt:\n{}".format(tokenizer.decode(example["prompt_ids"], skip_special_tokens=False)))
+        print("chosen_ids:\n{}".format(example["chosen_ids"]))
+        print("chosen:\n{}".format(tokenizer.decode(example["chosen_ids"], skip_special_tokens=False)))
+        print("rejected_ids:\n{}".format(example["rejected_ids"]))
+        print("rejected:\n{}".format(tokenizer.decode(example["rejected_ids"], skip_special_tokens=False)))
 
     def print_unsupervised_dataset_example(example):
         print("input_ids:\n{}".format(example["input_ids"]))
         print("inputs:\n{}".format(tokenizer.decode(example["input_ids"], skip_special_tokens=False)))
 
     if stage == "pt":
+        dataset = dataset.filter(lambda example: example["prompt"])
         preprocess_function = preprocess_pretrain_dataset
-    elif stage == "sft":
-        preprocess_function = preprocess_unsupervised_dataset \
-            if training_args.predict_with_generate else preprocess_supervised_dataset
+        print_function = print_unsupervised_dataset_example
+    elif stage == "sft" and not training_args.predict_with_generate:
+        dataset = dataset.filter(lambda example: example["prompt"] and example["response"])
+        preprocess_function = preprocess_supervised_dataset
+        print_function = print_supervised_dataset_example
     elif stage == "rm":
+        dataset = dataset.filter(lambda example: example["prompt"] and len(example["response"]) > 1)
         preprocess_function = preprocess_pairwise_dataset
-    elif stage == "ppo":
+        print_function = print_pairwise_dataset_example
+    else:
+        dataset = dataset.filter(lambda example: example["prompt"])
         preprocess_function = preprocess_unsupervised_dataset
+        print_function = print_unsupervised_dataset_example
 
     with training_args.main_process_first(desc="dataset map pre-processing"):
+        kwargs = {}
+        if not data_args.streaming:
+            kwargs = dict(
+                num_proc=data_args.preprocessing_num_workers,
+                load_from_cache_file=not data_args.overwrite_cache,
+                desc="Running tokenizer on dataset"
+            )
+
         dataset = dataset.map(
             preprocess_function,
-            batched=True,
-            num_proc=data_args.preprocessing_num_workers,
+            batched=True,            
             remove_columns=column_names,
-            load_from_cache_file=not data_args.overwrite_cache,
-            desc="Running tokenizer on dataset"
+            **kwargs
         )
 
-        if stage == "pt":
-            print_unsupervised_dataset_example(dataset[0])
-        elif stage == "sft":
-            print_supervised_dataset_example(dataset[0])
-        elif stage == "rm":
-            print_pairwise_dataset_example(dataset[0])
-        elif stage == "ppo":
-            print_unsupervised_dataset_example(dataset[0])
-
+        print_function(next(iter(dataset)))
         return dataset

src/llmtuner/dsets/utils.py

@@ -0,0 +1,59 @@
+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
+    from transformers import TrainingArguments
+    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",
+    training_args: "TrainingArguments"
+) -> Dict[str, "Dataset"]:
+    if training_args.do_train:
+        if data_args.val_size > 1e-6: # Split the dataset
+            if data_args.streaming:
+                val_set = dataset.take(int(data_args.val_size))
+                train_set = dataset.skip(int(data_args.val_size))
+                dataset = dataset.shuffle(buffer_size=data_args.buffer_size, seed=training_args.seed)
+                return {"train_dataset": train_set, "eval_dataset": val_set}
+            else:
+                val_size = int(data_args.val_size) if data_args.val_size > 1 else data_args.val_size
+                dataset = dataset.train_test_split(test_size=val_size, seed=training_args.seed)
+                return {"train_dataset": dataset["train"], "eval_dataset": dataset["test"]}
+        else:
+            if data_args.streaming:
+                dataset = dataset.shuffle(buffer_size=data_args.buffer_size, seed=training_args.seed)
+            return {"train_dataset": dataset}
+    else: # do_eval or do_predict
+        return {"eval_dataset": dataset}

src/llmtuner/extras/callbacks.py

@@ -1,71 +1,148 @@
 import os
 import json
 import time
+from typing import TYPE_CHECKING
 from datetime import timedelta
 
-from transformers import (
-    TrainerCallback,
-    TrainerControl,
-    TrainerState,
-    TrainingArguments
-)
+from transformers import TrainerCallback
 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
+
+if TYPE_CHECKING:
+    from transformers import TrainingArguments, TrainerState, TrainerControl
+
+
+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):
         self.runner = runner
+        self.in_training = False
         self.start_time = time.time()
-        self.tracker = {}
+        self.cur_steps = 0
+        self.max_steps = 0
+        self.elapsed_time = ""
+        self.remaining_time = ""
 
-    def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+    def timing(self):
+        cur_time = time.time()
+        elapsed_time = cur_time - self.start_time
+        avg_time_per_step = elapsed_time / self.cur_steps if self.cur_steps != 0 else 0
+        remaining_time = (self.max_steps - self.cur_steps) * avg_time_per_step
+        self.elapsed_time = str(timedelta(seconds=int(elapsed_time)))
+        self.remaining_time = str(timedelta(seconds=int(remaining_time)))
+
+    def on_train_begin(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
         r"""
         Event called at the beginning of training.
         """
-        self.start_time = time.time()
+        if state.is_local_process_zero:
+            self.in_training = True
+            self.start_time = time.time()
+            self.max_steps = state.max_steps
+            if os.path.exists(os.path.join(args.output_dir, LOG_FILE_NAME)):
+                logger.warning("Previous log file in this folder will be deleted.")
+                os.remove(os.path.join(args.output_dir, LOG_FILE_NAME))
 
-    def on_step_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+    def on_train_end(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
         r"""
-        Event called at the beginning of a training step. If using gradient accumulation, one training step
-        might take several inputs.
+        Event called at the end of training.
         """
-        if self.runner is not None and self.runner.aborted:
-            control.should_epoch_stop = True
-            control.should_training_stop = True
+        if state.is_local_process_zero:
+            self.in_training = False
+            self.cur_steps = 0
+            self.max_steps = 0
 
-    def on_substep_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+    def on_substep_end(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
         r"""
         Event called at the end of an substep during gradient accumulation.
         """
-        if self.runner is not None and self.runner.aborted:
+        if state.is_local_process_zero and self.runner is not None and self.runner.aborted:
             control.should_epoch_stop = True
             control.should_training_stop = True
 
-    def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs) -> None:
+    def on_step_end(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+        r"""
+        Event called at the end of a training step.
+        """
+        if state.is_local_process_zero:
+            self.cur_steps = state.global_step
+            self.timing()
+            if self.runner is not None and self.runner.aborted:
+                control.should_epoch_stop = True
+                control.should_training_stop = True
+
+    def on_evaluate(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+        r"""
+        Event called after an evaluation phase.
+        """
+        if state.is_local_process_zero and not self.in_training:
+            self.cur_steps = 0
+            self.max_steps = 0
+
+    def on_predict(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", *other, **kwargs):
+        r"""
+        Event called after a successful prediction.
+        """
+        if state.is_local_process_zero and not self.in_training:
+            self.cur_steps = 0
+            self.max_steps = 0
+
+    def on_log(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs) -> None:
         r"""
         Event called after logging the last logs.
         """
-        cur_time = time.time()
-        cur_steps = state.log_history[-1].get("step")
-        elapsed_time = cur_time - self.start_time
-        avg_time_per_step = elapsed_time / cur_steps if cur_steps != 0 else 0
-        remaining_steps = state.max_steps - cur_steps
-        remaining_time = remaining_steps * avg_time_per_step
-        self.tracker = {
-            "current_steps": cur_steps,
-            "total_steps": state.max_steps,
-            "loss": state.log_history[-1].get("loss", None),
-            "eval_loss": state.log_history[-1].get("eval_loss", None),
-            "predict_loss": state.log_history[-1].get("predict_loss", None),
-            "reward": state.log_history[-1].get("reward", None),
-            "learning_rate": state.log_history[-1].get("learning_rate", None),
-            "epoch": state.log_history[-1].get("epoch", None),
-            "percentage": round(cur_steps / state.max_steps * 100, 2) if state.max_steps != 0 else 100,
-            "elapsed_time": str(timedelta(seconds=int(elapsed_time))),
-            "remaining_time": str(timedelta(seconds=int(remaining_time)))
-        }
+        if not state.is_local_process_zero:
+            return
+
+        logs = dict(
+            current_steps=self.cur_steps,
+            total_steps=self.max_steps,
+            loss=state.log_history[-1].get("loss", None),
+            eval_loss=state.log_history[-1].get("eval_loss", None),
+            predict_loss=state.log_history[-1].get("predict_loss", None),
+            reward=state.log_history[-1].get("reward", None),
+            learning_rate=state.log_history[-1].get("learning_rate", None),
+            epoch=state.log_history[-1].get("epoch", None),
+            percentage=round(self.cur_steps / self.max_steps * 100, 2) if self.max_steps != 0 else 100,
+            elapsed_time=self.elapsed_time,
+            remaining_time=self.remaining_time
+        )
         os.makedirs(args.output_dir, exist_ok=True)
         with open(os.path.join(args.output_dir, "trainer_log.jsonl"), "a", encoding="utf-8") as f:
-            f.write(json.dumps(self.tracker) + "\n")
+            f.write(json.dumps(logs) + "\n")
+
+    def on_prediction_step(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
+        r"""
+        Event called after a prediction step.
+        """
+        eval_dataloader = kwargs.pop("eval_dataloader", None)
+        if state.is_local_process_zero and has_length(eval_dataloader) and not self.in_training:
+            if self.max_steps == 0:
+                self.max_steps = len(eval_dataloader)
+            self.cur_steps += 1
+            self.timing()

src/llmtuner/extras/constants.py

@@ -1,7 +1,82 @@
 IGNORE_INDEX = -100
 
-VALUE_HEAD_FILE_NAME = "value_head.bin"
+LOG_FILE_NAME = "trainer_log.jsonl"
 
-FINETUNING_ARGS_NAME = "finetuning_args.json"
+LAYERNORM_NAMES = ["norm", "ln_f", "ln_attn", "ln_mlp"]
 
-LAYERNORM_NAMES = ["norm", "ln_f", "ln_attn", "ln_mlp"] # for LLaMA, BLOOM and Falcon settings
+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",
+    "LLaMA-30B": "huggyllama/llama-30b",
+    "LLaMA-65B": "huggyllama/llama-65b",
+    "LLaMA2-7B": "meta-llama/Llama-2-7b-hf",
+    "LLaMA2-13B": "meta-llama/Llama-2-13b-hf",
+    "LLaMA2-70B": "meta-llama/Llama-2-70b-hf",
+    "LLaMA2-7B-Chat": "meta-llama/Llama-2-7b-chat-hf",
+    "LLaMA2-13B-Chat": "meta-llama/Llama-2-13b-chat-hf",
+    "LLaMA2-70B-Chat": "meta-llama/Llama-2-70b-chat-hf",
+    "ChineseLLaMA2-7B": "ziqingyang/chinese-llama-2-7b",
+    "ChineseLLaMA2-13B": "ziqingyang/chinese-llama-2-13b",
+    "ChineseLLaMA2-7B-Chat": "ziqingyang/chinese-alpaca-2-7b",
+    "ChineseLLaMA2-13B-Chat": "ziqingyang/chinese-alpaca-2-13b",
+    "BLOOM-560M": "bigscience/bloom-560m",
+    "BLOOM-3B": "bigscience/bloom-3b",
+    "BLOOM-7B1": "bigscience/bloom-7b1",
+    "BLOOMZ-560M": "bigscience/bloomz-560m",
+    "BLOOMZ-3B": "bigscience/bloomz-3b",
+    "BLOOMZ-7B1-mt": "bigscience/bloomz-7b1-mt",
+    "Falcon-7B": "tiiuae/falcon-7b",
+    "Falcon-7B-Chat": "tiiuae/falcon-7b-instruct",
+    "Falcon-40B": "tiiuae/falcon-40b",
+    "Falcon-40B-Chat": "tiiuae/falcon-40b-instruct",
+    "Baichuan-7B": "baichuan-inc/Baichuan-7B",
+    "Baichuan-13B": "baichuan-inc/Baichuan-13B-Base",
+    "Baichuan-13B-Chat": "baichuan-inc/Baichuan-13B-Chat",
+    "Baichuan2-7B": "baichuan-inc/Baichuan2-7B-Base",
+    "Baichuan2-13B": "baichuan-inc/Baichuan2-13B-Base",
+    "Baichuan2-7B-Chat": "baichuan-inc/Baichuan2-7B-Chat",
+    "Baichuan2-13B-Chat": "baichuan-inc/Baichuan2-13B-Chat",
+    "InternLM-7B": "internlm/internlm-7b",
+    "InternLM-7B-Chat": "internlm/internlm-chat-7b",
+    "Qwen-7B": "Qwen/Qwen-7B",
+    "Qwen-7B-Chat": "Qwen/Qwen-7B-Chat",
+    "XVERSE-13B": "xverse/XVERSE-13B",
+    "XVERSE-13B-Chat": "xverse/XVERSE-13B-Chat",
+    "ChatGLM2-6B-Chat": "THUDM/chatglm2-6b"
+}
+
+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"
+}

src/llmtuner/extras/logging.py

@@ -2,8 +2,33 @@ import sys
 import logging
 
 
-def get_logger(name: str) -> logging.Logger:
+class LoggerHandler(logging.Handler):
 
+    def __init__(self):
+        super().__init__()
+        self.log = ""
+
+    def reset(self):
+        self.log = ""
+
+    def emit(self, record):
+        if record.name == "httpx":
+            return
+        log_entry = self.format(record)
+        self.log += log_entry
+        self.log += "\n\n"
+
+
+def reset_logging():
+    r"""
+    Removes basic config of root logger
+    """
+    root = logging.getLogger()
+    list(map(root.removeHandler, root.handlers))
+    list(map(root.removeFilter, root.filters))
+
+
+def get_logger(name: str) -> logging.Logger:
     formatter = logging.Formatter(
         fmt="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
         datefmt="%m/%d/%Y %H:%M:%S"

src/llmtuner/extras/misc.py

@@ -1,11 +1,10 @@
+import gc
 import torch
-from typing import List, Optional
+from typing import TYPE_CHECKING, Tuple
+from transformers import InfNanRemoveLogitsProcessor, LogitsProcessorList
 
-from transformers.modeling_utils import PreTrainedModel
-from transformers.generation.utils import LogitsProcessorList
-from transformers.generation.logits_process import LogitsProcessor
-
-from llmtuner.extras.constants import LAYERNORM_NAMES
+if TYPE_CHECKING:
+    from transformers.modeling_utils import PreTrainedModel
 
 
 class AverageMeter:
@@ -28,78 +27,64 @@ class AverageMeter:
         self.avg = self.sum / self.count
 
 
-# Avoid runtime error in model.generate(do_sample=True).
-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
-
-
-def print_trainable_params(model: torch.nn.Module) -> None:
+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.
+    """
     trainable_params, all_param = 0, 0
     for param in model.parameters():
         num_params = param.numel()
         # if using DS Zero 3 and the weights are initialized empty
         if num_params == 0 and hasattr(param, "ds_numel"):
             num_params = param.ds_numel
+
+        # Due to the design of 4bit linear layers from bitsandbytes, multiply the number of parameters by 2
+        if param.__class__.__name__ == "Params4bit":
+            num_params = num_params * 2
+
         all_param += num_params
         if param.requires_grad:
             trainable_params += num_params
-    print("trainable params: {:d} || all params: {:d} || trainable%: {:.4f}".format(
-                trainable_params, all_param, 100 * trainable_params / all_param))
+
+    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_embedding_layer_name: Optional[str] = "lm_head",
-    use_gradient_checkpointing: Optional[bool] = True,
-    layer_norm_names: Optional[List[str]] = LAYERNORM_NAMES
-) -> PreTrainedModel:
+def get_logits_processor() -> LogitsProcessorList:
+    logits_processor = LogitsProcessorList()
+    logits_processor.append(InfNanRemoveLogitsProcessor())
+    return logits_processor
 
-    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_embedding_layer_name):
-        output_embedding_layer: torch.nn.Linear = getattr(model, output_embedding_layer_name)
-        input_dtype = output_embedding_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_embedding_layer_name, CastOutputToFloat(output_embedding_layer))
-
-    return model
 
 def torch_gc() -> None:
     r"""
     Collects GPU memory.
     """
+    gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
         torch.cuda.ipc_collect()
+
+
+def dispatch_model(model: "PreTrainedModel") -> "PreTrainedModel":
+    r"""
+    Dispatches a pre-trained model to GPUs with balanced memory.
+    Borrowed from: https://github.com/huggingface/transformers/blob/v4.31.0/src/transformers/modeling_utils.py#L2803
+    """
+    if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False): # do nothing
+        return model
+
+    if torch.cuda.device_count() > 1:
+        from accelerate import dispatch_model
+        from accelerate.utils import infer_auto_device_map, get_balanced_memory
+
+        if model._no_split_modules is None:
+            raise ValueError("The model class needs to implement the `_no_split_modules` attribute.")
+
+        kwargs = {"dtype": model.dtype, "no_split_module_classes": model._no_split_modules}
+        max_memory = get_balanced_memory(model, **kwargs)
+        # Make sure tied weights are tied before creating the device map.
+        model.tie_weights()
+        device_map = infer_auto_device_map(model, max_memory=max_memory, **kwargs)
+        return dispatch_model(model, device_map)
+    else:
+        return model.cuda()

src/llmtuner/extras/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

src/llmtuner/extras/ploting.py

@@ -1,4 +1,5 @@
 import os
+import math
 import json
 import matplotlib.pyplot as plt
 from typing import List, Optional
@@ -10,12 +11,13 @@ from llmtuner.extras.logging import get_logger
 logger = get_logger(__name__)
 
 
-def smooth(scalars: List[float], weight: Optional[float] = 0.9) -> List[float]:
+def smooth(scalars: List[float]) -> List[float]:
     r"""
     EMA implementation according to TensorBoard.
     """
     last = scalars[0]
     smoothed = list()
+    weight = 1.8 * (1 / (1 + math.exp(-0.05 * len(scalars))) - 0.5) # a sigmoid function
     for next_val in scalars:
         smoothed_val = last * weight + (1 - weight) * next_val
         smoothed.append(smoothed_val)

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]: # get state dict containing trainable parameters
-    state_dict = 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

src/llmtuner/extras/template.py

@@ -1,181 +1,603 @@
-from typing import List, Optional, Tuple
+import tiktoken
 from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+
+from llmtuner.extras.logging import get_logger
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedTokenizer
+
+
+logger = get_logger(__name__)
 
 
 @dataclass
 class 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
+    efficient_eos: bool
 
-    def __post_init__(self):
-
-        if self.name == "vanilla":
-            r"""
-            Supports language model inference without histories.
-            """
-            self._register_template(
-                prefix="",
-                prompt="{query}",
-                sep="",
-                use_history=False
-            )
-
-        elif self.name == "default":
-            r"""
-            Default template.
-            """
-            self._register_template(
-                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.",
-                prompt="Human: {query}\nAssistant: ",
-                sep="\n",
-                use_history=True
-            )
-
-        elif self.name == "alpaca":
-            r"""
-            Supports: https://huggingface.co/tatsu-lab/alpaca-7b-wdiff
-                      https://github.com/ymcui/Chinese-LLaMA-Alpaca
-            """
-            self._register_template(
-                prefix="Below is an instruction that describes a task. "
-                       "Write a response that appropriately completes the request.",
-                prompt="### Instruction:\n{query}\n\n### Response:\n",
-                sep="\n\n",
-                use_history=True
-            )
-
-        elif self.name == "vicuna":
-            r"""
-            Supports: https://huggingface.co/lmsys/vicuna-7b-delta-v1.1
-                      https://huggingface.co/lmsys/vicuna-13b-delta-v1.1
-            """
-            self._register_template(
-                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.",
-                prompt="USER: {query} ASSISTANT: ",
-                sep="</s>",
-                use_history=True
-            )
-
-        elif self.name == "belle":
-            r"""
-            Supports: https://huggingface.co/BelleGroup/BELLE-LLaMA-EXT-13B
-            """
-            self._register_template(
-                prefix="",
-                prompt="Human: {query}\n\nBelle: ",
-                sep="\n\n",
-                use_history=True
-            )
-
-        elif self.name == "linly":
-            r"""
-            Supports: https://github.com/CVI-SZU/Linly
-            """
-            self._register_template(
-                prefix="",
-                prompt="User: {query}\nBot: ",
-                sep="\n",
-                use_history=True
-            )
-
-        elif self.name == "billa":
-            r"""
-            Supports: https://github.com/Neutralzz/BiLLa
-            """
-            self._register_template(
-                prefix="",
-                prompt="Human: {query}\nAssistant: ",
-                sep="\n",
-                use_history=True
-            )
-
-        elif self.name == "ziya":
-            r"""
-            Supports: https://huggingface.co/IDEA-CCNL/Ziya-LLaMA-13B-v1
-            """
-            self._register_template(
-                prefix="",
-                prompt="<human>:{query}\n<bot>:",
-                sep="\n",
-                use_history=True
-            )
-
-        elif self.name == "aquila":
-            r"""
-            Supports: https://huggingface.co/qhduan/aquilachat-7b
-            """
-            self._register_template(
-                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.",
-                prompt="Human: {query}###Assistant: ",
-                sep="###",
-                use_history=True
-            )
-
-        elif self.name == "intern":
-            r"""
-            Supports: https://huggingface.co/internlm/internlm-chat-7b
-            """
-            self._register_template(
-                prefix="",
-                prompt="<|User|>:{query}<eoh>\n<|Bot|>:",
-                sep="<eoa>\n",
-                use_history=True
-            )
-
-        elif self.name == "baichuan":
-            r"""
-            Supports: https://huggingface.co/baichuan-inc/Baichuan-13B-Chat
-            """
-            self._register_template(
-                prefix="",
-                prompt="<reserved_102>{query}<reserved_103>",
-                sep="",
-                use_history=True
-            )
-
-        else:
-            raise ValueError("Template {} does not exist.".format(self.name))
-
-    def get_prompt(
-        self, query: str, history: Optional[List[Tuple[str, str]]] = None, prefix: Optional[str] = ""
-    ) -> str:
+    def encode_oneturn(
+        self,
+        tokenizer: "PreTrainedTokenizer",
+        query: str,
+        resp: str,
+        history: Optional[List[Tuple[str, str]]] = None,
+        system: Optional[str] = None
+    ) -> Tuple[List[int], List[int]]:
         r"""
-        Returns a string containing prompt without response.
+        Returns a single pair of token ids representing prompt and response respectively.
         """
-        return "".join(self._format_example(query, history, prefix))
+        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, answer_ids = prompt_ids + encoded_pairs[-1][0], encoded_pairs[-1][1]
+        return prompt_ids, answer_ids
 
-    def get_dialog(
-        self, query: str, resp: str, history: Optional[List[Tuple[str, str]]] = None, prefix: Optional[str] = ""
-    ) -> List[str]:
+    def encode_multiturn(
+        self,
+        tokenizer: "PreTrainedTokenizer",
+        query: str,
+        resp: str,
+        history: Optional[List[Tuple[str, str]]] = None,
+        system: Optional[str] = None
+    ) -> List[Tuple[List[int], List[int]]]:
         r"""
-        Returns a list containing 2 * n elements where the 2k-th is a query and the (2k+1)-th is a response.
+        Returns multiple pairs of token ids representing prompts and responses respectively.
         """
-        return self._format_example(query, history, prefix) + [resp]
+        system, history = self._format(query, resp, history, system)
+        encoded_pairs = self._encode(tokenizer, system, history)
+        return encoded_pairs
 
-    def _register_template(
-        self, prefix: str, prompt: str, sep: str, use_history: Optional[bool] = True
-    ) -> None:
-        self.prefix = prefix
-        self.prompt = prompt
-        self.sep = sep
-        self.use_history = use_history
-
-    def _format_example(
-        self, query: str, history: Optional[List[Tuple[str, str]]] = None, prefix: Optional[str] = ""
-    ) -> List[str]:
-        prefix = prefix if prefix else self.prefix # use prefix if provided
-        prefix = prefix + self.sep if prefix else "" # add separator for non-empty prefix
+    def _format(
+        self,
+        query: str,
+        resp: str,
+        history: Optional[List[Tuple[str, str]]] = None,
+        system: Optional[str] = None
+    ) -> Tuple[str, List[Tuple[str, str]]]:
+        r"""
+        Aligns inputs to the standard format.
+        """
+        system = system or self.system # use system if provided
         history = history if (history and self.use_history) else []
-        history = history + [(query, "<dummy>")]
-        convs = []
-        for turn_idx, (user_query, bot_resp) in enumerate(history):
+        history = history + [(query, resp)]
+        return system, history
+
+    def _get_special_ids(
+        self,
+        tokenizer: "PreTrainedTokenizer"
+    ) -> Tuple[List[int], List[int]]:
+        if tokenizer.bos_token_id is not None and getattr(tokenizer, "add_bos_token", True):
+            bos_ids = [tokenizer.bos_token_id]
+        else: # baichuan, qwen and gpt2 models have no bos token
+            bos_ids = []
+
+        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",
+        system: str,
+        history: List[Tuple[str, str]]
+    ) -> List[Tuple[List[int], List[int]]]:
+        r"""
+        Encodes formatted inputs to pairs of token ids.
+        Turn 0: bos + prefix + sep + query    resp + eos
+        Turn t: sep + bos + query             resp + eos
+        """
+        bos_ids, eos_ids = self._get_special_ids(tokenizer)
+        sep_ids = self._convert_inputs_to_ids(tokenizer, context=self.sep)
+        encoded_pairs = []
+        for turn_idx, (query, resp) in enumerate(history):
             if turn_idx == 0:
-                convs.append(prefix + self.prompt.format(query=user_query))
-                convs.append(bot_resp)
+                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:
-                convs.append(self.sep + self.prompt.format(query=user_query))
-                convs.append(bot_resp)
-        return convs[:-1] # drop last
+                prefix_ids = sep_ids + bos_ids
+
+            query_ids = self._convert_inputs_to_ids(tokenizer, context=self.prompt, query=query, idx=str(turn_idx))
+            resp_ids = self._convert_inputs_to_ids(tokenizer, context=[resp])
+            encoded_pairs.append((prefix_ids + query_ids, resp_ids + eos_ids))
+        return encoded_pairs
+
+    def _convert_inputs_to_ids(
+        self,
+        tokenizer: "PreTrainedTokenizer",
+        context: List[Union[str, Dict[str, str]]],
+        system: Optional[str] = None,
+        query: Optional[str] = None,
+        idx: Optional[str] = None
+    ) -> List[int]:
+        r"""
+        Converts context to token ids.
+        """
+        if isinstance(getattr(tokenizer, "tokenizer", None), tiktoken.Encoding): # for tiktoken tokenizer (Qwen)
+            kwargs = dict(allowed_special="all")
+        else:
+            kwargs = dict(add_special_tokens=False)
+
+        token_ids = []
+        for elem in context:
+            if isinstance(elem, str):
+                if len(elem) == 0:
+                    continue
+                elem = elem.replace("{{system}}", system, 1) if system is not None else elem
+                elem = elem.replace("{{query}}", query, 1) if query is not None else elem
+                elem = elem.replace("{{idx}}", idx, 1) if idx is not None else elem
+                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
+
+
+@dataclass
+class Llama2Template(Template):
+
+    def _encode(
+        self,
+        tokenizer: "PreTrainedTokenizer",
+        system: str,
+        history: List[Tuple[str, str]]
+    ) -> List[Tuple[List[int], List[int]]]:
+        r"""
+        Encodes formatted inputs to pairs of token ids.
+        Turn 0: bos + prefix + query    resp + eos
+        Turn t: bos + query             resp + eos
+        """
+        bos_ids, eos_ids = self._get_special_ids(tokenizer)
+        encoded_pairs = []
+        for turn_idx, (query, resp) in enumerate(history):
+            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))
+        return encoded_pairs
+
+
+templates: Dict[str, Template] = {}
+
+
+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: 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,
+        efficient_eos=efficient_eos
+    )
+
+
+def get_template_and_fix_tokenizer(
+    name: str,
+    tokenizer: "PreTrainedTokenizer"
+) -> Template:
+    if tokenizer.eos_token_id is None:
+        tokenizer.eos_token = "<|endoftext|>"
+        logger.info("Add eos token: {}".format(tokenizer.eos_token))
+
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        logger.info("Add pad token: {}".format(tokenizer.pad_token))
+
+    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
+
+
+r"""
+Supports language model inference without histories.
+"""
+register_template(
+    name="vanilla",
+    prefix=[],
+    prompt=[
+        "{{query}}"
+    ],
+    system="",
+    sep=[],
+    use_history=False
+)
+
+
+r"""
+Default template.
+"""
+register_template(
+    name="default",
+    prefix=[
+        "{{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"
+    ]
+)
+
+
+r"""
+Supports: https://huggingface.co/meta-llama/Llama-2-7b-chat-hf
+          https://huggingface.co/meta-llama/Llama-2-13b-chat-hf
+          https://huggingface.co/meta-llama/Llama-2-70b-chat-hf
+"""
+register_template(
+    name="llama2",
+    prefix=[
+        "<<SYS>>\n{{system}}\n<</SYS>>\n\n"
+    ],
+    prompt=[
+        "[INST] {{query}} [/INST] "
+    ],
+    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>>\n{{system}}\n<</SYS>>\n\n"
+    ],
+    prompt=[
+        "[INST] {{query}} [/INST] "
+    ],
+    system="You are a helpful assistant. 你是一个乐于助人的助手。",
+    sep=[]
+)
+
+
+r"""
+Supports: https://huggingface.co/tatsu-lab/alpaca-7b-wdiff
+          https://github.com/ymcui/Chinese-LLaMA-Alpaca
+"""
+register_template(
+    name="alpaca",
+    prefix=[
+        "{{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"
+    ]
+)
+
+
+r"""
+Supports: https://huggingface.co/lmsys/vicuna-7b-delta-v1.1
+          https://huggingface.co/lmsys/vicuna-13b-delta-v1.1
+"""
+register_template(
+    name="vicuna",
+    prefix=[
+        "{{system}}"
+    ],
+    prompt=[
+        "USER: {{query}} ASSISTANT: "
+    ],
+    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=[]
+)
+
+
+r"""
+Supports: https://huggingface.co/BelleGroup/BELLE-LLaMA-EXT-13B
+"""
+register_template(
+    name="belle",
+    prefix=[
+        "{{system}}"
+    ],
+    prompt=[
+        "Human: {{query}}\n\nBelle: "
+    ],
+    system="",
+    sep=[
+        "\n\n"
+    ]
+)
+
+
+r"""
+Supports: https://github.com/CVI-SZU/Linly
+"""
+register_template(
+    name="linly",
+    prefix=[
+        "{{system}}"
+    ],
+    prompt=[
+        "User: {{query}}\nBot: "
+    ],
+    system="",
+    sep=[
+        "\n"
+    ]
+)
+
+
+r"""
+Supports: https://github.com/Neutralzz/BiLLa
+"""
+register_template(
+    name="billa",
+    prefix=[
+        "{{system}}"
+    ],
+    prompt=[
+        "Human: {{query}}\nAssistant: "
+    ],
+    system="",
+    sep=[
+        "\n"
+    ]
+)
+
+
+r"""
+Supports: https://huggingface.co/IDEA-CCNL/Ziya-LLaMA-13B-v1
+"""
+register_template(
+    name="ziya",
+    prefix=[
+        "{{system}}"
+    ],
+    prompt=[
+        {"token": "<human>"},
+        ":{{query}}\n",
+        {"token": "<bot>"},
+        ":"
+    ],
+    system="",
+    sep=[
+        "\n"
+    ]
+)
+
+
+r"""
+Supports: https://huggingface.co/qhduan/aquilachat-7b
+"""
+register_template(
+    name="aquila",
+    prefix=[
+        "{{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=[
+        "###"
+    ]
+)
+
+
+r"""
+Supports: https://huggingface.co/internlm/internlm-chat-7b
+"""
+register_template(
+    name="intern",
+    prefix=[
+        "{{system}}"
+    ],
+    prompt=[
+        "<|User|>:{{query}}",
+        {"token": "<eoh>"},
+        "\n<|Bot|>:"
+    ],
+    system="",
+    sep=[
+        {"token": "<eoa>"},
+        "\n"
+    ],
+    stop_words=[
+        "<eoa>"
+    ],
+    efficient_eos=True
+)
+
+
+r"""
+Supports: https://huggingface.co/baichuan-inc/Baichuan-13B-Chat
+"""
+register_template(
+    name="baichuan",
+    prefix=[
+        "{{system}}"
+    ],
+    prompt=[
+        {"token": "<reserved_102>"}, # user token
+        "{{query}}",
+        {"token": "<reserved_103>"}  # assistant token
+    ],
+    system="",
+    sep=[],
+    efficient_eos=True
+)
+
+
+r"""
+Supports: https://huggingface.co/baichuan-inc/Baichuan2-7B-Chat
+          https://huggingface.co/baichuan-inc/Baichuan2-13B-Chat
+"""
+register_template(
+    name="baichuan2",
+    prefix=[
+        "{{system}}"
+    ],
+    prompt=[
+        {"token": "<reserved_106>"}, # user token
+        "{{query}}",
+        {"token": "<reserved_107>"}  # assistant token
+    ],
+    system="",
+    sep=[],
+    efficient_eos=True
+)
+
+
+r"""
+Supports: https://huggingface.co/HuggingFaceH4/starchat-alpha
+          https://huggingface.co/HuggingFaceH4/starchat-beta
+"""
+register_template(
+    name="starchat",
+    prefix=[
+        {"token": "<|system|>"},
+        "\n{{system}}",
+    ],
+    prompt=[
+        {"token": "<|user|>"},
+        "\n{{query}}",
+        {"token": "<|end|>"},
+        "\n",
+        {"token": "<|assistant|>"}
+    ],
+    system="",
+    sep=[
+        {"token": "<|end|>"},
+        "\n"
+    ],
+    stop_words=[
+        "<|end|>"
+    ],
+    efficient_eos=True
+)
+
+
+r"""
+Supports: https://huggingface.co/Qwen/Qwen-7B-Chat
+"""
+register_template(
+    name="chatml",
+    prefix=[
+        {"token": "<|im_start|>"},
+        "system\n{{system}}"
+    ],
+    prompt=[
+        {"token": "<|im_start|>"},
+        "user\n{{query}}",
+        {"token": "<|im_end|>"},
+        "\n",
+        {"token": "<|im_start|>"},
+        "assistant\n"
+    ],
+    system="You are a helpful assistant.",
+    sep=[
+        {"token": "<|im_end|>"},
+        "\n"
+    ],
+    stop_words=[
+        "<|im_end|>"
+    ],
+    efficient_eos=True
+)
+
+
+r"""
+Supports: https://huggingface.co/THUDM/chatglm2-6b
+"""
+register_template(
+    name="chatglm2",
+    prefix=[
+        {"token": "[gMASK]"},
+        {"token": "sop"},
+        "{{system}}"
+    ],
+    prompt=[
+        "[Round {{idx}}]\n\n问：{{query}}\n\n答："
+    ],
+    system="",
+    sep=[
+        "\n\n"
+    ],
+    efficient_eos=True
+)
+
+
+r"""
+Supports: https://huggingface.co/xverse/XVERSE-13B-Chat
+"""
+register_template(
+    name="xverse",
+    prefix=[
+        "{{system}}"
+    ],
+    prompt=[
+        "Human: {{query}}\n\nAssistant: "
+    ],
+    system="",
+    sep=[]
+)

src/llmtuner/hparams/data_args.py

@@ -1,6 +1,6 @@
 import os
 import json
-from typing import List, Optional
+from typing import List, Literal, Optional
 from dataclasses import dataclass, field
 
 
@@ -10,25 +10,28 @@ 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_column = "instruction"
-        self.query_column = "input"
-        self.response_column = "output"
-        self.history_column = None
-
 
 @dataclass
 class DataArguments:
-    """
+    r"""
     Arguments pertaining to what data we are going to input our model for training and evaluation.
     """
+    template: Optional[str] = field(
+        default=None,
+        metadata={"help": "Which template to use for constructing prompts in training and inference."}
+    )
     dataset: Optional[str] = field(
-        default="alpaca_zh",
+        default="alpaca_en",
         metadata={"help": "The name of provided dataset(s) to use. Use commas to separate multiple datasets."}
     )
     dataset_dir: Optional[str] = field(
@@ -39,6 +42,22 @@ class DataArguments:
         default="train",
         metadata={"help": "Which dataset split to use for training and evaluation."}
     )
+    streaming: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Enable streaming mode."}
+    )
+    buffer_size: Optional[int] = field(
+        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(
+        default="concat",
+        metadata={"help": "Strategy to use in dataset mixing."}
+    )
+    interleave_probs: Optional[str] = field(
+        default=None,
+        metadata={"help": "Probabilities to sample data from datasets. Use commas to separate multiple datasets."}
+    )
     overwrite_cache: Optional[bool] = field(
         default=False,
         metadata={"help": "Overwrite the cached training and evaluation sets."}
@@ -67,17 +86,13 @@ 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."}
     )
-    dev_ratio: Optional[float] = field(
+    val_size: Optional[float] = field(
         default=0,
-        metadata={"help": "Proportion of the dataset to include in the development set, should be between 0.0 and 1.0."}
-    )
-    prompt_template: Optional[str] = field(
-        default="default",
-        metadata={"help": "Which template to use for constructing prompts in training and inference."}
+        metadata={"help": "Size of the development set, should be an integer or a float in range `[0,1)`."}
     )
 
     def init_for_training(self): # support mixing multiple datasets
@@ -85,12 +100,12 @@ 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(",")]
 
         self.dataset_list: List[DatasetAttr] = []
         for i, name in enumerate(dataset_names):
@@ -108,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_column = dataset_info[name]["columns"].get("prompt", None)
-                dataset_attr.query_column = dataset_info[name]["columns"].get("query", None)
-                dataset_attr.response_column = dataset_info[name]["columns"].get("response", None)
-                dataset_attr.history_column = dataset_info[name]["columns"].get("history", None)
+                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)
 
-            self.dataset_list.append(dataset_attr)
+            dataset_attr.ranking = dataset_info[name].get("ranking", False)
+            dataset_attr.system_prompt = prompt_list[i]
+            self.dataset_list.append(dataset_attr)

src/llmtuner/hparams/finetuning_args.py

@@ -5,31 +5,37 @@ from dataclasses import asdict, dataclass, field
 
 @dataclass
 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."}
     )
     num_hidden_layers: Optional[int] = field(
         default=32,
-        metadata={"help": "Number of decoder blocks in the model. \
+        metadata={"help": "Number of decoder blocks in the model for partial-parameter (freeze) fine-tuning. \
                   LLaMA choices: [\"32\", \"40\", \"60\", \"80\"], \
+                  LLaMA-2 choices: [\"32\", \"40\", \"80\"], \
                   BLOOM choices: [\"24\", \"30\", \"70\"], \
                   Falcon choices: [\"32\", \"60\"], \
-                  Baichuan choices: [\"32\"]"}
+                  Baichuan choices: [\"32\", \"40\"] \
+                  Qwen choices: [\"32\"], \
+                  XVERSE choices: [\"40\"], \
+                  ChatGLM2 choices: [\"28\"]"}
     )
     num_layer_trainable: Optional[int] = field(
         default=3,
-        metadata={"help": "Number of trainable layers for Freeze fine-tuning."}
+        metadata={"help": "Number of trainable layers for partial-parameter (freeze) fine-tuning."}
     )
     name_module_trainable: Optional[Literal["mlp", "self_attn", "self_attention"]] = field(
         default="mlp",
-        metadata={"help": "Name of trainable modules for Freeze fine-tuning. \
+        metadata={"help": "Name of trainable modules for partial-parameter (freeze) fine-tuning. \
                   LLaMA choices: [\"mlp\", \"self_attn\"], \
-                  BLOOM & Falcon choices: [\"mlp\", \"self_attention\"], \
-                  Baichuan choices: [\"mlp\", \"self_attn\"]"}
+                  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."}
     )
     lora_rank: Optional[int] = field(
         default=8,
@@ -44,11 +50,25 @@ 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\"], \
-                  Baichuan choices: [\"W_pack\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"]"}
+                  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."}
+    )
+    resume_lora_training: Optional[bool] = field(
+        default=True,
+        metadata={"help": "Whether to resume training from the last LoRA weights or create new weights after merging them."}
+    )
+    ppo_score_norm: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Use score normalization in PPO Training."}
+    )
+    dpo_beta: Optional[float] = field(
+        default=0.1,
+        metadata={"help": "The beta parameter for the DPO loss."}
     )
 
     def __post_init__(self):
@@ -62,17 +82,17 @@ 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):
-        """Saves the content of this instance in JSON format inside `json_path`."""
+        r"""Saves the content of this instance in JSON format inside `json_path`."""
         json_string = json.dumps(asdict(self), indent=2, sort_keys=True) + "\n"
         with open(json_path, "w", encoding="utf-8") as f:
             f.write(json_string)
 
     @classmethod
     def load_from_json(cls, json_path: str):
-        """Creates an instance from the content of `json_path`."""
+        r"""Creates an instance from the content of `json_path`."""
         with open(json_path, "r", encoding="utf-8") as f:
             text = f.read()
         return cls(**json.loads(text))

src/llmtuner/hparams/general_args.py

@@ -4,10 +4,10 @@ from dataclasses import dataclass, field
 
 @dataclass
 class GeneralArguments:
+    r"""
+    Arguments pertaining to which stage we are going to perform.
     """
-    Arguments pertaining to which techniques we are going to fine-tuning with.
-    """
-    stage: Optional[Literal["pt", "sft", "rm", "ppo"]] = field(
+    stage: Optional[Literal["pt", "sft", "rm", "ppo", "dpo"]] = field(
         default="sft",
         metadata={"help": "Which stage will be performed in training."}
     )

src/llmtuner/hparams/generating_args.py

@@ -4,7 +4,7 @@ from dataclasses import asdict, dataclass, field
 
 @dataclass
 class GeneratingArguments:
-    """
+    r"""
     Arguments pertaining to specify the decoding parameters.
     """
     do_sample: Optional[bool] = field(

src/llmtuner/hparams/model_args.py

@@ -5,7 +5,7 @@ from dataclasses import dataclass, field
 
 @dataclass
 class ModelArguments:
-    """
+    r"""
     Arguments pertaining to which model/config/tokenizer we are going to fine-tune.
     """
     model_name_or_path: str = field(
@@ -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,
@@ -55,18 +55,33 @@ class ModelArguments:
         default=None,
         metadata={"help": "Path to the directory containing the checkpoints of the reward model."}
     )
-    resume_lora_training: Optional[bool] = field(
-        default=True,
-        metadata={"help": "Whether to resume training from the last LoRA weights or create new weights after merging them."}
-    )
     plot_loss: Optional[bool] = field(
         default=False,
         metadata={"help": "Whether to plot the training loss after fine-tuning or not."}
     )
+    hf_auth_token: Optional[str] = field(
+        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(",")]
 
         if self.quantization_bit is not None:
             assert self.quantization_bit in [4, 8], "We only accept 4-bit or 8-bit quantization."
+
+        if self.use_auth_token == True and self.hf_auth_token is not None:
+            from huggingface_hub.hf_api import HfFolder # lazy load
+            HfFolder.save_token(self.hf_auth_token)

src/llmtuner/tuner/__init__.py

@@ -1,5 +1 @@
-from llmtuner.tuner.core import get_train_args, get_infer_args, load_model_and_tokenizer
-from llmtuner.tuner.pt import run_pt
-from llmtuner.tuner.sft import run_sft
-from llmtuner.tuner.rm import run_rm
-from llmtuner.tuner.ppo import run_ppo
+from llmtuner.tuner.tune import export_model, run_exp

src/llmtuner/tuner/core/adapter.py

@@ -1,7 +1,7 @@
 import os
 import torch
+from typing import TYPE_CHECKING
 
-from transformers.modeling_utils import PreTrainedModel
 from peft import (
     PeftModel,
     TaskType,
@@ -11,20 +11,23 @@ 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.hparams import ModelArguments, FinetuningArguments
+from llmtuner.tuner.core.utils import find_all_linear_modules
+
+if TYPE_CHECKING:
+    from transformers.modeling_utils import PreTrainedModel
+    from llmtuner.hparams import ModelArguments, FinetuningArguments
 
 
 logger = get_logger(__name__)
 
 
 def init_adapter(
-    model: PreTrainedModel,
-    model_args: ModelArguments,
-    finetuning_args: FinetuningArguments,
+    model: "PreTrainedModel",
+    model_args: "ModelArguments",
+    finetuning_args: "FinetuningArguments",
     is_trainable: bool,
     is_mergeable: bool
-) -> PreTrainedModel:
+) -> "PreTrainedModel":
     r"""
     Initializes the adapters.
 
@@ -36,7 +39,7 @@ def init_adapter(
     if finetuning_args.finetuning_type == "none" and is_trainable:
         raise ValueError("You cannot use finetuning_type=none while training.")
 
-    if finetuning_args.finetuning_type == "full":
+    if finetuning_args.finetuning_type == "full" and is_trainable:
         logger.info("Fine-tuning method: Full")
         model = model.float()
 
@@ -49,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
@@ -62,7 +62,7 @@ def init_adapter(
             assert os.path.exists(os.path.join(model_args.checkpoint_dir[0], CONFIG_NAME)), \
                 "The given checkpoint may be not a LoRA checkpoint, please specify `--finetuning_type full/freeze` instead."
 
-            if (is_trainable and model_args.resume_lora_training) or (not is_mergeable): # continually train on the lora weights
+            if (is_trainable and finetuning_args.resume_lora_training) or (not is_mergeable): # continually fine-tuning
                 checkpoints_to_merge, latest_checkpoint = model_args.checkpoint_dir[:-1], model_args.checkpoint_dir[-1]
             else:
                 checkpoints_to_merge = model_args.checkpoint_dir
@@ -78,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)
 

src/llmtuner/tuner/core/loader.py

@@ -1,42 +1,56 @@
 import os
+import math
 import torch
-from typing import Literal, Optional, Tuple
+from types import MethodType
+from typing import TYPE_CHECKING, Literal, Optional, Tuple
 
+import transformers
 from transformers import (
     AutoConfig,
     AutoModelForCausalLM,
     AutoTokenizer,
-    BitsAndBytesConfig
+    BitsAndBytesConfig,
+    PretrainedConfig,
+    PreTrainedModel,
+    PreTrainedTokenizerBase
 )
 from transformers.utils import check_min_version
 from transformers.utils.versions import require_version
-from transformers.modeling_utils import PretrainedConfig, PreTrainedModel
-from transformers.tokenization_utils import PreTrainedTokenizer
 from trl import AutoModelForCausalLMWithValueHead
 
-from llmtuner.extras.logging import get_logger
-from llmtuner.extras.misc import prepare_model_for_training, print_trainable_params
+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
 from llmtuner.extras.save_and_load import load_valuehead_params
-from llmtuner.hparams import ModelArguments, FinetuningArguments
+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
+    from llmtuner.hparams import ModelArguments
 
 
 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.19.0", "To fix: pip install accelerate>=0.19.0")
-require_version("peft>=0.3.0", "To fix: pip install peft>=0.3.0")
-require_version("trl>=0.4.4", "To fix: pip install trl>=0.4.4")
+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.7.1", "To fix: pip install trl>=0.7.1")
 
 
 def load_model_and_tokenizer(
-    model_args: ModelArguments,
-    finetuning_args: FinetuningArguments,
+    model_args: "ModelArguments",
+    finetuning_args: "FinetuningArguments",
     is_trainable: Optional[bool] = False,
     stage: Optional[Literal["pt", "sft", "rm", "ppo"]] = "sft"
-) -> Tuple[PreTrainedModel, PreTrainedTokenizer]:
+) -> Tuple[PreTrainedModel, "PreTrainedTokenizer"]:
     r"""
     Loads pretrained model and tokenizer.
 
@@ -46,9 +60,6 @@ def load_model_and_tokenizer(
         logger.warning("Checkpoint is not found at evaluation, load the original model.")
         finetuning_args = FinetuningArguments(finetuning_type="none")
 
-    assert stage in ["pt", "sft"] or finetuning_args.finetuning_type == "lora", \
-        "RM and PPO training can only be performed with the LoRA method."
-
     config_kwargs = {
         "trust_remote_code": True,
         "cache_dir": model_args.cache_dir,
@@ -59,30 +70,91 @@ 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 tokenizer.pad_token_id is None or tokenizer.pad_token_id == 64000: # 64000 for baichuan model (older version)
-        tokenizer.pad_token_id = 0 # set as the <unk> token
 
-    config = AutoConfig.from_pretrained(model_args.model_name_or_path, **config_kwargs)
-    is_mergeable = True
+    # 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)
+
+    # 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
-            config_kwargs["quantization_config"] = BitsAndBytesConfig(
-                load_in_8bit=True,
-                llm_int8_threshold=6.0
-            )
+            config_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
 
         elif model_args.quantization_bit == 4:
             require_version("bitsandbytes>=0.39.0", "To fix: pip install bitsandbytes>=0.39.0")
-            require_version("transformers>=4.30.1", "To fix: pip install transformers>=4.30.1")
-            require_version("accelerate>=0.20.3", "To fix: pip install accelerate>=0.20.3")
-            require_version("peft>=0.4.0.dev0", "To fix: pip install git+https://github.com/huggingface/peft.git")
             config_kwargs["load_in_4bit"] = True
             config_kwargs["quantization_config"] = BitsAndBytesConfig(
                 load_in_4bit=True,
@@ -92,41 +164,44 @@ 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))
 
-    if not is_trainable: # `device_map=auto` should be used for inference only
-        config_kwargs["device_map"] = "auto"
-
-    if model_args.checkpoint_dir is not None and finetuning_args.finetuning_type == "full":
-        model_to_load = model_args.checkpoint_dir[0]
-    else:
-        model_to_load = model_args.model_name_or_path
-
-    # 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,
-        low_cpu_mem_usage=True,
+        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 hasattr(config, "auto_map") and "AutoConfig" in config.auto_map and isinstance(config, PretrainedConfig):
+    if isinstance(config, PretrainedConfig) and "AutoConfig" in getattr(config, "auto_map", {}):
         config.__class__.register_for_auto_class()
-    if hasattr(config, "auto_map") and "AutoTokenizer" in config.auto_map and isinstance(tokenizer, PreTrainedTokenizer):
-        tokenizer.__class__.register_for_auto_class()
-    if hasattr(config, "auto_map") and "AutoModelForCausalLM" in config.auto_map and isinstance(model, PreTrainedModel):
+    if isinstance(model, PreTrainedModel) and "AutoModelForCausalLM" in getattr(config, "auto_map", {}):
         model.__class__.register_for_auto_class()
+    if isinstance(tokenizer, PreTrainedTokenizerBase) and "AutoTokenizer" in tokenizer.init_kwargs.get("auto_map", {}):
+        tokenizer.__class__.register_for_auto_class()
 
     # Initialize adapters
     model = prepare_model_for_training(model, 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.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]):
@@ -136,16 +211,19 @@ def load_model_and_tokenizer(
                 })
 
         if stage == "ppo": # load reward model
-            assert is_trainable, "PPO stage cannot be performed at evaluation."
-            assert model_args.reward_model is not None, "Reward model is necessary for PPO training."
             logger.info("Load reward model from {}".format(model_args.reward_model))
             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
 
-    print_trainable_params(model)
+    trainable_params, all_param = count_parameters(model)
+    logger.info("trainable params: {:d} || all params: {:d} || trainable%: {:.4f}".format(
+        trainable_params, all_param, 100 * trainable_params / all_param
+    ))
 
     return model, tokenizer

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 (
@@ -19,20 +21,66 @@ from llmtuner.hparams import (
 logger = get_logger(__name__)
 
 
+def _parse_args(parser: HfArgumentParser, args: Optional[Dict[str, Any]] = None) -> Tuple[Any]:
+    if args is not None:
+        return parser.parse_dict(args)
+    elif len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"):
+        return parser.parse_yaml_file(os.path.abspath(sys.argv[1]))
+    elif len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        return parser.parse_json_file(os.path.abspath(sys.argv[1]))
+    else:
+        return parser.parse_args_into_dataclasses()
+
+
+def parse_train_args(
+    args: Optional[Dict[str, Any]] = None
+) -> Tuple[
+    ModelArguments,
+    DataArguments,
+    Seq2SeqTrainingArguments,
+    FinetuningArguments,
+    GeneratingArguments,
+    GeneralArguments
+]:
+    parser = HfArgumentParser((
+        ModelArguments,
+        DataArguments,
+        Seq2SeqTrainingArguments,
+        FinetuningArguments,
+        GeneratingArguments,
+        GeneralArguments
+    ))
+    return _parse_args(parser, args)
+
+
+def parse_infer_args(
+    args: Optional[Dict[str, Any]] = None
+) -> Tuple[
+    ModelArguments,
+    DataArguments,
+    FinetuningArguments,
+    GeneratingArguments
+]:
+    parser = HfArgumentParser((
+        ModelArguments,
+        DataArguments,
+        FinetuningArguments,
+        GeneratingArguments
+    ))
+    return _parse_args(parser, args)
+
+
 def get_train_args(
     args: Optional[Dict[str, Any]] = None
-) -> Tuple[ModelArguments, DataArguments, Seq2SeqTrainingArguments, FinetuningArguments, GeneralArguments]:
-
-    parser = HfArgumentParser((ModelArguments, DataArguments, Seq2SeqTrainingArguments, FinetuningArguments, GeneralArguments))
-
-    if args is not None:
-        model_args, data_args, training_args, finetuning_args, general_args = parser.parse_dict(args)
-    elif len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"):
-        model_args, data_args, training_args, finetuning_args, general_args = parser.parse_yaml_file(os.path.abspath(sys.argv[1]))
-    elif len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
-        model_args, data_args, training_args, finetuning_args, general_args = parser.parse_json_file(os.path.abspath(sys.argv[1]))
-    else:
-        model_args, data_args, training_args, finetuning_args, general_args = parser.parse_args_into_dataclasses()
+) -> Tuple[
+    ModelArguments,
+    DataArguments,
+    Seq2SeqTrainingArguments,
+    FinetuningArguments,
+    GeneratingArguments,
+    GeneralArguments
+]:
+    model_args, data_args, training_args, finetuning_args, generating_args, general_args = parse_train_args(args)
 
     # Setup logging
     if training_args.should_log:
@@ -48,87 +96,145 @@ def get_train_args(
     # Check arguments (do not check finetuning_args since it may be loaded from checkpoints)
     data_args.init_for_training()
 
-    assert general_args.stage == "sft" or (not training_args.predict_with_generate), \
-        "`predict_with_generate` cannot be set as True at PT, RM and PPO stages."
+    if general_args.stage != "pt" and data_args.template is None:
+        raise ValueError("Please specify which `template` to use.")
 
-    assert not (training_args.do_train and training_args.predict_with_generate), \
-        "`predict_with_generate` cannot be set as True while training."
+    if general_args.stage != "sft" and training_args.predict_with_generate:
+        raise ValueError("`predict_with_generate` cannot be set as True except SFT.")
 
-    assert (not training_args.do_predict) or training_args.predict_with_generate, \
-        "Please enable `predict_with_generate` to save model predictions."
+    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.")
 
-    assert model_args.quantization_bit is None or finetuning_args.finetuning_type == "lora", \
-        "Quantization is only compatible with the LoRA method."
+    if general_args.stage in ["rm", "ppo"] and finetuning_args.finetuning_type != "lora":
+        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 stages can only be performed at training.")
+
+    if general_args.stage in ["rm", "dpo"]:
+        for dataset_attr in data_args.dataset_list:
+            if not dataset_attr.ranking:
+                raise ValueError("Please use ranked datasets for reward modeling or DPO training.")
+
+    if general_args.stage == "ppo" and model_args.reward_model is None:
+        raise ValueError("Reward model is necessary for PPO training.")
+
+    if general_args.stage == "ppo" and training_args.deepspeed is not None:
+        raise ValueError("PPO training is incompatible with DeepSpeed, use Accelerate instead.")
+
+    if general_args.stage == "ppo" and data_args.streaming:
+        raise ValueError("Streaming mode does not suppport PPO training currently.")
+
+    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.")
 
     if model_args.checkpoint_dir is not None:
         if finetuning_args.finetuning_type != "lora":
-            assert len(model_args.checkpoint_dir) == 1, "Only LoRA tuning accepts multiple checkpoints."
-        else:
-            assert model_args.quantization_bit is None or len(model_args.checkpoint_dir) == 1, \
-                "Quantized model only accepts a single checkpoint."
+            if len(model_args.checkpoint_dir) != 1:
+                raise ValueError("Only LoRA tuning accepts multiple checkpoints.")
+        elif model_args.quantization_bit is not None and len(model_args.checkpoint_dir) != 1:
+                raise ValueError("Quantized model only accepts a single checkpoint.")
 
     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.")
 
-    if data_args.prompt_template == "default":
-        logger.warning("Please specify `prompt_template` if you are using other pre-trained models.")
+    # 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
 
-    if training_args.local_rank != -1 and training_args.ddp_find_unused_parameters is None:
-        logger.warning("`ddp_find_unused_parameters` needs to be set as False in DDP training.")
-        training_args.ddp_find_unused_parameters = False
+    # 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_dict = training_args.to_dict()
+        training_args_dict.update(dict(ddp_find_unused_parameters=False))
+        training_args = Seq2SeqTrainingArguments(**training_args_dict)
 
-    training_args.optim = "adamw_torch" if training_args.optim == "adamw_hf" else training_args.optim # suppress warning
+    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`.")
 
-    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
+        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."
+            )
+
+    # 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(
-        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}\n"
-        + f"  distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
-    )
+    logger.info("Process rank: {}, device: {}, n_gpu: {}\n  distributed training: {}, compute dtype: {}".format(
+        training_args.local_rank, training_args.device, training_args.n_gpu,
+        bool(training_args.local_rank != -1), str(model_args.compute_dtype)
+    ))
     logger.info(f"Training/evaluation parameters {training_args}")
 
     # Set seed before initializing model.
     transformers.set_seed(training_args.seed)
 
-    return model_args, data_args, training_args, finetuning_args, general_args
+    return model_args, data_args, training_args, finetuning_args, generating_args, general_args
 
 
 def get_infer_args(
     args: Optional[Dict[str, Any]] = None
-) -> Tuple[ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments]:
+) -> Tuple[
+    ModelArguments,
+    DataArguments,
+    FinetuningArguments,
+    GeneratingArguments
+]:
+    model_args, data_args, finetuning_args, generating_args = parse_infer_args(args)
 
-    parser = HfArgumentParser((ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments))
+    if data_args.template is None:
+        raise ValueError("Please specify which `template` to use.")
 
-    if args is not None:
-        model_args, data_args, finetuning_args, generating_args = parser.parse_dict(args)
-    elif len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"):
-        model_args, data_args, finetuning_args, generating_args = parser.parse_yaml_file(os.path.abspath(sys.argv[1]))
-    elif len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
-        model_args, data_args, finetuning_args, generating_args = parser.parse_json_file(os.path.abspath(sys.argv[1]))
-    else:
-        model_args, data_args, finetuning_args, generating_args = parser.parse_args_into_dataclasses()
-
-    assert model_args.quantization_bit is None or finetuning_args.finetuning_type == "lora", \
-        "Quantization is only compatible with the LoRA method."
+    if model_args.quantization_bit is not None and finetuning_args.finetuning_type != "lora":
+        raise ValueError("Quantization is only compatible with the LoRA method.")
 
     if model_args.checkpoint_dir is not None:
         if finetuning_args.finetuning_type != "lora":
-            assert len(model_args.checkpoint_dir) == 1, "Only LoRA tuning accepts multiple checkpoints."
-        else:
-            assert model_args.quantization_bit is None or len(model_args.checkpoint_dir) == 1, \
-                "Quantized model only accepts a single checkpoint."
-
-    if data_args.prompt_template == "default":
-        logger.warning("Please specify `prompt_template` if you are using other pre-trained models.")
+            if len(model_args.checkpoint_dir) != 1:
+                raise ValueError("Only LoRA tuning accepts multiple checkpoints.")
+        elif model_args.quantization_bit is not None and len(model_args.checkpoint_dir) != 1:
+                raise ValueError("Quantized model only accepts a single checkpoint.")
 
     return model_args, data_args, finetuning_args, generating_args

src/llmtuner/tuner/core/trainer.py

@@ -1,88 +0,0 @@
-import os
-import torch
-from typing import Dict, Optional
-
-from transformers import Seq2SeqTrainer
-from transformers.trainer import TRAINING_ARGS_NAME
-from transformers.modeling_utils import unwrap_model
-
-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, load_valuehead_params
-from llmtuner.hparams import FinetuningArguments
-
-
-logger = get_logger(__name__)
-
-
-class PeftTrainer(Seq2SeqTrainer):
-    r"""
-    Inherits Seq2SeqTrainer to support parameter-efficient checkpoints.
-    """
-
-    def __init__(self, finetuning_args: FinetuningArguments, **kwargs):
-        super().__init__(**kwargs)
-        self.finetuning_args = finetuning_args
-        self._remove_log()
-
-    def _remove_log(self):
-        if self.is_world_process_zero() and os.path.exists(os.path.join(self.args.output_dir, "trainer_log.jsonl")):
-            logger.warning("Previous log file in this folder will be deleted.")
-            os.remove(os.path.join(self.args.output_dir, "trainer_log.jsonl"))
-
-    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 hasattr(model, "pretrained_model"): # for models with valuehead (currently using LoRA only)
-            backbone_model = getattr(model, "pretrained_model")
-            torch.save(get_state_dict(getattr(model, "v_head")), os.path.join(output_dir, VALUE_HEAD_FILE_NAME))
-        else:
-            backbone_model = model
-
-        if self.finetuning_args.finetuning_type == "lora":
-            backbone_model.save_pretrained(output_dir, state_dict=get_state_dict(backbone_model))
-        else: # freeze/full tuning
-            backbone_model.config.use_cache = True
-            backbone_model.save_pretrained(
-                output_dir,
-                state_dict=get_state_dict(backbone_model),
-                safe_serialization=self.args.save_safetensors
-            )
-            backbone_model.config.use_cache = False
-            if self.tokenizer is not None:
-                self.tokenizer.save_pretrained(output_dir)
-
-        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)
-        backbone_model = getattr(model, "pretrained_model") if hasattr(model, "pretrained_model") else model
-
-        if self.finetuning_args.finetuning_type == "lora":
-            backbone_model.load_adapter(self.state.best_model_checkpoint, getattr(backbone_model, "active_adapter"))
-            if hasattr(model, "v_head") and load_valuehead_params(model, self.state.best_model_checkpoint):
-                model.v_head.load_state_dict({
-                    "summary.weight": getattr(model, "reward_head_weight"),
-                    "summary.bias": getattr(model, "reward_head_bias")
-                })
-        else: # freeze/full-tuning
-            load_trainable_params(backbone_model, self.state.best_model_checkpoint)

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

src/llmtuner/tuner/dpo/__init__.py

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

src/llmtuner/tuner/dpo/collator.py

@@ -0,0 +1,51 @@
+import torch
+from dataclasses import dataclass
+from typing import Any, Dict, List, Sequence, Tuple
+from transformers import DataCollatorForSeq2Seq
+
+
+@dataclass
+class DPODataCollatorWithPadding(DataCollatorForSeq2Seq):
+    r"""
+    Data collator for pairwise data.
+    """
+
+    def _pad_labels(self, batch: torch.Tensor, positions: List[Tuple[int, int]]) -> torch.Tensor:
+        padded_labels = []
+        for feature, (prompt_len, answer_len) in zip(batch, positions):
+            if self.tokenizer.padding_side == "left":
+                start, end = feature.size(0) - answer_len, feature.size(0)
+            else:
+                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)
+        return torch.stack(padded_labels, dim=0).contiguous() # in contiguous memory
+
+    def __call__(self, features: Sequence[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
+        r"""
+        Pads batched data to the longest sequence in the batch.
+
+        We generate 2 * n examples where the first n examples represent chosen examples and
+        the last n examples represent rejected examples.
+        """
+        concatenated_features = []
+        label_positions = []
+        for key in ("chosen_ids", "rejected_ids"):
+            for feature in features:
+                prompt_len, answer_len = len(feature["prompt_ids"]), len(feature[key])
+                concatenated_features.append({
+                    "input_ids": feature["prompt_ids"] + feature[key],
+                    "attention_mask": [1] * (prompt_len + answer_len)
+                })
+                label_positions.append((prompt_len, answer_len))
+
+        batch = self.tokenizer.pad(
+            concatenated_features,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors=self.return_tensors,
+        )
+        batch["labels"] = self._pad_labels(batch["input_ids"], label_positions)
+        return batch

src/llmtuner/tuner/dpo/trainer.py

@@ -0,0 +1,69 @@
+import torch
+from collections import defaultdict
+from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
+from transformers import BatchEncoding, Trainer
+from trl import DPOTrainer
+from trl.trainer.utils import disable_dropout_in_model
+
+from llmtuner.extras.constants import IGNORE_INDEX
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel
+
+
+class CustomDPOTrainer(DPOTrainer):
+
+    def __init__(
+        self,
+        beta: float,
+        model: Union["PreTrainedModel", torch.nn.Module],
+        ref_model: Optional[Union["PreTrainedModel", torch.nn.Module]] = None,
+        disable_dropout: Optional[bool] = True,
+        **kwargs
+    ):
+        if disable_dropout:
+            disable_dropout_in_model(model)
+            if ref_model is not None:
+                disable_dropout_in_model(ref_model)
+
+        self.is_encoder_decoder = model.config.is_encoder_decoder
+        self.ref_model = ref_model
+        self.use_dpo_data_collator = True # hack to avoid warning
+        self.label_pad_token_id = IGNORE_INDEX
+        self.padding_value = 0
+        self.beta = beta
+        self._stored_metrics = defaultdict(lambda: defaultdict(list))
+
+        Trainer.__init__(self, model=model, **kwargs)
+        if not hasattr(self, "accelerator"):
+            raise AttributeError("Please update `transformers`.")
+
+        if ref_model is not None:
+            if self.is_deepspeed_enabled:
+                self.ref_model, = self.accelerator._prepare_deepspeed(self.ref_model)
+                self.ref_model.eval()
+            else:
+                self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
+
+    def concatenated_forward(
+        self,
+        model: Optional[torch.nn.Module] = None,
+        batch: Optional[Dict[str, torch.Tensor]] = None
+    ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
+        batch_copied = BatchEncoding({k: v.detach().clone() for k, v in batch.items()}) # avoid error
+
+        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,
+            batch["labels"],
+            average_log_prob=False
+        )
+        batch_size = batch["input_ids"].size(0) // 2
+        chosen_logps, rejected_logps = all_logps.split(batch_size, dim=0)
+        chosen_logits, rejected_logits = all_logits.split(batch_size, dim=0)
+        return chosen_logps, rejected_logps, chosen_logits, rejected_logits

src/llmtuner/tuner/dpo/workflow.py

@@ -0,0 +1,59 @@
+# Inspired by: https://github.com/huggingface/trl/blob/main/examples/research_projects/stack_llama_2/scripts/dpo_llama2.py
+
+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 CustomDPOTrainer
+
+if TYPE_CHECKING:
+    from transformers import TrainerCallback
+    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
+
+
+def run_dpo(
+    model_args: "ModelArguments",
+    data_args: "DataArguments",
+    training_args: "Seq2SeqTrainingArguments",
+    finetuning_args: "FinetuningArguments",
+    callbacks: Optional[List["TrainerCallback"]] = None
+):
+    dataset = get_dataset(model_args, data_args)
+    model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="sft")
+    dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="rm")
+    data_collator = DPODataCollatorWithPadding(
+        tokenizer=tokenizer,
+        label_pad_token_id=IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
+    )
+
+    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 = 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,
+        callbacks=callbacks,
+        **split_dataset(dataset, data_args, training_args)
+    )
+
+    # Training
+    if training_args.do_train:
+        train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
+        trainer.log_metrics("train", train_result.metrics)
+        trainer.save_metrics("train", train_result.metrics)
+        trainer.save_state()
+        trainer.save_model()
+        if trainer.is_world_process_zero() and model_args.plot_loss:
+            plot_loss(training_args.output_dir, keys=["loss", "eval_loss"])

src/llmtuner/tuner/ppo/trainer.py

@@ -2,56 +2,62 @@ import os
 import math
 import torch
 from tqdm import tqdm
-from typing import Callable, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple
 
-from transformers import Seq2SeqTrainingArguments, TrainerState, TrainerControl
-from transformers.modeling_utils import PreTrainedModel
+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.callbacks import LogCallback
 from llmtuner.extras.logging import get_logger
-from llmtuner.extras.misc import AverageMeter, get_logits_processor
-from llmtuner.hparams import FinetuningArguments
-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, TrainerCallback
+    from trl import AutoModelForCausalLMWithValueHead
+    from llmtuner.hparams import GeneratingArguments
+
 
 logger = get_logger(__name__)
 
 
-class PPOPeftTrainer(PPOTrainer, PeftTrainer):
+class CustomPPOTrainer(PPOTrainer, Trainer):
     r"""
     Inherits PPOTrainer.
     """
 
     def __init__(
         self,
-        training_args: Seq2SeqTrainingArguments,
-        finetuning_args: FinetuningArguments,
-        callbacks: List[LogCallback],
+        training_args: "Seq2SeqTrainingArguments",
+        generating_args: "GeneratingArguments",
+        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.log_callback = callbacks[0]
+        self.generating_args = generating_args
+        self.log_callback, self.save_callback = callbacks[0], callbacks[1]
+        self.compute_dtype = compute_dtype
         self.state = TrainerState()
         self.control = TrainerControl()
-        self.data_collator = self.accelerator.prepare(kwargs["data_collator"]) # override the data collator of PPOTrainer
-        self._remove_log()
 
     def ppo_train(self, max_target_length: int) -> None:
         r"""
         Implements training loop for the PPO stage, like _inner_training_loop() in Huggingface's Trainer.
         """
-        total_train_batch_size = self.config.batch_size * self.config.gradient_accumulation_steps * self.args.world_size
+        total_train_batch_size = (
+            self.args.per_device_train_batch_size * self.args.gradient_accumulation_steps * self.args.world_size
+        )
         len_dataloader = len(self.dataloader)
-        num_steps_per_epoch = max(len_dataloader // self.config.gradient_accumulation_steps, 1)
         num_examples = len(self.dataset)
         num_train_epochs = self.args.num_train_epochs
-        max_steps = math.ceil(num_train_epochs * num_steps_per_epoch)
+        max_steps = math.ceil(num_train_epochs * len_dataloader)
 
         self.state.max_steps = max_steps
         self.state.num_train_epochs = num_train_epochs
@@ -62,23 +68,21 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
             logger.info("***** Running training *****")
             logger.info(f"  Num examples = {num_examples}")
             logger.info(f"  Num Epochs = {num_train_epochs}")
-            logger.info(f"  Instantaneous batch size per device = {self.config.batch_size}")
+            logger.info(f"  Instantaneous batch size per device = {self.args.per_device_train_batch_size}")
             logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size}")
-            logger.info(f"  Gradient Accumulation steps = {self.config.gradient_accumulation_steps}")
+            logger.info(f"  Gradient Accumulation steps = {self.args.gradient_accumulation_steps}")
             logger.info(f"  Total optimization steps = {max_steps}")
-            logger.info(f"  Number of trainable parameters = {sum(p.numel() for p in self.model.parameters() if p.requires_grad)}")
+            logger.info(f"  Number of trainable parameters = {count_parameters(self.model)[0]}")
 
         # Keyword arguments for `model.generate`
-        gen_kwargs = {
-            "top_k": 0.0,
-            "top_p": 1.0,
-            "do_sample": True,
-            "pad_token_id": self.tokenizer.pad_token_id,
-            "eos_token_id": self.tokenizer.eos_token_id,
-            "logits_processor": get_logits_processor()
-        }
-        output_length_sampler = LengthSampler(max_target_length // 2, max_target_length)
-        unwrapped_model: PreTrainedModel = self.accelerator.unwrap_model(self.model)
+        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)
 
         dataiter = iter(self.dataloader)
         steps_trained = 0
@@ -86,61 +90,42 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
         reward_meter = AverageMeter()
         self.log_callback.on_train_begin(self.args, self.state, self.control)
 
-        for step in tqdm(range(max_steps), disable=not self.is_world_process_zero(), leave=False):
+        for step in tqdm(range(max_steps), disable=not self.is_local_process_zero()):
+            batch = next(dataiter)
+            steps_trained += 1
 
-            for _ in range(self.config.gradient_accumulation_steps):
+            # Cast to inference mode
+            unwrapped_model.gradient_checkpointing_disable()
+            unwrapped_model.config.use_cache = True
+            self.model.eval()
 
-                batch = next(dataiter)
-                steps_trained += 1
+            # Get inputs
+            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)
 
-                unwrapped_model.gradient_checkpointing_disable()
-                unwrapped_model.config.use_cache = True
+            # Cast to training mode
+            unwrapped_model.gradient_checkpointing_enable()
+            unwrapped_model.config.use_cache = False
+            self.model.train()
 
-                # Get response from model
-                query_tensors: torch.Tensor = batch["input_ids"]
-                response_tensors = self.generate(batch, length_sampler=output_length_sampler, return_prompt=False, **gen_kwargs)
+            # 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))
 
-                queries: List[torch.Tensor] = []
-                responses: List[torch.Tensor] = []
-                for i in range(len(query_tensors)):
-                    query_length = (query_tensors[i] != self.tokenizer.pad_token_id).nonzero()[0]
-                    response_length = (response_tensors[i] != self.tokenizer.pad_token_id).nonzero()[-1] + 1
-                    queries.append(query_tensors[i, query_length:]) # remove padding from left
-                    if response_length < 2: # make response have at least 2 tokens
-                        responses.append(response_tensors.new_empty(2).fill_(self.tokenizer.eos_token_id))
-                    else:
-                        responses.append(response_tensors[i, :response_length]) # remove padding from right
+            self.state.global_step += 1
+            self.log_callback.on_step_end(self.args, self.state, self.control)
 
-                # Compute rewards
-                replace_model(unwrapped_model, target="reward")
-                _, _, values = self.model(**self.prepare_model_inputs(queries, responses))
-                rewards = [reward for reward in values[:, -1].to(torch.float32)] # use float32 type
-                replace_model(unwrapped_model, target="default") # make sure the model is default at the end
-
-                # Run PPO step
-                unwrapped_model.gradient_checkpointing_enable()
-                unwrapped_model.config.use_cache = False
-
-                stats = self.step(queries, responses, rewards)
-
-                loss_meter.update(stats["ppo/loss/total"], n=len(rewards))
-                reward_meter.update(torch.stack(rewards).mean().item(), n=len(rewards))
-
-                if self.control.should_epoch_stop or self.control.should_training_stop:
-                    break
-
-                if steps_trained == len_dataloader:
-                    dataiter = iter(self.dataloader)
-                    steps_trained = 0
-
-            if self.is_world_process_zero() and (step+1) % self.args.logging_steps == 0:
-                logs = {
-                    "loss": round(loss_meter.avg, 4),
-                    "reward": round(reward_meter.avg, 4),
-                    "learning_rate": stats["ppo/learning_rate"],
-                    "epoch": round(step / num_steps_per_epoch, 2)
-                }
-                print(logs)
+            if self.is_local_process_zero() and (step+1) % self.args.logging_steps == 0:
+                logs = dict(
+                    loss=round(loss_meter.avg, 4),
+                    reward=round(reward_meter.avg, 4),
+                    learning_rate=stats["ppo/learning_rate"],
+                    epoch=round(step / len_dataloader, 2)
+                )
+                tqdm.write(str(logs))
                 logs["step"] = step
                 self.state.log_history.append(logs)
                 self.log_callback.on_log(self.args, self.state, self.control)
@@ -148,43 +133,164 @@ 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_training_stop:
+            if self.control.should_epoch_stop or self.control.should_training_stop:
                 break
 
+            if steps_trained == len_dataloader:
+                dataiter = iter(self.dataloader)
+                steps_trained = 0
+
+        self.log_callback.on_train_end(
+            self.args, self.state, self.control, model=self.accelerator.unwrap_model(self.model)
+        )
+
     @torch.no_grad()
-    def generate(
+    def get_inputs(
         self,
-        inputs: Dict[str, torch.Tensor],
-        length_sampler: Optional[Callable] = None,
-        return_prompt: Optional[bool] = True,
-        **generation_kwargs
-    ) -> torch.Tensor:
+        batch: Dict[str, torch.Tensor],
+        length_sampler: Callable,
+        generating_args: Dict[str, Any]
+    ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
         r"""
         Generates model's responses given queries.
+        """
+        generating_args["max_new_tokens"] = length_sampler()
+        gen_kwargs = dict(
+            generation_config=GenerationConfig(**generating_args),
+            logits_processor=get_logits_processor(),
+            **batch
+        )
+
+        input_ids = batch["input_ids"]
+        self.model, layer_norm_params = cast_layernorm_dtype(self.model, self.compute_dtype)
+        unwrapped_model: "AutoModelForCausalLMWithValueHead" = self.accelerator.unwrap_model(self.model)
+        response: torch.Tensor = unwrapped_model.generate(**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 = [], []
+        for i in range(len(query)):
+            query_length = (query[i] != self.tokenizer.pad_token_id).nonzero()[0]
+            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
+
+        return queries, responses
+
+    @torch.no_grad()
+    def get_rewards(
+        self,
+        queries: List[torch.Tensor],
+        responses: List[torch.Tensor],
+        unwrapped_model: "AutoModelForCausalLMWithValueHead"
+    ) -> List[torch.Tensor]:
+        r"""
+        Computes scores using given reward model.
+        """
+        replace_model(unwrapped_model, target="reward")
+        batch = self.prepare_model_inputs(queries, responses)
+
+        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.
         """
-        self.model, layer_norm_params = cast_layernorm_dtype(self.model)
+        bs = len(queries)
+        fbs = self.config.mini_batch_size
+        all_logprobs = []
+        all_logits = []
+        all_masks = []
+        all_values = []
 
-        if length_sampler is not None:
-            generation_kwargs["max_new_tokens"] = length_sampler()
+        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"]
 
-        unwrapped_model = self.accelerator.unwrap_model(self.model)
+            with torch.cuda.amp.autocast(dtype=self.compute_dtype): # support bf16
+                logits, _, values = model(**input_kwargs)
 
-        response = unwrapped_model.generate(**inputs, **generation_kwargs)
+            if values.size(0) != input_ids.size(0): # adapt to chatglm2
+                values = torch.transpose(values, 0, 1)
 
-        # 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
+            logprobs = logprobs_from_logits(logits[:, :-1, :], input_ids[:, 1:])
+            masks = torch.zeros_like(attention_mask)
+            masks[:, :-1] = attention_mask[:, 1:]
 
-        self.model, _ = cast_layernorm_dtype(self.model, layer_norm_params)
+            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 not return_prompt and not self.is_encoder_decoder:
-            return response[:, inputs["input_ids"].size(1):]
-        return response
+                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"""

src/llmtuner/tuner/ppo/utils.py

@@ -1,11 +1,13 @@
 import torch
-from typing import Dict, List, Literal, Optional, Tuple
-from trl import AutoModelForCausalLMWithValueHead
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple
 
 from llmtuner.extras.constants import LAYERNORM_NAMES
 
+if TYPE_CHECKING:
+    from trl import AutoModelForCausalLMWithValueHead
 
-def replace_model(model: AutoModelForCausalLMWithValueHead, target: Literal["default", "reward"]) -> None:
+
+def replace_model(model: "AutoModelForCausalLMWithValueHead", target: Literal["default", "reward"]) -> None:
     if target == "reward": # save default head temporarily
         valuehead_state_dict = model.v_head.state_dict()
         setattr(model, "default_head_weight", valuehead_state_dict["summary.weight"])
@@ -19,19 +21,20 @@ def replace_model(model: AutoModelForCausalLMWithValueHead, target: Literal["def
 
 
 def cast_layernorm_dtype(
-    model: AutoModelForCausalLMWithValueHead,
-    layer_norm_names: List[str] = LAYERNORM_NAMES,
-    layer_norm_params: Optional[Dict[str, torch.Tensor]] = None
-) -> Tuple[AutoModelForCausalLMWithValueHead, Dict[str, torch.Tensor]]:
+    model: "AutoModelForCausalLMWithValueHead",
+    compute_dtype: torch.dtype,
+    layer_norm_params: Optional[Dict[str, torch.Tensor]] = None,
+    layer_norm_names: Optional[List[str]] = LAYERNORM_NAMES
+) -> Tuple["AutoModelForCausalLMWithValueHead", Dict[str, torch.Tensor]]:
 
     layer_norm_state_dict = {}
 
     for name, param in model.named_parameters():
         if param.ndim == 1 and any(layer_norm_name in name for layer_norm_name in layer_norm_names):
-            if layer_norm_params is 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

src/llmtuner/tuner/ppo/workflow.py

@@ -1,56 +1,72 @@
-# Inspired by:
-# https://github.com/lvwerra/trl/blob/main/examples/sentiment/scripts/gpt-neox-20b_peft/gpt-neo-20b_sentiment_peft.py
+# Inspired by: https://github.com/lvwerra/trl/blob/main/examples/research_projects/stack_llama/scripts/rl_training.py
 
 import math
 from trl import PPOConfig
 from torch.optim import AdamW
-from transformers import DataCollatorForSeq2Seq, Seq2SeqTrainingArguments
+from typing import TYPE_CHECKING, Optional, List
+from transformers import DataCollatorWithPadding
 from transformers.optimization import get_scheduler
 
 from llmtuner.dsets import get_dataset, preprocess_dataset
-from llmtuner.extras.callbacks import LogCallback
+from llmtuner.extras.callbacks import SavePeftModelCallback
 from llmtuner.extras.ploting import plot_loss
-from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
 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
+    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments
 
 
 def run_ppo(
-    model_args: ModelArguments,
-    data_args: DataArguments,
-    training_args: Seq2SeqTrainingArguments,
-    finetuning_args: FinetuningArguments
+    model_args: "ModelArguments",
+    data_args: "DataArguments",
+    training_args: "Seq2SeqTrainingArguments",
+    finetuning_args: "FinetuningArguments",
+    generating_args: "GeneratingArguments",
+    callbacks: Optional[List["TrainerCallback"]] = None
 ):
     dataset = get_dataset(model_args, data_args)
     model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="ppo")
     dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="ppo")
-    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,
         learning_rate=training_args.learning_rate,
         mini_batch_size=training_args.per_device_train_batch_size,
-        batch_size=training_args.per_device_train_batch_size,
+        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
     )
 
-    optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=ppo_config.learning_rate)
-    total_train_batch_size = \
+    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
+    )
+    num_training_steps = training_args.num_train_epochs * math.ceil(len(dataset) / total_train_batch_size)
     lr_scheduler = get_scheduler(
         training_args.lr_scheduler_type,
         optimizer=optimizer,
-        num_warmup_steps=training_args.warmup_steps,
-        num_training_steps=(training_args.num_train_epochs * math.ceil(len(dataset) / total_train_batch_size))
+        num_warmup_steps=training_args.get_warmup_steps(num_training_steps),
+        num_training_steps=num_training_steps
     )
 
     # Initialize our Trainer
-    ppo_trainer = PPOPeftTrainer(
+    ppo_trainer = CustomPPOTrainer(
         training_args=training_args,
-        finetuning_args=finetuning_args,
-        callbacks=[LogCallback()],
+        generating_args=generating_args,
+        callbacks=callbacks + [SavePeftModelCallback()],
+        compute_dtype=model_args.compute_dtype,
         config=ppo_config,
         model=model,
         ref_model=None,
@@ -61,8 +77,10 @@ def run_ppo(
         lr_scheduler=lr_scheduler
     )
 
-    ppo_trainer.ppo_train(max_target_length=data_args.max_target_length)
-    ppo_trainer.save_model()
-    ppo_trainer.save_state() # must be after save_model
-    if ppo_trainer.is_world_process_zero() and model_args.plot_loss:
-        plot_loss(training_args.output_dir, keys=["loss", "reward"])
+    # Training
+    if training_args.do_train:
+        ppo_trainer.ppo_train(max_target_length=data_args.max_target_length)
+        ppo_trainer.save_model()
+        ppo_trainer.save_state() # must be called after save_model to have a folder
+        if ppo_trainer.is_world_process_zero() and model_args.plot_loss:
+            plot_loss(training_args.output_dir, keys=["loss", "reward"])

src/llmtuner/tuner/pt/workflow.py

@@ -1,57 +1,43 @@
 # Inspired by: https://github.com/huggingface/transformers/blob/v4.29.2/examples/pytorch/language-modeling/run_clm.py
 
 import math
-from typing import Optional, List
-from transformers import Seq2SeqTrainingArguments, DataCollatorForSeq2Seq, TrainerCallback
+from typing import TYPE_CHECKING, Optional, List
+from transformers import DataCollatorForLanguageModeling, Trainer
 
-from llmtuner.dsets import get_dataset, preprocess_dataset
-from llmtuner.extras.callbacks import LogCallback
-from llmtuner.extras.constants import IGNORE_INDEX
+from llmtuner.dsets import get_dataset, preprocess_dataset, split_dataset
 from llmtuner.extras.ploting import plot_loss
-from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
 from llmtuner.tuner.core import load_model_and_tokenizer
-from llmtuner.tuner.core.trainer import PeftTrainer
+
+if TYPE_CHECKING:
+    from transformers import Seq2SeqTrainingArguments, TrainerCallback
+    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
 
 
 def run_pt(
-    model_args: ModelArguments,
-    data_args: DataArguments,
-    training_args: Seq2SeqTrainingArguments,
-    finetuning_args: FinetuningArguments,
-    callbacks: Optional[List[TrainerCallback]] = [LogCallback()]
+    model_args: "ModelArguments",
+    data_args: "DataArguments",
+    training_args: "Seq2SeqTrainingArguments",
+    finetuning_args: "FinetuningArguments",
+    callbacks: Optional[List["TrainerCallback"]] = None
 ):
     dataset = get_dataset(model_args, data_args)
     model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="pt")
     dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="pt")
-    data_collator = DataCollatorForSeq2Seq(
-        tokenizer=tokenizer,
-        label_pad_token_id=IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
-    )
-
-    # Split the dataset
-    if training_args.do_train:
-        if data_args.dev_ratio > 1e-6:
-            dataset = dataset.train_test_split(test_size=data_args.dev_ratio)
-            trainer_kwargs = {"train_dataset": dataset["train"], "eval_dataset": dataset["test"]}
-        else:
-            trainer_kwargs = {"train_dataset": dataset}
-    else: # do_eval or do_predict
-        trainer_kwargs = {"eval_dataset": dataset}
+    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,
         data_collator=data_collator,
         callbacks=callbacks,
-        **trainer_kwargs
+        **split_dataset(dataset, data_args, training_args)
     )
 
     # Training
     if training_args.do_train:
-        train_result = trainer.train()
+        train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
         trainer.log_metrics("train", train_result.metrics)
         trainer.save_metrics("train", train_result.metrics)
         trainer.save_state()
@@ -68,6 +54,5 @@ def run_pt(
             perplexity = float("inf")
 
         metrics["perplexity"] = perplexity
-
         trainer.log_metrics("eval", metrics)
         trainer.save_metrics("eval", metrics)

src/llmtuner/tuner/rm/collator.py

@@ -1,8 +1,10 @@
 import torch
+from dataclasses import dataclass
 from typing import Any, Dict, Sequence
 from transformers import DataCollatorWithPadding
 
 
+@dataclass
 class PairwiseDataCollatorWithPadding(DataCollatorWithPadding):
     r"""
     Data collator for pairwise data.
@@ -15,5 +17,11 @@ class PairwiseDataCollatorWithPadding(DataCollatorWithPadding):
         We generate 2 * n examples where the first n examples represent chosen examples and
         the last n examples represent rejected examples.
         """
-        features = [{"input_ids": feature[key]} for key in ("accept_ids", "reject_ids") for feature in features]
+        features = [
+            {
+                "input_ids": feature["prompt_ids"] + feature[key],
+                "attention_mask": [1] * (len(feature["prompt_ids"]) + len(feature[key]))
+            }
+            for key in ("chosen_ids", "rejected_ids") for feature in features
+        ]
         return super().__call__(features)

src/llmtuner/tuner/rm/trainer.py

@@ -1,11 +1,20 @@
+import os
+import json
 import torch
-from typing import Dict, List, Optional, Tuple, Union
-from transformers.modeling_utils import PreTrainedModel
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+from transformers import Trainer
 
-from llmtuner.tuner.core.trainer import PeftTrainer
+from llmtuner.extras.logging import get_logger
+
+if TYPE_CHECKING:
+    from transformers.trainer import PredictionOutput
+    from transformers.modeling_utils import PreTrainedModel
 
 
-class PairwisePeftTrainer(PeftTrainer):
+logger = get_logger(__name__)
+
+
+class PairwiseTrainer(Trainer):
     r"""
     Inherits PeftTrainer to compute pairwise loss.
     """
@@ -16,23 +25,81 @@ class PairwisePeftTrainer(PeftTrainer):
 
     def compute_loss(
         self,
-        model: PreTrainedModel,
+        model: "PreTrainedModel",
         inputs: Dict[str, torch.Tensor],
         return_outputs: Optional[bool] = False
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[torch.Tensor]]]:
         r"""
         Computes pairwise loss. The first n examples are chosen and the last n examples are rejected.
 
-        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
         """
+        # Compute rewards
+        _, _, values = model(**inputs, output_hidden_states=True, return_dict=True)
+        if values.size(0) != inputs["input_ids"].size(0): # adapt to chatglm2
+            values = torch.transpose(values, 0, 1)
+
+        # Split the inputs and rewards into two parts, chosen and rejected
         batch_size = inputs["input_ids"].size(0) // 2
-        _, _, values = model(**inputs)
-        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
+        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,
+        predict_results: "PredictionOutput"
+    ) -> None:
+        r"""
+        Saves model predictions to `output_dir`.
+
+        A custom behavior that not contained in Seq2SeqTrainer.
+        """
+        if not self.is_world_process_zero():
+            return
+
+        output_prediction_file = os.path.join(self.args.output_dir, "generated_predictions.jsonl")
+        logger.info(f"Saving prediction results to {output_prediction_file}")
+
+        chosen_scores, rejected_scores = predict_results.predictions
+
+        with open(output_prediction_file, "w", encoding="utf-8") as writer:
+            res: List[str] = []
+            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))

src/llmtuner/tuner/rm/workflow.py

@@ -1,52 +1,47 @@
 # 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
-from llmtuner.extras.callbacks import LogCallback
+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.hparams import ModelArguments, DataArguments, FinetuningArguments
 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 TrainerCallback
+    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
 
 
 def run_rm(
-    model_args: ModelArguments,
-    data_args: DataArguments,
-    training_args: Seq2SeqTrainingArguments,
-    finetuning_args: FinetuningArguments
+    model_args: "ModelArguments",
+    data_args: "DataArguments",
+    training_args: "Seq2SeqTrainingArguments",
+    finetuning_args: "FinetuningArguments",
+    callbacks: Optional[List["TrainerCallback"]] = None
 ):
     dataset = get_dataset(model_args, data_args)
     model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="rm")
     dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="rm")
     data_collator = PairwiseDataCollatorWithPadding(tokenizer)
 
-    training_args.remove_unused_columns = False # important for pairwise dataset
-
-    # Split the dataset
-    if training_args.do_train:
-        if data_args.dev_ratio > 1e-6:
-            dataset = dataset.train_test_split(test_size=data_args.dev_ratio)
-            trainer_kwargs = {"train_dataset": dataset["train"], "eval_dataset": dataset["test"]}
-        else:
-            trainer_kwargs = {"train_dataset": dataset}
-    else: # do_eval or do_predict
-        trainer_kwargs = {"eval_dataset": 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=[LogCallback()],
+        callbacks=callbacks + [SavePeftModelCallback()],
         compute_metrics=compute_accuracy,
-        **trainer_kwargs
+        **split_dataset(dataset, data_args, training_args)
     )
 
     # Training
@@ -64,3 +59,10 @@ def run_rm(
         metrics = trainer.evaluate(metric_key_prefix="eval")
         trainer.log_metrics("eval", metrics)
         trainer.save_metrics("eval", metrics)
+
+    # Predict
+    if training_args.do_predict:
+        predict_results = trainer.predict(dataset, metric_key_prefix="predict")
+        trainer.log_metrics("predict", predict_results.metrics)
+        trainer.save_metrics("predict", predict_results.metrics)
+        trainer.save_predictions(predict_results)

src/llmtuner/tuner/sft/metric.py

@@ -1,7 +1,6 @@
 import numpy as np
 from dataclasses import dataclass
-from typing import Dict, Sequence, Tuple, Union
-from transformers.tokenization_utils import PreTrainedTokenizer
+from typing import TYPE_CHECKING, Dict, Sequence, Tuple, Union
 
 import jieba
 from rouge_chinese import Rouge
@@ -9,6 +8,9 @@ from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
 
 from llmtuner.extras.constants import IGNORE_INDEX
 
+if TYPE_CHECKING:
+    from transformers.tokenization_utils import PreTrainedTokenizer
+
 
 @dataclass
 class ComputeMetrics:
@@ -16,7 +18,7 @@ class ComputeMetrics:
     Wraps the tokenizer into metric functions, used in Seq2SeqPeftTrainer.
     """
 
-    tokenizer: PreTrainedTokenizer
+    tokenizer: "PreTrainedTokenizer"
 
     def __call__(self, eval_preds: Sequence[Union[np.ndarray, Tuple[np.ndarray]]]) -> Dict[str, float]:
         r"""

src/llmtuner/tuner/sft/trainer.py

@@ -3,18 +3,20 @@ import json
 import torch
 import numpy as np
 import torch.nn as nn
-from typing import Any, Dict, List, Optional, Tuple, Union
-from transformers.trainer import PredictionOutput
+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
 
 
 logger = get_logger(__name__)
 
 
-class Seq2SeqPeftTrainer(PeftTrainer):
+class CustomSeq2SeqTrainer(Seq2SeqTrainer):
     r"""
     Inherits PeftTrainer to compute generative metrics such as BLEU and ROUGE.
     """
@@ -31,21 +33,49 @@ class Seq2SeqPeftTrainer(PeftTrainer):
 
         Subclass and override to inject custom behavior.
         """
-        prompt_len, label_len = inputs["input_ids"].size(-1), inputs["labels"].size(-1)
-        if self.tokenizer.padding_side == "right": # pads the labels to the same length as the inputs
-            inputs["labels"] = torch.cat((inputs["labels"], torch.zeros_like(inputs["input_ids"])[:, label_len:]), dim=-1)
-        else:
-            inputs["labels"] = torch.cat((torch.zeros_like(inputs["input_ids"])[:, label_len:], inputs["labels"]), dim=-1)
+        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[:, prompt_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,
+        src_tensor: torch.Tensor,
+        tgt_tensor: torch.Tensor,
+        pad_token_id: Optional[int] = None
+    ) -> torch.Tensor:
+        r"""
+        Pads the tensor to the same length as the target tensor.
+        """
+        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
 
     def save_predictions(
         self,
-        predict_results: PredictionOutput
+        predict_results: "PredictionOutput"
     ) -> None:
         r"""
         Saves model predictions to `output_dir`.

src/llmtuner/tuner/sft/workflow.py

@@ -1,74 +1,69 @@
 # Inspired by: https://github.com/huggingface/transformers/blob/v4.29.2/examples/pytorch/summarization/run_summarization.py
 
-from typing import Optional, List
-from transformers import Seq2SeqTrainingArguments, DataCollatorForSeq2Seq, TrainerCallback
+from typing import TYPE_CHECKING, Optional, List
+from transformers import DataCollatorForSeq2Seq, Seq2SeqTrainingArguments
 
-from llmtuner.dsets import get_dataset, preprocess_dataset
-from llmtuner.extras.callbacks import LogCallback
+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
 from llmtuner.extras.ploting import plot_loss
-from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
 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 TrainerCallback
+    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments
 
 
 def run_sft(
-    model_args: ModelArguments,
-    data_args: DataArguments,
-    training_args: Seq2SeqTrainingArguments,
-    finetuning_args: FinetuningArguments,
-    callbacks: Optional[List[TrainerCallback]] = [LogCallback()]
+    model_args: "ModelArguments",
+    data_args: "DataArguments",
+    training_args: "Seq2SeqTrainingArguments",
+    finetuning_args: "FinetuningArguments",
+    generating_args: "GeneratingArguments",
+    callbacks: Optional[List["TrainerCallback"]] = None
 ):
     dataset = get_dataset(model_args, data_args)
     model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="sft")
     dataset = preprocess_dataset(dataset, tokenizer, data_args, training_args, stage="sft")
+
+    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
-
-    # Split the dataset
-    if training_args.do_train:
-        if data_args.dev_ratio > 1e-6:
-            dataset = dataset.train_test_split(test_size=data_args.dev_ratio)
-            trainer_kwargs = {"train_dataset": dataset["train"], "eval_dataset": dataset["test"]}
-        else:
-            trainer_kwargs = {"train_dataset": dataset}
-    else: # do_eval or do_predict
-        trainer_kwargs = {"eval_dataset": dataset}
+    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,
         data_collator=data_collator,
         callbacks=callbacks,
         compute_metrics=ComputeMetrics(tokenizer) if training_args.predict_with_generate else None,
-        **trainer_kwargs
+        **split_dataset(dataset, data_args, training_args)
     )
 
     # Keyword arguments for `model.generate`
-    gen_kwargs = {
-        "do_sample": True,
-        "top_p": 0.7,
-        "max_new_tokens": data_args.max_target_length + 1,
-        "temperature": 0.95,
-        "logits_processor": get_logits_processor()
-    }
+    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()
 
     # 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()

src/llmtuner/tuner/tune.py

@@ -0,0 +1,48 @@
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from llmtuner.extras.callbacks import LogCallback
+from llmtuner.extras.logging import get_logger
+from llmtuner.tuner.core import get_train_args, load_model_and_tokenizer
+from llmtuner.tuner.pt import run_pt
+from llmtuner.tuner.sft import run_sft
+from llmtuner.tuner.rm import run_rm
+from llmtuner.tuner.ppo import run_ppo
+from llmtuner.tuner.dpo import run_dpo
+
+if TYPE_CHECKING:
+    from transformers import TrainerCallback
+
+
+logger = get_logger(__name__)
+
+
+def run_exp(args: Optional[Dict[str, Any]] = None, callbacks: Optional[List["TrainerCallback"]] = None):
+    model_args, data_args, training_args, finetuning_args, generating_args, general_args = get_train_args(args)
+    callbacks = [LogCallback()] if callbacks is None else callbacks
+
+    if general_args.stage == "pt":
+        run_pt(model_args, data_args, training_args, finetuning_args, callbacks)
+    elif general_args.stage == "sft":
+        run_sft(model_args, data_args, training_args, finetuning_args, generating_args, callbacks)
+    elif general_args.stage == "rm":
+        run_rm(model_args, data_args, training_args, finetuning_args, callbacks)
+    elif general_args.stage == "ppo":
+        run_ppo(model_args, data_args, training_args, finetuning_args, generating_args, callbacks)
+    elif general_args.stage == "dpo":
+        run_dpo(model_args, data_args, training_args, finetuning_args, callbacks)
+    else:
+        raise ValueError("Unknown task.")
+
+
+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, tokenizer = load_model_and_tokenizer(model_args, finetuning_args)
+    model.save_pretrained(training_args.output_dir, max_shard_size=max_shard_size)
+    try:
+        tokenizer.save_pretrained(training_args.output_dir)
+    except:
+        logger.warning("Cannot save tokenizer, please copy the files manually.")
+
+
+if __name__ == "__main__":
+    run_exp()

src/llmtuner/webui/__init__.py

@@ -0,0 +1 @@
+from llmtuner.webui.interface import create_ui, create_web_demo

src/llmtuner/webui/chat.py

@@ -0,0 +1,97 @@
+import os
+from typing import Any, Dict, List, Optional, Tuple
+
+from llmtuner.chat.stream_chat import ChatModel
+from llmtuner.extras.misc import torch_gc
+from llmtuner.hparams import GeneratingArguments
+from llmtuner.webui.common import get_model_path, get_save_dir
+from llmtuner.webui.locales import ALERTS
+
+
+class WebChatModel(ChatModel):
+
+    def __init__(self, args: Optional[Dict[str, Any]] = None, lazy_init: Optional[bool] = True) -> None:
+        if lazy_init:
+            self.model = None
+            self.tokenizer = None
+            self.generating_args = GeneratingArguments()
+        else:
+            super().__init__(args)
+
+    def load_model(
+        self,
+        lang: str,
+        model_name: str,
+        checkpoints: List[str],
+        finetuning_type: str,
+        quantization_bit: str,
+        template: str,
+        system_prompt: str
+    ):
+        if self.model is not None:
+            yield ALERTS["err_exists"][lang]
+            return
+
+        if not model_name:
+            yield ALERTS["err_no_model"][lang]
+            return
+
+        model_name_or_path = get_model_path(model_name)
+        if not model_name_or_path:
+            yield ALERTS["err_no_path"][lang]
+            return
+
+        if checkpoints:
+            checkpoint_dir = ",".join(
+                [os.path.join(get_save_dir(model_name), finetuning_type, checkpoint) for checkpoint in checkpoints]
+            )
+        else:
+            checkpoint_dir = None
+
+        yield ALERTS["info_loading"][lang]
+        args = dict(
+            model_name_or_path=model_name_or_path,
+            checkpoint_dir=checkpoint_dir,
+            finetuning_type=finetuning_type,
+            quantization_bit=int(quantization_bit) if quantization_bit and quantization_bit != "None" else None,
+            template=template,
+            system_prompt=system_prompt
+        )
+        super().__init__(args)
+
+        yield ALERTS["info_loaded"][lang]
+
+    def unload_model(self, lang: str):
+        yield ALERTS["info_unloading"][lang]
+        self.model = None
+        self.tokenizer = None
+        torch_gc()
+        yield ALERTS["info_unloaded"][lang]
+
+    def predict(
+        self,
+        chatbot: List[Tuple[str, str]],
+        query: str,
+        history: List[Tuple[str, str]],
+        system: str,
+        max_new_tokens: int,
+        top_p: float,
+        temperature: float
+    ):
+        chatbot.append([query, ""])
+        response = ""
+        for new_text in self.stream_chat(
+            query, history, system, max_new_tokens=max_new_tokens, top_p=top_p, temperature=temperature
+        ):
+            response += new_text
+            response = self.postprocess(response)
+            new_history = history + [(query, response)]
+            chatbot[-1] = [query, response]
+            yield chatbot, new_history
+
+    def postprocess(self, response: str) -> str:
+        blocks = response.split("```")
+        for i, block in enumerate(blocks):
+            if i % 2 == 0:
+                blocks[i] = block.replace("<", "&lt;").replace(">", "&gt;")
+        return "```".join(blocks)

src/llmtuner/webui/common.py

@@ -0,0 +1,87 @@
+import json
+import os
+from typing import Any, Dict, Optional
+
+import gradio as gr
+from peft.utils import WEIGHTS_NAME as PEFT_WEIGHTS_NAME
+from transformers.trainer import WEIGHTS_NAME, WEIGHTS_INDEX_NAME
+
+from llmtuner.extras.constants import DEFAULT_TEMPLATE, SUPPORTED_MODELS, TRAINING_STAGES
+
+
+DEFAULT_CACHE_DIR = "cache"
+DEFAULT_DATA_DIR = "data"
+DEFAULT_SAVE_DIR = "saves"
+USER_CONFIG = "user.config"
+DATA_CONFIG = "dataset_info.json"
+
+
+def get_save_dir(*args) -> os.PathLike:
+    return os.path.join(DEFAULT_SAVE_DIR, *args)
+
+
+def get_config_path() -> os.PathLike:
+    return os.path.join(DEFAULT_CACHE_DIR, USER_CONFIG)
+
+
+def load_config() -> Dict[str, Any]:
+    try:
+        with open(get_config_path(), "r", encoding="utf-8") as f:
+            return json.load(f)
+    except:
+        return {"lang": None, "last_model": None, "path_dict": {}, "cache_dir": 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["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)
+
+
+def get_model_path(model_name: str) -> str:
+    user_config = load_config()
+    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 = get_save_dir(model_name, finetuning_type)
+    if save_dir and os.path.isdir(save_dir):
+        for checkpoint in os.listdir(save_dir):
+            if (
+                os.path.isdir(os.path.join(save_dir, checkpoint))
+                and any([
+                    os.path.isfile(os.path.join(save_dir, checkpoint, name))
+                    for name in (WEIGHTS_NAME, WEIGHTS_INDEX_NAME, PEFT_WEIGHTS_NAME)
+                ])
+            ):
+                checkpoints.append(checkpoint)
+    return gr.update(value=[], choices=checkpoints)
+
+
+def load_dataset_info(dataset_dir: str) -> Dict[str, Any]:
+    try:
+        with open(os.path.join(dataset_dir, DATA_CONFIG), "r", encoding="utf-8") as f:
+            return json.load(f)
+    except:
+        return {}
+
+
+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)
+    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)

src/llmtuner/webui/components/__init__.py

@@ -0,0 +1,6 @@
+from llmtuner.webui.components.top import create_top
+from llmtuner.webui.components.train import create_train_tab
+from llmtuner.webui.components.eval import create_eval_tab
+from llmtuner.webui.components.infer import create_infer_tab
+from llmtuner.webui.components.export import create_export_tab
+from llmtuner.webui.components.chatbot import create_chat_box

src/llmtuner/webui/components/chatbot.py

@@ -0,0 +1,51 @@
+from typing import TYPE_CHECKING, Dict, Optional, Tuple
+
+import gradio as gr
+
+if TYPE_CHECKING:
+    from gradio.blocks import Block
+    from gradio.components import Component
+    from llmtuner.webui.chat import WebChatModel
+
+
+def create_chat_box(
+    chat_model: "WebChatModel",
+    visible: Optional[bool] = False
+) -> Tuple["Block", "Component", "Component", Dict[str, "Component"]]:
+    with gr.Box(visible=visible) as chat_box:
+        chatbot = gr.Chatbot()
+
+        with gr.Row():
+            with gr.Column(scale=4):
+                system = gr.Textbox(show_label=False)
+                query = gr.Textbox(show_label=False, lines=8)
+                submit_btn = gr.Button(variant="primary")
+
+            with gr.Column(scale=1):
+                clear_btn = gr.Button()
+                max_new_tokens = gr.Slider(10, 2048, value=chat_model.generating_args.max_new_tokens, step=1)
+                top_p = gr.Slider(0.01, 1, value=chat_model.generating_args.top_p, step=0.01)
+                temperature = gr.Slider(0.01, 1.5, value=chat_model.generating_args.temperature, step=0.01)
+
+    history = gr.State([])
+
+    submit_btn.click(
+        chat_model.predict,
+        [chatbot, query, history, system, max_new_tokens, top_p, temperature],
+        [chatbot, history],
+        show_progress=True
+    ).then(
+        lambda: gr.update(value=""), outputs=[query]
+    )
+
+    clear_btn.click(lambda: ([], []), outputs=[chatbot, history], show_progress=True)
+
+    return chat_box, chatbot, history, dict(
+        system=system,
+        query=query,
+        submit_btn=submit_btn,
+        clear_btn=clear_btn,
+        max_new_tokens=max_new_tokens,
+        top_p=top_p,
+        temperature=temperature
+    )

src/llmtuner/webui/components/data.py

@@ -0,0 +1,21 @@
+import gradio as gr
+from typing import TYPE_CHECKING, Tuple
+
+if TYPE_CHECKING:
+    from gradio.blocks import Block
+    from gradio.components import Component
+
+
+def create_preview_box() -> Tuple["Block", "Component", "Component", "Component"]:
+    with gr.Box(visible=False, elem_classes="modal-box") as preview_box:
+        with gr.Row():
+            preview_count = gr.Number(interactive=False)
+
+        with gr.Row():
+            preview_samples = gr.JSON(interactive=False)
+
+        close_btn = gr.Button()
+
+    close_btn.click(lambda: gr.update(visible=False), outputs=[preview_box], queue=False)
+
+    return preview_box, preview_count, preview_samples, close_btn

src/llmtuner/webui/components/eval.py

@@ -0,0 +1,90 @@
+from typing import TYPE_CHECKING, Dict
+import gradio as gr
+
+from llmtuner.webui.common import list_dataset, DEFAULT_DATA_DIR
+from llmtuner.webui.components.data import create_preview_box
+from llmtuner.webui.utils import can_preview, get_preview
+
+if TYPE_CHECKING:
+    from gradio.components import Component
+    from llmtuner.webui.runner import Runner
+
+
+def create_eval_tab(top_elems: Dict[str, "Component"], runner: "Runner") -> Dict[str, "Component"]:
+    with gr.Row():
+        dataset_dir = gr.Textbox(value=DEFAULT_DATA_DIR, scale=2)
+        dataset = gr.Dropdown(multiselect=True, scale=4)
+        data_preview_btn = gr.Button(interactive=False, scale=1)
+
+    preview_box, preview_count, preview_samples, close_btn = create_preview_box()
+
+    dataset_dir.change(list_dataset, [dataset_dir], [dataset])
+    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],
+        queue=False
+    )
+
+    with gr.Row():
+        max_source_length = gr.Slider(value=512, minimum=4, maximum=4096, step=1)
+        max_target_length = gr.Slider(value=512, minimum=4, maximum=4096, step=1)
+        max_samples = gr.Textbox(value="100000")
+        batch_size = gr.Slider(value=8, minimum=1, maximum=512, step=1)
+        predict = gr.Checkbox(value=True)
+
+    with gr.Row():
+        cmd_preview_btn = gr.Button()
+        start_btn = gr.Button()
+        stop_btn = gr.Button()
+
+    with gr.Row():
+        process_bar = gr.Slider(visible=False, interactive=False)
+
+    with gr.Box():
+        output_box = gr.Markdown()
+
+    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,
+        data_preview_btn=data_preview_btn,
+        preview_count=preview_count,
+        preview_samples=preview_samples,
+        close_btn=close_btn,
+        max_source_length=max_source_length,
+        max_target_length=max_target_length,
+        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
+    )

src/llmtuner/webui/components/export.py

@@ -0,0 +1,37 @@
+from typing import TYPE_CHECKING, Dict
+import gradio as gr
+
+from llmtuner.webui.utils import save_model
+
+if TYPE_CHECKING:
+    from gradio.components import Component
+
+
+def create_export_tab(top_elems: Dict[str, "Component"]) -> Dict[str, "Component"]:
+    with gr.Row():
+        save_dir = gr.Textbox()
+        max_shard_size = gr.Slider(value=10, minimum=1, maximum=100)
+
+    export_btn = gr.Button()
+    info_box = gr.Textbox(show_label=False, interactive=False)
+
+    export_btn.click(
+        save_model,
+        [
+            top_elems["lang"],
+            top_elems["model_name"],
+            top_elems["checkpoints"],
+            top_elems["finetuning_type"],
+            top_elems["template"],
+            max_shard_size,
+            save_dir
+        ],
+        [info_box]
+    )
+
+    return dict(
+        save_dir=save_dir,
+        max_shard_size=max_shard_size,
+        export_btn=export_btn,
+        info_box=info_box
+    )

src/llmtuner/webui/components/infer.py

@@ -0,0 +1,51 @@
+from typing import TYPE_CHECKING, Dict
+
+import gradio as gr
+
+from llmtuner.webui.chat import WebChatModel
+from llmtuner.webui.components.chatbot import create_chat_box
+
+if TYPE_CHECKING:
+    from gradio.components import Component
+
+
+def create_infer_tab(top_elems: Dict[str, "Component"]) -> Dict[str, "Component"]:
+    with gr.Row():
+        load_btn = gr.Button()
+        unload_btn = gr.Button()
+
+    info_box = gr.Textbox(show_label=False, interactive=False)
+
+    chat_model = WebChatModel(lazy_init=True)
+    chat_box, chatbot, history, chat_elems = create_chat_box(chat_model)
+
+    load_btn.click(
+        chat_model.load_model,
+        [
+            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"]
+        ],
+        [info_box]
+    ).then(
+        lambda: gr.update(visible=(chat_model.model is not None)), outputs=[chat_box]
+    )
+
+    unload_btn.click(
+        chat_model.unload_model, [top_elems["lang"]], [info_box]
+    ).then(
+        lambda: ([], []), outputs=[chatbot, history]
+    ).then(
+        lambda: gr.update(visible=(chat_model.model is not None)), outputs=[chat_box]
+    )
+
+    return dict(
+        info_box=info_box,
+        load_btn=load_btn,
+        unload_btn=unload_btn,
+        **chat_elems
+    )

src/llmtuner/webui/components/top.py

@@ -0,0 +1,66 @@
+from typing import TYPE_CHECKING, Dict
+
+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, get_template, save_config
+from llmtuner.webui.utils import can_quantize
+
+if TYPE_CHECKING:
+    from gradio.components import Component
+
+
+def create_top() -> Dict[str, "Component"]:
+    available_models = list(SUPPORTED_MODELS.keys()) + ["Custom"]
+
+    with gr.Row():
+        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(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(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, [lang, model_name, model_path])
+
+    finetuning_type.change(
+        list_checkpoint, [model_name, finetuning_type], [checkpoints]
+    ).then(
+        can_quantize, [finetuning_type], [quantization_bit]
+    )
+
+    refresh_btn.click(
+        list_checkpoint, [model_name, finetuning_type], [checkpoints], queue=False
+    )
+
+    return dict(
+        lang=lang,
+        model_name=model_name,
+        model_path=model_path,
+        finetuning_type=finetuning_type,
+        checkpoints=checkpoints,
+        refresh_btn=refresh_btn,
+        advanced_tab=advanced_tab,
+        quantization_bit=quantization_bit,
+        template=template,
+        system_prompt=system_prompt
+    )

src/llmtuner/webui/components/train.py

@@ -0,0 +1,184 @@
+from typing import TYPE_CHECKING, Dict
+from transformers.trainer_utils import SchedulerType
+
+import gradio as gr
+
+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
+
+if TYPE_CHECKING:
+    from gradio.components import Component
+    from llmtuner.webui.runner import Runner
+
+
+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)
+        data_preview_btn = gr.Button(interactive=False, scale=1)
+
+    preview_box, preview_count, preview_samples, close_btn = create_preview_box()
+
+    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],
+        queue=False
+    )
+
+    with gr.Row():
+        max_source_length = gr.Slider(value=512, minimum=4, maximum=4096, step=1)
+        max_target_length = gr.Slider(value=512, minimum=4, maximum=4096, step=1)
+        learning_rate = gr.Textbox(value="5e-5")
+        num_train_epochs = gr.Textbox(value="3.0")
+        max_samples = gr.Textbox(value="100000")
+
+    with gr.Row():
+        batch_size = gr.Slider(value=4, minimum=1, maximum=512, step=1)
+        gradient_accumulation_steps = gr.Slider(value=4, minimum=1, maximum=512, step=1)
+        lr_scheduler_type = gr.Dropdown(
+            choices=[scheduler.value for scheduler in SchedulerType], value="cosine"
+        )
+        max_grad_norm = gr.Textbox(value="1.0")
+        val_size = gr.Slider(value=0, minimum=0, maximum=1, step=0.001)
+
+    with gr.Accordion(label="Advanced config", open=False) as advanced_tab:
+        with gr.Row():
+            logging_steps = gr.Slider(value=5, minimum=5, maximum=1000, step=5)
+            save_steps = gr.Slider(value=100, minimum=10, maximum=5000, step=10)
+            warmup_steps = gr.Slider(value=0, minimum=0, maximum=5000, step=1)
+            compute_type = gr.Radio(choices=["fp16", "bf16"], value="fp16")
+
+    with gr.Accordion(label="LoRA config", open=False) as lora_tab:
+        with gr.Row():
+            lora_rank = gr.Slider(value=8, minimum=1, maximum=1024, step=1, scale=1)
+            lora_dropout = gr.Slider(value=0.1, minimum=0, maximum=1, step=0.01, scale=1)
+            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()
+
+    with gr.Row():
+        with gr.Column(scale=3):
+            with gr.Row():
+                output_dir = gr.Textbox()
+
+            with gr.Row():
+                process_bar = gr.Slider(visible=False, interactive=False)
+
+            with gr.Box():
+                output_box = gr.Markdown()
+
+        with gr.Column(scale=1):
+            loss_viewer = gr.Plot()
+
+    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(
+        gen_plot, [top_elems["model_name"], top_elems["finetuning_type"], output_dir], loss_viewer, queue=False
+    )
+
+    return dict(
+        training_stage=training_stage,
+        dataset_dir=dataset_dir,
+        dataset=dataset,
+        data_preview_btn=data_preview_btn,
+        preview_count=preview_count,
+        preview_samples=preview_samples,
+        close_btn=close_btn,
+        max_source_length=max_source_length,
+        max_target_length=max_target_length,
+        learning_rate=learning_rate,
+        num_train_epochs=num_train_epochs,
+        max_samples=max_samples,
+        batch_size=batch_size,
+        gradient_accumulation_steps=gradient_accumulation_steps,
+        lr_scheduler_type=lr_scheduler_type,
+        max_grad_norm=max_grad_norm,
+        val_size=val_size,
+        advanced_tab=advanced_tab,
+        logging_steps=logging_steps,
+        save_steps=save_steps,
+        warmup_steps=warmup_steps,
+        compute_type=compute_type,
+        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,
+        output_box=output_box,
+        loss_viewer=loss_viewer
+    )

src/llmtuner/webui/css.py

@@ -0,0 +1,18 @@
+CSS = r"""
+.modal-box {
+  position: fixed !important;
+  top: 50%;
+  left: 50%;
+  transform: translate(-50%, -50%); /* center horizontally */
+  max-width: 1000px;
+  max-height: 750px;
+  overflow-y: scroll !important;
+  background-color: var(--input-background-fill);
+  border: 2px solid black !important;
+  z-index: 1000;
+}
+
+.dark .modal-box {
+  border: 2px solid white !important;
+}
+"""

src/llmtuner/webui/interface.py

@@ -0,0 +1,78 @@
+import gradio as gr
+from transformers.utils.versions import require_version
+
+from llmtuner.webui.components import (
+    create_top,
+    create_train_tab,
+    create_eval_tab,
+    create_infer_tab,
+    create_export_tab,
+    create_chat_box
+)
+from llmtuner.webui.chat import WebChatModel
+from llmtuner.webui.css import CSS
+from llmtuner.webui.manager import Manager
+from llmtuner.webui.runner import Runner
+
+
+require_version("gradio>=3.36.0", "To fix: pip install gradio>=3.36.0")
+
+
+def create_ui() -> gr.Blocks:
+    runner = Runner()
+
+    with gr.Blocks(title="Web Tuner", css=CSS) as demo:
+        top_elems = create_top()
+
+        with gr.Tab("Train"):
+            train_elems = create_train_tab(top_elems, runner)
+
+        with gr.Tab("Evaluate"):
+            eval_elems = create_eval_tab(top_elems, runner)
+
+        with gr.Tab("Chat"):
+            infer_elems = create_infer_tab(top_elems)
+
+        with gr.Tab("Export"):
+            export_elems = create_export_tab(top_elems)
+
+        elem_list = [top_elems, train_elems, eval_elems, infer_elems, export_elems]
+        manager = Manager(elem_list)
+
+        demo.load(
+            manager.gen_label,
+            [top_elems["lang"]],
+            [elem for elems in elem_list for elem in elems.values()],
+        )
+
+        top_elems["lang"].change(
+            manager.gen_label,
+            [top_elems["lang"]],
+            [elem for elems in elem_list for elem in elems.values()],
+            queue=False
+        )
+
+    return demo
+
+
+def create_web_demo() -> gr.Blocks:
+    chat_model = WebChatModel(lazy_init=False)
+
+    with gr.Blocks(title="Web Demo", css=CSS) as demo:
+        lang = gr.Dropdown(choices=["en", "zh"], value="en")
+
+        _, _, _, chat_elems = create_chat_box(chat_model, visible=True)
+
+        manager = Manager([{"lang": lang}, chat_elems])
+
+        demo.load(manager.gen_label, [lang], [lang] + list(chat_elems.values()))
+
+        lang.select(manager.gen_label, [lang], [lang] + list(chat_elems.values()), queue=False)
+
+    return demo
+
+
+if __name__ == "__main__":
+    demo = create_ui()
+    demo.queue()
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False, inbrowser=True)

src/llmtuner/webui/locales.py

@@ -0,0 +1,612 @@
+LOCALES = {
+    "lang": {
+        "en": {
+            "label": "Lang"
+        },
+        "zh": {
+            "label": "语言"
+        }
+    },
+    "model_name": {
+        "en": {
+            "label": "Model name"
+        },
+        "zh": {
+            "label": "模型名称"
+        }
+    },
+    "model_path": {
+        "en": {
+            "label": "Model path",
+            "info": "Path to pretrained model or model identifier from Hugging Face."
+        },
+        "zh": {
+            "label": "模型路径",
+            "info": "本地模型的文件路径或 Hugging Face 的模型标识符。"
+        }
+    },
+    "finetuning_type": {
+        "en": {
+            "label": "Finetuning method"
+        },
+        "zh": {
+            "label": "微调方法"
+        }
+    },
+    "checkpoints": {
+        "en": {
+            "label": "Checkpoints"
+        },
+        "zh": {
+            "label": "模型断点"
+        }
+    },
+    "refresh_btn": {
+        "en": {
+            "value": "Refresh checkpoints"
+        },
+        "zh": {
+            "value": "刷新断点"
+        }
+    },
+    "advanced_tab": {
+        "en": {
+            "label": "Advanced configurations"
+        },
+        "zh": {
+            "label": "高级设置"
+        }
+    },
+    "quantization_bit": {
+        "en": {
+            "label": "Quantization bit (optional)",
+            "info": "Enable 4/8-bit model quantization."
+        },
+        "zh": {
+            "label": "量化等级（非必填）",
+            "info": "启用 4/8 比特模型量化。"
+        }
+    },
+    "template": {
+        "en": {
+            "label": "Prompt template",
+            "info": "The template used in constructing prompts."
+        },
+        "zh": {
+            "label": "提示模板",
+            "info": "构建提示词时使用的模板"
+        }
+    },
+    "system_prompt": {
+        "en": {
+            "label": "System prompt (optional)",
+            "info": "A sequence used as the default system prompt."
+        },
+        "zh": {
+            "label": "系统提示词（非必填）",
+            "info": "默认使用的系统提示词"
+        }
+    },
+    "training_stage": {
+        "en": {
+            "label": "Stage",
+            "info": "The stage to perform in training."
+        },
+        "zh": {
+            "label": "训练阶段",
+            "info": "目前采用的训练方式。"
+        }
+    },
+    "dataset_dir": {
+        "en": {
+            "label": "Data dir",
+            "info": "Path of the data directory."
+        },
+        "zh": {
+            "label": "数据路径",
+            "info": "数据文件夹的路径。"
+        }
+    },
+    "dataset": {
+        "en": {
+            "label": "Dataset"
+        },
+        "zh": {
+            "label": "数据集"
+        }
+    },
+    "data_preview_btn": {
+        "en": {
+            "value": "Preview dataset"
+        },
+        "zh": {
+            "value": "预览数据集"
+        }
+    },
+    "preview_count": {
+        "en": {
+            "label": "Count"
+        },
+        "zh": {
+            "label": "数量"
+        }
+    },
+    "preview_samples": {
+        "en": {
+            "label": "Samples"
+        },
+        "zh": {
+            "label": "样例"
+        }
+    },
+    "close_btn": {
+        "en": {
+            "value": "Close"
+        },
+        "zh": {
+            "value": "关闭"
+        }
+    },
+    "max_source_length": {
+        "en": {
+            "label": "Max source length",
+            "info": "Max tokens in source sequence."
+        },
+        "zh": {
+            "label": "输入序列最大长度",
+            "info": "输入序列分词后的最大长度。"
+        }
+    },
+    "max_target_length": {
+        "en": {
+            "label": "Max target length",
+            "info": "Max tokens in target sequence."
+        },
+        "zh": {
+            "label": "输出序列最大长度",
+            "info": "输出序列分词后的最大长度。"
+        }
+    },
+    "learning_rate": {
+        "en": {
+            "label": "Learning rate",
+            "info": "Initial learning rate for AdamW."
+        },
+        "zh": {
+            "label": "学习率",
+            "info": "AdamW 优化器的初始学习率。"
+        }
+    },
+    "num_train_epochs": {
+        "en": {
+            "label": "Epochs",
+            "info": "Total number of training epochs to perform."
+        },
+        "zh": {
+            "label": "训练轮数",
+            "info": "需要执行的训练总轮数。"
+        }
+    },
+    "max_samples": {
+        "en": {
+            "label": "Max samples",
+            "info": "Maximum samples per dataset."
+        },
+        "zh": {
+            "label": "最大样本数",
+            "info": "每个数据集最多使用的样本数。"
+        }
+    },
+    "batch_size": {
+        "en": {
+            "label": "Batch size",
+            "info": "Number of samples to process per GPU."
+        },
+        "zh":{
+            "label": "批处理大小",
+            "info": "每块 GPU 上处理的样本数量。"
+        }
+    },
+    "gradient_accumulation_steps": {
+        "en": {
+            "label": "Gradient accumulation",
+            "info": "Number of gradient accumulation steps."
+        },
+        "zh": {
+            "label": "梯度累积",
+            "info": "梯度累积的步数。"
+        }
+    },
+    "lr_scheduler_type": {
+        "en": {
+            "label": "LR Scheduler",
+            "info": "Name of learning rate scheduler.",
+        },
+        "zh": {
+            "label": "学习率调节器",
+            "info": "采用的学习率调节器名称。"
+        }
+    },
+    "max_grad_norm": {
+        "en": {
+            "label": "Maximum gradient norm",
+            "info": "Norm for gradient clipping.."
+        },
+        "zh": {
+            "label": "最大梯度范数",
+            "info": "用于梯度裁剪的范数。"
+        }
+    },
+    "val_size": {
+        "en": {
+            "label": "Val size",
+            "info": "Proportion of data in the dev set."
+        },
+        "zh": {
+            "label": "验证集比例",
+            "info": "验证集占全部样本的百分比。"
+        }
+    },
+    "logging_steps": {
+        "en": {
+            "label": "Logging steps",
+            "info": "Number of steps between two logs."
+        },
+        "zh": {
+            "label": "日志间隔",
+            "info": "每两次日志输出间的更新步数。"
+        }
+    },
+    "save_steps": {
+        "en": {
+            "label": "Save steps",
+            "info": "Number of steps between two checkpoints."
+        },
+        "zh": {
+            "label": "保存间隔",
+            "info": "每两次断点保存间的更新步数。"
+        }
+    },
+    "warmup_steps": {
+        "en": {
+            "label": "Warmup steps",
+            "info": "Number of steps used for warmup."
+        },
+        "zh": {
+            "label": "预热步数",
+            "info": "学习率预热采用的步数。"
+        }
+    },
+    "compute_type": {
+        "en": {
+            "label": "Compute type",
+            "info": "Whether to use fp16 or bf16 mixed precision training."
+        },
+        "zh": {
+            "label": "计算类型",
+            "info": "是否启用 FP16 或 BF16 混合精度训练。"
+        }
+    },
+    "lora_tab": {
+        "en": {
+            "label": "LoRA configurations"
+        },
+        "zh": {
+            "label": "LoRA 参数设置"
+        }
+    },
+    "lora_rank": {
+        "en": {
+            "label": "LoRA rank",
+            "info": "The rank of LoRA matrices."
+        },
+        "zh": {
+            "label": "LoRA 秩",
+            "info": "LoRA 矩阵的秩。"
+        }
+    },
+    "lora_dropout": {
+        "en": {
+            "label": "LoRA Dropout",
+            "info": "Dropout ratio of LoRA weights."
+        },
+        "zh": {
+            "label": "LoRA 随机丢弃",
+            "info": "LoRA 权重随机丢弃的概率。"
+        }
+    },
+    "lora_target": {
+        "en": {
+            "label": "LoRA modules (optional)",
+            "info": "The name(s) of target modules to apply LoRA. Use commas to separate multiple modules."
+        },
+        "zh": {
+            "label": "LoRA 作用层（非必填）",
+            "info": "应用 LoRA 的线性层名称。使用英文逗号分隔多个名称。"
+        }
+    },
+    "resume_lora_training": {
+        "en": {
+            "label": "Resume LoRA training",
+            "info": "Whether to resume training from the last LoRA weights or create new lora weights."
+        },
+        "zh": {
+            "label": "继续上次的训练",
+            "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"
+        },
+        "zh": {
+            "value": "开始"
+        }
+    },
+    "stop_btn": {
+        "en": {
+            "value": "Abort"
+        },
+        "zh": {
+            "value": "中断"
+        }
+    },
+    "output_dir": {
+        "en": {
+            "label": "Checkpoint name",
+            "info": "Directory to save checkpoint."
+        },
+        "zh": {
+            "label": "断点名称",
+            "info": "保存模型断点的文件夹名称。"
+        }
+    },
+    "output_box": {
+        "en": {
+            "value": "Ready."
+        },
+        "zh": {
+            "value": "准备就绪。"
+        }
+    },
+    "loss_viewer": {
+        "en": {
+            "label": "Loss"
+        },
+        "zh": {
+            "label": "损失"
+        }
+    },
+    "predict": {
+        "en": {
+            "label": "Save predictions"
+        },
+        "zh": {
+            "label": "保存预测结果"
+        }
+    },
+    "load_btn": {
+        "en": {
+            "value": "Load model"
+        },
+        "zh": {
+            "value": "加载模型"
+        }
+    },
+    "unload_btn": {
+        "en": {
+            "value": "Unload model"
+        },
+        "zh": {
+            "value": "卸载模型"
+        }
+    },
+    "info_box": {
+        "en": {
+            "value": "Model unloaded, please load a model first."
+        },
+        "zh": {
+            "value": "模型未加载，请先加载模型。"
+        }
+    },
+    "system": {
+        "en": {
+            "placeholder": "System prompt (optional)"
+        },
+        "zh": {
+            "placeholder": "系统提示词（非必填）"
+        }
+    },
+    "query": {
+        "en": {
+            "placeholder": "Input..."
+        },
+        "zh": {
+            "placeholder": "输入..."
+        }
+    },
+    "submit_btn": {
+        "en": {
+            "value": "Submit"
+        },
+        "zh": {
+            "value": "提交"
+        }
+    },
+    "clear_btn": {
+        "en": {
+            "value": "Clear history"
+        },
+        "zh": {
+            "value": "清空历史"
+        }
+    },
+    "max_length": {
+        "en": {
+            "label": "Maximum length"
+        },
+        "zh": {
+            "label": "最大长度"
+        }
+    },
+    "max_new_tokens": {
+        "en": {
+            "label": "Maximum new tokens"
+        },
+        "zh": {
+            "label": "最大生成长度"
+        }
+    },
+    "top_p": {
+        "en": {
+            "label": "Top-p"
+        },
+        "zh": {
+            "label": "Top-p 采样值"
+        }
+    },
+    "temperature": {
+        "en": {
+            "label": "Temperature"
+        },
+        "zh": {
+            "label": "温度系数"
+        }
+    },
+    "save_dir": {
+        "en": {
+            "label": "Export dir",
+            "info": "Directory to save exported model."
+        },
+        "zh": {
+            "label": "导出目录",
+            "info": "保存导出模型的文件夹路径。"
+        }
+    },
+    "max_shard_size": {
+        "en": {
+            "label": "Max shard size (GB)",
+            "info": "The maximum size for a model file."
+        },
+        "zh": {
+            "label": "最大分块大小（GB）",
+            "info": "模型文件的最大大小。"
+        }
+    },
+    "export_btn": {
+        "en": {
+            "value": "Export"
+        },
+        "zh": {
+            "value": "开始导出"
+        }
+    }
+}
+
+
+ALERTS = {
+    "err_conflict": {
+        "en": "A process is in running, please abort it firstly.",
+        "zh": "任务已存在，请先中断训练。"
+    },
+    "err_exists": {
+        "en": "You have loaded a model, please unload it first.",
+        "zh": "模型已存在，请先卸载模型。"
+    },
+    "err_no_model": {
+        "en": "Please select a model.",
+        "zh": "请选择模型。"
+    },
+    "err_no_path": {
+        "en": "Model not found.",
+        "zh": "模型未找到。"
+    },
+    "err_no_dataset": {
+        "en": "Please choose a dataset.",
+        "zh": "请选择数据集。"
+    },
+    "err_no_checkpoint": {
+        "en": "Please select a checkpoint.",
+        "zh": "请选择断点。"
+    },
+    "err_no_save_dir": {
+        "en": "Please provide export dir.",
+        "zh": "请填写导出目录"
+    },
+    "err_failed": {
+        "en": "Failed.",
+        "zh": "训练出错。"
+    },
+    "info_aborting": {
+        "en": "Aborted, wait for terminating...",
+        "zh": "训练中断，正在等待线程结束……"
+    },
+    "info_aborted": {
+        "en": "Ready.",
+        "zh": "准备就绪。"
+    },
+    "info_finished": {
+        "en": "Finished.",
+        "zh": "训练完毕。"
+    },
+    "info_loading": {
+        "en": "Loading model...",
+        "zh": "加载中……"
+    },
+    "info_unloading": {
+        "en": "Unloading model...",
+        "zh": "卸载中……"
+    },
+    "info_loaded": {
+        "en": "Model loaded, now you can chat with your model!",
+        "zh": "模型已加载，可以开始聊天了！"
+    },
+    "info_unloaded": {
+        "en": "Model unloaded.",
+        "zh": "模型已卸载。"
+    },
+    "info_exporting": {
+        "en": "Exporting model...",
+        "zh": "正在导出模型……"
+    },
+    "info_exported": {
+        "en": "Model exported.",
+        "zh": "模型导出完成。"
+    }
+}

src/llmtuner/webui/manager.py

@@ -0,0 +1,46 @@
+import gradio as gr
+from gradio.components import Component
+from typing import Any, Dict, List
+
+from llmtuner.webui.common import get_model_path, list_dataset, load_config
+from llmtuner.webui.locales import LOCALES
+from llmtuner.webui.utils import get_time
+
+
+class Manager:
+
+    def __init__(self, elem_list: List[Dict[str, Component]]):
+        self.elem_list = elem_list
+
+    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 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"])}
+
+        return refresh_dict
+
+    def gen_label(self, lang: str) -> Dict[Component, Dict[str, Any]]: # cannot use TYPE_CHECKING
+        update_dict = {}
+        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][refresh_dict["lang"]["value"]], **refresh_dict.get(name, {})
+                )
+
+        return update_dict

src/llmtuner/webui/runner.py

@@ -0,0 +1,282 @@
+import gradio as gr
+import logging
+import os
+import threading
+import time
+import transformers
+from transformers.trainer import TRAINING_ARGS_NAME
+from typing import Any, Dict, Generator, List, Tuple
+
+from llmtuner.extras.callbacks import LogCallback
+from llmtuner.extras.constants import DEFAULT_MODULE, 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, load_config
+from llmtuner.webui.locales import ALERTS
+from llmtuner.webui.utils import gen_cmd, get_eval_results, update_process_bar
+
+
+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(
+        self, lang: str, model_name: str, dataset: List[str]
+    ) -> str:
+        if self.running:
+            return ALERTS["err_conflict"][lang]
+
+        if not model_name:
+            return ALERTS["err_no_model"][lang]
+
+        if not get_model_path(model_name):
+            return ALERTS["err_no_path"][lang]
+
+        if len(dataset) == 0:
+            return ALERTS["err_no_dataset"][lang]
+
+        self.aborted = False
+        self.logger_handler.reset()
+        self.trainer_callback = LogCallback(self)
+        return ""
+
+    def _finalize(
+        self, lang: str, finish_info: str
+    ) -> str:
+        self.running = False
+        torch_gc()
+        if self.aborted:
+            return ALERTS["info_aborted"][lang]
+        else:
+            return finish_info
+
+    def _parse_train_args(
+        self,
+        lang: str,
+        model_name: str,
+        checkpoints: List[str],
+        finetuning_type: str,
+        quantization_bit: str,
+        template: str,
+        system_prompt: str,
+        training_stage: str,
+        dataset_dir: str,
+        dataset: List[str],
+        max_source_length: int,
+        max_target_length: int,
+        learning_rate: str,
+        num_train_epochs: str,
+        max_samples: str,
+        batch_size: int,
+        gradient_accumulation_steps: int,
+        lr_scheduler_type: str,
+        max_grad_norm: str,
+        val_size: float,
+        logging_steps: int,
+        save_steps: int,
+        warmup_steps: int,
+        compute_type: str,
+        lora_rank: int,
+        lora_dropout: float,
+        lora_target: str,
+        resume_lora_training: bool,
+        dpo_beta: float,
+        reward_model: str,
+        output_dir: str
+    ) -> Tuple[str, str, List[str], str, Dict[str, Any]]:
+        if checkpoints:
+            checkpoint_dir = ",".join(
+                [get_save_dir(model_name, finetuning_type, ckpt) for ckpt in checkpoints]
+            )
+        else:
+            checkpoint_dir = None
+
+        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=TRAINING_STAGES[training_stage],
+            model_name_or_path=get_model_path(model_name),
+            do_train=True,
+            overwrite_cache=False,
+            cache_dir=cache_dir,
+            checkpoint_dir=checkpoint_dir,
+            finetuning_type=finetuning_type,
+            quantization_bit=int(quantization_bit) if quantization_bit in ["8", "4"] else None,
+            template=template,
+            system_prompt=system_prompt,
+            dataset_dir=dataset_dir,
+            dataset=",".join(dataset),
+            max_source_length=max_source_length,
+            max_target_length=max_target_length,
+            learning_rate=float(learning_rate),
+            num_train_epochs=float(num_train_epochs),
+            max_samples=int(max_samples),
+            per_device_train_batch_size=batch_size,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            lr_scheduler_type=lr_scheduler_type,
+            max_grad_norm=float(max_grad_norm),
+            logging_steps=logging_steps,
+            save_steps=save_steps,
+            warmup_steps=warmup_steps,
+            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=(
+                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
+            args["evaluation_strategy"] = "steps"
+            args["eval_steps"] = save_steps
+            args["load_best_model_at_end"] = True
+
+        return lang, model_name, dataset, output_dir, args
+
+    def _parse_eval_args(
+        self,
+        lang: str,
+        model_name: str,
+        checkpoints: List[str],
+        finetuning_type: str,
+        quantization_bit: str,
+        template: str,
+        system_prompt: str,
+        dataset_dir: str,
+        dataset: List[str],
+        max_source_length: int,
+        max_target_length: int,
+        max_samples: str,
+        batch_size: int,
+        predict: bool
+    ) -> Tuple[str, str, List[str], str, Dict[str, Any]]:
+        if checkpoints:
+            checkpoint_dir = ",".join(
+                [get_save_dir(model_name, finetuning_type, ckpt) for ckpt in checkpoints]
+            )
+            output_dir = get_save_dir(model_name, finetuning_type, "eval_" + "_".join(checkpoints))
+        else:
+            checkpoint_dir = None
+            output_dir = get_save_dir(model_name, finetuning_type, "eval_base")
+
+        user_config = load_config()
+        cache_dir = user_config.get("cache_dir", None)
+
+        args = dict(
+            stage="sft",
+            model_name_or_path=get_model_path(model_name),
+            do_eval=True,
+            overwrite_cache=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 in ["8", "4"] else None,
+            template=template,
+            system_prompt=system_prompt,
+            dataset_dir=dataset_dir,
+            dataset=",".join(dataset),
+            max_source_length=max_source_length,
+            max_target_length=max_target_length,
+            max_samples=int(max_samples),
+            per_device_eval_batch_size=batch_size,
+            output_dir=output_dir
+        )
+
+        if predict:
+            args.pop("do_eval", None)
+            args["do_predict"] = True
+
+        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()
+
+        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, 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)

src/llmtuner/webui/utils.py

@@ -0,0 +1,159 @@
+import os
+import json
+import gradio as gr
+import matplotlib.figure
+import matplotlib.pyplot as plt
+from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional, Tuple
+from datetime import datetime
+
+from llmtuner.extras.ploting import smooth
+from llmtuner.tuner import export_model
+from llmtuner.webui.common import get_model_path, get_save_dir, DATA_CONFIG
+from llmtuner.webui.locales import ALERTS
+
+if TYPE_CHECKING:
+    from llmtuner.extras.callbacks import LogCallback
+
+
+def update_process_bar(callback: "LogCallback") -> Dict[str, Any]:
+    if not callback.max_steps:
+        return gr.update(visible=False)
+
+    percentage = round(100 * callback.cur_steps / callback.max_steps, 0) if callback.max_steps != 0 else 100.0
+    label = "Running {:d}/{:d}: {} < {}".format(
+        callback.cur_steps,
+        callback.max_steps,
+        callback.elapsed_time,
+        callback.remaining_time
+    )
+    return gr.update(label=label, value=percentage, visible=True)
+
+
+def get_time() -> str:
+    return datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
+
+
+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]]
+        and os.path.isfile(os.path.join(dataset_dir, dataset_info[dataset[0]]["file_name"]))
+    ):
+        return gr.update(interactive=True)
+    else:
+        return gr.update(interactive=False)
+
+
+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: str = dataset_info[dataset[0]]["file_name"]
+    with open(os.path.join(dataset_dir, data_file), "r", encoding="utf-8") as f:
+        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="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)
+    return "```json\n{}\n```\n".format(result)
+
+
+def gen_plot(base_model: str, finetuning_type: str, output_dir: str) -> matplotlib.figure.Figure:
+    log_file = get_save_dir(base_model, finetuning_type, output_dir, "trainer_log.jsonl")
+    if not os.path.isfile(log_file):
+        return None
+
+    plt.close("all")
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+    steps, losses = [], []
+    with open(log_file, "r", encoding="utf-8") as f:
+        for line in f:
+            log_info = json.loads(line)
+            if log_info.get("loss", None):
+                steps.append(log_info["current_steps"])
+                losses.append(log_info["loss"])
+
+    if len(losses) == 0:
+        return None
+
+    ax.plot(steps, losses, alpha=0.4, label="original")
+    ax.plot(steps, smooth(losses), label="smoothed")
+    ax.legend()
+    ax.set_xlabel("step")
+    ax.set_ylabel("loss")
+    return fig
+
+
+def save_model(
+    lang: str,
+    model_name: str,
+    checkpoints: List[str],
+    finetuning_type: str,
+    template: str,
+    max_shard_size: int,
+    save_dir: str
+) -> Generator[str, None, None]:
+    if not model_name:
+        yield ALERTS["err_no_model"][lang]
+        return
+
+    model_name_or_path = get_model_path(model_name)
+    if not model_name_or_path:
+        yield ALERTS["err_no_path"][lang]
+        return
+
+    if not checkpoints:
+        yield ALERTS["err_no_checkpoint"][lang]
+        return
+
+    checkpoint_dir = ",".join(
+            [get_save_dir(model_name, finetuning_type, ckpt) for ckpt in checkpoints]
+        )
+
+    if not save_dir:
+        yield ALERTS["err_no_save_dir"][lang]
+        return
+
+    args = dict(
+        model_name_or_path=model_name_or_path,
+        checkpoint_dir=checkpoint_dir,
+        finetuning_type=finetuning_type,
+        template=template,
+        output_dir=save_dir
+    )
+
+    yield ALERTS["info_exporting"][lang]
+    export_model(args, max_shard_size="{}GB".format(max_shard_size))
+    yield ALERTS["info_exported"][lang]

src/train_bash.py

@@ -1,17 +1,8 @@
-from llmtuner import get_train_args, run_pt, run_sft, run_rm, run_ppo
+from llmtuner import run_exp
 
 
 def main():
-    model_args, data_args, training_args, finetuning_args, general_args = get_train_args()
-
-    if general_args.stage == "pt":
-        run_pt(model_args, data_args, training_args, finetuning_args)
-    elif general_args.stage == "sft":
-        run_sft(model_args, data_args, training_args, finetuning_args)
-    elif general_args.stage == "rm":
-        run_rm(model_args, data_args, training_args, finetuning_args)
-    elif general_args.stage == "ppo":
-        run_ppo(model_args, data_args, training_args, finetuning_args)
+    run_exp()
 
 
 def _mp_fn(index):

src/train_web.py

@@ -0,0 +1,11 @@
+from llmtuner import create_ui
+
+
+def main():
+    demo = create_ui()
+    demo.queue()
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False, inbrowser=True)
+
+
+if __name__ == "__main__":
+    main()

src/web_demo.py

@@ -1,95 +1,11 @@
-# coding=utf-8
-# Implements user interface in browser for fine-tuned models.
-# Usage: python web_demo.py --model_name_or_path path_to_model --checkpoint_dir path_to_checkpoint
-
-import gradio as gr
-from threading import Thread
-from transformers import TextIteratorStreamer
-from transformers.utils.versions import require_version
-
-from llmtuner import Template, get_infer_args, load_model_and_tokenizer, get_logits_processor
+from llmtuner import create_web_demo
 
 
-require_version("gradio>=3.30.0", "To fix: pip install gradio>=3.30.0")
+def main():
+    demo = create_web_demo()
+    demo.queue()
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False, inbrowser=True)
 
 
-model_args, data_args, finetuning_args, generating_args = get_infer_args()
-model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args)
-
-prompt_template = Template(data_args.prompt_template)
-source_prefix = data_args.source_prefix if data_args.source_prefix else ""
-
-
-def predict(query, chatbot, max_new_tokens, top_p, temperature, history):
-    chatbot.append((query, ""))
-
-    input_ids = tokenizer([prompt_template.get_prompt(query, history, source_prefix)], return_tensors="pt")["input_ids"]
-    input_ids = input_ids.to(model.device)
-
-    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
-
-    gen_kwargs = generating_args.to_dict()
-    gen_kwargs.update({
-        "input_ids": input_ids,
-        "top_p": top_p,
-        "temperature": temperature,
-        "max_new_tokens": max_new_tokens,
-        "logits_processor": get_logits_processor(),
-        "streamer": streamer
-    })
-
-    thread = Thread(target=model.generate, kwargs=gen_kwargs)
-    thread.start()
-
-    response = ""
-    for new_text in streamer:
-        response += new_text
-        new_history = history + [(query, response)]
-        chatbot[-1] = (query, response)
-        yield chatbot, new_history
-
-
-def reset_user_input():
-    return gr.update(value="")
-
-
-def reset_state():
-    return [], []
-
-
-with gr.Blocks() as demo:
-
-    gr.HTML("""
-    <h1 align="center">
-        <a href="https://github.com/hiyouga/LLaMA-Efficient-Tuning" target="_blank">
-            LLaMA Efficient Tuning
-        </a>
-    </h1>
-    """)
-
-    chatbot = gr.Chatbot()
-
-    with gr.Row():
-        with gr.Column(scale=4):
-            with gr.Column(scale=12):
-                user_input = gr.Textbox(show_label=False, placeholder="Input...", lines=10).style(container=False)
-            with gr.Column(min_width=32, scale=1):
-                submitBtn = gr.Button("Submit", variant="primary")
-
-        with gr.Column(scale=1):
-            emptyBtn = gr.Button("Clear History")
-            max_new_tokens = gr.Slider(10, 2048, value=generating_args.max_new_tokens, step=1.0,
-                                       label="Maximum new tokens", interactive=True)
-            top_p = gr.Slider(0.01, 1, value=generating_args.top_p, step=0.01,
-                              label="Top P", interactive=True)
-            temperature = gr.Slider(0.01, 1.5, value=generating_args.temperature, step=0.01,
-                                    label="Temperature", interactive=True)
-
-    history = gr.State([])
-
-    submitBtn.click(predict, [user_input, chatbot, max_new_tokens, top_p, temperature, history], [chatbot, history], show_progress=True)
-    submitBtn.click(reset_user_input, [], [user_input])
-
-    emptyBtn.click(reset_state, outputs=[chatbot, history], show_progress=True)
-
-demo.queue().launch(server_name="0.0.0.0", share=True, inbrowser=True)
+if __name__ == "__main__":
+    main()

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)

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)
@@ -300,10 +301,49 @@ class BaichuanPreTrainedModel(PreTrainedModel):
         if isinstance(module, BaichuanModel):
             module.gradient_checkpointing = value
 
+    @staticmethod
+    def _convert_to_standard_cache(
+        past_key_value: Tuple[Tuple[torch.Tensor, torch.Tensor]], batch_size: int
+    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Standardizes the format of the cache so as to match most implementations, i.e. to tuple(tuple([batch_size,
+        num_heads, ...]))
+        """
+        batch_size_times_num_heads, head_dim, seq_length = past_key_value[0][0].shape
+        num_heads = batch_size_times_num_heads // batch_size
+        # key: [batch_size * num_heads, head_dim, seq_length] -> [batch_size, num_heads, head_dim, seq_length]
+        # value: [batch_size * num_heads, seq_length, head_dim] -> [batch_size, num_heads, seq_length, head_dim]
+        return tuple(
+            (
+                layer_past[0].view(batch_size, num_heads, head_dim, seq_length),
+                layer_past[1].view(batch_size, num_heads, seq_length, head_dim),
+            )
+            for layer_past in past_key_value
+        )
+
+    @staticmethod
+    def _convert_to_baichuan_cache(
+        past_key_value: Tuple[Tuple[torch.Tensor, torch.Tensor]]
+    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Converts the cache to the format expected by Baichuan, i.e. to tuple(tuple([batch_size * num_heads, ...]))
+        """
+        batch_size, num_heads, head_dim, seq_length = past_key_value[0][0].shape
+        batch_size_times_num_heads = batch_size * num_heads
+        # key:  [batch_size, num_heads, head_dim, seq_length] -> [batch_size * num_heads, head_dim, seq_length]
+        # value: [batch_size, num_heads, seq_length, head_dim] -> [batch_size * num_heads, seq_length, head_dim]
+        return tuple(
+            (
+                layer_past[0].view(batch_size_times_num_heads, head_dim, seq_length),
+                layer_past[1].view(batch_size_times_num_heads, seq_length, head_dim),
+            )
+            for layer_past in past_key_value
+        )
+
 
 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
@@ -318,9 +358,9 @@ class BaichuanModel(BaichuanPreTrainedModel):
 
     def get_input_embeddings(self):
         return self.embed_tokens
-        
+
     def set_input_embeddings(self, value):
-        self.embed_tokens = value  
+        self.embed_tokens = value
 
     def build_alibi_tensor(self, attention_mask: torch.Tensor, num_heads: int, dtype: torch.dtype) -> torch.Tensor:
         return build_alibi_tensor(attention_mask, num_heads, dtype)
@@ -468,7 +508,7 @@ class BaichuanModel(BaichuanPreTrainedModel):
             hidden_states=all_hidden_states,
             attentions=all_self_attns,
         )
-    
+
 
 class BaichuanForCausalLM(BaichuanPreTrainedModel):
 
@@ -498,7 +538,7 @@ class BaichuanForCausalLM(BaichuanPreTrainedModel):
 
     def get_decoder(self):
         return self.model
-        
+
     def forward(
         self,
         input_ids: torch.LongTensor = None,
@@ -528,7 +568,7 @@ class BaichuanForCausalLM(BaichuanPreTrainedModel):
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
-        )   
+        )
 
         hidden_states = outputs[0]
         logits = self.lm_head(hidden_states)
@@ -559,11 +599,20 @@ class BaichuanForCausalLM(BaichuanPreTrainedModel):
         )
 
     def prepare_inputs_for_generation(
-        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
-    ):  
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        **kwargs
+    ) -> dict:
         if past_key_values:
             input_ids = input_ids[:, -1:]
 
+            # the cache may be in the standard format (e.g. in contrastive search)
+            if past_key_values[0][0].shape[0] == input_ids.shape[0]:
+                past_key_values = self._convert_to_baichuan_cache(past_key_values)
+
         # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
         if inputs_embeds is not None and past_key_values is None:
             model_inputs = {"inputs_embeds": inputs_embeds}
@@ -571,108 +620,35 @@ class BaichuanForCausalLM(BaichuanPreTrainedModel):
             model_inputs = {"input_ids": input_ids}
 
         model_inputs.update(
-            {   
+            {
                 "past_key_values": past_key_values,
                 "use_cache": kwargs.get("use_cache"),
                 "attention_mask": attention_mask,
-            }   
-        )   
+            }
+        )
         return model_inputs
 
-    @staticmethod
-    def _reorder_cache(past_key_values, beam_idx):
-        return tuple(
-            tuple(past_state.index_select(0, beam_idx) for past_state in layer_past)
-            for layer_past in past_key_values
+    def _reorder_cache(
+        self, past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor
+    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:
+        """
+        This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
+        [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
+        beam_idx at every generation step.
+
+        Output shares the same memory storage as `past`.
+        """
+        standardized_past = self._convert_to_standard_cache(past, batch_size=len(beam_idx))
+
+        # Get a copy of `beam_idx` on all the devices where we need those indices.
+        device_to_beam_idx = {
+            past_state.device: beam_idx.to(past_state.device) for layer_past in past for past_state in layer_past
+        }
+        reordered_past = tuple(
+            (
+                layer_past[0].index_select(0, device_to_beam_idx[layer_past[0].device]),
+                layer_past[1].index_select(0, device_to_beam_idx[layer_past[0].device]),
+            )
+            for layer_past in standardized_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
+        return self._convert_to_baichuan_cache(reordered_past)