【自用】NLP算法面经（4）

一、deepseek

1、MLA

（1）LLM推理过程

prefill阶段：模型对全部的prompt tokens一次性并行计算，最终生成第一个输出token。
decode阶段：每次生成一个token，直到生成EOS（end-of-sequence）token，产出最终的response。

在LLM生成过程中，是一个基于前向序列token预测下一个token的过程，序列中的token（无论是prefill阶段，还是decode阶段）只与它前面的token交互来计算attention。矩阵计算上通过一个下三角的causal attention mask来实现token交互只感知前向序列。
在这里插入图片描述
以一个序列的t位置的token为例，计算一层transformer的attention过程，如下列公式所示：

从（7）可以看到，在计算attention时，t位置的q只与t位置前的k,v做计算，所以：

计算前面的k,v并不受后面token的影响。
后面计算t+1,t+2,…,t+n位置的attention，要使用前序的1->t位置的k,v的值是始终不变的。
因此为了加速训练和推理效率，在token-by-token生成过程中，避免重复计算前序的k,v。目前主流的kv cache就是把前序计算好的k,v缓存起来，本质是通过空间换时间的方法，因此kv cache势必会带来访存的瓶颈。即如果不用kv cache直接重复计算前序kv，是一个计算密集型任务；增加了kv cache，现在kv不是通过计算得到的，而是从“存储介质”中读取得到的，GPT内核和存储介质之间要频繁读写，这样就变成了一个访存密集型任务。

（2）LLM推理阶段显存使用情况

（2.1）访存速率分级

为了直观理解访存的速率，以一个分布式推理架构为例。

假如有两台机器，每台机器有8张A100，则在这个系统中，卡内、单卡卡间、机器之间的数据访问效率如下图所示。
GPU的存储介质除了HBM（显存），还有SRAM和DRAM。SRAM也被称为片上存储，是GPU计算单元上即时访问更快的存储，所有的计算都要先调度到SRAM才能计算，一般只有几十M大小，带宽可达到20T/s左右，SRAM是和计算单元强绑定的，推理阶段一般不考虑将SRAM作为存储单元使用。而DRAM就是CPU的内存，由于访问速率较慢，推理阶段一般也不考虑使用。所以推理存储介质一般就是HBM。

由上图的访存带宽可知，卡内的带宽是单机卡间带宽的6倍，是跨机带宽的20倍，因此对于存储的数据优先放到卡内，再到卡间，最后才考虑跨机存储。

（2.2）模型推理阶段显存分配

推理阶段主要有三部分数据会放到显存。

KV Cache：前序token序列计算的k,v结果，会随着后面token推理过程逐步存到显存中。存储的量随着batch,sequence_len长度变换。
模型参数：包括Transformer，embedding等模型参数会存到显存中。模型大小固定后，这个存储空间是固定的。
运行时中间数据：推理过程中产出的一些中间数据会临时存到显存，即用即释放，一般占用空间较小。

因此推理阶段主要存储消耗为两部分：kv cache和模型参数，它们分别各占多少？
以一个token的计算过程为例，使用Qwen-72B为例。

模型共80层，每层有64个头，每个头的向量维度是128。
$l=80,n_h=64,d_h=128$
ps：先不考虑Qwen-72B GQA的设置（压缩了kv cache），只考虑一般的MHA。

计算一个token，每个层的每个头都要保存一对k和v。
因此对于一个token，缓存的k,v数据总量：
$num_{kv} = 2 \times l \times n_h = 2 \times 80 \times 64 = 10240$
一个token就要缓存10240个k,v，那这么多k,v占了多少存储？假设模型推理阶段是半精度（bf16）参数，每个参数占2字节，最终一个token的存储占用为：
$token\_mem_{kv} = 2 \times num_{kv} \times d_h = \frac{2 \times 10240 \times 128}{1024 \times 1024} MB = 2.5MB$
对于一个实际的推理场景，还要考虑batch和sequence_len两个维度来确认整体kv cache的存储消耗。这两个维度通常是可以动态变化的，下面有两个场景：
场景1：单条短文本场景
batch和sequence_len设置：B=1，S=2048。此时k v cache总量：
$ mem_{kv}=1token_mem_{kv} \times B \times S = 2.5MB \times 1 \times 2048 = 5GB$.
场景2：并发长文本场景
batch和sequence_len设置：B=32，S=4096。此时k v cache总量：
$ mem_{kv}=1token_mem_{kv} \times B \times S = 2.5MB \times 32 \times 4096 = 320GB$
除了k v cache消耗存储空间，模型参数也要占用存储，推理阶段模型参数占用的存储空间是固定的。假设模型参数量为 $\Phi$ ，以bf16半精度做推理，则参数量为 $2\Phi$ （Byte）。还是以Qwen-72b为例，参数占用存储空间：
$mem_p = 2 * \Phi= 2 * 72 = 134GB$
结合上面两个场景，查看显存的整体分配：

