update interactive notebook

pull/214/head
mingkun 2025-02-25 23:00:46 +08:00
parent 9f405cafe3
commit 534d4e2c1a
12 changed files with 2583 additions and 507 deletions

View File

@ -10,9 +10,7 @@
"source": [
"# 0-Tokenizer\n",
"\n",
"对一段文本进行分词就是将其分割成单词或子词,然后通过查找表将这些单词或子词转换为对应的编号。\n",
"\n",
"分词器的就是把一段文本分割成单词或子词的工具。 "
"Tokenizer(分词器)在 NLP 领域扮演着基础且关键的作用,它将文本分割成单词或子词并转化为数组编号,为模型提供可处理的输入,在文本预处理、语义理解及适配不同语言和任务等方面奠定基础,是连接自然语言文本与计算机可处理数据的重要桥梁."
]
},
{
@ -33,19 +31,19 @@
"\n",
"### BPE\n",
"\n",
"BPE 是一种简单的数据压缩技术,它会迭代地替换序列中最频繁出现的字节对。BPE 依赖一个预分词器,该预分词器会将训练数据分割成单词(在本项目中,我们使用按空格分词的方法作为预分词方法)。\n",
"BPE 是一种简单的数据压缩技术,它会迭代地替换序列中最频繁出现的字节对。BPE 依赖一个预分词器,该预分词器会将训练数据分割成单词(在本项目中,我们使用按空格分词的方法作为预分词方法).\n",
"\n",
"在预分词之后,会创建一组唯一的单词,并确定它们在数据中的出现频率。接下来,BPE 会创建一个基础词表,该词表包含预分词器最初生成的数据中所有唯一单词的符号。然后,会将这对符号从词表中移除,新形成的符号将加入词表。在迭代过程中,BPE 算法会合并频繁出现的符号对。\n",
"在预分词之后,会创建一组唯一的单词,并确定它们在数据中的出现频率。接下来,BPE 会创建一个基础词表,该词表包含预分词器最初生成的数据中所有唯一单词的符号。然后,会将这对符号从词表中移除,新形成的符号将加入词表。在迭代过程中,BPE 算法会合并频繁出现的符号对.\n",
"\n",
"给定词表的大小,BPE(字节对编码)算法最终会合并出现频率最高的符号对,直到收敛到该大小。\n",
"给定词表的大小,BPE(字节对编码)算法最终会合并出现频率最高的符号对,直到收敛到该大小.\n",
"\n",
"### WordPiece\n",
"\n",
"WordPiece 算法与 BPE 非常相似。WordPiece 首先将词表初始化为包含训练数据中出现的每个字符,然后逐步学习给定数量的合并规则。与 BPE 不同的是,WordPiece 并不选择最频繁出现的符号对,而是选择那个加入词表后能使训练数据出现的可能性最大化的符号对。\n",
"WordPiece 算法与 BPE 非常相似。WordPiece 首先将词表初始化为包含训练数据中出现的每个字符,然后逐步学习给定数量的合并规则. 与 BPE 不同的是,WordPiece 并不选择最频繁出现的符号对,而是选择那个加入词表后能使训练数据出现的可能性最大化的符号对.\n",
"\n",
"### Unigram\n",
"\n",
"Unigram 算法将其基础词表初始化为大量的符号,然后逐步削减每个符号,以获得一个更小的词表。它会在训练数据上定义一个对数似然损失,以此来确定是否从词表中移除某个符号。"
"Unigram 算法将其基础词表初始化为大量的符号,然后逐步削减每个符号,以获得一个更小的词表。它会在训练数据上定义一个对数似然损失,以此来确定是否从词表中移除某个符号."
]
},
{
@ -54,7 +52,7 @@
"source": [
"## 训练一个最简单的分词器\n",
"\n",
"在本节中,我们将学习基于 transformers 库来训练你自己的分词器。"
"在本节中,我们将学习基于 transformers 库来训练你自己的分词器."
]
},
{
@ -63,7 +61,7 @@
"source": [
"### 初始化\n",
"\n",
"首先,我们应该初始化我们的分词器,并确定选择哪种方法。我们将使用字节对编码(BPE)算法。"
"首先,我们应该初始化我们的分词器,并确定选择哪种方法。我们将使用字节对编码(BPE)算法."
]
},
{
@ -92,7 +90,7 @@
"source": [
"### 定义特殊标记\n",
"\n",
"数据集中存在一些我们不希望被分词的特殊标记,我们会将这些标记定义为特殊标记,并将它们传递给分词器训练器,以防止出现错误的分词情况。 "
"数据集中存在一些我们不希望被分词的特殊标记,我们会将这些标记定义为特殊标记,并将它们传递给分词器训练器,以防止出现错误的分词情况."
]
},
{
@ -116,7 +114,8 @@
"metadata": {},
"source": [
"### 从文件中读取数据\n",
"我们使用 JSON Lines(jsonl)格式来存储训练数据,让我们编写函数来处理这些数据。"
"\n",
"在本次实验中,我们使用 JSON Lines(jsonl)格式来存储 Tokenizer 训练数据,分词器内置的训练函数要求训练数据以迭代器的形式传入,因此,我们首先获取一个数据读取的生成器."
]
},
{
@ -142,19 +141,16 @@
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'<s>近年来,人工智能技术迅速发展,深刻改变了各行各业的面貌。机器学习、自然语言处理、计算机视觉等领域的突破性进展,使得智能产品和服务越来越普及。从智能家居到自动驾驶,再到智能医疗,AI的应用场景正在快速拓展。随着技术的不断进步,未来的人工智能将更加智能、更加贴近人类生活。</s>'"
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
"name": "stdout",
"output_type": "stream",
"text": [
"Row 1: <s>近年来,人工智能技术迅速发展,深刻改变了各行各业的面貌。机器学习、自然语言处理、计算机视觉等领域的突破性进展,使得智能产品和服务越来越普及。从智能家居到自动驾驶,再到智能医疗,AI的应用场景正在快速拓展。随着技术的不断进步,未来的人工智能将更加智能、更加贴近人类生活。</s>\n"
]
}
],
"source": [
"data_iter = read_texts_from_jsonl(data_path)\n",
"next(data_iter)"
"print(f'Row 1: {next(data_iter)}')"
]
},
{
@ -163,7 +159,7 @@
"source": [
"### 开始训练!\n",
"\n",
"我们使用分词器的内置函数来训练分词器,轻松愉快。"
"我们使用分词器的内置函数 `tokenizer.train_from_iterator` 来训练分词器."
]
},
{
@ -316,7 +312,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"现在我们已经训练了一个简单的分词器,并将其进行保存,接下来,我们试着加载它,并使用其帮助我们对文本进行编解码"
"现在我们已经训练了一个简单的分词器,并将其进行保存,接下来,我们试着加载它,并使用其帮助我们对文本进行编解码."
]
},
{
@ -325,14 +321,12 @@
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'<s>失去的东西就要学着去接受,学着放下。</s>'"
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
"name": "stdout",
"output_type": "stream",
"text": [
"原始文本:[{'text': '失去的东西就要学着去接受,学着放下。'}]\n",
"修改文本:[{'text': '失去的东西就要学着去接受,学着放下。'}] (添加自定义聊天模板)\n"
]
}
],
"source": [
@ -344,7 +338,8 @@
" msg,\n",
" tokenize=False\n",
")\n",
"new_msg"
"print(f'原始文本:{msg}')\n",
"print(f'修改文本:{msg} (添加自定义聊天模板)')"
]
},
{
@ -356,13 +351,12 @@
"name": "stdout",
"output_type": "stream",
"text": [
"actural length of tokenizer: 259\n"
"分词器词表大小:259\n"
]
}
],
"source": [
"actual_vocab_size = len(tokenizer)\n",
"print('actural length of tokenizer:', actual_vocab_size)"
"print(f'分词器词表大小:{tokenizer.vocab_size}')"
]
},
{
@ -371,19 +365,17 @@
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'input_ids': [1, 164, 100, 112, 164, 239, 122, 166, 251, 229, 163, 119, 253, 167, 101, 126, 164, 111, 112, 167, 102, 226, 164, 258, 102, 166, 254, 225, 164, 239, 122, 165, 239, 101, 164, 240, 248, 174, 123, 237, 164, 258, 102, 166, 254, 225, 165, 245, 125, 163, 119, 236, 162, 225, 227, 2], 'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}"
]
},
"execution_count": 12,
"metadata": {},
"output_type": "execute_result"
"name": "stdout",
"output_type": "stream",
"text": [
"查看分词结果:\n",
"{'input_ids': [1, 164, 100, 112, 164, 239, 122, 166, 251, 229, 163, 119, 253, 167, 101, 126, 164, 111, 112, 167, 102, 226, 164, 258, 102, 166, 254, 225, 164, 239, 122, 165, 239, 101, 164, 240, 248, 174, 123, 237, 164, 258, 102, 166, 254, 225, 165, 245, 125, 163, 119, 236, 162, 225, 227, 2], 'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}\n"
]
}
],
"source": [
"model_inputs = tokenizer(new_msg)\n",
"model_inputs"
"print(f'查看分词结果:\\n{model_inputs}')"
]
},
{
@ -392,19 +384,16 @@
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'<s>失去的东西就要学着去接受,学着放下。</s>'"
]
},
"execution_count": 13,
"metadata": {},
"output_type": "execute_result"
"name": "stdout",
"output_type": "stream",
"text": [
"对分词结果进行解码:<s>失去的东西就要学着去接受,学着放下。</s> (保留特殊字符)\n"
]
}
],
"source": [
"response = tokenizer.decode(model_inputs['input_ids'], skip_special_tokens=False)\n",
"response"
"print(f'对分词结果进行解码:{response} (保留特殊字符)' )"
]
},
{
@ -413,19 +402,25 @@
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'失去的东西就要学着去接受,学着放下。'"
]
},
"execution_count": 14,
"metadata": {},
"output_type": "execute_result"
"name": "stdout",
"output_type": "stream",
"text": [
"对分词结果进行解码:失去的东西就要学着去接受,学着放下。 (移除特殊字符)\n"
]
}
],
"source": [
"response = tokenizer.decode(model_inputs['input_ids'], skip_special_tokens=True)\n",
"response"
"print(f'对分词结果进行解码:{response} (移除特殊字符)' )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 参考资料\n",
"\n",
"- [Hugging Face NLP Course](https://huggingface.co/learn/nlp-course/zh-CN/chapter2/4)"
]
}
],

