From c3580f59835302645e815d0aff0809c72edb14dd Mon Sep 17 00:00:00 2001 From: Younes Belkada <49240599+younesbelkada@users.noreply.github.com> Date: Tue, 19 Mar 2024 12:47:04 +0100 Subject: [PATCH 1/5] Update README.md --- README.md | 92 ++++++++++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 81 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 1ea8b9d86d..7855a3b8d5 100644 --- a/README.md +++ b/README.md @@ -24,30 +24,40 @@ -`trl` is a full stack library where we provide a set of tools to train transformer language models and stable diffusion models with Reinforcement Learning, from the Supervised Fine-tuning step (SFT), Reward Modeling step (RM) to the Proximal Policy Optimization (PPO) step. The library is built on top of the [`transformers`](https://github.com/huggingface/transformers) library by 🤗 Hugging Face. Therefore, pre-trained language models can be directly loaded via `transformers`. At this point, most of decoder architectures and encoder-decoder architectures are supported. Refer to the documentation or the `examples/` folder for example code snippets and how to run these tools. +`trl` is a full stack library where we provide a set of tools to train transformer language models and stable diffusion models with Reinforcement Learning, from the Supervised Fine-tuning step (SFT), Reward Modeling step (RM) to the Proximal Policy Optimization (PPO) step as well as Direct Preference Optimization algorithm (DPO). The library is built on top of the [`transformers`](https://github.com/huggingface/transformers) library by 🤗 Hugging Face. Therefore, pre-trained language models can be directly loaded via `transformers`. At this point, most of decoder architectures and encoder-decoder architectures are supported. Refer to the documentation or the `examples/` folder for example code snippets and how to run these tools. **Highlights:** - [`SFTTrainer`](https://huggingface.co/docs/trl/sft_trainer): A light and friendly wrapper around `transformers` Trainer to easily fine-tune language models or adapters on a custom dataset. - [`RewardTrainer`](https://huggingface.co/docs/trl/reward_trainer): A light wrapper around `transformers` Trainer to easily fine-tune language models for human preferences (Reward Modeling). - [`PPOTrainer`](https://huggingface.co/docs/trl/trainer#trl.PPOTrainer): A PPO trainer for language models that just needs (query, response, reward) triplets to optimise the language model. +- [`DPOTrainer`](https://huggingface.co/docs/trl/trainer#trl.DPOTrainer): A DPO trainer for language models to use [Direct Preference Optimization algorithm.](https://github.com/eric-mitchell/direct-preference-optimization) - [`AutoModelForCausalLMWithValueHead`](https://huggingface.co/docs/trl/models#trl.AutoModelForCausalLMWithValueHead) & [`AutoModelForSeq2SeqLMWithValueHead`](https://huggingface.co/docs/trl/models#trl.AutoModelForSeq2SeqLMWithValueHead): A transformer model with an additional scalar output for each token which can be used as a value function in reinforcement learning. - [Examples](https://github.com/huggingface/trl/tree/main/examples): Train GPT2 to generate positive movie reviews with a BERT sentiment classifier, full RLHF using adapters only, train GPT-j to be less toxic, [Stack-Llama example](https://huggingface.co/blog/stackllama), etc. -## How PPO works -Fine-tuning a language model via PPO consists of roughly three steps: +## Quickstart -1. **Rollout**: The language model generates a response or continuation based on query which could be the start of a sentence. -2. **Evaluation**: The query and response are evaluated with a function, model, human feedback or some combination of them. The important thing is that this process should yield a scalar value for each query/response pair. -3. **Optimization**: This is the most complex part. In the optimisation step the query/response pairs are used to calculate the log-probabilities of the tokens in the sequences. This is done with the model that is trained and a reference model, which is usually the pre-trained model before fine-tuning. The KL-divergence between the two outputs is used as an additional reward signal to make sure the generated responses don't deviate too far from the reference language model. The active language model is then trained with PPO. +You can use TRL Command Line Interface (CLI) to quickly get started with Supervised Fine-tuning (SFT), Direct Preference Optimization (DPO) and test your aligned model with `trl chat`: -This process is illustrated in the sketch below: +### SFT CLI: +```bash +trl sft --help +``` -
- -

Figure: Sketch of the workflow.

-
+### DPO CLI + +```bash +trl dpo --help +``` + +### Chat CLI + +```bash +trl chat --help +``` + +Read more about CLIs on the [relevant documentation section](https://huggingface.co/docs/trl/main/en/clis) for more details. ## Installation @@ -70,6 +80,37 @@ If you wish to develop TRL, you should install in editable mode: pip install -e . ``` +## Implementation details + +### How PPO works +Fine-tuning a language model via PPO consists of roughly three steps: + +1. **Rollout**: The language model generates a response or continuation based on query which could be the start of a sentence. +2. **Evaluation**: The query and response are evaluated with a function, model, human feedback or some combination of them. The important thing is that this process should yield a scalar value for each query/response pair. +3. **Optimization**: This is the most complex part. In the optimisation step the query/response pairs are used to calculate the log-probabilities of the tokens in the sequences. This is done with the model that is trained and a reference model, which is usually the pre-trained model before fine-tuning. The KL-divergence between the two outputs is used as an additional reward signal to make sure the generated responses don't deviate too far from the reference language model. The active language model is then trained with PPO. + +This process is illustrated in the sketch below: + +
+ +

Figure: Sketch of the workflow.

+
+ +### How DPO works + +Fine-tuning a language model via DPO consists of two steps and is easier than PPO: + +1. **Data collection**: Gather a preference dataset with positive and negative selected pairs of generation, given a prompt. +2. **Optimization**: Maximize the log-likelihood of the DPO loss directly. + +DPO-compatible datasets can be found with [the tag `dpo` on Hugging Face Hub](https://huggingface.co/datasets?other=dpo). + +This process is illustrated in the sketch below (from [figure 1 of the original paper](https://arxiv.org/pdf/2305.18290.pdf)): + +Screenshot 2024-03-19 at 12 39 41 + +Read more about DPO algorithm in the original paper. + ## How to use ### `SFTTrainer` @@ -161,11 +202,40 @@ reward = [torch.tensor(1.0)] train_stats = ppo_trainer.step([query_tensor[0]], [response_tensor[0]], reward) ``` +### `DPOTrainer` + +`DPOTrainer` is a trainer that uses [Direct Preference Optimization algorithm](https://arxiv.org/abs/2305.18290). This is a basic example on how to use the `DPOTrainer` from the library. The `DPOTrainer` is a wrapper around the `transformers` Trainer to easily fine-tune reward models or adapters on a custom preference dataset. + +```python +# imports +from transformers import AutoModelForCausalLM, AutoTokenizer +from trl import RewardTrainer + +# load model and dataset - dataset needs to be in a specific format +model = AutoModelForCausalLM.from_pretrained("gpt2") +tokenizer = AutoTokenizer.from_pretrained("gpt2") + +... + +# load trainer +trainer = DPOTrainer( + model=model, + tokenizer=tokenizer, + train_dataset=dataset, +) + +# train +trainer.train() +``` + ## References ### Proximal Policy Optimisation The PPO implementation largely follows the structure introduced in the paper **"Fine-Tuning Language Models from Human Preferences"** by D. Ziegler et al. \[[paper](https://arxiv.org/pdf/1909.08593.pdf), [code](https://github.com/openai/lm-human-preferences)]. +### Direct Preference Optimization +DPO original code and paper can be found here: https://github.com/eric-mitchell/direct-preference-optimization + ### Language models The language models utilize the `transformers` library by 🤗 Hugging Face. Currently, `trl` only supports `transformers` models **for text**. From 429758eba1e25da0138474234f5c816c96c1ee74 Mon Sep 17 00:00:00 2001 From: Younes Belkada <49240599+younesbelkada@users.noreply.github.com> Date: Tue, 19 Mar 2024 13:33:54 +0100 Subject: [PATCH 2/5] Update README.md --- README.md | 4 ---- 1 file changed, 4 deletions(-) diff --git a/README.md b/README.md index 7855a3b8d5..59f8b0867f 100644 --- a/README.md +++ b/README.md @@ -20,10 +20,6 @@ ## What is it? -
- -
- `trl` is a full stack library where we provide a set of tools to train transformer language models and stable diffusion models with Reinforcement Learning, from the Supervised Fine-tuning step (SFT), Reward Modeling step (RM) to the Proximal Policy Optimization (PPO) step as well as Direct Preference Optimization algorithm (DPO). The library is built on top of the [`transformers`](https://github.com/huggingface/transformers) library by 🤗 Hugging Face. Therefore, pre-trained language models can be directly loaded via `transformers`. At this point, most of decoder architectures and encoder-decoder architectures are supported. Refer to the documentation or the `examples/` folder for example code snippets and how to run these tools. **Highlights:** From b95e88fe28afa89a0f7d50376daf3e7f06c18d12 Mon Sep 17 00:00:00 2001 From: leandro Date: Tue, 19 Mar 2024 16:23:42 +0100 Subject: [PATCH 3/5] move dpo/ppo description to docs --- docs/source/dpo_trainer.mdx | 17 ++++++++++++++++- docs/source/ppo_trainer.mdx | 15 +++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/docs/source/dpo_trainer.mdx b/docs/source/dpo_trainer.mdx index c929885348..8121cf5b33 100644 --- a/docs/source/dpo_trainer.mdx +++ b/docs/source/dpo_trainer.mdx @@ -2,9 +2,24 @@ TRL supports the DPO Trainer for training language models from preference data, as described in the paper [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://arxiv.org/abs/2305.18290) by Rafailov et al., 2023. For a full example have a look at [`examples/scripts/dpo.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/dpo.py). - The first step as always is to train your SFT model, to ensure the data we train on is in-distribution for the DPO algorithm. +## How DPO works + +Fine-tuning a language model via DPO consists of two steps and is easier than PPO: + +1. **Data collection**: Gather a preference dataset with positive and negative selected pairs of generation, given a prompt. +2. **Optimization**: Maximize the log-likelihood of the DPO loss directly. + +DPO-compatible datasets can be found with [the tag `dpo` on Hugging Face Hub](https://huggingface.co/datasets?other=dpo). + +This process is illustrated in the sketch below (from [figure 1 of the original paper](https://arxiv.org/pdf/2305.18290.pdf)): + +Screenshot 2024-03-19 at 12 39 41 + +Read more about DPO algorithm in the [original paper](https://arxiv.org/pdf/2305.18290.pdf). + + ## Expected dataset format The DPO trainer expects a very specific format for the dataset. Since the model will be trained to directly optimize the preference of which sentence is the most relevant, given two sentences. We provide an example from the [`Anthropic/hh-rlhf`](https://huggingface.co/datasets/Anthropic/hh-rlhf) dataset below: diff --git a/docs/source/ppo_trainer.mdx b/docs/source/ppo_trainer.mdx index 5de79e9491..e06a73d1b4 100644 --- a/docs/source/ppo_trainer.mdx +++ b/docs/source/ppo_trainer.mdx @@ -4,6 +4,21 @@ TRL supports the [PPO](https://arxiv.org/abs/1707.06347) Trainer for training la The first step is to train your SFT model (see the [SFTTrainer](sft_trainer)), to ensure the data we train on is in-distribution for the PPO algorithm. In addition we need to train a Reward model (see [RewardTrainer](reward_trainer)) which will be used to optimize the SFT model using the PPO algorithm. +## How PPO works + +Fine-tuning a language model via PPO consists of roughly three steps: + +1. **Rollout**: The language model generates a response or continuation based on query which could be the start of a sentence. +2. **Evaluation**: The query and response are evaluated with a function, model, human feedback or some combination of them. The important thing is that this process should yield a scalar value for each query/response pair. +3. **Optimization**: This is the most complex part. In the optimisation step the query/response pairs are used to calculate the log-probabilities of the tokens in the sequences. This is done with the model that is trained and a reference model, which is usually the pre-trained model before fine-tuning. The KL-divergence between the two outputs is used as an additional reward signal to make sure the generated responses don't deviate too far from the reference language model. The active language model is then trained with PPO. + +This process is illustrated in the sketch below: + +
+ +

Figure: Sketch of the workflow.

+
+ ## Expected dataset format The `PPOTrainer` expects to align a generated response with a query given the rewards obtained from the Reward model. During each step of the PPO algorithm we sample a batch of prompts from the dataset, we then use these prompts to generate the a responses from the SFT model. Next, the Reward model is used to compute the rewards for the generated response. Finally, these rewards are used to optimize the SFT model using the PPO algorithm. From 0421325805c2eb63268b128d40d0e15c943f4e69 Mon Sep 17 00:00:00 2001 From: leandro Date: Tue, 19 Mar 2024 16:24:23 +0100 Subject: [PATCH 4/5] rework readme --- README.md | 109 +++++++++++++++++++++++------------------------------- 1 file changed, 46 insertions(+), 63 deletions(-) diff --git a/README.md b/README.md index 59f8b0867f..7d32018944 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@ # TRL - Transformer Reinforcement Learning -> Full stack transformer language models with reinforcement learning. +> Full stack library to fine-tune and align large language models.

@@ -20,95 +20,70 @@ ## What is it? -`trl` is a full stack library where we provide a set of tools to train transformer language models and stable diffusion models with Reinforcement Learning, from the Supervised Fine-tuning step (SFT), Reward Modeling step (RM) to the Proximal Policy Optimization (PPO) step as well as Direct Preference Optimization algorithm (DPO). The library is built on top of the [`transformers`](https://github.com/huggingface/transformers) library by 🤗 Hugging Face. Therefore, pre-trained language models can be directly loaded via `transformers`. At this point, most of decoder architectures and encoder-decoder architectures are supported. Refer to the documentation or the `examples/` folder for example code snippets and how to run these tools. +The `trl` library is a full stack tool to fine-tune and align transformer language and diffusion models using methods such as Supervised Fine-tuning step (SFT), Reward Modeling (RM) and the Proximal Policy Optimization (PPO) as well as Direct Preference Optimization (DPO). -**Highlights:** +The library is built on top of the [`transformers`](https://github.com/huggingface/transformers) library and thus allows to use any model architecture available there. -- [`SFTTrainer`](https://huggingface.co/docs/trl/sft_trainer): A light and friendly wrapper around `transformers` Trainer to easily fine-tune language models or adapters on a custom dataset. -- [`RewardTrainer`](https://huggingface.co/docs/trl/reward_trainer): A light wrapper around `transformers` Trainer to easily fine-tune language models for human preferences (Reward Modeling). -- [`PPOTrainer`](https://huggingface.co/docs/trl/trainer#trl.PPOTrainer): A PPO trainer for language models that just needs (query, response, reward) triplets to optimise the language model. -- [`DPOTrainer`](https://huggingface.co/docs/trl/trainer#trl.DPOTrainer): A DPO trainer for language models to use [Direct Preference Optimization algorithm.](https://github.com/eric-mitchell/direct-preference-optimization) -- [`AutoModelForCausalLMWithValueHead`](https://huggingface.co/docs/trl/models#trl.AutoModelForCausalLMWithValueHead) & [`AutoModelForSeq2SeqLMWithValueHead`](https://huggingface.co/docs/trl/models#trl.AutoModelForSeq2SeqLMWithValueHead): A transformer model with an additional scalar output for each token which can be used as a value function in reinforcement learning. -- [Examples](https://github.com/huggingface/trl/tree/main/examples): Train GPT2 to generate positive movie reviews with a BERT sentiment classifier, full RLHF using adapters only, train GPT-j to be less toxic, [Stack-Llama example](https://huggingface.co/blog/stackllama), etc. -## Quickstart +## Highlights -You can use TRL Command Line Interface (CLI) to quickly get started with Supervised Fine-tuning (SFT), Direct Preference Optimization (DPO) and test your aligned model with `trl chat`: - -### SFT CLI: - -```bash -trl sft --help -``` - -### DPO CLI - -```bash -trl dpo --help -``` - -### Chat CLI - -```bash -trl chat --help -``` - -Read more about CLIs on the [relevant documentation section](https://huggingface.co/docs/trl/main/en/clis) for more details. +- **`Efficient and scalable`**: + - [`accelerate`](https://github.com/huggingface/accelerate) is the backbone of `trl` which allows to scale model training from a single GPU to a large scale multi-node cluster with methods such as DDP and DeepSpeed. + - [`PEFT`](https://github.com/huggingface/peft) is fully integrated and allows to train even the largest models on modest hardware with quantisation and methods such as LoRA or QLoRA. + - [`unsloth`](https://github.com/unslothai/unsloth) is also integrated and allows to significantly speed up training with dedicated kernels. +- **`CLI`**: With the [CLI](https://huggingface.co/docs/trl/clis) you can fine-tune and chat with LLMs without writing any code using a single command and a flexible config system. +- **`Trainers`**: The Trainer classes are an abstraction to apply many fine-tuning methods with ease such as the [`SFTTrainer`](https://huggingface.co/docs/trl/sft_trainer), [`DPOTrainer`](https://huggingface.co/docs/trl/trainer#trl.DPOTrainer), [`RewardTrainer`](https://huggingface.co/docs/trl/reward_trainer), and [`PPOTrainer`](https://huggingface.co/docs/trl/trainer#trl.PPOTrainer). +- **`AutoModels`**: The [`AutoModelForCausalLMWithValueHead`](https://huggingface.co/docs/trl/models#trl.AutoModelForCausalLMWithValueHead) & [`AutoModelForSeq2SeqLMWithValueHead`](https://huggingface.co/docs/trl/models#trl.AutoModelForSeq2SeqLMWithValueHead) classes add an additional value head to the model which allows to train them with RL algorithms such as PPO. +- **`Examples`**: Train GPT2 to generate positive movie reviews with a BERT sentiment classifier, full RLHF using adapters only, train GPT-j to be less toxic, [StackLlama example](https://huggingface.co/blog/stackllama), etc. following the [examples](https://github.com/huggingface/trl/tree/main/examples). ## Installation ### Python package -Install the library with pip: +Install the library with `pip`: ```bash pip install trl ``` ### From source -If you want to run the examples in the repository a few additional libraries are required. Clone the repository and install it with pip: +If you want to use the latest features before an official release you can install from source: ```bash -git clone https://github.com/huggingface/trl.git -cd trl/ -pip install . +pip install git+https://github.com/huggingface/trl.git ``` -If you wish to develop TRL, you should install in editable mode: +### Repository +If you want to use the examples you can clone the repository with the following command: ```bash -pip install -e . +git clone https://github.com/huggingface/trl.git ``` -## Implementation details +## Command Line Interface (CLI) -### How PPO works -Fine-tuning a language model via PPO consists of roughly three steps: +You can use TRL Command Line Interface (CLI) to quickly get started with Supervised Fine-tuning (SFT), Direct Preference Optimization (DPO) and test your aligned model with the chat CLI: -1. **Rollout**: The language model generates a response or continuation based on query which could be the start of a sentence. -2. **Evaluation**: The query and response are evaluated with a function, model, human feedback or some combination of them. The important thing is that this process should yield a scalar value for each query/response pair. -3. **Optimization**: This is the most complex part. In the optimisation step the query/response pairs are used to calculate the log-probabilities of the tokens in the sequences. This is done with the model that is trained and a reference model, which is usually the pre-trained model before fine-tuning. The KL-divergence between the two outputs is used as an additional reward signal to make sure the generated responses don't deviate too far from the reference language model. The active language model is then trained with PPO. +**SFT:** -This process is illustrated in the sketch below: - -

- -

Figure: Sketch of the workflow.

-
- -### How DPO works - -Fine-tuning a language model via DPO consists of two steps and is easier than PPO: +```bash +trl sft --model_name_or_path facebook/opt-125m --dataset_name imdb --output_dir opt-sft-imdb +``` -1. **Data collection**: Gather a preference dataset with positive and negative selected pairs of generation, given a prompt. -2. **Optimization**: Maximize the log-likelihood of the DPO loss directly. +**DPO:** -DPO-compatible datasets can be found with [the tag `dpo` on Hugging Face Hub](https://huggingface.co/datasets?other=dpo). +```bash +trl dpo --model_name_or_path facebook/opt-125m --dataset_name trl-internal-testing/Anthropic-hh-rlhf-processed --output_dir opt-sft-hh-rlhf +``` -This process is illustrated in the sketch below (from [figure 1 of the original paper](https://arxiv.org/pdf/2305.18290.pdf)): +**Chat:** -Screenshot 2024-03-19 at 12 39 41 +```bash +trl chat --model_name_or_path Qwen/Qwen1.5-0.5B-Chat +``` -Read more about DPO algorithm in the original paper. +Read more about CLI in the [relevant documentation section](https://huggingface.co/docs/trl/main/en/clis) or use `--help` for more details. ## How to use +For more flexibility and control over the training you can use the dedicated trainer classes to fine-tune the model in Python. + ### `SFTTrainer` This is a basic example on how to use the `SFTTrainer` from the library. The `SFTTrainer` is a light wrapper around the `transformers` Trainer to easily fine-tune language models or adapters on a custom dataset. @@ -224,16 +199,24 @@ trainer = DPOTrainer( trainer.train() ``` +## Development + +If you want to contribute to `trl` or customizing it to your needs make sure to read the [contribution guide](https://github.com/huggingface/trl/blob/main/CONTRIBUTING.md) and make sure you make a dev install: + +```bash +git clone https://github.com/huggingface/trl.git +cd trl/ +make dev +``` + ## References ### Proximal Policy Optimisation The PPO implementation largely follows the structure introduced in the paper **"Fine-Tuning Language Models from Human Preferences"** by D. Ziegler et al. \[[paper](https://arxiv.org/pdf/1909.08593.pdf), [code](https://github.com/openai/lm-human-preferences)]. ### Direct Preference Optimization -DPO original code and paper can be found here: https://github.com/eric-mitchell/direct-preference-optimization +DPO is based on the original implementation of **"Direct Preference Optimization: Your Language Model is Secretly a Reward Model"** by E. Mitchell et al. \[[paper](), [code](https://github.com/eric-mitchell/direct-preference-optimization)] -### Language models -The language models utilize the `transformers` library by 🤗 Hugging Face. Currently, `trl` only supports `transformers` models **for text**. ## Citation From f44c05835520131ac716e909967b9340594434ea Mon Sep 17 00:00:00 2001 From: Younes Belkada <49240599+younesbelkada@users.noreply.github.com> Date: Tue, 19 Mar 2024 16:27:13 +0100 Subject: [PATCH 5/5] Update README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 7d32018944..3f55577cc2 100644 --- a/README.md +++ b/README.md @@ -180,7 +180,7 @@ train_stats = ppo_trainer.step([query_tensor[0]], [response_tensor[0]], reward) ```python # imports from transformers import AutoModelForCausalLM, AutoTokenizer -from trl import RewardTrainer +from trl import DPOTrainer # load model and dataset - dataset needs to be in a specific format model = AutoModelForCausalLM.from_pretrained("gpt2")