{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f53f31fb-4795-49cf-9a68-a13f8f7e5cae",
   "metadata": {},
   "source": [
    "# 5-DPO\n",
    "\n",
    "直接偏好优化（Direct Preference Optimization，DPO）是后训练阶段中，使用正反样例激励大模型产生符合人类偏好的回答的策略，为人类反馈强化学习（Reinforcement Learning from Human Feedback, RLHF）提供了一个高效简化的替代方案。通过这一阶段的训练，大模型将会学会依照人类的喜好生成回复.\n",
    "\n",
    "在这个笔记本中，我们仅对 DPO 的训练流程进行展示和学习，因此只给出必要的代码片段，如 wandb 和 ddp 不会在此笔记本中涉及.\n",
    "\n",
    "此笔记本的完整实现见主仓库 `/minimind/train_dpo.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3e09f234-5b8b-4bfb-8d42-bb766a457acc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入依赖\n",
    "import os\n",
    "import platform\n",
    "import argparse\n",
    "import time\n",
    "import math\n",
    "import warnings\n",
    "\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.distributed as dist\n",
    "from contextlib import nullcontext\n",
    "\n",
    "from torch import optim, nn\n",
    "from torch.nn.parallel import DistributedDataParallel\n",
    "from torch.utils.data import DataLoader, DistributedSampler\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "from model.model import MiniMindLM\n",
    "from model.LMConfig import LMConfig\n",
    "from model.dataset import DPODataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "eb946301-eb7e-45c6-9e1c-16cc0ee4db82",
   "metadata": {},
   "outputs": [],
   "source": [
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d70c65d3-14f7-4569-bebd-cddb48f1851f",
   "metadata": {},
   "source": [
    "## 可选参数设置\n",
    "\n",
    "首先，查看训练的可选参数，这些参数在实际使用时通过命令行导入，为了保持笔记本的易用性，选择用 class 进行包装."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b78b11a1-9189-4eaf-a094-07c0497c4d38",
   "metadata": {},
   "outputs": [],
   "source": [
    "class args:\n",
    "    # out_dir: str = \"out\" # pytorch 格式权重文件保存位置 我们只展示训练过程 所以不使用\n",
    "    epochs: int = 1 # 训练轮数\n",
    "    batch_size: int = 2 # pretrain 数据集仅两个样本，设置 batch 为 2\n",
    "    # sft阶段学习率为 「5e-6」->「5e-7」长度512，建议离线正负样本「概率」偏好对齐阶段lr <=「1e-8」长度3000，否则很容易遗忘训坏\n",
    "    learning_rate: float = 5e-4 # 学习率\n",
    "    device: str = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    dtype: str = 'bfloat16' # 16 bit 浮点数：8 bit 指数 + 7 bit 尾数\n",
    "    # use_wandb: bool = False # 是否使用 wandb 我们不使用\n",
    "    wandb_project: str = 'MiniMind-Notebook'\n",
    "    num_workers: int = 1 # 工作进程数\n",
    "    # ddp：bool = False # 单机多卡\n",
    "    accumulation_steps: int = 1 # 梯度累积步数\n",
    "    grad_clip: float = 1.0 # 梯度剪裁\n",
    "    warmup_iters: int = 0 # 学习率热启动\n",
    "    log_interval: int = 1 # 每一步打印日志 仅用于观察\n",
    "    # save_interval: int = 100 # checkpoint 保存点 我们不使用\n",
    "    local_rank: int = 1 # device 设备号\n",
    "    dim: int = 512 # 词嵌入维度 模型超参数\n",
    "    n_layers: int = 1 # MiniMind Block 数量 模型超参数 | 由于 dpo 要加载两个模型 我们出于演示目的设定 n_layers = 1\n",
    "    max_seq_len: int = 512 # 序列长度阈值\n",
    "    use_moe: bool = False # 是否启用混合专家\n",
    "    data_path: str = './toydata/dpo_data.jsonl' # 数据集路径"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0a72c8aa-ca60-4b0d-aabb-2c9d16b48175",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "查看工作设备 cuda\n"
     ]
    }
   ],
   "source": [
    "print(f'查看工作设备 {args.device}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8fb31ff-524b-4788-a25f-699edb887704",
   "metadata": {},
   "source": [
    "## 初始化训练\n",
    "\n",
    "接下来，我们对一些重要模块进行初始化，我们已经了解过，分词器，模型和数据集是大模型的基本组件，我们对其进行初始化.\n",
    "\n",
    "> 在这一阶段 我们调整的是大模型的问答偏好 因此与 sft 阶段同理 我们需要载入在 sft 阶段微调好的问答模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5d5971fc-4794-4903-b418-76f266bf19d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_model(lm_config):\n",
    "    \"\"\"注意 在此处我们注释了加载模型部分代码 但应当认识到 actor 和 ref 两个模型都具备 sft 阶段获得的权重\"\"\"\n",
    "    tokenizer = AutoTokenizer.from_pretrained('../model/minimind_tokenizer')\n",
    "    model = MiniMindLM(lm_config)\n",
    "    moe_path = '_moe' if lm_config.use_moe else ''\n",
    "    # ckp = f'./out/full_sft_{lm_config.dim}{moe_path}.pth' # 指示上一阶段训练保存的模型文件位置\n",
    "    # state_dict = torch.load(ckp, map_location=args.device) # 载入模型状态字典\n",
    "    # model.load_state_dict(state_dict, strict=False) # 装入模型\n",
    "    # 初始化参考模型\n",
    "    ref_model = MiniMindLM(lm_config)\n",
    "    # ref_model.load_state_dict(state_dict, strict=False)\n",
    "    ref_model.eval()\n",
    "    ref_model.requires_grad_(False)\n",
    "\n",
    "    print(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')\n",
    "    model = model.to(args.device)\n",
    "    ref_model = ref_model.to(args.device)\n",
    "\n",
    "    return model, ref_model, tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "fbd7bd15-a060-4aba-8c5d-277705ec6b52",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LLM总参数量：6.096 百万\n",
      "模型位于设备：cuda:0, 词表长度：6400, DataLoader：<torch.utils.data.dataloader.DataLoader object at 0x000001AF0BBB71C0>\n"
     ]
    }
   ],
   "source": [
    "lm_config = LMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, use_moe=args.use_moe)\n",
    "model, ref_model ,tokenizer = init_model(lm_config)\n",
    "\n",
    "train_ds = DPODataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)\n",
    "\n",
    "train_loader = DataLoader(\n",
    "    train_ds,\n",
    "    batch_size=args.batch_size,\n",
    "    pin_memory=True,\n",
    "    drop_last=False,\n",
    "    shuffle=False,\n",
    "    num_workers=args.num_workers,\n",
    ")\n",
    "\n",
    "print(f'模型位于设备：{model.device}, 词表长度：{tokenizer.vocab_size}, DataLoader：{train_loader}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d807a2af-41e9-4115-b470-a05f6a04613c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "打印一个 iter 的数据:\n",
      "{'x_chosen': tensor([[  1,  85, 736,  ...,   0,   0,   0],\n",
      "        [  1,  85, 736,  ...,   0,   0,   0]]), 'y_chosen': tensor([[ 85, 736, 201,  ...,   0,   0,   0],\n",
      "        [ 85, 736, 201,  ...,   0,   0,   0]]), 'mask_chosen': tensor([[0, 0, 0,  ..., 0, 0, 0],\n",
      "        [0, 0, 0,  ..., 0, 0, 0]]), 'x_rejected': tensor([[  1,  85, 736,  ...,   0,   0,   0],\n",
      "        [  1,  85, 736,  ...,   0,   0,   0]]), 'y_rejected': tensor([[ 85, 736, 201,  ...,   0,   0,   0],\n",
      "        [ 85, 736, 201,  ...,   0,   0,   0]]), 'mask_rejected': tensor([[0, 0, 0,  ..., 0, 0, 0],\n",
      "        [0, 0, 0,  ..., 0, 0, 0]])}\n",
      "\n",
      "数据集大小：2, DataLoader 大小：1\n"
     ]
    }
   ],
   "source": [
    "loader = iter(train_loader)\n",
    "print(f'打印一个 iter 的数据:\\n{next(loader)}\\n')\n",
    "print(f'数据集大小：{len(train_ds)}, DataLoader 大小：{len(loader)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "665ad506-ab12-4186-ad63-ada143a33d52",
   "metadata": {},
   "source": [
    "我们发现，train loader 的每一个 iter 都包含一个拥有六个键值对的字典，这是因为 train_dataset 每一次取数据都会返回:\n",
    "\n",
    "- chosen 样本 X: 包含 \\<bos> 在内的输入 content\n",
    "- chosen 标签 Y: 包含 \\<eos> 在内的输出 content\n",
    "- chosen 掩码 loss_mask: 指示需要计算损失的 token 位置\n",
    "- rejected 样本 X: 包含 \\<bos> 在内的输入 content\n",
    "- rejected 标签 Y: 包含 \\<eos> 在内的输出 content\n",
    "- rejected 掩码 loss_mask: 指示需要计算损失的 token 位置\n",
    "\n",
    "由于我们的数据集只有两条数据，而 batch size 设置为 2，因此我们的 dataloader 只有一个 iter."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32462eb9-7008-4389-a54b-47ef1c1c4732",
   "metadata": {},
   "source": [
    "# 启动训练\n",
    "\n",
    "训练一个深度学习模型，还涉及到了优化器，损失函数和学习率调度. 接下来，我们查看 MiniMind 训练部分的代码，并进行一轮简单的训练.\n",
    "\n",
    "> DPO 阶段涉及 DPO 损失函数涉及 因此与前两个阶段相比内容略有增加 不过整体流程与逻辑类似"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1752f117-33c0-4810-886f-fbd1b4d44c28",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 学习率调度方面 采用余弦退火学习率\n",
    "def get_lr(current_step, total_steps, lr):\n",
    "    return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))\n",
    "\n",
    "# 优化器方面 选择 AdamW 优化器 并在混精度场景下创建 scaler 进行梯度缩放避免数值下溢\n",
    "scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))\n",
    "optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)\n",
    "\n",
    "device_type = \"cuda\" if \"cuda\" in args.device else \"cpu\"\n",
    "ctx = nullcontext() if device_type == \"cpu\" else torch.cuda.amp.autocast() # 在 cuda 上启动混精度训练，否则空白上下文"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f2c3ac1-92bf-4eae-8d22-7fac4cfa69df",
   "metadata": {},
   "source": [
    "DPO 的原理是增加偏好样本的对数概率与减小非偏好样本响应的对数概率.\n",
    "\n",
    "该阶段引入 DPO 损失函数，通过计算选择样本和拒绝样本的对数比率，然后基于这些比率计算 DPO 损失，适用于偏好学习任务."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2c07f54a-496e-4b1d-9394-dd2c266ab5d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def logits_to_probs(logits, labels):\n",
    "    # logits shape: (batch_size, seq_len, vocab_size)\n",
    "    # labels shape: (batch_size, seq_len)\n",
    "    # probs shape: (batch_size, seq_len)\n",
    "    log_probs = F.log_softmax(logits, dim=2)\n",
    "    probs = torch.gather(log_probs, dim=2, index=labels.unsqueeze(2)).squeeze(-1)\n",
    "    return probs\n",
    "\n",
    "\n",
    "def dpo_loss(ref_probs, probs, beta):\n",
    "    # ref_probs 和 probs 都是 shape: (batch_size, seq_len)\n",
    "    # 计算每个样本的平均概率\n",
    "    ref_probs = ref_probs.mean(dim=1)\n",
    "    probs = probs.mean(dim=1)\n",
    "\n",
    "    # 将 chosen 和 rejected 数据分开\n",
    "    batch_size = ref_probs.shape[0]\n",
    "    chosen_ref_probs = ref_probs[:batch_size // 2]\n",
    "    reject_ref_probs = ref_probs[batch_size // 2:]\n",
    "    chosen_probs = probs[:batch_size // 2]\n",
    "    reject_probs = probs[batch_size // 2:]\n",
    "\n",
    "    # 计算对数比率，比较偏好差异\n",
    "    pi_logratios = chosen_probs - reject_probs\n",
    "    ref_logratios = chosen_ref_probs - reject_ref_probs\n",
    "    logits = pi_logratios - ref_logratios\n",
    "    loss = -F.logsigmoid(beta * logits)\n",
    "    return loss.mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e743562c-6e34-488d-bd84-f62eaee0ae20",
   "metadata": {},
   "source": [
    "接下来，我们来看看 MiniMind 的训练函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "493b7368-d82f-48a6-b5bd-7737b7a23bee",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_epoch(epoch):\n",
    "    start_time = time.time()\n",
    "    for step, batch in enumerate(train_loader):\n",
    "        # 提取数据\n",
    "        x_chosen = batch['x_chosen'].to(args.device)\n",
    "        x_rejected = batch['x_rejected'].to(args.device)\n",
    "        y_chosen = batch['y_chosen'].to(args.device)\n",
    "        y_rejected = batch['y_rejected'].to(args.device)\n",
    "        mask_chosen = batch['mask_chosen'].to(args.device)\n",
    "        mask_rejected = batch['mask_rejected'].to(args.device)\n",
    "        # 正反例拼接\n",
    "        x = torch.cat([x_chosen, x_rejected], dim=0)\n",
    "        y = torch.cat([y_chosen, y_rejected], dim=0)\n",
    "        mask = torch.cat([mask_chosen, mask_rejected], dim=0)\n",
    "\n",
    "        lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate)\n",
    "        for param_group in optimizer.param_groups:\n",
    "            param_group['lr'] = lr\n",
    "\n",
    "        with ctx:\n",
    "            with torch.no_grad(): # 计算 ref 模型输出\n",
    "                ref_outputs = ref_model(x)\n",
    "                ref_logits = ref_outputs.logits\n",
    "            ref_probs = logits_to_probs(ref_logits, y)\n",
    "            ref_probs = ref_probs * mask # 得到 ref 概率\n",
    "            outputs = model(x) # 计算 actor 模型输出\n",
    "            logits = outputs.logits\n",
    "            probs = logits_to_probs(logits, y)\n",
    "            probs = probs * mask # 得到 actor 概率\n",
    "            loss = dpo_loss(ref_probs, probs, beta=0.1) # dpo 损失\n",
    "            loss = loss / args.accumulation_steps\n",
    "\n",
    "        scaler.scale(loss).backward() # 梯度缩放\n",
    "\n",
    "        if (step + 1) % args.accumulation_steps == 0:\n",
    "            scaler.unscale_(optimizer)\n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) # 梯度剪裁\n",
    "            scaler.step(optimizer)\n",
    "            scaler.update()\n",
    "            optimizer.zero_grad(set_to_none=True)\n",
    "\n",
    "        if step % args.log_interval == 0:\n",
    "            spend_time = time.time() - start_time\n",
    "            print(\n",
    "                'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} epoch_Time:{}min:'.format(\n",
    "                    epoch + 1,\n",
    "                    args.epochs,\n",
    "                    step,\n",
    "                    iter_per_epoch,\n",
    "                    loss.item(),\n",
    "                    optimizer.param_groups[-1]['lr'],\n",
    "                    spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))\n",
    "\n",
    "        # 到达指定保存步数时，save as PyTorch\n",
    "        # if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):\n",
    "        #     model.eval()\n",
    "        #     moe_path = '_moe' if lm_config.use_moe else ''\n",
    "        #     ckp = f'{args.save_dir}/rlhf_{lm_config.dim}{moe_path}.pth'\n",
    "\n",
    "        #     if isinstance(model, torch.nn.parallel.DistributedDataParallel):\n",
    "        #         state_dict = model.module.state_dict()\n",
    "        #     else:\n",
    "        #         state_dict = model.state_dict()\n",
    "\n",
    "        #     torch.save(state_dict, ckp)\n",
    "        #     model.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49386dc7-2919-428a-a073-f37d576afbcf",
   "metadata": {},
   "source": [
    "准备完毕，我们尝试一轮长度 1 个 iter 的训练."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0ac86572-ee04-43b5-b5fa-b1c65cbaebec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:[1/1](0/1) loss:0.694 lr:0.000550000000 epoch_Time:0.0min:\n"
     ]
    }
   ],
   "source": [
    "iter_per_epoch = len(train_loader)\n",
    "for epoch in range(args.epochs):\n",
    "    train_epoch(epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "5215ef67-2647-4519-966b-074c46c225ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "del model"
   ]
  }
 ],
 "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}