Splitwise: Efficient Generative LLM Inference Using Phase Splitting

https://arxiv.org/pdf/2311.18677

https://yiyibooks.cn/arxiv/2311.18677v2/index.html

论文的核心思想:不管是 Ocra、还是 ExeGPT、还是 SplitWise,论文都提到一些关键的特点,整个推理可以分为2部分,提示阶段(首token计算阶段)和词生成,这2个阶段的负载类型是截然不同的,对算力的要求也不一样,因此,通过拆分这2阶段的计算,未来资源池里面,一部分机器用来跑首 token 计算,一部分机器用来跑词生成,通过这个方式,整个推理的吞吐提升了2倍左右

1. 论文的一些洞察

1.1. 【关键】提示阶段和词生成阶段的token分布

由于编码服务的常见情况只需要在用户键入时在程序中生成接下来的几个单词,因此输出词符的中位数为 13 个标记。另一方面,对话服务几乎呈双峰分布,生成的 Token 中位数为 129 个。洞察1:不同的推理服务可能具有截然不同的token数量分布(提示阶段、词生成阶段)

1.2. 【关键】批处理利用率

batch size = 30,但不代表每时每刻都有 30 个请求在处理,实际情况有很多管道气泡

如果一个请求,在做词生成,那标记为1个活跃的token。上面这个图就是 GPU 活跃 token 数量的一个累计分布函数从这个图可以看出:

  1. 对 chat 来说:GPU 大概有 60% 的时间,在处理 < 20 的 token。
  2. 对于 code 来说:这个情况更差,大概有 30% 的时间,在处理 1 个 token

也就是说,GPU 实际运行过程中,是有非常多管道气泡的。比如 batch size 是 30,但是还是有很多时间,只有1个token在处理

推测:

Q:为什么会出现这个问题?

A:由于请求和 seq 长度存在一定的分布,所以必然会存在很多时刻,很多流水线是处于 “短暂” 的空闲的,这种空闲的时间片,就是管道气泡

Q:能不能通过缩小 Batch size 来降低管道气泡?

A:不能,因为缩小 Batch size 会推高长尾延迟,整体吞吐下降


Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache

Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache

1. 问题

模型的架构:

  1. QKV Linear layer
  2. Multi-Head Self-Attention Mechanism, or the attention module
  3. FNN,Feed-Forward Neural Network

超长文推理的挑战,2k -> 256k ?

有2个:

  1. 不同的层对显存的需求差异很大(注意力层的显存和输入输出是成比例的,但是其他层不是),这就限制了并行的效率(类似木桶效应??)
  2. KV cache 的大小不可预估,显存容量不可规划

这2个问题,最终也极大的限制了超长文推理的效果


SpotServe – Serving Generative Large Language Models on Preemptible Instance

https://mp.weixin.qq.com/s/uK1owRF53FW4WI7nldYadg
https://arxiv.org/pdf/2311.15566.pdf
业界首个,在可抢占式实例上运行的分布式大语言模型(LLM)服务系统
目标:LLM 降本 -> 抢占式 GPU 实例 -> 怎么更好的用这些抢占实例
相比直接使用抢占式GPU实例,SpotServe 可以将推理引擎的长尾延迟降低 2.4x – 9.1x
LLM 特点:

  1. 高计算量
  2. 大内存占用

无论是哪一点,都意味着成本昂贵
把 LLM 运行在随时可抢占的 GPU 实例上
传统的方法,MArk、Cocktail,单实例多GPU卡,运行一个 small DNN 模型,通过请求重定向或者冗余计算来处理抢占,但是这种方式只是和数据并行这种小的DNN模型,不适合LLM
LLM 会同时使用数据并行、模型并行、流水线并行多种技术,单个实例抢占会影响整个多个实例的计算结果
所以需要有更有效的方法
亮点:第一个做推理容错的论文(其他的都是在做训练的容错)
SpotServe 的创新点:

  1. 动态配置并行度
  2. 实例迁移优化(复用模型参数、中间结果,减少迁移后的传输数据量):aws 抢占GPU实例,迁移后冷启动需要2m
  3. 高效利用宽限期(30s,尽量不中断推理)

1. 背景

1.1. Generative LLM Inference

