训练时Transformer KV Cache共享 --小路纯手工教学文章,通俗易懂 这篇适合有对LLM机制感兴趣的非AI科班的读者看,我会写的通俗易懂,任… · @Zen_with_AI · TwKit
@Zen_with_AI
训练时Transformer KV Cache共享 --小路纯手工教学文章,通俗易懂 这篇适合有对LLM机制感兴趣的非AI科班的读者看,我会写的通俗易懂,任何专业词汇都会有解释。 好,普通transformer (基于自注意力机制的神经网络架构)有L层,每一层都有自己的参数。层数越深,参数越多。参数就是我们说的1b 7b 70b那个数字。 普通 transformer 每层参数各不相同;后面讲的looped transformer 则是同一组层的参数被重复使用。 另外一个必须了解的概念,方便后面理解,是kv cache, 普遍的说,是在推理时把Key和Value缓存起来,避免每个token都重新被计算,可以大幅度提高推理速度。 普通做法是每一层有自己的kv cache. 那聪明的读者你应该想到,有一种优化就是共享kv cache给多层甚至所有层。这就是cross-layer kv sharing, 跨层共享kv. 但注意,这里讲的是推理时的kv cache共享. 一些比较有代表性的演化有 YOCO, LCKV, CLA, KVSharer, xKV等。这些名字不需要懂,知道他们是推理时kv cache共享的不同算法就可以了。 那回到这次要讲的,是训练时kv cache多层共享。也就是looped transformer,循环transformer. 训练时 transformer反复循环一组权重,每次循环是一个recursion,这样参数很小,却能获得比较深的计算,因为一直Loop. 可以理解成,一个人把一个问题反复深刻的思考几遍。 那聪明的读者你又猜到了,推理时kv cache很吃显存,传统做法是一层一个Kv cache,这里的做法是多层共享一个kv cache.循环5次就是省了4个kv cache显存空间. 说详细一点,循环第一轮的时候,保存完整cache, 后续几轮,各自之保留很小的本地窗口,然后在第一轮的kv cache上修改。好比五个人读一本书,第一个人做完整笔记,后面四个人在第一个人的笔记上修改。 这个算法的结果是非常不错的,paper作者用1b参数,循环5次的模型,context节约了76%. 同样的显卡能跑3.6倍batch. 这个算法的缺点是,循环要比普通训练多4倍,比标准transformer费算力,省了显存,费了算力,其实不是太划算。 另外就是论文测试只用了1b模型,更大的模型效果未知。 希望大家看完这个文章,初步理解了transformer, kv cache, 训练和推理的区别,还有研究员是如何在前人的优化基础上,产生新的想法去实验去设计验证的过程。 读完上文你能看着下面的逻辑链条回忆一下学到的概念然后自己能回答出最后的问题: 普通 transformer → 参数与层数 → 推理时的 KV cache 是什么 → 为什么 KV cache 吃显存 → 推理阶段已有的跨层共享思路 → looped transformer(权重复用)→ 循环带来的 KV cache 膨胀问题 → 这篇论文怎么解决的->结果和代价? 论文https://arxiv.org/pdf/2610.02383
曝光 295 · 评论 0 · 点赞 0 · 书签 2
查看原帖