为了帮助你更好地理解,我们先来看看大模型是怎么生成文本的。

大模型如何生成文本?

我们可以先来用黑盒思维来看看这个问题。把大模型当作是一个黑盒,它可以根据我们输入的一段话,不断地进行自我补全,直到输出一个结束标记。
比如说,下面这句话:
我很熟悉的一个人是小明,他是我最
如果交给人脑,我们很容易就能推断出后面要续写的应该是“好的朋友”,或者是“亲密的伙伴”这样的词,如果问为什么是这些词,你可能会回答,这是我们根据前面已经有的信息去总结归纳得出来的。大模型干的也是同一件事,而这种”根据前面的信息去进行补全“的能力,在大模型里面通过Attention机制去实现。
这种能力,可以理解为一种记忆能力,说点题外话,如何获得这种能力,长久以来一直是人们很头疼的一个问题。你可以想象,如果你输入的不是短短的一句话,而是一部推理小说的前半部分,最后要补全的句子是”凶手是“,这让人类来补全,都得往回再翻几页,才能得出一个不太确定的答案。
以前的人们为了让文本分析模型获得这种能力,提出了很多种记忆框架,比如RNN,LSTM,但是它们都有不是这样,就是那样的问题,换句话来说,就是效果都不太好,这也是为什么之前的文本生成模型一直没有在工业界掀起什么风浪。
读到这里你一定好奇,那么注意力机制又有什么魔力呢?请听我娓娓道来:

注意力机制

嵌入向量

我们要逐渐把大模型的盒子拆开。在计算机的世界,万物皆数,在大模型的世界,可以说万物皆向量,所以我们输入的句子,会被拆开成一个个词元(即token),然后每个token会被转化成一个向量,这个过程,和这个向量本身,都被称为embedding(不妨称为嵌入向量)。
这么说来可能有点抽象,但你可以把这个嵌入向量的每个维度看作是这个词元在不同方面的程度的表示,比如说对于小明这个词,假设它对应的嵌入向量可能是[0.1, 0, -99, …],每个数的范围是从负无穷到正无穷,那么对于第一个维度,可以是形容这个词是一条狗的名字的程度,那么0.1就是可能性是有,但是实在是不大;
第二个维度,可以是形容这个词是智商高低的程度,小明这个名字比较中性,看不出来,所以是0;
第三个维度,可以是形容这个词是男孩子名字的程度,99表示还是挺有可能的…
如此下去,这个向量就是一个综合的,对于小明这个词的一个词性的判断。当然,实际上嵌入向量的每个维度,并不是有那么具象的含义,每个维度都是计算机训练出来的,对计算机有着抽象意义的含义。但我们可以知道,把一个词转成一个由一系列数字组成的序列的过程,就是便于计算机去进行下一步计算的,一个翻译的过程,而这个过程被称为嵌入。

生成出下一个词的进一步解释

了解了嵌入向量,让我们先回到:
我很熟悉的一个人是小明,他是我最
这句不完整的话。
假设我们要把这段话输入到大模型,那么第一步,就是先给这段话给切成一段段token,这个过程大模型会帮我们去做,我们假设切成下面的样子:
我/很/熟悉/的/一个 /人/是/小明 , 他/是/我/最
然后会把上面的一个个词元翻译成机器能看懂的一个个向量。然后,做什么呢?机器实际上会先盯着最后一个词“最”对应的向量。

计算机的提问方式

类比成人类,就会像我们在盯着这句话的最后一个字,自言自语,“最…最…”,然后在问自己:“最什么呢…?”,是啊,“最”什么呢?后面跟什么单词比较好呢?这是我们迫在眉睫的一个问题。
而机器实际上也会有这么一个“发问”的过程,只不过机器只会做计算,对于机器的提问,是把“最”这个词对应的嵌入向量,给乘与这个大模型里面已经给定了的,一个称为”询问权重“的一个矩阵Wq,与这个矩阵相乘的结果,就是把”最“这个字,给换成”最什么呢?“这样的一句疑问,而它也是一个向量。我们把这个向量称为Q,用公式写出来就是
Q = Wq X 嵌入向量