File diff suppressed because it is too large Load Diff

View File

@ -173,13 +173,13 @@
"name": "stdout",
"output_type": "stream",
"text": [
"2\n"
"预训练数据集长度2\n"
]
}
],
"source": [
"pretrain_dataset = PretrainDataset(path_pretrain, tokenizer)\n",
"print(len(pretrain_dataset))\n",
"print(f'预训练数据集长度{len(pretrain_dataset)}')\n",
"# x, y, lm = pretrain_dataset[0]\n",
"# print(x.shape, y.shape, lm.shape)"
]
@ -306,10 +306,10 @@
" def __getitem__(self, index):\n",
" sample = self.samples[index]\n",
" prompt = self._create_chat_prompt(sample['conversations'])\n",
" print(prompt) # uncomment to see formating prompt\n",
" # print(prompt) # uncomment to see formating prompt\n",
" input_ids = self.tokenizer(prompt).input_ids[:self.max_length] # encode input\n",
" input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids)) # encode padding\n",
" print(input_ids) # uncomment to see encoded prompt\n",
" # print(input_ids) # uncomment to see encoded prompt\n",
" \n",
" loss_mask = self._generate_loss_mask(input_ids)\n",
"\n",
@ -331,19 +331,7 @@
"output_type": "stream",
"text": [
"2\n",
"<s>system\n",
"你是 MiniMind,是一个有用的人工智能助手。</s>\n",
"<s>user\n",
"你好吗?</s>\n",
"<s>assistant\n",
"我很好,谢谢!你呢?</s>\n",
"<s>user\n",
"我也很好,谢谢!</s>\n",
"<s>assistant\n",
"太好了!祝你今天愉快!</s>\n",
"\n",
"[1, 85, 736, 201, 608, 345, 562, 261, 75, 47, 807, 270, 1589, 400, 411, 1946, 740, 1728, 945, 1184, 286, 2, 201, 1, 320, 275, 201, 5134, 2235, 814, 2, 201, 1, 1078, 538, 501, 201, 397, 5675, 270, 4191, 2207, 608, 3274, 814, 2, 201, 1, 320, 275, 201, 397, 698, 5675, 270, 4191, 2207, 2, 201, 1, 1078, 538, 501, 201, 1633, 587, 446, 2207, 5625, 608, 4668, 5295, 1693, 2207, 2, 201, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
"torch.Size([511]) torch.Size([511]) torch.Size([511])\n"
"样本 shape = torch.Size([511]), 标签 shape = torch.Size([511]), loss_mask shape torch.Size([511])\n"
]
}
],
@ -351,7 +339,7 @@
"sft_dataset = SFTDataset(path_sft, tokenizer)\n",
"print(len(sft_dataset))\n",
"x, y, lm = sft_dataset[0]\n",
"print(x.shape, y.shape, lm.shape)\n",
"print(f'样本 shape = {x.shape}, 标签 shape = {y.shape}, loss_mask shape {lm.shape}')\n",
"# print(lm) # 打印 loss mask,你会发现在序列中有两处以 1 填充的序列,这是因为我们在一个 conversation 中开展了两轮对话,其中只有 assistant 回复计算损失"
]
},
@ -516,7 +504,7 @@
"name": "stdout",
"output_type": "stream",
"text": [
"2\n",
"DPO 数据集长度:2\n",
"<s>system\n",
"你是 MiniMind,是一个有用的人工智能助手。</s>\n",
"<s>user\n",
@ -544,7 +532,7 @@
],
"source": [
"dpo_dataset = DPODataset(path_dpo, tokenizer)\n",
"print(len(dpo_dataset))\n",
"print(f'DPO 数据集长度:{len(dpo_dataset)}')\n",
"res = dpo_dataset[0]\n",
"# 如有需要,请自主查看 res 中的元素"
]

