diff --git a/notebook/0_tokenizer.ipynb b/notebook/0_tokenizer.ipynb new file mode 100644 index 0000000..20e4b30 --- /dev/null +++ b/notebook/0_tokenizer.ipynb @@ -0,0 +1,453 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "source": [ + "# 0-Tokenizer\n", + "\n", + "对一段文本进行分词就是将其分割成单词或子词，然后通过查找表将这些单词或子词转换为对应的编号。\n", + "\n", + "分词器的就是把一段文本分割成单词或子词的工具。 " + ] + }, + { + "cell_type": "markdown", + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "source": [ + "## 子词分词算法\n", + "\n", + "常见的子词分词算法有三种：\n", + "\n", + "1. 字节对编码（Byte Pair Encoding，BPE）\n", + "2. WordPiece\n", + "3. Unigram\n", + "\n", + "### BPE\n", + "\n", + "BPE 是一种简单的数据压缩技术，它会迭代地替换序列中最频繁出现的字节对。BPE 依赖一个预分词器，该预分词器会将训练数据分割成单词（在本项目中，我们使用按空格分词的方法作为预分词方法）。\n", + "\n", + "在预分词之后，会创建一组唯一的单词，并确定它们在数据中的出现频率。接下来，BPE 会创建一个基础词表，该词表包含预分词器最初生成的数据中所有唯一单词的符号。然后，会将这对符号从词表中移除，新形成的符号将加入词表。在迭代过程中，BPE 算法会合并频繁出现的符号对。\n", + "\n", + "给定词表的大小，BPE（字节对编码）算法最终会合并出现频率最高的符号对，直到收敛到该大小。\n", + "\n", + "### WordPiece\n", + "\n", + "WordPiece 算法与 BPE 非常相似。WordPiece 首先将词表初始化为包含训练数据中出现的每个字符，然后逐步学习给定数量的合并规则。与 BPE 不同的是，WordPiece 并不选择最频繁出现的符号对，而是选择那个加入词表后能使训练数据出现的可能性最大化的符号对。\n", + "\n", + "### Unigram\n", + "\n", + "Unigram 算法将其基础词表初始化为大量的符号，然后逐步削减每个符号，以获得一个更小的词表。它会在训练数据上定义一个对数似然损失，以此来确定是否从词表中移除某个符号。" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 训练一个最简单的分词器\n", + "\n", + "在本节中，我们将学习基于 transformers 库来训练你自己的分词器。" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 初始化\n", + "\n", + "首先，我们应该初始化我们的分词器，并确定选择哪种方法。我们将使用字节对编码（BPE）算法。" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "from tokenizers import (\n", + " decoders,\n", + " models,\n", + " normalizers,\n", + " pre_tokenizers,\n", + " processors,\n", + " trainers,\n", + " Tokenizer,\n", + ")\n", + "\n", + "tokenizer = Tokenizer(models.BPE())\n", + "tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 定义特殊标记\n", + "\n", + "数据集中存在一些我们不希望被分词的特殊标记，我们会将这些标记定义为特殊标记，并将它们传递给分词器训练器，以防止出现错误的分词情况。 " + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "special_tokens = [\"\", \"\", \"\"]\n", + "\n", + "trainer = trainers.BpeTrainer(\n", + " vocab_size=256,\n", + " special_tokens=special_tokens,\n", + " show_progress=True,\n", + " initial_alphabet=pre_tokenizers.ByteLevel.alphabet()\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 从文件中读取数据\n", + "我们使用 JSON Lines（jsonl）格式来存储训练数据，让我们编写函数来处理这些数据。" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "\n", + "def read_texts_from_jsonl(file_path):\n", + " with open(file_path, 'r', encoding='utf-8') as f:\n", + " for line in f:\n", + " data = json.loads(line)\n", + " yield data['text']\n", + "\n", + "data_path = './demo/tokenizer_data.jsonl'" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'近年来，人工智能技术迅速发展，深刻改变了各行各业的面貌。机器学习、自然语言处理、计算机视觉等领域的突破性进展，使得智能产品和服务越来越普及。从智能家居到自动驾驶，再到智能医疗，AI的应用场景正在快速拓展。随着技术的不断进步，未来的人工智能将更加智能、更加贴近人类生活。'" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "data_iter = read_texts_from_jsonl(data_path)\n", + "next(data_iter)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 开始训练!\n", + "\n", + "我们使用分词器的内置函数来训练分词器，轻松愉快。" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [], + "source": [ + "tokenizer.train_from_iterator(data_iter, trainer=trainer)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 设置解码器" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [], + "source": [ + "tokenizer.decoder = decoders.ByteLevel()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "接下来，检查一下特殊标记是否得到了妥善处理。" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "assert tokenizer.token_to_id('') == 0\n", + "assert tokenizer.token_to_id('') == 1\n", + "assert tokenizer.token_to_id('') == 2" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 将训练好的分词器保存到磁盘" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "['./demo/tokenizer_demo\\\\vocab.json', './demo/tokenizer_demo\\\\merges.txt']" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "import os\n", + "\n", + "tokenizer_dir = \"./demo/tokenizer_demo\"\n", + "os.makedirs(tokenizer_dir, exist_ok=True)\n", + "tokenizer.save(os.path.join(tokenizer_dir, \"tokenizer.json\")) # At this point, you will see a file named tokenizer.json under tokenizer_dir\n", + "tokenizer.model.save(tokenizer_dir) # generate vocab.json & merges.txt" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 手动创建一份配置文件" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Tokenizer training completed and saved.\n" + ] + } + ], + "source": [ + "config = {\n", + " \"add_bos_token\": False,\n", + " \"add_eos_token\": False,\n", + " \"add_prefix_space\": False,\n", + " \"added_tokens_decoder\": {\n", + " \"0\": {\n", + " \"content\": \"\",\n", + " \"lstrip\": False,\n", + " \"normalized\": False,\n", + " \"rstrip\": False,\n", + " \"single_word\": False,\n", + " \"special\": True\n", + " },\n", + " \"1\": {\n", + " \"content\": \"\",\n", + " \"lstrip\": False,\n", + " \"normalized\": False,\n", + " \"rstrip\": False,\n", + " \"single_word\": False,\n", + " \"special\": True\n", + " },\n", + " \"2\": {\n", + " \"content\": \"\",\n", + " \"lstrip\": False,\n", + " \"normalized\": False,\n", + " \"rstrip\": False,\n", + " \"single_word\": False,\n", + " \"special\": True\n", + " }\n", + " },\n", + " \"additional_special_tokens\": [],\n", + " \"bos_token\": \"\",\n", + " \"clean_up_tokenization_spaces\": False,\n", + " \"eos_token\": \"\",\n", + " \"legacy\": True,\n", + " \"model_max_length\": 32768,\n", + " \"pad_token\": \"\",\n", + " \"sp_model_kwargs\": {},\n", + " \"spaces_between_special_tokens\": False,\n", + " \"tokenizer_class\": \"PreTrainedTokenizerFast\",\n", + " \"unk_token\": \"\",\n", + " \"chat_template\": \"{{ '' + messages[0]['text'] + '' }}\"\n", + "}\n", + "\n", + "with open(os.path.join(tokenizer_dir, \"tokenizer_config.json\"), \"w\", encoding=\"utf-8\") as config_file:\n", + " json.dump(config, config_file, ensure_ascii=False, indent=4)\n", + "\n", + "print(\"Tokenizer training completed and saved.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "现在我们已经训练了一个简单的分词器，并将其进行保存，接下来，我们试着加载它，并使用其帮助我们对文本进行编解码" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'失去的东西就要学着去接受，学着放下。'" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from transformers import AutoTokenizer\n", + "\n", + "tokenizer = AutoTokenizer.from_pretrained(\"./demo/tokenizer_demo\")\n", + "msg = [{\"text\": \"失去的东西就要学着去接受，学着放下。\"}]\n", + "new_msg = tokenizer.apply_chat_template(\n", + " msg,\n", + " tokenize=False\n", + ")\n", + "new_msg" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "actural length of tokenizer： 259\n" + ] + } + ], + "source": [ + "actual_vocab_size = len(tokenizer)\n", + "print('actural length of tokenizer：', actual_vocab_size)" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "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" + } + ], + "source": [ + "model_inputs = tokenizer(new_msg)\n", + "model_inputs" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'失去的东西就要学着去接受，学着放下。'" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "response = tokenizer.decode(model_inputs['input_ids'], skip_special_tokens=False)\n", + "response" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'失去的东西就要学着去接受，学着放下。'" + ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "response = tokenizer.decode(model_inputs['input_ids'], skip_special_tokens=True)\n", + "response" + ] + } + ], + "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": 4 +} diff --git a/notebook/1_model.ipynb b/notebook/1_model.ipynb new file mode 100644 index 0000000..b9ce3b9 --- /dev/null +++ b/notebook/1_model.ipynb @@ -0,0 +1,728 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "f359c4f9-9082-4843-8874-4e7c6f3c7404", + "metadata": {}, + "source": [ + "# 1-Model" + ] + }, + { + "cell_type": "markdown", + "id": "2f5a93fb-2d09-4af4-9782-578357927842", + "metadata": {}, + "source": [ + "如今的主流大模型的架构大致分为稠密模型(Dense)和混合专家模型(MOE), 在这一板块，我们将围绕 MiniMind 系列模型的源代码展开学习。\n", + "\n", + "## MiniMind Dense Model\n", + "\n", + "MiniMind 系列模型在 Llama 3.1 的基础上设计，基于经典的 Transformer Deocder-Only 架构，其主要特点如下：\n", + "1. 采用 Pre-Norm 归一化方法，使用 RMSNorm 归一化函数。\n", + "2. 使用 SwiGLU 激活函数。\n", + "3. 旋转位置嵌入 (RoPE)\n", + "\n", + "作者提供了对其 MiniMind 模型结构的可视化：\n", + "\n", + "![](../images/LLM-structure.png)\n", + "\n", + "我们在上一节中完成了 Tokenizer 的学习，这一节我们关注 LLM 模型架构的具体实现。" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "6cd2aa37-3460-4a64-8285-f888d62b99d5", + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "import torch.nn as nn\n", + "import torch.nn.functional as F" + ] + }, + { + "cell_type": "markdown", + "id": "acbb3ad8-5143-45ff-af16-8dcc0cc2cf53", + "metadata": {}, + "source": [ + "### Root Mean Square Layer Normalization (RMSNorm)\n", + "\n", + "RMSNorm 是对 LayerNorm 的一个改进, 没有做 re-center 操作（移除了均值项），可以看作 LayerNorm 在均值为零时的特例，使用平方根均值归一化降低噪声影响。\n", + "\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", + "\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", + "\n", + "假设输入张量形状为 (batch_size, sequence_length, embedding_dim)，RMS归一化同样对 embedding_dim 维度进行归一化，其中，$g_i$ 为可学习参数，用于对均方根归一化结果进行加权。" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "c52df1c5-1047-48d8-aa00-c0b576a87632", + "metadata": {}, + "outputs": [], + "source": [ + "class RMSNorm(torch.nn.Module):\n", + " def __init__(self, dim: int, eps: float):\n", + " super().__init__()\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)" + ] + }, + { + "cell_type": "markdown", + "id": "6207c186-9816-41fd-872d-d5116257b274", + "metadata": {}, + "source": [ + "### Rotary Position Embedding, RoPE\n", + "\n", + "旋转位置编码是一种能将相对位置信息集成到 self-attentio 中, 进而提升 transformer 架构性能的位置编码方式, 和绝对位置编码相比, RoPE 具有很好的外推性, 是目前的主流位置编码方式。\n", + "\n", + "外推性的解释, 通俗来说就是训练的时候限制了 512 的上下文长度，那么推理时如果面对超过该长度的文本，LLM 可能无法正确处理。\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", + "\n", + "其中编码向量 $p_i$ 的计算使用 Sinusoidal 函数：\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", + "\n", + "- 旋转位置编码\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", + "$$f_{\\{q,k\\}}\\left(\\boldsymbol{x}_m,m\\right)=\\boldsymbol{R}_{\\Theta,m}^d\\boldsymbol{W}_{\\{q,k\\}}\\boldsymbol{x}_m$$\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", + "$$\\Theta=\\left\\{\\theta_i=10000^{-2(i-1)/d},i\\in[1,2,\\ldots,d/2]\\right\\}$$\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", + "$$" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "5aae4194-3738-44a7-9c33-bbf964067f8a", + "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", + " 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", + " return pos_cis\n", + "\n", + "def apply_rotary_emb(xq, xk, pos_cis):\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", + " 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_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)" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "5f8376c4-5b00-4c72-9f23-912478db581c", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "(torch.Size([16, 32]), tensor(1.+0.j))" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "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]" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "b40d0ec8-8752-4db9-a4fe-e746d39dc970", + "metadata": {}, + "outputs": [ + { + "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" + ] + } + ], + "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" + ] + }, + { + "cell_type": "markdown", + "id": "fa47eeb0-874e-4df3-bf91-5f79109af8ca", + "metadata": {}, + "source": [ + "### Attention\n", + "\n", + "Attention 是 Transformer 架构中的重要机制，它可以捕获长序列间的依赖关系，并依据注意力得分进行重要性建模。\n", + "\n", + "在 MiniMind 模型中，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", + "\n", + "Group Querey Attention (GQA) 是对多头自注意力机制的扩展，通过提供计算效率和模型表达能力的灵活权衡，实现了查询头的分组。\n", + "\n", + "具体来说，GQA 将查询头分为 G 个组，每个组共享一个公共的键和值。\n", + "\n", + "![img](./demo/gqa.png)\n", + "\n", + "- KV Cache\n", + "\n", + "KV Cache 能够有效压缩大模型推理时的显存占用，在推理时，前面生成的字符不需要与后面的字符产生 attention，从而使得前面已经计算的 K 和 V 可以缓存起来。\n", + "\n", + "- SwiGLU\n", + "\n", + "$$\\text{SwiGLU}(x,W,V,b,c)=\\text{Swish}_1(xW+b)\\otimes(xV+c)$$" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "e61394db-5afa-4421-ba77-30a9c50f0603", + "metadata": {}, + "outputs": [], + "source": [ + "from demo.LMConfig import LMConfig\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", + " 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", + " )\n", + "\n", + "class Attention(nn.Module):\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", + " self.n_local_heads = args.n_heads\n", + " 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", + " 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", + "\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", + " ############## 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", + " ################### 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", + " )\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", + " 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 = self.attn_dropout(scores)\n", + " output = scores @ xv\n", + " ################################################\n", + " output = output.transpose(1, 2).reshape(bsz, seq_len, -1)\n", + " output = self.resid_dropout(self.wo(output))\n", + " return output, past_kv" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "ca9f6675-7d70-49b4-ba02-f91109869fd8", + "metadata": {}, + "outputs": [ + { + "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" + ] + } + ], + "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)" + ] + }, + { + "cell_type": "markdown", + "id": "495f1e81-36de-4e94-8ae4-69490f62d15f", + "metadata": {}, + "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": "4fe2abd9-0c40-44be-b64d-8b0caa28604b", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "torch.Size([4, 16, 512])" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "ffn = FeedForward(LMConfig())\n", + "x = torch.randn((4, 16, 512))\n", + "output = ffn(x)\n", + "output.shape" + ] + }, + { + "cell_type": "markdown", + "id": "2317829c-39b8-44d5-beb2-1b9727bc35c8", + "metadata": {}, + "source": [ + "### MiniMind Block\n", + "\n", + "到目前为止，，已经完成了 Attention Layer 和 FeedForward Layer 的构建，所有必须的组件都已经具备，我们着手构建一个 MiniMind Block" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "d6c224a4-7ba5-4cf8-9297-311c660f5764", + "metadata": {}, + "outputs": [], + "source": [ + "class MiniMindBlock(nn.Module):\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.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.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", + " 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" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "d23b4461-534c-4f6d-a00b-cdcfbd4d7bec", + "metadata": {}, + "outputs": [ + { + "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" + ] + }, + { + "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" + ] + }, + { + "cell_type": "markdown", + "id": "f9bb3b2e-4d1d-4721-9908-546f3c49dde3", + "metadata": {}, + "source": [ + "### MiniMindLM\n", + "\n", + "以 MiniMind Block 为基本组件，我们可以对该 LLM 进行最后组装！" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "7267245b-f13a-4bcf-9358-5fdd89106153", + "metadata": {}, + "outputs": [], + "source": [ + "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", + " self.params = params or LMConfig()\n", + " super().__init__(self.params)\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.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.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", + " 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", + " **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", + " 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", + " use_cache=use_cache\n", + " )\n", + " past_kvs.append(past_kv)\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", + " 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", + " if stream:\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", + " 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", + " for seq in generated\n", + " ]\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", + " 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", + " 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_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", + " 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", + " if input_ids_next.item() == eos_token_id:\n", + " break" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "dc7b4d79-a3aa-4143-aa5c-2abce34bc96c", + "metadata": {}, + "outputs": [ + { + "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" + ] + }, + { + "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)" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "3480a860-e934-43de-a621-8063f0bdea2f", + "metadata": {}, + "outputs": [ + { + "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" + ] + } + ], + "source": [ + "# 我们让 MiniMind 根据我们设计的 四个输入 token 生成输出\n", + "out = MiniMind.generate(input_ids, max_new_tokens=8, use_cache=True)\n", + "print(out.shape, out)" + ] + }, + { + "cell_type": "markdown", + "id": "b0e4ef5f-1358-4ffb-a150-61141ac3062f", + "metadata": {}, + "source": [ + "## Minimind MOE Model" + ] + }, + { + "cell_type": "markdown", + "id": "e54ced18-db3a-478b-9578-9825634f3d66", + "metadata": {}, + "source": [ + "@TODO" + ] + } + ], + "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/2_dataset.ipynb b/notebook/2_dataset.ipynb new file mode 100644 index 0000000..3a4499a --- /dev/null +++ b/notebook/2_dataset.ipynb @@ -0,0 +1,574 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "44db347d-fef7-4e56-aeab-ce55bdab8eab", + "metadata": {}, + "source": [ + "# 2-Dataset\n", + "\n", + "到这里我们便完成了对于 MiniMind Tokenizer 和 Model 部分的全部了解，我们所熟悉的大语言模型正是由这个组件构成的，接下来，我们需要对大模型训练所使用的数据集结构有个基本的认识。\n", + "\n", + "想要训练一个能够正常对话，并且符合人类对话偏好的大模型一般需要经过以下几个训练阶段：\n", + "\n", + "- 预训练（Pre-training）\n", + "- 有监督微调（Supervised Fine-tuning，SFT）\n", + "- 人类反馈强化学习（Reinforcement Learning from Human Feedback，RLHF）\n", + "\n", + "在不同训练阶段使用的数据集有所不同，下面会从 MiniMind 代码出发进行介绍和解读。" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "e41156fd-ca0e-4083-a7a4-a5f98019ab88", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import random\n", + "import re\n", + "\n", + "import pandas as pd\n", + "import numpy as np\n", + "from torch.utils.data import Dataset, DataLoader\n", + "import torch\n", + "from sklearn.model_selection import train_test_split\n", + "import os\n", + "import ast\n", + "from transformers import AutoTokenizer\n", + "\n", + "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "27d18ad2-3f71-435c-af03-ad9d347aa048", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "6400\n" + ] + } + ], + "source": [ + "# 从 ../model 目录加载分词器\n", + "\n", + "tokenizer = AutoTokenizer.from_pretrained('../model/minimind_tokenizer')\n", + "print(tokenizer.vocab_size)" + ] + }, + { + "cell_type": "markdown", + "id": "ad08fe92-1d04-4b9c-b689-6ad9d763a7c2", + "metadata": {}, + "source": [ + "## 预训练数据集\n", + "\n", + "预训练是模型在大规模语料上进行无监督学习的训练阶段，在该阶段，模型主要学习下一词预测的能力，简单的来说就是学会说话，而不是胡言乱语。因此，该阶段训练的模型不会具有问答能力，而是根据用户输入进行简单的词语接龙。\n", + "\n", + "我们可以看一看预训练的数据集格式：\n", + "\n", + "```\n", + "{\"text\": \"如何才能摆脱拖延症？ 治愈拖延症并不容易，但以下建议可能有所帮助...\"}\n", + "```\n", + "\n", + "为了降低该 demo 的运行门槛，在 `./demo` 文件夹下提供了包含两条训练数据的 `pretrain_data.jsonl` 文件作为熟悉训练流程的数据集 demo。" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "cdcac4aa-0ced-4a74-bba0-bc248923c7bc", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Row 1: {'text': 'LLM首先要学习的并非直接与人交流，而是让网络参数中充满知识的墨水，“墨水” 理论上喝的越饱越好，产生大量的对世界的知识积累。 预训练就是让Model先埋头苦学大量基本的知识，例如从Wiki百科、新闻、书籍整理大规模的高质量训练数据。 这个过程是“无监督”的，即人类不需要在过程中做任何“有监督”的校正，而是由模型自己从大量文本中总结规律学习知识点。 模型此阶段目的只有一个：学会词语接龙。例如我们输入“秦始皇”四个字，它可以接龙“是中国的第一位皇帝”。'}\n", + "\n", + "Row 2: {'text': '经过预训练，LLM此时已经掌握了大量知识，然而此时它只会无脑地词语接龙，还不会与人聊天。 SFT阶段就需要把半成品LLM施加一个自定义的聊天模板进行微调。 例如模型遇到这样的模板【问题->回答，问题->回答】后不再无脑接龙，而是意识到这是一段完整的对话结束。 称这个过程为指令微调，就如同让已经学富五车的「牛顿」先生适应21世纪智能手机的聊天习惯，学习屏幕左侧是对方消息，右侧是本人消息这个规律。 在训练时，MiniMind的指令和回答长度被截断在512，是为了节省显存空间。就像我们学习时，会先从短的文章开始，当学会写作200字作文后，800字文章也可以手到擒来。 在需要长度拓展时，只需要准备少量的2k/4k/8k长度对话数据进行进一步微调即可（此时最好配合RoPE-NTK的基准差值）。'}\n", + "\n" + ] + } + ], + "source": [ + "# 我们可以查看一下 demo 中提供的数据\n", + "path_pretrain = './demo/pretrain_data.jsonl'\n", + "\n", + "with open(path_pretrain, 'r', encoding='utf-8') as f:\n", + " for line_num, line in enumerate(f, 1):\n", + " data = json.loads(line.strip())\n", + " print(f'Row {line_num}: {data}\\n')" + ] + }, + { + "cell_type": "markdown", + "id": "0c0707e3-3ff9-424b-a5be-dab9db55aa44", + "metadata": {}, + "source": [ + "我们知道，构建一个深度学习数据集需要继承 `torch.utils.data.dataset`，并构建 DataLoader 数据迭代器进行迭代访问。下面，我们来看看 MiniMind 是如何抽象一个预训练数据集的。" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "fb2458e3-3661-40de-96be-aa036291a8fb", + "metadata": {}, + "outputs": [], + "source": [ + "class PretrainDataset(Dataset):\n", + " def __init__(self, data_path, tokenizer, max_length=512):\n", + " super().__init__()\n", + " self.tokenizer = tokenizer\n", + " self.max_length = max_length\n", + " self.samples = self.load_data(data_path)\n", + "\n", + " def load_data(self, path):\n", + " \"\"\"按行读取 jsonl 文件，并存储在列表中\"\"\"\n", + " samples = []\n", + " with open(path, 'r', encoding='utf-8') as f:\n", + " for line_num, line in enumerate(f, 1):\n", + " data = json.loads(line.strip())\n", + " samples.append(data)\n", + " return samples\n", + "\n", + " def __len__(self):\n", + " return len(self.samples)\n", + "\n", + " def __getitem__(self, index):\n", + " sample = self.samples[index]\n", + "\n", + " text = f\"{self.tokenizer.bos_token}{str(sample['text'])}{self.tokenizer.eos_token}\"\n", + " # print(text) # uncomment to see formating prompt\n", + " encoding = self.tokenizer(\n", + " text,\n", + " max_length=self.max_length,\n", + " padding='max_length',\n", + " truncation=True,\n", + " return_tensors='pt'\n", + " )\n", + " # print(encoding) # uncomment to see encoding result\n", + " input_ids = encoding.input_ids.squeeze()\n", + " loss_mask = (input_ids != self.tokenizer.pad_token_id) # mask to ignore loss on pad token\n", + "\n", + " X = torch.tensor(input_ids[:-1], dtype=torch.long) # token not included\n", + " Y = torch.tensor(input_ids[1:], dtype=torch.long) # token not included\n", + " loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long) # align to tensor X\n", + " return X, Y, loss_mask" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "677a1bbb-0df7-45e5-a3ed-d4431cddf43e", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "2\n" + ] + } + ], + "source": [ + "pretrain_dataset = PretrainDataset(path_pretrain, tokenizer)\n", + "print(len(pretrain_dataset))\n", + "# x, y, lm = pretrain_dataset[0]\n", + "# print(x.shape, y.shape, lm.shape)" + ] + }, + { + "cell_type": "markdown", + "id": "ebb160cb-0e96-45a7-a1c0-be145895e70f", + "metadata": {}, + "source": [ + "## 有监督微调数据集\n", + "\n", + "有监督微调（Supervised Fine Tuning，SFT）对预训练后得到的基座 LLM 施加一个自定义聊天模板进行微调，由于在这一阶段，模型训练的目标是根据用户指令生成响应（构建问答体系），故又称为指令微调。\n", + "\n", + "我们可以看一看有监督微调的数据集格式：\n", + "\n", + "```\n", + "{\n", + " \"conversations\": [\n", + " {\"role\": \"user\", \"content\": \"你好\"},\n", + " {\"role\": \"assistant\", \"content\": \"你好！\"},\n", + " {\"role\": \"user\", \"content\": \"再见\"},\n", + " {\"role\": \"assistant\", \"content\": \"再见！\"}\n", + " ]\n", + "}\n", + "```\n", + "\n", + "为了降低该 demo 的运行门槛，在 `./demo` 文件夹下提供了包含两条 conversation 问答数据的 `sft_data.jsonl` 文件作为熟悉训练流程的数据集 demo。" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "7e025e25-c937-4468-ae66-b32447440658", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Row 1: {'conversations': [{'role': 'user', 'content': '你好吗？'}, {'role': 'assistant', 'content': '我很好，谢谢！你呢？'}, {'role': 'user', 'content': '我也很好，谢谢！'}, {'role': 'assistant', 'content': '太好了！祝你今天愉快！'}]}\n", + "\n", + "Row 2: {'conversations': [{'role': 'user', 'content': '你喜欢什么运动？'}, {'role': 'assistant', 'content': '我喜欢跑步和游泳。你呢？'}, {'role': 'user', 'content': '我喜欢打篮球！'}, {'role': 'assistant', 'content': '篮球很棒！是一个很好的团队运动。'}]}\n", + "\n" + ] + } + ], + "source": [ + "# 我们可以查看一下 demo 中提供的数据\n", + "path_sft = './demo/sft_data.jsonl'\n", + "\n", + "with open(path_sft, 'r', encoding='utf-8') as f:\n", + " for line_num, line in enumerate(f, 1):\n", + " data = json.loads(line.strip())\n", + " print(f'Row {line_num}: {data}\\n')" + ] + }, + { + "cell_type": "markdown", + "id": "5f988233-530f-4ddf-8d64-9532442a2535", + "metadata": {}, + "source": [ + "接下来，我们尝试构造一个数据集对象，实现对 sft 格式数据的读取与处理。" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "8b25f0c6-b5b5-47b1-974f-15a2f0d2760e", + "metadata": {}, + "outputs": [], + "source": [ + "class SFTDataset(Dataset):\n", + " def __init__(self, jsonl_path, tokenizer, max_length=512):\n", + " super().__init__()\n", + " self.tokenizer = tokenizer\n", + " self.max_length = max_length\n", + " self.samples = self.load_data(jsonl_path)\n", + " self.bos_id = tokenizer('assistant\\n', add_special_tokens=False).input_ids # set bos token\n", + " self.eos_id = tokenizer('\\n', add_special_tokens=False).input_ids # set eos token\n", + "\n", + " def load_data(self, path):\n", + " samples = []\n", + " with open(path, 'r', encoding='utf-8') as f:\n", + " for line_num, line in enumerate(f, 1):\n", + " data = json.loads(line.strip())\n", + " samples.append(data)\n", + " return samples\n", + " \n", + " def __len__(self):\n", + " return len(self.samples)\n", + "\n", + " def _create_chat_prompt(self, conversations):\n", + " \"\"\"构建符合 ChatML 格式的对话\"\"\"\n", + " messages = []\n", + " for i, turn in enumerate(conversations): # for each speaker in one conversation\n", + " role = 'user' if i % 2 == 0 else 'assistant'\n", + " messages.append({\"role\": role, \"content\": turn['content']})\n", + " return self.tokenizer.apply_chat_template(\n", + " messages,\n", + " tokenize=False,\n", + " add_generation_prompt=False\n", + " )\n", + "\n", + " def _generate_loss_mask(self, input_ids):\n", + " loss_mask = [0] * len(input_ids)\n", + " i = 0\n", + " while i < len(input_ids):\n", + " if input_ids[i:i + len(self.bos_id)] == self.bos_id: # check if reach bos token\n", + " ########### find content start & end point ###########\n", + " start = i + len(self.bos_id) # set start point at the end of bos token\n", + " end = start\n", + " while end < len(input_ids):\n", + " if input_ids[end:end + len(self.eos_id)] == self.eos_id: # check if reach eos token\n", + " break\n", + " end += 1\n", + " ######################################################\n", + " for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)): # ignore tokens that reach max input length\n", + " loss_mask[j] = 1\n", + " i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids) # update i to exit current conversation turn\n", + " else:\n", + " i += 1\n", + " return loss_mask\n", + " \n", + " 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", + " 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", + " \n", + " loss_mask = self._generate_loss_mask(input_ids)\n", + "\n", + " X = torch.tensor(input_ids[:-1], dtype=torch.long)\n", + " Y = torch.tensor(input_ids[1:], dtype=torch.long)\n", + " loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long) # align with pred pos\n", + "\n", + " return X, Y, loss_mask" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "67a62dbf-6a50-4faf-8ad1-5228eb8aa7c9", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "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" + ] + } + ], + "source": [ + "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(lm) # 打印 loss mask，你会发现在序列中有两处以 1 填充的序列，这是因为我们在一个 conversation 中开展了两轮对话，其中只有 assistant 回复计算损失" + ] + }, + { + "cell_type": "markdown", + "id": "a7094ff6-4657-4483-8c34-785ffe9adbd4", + "metadata": {}, + "source": [ + "## 人类反馈强化学习数据集\n", + "\n", + "在 MiniMind 项目中，采用直接偏好优化（Direct Parameter Optimization，DPO）训练大模型对齐人类偏好。在这一训练阶段，模型将会根据提供的问答正反例进行偏好优化，从而降低让人类不满意的答案出现的几率。\n", + "\n", + "与PPO(Proximal Policy Optimization)这种需要奖励模型、价值模型的RL算法不同； DPO通过推导PPO奖励模型的显式解，把在线奖励模型换成离线数据，Ref模型输出可以提前保存。 DPO性能几乎不变，只用跑 actor_model 和 ref_model 两个模型，大大节省显存开销和增加训练稳定性。\n", + "\n", + "我们可以看一看有监督微调的数据集格式：\n", + "\n", + "```\n", + "{\n", + " \"chosen\": [\n", + " {\"content\": \"Query\", \"role\": \"user\"}, \n", + " {\"content\": \"good answer\", \"role\": \"assistant\"}\n", + " ], \n", + " \"rejected\": [\n", + " {\"content\": \"Query\", \"role\": \"user\"}, \n", + " {\"content\": \"bad answer\", \"role\": \"assistant\"}\n", + " ]\n", + "}\n", + "```\n", + "\n", + "为了降低该 demo 的运行门槛，在 ./demo 文件夹下提供了包含两条 conversation 问答数据的 sft_data.jsonl 文件作为熟悉训练流程的数据集 demo。" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "81682b5f-0d3e-4fa2-aec4-0be2b19d5c05", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Row 1: {'conversations': [{'role': 'user', 'content': '你好吗？'}, {'role': 'assistant', 'content': '我很好，谢谢！你呢？'}, {'role': 'user', 'content': '我也很好，谢谢！'}, {'role': 'assistant', 'content': '太好了！祝你今天愉快！'}]}\n", + "\n", + "Row 2: {'conversations': [{'role': 'user', 'content': '你喜欢什么运动？'}, {'role': 'assistant', 'content': '我喜欢跑步和游泳。你呢？'}, {'role': 'user', 'content': '我喜欢打篮球！'}, {'role': 'assistant', 'content': '篮球很棒！是一个很好的团队运动。'}]}\n", + "\n" + ] + } + ], + "source": [ + "# 我们可以查看一下 demo 中提供的数据\n", + "path_dpo = './demo/dpo_data.jsonl'\n", + "\n", + "with open(path_sft, 'r', encoding='utf-8') as f:\n", + " for line_num, line in enumerate(f, 1):\n", + " data = json.loads(line.strip())\n", + " print(f'Row {line_num}: {data}\\n')" + ] + }, + { + "cell_type": "markdown", + "id": "ebb54557-54f4-4004-8a7b-443b8671002b", + "metadata": {}, + "source": [ + "接下来，我们尝试构造 json 对象，实现对 dpo 格式数据的读取和处理" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "efa98777-5d3a-4b3f-a9fb-e0f05e4af81d", + "metadata": {}, + "outputs": [], + "source": [ + "class DPODataset(Dataset):\n", + " def __init__(self, file_path, tokenizer, max_length=512):\n", + " super().__init__()\n", + " self.tokenizer = tokenizer\n", + " self.max_length = max_length\n", + " self.padding = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0\n", + " self.bos_id = tokenizer('assistant\\n', add_special_tokens=False).input_ids # content prefix\n", + " self.eos_id = tokenizer('\\n', add_special_tokens=False).input_ids # content suffix\n", + " with open(file_path, 'r', encoding='utf-8') as f: # load data\n", + " self.data = []\n", + " for line in f:\n", + " line = line.strip()\n", + " obj = json.loads(line)\n", + " self.data.append(obj)\n", + "\n", + " def __len__(self):\n", + " return len(self.data)\n", + "\n", + " def _generate_loss_mask(self, input_ids):\n", + " \"\"\"此处的损失掩码生成函数与 SFT 阶段逻辑一致\"\"\"\n", + " loss_mask = [0] * len(input_ids)\n", + " i = 0\n", + " while i < len(input_ids):\n", + " if input_ids[i:i + len(self.bos_id)] == self.bos_id:\n", + " start = i + len(self.bos_id)\n", + " end = start\n", + " while end < len(input_ids):\n", + " if input_ids[end:end + len(self.eos_id)] == self.eos_id:\n", + " break\n", + " end += 1\n", + " for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)):\n", + " loss_mask[j] = 1\n", + " i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)\n", + " else:\n", + " i += 1\n", + " return loss_mask\n", + " \n", + " def __getitem__(self, index):\n", + " item = self.data[index]\n", + " chosen = item['chosen'] # 一个 list，里面包含若干 {role, content}\n", + " rejected = item['rejected'] # 同上\n", + " # format prompt\n", + " chosen_prompt = self.tokenizer.apply_chat_template(\n", + " chosen, tokenize=False, add_generation_prompt=False\n", + " )\n", + " print(chosen_prompt) # uncomment to see formating prompt\n", + " rejected_prompt = self.tokenizer.apply_chat_template(\n", + " rejected, tokenize = False, add_gerneration_prompt=False\n", + " )\n", + " print(rejected_prompt) # uncomment to see formating prompt\n", + " # tokenize\n", + " chosen_encoding = self.tokenizer(\n", + " chosen_prompt, truncation=True, max_length=self.max_length, padding='max_length'\n", + " )\n", + " rejected_encoding = self.tokenizer(\n", + " rejected_prompt, truncation=True, max_length=self.max_length, padding='max_length'\n", + " )\n", + " # generate loss mask\n", + " chosen_input_ids = chosen_encoding['input_ids']\n", + " chosen_loss_mask = self._generate_loss_mask(chosen_input_ids)\n", + " rejected_input_ids = rejected_encoding['input_ids']\n", + " rejected_loss_mask = self._generate_loss_mask(rejected_input_ids)\n", + " # same as sft / pretrain\n", + " x_chosen = torch.tensor(chosen_input_ids[:-1], dtype=torch.long)\n", + " y_chosen = torch.tensor(chosen_input_ids[1:], dtype=torch.long)\n", + " mask_chosen = torch.tensor(chosen_loss_mask[1:], dtype=torch.long)\n", + " x_rejected = torch.tensor(rejected_input_ids[:-1], dtype=torch.long)\n", + " y_rejected = torch.tensor(rejected_input_ids[1:], dtype=torch.long)\n", + " mask_rejected = torch.tensor(rejected_loss_mask[1:], dtype=torch.long)\n", + "\n", + " return {\n", + " 'x_chosen': x_chosen,\n", + " 'y_chosen': y_chosen,\n", + " 'mask_chosen': mask_chosen,\n", + " 'x_rejected': x_rejected,\n", + " 'y_rejected': y_rejected,\n", + " 'mask_rejected': mask_rejected\n", + " }" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "5ba205c4-0dd2-48fd-8aa2-26a42afc90d8", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "2\n", + "system\n", + "你是 MiniMind，是一个有用的人工智能助手。\n", + "user\n", + "你好吗？\n", + "assistant\n", + "我很好，谢谢！你呢？\n", + "user\n", + "今天过得怎么样？\n", + "assistant\n", + "挺好的，去跑步了，心情不错。\n", + "\n", + "system\n", + "你是 MiniMind，是一个有用的人工智能助手。\n", + "user\n", + "你好吗？\n", + "assistant\n", + "不好，我很累。\n", + "user\n", + "你喜欢什么运动？\n", + "assistant\n", + "我不喜欢运动，没兴趣。\n", + "\n" + ] + } + ], + "source": [ + "dpo_dataset = DPODataset(path_dpo, tokenizer)\n", + "print(len(dpo_dataset))\n", + "res = dpo_dataset[0]\n", + "# 如有需要，请自主查看 res 中的元素" + ] + } + ], + "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/3_pretrain.ipynb b/notebook/3_pretrain.ipynb new file mode 100644 index 0000000..5f3b82a --- /dev/null +++ b/notebook/3_pretrain.ipynb @@ -0,0 +1,372 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "aba635c1-cdcd-4d00-a123-00045edbd10f", + "metadata": {}, + "source": [ + "# 3-Pretrain\n", + "\n", + "预训练是模型经历的第一个阶段，在该阶段，模型将会吸收知识，学习尽可能正确的下一词语预测范式。在这个笔记本中，我们仅对预训练的训练流程进行展示和学习，因此只给出必要的代码片段，如 wandb 和 ddp 不会在此笔记本中涉及。\n", + "\n", + "此笔记本的完整实现见主仓库 `/minimind/train_pretrain.py`" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "98cb706f-2e03-4d8e-8e45-3663c9ab3b64", + "metadata": {}, + "outputs": [], + "source": [ + "# 导入依赖\n", + "import os\n", + "import platform\n", + "import argparse\n", + "import time\n", + "import math\n", + "import warnings\n", + "import pandas as pd\n", + "import torch\n", + "import torch.distributed as dist\n", + "from torch import optim, nn\n", + "from torch.nn.parallel import DistributedDataParallel\n", + "from torch.optim.lr_scheduler import CosineAnnealingLR\n", + "from torch.utils.data import DataLoader, DistributedSampler\n", + "from contextlib import nullcontext\n", + "\n", + "from transformers import AutoTokenizer\n", + "\n", + "from demo.model import MiniMindLM\n", + "from demo.LMConfig import LMConfig\n", + "from demo.dataset import PretrainDataset" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "c3b6f015-0e5c-4e88-bb0d-32092993933f", + "metadata": {}, + "outputs": [], + "source": [ + "warnings.filterwarnings('ignore')" + ] + }, + { + "cell_type": "markdown", + "id": "a46c9b5b-12eb-4cad-a859-c5101cdcba82", + "metadata": {}, + "source": [ + "## 可选参数设置\n", + "\n", + "首先，查看训练的可选参数，这些参数在实际使用时通过命令行导入，为了保持笔记本的易用性，选择用 class 进行包装" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "3f07e381-961e-4a4f-9f37-ff2f0bc2918c", + "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/pretrain_data.jsonl' # 数据集路径" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "8f73286e-48a0-4a5b-b7c4-8c5fefea5f14", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'cuda'" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "args.device" + ] + }, + { + "cell_type": "markdown", + "id": "7e4b0ea8-0f41-47fe-8125-93629d03d451", + "metadata": {}, + "source": [ + "## 初始化训练\n", + "\n", + "\n", + "接下来，我们对一些重要模块进行初始化，我们已经了解过，分词器，模型和数据集是大模型的基本组件，我们对其进行初始化。" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "9818e376-7f1e-4347-b8c3-c8e30f7d587f", + "metadata": {}, + "outputs": [], + "source": [ + "def init_model(lm_config):\n", + " tokenizer = AutoTokenizer.from_pretrained('../model/minimind_tokenizer')\n", + " model = MiniMindLM(lm_config).to(args.device)\n", + " print(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')\n", + " return model, tokenizer" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "24d15001-b670-4333-9fbe-b224ab1993c5", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "LLM总参数量：8.915 百万\n", + "cuda:0 6400 \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 = PretrainDataset(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(model.device, tokenizer.vocab_size, train_loader)" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "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" + } + ], + "source": [ + "loader = iter(train_loader)\n", + "next(loader), len(train_ds), len(loader)" + ] + }, + { + "cell_type": "markdown", + "id": "0d91904d-7ddf-4105-a5f5-65b84315a73b", + "metadata": {}, + "source": [ + "我们发现，train loader 的每一个 iter 都包含一个长度为 3 的张量列表，这是因为 train_dataset 每一次取数据都会返回三个张量，分别为:\n", + "\n", + "- 样本 X: 包含 \\ 在内的输入 content\n", + "- 标签 Y: 包含 \\ 在内的输出 content\n", + "- 掩码 loss_mask: 指示需要计算损失的 token 位置\n", + "\n", + "由于我们的数据集只有两条数据，而 batch size 设置为 2，因此我们的 dataloader 只有一个 iter." + ] + }, + { + "cell_type": "markdown", + "id": "40764309-9c98-45d5-ada5-e31c1aaa02c5", + "metadata": {}, + "source": [ + "## 启动训练\n", + "\n", + "训练一个深度学习模型，还涉及到了优化器，损失函数和学习率调度。接下来，我们查看 MiniMind 训练部分的代码，并进行一轮简单的训练。" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "4f70b9cf-fcff-4242-bd96-ce34a5da55e4", + "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", + "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": "c101ffb4-1a97-4e29-8612-caf40a8aa05d", + "metadata": {}, + "source": [ + "接下来，我们来看看 MiniMind 的训练函数" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "03043d8e-e90b-42a8-ba4d-de6187f6eb81", + "metadata": {}, + "outputs": [], + "source": [ + "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", + "\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()) # 取生成的最后一个 token 的 logits 计算损失\n", + " loss = (loss * loss_mask).sum() / loss_mask.sum()\n", + " loss += res.aux_loss # 若为混合专家 则将 MOE 辅助损失纳入考虑\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() * args.accumulation_steps,\n", + " optimizer.param_groups[-1]['lr'],\n", + " spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))\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.use_moe else ''\n", + " # ckp = f'{args.save_dir}/pretrain_{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": "7792bf9c-95a1-4c0d-bb53-fac2a1132ee0", + "metadata": {}, + "source": [ + "准备完毕，我们进行一轮长度 1 个 iter 的训练看看怎么个事。" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "17f51549-2009-4322-86c8-7aa783169e0b", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Epoch:[1/1](0/1) loss:8.889 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)" + ] + } + ], + "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/4_sft.ipynb b/notebook/4_sft.ipynb new file mode 100644 index 0000000..8313ca6 --- /dev/null +++ b/notebook/4_sft.ipynb @@ -0,0 +1,382 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "70560a4d-ff57-4699-afb2-52f6a145b97e", + "metadata": {}, + "source": [ + "# 4-SFT\n", + "\n", + "在预训练阶段后，我们应该能够获得一个下一词预测模型，此时的模型已经掌握了大量的知识。不过，仅仅具备下一词预测能力是不够的，我们希望大模型能够获得问答能力，这一能力便是在有监督微调（Supervised Fine Tuning，SFT）阶段获得的。\n", + "\n", + "在这个笔记本中，我们仅对 SFT 的训练流程进行展示和学习，因此只给出必要的代码片段，如 wandb 和 ddp 不会在此笔记本中涉及。\n", + "\n", + "此笔记本的完整实现见主仓库 `/minimind/train_full_sft.py`" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "2508c8c7-eb44-4ac0-918c-41479f41d1ac", + "metadata": {}, + "outputs": [], + "source": [ + "# 导入依赖\n", + "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" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "4a79e140-23c7-4498-9a4d-ef7459692996", + "metadata": {}, + "outputs": [], + "source": [ + "warnings.filterwarnings('ignore')" + ] + }, + { + "cell_type": "markdown", + "id": "dbe1d40e-e074-4856-8c67-da8019f250ab", + "metadata": {}, + "source": [ + "## 可选参数设置\n", + "\n", + "首先，查看训练的可选参数，这些参数在实际使用时通过命令行导入，为了保持笔记本的易用性，选择用 class 进行包装" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "98cbd6cf-081f-4c31-8dce-52566004801a", + "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": 4, + "id": "642f5107-aa9b-49cc-9759-e8731326edc3", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'cuda'" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "args.device" + ] + }, + { + "cell_type": "markdown", + "id": "b70fadb6-f78f-45a0-ad3e-d41d7253c72b", + "metadata": {}, + "source": [ + "## 初始化训练\n", + "\n", + "接下来，我们对一些重要模块进行初始化，我们已经了解过，分词器，模型和数据集是大模型的基本组件，我们对其进行初始化。\n", + "\n", + "> 注意 与预训练阶段不同的是 在 sft 阶段 我们实际上是在上一阶段训练获得的模型的基础上修改数据集进行接续训练 因此需要载入上一阶段的模型权重 出于展示的目的 载入权重的代码在此笔记本中只作展示 并不执行" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "c64a0e37-a9d0-4ce6-b5ee-7b2728de879b", + "metadata": {}, + "outputs": [], + "source": [ + "def init_model(lm_config):\n", + " tokenizer = AutoTokenizer.from_pretrained('../model/minimind_tokenizer')\n", + " model = MiniMindLM(lm_config).to(args.device)\n", + " moe_path = '_moe' if lm_config.use_moe else ''\n", + " # ckp = f'./out/pretrain_{lm_config.dim}{moe_path}.pth' # 指示上一阶段训练保存的模型文件位置\n", + " # state_dict = torch.load(ckp, map_location=args.device) # 载入模型状态字典\n", + " # model.laod_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": 6, + "id": "2eec33f3-2ac2-4c75-b58e-844d89d8add9", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "LLM总参数量：8.915 百万\n", + "cuda:0 6400 \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", + " num_workers=args.num_workers,\n", + ")\n", + "\n", + "print(model.device, tokenizer.vocab_size, train_loader)" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "b02644e0-e61d-4df0-9203-fa841cd83635", + "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" + } + ], + "source": [ + "loader = iter(train_loader)\n", + "next(loader), len(train_ds), len(loader)" + ] + }, + { + "cell_type": "markdown", + "id": "13ac8a47-d26a-49da-82fb-468ef5ac91a4", + "metadata": {}, + "source": [ + "我们发现，train loader 的每一个 iter 都包含一个长度为 3 的张量列表，这是因为 train_dataset 每一次取数据都会返回三个张量，分别为:\n", + "\n", + "- 样本 X: 包含 \\ 在内的输入 conversation\n", + "- 标签 Y: 包含 \\ 在内的输出 conversation\n", + "- 掩码 loss_mask: 指示需要计算损失的 token 位置\n", + "\n", + "由于我们的数据集只有两条数据，而 batch size 设置为 2，因此我们的 dataloader 只有一个 iter." + ] + }, + { + "cell_type": "markdown", + "id": "ee23f937-f333-4c02-a3b4-a768818c0e2b", + "metadata": {}, + "source": [ + "## 启动训练\n", + "\n", + "训练一个深度学习模型，还涉及到了优化器，损失函数和学习率调度。接下来，我们查看 MiniMind 训练部分的代码，并进行一轮简单的训练。\n", + "\n", + "> 不难发现 pretrain 阶段和 sft 阶段的训练主体差不多 因为这两个阶段的差异体现在数据集格式 而数据集在经过 chat template 格式化后差异小了很多" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "5c32a59a-be0b-4426-b47e-c522322bef69", + "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", + "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": "13963ec0-d22e-4120-b0f2-8f5e5abfcda1", + "metadata": {}, + "source": [ + "接下来，我们来看看 MiniMind 的训练函数" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "74499592-2065-4fff-97d1-762cd08b77b6", + "metadata": {}, + "outputs": [], + "source": [ + "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", + "\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()) # 取生成的最后一个 token 的 logits 计算损失\n", + " loss = (loss * loss_mask).sum() / loss_mask.sum()\n", + " loss += res.aux_loss # 若为混合专家 则将 MOE 辅助损失纳入考虑\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() * args.accumulation_steps,\n", + " optimizer.param_groups[-1]['lr'],\n", + " spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))\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.use_moe else ''\n", + " # ckp = f'{args.save_dir}/pretrain_{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": "546f28c2-418e-453c-88aa-695468f9216c", + "metadata": {}, + "source": [ + "准备完毕，我们进行一轮长度 1 个 iter 的训练看看怎么个事。" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "341200c9-0d3a-4ed1-9dd2-d2fc35d02564", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Epoch:[1/1](0/1) loss:8.865 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)" + ] + } + ], + "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/5_dpo.ipynb b/notebook/5_dpo.ipynb new file mode 100644 index 0000000..9b72797 --- /dev/null +++ b/notebook/5_dpo.ipynb @@ -0,0 +1,456 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "f53f31fb-4795-49cf-9a68-a13f8f7e5cae", + "metadata": {}, + "source": [ + "# 5-DPO\n", + "\n", + "直接偏好优化（Direct Preference Optimization，DPO）是后训练阶段中，使用正反样例激励大模型产生符合人类偏好的回答的策略，为人类反馈强化学习（Reinforcement Learning from Human Feedback, RLHF）提供了一个高效简化的替代方案。通过这一阶段的训练，大模型将会学会依照人类的喜好生成回复。\n", + "\n", + "在这个笔记本中，我们仅对 DPO 的训练流程进行展示和学习，因此只给出必要的代码片段，如 wandb 和 ddp 不会在此笔记本中涉及。\n", + "\n", + "此笔记本的完整实现见主仓库 `/minimind/train_dpo.py`" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "3e09f234-5b8b-4bfb-8d42-bb766a457acc", + "metadata": {}, + "outputs": [], + "source": [ + "# 导入依赖\n", + "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 DPODataset" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "eb946301-eb7e-45c6-9e1c-16cc0ee4db82", + "metadata": {}, + "outputs": [], + "source": [ + "warnings.filterwarnings('ignore')" + ] + }, + { + "cell_type": "markdown", + "id": "d70c65d3-14f7-4569-bebd-cddb48f1851f", + "metadata": {}, + "source": [ + "## 可选参数设置\n", + "\n", + "首先，查看训练的可选参数，这些参数在实际使用时通过命令行导入，为了保持笔记本的易用性，选择用 class 进行包装" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "b78b11a1-9189-4eaf-a094-07c0497c4d38", + "metadata": {}, + "outputs": [], + "source": [ + "class args:\n", + " # out_dir: str = \"out\" # pytorch 格式权重文件保存位置 我们只展示训练过程 所以不使用\n", + " epochs: int = 1 # 训练轮数\n", + " batch_size: int = 2 # pretrain 数据集仅两个样本，设置 batch 为 2\n", + " # sft阶段学习率为 「5e-6」->「5e-7」长度512，建议离线正负样本「概率」偏好对齐阶段lr <=「1e-8」长度3000，否则很容易遗忘训坏\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 = 1 # MiniMind Block 数量 模型超参数 | 由于 dpo 要加载两个模型 我们出于演示目的设定 n_layers = 1\n", + " max_seq_len: int = 512 # 序列长度阈值\n", + " use_moe: bool = False # 是否启用混合专家\n", + " data_path: str = './demo/dpo_data.jsonl' # 数据集路径" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "0a72c8aa-ca60-4b0d-aabb-2c9d16b48175", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'cuda'" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "args.device" + ] + }, + { + "cell_type": "markdown", + "id": "f8fb31ff-524b-4788-a25f-699edb887704", + "metadata": {}, + "source": [ + "## 初始化训练\n", + "\n", + "接下来，我们对一些重要模块进行初始化，我们已经了解过，分词器，模型和数据集是大模型的基本组件，我们对其进行初始化。\n", + "\n", + "> 在这一阶段 我们调整的是大模型的问答偏好 因此与 sft 阶段同理 我们需要载入在 sft 阶段微调好的问答模型" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "5d5971fc-4794-4903-b418-76f266bf19d9", + "metadata": {}, + "outputs": [], + "source": [ + "def init_model(lm_config):\n", + " \"\"\"注意 在此处我们注释了加载模型部分代码 但应当认识到 actor 和 ref 两个模型都具备 sft 阶段获得的权重\"\"\"\n", + " tokenizer = AutoTokenizer.from_pretrained('../model/minimind_tokenizer')\n", + " model = MiniMindLM(lm_config)\n", + " moe_path = '_moe' if lm_config.use_moe else ''\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", + " # 初始化参考模型\n", + " ref_model = MiniMindLM(lm_config)\n", + " # ref_model.load_state_dict(state_dict, strict=False)\n", + " ref_model.eval()\n", + " ref_model.requires_grad_(False)\n", + "\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", + " ref_model = ref_model.to(args.device)\n", + "\n", + " return model, ref_model, tokenizer" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "fbd7bd15-a060-4aba-8c5d-277705ec6b52", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "LLM总参数量：6.096 百万\n", + "cuda:0 6400 \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, ref_model ,tokenizer = init_model(lm_config)\n", + "\n", + "train_ds = DPODataset(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(model.device, tokenizer.vocab_size, train_loader)" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "d807a2af-41e9-4115-b470-a05f6a04613c", + "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" + } + ], + "source": [ + "loader = iter(train_loader)\n", + "next(loader), len(train_ds), len(loader)" + ] + }, + { + "cell_type": "markdown", + "id": "665ad506-ab12-4186-ad63-ada143a33d52", + "metadata": {}, + "source": [ + "我们发现，train loader 的每一个 iter 都包含一个拥有六个键值对的字典，这是因为 train_dataset 每一次取数据都会返回:\n", + "\n", + "- chosen 样本 X: 包含 \\ 在内的输入 content\n", + "- chosen 标签 Y: 包含 \\ 在内的输出 content\n", + "- chosen 掩码 loss_mask: 指示需要计算损失的 token 位置\n", + "- rejected 样本 X: 包含 \\ 在内的输入 content\n", + "- rejected 标签 Y: 包含 \\ 在内的输出 content\n", + "- rejected 掩码 loss_mask: 指示需要计算损失的 token 位置\n", + "\n", + "由于我们的数据集只有两条数据，而 batch size 设置为 2，因此我们的 dataloader 只有一个 iter." + ] + }, + { + "cell_type": "markdown", + "id": "32462eb9-7008-4389-a54b-47ef1c1c4732", + "metadata": {}, + "source": [ + "# 启动训练\n", + "\n", + "训练一个深度学习模型，还涉及到了优化器，损失函数和学习率调度。接下来，我们查看 MiniMind 训练部分的代码，并进行一轮简单的训练。\n", + "\n", + "> DPO 阶段涉及 DPO 损失函数涉及 因此与前两个阶段相比内容略有增加 不过整体流程与逻辑类似" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "1752f117-33c0-4810-886f-fbd1b4d44c28", + "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", + "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": "5f2c3ac1-92bf-4eae-8d22-7fac4cfa69df", + "metadata": {}, + "source": [ + "DPO 的原理是增加偏好样本的对数概率与减小非偏好样本响应的对数概率。\n", + "\n", + "该阶段引入 DPO 损失函数，通过计算选择样本和拒绝样本的对数比率，然后基于这些比率计算 DPO 损失，适用于偏好学习任务。" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "2c07f54a-496e-4b1d-9394-dd2c266ab5d3", + "metadata": {}, + "outputs": [], + "source": [ + "def logits_to_probs(logits, labels):\n", + " # logits shape: (batch_size, seq_len, vocab_size)\n", + " # labels shape: (batch_size, seq_len)\n", + " # probs shape: (batch_size, seq_len)\n", + " log_probs = F.log_softmax(logits, dim=2)\n", + " probs = torch.gather(log_probs, dim=2, index=labels.unsqueeze(2)).squeeze(-1)\n", + " return probs\n", + "\n", + "\n", + "def dpo_loss(ref_probs, probs, beta):\n", + " # ref_probs 和 probs 都是 shape: (batch_size, seq_len)\n", + " # 计算每个样本的平均概率\n", + " ref_probs = ref_probs.mean(dim=1)\n", + " probs = probs.mean(dim=1)\n", + "\n", + " # 将 chosen 和 rejected 数据分开\n", + " batch_size = ref_probs.shape[0]\n", + " chosen_ref_probs = ref_probs[:batch_size // 2]\n", + " reject_ref_probs = ref_probs[batch_size // 2:]\n", + " chosen_probs = probs[:batch_size // 2]\n", + " reject_probs = probs[batch_size // 2:]\n", + "\n", + " # 计算对数比率，比较偏好差异\n", + " pi_logratios = chosen_probs - reject_probs\n", + " ref_logratios = chosen_ref_probs - reject_ref_probs\n", + " logits = pi_logratios - ref_logratios\n", + " loss = -F.logsigmoid(beta * logits)\n", + " return loss.mean()" + ] + }, + { + "cell_type": "markdown", + "id": "e743562c-6e34-488d-bd84-f62eaee0ae20", + "metadata": {}, + "source": [ + "接下来，我们来看看 MiniMind 的训练函数" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "493b7368-d82f-48a6-b5bd-7737b7a23bee", + "metadata": {}, + "outputs": [], + "source": [ + "def train_epoch(epoch):\n", + " start_time = time.time()\n", + " for step, batch in enumerate(train_loader):\n", + " # 提取数据\n", + " x_chosen = batch['x_chosen'].to(args.device)\n", + " x_rejected = batch['x_rejected'].to(args.device)\n", + " y_chosen = batch['y_chosen'].to(args.device)\n", + " y_rejected = batch['y_rejected'].to(args.device)\n", + " mask_chosen = batch['mask_chosen'].to(args.device)\n", + " mask_rejected = batch['mask_rejected'].to(args.device)\n", + " # 正反例拼接\n", + " x = torch.cat([x_chosen, x_rejected], dim=0)\n", + " y = torch.cat([y_chosen, y_rejected], dim=0)\n", + " mask = torch.cat([mask_chosen, mask_rejected], dim=0)\n", + "\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", + " with torch.no_grad(): # 计算 ref 模型输出\n", + " ref_outputs = ref_model(x)\n", + " ref_logits = ref_outputs.logits\n", + " ref_probs = logits_to_probs(ref_logits, y)\n", + " ref_probs = ref_probs * mask # 得到 ref 概率\n", + " outputs = model(x) # 计算 actor 模型输出\n", + " logits = outputs.logits\n", + " probs = logits_to_probs(logits, y)\n", + " probs = probs * mask # 得到 actor 概率\n", + " loss = dpo_loss(ref_probs, probs, beta=0.1) # dpo 损失\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", + " 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:{:.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 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.use_moe else ''\n", + " # ckp = f'{args.save_dir}/rlhf_{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": "49386dc7-2919-428a-a073-f37d576afbcf", + "metadata": {}, + "source": [ + "准备完毕，我们进行一轮长度 1 个 iter 的训练看看怎么个事。" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "0ac86572-ee04-43b5-b5fa-b1c65cbaebec", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Epoch:[1/1](0/1) loss:0.693 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)" + ] + } + ], + "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/LMConfig.py b/notebook/demo/LMConfig.py new file mode 100644 index 0000000..323c110 --- /dev/null +++ b/notebook/demo/LMConfig.py @@ -0,0 +1,61 @@ +from transformers import PretrainedConfig +from typing import List + + +class LMConfig(PretrainedConfig): + model_type = "minimind" + + def __init__( + self, + dim: int = 512, + n_layers: int = 1, # n_layers set to 1 to enable local practice + n_heads: int = 8, + n_kv_heads: int = 2, + vocab_size: int = 6400, + hidden_dim: int = None, + multiple_of: int = 64, + norm_eps: float = 1e-5, + max_seq_len: int = 8192, + rope_theta: int = 1e6, + dropout: float = 0.0, + flash_attn: bool = True, + #################################################### + # Here are the specific configurations of MOE + # When use_moe is false, the following is invalid + #################################################### + use_moe: bool = False, + #################################################### + num_experts_per_tok: int = 2, + n_routed_experts: int = 4, + n_shared_experts: bool = True, + scoring_func: str = 'softmax', + aux_loss_alpha: float = 0.1, + seq_aux: bool = True, + norm_topk_prob: bool = True, + **kwargs, + ): + self.dim = dim + self.n_layers = n_layers + self.n_heads = n_heads + self.n_kv_heads = n_kv_heads + self.vocab_size = vocab_size + self.hidden_dim = hidden_dim + self.multiple_of = multiple_of + self.norm_eps = norm_eps + self.max_seq_len = max_seq_len + self.rope_theta = rope_theta + self.dropout = dropout + self.flash_attn = flash_attn + #################################################### + # Here are the specific configurations of MOE + # When use_moe is false, the following is invalid + #################################################### + self.use_moe = use_moe + self.num_experts_per_tok = num_experts_per_tok # 每个token选择的专家数量 + self.n_routed_experts = n_routed_experts # 总的专家数量 + self.n_shared_experts = n_shared_experts # 共享专家 + self.scoring_func = scoring_func # 评分函数，默认为'softmax' + self.aux_loss_alpha = aux_loss_alpha # 辅助损失的alpha参数 + self.seq_aux = seq_aux # 是否在序列级别上计算辅助损失 + self.norm_topk_prob = norm_topk_prob # 是否标准化top-k概率 + super().__init__(**kwargs) diff --git a/notebook/demo/dataset.py b/notebook/demo/dataset.py new file mode 100644 index 0000000..7750789 --- /dev/null +++ b/notebook/demo/dataset.py @@ -0,0 +1,200 @@ +import json +import random +import re + +import pandas as pd +import numpy as np +from torch.utils.data import Dataset, DataLoader +import torch +from sklearn.model_selection import train_test_split +import os +import ast + +os.environ["TOKENIZERS_PARALLELISM"] = "false" + + +class PretrainDataset(Dataset): + def __init__(self, data_path, tokenizer, max_length=512): + super().__init__() + self.tokenizer = tokenizer + self.max_length = max_length + self.samples = self.load_data(data_path) + + def load_data(self, path): + samples = [] + with open(path, 'r', encoding='utf-8') as f: + for line_num, line in enumerate(f, 1): + data = json.loads(line.strip()) + samples.append(data) + return samples + + def __len__(self): + return len(self.samples) + + def __getitem__(self, index): + sample = self.samples[index] + + # 构建输入文本 + text = f"{self.tokenizer.bos_token}{str(sample['text'])}{self.tokenizer.eos_token}" + encoding = self.tokenizer( + text, + max_length=self.max_length, + padding='max_length', + truncation=True, + return_tensors='pt' + ) + input_ids = encoding.input_ids.squeeze() + loss_mask = (input_ids != self.tokenizer.pad_token_id) + + X = torch.tensor(input_ids[:-1], dtype=torch.long) + Y = torch.tensor(input_ids[1:], dtype=torch.long) + loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long) + return X, Y, loss_mask + + +class SFTDataset(Dataset): + def __init__(self, jsonl_path, tokenizer, max_length=1024): + super().__init__() + self.tokenizer = tokenizer + self.max_length = max_length + self.samples = self.load_data(jsonl_path) + self.bos_id = tokenizer('assistant\n', add_special_tokens=False).input_ids + self.eos_id = tokenizer('\n', add_special_tokens=False).input_ids + + def __len__(self): + return len(self.samples) + + def load_data(self, path): + samples = [] + with open(path, 'r', encoding='utf-8') as f: + for line_num, line in enumerate(f, 1): + data = json.loads(line.strip()) + samples.append(data) + return samples + + def _create_chat_prompt(self, conversations): + """构建符合ChatML格式的对话""" + messages = [] + for i, turn in enumerate(conversations): + role = 'user' if i % 2 == 0 else 'assistant' + messages.append({"role": role, "content": turn['content']}) + return self.tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=False + ) + + def _generate_loss_mask(self, input_ids): + loss_mask = [0] * len(input_ids) + i = 0 + while i < len(input_ids): + if input_ids[i:i + len(self.bos_id)] == self.bos_id: + start = i + len(self.bos_id) + end = start + while end < len(input_ids): + if input_ids[end:end + len(self.eos_id)] == self.eos_id: + break + end += 1 + for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)): + loss_mask[j] = 1 + i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids) + else: + i += 1 + return loss_mask + + def __getitem__(self, index): + sample = self.samples[index] + # 构建对话提示 + prompt = self._create_chat_prompt(sample['conversations']) + input_ids = self.tokenizer(prompt).input_ids[:self.max_length] + input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids)) + + # 生成动态损失掩码 + loss_mask = self._generate_loss_mask(input_ids) + + # 构建训练数据 + X = torch.tensor(input_ids[:-1], dtype=torch.long) + Y = torch.tensor(input_ids[1:], dtype=torch.long) + loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long) # 对齐预测位置 + + return X, Y, loss_mask + + +class DPODataset(Dataset): + def __init__(self, file_path, tokenizer, max_length=4096): + super().__init__() + self.tokenizer = tokenizer + self.max_length = max_length + self.padding = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0 + self.bos_id = tokenizer('assistant\n', add_special_tokens=False).input_ids + self.eos_id = tokenizer('\n', add_special_tokens=False).input_ids + with open(file_path, 'r', encoding='utf-8') as f: + self.data = [] + for line in f: + line = line.strip() + obj = json.loads(line) + self.data.append(obj) + + def __len__(self): + return len(self.data) + + def __getitem__(self, index): + item = self.data[index] + chosen = item['chosen'] # 是一个 list，里面包含若干 {role, content} + rejected = item['rejected'] # 同上 + chosen_prompt = self.tokenizer.apply_chat_template( + chosen, tokenize=False, add_generation_prompt=False + ) + + rejected_prompt = self.tokenizer.apply_chat_template( + rejected, tokenize=False, add_generation_prompt=False + ) + chosen_encoding = self.tokenizer( + chosen_prompt, truncation=True, max_length=self.max_length, padding='max_length' + ) + rejected_encoding = self.tokenizer( + rejected_prompt, truncation=True, max_length=self.max_length, padding='max_length' + ) + + chosen_input_ids = chosen_encoding['input_ids'] + chosen_loss_mask = self._generate_loss_mask(chosen_input_ids) + + rejected_input_ids = rejected_encoding['input_ids'] + rejected_loss_mask = self._generate_loss_mask(rejected_input_ids) + x_chosen = torch.tensor(chosen_input_ids[:-1], dtype=torch.long) + y_chosen = torch.tensor(chosen_input_ids[1:], dtype=torch.long) + mask_chosen = torch.tensor(chosen_loss_mask[1:], dtype=torch.long) + x_rejected = torch.tensor(rejected_input_ids[:-1], dtype=torch.long) + y_rejected = torch.tensor(rejected_input_ids[1:], dtype=torch.long) + mask_rejected = torch.tensor(rejected_loss_mask[1:], dtype=torch.long) + + return { + 'x_chosen': x_chosen, + 'y_chosen': y_chosen, + 'mask_chosen': mask_chosen, + 'x_rejected': x_rejected, + 'y_rejected': y_rejected, + 'mask_rejected': mask_rejected + } + + def _generate_loss_mask(self, input_ids): + loss_mask = [0] * len(input_ids) + i = 0 + while i < len(input_ids): + if input_ids[i:i + len(self.bos_id)] == self.bos_id: + start = i + len(self.bos_id) + end = start + while end < len(input_ids): + if input_ids[end:end + len(self.eos_id)] == self.eos_id: + break + end += 1 + for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)): + loss_mask[j] = 1 + i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids) + else: + i += 1 + return loss_mask + + +if __name__ == "__main__": + pass diff --git a/notebook/demo/dpo_data.jsonl b/notebook/demo/dpo_data.jsonl new file mode 100644 index 0000000..c465539 --- /dev/null +++ b/notebook/demo/dpo_data.jsonl @@ -0,0 +1,2 @@ +{"chosen":[{"content":"你好吗？","role":"user"},{"content":"我很好，谢谢！你呢？","role":"assistant"},{"content":"今天过得怎么样？","role":"user"},{"content":"挺好的，去跑步了，心情不错。","role":"assistant"}],"rejected":[{"content":"你好吗？","role":"user"},{"content":"不好，我很累。","role":"assistant"},{"content":"你喜欢什么运动？","role":"user"},{"content":"我不喜欢运动，没兴趣。","role":"assistant"}]} +{"chosen":[{"content":"你喜欢旅行吗？","role":"user"},{"content":"我喜欢，尤其是去海边。","role":"assistant"},{"content":"你去过哪些地方？","role":"user"},{"content":"我去过欧洲、美国和日本，都是很棒的经历。","role":"assistant"}],"rejected":[{"content":"你喜欢旅行吗？","role":"user"},{"content":"不太喜欢，觉得太麻烦。","role":"assistant"},{"content":"你去过哪些地方？","role":"user"},{"content":"没去过太多地方，主要是待在家里。","role":"assistant"}]} \ No newline at end of file diff --git a/notebook/demo/gqa.png b/notebook/demo/gqa.png new file mode 100644 index 0000000..e6916e0 Binary files /dev/null and b/notebook/demo/gqa.png differ diff --git a/notebook/demo/model.py b/notebook/demo/model.py new file mode 100644 index 0000000..1dc6478 --- /dev/null +++ b/notebook/demo/model.py @@ -0,0 +1,376 @@ +import math +import struct +import inspect +import time + +from .LMConfig import LMConfig +from typing import Any, Optional, Tuple, List +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn +from transformers import PreTrainedModel +from transformers.modeling_outputs import CausalLMOutputWithPast + + +class RMSNorm(torch.nn.Module): + def __init__(self, dim: int, eps: float): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + return self.weight * (x.float() * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)).type_as(x) + + +def precompute_pos_cis(dim: int, end: int = int(32 * 1024), theta: float = 1e6): + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) + t = torch.arange(end, device=freqs.device) # type: ignore + freqs = torch.outer(t, freqs).float() # type: ignore + pos_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 + return pos_cis + + +def apply_rotary_emb(xq, xk, pos_cis): + def unite_shape(pos_cis, x): + ndim = x.ndim + assert 0 <= 1 < ndim + assert pos_cis.shape == (x.shape[1], x.shape[-1]) + shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] + return pos_cis.view(*shape) + + xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) + xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) + pos_cis = unite_shape(pos_cis, xq_) + xq_out = torch.view_as_real(xq_ * pos_cis).flatten(3) + xk_out = torch.view_as_real(xk_ * pos_cis).flatten(3) + return xq_out.type_as(xq), xk_out.type_as(xk) + + +def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor: + """torch.repeat_interleave(x, dim=2, repeats=n_rep)""" + bs, slen, n_kv_heads, head_dim = x.shape + if n_rep == 1: + return x + return ( + x[:, :, :, None, :] + .expand(bs, slen, n_kv_heads, n_rep, head_dim) + .reshape(bs, slen, n_kv_heads * n_rep, head_dim) + ) + + +class Attention(nn.Module): + def __init__(self, args: LMConfig): + super().__init__() + self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads + assert args.n_heads % self.n_kv_heads == 0 + self.n_local_heads = args.n_heads + self.n_local_kv_heads = self.n_kv_heads + self.n_rep = self.n_local_heads // self.n_local_kv_heads + self.head_dim = args.dim // args.n_heads + self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False) + self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False) + self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False) + self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False) + self.attn_dropout = nn.Dropout(args.dropout) + self.resid_dropout = nn.Dropout(args.dropout) + self.dropout = args.dropout + self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn + # print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0") + mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf")) + mask = torch.triu(mask, diagonal=1) + self.register_buffer("mask", mask, persistent=False) + + def forward(self, + x: torch.Tensor, + pos_cis: torch.Tensor, + past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + use_cache=False): + bsz, seq_len, _ = x.shape + xq, xk, xv = self.wq(x), self.wk(x), self.wv(x) + xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim) + xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim) + xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim) + + xq, xk = apply_rotary_emb(xq, xk, pos_cis) + # kv_cache实现 + if past_key_value is not None: + xk = torch.cat([past_key_value[0], xk], dim=1) + xv = torch.cat([past_key_value[1], xv], dim=1) + past_kv = (xk, xv) if use_cache else None + + xq, xk, xv = ( + xq.transpose(1, 2), + repeat_kv(xk, self.n_rep).transpose(1, 2), + repeat_kv(xv, self.n_rep).transpose(1, 2) + ) + if self.flash and seq_len != 1: + dropout_p = self.dropout if self.training else 0.0 + output = F.scaled_dot_product_attention( + xq, xk, xv, + attn_mask=None, + dropout_p=dropout_p, + is_causal=True + ) + else: + scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim) + scores += self.mask[:, :, :seq_len, :seq_len] + scores = F.softmax(scores.float(), dim=-1).type_as(xq) + scores = self.attn_dropout(scores) + output = scores @ xv + + output = output.transpose(1, 2).reshape(bsz, seq_len, -1) + output = self.resid_dropout(self.wo(output)) + return output, past_kv + + +class FeedForward(nn.Module): + def __init__(self, config: LMConfig): + super().__init__() + if config.hidden_dim is None: + hidden_dim = 4 * config.dim + hidden_dim = int(2 * hidden_dim / 3) + config.hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of) + self.w1 = nn.Linear(config.dim, config.hidden_dim, bias=False) + self.w2 = nn.Linear(config.hidden_dim, config.dim, bias=False) + self.w3 = nn.Linear(config.dim, config.hidden_dim, bias=False) + self.dropout = nn.Dropout(config.dropout) + + def forward(self, x): + return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x))) + + +class MoEGate(nn.Module): + def __init__(self, config: LMConfig): + super().__init__() + self.config = config + self.top_k = config.num_experts_per_tok + self.n_routed_experts = config.n_routed_experts + + self.scoring_func = config.scoring_func + self.alpha = config.aux_loss_alpha + self.seq_aux = config.seq_aux + + self.norm_topk_prob = config.norm_topk_prob + self.gating_dim = config.dim + self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim))) + self.reset_parameters() + + def reset_parameters(self) -> None: + import torch.nn.init as init + init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + + def forward(self, hidden_states): + bsz, seq_len, h = hidden_states.shape + hidden_states = hidden_states.view(-1, h) + logits = F.linear(hidden_states, self.weight, None) + if self.scoring_func == 'softmax': + scores = logits.softmax(dim=-1) + else: + raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}') + + topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False) + + if self.top_k > 1 and self.norm_topk_prob: + denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20 + topk_weight = topk_weight / denominator + + if self.training and self.alpha > 0.0: + scores_for_aux = scores + aux_topk = self.top_k + topk_idx_for_aux_loss = topk_idx.view(bsz, -1) + if self.seq_aux: + scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1) + ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device) + ce.scatter_add_(1, topk_idx_for_aux_loss, + torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device)).div_( + seq_len * aux_topk / self.n_routed_experts) + aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha + else: + mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts) + ce = mask_ce.float().mean(0) + Pi = scores_for_aux.mean(0) + fi = ce * self.n_routed_experts + aux_loss = (Pi * fi).sum() * self.alpha + else: + aux_loss = 0 + return topk_idx, topk_weight, aux_loss + + +class MOEFeedForward(nn.Module): + def __init__(self, config: LMConfig): + super().__init__() + self.config = config + self.experts = nn.ModuleList([ + FeedForward(config) + for _ in range(config.n_routed_experts) + ]) + self.gate = MoEGate(config) + if config.n_shared_experts is not None: + self.shared_experts = FeedForward(config) + + def forward(self, x): + identity = x + orig_shape = x.shape + bsz, seq_len, _ = x.shape + # 使用门控机制选择专家 + topk_idx, topk_weight, aux_loss = self.gate(x) + x = x.view(-1, x.shape[-1]) + flat_topk_idx = topk_idx.view(-1) + if self.training: + # 训练模式下，重复输入数据 + x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0) + y = torch.empty_like(x, dtype=torch.float16) + for i, expert in enumerate(self.experts): + y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(y.dtype) # 确保类型一致 + y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1) + y = y.view(*orig_shape) + else: + # 推理模式下，只选择最优专家 + y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape) + if self.config.n_shared_experts is not None: + y = y + self.shared_experts(identity) + self.aux_loss = aux_loss + return y + + @torch.no_grad() + def moe_infer(self, x, flat_expert_indices, flat_expert_weights): + expert_cache = torch.zeros_like(x) + idxs = flat_expert_indices.argsort() + tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0) + token_idxs = idxs // self.config.num_experts_per_tok + # 例如当tokens_per_expert=[6, 15, 20, 26, 33, 38, 46, 52] + # 当token_idxs=[3, 7, 19, 21, 24, 25, 4, 5, 6, 10, 11, 12...] + # 意味着当token_idxs[:6] -> [3, 7, 19, 21, 24, 25, 4]位置的token都由专家0处理，token_idxs[6:15]位置的token都由专家1处理...... + for i, end_idx in enumerate(tokens_per_expert): + start_idx = 0 if i == 0 else tokens_per_expert[i - 1] + if start_idx == end_idx: + continue + expert = self.experts[i] + exp_token_idx = token_idxs[start_idx:end_idx] + expert_tokens = x[exp_token_idx] + expert_out = expert(expert_tokens).to(expert_cache.dtype) + expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]]) + # 使用 scatter_add_ 进行 sum 操作 + expert_cache.scatter_add_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out) + + return expert_cache + + +class MiniMindBlock(nn.Module): + def __init__(self, layer_id: int, config: LMConfig): + super().__init__() + self.n_heads = config.n_heads + self.dim = config.dim + self.head_dim = config.dim // config.n_heads + self.attention = Attention(config) + + self.layer_id = layer_id + self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps) + self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps) + self.feed_forward = FeedForward(config) if not config.use_moe else MOEFeedForward(config) + + def forward(self, x, pos_cis, past_key_value=None, use_cache=False): + h_attn, past_kv = self.attention( + self.attention_norm(x), + pos_cis, + past_key_value=past_key_value, + use_cache=use_cache + ) + h = x + h_attn + out = h + self.feed_forward(self.ffn_norm(h)) + return out, past_kv + + +class MiniMindLM(PreTrainedModel): + config_class = LMConfig + + def __init__(self, params: LMConfig = None): + self.params = params or LMConfig() + super().__init__(self.params) + self.vocab_size, self.n_layers = params.vocab_size, params.n_layers + self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim) + self.dropout = nn.Dropout(params.dropout) + self.layers = nn.ModuleList([MiniMindBlock(l, params) for l in range(self.n_layers)]) + self.norm = RMSNorm(params.dim, eps=params.norm_eps) + self.output = nn.Linear(params.dim, params.vocab_size, bias=False) + self.tok_embeddings.weight = self.output.weight + self.register_buffer("pos_cis", + precompute_pos_cis(dim=params.dim // params.n_heads, theta=params.rope_theta), + persistent=False) + self.OUT = CausalLMOutputWithPast() + + def forward(self, + input_ids: Optional[torch.Tensor] = None, + past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None, + use_cache: bool = False, + **args): + past_key_values = past_key_values or [None] * len(self.layers) + start_pos = args.get('start_pos', 0) + h = self.dropout(self.tok_embeddings(input_ids)) + pos_cis = self.pos_cis[start_pos:start_pos + input_ids.size(1)] + past_kvs = [] + for l, layer in enumerate(self.layers): + h, past_kv = layer( + h, pos_cis, + past_key_value=past_key_values[l], + use_cache=use_cache + ) + past_kvs.append(past_kv) + logits = self.output(self.norm(h)) + aux_loss = sum(l.feed_forward.aux_loss for l in self.layers if isinstance(l.feed_forward, MOEFeedForward)) + self.OUT.__setitem__('logits', logits) + self.OUT.__setitem__('aux_loss', aux_loss) + self.OUT.__setitem__('past_key_values', past_kvs) + return self.OUT + + @torch.inference_mode() + def generate(self, input_ids, eos_token_id=2, max_new_tokens=1024, temperature=0.75, top_p=0.90, + stream=False, rp=1., use_cache=True, pad_token_id=0, **args): + # 流式生成 + if stream: + return self._stream(input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args) + + # 直接生成 + generated = [] + for i in range(input_ids.size(0)): + non_pad = input_ids[i][input_ids[i] != pad_token_id].unsqueeze(0) + out = self._stream(non_pad, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args) + tokens_list = [tokens[:, -1:] for tokens in out] + gen = torch.cat(tokens_list, dim=-1) if tokens_list else non_pad + full_sequence = torch.cat([non_pad, gen], dim=-1) + generated.append(full_sequence) + max_length = max(seq.size(1) for seq in generated) + generated = [ + torch.cat( + [seq, torch.full((1, max_length - seq.size(1)), pad_token_id, dtype=seq.dtype, device=seq.device)], + dim=-1) + for seq in generated + ] + return torch.cat(generated, dim=0) + + def _stream(self, input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args): + start, first_seq, past_kvs = input_ids.shape[1], True, None + while input_ids.shape[1] < max_new_tokens - 1: + if first_seq or not use_cache: + out, first_seq = self(input_ids, past_key_values=past_kvs, use_cache=use_cache, **args), False + else: + out = self(input_ids[:, -1:], past_key_values=past_kvs, use_cache=use_cache, + start_pos=input_ids.shape[1] - 1, **args) + logits, past_kvs = out.logits[:, -1, :], out.past_key_values + logits[:, list(set(input_ids.tolist()[0]))] /= rp + logits /= (temperature + 1e-9) + if top_p is not None and top_p < 1.0: + sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1) + sorted_probs = F.softmax(sorted_logits, dim=-1) + cumulative_probs = torch.cumsum(sorted_probs, dim=-1) + sorted_indices_to_remove = cumulative_probs > top_p + sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone() + sorted_indices_to_remove[:, 0] = False + indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) + logits[indices_to_remove] = -float('Inf') + input_ids_next = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1) + input_ids = torch.cat((input_ids, input_ids_next), dim=1) + yield input_ids[:, start:] + if input_ids_next.item() == eos_token_id: + break diff --git a/notebook/demo/pretrain_data.jsonl b/notebook/demo/pretrain_data.jsonl new file mode 100644 index 0000000..c2c691c --- /dev/null +++ b/notebook/demo/pretrain_data.jsonl @@ -0,0 +1,2 @@ +{"text": "LLM首先要学习的并非直接与人交流，而是让网络参数中充满知识的墨水，“墨水” 理论上喝的越饱越好，产生大量的对世界的知识积累。 预训练就是让Model先埋头苦学大量基本的知识，例如从Wiki百科、新闻、书籍整理大规模的高质量训练数据。 这个过程是“无监督”的，即人类不需要在过程中做任何“有监督”的校正，而是由模型自己从大量文本中总结规律学习知识点。 模型此阶段目的只有一个：学会词语接龙。例如我们输入“秦始皇”四个字，它可以接龙“是中国的第一位皇帝”。"} +{"text": "经过预训练，LLM此时已经掌握了大量知识，然而此时它只会无脑地词语接龙，还不会与人聊天。 SFT阶段就需要把半成品LLM施加一个自定义的聊天模板进行微调。 例如模型遇到这样的模板【问题->回答，问题->回答】后不再无脑接龙，而是意识到这是一段完整的对话结束。 称这个过程为指令微调，就如同让已经学富五车的「牛顿」先生适应21世纪智能手机的聊天习惯，学习屏幕左侧是对方消息，右侧是本人消息这个规律。 在训练时，MiniMind的指令和回答长度被截断在512，是为了节省显存空间。就像我们学习时，会先从短的文章开始，当学会写作200字作文后，800字文章也可以手到擒来。 在需要长度拓展时，只需要准备少量的2k/4k/8k长度对话数据进行进一步微调即可（此时最好配合RoPE-NTK的基准差值）。"} \ No newline at end of file diff --git a/notebook/demo/sft_data.jsonl b/notebook/demo/sft_data.jsonl new file mode 100644 index 0000000..e037570 --- /dev/null +++ b/notebook/demo/sft_data.jsonl @@ -0,0 +1,2 @@ +{"conversations":[{"role":"user","content":"你好吗？"},{"role":"assistant","content":"我很好，谢谢！你呢？"},{"role":"user","content":"我也很好，谢谢！"},{"role":"assistant","content":"太好了！祝你今天愉快！"}]} +{"conversations":[{"role":"user","content":"你喜欢什么运动？"},{"role":"assistant","content":"我喜欢跑步和游泳。你呢？"},{"role":"user","content":"我喜欢打篮球！"},{"role":"assistant","content":"篮球很棒！是一个很好的团队运动。"}]} diff --git a/notebook/demo/tokenizer_data.jsonl b/notebook/demo/tokenizer_data.jsonl new file mode 100644 index 0000000..26e183d --- /dev/null +++ b/notebook/demo/tokenizer_data.jsonl @@ -0,0 +1,3 @@ +{"text": "近年来，人工智能技术迅速发展，深刻改变了各行各业的面貌。机器学习、自然语言处理、计算机视觉等领域的突破性进展，使得智能产品和服务越来越普及。从智能家居到自动驾驶，再到智能医疗，AI的应用场景正在快速拓展。随着技术的不断进步，未来的人工智能将更加智能、更加贴近人类生活。"} +{"text": "在数据时代，信息的流通和处理速度决定了一个国家的竞争力。随着大数据技术和云计算的普及，企业和政府能够更高效地利用信息资源，实现精细化管理和决策，1234567890qwertyuiopasdfhgjklzxbcdljfasjdfaflsdjffa深度发掘的链接发的饭卡京东方卡拉快递费发动快速开发设计大大大大大大顶顶顶。"} +{"text": "一个卖瓷碗的老人挑着扁担在路上走着，突然一个瓷碗掉到地上摔碎了，但是老人头也不回继续向前走。路人看到很奇怪，便问：“为什么你的碗摔碎了你却不看一下呢？”老人答到：“我再怎么回头看，碗是碎的。”失去的东西就要学着去接受，学着放下。很多事并不会因为你的悲伤就会回来，结果就会被改变。"} \ No newline at end of file