大模型推理 & memory bandwidth bound (1) - 性能瓶颈与优化概述_大模型计算中常见的带宽瓶颈-CSDN博客

本文链接：https://blog.csdn.net/daihaoguang/article/details/143671796

系列文章目录

大模型推理 & memory bandwidth bound (1) - 性能瓶颈与优化概述
 大模型推理 & memory bandwidth bound (2) - Multi-Query Attention
大模型推理 & memory bandwidth bound (3) - MLA

文章目录

系列文章目录
前言
一、重要概念
二、Roofline Model
三、性能瓶颈及优化策略
四、GPU简述
- 1.GPU vs CPU
- 2.内存结构
总结
参考文献

前言

随着大模型参数量的增加，其推理加速成为一个重要的研究方向。比如我们在vLLM系列中有讲到，vLLM在内存（显存）管理上使用了PagedAttention技术，再结合continuous batching的调度策略，大大提高了在面对多个请求时GPU的使用率，这种系统层级上的优化提高了吞吐，降低了延迟。
大模型推理之所以低效，主要是因为自回归解码过程受到内存带宽限制（memory-bandwidth-bound）。在解码阶段，模型需要频繁从内存中读取和写入数据，而内存带宽又比较有限，因此有较高的延迟。为了突破该性能瓶颈，出现了各种有意思的工作，比如Multi-Query attention、Flash Attention系列、Speculative Decoding以及MEDUSA等等。本系列会对这些知名工作做一些解读，加深对大模型推理加速的理解。
作为开篇，本文主要围绕性能瓶颈等基础问题展开，包括但不限于memory-bandwidth-bound。

一、重要概念

出于后文对大模型推理性能瓶颈分析的需要，在此处先讲一些相关的重要概念。

FLOPS：同FLOP/s，全称为Floating Point Operations Per Second，表示每秒钟执行的浮点数操作次数，是衡量硬件性能的指标；
FLOPs：Floating Point Operations，表示浮点数运算次数，用于衡量模型或算法的复杂度；
MOPs：Memory Operations，表示用于数据读写的内存操作数（文献[1]）；
Arithmetic Intensity：算数强度（文献[1][2]），表示每单位数据处理（内存操作数）所执行的浮点运算数量，用于衡量计算任务的计算密集程度，可以用如下公式表达。为了充分发挥GPU的性能，需要确保计算任务具有较高的Arithmetic Intensity。

$\text{Arithmetic Intensity} = \frac{\text{FLOPs}}{\text{MOPs}}$

Operational Intensity：计算强度/操作强度，每秒浮点运算次数与DRAM实际数据传输速率之比，可用如下公式表示。与Arithmetic Intensity的差异如下：Arithmetic Intensity关注的是处理器processor和缓存cache之间的流量；而Operational Intensity关注的是缓存cache和主内存DRAM之间的流量。
$\text {Operational Intensity} = \frac {\text {FLOPS} \ (FLOP/s)} {\text {Effective DRAM bandwidth}\ (Bytes/s)}$
Operational Intensity vs Arithmetic Intensity：Operational Intensity这个概念是在Roofline Model（文献[3]）中提出的，但有意思的是后面的论文（比如文献[4]及其引用文献）在讨论Roofline Model时都在使用Arithmetic Intensity而非Operational Intensity。

二、Roofline Model

Samuel Williams等人在论文Roofline: An Insightful Visual Performance Model for Floating-Point Programs and Multicore Architectures中提出了一种直观的性能模型，帮助开发人员识别和解决性能瓶颈，优化程序性能，该模型就是Roofline Model。
下图是论文中给出的AMD Opteron X2的Roofline Model。其中横轴代表Operational Intensity，纵轴表示可获得的最大性能（每秒最大浮点运算次数）Attainable Performance。整张图分为左侧黑色斜线（屋檐）区以及右侧黑色水平线（屋顶）区两个部分，让我们来解释一下：
在这里插入图片描述

