GRPO-Loss-Pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ba303ea0-7339-444f-a85a-4389873b5bf5",
   "metadata": {},
   "source": [
    "# GRPO-Loss-pytorch\n",
    "\n",
    "author: xiaodongguaAIGC\n",
    "\n",
    "git: [dhcode-cpp](https://github.com/dhcode-cpp)\n",
    "\n",
    "blog: [【手撕LLM-GRPO】你只管给Reward, 剩下的交给RL（附代码）](https://zhuanlan.zhihu.com/p/20812786520)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "37d15458-8a27-4baf-afc7-6b6187013547",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b0531f35-3559-4c4f-9225-4fdc20a647e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def grpo_kl(pi_logprob, pi_ref_logprob):\n",
    "    return pi_ref_logprob.exp() / pi_logprob.exp()- (pi_ref_logprob - pi_logprob) - 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "74242e19-4f3c-4a58-9413-cc3d6f1863db",
   "metadata": {},
   "outputs": [],
   "source": [
    "def grpo_loss(pi_logprob, pi_old_logprob, pi_ref_logprob, advantage, input_len, len_oi):\n",
    "    epsilon = 0.2\n",
    "    beta = 0.01\n",
    "\n",
    "    bs, seq_len = pi_logprob.shape\n",
    "    # skip计算采样的每条采样长度\n",
    "    len_oi = torch.tensor([len_oi] * bs, dtype = torch.long)\n",
    "    # 设定mask, 仅对response 为 1， 算loss\n",
    "    mask = torch.zeros(bs, seq_len)\n",
    "    mask[:, input_len:] = 1\n",
    "\n",
    "    # GRPO loss\n",
    "    ratio = torch.exp(pi_logprob - pi_old_logprob)\n",
    "    ratio_clip = torch.clamp(ratio, 1 - epsilon, 1 + epsilon)\n",
    "    advantage = advantage.unsqueeze(dim = 1) # [a, b ,c] -> [[a], [b], [c]]\n",
    "    policy_gradient = torch.minimum(ratio * advantage , ratio_clip * advantage)\n",
    "    kl = grpo_kl(pi_logprob, pi_ref_logprob)\n",
    "\n",
    "    loss = (policy_gradient -  beta * kl) * mask\n",
    "    loss = (-1 / bs ) * (1/len_oi.unsqueeze(dim = 1)) * loss  \n",
    "    loss = loss.sum()\n",
    "\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "50164e56-71bf-4330-80e7-1e9606c47b18",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.4713</span><span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m-0.4713\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 输出分布\n",
    "pi_logits = torch.randn(3, 5, 32) # batch, seq_len, vocab_size\n",
    "pi_ref_logits = torch.randn(3, 5, 32)\n",
    "pi_old_logits = torch.randn(3, 5, 32)\n",
    "\n",
    "# 获取log prob\n",
    "pi_logprob = F.log_softmax(pi_logits, dim = -1)\n",
    "pi_ref_logprob = F.log_softmax(pi_ref_logits, dim = -1)\n",
    "pi_old_logprob = F.log_softmax(pi_old_logits, dim = -1)\n",
    "\n",
    "# group data\n",
    "token_ids = torch.tensor([[11, 12, 13, 14, 15], # 输入为11,12,13, 输出为:14, 15\n",
    "                          [11, 12, 13, 15, 16],\n",
    "                          [11, 12, 13, 16, 17],])\n",
    "\n",
    "# 获取policy\n",
    "pi_logprob = torch.gather(pi_logprob, dim=-1, index=token_ids.unsqueeze(-1)).squeeze(-1)\n",
    "pi_ref_logprob = torch.gather(pi_ref_logprob, dim=-1, index=token_ids.unsqueeze(-1)).squeeze(-1)\n",
    "pi_old_logprob = torch.gather(pi_old_logprob, dim=-1, index=token_ids.unsqueeze(-1)).squeeze(-1)\n",
    "loss = grpo_loss(pi_logprob, pi_old_logprob, pi_ref_logprob, torch.tensor([-1, 2, 1]), 3, 2)\n",
    "print(loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf925213-e779-4155-b85d-fff71b4ec661",
   "metadata": {},
   "source": [
    "## Trl Implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c4b0d89-6626-4e0e-99a2-57ce1f0f2544",
   "metadata": {},
   "source": [
    "- ppo clip ratio\n",
    "- grpo clip ratio\n",
    "- trl \"not\" clip ratio, it haven't minibatch, ` exp( logprob - logprob.detach()` always equal `1`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "60b46980-106c-4fb0-9ca7-5f0fe18156ef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>.<span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m1\u001b[0m.\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>.<span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m1\u001b[0m.\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "policy = torch.tensor([0.5])\n",
    "old_policy = torch.tensor([0.5])\n",
    "ratio = policy/old_policy\n",
    "print(ratio)\n",
    "\n",
    "ratio = torch.exp( policy.log() - old_policy.log())\n",
    "print(ratio)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1109e320-5872-4ba4-af99-05dee46d4a93",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.4000</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m0.4000\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "gradient = -0.2\n",
    "policy_gradient = - gradient * ( 1 / old_policy)\n",
    "print(policy_gradient)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}