场景一：模型参数存储134GB，kv存储5GB，模型参数存储占主导，若使用80G的A100，至少需要2张卡做推理。
场景二：模型参数存储134GB，kv存储320GB，kv存储占主导，若使用80G的A100，至少需要6张卡做推理。

当前大模型都比较大，而访存的容量和访存的速率有分级的特点。因此在推理过程中，减少跨卡、跨机的访存读写是优化推理性能的一个有效路径。

（3）减小KV cache的方法

（3.1）KV cache优化方法汇总

共享KV：多个头共享一组k和v，将原来每个头一个kv，变成1组头一个kv，来压缩kv的存储。代表方法：GQA、MQA。
窗口KV：针对长序列控制一个计算KV的窗口，KV cache只保存窗口内的结果（窗口长度远小于序列长度），超出窗口的KV会被丢弃，通过这种方法能减少KV的存储，也会损失一定的长文推理效果。代表方法：Longformer。
量化压缩：基于量化的方法，通过更低的bit位来保存KV，将单KV结果进一步压缩，代表方法：INT8。
计算优化：通过优化计算过程，减少访存换入换出的次数，让更多计算在片上存储SRAM进行，以提升推理性能，代表方法：FlashAttention。

（3.2）共享KV优化显存方法

（3.2.1）MQA（Multi-Query Attention）

每一层的所有头共享一组KV来计算Attention。相对于MHA的单个token需要保存的KV数（ $2 * l * n_h$ ）个减少到了 $2 * l$ 个。

（3.2.2）GQA（Group-Query Attention）

对所有头分组，比如分组数为g，则 $\frac{n_h}{g}$ 个头共享一个KV。当g=1时，GQA就等价于MQA，当 $g=n_h$ 时，GQA就等价于MHA。

（4）MLA（Multi-Head Linear Attention）

在这里插入图片描述

（4.1）MLA原理解读

$d_c$ ：MLA低秩压缩的维度，论文中取值： $d_c = 4 \times d_h$
$d_h$ :单个头的向量维度
$n_h$ ：每层头的数量
$d$ :隐层维度， $d=d_h \times n_h$
$W^{DKV} \in \mathbb{R}^{d_c \times d}$ 是低秩变换矩阵
①KV计算过程
首先公式（41）对输入 $h_t$ 做一个低秩压缩，将d维的输入经过 $W^{DKV}$ 变换后压缩成 $d_c$ 维的 $c_t^{KV}$ 。DeepSeek-V3中 $d-7168,d_c=512$ 。

然后通过公式（42）和（45）两个变换矩阵（ $W^{UK},W^{UV} \in \mathbb{R}^{d_hn_h \times d_c}$ ），将KV的维度扩展回 $d = d_hn_h$ ，即每个头有一个单独的k,v（跟MHA的kv数量一致）。

经过上述变换，非常类似lora做低参数微调的过程。通过两个低秩矩阵先做压缩，再做扩展，最终能降低参数的数量。但MLA本质时要做到减少KV Cache的存储。lora强调的是参数量的减少，类似MLA的操作确实减少了参数量，按DeepSeek-V3的参数配置，两个低秩矩阵配置： $\times d_c \times d=2 \times 512 \times 7168$ ，而正常MHA的参数矩阵参数量： $\times d = 7168 \times 7168$ 。但MLA强调的是kv cache的减少，即kv的激活值减少。

②Q的计算过程
公式（37）（38）类似KV的逻辑，通过两个矩阵（ $W^{DQ},W^{UQ} \in \mathbb{R}^{d_hn_h \times d_q}$ ）也做了一层低秩变换，这一步Q的变换是为了减少模型参数的数量。在DeepSeek-V3中， $d_q=1536$ ，是KV压缩维度 $d_c$ 的3倍，但相对于d=7168还是压缩了不少。
在这里插入图片描述

③q,k增加rope位置编码
增加rope位置编码并没有在上述计算出的 $q_t^C,k_t^C$ 的基础上乘以rope的对角矩阵，而是单独计算了两个带着位置编码的 $q_t^R,k_t^R$ 。
在这里插入图片描述

$q_t^R,k_t^R$ 的向量维度 $d_h^R$ 是一个比较小的维度，DeepSeek设置为单attention head维度的一半： $d_h^R = d_h / 2 = 64$ ；
这部分计算 $k_t^R$ 实际是个MQA的计算方式，同一层中，所有头共享同一个k。

然后按照公式（40）（44）跟已经计算的 $q_t^C,k_t^C$ 拼接，构成完整的 $q_t,k_t$ 向量。
到目前为止，得到的q,k包括两部分拼接而成，一部分那是做了低秩压缩得到的q,k向量，一部分是增加了rope位置编码的q,k向量。

DeepSeek原论文对上述操作过程有一段解释：

位置编码使用rope，但rope与低秩KV不兼容。具体来说，rope对q和k都是位置敏感的。若我们为 $k_t^C$ 应用rope，那么公式(42)的W（k的权重矩阵）将与位置敏感的rope矩阵耦合。因此，在推理过程中， $W^{UK}$ 无法再被吸收到 $W^{UQ}$ 中，因为与当前生成token相关的rope矩阵位于 $W^{UK}$ 和 $W^{UQ}$ 之间，而矩阵乘法不满足交换律。因此，我们必须在推理过程中重新计算所有前缀token的k，这将极大地降低推理效率。