View File

@ -59,7 +59,7 @@
"source": [
"## 可选参数设置\n",
"\n",
"首先,查看训练的可选参数,这些参数在实际使用时通过命令行导入,为了保持笔记本的易用性,选择用 class 进行包装"
"首先,查看训练的可选参数,这些参数在实际使用时通过解析命令行进行导入,我们用 class 进行包装."
]
},
{
@ -100,18 +100,15 @@
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'cuda'"
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
"name": "stdout",
"output_type": "stream",
"text": [
"查看工作设备 cuda\n"
]
}
],
"source": [
"args.device"
"print(f'查看工作设备 {args.device}')"
]
},
{
@ -150,7 +147,7 @@
"output_type": "stream",
"text": [
"LLM总参数量:8.915 百万\n",
"cuda:0 6400 <torch.utils.data.dataloader.DataLoader object at 0x0000020C10DF2200>\n"
"模型位于设备:cuda:0, 词表长度:6400, DataLoader:<torch.utils.data.dataloader.DataLoader object at 0x000001E4890FE1D0>\n"
]
}
],
@ -169,36 +166,33 @@
" num_workers=args.num_workers,\n",
")\n",
"\n",
"print(model.device, tokenizer.vocab_size, train_loader)"
"print(f'模型位于设备:{model.device}, 词表长度:{tokenizer.vocab_size}, DataLoader:{train_loader}')"
]
},
{
"cell_type": "code",
"execution_count": 7,
"execution_count": 11,
"id": "60c8c970-f448-4ea0-b21b-8f929e03f6fa",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"([tensor([[ 1, 46, 46, ..., 0, 0, 0],\n",
" [ 1, 5349, 1619, ..., 0, 0, 0]]),\n",
" tensor([[ 46, 46, 47, ..., 0, 0, 0],\n",
" [5349, 1619, 2875, ..., 0, 0, 0]]),\n",
" tensor([[1, 1, 1, ..., 0, 0, 0],\n",
" [1, 1, 1, ..., 0, 0, 0]])],\n",
" 2,\n",
" 1)"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
"name": "stdout",
"output_type": "stream",
"text": [
"打印一个 iter 的数据:\n",
"[tensor([[ 1, 46, 46, ..., 0, 0, 0],\n",
" [ 1, 5349, 1619, ..., 0, 0, 0]]), tensor([[ 46, 46, 47, ..., 0, 0, 0],\n",
" [5349, 1619, 2875, ..., 0, 0, 0]]), tensor([[1, 1, 1, ..., 0, 0, 0],\n",
" [1, 1, 1, ..., 0, 0, 0]])]\n",
"\n",
"数据集大小:2, DataLoader 大小:1\n"
]
}
],
"source": [
"loader = iter(train_loader)\n",
"next(loader), len(train_ds), len(loader)"
"print(f'打印一个 iter 的数据:\\n{next(loader)}\\n')\n",
"print(f'数据集大小:{len(train_ds)}, DataLoader 大小:{len(loader)}')"
]
},
{
@ -324,7 +318,7 @@
"id": "7792bf9c-95a1-4c0d-bb53-fac2a1132ee0",
"metadata": {},
"source": [
"准备完毕,我们进行一轮长度 1 个 iter 的训练看看怎么个事。"
"准备完毕,我们尝试一轮长度 1 个 iter 的训练."
]
},
{
@ -337,7 +331,7 @@
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch:[1/1](0/1) loss:8.889 lr:0.000550000000 epoch_Time:0.0min:\n"
"Epoch:[1/1](0/1) loss:8.853 lr:0.000550000000 epoch_Time:0.0min:\n"
]
}
],
@ -346,6 +340,16 @@
"for epoch in range(args.epochs):\n",
" train_epoch(epoch)"
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "6f257b9e-3a1e-4686-9516-3f9eeb7d0654",
"metadata": {},
"outputs": [],
"source": [
"del model"
]
}
],
"metadata": {

View File

@ -7,9 +7,9 @@
"source": [
"# 4-SFT\n",
"\n",
"在预训练阶段后,我们应该能够获得一个下一词预测模型,此时的模型已经掌握了大量的知识。不过,仅仅具备下一词预测能力是不够的,我们希望大模型能够获得问答能力,这一能力便是在有监督微调(Supervised Fine Tuning,SFT)阶段获得的。\n",
"在预训练阶段后,我们应该能够获得一个下一词预测模型,此时的模型已经掌握了大量的知识。不过,仅仅具备下一词预测能力是不够的,我们希望大模型能够获得问答能力,这一能力便是在有监督微调(Supervised Fine Tuning,SFT)阶段获得的.\n",
"\n",
"在这个笔记本中,我们仅对 SFT 的训练流程进行展示和学习,因此只给出必要的代码片段,如 wandb 和 ddp 不会在此笔记本中涉及。\n",
"在这个笔记本中,我们仅对 SFT 的训练流程进行展示和学习,因此只给出必要的代码片段,如 wandb 和 ddp 不会在此笔记本中涉及.\n",
"\n",
"此笔记本的完整实现见主仓库 `/minimind/train_full_sft.py`"
]
@ -61,7 +61,7 @@
"source": [
"## 可选参数设置\n",
"\n",
"首先,查看训练的可选参数,这些参数在实际使用时通过命令行导入,为了保持笔记本的易用性,选择用 class 进行包装"
"首先,查看训练的可选参数,这些参数在实际使用时通过命令行导入,为了保持笔记本的易用性,选择用 class 进行包装."
]
},
{
@ -102,18 +102,15 @@
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'cuda'"
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
"name": "stdout",
"output_type": "stream",
"text": [
"查看工作设备 cuda\n"
]
}
],
"source": [
"args.device"
"print(f'查看工作设备 {args.device}')"
]
},
{
@ -123,7 +120,7 @@
"source": [
"## 初始化训练\n",
"\n",
"接下来,我们对一些重要模块进行初始化,我们已经了解过,分词器,模型和数据集是大模型的基本组件,我们对其进行初始化。\n",
"接下来,我们对一些重要模块进行初始化,我们已经了解过,分词器,模型和数据集是大模型的基本组件,我们对其进行初始化.\n",
"\n",
"> 注意 与预训练阶段不同的是 在 sft 阶段 我们实际上是在上一阶段训练获得的模型的基础上修改数据集进行接续训练 因此需要载入上一阶段的模型权重 出于展示的目的 载入权重的代码在此笔记本中只作展示 并不执行"
]
@ -158,7 +155,7 @@
"output_type": "stream",
"text": [
"LLM总参数量:8.915 百万\n",
"cuda:0 6400 <torch.utils.data.dataloader.DataLoader object at 0x00000184B5C6B4F0>\n"
"模型位于设备:cuda:0, 词表长度:6400, DataLoader:<torch.utils.data.dataloader.DataLoader object at 0x000002D0DD01D600>\n"
]
}
],
@ -177,7 +174,7 @@
" num_workers=args.num_workers,\n",
")\n",
"\n",
"print(model.device, tokenizer.vocab_size, train_loader)"
"print(f'模型位于设备:{model.device}, 词表长度:{tokenizer.vocab_size}, DataLoader:{train_loader}')"
]
},
{
@ -187,26 +184,23 @@
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"([tensor([[ 1, 85, 736, ..., 0, 0, 0],\n",
" [ 1, 85, 736, ..., 0, 0, 0]]),\n",
" tensor([[ 85, 736, 201, ..., 0, 0, 0],\n",
" [ 85, 736, 201, ..., 0, 0, 0]]),\n",
" tensor([[0, 0, 0, ..., 0, 0, 0],\n",
" [0, 0, 0, ..., 0, 0, 0]])],\n",
" 2,\n",
" 1)"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
"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",
"next(loader), len(train_ds), len(loader)"
"print(f'打印一个 iter 的数据:\\n{next(loader)}\\n')\n",
"print(f'数据集大小:{len(train_ds)}, DataLoader 大小:{len(loader)}')"
]
},
{
@ -230,7 +224,7 @@
"source": [
"## 启动训练\n",
"\n",
"训练一个深度学习模型,还涉及到了优化器,损失函数和学习率调度。接下来,我们查看 MiniMind 训练部分的代码,并进行一轮简单的训练。\n",
"训练一个深度学习模型,还涉及到了优化器,损失函数和学习率调度. 接下来,我们查看 MiniMind 训练部分的代码,并进行一轮简单的训练.\n",
"\n",
"> 不难发现 pretrain 阶段和 sft 阶段的训练主体差不多 因为这两个阶段的差异体现在数据集格式 而数据集在经过 chat template 格式化后差异小了很多"
]
@ -334,7 +328,7 @@
"id": "546f28c2-418e-453c-88aa-695468f9216c",
"metadata": {},
"source": [
"准备完毕,我们进行一轮长度 1 个 iter 的训练看看怎么个事。"
"准备完毕,我们尝试一轮长度 1 个 iter 的训练."
]
},
{
@ -347,7 +341,7 @@
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch:[1/1](0/1) loss:8.865 lr:0.000550000000 epoch_Time:0.0min:\n"
"Epoch:[1/1](0/1) loss:8.869 lr:0.000550000000 epoch_Time:0.0min:\n"
]
}
],
@ -356,6 +350,16 @@
"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": {

View File

@ -7,9 +7,9 @@
"source": [
"# 5-DPO\n",
"\n",
"直接偏好优化(Direct Preference Optimization,DPO)是后训练阶段中,使用正反样例激励大模型产生符合人类偏好的回答的策略,为人类反馈强化学习(Reinforcement Learning from Human Feedback, RLHF)提供了一个高效简化的替代方案。通过这一阶段的训练,大模型将会学会依照人类的喜好生成回复。\n",
"直接偏好优化(Direct Preference Optimization,DPO)是后训练阶段中,使用正反样例激励大模型产生符合人类偏好的回答的策略,为人类反馈强化学习(Reinforcement Learning from Human Feedback, RLHF)提供了一个高效简化的替代方案。通过这一阶段的训练,大模型将会学会依照人类的喜好生成回复.\n",
"\n",
"在这个笔记本中,我们仅对 DPO 的训练流程进行展示和学习,因此只给出必要的代码片段,如 wandb 和 ddp 不会在此笔记本中涉及。\n",
"在这个笔记本中,我们仅对 DPO 的训练流程进行展示和学习,因此只给出必要的代码片段,如 wandb 和 ddp 不会在此笔记本中涉及.\n",
"\n",
"此笔记本的完整实现见主仓库 `/minimind/train_dpo.py`"
]
@ -61,7 +61,7 @@
"source": [
"## 可选参数设置\n",
"\n",
"首先,查看训练的可选参数,这些参数在实际使用时通过命令行导入,为了保持笔记本的易用性,选择用 class 进行包装"
"首先,查看训练的可选参数,这些参数在实际使用时通过命令行导入,为了保持笔记本的易用性,选择用 class 进行包装."
]
},
{
@ -103,18 +103,15 @@
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'cuda'"
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
"name": "stdout",
"output_type": "stream",
"text": [
"查看工作设备 cuda\n"
]
}
],
"source": [
"args.device"
"print(f'查看工作设备 {args.device}')"
]
},
{
@ -124,7 +121,7 @@
"source": [
"## 初始化训练\n",
"\n",
"接下来,我们对一些重要模块进行初始化,我们已经了解过,分词器,模型和数据集是大模型的基本组件,我们对其进行初始化。\n",
"接下来,我们对一些重要模块进行初始化,我们已经了解过,分词器,模型和数据集是大模型的基本组件,我们对其进行初始化.\n",
"\n",
"> 在这一阶段 我们调整的是大模型的问答偏好 因此与 sft 阶段同理 我们需要载入在 sft 阶段微调好的问答模型"
]
@ -168,7 +165,7 @@
"output_type": "stream",
"text": [
"LLM总参数量:6.096 百万\n",
"cuda:0 6400 <torch.utils.data.dataloader.DataLoader object at 0x00000163D4EEF640>\n"
"模型位于设备:cuda:0, 词表长度:6400, DataLoader:<torch.utils.data.dataloader.DataLoader object at 0x00000264370DC7F0>\n"
]
}
],
@ -187,7 +184,7 @@
" num_workers=args.num_workers,\n",
")\n",
"\n",
"print(model.device, tokenizer.vocab_size, train_loader)"
"print(f'模型位于设备:{model.device}, 词表长度:{tokenizer.vocab_size}, DataLoader:{train_loader}')"
]
},
{
@ -197,32 +194,26 @@
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"({'x_chosen': tensor([[ 1, 85, 736, ..., 0, 0, 0],\n",
" [ 1, 85, 736, ..., 0, 0, 0]]),\n",
" 'y_chosen': tensor([[ 85, 736, 201, ..., 0, 0, 0],\n",
" [ 85, 736, 201, ..., 0, 0, 0]]),\n",
" 'mask_chosen': tensor([[0, 0, 0, ..., 0, 0, 0],\n",
" [0, 0, 0, ..., 0, 0, 0]]),\n",
" 'x_rejected': tensor([[ 1, 85, 736, ..., 0, 0, 0],\n",
" [ 1, 85, 736, ..., 0, 0, 0]]),\n",
" 'y_rejected': tensor([[ 85, 736, 201, ..., 0, 0, 0],\n",
" [ 85, 736, 201, ..., 0, 0, 0]]),\n",
" 'mask_rejected': tensor([[0, 0, 0, ..., 0, 0, 0],\n",
" [0, 0, 0, ..., 0, 0, 0]])},\n",
" 2,\n",
" 1)"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
"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",
"next(loader), len(train_ds), len(loader)"
"print(f'打印一个 iter 的数据:\\n{next(loader)}\\n')\n",
"print(f'数据集大小:{len(train_ds)}, DataLoader 大小:{len(loader)}')"
]
},
{
@ -249,7 +240,7 @@
"source": [
"# 启动训练\n",
"\n",
"训练一个深度学习模型,还涉及到了优化器,损失函数和学习率调度。接下来,我们查看 MiniMind 训练部分的代码,并进行一轮简单的训练。\n",
"训练一个深度学习模型,还涉及到了优化器,损失函数和学习率调度. 接下来,我们查看 MiniMind 训练部分的代码,并进行一轮简单的训练.\n",
"\n",
"> DPO 阶段涉及 DPO 损失函数涉及 因此与前两个阶段相比内容略有增加 不过整体流程与逻辑类似"
]
@ -278,9 +269,9 @@
"id": "5f2c3ac1-92bf-4eae-8d22-7fac4cfa69df",
"metadata": {},
"source": [
"DPO 的原理是增加偏好样本的对数概率与减小非偏好样本响应的对数概率。\n",
"DPO 的原理是增加偏好样本的对数概率与减小非偏好样本响应的对数概率.\n",
"\n",
"该阶段引入 DPO 损失函数,通过计算选择样本和拒绝样本的对数比率,然后基于这些比率计算 DPO 损失,适用于偏好学习任务。"
"该阶段引入 DPO 损失函数,通过计算选择样本和拒绝样本的对数比率,然后基于这些比率计算 DPO 损失,适用于偏好学习任务."
]
},
{
@ -408,7 +399,7 @@
"id": "49386dc7-2919-428a-a073-f37d576afbcf",
"metadata": {},
"source": [
"准备完毕,我们进行一轮长度 1 个 iter 的训练看看怎么个事。"
"准备完毕,我们尝试一轮长度 1 个 iter 的训练."
]
},
{
@ -421,7 +412,7 @@
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch:[1/1](0/1) loss:0.693 lr:0.000550000000 epoch_Time:0.0min:\n"
"Epoch:[1/1](0/1) loss:0.694 lr:0.000550000000 epoch_Time:0.0min:\n"
]
}
],
@ -430,6 +421,16 @@
"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": {

597
notebook/6_lora.ipynb Normal file
View File

@ -0,0 +1,597 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "fd56e0ed-d6e3-4fdd-bed2-2a0c32fb61f9",
"metadata": {},
"source": [
"# 6-LoRA\n",
"\n",
"大语言模型的低秩适应(LoRA, Low-Rank Adaptation of Large Language Models) 是一项大模型参数高效微调技术,其可以显著减少可训练参数的数量.\n",
"\n",
"由于大模型参数量较大,直接进行全参微调需要消耗大量硬件资源,LoRA 的工作原理是将少量的新权重插入倒模型中,并且仅训练这些权重,这使得使用 LoRA 进行训练的速度更快、内存效率更高,并生成更小的模型权重.\n",
"\n",
"具体来说,LoRA 它冻结了预训练模型 W 的权重,并注入可训练的秩分解矩阵 A 与 B,在微调时,只训练降维矩阵 A 和 升维矩阵 B,微调结束后,将 AB 与 W 进行叠加.\n",
"\n",
"![image](./demo/lora.png)\n",
"\n",
"其中,用随机高斯分布进行初始化 A,用 0 矩阵初始化 B,从而保证训练开始时旁路矩阵为 0 矩阵.\n",
"\n",
"具体来看,假设模型经过预训练主干的输出为 $W_0 x$,在 LoRA 微调阶段,我们可以用如下形式对输出进行表示.\n",
"\n",
"$$h=W_0x + \\Delta Wx = W_0x + BA x=(W_0 + BA)x$$\n",
"\n",
"其中, $B \\in \\mathbb{R}^{d \\times r},A \\in \\mathbb{R}^{r\\times k}$,r 为 LoRA 低秩矩阵的维数,$r << min(d, k)$."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "dad11c0c-043d-44cb-a9a7-4c36c2e909da",
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"from torch import optim, nn"
]
},
{
"cell_type": "markdown",
"id": "3031399b-9342-454f-86df-9a96b8c6ff50",
"metadata": {},
"source": [
"## LoRA Adapter\n",
"\n",
"简单的来说,LoRA 矩阵就是具有一个隐藏层的全连接网络,其挂接在主干网络边侧进行参数更新,我们来看看 MiniMind 模型是如何在主干网络外部定义 LoRA 网络结构的."
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "2f9bc623-f13d-4d67-bf2e-dce0784b43f8",
"metadata": {},
"outputs": [],
"source": [
"class LoRA(nn.Module):\n",
" def __init__(self, in_features, out_features, rank):\n",
" super().__init__()\n",
" self.rank = rank # LoRA 秩 控制低秩矩阵大小\n",
" self.A = nn.Linear(in_features, rank, bias=False)\n",
" self.B = nn.Linear(rank, out_features, bias=False)\n",
" # 矩阵 A 高斯分布初始化\n",
" self.A.weight.data.normal_(mean=0.0, std=0.02)\n",
" # 矩阵 B 全零初始化\n",
" self.B.weight.data.zero_()\n",
"\n",
" def forward(self, x):\n",
" return self.B(self.A(x))"
]
},
{
"cell_type": "markdown",
"id": "4a3557bd-d34d-4a15-b789-22afea60a87d",
"metadata": {},
"source": [
"可以看到,LoRA 的网络结构非常简单直观,我们接下来定义一个方法,将 LoRA 网络应用到 MiniMind 模型的特定线性层."
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "4d52e346-af62-40e2-a5b0-5cc0cf8ae6e8",
"metadata": {},
"outputs": [],
"source": [
"def apply_lora(model, rank=16):\n",
" \"\"\"将 LoRA 模块与目标模块进行绑定\"\"\"\n",
" for name, module in model.named_modules():\n",
" if isinstance(module, nn.Linear) and module.weight.shape[0] == module.weight.shape[1]:\n",
" lora = LoRA(module.weight.shape[0], module.weight.shape[1], rank=rank).to(model.device)\n",
" setattr(module, 'lora', lora) # 显式添加 LoRA 模块\n",
" original_forward = module.forward\n",
"\n",
" # 修改目标模块的 forward 函数\n",
" def forward_with_lora(x, layer1=original_forward, layer2=lora):\n",
" return layer1(x) + layer2(x)\n",
" \n",
" module.forward = forward_with_lora\n",
" # 打印 LoRA 绑定的模块名称\n",
" print(f'apply lora on module: {name}')"
]
},
{
"cell_type": "markdown",
"id": "753a54c6-a442-4ce5-be89-3bc601bc8d26",
"metadata": {},
"source": [
"我们可以声明一个小模型,对于 LoRA 的绑定进行测试."
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "f830221b-891e-4e27-b401-20b6fe5d3c59",
"metadata": {},
"outputs": [],
"source": [
"class TestModel(nn.Module):\n",
" def __init__(self):\n",
" super().__init__()\n",
" self.linear1 = nn.Linear(64, 512)\n",
" self.linear2 = nn.Linear(512, 512)\n",
" self.linear3 = nn.Linear(512, 64)\n",
"\n",
" @property\n",
" def device(self):\n",
" return next(self.parameters()).device\n",
" \n",
" def forward(self, x):\n",
" out = self.linear3(self.linear2(self.linear1))\n",
" return out"
]
},
{
"cell_type": "markdown",
"id": "1a8edbc2-f36c-4127-be3b-3c598872621a",
"metadata": {},
"source": [
"按照 apply_lora 的函数逻辑,LoRA 模块会应用在主干网络中满足 input_feature == output_feature 的模块上."
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "eb24e24f-da09-4cd7-b4ce-7fe482194f75",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"apply lora on module: linear2\n",
"TestModel(\n",
" (linear1): Linear(in_features=64, out_features=512, bias=True)\n",
" (linear2): Linear(\n",
" in_features=512, out_features=512, bias=True\n",
" (lora): LoRA(\n",
" (A): Linear(in_features=512, out_features=16, bias=False)\n",
" (B): Linear(in_features=16, out_features=512, bias=False)\n",
" )\n",
" )\n",
" (linear3): Linear(in_features=512, out_features=64, bias=True)\n",
")\n"
]
}
],
"source": [
"test_model = TestModel()\n",
"apply_lora(test_model)\n",
"print(test_model)"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "f882e6c0-65a6-47b2-9ee6-4a3fca988636",
"metadata": {},
"outputs": [],
"source": [
"del test_model"
]
},
{
"cell_type": "markdown",
"id": "9771475e-0c60-42db-b0cd-7e193de23f7c",
"metadata": {},
"source": [
"完成了 LoRA 模块在主干网络特定模块的绑定后,我们便可以冻结主干网络参数进行微调了,不过,考虑到主干网络权重在训练过程中并不会做任何参数更新,我们可以只保存 LoRA 模块的参数来节省内存,下面给出加载/保存 LoRA 权重的方法."
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "95a453d7-2011-4575-9a51-36d91876874f",
"metadata": {},
"outputs": [],
"source": [
"def load_lora(model, path):\n",
" state_dict = torch.load(path, map_location=model.device)\n",
" for name, module in model.named_modules():\n",
" if hasattr(module, 'lora'):\n",
" lora_state = {k.replace(f'{name}.lora.', ''): v for k, v in state_dict.items() if f'{name}.lora.' in k}\n",
" module.lora.load_state_dict(lora_state)\n",
"\n",
"\n",
"def save_lora(model, path):\n",
" state_dict = {}\n",
" for name, module in model.named_modules():\n",
" if hasattr(module, 'lora'):\n",
" lora_state = {f'{name}.lora.{k}': v for k, v in module.lora.state_dict().items()}\n",
" state_dict.update(lora_state)\n",
" torch.save(state_dict, path)"
]
},
{
"cell_type": "markdown",
"id": "854b8afb-39b1-4175-8e69-4a9d5b25a3e5",
"metadata": {},
"source": [
"## Fine-Tuning MiniMind with LoRA"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "10baf158-6653-49e3-9ba3-22b2e341730f",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import platform\n",
"import argparse\n",
"import random\n",
"import time\n",
"import math\n",
"import warnings\n",
"import torch.distributed as dist\n",
"from contextlib import nullcontext\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": "f0b93272-80b3-45ce-9dec-4583fca8f87a",
"metadata": {},
"source": [
"### 可选参数设置\n",
"\n",
"首先,查看训练的可选参数,这些参数在实际使用时通过解析命令行进行导入,我们用 class 进行包装."
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "bb9d4bac-8a30-4ca9-b293-94df3f880826",
"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/lora_data.jsonl' # 数据集路径\n",
" lora_name: str = 'lora_identity' # 根据任务保存成lora_(英文/医学/心理...)"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "3b496134-ac9a-4146-a236-fd8867f45bc6",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"查看工作设备 cuda\n"
]
}
],
"source": [
"print(f'查看工作设备 {args.device}')"
]
},
{
"cell_type": "markdown",
"id": "32b956e6-e81b-49d9-bbf7-b78fbbc7eff5",
"metadata": {},
"source": [
"接下来,我们对分词器、MiniMindLM 和数据迭代器执行初始化."
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "2fb7c880-b86e-4329-aeef-2827837425a3",
"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",
" return model.to(args.device), tokenizer"
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "7f0c74d0-a191-40c1-b922-104f241aa879",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"apply lora on module: layers.0.attention.wq\n",
"apply lora on module: layers.0.attention.wo\n",
"apply lora on module: layers.1.attention.wq\n",
"apply lora on module: layers.1.attention.wo\n",
"模型位于设备:cuda:0, 词表长度:6400, DataLoader:<torch.utils.data.dataloader.DataLoader object at 0x0000018FB27CA050>\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",
"apply_lora(model)\n",
"\n",
"# 由于 MiniMind 用于 LoRA 微调的数据集和 SFT 数据集格式相同,可以用 SFT 数据集进行加载\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": "markdown",
"id": "c35bb3ae-0b31-42ff-88d1-04d7e945bd31",
"metadata": {},
"source": [
"可以看到,LoRA 模块挂接在 Attention Block 的 Query 与 Output 线性层上,下面查看 LoRA 微调下可学习参数的占比:"
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "766ef2d5-1f24-4efb-9eb9-f969a57a85b0",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"LLM 总参数量: 8980992\n",
"LoRA 参数量: 65536\n",
"LoRA 参数占比: 0.73%\n"
]
}
],
"source": [
"total_params = sum(p.numel() for p in model.parameters()) # 总参数数量\n",
"lora_params_count = sum(p.numel() for name, p in model.named_parameters() if 'lora' in name) # LoRA 参数数量\n",
"print(f\"LLM 总参数量: {total_params}\")\n",
"print(f\"LoRA 参数量: {lora_params_count}\")\n",
"print(f\"LoRA 参数占比: {lora_params_count / total_params * 100:.2f}%\")"
]
},
{
"cell_type": "markdown",
"id": "8a487807-65da-4e19-ac2f-fce7077494c0",
"metadata": {},
"source": [
"接下来,冻结 MiniMindLM 主干网络的参数,做好 LoRA 微调准备."
]
},
{
"cell_type": "code",
"execution_count": 14,
"id": "db99fd52-4f07-4278-939f-47b91aeeedbe",
"metadata": {},
"outputs": [],
"source": [
"for name, param in model.named_parameters():\n",
" if 'lora' not in name:\n",
" param.requires_grad = False\n",
"lora_params = [] # 收集 LoRA 模块的可学习参数, 提供给优化器\n",
"for name, param in model.named_parameters():\n",
" if 'lora' in name:\n",
" lora_params.append(param)"
]
},
{
"cell_type": "markdown",
"id": "6f865b22-35a7-4c2c-85a1-a13d1ada12a7",
"metadata": {},
"source": [
"### 启动训练\n",
"\n",
"接下来,我们定义 MiniMind LoRA 微调所使用的优化器,损失函数和学习率调度,并进行一轮简单的训练."
]
},
{
"cell_type": "code",
"execution_count": 15,
"id": "2b57beaf-c04a-4b88-9617-7142054f2927",
"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",
"# 只训练 LoRA 模块参数\n",
"optimizer = optim.AdamW(lora_params, 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": "8037d2b2-1fa3-4e19-9ad5-9674ff866a36",
"metadata": {},
"source": [
"接下来,我们来看训练函数."
]
},
{
"cell_type": "code",
"execution_count": 16,
"id": "8f50d42d-da6d-4f8b-aa35-853b3900f8b5",
"metadata": {},
"outputs": [],
"source": [
"# 代码和 SFT 几乎一致\n",
"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",
" 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",
" 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_(lora_params, 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",
" # 只保存 LoRA 模块的权重\n",
" # if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):\n",
" # model.eval()\n",
" # # 【区别1】只保存lora权重即可\n",
" # save_lora(model, f'{args.save_dir}/lora/{args.lora_name}_{lm_config.dim}.pth')\n",
" # model.train()"
]
},
{
"cell_type": "markdown",
"id": "618876fa-7f0f-43b3-9eed-384be6cd66be",
"metadata": {},
"source": [
"接下来,我们启动一个 Epoch 的训练进行观察."
]
},
{
"cell_type": "code",
"execution_count": 17,
"id": "38584341-9ef6-4ebd-9024-0c44341de7c8",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch:[1/1](0/1) loss:8.894 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": 18,
"id": "61150c72-39eb-4b90-b321-d4c7aa6c885a",
"metadata": {},
"outputs": [],
"source": [
"del model"
]
},
{
"cell_type": "markdown",
"id": "dd5e9aa8-b5a5-4380-989d-e54948f33a1b",
"metadata": {},
"source": [
"## 参考资料\n",
"\n",
"- [LoRA: Low-Rank Adaptation of Large Language Models](https://arxiv.org/abs/2106.09685)\n",
"- [HuggingFace LoRA](https://huggingface.co/docs/diffusers/en/training/lora)\n",
"- [LoRA 微调和低秩矩阵](https://www.cnblogs.com/ghj1976/p/18032882/lora-finetuning-he-di-zhi-ju-zhen)"
]
}
],
"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
}

