如何看待 DeepSeek 发布的 MoE 大模型 DeepSeek-V2？（从推理角度分析）

作者丨233

来源丨https://www.zhihu.com/question/655172528/answer/3491439374

编辑丨GiantPandaCV

计算量

DeepSeek-V2在decode阶段，只需要4K上下文就可以使Attention的计算量（只计SPDA算子，不计其中的Linear层）超过Linear层：

NoPE部分qK由于W_K与q而非latent结合，每head是512维内积；（注：q乘W_K.T的计算量归入Linear层）

RoPE部分每head是64维内积；

Attention Weight每head需要对512维latent加权求和。

因此每层需要128×(512+64+512)×4K=557056K=544M次MAC，60层总共需要31.875G MAC。

21B激活参数却只需要21G MAC。

推理成本分析

假设memory bound（但我怀疑已经compute bound，至少局部compute bound）。

DeepSeek-V2是60层，而Llama 3 70B是80层，而MLA是GQA(8 group)的9/32，因此context_size是27/128。

27/128的context确实可以实现128/27倍的batch_size，但routed expert是6/160，实际均摊带宽的token数应该要乘6/160。另一方面，attention和shared expert又是全量均摊带宽的，从这个角度计算比较麻烦。

可以干脆从另一个角度考虑：足够大batch_size下所有参数是全量激活的，那也就是128/27倍的batch_size去均摊236B参数，约等于50B参数，在平均每token的参数带宽需求相比GQA Dense 70B优势不显著。

但是：大batch_size下，带宽需求从参数占主导转变为context占主导。DeepSeek-V2是每层576维，60层，6bit量化。假设3K个请求，平均每个请求4K上下文，就需要576×60×4K×3K×0.75Byte等于135M×3K×0.75Byte=405G×0.75Byte=303.75GiBytes，已经超过FP8精度的参数。

考虑到均摊带宽需求等于weight_size / batch_size + context_size，如果第二项占主导，那么DeepSeek-V2的推理成本应该接近GQA Dense 70B的27/128（但比这更高）。

假设平均4K上下文、d_h=128和6bit量化，GQA Dense 70B（80层）双实例用70×2+640×0.75=620G VRAM可以服务1K请求（2×8×128×80×4K×1K=640M×1K=640G），单次Forward线性层的MAC为70T，qKV的MAC为8×2×320G（8 head/group）等于5T。DeepSeek-V2用236+405×0.75=539.75G VRAM（留一些余量，部分权重可能有多个副本）可服务3K请求，单次Forward线性层的MAC为63T，qKV的MAC为128×(1+512/576)×405G（128 head）等于95.625T。

如果都用H100，推测成本在1/3左右。

具体的带宽需求比较复杂：因为大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 shared expert的M=3K不太可能小于BLOCK_M（此时3K乘5K的activation比5K乘1.5K的weight大，且应是compute bound）。不过shared expert应该DP=8，但384仍然不小。

Llama 3 70B的M=1K（或者DP=2下的M=512）则不太可能小于BLOCK_M（此时应该接近或达到compute bound）。

Attention SPDA部分，DeepSeek-V2 M=128（等于总head数），Llama 3 70B M=8（等于每group的head数），L3显然可以只加载一次，DS2则有点难度。

Embedding比较特殊。

Attention中的各projection以及LM Head，则与shared expert情况类似。

93.3% KV Cache节省的由来

单层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，推理已经是Compute Bound。

推理时巨大的batch_size允许细粒度MoE（2个固定+160选6个）保持较好的计算访存比和负载均衡，而细粒度MoE反过来大幅降低了训练计算成本！

注：160选6个做D-way专家并行（Device-Limit为M）的参数访存量是1，计算量是6/160，all-to-all通讯量是M。作为对比，160个全选（相当于D-way张量并行）的参数访存量是1，计算量是1，all-to-all通讯量是D（张量并行partition策略和通讯策略很多，仅举一例）。可见细粒度MoE+专家并行相比Dense+张量并行会降低计算访存比，但也会降低通讯带宽需求。

为了保证细粒度MoE的推理效率，在训练阶段引入了许多额外设计和辅助损失，如Device-Limited Routing，Communication Balance Loss（和Device-Level Balance Loss类似，但只统计从其他device来的token），Token-Dropping等。

架构部分

MLA：(latent) x cache，把WK转置后apply到q上，因此无需计算K，算出attention weighted average (latent) x后再apply 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。

注：以上计算假设d = n_h d_h，但实际上DeepSeek-V2设定d = 5120, n_h = 128, d_h = 128。也就是说，26+16.5=42.5这个常数需要改写为26×3.2+16.5=99.7。

其它

之前也想过为什么不cache x，然后应用结合律把W_K“吸收”进W_Q（不一定要materialize复合后的矩阵，只需要把WK转置后apply到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倒不成问题（只要不materialize），attn乘X的结果乘W_V，是(n_h, d)乘(n_h, d, d_h)的Batch GEMV，主要开销依然是读取W_v，和KV Cache方法更新V的代价相似。

“吸收”但不materialize，是指仅在逻辑上融合，类似LoRA：

def W_Q_matmul_W_K_transposed(x): return (x @ W_Q) @ W_K.T

当然实际上并不是这么乘（除非把n_h这个轴vmap了），应该是

def W_Q_matmul_W_K_transposed(x):
    # W_Q: (n_h, d, d)
    # W_K: (n_h, d, d)
    # x: (d,)
    q = einsum('hij,j->hi', W_Q, x)
    return einsum('hij,hi->hj', W_K, q)

注意：如果真的materialize吸收后的矩阵（或者说张量），尺寸会变成(n_h, d, d)。即使是MLA，也会变成(n_h, d_c, d)，由于d_c是4倍d_h，内存消耗和计算量都会是4 d d，而不合并是(n_h, d_h, d)加(n_h, d_c, d_h)等于(d + 4 d_h)d。

实际的DeepSeek-V2设定d = 5120, n_h = 128, d_h = 128, d_c = 512, d_c' = 1536，其中d_c'是query的latent维度。

这样一来，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不仅需要更大的memory容量和带宽，还需要更多的FLOP。

P.S.：可能之前也没想过应用结合律避免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），情况就截然不同了，不仅比cache x更省空间，而且吸收WQ后FLOPs可能反而比原版更低——这时候还能发挥GEMM相比于Batch GEMV的优势。

当时想到cache x并应用结合律避免重结算的灵感来源可能是陈怡然老师的ReTransformer，当时还闹了个笑话，我说怎么不把WK^T和WQ直接乘起来，还要让X先乘WQ再乘WK^T，完全没注意到d_h远小于d。（P.S.：当时搜索忆阻器做Transformer的研究，发现即使最近的研究也完全没考虑推理场景的KV Cache，我还疯狂吐槽了一波）