论文中提到了“矩阵吸收计算”，这个概念对理解MLA比较重要，用一个简单的例子理解：

假设有两个向量变量 $x_1,x_2 \in R^{3 \times 1}$ 都是三维向量。有两个固定的变换矩阵 $\in R^{2 \times 3}$ 分别对 $x_1,x_2$ 做线性变换得到新的向量 $x'_1,x'_2$ .。最终求 $x'_1,x'_2$ 两个向量的乘积。
方法1：常规计算
$\begin{equation} x'_1 = Px_1\end{equation}$ $\begin{equation} x'_2 = Qx_2\end{equation}$ $\begin{equation} x'^T_1x'_2= (Px_1)^T * (Qx_2) = x^T_1P^T Qx_2\end{equation}$
方法2：矩阵吸收计算
矩阵乘法满足结合律，对于公式(3)可以先计算好两个变换矩阵的乘积：
$\begin{equation} Q'= P^TQ \end{equation}$
然后通过Q’和x2现成，计算出x2’‘，而x1不做任何操作：
$\begin{equation}x''_2=Q'x_2 \end{equation}$
再计算x1和x2’'乘积：
$\begin{equation}x^T_1x''_2=x^T_1Q'x_2=x^T_1P^TQx_2=x'^T_1x'_2\end{equation}$

理解了上面的例子，再来看原文中的“RoPE与低秩KV不兼容，没法做矩阵吸收计算”的问题。

a)不加rope
假设当前不添加rope，那么q,k乘积计算如下，其中(i)表示变换矩阵第i个头的切片：
$q^T_{t,i} \times k_{t,i}=(W^{UQ}_{(i)}c^Q_t)^T \times W^{UK}_{(i)}c^{KV}_j=(c^Q_t)^T \times (W^{UQ}_{(i)})^TW^{UK}_{(i)} \times c^{KV}_j$
不加rope，可以提前计算 $(W^{UQ}_{(i)})^TW^{UK}_{(i)}$ ，即 $W^{UK}$ 吸收到 $W^{UQ}$ 中，这样在做q的变换时，也就同时计算了 $W^{UK}$ 矩阵的乘法。
这样做的好处是，只需要缓存 $c^{KV}_j$ ，而不是缓存 $W^{UK}_{(i)} \times c^{KV}_j$ 的结果。 $c^{KV}_j$ 只有 $4d_h$ 的长度，而 $W^{UK}_{(i)} \times c^{KV}_j$ 是个 $4d_h->d$ 的变换，即完全恢复了隐层维度 $d = n_h * d_h=64d_h$ 。

b)假设更加rope
加上rope后，计算q,k乘积，会在 $(W^{UQ}_{(i)})^T$ 和 $W^{UK}_{(i)}$ 之间，增加一个融合了相对位置的变量 $R_{t-j}$
在这里插入图片描述
中间这个变量 $(W^{UQ}_{(i)}R_{t-j}W^{UK}_{(i)})$ 是随相对位置变化的，并不是固定的矩阵，因此并不能提前计算好，所以论文说rope与低秩变换不兼容。

c)通过增加一个很小的q,k分量，引入rope
为了引入位置编码，作者在一个很小的维度下，用MQA计算了q,k，即在每层网络中，所有head只计算一个k。引入位置编码的向量维度取的比较小为： $d_h / 2 = 128 / 2 = 64$ 。
所以最终q,k向量通过两部分拼接而成，计算权重时，由前后两部分分别相乘再相加得到，如下公式所示：
在这里插入图片描述
MLA实际缓存的向量：

$c^{KV}_t$ :维度为 $\times d_h = 512$
$k^R_t$ :维度为 $d_h / 2=64$

$c^{KV}_t$ 是低秩压缩的向量， $k^R_t$ 是引入位置编码的MQA范式计算的共享k。

自己的话总结！！
MLA的提出同样是为了减少KV Cache，MLA的方法本质上是对原本MHA的KV Cache做低秩分解，得到一个低维的隐向量，在推理阶段，MLA只需要缓存这个隐向量，有次大大降低需要缓存的数据量。
首先，对于KV的处理，先对输入做一个低秩压缩，将d维的输入经过低秩变换矩阵后压缩为 $d_c$ 维的向量，其中输入的d维为7168，而压缩后的 $d_c$ 维为512。接着，再通过两个变换矩阵将KV的维度扩展回原始的d维。
而对于Q的处理，同样通过两个矩阵做了一层低秩变换，以减少模型参数量，在deepseek-v3中，从原始d=7168压缩到了 $d_q=1536$ 。
以上操作还需要针对旋转位置编码进行处理，因为如果隐向量中包含rope，经过升降维操作后，会对位置信息造成破坏，为了解决这个问题，MLA提出解耦rope的方法，即对于隐向量，不将rope包含在其中，而是专门为注意力头的query和key新增一个小的向量维度，deepseek-v3中为头的维度 $d_h$ 的一半64维，以添加rope的位置信息。
最后，将做了低秩压缩得到的q,k向量和增加了rope的q,k向量进行拼接得到最终的q,k向量。
因此，在MLA中，全头只缓存一个维度为512的低秩压缩向量和一个引入位置编码的共享k向量。

