资讯
展览资讯 大型展会 灯光节 大型盛典 赛事 中标捷报 产品快讯 热门话题 艺术节 活动 演出 新闻 数艺报道 俱乐部签约
观点
大咖专访 观点洞察 书籍推荐 吐槽 设计观点 企业访谈 问答 趋势创新 论文 职场方法 薪资报价 数艺专访
幕后
幕后故事 团队访谈 经验分享 解密 评测 数艺访谈
干货
设计方案 策划方案 素材资源 教程 文案资源 PPT下载 入门宝典 作品合集 产品手册 电子书 项目对接
  • 0
  • 0
  • 0

分享

推理飙升2倍!普林斯顿北大校友祭出多头「美杜莎」,33B模型与13B一样快

2023-09-12



  新智元报道  

编辑:桃子 好困
【新智元导读】大模型推理加速如何破局?普林斯顿、UIUC等机构华人团队提出一个简单框架Medusa,训练多个解码头,直接让LLM推理速度飙升2倍。

由于LLM架构固有的内存限制,使得生成又慢又贵。

对此,很多大佬都在寻找一种挽救的方法。Karpathy曾多次提出,大模型「投机采样」对于推理时间优化是一个出色的方案。

但是,尽管投机解码能够加速生成速度,但因其太过复杂,并没有被许多开发者采用。

今天,来自普林斯顿、UIUC等机构的华人团队提出了全新的简单框架:Medusa(美杜莎)。

没有额外的草稿模型,研究人员只是引入了几个额外的解码头,微调出「美杜莎头」,能够在单个A100-8G GPU,一天时间内完成训练。

结果发现,Medusa直接让模型推理加速约2倍。

Vicuna-7b与Medusa

为什么LLM生成效率低?


从系统角度来看,LLM生成遵循一种「以内存为限制」的计算模式,主要的延迟瓶颈来自内存读取/写入,而非算术计算。

这一问题的根源在于,自回归解码过程中固有的顺序性。

即每次前向传递都需要将整个模型的参数,从高带宽内存(HBM)传输到加速器的计算单元。

尽管每个样本只产生一个token,但这个操作未能充分利用现代加速器的算术计算能力,由此导致模型效率低下。

在LLM崛起之前,应对这种低效率的常见方法是,简单地「增加批大小」,从而实现更多token的并行生成。

但大模型的不断涌现,让情况变得更加复杂。

在这种情况下,增加批大小不仅会带来更高的延迟,还会大大增加Transformer模型的键-值缓存的内存需求。

此外,这种低效率也体现在「成本结构」上。

截止到2023年9月,与仅处理提示相比,GPT-4的生成成本约高出2倍,Claude 2的生成成本大约高出3倍。

研究人员主要关注的重点是,改善LLM生成的延迟,同时Medusa也可以适用于需要平衡延迟和吞吐量的LLM服务。

每个token的价格

「投机解码」是终极解决方案?


基于上述的挑战,加速文本生成的一种吸引人策略是:更高效地利用计算资源。

具体来说,通过并行处理更多的token。

这种方法,采用了一个简化的「草稿」模型,每一步都能快速生成一批token的候选项。

然后,这些候选token将通过原始的、全尺寸的语言模型进行验证,以确定最合理的文本延续。

这一基本逻辑基于一个有趣的假设:「草稿」模型虽然小,但应该足够娴熟,能够生成原始模型可以接受的序列。

如果这个假设成立,「草稿」模型可以迅速生成token序列,同时原始模型可以高效地并行验证多个token,从而最大化计算吞吐量。

最近的研究表明,通过精心调整的草稿模型,投机解码可以将延迟降低高达2.5倍,令人印象深刻。

然而,这种方法并非没有挑战:

1. 寻找理想的「草稿模型」(Draft Model):找到一个「小而强大」的草稿模型,与原始模型很好地协调,说起来容易,做起来难。

2. 系统复杂性:在一个系统中托管两个不同的模型会引入多层的复杂性,不论是计算还是操作,尤其是在分布式环境中。

3. 采样效率低:使用投机解码进行采样时,需要使用一种重要性采样方案。这会带来额外的生成开销,尤其是在较高的采样温度下。

这些复杂性和权衡限制了投机解码的广泛采用。因此,虽然投机解码前景广阔,但并未被广泛采用。

Medusa:将简单与高效融合


为了满足对更加用户友好,且功能强大的解决方案的需求,普林斯顿研究团队推出了创新框架Medusa(美杜莎)。

它不仅加速了模型的生成,甚至让LLM能够让更多人去访问和使用。

据介绍,最新方法重新回顾了「Blockwise Parallel Decoding for Deep Autoregressive Models」论文中,一个被低估的宝藏:

回溯Transformer模型的发明,与其引入一个全新的「草稿」模型来预测后续token,为什么不简单地扩展原始模型本身呢?

论文地址:https://arxiv.org/abs/1811.03115

这就是「Medusa head」(美杜莎头)发挥作用的地方。

这些额外的解码头与原始模型无缝集成在一起,在每个生成的关键点产生token块。

与草稿模型不同的是,Medusa head可以与原始模型一起训练,而原始模型在训练期间保持冻结状态。

这种方法允许研究人员在单个GPU上微调大模型,充分利用强大的基础模型学习到的表征。

此外,由于新的头仅由一个与原始语言模型头类似的层构成,因此Medusa不会增加服务系统设计的复杂性,并且适用于分布式环境。

单靠Medusa head,并不能达到将处理速度提高一倍的目标。

但这里有一个小技巧:将其与基于树状注意力机制配对使用时,就可以并行验证由Medusa head生成的多个候选项。

这样一来,Medusa head的预测能力速度提升2倍-3倍。

另外,研究人员摒弃了传统的重要性采样方案,专门为Medusa head生成创建了一种高效且高质量的替代方案。

种新方法完全绕过了采样开销,甚至进一步提升了Medusa的加速步骤。

简而言之,研究人员用一个简单的系统解决了投机解码的挑战:

1. 没有独立的模型:不是引入一个新的草稿模型,而是在同一个模型上训练多个解码头。

2. 轻松集成到现有系统中:训练参数效率高,即使是GPU性能较差的情况下也可以进行。而且由于没有额外的模型,无需调整分布式计算设置。

3. 将采样视为一种放松:放宽与原始模型分布相匹配的要求,使得「非贪心生成」甚至比「贪心解码」更快。

Medusa概述

具体来说,Medusa在LLM的最后隐藏状态之上引入了多个头,使其能够并行预测多个后续token。

在使用Medusa head扩充模型时,原始模型在训练期间被冻结,只有Medusa head经过微调。这种方法使得在单个GPU上对大型模型进行微调成为可能。

在推理过程中,每个头为其指定的位置生成多个顶级预测。这些预测被组合成候选项,并使用基于树状注意力机制并行处理。

最后一步是,使用典型接受方案选择合理的延续,被接受的最长候选项prefix将用于下一阶段的解码。

这样,Medusa通过同时接受更多token,从而减少所需的解码步骤,提高了解码过程的效率。

接下来,让我们深入了解Medusa的三个组成部分:Medusa head(美杜莎头)、tree attention(树状注意力机制)和典型接受方案。

Medusa head(美杜莎头)


那么,Medusa head究竟是什么呢?

它们类似于原始架构中的语言模型头(因果Transformer模型的最后一层),但有一个变化:它们预测多个即将出现的token,而不仅仅是下一个token。

受到块状并行解码方法的启发,研究人员将每个Medusa head作为单层前馈网络来实现,并增加了一个残差连接。

训练这些头非常简单。你可以使用训练原始模型的相同语料库,或者使用模型本身生成新的语料库。

在这个训练阶段,原始模型保持不变;只有Medusa head经过微调。

这种有针对性的训练会带来一个参数效率极高的过程,并能迅速达到收敛。

尤其是,与在推测解码方法中训练独立的草稿模型的计算复杂性相比,优势更加突出。

在研究人员测试的Vicuna模型上,Medusa head在预测next-next token方面准确率达到60%,位列第一。同时,它仍然有改进的空间。

tree attention(树状注意力机制)


测试中,团队发现一些引人注目的指标:尽管对于预测next-next token的第1名准确率大约在60%左右,但第5名的准确率却超过了80%。

这个显著的提高表明,如果可以有效利用Medusa head产生多个排名靠前的预测,就可以增加每个解码步骤生成的token数量。

为了实现这个目标,研究人员首先通过从每个Medusa head的顶级预测中,获取笛卡尔积来创建一组候选项。

然后,根据图神经网络的思路,将依赖关系图编码到注意力机制中,以便可以并行处理多个候选项。

比如,使用第一个Medusa head的前2个预测,以及第二个Medusa head的前3个预测,如下所示。

树状注意力

上图中的可视化效果,展示了使用树状注意力同时处理多个候选项的过程。

在这种情况下,第一个头的任何预测都可以与第二个头的任何预测配对,形成一个多级树结构。

这个树的每一层,都对应着一个Medusa head的预测。由此,能够产生2*3=6个候选项。

这些候选项中的每一个对应于树结构内的一个不同的分支。

为了确保每个token只访问其前置token,研究人员设计了一种注意力掩码,专门允许注意流从当前token返回到其前置token。

通过这样做,并相应地位位置编码设置位置索引,可以在不需要增加批大小的情况下,同时处理各种候选项。

研究人员还指出,一些研究也采用了非常相似的树状注意力思想。

与它们相比,最新的方法更倾向于一种更简单的树状注意力形式,在推理期间,树状模式是规则的且固定的,使得树状注意力掩码进行预处理,进一步提高了效率。

典型接受


在早期关于投机解码的研究中,「重要性采样」被用来产生与原始模型预测密切相关的各种输出结果。

然而,后来的研究表明,随着采样温度升高,这种方法的效率往往会降低。

简单来说,如果你的草稿模型和你原来的模型一样好,理想情况下,你应该接受它的所有输出,从而使这个过程变得超级高效。

然而,重要性采样可能会在中间环节,拒绝这个解决方案。

在现实世界中,我们经常调整采样温度,只是为了控制模型的创造力,不一定是为了匹配原始模型的分布。

那么为什么不把重点放在,可接受可信的候选项上呢?

对此,研究人员引入了「典型接受」方案。

从现有的截断采样(truncation sampling)中汲取灵感,普林斯顿研究人员目标是根据原始模型选择足够可能的候选项。

研究人员根据根据原始模型的预测概率设置一个阈值,如果候选项超过这个阈值,就会被接受。

用专业术语来说,研究人员采用硬阈值,以及熵相关阈值中的最小值,来决定是否像截断采样那样接受候选项。

这确保了在解码过程中,选择有意义的token和合理延续。

研究人员总是使用贪心解码接受第一个token,确保在每个步骤中至少生成一个token。最终输出是通过研究接受测试的最长序列。

这一方法的高明之处在于,适应性强。

如果你将采样温度设置为0,它就会恢复到最有效的形式——贪心解码。当温度升高时,最新研究的方法变得更加高效,可以接受更长的序列。

并且,研究人员还已经通过严格的测试证实了这一说法。

因此,从本质上讲,最新的典型接受方案提供了一种更有效的方式,让LLM进行更创造性地输出。

Llama「吐口水」可以有多快?


研究人员使用了专门针对聊天应用进行微调的Vicuna模型对Medusa进行了测试。

这些模型的大小不同,参数数量分别为7B、13B和33B。

研究的目标是,测量Medusa在实际聊天机器人环境中如何加速这些模型。

在训练Medusa head时,研究人员采用了一种简单的方法,使用了公开可用的ShareGPT数据集。这是最初用于Vicuna模型的训练数据的子集,仅进行了一个epoch的训练。

而且,关键是,根据模型参数的大小,整个训练过程可以在几小时到一天之内完成,仅需要一块A100-80G GPU。

值得注意的是,Medusa可以轻松与量化的基础模型结合使用,以减少内存需求。

研究人员利用这一优势,在训练33B模型时使用了8位量化。为了模拟实际环境,他们使用了MT bench进行评估。

凭借其简单的设计,Medusa在各种用例中,始终能够实现约2倍的绝对时间加速。