当前类似的系统:

  1. FasterTransformer
  2. Orca
  3. FairSeq
  4. Megatron-LM

单次推理的总耗时,分析


texe(Sin) 是输入序列的解码时间
Sout 是生成的 token 数量,texe(1) 是生成每个 token 的时间
KV cache 技术可以将每个 token 的生成时间优化到接近常量 (i.e., 𝑡𝑒𝑥𝑒 (1) in E.q.(2) and Figure 1a).


Microsecond-scale Preemption for Concurrent GPU-accelerated DNN Inferences

这是到目前为止(2023/10)应该是唯一一篇在认真思考并且实现 GPU 抢占的论文了。
之前很多论文都研究过 GPU 怎么混部,怎么做算力和显存的隔离,但是目标都是提升利用率,但是很难保证长尾延迟。比较适合训练,但不适合推理。
但是这篇论文其实也是有很多限制的:
  1. 不支持 nvidia GPU,只支持 amd
  2. 模型需要经过 tvm 的编译,因为抢占的实现依赖了编译技术
这篇论文主要的工作就2个事情:
  1. GPU 抢占:real-time任务启动的时候,best-effort 任务怎么做到快速的退出。
  2. dynamic kernel padding:怎么在尽可能保证 rt 任务 latency 的情况下,混部 be 的任务
最终效果:
  1. 2% 的延迟损失,换取 7.7x 倍的性能提升

1. 系统架构

0
reef 在 GPU runtime 的基础上扩展了4个核心组件:
  1. Task Queues:一个rt队列和多个be队列,每个队列绑定一个 GPU Stream
  2. Scheduler:一个调度器,2种模式,rt模式,normal模式,如果当前没有rt任务,调度器会切换到normal模式,如果有rt任务,会切换到rt模式。这个地方应该是为了和下面的抢占进行协同
  3. Preemption:抢占模块,实现rt抢占be
  4. DKP:dynamic kernel padding
前面2个没啥东西,重点讲一下抢占是怎么实现的吧,这个比较有意思


深入浅出 tvm – (16) Tensor Expression 详解

TVM addresses these challenges with three key modules.
(1) We introduce a tensor expression languageto build operators and provide program transformation primitives that generate different versions of the program with various optimizations.
tensor expression 是 tvm 引入的一门面向张量计算的函数式编程语言。前面我们知道 topi 就是基于 te 实现了一系列常见的张量算子库
在 tvm 里面,每个算子都可以表达为一个 te 表达式,te 接受张量作为输入,并返回一个新的张量
这个语言本身是非常简单的:
  1. Type:因为处理的输入输出,都是 tensor,所以只有一种数据类型(不像 relay ir,有 Var、Constant、Function、Call 等各种类型)
  2. Op:操作符也很简单,最常见的就2个 placeholder_op,compute_op
te 的代码主要在 include/tvm/te/ src/te/ 这2个目录,全部加起来才1.3w行,所以说是并不复杂
这篇文章,我们从编译器的视角来理解一下 tensor expression
参考资料:

1. te 的设计目标

te 可能应该有其他更多的设计目标
但从我目前了解到的来看,te 最大的作用应该是简化tensor 算子的编写难度,并支持进一步的调度和优化
比如向量加法
1)TVMScript 版本
在没有 te 之前,比如只有 tir 的时候,用户只能用 TVM script 来编写tensor 算子。如下:
@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def main(a: T.handle, b: T.handle):
        # 我们通过 T.handle 进行数据交换,类似于内存指针
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # 通过 handle 创建 Buffer
        A = T.match_buffer(a, (8,), dtype="float32")
        B = T.match_buffer(b, (8,), dtype="float32")
        for i in range(8):
            # block 是针对计算的抽象
            with T.block("B"):
                # 定义一个空间(可并行)block 迭代器,并且将它的值绑定成 i
                vi = T.axis.spatial(8, i)
                B[vi] = A[vi] + 1.0

ir_module = MyModule
2)tensor expression 版本
但是如果用 te 来描述向量加法的话,只需要2行,并且通过script() 打印出来会发现和 TVM script 是一样的
from tvm import te

A = te.placeholder((8,), dtype="float32", name="A")
B = te.compute((8,), lambda *i: A(*i) + 1.0, name="B")