2、MTP（Multi-Token Prediction）

（1）作用

训练阶段：通过预测多步token，迫使模型学到更长的token依赖关系，从而更好理解上下文，避免陷入局部决策的学习模式。同时一次预测多个token，可大大提高样本的利用效率，相当于一次评估可生成多个<predict,label>样本来评估模型，有助于模型加速收敛；
推理阶段：并行预估多个token，提升推理速度。

（2）DeepSeek MTP

在这里插入图片描述

(2.1)MTP模块实现细节

用D个顺序的模块，预测D个token，每个MTP模块的具体结构：

输入token首先接入一层共享的embedding layer；
对于第i个token $t_i$ 和第k个预测深度
- 首先将第k-1层的隐层输出 $h^{k - 1} \in \mathbb{R}^d$ 做归一化处理 $RMSNorm(h^{k-1}_i)$
- 再对第i + k位置的token embedding: $Emb(t_{i+k}) \in \mathbb{R}^d$ 做归一化处理 $RMSNorm(Emb(t_{i + k}))$
- 将上述两个结果拼接后，通过投影矩阵 $M_k \in \mathbb{R}^{d \times 2d}$ 做一层线性变换得到 $h'^k_i \in \mathbb{R}^d$
- 上述过程如下公式所示（当k=1时， $h_i^{k-1}$ 对main model的隐层表征）
- 再将 $h'^k_i$ 输入到Transformer层，获得第k个预测深度的输出 $h^k_i$ ，如公式所示：
- 最后将 $h^k_i$ 通过一个各module共享的映射矩阵 $\in \mathbb{R}^{V \times d}$ 变换，再过softmax()处理，计算出词表V维度的输出概率。注意： $h^k_i$ 的label是对应i+1+k位置的token。如公式所示：

$h_i^k$ 是第i个token在第k个预测深度上输出的表征，是要预测序列中第i+k位置的token的。由于序列总长度为T，所以第k个预测深度最长处理的输入token位置i应该满足 $\leq T$ ，所以第k预测头能接受的i的范围为： $\leq T-k$ ，也就是 $\in [1,T-k]$ ，也是(22)表示的切片范围。
举个简单的例子：T=10，对于k预测深度，模型训练期间样本构建方式，如图所示，main model是预测next token，所以input和label序列错开1位，MTP Module1是预测next next token，input和label序列错开2位，在T+1总长度下，输入的后续token和输出的前序token都要按错位裁剪。

（2.2）MTP模型训练

通过交叉熵损失计算每个MTP Module Head的损失，如公式所示：
在这里插入图片描述

2+k:T+1表示label范围的下标。
起始下标2+k：MTP Module 1是预测next next token，即输入第一个token是t1，预测第一个label token是 $t_{1+2}=t_3$ ，以此类推，MTP Module k，输入第一个token是t1，预测第一个token是 $t_{1+k}$ 。
结束下标T+1：所有序列样本默认在原序列上额外增加一个eos token，所以token下标为序列长度T+1。

有一个问题，参考公式（23）的表述，第k预测深度是输入 $t_i$ 来预测 $t_{i + k+ 1}$ 。比如MTP Module 1，输入第一个token $t_1$ 来预测 $t_3$ 。但MTP Module 1的输入明明还有一个是t2，这怎么理解？

这是处理序列建模任务中典型的Teacher forcing模式。正常应该是拿上一个状态的输出（也就是图中的 $t'_2$ ）作为输入，但在序列建模训练中，直接用样本的ground truth作为输入，效果会更好。因为如果拿预估的状态 $t'_2$ 作为输入，随着时间的推移，预估错误会持续累加，导致效果有损。

与Teacher forcing模式相对应的事free-running模式，free-running是直接用上一个状态的输出，来作为下一个状态的输入。

（2.3）MTP模型推理

DeepSeek-V3推理有两种方法：
**方法1：**直接把MTP Model头全部删掉，模型变成了一个predict next token的main model。然后部署模型做推理，这就跟正常LLM推理一样没有什么加速效果。
**方法2：**保留MTP Model做self-speculative decoding，这样充分使用多头预测能力，提升推理加速性能。

阶段1：predict，利用k个头一次生成k个token，每个头生成一个token；
阶段2：verify：将原始的序列和生成的token拼接，组成多个 $Pair <sequence_input,label>$ ，将组装的多 $Pair <sequence_input,label>$ 组成一个batch，一次发给main model做校验；
阶段3：accept：选择 $Head_1$ 预估token与label一致的最长k作为可接受的结果。

