[Paper Reading] FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU

Paper Source: FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU (arxiv.org)

Background

LLM推理的相关研究一般从系统和算法层面同时进行考虑，其中稀疏和量化属于算法层面的研究。

token生成过程可分为2个阶段：prefill phase, decoding phase。

prefill phase，主要是input sequence(即prompt)进入模型后的相关计算，包括KV cache的计算以及attention相关的计算。
decode phase，主要是为了输出generated token而在模型里做的相关计算，包括KV cache的更新以及attention相关的计算。

本文目标在于设计出一个在单一消费级GPU上的高吞吐 offloading 策略。

Hardware

Offloading Strategy

Problem Formulation

three-level memory hierarchy: GPU-CPU-Disk

其中，GPU内存最小但最快，磁盘最大但最慢

当 GPU 无法完全容纳 LLM 时，我们需要将其卸载到二级存储中，并通过部分加载 LLM 来逐个执行计算。

将模型推理建模为对计算图的遍历

Computational Graph

Search Space

基于某种数学模型约束各个参数取值范围的调度算法，也就是要在数学模型限定的取值范围这个space里search出最优的参数组合。

调度算法描述

很容易想到的遍历方式是：

row-by-row: 对一个batch来说，这种方式最快；且在一行遍历结束后 KV Cache 可以立刻释放。但由于相邻层的权重并不共享，因此需要不断加重权重，导致大量I/O开销。
col-by-col: 按列遍历时，吞吐高，同一层的权重可以重复使用，因此只需要加载激活值和 KV Cache。但由于激活值和 KV Cache 需要存储下来，存储压力大，不可能遍历完整列。

为了平衡上面的两种方式，引入一个新的粒度 block

Computational Schedule

作者简单起见设计了这种方式，这种方式不是最优的，但文中证明该方式至多比最优的差两倍。

在此基础上，设计了一个block的计算调度算法

Algorithm

上述算法中最内层循环里的前6个函数可以用6个线程做并行处理，从而把计算和IO做了overlap，也就是利用overlapping隐藏了CPU和Disk的访存延迟。

Tensor Placement

各类tensor，即weights，activations，KV Cache分别按比例划分，保存在不同存储介质上（GPU、CPU(内存)、Disk），即所谓的Tensor
Placement

wg, wc, wd    # weights 
hg, hc, hd    # hidden state, 也就是activation 
cg, cc, cd    # cache, 即kv cache  

wg + wc + wd = 1    # 即100% 
hg + hc + hd = 1 
cg + cc + cd = 1

需要找到一个最优参数组合使得整体吞吐率性能最优。

Computation delegation

虽然用GPU算会很快，但当KV Cache很大的时候，是把一大坨KV Cache从内存挪到显存用GPU算呢？还是只把激活从显存挪到内存，用CPU算呢？

如果把数据从主机侧传输到设备侧(GPU的HBM)的IO开销太大(数据搬移太慢)，那干脆不做数据搬移，不用GPU计算，而直接把数据从CPU内存(DRAM)搬移到CPU里计算，即所谓的Computation delegation，也即CPU delegation。

比如可以直接用CPU计算attention score，因为decoding阶段 attention score 的计算实际是IO-bounded，也就是从CPU侧搬移数据到GPU的时间比直接在CPU侧计算的时间还要长，即通信开销 > 计算开销。

性能预估

Cost Model

所谓Cost Model就是一个用来预估推理时延和内存占用的数学模型，注意这里的model是数学模型而非 NN model，但cost model的获得也是可以基于机器学习技术的，比如 TVM compiler 里的 cost model 就可以使用 XGBoost 算法来获得。计算一个block的 total latency 可以表示为T。prefill phase 和 decoding phase 两个阶段的跑完一个 layer 的 latency 分别表示为 T_pre 和 T_gen 。

1
2
3

T = T_pre*layers + T_gen*(n-1)*layers = (T_pre + T_gen*(n-1))*layers  #其中n是生成的token数 
T_pre = max(ctog, gtoc, dtoc, ctod, comp)  #其中c,g,d分别代表cpu,gpu,disk。comp代表计算 
T_gen = max(ctog, gtoc, dtoc, ctod, comp)

上述ctog, gtoc, dtoc, ctod都是latency，对应了各个不同的数据搬移的过程。具体的数据则包括了weights, activations, KV Cache三种，占用的存储空间可以这样估算：

#(1) weights, 即所有transformer层的weights占用的存储空间(单位Bytes, 精度FP16): 
part1 = (8*h1*h1 + 4*h1*h2) * layers   #其中h1是hidden size, h2是2nd MLP层的hidden size 
#(2) activations 
part2 = (2*bls*h1) * layers            #其中bls代表block size, 即effective batch size 
#(3) kv cache 
part3 = (4*bls*(s+n/2)*h1) * layers    #其中s代表输入prompt长度，n代表输出tokens个数 
#总存储空间占用 
total bytes = part1 + part2 + part3

Policy Search

一个policy全称叫offloading policy，包含11个变量：

bls, gbs      # 分别是block size即effective bs, gpu bs 
wg, wc, wd    # w代表weights, g,c,d分别代表gpu,cpu,disk 
hg, hc, hd    # h代表hidden, 也就是activation 
cg, cc, cd    # c代表cache, 即kv cache

其中bls和gbs可以组成一个tuple，这个tuple的组合并不多，即可选的bls, gbs值有限，因此可以先固定(bls, gbs), 然后找到最佳的tensor placement, 这其实就变成了一个线性规划(linear programming)问题，目标是找到某个placement，使得T/bls最小，即处理一个block size的latency最小。