LLM时代中的AI推理优化
问题与挑战 毫无疑问，AI是当下最热的话题之一，而大模型又是当前AI的主角。几年前，正当深度学习进入瓶颈时，以GP 2024-6-15 03:5:14 Author: blog.csdn.net(查看原文) 阅读量:0 收藏

问题与挑战

毫无疑问，AI是当下最热的话题之一，而大模型又是当前AI的主角。几年前，正当深度学习进入瓶颈时，以GPT为首的LLM的横空出世让之似乎又找到了“第二增长曲线”。当模型规模大到一定程度时，它所表现出来的涌现能力（Emergent ability）是之前在小模型中所不曾见过的。这种大模型所特有的推理、计算等能力给我们带来了无穷的想象空间。

但是，它的代价是模型和以往模型相比增大了成百上千倍。要玩大模型十亿参数基本是个入门级门槛，上百亿才算像点样。各个大公司为了争夺大模型的话语权，更是将大模型越“卷”越大。模型规模增大所带来的挑战是多方面的，这里我们主要关注它在计算，尤其是推理时带来的挑战。提高LLM推理性能，降低LLM推理成本是加速其相关应用产品化的必经之路。

具体来说，LLM推理计算的主要挑战来自几个方面：

  • 内存用量大：模型参数多，导致模型本身及中间张量的内存使用量增大。就算以FP16存储，模型的大小基本上也是参数量的两倍。一个模型本身动辄占用几十、几百个G，就已经数倍于主流GPU的内存了。再加上推理时其它一些中间张量，及KV Cache等，所需内存往往需要参数本身的几倍。因此，内存就成为了非常严峻的问题，甚至决定了能否跑得起来的关键因素。

  • 硬件利用率低：目前主流的LLM推理采用auto-regressive方式。它是一个迭代的过程，每次迭代会产生一个token。其工作范式分为两个阶段：第一个阶段称为Prefill或prompt Processing。它根据给定的prompt计算KV Cache，并产生第一个输出token。第二个阶段称为decoding或token generation。它迭代式地逐个生成token，同时更新KV Cache用于下一次迭代。两者的计算模式有很大不同，前者通常为compute-bound；后者由于arithmetic density较低，通常为memory-bound。它难以并行，导致硬件利用率低。

优化技术

业界提出了很多LLM的推理优化技术。本文无法全部覆盖，主要涉及以下几个方面：

  • Attention Optimization
  • Approximate Attention
  • KV Cache
  • Decoding
  • Batching
  • Distributed
  • Offloading
  • Compression
  • Others

Attention优化

目前主流LLM模型几乎都是基于Transformer架构。该网络架构除了头尾，其它都是Transformer layer（或称为Transformer block）的重复堆叠。Transformer block包含三部分：dense layer projection，self-attention, feed-forward。它们之中，除了self-attention其它其实都是GEMM。GEMM的优化技术被研究了几十年，相对成熟，不用多说。而attention的结构相对更复杂，更难并行，而且naive的实现中内存占用与context长度是平方关系。这使它成为耗时与内存占用的瓶颈，因此是优化的重点。

其中一个优化角度是Sharing-based Attention优化，即修改Attention结构使多个头间共享部分数据。如论文《Fast Transformer Decoding: One Write-Head is All You Need》中提出的MQA（Multi-Query Attention）与论文《GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints》中提出的GQA（Group-Query Attention）。它们被用于在PaLM, StarCoder, LLaMA2等模型中。但由于涉及网络结构本身的改变，这里不作展开。

FlashAttention & FlashAttention2

FlashAttention系列可能是这几年中关于Attention最重要的优化。论文《FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness》与《FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning》分别提出了FlashAttention（FA）和FlashAttention 2（FA2）。由于出色的性能表现，目前已成为业界的标配。

记Attention的输入为 Q , K , V ∈ R N × d \mathbf{Q,K,V} \in \mathbb{R}^{N \times d} Q,K,V∈RN×d，其计算过程的数学描述为：
S = Q K T ∈ R N × N , P = softmax ( S ) ∈ R N × N , O = P V ∈ R N × d \mathbf{S = QK}^T \in \mathbb{R}^{N \times N}, \quad \mathbf{P} = \text{softmax}(\mathbf{S}) \in \mathbb{R}^{N \times N}, \quad \mathbf{O = PV} \in \mathbb{R}^{N \times d} S=QKT∈RN×N,P=softmax(S)∈RN×N,O=PV∈RN×d
其中矩阵相乘 Q K T \mathbf{QK}^T QKT产生shape为[batch, head num, seq len, seq len]的中间张量， P \mathbf{P} P也一样。而最终需要的 O \mathbf{O} O是shape为[batch, head num, seq len, hidden size of one head]的张量。由于seq len一般会很长（现在的模型中支持的seq len越来越长，如GPT-3为2K, LLaMA-2为4K到32K，GPT-4有32k，CodeLlama有100k，Gemini 1.5有1M），因此前者的中间张量会非常大。这也意味着如果按照上面的数学定义来计算，中间张量大概率会把内存撑爆。FlashAttention使用tiling解决了该问题，同时减少了HBM与SRAM的读写。谈到tiling，很多文献是针对GEMM或者Conv这样的算子，但要应用于softmax咋一看会有些棘手，因为它本身包含了reduction的语义。对于这个问题，FlashAttention使用论文《Online normalizer calculation for softmax》中的online softmax技术，从而避免了构建巨大的attention matrix，同时也使得softmax可以被并行（原文3.1节）。该算法使内存的复杂度降低到线性，极大降低了它的memory footprint。另外，它还通过kernel fusion，将attention所有操作融合到一个CUDA kernel中，从而减少了kernel launch与访存的开销。

FlashAttention可达理论FLOPS的25-40%，比标准的attention快2-4倍，节省10-20倍的内存。但是它在thread block与warp的任务切分上还不是最优的，会导致occupancy低下，以及不必要的shared memory读写。另外，近年来context长度越来越长，FA在这方面没有专门的优化。在这样的背景下，FlashAttention 2对它进行了优化，改进了计算任务的划分，提升了并行性，使之能达到理论FLOPS的50~73%，性能提升2倍。它采用的主要技术包括：

  • 通过online softmax与casual mask计算的调整，减少了非GEMM运算（由于GEMM计算有Tensor Cores加持，相比而言厚厚GEMM计算虽然FLOPs占比不大，但耗时占比大）
  • 除了FA中对于batch与head并行外，还在seq维度上并行，从而提升在长seq length时的性能。
  • FA中的split-K scheme需要所有的warps都将中间结果写入shared memory，同步后累加。FA2中变换了切分策略，使得warps间不需要通信，从而减少shared memory读写。

论文《A Case Study in CUDA Kernel Fusion: Implementing FlashAttention-2 on NVIDIA Hopper Architecture using the CUTLASS Library》针对NVIDIA Hopper架构改进了Flash-Attention-2实现。它基于CUTLASS实现，利用了Hopper架构新引入的特性-Tensor Memory Accelerator（TMA）和Warpgroup Matrix-Multiply-Accumulate（WGMMA）指令。通过将TMA copy与WGMMA指令进行overlap，同时选取最优的tile size以平衡寄存器压力与shred memory的使用，相比Ampere架构上的FA 2实现有20-50%的性能提升。

FlashDecoding & FlashDecoding++

文章"Flash-Decoding for long-context inference"和论文《FlashDecoding++: Faster Large Language Model Inference on GPUs》提出了Flash-Decoding（FD）与FlashDecoding++（FD++），也称为FlashAttention-V3和V4。前面的FA是训练与推理通用的，而FD，就像它的名字一样，是针对推理场景中的decoding的。

Attention的性能瓶颈在于对中间张量 Q K T QK^T QKT的读和写。前面的FA在batch size与query length维度上做并行。训练阶段由于batch大，容易并行。但推理阶段query length一般是1，同时推理时batch一般还比较小。这意味着如果batch数小于SM数量，就打不满GPU，硬件资源会被浪费。这在长序列时尤为严重，因为序列长一般意味着batch较小。因此，FA的优化不太适用于推理。

FD在FA的基础上，对于长序列推理有8x的提速。它的主要思想将并行化维度扩展到key和value，使得它们可以并行地加载和计算。其代价是最后需要做reduction。其工作流程有三步：1) 将keys/values切成更小的chunk。这一步无需GPU操作。2) 使用FlashAttention并行地为每个split计算query的attention。3) 最后，对所有split的结果做reduction。

FlashDecoding++通过进一步实现了Attention计算的并行，并针对扁平矩阵的乘法进行了优化。它针对LLM推理中存在的三个问题，在FD的基础上做了如下改进：

挑战方案
Synchronized partial softmax updateAsynchronized softmax with unified max value
Computation resources is under-utilized for the flat GEMM operationFlat GEMM optimization with double buffering
The performance of LLM inference suffers from the static dataflowHeuristic dataflow with hardware resource adaption