阶段1不实用teacher forcing模式，因为teacher forcing模式只用于训练阶段，推理阶段要用上一个状态的预估值作为下一个状态的输入（free-running模式）。
在这里插入图片描述
自己的话总结！！
MTP的目的是：在训练阶段，通过预测多步token，让模型能够学习到更长的token依赖关系，从而更好地理解上下文，避免陷入局部决策的学习模式。同时也能进一步提高训练样本的利用效率；在推理阶段，并行预估多个token，提高推理速度。
MTP是由一个原始预训练模型和K个MTP module组成的。
对于序列中第i个token和第k个预测深度，MTP训练流程为：

首先，token经过一层共享的embedding layer；
对第k-1层的隐式输出做均方根归一化；
对第i+k位置的token embedding做均方根归一化；
将以上两部分进行拼接后，通过一个投影矩阵做一层线性变换后输入transformer层，获得第k个预测深度的输出；
最后将第k个预测深度的输出经过一个各module共享的投影矩阵变化，在经过softmax处理，计算出词表V维度的输出概率。因为deepseek-v3中预训练模型预测的是next token，第一个MTP Module预测的是next next token，因此第k个MTP Module预测的是第k个next token，因此第k个预测深度的输出的label是对应于第i + k位置的token。
获得输出概率后，通过交叉熵损失函数计算每个MTP Module的损失。

MTP推理流程为：
在推理阶段包括3步，分别是predict，verify和accept。

在predict阶段，利用k个head一次生成k个token，其中每个头生成一个token；
在verify阶段，将原始序列和生成token拼接，组成多个<sequence_input,label>对，其中sequence_input就是原始序列，label就是生成的token，然后将多个队组成一个batch，一次发给预训练模型做校验；
在accept阶段，选择预训练模型预估的token和在predict阶段生成的k个token中一致的最长token序列作为可接受的结果。

其中，在训练阶段采用的是teacher forcing模式，即在序列建模任务重，下一个状态的输入并不是上一个状态的输出，而是将样本的ground truth作为输入，以避免随着时间的推移，预估错误的持续累加导致的效果受损。
在推理阶段采用的是free-running模式，即将上一个状态输出的预估值作为下一个状态的输入。

3、MoE（Mixture-of-Experts）

（1）MoE的发展历程

在这里插入图片描述
MoE是一种网络层结构，网络层主要包括三部分：

专家网络：是一个前馈网络，逻辑上一个专家网络擅长处理一类专项的子任务，所有专家接受相同的输入，来做特定计算处理，产出不同输出；
门控网络：跟专家网络接收一样的输入，负责产出专家偏好的权重，来指示对于一个输入，不同专家的重要程度。
选择器：是一种根据专家权重来做专家选择的策略。可以选择权重最高的top1专家或选择topk专家来融合得到最终的结果。

Transformer MoE:MoE层替换了Transformer的前馈神经网络层，计算逻辑：对于一个token，分别通过门控网络和专家网络计算门控值和专家输出，然后用门控值加权多个专家输出来产出最终结果。

门控计算：
专家计算：

多专家结果加权求和得到MoE的输出

这里的专家是token级专家，而不是样本粒度，每个token都会做专家路由。此外专家是稀疏激活的，是根据门控值取top k 个专家来融合计算最终的结果。GShard最多激活权重最高的2个专家。

负载均衡-辅助损失：引入负载均衡损失，目的是解决多专家token分布不均的问题。因为如果完全按门控权重选取top k个专家，容易导致训练过程中出现负载不均衡的问题。比如：大多数token被分配到少数几个专家，导致只有少数专家数据通信繁忙造成拥堵，从而减缓训练速度；也会导致其他专家得不到充分训练。为了解决这个问题，定义了一个辅助损失（aux_loss）来解决负载不均衡的问题。

如何定义负载均衡的辅助损失？
考虑每个专家收到的token占总token的比例，分别为 $\frac{c_1}{S},\frac{2_1}{S},\dots,\frac{c_E}{S}$ ,S表示token的总数量，{1,2,…,E}表示专家集合， $c_e$ 表示第e个专家接受的token数量。如果是负载均衡的，则每个专家收到的token一样多，token比例 $\frac{c_i}{S}$ 值一样。
可以用每个专家收到的token比例的平方和来描述负载均衡损失，如下公式所示，当所有专家收到token比例都相等时， $l_{aux}$ 取最小值。
在这里插入图片描述
但由于公式（1）是参数无关的量，不可梯度更新。作者用每个专家的门控权重的均值 $m_e$ 作为 $\frac{c_e}{S}$ 的近似。如下公式所示：

其中 $g_{s,e}$ 为公式(1)针对token s计算的专家e的门控权重。

