我自己的原文哦~ https://blog.51cto.com/whaosoft/12459590
#大模型~微调~用带反馈的自训练
面对当前微调大模型主要依赖人类生成数据的普遍做法,谷歌 DeepMind 探索出了一种减少这种依赖的更高效方法。大模型微调非得依赖人类数据吗?用带反馈的自训练更好
如你我所见,大语言模型(LLM)正在改变深度学习的格局,在生成人类质量的文本和解决各种语言任务方面展现出了卓越的能力。虽然业界通过对人类收集的数据进行监督微调进一步提升了在具体任务上的性能,但获取高质量人类数据却面临着重大瓶颈。这对于要解决复杂问题的任务来说尤为明显,需要大量资源和专业知识。
怎么解决呢?模型生成得合成数据是一种有潜力的替代方案,只要能保证数据的质量,就能实现可扩展性和成本效益。
虽然 LLM 能够自我评估生成的数据,但在本文中,谷歌 DeepMind 探索了一种更简单的设置,将外部标量反馈信号用作每个生成样本的质量指标。
论文地址:https://arxiv.org/pdf/2312.06585.pdf
为了研究在模型生成数据上的训练,研究者考虑了一种简单但强大的语言模型自训练方法,仅需要两项功能,一是基于模型生成样本,二是利用评分机制对这些样本进行评估。
为了确保清晰度和一致性,研究者采用了一种强化自训练方法 ReST^𝐸𝑀,并证明该方法可以将期望最大化(expectation-maximization,EM)用于强化学习。具体来讲,ReST^𝐸𝑀在期望和最大化步骤之间交替进行。
- 生成(E-step):语言模型为每个输入上下文生成多个输出样本,然后使用二元奖励过滤这些样本以收集训练数据集。
- 改进(M-step):原始语言模型在来自前一个 E-step 的训练数据集上进行监督微调,然后在下一个 E-step 中使用。
研究者证实,ReST^𝐸𝑀及变体在增强各个领域的语言模型方面取得了成功,包括机器翻译、语义分析、偏好对齐和基础推理。
此外,以往工作主要将 ReST^𝐸𝑀用于相对较小的模型(最高 70 亿参数),对于较大模型的可扩展性受限。因此,本文旨在探究模型生成的合成数据与人类生成的数据在以下两个具有挑战性但研究较少领域的有效性和可扩展性,这两个领域分别是竞争水平数学解题(MATH)和代码生成(APPS)。
实证结果表明,当将 ReST^𝐸𝑀用于不同规模的 PaLM 2 模型时,在数学推理和代码生成任务中实现了显著的能力改进。与在人类编写数据上训练的模型相比,在模型生成的合成数据上微调的模型取得了更大的性能增益。有趣的是,超过了一定数量的 ReST^𝐸𝑀 迭代后,性能会降低,这表明了在少量训练问题上可能会出现过拟合。
此外,使用 ReST^𝐸𝑀微调的模型提升了 pass@k 指标和多数投票性能。这些微调后的模型在相关但 held-out 的基准上也表现出了性能增强,包括数学题(GSM8K 和 Hungarian HS finals)、编码(HumanEval)和 Big-Bench Hard 任务。
总之,本文研究结果表明,具有反馈的自训练是减少对人类数据依赖的一种有潜力的方法。
用于强化自训练的期望最大值(EM)
首先,该研究基于 Dayan 和 Hinton 之前的研究,用语言模型描述了基于 EM 的强化学习框架。具体而言,他们先是定义了一个二进制最优变量 O,使得𝑝(𝑂= 1|𝒙,𝒚)∝𝑓(𝑟(𝒙,𝒚));然后对非递减函数 𝑓 : ℝ → ℝ+ ,实现最大化观察𝑂= 1(获得高奖励),得到如下公式:
然而,求解上式中的序列 𝒚 的和很棘手。因而本文考虑相对于参数 𝜃 和变分分布 𝑞( 𝑦|𝑥) 最大化其 ELBO 𝐿( 𝑝𝜃, 𝑞),而不是最大化 log 𝑝(𝑂 = 1; 𝒙)。具体来说:
公式(2)中的 EM 算法在 E-step(Expectation) 和 M-step(Maximization)之间交替进行。
ReST^𝐸𝑀:受 EM 框架的启发,接下来论文讨论了 Gulcehre 等人提出的 ReST 方法的简化版本。为了清楚起见,本文将这种方法称为 ReST^𝐸𝑀,它将 RL pipeline 中的数据收集 (E-step) 和策略优化 (M-step) 进行解耦。如算法 1 所示:
实验和分析
本文进行实验的主要目标是回答以下问题:
- 与人类生成的数据进行微调相比,ReST^𝐸𝑀的效果如何?
- 需要多少次迭代才能获得最佳性能?ReST^𝐸𝑀多长时间会导致训练集过度拟合?
- ReST^𝐸𝑀如何影响 pass@k 和多数投票表现?
- 如果用户在特定任务上使用模型生成的数据进行微调,是否会迁移到其他任务上?在广泛的任务中评估本文的微调模型时,与基本模型相比,性能是否会下降?
- 大约需要多少输入数据才能从 ReST^𝐸𝑀 获得大部分性能提升?ReST^𝐸𝑀的一次迭代是否足够?
该研究使用 PaLM 2 模型和 Google Cloud 上的公共 API 进行实验,包括 PaLM 2-S (Bison)、PaLM 2-S* (Codey) 和 PaLM 2-L (Unicorn)。训练数据集采用 MATH 数据集和 APPS 数据集。
图 2 和图 3 分别显示了 ReST^𝐸𝑀在 MATH 和 APPS 数据集上训练的性能。可以得出 MATH 受益于 ReST^𝐸𝑀 的多次迭代,无论是在 MATH 测试集上的性能还是迁移到 GSM8K 方面。另一方面可以看到 APPS 的大部分收益来自第一次迭代,而执行更多次迭代会导致 APPS 和 HumanEval 的性能下降。
训练和测试性能的差距。图 4 显示,虽然训练集性能随着 ReST^𝐸𝑀迭代次数线性增加,但测试集性能却没有。对于 MATH,第一次迭代后测试性能改进很小,而对于 APPS,在第二次迭代中观察到性能回归。该研究猜测性能的回归可能是由于过度拟合造成的。由于 APPS 数据集的大小约为 MATH 数据集的三分之一,因此它更容易受到此问题的影响。
图 5 显示了 Palm-2-L 模型在 pass@K 指标上的性能。结果显示,微调后获得的 ReST^𝐸𝑀 模型对于所有 K 值都更强,其中性能差距通常在 K=1 时最大。
#大模型训练~数据并行
大模型场景里巨大的存储和GPU间通讯量是系统设计时需要考虑的重点,本文递进介绍了三种主流数据并行的实现方法:DP、DD皮、ZeRo。
当模型太大,一块GPU放不下时,流水线并行将模型的不同层放到不同的GPU上,通过切割mini-batch实现对训练数据的流水线处理,提升GPU计算通讯比。同时通过re-materialization机制降低显存消耗。
但在实际应用中,流水线并行并不特别流行,主要原因是模型能否均匀切割,影响了整体计算效率,这就需要算法工程师做手调。因此,今天我们来介绍一种应用最广泛,最易于理解的并行范式:数据并行。
数据并行的核心思想是:在各个GPU上都拷贝一份完整模型,各自吃一份数据,算一份梯度,最后对梯度进行累加来更新整体模型。理念不复杂,但到了大模型场景,巨大的存储和GPU间的通讯量,就是系统设计要考虑的重点了。在本文中,我们将递进介绍三种主流数据并行的实现方式:
DP(Data Parallelism):最早的数据并行模式,一般采用参数服务器(Parameters Server)这一编程框架。实际中多用于单机多卡
DDP(Distributed Data Parallelism):分布式数据并行,采用Ring AllReduce的通讯方式,实际中多用于多机场景
ZeRO:零冗余优化器。由微软推出并应用于其DeepSpeed框架中。严格来讲ZeRO采用数据并行+张量并行的方式,旨在降低存储。
数据并行(DP)整体架构
一个经典数据并行的过程如下:
若干块计算GPU,如图中GPU0~GPU2;1块梯度收集GPU,如图中AllReduce操作所在GPU。
在每块计算GPU上都拷贝一份完整的模型参数。
把一份数据X(例如一个batch)均匀分给不同的计算GPU。
每块计算GPU做一轮FWD和BWD后,算得一份梯度G。
每块计算GPU将自己的梯度push给梯度收集GPU,做聚合操作。这里的聚合操作一般指梯度累加。当然也支持用户自定义。
梯度收集GPU聚合完毕后,计算GPU从它那pull下完整的梯度结果,用于更新模型参数W。更新完毕后,计算GPU上的模型参数依然保持一致。
聚合再下发梯度的操作,称为AllReduce。
前文说过,实现DP的一种经典编程框架叫“参数服务器”,在这个框架里,计算GPU称为Worker,梯度聚合GPU称为Server。在实际应用中,为了尽量减少通讯量,一般可选择一个Worker同时作为Server。比如可把梯度全发到GPU0上做聚合。需要再额外说明几点:
1个Worker或者Server下可以不止1块GPU。
Server可以只做梯度聚合,也可以梯度聚合+全量参数更新一起做
在参数服务器的语言体系下,DP的过程又可以被描述下图:
通讯瓶颈与梯度异步更新
DP的框架理解起来不难,但实战中确有两个主要问题:
存储开销大。每块GPU上都存了一份完整的模型,造成冗余。关于这一点的优化,我们将在后文ZeRO部分做讲解。
通讯开销大。Server需要和每一个Worker进行梯度传输。当Server和Worker不在一台机器上时,Server的带宽将会成为整个系统的计算效率瓶颈。
我们对通讯开销再做详细说明。如果将传输比作一条马路,带宽就是马路的宽度,它决定每次并排行驶的数据量。例如带宽是100G/s,但每秒却推给Server 1000G的数据,消化肯定需要时间。那么当Server在搬运数据,计算梯度的时候,Worker们在干嘛呢?当然是在:
人类老板不愿意了:“打工系统里不允许有串行存在的任务!”,于是梯度异步更新这一管理层略诞生了。
上图刻画了在梯度异步更新的场景下,某个Worker的计算顺序为:
在第10轮计算中,该Worker正常计算梯度,并向Server发送push&pull梯度请求。
但是,该Worker并不会实际等到把聚合梯度拿回来,更新完参数W后再做计算。而是直接拿旧的W,吃新的数据,继续第11轮的计算。这样就保证在通讯的时间里,Worker也在马不停蹄做计算,提升计算通讯比。
当然,异步也不能太过份。只计算梯度,不更新权重,那模型就无法收敛。图中刻画的是延迟为1的异步更新,也就是在开始第12轮对的计算时,必须保证W已经用第10、11轮的梯度做完2次更新了。
参数服务器的框架下,延迟的步数也可以由用户自己决定,下图分别刻划了几种延迟情况:
(a) 无延迟
(b) 延迟但不指定延迟步数。也即在迭代2时,用的可能是老权重,也可能是新权重,听天由命。
(c) 延迟且指定延迟步数为1。例如做迭代3时,可以不拿回迭代2的梯度,但必须保证迭代0、1的梯度都已拿回且用于参数更新。
总结一下,异步很香,但对一个Worker来说,只是等于W不变,batch的数量增加了而已,在SGD下,会减慢模型的整体收敛速度。异步的整体思想是,比起让Worker闲着,倒不如让它多吃点数据,虽然反馈延迟了,但只要它在干活在学习就行。
batch就像活,异步就像画出去的饼,且往往不指定延迟步数,每个Worker干越来越多的活,但模型却没收敛取效,这又是刺伤了哪些打工仔们的心(狗头
分布式数据并行(DDP)
受通讯负载不均的影响,DP一般用于单机多卡场景。因此,DDP作为一种更通用的解决方案出现了,既能多机,也能单机。DDP首先要解决的就是通讯问题:将Server上的通讯压力均衡转到各个Worker上。实现这一点后,可以进一步去Server,留Worker。
前文我们说过,聚合梯度 + 下发梯度这一轮操作,称为AllReduce。接下来我们介绍目前最通用的AllReduce方法:Ring-AllReduce。它由百度最先提出,非常有效地解决了数据并行中通讯负载不均的问题,使得DDP得以实现。
Ring-AllReduce
如下图,假设有4块GPU,每块GPU上的数据也对应被切成4份。AllReduce的最终目标,就是让每块GPU上的数据都变成箭头右边汇总的样子。
Ring-ALLReduce则分两大步骤实现该目标:Reduce-Scatter和All-Gather。
- Reduce-Scatter
定义网络拓扑关系,使得每个GPU只和其相邻的两块GPU通讯。每次发送对应位置的数据进行累加。每一次累加更新都形成一个拓扑环,因此被称为Ring。看到这觉得困惑不要紧,我们用图例把详细步骤画出来。
一次累加完毕后,蓝色位置的数据块被更新,被更新的数据块将成为下一次更新的起点,继续做累加操作。
3次更新之后,每块GPU上都有一块数据拥有了对应位置完整的聚合(图中红色)。此时,Reduce-Scatter阶段结束。进入All-Gather阶段。目标是把红色块的数据广播到其余GPU对应的位置上。
- All-Gather
如名字里Gather所述的一样,这操作里依然按照“相邻GPU对应位置进行通讯”的原则,但对应位置数据不再做相加,而是直接替换。All-Gather以红色块作为起点。
以此类推,同样经过3轮迭代后,使得每块GPU上都汇总到了完整的数据,变成如下形式:
建议读者们手动推一次,加深理解。
Ring-AllReduce通讯量分析
在多Server的模式下,进一步,每个Server可以只负责维护和更新某一块梯度(也可以某块梯度+参数一起维护),此时虽然每个Server仍然需要和所有Worker通讯,但它的带宽压力会小非常多。经过调整设计后,依然可以用来做DDP。虽然这篇文章是用递进式的方式来介绍两者,但不代表两者间一定要决出优劣。我想表达的观点是,方法是多样性的。 对参数服务器有兴趣的朋友,可以阅读参考的第1个链接。
最后,请大家记住Ring-AllReduce的方法,因为在之后的ZeRO,Megatron-LM中,它将频繁地出现,是分布式训练系统中重要的算子。
总结
1、在DP中,每个GPU上都拷贝一份完整的模型,每个GPU上处理batch的一部分数据,所有GPU算出来的梯度进行累加后,再传回各GPU用于更新参数
2、DP多采用参数服务器这一编程框架,一般由若个计算Worker和1个梯度聚合Server组成。Server与每个Worker通讯,Worker间并不通讯。因此Server承担了系统所有的通讯压力。基于此DP常用于单机多卡场景。
3、异步梯度更新是提升计算通讯比的一种方法,延迟更新的步数大小决定了模型的收敛速度。
4、Ring-AllReduce通过定义网络环拓扑的方式,将通讯压力均衡地分到每个GPU上,使得跨机器的数据并行(DDP)得以高效实现。
5、DP和DDP的总通讯量相同,但因负载不均的原因,DP需要耗费更多的时间搬运数据
由微软开发的ZeRO(零冗余优化),它是DeepSpeed这一分布式训练框架的核心,被用来解决大模型训练中的显存开销问题。ZeRO的思想就是用通讯换显存。 如果初读ZeRO,觉得它逻辑跳跃,晦涩难懂,那么下文或许可以帮到你~
存储消耗
存储分类
首先,我们来看在大模型训练的过程中,GPU都需要存什么内容。
存储主要分为两大块:Model States和Residual StatesModel States指和模型本身息息相关的,必须存储的内容,具体包括:
- optimizer states:Adam优化算法中的momentum和variance
- gradients:模型梯度
- parameters:模型参数W
Residual States指并非模型必须的,但在训练过程中会额外产生的内容,具体包括:
- activation:激活值。在流水线并行中我们曾详细介绍过。在backward过程中使用链式法则计算梯度时会用到。有了它算梯度会更快,但它不是必须存储的,因为可以通过重新做Forward来算它。
- temporary buffers: 临时存储。例如把梯度发送到某块GPU上做加总聚合时产生的存储。
- unusable fragment memory:碎片化的存储空间。虽然总存储空间是够的,但是如果取不到连续的存储空间,相关的请求也会被fail掉。对这类空间浪费可以通过内存整理来解决。
精度混合训练
在分析这个问题前,我们需要来了解精度混合训练。
对于模型,我们肯定希望其参数越精准越好,也即我们用fp32(单精度浮点数,存储占4byte)来表示参数W。但是在forward和backward的过程中,fp32的计算开销也是庞大的。那么能否在计算的过程中,引入fp16或bf16(半精度浮点数,存储占2byte),来减轻计算压力呢?于是,混合精度训练就产生了,它的步骤如下图:
- 存储一份fp32的parameter,momentum和variance(统称model states)
- 在forward开始之前,额外开辟一块存储空间,将fp32 parameter减半到fp16 parameter。
- 正常做forward和backward,在此之间产生的activation和gradients,都用fp16进行存储。
- 用fp16 gradients去更新fp32下的model states。
- 当模型收敛后,fp32的parameter就是最终的参数输出。
通过这种方式,混合精度训练在计算开销和模型精度上做了权衡。如果不了解fp32,fp16和bf16的细节也没关系,不影响下文的阅读。只要记住它们所占的存储空间和精度表达上的差异即可。
存储大小
另外,这里暂不将activation纳入统计范围,原因是:
- activation不仅与模型参数相关,还与batch size相关
- activation的存储不是必须的。存储activation只是为了在用链式法则做backward的过程中,计算梯度更快一些。但你永远可以通过只保留最初的输入X,重新做forward来得到每一层的activation(虽然实际中并不会这么极端)。
- 因为activation的这种灵活性,纳入它后不方便衡量系统性能随模型增大的真实变动情况。因此在这里不考虑它,在后面会单开一块说明对activation的优化。
ZeRO-DP
知道了什么东西会占存储,以及它们占了多大的存储之后,我们就可以来谈如何优化存储了。
注意到,在整个训练中,有很多states并不会每时每刻都用到,举例来说;
- Adam优化下的optimizer states只在最终做update时才用到
- 数据并行中,gradients只在最后做AllReduce和updates时才用到
- 参数W只在做forward和backward的那一刻才用到
- 诸如此类
所以,ZeRO想了一个简单粗暴的办法:如果数据算完即废,等需要的时候,我再想办法从个什么地方拿回来,那不就省了一笔存储空间吗?
沿着这个思路,我们逐一来看ZeRO是如何递进做存储优化的。
(3)得到完整梯度G,就可以对W做更新。我们知道W的更新由optimizer states和梯度共同决定。由于每块GPU上只保管部分optimizer states,因此只能将相应的W(蓝色部分)进行更新。(2)和(3)可以用下图表示:
此时,数据并行的整体流程如下:
(1)每块GPU上存一份完整的参数W。将一个batch的数据分成3份,每块GPU各吃一份,做完一轮foward和backward后,算得一份完整的梯度(下图中绿色+白色)。
再次比对下显存和通讯量:
和朴素DP相比,存储降了8倍,单卡通讯量持平,好像更牛皮了呢!那么,还可以优化吗?
数据并行的流程如下:
(1)每块GPU上只保存部分参数W。将一个batch的数据分成3份,每块GPU各吃一份。
到这一步,我们用1.5倍的通讯开销,换回近120倍的显存。只要梯度计算和异步更新做的好,通讯时间大部分可以被计算时间隐藏,因此这样的额外通讯开销,也是划算的。
到这里,我们可以放出原始论文中的说明图了,经过以上分析,这张说明图是不是瞬间就能看懂了。不得不吐槽下,虽然ZeRO的设计不复杂,但对应论文写得真是逻辑跳跃,晦涩难懂....
仔细一想,ZeRO其实掌握了降本增效的精髓:用完即弃,需要再补。反正我补一个和你差不多的,也不会花费很多通(找)讯(人)时间,还大大降低了我的成本。模型的每一层多算(造)几(轮)遍(子)有啥关系呢,反正在我的预算里每个人都一刻不停地干活,就行啦!
ZeRO VS 模型并行
知道模型并行的朋友,可能会想,既然ZeRO都把参数W给切了,那它应该是个模型并行呀?为什么要归到数据并行里呢?
其实ZeRO是模型并行的形式,数据并行的实质。
模型并行,是指在forward和backward的过程中,我只需要用自己维护的那块W来计算就行。即同样的输入X,每块GPU上各算模型的一部分,最后通过某些方式聚合结果。
但对ZeRO来说,它做forward和backward的时候,是需要把各GPU上维护的W聚合起来的,即本质上还是用完整的W进行计算。它是不同的输入X,完整的参数W,最终再做聚合。
因为下一篇要写模型并行Megatron-LM,因此现在这里罗列一下两者的对比。
ZeRO-R
说完了以上对model states的显存优化,现在来看对residual states的优化。
在前文提过,设置机制,对碎片化的存储空间进行重新整合,整出连续的存储空间。防止出现总存储足够,但连续存储不够而引起的存储请求fail
ZeRO-Offload与ZeRO-Infinity
最后,简单介绍一下ZeRO-Offload。它的核心思想是:显存不够,内存来凑。如果我把要存储的大头卸载(offload)到CPU上,而把计算部分放到GPU上,这样比起跨机,是不是能既降显存,也能减少一些通讯压力呢? ZeRO-Offload的做法是:
- forward和backward计算量高,因此和它们相关的部分,例如参数W(fp16),activation,就全放入GPU。
- update的部分计算量低,因此和它相关的部分,全部放入CPU中。例如W(fp32),optimizer states(fp32)和gradients(fp16)等。
具体切分如下图:
ZeRO-infinity也是同理,它们在解决的事情都是:找个除GPU之外的地方,存数据。感兴趣的朋友可以深入研究,这里就不展开了。
#大模型训练~显卡
为什么用 A100 不用 4090? 结论,大模型的训练用 4090 是不行的,但推理(inference/serving)用 4090 不仅可行,在性价比上还能比 H100 稍高。4090 如果极致优化,性价比甚至可以达到 H100 的 2 倍。
事实上,H100/A100 和 4090 最大的区别就在通信和内存上,算力差距不大。
NVIDIA 的算力表里面油水很多,比如 H100 TF16 算力写的是 1979 Tflops,但那是加了 sparsity(稀疏)的,稠密的算力只有一半;4090 官方宣传 Tensor Core 算力高达 1321 Tflops,但那是 int8 的,FP16 直只有 330 Tflops。这篇文章的第一版就是用了错的数据,H100 和 4090 的数据都用错了,得到的结论非常离谱。
H100 这个售价其实是有 10 倍以上油水的。
2016 年我在 MSRA 的时候,见证了微软给每块服务器部署了 FPGA,把 FPGA 打到了沙子的价格,甚至成为了供应商 Altera 被 Intel 收购的重要推手。2017 年我还自己挖过矿,知道什么显卡最划算。后来在华为,我也是鲲鹏、昇腾生态软件研发的核心参与者。因此,一个芯片成本多少,我心里大概是有数的。
鲲鹏的首席架构师夏 Core 有一篇知名文章《谈一下英伟达帝国的破腚》,很好的分析了 H100 的成本:
据说微软和 OpenAI 包下了 H100 2024 年产能的一半,猜猜他们会不会发挥当年跟 Altera 砍价的传统艺能?会真的花 $40,000 * 500,000 = 200 亿美金去买卡?
可以说,H100 就像是中国一线城市的房子,本身钢筋水泥不值多少钱,房价完全是被供求关系吹起来的。我在 LA 已经住了两周,公司租的房子使用面积是我北京房子的 4 倍,但售价只贵了 30%,还带个小院,相当于单位面积的房价是北京的 1/3。我跟本地的老外聊天,他们都很吃惊,你们的平均收入水平比 LA 低这么多,怎么买得起北京的房子的?
问题来了,如果 4090 这么香的话,为啥大家还要争着买 H100,搞得 H100 都断货了?甚至 H100 都要对华禁售,搞出个 H800 的阉割版?
大模型训练为什么不能用 4090
GPU 训练性能和成本对比
LambdaLabs 有个很好的 GPU 单机训练性能和成本对比,在此摘录如下。
首先看吞吐量,看起来没有什么违和的,在单卡能放下模型的情况下,确实是 H100 的吞吐量最高,达到 4090 的两倍。看算力和内存也能看出来,H100 的 FP16 算力大约是 4090 的 3 倍,内存带宽是 3.35 倍,训练过程中由于 batch size 比较大,大多数算子是 compute bound(计算密集型),少数算子是 memory bound(内存密集型),这个结果是不意外的。
LambdaLabs PyTorch 单卡训练吞吐量对比图
LambdaLabs PyTorch 单卡训练吞吐量对比表
然后看性价比,就有意思了,原来排在榜首的 H100 现在几乎垫底了,而且 4090 和 H100 的差距高达接近 10 倍。这就是因为 H100 比 4090 贵太多了。
由于 H100 货源紧张,云厂商的 H100 租用价格就更黑了,按照标价大约 7 个月就可以回本。就算大客户价能便宜一半,一年半也足够回本了。
在价格战中过惯了苦日子的 IaaS 云服务商看到这样的 H100 回本速度,估计要感叹,这真是比区块链回本还快呐。
LambdaLabs PyTorch 单卡训练单位成本吞吐量对比图
LambdaLabs PyTorch 单卡训练单位成本吞吐量对比表
大模型训练的算力需求
既然 4090 单卡训练的性价比这么高,为啥不能用来做大模型训练呢?抛开不允许游戏显卡用于数据中心这样的许可证约束不谈,从技术上讲,根本原因是大模型训练需要高性能的通信,但 4090 的通信效率太低。
大模型训练需要多少算力?训练总算力(Flops)= 6 * 模型的参数量 * 训练数据的 token 数。
我今年初第一次看到有人煞有介事地讲这个公式的时候,觉得这不是显然的吗?又看到 OpenAI 的高级工程师能拿 90 多万美金的年薪,顿时整个人都不好了,还是 AI 香呀。之前我也面试过一些做 AI 的工程师,包括一些做 AI 系统优化的专家,连 Q、K、V 是啥都说不清楚,LLaMA 每个 tensor 的大小也算不出来,就这样还能拿到 offer。
APNet 2023 panel 的主题是 Network, AI, and Foundational Models: Opportunties and Challenges。前面几个问题都中规中矩的,panelists 有点放不开,我就提了一个问题,网络历史上的重要成就基本上都基于对应用场景深刻的理解,但我们现在做网络的很多都不了解 AI,甚至连每个 tensor 的大小和每个 step 传输的数据量都不知道,如何让 network community 更了解 AI 呢?
这下热闹了,台下的谭博首先发言,说我在华为肯定能知道所有这些东西;然后传雄老师也跟了一句,要是做网络的懂了太多 AI,那可能他就变成一个 AI guy 了。接着主持人陈凯教授问,你们有谁真的训练过大模型?沉默了一会儿,阿里的兄弟先说,我算是半个训练过大模型的,我们做的东西是支撑阿里大模型 infra 的。后面又有 panelist 说,做 AI 系统的网络优化是否有必要自己懂 AI 呢,是不是只要会做 profiling 就行了?
我个人观点仍然是,AI 并不难学,要想做好 AI 系统优化,可以不懂 attention 的 softmax 里面为什么要除以 sqrt(d_k),但不能不会计算模型所需的算力、内存带宽、内存容量和通信数据量。Jeff Dean 就有个很有名的 Numbers Every Programmer Should Know,数量级的估算对任何系统优化来说都很关键,不然根本不知道瓶颈在哪里。
回到大模型训练所需的总算力,其实很简单,6 * 模型的参数量 * 训练数据的 token 数就是所有训练数据过一遍所需的算力。这里的 6 就是每个 token 在模型正向传播和反向传播的时候所需的乘法、加法计算次数。
一堆矩阵相乘,简单来想就是左边若干个神经元,右边若干个神经元,组成一个完全二分图。选出其中任意一个左边的神经元 l 和右边的神经元 r,正向传播的时候:
- l 把它的输出乘上 l 和 r 之间的权重 w,发给 r;
- r 不可能只连一个神经元吧,总要把多个 l 的加到一起,这就是 reduce,需要一次加法。
反向传播的时候:
- r 把它收到的梯度乘上 l 和 r 之间的权重 w,发给 l;
- l 也不可能只连一个 r,需要把梯度 reduce 一下,做个加法;
- 别忘了权重 w 需要更新,那就要计算 w 的梯度,把 r 收到的梯度乘上 l 正向传播的输出(activation);
- 一个 batch 一般有多个 sample,权重 w 的更新需要把这些 sample 的梯度加到一起。
一共 3 次乘法,3 次加法,不管 Transformer 多复杂,矩阵计算就是这么简单,其他的向量计算、softmax 之类的都不是占算力的主要因素,估算的时候可以忽略。
想起来我 2019 年刚加入 MindSpore 团队的时候,领导让我开发一个正向算子的反向版本,我求导给求错了,搞得算子的计算结果总是不对,还以为是我们的编译器出 bug 了。当发现求导求错的时候,领导像以为我没学过微积分一样看着我,确实我的微积分学的不好,这也是我从数学专业转到计算机专业的原因之一。
在 MindSpore 的时候,自动微分一共就不到 1000 行代码,按照微分公式递归计算下去就行了,但自动微分作为一个重要特性被吹了半天,我都感觉不好意思了。
模型的参数量和训练数据的 token 数之间也有个比例关系,这也很容易理解,只要把模型想象成数据的压缩版本就行了,压缩比总是有极限的。模型的参数量太小,就吃不下训练数据里面所有的知识;模型的参数量如果大于训练数据的 token 数,那又浪费,还容易导致 over-fitting。
训练 LLaMA-2 70B 需要多少张卡
有了模型训练所需的总算力,除以每个 GPU 的理论算力,再除以 GPU 的有效算力利用比例,就得到了所需的 GPU-hours,这块已经有很多开源数据。LLaMA 2 70B 训练需要 1.7M GPU hours(A100),要是用 1 个 GPU,那得算 200 年。要在一个月这种比较能接受的时间周期内训练出来,就得至少有 2400 块 A100。
如果用 4090,单卡 FP16 算力是跟 A100 差不多(330 vs 312 Tflops),但是内存带宽比 A100 低一半(1 vs 2 TB/s),内存容量更是差好几倍(24 vs 80 GB),计算梯度时需要使用的 TF32 算力也低一半(83 vs 156 Tflops),综合起来 4090 单卡的训练速度还比 A100 稍低(参考前面 LambdaLabs 的评测)。
就按照 2048 块 4090 算吧,这 2048 块 4090 之间的通信就成了最大的问题。
为什么?一般有 tensor parallelism、pipeline parallelism、data parallelism 几种并行方式,分别在模型的层内、模型的层间、训练数据三个维度上对 GPU 进行划分。三个并行度乘起来,就是这个训练任务总的 GPU 数量。
三种并行方式从三个维度划分计算空间的示意图,来源:DeepSpeed
Data parallelism(数据并行)
数据并行是最容易想到的并行方式。每个 GPU 分别计算不同的输入数据,计算各自的梯度(也就是模型参数的改变量),再把梯度汇总起来,取个平均值,广播给各个 GPU 分别更新。
Data Parallelism 示意图,来源:Colossal AI
但只用数据并行是肯定不行的,因为一块 GPU 放不下整个 LLaMA 70B 模型。
就模型训练需要多少 GPU 内存,我发现能算清楚的人就不多。有的人甚至以为只需要把模型的参数和反向传播的梯度存下来就够了。事实上,训练需要的内存包括模型参数、反向传播的梯度、优化器所用的内存、正向传播的中间状态(activation)。
优化器所用的内存其实也很简单,如果用最经典的 Adam 优化器,它需要用 32 位浮点来计算,否则单纯使用 16 位浮点来计算的误差太大,模型容易不收敛。因此,每个参数需要存 4 字节的 32 位版本(正向传播时用 16 位版本,优化时用 32 位版本,这叫做 mixed-precision),还需要存 4 字节的 momentum 和 4 字节的 variance,一共 12 字节。如果是用类似 SGD 的优化器,可以不存 variance,只需要 8 字节。
正向传播的中间状态(activation)是反向传播时计算梯度必需的,而且跟 batch size 成正比。Batch size 越大,每次读取模型参数内存能做的计算就越多,这样对 GPU 内存带宽的压力就越小。可是不要忘了,正向传播的中间状态数量是跟 batch size 成正比的,GPU 内存容量又会成为瓶颈。
大家也发现正向传播中间状态占的内存太多了,可以玩一个用算力换内存的把戏,就是不要存储那么多梯度和每一层的正向传播的中间状态,而是在计算到某一层的时候再临时从头开始重算正向传播的中间状态,这样这层的正向传播中间状态就不用保存了。如果每一层都这么干,那么就只要 2 个字节来存这一层的梯度。但是计算中间状态的算力开销会很大。因此实际中一般是把整个 Transformer 分成若干组,一组有若干层,只保存每组第一层的中间状态,后面的层就从该组第一层开始重新计算,这样就平衡了算力和内存的开销。
如果还是算不清楚,可以读读这篇论文:Reducing Activation Recomputation in Large Transformer Models。
当然有人说,GPU 内存放不下可以换出到 CPU 内存,但是就目前的 PCIe 速度,换出到 CPU 内存的代价有时候还不如在 GPU 内存里重算。如果是像 Grace Hopper 那种极高带宽的统一内存,那么换入换出倒是一个不错的主意,不管训练的正向传播中间状态还是 KV Cache,都有很多优化的空间。
Pipeline parallelism(流水线并行)
既然一块 GPU 放不下,用多块 GPU 总行了吧?这就是 model parallelism(模型并行),可以大致分为 pipeline parallelism 和 tensor parallelism。
大家最容易想到的并行方式就是 pipeline parallelism,模型不是有很多层吗,那就分成几组,每组算连续的几层,穿成一条链。
Pipeline Parallelism 示意图,来源:Colossal AI
这样就有个问题,一条链上只有一个 GPU 在干活,剩下的都在干等。当然聪明的你一定也想到了,既然叫 pipeline,那就可以流水线处理,可以把一个 batch 分为若干个 mini-batch,每个 mini-batch 分别计算。
Pipeline Parallelism 示意图,来源:GPipe
这可好,是不是把 pipeline 搞的越深越好,每个 GPU 只算一层?
首先,正向传播中间状态(activation)的存储容量会成倍增加,加剧内存容量不足的问题。比如流水线的第一级算出了正向传播的中间状态,如果有 N 个流水级,那就要正向流过后面的 N - 1 个流水级,再等反向传播 N - 1 个流水级,也就是 2N - 2 轮之后才能用到这个正向传播的中间状态。不要忘了每一轮都会产生这么多中间状态,因此一共是保存了 2N - 1 个中间状态。如果 N 比较大,这个存储容量是非常恐怖的。
其次,pipeline 的相邻流水级(pipeline stage)之间是要通信的,级数越多,通信的总数据量和总时延就越高。
最后,要让这样的 pipeline 流起来,batch size 需要等于 Transformer 里面的层数,一般是几十,再乘以 data parallelism 的并行数,batch size 会很大,影响模型收敛的速度或模型收敛后的精度。
因此,在内存容量足够的情况下,最好还是少划分一些流水级。
对于 LLaMA-2 70B 模型,模型参数需要 140 GB,反向传播的梯度需要 140 GB,优化器的状态(如果用 Adam)需要 840 GB。
正向传播的中间状态跟 batch size 和选择性重新计算的配置有关,我们在算力和内存之间取一个折中,那么正向传播的中间状态需要 token 长度 * batch size * hidden layer 的神经元数量 * 层数 * (10 + 24/张量并行度) 字节。假设 batch size = 8,不用张量并行,那么 LLaMA-2 70B 模型的正向传播中间状态需要 4096 * 8 * 8192 * 80 * (10 + 24) byte = 730 GB,是不是很大?
总共需要 140 + 140 + 840 + 730 = 1850 GB,这可比单放模型参数的 140 GB 大多了。一张 A100/H100 卡也只有 80 GB 内存,这就至少要 24 张卡;如果用 4090,一张卡 24 GB 内存,就至少需要 78 张卡。
LLaMA-2 模型一共就只有 80 层,一张卡放一层,是不是正好?这样就有 80 个流水级,单是流水线并行就有 80 个并行的 batch 才能填满流水线。
这样,正向传播的中间状态存储就会大到无法忍受,这可是 80 * 2 = 160 轮的中间状态,翻了 160 倍。就算是使用选择性重新计算,比如把 80 层分成 8 组,每组 10 层,中间状态存储仍然是翻了 16 倍。
除非是用最极端的完全重新计算,反向传播到每一层都重新从头开始计算正向传播的中间结果,但这样计算开销可是随模型层数平方级别的增长,第 1 层算 1 层,第 2 层算 2 层,一直到第 80 层算 80 层,一共算了 3240 层,计算开销可是比正常算一次 80 层翻了 40 倍,这还能忍?
中间状态存储的问题就已经够大了,再看这 2048 张卡之间的通信开销。按照一张卡放一层,并且用不同的输入数据让它完全流水起来的做法,这 2048 张卡分别在计算自己的 mini-batch,可以认为是独立参与到 data parallelism 里面了。前面讲过,在数据并行中,每一轮需要传输的是它计算出的梯度和全局平均后的梯度,梯度的数据量就等于模型的参数数量。
把 70B 模型分成 80 层,每一层大约有 1B 参数,由于优化器用的是 32 bit 浮点数,这就需要传输 4 GB 数据。那么一轮计算需要多久呢?总的计算量 = batch size * token 数量 * 6 * 参数量 = 8 * 4096 * 6 * 1B = 196 Tflops,在 4090 上如果假定算力利用率 100%,只需要 0.6 秒。而通过 PCIe Gen4 传输这 4 GB 数据就已经至少需要 0.12 秒了,还需要传两遍,也就是先传梯度,再把平均梯度传过来,这 0.24 秒的时间相比 0.6 秒来说,是占了比较大的比例。
当然我们也可以做个优化,让每个 GPU 在 pipeline parallelism 中处理的 80 组梯度数据首先在内部做个聚合,这样理论上一个 training step 就需要 48 秒,通信占用的时间不到 1 秒,通信开销就可以接受了。当然,通信占用时间不到 1 秒的前提是机器上插了足够多的网卡,能够把 PCIe Gen4 的带宽都通过网络吐出去,否则网卡就成了瓶颈。假如一台机器上插了 8 块 GPU,这基本上需要 8 块 ConnectX-6 200 Gbps RDMA 网卡才能满足我们的需求。
最后再看 batch size,整个 2048 张卡的集群跑起来,每个 GPU 的 mini-batch 我们刚才设置为 8,那可真是 batch size = 16384,已经是大规模训练中比较大的 batch size 了,如果再大,可能就影响模型的收敛速度或收敛后的精度了。
因此,单纯使用流水线并行和数据并行训练大模型的最大问题在于流水线并行级数过多,导致正向传播中间状态(activation)存储容量不足。
Tensor parallelism(张量并行)
那就没办法了吗?我们还有最后一招,就是 Tensor parallelism(张量并行)。它也是模型并行的一种,但不像流水线并行那样是在模型的层间划分,而是在模型的层内划分,也就是把一层内的 attention 计算和 Feed Forward Network 划分到多个 GPU 上处理。
有了张量并行,就可以缓解 GPU 放不下模型导致的流水级太多的问题。分到 80 个 GPU 才能放下的模型,如果用单机 8 卡张量并行,就只需要划分 10 个流水级。同时,张量并行还可以降低 batch size,因为张量并行的几个 GPU 是在算同一个输入数据。
Tensor、Pipeline、Data 三种并行方式从模型层内、模型层间、训练数据三个维度上划分计算空间,来源:DeepSpeed
Attention 的计算过程是比较容易并行的,因为有多个 head,用来关注输入序列中的不同位置的,那么把这些 head 分别拆开就行了。
Attention 的计算过程,来源:The Illustrated Transformer
但是我们做任何并行计算的时候都不要忘记通信开销。
每个 head 里面的 Q、K 两个矩阵的大小是 batch size * token 长度 * key 的大小,V 矩阵的大小是 batch size * token 长度 * value 的大小。key/value 的大小一般等于 embedding size / heads 数量,例如在 LLaMA-2 70B 中就是 8192 / 64 = 128,矩阵大小是 batch size * 4096 * 8192 / 64(注意,这只是一个 head 的)。而 Q、K、V 参数矩阵在每个 head 上的大小是 embedding size * embedding size / heads num = 8192 * 8192 / 64。
我们前面推导过,正向的计算量基本上就是每个 token 过一遍所有参数的计算量,2 * 3 (Q, K, V) * batch size * token 长度 * 参数个数 = 2 * 3 * batch size * 4096 * 8192 * 8192 / 64。可以跟矩阵的大小对一下,看看有没有算错。
那么通信量是多少呢?输出矩阵 Z 是由每个 head 拼起来的,每个 head 的大小是 batch size * token 长度 * embedding size / heads num = batch size * 4096 * 8192 / 64。输入矩阵 X 的大小是 batch size * token 长度 * embedding size = batch size * 4096 * 8192。注意这里的 X 大小跟所有 heads 合并在一起后的 Z 大小是一致的,而我们在这里算的是每个 head 的 Z 大小。这里的单位是参数数量,如果按照字节算,还要乘以每个参数的大小。
如果我们采用最极端的方式,每个 head 交给一个 GPU 去算,那么计算量和通信量的比例是多少?大概是 2 * 3 * embedding size / heads num / bytes per param = 2 * 3 * 8192 / 64 / 2 = 384。代入 4090 的 330 Tflops,如果想让通信不成为瓶颈,那么通信带宽至少需要是 330T / 384 = 859 GB/s,发送接收双向还得乘以 2,就是 1.7 TB/s。太大了,远远超过 PCIe Gen4 x16 的 64 GB/s,就算 NVLink 的 900 GB/s 都撑不住。
所以,tensor parallelism 不能切得太细,每个 GPU 需要多算几个 heads。如果每个 GPU 多算几个 attention heads,输入矩阵 X 就是这些 heads 共享的了,因此输入矩阵的通信开销就被多个 heads 平摊了,计算量和通信量的比例就可以提高。
还是按照 4090 的算力 / 单向通信带宽 = 330T / (64GB/s / 2) 来算,计算量和通信量的比例最少需要是 10000,也就是 2 * 3 * (embedding size / 张量并行 GPU 数量) / bytes per param = 2 * 3 * 8192 / 张量并行 GPU 数量 / 2 >= 10000,解得:张量并行 GPU 数量 <= 2.4。也就是告诉你,要是用了张量并行,最多用 2 个 GPU,如果用更多的 GPU,算力就肯定跑不满理论值。这让我怎么玩?
但是,如果把 H100 的参数代入进去,马上就不一样了。H100 的峰值算力是 989 Tflops,NVLink 双向带宽是 900 GB/s,计算量和通信量的比例最少需要是 1100,也就是 2 * 3 * (embedding size / 张量并行 GPU 数量) / bytes per param = 2 * 3 * 8192 / 张量并行 GPU 数量 / 2 >= 1100,解得:张量并行 GPU 数量 <= 11,也就是单机 8 卡做张量并行,对于 embedding size = 8192 的模型,刚刚好,通信不会成为瓶颈!
阉割版的 H800 相比 H100 卡的就是网络带宽,把网络带宽从 900 GB/s 降到 400 GB/s 了。我们再代入一次,计算量和通信量比例最少需要是 5000,那么张量并行 GPU 数量 <= 4.8。这样单机 8 卡做张量并行,就会导致网络成为瓶颈。当然,计算量 989 Tflops 是理论值,并行切分方式也可以优化,因此实际训练 70B 的模型 8 卡 H800 网络不一定真的是瓶颈。这就是 H800 精准打击大模型训练,让张量并行过得不舒服。
Feed Forward Network 的计算过程,虽然这是 encoder 的,但 decoder 也差不多,来源:Step-by-Step Illustrated Explanations of Transformer
如果在 Feed Forward Network 这里做张量并行,也是可以做类似的推导,在这里就不赘述了。大凡神经网络里的矩阵乘法,M*N 的矩阵乘上 N*K 的矩阵,总的计算量是 M*N*K,输入输出的总大小是 (M*N + N*K),多摞几个矩阵那也是常数(就像 Q、K、V),也就是计算和通信的比例跟矩阵的边长(dimension)是一个量级的。
这么分析完了,如果你是要做大规模大模型训练,你还会买 A100/H100/H800 的 PCIe 版吗?PCIe Gen5 虽然比 Gen 4 快一倍,但对 H100 而言,计算量和通信量的比例仍然最少需要是 989T / (128G / 2) = 15000,解出来张量并行 GPU 数量 <= 1.6,也就是只要用了张量并行,就是损失算力的!
等到 H100 的下一代出来了,比如 GH200,算力又翻了一倍,NVLink 还是 900 GB/s,这时候 NVLink 就也开始有点吃力了。所以 GH200 不失时机的推出了统一大内存,号称 144 TB,就是为了更好的做换入换出,用内存换网络通信。如果禁令保持不变,国内版本还是卡住 400 GB/s 的通信,那性能差距会有多大?
上面的推导当然都是简化的,实际上可能不会这么夸张,但数量级是差不多的。
训练部分小结
4090 不容易做大模型训练的原因除了前面分析的内存小,通信慢,license 不支持数据中心,还有很多其他问题。
比如,A100/H100 支持 ECC 显存容错,据说 4090 也支持 ECC,但是不知道故障率会不会比 A100/H100 更高。不要小看了容错,2048 张卡的集群就算每张卡 1 个月出一次故障,平均 20 分钟就会有一张卡出故障!要是没有自动化的故障恢复方式,炼丹师就别想睡觉了。
就算是自动从上一个 checkpoint 恢复,这可是要时间的,如果不考虑丢弃故障 GPU 梯度这种比较暴力的方式,当前这个 step 就算是白算了,还要从上一个 checkpoint 加载梯度,一般需要 10 来分钟的时间才能搞定。这样,每 20 分钟就浪费 10 分钟,这 10 分钟恢复过程中可能又有新的卡故障,总的算下来要浪费掉一半的有效算力。
因此,保持大规模训练集群的低故障率是非常重要的,这些 GPU 卡都非常金贵,可不能像挖矿机房那样,动不动就过热死机了。
据说 3090 是支持 NVLink 的,但 4090 就把 NVLink 给砍掉了。更老的卡,甚至还有支持 PCIe P2P 的,现在也都被砍掉了。谁感兴趣可以测一测 3090 的 NVLink 性能怎么样,是不是真的能达到标称的 600 GB/s,如果真的能达到的话,是否又可以用来做大模型训练了呢。
我们年会的时候,海哥讲了个段子,我们找老婆都希望又漂亮,又能挣钱,还一心一意爱自己。可同时满足这三个条件的老婆就很难找到了。类似的,在分布式系统中,我们都希望性能又高,通用性又强,成本还低。这三个条件的交集也很小。海哥讲到这里,谭博补充了一句,同时满足这三个条件的分布式系统根本就不存在。
Tensor、Pipeline、Data Parallelism 就像是这样的不可能三角,相互牵制,只要集群规模够大,模型结构仍然是 Transformer,就很难逃出内存容量和网络带宽的魔爪。
大模型推理为什么 4090 很香
推理和训练有什么区别?
首先,训练不仅需要存储模型参数,还需要存储梯度、优化器状态、正向传播每一层的中间状态(activation),后面几个比参数更大,对模型内存的需求量也更大。
其次,训练任务是一个整体,流水线并行的正向传播中间结果是需要存下来给反向传播用的。为了节约内存而使用流水线并行,流水级越多,要存储的中间状态也就更多,反而加剧内存的不足。而推理任务中的各个输入数据之间并没有关系,正向传播每一层的中间状态也不需要保存下来,因此流水线并行不需要存储很多中间状态。
首先我们需要计算一下推理需要多少算力。前面针对训练算力的估算,为了简单起见,忽略了两个事情,首先是没有考虑 KV Cache,其次是没有考虑内存带宽。
KV Cache
什么是 KV Cache?对于每个输入的 prompt,在计算第一个 token 输出的时候,每个 token 的 attention 肯定是都要从头计算。但是在后续 token 的生成中,都需要计算 self-attention,也就是输入 prompt 以及前面输出的 token 的 attention。这是就需要用到前面每一个 token 的 K 和 V,由于每一层的参数矩阵是不变的,此时只有刚生成的那个 token 的 K 和 V 需要从头计算,输入 prompt 和之前生成的 token 的 K 和 V 其实是跟上一轮一样的。
这时,我们就可以把每一层的 K、V 矩阵缓存起来,生成下一个 token 的时候不再需要重新计算,这就是所谓的 KV Cache。Q 矩阵每次都不一样,没有缓存的价值。前面讲的训练中的选择性保存正向 activation 是个拿计算换内存的把戏,这里的 KV Cache 就是一个拿内存换计算的把戏。
KV Cache 需要多少存储容量呢?每一层,每个 token 的 K、V 矩阵都是 embedding size 这么大,再乘上 token 数量和 batch size,就是这一层的 KV Cache 所需的存储容量了。一定要记住 batch size,在正向和反向传播的几乎所有阶段,都不会涉及到对 batch size 中各个 sample 的合并处理,因此它始终是存储量和计算量计算中的一个系数。
例如,如果 batch size = 8,在 LLaMA 2 70B 中,假设输入和输出的 token 数量达到了模型的极限 4096,80 层的 KV Cache 一共需要 2 (K, V) * 80 * 8192 * 4096 * 8 * 2B = 80 GB。如果 batch size 更大,那么 KV Cache 占据的空间将超过参数本身占的 140 GB。
KV Cache 能省下来多少计算量?每一层计算 K、V 矩阵一共需要 2 (K, V) * 2 (mult, add) * embedding size * embedding size = 4 * 8192 * 8192 这么多计算量,乘以之前输入过的 token 数量、层数和 batch size,就是 4096 * 80 * 8 * 4 * 8192 * 8192 = 640 Tflops。相当于每存储 1 个字节,节约了 16K 次计算,还是很划算的。
事实上,KV Cache 节约的远远不止这些。计算 K、V 矩阵的过程是个典型的内存密集型过程,它需要加载每一层的 K、V 参数矩阵。也就是如果不做任何缓存,假设 prompt 长度很短而输出长度接近 token 的最大长度 4096,到了最后一个 token 的时候,单是重复计算前面每个 token 的 K、V 矩阵,就需要读取内存 4096 * 80 * 2 * 8192 * 8192 = 40T 次,每次 2 个字节,要知道 H100 的内存带宽只有 3.35 TB/s,4090 更是只有 1 TB/s,这单是最后一个 token 就得耗掉一张卡几十秒的时间来做重复计算。这样,token 的输出就会越来越慢,整个输出时间是输出长度平方级别的,根本没法用。
推理是计算密集还是存储密集
接下来我们就可以计算推理所需的计算量了。总的算力很好算,前面讲过,大概就是 2 * 输出 token 数量 * 参数数量 flops。如果想看细节,可以看下面这张图,[来源:https://le.qun.ch/en/blog/2023/05/13/transformer-batching/]。
Transformer 推理过程中每一步的矩阵形状、所需算力和内存访问量,来源:Lequn Chen,Dissecting Batching Effects in GPT Inference
但算力并不能说明一切,模型还需要访问 GPU 内存,内存带宽也可能成为瓶颈。至少需要把参数从内存里面读出来吧?事实上,内存带宽的估算就这么简单,内存访问量 = 参数数量 * 2 bytes。中间结果有一部分是可以放在缓存里面的,缓存放不下的部分也需要占内存带宽,我们先不算。
如果不做任何批量输入,也就是模型专门服务一个 prompt,batch size = 1,整个 context 的长度很短(例如只有 128),那么整个推理过程中,每载入一个参数(2 字节),就只进行 128 次乘法和加法计算,那么计算 flops 和访问内存 bytes 的比例就只有 128。基本上任何 GPU 在这种情况下都会变成 memory bound,时间都耗在加载内存上了。
对于 4090 来说,计算 flops 和内存带宽之比是 330 / 1 = 330;对于 H100 来说,计算 flops 和内存带宽之比是 989 / 3.35 = 295。也就是说,如果 context 中的 token 数量小于 330 或者 295,那么内存访问就会成为瓶颈。
虽然 LLaMA 2 的理论上限是 4096 个 token,但很多输入 prompt 用不了这么多,因此内存访问是有可能成为瓶颈的。此时,就需要靠 batch size 来补足了。推理中的批量处理,就是把几乎同时到达后端服务的 prompt 放到一起处理。不用担心,batch 里面的不同 prompt 的处理是完全独立的,不用担心会互相干扰。但这些 prompt 的输出是步调整齐划一的,每一轮整个 batch 中的每个 prompt 都会输出一个 token,因此如果有的 prompt 先输出完了,那就只能等其他的输出结束,造成一定的算力浪费。
有的人问,批量处理所需的算力跟分别单独处理所需的算力是一样的呀,那推理时为什么需要批量处理?答案就在访问内存的带宽上。
如果同时到达服务器的 prompt 很多,是不是 batch size 越大越好?也不是,因为 KV Cache 的大小可是正比于 batch size 的,batch size 大了,KV Cache 占据的 GPU 内存容量就很可观,比如在 LLaMA-2 70B 中,每个 prompt 都要占据 5 GB 的 KV Cache,如果 batch size 搞到 32,那么 KV Cache 就会占掉 160 GB 的 GPU 内存,比参数都大了。
70B 推理需要多少张卡?
总的存储容量也很好算,推理的时候最主要占内存的就是参数、KV Cache 和当前层的中间结果。当 batch size = 8 时,中间结果所需的大小是 batch size * token length * embedding size = 8 * 4096 * 8192 * 2B = 0.5 GB,相对来说是很小的。
70B 模型的参数是 140 GB,不管 A100/H100 还是 4090 都是单卡放不下的。那么 2 张 H100 够吗?看起来 160 GB 是够了,但是剩下的 20 GB 如果用来放 KV Cache,要么把 batch size 压缩一半,要么把 token 最大长度压缩一半,听起来是不太明智。因此,至少需要 3 张 H100。
对于 4090,140 GB 参数 + 40 GB KV Cache = 180 GB,每张卡 24 GB,8 张卡刚好可以放下。
推理用流水线并行可以吗?
推理使用流水线并行,最主要的问题是串行处理的推理延迟,网络延迟倒是小问题。
首先是推理延迟。虽然流水线的不同阶段可以塞进不同的 prompt,但同一个 prompt 的处理仍然永远在单个 GPU 上轮转,这样相比 Tensor parallelism 而言,单个 prompt 的延迟就增大了。
对于很小的 batch size,GPU 内存带宽是瓶颈,此时每张卡计算每个 token 的时延就是 2 byte * 参数量 / 卡的数量 / 内存带宽,例如 8 卡 4090 跑 LLaMA-2 70B,就是 2 * 70G / 8 / 1 TB/s = 0.0175 秒。这里没有考虑 KV Cache 带来的节约。注意,8 张卡是串行处理的,因此每个 token 的时延还要乘以 8,也就是 0.14 秒。每秒只能输出 7 个 token,对于 70B 这么小的模型来说是有点慢了。
对于很大的 batch size,GPU 算力是瓶颈,此时每张卡计算每个 token 的时延就是 batch size * 2 * 参数量 / 卡的数量 / 算力,例如 batch size = 1024,同样的 8 卡例子,就是 1024 * 2 * 70G / 8 / 330 Tflops = 0.0543 秒。事实上,对于这么大的 batch size,KV Cache 和正向传播的中间结果先把 GPU 内存给吃满了。
那么要平衡利用 GPU 算力和内存带宽,batch size 需要是多少呢?这就是 2 byte * 参数量 / 卡的数量 / 内存带宽 = batch size * 2 * 参数量 / 卡的数量 / 算力,左右两边参数量和卡的数量互相抵消,得到 batch size = 算力 / 内存带宽。对于 4090,就是 330 / 1 = 330;对于 H100,就是 989 / 3.35 = 295。也就是说,对 4090 而言,batch size 小于 330 的时候 GPU 内存带宽是瓶颈,大于 330 的时候 GPU 算力是瓶颈。当 batch size = 330 的时候,理想情况下,内存带宽和算力恰好都打满,每张卡处理每个 token 的时间就是 17.5 ms。
其次是网络延迟。流水线并行相比张量并行的优点就是网络传输量小,流水级之间只需要传输 batch size * embedding size 这么多数据。例如 batch size = 8,embedding size = 8192,只需要传输 128 KB 数据,在 32 GB/s 的 PCIe Gen4 x16 上,只需要 4 us 就可以传输完成。当然,还需要考虑到通信库本身的开销,加上 4090 不支持 GPU 之间 P2P 传输,需要通过 CPU 中转,实际上需要几十 us 的时间,相比计算部分动辄几十 ms 的时延,可以忽略不计。
即使 batch size = 330,这 5.28 MB 数据在 PCIe 上也只需要传输 0.16 ms,相比计算部分的 17.5 ms 仍然可以忽略不计。
如果可以忍受流水线并行的推理延迟,甚至可以用多台主机来做流水线并行。我们假设主机间只有 1 Gbps 的普通以太网络,每台主机只有一张 4090。对于 batch size = 1,16 KB 数据需要 0.25 ms 才能传输完成,再加上 0.25 ms 两端网络协议栈的处理时间,每个流水级就需要 0.5 ms 的时延,8 张卡花在通信上的时间只有 4 ms,相比整体计算时延 140 ms 来说可以忽略,不会显著影响系统的推理延迟。
当 batch size 很小时,流水线推理中的网络流量是突发性(bursty)的,每过 18 ms 只会进行 0.25 ms 数据传输,只有 1/72 的占空比,不用担心流水线推理把局域网全部给占满了,搞得没法正常上网了。
如果为了充分利用算力,把 batch size 设置得很大,比如 330,那么 16 KB * 330 = 5.28 MB 数据需要传输 41 ms,8 张卡花在通信上的时间高达 0.33 秒,这样就只有 3 token/s 的输出速度了,难以忍受。因此,如果用主机间通信来做流水线并行,主机间又没有很高的通信带宽,就势必需要牺牲一定的吞吐量。
例如,我们设置输出速度不小于 5 token/s,这时留给通信的时间是 60 ms,每个流水级至多 7.5 ms,1 Gbps 网络可以传输 960 KB 数据,这时 batch size 至多设置为 60,也就是这 8 张 4090 的总吞吐量是 2400 token/s。此时的有效算力利用率只有不到 20%。
最近有一个比较火的 [Petals 开源项目:https://github.com/bigscience-workshop/petals],就是利用流水线并行,把 GPU 做成了一个类似 BitTorrent 的分布式网络。虽然推理延迟确实比较高,但至少说明了分布式 GPU 推理的可行性。
推理用张量并行怎么样?
前面讲到,流水线并行的最大缺点是 GPU 串行处理,延迟较高,导致输出 token 比较慢。而张量并行的最大缺点是传输数据量大,网络带宽低的设备不一定 hold 得住。
但是推理要传输的数据量跟训练要传输的数据量可不是一回事啊!推理只需要传输正向传播的中间结果(activation),而训练还需要传输所有参数的梯度,梯度才是数据量的大头。
在推理中,如果使用张量并行,Transformer 的每一层都需要传输把自己负责的结果向量(大小为 batch size * embedding size / num GPUs)广播给其他所有 GPU,并接受来自所有其他 GPU 广播来的数据。计算 attention 的时候需要传输一次,计算 feed-forward network 的时候又需要传输一次,也就是总共需要传输 2 * 层数这么多次。
每次发送就是 batch size * embedding size(发送和接收是不同的方向,不能算两次),对于 batch size = 1, embedding size = 8192,只需要传输 16 KB 数据,在 32 GB/s 的 PCIe Gen4 上传输只需要 1 us。当然,考虑到前面讨论的 CPU 中转开销,还是需要大约 30 us 的。一共 160 次传输,需要 4.8 ms。
我们再考虑计算的开销。还是考虑 batch size = 1 的情形,GPU 内存带宽是瓶颈,此时每张卡计算每个 token 的时延就是 2 byte * 参数量 / 卡的数量 / 内存带宽,代入我们前面的数值,仍然是 17.5 ms。但是这里 8 张卡是并行处理的,因此总的处理时长就是计算时间 + 通信时间 = 17.5 ms + 4.8 ms = 22.3 ms。这就意味着每秒可以生成 45 个 token,这个 token 生成速度已经很不错了,至少人类的阅读速度是很难赶上生成的速度了。
如果 batch size 更大会怎样?例如 batch size = 330,把 GPU 算力和内存带宽都充分利用起来,每次需要传输的数据量是 330 * 8192 * 2 = 5.4 MB,在 32 GB/s 的 PCIe Gen4 上需要 0.17 ms。一共 160 次传输,就是 27 ms。这下网络通信开销成了延迟的大头,总处理时长为 27 + 17.5 = 44.5 ms,每秒只能生成 22 个 token 了,但也不算慢。
注意,不管用多少个 GPU 做并行推理,只要用的是张量并行,网络传输的总数据量是相同的,因此增加 GPU 的数量只能加速计算,不能加速通信。
因此,A100/H100 的 NVLink 在降低推理延迟方面还是有很大作用的。如果用 H100,取 batch size = 295 达到算力和带宽的平衡利用,这 4.72 MB 数据只需要 4.72 MB / 450 GB/s = 0.01 ms。一共 160 次传输,也只有 1.6 ms。由于内存带宽大了,计算时间也可以大幅缩短,例如 H100 的计算时间为 2 * 70G / 8 / 3.35 TB/s = 5.2 ms。总处理时长只有 5.2 ms + 1.6 ms = 6.8 ms,每秒可以生成 147 个 token,非常棒!
可以说,如果论单个 prompt 的 token 生成速度,无论用多少块 4090 也追不上 8 卡 H100。
用 4090 做推理的成本怎么样?
对于推理,不管用流水线并行还是张量并行,batch size 不算高到太离谱的情况下内存带宽都是瓶颈。
假如 batch size 能够高到把算力 100% 利用起来,并且还能解决 KV Cache 不够大的问题,能解决中间结果占用内存过多的问题,那么这 8 张 4090 可以达到多少吞吐量?
当然,这两个问题都不好解决,因此推理优化才是一个热门的研究领域,存在很多的 trade-off 和奇技淫巧。如果只是用标准的 PyTorch,那推理性能距离把算力 100% 利用起来还远得很呐。
假设都解决了,在张量并行的通信过程中我们可以利用 double buffer 做另外一个 batch 的计算,也就是计算和通信并行,进一步提高吞吐量。通信和计算分别是 27 ms 和 17.5 ms,传输的 27 ms 是瓶颈,也就是每 27 ms 输出一组 token,一个 batch 330 个 prompt,那这 8 张 4090 真是可以达到每秒 330 / 0.027 = 12.2K token 的吞吐量。
8 张 4090 的成本是 12800 美金,8 卡 PCIe Gen4 服务器本身要 2 万美金,加上网络设备,平均每台 4 万美金的设备成本。固定资产按照 3 年摊销,每小时 1.52 美元。整机功耗大约 400W * 8 + 2 kW = 5 kW,按照 0.1 美元一度电算,每小时 0.5 美元。一个机架可以放 4 台这样的 8 卡服务器,数据中心机柜租用成本(不含电费)一个月 1500 美元算贵的了,合每小时 0.5 美元。这 2.5 美元一小时的机器,满打满算能生成 12.2K * 3600 = 44M tokens,也就是说 1 美元能生成 17.6M tokens。
是不是比 GPT-3.5 Turbo 的 $0.002 / 1K tokens,也就是 1 美元 0.5M tokens 便宜 35 倍?当然,账不能这么算。
- 首先,GPU 的算力利用率到不了 100%;
- 其次,如同所有 SaaS 服务一样,用户的请求数量有波峰有波谷,用户是按量付费的,平台提供方可是不管有没有人用都在烧钱的;
- 此外,每个 batch 中不同 prompt 的长度和响应 token 数量都不同,消耗的算力是 batch 中最大的那个,但收的钱是用户实际用的 token 数;
- 再次,GPT-3.5 是 175B 的模型,比 70B 的 LLaMA 很可能推理成本更高;
- 最后,OpenAI 开发 GPT-3.5 是烧了不知道多少钱的,人家至少要赚回训练成本和研发人员的工资吧。
其实 GPT-3.5 Turbo 的 $0.002 / 1K tokens 真的挺良心的,有的卖 API 的,LLaMA-2 70B 都敢比 GPT-3.5 Turbo 卖得贵。
如果换成用 H100 做推理,重新算一下这笔账。一张 H100 至少要 3 万美金,一台 8 卡 H100 高配服务器加上配套的 IB 网络,起码要 30 万美金,同样按照 3 年摊销,每小时 11.4 美元。10 kW 功耗,电费每小时 1 美元。一个普通供电和散热的机架只能放 2 台 8 卡 H100,机柜租用成本(不含电费)还按 1500 美元算,合每小时 1 美元。一共 13.4 美元一小时。
这其实已经是非常良心的价格了,你在任何云服务商都不可能租得到这么便宜的 8 卡 H100。所以说从云服务商租卡卖没有护城河的 SaaS 服务,比如开源模型的推理 API,除非有一种提高推理性能的独门绝技,很难赚得了什么大钱,二房东的生意不是这么好做的。
再算算这台 8 卡 H100 机器的吞吐量,张量并行也采用传输和计算并行,H100 的通信比较快,因此计算是瓶颈,每 5.2 ms 可以输出一组 token,一个 batch 295 个 prompt,满打满算可以达到每秒 295 / 0.0052 = 56K token 的吞吐量。理想情况下,一小时能生成 204M tokens,也就是 1 美元能生成 15.2M tokens,H100 单位 token 的成本比 4090 仅仅高 16%,可以算打个平手吧。
为什么 8 卡 H100 机器是 4090 机器生命周期价格的 5 倍,性价比却跟 4090 差不多?因为一张 H100 的算力是 4090 的 3 倍,内存带宽是 4090 的 3.35 倍,不管按延迟还是按带宽算,单卡的性能就基本上是 3 倍。而且,H100 比 4090 的网络带宽强太多了,导致 4090 在张量并行中网络通信成了瓶颈,浪费了有效算力。因此,同样的 8 卡机器吞吐量可以达到 4090 的 4.6 倍。虽然一张 H100 卡的价格是 4090 的 20 倍以上,但算上服务器本身的成本、电费和数据中心托管费用,整机的成本只是 5 倍左右。
用最便宜的设备搞出最高的推理性能
我们发现在 8 卡 4090 机器中,3 万美金的设备成本,GPU 卡只占了 1.28 万美金,不像 H100 机器那样 GPU 成本占了大头。还有办法进一步降低吗?
如果我们可以忍受 5 token/s 的输出速度,甚至可以利用流水线并行,用家用台式机和 4090 攒出个推理集群来。
遥想我当年在 MSRA 的时候,在一台只用 1000 美金攒出来的机器上插了 10 块 FPGA ,做出个世界最快的 Key-Value Store。其实如果让我去设计一个性价比最高的 4090 推理集群,有很多种方案可以尝试:
- 用流水线并行,台式机 + 10 Gbps 网卡,足够在 5 ms 内传输 batch size = 330 的 5.28 MB 数据了,通信 40 ms,计算 140 ms,达到 5 token/s 的单 prompt 输出速度,同时又能充分利用 4090 的算力。10 Gbps 的网卡和交换机都很便宜,Intel X710 网卡只要 150 美金,20 口交换机只要 1500 美金(每 8 个口 750 美金),一台家用台式机 700 美金,这只要 2 万美金就可以搞定原本需要 4 万美金的设备。
- 用张量并行,台式机 + 200 Gbps ConnectX-6 网卡,上 RoCE,可以把 batch size = 330 的 5.28 MB 数据在 0.22 ms 内传完,160 次传输是 35 ms,加上计算的 17.5 ms,一个 token 52.5 ms,可以达到 19 token/s 的单 prompt 输出速度,这个速度已经不错了。网卡 1000 美金,200G 交换机 2 万美金 40 个端口,平均每 8 个端口 4000 美金,一台家用台式机 700 美金,这只要 3 万美金就能搞定原本 4 万美金的设备。
- 主机内用张量并行,主机间用流水线并行,4 卡 PCIe Gen4 服务器主板只要 1000 美金而且能跑满 PCIe 带宽(因为 8 卡就需要 PCIe switch 了,价格会贵很多),两台主机之间用 25 Gbps 网卡直连,主机内张量并行的时延是 27 ms,主机间流水线并行只需 2 次 8 ms 的传输(注意 25G 的网络带宽是 4 张 GPU 卡共享的),加上两次流水线计算各 17.5 ms,总共 78 ms,可以达到 13 token/s 的单 prompt 输出速度。网卡 300 美金 * 2,服务器 3000 美金 * 2,这只要 1.95 万美金就可以搞定原本需要 4 万美金的设备。
2 万美金按照 3 年摊销是每小时 0.76 美元。按照 0.1 美元/度的电价,每小时的电费都要 0.5 美元,接近设备成本了,这有点挖矿的味道了。矿场里面可没有中央空调和 UPS,只有暴力风扇,托管费用比数据中心低很多,整机的成本是有可能压到 1.5 美元/小时的。如果跑满了 44M tokens 的吞吐量,1 美元能生成 30M tokens,正好是 8 卡 H100 的 15M token per dollar 的 2 倍。
为什么 H100 以 20 倍于 4090 的 GPU 价格,性价比却只差一倍?首先是因为能耗成本更低,8 卡 H100 的功耗是 10 kW,但 9 台 8 卡 4090 的功耗是 45 kW;其次是因为主机和网络设备成本更低,一台 8 卡 H100 准系统虽然贵,但只占整机价格的 20% 左右;但 4090 因为卡多,除非像 GPU 矿机那样压成本,只要还是用数据中心级的设备,准系统价格就要占到 35% 以上。
其实,这个世界上不止有 A100/H100 和 4090,还有 A10、A40 等计算卡和 3090 等游戏卡,还有 AMD 的 GPU 和很多其他厂商的 AI 芯片。H100 和 4090 大概率都不是性价比的最优解,例如 A10、A40 和 AMD GPU 的性价比有可能就更高。
我都想搞一个推理性价比挑战赛,看谁能用最便宜的设备搞出最强的推理吞吐量,同时延迟不能太高;或者用最便宜的设备搞出最低的推理延迟,同时吞吐量不能太低。
这一切都是在假设使用 LLaMA-2 70B 模型,没有做量化压缩的前提下。如果做了量化压缩,那性能就更高,甚至在 Unified Memory 够大的 MacBook Pro 上都能单机跑了。
License 问题怎么办?
我把这个问题放到最后。[NVIDIA Geforce driver 的 License:https://www.nvidia.com/en-us/drivers/geforce-license/] 里写道:
No Datacenter Deployment. The SOFTWARE is not licensed for datacenter deployment, except that blockchain processing in a datacenter is permitted.
既然机器都是用台式机攒起来的,这能叫 data center 吗?还是叫矿场比较合适吧。人家也说了,4090 用来做区块链是允许的。
我有一个大胆的想法,如果未来的区块链不再用挖矿来做 proof of work,而是用大模型推理来做 proof of work,这是不是很有意思?每个人买几块显卡,接到矿池上,既可以自己用来玩游戏,闲时又可以贡献算力。矿池直接就是个卖大模型推理 SaaS 服务的公司,提供前所未有的低价 API。甚至需要大模型推理服务的人可以在区块链里自己 P2P 玩起来,谁要用大模型就付点 gas。
当然,目前的 proof of work 都是计算很复杂,验证很简单的。如果真用大模型推理做 proof of work,必须防止用户随意编造一个结果交上去。当然这也是有解决方案的,就像 BitTorrent 和其他一些去中心化网络一样,采用信用机制,新人只能做验证别人计算结果的工作,积攒信用;老人每次算错了,都有比较严厉的惩罚。
从另一个角度看,家庭局域网络的速度也越来越快,比如我家就自己部署了 10 Gbps 的网络。家中的智能设备越来越多,算力越来越强。光纤入户也越来越普遍,小区和城市的运营商机房里部署了越来越多的边缘计算节点。前面我们用 1 Gbps 的网络就足以把多台主机上的 GPU 组成流水线并行,那么在未来的家庭高速网络中,流水线并行甚至张量并行都将成为可能。
大多数搞 AI 推理的都只关心数据中心,忽略了家中的分布式算力。只要解决了安全、隐私和经济动机问题,我家的 Siri,也许就跑在邻居家里的 GPU 上。
很多人都在说要 democratize AI。但现在大模型平民化的最大障碍就是成本,而成本最大的来源又是 GPU 市场上计算卡和游戏卡价格的剪刀差。这并不是指责某家公司,其他做 AI 芯片的公司,AI 芯片的算力也并不便宜。毕竟芯片、软件和生态的研发都是白花花的银子。
就像本文开头提到的微软给每台服务器部署 FPGA 一样,大规模量产的芯片价格就像沙子一样。到时候,能限制大模型推理算力的就只有能源了,就像区块链挖矿和通用 CPU 的云计算一样,都在找最便宜的电力供应。我在之前的一个采访中就表示,长期来看,能源和材料可能是制约大模型发展的关键。让我们期待廉价的大模型走进千家万户,真正改变人们的生活。
#大模型~简化开发过程
史蒂夫・乔布斯曾经把计算机称作 “心灵之自行车”。不过,人们对他这个比喻的背景知之甚少,他是在谈及地球上所有物种移动效率的时候提到的。
由 DALL·E 3 生成的图片,提示 “将计算机想象成心灵的自行车”
秃鹫赢了,位居榜首,超过了其他所有物种。人类排在榜单大约三分之一的位置…… 但是,一旦人类骑上自行车,就能远远超越秃鹫,登顶榜首。这让我深受启发,人类是工具制造者,我们可以制造出将这些固有能力放大到惊人程度的工具。对我来说,计算机一直是思维的自行车,它让我们远远超越了固有的能力。我认为我们只是处于这个工具的早期阶段,非常早期的阶段。我们只走了很短的一段距离,它仍处于形成阶段,但我们已经看到了巨大的变化。我认为,与未来 100 年发生的事情相比,这算不了什么。
—— 史蒂夫・乔布斯(1990)
#01 谨慎乐观
LLM 在加速软件开发方面的作用引发了广泛讨论。有人认为,自动生成的代码质量过低,以至于使用这些代码产生的是负面效果。而另一方面,许多人声称编程的时代已经结束。已经有众多研究试图客观评估 LLM 在诸如 HumanEval 或 MBPP 这样的代码质量基准数据集上的表现。这些评估对于该领域的发展至关重要,但这并非本文的焦点。
本文旨在为开发者,尤其是那些对利用这些模型持保留态度的开发者,提供实用的建议。要注意的是,我相信通过恰当地使用这项技术,一个人的工作效率确实能够实现质的飞跃。
需要明确的是,使用 LLM 并不意味着一定能产出高性能的代码。即便是最先进的模型也会犯下严重的错误。但重要的是要记住,最优秀的人类工程师也难免犯错。这就是我们为什么不会在没有设置防护措施的情况下发布代码,比如进行自动化测试和同行评审。LLM 没有改变最佳实践的本质,但可能改变了实现这些最佳实践的容易度。
编写高性能代码的速度真的那么重要吗?
如果大多数专业软件工程师的大部分时间都花在了键盘敲击上,我会觉得非常意外。如果真是这样,我们或许会更认真地把 “代码行数” 当作衡量工作表现的一个标准。我认为,工程师实际上更多的时间是用来决定首先需要编写哪些代码。目前,LLM 还没有能力进行这种高级的抽象思维 ——OpenAI 和许多其他机构正在努力使这些步骤成为可能,但这很可能需要的不仅仅是增加模型的参数或扩大训练数据,而是需要一种全新的思维方式。然而,LLM 已经能够在很大程度上加速我们的工作中的很大一部分了。
#02 编写高效的函数
让我举个例子来说明,当 LLM 生成了一小段代码后,我们需要采取哪些步骤来优化它,使之变得实用。在地球上的某个点和另一组坐标之间测量距离是地理信息分析中的一个常见任务。想象一下这样的情境。
首先看到的代码输出是这样的:
import math
def haversine(lat1, lon1, lat2, lon2):
# Radius of the Earth in kilometers
R = 6371.0
# Convert degrees to radians
lat1_rad = math.radians(lat1)
lon1_rad = math.radians(lon1)
lat2_rad = math.radians(lat2)
lon2_rad = math.radians(lon2)
# Differences in coordinates
dlat = lat2_rad - lat1_rad
dlon = lon2_rad - lon1_rad
# Haversine formula
a = math.sin(dlat / 2)**2 + math.cos(lat1_rad) * math.cos(lat2_rad) * math.sin(dlon / 2)**2
c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))
# Distance in kilometers
distance = R * c
return distance
def find_closest_oil_platform(vessel_lat, vessel_lon, platforms):
closest_platform = None
min_distance = float('inf')
for platform in platforms:
distance = haversine(vessel_lat, vessel_lon, platform['latitude'], platform['longitude'])
if distance < min_distance:
min_distance = distance
closest_platform = platform
return closest_platform, min_distance
那么,这段代码的质量怎么样呢?
这段代码的质量并不理想。它的问题不少,最明显的是它难以阅读,因为它使用了难以理解和不一致的变量命名,包括大小写的单字符命名。此外,它缺少了文档说明和类型提示,使用了一个不必要且低效的 for 循环,也没有处理一些特殊情况,比如相等距离的问题。
但经过修改之后:
from typing import Tuple, List
import numpy as np
def haversine_vectorized(vessel_lat: float, vessel_lon: float, platform_lats: np.ndarray, platform_lons: np.ndarray) -> np.ndarray:
"""
Calculate the Haversine distance between a vessel and multiple platforms in a vectorized manner.
Parameters:
- vessel_lat: Latitude of the vessel in decimal degrees.
- vessel_lon: Longitude of the vessel in decimal degrees.
- platform_lats: Numpy array of latitudes of oil platforms in decimal degrees.
- platform_lons: Numpy array of longitudes of oil platforms in decimal degrees.
Returns:
- distances: Numpy array of distances from the vessel to each platform in kilometers.
"""
# Convert decimal degrees to radians
lat1, lon1, lat2, lon2 = map(np.radians, [vessel_lat, vessel_lon, platform_lats, platform_lons])
# Haversine formula
dlat = lat2 - lat1
dlon = lon2 - lon1
a = np.sin(dlat/2)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon/2)**2
c = 2 * np.arctan2(np.sqrt(a), np.sqrt(1-a))
r = 6371 # Radius of Earth in kilometers
return c * r
def find_closest_platform(vessel_lat: float, vessel_lon: float, platform_lats: np.ndarray, platform_lons: np.ndarray) -> Tuple[List[int], float]:
"""
Finds the closest oil platform(s) to a vessel given arrays of platform latitudes and longitudes, handling equidistant platforms.
Parameters:
- vessel_lat: Latitude of the vessel in decimal degrees.
- vessel_lon: Longitude of the vessel in decimal degrees.
- platform_lats: Numpy array of latitudes for oil platforms.
- platform_lons: Numpy array of longitudes for oil platforms.
Returns:
- A tuple containing a list of indices of the closest platforms and the distance to them in kilometers.
"""
# Calculate distances to all platforms
distances = haversine_vectorized(vessel_lat, vessel_lon, platform_lats, platform_lons)
# Find the minimum distance
min_distance = np.min(distances)
# Find all indices with the minimum distance
closest_indices = np.where(distances == min_distance)[0].tolist()
# Return the indices of all closest platforms and the minimum distance
return closest_indices, min_distance
改进后的代码有了明显提升。它变得更容易阅读了,增加了文档说明和类型提示,并且用更高效的向量计算方式替换了原有的 for 循环。
但是,代码的 “好坏”,更重要的是,它是否满足需求这些都取决于代码将要运行的具体环境。要知道,我们无法仅凭几行代码就能有效评估其质量,这一点对人类如此,对 LLM 也是如此。
比如说,这段代码的准确度是否满足用户的预期?它会被频繁运行吗?是一年一次,还是每微秒一次?使用的硬件条件如何?预期的使用量和规模是否值得我们去追求那些细小的优化?在考虑到你的薪资之后,这样做是否划算?
让我们在上述因素的基础上来评估这段代码。
在准确性方面,虽然半正矢公式(haversine formula)表现不错,但并非最佳选择,因为它将地球视为一个完美的球体,而实际上地球更接近于一个扁球体。在需要跨越巨大距离进行毫米级精确测量时,这种差异变得非常重要。如果真的需要这样的精确度,虽然有更精确的公式(如 Vincenty 公式)可用,但这会带来性能上的折中。因为对于这段代码的用户而言,毫米级的精确度并不是必须的(事实上,由于卫星图像导出的船舶坐标的误差,这种精度也并不相关),所以在准确性方面,半正弦函数是一个合理的选择。
代码运行得够快吗?考虑到只需要对几千个海上石油平台计算距离,特别是通过向量计算方法,这种计算是非常高效的。但如果应用场景变成了计算与岸边任意点的距离(岸线上有数以亿计的点),那么采用 “分而治之” 的策略可能会更加合适。在实际应用中,考虑到节约计算成本的需要,这个函数设计为在一个尽可能配置低的虚拟机上每天运行约 1 亿次。
基于这些详细的背景信息,我们可以认为上面的代码实现是合理的。这也意味着,在代码最终合并前,它应该先经过测试(我通常不推荐仅依赖 LLM 进行测试)和人工同行评审。
#03 加速前进
像之前那样利用 LLM 自动生成实用的函数不仅可以节省时间,而且当你开始利用它们来生成整套的库、处理模块间的依赖、撰写文档、实现可视化(通过多模态能力)、编写 README 文件、开发命令行接口等时,它们带来的价值将会成倍增长。
我们来试着从零开始,借助 LLM 的广泛辅助,创建、训练、评估并推断一个全新的计算机视觉模型。以一篇最近发表的论文为例,“通过深度学习识别 Sentinel-2 图像中船舶尾迹组件的关键点方法”(Del Prete 等人,IEEE GRSL,2023),这篇论文就是我们前进的动力和灵感来源。
论文链接:https://www.semanticscholar.org/paper/Keypoints-Method-for-Recognition-of-Ship-Wake-in-by-Prete-Graziano/a38d19b5ebaa2441e1bef2af0ecf24332bd6ca5b
为什么我们需要关心船舶在卫星图像中的行进方向,这项任务有什么难点呢?
通过静态图像识别船只的航行方向,对于那些需要监控水域中人类活动的组织来说,是极其宝贵的信息。比如,如果一艘船正朝向一个海洋保护区行进,这可能意味着需要警觉或者采取拦截措施。通常,全球范围内公开的卫星图像的分辨率不足以精确判断一艘船的朝向,尤其是那些在图像上只占据几个像素的小型船只(例如,Sentinel-2 的图像分辨率为 10 米 / 像素)。然而,即便是小型船只留下的水波纹也可能相当明显,这就为我们提供了一个判断船只朝向和行进方向的线索,即使船的尾部无法直接识别。
这项研究之所以引人注目,是因为它采用的模型基于 EfficientNetB0,这是一个足够小的模型,能够在不花费太多计算资源的情况下进行大规模应用。虽然我没有找到具体的代码实现,但作者公开了包括标注在内的数据集,这是值得赞赏的一步。
https://zenodo.org/records/7947694
开始我们的探索吧!
如同启动任何新的机器学习项目一样,首先对数据进行可视化是极富启发性的一步。
import os
import json
from PIL import Image, ImageDraw
import matplotlib.pyplot as plt
import seaborn as sns
# Define the path to your data directory
data_dir = "/path/to/your/data" # Adjust this to the path of your data directory
annotations_dir = os.path.join(data_dir, "annotations")
images_dir = os.path.join(data_dir, "imgs")
# Initialize Seaborn for better visual aesthetics
sns.set(style="whitegrid", palette="muted")
# Create a list to hold file paths for images and their corresponding annotations
image_files = []
annotation_files = []
# Loop through the annotations directory to get the list of annotation files
for annotation_file in os.listdir(annotations_dir):
if annotation_file.endswith(".json"):
annotation_path = os.path.join(annotations_dir, annotation_file)
image_file = annotation_file.replace(".json", ".png") # Assuming image file names match annotation file names
image_path = os.path.join(images_dir, image_file)
# Check if the corresponding image file exists
if os.path.exists(image_path):
annotation_files.append(annotation_path)
image_files.append(image_path)
# Plotting
num_examples = min(len(image_files), 10) # Limiting to 10 examples for visualization
fig, axes = plt.subplots(2, 5, figsize=(20, 8))
for idx, (image_path, annotation_path) in enumerate(zip(image_files[:num_examples], annotation_files[:num_examples])):
# Load the image
img = Image.open(image_path).convert("RGB") # Ensure the image is treated as RGB
draw = ImageDraw.Draw(img)
# Load the corresponding annotations and draw keypoints
with open(annotation_path, 'r') as f:
annotations = json.load(f)
for point in annotations["tooltips"]:
x, y = point["x"], point["y"]
# Draw keypoints in cyan for visibility
draw.ellipse([(x-10, y-10), (x+10, y+10)], fill='cyan', outline='black')
# Plot the image with keypoints
ax = axes[idx // 5, idx % 5]
ax.imshow(img)
ax.axis('off')
ax.set_title(f"Image {idx+1}")
plt.tight_layout()
plt.show()
结果如下:
Sentinel-2 卫星图像中的船只和尾流关键点注释(青色)样本。
结果展示了几个关键的发现:首先,不同图片中标注的关键点(青色点)数量不一;其次,这些图像都是灰度的,并且形状各异;第三,尽管样本量不大,但其内部差异性却相当大。在明确了这些基础之后,我们接下来将编写用于加载数据集、定义神经网络、进行训练和验证的代码。在请求代码之前,我们会先向 LLM 提供一些背景信息。
这导致了以下输出:
import os
import json
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torchvision.io import read_image
from typing import Tuple, List
class WakeDataset(Dataset):
"""
Custom Dataset for loading wake detection imagery and annotations, with preprocessing for EfficientNet.
Images are grayscale and will be converted to three-channel images for model compatibility.
"""
def __init__(self, annotations_dir: str, images_dir: str, transform=None):
"""
Initializes the dataset with paths to annotations and images directories, and a transform.
Parameters:
- annotations_dir: Path to the directory containing annotation JSON files.
- images_dir: Path to the directory containing grayscale image files.
- transform: A torchvision.transforms transform to be applied on a sample.
"""
self.annotations_dir = annotations_dir
self.images_dir = images_dir
self.transform = transform if transform else self.default_transform()
self.annotations = self._load_annotations()
def default_transform(self):
"""
Defines the default transform to resize images, convert them to three-channel format, and to tensors.
Returns:
- A torchvision.transforms composition for image preprocessing.
"""
return transforms.Compose([
transforms.Resize((224, 224)), # Resize to the input size for EfficientNet
transforms.Grayscale(num_output_channels=3), # Convert grayscale to three-channel
transforms.ConvertImageDtype(torch.float), # Ensures image is a float tensor
])
def _load_annotations(self) -> List[Tuple[str, List[int]]]:
"""
Loads and filters annotations with at least 2 keypoints.
Returns:
- A list of tuples, where each tuple contains the image filename and its keypoints.
"""
annotations = []
for filename in os.listdir(self.annotations_dir):
if not filename.endswith('.json'):
continue
filepath = os.path.join(self.annotations_dir, filename)
with open(filepath, 'r') as f:
data = json.load(f)
if len(data['tooltips']) >= 2: # Filter annotations with at least 2 keypoints
keypoints = [item for sublist in [[kp['x'], kp['y']] for kp in data['tooltips']] for item in sublist]
annotations.append((filename.replace('.json', '.png'), keypoints))
return annotations
def __len__(self) -> int:
"""
Returns the number of items in the dataset.
"""
return len(self.annotations)
def __getitem__(self, idx) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Retrieves an item from the dataset at the specified index, with preprocessing for EfficientNet.
Parameters:
- idx: Index of the item to retrieve.
Returns:
- A tuple containing the preprocessed image tensor and its corresponding keypoints tensor.
"""
img_name, keypoints = self.annotations[idx]
img_path = os.path.join(self.images_dir, img_name)
image = read_image(img_path) # Load image
keypoints = torch.tensor(keypoints + [0] * (6 - len(keypoints)), dtype=torch.float) # Pad keypoints
if self.transform:
image = self.transform(image)
return image, keypoints
最终得到的代码看起来还算合理,并且似乎注意到了问题的一些细节。虽然不是完美的(例如,选择在 default_transform 中将图像从灰度转为 RGB 颜色模式是个奇怪的决定,我个人更倾向于使用 pathlib 而非 os。listdir),但我们现在可以先放下这些细节不谈。
import torch
import torch.nn as nn
from torchvision import models
from typing import Tuple
class EfficientNetB0KeypointDetector(nn.Module):
"""
A custom EfficientNet-B0 model for keypoint detection in grayscale images.
The network modifies the input layer to accept single-channel images and
adds a custom output layer for predicting a fixed number of keypoints.
"""
def __init__(self, num_keypoints: int = 3):
"""
Initializes the model.
Parameters:
- num_keypoints: The number of keypoints to predict. Default is 3,
corresponding to the maximum number of keypoints in the dataset.
"""
super(EfficientNetB0KeypointDetector, self).__init__()
self.num_keypoints = num_keypoints
# Load a pre-trained EfficientNet-B0
self.efficientnet_b0 = models.efficientnet_b0(pretrained=True)
# Modify the first convolutional layer to accept single-channel (grayscale) images
original_first_conv = self.efficientnet_b0.features[0][0]
self.efficientnet_b0.features[0][0] = nn.Conv2d(
in_channels=1,
out_channels=original_first_conv.out_channels,
kernel_size=original_first_conv.kernel_size,
stride=original_first_conv.stride,
padding=original_first_conv.padding,
bias=False)
# Modify the classifier to predict keypoints (num_keypoints * 2 because x, y for each keypoint)
self.efficientnet_b0.classifier = nn.Linear(
in_features=self.efficientnet_b0.classifier[1].in_features,
out_features=num_keypoints * 2)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Forward pass of the network.
Parameters:
- x: Input tensor of shape (batch_size, 1, 224, 224)
Returns:
- A tensor of shape (batch_size, num_keypoints * 2) representing the predicted keypoints.
"""
return self.efficientnet_b0(x)
# Example usage
if __name__ == "__main__":
model = EfficientNetB0KeypointDetector()
print(model)
# Example input tensor (batch_size, channels, height, width)
example_input = torch.randn(1, 1, 224, 224)
output = model(example_input)
print(output.shape) # Expected shape: (1, 6) for 3 keypoints (each with x, y coordinates)
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from dataset import WakeDataset, get_dataloader # Assuming dataset.py is in the same directory
from net import EfficientNetB0KeypointDetector
from torchvision import transforms
import wandb
# Initialize Weights & Biases
wandb.init(project="wake_model_llm_assist", entity="your_wandb_entity", mode="offline")
def custom_collate_fn(batch):
"""
Custom collate function to pad keypoints to a fixed size.
Parameters:
- batch: A list of tuples (image, keypoints).
Returns:
- Batched images and keypoints with padding as necessary.
"""
images, keypoints = zip(*batch)
images = torch.stack(images, 0)
max_keypoints = 3 * 2 # Maximum number of keypoints (3) times 2 for (x, y)
padded_keypoints = [torch.cat((kp, torch.zeros(max_keypoints - len(kp)))) for kp in keypoints]
keypoints = torch.stack(padded_keypoints, 0)
return images, keypoints
def train_one_epoch(model: nn.Module, dataloader: DataLoader, optimizer: optim.Optimizer, device: torch.device):
"""
Trains the model for one epoch.
Parameters:
- model: The neural network model.
- dataloader: DataLoader providing the training data.
- optimizer: Optimizer used for model training.
- device: The device to train on.
"""
model.train()
total_loss = 0.0
for images, keypoints in dataloader:
images, keypoints = images.to(device), keypoints.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = nn.MSELoss()(outputs, keypoints)
loss.backward()
optimizer.step()
total_loss += loss.item()
average_loss = total_loss / len(dataloader)
wandb.log({"train_loss": average_loss})
def evaluate(model: nn.Module, dataloader: DataLoader, device: torch.device):
"""
Evaluates the model on the validation set.
Parameters:
- model: The neural network model.
- dataloader: DataLoader providing the validation data.
- device: The device to evaluate on.
"""
model.eval()
total_loss = 0.0
with torch.no_grad():
for images, keypoints in dataloader:
images, keypoints = images.to(device), keypoints.to(device)
outputs = model(images)
loss = nn.MSELoss()(outputs, keypoints)
total_loss += loss.item()
average_loss = total_loss / len(dataloader)
wandb.log({"val_loss": average_loss})
def main():
# Setup
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = EfficientNetB0KeypointDetector().to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# DataLoaders
train_dataloader = get_dataloader('path/to/train/annotations', 'path/to/train/images', batch_size=4, transform=transforms.ToTensor(), collate_fn=custom_collate_fn)
val_dataloader = get_dataloader('path/to/val/annotations', 'path/to/val/images', batch_size=4, transform=transforms.ToTensor(), collate_fn=custom_collate_fn)
# Training loop
for epoch in range(num_epochs):
train_one_epoch(model, train_dataloader, optimizer, device)
evaluate(model, val_dataloader, device)
# Log model checkpoints, images, or other metrics as needed
wandb.log({})
if __name__ == "__main__":
main()
当我开始尝试运行代码时,遇到了一系列与张量形状不匹配相关的问题(想起了之前 RGB 与灰度图像的差异以及自定义 collate 函数的问题)。我尝试了几分钟的调试,但最终选择将所有模块代码复制到提示框里,让它帮我一起找出问题所在。
这样做有效解决了所有挂起的问题,我随即开始了模型的训练…… 但训练速度慢得出奇。这时我才意识到我没有利用 Apple 芯片上的 Metal 性能着色器(MPS)后端。虽然我平时不太在个人电脑上训练模型,MPS 对我还比较新,但我还是决定添加一个条件判断来使用它。
鉴于训练数据量适中(共 581 张图片),加上 EfficientNet 已经在 ImageNet 上进行了预训练,我决定进行 1000 个周期的训练。
500 个周期后,train_loss 仍在减少,但验证损失似乎已经收敛(至少对于快速评估而言足够收敛)。图像是根据权重和偏差复制的。
值得注意的是,尽管存在一些已经在卫星图像上进行预训练的基础模型(相对于 ImageNet),这些模型很可能对这项任务更有效,但这些网络体积比 EfficientNet 大得多,因此训练速度更慢(而且太新,还没被 LLM 训练数据集包括)。
现在,让我们编写一个推理脚本和一个命令行界面(CLI),来测试我们刚训练好的模型。
import torch
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt
import click
# Assuming net.py contains the EfficientNetB0KeypointDetector class
# and dataset.py contains the transforms used during training
from net import EfficientNetB0KeypointDetector
from dataset import WakeDataset # Modify as necessary to import transforms
def load_model(model_path):
"""Load the trained model from a file."""
model = EfficientNetB0KeypointDetector()
model.load_state_dict(torch.load(model_path))
model.eval() # Set the model to inference mode
return model
def preprocess_image(image_path, transform):
"""Load and preprocess an image."""
image = Image.open(image_path).convert("L") # Assuming grayscale conversion as in your dataset
image = transform(image)
# Add batch dimension (BxCxHxW)
image = image.unsqueeze(0)
return image
def plot_keypoints(image, keypoints):
"""Plot keypoints on the image."""
plt.imshow(image.squeeze(), cmap='gray') # Remove batch dimension and show image
plt.scatter(keypoints[:, 0], keypoints[:, 1], s=50, marker='.', c='red')
plt.show()
@click.command()
@click.argument('model_path', type=click.Path(exists=True))
@click.argument('image_path', type=click.Path(exists=True))
def run_inference(model_path, image_path):
"""Run inference on an image using a trained model."""
# Use the same transforms as during training
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Grayscale(num_output_channels=3),
])
model = load_model(model_path)
image = preprocess_image(image_path, transform)
# Perform inference
with torch.no_grad():
keypoints = model(image)
keypoints = keypoints.view(-1, 2).cpu().numpy() # Reshape and convert to numpy for plotting
# Load original image for plotting
original_image = Image.open(image_path).convert("L")
plot_keypoints(original_image, keypoints)
if __name__ == '__main__':
run_inference()
让我们开始吧!
虽不完美,但对于第一次通过来说是合理的。
你可以在 GitHub 上找到包括所有模块、模型及权重(第 500 周期的)和一个 readme 的完整代码。我花了不到一个小时就生成了整个库,这个过程比写这篇文章花费的时间要少得多。所有这些工作都是在我的个人开发环境中完成的:MacBook Air M2 + VS Code + Copilot + 保存时自动格式化(使用 black、isort 等)+ 一个 Python 3.9.6 的虚拟环境(.venv)。
GitHub:https://github.com/pbeukema/wakemodel_llmassist
学到的教训
- 向模型提供尽可能多的相关上下文,帮助其解决任务。要记住,模型缺少许多你可能认为理所当然的假设。
- LLM 生成的代码通常远非完美,预测其失败的方式也颇具挑战。因此,在 IDE 中有一个辅助工具(比如 Copilot)非常有帮助。
- 当你的代码高度依赖 LLM 时,要记得编写代码的速度往往是限制因素。避免请求重复且不需要任何改动的代码,这不仅浪费能源,也会拖慢你的进度。
- LLM 很难 “记住” 它们输出的每一行代码,经常需要提醒它们当前的状态(特别是当存在跨多个模块的依赖时)。
- 对 LLM 生成的代码保持怀疑态度。尽可能多地进行验证,使用测试、可视化等手段。并且在重要的地方投入时间。相比于神经网络部分,我在 haversine 函数上花费了更多的时间(因为预期规模对性能的要求较高),对于神经网络,我更关注的是快速发现失败。
#04 LLM 与工程领域的未来
唯有变化是永恒的。
—— 赫拉克利特
在 LLM 引发的热潮和巨额资金流动的背景下,人们很容易一开始就期待完美。然而,有效利用这些工具,需要我们勇于尝试、学习并做出调整。
LLM 是否会改变软件工程团队的根本结构呢?可能吧,我们现在只是新世界的门前小道。但 LLM 已经使代码的获取变得更加民主化了。即使是没有编程经验的人,也能快速而容易地构建出功能性原型。如果你有严格的需求,将 LLM 应用在你已经熟悉的领域或许更为明智。根据我个人的经验,LLM 能够使得编写高效代码所需的时间缩短约 90%。如果你发现它们一直输出低质量的代码,那么也许是时候重新审视你的输入了。
原文链接:https://towardsdatascience.com/accelerating-engineering-with-llms-e83a524a5a13
#Mixtral 8x7B Instruct 与 Flash Attention 2 结合
感兴趣的小伙伴,可以跟着操作过程试一试。
前段时间,Mistral AI 公布的 Mixtral 8x7B 模型爆火整个开源社区,其架构与 GPT-4 非常相似,很多人将其形容为 GPT-4 的「缩小版」。
我们都知道,OpenAI 团队一直对 GPT-4 的参数量和训练细节守口如瓶。Mistral 8x7B 的放出,无疑给广大开发者提供了一种「非常接近 GPT-4」的开源选项。
在基准测试中,Mistral 8x7B 的表现优于 Llama 2 70B,在大多数标准基准测试上与 GPT-3.5 不相上下,甚至略胜一筹。
图源:https://mistral.ai/news/mixtral-of-experts/
随着这项研究的出现,很多人表示:「闭源大模型已经走到了结局。」
短短几周的时间,机器学习爱好者 Vaibhav (VB) Srivastav 表示:随着 AutoAWQ(支持 Mixtral、LLaVa 等模型的量化)最新版本的发布,现在用户可以将 Mixtral 8x7B Instruct 与 Flash Attention 2 结合使用,达到快速推理的目的,实现这一功能大约只需 24GB GPU VRAM、不到十行代码。
图源:https://twitter.com/reach_vb/status/1741175347821883502
AutoAWQ 地址:https://github.com/casper-hansen/AutoAWQ
操作过程是这样的:
首先是安装 AutoAWQ 以及 transformers:
pip install autoawq git+https://github. com/huggingface/transformers.git
第二步是初始化 tokenizer 和模型:
第三步是初始化 TextStreamer:
第四步对输入进行 Token 化:
第五步生成:
当你配置好项目后,就可以与 Mixtral 进行对话,例如对于用户要求「如何做出最好的美式咖啡?通过简单的步骤完成」,Mixtral 会按照 1、2、3 等步骤进行回答。
项目中使用的代码:
Srivastav 表示上述实现也意味着用户可以使用 AWQ 运行所有的 Mixtral 微调,并使用 Flash Attention 2 来提升它们。
看到这项研究后,网友不禁表示:真的很酷。
更多相关链接,请参考:
模型地址:https://huggingface.co/models?search=mixtral%20awq
Transformer 中量化技术:https://huggingface.co/docs/transformers/main/en/quantization
#Lightning Attention-2
Lightning Attention-2 是一种新型的线性注意力机制,让长序列的训练和推理成本与 1K 序列长度的一致。无限序列长度、恒定算力开销、更高建模精度
大语言模型序列长度的限制,极大地制约了其在人工智能领域的应用,比如多轮对话、长文本理解、多模态数据的处理与生成等。造成这一限制的根本原因在于当前大语言模型均采用的 Transformer 架构有着相对于序列长度的二次计算复杂度。这意味着随着序列长度的增加,需要的计算资源成几何倍数提升。如何高效地处理长序列一直是大语言模型的挑战之一。
之前的方法往往集中在如何让大语言模型在推理阶段适应更长的序列。比如采用 Alibi 或者类似的相对位置编码的方式来让模型自适应不同的输入序列长度,亦或采用对 RoPE 等类似的相对位置编码进行差值的方式,在已经完成训练的模型上再进行进一步的短暂精调来达到扩增序列长度的目的。这些方法只是让大模型具有了一定的长序列建模能力,但实际训练和推理的开销并没有减少。
OpenNLPLab 团队尝试一劳永逸地解决大语言模型长序列问题。他们提出并开源了 Lightning Attention-2—— 一种新型的线性注意力机制,让长序列的训练和推理成本与 1K 序列长度的一致。在遇到显存瓶颈之前,无限地增大序列长度并不会对于模型训练速度产生负面影响。这让无限长度预训练成为了可能。同时,超长文本的推理成本也与 1K Tokens 的成本一致甚至更少,这将极大地减少当前大语言模型的推理成本。如下图所示,在 400M、1B、3B 的模型大小下,随着序列长度的增加,FlashAttention2 加持的 LLaMA 的训练速度开始快速下降,然而 Lightning Attention-2 加持的 TansNormerLLM 的速度几无变化。
- 论文:Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models
- 论文地址:https://arxiv.org/pdf/2401.04658.pdf
- 开源地址:https://github.com/OpenNLPLab/lightning-attention
Lightning Attention-2 简介
让大模型的预训练速度在不同序列长度下保持一致,这听起来是一个不可能的任务。事实上,如果一个注意力机制的计算复杂度相对于序列长度保持线性关系的话,就可以实现这一点。自 2020 年线性注意力【https://arxiv.org/abs/2006.16236】横空出世以来,研究人员一直在为了线性注意力的实际效率符合它的理论线性计算复杂度而努力。在 2023 年之前,大多数的关于线性注意力的工作均集中在对齐它们与 Transformer 的精度上。终于在 2023 年中期,改进的线性注意力机制【https://arxiv.org/abs/2307.14995】在精度上可以与最先进的 Transformer 架构对齐。然而,线性注意力中将计算复杂度变成线性的最关键的 “左乘变右乘” 的计算 Trick (如下图所示),在实际实现中远慢于直接左乘的算法。其原因在于右乘的实现需要用到包含大量循环操作的累积求和(cumsum),大量的 IO 操作使得右乘的效率远低于左乘。
为了更好的理解 Lightning Attention-2 的思路,让我们先回顾下传统 softmax attention 的计算公式:O=softmax ((QK^T)⊙M_) V,其中 Q, K, V, M, O 分别为 query, key, value, mask 和输出矩阵,这里的 M 在单向任务(如 GPT)中是一个下三角的全 1 矩阵,在双向任务(如 Bert)中则可以忽略,即双向任务没有 mask 矩阵。
作者将 Lightning Attention-2 的整体思路总结为以下三点进行解释:
1. Linear Attention 的核心思想之一就是去除了计算成本高昂的 softmax 算子,使 Attention 的计算公式可以写为 O=((QK^T)⊙M_) V。但由于单向任务中 mask 矩阵 M 的存在,使得该形式依然只能进行左乘计算,从而不能获得 O (N) 的复杂度。但对于双向任务,由于没有没有 mask 矩阵,Linear Attention 的计算公式可以进一步简化为 O=(QK^T) V。Linear Attention 的精妙之处在于,仅仅利用简单的矩阵乘法结合律,其计算公式就可以进一步转化为:O=Q (K^T V),这种计算形式被称为右乘,相对应的前者为左乘。通过图 2 可以直观地理解到 Linear Attention 在双向任务中可以达到诱人的 O (N) 复杂度!
2. 但是随着 decoder-only 的 GPT 形式的模型逐渐成为 LLM 的事实标准,如何利用 Linear Attention 的右乘特性加速单向任务成为了亟待解决的难题。为了解决这个问题,本文作者提出了利用 “分而治之” 的思想,将注意力矩阵的计算分为对角阵和非对角阵两种形式,并采用不同的方式对他们进行计算。如图 3 所示,Linear Attention-2 利用计算机领域常用的 Tiling 思想,将 Q, K, V 矩阵分别切分为了相同数量的块 (blocks)。其中 block 自身(intra-block)的计算由于 mask 矩阵的存在,依然保留左乘计算的方式,具有 O (N^2) 的复杂度;而 block 之间(inter-block)的计算由于没有 mask 矩阵的存在,可以采用右乘计算方式,从而享受到 O (N) 的复杂度。两者分别计算完成后,可以直接相加得到对应第 i 块的 Linear Attention 输出 Oi。同时,通过 cumsum 对 KV 的状态进行累积以在下一个 block 的计算中使用。这样就得到了整个 Lightning Attention-2 的算法复杂度为 intra-block 的 O (N^2) 和 inter-block 的 O (N) 的 Trade-off。怎么取得更好的 Trade-off 则是由 Tiling 的 block size 决定的。
3. 细心的读者会发现,以上的过程只是 Lightning Attention-2 的算法部分,之所以取名 Lightning 是因为作者充分考虑了该算法过程在 GPU 硬件执行过程中的效率问题。受到 FlashAttention 系列工作的启发,实际在 GPU 上进行计算的时候,作者将切分后的 Q_i, K_i, V_i 张量从 GPU 内部速度更慢容量更大的 HBM 搬运到速度更快容量更小的 SRAM 上进行计算,从而减少大量的 memory IO 开销。当该 block 完成 Linear Attention 的计算之后,其输出结果 O_i 又会被搬回至 HBM。重复这个过程直到所有 block 被处理完毕即可。
想要了解更多细节的读者可以仔细阅读本文中的 Algorithm 1 和 Algorithm 2,以及论文中的详细推导过程。Algorithm 以及推导过程都对 Lightning Attention-2 的前向和反向过程进行了区分,可以帮助读者有更深入的理解。
Lightning Attention-2 精度对比
研究人员首先在小规模(400M)参数模型上对比了 Lightning Attention-2 与 Lightning Attention-1 的精度区别,如下图所示,二者几无差别。
随后研究人员在 1B、3B 上将 Lightning Attention-2 加持的 TransNormerLLM(TNL-LA2)与其它先进的非 Transformer 架构的网络以及 FlashAttention2 加持的 LLaMA 在相同的语料下做了对比。如下图所示,TNL-LA2 与 LLaMA 保持了相似的趋势,并且 loss 的表现更优。这个实验表明,Lightning Attention-2 在语言建模方面有着不逊于最先进的 Transformer 架构的精度表现。
在大语言模型任务中,研究人员对比了 TNL-LA2 15B 与 Pythia 在类似大小下的大模型常见 Benchmark 的结果。如下表所示,在吃掉了相同 tokens 的条件下,TNL-LA2 在常识推理和多项选择综合能力上均略高于基于 Softmax 的注意力的 Pythia 模型。
Lightning Attention-2 速度对比
研究人员对 Lightning Attention-2 与 FlashAttention2 进行了单模块速度与显存占用对比。如下图所示,相比于 Lightning Attention-1 和 FlashAttention2,在速度上,Lightning Attention-2 表现出了相比于序列长度的严格线性增长。在显存占用上,三者均显示出了类似的趋势,但 Lightning Attention-2 的显存占用更小。这个的原因是 FlashAttention2 和 Lightning Attention-1 的显存占用也是近似线性的。
笔者注意到,这篇文章主要关注点在解决线性注意力网络的训练速度上,并实现了任意长度的长序列与 1K 序列相似的训练速度。在推理速度上,并没有过多的介绍。这是因为线性注意力在推理的时候可以无损地转化为 RNN 模式,从而达到类似的效果,即推理单 token 的速度恒定。对于 Transformer 来说,当前 token 的推理速度与它之前的 token 数量相关。
笔者测试了 Lightning Attention-1 加持的 TransNormerLLM-7B 与常见的 7B 模型在推理速度上的对比。如下图所示,在近似参数大小下,Lightning Attention-1 的吞吐速度是百川的 4 倍,ChatGLM 的 3.5 倍以上,显示出了优异的推理速度优势。
小结
Lightning Attention-2 代表了线性注意力机制的重大进步,使其无论在精度还是速度上均可以完美的替换传统的 Softmax 注意力,为今后越来越大的模型提供了可持续扩展的能力,并提供了一条以更高效率处理无限长序列的途径。OpenNLPLab 团队在未来将研究基于线性注意力机制的序列并行算法,以解决当前遇到的显存屏障问题。
#FontDiffuser
论文FontDiffuser: One-Shot Font Generation via Denoising Diffusion with Multi-Scale Content Aggregation and Style Contrastive Learning
,Diffusion 扩散模型用于生成任意风格的复杂字的使用配方。
一种基于扩散的Image-to-Image字体生成方法“FontDiffuser”,它创新性地将字体模仿任务建模为噪声到噪声范式。作者引入了多尺度内容聚合(MCA)模块,它有效地结合了不同尺度的全局和局部内容线索,从而增强了对复杂字的复杂笔画的保留。
- 作者: 杨振华、彭德智.
- 单位: 华南理工大学DLVC实验室、阿里巴巴
论文链接: https://arxiv.org/abs/2312.12142 - 项目链接: https://yeungchenwa.github.io/fontdiffuser-homepage/
- HuggingFace Demo链接: https://huggingface.co/spaces/yeungchenwa/FontDiffuser-Gradio
图1 FontDiffuser的输出结果 (复杂字+风格跨度大
图2 FontDiffuser与InstructPix2Pix的联合使用效果
摘要
字体生成是一项模仿任务,其目的是创建一个字体库,模仿参考图像(风格图像)的风格,同时保留源图像(内容图像)的内容,如下图3所示。虽然现有方法已经取得了令人满意的性能,但它们在处理复杂字符和风格变化较大的字符时仍然很吃力。
为了解决这些问题,我们提出了一种基于扩散的Image-to-Image字体生成方法“FontDiffuser”,它创新性地将字体模仿任务建模为噪声到噪声范式。在我们的方法中,我们引入了多尺度内容聚合(MCA)模块,它有效地结合了不同尺度的全局和局部内容线索,从而增强了对复杂字的复杂笔画的保留。
此外,为了更好地处理具有巨大差异的风格转换,我们提出了风格对比细化(SCR)模块,这是一种新颖的风格表征对比学习策略,它利用风格提取器将风格从图像中分离出来,然后通过精心设计的风格对比损失对扩散模型进行监督。
广泛的实验证明,FontDiffuser 在生成多样化字符和风格方面具有一流的性能。与以前的方法相比,它在复杂字符和大的风格变化方面始终表现出色。
图3 字体生成任务定义
研究动机
图4 其他方法在复杂字和风格跨度大上的生成效果
存在的问题
现有的字体生成方法虽然取得了令人满意的性能,但在处理复杂字和风格变化较大的字符(尤其是中文字符)时,仍会出现严重的笔画缺失、伪影、模糊、结构布局错误和风格不一致等问题,如上图4所示。
原因分析
- 大多数方法都采用基于 GAN 的框架,由于其对抗训练的性质,可能会出现训练不稳定的问题。
- 这些方法大多只通过单一尺度的高维特征来感知内容信息,而忽略了对保留源内容(尤其是复杂字符)的细粒度细节。
- 许多方法利用先验知识来帮助字体生成,例如字符的笔画或部件组成;然而,对于复杂的字符来说,获取这些细粒度信息的成本很高;
- 在过去的方法中,目标风格通常由一个简单的分类器或判别器来进行特征表示学习,这种分类器或判别器很难学习到合适的风格,在一定程度上阻碍了在风格变化较大时的风格转换。
采用策略
- 我们提出了一种基于扩散的图像到图像(Image-to-Image)的字体生成方法,名字叫FontDiffuser,它将字体生成学习建模为加噪到去噪的范式,能够生成未见过的字符和风格。
- (1) 我们引入了多尺度内容聚合(MCA)模块,它利用了不同尺度的全局和局部内容特征。
- (2) 我们引入了一种新颖的风格表征对比学习策略,通过利用风格对比细化(SCR)模块来增强生成器模仿风格的能力。
FontDiffuser模型框架
图5 FontDiffuser的模型框架图
FontDiffuser的模型框架如图5所示,其中包括字体生成条件扩散模型和风格细化对比模块两部分。在条件扩散字体生成模型部分,我们的目标是,在给定一张源图片(内容图片)和一张参考图片(风格图片),该模型的输出结果能够既与源图片的字符内容相同,也能与参考图片的风格保持一致;在风格细化对比学习模块,起主要目标是在训练的过程中,能够将一组图片的风格分离出来并完成辨别,最后通过风格对比损失函数来为条件扩散模型提供监督信息。
条件扩散字体生成模型
- 多尺度内容聚合(MCA)模块
生成复杂的字符一直是一项极具挑战性的任务,现有的许多方法仅依赖于单一尺度的内容特征,而忽略了复杂的细节,如笔画和部首。如图 6 所示,大尺度特征保留了大量细节信息,而小尺度特征则缺乏这些信息。
图6 在内容编码器中不同尺度的特征图
因此,我们设计了多尺度内容聚合(MCA)模块,将不同尺度的全局和局部内容特征注入到扩散模型的 UNet 中,如下图7所示。其中包括了一个通道注意力机制来提高自适应通道选择能力和一个Cross-Attention来引入风格特征。
图7 多尺度内容聚合(MCA)模块
- Reference-Structure Interaction (RSI)交互模块
源图像(内容图像)和目标图像之间存在结构差异(如字体大小),为了解决这个问题,我们提出了Reference-Structure(RSI)交互模块,它采用可变形卷积网络DCN对UNet 的skip-connection部分进行结构变形。具体过程可以使用以下公式表
风格细化对比学习模块
不管源图像(内容图像)和参考图像(风格图像)之间的风格差异,字体生成的目的之一是达到预期的风格模仿效果。其中的一种可行的解决方案是找到合适的风格特征表示,并进一步为我们的生成器提供监督信息。因此,我们提出了风格细化对比学习(SCR)模块,这是一个字体风格表示学习模块,它能从一组样本图像中分离出风格,并结合风格对比损失来监督我们的条件扩散模型,能够较好地与扩散模型进行较好的结合,确保生成的风格在全局和局部层面上与目标保持一致,具体的做法如上图5(b)所示,用公式表示如下:
训练损失函数
我们的训练采用了一个从粗到细的两阶段策略,具体包括:
- Phase 1
我们主要使用标准的 MSE 扩散损耗来优化 FontDiffuser,而不使用 SCR 模块。这确保我们的生成器具备字体重构的基本能力:
更多可视化结果
图8 困难难度字符
中等难度字符 (medium)
图9 中等难度字符
容易难度字符 (easy)
图10 容易难度字符
跨语言生成 (中文到韩文)
图11 跨语言生成 (中文到韩文)
实验
为了验证生成不同复杂度字符的有效性,我们根据中文字符的笔画数将其分为三个复杂度级别(easy, medium, hard),并在每个级别上分别测试我们的方法。不同复杂度的中文字符样例如下:
图12 三种不同复杂度的中文字符样例
定量结果 (Quantitative Results)
图13 在SFUC和UFUC上的定量结果
图14 在UFSC上的定量结果
定性结果 (Qualitative Results)
图15 在SFUC和UFUC上的定性结果
跨语言生成 (中文到韩文)
FontDiffuser还能在仅有中文数据集的训练的情况下,用于韩文的生成,结果如下图11所示:
图16 跨语言生成 (中文到韩文)
结论分析
我们的方法在不同的在三个复杂度级别的字符上,FontDiffuser 的表现优于目前的其他方法。此外,FontDiffuser 还证明了它在跨语言字体生成任务(如中文到韩文)中的适用性,凸显了其良好的跨领域能力。
消融实验
- 不同模块的有效性
图17 使用不同模块后的性能对比
图18 使用不同模块后的可视化效果对比
- SCR模块中数据增广的有效性
图19 SCR模块中数据增广的有效性
- RSI模块中基于cross-attention和CNN的交互方式的性能对比
图20 RSI模块中基于cross-attention和CNN的交互方式的性能对比
总结
在本文中,我们提出了一种基于扩散的图像到图像(Image-to-Image)的字体生成方法,称为FontDiffuser,它在生成复杂字符和处理差异较大的风格转换方面表现出色。
此外,我们还提出了一种新颖的风格特征对比学习策略,该策略提出了一个 SCR 模块,并使用风格对比损失来监督我们的扩散模型。此外,我们还采用了 RSI 模块,利用参考图像的结构特征促进结构变形。
大量实验证明,在三个复杂度级别的字符上,FontDiffuser 的表现优于目前的方法。此外,FontDiffuser 还证明了它在跨语言字体生成任务(如中文到韩文)中的适用性,凸显了其良好的跨领域能力。
demo
图21 demo演示
#AppAgent
近日,一项名为 AppAgent 的创新技术引起了广泛关注。这就是2024年的新趋势吗?大模型竟然能玩手机了,还能用软件修图
简单来说,AppAgent 的智能代理能力可以用于操作任何 App,它在 50 个复杂手机任务上展示了强大的能力。
AppAgent 幕后团队来自腾讯,相关技术论文也已经公布。
- 论文地址:https://arxiv.org/abs/2312.13771
- 主页地址:https://appagent-official.github.io/
据论文介绍,这项技术通过引入一种基于大型语言模型(LLMs)的多模态智能代理(Agent)框架,赋予了智能体操作智能手机应用的能力。与传统的智能助手如 Siri 不同,AppAgent 不依赖于系统后端访问,而是通过模拟人类的点击和滑动等操作,直接与手机应用的图形用户界面(GUI)互动。这种独特的方法不仅提高了安全性和隐私性,还确保了智能体能够适应应用界面的变化和更新。
AppAgent 的核心在于其创新的学习方式。智能体可以通过自主探索或观察人类演示来学习如何导航和使用新应用。在自主探索过程中,智能体通过一系列预定义的动作与应用互动,观察每个动作带来的界面变化,从而构建知识库。这个过程还可以通过观察少数几个人类演示来加速,使智能体能够更快地理解复杂功能。
为了验证其实用性,研究团队对 AppAgent 进行了广泛测试,覆盖了社交媒体、电子邮件、地图、购物以及复杂的图像编辑工具等 10 种不同的应用中的 50 个任务。测试结果证明,AppAgent 在处理各种高级任务方面表现出色,显示出其适应性、用户友好性以及高效的学习和操作能力。
推特大 V Andrew Torba 评价道:「一切都结束了,自 2017 年以来整个互联网都已经是 AI 的天下。互联网已死的理论是真实的。构建并推动尽可能多的基于 AI 的机器人来充斥网络空间,用红色药丸 (来自电影《黑客帝国》) 淹没它们是唯一的出路。」
推特用户 Ate-a-Pi 评价道:「它来了!苹果的创新者困境。如果代理完成所有工作,那么优秀的智能手机用户界面有什么意义呢? 如果代理使得智能手机操作系统变得多余,苹果将怎么做?他们有多长时间来解决这个问题?」
推特用户 Morally Clear 评价道:「AI 代理将会成为 2024 年的 ChatGPT。」
也有很多用户表示了对此项技术的担心。
推特用户 Benjamin Cox 评价道:「这太疯狂了。我们所知的应用程序可能仅仅会变成集成。此外,机器人用户参与度指标将变得完全不可靠。」
推特用户 Karma 表示了这项技术可能被用来非法牟利:“建立一个拥有 500,000 台这样的服务器农场。点击你在 Instagram 或 Twitter 上看到的每个广告。短期投资你点击的广告公司并购买 Meta、X 和其他社交媒体应用的看涨期权?利润?谁在进行这项工作?”
推特用户 Logan Thorneloe 评价道:“这太酷了!但是我知道它会被用来制造水军机器人,这让我有些担心。”
方法概述
环境搭建
AppAgent 的实验环境是基于命令行界面(CLI),使得代理能与安卓系统上的智能手机应用进行交互。代理接收两种关键输入:一是展示应用界面的实时屏幕截图,二是详细描述交互元素的 XML 文件。为了提升代理识别和交互这些元素的能力,每个元素都被赋予了一个唯一标识符。这些标识符要么来源于 XML 文件中的资源 ID(如果提供),要么通过结合元素的类名、大小和内容来构建。这些以半透明数字形式覆盖在屏幕截图上的元素,帮助代理在不需指定精确位置的情况下准确互动,从而提升控制手机的精确度。
动作空间模拟了人类与智能手机的常见交互方式,包括点击和滑动。设计了四个基本功能:点击(Tap)、长按(Long_press)、滑动(Swipe)和文本输入(Text),以及两个系统级功能:返回(Back)和退出(Exit)。这些预定义动作旨在简化代理的交互流程,并减少对精确屏幕坐标的依赖,解决了语言模型在准确预测中可能遇到的挑战。
App 探索阶段
探索阶段是 AppAgent 框架的核心,代理通过自主交互或观察人类演示来学习应用程序的功能和特性。在自主交互模式下,代理被分配一个任务并开始与 UI 元素进行自主互动。它尝试不同的动作,并观察应用界面的变化以理解其工作原理。代理通过分析每个动作前后的屏幕截图,尝试弄清楚 UI 元素的功能和特定动作的效果,并将这些信息编译成文档,记录下不同元素所执行动作的效果。当一个 UI 元素被多次操作时,代理会根据之前的文档和当前的观察来更新信息,以提高认知质量。
为了提高探索效率,如果当前 UI 页面似乎与应用的主要任务无关(如广告页面),代理将停止进一步探索并使用 Android 系统的返回功能返回到前一个 UI 页面。这种目标导向的探索方法,相比随机探索(如深度优先搜索和广度优先搜索),确保代理专注于对应用有效操作至关重要的元素。此外,代理还利用语言模型的现有关于用户界面的知识来提高探索效率,直到完成分配的任务。
在通过观察人类演示进行探索的方式中,代理通过观察人类用户的操作来学习应用的复杂功能,这对于那些可能难以通过自主交互发现的功能尤其有效。在此方法中,代理记录人类使用的元素和动作,这种策略缩小了探索空间,并阻止代理与无关的应用页面进行交互,从而比自主交互更为高效和有条理。
AppAgent 通过观察 app 的操作响应来理解 UI 界面功能和操作逻辑,并整理成文档
部署阶段
经过探索阶段的训练,AppAgent 已准备好执行基于其累积经验的复杂任务。在部署阶段,代理遵循逐步方法,每一步都包括获取当前用户界面的屏幕截图和一个动态生成的文档,详细描述了用户界面元素的功能和当前用户界面页面上动作的效果。代理首先对当前用户界面进行观察,然后阐述其关于任务和当前观察的思考过程。随后,代理通过调用可用函数执行动作。每次动作后,代理总结交互历史和当前步骤中采取的动作。这些信息被整合进下一个提示中,为代理提供了一种记忆形式。这种细致的方法提高了代理动作的可靠性和可解释性,从而促进了更明智的决策。部署阶段持续进行,直到代理确定任务已经完成,此时它可以通过执行退出(Exit)动作结束过程。
部署阶段,AppAgent 参考整理好的文档,不断的观察思考总结来完成复杂的任务。
实验结论
AppAgent 通过一系列的量化和定性实验,证明了其在操控多样化智能手机应用方面的显著能力。该评估涉及了 10 种广泛使用的应用,覆盖了从社交媒体、地图导航到音乐播放和图片编辑等多个领域。通过特别针对 Adobe Lightroom 这一图像编辑应用的深入案例研究,AppAgent 的视觉处理能力得到了详尽的考察。
AppAgent 采用了先进的多模态大型语言模型 GPT-4 来处理交错的图像和文本输入,这种独特的能力使其能够无缝地解释和互动应用中的视觉及文本信息。实验结果表明,通过简化动作空间的设计,AppAgent 在准确性和效率上大幅度超越了原始的 GPT-4 模型。这主要是因为简化动作空间消除了对生成精确 xy 坐标的需求,这是传统语言模型的一个挑战点。在成功率、奖励和平均步骤数三个关键性能指标上,AppAgent 表现优异。即便在任务步骤上遇到失败,它也能根据其最终状态获得一定的奖励分数,这反映出其在理解和执行任务方面的适应性和韧性。特别是在多样化的应用中,AppAgent 通过有效地完成任务,展现了其出色的操作能力。
通过自主探索和观察人类演示生成的文档,AppAgent 的性能显著优于仅依赖原始 GPT-4 模型的基线。这些文档的有效性与人工编写的文档相媲美,凸显了设计在增强代理跨多样化应用表现的有效性。定性结果进一步证明了 AppAgent 在准确感知、推理和响应任务要求方面的能力。
在 Adobe Lightroom 的案例研究中,AppAgent 对具有不同视觉问题的图片进行编辑,展现了其处理视觉任务的高级能力。用户研究的结果显示,相较于 GPT-4 基线,AppAgent 在图像编辑质量上有了明显提升。尤其是在使用观察演示生成的文档时,AppAgent 倾向于使用更多工具来提升图像质量,而 GPT-4 基线通常使用较少的工具。
总而言之,AppAgent 在多项任务中展示了出色的性能和适应性。其创新的多模态框架和有效的探索策略,不仅验证了 AppAgent 的实用性,还为智能代理在理解和操作多样化智能手机应用方面的研究提供了宝贵的见解和基准。
用 AppAgent 操作 Lightroom 修图软件修图的例子
总结
腾讯的 AppAgent 技术标志着智能代理领域的一个重大进步,不仅提升了多模态大型语言模型在实际应用中的能力,而且开辟了未来智能手机应用操作和交互的新境界。这项技术的核心优势在于其创新性和灵活性:通过模拟人类的互动方式,AppAgent 能够直接与应用的图形用户界面进行交互,从而摆脱了对传统后端接口的依赖。这不仅提高了操作的安全性和隐私性,而且使得智能体能够更加灵活地适应应用界面的变化和更新。
AppAgent 的多模态学习框架,结合自主探索和观察人类演示的方式,使其能够快速适应和掌握新应用,无论是处理复杂的图像编辑任务,还是导航复杂的社交媒体平台。这种学习方法的高效性和适应性,证明了 AppAgent 在理解和执行多样化任务方面的高级能力,同时也突显了其在未来应用开发和用户体验改善中的巨大潜力。
随着技术的不断进步和优化,我们可以期待 AppAgent 在多个领域中的应用,如提高工作效率、改善用户体验、辅助残障人士等。AppAgent 的成功不仅为智能代理技术的发展树立了新的里程碑,而且为我们如何与智能手机和其他智能设备互动提供了全新的视角。它不仅有望改变个人用户的日常生活,还可能对商业、教育和互联网等行业产生深远影响。
#DDPM~2
第2集 详解扩散模型的基石DDPM,附详细模型架构图以及模型运作流程的详细解说。
大家好,生成式大模型的热度,从今年3月开始已经燃了一个多季度了。在这个季度中,相信大家肯定看过很多AI产生的有趣内容,比如著名的抓捕川普现场与监狱风云 [现在来看不仅是画得像而已],AI换声孙燕姿等。这背后都用到了一个强大的模型:Diffusion Model。所以,在这个系列中,我们将从原理到源码,从基石DDPM到DALLE2,Imagen与Stable Diffusion,通过详细的图例和解说,和大家一起来了解扩散模型的奥秘。同时,也会穿插对经典的GAN,VAE等模型的解读,敬请期待~
本篇将和大家一起解读扩散模型的基石:DDPM(Denoising Diffusion Probalistic Models)。 扩散模型的研究并不始于DDPM,但DDPM的成功对扩散模型的发展起到至关重要的作用。在这个系列里我们也会看到,后续一连串效果惊艳的模型,都是在DDPM的框架上迭代改进而来。所以,我把DDPM放在这个系列的第一篇进行讲解。初读DDPM论文的朋友,可能有以下两个痛点:
- 论文花极大篇幅讲数学推导,可是我看不懂。
- 论文没有给出模型架构图和详细的训练解说,而这是我最关心的部分。
针对这些痛点,DDPM系列将会出如下三篇文章:
- DDPM(模型架构篇):也就是本篇文章。在阅读源码的基础上,本文绘制了详细的DDPM模型架构图,同时附上关于模型运作流程的详细解说。本文不涉及数学知识,直观帮助大家了解DDPM怎么用,为什么好用。
- DDPM(人人都能看懂的数学推理篇):DDPM的数学推理可能是很多读者头疼的部分。我尝试跳出原始论文的推导顺序和思路,从更符合大家思维模式的角度入手,把整个推理流程串成一条完整的逻辑线。同样,我也会配上大量的图例,方便大家理解数学公式。如果你不擅长数学推导,这篇文章可以帮助你从直觉上了解DDPM的数学有效性;如果你更关注推导细节,这篇文章中也有详细的推导中间步骤。
- DDPM(源码解读篇):在前两篇的基础上,我们将配合模型架构图,一起阅读DDPM源码,并实操跑一次,观测训练过程里的中间结果。
一、DDPM在做一件什么事
假设你想做一个以文生图的模型,你的目的是给一段文字,再随便给一张图(比如一张噪声),这个模型能帮你产出符合文字描述的逼真图片,例如:
文字描述就像是一个指引(guidance),帮助模型去产生更符合语义信息的图片。但是,毕竟语义学习是复杂的。我们能不能先退一步,先让模型拥有产生逼真图片的能力?
比如说,你给模型喂一堆cyperpunk风格的图片,让模型学会cyberpunk风格的分布信息,然后喂给模型一个随机噪音,就能让模型产生一张逼真的cyberpunk照片。或者给模型喂一堆人脸图片,让模型产生一张逼真的人脸。同样,我们也能选择给训练好的模型喂带点信息的图片,比如一张夹杂噪音的人脸,让模型帮我们去噪。
具备了产出逼真图片的能力,模型才可能在下一步中去学习语义信息(guidance),进一步产生符合人类意图的图片。而DDPM的本质作用,就是学习训练数据的分布,产出尽可能符合训练数据分布的真实图片。所以,它也成为后续文生图类扩散模型框架的基石。
二、DDPM训练流程
理解DDPM的目的,及其对后续文生图的模型的影响,现在我们可以更好来理解DDPM的训练过程了。总体来说,DDPM的训练过程分为两步:
- Diffusion Process (又被称为Forward Process)
- Denoise Process(又被称为Reverse Process)
前面说过,DDPM的目的是要去学习训练数据的分布,然后产出和训练数据分布相似的图片。那怎么“迫使”模型去学习呢?
一个简单的想法是,我拿一张干净的图,每一步(timestep)都往上加一点噪音,然后在每一步里,我都让模型去找到加噪前图片的样子,也就是让模型学会去噪。这样训练完毕后,我再塞给模型一个纯噪声,它不就能一步步帮我还原出原始图片的分布了吗?
一步步加噪的过程,就被称为Diffusion Process;一步步去噪的过程,就被称为Denoise Process。我们来详细看这两步。
2.1 Diffusion Process
Diffusion Process的命名受到热力学中分子扩散的启发:分子从高浓度区域扩散至低浓度区域,直至整个系统处于平衡。加噪过程也是同理,每次往图片上增加一些噪声,直至图片变为一个纯噪声为止。整个过程如下:
如图所示,我们进行了1000步的加噪,每一步我们都往图片上加入一个高斯分布的噪声,直到图片变为一个纯高斯分布的噪声。
我们记:
2.2 Denoise Process
你可能想问:
- 为什么我们的输入中要包含time_step?
- 为什么通过预测噪声的方式,就能让模型学得训练数据的分布,进而产生逼真的图片?
第二个问题的答案我们同样放在下一篇(数学推理篇)中进行详解。而对于第一个问题,由于模型每一步的去噪都用的是同一个模型,所以我们必须告诉模型,现在进行的是哪一步去噪。因此我们要引入timestep。timestep的表达方法类似于Transformer中的位置编码,将一个常数转换为一个向量,再和我们的输入图片进行相加。
注意到,UNet模型是DDPM的核心架构,我们将关于它的介绍放在本文的第四部分。
到这里为止,如果不考虑整个算法在数学上的有效性,我们已经能从直觉上理解扩散模型的运作流程了。那么,我们就可以对它的训练和推理过程来做进一步总结了。
三、DDPM的Training与Sampling过程
3.1 DDPM Training
上图给出了DDPM论文中对训练步骤的概述,我们来详细解读它。
前面说过,DDPM模型训练的目的,就是给定time_step和输入图片,结合这两者去预测图片中的噪声。
上面演示的是单条数据计算loss的过程,当然,整个过程也可以在batch范围内做,batch中单条数据计算loss的方法不变。
3.2 DDPM的Sampling
到目前为止,我们已经把整个DDPM的核心运作方法讲完了。接下来,我们来看DDPM用于预测噪声的核心模型:UNet,到底长成什么样。我在学习DDPM的过程中,在网上几乎找不到关于DDPM UNet的详细模型解说,或者一张清晰的架构图,这给我在源码阅读过程中增加了难度。所以在读完源码并进行实操训练后,我干脆自己画一张出来,也借此帮助自己更好理解DDPM。
四、DDPM中的Unet架构
UNet模型最早提出时,是用于解决医疗影响诊断问题的。总体上说,它分成两个部分:
- Encoder
- Decoder
在Encoder部分中,UNet模型会逐步压缩图片的大小;在Decoder部分中,则会逐步还原图片的大小。 同时在Encoder和Deocder间,还会使用 “残差连接”,确保Decoder部分在推理和还原图片信息时,不会丢失掉之前步骤的信息。整体过程示意图如下,因为压缩再放大的过程形似"U"字,因此被称为UNet:
那么DDPM中的UNet,到底长什么样子呢?我们假设输入为一张32_32_3大小的图片,来看一下DDPM UNet运作的完整流程:
如图,左半边为UNet的Encoder部分,右半边为UNet的Deocder部分,最下面为MiddleBlock。我们以从上往下数第二行来分析UNet的运作流程。
在Encoder部分的第二行,输入是一个16*16*64
的图片,它是由上一行最右侧32*32*64
的图片压缩而来(DownSample)。对于这张16*16*64
大小的图片,在引入time_embedding后,让它们一起过一层DownBlock,得到大小为16*16*128
的图片。再引入time_embedding,再过一次DownBlock,得到大小同样为16*16*128
的图片。对该图片做DowSample,就可以得到第三层的输入,也就是大小为8*8*128
的图片。由此不难知道,同层间只做channel上的变化,不同层间做图片的压缩处理。至于每一层channel怎么变,层间size如何调整,就取决于实际训练中对模型的设定了。Decoder层也是同理。其余的信息可以参见图片,这里不再赘述。
我们再详细来看右下角箭头所表示的那些模型部分,具体架构长什么样:
4.1 DownBlock和UpBlock
如果你曾在学习DDPM的过程中,困惑time_embedding要如何与图片相加,Attention要在哪里做,那么这张图可以帮助你解答这些困惑。TimeEmbedding层采用和Transformer一致的三角函数位置编码,将常数转变为向量。Attention层则是沿着channel维度将图片拆分为token,做完attention后再重新组装成图片(注意Attention层不是必须的,是可选的,可以根据需要选择要不要上attention)。
你可能想问:一定要沿着channel方向拆分图片为token吗?我可以选择VIT那样以patch维度拆分token,节省计算量吗?当然没问题,你可以做各种实验,这只是提供DDPM对图片做attention的一种方法。
4.2 DownSample和UpSample
这个模块很简单,就是压缩(Conv)和放大(ConvT)图片的过程。对ConvT原理不熟悉的朋友们,可以参考这篇文章()。
4.3 MiddleBlock
和DownBlock与UpBlock的过程相似,不再赘述。
到这一步,我们就把DDPM的模型核心给讲完啦。在第三篇源码解读中,我们会结合这些架构图,来一起阅读DDPM training和sampling代码。
五、文生图模型的一般公式
讲完了DDPM,让我们再回到开头,看看最初我们想训练的那个“以文生图”模型吧!
当我们拥有了能够产生逼真图片的模型后,我们现在能进一步用文字信息去引导它产生符合我们意图的模型了。通常来说,文生图模型遵循以下公式(图片来自李宏毅老师课堂PPT):
- Text Encoder: 一个能对输入文字做语义解析的Encoder,一般是一个预训练好的模型。在实际应用中,CLIP模型由于在训练过程中采用了图像和文字的对比学习,使得学得的文字特征对图像更加具有鲁棒性,因此它的text encoder常被直接用来做文生图模型的text encoder(比如DALLE2)
- Generation Model:输入为文字token和图片噪声,输出为一个关于图片的压缩产物(latent space)。这里通常指的就是扩散模型,采用文字作为引导(guidance)的扩散模型原理,我们将在这个系列的后文中出讲解。
- Decoder:用图片的中间产物作为输入,产出最终的图片。Decoder的选择也有很多,同样也能用一个扩散模型作为Decoder。
5.1 DALLE2
DALLE2就套用了这个公式。它曾尝试用Autoregressive和Diffusion分别来做Generation Model,但实验发现Diffusion的效果更好。所以最后它的2和3都是一个Diffusion Model。
Stable Diffusion
大名鼎鼎Stable Diffsuion也能按这个公式进行拆解。
5.3 Imagen
Google的Imagen,小图生大图,遵循的也是这个公式。
按这个套路一看,是不是文生图模型,就不难理解了呢?我们在这个系列后续文章中,也会对这些效果惊艳的模型,进行解读。
参考
1、https://arxiv.org/abs/2006.11239
2、https://arxiv.org/abs/2204.06125
3、https://arxiv.org/abs/2112.10752
4、https://arxiv.org/abs/2205.11487
5、https://speech.ee.ntu.edu.tw/~hylee/ml/ml2023-course-data/StableDiffusion%20(v2).pdf
6、https://speech.ee.ntu.edu.tw/~hylee/ml/ml2023-course-data/StableDiffusion%20(v2).pdf
7、https://github.com/labmlai/annotated_deep_learning_paper_implementations/tree/master/labml_nn/diffusion