基于这些优化，FlashDecoding++与FlashDeocding相比提速1.37x。

  • Asynchronous Softmax
    为了数值稳定性，避免溢出，在计算softmax一般会先将数值减去最大值。以往FA中采用partial softmax的方式，将数据切块，分别计算块内softmax。虽然达到了一定程度的并行，但由于依赖于其它的partial结果，因此最后仍需要以同步的方式更新。这种synchronized partial softmax update的方式占到Attention计算的相当比重（Llama2-7B中18.8%）。文中指出LLM中的softmax的输入值域99.99%都在特定区间中，因此可以使用一个统一的值来取代这个最大值，从而避免同步，让这部分计算完全并行起来。如果发生溢出，再用同步的方法进行重计算（recomputation）。

  • Flat GEMM
    Decode阶段batch size与seq len一般较小，因此GEMM操作中的矩阵比较扁平（即flat GEMM）。当batch size为1时甚至退化成GEMV。另外当prefill阶段的seq len较小时也会产生flat GEMM。这些情况下会导致硬件利用率很低。LLM推理中广泛使用cuBLAS或CUTLASS这样的库利用Tensor Core进行加速。尽管Tensor Core支持GEMM中的M=8，但库中一般会在M维度上将tile size设为64以hide memory latency。但是，decode阶段中GEMM的M常常远小于64，这样会需要padding从而导致浪费。 FD++将矩阵padding到8（而非像之前设计中的64）的倍数以提升利用率。文中指出不同shape的flat GEMM会面临不同的瓶颈（当N较小时是parallelism-bound，N较大时是memory-bound）。FD++在N较大时通过引入double buffering来hide memory latency。即在shared memory中分配两块buffer：一个buffer中的tile执行GEMM操作，另一个块buffer为下一次GEMM操作加载tile。这样，计算与访存就可以overlap。

  • Heuristic Dataflow
    LLM推理中的GEMM是多样的，可能是GEMV（decode阶段batch size=1的情况），可能是flat GEMM（decode阶段小batch size情况，或是prefill阶段小seq len情况），也可能是传统的GEMM（prefill阶段当batch size与seq len都较大时）。当前的框架（如FasterTransformer，DeeSpeed）使用cuBLAS中的高性能GEMM实现处理不同的workload。这种统一的处理方式并不能保证达到最好的性能。影响linear workload性能的因素有很多，有input dynamics（batch size, seq len），Model多样性（结构和size），GPU能力（memory bandwidth, cache size等），工程effort等。对不同的输入与硬件采用静态的方案会让不同shape下GEMM性能有50.25%的性能损失。因此，FD++考虑输入的多样性，用启发式的方式使用不同的硬件资源（Tensor Core或CUDA core）进行优化。经过分析可知，对于一个特定LLM模型，其GEMV/GEMM操作的[K,N]组合只有4种。对于这几种情况，考虑三种实现：FastGEMV，自定义的GEMM实现和CUTLASS的传统GEMM实现。M较小时考虑前两种，M较大时考虑后两种。具体用哪种通过profile的方式确定。

Blockwise Parallel Transformer

论文《Blockwise Parallel Transformer for Large Context Models》提出了Blockwise Parallel Transformer（BPT）方法对self-attention与FF模块都用blockwise的方式计算，并将它们融合。这样可以节省构建中间张量所需的内存。对于GPT2模型，能训练的seq长度比标准的attention长32倍，比FlashAttention长4倍。

FlashInfer

由于attention在LLM中独特的重要性，业界也涌现出一些优秀的计算库，如FlashInfer。它是一个LLM推理的高性能GPU kernel（如FlashAttention, PageAttention和LoRA）实现库。它包含single-request和batching的prefill, decode和append kernel。KV-Cache支持Padded tensor, Ragged Tensor, Page Table。另外，它还支持shared-prefix batch decoding（详见"Cascade Inference: Memory Bandwidth Efficient Shared Prefix Batch Decoding"），优化GQA, Fused-RoPE Attention和Quantized Attention等特性。

Approximate Attention

为了降低计算复杂度，有一类方法研究用稀疏矩阵或者低轶矩阵近似attention的中间矩阵。利用attention矩阵的稀疏性提升性能的方法如：

Sparse Transformer

论文《Generating Long Sequences with Sparse Transformers》基于很多层有sparse attention pattern的观察，提出Factorized Self-Attention，将full self-attention操作分解成多维护更快的操作。每个attention head，会使用pattern选取一些位置的子集。对于图像和音乐等具备周期性的数据，可以使用strided attention pattern；而对于其它如文本数据，可以使用fixed attention pattern。该方法的复杂度为 O ( n n ) O(n \sqrt{n}) O(nn ​)。实现上，它还使用了sparse attention kernels用于计算attention matrix的子集。

Reformer

论文《Reformer: The Efficient Transformer》采用了以下技术：1) 使用locality-sensitive hashing（LSH）替换dot-product attention。使其复杂度从 O ( L ) O(L) O(L)降到 O ( L log ⁡ L ) O(L \log L) O(LlogL)。LSH是一种在高维空间中寻找nearest neighbor的方法。这里用来在attention中找序列中相关的context。2) 它还使用了reversible residual layer来替代原来的residual，它使得训练中的activation只需存一次而不是N次（N为layer num）。3) 将FF层中的activation切分并以chunk为单位进行处理，从而减少FF层中的memory。

Routing Transformer

论文《Efficient Content-Based Sparse Attention with Routing Transformers》提出学习dynamic sparse attention patterns，避免为与query无关的内容浪费memory与计算。它基于local attention与content-based sparse两方面的思想。Routing Transformer在self-attention上加了基于online k-means的sparse routing module。每一步的query都被route到部分cluster assignment所指定的context元素。它的启发来源于spherical k-means clustering在Maximum Inner Product Search（MIPS）问题上的应用。该方法避免了实例化dense attention matrix，将复杂度从 O ( n 2 d ) O(n^2d) O(n2d)降低到了 O ( n 1.5 d ) O(n^{1.5}d) O(n1.5d)。

HyperAttention

以往的Approximate Attention些方法普遍不支持causal masking。以往工作表明Attention在最坏情况下，需要quadratic time，除非attention matrix的元素有界或者矩阵rank较小。论文《HyperAttention: Long-context Attention in Near-Linear Time》中提出的HyperAttention是一种approximate attention机制，用来解决LLM中长上下文带来的计算挑战。该方法采用Locality Sensitive Hashing（Hamming sorted LSH）与Approximate Matrix Multiplication（AMM）对Attention的计算 A t t n = D − 1 A V \mathbf{Attn = D^{-1}AV} Attn=D−1AV进行近似。它可以达到接近线性复杂度，且支持casual mask。

KV Cache

Auto-regressive的token生成方式中，每一个token理论上需要基于前面的token。Naive的LLM推理会重复计算前面token的Keys与values（上面公式中K和V两个张量）。为了避免了这种重复计算，我们可以将前面token计算好的K和V缓存起来，该方法称为KV Cache。它也是目前LLM推理的主要优化方法之一。这本质上是一种“空间换时间”的优化，因此其代价是会占用很多内存。在实际场景中所占内存相当大，可达内存总用量的30%左右。因此，业界提出了不少关于KV Cache优化的工作。它们可大致分为几个方向：

Granularity

Naive的KV Cache以最长sequence为粒度分配和管理，会带来巨大的浪费。以更小的粒度管理通常可以节省大量的内存。但过小的粒度又会因为物理不连续带来性能负面影响。因此，我们需要根据实际需求在两者之前平衡。

PagedAttention

论文《Efficient Memory Management for Large Language Model Serving with PagedAttention》中提出的PagedAttention使用类似于操作系统中的虚拟内存机制对KV Cache进行管理。KV Cache的size会随着token产生的过程动态变化。一种简单的方式是提前按最大可能长度预分配连续的内存。但是，这会导致三种内存浪费：预留空间，内部碎片化，外部碎片化。其中有效利用的内存可能只有20%左右。另一方面，这种方式难以实现KV Cache的内存共享（Memory sharing）。在一些场景中，我们需要对于一个request生成多个序列，或者在做beam search时，其prompt或者不同beam candidate的序列的KV cache是可以共享的。

PagedAttention允许在不连续的内存空间中存储连续的K和V。它将KV Cache分成逻辑块（Logical block），每个逻辑块包含固定token数量的K和V。这些逻辑块，只有当真正用到时才分配物理块。这种方式使浪费率低于4%。这样做的另一好处是可以实现内存的共享，即将不同序列的逻辑块映射到同一物理块。它引入了reference count和copy-on-write机制来进行对这些内存块进行管理。

vLLM是基于PagedAttention的LLM推理系统。vLLM采用了centralized scheduler来协调分布式GPU worker的执行。KV Cache manager通过来自centralized scheduler的指令管理物理的KV Cache 内存。在调度方面，它采用FCFS调度策略，保证公平防止“饥饿”，同时利用fine-grained batching mechanism（即以iteration为粒度调度）实现了preemptive request scheduling。在cache的eviction策略上，它根据LLM的特点（一个序列的所有block同时使用）实现了all-or-nothing eviction policy。同时支持swapping和re-computation两种recovery机制。前者适用于block size大时，后者适用于block size小时。在kernel级别的优化上，它还实现了Fused reshape and block write，block read and attention, fused block copy等GPU kernel。与FasterTransformer与Orca相比，它在latency不变的情况下将throughput提升2-4x。

