微软打破Decoder-Only架构!大幅降低GPU内存需求,网友:把Llama3 70B弄20GB GPU上运行
西风 发自 凹非寺
量子位 | 公众号 QbitAI
微软&清华最新研究,打破GPT系列开创的Decoder-Only架构——
提出Decoder-Decoder新型架构,名为YOCO(You Only Cache Once)。
YOCO仅缓存一次键值对,可大幅降低GPU内存需求,且保留全局注意力能力。
一张图来看YOCO和标准Transformer的比较。
在处理512K上下文长度时,标准Transformer内存使用是YOCO的6.4倍,预填充延迟是YOCO的30.3倍,而YOCO的吞吐量提升到标准Transformer的9.6倍。
去年一张“大语言模型进化树”动图在学术圈疯转,模型架构还只有三大类:Decoder-Only、Encoder-Only、Encoder-Decoder。
那么这个新出的Decoder-Decoder架构到底长啥样?
嗯,如网友所言,要读的论文又增加了。
话不多说,一起来看。
打破Decoder-Only
YOCO整体架构设计如下,分为自解码器(Self-Decoder)和交叉解码器(Cross-Decoder)两部分。
具体来说,YOCO由L个块堆叠而成,其中前L/2层是自解码器,其余模块是交叉解码器。
自解码器利用高效自注意力(efficient self-attention)机制来获取键值(KV)缓存:
接收输入序列的嵌入表示,并使用高效自注意力来生成中间向量表示;使用因果掩码(causal masking)保证解码的自回归特性;自解码器的输出用于生成全局KV缓存。
而交叉解码器使用交叉注意力(cross-attention)来重用自解码器生成的共享KV缓存:
在自解码器生成的KV缓存基础上进行堆叠,以获得最终的输出向量;同样使用因果掩码来维持自回归生成;允许交叉解码器层间高效地重用KV缓存,减少了对GPU内存的需求。
总的来说,自解码器和交叉解码器的模块设计与Transformer的解码器层类似,包含交错注意力和前馈网络子层。不过,研究人员还进行了预RMSNorm、SwiGLU和分组查询注意力等改进。
两部分之间的区别在于注意力模块。
自解码器使用高效自注意力,如滑动窗口注意力(Sliding-Window Attention)或门控保留(gated retention)。
而交叉解码器使用标准的多头交叉注意力,Query向量通过注意力与自解码器产生的全局键值缓存相关联。
推理大幅度省 省 省
实验阶段,研究人员将YOCO模型与同体量的Transformer模型进行比较。
分析维度有四个:语言建模评估、与Transformer比较的可扩展性、长上下文评估、推理优势。
语言建模评估
研究人员训练了一个3B参数的YOCO语言模型,并根据训练token数量(1T和1.6T)进行评估。
在LM Eval Harness的多个下游任务上,YOCO与Transformer模型OpenLLaMA-3B-v2、StableLM-base-alpha-3B-v2、StableLM-3B-4E1T打得有来有回。
可扩展性对比
接着,研究人员在160M到13B参数规模范围内,分别训练了YOCO(门控保留和滑动窗口注意力版本)和Transformer语言模型。
对比了它们在验证集上的语言模型损失,YOCO的表现与Transformer基本持平:
结果证明YOCO在模型大小扩展方面具有很强的可扩展性。
长上下文评估
将3B的YOCO模型扩展到上下文为1M,在“大海捞针”等长序列的needle retrieval任务上,YOCO-3B-1M的准确率接近100%。
在多针检索任务上,YOCO-3B-1M的性能优于一些超3B的Transformer模型:
此外,YOCO模型在长序列上的NLL随着上下文长度的增加而一致下降,表明YOCO能够有效地利用长距离依赖信息进行语言建模:
综上,可见YOCO在性能上完全不输Transformer,关键来看YOCO在推理效率上取得的显著提升。
推理优势
研究人员评估了YOCO在GPU内存占用、prefilling延迟、吞吐量和服务容量等方面的优势,评估上下文范围为32K至1M。
如下图所示,与Transformer相比,YOCO大幅度降低了GPU内存占用,且YOCO的内存消耗随上下文长度增加,增长幅度很小。
例如,在1M长度下,整体推理内存使用量仅为12.4GB,而传统的Transformer则占用了9.38倍的GPU内存。
下面展示了token的KV缓存对GPU内存的占用情况。
YOCO模型只缓存一层全局的键值对,因此与Transformer模型相比,它需要的内存约少了L(指模型的层数)倍。
例如,YOCO模型可以使用1GB的GPU内存来处理128K token。而具有GQA的Transformer 65B大小模型,仅能支持1.6K token。
也就是说,模型越大,YOCO可以节省更多。
在预填充阶段,模型并行编码输入token。对于512K和1M长度的输入,Transformer分别需要大约180秒和300秒。Transformer的计算复杂度为O(N^2),处理长上下文需要大量的浮点运算操作。
相比之下,YOCO的预填充时间为O(N),随序列长度线性增长。
YOCO将Transformer的512K上下文预填充时间从180秒减少到不到6秒。
预填充阶段可以在进入交叉解码器之前提前退出。因此,即使对于短上下文,预填充延迟的加速至少是两倍。例如,对于32K长度,YOCO比Transformer快2.87倍。
吞吐量表示模型每秒可以处理多少个token,涵盖了预填充和生成时间。如下图所示,与Transformer相比,YOCO在不同上下文长度下实现了更高的吞吐量。
以512K查询为例,Transformer的吞吐量为4.5 token/秒,而YOCO达到了43.1token/秒,即实现了9.6倍的加速。
吞吐量提高的原因如前所述,YOCO减少了预填充所需的时间。其次,由于内存消耗减少,因此可以在推理时使用更大的批量大小,这也有助于提高吞吐量。
详细细节,感兴趣的家人们可以查看原论文。
论文链接:https://arxiv.org/abs/2405.05254
点这里👇关注我,记得标星哦~
科技前沿进展日日相见 ~
- 大模型“取长补短”新思路入选NeurIPS’24,显著优于现有路由方法,南科大港科大出品2024-11-14
- 稚晖君后宇树也来玩开源了:机器人操作数据集,采用抱抱脸LeRobot训练测试,网友:泰裤辣!2024-11-13
- 高瓴、蓝驰领投灵初智能,致力于打造业界领先通用灵巧操作智能体2024-11-13
- 智源举办2024具身与世界模型专题峰会,产学研共促技术创新与产业应用2024-11-08