# PPO Trainer

## PPO Trainer

TRL supports the [PPO](https://arxiv.org/abs/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](https://huggingface.co/docs/trl/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](https://huggingface.co/docs/trl/reward_trainer)) which will be used to optimize the SFT model using the PPO algorithm.

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

Copied

```
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:

Copied

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

Copied

```
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:

Copied

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

Copied

```
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:

Copied

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

Copied

```
from trl import PPOTrainer

ppo_trainer = PPOTrainer(
    model=model,
    config=config,
    train_dataset=train_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](https://huggingface.co/docs/trl/how_to_train#how-to-generate-text-for-training).

Copied

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

Copied

```
from tqdm import tqdm

for epoch, batch in tqdm(enumerate(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_model("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

#### class trl.PPOTrainer

[\<source>](https://github.com/huggingface/trl/blob/v0.7.2/trl/trainer/ppo_trainer.py#L107)

( config: PPOConfig = Nonemodel: PreTrainedModelWrapper = Noneref\_model: typing.Optional\[trl.models.modeling\_base.PreTrainedModelWrapper] = Nonetokenizer: PreTrainedTokenizerBase = Nonedataset: typing.Union\[torch.utils.data.dataset.Dataset, datasets.arrow\_dataset.Dataset, NoneType] = Noneoptimizer: typing.Optional\[torch.optim.optimizer.Optimizer] = Nonedata\_collator: typing.Optional\[typing.Callable] = Nonenum\_shared\_layers: typing.Optional\[int] = Nonelr\_scheduler: typing.Optional\[torch.optim.lr\_scheduler.\_LRScheduler] = None )

Parameters

* \***\*config\*\*** (`PPOConfig`) — Configuration object for PPOTrainer. Check the documentation of `PPOConfig` for more — details.
* \***\*model\*\*** (`PreTrainedModelWrapper`) — Model to be optimized, BOINC AI transformer model with a value head. — Check the documentation of `PreTrainedModelWrapper` for more details.
* \***\*ref\_model\*\*** (`PreTrainedModelWrapper`, *optional*) — Reference model to be used for KL penalty, BOINC AI — transformer model with a casual language modelling head. Check the documentation of `PreTrainedModelWrapper` for more details. If no reference model is provided, the trainer will create a reference model with the same architecture as the model to be optimized with shared layers.
* \***\*tokenizer\*\*** (`PreTrainedTokenizerBase`) — Tokenizer to be used for encoding the — data. Check the documentation of `transformers.PreTrainedTokenizer` and `transformers.PreTrainedTokenizerFast` for more details.
* \***\*dataset\*\*** (Union\[`torch.utils.data.Dataset`, `datasets.Dataset`], *optional*) — PyTorch dataset or BOINC AI dataset. This is used to create a PyTorch dataloader. If no dataset is provided, the dataloader must be created outside the trainer users needs to design their own dataloader and make sure the batch size that is used is the same as the one specified in the configuration object.
* \***\*optimizer\*\*** (`torch.optim.Optimizer`, *optional*) — Optimizer to be used for training. If no optimizer is — provided, the trainer will create an Adam optimizer with the learning rate specified in the configuration object.
* \***\*data\_collator\*\*** (DataCollatorForLanguageModeling, *optional*) — Data collator to be used for training and — passed along the dataloader
* \***\*num\_shared\_layers\*\*** (int, *optional*) — Number of layers to be shared between the model and the reference — model, if no reference model is passed. If no number is provided, all the layers will be shared.
* \***\*lr\_scheduler\*\*** (`torch.optim.lr_scheduler`, *optional*) — Learning rate scheduler to be used for training. —

The PPOTrainer uses Proximal Policy Optimization to optimise language models. Note, this trainer is heavily inspired by the original OpenAI learning to summarize work here: <https://github.com/openai/summarize-from-feedback>

**batched\_forward\_pass**

[\<source>](https://github.com/huggingface/trl/blob/v0.7.2/trl/trainer/ppo_trainer.py#L912)

( model: PreTrainedModelWrapperqueries: Tensorresponses: Tensormodel\_inputs: dictreturn\_logits: bool = Falseresponse\_masks: typing.Optional\[torch.Tensor] = None ) → (tuple)

Parameters

* **queries** (`torch.LongTensor`) — List of tensors containing the encoded queries, shape (`batch_size`, `query_length`)
* **responses** (`torch.LongTensor`) — List of tensors containing the encoded responses, shape (`batch_size`, `response_length`)
* **return\_logits** (`bool`, *optional*, defaults to `False`) — Whether to return all\_logits. Set to `False` if logits are not needed to reduce memory consumption.

Returns

(tuple)

* all\_logprobs (`torch.FloatTensor`): Log probabilities of the responses, shape (`batch_size`, `response_length`)
* all\_ref\_logprobs (`torch.FloatTensor`): Log probabilities of the responses, shape (`batch_size`, `response_length`)
* all\_values (`torch.FloatTensor`): Values of the responses, shape (`batch_size`, `response_length`)

Calculate model outputs in multiple batches.

**compute\_rewards**

[\<source>](https://github.com/huggingface/trl/blob/v0.7.2/trl/trainer/ppo_trainer.py#L1049)

( scores: FloatTensorlogprobs: FloatTensorref\_logprobs: FloatTensormasks: LongTensor )

Parameters

* **scores** (`torch.FloatTensor`) — Scores from the reward model, shape (`batch_size`)
* **logprobs** (`torch.FloatTensor`) — Log probabilities of the model, shape (`batch_size`, `response_length`)
* **ref\_logprobs** (`torch.FloatTensor`) — Log probabilities of the reference model, shape (`batch_size`, `response_length`)

Compute per token rewards from scores and KL-penalty.

**create\_model\_card**

[\<source>](https://github.com/huggingface/trl/blob/v0.7.2/trl/trainer/ppo_trainer.py#L1355)

( path: strmodel\_name: typing.Optional\[str] = 'TRL Model' )

Parameters

* **path** (`str`) — The path to save the model card to.
* **model\_name** (`str`, *optional*) — The name of the model, defaults to `TRL Model`.

Creates and saves a model card for a TRL model.

**gather\_stats**

[\<source>](https://github.com/huggingface/trl/blob/v0.7.2/trl/trainer/ppo_trainer.py#L868)

( stats ) → `dict[str, Any]`

Parameters

* **stats** (dict\[str, Any]) —
* **a** dictionary of stats to be gathered. The stats should contain torch tensors. —

Returns

`dict[str, Any]`

A dictionary of stats with the tensors gathered.

Gather stats from all processes. Useful in the context of distributed training.

**generate**

[\<source>](https://github.com/huggingface/trl/blob/v0.7.2/trl/trainer/ppo_trainer.py#L417)

( query\_tensor: typing.Union\[torch.Tensor, typing.List\[torch.Tensor]]length\_sampler: typing.Callable = Nonebatch\_size: int = 4return\_prompt: bool = True\*\*generation\_kwargs ) → `torch.LongTensor`

Parameters

* **query\_tensor** (`torch.LongTensor`) — A tensor of shape (`seq_len`) containing query tokens or a list of tensors of shape (`seq_len`).
* **generation\_kwargs** (dict\[str, Any]) — Keyword arguments for generation.
* **length\_sampler** (`Callable`, *optional*) — Callable that returns the number of newly generated tokens.
* **batch\_size** (`int`, \*optional) — Batch size used for generation, defaults to `4`.
* **return\_prompt** (`bool`, *optional*) — If set to `False` the prompt is not returned but only the newly generated tokens, defaults to `True`.

Returns

`torch.LongTensor`

A tensor of shape (`batch_size`, `gen_len`) containing response tokens.

Generate response with the model given the query tensor. call the `generate` method of the model.

**log\_stats**

[\<source>](https://github.com/huggingface/trl/blob/v0.7.2/trl/trainer/ppo_trainer.py#L1274)

( stats: dictbatch: dictrewards: typing.List\[torch.FloatTensor]columns\_to\_log: typing.List\[str] = \['query', 'response'] )

Parameters

* **stats** (dict\[str, Any]) — A dictionary of training stats.
* **batch** (dict\[str, Any]) — A dictionary of batch data, this contains the queries and responses.
* **rewards** (`List[torch.FloatTensor]`) — A tensor of rewards.

A function that logs all the training stats. Call it at the end of each epoch.

**loss**

[\<source>](https://github.com/huggingface/trl/blob/v0.7.2/trl/trainer/ppo_trainer.py#L1122)

( old\_logprobs: FloatTensorvalues: FloatTensorlogits: FloatTensorvpreds: FloatTensorlogprobs: FloatTensormask: LongTensoradvantages: FloatTensorreturns: FloatTensor )

Parameters

* **old\_logprobs** (`torch.FloatTensor`) — Log probabilities of the model, shape (`batch_size`, `response_length`)
* **values** (`torch.FloatTensor`) — Values of the value head, shape (`batch_size`, `response_length`)
* **rewards** (`torch.FloatTensor`) — Rewards from the reward model, shape (`batch_size`, `response_length`)
* **logits** (`torch.FloatTensor`) — Logits of the model, shape (`batch_size`, `response_length`, `vocab_size`)
* **v\_pred** (`torch.FloatTensor`) — Values of the value head, shape (`batch_size`, `response_length`)
* **logprobs** (`torch.FloatTensor`) — Log probabilities of the model, shape (`batch_size`, `response_length`)

Calculate policy and value losses.

**prepare\_dataloader**

[\<source>](https://github.com/huggingface/trl/blob/v0.7.2/trl/trainer/ppo_trainer.py#L362)

( dataset: typing.Union\[torch.utils.data.dataset.Dataset, datasets.arrow\_dataset.Dataset]data\_collator = None ) → `torch.utils.data.DataLoader`

Parameters

* **dataset** (Union\[`torch.utils.data.Dataset`, `datasets.Dataset`]) — PyTorch dataset or BOINC AI dataset. If a BOINC AI dataset is passed, the dataset will be preprocessed by removing the columns that are not used by the model.
* **data\_collator** (Optional\[function]) — Data collator function.

Returns

`torch.utils.data.DataLoader`

PyTorch dataloader

Prepare the dataloader for training.

**record\_step\_stats**

[\<source>](https://github.com/huggingface/trl/blob/v0.7.2/trl/trainer/ppo_trainer.py#L1211)

( kl\_coef: float\*\*data ) → stats (`dict`)

Parameters

* **kl\_coef** (`float`) — KL coefficient
* **data** (`dict`) — Dictionary of training step data

Returns

stats (`dict`)

Dictionary of training step statistics

Record training step statistics.

**step**

[\<source>](https://github.com/huggingface/trl/blob/v0.7.2/trl/trainer/ppo_trainer.py#L579)

( queries: typing.List\[torch.LongTensor]responses: typing.List\[torch.LongTensor]scores: typing.List\[torch.FloatTensor]response\_masks: typing.Optional\[typing.List\[torch.LongTensor]] = None ) → `dict[str, Any]`

Parameters

* **queries** (List`torch.LongTensor`) — List of tensors containing the encoded queries of shape (`query_length`)
* **responses** (List`torch.LongTensor`) — List of tensors containing the encoded responses of shape (`response_length`)
* **scores** (List`torch.FloatTensor`) — List of tensors containing the scores.
* **response\_masks** (List`torch.FloatTensor`, *optional*)) — List of tensors containing masks of the response tokens.

Returns

`dict[str, Any]`

A summary of the training statistics

Run a PPO optimisation step given a list of queries, model responses, and rewards.

**train\_minibatch**

[\<source>](https://github.com/huggingface/trl/blob/v0.7.2/trl/trainer/ppo_trainer.py#L1003)

( old\_logprobs: FloatTensorvalues: FloatTensorlogprobs: FloatTensorlogits: FloatTensorvpreds: FloatTensormask: LongTensoradvantages: FloatTensorreturns: FloatTensor ) → train\_stats (dict\[str, `torch.Tensor`])

Parameters

* **logprobs** (`torch.FloatTensor`) — Log probabilities of the model, shape \[batch\_size, response\_length]
* **values** (`torch.FloatTensor`) — Values of the value head, shape \[batch\_size, response\_length]
* **query** (`torch.LongTensor`) — Encoded queries, shape \[batch\_size, query\_length]
* **response** (`torch.LongTensor`) — Encoded responses, shape \[batch\_size, response\_length]
* **model\_input** (`torch.LongTensor`) — Concatenated queries and responses, shape \[batch\_size, query\_length+response\_length]

Returns

train\_stats (dict\[str, `torch.Tensor`])

Dictionary of training statistics

Train one PPO minibatch

#### class trl.PPOConfig

[\<source>](https://github.com/huggingface/trl/blob/v0.7.2/trl/trainer/ppo_config.py#L30)

( exp\_name: str = 'doc-buil'seed: int = 0log\_with: typing.Union\[typing.Literal\['wandb', 'tensorboard'], NoneType] = Nonetask\_name: typing.Optional\[str] = Nonemodel\_name: typing.Optional\[str] = Nonequery\_dataset: typing.Optional\[str] = Nonereward\_model: typing.Optional\[str] = Noneremove\_unused\_columns: bool = Truetracker\_kwargs: dict = \<factory>accelerator\_kwargs: dict = \<factory>project\_kwargs: dict = \<factory>tracker\_project\_name: str = 'trl'push\_to\_hub\_if\_best\_kwargs: dict = \<factory>steps: int = 20000learning\_rate: float = 1e-05adap\_kl\_ctrl: bool = Trueinit\_kl\_coef: typing.Optional\[float] = 0.2kl\_penalty: typing.Literal\['kl', 'abs', 'mse', 'full'] = 'kl'target: typing.Optional\[float] = 6horizon: typing.Optional\[float] = 10000gamma: float = 1lam: float = 0.95cliprange: float = 0.2cliprange\_value: float = 0.2vf\_coef: float = 0.1batch\_size: int = 256forward\_batch\_size: typing.Optional\[int] = Nonemini\_batch\_size: int = 1gradient\_accumulation\_steps: int = 1world\_size: typing\_extensions.Annotated\[int, Suppress] = Noneppo\_epochs: int = 4max\_grad\_norm: typing.Optional\[float] = Noneoptimize\_cuda\_cache: bool = Falseearly\_stopping: bool = Falsetarget\_kl: float = 1compare\_steps: int = 1ratio\_threshold: float = 10.0use\_score\_scaling: bool = Falseuse\_score\_norm: bool = Falsescore\_clip: typing.Optional\[float] = Noneis\_encoder\_decoder: typing.Union\[typing\_extensions.Annotated\[bool, Suppress], NoneType] = Noneis\_peft\_model: typing.Union\[typing\_extensions.Annotated\[bool, Suppress], NoneType] = Nonebackward\_batch\_size: typing\_extensions.Annotated\[int, Suppress] = Noneglobal\_backward\_batch\_size: typing\_extensions.Annotated\[int, Suppress] = Noneglobal\_batch\_size: typing\_extensions.Annotated\[int, Suppress] = None )

Configuration class for PPOTrainer\\