对于request间有相同的prefix的情况，将KV Cache缓存下来可以避免重复计算。这就是Prefix caching的基本思想。比如论文《Improving Large Language Model Throughput with Efficient Long-Term Memory Management》基于PagedAttention提出prefix cache，该prefix cache可被swap到CPU与disk。

TokenAttention

PagedAttention以block（一个block可以存放多个token的KV Cache）为单位对KV Cache进行管理，仍然会有内存浪费。基于Python的LLM推理与服务框架LightLLM中采用了"TokenAttention"以更小的粒度，即token为粒度对KV Cache进行管理。它能最小化碎片化，并使得内存的分配与释放更加高效。其工作流程可分几步：1) 模型初始化时，基于用户设定的最大token数量分配KV Cache，同时创建token table用于记录物理内存。2) 当新request到来时，先检查在预分配的内存中是否有可用的连续空间。如果没有，它会分配非连续的GPU memory，并记录在token table中。3) 对于新产生的token，从预分配的token cache中找到空闲的空间，并加到token table中。4) 当一个request完成后，可以将token table中对应的记录删除。从而可以让内存被新的request所使用。

RaggedAttention

前面的PagedAttention与TokenAttention都是以减少内存浪费为目标，但以性能为代价。因为内存不物理连续时，其性能也受到一定程度的影响。在Continuous batching技术（后面会提到）中，同时处理的序列的长度可以是不同的，这样就形成了Ragged tensor。"PAI BladeLLM 推理引擎: 超长上下文、更高性能"中的RaggedAttention通过保证同一序列的KV是连续存储的，这样可以提升kernel的访存效率，从而提升性能。也就是说，它没有PagedAttention和TokenAttention那么极端，但其代价是会有一些内存浪费。

Platform Feature

平台特性可以被用来优化KV Cache。

vAttention

论文《vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention》指出PagedAttention为KV Cache使用了非连续的virtual memory，导致软件复杂、移植性、冗余和低效问题。vAttention为KV Cache使用连续的virtual memory，使用low-level系统中已有的demand paging支持（CUDA 10.2中引入的cuMemCreate, cuMemMap, cuMemAddressReserve等API用于分离virtual与physical memory的分配，详细请见https://developer.nvidia.com/blog/introducing-low-level-gpu-virtual-memory-management/）。另外，它引入两个优化：一方面，修改NVIDIA driver（开源部分）增加page size选择，以减少内部碎片，和hide内存分配延迟。另一方面，它使用三种技术hide内存分配的开销：1) 在background thread中分配新页，使其与计算可以overlap。2) Deferred reclamation使batch中的request的KV cache可以重用。3) Eager allocation在physical page在需要前分配好，比如新的prefill request到来前。

Shared Prefix

在prompt工程中，我们会为用户的prompt前加一些system prompt。它包含了任务描述、上下文信息、样例等信息。System prompt在多个request间是相同的，并且可能会很长。如果对于每个request，都去重新计算就会带来巨大的浪费，导致Time to first token (TTFT)非常长。

RadixAttention

"Fast and Expressive LLM Inference with RadixAttention and SGLang"一文指出现有支持KV Cache共享的系统需要手工配置和调整，无法自动化。它提出的RadixAttention是一种运行时的自动高效的跨多LLM generation调用的KV Cache reuse机制。它在当request的处理完成后，不会丢弃其KV Cache，而是将token序列到KV Cache的映射用radix tree管理起来。这样prefix search，insertion和eviction操作会非常高效。另外它还实现了LRU eviction policy和cache-aware scheduling policy来提升cache hit rate。

ChunkAttention

论文《ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition》也是用于多request的prompt有共同prefix的场景（如system prompt与task-specific input组合成最终prompt的情况，或者使用templated requests的情况）。它提出的ChunkAttention是一种新的self-attention module。一方面，它采用Prefix-Aware KV Cache（PAKV），将KV Cache切成小的chunk，将它们以prefix tree的方式组织。具有相同prompt prefix的序列的query张量会被batch到一起。另一方面，它引入two-phase partition（TPP）算法加速在PAKV上的计算。它的优点是可以在运行时自动与动态地检测相同的prefix并优化。

Eviction Policy

KV Cache和其它cache一样，当内存不够时也需要考虑eviction。但它与其它cache相比，又有自身的一些特点。因此，传统cache的eviction policy对它未必适用。

StreamingLLM

论文《Efficient Streaming Language Models with Attention Sinks》中它提出的StreamingLLM通过保留特定token的KV Cache改进了LLM在streaming应用中的表现。它指出将LLM用于streaming应用（如多轮对话）会面临两大挑战：1）在decoding阶段，KV Cache会消耗大量内存。2）流行的LLM无法应用到比训练时所用序列更长的场景。流行的Window attention方法只保留最近的固定长度的sliding window中的token的KV Cache。但这种方式在序列长度超过cache size会失效。业界涌现出不少方法进行改进。总得来说，将LLM扩展到lengthy text的方法有三个方向：

  • Length Extrapolation：目标是让在短文本上训练的模型能在推理时处理更长的文本。比如relative position encoding methods，像RoPE，ALiBi。虽然extrapolation能力有所提升，但当文本长度远远长于训练时的window，表现也不尽如人意。
  • Context Window Extension：研究如何扩展LLM的context window，使得forward可以处理更多token，从而提升训练效率。比如FlashAttention和approximate attention methods。
  • Improving LLMs’ Utilization of Long Text：使LLM能更好地捕捉与使用上下文中的内容。
    前两种难以用于streaming应用，即无限长序列的场景。

本文中指出在auto-regressive LLM中，大量的attention score会位于几个initial tokens上，即使它们在语义上并不重要。这些tokens称为attention sinks。直观上，在auto-regressive LLM中，初始的token会对后面所有的token产生影响。基于该观察，StreamingLLM除了保留最近的token对应的KV Cache，还保留initial tokens的KV Cache。它可以让LLM在训练时使用有限长的attention window，而在推理时无需fine-tuning便可以用于无限长的情况。实验证明在序列长度达到120K tokens时仍能保持合理的准确率。

H2O

KV Cache虽然也是cache，但由于它的特殊性，传统用于Cache的LRU/LFU策略可能并不是最适合的。论文《H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models》中提出的H2O是一种用于减少memory footprint的KV Cache eviction policy。该方法基本这样的现象：一小部分的token在计算attention score的过程中有大部分的贡献，去掉它们会对结果有很大影响。LLM中的attention矩阵在推理时稀疏度可达95%（也就是说5%的KV Cache其实就足够用于得到相同的输出）。这部分token称为Heavy Hitters（H2）。统计前置token的attention score的累加可以用于找出这些Heavy Hitters。根据这个观察，文中提出Heavy Hitter Oracle（H2O）框架。它是一种可以动态地在最近的token与H2 token间保持平衡的KV Cache eviction policy。该方法可以在不显著影响准确率的前提下将memory footprint降低5-10x。

Decoding优化

前面提到，LLM推理中的decoding阶段相对来说难以打满GPU，因此往往是优化的重点。业界有几种优化思路：

Token-level early exit

网络模型对于不同的输入，需要的计算是不同的。因此，我们可以用adaptive compute来减少计算量。Adaptive compute主要依赖early-exit策略。Transformer网络是一个多层结构。意味着对于简单的输入，可能只要算前面几层就行。Early exiting可以根据中间层的表征直接输出新token，跳过网络中后面的部分，从而避免计算整个模型。那现在主要的问题就是如何确定exit point，也就是如何判断是否是“难”的。对于这个问题大致有几类方法：1) Heuristic metrics：如entropy, maximum softmax score, 最上两层softmax的差，或者cosine similarity等。但它们的泛化性较差，且threshold需要tuning。2) Learn-to-exit：泛化性好，且不需要threshold tuning。如论文《Confident adaptive language modeling》中就用了这两种方式。3) Hash-function：如论文《A simple hash-based early exiting approach for language understanding and generation》中提出Hash-based Early Exiting（HASHEE）方法用于替代learn-to-exit。这里的hash function用于构建token到exiting layer的映射。基本思想是如果训练中一个样本是容易的，那推理时与它相似的样本也应该是容易的。

论文《SkipDecode: Autoregressive Skip Decoding with Batching and Caching for Efficient LLM Inference》是一种token-level early exit方法，使之能用于batching inferencing和Key-Value Caching。Batching带来的难点是需要等batch中最后的token被处理。KV Cache带来的难点在于如果当前token比其它token退出更晚则需要更新前面token的KV Cache。SkipDecode为batch中的每个sequence position设立单独的exit point，从而克服了这些缺点。它利用了序列结尾更容易预测的特点，序列越到后面所需计算越少（可参见论文Figure 2）。这样做也使前面token的KV Cache不需重新计算。

