360 lines
13 KiB
Plaintext
360 lines
13 KiB
Plaintext
{
|
||
"cells": [
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "2edfeb4e-03ea-406a-a345-6fe19d509c37",
|
||
"metadata": {},
|
||
"source": [
|
||
"# 8-Reason\n",
|
||
"\n",
|
||
"参数太小的模型直接通过冷启动SFT+GRPO几乎不可能获得任何推理效果，因此，使用冷启动 SFT + GRPO 训练方法对小模型推理能力的作用有限.因此，MiniMind 项目作者使用推理数据集对 MiniMind 系列模型进行黑盒蒸馏来训练推理模型.\n",
|
||
"\n",
|
||
"使用的推理数据格式:\n",
|
||
"\n",
|
||
"```\n",
|
||
"{\n",
|
||
" \"conversations\": [\n",
|
||
" {\"role\": \"user\", \"content\": \"Q1?\"},\n",
|
||
" {\"role\": \"assistant\", \"content\": \"<think>T1</think>\\n<answer>A1</answer>\"},\n",
|
||
" {\"role\": \"user\", \"content\": \"Q2?\"},\n",
|
||
" {\"role\": \"assistant\", \"content\": \"<think>T2</think>\\n<answer>A2</answer>\"}\n",
|
||
" ]\n",
|
||
"}\n",
|
||
"```\n",
|
||
"\n",
|
||
"此笔记本的完整实现见主仓库 `/minimind/train_distill_reason.py`"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 1,
|
||
"id": "f7eb789f-18d9-45c5-acab-ca3167950911",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"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\n",
|
||
"\n",
|
||
"warnings.filterwarnings('ignore')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "24cefe29-fdf0-44c7-86fa-97206bc18161",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 可选参数设置\n",
|
||
"\n",
|
||
"首先，查看训练的可选参数，这些参数在实际使用时通过解析命令行进行导入，我们用 class 进行包装."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 2,
|
||
"id": "22643ad2-fc7a-4efd-ac53-81c2bcc88cd1",
|
||
"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/r1_data.jsonl' # 数据集路径"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 3,
|
||
"id": "d2dc3bf2-215f-4fc0-b7ce-6b4a09b0f2d9",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"查看工作设备 cuda\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"print(f'查看工作设备 {args.device}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "82a5bce8-a578-4e56-9dd5-9cb43321da0c",
|
||
"metadata": {},
|
||
"source": [
|
||
"接下来，我们对分词器、MiniMind 学生模型以及数据迭代器执行初始化."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 4,
|
||
"id": "db202500-e265-434e-b0a4-c352854de648",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"def init_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/rlhf_{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"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 5,
|
||
"id": "e1948aa4-381d-485a-9b19-e0e566e0b5da",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"LLM总参数量：8.915 百万\n",
|
||
"模型位于设备：cuda:0, 词表长度：6400, DataLoader：<torch.utils.data.dataloader.DataLoader object at 0x0000021EAB98B940>\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",
|
||
")\n",
|
||
"\n",
|
||
"print(f'模型位于设备：{model.device}, 词表长度：{tokenizer.vocab_size}, DataLoader：{train_loader}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "9be67e55-0d1b-4c5d-9304-3b65d5b81c4d",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 启动训练\n",
|
||
"\n",
|
||
"接下来，我们定义 MiniMind LoRA 微调所使用的优化器，损失函数和学习率调度，并进行一轮简单的训练."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 6,
|
||
"id": "59084328-05e4-49d3-8459-54c549131511",
|
||
"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": "7a241f42-32ec-4070-9fda-1cb39860089f",
|
||
"metadata": {},
|
||
"source": [
|
||
"接下来，我们来看看训练函数.\n",
|
||
"\n",
|
||
"蒸馏思考数据集的训练过程与 SFT 类似，区别在于模型生成序列中，思考标签位置的预测错误惩罚被放大."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 7,
|
||
"id": "e705c796-938e-4e0e-b813-3d6ffc68635e",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"def train_epoch(epoch):\n",
|
||
" # 思考标签占位符\n",
|
||
" start_of_think_ids = tokenizer('<think>').input_ids\n",
|
||
" end_of_think_ids = tokenizer('</think>').input_ids\n",
|
||
" start_of_answer_ids = tokenizer('<answer>').input_ids\n",
|
||
" end_of_answer_ids = tokenizer('</answer>').input_ids\n",
|
||
" loss_fct = nn.CrossEntropyLoss(reduction='none') # ce 损失\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",
|
||
" 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())\n",
|
||
" # 判断思考标签是否存在于 Y 中，并返回一个 bool 类型张量，指示思考标签在 Y 中的位置\n",
|
||
" sp_ids = torch.isin(Y.view(-1),\n",
|
||
" torch.tensor(start_of_think_ids + end_of_think_ids\n",
|
||
" + start_of_answer_ids + end_of_answer_ids\n",
|
||
" ).to(args.device))\n",
|
||
" # 在 sp_ids 对应的位置增加额外的惩罚\n",
|
||
" loss_mask = loss_mask.view(-1)\n",
|
||
" loss_mask_sum = loss_mask.sum()\n",
|
||
" loss_mask[sp_ids] = 10\n",
|
||
" loss_mask = loss_mask.view(Y.size())\n",
|
||
" loss = (loss * loss_mask).sum() / loss_mask_sum # 思考标签相应位置的惩罚被放大十倍\n",
|
||
" loss += res.aux_loss # 负载均衡损失\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(),\n",
|
||
" optimizer.param_groups[-1]['lr'],\n",
|
||
" spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))\n",
|
||
"\n",
|
||
" # 到达指定步数后 save as torch\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}/reason_{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": "a3b5aaf7-feac-42d9-adfe-9a583f26879d",
|
||
"metadata": {},
|
||
"source": [
|
||
"接下来，我们启动一个 Epoch 的训练进行观察."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 8,
|
||
"id": "f021611c-6291-4c50-acf0-9a24dde87e77",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"Epoch:[1/1](0/1) loss:11.907 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": 9,
|
||
"id": "fba06e90-680e-47ff-99e3-49fd91d74dc9",
|
||
"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
|
||
}
|