演讲嘉宾|张君
编辑 |李忠良
策划 |AICon 全球人工智能开发与应用大会
随着大模型技术的快速发展,其在 LLM、多模态融合等领域的应用越来越广泛。然而,大模型的高效推理仍然是一个关键挑战,从计算复杂度、内存占用、通信技术等各个技术层面展开,如何在保证性能的同时降低计算成本、提升推理效率成为了关键挑战。
在 InfoQ 举办的 AICon 全球人工智能开发与应用大会上华为高级开发工程师张君做了专题演讲“华为昇腾推理技术的优化实践”,演讲围绕大模型推理优化的技术发展方向,围绕模型层、推理框架层、算子层这 3 个方面展开,并结合实践案例,阐述相关的技术方案和选型,帮助听众更好地理解和应用大模型推理技术。
内容亮点
了解当前华为昇腾推理技术的优化实践
以下为演讲内容整理。
大模型推理的现状及挑战
大模型推理的现状及挑战主要分为两个部分:prefill 阶段和 decode 阶段。在 prefill 阶段,计算资源是受限因素,而在解码阶段,缓存资源则成为瓶颈。这两个阶段存在一定的权衡,主要体现在时延和吞吐量的问题上。
具体来说,解码阶段需要更快的速度以降低时延,而为了提高吞吐量,我们则需要增大 batch size,即在预览阶段让更多的序列或请求进入。然而,在这两个部分之间寻求平衡时,又会面临一些问题,主要集中在两大方面。
现状 1:模型规模增大及自回归解码带来访存及算力利用率压力
当前的模型规模较大,这导致内存容量和缓存成为瓶颈,影响了解码阶段的性能。当内存不足时,虽然有一些常见的解决方案,例如在模型参数增大时,KV Cache 会随着 Batch size 或序列长度的增长而占用更多内存,但这又带来了新的挑战。
例如,我们可能会采用多卡多机的分布式计算方式,或者面临访谈断宽的问题,即每秒 token 数达到 50 时,虽然从人的视觉上看不会感觉卡顿,但实际上会增加多路并发的带宽压力,进而导致吞吐量的延时增加。
此外,自回归模型在低时延时难以兼顾算力利用率的问题也较为突出,即在预览和解码两个阶段,由于阶段差异较大,解码阶段难以充分利用算力资源,而很多算法的重点就在于如何充分利用算力以提高速度。再者,解码阶段每 token 串行解码的算力利用率较低,主要是以 GEMV 为主,计算访存比也相对较低。
现状 2:KV Cache 导致"内存墙”瓶颈进一步加剧
KV Cache 是模型推理中的关键部分,但当其数量增多时,会进一步加剧内存瓶颈。如果不采用 KV Cache,而是进行全量计算,随着序列长度的增加,计算量会呈指数级增长。
例如,attention 的计算量与序列长度呈平方关系。以 Llama 2 70B 模型为例,当序列长度在 1 兆以下时,若将延时控制在 52 秒以内,算力利用率可达 50%,但需要更多的卡来实现。
而如果采用 KV Cache,推理的内存开销会呈线性增长,KV Cache 越大,内存占用率越高。同样以 Llama 72B 模型为例,在 1 兆以下序列长度时,使用 KV Cache 并将延时控制在 52 秒以内,整个卡的消耗大约在 18 NPU 卡左右,但显存占用率会显著上升。
大模型推理常用加速技术
针对上述问题,业界已经发展出一些加速技术。在算子层,常见的有 QKV 大融合算子和 Flush attention 等融合算子。在算法层,有分片策略优化和投机推理等方法。量化也是常用的手段,效果较好。
在框架层有 Page attention、Continuous batch 等优化措施,以及 PD 分离部署的策略。由于我们在算子层有自己的算子开发语言,能够较好地掌控融合算子等操作,因此本次重点介绍算子层的两个优秀实践案例。
昇腾硬件亲和的 FA 融合算子性能优化实践
在昇腾硬件环境下,我们面临着 Victor 能力不足的问题,而需要借助 Cube 来补齐。这一问题的来源背景是,我们所采用的 FA 或 PA 算法依赖于 online Softmax,这是一种切块或动态的方法,用于计算序列的一部分。该算法通常与 FA 联合使用,主要目的是对 KV Cache 进行计算,以提高计算效率并减少内存占用。
online Softmax 涉及一些小算子,如 exp、sub 或 Mul 等,其中 exp 和 sub 属于向量操作。在生成过程中,这些向量操作通常会在 Victor 计算单元中进行。但在我们的奥特莱斯 300I 卡上,Victor 的算力相对较低,导致 FA 的性能受到影响。
经过算子性能分析,我们发现 Victor bound,即 Victor 的性能瓶颈较大,Victor 占用的时间占到了总时间的 90% 左右。通过进一步分析,我们注意到 mul 和 add 操作可以放在矩阵上进行,也就是在 cube 单元上执行,因为 cube 单元的计算能力较强。因此,我们考虑在 Victor 能力不足的情况下,改造算子或算法的实现,将其提升到 Cube 上运行。
我介绍一下我们的 AI Core 架构。在奥特莱斯 300I 上,我们的架构如图所示,Cube 和 Victor 是同核部署的。一个 AI Core 包含 Cube 计算单元、Scala 单元和 Victor 计算单元。Cube 计算单元和 Victor 计算单元共享一个 Scala 单元,这种架构是耦合的,与 Atlas 800 A2 上的分离架构不同。
Scala 负责各类数据类型的基本运算和程序流程控制,例如 if else 等,可以看作是一个很小的 CPU,进行流程的分发。Victor 则主要用于向量运算,如向量加法、乘法等,能够快速完成例如 FB16 类型的相加或相乘操作,并支持多迭代执行。
Cube 负责矩阵运算,在一个时钟周期内可以完成矩阵 m 乘 k 或矩阵 k 乘 n 的运算。其主要单元包括 L0A、L0B 和 L0C,分别用于存储左矩阵、右矩阵和结果矩阵或中间结果。
回到我们的问题,即在 Victor 算力不足的情况下,如何用 Cube 代替。我们先来看 Flash attention 算法中的 online Softmax 公式:
在实施第一步后,我们发现了一个“假”VecBound 场景,即 Victor 的耗时不降反增。经过分析,虽然 Victor 的 online Softmax 耗时下降了 300,但我们在构造矩阵或对角矩阵时引入了额外成本,导致时间增加。
这说明,虽然我们把部分运算移到了 Cube 上,但 Victor 的构造逻辑也需要优化。而且,我们还发现总耗时大于 Victor 耗时的增加时间,优化后整个过程增加了 60 毫秒左右。这表明,除了优化对角矩阵的构造,我们还需要关注方案下的流水调整,无论是 Cube 还是 MTE,都需要进行 Victor 的覆盖。
在对代码和流水进行深入分析后,我发现原来的代码结构在开启双缓冲区的情况下存在一些问题。双缓冲区的原理是将数据流分为两部分,例如 Tensor 1 和 Tensor 2。在计算 Tensor 1 的同时,我们可以进行 Tensor 2 的数据搬运。
具体来说,当处理 Tensor 1 时,我们先将内存中的数据拷贝过来进行计算,计算完成后,再将 Tensor 1 的结果返回。而双缓冲区的作用在于,我们可以有两个缓冲区,一个用于存放 Tensor 1,另一个用于存放 Tensor 2。这样,在计算 Tensor 1 的过程中,我们可以提前完成 Tensor 2 的数据拷贝,待 Cube 计算完成后,直接对 Tensor 2 进行计算。这种策略类似于双缓存,是一种优化手段。
在实际操作中,我们首先执行 Cube 1 的 Ping-Pong 计算,计算完成后进行 Softmax 操作,然后再计算 Cube 2。这里需要注意的是,Cube 和 Victor 是分开计算的。在之前的流程中,我们没有发现整体耗时增加的问题,原因在于 Softmax 的执行时间较长。
由于 Cube 和 Victor 是独立的计算单元,它们会相互掩盖。当 Softmax 阶段时间较长时,这种掩盖并不会导致问题,因为流水被掩盖掉了。然而,当我们把 Softmax 移到 Cube 之后,流水排布发生了两次变化:一是 Cube 2 中增加了 online Softmax 的步骤,导致 Cube 2 的执行时间变长;二是 SoftMax2 中减少了 online Softmax 的步骤,导致耗时变短。这种变化在流水上是可以观察到的。
发现问题后,我们进行了优化实践。通过分析,我们决定交换前面的双缓冲区,将 Softmax 1 的 Pong 和 Cube 2 的 Ping 的执行顺序进行调整。这种调整类似于优化前后的步骤。交换完成后,首先确保不影响结果的正确性。在 Cube 2 Ping 阶段,我们只依赖于 Softmax 1 Pong 的计算结果,不依赖 Softmax1 Pong 的计算结果,从而保证计算结构的正确性。
此外,这种调整的好处在于,它使 Cube2 Ping 的提前执行成为可能,从而更早地被 Softmax 1 Ping 所掩盖。由于 Softmax1 Pong 提前执行结束,依赖于 Softmax 2 Ping 的操作也能提前执行,进而减少 Victor 之间的间隔。调整后,我们发现总耗时下降了 5%,性能提升了大约 8%。
在特定场景下或遇到问题时,我们有一种思考方式。首先,要考虑不同计算单元或模块之间的掩盖问题。例如,在 GPU 上,Tensor Core 擅长矩阵计算,而 Victor 计算则有所不同。
在 CPU 上或 Scala 计算中,也存在类似的情况。当遇到瓶颈时,我们可以考虑将计算能力较弱的部分转移到计算能力较强的模块上。其次,性能优化需要有整体规划或整体视角,从计算、内存搬运和网络等多个维度去考虑问题。最后,性能流水中很重要的一点是流水掩盖。不仅昇腾有,英伟达也有各级流水。我们需要考虑如何进行流水掩盖,以提升性能。如果流水做好了,我们也就完成了流程优化的最重要的一部分。
基于 Ascend C 的通算融合算子性能优化
随着模型规模的指数级增长,单设备在计算能力、存储容量以及能效方面都面临着根本性的瓶颈。以拥有 1750 亿参数的 GPT-3 为例,至少需要数百 GB 的内存或显存资源才能满足需求。
因此,分布式推理与分布式训练成为了必然选择。在实际应用中,混合并行策略,包括模型并行、数据并行和流水线并行,已被广泛应用于各种框架中。我们将并行计算与通信算子相结合,形成了所谓的通算融合,旨在通过计算和通信的流水并行来提升性能。在分布式推理和训练框架中,MC²的性能优化是一个关键挑战,核心在于如何平衡计算与通信。
MC²通算融合算子的性能收益主要来源于通过合理切分 Matmul 计算,使得下一个数据块的 Matmul 计算与当前数据块的通信任务并行执行,从而隐藏通信时间。例如,对 Matmul 进行 m 轴切分,第二块数据的计算可以与第一块数据的通信并行进行,进而隐藏通信时间,提升算子性能。然而,这种优化方法也存在瓶颈。当计算和通信任务时间相差不大时,性能收益较为显著;但当两者时间差距较大,如 Matmul 计算时间远小于通信时间时,切分后的性能提升效果则会大打折扣。此外,数据切分可能导致计算或通信执行时间膨胀。如果切片数据量过小,可能会导致计算单元未对齐、计算量不足或通信效率降低等问题。同时,切分速度过快可能引入额外的调度开销,而并行化后计算和通信时间对 L2 缓存的访问冲突也可能影响性能。
在实际优化过程中,我以 MatmulAllreduce 算子为例进行了深入探索。首先,通过分析原始矩阵的计算和通信任务,判定其 bound 场景。当 k 轴较大时,计算量大,执行时间大于通信时间,此时算子为计算 bound;而当 k 轴较小时,计算量小,执行时间小于通信时间,算子则为通信 bound。在设计数据切分策略时,我选择按行切分 m 轴,以保证每行数据的内存连续性,满足通信算法的要求。
在切分过程中,我将数据块分为长块和短块,根据计算和通信时间的相对大小,合理安排数据块的顺序。在计算 bound 场景下,将长块放前面,短块放后面;而在通信 bound 场景下,则将短块放前面,长块放后面,以实现最佳的流水掩盖效果。
前面前置分析之后,那我们接着下来,我们要做一些这个方案上面的一个验证,这里还是以这个 MatmulAllreducereduce 这个为例,比如说我们这个输入指针是为 M=4096,K=3072,N=8192,数据类型为 half,分核数为 8,通过融合前 msProf 工具采集,得到该输入的 Matmul 计算执行时间为 803us,AllReduce 通信执行时间为 1071us,总耗时 1874us,属于通信 bound 场景。按照上述的切分算法,本案例的具体切分情况如下:
根据经验值选定断块(bound 场景中断块为头块,即切分的第一个数据块)M 方向长度 m0 为 384,则通信数据量 x 为:384*8192*2 / 1024 / 1024 = 6MB。按通信拟合公式估算通信的执行时间为 143us。考虑可能会发生内存带宽冲突,因此再乘以 1.15 的系数,得出通信执行时间 164us。
短块经验值公式,a、b、c 表示根据经验值给出的短块长度的备选。
a * K * N >= 4 * 1024 * 1024 * 1024,a取不等式的最小值;
b * K * N / 1024 + b * N >= 6 * 1024 * 1024,b取不等式的最小值;
c >= 3 * 128,c取不等式的最小值
m0 = min(a, b, c)
计算 M 方向长块长度 m1。根据短块通信时间,配平长块的计算执 bound 行时间同样为 164us,即将 t1 作为长块的计算执行时间, 带入 t1 = CostMM(m1) 公式,按计算拟合公式估算出该长度 m1 为 768。
根据 M=4096、m0=384、m1=768,计算长块个数:(4096 - 384) / 768 = 4.83,向下取整为 4。
根据短块长度 m0=384,长块个数 4,调整长块 m1 长度:(4096 - 384) / 4 = 928,向下按 128 对齐,调整 m1 为 896。 5. 根据长块长度 m1=896,长块个数 4,调整短块 m0 长度:4096 - 896 * 4 = 512。
最终得到将原始输入矩阵切分为 5 个数据块,长度分别为:{512,896,896,896,896}。
在算子的 Tiling 代码中设置制定好的切分策略。按该 切分策略测试,融合后该算子的执行时间为 1262us,则融 合算子的性能收益为 (1874 - 1262) / 1874 = 32.7%。
总之,在推理过程中,我们确实需要考虑计算与通信并行的问题。随着硬件设备的不断发展,计算单元数量不断增加,计算与通信并行的需求也日益迫切。
目前,虽然我们在通算融合方面已经做了一些工作,但仍有进一步优化的空间。在实际操作中,切分策略是一个关键因素。切分策略的选择与算子的实现密切相关。有些算子可能有专门的实现方式,这就需要我们根据具体情况来确定合适的切分方式。
此外,我们还需要动态调整切分策略。在实际应用中,可能会遇到一些性能膨胀的问题。虽然我们有拟合公式来指导切分策略,但这些公式并非完全准确,可能会存在一些偏差。因此,我们需要根据具体场景对切分策略进行校正或修正,以确保其适应性和有效性。
总结与展望
大模型推理加速是一项复杂的系统工程,它涵盖了算子、算法、框架、资源调度以及底层芯片等全栈综合能力,需要我们在不断的探索与实践中逐步推进。
这项技术的核心目标在于提升硬件资源的利用率,减少不必要的计算量,降低通信开销,尤其是在分布式环境中,最终实现端到端性能的提升以及推理成本的有效降低。
展望未来,专用硬件加速器的发展将遵循软硬协同设计的方法。例如,存算一体芯片和计算型存储盘等创新技术将使存储单元更加接近计算单元,从而减少数据传输的延迟和能耗。此外,针对大语言模型算法数据流的优化芯片架构也将成为研究的重点,以更好地适应大模型的计算需求。
在长上下文或序列场景的优化方面,随着应用场景的日益复杂,处理更长的上下文或序列的需求也在不断增长。对于服务长序列负载的大语言模型来说,需要在算法和系统两个层面解决诸多挑战。
在算法层面,目前主要面临长度泛化失效的问题,即模型在处理超出训练时序列长度的输入时性能下降。为了解决这一问题,目前的策略包括召回增强、序列压缩和缓存等技术,目的是尽可能缩短序列长度,同时保留关键信息,以提高模型在长序列上的性能表现。
嘉宾介绍
张君,作为核心开发者参与 AI 框架 (昇思) 的开发,并负责动态图的自动微分以及动静结合模块。目前主要参与大模型推理在昇腾硬件上的相关开发和优化工作,致力于通过优化推理框架、模型算法和算子加速库等层面,进一步提升大模型推理的性能。
会议推荐
首届 AICon 全球人工智能开发与应用大会(深圳站)将于 8 月 22-23 日正式举行!本次大会以 “探索 AI 应用边界” 为主题,聚焦 Agent、多模态、AI 产品设计等热门方向,围绕企业如何通过大模型降低成本、提升经营效率的实际应用案例,邀请来自头部企业、大厂以及明星创业公司的专家,带来一线的大模型实践经验和前沿洞察。一起探索 AI 应用的更多可能,发掘 AI 驱动业务增长的新路径!