Parallel Decoding

Decoding阶段的主要挑战在于每个token依赖前面的token，它是迭代式地吐结果，因此整个过程是串行的。想要提高性能，并行化是一个重要手段。Decoding阶段并行化有几种常见手段：

Speculative decoding

LLM推理的性能瓶颈在于auto-regressive decoding。该过程对GPU不友好，因为auto-regressive这种逐个输出token的方式不易并行。Speculative decoding是近几年比较火的一类方法。它可以打破每次只输出一个token的限制。这类方法采用的是Guess + Verify的范式。它分为两步：第一步使用draft model（或称small speculative model，通常是一个计算量小的模型）产生一定长度的token，称为draft（可通过non-autoregressive方式，也可使用更快的autoregressive模型）。第二步对前面产生的draft用LLM模型（或称为target model，verifier model）进行验证或评分。它最大的好处就是这两个阶段都可以并行化，而且可以保证结果与LLM模型产生的一致。

其中draft model可以是target model的一部分（一起从头训练），也可是模型压缩方法得到，或是target model的部分activation作为draft model的输入，然后训练得到。可以看到，这种方法要想高效，很大程度上取决于draft model的能力。它的准确度越高，或者说与target model的输出越吻合，整个方法的提速就越明显。但是如果为了它的准确度过分增大其计算量，又得不偿失。因此，我们需要从其它方面去想办法，比如一次产生多个draft等。

在这种Guess + Verify范式下，衍生出来很多玩法，比如：

  • Genearalizied Aggressive decoding（GAD）：来自论文《Lossless Speedup of Autoregressive Translation with Generalized Aggressive Decoding》。它是一种lossless的auto-regressive translation加速方法。它使用NAT同时得到多个token作为draft，然后用AT进行验证，发现分歧就丢弃。
  • Speculative sampling（SpS）：论文《Accelerating Large Language Model Decoding with Speculative Sampling》中提出。它用小型的draft model产生token序列（可用并行的方式，也可用auto-regressive方式），然后用target model给这段序列打分。最后用reject sampling方法，从左到右判断是否接受，从而还原target model的分布。
  • SpecInfer：论文《SpecInfer: Accelerating Generative Large Language Model Serving with Speculative Inference and Token Tree Verification》中提出。它通过小模型SSM（Small speculative model）进行投机式推理，每次试探推理多步，然后将多个SSM的推理结果汇聚成一个speculated token tree，交由LLM验证，通过高效的树形解码算子实现并行化推理。
  • Lookahead：论文《Lookahead: An Inference Acceleration Framework for Large Language Model with Lossless Generation Accuracy》提出的lookahead方法采用了multi-branch策略。与single-branch策略每次只考虑一个draft不同，它同时取多个drafts（当context length固定的情况下，decoding length增加在前期对于耗时影响很小），再用并行地对其进行验证，最后确定最长的子序列作为输出。由于drafts通常会有共同的prefix tokens，因此它利用trie tree对其进行组织，提出hierarchical multi-branch draft。
  • Minions：论文《Minions: Accelerating Large Language Model Inference with Adaptive and Collective Speculative Decoding》中提出的Minions是基于speculative decoding的LLM推理系统。它通过一些手段提高draft model的低acceptance rate，和降低验证时的开销：1) 为了提高acceptance rate，它提出的majority-voted机制使用树结构表达多个不同draft model的输出序列，并结合draft model过往的表现（average acceptance rate）得到majority-approved output。2) Draft model的speculation length影响target model的总推理时间，而最优的speculation length被数据集、batch size等多种因素影响。系统基于验证时的运行时信息，使用启发式搜索算法动态调节该参数。3) 通过intermidiate resulting pool将draft model与target model的执行解耦，使它们组成pipeline，减少idle time，提高吞吐。

还有一些工作致力于将之应用于小型设备，比如：

  • Staged speculative decoding：论文《Accelerating LLM Inference with Staged Speculative Decoding》主要针对small-batch，on-device的场景。它提出两个改进：一个是tree-structured batches，它生成多个可能的序列，并将之动态构建成树。另一个是对draft model的计算也用speculative decoding加速。也就是说，系统中会有两个draft model。
  • LLMCad：论文《LLMCad: Fast and Scalable On-device Large Language Model Inference》 目标是将LLM布署在如手机，IoT等移动设备上。它将一个小型的LLM放在内存中用于构建token tree。另外用一个高精度的LLM用来验证前者产生的token和修正错误。它主要用到的技术有token tree generation and verification(与token sequence不同，token tree中每个token可以有多个后继候选), Self-adaptive fallback strategy(提出使用tree-cumulative confidence指标来确定是否用target model进行验证),和speculative generation pipeline（让小型模型的产生过程与大型模型的验证过程并行）。

另一类方法不需要额外的draft model，称为Model-free prediction strategies。比如：

  • 论文《Blockwise Parallel Decoding for Deep Autoregressive Models》提出一种blockwise parallel decoding scheme。它采用一些proposal models独立并行地做接下来几个位置的预测，然后通过scoring model选取最长的prefix。为了将scoring和prediciton模型整合在一个模型中，它改造了Transformer，在原decoder output layer后插入多输出的FF层。
  • LLMA：在一些任务（如retrieval-augmented generation, cache-assisted generation和multi-turn conversions）中，in-context reference与LLM产生的输出序列会有很多重复。基于这个观察，论文《Inference with Reference: Lossless Acceleration of Large Language Models》提出了inference-with-reference decoding机制，它先从reference中选出一段文字，然后将之拷贝到LLM decoder，检查它能否接受。该过程可被并行。
  • Aggressive Decoding：论文《Lossless Acceleration for Seq2seq Generation with Aggressive Decoding》提出的Aggressive Decoding方法是一种lossless的decoding算法 。它应用于像Grammatical Error Correction这种输入与输出非常相似的任务。它从输入序列拷贝过来作为drafted decoded tokens，然后并行地对其进行verify。对于翻译任务，Generalized Aggressive Decoding会采用额外的NAT decoding模型，然后用auto-regressive的方式并行验证。
  • Medusa：论文《Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads》使用multiple decoding heads技术提高decode的并行度。它不需要单独的draft model，而是在LLM的last hidden layer上添加多个额外的head，并行地预测后续多个token。对于这些多个候选使用tree-based attention mechanism进行验证。对于这些head的训练，它提供了两种fine-tuning方法：一种是backbone LLM保持不动，另一种是与LLM一起联合训练。后者有更好的性能，但是需要一些手段（combined loss, differential learning rates与heads warmup）来避免破坏LLM的能力。
  • Lookahead decoding：论文《Break the Sequential Dependency of LLM Inference Using Lookahead Decoding》中的lookahead deocding是一种不需要draft model的用于加速LLM推理的并行decode算法。它使用Jacobi iteration method并行的提取和验证多个n-grams。LLM可并行地产生多个disjoint n-grams。这些n-grams可以用于所产生序列。它将auto-regressive decoding建模成一个非线性方程组的求解过程，然后求解得到n-grams，通过验证后的为可用于最终的序列。该非线性系统它可以通过Jacobi iteration method来求解。该方法可以并行化，且可保证在m步内完成（m为变量个数）。
Non-autoregressive Transformers

既然auto-regressive的方式对性能不友好，那还有一个思路就是使用Non-autoregressive（NAT）的方式。NAT以迭代的方式迭代的方式将所有的token一起decode出来。NAR（Non-autoregressive）最开始在NMT（Neural machine translation）领域被提出。与auto-regressive方式相比，它可以并行产生token，因此耗时更少，但是代价是准确率更低。其主要原因是"failure of capturing the target side dependency"。它不像auto-regressive方法一样，产生第t个token时有前面t-1个上下文token信息。业界有一种折中的做法是iteration-based NAT模型。即每一次产生的token再给decoder做refinement，迭代数次。相关方法详细可参见论文《A Survey on Non-Autoregressive Generation for Neural Machine Translation and Beyond》。其中讨论了包括data manipulation, modeling method, traning critierion, decoding algorithm和pre-trained model等方面的工作。

Batching

众所周知，GPU是massively-parallel架构。通常batch越大，可并行的计算任务越多，计算任务更倾向于compute-bound而非memory-bound，也就越容易提升硬件利用率。另一方面，同一batch中的推理无需重新加载模型权重，省去了从HBM到on-chip SRAM的数据搬运开销。但是推理场景下，batch一般较小，因此一般硬件利用率也比训练场景下更低。因此，推理场景下，我们希望将更多任务batching在一些，从而提升吞吐。根据batching的方式不同，可分为static batching, dynamic batching和continuous batching三种。

最简单的方式是static batching。这种方式中，client端就将request进行batching，发到服务端。但作为服务端，假设client做batching太过理想，而且很多时候request来自于不同的client。另外，不同request产生的序列长度变化可能很大。因此，在LLM中实用性不大。

