nbs/02-ppo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PPO for transformer models\n",
    "> A Pytorch implementation of Proximal Policy Optimization for transfomer models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This follows the language model approach proposed in paper [\"Fine-Tuning Language Models from Human Preferences\"](\n",
    "https://arxiv.org/pdf/1909.08593.pdf) and is similar to the [original implementation](https://github.com/openai/lm-human-preferences). The two main differences are 1) the method is implemented in Pytorch and 2) works with the `transformer` library by Hugging Face."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp ppo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "import numpy as np\n",
    "import torch.nn.functional as F\n",
    "from torch.optim import Adam\n",
    "import torch\n",
    "import collections\n",
    "import time\n",
    "import random\n",
    "\n",
    "from transformers import DataCollatorForLanguageModeling\n",
    "\n",
    "from trl.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",
    "                      WANDB_PADDING)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## KL-controllers\n",
    "To ensure that the learned policy does not deviate to much from the original language model the KL divergence between the policy and a reference policy (the language model before PPO training) is used as an additional reward signal. Large KL-divergences are punished and staying close to the reference is rewarded.\n",
    "\n",
    "Two controllers are presented in the paper: an adaptive log-space proportional controller and a fixed controller."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# exports\n",
    "\n",
    "class 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# exports \n",
    "\n",
    "class 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# exports \n",
    "\n",
    "class 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, tokenizer, **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",
    "            tokenizer (tokenizer): Hugging Face tokenizer\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.tokenizer = tokenizer\n",
    "        self.data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)\n",
    "        \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, queries, responses, scores):\n",
    "        \"\"\"\n",
    "        Run a PPO optimisation step.\n",
    "        \n",
    "        args:\n",
    "            queries (List): List of tensors containing the encoded queries, shape [query_length]\n",
    "            responses (List): List of tensors containing the encoded responses, shape [response_length]\n",
    "            scores (List): 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",
    "        assert bs == len(queries), f\"Batch size ({bs}) does not match number of examples ({len(queries)})\"\n",
    "        \n",
    "        timing = dict()\n",
    "        t0 = time.time()\n",
    "        \n",
    "        response_lengths = [len(r) for r in responses]\n",
    "        \n",
    "        t = time.time()\n",
    "        logprobs, ref_logprobs, values = self.batched_forward_pass(queries, responses)\n",
    "        timing['time/ppo/forward_pass'] = time.time()-t\n",
    "\n",
    "        t = time.time()\n",
    "        rewards, non_score_reward = 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].unsqueeze(0), values[idx].unsqueeze(0),\n",
    "                                                   rewards[idx].unsqueeze(0), queries[idx].unsqueeze(0),\n",
    "                                                   responses[idx].unsqueeze(0),\n",
    "                                                   torch.cat([queries[idx],responses[idx]]).unsqueeze(0))\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/advantages'] = torch.nan_to_num(train_stats['policy/advantages'], WANDB_PADDING)\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=self.kl_ctl.value)\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, queries, responses):\n",
    "        \"\"\"Calculate model outputs in multiple batches.\"\"\"\n",
    "        bs = self.ppo_params['batch_size']\n",
    "        fbs = self.ppo_params['forward_batch_size']\n",
    "        all_logprobs = []\n",
    "        all_ref_logprobs = []\n",
    "        all_values = []\n",
    "        \n",
    "        for i in range(int(bs/fbs)):\n",
    "            query_batch = queries[i*fbs:(i+1)*fbs]\n",
    "            response_batch = responses[i*fbs:(i+1)*fbs]\n",
    "            input_ids = self.data_collator([torch.cat([q, r]) for q, r in zip(query_batch, response_batch)])[\"input_ids\"]\n",
    "            with torch.no_grad():\n",
    "                logits, _, v = self.model(input_ids)\n",
    "                ref_logits, _, _ = self.ref_model(input_ids)\n",
    "            logprobs = logprobs_from_logits(logits[:,:-1,:], input_ids[:,1:])\n",
    "            ref_logprobs = logprobs_from_logits(ref_logits[:,:-1,:], input_ids[:,1:])\n",
    "            for j in range(fbs):\n",
    "                start = len(query_batch[j])-1\n",
    "                end = len(query_batch[j]) + len(response_batch[j])-1\n",
    "                all_values.append(v[j, start-1:end-1])\n",
    "                all_logprobs.append(logprobs[j, start:end])\n",
    "                all_ref_logprobs.append(ref_logprobs[j, start:end])\n",
    "        return all_logprobs, all_ref_logprobs, all_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",
    "        rewards, non_score_rewards = [], []\n",
    "        for score, logprob, ref_logprob in zip(scores, logprobs, ref_logprobs):\n",
    "            kl = logprob - ref_logprob\n",
    "            non_score_reward = -self.kl_ctl.value * kl\n",
    "            non_score_rewards.append(non_score_reward)\n",
    "            reward = non_score_reward.clone()\n",
    "            reward[-1] += score\n",
    "            rewards.append(reward)\n",
    "        return rewards, non_score_rewards\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_list = [logprobs-ref_logprobs for logprobs, ref_logprobs in zip(data['logprobs'], data['ref_logprobs'])] \n",
    "        mean_kl = torch.mean(torch.stack([torch.sum(kl) for kl in kl_list]))\n",
    "        mean_entropy = torch.mean(torch.stack([torch.sum(-log_probs) for log_probs in data['logprobs']]))\n",
    "        mean_non_score_reward =torch.mean(torch.stack([torch.sum(non_score_reward) for non_score_reward in data['non_score_reward']]))\n",
    "        stats = {\n",
    "            'objective/kl': mean_kl,\n",
    "            'objective/kl_dist': kl_list,\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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tensor shapes and contents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Debugging tensor shapes and contents usually involves inserting a lot of print statements in the code. To avoid this in the future I add a list of the tensor shapes and contents for reference. If the tensors are sliced or reshaped I list the last shape.\n",
    "\n",
    "| Name  | Shape   | Content |\n",
    "|-------|---------|---------|\n",
    "| `query` | `[batch_size, query_length]`| contains token ids of query|\n",
    "| `response`| `[batch_size, response_length]`| contains token ids of responses|\n",
    "| `scores`| `[batch_size]`| rewards of each query/response pair|\n",
    "| `model_input`| `[batch_size, query_length + response_length]`| combined query and response tokens|\n",
    "| `m_input`|`[forward_batch_size, query_length + response_length]`| small forward batch of model_input|\n",
    "| `logits` | `[forward_batch_size, query_length + response_length, vocab_size]`| logits from model outputs|\n",
    "| `ref_logits`|`[forward_batch_size, query_length + response_length, vocab_size]`| logits from ref_model outputs|\n",
    "| `logprobs`| `[batch_size, response_length]`| log-probabilities of response tokens |\n",
    "| `ref_logprobs`| `[batch_size, response_length]`| reference log-probabilities of response tokens |\n",
    "| `rewards`| `[batch_size, response_length]`| the model rewards incl. kl-score for each token|\n",
    "| `non_score_reward`| `[batch_size, response_length]`| the model kl-score for each token|"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model output alignments\n",
    "Some notes on output alignments, since I spent a considerable time debugging this. All model outputs are shifted by 1 to the model inputs. That means that the logits are shifted by one as well as values. For this reason the logits and values are always shifted one step to the left. This also means we don't have logits for the first input element and so we delete the first input token when calculating the softmax, since we don't have logits predictions. The same applies for the values and we shift them by index one to the left."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## KL-divergence\n",
    "One question that came up during the implementation was \"Why is the KL-divergence just the difference of the log-probs? Where is the probability in front of the log term?\". The answer can be found in Sergey Levine's [lecture slides](http://rll.berkeley.edu/deeprlcourse/docs/week_3_lecture_1_dynamics_learning.pdf): To calculate the KL divergence we calculate the expected value of the log term. The probability usually in front of the log-term comes from that expected value and for a set of trajectories we can simply take the mean over the sampled trajectories."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}