计算机找到答案的方式

”最什么呢?“如果是人类,我们发出这个疑问之后就会意识到一个问题,光盯着这个字看,即使提了个问,没有背景,我们也得不出什么答案啊,我们不能光盯着这个字看,我们要回想起前面的内容,然后看看有什么线索。
那么,计算机又是怎么做到”回想起前面的内容,然后看看有什么线索“的呢?
用人类的思维的话,那就是,我们也许可以看看能回答”最什么呢“这个问题的词有哪些,然后借用这些词的特性,去找出下一个答案。
那么对于计算机,那就是怎么去衡量所有词”回答这个问题的能力“呢?注意力框架给出的一个答案是,训练出一个称为”关键字权重“的一个矩阵,然后像上面把一个词的嵌入向量转化成一个疑问的过程一样,把一个嵌入向量转化成一个”关键字“向量,然后计算当前的问题,和每个词的关键字向量的相似度,相似度越大,就表示这个词语的特性越能回答这个问题。把这个关键字矩阵称为Wk,把生成出来的关键字向量称为K,用公式写出来就是
K = Wk X 嵌入向量
所以,带着“最”什么呢?这个问题代表的向量Q,我们对所有的词元(包括“最”这个字本身)进行回答能力的匹配,我们会得出一系列的匹配程度的数值,可能是这样的:
我-2/很-8/熟悉-88/的-0.1/一个-0.1 /人-4/是-0.1/小明-3.8 , 他-5.5/是-0.2/我-3/最-4
同样的,对于如何衡量“词语的特性”这个问题,我们也有相应的转化矩阵,称为值矩阵,把值矩阵称为Wv,把生成出来的“词语的特性”向量称为V,用公式写出来就是
V = Wv X 嵌入向量
注意,“词语的特性”向量叫值向量,它跟我们的嵌入向量,虽然从意义上看都是表达词语的意思,但是它们的用途不同,值向量后面会被用来推断对应的词元应该是什么,这个要区分开。

结合特性,随机挑选答案

我-2/很-8/熟悉-88/的-0.1/一个-0.1 /人-4/是-0.1/小明-3.8 , 他-5.5/是-0.2/我-3/最-4
我们得出了每个词元对于回答最后面应该填什么这个问题的能力,而且对于每个词元,我们还有另外的回答问题的特性的向量,那么我们应该做的,应该是根据能力大小,去决定怎么去把这些词元的特性去结合起来,从而找到问题的答案,那么我们可以把所有的能力值加起来,把这个总和去作为分母去除以上面各个词元的能力值,就会得出一系列新值,它们的总和为1,这样的把不同的值转化成为总和为1的过程,被叫做“正规化”。
上面提到的正规化,是其中一种方式。实际上,attention采取的不是上面提到的方式,是一种更温和的,被称为softmax的方式,这里不多介绍,总而言之,经过softmax的正规化,我们也能得出一系列总和为1的能力值。经过这样的正规化,我们就可以把正规化后的数字作为份额,去乘以各自的值向量并相加,就能得出一个新的值向量了!这个新的值向量被称作Z矩阵,蕴含了之前词语的上下文信息,会作为attention的最终输出,去输入到下一个结构,帮助我们去得出下一个词语的概率分布,也就是我们会得出可能是这样的结果:
34%的概率是“好”,32%的概率是“亲密”,12%的概率是“熟悉”…计算机做的就是根据对应的概率,在这些词语里面随机选择,这也就是为什么面对同一句话,计算机总是会随机输出不一样的答案——然后,周而复始,假设这里选择的是”好“,那么“好”这个词会被再输入大模型,直到大模型发出”生成完毕“的信号,程序终止。

多头注意(MHA)