另一种常见的是dynamic batching。比如Triton Inference Server可以动态地将客户端的请示在服务端组成一个大的batch。为了控制对于latency的影响，它可以配置latency的约束（最长不超过多长时间做一次batching）。但是，对于input shape会变化的场景（比如语言类模型），需要对input tensor进行padding。而padding会带来资源的浪费与性能的损失。为了避免这个损耗，就需要对输入数据进行"压缩"。比如论文《Prepacking: A Simple Method for Fast Prefilling and Increased Throughput in Large Language Models》提出batch当中prompt长度差很多时，padding的方式浪费资源严重。Prepacking优化prefill计算。它避免了pad token的计算。它使用bin-packing算法，将不同长度的prompt打包进一个序列。然后修改attention mask和positional encoding来对单个序列中的多个prompt进行计算。

这种长度不规则的张量称为ragged tensor。Ragged batching让用户指定输入中的哪些是有效的数据，从而可以避免显式的padding。接下来的问题就是还需要有相应的kernel来高效处理这种ragged tensor。一种是手写，比如Effective Transformer中的Remove padding算法可用于减少padding部分的计算量。该算法后来也被集成到FasterTransformer中。另一种是自动生成，比如论文《The CoRa Tensor Compiler: Compilation for Ragged Tensors with Minimal Padding》中的CoRA是一个张量编译器，用于为CPU与GPU平台自动生成ragged tensor的高效代码。

如果batch中的request的序列长度（包括prompt与生成的序列）相差很大，那些早结束的request需要等最长的序列完成，导致计算资源被浪费。因此，如果我们将长度相差不多的request放在一起，对性能是更好的。对于一个request，在没有到达最长限制前，它会一直生成token直到得到遇到EOS为止。这里最大的难点在于它会生成多长的序列是不确定的，因此我们很难提前将生成差不多长度输出的request放在一块做。一种直观的思路就是预测其输出序列长度，业界提出的一些相关方法包括：

  • Response Length Perception and Sequence Scheduling：论文《Response Length Perception and Sequence Scheduling: An LLM-Empowered LLM Inference Pipeline》使用length predictor预测response的长度，然后将response长度差不多的放到一个micro-batch中，从而加速推理过程。对于预测偏离较大的情况，提出FCR（Failure Collection and Recomputation）机制，它将batch中最大预测长度作为新产生token数量的上限。如果超过，就会判定为失败，并在后面重计算。VBS（Variable Batch Size）让短的response的batch有更大的batch size。
  • S3：论文《S3: Increasing GPU Utilization during Generative Inference for Higher Throughput》中提出的方法全称Scheduling sequences with speculation。它是一个通过预测输出序列长度来优化吞吐，减少内存浪费的框架。文中指出像FasterTransformer会按最大的序列长度预留内存，导致内存的浪费，并限制batch size，降低吞吐。该方法预测输出的序列长度，并基于预测进行调度从而提高硬件利用率与吞吐。它包含三个部分：Predictor用于预测输出序列长度；Scheduler根据预测与HBM使用信息将batch派发到GPU计算文本生成（视作bin packing problem用decreasing first fit算法解决）；Supervisor在后台运行，监控GPU利用率，可用HBM并将信息给scheduler。对于预测错误的情况。它会抢占那些实际生成序列超过预测值的case，为其增加内存并计算。

前面一些方法使用预测的方法来将输出序列长度差不多的request放在一块。但这样就依赖于预测的准确性。如果预测不准，不仅起不到预期的效果，还搭上了预测本身的计算开销。

回过头来思考，之所以输出序列长度差异会导致浪费，更本质的原因是因为传统采用的是request-level batching，即一个batch中的request在batch没全部完成前是不会变的。这样，就衍生出另一种策略。就是根据LLM中迭代式产生token的特征，每当有request完成，就把它移出batch，然后把新的request放入batch。它有很多名字，如continuous batching（vLLM中的叫法）, iteration-level scheduling（Orca中的叫法）, in-flight batching（TensorRT-LLM中的叫法）。很多框架中都有实现：

  • Orca：论文《Orca: A Distributed Serving System for Transformer-Based Generative Models》中提出的Orca首先研究了这个问题，并实现了iteration-level scheduling的机制。它在每次迭代后都会进行调度，选取要处理的request交给engine。当batch中有序列完成token的产生，就会将新的request放进来。另外，对于naive的batching，有三种情况不适用：都是prefill，但input token数量不一样；都是decode，但token index不一样；prefill与decode混合。因此，它提出selective batching，即选择性地对部分操作进行batching：对线性操作，直接处理；对Attention模块先按request切分，单独处理再合起来。
  • Hugging Face在text-generation-inference（TGI）中实现了continuous batching。
  • vLLM中实现了continuous batching（或称batching with iteration-level scheduling）。它比naive batching有3x的吞吐提升。详细介绍可参见"How continuous batching enables 23x throughput in LLM inference while reducing p50 latency"及源码。
  • TensorRT-LLM中包含称为Batch Manager的组件，用于支持in-flight batching。这部分没有 随着项目开源，相关介绍可参见：https://nvidia.github.io/TensorRT-LLM/advanced/batch-manager.html。
  • DeepSpeed-FastGen：论文《DeepSpeed-FastGen: High-throughput Text Generation for LLMs via MII and DeepSpeed-Inference》中的DeepSpeed-FastGen基于DeepSpeed-MII与DeepSpeed-Inference提供了高效易用的LLM serving系统。它采用了Dynamic SplitFuse技术，可以将prompt与generation阶段放到一起。比vLLM在吞吐与延迟上都有2倍多的提升。它观察到batch相对于token数量，对于latency影响很小。同时token-throughput是个凹函数。根据这些现象，Dynamic SplitFuse将长prompt切成小的chunk，短prompt进行合成，从而使得每次forward计算的batch size恒定。这样达到更好的响应，高效与低波动的目的。
  • SARATHI：论文《SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills》指出由于prefill与decode阶段的计算模式差异很大，导致micro-batch间计算时间差很大，从而产生pipeline bubble，导致硬件利用率低。而目前的batch调度方式中，像FT的request-level inference scheduling以request为粒度，只有当batch中的所有request结束才跑下一个batch。Orca, vLLM, TGI中的iteration-level scheduling虽然让request可以动态地进入与离开batch，但是对于batch中执行时间差异大的问题没有关注，导致bursty resource utilization与pipeline bubble。基于两点观察（一方面，对于给定模型与GPU，增加prefill token数量的边际收益在超过一定值后是减少的。另一方面，随着hidden dim增加，使GPU计算饱和所需的chunk size减少），SARATHI采用了chunked-prefill技术将prefill请求切分成相同大小的chunk，利用attention mask实现。另外它使用decode-maximal batching技术将单个prefill chunk与decode组成batch。这样，prefill chunk可以充分利用GPU算力的同时，也把memory-bound的decode做了，从而减小了decode的开销。
  • SARATHI-Serve：论文《Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve》基于SARATHI给出了更详细具体的调度策略。它指出当前的LLM推理scheduler分两类：prefill-prioritizing和decode-prioritizing。前者如iteration-level batching系统vLLM, Orca；后者如FT，Triton Inference Server采用request-level batching。在已运行的request的decode阶段没完成之前不会调度新的request。前者相比后者吞吐更高，因为后续decode迭代的batch size更高，但代价是latency变长。SARATHI-Serve基于SARATHI的技术，提出了stall-free batching，避免decode因为一起运行的prefill chunk导致generation stall，使throughput与latency同时得到优化。SARATHI-Serve先计算能在一个batch所能计算的最多token所需的budget，然后在每次调度迭代中，先选择所有运行中的decode，再选择那些跑了一半的prefill。最后再考虑新的request。

上面的工作主要想着怎么把prefill与decode两个阶段放在一起，从而提高计算密度。业界现在还有一种看似有些相反的思路是既然它们计算模式如此不同，何不把它们尽可能分开，以减少相互影响，方便调度。

  • Splitwise：论文《Splitwise: Efficient Generative LLM Inference Using Phase Splitting》中认为LLM中两个阶段prompt computation和token generation计算模式差别很大，一个是compute intensive，一个是memory-intensive。既然如此不同，那为何不干脆将它们分到不同的机器上，这样可以使用适合每个阶段的硬件。论文基于这个思想使用Splitwise技术设计了LLM推理集群。
  • TetriInfer：论文《Inference without Interference: Disaggregate LLM Inference for Mixed Downstream Workloads》指出现有系统忽略prefill与decode阶段的区别，导致它们相互影响。文中提出的TetriInfer将prompt切分成固定大小的chunk方便accelerator用满计算资源。同时它将prefill与decode分开，独立运行。另外，它使用两层的调度算法，考虑预测的资源使用，避免了decode scheduling hotspot。
  • DistServe：论文《DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving》指出现有系统为了latency需求，需要调整优先级，或者over-provision计算资源。DistServe将prefill与decoding阶段放到不同的GPU上，以避免它们相互影响，为每个阶段联合优化资源分配与并行策略。它的目标是提升goodput，即满足latency要求下的throughput，而非单一性能指标（TTFT或TPOP）。

