RWKV，引领大模型架构变更的新型RNN

RWKV，引领大模型架构变更的新型RNN林玥煜 元始智能演讲嘉宾林玥煜 元始智能算法工程VP原始智能算法工程VP，曾任大数医达科技有限公司算法总监，阿里巴巴数据事业部系统架构师，多年来深耕大数据、人工智能在工业界应用和开发管理。对大语言模型在严肃医疗场景的应用、开发拥有丰富的实战经验。目 录CONTENTS1.RWKV的历史2.RWKV的架构特点3.RWKV的基础模型4.RWKV的落地场景5.RWKV的未来发展方向6.RWKV的评测结果新一代模型架构/超越TransformerRWKV推理速度、内存恒定适合长文本处理、多轮对话等只做矩阵乘矢量，无 KV CacheApache 2.0 协议计算效率高无限上下文对芯片友好全球开源开放项目历史2023.12一个人的开源项目2020年初2022.11RWKV-4 7B模型开源2023.03RWKV-4 14B模型开源2023.05RWKV-4论文公开RWKV-5 7B模型开源种子轮奇绩创坛2024.01RWKV-5和RWKV-6论文公开2024.04RWKV-6 7B 模型开源2024.052023.06成立元始智能2023.08第一个商业客户2023.102023.07RWKV-6代码开源全球开发者超过2万2023.10高通全球合作2024.2第一个toC 端侧应用RWKV要解决的问题⚠ Transformer 是死胡同算力需求巨大，Scaling-law 失效RWKV 正引领大模型的架构迁移RWKV 开始于 2020 年初，正在研发 RWKV-7架构名称作者和论文地址架构版本阶段算法复杂度最大模型参数最大训练TOKENRWKVBo PENGhttps://arxiv.org/abs/2305.13048RWKV-6商用O(N)14 B2.5 T(SlimPajama+pile+全球语言+代码)Mamba CMU，Princetonhttps://arxiv.org/abs/2312.00752接近 RWKV-6发展O(N)6.7 B0.627 T(SlimPajama)Gated Linear AttentionMIThttps://arxiv.org/abs/2312.06635接近 RWKV-6研究O(N)1.3 B0.1 TStriped Hyena Together, Stanfordhttps://arxiv.org/abs/2302.10866接近 RWKV-4.5 与 Llama2 的混合发展O(N logN) 与 O(N^2) 之间7 B1 T+xLSTMLSTM 作者https://arxiv.org/abs/2405.04517接近 RWKV-6研究O(N)1.3 B0.3 TRetNet微软亚洲研究院，清华大学https://arxiv.org/abs/2307.08621接近 RWKV-5研究O(N)6.7 B0.1 TTransnormerLLM上海人工智能实验室，OpenNLPLabhttps://arxiv.org/abs/2307.14995接近 RWKV-5发展O(N)6.8B1.4TRWKV兼具Transformer高效训练和RNN高效推理的新架构我们是怎么做的？RNN和Transformer各⾃的局限性ØRNN 在训练长序列时容易出现梯度消失问题。ØRNN 在训练过程中无法在时间维度上进行并行化，限制了其可扩展性。ØTransformer 具有二次复杂度, 长序列任务中计算成本高和占用内存多。时间和空间复杂度⽐较我们是怎么做的？RNN 结构示意图相比较，Transformer 一次处理一整句话，或一整段话，可以并行训练RNN最简单有效的基本形式•ht = αt ⊙ ht−1 + (1 − αt) ⊙ xt•RNN 一步一步执行，每次仅处理一个字或一个词•内存占用小，计算量小•对前一步结果的依赖，使得 RNN 无法并行化训练，极大限制了 RNN 的可扩展性我们是怎么做的？RWKV的由来RWKV与QKV相对，贯穿整个RWKV系列虚线代表RWKV-6中有RWKV-5中没有RVKWReceptance作为过去信息的接受程度的接受向量值（Value）类似于传统注意力中 V 的向量Weight位置权重衰减向量，可训练的模型参数键（Key）类似于传统注意力中 K 的向量时间混合模块的演进Ø RWKV-6引入了channel-wise的衰减率wtRWKV-5 time-mixing时间混合模块RWKV-6 time-mixing时间混合模块Ø RWKV-5中消除了归一化项（分母）Ø RWKV-5引入了矩阵值状态, k,v的维度从D-> (D/H, D/H)time-mixing时间混合模块的演进RWKV的RNN视角RWKV5/6写成递归形式:虽然递归形式一样，但是RWKV-5中的w是data-independent 的 ， 而 RWKV-6 中 的 w 是 data-dependent的wt虚线代表RWKV-6中有，RWKV-5中没有# r, k, v parameter shape (B,H,1,D//H)# w parameter of shape (1,H,1,D//H) for Eagle (RWKV-5),# (B,H,1,D//H) for Finch (RWKV-6)# u parameter of shape (1,H,1,D//H)def rwkv_5_or_6_recurrent(r, k, v, w, u, wkv_state): kv = k.mT @ v # x.mT is equivalent to x.transpose(-2, -1) out = r @ (wkv_state + u.mT * kv) wkv_state = w.mT * wkv_state + kv # (B,H,D//H,D//H) return out, wkv_statek vdiag w SSdiag(u) k vSwkvTttΤt×+×=××+=--11t)(代码形式:WKV模块的改进：RWKV-5: 通过学习得到的通道衰减率来替代RetNet中的静态衰减率。RWKV-6: 通过动态生成依赖于数据的token-shift量和衰减率。RWKV架构相对应传统RNN的改造1.把每⼀个Block拆成若⼲个部分，在训练/预测的时候，不互相依赖的模块可以相互并⾏计算。2.在需要状态传递的Time Mixer模块，通过CUDA/FLA扩展，在Channel Wise+Head Wise并⾏处理。由于Channel和Head的数⽬很多，通常都超过了⼀个GPU所拥有的Tensor core的数⽬，我们在Time Mixer模块也能充分利⽤GPU的并⾏计算能⼒。* 此处RWKV6非完全优化版本，后续会有更高效的RWKV6实现modelblock_size batch_size n_layer n_head n_embd parameters time(ms)GPT256166322048302.15M60.13GPT512166322048302.15M108