PPO Trainer
Source: trl

# PPO TrainerTRL 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 worksFine-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 formatThe `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:```pyfrom datasets import load_datasetdataset = 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:```pyppo_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.```pyfrom trl import PPOConfigconfig = 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:```pyfrom transformers import AutoTokenizerfrom trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainermodel = 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.```pyfrom transformers import pipelinereward_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:```pydef tokenize(sample):    sample["input_ids"] = tokenizer.encode(sample["query"])    return sampledataset = dataset.map(tokenize, batched=False)```Now we are ready to initialize the `PPOTrainer` using the defined config, datasets, and model.```pyfrom trl import PPOTrainerppo_trainer = PPOTrainer(    model=model,    config=config,    dataset=dataset,    tokenizer=tokenizer,)```### Starting the training loopBecause 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).```pygeneration_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.```pyfrom tqdm import tqdmepochs = 10for 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 modelppo_trainer.save_pretrained("my_ppo_model")```## LoggingWhile 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