Distributed

多数分布式AI计算相关工作和是针对训练，感兴趣的可以看之前写的文章《LLM时代中的分布式AI》。但其中一些思路其实对推理也适用，如最为常见的TP，PP等。很多框架如FasterTransform，Megatron-LM，DeepSpeed, Petals等都对分布式推理有支持。

一些可用于AI推理的分布式方法如：

  • Petals：论文《Petals: Collaborative Inference and Fine-tuning of Large Models》中提出的Petals是一个用于LLM的collaborative inference与fine-tuning系统。它由crowd-sourced distributed training启发，将互联网上多方的资源整合起来，将模型的不同部分分配到这些计算资源上。为了减少传输量，它使用dynamic blockwise quantization压缩hidden state，并将权重压缩到8-bit。另外它还设计了server端的load balancing机制（每个server选取最差吞吐的block，并周期性将该信息报告给distributed hash table。同时周期性检查是否启动re-balancing过程），以及client端的routing机制（ping最近的server，并通过beam search找到最优的路径）。
  • 论文《Efficiently Scaling Transformer Inference》研究大模型在有latency要求以及长序列下的挑战。它提出构建优化LLM推理的partitioning框架可以提升scalability。比较挑战的地方在于不同的情况其访存与计算的性能瓶颈是不一样的。比如在小batch size与seq len情况下，加载权重是主要耗时点；而大batch size与长seq len情况下，加载KV Cache的耗时会占大头。本文方法构建一个解析模型，用来挑选TPU v4集群上的最优多维切分策略。另一方面，它将内存优化应用于PaLM模型的MQA结构，减少不必要的开销，从而增大batch size和提高吞吐。在low-level优化上，利用Looped CollectiveEinsum技术使通信与计算并行。
  • Ring Attention： 论文《Ring Attention with Blockwise Transformers for Near-infinite Context》基于BPT（blockwise parallel transformer），将长序列划分成块，然后将块的计算分布到多个host上。这些host设备组成逻辑上的ring，每个device会将它的key-value块拷贝传到下一个device上，同时接收前一个device来的key-value块。只要块的计算时间长于传输时间，数据传输就可以被hide，不会增加额外的开销。
  • Infinite-LLM：论文《Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache》中提出的DistAttention是一种分布式的attention算法。为了克服传统模型并行应用在attention上的资源利用率低的问题，DistAttention将传统的Attention切成更小的称为macro-attention的单元，以及对应的KV Cache(称为rBlock)。每个rBlock包含固定个数token对应的KV Cache以及元信息。每个LLM service instance有一个rBlock manager（称为rManager），负责监管本地的rBlocks。DistKV-LLM是一个基于该算法的分布式LLM推理系统，它可以调度跨数据中心的GPU与CPU内存。另外，它还提出一种称为DGFM的算法可以解决分布式KV Cche环境中的内存碎片的问题。
  • ExeGPT：论文《ExeGPT: Constraint-Aware Resource Scheduling for LLM Inference》中提出的ExeGPT用于在分布式LLM推理系统中找到与执行最优的执行调度，从而在满足latency约束的前提下最大化吞吐。它会确定batch size，partial tensor parallelism等参数。另外，它引入基于Round-Rogin Allocation与Workload-Aware Allocation的调度策略，用于不同的NLP workload。整个系统包含四个组件：XProfiler用于测量所有配置下层的执行时间, XSimulator构建XScheduler所需的timeline, XScheduler基于simulator的估计找到最优的配置参数，XRunner用这些配置参数执行。
  • SpotServe：论文《SpotServe: Serving Generative Large Language Models on Preemptible Instances》用于preemptible instance的分布式LLM serving system。一般云厂商会提供一些preemptible GPU instance（如AWS spot instance，Azure spot VMs）。它们价格比on-demand instance低90%，但缺点是当资源紧张时可能在任何时候被抢占。SpotServe根据instance的可用情况与workload的波动动态调整LLM的并行配置，以达到throughput，latency和开销的balance。另外，它将instance迁移规划问题建模成bipartile graph matching问题进行优化，从而减少开销。最后，它利用云厂商提供的grade period（抢占前的准备时间）引入了stateful inference recovery机制。可以更细粒度的commit inference进度，便于resume inference。

分布式计算意味着会有很多通信算子。这些算子容易成为性能瓶颈。因此，减少它们的开销是优化的重点。使用NCCL这样的三方库优点是方便，大部分情况下性能也不错。要进一步深入优化，就需要让通信与计算算子融合。一些工作比如：

  • T3：论文《T3: Transparent Tracking & Triggering for Fine-grained Overlap of Compute & Collectives》中提出的的T3通过配置producer的输出地址空间，以一种“透明”的方式将producer操作与后面的通信操作融合起来。该方法还加了轻量级的track和trigger机制，调度producer的计算与通信。另外，它还将compute-enhanced memory用在通信计算上。
  • Centauri：论文《Centauri: Enabling Efficient Scheduling for Communication-Computation Overlap in Large Model Training via Communication Partitioning》指出现有系统中对于通信与计算overlap的优化采用fine-grained kernel fusion或有限操作调度的方式在异构训练环境中优化空间有限。因此，Centauri将通信切分，并用层次化的调度机制优化overlap。它将切分空间抽象成三个维度：primitive substitution, topology-aware group partitioning和workload partitioning。
  • FLUX：论文《FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion》用于LLM在训练与推理时的计算与通信overlapping。在有TP的情况下，以往的工作切分比较粗粒度，并依赖于精巧的调度，不太适用于GPU。Flux（Fine-grained Communication Overlapping）可用于AllGather - GEMM，GEMM - ReduceScatter这样的pattern。与以往工作相比，它使用细粒度的切分，并将之融合成一个大的kernel。此外，它还结合了tile coordinate swizzling, GPU instruction selection, communicataion order selection等优化技术。实现上，它基于CUTLASS项目，使用auto-tuning选择transfer method及模板参数。

Offloading

对于内存不足的问题，除了优化内存占用外，另一个常用的方法是利用memory hierarchy中更下一级的存储，比如CPU memory，或者磁盘。这方面业界相关的可用于LLM的工作比如：

  • DeepSpeed Inference：论文《DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale》中提出的DeepSpeed Inference是一种支持单卡、多卡和异构的transformer模型推理解决方案。DeepSpeed Inference包含DeepSpeed Transfomer与ZeRO-Inference两个部分。前者为三层架构，包含单GPU transfomer kernel，多GPU dense transformer layer与多GPU sparse transformer layer。它通过Deep-Fusion算子融合和CUDA-Graph来减少访存和kernel launch开销。对于多GPU场景，它使用TP, PP并行方式。后是基于GPU+CPU+NVMe的异构方案，支持多种配置，如dense or sparse(MoE), small or lage batches, billions to trilliions参数级别，单GPU或多GPU。为了减少减少从DRAM或NVMe拿模型权重的开销，它采用了Prefetching与Multi-GPU PCI-e bandwidth utilization技术进行优化。相比GPU-only的方案，它所能支持的模型大25x，并达到50%以上的峰值性能。
  • FlexGen：论文《FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU》中提出的FlexGen是面向高吞吐批处理的LLM推理框架。它整合了GPU，CPU和磁盘的存储和计算资源，并调度I/O操作使之更加高效，从而实现高吞吐。再加上压缩方法（group-wise quantization和sparse attention）与分布式pipeline并行。它将offloading stragey的优化问题建模为graph traversal problem，定义offloading strategy的搜索空间，它考虑computation schedule, tensor placement和computation delegation，然后构建解析的cost model，并通过线性规划方法找到解。
  • LLM in a flash：论文《LLM in a flash: Efficient Large Language Model Inference with Limited Memory》中提出的方法将模型参数放在flash memory中，然后按需放入DRAM。它引入cost model来为数据传输优化提供指导。利用稀疏性选择性地从flash memory只加载那些输入非零或者预测输出非零的相关参数。减少数据加载和内存使用效率使用了两个技术：windowing和row-column bundling。前者只加载和临时缓存最近token的；后者将up-projection和down-project层的行与列一起存放，有利于提高吞吐。
  • PowerInfer：论文《PowerInfer: Fast Large Language Model Serving with a Consumer-grade GPU》中提出的PowerInfer是一种GPU-CPU混合推理引擎。网络中一小部分神经元，称为hot neuron几乎会被所有的输入激活。而大部分则为cold neuron，只在特定输入时被激活。Hot neuron会被预加载到GPU，而cold neuron在CPU上计算。这样减少了GPU的内存需求和CPU-GPU的数据传输。PowerInfer还引入adaptive predictor和neuron-aware sparse operator提升计算效率。
  • Hugging Face Accelerate：支持offload一部分权重。它使用PyTorch中的hooks机制，在forward pass的开始前将CPU或者disk上的权重加载到GPU上，结束后再将之放回原地。

Compression

LLM的常见压缩方法和传统神经网络模型类似，有pruning, knowledge distillation, quantization, low-rank factorization这几大类。LLM会给这些技术带来一些新的挑战。

