{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e0f42399-b424-41a3-b179-a6f1165a92e7",
   "metadata": {},
   "source": [
    "# 7-Distill\n",
    "\n",
    "模型蒸馏 (Knowledge Distillation, KD) 是一种机器学习模型压缩方法，它用于将大型模型（教师模型）的知识迁移到较小的模型（学生模型）中.\n",
    "\n",
    "KD 背后的核心思想是将教师模型的综合知识转化为更精简、更有效的表示. 学生模型是一个较小的模型，目标是学习教师模型的行为，而不是直接从原始数据中学习.\n",
    "\n",
    "大模型的 KD 有白盒蒸馏与黑盒蒸馏两个派别，对于本次实验代码中两个模型均为 MiniMind 开源模型，支持对教师模型内部结构的访问，因此在训练过程中，我们能够获取教师模型的 softmax 概率分布并用作软标签（soft labels），让小模型学习软标签，并使用 KL-Loss 来优化模型的参数，而不是直接学习输出 Token 的硬标签. 对于下一章蒸馏推理模型中，由于我们面向推理数据集进行蒸馏，并不存在输出 Token 的概率分布让我们学习，这种面向输出数据学习的蒸馏方式被称为黑盒蒸馏.\n",
    "\n",
    "此笔记本的完整实现见主仓库 `/minimind/train_distillation.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fb823b9f-c0d3-4566-ad9c-f9ec6f95c446",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\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 SFTDataset\n",
    "\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f4546ae-9520-4f67-8530-0dfed996fa35",
   "metadata": {},
   "source": [
    "## 可选参数设置\n",
    "\n",
    "首先，查看训练的可选参数，这些参数在实际使用时通过解析命令行进行导入，我们用 class 进行包装."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c7907dbe-ca43-4904-b598-ab9e5f019982",
   "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 = './toydata/sft_data.jsonl' # 数据集路径"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "044bb66f-f878-4785-b860-a9a89b52f94c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "查看工作设备 cuda\n"
     ]
    }
   ],
   "source": [
    "print(f'查看工作设备 {args.device}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68f3ea54-b23c-4d2c-abe7-63bbe9f1e5dd",
   "metadata": {},
   "source": [
    "接下来，我们对分词器、MiniMind 教师/学生模型以及数据迭代器执行初始化."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f355d146-ee89-41c6-ba79-7aba75e3d151",
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_student_model(lm_config):\n",
    "    tokenizer = AutoTokenizer.from_pretrained('../model/minimind_tokenizer')\n",
    "    model = MiniMindLM(lm_config)\n",
    "    moe_path = '_moe' if lm_config.use_moe else ''\n",
    "    # 学生模型热启动\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",
    "    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\n",
    "\n",
    "\n",
    "def init_teacher_model(lm_config):\n",
    "    model = MiniMindLM(lm_config)\n",
    "    moe_path = '_moe' if lm_config.use_moe else ''\n",
    "    # 教师模型热启动\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",
    "    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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5592c247-4424-45e0-949a-0a56b1604f83",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "学生模型(LLM)总参数量：6.096 百万\n",
      "教师模型(LLM)总参数量：17.305 百万\n",
      "模型位于设备：cuda:0, 词表长度：6400, DataLoader：<torch.utils.data.dataloader.DataLoader object at 0x0000021E5B6E7010>\n"
     ]
    }
   ],
   "source": [
    "# 初始化模型配置，一般学生模型比较小，从大的教师模型那里学知识\n",
    "lm_config_student = LMConfig(dim=512, n_layers=1, max_seq_len=512)\n",
    "lm_config_teacher = LMConfig(dim=768, n_layers=2, max_seq_len=512)\n",
    "\n",
    "model, tokenizer = init_student_model(lm_config_student)\n",
    "teacher_model = init_teacher_model(lm_config_teacher)\n",
    "\n",
    "train_ds = SFTDataset(args.data_path, tokenizer, max_length=lm_config_student.max_seq_len)\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": "markdown",
   "id": "9af25f3d-6127-4053-92e6-533bd889a917",
   "metadata": {},
   "source": [
    "## 启动训练\n",
    "\n",
    "接下来，我们定义 MiniMind LoRA 微调所使用的优化器，损失函数和学习率调度，并进行一轮简单的训练."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ae344278-4f4c-495e-8c1b-e5c542ae7a92",
   "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",
    "# 优化学生模型参数\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": "ceadf787-ab7c-4e3e-958a-4b3cd3126a08",
   "metadata": {},
   "source": [
    "损失函数方面，使用 KL Loss 方法. \n",
    "\n",
    "KL Loss 中，损失是 KL 散度，衡量学生模型和教师模型在面对相同输入时，在输出层产生的分类 logits 分布之间的距离. 直观理解上，就是让学生模型的输出尽量向教师模型的输出概率靠近.\n",
    "\n",
    "$$D_{KL}(P||Q)=\\sum_i P(i)\\log\\frac{P(i)}{Q(i)}$$\n",
    "\n",
    "其中，$P(i)$ 代表教师模型的概率分布，$Q(i)$ 代表学生模型的预测分布."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d273971f-50ec-4319-807d-1a8e59ea9813",
   "metadata": {},
   "outputs": [],
   "source": [
    "def distillation_loss_fn(student_logits, teacher_logits, temperature=1.0, reduction='batchmean'):\n",
    "    with torch.no_grad():\n",
    "        teacher_probs = F.softmax(teacher_logits / temperature, dim=-1).detach()\n",
    "\n",
    "    student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)\n",
    "\n",
    "    kl = F.kl_div(\n",
    "        student_log_probs,\n",
    "        teacher_probs,\n",
    "        reduction=reduction  # 对各批次损失求平均值\n",
    "    )\n",
    "    return (temperature ** 2) * kl # 尺度不变"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ce6f00d-1cbd-47ab-8736-991eeb904cd9",
   "metadata": {},
   "source": [
    "接下来，我们来看训练函数."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "cd01790d-6a18-40ab-b133-da31b3ca6b66",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 作者在训练函数上的注释浅显易懂 故不作额外注释\n",
    "def train_epoch(epoch, alpha=0.0, temperature=1.0):\n",
    "    start_time = time.time()\n",
    "\n",
    "    if teacher_model is not None: # 禁用教师模型梯度\n",
    "        teacher_model.eval()\n",
    "        teacher_model.requires_grad_(False)\n",
    "\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",
    "        lr = get_lr(epoch * iter_per_epoch + step,\n",
    "                    args.epochs * iter_per_epoch,\n",
    "                    args.learning_rate)\n",
    "        for param_group in optimizer.param_groups:\n",
    "            param_group['lr'] = lr\n",
    "\n",
    "        # 前向传播（学生模型）\n",
    "        with ctx:\n",
    "            res = model(X)\n",
    "            student_logits = res.logits\n",
    "\n",
    "        # 教师模型前向传播（只在eval & no_grad）\n",
    "        if teacher_model is not None:\n",
    "            with torch.no_grad():\n",
    "                teacher_logits = teacher_model(X).logits\n",
    "                vocab_size_student = student_logits.size(-1)  # N\n",
    "                teacher_logits = teacher_logits[..., :vocab_size_student] # ... 保留除了最后一个维度外的所有维度\n",
    "\n",
    "        # ========== 计算损失 ==========\n",
    "        # 1) Ground-Truth CE Loss（可选）\n",
    "        loss_mask_flat = loss_mask.view(-1)\n",
    "        ce_loss = F.cross_entropy(\n",
    "            student_logits.view(-1, student_logits.size(-1)),\n",
    "            Y.view(-1),\n",
    "            ignore_index=0,\n",
    "            reduction='none'\n",
    "        )\n",
    "        ce_loss = torch.sum(ce_loss * loss_mask_flat) / loss_mask_flat.sum()\n",
    "        if lm_config_student.use_moe:\n",
    "            ce_loss += res.aux_loss\n",
    "\n",
    "        # 2) Distillation Loss（可选）\n",
    "        if teacher_model is not None:\n",
    "            # 只在有效token位置做蒸馏\n",
    "            distill_loss = distillation_loss_fn(\n",
    "                student_logits.view(-1, student_logits.size(-1))[loss_mask_flat == 1],\n",
    "                teacher_logits.view(-1, teacher_logits.size(-1))[loss_mask_flat == 1],\n",
    "                temperature=temperature\n",
    "            )\n",
    "        else:\n",
    "            distill_loss = torch.tensor(0.0, device=args.device)\n",
    "\n",
    "        # 3) 总损失 = alpha * CE + (1-alpha) * Distill\n",
    "        loss = alpha * ce_loss + (1 - alpha) * distill_loss\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:{:.4f} lr:{:.12f} epoch_Time:{}min:'.format(\n",
    "                    epoch,\n",
    "                    args.epochs - 1,\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",
    "            )\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_student.use_moe else ''\n",
    "        #     ckp = f'{args.save_dir}/full_dist_{lm_config_student.dim}{moe_path}.pth'\n",
    "        #     if isinstance(model, torch.nn.parallel.DistributedDataParallel):\n",
    "        #         state_dict = model.module.state_dict()\n",
    "        #     else:\n",
    "        #         state_dict = model.state_dict()\n",
    "        #     torch.save(state_dict, ckp)\n",
    "        #     model.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5369c12-e315-4a21-8f23-42dc0151909e",
   "metadata": {},
   "source": [
    "接下来，我们启动一个 Epoch 的训练进行观察."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e73d59fc-677a-4795-8377-cdef06e8e1d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:[0/0](0/1) loss:0.3257 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": 10,
   "id": "94cd3e5c-e37b-4b37-919f-7a1066b0f482",
   "metadata": {},
   "outputs": [],
   "source": [
    "del model, teacher_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3ff1681-962d-4f4b-a75d-e4673978e5d3",
   "metadata": {},
   "source": [
    "## 参考资料\n",
    "\n",
    "- [大模型知识蒸馏概述](https://zhuanlan.zhihu.com/p/659943824)\n",
    "- [使用知识蒸馏将大模型能力克隆到小模型](https://zhuanlan.zhihu.com/p/691672620)\n",
    "- [理解知识蒸馏中的散度损失函数](https://deepseek.csdn.net/67ab1c3f79aaf67875cb9664.html)"
   ]
  }
 ],
 "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
}