├── .github └── FUNDING.yml ├── README.md ├── assets └── cover.png ├── chapter1 ├── 01-1.png ├── 01-2.png └── README.md ├── chapter2 ├── 02-1.png ├── 02-2.png ├── 02-3.png ├── README.md ├── mh_causal_selfattention.py └── selfattention.py └── chapter3 ├── 03-1.png ├── 03-2.png ├── 03-3.png └── README.md /.github/FUNDING.yml: -------------------------------------------------------------------------------- 1 | # These are supported funding model platforms 2 | 3 | github: [KaihuaTang] 4 | patreon: # Replace with a single Patreon username 5 | open_collective: # Replace with a single Open Collective username 6 | ko_fi: # Replace with a single Ko-fi username 7 | tidelift: # Replace with a single Tidelift platform-name/package-name e.g., npm/babel 8 | community_bridge: # Replace with a single Community Bridge project-name e.g., cloud-foundry 9 | liberapay: # Replace with a single Liberapay username 10 | issuehunt: # Replace with a single IssueHunt username 11 | lfx_crowdfunding: # Replace with a single LFX Crowdfunding project-name e.g., cloud-foundry 12 | polar: # Replace with a single Polar username 13 | buy_me_a_coffee: tkhchipaomg 14 | thanks_dev: # Replace with a single thanks.dev username 15 | custom: ['https://kaihuatang.github.io/donate'] 16 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # 如何从零构建“小”大语言模型 2 | 3 | -- Build a small LLM from scratch: a tutorial -- 4 | 5 |
6 | logo 7 |
8 | 9 | 申明：本教程的所有内容(文字，图片，代码等)可以用于非盈利目的个人使用和分享。但如果用于盈利目的，包括但不限于卖课，公众号，视频号等需要经由作者的批准。谢谢理解。 10 | 11 | - [\[知乎链接\]](https://zhuanlan.zhihu.com/p/19275166926) 12 | - [\[B站链接 (正在建设中)\]]() 13 | - [\[Substack (under construction)\]]() 14 | 15 | ## 前言 16 | 进入工业界两年，一回头发现和在学校相比，时间似乎溜走地更悄无声息了。没有论文来总结自己每个阶段的思考和成果，似乎我的价值只存在于这六七人小团队的梦呓呢喃中，一旦离开了屈指可数的小圈子，自己这两年的所有足迹便被一个响指抹平了。本着知识就应该被分享的开源精神和一点无法被公司满足的小虚荣心，我决定写一个系列分享一下自己这两年从事大语言模型的一些理解和思考。从23年7月开始到24年底，我在公司主要做了两个和大语言模型相关的项目。其一是从23年7月开始为期半年的和中国移动合作的端侧小模型项目，在这项目中我们算法团队四个人从零开始参考[GPT2](https://github.com/openai/gpt-2)和[ChatGLM](https://github.com/THUDM/ChatGLM-6B)，训练了一个0.34B的中文语言模型（出于端侧芯片算力和我们自身训练资源和时间的考量，我们在项目要求时限内仅能训练GPT2-medium参数量的小模型），并在自建的家庭对话垂域数据上进行微调，最后通过ONNX部署在移动端的安卓智慧屏，这个项目参展了2023年中国移动全球合作伙伴大会，到24年初我们又更新了一版1B的模型进一步优化了聊天效果。第二个项目是24年5月开始的对公司内某个图文多模态大模型进行算力优化的项目，我们参考了一些开源的论文，通过对网络结构和推理逻辑的调整在量化模型的基础上进一步提升了30%的推理速度。在这两个项目中，虽然训练规模有限，但我也算是完整地了解并实践了大语言模型和图文多模态大模型的网络结构设计和训练流程，抛开那些无法分享的公司项目细节，我打算整理一份比较基础的，从零开始实现大语言模型的教程，让新入门的同学们可以更快的了解必要的知识，少走弯路。当然同时也可以作为一个记录我思考的笔记本，供该领域的从业者们参考和交流，也请大家检验下我自己的认知和理解是否存在偏差。 17 | 18 | ## 涵盖范围 19 | 20 | 1. 该系列的目的是让读者可以在基础的pytorch上，不依赖任何其他现成的外部库，从零开始理解并实现一个大语言模型的所有组成部分，以及训练微调代码，因此读者仅需python，pytorch，numpy和最基础深度学习背景知识即可。 21 | 22 | 2. 为了让读者能够亲手实践，我们希望能从小模型开始，让大家尽可能可以自己在本地跑起来 23 | 24 | 3. 后续也考虑拓展到图文多模态大语言模型。 25 | 26 | 4. 考虑到国内外网上已经有大量现成的大语言模型教程和代码，本系列除了单纯的梳理知识点外，也记录了自己在实践中的思考和做项目时遇到的具体业界问题。 27 | 28 | 5. 还有一些最新热点模型的知识点，比如deepseek的网络结构和正常的大语言模型有什么区别，为什么生成式语言模型一定要因果注意力（causal attention），pytorch的动态长度推理怎么转换为需要静态张量形状的ONNX推理格式，如何简单有效地加速首轮问答响应时间，RMSNorm归一化层在只支持FP16计算的NPU芯片上怎么解决值域越界，tokenizer分词器词表裁剪等。 29 | 30 | ## 章节链接 31 | 本系列目前计划将内容分为如下章节： 32 | 33 | 1. [大语言模型结构概览](chapter1/README.md) 34 | 2. [注意力模块与KV Cache](chapter2/README.md) 35 | 3. [DeepSeekV3的注意力优化](chapter3/README.md) 36 | 3. 旋转位置编码 (待更新) 37 | 4. 前馈网络 (待更新) 38 | 5. 归一化层 (待更新) 39 | 6. tokenizer分词器 (待更新) 40 | 7. 文本预训练 (待更新) 41 | 8. 对话数据微调 (待更新) 42 | 9. LoRA高效微调 (待更新) 43 | 10. 视觉Transformer网络 (？待定) 44 | 11. 图文多模态网络 (？待定) 45 | 12. 张量并行多卡部署 (？待定) 46 | 47 | 具体章节名称和数量在实际撰写时可能进行调整。 48 | 49 | ## 引用链接 50 | 51 | ``` 52 | @misc{tang2025all, 53 | title = {Building a Small LLM from Scratch: a tutorial}, 54 | author = {Tang, Kaihua and Zhang, Huaizheng}, 55 | year = {2025}, 56 | note = {\url{https://github.com/KaihuaTang/Building-a-Small-LLM-from-Scratch}}, 57 | } 58 | ``` 59 | -------------------------------------------------------------------------------- /assets/cover.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KaihuaTang/Building-a-Small-LLM-from-Scratch/b63c6353fa87896a446dc6ec08f792f14aebd72b/assets/cover.png -------------------------------------------------------------------------------- /chapter1/01-1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KaihuaTang/Building-a-Small-LLM-from-Scratch/b63c6353fa87896a446dc6ec08f792f14aebd72b/chapter1/01-1.png -------------------------------------------------------------------------------- /chapter1/01-2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KaihuaTang/Building-a-Small-LLM-from-Scratch/b63c6353fa87896a446dc6ec08f792f14aebd72b/chapter1/01-2.png -------------------------------------------------------------------------------- /chapter1/README.md: -------------------------------------------------------------------------------- 1 | # 01 大语言模型结构概览 2 | 3 | 申明：本教程的所有内容(文字，图片，代码等)可以用于非盈利目的个人使用和分享。但如果用于盈利目的，包括但不限于卖课，公众号，视频号等需要经由作者的批准。谢谢理解。 4 | 5 | [\[主目录链接\]](https://github.com/KaihuaTang/All-you-need-to-know-about-LLM#章节链接) 6 | 7 | 8 | ## 前言 9 | 在具体深入了解每个网络模块之前，让我们先整体了解一下一个主流大语言模型的网络结构都有哪些组成部分。这里的主流结构主要是指的Qwen系列，LLaMA系列和DeepSeek系列等。最近比较火的国产DeepSeek大模型虽然对注意力模块(Attention)和前馈网络(FFN)的设计有调整，但整体架构依然符合我们图1所示的主流结构，在之后的章节中我也会把DeepSeek的改进作为拓展知识进行简单的介绍。在本系列中，为了不至于偏离主题陷入无止境背景知识罗列，我会本着**2大原则** 10 | 11 | - 最小化地收录实现大语言模型所需的知识，避免不必要的展开(展开太多我也不懂)。 12 | - 但同时我也会尽量保证知识的完备性，让大家不需要看的时候反复跳转出去搜索不理解的背景知识。 13 | 14 | 下面我就从文本的输入到输出，端到端地讲解下主流大语言模型的整体架构。 15 | 16 | ## 一. 模型结构概览 17 |
18 | logo 19 | 图1：主流大语言模型输入到输出的流程图 20 |
21 | 22 | 上图是一个主流大语言模型(如[Qwen2](https://huggingface.co/docs/transformers/en/model_doc/qwen2),[LLaMA3](https://github.com/meta-llama/llama3),[DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3)等)从文本输入到文本输出的完整流程。可以看到其并不复杂。分成了 前处理(Tokenizer) -> **大模型：嵌入层(Embedding Layer)** -> **Transformer模块** -> **输出层(Output Layer)** -> 后处理(Tokenizer) 几个部分。 23 | 24 | ### 0. Tokenizer (分词数字化) 25 | 26 | 其中预处理和后处理使用的是名为tokenizer的文字编码解码模块，由于其并非神经网络，而是基于规则的传统分词器，因此我们把这部分的介绍放到网络结构后面。在开始阶段我们仅需知道tokenizer是一个给文本进行编号的模块，比如"你"=1，"好"=2，"世"=3，"界"=4。通常如果"你好"二字在数据集中经常同时出现，为了提升编码效率，我们会赋予"你好"一个独立的编号，例如"你好"=5。这样当"你好世界"出现时，我们优先将其编码为三位的[5,3,4]而非四位的[1,2,3,4]以此来提升编码效率，降低模型推理长度。 27 | 28 | 大家可以去这个网站，快速直观理解tokenizer的作用： 29 | 30 | - [Tokenizer可视化](https://tiktokenizer.vercel.app/) 31 | 32 | 除去tokenizer的编码和解码，上图的各个黄底小模块才是我们平时说的大语言模型的本体，也就是可学习的带参数神经网络模块。其中包含嵌入层(Embedding Layer)，若干Transformer模块，和输出层(Output Layer)。 33 | 34 | ### 1. 嵌入层(Embedding Layer) 35 | 我们可以将嵌入层理解为给tokenizer之后，得到的每个token编号，都学习了一个特征向量来表达该token所包含的信息。嵌入层(Embedding Layer)的输入是tokens，也就是一连串离散的数字编号(通常数据类型为LongTensor)，而其输出则是连续的浮点向量(通常为torch.float32, torch.float16或torch.bfloat16，具体类型取决于模型的精度)。嵌入层包含一个Weight权重，Weight的张量形状为(vocab_size, embedding_dim), 前者为tokenzier中字典的长度，后者为每个token在计算时的特征维度。 36 | 37 | 嵌入层的权重虽然也和线性层Linear一样叫Weight，但其的计算不同于正常的pytorch的线性层Linear，不是进行矩阵计算，而是直接将输入看作索引去提取Weight对应位置的特征向量，**可以理解为根据输入的token的数字，去查字典对应的特征向量的过程**。比如下面例子中，输入的[[1,3,5,7]]就是分别提取了Embedding权重的第1行，第3行，第5行，和第7行。每行是一个长度为embedding_dim的特征向量。 38 | 39 | ``` 40 | import torch 41 | import torch.nn as nn 42 | 43 | # Define the embedding layer 44 | vocab_size = 30000 # Number of unique categories (e.g., vocabulary size) 45 | embedding_dim = 1024 # Dimension of each embedding vector 46 | 47 | embedding_layer = nn.Embedding(vocab_size, embedding_dim) 48 | 49 | # Example input (batch of tokens) 50 | batch_tokens = torch.tensor([[1, 3, 5, 7]]) # tensor shape [1, 4] 51 | 52 | # Forward pass to obtain embeddings 53 | output_embeddings = embedding_layer(batch_tokens) # tensor shape [1, 4, 1024] 54 | ``` 55 | 56 | ### 2. Transformer模块 57 | 大语言模型网络在嵌入层和输出层之间会串联若干个Transformer模块作为网络的主体，通常情况下同一个网络的所有Transformer模块都是相同的结构，但也有一些尝试串联不同网络模块的探索，这里暂不展开。 58 | 59 | 每个Transformer模块内部由注意力模块(Attention)，前馈网络(Feed-Forward Network / Multi-Layer Perceptron)，归一化层(Normalization)组成。每个模块的输入和输出的张量数据形状都是[batch_size, sequence_length, embedding_dim]。其中通常包含两个残差连接(Residual Connection)，即类似x=x+f(x)的网络结构，也就是图1中每个模块在Normalization+Attention前后的残差连接和Normalization+FFN前后的残差连接。这里我们图1展示的结构为主流的[Qwen2](https://huggingface.co/docs/transformers/en/model_doc/qwen2)与[LLaMA3](https://github.com/meta-llama/llama3)，其主干上全是残差连接，两个残差连接之间不会有额外的模块，这通常被认为可以将网络训练的更深，当然Transformer模块内部也有不同于该主流形式的结构，我们会放到归一化层Normalization章节进行讲解。 60 | 61 | 作为大语言模型的主体，我们会将Transformer模块的各个子模块拆开在后续章节一一讲解，这里就不更进一步展开了，我们仅需将Transformer模块看作一个学习token与token之间联系并存储知识的神经网络模块即可。 62 | 63 | ### 3. 输出层(Output Layer) 64 | 经过若干层Transformer模块提取token之间的相关性和网络本身学到的知识，我们最终就会进入输出层。输出层本质就是一个线性层网络，其通用数学形式为 $y=xW^T+b$ 。而大语言模型的输出层一般会省略$b$来减少模型的偏见。输出层Weight的张量形状为(vocab_size, embedding_dim)，进入这个输出层的张量形状为[batch_size, sequence_length, embedding_dim]，通过输出层后，得到的张量形状为[batch_size, sequence_length, vocab_size]。 65 | 66 | token在输出层的输出会用于预测其相邻的下一个token。具体的预测方式有贪婪预测和采样预测两种： 67 | - 其一的贪婪预测直接看作分类任务，取输出在最后一维的最大值的位置，即argmax计算，由于最后一维大小为vocab_size，所以刚好对应了tokenizer的词表，可以解码为对应的字符串。 68 | - 其二的的采样预测则需要将输出通过softmax归一化变成输出每个token的概率，然后基于其概率进行采样，采样策略有很多不同方案，比如top-k, top-p, temperature和其组合等，这里就不一一展开。 69 | 70 | ### :bulb: 扩展知识：嵌入层输出层共享权重 71 | 这里扩展一个额外的知识点，由于嵌入层和输出层的本质都是用作将token从离散的编号空间和连续的特征空间之间进行映射，只是方向不同，但本质是做的同样的事情，因此理论上可以用同一套网络权重来做这事。即共享嵌入层和输出层的Weight。 72 | 73 | 这种共享权重的操作一般并不能提升网络的效果，但有2大好处 74 | 75 | 1. 降低参数量，也就降低了模型大小。所以通常被用在较小的大语言模型上，比如[1B参数量的大语言模型](https://arxiv.org/pdf/2402.02791)等。 76 | 2. 节省可训练参数。因为vocab_size很大（通常有几万到十几万），嵌入层和输出层其实参数量其实并不小，将他们的权重共享，可以空出更多的参数用到Transformer模块上。 77 | 78 | 但需要注意的是，其计算量并不会减少，因为嵌入层和输出层的输入输出和权重的形状都没有变，还是会分别计算的。 79 | 80 | ## 二. 大语言模型推理的两个阶段 81 | 82 |
83 | logo 84 | 图2：大语言模型推理的两个阶段: 预装填(Prefilling)阶段与解码(Decoding)阶段 85 |
86 | 87 | 88 | 目前主流的大语言模型框架又被称为decoder-only架构(仅解码器架构)，也就是没有encoder只有decoder。那为什么又会说大语言模型的推理有预装填(Prefilling)与解码(Decoding)两个阶段呢？不应该只有decoding吗？ 89 | 90 | 在回答上述疑问之前我们先说说什么样的Transformer网络可以被称为decoder解码器。顾名思义，decoder是用来产生输出的，对于文本生成这个任务来说，我们无法一次性穷举文本的所有可能性，比如当我们将tokenizer的字典集扩展到十几万时，可能最长的一个token也就四五个汉字，如果要用一个token表达一个完整的句子，可能需要兆亿的( $100000^N$ )字典集来穷举所有文字的可能组合，这显然是不合理的。因此我们将生成一个句子，拆成生成多个token的任务，每次仅生成一个token(2-3个汉字)，不停迭代直到输出完整的一句话（也就是*自回归*）。在这种结构里，生成的第一个token只能看到用户的输入，而生成的第二个token可以看到用户的输入加已经生成的第一个token，以此类推，越往后的token能看到输入越多（因为已经生成的token越多），这就天然地形成了一个叫causal attention因果注意力的结构，下章会详细展开。我们这里需要知道的是，因果注意力的特性就是后面的token可以看到前面的，前面的看不到后面的，例如第3个token计算的时候可以利用第1-3个token的信息，但不能利用第4个和往后的token的信息，因为这时候后面的token还没产生。这样的网络就是一个decoder transformer网络。 91 | 92 | 而大语言模型框架被称为decoder-only架构是因为不管是输入的问题还是输出的回答，都统一按照decoder的形式处理，即前面的token看不到后面的。所以大语言模型只有一种网络结构，也只有一套网络参数。不管是处理问题还是生成答案，数学上都是走的同一套网络的计算形式，也利用的是同一套网络参数。 93 | 94 | ### 预装填(Prefilling)阶段 95 | 虽然计算方式和计算参数相同，但问题和答案终究是有区别的，区别在于问题是一次性完整的丢给模型的，答案是需要一个一个生成的。即把问题丢给网络时，虽然我们处理第1个token时不会利用第2个token和往后token的信息，但其实我们已经知道第2个token了，而事实上第1个token的输出我们也并不会使用，不然如上图2所示，第1个token的输出很可能是“人民”，可问题中第2个token是“首都”，所以我们肯定是以输入的问题为准。在整个问题的token序列中，我们真正在意的也是会看作输出的其实只有问题的最后一个token的输出，因为它给出了答案的第一个token。 96 | 97 | 所以综上，预装填阶段除了最后一个token需要输出外，**其他token的输出均可丢弃**，因为输入问题的所有token已经知晓，无需模型预测。这带来的一个好处就是，我们可以一次性把整个问题丢进去，比如问题有N个token，预装填阶段就输入N个token，然后输出最后一个token的输出。 98 | 99 | *代码可以参考：* 100 | 101 | ### 解码(Decoding)阶段 102 | 接着上面的例子，在解码阶段，数学上等价的计算其实是：第一次解码输入N+1个token（完整的问题+预测的一个token）然后取最后一个token 的输出，第二次解码输入N+2个token，然后继续取最后一个token 的输出，直到最后模型预测出了结束符，标志着回答结束。但这似乎就和预装填没什么区别了。这里就得提前说一个概念叫KV Cache的概念了，这个概念是大语言模型推理优化的一个核心概念，可以大大加速模型推理。我们会在介绍完注意力模块的章节之后完整地介绍KV Cache，这里仅需知道有了KV Cache之后，我们在解码阶段，网络每次仅需输入1个token在加上历史的KV Cache，就可以预测下一个token，而不需要重复计算完整的N+1个token。也因此让预装填和解码两个阶段有了不同的网络调用形式。 103 | 104 | 但我们需要理解的是，归根结底，我们在预装填和解码两个阶段调用的都是同一个模型，因此数学上解码阶段是可以被等价改写成预装填的计算形式的。 105 | 106 | 课外思考：通过KV Cache这个名字，聪明的同学不知道能否想到其原理呢？我可以剧透的是，这是由于因果注意力结构，让前面的token不会随后面的token生成而更新，在这种情况下，不知道大家是否可以想到为什么有了KV Cache我们就可以仅需要1个token+前N个token的KV Cache就实现等同于输入N+1个token同时计算的效果了？以及这到底能节省多少计算？我会在后面的章节回答这个问题。 107 | 108 | ### :bulb: 扩展知识: Encoder-Decoder 与 Prefill-Decode区别 109 | 在我个人的理解里：Encoder-Decoder架构的本质在于不等价性，其输入和输出是走的两套网络。通常这意味这处理问题和生成答案的两个网络有不同的网络结构，比如问题用双向自注意力结构，答案用交叉注意力结构。或者即便是相同的网络参数，处理问题用双向注意力，处理答案用因果注意力[(如早期版ChatGLM)](https://github.com/THUDM/ChatGLM-6B)，我认为只要体现了数学上的不等价性这也都可以算是Encoder-Decoder架构。 110 | 111 | 而Prefill-Decode和Encoder-Decoder的本质区别我认为是在于Prefill-Decode的等价性。上面我们也说了，大语言模型的Prefill和Decode在计算上是等价的，Decode过程也可以数学上写成或代码上实现成Prefill的形式，只是因为Decode阶段用KV Cache做加速更高效（*利用了系统优化技巧，不改变其本质数学意义*），所以才让他们看似走了不同的计算流程。 112 | 113 | --------- 114 | 115 | [\[主目录链接\]](https://github.com/KaihuaTang/All-you-need-to-know-about-LLM#章节链接) 116 | 117 | ## 引用链接 118 | 119 | ``` 120 | @misc{tang2025all, 121 | title = {Building a Small LLM from Scratch: a tutorial}, 122 | author = {Tang, Kaihua and Zhang, Huaizheng}, 123 | year = {2025}, 124 | note = {\url{https://github.com/KaihuaTang/Building-a-Small-LLM-from-Scratch}}, 125 | } 126 | ``` -------------------------------------------------------------------------------- /chapter2/02-1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KaihuaTang/Building-a-Small-LLM-from-Scratch/b63c6353fa87896a446dc6ec08f792f14aebd72b/chapter2/02-1.png -------------------------------------------------------------------------------- /chapter2/02-2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KaihuaTang/Building-a-Small-LLM-from-Scratch/b63c6353fa87896a446dc6ec08f792f14aebd72b/chapter2/02-2.png -------------------------------------------------------------------------------- /chapter2/02-3.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KaihuaTang/Building-a-Small-LLM-from-Scratch/b63c6353fa87896a446dc6ec08f792f14aebd72b/chapter2/02-3.png -------------------------------------------------------------------------------- /chapter2/README.md: -------------------------------------------------------------------------------- 1 | # 02 注意力模块与KV Cache 2 | 3 | 申明：本教程的所有内容(文字，图片，代码等)可以用于非盈利目的个人使用和分享。但如果用于盈利目的，包括但不限于卖课，公众号，视频号等需要经由作者的批准。谢谢理解。[\[知乎链接\]](https://zhuanlan.zhihu.com/p/19275166926) 4 | 5 | [\[主目录链接\]](https://github.com/KaihuaTang/All-you-need-to-know-about-LLM#章节链接) 6 | 7 | 8 | ## 前言 9 | 作为大语言模型中核心的核心，我将注意力模块排在了其他模块之前放在最前面讲解。我们在本章里会从其原理，结构，各种优化版本讲到目前主流开源大语言模型的具体代码(Qwen2/LLaMA3, DeepSeek-V3)。但本章节仅限于对注意力结构本身原理的阐述，并不会太涉及优化，比如目前主流的[FlashAttention-v1/v2/v3](https://github.com/Dao-AILab/flash-attention)或者一些[线性注意力架构](https://arxiv.org/abs/2401.04658)，这些要么就是基于硬件做的数学等价优化，要么就是完全改变了传统注意力计算形式尚没有被主流认可。不过我相信只要理解了注意力结构的基本原理，后续优化相关工作也很容易入门了 10 | 11 | ## 一. 注意力原理 12 | 和传统的卷积网络(CNN)或者多层感知器(MLP)直接对输入特征做映射不同，要想理解注意力模块的运行原理，我们不能将其看作对输入的单纯映射，即y=f(x)。一次最基本的注意力计算，需要三个输入QKV：Query，Key，Value。其中Query表示当前想要检索或关注的信息，可以理解为模型“要解决的问题”或“要检索的关键词”。Key表示已知信息的索引或标识，用于和 Query 进行匹配，看看是否匹配。而每个Value都和一个Key一一对应，表示其指代的具体的内容或信息，当 Key 与 Query 匹配时，就会将这个Key对应的Value 视为“回答”或“检索出的信息”。注意，虽然Key和Value必须一一对应，但Query则可以有任意的数量。 13 | 14 | 我们可以通过下面例子理解注意力模块的计算方式：假设我们需要查询班上同学的身高。我们对班上的四个同学小明，小凯，小丽，小花先构建他们对应的Key和Value信息字典。其中每个Key指代一个同学，Value则是其对应的身高。而Query则是一个个具体需要解决的问题。比如我们想要知道小明的身高，则Query1就是小明。我们还想知道班上女生的身高(平均身高)，Query2就是女生。而Query和Key之间的匹配程度就是注意力图，也叫Attention Map，一般注意力图会在数值上利用softmax非线性函数上进行归一化，保证每一行和为1。然后再利用注意力图作为权重从Value中提取对应的加权信息，最后得到小明的身高175 x 1=175cm和班上女生的平均身高166 x 0.5 + 170 x 0.5 = 168cm。 15 | 16 |
17 | logo 18 | 图1：注意力原理示意图 19 |
20 | 21 | 因此，不同于单纯的映射。注意力模块还能计算信息之间的交互，这赋予了模型更大的表达能力和推理能力，让模型在处理文本、图像等复杂数据时，能够依据不同的任务“关注”到对应的重要信息，大大提升了处理效率和准确度。 22 | 23 | ## 二. 自注意力结构(Self-Attention) 24 | 25 | ### 1. 基础自注意力结构 26 | 广义的注意力结构仅仅要求K(Key)和V(Value)有同一输入序列来源，Q(Query)则可以来自其他序列。而QKV均来自于同一输入序列，且三者全部一一一对应的就是自注意力结构(Self-Attention)。注意当QKV均来自于同一输入序列时，他们的区别仅仅就是通过了不同的线性层，即Q=Linear(X), K=Linear(X), V=Linear(X)，三个线性层权重独立。因此最简单的自注意力结构可以通过如下代码实现： 27 | 28 | ``` 29 | import torch 30 | import torch.nn as nn 31 | import math 32 | 33 | class SelfAttention(nn.Module): 34 | def __init__(self, input_dim, hidden_dim): 35 | """ 36 | 简易自注意力层 37 | :param input_dim: 输入维度 38 | :param hidden_dim: 隐藏层维度 39 | """ 40 | super(SelfAttention, self).__init__() 41 | 42 | # 定义线性层，用于生成 Q、K、V 43 | self.query = nn.Linear(input_dim, hidden_dim) 44 | self.key = nn.Linear(input_dim, hidden_dim) 45 | self.value = nn.Linear(input_dim, hidden_dim) 46 | 47 | # 输出变换 48 | self.out = nn.Linear(hidden_dim, input_dim) 49 | 50 | def forward(self, x): 51 | """ 52 | :param x: 输入张量，形状 [batch_size, seq_len, hidden_dim] 53 | :output: 自注意力计算后的输出，形状 [batch_size, seq_len, hidden_dim] 54 | """ 55 | batch_size, seq_len, hidden_dim = x.shape 56 | 57 | # 线性变换得到 Q, K, V 58 | # 形状: [batch_size, seq_len, hidden_dim] 59 | Q = self.query(x) 60 | K = self.key(x) 61 | V = self.value(x) 62 | 63 | # 计算注意力分数: Q @ K^T / sqrt(hidden_dim) 64 | # Q: [batch_size, seq_len, hidden_dim] 65 | # K^T: [batch_size, hidden_dim, seq_len] 66 | # scores: [batch_size, seq_len, seq_len] 67 | scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(hidden_dim) 68 | 69 | # 通过 softmax 得到注意力分布 70 | attn_weights = torch.softmax(scores, dim=-1) 71 | 72 | # 注意力加权 V 73 | # [batch_size, seq_len, seq_len] x [batch_size, seq_len, hidden_dim] 74 | # => [batch_size, seq_len, hidden_dim] 75 | attn_output = torch.matmul(attn_weights, V) 76 | 77 | # 输出层 78 | output = self.out(attn_output) 79 | 80 | return output 81 | 82 | if __name__ == "__main__": 83 | # 测试代码 84 | batch_size = 2 85 | seq_len = 5 86 | input_dim = 256 87 | hidden_dim = 256 88 | x = torch.randn(batch_size, seq_len, input_dim) 89 | 90 | attention = SelfAttention(input_dim, hidden_dim) 91 | out = attention(x) 92 | print("输入形状:", x.shape) 93 | print("输出形状:", out.shape) 94 | ``` 95 | 96 | 上述代码仅需一个输入序列x即可完成注意力的计算。通过三个线性层得到Q,K,V后，我们会计算注意力分数，也就是Q和K^T的矩阵乘，注意这里torch.matmul(Q, K.transpose(-2, -1))还需除以一个系数math.sqrt(hidden_dim)，这么做的原因我认为有两个：1）一个是为了归一化，随着hidden_dim的维度增加，注意力分数很容易出现极大值，这在后续的softmax计算中很容易出现数值溢出，因为softmax中会有exp计算，会进一步放大极大值，导致梯度消失或者梯度爆炸让训练极不稳定；2）这涉及到softmax非线性函数中引入temperature的概念，通过增加参数t，softmax(x/t)可以控制输出的分布，当t=math.sqrt(hidden_dim)时，本质是防止softmax后的分布过于“尖锐”（即一个位置可能占据了绝大部分注意力），模型在训练初期就过早地偏向了某些token可能会导致对大多数token训练不足，影响模型训练效果。如下图2所示，当hidden_dim=256时，math.sqrt(hidden_dim)=16.0，同一组输入[1, 2, 7, 12, 8, 5, 2, 1]在添加归一化系数t=16.0和不添加的情况下softmax结果差异巨大，不归一化的话大部分位置的注意力权重都几乎为0，只能学到值为12的位置。添加归一化后大部分位置的注意力分数在0.1-0.2的区间，都能学到。 97 | 98 |
99 | logo 100 | 图2：softmax前添加temperature归一化系数与不添加的区别。 101 |
102 | 103 | 获得注意力分数后，我们还需要进行softmax归一化，softmax的公式如下。softmax归一化主要是为了保证每一个Query对所有Key的分数的总和为1.0，这样当我们将注意力分数作为加权系数乘Value求和后输出的向量整体分布稳定。如图1的例子，如果不做归一化，Query2女生与Key3小丽(女)和Key3小花(女)均匹配成功（即系数为1），那女生的平均身高就变成了166 x 1.0 + 170 x 1.0 = 336cm。如果说把匹配成功的分数改成0.5不就行了，这样的话Query1小明的身高不就是175 x 0.5 = 87.5cm了，所以还是需要归一化来决定什么时候系数时1.0，什么时候是0.5，这就是引入softmax非线性层的原因。 104 | 105 | $ softmax(Z)_i = \frac{e^{z_i}}{\sum_{k=i}^n e^{z_k}}, Z=(z_0, z_1, ..., z_n) $ 106 | 107 | 而后归一化的注意力权重attn_weights就会乘以value得到注意力计算的输出。我们一般还会再添加一层线性层(Output Projection)，作为注意力模块的最终输出。线性层主要就是增加模型的表达能力并且存储知识，这就没什么可以解释的了。 108 | 109 | ### 2. 多头因果自注意力结构(Multi-Head Causal Self-Attention) 110 | 但这还远不能算作自注意力结构的完整形态，在实际大语言模型的注意力结构中，我们还需引入两个重要的改动：多头注意力(multi-head attention)与因果掩码(causal mask)。其中关于多头注意力，如果说我们平时堆叠网络的模块数是增加网络的深度(串联)，那么多头注意力其实就是增加网络的广度(并联)，一个8头的注意力层相当于并联了8个独立权重的注意力层。因果掩码则和下面要说的KV Cache缓存机制息息相关，因果掩码不仅确保了输入问题和答案生成的注意计算形式完全一致，还让生成答案时可以大幅降低计算量。 111 | 112 |
113 | logo 114 | 图3：多头因果自注意力结构原理图。 115 |
116 | 117 | 如上图所示，首先多头因果自注意力结构在计算softmax前会加上一个掩码，该掩码白色部分为-inf，这样计算softmax时，-inf对应位置的输出就为0，因为exp(-inf)=0。这样可以保证第K个Query提取信息时，仅能访问第0-K个Key和Value，无法获取第K往后的Token的信息。这也就是因果注意力中“因果(Causal)”的由来，无法提取未来Token的信息。 118 | 119 | 而我图中多头注意力的呈现方式与大多数Transformer论文的配图略有不同，我认为这样更便于理解，而且数学上也完全是等价的。我们完全把Attention分成了独立的N份，每份的中间为度从hidden_dim降到了hidden_dim / N，这里的N就是头数(num_head)，这样保证了模型参数量不变的情况下，模型变成独立参数的N份，这样就实现Ensemble，Ensemble是神经网络中一种常见的提升模型能力的方式，或者说奇技淫巧(trick)，相当于三个臭皮匠胜过诸葛亮，三个独立的小模型的结果求平均会比一个大模型更好。因为这可以防止一个模型不小心陷入局部次优解。而在注意力中，多头还引入了一个新的意义，就是不同的头可以独立的关注到不同区域，也就是每个头softmax后的注意力图都不一样，增加了模型的学习和表达能力。 120 | 121 | 注意：多头注意力的实现并不需要真的维护N个子网络，而是通过reshape将hidden_dim维度拆分成两个维度num_head和head_dim维度即可，而且最终先concate再过一个大的输出线性层和过N个小的输出线性层再求和在数学上也是等价的。具体实现代码如下： 122 | 123 | ``` 124 | class MultiHeadCausalSelfAttention(nn.Module): 125 | def __init__(self, input_dim, hidden_dim, num_heads=8): 126 | """ 127 | 简易自注意力层 128 | :param input_dim: 输入维度 129 | :param hidden_dim: 隐藏层维度 130 | :param num_heads: 多头注意力的头数，这里可以指定为8表示单头注意力 131 | """ 132 | super(MultiHeadCausalSelfAttention, self).__init__() 133 | self.input_dim = input_dim 134 | self.hidden_dim = hidden_dim 135 | self.num_heads = num_heads 136 | # 将嵌入维度平分到各个头上 137 | # 注意实际应用中需要确保 embed_dim % num_heads == 0 138 | self.head_dim = hidden_dim // num_heads 139 | 140 | # 定义线性层，用于生成 Q、K、V 141 | self.query = nn.Linear(input_dim, hidden_dim) 142 | self.key = nn.Linear(input_dim, hidden_dim) 143 | self.value = nn.Linear(input_dim, hidden_dim) 144 | 145 | # 输出变换 146 | self.out = nn.Linear(hidden_dim, input_dim) 147 | 148 | def forward(self, x, mask=None): 149 | """ 150 | :param x: 输入张量，形状 [batch_size, seq_len, input_dim] 151 | :param mask: 可选的掩码（mask），形状与注意力矩阵匹配，如 [batch_size, 1, seq_len, seq_len] 152 | :return: 自注意力计算后的输出，形状 [batch_size, seq_len, input_dim] 153 | """ 154 | batch_size, seq_len, input_dim = x.shape 155 | 156 | # 线性变换得到 Q, K, V 157 | # 形状: [batch_size, seq_len, input_dim] 158 | Q = self.query(x) 159 | K = self.key(x) 160 | V = self.value(x) 161 | 162 | # 多头展开 163 | # 变换后形状: [batch_size, seq_len, num_heads, head_dim] 164 | Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim) 165 | K = K.view(batch_size, seq_len, self.num_heads, self.head_dim) 166 | V = V.view(batch_size, seq_len, self.num_heads, self.head_dim) 167 | 168 | # 将 [batch_size, seq_len, num_heads, head_dim] 转成 [batch_size, num_heads, seq_len, head_dim] 169 | Q = Q.permute(0, 2, 1, 3) # [batch_size, num_heads, seq_len, head_dim] 170 | K = K.permute(0, 2, 1, 3) 171 | V = V.permute(0, 2, 1, 3) 172 | 173 | # 计算注意力分数: Q @ K^T / sqrt(head_dim) 174 | # Q: [batch_size, num_heads, seq_len, head_dim] 175 | # K^T: [batch_size, num_heads, head_dim, seq_len] 176 | # scores: [batch_size, num_heads, seq_len, seq_len] 177 | scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim) 178 | 179 | # 如果有 mask，则在计算分数时将被mask的部分赋予一个很大的负数，以避免注意力 180 | # mask 形状一般为 [batch_size, 1, seq_len, seq_len] 或 [batch_size, num_heads, seq_len, seq_len] 181 | if mask is not None: 182 | scores = scores + mask 183 | 184 | # 通过 softmax 得到注意力分布 185 | attn_weights = torch.softmax(scores, dim=-1) 186 | print("注意力权重分布:", attn_weights) 187 | 188 | # 注意力加权 V 189 | # [batch_size, num_heads, seq_len, seq_len] x [batch_size, num_heads, seq_len, head_dim] 190 | # => [batch_size, num_heads, seq_len, head_dim] 191 | attn_output = torch.matmul(attn_weights, V) 192 | 193 | # 把多头重新拼接回原始形状 194 | # [batch_size, num_heads, seq_len, head_dim] -> [batch_size, seq_len, num_heads, head_dim] 195 | attn_output = attn_output.permute(0, 2, 1, 3).contiguous() 196 | # 拼接头部维度 197 | # => [batch_size, seq_len, input_dim] 198 | attn_output = attn_output.view(batch_size, seq_len, input_dim) 199 | 200 | # 输出层 201 | output = self.out(attn_output) 202 | 203 | return output 204 | 205 | if __name__ == "__main__": 206 | # 测试代码 207 | batch_size = 1 208 | seq_len = 5 209 | input_dim = 256 210 | hidden_dim = 256 211 | num_heads = 2 212 | x = torch.randn(batch_size, seq_len, input_dim) 213 | causal_mask = torch.triu(torch.ones(5, 5, device=x.device, dtype=torch.float), 214 | diagonal=1 215 | ) 216 | causal_mask = causal_mask.masked_fill(causal_mask == 1, float('-inf')) 217 | 218 | attention = MultiHeadCausalSelfAttention(input_dim, hidden_dim, num_heads=num_heads) 219 | out = attention(x, mask=causal_mask) 220 | print("输入形状:", x.shape) 221 | print("输出形状:", out.shape) 222 | ``` 223 | 224 | ## 三. KV Cache缓存机制 225 | 假设我们有两个token在某一层的输入特征(hidden states): $[x_1, x_2]$, 经过三个线性层，我们得到他们分别的QKV: $[q_1, q_2], [k_1, k_2], [v_1, v_2]$, Attention后对应的输出为： $[y_1, y_2]$ 。在因果掩码(Causal Mask)下，Attention的计算公式应该如下，其中OutLinear只是对输出做映射的线性层： 226 | 227 | $y_1 = OutLinear( softmax(q_1 @ [k_1]^T) @ [v_1] )$ 228 | 229 | $y_2 = OutLinear( softmax(q_2 @ [k_1,k_2]^T) @ [v_1,v_2] )$ 230 | 231 | 由于因果掩码让每个token的query只能看到其自己和其前面的token的key-value，当我们计算完这个两个token要去推理第三个token的时候，计算公式如下： 232 | 233 | $y_1 = OutLinear( softmax(q_1 @ [k_1]^T) @ [v_1] )$ 234 | 235 | $y_2 = OutLinear( softmax(q_2 @ [k_1,k_2]^T) @ [v_1,v_2] )$ 236 | 237 | $y_3 = OutLinear( softmax(q_3 @ [k_1,k_2, k_3]^T) @ [v_1,v_2,v_3] )$ 238 | 239 | 这时候我们会发现，当推理第三个token时，前两个token的计算完全没有变，也就是说 $[y_1, y_2]$ 的结果和推理第二个token时一摸一样。也许有人要说大模型除了注意力还有前馈网络FFN呀，注意，FFN的计算中每个token都是独立的，不会相互影响，只要attention的输出结果一样，后续FFN的输出结果必然一样。这时候我们就会发现，一旦有了因果掩码，当推理第N个token的时候，再去计算前0到N-1个token就完全是浪费了，因为计算结果完全一样。我们再看如果推理第N个token的时候只计算 $x_N$ 本身，他的 $q_N,k_N,v_N$ 都可以直接通过 $x_N$ 过线性层得到，其需要额外进入的就只有 $[k_1,k_2,...,k_{N-1}],[v_1,v_2,...,v_{N-1}]$ 。而这些Key和Value完全没必要重走一边线性层的计算，因为推理第N-1个token的时候已经得到过一次了，所以当时只要存下来这是再读取拿来用就好了，这也就是KV Cache，用存储和加载的带宽增加换计算节省的一种优化方式。 240 | 241 | ### 扩展知识1：为什么生成式语言模型一定要因果注意力（causal attention）？ 242 | 243 | 如果没有因果掩码，如下面公式所示，一旦推理到第3个token的时候，第1和第2个token可以看到后面的token，那么 $[y_1,y_2,y_3]$ 三个值都会更新，这时候就需要重算所有的token了。以此类推，推理到第N个token的时候，前N-1个token还得全部输入重新计算attention。这就没法用KV Cache了，因为一旦attention的输出改变了，前N-1个token的FFN也需要重新计算。 244 | 245 | $y_1 = OutLinear( softmax(q_1 @ [k_1,k_2, k_3]^T) @ [v_1,v_2,v_3] )$ 246 | 247 | $y_2 = OutLinear( softmax(q_2 @ [k_1,k_2, k_3]^T) @ [v_1,v_2,v_3] )$ 248 | 249 | $y_3 = OutLinear( softmax(q_3 @ [k_1,k_2, k_3]^T) @ [v_1,v_2,v_3] )$ 250 | 251 | 所以说因果掩码不仅仅是能让模型的推理过程更符合人类接受和输出信息的时序性。我认为其最大的价值是，引入了KV-Cache这个优化方案。这是只有在因果掩码注意力下才可以使用的。 252 | 253 | ### 扩展知识2：KV-Cache的优化能省下多少计算力？ 254 | 255 | 我们来估算下计算复杂度，我们假设特征的维度d是常数，仅考虑推理token数N的变化，那么注意力的计算复杂度其实就是来自于QKV计算中的两个矩阵乘，也就是 $O(N^2)$ ，因为线性层的复杂度此时只有 $O(N)$ 。而一旦引入了KV Cache，我们QKV的计算复杂度是要输入一个query，所以复杂度就缩减到了 $O(N)$ 。一个KV-Cache就可以将计算复杂度从$O(N^2)$降低到$O(N)$还能保证数学等价性，简直血赚。 256 | 257 | ### 扩展知识3：如何简单有效地加速首轮(Prefilling)问答响应时间？ 258 | 259 | 首轮Prefilling回答速度之所以会慢，是因为首轮需要一次性把问题的prompt全部输入，同时推理所有token还是 $O(N^2)$ 的复杂度，后续增量decoding每次只要推理一个token就快很多只要 $O(N)$ 的复杂度。 260 | 261 | 优化首轮回答响应速度是工业界一个非常重要的研究课题，尤其是超长文本输入的时候，比如把整个几十万字的一本书全部作为输入来询问关于这本书的问题时。我这里提供一个简单的思路，也是我2023年末参加中国移动全球合作伙伴大会演示时想到的一个非常简单的技巧，就是把系统prompt提前计算好存成KV-Cache。因为大模型使用时除了问题本身，我们往往还会增加一个固定的系统prompt在问题前面，这部分的token是不会变的，所以我们完全可以提前离线计算好。这个小技巧让我们当时首轮平均响应速度从7秒降低到了3-4秒（具体提升比例受实际问题长短影响）。所以说KV-Cache真是个妙不可言的好东西。我2023年第一次理解了KV-Cache的原理时，深深的感受到了什么叫工程的美感，人类智慧的结晶。 262 | 263 | ### 扩展知识4：pytorch的动态长度推理怎么转换为需要静态张量形状的ONNX推理格式？ 264 | 265 | 我们2023年部署端侧大语言模型参展时另外遇到的一个问题就是动态转静态的问题。我们当时的安卓部署平台仅仅支持老版本的ONNX(一种工业界AI模型的存储格式，包含了模型结构与权重，可以直接拿去运行推理)，不支持动态轴，因为我们项目周期只有不足4个月，我们没有时间和人力去修改部署平台底层，因此我想到了一个取巧的办法，通过一定的冗余来将动态推理转换为静态。 266 | 267 | 首先ONNX的静态推理意味着，模型的所有计算都必须有固定的张量形状(tensor shape)，可在大语言模型的推理中token数的N是变化的呀，因此我只能将token数的维度固定为2048，永远计算2048个token。好在有KV-Cache，计算复杂度只是线性提升，不是平方提升。然后实际推理第k个token的时候 $(k < N)$ ，把k+1到N个token的注意力部分给mask掉，也就是赋予-inf的值，让其softmax后为0，防止其参与计算。 268 | 269 | 这个部署策略我们后来也看到其他公司和其他人使用过，但我觉得这毕竟是无奈之举，会有计算浪费，还希望各个公司早日把推理库的底层改成支持动态轴(有一些张量维度的长度可以变)，2025年了，别再拖着了。 270 | 271 | ------------- 272 | 273 | 注意力模块相关的内容确实比较多，我决定把注意力模型的优化单开一章节，尤其是着重讲一下最近特别火的DeepSeek开源模块对注意力模块的修改。 274 | 275 | [\[主目录链接\]](https://github.com/KaihuaTang/All-you-need-to-know-about-LLM#章节链接) 276 | 277 | 278 | 279 | ## 引用链接 280 | 281 | ``` 282 | @misc{tang2025all, 283 | title = {Building a Small LLM from Scratch: a tutorial}, 284 | author = {Tang, Kaihua and Zhang, Huaizheng}, 285 | year = {2025}, 286 | note = {\url{https://github.com/KaihuaTang/Building-a-Small-LLM-from-Scratch}}, 287 | } 288 | ``` -------------------------------------------------------------------------------- /chapter2/mh_causal_selfattention.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import math 4 | 5 | class MultiHeadCausalSelfAttention(nn.Module): 6 | def __init__(self, input_dim, hidden_dim, num_heads=8): 7 | """ 8 | 简易自注意力层 9 | :param input_dim: 输入维度 10 | :param hidden_dim: 隐藏层维度 11 | :param num_heads: 多头注意力的头数，这里可以指定为8表示单头注意力 12 | """ 13 | super(MultiHeadCausalSelfAttention, self).__init__() 14 | self.input_dim = input_dim 15 | self.hidden_dim = hidden_dim 16 | self.num_heads = num_heads 17 | # 将嵌入维度平分到各个头上 18 | # 注意实际应用中需要确保 embed_dim % num_heads == 0 19 | self.head_dim = hidden_dim // num_heads 20 | 21 | # 定义线性层，用于生成 Q、K、V 22 | self.query = nn.Linear(input_dim, hidden_dim) 23 | self.key = nn.Linear(input_dim, hidden_dim) 24 | self.value = nn.Linear(input_dim, hidden_dim) 25 | 26 | # 输出变换 27 | self.out = nn.Linear(hidden_dim, input_dim) 28 | 29 | def forward(self, x, mask=None): 30 | """ 31 | :param x: 输入张量，形状 [batch_size, seq_len, input_dim] 32 | :param mask: 可选的掩码（mask），形状与注意力矩阵匹配，如 [batch_size, 1, seq_len, seq_len] 33 | :return: 自注意力计算后的输出，形状 [batch_size, seq_len, input_dim] 34 | """ 35 | batch_size, seq_len, input_dim = x.shape 36 | 37 | # 线性变换得到 Q, K, V 38 | # 形状: [batch_size, seq_len, input_dim] 39 | Q = self.query(x) 40 | K = self.key(x) 41 | V = self.value(x) 42 | 43 | # 多头展开 44 | # 变换后形状: [batch_size, seq_len, num_heads, head_dim] 45 | Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim) 46 | K = K.view(batch_size, seq_len, self.num_heads, self.head_dim) 47 | V = V.view(batch_size, seq_len, self.num_heads, self.head_dim) 48 | 49 | # 将 [batch_size, seq_len, num_heads, head_dim] 转成 [batch_size, num_heads, seq_len, head_dim] 50 | Q = Q.permute(0, 2, 1, 3) # [batch_size, num_heads, seq_len, head_dim] 51 | K = K.permute(0, 2, 1, 3) 52 | V = V.permute(0, 2, 1, 3) 53 | 54 | # 计算注意力分数: Q @ K^T / sqrt(head_dim) 55 | # Q: [batch_size, num_heads, seq_len, head_dim] 56 | # K^T: [batch_size, num_heads, head_dim, seq_len] 57 | # scores: [batch_size, num_heads, seq_len, seq_len] 58 | scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim) 59 | 60 | # 如果有 mask，则在计算分数时将被mask的部分赋予一个很大的负数，以避免注意力 61 | # mask 形状一般为 [batch_size, 1, seq_len, seq_len] 或 [batch_size, num_heads, seq_len, seq_len] 62 | if mask is not None: 63 | scores = scores + mask 64 | 65 | # 通过 softmax 得到注意力分布 66 | attn_weights = torch.softmax(scores, dim=-1) 67 | print("注意力权重分布:", attn_weights) 68 | 69 | # 注意力加权 V 70 | # [batch_size, num_heads, seq_len, seq_len] x [batch_size, num_heads, seq_len, head_dim] 71 | # => [batch_size, num_heads, seq_len, head_dim] 72 | attn_output = torch.matmul(attn_weights, V) 73 | 74 | # 把多头重新拼接回原始形状 75 | # [batch_size, num_heads, seq_len, head_dim] -> [batch_size, seq_len, num_heads, head_dim] 76 | attn_output = attn_output.permute(0, 2, 1, 3).contiguous() 77 | # 拼接头部维度 78 | # => [batch_size, seq_len, input_dim] 79 | attn_output = attn_output.view(batch_size, seq_len, input_dim) 80 | 81 | # 输出层 82 | output = self.out(attn_output) 83 | 84 | return output 85 | 86 | if __name__ == "__main__": 87 | # 测试代码 88 | batch_size = 1 89 | seq_len = 5 90 | input_dim = 256 91 | hidden_dim = 256 92 | num_heads = 2 93 | x = torch.randn(batch_size, seq_len, input_dim) 94 | causal_mask = torch.triu(torch.ones(5, 5, device=x.device, dtype=torch.float), 95 | diagonal=1 96 | ) 97 | causal_mask = causal_mask.masked_fill(causal_mask == 1, float('-inf')) 98 | 99 | attention = MultiHeadCausalSelfAttention(input_dim, hidden_dim, num_heads=num_heads) 100 | out = attention(x, mask=causal_mask) 101 | print("输入形状:", x.shape) 102 | print("输出形状:", out.shape) -------------------------------------------------------------------------------- /chapter2/selfattention.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import math 4 | 5 | class SelfAttention(nn.Module): 6 | def __init__(self, input_dim, hidden_dim): 7 | """ 8 | 简易自注意力层 9 | :param input_dim: 输入维度 10 | :param hidden_dim: 隐藏层维度 11 | """ 12 | super(SelfAttention, self).__init__() 13 | 14 | # 定义线性层，用于生成 Q、K、V 15 | self.query = nn.Linear(input_dim, hidden_dim) 16 | self.key = nn.Linear(input_dim, hidden_dim) 17 | self.value = nn.Linear(input_dim, hidden_dim) 18 | 19 | # 输出变换 20 | self.out = nn.Linear(hidden_dim, input_dim) 21 | 22 | def forward(self, x): 23 | """ 24 | :param x: 输入张量，形状 [batch_size, seq_len, hidden_dim] 25 | :output: 自注意力计算后的输出，形状 [batch_size, seq_len, hidden_dim] 26 | """ 27 | batch_size, seq_len, hidden_dim = x.shape 28 | 29 | # 线性变换得到 Q, K, V 30 | # 形状: [batch_size, seq_len, hidden_dim] 31 | Q = self.query(x) 32 | K = self.key(x) 33 | V = self.value(x) 34 | 35 | # 计算注意力分数: Q @ K^T / sqrt(hidden_dim) 36 | # Q: [batch_size, seq_len, hidden_dim] 37 | # K^T: [batch_size, hidden_dim, seq_len] 38 | # scores: [batch_size, seq_len, seq_len] 39 | scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(hidden_dim) 40 | 41 | # 通过 softmax 得到注意力分布 42 | attn_weights = torch.softmax(scores, dim=-1) 43 | 44 | # 注意力加权 V 45 | # [batch_size, seq_len, seq_len] x [batch_size, seq_len, hidden_dim] 46 | # => [batch_size, seq_len, hidden_dim] 47 | attn_output = torch.matmul(attn_weights, V) 48 | 49 | # 输出层 50 | output = self.out(attn_output) 51 | 52 | return output 53 | 54 | if __name__ == "__main__": 55 | # 测试代码 56 | batch_size = 2 57 | seq_len = 5 58 | input_dim = 256 59 | hidden_dim = 256 60 | x = torch.randn(batch_size, seq_len, input_dim) 61 | 62 | attention = SelfAttention(input_dim, hidden_dim) 63 | out = attention(x) 64 | print("输入形状:", x.shape) 65 | print("输出形状:", out.shape) -------------------------------------------------------------------------------- /chapter3/03-1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KaihuaTang/Building-a-Small-LLM-from-Scratch/b63c6353fa87896a446dc6ec08f792f14aebd72b/chapter3/03-1.png -------------------------------------------------------------------------------- /chapter3/03-2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KaihuaTang/Building-a-Small-LLM-from-Scratch/b63c6353fa87896a446dc6ec08f792f14aebd72b/chapter3/03-2.png -------------------------------------------------------------------------------- /chapter3/03-3.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/KaihuaTang/Building-a-Small-LLM-from-Scratch/b63c6353fa87896a446dc6ec08f792f14aebd72b/chapter3/03-3.png -------------------------------------------------------------------------------- /chapter3/README.md: -------------------------------------------------------------------------------- 1 | 2 | # 03 DeepSeekV3的注意力优化 3 | 4 | 申明：本教程的所有内容(文字，图片，代码等)可以用于非盈利目的个人使用和分享。但如果用于盈利目的，包括但不限于卖课，公众号，视频号等需要经由作者的批准。谢谢理解。[\[知乎链接\]](https://zhuanlan.zhihu.com/p/19275166926) 5 | 6 | [\[主目录链接\]](https://github.com/KaihuaTang/All-you-need-to-know-about-LLM#章节链接) 7 | 8 | 9 | ## 前言 10 | 书接上文[注意力模块与KV Cache](chapter2/README.md)，介绍完了基本的大语言模型的注意力模块与相关的KV-Cache，我们本章着重展开讲讲主流开源大模型注意力模块的后续改良与优化。本章我们不仅会介绍最近当红的[DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3)注意力优化原理，更会直接深扒一下Qwen2/LLaMA3和DeepSeekV3具体的注意力模块的代码，深入讲解每一行代码对应的功能和原理。(DeepSeek太火了，被拉回国项目攻关了，更新耽误了两周) 11 | 12 | ## 一. 注意力优化 13 | 主流的开源大模型网络的注意力计算机制除了上一章介绍的多头注意力Multi-Head Attention(MHA)以外，最近也有了新的变种，主要包括Multi-Query Attention (MQA)，Grouped-Query Attention (GQA)和最近当红的DeepSeek的Multi-head Latent Attention (MLA)而他们优化的方向其实是一致的，就是极致的压缩KV的大小，因为这样KV-Cache可以加载的更快。毕竟现在都在说超长上下文，token数N长了KV-Cache优化后加载KV的传输带宽开销可也不小啊。 14 | 15 |
16 | logo 17 | 图1：各种注意力优化方案。 18 |
19 | 20 | ### 1. Multi-Query Attention (MQA) / 多查询注意力 21 | 22 | 参考上一章的标准多头注意力Multi-Head Attention(MHA)的代码，其QKV构成如下： 23 | ``` 24 | # x为attention模块输入的所有token的hidden states 25 | batch_size, seq_len, input_dim = x.shape 26 | 27 | # 线性变换得到 Q, K, V 28 | Q = self.query(x) 29 | K = self.key(x) 30 | V = self.value(x) 31 | 32 | # 多头展开 33 | # 变换后形状: [batch_size, seq_len, num_heads, head_dim] 34 | Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim) 35 | K = K.view(batch_size, seq_len, self.num_heads, self.head_dim) 36 | V = V.view(batch_size, seq_len, self.num_heads, self.head_dim) 37 | 38 | # 此处省略后续attention计算 39 | ``` 40 | 其中K和V对应的线性层self.key和self.value都必须将token的特征从input_dim维度映射到num_heads * head_dim维度。也就是说这两层的线性层权重的张量形状为[num_heads * head_dim, input_dim]。而同时KV Cache也必须存储两个[batch_size, seq_len, num_heads, head_dim]大小的张量。 41 | 42 | 而多查询注意力Multi-Query Attention (MQA)做的优化直白来来讲就是只对Query做多头注意力，而Key和Query只生成单头。此时self.key和self.value线性层权重的张量形状为[head_dim, input_dim]，这不仅让线性层计算量缩小num_head倍，KV Cache的大小也缩小了同样的倍数，计算量和带宽双收益。当然天下没有免费的午餐，由于Key和Value仅有一个头，其表达能力肯定是有所损失的。具体MQA参考代码如下： 43 | 44 | ``` 45 | # x为attention模块输入的所有token的hidden states 46 | batch_size, seq_len, input_dim = x.shape 47 | 48 | # 线性变换得到 Q, K, V 49 | Q = self.query(x) 50 | K = self.key(x) # 注意此处的key线性层输出维度小于MHA 51 | V = self.value(x) # 注意此处的key线性层输出维度小于MHA 52 | 53 | # 多头展开 54 | Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim) 55 | K = K.view(batch_size, seq_len, 1, self.head_dim) 56 | V = V.view(batch_size, seq_len, 1, self.head_dim) 57 | 58 | # 将K和V复制若干份，扩展至: [batch_size, seq_len, num_heads, head_dim] 59 | K = K.repeat(1, 1, self.num_heads, 1) 60 | V = V.repeat(1, 1, self.num_heads, 1) 61 | 62 | # 此处省略后续attention计算 63 | ``` 64 | 65 | 66 | ### 2. Grouped-Query Attention (GQA) / 组查询注意力 67 | 68 | 正如上文说的，MQA必然带来注意力层表达能力的下降。因此就有了组查询注意力Grouped-Query Attention (GQA)，这其实是MHA和MQA的折中，就是Query保留所有head数的情况下，Key和Value不止一个头，而是保留num_group数的头(num_groups <= num_heads且num_heads可以被num_groups整除)，我们不难发现GQA是MHA和MQA的一种泛化形式，num_groups=1时就是MQA，num_groups=num_heads时就是MHA。可以说是万金油的表达形式。因为GQA是更泛化的表达形式，同时也有个额外的参数num_groups（有的代码中也叫num_key_value_heads）可以调，因此GQA往往可以调到与MHA差不多的性能，同时又能有KV Cache和线性层计算减少的收益。在主流的Qwen2和LLaMA3的代码中，一般也都支持GQA的配置。 69 | 70 | 具体GQA参考代码如下： 71 | ``` 72 | # x为attention模块输入的所有token的hidden states 73 | batch_size, seq_len, input_dim = x.shape 74 | 75 | # 线性变换得到 Q, K, V 76 | Q = self.query(x) 77 | K = self.key(x) 78 | V = self.value(x) 79 | 80 | # 多头展开 81 | # num_groups <= num_heads且num_heads可以被num_groups整除 82 | Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim) 83 | K = K.view(batch_size, seq_len, 1, self.num_groups, self.head_dim) 84 | V = V.view(batch_size, seq_len, 1, self.num_groups, self.head_dim) 85 | 86 | # 将K和V复制若干份，扩展至: [batch_size, seq_len, num_heads, head_dim] 87 | K = K.repeat(1, 1, self.num_heads // self.num_groups, 1, 1).view(batch_size, seq_len, self.num_heads, self.head_dim) 88 | V = V.repeat(1, 1, self.num_heads // self.num_groups, 1, 1).view(batch_size, seq_len, self.num_heads, self.head_dim) 89 | 90 | # 此处省略后续attention计算 91 | ``` 92 | 93 | ### 3. Multi-head Latent Attention (MLA) / 多头潜在注意力 94 | 95 | 最近大火的[DeepSeek系列(从V2到V3)](https://github.com/deepseek-ai/DeepSeek-V3/tree/main)则采用了一种比GQA更极致的压缩方式，不仅进一步减少了注意力层中线性层的理论算力，更把KV Cache压缩到了新的境界。在介绍MLA之前，先介绍一个低秩分解在大模型上应用的例子，后面我们可能也会单独详细讲讲，就是[LoRA: Low-Rank Adaptation](https://arxiv.org/abs/2106.09685)。这是大模型常用的一种高性能微调方法，其具体概念就是，如果线性层的权重$W$太大(假设其张量形状为[out_dim, in_dim])，训练这个权重太耗显存了，我们可以训练两个更小的权重$W_a$和$W_b$(形状分别为[K, in_dim]和[out_dim, K]，K << in_dim, K << out_dim)。由于K远远小于in_dim和out_dim，这两个权重加起来也远远小于原始的$W$。参考概念如下图2。 96 | 97 |
98 | logo 99 | 图2：LoRA微调概念图。 100 |
101 | 102 | 当我第一次看到DeepSeek的多头潜在注意力Multi-head Latent Attention (MLA)，我首先映入脑袋的便是LoRA，区别是在MLA中并不是额外学两个小的线性权重，而是用直接用两个小的线性权重取代一个完整的线性层。具体MLA的网络结构如下图3（注意MLA在DeepSeek中有两种形式，一种是只有KV的线性层运用了低秩分解，一种是Q和KV都利用了低秩分解，后者也是最近大火的[DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3/tree/main)和[DeepSeek-R1](https://github.com/deepseek-ai/DeepSeek-R1/tree/main)的网络结构，因此我们这里以后者为例）。 103 | 104 |
105 | logo 106 | 图3：多头潜在注意力结构，参数规模参考DeepSeek-V3与R1。 107 |
108 | 109 | 上图是MLA如何生成Query，Key和Value的流程图。MLA结构中Query和Key都分为nope部分和rope部分，前者是指No Position Embedding即无需位置编码，后者指Rotary Position Embedding即需要旋转位置编码，而旋转位置编码则是图上apply rotary pos embed模块，该模块主要为了给token提供其在序列中与其他token的相对位置信息，下一章我们会展开详细讲解。 110 | 111 | 我们可以发现，原始的注意力仅需简单的三个Linear加上Reshape就可以生成Query，Key和Value，但MLA似乎让网络变得更复杂了。这是为什么呢？因为MLA利用了低秩分解的概念，将一个大的矩阵乘拆解为两个小的矩阵乘加一个归一化层。我们可以通过参数量估算发现，如果用三个线性层直接生成同样的Query, Key, Value维度，参数量需要扩大6.7倍。 112 | 113 | 除了参数量的降低，MLA更可以大幅降低KV Cache，如上图DeepSeek-V3与R1网络中，其每个token仅需保留64维的k_pe + 右侧RMSNorm后的512维特征，总计576维每token的Cache即可。因为其余的k_nope和Value都可以直接通过512维的特征再经过一个线性层得到。而此前其他大语言模型每个token需要多少维度的特征呢？以[72B的Qwen2.5](https://huggingface.co/Qwen/Qwen2.5-72B-Instruct/tree/main)模型为例，即便已经使用了组查询注意力GQA，每个token依然需要2048维(128 x 8 x 2)。DeepSeek的MLA将KV Cache压缩了3.5倍，当然这里存的已经不是标准的Key和Value了，需要引入额外的线性层计算才能转换为Key和Value，这就涉及到更复杂的算力和带宽的权衡了。 114 | 115 | 116 | ## 二. 大语言模型中的注意力 117 | 下面，让我们切切实实的看一看真实的大模型网络结构里注意力都长什么样。下面我会拿Qwen2/LLaMA3（这两个网络的注意力部分非常相似，因此我仅展示一个）与DeepSeekV3的实际代码进行演示。我会尽可能保留原始代码，仅出于可读性做一些修改，然后通过详细的注释来阐释每一块的作用。 118 | 119 | ### 1. Qwen2/LLaMA3的注意力代码详解 120 | 121 | 以transformers库v4.49.0版的代码为例，我们看一下Qwen2注意力模块的实现（LLaMA与他也是大同小异）。完整代码参考链接：[https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/models/qwen2/modeling_qwen2.py](https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/models/qwen2/modeling_qwen2.py) 122 | 123 | ``` 124 | class Qwen2Attention(nn.Module): 125 | """Multi-headed attention from 'Attention Is All You Need' paper""" 126 | 127 | def __init__(self, config: Qwen2Config, layer_idx: int): 128 | super().__init__() 129 | self.config = config # 获取配置参数 130 | self.layer_idx = layer_idx # 获取当前是第几层网络 131 | # 获取注意力头每个头的维度，如果不直接提供head_dim参数则通过hidden_size除以注意力头数获得。 132 | self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) 133 | # 这里的GQA的num_key_value_groups是要复制几遍KV的意思，我上文说的group则是kv本身有几个头。 134 | self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads 135 | 136 | self.scaling = self.head_dim**-0.5 # attention计算时的缩放参数 137 | self.attention_dropout = config.attention_dropout # attention是否开启dropout 138 | self.is_causal = True # 是否是因果注意力 139 | # 一些线性层的初始化 140 | self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=True) 141 | self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True) 142 | self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True) 143 | self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) 144 | 145 | def forward( 146 | self, 147 | hidden_states: torch.Tensor, 148 | position_embeddings: Tuple[torch.Tensor, torch.Tensor], 149 | attention_mask: Optional[torch.Tensor], 150 | past_key_value: Optional[Cache] = None, 151 | cache_position: Optional[torch.LongTensor] = None, 152 | **kwargs: Unpack[FlashAttentionKwargs], 153 | ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: 154 | input_shape = hidden_states.shape[:-1] 155 | hidden_shape = (*input_shape, -1, self.head_dim) 156 | # 标准的通过线性层生成query, key和value。 hidden_shape将头数的维度设为动态的。 157 | query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) 158 | key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) 159 | value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) 160 | 161 | # 下面两行为位置编码相关，后续章节会详细讲 162 | cos, sin = position_embeddings 163 | query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) 164 | 165 | # 下面部分为更新kv cache 166 | if past_key_value is not None: 167 | # sin and cos are specific to RoPE models; cache_position needed for the static cache 168 | cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} 169 | key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) 170 | 171 | # sliding window是一种特殊的注意力稀疏机制，可以减少计算量增加推理长度，比较进阶。可以忽略下面sliding window相关。 172 | sliding_window = None 173 | if ( 174 | self.config.use_sliding_window 175 | and getattr(self.config, "sliding_window", None) is not None 176 | and self.layer_idx >= self.config.max_window_layers 177 | ): 178 | sliding_window = self.config.sliding_window 179 | 180 | # attention的计算有很多优化库，比如flash attention等。attention_interface决定了调用那种实现。 181 | # 为了便于讲解，我们默认attention_interface = eager_attention_forward，也就是下面的pytorch实现。 182 | attention_interface: Callable = eager_attention_forward 183 | if self.config._attn_implementation != "eager": 184 | if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False): 185 | logger.warning_once( 186 | "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to " 187 | 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' 188 | ) 189 | else: 190 | attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] 191 | 192 | # 调用eager_attention_forward，计算attention 193 | attn_output, attn_weights = attention_interface( 194 | self, 195 | query_states, 196 | key_states, 197 | value_states, 198 | attention_mask, 199 | dropout=0.0 if not self.training else self.attention_dropout, 200 | scaling=self.scaling, 201 | sliding_window=sliding_window, # main diff with Llama 202 | **kwargs, 203 | ) 204 | 205 | # 通过reshape和输出线性层o_proj, 得到注意力层的最后输出结果。 206 | attn_output = attn_output.reshape(*input_shape, -1).contiguous() 207 | attn_output = self.o_proj(attn_output) 208 | return attn_output, attn_weights 209 | 210 | 211 | # 如果涉及上述说到的MQA和GQA，就需要这里的repeat_kv函数，将key和value复制多份，直到与query的头数相等。 212 | def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: 213 | """ 214 | This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, 215 | num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) 216 | """ 217 | batch, num_key_value_heads, slen, head_dim = hidden_states.shape 218 | if n_rep == 1: 219 | return hidden_states 220 | hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) 221 | return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) 222 | 223 | 224 | def eager_attention_forward( 225 | module: nn.Module, 226 | query: torch.Tensor, 227 | key: torch.Tensor, 228 | value: torch.Tensor, 229 | attention_mask: Optional[torch.Tensor], 230 | scaling: float, 231 | dropout: float = 0.0, 232 | **kwargs, 233 | ): 234 | # 在MQA和GQA中，将key和value复制多份，直到与query的头数相等。 235 | key_states = repeat_kv(key, module.num_key_value_groups) 236 | value_states = repeat_kv(value, module.num_key_value_groups) 237 | 238 | # 下面就是标准的多头因果注意力，和我们上一章讲解的大同小异。 239 | attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling 240 | if attention_mask is not None: 241 | causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] 242 | attn_weights = attn_weights + causal_mask 243 | 244 | attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) 245 | # attention的droupout实现，会丢掉一些attention map中的值。 246 | attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) 247 | attn_output = torch.matmul(attn_weights, value_states) 248 | attn_output = attn_output.transpose(1, 2).contiguous() 249 | 250 | return attn_output, attn_weights 251 | ``` 252 | 253 | 254 | 255 | ### 2. DeepSeekV3的注意力代码详解 256 | 257 | DeepSeek-V3的网络结构在写本文时还没有集成如transformers库，所以我们参考的是他官方checkpoint中的网络结构源码：[https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/main/modeling_deepseek.py](https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/main/modeling_deepseek.py) 258 | 259 | ``` 260 | class DeepseekV3Attention(nn.Module): 261 | """Multi-headed attention from 'Attention Is All You Need' paper""" 262 | 263 | def __init__(self, config: DeepseekV3Config, layer_idx: Optional[int] = None): 264 | super().__init__() 265 | self.config = config # 获取配置参数 266 | self.layer_idx = layer_idx # 获取当前第几层 267 | if layer_idx is None: 268 | logger.warning_once( 269 | f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " 270 | "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` " 271 | "when creating this class." 272 | ) 273 | 274 | # 提取配置文件中的各种参数，与Qwen的大同小异。 275 | self.attention_dropout = config.attention_dropout 276 | self.hidden_size = config.hidden_size 277 | self.num_heads = config.num_attention_heads 278 | 279 | self.max_position_embeddings = config.max_position_embeddings # 最大位置编码 280 | self.rope_theta = config.rope_theta # 位置编码相关参数 281 | 282 | # 下面的参数就是和Qwen/LLaMA不同的部分了，因为生成QKV的Linear被拆成了两个小Linear + 归一化层的形式，所以需要额外的中间维度参数。 283 | self.q_lora_rank = config.q_lora_rank 284 | self.qk_rope_head_dim = config.qk_rope_head_dim 285 | self.kv_lora_rank = config.kv_lora_rank 286 | self.v_head_dim = config.v_head_dim 287 | self.qk_nope_head_dim = config.qk_nope_head_dim 288 | self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim 289 | 290 | self.is_causal = True 291 | 292 | # 如果不设置q_lora_rank参数，则query还是通过一个线性层直接生成。 293 | if self.q_lora_rank is None: 294 | self.q_proj = nn.Linear( 295 | self.hidden_size, self.num_heads * self.q_head_dim, bias=False 296 | ) 297 | # DeepSeek-V3与DeepSeek-R1设置了q_lora_rank参数，因此变成两个Linear + 一个RMSNorm的形式，节省了计算量 298 | else: 299 | self.q_a_proj = nn.Linear( 300 | self.hidden_size, config.q_lora_rank, bias=config.attention_bias 301 | ) 302 | self.q_a_layernorm = DeepseekV3RMSNorm(config.q_lora_rank) 303 | self.q_b_proj = nn.Linear( 304 | config.q_lora_rank, self.num_heads * self.q_head_dim, bias=False 305 | ) 306 | 307 | # Key和Value被作为一个整体通过两个Linear + 一个RMSNorm后再split成key和Value 308 | self.kv_a_proj_with_mqa = nn.Linear( 309 | self.hidden_size, 310 | config.kv_lora_rank + config.qk_rope_head_dim, 311 | bias=config.attention_bias, 312 | ) 313 | self.kv_a_layernorm = DeepseekV3RMSNorm(config.kv_lora_rank) 314 | self.kv_b_proj = nn.Linear( 315 | config.kv_lora_rank, 316 | self.num_heads 317 | * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), 318 | bias=False, 319 | ) 320 | 321 | # 正常的输出线性层 322 | self.o_proj = nn.Linear( 323 | self.num_heads * self.v_head_dim, 324 | self.hidden_size, 325 | bias=config.attention_bias, 326 | ) 327 | self._init_rope() 328 | 329 | self.softmax_scale = self.q_head_dim ** (-0.5) 330 | 331 | # 位置编码相关，暂时跳过。 332 | if self.config.rope_scaling is not None: 333 | mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0) 334 | scaling_factor = self.config.rope_scaling["factor"] 335 | if mscale_all_dim: 336 | mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) 337 | self.softmax_scale = self.softmax_scale * mscale * mscale 338 | 339 | # 初始化位置编码，包含多种位置编码实现，此处跳过。 340 | def _init_rope(self): 341 | if self.config.rope_scaling is None: 342 | self.rotary_emb = DeepseekV3RotaryEmbedding( 343 | self.qk_rope_head_dim, 344 | max_position_embeddings=self.max_position_embeddings, 345 | base=self.rope_theta, 346 | ) 347 | else: 348 | scaling_type = self.config.rope_scaling["type"] 349 | scaling_factor = self.config.rope_scaling["factor"] 350 | if scaling_type == "linear": 351 | self.rotary_emb = DeepseekV3LinearScalingRotaryEmbedding( 352 | self.qk_rope_head_dim, 353 | max_position_embeddings=self.max_position_embeddings, 354 | scaling_factor=scaling_factor, 355 | base=self.rope_theta, 356 | ) 357 | elif scaling_type == "dynamic": 358 | self.rotary_emb = DeepseekV3DynamicNTKScalingRotaryEmbedding( 359 | self.qk_rope_head_dim, 360 | max_position_embeddings=self.max_position_embeddings, 361 | scaling_factor=scaling_factor, 362 | base=self.rope_theta, 363 | ) 364 | elif scaling_type == "yarn": 365 | kwargs = { 366 | key: self.config.rope_scaling[key] 367 | for key in [ 368 | "original_max_position_embeddings", 369 | "beta_fast", 370 | "beta_slow", 371 | "mscale", 372 | "mscale_all_dim", 373 | ] 374 | if key in self.config.rope_scaling 375 | } 376 | self.rotary_emb = DeepseekV3YarnRotaryEmbedding( 377 | self.qk_rope_head_dim, 378 | max_position_embeddings=self.max_position_embeddings, 379 | scaling_factor=scaling_factor, 380 | base=self.rope_theta, 381 | **kwargs, 382 | ) 383 | else: 384 | raise ValueError(f"Unknown RoPE scaling type {scaling_type}") 385 | 386 | # reshape函数 387 | def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): 388 | return ( 389 | tensor.view(bsz, seq_len, self.num_heads, self.v_head_dim) 390 | .transpose(1, 2) 391 | .contiguous() 392 | ) 393 | 394 | def forward( 395 | self, 396 | hidden_states: torch.Tensor, 397 | attention_mask: Optional[torch.Tensor] = None, 398 | position_ids: Optional[torch.LongTensor] = None, 399 | past_key_value: Optional[Cache] = None, 400 | output_attentions: bool = False, 401 | use_cache: bool = False, 402 | **kwargs, 403 | ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: 404 | if "padding_mask" in kwargs: 405 | warnings.warn( 406 | "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" 407 | ) 408 | bsz, q_len, _ = hidden_states.size() 409 | 410 | # 通过标准的Linear或我们说的LoRA形式生成query 411 | if self.q_lora_rank is None: 412 | q = self.q_proj(hidden_states) 413 | else: 414 | q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) 415 | q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2) 416 | 417 | # 将query每个头的特征维度分成需要位置编码的部分和不需要的部分 418 | q_nope, q_pe = torch.split( 419 | q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 420 | ) 421 | 422 | # 先通过一个线性层，生成压缩后的kv特征：512+64。 423 | compressed_kv = self.kv_a_proj_with_mqa(hidden_states) 424 | # 将其中key的需要位置编码部分的64维单独拿出来 425 | compressed_kv, k_pe = torch.split( 426 | compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 427 | ) 428 | k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2) 429 | 430 | # 生成完整的key和value 431 | kv = ( 432 | self.kv_b_proj(self.kv_a_layernorm(compressed_kv)) 433 | .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) 434 | .transpose(1, 2) 435 | ) 436 | 437 | # 将key不需要位置编码的部分，和value部分分离 438 | k_nope, value_states = torch.split( 439 | kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1 440 | ) 441 | kv_seq_len = value_states.shape[-2] 442 | # 通过kvcache状态判断key和value的长度 443 | if past_key_value is not None: 444 | if self.layer_idx is None: 445 | raise ValueError( 446 | f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " 447 | "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " 448 | "with a layer index." 449 | ) 450 | kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) 451 | 452 | # query和key的旋转位置编码部分 453 | cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) 454 | q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) 455 | 456 | # 分别融合query和key的 无位置编码nope和选择位置编码pe 457 | query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) 458 | query_states[:, :, :, : self.qk_nope_head_dim] = q_nope 459 | query_states[:, :, :, self.qk_nope_head_dim :] = q_pe 460 | 461 | key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) 462 | key_states[:, :, :, : self.qk_nope_head_dim] = k_nope 463 | key_states[:, :, :, self.qk_nope_head_dim :] = k_pe 464 | 465 | # 更新kv cache 466 | if past_key_value is not None: 467 | cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models 468 | key_states, value_states = past_key_value.update( 469 | key_states, value_states, self.layer_idx, cache_kwargs 470 | ) 471 | 472 | # 后续都是正常的attention计算。 473 | # mla比较特别的就是query，key和value的维度数不一定一致。 474 | # 因为只要query和key对齐就可以进行attention了，value其实不需要对齐每个头的维度，仅需要对齐长度即可。 475 | attn_weights = ( 476 | torch.matmul(query_states, key_states.transpose(2, 3)) * self.softmax_scale 477 | ) 478 | 479 | if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): 480 | raise ValueError( 481 | f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" 482 | f" {attn_weights.size()}" 483 | ) 484 | assert attention_mask is not None 485 | if attention_mask is not None: 486 | if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): 487 | raise ValueError( 488 | f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" 489 | ) 490 | attn_weights = attn_weights + attention_mask 491 | 492 | # upcast attention to fp32 493 | attn_weights = nn.functional.softmax( 494 | attn_weights, dim=-1, dtype=torch.float32 495 | ).to(query_states.dtype) 496 | attn_weights = nn.functional.dropout( 497 | attn_weights, p=self.attention_dropout, training=self.training 498 | ) 499 | attn_output = torch.matmul(attn_weights, value_states) 500 | 501 | if attn_output.size() != (bsz, self.num_heads, q_len, self.v_head_dim): 502 | raise ValueError( 503 | f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.v_head_dim)}, but is" 504 | f" {attn_output.size()}" 505 | ) 506 | 507 | attn_output = attn_output.transpose(1, 2).contiguous() 508 | 509 | attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim) 510 | 511 | attn_output = self.o_proj(attn_output) 512 | 513 | if not output_attentions: 514 | attn_weights = None 515 | 516 | return attn_output, attn_weights, past_key_value 517 | ``` 518 | 519 | 520 | ------------- 521 | 522 | [\[主目录链接\]](https://github.com/KaihuaTang/All-you-need-to-know-about-LLM#章节链接) 523 | 524 | 525 | 526 | ## 引用链接 527 | 528 | ``` 529 | @misc{tang2025all, 530 | title = {Building a Small LLM from Scratch: a tutorial}, 531 | author = {Tang, Kaihua and Zhang, Huaizheng}, 532 | year = {2025}, 533 | note = {\url{https://github.com/KaihuaTang/Building-a-Small-LLM-from-Scratch}}, 534 | } 535 | ``` 536 | --------------------------------------------------------------------------------