func = te.create_prim_func([A, B])
ir_module_from_te = IRModule({"main": func})
print(ir_module_from_te.script())


使用 ftrace 跟踪内核函数调用栈

搞虚拟化的,经常会遇到很多和内核相关的线上问题。比如最近我们遇到一个 k8s 的 pod 在删除的时候 pvc 卸载不掉的情况。

umount 的时候提示挂载点已经 not mounted 了,但是在 /proc/mounts 里仍然存在。statvfs 看了一下,这个挂载点的device id和父目录的device id 是一样的,说明已经不是挂载点了(因为看 /proc/mounts,这个挂载点的src是在另外一个设备上)

特别注意:如果 /proc/mounts 里,src 和 dst 是同一个设备,那就是 –bind 的方式,这种情况下就很难判断挂载点是否已经真正的卸载掉了

针对这个问题,社区已经有了一个比较 hack 的解决方案:https://github.com/kubernetes/kubernetes/issues/114546,就是忽略 /proc/mounts 残留,允许 pod 删除流程继续往下走

但是根源还是内核的存在脏数据(大概率猜测可能是引用计数泄露导致),还是得从内核层面定位这种问题,而且这种问题,很多时候我们得知道内核代码的执行路径是什么样的。这个时候用 ftrace 来定位就很方便了,比 eBPF 好使,ftrace 能够直接跟踪一个内核函数的调用栈


Pathways – Asynchronous Distributed Dataflow for ML

其他参考:
xx

1. 论文背景、动机、贡献

在探讨现有 AI 模型的局限时,Jeff Dean 曾经说过,今天的人工智能系统总是从头开始学习新问题,我们为数千个单独的任务开发了数千个单独的模型,这种方式需要大量的数据和算力,是非常低效的
理想的发展模式应该是训练1个模型来干多个事情
0
而导致这个现象的根本原因,是因为过去10年的机器学习,总是 算法&硬件&系统 一起演进的,这种 co-evolution 的方式有一定的隐患:系统过于针对当前的任务,而不能很好的适应未来的需求
未来 ML 的需求是什么?
  1. 模型会越来越大:
    1. 多模态:复杂,参数多。纯粹的数据并行已经不行了,开始引入流水线并行、模型并行等新的并行训练方式
    2. 计算稀疏:大模型特有的一个问题
  2. 硬件异构:SPMD 不适用了,因为我们需要在不同的“硬件”上运行不同的代码,需要转向 MPMD
  3. 训推一体:提升硬件利用率
为了实现这个愿景,论文提出了一种叫 [Pathways] 的通用AI架构。Pathways 旨在用一个架构同时处理多项任务,并且拥有快速学习新任务、更好地理解世界的能力


深入浅出 tvm – (15) TVM Operator Inventory (TOPI)

topi 是 tvm 的一个张量算子库,提供了很多常见的算子操作,例如 conv2D,transpose,等等。
tvm 在编译计算图的时候,会先将计算图从 relay ir 翻译成 tir,再将 tir 翻译成目标设备代码(比如llvm,比如c,比如cuda),但是 relay ir 和 tir 之间还有一种中间语言叫 tensor expression,这是一种专门为张量计算而设计的语言。relay ir 的绝大多数部分 op 其实都是通过 tensor expression 这种语言来实现的
An operator is a primitive operation, such as add or conv2d, not defined in the Relay language. Operators are declared in the global operator registry in C++. Many common operators are backed by TVM’s Tensor Operator Inventory.
topi 和 te 的关系:
  1. te( tensor expression ) 是一门函数式编程语言,面向张量计算的
  2. topi 是基于 te 实现的张量算子库
类似于编程语言和标准库之间的关系

1. topi 算子列表

