{ "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 }