From 534d4e2c1acbc42ceceff26aa1b70b8998d23793 Mon Sep 17 00:00:00 2001 From: mingkun Date: Tue, 25 Feb 2025 23:00:46 +0800 Subject: [PATCH] update interactive notebook --- notebook/0_tokenizer.ipynb | 119 ++- notebook/1_model.ipynb | 1330 +++++++++++++++++++++++++-------- notebook/2_dataset.ipynb | 28 +- notebook/3_pretrain.ipynb | 66 +- notebook/4_sft.ipynb | 72 +- notebook/5_dpo.ipynb | 85 +-- notebook/6_lora.ipynb | 597 +++++++++++++++ notebook/7_distill.ipynb | 430 +++++++++++ notebook/8_reason.ipynb | 359 +++++++++ notebook/demo/lora.png | Bin 0 -> 11765 bytes notebook/demo/lora_data.jsonl | 2 + notebook/demo/r1_data.jsonl | 2 + 12 files changed, 2583 insertions(+), 507 deletions(-) create mode 100644 notebook/6_lora.ipynb create mode 100644 notebook/7_distill.ipynb create mode 100644 notebook/8_reason.ipynb create mode 100644 notebook/demo/lora.png create mode 100644 notebook/demo/lora_data.jsonl create mode 100644 notebook/demo/r1_data.jsonl diff --git a/notebook/0_tokenizer.ipynb b/notebook/0_tokenizer.ipynb index 20e4b30..8bde653 100644 --- a/notebook/0_tokenizer.ipynb +++ b/notebook/0_tokenizer.ipynb @@ -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": [ - "'近年来，人工智能技术迅速发展，深刻改变了各行各业的面貌。机器学习、自然语言处理、计算机视觉等领域的突破性进展，使得智能产品和服务越来越普及。从智能家居到自动驾驶，再到智能医疗，AI的应用场景正在快速拓展。随着技术的不断进步，未来的人工智能将更加智能、更加贴近人类生活。'" - ] - }, - "execution_count": 4, - "metadata": {}, - "output_type": "execute_result" + "name": "stdout", + "output_type": "stream", + "text": [ + "Row 1: 近年来，人工智能技术迅速发展，深刻改变了各行各业的面貌。机器学习、自然语言处理、计算机视觉等领域的突破性进展，使得智能产品和服务越来越普及。从智能家居到自动驾驶，再到智能医疗，AI的应用场景正在快速拓展。随着技术的不断进步，未来的人工智能将更加智能、更加贴近人类生活。\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": [ - "'失去的东西就要学着去接受，学着放下。'" - ] - }, - "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": [ - "'失去的东西就要学着去接受，学着放下。'" - ] - }, - "execution_count": 13, - "metadata": {}, - "output_type": "execute_result" + "name": "stdout", + "output_type": "stream", + "text": [ + "对分词结果进行解码：失去的东西就要学着去接受，学着放下。 (保留特殊字符)\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)" ] } ], diff --git a/notebook/1_model.ipynb b/notebook/1_model.ipynb index b9ce3b9..5ca2d42 100644 --- a/notebook/1_model.ipynb +++ b/notebook/1_model.ipynb @@ -13,20 +13,21 @@ "id": "2f5a93fb-2d09-4af4-9782-578357927842", "metadata": {}, "source": [ - "如今的主流大模型的架构大致分为稠密模型(Dense)和混合专家模型(MOE), 在这一板块，我们将围绕 MiniMind 系列模型的源代码展开学习。\n", + "在当前人工智能领域, 主流大模型从架构上大致可分为稠密（Dense）模型和混合专家（Mixture of Expert, MoE）模型 。稠密模型中所有参数在每次计算时都会参与运算；混合专家模型则将不同的 “专家” 模块组合, 根据输入选择合适的专家处理, 能在保证效果的同时减少计算量和参数量.\n", "\n", + "MiniMind 系列模型在 Llama 3.1 的基础上设计, 基于经典的 Transformer Deocder-Only 架构，在这一板块, 我们将围绕 MiniMind 系列模型的源代码展开学习." + ] + }, + { + "cell_type": "markdown", + "id": "f12e2248-7b71-442f-96a0-01a5cb4a1171", + "metadata": {}, + "source": [ "## MiniMind Dense Model\n", "\n", - "MiniMind 系列模型在 Llama 3.1 的基础上设计，基于经典的 Transformer Deocder-Only 架构，其主要特点如下：\n", - "1. 采用 Pre-Norm 归一化方法，使用 RMSNorm 归一化函数。\n", - "2. 使用 SwiGLU 激活函数。\n", - "3. 旋转位置嵌入 (RoPE)\n", + "作者提供了对其 MiniMind Dense Model 模型结构的可视化：\n", "\n", - "作者提供了对其 MiniMind 模型结构的可视化：\n", - "\n", - "![](../images/LLM-structure.png)\n", - "\n", - "我们在上一节中完成了 Tokenizer 的学习，这一节我们关注 LLM 模型架构的具体实现。" + "![](../images/LLM-structure.png)" ] }, { @@ -36,9 +37,19 @@ "metadata": {}, "outputs": [], "source": [ + "import math\n", + "import struct\n", + "import inspect\n", + "import time\n", + "\n", + "from demo.LMConfig import LMConfig\n", + "from typing import Any, Optional, Tuple, List\n", + "import numpy as np\n", "import torch\n", - "import torch.nn as nn\n", - "import torch.nn.functional as F" + "import torch.nn.functional as F\n", + "from torch import nn\n", + "from transformers import PreTrainedModel\n", + "from transformers.modeling_outputs import CausalLMOutputWithPast" ] }, { @@ -46,21 +57,23 @@ "id": "acbb3ad8-5143-45ff-af16-8dcc0cc2cf53", "metadata": {}, "source": [ - "### Root Mean Square Layer Normalization (RMSNorm)\n", + "### 均方根层归一化 (Root Mean Square Layer Normalization, RMSNorm)\n", "\n", - "RMSNorm 是对 LayerNorm 的一个改进, 没有做 re-center 操作（移除了均值项），可以看作 LayerNorm 在均值为零时的特例，使用平方根均值归一化降低噪声影响。\n", + "RMSNorm 是对 LayerNorm 的一个改进, 没有做 re-center 操作（移除了均值项）, 可以看作 LayerNorm 在均值为零时的特例, 使用平方根均值归一化降低噪声影响。\n", "\n", - "- Layer Norm\n", + "- **Layer Norm**\n", "\n", "$$y = \\frac{x-E(x)}{\\sqrt{Var(x) + \\epsilon}} * \\gamma + \\beta$$\n", "\n", - "假设输入张量形状为 (batch_size, sequence_length, embedding_dim)，层归一化会对 embedding_dim 维度进行归一化操作，其中， $\\epsilon$ 是一个超参数，用于防止非法运算（分母为零）， $\\gamma$, $\\beta$ 均为可学习参数。\n", + "假设输入张量形状为 (batch_size, sequence_length, embedding_dim), 层归一化对 embedding_dim 维度进行归一化操作, 其中, $\\epsilon$ 是一个超参数, 用于防止分母为零导致结果上溢, $\\gamma$, $\\beta$ 均为可学习参数。\n", "\n", - "- RMS Norm\n", + "- **RMS Norm**\n", "\n", - "$$a_i=\\frac{a_i}{RMS(a) g_i}, where RMS(a) = \\sqrt{\\frac{1}{n}\\sum^n_{i=1}a^2_i}.$$\n", + "$$a_i=\\frac{a_i}{RMS(a) + \\epsilon} * \\gamma, \\quad where \\quad RMS(a) = \\sqrt{\\frac{1}{n}\\sum^n_{i=1}a^2_i}.$$\n", "\n", - "假设输入张量形状为 (batch_size, sequence_length, embedding_dim)，RMS归一化同样对 embedding_dim 维度进行归一化，其中，$g_i$ 为可学习参数，用于对均方根归一化结果进行加权。" + "假设输入张量形状为 (batch_size, sequence_length, embedding_dim), RMS Norm 对 embedding_dim 维度进行归一化,其中, 其中, $\\epsilon$ 是一个超参数, 用于防止分母为零导致结果上溢, $\\gamma$ 为可学习参数.\n", + "\n", + "不难发现, 当均值为零时, Layer Norm 退化为 RMS Norm. 这是因为 RMS Norm 在 Layer Norm 的基础上舍弃了中心化操作, 仅用缩放进行归一化, 其不改变数据原本的分布, 有利于激活函数输出的稳定." ] }, { @@ -71,13 +84,13 @@ "outputs": [], "source": [ "class RMSNorm(torch.nn.Module):\n", - " def __init__(self, dim: int, eps: float):\n", + " def __init__(self, dim: int, eps: float):\n", " super().__init__()\n", - " self.eps = eps\n", - " self.weight = nn.Parameter(torch.ones(dim))\n", + " self.eps = eps # 防止分母等于零导致结果上溢\n", + " self.weight = nn.Parameter(torch.ones(dim)) # 可学习参数\n", "\n", - " def forward(self, x):\n", - " return self.weight * (x.float() * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)).type_as(x)" + " def forward(self, x):\n", + " return self.weight * (x.float() * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)).type_as(x)" ] }, { @@ -87,35 +100,47 @@ "source": [ "### Rotary Position Embedding, RoPE\n", "\n", - "旋转位置编码是一种能将相对位置信息集成到 self-attentio 中, 进而提升 transformer 架构性能的位置编码方式, 和绝对位置编码相比, RoPE 具有很好的外推性, 是目前的主流位置编码方式。\n", + "旋转位置编码是一种能将相对位置信息集成到 self-attention 中, 进而提升 transformer 架构性能的位置编码方式, 和绝对位置编码相比, RoPE 具有很好的外推性, 是目前的主流位置编码方式.\n", "\n", - "外推性的解释, 通俗来说就是训练的时候限制了 512 的上下文长度，那么推理时如果面对超过该长度的文本，LLM 可能无法正确处理。\n", + "外推性的解释, 通俗来说就是训练的时候限制了 512 的上下文长度，那么推d理时如果面对超过该长度的文本，LLM 可能无法正确处理.\n", "\n", - "- 绝对位置编码\n", + "- **绝对位置编码**\n", "\n", - "$$f_{t:t\\in\\{q,k,v\\}}(\\boldsymbol{x}_i,i):=\\boldsymbol{W}_{t:t\\in\\{q,k,v\\}}(\\boldsymbol{x}_i+\\boldsymbol{p}_i)$$\n", + "绝对位置编码是早期 Transformer 架构采用的绝对位置编码方案，及那个每个位置映射为固定的向量表示.\n", "\n", - "其中编码向量 $p_i$ 的计算使用 Sinusoidal 函数：\n", + "$$f_{t:t\\in\\{q,k,v\\}}(\\boldsymbol{x}_i,i)=\\boldsymbol{W}_{t:t\\in\\{q,k,v\\}}(\\boldsymbol{x}_i+\\boldsymbol{p}_i)$$\n", + "\n", + "其中编码向量 $p_i$ 的计算使用如下公式：\n", "\n", "$$\\boldsymbol{p}_{i,2t}=\\sin\\left(k/1000^{2t/d}\\right), \\boldsymbol{p}_{i,2t+1}=\\cos\\left(k/1000^{2t/d}\\right)$$\n", "\n", - "正如其名，绝对位置编码只考虑了输入序列中的绝对位置关系，对于 token 之间的相对信息则没有纳入考虑。\n", + "正如其名，绝对位置编码只考虑了输入序列中的绝对位置关系，对于 token 之间的相对信息则没有纳入考虑.\n", "\n", - "- 旋转位置编码\n", + "- **旋转位置编码**\n", "\n", - "假定 query 和 key 的内积操作可以被函数 g 表示（有点像核函数），该函数 g 的输入是词嵌入向量 $x_m, x_n$ 和它们之间的相对位置 $m-n$:\n", + "假定 query 和 key 的内积操作可以被函数 g 表示，该函数 g 的输入是词嵌入向量 $x_m, x_n$ 和它们之间的相对位置 $m-n$:\n", "\n", "$$=g(x_m, x_n, m, n)$$\n", "\n", - "旋转位置编码就是找到一个使上式成立的位置编码方式。\n", + "旋转位置编码就是找到一个使上式成立的位置编码方式. \n", "\n", - "$$f_{\\{q,k\\}}\\left(\\boldsymbol{x}_m,m\\right)=\\boldsymbol{R}_{\\Theta,m}^d\\boldsymbol{W}_{\\{q,k\\}}\\boldsymbol{x}_m$$\n", + "出于认识的目的，我们省略复杂的数学推导，直接看 RoPE 的的结论：\n", "\n", - "$$\\boldsymbol{R}_{\\Theta,m}^d=\\underbrace{\\begin{pmatrix}\\cos m\\theta_0&-\\sin m\\theta_0&0&0&\\cdots&0&0\\\\\\sin m\\theta_0&\\cos m\\theta_0&0&0&\\cdots&0&0\\\\0&0&\\cos m\\theta_1&-\\sin m\\theta_1&\\cdots&0&0\\\\0&0&\\sin m\\theta_1&\\cos m\\theta_1&\\cdots&0&0\\\\\\vdots&\\vdots&\\vdots&\\vdots&\\ddots&\\vdots&\\vdots\\\\0&0&0&0&\\cdots&\\cos m\\theta_{d/2-1}&-\\sin m\\theta_{d/2-1}&-\\sin m\\theta_{d/2-1}\\end{pmatrix}}_{\\boldsymbol{W}_m}(13)$$\n", + "存在这样一个正交矩阵：\n", "\n", - "$$\\Theta=\\left\\{\\theta_i=10000^{-2(i-1)/d},i\\in[1,2,\\ldots,d/2]\\right\\}$$\n", + "$$\\boldsymbol{R}_{\\Theta,m}^d=\\underbrace{\\begin{pmatrix}\\cos m\\theta_0&-\\sin m\\theta_0&0&0&\\cdots&0&0\\\\\\sin m\\theta_0&\\cos m\\theta_0&0&0&\\cdots&0&0\\\\0&0&\\cos m\\theta_1&-\\sin m\\theta_1&\\cdots&0&0\\\\0&0&\\sin m\\theta_1&\\cos m\\theta_1&\\cdots&0&0\\\\\\vdots&\\vdots&\\vdots&\\vdots&\\ddots&\\vdots&\\vdots\\\\0&0&0&0&\\cdots&\\cos m\\theta_{d/2-1}&-\\sin m\\theta_{d/2-1}&-\\sin m\\theta_{d/2-1}\\end{pmatrix}}_{\\boldsymbol{W}_m}$$\n", "\n", - "由于 $\\boldsymbol{R}_{\\Theta,m}^d$ 的稀疏性，直接使用矩阵乘法会浪费算力，因此采用下述方式实现：\n", + "其中，$\\Theta=\\left\\{\\theta_i=10000^{-2(i-1)/d},i\\in[1,2,\\ldots,d/2]\\right\\}$\n", + "\n", + "我们可以将 query 和 key 的内积操作转换为与原始向量 $x$ 相关的以下等价形式：\n", + "\n", + "$$\n", + "\\boldsymbol{q}_m^\\mathbf{T}\\boldsymbol{k}_n=\\left(\\boldsymbol{R}_{\\Theta,m}^d\\boldsymbol{W}_q\\boldsymbol{x}_m\\right)^\\mathbf{T}\\left(\\boldsymbol{R}_{\\Theta,n}^d\\boldsymbol{W}_k\\boldsymbol{x}_n\\right)=\\boldsymbol{x}_m^\\mathbf{T}\\boldsymbol{W}_q\\boldsymbol{R}_{\\Theta,n-m}^d\\boldsymbol{W}_k\\boldsymbol{x}_n\n", + "$$\n", + "\n", + "其中， $\\boldsymbol{R}_{\\Theta,n-m}^d=\\left(\\boldsymbol{R}_{\\Theta,m}^d\\right)^\\mathbf{T}\\boldsymbol{R}_{\\Theta,n}^d$.\n", + "\n", + "由于 $\\boldsymbol{R}_{\\Theta,m}^d$ 的稀疏性，直接使用矩阵乘法会浪费算力，因此代码中采用下述方式实现：\n", "\n", "$$\\boldsymbol{R}_{\\Theta,m}^{d}\\boldsymbol{x}=\\begin{pmatrix}x_{0}\\\\x_{1}\\\\x_{2}\\\\x_{3}\\\\\\vdots\\\\x_{d-2}\\\\x_{d-1}\\end{pmatrix}\\otimes\\begin{pmatrix}\\cos m\\theta_{0}\\\\\\cos m\\theta_{0}\\\\\\cos m\\theta_{1}\\\\\\cos m\\theta_{1}\\\\\\vdots\\\\\\cos m\\theta_{d/2-1}\\\\\\cos m\\theta_{d/2-1}\\end{pmatrix}+\\begin{pmatrix}-x_{1}\\\\x_{0}\\\\-x_{3}\\\\x_{2}\\\\\\vdots\\\\-x_{d-1}\\\\x_{d-2}\\end{pmatrix}\\otimes\\begin{pmatrix}\\sin m\\theta_{0}\\\\\\sin m\\theta_{0}\\\\\\sin m\\theta_{1}\\\\\\sin m\\theta_{1}\\\\\\vdots\\\\\\sin m\\theta_{d/2-1}\\\\\\sin m\\theta_{d/2-1}\\end{pmatrix}\n", "$$" @@ -128,34 +153,51 @@ "metadata": {}, "outputs": [], "source": [ - "# compute pe in complex form\n", - "def precompute_pos_cis(dim: int, end: int=int(32 * 1024), theta: float=1e6):\n", - " \"\"\"\n", - " compute pe frequency in complex form.\n", - " \"\"\"\n", + "def precompute_pos_cis(dim: int, end: int=int(32 * 1024), theta: float=1e6):\n", + " r\"\"\"生成旋转矩阵\"\"\"\n", + " # 计算词向量元素两两分组后 每组元素对应的旋转角度 \\theta_i 也就是上文公式中对应的旋转角度集合\n", " freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))\n", - " t = torch.arange(end, device=freqs.device)\n", - " freqs = torch.outer(t, freqs).float()\n", - " pos_cis = torch.polar(torch.ones_like(freqs), freqs)\n", + " # token 序列索引 t = [0, 1, ..., seq_len-1]\n", + " t = torch.arange(end, device=freqs.device)\n", + " # 计算 m * \\theta\n", + " freqs = torch.outer(t, freqs).float()\n", + " # 计算结果是个复数向量\n", + " # 假设 freqs = [x, y]\n", + " # 则 freqs_cis = [cos(x) + sin(x)i, cos(y) + sin(y)i]\n", + " pos_cis = torch.polar(torch.ones_like(freqs), freqs)\n", " return pos_cis\n", "\n", - "def apply_rotary_emb(xq, xk, pos_cis):\n", - " def unite_shape(pos_cis, x):\n", + "def apply_rotary_emb(xq, xk, pos_cis):\n", + " r\"\"\"应用 RoPE\"\"\"\n", + " def unite_shape(pos_cis, x):\n", " ndim = x.ndim\n", " assert 0 <= 1 < ndim\n", - " assert pos_cis.shape == (x.shape[1], x.shape[-1])\n", - " shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]\n", + " assert pos_cis.shape == (x.shape[1], x.shape[-1])\n", + " shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]\n", " return pos_cis.view(*shape)\n", "\n", - " xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))\n", - " print('original shape xq: {}'.format(xq.shape)) # 查看 xq 的 shape\n", - " print('ajusted shape xq: {}'.format(xq_.shape)) # 查看经过复数表示后 xq 的 shape\n", - " xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))\n", - " pos_cis = unite_shape(pos_cis, xq_)\n", - " print('ajusted shape pos_cis： {}'.format(pos_cis.shape)) # 查看位置编码旋转角的形状\n", + " xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))\n", + " # print('original shape xq: {}'.format(xq.shape)) # 取消注释以查看 xq 的 shape\n", + " # print('ajusted shape xq: {}'.format(xq_.shape)) # 取消注释查看经过复数表示后 xq 的 shape\n", + " xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))\n", + " pos_cis = unite_shape(pos_cis, xq_)\n", + " # print('ajusted shape pos_cis： {}'.format(pos_cis.shape)) # 取消注释以查看位置编码旋转角的形状\n", + " # 保留实部\n", " xq_out = torch.view_as_real(xq_ * pos_cis).flatten(3)\n", " xk_out = torch.view_as_real(xk_ * pos_cis).flatten(3)\n", - " return xq_out.type_as(xq), xk_out.type_as(xk)" + " return xq_out.type_as(xq), xk_out.type_as(xk)" + ] + }, + { + "cell_type": "markdown", + "id": "eff57879-1f39-404d-9501-3c99e2e67c41", + "metadata": {}, + "source": [ + "我们知道，RoPE 是在 Attention 阶段生成 Query 和 Key 向量后，对这两个向量进行位置编码的.\n", + "\n", + "对于 MiniMindLM，嵌入维度为 512，注意力头数量为 8，故每一个注意力头的维度应该为 512 / 8 = 64.\n", + "\n", + "我们使用固定形状的张量代表 Query 和 Key 向量，对 RoPE 的应用过程进行观察." ] }, { @@ -165,20 +207,17 @@ "metadata": {}, "outputs": [ { - "data": { - "text/plain": [ - "(torch.Size([16, 32]), tensor(1.+0.j))" - ] - }, - "execution_count": 4, - "metadata": {}, - "output_type": "execute_result" + "name": "stdout", + "output_type": "stream", + "text": [ + "pos_cis 的形状为 torch.Size([16, 32]), 其中 [0, 0] 下标元素为 (1+0j)\n" + ] } ], "source": [ - "xq, xk = torch.randn((2, 16, 4, 64)), torch.randn((2, 16, 4, 64)) # (batch_size, sequence_length, num_heads, head_dim)\n", - "pos_cis = precompute_pos_cis(64, 16) # 计算旋转位置编码的旋转角（复数表示）\n", - "pos_cis.shape, pos_cis[0, 0]" + "xq, xk = torch.randn((2, 16, 4, 64)), torch.randn((2, 16, 4, 64)) # (batch_size, sequence_length, num_heads, head_dim)\n", + "pos_cis = precompute_pos_cis(64, 16) # 计算旋转位置编码的旋转角（复数表示）\n", + "print(f'pos_cis 的形状为 {pos_cis.shape}, 其中 [0, 0] 下标元素为 {pos_cis[0, 0]}')" ] }, { @@ -191,35 +230,16 @@ "name": "stdout", "output_type": "stream", "text": [ - "original shape xq: torch.Size([2, 16, 4, 64])\n", - "ajusted shape xq: torch.Size([2, 16, 4, 32])\n", - "ajusted shape pos_cis： torch.Size([1, 16, 1, 32])\n" + "经过 RoPE 编码后的 Query 与 Key 的形状为 torch.Size([2, 16, 4, 64]), torch.Size([2, 16, 4, 64])\n" ] } ], "source": [ - "xq_rope, xk_rope = apply_rotary_emb(xq, xk, pos_cis)" - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "id": "c03dd09c-9ccf-4716-bbd3-58a1314c05d5", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "(torch.Size([2, 16, 4, 64]), torch.Size([2, 16, 4, 64]))" - ] - }, - "execution_count": 6, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "xq_rope.shape, xk_rope.shape" + "xq_rope, xk_rope = apply_rotary_emb(xq, xk, pos_cis)\n", + "# original shape xq: torch.Size([2, 16, 4, 64])\n", + "# ajusted shape xq: torch.Size([2, 16, 4, 32])\n", + "# ajusted shape pos_cis： torch.Size([1, 16, 1, 32])\n", + "print(f'经过 RoPE 编码后的 Query 与 Key 的形状为 {xq_rope.shape}, {xk_rope.shape}')" ] }, { @@ -229,57 +249,86 @@ "source": [ "### Attention\n", "\n", - "Attention 是 Transformer 架构中的重要机制，它可以捕获长序列间的依赖关系，并依据注意力得分进行重要性建模。\n", + "注意力机制（Attention Mechanism）是Transformer架构的核心组件，能够有效捕捉长序列内各元素间的依赖关系. 该机制通过计算输入序列中不同位置元素间的注意力得分，对其重要性进行精准建模，使模型在处理信息时能够聚焦于关键部分，从而显著提升对长序列数据的理解与处理能力.\n", "\n", - "在 MiniMind 模型中，Attention Block 涉及几个重要机制。\n", + "在 MiniMindLM 模型中, Attention Block 包含以下机制和模块:\n", "\n", "1. GQA\n", "2. KV Cache\n", - "4. Pre-RMSNorm & RoPE\n", "5. SwiGLU\n", - "6. FFN\n", "\n", - "- GQA\n", + "- **GQA**\n", "\n", - "Group Querey Attention (GQA) 是对多头自注意力机制的扩展，通过提供计算效率和模型表达能力的灵活权衡，实现了查询头的分组。\n", + "分组查询注意力 (Group Querey Attention, GQA) 是对多头自注意力机制的扩展, 通过提供计算效率和模型表达能力的灵活权衡, 实现了查询头的分组.\n", "\n", - "具体来说，GQA 将查询头分为 G 个组，每个组共享一个公共的键和值。\n", + "具体来说, GQA 中将 h 个查询头分为 G 组, 每组 包含 h / G 个查询头，并共享一个公共的键和值.\n", "\n", "![img](./demo/gqa.png)\n", "\n", - "- KV Cache\n", + "GQA 相比传统的 MHA， 减少了键和值的数量，降低了计算量和内存开销，提高了推理速度.\n", "\n", - "KV Cache 能够有效压缩大模型推理时的显存占用，在推理时，前面生成的字符不需要与后面的字符产生 attention，从而使得前面已经计算的 K 和 V 可以缓存起来。\n", + "- **KV Cache**\n", "\n", - "- SwiGLU\n", + "在语言模型生成文本的过程中，每生成一个新的 token，模型都需要计算注意力得分，以确定当前位置与之前所有位置的相关性.\n", "\n", - "$$\\text{SwiGLU}(x,W,V,b,c)=\\text{Swish}_1(xW+b)\\otimes(xV+c)$$" + "比如以下内容：\n", + "\n", + "1. *seq = \\[tok1]:*\n", + "\n", + " attn_11 = softmax(Q1 * K1.T / sqrt(dim)) * V1\n", + " \n", + "3. *seq = \\[tok1, tok2]:*\n", + "\n", + " attn_11 = softmax(Q1 * K1.T / sqrt(dim)) * V1, attn_12 = 0 (masked)\n", + " \n", + " attn_21 = softmax(Q2 * K1.T / sqrt(dim)) * V2, attn_22 = softmax(Q2 * K2.T / sqrt(dim)) * V2\n", + "4. *seq = \\[tok1, tok2, tok3]:*\n", + "\n", + " attn_11 = softmax(Q1 * K1.T / sqrt(dim)) * V1, attn_12 = 0 (masked), attn_13 = 0 (masked)\n", + " \n", + " attn_21 = softmax(Q2 * K1.T / sqrt(dim)) * V2, attn_22 = softmax(Q2 * K2.T / sqrt(dim)) * V2, attn_23 = 0 (masked)\n", + " \n", + " attn_31 = softmax(Q3 * K1.T / sqrt(dim)) * V3, attn_32 = softmax(Q3 * K2.T / sqrt(dim)) * V3, attn_33 = softmax(Q3 * K3.T / sqrt(dim)) * V3\n", + "5. ··· ···\n", + "\n", + "不难发现，大模型生成一个 token 后的注意力计算中，总会用到 token 序列的历史 KV 值，导致重复计算，KV Cache 的设计正是为了通过缓存历史 KV 值，节省计算开销.\n", + "\n", + "KV Cache 能够有效压缩大模型推理时的显存占用.\n", + "\n", + "- **SwiGLU**\n", + "\n", + "SwiGLU 是一种在深度学习中用于神经网络架构的激活函数变体：\n", + "\n", + "$$\\text{SwiGLU}(x, W, V, b, c)=\\text{Swish}_1(xW+b)\\otimes(xV+c)$$\n", + "\n", + "与传统的 ReLU 激活函数相比，SwiGLU 具有更好的平滑性和非线性表达能力，由于其门控机制，在处理信息筛选和流动方面有独特的优势." ] }, { "cell_type": "code", - "execution_count": 7, + "execution_count": 6, "id": "e61394db-5afa-4421-ba77-30a9c50f0603", "metadata": {}, "outputs": [], "source": [ "from demo.LMConfig import LMConfig\n", - "from typing import Any, Optional, Tuple, List\n", + "from typing import Any, Optional, Tuple, List\n", "import torch.nn as nn\n", "import math\n", "\n", - "def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:\n", - " bs, slen, n_kv_heads, head_dim = x.shape\n", + "def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:\n", + " r\"\"\"使得 KV 头数适应 Query 头数\"\"\"\n", + " bs, slen, n_kv_heads, head_dim = x.shape\n", " if n_rep == 1:\n", " return x\n", " return (\n", - " x[:, :, :, None, :]\n", - " .expand(bs, slen, n_kv_heads, n_rep, head_dim)\n", - " .reshape(bs, slen, n_kv_heads * n_rep, head_dim)\n", + " x[:, :, :, None, :]\n", + " .expand(bs, slen, n_kv_heads, n_rep, head_dim)\n", + " .reshape(bs, slen, n_kv_heads * n_rep, head_dim)\n", " )\n", "\n", "class Attention(nn.Module):\n", - " def __init__(self, args: LMConfig):\n", + " def __init__(self, args: LMConfig):\n", " super().__init__()\n", " self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads\n", " assert args.n_heads % self.n_kv_heads == 0\n", @@ -287,59 +336,79 @@ " self.n_local_kv_heads = args.n_kv_heads\n", " self.n_rep = self.n_local_heads // self.n_local_kv_heads\n", " self.head_dim = args.dim // args.n_heads\n", - " self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)\n", - " self.wk = nn.Linear(args.dim, args.n_kv_heads * self.head_dim, bias=False)\n", - " self.wv = nn.Linear(args.dim, args.n_kv_heads * self.head_dim, bias=False)\n", - " self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)\n", + " # q k v o projection\n", + " self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)\n", + " self.wk = nn.Linear(args.dim, args.n_kv_heads * self.head_dim, bias=False)\n", + " self.wv = nn.Linear(args.dim, args.n_kv_heads * self.head_dim, bias=False)\n", + " self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)\n", " self.attn_dropout = nn.Dropout(args.dropout)\n", " self.resid_dropout = nn.Dropout(args.dropout)\n", " self.dropout = args.dropout\n", - " self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn\n", - " mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float(\"-inf\"))\n", - " mask = torch.triu(mask, diagonal=1)\n", - " self.register_buffer(\"mask\", mask, persistent=False)\n", + " self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn\n", + " mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float(\"-inf\"))\n", + " mask = torch.triu(mask, diagonal=1)\n", + " self.register_buffer(\"mask\", mask, persistent=False)\n", "\n", - " def forward(self,\n", - " x: torch.Tensor,\n", - " pos_cis: torch.Tensor,\n", - " past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,\n", + " def forward(self, \n", + " x: torch.Tensor, \n", + " pos_cis: torch.Tensor, \n", + " past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, \n", " use_cache=False):\n", - " bsz, seq_len, _ = x.shape\n", + " bsz, seq_len, _ = x.shape\n", " ############## Forward QKV & RoPE ##############\n", - " xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)\n", - " xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)\n", - " xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)\n", - " xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)\n", - " xq, xv = apply_rotary_emb(xq, xk, pos_cis)\n", + " xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)\n", + " xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)\n", + " xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)\n", + " xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)\n", + " xq, xv = apply_rotary_emb(xq, xk, pos_cis)\n", " ################### KV Cache ###################\n", " if past_key_value is not None:\n", - " xk = torch.cat([past_key_value[0], xk], dim=1)\n", - " xv = torch.cat([past_key_value[1], xv], dim=1)\n", - " past_kv = (xk, xv) if use_cache else None\n", - " xq, xk, xv = (\n", - " xq.transpose(1, 2),\n", - " repeat_kv(xk, self.n_rep).transpose(1, 2),\n", - " repeat_kv(xv, self.n_rep).transpose(1, 2)\n", + " xk = torch.cat([past_key_value[0], xk], dim=1)\n", + " xv = torch.cat([past_key_value[1], xv], dim=1)\n", + " past_kv = (xk, xv) if use_cache else None\n", + " xq, xk, xv = (\n", + " xq.transpose(1, 2), \n", + " repeat_kv(xk, self.n_rep).transpose(1, 2), \n", + " repeat_kv(xv, self.n_rep).transpose(1, 2)\n", " )\n", " ############ Scaled Dot Production #############\n", " if self.flash and seq_len != 1:\n", " dropout_p = self.dropout if self.training else 0.0\n", " output = F.scaled_dot_product_attention(\n", - " xq, xk, xv,\n", - " attn_mask=None,\n", - " dropout_p=dropout_p,\n", + " xq, xk, xv, \n", + " attn_mask=None, \n", + " dropout_p=dropout_p, \n", " is_causal=True\n", " )\n", " else:\n", - " scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)\n", - " scores += self.mask[:, :, :seq_len, :seq_len]\n", - " scores = F.softmax(scores.float(), dim=-1).type_as(xq)\n", + " scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)\n", + " scores += self.mask[:, :, :seq_len, :seq_len]\n", + " scores = F.softmax(scores.float(), dim=-1).type_as(xq)\n", " scores = self.attn_dropout(scores)\n", " output = scores @ xv\n", " ################################################\n", - " output = output.transpose(1, 2).reshape(bsz, seq_len, -1)\n", + " output = output.transpose(1, 2).reshape(bsz, seq_len, -1)\n", " output = self.resid_dropout(self.wo(output))\n", - " return output, past_kv" + " return output, past_kv" + ] + }, + { + "cell_type": "markdown", + "id": "7420b47b-b38e-442a-84a2-e307f82584d3", + "metadata": {}, + "source": [ + "同样的，我们假设一批 batch size = 4， seq len = 16 的 token 序列通过这个 Attention 块，在输入前，它的 input id 会被投影到 dim = 512 维." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "e1519451-3dc8-498e-96a7-db8b32a71af1", + "metadata": {}, + "outputs": [], + "source": [ + "# 创建 Config，为了简单起见，我们设定 num_layers = 2\n", + "LMConfig_Dense = LMConfig(n_layers=2)" ] }, { @@ -352,23 +421,28 @@ "name": "stdout", "output_type": "stream", "text": [ - "torch.Size([4, 16, 512])\n", - "torch.Size([16, 32])\n", - "original shape xq: torch.Size([4, 16, 8, 64])\n", - "ajusted shape xq: torch.Size([4, 16, 8, 32])\n", - "ajusted shape pos_cis： torch.Size([1, 16, 1, 32])\n", - "torch.Size([4, 16, 512]) torch.Size([4, 16, 2, 64]) torch.Size([4, 16, 2, 64])\n" + "输入张量 x ：size = torch.Size([4, 16, 512])，RoPE 旋转角： size = torch.Size([16, 32])\n", + "输出 output: size = torch.Size([4, 16, 512]), kv_cache 基本信息：size_key = torch.Size([4, 16, 2, 64]), size_value = torch.Size([4, 16, 2, 64])\n" ] } ], "source": [ - "attn = Attention(LMConfig())\n", - "x = torch.randn((4, 16, 512))\n", - "print(x.shape)\n", - "pos_cis = precompute_pos_cis(64, 16)\n", - "print(pos_cis.shape)\n", - "output, past_kv = attn(x, pos_cis=pos_cis, use_cache=True)\n", - "print(output.shape, past_kv[0].shape, past_kv[1].shape)" + "attn = Attention(LMConfig_Dense)\n", + "x = torch.randn((4, 16, 512)) # (batch size, seq len, embed dim)\n", + "pos_cis = precompute_pos_cis(64, 16) # (head dim, batch size) 其中 head dim = embed dim / num heads\n", + "output, past_kv = attn(x, pos_cis=pos_cis, use_cache=True)\n", + "print(f'输入张量 x ：size = {x.shape}，RoPE 旋转角： size = {pos_cis.shape}')\n", + "print(f'输出 output: size = {output.shape}, kv_cache 基本信息：size_key = {past_kv[0].shape}, size_value = {past_kv[1].shape}')" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "b53cee91-ee0d-4a61-8a7b-6feac5ba1b40", + "metadata": {}, + "outputs": [], + "source": [ + "del attn" ] }, { @@ -378,54 +452,77 @@ "source": [ "### FeedForward Network\n", "\n", - "前向传播神经网络接收来自注意力层的输出，并对其做进一步的线性变换，以捕获更复杂的特征和表示。" - ] - }, - { - "cell_type": "code", - "execution_count": 9, - "id": "a16f3ece-3b4e-42aa-9b88-1a7193dd2bf0", - "metadata": {}, - "outputs": [], - "source": [ - "class FeedForward(nn.Module):\n", - " def __init__(self, config: LMConfig):\n", - " super().__init__()\n", - " if config.hidden_dim is None:\n", - " hidden_dim = 4 * config.dim\n", - " hidden_dim = int(2 * hidden_dim / 3)\n", - " config.hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of)\n", - " self.w1 = nn.Linear(config.dim, config.hidden_dim, bias=False)\n", - " self.w2 = nn.Linear(config.hidden_dim, config.dim, bias=False)\n", - " self.w3 = nn.Linear(config.dim, config.hidden_dim, bias=False)\n", - " self.dropout = nn.Dropout(config.dropout)\n", - "\n", - " def forward(self, x):\n", - " return self.dropout(self.w2(F.silu(self.w1(x) * self.w3(x))))" + "前馈神经网络接收来自注意力机制层的输出结果，随后对该输出执行进一步的线性变换. 通过这种方式，网络能够深入挖掘并捕获更为复杂、抽象的特征." ] }, { "cell_type": "code", "execution_count": 10, + "id": "a16f3ece-3b4e-42aa-9b88-1a7193dd2bf0", + "metadata": {}, + "outputs": [], + "source": [ + "class FeedForward(nn.Module):\n", + " def __init__(self, config: LMConfig):\n", + " super().__init__()\n", + " if config.hidden_dim is None:\n", + " hidden_dim = 4 * config.dim\n", + " hidden_dim = int(2 * hidden_dim / 3)\n", + " config.hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of)\n", + " self.w1 = nn.Linear(config.dim, config.hidden_dim, bias=False)\n", + " self.w2 = nn.Linear(config.hidden_dim, config.dim, bias=False)\n", + " self.w3 = nn.Linear(config.dim, config.hidden_dim, bias=False)\n", + " self.dropout = nn.Dropout(config.dropout)\n", + "\n", + " def forward(self, x):\n", + " return self.dropout(self.w2(F.silu(self.w1(x) * self.w3(x))))" + ] + }, + { + "cell_type": "markdown", + "id": "8ddb2503-018c-4486-ae45-d9d0a192490a", + "metadata": {}, + "source": [ + "设置输入 x 为 batch size = 4，seq len = 16 的 token 序列投影向量，观察 x 在 MiniMind Block 的前向传播过程." + ] + }, + { + "cell_type": "code", + "execution_count": 11, "id": "4fe2abd9-0c40-44be-b64d-8b0caa28604b", "metadata": {}, "outputs": [ { - "data": { - "text/plain": [ - "torch.Size([4, 16, 512])" - ] - }, - "execution_count": 10, - "metadata": {}, - "output_type": "execute_result" + "name": "stdout", + "output_type": "stream", + "text": [ + "给定输入 x: size = torch.Size([4, 16, 512]) 下的输出 output：size = torch.Size([4, 16, 512])\n" + ] } ], "source": [ - "ffn = FeedForward(LMConfig())\n", - "x = torch.randn((4, 16, 512))\n", + "ffn = FeedForward(LMConfig_Dense)\n", + "x = torch.randn((4, 16, 512))\n", "output = ffn(x)\n", - "output.shape" + "print(f'给定输入 x: size = {x.shape} 下的输出 output：size = {output.shape}')" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "5ae0c7e5-67e9-4dc4-9ea5-4ec20667c634", + "metadata": {}, + "outputs": [], + "source": [ + "del ffn" + ] + }, + { + "cell_type": "markdown", + "id": "efe25d3b-2af9-4fe1-8eb0-7522324ec2c6", + "metadata": {}, + "source": [ + "这就是 Transformer 结构的精妙之处，张量在模型内部进行了各种复杂的投影变形，但是输入输出的张量形状不发生改变!" ] }, { @@ -435,44 +532,52 @@ "source": [ "### MiniMind Block\n", "\n", - "到目前为止，，已经完成了 Attention Layer 和 FeedForward Layer 的构建，所有必须的组件都已经具备，我们着手构建一个 MiniMind Block" + "到目前为止, , 已经完成了 Attention Layer 和 FeedForward Layer 的构建, 所有必须的组件都已经具备, 我们着手构建一个 MiniMind Block" ] }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 13, "id": "d6c224a4-7ba5-4cf8-9297-311c660f5764", "metadata": {}, "outputs": [], "source": [ "class MiniMindBlock(nn.Module):\n", - " def __init__(self, layer_id: int, config: LMConfig):\n", + " def __init__(self, layer_id: int, config: LMConfig):\n", " super().__init__()\n", " self.n_heads = config.n_heads\n", " self.dim = config.dim\n", - " self.head_dim = config.dim // config.n_heads # number of GQA heads\n", + " self.head_dim = config.dim // config.n_heads # head dim\n", " self.attention = Attention(config)\n", "\n", - " self.layer_id = layer_id\n", - " self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)\n", - " self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)\n", + " self.layer_id = layer_id # layer id 指向其维护的 KV Cache\n", + " self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)\n", + " self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)\n", " self.feed_forward = FeedForward(config)\n", "\n", - " def forward(self, x, pos_cis, past_key_value=None, use_cache=False):\n", - " h_attn, past_kv = self.attention(\n", - " self.attention_norm(x), # pre-normed x\n", - " pos_cis,\n", - " past_key_value=past_key_value,\n", + " def forward(self, x, pos_cis, past_key_value=None, use_cache=False):\n", + " h_attn, past_kv = self.attention(\n", + " self.attention_norm(x), # pre-normed x\n", + " pos_cis, \n", + " past_key_value=past_key_value, \n", " use_cache=use_cache\n", " )\n", " h = x + h_attn # residual connection\n", " out = h + self.feed_forward(self.ffn_norm(h)) # feed forward + residual connection\n", - " return out, past_kv" + " return out, past_kv" + ] + }, + { + "cell_type": "markdown", + "id": "f627c780-f07e-41c3-8b44-7a819e80a2a4", + "metadata": {}, + "source": [ + "我们依然设置输入 x 为 batch size = 4，seq len = 16 的 token 序列投影向量，观察 x 在 MiniMind Block 的前向传播过程." ] }, { "cell_type": "code", - "execution_count": 12, + "execution_count": 14, "id": "d23b4461-534c-4f6d-a00b-cdcfbd4d7bec", "metadata": {}, "outputs": [ @@ -480,30 +585,27 @@ "name": "stdout", "output_type": "stream", "text": [ - "original shape xq: torch.Size([4, 16, 8, 64])\n", - "ajusted shape xq: torch.Size([4, 16, 8, 32])\n", - "ajusted shape pos_cis： torch.Size([1, 16, 1, 32])\n" + "输出 output 信息: size = torch.Size([4, 16, 512])\n", + "该 Block 维护的 KV Cache 信息：size_key = torch.Size([4, 16, 2, 64]), size_value = torch.Size([4, 16, 2, 64])\n" ] - }, - { - "data": { - "text/plain": [ - "(torch.Size([4, 16, 512]),\n", - " torch.Size([4, 16, 2, 64]),\n", - " torch.Size([4, 16, 2, 64]))" - ] - }, - "execution_count": 12, - "metadata": {}, - "output_type": "execute_result" } ], "source": [ - "miniblock = MiniMindBlock(1, LMConfig())\n", - "x = torch.randn((4, 16, 512))\n", - "pos_cis = precompute_pos_cis(64, 16)\n", - "out, past_kv = miniblock(x, pos_cis, use_cache=True)\n", - "out.shape, past_kv[0].shape, past_kv[1].shape" + "miniblock = MiniMindBlock(1, LMConfig_Dense)\n", + "x = torch.randn((4, 16, 512))\n", + "pos_cis = precompute_pos_cis(64, 16)\n", + "out, past_kv = miniblock(x, pos_cis, use_cache=True)\n", + "print(f'输出 output 信息: size = {out.shape}\\n该 Block 维护的 KV Cache 信息：size_key = {past_kv[0].shape}, size_value = {past_kv[1].shape}')" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "bb9745cf-2010-49b3-90ad-a1d59f76de89", + "metadata": {}, + "outputs": [], + "source": [ + "del miniblock" ] }, { @@ -511,14 +613,14 @@ "id": "f9bb3b2e-4d1d-4721-9908-546f3c49dde3", "metadata": {}, "source": [ - "### MiniMindLM\n", + "### MiniMindLM (Dense)\n", "\n", - "以 MiniMind Block 为基本组件，我们可以对该 LLM 进行最后组装！" + "以 MiniMind Block 为基本组件, 我们对 MiniMindLM 进行最后组装！" ] }, { "cell_type": "code", - "execution_count": 13, + "execution_count": 16, "id": "7267245b-f13a-4bcf-9358-5fdd89106153", "metadata": {}, "outputs": [], @@ -526,108 +628,119 @@ "from transformers import PreTrainedModel\n", "from transformers.modeling_outputs import CausalLMOutputWithPast\n", "\n", - "class MOEFeedForward():\n", - " \"\"\"Place holder\"\"\"\n", - " pass\n", - "\n", "class MiniMindLM(PreTrainedModel):\n", - " \"\"\"Notice: To simplify notebook run, set n_layers to 2\"\"\"\n", " config_class = LMConfig\n", "\n", - " def __init__(self, params: LMConfig = None):\n", + " def __init__(self, params: LMConfig = None):\n", " self.params = params or LMConfig()\n", " super().__init__(self.params)\n", - " self.vocab_size, self.n_layers = params.vocab_size, params.n_layers\n", + " self.vocab_size, self.n_layers = params.vocab_size, params.n_layers\n", " # 映射：词表维度 -> 嵌入维度\n", - " self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)\n", + " self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)\n", " self.dropout = nn.Dropout(params.dropout)\n", - " self.layers = nn.ModuleList([MiniMindBlock(1, params) for l in range(self.n_layers)])\n", - " self.norm = RMSNorm(params.dim, eps = params.norm_eps)\n", - " self.output = nn.Linear(params.dim, params.vocab_size, bias=False)\n", + " self.layers = nn.ModuleList([MiniMindBlock(1, params) for l in range(self.n_layers)])\n", + " self.norm = RMSNorm(params.dim, eps = params.norm_eps)\n", + " self.output = nn.Linear(params.dim, params.vocab_size, bias=False)\n", " self.tok_embeddings.weight = self.output.weight\n", " self.register_buffer(\n", - " \"pos_cis\",\n", - " precompute_pos_cis(dim=params.dim // params.n_heads, theta=params.rope_theta),\n", + " \"pos_cis\", \n", + " precompute_pos_cis(dim=params.dim // params.n_heads, theta=params.rope_theta), \n", " persistent=False\n", " )\n", " self.OUT = CausalLMOutputWithPast()\n", "\n", - " def forward(self,\n", - " input_ids: Optional[torch.Tensor] = None,\n", - " past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,\n", - " use_cache: bool = False,\n", + " def forward(self, \n", + " input_ids: Optional[torch.Tensor] = None, \n", + " past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None, \n", + " use_cache: bool = False, \n", " **args):\n", " past_key_values = past_key_values or [None] * len(self.layers) # save history KV\n", - " start_pos = args.get('start_pos', 0)\n", + " start_pos = args.get('start_pos', 0)\n", " h = self.dropout(self.tok_embeddings(input_ids))\n", " pos_cis = self.pos_cis[start_pos: start_pos + input_ids.size(1)]\n", - " past_kvs = []\n", - " for l, layer in enumerate(self.layers):\n", - " h, past_kv = layer(\n", - " h, pos_cis,\n", - " past_key_value=past_key_values[l],\n", + " past_kvs = [] # KV Cache 列表 按照索引维护对应 MiniMind Block 模块的 KV Cache\n", + " print(f'====> forward propagation started, num_minimind_blocks = {len(self.layers)}')\n", + " for l, layer in enumerate(self.layers):\n", + " print(f'------------> entering minimind block: id = {l}')\n", + " h, past_kv = layer(\n", + " h, pos_cis, \n", + " past_key_value=past_key_values[l], \n", " use_cache=use_cache\n", " )\n", + " print(f'<------------ finished, size_cache_k = {past_kv[0].shape}, size_cache_v = {past_kv[1].shape}')\n", " past_kvs.append(past_kv)\n", + " print(f'<==== forward propagation completed, num_kv_cache = {len(past_kvs)}\\n')\n", " logits = self.output(self.norm(h))\n", - " aux_loss = sum(l.feed_forward.aux_loss for l in self.layers if isinstance(l.feed_forward, MOEFeedForward)) # MOE 辅助损失\n", - " self.OUT.__setitem__('logits', logits)\n", - " self.OUT.__setitem__('aux_loss', aux_loss)\n", - " self.OUT.__setitem__('past_key_value', past_kvs)\n", + " aux_loss = 0 # MoE 辅助损失，我们在 Dense Model 中不纳入考虑\n", + " self.OUT.__setitem__('logits', logits) # token 分类 logits\n", + " self.OUT.__setitem__('aux_loss', aux_loss) # MoE 辅助损失\n", + " self.OUT.__setitem__('past_key_value', past_kvs) # 当前序列的历史 KV_cache 列表\n", " return self.OUT\n", "\n", " @torch.inference_mode()\n", - " def generate(self, input_ids, eos_token_id=2, max_new_tokens=512, temperature=0.75, top_p=0.90,\n", - " stream=False, rp=1, use_cache=True, pad_token_id=0, **args):\n", - " # 流式生成 （返回生成器，自由控制输出）\n", + " def generate(self, input_ids, eos_token_id=2, max_new_tokens=512, temperature=0.75, top_p=0.90, \n", + " stream=False, rp=1, use_cache=True, pad_token_id=0, **args):\n", + " # 流式生成 （返回生成器, 自由控制输出）\n", " if stream:\n", - " return self._stream(input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args)\n", + " return self._stream(input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args)\n", " # 直接生成 （一步到位）\n", " generated = []\n", " for i in range(input_ids.size(0)):\n", " non_pad = input_ids[i][input_ids[i] != pad_token_id].unsqueeze(0)\n", - " out = self._stream(non_pad, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args)\n", - " tokens_list = [tokens[:, -1:] for tokens in out]\n", - " gen = torch.cat(tokens_list, dim=-1) if tokens_list else non_pad\n", - " full_sequence = torch.cat([non_pad, gen], dim=-1)\n", + " out = self._stream(non_pad, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args)\n", + " tokens_list = [tokens[:, -1:] for tokens in out]\n", + " print(f'new tokens list :{tokens_list}\\n')\n", + " gen = torch.cat(tokens_list, dim=-1) if tokens_list else non_pad\n", + " full_sequence = torch.cat([non_pad, gen], dim=-1)\n", " generated.append(full_sequence)\n", " max_length = max(seq.size(1) for seq in generated)\n", " generated = [\n", - " torch.cat([seq, torch.full((1, max_length - seq.size(1)), pad_token_id, dtype=seq.dtype, device=seq.device)],dim=-1) \n", + " torch.cat([seq, torch.full((1, max_length - seq.size(1)), pad_token_id, dtype=seq.dtype, device=seq.device)], dim=-1) \n", " for seq in generated\n", " ]\n", - " return torch.cat(generated, dim=0)\n", + " return torch.cat(generated, dim=0)\n", "\n", - " def _stream(self, input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args):\n", - " start, first_seq, past_kvs = input_ids.shape[1], True, None\n", + " def _stream(self, input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args):\n", + " start, first_seq, past_kvs = input_ids.shape[1], True, None\n", + " new_token_idx = 0 # new token 计数器\n", " while input_ids.shape[1] < max_new_tokens - 1:\n", - " if first_seq or not use_cache: # 若第一次生成序列 or 无 KV Cache, 每次生成传入整个 token id 序列\n", - " out, first_seq = self(input_ids, past_key_values=past_kvs, use_cache=use_cache, **args), False\n", - " else: # 若非第一次生成 and 有 KV Cache，每次传入最后一个 token id 与 KV Cache 进行推理加速\n", - " out = self(input_ids[:, -1:], past_key_values=past_kvs, use_cache=use_cache,\n", - " start_pos=input_ids.shape[1] - 1, **args)\n", - " logits, past_kvs = out.logits[:, -1, :], out.past_key_values # logits.shape: (batch_size, seq_len, embed_dim)，获取最后一位 logits\n", - " logits[:, list(set(input_ids.tolist()[0]))] /= rp # 对生成 token 进行惩罚，降低后续重复生成几率\n", - " logits /= (temperature + 1e-9) # 调整温度，控制生成多样性\n", + " print(f'gernerating new token: idx = {start + new_token_idx}')\n", + " if first_seq or not use_cache: # 若第一次生成序列 or 无 KV Cache, 每次生成传入整个 token id 序列\n", + " out, first_seq = self(input_ids, past_key_values=past_kvs, use_cache=use_cache, **args), False\n", + " else: # 若非第一次生成 and 有 KV Cache, 每次传入最后一个 token id 与 KV Cache 进行推理加速\n", + " out = self(input_ids[:, -1:], past_key_values=past_kvs, use_cache=use_cache, \n", + " start_pos=input_ids.shape[1] - 1, **args)\n", + " logits, past_kvs = out.logits[:, -1, :], out.past_key_values # logits.shape: (batch_size, seq_len, embed_dim), 获取最后一位 logits\n", + " logits[:, list(set(input_ids.tolist()[0]))] /= rp # 对生成 token 进行惩罚, 降低后续重复生成几率\n", + " logits /= (temperature + 1e-9) # 调整温度, 控制生成多样性\n", " if top_p is not None and top_p < 1.0: # top-p 采样\n", - " sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)\n", - " sorted_probs = F.softmax(sorted_logits, dim=-1)\n", - " cumulative_probs = torch.cumsum(sorted_probs, dim=-1)\n", + " sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)\n", + " sorted_probs = F.softmax(sorted_logits, dim=-1)\n", + " cumulative_probs = torch.cumsum(sorted_probs, dim=-1)\n", " sorted_indices_to_remove = cumulative_probs > top_p\n", - " sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()\n", - " sorted_indices_to_remove[:, 0] = False\n", - " indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)\n", + " sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()\n", + " sorted_indices_to_remove[:, 0] = False\n", + " indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)\n", " logits[indices_to_remove] = -float('Inf')\n", - " input_ids_next = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1) # 从保留的 token 中采样\n", - " input_ids = torch.cat((input_ids, input_ids_next), dim=1)\n", - " yield input_ids[:, start:]\n", + " input_ids_next = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1) # 从保留的 token 中采样\n", + " input_ids = torch.cat((input_ids, input_ids_next), dim=1)\n", + " new_token_idx += 1\n", + " yield input_ids[:, start:]\n", " if input_ids_next.item() == eos_token_id:\n", " break" ] }, + { + "cell_type": "markdown", + "id": "bdee2ebc-f98c-49cc-b5fc-c6cd0056f4d9", + "metadata": {}, + "source": [ + "接下来，我们设置一条长度为 4 的 token 序列，使用 MiniMindLM 对其执行一次前向传播，观察传播过程与返回值." + ] + }, { "cell_type": "code", - "execution_count": 14, + "execution_count": 17, "id": "dc7b4d79-a3aa-4143-aa5c-2abce34bc96c", "metadata": {}, "outputs": [ @@ -635,32 +748,27 @@ "name": "stdout", "output_type": "stream", "text": [ - "original shape xq: torch.Size([1, 4, 8, 64])\n", - "ajusted shape xq: torch.Size([1, 4, 8, 32])\n", - "ajusted shape pos_cis： torch.Size([1, 4, 1, 32])\n" + "====> forward propagation started, num_minimind_blocks = 2\n", + "------------> entering minimind block: id = 0\n", + "<------------ finished, size_cache_k = torch.Size([1, 4, 2, 64]), size_cache_v = torch.Size([1, 4, 2, 64])\n", + "------------> entering minimind block: id = 1\n", + "<------------ finished, size_cache_k = torch.Size([1, 4, 2, 64]), size_cache_v = torch.Size([1, 4, 2, 64])\n", + "<==== forward propagation completed, num_kv_cache = 2\n", + "\n", + "返回 logits：size = torch.Size([1, 4, 6400]), 返回 aux_loss: 0, 返回 KV Cache List： len = 2\n" ] - }, - { - "data": { - "text/plain": [ - "(torch.Size([1, 4, 6400]), 0, 1)" - ] - }, - "execution_count": 14, - "metadata": {}, - "output_type": "execute_result" } ], "source": [ - "MiniMind = MiniMindLM(LMConfig())\n", - "input_ids = torch.Tensor([1, 3, 5, 7]).long().reshape(1, 4)\n", - "OUT = MiniMind(input_ids, use_cache=True)\n", - "OUT.logits.shape, OUT.aux_loss, len(OUT.past_key_value)" + "MiniMind_Dense = MiniMindLM(LMConfig_Dense)\n", + "input_ids = torch.Tensor([1, 3, 5, 7]).long().reshape(1, 4)\n", + "OUT = MiniMind_Dense(input_ids, use_cache=True)\n", + "print(f'返回 logits：size = {OUT.logits.shape}, 返回 aux_loss: {OUT.aux_loss}, 返回 KV Cache List： len = {len(OUT.past_key_value)}')" ] }, { "cell_type": "code", - "execution_count": 15, + "execution_count": 18, "id": "3480a860-e934-43de-a621-8063f0bdea2f", "metadata": {}, "outputs": [ @@ -668,23 +776,50 @@ "name": "stdout", "output_type": "stream", "text": [ - "original shape xq: torch.Size([1, 4, 8, 64])\n", - "ajusted shape xq: torch.Size([1, 4, 8, 32])\n", - "ajusted shape pos_cis： torch.Size([1, 4, 1, 32])\n", - "original shape xq: torch.Size([1, 1, 8, 64])\n", - "ajusted shape xq: torch.Size([1, 1, 8, 32])\n", - "ajusted shape pos_cis： torch.Size([1, 1, 1, 32])\n", - "original shape xq: torch.Size([1, 1, 8, 64])\n", - "ajusted shape xq: torch.Size([1, 1, 8, 32])\n", - "ajusted shape pos_cis： torch.Size([1, 1, 1, 32])\n", - "torch.Size([1, 7]) tensor([[ 1, 3, 5, 7, 4867, 4289, 6268]])\n" + "gernerating new token: idx = 4\n", + "====> forward propagation started, num_minimind_blocks = 2\n", + "------------> entering minimind block: id = 0\n", + "<------------ finished, size_cache_k = torch.Size([1, 4, 2, 64]), size_cache_v = torch.Size([1, 4, 2, 64])\n", + "------------> entering minimind block: id = 1\n", + "<------------ finished, size_cache_k = torch.Size([1, 4, 2, 64]), size_cache_v = torch.Size([1, 4, 2, 64])\n", + "<==== forward propagation completed, num_kv_cache = 2\n", + "\n", + "gernerating new token: idx = 5\n", + "====> forward propagation started, num_minimind_blocks = 2\n", + "------------> entering minimind block: id = 0\n", + "<------------ finished, size_cache_k = torch.Size([1, 1, 2, 64]), size_cache_v = torch.Size([1, 1, 2, 64])\n", + "------------> entering minimind block: id = 1\n", + "<------------ finished, size_cache_k = torch.Size([1, 1, 2, 64]), size_cache_v = torch.Size([1, 1, 2, 64])\n", + "<==== forward propagation completed, num_kv_cache = 2\n", + "\n", + "gernerating new token: idx = 6\n", + "====> forward propagation started, num_minimind_blocks = 2\n", + "------------> entering minimind block: id = 0\n", + "<------------ finished, size_cache_k = torch.Size([1, 1, 2, 64]), size_cache_v = torch.Size([1, 1, 2, 64])\n", + "------------> entering minimind block: id = 1\n", + "<------------ finished, size_cache_k = torch.Size([1, 1, 2, 64]), size_cache_v = torch.Size([1, 1, 2, 64])\n", + "<==== forward propagation completed, num_kv_cache = 2\n", + "\n", + "new tokens list :[tensor([[2747]]), tensor([[2065]]), tensor([[5211]])]\n", + "\n", + "生成结果：tensor([[ 1, 3, 5, 7, 2747, 2065, 5211]])\n" ] } ], "source": [ "# 我们让 MiniMind 根据我们设计的 四个输入 token 生成输出\n", - "out = MiniMind.generate(input_ids, max_new_tokens=8, use_cache=True)\n", - "print(out.shape, out)" + "out = MiniMind_Dense.generate(input_ids, max_new_tokens=8, use_cache=True)\n", + "print(f'生成结果：{out}')" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "id": "4509ff58-0439-42b5-b303-d308b502e821", + "metadata": {}, + "outputs": [], + "source": [ + "del MiniMind_Dense" ] }, { @@ -692,7 +827,7 @@ "id": "b0e4ef5f-1358-4ffb-a150-61141ac3062f", "metadata": {}, "source": [ - "## Minimind MOE Model" + "## Minimind MoE Model" ] }, { @@ -700,7 +835,566 @@ "id": "e54ced18-db3a-478b-9578-9825634f3d66", "metadata": {}, "source": [ - "@TODO" + "作者提供了 MiniMind MoE Model 的可视化.\n", + "\n", + "![image](../images/LLM-structure-moe.png)" + ] + }, + { + "cell_type": "markdown", + "id": "82682a98-cae7-49c2-be2f-12dd1adb8deb", + "metadata": {}, + "source": [ + "可以看到，Dense Model 和 MoE Model 的差异在于 FFN 层. MoE 模型将稠密连接的 FFN 层置换为 M x Expert 层，每次前向传播时只激活部分 Expert.\n", + "\n", + "其组成可以分为以下部分：\n", + "\n", + "- Experts: MoE 架构单元，每个专家本质上是一个独立的神经网络模块，负责处理特定类型或范围的数据.\n", + "- Router： 控制信息流动，决定每次前向传播激活的 Experts 模块以及输入数据在这些模块的分配组合.\n", + "\n", + "在 MoE 网络中，为了平衡专家的重要性，我们需要关注路由，它是决定在特定时间选择哪些专家的重要组件.\n", + "\n", + "### 辅助损失\n", + "\n", + "为了让训练过程中实现专家的更均匀分布，辅助损失（又称负载均衡损失）被添加到网络的常规损失中，它增加了一个约束，迫使专家具有相等的重要性.\n", + "\n", + "假设有输入序列 \\[What is Mixture of Experts], *Prob* (·) 表示每一个 token 激活的专家概率分布:\n", + "\n", + "- **Step 1**：在整个批次中对每个专家的路由值进行求和.\n", + "\n", + " Importance_per_token = *Prob* (What) + *Prob* (is) + *Prob* (Mixture) + *Prob* (of) + *Prob* (Experts)\n", + "\n", + " 这个指标反映了 batch 维度上每个专家的重要性分数.\n", + "\n", + "- **Step 2**: 计算变异系数\n", + "\n", + " 我们希望专家之间的重要性尽可能靠近，为了衡量专家得分之间的差异程度，引入变异系数指标（Coefficient Variation, CV）\n", + "\n", + " $$Coeifficient \\, Variation (CV) = \\dfrac{standard \\, deviation (\\sigma)}{mean(\\mu)}$$\n", + "\n", + " 如果专家的重要性分数相似，变异系数会降到很低（这是我们期望的）\n", + "\n", + "- **Step 3**: 计算负载均衡损失\n", + "\n", + " $$uxiliary Loss = \\alpha * CV$$\n", + "\n", + " 其中 $\\alpha$ 是缩放系数." + ] + }, + { + "cell_type": "markdown", + "id": "26dffcd0-7c9e-40bf-9c91-1636aa571d46", + "metadata": {}, + "source": [ + "## MoE Gate\n", + "\n", + "MoE 门控单元决定每次前向传播激活的 Experts 模块及其权重，同时计算 MoE 辅助损失." + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "id": "cba5532c-e8d2-40ff-9ed6-b8b33b4a31ca", + "metadata": {}, + "outputs": [], + "source": [ + "class MoEGate(nn.Module):\n", + " def __init__(self, config: LMConfig):\n", + " super().__init__()\n", + " self.config = config\n", + " self.top_k = config.num_experts_per_tok\n", + " self.n_routed_experts = config.n_routed_experts # 总专家数量\n", + "\n", + " self.scoring_func = config.scoring_func # 评分函数\n", + " self.alpha = config.aux_loss_alpha # 辅助损失的 alpha 参数\n", + " self.seq_aux = config.seq_aux # 是否启用序列辅助损失\n", + "\n", + " self.norm_topk_prob = config.norm_topk_prob\n", + " self.gating_dim = config.dim\n", + " self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim))) # 混合专家层 -> embeding dim 层\n", + " self.reset_parameter()\n", + "\n", + " def reset_parameter(self):\n", + " import torch.nn.init as init\n", + " init.kaiming_uniform_(self.weight, a=math.sqrt(5))\n", + "\n", + " def forward(self, hidden_states):\n", + " bsz, seq_len, h = hidden_states.shape\n", + " hidden_states = hidden_states.view(-1, h) # (total_tokens, hidden_dim) 其中 total_tokens = bsz * seq_len\n", + " logits = F.linear(hidden_states, self.weight, None) # logits = hidden_states · self.weight.T\n", + " # print(f'shape of logits computed: {logits.shape}') # (total_tokens, n_routed_experts)\n", + "\n", + " if self.scoring_func == 'softmax':\n", + " scores = logits.softmax(dim=-1)\n", + " else:\n", + " raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}')\n", + "\n", + " topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)\n", + " if self.top_k > 1 and self.norm_topk_prob: # prob 归一化\n", + " denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e20\n", + " topk_weight = topk_weight / denominator\n", + "\n", + " if self.training and self.alpha > 0.0:\n", + " scores_for_aux = scores\n", + " aux_topk = self.top_k\n", + " topk_idx_for_aux_loss = topk_idx.view(bsz, -1) # (bsz, seq_len * num_experts_per_tok)\n", + " if self.seq_aux:\n", + " ################# 对所有批次计算 Experts 重要性得分 #################\n", + " scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)\n", + " ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)\n", + " # 根据索引将 torch.ones 累加到 ce 中，并执行归一化\n", + " ce.scatter_add_(1, topk_idx_for_aux_loss,\n", + " torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device)).div_(\n", + " seq_len * aux_topk / self.n_routed_experts)\n", + " ######################### 计算负载均衡损失 ##########################\n", + " aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha\n", + " else:\n", + " mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)\n", + " ce = mask_ce.float().mean(0)\n", + " Pi = scores_for_aux.mean(0)\n", + " fi = ce * self.n_routed_experts\n", + " aux_loss = (Pi * fi).sum() * self.alpha\n", + " else:\n", + " aux_loss = 0\n", + " return topk_idx, topk_weight, aux_loss" + ] + }, + { + "cell_type": "markdown", + "id": "455f5c52-aad9-4de8-a571-d47de6f0d7d3", + "metadata": {}, + "source": [ + "接下来，我们设置一个 batch size = 4, seq len =16, emb dim = 512 的输入向量，在 MoE 门控单元前向传播进行观察" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "id": "5b3acf9e-b814-49f7-954f-246c33daa23d", + "metadata": {}, + "outputs": [], + "source": [ + "# 重新创建一个 LMConfig 使得 use MoE 为 True\n", + "LMConfig_MoE = LMConfig(n_layers=2, use_moe=True)" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "id": "534e4ffa-2a4d-4a0d-8906-997f23ec5098", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "shape: topk_idx = torch.Size([64, 2]), topk_weight = torch.Size([64, 2]), aux_loss = 0.10118068754673004\n", + "token 0 选择的专家：idx = tensor([3, 2]), weight = tensor([3.7838e-21, 2.9849e-21], grad_fn=)\n", + "辅助损失：aux_loss = 0.10118068754673004\n" + ] + } + ], + "source": [ + "gate = MoEGate(LMConfig_MoE)\n", + "hidden_states = torch.randn((4, 16, 512))\n", + "topk_idx, topk_weight, aux_loss = gate(hidden_states)\n", + "print(f'shape: topk_idx = {topk_idx.shape}, topk_weight = {topk_weight.shape}, aux_loss = {aux_loss}')\n", + "print(f'token 0 选择的专家：idx = {topk_idx[0]}, weight = {topk_weight[0]}')\n", + "print(f'辅助损失：aux_loss = {aux_loss}')" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "id": "521d65d2-bbcb-4b5c-8e1e-d532ae79b6de", + "metadata": {}, + "outputs": [], + "source": [ + "del gate" + ] + }, + { + "cell_type": "markdown", + "id": "b3773961-8e78-4bd1-91af-8d8384bd31d3", + "metadata": {}, + "source": [ + "### MoE Feed Forward NetWork\n", + "\n", + "完成 MoE 门控单元的设计后，我们可以对 MoE 前向传播网络进行重新设计." + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "id": "31c28178-7faf-49aa-a813-5105129eefcb", + "metadata": {}, + "outputs": [], + "source": [ + "class MoEFeedForward(nn.Module):\n", + " def __init__(self, config: LMConfig):\n", + " super().__init__()\n", + " self.config = config\n", + " # 可选专家层\n", + " self.experts = nn.ModuleList([\n", + " FeedForward(config)\n", + " for _ in range(config.n_routed_experts)\n", + " ])\n", + " self.gate = MoEGate(config)\n", + " # 共享专家层d\n", + " if config.n_shared_experts is not None:\n", + " self.shared_experts = FeedForward(config)\n", + "\n", + " def forward(self, x):\n", + " identity = x\n", + " orig_shape = x.shape\n", + " bsz, seq_len, _ = x.shape\n", + " # 使用门控机制选择专家\n", + " topk_idx, topk_weight, aux_loss = self.gate(x)\n", + " x = x.view(-1, x.shape[-1])\n", + " flat_topk_idx = topk_idx.view(-1)\n", + " if self.training:\n", + " # 训练模式下，重复输入数据\n", + " x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)\n", + " y = torch.empty_like(x, dtype=torch.float16)\n", + " # 将输入数据中选择第 i 个专家的部分输入到该专家网络中进行处理，并将输出结果存储到 y 中对应位置的元素\n", + " for i, expert in enumerate(self.experts):\n", + " y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(y.dtype) # 确保类型一致\n", + " y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1) # 加权求和\n", + " y = y.view(*orig_shape) # 恢复原始形状\n", + " else:\n", + " # 推理模式下，只选择最优专家\n", + " y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)\n", + " if self.config.n_shared_experts is not None:\n", + " y = y + self.shared_experts(identity)\n", + " self.aux_loss = aux_loss\n", + " return y\n", + "\n", + " @torch.no_grad()\n", + " def moe_infer(self, x, flat_expert_indices, flat_expert_weights):\n", + " expert_cache = torch.zeros_like(x)\n", + " idxs = flat_expert_indices.argsort() # [exp_tok_11, exp_tok_12, exp_tok_21, exp_tok_22] -> idxs[2 * token_idxs] = moe_expert_name\n", + " tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0) # [1, 3, 1, 2] -> bincount -> [0, 2, 1, 1] -> cumsum -> [0, 2, 3, 4]\n", + " token_idxs = idxs // self.config.num_experts_per_tok\n", + " # 例如当tokens_per_expert=[6, 15, 20, 26, 33, 38, 46, 52]\n", + " # 当token_idxs=[3, 7, 19, 21, 24, 25, 4, 5, 6, 10, 11, 12...]\n", + " # 意味着当token_idxs[:6] -> [3, 7, 19, 21, 24, 25, 4]位置的token都由专家0处理，token_idxs[6:15]位置的token都由专家1处理......\n", + " for i, end_idx in enumerate(tokens_per_expert):\n", + " start_idx = 0 if i == 0 else tokens_per_expert[i - 1]\n", + " if start_idx == end_idx:\n", + " continue\n", + " expert = self.experts[i]\n", + " exp_token_idx = token_idxs[start_idx:end_idx]\n", + " expert_tokens = x[exp_token_idx]\n", + " expert_out = expert(expert_tokens).to(expert_cache.dtype)\n", + " expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]]) # 专家网络输出加权\n", + " # 使用 scatter_add_ 进行 sum 操作\n", + " expert_cache.scatter_add_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out)\n", + "\n", + " return expert_cache" + ] + }, + { + "cell_type": "markdown", + "id": "c1f42fce-8bca-4853-8367-929c200f1b33", + "metadata": {}, + "source": [ + "接下来，我们设置一个 batch size = 4, seq len =16, emb dim = 512 的输入向量，在 MoE 门控单元前向传播进行观察" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "id": "6c7fb409-be15-4ee9-8a77-edacfacb3743", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "输出张量：shape = torch.Size([4, 16, 512]), 辅助损失：aux_loss = 0\n" + ] + } + ], + "source": [ + "moe_ffn = MoEFeedForward(LMConfig_MoE).eval()\n", + "x = torch.randn((4, 16, 512))\n", + "output = moe_ffn(x)\n", + "print(f'输出张量：shape = {output.shape}, 辅助损失：aux_loss = {moe_ffn.aux_loss}')" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "id": "698e4148-8f68-4f05-a33a-a76f70e20760", + "metadata": {}, + "outputs": [], + "source": [ + "del moe_ffn" + ] + }, + { + "cell_type": "markdown", + "id": "bec4c71c-00a3-48bc-abf1-d88b6fb94444", + "metadata": {}, + "source": [ + "之前提到过，MiniMind MoE 和 MiniMind Dense 的最大差别在于 FFN 模块的不同，我们可以对之前声明的 MiniMindBlock 进行继承，添加 MoE FFN 选项." + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "id": "3bad938b-bcf6-4bc5-ba2c-9c9689656959", + "metadata": {}, + "outputs": [], + "source": [ + "# DM stands for Dense & Moe\n", + "class MiniMindBlock_DM(MiniMindBlock):\n", + " def __init__(self, layer_id: int, config: LMConfig):\n", + " super().__init__(layer_id, config)\n", + " self.feed_forward = FeedForward(config) if not config.use_moe else MoEFeedForward(config)" + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "id": "778a8697-458a-47a4-b2bd-9bbdcaf64d24", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "类变量属性检查：layer_id = 1, 类函数属性检查：forward func = \n" + ] + } + ], + "source": [ + "# 检查继承是否正常\n", + "\n", + "miniblock_dm = MiniMindBlock_DM(1, LMConfig_MoE)\n", + "print(f'类变量属性检查：layer_id = {miniblock_dm.layer_id}, 类函数属性检查：forward func = {miniblock_dm.forward}')" + ] + }, + { + "cell_type": "markdown", + "id": "433c99c0-ab12-49e3-805a-493de3cd3613", + "metadata": {}, + "source": [ + "我们仍然设置一个 batch size = 4, seq len =16, emb dim = 512 的输入向量，在 MoE 门控单元前向传播进行观察." + ] + }, + { + "cell_type": "code", + "execution_count": 29, + "id": "0f97a54d-aa1a-45f2-ba86-4e31f17f9cb6", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "输出张量：shape = torch.Size([4, 16, 512]), KV Cache: shape Key = torch.Size([4, 16, 2, 64]), shape Value = torch.Size([4, 16, 2, 64])\n", + "辅助损失：aux_loss = 0\n" + ] + } + ], + "source": [ + "miniblock_dm.eval() # 冻结参数，直接 moe infer\n", + "x = torch.randn((4, 16, 512))\n", + "pos_cis = precompute_pos_cis(64, 16)\n", + "output, past_kv = miniblock_dm(x, pos_cis, use_cache=True)\n", + "print(f'输出张量：shape = {output.shape}, KV Cache: shape Key = {past_kv[0].shape}, shape Value = {past_kv[1].shape}')\n", + "print(f'辅助损失：aux_loss = {miniblock_dm.feed_forward.aux_loss}')" + ] + }, + { + "cell_type": "code", + "execution_count": 30, + "id": "925c973b-44d0-4774-a515-404f76a43a3f", + "metadata": {}, + "outputs": [], + "source": [ + "del miniblock_dm" + ] + }, + { + "cell_type": "markdown", + "id": "a9e64f09-ab3f-42a5-ac89-3809c1b99c26", + "metadata": {}, + "source": [ + "如此，我们便完成了包含 MoE 和 Dense 两种 FFN 选项的 MiniMind Block 的定义，我们对此前声明的 MiniMindLM 作继承修改，使其具备 MoE 架构." + ] + }, + { + "cell_type": "code", + "execution_count": 31, + "id": "cec79087-53e4-4925-95e2-0c46ad28aba6", + "metadata": {}, + "outputs": [], + "source": [ + "class MiniMindLM_DM(MiniMindLM):\n", + " def __init__(self, params: LMConfig = None):\n", + " super().__init__(params)\n", + " structure = 'MoE' if params.use_moe else 'Dense'\n", + " print(f'Initializing MiniMind {structure} Model...\\n')\n", + " self.layers = nn.ModuleList([MiniMindBlock_DM(l, params) for l in range(self.n_layers)])" + ] + }, + { + "cell_type": "code", + "execution_count": 32, + "id": "cc4517ca-b5ea-4756-9c51-7da415140e12", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Initializing MiniMind Dense Model...\n", + "\n", + "Initializing MiniMind MoE Model...\n", + "\n" + ] + } + ], + "source": [ + "# 检查一下\n", + "a = MiniMindLM_DM(LMConfig_Dense)\n", + "b = MiniMindLM_DM(LMConfig_MoE)\n", + "del a, b" + ] + }, + { + "cell_type": "markdown", + "id": "5a586a43-6f9b-4766-8d5d-24980d1a7241", + "metadata": {}, + "source": [ + "接下来，我们设置一条长度为 4 的 token 序列，使用 MiniMindLM 对其执行一次前向传播，观察传播过程与返回值." + ] + }, + { + "cell_type": "code", + "execution_count": 33, + "id": "f5009f2b-600c-4c30-b57a-85757301e6de", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Initializing MiniMind MoE Model...\n", + "\n", + "====> forward propagation started, num_minimind_blocks = 2\n", + "------------> entering minimind block: id = 0\n", + "<------------ finished, size_cache_k = torch.Size([1, 4, 2, 64]), size_cache_v = torch.Size([1, 4, 2, 64])\n", + "------------> entering minimind block: id = 1\n", + "<------------ finished, size_cache_k = torch.Size([1, 4, 2, 64]), size_cache_v = torch.Size([1, 4, 2, 64])\n", + "<==== forward propagation completed, num_kv_cache = 2\n", + "\n", + "返回 logits：size = torch.Size([1, 4, 6400]), 返回 aux_loss: 0, 返回 KV Cache List： len = 2\n" + ] + } + ], + "source": [ + "MiniMind_MoE = MiniMindLM_DM(LMConfig_MoE)\n", + "input_ids = torch.Tensor([1, 3, 5, 7]).long().reshape(1, 4)\n", + "OUT = MiniMind_MoE(input_ids, use_cache=True)\n", + "print(f'返回 logits：size = {OUT.logits.shape}, 返回 aux_loss: {OUT.aux_loss}, 返回 KV Cache List： len = {len(OUT.past_key_value)}')" + ] + }, + { + "cell_type": "code", + "execution_count": 34, + "id": "c0fe2fa3-f377-45da-aee1-b8f46341a470", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "gernerating new token: idx = 4\n", + "====> forward propagation started, num_minimind_blocks = 2\n", + "------------> entering minimind block: id = 0\n", + "<------------ finished, size_cache_k = torch.Size([1, 4, 2, 64]), size_cache_v = torch.Size([1, 4, 2, 64])\n", + "------------> entering minimind block: id = 1\n", + "<------------ finished, size_cache_k = torch.Size([1, 4, 2, 64]), size_cache_v = torch.Size([1, 4, 2, 64])\n", + "<==== forward propagation completed, num_kv_cache = 2\n", + "\n", + "gernerating new token: idx = 5\n", + "====> forward propagation started, num_minimind_blocks = 2\n", + "------------> entering minimind block: id = 0\n", + "<------------ finished, size_cache_k = torch.Size([1, 1, 2, 64]), size_cache_v = torch.Size([1, 1, 2, 64])\n", + "------------> entering minimind block: id = 1\n", + "<------------ finished, size_cache_k = torch.Size([1, 1, 2, 64]), size_cache_v = torch.Size([1, 1, 2, 64])\n", + "<==== forward propagation completed, num_kv_cache = 2\n", + "\n", + "gernerating new token: idx = 6\n", + "====> forward propagation started, num_minimind_blocks = 2\n", + "------------> entering minimind block: id = 0\n", + "<------------ finished, size_cache_k = torch.Size([1, 1, 2, 64]), size_cache_v = torch.Size([1, 1, 2, 64])\n", + "------------> entering minimind block: id = 1\n", + "<------------ finished, size_cache_k = torch.Size([1, 1, 2, 64]), size_cache_v = torch.Size([1, 1, 2, 64])\n", + "<==== forward propagation completed, num_kv_cache = 2\n", + "\n", + "new tokens list :[tensor([[1233]]), tensor([[2740]]), tensor([[2772]])]\n", + "\n", + "生成结果：tensor([[ 1, 3, 5, 7, 1233, 2740, 2772]])\n" + ] + } + ], + "source": [ + "# 我们让 MiniMind 根据我们设计的 四个输入 token 生成输出\n", + "out = MiniMind_MoE.generate(input_ids, max_new_tokens=8, use_cache=True)\n", + "print(f'生成结果：{out}')" + ] + }, + { + "cell_type": "code", + "execution_count": 35, + "id": "4cdccdac-930e-4670-8605-2963d0e62d1e", + "metadata": {}, + "outputs": [], + "source": [ + "del MiniMind_MoE" + ] + }, + { + "cell_type": "markdown", + "id": "903f2172-f021-442c-ae02-611a30506ebd", + "metadata": {}, + "source": [ + "## 参考资料\n", + "\n", + "- [十分钟读懂旋转编码（RoPE）](https://www.zhihu.com/tardis/zm/art/647109286?source_id=1003)\n", + "- [混合专家模型 MoE 的全面指南](https://blog.csdn.net/Code1994/article/details/145209710#:~:text=%E4%B8%BA%E4%BA%86%E5%9C%A8%E8%AE%AD%E7%BB%83%E8%BF%87%E7%A8%8B%E4%B8%AD%E5%AE%9E%E7%8E%B0%E4%B8%93%E5%AE%B6%E7%9A%84%E6%9B%B4%E5%9D%87%E5%8C%80%E5%88%86%E5%B8%83%EF%BC%8C%E8%BE%85%E5%8A%A9%E6%8D%9F%E5%A4%B1%EF%BC%88%E4%B9%9F%E7%A7%B0%E4%B8%BA%E8%B4%9F%E8%BD%BD%E5%9D%87%E8%A1%A1%E6%8D%9F%E5%A4%B1%EF%BC%89%E8%A2%AB%E6%B7%BB%E5%8A%A0%E5%88%B0%E4%BA%86%E7%BD%91%E7%BB%9C%E7%9A%84%E5%B8%B8%E8%A7%84%E6%8D%9F%E5%A4%B1%E4%B8%AD%E3%80%82%20%E5%AE%83%E5%A2%9E%E5%8A%A0%E4%BA%86%E4%B8%80%E4%B8%AA%E7%BA%A6%E6%9D%9F%EF%BC%8C%E8%BF%AB%E4%BD%BF%E4%B8%93%E5%AE%B6%E5%85%B7%E6%9C%89%E7%9B%B8%E7%AD%89%E7%9A%84%E9%87%8D%E8%A6%81%E6%80%A7%E3%80%82,%E8%BF%99%E4%B8%AA%E8%BE%85%E5%8A%A9%E6%8D%9F%E5%A4%B1%E7%9A%84%E7%AC%AC%E4%B8%80%E4%B8%AA%E7%BB%84%E6%88%90%E9%83%A8%E5%88%86%E6%98%AF%E5%9C%A8%E6%95%B4%E4%B8%AA%E6%89%B9%E6%AC%A1%E4%B8%AD%E5%AF%B9%E6%AF%8F%E4%B8%AA%E4%B8%93%E5%AE%B6%E7%9A%84%E8%B7%AF%E7%94%B1%E5%99%A8%E5%80%BC%E8%BF%9B%E8%A1%8C%E6%B1%82%E5%92%8C%EF%BC%9A%20%E8%BF%99%E4%B8%BA%E6%88%91%E4%BB%AC%E6%8F%90%E4%BE%9B%E4%BA%86%E6%AF%8F%E4%B8%AA%E4%B8%93%E5%AE%B6%E7%9A%84%E9%87%8D%E8%A6%81%E6%80%A7%E5%88%86%E6%95%B0%EF%BC%8C%E8%BF%99%E4%BA%9B%E5%88%86%E6%95%B0%E8%A1%A8%E7%A4%BA%E6%97%A0%E8%AE%BA%E8%BE%93%E5%85%A5%E6%98%AF%E4%BB%80%E4%B9%88%EF%BC%8C%E7%BB%99%E5%AE%9A%E4%B8%93%E5%AE%B6%E8%A2%AB%E9%80%89%E4%B8%AD%E7%9A%84%E5%8F%AF%E8%83%BD%E6%80%A7%E3%80%82%20%E6%88%91%E4%BB%AC%E5%8F%AF%E4%BB%A5%E5%88%A9%E7%94%A8%E8%BF%99%E4%BA%9B%E5%88%86%E6%95%B0%E6%9D%A5%E8%AE%A1%E7%AE%97%E5%8F%98%E5%BC%82%E7%B3%BB%E6%95%B0%EF%BC%88CV%EF%BC%89%EF%BC%8C%E5%AE%83%E5%91%8A%E8%AF%89%E6%88%91%E4%BB%AC%E4%B8%93%E5%AE%B6%E4%B9%8B%E9%97%B4%E7%9A%84%E9%87%8D%E8%A6%81%E6%80%A7%E5%88%86%E6%95%B0%E7%9A%84%E5%B7%AE%E5%BC%82%E7%A8%8B%E5%BA%A6%E3%80%82)" ] } ], diff --git a/notebook/2_dataset.ipynb b/notebook/2_dataset.ipynb index 3a4499a..2fbd650 100644 --- a/notebook/2_dataset.ipynb +++ b/notebook/2_dataset.ipynb @@ -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", - "system\n", - "你是 MiniMind，是一个有用的人工智能助手。\n", - "user\n", - "你好吗？\n", - "assistant\n", - "我很好，谢谢！你呢？\n", - "user\n", - "我也很好，谢谢！\n", - "assistant\n", - "太好了！祝你今天愉快！\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", "system\n", "你是 MiniMind，是一个有用的人工智能助手。\n", "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 中的元素" ] diff --git a/notebook/3_pretrain.ipynb b/notebook/3_pretrain.ipynb index 5f3b82a..d6fb2d9 100644 --- a/notebook/3_pretrain.ipynb +++ b/notebook/3_pretrain.ipynb @@ -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 \n" + "模型位于设备：cuda:0, 词表长度：6400, DataLoader：\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": { diff --git a/notebook/4_sft.ipynb b/notebook/4_sft.ipynb index 8313ca6..235f895 100644 --- a/notebook/4_sft.ipynb +++ b/notebook/4_sft.ipynb @@ -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 \n" + "模型位于设备：cuda:0, 词表长度：6400, DataLoader：\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": { diff --git a/notebook/5_dpo.ipynb b/notebook/5_dpo.ipynb index 9b72797..bab8cea 100644 --- a/notebook/5_dpo.ipynb +++ b/notebook/5_dpo.ipynb @@ -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 \n" + "模型位于设备：cuda:0, 词表长度：6400, DataLoader：\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": { diff --git a/notebook/6_lora.ipynb b/notebook/6_lora.ipynb new file mode 100644 index 0000000..bb98b13 --- /dev/null +++ b/notebook/6_lora.ipynb @@ -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：\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 +} diff --git a/notebook/7_distill.ipynb b/notebook/7_distill.ipynb new file mode 100644 index 0000000..2d2744b --- /dev/null +++ b/notebook/7_distill.ipynb @@ -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：\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 +} diff --git a/notebook/8_reason.ipynb b/notebook/8_reason.ipynb new file mode 100644 index 0000000..4ee3e3b --- /dev/null +++ b/notebook/8_reason.ipynb @@ -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\": \"T1\\nA1\"},\n", + " {\"role\": \"user\", \"content\": \"Q2?\"},\n", + " {\"role\": \"assistant\", \"content\": \"T2\\nA2\"}\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：\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('').input_ids\n", + " end_of_think_ids = tokenizer('').input_ids\n", + " start_of_answer_ids = tokenizer('').input_ids\n", + " end_of_answer_ids = tokenizer('').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 +} diff --git a/notebook/demo/lora.png b/notebook/demo/lora.png new file mode 100644 index 0000000000000000000000000000000000000000..4a22fc01f791cf97111bad274942b8402bbdbb51 GIT binary patch literal 11765 zcmX}SWmH^E)3Cdd1cEyR4TRwCF2UV{OK^AhKp?^0-QC>@8r)^j;O_3a-sIoOrp7%t`Hs@IxKtWiDlUnw0sMZ+AFT=(!HgM8y6(!@13PMaKSRzt-YbWrygSLDLyRBG%ud>9mb z-PqWeDpJT&EBlrrOHoUo`{^U6lB=bqrLgdy6jn2V`+tE|U9X;JYaPbgqQWe;xyMR* z;zWGzOgucEbERs`%*=4o_t4!*)VO~e`^R$H(Gl+Q?;0Tztcb9DiT;CR_(}+zd3QA0o1eW8G6G== z*4&KQ0VWnQp>Cs~U%$T?=JZM6QD7sw)79D>rocb^!xNE-qEtJ+K0 zq=%;0dM4mJmEj~lyUI$ohgN|7moC);W~|}uq$`;b)`Xzi)C{``$Jbq*0%lqLgz|a( zd!28|vjhOSKo3Ad3WSFuJ(KNhp3+|pHO3HnR@Fiom^>%5CWSUqnpJ5N%lvj%mZiE- zdR}_bwWvchxwP@BO(QC=@kvY6w0|%yIf*$(yJ+31Au5&ZHY#a&zqwmq(605d>$dxG z2eT)C6GeB@W`Az@m1XCvM*|JdntZRf#J3r2xnj2W3G>yRDY6v{>HN&O4=@{HByBd` z+veYX_zyb2nniBy7Yu*1K&tX{mMvh077wdn|a3hbr+P%d;|+U2{l!}4`sKnU7rm=nqP7#8J<=+s@!$&@ zSwy2fO3YQJOZLpf(CPP23_a#|o*3BlDE`$%nylh~M6dQ5f!>@_Wl&%8!KX`M*8$!8 zCGT48A}?taf9>amBt{7_G#k!`Z7_8_e;JKAyM{sZD8De12w)g=NhaJdgkgG!>3MPnLrv_ z|7Dssoyz!oc+;jy^wD-zEYN$mlP6cTao83*w$hJnh9IDsFgc0}U5@aa;wWSN2bWU0 zY#e3HpO;{QSUj#Sp^2^T65oNQ0mH-eStrxMDkoawj5FrTEm497yCv(^`q$9@tX-QEj4*)()DJkz|@tE zF+ak)Z1uiM(70eIFC{(`Ty)zwDX-F#np!gKB~e(*H~1KW2FK^Lh)x??!f~^$`*2&d zmv4eC8~NeCA9g@MBU8m;va*9(8!k(*{TMSz+m22P23NUGoPDPu`F4r)<|F^O76Bj9xy8LvtuvD1c#?L{JGyoAEY9*V69`iWA+ zu`4yI(x?@?9EmfT{`818NiM@|>T>UMe|nm4w5rImTP(J1C|XcbbBCN>9FHs*NgyF$ zFvC4LdM=W4e!6gfXvbBb@NkV;RiGA3lNBYK;uIuYNmi3+B&{60^KiNL|6^zE*;N2- z+bOX8@R6AHwBpXSC*s7IV3Q@?EBn2Fm0fCWRbiuYxH8AXsE6y+=AHgK|9ptB@$x1# zd1mo@|8`!M7|B^_W@(5#6h9RL&^gi8WS&e3Y2;!@_8f`1&VtHz(S#|mru@FJrV9Kf z`pDg4j^NuV=k{rWCe@x%Ac{8U+vXNVUcp$uXa!O@cGg-M2 zlN2K)GXJgp*xn;>?Q0{w{C}eFtoXLw5}KQNNJGZ^&yyYO%5}5c(^Ic~!qe1RNy4J- zP&3TmsN-uhB6q56;EDKMmhDa(OKgiQCbbiZX_nJ6A*NMRIN+cER`SEy=KNXTQ)1!Q(Z4)-B}prnGsgrMHQe+|{#oPaO8In*4g8QbpYF}T(hX*j+}VBzI5t*{Wc^WF zuMBF!&F6AIRG9Rovg}wb4O%J z-2nx_XAFh-NV(l=cN2GMn>*ho_m3&`#j6bvbS)AKzcW<*7crLv15~PSX}udO$||@O z-b3QO6MAuhO7v-Ia1Xc@X@H!F{mq#mAl z$*?Bj?#uhF)x$<^c_w7|iQB{lzz@Oz_D)8vLMz2$$t&ZUE_KpL?zucqmp7Hy_%;p( zX&v7^w*o)(p7T^Z9P9R<3>OMUY@tnWN#s@y1eI+tc74){%k1?pO=M2GM|`a^lZQ7x?56KQ!3m*1Iq+HPy|zI zG>Vqsd-YGhm&S*d-O%OK-nUx*0O~Se5Umrp*a}JvJ#td{sXC7WN@(-q61&QVka(st zI>2>P$mA#-2jEK4hXV$y7Oc~$GX$rLLTW9j4F3@CXovulgHcQe#r|?e%*XC>a^p6a z^TvPfUcZ~Q5QYM#j3wGvrP{452WZR)z3-)GGL7DcBopx@t+^$)Y`?!q#lAr*j7_6G zW1|jCLmIA1Vi#=k;NV*+W)=Kv@>5is=4=w?y%92Y;$bcd1O^~r9e#|;YVp1~BI2)c z@hpZ5fx-G_k#%t;j@IoUnNFyolNaadFmR2TIa;;smoptG;^)xrr>ZZ2EQUba@b=Qt zu4NHUMCHk`1P838;~Nj0clfF5#ZafEd$Fl^y4SCtcX|_teXC+G+M6@RVq=I|ZZQA% zikN^Kme6;^Z#F`SgcX5=L*ugB%YDh(L_k2(riMv^jEX@=v_*--G<=9We>pB>6yEc! zX%Y75S-pJpm`CO}e-26h_9vqg=lca3j>65@hygP@V@TRE!9pJQE2-bB7u9+g4q6F! z^9P<}#Ms}Il24CmC>e9ap0?Jahy97P5q-FOb>?-yJa8BG-HiFzl9>Q-sqlcGphiBY zk<8$Ig&)oSNuXHE7fZz8EwhOGVP>=Q+jS)w>W<8 z@Tg(;`l!KQq#Ad*lHbgB?#;R2QMlr@Eldqy>Q0h*wrjEkGY=Ju(fyk*qn~2g1dZzR zom-K1dkyGGIeA64qg0`R;79NK+3Fx{!tKPV9ODw(|9 zH-&v$1|ARiTgDA+pXs5Ws)0`eOM|!03VLCUkW55StB+JA-P+yTOp1j5QcZ$qiY3;^ zf%s#|^XBN)78U9zjD`ub2Xcu6gJn7h6D%E{nRAh?&0SFb{(2TtS`*tm0Q(t(eBA*EsHWWe4U% zwMtT@MAA#v7s=L%d6=IUR#- zyO0i5=oVfm>l_#_a0edU)aP9AXFmB6Pwl(|&=;tE%XWW@=5`#U)Z_)5(xY1V*~N-B zzI=lmfE(`?mQyt6@T_Bk1w_M-kr)Q}N$7bjgia&|2$j_R17a{T{pBzhpXxr2xswoQ zjEEkvZMP_OI%^Fk3e_v$1U*SZrE&QEM9mTlI~-AP?;!!CdAZ8hSS1UT$;AIWz}WjT zSl^4x70vul9VgKRI}+T%TaBomb?@G`XM2+}s5)-Zhnb|M(9SRZp?xV9V z86R!)bKrq`w!0sK4D8T2MA(kz$Xaz57s~-7q957nh0Pkh7@$9GPDkMA7*$Uy!uH~K zVPpg+zqLo6SOM92zu$Qsr!}dvr~SCptk;V8F!&>qY*)TzZ6V6{6+(bWAzz=>+fJa# zad%*^`#OwgwxS!d&wk49S{Q~{rCmgMTT)}-E~DaxGa6YA)on1aO%qq6X;sB;__O>Z zGYn?dYOK!@MZ0538uEP4KBh54&{|M>Ls2enLniGeKhHk|9z^_N5d9dA7cqh}Nl)9J zv;`ibL`-bLRPG^ItCC9xPOI2?lM&>^jQX9Hatr^L_}v13D@i`L410LSt>XQc)J@t* zdyNKc9<~KXP+UzRx=kW|n8*B&R;UX{`9}3vy2j&mVWD>GvV65(QwA)Y+qRS1uPJIt z+8QfYYAVC7+G~r~uPJ}ecYiSPf583ZG=hnu^>LafzIHWb=qfSR2HK*zCF{kyUO`6Y zjcSDeBOc7j=kdZm)=_jMXI++t{wFc;AI>t&z9wJejA2$=cI3iVbt8M& zNrfRE1&3&MC`pj!c8hFAxZXbj8zRu|^DHq??d{49Cz^da9oS@}HRl%w_unvjvd08A z>&(A;$&{Oqg()|GBESArl~kG(ej5>eUimj3K--vT(=ITXj3)izP4>DFHh7N7a+g>N zi)Rw*c$tQ;DPk>95M!6lOvs3Ts`um1g8 zzdo2-TqejVILz(&@*LbsA?8{82Ex->BMndnukaTa1PyXMByUe%9?w)H97Hc0j_l}0 zZ_fF9fd2S8D)z_GMXHi&#D{LGbZEx`@k{Bs7z|>d`eX5Im`~YPrw-N6+mnIqivSw+ z2DYlX98P*L8EddR(Rj>yNpIT+hj?_9_jbnQI1 zv$NVY7xF@)0z6FVYD}w{jlWdv3-)PlI+i^R!t*s2>udG1@e@7)|+hvX$ zGNwrg)$!s%kQ!!CdO%27^B(x*mpK#wrc{##o`f|PwrErU9Yaqq1Oe!SejgVZiFl*I z`X%~)&}f>^Bb-8A#a#B4j3O)&x*O1sgaj;D5ckt9A^Ta&o96U{P2AtD_^N#7)1;3Y zy2h}yvckl~q?Auj&dhurrmYGm;AHsxInX+f=RG$oyJ#FAx&IlX;acS^2Jy)l%gsMv ztpTo7K!<0_$9Pf^XjaRq5wF#0uJTqavvXwEtkCPr@grRC61O=#8d~$pnV$U`wk*x7 zqgB;Yc{f*fcy(FFM4;QE&h^H+P%c+3oo` z9*;8>1jhw$nz+s5%OUBH=T8ca?BXxtDXEeF85+G39@v*#BoRZNK)OP~c0z+qc6Ju6 z?p9c2TUU=*Vlz?9G7I*QfJ1y^`RgQFoz4H!^*)2U8g&GW9ds5b1Fig`s^|}ZPqQ=Z z4=uSiTU5^7jox2036{h0)W!tUQj*%7vD&$MEK3RnoPy<6*9=%%b>-g}Zw&Rl?#;=) z)=3gZ1|Drvzq|wRL*r?Y+Q!);y);A!NahbO6r-j+`H95%xUSV|pv=zLy}df2kAnoY zRO(D|Xa=sY79Upp9fLN&)0UyYR^CGn7YaOd*IH@U#l5*F>uBFbBztR7wLUs=J1*IV)Pi*4Q&lig2uZ@;J} z<=Mq<2a~_T#kv&c=R?dZz^{RzIa&jf*?FTD-$X*!MKP6 zj{Ll&r;~yC<{_m#A|^gFCLAYw&k;NVzcSTl3v?`dzqWm^x|bJffSiuzthd&~4AmbBzZ|eYcc9!dC*vpEVfUb1iO+@9CT#>QEiGy5-X>?>luNCS)r@;~ zfgi0B1sCEHZuLwih#2(vnX{|R;2hCpIu@@ae(y!u53J&Rz(VLx-P@xWJO8rUVcV1Y zJ#+V0%TSijq7ggzFMW?}e1gNA$yC4}Vy$qrD**_7D+6x@-ib~j1Btl1a{73osJA=O zHouA1vNoR%B5k;n#_wcZIx69~JG!r1p*k(=E%hSJ}27CCw zsCfUzo$UQ#OJ)ZC(pt!_Z0!BH#%8L}toa4^0gnw90Mp4%C<8qln34mZJsKtG@O-PZ z-~i+(#`HUMrUNQy;I7o=C=+4WV==jQue*Pkv*=w8yd#$cG0pC?O~CeyzD#+o9@DhM z;h2E4oIqg~)90>A+CrYl&w@UFyOVAfGGG64zYe?M^iQN+Q`?ScHuxb++8xd9;p_ zX}T3GINvN!Rr!~qKPmSh-fYY+ndj&I|uyxYyjQ8sJO8?)%o_+K?vmQ9s6#I!^MvY`2>x& zap0uV-A}Eg`i(|scIw(aA*J9??LBwA-+aJ4?(ekkYCsCi&sL<`EDT@gWbB4`M$p%= zh^M>raXwsSjHq{=Hfs^i16V z(7x&G3cv|SN|mUPkZeZBI3*DqX3}SHAYTl#>1Y_Zy<$qfsC|?`I6dpty!dIXvgD%= z^_gLK*^_TEOjo%I4%~#p4C6CHA8~=f!G7hn_4!JA3KYcic`E^S{G#ALWXq}0* zJJ@gEE?J3qVf%a|(x2YxqtPbcgojJNIh_=Fv;5-4UDG8&`Keaq12(&gyLjX#AQf>T%hd%I@(Z9Ca(`FTs$%Qk;@MMGvB zt$rIX!Xk=461G#=eeqd4?+n0RU*TDR!1GZcx2RXNsaK>4_UzqfQPYHUS|1B39Scn< zsFxM29nWFHz^~rjZTNbZ{yeXc!+DNV`SO(`VlkBZDUwWJ$!t|c*rUKwvw}f&^aHk2 zy6C*-N6OePX=(6efQfXwQ%=wyk;bwnyx^vjwt7E^itV(BRG+OQQtt|&oetjon;c)W zvZC6`C(b2#5+1^y{UuB56p2%Gu747SaqM=kg^m$9?0 zX;a$GFX{D}58UdPlj)I+$GpYrm$7~avxIRx+k}`T&VSYGFNk%Cyf+4tN2Bmp&6;X0 zrv00`UMS8_{X}!71g|U96rh!t_IjBJ0i5vbhYd%Uw1@OQXz?pQ`jh_e{Q(;(O-E0+ z-_^>(=VY{#Gg9KvAOOL{!tj*8=1WFD~1csQ<>pd0hGAY%W}X^200${c@auE$ zPU9D^JXxaC!3G{CqMbbgc58}=v3#Pjgh7co66U!Ma=+W0hwrDucze2?)1<%Gw!tmj z^~_KR?h(@0NU%W6w`4;v)Yqy}PFHE+70_Dyr`e(q3-E2#($&dpx@B(=4_i$kddwQ( zQ-l>po5ZgvF!BcewS^_jkS#B!tv38^aekd4TAs#Oi+zehM&F@PXl z2i3f;r&r+-lLP+#0(2BaA1+K`gdF7d>b`J)KnOA%FaKyU@vgPiYc*Z0m(R|0#m{WD z?9BG9%eEyu&-E&M1}(SEw1e;050?GnX%p5@u0f|uV@8Yq*>?jqtGylX0UF7U6(xs!D64T6@8VVi`pf~8X;TdQ7_H>} zg5Ad)qx%bz+TaNal+%oGd9oFd5vcx;1t6(^7*lWk!qjqCIpInlE4Z%dd=Vc9>kkF= z93L7B6jrJgSu^SI8vf-LNBtu{4>sJuC+dX{rc#=2t z&&H~_W!Z3I9}7ChF^ORQI>+@H-RDd#&UooVTq#=L z99#0g7cCEG z3fJ|^R>l#(jed&eW&Y4Df&PAV@JH!zHo|%RwcW$dHcrCG-mH8k^SSSu>DTUDGdug? z#-N{V8^?>t)(IW(kFSOB<{bQbQO6~!%ie2{R@{B!S;TL{%SjUoNS3iJR^hNOs?z20 zX>U2+1$Qgg9T}xzS66UT*XInia-$~c*<|jaz(cGkhNHqRPg5t($hOAy0T>{l1Vx@F=oWwvt#*gk z_Gcm}qUKW3W5&*8U*(p$BAPT0k*eZv31mHK$WMvgu8OT?k#~NSxoI_LUNF2mVvNF` zoq(@9({nP^pMJ|6CFCL6SXANwH>&uv)#v8dXMr{^wJ8zt6q(Ajk;@d3toe-BbYi|9 z9nYCY02CiiZa(uq**y%Queh6_j1dmdsTycNNL3J&ISJfVH$}pmAnroCITyyWmR0&r z(^@?897>EfMBFhs%@|CS=3i!|_C&B?nMSfZ)gEB3Cu|Yc7GL~Rx$(+zyXcHHs-PYV z5WhnBdannrfAm~q);7&%a??2>V~bwsjCb~#30LLY#cpRA?+Z8or>|aHNrRTxCMtddh^R9 z4BXMgSuZ{fI>oPDNYXBS!{4)I9HF4j=GNCKuTdlbK4NO|btbuaX#AMv?tpXd*IJD# zE);OM5rxpNA)f^WHSdbH4f=Py#}W z(4LdkyC(ODYef$~IRw|vcMQ8zZF=PL^4bNGbD;s*CO@Od-*vM|7!hUIE7q_)Gmkhb z-=HE%y%`l3phL+Ba|#+QRv2i|KFd~Wx?%&0F|fi1&7Beb_JBXr;`v^l!4%Vk)1XtF z;FJA>6~PkU|9Wqd1QsIZNT=6LszZ!Mhgv}QYy*6g!jZ z9{`L7q~YVTSBsCc<2+lsm})`oczs7sm~T2$xQJM;p|?E#xxu`(#J7L|cI5IR>Pe}F zvzmf-=?g+cm|W;V@B76kcSj(EO1P5W3x}iPl^DnRYs@%8GyF)cl~o^20u%}+Ao-LX zS-E}|3kIdw?9(@OX~HSOw`jpP{`fO7#NV@oE27S_DKRQcV()USWmU{Z-x7FTkh7_1 zV$CBv7Yg~|1Ldwhu=D?|AwJN_#gM2V$f+)PIT8nh$XctR`PWSFCqh)6s(Ko^kbehSH%@TL^4i;cr96W&F(H86189LDRQ<=oc9vA1>MuX#7uqu5hux4>6swd@F zn3f$XBC)oht}nL{^qRW9Xj$)=Uj25%zy~Foo0LizKV)3X`|byfx2q0iKNY9H0J(f;^_O)Jm>D3Bde!D z#~43khdWY*6JHTd(!t`31QF-{j8Svmw#7{op;WeXdeoy{NgifjcTN=sOS=tMiP?j z5^;ttVfhM0>B7^nn7x!?Asnvn%d$>-h8g0Aa4Y$l98|j~PUJiqSU~Crgd##CkCiK- zo6Ni%Z%+&#PU)jeZmFui`NlYsc1!H%M)8)_>Mjy|hzo@>u-dGrKvnGVQMG8TAIx8_ z1JM}ztNE+C)WY(_dw7tEi8IpiEnN>vKn6a4h6dF@VRLmXg^oe8QJ5cIKN_&xTBP{@ zVM)hpELrA6sfr-%!hfpOe7pF!_|Fel-utPTAJn{wu=nsqRnp__s7eK9-p_WDmu`WqgQGaHSmx*~wPYzi5?(CgS^zm&Evv z003)=z2ww_nwiDEqTP2S#oLDNhBrZR&$82#A|nR=45WQ-NhYL5+8fv%4)YV?DnL!i zNqvN7D1+m(YtVYx3}Zz3$N~#A+AQDBCRJNGH9%yQvDTCcRneO8*SGs z|A#xBaxH%=*I;%(|HrxRt&*eY?%^>vJ9~X#M46(}G?Pj4XDMeTm22$K<4j?K+o_~@ zXeGkxqBtjUbZHnsg@&+30~rcOuSCsYneB-!#g=Ubo*mg+5> z3Pbz%Ni>6mm35RxTljKpo%R_^DcisoR`7nK-{vYulW_Uq4RC>?H#-phm0A#Mq>z^z zp{wRaTsHkRZIeje++M4lJ9g0tDz$7=f_HqbQG*UI7^-`&Or4dLH7_?eG%U<(;7=MB zCdj(^r8wpcU4UUeqy7bim^cPsSAe9uy5&CGq>GVwL@?TUW^+`_=9H9_0^cL2ayu33 zFbKLHFSdJJaaiwbl&;8wP7g*;qx^7voyfuaW(CyRj78UvAdU!Pd@8PVuGfq3!QbBp z{^;HK$5lCV(7YmV`+1$4g%R6EW*h{_el|BPAFH*&b;v^uX57o_gPTLpQHZc>3W{Z6 zfE+0DaI~)1sT#+Y+IFp-Te*gWwPI^%!)IQT!C2xl=Bi>SAUMdgUbA)?bP!Ni3r&1C z`@SG}dB^w5w`_8MR3__CrjfJ*+U+-=J$C8g@m_ABJ=JNIb}|L2AY{ikJe=WC#g1y=Xq>JaDo#M5EG*SQc_Y1{5P{P zpCxujQ;g_8V}#T?k{EjGWhtdYJoLH7R%-PUTCg% z_GN@(Gv$dg$Gf1Tul${$Qo%y-DVS;z?zZ2c`f3+K7r5>0K%5(p5|bCL{Av*J{{W{e Ban=9; literal 0 HcmV?d00001 diff --git a/notebook/demo/lora_data.jsonl b/notebook/demo/lora_data.jsonl new file mode 100644 index 0000000..cb8e2bf --- /dev/null +++ b/notebook/demo/lora_data.jsonl @@ -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?"}]} \ No newline at end of file diff --git a/notebook/demo/r1_data.jsonl b/notebook/demo/r1_data.jsonl new file mode 100644 index 0000000..a1d222a --- /dev/null +++ b/notebook/demo/r1_data.jsonl @@ -0,0 +1,2 @@ +{"conversations": [{"role": "user", "content": "请用一段话描述阿里巴巴集团的企业文化。"}, {"role": "assistant", "content": "\n嗯，用户让我用一段话描述阿里巴巴集团的企业文化。首先，我需要明确阿里巴巴集团的平均文化是什么。企业文化的定义是企业内部发生、指导管理方向的准则和规范性的精神或规范。\n\n阿里巴巴集团在2000年代中期开始发展，之前也有类似的概念，比如微软的人工智能理念。所以，我可以从这个发展中来寻找关键词和结构。\n\n接下来，我要考虑用户可能的意图。用户可能想了解阿里巴巴的企业文化，或者需要用于某个报告、宣传材料或者内部交流中。因此，我需要确保描述既全面又有深度，涵盖战略、创新、效率等各个方面。\n\n我还需要用简洁的语言来表达，避免过于复杂的词汇，让段落易于理解。同时，要注意逻辑连贯，段落之间要有过渡，让读者能够顺畅地跟随观点。\n\n可能还需要引用一些关键点，比如战略协作、用户体验、创新、效率、数字化转型、客户支持、全球化等，这些都是阿里巴巴成功的关键点，能够帮助用户全面了解企业文化。\n\n最后，检查语法和结构是否正确，确保专业性和准确性，这样用户可以直接使用这段文字，或者稍作修改，就能得到满意的效果。\n\n\n阿里巴巴集团的企业文化以战略协作为核心，倡导员工之间的资源整合与创新思维的碰撞，旨在实现企业价值的提升和持续发展。阿里巴巴将其视为一个协同共生的整体，致力于通过技术创新、用户体验优化和数字化转型，致力于为客户创造最大价值。企业内外部力量协同作战，推动企业从单纯的业务拓展延伸至价值创造和社会服务，并在全球范围内进行全方位的数字化转型，以满足多样化、个性化、高端化的客户需求。阿里巴巴集团ix platform的建立旨在帮助员工实现高效协作，激发创新精神，吸引更多优秀人才加入，共同推动企业不断向前发展。\n"}]} +{"conversations": [{"role": "user", "content": "你好，我是张华，今年15岁，是北京四中的学生。我喜欢编程和篮球，我梦想成为一名软件工程师。很高兴认识你！"}, {"role": "assistant", "content": "\n好，用户是张华，15岁，来自北京四中，喜欢编程和篮球，梦想成为软件工程师。我需要回应他的兴趣，强调他的优势，比如编程基础和对篮球的热爱，有针对性地引导他规划未来。可以建议他关注 Font Awesome，展示他的技术背景，同时鼓励他扩展兴趣，并辅以鼓励自己。\n\n\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"}]} \ No newline at end of file