写这篇文章有四个理由。

周一下午一个僵死一年的传输优化群突然有人问 BBR 的问题,没什么人接话;

有位大模型领域的年轻群友加我,曾看了我很多文章,从我这里学到很多,本着 “交易”原则,我请教他大模型的问题就不算白嫖,但仍感谢他给我指了路;

我每周五晚上到周日早晨都会写网络方面的随笔记录并分享在 pyq,经常有人评论 AI 如何,我颇为不屑;

周四早上看 https://kernelnewbies.org/Linux_6.18#Networking,已经很少有我精通的了;

我并不好奇当今几乎所有的招聘以及各个技术群的话题都与大模型相关,我担心的是如果不去擦点边,将来会很艰难,甚至群里瞎扯淡都插不上话,所以我花了点时间浏览了 https://magazine.sebastianraschka.com,并且购买了塞巴斯蒂安·拉施卡(Sebastian Raschka)的两本书,主要看了《从零构建大模型》第三章,写一点浅见。

大模型与“小模型”

从 “小模型” 可以理解大模型如何做事。所以本文实际上是基于小模型来谈。

不管大小,模型的训练过程就是待定一个一次方程组的参数,使一个损失函数值最小,表面上看就是在通过线性变换求最值。

但大模型需要待定的参数数量非常多,多达百亿千亿,当参数规模如此之大,计算量的规模也将变得即使现代计算机也难以轻易完成,和人脑有大不同,大模型随参数规模不可扩展,因此其复杂性体现在核心架构设计,前沿探索,如MoE,更高效的注意力机制等,解决的是“如何造出一个更好的模型”的根本问题,最终,落实到根本的是当计算量如此之大时所需的算法和技巧。

计算的不可扩展体现在,计算时间或能耗随计算规模而以超过 O(n) 的水平增加,这催生了向量,矩阵和 GPU,TPU 的搭配,于是就有了两类工作,一类与高维,非线性,复杂函数优化的数学理论与算法以及神经网络结构的设计相关,另一类与支撑这些计算的软硬 infra 相关,关联张量并行,GPU/TPU/NPU/JPU(经理处理单元,只要多个经理一人一句话即可)/xPU,博通,英伟达,黄教主,MLX,CUDA,ASIC,万卡集群,NCCL 等磨耳根子的烦人术语,总之,与传统 infra 高可用,高性能,高扩展已经不同,算力,通信,存储成了新的博弈三角。

同时,作为一个黑盒子的大模型训练工作也显得异常重要,这需要深刻理解数据、损失函数和优化器之间的相互作用,N 元一次的简单线性变换,N = 1,2,… 1000 时,小学生理解每一个细节没什么问题,当 N 大于一定阈值,神奇出自不明觉厉,没人知道会发生什么,这就是涌现,因此,这是一个高度经验性的工作。

涌现是什么,我在高中时读过一些复杂系统和混沌理论的书,但我觉得读一些突变理论的资料是有益的。突变理论是由法国数学家勒内·托姆在20世纪60年代提出的。突变理论是个数学框架,用于建模系统状态发生突然的,不连续的跳跃式变化,这种跳跃,是系统从一种稳定形态切换到另一种稳定形态的过程。

当然,本文的 “小模型” 与涌现无关。

此外,针对特定场景的提示工程(prompt),上下文管理更像是在心理学引导,需要大量的实验,直觉和经验,而不是确定的编程逻辑。所感叹的是虽然快排,红黑树,TCP BBR 不再那么重要,但 Python/Java/Go/Rust 等后端开发能力依然是必须。

大模型与一次方程

当 ChatGPT 等大模型生成文本时,它本质上是在求解一个巨大无比的 N 元一次方程组,只是这个 N 可能达到千亿级别!