屋檐区：Operational Intensity较小，增加memory bandwidth即可获得计算性能的收益，所以此处打满，黑色斜线的斜率就是最大内存带宽peak memory bandwidth；由于可获得的最大计算性能受制于memory bandwidth，因此该区域为memory bound，更严谨的说，应该是memory bandwidth bound；
屋顶区：硬件存在一个最大浮点计算性能peak floating-point performance，如黑色水平线所示，当Operational Intensity增大至一定限度，即便增大使用peak memory bandwidth也不能获得计算性能收益；此区受制于硬件计算性能，因此被称为compute bound。
现代硬件的计算性能增速是远超内存带宽的，加之大模型自回归解码需要频繁访问内存，可以想见我们面对的问题主要集中在memory-bandwidth-bound。
题外话：1）Roofline Model与深度学习模型的性能分析 - 知乎这篇讲的挺好，如果还没有理解可以看一下；2）上图只是一个简化的理想的模型图，前提条件是已经使用了多种性能优化的手段，比如ILP和SIMD等等，感兴趣的推荐阅读原论文。

三、性能瓶颈及优化策略

模型性能瓶颈主要有4种类型，分别是计算受限（compute-bound），内存带宽受限（memory-bandwidth-bound）、通信受限（communications-bound）以及开销受限（overhead-bound）。现在让我们来对这些性能瓶颈做一个说明。尽管本系列主要探究的是memory-bandwidth-bound，但还是希望能完整的介绍这4种性能瓶颈，以及在模型推理场景下相应的优化策略。
这一部分主要参考了Dissecting model performance、Making Deep Learning Go Brrrr From First Principles以及MEDUSA。

1.memory-bandwidth-bound

memory-bandwidth-bound就是前面提到的计算性能受到内存带宽制约的情况。需要说明的是，在Roofline Model论文中使用的是memory bound，PagedAttention论文中也采用memory-bound的说法，但是很多大模型推理优化的论文中都是用memory-bandwidth-bound一词；个人觉得后者表达更准确，推荐采用memory-bandwidth-bound的说法。
memory-bandwidth-bound的典型操作包括：elementwise操作，比如activation, dropout，mask； reduction操作，比如sum, softmax和batch norm等归一化操作。这些操作计算快，内存读写慢。
针对大模型推理场景，我们应该怎么解决这个问题？从Roofline Model来看，我们需要在更少的内存操作下进行更多的浮点运算，因此存在如下这些优化策略：

使用更大的batch_size批量推理，增加算数强度，充分发挥GPU的并行计算能力。为了使用更大的batch_size，可以选择如下方式：
a. 减少KV Cache：比如Multi-query attention和Grouped-query attention减少了KV Cache对内存的占用；PagedAttention则是在系统层级减少了KV cache的内存碎片；
b. 量化：使用量化等模型压缩技术能减少内存消耗，增大batch_size，包括权重量化，以及最近流行的KV Cache量化；
c. continuous batching也是一种充分利用批量推理的优化方式；
投机采样：投机采样借助小模型（draft model），每个step能并行推理多个tokens，从而加速了推理；
减少内存操作次数，可选方案有：
a. 算法优化：比如Online Softmax优化了Softmax的计算，减少了内存访问次数；
b. kernel fusion (operator fusion / layer fusion)：如下图所示，经过内核融合（算子融合/层融合），一次性计算多个算子的计算，避免中间结果在HBM (global memory)到SRAM的反复读写带来的额外开销；

2.compute-bound

compute-bound对应Roofline Model的屋顶区。compute-bound的典型操作包括：大矩阵相乘和多通道卷积，这些操作计算慢，内存读写快。
针对大模型推理场景，有如下优化策略：

减少模型推理的浮点运算数量：比如模型剪枝和知识蒸馏；
使用更低精度，运算速度更快的数据类型，如下图所示，在NVIDIA H100中，8-bit Tensor Cores>16-bit Tensor Core>32-bit Tensor Core，而且都是约两倍的计算性能提升。