为什么 $m_e$ 可以看作是 $\frac{c_e}{S}$ 的近似？
假设极端情况下每个token最多分配给1个专家，那么可以假设被激活的专家的权重可以是1，其他专家权重为0，这样对于专家e来说，可以计算得到 $m_e = \frac{c_e}{S}$ 。另外从定义上 $m_e=\frac{1}{S} \sum_{s=1}^Sg_{s,e}$ 表示token集合S被分配给专家e的概率，如果不考虑token分配的完整性，这其实就是 $\frac{c_e}{S}$ 的定义。只不过 $m_e$ 是个soft的计算方式，而 $\frac{c_e}{S}$ 是取top k的hard计算的。
比如，图示展示了6个专家，在6个token上计算，取top 1专家激活。左图每一行是多专家softmax的结果，，左边按列相加计算 $\sum _{s=1}^Sg_{s,e}=m_e \times S$ 其实是计算分配给专家e的token数，是soft的计算；右边按列加和计算 $c_e$ 也是计算分配给专家e的token数，是hard的计算。 $c_i$ 和 $m_e \times S$ 的值是近似的。

用 $m_e$ 把公式（4）改造一下，将平方项的一个分量替换成 $m_e$ ，如公式（6）：
在这里插入图片描述
公式（6）就是经常看到的负载均衡loss形式。对于专家级的负载均衡的loss是加到每个MoE层的，每层都有一个 $l_{aux}$ 辅助损失。

（2）DeepSeek-MoE(V1)

14年1月DeepSeek发布V1版MoE模型，作者指出当前方法存在两方面问题：

知识混合性：现有的MoE模型通常使用数量有限的专家（如8个或16个），由于token的知识是丰富多样的，将多样的知识分配给有限的专家，会导致特定专家的token很可能涵盖多样化的知识，而使得专家变成一个杂糅多知识的专家，这样不能充分发挥专家的专业效果。
知识冗余性：分配给不同专家的token可能存在共同知识。因此，多个专家可能会在其各自的参数中学习到共享知识，从而导致专家参数存在冗余。
针对上述问题，DeepSeek引入一种实现了专家专业化而设计的创新MoE架构。架构主要包含两方面优化：
细粒度专家分割：在保持参数量不变的情况下，作者通过分割前馈神经网络中间隐藏维度来将专家分割成更细的粒度。相应地，在保持计算成本不变的情况下，可激活更多细粒度的专家，以实现激活专家组合的更高灵活性。细粒度专家分割使得多样化只是能够被更细致地分解，并更精确地学习到不同的专家中，每个专家将保持更高的专业化水平。
共享专家隔离：将某些专家隔离出来，作为始终激活的共享专家，旨在捕获不同上下文中的共同知识。通过将共同知识压缩到这些共享专家中，可以减轻其他路由专家之间的冗余。

DeepSeek MoE架构的公式：
在这里插入图片描述
其中 $K_s$ 是共享专家数量， $mN-K_s$ 是路由专家数量， $u_t^l$ 是第l层第t个token的输入，计算 $h_l^t$ 的三个因子分别是：贡献专家结果、路由专家结果、残差连接。 $e_i^l$ 是l层专家i可学习的参数， $s_{i,t}$ 表示第t个token在第i个专家上的打分， $g_{i,t}$ 表示取top k高的专家权重。
除了在模型架构上的改进，随着deepseek从v1到v3的演进，在负载均衡上，做了较多工作，先看v1额负载均衡的优化，主要是在计算负载均衡上做了优化，包括两个负载均衡的设置：
①专家级负载loss
loss计算如下：
在这里插入图片描述
其中， $\alpha_1$ 是超参数，用来调节与主网络loss的权重；T是专家要处理的全部token数； $N'=mN-K_s$ 表示去掉共享专家后的路由专家的数量； $K'=mK-K_s$ 表示激活路由专家的数量； $\mathbb{I}$ 是指示函数。
针对上述公式(13) $f_i$ 的计算，若参照公式(6)，计算 $f_i$ 应该为：
在这里插入图片描述

（13）相比于（15），分子多乘了路由专家数，分母多除了激活路由专家数。

为什么要乘以N’并除以K’？
是为了保持计算损失的恒定，不随专家数量的变化而变化。