Quantization

我们知道，模型的量化可分为PTQ（Post-training quantization）和QAT（Quantization-aware training）两种。LLM因其训练成本很大，因此QAT很多时候不实际，而PTQ要实用地多。但是，PTQ往往会导致准确率下降很多。因此，LLM量化中主要目标之一是减少PTQ的准确率损失。量化本质上是用更小的位数去表达原数据，需要在范围与精度间做取舍。当网络参数量大到一定程度，activation就容易出现outlier。它会使量化范围增大，从而使得量化误差变大，最终降低准确率。一些业界的一些量化方法包括：

  • GPTQ：论文《GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers》中的GPTQ方法是一种基于OBQ（Optimal Brain Quantization）方法（使用了二阶信息）的权重的PTQ量化方法。OBQ方法采用layer-wise的方式，对每层权重，独立处理每一行。每次量化一个权重，然后更新其它权重以弥补该权重量化带来的误差。本文为了将之用于LLM，提出了几项改进：Arbitrary Order Insight，Lazy Batch-Update, Cholesky Reformulation。经过改进后，它可用于LLM（如OPT-175B，BLOOM-176B）的近于无损的4-bit量化，使得用单张GPU跑LLM成为可能。缺点是由于硬件不支持混合精度（FP16 x INT4），难以加速。
  • ZeroQuant: 是一种流行的PTQ方法，可以将大型transformer-based模型（如BERT，GPT3）的weight与activation压缩至INT8。它有几个主要的变体：
    • ZeroQuant：论文《ZeroQuant: Efficient and Affordable Post-Training Quantization for Large-Scale Transformers》提出的ZeroQuant是用于INT8和INT4/INT8混合精度量化的end-to-end的PTQ和推理pipeline，可用于weight与activation。它采用的fine-grained hardware-friendly quantization scheme主要包含group-wise weight quantization与token-wise quantization for activation两项技术。前者将weight矩阵分成若干group，然后对每个group单独量化。同时考虑了GPU上的硬件约束。后者动态地计算每个token的范围，从而减少量化误差。为了减少额外操作带来的开销，它采用kernel fusion将量化算子与前面的算子（如LayerNorm，bias-ad, GeLU）融合。此外，它引入了CUTLASS的INT8 GEMM实现对量化模型进行加速。
    • ZeroQuant-V2：论文《ZeroQuant-V2: Exploring Post-training Quantization in LLMs from Comprehensive Study to Low Rank Compensation》对round-to-nearest (RTN)，GPTQ，ZeroQuant等PTQ方法分别研究了weight-only，activation-only和weight-and-activation下的效果。基于结果，它观察到当前的方法均无法在INT4-weight或W4A8下达到原模型质量，于是提出了Low-Rank compensation (LoRC)方法，使用低轶分解来近似量化误差矩阵的。该方法以少量的model size与计算量增长为代价提升量化后的准确率。
    • ZeroQuant-FP：论文《ZeroQuant-FP: A Leap Forward in LLMs Post-Training W4A8 Quantization Using Floating-Point Formats》基于H100 GPU引入的FP8和FP4格式，将W4A8的PTQ应用于LLM。
  • LLM.int8：论文《LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale》中的LLM.int8引入mixed-precision decomposition，它对矩阵相乘中的每个内积用单独的normalization constoant进行量化，对于outlier，它使用FP16精度，而其它的以8-bit相乘。缺点是在加速硬件上较难高效实现。
  • OdysseyLLM：论文《A Speed Odyssey for Deployable Quantization of LLMs》是考虑硬件加速的W4A8解决方案。它包含Symmetric learnable weight clipping（LWC），iterative Hessian-based compensation，和用于W4A8的kernel实现（FastGEMM）。在混合精度GEMM中，它使用了Kernel fusion（现代GPU只支持同类型的GEMM计算，因此将SINT4toS8与GEMM融合进一个kernel），Removal of INT8 subtraction（现在GPU没有signed 8-bit整数的减法，因此使用对称量化）和Reusing the sign bit（优化int4到int8转换的开销）技术。
  • RPTQ：论文《RPTQ: Reorder-based Post-training Quantization for Large Language Models》指出LLM量化activation的挑战除了outlier外，还在于不同channel的范围差异很大。因此它将channel重新排列，并按cluster进行量化。不同cluster内使用不同的量化参数。同时为了减少reoreder的开销，它将reorder操作融入layer norm和linear layer。RPTQ对LLM的activation进行3-bit的量化。
  • SmoothQuant：论文《SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models》提出的SmoothQuant方法是一个training-free, accuracy-preserving的通用PTQ方法。它可用于LLM的W8A8（8位权重，8位activation）量化。主要应用于LLM中的GEMM。基于权重比activation更少outlier，因此也更容易量化的事实，SmoothQuant通过一个数学上的等价变换以离线方式将activation的困难转移到weight上。它让输入activatino除以一个per-channel smoothing factor，同时权重乘以相同的系数，就可以保证结果不变。
  • AWQ：论文《AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration》中的AWQ（Activation-aware Weight Quantization）方法是一种对硬件友好的，用于LLM中low-bit weight-only量化的方法。它基于权重不是同等重要的观察，即保护很小部分（0.1%-1%）的主要权重可以有效减少量化误差。它通过在calibration set中activation的值来判断是否为主要权重，即对应绝对值较大的activation的权重更重要。
    通过分析得出将salient channel的值放大可以减少相对量化误差，因此它设计了一种per-channel scaling方法自动找寻最优的scaling来最小化量化误差。AWQ可使LLM被部署于desktop与mobile GPU。
  • SqueezeLLM：论文《SqueezeLLM: Dense-and-Sparse Quantization》提出的SqueezeLLM是一个3-bit的超低精度PTQ框架。它的两个核心思想是：1) Sensitivity-based non-uniform quantization：由于权重的分布不是均匀的，因此采用非均匀量化。找寻最优的非均匀量化配置转为k-means问题。二阶层数信息可以作为权重要度的测度，因此用它作为k-means clustering问题中的样本权重。2) Dense-and-Sparse decomposition：由于权重的值分布中绝大部分的值聚集在小部分区间（如LLaMA-7B中99.9%的权重在10%的分布区间中），因此如果将小部分outlier去除，我们便能将权重的范围缩小很多倍。因此，可以用稀疏格式（CSR）存放outlier和sensitive weight values。
  • KVQuant：论文《KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization》考虑到长序列下KV Cache将是内存使用中的瓶颈，因此主要考虑KV Cache的量化。该方法将KV Cache量化到4位以下。它采用以下技术：Per-Channel Key Quantization（调整Key activation的量化维度使之更贴合分布），Pre-RoPE Key quantization（在RoPE前量化Key activations减小量化影响），Non-Uniform KV Cache quantization（per-layer sensitivity-weighted non-uniform datatype更好表达分布），Per-Vector Dense-and-Sparse Quantization（将outlier独立出来减少量化范围的skew）, Q-Norm（normalize量化中心点避免distribution shift）。该方法用于LLaMA，LLaMA-2和Mistral模型达到3-bit量化。
  • SpQR：论文《SpQR: A Sparse-Quantized Representation for Near-Lossless LLM Weight Compression》提出Sparse-Quantized Representation（SpQR）是一种新的压缩格式和量化技术，可以接近无损的压缩LLM。它主要用来解决量化至低bit时导致的accuracy loss问题。该方法先找出那些会导致量化误差的outlier weight并将之分离出来，然后将它们以高精度来存放，同时将其它的权重压缩至3-4位。同时，它还提供了编码与解码的高效算法，与针对SpQR的高效的GPU推理算法。

Pruning