源码目录:src/topi
src/topi 只是把算子通过 TVM_REGISTER_GLOBAL 注册到 tvm 的全局函数中,并没有太多细节
函数的实现,全部是以 inline 函数的形式,实现在 include/tvm/topi 中,由于算子都是通过 tensor expression 来实现的,因此所有算子的实现最终都会调用 tvm::te::compute 来进行计算
最常见的有4类操作:broadcast,elemwise,nn,reduction
broadcast
elemwise
nn
reduction
topi.add
topi.subtract
topi.multiply
topi.divide
topi.floor_divide
topi.mod
topi.floor_mod
topi.maximum
topi.minimum
topi.power
topi.left_shift
topi.logical_and
topi.logical_or
topi.logical_xor
topi.bitwise_and
topi.bitwise_or
topi.bitwise_xor
topi.right_shift
topi.greater
topi.less
topi.equal
topi.not_equal
topi.greater_equal
topi.less_equal
topi.broadcast_to
topi.acos
topi.acosh
topi.asin
topi.asinh
topi.atanh
topi.exp
topi.fast_exp
topi.erf
topi.fast_erf
topi.tan
topi.cos
topi.cosh
topi.sin
topi.sinh
topi.tanh
topi.fast_tanh
topi.atan
topi.sigmoid
topi.sqrt
topi.rsqrt
topi.log
topi.log2
topi.log10
topi.identity
topi.negative
topi.clip
topi.cast
topi.reinterpret
topi.elemwise_sum
topi.sign
topi.full
topi.full_like
topi.logical_not
topi.bitwise_not
topi.nn.relu
topi.nn.leaky_relu
topi.nn.prelu
topi.nn.pad
topi.nn.space_to_batch_nd
topi.nn.batch_to_space_nd
topi.nn.nll_loss
topi.nn.dense
topi.nn.bias_add
topi.nn.dilate
topi.nn.flatten
topi.nn.scale_shift_nchw
topi.nn.scale_shift_nhwc
topi.nn.pool_grad
topi.nn.global_pool
topi.nn.adaptive_pool
topi.nn.adaptive_pool3d
topi.nn.pool1d
topi.nn.pool2d
topi.nn.pool3d
topi.nn.softmax
topi.nn.log_softmax
topi.nn.lrn
topi.nn.binarize_pack
topi.nn.binary_dense
topi.sum
topi.min
topi.max
topi.argmin
topi.argmax
topi.prod
topi.all
topi.any

2. topi 编程示例

由于 topi 实现的所有算子,最终都会注册到 tvm 的全局函数中,并已 relay.op 的形式存在
因此,我们在描述计算图的时候,可以直接调用相关的算子,相当于我们平时写 c 代码的时候,有很多编译器的 builtin 函数一样
比如 topi.add 对应的 relay.op 是 relay.add,其他类推
a1 = relay.var("a1", shape=(1,), dtype="float32")
c1 = relay.const(10, 'float32')
c2 = relay.add(c1, a1)
xx


深入浅出 tvm – (13) Relay Pass 之常量折叠

和算子融合一样,常量折叠是编译领域里最常见的一个优化,简单来说,就是把常量表达式前置计算,在编译阶段就计算好,然后以常量的形式翻译成底层机器码,以提高执行效率,减少计算量
实际上大部分的编译器,常量折叠一般包含2种优化技术:常量折叠和常量传播

1. 基本概念

1.1. 常量折叠

constant folding,常量折叠,编译器优化技术之一,通过对编译时常量或常量表达式进行计算来简化代码。以下面的代码为例:
i = 320 * 200 * 32;
上面的代码中,编译器通常会在编译过程中直接对表达式进行求值,计算出320 * 200 * 32的结果,而不会生成2个乘法指令。
还有一些更复杂(但不清楚tvm是否支持,后面验证下)。比如,在执行一些复杂表达式的计算时,我们可以将表达式内部一些常量运算合并,最终起到简化的效果,如下
优化前(左边)每个表达式的运算量是 8 flop,优化后(右边)的运算量是 2 flop,运算效率极大提升了
0
不过,实际上 tvm 并没有做的那么高级,tvm 的常亮折叠只用来处理一些比较简单的场景

1.2. 常量传播

constant propagation,常量传播,同样也是编译器最常见的优化技术之一,在编译的过程中,对常量依赖做一些基本的推演和提前计算,再使用常量折叠技术来简化代码。如下:
int x = 14;
int y = 7 - x / 2;
return y * (28 / x + 2);

//常量传播

int x = 14;
int y = 7 - 14 / 2;
return y * (28 / 14 + 2);

//常量折叠

int x = 14;
int y = 0;
return 0;