当token分配均匀的情况下，也就是说T个token，每个token激活 $K^{'}$ 个专家。共需要分配TK’个token，平均分配给N’个专家，则每个专家被分配的token数为TK’/N’。那么按照(15) $f_i$ 计算的专家i分配的token率为: $f_i = K'/N'$ 。
考虑 $P_i$ 计算，当token均匀分配，且 $s_{i,t}$ 计算softmax保持将权重均匀分配给K’个激活的专家，即计算的权重类似于 $[0,1/K',0,0,1/K',\dots]_{N'}$ 。softmax计算后的权重向量维度为N’，其中有K’个位置为1/K’，其他位置都为0。token均匀分配的情况下，每个专家有非零权重的token数为TK’/N’。按 $P_i$ 公式计算由：
$P_i = \frac{1}{T} \times \frac{TK'}{N'} \times \frac{1}{K'}= 1/N'$
所以loss计算： $\mathcal{L}_{ExpBal}=\alpha_1\sum_{i=1}^{N'}(K'/N' \times 1/N')=\alpha_1 \times K'/N'$ 。在最终loss里面有个K’/N’项，是随着路由专家数（N’）和激活路由专家数（K’）动态变化而变化的。为了去掉这个动态变化的项，让Loss维持一个恒定的量级，对辅助loss整体乘以N’/K’，以保持loss计算是不随专家数变化而变化的。

为什么保持loss的计算不随专家的数量变化？
①超参 $\alpha_1$ 的调整简单。超参 $\alpha_1$ 是平衡主loss和辅助loss的超参，既不能太大，也不能太小，太大会干扰主loss的收敛效果，太小会达不到负载平衡的目标。所以如果辅助loss随专家数变化，那么调整超参 $\alpha_1$ 会比较复杂。
②做专家数对比消融实验时，如果loss不受专家数设置影响，那么loss收敛的绝对值是有可比性的。尤其在做细粒度专家效果对比时，不同实验的绝对loss值是有参考意义的，一组实验的loss的绝对值低，能说明效果是更好的。

②设备级负载loss
将专家分为D组 ${\mathcal{E}_1,\mathcal{E}_2,\dots \mathcal{E}_D}$ ，每个专家放在一个设备上，为了保证设备间的负载均衡，引入设备级负载loss。设备级负载loss比专家级粒度更大，相当于在多组专家间做负载均衡，主要用来平衡不同设备的计算负载。如以下公式所示：
在这里插入图片描述
在公式中T表示要处理的总token量，在实际模型训练中，模型是按batch接受输入的，那这个T总token量，到底是什么口径？
从V1的源码看，是以每个Batch为一组token计算负载loss的，T就是一个batch的总token量。

class MoEGate(nn.Module):
    def forward(self, hidden_states):
        bsz, seq_len, h = hidden_states.shape
        ############################
        # 这里的hidden_states就是公式里的T，是一个Batch数据的全部token做计算，每个Batch会重新计算
        ############################
        hidden_states = hidden_states.view(-1, h)

        logits = F.linear(hidden_states. self.weight, None)
        scores_for_aux = logits.softmax(dim=-1)

        topk_weight, topk_idx = torch.topk(scores_for_aux, k=self.top_k, dim=-1, sorted=False)
        topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
        mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
        ce = mask_ce.float().mean(0)
        ############################
        # 计算Pi，fi 和 aux_loss。这里的计算并没有跨Batch累积，每个Batch单独计算
        ############################
        Pi = scores_for_aux.mean(0)
        fi = ce * self.n_routed_experts
        aux_loss = (Pi * fi).sum() * self.alpha

（3）DeepSeek V2 MoE升级

①设备受限的专家路由机制
随着LLM的size越来越大，对MOE模型的训练，一般要采用专家并行来分布式加载模型，即对于网络的一个MoE层的多个专家，分配到多个设备上，来并行训练。由于deepseek的moe做了细粒度专家的设计，通常专家会很多（V2模型的路由专家数有160个，激活专家数6个）。在moe层多专家的输入是一样的，由当前层的自注意力输出的隐层激活值作为moe的输入。如果被激活的专家分布在多个机器上，那么要把输入传输到多个机器，势必会带来成倍的通讯成本。
为了解决这个问题，deepseek v2引入了设备受限的专家路由机制，具体就是保证每个token的激活专家，最多分布到M个设备上（M < TopK），这样来控制通信成本。

对于每个token，首先选择门控分数（ $s_{i,t}$ ）最高的专家所在的M个设备；
然后把M个设备上的所有专家作为备选集合，选择Top K个专家。

deepseek实际验证出，当 $\geq 3$ 时，这种受限的选TopK的操作，与不受限的全局选Top K的操作，模型效果上是大致相当的。所以在V2模型上，选择的topk=6,M=3.

②增加通信负载均衡loss
通过设备受限的路由机制可以减轻从输入侧将数据分发到多设备，减少数据在不同设备之间的分发和通信量。但在设备接收侧可能还会出现集中几个设备的专家激活的问题，导致通信拥堵问题。所以V2相对于V1增加了通信负载均衡loss
在这里插入图片描述
其中， $\mathcal{E}_i$ 表示第i个设备的一组专家，D是设备数，M是受限路由的设备数，T是一个batch的token数， $\alpha_3$ 是该辅助loss的超参。对于 $f''_i$ 的计算，乘以D再除以M也是为了保证loss不随设备的增减或限制路由配置而动态变化。

设备受限的专家路由机制和通信负载均衡loss，都是为了解决通信负载平衡的方法。不同的是：设备受限的专家路由机制是在通信分发端确保分发的一个上限；而通信负载均衡loss是在通信接收端确保接收的平衡，鼓励每个设备接收等量的token。所以通过这两种方法，可以确保设备输入、输出的通信负载均衡。

③设备级token丢弃策略
虽然多个负载均衡的loss（包括专家、设备、通信）能引导模型做出通信和计算的平衡，但并不能严格做到负载均衡，为了进一步做计算的负载均衡，引入设备级的token丢弃策略。具体做法：

首先对于一个batch输入token，算出每个设备的平均接收的token量，即设备的容量C；
对于每个设备实际分配的token量 $T_d$ ，按照路由打分降序排列；
如果 $T_d > C$ ，则将超过容量C的尾部token丢弃掉，不进行专家网络计算。

这里的丢弃token，只是在单MoE层对token不做计算，但这个token会通过残差继续传入上层Transformer网络，参与计算。所以被丢弃的token依然是由hidden_state表征的，只是这个表征不是专家输出+残差结合的结果，而是只有残差部分的结果。而且多层Transformer MoE的专家是不耦合的，在某些层可能丢弃，在另一些层参与专家计算。

作者为了保持推理和训练的一致性，在训练阶段也保持有10%的样本是不做Token丢弃的，来保证在推理阶段不做token丢弃的效果。

（3）DeepSeek V3 MoE升级

首先在基本的MoE框架上，延续了细粒度专家（finer-grained experts）和共享专家（Shared Expert Isolation）的设计。在门控网络和负载均衡方面都做了些改进。

MoE门控计算softmax->sigmoid
在这里插入图片描述
V3版相对于V2版的专家设置发生了哪些变化：

V2版：路由专家数：160，激活专家数：6，模型总参数67B，激活参数21B；
V3版：路由专家数：256，激活专家数：8，模型总参数671B，激活参数37B

V3相较于V2的路由专家数增加了近100个，我们考虑在计算一个较大维度的softmax操作，softmax要在内部对所有维度的值做归一化处理，维度越大，会趋向于计算出的每个维度的值越小，因为所有维度加和要等于1，所以维度越大，每个维度值理论上分配的值就越小，这样在选取top k个最大值时，对更小的小数位会敏感，会有数据区分度不高的问题，维度越大，问题越严重。而选择sigmoid函数，它是对每个专家分别计算一个[0,1]的打分，并不是随着专家维度的变化而变化，理论上计算的打分值域更宽，区分度更高。

②无辅助损失负载均衡
DeepSeek在V1，V2版MoE模型中，增加了专家级，设备级和设备通信级等平衡负载辅助loss。这些辅助loss只是为了做计算、通讯的负载均衡，对模型的效果调优并没有帮助。甚至这些辅助loss增加过多，loss太大会对主模型造成影响，导致主模型的效果有损。为了减轻多辅助负载均衡的loss对主模型的影响，在V3版把多辅助loss都精简掉了，通过引入一个可动态调节的bias来做到负载均衡。
具体方法：V2选择专家是通过专家的门控权重 $s_{i,t}$ 来选取TopK，V3对每个专家维护一个可动态调节的bias（ $b_i$ ），现在选择专家通过 $s_{i,t} + b_i$ 来选择topk个专家。当我们检测到专家是过载状态时，减小该专家的 $b_i$ ，来降低门控权重，从而减少路由到该专家的token量；当我们检测到专家负载不足时，我们增加该专家的bias来提升门控权重，从而增加路由到该专家的token量。

③sequence粒度的负载均衡损失
DeepSeek V3也增加了一个sequence粒度的负载均衡损失，来平衡单个sequence的token分配给每个专家。如下图公式所示
在这里插入图片描述
相对于V1版的专家级辅助损失(Expert-Level Balance Loss)其实就是作用粒度不一样，Sequence-Wise的粒度是单条样本粒度的token做计算。Expert-Level Balance是一个Batch的多Sequence的token做计算。公式的计算形式并没有什么差异。

（4）总结

V1版为了兼顾对通用知识和细粒度领域知识的建模，引入了共享专家和细粒度专家。同时为了平衡各个专家的计算负载，引入了专家级负载loss和设备级负载loss。
V2版主要在通信负载上做了优化，通过引入设备首先的专家路由机制和通信负载均衡loss确保设备输入、输出的通信负载均衡。
V3版考虑负载loss对主模型的优化会有影响，将辅助负载loss做了精简，通过在门控权重增加一个可调的bias来解决通信和计算的负载。也引入一个更细粒度的sequence负载均衡loss。同时考虑随着路由专家数增加到256个，在门控权重计算上选择了值域更宽、打分差异更显著的sigmoid函数替换了原来的softmax函数。

4、GRPO（群体相对策略优化）

GRPO可以算作是PPO的计算效率优化版本，在保持效果的同时，降低计算资源消耗。
PPO采用了Actor-Critic架构，使用了四个模型：

Policy模型（又称Actor）:输入一段上文，输出下一个token的概率分布。该模型需要训练，是我们最终得到的模型。输出下一个token即为policy模型的“行为”。
Value模型（又称Critic）：用于预估当前模型回复的总收益。该总收益不仅考虑当前token的质量，还需要衡量当前token对后续文本生成的影响。该模型需要训练。
Reward模型：事先用偏好数据进行训练，用于对Policy模型的预测进行打分，评估模型对于当前输出的及时效益。
Reference模型：与Policy模型相同，但在训练过程中不进行优化更新，用于维持模型在训练中的表现，防止在更新过程中出现较大偏差。

GRPO的方法是通过，大模型根据当前的上文输入进行多次采样，生成多个预测结果，并分别使用reward模型对这些预测结果进行评分，最后取这些评分的平均值来替代value模型的预期总收益估计。通过这种方式，GRPO在训练过程中可以减少一个模型的前向和反向传播计算，从而降低计算资源的消耗。