Pruning通过裁剪网络进行模型压缩。对LLM进行pruning时，需要在减小模型规模的同时保证原模型的能力尽可能地保留。

  • SparseGPT：论文《SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shot》是一种Post-training pruning方法。它指出大的模型更易被压缩，GPT系的模型参数至少可被裁减一半，而不需要retraining，且精度损失很少。传统的pruning方法，如AdaPrune，耗时很长（175B Transformers需要数百小时）。SparseGPT可用于GPT系模型，如将OPT-175B和BLOOM-176B在4.5小时内达到60%的unstructured sparsity，且准确率与perplexity损失很小。它基于mask selection + weight reconstruction的范式，并通过Fast approximation reconstruction（优化Hessian矩阵计算）adaptive mask selection（使用iterative blocking）提高了效率。它还可适用于semi-structured pattern（n:m sparsity format），比如结合Ampere NV GPU中的2:4 sparsity。
  • LLM-Pruner：论文《LLM-Pruner: On the Structural Pruning of Large Language Models》提出的LLM-Pruner是一种structural pruning方法，它基于梯度信息选择性地删除一些不重要的结构。可以做到task-agnostic，并不依赖原始训练数据集。之后用LoRA可以用少量数据快速恢复准确率。
  • DeJavu：论文《Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time》中提出的Contextual sparsity定义为attention head与MLP参数的子集。对于给定的输入，它可以获得与dense model相似的输出。DeJa Vu是一个使用low-cost算法对于每layer的指定输入预测contextual sparsity的系统。另外它设计了asynchronous predictor减少相关开销。
  • Flash-LLM: 论文《Flash-LLM: Enabling Cost-Effective and Highly-Efficient Large Generative Model Inference with Unstructured Sparsity》指出产生式模型推理的性能瓶颈在于skinny matrix multiplication，因为它无法用Tensor Cores。Unstructured pruning往往比structured pruning有更高的准确率，但unstructured sparse MatMul（SpMM）无法利用tensor core，现有的unstructured SpMM实现（如cuSPARSE, Sputnik）除非sparsity达到很高（98%，86%）才能胜过cuBLAS。文中提出的Flash-LLM是一个高效的GPU库用于大型生成模型推理，它支持tensor core上unstructured sparsity。它的核心是采用Load-as-Sparse and Compute-as-Dense方法用于unstructured sparse matrix multiplication（SpMM）。
  • Sheared LLAMA: 论文《Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning》中提出的Sheared LLaMA是一种Structured pruning方法，它主要采用两个技术：1) Targeted structured pruning：它在不同粒度为模型参数学习mask，该mask与模型参数以min-max目标函数被联合优化。该mask决定了对应的子结构是否被裁剪。得到该裁剪后的目标结构后，通过continuing pre-training该裁剪后的网络提升表现。2) Dynamic batch loading：观察到pruning会相对保留那些专用（low-entry）领域的知识。因此将预测loss作为参考，然后在训练过程中根据loss与参考值的差距更新各domain的权重，再根据该权重来采样训练数据。

Low-Rank

Low-Rank Decomposition基本思想是将一个大型矩阵分成两个更小矩阵的积，从而减少参数量。它有好处是处理后仍是dense matrix，因此可以充分利用硬件与为之写的高效的dense matrix计算库。相关工作如：

  • LoRA：论文《LoRA: Low-Rank Adaptation of Large Language Models》指出大模型动不动就是成百上千GPU训练几周甚至几月。因此对于很多任务会选择fine-tuning，而不是重头训练。于是就诞生了Parameter-Efficient Fine-Tuning (PEFT)方法，研究如何高效的做fine-tuning。这类方法中近两年比较火的就是LoRA（Low-Rank Adaption）。以往研究发现那些巨量参数模型其实都存在在低维度上，因此在为downstream task训练时会在Transfomer模型的每层的dense layer插入可训练的rank decomposition矩阵（即low-rank矩阵）。这样，权重可写成 W = W 0 + B A W=W_0 + BA W=W0​+BA。在fine-tuning时，pre-trained model权重不变，只训练那插入的low-rank矩阵。实验分析发现这个增量会为特定的downstream任务放大那些重要的但在pre-training过程中没有被强调的feature。它能使可训练参数减少万倍，GPU memory需求减少3倍。
  • LORD（Low Rank Decomposition）：论文《LORD: Low Rank Decomposition Of Monolingual Code LLMs For One-Shot Compression》中使用Low Rank Decomposition（LoRD）方法来减少模型size。SVD可以用来生成r-rank的近似矩阵 W = U S V T W=USV^T W=USVT。对于对称矩阵可以用特征值分解 W = Q Λ Q T W=Q \Lambda Q^T W=QΛQT。但是，SVD没有考虑输入与输出数据的分布，同时计算还非常费时。LORD利用activation是low-ranked的特点，使用Atomic Feature Mimicking（AFM）方法，基于小量calibration数据产生low-rank矩阵。该方法可应用在像CodeGen和StarCoder这样的代码生成模型上。
    TensorGPT：论文《TensorGPT: Efficient Compression of the Embedding Layer in LLMs based on the Tensor-Train Decomposition》将LLM模型中Embedding Layer（在GPT-2中占31.02%的参数）压缩成low-rank的张量格式（MPS）。它基于TTD（Tensor-Train Decomposition）方法，其最常见的形式是MPS（Matrix Product State）。它使用TT-SVD（Tensor-Train Singular Value Decomposition）算法将一个大型的N阶张量分解成N个2阶或3阶的张量（称为core tensor）。基于该方法，TensorGPT将embedding矩阵中的每个token embedding生成MPS（Matrix Product State）。该格式可被分布式计算。GPT-2中embedding layer可在准确率不受影响的前提下压缩3.31x倍。

LoRA还可以与量化方法相结合，如：

  • QLoRA：论文《QLoRA: Efficient Finetuning of Quantized LLMs》中的QLoRA在4-bit量化的pretrained模型上做LoRA的训练（计算用BFloat16）。它的主要技术有：1) 使用4-bit NormalFloat（NF4）：适合正态分布权重的数据类型。它可以为每个参数平均减少0.37 bits。2) Double Quantization：两次量化，将第一次量化的量化常量作为第二次量化的输入，减少常量的memory footprint。3) Paged Optimizers：使用unified memory，避免长序列时的做gradient checkpointing时产生的memory spike。该特性为optimizer state分配paged memory，当GPU memory用尽时会evict到CPU RAM。当optimizer update需要时再移入。
  • QA-LoRA: 论文《Quantization-Aware Low-Rank Adaptation of Large Language Models》中提出QA-LoRA（quantization-aware low-rank adaption）算法。它关注量化和adaption自由度的不均衡问题，即权重矩阵的每一列只有一对量化参数，但有很多LoRA参数。针对该问题，它使用group-wise operators一方面让每组单独量化从而增加量化的自由度，另一方面，组内共享LoRA参数从而减少adaption自由度。它为LoRA增加了两种能力：在fine-tuning阶段，LLM权重量化成INT4。在fine-tuning后，LLM和附加的权重被整合进量化模型，而没有准确率损失。

Distillation

因为很多LLM模型是通用的，我们可以将LLM中对于特定任务的能力用更小的模型抽取出来。用小型模型来（称为student model）进行训练，学习原LLM模型（称为teacher model）的概率分布，这种方法称为知识蒸馏。在LLM之前，该方法就被用于LM，如DistillBERT。在LLM下的相关工作如：

  • MiniLM：论文《MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers》中提出的MiniLM方法让Student model通过训练模仿teacher model中的self-attention module。它还使用了Teature assistant技术，即先将teature model distill到teacher assistant（它与teacher model的layer num一样，与student model的hidden size一样）中，再将之作为teacher model训练student model。
  • Distilling step-by-step：论文《Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data and Smaller Model Sizes》中提出的Distilling step-by-step范式是一种可以用更少训练数据训练task-specific的小模型。它的主要思想是利用LLM的推理能力以更加data-effient的方式训练更小的模型。其过程分两步：首先给LLM未标注的数据，再通过CoT（Chain-of-Thought）让其生成结果以及相应的rationales（以自然语言描述的对最终输出标签的解释）。然后，利用这些rationals连同标签一起作为小模型的输入进行训练。这些rationals能够提供更丰富更细节的信息，因此可以使小模型的训练更高效。

Others

业界还有一些从其它角度进行优化的工作，比如：

  • FastServe：大多数工作的目标是提高吞吐，但在一些推理场景中latency很重要。论文《Fast Distributed Inference Serving for Large Language Models》中的FastServe挖掘LLM推理中的auto-regressive模式，让其能以output token为粒度抢占。它使用可抢占调度来最小化JCT。
  • LongLLMLingua：论文《LongLLMLingua: Accelerating and Enhancing LLMs in Long Context Scenarios via Prompt Compression》指出LLM推理的表现决定于输入prompt中关键信息的密度与位置。它提出的LongLLMLingua基于LLMLingua框架，是一种应用于长上下文的prompt compression方法。首先，它使用一些小的语言模型，如GPT2-small或者LLaMA-7B从prompt中识别和删除一些不重要的token。然后，再用LLM对压缩后的prompt进行推理。为了有效地进行prompt压缩，关键信息密度，它使用了question-aware coarse-to-fine压缩方法，document reordering机制，dynamic compression ratio，post-compression subsequence recovery策略等技术。
  • EcoOptiGen：超参数对于框架性能的重要性不言而喻。但通常系统中的超参很多且相互影响，这给超参数的确定带来了很大的难度。很多时候人们会通过经验或者一些规则来确定它们，代价是会让性能打一些折扣。论文《Cost-Effective Hyperparameter Optimization for Large Language Model Generation Inference》中的EcoOptiGen框架使用超参数优化和cost-based pruning提升token generation的性能。它着眼于优化LLM中的一些超参数，比如number of responses，temperature(控制生成文本的随机性)，max token(output中最多生成token数量)，stop（使token产生停止的字符串列表）等。考虑到这些参数间的相互影响，文中方法对它们进行联合优化。它使用黑盒优化方法BlendSearch。该方法结合了Bayesian optimization和local search。前者用于产生后者的起始点，后者会在保证收敛速度和开销上界的前提下进行随机方向的搜索。

文章来源: https://blog.csdn.net/ariesjzj/article/details/139693922
如有侵权请联系:admin#unsafe.sh