值得注意的是,通过Medusa的优化,33B参数的Vicuna模型可以像13B模型一样迅速运行。

消融实验


Medusa head的配置


在利用美杜莎头的预测能力时,研究人员灵活地选择了每个头应考虑多少个顶级候选项。

例如可以选择第一个头的前3个预测和第二个头的前2个预测。当对这些顶级候选项进行笛卡尔积运算时,就生成了6个延续选项供模型评估。

这种可配置性,是需要进行权衡的。

一方面,选择更多的顶级预测会增加模型接受生成token的可能性。另一方面,它还会增加每个解码步骤的计算开销。

为了找到最佳平衡,研究人员尝试了各种配置,并确定了最有效的设置,如附图所示。

典型接受的阈值


在典型接受方案中,一个关键的超参数,称为「阈值」,帮助研究人员根据模型自身的预测,来确定生成的token是否合理。

阈值越高,接受标准越严格,从而影响通过这种方法获得的整体加速。

研究人员通过在MT bench的2个以创造力为导向的任务上进行实验,来探讨质量和加速之间的这种权衡。

图中显示的结果表明,与贪心解码方法相比,典型接受能够加速10%。这种加速明显优于使用随机采样的投机解码方法,后者与贪心解码相比实际减慢了过程。

作者介绍


Tianle Cai(蔡天乐)


共同一作蔡天乐是普林斯顿大学的博士生,导师为Kai Li、Jason D. Lee。

就读博士期间,他在Xuezhi Wang和Denny Zhou指导下在Google DeepMind实习,还在Sébastien Bubeck和Debadeepta Dey的指导下在Microsoft Research完成了实习。

他曾在北大拿到了本科学位,主修应用数学,同时也主修计算机科学双学位,在Liwei Wang教授的指导下,开始了机器学习领域的研究。

蔡天乐的学术兴趣跨越了机器学习的广泛领域,包括优化、表征学习、架构设计(重点是Transfomer、图神经网络等),以及最近的系统架构协同设计。

Yuhong Li


共同一作Yuhong (Jesse) Li是伊利诺伊大学厄巴纳-香槟分校(UIUC)的ECE博士生,导师是Deming Chen教授。

此前,他在北京邮电大学获得了学士学位,兴趣是高效机器学习。

Zhengyang Geng(耿正阳)


Zhengyang Geng是卡内基梅隆大学(CMU)的计算机科学博士生,由J. Zico Kolter指导。

之前,他在北京大学做研究助理,由Zhouchen Lin指导。致力于识别和开发能够自组织复杂系统的结构。

Hongwu Peng


Hongwu Peng是康涅狄格大学计算机科学与工程系的博士生。

此前,他于2018年获得华中科技大学电气工程学士学位,于2020年获得阿肯色大学电气工程硕士学位。

Tri Dao


Tri Dao是生成式AI初创公司Together AI的首席科学家。2024年9月起,将出任普林斯顿大学计算机科学助理教授。

此前,他在斯坦福大学获得了计算机博士学位,导师是Christopher Ré和Stefano Ermon。

Tri Dao的研究兴趣在于机器学习和系统,重点关注:高效Transformer训练和推理;远程记忆的序列模型;紧凑型深度学习模型的结构化稀疏性。

项目鸣谢:Zhuohan Li,Shaojie Bai,Denny Zhou,Yanping Huang,stability.ai,together.ai,ChatGPT。

参考资料:
https://twitter.com/tianle_cai/status/1701272996328120408
https://sites.google.com/view/medusa-llm

阅读原文

* 文章为作者独立观点,不代表数艺网立场转载须知

本文内容由数艺网收录采集自微信公众号新智元 ,并经数艺网进行了排版优化。转载此文章请在文章开头和结尾标注“作者”、“来源:数艺网” 并附上本页链接: 如您不希望被数艺网所收录,感觉到侵犯到了您的权益,请及时告知数艺网,我们表示诚挚的歉意,并及时处理或删除。

数字媒体艺术 新媒体艺术 科技艺术

9813 举报
  0
登录| 注册 后参与评论