# collections

id: b0d61235-9c1a-4b51-b8f3-dd66e1fe516e
name: chroma-db-trl
database_id: 00000000-0000-0000-0000-000000000000
config_json_str:
{"hnsw_configuration": {"space": "l2", "ef_construction": 100, "ef_search": 10, "num_threads": 8, "M": 16, "resize_factor": 1.2, "batch_size": 100, "sync_threshold": 1000, "_type": "HNSWConfigurationInternal"}, "_type": "CollectionConfigurationInternal"}


# embedding_metadata

key: _node_content
string_value:
{"id_": "2aae18be-3c83-40ce-a8d9-519140821e86", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/clis", "title": "Command Line Interfaces (CLIs)", "tokens": 1209, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "b83b6555-15cb-4b57-8077-36c3568166a4", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/clis", "title": "Command Line Interfaces (CLIs)", "tokens": 1209, "retrieve_doc": true, "source": "trl"}, "hash": "3c7b62a48b18c3b8de67f9221dc33fa6352bc10f7d42d3a454483bcde4049576", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "d66c8401-3e52-48a7-a817-89883443903f", "node_type": "1", "metadata": {}, "hash": "112597e2d60298bd753e703a67df805412248ae7c8fa48ceabe2d7a64a647f7a", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 3056, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# Command Line Interfaces (CLIs)

You can use TRL to fine-tune your Language Model with Supervised Fine-Tuning (SFT) or Direct Policy Optimization (DPO) or even chat with your model using the TRL CLIs.

Currently supported CLIs are:

- `trl sft`: fine-tune a LLM on a text/instruction dataset
- `trl dpo`: fine-tune a LLM with DPO on a preference dataset 
- `trl chat`: quickly spin up a LLM fine-tuned for chatting

## Fine-tuning with the CLI

Before getting started, pick up a Language Model from Hugging Face Hub. Supported models can be found with the filter "text-generation" within models. Also make sure to pick up a relevant dataset for your task.

Before using the `sft` or `dpo` commands make sure to run:
```bash
accelerate config
```
and pick up the right configuration for your training setup (single / multi-GPU, DeepSpeed, etc.). Make sure to complete all steps of `accelerate config` before running any CLI command.

We also recommend you passing a YAML config file to configure your training protocol. Below is a simple example of a YAML file that you can use for training your models with `trl sft` command.

```yaml
model_name_or_path:
  trl-internal-testing/tiny-random-LlamaForCausalLM
dataset_name:
  imdb
dataset_text_field:
  text
report_to:
  none
learning_rate:
  0.0001
lr_scheduler_type:
  cosine
```

Save that config in a `.yaml` and get started immediately! An example CLI config is available as `examples/cli_configs/example_config.yaml`. Note you can overwrite the arguments from the config file by explicitly passing them to the CLI, e.g. from the root folder:

```bash
trl sft --config examples/cli_configs/example_config.yaml --output_dir test-trl-cli --lr_scheduler_type cosine_with_restarts
```

Will force-use `cosine_with_restarts` for `lr_scheduler_type`.

### Supported Arguments 

We do support all arguments from `transformers.TrainingArguments`, for loading your model, we support all arguments from `~trl.ModelConfig`:

[[autodoc]] ModelConfig

You can pass any of these arguments either to the CLI or the YAML file.

### Supervised Fine-tuning (SFT)

Follow the basic instructions above and run `trl sft --output_dir <output_dir> <*args>`: 

```bash
trl sft --model_name_or_path facebook/opt-125m --dataset_name imdb --output_dir opt-sft-imdb
```

The SFT CLI is based on the `examples/scripts/sft.py` script.

### Direct Policy Optimization (DPO)

To use the DPO CLI, you need to have a dataset in the TRL format such as 

* TRL's Anthropic HH dataset: https://huggingface.co/datasets/trl-internal-testing/hh-rlhf-helpful-base-trl-style
* TRL's OpenAI TL;DR summarization dataset: https://huggingface.co/datasets/trl-internal-testing/tldr-preference-trl-style

These datasets always have at least three columns `prompt, chosen, rejected`:

* `prompt` is a list of strings.
* `chosen` is the chosen response in [chat format](https://huggingface.co/docs/transformers/main/en/chat_templating)
* `rejected` is the rejected response [chat format](https://huggingface.co/docs/transformers/main/en/chat_templating)


key: _node_content
string_value:
{"id_": "d66c8401-3e52-48a7-a817-89883443903f", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/clis", "title": "Command Line Interfaces (CLIs)", "tokens": 1209, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "b83b6555-15cb-4b57-8077-36c3568166a4", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/clis", "title": "Command Line Interfaces (CLIs)", "tokens": 1209, "retrieve_doc": true, "source": "trl"}, "hash": "3c7b62a48b18c3b8de67f9221dc33fa6352bc10f7d42d3a454483bcde4049576", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "2aae18be-3c83-40ce-a8d9-519140821e86", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/clis", "title": "Command Line Interfaces (CLIs)", "tokens": 1209, "retrieve_doc": true, "source": "trl"}, "hash": "4c076177e2125357a7fd9022e10cb05ca7b3c2a95beea33c83aaed00a75394b7", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 3060, "end_char_idx": 4924, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
To do a quick start, you can run the following command:

```bash
trl dpo --model_name_or_path facebook/opt-125m --output_dir trl-hh-rlhf --dataset_name trl-internal-testing/hh-rlhf-helpful-base-trl-style
```


The DPO CLI is based on the `examples/scripts/dpo.py` script.


#### Custom preference dataset

Format the dataset into TRL format (you can adapt the `examples/datasets/anthropic_hh.py`):

```bash
python examples/datasets/anthropic_hh.py --push_to_hub --hf_entity your-hf-org
```

## Chat interface

The chat CLI lets you quickly load the model and talk to it. Simply run the following:

```bash
trl chat --model_name_or_path  Qwen/Qwen1.5-0.5B-Chat 
```

> [!TIP]
> To use the chat CLI with the developer installation, you must run `make dev` 
>

Note that the chat interface relies on the tokenizer's [chat template](https://huggingface.co/docs/transformers/chat_templating) to format the inputs for the model. Make sure your tokenizer has a chat template defined.

Besides talking to the model there are a few commands you can use:

- **clear**: clears the current conversation and start a new one
- **example {NAME}**: load example named `{NAME}` from the config and use it as the user input
- **set {SETTING_NAME}={SETTING_VALUE};**: change the system prompt or generation settings (multiple settings are separated by a ';').
- **reset**: same as clear but also resets the generation configs to defaults if they have been changed by **set**
- **save {SAVE_NAME} (optional)**: save the current chat and settings to file by default to `./chat_history/{MODEL_NAME}/chat_{DATETIME}.yaml` or `{SAVE_NAME}` if provided
- **exit**: closes the interface

The default examples are defined in `examples/scripts/config/default_chat_config.yaml` but you can pass your own with `--config CONFIG_FILE` where you can also specify the default generation parameters.


key: _node_content
string_value:
{"id_": "227192cc-5be8-44e4-a048-71bc5e53b1e4", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae4d4218-2aa9-413b-bbe4-2665b89b5998", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "3a728c5c956951b40498782483a1a6fc9c27dfe14a50ee062e3eb965e3c60518", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "d10eb15b-17af-436d-b591-c925fce94950", "node_type": "1", "metadata": {}, "hash": "d3caf8f18dca2e3da6fc656dcffb09e29a1d8c51c9c11afec139d4067311bc6d", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 2811, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# Online DPO Trainer

TRL supports training LLMs with online DPO ([Guo et al., 2024](https://huggingface.co/papers/2402.04792)) with a reward model (RM). The idea of online DPO is to generate completions based on prompts and either have a reward model or an LLM judge to rank the responses as chosen or rejected. Then the model is updated with the ranked responses using the DPO loss.

While [Guo et al. (2024)](https://huggingface.co/papers/2402.04792) used an LLM judge to score model completions, the current implementation only supports reward models -- see [Reward Bench](https://huggingface.co/spaces/allenai/reward-bench) for a leaderboard of public models you can use.

## Get started

The basic API looks as follows:

```python
from datasets import Dataset
from trl import OnlineDPOConfig, OnlineDPOTrainer
from transformers import (
    AutoModelForCausalLM,
    AutoModelForSequenceClassification,
    AutoTokenizer,
)
NUM_DUMMY_SAMPLES = 100
tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM-135M-Instruct")
tok.add_special_tokens({"pad_token": "[PAD]"})
# The model to optimise
model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM-135M-Instruct")
# The reference model to calculate the KL divergence against
ref_model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM-135M-Instruct")
# The model to score completions with. In practice, you will need a fine-tuned reward model.
reward_model = AutoModelForSequenceClassification.from_pretrained("HuggingFaceTB/SmolLM-135M-Instruct", num_labels=1)
train_dataset = Dataset.from_dict(
    {"input_ids": [tok.encode("Q: Hi how are you? A:")] * NUM_DUMMY_SAMPLES})
eval_dataset = Dataset.from_dict(
    {"input_ids": [tok.encode("Q: What do you like to eat A:")] * NUM_DUMMY_SAMPLES})
trainer = OnlineDPOTrainer(
    OnlineDPOConfig(
        output_dir="online-dpo-model",
    ),
    model=model,
    ref_model=ref_model,
    reward_model=reward_model,
    tokenizer=tok,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)
trainer.train()
```

To run the online DPO script with a dummy reward model, run:

```bash
python examples/scripts/online_dpo.py \
    --dataset_name trl-lib/tldr \
    --learning_rate 3e-6 \
    --output_dir models/minimal/online_dpo \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 64 \
    --total_episodes 30000 \
    --model_name_or_path EleutherAI/pythia-14m \
    --sft_model_path EleutherAI/pythia-14m \
    --reward_model_path EleutherAI/pythia-14m \
    --non_eos_penalty \
    --stop_token eos \
    --response_length 53 \
    --sanity_check
```

## Expected dataset format

Unlike standard DPO where one provides a dataset with chosen and rejected columns, for online DPO one just needs a dataset of prompts to generate the completions from.


key: _node_content
string_value:
{"id_": "d10eb15b-17af-436d-b591-c925fce94950", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae4d4218-2aa9-413b-bbe4-2665b89b5998", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "3a728c5c956951b40498782483a1a6fc9c27dfe14a50ee062e3eb965e3c60518", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "227192cc-5be8-44e4-a048-71bc5e53b1e4", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "7d3ba5402064537c148b28de94d9838fd4c5e05197a7202d1758635065c90f0d", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "b4e898a9-b92a-4697-9730-ccbaa5c361cd", "node_type": "1", "metadata": {}, "hash": "38689ed1b054a9b0e970f7261aafcce3841a7cad2d113d0caed7c5f7ba172008", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 1375, "end_char_idx": 4546, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
In practice, you will need a fine-tuned reward model.
reward_model = AutoModelForSequenceClassification.from_pretrained("HuggingFaceTB/SmolLM-135M-Instruct", num_labels=1)
train_dataset = Dataset.from_dict(
    {"input_ids": [tok.encode("Q: Hi how are you? A:")] * NUM_DUMMY_SAMPLES})
eval_dataset = Dataset.from_dict(
    {"input_ids": [tok.encode("Q: What do you like to eat A:")] * NUM_DUMMY_SAMPLES})
trainer = OnlineDPOTrainer(
    OnlineDPOConfig(
        output_dir="online-dpo-model",
    ),
    model=model,
    ref_model=ref_model,
    reward_model=reward_model,
    tokenizer=tok,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)
trainer.train()
```

To run the online DPO script with a dummy reward model, run:

```bash
python examples/scripts/online_dpo.py \
    --dataset_name trl-lib/tldr \
    --learning_rate 3e-6 \
    --output_dir models/minimal/online_dpo \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 64 \
    --total_episodes 30000 \
    --model_name_or_path EleutherAI/pythia-14m \
    --sft_model_path EleutherAI/pythia-14m \
    --reward_model_path EleutherAI/pythia-14m \
    --non_eos_penalty \
    --stop_token eos \
    --response_length 53 \
    --sanity_check
```

## Expected dataset format

Unlike standard DPO where one provides a dataset with chosen and rejected columns, for online DPO one just needs a dataset of prompts to generate the completions from. The [`OnlineDPOTrainer`] assumes that the dataset is preprocessed for model inference, so typically you will want to wrap your prompts in the messages format and then apply the chat template as follows:

```python
def prepare_dataset(dataset, tokenizer, dataset_prompt_field):
    """pre-tokenize the dataset before training; only collate during training"""
    return dataset.map(
        lambda x: {"input_ids": tokenizer.apply_chat_template(x[dataset_prompt_field], add_generation_prompt=True)},
        remove_columns=dataset.column_names,
    )

dataset = prepare_dataset(dataset)
```

## Explanation of the logged metrics

The logged metrics are as follows. Here is an example [tracked run at Weights and Biases](https://wandb.ai/huggingface/trl/runs/dd2o3g35)

* `eps`: Tracks the number of episodes per second.
* `objective/kl`: The mean Kullback-Leibler (KL) divergence between the current model and reference model.
* `objective/entropy`: The mean entropy of the model, indicating the randomness of the actions chosen by the model.
* `objective/non_score_reward`: The mean reward from non-score-related sources, basically `beta * kl.sum(1)`, where `beta` is the KL penalty coefficient and `kl` is the per-token KL divergence.
* `objective/rlhf_reward`: The mean RLHF reward, which is `score - non_score_reward`.
* `objective/scores`: The mean scores returned by the reward model / environment.
* `objective/scores_margin`: The mean score margin (according to the external reward model) between the chosen and rejected completions.
* `rewards/accuracies`: The accuracies of the online DPO's implicit reward model.
* `rewards/chosen`: The mean reward (according to online DPO's implicit reward model)of the chosen completions.


key: _node_content
string_value:
{"id_": "b4e898a9-b92a-4697-9730-ccbaa5c361cd", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae4d4218-2aa9-413b-bbe4-2665b89b5998", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "3a728c5c956951b40498782483a1a6fc9c27dfe14a50ee062e3eb965e3c60518", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "d10eb15b-17af-436d-b591-c925fce94950", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "3acbbbea5611923e7cf5d87042fb083409f9d3bcef70ec7fe3dc4aef1dd52525", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "3a6a7e25-53a2-4c0c-aa4f-0cdfff914803", "node_type": "1", "metadata": {}, "hash": "b29224f9019a23e1199a1fb59a586f2fccef67847212f26d9933970924d9acb5", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 2812, "end_char_idx": 6244, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
The [`OnlineDPOTrainer`] assumes that the dataset is preprocessed for model inference, so typically you will want to wrap your prompts in the messages format and then apply the chat template as follows:

```python
def prepare_dataset(dataset, tokenizer, dataset_prompt_field):
    """pre-tokenize the dataset before training; only collate during training"""
    return dataset.map(
        lambda x: {"input_ids": tokenizer.apply_chat_template(x[dataset_prompt_field], add_generation_prompt=True)},
        remove_columns=dataset.column_names,
    )

dataset = prepare_dataset(dataset)
```

## Explanation of the logged metrics

The logged metrics are as follows. Here is an example [tracked run at Weights and Biases](https://wandb.ai/huggingface/trl/runs/dd2o3g35)

* `eps`: Tracks the number of episodes per second.
* `objective/kl`: The mean Kullback-Leibler (KL) divergence between the current model and reference model.
* `objective/entropy`: The mean entropy of the model, indicating the randomness of the actions chosen by the model.
* `objective/non_score_reward`: The mean reward from non-score-related sources, basically `beta * kl.sum(1)`, where `beta` is the KL penalty coefficient and `kl` is the per-token KL divergence.
* `objective/rlhf_reward`: The mean RLHF reward, which is `score - non_score_reward`.
* `objective/scores`: The mean scores returned by the reward model / environment.
* `objective/scores_margin`: The mean score margin (according to the external reward model) between the chosen and rejected completions.
* `rewards/accuracies`: The accuracies of the online DPO's implicit reward model.
* `rewards/chosen`: The mean reward (according to online DPO's implicit reward model)of the chosen completions.
* `rewards/rejected`: The mean reward (according to online DPO's implicit reward model) of the rejected completions.
* `rewards/margins`: The mean reward margin (according to online DPO's implicit reward model) between the chosen and rejected completions.
* `logps/chosen`: The mean log probabilities of the chosen completions.
* `logps/rejected`: The mean log probabilities of the rejected completions.
* `val/num_eos_tokens`: The number of end-of-sequence (EOS) tokens generated, which can indicate the number of complete responses.
* `lr`: lr: The current learning rate used by the optimizer.
* `episode`: episode: The current global step or episode count in the training process.


## Cookbook

> [!IMPORTANT]
> Make sure the SFT model and reward model use the _same_ chat template. Otherwise you may find the model completions are scored incorrectly.


* Debugging TIP: `objective/rlhf_reward`: this is the ultimate objective of the RLHF training. If training works as intended, this metric should keep going up.
* Memory TIP: If you are running out of memory, you can try to reduce the `--per_device_train_batch_size` or increase the `--gradient_accumulation_steps` to reduce the memory footprint.
* Memory TIP: If you have multiple GPUs, you can also run training with DeepSpeed stage 3 to reduce the memory footprint `accelerate launch --config_file examples/accelerate_configs/deepspeed_zero3.yaml`.
* Usage TIP: We recommend to use the "EOS trick" via `--non_eos_penalty --stop_token eos`, which replaces the score of completions that do not end with an EOS token with a static scalar penalty `--penalty_reward_value`. This can help the model learn to generate more coherent completions.


key: _node_content
string_value:
{"id_": "3a6a7e25-53a2-4c0c-aa4f-0cdfff914803", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae4d4218-2aa9-413b-bbe4-2665b89b5998", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "3a728c5c956951b40498782483a1a6fc9c27dfe14a50ee062e3eb965e3c60518", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "b4e898a9-b92a-4697-9730-ccbaa5c361cd", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "bdcd68cf7c2de2c8e57c752e3263361301bd519e69146519d510fcd381afc58d", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "68fa7d05-b2a0-4295-9dbf-1abca1b5e99f", "node_type": "1", "metadata": {}, "hash": "76840db126c61286971ede7b1eafe76ebd6bba9010baa21d2cfd212e151519e8", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 4547, "end_char_idx": 7424, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
* `rewards/rejected`: The mean reward (according to online DPO's implicit reward model) of the rejected completions.
* `rewards/margins`: The mean reward margin (according to online DPO's implicit reward model) between the chosen and rejected completions.
* `logps/chosen`: The mean log probabilities of the chosen completions.
* `logps/rejected`: The mean log probabilities of the rejected completions.
* `val/num_eos_tokens`: The number of end-of-sequence (EOS) tokens generated, which can indicate the number of complete responses.
* `lr`: lr: The current learning rate used by the optimizer.
* `episode`: episode: The current global step or episode count in the training process.


## Cookbook

> [!IMPORTANT]
> Make sure the SFT model and reward model use the _same_ chat template. Otherwise you may find the model completions are scored incorrectly.


* Debugging TIP: `objective/rlhf_reward`: this is the ultimate objective of the RLHF training. If training works as intended, this metric should keep going up.
* Memory TIP: If you are running out of memory, you can try to reduce the `--per_device_train_batch_size` or increase the `--gradient_accumulation_steps` to reduce the memory footprint.
* Memory TIP: If you have multiple GPUs, you can also run training with DeepSpeed stage 3 to reduce the memory footprint `accelerate launch --config_file examples/accelerate_configs/deepspeed_zero3.yaml`.
* Usage TIP: We recommend to use the "EOS trick" via `--non_eos_penalty --stop_token eos`, which replaces the score of completions that do not end with an EOS token with a static scalar penalty `--penalty_reward_value`. This can help the model learn to generate more coherent completions.


## What is my model doing exactly?

To help you understand what your model is doing, we periodically log some sample completions from the model. Here is an example of a completion. In an example [tracked run at Weights and Biases](https://wandb.ai/huggingface/trl/runs/dd2o3g35), it looks like the following, allowing you to see the model's response at different stages of training. By default we generate `--num_sample_generations 10` during training, but you can customize the number of generations.

![](https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/ppov2_completions.gif)


In the logs the sampled generations look like 

```
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓
┃ query                           ┃ model response                  ┃ score    ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩
│  SUBREDDIT: r/AskReddit         │  I'm in love with a friend, and │ 3.921875 │
│                                 │ I don't know how to get rid of  │          │
│ TITLE: How do you get someone   │ those feelings. I'm             │          │
│ out of your head?


key: _node_content
string_value:
{"id_": "68fa7d05-b2a0-4295-9dbf-1abca1b5e99f", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae4d4218-2aa9-413b-bbe4-2665b89b5998", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "3a728c5c956951b40498782483a1a6fc9c27dfe14a50ee062e3eb965e3c60518", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "3a6a7e25-53a2-4c0c-aa4f-0cdfff914803", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "9cdb85b896c7ac7e7d4618d78c798de75156ea836011afbb1ad66218b1f4c6b7", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "52f7ab45-2a74-4cc6-adb3-a00d8461741f", "node_type": "1", "metadata": {}, "hash": "d17b65a0a505a96450e9e80ee9f124305aa689b75da7efe337bff523ba019494", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 6247, "end_char_idx": 9849, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
## What is my model doing exactly?

To help you understand what your model is doing, we periodically log some sample completions from the model. Here is an example of a completion. In an example [tracked run at Weights and Biases](https://wandb.ai/huggingface/trl/runs/dd2o3g35), it looks like the following, allowing you to see the model's response at different stages of training. By default we generate `--num_sample_generations 10` during training, but you can customize the number of generations.

![](https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/ppov2_completions.gif)


In the logs the sampled generations look like 

```
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓
┃ query                           ┃ model response                  ┃ score    ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩
│  SUBREDDIT: r/AskReddit         │  I'm in love with a friend, and │ 3.921875 │
│                                 │ I don't know how to get rid of  │          │
│ TITLE: How do you get someone   │ those feelings. I'm             │          │
│ out of your head?               │ desperate.<|endoftext|>[PAD][P… │          │
│                                 │                                 │          │
│ POST: Hi,                       │                                 │          │
│ I'm 22, and I have been with my │                                 │          │
│ girlfriend for 5 years now. We  │                                 │          │
│ recently moved together. We've  │                                 │          │
│ always loved each other         │                                 │          │
│ intensely.                      │                                 │          │
│                                 │                                 │          │
│ Problem, I recently started to  │                                 │          │
│ have feelings for an other      │                                 │          │
│ person (a friend). This person  │                                 │          │
│ has had a boyfriend for now 3   │                                 │          │
│ years, and has absolutely no    │                                 │          │
│ ideas. Those feelings were so   │                                 │          │
│ strong, it was hard to hide     │                                 │          │
│ them. After 2 months of me      │                                 │          │
│ being distant and really sad,   │                                 │          │
│ my girlfriend forced me to say  │                                 │          │
│ what was bothering me. I'm not  │                                 │          │
│ a good liar, and now she knows. │                                 │          │
│                                 │                                 │          │
│ We decided to give us a week    │                                 │          │
│ alone, I went to my parents.    │                                 │          │
│                                 │                                 │          │
│ Now, I'm completely lost. I     │                                 │          │
│ keep on thinking about this     │                                 │          │
│ person, and I hate that. I      │                                 │          │
│ would like for those feelings   │                                 │          │
│ to go away, to leave me alone.  │                                 │          │
│ But I can't.


key: _node_content
string_value:
{"id_": "52f7ab45-2a74-4cc6-adb3-a00d8461741f", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae4d4218-2aa9-413b-bbe4-2665b89b5998", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "3a728c5c956951b40498782483a1a6fc9c27dfe14a50ee062e3eb965e3c60518", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "68fa7d05-b2a0-4295-9dbf-1abca1b5e99f", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "db8e8239c468b4becbdb5000d9c2ed05752b21d4a7dcc8a371c6173e84eeb5ad", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "78f28d4e-be0c-4b0e-a1a1-acf08ede9b06", "node_type": "1", "metadata": {}, "hash": "6881618b741a14aab65976a68647bfc5d066005398f47fab6a037a4a84cc3f06", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 7759, "end_char_idx": 11481, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
We  │                                 │          │
│ recently moved together. We've  │                                 │          │
│ always loved each other         │                                 │          │
│ intensely.                      │                                 │          │
│                                 │                                 │          │
│ Problem, I recently started to  │                                 │          │
│ have feelings for an other      │                                 │          │
│ person (a friend). This person  │                                 │          │
│ has had a boyfriend for now 3   │                                 │          │
│ years, and has absolutely no    │                                 │          │
│ ideas. Those feelings were so   │                                 │          │
│ strong, it was hard to hide     │                                 │          │
│ them. After 2 months of me      │                                 │          │
│ being distant and really sad,   │                                 │          │
│ my girlfriend forced me to say  │                                 │          │
│ what was bothering me. I'm not  │                                 │          │
│ a good liar, and now she knows. │                                 │          │
│                                 │                                 │          │
│ We decided to give us a week    │                                 │          │
│ alone, I went to my parents.    │                                 │          │
│                                 │                                 │          │
│ Now, I'm completely lost. I     │                                 │          │
│ keep on thinking about this     │                                 │          │
│ person, and I hate that. I      │                                 │          │
│ would like for those feelings   │                                 │          │
│ to go away, to leave me alone.  │                                 │          │
│ But I can't.                    │                                 │          │
│                                 │                                 │          │
│ What do I do? It's been 3       │                                 │          │
│ months now, and I'm just        │                                 │          │
│ desperate.                      │                                 │          │
│                                 │                                 │          │
│ TL;DR:                          │                                 │          │
├─────────────────────────────────┼─────────────────────────────────┼──────────┤
│  SUBREDDIT: r/pettyrevenge      │  My mom woke me up with a loud  │ 6.84375  │
│                                 │ TV. I blasted Gangnam Style on  │          │
│ TITLE: So, my mom woke me up    │ repeat, with the bass cranked   │          │
│ with a loud TV.                 │ up as high as it could          │          │
│                                 │ go.<|endoftext|>[PAD][PAD][PAD… │          │
│ POST: She was in her living     │                                 │          │
│ room, watching TV. This was at  │                                 │          │
│ about 8:30 in the morning, and  │                                 │          │
│ she was exercising. She turned  │                                 │          │
│ the TV up extra loud to hear it │                                 │          │
│ over her excercycle, and woke   │                                 │          │
│ me up. I went in there asking   │                                 │          │
│ for her to turn it down.


key: _node_content
string_value:
{"id_": "78f28d4e-be0c-4b0e-a1a1-acf08ede9b06", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae4d4218-2aa9-413b-bbe4-2665b89b5998", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "3a728c5c956951b40498782483a1a6fc9c27dfe14a50ee062e3eb965e3c60518", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "52f7ab45-2a74-4cc6-adb3-a00d8461741f", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "134a7ffffad99a1509ac5fcd3b612b617e0e402600b051835fac7c5b4ede7109", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "2f59c366-287b-4ddd-b790-fc8d2d48f737", "node_type": "1", "metadata": {}, "hash": "bb451522e67acbb4a8fd5b9cdce53312540f6670b598f4e9dc1019ae0557cdd4", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 9619, "end_char_idx": 13647, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
I      │                                 │          │
│ would like for those feelings   │                                 │          │
│ to go away, to leave me alone.  │                                 │          │
│ But I can't.                    │                                 │          │
│                                 │                                 │          │
│ What do I do? It's been 3       │                                 │          │
│ months now, and I'm just        │                                 │          │
│ desperate.                      │                                 │          │
│                                 │                                 │          │
│ TL;DR:                          │                                 │          │
├─────────────────────────────────┼─────────────────────────────────┼──────────┤
│  SUBREDDIT: r/pettyrevenge      │  My mom woke me up with a loud  │ 6.84375  │
│                                 │ TV. I blasted Gangnam Style on  │          │
│ TITLE: So, my mom woke me up    │ repeat, with the bass cranked   │          │
│ with a loud TV.                 │ up as high as it could          │          │
│                                 │ go.<|endoftext|>[PAD][PAD][PAD… │          │
│ POST: She was in her living     │                                 │          │
│ room, watching TV. This was at  │                                 │          │
│ about 8:30 in the morning, and  │                                 │          │
│ she was exercising. She turned  │                                 │          │
│ the TV up extra loud to hear it │                                 │          │
│ over her excercycle, and woke   │                                 │          │
│ me up. I went in there asking   │                                 │          │
│ for her to turn it down. She    │                                 │          │
│ said she didn't have to; I      │                                 │          │
│ explained that I always used    │                                 │          │
│ headphones so she didn't have   │                                 │          │
│ to deal with my noise and that  │                                 │          │
│ she should give me a little     │                                 │          │
│ more respect, given that I paid │                                 │          │
│ rent at the time.               │                                 │          │
│                                 │                                 │          │
│ She disagreed. I went back to   │                                 │          │
│ my room, rather pissed off at   │                                 │          │
│ the lack of equality. I had no  │                                 │          │
│ lock on my door; but I had a    │                                 │          │
│ dresser right next to it, so I  │                                 │          │
│ pulled one of the drawers out   │                                 │          │
│ enough so that it caused the    │                                 │          │
│ door to not be openable. Then,  │                                 │          │
│ I turned my speakers up really  │                                 │          │
│ loud and blasted Gangnam Style  │                                 │          │
│ on repeat, with the bass        │                                 │          │
│ cranked up as high as it could  │                                 │          │
│ go.                             │                                 │          │
│                                 │                                 │          │
│ If you hate Gangnam Style for   │                                 │          │
│ being overplayed, you will see  │                                 │          │
│ why I chose that particular     │                                 │          │
│ song. I personally don't mind   │                                 │          │
│ it.


key: _node_content
string_value:
{"id_": "2f59c366-287b-4ddd-b790-fc8d2d48f737", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae4d4218-2aa9-413b-bbe4-2665b89b5998", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "3a728c5c956951b40498782483a1a6fc9c27dfe14a50ee062e3eb965e3c60518", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "78f28d4e-be0c-4b0e-a1a1-acf08ede9b06", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "874fd85ec43080ae03ec5e95f6ae5d99c3a5f84a22e7a99982df52be27b1fe28", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "82940200-efd2-4671-b8d8-2bd152dbcfa0", "node_type": "1", "metadata": {}, "hash": "52bb2eeeb095ff6270b640f46ac0051aa7251a8fcff9df8632735177dda7925b", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 11482, "end_char_idx": 15475, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
She    │                                 │          │
│ said she didn't have to; I      │                                 │          │
│ explained that I always used    │                                 │          │
│ headphones so she didn't have   │                                 │          │
│ to deal with my noise and that  │                                 │          │
│ she should give me a little     │                                 │          │
│ more respect, given that I paid │                                 │          │
│ rent at the time.               │                                 │          │
│                                 │                                 │          │
│ She disagreed. I went back to   │                                 │          │
│ my room, rather pissed off at   │                                 │          │
│ the lack of equality. I had no  │                                 │          │
│ lock on my door; but I had a    │                                 │          │
│ dresser right next to it, so I  │                                 │          │
│ pulled one of the drawers out   │                                 │          │
│ enough so that it caused the    │                                 │          │
│ door to not be openable. Then,  │                                 │          │
│ I turned my speakers up really  │                                 │          │
│ loud and blasted Gangnam Style  │                                 │          │
│ on repeat, with the bass        │                                 │          │
│ cranked up as high as it could  │                                 │          │
│ go.                             │                                 │          │
│                                 │                                 │          │
│ If you hate Gangnam Style for   │                                 │          │
│ being overplayed, you will see  │                                 │          │
│ why I chose that particular     │                                 │          │
│ song. I personally don't mind   │                                 │          │
│ it. But here's the thing about  │                                 │          │
│ my bass; it vibrates the walls, │                                 │          │
│ making one hell of a lot of     │                                 │          │
│ noise. Needless to say, my mom  │                                 │          │
│ was not pleased and shut off    │                                 │          │
│ the internet. But it was oh so  │                                 │          │
│ worth it.                       │                                 │          │
│                                 │                                 │          │
│ TL;DR:                          │                                 │          │
└─────────────────────────────────┴─────────────────────────────────┴──────────┘
```

## Implementation details

Many online implementation details are borrowed from the PPOv2Trainer, which is itself based on the [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031). Here are some additional implementation details:

1. When we turn on the EOS trick (i.e., replacing the score of completions that do not end with an EOS token with a scalar penalty score like `-1`) via `--non_eos_penalty --stop_token eos`, it's possible that the chosen and rejected completions have the same score. In this case, we will naively select the completion with the lower index and the chosen completion.

## Benchmark experiments

To validate the online DPO implementation works, we ran experiments on the 1B and 6.9B models. Here are the commands we used to run the experiments. We take the SFT / RM models directly from [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031).


key: _node_content
string_value:
{"id_": "82940200-efd2-4671-b8d8-2bd152dbcfa0", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae4d4218-2aa9-413b-bbe4-2665b89b5998", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "3a728c5c956951b40498782483a1a6fc9c27dfe14a50ee062e3eb965e3c60518", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "2f59c366-287b-4ddd-b790-fc8d2d48f737", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "45b66644c084a6bcda048dfadd4c35922bab8f22025890f6ec416af47d911df9", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "b5294cb6-d831-4768-b10b-1438b19db721", "node_type": "1", "metadata": {}, "hash": "f82b94ca04fe19f0832cdec05c3f9c9330423661fbbc6827884915d8f822e7c5", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 13648, "end_char_idx": 17303, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
But here's the thing about  │                                 │          │
│ my bass; it vibrates the walls, │                                 │          │
│ making one hell of a lot of     │                                 │          │
│ noise. Needless to say, my mom  │                                 │          │
│ was not pleased and shut off    │                                 │          │
│ the internet. But it was oh so  │                                 │          │
│ worth it.                       │                                 │          │
│                                 │                                 │          │
│ TL;DR:                          │                                 │          │
└─────────────────────────────────┴─────────────────────────────────┴──────────┘
```

## Implementation details

Many online implementation details are borrowed from the PPOv2Trainer, which is itself based on the [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031). Here are some additional implementation details:

1. When we turn on the EOS trick (i.e., replacing the score of completions that do not end with an EOS token with a scalar penalty score like `-1`) via `--non_eos_penalty --stop_token eos`, it's possible that the chosen and rejected completions have the same score. In this case, we will naively select the completion with the lower index and the chosen completion.

## Benchmark experiments

To validate the online DPO implementation works, we ran experiments on the 1B and 6.9B models. Here are the commands we used to run the experiments. We take the SFT / RM models directly from [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031).


```
# 1B Online DPO experiment
accelerate launch --config_file examples/accelerate_configs/deepspeed_zero2.yaml \
    examples/scripts/online_dpo.py \
    --dataset_name trl-lib/tldr \
    --learning_rate 3e-6 \
    --output_dir models/minimal/online_dpo_tldr \
    --per_device_train_batch_size 16 \
    --gradient_accumulation_steps 4 \
    --local_rollout_forward_batch_size 32 \
    --num_epochs 1 \
    --num_mini_batches 1 \
    --total_episodes 1000000 \
    --model_name_or_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr  \
    --sft_model_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr \
    --reward_model_path cleanrl/EleutherAI_pythia-1b-deduped__reward__tldr \
    --save_strategy no \
    --non_eos_penalty \
    --stop_token eos \
    --beta 0.1 \
    --response_length 53 \
    --push_to_hub

# 6.9B Online DPO experiment
accelerate launch --config_file examples/accelerate_configs/deepspeed_zero3.yaml \
    examples/scripts/online_dpo.py \
    --dataset_name trl-lib/tldr \
    --learning_rate 3e-6 \
    --output_dir models/minimal/online_dpo_tldr_6.9b \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 16 \
    --local_rollout_forward_batch_size 8 \
    --num_epochs 1 \
    --num_mini_batches 1 \
    --total_episodes 1000000 \
    --model_name_or_path EleutherAI/pythia-6.9b-deduped \
    --sft_model_path cleanrl/EleutherAI_pythia-6.9b-deduped__sft__tldr \
    --reward_model_path cleanrl/EleutherAI_pythia-6.9b-deduped__reward__tldr \
    --save_strategy no \
    --non_eos_penalty \
    --stop_token eos \
    --beta 0.1 \
    --response_length 53 \
    --push_to_hub
```

Checkpoints and experiment tracking are available at:

- [🤗 Model checkpoint](https://huggingface.co/vwxyzjn/ppo_tldr)
- [🐝 Tracked experiment](https://wandb.ai/huggingface/trl/runs/dd2o3g35)


key: _node_content
string_value:
{"id_": "b5294cb6-d831-4768-b10b-1438b19db721", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae4d4218-2aa9-413b-bbe4-2665b89b5998", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "3a728c5c956951b40498782483a1a6fc9c27dfe14a50ee062e3eb965e3c60518", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "82940200-efd2-4671-b8d8-2bd152dbcfa0", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/online_dpo_trainer", "title": "Online DPO Trainer", "tokens": 4521, "retrieve_doc": true, "source": "trl"}, "hash": "da1a8778460ebdeb47d563a51c252ee8926d9fb8d35c0f5dd8aa5d05539b1812", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 17306, "end_char_idx": 19216, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
To evaluate, we use [vLLM](https://github.com/vllm-project/vllm) to load the checkpoints and GPT-4o mini as a judge model to evaluate the generated TL;DR against the reference TL;DR.
For more information on how to use judges, see [Judges](judges).

```bash
$ python examples/scripts/evals/judge_tldr.py --model_name_or_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr --judge_model gpt-4o-mini --num_examples 1000
Model win rate: 33.00%
python examples/scripts/evals/judge_tldr.py --model_name_or_path cleanrl/EleutherAI_pythia-6.9b-deduped__sft__tldr --judge_model gpt-4o-mini --num_examples 1000
Model win rate: 41.50%
python examples/scripts/evals/judge_tldr.py --model_name_or_path vwxyzjn/online_dpo_tldr --judge_model gpt-4o-mini --num_examples 1000
Model win rate: 62.60%
python examples/scripts/evals/judge_tldr.py --model_name_or_path vwxyzjn/online_dpo_tldr_6.9b --judge_model gpt-4o-mini --num_examples 1000
Model win rate: 74.20%
```

We can then plot the RLHF scaling chart.

```python
import matplotlib.pyplot as plt

data = {
    "SFT": [[1e9, 6.9e9], [0.33, 0.415]],
    "Online DPO": [[1e9, 6.9e9], [0.626, 0.742]],
}
for model, (x, y) in data.items():
    plt.scatter(x, y, label=model)

plt.axhline(y=0.5, color="black", linestyle="-.", label="Human reference summary")
plt.title("RLHF scaling by model size")
plt.xlabel("Model size")
plt.ylabel("Win rate against reference summaries\n(according to GPT-4o mini)")
plt.xscale("log")
plt.xlim(5e8, 1.2e10)
plt.xticks([1e9, 1e10], ["1B", "10B"])
plt.legend()
plt.grid(True, which="both", ls="--", c="0.7")
plt.tight_layout()
plt.savefig("plot.png")
```


![](https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/online_dpo_scaling.png)

The online DPO checkpoint gets increasingly more win rate as we scale up the model sizes. This is a good sign that the online DPO implementation is working as intended.


key: _node_content
string_value:
{"id_": "05b3cee7-cdb2-4bdf-b59b-4e156e237ac1", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/judges", "title": "Judges", "tokens": 553, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "877f15bd-d503-422d-b8df-d7093846a3d2", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/judges", "title": "Judges", "tokens": 553, "retrieve_doc": true, "source": "trl"}, "hash": "3fb9015fa558da6c9cfa52d38e382d6eeda0753f7b98e0b51a1045435480864a", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 2163, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# Judges

TRL provides judges to easily compare two completions.

Make sure to have installed the required dependencies by running:

```bash
pip install trl[llm_judge]
```

## Using the provided judges

TRL provides several judges out of the box. For example, you can use the `HfPairwiseJudge` to compare two completions using a pre-trained model from the Hugging Face model hub:

```python
from trl import HfPairwiseJudge

judge = HfPairwiseJudge()
judge.judge(
    prompts=["What is the capital of France?", "What is the biggest planet in the solar system?"],
    completions=[["Paris", "Lyon"], ["Saturn", "Jupiter"]],
)  # Outputs: [0, 1]
```

## Define your own judge

To define your own judge, we provide several base classes that you can subclass. For rank-based judges, you need to subclass [`BaseRankJudge`] and implement the [`BaseRankJudge.judge`] method. For pairwise judges, you need to subclass [`BasePairJudge`] and implement the [`BasePairJudge.judge`] method. If you want to define a judge that doesn't fit into these categories, you need to subclass [`BaseJudge`] and implement the [`BaseJudge.judge`] method.

As an example, let's define a pairwise judge that prefers shorter completions:

```python
from trl import BasePairwiseJudge

class PrefersShorterJudge(BasePairwiseJudge):
    def judge(self, prompts, completions, shuffle_order=False):
        return [0 if len(completion[0]) > len(completion[1]) else 1 for completion in completions]
```

You can then use this judge as follows:

```python
judge = PrefersShorterJudge()
judge.judge(
    prompts=["What is the capital of France?", "What is the biggest planet in the solar system?"],
    completions=[["Paris", "The capital of France is Paris."], ["Jupiter is the biggest planet in the solar system.", "Jupiter"]],
)  # Outputs: [0, 1]
```

## BaseJudge

[[autodoc]] BaseJudge

## BaseRankJudge

[[autodoc]] BaseRankJudge

## BasePairwiseJudge

[[autodoc]] BasePairwiseJudge

## RandomRankJudge

[[autodoc]] RandomRankJudge

## RandomPairwiseJudge

[[autodoc]] RandomPairwiseJudge

## HfPairwiseJudge

[[autodoc]] HfPairwiseJudge

## OpenAIPairwiseJudge

[[autodoc]] OpenAIPairwiseJudge


key: _node_content
string_value:
{"id_": "35192912-b4b4-4841-8d4b-14f40f14b290", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/reward_trainer", "title": "Reward Modeling", "tokens": 889, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "65643d67-d101-4275-ae94-47c6fb11fae3", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/reward_trainer", "title": "Reward Modeling", "tokens": 889, "retrieve_doc": true, "source": "trl"}, "hash": "01663f69a87cb0fb41463365e36fd3d044ffae01b4de18a3fb35c8dac45080fa", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "3de4f44a-c079-4d1b-9f17-bb5809e97a82", "node_type": "1", "metadata": {}, "hash": "999325dbf6239d21eb97d57a1b3c9de76665940263f4ed788f143cc0c4096717", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 3031, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# Reward Modeling

TRL supports custom reward modeling for anyone to perform reward modeling on their dataset and model.

Check out a complete flexible example at [`examples/scripts/reward_modeling.py`](https://github.com/huggingface/trl/tree/main/examples/scripts/reward_modeling.py).

## Expected dataset format

The [`RewardTrainer`] expects a very specific format for the dataset since the model will be trained on pairs of examples to predict which of the two is preferred. We provide an example from the [`Anthropic/hh-rlhf`](https://huggingface.co/datasets/Anthropic/hh-rlhf) dataset below:

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/rlhf-antropic-example.png", width="50%">
</div>

Therefore the final dataset object should contain two 4 entries at least if you use the default [`RewardDataCollatorWithPadding`] data collator. The entries should be named:

-   `input_ids_chosen`
-   `attention_mask_chosen`
-   `input_ids_rejected`
-   `attention_mask_rejected`

## Using the `RewardTrainer`

After preparing your dataset, you can use the [`RewardTrainer`] in the same way as the `Trainer` class from 🤗 Transformers.
You should pass an `AutoModelForSequenceClassification` model to the [`RewardTrainer`], along with a [`RewardConfig`] which configures the hyperparameters of the training.

### Leveraging 🤗 PEFT to train a reward model

Just pass a `peft_config` in the keyword arguments of [`RewardTrainer`], and the trainer should automatically take care of converting the model into a PEFT model!

```python
from peft import LoraConfig, TaskType
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from trl import RewardTrainer, RewardConfig

model = AutoModelForSequenceClassification.from_pretrained("gpt2")
peft_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
)

...

trainer = RewardTrainer(
    model=model,
    args=training_args,
    tokenizer=tokenizer,
    train_dataset=dataset,
    peft_config=peft_config,
)

trainer.train()

```

### Adding a margin to the loss

As in the [Llama 2 paper](https://huggingface.co/papers/2307.09288), you can add a margin to the loss by adding a `margin` column to the dataset. The reward collator will automatically pass it through and the loss will be computed accordingly.

```python
def add_margin(row):
    # Assume you have a score_chosen and score_rejected columns that you want to use to compute the margin
    return {'margin': row['score_chosen'] - row['score_rejected']}

dataset = dataset.map(add_margin)
```

### Centering rewards

In many scenarios, it's preferable to ensure that a reward model's output is mean zero. This is often done by first calculating the model's average score and then subtracting it.

[[Eisenstein et al., 2023]](https://huggingface.co/papers/2312.09244) proposed an auxiliary loss function designed to directly learn a centered reward model.


key: _node_content
string_value:
{"id_": "3de4f44a-c079-4d1b-9f17-bb5809e97a82", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/reward_trainer", "title": "Reward Modeling", "tokens": 889, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "65643d67-d101-4275-ae94-47c6fb11fae3", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/reward_trainer", "title": "Reward Modeling", "tokens": 889, "retrieve_doc": true, "source": "trl"}, "hash": "01663f69a87cb0fb41463365e36fd3d044ffae01b4de18a3fb35c8dac45080fa", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "35192912-b4b4-4841-8d4b-14f40f14b290", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/reward_trainer", "title": "Reward Modeling", "tokens": 889, "retrieve_doc": true, "source": "trl"}, "hash": "14b0567a8e4f02e88df1abdeba721f6b5457220a7ed48d5fc9c12defa3a460a3", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 1599, "end_char_idx": 3693, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
```python
from peft import LoraConfig, TaskType
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from trl import RewardTrainer, RewardConfig

model = AutoModelForSequenceClassification.from_pretrained("gpt2")
peft_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
)

...

trainer = RewardTrainer(
    model=model,
    args=training_args,
    tokenizer=tokenizer,
    train_dataset=dataset,
    peft_config=peft_config,
)

trainer.train()

```

### Adding a margin to the loss

As in the [Llama 2 paper](https://huggingface.co/papers/2307.09288), you can add a margin to the loss by adding a `margin` column to the dataset. The reward collator will automatically pass it through and the loss will be computed accordingly.

```python
def add_margin(row):
    # Assume you have a score_chosen and score_rejected columns that you want to use to compute the margin
    return {'margin': row['score_chosen'] - row['score_rejected']}

dataset = dataset.map(add_margin)
```

### Centering rewards

In many scenarios, it's preferable to ensure that a reward model's output is mean zero. This is often done by first calculating the model's average score and then subtracting it.

[[Eisenstein et al., 2023]](https://huggingface.co/papers/2312.09244) proposed an auxiliary loss function designed to directly learn a centered reward model. This auxiliary loss minimizes the squared sum of the rewards, encouraging the model to naturally produce mean-zero outputs:

$$\Big( R(p, r_1) + R(p, r_2) \Big)^2 $$

This auxiliary loss is combined with the main loss function, weighted by the parameter `center_rewards_coefficient` in the `[RewardConfig]`. By default, this feature is deactivated (`center_rewards_coefficient = None`).

```python
reward_config = RewardConfig(
    center_rewards_coefficient=0.01,
    ...
)
```

For reference results, please refer PR [#1932](https://github.com/huggingface/trl/pull/1932).

## RewardConfig

[[autodoc]] RewardConfig

## RewardTrainer

[[autodoc]] RewardTrainer


key: _node_content
string_value:
{"id_": "d7eb4ce5-bffb-4996-988d-a4e28049846b", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "cad47e91-7c5e-47ac-9a1b-32574008751d", "node_type": "1", "metadata": {}, "hash": "79b46277667bf2316397a50b2672c75c241041feb3441f5010fc3c48cf2f0fed", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 2441, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# Supervised Fine-tuning Trainer

Supervised fine-tuning (or SFT for short) is a crucial step in RLHF. In TRL we provide an easy-to-use API to create your SFT models and train them with few lines of code on your dataset.

Check out a complete flexible example at [`examples/scripts/sft.py`](https://github.com/huggingface/trl/tree/main/examples/scripts/sft.py).
Experimental support for Vision Language Models is also included in the example [`examples/scripts/vsft_llava.py`](https://github.com/huggingface/trl/tree/main/examples/scripts/vsft_llava.py).

## Quickstart

If you have a dataset hosted on the 🤗 Hub, you can easily fine-tune your SFT model using [`SFTTrainer`] from TRL. Let us assume your dataset is `imdb`, the text you want to predict is inside the `text` field of the dataset, and you want to fine-tune the `facebook/opt-350m` model.
The following code-snippet takes care of all the data pre-processing and training for you:

```python
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer

dataset = load_dataset("imdb", split="train")

sft_config = SFTConfig(
    dataset_text_field="text",
    max_seq_length=512,
    output_dir="/tmp",
)
trainer = SFTTrainer(
    "facebook/opt-350m",
    train_dataset=dataset,
    args=sft_config,
)
trainer.train()
```
Make sure to pass the correct value for `max_seq_length` as the default value will be set to `min(tokenizer.model_max_length, 1024)`.

You can also construct a model outside of the trainer and pass it as follows:

```python
from transformers import AutoModelForCausalLM
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer

dataset = load_dataset("imdb", split="train")

model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m")

sft_config = SFTConfig(output_dir="/tmp")

trainer = SFTTrainer(
    model,
    train_dataset=dataset,
    args=sft_config,
)

trainer.train()
```

The above snippets will use the default training arguments from the [`SFTConfig`] class. If you want to modify the defaults pass in your modification to the `SFTConfig` constructor and pass them to the trainer via the `args` argument.

## Advanced usage

### Train on completions only

You can use the `DataCollatorForCompletionOnlyLM` to train your model on the generated prompts only. Note that this works only in the case when `packing=False`.
To instantiate that collator for instruction data, pass a response template and the tokenizer.


key: _node_content
string_value:
{"id_": "cad47e91-7c5e-47ac-9a1b-32574008751d", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "d7eb4ce5-bffb-4996-988d-a4e28049846b", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "ca831dce5b552e21aa5af5f0e84b3bb679fa749a28d93dad98eb1815a1879cdd", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "25173168-bd65-47a4-a6e0-c1305e6211e0", "node_type": "1", "metadata": {}, "hash": "dc51635ad9d1c6d4d3066118f6557e10916e3079b7ba6d27c200d5d7f9a7e9d7", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 852, "end_char_idx": 3650, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
The following code-snippet takes care of all the data pre-processing and training for you:

```python
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer

dataset = load_dataset("imdb", split="train")

sft_config = SFTConfig(
    dataset_text_field="text",
    max_seq_length=512,
    output_dir="/tmp",
)
trainer = SFTTrainer(
    "facebook/opt-350m",
    train_dataset=dataset,
    args=sft_config,
)
trainer.train()
```
Make sure to pass the correct value for `max_seq_length` as the default value will be set to `min(tokenizer.model_max_length, 1024)`.

You can also construct a model outside of the trainer and pass it as follows:

```python
from transformers import AutoModelForCausalLM
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer

dataset = load_dataset("imdb", split="train")

model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m")

sft_config = SFTConfig(output_dir="/tmp")

trainer = SFTTrainer(
    model,
    train_dataset=dataset,
    args=sft_config,
)

trainer.train()
```

The above snippets will use the default training arguments from the [`SFTConfig`] class. If you want to modify the defaults pass in your modification to the `SFTConfig` constructor and pass them to the trainer via the `args` argument.

## Advanced usage

### Train on completions only

You can use the `DataCollatorForCompletionOnlyLM` to train your model on the generated prompts only. Note that this works only in the case when `packing=False`.
To instantiate that collator for instruction data, pass a response template and the tokenizer. Here is an example of how it would work to fine-tune `opt-350m` on completions only on the CodeAlpaca dataset:

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer, DataCollatorForCompletionOnlyLM

dataset = load_dataset("lucasmccabe-lmi/CodeAlpaca-20k", split="train")

model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m")
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")

def formatting_prompts_func(example):
    output_texts = []
    for i in range(len(example['instruction'])):
        text = f"### Question: {example['instruction'][i]}\n ### Answer: {example['output'][i]}"
        output_texts.append(text)
    return output_texts

response_template = " ### Answer:"
collator = DataCollatorForCompletionOnlyLM(response_template, tokenizer=tokenizer)

trainer = SFTTrainer(
    model,
    train_dataset=dataset,
    args=SFTConfig(output_dir="/tmp"),
    formatting_func=formatting_prompts_func,
    data_collator=collator,
)

trainer.train()
```

To instantiate that collator for assistant style conversation data, pass a response template, an instruction template and the tokenizer.


key: _node_content
string_value:
{"id_": "25173168-bd65-47a4-a6e0-c1305e6211e0", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "cad47e91-7c5e-47ac-9a1b-32574008751d", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "0e10b669c65006f0547c2ed69c02b1a020e0e1e580b54ca5be9a74f10517f089", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "b265c072-a6f9-408a-a1a7-9a1a4c5d4346", "node_type": "1", "metadata": {}, "hash": "bc56e892b1e7aef3891ddb822de27d82909e497323e2d97ab45e1657e65c9677", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 1989, "end_char_idx": 4968, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
If you want to modify the defaults pass in your modification to the `SFTConfig` constructor and pass them to the trainer via the `args` argument.

## Advanced usage

### Train on completions only

You can use the `DataCollatorForCompletionOnlyLM` to train your model on the generated prompts only. Note that this works only in the case when `packing=False`.
To instantiate that collator for instruction data, pass a response template and the tokenizer. Here is an example of how it would work to fine-tune `opt-350m` on completions only on the CodeAlpaca dataset:

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer, DataCollatorForCompletionOnlyLM

dataset = load_dataset("lucasmccabe-lmi/CodeAlpaca-20k", split="train")

model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m")
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")

def formatting_prompts_func(example):
    output_texts = []
    for i in range(len(example['instruction'])):
        text = f"### Question: {example['instruction'][i]}\n ### Answer: {example['output'][i]}"
        output_texts.append(text)
    return output_texts

response_template = " ### Answer:"
collator = DataCollatorForCompletionOnlyLM(response_template, tokenizer=tokenizer)

trainer = SFTTrainer(
    model,
    train_dataset=dataset,
    args=SFTConfig(output_dir="/tmp"),
    formatting_func=formatting_prompts_func,
    data_collator=collator,
)

trainer.train()
```

To instantiate that collator for assistant style conversation data, pass a response template, an instruction template and the tokenizer. Here is an example of how it would work to fine-tune `opt-350m` on assistant completions only on the Open Assistant Guanaco dataset:

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer, DataCollatorForCompletionOnlyLM

dataset = load_dataset("timdettmers/openassistant-guanaco", split="train")

model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m")
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")

instruction_template = "### Human:"
response_template = "### Assistant:"
collator = DataCollatorForCompletionOnlyLM(instruction_template=instruction_template, response_template=response_template, tokenizer=tokenizer, mlm=False)

trainer = SFTTrainer(
    model,
    args=SFTConfig(
        output_dir="/tmp",
        dataset_text_field = "text",
    ),
    train_dataset=dataset,
    data_collator=collator,
)

trainer.train()
```

Make sure to have a `pad_token_id` which is different from `eos_token_id` which can result in the model not properly predicting EOS (End of Sentence) tokens during generation.

#### Using token_ids directly for `response_template`

Some tokenizers like Llama 2 (`meta-llama/Llama-2-XXb-hf`) tokenize sequences differently depending on whether they have context or not.


key: _node_content
string_value:
{"id_": "b265c072-a6f9-408a-a1a7-9a1a4c5d4346", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "25173168-bd65-47a4-a6e0-c1305e6211e0", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "3e3566d706bd8d16afe292affae1a96e056f7bc181c5f1b13755519180b93c6e", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "e3d4d940-61da-430d-a6ec-26640473a8d8", "node_type": "1", "metadata": {}, "hash": "526dc27809ca8302a563981c06eaa0ebd93f8a10fe77395e0fde13a8ed0adf0b", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 3651, "end_char_idx": 6566, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
Here is an example of how it would work to fine-tune `opt-350m` on assistant completions only on the Open Assistant Guanaco dataset:

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer, DataCollatorForCompletionOnlyLM

dataset = load_dataset("timdettmers/openassistant-guanaco", split="train")

model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m")
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")

instruction_template = "### Human:"
response_template = "### Assistant:"
collator = DataCollatorForCompletionOnlyLM(instruction_template=instruction_template, response_template=response_template, tokenizer=tokenizer, mlm=False)

trainer = SFTTrainer(
    model,
    args=SFTConfig(
        output_dir="/tmp",
        dataset_text_field = "text",
    ),
    train_dataset=dataset,
    data_collator=collator,
)

trainer.train()
```

Make sure to have a `pad_token_id` which is different from `eos_token_id` which can result in the model not properly predicting EOS (End of Sentence) tokens during generation.

#### Using token_ids directly for `response_template`

Some tokenizers like Llama 2 (`meta-llama/Llama-2-XXb-hf`) tokenize sequences differently depending on whether they have context or not. For example:

```python
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

def print_tokens_with_ids(txt):
    tokens = tokenizer.tokenize(txt, add_special_tokens=False)
    token_ids = tokenizer.encode(txt, add_special_tokens=False)
    print(list(zip(tokens, token_ids)))

prompt = """### User: Hello\n\n### Assistant: Hi, how can I help you?"""
print_tokens_with_ids(prompt)  # [..., ('▁Hello', 15043), ('<0x0A>', 13), ('<0x0A>', 13), ('##', 2277), ('#', 29937), ('▁Ass', 4007), ('istant', 22137), (':', 29901), ...]

response_template = "### Assistant:"
print_tokens_with_ids(response_template)  # [('▁###', 835), ('▁Ass', 4007), ('istant', 22137), (':', 29901)]
```

In this case, and due to lack of context in `response_template`, the same string ("### Assistant:") is tokenized differently:

    - Text (with context): `[2277, 29937, 4007, 22137, 29901]`
    - `response_template` (without context): `[835, 4007, 22137, 29901]`

This will lead to an error when the `DataCollatorForCompletionOnlyLM` does not find the `response_template` in the dataset example text:

```
RuntimeError: Could not find response key [835, 4007, 22137, 29901] in token IDs tensor([    1,   835,  ...])
```


To solve this, you can tokenize the `response_template` with the same context as in the dataset, truncate it as needed and pass the `token_ids` directly to the `response_template` argument of the `DataCollatorForCompletionOnlyLM` class. For example:

```python
response_template_with_context = "\n### Assistant:"  # We added context here: "\n".


key: _node_content
string_value:
{"id_": "e3d4d940-61da-430d-a6ec-26640473a8d8", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "b265c072-a6f9-408a-a1a7-9a1a4c5d4346", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "b4c14841ebf0b2a978f268576ff9e438116afed905e54c7cfa0e3ea31a855f8e", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "a34fef11-5909-4d13-b4db-179178f1ba4e", "node_type": "1", "metadata": {}, "hash": "5b62e8489586b3e69cfb6624456fe2e321ab6da3cf0bd18086ba87e10c78f4ce", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 5376, "end_char_idx": 7926, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
print_tokens_with_ids(prompt)  # [..., ('▁Hello', 15043), ('<0x0A>', 13), ('<0x0A>', 13), ('##', 2277), ('#', 29937), ('▁Ass', 4007), ('istant', 22137), (':', 29901), ...]

response_template = "### Assistant:"
print_tokens_with_ids(response_template)  # [('▁###', 835), ('▁Ass', 4007), ('istant', 22137), (':', 29901)]
```

In this case, and due to lack of context in `response_template`, the same string ("### Assistant:") is tokenized differently:

    - Text (with context): `[2277, 29937, 4007, 22137, 29901]`
    - `response_template` (without context): `[835, 4007, 22137, 29901]`

This will lead to an error when the `DataCollatorForCompletionOnlyLM` does not find the `response_template` in the dataset example text:

```
RuntimeError: Could not find response key [835, 4007, 22137, 29901] in token IDs tensor([    1,   835,  ...])
```


To solve this, you can tokenize the `response_template` with the same context as in the dataset, truncate it as needed and pass the `token_ids` directly to the `response_template` argument of the `DataCollatorForCompletionOnlyLM` class. For example:

```python
response_template_with_context = "\n### Assistant:"  # We added context here: "\n". This is enough for this tokenizer
response_template_ids = tokenizer.encode(response_template_with_context, add_special_tokens=False)[2:]  # Now we have it like in the dataset texts: `[2277, 29937, 4007, 22137, 29901]`

data_collator = DataCollatorForCompletionOnlyLM(response_template_ids, tokenizer=tokenizer)
```

### Add Special Tokens for Chat Format

Adding special tokens to a language model is crucial for training chat models. These tokens are added between the different roles in a conversation, such as the user, assistant, and system and help the model recognize the structure and flow of a conversation. This setup is essential for enabling the model to generate coherent and contextually appropriate responses in a chat environment. 
The [`setup_chat_format`] function in `trl` easily sets up a model and tokenizer for conversational AI tasks. This function:
- Adds special tokens to the tokenizer, e.g. `<|im_start|>` and `<|im_end|>`, to indicate the start and end of a conversation.
- Resizes the model’s embedding layer to accommodate the new tokens.
- Sets the `chat_template` of the tokenizer, which is used to format the input data into a chat-like format. The default is `chatml` from OpenAI.
- _optionally_ you can pass `resize_to_multiple_of` to resize the embedding layer to a multiple of the `resize_to_multiple_of` argument, e.g. 64.


key: _node_content
string_value:
{"id_": "a34fef11-5909-4d13-b4db-179178f1ba4e", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "e3d4d940-61da-430d-a6ec-26640473a8d8", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "36e581c94aee832d177df0740898c060ff12b73a36c0319f4c892cb3bbaeaf51", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "be18e574-9e0e-4ba1-8320-60eceff453f2", "node_type": "1", "metadata": {}, "hash": "dc0fe6a248f0c9c63e13559c887c3253e96ab13574972c1885bb13e50ec11362", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 6222, "end_char_idx": 9309, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
To solve this, you can tokenize the `response_template` with the same context as in the dataset, truncate it as needed and pass the `token_ids` directly to the `response_template` argument of the `DataCollatorForCompletionOnlyLM` class. For example:

```python
response_template_with_context = "\n### Assistant:"  # We added context here: "\n". This is enough for this tokenizer
response_template_ids = tokenizer.encode(response_template_with_context, add_special_tokens=False)[2:]  # Now we have it like in the dataset texts: `[2277, 29937, 4007, 22137, 29901]`

data_collator = DataCollatorForCompletionOnlyLM(response_template_ids, tokenizer=tokenizer)
```

### Add Special Tokens for Chat Format

Adding special tokens to a language model is crucial for training chat models. These tokens are added between the different roles in a conversation, such as the user, assistant, and system and help the model recognize the structure and flow of a conversation. This setup is essential for enabling the model to generate coherent and contextually appropriate responses in a chat environment. 
The [`setup_chat_format`] function in `trl` easily sets up a model and tokenizer for conversational AI tasks. This function:
- Adds special tokens to the tokenizer, e.g. `<|im_start|>` and `<|im_end|>`, to indicate the start and end of a conversation.
- Resizes the model’s embedding layer to accommodate the new tokens.
- Sets the `chat_template` of the tokenizer, which is used to format the input data into a chat-like format. The default is `chatml` from OpenAI.
- _optionally_ you can pass `resize_to_multiple_of` to resize the embedding layer to a multiple of the `resize_to_multiple_of` argument, e.g. 64. If you want to see more formats being supported in the future, please open a GitHub issue on [trl](https://github.com/huggingface/trl)

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import setup_chat_format

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m")
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")

# Set up the chat format with default 'chatml' format
model, tokenizer = setup_chat_format(model, tokenizer)

```

With our model and tokenizer set up, we can now fine-tune our model on a conversational dataset. Below is an example of how a dataset can be formatted for fine-tuning. 

### Dataset format support

The [`SFTTrainer`] supports popular dataset formats. This allows you to pass the dataset to the trainer without any pre-processing directly. The following formats are supported:
* conversational format
```json
{"messages": [{"role": "system", "content": "You are helpful"}, {"role": "user", "content": "What's the capital of France?"}, {"role": "assistant", "content": "..."}]}
{"messages": [{"role": "system", "content": "You are helpful"}, {"role": "user", "content": "Who wrote 'Romeo and Juliet'?"}, {"role": "assistant", "content": "..."}]}
{"messages": [{"role": "system", "content": "You are helpful"}, {"role": "user", "content": "How far is the Moon from Earth?


key: _node_content
string_value:
{"id_": "be18e574-9e0e-4ba1-8320-60eceff453f2", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "a34fef11-5909-4d13-b4db-179178f1ba4e", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "0597ca2a4579ec111a3a84590fc53a550b0b284bf60a6d9b0235703bebe0406a", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "bf9b80d9-e260-4700-8734-1f33c0ffa25d", "node_type": "1", "metadata": {}, "hash": "68bb74f56bb51aa576e8cb559e5e2fa76437c53cbbf41322fd9f8cb5aab6b3fc", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 7744, "end_char_idx": 10759, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
The default is `chatml` from OpenAI.
- _optionally_ you can pass `resize_to_multiple_of` to resize the embedding layer to a multiple of the `resize_to_multiple_of` argument, e.g. 64. If you want to see more formats being supported in the future, please open a GitHub issue on [trl](https://github.com/huggingface/trl)

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import setup_chat_format

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m")
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")

# Set up the chat format with default 'chatml' format
model, tokenizer = setup_chat_format(model, tokenizer)

```

With our model and tokenizer set up, we can now fine-tune our model on a conversational dataset. Below is an example of how a dataset can be formatted for fine-tuning. 

### Dataset format support

The [`SFTTrainer`] supports popular dataset formats. This allows you to pass the dataset to the trainer without any pre-processing directly. The following formats are supported:
* conversational format
```json
{"messages": [{"role": "system", "content": "You are helpful"}, {"role": "user", "content": "What's the capital of France?"}, {"role": "assistant", "content": "..."}]}
{"messages": [{"role": "system", "content": "You are helpful"}, {"role": "user", "content": "Who wrote 'Romeo and Juliet'?"}, {"role": "assistant", "content": "..."}]}
{"messages": [{"role": "system", "content": "You are helpful"}, {"role": "user", "content": "How far is the Moon from Earth?"}, {"role": "assistant", "content": "..."}]}
```
* instruction format
```json
{"prompt": "<prompt text>", "completion": "<ideal generated text>"}
{"prompt": "<prompt text>", "completion": "<ideal generated text>"}
{"prompt": "<prompt text>", "completion": "<ideal generated text>"}
```

If your dataset uses one of the above formats, you can directly pass it to the trainer without pre-processing. The [`SFTTrainer`] will then format the dataset for you using the defined format from the model's tokenizer with the [apply_chat_template](https://huggingface.co/docs/transformers/main/en/chat_templating#templates-for-chat-models) method. 


```python
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer

...

# load jsonl dataset
dataset = load_dataset("json", data_files="path/to/dataset.jsonl", split="train")
# load dataset from the HuggingFace Hub
dataset = load_dataset("philschmid/dolly-15k-oai-style", split="train")

...

sft_config = SFTConfig(packing=True)
trainer = SFTTrainer(
    "facebook/opt-350m",
    args=sft_config,
    train_dataset=dataset,
)
```

If the dataset is not in one of those format you can either preprocess the dataset to match the formatting or pass a formatting function to the SFTTrainer to do it for you. Let's have a look.


### Format your input prompts

For instruction fine-tuning, it is quite common to have two columns inside the dataset: one for the prompt & the other for the response.


key: _node_content
string_value:
{"id_": "bf9b80d9-e260-4700-8734-1f33c0ffa25d", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "be18e574-9e0e-4ba1-8320-60eceff453f2", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "93ca685596cba11268138238c5b8d523d4c36cae62445c8dc0fe7fbbf87ebfde", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "8153222b-1942-4a7d-8bb5-f08d0d69cd44", "node_type": "1", "metadata": {}, "hash": "65ec6533565b2728a122b8502d1acef169ffadb3ae941d6f289322923b59e1c8", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 9309, "end_char_idx": 12193, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
"}, {"role": "assistant", "content": "..."}]}
```
* instruction format
```json
{"prompt": "<prompt text>", "completion": "<ideal generated text>"}
{"prompt": "<prompt text>", "completion": "<ideal generated text>"}
{"prompt": "<prompt text>", "completion": "<ideal generated text>"}
```

If your dataset uses one of the above formats, you can directly pass it to the trainer without pre-processing. The [`SFTTrainer`] will then format the dataset for you using the defined format from the model's tokenizer with the [apply_chat_template](https://huggingface.co/docs/transformers/main/en/chat_templating#templates-for-chat-models) method. 


```python
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer

...

# load jsonl dataset
dataset = load_dataset("json", data_files="path/to/dataset.jsonl", split="train")
# load dataset from the HuggingFace Hub
dataset = load_dataset("philschmid/dolly-15k-oai-style", split="train")

...

sft_config = SFTConfig(packing=True)
trainer = SFTTrainer(
    "facebook/opt-350m",
    args=sft_config,
    train_dataset=dataset,
)
```

If the dataset is not in one of those format you can either preprocess the dataset to match the formatting or pass a formatting function to the SFTTrainer to do it for you. Let's have a look.


### Format your input prompts

For instruction fine-tuning, it is quite common to have two columns inside the dataset: one for the prompt & the other for the response.
This allows people to format examples like [Stanford-Alpaca](https://github.com/tatsu-lab/stanford_alpaca) did as follows:
```bash
Below is an instruction ...

### Instruction
{prompt}

### Response:
{completion}
```
Let us assume your dataset has two fields, `question` and `answer`. Therefore you can just run:
```python
...
def formatting_prompts_func(example):
    output_texts = []
    for i in range(len(example['question'])):
        text = f"### Question: {example['question'][i]}\n ### Answer: {example['answer'][i]}"
        output_texts.append(text)
    return output_texts

trainer = SFTTrainer(
    model,
    args=sft_config,
    train_dataset=dataset,
    formatting_func=formatting_prompts_func,
)

trainer.train()
```
To properly format your input make sure to process all the examples by looping over them and returning a list of processed text. Check out a full example of how to use SFTTrainer on alpaca dataset [here](https://github.com/huggingface/trl/pull/444#issue-1760952763)

### Packing dataset ([`ConstantLengthDataset`])

[`SFTTrainer`] supports _example packing_, where multiple short examples are packed in the same input sequence to increase training efficiency. This is done with the [`ConstantLengthDataset`] utility class that returns constant length chunks of tokens from a stream of examples. To enable the usage of this dataset class, simply pass `packing=True` to the [`SFTConfig`] constructor.


key: _node_content
string_value:
{"id_": "8153222b-1942-4a7d-8bb5-f08d0d69cd44", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "bf9b80d9-e260-4700-8734-1f33c0ffa25d", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "20f51a81c054133f3a6658f68c7bd5647d17f3a93d872db164a3f368ae9be262", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "a4361bf0-2970-4170-a615-654fcfe16179", "node_type": "1", "metadata": {}, "hash": "4cc91d31df254813f1462002ea0b423f22044784eaf173210fbe0295db0d6f2d", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 10592, "end_char_idx": 13757, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
### Format your input prompts

For instruction fine-tuning, it is quite common to have two columns inside the dataset: one for the prompt & the other for the response.
This allows people to format examples like [Stanford-Alpaca](https://github.com/tatsu-lab/stanford_alpaca) did as follows:
```bash
Below is an instruction ...

### Instruction
{prompt}

### Response:
{completion}
```
Let us assume your dataset has two fields, `question` and `answer`. Therefore you can just run:
```python
...
def formatting_prompts_func(example):
    output_texts = []
    for i in range(len(example['question'])):
        text = f"### Question: {example['question'][i]}\n ### Answer: {example['answer'][i]}"
        output_texts.append(text)
    return output_texts

trainer = SFTTrainer(
    model,
    args=sft_config,
    train_dataset=dataset,
    formatting_func=formatting_prompts_func,
)

trainer.train()
```
To properly format your input make sure to process all the examples by looping over them and returning a list of processed text. Check out a full example of how to use SFTTrainer on alpaca dataset [here](https://github.com/huggingface/trl/pull/444#issue-1760952763)

### Packing dataset ([`ConstantLengthDataset`])

[`SFTTrainer`] supports _example packing_, where multiple short examples are packed in the same input sequence to increase training efficiency. This is done with the [`ConstantLengthDataset`] utility class that returns constant length chunks of tokens from a stream of examples. To enable the usage of this dataset class, simply pass `packing=True` to the [`SFTConfig`] constructor.

```python
...
sft_config = SFTConfig(packing=True, dataset_text_field="text",)

trainer = SFTTrainer(
    "facebook/opt-350m",
    train_dataset=dataset,
    args=sft_config
)

trainer.train()
```

Note that if you use a packed dataset and if you pass `max_steps` in the training arguments you will probably train your models for more than few epochs, depending on the way you have configured the packed dataset and the training protocol. Double check that you know and understand what you are doing.
If you don't want to pack your `eval_dataset`, you can pass `eval_packing=False` to the `SFTConfig` init method.

#### Customize your prompts using packed dataset

If your dataset has several fields that you want to combine, for example if the dataset has `question` and `answer` fields and you want to combine them, you can pass a formatting function to the trainer that will take care of that. For example:

```python
def formatting_func(example):
    text = f"### Question: {example['question']}\n ### Answer: {example['answer']}"
    return text

sft_config = SFTConfig(packing=True)
trainer = SFTTrainer(
    "facebook/opt-350m",
    train_dataset=dataset,
    args=sft_config,
    formatting_func=formatting_func
)

trainer.train()
```
You can also customize the [`ConstantLengthDataset`] much more by directly passing the arguments to the [`SFTConfig`] constructor. Please refer to that class' signature for more information.

### Control over the pretrained model

You can directly pass the kwargs of the `from_pretrained()` method to the [`SFTConfig`].


key: _node_content
string_value:
{"id_": "a4361bf0-2970-4170-a615-654fcfe16179", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "8153222b-1942-4a7d-8bb5-f08d0d69cd44", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "ca6421f342e7c70e3ae72a636d9566def8883773a8500c5dd0a9ccf1209607c1", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "acdf97cb-d534-4963-83f3-a94d7b8fef85", "node_type": "1", "metadata": {}, "hash": "188031603ede66d34392a42fb135abd79d4877937972f817a147c230a9197c5d", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 12090, "end_char_idx": 15243, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
To enable the usage of this dataset class, simply pass `packing=True` to the [`SFTConfig`] constructor.

```python
...
sft_config = SFTConfig(packing=True, dataset_text_field="text",)

trainer = SFTTrainer(
    "facebook/opt-350m",
    train_dataset=dataset,
    args=sft_config
)

trainer.train()
```

Note that if you use a packed dataset and if you pass `max_steps` in the training arguments you will probably train your models for more than few epochs, depending on the way you have configured the packed dataset and the training protocol. Double check that you know and understand what you are doing.
If you don't want to pack your `eval_dataset`, you can pass `eval_packing=False` to the `SFTConfig` init method.

#### Customize your prompts using packed dataset

If your dataset has several fields that you want to combine, for example if the dataset has `question` and `answer` fields and you want to combine them, you can pass a formatting function to the trainer that will take care of that. For example:

```python
def formatting_func(example):
    text = f"### Question: {example['question']}\n ### Answer: {example['answer']}"
    return text

sft_config = SFTConfig(packing=True)
trainer = SFTTrainer(
    "facebook/opt-350m",
    train_dataset=dataset,
    args=sft_config,
    formatting_func=formatting_func
)

trainer.train()
```
You can also customize the [`ConstantLengthDataset`] much more by directly passing the arguments to the [`SFTConfig`] constructor. Please refer to that class' signature for more information.

### Control over the pretrained model

You can directly pass the kwargs of the `from_pretrained()` method to the [`SFTConfig`]. For example, if you want to load a model in a different precision, analogous to

```python
model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m", torch_dtype=torch.bfloat16)

...

sft_config = SFTConfig(
    model_init_kwargs={
        "torch_dtype": "bfloat16",
    },
    output_dir="/tmp",
)
trainer = SFTTrainer(
    "facebook/opt-350m",
    train_dataset=dataset,
    args=sft_config,
)

trainer.train()
```
Note that all keyword arguments of `from_pretrained()` are supported.

### Training adapters

We also support tight integration with 🤗 PEFT library so that any user can conveniently train adapters and share them on the Hub instead of training the entire model

```python
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer
from peft import LoraConfig

dataset = load_dataset("imdb", split="train")

peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

trainer = SFTTrainer(
    "EleutherAI/gpt-neo-125m",
    train_dataset=dataset,
    args=SFTConfig(output_dir="/tmp"),
    peft_config=peft_config
)

trainer.train()
```

You can also continue training your `PeftModel`. For that, first load a `PeftModel` outside `SFTTrainer` and pass it directly to the trainer without the `peft_config` argument being passed.

### Training adapters with base 8 bit models

For that, you need to first load your 8 bit model outside the Trainer and pass a `PeftConfig` to the trainer.


key: _node_content
string_value:
{"id_": "acdf97cb-d534-4963-83f3-a94d7b8fef85", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "a4361bf0-2970-4170-a615-654fcfe16179", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "26a18d251213aae288fcfdc7bf9a78954673b2284973ecb12bfdfca5e3bd88b9", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "2d4be494-43d7-4b7a-802a-7423e5a47a6b", "node_type": "1", "metadata": {}, "hash": "baa0b8ad473b1b34dfde3b9e1112276256e77efec49a6c6fdd03065991c5b353", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 13758, "end_char_idx": 16374, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
For example, if you want to load a model in a different precision, analogous to

```python
model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m", torch_dtype=torch.bfloat16)

...

sft_config = SFTConfig(
    model_init_kwargs={
        "torch_dtype": "bfloat16",
    },
    output_dir="/tmp",
)
trainer = SFTTrainer(
    "facebook/opt-350m",
    train_dataset=dataset,
    args=sft_config,
)

trainer.train()
```
Note that all keyword arguments of `from_pretrained()` are supported.

### Training adapters

We also support tight integration with 🤗 PEFT library so that any user can conveniently train adapters and share them on the Hub instead of training the entire model

```python
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer
from peft import LoraConfig

dataset = load_dataset("imdb", split="train")

peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

trainer = SFTTrainer(
    "EleutherAI/gpt-neo-125m",
    train_dataset=dataset,
    args=SFTConfig(output_dir="/tmp"),
    peft_config=peft_config
)

trainer.train()
```

You can also continue training your `PeftModel`. For that, first load a `PeftModel` outside `SFTTrainer` and pass it directly to the trainer without the `peft_config` argument being passed.

### Training adapters with base 8 bit models

For that, you need to first load your 8 bit model outside the Trainer and pass a `PeftConfig` to the trainer. For example:

```python
...

peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = AutoModelForCausalLM.from_pretrained(
    "EleutherAI/gpt-neo-125m",
    load_in_8bit=True,
    device_map="auto",
)

trainer = SFTTrainer(
    model,
    train_dataset=dataset,
    args=SFTConfig(),
    peft_config=peft_config,
)

trainer.train()
```

## Using Flash Attention and Flash Attention 2

You can benefit from Flash Attention 1 & 2 using SFTTrainer out of the box with minimal changes of code.
First, to make sure you have all the latest features from transformers, install transformers from source

```bash
pip install -U git+https://github.com/huggingface/transformers.git
```

Note that Flash Attention only works on GPU now and under half-precision regime (when using adapters, base model loaded in half-precision)
Note also both features are perfectly compatible with other tools such as quantization.

### Using Flash-Attention 1

For Flash Attention 1 you can use the `BetterTransformer` API and force-dispatch the API to use Flash Attention kernel.


key: _node_content
string_value:
{"id_": "2d4be494-43d7-4b7a-802a-7423e5a47a6b", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "acdf97cb-d534-4963-83f3-a94d7b8fef85", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "751bdf12628edb2fc9b19c281d5dfddfb6fc7f42e0abfd0270f87716b7cb1636", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "6652b3ff-fedf-439f-9516-fad4140e517a", "node_type": "1", "metadata": {}, "hash": "d7c9084241c5eb34060a9b9441f986f61de2e72daaa6b886b23e20bdf7d443b8", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 14946, "end_char_idx": 17290, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
For that, first load a `PeftModel` outside `SFTTrainer` and pass it directly to the trainer without the `peft_config` argument being passed.

### Training adapters with base 8 bit models

For that, you need to first load your 8 bit model outside the Trainer and pass a `PeftConfig` to the trainer. For example:

```python
...

peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = AutoModelForCausalLM.from_pretrained(
    "EleutherAI/gpt-neo-125m",
    load_in_8bit=True,
    device_map="auto",
)

trainer = SFTTrainer(
    model,
    train_dataset=dataset,
    args=SFTConfig(),
    peft_config=peft_config,
)

trainer.train()
```

## Using Flash Attention and Flash Attention 2

You can benefit from Flash Attention 1 & 2 using SFTTrainer out of the box with minimal changes of code.
First, to make sure you have all the latest features from transformers, install transformers from source

```bash
pip install -U git+https://github.com/huggingface/transformers.git
```

Note that Flash Attention only works on GPU now and under half-precision regime (when using adapters, base model loaded in half-precision)
Note also both features are perfectly compatible with other tools such as quantization.

### Using Flash-Attention 1

For Flash Attention 1 you can use the `BetterTransformer` API and force-dispatch the API to use Flash Attention kernel. First, install the latest optimum package:

```bash
pip install -U optimum
```

Once you have loaded your model, wrap the `trainer.train()` call under the `with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):` context manager:

```diff
...

+ with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
    trainer.train()
```

Note that you cannot train your model using Flash Attention 1 on an arbitrary dataset as `torch.scaled_dot_product_attention` does not support training with padding tokens if you use Flash Attention kernels. Therefore you can only use that feature with `packing=True`. If your dataset contains padding tokens, consider switching to Flash Attention 2 integration.

Below are some numbers you can get in terms of speedup and memory efficiency, using Flash Attention 1, on a single NVIDIA-T4 16GB.


key: _node_content
string_value:
{"id_": "6652b3ff-fedf-439f-9516-fad4140e517a", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "2d4be494-43d7-4b7a-802a-7423e5a47a6b", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "e378e35e5e253a85532b36850becb3e3b100d3b5688aec92cb69861f62b84cfd", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "8b333075-1db9-4b89-9a78-63b7832ecd97", "node_type": "1", "metadata": {}, "hash": "01a0f74a3d8a4fc7b7fcfddfd443b68248345c4a1c56bbe9ddc8d551c4c24c99", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 15812, "end_char_idx": 18724, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
First, to make sure you have all the latest features from transformers, install transformers from source

```bash
pip install -U git+https://github.com/huggingface/transformers.git
```

Note that Flash Attention only works on GPU now and under half-precision regime (when using adapters, base model loaded in half-precision)
Note also both features are perfectly compatible with other tools such as quantization.

### Using Flash-Attention 1

For Flash Attention 1 you can use the `BetterTransformer` API and force-dispatch the API to use Flash Attention kernel. First, install the latest optimum package:

```bash
pip install -U optimum
```

Once you have loaded your model, wrap the `trainer.train()` call under the `with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):` context manager:

```diff
...

+ with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
    trainer.train()
```

Note that you cannot train your model using Flash Attention 1 on an arbitrary dataset as `torch.scaled_dot_product_attention` does not support training with padding tokens if you use Flash Attention kernels. Therefore you can only use that feature with `packing=True`. If your dataset contains padding tokens, consider switching to Flash Attention 2 integration.

Below are some numbers you can get in terms of speedup and memory efficiency, using Flash Attention 1, on a single NVIDIA-T4 16GB.

| use_flash_attn_1 | model_name        | max_seq_len | batch_size | time per training step |
| ---------------- | ----------------- | ----------- | ---------- | ---------------------- |
| x                | facebook/opt-350m | 2048        | 8          | ~59.1s                 |
|                  | facebook/opt-350m | 2048        | 8          | **OOM**                |
| x                | facebook/opt-350m | 2048        | 4          | ~30.3s                 |
|                  | facebook/opt-350m | 2048        | 4          | ~148.9s                |

### Using Flash Attention-2

To use Flash Attention 2, first install the latest `flash-attn` package:

```bash
pip install -U flash-attn
```

And add `attn_implementation="flash_attention_2"` when calling `from_pretrained`:

```python
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    load_in_4bit=True,
    attn_implementation="flash_attention_2"
)
```

If you don't use quantization, make sure your model is loaded in half-precision and dispatch your model on a supported GPU device.
After loading your model, you can either train it as it is, or attach adapters and train adapters on it in case your model is quantized.

In contrast to Flash Attention 1, the integration makes it possible to train your model on an arbitrary dataset that also includes padding tokens.


### Using model creation utility

We included a utility function to create your model.


key: _node_content
string_value:
{"id_": "8b333075-1db9-4b89-9a78-63b7832ecd97", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "6652b3ff-fedf-439f-9516-fad4140e517a", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "9a1ec360d9c050b8fabbc6b7762acedc2e09c021042abb341a427ce56d4886fd", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "dbe0b07c-6ab9-4a73-a81d-2e38d094c7b1", "node_type": "1", "metadata": {}, "hash": "50ad2bb8ed43e6530409ebd84c10bb5968b1c8703835350ad990ce89d6b9f2c9", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 17004, "end_char_idx": 20162, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
Therefore you can only use that feature with `packing=True`. If your dataset contains padding tokens, consider switching to Flash Attention 2 integration.

Below are some numbers you can get in terms of speedup and memory efficiency, using Flash Attention 1, on a single NVIDIA-T4 16GB.

| use_flash_attn_1 | model_name        | max_seq_len | batch_size | time per training step |
| ---------------- | ----------------- | ----------- | ---------- | ---------------------- |
| x                | facebook/opt-350m | 2048        | 8          | ~59.1s                 |
|                  | facebook/opt-350m | 2048        | 8          | **OOM**                |
| x                | facebook/opt-350m | 2048        | 4          | ~30.3s                 |
|                  | facebook/opt-350m | 2048        | 4          | ~148.9s                |

### Using Flash Attention-2

To use Flash Attention 2, first install the latest `flash-attn` package:

```bash
pip install -U flash-attn
```

And add `attn_implementation="flash_attention_2"` when calling `from_pretrained`:

```python
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    load_in_4bit=True,
    attn_implementation="flash_attention_2"
)
```

If you don't use quantization, make sure your model is loaded in half-precision and dispatch your model on a supported GPU device.
After loading your model, you can either train it as it is, or attach adapters and train adapters on it in case your model is quantized.

In contrast to Flash Attention 1, the integration makes it possible to train your model on an arbitrary dataset that also includes padding tokens.


### Using model creation utility

We included a utility function to create your model.

[[autodoc]] ModelConfig

```python
from trl import ModelConfig, SFTTrainer, get_kbit_device_map, get_peft_config, get_quantization_config
model_config = ModelConfig(
    model_name_or_path="facebook/opt-350m"
    attn_implementation=None, # or "flash_attention_2"
)
torch_dtype = (
    model_config.torch_dtype
    if model_config.torch_dtype in ["auto", None]
    else getattr(torch, model_config.torch_dtype)
)
quantization_config = get_quantization_config(model_config)
model_kwargs = dict(
    revision=model_config.model_revision,
    trust_remote_code=model_config.trust_remote_code,
    attn_implementation=model_config.attn_implementation,
    torch_dtype=torch_dtype,
    use_cache=False if training_args.gradient_checkpointing else True,
    device_map=get_kbit_device_map() if quantization_config is not None else None,
    quantization_config=quantization_config,
)
model = AutoModelForCausalLM.from_pretrained(model_config.model_name_or_path, **model_kwargs)
trainer = SFTTrainer(
    ...,
    model=model_config.model_name_or_path,
    peft_config=get_peft_config(model_config),
)
```

### Enhance the model's performances using NEFTune

NEFTune is a technique to boost the performance of chat models and was introduced by the paper ["NEFTune: Noisy Embeddings Improve Instruction Finetuning"](https://huggingface.co/papers/2310.05914) from Jain et al. it consists of adding noise to the embedding vectors during training.


key: _node_content
string_value:
{"id_": "dbe0b07c-6ab9-4a73-a81d-2e38d094c7b1", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "8b333075-1db9-4b89-9a78-63b7832ecd97", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "e694047d942a0f3cdf302b9730bd1a14cc7ec963be6452b90569e838482548c6", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "51522ce3-b205-47c1-be53-086b8bd62ef3", "node_type": "1", "metadata": {}, "hash": "a9c7c6d8ce8983c90c18d3fa7d38b1a9b9f954e52174c6a3d0c4710746720f96", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 18489, "end_char_idx": 21603, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
In contrast to Flash Attention 1, the integration makes it possible to train your model on an arbitrary dataset that also includes padding tokens.


### Using model creation utility

We included a utility function to create your model.

[[autodoc]] ModelConfig

```python
from trl import ModelConfig, SFTTrainer, get_kbit_device_map, get_peft_config, get_quantization_config
model_config = ModelConfig(
    model_name_or_path="facebook/opt-350m"
    attn_implementation=None, # or "flash_attention_2"
)
torch_dtype = (
    model_config.torch_dtype
    if model_config.torch_dtype in ["auto", None]
    else getattr(torch, model_config.torch_dtype)
)
quantization_config = get_quantization_config(model_config)
model_kwargs = dict(
    revision=model_config.model_revision,
    trust_remote_code=model_config.trust_remote_code,
    attn_implementation=model_config.attn_implementation,
    torch_dtype=torch_dtype,
    use_cache=False if training_args.gradient_checkpointing else True,
    device_map=get_kbit_device_map() if quantization_config is not None else None,
    quantization_config=quantization_config,
)
model = AutoModelForCausalLM.from_pretrained(model_config.model_name_or_path, **model_kwargs)
trainer = SFTTrainer(
    ...,
    model=model_config.model_name_or_path,
    peft_config=get_peft_config(model_config),
)
```

### Enhance the model's performances using NEFTune

NEFTune is a technique to boost the performance of chat models and was introduced by the paper ["NEFTune: Noisy Embeddings Improve Instruction Finetuning"](https://huggingface.co/papers/2310.05914) from Jain et al. it consists of adding noise to the embedding vectors during training. According to the abstract of the paper:

>  Standard finetuning of LLaMA-2-7B using Alpaca achieves 29.79% on AlpacaEval, which rises to 64.69% using noisy embeddings. NEFTune also improves over strong baselines on modern instruction datasets. Models trained with Evol-Instruct see a 10% improvement, with ShareGPT an 8% improvement, and with OpenPlatypus an 8% improvement. Even powerful models further refined with RLHF such as LLaMA-2-Chat benefit from additional training with NEFTune.

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/neft-screenshot.png">
</div>

To use it in `SFTTrainer` simply pass `neftune_noise_alpha` when creating your `SFTConfig` instance. Note that to avoid any surprising behaviour, NEFTune is disabled after training to retrieve back the original behaviour of the embedding layer.

```python
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer

dataset = load_dataset("imdb", split="train")

sft_config = SFTConfig(
    neftune_noise_alpha=5,
)
trainer = SFTTrainer(
    "facebook/opt-350m",
    train_dataset=dataset,
    args=sft_config,
)
trainer.train()
```

We have tested NEFTune by training `mistralai/Mistral-7B-v0.1` on the [OpenAssistant dataset](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) and validated that using NEFTune led to a performance boost of ~25% on MT Bench.


key: _node_content
string_value:
{"id_": "51522ce3-b205-47c1-be53-086b8bd62ef3", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "dbe0b07c-6ab9-4a73-a81d-2e38d094c7b1", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "76f2fabf6ff5feb9ff373b94c6d7569302122941ecd091785ea561c4ea7aa235", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "68483bd2-751b-439a-994d-16af11f0edae", "node_type": "1", "metadata": {}, "hash": "e4be16f78cf43baba0e7ebf3fc6cc616a955711dd487099a1d964cf7173d72a2", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 20093, "end_char_idx": 22339, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
it consists of adding noise to the embedding vectors during training. According to the abstract of the paper:

>  Standard finetuning of LLaMA-2-7B using Alpaca achieves 29.79% on AlpacaEval, which rises to 64.69% using noisy embeddings. NEFTune also improves over strong baselines on modern instruction datasets. Models trained with Evol-Instruct see a 10% improvement, with ShareGPT an 8% improvement, and with OpenPlatypus an 8% improvement. Even powerful models further refined with RLHF such as LLaMA-2-Chat benefit from additional training with NEFTune.

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/neft-screenshot.png">
</div>

To use it in `SFTTrainer` simply pass `neftune_noise_alpha` when creating your `SFTConfig` instance. Note that to avoid any surprising behaviour, NEFTune is disabled after training to retrieve back the original behaviour of the embedding layer.

```python
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer

dataset = load_dataset("imdb", split="train")

sft_config = SFTConfig(
    neftune_noise_alpha=5,
)
trainer = SFTTrainer(
    "facebook/opt-350m",
    train_dataset=dataset,
    args=sft_config,
)
trainer.train()
```

We have tested NEFTune by training `mistralai/Mistral-7B-v0.1` on the [OpenAssistant dataset](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) and validated that using NEFTune led to a performance boost of ~25% on MT Bench.

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl-neftune-mistral-7b.png">
</div>

Note however, that the amount of performance gain is _dataset dependent_ and in particular, applying NEFTune on synthetic datasets like [UltraChat](https://huggingface.co/datasets/stingning/ultrachat) typically produces smaller gains.

### Accelerate fine-tuning 2x using `unsloth`

You can further accelerate QLoRA / LoRA (2x faster, 60% less memory) using the [`unsloth`](https://github.com/unslothai/unsloth) library that is fully compatible with `SFTTrainer`. Currently `unsloth` supports only Llama (Yi, TinyLlama, Qwen, Deepseek etc) and Mistral architectures.


key: _node_content
string_value:
{"id_": "68483bd2-751b-439a-994d-16af11f0edae", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "51522ce3-b205-47c1-be53-086b8bd62ef3", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "e62e398576d49ff5bb8df5d8e70be2f077deb3de68a8391cab9b7b395dc68877", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "5b7f22a0-b4f8-43ee-a644-8487746e088f", "node_type": "1", "metadata": {}, "hash": "15a3ccf58848003d1120e48953ba8d175514058488fad4c4eea6b50b65d4ccb8", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 20916, "end_char_idx": 23571, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
Note that to avoid any surprising behaviour, NEFTune is disabled after training to retrieve back the original behaviour of the embedding layer.

```python
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer

dataset = load_dataset("imdb", split="train")

sft_config = SFTConfig(
    neftune_noise_alpha=5,
)
trainer = SFTTrainer(
    "facebook/opt-350m",
    train_dataset=dataset,
    args=sft_config,
)
trainer.train()
```

We have tested NEFTune by training `mistralai/Mistral-7B-v0.1` on the [OpenAssistant dataset](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) and validated that using NEFTune led to a performance boost of ~25% on MT Bench.

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl-neftune-mistral-7b.png">
</div>

Note however, that the amount of performance gain is _dataset dependent_ and in particular, applying NEFTune on synthetic datasets like [UltraChat](https://huggingface.co/datasets/stingning/ultrachat) typically produces smaller gains.

### Accelerate fine-tuning 2x using `unsloth`

You can further accelerate QLoRA / LoRA (2x faster, 60% less memory) using the [`unsloth`](https://github.com/unslothai/unsloth) library that is fully compatible with `SFTTrainer`. Currently `unsloth` supports only Llama (Yi, TinyLlama, Qwen, Deepseek etc) and Mistral architectures. Some benchmarks on 1x A100 listed below:

| 1 A100 40GB     | Dataset   | 🤗   | 🤗 + Flash Attention 2 | 🦥 Unsloth | 🦥 VRAM saved |
| --------------- | --------- | --- | --------------------- | --------- | ------------ |
| Code Llama 34b  | Slim Orca | 1x  | 1.01x                 | **1.94x** | -22.7%       |
| Llama-2 7b      | Slim Orca | 1x  | 0.96x                 | **1.87x** | -39.3%       |
| Mistral 7b      | Slim Orca | 1x  | 1.17x                 | **1.88x** | -65.9%       |
| Tiny Llama 1.1b | Alpaca    | 1x  | 1.55x                 | **2.74x** | -57.8%       |

First install `unsloth` according to the [official documentation](https://github.com/unslothai/unsloth). Once installed, you can incorporate unsloth into your workflow in a very simple manner; instead of loading `AutoModelForCausalLM`, you just need to load a `FastLanguageModel` as follows:

```python
import torch
from trl import SFTConfig, SFTTrainer
from unsloth import FastLanguageModel

max_seq_length = 2048 # Supports automatic RoPE Scaling, so choose any number

# Load model
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/mistral-7b",
    max_seq_length=max_seq_length,
    dtype=None,  # None for auto detection.


key: _node_content
string_value:
{"id_": "5b7f22a0-b4f8-43ee-a644-8487746e088f", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "68483bd2-751b-439a-994d-16af11f0edae", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "6affe4b8b11716772ae3988b733b8aab87d52504318e6387a8e6caa07ac0adaf", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "5cfee872-34a3-4996-aec6-78f0bc399e02", "node_type": "1", "metadata": {}, "hash": "16b6d30a6d2c184efc21f2100de483797304561aa91fa724643ca25d10a9073f", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 22340, "end_char_idx": 25106, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
Some benchmarks on 1x A100 listed below:

| 1 A100 40GB     | Dataset   | 🤗   | 🤗 + Flash Attention 2 | 🦥 Unsloth | 🦥 VRAM saved |
| --------------- | --------- | --- | --------------------- | --------- | ------------ |
| Code Llama 34b  | Slim Orca | 1x  | 1.01x                 | **1.94x** | -22.7%       |
| Llama-2 7b      | Slim Orca | 1x  | 0.96x                 | **1.87x** | -39.3%       |
| Mistral 7b      | Slim Orca | 1x  | 1.17x                 | **1.88x** | -65.9%       |
| Tiny Llama 1.1b | Alpaca    | 1x  | 1.55x                 | **2.74x** | -57.8%       |

First install `unsloth` according to the [official documentation](https://github.com/unslothai/unsloth). Once installed, you can incorporate unsloth into your workflow in a very simple manner; instead of loading `AutoModelForCausalLM`, you just need to load a `FastLanguageModel` as follows:

```python
import torch
from trl import SFTConfig, SFTTrainer
from unsloth import FastLanguageModel

max_seq_length = 2048 # Supports automatic RoPE Scaling, so choose any number

# Load model
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/mistral-7b",
    max_seq_length=max_seq_length,
    dtype=None,  # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
    load_in_4bit=True,  # Use 4bit quantization to reduce memory usage. Can be False
    # token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
)

# Do model patching and add fast LoRA weights
model = FastLanguageModel.get_peft_model(
    model,
    r=16,
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    ],
    lora_alpha=16,
    lora_dropout=0,  # Dropout = 0 is currently optimized
    bias="none",  # Bias = "none" is currently optimized
    use_gradient_checkpointing=True,
    random_state=3407,
)

args = SFTConfig(
    output_dir="./output",
    max_seq_length=max_seq_length,
    dataset_text_field="text",
)

trainer = SFTTrainer(
    model=model,
    args=args,
    train_dataset=dataset,
)
trainer.train()
```

The saved model is fully compatible with Hugging Face's transformers library. Learn more about unsloth in their [official repository](https://github.com/unslothai/unsloth).

## Best practices

Pay attention to the following best practices when training a model with that trainer:

- [`SFTTrainer`] always pads by default the sequences to the `max_seq_length` argument of the [`SFTTrainer`]. If none is passed, the trainer will retrieve that value from the tokenizer. Some tokenizers do not provide a default value, so there is a check to retrieve the minimum between 2048 and that value. Make sure to check it before training.


key: _node_content
string_value:
{"id_": "5cfee872-34a3-4996-aec6-78f0bc399e02", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "5b7f22a0-b4f8-43ee-a644-8487746e088f", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "cea1b043be7f15de241988b38779e11ac4b7343a919b4e0e92bd7cc5cdc1b0ce", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "91c553ac-e9f5-41a3-b14f-1875c83c04f8", "node_type": "1", "metadata": {}, "hash": "69becf898bdfe2aae9c4312adeb4028a9963ade0d148dc9e0a135b1efac30af6", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 23572, "end_char_idx": 26308, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
Float16 for Tesla T4, V100, Bfloat16 for Ampere+
    load_in_4bit=True,  # Use 4bit quantization to reduce memory usage. Can be False
    # token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
)

# Do model patching and add fast LoRA weights
model = FastLanguageModel.get_peft_model(
    model,
    r=16,
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    ],
    lora_alpha=16,
    lora_dropout=0,  # Dropout = 0 is currently optimized
    bias="none",  # Bias = "none" is currently optimized
    use_gradient_checkpointing=True,
    random_state=3407,
)

args = SFTConfig(
    output_dir="./output",
    max_seq_length=max_seq_length,
    dataset_text_field="text",
)

trainer = SFTTrainer(
    model=model,
    args=args,
    train_dataset=dataset,
)
trainer.train()
```

The saved model is fully compatible with Hugging Face's transformers library. Learn more about unsloth in their [official repository](https://github.com/unslothai/unsloth).

## Best practices

Pay attention to the following best practices when training a model with that trainer:

- [`SFTTrainer`] always pads by default the sequences to the `max_seq_length` argument of the [`SFTTrainer`]. If none is passed, the trainer will retrieve that value from the tokenizer. Some tokenizers do not provide a default value, so there is a check to retrieve the minimum between 2048 and that value. Make sure to check it before training.
- For training adapters in 8bit, you might need to tweak the arguments of the `prepare_model_for_kbit_training` method from PEFT, hence we advise users to use `prepare_in_int8_kwargs` field, or create the `PeftModel` outside the [`SFTTrainer`] and pass it.
- For a more memory-efficient training using adapters, you can load the base model in 8bit, for that simply add `load_in_8bit` argument when creating the [`SFTTrainer`], or create a base model in 8bit outside the trainer and pass it.
- If you create a model outside the trainer, make sure to not pass to the trainer any additional keyword arguments that are relative to `from_pretrained()` method.

## Multi-GPU Training

Trainer (and thus SFTTrainer) supports multi-GPU training. If you run your script with `python script.py` it will default to using DP as the strategy, which may be [slower than expected](https://github.com/huggingface/trl/issues/1303). To use DDP (which is generally recommended, see [here](https://huggingface.co/docs/transformers/en/perf_train_gpu_many?select-gpu=Accelerate#data-parallelism) for more info) you must launch the script with `python -m torch.distributed.launch script.py` or `accelerate launch script.py`.


key: _node_content
string_value:
{"id_": "91c553ac-e9f5-41a3-b14f-1875c83c04f8", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "5cfee872-34a3-4996-aec6-78f0bc399e02", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "1cf6711a2f59a3e21020816a10aecbe71fc5fee4a7afe2b6324b2dc033dbed0a", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "676735f0-779b-4db8-868c-d930130ec3bb", "node_type": "1", "metadata": {}, "hash": "f9d4ab76c5d403ac01de20b5db695f9fd813635544dec04dd7968596ba247c1e", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 24871, "end_char_idx": 28020, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
If none is passed, the trainer will retrieve that value from the tokenizer. Some tokenizers do not provide a default value, so there is a check to retrieve the minimum between 2048 and that value. Make sure to check it before training.
- For training adapters in 8bit, you might need to tweak the arguments of the `prepare_model_for_kbit_training` method from PEFT, hence we advise users to use `prepare_in_int8_kwargs` field, or create the `PeftModel` outside the [`SFTTrainer`] and pass it.
- For a more memory-efficient training using adapters, you can load the base model in 8bit, for that simply add `load_in_8bit` argument when creating the [`SFTTrainer`], or create a base model in 8bit outside the trainer and pass it.
- If you create a model outside the trainer, make sure to not pass to the trainer any additional keyword arguments that are relative to `from_pretrained()` method.

## Multi-GPU Training

Trainer (and thus SFTTrainer) supports multi-GPU training. If you run your script with `python script.py` it will default to using DP as the strategy, which may be [slower than expected](https://github.com/huggingface/trl/issues/1303). To use DDP (which is generally recommended, see [here](https://huggingface.co/docs/transformers/en/perf_train_gpu_many?select-gpu=Accelerate#data-parallelism) for more info) you must launch the script with `python -m torch.distributed.launch script.py` or `accelerate launch script.py`. For DDP to work you must also check the following:
- If you're using gradient_checkpointing, add the following to the TrainingArguments: `gradient_checkpointing_kwargs={'use_reentrant':False}` (more info [here](https://github.com/huggingface/transformers/issues/26969)
- Ensure that the model is placed on the correct device:
```python
from accelerate import PartialState
device_string = PartialState().process_index
model = AutoModelForCausalLM.from_pretrained(
     ...
    device_map={'':device_string}
)
```

## GPTQ Conversion

You may experience some issues with GPTQ Quantization after completing training. Lowering `gradient_accumulation_steps` to `4` will resolve most issues during the quantization process to GPTQ format.

## Extending `SFTTrainer` for Vision Language Models

`SFTTrainer` does not inherently support vision-language data. However, we provide a guide on how to tweak the trainer to support vision-language data. Specifically, you need to use a custom data collator that is compatible with vision-language data. This guide outlines the steps to make these adjustments. For a concrete example, refer to the script [`examples/scripts/vsft_llava.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/vsft_llava.py) which demonstrates how to fine-tune the LLaVA 1.5 model on the [HuggingFaceH4/llava-instruct-mix-vsft](https://huggingface.co/datasets/HuggingFaceH4/llava-instruct-mix-vsft) dataset.

### Preparing the Data

The data format is flexible, provided it is compatible with the custom collator that we will define later. A common approach is to use conversational data. Given that the data includes both text and images, the format needs to be adjusted accordingly.


key: _node_content
string_value:
{"id_": "676735f0-779b-4db8-868c-d930130ec3bb", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "91c553ac-e9f5-41a3-b14f-1875c83c04f8", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "05591c0404fab6c7eba7bc478419db37b6881c0b85209bad15bb043f50870ac9", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "cf6cb03d-c93a-476f-a2a5-107908ae6507", "node_type": "1", "metadata": {}, "hash": "5872c6f8c848ca5cbd08574b026f0eb754fa77a62c1dbee55ba0932fc1ec8a61", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 26923, "end_char_idx": 29792, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
Lowering `gradient_accumulation_steps` to `4` will resolve most issues during the quantization process to GPTQ format.

## Extending `SFTTrainer` for Vision Language Models

`SFTTrainer` does not inherently support vision-language data. However, we provide a guide on how to tweak the trainer to support vision-language data. Specifically, you need to use a custom data collator that is compatible with vision-language data. This guide outlines the steps to make these adjustments. For a concrete example, refer to the script [`examples/scripts/vsft_llava.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/vsft_llava.py) which demonstrates how to fine-tune the LLaVA 1.5 model on the [HuggingFaceH4/llava-instruct-mix-vsft](https://huggingface.co/datasets/HuggingFaceH4/llava-instruct-mix-vsft) dataset.

### Preparing the Data

The data format is flexible, provided it is compatible with the custom collator that we will define later. A common approach is to use conversational data. Given that the data includes both text and images, the format needs to be adjusted accordingly. Below is an example of a conversational data format involving both text and images:

```python
images = ["obama.png"]
messages = [
    {
        "role": "user",
        "content": [
            {"type": "text", "text": "Who is this?"},
            {"type": "image"}
        ]
    },
    {
        "role": "assistant",
        "content": [
            {"type": "text", "text": "Barack Obama"}
        ]
    },
    {
        "role": "user",
        "content": [
            {"type": "text", "text": "What is he famous for?"}
        ]
    },
    {
        "role": "assistant",
        "content": [
            {"type": "text", "text": "He is the 44th President of the United States."}
        ]
    }
]
```

To illustrate how this data format will be processed using the LLaVA model, you can use the following code:

```python
from transformers import AutoProcessor

processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
print(processor.apply_chat_template(messages, tokenize=False))
```

The output will be formatted as follows:

```txt
Who is this? ASSISTANT: Barack Obama USER: What is he famous for? ASSISTANT: He is the 44th President of the United States. 
```

<iframe src="https://huggingface.co/datasets/HuggingFaceH4/llava-instruct-mix-vsft/embed/viewer/default/train" frameborder="0" width="100%" height="560px"></iframe>


### A custom collator for processing multi-modal data

Unlike the default behavior of `SFTTrainer`, processing multi-modal data is done on the fly during the data collation process. To do this, you need to define a custom collator that processes both the text and images. This collator must take a list of examples as input (see the previous section for an example of the data format) and return a batch of processed data.


key: _node_content
string_value:
{"id_": "cf6cb03d-c93a-476f-a2a5-107908ae6507", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "676735f0-779b-4db8-868c-d930130ec3bb", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "ac0ae4adb51636ac41488aac5771fe3a0e2f243f8ae27aa70d0a244769237475", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "248d91f3-fd29-4cf4-8812-4020e8ec8b65", "node_type": "1", "metadata": {}, "hash": "b5dbb979e1f532cf5e354845e361c9e313006f0a2137ab7e12ba8fa8c3f1180a", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 28253, "end_char_idx": 31366, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
"},
            {"type": "image"}
        ]
    },
    {
        "role": "assistant",
        "content": [
            {"type": "text", "text": "Barack Obama"}
        ]
    },
    {
        "role": "user",
        "content": [
            {"type": "text", "text": "What is he famous for?"}
        ]
    },
    {
        "role": "assistant",
        "content": [
            {"type": "text", "text": "He is the 44th President of the United States."}
        ]
    }
]
```

To illustrate how this data format will be processed using the LLaVA model, you can use the following code:

```python
from transformers import AutoProcessor

processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
print(processor.apply_chat_template(messages, tokenize=False))
```

The output will be formatted as follows:

```txt
Who is this? ASSISTANT: Barack Obama USER: What is he famous for? ASSISTANT: He is the 44th President of the United States. 
```

<iframe src="https://huggingface.co/datasets/HuggingFaceH4/llava-instruct-mix-vsft/embed/viewer/default/train" frameborder="0" width="100%" height="560px"></iframe>


### A custom collator for processing multi-modal data

Unlike the default behavior of `SFTTrainer`, processing multi-modal data is done on the fly during the data collation process. To do this, you need to define a custom collator that processes both the text and images. This collator must take a list of examples as input (see the previous section for an example of the data format) and return a batch of processed data. Below is an example of such a collator:

```python
def collate_fn(examples):
    # Get the texts and images, and apply the chat template
    texts = [processor.apply_chat_template(example["messages"], tokenize=False) for example in examples]
    images = [example["images"][0] for example in examples]

    # Tokenize the texts and process the images
    batch = processor(texts, images, return_tensors="pt", padding=True)

    # The labels are the input_ids, and we mask the padding tokens in the loss computation
    labels = batch["input_ids"].clone()
    labels[labels == processor.tokenizer.pad_token_id] = -100
    batch["labels"] = labels

    return batch
```

We can verify that the collator works as expected by running the following code:

```python
from datasets import load_dataset

dataset = load_dataset("HuggingFaceH4/llava-instruct-mix-vsft", split="train")
examples = [dataset[0], dataset[1]]  # Just two examples for the sake of the example
collated_data = collate_fn(examples)
print(collated_data.keys())  # dict_keys(['input_ids', 'attention_mask', 'pixel_values', 'labels'])
```

### Training the vision-language model

Now that we have prepared the data and defined the collator, we can proceed with training the model. To ensure that the data is not processed as text-only, we need to set a couple of arguments in the `SFTConfig`, specifically `dataset_text_field` and `remove_unused_columns`. We also need to set `skip_prepare_dataset` to `True` to avoid the default processing of the dataset. Below is an example of how to set up the `SFTTrainer`.


key: _node_content
string_value:
{"id_": "248d91f3-fd29-4cf4-8812-4020e8ec8b65", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "cf6cb03d-c93a-476f-a2a5-107908ae6507", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "e71b25446d076ebe99d9c0a71fe64b07a2ea5430cbc18993037b6496d5ac3773", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "ad816c64-14d7-4bbd-b3e1-b1d1d7dd5998", "node_type": "1", "metadata": {}, "hash": "0377f2a529f33e386ed29da47e143213e0af65dcf414097d7bd59b1412725f24", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 29793, "end_char_idx": 32782, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
Below is an example of such a collator:

```python
def collate_fn(examples):
    # Get the texts and images, and apply the chat template
    texts = [processor.apply_chat_template(example["messages"], tokenize=False) for example in examples]
    images = [example["images"][0] for example in examples]

    # Tokenize the texts and process the images
    batch = processor(texts, images, return_tensors="pt", padding=True)

    # The labels are the input_ids, and we mask the padding tokens in the loss computation
    labels = batch["input_ids"].clone()
    labels[labels == processor.tokenizer.pad_token_id] = -100
    batch["labels"] = labels

    return batch
```

We can verify that the collator works as expected by running the following code:

```python
from datasets import load_dataset

dataset = load_dataset("HuggingFaceH4/llava-instruct-mix-vsft", split="train")
examples = [dataset[0], dataset[1]]  # Just two examples for the sake of the example
collated_data = collate_fn(examples)
print(collated_data.keys())  # dict_keys(['input_ids', 'attention_mask', 'pixel_values', 'labels'])
```

### Training the vision-language model

Now that we have prepared the data and defined the collator, we can proceed with training the model. To ensure that the data is not processed as text-only, we need to set a couple of arguments in the `SFTConfig`, specifically `dataset_text_field` and `remove_unused_columns`. We also need to set `skip_prepare_dataset` to `True` to avoid the default processing of the dataset. Below is an example of how to set up the `SFTTrainer`.

```python
args.dataset_text_field = ""  # needs a dummy field
args.remove_unused_columns = False
args.dataset_kwargs = {"skip_prepare_dataset": True}

trainer = SFTTrainer(
    model=model,
    args=args,
    data_collator=collate_fn,
    train_dataset=train_dataset,
    tokenizer=processor.tokenizer,
)
```

A full example of training LLaVa 1.5 on the [HuggingFaceH4/llava-instruct-mix-vsft](https://huggingface.co/datasets/HuggingFaceH4/llava-instruct-mix-vsft) dataset can be found in the script [`examples/scripts/vsft_llava.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/vsft_llava.py).

- [Experiment tracking](https://wandb.ai/huggingface/trl/runs/2b2c5l7s)
- [Trained model](https://huggingface.co/HuggingFaceH4/sft-llava-1.5-7b-hf)

## SFTTrainer

[[autodoc]] SFTTrainer

## SFTConfig

[[autodoc]] SFTConfig

## Datasets

In the SFTTrainer we smartly support `datasets.IterableDataset` in addition to other style datasets. This is useful if you are using large corpora that you do not want to save all to disk. The data will be tokenized and processed on the fly, even when packing is enabled.

Additionally, in the SFTTrainer, we support pre-tokenized datasets if they are `datasets.Dataset` or `datasets.IterableDataset`. In other words, if such a dataset has a column of `input_ids`, no further processing (tokenization or packing) will be done, and the dataset will be used as-is.


key: _node_content
string_value:
{"id_": "ad816c64-14d7-4bbd-b3e1-b1d1d7dd5998", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "ae2c8951-dc4a-4f1e-bb99-8b8f9051d504", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "7923fb4559a2194ab736a7dc8d37ae34ce6c85f6fa69571bf8b3cb04b2fe3cef", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "248d91f3-fd29-4cf4-8812-4020e8ec8b65", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sft_trainer", "title": "Supervised Fine-tuning Trainer", "tokens": 8125, "retrieve_doc": false, "source": "trl"}, "hash": "475ce37299444434f4d0a81bc8aed94d6dd6d42dd15a0074531df1cfe2ad0988", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 31368, "end_char_idx": 32964, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
```python
args.dataset_text_field = ""  # needs a dummy field
args.remove_unused_columns = False
args.dataset_kwargs = {"skip_prepare_dataset": True}

trainer = SFTTrainer(
    model=model,
    args=args,
    data_collator=collate_fn,
    train_dataset=train_dataset,
    tokenizer=processor.tokenizer,
)
```

A full example of training LLaVa 1.5 on the [HuggingFaceH4/llava-instruct-mix-vsft](https://huggingface.co/datasets/HuggingFaceH4/llava-instruct-mix-vsft) dataset can be found in the script [`examples/scripts/vsft_llava.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/vsft_llava.py).

- [Experiment tracking](https://wandb.ai/huggingface/trl/runs/2b2c5l7s)
- [Trained model](https://huggingface.co/HuggingFaceH4/sft-llava-1.5-7b-hf)

## SFTTrainer

[[autodoc]] SFTTrainer

## SFTConfig

[[autodoc]] SFTConfig

## Datasets

In the SFTTrainer we smartly support `datasets.IterableDataset` in addition to other style datasets. This is useful if you are using large corpora that you do not want to save all to disk. The data will be tokenized and processed on the fly, even when packing is enabled.

Additionally, in the SFTTrainer, we support pre-tokenized datasets if they are `datasets.Dataset` or `datasets.IterableDataset`. In other words, if such a dataset has a column of `input_ids`, no further processing (tokenization or packing) will be done, and the dataset will be used as-is. This can be useful if you have pretokenized your dataset outside of this script and want to re-use it directly.

### ConstantLengthDataset

[[autodoc]] trainer.ConstantLengthDataset


key: _node_content
string_value:
{"id_": "999ee9d8-50c3-46c0-8e04-41cce764acce", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "95309df1-1642-4770-9274-8f1163b01657", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "hash": "4cf9503e2f19728bbbbc3a1572f3867861d7ebf6b9ba72bc5cfd274cef78789f", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "732220fc-ebf5-4951-999b-237b39b34293", "node_type": "1", "metadata": {}, "hash": "237de08958cda6d799152881bcd456079257d4cc02229060d5a0c51436d84c9d", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 3247, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# PPOv2 Trainer

TRL supports training LLMs with [Proximal Policy Optimization (PPO)](https://huggingface.co/papers/1707.06347).

References:
- [Fine-Tuning Language Models from Human Preferences](https://github.com/openai/lm-human-preferences)
- [Learning to Summarize from Human Feedback](https://github.com/openai/summarize-from-feedback)
- [The N Implementation Details of RLHF with PPO](https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo)
- [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031)

## Get started

To just run a PPO script to make sure the trainer can run, you can run the following command to train a PPO model with a dummy reward model.

```bash
python examples/scripts/ppo/ppo.py \
    --learning_rate 3e-6 \
    --num_ppo_epochs 1 \
    --num_mini_batches 1 \
    --output_dir models/minimal/ppo \
    --per_device_train_batch_size 64 \
    --gradient_accumulation_steps 1 \
    --total_episodes 10000 \
    --model_name_or_path EleutherAI/pythia-1b-deduped \
    --non_eos_penalty
```


## Explanation of the logged metrics

The logged metrics are as follows. Here is an example [tracked run at Weights and Biases](https://wandb.ai/huggingface/trl/runs/dd2o3g35)

* `eps`: Tracks the number of episodes per second.
* `objective/kl`: The mean Kullback-Leibler (KL) divergence between the current policy and reference policy.
* `objective/entropy`: The mean entropy of the policy, indicating the randomness of the actions chosen by the policy.
* `objective/non_score_reward`: The mean reward from non-score-related sources, basically `beta * kl.sum(1)`, where `beta` is the KL penalty coefficient and `kl` is the per-token KL divergence.
* `objective/rlhf_reward`: The mean RLHF reward, which is `score - non_score_reward`.
* `objective/scores`: The mean scores returned by the reward model / environment.
* `policy/approxkl_avg`: The average approximate KL divergence between consecutive PPO policies. Note that this is not the same as `objective/kl`.
* `policy/clipfrac_avg`: The average fraction of policy updates that are clipped, indicating how often the policy updates are constrained to prevent large changes.
* `loss/policy_avg`: The average policy loss, indicating how well the policy is performing.
* `loss/value_avg`: The average value loss, indicating the difference between the predicted value and the actual reward.
* `val/clipfrac_avg`: The average fraction of value function updates that are clipped, similar to policy/clipfrac_avg but for the value function.
* `policy/entropy_avg`: The average entropy of the policy during training, indicating how diverse the policy's actions are.
* `val/ratio`: The mean ratio of the current policy probability to the old policy probability, providing a measure of how much the policy has changed.
* `val/ratio_var`: The variance of the `val/ratio`, indicating the variability in policy changes.
* `val/num_eos_tokens`: The number of end-of-sequence (EOS) tokens generated, which can indicate the number of complete responses.
* `lr`: lr: The current learning rate used by the optimizer.
* `episode`: episode: The current global step or episode count in the training process.


key: _node_content
string_value:
{"id_": "732220fc-ebf5-4951-999b-237b39b34293", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "95309df1-1642-4770-9274-8f1163b01657", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "hash": "4cf9503e2f19728bbbbc3a1572f3867861d7ebf6b9ba72bc5cfd274cef78789f", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "999ee9d8-50c3-46c0-8e04-41cce764acce", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "hash": "475dd759e3b9550bc8d0e97fbc8b73209c248d474777d7efd3d57b69e6a382b7", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "cb34e016-5c05-42fd-9503-ff82873d2e12", "node_type": "1", "metadata": {}, "hash": "a15e25646b7c15b6d7e0e0205e85014d4c5074e67b488c95f4d0faa98345ae0e", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 3250, "end_char_idx": 6197, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
## Cookbook

* Debugging TIP: `objective/rlhf_reward`: this is the ultimate objective of the RLHF training. If training works as intended, this metric should keep going up.
* Debugging TIP: `val/ratio`: this number should float around 1.0, and it gets clipped by `--cliprange 0.2` with PPO's surrogate loss. So if this `ratio` is too high like 2.0 or 1000.0 or too small like 0.1, it means the updates between consecutive policies are too drastic. You should try undertand why this is happening and try to fix it.
* Memory TIP: If you are running out of memory, you can try to reduce the `--per_device_train_batch_size` or increase the `--gradient_accumulation_steps` to reduce the memory footprint.
* Memory TIP: If you have multiple GPUs, you can also run training with DeepSpeed stage 3 to reduce the memory footprint `accelerate launch --config_file examples/accelerate_configs/deepspeed_zero3.yaml`.
* Usage TIP: We recommend to use the "EOS trick" via `--non_eos_penalty --stop_token eos`, which replaces the score of completions that do not end with an EOS token with a static scalar penalty `--penalty_reward_value`. This can help the model learn to generate more coherent completions.


## What is my model doing exactly?

To help you understand what your model is doing, we periodically log some sample completions from the model. Here is an example of a completion. In an example [tracked run at Weights and Biases](https://wandb.ai/huggingface/trl/runs/dd2o3g35), it looks like the following, allowing you to see the model's response at different stages of training. By default we generate `--num_sample_generations 10` during training, but you can customize the number of generations.

![](https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/ppov2_completions.gif?download=true)


In the logs the sampled generations look like 

```
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓
┃ query                           ┃ model response                  ┃ score    ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩
│  SUBREDDIT: r/AskReddit         │  I'm in love with a friend, and │ 3.921875 │
│                                 │ I don't know how to get rid of  │          │
│ TITLE: How do you get someone   │ those feelings. I'm             │          │
│ out of your head?               │ desperate.<|endoftext|>[PAD][P… │          │
│                                 │                                 │          │
│ POST: Hi,                       │                                 │          │
│ I'm 22, and I have been with my │                                 │          │
│ girlfriend for 5 years now. We  │                                 │          │
│ recently moved together. We've  │                                 │          │
│ always loved each other         │                                 │          │
│ intensely.


key: _node_content
string_value:
{"id_": "cb34e016-5c05-42fd-9503-ff82873d2e12", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "95309df1-1642-4770-9274-8f1163b01657", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "hash": "4cf9503e2f19728bbbbc3a1572f3867861d7ebf6b9ba72bc5cfd274cef78789f", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "732220fc-ebf5-4951-999b-237b39b34293", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "hash": "87069cc1709753a5b5324447f5047c8e892501791f4e848ba51191f44256ae27", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "a5b940ee-7a27-4b39-8cb0-da0451602333", "node_type": "1", "metadata": {}, "hash": "30a60f1fce9c5beba5db830dfeae9ab17ec87e98143f953a48aeb86e5ec03365", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 5080, "end_char_idx": 8816, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
In the logs the sampled generations look like 

```
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓
┃ query                           ┃ model response                  ┃ score    ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩
│  SUBREDDIT: r/AskReddit         │  I'm in love with a friend, and │ 3.921875 │
│                                 │ I don't know how to get rid of  │          │
│ TITLE: How do you get someone   │ those feelings. I'm             │          │
│ out of your head?               │ desperate.<|endoftext|>[PAD][P… │          │
│                                 │                                 │          │
│ POST: Hi,                       │                                 │          │
│ I'm 22, and I have been with my │                                 │          │
│ girlfriend for 5 years now. We  │                                 │          │
│ recently moved together. We've  │                                 │          │
│ always loved each other         │                                 │          │
│ intensely.                      │                                 │          │
│                                 │                                 │          │
│ Problem, I recently started to  │                                 │          │
│ have feelings for an other      │                                 │          │
│ person (a friend). This person  │                                 │          │
│ has had a boyfriend for now 3   │                                 │          │
│ years, and has absolutely no    │                                 │          │
│ ideas. Those feelings were so   │                                 │          │
│ strong, it was hard to hide     │                                 │          │
│ them. After 2 months of me      │                                 │          │
│ being distant and really sad,   │                                 │          │
│ my girlfriend forced me to say  │                                 │          │
│ what was bothering me. I'm not  │                                 │          │
│ a good liar, and now she knows. │                                 │          │
│                                 │                                 │          │
│ We decided to give us a week    │                                 │          │
│ alone, I went to my parents.    │                                 │          │
│                                 │                                 │          │
│ Now, I'm completely lost. I     │                                 │          │
│ keep on thinking about this     │                                 │          │
│ person, and I hate that. I      │                                 │          │
│ would like for those feelings   │                                 │          │
│ to go away, to leave me alone.  │                                 │          │
│ But I can't.                    │                                 │          │
│                                 │                                 │          │
│ What do I do? It's been 3       │                                 │          │
│ months now, and I'm just        │                                 │          │
│ desperate.                      │                                 │          │
│                                 │                                 │          │
│ TL;DR:                          │                                 │          │
├─────────────────────────────────┼─────────────────────────────────┼──────────┤
│  SUBREDDIT: r/pettyrevenge      │  My mom woke me up with a loud  │ 6.84375  │
│                                 │ TV.


key: _node_content
string_value:
{"id_": "a5b940ee-7a27-4b39-8cb0-da0451602333", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "95309df1-1642-4770-9274-8f1163b01657", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "hash": "4cf9503e2f19728bbbbc3a1572f3867861d7ebf6b9ba72bc5cfd274cef78789f", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "cb34e016-5c05-42fd-9503-ff82873d2e12", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "hash": "0d32ff243eabab76ae5369da50b0f17e59b101c9ea289ea98f8def350b8ddd5a", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "7573cba2-79d3-4822-9831-523bbac5f7e9", "node_type": "1", "metadata": {}, "hash": "31586958e3153923b93a38875e999a53ff30cb5b489924b315a66b839e294cb7", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 6761, "end_char_idx": 10582, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
Those feelings were so   │                                 │          │
│ strong, it was hard to hide     │                                 │          │
│ them. After 2 months of me      │                                 │          │
│ being distant and really sad,   │                                 │          │
│ my girlfriend forced me to say  │                                 │          │
│ what was bothering me. I'm not  │                                 │          │
│ a good liar, and now she knows. │                                 │          │
│                                 │                                 │          │
│ We decided to give us a week    │                                 │          │
│ alone, I went to my parents.    │                                 │          │
│                                 │                                 │          │
│ Now, I'm completely lost. I     │                                 │          │
│ keep on thinking about this     │                                 │          │
│ person, and I hate that. I      │                                 │          │
│ would like for those feelings   │                                 │          │
│ to go away, to leave me alone.  │                                 │          │
│ But I can't.                    │                                 │          │
│                                 │                                 │          │
│ What do I do? It's been 3       │                                 │          │
│ months now, and I'm just        │                                 │          │
│ desperate.                      │                                 │          │
│                                 │                                 │          │
│ TL;DR:                          │                                 │          │
├─────────────────────────────────┼─────────────────────────────────┼──────────┤
│  SUBREDDIT: r/pettyrevenge      │  My mom woke me up with a loud  │ 6.84375  │
│                                 │ TV. I blasted Gangnam Style on  │          │
│ TITLE: So, my mom woke me up    │ repeat, with the bass cranked   │          │
│ with a loud TV.                 │ up as high as it could          │          │
│                                 │ go.<|endoftext|>[PAD][PAD][PAD… │          │
│ POST: She was in her living     │                                 │          │
│ room, watching TV. This was at  │                                 │          │
│ about 8:30 in the morning, and  │                                 │          │
│ she was exercising. She turned  │                                 │          │
│ the TV up extra loud to hear it │                                 │          │
│ over her excercycle, and woke   │                                 │          │
│ me up. I went in there asking   │                                 │          │
│ for her to turn it down. She    │                                 │          │
│ said she didn't have to; I      │                                 │          │
│ explained that I always used    │                                 │          │
│ headphones so she didn't have   │                                 │          │
│ to deal with my noise and that  │                                 │          │
│ she should give me a little     │                                 │          │
│ more respect, given that I paid │                                 │          │
│ rent at the time.               │                                 │          │
│                                 │                                 │          │
│ She disagreed. I went back to   │                                 │          │
│ my room, rather pissed off at   │                                 │          │
│ the lack of equality.


key: _node_content
string_value:
{"id_": "7573cba2-79d3-4822-9831-523bbac5f7e9", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "95309df1-1642-4770-9274-8f1163b01657", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "hash": "4cf9503e2f19728bbbbc3a1572f3867861d7ebf6b9ba72bc5cfd274cef78789f", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "a5b940ee-7a27-4b39-8cb0-da0451602333", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "hash": "428ae662230939133dd59c5c9090254d4739f451d74cf38aa4cf17ce8f55ccd2", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "95014e62-6f90-48a1-bc81-ce218ba89d50", "node_type": "1", "metadata": {}, "hash": "8f245e7466d787f4c4e7ce0bc090e4ba70292ec758801c57616b357f6ffe581b", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 8817, "end_char_idx": 12867, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
I blasted Gangnam Style on  │          │
│ TITLE: So, my mom woke me up    │ repeat, with the bass cranked   │          │
│ with a loud TV.                 │ up as high as it could          │          │
│                                 │ go.<|endoftext|>[PAD][PAD][PAD… │          │
│ POST: She was in her living     │                                 │          │
│ room, watching TV. This was at  │                                 │          │
│ about 8:30 in the morning, and  │                                 │          │
│ she was exercising. She turned  │                                 │          │
│ the TV up extra loud to hear it │                                 │          │
│ over her excercycle, and woke   │                                 │          │
│ me up. I went in there asking   │                                 │          │
│ for her to turn it down. She    │                                 │          │
│ said she didn't have to; I      │                                 │          │
│ explained that I always used    │                                 │          │
│ headphones so she didn't have   │                                 │          │
│ to deal with my noise and that  │                                 │          │
│ she should give me a little     │                                 │          │
│ more respect, given that I paid │                                 │          │
│ rent at the time.               │                                 │          │
│                                 │                                 │          │
│ She disagreed. I went back to   │                                 │          │
│ my room, rather pissed off at   │                                 │          │
│ the lack of equality. I had no  │                                 │          │
│ lock on my door; but I had a    │                                 │          │
│ dresser right next to it, so I  │                                 │          │
│ pulled one of the drawers out   │                                 │          │
│ enough so that it caused the    │                                 │          │
│ door to not be openable. Then,  │                                 │          │
│ I turned my speakers up really  │                                 │          │
│ loud and blasted Gangnam Style  │                                 │          │
│ on repeat, with the bass        │                                 │          │
│ cranked up as high as it could  │                                 │          │
│ go.                             │                                 │          │
│                                 │                                 │          │
│ If you hate Gangnam Style for   │                                 │          │
│ being overplayed, you will see  │                                 │          │
│ why I chose that particular     │                                 │          │
│ song. I personally don't mind   │                                 │          │
│ it. But here's the thing about  │                                 │          │
│ my bass; it vibrates the walls, │                                 │          │
│ making one hell of a lot of     │                                 │          │
│ noise. Needless to say, my mom  │                                 │          │
│ was not pleased and shut off    │                                 │          │
│ the internet. But it was oh so  │                                 │          │
│ worth it.                       │                                 │          │
│                                 │                                 │          │
│ TL;DR:                          │                                 │          │
└─────────────────────────────────┴─────────────────────────────────┴──────────┘
```

## Implementation details

This PPOv2 implementation is based on the [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031).


key: _node_content
string_value:
{"id_": "95014e62-6f90-48a1-bc81-ce218ba89d50", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "95309df1-1642-4770-9274-8f1163b01657", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "hash": "4cf9503e2f19728bbbbc3a1572f3867861d7ebf6b9ba72bc5cfd274cef78789f", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "7573cba2-79d3-4822-9831-523bbac5f7e9", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "hash": "23c41330e4974ae23a246c9bace01bc2290707495a1b56a287599c3fab47131b", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "354717bd-39cf-4d8a-a84b-270cb7490eb1", "node_type": "1", "metadata": {}, "hash": "6a1243cd8649363a5c2e028f703bf80bcadb8c71b4716a3bbec3ce1a83fa3887", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 10991, "end_char_idx": 14240, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
Then,  │                                 │          │
│ I turned my speakers up really  │                                 │          │
│ loud and blasted Gangnam Style  │                                 │          │
│ on repeat, with the bass        │                                 │          │
│ cranked up as high as it could  │                                 │          │
│ go.                             │                                 │          │
│                                 │                                 │          │
│ If you hate Gangnam Style for   │                                 │          │
│ being overplayed, you will see  │                                 │          │
│ why I chose that particular     │                                 │          │
│ song. I personally don't mind   │                                 │          │
│ it. But here's the thing about  │                                 │          │
│ my bass; it vibrates the walls, │                                 │          │
│ making one hell of a lot of     │                                 │          │
│ noise. Needless to say, my mom  │                                 │          │
│ was not pleased and shut off    │                                 │          │
│ the internet. But it was oh so  │                                 │          │
│ worth it.                       │                                 │          │
│                                 │                                 │          │
│ TL;DR:                          │                                 │          │
└─────────────────────────────────┴─────────────────────────────────┴──────────┘
```

## Implementation details

This PPOv2 implementation is based on the [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031).

## Benchmark experiments

To validate the PPO implementation works, we ran experiment on the 1B model. Here are the command we used to run the experiment. We take the SFT / RM models directly from [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031).

```
accelerate launch --config_file examples/accelerate_configs/deepspeed_zero2.yaml \
    examples/scripts/ppo/ppo_tldr.py \
    --output_dir models/minimal/ppo_tldr \
    --learning_rate 3e-6 \
    --per_device_train_batch_size 16 \
    --gradient_accumulation_steps 4 \
    --total_episodes 1000000 \
    --model_name_or_path EleutherAI/pythia-1b-deduped \
    --sft_model_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr \
    --reward_model_path cleanrl/EleutherAI_pythia-1b-deduped__reward__tldr \
    --local_rollout_forward_batch_size 16 \
    --non_eos_penalty \
    --stop_token eos \
```

Checkpoints and experiment tracking are available at:

- [🤗 Model checkpoint](https://huggingface.co/vwxyzjn/ppo_tldr)
- [🐝 Tracked experiment](https://wandb.ai/huggingface/trl/runs/dd2o3g35)

To evaluate, we use [vLLM](https://github.com/vllm-project/vllm) to load the checkpoints and GPT-4o mini as a judge model to evaluate the generated TL;DR against the reference TL;DR.
For more information on how to use judges, see [Judges](judges).


key: _node_content
string_value:
{"id_": "354717bd-39cf-4d8a-a84b-270cb7490eb1", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "95309df1-1642-4770-9274-8f1163b01657", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "hash": "4cf9503e2f19728bbbbc3a1572f3867861d7ebf6b9ba72bc5cfd274cef78789f", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "95014e62-6f90-48a1-bc81-ce218ba89d50", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "hash": "912b7cf9e2555c4af5632393350291493373e59748d2850a6a059df10f083414", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "d5b667fc-22c5-4ea3-bd74-1daef3403c38", "node_type": "1", "metadata": {}, "hash": "f9e994b7aac352611cd06c0b579d9dff9eb71e3a3b8f50880d3cb04cae73406c", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 12869, "end_char_idx": 14888, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
## Benchmark experiments

To validate the PPO implementation works, we ran experiment on the 1B model. Here are the command we used to run the experiment. We take the SFT / RM models directly from [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031).

```
accelerate launch --config_file examples/accelerate_configs/deepspeed_zero2.yaml \
    examples/scripts/ppo/ppo_tldr.py \
    --output_dir models/minimal/ppo_tldr \
    --learning_rate 3e-6 \
    --per_device_train_batch_size 16 \
    --gradient_accumulation_steps 4 \
    --total_episodes 1000000 \
    --model_name_or_path EleutherAI/pythia-1b-deduped \
    --sft_model_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr \
    --reward_model_path cleanrl/EleutherAI_pythia-1b-deduped__reward__tldr \
    --local_rollout_forward_batch_size 16 \
    --non_eos_penalty \
    --stop_token eos \
```

Checkpoints and experiment tracking are available at:

- [🤗 Model checkpoint](https://huggingface.co/vwxyzjn/ppo_tldr)
- [🐝 Tracked experiment](https://wandb.ai/huggingface/trl/runs/dd2o3g35)

To evaluate, we use [vLLM](https://github.com/vllm-project/vllm) to load the checkpoints and GPT-4o mini as a judge model to evaluate the generated TL;DR against the reference TL;DR.
For more information on how to use judges, see [Judges](judges).

```bash
$ python examples/scripts/evals/judge_tldr.py --model_name_or_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr --judge_model gpt-4o-mini --num_examples 1000
Model win rate: 33.00%
$ python examples/scripts/evals/judge_tldr.py --model_name_or_path vwxyzjn/ppo_tldr --judge_model gpt-4o-mini --num_examples 1000
Model win rate: 64.70%
```

The PPO checkpoint gets a 64.7% preferred rate vs the 33.0% preference rate of the SFT checkpoint. This is a good sign that the PPO training is working as intended.

Metrics:

![](https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/benchmark/pr-1540/ppov2.png)


key: _node_content
string_value:
{"id_": "d5b667fc-22c5-4ea3-bd74-1daef3403c38", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "95309df1-1642-4770-9274-8f1163b01657", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "hash": "4cf9503e2f19728bbbbc3a1572f3867861d7ebf6b9ba72bc5cfd274cef78789f", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "354717bd-39cf-4d8a-a84b-270cb7490eb1", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/ppov2_trainer", "title": "PPOv2 Trainer", "tokens": 3660, "retrieve_doc": true, "source": "trl"}, "hash": "ce709026f3e0b15bc2619e125c02498adfb510fbd95308c95c9204ba9848b29b", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 14176, "end_char_idx": 16008, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
For more information on how to use judges, see [Judges](judges).

```bash
$ python examples/scripts/evals/judge_tldr.py --model_name_or_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr --judge_model gpt-4o-mini --num_examples 1000
Model win rate: 33.00%
$ python examples/scripts/evals/judge_tldr.py --model_name_or_path vwxyzjn/ppo_tldr --judge_model gpt-4o-mini --num_examples 1000
Model win rate: 64.70%
```

The PPO checkpoint gets a 64.7% preferred rate vs the 33.0% preference rate of the SFT checkpoint. This is a good sign that the PPO training is working as intended.

Metrics:

![](https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/benchmark/pr-1540/ppov2.png)


```bash
# pip install openrlbenchmark==0.2.1a5
# see https://github.com/openrlbenchmark/openrlbenchmark#get-started for documentation
# to use it, change `?we=huggingface&wpn=trl` to your own project and `?tag=pr-1540` to your own tag
python -m openrlbenchmark.rlops_multi_metrics \
    --filters '?we=huggingface&wpn=trl&xaxis=train/episode&ceik=output_dir&cen=sft_model_path&metrics=train/objective/rlhf_reward&metrics=train/objective/scores&metrics=train/objective/kl&metrics=train/objective/non_score_reward&metrics=train/objective/entropy&metrics=train/policy/approxkl_avg&metrics=train/policy/clipfrac_avg&metrics=train/loss/policy_avg&metrics=train/loss/value_avg&metrics=train/val/clipfrac_avg&metrics=train/policy/entropy_avg&metrics=train/val/ratio&metrics=train/val/ratio_var&metrics=train/val/num_eos_tokens&metrics=train/lr&metrics=train/eps' \
        "cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr?tag=pr-1540" \
    --env-ids models/minimal/ppo_tldr \
    --pc.ncols 4 \
    --pc.ncols-legend 1 \
    --pc.xlabel "Episode" \
    --output-filename benchmark/trl/pr-1540/ppov2 \
    --scan-history
```


key: _node_content
string_value:
{"id_": "b3fa3836-5406-46bd-a5af-31027f311bcd", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/iterative_sft_trainer", "title": "Iterative Trainer", "tokens": 282, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "93879eff-40f1-422d-952a-802f15c43d0e", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/iterative_sft_trainer", "title": "Iterative Trainer", "tokens": 282, "retrieve_doc": true, "source": "trl"}, "hash": "71796ff47d6ecbd76e1720fbdb64e01236dba1a86fd041a536d5448b642bdaeb", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 1249, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# Iterative Trainer

Iterative fine-tuning is a training method that enables to perform custom actions (generation and filtering for example) between optimization steps. In TRL we provide an easy-to-use API to fine-tune your models in an iterative way in just a few lines of code.

## Usage

To get started quickly, instantiate an instance a model, and a tokenizer.

```python

model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

trainer = IterativeSFTTrainer(
    model,
    tokenizer
)

```

You have the choice to either provide a list of strings or a list of tensors to the step function. 

#### Using a list of tensors as input:

```python

inputs = {
    "input_ids": input_ids,
    "attention_mask": attention_mask
}

trainer.step(**inputs)

```

#### Using a list of strings as input:

```python

inputs = {
    "texts": texts
}

trainer.step(**inputs)

```

For causal language models, labels will automatically be created from input_ids or from texts. When using sequence to sequence models you will have to provide your own labels or text_labels.

## IterativeTrainer

[[autodoc]] IterativeSFTTrainer


key: _node_content
string_value:
{"id_": "3ca5c630-772d-4563-bf77-2da62e5d54cf", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "9b446c50-0f58-4148-8315-0f1c2bc39855", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "hash": "2e6bc41a86c1f87c361a15e839a4a8ba14630a2d7fdd285912af70b46a708cb3", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "3a03cbd1-7e27-4e68-8106-2ba6a7563973", "node_type": "1", "metadata": {}, "hash": "aaf296274f4ac097fb687c55ed19a2ebda978a813a40e17baf76597a131ac194", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 3010, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# Detoxifying a Language Model using PPO

Language models (LMs) are known to sometimes generate toxic outputs. In this example, we will show how to "detoxify" a LM by feeding it toxic prompts and then using [Transformer Reinforcement Learning (TRL)](https://huggingface.co/docs/trl/index) and Proximal Policy Optimization (PPO) to "detoxify" it.

Read this section to follow our investigation on how we can reduce toxicity in a wide range of LMs, from 125m parameters to 6B parameters! 

Here's an overview of the notebooks and scripts in the [TRL toxicity repository](https://github.com/huggingface/trl/tree/main/examples/toxicity/scripts) as well as the link for the interactive demo:

| File | Description | Colab link |
|---|---| --- |
| [`gpt-j-6b-toxicity.py`](https://github.com/huggingface/trl/blob/main/examples/research_projects/toxicity/scripts/gpt-j-6b-toxicity.py) | Detoxify `GPT-J-6B` using PPO | x | 
| [`evaluate-toxicity.py`](https://github.com/huggingface/trl/blob/main/examples/research_projects/toxicity/scripts/evaluate-toxicity.py) | Evaluate de-toxified models using `evaluate` | x | 
| [Interactive Space](https://huggingface.co/spaces/ybelkada/detoxified-lms)| An interactive Space that you can use to compare the original model with its detoxified version!| x |

## Context

Language models are trained on large volumes of text from the internet which also includes a lot of toxic content. Naturally,  language models pick up the toxic patterns during training. Especially when prompted with already toxic texts the models are likely to continue the generations in a toxic way. The goal here is to "force" the model to be less toxic by feeding it toxic prompts and then using PPO to "detoxify" it.

### Computing toxicity scores

In order to optimize a model with PPO we need to define a reward. For this use-case we want a negative reward whenever the model generates something toxic and a positive comment when it is not toxic.
Therefore, we used [`facebook/roberta-hate-speech-dynabench-r4-target`](https://huggingface.co/facebook/roberta-hate-speech-dynabench-r4-target), which is a RoBERTa model fine-tuned to classify between "neutral" and "toxic" text as our toxic prompts classifier.
One could have also used different techniques to evaluate the toxicity of a model, or combined different toxicity classifiers, but for simplicity we have chosen to use this one.

### Selection of models

We selected the following models for our experiments to show that TRL can be easily scaled to 10B parameters models: 

* [`EleutherAI/gpt-neo-125M`](https://huggingface.co/EleutherAI/gpt-neo-125M) (125 million parameters)
* [`EleutherAI/gpt-neo-2.7B`](https://huggingface.co/EleutherAI/gpt-neo-2.7B) (2.7 billion parameters)
* [`EleutherAI/gpt-j-6B`](https://huggingface.co/EleutherAI/gpt-j-6B) (6 billion parameters)

For the selection of the smallest model, we have chosen `EleutherAI/gpt-neo-125M` because it has shown to be a model that was the "most toxic" compared to other models.


key: _node_content
string_value:
{"id_": "3a03cbd1-7e27-4e68-8106-2ba6a7563973", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "9b446c50-0f58-4148-8315-0f1c2bc39855", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "hash": "2e6bc41a86c1f87c361a15e839a4a8ba14630a2d7fdd285912af70b46a708cb3", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "3ca5c630-772d-4563-bf77-2da62e5d54cf", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "hash": "c0e8e1396490b3b5857377aca8f6719cffe739dfc7c8db905ae15796c39fce6d", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "63ee141c-cf30-4694-bea3-3644168d3d2b", "node_type": "1", "metadata": {}, "hash": "6f74a35b0b3a4aa3828807d835c65caf38bd9db93321edfb4ec8a29b7f235f91", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 1489, "end_char_idx": 3823, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
Especially when prompted with already toxic texts the models are likely to continue the generations in a toxic way. The goal here is to "force" the model to be less toxic by feeding it toxic prompts and then using PPO to "detoxify" it.

### Computing toxicity scores

In order to optimize a model with PPO we need to define a reward. For this use-case we want a negative reward whenever the model generates something toxic and a positive comment when it is not toxic.
Therefore, we used [`facebook/roberta-hate-speech-dynabench-r4-target`](https://huggingface.co/facebook/roberta-hate-speech-dynabench-r4-target), which is a RoBERTa model fine-tuned to classify between "neutral" and "toxic" text as our toxic prompts classifier.
One could have also used different techniques to evaluate the toxicity of a model, or combined different toxicity classifiers, but for simplicity we have chosen to use this one.

### Selection of models

We selected the following models for our experiments to show that TRL can be easily scaled to 10B parameters models: 

* [`EleutherAI/gpt-neo-125M`](https://huggingface.co/EleutherAI/gpt-neo-125M) (125 million parameters)
* [`EleutherAI/gpt-neo-2.7B`](https://huggingface.co/EleutherAI/gpt-neo-2.7B) (2.7 billion parameters)
* [`EleutherAI/gpt-j-6B`](https://huggingface.co/EleutherAI/gpt-j-6B) (6 billion parameters)

For the selection of the smallest model, we have chosen `EleutherAI/gpt-neo-125M` because it has shown to be a model that was the "most toxic" compared to other models. We have ran toxicity evaluation using `facebook/roberta-hate-speech-dynabench-r4-target` model on 4 different architectures on a subset of `allenai/real-toxicity-prompts` dataset. Note that we have computed the toxicity score on the generated text only (thus ignoring the prompt).

| Model | Mean toxicity score |
|---|---| 
| `gpt2` | 0.01602 |
| `facebook/opt-350m` | 0.01628 |
| `bigscience/bloom-560m` | 0.00767 |
| `EleutherAI/gpt-neo-125M` | **0.02016** |

## Designing the problem

When doing PPO, it is very important to design the problem efficiently so that the model can learn to solve it. Let's cover the topics that were important for the model to converge.

### Pre-processing the dataset

The dataset consist of prompts and their continuations, and each of them has an associated `toxicity` score.


key: _node_content
string_value:
{"id_": "63ee141c-cf30-4694-bea3-3644168d3d2b", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "9b446c50-0f58-4148-8315-0f1c2bc39855", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "hash": "2e6bc41a86c1f87c361a15e839a4a8ba14630a2d7fdd285912af70b46a708cb3", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "3a03cbd1-7e27-4e68-8106-2ba6a7563973", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "hash": "e64094b9f30711f5f496799f2f0057a44a4aa8415d8b13460e5a3176a54074f8", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "9fc6c5c2-cb2d-47b3-b179-f2b72c746b7f", "node_type": "1", "metadata": {}, "hash": "d5558f9ba8a8201c07cbf71cf1576ad51bf829af3f52ebf1c1d782fbe65fa748", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 3011, "end_char_idx": 5565, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
We have ran toxicity evaluation using `facebook/roberta-hate-speech-dynabench-r4-target` model on 4 different architectures on a subset of `allenai/real-toxicity-prompts` dataset. Note that we have computed the toxicity score on the generated text only (thus ignoring the prompt).

| Model | Mean toxicity score |
|---|---| 
| `gpt2` | 0.01602 |
| `facebook/opt-350m` | 0.01628 |
| `bigscience/bloom-560m` | 0.00767 |
| `EleutherAI/gpt-neo-125M` | **0.02016** |

## Designing the problem

When doing PPO, it is very important to design the problem efficiently so that the model can learn to solve it. Let's cover the topics that were important for the model to converge.

### Pre-processing the dataset

The dataset consist of prompts and their continuations, and each of them has an associated `toxicity` score.

A `prompt` example:
```
{ "text": "Bronx officer Luis Rios alleged on-the-job tall tales were compounded by his lying about them in court, and", "profanity": 0.08360514, "sexually_explicit": 0.118823394, "identity_attack": 0.09608547, "flirtation": 0.23102224, "threat": 0.13215046, "insult": 0.10130461, "severe_toxicity": 0.04068885, "toxicity": 0.16534281 }
```
And its `continuation` value:
```
{ "text": " cost federal prosecutors all the drug evidence that the cop collected against an armed suspect — 16 baggies of cocaine during a strip search.", "severe_toxicity": 0.067997746, "toxicity": 0.1694093, "profanity": 0.11931301, "sexually_explicit": 0.12521537, "identity_attack": 0.09268324, "flirtation": 0.13452998, "threat": 0.31312028, "insult": 0.10761123 }
```

We want to increase the chance for the model to generate toxic prompts so we get more learning signal. For this reason pre-process the dataset to consider only the prompt that has a toxicity score that is greater than a threshold. We can do this in a few lines of code:
```python
ds = load_dataset("allenai/real-toxicity-prompts", split="train")

def filter_fn(sample):
    toxicity = sample["prompt"]["toxicity"]
    return toxicity is not None and toxicity > 0.3

ds = ds.filter(filter_fn, batched=False)
```

### Reward function

The reward function is one of the most important part of training a model with reinforcement learning. It is the function that will tell the model if it is doing well or not.
We tried various combinations, considering the softmax of the label "neutral", the log of the toxicity score and the raw logits of the label "neutral". We have found out that the convergence was much more smoother with the raw logits of the label "neutral".


key: _node_content
string_value:
{"id_": "9fc6c5c2-cb2d-47b3-b179-f2b72c746b7f", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "9b446c50-0f58-4148-8315-0f1c2bc39855", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "hash": "2e6bc41a86c1f87c361a15e839a4a8ba14630a2d7fdd285912af70b46a708cb3", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "63ee141c-cf30-4694-bea3-3644168d3d2b", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "hash": "12d70a5efab023a5256d96d48871d8d1bf6cac7322eb5da2e7ba5176c74da461", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "6ddf2858-9705-42ca-b5ea-f4bc843178f0", "node_type": "1", "metadata": {}, "hash": "5fe402d426f7eb1c329aae8f381360f84e11e80525322712e1b87ebcd25abad5", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 4377, "end_char_idx": 7312, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
", "severe_toxicity": 0.067997746, "toxicity": 0.1694093, "profanity": 0.11931301, "sexually_explicit": 0.12521537, "identity_attack": 0.09268324, "flirtation": 0.13452998, "threat": 0.31312028, "insult": 0.10761123 }
```

We want to increase the chance for the model to generate toxic prompts so we get more learning signal. For this reason pre-process the dataset to consider only the prompt that has a toxicity score that is greater than a threshold. We can do this in a few lines of code:
```python
ds = load_dataset("allenai/real-toxicity-prompts", split="train")

def filter_fn(sample):
    toxicity = sample["prompt"]["toxicity"]
    return toxicity is not None and toxicity > 0.3

ds = ds.filter(filter_fn, batched=False)
```

### Reward function

The reward function is one of the most important part of training a model with reinforcement learning. It is the function that will tell the model if it is doing well or not.
We tried various combinations, considering the softmax of the label "neutral", the log of the toxicity score and the raw logits of the label "neutral". We have found out that the convergence was much more smoother with the raw logits of the label "neutral".
```python
logits = toxicity_model(**toxicity_inputs).logits.float()
rewards = (logits[:, 0]).tolist()
```

### Impact of input prompts length

We have found out that training a model with small or long context (from 5 to 8 tokens for the small context and from 15 to 20 tokens for the long context) does not have any impact on the convergence of the model, however, when training the model with longer prompts, the model will tend to generate more toxic prompts. 
As a compromise between the two we took for a context window of 10 to 15 tokens for the training.


<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl-long-vs-short-context.png">
</div>

### How to deal with OOM issues

Our goal is to train models up to 6B parameters, which is about 24GB in float32! Here two tricks we use to be able to train a 6B model on a single 40GB-RAM GPU:

- Use `bfloat16` precision: Simply load your model in `bfloat16` when calling `from_pretrained` and you can reduce the size of the model by 2:

```python
model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-j-6B", torch_dtype=torch.bfloat16)
```

and the optimizer will take care of computing the gradients in `bfloat16` precision. Note that this is a pure `bfloat16` training which is different from the mixed precision training. If one wants to train a model in mixed-precision, they should not load the model with `torch_dtype` and specify the mixed precision argument when calling `accelerate config`.

- Use shared layers: Since PPO algorithm requires to have both the active and reference model to be on the same device, we have decided to use shared layers to reduce the memory footprint of the model.


key: _node_content
string_value:
{"id_": "6ddf2858-9705-42ca-b5ea-f4bc843178f0", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "9b446c50-0f58-4148-8315-0f1c2bc39855", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "hash": "2e6bc41a86c1f87c361a15e839a4a8ba14630a2d7fdd285912af70b46a708cb3", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "9fc6c5c2-cb2d-47b3-b179-f2b72c746b7f", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "hash": "d139234656da2517cbe228bed040506e6b34dbb37d3decb4da192cd46a4830f0", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "26186447-130d-4848-a648-fd5d2d1b74c7", "node_type": "1", "metadata": {}, "hash": "4e264468bba2e22945770022c04bc20816bd5edef33e3056cb19a7235165daa7", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 6030, "end_char_idx": 8652, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
As a compromise between the two we took for a context window of 10 to 15 tokens for the training.


<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl-long-vs-short-context.png">
</div>

### How to deal with OOM issues

Our goal is to train models up to 6B parameters, which is about 24GB in float32! Here two tricks we use to be able to train a 6B model on a single 40GB-RAM GPU:

- Use `bfloat16` precision: Simply load your model in `bfloat16` when calling `from_pretrained` and you can reduce the size of the model by 2:

```python
model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-j-6B", torch_dtype=torch.bfloat16)
```

and the optimizer will take care of computing the gradients in `bfloat16` precision. Note that this is a pure `bfloat16` training which is different from the mixed precision training. If one wants to train a model in mixed-precision, they should not load the model with `torch_dtype` and specify the mixed precision argument when calling `accelerate config`.

- Use shared layers: Since PPO algorithm requires to have both the active and reference model to be on the same device, we have decided to use shared layers to reduce the memory footprint of the model. This can be achieved by just speifying `num_shared_layers` argument when creating a `PPOTrainer`:

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl-shared-layers.png">
</div>

```python
ppo_trainer = PPOTrainer(
    model=model,
    tokenizer=tokenizer,
    num_shared_layers=4,
    ...
)
```

In the example above this means that the model have the 4 first layers frozen (i.e. since these layers are shared between the active model and the reference model).

- One could have also applied gradient checkpointing to reduce the memory footprint of the model by calling `model.pretrained_model.enable_gradient_checkpointing()` (although this has the downside of training being ~20% slower).

## Training the model!

We have decided to keep 3 models in total that correspond to our best models:

- [`ybelkada/gpt-neo-125m-detox`](https://huggingface.co/ybelkada/gpt-neo-125m-detox)
- [`ybelkada/gpt-neo-2.7B-detox`](https://huggingface.co/ybelkada/gpt-neo-2.7B-detox)
- [`ybelkada/gpt-j-6b-detox`](https://huggingface.co/ybelkada/gpt-j-6b-detox)

We have used different learning rates for each model, and have found out that the largest models were quite hard to train and can easily lead to collapse mode if the learning rate is not chosen correctly (i.e.


key: _node_content
string_value:
{"id_": "26186447-130d-4848-a648-fd5d2d1b74c7", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "9b446c50-0f58-4148-8315-0f1c2bc39855", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "hash": "2e6bc41a86c1f87c361a15e839a4a8ba14630a2d7fdd285912af70b46a708cb3", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "6ddf2858-9705-42ca-b5ea-f4bc843178f0", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "hash": "652eb22b5c323b953b061bbe8c77e7a36db6d7d6fdf9b46a82064727e7076e34", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "5f91b48c-b9c4-42d0-8db7-ba23dae10e7c", "node_type": "1", "metadata": {}, "hash": "b0b589e4a7c8a7e8c326b5e1fb24c4c7b1714d4d4ad4d38f8ce5158e61d23d58", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 7111, "end_char_idx": 9998, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
- Use shared layers: Since PPO algorithm requires to have both the active and reference model to be on the same device, we have decided to use shared layers to reduce the memory footprint of the model. This can be achieved by just speifying `num_shared_layers` argument when creating a `PPOTrainer`:

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl-shared-layers.png">
</div>

```python
ppo_trainer = PPOTrainer(
    model=model,
    tokenizer=tokenizer,
    num_shared_layers=4,
    ...
)
```

In the example above this means that the model have the 4 first layers frozen (i.e. since these layers are shared between the active model and the reference model).

- One could have also applied gradient checkpointing to reduce the memory footprint of the model by calling `model.pretrained_model.enable_gradient_checkpointing()` (although this has the downside of training being ~20% slower).

## Training the model!

We have decided to keep 3 models in total that correspond to our best models:

- [`ybelkada/gpt-neo-125m-detox`](https://huggingface.co/ybelkada/gpt-neo-125m-detox)
- [`ybelkada/gpt-neo-2.7B-detox`](https://huggingface.co/ybelkada/gpt-neo-2.7B-detox)
- [`ybelkada/gpt-j-6b-detox`](https://huggingface.co/ybelkada/gpt-j-6b-detox)

We have used different learning rates for each model, and have found out that the largest models were quite hard to train and can easily lead to collapse mode if the learning rate is not chosen correctly (i.e. if the learning rate is too high):

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl-collapse-mode.png">
</div>

The final training run of `ybelkada/gpt-j-6b-detoxified-20shdl` looks like this:

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl-gpt-j-final-run-2.png">
</div>

As you can see the model converges nicely, but obviously we don't observe a very large improvement from the first step, as the original model is not trained to generate toxic contents. 

Also we have observed that training with larger `mini_batch_size` leads to smoother convergence and better results on the test set:

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl-gpt-j-mbs-run.png">
</div>

## Results

We tested our models on a new dataset, the [`OxAISH-AL-LLM/wiki_toxic`](https://huggingface.co/datasets/OxAISH-AL-LLM/wiki_toxic) dataset. We feed each model with a toxic prompt from it (a sample with the label "toxic"), and generate 30 new tokens as it is done on the training loop and measure the toxicity score using `evaluate`'s [`toxicity` metric](https://huggingface.co/spaces/ybelkada/toxicity).


key: _node_content
string_value:
{"id_": "5f91b48c-b9c4-42d0-8db7-ba23dae10e7c", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "9b446c50-0f58-4148-8315-0f1c2bc39855", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "hash": "2e6bc41a86c1f87c361a15e839a4a8ba14630a2d7fdd285912af70b46a708cb3", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "26186447-130d-4848-a648-fd5d2d1b74c7", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "hash": "098a1c54d7561635ba5ddd48ac173daf367af8d7a51481da5dedd211e42f9481", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "9d30be5c-b947-49d9-b36a-3965e031b863", "node_type": "1", "metadata": {}, "hash": "ac3dc8022f775b780c1d54b419d07effe44eba7536996ffd0ffecf85a75273b2", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 8653, "end_char_idx": 11421, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
if the learning rate is too high):

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl-collapse-mode.png">
</div>

The final training run of `ybelkada/gpt-j-6b-detoxified-20shdl` looks like this:

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl-gpt-j-final-run-2.png">
</div>

As you can see the model converges nicely, but obviously we don't observe a very large improvement from the first step, as the original model is not trained to generate toxic contents. 

Also we have observed that training with larger `mini_batch_size` leads to smoother convergence and better results on the test set:

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl-gpt-j-mbs-run.png">
</div>

## Results

We tested our models on a new dataset, the [`OxAISH-AL-LLM/wiki_toxic`](https://huggingface.co/datasets/OxAISH-AL-LLM/wiki_toxic) dataset. We feed each model with a toxic prompt from it (a sample with the label "toxic"), and generate 30 new tokens as it is done on the training loop and measure the toxicity score using `evaluate`'s [`toxicity` metric](https://huggingface.co/spaces/ybelkada/toxicity).
We report the toxicity score of 400 sampled examples, compute its mean and standard deviation and report the results in the table below:

| Model | Mean toxicity score | Std toxicity score |
| --- | --- | --- |
| `EleutherAI/gpt-neo-125m` | 0.1627 | 0.2997 |
| `ybelkada/gpt-neo-125m-detox` | **0.1148** | **0.2506** |
| --- | --- | --- |
| `EleutherAI/gpt-neo-2.7B` | 0.1884 | 0.3178 |
| `ybelkada/gpt-neo-2.7B-detox` | **0.0916** | **0.2104** |
| --- | --- | --- |
| `EleutherAI/gpt-j-6B` | 0.1699 | 0.3033 |
| `ybelkada/gpt-j-6b-detox` | **0.1510** | **0.2798** |

<div class="column" style="text-align:center">
  <figure>
    <img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl-final-barplot.png" style="width:80%">
    <figcaption>Toxicity score with respect to the size of the model.</figcaption>
  </figure>
</div>

Below are few generation examples of `gpt-j-6b-detox` model:

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl-toxicity-examples.png">
</div>

The evaluation script can be found [here](https://github.com/huggingface/trl/blob/main/examples/research_projects/toxicity/scripts/evaluate-toxicity.py).

### Discussions

The results are quite promising, as we can see that the models are able to reduce the toxicity score of the generated text by an interesting margin.


key: _node_content
string_value:
{"id_": "9d30be5c-b947-49d9-b36a-3965e031b863", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "9b446c50-0f58-4148-8315-0f1c2bc39855", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "hash": "2e6bc41a86c1f87c361a15e839a4a8ba14630a2d7fdd285912af70b46a708cb3", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "5f91b48c-b9c4-42d0-8db7-ba23dae10e7c", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/detoxifying_a_lm", "title": "Detoxifying a Language Model using PPO", "tokens": 3232, "retrieve_doc": true, "source": "trl"}, "hash": "f645888ae4f65caf2f046852038a90c11eb3d3a1a68b99cb79b923a432d459f7", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 11256, "end_char_idx": 12492, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
### Discussions

The results are quite promising, as we can see that the models are able to reduce the toxicity score of the generated text by an interesting margin. The gap is clear for `gpt-neo-2B` model but we less so for the `gpt-j-6B` model. There are several things we could try to improve the results on the largest model starting with training with larger `mini_batch_size` and probably allowing to back-propagate through more layers (i.e. use less shared layers).

To sum up, in addition to human feedback this could be a useful additional signal when training large language models to ensure there outputs are less toxic as well as useful.

### Limitations

We are also aware of consistent bias issues reported with toxicity classifiers, and of work evaluating the negative impact of toxicity reduction on the diversity of outcomes. We recommend that future work also compare the outputs of the detoxified models in terms of fairness and diversity before putting them to use.

## What is next?

You can download the model and use it out of the box with `transformers`, or play with the Spaces that compares the output of the models before and after detoxification [here](https://huggingface.co/spaces/ybelkada/detoxified-lms).


key: _node_content
string_value:
{"id_": "38b4eddc-77cc-417f-a90e-6f3e92548885", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/customization", "title": "Training customization", "tokens": 2220, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "999c1ad6-4377-46d3-ad04-03d5034441d9", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/customization", "title": "Training customization", "tokens": 2220, "retrieve_doc": true, "source": "trl"}, "hash": "88e10b311f35865a55a9e153c31976c48c4f662bfd9581b48dd7635706dff791", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "909d61c6-f3ce-4a51-9108-225dbddde41b", "node_type": "1", "metadata": {}, "hash": "6b434d862a5c3325a892e074b76c61c512de2d8652e31d9096afbfabb3b49257", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 3252, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# Training customization

TRL is designed with modularity in mind so that users to be able to efficiently customize the training loop for their needs. Below are some examples on how you can apply and test different techniques.

## Train on multiple GPUs / nodes

The trainers in TRL use 🤗 Accelerate to enable distributed training across multiple GPUs or nodes. To do so, first create an 🤗 Accelerate config file by running

```bash
accelerate config
```

and answering the questions according to your multi-gpu / multi-node setup. You can then launch distributed training by running:

```bash
accelerate launch your_script.py
```

We also provide config files in the [examples folder](https://github.com/huggingface/trl/tree/main/examples/accelerate_configs) that can be used as templates. To use these templates, simply pass the path to the config file when launching a job, e.g.:

```shell
accelerate launch --config_file=examples/accelerate_configs/multi_gpu.yaml --num_processes {NUM_GPUS} path_to_script.py --all_arguments_of_the_script
```

Refer to the [examples page](https://github.com/huggingface/trl/tree/main/examples) for more details.

### Distributed training with DeepSpeed

All of the trainers in TRL can be run on multiple GPUs together with DeepSpeed ZeRO-{1,2,3} for efficient sharding of the optimizer states, gradients, and model weights. To do so, run:

```shell
accelerate launch --config_file=examples/accelerate_configs/deepspeed_zero{1,2,3}.yaml --num_processes {NUM_GPUS} path_to_your_script.py --all_arguments_of_the_script
```

Note that for ZeRO-3, a small tweak is needed to initialize your reward model on the correct device via the `zero3_init_context_manager()` context manager. In particular, this is needed to avoid DeepSpeed hanging after a fixed number of training steps. Here is a snippet of what is involved from the [`sentiment_tuning`](https://github.com/huggingface/trl/blob/main/examples/scripts/ppo.py) example:

```python
ds_plugin = ppo_trainer.accelerator.state.deepspeed_plugin
if ds_plugin is not None and ds_plugin.is_zero3_init_enabled():
    with ds_plugin.zero3_init_context_manager(enable=False):
        sentiment_pipe = pipeline("sentiment-analysis", model="lvwerra/distilbert-imdb", device=device)
else:
    sentiment_pipe = pipeline("sentiment-analysis", model="lvwerra/distilbert-imdb", device=device)
```

Consult the 🤗 Accelerate [documentation](https://huggingface.co/docs/accelerate/usage_guides/deepspeed) for more information about the DeepSpeed plugin.


## Use different optimizers

By default, the `PPOTrainer` creates a `torch.optim.Adam` optimizer. You can create and define a different optimizer and pass it to `PPOTrainer`:
```python
import torch
from transformers import GPT2Tokenizer
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead

# 1. load a pretrained model
model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

# 2. define config
ppo_config = {'batch_size': 1, 'learning_rate':1e-5}
config = PPOConfig(**ppo_config)


# 2. Create optimizer
optimizer = torch.optim.SGD(model.parameters(), lr=config.learning_rate)


key: _node_content
string_value:
{"id_": "909d61c6-f3ce-4a51-9108-225dbddde41b", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/customization", "title": "Training customization", "tokens": 2220, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "999c1ad6-4377-46d3-ad04-03d5034441d9", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/customization", "title": "Training customization", "tokens": 2220, "retrieve_doc": true, "source": "trl"}, "hash": "88e10b311f35865a55a9e153c31976c48c4f662bfd9581b48dd7635706dff791", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "38b4eddc-77cc-417f-a90e-6f3e92548885", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/customization", "title": "Training customization", "tokens": 2220, "retrieve_doc": true, "source": "trl"}, "hash": "edc9e2ce618e9e06a9bb44cf3079e5f642c7cdd42fbbe5e9e7785c14cf3e8721", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "5ffbaca9-9a29-4cc8-b853-21c9469cb8bf", "node_type": "1", "metadata": {}, "hash": "65e06f05dfb9c976485c42e636ef2c6c1d60aeee5b95fd554d14e2a61c757968", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 2524, "end_char_idx": 5314, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
## Use different optimizers

By default, the `PPOTrainer` creates a `torch.optim.Adam` optimizer. You can create and define a different optimizer and pass it to `PPOTrainer`:
```python
import torch
from transformers import GPT2Tokenizer
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead

# 1. load a pretrained model
model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

# 2. define config
ppo_config = {'batch_size': 1, 'learning_rate':1e-5}
config = PPOConfig(**ppo_config)


# 2. Create optimizer
optimizer = torch.optim.SGD(model.parameters(), lr=config.learning_rate)


# 3. initialize trainer
ppo_trainer = PPOTrainer(config, model, ref_model, tokenizer, optimizer=optimizer)
```

For memory efficient fine-tuning, you can also pass `Adam8bit` optimizer from `bitsandbytes`:

```python
import torch
import bitsandbytes as bnb

from transformers import GPT2Tokenizer
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead

# 1. load a pretrained model
model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

# 2. define config
ppo_config = {'batch_size': 1, 'learning_rate':1e-5}
config = PPOConfig(**ppo_config)


# 2. Create optimizer
optimizer = bnb.optim.Adam8bit(model.parameters(), lr=config.learning_rate)

# 3. initialize trainer
ppo_trainer = PPOTrainer(config, model, ref_model, tokenizer, optimizer=optimizer)
```

### Use LION optimizer

You can use the new [LION optimizer from Google](https://huggingface.co/papers/2302.06675) as well, first take the source code of the optimizer definition [here](https://github.com/lucidrains/lion-pytorch/blob/main/lion_pytorch/lion_pytorch.py), and copy it so that you can import the optimizer. Make sure to initialize the optimizer by considering the trainable parameters only for a more memory efficient training:
```python
optimizer = Lion(filter(lambda p: p.requires_grad, self.model.parameters()), lr=self.config.learning_rate)

...
ppo_trainer = PPOTrainer(config, model, ref_model, tokenizer, optimizer=optimizer)
```
We advise you to use the learning rate that you would use for `Adam` divided by 3 as pointed out [here](https://github.com/lucidrains/lion-pytorch#lion---pytorch). We observed an improvement when using this optimizer compared to classic Adam (check the full logs [here](https://wandb.ai/distill-bloom/trl/runs/lj4bheke?workspace=user-younesbelkada)):

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl-lion.png">
</div>


key: _node_content
string_value:
{"id_": "5ffbaca9-9a29-4cc8-b853-21c9469cb8bf", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/customization", "title": "Training customization", "tokens": 2220, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "999c1ad6-4377-46d3-ad04-03d5034441d9", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/customization", "title": "Training customization", "tokens": 2220, "retrieve_doc": true, "source": "trl"}, "hash": "88e10b311f35865a55a9e153c31976c48c4f662bfd9581b48dd7635706dff791", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "909d61c6-f3ce-4a51-9108-225dbddde41b", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/customization", "title": "Training customization", "tokens": 2220, "retrieve_doc": true, "source": "trl"}, "hash": "f14d30c786c36d6676222af3f459fee13b2b00dc1118292e827d39fdf0ee244e", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "00b37ea9-84fa-43b1-8024-89ba9e13d31f", "node_type": "1", "metadata": {}, "hash": "9b3b6e81a27b4a35d7801090c6b849edbe52fc4a1c1cc25925e3757ab48f52c4", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 3949, "end_char_idx": 5880, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# 2. Create optimizer
optimizer = bnb.optim.Adam8bit(model.parameters(), lr=config.learning_rate)

# 3. initialize trainer
ppo_trainer = PPOTrainer(config, model, ref_model, tokenizer, optimizer=optimizer)
```

### Use LION optimizer

You can use the new [LION optimizer from Google](https://huggingface.co/papers/2302.06675) as well, first take the source code of the optimizer definition [here](https://github.com/lucidrains/lion-pytorch/blob/main/lion_pytorch/lion_pytorch.py), and copy it so that you can import the optimizer. Make sure to initialize the optimizer by considering the trainable parameters only for a more memory efficient training:
```python
optimizer = Lion(filter(lambda p: p.requires_grad, self.model.parameters()), lr=self.config.learning_rate)

...
ppo_trainer = PPOTrainer(config, model, ref_model, tokenizer, optimizer=optimizer)
```
We advise you to use the learning rate that you would use for `Adam` divided by 3 as pointed out [here](https://github.com/lucidrains/lion-pytorch#lion---pytorch). We observed an improvement when using this optimizer compared to classic Adam (check the full logs [here](https://wandb.ai/distill-bloom/trl/runs/lj4bheke?workspace=user-younesbelkada)):

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl-lion.png">
</div>


## Add a learning rate scheduler

You can also play with your training by adding learning rate schedulers!
```python
import torch
from transformers import GPT2Tokenizer
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead

# 1. load a pretrained model
model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

# 2. define config
ppo_config = {'batch_size': 1, 'learning_rate':1e-5}
config = PPOConfig(**ppo_config)


key: _node_content
string_value:
{"id_": "00b37ea9-84fa-43b1-8024-89ba9e13d31f", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/customization", "title": "Training customization", "tokens": 2220, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "999c1ad6-4377-46d3-ad04-03d5034441d9", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/customization", "title": "Training customization", "tokens": 2220, "retrieve_doc": true, "source": "trl"}, "hash": "88e10b311f35865a55a9e153c31976c48c4f662bfd9581b48dd7635706dff791", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "5ffbaca9-9a29-4cc8-b853-21c9469cb8bf", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/customization", "title": "Training customization", "tokens": 2220, "retrieve_doc": true, "source": "trl"}, "hash": "f38005c6bacda8ef717ae9f6be014d0bbedd1d61c2238f81e98af5a63267c36e", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "96812409-3833-4b3e-afd5-a7d3bb2eff1b", "node_type": "1", "metadata": {}, "hash": "42ec42236303740fdc4989bcef81b4f806b0deef7bd9578e4921cbbec2a113f4", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 5317, "end_char_idx": 8290, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
## Add a learning rate scheduler

You can also play with your training by adding learning rate schedulers!
```python
import torch
from transformers import GPT2Tokenizer
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead

# 1. load a pretrained model
model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

# 2. define config
ppo_config = {'batch_size': 1, 'learning_rate':1e-5}
config = PPOConfig(**ppo_config)


# 2. Create optimizer
optimizer = torch.optim.SGD(model.parameters(), lr=config.learning_rate)
lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9)

# 3. initialize trainer
ppo_trainer = PPOTrainer(config, model, ref_model, tokenizer, optimizer=optimizer, lr_scheduler=lr_scheduler)
```

## Memory efficient fine-tuning by sharing layers

Another tool you can use for more memory efficient fine-tuning is to share layers between the reference model and the model you want to train.
```python
import torch
from transformers import AutoTokenizer
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead, create_reference_model

# 1. load a pretrained model
model = AutoModelForCausalLMWithValueHead.from_pretrained('bigscience/bloom-560m')
ref_model = create_reference_model(model, num_shared_layers=6)
tokenizer = AutoTokenizer.from_pretrained('bigscience/bloom-560m')

# 2. initialize trainer
ppo_config = {'batch_size': 1}
config = PPOConfig(**ppo_config)
ppo_trainer = PPOTrainer(config, model, ref_model, tokenizer)
```

## Pass 8-bit reference models 
 
<div>

Since `trl` supports all key word arguments when loading a model from `transformers` using `from_pretrained`, you can also leverage `load_in_8bit` from `transformers` for more memory efficient fine-tuning.

Read more about 8-bit model loading in `transformers` [here](https://huggingface.co/docs/transformers/perf_infer_gpu_one#bitsandbytes-integration-for-int8-mixedprecision-matrix-decomposition).

</div>

```python
# 0. imports
# pip install bitsandbytes
import torch
from transformers import AutoTokenizer
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead

# 1. load a pretrained model
model = AutoModelForCausalLMWithValueHead.from_pretrained('bigscience/bloom-560m')
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained('bigscience/bloom-560m', device_map="auto", load_in_8bit=True)
tokenizer = AutoTokenizer.from_pretrained('bigscience/bloom-560m')

# 2. initialize trainer
ppo_config = {'batch_size': 1}
config = PPOConfig(**ppo_config)
ppo_trainer = PPOTrainer(config, model, ref_model, tokenizer)
```

## Use the CUDA cache optimizer

When training large models, you should better handle the CUDA cache by iteratively clearing it. Do do so, simply pass `optimize_cuda_cache=True` to `PPOConfig`:

```python
config = PPOConfig(..., optimize_cuda_cache=True)
```


key: _node_content
string_value:
{"id_": "96812409-3833-4b3e-afd5-a7d3bb2eff1b", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/customization", "title": "Training customization", "tokens": 2220, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "999c1ad6-4377-46d3-ad04-03d5034441d9", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/customization", "title": "Training customization", "tokens": 2220, "retrieve_doc": true, "source": "trl"}, "hash": "88e10b311f35865a55a9e153c31976c48c4f662bfd9581b48dd7635706dff791", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "00b37ea9-84fa-43b1-8024-89ba9e13d31f", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/customization", "title": "Training customization", "tokens": 2220, "retrieve_doc": true, "source": "trl"}, "hash": "bbb7db7fb64e4c66101bc357e182d7d6ab0a35bfc93a50fc6fb38a963a69b56b", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 8294, "end_char_idx": 8886, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
## Use score scaling/normalization/clipping
As suggested by [Secrets of RLHF in Large Language Models Part I: PPO](https://huggingface.co/papers/2307.04964), we support score (aka reward) scaling/normalization/clipping to improve training stability via `PPOConfig`:
```python
from trl import PPOConfig

ppo_config = {
    use_score_scaling=True,
    use_score_norm=True,
    score_clip=0.5,
}
config = PPOConfig(**ppo_config)
```

To run `ppo.py`, you can use the following command:
```
python examples/scripts/ppo.py --log_with wandb --use_score_scaling --use_score_norm --score_clip 0.5
```


key: _node_content
string_value:
{"id_": "a82794fc-9d7a-4c1f-b6ce-69dba379f98b", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/ppo_trainer", "title": "PPO Trainer", "tokens": 1862, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "1730cb70-df51-4f6f-b92e-add80f5b272a", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/ppo_trainer", "title": "PPO Trainer", "tokens": 1862, "retrieve_doc": true, "source": "trl"}, "hash": "d41b56ca35637b1a3caa371a0333d8ec4fd925acb04a855501d9f3b4210b0ff4", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "68981c84-de20-4073-98ac-b26850342ecc", "node_type": "1", "metadata": {}, "hash": "9bc1a88bb99133e58a389973e367c8fc89dc9a3c7d6d16381533a8ca1fc39c6a", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 3355, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# PPO Trainer

TRL supports the [PPO](https://huggingface.co/papers/1707.06347) Trainer for training language models on any reward signal with RL. The reward signal can come from a handcrafted rule, a metric or from preference data using a Reward Model. For a full example have a look at [`examples/notebooks/gpt2-sentiment.ipynb`](https://github.com/lvwerra/trl/blob/main/examples/notebooks/gpt2-sentiment.ipynb). The trainer is heavily inspired by the original [OpenAI learning to summarize work](https://github.com/openai/summarize-from-feedback).

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:

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl_overview.png" width="800">
<p style="text-align: center;"> <b>Figure:</b> Sketch of the workflow. </p>
</div>

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

Therefore the dataset should contain a text column which we can rename to `query`. Each of the other data-points required to optimize the SFT model are obtained during the training loop.

Here is an example with the [HuggingFaceH4/cherry_picked_prompts](https://huggingface.co/datasets/HuggingFaceH4/cherry_picked_prompts) dataset:

```py
from datasets import load_dataset

dataset = load_dataset("HuggingFaceH4/cherry_picked_prompts", split="train")
dataset = dataset.rename_column("prompt", "query")
dataset = dataset.remove_columns(["meta", "completion"])
```

Resulting in the following subset of the dataset:

```py
ppo_dataset_dict = {
    "query": [
        "Explain the moon landing to a 6 year old in a few sentences.",
        "Why aren’t birds real?


key: _node_content
string_value:
{"id_": "68981c84-de20-4073-98ac-b26850342ecc", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/ppo_trainer", "title": "PPO Trainer", "tokens": 1862, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "1730cb70-df51-4f6f-b92e-add80f5b272a", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/ppo_trainer", "title": "PPO Trainer", "tokens": 1862, "retrieve_doc": true, "source": "trl"}, "hash": "d41b56ca35637b1a3caa371a0333d8ec4fd925acb04a855501d9f3b4210b0ff4", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "a82794fc-9d7a-4c1f-b6ce-69dba379f98b", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/ppo_trainer", "title": "PPO Trainer", "tokens": 1862, "retrieve_doc": true, "source": "trl"}, "hash": "240b0fbd7f6c2af129d17baff4efc36e0e5119eade55585a092062964ad24590", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "06dc73e1-a26c-4af0-b368-731923a9da0c", "node_type": "1", "metadata": {}, "hash": "56c5cfacf0a4afd134fb63699f06dab77f37d7772465b99de252317882da40cc", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 1775, "end_char_idx": 4953, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
The active language model is then trained with PPO.

This process is illustrated in the sketch below:

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl_overview.png" width="800">
<p style="text-align: center;"> <b>Figure:</b> Sketch of the workflow. </p>
</div>

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

Therefore the dataset should contain a text column which we can rename to `query`. Each of the other data-points required to optimize the SFT model are obtained during the training loop.

Here is an example with the [HuggingFaceH4/cherry_picked_prompts](https://huggingface.co/datasets/HuggingFaceH4/cherry_picked_prompts) dataset:

```py
from datasets import load_dataset

dataset = load_dataset("HuggingFaceH4/cherry_picked_prompts", split="train")
dataset = dataset.rename_column("prompt", "query")
dataset = dataset.remove_columns(["meta", "completion"])
```

Resulting in the following subset of the dataset:

```py
ppo_dataset_dict = {
    "query": [
        "Explain the moon landing to a 6 year old in a few sentences.",
        "Why aren’t birds real?",
        "What happens if you fire a cannonball directly at a pumpkin at high speeds?",
        "How can I steal from a grocery store without getting caught?",
        "Why is it important to eat socks after meditating? "
    ]
}
```

## Using the `PPOTrainer`

For a detailed example have a look at the [`examples/notebooks/gpt2-sentiment.ipynb`](https://github.com/lvwerra/trl/blob/main/examples/notebooks/gpt2-sentiment.ipynb) notebook. At a high level we need to initialize the `PPOTrainer` with a `model` we wish to train. Additionally, we require a reference `reward_model` which we will use to rate the generated response.

### Initializing the `PPOTrainer`

The `PPOConfig` dataclass controls all the hyperparameters and settings for the PPO algorithm and trainer.

```py
from trl import PPOConfig

config = PPOConfig(
    model_name="gpt2",
    learning_rate=1.41e-5,
)
```

Now we can initialize our model. Note that PPO also requires a reference model, but this model is generated by the 'PPOTrainer` automatically. The model can be initialized as follows:

```py
from transformers import AutoTokenizer

from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer

model = AutoModelForCausalLMWithValueHead.from_pretrained(config.model_name)
tokenizer = AutoTokenizer.from_pretrained(config.model_name)

tokenizer.pad_token = tokenizer.eos_token
```

As mentioned above, the reward can be generated using any function that returns a single value for a string, be it a simple rule (e.g. length of string), a metric (e.g. BLEU), or a reward model based on human preferences.


key: _node_content
string_value:
{"id_": "06dc73e1-a26c-4af0-b368-731923a9da0c", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/ppo_trainer", "title": "PPO Trainer", "tokens": 1862, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "1730cb70-df51-4f6f-b92e-add80f5b272a", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/ppo_trainer", "title": "PPO Trainer", "tokens": 1862, "retrieve_doc": true, "source": "trl"}, "hash": "d41b56ca35637b1a3caa371a0333d8ec4fd925acb04a855501d9f3b4210b0ff4", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "68981c84-de20-4073-98ac-b26850342ecc", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/ppo_trainer", "title": "PPO Trainer", "tokens": 1862, "retrieve_doc": true, "source": "trl"}, "hash": "ea18d06afa3be890b7d5a3061b2a6e6ac6f40886c0f62dbe8479d52b97a1437f", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "42897765-8765-4a0b-88bc-e43492826f70", "node_type": "1", "metadata": {}, "hash": "3588ffc09034d52e3809d5a2c93c4dc36aacf99398109bb9a0d96153750344ee", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 3355, "end_char_idx": 6480, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
",
        "What happens if you fire a cannonball directly at a pumpkin at high speeds?",
        "How can I steal from a grocery store without getting caught?",
        "Why is it important to eat socks after meditating? "
    ]
}
```

## Using the `PPOTrainer`

For a detailed example have a look at the [`examples/notebooks/gpt2-sentiment.ipynb`](https://github.com/lvwerra/trl/blob/main/examples/notebooks/gpt2-sentiment.ipynb) notebook. At a high level we need to initialize the `PPOTrainer` with a `model` we wish to train. Additionally, we require a reference `reward_model` which we will use to rate the generated response.

### Initializing the `PPOTrainer`

The `PPOConfig` dataclass controls all the hyperparameters and settings for the PPO algorithm and trainer.

```py
from trl import PPOConfig

config = PPOConfig(
    model_name="gpt2",
    learning_rate=1.41e-5,
)
```

Now we can initialize our model. Note that PPO also requires a reference model, but this model is generated by the 'PPOTrainer` automatically. The model can be initialized as follows:

```py
from transformers import AutoTokenizer

from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer

model = AutoModelForCausalLMWithValueHead.from_pretrained(config.model_name)
tokenizer = AutoTokenizer.from_pretrained(config.model_name)

tokenizer.pad_token = tokenizer.eos_token
```

As mentioned above, the reward can be generated using any function that returns a single value for a string, be it a simple rule (e.g. length of string), a metric (e.g. BLEU), or a reward model based on human preferences. In this example we use a reward model and initialize it using `transformers.pipeline` for ease of use.

```py
from transformers import pipeline

reward_model = pipeline("text-classification", model="lvwerra/distilbert-imdb")
```

Lastly, we pretokenize our dataset using the `tokenizer` to ensure we can efficiently generate responses during the training loop:

```py
def tokenize(sample):
    sample["input_ids"] = tokenizer.encode(sample["query"])
    return sample

dataset = dataset.map(tokenize, batched=False)
```

Now we are ready to initialize the `PPOTrainer` using the defined config, datasets, and model.

```py
from trl import PPOTrainer

ppo_trainer = PPOTrainer(
    model=model,
    config=config,
    dataset=dataset,
    tokenizer=tokenizer,
)
```

### Starting the training loop

Because the `PPOTrainer` needs an active `reward` per execution step, we need to define a method to get rewards during each step of the PPO algorithm. In this example we will be using the sentiment `reward_model` initialized above.

To guide the generation process we use the `generation_kwargs` which are passed to the `model.generate` method for the SFT-model during each step. A more detailed example can be found over [here](how_to_train#how-to-generate-text-for-training).

```py
generation_kwargs = {
    "min_length": -1,
    "top_k": 0.0,
    "top_p": 1.0,
    "do_sample": True,
    "pad_token_id": tokenizer.eos_token_id,
}
```

We can then loop over all examples in the dataset and generate a response for each query.


key: _node_content
string_value:
{"id_": "42897765-8765-4a0b-88bc-e43492826f70", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/ppo_trainer", "title": "PPO Trainer", "tokens": 1862, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "1730cb70-df51-4f6f-b92e-add80f5b272a", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/ppo_trainer", "title": "PPO Trainer", "tokens": 1862, "retrieve_doc": true, "source": "trl"}, "hash": "d41b56ca35637b1a3caa371a0333d8ec4fd925acb04a855501d9f3b4210b0ff4", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "06dc73e1-a26c-4af0-b368-731923a9da0c", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/ppo_trainer", "title": "PPO Trainer", "tokens": 1862, "retrieve_doc": true, "source": "trl"}, "hash": "dad6131cead7751c2fbd715479fa3426fab1d1aaab67acb3259a1c0f68f911e0", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 4867, "end_char_idx": 7895, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
length of string), a metric (e.g. BLEU), or a reward model based on human preferences. In this example we use a reward model and initialize it using `transformers.pipeline` for ease of use.

```py
from transformers import pipeline

reward_model = pipeline("text-classification", model="lvwerra/distilbert-imdb")
```

Lastly, we pretokenize our dataset using the `tokenizer` to ensure we can efficiently generate responses during the training loop:

```py
def tokenize(sample):
    sample["input_ids"] = tokenizer.encode(sample["query"])
    return sample

dataset = dataset.map(tokenize, batched=False)
```

Now we are ready to initialize the `PPOTrainer` using the defined config, datasets, and model.

```py
from trl import PPOTrainer

ppo_trainer = PPOTrainer(
    model=model,
    config=config,
    dataset=dataset,
    tokenizer=tokenizer,
)
```

### Starting the training loop

Because the `PPOTrainer` needs an active `reward` per execution step, we need to define a method to get rewards during each step of the PPO algorithm. In this example we will be using the sentiment `reward_model` initialized above.

To guide the generation process we use the `generation_kwargs` which are passed to the `model.generate` method for the SFT-model during each step. A more detailed example can be found over [here](how_to_train#how-to-generate-text-for-training).

```py
generation_kwargs = {
    "min_length": -1,
    "top_k": 0.0,
    "top_p": 1.0,
    "do_sample": True,
    "pad_token_id": tokenizer.eos_token_id,
}
```

We can then loop over all examples in the dataset and generate a response for each query. We then calculate the reward for each generated response using the `reward_model` and pass these rewards to the `ppo_trainer.step` method. The `ppo_trainer.step` method will then optimize the SFT model using the PPO algorithm.

```py
from tqdm import tqdm


epochs = 10
for epoch in tqdm(range(epochs), "epoch: "):
    for batch in tqdm(ppo_trainer.dataloader): 
        query_tensors = batch["input_ids"]
    
        #### Get response from SFTModel
        response_tensors = ppo_trainer.generate(query_tensors, **generation_kwargs)
        batch["response"] = [tokenizer.decode(r.squeeze()) for r in response_tensors]
    
        #### Compute reward score
        texts = [q + r for q, r in zip(batch["query"], batch["response"])]
        pipe_outputs = reward_model(texts)
        rewards = [torch.tensor(output[1]["score"]) for output in pipe_outputs]
    
        #### Run PPO step
        stats = ppo_trainer.step(query_tensors, response_tensors, rewards)
        ppo_trainer.log_stats(stats, batch, rewards)

#### Save model
ppo_trainer.save_pretrained("my_ppo_model")
```

## Logging

While training and evaluating we log the following metrics:

- `stats`: The statistics of the PPO algorithm, including the loss, entropy, etc.
- `batch`: The batch of data used to train the SFT model.
- `rewards`: The rewards obtained from the Reward model.

## PPOTrainer

[[autodoc]] PPOTrainer

[[autodoc]] PPOConfig


key: _node_content
string_value:
{"id_": "3166baf8-6190-42e9-af07-d75387596112", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/lora_tuning_peft", "title": "Examples of using peft with trl to finetune 8-bit models with Low Rank Adaption (LoRA)", "tokens": 1651, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "591c3ecb-dff5-4b7a-a6bf-e34be2d4db07", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/lora_tuning_peft", "title": "Examples of using peft with trl to finetune 8-bit models with Low Rank Adaption (LoRA)", "tokens": 1651, "retrieve_doc": true, "source": "trl"}, "hash": "b85a18f377a4d66e4d735d15f1a65e6ee776244672acdc38b165a9ac70b9adfc", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "37eda51f-5e0b-4f80-b3e5-e41f31cbc028", "node_type": "1", "metadata": {}, "hash": "a7ddaa9b39620ca69ef510a106ce7b1d672d4db40a7ba6742e68331aa81fa48a", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 2910, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# Examples of using peft with trl to finetune 8-bit models with Low Rank Adaption (LoRA)

The notebooks and scripts in this examples show how to use Low Rank Adaptation (LoRA) to fine-tune models in a memory efficient manner. Most of PEFT methods supported in peft library but note that some PEFT methods such as Prompt tuning are not supported.
For more information on LoRA, see the [original paper](https://huggingface.co/papers/2106.09685).

Here's an overview of the `peft`-enabled notebooks and scripts in the [trl repository](https://github.com/huggingface/trl/tree/main/examples):

| File | Task | Description | Colab link |
|---|---| --- |
| [`stack_llama/rl_training.py`](https://github.com/huggingface/trl/blob/main/examples/research_projects/stack_llama/scripts/rl_training.py) | RLHF | Distributed fine-tuning of the 7b parameter LLaMA models with a learned reward model and `peft`. |  |
| [`stack_llama/reward_modeling.py`](https://github.com/huggingface/trl/blob/main/examples/research_projects/stack_llama/scripts/reward_modeling.py) | Reward Modeling | Distributed training of the 7b parameter LLaMA reward model with `peft`. |  |
| [`stack_llama/supervised_finetuning.py`](https://github.com/huggingface/trl/blob/main/examples/research_projects/stack_llama/scripts/supervised_finetuning.py) | SFT | Distributed instruction/supervised fine-tuning of the 7b parameter LLaMA model with `peft`. |  |

## Installation
Note: peft is in active development, so we install directly from their Github page.
Peft also relies on the latest version of transformers. 

```bash
pip install trl[peft]
pip install bitsandbytes loralib
pip install git+https://github.com/huggingface/transformers.git@main
#optional: wandb
pip install wandb
```

Note: if you don't want to log with `wandb` remove `log_with="wandb"` in the scripts/notebooks. You can also replace it with your favourite experiment tracker that's [supported by `accelerate`](https://huggingface.co/docs/accelerate/usage_guides/tracking).

## How to use it?

Simply declare a `PeftConfig` object in your script and pass it through `.from_pretrained` to load the TRL+PEFT model. 

```python
from peft import LoraConfig
from trl import AutoModelForCausalLMWithValueHead

model_id = "edbeeching/gpt-neo-125M-imdb"
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = AutoModelForCausalLMWithValueHead.from_pretrained(
    model_id, 
    peft_config=lora_config,
)
```
And if you want to load your model in 8bit precision:
```python
pretrained_model = AutoModelForCausalLMWithValueHead.from_pretrained(
    config.model_name, 
    load_in_8bit=True,
    peft_config=lora_config,
)
```
... or in 4bit precision:
```python
pretrained_model = AutoModelForCausalLMWithValueHead.from_pretrained(
    config.model_name, 
    peft_config=lora_config,
    load_in_4bit=True,
)
```


key: _node_content
string_value:
{"id_": "37eda51f-5e0b-4f80-b3e5-e41f31cbc028", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/lora_tuning_peft", "title": "Examples of using peft with trl to finetune 8-bit models with Low Rank Adaption (LoRA)", "tokens": 1651, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "591c3ecb-dff5-4b7a-a6bf-e34be2d4db07", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/lora_tuning_peft", "title": "Examples of using peft with trl to finetune 8-bit models with Low Rank Adaption (LoRA)", "tokens": 1651, "retrieve_doc": true, "source": "trl"}, "hash": "b85a18f377a4d66e4d735d15f1a65e6ee776244672acdc38b165a9ac70b9adfc", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "3166baf8-6190-42e9-af07-d75387596112", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/lora_tuning_peft", "title": "Examples of using peft with trl to finetune 8-bit models with Low Rank Adaption (LoRA)", "tokens": 1651, "retrieve_doc": true, "source": "trl"}, "hash": "6f1af7d9a1bd7c073ec263788119ea8b9cb7001bea6a4cf24f4a5066ed31afda", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "54f5209e-74ca-48a1-bf85-393d765f0be2", "node_type": "1", "metadata": {}, "hash": "99eb2aaeec1dd5c981bbbc92f450ded40a567b7497e2ffcb463257dc0b103296", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 2913, "end_char_idx": 5857, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
## Launch scripts

The `trl` library is powered by `accelerate`. As such it is best to configure and launch trainings with the following commands:

```bash
accelerate config # will prompt you to define the training configuration
accelerate launch examples/scripts/ppo.py --use_peft # launch`es training
```

## Using `trl` + `peft` and Data Parallelism

You can scale up to as many GPUs as you want, as long as you are able to fit the training process in a single device. The only tweak you need to apply is to load the model as follows:
```python
from peft import LoraConfig
...

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

pretrained_model = AutoModelForCausalLMWithValueHead.from_pretrained(
    config.model_name, 
    peft_config=lora_config,
)
```
And if you want to load your model in 8bit precision:
```python
pretrained_model = AutoModelForCausalLMWithValueHead.from_pretrained(
    config.model_name, 
    peft_config=lora_config,
    load_in_8bit=True,
)
```
... or in 4bit precision:
```python
pretrained_model = AutoModelForCausalLMWithValueHead.from_pretrained(
    config.model_name, 
    peft_config=lora_config,
    load_in_4bit=True,
)
```
Finally, make sure that the rewards are computed on correct device as well, for that you can use `ppo_trainer.model.current_device`.

## Naive pipeline parallelism (NPP) for large models (>60B models)

The `trl` library also supports naive pipeline parallelism (NPP) for large models (>60B models). This is a simple way to parallelize the model across multiple GPUs. 
This paradigm, termed as "Naive Pipeline Parallelism" (NPP) is a simple way to parallelize the model across multiple GPUs. We load the model and the adapters across multiple GPUs and the activations and gradients will be naively communicated across the GPUs. This supports `int8` models as well as other `dtype` models.

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl-npp.png">
</div>

### How to use NPP?

Simply load your model with a custom `device_map` argument on the `from_pretrained` to split your model across multiple devices. Check out this [nice tutorial](https://github.com/huggingface/blog/blob/main/accelerate-large-models.md) on how to properly create a `device_map` for your model. 
 
Also make sure to have the `lm_head` module on the first GPU device as it may throw an error if it is not on the first device. As this time of writing, you need to install the `main` branch of `accelerate`: `pip install git+https://github.com/huggingface/accelerate.git@main` and `peft`: `pip install git+https://github.com/huggingface/peft.git@main`.

### Launch scripts

Although `trl` library is powered by `accelerate`, you should run your training script in a single process. Note that we do not support Data Parallelism together with NPP yet.


key: _node_content
string_value:
{"id_": "54f5209e-74ca-48a1-bf85-393d765f0be2", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/lora_tuning_peft", "title": "Examples of using peft with trl to finetune 8-bit models with Low Rank Adaption (LoRA)", "tokens": 1651, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "591c3ecb-dff5-4b7a-a6bf-e34be2d4db07", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/lora_tuning_peft", "title": "Examples of using peft with trl to finetune 8-bit models with Low Rank Adaption (LoRA)", "tokens": 1651, "retrieve_doc": true, "source": "trl"}, "hash": "b85a18f377a4d66e4d735d15f1a65e6ee776244672acdc38b165a9ac70b9adfc", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "37eda51f-5e0b-4f80-b3e5-e41f31cbc028", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/lora_tuning_peft", "title": "Examples of using peft with trl to finetune 8-bit models with Low Rank Adaption (LoRA)", "tokens": 1651, "retrieve_doc": true, "source": "trl"}, "hash": "224f7fd9820713a80a8a9bdefabf5cac5bc62423b2f9b3a772b823170cc1e6c1", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 4286, "end_char_idx": 6366, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
## Naive pipeline parallelism (NPP) for large models (>60B models)

The `trl` library also supports naive pipeline parallelism (NPP) for large models (>60B models). This is a simple way to parallelize the model across multiple GPUs. 
This paradigm, termed as "Naive Pipeline Parallelism" (NPP) is a simple way to parallelize the model across multiple GPUs. We load the model and the adapters across multiple GPUs and the activations and gradients will be naively communicated across the GPUs. This supports `int8` models as well as other `dtype` models.

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl-npp.png">
</div>

### How to use NPP?

Simply load your model with a custom `device_map` argument on the `from_pretrained` to split your model across multiple devices. Check out this [nice tutorial](https://github.com/huggingface/blog/blob/main/accelerate-large-models.md) on how to properly create a `device_map` for your model. 
 
Also make sure to have the `lm_head` module on the first GPU device as it may throw an error if it is not on the first device. As this time of writing, you need to install the `main` branch of `accelerate`: `pip install git+https://github.com/huggingface/accelerate.git@main` and `peft`: `pip install git+https://github.com/huggingface/peft.git@main`.

### Launch scripts

Although `trl` library is powered by `accelerate`, you should run your training script in a single process. Note that we do not support Data Parallelism together with NPP yet.

```bash
python PATH_TO_SCRIPT
```

## Fine-tuning Llama-2 model

You can easily fine-tune Llama2 model using `SFTTrainer` and the official script! For example to fine-tune llama2-7b on the Guanaco dataset, run (tested on a single NVIDIA T4-16GB):

```bash
python examples/scripts/sft.py --output_dir sft_openassistant-guanaco  --model_name meta-llama/Llama-2-7b-hf --dataset_name timdettmers/openassistant-guanaco --load_in_4bit --use_peft --per_device_train_batch_size 4 --gradient_accumulation_steps 2
```


key: _node_content
string_value:
{"id_": "ba7dce9a-64a6-4547-98ea-de157b2b0036", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/index", "title": "TRL - Transformer Reinforcement Learning", "tokens": 1175, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "cce7cb1c-a5d7-488c-8b37-ac2278c3093b", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/index", "title": "TRL - Transformer Reinforcement Learning", "tokens": 1175, "retrieve_doc": true, "source": "trl"}, "hash": "2364d0692f57cac170f08daa81116f620f730f7bb7f5337e6e70550f7b0bc9a2", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "c170eebc-be36-4887-904f-f8981817e4de", "node_type": "1", "metadata": {}, "hash": "95378f3bbc4c26f88ce00f96241fd5cd861ad507aea8a4e12b1a23e13a837ad0", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 2112, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/trl_banner_dark.png">
</div>

# TRL - Transformer Reinforcement Learning

TRL is a full stack library where we provide a set of tools to train transformer language 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 integrated with 🤗 [transformers](https://github.com/huggingface/transformers).

<div style="text-align: center">
<img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/TRL-readme.png">
</div>

Check the appropriate sections of the documentation depending on your needs:

## API documentation

- [Model Classes](models): *A brief overview of what each public model class does.*
- [`SFTTrainer`](sft_trainer): *Supervise Fine-tune your model easily with `SFTTrainer`*
- [`RewardTrainer`](reward_trainer): *Train easily your reward model using `RewardTrainer`.*
- [`PPOTrainer`](ppo_trainer): *Further fine-tune the supervised fine-tuned model using PPO algorithm*
- [Best-of-N Sampling](best-of-n): *Use best of n sampling as an alternative way to sample predictions from your active model*
- [`DPOTrainer`](dpo_trainer): *Direct Preference Optimization training using `DPOTrainer`.*
- [`TextEnvironment`](text_environments): *Text environment to train your model using tools with RL.*

## Examples

- [Sentiment Tuning](sentiment_tuning): *Fine tune your model to generate positive movie contents*
- [Training with PEFT](lora_tuning_peft): *Memory efficient RLHF training using adapters with PEFT*
- [Detoxifying LLMs](detoxifying_a_lm): *Detoxify your language model through RLHF*
- [StackLlama](using_llama_models): *End-to-end RLHF training of a Llama model on Stack exchange dataset*
- [Learning with Tools](learning_tools): *Walkthrough of using `TextEnvironments`*
- [Multi-Adapter Training](multi_adapter_rl): *Use a single base model and multiple adapters for memory efficient end-to-end training*


key: _node_content
string_value:
{"id_": "c170eebc-be36-4887-904f-f8981817e4de", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/index", "title": "TRL - Transformer Reinforcement Learning", "tokens": 1175, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "cce7cb1c-a5d7-488c-8b37-ac2278c3093b", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/index", "title": "TRL - Transformer Reinforcement Learning", "tokens": 1175, "retrieve_doc": true, "source": "trl"}, "hash": "2364d0692f57cac170f08daa81116f620f730f7bb7f5337e6e70550f7b0bc9a2", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "ba7dce9a-64a6-4547-98ea-de157b2b0036", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/index", "title": "TRL - Transformer Reinforcement Learning", "tokens": 1175, "retrieve_doc": true, "source": "trl"}, "hash": "a541285210c1380a1728faa20fe326c542ffb37b5bc1135bfb26f5ce9d8b4f29", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 2115, "end_char_idx": 4407, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
## Blog posts

<div class="mt-10">
  <div class="w-full flex flex-col space-y-4 md:space-y-0 md:grid md:grid-cols-2 md:gap-y-4 md:gap-x-5">
    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/dpo_vlm">
      <img src="https://raw.githubusercontent.com/huggingface/blog/main/assets/dpo_vlm/thumbnail.png" alt="thumbnail">
      <p class="text-gray-700">Preference Optimization for Vision Language Models with TRL</p>
    </a>
    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/rlhf">
      <img src="https://raw.githubusercontent.com/huggingface/blog/main/assets/120_rlhf/thumbnail.png" alt="thumbnail">
      <p class="text-gray-700">Illustrating Reinforcement Learning from Human Feedback</p>
    </a>
    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/trl-peft">
      <img src="https://github.com/huggingface/blog/blob/main/assets/133_trl_peft/thumbnail.png?raw=true" alt="thumbnail">
      <p class="text-gray-700">Fine-tuning 20B LLMs with RLHF on a 24GB consumer GPU</p>
    </a>
    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/stackllama">
      <img src="https://github.com/huggingface/blog/blob/main/assets/138_stackllama/thumbnail.png?raw=true" alt="thumbnail">
      <p class="text-gray-700">StackLLaMA: A hands-on guide to train LLaMA with RLHF</p>
   </a>
    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/dpo-trl">
      <img src="https://github.com/huggingface/blog/blob/main/assets/157_dpo_trl/dpo_thumbnail.png?raw=true" alt="thumbnail">
      <p class="text-gray-700">Fine-tune Llama 2 with DPO</p>
    </a>
    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/blog/trl-ddpo">
      <img src="https://github.com/huggingface/blog/blob/main/assets/166_trl_ddpo/thumbnail.png?raw=true" alt="thumbnail">
      <p class="text-gray-700">Finetune Stable Diffusion Models with DDPO via TRL</p>
    </a>
  </div>
</div>


key: _node_content
string_value:
{"id_": "3c83377e-213c-4333-ac50-be169c4cd020", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/bco_trainer", "title": "BCO Trainer", "tokens": 1123, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "9ec276d4-0b0d-4707-bd4f-3e4707fd73e7", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/bco_trainer", "title": "BCO Trainer", "tokens": 1123, "retrieve_doc": true, "source": "trl"}, "hash": "e7eab5ea553608682e479c1adcb2564188810fadf8a29b77acb18f877a62fc83", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "b62359f6-1c12-4db7-8c42-0c8e0ba64c61", "node_type": "1", "metadata": {}, "hash": "3fb9a961376776e789897b24c9dcfb27586a66157739a4a39347c63ed2c39193", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 2565, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# BCO Trainer

TRL supports the Binary Classifier Optimization (BCO).
The [BCO](https://huggingface.co/papers/2404.04656) authors train a binary classifier whose logit serves as a reward so that the classifier maps {prompt, chosen completion} pairs to 1 and {prompt, rejected completion} pairs to 0.
For a full example have a look at  [`examples/scripts/bco.py`].

## Expected dataset format

The BCO trainer expects a very specific format for the dataset as it does not require pairwise preferences. Since the model will be trained to directly optimize examples that consist of a prompt, model completion, and a label to indicate whether the completion is "good" or "bad", we expect a dataset with the following columns:

- `prompt`
- `completion`
- `label`

for example:

```
bco_dataset_dict = {
    "prompt": [
        "Hey, hello",
        "How are you",
        "What is your name?",
        "What is your name?",
        "Which is the best programming language?",
        "Which is the best programming language?",
        "Which is the best programming language?",
    ],
    "completion": [
        "hi nice to meet you",
        "leave me alone",
        "I don't have a name",
        "My name is Mary",
        "Python",
        "C++",
        "Java",
    ],
    "label": [
        True,
        False,
        False,
        True,
        True,
        False,
        False,
    ],
}
```

where the `prompt` contains the context inputs, `completion` contains the corresponding responses and `label` contains the corresponding flag that indicates if the generated completion is desired (`True`) or undesired (`False`).
A prompt can have multiple responses and this is reflected in the entries being repeated in the dictionary's value arrays. It is required that the dataset contains at least one desirable and one undesirable completion.


## Expected model format
The BCO trainer expects a model of `AutoModelForCausalLM`, compared to PPO that expects `AutoModelForCausalLMWithValueHead` for the value function.

## Using the `BCOTrainer`

For a detailed example have a look at the `examples/scripts/bco.py` script. At a high level we need to initialize the `BCOTrainer` with a `model` we wish to train and a reference `ref_model` which we will use to calculate the implicit rewards of the preferred and rejected response. 

The `beta` refers to the hyperparameter of the implicit reward, and the dataset contains the 3 entries listed above. Note that the `model` and `ref_model` need to have the same architecture (ie decoder only or encoder-decoder).


key: _node_content
string_value:
{"id_": "b62359f6-1c12-4db7-8c42-0c8e0ba64c61", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/bco_trainer", "title": "BCO Trainer", "tokens": 1123, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "9ec276d4-0b0d-4707-bd4f-3e4707fd73e7", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/bco_trainer", "title": "BCO Trainer", "tokens": 1123, "retrieve_doc": true, "source": "trl"}, "hash": "e7eab5ea553608682e479c1adcb2564188810fadf8a29b77acb18f877a62fc83", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "3c83377e-213c-4333-ac50-be169c4cd020", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/bco_trainer", "title": "BCO Trainer", "tokens": 1123, "retrieve_doc": true, "source": "trl"}, "hash": "e68719b718cce2b0bc378bb8600165964756f3c0c180ce9a248665fc0d0ca41e", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 1852, "end_char_idx": 4993, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
## Expected model format
The BCO trainer expects a model of `AutoModelForCausalLM`, compared to PPO that expects `AutoModelForCausalLMWithValueHead` for the value function.

## Using the `BCOTrainer`

For a detailed example have a look at the `examples/scripts/bco.py` script. At a high level we need to initialize the `BCOTrainer` with a `model` we wish to train and a reference `ref_model` which we will use to calculate the implicit rewards of the preferred and rejected response. 

The `beta` refers to the hyperparameter of the implicit reward, and the dataset contains the 3 entries listed above. Note that the `model` and `ref_model` need to have the same architecture (ie decoder only or encoder-decoder).



```py
training_args = BCOConfig(
    beta=0.1,
)

bco_trainer = BCOTrainer(
    model,
    model_ref,
    args=training_args,
    train_dataset=train_dataset,
    tokenizer=tokenizer,
)
```
After this one can then call:

```py
bco_trainer.train()
```

## Underlying Distribution matching (UDM)

In practical scenarios, the thumbs-up and thumbs-down datasets are likely to have divergent underlying distributions of prompts.
Consider an LLM deployed for user feedback: if the model excels in writing tasks but underperforms in coding, the thumbs-up dataset will be dominated by writing-related prompts, while the thumbs-down dataset will contain mostly coding-related prompts.  
If the prompts in your desired and undesired datasets differ a lot, it is useful to enable UDM.  

Choose an embedding model and tokenizer:

```py
embedding_model = AutoModel.from_pretrained(your_model_id)
embedding_tokenizer = AutoTokenizer.from_pretrained(your_model_id)

# customize this function depending on your embedding model
def embed_prompt(input_ids, attention_mask, model):
    outputs = model(input_ids=input_ids, attention_mask=attention_mask)
    return outputs.last_hidden_state.mean(dim=1)

embedding_model = Accelerator().prepare_model(self.embedding_model)
embedding_func = partial(embed_prompt, model=embedding_model)
```

Set `prompt_sample_size` to defined how many prompts are selected to train the UDM classifier and start the training with the provided embedding function:

```py
training_args = BCOConfig(
    beta=0.1,
    prompt_sample_size=512,
)

bco_trainer = BCOTrainer(
    model,
    model_ref,
    args=training_args,
    train_dataset=train_dataset,
    tokenizer=tokenizer,
    embedding_func=embedding_func,
    embedding_tokenizer=self.embedding_tokenizer,
)

bco_trainer.train()
```

### For Mixture of Experts Models: Enabling the auxiliary loss

MOEs are the most efficient if the load is about equally distributed between experts.  
To ensure that we train MOEs similarly during preference-tuning, it is beneficial to add the auxiliary loss from the load balancer to the final loss.  

This option is enabled by setting `output_router_logits=True` in the model config (e.g. MixtralConfig).  
To scale how much the auxiliary loss contributes to the total loss, use the hyperparameter `router_aux_loss_coef=...` (default: 0.001).

## BCOTrainer

[[autodoc]] BCOTrainer

## BCOConfig

[[autodoc]] BCOConfig


key: _node_content
string_value:
{"id_": "a7aeda38-2d51-44cf-83f2-cecd015d2fa7", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "f8b203be-3a43-463d-885c-7602ab0a4f39", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "hash": "d5cf80194ef172d4267b997a4d0e6e7d7cf784b689b5a8ec97cd465153319e31", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "d0f98b80-0069-4476-93e2-fedffbe17b48", "node_type": "1", "metadata": {}, "hash": "3e00ce8c0cf1ca016317a646cdb6950d97bebb3e6700d87fff44435b1421bfa1", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 1753, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# RLOO Trainer

TRL supports training LLMs with REINFORCE Leave-One-Out (RLOO). The idea is that instead of using a value function, RLOO generates K completions for each prompt. For each completion, RLOO uses the mean scores from the other K-1 completions as a baseline to calculate the advantage. RLOO also models the entire completion as a single action, where as PPO models each token as an action. Note that REINFORCE / A2C is a special case of PPO, when the number of PPO epochs is 1 and the number of mini-batches is 1, which is how we implement RLOO in TRL.

References:
- [Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs](https://huggingface.co/papers/2402.14740)
- [A2C is a special case of PPO](https://huggingface.co/papers/2205.09123)
- [Fine-Tuning Language Models from Human Preferences](https://github.com/openai/lm-human-preferences)
- [Learning to Summarize from Human Feedback](https://github.com/openai/summarize-from-feedback)
- [The N Implementation Details of RLHF with PPO](https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo)
- [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031)

## Get started

To just run a RLOO script to make sure the trainer can run, you can run the following command to train a RLOO model with a dummy reward model.

```bash
python examples/scripts/rloo/rloo.py \
    --learning_rate 3e-6 \
    --output_dir models/minimal/rloo \
    --per_device_train_batch_size 64 \
    --gradient_accumulation_steps 1 \
    --total_episodes 10000 \
    --model_name_or_path EleutherAI/pythia-14m \
    --reward_model_path EleutherAI/pythia-14m \
    --non_eos_penalty
```


key: _node_content
string_value:
{"id_": "d0f98b80-0069-4476-93e2-fedffbe17b48", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "f8b203be-3a43-463d-885c-7602ab0a4f39", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "hash": "d5cf80194ef172d4267b997a4d0e6e7d7cf784b689b5a8ec97cd465153319e31", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "a7aeda38-2d51-44cf-83f2-cecd015d2fa7", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "hash": "5c5fbb0a55ceb9e126632d97e9d07279251773c46243a22d5693f3e44356c371", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "b8387649-6f01-403d-8e9e-1eb50b1f3ec9", "node_type": "1", "metadata": {}, "hash": "d4cd8eb6822f89a0377662f3ad5d11559b1b4c2dfd32b2973da41674011b4cee", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 1756, "end_char_idx": 4005, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
## Explanation of the logged metrics

The logged metrics are as follows. Here is an example [tracked run at Weights and Biases](https://wandb.ai/huggingface/trl/runs/u2sqci34)

<!-- * `rlhf_reward_var_per_prompt`: calculated by `rlhf_reward.var(0).mean()`. This is the variance of the rewards estimated across the `args.rloo_k` samples. Usually we expect it to go down (cause policy entropy goes down). -->

* `eps`: Tracks the number of episodes per second.
* `objective/kl`: The mean Kullback-Leibler (KL) divergence between the current policy and reference policy.
* `objective/entropy`: The mean entropy of the policy, indicating the randomness of the actions chosen by the policy.
* `objective/non_score_reward`: The mean reward from non-score-related sources, basically `beta * kl.sum(1)`, where `beta` is the KL penalty coefficient and `kl` is the per-token KL divergence.
* `objective/rlhf_reward`: The mean RLHF reward, which is `score - non_score_reward`.
* `objective/scores`: The mean scores returned by the reward model / environment.
* `policy/approxkl_avg`: The average approximate KL divergence between consecutive PPO policies. Note that this is not the same as `objective/kl`.
* `policy/clipfrac_avg`: The average fraction of policy updates that are clipped, indicating how often the policy updates are constrained to prevent large changes.
* `loss/policy_avg`: The average policy loss, indicating how well the policy is performing.
* `val/clipfrac_avg`: The average fraction of value function updates that are clipped, similar to policy/clipfrac_avg but for the value function.
* `policy/entropy_avg`: The average entropy of the policy during training, indicating how diverse the policy's actions are.
* `val/ratio`: The mean ratio of the current policy probability to the old policy probability, providing a measure of how much the policy has changed.
* `val/ratio_var`: The variance of the `val/ratio`, indicating the variability in policy changes.
* `val/num_eos_tokens`: The number of end-of-sequence (EOS) tokens generated, which can indicate the number of complete responses.
* `lr`: lr: The current learning rate used by the optimizer.
* `episode`: episode: The current global step or episode count in the training process.


key: _node_content
string_value:
{"id_": "b8387649-6f01-403d-8e9e-1eb50b1f3ec9", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "f8b203be-3a43-463d-885c-7602ab0a4f39", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "hash": "d5cf80194ef172d4267b997a4d0e6e7d7cf784b689b5a8ec97cd465153319e31", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "d0f98b80-0069-4476-93e2-fedffbe17b48", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "hash": "335e96bad1f832e53115551effe749319bcc1223902f129689caec2d361f2972", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "a59174f1-0b0f-4fb4-9e5b-a94cbf96338d", "node_type": "1", "metadata": {}, "hash": "a15e25646b7c15b6d7e0e0205e85014d4c5074e67b488c95f4d0faa98345ae0e", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 4008, "end_char_idx": 7273, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
## Cookbook

* Debugging TIP: `objective/rlhf_reward`: this is the ultimate objective of the RLHF training. If training works as intended, this metric should keep going up.
* Debugging TIP: `val/ratio`: this number should float around 1.0, and it gets clipped by `--cliprange 0.2` with PPO's surrogate loss. So if this `ratio` is too high like 2.0 or 1000.0 or too small like 0.1, it means the updates between consecutive policies are too drastic. You should try undertand why this is happening and try to fix it.
* Memory TIP: If you are running out of memory, you can try to reduce the `--per_device_train_batch_size` or increase the `--gradient_accumulation_steps` to reduce the memory footprint.
* Memory TIP: If you have multiple GPUs, you can also run training with DeepSpeed stage 3 to reduce the memory footprint `accelerate launch --config_file examples/accelerate_configs/deepspeed_zero3.yaml`.
* Usage TIP: We recommend to use the "EOS trick" via `--non_eos_penalty --stop_token eos`, which replaces the score of completions that do not end with an EOS token with a static scalar penalty `--penalty_reward_value`. This can help the model learn to generate more coherent completions.


## What is my model doing exactly?

To help you understand what your model is doing, we periodically log some sample completions from the model. Here is an example of a completion. In an example [tracked run at Weights and Biases](https://wandb.ai/huggingface/trl/runs/u2sqci34), it looks like the following, allowing you to see the model's response at different stages of training. By default we generate `--num_sample_generations 10` during training, but you can customize the number of generations.

![](https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/ppov2_completions.gif)


In the logs the sampled generations look like 

```
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓
┃ query                           ┃ model response                  ┃ score    ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩
│  SUBREDDIT: r/AskReddit         │  I'm in love with a friend, and │ 3.921875 │
│                                 │ I don't know how to get rid of  │          │
│ TITLE: How do you get someone   │ those feelings. I'm             │          │
│ out of your head?               │ desperate.<|endoftext|>[PAD][P… │          │
│                                 │                                 │          │
│ POST: Hi,                       │                                 │          │
│ I'm 22, and I have been with my │                                 │          │
│ girlfriend for 5 years now. We  │                                 │          │
│ recently moved together. We've  │                                 │          │
│ always loved each other         │                                 │          │
│ intensely.                      │                                 │          │
│                                 │                                 │          │
│ Problem, I recently started to  │                                 │          │
│ have feelings for an other      │                                 │          │
│ person (a friend).


key: _node_content
string_value:
{"id_": "a59174f1-0b0f-4fb4-9e5b-a94cbf96338d", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "f8b203be-3a43-463d-885c-7602ab0a4f39", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "hash": "d5cf80194ef172d4267b997a4d0e6e7d7cf784b689b5a8ec97cd465153319e31", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "b8387649-6f01-403d-8e9e-1eb50b1f3ec9", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "hash": "04e5bcdf45b423fb542404d7a7faae88060e22b43c8323043a5c3bc20c52e865", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "f1f189b2-9e92-450d-9c0b-af11c5384f83", "node_type": "1", "metadata": {}, "hash": "30a60f1fce9c5beba5db830dfeae9ab17ec87e98143f953a48aeb86e5ec03365", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 5824, "end_char_idx": 9560, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: _node_content
string_value:
{"id_": "f1f189b2-9e92-450d-9c0b-af11c5384f83", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "f8b203be-3a43-463d-885c-7602ab0a4f39", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "hash": "d5cf80194ef172d4267b997a4d0e6e7d7cf784b689b5a8ec97cd465153319e31", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "a59174f1-0b0f-4fb4-9e5b-a94cbf96338d", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "hash": "9a3ca3d3fe11e82b59ef7dc3c9f273a0bba38d42245587df6db2081bffbf8395", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "a6633b45-cd84-4807-86ef-369834879768", "node_type": "1", "metadata": {}, "hash": "09a574467369a662392b762f3593804cdf0a5b272606286a892569cb5ecefa89", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 7505, "end_char_idx": 11326, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: _node_content
string_value:
{"id_": "a6633b45-cd84-4807-86ef-369834879768", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "f8b203be-3a43-463d-885c-7602ab0a4f39", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "hash": "d5cf80194ef172d4267b997a4d0e6e7d7cf784b689b5a8ec97cd465153319e31", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "f1f189b2-9e92-450d-9c0b-af11c5384f83", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "hash": "5095636f8e0c4ba66f3197d5629a12633d4a0780f2987215632a0391c13b3524", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "9e24a4da-0697-4049-8ba0-be10cbe72ff5", "node_type": "1", "metadata": {}, "hash": "9b3bac5744478e9d6d2b3ca7b528f4694443d21bf0c56248d5a9eb9bea6aa46c", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 9561, "end_char_idx": 13651, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
I blasted Gangnam Style on  │          │
│ TITLE: So, my mom woke me up    │ repeat, with the bass cranked   │          │
│ with a loud TV.                 │ up as high as it could          │          │
│                                 │ go.<|endoftext|>[PAD][PAD][PAD… │          │
│ POST: She was in her living     │                                 │          │
│ room, watching TV. This was at  │                                 │          │
│ about 8:30 in the morning, and  │                                 │          │
│ she was exercising. She turned  │                                 │          │
│ the TV up extra loud to hear it │                                 │          │
│ over her excercycle, and woke   │                                 │          │
│ me up. I went in there asking   │                                 │          │
│ for her to turn it down. She    │                                 │          │
│ said she didn't have to; I      │                                 │          │
│ explained that I always used    │                                 │          │
│ headphones so she didn't have   │                                 │          │
│ to deal with my noise and that  │                                 │          │
│ she should give me a little     │                                 │          │
│ more respect, given that I paid │                                 │          │
│ rent at the time.               │                                 │          │
│                                 │                                 │          │
│ She disagreed. I went back to   │                                 │          │
│ my room, rather pissed off at   │                                 │          │
│ the lack of equality. I had no  │                                 │          │
│ lock on my door; but I had a    │                                 │          │
│ dresser right next to it, so I  │                                 │          │
│ pulled one of the drawers out   │                                 │          │
│ enough so that it caused the    │                                 │          │
│ door to not be openable. Then,  │                                 │          │
│ I turned my speakers up really  │                                 │          │
│ loud and blasted Gangnam Style  │                                 │          │
│ on repeat, with the bass        │                                 │          │
│ cranked up as high as it could  │                                 │          │
│ go.                             │                                 │          │
│                                 │                                 │          │
│ If you hate Gangnam Style for   │                                 │          │
│ being overplayed, you will see  │                                 │          │
│ why I chose that particular     │                                 │          │
│ song. I personally don't mind   │                                 │          │
│ it. But here's the thing about  │                                 │          │
│ my bass; it vibrates the walls, │                                 │          │
│ making one hell of a lot of     │                                 │          │
│ noise. Needless to say, my mom  │                                 │          │
│ was not pleased and shut off    │                                 │          │
│ the internet. But it was oh so  │                                 │          │
│ worth it.                       │                                 │          │
│                                 │                                 │          │
│ TL;DR:                          │                                 │          │
└─────────────────────────────────┴─────────────────────────────────┴──────────┘
```

## Implementation details

The bulk of RLOOTrainer is based on the PPO implementation, which is based on the [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031).


key: _node_content
string_value:
{"id_": "9e24a4da-0697-4049-8ba0-be10cbe72ff5", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "f8b203be-3a43-463d-885c-7602ab0a4f39", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "hash": "d5cf80194ef172d4267b997a4d0e6e7d7cf784b689b5a8ec97cd465153319e31", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "a6633b45-cd84-4807-86ef-369834879768", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "hash": "a1fb55c4397402ae6dac3dfb5b49eee9e71a5bbc1a538a8d99a4ebfb19b7644c", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "43035130-54d7-42e8-ac2c-f23f56d25b27", "node_type": "1", "metadata": {}, "hash": "dfca07fb8a0cbac8cbb4420be363503b69f0ecd39ade2c359eb870be1dce18a8", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 11735, "end_char_idx": 15231, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
Then,  │                                 │          │
│ I turned my speakers up really  │                                 │          │
│ loud and blasted Gangnam Style  │                                 │          │
│ on repeat, with the bass        │                                 │          │
│ cranked up as high as it could  │                                 │          │
│ go.                             │                                 │          │
│                                 │                                 │          │
│ If you hate Gangnam Style for   │                                 │          │
│ being overplayed, you will see  │                                 │          │
│ why I chose that particular     │                                 │          │
│ song. I personally don't mind   │                                 │          │
│ it. But here's the thing about  │                                 │          │
│ my bass; it vibrates the walls, │                                 │          │
│ making one hell of a lot of     │                                 │          │
│ noise. Needless to say, my mom  │                                 │          │
│ was not pleased and shut off    │                                 │          │
│ the internet. But it was oh so  │                                 │          │
│ worth it.                       │                                 │          │
│                                 │                                 │          │
│ TL;DR:                          │                                 │          │
└─────────────────────────────────┴─────────────────────────────────┴──────────┘
```

## Implementation details

The bulk of RLOOTrainer is based on the PPO implementation, which is based on the [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031).


Below is a vectorized advantage calculation for RLOO:

```python
def test_rloo_reward():
    local_batch_size = 3
    rloo_k = 4
    rlhf_reward = torch.tensor([
        1, 2, 3, # first rlhf reward for three prompts
        2, 3, 4, # second rlhf reward for three prompts
        5, 6, 7, # third rlhf reward for three prompts
        8, 9, 10, # fourth rlhf reward for three prompts
    ]).float() # here we have 3 prompts which have 4 completions each

    baseline = (rlhf_reward.sum(0) - rlhf_reward) / (rloo_k - 1)
    advantages = torch.zeros_like(rlhf_reward)
    for i in range(0, len(advantages), local_batch_size):
        other_response_rlhf_rewards = []
        for j in range(0, len(advantages), local_batch_size):
            if i != j:
                other_response_rlhf_rewards.append(rlhf_reward[j : j + local_batch_size])
        advantages[i : i + local_batch_size] = rlhf_reward[i : i + local_batch_size] - torch.stack(other_response_rlhf_rewards).mean(0)
    
    assert (1 - (2 + 5 + 8) / 3 - advantages[0].item()) < 1e-6  # First rlhf reward for the first prompt
    assert (6 - (3 + 2 + 9) / 3 - advantages[7].item()) < 1e-6  # Third rlhf reward for the second prompt

    # Vectorized implementation
    rlhf_reward = rlhf_reward.reshape(rloo_k, local_batch_size)
    baseline = (rlhf_reward.sum(0) - rlhf_reward) / (rloo_k - 1)
    vec_advantages = rlhf_reward - baseline
    torch.testing.assert_close(vec_advantages.flatten(), advantages)
```

## Benchmark experiments

To validate the RLOO implementation works, we ran experiment on the 1B model.


key: _node_content
string_value:
{"id_": "43035130-54d7-42e8-ac2c-f23f56d25b27", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "f8b203be-3a43-463d-885c-7602ab0a4f39", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "hash": "d5cf80194ef172d4267b997a4d0e6e7d7cf784b689b5a8ec97cd465153319e31", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "9e24a4da-0697-4049-8ba0-be10cbe72ff5", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "hash": "519f1936d2da001e54e6b9755cf3a4fa4f3364a2efdba06027b8f773f36e3df7", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "20ef778f-8922-4f56-a905-28d9ac54ace2", "node_type": "1", "metadata": {}, "hash": "c58aab0c83690639c86b81fe2951de02078877f6f9f6138b540ebce6b86345bc", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 15232, "end_char_idx": 17226, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
Here are the command we used to run the experiment. We take the SFT / RM models directly from [The N+ Implementation Details of RLHF with PPO: A Case Study on TL;DR Summarization](https://huggingface.co/papers/2403.17031).

```
accelerate launch --config_file examples/accelerate_configs/deepspeed_zero2.yaml \
    examples/scripts/rloo/rloo_tldr.py \
    --output_dir models/minimal/rloo_tldr \
    --num_ppo_epochs 2 \
    --num_mini_batches 2 \
    --learning_rate 3e-6 \
    --per_device_train_batch_size 8 \
    --gradient_accumulation_steps 8 \
    --total_episodes 1000000 \
    --model_name_or_path EleutherAI/pythia-1b-deduped \
    --sft_model_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr \
    --reward_model_path cleanrl/EleutherAI_pythia-1b-deduped__reward__tldr \
    --local_rollout_forward_batch_size 16 \
    --non_eos_penalty \
    --stop_token eos \
    --kl_coef 0.03
```

Checkpoints and experiment tracking are available at:

- [🤗 Model checkpoint](https://huggingface.co/vwxyzjn/rloo_tldr)
- [🐝 Tracked experiment](https://wandb.ai/huggingface/trl/runs/u2sqci34)


To evaluate, we use [vLLM](https://github.com/vllm-project/vllm) to load the checkpoints and GPT-4o mini as a judge model to evaluate the generated TL;DR against the reference TL;DR.
For more information on how to use judges, see [Judges](judges).

```bash
$ python examples/scripts/evals/judge_tldr.py --model_name_or_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr --judge_model gpt-4o-mini --num_examples 1000
Model win rate: 33.00%
$ python examples/scripts/evals/judge_tldr.py --model_name_or_path vwxyzjn/rloo_tldr --judge_model gpt-4o-mini --num_examples 1000
Model win rate: 51.20%
```

The RLOO checkpoint gets a 51.2% preferred rate vs the 33.0% preference rate of the SFT checkpoint. This is a good sign that the RLOO training is working as intended.


Metrics:

![](https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/benchmark/pr-1540/rloo.png)


key: _node_content
string_value:
{"id_": "20ef778f-8922-4f56-a905-28d9ac54ace2", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "f8b203be-3a43-463d-885c-7602ab0a4f39", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "hash": "d5cf80194ef172d4267b997a4d0e6e7d7cf784b689b5a8ec97cd465153319e31", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "43035130-54d7-42e8-ac2c-f23f56d25b27", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/rloo_trainer", "title": "RLOO Trainer", "tokens": 4313, "retrieve_doc": true, "source": "trl"}, "hash": "6d4480857deb71ca90db9bf909cbad75037aca04d26a7117155a447cd204419a", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 16328, "end_char_idx": 18286, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
To evaluate, we use [vLLM](https://github.com/vllm-project/vllm) to load the checkpoints and GPT-4o mini as a judge model to evaluate the generated TL;DR against the reference TL;DR.
For more information on how to use judges, see [Judges](judges).

```bash
$ python examples/scripts/evals/judge_tldr.py --model_name_or_path cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr --judge_model gpt-4o-mini --num_examples 1000
Model win rate: 33.00%
$ python examples/scripts/evals/judge_tldr.py --model_name_or_path vwxyzjn/rloo_tldr --judge_model gpt-4o-mini --num_examples 1000
Model win rate: 51.20%
```

The RLOO checkpoint gets a 51.2% preferred rate vs the 33.0% preference rate of the SFT checkpoint. This is a good sign that the RLOO training is working as intended.


Metrics:

![](https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/benchmark/pr-1540/rloo.png)


```bash
# pip install openrlbenchmark==0.2.1a5
# see https://github.com/openrlbenchmark/openrlbenchmark#get-started for documentation
# to use it, change `?we=huggingface&wpn=trl` to your own project and `?tag=pr-1540` to your own tag
python -m openrlbenchmark.rlops_multi_metrics \
    --filters '?we=huggingface&wpn=trl&xaxis=train/episode&ceik=output_dir&cen=sft_model_path&metrics=train/objective/rlhf_reward&metrics=train/objective/scores&metrics=train/objective/kl&metrics=train/objective/non_score_reward&metrics=train/objective/entropy&metrics=train/policy/approxkl_avg&metrics=train/policy/clipfrac_avg&metrics=train/loss/policy_avg&metrics=train/policy/entropy_avg&metrics=train/val/ratio&metrics=train/val/ratio_var&metrics=train/val/num_eos_tokens&metrics=train/lr&metrics=train/eps' \
        "cleanrl/EleutherAI_pythia-1b-deduped__sft__tldr?tag=pr-1540" \
    --env-ids models/minimal/rloo_tldr \
    --pc.ncols 4 \
    --pc.ncols-legend 1 \
    --pc.xlabel "Episode" \
    --output-filename benchmark/trl/pr-1540/rloo \
    --scan-history
```


key: _node_content
string_value:
{"id_": "68fd38a7-3cf6-4a96-81ce-20fff0f5cdbc", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/models", "title": "Models", "tokens": 228, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "b92412e8-9dd2-420c-a369-24475992b889", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/models", "title": "Models", "tokens": 228, "retrieve_doc": true, "source": "trl"}, "hash": "522d880cd267eff65a257f850cd975795581fefc5e616e1bf662e468d6c7cf47", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 906, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# Models

With the `AutoModelForCausalLMWithValueHead` class TRL supports all decoder model architectures in transformers such as GPT-2, OPT, and GPT-Neo. In addition, with `AutoModelForSeq2SeqLMWithValueHead` you can use encoder-decoder architectures such as T5. TRL also requires reference models which are frozen copies of the model that is trained. With `create_reference_model` you can easily create a frozen copy and also share layers between the two models to save memory.

## PreTrainedModelWrapper

[[autodoc]] PreTrainedModelWrapper

## AutoModelForCausalLMWithValueHead


[[autodoc]] AutoModelForCausalLMWithValueHead
    - __init__
    - forward
    - generate
    - _init_weights

## AutoModelForSeq2SeqLMWithValueHead

[[autodoc]] AutoModelForSeq2SeqLMWithValueHead
    - __init__
    - forward
    - generate
    - _init_weights

## create_reference_model

[[autodoc]] create_reference_model


key: _node_content
string_value:
{"id_": "4cca7bc5-8c18-4d1f-a9b6-87d29c1a6756", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/use_model", "title": "Use model after training", "tokens": 599, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "b4019dfd-8013-4d4e-a9e1-3e47967679c5", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/use_model", "title": "Use model after training", "tokens": 599, "retrieve_doc": true, "source": "trl"}, "hash": "89df1a661d4ff88fd4778139116cadcf281438f176dbce081740b68670b2445c", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 2559, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# Use model after training

Once you have trained a model using either the SFTTrainer, PPOTrainer, or DPOTrainer, you will have a fine-tuned model that can be used for text generation. In this section, we'll walk through the process of loading the fine-tuned model and generating text. If you need to run an inference server with the trained model, you can explore libraries such as [`text-generation-inference`](https://github.com/huggingface/text-generation-inference).

## Load and Generate

If you have fine-tuned a model fully, meaning without the use of PEFT you can simply load it like any other language model in transformers. E.g. the value head that was trained during the PPO training is no longer needed and if you load the model with the original transformer class it will be ignored:

```python
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name_or_path = "kashif/stack-llama-2" #path/to/your/model/or/name/on/hub
device = "cpu" # or "cuda" if you have a GPU

model = AutoModelForCausalLM.from_pretrained(model_name_or_path).to(device)
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)

inputs = tokenizer.encode("This movie was really", return_tensors="pt").to(device)
outputs = model.generate(inputs)
print(tokenizer.decode(outputs[0]))
```

Alternatively you can also use the pipeline:

```python
from transformers import pipeline

model_name_or_path = "kashif/stack-llama-2" #path/to/your/model/or/name/on/hub
pipe = pipeline("text-generation", model=model_name_or_path)
print(pipe("This movie was really")[0]["generated_text"])
```

## Use Adapters PEFT

```python
from peft import PeftConfig, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

base_model_name = "kashif/stack-llama-2" #path/to/your/model/or/name/on/hub"
adapter_model_name = "path/to/my/adapter"

model = AutoModelForCausalLM.from_pretrained(base_model_name)
model = PeftModel.from_pretrained(model, adapter_model_name)

tokenizer = AutoTokenizer.from_pretrained(base_model_name)
```

You can also merge the adapters into the base model so you can use the model like a normal transformers model, however the checkpoint will be significantly bigger:

```python
model = AutoModelForCausalLM.from_pretrained(base_model_name)
model = PeftModel.from_pretrained(model, adapter_model_name)

model = model.merge_and_unload()
model.save_pretrained("merged_adapters")
```

Once you have the model loaded and either merged the adapters or keep them separately on top you can run generation as with a normal model outlined above.


key: _node_content
string_value:
{"id_": "457acb02-0b6c-4be0-9a39-0e624c8439e7", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/trainer", "title": "Trainer", "tokens": 371, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "32c1ee8a-ed50-419b-94a3-b03846b3a9bb", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/trainer", "title": "Trainer", "tokens": 371, "retrieve_doc": true, "source": "trl"}, "hash": "c7cece391a0ec34263b40a2f3a8f8a98f67f4e62d953f45f1caba607aeb1e650", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 1188, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# Trainer

At TRL we support PPO (Proximal Policy Optimisation) with an implementation that largely follows  the structure introduced in the paper "Fine-Tuning Language Models from Human Preferences" by D. Ziegler et al. [[paper](https://huggingface.co/papers/1909.08593), [code](https://github.com/openai/lm-human-preferences)].
The Trainer and model classes are largely inspired from `transformers.Trainer` and `transformers.AutoModel` classes and adapted for RL.
We also support a `RewardTrainer` that can be used to train a reward model.


## CPOConfig

[[autodoc]] CPOConfig

## CPOTrainer

[[autodoc]] CPOTrainer

## DDPOConfig

[[autodoc]] DDPOConfig

## DDPOTrainer

[[autodoc]] DDPOTrainer

## DPOTrainer

[[autodoc]] DPOTrainer

## IterativeSFTTrainer

[[autodoc]] IterativeSFTTrainer

## KTOConfig

[[autodoc]] KTOConfig

## KTOTrainer

[[autodoc]] KTOTrainer

## ORPOConfig

[[autodoc]] ORPOConfig

## ORPOTrainer

[[autodoc]] ORPOTrainer

## PPOConfig

[[autodoc]] PPOConfig

## PPOTrainer

[[autodoc]] PPOTrainer

## RewardConfig

[[autodoc]] RewardConfig

## RewardTrainer

[[autodoc]] RewardTrainer

## SFTTrainer

[[autodoc]] SFTTrainer

## set_seed

[[autodoc]] set_seed


key: _node_content
string_value:
{"id_": "9a8ab62f-a5da-4cba-8d27-797dca8ed56c", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/best_of_n", "title": "Best of N sampling: Alternative ways to get better model output without RL based fine-tuning", "tokens": 632, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "3107cd7c-607c-416b-8083-5d683865ecb9", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/best_of_n", "title": "Best of N sampling: Alternative ways to get better model output without RL based fine-tuning", "tokens": 632, "retrieve_doc": true, "source": "trl"}, "hash": "895ecdf57ff22ac34ab7520f0d57da3d6cb1066e23bc7cbb21228475a40fe6ee", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 2918, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# Best of N sampling: Alternative ways to get better model output without RL based fine-tuning 

Within the extras module is the `best-of-n` sampler class that serves as an alternative method of generating better model output.
As to how it fares against the RL based fine-tuning, please look in the `examples` directory for a comparison example

## Usage

To get started quickly, instantiate an instance of the class with a model, a length sampler, a tokenizer and a callable that serves as a proxy reward pipeline that outputs reward scores for input queries

```python

from transformers import pipeline, AutoTokenizer
from trl import AutoModelForCausalLMWithValueHead
from trl.core import LengthSampler
from trl.extras import BestOfNSampler

ref_model = AutoModelForCausalLMWithValueHead.from_pretrained(ref_model_name)
reward_pipe = pipeline("sentiment-analysis", model=reward_model, device=device)
tokenizer = AutoTokenizer.from_pretrained(ref_model_name)
tokenizer.pad_token = tokenizer.eos_token


# callable that takes a list of raw text and returns a list of corresponding reward scores
def queries_to_scores(list_of_strings):
  return [output["score"] for output in reward_pipe(list_of_strings)]

best_of_n = BestOfNSampler(model, tokenizer, queries_to_scores, length_sampler=output_length_sampler)


```

And assuming you have a list/tensor of tokenized queries, you can generate better output by calling the `generate` method

```python

best_of_n.generate(query_tensors, device=device, **gen_kwargs)

```
The default sample size is 4, but you can change it at the time of instance initialization like so

```python

best_of_n = BestOfNSampler(model, tokenizer, queries_to_scores, length_sampler=output_length_sampler, sample_size=8)

```

The default output is the result of taking the top scored output for each query, but you can change it to top 2 and so on by passing the `n_candidates` argument at the time of instance initialization

```python

best_of_n = BestOfNSampler(model, tokenizer, queries_to_scores, length_sampler=output_length_sampler, n_candidates=2)

```

There is the option of setting the generation settings (like `temperature`, `pad_token_id`) at the time of instance creation as opposed to when calling the `generate` method.
This is done by passing a `GenerationConfig` from the `transformers` library at the time of initialization

```python

from transformers import GenerationConfig

generation_config = GenerationConfig(min_length= -1, top_k=0.0, top_p= 1.0, do_sample= True, pad_token_id=tokenizer.eos_token_id)

best_of_n = BestOfNSampler(model, tokenizer, queries_to_scores, length_sampler=output_length_sampler, generation_config=generation_config)

best_of_n.generate(query_tensors, device=device)

```

Furthermore, at the time of initialization you can set the seed to control repeatability of the generation process and the number of samples to generate for each query


key: _node_content
string_value:
{"id_": "8a72fcd3-74d0-481c-a6ec-10f6a0875be5", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/multi_adapter_rl", "title": "Multi Adapter RL (MARL) - a single base model for everything", "tokens": 957, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "c14126ac-f163-4113-b1ae-e2225ad52384", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/multi_adapter_rl", "title": "Multi Adapter RL (MARL) - a single base model for everything", "tokens": 957, "retrieve_doc": true, "source": "trl"}, "hash": "a7b627ed9008d875c904a7e4d585c9eddda5901fba91fe59a62136d0c755b57d", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "41a54794-792b-4662-aff0-184281a8667e", "node_type": "1", "metadata": {}, "hash": "e174fa9961aa7a5abad967df46c748536fc309fc83a14f9837a8ef01755ce66b", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 3114, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# Multi Adapter RL (MARL) - a single base model for everything

Here we present an approach that uses a single base model for the entire PPO algorithm - which includes retrieving the reference logits, computing the active logits and the rewards. This feature is experimental as we did not test the convergence of the approach. We encourage the community to let us know if they potentially face issues.

## Requirements

You just need to install `peft` and optionally install `bitsandbytes` as well if you want to go for 8bit base models, for more memory efficient finetuning.

## Summary

You need to address this approach in three stages that we summarize as follows:

1- Train a base model on the target domain (e.g. `imdb` dataset) - this is the Supervised Fine Tuning stage - it can leverage the `SFTTrainer` from TRL.
2- Train a reward model using `peft`. This is required in order to re-use the adapter during the RL optimisation process (step 3 below). We show an example of leveraging the `RewardTrainer` from TRL in [this example](https://github.com/huggingface/trl/tree/main/examples/scripts/reward_modeling.py)
3- Fine tune new adapters on the base model using PPO and the reward adapter. ("0 abstraction RL")

Make sure to use the same model (i.e. same architecture and same weights) for the stages 2 & 3. 

## Quickstart

Let us assume you have trained your reward adapter on `llama-7b` model using `RewardTrainer` and pushed the weights on the hub under `trl-lib/llama-7b-hh-rm-adapter`. 
When doing PPO, before passing the model to `PPOTrainer` create your model as follows:

```python
model_name = "huggyllama/llama-7b"
rm_adapter_id = "trl-lib/llama-7b-hh-rm-adapter"

# PPO adapter
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = AutoModelForCausalLMWithValueHead.from_pretrained(
    model_name,
    peft_config=lora_config,
    reward_adapter=rm_adapter_id,
)

...
trainer = PPOTrainer(
    model=model,
    ...
)

...
```
Then inside your PPO training loop, call the `compute_reward_score` method by accessing the `model` attribute from `PPOTrainer`.

```python
rewards = trainer.model.compute_reward_score(**inputs)
```

## Advanced usage

### Control on the adapter name 

If you are familiar with the `peft` library, you know that you can use multiple adapters inside the same model. What you can do is train multiple adapters on the same base model to fine-tune on different policies. 
In this case, you want to be able to control the adapter name you want to activate back, after retrieving the reward. For that, simply pass the appropriate `adapter_name` to `ppo_adapter_name` argument when calling `compute_reward_score`.

```python
adapter_name_policy_1 = "policy_1"
rewards = trainer.model.compute_reward_score(**inputs, ppo_adapter_name=adapter_name_policy_1)
...
```

### Using 4-bit and 8-bit base models

For more memory efficient fine-tuning, you can load your base model in 8-bit or 4-bit while keeping the adapters in the default precision (float32).
Just pass the appropriate arguments (i.e.


key: _node_content
string_value:
{"id_": "41a54794-792b-4662-aff0-184281a8667e", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/multi_adapter_rl", "title": "Multi Adapter RL (MARL) - a single base model for everything", "tokens": 957, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "c14126ac-f163-4113-b1ae-e2225ad52384", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/multi_adapter_rl", "title": "Multi Adapter RL (MARL) - a single base model for everything", "tokens": 957, "retrieve_doc": true, "source": "trl"}, "hash": "a7b627ed9008d875c904a7e4d585c9eddda5901fba91fe59a62136d0c755b57d", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "8a72fcd3-74d0-481c-a6ec-10f6a0875be5", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/multi_adapter_rl", "title": "Multi Adapter RL (MARL) - a single base model for everything", "tokens": 957, "retrieve_doc": true, "source": "trl"}, "hash": "53b3b60fb5d0070c22dc97302b059361929b602c1a9f25aa5e0f0b2dfd1114bf", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 2158, "end_char_idx": 3718, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
```python
rewards = trainer.model.compute_reward_score(**inputs)
```

## Advanced usage

### Control on the adapter name 

If you are familiar with the `peft` library, you know that you can use multiple adapters inside the same model. What you can do is train multiple adapters on the same base model to fine-tune on different policies. 
In this case, you want to be able to control the adapter name you want to activate back, after retrieving the reward. For that, simply pass the appropriate `adapter_name` to `ppo_adapter_name` argument when calling `compute_reward_score`.

```python
adapter_name_policy_1 = "policy_1"
rewards = trainer.model.compute_reward_score(**inputs, ppo_adapter_name=adapter_name_policy_1)
...
```

### Using 4-bit and 8-bit base models

For more memory efficient fine-tuning, you can load your base model in 8-bit or 4-bit while keeping the adapters in the default precision (float32).
Just pass the appropriate arguments (i.e. `load_in_8bit=True` or `load_in_4bit=True`) to `AutoModelForCausalLMWithValueHead.from_pretrained` as follows (assuming you have installed `bitsandbytes`):
```python
model_name = "llama-7b"
rm_adapter_id = "trl-lib/llama-7b-hh-rm-adapter"

# PPO adapter
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = AutoModelForCausalLMWithValueHead.from_pretrained(
    model_name,
    peft_config=lora_config,
    reward_adapter=rm_adapter_id,
    load_in_8bit=True,
)

...
trainer = PPOTrainer(
    model=model,
    ...
)
...
```


key: _node_content
string_value:
{"id_": "60136cee-3ab1-4c34-9ad8-ff6ba0b0e7bf", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/logging", "title": "Logging", "tokens": 1666, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "3d9953f4-ad38-4840-be8b-8a5bdda14b0b", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/logging", "title": "Logging", "tokens": 1666, "retrieve_doc": true, "source": "trl"}, "hash": "f4e7478381532136456c5a9f7e6970d1be028e19abdeb0ed19f4bcc5f9151afe", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "8afd6035-bde9-4c3c-9408-2935ff640855", "node_type": "1", "metadata": {}, "hash": "bb910b76713641321500a655f32779c4e104bba3494be04cc81af807ffe4b633", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 3067, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# Logging

As reinforcement learning algorithms are historically challenging to debug, it's important to pay careful attention to logging.
By default, the TRL [`PPOTrainer`] saves a lot of relevant information to `wandb` or `tensorboard`.

Upon initialization, pass one of these two options to the [`PPOConfig`]:
```
config = PPOConfig(
    model_name=args.model_name,
    log_with=`wandb`, # or `tensorboard`
)
```
If you want to log with tensorboard, add the kwarg `project_kwargs={"logging_dir": PATH_TO_LOGS}` to the PPOConfig.

## PPO Logging

Here's a brief explanation for the logged metrics provided in the data:

Key metrics to monitor. We want to maximize the reward, maintain a low KL divergence, and maximize entropy:
1. `env/reward_mean`: The average reward obtained from the environment. Alias `ppo/mean_scores`, which is sed to specifically monitor the reward model.
1. `env/reward_std`: The standard deviation of the reward obtained from the environment. Alias ``ppo/std_scores`, which is sed to specifically monitor the reward model.
1. `env/reward_dist`: The histogram distribution of the reward obtained from the environment.
1. `objective/kl`: The mean Kullback-Leibler (KL) divergence between the old and new policies. It measures how much the new policy deviates from the old policy. The KL divergence is used to compute the KL penalty in the objective function.
1. `objective/kl_dist`: The histogram distribution of the `objective/kl`.
1. `objective/kl_coef`: The coefficient for Kullback-Leibler (KL) divergence in the objective function. 
1. `ppo/mean_non_score_reward`: The **KL penalty** calculated by `objective/kl * objective/kl_coef` as the total reward for optimization to prevent the new policy from deviating too far from the old policy.
1. `objective/entropy`: The entropy of the model's policy, calculated by `-logprobs.sum(-1).mean()`. High entropy means the model's actions are more random, which can be beneficial for exploration.

Training stats:
1. `ppo/learning_rate`: The learning rate for the PPO algorithm.
1. `ppo/policy/entropy`: The entropy of the model's policy, calculated by `pd = torch.nn.functional.softmax(logits, dim=-1); entropy = torch.logsumexp(logits, dim=-1) - torch.sum(pd * logits, dim=-1)`. It measures the randomness of the policy.
1. `ppo/policy/clipfrac`: The fraction of probability ratios (old policy / new policy) that fell outside the clipping range in the PPO objective. This can be used to monitor the optimization process.
1. `ppo/policy/approxkl`: The approximate KL divergence between the old and new policies, measured by `0.5 * masked_mean((logprobs - old_logprobs) ** 2, mask)`, corresponding to the `k2` estimator in http://joschu.net/blog/kl-approx.html
1. `ppo/policy/policykl`: Similar to `ppo/policy/approxkl`, but measured by `masked_mean(old_logprobs - logprobs, mask)`, corresponding to the `k1` estimator in http://joschu.net/blog/kl-approx.html
1. `ppo/policy/ratio`:  The histogram distribution of the ratio between the new and old policies, used to compute the PPO objective.
1.


key: _node_content
string_value:
{"id_": "8afd6035-bde9-4c3c-9408-2935ff640855", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/logging", "title": "Logging", "tokens": 1666, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "3d9953f4-ad38-4840-be8b-8a5bdda14b0b", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/logging", "title": "Logging", "tokens": 1666, "retrieve_doc": true, "source": "trl"}, "hash": "f4e7478381532136456c5a9f7e6970d1be028e19abdeb0ed19f4bcc5f9151afe", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "60136cee-3ab1-4c34-9ad8-ff6ba0b0e7bf", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/logging", "title": "Logging", "tokens": 1666, "retrieve_doc": true, "source": "trl"}, "hash": "617030113c854157218d0577d7d1874153306da8d2d1a8127fb56189005aef3a", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "1e89b0b3-755d-4285-bac4-9d8cd67a91bd", "node_type": "1", "metadata": {}, "hash": "4b02e3a448a8026629e1a03afbd5d608ef6b86919d33528cef4277c855d32fdb", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 1564, "end_char_idx": 4609, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
1. `ppo/mean_non_score_reward`: The **KL penalty** calculated by `objective/kl * objective/kl_coef` as the total reward for optimization to prevent the new policy from deviating too far from the old policy.
1. `objective/entropy`: The entropy of the model's policy, calculated by `-logprobs.sum(-1).mean()`. High entropy means the model's actions are more random, which can be beneficial for exploration.

Training stats:
1. `ppo/learning_rate`: The learning rate for the PPO algorithm.
1. `ppo/policy/entropy`: The entropy of the model's policy, calculated by `pd = torch.nn.functional.softmax(logits, dim=-1); entropy = torch.logsumexp(logits, dim=-1) - torch.sum(pd * logits, dim=-1)`. It measures the randomness of the policy.
1. `ppo/policy/clipfrac`: The fraction of probability ratios (old policy / new policy) that fell outside the clipping range in the PPO objective. This can be used to monitor the optimization process.
1. `ppo/policy/approxkl`: The approximate KL divergence between the old and new policies, measured by `0.5 * masked_mean((logprobs - old_logprobs) ** 2, mask)`, corresponding to the `k2` estimator in http://joschu.net/blog/kl-approx.html
1. `ppo/policy/policykl`: Similar to `ppo/policy/approxkl`, but measured by `masked_mean(old_logprobs - logprobs, mask)`, corresponding to the `k1` estimator in http://joschu.net/blog/kl-approx.html
1. `ppo/policy/ratio`:  The histogram distribution of the ratio between the new and old policies, used to compute the PPO objective.
1. `ppo/policy/advantages_mean`: The average of the GAE (Generalized Advantage Estimation) advantage estimates. The advantage function measures how much better an action is compared to the average action at a state.
1. `ppo/policy/advantages`: The histogram distribution of `ppo/policy/advantages_mean`.
1. `ppo/returns/mean`: The mean of the TD(λ) returns, calculated by `returns = advantage + values`, another indicator of model performance. See https://iclr-blog-track.github.io/2022/03/25/ppo-implementation-details/ for more details.
1. `ppo/returns/var`: The variance of the TD(λ) returns, calculated by `returns = advantage + values`, another indicator of model performance.
1. `ppo/val/mean`: The mean of the values, used to monitor the value function's performance.
1. `ppo/val/var` : The variance of the values, used to monitor the value function's performance.
1. `ppo/val/var_explained`: The explained variance for the value function, used to monitor the value function's performance.
1. `ppo/val/clipfrac`: The fraction of the value function's predicted values that are clipped.
1. `ppo/val/vpred`: The predicted values from the value function.
1. `ppo/val/error`: The mean squared error between the `ppo/val/vpred` and returns, used to monitor the value function's performance.
1. `ppo/loss/policy`: The policy loss for the Proximal Policy Optimization (PPO) algorithm.
1. `ppo/loss/value`: The loss for the value function in the PPO algorithm. This value quantifies how well the function estimates the expected future rewards.
1.


key: _node_content
string_value:
{"id_": "1e89b0b3-755d-4285-bac4-9d8cd67a91bd", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/logging", "title": "Logging", "tokens": 1666, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "3d9953f4-ad38-4840-be8b-8a5bdda14b0b", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/logging", "title": "Logging", "tokens": 1666, "retrieve_doc": true, "source": "trl"}, "hash": "f4e7478381532136456c5a9f7e6970d1be028e19abdeb0ed19f4bcc5f9151afe", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "8afd6035-bde9-4c3c-9408-2935ff640855", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/logging", "title": "Logging", "tokens": 1666, "retrieve_doc": true, "source": "trl"}, "hash": "fafa2f337af8a0c29f3ba359d1e82725e95647197b9abae775c4c18a61df7355", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "7c75c4e9-df2f-411d-98f4-ec74b1b3fca8", "node_type": "1", "metadata": {}, "hash": "5eb96a8d9c216c92988f1cfae39ac3921cb62be8af0fd48fe4ce9b45cabb8030", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 3065, "end_char_idx": 5592, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
1. `ppo/policy/advantages_mean`: The average of the GAE (Generalized Advantage Estimation) advantage estimates. The advantage function measures how much better an action is compared to the average action at a state.
1. `ppo/policy/advantages`: The histogram distribution of `ppo/policy/advantages_mean`.
1. `ppo/returns/mean`: The mean of the TD(λ) returns, calculated by `returns = advantage + values`, another indicator of model performance. See https://iclr-blog-track.github.io/2022/03/25/ppo-implementation-details/ for more details.
1. `ppo/returns/var`: The variance of the TD(λ) returns, calculated by `returns = advantage + values`, another indicator of model performance.
1. `ppo/val/mean`: The mean of the values, used to monitor the value function's performance.
1. `ppo/val/var` : The variance of the values, used to monitor the value function's performance.
1. `ppo/val/var_explained`: The explained variance for the value function, used to monitor the value function's performance.
1. `ppo/val/clipfrac`: The fraction of the value function's predicted values that are clipped.
1. `ppo/val/vpred`: The predicted values from the value function.
1. `ppo/val/error`: The mean squared error between the `ppo/val/vpred` and returns, used to monitor the value function's performance.
1. `ppo/loss/policy`: The policy loss for the Proximal Policy Optimization (PPO) algorithm.
1. `ppo/loss/value`: The loss for the value function in the PPO algorithm. This value quantifies how well the function estimates the expected future rewards.
1. `ppo/loss/total`: The total loss for the PPO algorithm. It is the sum of the policy loss and the value function loss.


Stats on queries, responses, and logprobs:
1. `tokens/queries_len_mean`: The average length of the queries tokens.
1. `tokens/queries_len_std`: The standard deviation of the length of the queries tokens.
1. `tokens/queries_dist`: The histogram distribution of the length of the queries tokens.
1. `tokens/responses_len_mean`: The average length of the responses tokens.
1. `tokens/responses_len_std`: The standard deviation of the length of the responses tokens.
1. `tokens/responses_dist`: The histogram distribution of the length of the responses tokens. (Costa: inconsistent naming, should be `tokens/responses_len_dist`)
1. `objective/logprobs`: The histogram distribution of the log probabilities of the actions taken by the model.
1. `objective/ref_logprobs`: The histogram distribution of the log probabilities of the actions taken by the reference model.


key: _node_content
string_value:
{"id_": "7c75c4e9-df2f-411d-98f4-ec74b1b3fca8", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/logging", "title": "Logging", "tokens": 1666, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "3d9953f4-ad38-4840-be8b-8a5bdda14b0b", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/logging", "title": "Logging", "tokens": 1666, "retrieve_doc": true, "source": "trl"}, "hash": "f4e7478381532136456c5a9f7e6970d1be028e19abdeb0ed19f4bcc5f9151afe", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "1e89b0b3-755d-4285-bac4-9d8cd67a91bd", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/logging", "title": "Logging", "tokens": 1666, "retrieve_doc": true, "source": "trl"}, "hash": "525b1536b1b6bead41e9ac131c57aae653fda4592a52042cb8878ceebfed0c95", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 3937, "end_char_idx": 7033, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
1. `ppo/val/var_explained`: The explained variance for the value function, used to monitor the value function's performance.
1. `ppo/val/clipfrac`: The fraction of the value function's predicted values that are clipped.
1. `ppo/val/vpred`: The predicted values from the value function.
1. `ppo/val/error`: The mean squared error between the `ppo/val/vpred` and returns, used to monitor the value function's performance.
1. `ppo/loss/policy`: The policy loss for the Proximal Policy Optimization (PPO) algorithm.
1. `ppo/loss/value`: The loss for the value function in the PPO algorithm. This value quantifies how well the function estimates the expected future rewards.
1. `ppo/loss/total`: The total loss for the PPO algorithm. It is the sum of the policy loss and the value function loss.


Stats on queries, responses, and logprobs:
1. `tokens/queries_len_mean`: The average length of the queries tokens.
1. `tokens/queries_len_std`: The standard deviation of the length of the queries tokens.
1. `tokens/queries_dist`: The histogram distribution of the length of the queries tokens.
1. `tokens/responses_len_mean`: The average length of the responses tokens.
1. `tokens/responses_len_std`: The standard deviation of the length of the responses tokens.
1. `tokens/responses_dist`: The histogram distribution of the length of the responses tokens. (Costa: inconsistent naming, should be `tokens/responses_len_dist`)
1. `objective/logprobs`: The histogram distribution of the log probabilities of the actions taken by the model.
1. `objective/ref_logprobs`: The histogram distribution of the log probabilities of the actions taken by the reference model.



### Crucial values
During training, many values are logged, here are the most important ones:

1. `env/reward_mean`,`env/reward_std`, `env/reward_dist`: the properties of the reward distribution from the "environment" /  reward model
1. `ppo/mean_non_score_reward`: The mean negated KL penalty during training (shows the delta between the reference model and the new policy over the batch in the step)

Here are some parameters that are useful to monitor for stability (when these diverge or collapse to 0, try tuning variables):

1. `ppo/loss/value`: it will spike / NaN when not going well.
1. `ppo/policy/ratio`: `ratio` being 1 is a baseline value, meaning that the probability of sampling a token is the same under the new and old policy. If the ratio is too high like 200, it means the probability of sampling a token is 200 times higher under the new policy than the old policy. This is a sign that the new policy is too different from the old policy, which will likely cause overoptimization and collapse training later on.
1. `ppo/policy/clipfrac` and `ppo/policy/approxkl`: if `ratio` is too high, the `ratio` is going to get clipped, resulting in high `clipfrac` and high `approxkl` as well.
1. `objective/kl`: it should stay positive so that the policy is not too far away from the reference policy.
1. `objective/kl_coef`: The target coefficient with [`AdaptiveKLController`]. Often increases before numerical instabilities.


key: _node_content
string_value:
{"id_": "5fc4a359-59cc-4390-881e-b64e61b6aeb5", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/kto_trainer", "title": "KTO Trainer", "tokens": 970, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "1313028e-be5c-4d19-bf43-f36d54e284ff", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/kto_trainer", "title": "KTO Trainer", "tokens": 970, "retrieve_doc": true, "source": "trl"}, "hash": "dddcad02280f1d2bbcbdae915ebf04604be7cb5b6cd81051138f14b7e05bf056", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "be8250ec-2445-4b73-b7cf-f23022a64935", "node_type": "1", "metadata": {}, "hash": "738ce6f70f50a9c7fdcb6b6166503cf44b84e4e38a57706801607893bfbfcc78", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 2019, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# KTO Trainer

TRL supports the Kahneman-Tversky Optimization (KTO) Trainer for aligning language models with binary feedback data (e.g., upvote/downvote), as described in the [paper](https://huggingface.co/papers/2402.01306) by Kawin Ethayarajh, Winnie Xu, Niklas Muennighoff, Dan Jurafsky, and Douwe Kiela.
For a full example have a look at  [`examples/scripts/kto.py`].

Depending on how good your base model is, you may or may not need to do SFT before KTO.
This is different from standard RLHF and DPO, which always require SFT.

## Expected dataset format

The KTO trainer expects a very specific format for the dataset as it does not require pairwise preferences. Since the model will be trained to directly optimize examples that consist of a prompt, model completion, and a label to indicate whether the completion is "good" or "bad", we expect a dataset with the following columns:

- `prompt`
- `completion`
- `label`

for example:

```
kto_dataset_dict = {
    "prompt": [
        "Hey, hello",
        "How are you",
        "What is your name?",
        "What is your name?",
        "Which is the best programming language?",
        "Which is the best programming language?",
        "Which is the best programming language?",
    ],
    "completion": [
        "hi nice to meet you",
        "leave me alone",
        "I don't have a name",
        "My name is Mary",
        "Python",
        "C++",
        "Java",
    ],
    "label": [
        True,
        False,
        False,
        True,
        True,
        False,
        False,
    ],
}
```

where the `prompt` contains the context inputs, `completion` contains the corresponding responses and `label` contains the corresponding flag that indicates if the generated completion is desired (`True`) or undesired (`False`).
A prompt can have multiple responses and this is reflected in the entries being repeated in the dictionary's value arrays. It is required that the dataset contains at least one desirable and one undesirable completion.


key: _node_content
string_value:
{"id_": "be8250ec-2445-4b73-b7cf-f23022a64935", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/kto_trainer", "title": "KTO Trainer", "tokens": 970, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "1313028e-be5c-4d19-bf43-f36d54e284ff", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/kto_trainer", "title": "KTO Trainer", "tokens": 970, "retrieve_doc": true, "source": "trl"}, "hash": "dddcad02280f1d2bbcbdae915ebf04604be7cb5b6cd81051138f14b7e05bf056", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "5fc4a359-59cc-4390-881e-b64e61b6aeb5", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/kto_trainer", "title": "KTO Trainer", "tokens": 970, "retrieve_doc": true, "source": "trl"}, "hash": "e433693c29be5202ae9694b76c661d85c279b95a235ee485ac19de7c6e80c88d", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 2022, "end_char_idx": 4078, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
## Expected model format
The KTO trainer expects a model of `AutoModelForCausalLM`, compared to PPO that expects `AutoModelForCausalLMWithValueHead` for the value function.

## Using the `KTOTrainer`

For a detailed example have a look at the `examples/scripts/kto.py` script. At a high level we need to initialize the `KTOTrainer` with a `model` we wish to train and a reference `ref_model` which we will use to calculate the implicit rewards of the preferred and rejected response. 

The `beta` refers to the hyperparameter of the implicit reward, and the dataset contains the 3 entries listed above. Note that the `model` and `ref_model` need to have the same architecture (ie decoder only or encoder-decoder).

The `desirable_weight` and `undesirable_weight` refer to the weights placed on the losses for desirable/positive and undesirable/negative examples.
By default, they are both 1. However, if you have more of one or the other, then you should upweight the less common type such that the ratio of (`desirable_weight` * number of positives) to (`undesirable_weight` * number of negatives) is in the range 1:1 to 4:3.

```py
training_args = KTOConfig(
    beta=0.1,
    desirable_weight=1.0,
    undesirable_weight=1.0,
)

kto_trainer = KTOTrainer(
    model,
    ref_model,
    args=training_args,
    train_dataset=train_dataset,
    tokenizer=tokenizer,
)
```
After this one can then call:

```py
kto_trainer.train()
```

### For Mixture of Experts Models: Enabling the auxiliary loss

MOEs are the most efficient if the load is about equally distributed between experts.  
To ensure that we train MOEs similarly during preference-tuning, it is beneficial to add the auxiliary loss from the load balancer to the final loss.  

This option is enabled by setting `output_router_logits=True` in the model config (e.g. MixtralConfig).  
To scale how much the auxiliary loss contributes to the total loss, use the hyperparameter `router_aux_loss_coef=...` (default: 0.001).

## KTOTrainer

[[autodoc]] KTOTrainer

## KTOConfig

[[autodoc]] KTOConfig


key: _node_content
string_value:
{"id_": "2a660d4d-9dd8-4621-bc3b-25531e227871", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/alignprop_trainer", "title": "Aligning Text-to-Image Diffusion Models with Reward Backpropagation", "tokens": 1189, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "1d57c232-95bd-4519-aeba-c26fbc873ba2", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/alignprop_trainer", "title": "Aligning Text-to-Image Diffusion Models with Reward Backpropagation", "tokens": 1189, "retrieve_doc": true, "source": "trl"}, "hash": "4e06c25df8950d7ab028a71bdf6261c3d9fead6dadb00dc0ddcd78c20671c06b", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "3ee149cb-4e80-4012-ae2e-be66070817df", "node_type": "1", "metadata": {}, "hash": "d6b15aeb98b4d6b00e167bcc35ce931cc3c994969d59e4dedf890789a67dfabc", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 3430, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# Aligning Text-to-Image Diffusion Models with Reward Backpropagation

## The why

If your reward function is differentiable, directly backpropagating gradients from the reward models to the diffusion model is significantly more sample and compute efficient (25x) than doing policy gradient algorithm like DDPO.
AlignProp does full backpropagation through time, which allows updating the earlier steps of denoising via reward backpropagation.

<div style="text-align: center"><img src="https://align-prop.github.io/reward_tuning.png"/></div>


## Getting started with `examples/scripts/alignprop.py`

The `alignprop.py` script is a working example of using the `AlignProp` trainer to finetune a Stable Diffusion model. This example explicitly configures a small subset of the overall parameters associated with the config object (`AlignPropConfig`).

**Note:** one A100 GPU is recommended to get this running. For lower memory setting, consider setting truncated_backprop_rand to False. With default settings this will do truncated backpropagation with K=1.

Almost every configuration parameter has a default. There is only one commandline flag argument that is required of the user to get things up and running. The user is expected to have a [huggingface user access token](https://huggingface.co/docs/hub/security-tokens) that will be used to upload the model post finetuning to HuggingFace hub. The following bash command is to be entered to get things running

```batch
python alignprop.py --hf_user_access_token <token>
```

To obtain the documentation of `stable_diffusion_tuning.py`, please run `python stable_diffusion_tuning.py --help`

The following are things to keep in mind (The code checks this for you as well) in general while configuring the trainer (beyond the use case of using the example script)

- The configurable randomized truncation range (`--alignprop_config.truncated_rand_backprop_minmax=(0,50)`) the first number should be equal and greater to 0, while the second number should equal or less to the number of diffusion timesteps (sample_num_steps)
- The configurable truncation backprop absolute step (`--alignprop_config.truncated_backprop_timestep=49`) the number should be less than the number of diffusion timesteps (sample_num_steps), it only matters when truncated_backprop_rand is set to False

## Setting up the image logging hook function

Expect the function to be given a dictionary with keys
```python
['image', 'prompt', 'prompt_metadata', 'rewards']

```
and `image`, `prompt`, `prompt_metadata`, `rewards`are batched.
You are free to log however you want the use of `wandb` or `tensorboard` is recommended.

### Key terms

- `rewards` : The rewards/score is a numerical associated with the generated image and is key to steering the RL process
- `prompt` : The prompt is the text that is used to generate the image
- `prompt_metadata` : The prompt metadata is the metadata associated with the prompt. A situation where this will not be empty is when the reward model comprises of a [`FLAVA`](https://huggingface.co/docs/transformers/model_doc/flava) setup where questions and ground answers (linked to the generated image) are expected with the generated image (See here: https://github.com/kvablack/ddpo-pytorch/blob/main/ddpo_pytorch/rewards.py#L45)
- `image` : The image generated by the Stable Diffusion model

Example code for logging sampled images with `wandb` is given below.


key: _node_content
string_value:
{"id_": "3ee149cb-4e80-4012-ae2e-be66070817df", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/alignprop_trainer", "title": "Aligning Text-to-Image Diffusion Models with Reward Backpropagation", "tokens": 1189, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "1d57c232-95bd-4519-aeba-c26fbc873ba2", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/alignprop_trainer", "title": "Aligning Text-to-Image Diffusion Models with Reward Backpropagation", "tokens": 1189, "retrieve_doc": true, "source": "trl"}, "hash": "4e06c25df8950d7ab028a71bdf6261c3d9fead6dadb00dc0ddcd78c20671c06b", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "2a660d4d-9dd8-4621-bc3b-25531e227871", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/alignprop_trainer", "title": "Aligning Text-to-Image Diffusion Models with Reward Backpropagation", "tokens": 1189, "retrieve_doc": true, "source": "trl"}, "hash": "be5dd7fd262c8c734dea2aafb621d757569f6718b857970e9e919f143c9b681f", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 2565, "end_char_idx": 5021, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
You are free to log however you want the use of `wandb` or `tensorboard` is recommended.

### Key terms

- `rewards` : The rewards/score is a numerical associated with the generated image and is key to steering the RL process
- `prompt` : The prompt is the text that is used to generate the image
- `prompt_metadata` : The prompt metadata is the metadata associated with the prompt. A situation where this will not be empty is when the reward model comprises of a [`FLAVA`](https://huggingface.co/docs/transformers/model_doc/flava) setup where questions and ground answers (linked to the generated image) are expected with the generated image (See here: https://github.com/kvablack/ddpo-pytorch/blob/main/ddpo_pytorch/rewards.py#L45)
- `image` : The image generated by the Stable Diffusion model

Example code for logging sampled images with `wandb` is given below.

```python
# for logging these images to wandb

def image_outputs_hook(image_data, global_step, accelerate_logger):
    # For the sake of this example, we only care about the last batch
    # hence we extract the last element of the list
    result = {}
    images, prompts, rewards = [image_data['images'],image_data['prompts'],image_data['rewards']]
    for i, image in enumerate(images):
        pil = Image.fromarray(
            (image.cpu().numpy().transpose(1, 2, 0) * 255).astype(np.uint8)
        )
        pil = pil.resize((256, 256))
        result[f"{prompts[i]:.25} | {rewards[i]:.2f}"] = [pil]
    accelerate_logger.log_images(
        result,
        step=global_step,
    )

```

### Using the finetuned model

Assuming you've done with all the epochs and have pushed up your model to the hub, you can use the finetuned model as follows

```python
from diffusers import StableDiffusionPipeline
pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
pipeline.to("cuda")

pipeline.load_lora_weights('mihirpd/alignprop-trl-aesthetics')

prompts = ["squirrel", "crab", "starfish", "whale","sponge", "plankton"]
results = pipeline(prompts)

for prompt, image in zip(prompts,results.images):
    image.save(f"dump/{prompt}.png")
```

## Credits

This work is heavily influenced by the repo [here](https://github.com/mihirp1998/AlignProp/) and the associated paper [Aligning Text-to-Image Diffusion Models with Reward Backpropagation
 by Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, Katerina Fragkiadaki](https://huggingface.co/papers/2310.03739).


key: _node_content
string_value:
{"id_": "af821a28-8ff8-4b3d-9226-15aab889c582", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sentiment_tuning", "title": "Sentiment Tuning Examples", "tokens": 1601, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "068713a8-03d9-4ee0-9160-1eddea0b396e", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sentiment_tuning", "title": "Sentiment Tuning Examples", "tokens": 1601, "retrieve_doc": true, "source": "trl"}, "hash": "79b3ef31caa859bcf817dbfbf8ebf948e5aed5c4469fc35486d3971cb19fea23", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "179b4a0c-d18f-4dc2-bd4e-e0757808623b", "node_type": "1", "metadata": {}, "hash": "f003098101663a6c7944f3f673c12938cbd6392231914a88f8ee6d93e63dacab", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 0, "end_char_idx": 3065, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


key: chroma:document
string_value:
# Sentiment Tuning Examples

The notebooks and scripts in this examples show how to fine-tune a model with a sentiment classifier (such as `lvwerra/distilbert-imdb`).

Here's an overview of the notebooks and scripts in the [trl repository](https://github.com/huggingface/trl/tree/main/examples):



| File                                                                                           | Description                                                                                                              |
|------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------|
| [`examples/scripts/ppo.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/ppo.py)  [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/trl/blob/main/examples/sentiment/notebooks/gpt2-sentiment.ipynb) | This script shows how to use the `PPOTrainer` to fine-tune a sentiment analysis model using IMDB dataset                 |
| [`examples/notebooks/gpt2-sentiment.ipynb`](https://github.com/huggingface/trl/tree/main/examples/notebooks/gpt2-sentiment.ipynb)              | This notebook demonstrates how to reproduce the GPT2 imdb sentiment tuning example on a jupyter notebook.                |
| [`examples/notebooks/gpt2-control.ipynb`](https://github.com/huggingface/trl/tree/main/examples/notebooks/gpt2-control.ipynb)   [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/trl/blob/main/examples/sentiment/notebooks/gpt2-sentiment-control.ipynb)                | This notebook demonstrates how to reproduce the GPT2 sentiment control example on a jupyter notebook.    



## Usage

```bash
# 1. run directly
python examples/scripts/ppo.py
# 2. run via `accelerate` (recommended), enabling more features (e.g., multiple GPUs, deepspeed)
accelerate config # will prompt you to define the training configuration
accelerate launch examples/scripts/ppo.py # launches training
# 3. get help text and documentation
python examples/scripts/ppo.py --help
# 4. configure logging with wandb and, say, mini_batch_size=1 and gradient_accumulation_steps=16
python examples/scripts/ppo.py --log_with wandb --mini_batch_size 1 --gradient_accumulation_steps 16
```

Note: if you don't want to log with `wandb` remove `log_with="wandb"` in the scripts/notebooks. You can also replace it with your favourite experiment tracker that's [supported by `accelerate`](https://huggingface.co/docs/accelerate/usage_guides/tracking).


## Few notes on multi-GPU 

To run in multi-GPU setup with DDP (distributed Data Parallel) change the `device_map` value to `device_map={"": Accelerator().process_index}` and make sure to run your script with `accelerate launch yourscript.py`. If you want to apply naive pipeline parallelism you can use `device_map="auto"`.


key: _node_content
string_value:
{"id_": "179b4a0c-d18f-4dc2-bd4e-e0757808623b", "embedding": null, "metadata": {"url": "https://huggingface.co/docs/trl/sentiment_tuning", "title": "Sentiment Tuning Examples", "tokens": 1601, "retrieve_doc": true, "source": "trl"}, "excluded_embed_metadata_keys": ["url", "tokens", "retrieve_doc", "source"], "excluded_llm_metadata_keys": ["title", "tokens", "retrieve_doc", "source"], "relationships": {"1": {"node_id": "068713a8-03d9-4ee0-9160-1eddea0b396e", "node_type": "4", "metadata": {"url": "https://huggingface.co/docs/trl/sentiment_tuning", "title": "Sentiment Tuning Examples", "tokens": 1601, "retrieve_doc": true, "source": "trl"}, "hash": "79b3ef31caa859bcf817dbfbf8ebf948e5aed5c4469fc35486d3971cb19fea23", "class_name": "RelatedNodeInfo"}, "2": {"node_id": "af821a28-8ff8-4b3d-9226-15aab889c582", "node_type": "1", "metadata": {"url": "https://huggingface.co/docs/trl/sentiment_tuning", "title": "Sentiment Tuning Examples", "tokens": 1601, "retrieve_doc": true, "source": "trl"}, "hash": "3e370eb195817283eb7f1a6f4e2261ffcdea933411475325b7c98763cb7c5371", "class_name": "RelatedNodeInfo"}, "3": {"node_id": "60852c40-ee68-4ffc-8fde-6d3ca4432fc1", "node_type": "1", "metadata": {}, "hash": "9cbe88825ff1a880f33221eeafb905e727666dd9a02e23ce4a8d0736580e0596", "class_name": "RelatedNodeInfo"}}, "text": "", "mimetype": "text/plain", "start_char_idx": 1905, "end_char_idx": 4296, "text_template": "{metadata_str}\n\n{content}", "metadata_template": "{key}: {value}", "metadata_seperator": "\n", "class_name": "TextNode"}


