Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

huggingface 上 deepseek-r1 源码中采用 MLA 架构的 KV Cache 压缩存储策略的实现似乎与文中说的不一致,这是为什么?代码中似乎没实现这个大优化 #238

Open
free-stardust opened this issue Feb 2, 2025 · 0 comments

Comments

@free-stardust
Copy link

首先,感谢 DeepSeek 的开源和贡献,DeepSeek 确实很棒,所以最近也好好拜读了下这边的工作,然后恰巧最近在做一些 kv cache 压缩的工作,所以重点看了下这部分,然后发现了一些问题,来请教一下。

问题:在 DeepSeek-R1 论文中说采用了类似 DeepSeek-V3 中的 MLA 架构,文中说在这个架构中使用了低秩的方式存储 Key 和 Value,但是我看了 huggingface 中的代码,发现和论文中说的不一致,这是为什么?

代码地址:https://huggingface.co/deepseek-ai/DeepSeek-R1/blob/main/modeling_deepseek.py

详细说明:

按 DeepSeek-V3 一文中对该架构的说明,假设输入的序列长度为 n,hidden_dim 维度为 d,即为一个 n×d 的 hidden_states,如果低秩的维度为 d_c,位置编码的维度为 d_rope, 那么按文中的意思最后存储的 kv 大小为 n×d_rope + n×d_c,但是在代码中,存储的并不是这么大,实际存储的要比这个大,实际存储的是 key 为 n×d,value 为 n×d_v,即存储的 kv 大小为 n×d + n×d_v(对应文中图的话应该是 compress_kv 映射后上面没有位置编码的 k 和 v 外加有位置编码的 k 会被存储,即 n×d_rope + n×d_nope + n×d_c,其中 d_v 为一个可配置参数)。即对应文中说的 compress_kv 是实际存的 kv cache,实际并没存,并且这个的作用实际只是作为一个中间计算的过渡。

另外,从 config 中查看相关参数得知,d=7168,heads = 128,d_v_head = 128,故 d_v = 16384,这比 d 还要大,且 d_rope_head = 64,d_nope_head = 128, 则 d_k = 24576,这样一来,存储的 kv 相比 kv 同维度大小(同为24576)要小,但也仅仅是压缩了约 16.67%。

但对于论文中说的方式,存储的 kv 相关参数为 d_c = 512,d_rope_head = 64,则如果以 key 和 value 都为 24576 的未压缩状态计算,完全的 kv cache 大小为 49152n,论文中低秩存储时,存储的大小对应为 8704n,则压缩了约 82.29%。

然而,代码中的实现方式为压缩了 16.67% 的这种方式,并非为压缩了 82.29% 这种实现方式,且这种中间采用低秩计算这种方式,实际也增加了计算开销,并且如果仅仅是为了降低 value 的维度,以这个注意力层的隐维度大小,不采用这种低秩过度,直接采用 4 个矩阵映射的方式不是也可以吗?我大概算了下,这两种实现方式的权重大小基本差不多,这种采用多步计算的方式权重倒是多一点,还是说,采用这种计算方式,是为了降低计算复杂度?

以上,就是我对这块代码种实现和文中描述不一致地方的疑问,感觉 kv cache 这块在在实现上还未完全优化,但也不确定是不是我理解的有问题,所以特来请教。

@free-stardust free-stardust changed the title huggingface 上 deepseek-r1 源码中采用 MLA 架构的 KV Cache 压缩存储策略的实现似乎与文中说的不一致,这是为什么? huggingface 上 deepseek-r1 源码中采用 MLA 架构的 KV Cache 压缩存储策略的实现似乎与文中说的不一致,这是为什么?代码中似乎没实现这个大优化 Feb 2, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant