{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "70560a4d-ff57-4699-afb2-52f6a145b97e",
   "metadata": {},
   "source": [
    "# 4-SFT\n",
    "\n",
    "在预训练阶段后，我们应该能够获得一个下一词预测模型，此时的模型已经掌握了大量的知识。不过，仅仅具备下一词预测能力是不够的，我们希望大模型能够获得问答能力，这一能力便是在有监督微调（Supervised Fine Tuning，SFT）阶段获得的.\n",
    "\n",
    "在这个笔记本中，我们仅对 SFT 的训练流程进行展示和学习，因此只给出必要的代码片段，如 wandb 和 ddp 不会在此笔记本中涉及.\n",
    "\n",
    "此笔记本的完整实现见主仓库 `/minimind/train_full_sft.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2508c8c7-eb44-4ac0-918c-41479f41d1ac",
   "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 demo.model import MiniMindLM\n",
    "from demo.LMConfig import LMConfig\n",
    "from demo.dataset import SFTDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4a79e140-23c7-4498-9a4d-ef7459692996",
   "metadata": {},
   "outputs": [],
   "source": [
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbe1d40e-e074-4856-8c67-da8019f250ab",
   "metadata": {},
   "source": [
    "## 可选参数设置\n",
    "\n",
    "首先，查看训练的可选参数，这些参数在实际使用时通过命令行导入，为了保持笔记本的易用性，选择用 class 进行包装."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "98cbd6cf-081f-4c31-8dce-52566004801a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class args:\n",
    "    # out_dir: str = \"out\" # pytorch 格式权重文件保存位置 我们只展示训练过程 所以不使用\n",
    "    epochs: int = 1 # 训练轮数\n",
    "    batch_size: int = 2 # pretrain 数据集仅两个样本，设置 batch 为 2\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 = 2 # MiniMind Block 数量 模型超参数\n",
    "    max_seq_len: int = 512 # 序列长度阈值\n",
    "    use_moe: bool = False # 是否启用混合专家\n",
    "    data_path: str = './demo/sft_data.jsonl' # 数据集路径"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "642f5107-aa9b-49cc-9759-e8731326edc3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "查看工作设备 cuda\n"
     ]
    }
   ],
   "source": [
    "print(f'查看工作设备 {args.device}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b70fadb6-f78f-45a0-ad3e-d41d7253c72b",
   "metadata": {},
   "source": [
    "## 初始化训练\n",
    "\n",
    "接下来，我们对一些重要模块进行初始化，我们已经了解过，分词器，模型和数据集是大模型的基本组件，我们对其进行初始化.\n",
    "\n",
    "> 注意 与预训练阶段不同的是 在 sft 阶段 我们实际上是在上一阶段训练获得的模型的基础上修改数据集进行接续训练 因此需要载入上一阶段的模型权重 出于展示的目的 载入权重的代码在此笔记本中只作展示 并不执行"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c64a0e37-a9d0-4ce6-b5ee-7b2728de879b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_model(lm_config):\n",
    "    tokenizer = AutoTokenizer.from_pretrained('../model/minimind_tokenizer')\n",
    "    model = MiniMindLM(lm_config).to(args.device)\n",
    "    moe_path = '_moe' if lm_config.use_moe else ''\n",
    "    # ckp = f'./out/pretrain_{lm_config.dim}{moe_path}.pth' # 指示上一阶段训练保存的模型文件位置\n",
    "    # state_dict = torch.load(ckp, map_location=args.device) # 载入模型状态字典\n",
    "    # model.laod_state_dict(state_dict, strict=False) # 装入模型\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",
    "    return model, tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2eec33f3-2ac2-4c75-b58e-844d89d8add9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LLM总参数量：8.915 百万\n",
      "模型位于设备：cuda:0, 词表长度：6400, DataLoader：<torch.utils.data.dataloader.DataLoader object at 0x000002D0DD01D600>\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, tokenizer = init_model(lm_config)\n",
    "\n",
    "train_ds = SFTDataset(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": "b02644e0-e61d-4df0-9203-fa841cd83635",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "打印一个 iter 的数据:\n",
      "[tensor([[  1,  85, 736,  ...,   0,   0,   0],\n",
      "        [  1,  85, 736,  ...,   0,   0,   0]]), tensor([[ 85, 736, 201,  ...,   0,   0,   0],\n",
      "        [ 85, 736, 201,  ...,   0,   0,   0]]), 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": "13ac8a47-d26a-49da-82fb-468ef5ac91a4",
   "metadata": {},
   "source": [
    "我们发现，train loader 的每一个 iter 都包含一个长度为 3 的张量列表，这是因为 train_dataset 每一次取数据都会返回三个张量，分别为:\n",
    "\n",
    "- 样本 X: 包含 \\<bos> 在内的输入 conversation\n",
    "- 标签 Y: 包含 \\<eos> 在内的输出 conversation\n",
    "- 掩码 loss_mask: 指示需要计算损失的 token 位置\n",
    "\n",
    "由于我们的数据集只有两条数据，而 batch size 设置为 2，因此我们的 dataloader 只有一个 iter."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee23f937-f333-4c02-a3b4-a768818c0e2b",
   "metadata": {},
   "source": [
    "## 启动训练\n",
    "\n",
    "训练一个深度学习模型，还涉及到了优化器，损失函数和学习率调度. 接下来，我们查看 MiniMind 训练部分的代码，并进行一轮简单的训练.\n",
    "\n",
    "> 不难发现 pretrain 阶段和 sft 阶段的训练主体差不多 因为这两个阶段的差异体现在数据集格式 而数据集在经过 chat template 格式化后差异小了很多"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5c32a59a-be0b-4426-b47e-c522322bef69",
   "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": "13963ec0-d22e-4120-b0f2-8f5e5abfcda1",
   "metadata": {},
   "source": [
    "接下来，我们来看看 MiniMind 的训练函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "74499592-2065-4fff-97d1-762cd08b77b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_epoch(epoch):\n",
    "    loss_fct = nn.CrossEntropyLoss(reduction='none') # 损失函数 采用交叉熵损失\n",
    "    start_time = time.time()\n",
    "    for step, (X, Y, loss_mask) in enumerate(train_loader):\n",
    "        X = X.to(args.device)\n",
    "        Y = Y.to(args.device)\n",
    "        loss_mask = loss_mask.to(args.device)\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",
    "            res = model(X) # 前向推理\n",
    "            loss = loss_fct(\n",
    "                res.logits.view(-1, res.logits.size(-1)),\n",
    "                Y.view(-1)\n",
    "            ).view(Y.size()) # 取生成的最后一个 token 的 logits 计算损失\n",
    "            loss = (loss * loss_mask).sum() / loss_mask.sum()\n",
    "            loss += res.aux_loss # 若为混合专家 则将 MOE 辅助损失纳入考虑\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",
    "\n",
    "            scaler.step(optimizer)\n",
    "            scaler.update()\n",
    "\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() * args.accumulation_steps,\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}/pretrain_{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": "546f28c2-418e-453c-88aa-695468f9216c",
   "metadata": {},
   "source": [
    "准备完毕，我们尝试一轮长度 1 个 iter 的训练."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "341200c9-0d3a-4ed1-9dd2-d2fc35d02564",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:[1/1](0/1) loss:8.869 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": 11,
   "id": "b0480212-1500-49fb-b679-7bdf2c100c87",
   "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
}