就像我们初中学习的 { a 0 ⋅ x + b 0 ⋅ y = c 0 a 1 ⋅ x + b 1 ⋅ y = c 1 \left\{ \begin{aligned} &a_0\cdot x+b_0\cdot y=c_0\\ &a_1\cdot x+b_1\cdot y=c_1 \end{aligned} \right. {a0x+b0y=c0a1x+b1y=c1,它适合拟合平面上的一个点,大模型无非是在解:

{ w 00 x 1 + w 01 x 2 + ⋯ + w 0 n x n = 0 . . . w n 1 x 1 + w n 2 x 2 + ⋯ + w n n x n = 0 \left\{ \begin{aligned} &w_{00}x_1 + w_{01}x_2 + \cdots + w_{0n}x_n = 0\\ &...\\ &w_{n1}x_1 + w_{n2}x_2 + \cdots + w_{nn}x_n = 0 \end{aligned} \right. w00x1+w01x2++w0nxn=0...wn1x1+wn2x2++wnnxn=0

只不过方程数量和变量规模超出了人类的直观想象,它拟合的是超高维空间里的点,也就是一个 “意义”。

既然如此,N = 3 和 N = 30000000000 仅在计算上,就没有本质区别,区别只在于量变引起质变后的涌现。如果只为理解 KV Cache,N = 3 足够了。

3 个参数简化 300 亿参数的注意力机制

由于大多数大模型都基于 Transformer 神经网络(2017 年 Google 论文 “Attension Is All You Need” 中提出的神经网络架构)构建,本文就以它为例。

让我们把问题极度简化。假设我们的 “小模型” 只有 3 个参数(后面会将为什么是如此命名的 3 个参数): W K W_K WK W V W_V WV W Q W_Q WQ,分别代表 Key,Value,Query 的权重。做如下设定:

词嵌入,即把输入的原材料词语变成可计算的数字,第 t 个位置的词嵌入为 a t a_t at。由于标量,向量,张量,矩阵在大模型上下文中体现只是运算技巧和效率的不同,与说明原理无关,由于我已经将参数数量简化成了可怜的区区 3 个,也就没有必要上向量,张量,矩阵这些概念,相关本来要表示为向量,矩阵的变量,全用一个数字标量替代,因此:

  • “下雨”: a 0 = 1.0 a_0=1.0 a0=1.0
  • “进水” : a 1 = 2.0 a_1=2.0 a1=2.0
  • "不会 : a 2 = 3.0 a_2=3.0 a2=3.0

参数设定:

  • W K = 2.0 W_K = 2.0 WK=2.0,Key 权重;
  • W V = 1.5 W_V = 1.5 WV=1.5,Value 权重;
  • W Q = 1.0 W_Q = 1.0 WQ=1.0,Query 权重;

参数 Key、Value、Query 的分工

计算之前,需要理解这三个看似并列(3 元一次 x,y,z 的参数)的参数称呼背后的哲学。

Key,Value,Query 的原始含义如下:

  • Query:“要找什么”,即发起查询一方的输入序列,就像搜索引擎中的关键词,其过程表达为 q t = a t × W Q q_t = a_t \times W_Q qt=at×WQ,它是当前词元想要获取信息的 “需求表示”;
  • Key:“被查询对象是谁”,即被查询对象的标识符, 类似匹配搜索关键词的网页标签,其过程表达为 k t = a t × W K k_t = a_t \times W_K kt=at×WK,它是每个词元用于被匹配的 “身份表示”;
  • Value:“有什么实际内容”,即被查询对象的 “实际信息”,就像网页的实际内容,过程表达 v t = a t × W V v_t = a_t \times W_V vt=at×WV,它是每个词元真正携带的 “信息表示”;

实际的 Transformer 中,输入词嵌入,Key,Value,Query 分别展现为向量或矩阵:

  • 输入词嵌入:如 GPT-3 中,其为 d_model = 12288 或更大维度的向量;
  • Key, Value, Query:多维向量,形状为 ( d m o d e l , d v ) (d_{model}, d_v) (dmodel,dv) 矩阵;

不管向量还是矩阵,意义仅在于化简计算过程,整个计算过程是批量矩阵运算,充分利用 GPU 并行能力。由于本文用最简化的 3 元一次方程描述 “小模型”,就使用 3 个标量替代了作为向量和矩阵的 Key,Value,Query。

Key,Value,Query 相当于一件事的 3 个投影,如果只有一个投影,模型就无法区分 “用于匹配的特征” 和 “用于传递信息的特征”,3 个投影意味着:

  • Key 专注 “可匹配性”,学习如何被有效检索
  • Value 专注 “信息含量”,学习如何有效传递信息
  • Query 专注 “查询意图”,学习如何有效表达需求

可以类比找工作:

  • Key 即你的个人介绍,要便于 HR 根据筛选和检索;
  • Value 即你的工作能力和履历,这才是公司真正看重的;
  • Query 即 HR 的招聘需求,表达公司要找什么样的人;

注意力机制的 4 步计算

注意力机制的工作是,将输入词嵌入向量与 Key,Value,Query 相乘,得到 3 个新的 Key,Value,Query 矩阵,用新 Query 求得 S = Query × K T S=\text{Query}\times \text{K}^T S=Query×KT,再将 S 用 Softmax 归一化,获得注意力权重 W,最后加权求和获得输出 Output = W × Value \text{Output}=\text{W}\times \text{Value} Output=W×Value

以上, S [ i , j ] S[i, j] S[i,j] 为第 i 个词的 Query 与第 j 个词的 Key 的相似度,这是注意力的核心,从数学意义上也很容易理解,注意力,就是把不重要的抛之脑后,具体到运算就是相似度越低,权重越低。

由于已经将向量和矩阵简化为了标量,就只需要进行标量算术运算,如下 4 步:

第 1 步:计算 Key,Value,Query

对于每个词元,我们计算三个不同的表示:

  • k t = a t × W K k_t = a_t \times W_K kt=at×WK
  • v t = a t × W V v_t = a_t \times W_V vt=at×WV
  • q t = a t × W Q q_t = a_t \times W_Q qt=at×WQ

具体计算,我将 3 个词写在一起:

  • Key: [ 1.0 × 2.0 , 2.0 × 2.0 , 3.0 × 2.0 ] = [ 2.0 , 4.0 , 6.0 ] [1.0 \times 2.0, 2.0 \times 2.0, 3.0 \times 2.0] = [2.0, 4.0, 6.0] [1.0×2.0,2.0×2.0,3.0×2.0]=[2.0,4.0,6.0]
  • Value: [ 1.0 × 1.5 , 2.0 × 1.5 , 3.0 × 1.5 ] = [ 1.5 , 3.0 , 4.5 ] [1.0 \times 1.5, 2.0 \times 1.5, 3.0 \times 1.5] = [1.5, 3.0, 4.5] [1.0×1.5,2.0×1.5,3.0×1.5]=[1.5,3.0,4.5]
  • Query: [ 1.0 × 1.0 , 2.0 × 1.0 , 3.0 × 1.0 ] = [ 1.0 , 2.0 , 3.0 ] [1.0 \times 1.0, 2.0 \times 1.0, 3.0 \times 1.0] = [1.0, 2.0, 3.0] [1.0×1.0,2.0×1.0,3.0×1.0]=[1.0,2.0,3.0]
第 2 步:计算注意力分数

假设我们计算 “不会”(即第三个词)的注意力,用它的 Query 去匹配所有 Key:

  • s c o r e s = [ q 3 × k 1 , q 3 × k 2 , q 3 × k 3 ] scores = [q_3 \times k_1, q_3 \times k_2, q_3 \times k_3] scores=[q3×k1,q3×k2,q3×k3]
  • s c o r e s = [ 3.0 × 2.0 , 3.0 × 4.0 , 3.0 × 6.0 ] = [ 6.0 , 12.0 , 18.0 ] scores = [3.0 \times 2.0, 3.0 \times 4.0, 3.0 \times 6.0] = [6.0, 12.0, 18.0] scores=[3.0×2.0,3.0×4.0,3.0×6.0]=[6.0,12.0,18.0]
第 3 步:Softmax 归一化

通过 Softmax 将分数转换为权重:

  • W = softmax ( [ 6.0 , 12.0 , 18.0 ] ) ≈ [ 0 , 0 , 1 ] W = \text{softmax}([6.0, 12.0, 18.0]) \approx [0, 0, 1] W=softmax([6.0,12.0,18.0])[0,0,1]
第 4 步:加权求和

用权重对 Value 加权求和:

  • Output = 0 × 1.5 + 0 × 3.0 + 1 × 4.5 = 4.5 \text{Output} = 0 \times 1.5 + 0 \times 3.0 + 1 \times 4.5 = 4.5 Output=0×1.5+0×3.0+1×4.5=4.5

这个计算流程体现了 Key,Value,Query 的功能分工,用 Query 在一堆 Key 中找出相关的,然后去对应的 Value 那里获取实际信息。

自回归生成中的重复计算

当模型生成下一个词时,传统做法是重新计算所有内容:

输入序列:["下雨", "进水", "不会"] -> 输出:"胖"
新输入序列:["下雨", "进水", "不会", "胖"] -> 重新计算所有 Key,Value

关键问题是,前 3 个词的 Key,Value 需要重新计算,即使它们完全没有变化,而这是绝对冗余的操作,这就好比每次解方程时,都把前面已经算过的部分重新算一遍,造成了巨大的计算浪费。

KV Cache 如何避免重复计算

KV Cache 的洞见是,在生成过程中,历史词序列的 Key 和 Value 不会改变,不需要每次都重新计算。

KV Cache 的工作方式很容易理解,第一次计算时缓存历史信息:

K_cache = [2.0, 4.0, 6.0]
V_cache = [1.5, 3.0, 4.5]

生成新词时只需计算新增部分:

  • k 4 = a 4 × W K = 4.5 × 2.0 = 9.0 k_4 = a_4 \times W_K = 4.5 \times 2.0 = 9.0 k4=a4×WK=4.5×2.0=9.0
  • v 4 = a 4 × W V = 4.5 × 1.5 = 6.75 v_4 = a_4 \times W_V = 4.5 \times 1.5 = 6.75 v4=a4×WV=4.5×1.5=6.75

然后从缓存中读取历史信息,拼接成完整的序列。

KV Cache 可行的依据在于所有计算都是线性的,正交可分解,就像向量分解在每个维度分别进行一样,我用一个数字标量的加法算术简化了,但如果换回向量,矩阵代数运算,只要运算是线性的,它和标量加法就没有本质区别,向量,矩阵与标量的区别仅在于 “运算技术”的效率区别,对结果的性质没有改变,相关的 “运算”在大模型上下文中,不管技术如何,它们都满足:

KV Cache 可行的依据在于所有计算都是正交可分解,就像向量分解可分别进行一样:

  • 参数不变性,模型参数 W K W_K WK W V W_V WV 在推理过程中不变;
  • 历史不变性,已经生成的词元嵌入不变 ;
  • 计算局部性,只有当前词嵌入需要计算 Query;

这也多亏了简单的多元线性运算,这种设计将 Key,Value 的计算复杂度从 O ( n 2 ) O(n^2) O(n2) 降低到 O ( n ) O(n) O(n),实现了数量级的性能提升。

从 3 个参数回到千亿参数

在这个 3 元一次极简例子中,KV Cache 可能只节省了 3 次乘法。但在真实的大模型中:

  • 参数规模从 3 个参数变成百亿,万亿参数;
  • 序列长度从 3 个词变成数万个 token ;
  • 计算节约从节省 3 次计算变成节省数亿次计算;

对于生成 2000 个 token 的对话,KV Cache 可能减少 90%+ 的重复计算!

我将向量,矩阵简化为 3 个标量是合理的,因为只有 3 个参数,标量算术足够了。可以从下面的认知去理解这个理由,标量算术 vs. 矩阵运算本质上就是 CPU vs. GPU 在软件上的映射,都是串行和并行处理的区别,如果我不需要处理超大规模的运算量,我就不需要 GPU,用 CPU 就足够了,同理,只要 3 个参数,标量就足够了。

否则,本文的重心将偏离并陷入到矩阵运算操作的细节,后者更适合去学一门课程,而不是看一篇文章。

本质还是解方程

大语言模型看似复杂,但其数学本质依然是求解线性方程组,即使矩阵运算充斥着再巧妙无比的优化技巧,海量参数下运算量依然巨大。KV Cache 意味着即使是求解千亿级别方程组这种简单重复的事,也可以通过巧妙的缓存策略避免大量重复计算。

也恰恰是多亏了这种简单的一次线性变换让正交可分解变得可行,不会让这海量参数互相依赖而纠缠,而这种简单加法线性运算也正是机器擅长的,且非常容易并行化,这又是因为正交可分解。

Key,Value,Query 的设计体现了深度学习的精妙,在数学对称中实现功能分工,3 个投影在数学形式上完全并列,但在功能角色上各司其职,共同实现了精准的信息检索机制。

a x + b y = c ax + by =c ax+by=c 到千亿参数的复杂大模型,数学的简洁与优雅始终贯穿其中,最有效的优化,往往来自于对基础数学的深刻理解与巧妙应用。

KV Cache 利用了一次方程组的可正交化的易算特征,让大模型从 “重复解方程” 硬算变成了 “增量解方程”,这是 AI 工程领域的重要突破。参考杨辉三角二项式,但凡次数 > 1 不再线性,各项就会纠缠在一起,拆毁 KV Cache 的可行性根基。

KV Cache 改变了算力扩展性的复杂度, O ( n 2 ) O(n^2) O(n2) O ( n ) O(n) O(n) 的意义远大于 CDN 对带宽的意义。至于 KV Cache 的编码,传输和管理,就是网络和 Cache 管理的范畴了。

综合真实大模型,百亿千亿参数,虽然只是线性加法,问题在于一锅装不下,数据并行,张量并行,序列并行,连同 KV Cache,难题是如何算得快,如何传得快,如何存取得快,劳动密集型可以形容大模型的工作方式,就像造金字塔,每个工人本质上都只是搬砖并摞上去,但千万人一起做这事,就有了奇迹。

后面我会再谈谈算力,通信,存储的事。

浙江温州皮鞋湿,下雨进水不会胖。

更多推荐