729 lines
29 KiB
Plaintext
729 lines
29 KiB
Plaintext
{
|
||
"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",
|
||
"\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",
|
||
"$$<f_q(x_m ,m), f_k(x_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",
|
||
"\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
|
||
}
|