DeepSeek - V2的推理成本如何计算？

DeepSeek

RoPE部分中每个head的内积为64维。Attention Weight里每个head都要对512维的latent进行加权求和。所以每层所需的MAC次数为128×(512+64+512)×4K，其结果是557056K，也就是544M，60层一共需要31.875G MAC。21B的激活参数仅需21G MAC。假设存在内存限制（但我觉得已经是计算受限了，至少局部是计算受限的）。DeepSeek - V2有60层，Llama 3 70B为80层，MLA是GQA（8组）的9/32，所以context_size为27/128。

换个角度看，batch_size足够大时所有参数会被全量激活。用128/27倍的batch_size来分摊236B参数，大约相当于50B参数，这样在平均每个token的参数带宽需求方面，和GQA Dense 70B相比优势并不明显。

均摊带宽需求为weight_size / batch_size + context_size，若第二项起主导作用，那DeepSeek - V2的推理成本应接近GQA Dense 70B的27/128（且比这个数值更高）。

专家

若皆采用H100，预估成本约为三分之一。具体的带宽需求相当复杂。大batch_size下的带宽需求不同于小batch_size，不是简单等同于全参数加载1次。而且MLA和GQA的带宽需求也不一定如MHA那样，等同于全KV cache加载1次。DeepSeek - V2（3K请求）与Llama 3 70B（1K请求）的对比情况如下：160个routed expert，平均每个expert能分到115.2，将其取整为M = 128，这有可能小于X@W的BLOCK_M，进而只加载一次参数。对于2个共享专家（2 shared expert），M = 3K不太可能小于BLOCK_M。因为此时3K乘5K的激活值（activation）比5K乘1.5K的权重（weight）大，这种情况应该是计算受限（compute bound）的。另外，共享专家应该深度并行（DP）为8，不过384这个数值仍然不算小。Llama 3 70B中，M = 1K（或者在DP = 2时M = 512），这种情况下不太可能小于BLOCK_M（此时应该接近或者达到计算受限状态）。注意SPDA部分，DeepSeek - V2中M = 128（此为总head数），Llama 3 70B里M = 8（即每group的head数）。很明显，L3只需加载一次就行，而DS2加载起来就有点困难了。Embedding具有特殊性。Attention里的各个projection以及LM Head，和shared expert的情况相近。单层KV Cache等于576除以2048，即（4.5×128）除以（2×8×128），结果为2.25除以8。层数等于60除以95。6bit量化等同于6除以16。相乘之后，结果大约是0.066612。推理时的带宽和容量需求为weight_size加上batch_size与context_size的乘积。架构经过MLA改进后，context_size显著下降，这就使得batch_size能够大幅提高。如此一来，均摊带宽需求（weight_size除以batch_size再加上context_size）就会双重大幅降低。目测，要是没有MoE，推理就已经受限于计算了。在推理过程中，较大的batch_size能让细粒度MoE（2个固定+160选6个）维持较好的计算访存比与负载均衡，反过来，细粒度MoE也大大削减了训练计算成本。160选6个做D - way专家并行（设备限制为M）时，参数访存量为1，计算量是6/160，全对全（all - to - all）通讯量是M。与之对比，160个全选（类似D - way张量并行）的参数访存量为1，计算量是1，全对全通讯量是D（张量并行的分区策略和通讯策略众多，这里仅举一个例子）。显然，细粒度的MoE加上专家并行相较于Dense加上张量并行会使计算访存比降低，不过也会减少通讯带宽需求。在训练阶段，为确保细粒度MoE（混合专家模型）的推理效率，引入了诸多额外设计与辅助损失，像Device - Limited Routing（设备受限路由）、Communication Balance Loss（通信平衡损失，与设备级平衡损失类似，不过只对来自其他设备的token进行统计）以及Token - Dropping（丢弃token）等。MLA：（潜在的）×缓存，将WK转置后作用于q，这样就不用计算K了，算出注意力加权平均（潜在的）×之后再作用WV，所以也无需计算V。要引入RoPE的话，可以看作是额外做了个RoPE MQA（再单独弄一份用了RoPE的K，各个head共享），然后把RoPE MQA的qk score添加到MLA的qk score上（这就相当于对attention weight做乘性调制）。官方论文有更好的总结。

