We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
adap_kl_ctrl boolean is not used 
Hello,

Thanks for this implementation, is it possible that your kl adaptation is always on ?
https://github.com/lvwerra/trl/blob/750f5fd5329bb81c79b00243c4c8923ac14981d5/trl/ppo.py#L92

Best,
Thibaud

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `trl/ppo.py`
Content:
```
1 # AUTOGENERATED! DO NOT EDIT! File to edit: nbs/02-ppo.ipynb (unless otherwise specified).
2 
3 __all__ = ['AdaptiveKLController', 'FixedKLController', 'PPOTrainer']
4 
5 # Cell
6 import numpy as np
7 import torch.nn.functional as F
8 from torch.optim import Adam
9 import torch
10 import collections
11 import time
12 import random
13 
14 from .core import (logprobs_from_logits,
15                          whiten,
16                          clip_by_value,
17                          entropy_from_logits,
18                          flatten_dict,
19                          average_torch_dicts,
20                          stats_to_np,
21                          stack_dicts,
22                          add_suffix)
23 
24 # Cell
25 
26 class AdaptiveKLController:
27     """
28     Adaptive KL controller described in the paper:
29     https://arxiv.org/pdf/1909.08593.pdf
30     """
31     def __init__(self, init_kl_coef, target, horizon):
32         self.value = init_kl_coef
33         self.target = target
34         self.horizon = horizon
35 
36     def update(self, current, n_steps):
37         target = self.target
38         proportional_error = np.clip(current / target - 1, -0.2, 0.2)
39         mult = 1 + proportional_error * n_steps / self.horizon
40         self.value *= mult
41 
42 # Cell
43 
44 class FixedKLController:
45     """Fixed KL controller."""
46     def __init__(self, kl_coef):
47         self.value = kl_coef
48 
49     def update(self, current, n_steps):
50         pass
51 
52 # Cell
53 
54 class PPOTrainer:
55     """
56     The PPO_trainer uses Proximal Policy Optimization to optimise language models.
57     """
58 
59     default_params = {
60         "lr": 1.41e-5,
61         "adap_kl_ctrl": True,
62         "init_kl_coef":0.2,
63         "target": 6,
64         "horizon":10000,
65         "gamma":1,
66         "lam":0.95,
67         "cliprange": .2,
68         "cliprange_value":.2,
69         "vf_coef":.1,
70         "batch_size": 256,
71         "forward_batch_size": 16,
72         "ppo_epochs": 4,
73     }
74 
75     def __init__(self, model, ref_model, **ppo_params):
76         """
77         Initialize PPOTrainer.
78 
79         Args:
80             model (torch.model): Hugging Face transformer GPT2 model with value head
81             ref_model (torch.model): Hugging Face transformer GPT2 refrence model used for KL penalty
82             ppo_params (dict or None): PPO parameters for training. Can include following keys:
83                 'lr' (float): Adam learning rate, default: 1.41e-5
84                 'batch_size' (int): Number of samples per optimisation step, default: 256
85                 'forward_batch_size' (int): Number of samples forward passed through model at a time, default: 16
86                 'ppo_epochs' (int): Number of optimisation epochs per batch of samples, default: 4
87                 'gamma' (float)): Gamma parameter for advantage calculation, default: 1.
88                 'lam' (float): Lambda parameter for advantage calcualation, default: 0.95
89                 'cliprange_value' (float): Range for clipping values in loss calculation, default: 0.2
90                 'cliprange' (float): Range for clipping in PPO policy gradient loss, default: 0.2
91                 'vf_coef' (float): Scaling factor for value loss, default: 0.1
92                 'adap_kl_ctrl' (bool): Use adaptive KL control, otherwise linear, default: True
93                 'init_kl_coef' (float): Initial KL penalty coefficient (used for adaptive and linear control), default: 0.2
94                 'target' (float): Target KL value for adaptive KL control, default: 6.0
95                 'horizon' (float): Horizon for adaptive KL control, default: 10000
96 
97         """
98         self.ppo_params = self.default_params
99         self.ppo_params.update(ppo_params)
100 
101         self.ref_model = ref_model
102         self.model = model
103         self.optimizer = Adam(model.parameters(), lr=self.ppo_params['lr'])
104 
105         self.kl_ctl = AdaptiveKLController(self.ppo_params['init_kl_coef'],
106                                            self.ppo_params['target'],
107                                            self.ppo_params['horizon'])
108 
109 
110     def step(self, query, response, scores):
111         """
112         Run a PPO optimisation step.
113 
114         args:
115             query (torch.tensor): tensor containing the encoded queries, shape [batch_size, query_length]
116             response (torch.tensor): tensor containing the encoded responses, shape [batch_size, response_length]
117             scores (torch.tensor): tensor containing the scores, shape [batch_size]
118 
119         returns:
120             train_stats (dict): a summary of the training statistics
121         """
122 
123         bs = self.ppo_params['batch_size']
124         timing = dict()
125         t0 = time.time()
126 
127         gen_len = response.shape[1]
128         model_input = torch.cat((query, response), axis=1)
129 
130         t = time.time()
131         logprobs, ref_logprobs, values = self.batched_forward_pass(model_input, gen_len)
132         timing['time/ppo/forward_pass'] = time.time()-t
133 
134         t = time.time()
135         rewards, non_score_reward, kl_coef = self.compute_rewards(scores, logprobs, ref_logprobs)
136         timing['time/ppo/compute_rewards'] = time.time()-t
137 
138         t = time.time()
139         all_stats = []
140         idxs = list(range(bs))
141         for _ in range(self.ppo_params['ppo_epochs']):
142             random.shuffle(idxs)
143             for i in range(bs):
144                 idx = idxs[i]
145                 train_stats = self.train_minibatch(logprobs[idx:idx+1], values[idx:idx+1],
146                                                    rewards[idx:idx+1], query[idx:idx+1],
147                                                    response[idx:idx+1], model_input[idx:idx+1])
148                 all_stats.append(train_stats)
149         timing['time/ppo/optimize_step'] = time.time()-t
150 
151         t = time.time()
152         train_stats = stack_dicts(all_stats)
153 
154         # reshape advantages/ratios such that they are not averaged.
155         train_stats['policy/advantages'] = torch.flatten(train_stats['policy/advantages']).unsqueeze(0)
156         train_stats['policy/ratio'] = torch.flatten(train_stats['policy/ratio']).unsqueeze(0)
157 
158         stats = self.record_step_stats(scores=scores, logprobs=logprobs, ref_logprobs=ref_logprobs,
159                                        non_score_reward=non_score_reward, train_stats=train_stats,
160                                        kl_coef=kl_coef)
161         stats = stats_to_np(stats)
162         timing['time/ppo/calc_stats'] = time.time()-t
163 
164         self.kl_ctl.update(stats['objective/kl'], self.ppo_params['batch_size'])
165 
166         timing['time/ppo/total'] = time.time()-t0
167         stats.update(timing)
168         return stats
169 
170     def batched_forward_pass(self, model_input, gen_len):
171         """Calculate model outputs in multiple batches."""
172         bs = self.ppo_params['batch_size']
173         fbs = self.ppo_params['forward_batch_size']
174         logprobs = []
175         ref_logprobs = []
176         values = []
177 
178         for i in range(int(self.ppo_params['batch_size']/fbs)):
179             m_input = model_input[i*fbs:(i+1)*fbs]
180             logits, _, v = self.model(m_input)
181             ref_logits, _, _ = self.ref_model(m_input)
182 
183             values.append(v[:, -gen_len-1:-1].detach())
184             logprobs.append(logprobs_from_logits(logits[:,:-1,:], m_input[:,1:])[:, -gen_len:].detach())
185             ref_logprobs.append(logprobs_from_logits(ref_logits[:,:-1,:], m_input[:,1:])[:, -gen_len:].detach())
186 
187         return torch.cat(logprobs), torch.cat(ref_logprobs), torch.cat(values)
188 
189     def train_minibatch(self, logprobs, values, rewards, query, response, model_input):
190         """Train one PPO minibatch"""
191         loss_p, loss_v, train_stats  = self.loss(logprobs, values, rewards, query, response, model_input)
192         loss = loss_p + loss_v
193         self.optimizer.zero_grad()
194         loss.backward()
195         self.optimizer.step()
196         return train_stats
197 
198     def compute_rewards(self, scores, logprobs, ref_logprobs):
199         """Compute per token rewards from scores and KL-penalty."""
200         kl = logprobs - ref_logprobs
201         non_score_reward = -self.kl_ctl.value * kl
202         rewards = non_score_reward.clone().detach()
203         rewards[:, -1] += scores
204         return rewards, non_score_reward, self.kl_ctl.value
205 
206     def loss(self, old_logprobs, values, rewards, query, response, model_input):
207         """Calculate policy and value losses."""
208         lastgaelam = 0
209         advantages_reversed = []
210         gen_len = response.shape[1]
211 
212         for t in reversed(range(gen_len)):
213             nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0
214             delta = rewards[:, t] + self.ppo_params['gamma'] * nextvalues - values[:, t]
215             lastgaelam = delta + self.ppo_params['gamma'] * self.ppo_params['lam'] * lastgaelam
216             advantages_reversed.append(lastgaelam)
217         advantages = torch.stack(advantages_reversed[::-1]).transpose(0, 1)
218 
219         returns = advantages + values
220         advantages = whiten(advantages)
221         advantages = advantages.detach()
222 
223         logits, _, vpred = self.model(model_input)
224         logprob = logprobs_from_logits(logits[:,:-1,:], model_input[:, 1:])
225 
226         #only the generation part of the values/logprobs is needed
227         logprob, vpred = logprob[:, -gen_len:], vpred[:,-gen_len-1:-1]
228 
229         vpredclipped = clip_by_value(vpred,
230                                      values - self.ppo_params["cliprange_value"],
231                                      values + self.ppo_params["cliprange_value"])
232 
233         vf_losses1 = (vpred - returns)**2
234         vf_losses2 = (vpredclipped - returns)**2
235         vf_loss = .5 * torch.mean(torch.max(vf_losses1, vf_losses2))
236         vf_clipfrac =  torch.mean(torch.gt(vf_losses2, vf_losses1).double())
237 
238         ratio = torch.exp(logprob - old_logprobs)
239 
240         pg_losses = -advantages * ratio
241         pg_losses2 = -advantages * torch.clamp(ratio,
242                                                1.0 - self.ppo_params['cliprange'],
243                                                1.0 + self.ppo_params['cliprange'])
244 
245         pg_loss = torch.mean(torch.max(pg_losses, pg_losses2))
246         pg_clipfrac = torch.mean(torch.gt(pg_losses2, pg_losses).double())
247 
248         loss = pg_loss + self.ppo_params['vf_coef'] * vf_loss
249 
250         entropy = torch.mean(entropy_from_logits(logits))
251         approxkl = .5 * torch.mean((logprob - old_logprobs)**2)
252         policykl = torch.mean(logprob - old_logprobs)
253         return_mean, return_var = torch.mean(returns), torch.var(returns)
254         value_mean, value_var = torch.mean(values), torch.var(values)
255 
256         stats = dict(
257             loss=dict(policy=pg_loss, value=vf_loss, total=loss),
258             policy=dict(entropy=entropy, approxkl=approxkl,policykl=policykl, clipfrac=pg_clipfrac,
259                         advantages=advantages, advantages_mean=torch.mean(advantages), ratio=ratio),
260             returns=dict(mean=return_mean, var=return_var),
261             val=dict(vpred=torch.mean(vpred), error=torch.mean((vpred - returns) ** 2),
262                      clipfrac=vf_clipfrac, mean=value_mean, var=value_var),
263         )
264         return pg_loss, self.ppo_params['vf_coef'] * vf_loss, flatten_dict(stats)
265 
266 
267     def record_step_stats(self, kl_coef, **data):
268         """Record training step statistics."""
269         kl = data['logprobs'] - data['ref_logprobs']
270         mean_kl = torch.mean(torch.sum(kl, axis=-1))
271         mean_entropy = torch.mean(torch.sum(-data['logprobs'], axis=1))
272         mean_non_score_reward =torch.mean(torch.sum(data['non_score_reward'], axis=1))
273         stats = {
274             'objective/kl': mean_kl,
275             'objective/kl_dist': kl,
276             'objective/logprobs': data['logprobs'],
277             'objective/ref_logprobs': data['ref_logprobs'],
278             'objective/kl_coef': kl_coef,
279             'objective/entropy': mean_entropy,
280             'ppo/mean_non_score_reward': mean_non_score_reward,
281         }
282 
283         for k, v in data['train_stats'].items():
284             stats[f'ppo/{k}'] = torch.mean(v, axis=0)
285         stats['ppo/val/var_explained'] = 1 - stats['ppo/val/error'] / stats['ppo/returns/var']
286         return stats
287
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/trl/ppo.py b/trl/ppo.py
--- a/trl/ppo.py
+++ b/trl/ppo.py
@@ -102,9 +102,12 @@
         self.model = model
         self.optimizer = Adam(model.parameters(), lr=self.ppo_params['lr'])
 
-        self.kl_ctl = AdaptiveKLController(self.ppo_params['init_kl_coef'],
-                                           self.ppo_params['target'],
-                                           self.ppo_params['horizon'])
+        if self.ppo_params['adap_kl_ctrl']:
+            self.kl_ctl = AdaptiveKLController(self.ppo_params['init_kl_coef'],
+                                               self.ppo_params['target'],
+                                               self.ppo_params['horizon'])
+        else:
+            self.kl_ctl = FixedKLController(self.ppo_params['init_kl_coef'])
 
 
     def step(self, query, response, scores):

{"golden_diff": "diff --git a/trl/ppo.py b/trl/ppo.py\n--- a/trl/ppo.py\n+++ b/trl/ppo.py\n@@ -102,9 +102,12 @@\n         self.model = model\n         self.optimizer = Adam(model.parameters(), lr=self.ppo_params['lr'])\n \n-        self.kl_ctl = AdaptiveKLController(self.ppo_params['init_kl_coef'],\n-                                           self.ppo_params['target'],\n-                                           self.ppo_params['horizon'])\n+        if self.ppo_params['adap_kl_ctrl']:\n+            self.kl_ctl = AdaptiveKLController(self.ppo_params['init_kl_coef'],\n+                                               self.ppo_params['target'],\n+                                               self.ppo_params['horizon'])\n+        else:\n+            self.kl_ctl = FixedKLController(self.ppo_params['init_kl_coef'])\n \n \n     def step(self, query, response, scores):\n", "issue": "adap_kl_ctrl boolean is not used \nHello,\r\n\r\nThanks for this implementation, is it possible that your kl adaptation is always on ?\r\nhttps://github.com/lvwerra/trl/blob/750f5fd5329bb81c79b00243c4c8923ac14981d5/trl/ppo.py#L92\r\n\r\nBest,\r\nThibaud\n", "before_files": [{"content": "# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/02-ppo.ipynb (unless otherwise specified).\n\n__all__ = ['AdaptiveKLController', 'FixedKLController', 'PPOTrainer']\n\n# Cell\nimport numpy as np\nimport torch.nn.functional as F\nfrom torch.optim import Adam\nimport torch\nimport collections\nimport time\nimport random\n\nfrom .core import (logprobs_from_logits,\n                         whiten,\n                         clip_by_value,\n                         entropy_from_logits,\n                         flatten_dict,\n                         average_torch_dicts,\n                         stats_to_np,\n                         stack_dicts,\n                         add_suffix)\n\n# Cell\n\nclass AdaptiveKLController:\n    \"\"\"\n    Adaptive KL controller described in the paper:\n    https://arxiv.org/pdf/1909.08593.pdf\n    \"\"\"\n    def __init__(self, init_kl_coef, target, horizon):\n        self.value = init_kl_coef\n        self.target = target\n        self.horizon = horizon\n\n    def update(self, current, n_steps):\n        target = self.target\n        proportional_error = np.clip(current / target - 1, -0.2, 0.2)\n        mult = 1 + proportional_error * n_steps / self.horizon\n        self.value *= mult\n\n# Cell\n\nclass FixedKLController:\n    \"\"\"Fixed KL controller.\"\"\"\n    def __init__(self, kl_coef):\n        self.value = kl_coef\n\n    def update(self, current, n_steps):\n        pass\n\n# Cell\n\nclass PPOTrainer:\n    \"\"\"\n    The PPO_trainer uses Proximal Policy Optimization to optimise language models.\n    \"\"\"\n\n    default_params = {\n        \"lr\": 1.41e-5,\n        \"adap_kl_ctrl\": True,\n        \"init_kl_coef\":0.2,\n        \"target\": 6,\n        \"horizon\":10000,\n        \"gamma\":1,\n        \"lam\":0.95,\n        \"cliprange\": .2,\n        \"cliprange_value\":.2,\n        \"vf_coef\":.1,\n        \"batch_size\": 256,\n        \"forward_batch_size\": 16,\n        \"ppo_epochs\": 4,\n    }\n\n    def __init__(self, model, ref_model, **ppo_params):\n        \"\"\"\n        Initialize PPOTrainer.\n\n        Args:\n            model (torch.model): Hugging Face transformer GPT2 model with value head\n            ref_model (torch.model): Hugging Face transformer GPT2 refrence model used for KL penalty\n            ppo_params (dict or None): PPO parameters for training. Can include following keys:\n                'lr' (float): Adam learning rate, default: 1.41e-5\n                'batch_size' (int): Number of samples per optimisation step, default: 256\n                'forward_batch_size' (int): Number of samples forward passed through model at a time, default: 16\n                'ppo_epochs' (int): Number of optimisation epochs per batch of samples, default: 4\n                'gamma' (float)): Gamma parameter for advantage calculation, default: 1.\n                'lam' (float): Lambda parameter for advantage calcualation, default: 0.95\n                'cliprange_value' (float): Range for clipping values in loss calculation, default: 0.2\n                'cliprange' (float): Range for clipping in PPO policy gradient loss, default: 0.2\n                'vf_coef' (float): Scaling factor for value loss, default: 0.1\n                'adap_kl_ctrl' (bool): Use adaptive KL control, otherwise linear, default: True\n                'init_kl_coef' (float): Initial KL penalty coefficient (used for adaptive and linear control), default: 0.2\n                'target' (float): Target KL value for adaptive KL control, default: 6.0\n                'horizon' (float): Horizon for adaptive KL control, default: 10000\n\n        \"\"\"\n        self.ppo_params = self.default_params\n        self.ppo_params.update(ppo_params)\n\n        self.ref_model = ref_model\n        self.model = model\n        self.optimizer = Adam(model.parameters(), lr=self.ppo_params['lr'])\n\n        self.kl_ctl = AdaptiveKLController(self.ppo_params['init_kl_coef'],\n                                           self.ppo_params['target'],\n                                           self.ppo_params['horizon'])\n\n\n    def step(self, query, response, scores):\n        \"\"\"\n        Run a PPO optimisation step.\n\n        args:\n            query (torch.tensor): tensor containing the encoded queries, shape [batch_size, query_length]\n            response (torch.tensor): tensor containing the encoded responses, shape [batch_size, response_length]\n            scores (torch.tensor): tensor containing the scores, shape [batch_size]\n\n        returns:\n            train_stats (dict): a summary of the training statistics\n        \"\"\"\n\n        bs = self.ppo_params['batch_size']\n        timing = dict()\n        t0 = time.time()\n\n        gen_len = response.shape[1]\n        model_input = torch.cat((query, response), axis=1)\n\n        t = time.time()\n        logprobs, ref_logprobs, values = self.batched_forward_pass(model_input, gen_len)\n        timing['time/ppo/forward_pass'] = time.time()-t\n\n        t = time.time()\n        rewards, non_score_reward, kl_coef = self.compute_rewards(scores, logprobs, ref_logprobs)\n        timing['time/ppo/compute_rewards'] = time.time()-t\n\n        t = time.time()\n        all_stats = []\n        idxs = list(range(bs))\n        for _ in range(self.ppo_params['ppo_epochs']):\n            random.shuffle(idxs)\n            for i in range(bs):\n                idx = idxs[i]\n                train_stats = self.train_minibatch(logprobs[idx:idx+1], values[idx:idx+1],\n                                                   rewards[idx:idx+1], query[idx:idx+1],\n                                                   response[idx:idx+1], model_input[idx:idx+1])\n                all_stats.append(train_stats)\n        timing['time/ppo/optimize_step'] = time.time()-t\n\n        t = time.time()\n        train_stats = stack_dicts(all_stats)\n\n        # reshape advantages/ratios such that they are not averaged.\n        train_stats['policy/advantages'] = torch.flatten(train_stats['policy/advantages']).unsqueeze(0)\n        train_stats['policy/ratio'] = torch.flatten(train_stats['policy/ratio']).unsqueeze(0)\n\n        stats = self.record_step_stats(scores=scores, logprobs=logprobs, ref_logprobs=ref_logprobs,\n                                       non_score_reward=non_score_reward, train_stats=train_stats,\n                                       kl_coef=kl_coef)\n        stats = stats_to_np(stats)\n        timing['time/ppo/calc_stats'] = time.time()-t\n\n        self.kl_ctl.update(stats['objective/kl'], self.ppo_params['batch_size'])\n\n        timing['time/ppo/total'] = time.time()-t0\n        stats.update(timing)\n        return stats\n\n    def batched_forward_pass(self, model_input, gen_len):\n        \"\"\"Calculate model outputs in multiple batches.\"\"\"\n        bs = self.ppo_params['batch_size']\n        fbs = self.ppo_params['forward_batch_size']\n        logprobs = []\n        ref_logprobs = []\n        values = []\n\n        for i in range(int(self.ppo_params['batch_size']/fbs)):\n            m_input = model_input[i*fbs:(i+1)*fbs]\n            logits, _, v = self.model(m_input)\n            ref_logits, _, _ = self.ref_model(m_input)\n\n            values.append(v[:, -gen_len-1:-1].detach())\n            logprobs.append(logprobs_from_logits(logits[:,:-1,:], m_input[:,1:])[:, -gen_len:].detach())\n            ref_logprobs.append(logprobs_from_logits(ref_logits[:,:-1,:], m_input[:,1:])[:, -gen_len:].detach())\n\n        return torch.cat(logprobs), torch.cat(ref_logprobs), torch.cat(values)\n\n    def train_minibatch(self, logprobs, values, rewards, query, response, model_input):\n        \"\"\"Train one PPO minibatch\"\"\"\n        loss_p, loss_v, train_stats  = self.loss(logprobs, values, rewards, query, response, model_input)\n        loss = loss_p + loss_v\n        self.optimizer.zero_grad()\n        loss.backward()\n        self.optimizer.step()\n        return train_stats\n\n    def compute_rewards(self, scores, logprobs, ref_logprobs):\n        \"\"\"Compute per token rewards from scores and KL-penalty.\"\"\"\n        kl = logprobs - ref_logprobs\n        non_score_reward = -self.kl_ctl.value * kl\n        rewards = non_score_reward.clone().detach()\n        rewards[:, -1] += scores\n        return rewards, non_score_reward, self.kl_ctl.value\n\n    def loss(self, old_logprobs, values, rewards, query, response, model_input):\n        \"\"\"Calculate policy and value losses.\"\"\"\n        lastgaelam = 0\n        advantages_reversed = []\n        gen_len = response.shape[1]\n\n        for t in reversed(range(gen_len)):\n            nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0\n            delta = rewards[:, t] + self.ppo_params['gamma'] * nextvalues - values[:, t]\n            lastgaelam = delta + self.ppo_params['gamma'] * self.ppo_params['lam'] * lastgaelam\n            advantages_reversed.append(lastgaelam)\n        advantages = torch.stack(advantages_reversed[::-1]).transpose(0, 1)\n\n        returns = advantages + values\n        advantages = whiten(advantages)\n        advantages = advantages.detach()\n\n        logits, _, vpred = self.model(model_input)\n        logprob = logprobs_from_logits(logits[:,:-1,:], model_input[:, 1:])\n\n        #only the generation part of the values/logprobs is needed\n        logprob, vpred = logprob[:, -gen_len:], vpred[:,-gen_len-1:-1]\n\n        vpredclipped = clip_by_value(vpred,\n                                     values - self.ppo_params[\"cliprange_value\"],\n                                     values + self.ppo_params[\"cliprange_value\"])\n\n        vf_losses1 = (vpred - returns)**2\n        vf_losses2 = (vpredclipped - returns)**2\n        vf_loss = .5 * torch.mean(torch.max(vf_losses1, vf_losses2))\n        vf_clipfrac =  torch.mean(torch.gt(vf_losses2, vf_losses1).double())\n\n        ratio = torch.exp(logprob - old_logprobs)\n\n        pg_losses = -advantages * ratio\n        pg_losses2 = -advantages * torch.clamp(ratio,\n                                               1.0 - self.ppo_params['cliprange'],\n                                               1.0 + self.ppo_params['cliprange'])\n\n        pg_loss = torch.mean(torch.max(pg_losses, pg_losses2))\n        pg_clipfrac = torch.mean(torch.gt(pg_losses2, pg_losses).double())\n\n        loss = pg_loss + self.ppo_params['vf_coef'] * vf_loss\n\n        entropy = torch.mean(entropy_from_logits(logits))\n        approxkl = .5 * torch.mean((logprob - old_logprobs)**2)\n        policykl = torch.mean(logprob - old_logprobs)\n        return_mean, return_var = torch.mean(returns), torch.var(returns)\n        value_mean, value_var = torch.mean(values), torch.var(values)\n\n        stats = dict(\n            loss=dict(policy=pg_loss, value=vf_loss, total=loss),\n            policy=dict(entropy=entropy, approxkl=approxkl,policykl=policykl, clipfrac=pg_clipfrac,\n                        advantages=advantages, advantages_mean=torch.mean(advantages), ratio=ratio),\n            returns=dict(mean=return_mean, var=return_var),\n            val=dict(vpred=torch.mean(vpred), error=torch.mean((vpred - returns) ** 2),\n                     clipfrac=vf_clipfrac, mean=value_mean, var=value_var),\n        )\n        return pg_loss, self.ppo_params['vf_coef'] * vf_loss, flatten_dict(stats)\n\n\n    def record_step_stats(self, kl_coef, **data):\n        \"\"\"Record training step statistics.\"\"\"\n        kl = data['logprobs'] - data['ref_logprobs']\n        mean_kl = torch.mean(torch.sum(kl, axis=-1))\n        mean_entropy = torch.mean(torch.sum(-data['logprobs'], axis=1))\n        mean_non_score_reward =torch.mean(torch.sum(data['non_score_reward'], axis=1))\n        stats = {\n            'objective/kl': mean_kl,\n            'objective/kl_dist': kl,\n            'objective/logprobs': data['logprobs'],\n            'objective/ref_logprobs': data['ref_logprobs'],\n            'objective/kl_coef': kl_coef,\n            'objective/entropy': mean_entropy,\n            'ppo/mean_non_score_reward': mean_non_score_reward,\n        }\n\n        for k, v in data['train_stats'].items():\n            stats[f'ppo/{k}'] = torch.mean(v, axis=0)\n        stats['ppo/val/var_explained'] = 1 - stats['ppo/val/error'] / stats['ppo/returns/var']\n        return stats\n", "path": "trl/ppo.py"}], "after_files": [{"content": "# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/02-ppo.ipynb (unless otherwise specified).\n\n__all__ = ['AdaptiveKLController', 'FixedKLController', 'PPOTrainer']\n\n# Cell\nimport numpy as np\nimport torch.nn.functional as F\nfrom torch.optim import Adam\nimport torch\nimport collections\nimport time\nimport random\n\nfrom .core import (logprobs_from_logits,\n                         whiten,\n                         clip_by_value,\n                         entropy_from_logits,\n                         flatten_dict,\n                         average_torch_dicts,\n                         stats_to_np,\n                         stack_dicts,\n                         add_suffix)\n\n# Cell\n\nclass AdaptiveKLController:\n    \"\"\"\n    Adaptive KL controller described in the paper:\n    https://arxiv.org/pdf/1909.08593.pdf\n    \"\"\"\n    def __init__(self, init_kl_coef, target, horizon):\n        self.value = init_kl_coef\n        self.target = target\n        self.horizon = horizon\n\n    def update(self, current, n_steps):\n        target = self.target\n        proportional_error = np.clip(current / target - 1, -0.2, 0.2)\n        mult = 1 + proportional_error * n_steps / self.horizon\n        self.value *= mult\n\n# Cell\n\nclass FixedKLController:\n    \"\"\"Fixed KL controller.\"\"\"\n    def __init__(self, kl_coef):\n        self.value = kl_coef\n\n    def update(self, current, n_steps):\n        pass\n\n# Cell\n\nclass PPOTrainer:\n    \"\"\"\n    The PPO_trainer uses Proximal Policy Optimization to optimise language models.\n    \"\"\"\n\n    default_params = {\n        \"lr\": 1.41e-5,\n        \"adap_kl_ctrl\": True,\n        \"init_kl_coef\":0.2,\n        \"target\": 6,\n        \"horizon\":10000,\n        \"gamma\":1,\n        \"lam\":0.95,\n        \"cliprange\": .2,\n        \"cliprange_value\":.2,\n        \"vf_coef\":.1,\n        \"batch_size\": 256,\n        \"forward_batch_size\": 16,\n        \"ppo_epochs\": 4,\n    }\n\n    def __init__(self, model, ref_model, **ppo_params):\n        \"\"\"\n        Initialize PPOTrainer.\n\n        Args:\n            model (torch.model): Hugging Face transformer GPT2 model with value head\n            ref_model (torch.model): Hugging Face transformer GPT2 refrence model used for KL penalty\n            ppo_params (dict or None): PPO parameters for training. Can include following keys:\n                'lr' (float): Adam learning rate, default: 1.41e-5\n                'batch_size' (int): Number of samples per optimisation step, default: 256\n                'forward_batch_size' (int): Number of samples forward passed through model at a time, default: 16\n                'ppo_epochs' (int): Number of optimisation epochs per batch of samples, default: 4\n                'gamma' (float)): Gamma parameter for advantage calculation, default: 1.\n                'lam' (float): Lambda parameter for advantage calcualation, default: 0.95\n                'cliprange_value' (float): Range for clipping values in loss calculation, default: 0.2\n                'cliprange' (float): Range for clipping in PPO policy gradient loss, default: 0.2\n                'vf_coef' (float): Scaling factor for value loss, default: 0.1\n                'adap_kl_ctrl' (bool): Use adaptive KL control, otherwise linear, default: True\n                'init_kl_coef' (float): Initial KL penalty coefficient (used for adaptive and linear control), default: 0.2\n                'target' (float): Target KL value for adaptive KL control, default: 6.0\n                'horizon' (float): Horizon for adaptive KL control, default: 10000\n\n        \"\"\"\n        self.ppo_params = self.default_params\n        self.ppo_params.update(ppo_params)\n\n        self.ref_model = ref_model\n        self.model = model\n        self.optimizer = Adam(model.parameters(), lr=self.ppo_params['lr'])\n\n        if self.ppo_params['adap_kl_ctrl']:\n            self.kl_ctl = AdaptiveKLController(self.ppo_params['init_kl_coef'],\n                                               self.ppo_params['target'],\n                                               self.ppo_params['horizon'])\n        else:\n            self.kl_ctl = FixedKLController(self.ppo_params['init_kl_coef'])\n\n\n    def step(self, query, response, scores):\n        \"\"\"\n        Run a PPO optimisation step.\n\n        args:\n            query (torch.tensor): tensor containing the encoded queries, shape [batch_size, query_length]\n            response (torch.tensor): tensor containing the encoded responses, shape [batch_size, response_length]\n            scores (torch.tensor): tensor containing the scores, shape [batch_size]\n\n        returns:\n            train_stats (dict): a summary of the training statistics\n        \"\"\"\n\n        bs = self.ppo_params['batch_size']\n        timing = dict()\n        t0 = time.time()\n\n        gen_len = response.shape[1]\n        model_input = torch.cat((query, response), axis=1)\n\n        t = time.time()\n        logprobs, ref_logprobs, values = self.batched_forward_pass(model_input, gen_len)\n        timing['time/ppo/forward_pass'] = time.time()-t\n\n        t = time.time()\n        rewards, non_score_reward, kl_coef = self.compute_rewards(scores, logprobs, ref_logprobs)\n        timing['time/ppo/compute_rewards'] = time.time()-t\n\n        t = time.time()\n        all_stats = []\n        idxs = list(range(bs))\n        for _ in range(self.ppo_params['ppo_epochs']):\n            random.shuffle(idxs)\n            for i in range(bs):\n                idx = idxs[i]\n                train_stats = self.train_minibatch(logprobs[idx:idx+1], values[idx:idx+1],\n                                                   rewards[idx:idx+1], query[idx:idx+1],\n                                                   response[idx:idx+1], model_input[idx:idx+1])\n                all_stats.append(train_stats)\n        timing['time/ppo/optimize_step'] = time.time()-t\n\n        t = time.time()\n        train_stats = stack_dicts(all_stats)\n\n        # reshape advantages/ratios such that they are not averaged.\n        train_stats['policy/advantages'] = torch.flatten(train_stats['policy/advantages']).unsqueeze(0)\n        train_stats['policy/ratio'] = torch.flatten(train_stats['policy/ratio']).unsqueeze(0)\n\n        stats = self.record_step_stats(scores=scores, logprobs=logprobs, ref_logprobs=ref_logprobs,\n                                       non_score_reward=non_score_reward, train_stats=train_stats,\n                                       kl_coef=kl_coef)\n        stats = stats_to_np(stats)\n        timing['time/ppo/calc_stats'] = time.time()-t\n\n        self.kl_ctl.update(stats['objective/kl'], self.ppo_params['batch_size'])\n\n        timing['time/ppo/total'] = time.time()-t0\n        stats.update(timing)\n        return stats\n\n    def batched_forward_pass(self, model_input, gen_len):\n        \"\"\"Calculate model outputs in multiple batches.\"\"\"\n        bs = self.ppo_params['batch_size']\n        fbs = self.ppo_params['forward_batch_size']\n        logprobs = []\n        ref_logprobs = []\n        values = []\n\n        for i in range(int(self.ppo_params['batch_size']/fbs)):\n            m_input = model_input[i*fbs:(i+1)*fbs]\n            logits, _, v = self.model(m_input)\n            ref_logits, _, _ = self.ref_model(m_input)\n\n            values.append(v[:, -gen_len-1:-1].detach())\n            logprobs.append(logprobs_from_logits(logits[:,:-1,:], m_input[:,1:])[:, -gen_len:].detach())\n            ref_logprobs.append(logprobs_from_logits(ref_logits[:,:-1,:], m_input[:,1:])[:, -gen_len:].detach())\n\n        return torch.cat(logprobs), torch.cat(ref_logprobs), torch.cat(values)\n\n    def train_minibatch(self, logprobs, values, rewards, query, response, model_input):\n        \"\"\"Train one PPO minibatch\"\"\"\n        loss_p, loss_v, train_stats  = self.loss(logprobs, values, rewards, query, response, model_input)\n        loss = loss_p + loss_v\n        self.optimizer.zero_grad()\n        loss.backward()\n        self.optimizer.step()\n        return train_stats\n\n    def compute_rewards(self, scores, logprobs, ref_logprobs):\n        \"\"\"Compute per token rewards from scores and KL-penalty.\"\"\"\n        kl = logprobs - ref_logprobs\n        non_score_reward = -self.kl_ctl.value * kl\n        rewards = non_score_reward.clone().detach()\n        rewards[:, -1] += scores\n        return rewards, non_score_reward, self.kl_ctl.value\n\n    def loss(self, old_logprobs, values, rewards, query, response, model_input):\n        \"\"\"Calculate policy and value losses.\"\"\"\n        lastgaelam = 0\n        advantages_reversed = []\n        gen_len = response.shape[1]\n\n        for t in reversed(range(gen_len)):\n            nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0\n            delta = rewards[:, t] + self.ppo_params['gamma'] * nextvalues - values[:, t]\n            lastgaelam = delta + self.ppo_params['gamma'] * self.ppo_params['lam'] * lastgaelam\n            advantages_reversed.append(lastgaelam)\n        advantages = torch.stack(advantages_reversed[::-1]).transpose(0, 1)\n\n        returns = advantages + values\n        advantages = whiten(advantages)\n        advantages = advantages.detach()\n\n        logits, _, vpred = self.model(model_input)\n        logprob = logprobs_from_logits(logits[:,:-1,:], model_input[:, 1:])\n\n        #only the generation part of the values/logprobs is needed\n        logprob, vpred = logprob[:, -gen_len:], vpred[:,-gen_len-1:-1]\n\n        vpredclipped = clip_by_value(vpred,\n                                     values - self.ppo_params[\"cliprange_value\"],\n                                     values + self.ppo_params[\"cliprange_value\"])\n\n        vf_losses1 = (vpred - returns)**2\n        vf_losses2 = (vpredclipped - returns)**2\n        vf_loss = .5 * torch.mean(torch.max(vf_losses1, vf_losses2))\n        vf_clipfrac =  torch.mean(torch.gt(vf_losses2, vf_losses1).double())\n\n        ratio = torch.exp(logprob - old_logprobs)\n\n        pg_losses = -advantages * ratio\n        pg_losses2 = -advantages * torch.clamp(ratio,\n                                               1.0 - self.ppo_params['cliprange'],\n                                               1.0 + self.ppo_params['cliprange'])\n\n        pg_loss = torch.mean(torch.max(pg_losses, pg_losses2))\n        pg_clipfrac = torch.mean(torch.gt(pg_losses2, pg_losses).double())\n\n        loss = pg_loss + self.ppo_params['vf_coef'] * vf_loss\n\n        entropy = torch.mean(entropy_from_logits(logits))\n        approxkl = .5 * torch.mean((logprob - old_logprobs)**2)\n        policykl = torch.mean(logprob - old_logprobs)\n        return_mean, return_var = torch.mean(returns), torch.var(returns)\n        value_mean, value_var = torch.mean(values), torch.var(values)\n\n        stats = dict(\n            loss=dict(policy=pg_loss, value=vf_loss, total=loss),\n            policy=dict(entropy=entropy, approxkl=approxkl,policykl=policykl, clipfrac=pg_clipfrac,\n                        advantages=advantages, advantages_mean=torch.mean(advantages), ratio=ratio),\n            returns=dict(mean=return_mean, var=return_var),\n            val=dict(vpred=torch.mean(vpred), error=torch.mean((vpred - returns) ** 2),\n                     clipfrac=vf_clipfrac, mean=value_mean, var=value_var),\n        )\n        return pg_loss, self.ppo_params['vf_coef'] * vf_loss, flatten_dict(stats)\n\n\n    def record_step_stats(self, kl_coef, **data):\n        \"\"\"Record training step statistics.\"\"\"\n        kl = data['logprobs'] - data['ref_logprobs']\n        mean_kl = torch.mean(torch.sum(kl, axis=-1))\n        mean_entropy = torch.mean(torch.sum(-data['logprobs'], axis=1))\n        mean_non_score_reward =torch.mean(torch.sum(data['non_score_reward'], axis=1))\n        stats = {\n            'objective/kl': mean_kl,\n            'objective/kl_dist': kl,\n            'objective/logprobs': data['logprobs'],\n            'objective/ref_logprobs': data['ref_logprobs'],\n            'objective/kl_coef': kl_coef,\n            'objective/entropy': mean_entropy,\n            'ppo/mean_non_score_reward': mean_non_score_reward,\n        }\n\n        for k, v in data['train_stats'].items():\n            stats[f'ppo/{k}'] = torch.mean(v, axis=0)\n        stats['ppo/val/var_explained'] = 1 - stats['ppo/val/error'] / stats['ppo/returns/var']\n        return stats\n", "path": "trl/ppo.py"}]}