1422 lines
63 KiB
Plaintext
1422 lines
63 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）模型和混合专家（Mixture of Expert, MoE）模型 。稠密模型中所有参数在每次计算时都会参与运算；混合专家模型则将不同的 “专家” 模块组合, 根据输入选择合适的专家处理, 能在保证效果的同时减少计算量和参数量.\n",
|
||
"\n",
|
||
"MiniMind 系列模型在 Llama 3.1 的基础上设计, 基于经典的 Transformer Deocder-Only 架构，在这一板块, 我们将围绕 MiniMind 系列模型的源代码展开学习."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "f12e2248-7b71-442f-96a0-01a5cb4a1171",
|
||
"metadata": {},
|
||
"source": [
|
||
"## MiniMind Dense Model\n",
|
||
"\n",
|
||
"作者提供了对其 MiniMind Dense Model 模型结构的可视化：\n",
|
||
"\n",
|
||
""
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 1,
|
||
"id": "6cd2aa37-3460-4a64-8285-f888d62b99d5",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import math\n",
|
||
"import struct\n",
|
||
"import inspect\n",
|
||
"import time\n",
|
||
"\n",
|
||
"from model.LMConfig import LMConfig\n",
|
||
"from typing import Any, Optional, Tuple, List\n",
|
||
"import numpy as np\n",
|
||
"import torch\n",
|
||
"import torch.nn.functional as F\n",
|
||
"from torch import nn\n",
|
||
"from transformers import PreTrainedModel\n",
|
||
"from transformers.modeling_outputs import CausalLMOutputWithPast"
|
||
]
|
||
},
|
||
{
|
||
"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) + \\epsilon} * \\gamma, \\quad where \\quad RMS(a) = \\sqrt{\\frac{1}{n}\\sum^n_{i=1}a^2_i}.$$\n",
|
||
"\n",
|
||
"假设输入张量形状为 (batch_size, sequence_length, embedding_dim), RMS Norm 对 embedding_dim 维度进行归一化,其中, 其中, $\\epsilon$ 是一个超参数, 用于防止分母为零导致结果上溢, $\\gamma$ 为可学习参数.\n",
|
||
"\n",
|
||
"不难发现, 当均值为零时, Layer Norm 退化为 RMS Norm. 这是因为 RMS Norm 在 Layer Norm 的基础上舍弃了中心化操作, 仅用缩放进行归一化, 其不改变数据原本的分布, 有利于激活函数输出的稳定."
|
||
]
|
||
},
|
||
{
|
||
"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-attention 中, 进而提升 transformer 架构性能的位置编码方式, 和绝对位置编码相比, RoPE 具有很好的外推性, 是目前的主流位置编码方式.\n",
|
||
"\n",
|
||
"外推性的解释, 通俗来说就是训练的时候限制了 512 的上下文长度，那么推d理时如果面对超过该长度的文本，LLM 可能无法正确处理.\n",
|
||
"\n",
|
||
"- **绝对位置编码**\n",
|
||
"\n",
|
||
"绝对位置编码是早期 Transformer 架构采用的绝对位置编码方案，及那个每个位置映射为固定的向量表示.\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$ 的计算使用如下公式：\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",
|
||
"出于认识的目的，我们省略复杂的数学推导，直接看 RoPE 的的结论：\n",
|
||
"\n",
|
||
"存在这样一个正交矩阵：\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}$$\n",
|
||
"\n",
|
||
"其中，$\\Theta=\\left\\{\\theta_i=10000^{-2(i-1)/d},i\\in[1,2,\\ldots,d/2]\\right\\}$\n",
|
||
"\n",
|
||
"我们可以将 query 和 key 的内积操作转换为与原始向量 $x$ 相关的以下等价形式：\n",
|
||
"\n",
|
||
"$$\n",
|
||
"\\boldsymbol{q}_m^\\mathbf{T}\\boldsymbol{k}_n=\\left(\\boldsymbol{R}_{\\Theta,m}^d\\boldsymbol{W}_q\\boldsymbol{x}_m\\right)^\\mathbf{T}\\left(\\boldsymbol{R}_{\\Theta,n}^d\\boldsymbol{W}_k\\boldsymbol{x}_n\\right)=\\boldsymbol{x}_m^\\mathbf{T}\\boldsymbol{W}_q\\boldsymbol{R}_{\\Theta,n-m}^d\\boldsymbol{W}_k\\boldsymbol{x}_n\n",
|
||
"$$\n",
|
||
"\n",
|
||
"其中， $\\boldsymbol{R}_{\\Theta,n-m}^d=\\left(\\boldsymbol{R}_{\\Theta,m}^d\\right)^\\mathbf{T}\\boldsymbol{R}_{\\Theta,n}^d$.\n",
|
||
"\n",
|
||
"由于 $\\boldsymbol{R}_{\\Theta,m}^d$ 的稀疏性，直接使用矩阵乘法会浪费算力，因此代码中采用下述方式实现：\n",
|
||
"\n",
|
||
"$$\\boldsymbol{R}_{\\Theta,m}^{d}\\boldsymbol{x}=\\begin{pmatrix}x_{0}\\\\x_{1}\\\\x_{2}\\\\x_{3}\\\\\\vdots\\\\x_{d-2}\\\\x_{d-1}\\end{pmatrix}\\otimes\\begin{pmatrix}\\cos m\\theta_{0}\\\\\\cos m\\theta_{0}\\\\\\cos m\\theta_{1}\\\\\\cos m\\theta_{1}\\\\\\vdots\\\\\\cos m\\theta_{d/2-1}\\\\\\cos m\\theta_{d/2-1}\\end{pmatrix}+\\begin{pmatrix}-x_{1}\\\\x_{0}\\\\-x_{3}\\\\x_{2}\\\\\\vdots\\\\-x_{d-1}\\\\x_{d-2}\\end{pmatrix}\\otimes\\begin{pmatrix}\\sin m\\theta_{0}\\\\\\sin m\\theta_{0}\\\\\\sin m\\theta_{1}\\\\\\sin m\\theta_{1}\\\\\\vdots\\\\\\sin m\\theta_{d/2-1}\\\\\\sin m\\theta_{d/2-1}\\end{pmatrix}\n",
|
||
"$$"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 3,
|
||
"id": "5aae4194-3738-44a7-9c33-bbf964067f8a",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"def precompute_pos_cis(dim: int, end: int=int(32 * 1024), theta: float=1e6):\n",
|
||
" r\"\"\"生成旋转矩阵\"\"\"\n",
|
||
" # 计算词向量元素两两分组后 每组元素对应的旋转角度 \\theta_i 也就是上文公式中对应的旋转角度集合\n",
|
||
" freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))\n",
|
||
" # token 序列索引 t = [0, 1, ..., seq_len-1]\n",
|
||
" t = torch.arange(end, device=freqs.device)\n",
|
||
" # 计算 m * \\theta\n",
|
||
" freqs = torch.outer(t, freqs).float()\n",
|
||
" # 计算结果是个复数向量\n",
|
||
" # 假设 freqs = [x, y]\n",
|
||
" # 则 freqs_cis = [cos(x) + sin(x)i, cos(y) + sin(y)i]\n",
|
||
" pos_cis = torch.polar(torch.ones_like(freqs), freqs)\n",
|
||
" return pos_cis\n",
|
||
"\n",
|
||
"def apply_rotary_emb(xq, xk, pos_cis):\n",
|
||
" r\"\"\"应用 RoPE\"\"\"\n",
|
||
" def unite_shape(pos_cis, x):\n",
|
||
" ndim = x.ndim\n",
|
||
" assert 0 <= 1 < ndim\n",
|
||
" assert pos_cis.shape == (x.shape[1], x.shape[-1])\n",
|
||
" shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]\n",
|
||
" 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",
|
||
" # 保留实部\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": "markdown",
|
||
"id": "eff57879-1f39-404d-9501-3c99e2e67c41",
|
||
"metadata": {},
|
||
"source": [
|
||
"我们知道，RoPE 是在 Attention 阶段生成 Query 和 Key 向量后，对这两个向量进行位置编码的.\n",
|
||
"\n",
|
||
"对于 MiniMindLM，嵌入维度为 512，注意力头数量为 8，故每一个注意力头的维度应该为 512 / 8 = 64.\n",
|
||
"\n",
|
||
"我们使用固定形状的张量代表 Query 和 Key 向量，对 RoPE 的应用过程进行观察."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 4,
|
||
"id": "5f8376c4-5b00-4c72-9f23-912478db581c",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"pos_cis 的形状为 torch.Size([16, 32]), 其中 [0, 0] 下标元素为 (1+0j)\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"xq, xk = torch.randn((2, 16, 4, 64)), torch.randn((2, 16, 4, 64)) # (batch_size, sequence_length, num_heads, head_dim)\n",
|
||
"pos_cis = precompute_pos_cis(64, 16) # 计算旋转位置编码的旋转角（复数表示）\n",
|
||
"print(f'pos_cis 的形状为 {pos_cis.shape}, 其中 [0, 0] 下标元素为 {pos_cis[0, 0]}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 5,
|
||
"id": "b40d0ec8-8752-4db9-a4fe-e746d39dc970",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"经过 RoPE 编码后的 Query 与 Key 的形状为 torch.Size([2, 16, 4, 64]), torch.Size([2, 16, 4, 64])\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"xq_rope, xk_rope = apply_rotary_emb(xq, xk, pos_cis)\n",
|
||
"# original shape xq: torch.Size([2, 16, 4, 64])\n",
|
||
"# ajusted shape xq: torch.Size([2, 16, 4, 32])\n",
|
||
"# ajusted shape pos_cis： torch.Size([1, 16, 1, 32])\n",
|
||
"print(f'经过 RoPE 编码后的 Query 与 Key 的形状为 {xq_rope.shape}, {xk_rope.shape}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "fa47eeb0-874e-4df3-bf91-5f79109af8ca",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Attention\n",
|
||
"\n",
|
||
"注意力机制（Attention Mechanism）是Transformer架构的核心组件，能够有效捕捉长序列内各元素间的依赖关系. 该机制通过计算输入序列中不同位置元素间的注意力得分，对其重要性进行精准建模，使模型在处理信息时能够聚焦于关键部分，从而显著提升对长序列数据的理解与处理能力.\n",
|
||
"\n",
|
||
"在 MiniMindLM 模型中, Attention Block 包含以下机制和模块:\n",
|
||
"\n",
|
||
"1. GQA\n",
|
||
"2. KV Cache\n",
|
||
"5. SwiGLU\n",
|
||
"\n",
|
||
"- **GQA**\n",
|
||
"\n",
|
||
"分组查询注意力 (Group Querey Attention, GQA) 是对多头自注意力机制的扩展, 通过提供计算效率和模型表达能力的灵活权衡, 实现了查询头的分组.\n",
|
||
"\n",
|
||
"具体来说, GQA 中将 h 个查询头分为 G 组, 每组 包含 h / G 个查询头，并共享一个公共的键和值.\n",
|
||
"\n",
|
||
"\n",
|
||
"\n",
|
||
"GQA 相比传统的 MHA， 减少了键和值的数量，降低了计算量和内存开销，提高了推理速度.\n",
|
||
"\n",
|
||
"- **KV Cache**\n",
|
||
"\n",
|
||
"在语言模型生成文本的过程中，每生成一个新的 token，模型都需要计算注意力得分，以确定当前位置与之前所有位置的相关性.\n",
|
||
"\n",
|
||
"比如以下内容：\n",
|
||
"\n",
|
||
"1. *seq = \\[tok1]:*\n",
|
||
"\n",
|
||
" attn_11 = softmax(Q1 * K1.T / sqrt(dim)) * V1\n",
|
||
" \n",
|
||
"3. *seq = \\[tok1, tok2]:*\n",
|
||
"\n",
|
||
" attn_11 = softmax(Q1 * K1.T / sqrt(dim)) * V1, attn_12 = 0 (masked)\n",
|
||
" \n",
|
||
" attn_21 = softmax(Q2 * K1.T / sqrt(dim)) * V2, attn_22 = softmax(Q2 * K2.T / sqrt(dim)) * V2\n",
|
||
"4. *seq = \\[tok1, tok2, tok3]:*\n",
|
||
"\n",
|
||
" attn_11 = softmax(Q1 * K1.T / sqrt(dim)) * V1, attn_12 = 0 (masked), attn_13 = 0 (masked)\n",
|
||
" \n",
|
||
" attn_21 = softmax(Q2 * K1.T / sqrt(dim)) * V2, attn_22 = softmax(Q2 * K2.T / sqrt(dim)) * V2, attn_23 = 0 (masked)\n",
|
||
" \n",
|
||
" attn_31 = softmax(Q3 * K1.T / sqrt(dim)) * V3, attn_32 = softmax(Q3 * K2.T / sqrt(dim)) * V3, attn_33 = softmax(Q3 * K3.T / sqrt(dim)) * V3\n",
|
||
"5. ··· ···\n",
|
||
"\n",
|
||
"不难发现，大模型生成一个 token 后的注意力计算中，总会用到 token 序列的历史 KV 值，导致重复计算，KV Cache 的设计正是为了通过缓存历史 KV 值，节省计算开销.\n",
|
||
"\n",
|
||
"KV Cache 能够有效压缩大模型推理时的显存占用.\n",
|
||
"\n",
|
||
"- **SwiGLU**\n",
|
||
"\n",
|
||
"SwiGLU 是一种在深度学习中用于神经网络架构的激活函数变体：\n",
|
||
"\n",
|
||
"$$\\text{SwiGLU}(x, W, V, b, c)=\\text{Swish}_1(xW+b)\\otimes(xV+c)$$\n",
|
||
"\n",
|
||
"与传统的 ReLU 激活函数相比，SwiGLU 具有更好的平滑性和非线性表达能力，由于其门控机制，在处理信息筛选和流动方面有独特的优势."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 6,
|
||
"id": "e61394db-5afa-4421-ba77-30a9c50f0603",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"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",
|
||
" r\"\"\"使得 KV 头数适应 Query 头数\"\"\"\n",
|
||
" bs, slen, n_kv_heads, head_dim = x.shape\n",
|
||
" if n_rep == 1:\n",
|
||
" return x\n",
|
||
" return (\n",
|
||
" x[:, :, :, None, :]\n",
|
||
" .expand(bs, slen, n_kv_heads, n_rep, head_dim)\n",
|
||
" .reshape(bs, slen, n_kv_heads * n_rep, head_dim)\n",
|
||
" )\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",
|
||
" # q k v o projection\n",
|
||
" self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)\n",
|
||
" self.wk = nn.Linear(args.dim, args.n_kv_heads * self.head_dim, bias=False)\n",
|
||
" self.wv = nn.Linear(args.dim, args.n_kv_heads * self.head_dim, bias=False)\n",
|
||
" self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)\n",
|
||
" self.attn_dropout = nn.Dropout(args.dropout)\n",
|
||
" self.resid_dropout = nn.Dropout(args.dropout)\n",
|
||
" self.dropout = args.dropout\n",
|
||
" self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn\n",
|
||
" mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float(\"-inf\"))\n",
|
||
" mask = torch.triu(mask, diagonal=1)\n",
|
||
" self.register_buffer(\"mask\", mask, persistent=False)\n",
|
||
"\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": "markdown",
|
||
"id": "7420b47b-b38e-442a-84a2-e307f82584d3",
|
||
"metadata": {},
|
||
"source": [
|
||
"同样的，我们假设一批 batch size = 4， seq len = 16 的 token 序列通过这个 Attention 块，在输入前，它的 input id 会被投影到 dim = 512 维."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 7,
|
||
"id": "e1519451-3dc8-498e-96a7-db8b32a71af1",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# 创建 Config，为了简单起见，我们设定 num_layers = 2\n",
|
||
"LMConfig_Dense = LMConfig(n_layers=2)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 8,
|
||
"id": "ca9f6675-7d70-49b4-ba02-f91109869fd8",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"输入张量 x ：size = torch.Size([4, 16, 512])，RoPE 旋转角： size = torch.Size([16, 32])\n",
|
||
"输出 output: size = torch.Size([4, 16, 512]), kv_cache 基本信息：size_key = torch.Size([4, 16, 2, 64]), size_value = torch.Size([4, 16, 2, 64])\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"attn = Attention(LMConfig_Dense)\n",
|
||
"x = torch.randn((4, 16, 512)) # (batch size, seq len, embed dim)\n",
|
||
"pos_cis = precompute_pos_cis(64, 16) # (head dim, batch size) 其中 head dim = embed dim / num heads\n",
|
||
"output, past_kv = attn(x, pos_cis=pos_cis, use_cache=True)\n",
|
||
"print(f'输入张量 x ：size = {x.shape}，RoPE 旋转角： size = {pos_cis.shape}')\n",
|
||
"print(f'输出 output: size = {output.shape}, kv_cache 基本信息：size_key = {past_kv[0].shape}, size_value = {past_kv[1].shape}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 9,
|
||
"id": "b53cee91-ee0d-4a61-8a7b-6feac5ba1b40",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"del attn"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "495f1e81-36de-4e94-8ae4-69490f62d15f",
|
||
"metadata": {},
|
||
"source": [
|
||
"### FeedForward Network\n",
|
||
"\n",
|
||
"前馈神经网络接收来自注意力机制层的输出结果，随后对该输出执行进一步的线性变换. 通过这种方式，网络能够深入挖掘并捕获更为复杂、抽象的特征."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 10,
|
||
"id": "a16f3ece-3b4e-42aa-9b88-1a7193dd2bf0",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"class FeedForward(nn.Module):\n",
|
||
" def __init__(self, config: LMConfig):\n",
|
||
" super().__init__()\n",
|
||
" if config.hidden_dim is None:\n",
|
||
" hidden_dim = 4 * config.dim\n",
|
||
" hidden_dim = int(2 * hidden_dim / 3)\n",
|
||
" config.hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of)\n",
|
||
" self.w1 = nn.Linear(config.dim, config.hidden_dim, bias=False)\n",
|
||
" self.w2 = nn.Linear(config.hidden_dim, config.dim, bias=False)\n",
|
||
" self.w3 = nn.Linear(config.dim, config.hidden_dim, bias=False)\n",
|
||
" self.dropout = nn.Dropout(config.dropout)\n",
|
||
"\n",
|
||
" def forward(self, x):\n",
|
||
" return self.dropout(self.w2(F.silu(self.w1(x) * self.w3(x))))"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "8ddb2503-018c-4486-ae45-d9d0a192490a",
|
||
"metadata": {},
|
||
"source": [
|
||
"设置输入 x 为 batch size = 4，seq len = 16 的 token 序列投影向量，观察 x 在 MiniMind Block 的前向传播过程."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 11,
|
||
"id": "4fe2abd9-0c40-44be-b64d-8b0caa28604b",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"给定输入 x: size = torch.Size([4, 16, 512]) 下的输出 output：size = torch.Size([4, 16, 512])\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"ffn = FeedForward(LMConfig_Dense)\n",
|
||
"x = torch.randn((4, 16, 512))\n",
|
||
"output = ffn(x)\n",
|
||
"print(f'给定输入 x: size = {x.shape} 下的输出 output：size = {output.shape}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 12,
|
||
"id": "5ae0c7e5-67e9-4dc4-9ea5-4ec20667c634",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"del ffn"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "efe25d3b-2af9-4fe1-8eb0-7522324ec2c6",
|
||
"metadata": {},
|
||
"source": [
|
||
"这就是 Transformer 结构的精妙之处，张量在模型内部进行了各种复杂的投影变形，但是输入输出的张量形状不发生改变!"
|
||
]
|
||
},
|
||
{
|
||
"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": 13,
|
||
"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 # head dim\n",
|
||
" self.attention = Attention(config)\n",
|
||
"\n",
|
||
" self.layer_id = layer_id # layer id 指向其维护的 KV Cache\n",
|
||
" self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)\n",
|
||
" self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)\n",
|
||
" self.feed_forward = FeedForward(config)\n",
|
||
"\n",
|
||
" def forward(self, x, pos_cis, past_key_value=None, use_cache=False):\n",
|
||
" h_attn, past_kv = self.attention(\n",
|
||
" self.attention_norm(x), # pre-normed x\n",
|
||
" pos_cis, \n",
|
||
" past_key_value=past_key_value, \n",
|
||
" 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": "markdown",
|
||
"id": "f627c780-f07e-41c3-8b44-7a819e80a2a4",
|
||
"metadata": {},
|
||
"source": [
|
||
"我们依然设置输入 x 为 batch size = 4，seq len = 16 的 token 序列投影向量，观察 x 在 MiniMind Block 的前向传播过程."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 14,
|
||
"id": "d23b4461-534c-4f6d-a00b-cdcfbd4d7bec",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"输出 output 信息: size = torch.Size([4, 16, 512])\n",
|
||
"该 Block 维护的 KV Cache 信息：size_key = torch.Size([4, 16, 2, 64]), size_value = torch.Size([4, 16, 2, 64])\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"miniblock = MiniMindBlock(1, LMConfig_Dense)\n",
|
||
"x = torch.randn((4, 16, 512))\n",
|
||
"pos_cis = precompute_pos_cis(64, 16)\n",
|
||
"out, past_kv = miniblock(x, pos_cis, use_cache=True)\n",
|
||
"print(f'输出 output 信息: size = {out.shape}\\n该 Block 维护的 KV Cache 信息：size_key = {past_kv[0].shape}, size_value = {past_kv[1].shape}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 15,
|
||
"id": "bb9745cf-2010-49b3-90ad-a1d59f76de89",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"del miniblock"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "f9bb3b2e-4d1d-4721-9908-546f3c49dde3",
|
||
"metadata": {},
|
||
"source": [
|
||
"### MiniMindLM (Dense)\n",
|
||
"\n",
|
||
"以 MiniMind Block 为基本组件, 我们对 MiniMindLM 进行最后组装！"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 16,
|
||
"id": "7267245b-f13a-4bcf-9358-5fdd89106153",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from transformers import PreTrainedModel\n",
|
||
"from transformers.modeling_outputs import CausalLMOutputWithPast\n",
|
||
"\n",
|
||
"class MiniMindLM(PreTrainedModel):\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 = [] # KV Cache 列表 按照索引维护对应 MiniMind Block 模块的 KV Cache\n",
|
||
" print(f'====> forward propagation started, num_minimind_blocks = {len(self.layers)}')\n",
|
||
" for l, layer in enumerate(self.layers):\n",
|
||
" print(f'------------> entering minimind block: id = {l}')\n",
|
||
" h, past_kv = layer(\n",
|
||
" h, pos_cis, \n",
|
||
" past_key_value=past_key_values[l], \n",
|
||
" use_cache=use_cache\n",
|
||
" )\n",
|
||
" print(f'<------------ finished, size_cache_k = {past_kv[0].shape}, size_cache_v = {past_kv[1].shape}')\n",
|
||
" past_kvs.append(past_kv)\n",
|
||
" print(f'<==== forward propagation completed, num_kv_cache = {len(past_kvs)}\\n')\n",
|
||
" logits = self.output(self.norm(h))\n",
|
||
" aux_loss = 0 # MoE 辅助损失，我们在 Dense Model 中不纳入考虑\n",
|
||
" self.OUT.__setitem__('logits', logits) # token 分类 logits\n",
|
||
" self.OUT.__setitem__('aux_loss', aux_loss) # MoE 辅助损失\n",
|
||
" self.OUT.__setitem__('past_key_value', past_kvs) # 当前序列的历史 KV_cache 列表\n",
|
||
" return self.OUT\n",
|
||
"\n",
|
||
" @torch.inference_mode()\n",
|
||
" def generate(self, input_ids, eos_token_id=2, max_new_tokens=512, temperature=0.75, top_p=0.90, \n",
|
||
" stream=False, rp=1, use_cache=True, pad_token_id=0, **args):\n",
|
||
" # 流式生成 （返回生成器, 自由控制输出）\n",
|
||
" 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",
|
||
" print(f'new tokens list :{tokens_list}\\n')\n",
|
||
" gen = torch.cat(tokens_list, dim=-1) if tokens_list else non_pad\n",
|
||
" full_sequence = torch.cat([non_pad, gen], dim=-1)\n",
|
||
" generated.append(full_sequence)\n",
|
||
" max_length = max(seq.size(1) for seq in generated)\n",
|
||
" generated = [\n",
|
||
" torch.cat([seq, torch.full((1, max_length - seq.size(1)), pad_token_id, dtype=seq.dtype, device=seq.device)], dim=-1) \n",
|
||
" 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",
|
||
" new_token_idx = 0 # new token 计数器\n",
|
||
" while input_ids.shape[1] < max_new_tokens - 1:\n",
|
||
" print(f'gernerating new token: idx = {start + new_token_idx}')\n",
|
||
" if first_seq or not use_cache: # 若第一次生成序列 or 无 KV Cache, 每次生成传入整个 token id 序列\n",
|
||
" out, first_seq = self(input_ids, past_key_values=past_kvs, use_cache=use_cache, **args), False\n",
|
||
" else: # 若非第一次生成 and 有 KV Cache, 每次传入最后一个 token id 与 KV Cache 进行推理加速\n",
|
||
" out = self(input_ids[:, -1:], past_key_values=past_kvs, use_cache=use_cache, \n",
|
||
" start_pos=input_ids.shape[1] - 1, **args)\n",
|
||
" logits, past_kvs = out.logits[:, -1, :], out.past_key_values # logits.shape: (batch_size, seq_len, embed_dim), 获取最后一位 logits\n",
|
||
" logits[:, list(set(input_ids.tolist()[0]))] /= rp # 对生成 token 进行惩罚, 降低后续重复生成几率\n",
|
||
" logits /= (temperature + 1e-9) # 调整温度, 控制生成多样性\n",
|
||
" if top_p is not None and top_p < 1.0: # top-p 采样\n",
|
||
" sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)\n",
|
||
" sorted_probs = F.softmax(sorted_logits, dim=-1)\n",
|
||
" cumulative_probs = torch.cumsum(sorted_probs, dim=-1)\n",
|
||
" sorted_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",
|
||
" new_token_idx += 1\n",
|
||
" yield input_ids[:, start:]\n",
|
||
" if input_ids_next.item() == eos_token_id:\n",
|
||
" break"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "bdee2ebc-f98c-49cc-b5fc-c6cd0056f4d9",
|
||
"metadata": {},
|
||
"source": [
|
||
"接下来，我们设置一条长度为 4 的 token 序列，使用 MiniMindLM 对其执行一次前向传播，观察传播过程与返回值."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 17,
|
||
"id": "dc7b4d79-a3aa-4143-aa5c-2abce34bc96c",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"====> forward propagation started, num_minimind_blocks = 2\n",
|
||
"------------> entering minimind block: id = 0\n",
|
||
"<------------ finished, size_cache_k = torch.Size([1, 4, 2, 64]), size_cache_v = torch.Size([1, 4, 2, 64])\n",
|
||
"------------> entering minimind block: id = 1\n",
|
||
"<------------ finished, size_cache_k = torch.Size([1, 4, 2, 64]), size_cache_v = torch.Size([1, 4, 2, 64])\n",
|
||
"<==== forward propagation completed, num_kv_cache = 2\n",
|
||
"\n",
|
||
"返回 logits：size = torch.Size([1, 4, 6400]), 返回 aux_loss: 0, 返回 KV Cache List： len = 2\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"MiniMind_Dense = MiniMindLM(LMConfig_Dense)\n",
|
||
"input_ids = torch.Tensor([1, 3, 5, 7]).long().reshape(1, 4)\n",
|
||
"OUT = MiniMind_Dense(input_ids, use_cache=True)\n",
|
||
"print(f'返回 logits：size = {OUT.logits.shape}, 返回 aux_loss: {OUT.aux_loss}, 返回 KV Cache List： len = {len(OUT.past_key_value)}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 18,
|
||
"id": "3480a860-e934-43de-a621-8063f0bdea2f",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"gernerating new token: idx = 4\n",
|
||
"====> forward propagation started, num_minimind_blocks = 2\n",
|
||
"------------> entering minimind block: id = 0\n",
|
||
"<------------ finished, size_cache_k = torch.Size([1, 4, 2, 64]), size_cache_v = torch.Size([1, 4, 2, 64])\n",
|
||
"------------> entering minimind block: id = 1\n",
|
||
"<------------ finished, size_cache_k = torch.Size([1, 4, 2, 64]), size_cache_v = torch.Size([1, 4, 2, 64])\n",
|
||
"<==== forward propagation completed, num_kv_cache = 2\n",
|
||
"\n",
|
||
"gernerating new token: idx = 5\n",
|
||
"====> forward propagation started, num_minimind_blocks = 2\n",
|
||
"------------> entering minimind block: id = 0\n",
|
||
"<------------ finished, size_cache_k = torch.Size([1, 1, 2, 64]), size_cache_v = torch.Size([1, 1, 2, 64])\n",
|
||
"------------> entering minimind block: id = 1\n",
|
||
"<------------ finished, size_cache_k = torch.Size([1, 1, 2, 64]), size_cache_v = torch.Size([1, 1, 2, 64])\n",
|
||
"<==== forward propagation completed, num_kv_cache = 2\n",
|
||
"\n",
|
||
"gernerating new token: idx = 6\n",
|
||
"====> forward propagation started, num_minimind_blocks = 2\n",
|
||
"------------> entering minimind block: id = 0\n",
|
||
"<------------ finished, size_cache_k = torch.Size([1, 1, 2, 64]), size_cache_v = torch.Size([1, 1, 2, 64])\n",
|
||
"------------> entering minimind block: id = 1\n",
|
||
"<------------ finished, size_cache_k = torch.Size([1, 1, 2, 64]), size_cache_v = torch.Size([1, 1, 2, 64])\n",
|
||
"<==== forward propagation completed, num_kv_cache = 2\n",
|
||
"\n",
|
||
"new tokens list :[tensor([[5081]]), tensor([[4500]]), tensor([[850]])]\n",
|
||
"\n",
|
||
"生成结果：tensor([[ 1, 3, 5, 7, 5081, 4500, 850]])\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"# 我们让 MiniMind 根据我们设计的 四个输入 token 生成输出\n",
|
||
"out = MiniMind_Dense.generate(input_ids, max_new_tokens=8, use_cache=True)\n",
|
||
"print(f'生成结果：{out}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 19,
|
||
"id": "4509ff58-0439-42b5-b303-d308b502e821",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"del MiniMind_Dense"
|
||
]
|
||
},
|
||
{
|
||
"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": [
|
||
"作者提供了 MiniMind MoE Model 的可视化.\n",
|
||
"\n",
|
||
""
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "82682a98-cae7-49c2-be2f-12dd1adb8deb",
|
||
"metadata": {},
|
||
"source": [
|
||
"可以看到，Dense Model 和 MoE Model 的差异在于 FFN 层. MoE 模型将稠密连接的 FFN 层置换为 M x Expert 层，每次前向传播时只激活部分 Expert.\n",
|
||
"\n",
|
||
"其组成可以分为以下部分：\n",
|
||
"\n",
|
||
"- Experts: MoE 架构单元，每个专家本质上是一个独立的神经网络模块，负责处理特定类型或范围的数据.\n",
|
||
"- Router： 控制信息流动，决定每次前向传播激活的 Experts 模块以及输入数据在这些模块的分配组合.\n",
|
||
"\n",
|
||
"在 MoE 网络中，为了平衡专家的重要性，我们需要关注路由，它是决定在特定时间选择哪些专家的重要组件.\n",
|
||
"\n",
|
||
"### 辅助损失\n",
|
||
"\n",
|
||
"为了让训练过程中实现专家的更均匀分布，辅助损失（又称负载均衡损失）被添加到网络的常规损失中，它增加了一个约束，迫使专家具有相等的重要性.\n",
|
||
"\n",
|
||
"假设有输入序列 \\[What is Mixture of Experts], *Prob* (·) 表示每一个 token 激活的专家概率分布:\n",
|
||
"\n",
|
||
"- **Step 1**：在整个批次中对每个专家的路由值进行求和.\n",
|
||
"\n",
|
||
" Importance_per_token = *Prob* (What) + *Prob* (is) + *Prob* (Mixture) + *Prob* (of) + *Prob* (Experts)\n",
|
||
"\n",
|
||
" 这个指标反映了 batch 维度上每个专家的重要性分数.\n",
|
||
"\n",
|
||
"- **Step 2**: 计算变异系数\n",
|
||
"\n",
|
||
" 我们希望专家之间的重要性尽可能靠近，为了衡量专家得分之间的差异程度，引入变异系数指标（Coefficient Variation, CV）\n",
|
||
"\n",
|
||
" $$Coeifficient \\, Variation (CV) = \\dfrac{standard \\, deviation (\\sigma)}{mean(\\mu)}$$\n",
|
||
"\n",
|
||
" 如果专家的重要性分数相似，变异系数会降到很低（这是我们期望的）\n",
|
||
"\n",
|
||
"- **Step 3**: 计算负载均衡损失\n",
|
||
"\n",
|
||
" $$uxiliary Loss = \\alpha * CV$$\n",
|
||
"\n",
|
||
" 其中 $\\alpha$ 是缩放系数."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "26dffcd0-7c9e-40bf-9c91-1636aa571d46",
|
||
"metadata": {},
|
||
"source": [
|
||
"## MoE Gate\n",
|
||
"\n",
|
||
"MoE 门控单元决定每次前向传播激活的 Experts 模块及其权重，同时计算 MoE 辅助损失."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 20,
|
||
"id": "cba5532c-e8d2-40ff-9ed6-b8b33b4a31ca",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"class MoEGate(nn.Module):\n",
|
||
" def __init__(self, config: LMConfig):\n",
|
||
" super().__init__()\n",
|
||
" self.config = config\n",
|
||
" self.top_k = config.num_experts_per_tok\n",
|
||
" self.n_routed_experts = config.n_routed_experts # 总专家数量\n",
|
||
"\n",
|
||
" self.scoring_func = config.scoring_func # 评分函数\n",
|
||
" self.alpha = config.aux_loss_alpha # 辅助损失的 alpha 参数\n",
|
||
" self.seq_aux = config.seq_aux # 是否启用序列辅助损失\n",
|
||
"\n",
|
||
" self.norm_topk_prob = config.norm_topk_prob\n",
|
||
" self.gating_dim = config.dim\n",
|
||
" self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim))) # 混合专家层 -> embeding dim 层\n",
|
||
" self.reset_parameter()\n",
|
||
"\n",
|
||
" def reset_parameter(self):\n",
|
||
" import torch.nn.init as init\n",
|
||
" init.kaiming_uniform_(self.weight, a=math.sqrt(5))\n",
|
||
"\n",
|
||
" def forward(self, hidden_states):\n",
|
||
" bsz, seq_len, h = hidden_states.shape\n",
|
||
" hidden_states = hidden_states.view(-1, h) # (total_tokens, hidden_dim) 其中 total_tokens = bsz * seq_len\n",
|
||
" logits = F.linear(hidden_states, self.weight, None) # logits = hidden_states · self.weight.T\n",
|
||
" # print(f'shape of logits computed: {logits.shape}') # (total_tokens, n_routed_experts)\n",
|
||
"\n",
|
||
" if self.scoring_func == 'softmax':\n",
|
||
" scores = logits.softmax(dim=-1)\n",
|
||
" else:\n",
|
||
" raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}')\n",
|
||
"\n",
|
||
" topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)\n",
|
||
" if self.top_k > 1 and self.norm_topk_prob: # prob 归一化\n",
|
||
" denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e20\n",
|
||
" topk_weight = topk_weight / denominator\n",
|
||
"\n",
|
||
" if self.training and self.alpha > 0.0:\n",
|
||
" scores_for_aux = scores\n",
|
||
" aux_topk = self.top_k\n",
|
||
" topk_idx_for_aux_loss = topk_idx.view(bsz, -1) # (bsz, seq_len * num_experts_per_tok)\n",
|
||
" if self.seq_aux:\n",
|
||
" ################# 对所有批次计算 Experts 重要性得分 #################\n",
|
||
" scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)\n",
|
||
" ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)\n",
|
||
" # 根据索引将 torch.ones 累加到 ce 中，并执行归一化\n",
|
||
" ce.scatter_add_(1, topk_idx_for_aux_loss,\n",
|
||
" torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device)).div_(\n",
|
||
" seq_len * aux_topk / self.n_routed_experts)\n",
|
||
" ######################### 计算负载均衡损失 ##########################\n",
|
||
" aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha\n",
|
||
" else:\n",
|
||
" mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)\n",
|
||
" ce = mask_ce.float().mean(0)\n",
|
||
" Pi = scores_for_aux.mean(0)\n",
|
||
" fi = ce * self.n_routed_experts\n",
|
||
" aux_loss = (Pi * fi).sum() * self.alpha\n",
|
||
" else:\n",
|
||
" aux_loss = 0\n",
|
||
" return topk_idx, topk_weight, aux_loss"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "455f5c52-aad9-4de8-a571-d47de6f0d7d3",
|
||
"metadata": {},
|
||
"source": [
|
||
"接下来，我们设置一个 batch size = 4, seq len =16, emb dim = 512 的输入向量，在 MoE 门控单元前向传播进行观察"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 21,
|
||
"id": "5b3acf9e-b814-49f7-954f-246c33daa23d",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# 重新创建一个 LMConfig 使得 use MoE 为 True\n",
|
||
"LMConfig_MoE = LMConfig(n_layers=2, use_moe=True)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 22,
|
||
"id": "534e4ffa-2a4d-4a0d-8906-997f23ec5098",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"shape: topk_idx = torch.Size([64, 2]), topk_weight = torch.Size([64, 2]), aux_loss = 0.10163114219903946\n",
|
||
"token 0 选择的专家：idx = tensor([2, 0]), weight = tensor([5.8955e-21, 1.9281e-21], grad_fn=<SelectBackward0>)\n",
|
||
"辅助损失：aux_loss = 0.10163114219903946\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"gate = MoEGate(LMConfig_MoE)\n",
|
||
"hidden_states = torch.randn((4, 16, 512))\n",
|
||
"topk_idx, topk_weight, aux_loss = gate(hidden_states)\n",
|
||
"print(f'shape: topk_idx = {topk_idx.shape}, topk_weight = {topk_weight.shape}, aux_loss = {aux_loss}')\n",
|
||
"print(f'token 0 选择的专家：idx = {topk_idx[0]}, weight = {topk_weight[0]}')\n",
|
||
"print(f'辅助损失：aux_loss = {aux_loss}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 23,
|
||
"id": "521d65d2-bbcb-4b5c-8e1e-d532ae79b6de",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"del gate"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "b3773961-8e78-4bd1-91af-8d8384bd31d3",
|
||
"metadata": {},
|
||
"source": [
|
||
"### MoE Feed Forward NetWork\n",
|
||
"\n",
|
||
"完成 MoE 门控单元的设计后，我们可以对 MoE 前向传播网络进行重新设计."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 24,
|
||
"id": "31c28178-7faf-49aa-a813-5105129eefcb",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"class MoEFeedForward(nn.Module):\n",
|
||
" def __init__(self, config: LMConfig):\n",
|
||
" super().__init__()\n",
|
||
" self.config = config\n",
|
||
" # 可选专家层\n",
|
||
" self.experts = nn.ModuleList([\n",
|
||
" FeedForward(config)\n",
|
||
" for _ in range(config.n_routed_experts)\n",
|
||
" ])\n",
|
||
" self.gate = MoEGate(config)\n",
|
||
" # 共享专家层d\n",
|
||
" if config.n_shared_experts is not None:\n",
|
||
" self.shared_experts = FeedForward(config)\n",
|
||
"\n",
|
||
" def forward(self, x):\n",
|
||
" identity = x\n",
|
||
" orig_shape = x.shape\n",
|
||
" bsz, seq_len, _ = x.shape\n",
|
||
" # 使用门控机制选择专家\n",
|
||
" topk_idx, topk_weight, aux_loss = self.gate(x)\n",
|
||
" x = x.view(-1, x.shape[-1])\n",
|
||
" flat_topk_idx = topk_idx.view(-1)\n",
|
||
" if self.training:\n",
|
||
" # 训练模式下，重复输入数据\n",
|
||
" x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)\n",
|
||
" y = torch.empty_like(x, dtype=torch.float16)\n",
|
||
" # 将输入数据中选择第 i 个专家的部分输入到该专家网络中进行处理，并将输出结果存储到 y 中对应位置的元素\n",
|
||
" for i, expert in enumerate(self.experts):\n",
|
||
" y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(y.dtype) # 确保类型一致\n",
|
||
" y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1) # 加权求和\n",
|
||
" y = y.view(*orig_shape) # 恢复原始形状\n",
|
||
" else:\n",
|
||
" # 推理模式下，只选择最优专家\n",
|
||
" y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)\n",
|
||
" if self.config.n_shared_experts is not None:\n",
|
||
" y = y + self.shared_experts(identity)\n",
|
||
" self.aux_loss = aux_loss\n",
|
||
" return y\n",
|
||
"\n",
|
||
" @torch.no_grad()\n",
|
||
" def moe_infer(self, x, flat_expert_indices, flat_expert_weights):\n",
|
||
" expert_cache = torch.zeros_like(x)\n",
|
||
" idxs = flat_expert_indices.argsort() # [exp_tok_11, exp_tok_12, exp_tok_21, exp_tok_22] -> idxs[2 * token_idxs] = moe_expert_name\n",
|
||
" tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0) # [1, 3, 1, 2] -> bincount -> [0, 2, 1, 1] -> cumsum -> [0, 2, 3, 4]\n",
|
||
" token_idxs = idxs // self.config.num_experts_per_tok\n",
|
||
" # 例如当tokens_per_expert=[6, 15, 20, 26, 33, 38, 46, 52]\n",
|
||
" # 当token_idxs=[3, 7, 19, 21, 24, 25, 4, 5, 6, 10, 11, 12...]\n",
|
||
" # 意味着当token_idxs[:6] -> [3, 7, 19, 21, 24, 25, 4]位置的token都由专家0处理，token_idxs[6:15]位置的token都由专家1处理......\n",
|
||
" for i, end_idx in enumerate(tokens_per_expert):\n",
|
||
" start_idx = 0 if i == 0 else tokens_per_expert[i - 1]\n",
|
||
" if start_idx == end_idx:\n",
|
||
" continue\n",
|
||
" expert = self.experts[i]\n",
|
||
" exp_token_idx = token_idxs[start_idx:end_idx]\n",
|
||
" expert_tokens = x[exp_token_idx]\n",
|
||
" expert_out = expert(expert_tokens).to(expert_cache.dtype)\n",
|
||
" expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]]) # 专家网络输出加权\n",
|
||
" # 使用 scatter_add_ 进行 sum 操作\n",
|
||
" expert_cache.scatter_add_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out)\n",
|
||
"\n",
|
||
" return expert_cache"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "c1f42fce-8bca-4853-8367-929c200f1b33",
|
||
"metadata": {},
|
||
"source": [
|
||
"接下来，我们设置一个 batch size = 4, seq len =16, emb dim = 512 的输入向量，在 MoE 门控单元前向传播进行观察"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 25,
|
||
"id": "6c7fb409-be15-4ee9-8a77-edacfacb3743",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"输出张量：shape = torch.Size([4, 16, 512]), 辅助损失：aux_loss = 0\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"moe_ffn = MoEFeedForward(LMConfig_MoE).eval()\n",
|
||
"x = torch.randn((4, 16, 512))\n",
|
||
"output = moe_ffn(x)\n",
|
||
"print(f'输出张量：shape = {output.shape}, 辅助损失：aux_loss = {moe_ffn.aux_loss}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 26,
|
||
"id": "698e4148-8f68-4f05-a33a-a76f70e20760",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"del moe_ffn"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "bec4c71c-00a3-48bc-abf1-d88b6fb94444",
|
||
"metadata": {},
|
||
"source": [
|
||
"之前提到过，MiniMind MoE 和 MiniMind Dense 的最大差别在于 FFN 模块的不同，我们可以对之前声明的 MiniMindBlock 进行继承，添加 MoE FFN 选项."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 27,
|
||
"id": "3bad938b-bcf6-4bc5-ba2c-9c9689656959",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# DM stands for Dense & Moe\n",
|
||
"class MiniMindBlock_DM(MiniMindBlock):\n",
|
||
" def __init__(self, layer_id: int, config: LMConfig):\n",
|
||
" super().__init__(layer_id, config)\n",
|
||
" self.feed_forward = FeedForward(config) if not config.use_moe else MoEFeedForward(config)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 28,
|
||
"id": "778a8697-458a-47a4-b2bd-9bbdcaf64d24",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"类变量属性检查：layer_id = 1, 类函数属性检查：forward func = <bound method MiniMindBlock.forward of MiniMindBlock_DM(\n",
|
||
" (attention): Attention(\n",
|
||
" (wq): Linear(in_features=512, out_features=512, bias=False)\n",
|
||
" (wk): Linear(in_features=512, out_features=128, bias=False)\n",
|
||
" (wv): Linear(in_features=512, out_features=128, bias=False)\n",
|
||
" (wo): Linear(in_features=512, out_features=512, bias=False)\n",
|
||
" (attn_dropout): Dropout(p=0.0, inplace=False)\n",
|
||
" (resid_dropout): Dropout(p=0.0, inplace=False)\n",
|
||
" )\n",
|
||
" (attention_norm): RMSNorm()\n",
|
||
" (ffn_norm): RMSNorm()\n",
|
||
" (feed_forward): MoEFeedForward(\n",
|
||
" (experts): ModuleList(\n",
|
||
" (0-3): 4 x FeedForward(\n",
|
||
" (w1): Linear(in_features=512, out_features=1408, bias=False)\n",
|
||
" (w2): Linear(in_features=1408, out_features=512, bias=False)\n",
|
||
" (w3): Linear(in_features=512, out_features=1408, bias=False)\n",
|
||
" (dropout): Dropout(p=0.0, inplace=False)\n",
|
||
" )\n",
|
||
" )\n",
|
||
" (gate): MoEGate()\n",
|
||
" (shared_experts): FeedForward(\n",
|
||
" (w1): Linear(in_features=512, out_features=1408, bias=False)\n",
|
||
" (w2): Linear(in_features=1408, out_features=512, bias=False)\n",
|
||
" (w3): Linear(in_features=512, out_features=1408, bias=False)\n",
|
||
" (dropout): Dropout(p=0.0, inplace=False)\n",
|
||
" )\n",
|
||
" )\n",
|
||
")>\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"# 检查继承是否正常\n",
|
||
"\n",
|
||
"miniblock_dm = MiniMindBlock_DM(1, LMConfig_MoE)\n",
|
||
"print(f'类变量属性检查：layer_id = {miniblock_dm.layer_id}, 类函数属性检查：forward func = {miniblock_dm.forward}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "433c99c0-ab12-49e3-805a-493de3cd3613",
|
||
"metadata": {},
|
||
"source": [
|
||
"我们仍然设置一个 batch size = 4, seq len =16, emb dim = 512 的输入向量，在 MoE 门控单元前向传播进行观察."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 29,
|
||
"id": "0f97a54d-aa1a-45f2-ba86-4e31f17f9cb6",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"输出张量：shape = torch.Size([4, 16, 512]), KV Cache: shape Key = torch.Size([4, 16, 2, 64]), shape Value = torch.Size([4, 16, 2, 64])\n",
|
||
"辅助损失：aux_loss = 0\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"miniblock_dm.eval() # 冻结参数，直接 moe infer\n",
|
||
"x = torch.randn((4, 16, 512))\n",
|
||
"pos_cis = precompute_pos_cis(64, 16)\n",
|
||
"output, past_kv = miniblock_dm(x, pos_cis, use_cache=True)\n",
|
||
"print(f'输出张量：shape = {output.shape}, KV Cache: shape Key = {past_kv[0].shape}, shape Value = {past_kv[1].shape}')\n",
|
||
"print(f'辅助损失：aux_loss = {miniblock_dm.feed_forward.aux_loss}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 30,
|
||
"id": "925c973b-44d0-4774-a515-404f76a43a3f",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"del miniblock_dm"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "a9e64f09-ab3f-42a5-ac89-3809c1b99c26",
|
||
"metadata": {},
|
||
"source": [
|
||
"如此，我们便完成了包含 MoE 和 Dense 两种 FFN 选项的 MiniMind Block 的定义，我们对此前声明的 MiniMindLM 作继承修改，使其具备 MoE 架构."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 31,
|
||
"id": "cec79087-53e4-4925-95e2-0c46ad28aba6",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"class MiniMindLM_DM(MiniMindLM):\n",
|
||
" def __init__(self, params: LMConfig = None):\n",
|
||
" super().__init__(params)\n",
|
||
" structure = 'MoE' if params.use_moe else 'Dense'\n",
|
||
" print(f'Initializing MiniMind {structure} Model...\\n')\n",
|
||
" self.layers = nn.ModuleList([MiniMindBlock_DM(l, params) for l in range(self.n_layers)])"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 32,
|
||
"id": "cc4517ca-b5ea-4756-9c51-7da415140e12",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"Initializing MiniMind Dense Model...\n",
|
||
"\n",
|
||
"Initializing MiniMind MoE Model...\n",
|
||
"\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"# 检查一下\n",
|
||
"a = MiniMindLM_DM(LMConfig_Dense)\n",
|
||
"b = MiniMindLM_DM(LMConfig_MoE)\n",
|
||
"del a, b"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "5a586a43-6f9b-4766-8d5d-24980d1a7241",
|
||
"metadata": {},
|
||
"source": [
|
||
"接下来，我们设置一条长度为 4 的 token 序列，使用 MiniMindLM 对其执行一次前向传播，观察传播过程与返回值."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 33,
|
||
"id": "f5009f2b-600c-4c30-b57a-85757301e6de",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"Initializing MiniMind MoE Model...\n",
|
||
"\n",
|
||
"====> forward propagation started, num_minimind_blocks = 2\n",
|
||
"------------> entering minimind block: id = 0\n",
|
||
"<------------ finished, size_cache_k = torch.Size([1, 4, 2, 64]), size_cache_v = torch.Size([1, 4, 2, 64])\n",
|
||
"------------> entering minimind block: id = 1\n",
|
||
"<------------ finished, size_cache_k = torch.Size([1, 4, 2, 64]), size_cache_v = torch.Size([1, 4, 2, 64])\n",
|
||
"<==== forward propagation completed, num_kv_cache = 2\n",
|
||
"\n",
|
||
"返回 logits：size = torch.Size([1, 4, 6400]), 返回 aux_loss: 0, 返回 KV Cache List： len = 2\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"MiniMind_MoE = MiniMindLM_DM(LMConfig_MoE)\n",
|
||
"input_ids = torch.Tensor([1, 3, 5, 7]).long().reshape(1, 4)\n",
|
||
"OUT = MiniMind_MoE(input_ids, use_cache=True)\n",
|
||
"print(f'返回 logits：size = {OUT.logits.shape}, 返回 aux_loss: {OUT.aux_loss}, 返回 KV Cache List： len = {len(OUT.past_key_value)}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 34,
|
||
"id": "c0fe2fa3-f377-45da-aee1-b8f46341a470",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"gernerating new token: idx = 4\n",
|
||
"====> forward propagation started, num_minimind_blocks = 2\n",
|
||
"------------> entering minimind block: id = 0\n",
|
||
"<------------ finished, size_cache_k = torch.Size([1, 4, 2, 64]), size_cache_v = torch.Size([1, 4, 2, 64])\n",
|
||
"------------> entering minimind block: id = 1\n",
|
||
"<------------ finished, size_cache_k = torch.Size([1, 4, 2, 64]), size_cache_v = torch.Size([1, 4, 2, 64])\n",
|
||
"<==== forward propagation completed, num_kv_cache = 2\n",
|
||
"\n",
|
||
"gernerating new token: idx = 5\n",
|
||
"====> forward propagation started, num_minimind_blocks = 2\n",
|
||
"------------> entering minimind block: id = 0\n",
|
||
"<------------ finished, size_cache_k = torch.Size([1, 1, 2, 64]), size_cache_v = torch.Size([1, 1, 2, 64])\n",
|
||
"------------> entering minimind block: id = 1\n",
|
||
"<------------ finished, size_cache_k = torch.Size([1, 1, 2, 64]), size_cache_v = torch.Size([1, 1, 2, 64])\n",
|
||
"<==== forward propagation completed, num_kv_cache = 2\n",
|
||
"\n",
|
||
"gernerating new token: idx = 6\n",
|
||
"====> forward propagation started, num_minimind_blocks = 2\n",
|
||
"------------> entering minimind block: id = 0\n",
|
||
"<------------ finished, size_cache_k = torch.Size([1, 1, 2, 64]), size_cache_v = torch.Size([1, 1, 2, 64])\n",
|
||
"------------> entering minimind block: id = 1\n",
|
||
"<------------ finished, size_cache_k = torch.Size([1, 1, 2, 64]), size_cache_v = torch.Size([1, 1, 2, 64])\n",
|
||
"<==== forward propagation completed, num_kv_cache = 2\n",
|
||
"\n",
|
||
"new tokens list :[tensor([[4065]]), tensor([[431]]), tensor([[2541]])]\n",
|
||
"\n",
|
||
"生成结果：tensor([[ 1, 3, 5, 7, 4065, 431, 2541]])\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"# 我们让 MiniMind 根据我们设计的 四个输入 token 生成输出\n",
|
||
"out = MiniMind_MoE.generate(input_ids, max_new_tokens=8, use_cache=True)\n",
|
||
"print(f'生成结果：{out}')"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 35,
|
||
"id": "4cdccdac-930e-4670-8605-2963d0e62d1e",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"del MiniMind_MoE"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "903f2172-f021-442c-ae02-611a30506ebd",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 参考资料\n",
|
||
"\n",
|
||
"- [十分钟读懂旋转编码（RoPE）](https://www.zhihu.com/tardis/zm/art/647109286?source_id=1003)\n",
|
||
"- [混合专家模型 MoE 的全面指南](https://blog.csdn.net/Code1994/article/details/145209710#:~:text=%E4%B8%BA%E4%BA%86%E5%9C%A8%E8%AE%AD%E7%BB%83%E8%BF%87%E7%A8%8B%E4%B8%AD%E5%AE%9E%E7%8E%B0%E4%B8%93%E5%AE%B6%E7%9A%84%E6%9B%B4%E5%9D%87%E5%8C%80%E5%88%86%E5%B8%83%EF%BC%8C%E8%BE%85%E5%8A%A9%E6%8D%9F%E5%A4%B1%EF%BC%88%E4%B9%9F%E7%A7%B0%E4%B8%BA%E8%B4%9F%E8%BD%BD%E5%9D%87%E8%A1%A1%E6%8D%9F%E5%A4%B1%EF%BC%89%E8%A2%AB%E6%B7%BB%E5%8A%A0%E5%88%B0%E4%BA%86%E7%BD%91%E7%BB%9C%E7%9A%84%E5%B8%B8%E8%A7%84%E6%8D%9F%E5%A4%B1%E4%B8%AD%E3%80%82%20%E5%AE%83%E5%A2%9E%E5%8A%A0%E4%BA%86%E4%B8%80%E4%B8%AA%E7%BA%A6%E6%9D%9F%EF%BC%8C%E8%BF%AB%E4%BD%BF%E4%B8%93%E5%AE%B6%E5%85%B7%E6%9C%89%E7%9B%B8%E7%AD%89%E7%9A%84%E9%87%8D%E8%A6%81%E6%80%A7%E3%80%82,%E8%BF%99%E4%B8%AA%E8%BE%85%E5%8A%A9%E6%8D%9F%E5%A4%B1%E7%9A%84%E7%AC%AC%E4%B8%80%E4%B8%AA%E7%BB%84%E6%88%90%E9%83%A8%E5%88%86%E6%98%AF%E5%9C%A8%E6%95%B4%E4%B8%AA%E6%89%B9%E6%AC%A1%E4%B8%AD%E5%AF%B9%E6%AF%8F%E4%B8%AA%E4%B8%93%E5%AE%B6%E7%9A%84%E8%B7%AF%E7%94%B1%E5%99%A8%E5%80%BC%E8%BF%9B%E8%A1%8C%E6%B1%82%E5%92%8C%EF%BC%9A%20%E8%BF%99%E4%B8%BA%E6%88%91%E4%BB%AC%E6%8F%90%E4%BE%9B%E4%BA%86%E6%AF%8F%E4%B8%AA%E4%B8%93%E5%AE%B6%E7%9A%84%E9%87%8D%E8%A6%81%E6%80%A7%E5%88%86%E6%95%B0%EF%BC%8C%E8%BF%99%E4%BA%9B%E5%88%86%E6%95%B0%E8%A1%A8%E7%A4%BA%E6%97%A0%E8%AE%BA%E8%BE%93%E5%85%A5%E6%98%AF%E4%BB%80%E4%B9%88%EF%BC%8C%E7%BB%99%E5%AE%9A%E4%B8%93%E5%AE%B6%E8%A2%AB%E9%80%89%E4%B8%AD%E7%9A%84%E5%8F%AF%E8%83%BD%E6%80%A7%E3%80%82%20%E6%88%91%E4%BB%AC%E5%8F%AF%E4%BB%A5%E5%88%A9%E7%94%A8%E8%BF%99%E4%BA%9B%E5%88%86%E6%95%B0%E6%9D%A5%E8%AE%A1%E7%AE%97%E5%8F%98%E5%BC%82%E7%B3%BB%E6%95%B0%EF%BC%88CV%EF%BC%89%EF%BC%8C%E5%AE%83%E5%91%8A%E8%AF%89%E6%88%91%E4%BB%AC%E4%B8%93%E5%AE%B6%E4%B9%8B%E9%97%B4%E7%9A%84%E9%87%8D%E8%A6%81%E6%80%A7%E5%88%86%E6%95%B0%E7%9A%84%E5%B7%AE%E5%BC%82%E7%A8%8B%E5%BA%A6%E3%80%82)"
|
||
]
|
||
}
|
||
],
|
||
"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
|
||
}
|