431 lines
16 KiB
Plaintext
431 lines
16 KiB
Plaintext
{
|
||
"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
|
||
}
|