我们还可以引入多个注意力头,为什么要多头?关键在于:多头不是把一套Q/K/V投影矩阵切开用,而是拥有多套独立的Q/K/V矩阵(原始Transformer用了 8 套),每套随机初始化、独立训练,把输入投影到不同的表示子空间——通俗地说,每个头都在用自己的一套视角看这句话。这带来一个直接的好处:模型能同时从多个角度关注不同位置的词

设想翻译 “The animal didn’t cross the street because it was too tired”,我们需要知道”it”指代的是”animal”还是”street”。单个头在做加权时,结果容易被词本身的特征主导,未必能抓住这种跨词的指代关系;而多个头各看一面,即使一个头没抓到,另一个头也可能锁定 “it->animal” 这样的联系。多视角让模型不被任何单一倾向(包括过度关注词自身)带偏。

从数学上说,多头注意其实就是把嵌入向量A转化成的Q向量平均切为num_heads份,然后每一份用不同的Wk,Wv去进行计算,最后把得出来的结果concat成一个矩阵,再用一个特殊训练的Z0矩阵,去重新转化成一个不再是切片组合,而是信息重新组合的Z矩阵。这样做还有利于并行计算!多好的一个机制。

GQA

如果把上面的多个注意力头分成N组,每组共用一个Wk、Wv矩阵,既能减少模型本身的权重大小,还能减小KV Cache的大小(在后文,我们会了解到在并发下,KV Cache所占用的内存会成为多并发的一个瓶颈)。

计算模型所占显存

说了这么多前置知识,说回来怎么计算模型会占用的显存大小吧。只要前面的知识都掌握牢固了,其实很简单。

计算权重所占显存-以Llama-3-70B为例

模型参数大小*精度代表字节数=权重所占显存

比如说,Llama-3-70B,fp16,16代表16位,那就是一个参数占2个字节,2Byte * 70B 约等于 140 * 2^30 = 140GB。

所以你现在能明白为什么int4能这么省显存了吧,足足是原来的1/4呢。

计算KV-Cache所占显存-以Qwen2.5-72B-Instruct为例

计算KV-Cache的思路是,先算单个token单个kv头的KV Cache会用到多少(k,v,记得乘以2),然后乘以kv头的数量(注意这里的头应该是kv头的数量,而不是query头的数量,在GQA场景下,多个query头可以对应同一个kv头!),再乘以模型层数(也就是attention的层数)还有上下文的长度。

举例:下面是Qwen2.5-72B-Instruct的config.json,这里上下文长度按8K算

{
  "architectures": [
    "Qwen2ForCausalLM"
  ],
  "attention_dropout": 0.0,
  "bos_token_id": 151643,
  "eos_token_id": 151645,
  "hidden_act": "silu",
  "hidden_size": 8192,
  "initializer_range": 0.02,
  "intermediate_size": 29568,
  "max_position_embeddings": 32768,
  "max_window_layers": 70,
  "model_type": "qwen2",
  "num_attention_heads": 64,
  "num_hidden_layers": 80,
  "num_key_value_heads": 8,
  "rms_norm_eps": 1e-06,
  "rope_theta": 1000000.0,
  "sliding_window": 131072,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.43.1",
  "use_cache": true,
  "use_sliding_window": false,
  "vocab_size": 152064
}

那么可以看到hidden_size是8192,但是num_attention_heads是64,所以k、v向量的维度是8192/64=128,精度bf16,单个参数占用2字节,单个kv头单token是2*128*2=512字节;kv头8个,4KB;上下文8192=2^13,4KB*2^13=32MB;层数80层,32MB*80=2.56GB.所以对于单个对话窗口,占用2.56GB的KV-Cache。

每个对话并发会占用2.56GB的KV-Cache,真多啊!但是上面的模型用了GQA方法,如果是一般的MHA方法,也就是一个Attention Head对应一个KV Head,这个数字还要乘以8倍——20.48GB!所以说GQA真的是一个折中的,节省KV-Cache占用的好方法。