3.communications-bound

顾名思义，communications-bound是指数据传输和网络通信受限的情况，适用于计算和数据分布在多个芯片上的情况。优化策略为，在分布式系统中合理的数据划分，减少通信。

4.overhead-bound

overhead-bound是软件导致的瓶颈，这种情况花费了更多的时间在做一些与张量传输和计算无关的事情，比如Python Interpreter，Pytorch框架需要即时完成的操作，以及CUDA kernel launch（内核启动）。此时GPU由于CPU开销处于等待（闲置）状态。优化策略包括：

使用Graph Capture技术，将代码块中的所有GPU操作（如内核启动）记录下来，并将它们整理成一个有向图，以便将其作为一个整体一次性提交给GPU执行。vLLM中就集成了该技术，以提升Decode阶段的推理效率；
还有就是优化推理系统的任务提交，比如前不久（2024年10月）SGLang在LMSYS Online Meetup上提到的重要的优化方向：CPU overhead hiding。在大模型推理框架中，CPU负责调度工作（图中的Scheduler），而GPU负责模型推理（图中的Model Worker）。优化后的Overlapped Scheduler相较于Blocking Scheduler隐藏了CPU开销时间，解决了overhead-bound问题。

四、GPU简述

1.GPU vs CPU

在这里插入图片描述
上图是CPU和GPU的架构对比图，其中绿色部分代表计算单元或者计算核心，橙色部分表示内存（包含缓存），黄色部分表示控制单元。可以看到，尽管CPU的单核能力很强，但是GPU的核心足够多。GPU编程的并行处理模型是基于SIMT(Single Instruction, Multiple Threads)的，这使得它能够高效地处理大量并行线程，特别适合图形处理和计算密集型任务。

2.内存结构

和CPU类似，GPU也分为主内存DRAM以及缓存cache。
1）GPU的DRAM为HBM (High Bandwidth Memory)或GDDR (Graphics Double Data Rate)，比如A100使用的是HBM2，如下图所示（文献[11]）。它通常被称为Global Memory，也即是我们常说的显存。
2）cache呈现层级结构，分别为L0/L1/L2，他们的容量大小依次递增，但是带宽（即速度）依次增加；且他们的带宽大于HBM。
在这里插入图片描述
在Flash Attention论文中，给出了如下的内存层次结构：
1）SRAM速度最快，速度高达 19TB/s ，但是存储空间仅有 20MB ，对比上图你就发现它指的是L1 cache（关于哪部分属于SRAM在图解大模型计算加速系列：FlashAttention V1，从硬件到计算逻辑 - 知乎评论区有讨论）；HBM速度次之，为 1.5TB/s ，但是存储空间为 40GB ；还有CPU DRAM部分，存储空间最大，但带宽最小；
3）在使用GPU进行模型推理时，会将HBM上的数据加载至SRAM进行计算，完成后再将计算结果写回HBM；我们把SRAM称为on-chip memory，而把HBM称作off-chip memory；
4）回到本系列主题memory bandwidth bound，大模型推理（特别是Decode阶段），主要的性能瓶颈就来自于HBM有限的内存带宽（数据传输速度太慢）；
5）CPU和GPU之间通过PCIe或者NVLink互联；通常是在显存不够的时候，会将一部分数据暂存至CPU DRAM，比如vLLM中处于SWAPPED状态的序列，其KV Cache就是下放到CPU，当该序列再一次被调度执行解码时，这部分KV Cache会被重新加载到GPU HBM上（参见vLLM (3) - Sequence & SequenceGroup）。
在这里插入图片描述
本系列所需的GPU（内存结构）知识到此已经足够了，关于GPU更深入的知识可以参考NVIDIA H100白皮书，以及其他资料，我在【参考文献】部分也会放一些。