430
notebook/7_distill.ipynb Normal file
View File

@ -0,0 +1,430 @@
{
"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 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": "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 = './demo/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 0x000002AE220D7880>\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.3301 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
}

359
notebook/8_reason.ipynb Normal file
View File

@ -0,0 +1,359 @@
{
"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": 2,
"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": 3,
"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": 4,
"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": 6,
"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": 7,
"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 0x000002272AFA32B0>\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": 11,
"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": 15,
"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": 16,
"id": "f021611c-6291-4c50-acf0-9a24dde87e77",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch:[1/1](0/1) loss:11.873 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": 17,
"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
}

BIN
notebook/demo/lora.png Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 12 KiB

View File

@ -0,0 +1,2 @@
{"conversations": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "Hello! I am MiniMind, an AI assistant developed by Jingyao Gong. How can I assist you today?"}]}
{"conversations": [{"role": "user", "content": "hello"}, {"role": "assistant", "content": "Hello! I am MiniMind, an AI assistant developed by Jingyao Gong. How can I assist you today?"}]}

View File

@ -0,0 +1,2 @@
{"conversations": [{"role": "user", "content": "请用一段话描述阿里巴巴集团的企业文化。"}, {"role": "assistant", "content": "<think>\n嗯,用户让我用一段话描述阿里巴巴集团的企业文化。首先,我需要明确阿里巴巴集团的平均文化是什么。企业文化的定义是企业内部发生、指导管理方向的准则和规范性的精神或规范。\n\n阿里巴巴集团在2000年代中期开始发展,之前也有类似的概念,比如微软的人工智能理念。所以,我可以从这个发展中来寻找关键词和结构。\n\n接下来,我要考虑用户可能的意图。用户可能想了解阿里巴巴的企业文化,或者需要用于某个报告、宣传材料或者内部交流中。因此,我需要确保描述既全面又有深度,涵盖战略、创新、效率等各个方面。\n\n我还需要用简洁的语言来表达,避免过于复杂的词汇,让段落易于理解。同时,要注意逻辑连贯,段落之间要有过渡,让读者能够顺畅地跟随观点。\n\n可能还需要引用一些关键点,比如战略协作、用户体验、创新、效率、数字化转型、客户支持、全球化等,这些都是阿里巴巴成功的关键点,能够帮助用户全面了解企业文化。\n\n最后,检查语法和结构是否正确,确保专业性和准确性,这样用户可以直接使用这段文字,或者稍作修改,就能得到满意的效果。\n</think>\n<answer>\n阿里巴巴集团的企业文化以战略协作为核心,倡导员工之间的资源整合与创新思维的碰撞,旨在实现企业价值的提升和持续发展。阿里巴巴将其视为一个协同共生的整体,致力于通过技术创新、用户体验优化和数字化转型,致力于为客户创造最大价值。企业内外部力量协同作战,推动企业从单纯的业务拓展延伸至价值创造和社会服务,并在全球范围内进行全方位的数字化转型,以满足多样化、个性化、高端化的客户需求。阿里巴巴集团ix platform的建立旨在帮助员工实现高效协作,激发创新精神,吸引更多优秀人才加入,共同推动企业不断向前发展。\n</answer>"}]}
{"conversations": [{"role": "user", "content": "你好,我是张华,今年15岁,是北京四中的学生。我喜欢编程和篮球,我梦想成为一名软件工程师。很高兴认识你!"}, {"role": "assistant", "content": "<think>\n好,用户是张华,15岁,来自北京四中,喜欢编程和篮球,梦想成为软件工程师。我需要回应他的兴趣,强调他的优势,比如编程基础和对篮球的热爱,有针对性地引导他规划未来。可以建议他关注 Font Awesome,展示他的技术背景,同时鼓励他扩展兴趣,并辅以鼓励自己。\n</think>\n<answer>\n你好!很高兴认识你。既然你对编程和篮球都很感兴趣,结合你的梦想成为一名软件工程师,这可能是一个很好的切入点。你编程的能力很强,可能在某个领域表现出色,这也是你成为 software engineer的有利条件。\n\n其实,除了你的编程技能,你对篮球的热爱也是挺有趣的,如果你能将这些技能结合使用,可能是一个很好的方向。例如,你可以利用你的篮球知识来设计游戏或解决某种技术问题,或者利用你的编程能力来优化自己的篮球战术。\n\n如果你对 Font Awesome 这样的开源 hesitate,这将是你的一个很好的起点。它可以让你快速开发一个简单的网站,展示你的编程能力和对足球的热爱。这样的话,你不仅能展示你的技术能力,还能展示你的 Otherwise Half。\n\n说到 else half,你应该更 Later来扩展你的兴趣,比如学习 Game Dev, Game Engine, 项目管理等,这样可以让你的探索更加多元化。或许你对 AI 或数据科学也有一定的兴趣,如果有机会,可以加入一些相关的话题和团队,体验运动与技术结合的有趣场景。\n\n总之,你很有潜力成为 software工程师!比方说,你可以尝试以下步骤:\n\n1. 关注 Font Awesome,展示你的编程能力。\n2. 尝试开发某个功能,结合你的篮球知识。\n3. 现在是好好学习 else half 的机会。\n\n希望这是一个鼓励的方向。祝你未来一切顺利!\n</answer>"}]}