文章

DeepSeek 笔记

DeepSeek 笔记

Multi-Head Latent Attention

主要作用是加速推理。

kv cache

在之前的 transformer 架构中,我们了解到 Decoder 实际是一个 auto-regressive 的模型,推理时每个 token 生成时都需要前面所有已经生成的 token 信息,这样就避免不了大量的重复计算。这里的重复计算主要是指 token 变为 key 和 value 的过程,也就是 Embedding 和 Positional Encoding(见Transformer 架构图)。

通过将这些已经生成过的 key 和 value 矩阵保存下来用于之后的推理就是 kv cache。但是这种简单的解决方案无非时空间换时间,带来的问题就是消耗了宝贵的内存。

alt text

MLA 就是用来解决 kv cache 过大的问题,原理很简单,原来 Multi-Head Attention 中是把 query、key、value 分别投影到各自的若干个低维空间中(见Multi-Head Attention),假设 $h_t$ 表示输入序列中的第 t 个 token,由 Attention 的理念可知对于 query,key,value 来说这个值是相同的,由 MHA 的定义可知投影的参数不同,所以经过投影后的结果各不相同:

[ q_𝑡 = 𝑊^{𝑄}h_𝑡, ]

[ k_𝑡 = 𝑊^{𝐾}h_𝑡, ]

[ v_𝑡 = 𝑊^{𝑉}h_𝑡, ]

[ [q_{𝑡,1};q_{𝑡,2};\dots;q_{𝑡,𝑛_ℎ}]=q_𝑡, ]

[ [k_{𝑡,1};k_{𝑡,2};\dots;k_{𝑡,𝑛_ℎ}]=k_𝑡, ]

[ [v_{𝑡,1};v_{𝑡,2};\dots;v_{𝑡,𝑛_ℎ}]=v_𝑡, ]

[ \mathbf{o}{t,i} = \sum{j=1}^{t} \text{Softmax}j \left( \frac{\mathbf{q}{t,i}^T \mathbf{k}{j,i}}{\sqrt{d_h}} \right) \mathbf{v}{j,i}, ]

[ \mathbf{u}t = W^O [\mathbf{o}{t,1}; \mathbf{o}{t,2}; \dots; \mathbf{o}{t,n_h}]. ]

MLA 将其中的 $k_t$ 和 $v_t$ 分别再次进行了压缩,这样 key 和 value 作为 cache 时更小:

[ c^{𝐾𝑉}_𝑡 =𝑊^{𝐷𝐾𝑉}h_𝑡,(1) ]

[ k^{𝐶}_𝑡 =𝑊^{𝑈𝐾}c^{𝐾𝑉}_𝑡, ]

[ v^{𝐶}_𝑡 =𝑊^{𝑈V}c^{𝐾𝑉}_𝑡, ]

这里的(1)是将 MHA 的映射和压缩进行了整合,归结为 将$h_t$使用 $𝑊^{𝐷𝐾𝑉}$ 矩阵进行映射+压缩(就是降维,D 就是 down 降维)得到 $c_t$ (也就是 cache),计算权重时可以代替 $h_t$,这样就无需每次都通过 $h_t$ 计算 $k_t$ 和 $v_t$ 了,而且 key 和 value 使用同样的压缩值(解压参数不同,$𝑊^{𝑈𝐾}$ 和 $𝑊^{𝑈V}$),这样只要保存一份 cache 就行。

但是这么做带来了一个问题,MHA 中的 $k_t$ 实际是带了位置信息的,此时每次计算权重时不能直接利用之前的结果(也就是 cache)。DeepSeek 通过额外附加

DeepSeek 还对 query 也进行了压缩,方式和对 key 和 value 的压缩相似,主要是用于降低训练时的内存消耗。

DeepSeekMoE

分为 Shared Expert(常驻) 和 Router Expert(共享)

Router Expert 会被分为几组,通过 Router 进行负载均衡,将问题分配到合适的组中。这是为了使用多节点训练是避免通信消耗,一般一个节点存放一组专家,同一个问题放在一个节点上处理,避免了多节点间的通信。

多节点间的前向和反向传播(pipeline parallel 流水线并行):

alt text

数据就像流水线一样在不同节点间通过,每个节点处理完一个数据就可以立即处理下一个数据,就像工厂里的流水线一样,每道工序都饱和运转。

多节点间同 batch 反向传播时的梯度整合(Data parallel 数据并行):

alt text

Multi-Token Prediction

alt text

Token Prediction 的意图是通过当前 token 预测下一个(组) token,图中最左边输入为 $t_1$ 到 $t_4$ ,参考的结果为 $t_2$ 到 $t_5$,也就是用 $t_1$ 预测 $t_2$ ,以此类推。但仅有最左边部分无法做到让 $t_1$ 预测 $t_3$,换句话说,$t_1$ 无法为预测 $t_3$ 做出贡献。于是添加了后面几个层,这些层都是用前一个层的输出作为输入,所以 $t_1$ 的信息被不断传递,每个 Cross-Entropy Loss 都会被用于改进最左边的主模型的权重。

专家负载均衡

训练过程记录每个 Router Expert 的负载率,仅激活负载率较低且与问题较相关的专家。

强化学习

Reward Model

使用 DeepSeek v3 的一个检查点开始训练专用的 Reward Model ,用来给主模型进行强化训练。该模型训练时不仅包含 reward 结果还加入了得到 reward 结果的 chain-of-thought(思维链)。

Rule-Based RM

数学、代码等能使用明确规则验证的问题使用该 RM。

比如代码可以使用编译器编译后运行输出结果

Model-Based RM

自由形式但有明确答案的,使用该 RM 匹配答案。

没有明确答案的(如创意写作),RM 根据输入的问题和结果给予反馈。

agent 就是模拟人类解决现实问题的方式执行对应步骤解决问题,也就是说它也要能拆分问题、运用工具、查找资料、进行决策等

参考

本文由作者按照 CC BY 4.0 进行授权

© Kai. 保留部分权利。

浙ICP备20006745号-2,本站由 Jekyll 生成,采用 Chirpy 主题。