效果非常好，（参数效率）比MHA更优。下图是同参数量对比，注意：n_h很大时，MLA参数量大概是同形状MHA的25%。MLA的参数量为n_h(d_h d_c + d_h d_c + d_h d_c' + d_R d_c') + d_c d + d_c' d + d_R d + d n_h d_h，其结果等于(n_h + 26 + 16.5)d_h d；而MHA的参数量是n_h(d_h d + d_h d + d_h d) + d n_h d_h，其结果为4n_h d_h d。

之前也曾思考，为何不将x进行缓存，接着运用结合律让W_K被W_Q吸收（不一定要实例化复合后的矩阵，只需把W_K转置后应用于q就行），从而避免重复计算。但这么推理时会遇到一个问题，那就是qK^T的计算会从计算n_h个(1, d_h)乘以(d_h, T)的矩阵乘法（批量通用矩阵向量乘法，即Batch GEMV）变成计算1个(n_h, d)乘以(d, T)的矩阵乘法（通用矩阵乘法，即GEMM）。虽然是GEMM相对于Batch GEMV，但浮点运算次数（FLOPs）会大幅增加n_h倍（通常d = n_h乘以d_h），在访存量方面也没有优势。吸收W_V不是问题（只要不实例化就行），attn乘X的结果再乘W_V，这属于(n_h, d)乘(n_h, d, d_h)的批量矩阵 - 向量乘法（Batch GEMV），主要成本还是读取W_v，这和KV Cache方法更新V的代价差不多。吸收却不实体化，即仅逻辑融合，就像LoRA那样。当然，实际并非如此相乘（除非对n_h这个轴进行vmap），而应该是。要是真的将吸收后的矩阵（也可说是张量）实例化，其尺寸会变为(n_h, d, d)。就算是MLA，也会成为(n_h, d_c, d)。鉴于d_c是d_h的4倍，这样内存消耗与计算量都会达到4dd，要是不合并的话，那就是(n_h, d_h, d)加上(n_h, d_c, d_h)，结果为(d + 4d_h)d。DeepSeek - V2实际设定为：d = 5120，n_h = 128，d_h = 128，d_c = 512，d_c' = 1536，这里的d_c'为query的潜在维度。如此，materialize W_Q@W_K.T为(n_h, d_c', d_c)=(128, 1536, 512)，W_V@W_O为(n_h, d_c, d)=(128, 512, 5120)，总的开销是(3+10)乘以2的25次方。在不materialize的情形下，W_Q为(n_h, d_c', d_h)=(128, 1536, 128)，W_K是(n_h, d_c, d_h)=(128, 512, 128)，W_V为(n_h, d_c, d_h)=(128, 512, 128)，W_O是(n_h, d_h, d)=(128, 128, 5120)，总的开销为(3 + 1 + 1 + 10)乘以2的23次方。自动化求解时，materialize不但要有更大的内存容量与带宽，而且需要更多的浮点运算次数（FLOP）。

之前可能没考虑过运用结合律来避免x cache所需的对KV重计算。之前有一回我计算x cache每token的MACs是按照(dd_h+2Tdd_h+2Td_h)n_h来算的，要是运用结合律避免重计算的话应该是(3dd_h+2Td)n_h，相比之下，kv cache是(dd_h+2Td_h)n_h。但要是像DeepSeek V2那样多想一步，先对x进行降维（这也可视为对WK和WV做低秩分解，就和LoRA类似），那情况就大不一样了。这样做不但比缓存x更节省空间，而且在吸收WQ之后，其FLOPs可能比原版还低，此时还能展现出GEMM相对于Batch GEMV的优势。当时能想到cache x并运用结合律避免重新结算，灵感或许是源于陈怡然老师的ReTransformer。那时候还出了个糗事，我纳闷为啥不直接把WK^T和WQ相乘，而非要让X先乘WQ再乘WK^T，完全没察觉到d_h要远小于d。（附言：当时搜索忆阻器做Transformer的研究，发觉即便近期的研究也完全没有考虑推理场景的KV Cache，我还狠狠吐槽了一番。）