当前位置: 首页 > article >正文

w~大模型~合集27

我自己的原文哦~    https://blog.51cto.com/whaosoft/12898045

#vLLM~2

作者尽量少涉及对源码本身的解读,把源码中的信息总结出来,配合图例做整体介绍。 

大家好,这段时间精读了一下vLLM源码实现,打算开个系列来介绍它的源码,也把它当作我的总结和学习笔记。

整个vLLM代码读下来,给我最深的感觉就是:代码呈现上非常干净历练,但是逻辑比较复杂,环环嵌套,毕竟它是一个耦合了工程调度和模型架构改进的巨大工程。

所以在源码解读的第一篇,我想先写一下对整个代码架构的介绍。在本篇中,我特意少涉及对源码本身的解读,而是把源码中的信息总结出来,配合图例先做整体介绍。如果你不想阅读源码细节,但又想对vLLM代码有整体把握,方便后续能知道从哪里查bug的话,这篇文章或许可以帮到你。如果你后续想更深入阅读源码的话,这篇文章可以作为一个引子,后续的细节解读都将在本文的基础上扩展开。

一、调用vLLM的两种方式

根据vLLM的官方文档,它向用户提供了两种调用它的方法,分别是:

  • Offline Batched Inference同步,离线批处理)
  • API Server For Online Serving异步,在线推理服务),在这下面又提供了2种支持的API类型:
  • OpenAI-Compatible API Server(官方推荐):兼容了OpenAI请求格式的server,包括OpenAI Completions API和OpenAI Chat API。
  • Simple Demo API Server(测试开发用,官方不推荐,相关脚本也不再维护)

在代码实现上,vLLM首先实现了一个推理内核引擎(LLMEngine),在此基础上封装了上述两种调用方法。在本系列的讲解中,我们会先以“offline bacthed inference”作为入口,详细解说内核引擎LLMEngine的各块细节。在此基础上我们再来看“online serving”的运作流程。

现在,让我们来看这两种调用方法的具体例子。

1.1 Offline Batched Inference

from vllm import LLM, SamplingParams

# ===========================================================================
# batch prompts
# ===========================================================================
prompts = ["Hello, my name is",
           "The president of the United States is",
           "The capital of France is",
           "The future of AI is",]

# ===========================================================================
# 采样参数
# ===========================================================================
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

# ===========================================================================
# 初始化vLLM offline batched inference实例,并加载指定模型
# ===========================================================================
llm = LLM(model="facebook/opt-125m")

# ===========================================================================
# 推理
# ===========================================================================
outputs = llm.generate(prompts, sampling_params)

# ===========================================================================
# 对每一条prompt,打印其推理结果
# ===========================================================================
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

在传统离线批处理中,我们每次给模型发送推理请求时,都要:

  • 等一个batch的数据齐全后,一起发送
  • 整个batch的数据一起做推理
  • 等一个batch的数据全部推理完毕后,一起返回推理结果

这种“团体间等成员到齐,再一起行动”的行为,就被称为“同步”。

在vLLM中,当我们使用离线批处理模式时,表面上是在做“同步”推理,也即batch_size是静态固定的。但推理内核引擎(LLMEngine)在实际运作时,batch_size是可以动态变更的:在每一个推理阶段(prefill算1个推理阶段,每个decode各算1个推理阶段)处理的batch size可以根据当下显存的实际使用情况而变动。

举个例子来说:

  • 给定一个很大的batch,此时尽管vLLM采用了PagedAttention这样的显存优化技术,我们的gpu依然无法同时处理这么大的batch。
  • 所以batch中的每一条数据,会被先放到一个waiting队列中。vLLM会用自己的调度策略从waiting队列中依次取数,加入running队列中,直到它认为取出的这些数据将会打满它为1个推理阶段分配好的显存。此时waiting队列中可能还会剩一些数据。
  • 在每1个推理阶段,vLLM对running队列中的数据做推理。如果这1个推理阶段执行完毕后,有的数据已经完成了生成(比如正常遇到​​<eos>​​了),就将这些完成的数据从running队列中移开,并释放它占据的物理块显存。
  • 这时,waiting队列中的数据就可以继续append进running队列中,做下1个阶段的推理。
  • 因此在每1个推理阶段,vLLM处理的batch size可能会动态变更。
  • 将LLMEngine包装成离线批处理形式后,所有的数据必须等到一起做完推理才能返给我们。所以从体感上,我们可能很难感知到内核引擎的“动态”逻辑。

以上是一个浅显粗暴的例子,目的是帮助大家理解“在vLLM中,即使是同步形式的离线批处理,其背后的内核引擎也是按动态batch的形式来实现的”,实际的调度策略(Scheduler)要更加复杂,我们将在后续的解读中来具体看它。

也正是因为LLMEngine这种“动态处理”的特性,才使得它同时也能成为异步在线服务的内核引擎:当一条条请求发来时,它们都先进入LLMEngine调度器(Scheduler)的waiting队列中(实际并不是直接进入waiting队列中的,而是在传给LLMEngine前先进入asyncio.Queue()中,然后再由LLMEngine调度进waiting队列中的,这些细节我们也放在后面说,这里不影响理解就行)。此时模型正常执行它的1个推理阶段,调度器也正常处理新来的请求。当模型准备执行下1个推理阶段时,调度器再根据设定的策略,决定哪些数据可以进入running队列进行推理。由于在线服务是异步的,先推理完成的数据就可以先发给客户端了(如果采用流式传输,也可以生成多少先发多少)。

在这个过程中,vLLM通过PagedAttention技术和“先来先服务(FCFS),后来先抢占,gpu不够就先swap到cpu上”的调度策略,在1个推理阶段处理尽可能多的请求,解决高并发场景下的推理吞吐问题。这就是整个vLLM运作的核心思想。(对这行黑体字里的术语有疑惑的朋友,建议先看vLLM原理篇讲解)

1.2 API Server For Online Serving

# ===========================================================================
# Server:起服务
# ===========================================================================
$ python -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-2-7b-hf

# ===========================================================================
# Client:发请求(OpenAI API)
# ===========================================================================
$ curl http://localhost:8000/v1/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "meta-llama/Llama-2-7b-hf",
        "prompt": "San Francisco is a",
        "max_tokens": 7,
        "temperature": 0
    }'

vLLM在实现在线服务时,采用uvicorn部署fastapi app实例,以此实现异步的请求处理。而核心处理逻辑封装在​​AsyncLLMEngine​​类中(它继承自LLMEngine)。所以,只要我们搞懂了LLMEngine,对vLLM的这两种调用方式就能举一反三了。

1.3 总结

vLLM的两种调用方式与内核引擎LLMEngine的关系如下(图片来自vLLM团队2023 first meetup PPT):

图中左侧是用户使用界面,罗列了上述所说的两种调用方式(注意,如前文所说,做demo用的api server官方已经不再维护了,openai_api_server才是官方推荐的使用方式,user custom server目前还没有实现)。右侧则是开发者界面,不难发现LLMEngine是vLLM的核心逻辑。

我们来看开发者界面下的几个函数,先来看LLMEngine:

  • ​add_request()​​:该方法将每一个请求包装成vLLM能处理的数据类型(SequenceGroup,后面我们会详细解释),并将其加入调度器(Scheduler)的waiting队列中。在LLMEngine中,这个函数是按照“同步”的方式设计的,也就是它被设计为“遍历batch中的每条数据,然后做相应处理”。所以这个函数本身只适合批处理场景。在异步的online serving中将会把它重写成异步的形式。
  • ​abort_request​​:在推理过程中,并不是所有的请求都能有返回结果。比如客户端断开连接时,这个请求的推理就可以终止了(abort),这个函数就被用来做这个操作。
  • ​step()​​:负责执行1次推理过程(1个prefill算1个次推理,每个decode各算1次推理)。在这个函数中,vLLM的调度器会决定要送那些数据去执行本次推理,并负责给这些数据分配好物理块(这些信息都被作为metadata放在要送给模型做推理的数据中)。模型会根据这些信息,采用PagedAttention方法,实际完成推理。

​AsyncLLMEngine​​​下的函数也是同理类推,这里不赘述了。从上面的解读你可能发现了,其实只要掌握了​​add_request()​​​和​​step()​​这两个函数,就等于掌握LLMEngine的全部思想了!于是你兴奋地打开这两个函数,发现它们的实现代码只有十几行,你突然感觉自己好像是去项羽那吃席的刘邦,因为你渐渐发现:背后有万行代码逻辑正在等你😊所以说,vLLM真得是一个巨大的工程,它耦合了调度工程和模型本身的改造,代码中的嵌套非常复杂。所以在代码解读的第一篇,我们得先理清整个代码架构,然后再逐一攻克细节。   

二、vLLM代码整体架构

LLMEngine可以具体分成两个部分:

2.1 Centralized Controller

Centralized Controller,也就是前文我们所说的调度器。它和LLMEngine所在的进程是同一个,且两者都是在CPU上的。

  • 调度器的主要作用就是,在每1个推理阶段,决定要把哪些数据送给模型做推理,同时负责给这些模型分配KV Cache物理块。但要注意,它只是分配了物理块的id,而不是物理块本身。物理块的实际分配是模型在推理过程中根据物理块id来操作的,也就是CacheEngine做的事情。
  • 调度器下维护着BlockSpaceManager。它负责管理BlockAllocator(实际参与分配物理块的类)。BlockAllocator又分成gpu和cpu两种类型,分别管理这两类设备上的物理块。你可能会问,cpu上的物理块是什么呢?你还记得调度器有一个swap策略吗?当gpu上显存不足时,它会把后来的请求抢占,并将其相关的KV cache物理块全部都先swap(置换、卸载)在cpu上,等后续gpu显存充足时,再把它们加载回gpu上继续做相关请求的推理。所以在cpu上我们也需要一个管控物理块的BlockAllocator。实际代码实现时,Block相关的部分可不止这两个class,还有一些更复杂的逻辑细节。这个我们放在本系列后面的文章中讲解

2.2 Distributed Workers

Distributed Workers,也就是分布式系统,你可以将每个worker理解成一块gpu。它的作用是将我们要使用的模型load到各块卡上(目前对单卡装不下的模型,vLLM支持tp/pp推理),然后对Controller传来的数据做1次推理,返回相关结果。我们来细看下这块:

  • Distributed Workers:图中绘制为Distributed Workers这个绿色块,其实按vLLM的源码内容,写成Executor会更合适一些。它就是所有Workers的管控中心,它指定了用什么方法管控这些Workers,负责分布式环境的初始化,目前支持的方法有:
  • cpu_executor:(较少用),使用cpu做推理时可考虑
  • gpu_executor:单卡(world_size = 1)的情况下可用
  • ray_gpu_executor:使用ray这个分布式计算框架实现的executor,适用于多卡环境
  • Worker:在硬件上,它指gpu;在代码上,它指的是Worker实例(每个gpu上的进程维护自己的Worker实例)。在每个Worker实例中又管控着如下两个重要实例:
  • CacheEngine:负责管控gpu/cpu上的KV cache物理块(调度器的block manager只负责物理块id的分配,CacheEngine则是根据这个id分配结果实打实地在管理物理块中的数据)
  • Worker.model:根据vLLM代码,这里写成model_runner会更合适一些。它负责加载模型,并执行推理。调用PagedAttention的相关逻辑,就维护这个实例关联的代码下。

三、加载模型与预分配显存

现在你已经从代码层面知道vLLM的整体架构了,你是不是迫不及待想看看:当一条请求过来时,整个vLLM是怎么运作的呢?现在,我们就来解析这个流程。

在vLLM正式开始处理1条请求(也就是LLMEngine的调度器正式开始运作时),它需要做两件和初始化相关的事:

  • 加载模型
  • 预分配显存

我们分别来看这两个步骤。

3.1 加载模型

这里在做的事很直观:把你的base model加载到worker上。如果你是online加载的,vLLM默认使用HuggingFace,你也可以在环境变量中把相关配置改成ModelScope。

3.2 预分配显存

欸这个就非常有意思了。在模型部署的初始化阶段(推理正式开始前),vLLM会通过模拟实验的方式,来决定gpu/cpu上到底有多少个KV cache物理块可以分配给后续的请求们做推理。vLLM管这个步骤叫​profile_num_available_blocks​​。我们来看看这个模拟实验是怎么做的:

(1)杜撰假数据

首先,用户在初始化LLMEngine引擎时,会提供两个重要参数:

  • ​max_num_seqs​​:在1个推理阶段中,LLMEngine最多能处理的seq数量(1条seq就是指我们待推理的1条数据)。默认是256
  • ​max_num_batched_tokens​​:在1个推理阶段中,LLMEngine最多能处理的token数量。默认是2048

根据这两个参数,我们可以假设在模型推理中,平均一个seq要处理​​max_num_batched_tokens // max_num_seqs​​​个token,余数部分我们默认放在第一个seq中。例如,假设​​max_num_batched_tokens=10,max_num_seqs = 3​​,那么我们就能杜撰出3条seq,每个seq的长度分别为4,3,3

(2)用假数据模拟一次前向推理

我们现在想知道在1次推理过程中,可以分配多少的显存给KV cache。我们可以使用如下公式计算:

分配给KV cache显存 = gpu总显存 - 不使用KV cache情况下做1次FWD时的显存占用(包括模型本身和FWD过程中的中间数据)

对于“不使用KV cache做1次FWD时的显存占用”,我们就可以用杜撰出来的假数据模拟一次FWD来计算得出。在前向推理之后,我们把gpu上的缓存清一次,让它不要影响后续模型的正常推理。

(3)计算可分配的KV cache物理块总数

从(2)的模拟实验中,我们已经预估了一块卡上“分配给KV Cache的总显存”。现在,我们可以来计算总的物理块数量了。

我们易知:总物理块数量 = 分配给KV Cache的显存大小/ 物理块大小,其中“大小”的单位是bytes。

物理块大小(block_size)也是可以由用户自定义的,vLLM推荐的默认值是block_size = 16。

由大模型中KV值的定义,我们易知:​​K_cache_block_size = block_size * num_heads * head_size * num_layers * dtype_size​​其中dtype_size表示精度对应的大小,例如fp16就是2,fp32就是4

同理可知:​​V_cache_block_size = K_cache_block_size​

则最终一个物理块的大小为:

​cache_block_size = block_size * num_heads * head_size * num_layers * dtype_size * 2​

知道了物理块的大小,我们就能求出物理块的总数了。

CPU上物理块总数也是同理,但与GPU不同的是,它不需要做模拟实验。CPU上可用的内存总数是用户通过参数传进来的(默认是4G)。也就是我们认为只能在这4G的空间上做swap。将上面公式中“分配给KV Cache的显存大小”替换成4G,就能得到CPU上物理块的数量。

(4)将预分配的KV Cache加载到gpu上

当我们确定好KV Cache block的大小后,我们就可以创建empty tensor,将其先放置到gpu上,实现显存的预分配。以后这块显存就是专门用来做KV Cache的了。也正是因为这种预分配,你可能会发现在vLLM初始化后,显存的占用比你预想地要多(高过模型大小),这就是预分配起的作用。相关代码如下(帮助大家更好看一下KV cache tensor的shape):

def _allocate_kv_cache(
        self,
        num_blocks: int,
        device: str,
    ) -> List[torch.Tensor]:
        """Allocates KV cache on the specified device."""
        kv_cache_shape = self.attn_backend.get_kv_cache_shape(
            num_blocks, self.block_size, self.num_heads, self.head_size)
        pin_memory = is_pin_memory_available() if device == "cpu" else False
        kv_cache: List[torch.Tensor] = []
        # =======================================================================
        # kv_cache_shape: (2, num_blocks, block_size * num_kv_heads * head_size)
        # =======================================================================
        for _ in range(self.num_layers):
            kv_cache.append(
                torch.empty(kv_cache_shape,
                            dtype=self.dtype,
                            pin_memory=pin_memory,
                            device=device))
        return kv_cache

由于这篇文章的主要目的是帮大家了解vLLM代码框架,所以关于预分配的详细代码(注释版)的讲解,我们放在后面的系列中说。这篇文章会尽量少放代码。

整个预分配的过程,其实也是在提醒我们:当你发现vLLM推理吞吐量可能不及预期,或者出现难以解释的bug时,可以先查查输出日志中pending(waiting)/running/swapped的序列数量,以及此时KV Cache部分的显存利用程度,尝试分析下这些默认的预分配设置是不是很好契合你的推理场景,如果不行,可以先尝试调整这些参数进行解决。

四、Scheduler调度

好,目前为止,vLLM所有初始化的工作都完成了,我们现在可以来处理一条请求了。这就是我们调度器发挥作用的时候了,整个调度过程如下:

具体的内容我们在前文说了很多了。这里只提一点:你会发现这出现了叫swapped的队列,这是前文没有提过的。

如果你读过vLLM的原理篇,你可能记得vLLM的调度策略中有一项叫做:后来先抢占(Preemption)。它是指在准备执行当前这1个推理阶段时,如果gpu上没有足够的资源对running队列中的全部数据完成下1次推理,我们就取出running队列中最后来的数据,将它的KV Cache swapped到CPU上,同时将这个数据从running移到swapped中。我们重复执行这个步骤,直到当前gpu上有足够的KV Cache空间留给剩在running中的全部数据为止。

而存放在Swapped队列中的数据,也会在后续gpu上有足够空间时,被重新加入running计算。

详细的调度策略会更精细复杂,我们放在对Scheduler单独的代码解析中来说。

关于vLLM整体代码架构,我们就介绍到这了。对于这个精妙而复杂的系统,我们还有很多细节可以探索。在本系列后续的文章中,我们将来仔细研究这些细节。

参考

1、https://arxiv.org/pdf/2309.06180.pdf

2、https://github.com/vllm-project/vllm

3、https://docs.google.com/presentation/d/1QL-XPFXiFpDBh86DbEegFXBXFXjix4v032GhShbKf3s/edit?pli=1#slide=id.g24ad94a0065_0_162

#vLLM~3

图解大模型计算加速系列:vLLM源码解析2,调度器策略(Scheduler)

在本文中,作者从vLLM批处理的入口函数开始,介绍了其推理内核LLMEngine的两个重要函数add_request()和step()。当LLMEngine开始执行1次调度时(step),调度器策略(Scheduler)会根据实际gpu上KV Cache block的使用情况等要素,来选择要送哪些seq_group去做新一轮推理。

vLLM源码解读第二期更新了,本期我们一起来解读vLLM的调度器策略。

实话说,这真得是我写过最难的源码解读了。由于vLLM代码本身的复杂性,逻辑上的嵌套性,使得我在读源码时,先接收到的是碎片化的东西,当代码一长、细节一多时,就很难把碎片化的东西拼成全貌。所以在本系列对vLLM的介绍中,不管是哪一块,都会按照 “宏观(图解) -> 细节(配合源码) ”的方式,先理清vLLM在这里想做什么事,为什么要这么做,然后再一起来看各小块的代码实现。

前期提要与本期导览

在上一篇关于vLLM代码整体架构的文章中,我们提到过无论是“离线批处理(同步)”还是“在线流式服务(异步)”,它们都采用了同一个推理内核引擎LLMEngine,其整体架构如下:

其中:

  • 在每1个推理阶段中,调度器(Scheduler)决定哪些请求可以参与推理,并为这些请求做好逻辑块->物理块的映射。
  • 在每1个推理阶段中,分布式执行者(图中Distributed Workers部分,根据代码,我们将其命名为model_executor会更加合适)接收调度器传来的这些请求,分发到各个worker上去做推理。Worker中的CacheEngine负责实际管理KV Cache;Worker中的model负责加载模型、实行推理,PagedAttention相关的实现和调用就在model下。

这里,每1个推理阶段的定义是:prefill算1个推理阶段,每个decode各算1个推理阶段。在本文中,我们统一用step来表示“1个推理阶段”。

  • 在本文中,我们会详细解读调度器(Scheduler)全部细节;
  • 在下一篇文章中,我们会详细解读块管理(blockmanager)的全部细节,并以parallel sampling,beam search和prefix caching为例,将上图左半部分全部串一遍
  • 在后续文章中,我们会来解读上图右半部分细节(还没来得及拆逻辑,暂时不知道会写几篇)

由于块管理者和调度器在代码上逻辑层层嵌套,所以为了不影响大家对调度器的理解,涉及到块管理者的部分,本文也会给出尽量简明清晰的说明。

一、入口函数

在源码架构篇中我们提过,本系列的介绍思路是:以“离线批处理”作为入口,详细解说内核引擎LLMEngine的各块细节。在此基础上我们再来看“在线流式服务”的运作流程。所以现在,我们先来回顾下离线批处理的调用方式:

from vllm import LLM, SamplingParams

# ===========================================================================
# batch prompts
# ===========================================================================
prompts = ["Hello, my name is",
           "The president of the United States is",
           "The capital of France is",
           "The future of AI is",]

# ===========================================================================
# 采样参数
# ===========================================================================
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

# ===========================================================================
# 初始化vLLM offline batched inference实例,并加载指定模型
# ===========================================================================
llm = LLM(model="facebook/opt-125m")

# ===========================================================================
# 推理
# ===========================================================================
outputs = llm.generate(prompts, sampling_params)

# ===========================================================================
# 对每一条prompt,打印其推理结果
# ===========================================================================
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

有两点需要注意:

  • ​llm = LLM(model="facebook/opt-125m")​​:实例化了一个离线批处理的vLLM对象。其本质是实例化了一个内核引擎LLMEngine对象。在执行这个步骤时,LLMEngine会执行一次模拟实验(profiling),来判断需要在gpu上预留多少的显存空间给KV Cache block(模拟实验的流程参见源码篇1的3.2节,TODO,大家可以对照着来读源码,本文不再涉及这块源码细节)。
  • 推理入口在第24行​​outputs = llm.generate(prompts, sampling_params)​​​。现在我们进入LLM类下,来看这个​​generate​​函数,代码如下:
# vllm/entrypoints/llm.py
class LLM:
    """An LLM for generating texts from given prompts and sampling parameters.
       ...
    """

    def __init__(
        self,
        model: str,
        tokenizer: Optional[str] = None,
        tokenizer_mode: str = "auto",
        trust_remote_code: bool = False,
        tensor_parallel_size: int = 1,
        dtype: str = "auto",
        quantization: Optional[str] = None,
        revision: Optional[str] = None,
        tokenizer_revision: Optional[str] = None,
        seed: int = 0,
        gpu_memory_utilization: float = 0.9,
        swap_space: int = 4,
        enforce_eager: bool = False,
        max_context_len_to_capture: int = 8192,
        disable_custom_all_reduce: bool = True,
        **kwargs,
    ) -> None:
        ...
        # ==============================================================================
        # 使用配置好的engine参数,初始化LLMEngine实例
        # ==============================================================================
        self.llm_engine = LLMEngine.from_engine_args(
            engine_args, usage_context=UsageContext.LLM_CLASS)
        # ==============================================================================
        # 用于全局唯一的request_id,
        # 在vLLM中内核引擎的处理中,1个prompt视为1个request,分配全局唯一的request_id
        # ==============================================================================
        self.request_counter = Counter()
        
        ...

    def generate(
        self,
        prompts: Optional[Union[str, List[str]]] = None, 
        sampling_params: Optional[SamplingParams] = None,
        prompt_token_ids: Optional[List[List[int]]] = None, 
        use_tqdm: bool = True,
        lora_request: Optional[LoRARequest] = None,
        multi_modal_data: Optional[MultiModalData] = None,
    ) -> List[RequestOutput]:
        """Generates the completions for the input prompts.

        NOTE: This class automatically batches the given prompts, considering
        the memory constraint. For the best performance, put all of your prompts
        into a single list and pass it to this method.

        Args:
            prompts: prompts可以是str,也可以是list[str]
            sampling_params: 采样超参,例如温度、top_k等;如果为None则使用vLLM默认的参数
            prompt_token_ids: prompt对应的token_id,如果没有提供的话,vllm会调用tokenizer进行                               转换
            use_tqdm: 是否要展示process bar
            lora_request: 如果想请求特定的lora_adapter,可以将它的path等信息包装在该请求中,
                          但vLLM建议尽量不要使用这种方式,因为私有的lora adapter可能会带来一些
                          安全性的问题        
            multi_modal_data: 多模态相关的数据

        Returns:
            A list of `RequestOutput` objects containing the generated
            completions in the same order as the input prompts.
        """
        if prompts is None and prompt_token_ids is None:
            raise ValueError("Either prompts or prompt_token_ids must be "
                             "provided.")

        if isinstance(prompts, str):
            # Convert a single prompt to a list.
            prompts = [prompts]
        if (prompts is not None and prompt_token_ids is not None
                and len(prompts) != len(prompt_token_ids)):
            raise ValueError("The lengths of prompts and prompt_token_ids "
                             "must be the same.")
        
        if sampling_params is None:
            # Use default sampling params.
            sampling_params = SamplingParams()

        if multi_modal_data:
            multi_modal_data.data = multi_modal_data.data.to(torch.float16)

        # ============================================================================
        # 将request添加到engine中
        # 在vLLM内核运算逻辑中,1个prompt算1个request,需要有1个全局唯一的request_id
        # ============================================================================
        num_requests = len(prompts) if prompts is not None else len(
            prompt_token_ids)
        for i in range(num_requests):
            prompt = prompts[i] if prompts is not None else None
            token_ids = None if prompt_token_ids is None else prompt_token_ids[
                i]
            # =======================================================================
            # 将每个prompt添加进LLMEngine中,_add_request具体做了以下几件事:
            # - 将每个prompt处理成特定的输入类型(SequenceGroup实例,后文会细说)
            # - 将每个prompt加入Scheduler的waiting队列,等待处理
            # =======================================================================
            self._add_request(
                prompt,
                sampling_params,
                token_ids,
                lora_request=lora_request,
                # Get ith image while maintaining the batch dim.
                multi_modal_data=MultiModalData(
                    type=multi_modal_data.type,
                    data=multi_modal_data.data[i].unsqueeze(0))
                if multi_modal_data else None,
            )
        
        # ============================================================================
        # 把这个batch的所有prompt都添加完后,执行推理,详情参见_run_engine
        # ============================================================================
        return self._run_engine(use_tqdm)

    def _add_request(
        self,
        prompt: Optional[str],
        sampling_params: SamplingParams,
        prompt_token_ids: Optional[List[int]],
        lora_request: Optional[LoRARequest] = None,
        multi_modal_data: Optional[MultiModalData] = None,
    ) -> None:
        # 每个prompt赋1个request_id
        request_id = str(next(self.request_counter))
        self.llm_engine.add_request(request_id,
                                    prompt,
                                    sampling_params,
                                    prompt_token_ids,
                                    lora_request=lora_request,
                                    multi_modal_data=multi_modal_data)

    def _run_engine(self, use_tqdm: bool) -> List[RequestOutput]:
        # Initialize tqdm.
        if use_tqdm:
            num_requests = self.llm_engine.get_num_unfinished_requests()
            pbar = tqdm(total=num_requests,
                        desc="Processed prompts",
                        dynamic_ncols=True)
        
        # ===========================================================================
        # 如果当前调度器中还有没完成推理的请求(调度器中waiting/running/swapped任一队列非空)
        # ===========================================================================
        outputs: List[RequestOutput] = []
        while self.llm_engine.has_unfinished_requests():
            # =========================================================================
            # 执行1次推理调度(step),决定哪些请求的数据可以参与到这次推理中
            # =========================================================================
            step_outputs = self.llm_engine.step()
            for output in step_outputs:
                # =====================================================================
                # 如果本step后,有请求已经完成了推理,就将推理结果装进outputs中
                # =====================================================================
                if output.finished:
                    outputs.append(output)
                    if use_tqdm:
                        pbar.update(1)
        if use_tqdm:
            pbar.close()
        # Sort the outputs by request ID.
        # This is necessary because some requests may be finished earlier than
        # its previous requests.
        outputs = sorted(outputs, key=lambda x: int(x.request_id))
        return outputs

总结来说,当我们调用​​outputs = llm.generate(prompts, sampling_params)​​时,它实际做了两件事情

  • ​_add_request​​:将输入数据传给LLMEngine,它具体做了如下事情:
  • 把每1个prompt包装成一个SequenceGroup对象。从客户端角度看,1个请求可能包含多个prompts,例如离线批处理场景下你可以将1个batch理解成1个请求;但是从LLMEngine的角度看,1个prompt是1个请求,所以它会对输入数据进行预处理。在后文对SequenceGroup的讲解中,我们会来看vLLM这样做的意义。
  • 把包装成SequenceGroup对象的数据加入调度器(Scheduler)的waiting队列,等待处理。这一块相关的细节,我们放在后文说。
  • ​_run_engine​​:执行推理。只要调度器的waiting/running/swapped队列非空,我们就认为此时这批batch还没有做完推理,这时我们就会调用LLMEngine的​​step()​​函数,来完成1次调度以决定要送哪些数据去做推理。

所以,想要知道调度器的运作流程,我们只要从LLMEngineadd_request()step()两个函数入手就好了。不过在正式进入这两个函数的讲解之前,我们先来看和输入数据一个问题:为什么要把每个prompt都包装成一个SequenceGroup实例?SequenceGroup又长什么样呢?

二、SequenceGroup

2.1 原生输入

在一般的推理场景中,我们通常给模型传1个prompt及相关的采样参数,让模型来做推理。此时你的输入可能长下面这样:

("To be or not to be,",
SamplingParams(temperature=0.8, top_k=5, presence_penalty=0.2)),

但在其余的场景中,模型decoding的策略可能更加复杂,例如:

  • Parallel Sampling:你传给模型1个prompt,希望模型基于这个这个prompt,给出n种不同的output
  • Beam Search:你传给模型1个prompt,在采用Beam Search时,每个推理阶段你都会产出top k个output,其中k被称为Beam width(束宽)。

这些情况下,你传给模型的输入可能长下面这样:

# Parallel Sampling
("What is the meaning of life?",
SamplingParams(n=2, temperature=0.8, top_p=0.95, frequency_penalty=0.1))

# Beam Search (best_of = 束宽)
("It is only with the heart that one can see rightly",
SamplingParams(n=3, best_of=3, use_beam_search=True, temperature=0.0)),

【备注:SamplingParams遵从OpenAI API范式,对其中各种参数的解释可参见OpenAI官方文档】

总结来说,可能出现"1个prompt -> 多个outputs"的情况。那是否能设计一种办法,对1个prompt下所有的outputs进行集中管理,来方便vLLM更好做推理呢?

2.2 SequenceGroup的作用

  • "1个prompt -> 多个outputs"这样的结构组成一个SequenceGroup实例。
  • 其中每组"prompt -> output"组成一个序列(seq,属于Sequence实例),每个seq下有若干状态(status)属性,包括
  • ​FINISHED_STOPPED​​:正常执行完毕,例如碰到符号,该seq的推理正常结束了
  • ​FINISHED_LENGTH_CAPPED​​:因为seq的长度达到最大长度限制,而结束推理
  • ​FINISHED_ABORTED​​:因不正常状态,而被终止的推理。例如客户端断开连接,则服务器会终止相关seq的推理
  • ​FINISHED_IGNORED​​:因prompt过长而被终止执行的推理。本质上也是受到长度限制
  • ​WAITING​​:正在waiting队列中。waiting队列中的序列都没有做过prefill。
  • ​RUNNING​​:正在running队列中,即已经开始做推理。
  • ​SWAPPED​​:正在swapped队列中,表示此时gpu资源不足,相关的seq_group被抢占,导致其暂停推理,相关的KV block被置换到cpu上(swap out),等待gpu资源充足时再置换回来重新计算(swap in)。
  • 若干和Finish相关的状态,表示该seq推理已经结束,具体包括:
  • 在vLLM中有一个重要假设:一个seq_group中的所有seq共享1个prompt。

我们来通过一个具体的例子,更好感受一下SequenceGroup的作用:

  • 在推理开始之前,这个seq_group下只有1条seq,它就是prompt,状态为waiting。
  • 在第1个推理阶段,调度器选中了这个seq_group,由于它的采样参数中n = 4,所以在做完prefill之后,它会生成4个seq,它们的状态都是running。
  • 在若干个推理阶段后,gpu上的资源不够了,这个seq_group不幸被调度器抢占(preemption),它相关的KV block也被swap out到cpu上。此时所有seq的状态变为swapped。这里要注意,当一个seq_group被抢占时,对它的处理有两种方式:
  • ​Swap​​:如果该seq_group下的seq数量 > 1,此时会采取swap策略,即把seq_group下【所有】seq的KV block从gpu上卸载到cpu上。(seq数量比较多,直接把算出的KV block抛弃,比较可惜)
  • ​Recomputation​​:如果该seq_group下的seq数量 = 1,此时会采取recomputation策略,即把该seq_group相关的物理块都释放掉,然后将它重新放回waiting队列中。等下次它被选中推理时,就是从prefill阶段开始重新推理了,因此被称为“重计算”。(seq数量少,重新计算KV block的成本不高)

【注意,并不是每个seq_group都会经历抢占,具体要看调度器策略和gpu资源使用情况】

  • 又过了若干个推理阶段,gpu上的资源又充足了,此时执行swap in操作,将卸载到cpu上的KV block重新读到gpu上,继续对该seq_group做推理,此时seq的状态又变为running。
  • 又过了若干个推理阶段,该seq_group中有1个seq已经推理完成了,它的状态就被标记为finish,此后这条已经完成的seq将不参与调度。
  • 又过了若干个推理阶段,这个seq_group下所有的seq都已经完成推理了,这样就可以把它作为最终output返回了。

相信通过这个例子,我们已经能更好理解为什么vLLM要把1个prompt包装成SequenceGroup实例了。接下来我们就来看SequenceGroup实例的具体结构。

2.3 SequenceGroup的结构

SequenceGroup相关的脚本在​​vllm/sequence.py​​中,下图给出了SequenceGroup的结构图解(仅列出重要的属性和方法):

(1)结构总述

SequenceGroup:

  • ​self.seqs_dict​​:{seq_id: seq},其中每个seq是一个Sequence对象。正如我们前文介绍的那样,一个seq_group下包含若干seqs
  • ​self.sampling_params​​:采样参数
  • ​self.metrics​​:记录该seq_group相关的指标,例如该seq_group是什么时候被加入LLMEngine的(arrival_time),该seq_group第一次被调度器选中调度是什么时候等等。调度器在选择时,会参考seq_groups们的这些指标来做决策。
  • ​get_max_num_running_steps​​:该seq_group在剩余生命周期内并行running的最大seq数量。“剩余生命周期”指从此刻一直到seq_group中所有的seq都做完推理。举个例子来说,我们看2.2节配图中倒数第3个时刻,此时这个seq_group内所有的seq都还没结束推理,所以若调用这个方法,则返回值为4;再看倒数第2个时刻,此时有1个seq已经完成了推理,所以若调用这个方法,则返回值为3。在后续调度策略代码中,我们将经常看到这个方法被调用,目的是用于估计若当前对一个seq_group做推理,它将消耗多少gpu资源。

我们来详细看下​​get_max_num_running_steps​​代码实现(一切尽在注释中):

def get_max_num_running_seqs(self) -> int:
        """The maximum number of sequences running in parallel in the remaining
        lifetime of the request.
        返回请求在其剩余生命周期中并行运行的最大序列数。
        """
        # ============================================================================
        # 若采用beam search,每1个推理阶段都是best_of(束宽)个seq在running
        # ============================================================================
        if self.sampling_params.use_beam_search:
            return self.sampling_params.best_of
        # ============================================================================
        # 如果不采用beam search
        # ============================================================================
        else:
            # =========================================================================
            # 此时best_of默认和n一致,即表示我们希望1个prompt产出n个outputs。因此理论上,这个
            # seq_group下会维护best_of个seq(这就是self.num_seqs()的返回值)。
            # 如果出现best_of > self.num_seqs()的情况,说明该seq_group刚从waiting变成running
            # 准备做推理(参考2.2节配图中左侧第1个时刻),此时对于这个seq_group来说,
            # 其剩余生命周期并行运行的最大seq数量为best_of
            # =========================================================================
            if self.sampling_params.best_of > self.num_seqs():
                # At prompt stage, the sequence group is not yet filled up
                # and only have one sequence running. However, in the
                # generation stage, we will have `best_of` sequences running.
                return self.sampling_params.best_of
            
            # =========================================================================
            # 其余时刻(例如2.2节配图中非左侧第1个时刻的所有时刻)下,我们就返回这个seq_group中
            # 未完成推理的seq数量。根据2.2节介绍,我们知道一个seq的完成状态有四种:
            #   SequenceStatus.FINISHED_STOPPED,
            #   SequenceStatus.FINISHED_LENGTH_CAPPED,
            #   SequenceStatus.FINISHED_ABORTED,
            #   SequenceStatus.FINISHED_IGNORED
            # =========================================================================
            return self.num_unfinished_seqs()

Sequence:对于一个seq,我们重点来看它的属性self.logical_token_blocks(逻辑块)和方法_append_tokens_to_blocks(生成逻辑块的方法)。在vLLM中,每个seq都单独维护一份属于自己的逻辑块,不同的逻辑块可以指向同一个物理块(此刻你一定很关心逻辑块和物理块是如何做映射的,我们会循序渐进地讲解这点,现在你可以先忽略映射方法,把目光聚焦于“一个seq的逻辑块长什么样,怎么初始化它的逻辑块”)

(2)1个逻辑块的结构

我们先来回答“1个逻辑块长什么样”这个问题,逻辑块定义的代码比较简单,所以我们直接看代码(一切尽在注释中),代码路径​​vllm/block.py​

class LogicalTokenBlock:
    """A block that stores a contiguous chunk of tokens from left to right.

    Logical blocks are used to represent the states of the corresponding
    physical blocks in the KV cache.
    
    KV cache的逻辑块
    """

    def __init__(
        self,
        block_number: int, # 逻辑块的序号
        block_size: int, # 每个逻辑块中有多少个槽位(默认为16)
    ) -> None:
        self.block_number = block_number
        self.block_size = block_size

        # 逻辑块刚初始化时,将其中的每个token_id都初始化为_BLANK_TOKEN_ID(-1)
        self.token_ids = [_BLANK_TOKEN_ID] * block_size 
        # 当前逻辑块中已经装下的token的数量
        self.num_tokens = 0

    def is_empty(self) -> bool:
        """判断当前逻辑块是为空"""
        return self.num_tokens == 0

    def get_num_empty_slots(self) -> int:
        """当前逻辑块的空余槽位"""
        return self.block_size - self.num_tokens

    def is_full(self) -> bool:
        """判断当前逻辑块是否已经被装满"""
        return self.num_tokens == self.block_size

    def append_tokens(self, token_ids: List[int]) -> None:
        """将给定的一些token_ids装入当前逻辑块中"""
        # 给定的token_ids的长度必须 <= 当前逻辑块剩余的槽位
        assert len(token_ids) <= self.get_num_empty_slots()
        # 当前逻辑块第一个空槽的序号
        curr_idx = self.num_tokens
        # 将这些tokens装进去
        self.token_ids[curr_idx:curr_idx + len(token_ids)] = token_ids
        # 更新当前逻辑块中tokens的数量
        self.num_tokens += len(token_ids)

    def get_token_ids(self) -> List[int]:
        """获取当前逻辑块中所有被装满的位置的token_ids"""
        return self.token_ids[:self.num_tokens]

    def get_last_token_id(self) -> int:
        """获取当前逻辑块所所有被装满的位置的最后一个token_id"""
        assert self.num_tokens > 0
        return self.token_ids[self.num_tokens - 1]

 (3)再回到Sequence上来

知道了每个逻辑块的结构,我们现在来回答“怎么给一个seq分配逻辑块”这个问题,也就是回到2.3(1)中Sequence的​​_append_tokens_to_blocks​​方法上来:当一个seq只有prompt时,这个方法负责给prompt分配逻辑块;当这个seq开始产出output时,这个方法负责给每一个新生成的token分配逻辑块,整个过程如下图(图片来自vLLM论文,大家忽略图中block_table的部分):

代码如下(一切尽在注释中,​​/vllm/sequence.py​​):

def _append_tokens_to_blocks(self, token_ids: List[int]) -> None:
        """
        将token_ids动态填入逻辑块列表中
        Args:
            token_ids: prompt部分的token_ids
        """
        cursor = 0
        # 遍历prompt token_ids中的每一个token_id
        while cursor < len(token_ids):
            # 如果当前逻辑块列表(logical_token_blocks)为空
            if not self.logical_token_blocks:
                # 则先append一个逻辑块,该逻辑块index为0,大小为16,其中的每一个token_id为-1
                self._append_logical_block()

            # 取出逻辑块列表中的最后一个逻辑块
            last_block = self.logical_token_blocks[-1]
            # 如果这最后一个逻辑块中已经没有槽位
            if last_block.is_full():
                # 那么再append一个逻辑块,其大小为16,其中每一个token_id为-1
                self._append_logical_block()
                # 把这个新append的逻辑块取出来
                last_block = self.logical_token_blocks[-1]
            
            # 检查当前取出的逻辑块中空槽位的数量
            num_empty_slots = last_block.get_num_empty_slots()
            # 用当前的token_ids填充空槽位,直到无法填满为止
            last_block.append_tokens(token_ids[cursor:cursor +
                                               num_empty_slots])
            cursor += num_empty_slots

好,到目前为止,我们就把vLLM对输入数据做预处理的部分介绍完了,简单总结下:

  • 在vLLM内部计算逻辑中,1个prompt是1个request
  • 每个prompt将被包装成一个SequenceGroup实例提供给调度器做调度
  • 1个SequenceGroup实例下维护着若干个Sequence实例,对应着“1个prompt -> 多个outputs"这种更一般性的解码场景。
  • 1个Sequence实例下维护着属于自己的逻辑块列表,数据类型为List[LogicalTokenBlock]

三、add_request():将seq_group添加进调度器waiting队列

写了这么多,你是不是已经忘记上面都说了些什么了,不要紧,我们快速回顾下:

  • 首先,我们明确了vLLM最重要的推理内核引擎是LLMEngine
  • LLMEngine下有两个最重要的方法:add_request()和step()
  • add_request()负责将每个prompt都包装成一个SequenceGroup对象,送入调度器的waiting队列中等待调度
  • step()负责执行1次推理过程,在这个过程中,调度器首先决定哪些seq_group可以被送去推理,然后model_executor负责实际执行推理。

现在,在知道SequenceGroup相关定义的基础上,我们可以来看add_request()了,我们直接来看代码(一切尽在注释中,为了方便阅读,代码有所省略):

# vllm/engine/llm_engine.py
    def add_request(
        self,
        request_id: str, # 每个请求的唯一id
        prompt: Optional[str], # prompt(文字版)
        sampling_params: SamplingParams, # 用于采样的参数(温度、topk等)
        prompt_token_ids: Optional[List[int]] = None, # prompt(input_ids版)
        arrival_time: Optional[float] = None, # 请求到达的时间。如果是None,则用当前系统时间
        lora_request: Optional[LoRARequest] = None,  # 如果是用lora模型做推理,相关的lora请求
        multi_modal_data: Optional[MultiModalData] = None, # 每个请求的多模态数据
    ) -> None:
        """
        将request添加给LLMEngine

        Args:
            request_id: 在vLLM内部,1条prompt算1个请求,会附给1个请求id
            prompt: prompt(文字版)
            sampling_params: 采样参数(温度、topk等)
            prompt_token_ids: prompt(token_id版),没有提供的话vLLM会调用tokenizer来做
            arrival_time: 请求到达的时间。如果是None,则用当前系统时间
            multi_modal_data: 多模态数据(暂时忽略不看)
        """
        ...
        
        # ============================================================================
        # 设置该请求的到达时间
        # ============================================================================
        if arrival_time is None:
            arrival_time = time.time()
        ...

        # 每个KV cache block的大小(默认为16)
        block_size = self.cache_config.block_size
        # 当前seq的id(见后文讲解)
        seq_id = next(self.seq_counter)
        # 获取用于表示<eos>的token_id
        eos_token_id = self.tokenizer.get_lora_tokenizer(
            lora_request).eos_token_id
        
        # ============================================================================
        # 为当前序列创建Sequence对象,在Sequence对象中也包括对当前序列逻辑块们的管理
        # ============================================================================
        seq = Sequence(seq_id, prompt, prompt_token_ids, block_size,
                       eos_token_id, lora_request)
        ...
        # ============================================================================
        # 每个prompt被包装成一个SequenceGroup实例
        # ============================================================================
        seq_group = SequenceGroup(request_id, [seq], sampling_params,
                                  arrival_time, lora_request, multi_modal_data)

        # ============================================================================
        # 将seq_group中所有序列添加进scheduler的self.waiting队列中
        # self.waiting是一个双端队列实例,我们可以在队列的两端进行插入/删除操作
        # ============================================================================
        self.scheduler.add_seq_group(seq_group)

四、step():调度器策略

现在所有的seq_group都已经被送入调度器(Scheduler)的waiting队列中了,接下来我们就来看,在1个推理阶段中,调度器是通过什么策略来决定要送哪些seq_group去做推理的,这也是vLLM难啃的硬骨头之一。

调度器相关的代码都在vllm/core/scheduler.py中,由于代码逻辑嵌套比较复杂,所以我们依然先通过图解的方式把整个调度流程介绍一遍,然后再看关键的源码细节。

4.1 调度器结构

vLLM调度器维护的重要属性如上图所示:

  • ​self.waiting, self.running, self.swapped​​:这三个都是python的deque()实例(双端队列,允许你从队列两侧添加或删除元素)。
  • waiting队列用于存放所有还未开始做推理的seq_group,“未开始”指连prefill阶段都没有经历过。所以waiting队列中的seq_group只有一个seq,即是原始的prompt。
  • running队列用于存放当前正在做推理的seq_group。更准确地说,它存放的是上1个推理阶段被送去做推理的seq_group们,在开始新一轮推理阶段时,调度器会根据本轮的筛选结果,更新running队列,即决定本轮要送哪些seq_group去做推理。
  • swapped队列用于存放被抢占的seq_group。在2.2节中我们有提过,若一个seq_group被抢占,调度器会对它执行swap或recomputation操作,分别对应着将它送去swapped队列或waiting队列,在后文我们会详细分析抢占处理的代码
  • ​self.policy​​:是vLLM自定义的一个Policy实例,目标是根据调度器总策略(FCFS,First Come First Serve,先来先服务)原则,对各个队列里的seq_group按照其arrival time进行排序。相关代码比较好读,所以这里我们只概述它的作用,后续不再介绍它的代码实现。
  • ​self.prev_time​​:上一次调度发起的时间点,初始化为0。我们知道每执行1次推理阶段前,调度器都要做一次调度,这个变量存放的就是上次调度发起的时间点。
  • ​self.prev_prompt​​:取值为True/False,初始化为False。若上一次调度时,调度器有从waiting队列中取出seq_group做推理,即为True,否则为False
  • ​self.last_prompt_latency​​:记录“当前调度时刻(now) -  最后一次有从waiting队列中取数做推理的那个调度时刻”的差值(并不是每一次调度时,调度器一定都会从waiting队列中取seq_group,它可能依旧继续对running队列中的数据做推理),初始化为0。

目前你可能很难明白这三个属性的作用,不要着急,在后文讲解具体调度流程时,我们会再来看它们。这里只需记住它们的定义即可。

  • ​BlockManager​​:物理块管理器。这也是vLLM自定义的一个class。截止本文写作时,vLLM提供了​​BlockSpaceManagerV1​​​和​​BlockSpaceManagerV2​​两个版本的块管理器。V1是vLLM默认的版本,V2是改进版本(但还没开发完,例如不支持prefix caching等功能)。所以本文依然基于BlockSpaceManagerV1进行讲解。物理块管理器这个class下又维护着两个重要属性:
  • ​BlockAllocator​​:物理块分配者,负责实际为seq做物理块的分配、释放、拷贝等操作。这也是我们后文要解读的对象。其下又分成self.gpu_allocator和self.cpu_allocator两种类型,分别管理gpu和cpu上的物理块。
  • ​self.block_tables​​:负责维护每个seq下的物理块列表,本质上它是一个字典,形式如​​{seq_id: List[PhysicalTokenBlock]}​​。注意,这里维护者【所有】seq_group下seq的物理块,而不是单独某一个seq的。因为整个调度器都是全局的,其下的BlockManager自然也是全局的。

读到这里,你还记得2.3节中我们曾介绍过,每个Sequence实例中维护着属于这个seq的逻辑块吗?而我们从self.block_tables中,又能根据seq_id找到这个seq对应的物理块。这就实现了“逻辑块 -> 物理块”的映射。在刚开始读代码的时候,很多朋友从直觉上都会觉得​​BlockManager​​就是用来存储逻辑块和物理块映射的,其实它只负责管理和分配物理块,映射关系潜藏在seq中。理解这点对理解代码非常重要。

现在,我们就把调度器(Scheduler)的结构理清了。我知道你肯定还有很多疑惑。所以我们马上来看调度策略的具体流程:“对于装在waiting、running、swapped队列中的那些seq_group,是根据什么规则决定本次推理阶段该送谁去推理呢?”

4.2 整体调度流程

上图刻画了某次调度步骤中三个队列的情况,再复习一下:

  • waiting队列中的数据都没有做过prefill,每个seq_group下只有1个seq(prompt)
  • running队列中存放着上一个推理阶段被送去做推理的所有seq_group
  • swapped队列中存放着之前调度阶段中被抢占的seq_group

running队列中的seq_group不一定能继续在本次调度中被选中做推理,这是因为gpu上KV cache的使用情况一直在变动,以及waiting队列中持续有新的请求进来的原因。所以调度策略的职责就是要根据这些变动,对送入模型做推理的数据做动态规划。

根据源码,我将vLLM调度步骤整理成上述流程图。看着有点复杂是吧,不要担心,我们这就来拆解它。

总结来说:

  • 如果当前swapped队列为空,那就去检查是否能从waiting队列中调度seq_group,直到不满足调度条件为止(gpu空间不足,或waiting队列已为空等)。此时,1个推理阶段中,所有的seq_group都处在prefill阶段。
  • 如果当前swapped队列非空,或者无法从waiting队列中调度任何seq_group时
  • 检查是否能从running队列中调度seq_group,直到不满足调度条件为止。
  • 若本次无新的被抢占的seq_group,且swapped队列非空,就检查是否能从swapped队列中调度seq_group,直到不满足调度条件为止。

此时,1个推理阶段中,所有的seq_group要么全来自running队列,要么来自running + swapped队列,它们都处在decode阶段。

至此我们要记住vLLM调度中非常重要的一点:在1个推理阶段中,所有的seq_group要么全部处在prefill阶段。要么全部处在decode阶段。

你可能想问:为什么要以swapped是否非空为判断入口呢?

这是因为,如果当前调度步骤中swapped队列非空,说明在之前的调度步骤中这些可怜的seq_group因为资源不足被抢占,而停滞了推理。所以根据FCFS规则,当gpu上有充足资源时,我们应该先考虑它们,而不是考虑waiting队列中新来的那些seq_group

同理,在图中你会发现,当我们进入对running队列的调度时(图中红色分支),我们会根据“本次调度是否有新的被抢占的seq_group”,来决定要不要调度swapped队列中的数据。这个理由也很简单:在本次调度中,我就是因为考虑到gpu空间不足的风险,我才新抢占了一批序列。既然存在这个风险,我就最好不要再去已有的swapped队列中继续调度seq_group了。

到这里,我们已经把整个调度流程的关键点给说完了。接下来,我们会配合源码,对上图中的细节进行介绍。

4.3 _passed_delay:判断调度waiting队列的时间点

在4.2的流程图中,我们会看到进入waiting循环的判断条件之一是:waiting队列是否达到调度间隔阈值。这是个什么东西?又为什么要设置这样一个阈值呢?

我们知道模型在做推理时,waiting队列中是源源不断有seq_group进来的,一旦vLLM选择调度waiting队列,它就会停下对running/swapped中seq_group的decode处理,转而去做waiting中seq_group的prefill,也即vLLM必须在新来的seq_group和已经在做推理的seq_group间取得一种均衡:既不能完全不管新来的请求,也不能耽误正在做推理的请求。所以“waiting队列调度间隔阈值”就是来控制这种均衡的

  • 调度间隔设置得太小,每次调度都只关心waiting中的新请求,这样发送旧请求的用户就迟迟得不到反馈结果。且此时waiting队列中积累的新请求数量可能比较少,不利于做batching,浪费了并发处理的能力。
  • 调度间隔设置得太大,waiting中的请求持续挤压,同样对vLLM推理的整体吞吐有影响。

那这个阈值在代码中是怎么控制的呢?还记得4.1中我们画Scheduler的结构图时有三个乍一看比较难懂的属性吗(见下图),它们就是用来控制这个阈值的:

​vllm/core/scheduler.py​​​脚本的​​_passed_delay()​​函数写了阈值判断的相关逻辑,我们直接看代码(一切尽在注释中):

def _passed_delay(self, now: float) -> bool:
        """
        判断当下是否可以从waiting队列中调度新请求
        这个函数确保了在调度过程中不会频繁地处理新来的seq_group
        
        Args:
            now: 当前调度时间点
        """
        # =============================================================================
        # self.prev_prompt: True/False,记录上一次调度步骤中,是否选择了从waiting队列中做调度
        # self.prev_time:上次调度步骤时间点(不管是从哪个队列中调度,每次调度都会记录下时间点)
        # 若上个调度步骤中,我们选择从waiting队列中做调度,则计算两个调度时刻的间隔
        # ==============================================================================
        if self.prev_prompt:
            self.last_prompt_latency = now - self.prev_time
        
        # =============================================================================
        # 用当前调度时间更新prev_time
        # 由于目前还不知道本次是否会从waiting队列中调度,因此prev_prompt先设为False
        # =============================================================================
        self.prev_time, self.prev_prompt = now, False
        
        # =============================================================================
        # Delay scheduling prompts to let waiting queue fill up
        # delay_factor:用户配置的,用于调整调度间隔阈值的因子。大于0则意味着用户想开启阈值判断
        # =============================================================================
        if self.scheduler_config.delay_factor > 0 and self.waiting:
            # =========================================================================
            # 计算在waiting队列中,最早到达的seq_group的到达时间
            # =========================================================================
            earliest_arrival_time = min(
                [e.metrics.arrival_time for e in self.waiting])
            # =========================================================================
            # now - earliest_arrival_time:最早到达waiting队列的seq_group当前“实际”等待的时间
            # delay_factor*last_prompt_latency:最早到达waiting队列的请求当前“应该”等待的时间
            # 只要前者比后者大,或者此时running队列中根本没有请求在跑,就可以进行对waiting做调度
            # =========================================================================
            passed_delay = (
                (now - earliest_arrival_time) >
                (self.scheduler_config.delay_factor * self.last_prompt_latency)
                or not self.running)
        # =============================================================================
        # 如果你不想开启阈值判断,那就直接返回True
        # =============================================================================
        else:
            passed_delay = True
        return passed_delay

4.4 can_allocate:能否为seq_group分配物理块做prefill

通过了调度时间阈值的判断条件,现在我们顺利从waiting中取出一个seq_group,我们将对它进行prefill操作。所以这里我们必须先判断:gpu上是否有充足的空间为该seq_group分配物理块做prefill,根据4.1中绘制的调度器结构,这个操作当然是由我们的self.block_manager来做。

判断的入口代码为​​can_allocate = self.block_manager.can_allocate​​​(seq_group),配合上面图例,我们直接来看​​can_allocate​​函数的代码,(一切尽在注释中):

# vllm/core/block_manager_v1.py
    def can_allocate(self, seq_group: SequenceGroup) -> AllocStatus:
        """
        确实是否可以给这个seq_group分配物理块,返回结果有三种情况:
        - AllocStatus.NEVER:不分配;
        - AllocStatus.OK:可以分配;
        - AllocStatus.LATER:延迟分配
        """
        # FIXME(woosuk): Here we assume that all sequences in the group share
        # the same prompt. This may not be true for preempted sequences.
        # (这里我们假设一个seq_group下的所有序列的prompt都是相同的)
        
        # ===========================================================================
        # 取出这个seq_group下处于waiting状态的序列
        # ===========================================================================
        seq = seq_group.get_seqs(status=SequenceStatus.WAITING)[0]
        
        # ===========================================================================
        # 取出这个seq所有的逻辑块
        # ===========================================================================
        num_required_blocks = len(seq.logical_token_blocks)

        # ===========================================================================
        # block上的滑动窗口(可暂时假设其值为None,先忽略不看
        # ===========================================================================
        if self.block_sliding_window is not None:
            num_required_blocks = min(num_required_blocks,
                                      self.block_sliding_window)
        # ===========================================================================
        # 计算当前所有可用的物理块数量,List[PhysicalTokenBlock]
        # ===========================================================================
        num_free_gpu_blocks = self.gpu_allocator.get_num_free_blocks()

        # ===========================================================================
        # Use watermark to avoid frequent cache eviction.
        # 决定是否能为当前seq分配物理块
        # ===========================================================================
        # 如果设备中所有的物理块数量 - 该seq实际需要的物理块数量 < 水位线block数量,则不分配
        # (说明当前seq太长了)
        if (self.num_total_gpu_blocks - num_required_blocks <
                self.watermark_blocks):
            return AllocStatus.NEVER
        # 如果设备中可用的物理块数量 - 该seq实际需要的block数量 >= 水位线block数量,则分配
        if num_free_gpu_blocks - num_required_blocks >= self.watermark_blocks:
            return AllocStatus.OK
        # 否则,现在不能分配,但可以延迟分配
        else:
            return AllocStatus.LATER

我们对上述代码做一些额外的说明:

  • 代码第32行:​​num_free_gpu_blocks = self.gpu_allocator.get_num_free_blocks()​​。这里是在统计当前gpu上所有可用的物理块数数量(忘记gpu_allocator是什么的朋友,可以再回顾下4.1的调度器结构图)。在vLLM中,gpu_allocator的类型有两种:
  • ​CachedBlockAllocator​​:按照prefix caching的思想来分配和管理物理块。在原理篇中,我们提过又些prompts中可能含有类似system message(例如,“假设你是一个能提供帮助的行车导航”)E)等prefix信息,带有这些相同prefix信息的prompt完全可以共享用于存放prefix的物理块,这样既节省显存,也不用再对prefix做推理。
  • ​UncachedBlockAllocator​​:正常分配和管理物理块,没有额外实现prefix caching的功能。

关于这两种allocator的具体实现方式,我们将放在源码解读第3篇块管理来做讲解。这里大家只要明白大致定义即可,并不影响我们对调度策略的解读。

  • ​self.watermark_blocks​​:水位线block数量,它起的是一个预警和缓冲的作用,防止在1次调度中把gpu上预留给KV Cache的显存空间打得过满,出现一些意外风险(毕竟这个预留的显存空间也是我们估计出来的)。
  • ​NEVER和LATER的区别​​:这两者的相同之处在于,都是因为当前显存空间不够,而无法继续调度seq_group。区别在于,NEVER是因为这条seq实在太长(即prompt太长),长到动用了gpu上所有的block(num_total_gpu_blocks)都无法处理它,所以后续步骤中我们会直接把这个seq标记为完成,不再处理它;而LATER是因为之前可能已经调度了很多seq_group,它们占据了相当一部分显存空间,导致gpu上剩余的可用block(num_free_gpu_blocks)无法再处理它,所以我们延迟处理。   

4.5 can_append_slot:能否为seq_group分配物理块做decode

回顾4.2调度器的流程图,你会看到我们从running队列中调度seq_group时,我们也会判断是否能为该seq_group分配物理块。但这时,我们的物理块空间是用来做decode的(给每个seq分配1个token的位置),而不是用来做prefill的(给每个seq分配若干个token的位置),所以这里我们采取的是另一种判断方法​​can_append_slot​​。

更具体来说,running队列中seq_group下的n个seqs在上1个推理阶段共生成了n个token。在本次调度中,我们要先为这n个token分配物理块空间,用于存放它们在本次调度中即将产生的KV值。

好,我们再回到这个seq_group的n个seqs上来,我们知道:

  • 当往1个seq的物理块上添加1个token时,可能有两种情况:
  • 之前的物理块满了,所以我新开1个物理块给它
  • 之前的物理块没满,我直接添加在最后一个物理块的空槽位上
  • 所以,对于1个seq来说,最坏的情况就是添加1个物理块;对于n个seqs来说,最坏的情况就是添加n个物理块(想想原理篇中讲过的copy-on-write机制)
  • 对于1个seq_group,除了那些标记为“finish”的seq外,其余seqs要么一起送去推理,要么一起不送去推理。即它们是集体行动的

所以,判断能否对一个正在running的seq_group继续做推理的最保守的方式,就是判断当前可用的物理块数量是否至少为n

我们直接看代码(一切尽在注释中,编辑器有问题,大家看截图吧):

4.6 allocate与append_slot:为seq_group分配物理块

当我们判断当前有充足的gpu KV Cache空间给对应的seq_group做新一轮推理时,我们就可以实际给它分配物理块了。这一块的内容涉及的细节太多(不同的prefix caching方式,逻辑块到物理块的映射,物理块释放,物理块的refcount即copy-on-write机制等等),所以我们把这部分留在源码解读3:块管理中来详细说明。

跳过这块并不影响大家对调度器策略的解读。

4.7 preempt:抢占策略

纵观4.2的调度流程,现在我们只剩1个重点没讲了:抢占策略。

其实在2.2介绍SequenceGroup时,我们已经提到了抢占策略的核心逻辑,这里再复制一遍:

在若干个推理阶段后,gpu上的资源不够了,这个seq_group不幸被调度器抢占(preemption),它相关的KV block也被swap out到cpu上。此时所有seq的状态变为swapped。这里要注意,当一个seq_group被抢占时,对它的处理有两种方式

  • Swap:如果该seq_group剩余生命周期中并行运行的最大seq数量 > 1,此时会采取swap策略,即把seq_group下【所有】seq的KV block从gpu上卸载到cpu上。(seq数量比较多,直接把算出的KV block抛弃,比较可惜)
  • Recomputation:如果该seq_group剩余生命周期中并行运行的最大seq数量 = 1,此时会采取recomputation策略,即把该seq_group相关的物理块都释放掉,然后将它重新放回waiting队列中(放在最前面)。等下次它被选中推理时,就是从prefill阶段开始重新推理了,因此被称为“重计算”。(seq数量少,重新计算KV block的成本不高)

对“最大生命周期...”这里有疑惑的朋友,回顾下本文2.3(1)。

我们直接来看代码(一切尽在注释中)

# vllm/core/scheduler.py
    def _preempt(
        self,
        seq_group: SequenceGroup, # 被抢占的seq_group
        blocks_to_swap_out: Dict[int, int],
        preemption_mode: Optional[PreemptionMode] = None,
    ) -> None:
        """
        对被抢占的seq_group进行处理,包括修改其下seq状态,做好gpu到cpu块之间的映射等
        """
        # If preemption mode is not specified, we determine the mode as follows:
        # We use recomputation by default since it incurs lower overhead than
        # swapping. However, when the sequence group has multiple sequences
        # (e.g., beam search), recomputation is not currently supported. In
        # such a case, we use swapping instead.
        # FIXME(woosuk): This makes our scheduling policy a bit bizarre.
        # As swapped sequences are prioritized over waiting sequences,
        # sequence groups with multiple sequences are implicitly prioritized
        # over sequence groups with a single sequence.
        # TODO(woosuk): Support recomputation for sequence groups with multiple
        # sequences. This may require a more sophisticated CUDA kernel.
        
        # 如果没有指定被抢占的类型
        if preemption_mode is None:
            # 如果这个seq_group在剩余生命周期中并行运行的最大seq数为1
            if seq_group.get_max_num_running_seqs() == 1:
                # 就将抢占类型定位“recompute”
                preemption_mode = PreemptionMode.RECOMPUTE
            # 否则定为swap
            else:
                preemption_mode = PreemptionMode.SWAP
        
        # =======================================================================
        # 如果抢占类型是“RECOMPUTE”
        # 则去除该seq对对应物理块的引用,同时将该seq状态改为running,放入waiting队列最前面
        # (详情参见self._preempt_by_recompute)
        # =======================================================================
        if preemption_mode == PreemptionMode.RECOMPUTE:
            self._preempt_by_recompute(seq_group)
        # =======================================================================
        # 如果抢占类型是“SWAP“
        # 详情参见self._preempt_by_swap)
        # =======================================================================
        elif preemption_mode == PreemptionMode.SWAP:
            self._preempt_by_swap(seq_group, blocks_to_swap_out)
        else:
            raise AssertionError("Invalid preemption mode.")

    def _preempt_by_recompute(
        self,
        seq_group: SequenceGroup,
    ) -> None:
        # 获取这个seq_group下正在running的所有seqs,
        # preemption_mode是RECOMPUTE时需要满足正在running的seqs数量为1
        seqs = seq_group.get_seqs(status=SequenceStatus.RUNNING)
        assert len(seqs) == 1
        
        for seq in seqs:
            # 将这条seq的状态从running改成waiting(后续这条seq就要重计算了)
            seq.status = SequenceStatus.WAITING
            # 释放这条seq对应的物理块
            # 即将对应物理块的引用-1,如果此时引用数量为0,说明对应物理块完全自由了,需要再将其放入自由物理块列表中
            self.free_seq(seq)
            # 因为这条seq需要重计算了,所以将其data对象下_num_computed_tokens设置为0
            seq.reset_state_for_recompute()
        
        # NOTE: For FCFS, we insert the preempted sequence group to the front
        # of the waiting queue.
        # 将被抢占,且未来需要重计算的序列,放到waiting队列的最前面
        self.waiting.appendleft(seq_group)

    def _preempt_by_swap(
        self,
        seq_group: SequenceGroup,
        blocks_to_swap_out: Dict[int, int],
    ) -> None:
        # ======================================================================
        # - 释放该seq_group下所有seq的物理块,并为其分配对应的cpu物理块,
        # - 将seq的状态从running改成swapped
        # ======================================================================
        self._swap_out(seq_group, blocks_to_swap_out)
        # ======================================================================
        # 在scheduler的swapped队列中添加该seq_group
        # ======================================================================
        self.swapped.append(seq_group)
  
      def _swap_out(
        self,
        seq_group: SequenceGroup, # 需要被swap到cpu上的seq_group
        blocks_to_swap_out: Dict[int, int],
    ) -> None:
        # ======================================================================
        # 检查是否可以将当前seq_group对应的物理块swap到cpu上
        # 可以的条件:当前seq_group占用的gpu物理块数量 <= cpu上可用的物理块数量
        # ======================================================================
        if not self.block_manager.can_swap_out(seq_group):
            # FIXME(woosuk): Abort the sequence group instead of aborting the
            # entire engine.
            raise RuntimeError(
                "Aborted due to the lack of CPU swap space. Please increase "
                "the swap space to avoid this error.")
        # ======================================================================
        # 释放该seq_group下所有seq的gpu物理块,并为其创建对应的cpu块
        # mapping:{gpu物理块id:cpu物理块id}
        # ======================================================================
        mapping = self.block_manager.swap_out(seq_group)
        blocks_to_swap_out.update(mapping)
        
        # ======================================================================
        # 修改该seq_group下所有seq的状态:从running改成swapped
        # ======================================================================
        for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
            seq.status = SequenceStatus.SWAPPED

额外说明的一点是,一旦你决定执行swap out操作,你就做做好gpu物理块->cpu物理块之间的映射,这样等之后你想swap in时,你才知道去cpu上的哪里把这些物理块找回来。

swap更多的细节也会涉及到blockmanager,所以遗留的细节,我们也放在第三篇中说。

# vllm/core/block_manager_v1.py
def can_allocate(self, seq_group: SequenceGroup) -> AllocStatus:
"""
        """
        
        # ==============================================================================
        # blocks_to_swap_in:{cpu物理块id: gpu物理块id}
        # blocks_to_swap_out:{gpu物理块id: cpu物理块id}
        # blocks_to_copy: {旧物理块id:[由旧物理块copy-on-write而来的新物理块id]}
        # ==============================================================================
        blocks_to_swap_in: Dict[int, int] = {}
        blocks_to_swap_out: Dict[int, int] = {}
        blocks_to_copy: Dict[int, List[int]] = {}

        # ==============================================================================
        # Fix the current time.
        # 获取当下时间
        # ==============================================================================
        now = time.time()

        # ==============================================================================
        # Join waiting sequences if possible.
        # 如果swapped队列为空
        # ==============================================================================
        if not self.swapped:
            # ==========================================================================
            # ignored_seq_groups:记录因太长(所需的blocks和总blocks之间的差值超过阈值了),
            # 而无法继续做生成的seq_group,这些seq_group中的seq状态都会被标记为
            # FINISHED_IGNORED,表示直接不处理他们
            # ==========================================================================
            ignored_seq_groups: List[SequenceGroup] = []
            
            # ==========================================================================
            # 记录本次被调度的seq_group
            # ==========================================================================
            scheduled: List[SequenceGroup] = []
            
            # ==========================================================================
            # The total number of sequences on the fly, including the
            # requests in the generation phase.
            # 计算Scheduler running队列中还没有执行完的seq数量
            # ==========================================================================
            num_curr_seqs = sum(seq_group.get_max_num_running_seqs()
                                for seq_group in self.running)
            curr_loras = set(
                seq_group.lora_int_id
                for seq_group in self.running) if self.lora_enabled else None

            # ==========================================================================
            # Optimization: We do not sort the waiting queue since the preempted
            # sequence groups are added to the front and the new sequence groups
            # are added to the back.
            # lora相关的,可以暂时不看
            # ==========================================================================
            leftover_waiting_sequences = deque()
            
            # ==========================================================================
            # 本次调度处理的token总数
            # ==========================================================================
            num_batched_tokens = 0

            # ==========================================================================
            # 开启新一次调度(while循环不结束意味着本次调度不结束,
            # 跳出while循环时意味着本次调度结束了)
            # 开启新一次调度的条件:当waiting队列中有等待处理的请求,且当前时刻可以处理请求
            # ==========================================================================
            while self._passed_delay(now) and self.waiting:
                # =====================================================================
                # 取出waiting队列中的第一个请求,也即最早到达的请求(seq_group)
                # =====================================================================
                seq_group = self.waiting[0]
                
                # =====================================================================
                # 统计该seq_group中s处于waiting的seq的数量
                # =====================================================================
                waiting_seqs = seq_group.get_seqs(
                    status=SequenceStatus.WAITING)
                # =====================================================================
                # 从waiting队列中取出来的seq_group,其seq数量一定是1
                # =====================================================================
                assert len(waiting_seqs) == 1, (
                    "Waiting sequence group should have only one prompt "
                    "sequence.")
                
                # =====================================================================
                # 获取该seq的序列长度(如果该seq_group来自之前被抢占的请求,
                # 那么这个长度不仅包括prompt,
                # 还包括output)
                # =====================================================================                 
                num_prefill_tokens = waiting_seqs[0].get_len()
                
                # =====================================================================
                # 如果从waiting队列中取出的这条seq的长度 > 每次调度能处理的最大序列长度,
                # 那么就打印警告信息,同时把这条seq的状态置为FINISHED_IGNORED,
                # 并将对应seq_group装入ignored_seq_groups中,
                # 然后将其从waiting列表中移除,不再处理
                # =====================================================================
                if num_prefill_tokens > self.prompt_limit:
                    logger.warning(
                        f"Input prompt ({num_prefill_tokens} tokens) is too "
                        f"long and exceeds limit of {self.prompt_limit}")
                    for seq in waiting_seqs:
                        seq.status = SequenceStatus.FINISHED_IGNORED
                    ignored_seq_groups.append(seq_group)
                    self.waiting.popleft()
                    continue
                
                # =====================================================================
                # If the sequence group cannot be allocated, stop.
                # 决定是否能给当前seq_group分配物理块
                # can_allocate返回值可能有三种:
                #       AllocStatus.NEVER:不分配;
                #       AllocStatus.OK:可以分配;
                #       AllocStatus.LATER:延迟分配
                # =====================================================================
                can_allocate = self.block_manager.can_allocate(seq_group)
                # 若是延迟分配,则说明现在没有足够的block空间,所以跳出while循环(不继续对waiting队列中的数据做处理了)
                if can_allocate == AllocStatus.LATER:
                    break
                # 如果不分配,说明seq长得超出了vLLM的处理范围,则后续也不再处理它,直接将该seq状态标记为FINISHED_IGNORED
                elif can_allocate == AllocStatus.NEVER:
                    logger.warning(
                        f"Input prompt ({num_prefill_tokens} tokens) is too "
                        f"long and exceeds the capacity of block_manager")
                    for seq in waiting_seqs:
                        seq.status = SequenceStatus.FINISHED_IGNORED
                    ignored_seq_groups.append(seq_group) # 记录因为太长而无法处理的seq_group
                    self.waiting.popleft() # 将该seq_group从waiting队列中移除
                    continue

                # =====================================================================                 # lora推理相关的部分,可暂时忽略
                # =====================================================================
                lora_int_id = 0
                if self.lora_enabled:
                    lora_int_id = seq_group.lora_int_id
                    if (lora_int_id > 0 and lora_int_id not in curr_loras
                            and len(curr_loras) >= self.lora_config.max_loras):
                        # We don't have a space for another LoRA, so
                        # we ignore this request for now.
                        leftover_waiting_sequences.appendleft(seq_group)
                        self.waiting.popleft()
                        continue

                # =====================================================================
                # If the number of batched tokens exceeds the limit, stop.
                # max_num_batched_tokens:单次调度中最多处理的token数量
                # num_batched_tokens:本次调度中累计处理的token数量
                # 如果后者 > 前者,则结束本次调度
                # =====================================================================
                num_batched_tokens += num_prefill_tokens
                if (num_batched_tokens >
                        self.scheduler_config.max_num_batched_tokens):
                    break

                # =====================================================================
                # The total number of sequences in the RUNNING state should not
                # exceed the maximum number of sequences.
                # num_new_seqs: 当前seq_group中状态为“未执行完”的序列的数量
                # num_curr_seqs:当前调度轮次中,状态为"未执行完“的序列总数
                # 如果超过了我们对单次调度能执行的序列总数的阈值,就结束本次调度
                # =====================================================================
                num_new_seqs = seq_group.get_max_num_running_seqs()
                if (num_curr_seqs + num_new_seqs >
                        self.scheduler_config.max_num_seqs): # 单次迭代中最多处理多少个序列
                    break

                if lora_int_id > 0:
                    curr_loras.add(lora_int_id)
                
                # =====================================================================
                # 走到这一步时,说明当前seq_group已经通过上述种种验证,可以被加入本次调度中执行了
                # 先将其从waiting队列中移出
                # =====================================================================
                self.waiting.popleft()
                
                # =====================================================================
                # 为当前seq_group分配物理块
                # =====================================================================
                self._allocate(seq_group)
                
                # =====================================================================
                # 将当前seq_group放入running队列中
                # =====================================================================
                self.running.append(seq_group)
                
                # =====================================================================
                # 记录本次调度累计处理的序列数量
                # =====================================================================
                num_curr_seqs += num_new_seqs
                
                # =====================================================================
                # 记录本次被调度的seq_group
                # =====================================================================
                scheduled.append(
                    ScheduledSequenceGroup(
                        seq_group=seq_group,
                        token_chunk_size=num_prefill_tokens))
            
            # =====================================================================
            # 和lora相关的操作,暂时忽略
            # =====================================================================
            self.waiting.extendleft(leftover_waiting_sequences)

            # =====================================================================
            # 如果本次有被调度的seq_group(scheduled非空)
            # 或者本次有被设置为不再处理的seq_group(ignored_seq_groups非空)
            # 就将其包装成SchedulerOutputs对象
            # =====================================================================
            if scheduled or ignored_seq_groups:
                self.prev_prompt = True
                scheduler_outputs = SchedulerOutputs(
                    scheduled_seq_groups=scheduled,
                    prompt_run=True,
                    num_batched_tokens=num_batched_tokens,
                    blocks_to_swap_in=blocks_to_swap_in,
                    blocks_to_swap_out=blocks_to_swap_out,
                    blocks_to_copy=blocks_to_copy,
                    ignored_seq_groups=ignored_seq_groups,
                )
                return scheduler_outputs

        # ==============================================================================
        # NOTE(woosuk): Preemption happens only when there is no available slot
        # to keep all the sequence groups in the RUNNING state.
        # In this case, the policy is responsible for deciding which sequence
        # groups to preempt.
        # 如果swap队列非空,且本次没有新的需要被发起推理的seq_group,
        # 则对running队列中的seq_group,
        # 按照 "当前时间-该seq_group到达时间" ,从早到晚排列running队列中的seq_group
        # ==============================================================================
        self.running = self.policy.sort_by_priority(now, self.running)

        # ==============================================================================
        # Reserve new token slots for the running sequence groups.
        # 初始化一个新的running队列(deque())
        # 初始化一个抢占列表
        # ==============================================================================
        running: Deque[SequenceGroup] = deque()
        preempted: List[SequenceGroup] = []
        
        # ==============================================================================
        # 当running队列非空时
        # ==============================================================================
        while self.running:
            # 取出running队列中最早到来的seq_group
            seq_group = self.running.popleft()
            # =====================================================================
            # 对于running队列中这个最早到来的seq_group,检查对于其中的每一个seq,
            # 是否能至少分配一个物理块给它,如果不能的话
            # (说明要执行抢占操作了,否则马上会没有资源让这个最早到达的seq_group做完推理):
            # (注意,这里用了while...else,如果while条件正常结束,则进入else内容;
            #  如果被break,则不会执行else)
            # =====================================================================
            while not self.block_manager.can_append_slot(seq_group):
                # =====================================================================
                # 如果从running队列中取出最早达到的seq_group后,running队列还是非空
                # =====================================================================
                if self.running:
                    # ==============================================================
                    # 抢占running队列中最晚到来的seq_group(可怜的被害者)
                    # ==============================================================
                    victim_seq_group = self.running.pop()
                    
                    # ==============================================================
                    # 一个seq_group被抢占后,有2中处理方式:
                    # - 如果该seq_group下只有一个seq,执行【重计算】,
                    #   将其从running队列中移除,并清空它的物理块,
                    #   将其seq的状态从running->waiting,并加入waiting队列。后面将重新计算
                    #
                    # - 如果该seq_group下有多个seq,执行【swap】,
                    #   清空它的gpu物理块,并为这些物理块做好cpu物理块映射,
                    #   这些seq的block_table字典中({seq_id: block_table})的block_table
                    #   从gpu物理块改成cpu物理块
                    #   将其seqs状态从running -> swapped,加入swapped队列
                    # ==============================================================
                    self._preempt(victim_seq_group, blocks_to_swap_out)
                    preempted.append(victim_seq_group)
                # ==============================================================
                # 如果除这个最早到来的seq_group外,running队列中再没有别的seq_group了,
                # 且此时又没有足够的空间留给这个最早来的seq_group做推理了,那么只能抢占它
                # ==============================================================
                else:
                    # 那就只能抢占这个最早到达的seq_group了
                    # No other sequence groups can be preempted.
                    # Preempt the current sequence group.
                    self._preempt(seq_group, blocks_to_swap_out)
                    preempted.append(seq_group)
                    break
            # ==============================================================
            # 如果此时有足够的空间给running队列中最早来的seq_group做推理了
            # ==============================================================
            else:
                # ==============================================================
                #  Append new slots to the sequence group.
                # seq_group里的每个seq正常做推理。假设现在每个seq正常生成一个token,我们需要根据每个seq当前
                # 维护的最后一个物理块的情况,决定是否需要分配新的物理块,决定的结果可能如下:
                # - 物理块refcount = 1,且有充足槽位,则无需分配新物理块
                # - 物理块refcount = 1,且无充足槽位,分配新的物理块
                # - 物理块refcount > 1, 采用copy-on-write机制,分配新物理块,对该seq,
                #                      用新物理块替换掉其block_table中维护的最后一个物理块
                #                     (称为旧物理块)。释放旧物理块(令其refcount-1)。
                #                      同时记录下新旧物理块之间的映射,
                #   blocks_to_copy:{旧物理块id:[由旧物理块copy-on-write而来的新物理块id]}
                # ==============================================================
                self._append_slot(seq_group, blocks_to_copy)
                
                # ==============================================================
                # 自定义的running队列中添加这个seq_group
                # ==============================================================
                running.append(seq_group)
        
        # ==============================================================================
        # 最终还能在running队列中运行的seq_group
        # ==============================================================================
        self.running = running

        # ==============================================================================
        # Swap in the sequence groups in the SWAPPED state if possible.
        # 对于swapped队列中的seq_group,按照到达时间从早到晚排序
        # ==============================================================================
        self.swapped = self.policy.sort_by_priority(now, self.swapped)
        
        # ==============================================================================
        # 如果本次调度没有新安排的被抢占的seq_group(即preempted为空)
        # ==============================================================================
        if not preempted:
            # ==============================================================
            # 计算running队列中,所有seq_group下,“到生命周期结束为止最多运行的seq数量”的总和
            # ==============================================================
            num_curr_seqs = sum(seq_group.get_max_num_running_seqs()
                                for seq_group in self.running)
            # ==============================================================
            # lora部分,暂时忽略
            # ==============================================================
            curr_loras = set(
                seq_group.lora_int_id
                for seq_group in self.running) if self.lora_enabled else None

            # ==============================================================
            # lora相关的,可以暂时不看
            # ==============================================================
            leftover_swapped = deque()

            # ==============================================================
            # 当swapped队列非空时
            # ==============================================================
            while self.swapped:
                # ==============================================================
                # 取出swap队列中最早被抢占的seq_group
                # ==============================================================
                seq_group = self.swapped[0]
                # ==============================================================
                # lora相关,暂时不看
                # ==============================================================
                lora_int_id = 0
                if self.lora_enabled:
                    lora_int_id = seq_group.lora_int_id
                    if (lora_int_id > 0 and lora_int_id not in curr_loras
                            and len(curr_loras) >= self.lora_config.max_loras):
                        # We don't have a space for another LoRA, so
                        # we ignore this request for now.
                        leftover_swapped.appendleft(seq_group)
                        self.swapped.popleft()
                        continue

                # ==============================================================
                # If the sequence group cannot be swapped in, stop.
                # 判断一个被swap的seq_group现在是否能重新running起来
                # 【判断条件】:
                # 当前gpu上可用的物理块数量 - 重新跑起这个seq_group需要的物理块数量 
                # >= 水位线物理块数量
                # 其中:
                # 后者 = 在被swap之前它已经使用的物理块数量(去重过了) 
                #       + 若能再次跑起来它至少需要的物理块数量
                #(假设每个seq至少需要1个物理块)
                # ==============================================================
                # 如果不能,则意味着当前没有充足资源处理swap队列中的seq_group,则直接跳出循环
                if not self.block_manager.can_swap_in(seq_group):
                    break

                # ==============================================================
                # The total number of sequences in the RUNNING state should not
                # exceed the maximum number of sequences.
                # 如果对于swap队列中的这个seq_group,当前gpu上有充足资源可以让它重新跑起来的话:
                # ==============================================================
                # 取出这个seq_group在剩余生命周期内将并行运行的最大序列数
                num_new_seqs = seq_group.get_max_num_running_seqs()
                # 如果已超过一次调度中能处理的最大序列数,则不再对该seq_group进行处理
                if (num_curr_seqs + num_new_seqs >
                        self.scheduler_config.max_num_seqs):
                    break
                
                # lora部分暂时不看
                if lora_int_id > 0:
                    curr_loras.add(lora_int_id)
                
                # ==============================================================
                # 走到这一步,说明可以对swapped队列中的这个seq_group做相关处理了,
                # 先把它从队列中移出去
                # ==============================================================
                self.swapped.popleft()
                
                # ==============================================================
                # 将该seq_group下所有cpu块置换回gpu块,
                # 并将其下每个seq的状态从swapped改成running
                # ==============================================================
                self._swap_in(seq_group, blocks_to_swap_in)
                
                # ==============================================================
                # 假设其正常做推理了,假设现在生成了一个token,要如何分配物理块(参见上面注释)
                # ==============================================================
                self._append_slot(seq_group, blocks_to_copy)
                num_curr_seqs += num_new_seqs
                self.running.append(seq_group)

            self.swapped.extendleft(leftover_swapped)

        # ==============================================================================
        # 如果本次调度有新安排的被抢占的seq_group(即preempted不为空),那就准备将最终的running队列
        # 作为scheduleroutputs返回
        # ==============================================================================
        
        # Each sequence in the generation phase only takes one token slot.
        # Therefore, the number of batched tokens is equal to the number of
        # sequences in the RUNNING state.
        # 由于每个seq一次只生成1个token,因此num_batched_tokens = 状态为running的seq数量
        num_batched_tokens = sum(
            seq_group.num_seqs(status=SequenceStatus.RUNNING)
            for seq_group in self.running)

        # ==============================================================================
        # 构建Schduleroutputs
        # ==============================================================================
        scheduler_outputs = SchedulerOutputs(
            scheduled_seq_groups=[
                ScheduledSequenceGroup(seq_group=running_group,
                                       token_chunk_size=1)
                for running_group in self.running
            ],
            prompt_run=False,
            num_batched_tokens=num_batched_tokens,
            blocks_to_swap_in=blocks_to_swap_in,
            blocks_to_swap_out=blocks_to_swap_out,
            blocks_to_copy=blocks_to_copy,
            ignored_seq_groups=[],
        )
        return scheduler_outputs

五、总结

在本文中,我们:

  • 从vLLM批处理的入口函数开始,介绍了其推理内核LLMEngine的两个重要函数add_request()和step()
  • 在LLMEngine开始处理请求前(实例化阶段),它会先做一次模拟实验,来估计gpu上需要预留多少显存给KV Cache block。
  • 当LLMEngine开始处理请求时(add_request),它会把每个prompt当成一个请求,同时把它包装成一个SequenceGroup对象。
  • 当LLMEngine开始执行1次调度时(step),调度器策略(Scheduler)会根据实际gpu上KV Cache block的使用情况等要素,来选择要送哪些seq_group去做新一轮推理。注意,在1次推理中,所有seq_group要么一起做prefill,要么一起做decode。

到目前为止,我们遗留了以下问题

  • vLLM的物理块管理(block manager)的细节,包括物理块结构,逻辑块-物理块映射,物理块新增与释放,prefix caching等等
  • step()其余步骤:调度器只是决定了要送哪些seq_group去做推理,但是“每1个推理阶段结束后,如何根据output更新seq_group,将其送入下一次调度”这块不是调度器的职责,也是本文没涉及到的。

我们将在本系列后续的文章中,对块管理做详细讲解。在这之后,我们会分别以parallel sampling和beam search这两种decode方式为例,把整个流程传一遍,一起来更好理解vLLM背后的运作逻辑。

#Long-CLIP

本文介绍的工作在CLIP的基础上,提出了具有长文本能力的Long-CLIP,弥补了CLIP在长文本建模上的重大短板,并可以即插即用地利用在各种多模态任务中。

本文介绍了一个名为Long-CLIP的框架。Long-CLIP解决了CLIP有效长度不足、缺乏长文本能力的弊病,并在检索任务上获得了显著提升。此外,Long-CLIP保持了CLIP原始的特征空间,可以在图像生成等下游任务中即插即用地替换CLIP,以实现长文本细粒度图像生成。

图1 Long-CLIP使用场景总览

CLIP对齐了视觉与文本模态,拥有强大的zero-shot泛化能力。因此,CLIP被广泛应用在各种多模态任务中,如图像分类、文本图像检索、图像生成等。

然而,CLIP的一大弊病是在于长文本能力的缺失。首先,由于采用了绝对位置编码,CLIP的文本输入长度被限制在了77个token。不仅如此,实验发现CLIP真正的有效长度甚至不足20个token,远远不足以表征细粒度信息。文本端的长文本缺失也限制了视觉端的能力。由于仅包含短文本,CLIP的视觉编码器也只会提取一张图片中最主要的成分,而忽略了各种细节。这对跨模态检索等细粒度任务是十分不利的。同时,长文本的缺乏也使CLIP采取了类似bag-of-feature(BOF)的简单建模方式,不具备因果推理等复杂能力。

图2 CLIP的不足之处——缺乏长文本与复杂关系建模能力

为此,上海交通大学联合上海人工智能实验室的学者们提出了Long-CLIP模型。通过采用保留知识的位置编码扩充与加入核心成分对齐的微调策略,Long-CLIP模型仅仅额外采用ShareGPT4V数据集中的1M的(长文本,图片)数据对,通过不到100 GPU小时的微调,就可以在检索任务中获得显著提升(长文本-图像检索提升20%,短文本-图像检索提升6%)。不仅如此,Long-CLIP模型保持了CLIP原始的特征空间,因此可以在图像生成等下游任务上即插即用地替代原始的CLIP编码器,以实现长文本细粒度图像生成。以下,将介绍Long-CLIP采取的方法与应用场景。

  • 论文链接:https://arxiv.org/abs/2403.15378
  • 代码链接:https://github.com/beichenzbc/Long-CLIP

1.训练方法

一个简单的扩充输入长度、增强长文本能力的方法是先以固定的比率 λ₁ 对位置编码进行插值,再通过长文本进行微调。然而,这种策略会导致CLIP原始能力的急剧退化。前者破坏了CLIP充分建模的相对位置关系,而后者会使CLIP走入另一个极端:从仅仅关注最主要的特征变为以不论重要性一视同仁地涵盖所有细节。这导致模型在图片分类和短文本检索中的表现大幅下滑。

针对以上问题,研究者们提出了保留知识的位置编码扩充(Knowledge-Preserving Stretching of Positional Embedding)与加入核心成分对齐(Primary Component Matching)的微调策略,在保持甚至超过CLIP短文本能力的同时,解锁了了其长文本能力。

1.1保留知识的位置编码扩充

研究者们发现,CLIP的不同位置编码的训练程度是不同的。由于训练文本很可能以短文本为主,较低位的位置编码训练较为充分,能够精确地表征绝对位置,而较高位的位置编码则仅能表征其大致的相对位置。因此,对不同位置的编码进行插值的代价是不同的。基于以上观察,研究者保留了前20个位置编码,而对于剩下的57个位置编码,则以一个更大的比率 λ₂ 进行插值,计算公式可表示为:

实验表明,相较于直接插值,该策略可以在支持更长的总长度的同时大幅提升在各个任务上的性能。

1.2加入核心属性对齐的微调

仅仅引入长文本微调会使模型走入另一个误区,即一视同仁地囊括所有细节。针对这一问题,研究者们在微调中引入核心属性对齐这一策略。具体而言,研究者们利用主成分分析(PCA)算法,从细粒度的图像特征中提取核心属性,将其余属性过滤后重建粗粒度图像特征,并将其与概括性的短文本进行对齐。这一策略既要求模型不仅能够包含更多的细节(细粒度对齐),同时还能识别并建模其中最为核心的属性(核心成分提取与粗粒度对齐)。

图3 加入核心属性对齐的微调流程

2.应用场景

Long-CLIP在保留CLIP原始特征空间与能力的同时,大幅提升其长文本能力。因而,在图文检索、图像生成等领域,Long-CLIP可即插即用地替换CLIP。

2.1图文检索

Long-CLIP能够在图像与文本模态捕捉更多细粒度信息,从而可以增强相似图像和文本的区分能力,大幅提升图文检索的表现。无论是在传统的短文本检索(COCO、Flickr30k),还是在长文本检索任务上,我们的模型在召回率上均有显著提升。

图6 长文本-图像检索可视化,棕色文本为区分两张图片的关键细节

2.2图像生成

CLIP的文本编码器常被用于文本到图像生成模型中,如stable diffusion系列等。然而,由于长文本能力的缺失,用于生成图像的文本描述通常都十分简短,无法个性化地订制各种细节。而Long-CLIP无需任何训练,可以即插即用地替换CLIP作为文本编码器。Long-CLIP既可以突破77个token的限制,实现篇章级别的图像生成(右下),也可以在77个token内建模更多地细节,实现细粒度图像生成(右上)。而对于简单的短文本(左),由于Long-CLIP保持了CLIP的特征空间,可以和原始的CLIP生成相同的内容,图像质量不会产生退化。

图6 图像生成效果演示,棕色文本为CLIP生成时遗失的文本细节

结论:

本文介绍的工作在CLIP的基础上,提出了具有长文本能力的Long-CLIP,弥补了CLIP在长文本建模上的重大短板,并可以即插即用地利用在各种多模态任务中。

#OmniQuant

模型量化是模型压缩与加速中的一项关键技术,其将模型权重与激活值量化至低 bit,以允许模型占用更少的内存开销并加快推理速度。对于具有海量参数的大语言模型而言,模型量化显得更加重要。例如,GPT-3 模型的 175B 参数当使用 FP16 格式加载时,需消耗 350GB 的内存,需要至少 5 张 80GB 的 A100 GPU。

但若是可以将 GPT-3 模型的权重压缩至 3bit,则可以实现单张 A100-80GB 完成所有模型权重的加载。大语言模型权重、激活的全方位低bit可微量化,已集成进商用APP

现有的大语言模型后训练量化算法依赖于手工制定量化参数,优于缺乏相应的优化过程,导致面对低 bit 量化时,现有的方法都表现出显著的性能下降。尽管量化感知训练在确定最佳量化配置方面是有效的,但它需要引入大量额外的训练开销和训练数据。尤其是大语言模型本身的计算量进一步阻碍了量化感知训练在大预言模型量化上的应用。

这引出一个问题:我们能否在保持后训练量化的时间和数据效率的同时,达到量化感知训练的性能?

为了解决大语言模型后训练量化中的量化参数优化问题,来自上海人工智能实验室、香港大学、香港中文大学的研究者们提出了《OmniQuant: Omnidirectionally Calibrated Quantization for Large Language Models》。该算法同时支持大语言模型中的权重与激活值的量化,且覆盖多种量化 bit 位设置。

arXiv 论文地址:https://arxiv.org/abs/2308.13137

OpenReview 论文地址:https://openreview.net/forum?id=8Wuvhh0LYW

代码地址:https://github.com/OpenGVLab/OmniQuant

框架方法

如上图所示,OmniQuant 是一种针对大语言模型(LLM)的可微分量化技术,同时支持仅权重量化和权重激活值同时量化。并且,其在实现高性能量化模型的同时,保持了后训练量化的训练时间高效性和数据高效性。例如,OmniQuant 可在单卡 A100-40GB 上,在 1-16 小时内完成对 LLaMA-7B ~ LLaMA70B 模型量化参数的更新。

为了达到这个目标,OmniQuant 采用了一个 Block-wise 量化误差最小化框架。同时,OmniQuant 设计了两种新颖的策略来增加可学习的量化参数,包括可学习的权重裁剪(Learnable Weight Clipping,LWC),以减轻量化权重的难度,以及一个可学习的等价转换(Learnable Equivalent Transformation, LET),进一步将量化的挑战从激活值转移到权重。

此外,OmniQuant 引入的所有可学习参数在量化完成后可以被融合消除,量化模型可以基于现有工具完成在多平台的部署,包括 GPU、Android、IOS 等等。

Block-wise 量化误差最小化

OmniQuant 提出了一个新的优化流程,该流程采用 Block-wise 量化误差最小化,并且以可微分的方式优化额外的量化参数。其中,优化目标公式化如下:

可学习的权重裁剪 (LWC)

等价转换在模型权重和激活值之间进行量级迁移。OmniQuant 采用的可学习等价转换使得在参数优化过程中会使得模型权重的分布随着训练不断地发生改变。此前直接学习权重裁剪阈值的方法 [1,2] 只适用于权重分布不发生剧烈改变的情况,否则会难以收敛。基于此问题,与以往方法直接学习权重裁剪阈值不同,LWC 通过以下方式优化裁剪强度:

可学习的等价转换 (LET)

除了通过优化裁剪阈值来实现更适合量化的权重的 LWC 之外,OmniQuant 通过 LET 进一步降低激活值的量化难度。考虑到 LLM 激活值中的异常值是存在于特定通道,以前的方法如 SmoothQuant [3], Outlier Supression+[4] 通过数学上的等价转换将量化的难度从激活值转移到权重。

然而,手工选择或者贪心搜索得到的等价转换参数会限制量化模型的性能。得益于 Block-wise 量化误差最小化的引入,OmniQuant 的 LET 可以以一种可微分的方式确定最优的等价转换参数。受 Outlier Suppression+~\citep {outlier-plus} 的启发,采用了通道级的缩放和通道级的移位来操纵激活分布,为激活值中的异常值问题提供了一个有效的解决方案。具体来说,OmniQuant 探索了线性层和注意力操作中的等价转换。

其中 Q_a 是普通的 MinMax 量化器,Q_w 是带有可学习权重裁剪(即所提出的 LWC)的 MinMax 量化器。

注意力操作中的等价转换:除了线性层之外,注意力操作也占据了 LLM 的大部分计算。此外,LLM 的自回归推理模式需要为每个 token 存储键值(KV)缓存,这对于长序列来说导致了巨大的内存需求。因此,OmniQuant 也考虑将自主力计算中的 Q/K/V 矩阵量化为低位。具体来说,自注意力矩阵中的可学习等效变换可以写为:

伪代码

OmniQuant 的伪算法如上图所示。注意,LWC 与 LET 引入的额外参数在模型量化完后都可以被消除,即 OmniQuant 不会给量化模型引入任何额外开销,因此其可直接适配于现有的量化部署工具。

实验性能

上图显示了 OmniQuant 在 LLaMA 模型上仅权重量化结果的实验结果,更多 OPT 模型结果详见原文。可以看出,OmniQuant 在各种 LLM 模型(OPT、LLaMA-1、LLaMA-2)以及多样化的量化配置(包括 W2A16、W2A16g128、W2A16g64、W3A16、W3A16g128、W4A16 和 W4A16g128)中,始终优于以前的 LLM 仅权重量化方法。同时,这些实验表明了 OmniQuant 的通用性,能够适应多种量化配置。例如,尽管 AWQ [5] 在分组量化方面特别有效,但 OmniQuant 在通道级和分组级量化中均显示出更优的性能。此外,随着量化比特位数的减少,OmniQuant 的性能优势变得更加明显。   

在权重和激活值都量化的设置中中,实验主要关注点在于 W6A6 和 W4A4 量化。实验设置中排除了 W8A8 量化,因为与全精度模型相比,此前的 SmoothQuant 几乎可以实现无损的 W8A8 模型量化。上图显示了 OmniQuant 在 LLaMA 模型上权重激活值都量化结果的实验结果。值得注意的是,在 W4A4 量化的不同模型中,OmniQuant 显著提高了平均准确率,增幅在 + 4.99% ∼ +11.80% 之间。特别是在 LLaMA-7B 模型中,OmniQuant 甚至以 + 6.22% 的显著差距超越了最近的量化感知训练方法 LLM-QAT [6]。这一改进证明了引入额外可学习参数的有效性,这比量化感知训练所采用的全局权重调整更为有益。

同时,使用 OmniQuant 量化的模型可以在 MLC-LLM [7] 上实现无缝部署。上图展示了 LLaMA 系列量化模型在 NVIDIA A100-80G 上的内存需求和推理速度。

Weights Memory (WM) 代表量化权重存储,而 Running Memory (RM) 表示推理过程中的内存,后者更高是因为保留了某些激活值。推理速度是通过生成 512 个令牌来衡量的。显而易见,与 16 位全精度模型相比,量化模型显著减少了内存使用。而且,W4A16g128 和 W2A16g128 量化几乎使推理速度翻倍。

值得注意的是,MLC-LLM [7] 也支持 OmniQuant 量化模型在其余平台的部署,包括 Android 手机和 IOS 手机。如上图所示,近期的应用 Private LLM 即是利用 OmniQuant 算法来完成 LLM 在 iPhone、iPad,macOS 等多平台的内存高效部署。

总结

OmniQuant 是一种将量化推进到到低比特格式的先进大语言模型量化算法。OmniQuant 的核心原则是保留原始的全精度权重的同时添加可学习的量化参数。它利用可学习的权重才接和等价变换来优化权重和激活值的量化兼容性。在融合梯度更新的同时,OmniQuant 保持了与现有的 PTQ 方法相当的训练时间效率和数据效率。此外,OmniQuant 还确保了硬件兼容性,因为其添加的可训练参数可以被融合到原模型中不带来任何额外开销。

#Gemini1.5

今天,谷歌宣布推出 Gemini 1.5。

Gemini 1.5 建立在谷歌基础模型开发和基础设施的研究与工程创新的基础上,包括通过新的专家混合 (MoE) 架构使 Gemini 1.5 的训练和服务更加高效。

谷歌现在推出的是用于早期测试的 Gemini 1.5 的第一个版本 ——Gemini 1.5 Pro。它是一种中型多模态模型,针对多种任务的扩展进行了优化,其性能水平与谷歌迄今为止最大的模型 1.0 Ultra 类似,并引入了长上下文理解方面的突破性实验特征。

Gemini 1.5 Pro 配备了 128000 个 token 上下文窗口。但从今天开始,少数开发人员和企业客户可以通过 AI Studio 和 Vertex AI 的私人预览版在最多 100 万个 token 的上下文窗口中进行尝试。谷歌还进行了一些优化,以改善延迟、减少计算要求并增强用户体验。

谷歌 CEO Sundar Pichai 和谷歌 DeepMind CEO Demis Hassabis 对新模型进行了专门介绍。

领先基础模型的上下文长度

高效架构

Gemini 1.5 建立在谷歌对 Transformer 和 MoE 架构的领先研究之上。传统 Transformer 充当一个大型神经网络,而 MoE 模型则分为更小的 “专家” 神经网络。

根据给定输入的类型,MoE 模型学会选择性地仅激活其神经网络中最相关的专家路径。这种专业化极大地提高了模型的效率。通过稀疏门控 MoE、GShard-Transformer、Switch-Transformer、M4 等研究,Google 一直是深度学习 MoE 技术的早期采用者和先驱。

谷歌在模型架构方面的最新创新使 Gemini 1.5 能够更快地学习复杂任务并保持质量,同时更高效地训练和服务。这些效率正在帮助谷歌团队比以往更快地迭代、培训和交付更高级的 Gemini 版本,并且正在努力进一步优化。

更长的上下文,更有用的功能

人工智能模型的 “上下文窗口” 由 token 组成,token 是用于处理信息的构建块。token 可以是文字、图像、视频、音频或代码的整个部分或子部分。模型的上下文窗口越大,它在给定提示中可以接收和处理的信息就越多,从而使其输出更加一致、相关和有用。

通过一系列机器学习创新,谷歌增加了 1.5 Pro 的上下文窗口容量,远远超出了 Gemini 1.0 最初的 32,000 个 token。该大模型现在可以在生产环境中运行多达 100 万个 token。

这意味着 1.5 Pro 可以一次性处理大量信息,包括 1 小时的视频、11 小时的音频、超过 30,000 行代码或超过 700,000 个单词的代码库。在谷歌的研究中,还成功测试了多达 1000 万个 token。

对大量信息进行复杂推理

1.5 Pro 可以在给定提示内无缝分析、分类和总结大量内容。例如,当给出阿波罗 11 号登月任务的 402 页记录时,它可以推理整个文档中的对话、事件和细节。 

Gemini 1.5 Pro 可以理解、推理和识别阿波罗 11 号登月任务的 402 页记录中的好奇细节。

更好地理解和推理跨模态

1.5 Pro 可以针对包括视频在内的不同模式执行高度复杂的理解和推理任务。例如,当给定一部 44 分钟的巴斯特・基顿无声电影时,该模型可以准确分析各种情节点和事件,甚至推理出电影中容易被忽略的小细节。

当给出简单的线条图作为现实生活中物体的参考材料时,Gemini 1.5 Pro 可以识别 44 分钟的巴斯特基顿无声电影中的场景。

使用较长的代码块解决相关问题

1.5 Pro 可以跨较长的代码块执行更相关的问题解决任务。当给出超过 100,000 行代码的提示时,它可以更好地推理示例、建议有用的修改并解释代码不同部分的工作原理。

Gemini 1.5 Pro 可以推理 100,000 行代码,提供有用的解决方案、修改和注释

增强性能

在文本、代码、图像、音频、视频评估综合面板上进行测试时,1.5 Pro 在用于开发大型语言模型 (LLM) 的基准测试中,87% 的性能优于 1.0 Pro。在相同的基准测试中与 1.0 Ultra 相比,它的表现大致相似。

即使上下文窗口增加,Gemini 1.5 Pro 仍能保持高水平的性能。

在 NIAH 评估中,故意将包含特定事实或陈述的一小段文本放置在很长的文本块中,1.5 Pro 99% 的时间都能找到嵌入的文本,在数据块中如下只要 100 万个 token。

Gemini 1.5 Pro 还展示了令人印象深刻的 “上下文学习(in-context learning)” 技能,这意味着它可以从长提示中给出的信息中学习新技能,而不需要额外的微调。谷歌在 MTOB (Translation from One Book )基准测试中测试了这项技能,该基准显示了该模型从以前从未见过的信息中学习的能力。当给定卡拉芒语(一种全球使用人数不足 200 人的语言)的语法手册时,该模型可以学习将英语翻译成卡拉芒语,其水平与学习相同内容的人相似。

由于 1.5 Pro 的长上下文窗口是大型模型中的首创,因此谷歌正在不断开发新的评估和基准来测试其新颖的功能。

有关更多详细信息,请参阅 Gemini 1.5 Pro 技术报告。

技术报告地址:https://storage.googleapis.com/deepmind-media/gemini/gemini_v1_5_report.pdf

使用 Gemini 模型进行构建和实验

谷歌致力于负责任地将每个新一代 Gemini 模型带给全球数十亿人、开发者和企业用户使用。

从今天开始,谷歌将通过 AI Studio 和 Vertex AI 向开发者和企业客户提供 1.5 Pro 预览版。

未来,当模型进行更广泛的发布时,届时,谷歌将推出具有标准 128,000 个 token 上下文窗口的 1.5 Pro。很快,随着谷歌对模型的改进,谷歌计划引入从标准 128,000 个上下文窗口开始并扩展到 100 万个 token 的定价等级。

早期测试人员可以在测试期间免费尝试 100 万个 token 上下文窗口,速度的显着提高也即将到来。

有兴趣测试 1.5 Pro 的开发人员现在可以在 AI Studio 中注册,而企业客户可以联系他们的 Vertex AI 客户团队。

参考链接:https://blog.google/technology/ai/google-gemini-next-generation-model-february-2024/#sundar-note

#SimPO

一种简单却有效的离线偏好优化算法。 陈丹琦团队提出简单偏好优化SimPO,还炼出最强8B开源模型

为了将大型语言模型(LLM)与人类的价值和意图对齐,学习人类反馈至关重要,这能确保它们是有用的、诚实的和无害的。在对齐 LLM 方面,一种有效的方法是根据人类反馈的强化学习(RLHF)。尽管经典 RLHF 方法的结果很出色,但其多阶段的过程依然带来了一些优化难题,其中涉及到训练一个奖励模型,然后优化一个策略模型来最大化该奖励。

近段时间已有一些研究者探索了更简单的离线算法,其中之一便是直接偏好优化(DPO)。DPO 是通过参数化 RLHF 中的奖励函数来直接根据偏好数据学习策略模型,这样就无需显式的奖励模型了。该方法简单稳定,已经被广泛用于实践。

使用 DPO 时,得到隐式奖励的方式是使用当前策略模型和监督式微调(SFT)模型之间的响应似然比的对数 的对数比。但是,这种构建奖励的方式并未与引导生成的指标直接对齐,该指标大约是策略模型所生成响应的平均对数似然。训练和推理之间的这种差异可能导致性能不佳。

为此,弗吉尼亚大学的助理教授孟瑜与普林斯顿大学的在读博士夏梦舟和助理教授陈丹琦三人共同提出了 SimPO—— 一种简单却有效的离线偏好优化算法。

  • 论文标题:SimPO: Simple Preference Optimization with a Reference-Free Reward
  • 论文地址:https://arxiv.org/pdf/2405.14734
  • 代码 & 模型:https://github.com/princeton-nlp/SimPO

该算法的核心是将偏好优化目标中的奖励函数与生成指标对齐。SimPO 包含两个主要组件:(1)在长度上归一化的奖励,其计算方式是使用策略模型的奖励中所有 token 的平均对数概率;(2)目标奖励差额,用以确保获胜和失败响应之间的奖励差超过这个差额。

总结起来,SimPO 具有以下特点:

  • 简单:SimPO 不需要参考模型,因此比 DPO 等其它依赖参考模型的方法更轻量更容易实现。
  • 性能优势明显:尽管 SimPO 很简单,但其性能却明显优于 DPO 及其最新变体(比如近期的无参考式目标 ORPO)。如图 1 所示。并且在不同的训练设置和多种指令遵从基准(包括 AlpacaEval 2 和高难度的 Arena-Hard 基准)上,SimPO 都有稳定的优势。
  • 尽量小的长度利用:相比于 SFT 或 DPO 模型,SimPO 不会显著增加响应长度(见表 1),这说明其长度利用是最小的。

该团队进行了大量分析,结果表明 SimPO 能更有效地利用偏好数据,从而在验证集上对高质量和低质量响应的似然进行更准确的排序,这进一步能造就更好的策略模型。

如表 1 所示,该团队基于 Llama3-8B-instruct 构建了一个具有顶尖性能的模型,其在 AlpacaEval 2 上得到的长度受控式胜率为 44.7,在排行榜上超过了 Claude 3 Opus;另外其在 Arena-Hard 上的胜率为 33.8,使其成为了目前最强大的 8B 开源模型。

SimPO:简单偏好优化

为便于理解,下面首先介绍 DPO 的背景,然后说明 DPO 的奖励与生成所用的似然度量之间的差异,并提出一种无参考的替代奖励公式来缓解这一问题。最后,通过将目标奖励差额项整合进 Bradley-Terry 模型中,推导出 SimPO 目标。

背景:直接偏好优化(DPO)

DPO 是最常用的离线偏好优化方法之一。DPO 并不会学习一个显式的奖励模型,而是使用一个带最优策略的闭式表达式来对奖励函数 r 进行重新参数化:

其中 (x, y_w, y_l) 是由来自偏好数据集 D 的 prompt、获胜响应和失败响应构成的偏好对。

一种与生成结果对齐的简单无参考奖励

DPO 的奖励与生成之间的差异。使用 (1) 式作为隐式的奖励表达式有以下缺点:(1) 训练阶段需要参考模型 π_ref,这会带来额外的内存和计算成本;(2) 训练阶段优化的奖励与推理所用的生成指标之间存在差异。具体来说,在生成阶段,会使用策略模型 π_θ 生成一个能近似最大化平均对数似然的序列,定义如下:

其中 β 是控制奖励差异大小的常量。该团队发现,根据响应长度对奖励进行归一化非常关键;从奖励公式中移除长度归一化项会导致模型倾向于生成更长但质量更低的序列。这样一来,构建的奖励中就无需参考模型了,从而实现比依赖参考模型的算法更高的内存和计算效率。

SimPO 目标

目标奖励差额。另外,该团队还为 Bradley-Terry 目标引入了一个目标奖励差额项 γ > 0,以确保获胜响应的奖励 r (x, y_w) 超过失败响应的奖励 r (x, y_l) 至少 γ:

两个类之间的差额已知会影响分类器的泛化能力。在使用随机模型初始化的标准训练设置中,增加目标差额通常能提升泛化性能。在偏好优化中,这两个类别是单个输入的获胜或失败响应。

在实践中,该团队观察到随着目标差额增大,生成质量一开始会提升,但当这个差额变得过大时,生成质量就会下降。DPO 的一种变体 IPO 也构建了与 SimPO 类似的目标奖励差额,但其整体目标的效果不及 SimPO。

目标。最后,通过将 (4) 式代入到 (5) 式中,可以得到 SimPO 目标:

总结起来,SimPO 采用了与生成指标直接对齐的隐式奖励形式,从而消除了对参考模型的需求。此外,其还引入了一个目标奖励差额 γ 来分离获胜和失败响应。

实验设置

模型和训练设置。该团队的实验使用了 Base 和 Instruct 两种设置下的两类模型 Llama3-8B 和 Mistral-7B。

评估基准。该团队使用了三个最常用的开放式指令遵从基准:MT-Bench、AlpacaEval 2 和 Arena-Hard v0.1。这些基准可评估模型在各种查询上的多样化对话能力,并已被社区广泛采用。表 2 给出了一些细节。

基线方法。表 3 列出了与 SimPO 做对比的其它离线偏好优化方法。

实验结果

主要结果与消融研究

SimPO 的表现总是显著优于之前已有的偏好优化方法。如表 4 所示,尽管所有的偏好优化算法的表现都优于 SFT 模型,但简单的 SimPO 却在所有基准和设置上都取得了最佳表现。这样全面的大幅领先彰显了 SimPO 的稳健性和有效性。

基准质量各不相同。可以观察到,在 Arena-Hard 上的胜率明显低于在 AlpacaEval 2 上胜率,这说明 Arena-Hard 是更困难的基准

Instruct 设置会带来显著的性能增益。可以看到,Instruct 设置在所有基准上都全面优于 Base 设置。这可能是因为这些模型使用了更高质量的 SFT 模型来进行初始化以及这些模型生成的偏好数据的质量更高。

SimPO 的两种关键设计都很重要。表 5 展示了对 SimPO 的每种关键设计进行消融实验的结果。(1) 移除 (4) 式中的长度归一化(即 w/o LN);(2) 将 (6) 式中的目标奖励差额设置为 0(即 γ = 0)。

移除长度归一化对结果的影响最大。该团队研究发现,这会导致模型生成长且重复的模式,由此严重拉低输出的整体质量。将 γ 设为 0 也会导致 SimPO 的性能下降,这说明 0 并非最优的目标奖励差额。

有关这两项设计选择的更深度分析请参阅原论文。

深度对比 DPO 与 SimPO

最后,该团队还从四个角度全面比较了 DPO 与 SimPO:(1) 似然 - 长度相关性、(2) 奖励构建、(3) 奖励准确度、(4) 算法效率。结果表明 SimPO 在准确度和效率方面优于 DPO。

DPO 奖励会隐式地促进长度归一化。

DPO 奖励与生成似然不匹配。

DPO 在奖励准确度方面不及 SimPO。

图 4c 比较了 SimPO 和 DPO 的奖励准确度,这评估的是它们最终学习到的奖励与留存集上的偏好标签的对齐程度。可以观察到,SimPO 的奖励准确度高于 DPO,这说明 SimPO 的奖励设计有助于实现更有效的泛化和更高质量的生成。

SimPO 的内存效率和计算效率都比 DPO 高。

SimPO 的另一大优势是效率,毕竟它不使用参考模型。图 4d 给出了在 8×H100 GPU 上使用 Llama3-Base 设置时,SimPO 和 DPO 的整体运行时间和每台 GPU 的峰值内存使用量。相比于原版 DPO 实现,得益于消除了使用参考模型的前向通过,SimPO 可将运行时间降低约 20%,将 GPU 内存使用量降低约 10%。

#指令微调~

自 ChatGPT 等大型语言模型推出以来,为了提升模型效果,各种指令微调方法陆续被提出。本文中,普林斯顿博士生、陈丹琦学生高天宇汇总了指令微调领域的进展,包括数据、算法和评估等。

大型语言模型(LLM)很强大,但要想真正帮助我们处理各种日常和工作任务,指令微调就必不可少了。近日,普林斯顿大学博士生高天宇在自己的博客上总结了指令微调研究方向的近期进展并介绍了其团队的一项近期研究成果。

具有十亿级参数且使用万亿级 token 训练的大型语言模型(LLM)非常强大,直接就能用于解决大量不同的任务。但是,要用于真实世界应用以及作为通用任务求解机,LLM 就必须学会遵从用户指令并以一种连贯且有用的方式进行响应,而不是仅仅作为一只「随机鹦鹉」,学舌来自互联网的混乱语言模式。

因此,开放式指令微调(InstructGPT)变成了一种颇具潜力的方法,这种方法的目标是让 LLM 能遵从用户指令并以一种有助益、诚实且无害(即 Anthropic 的 HHH 指标)的方式给出响应。ChatGPT 取得巨大成功之后,人们对指令微调的兴趣进一步提升。开放式指令微调通常包含两个阶段:

  • 基于收集到的用户指令和标准响应对模型进行监督式微调(SFT)。   
  • 将模型与人类偏好对齐(这方面的主要方法是根据人类反馈的强化学习 / RLHF)。这通常需要人类偏好数据,其中包含一对响应以及一个标注(表明哪个响应更好)。

众所周知,收集监督式微调或偏好数据的成本非常高,因此一直以来都只有大企业承担得起;直到 2023 年,人们找到了更低成本的构建此类数据的方法。自此,许多用于开发指令微调模型的开源项目应势而生。下面将分四部分介绍这些项目:SFT 数据、偏好数据、算法和评估。最后,作者还将介绍他们在指令遵从评估方面的最新研究成果,其中表明:设置正确的评估器很重要,否则就可能得到误导性的结果。

监督式微调(SFT)数据

一般来说,监督式微调的目的有两种,分别对应于两种类型的数据。一种是进一步提升 LLM 的一般语言理解能力,HellaSwag 和 HellaSwag 等传统 NLP 基准便是为了这样的目的。另一种则是为了让 LLM 遵从指令、获得对话能力、变得有用和无害。

与第一种目的相对应的是多任务指令微调数据,这在 2020-2022 年间得到了人们的大力探索。这些数据是将数以千计的 NLP 任务组合起来并为每一个任务提供一个自然语言指令,然后人们就能在这个组合上以多任务的方式训练模型。Sebastian Ruder 在其博客中进行了更透彻详细的回顾。

博客地址:https://nlpnewsletter.substack.com/p/instruction-tuning-vol-1

代表性的数据集包括 Natural Instruction、T0、Flan。不同于开放式指令微调,这些数据集 / 模型更面向传统 NLP 任务(问答、自然语言推理等),其中的指令往往更短 / 更简单 / 种类更少 —— 想象一下「对这个句子进行情绪分类」与「使用 Jekyll 给我写一个与 OpenAI 博客风格类似的个人网页」之间的区别。因此,在这些数据集上训练的模型往往无法部署成如今的「指令微调」模型或聊天机器人,即便它们在 NLP 基准上表现很好。

Wang et al., 2023 (TÜLU) 表明:如果将这些数据集与新的开放式指令微调数据集组合起来,可以提升模型的一般语言理解能力和指令遵从能力。Mukherjee et al., 2023 (Orca) 发现如果使用这些数据作为种子,提示 GPT-4 来输出带解释的答案以及模仿 GPT-4 的响应,可以显著提升较弱模型的性能。Stable Beluga 等一些公共指令微调模型便采用了这种数据混合方法。

下面来谈谈开放式指令微调数据,其在 2023 年尤为兴盛(下面将使用 SFT 数据指代开放式指令微调数据)。人们普遍相信使用这些数据训练不会提升 LLM 的「知识」(反映为在传统基准上的分数),而只是会「引导」它们遵从指令遵从或对话格式、获取引人互动的语调、变得有礼貌等等(表面对齐假设;Zhou et al., 2023 的《LIMA: Less Is More for Alignment》)。

收集 SFT 数据的成本很高,因为这既需要收集用户指令,也需要标注演示数据。对开源模型而言,一种获取开放式指令微调数据的方法是从专有 LLM 中蒸馏获取。

最早的开源指令模型之一 Alpaca 使用了自指示(self-instruct)来为 text-davinci-003 (InstructGPT 模型的一个变体)构建 prompt 并生成伪 SFT 数据,然后再在其上对 LLaMA-7B 模型进行监督式指令微调;白泽(Baize)也使用了自指示,但却是通过让 ChatGPT 自我聊天来获取多轮数据;WizardLM 提升数据多样性的方法是使用 ChatGPT 来迭代式地重写 Alpaca 数据;UltraChat 先是使用不同的策略来自动构建问题,然后使用 ChatGPT 根据给定问题模拟对话。

Vicuna 和 Koala 探索了 ShareGPT,这是一个用户分享自己与 ChatGPT 的聊天记录的网站(https://sharegpt.com )—— 这些记录可作为 SFT 数据。近期还有一项类似的工作 WildChat,其会为在线用户提供免费的 ChatGPT 并收集对话记录,但其重心更偏向于研究有毒用例。尽管这是一种相对低成本的获取数据的方式,但有研究发现模仿专有 LLM 只能「模仿 ChatGPT 的风格而不是其事实性」,因此开源模型的能力范围就完全仰赖这些 SFT 数据了。

另一种收集 SFT 数据的方法是人工标注少量数据。Open Assistant 发起过一个请志愿者编写指令和响应的众包项目;Dolly 包含 1.5 万条 Databricks 的员工构建的数据(更倾向于基于维基百科的事实性问答)。LIMA(less is more for alignment)中包含 1000 条作者人工调整过的 SFT 数据(其分布严重倾向于 Stack Exchange 和 wikiHow),人们惊讶地发现其可有效地用于得到强大的指令模型。但是,我们仍然不清楚:相比于使用专门收集的大规模数据,我们是只需要 1000 个示例,还是可以使用互联网众包的数据,因为这方面还没有专门的对比研究。

尽管在这些模仿和人类 SFT 数据上训练的开源模型仍旧无法媲美 ChatGPT、GPT-4 或 Claude 等专有模型,但我们可以看到两个颇具希望的成果:

  1. 人类评估表明,LLaMA-2-70B-chat 比 ChatGPT 「更有帮助」,而前者是在闭源数据上微调过的开源模型 LLaMA-2-70B。这表明 LLaMA-2 基础模型有可能与 ChatGPT 的基础模型一样强大(在事实知识、常识、推理能力方面),从而可减轻「错误承诺」问题。
  2. (研究社区已经在「toy」或「实验室」数据上进行了一些激动人心的研究,比如更好的对齐算法。下面会谈到它们。

偏好数据又如何?

尽管开源 SFT 模型能提供让人印象深刻的「幻象」(事实上,它们启动了开源社区研发指令微调的趋势),但仅仅有 SFT 是不够的。要让模型成为更优秀的语言助手,将模型与人类偏好对齐至关重要。推断它的一种简单方法是思考如何「变得诚实」。SFT 几乎总是会鼓励模型给出答案,几乎不会教模型说「我不知道这一点」。

一些研究已经表明,对齐算法可以带来更好的「人类满意度」。但是,大多数开源模型并不会经历对齐阶段(RLHF),原因包括 (1) 强化学习的成本很高,(2) 调整 PPO(OpenAI 使用的强化学习算法)超参数时模型很脆弱,(3) 缺乏高质量偏好数据。数据的缺乏进一步阻碍了社区创造比强化学习更有效 / 更高效的(可能存在的)更优算法。

在开发对齐算法时,最常用的两个偏好数据集是 OpenAI 的 TL;DR 偏好数据集(摘要)和 Anthropic 的 HH-RLHF 数据集(人类 - 模型开放式对话)。尽管它们的质量都不错,但其指令的多样性和复杂性还不足以比肩如今的 SFT 数据集。

2023 年出现了许多新的偏好数据集。尽管对研究者而言它们可能是宝贵的资源,但我们还不清楚它们的质量是否足以用于对齐算法。其中一些是通过众包方式从人类收集的偏好数据,Open Assistant 和 Chatbot Arena 都在网上发起了一个偏好数据收集倡议,收集志愿者提供的偏好标签。

更多数据集采用的是模拟或启发式方法:SHP 是使用在 Reddit 上的点赞数而启发式地构建的一个合成偏好数据集;AlpacaFarm 和 UltraFeedback 使用了 GPT-4 作为标准标注者;Kim et al., 2023 的《Aligning Large Language Models through Synthetic Feedback》、Xu et al., 2023 的《Contrastive Post-training Large Language Models on Data Curriculum》、Yang et al., 2023 的《RLCD: Reinforcement Learning from Contrast Distillation for Language Model Alignment》使用的启发式方法包括:更强模型的输出结果应该更受偏爱,或来自一个「好」prompt 的输出结果应该更受偏爱。

有证据表明这里提到的大多数数据集都有助于强化学习或其它对齐算法,但尚无研究者对它们进行基准对比。Huggingface 近期发布了一个模型,其训练使用了 UltraChat(SFT)和 UltraFeedback(对齐,使用 DPO);结果表明其性能与使用闭源数据训练的 LLaMA-2-Chat-70B 相当。

不同于依靠人类偏好,另一个研究策略是使用「AI 反馈」—— 使用 LLM 来引导 LLM,而没有人类参与其中。这一思想不同于「使用 GPT-4 作为标注者」,因为 GPT-4 仍旧是使用人类偏好训练的,但这里的目标是在没有人类偏好数据的前提下用模型来引导。

Bai et al., 2022 的《Constitutional AI: Harmlessness from AI Feedback》最早提出了两个概念:「宪法 AI(Constitutional AI)」和「根据人工智能反馈的强化学习(RLAIF)」。

Constitutional AI 会定义一系列「原则」,其中包括好的生成结果应当遵守的原则以及为 SFT 模型提供 prompt 使之自我提升生成结果的原则(通过自我批评和修正)。

RLAIF 则是让 SFT 模型(而非人类)为输出结果对生成偏好。他们通过实验表明,如果一开始训练模型时仅使用「有助益的」人类监督,那么就有可能训练出「无害的」模型(没有在无害性上的人类监督)。

但是,Bai et al., 2022 中的流程一开始依然使用了一些人类偏好标签。Lee et al., 2023 的《RLAIF: Scaling Reinforcement Learning from Human Feedback with AI Feedback》表明:如果从一个 SFT 模型开始,RLAIF 能在一个摘要任务上达到与 RLHF 相当的水平,期间不涉及人类偏好标签。

RLAIF 这个研究方向引起了人们的极大热情和兴趣,因为其是「可扩展监督(scalable oversight)」的一个可行解决方案。可扩展监督面向的是要对齐的模型已超越人类能力的情形。但是,我们仍旧不清楚这些方法的表现究竟有多好,因为使用简单的启发式方法来构建数据(也没有人类参与)也能超越它们(RLCD)。

强化学习是唯一的方法吗?

使用 PPO 来执行 RLHF 已经成为对齐的主要方法,比如 InstructGPT 和 LLaMA-2-Chat 都使用了它,据信 ChatGPT 和 GPT-4 也使用了这种方法。其基本思想是首先在偏好数据上训练一个奖励模型,然后使用该奖励模型来提供反馈并使用强化学习对模型进行微调。

有关 RLHF 的文献汗牛充栋,HuggingFace 的这篇博客提供了更多细节:https://huggingface.co/blog/rlhf

RLHF 是有效的,但实现起来很复杂,容易出现优化不稳定问题,而且对超参数敏感。令人兴奋的是,研究者们已经提出了一些可将模型与偏好数据对齐的新方法,而且据称其中一些方法强于 RLHF。

best-of-n,即 n 个结果中取最佳。我们可以有一个直觉认识:在 SFT 之后,模型已经有希望生成好的输出结果了,我们只需要把它们找出来。在 WebGPT (Nakano et al., 2021)和根据人类反馈的摘要(Stiennon et al., 2022)中,作者探索了 best-of-n 采样法,即采样 n 个输出结果并使用奖励模型选出其中最好的一个,结果表明这种方法往往能实现与 RLHF 相近的性能。但是,OpenAI 指出,如果最终的最优策略与原始的 SFT 模型相距很远(n 显著增至最终策略与 SFT 模型之间的 KL),则 best-of-n 的效率会很低;更不要说就算 n 很小,其推理的效率也非常低。

专家迭代。还有在训练中使用 best-of-n 的方法 —— 在训练过程中大量采样(不涉及推理效率),选出最好的,然后基于它们执行 SFT。FeedME 是使用其自身采样的且人类标注者喜欢的输出来训练模型;OpenAI 使用该方法训练了 text-davinci-002。再进一步,这可以和在线采样 best-of-n 组合起来,即采样 n 个输出,由奖励模型选出其中最好的,然后在最好的输出上训练,如此反复。这本质上就是专家迭代。best-of-n 采样还能与自然语言反馈组合使用。

条件 token。另一个思路是「条件 token」:同时使用好示例和差示例在语言模型上执行 SFT,并在好示例前面加上「好」prompt,在差示例前面加上「差」prompt。在推理过程中,可以通过添加前缀「好」来为模型设定条件,并期望模型生成好的输出。

对比式方法。最后,还有一些新提出的方法非常近似对比学习思想:可以从模型获得好示例和差示例的概率,然后「促进」好示例,「压制」差示例。给定偏好数据,SLiC 和 RRHF 的做法是同时优化一个对比排名损失和一个正则化损失。举个例子,如下是 SLiC 损失函数:

其中 r_ϕ (x,y) 是奖励模型,σ(⋅) 是 sigmoid 函数,π_θ 是当前模型,π_ref 是 SFT 模型。

这些模型有一个缺点:它们要么从 SFT 模型中采样 y_w、y_l,要么就是直接从已有数据集中取出它们(因此是从其它模型采样的),由此会造成分布不匹配问题。Liu et al., 2023 (RSO) 提出可通过从最优策略 π^∗ 采样来解决这个问题 —— 通过使用奖励模型执行拒绝采样(reject sampling)。他们的研究表明,在 SLiC 或 DPO 上应用这样的采样策略可以提升最终模型的性能。

这些方法在近期吸引了不少关注,并且已经被多个研究团队证明是有效的。比如 Xu et al., 2023 的《Contrastive Post-training Large Language Models on Data Curriculum》表明 DPO 可以在 SFT 之上带来显著的提升,HuggingFace 的 Zephyr 模型也使用了 DPO 训练,其在 MT-Bench 上表现很好,甚至可以比肩 Llama-2-chat 和 GPT-3.5。所有这些方法的成本都比强化学习低很多,对研究和开源社区来说,这是个好消息,并且这也有望激励人们创造出更多更好的对齐算法。

另一方面,我们需要更好地理解使用对齐算法训练的模型的性质以及他们是否真的有助于学习有用的特征,举个例子,Singhal et al., 2023 的《A Long Way to Go: Investigating Length Correlations in RLHF》研究了几个常用数据集后发现:所学习到的奖励模型的偏好通常与长度高度相关,而使用长度进行 RLHF 就能恢复大部分性能提升。

评估

在开发开放式指令微调模型(或任何开放式生成方法)方面,一大重要难题是评估。人类评估依然是评估开放式对话模型的能力的「黄金标准」。但是,人类评估往往不可靠,尤其是当使用了 Amazon Mechanical Turk 等廉价众包平台来获取数据时。此外,人类评估的成本很高,也难以执行统一基准的比较。最近人们开始使用更强的 LLM(如 ChatGPT 或 GPT-4)来评估更弱的 LLM(如基于开源的 LLaMA 的模型),事实证明这是一个受欢迎的具有成本效益的替代方案。

乍一看,用模型评估模型听起来很荒谬。但是,对于开发开源模型和研究模型来说,这是有意义的:GPT-4 等专有模型的训练使用了远远更为强大的基础模型,并且其使用的指令数据的质量和数量都高得多,因此它们会比开源或研究模型更优秀。只要它们的能力存在巨大差异,GPT-4 这样的模型就足以胜任评估器。

一些使用 LLM 作为评估器的先驱研究给出了「让人心安的」结果:LLM 评估器通常与人类评估具有很高的一致性。另一方面,一些论文则表明 LLM 评估器往往对某些偏见非常敏感。

举个例子,如果你交换两个要比较的输出的位置,它们就可能改变自己的偏好。LLM 评估器也更偏爱更长的输出以及相似模型生成的输出。因此,人们提出了一些「元评估(meta-evaluation)」基准,用以评估 LLM 评估器的质量(通常的形式是看在人类偏好数据上的准确度),其中包括 FairEval、MT-Bench、LLMEval^2。尽管这些是帮助我们理解 LLM 评估器可靠程度的宝贵资源,但不同的评估器在这些基准上的分数往往差不多。

此外,这些基准的人类标注往往很多噪声且很主观,并且内在的人类一致率(human agreement rate)较低,比如 AlpacaFarm 报告的人类一致率为 66%、MT-Bench 报告的数据为 63%、FairEval 的为 71.7%。这样一来,我们就不清楚是否可以信任这些元评估基准和 LLM 评估器了。

LLMBar:LLM 评估器的更优元评估

最后,作者介绍了自己团队的一项研究成果《Evaluating Large Language Models at Evaluating Instruction Following》。他们在其中重新思考了元评估问题。他们认为之前的研究忽视了一个重要因素:人类偏好的固有主观性。

对于来自之前一个数据集的上述示例,即使这两者之间的质量差异可以区别,但人类标注者还是更偏爱更长的,这就将这种偏见加到了偏好数据集中。当我们基于这样的主观且带噪声的元基准评估 LLM 评估器时,我们无法保证得分高的评估器能可靠地评估输出长度等主观偏好之外的客观属性,比如指令遵从或事实正确性。

遵循这一路径,他们创造了一个新的元评估基准 LLMBar,其关注重点是一个客观指标 —— 指令遵从。他们选择指令遵从的原因包括:(1) 该能力能以客观的方式进行评估;(2) 其与有用性等人们期望获得的 LLM 性质直接相关;(3) 不同于可通过模仿学习轻松实现的表面质量,当前最强大的 LLM 也难以应对这一问题。下面给出了一个来自 LLMBar 的示例: 

即使很明显右侧的输出遵从指令,人类和 LLM 评估者都往往更偏爱左侧的输出,因为其语调让人更有参与感。如果我们不对评估者的能力进行严格的分析,以便分辨真正的指令遵从能力和表面线索,那就会存在这样的风险:先进模型只是擅长模仿对话助理,而不是执行所需任务。

LLMBar 的作者手动构建了 419 个实例,其中每一项都包含一个指令与配对的两个输出:一个忠实地遵从指令,另一个则偏离了,并且总是会存在一个客观偏好。得益于这个客观指标和手动调整,LLMBar 的人类一致率达到了 94%。他们在这些输出对上测试各种评估器,并比较了评估器偏好和标准参考标签。他们还精心构建了一个对抗集,其中的「差」输出往往具有一些表面上的吸引力(长度、参与性语调、由更好的语言模型生成等),可能会误导评估器。LLMBar 的结果令人惊讶:

尽管 ChatGPT、LLaMA2-70B-Chat、PaLM2-bison 和 GPT-4 在其它元评估基准上的表现相近,但它们在 LLMBar (adversarial) 上的表现则大不相同 ——ChatGPT 和 LLaMA2 的分数甚至低于随机乱猜,而 GPT-4 的准确度则远远胜过其它任何评估器。

除了不同的 LLM,他们的研究还表明不同的 prompt 也对评估器非常重要。之前探索这一方向的研究包括:Wang et al., 2023 的《Large Language Models are not Fair Evaluators》提出采样多个解释并将它们聚合成一个最终判断;Zheng et al., 2023 的《Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena》则提出了一种参照引导式方式,即先让 LLM 评估器根据指令生成自己的输出,然后将其用作参照;还有一些论文表明:部署多个评估器(不同的 LLM 或 prompt)并让它们互相交流或合成它们的判断,可以提升评估器的准确度。作者团队则提出将这些方法组合起来:指标 + 参照 + 规则(如下所示)。

首先,使用 prompt 让 LLM 生成三个特定于指令的指标,也就是 rubric,并还使用 prompt 让 LLM 生成一个参照输出。然后,将这些指标和参照输入 LLM,显式地列出规则(比如重点是指令遵从,忽略位置偏见等),并让模型给出判断。相比于 AlpacaFarm 使用的最基本的 prompt,新的 prompt 可以显著提升在 LLMBar 上的评估器性能(GPT-4 在对抗集上的性能可提升 10%)。研究者还进行了更多消融研究,并还有另一些有趣的结果,比如这样一个反直觉的发现:思维链大多数时候都会损害评估器准确度。

展望未来

开源指令微调数据、算法和模型的涌现是 2023 年 LLM 领域最激动人心的进展之一。这让研究者和开发者有机会将指令模型的训练、评估、互动和分析置于自己的完全控制之下(从参数到数据),之前这些都是黑盒。

过去几个月,这一领域也有点混乱,因为有数以百计的论文发布的结果使用了不同的数据、算法、基础模型甚至评估方法,使得人们难以对这些文献进行交叉比对。作者表示希望社区能很快聚合出一个标准的数据集 / 评估方法,让研究者能以一种更科学和可复现的方法开发出更好的指令微调模型。

博客地址:https://gaotianyu.xyz/blog/2023/11/30/instruction-tuning/

#Weaver

ChatGPT 等通用大模型支持的功能成百上千,但是对于普通日常用户来说,智能写作一定是最常见的,也是大模型最能真正帮上忙的使用场景之一。尽管大模型经常能写出看起来像模像样的文字,但是大多数情况下内容的创意程度和文风都经不起深究。尤其是在创作领域,大模型常见的 “GPT 文风” 更是让利用大模型进行创意写作看起来简单,实际却困难重重。中文创意写作能力超GPT-4,最会写的中文大模型来了

近日,波形智能的大模型团队发布了一款专精 AI 写作的专业大模型 Weaver。通过写作领域专业预训练和一套创新性的数据生成和 Alignment 算法,Weaver 在写作领域的各种任务上均取得了领先 GPT-4 和众多中文通用大模型的效果,尤其是在生成内容的创意性和文风质量上大幅领先,是一款更能写出 “人话” 的大模型。

  • 论文地址:https://arxiv.org/pdf/2401.17268.pdf
  • 在线 Demo:https://www.wawawriter.com/

ChatGPT 等大模型在通用指令跟随和问答任务中效果出色,但是将大模型应用于专业写作,尤其是需要创造性和个性化文风的创意写作领域却依然面临重重阻碍。其中最大的问题就是大模型生成内容风格过于平淡,或者说文风过于 “GPT”,缺少创造性。

为了解决这个问题,训练出更适合专业写作的大模型,波形智能的研究团队分析了为什么 GPT 和其他通用大模型都做不好创意写作类任务。首先,通用大模型的预训练过程,因为希望让模型在更多的数据中自监督学习,预训练的数据集中常常会包含非常多的低质量内容,真正由专业作家和内容创作者写作的高质量文本内容可能只占预训练数据总量的 0.1% 不到。因此,经过预训练后的语言模型在建模了整个互联网的文本分布之后,自然会倾向于输出较为普通的内容。而在模型的对齐阶段,OpenAI 等公司众包标注指令微调数据集的过程中的标注员的教育 / 写作水平有限,没有对标注者的写作 / 创作能力进行筛选。另外标注的过程中的标准也主要强调回答的无害性 (harmlessness) 和有效性 (helpfulness),而没有考虑回答内容的创造性和语言 / 写作风格。因此,经过指令微调的语言模型反而更容易生成平庸无趣的文字。最后,在 RLHF/DPO 等 alignment 算法中,模型的训练数据和 Reward Model 均由经过指令微调后的模型生成或训练得到,因此对于文风和创造性上,RLHF/DPO 的过程也只能是 “矮子里拔将军”,无法强化出真正擅长写作的大模型。

基于此观察,波形智能的大模型团队提出了一个尤其适合创意写作领域的垂域专业模型训练 pipeline,并基于此方案训练了 Weaver,一个全球领先的创意写作大模型。该方案覆盖了模型的 (持续) 预训练,指令微调 (instruction tuning),和对齐 (RLHF/DPO) 阶段。在预训练阶段,团队进行了非常仔细的数据筛选和过滤,利用人工 + 规则 + 机器学习模型协同的方案,从开源预训练数据集中找到了高质量的小说 / 短故事 / 创意文案等类别的文本内容,舍弃掉了大量的低质量内容和代码 / 广告等数据,并下采样了一部分高质量的新闻数据,同时结合了大规模的私有创作领域数据 (小说,短故事等),构建出了超过 200B 的可以让模型专注学习创作能力的预训练数据。

在指令微调阶段,波形智能的数据生成团队参考并改进了 Meta 提出的 LongForm 和 HumpBack 方案,构建了一套可以基于一段高质量内容,自动生成各种写作相关任务指令和对应的高质量输出的 Instruction Backtranslation 流水线。团队总结并定义了 “写内容”,“写大纲”,“扩写”,“润色”,“精简”,“风格迁移 (仿写)”,“审校”,“头脑风暴”,“起标题”,和 “写作相关对话” 十个类别的任务。对于一类任务,如 “润色”,标注 Prompt 中首先解释任务的定义和几个输入输出样例,之后给出一个从一段文本中自动挖掘润色任务指令 / 输入 / 输出的例子和标注的思考过程: “首先在文本中找到一段写的很好的句子,假设这句话是经过一次润色而来的,之后猜测在润色之前这句话会是什么样子,最后分析润色前后的变化,推理出润色的指令会是什么样子。” 之后标注的 Prompt 中输入需要标注的例子并指示大模型按照例子中的标注流程进行输出,最后 parse 出模型输出中标注的 “指令 / 输入 / 输出” 部分,组合成一条写作指令数据。

相比 OpenAI 等公司的标准众包标注指令数据的流程,波形智能的标注策略更高效 (众包标注者只需要挑选特定领域高质量的内容即可,后续标注流程由 AI 完成),而众包标注和目前常用的 self-instruct 类的全自动标注流程相比,波形智能的标注流程能够生成更高质量的数据 (因为输出是手工挑选的高质量内容或其中的一部分)。基于这个策略,波形智能的大模型团队收集了涵盖小说写作,创意写作,专业写作,营销文案写作这四大领域中高质量的内容并进行了自动化标注,产出了 100 万 + 高质量的写作领域指令微调数据集。

图 1: Weaver 训练数据分布和来源

接下来,在对齐 (Alignment) 阶段,波形智能的数据生成团队提出了 Constitutional DPO, 一套全新的,基于原则高效将模型和专业作家 / 创作者对齐的方案。和以往基于模型输出 + 人类 / 大模型评估的对齐策略不同。Constitutional DPO 以人类创作者创作的高质量的输出作为正样本,利用人类作家 / 编辑整理提炼出的各个领域写作的 “原则 (Principles)”,用这些原则去生成能够教会模型更好地遵守这些原则的负样本。具体来说,专业作家 / 编辑首先整理出四大领域十个任务中,好的内容需要遵循的共 200 余条原则。对于每一个原则,编辑总结出原则的详细解释和一对符合 / 违背该原则的例子,并用几句话解释出符合 / 违背原则的原因。之后,对于每一个正样本,负例生成的 prompt 中首先展示出领域 - 任务上的原则集合和原则对应的例子和解释,之后展示出正样本,要求大模型分析出正样本最符合哪几条原则,并推理出如何修改能够在作出较少改变的情况下让正样本转而违背这个原则,从而变成一条质量没那么好的输出。团队精选了各个领域高评分 / 高阅读量 / 高点赞评论数的内容作为正样本,通过 Consitutional DPO 的流水线生成出了数万条偏好数据 (preference data),并利用这些数据对模型利用 DPO 进行了对齐训练。

图 2 - Constitutional DPO 方法示意图

图 3 - 专家标注的写作原则

除此之外,波形智能的数据生成团队还设计了一套支持 RAG-aware training 的数据生成方案,过滤 / 精选出了一系列输出内容明显基于其他内容的样本,通过 10 余个常用的 RAG 模版,构造出了 10 万余条的 RAG 训练数据,使得 Weaver 模型能够原生支持 RAG,能够结合参考文献和范文进行高质量的创作 / 仿写。除此之外,团队还设计了一套让 Weaver 支持 Function Calling 的数据生成方案。最终 Weaver 的微调数据量总和达到了 100 万 + 量级。

Weaver 模型家族一共包括四个不同大小的模型,名字叫做 Weaver-mini/base/pro/ultra, 分别包括 18 亿,60 亿,140 亿和 340 亿参数。为了评估 Weaver 模型和通用大模型的写作能力,波形智能的模型评估团队构建了一个新的用户大模型专业写作能力评估的 Benchmark。Benchmark 中精选了涵盖四大写作领域 30 余个子领域的十项写作任务的有代表性指令,共包含 2000 + 条指令。团队收集了 Weaver 和 10 余个有代表性的开源 + 闭源模型在 Benchmark 上的输出,并分别进行了人工对比评估和基于 GPT4 的自动评估。

评估结果显示,Weaver Ultra 在 Benchmark 中对生成内容的新颖度和文风的评估中对比包括 GPT-4 在内的通用大模型均有显著领先,在生成内容的流畅性和切题程度上也和行业领先的 GPT-4 相当,领先其他开源 / 闭源模型。而其他较小的 Weaver 模型也都在各项指标中相比大 2-3 倍的通用大模型有明显优势。

图 4: Weaver 在 WriteBench 的评测结果

除了标准 Benchmark 的人工和自动评估以外,波形智能的模型评估团队还在包含人机交互的实际应用场景中对 Weaver Ultra 和 GPT-4 进行了用户体验测评。由 4 位人类写手在同样的 Chat Interface 分别使用 Weaver Ultra 和 GPT-4,以相同的主题分别创作一个短故事,一个小红书文案,一个商业计划书,和一个课程论文。测评结果显示,人类写手利用 Weaver 进行创作的效率相比使用 GPT-4 提升了约 40%,而专业编辑对创作内容的质量评比中也以 9:3 的比分更倾向于采用 Weaver 创作的文案。分析显示,Weaver 带来的效率提升主要来自于生成内容的文风更得体,需要的后编辑更少,以及创作过程中 Weaver 交互更加直接,不会输出无用的废话和疑问。而来自专业编辑的反馈主要集中在基于 Weaver 创作的作品风格往往更符合实用标准,以及创作的内容个新颖程度更高,不死板。

图 5: Weaver 和其他大模型在人工评测中的 ELO Rating

#xxxx

#xxxx

http://www.kler.cn/a/503879.html

相关文章:

  • SQL面试题1:连续登陆问题
  • Oracle Dataguard(主库为双节点集群)配置详解(5):将主库复制到备库并启动同步
  • Grails应用http.server.requests指标数据采集问题排查及解决
  • C++ STL之容器介绍(vector、list、set、map)
  • Web前端界面开发
  • linux的大内核锁与顺序锁
  • 托宾效应和托宾q理论。简单解释
  • uniapp 发布后原生img正常,image无法显示,img与uniapp image使用区别
  • 【Block总结】Conv2Former的Block,结合卷积网络和Transformer的优点|即插即用
  • 视频超分(VSR)论文阅读记录/idea积累(一)
  • 【学术会议指南】方向包括遥感、测绘、图像处理、信息化教育、计算机技术、通信、大数据、人工智能、机械设计、仿真...可线上参与
  • Oracle重启后业务连接大量library cache lock
  • 【web靶场】之upload-labs专项训练(基于BUUCTF平台)
  • 工程师 - Eclipse安装和UML插件
  • 代码随想录刷题day07|(数组篇)58.区间和
  • LeetCode 热题 100_从前序与中序遍历序列构造二叉树(47_105_中等_C++)(二叉树;递归)
  • AI-ANNE:探索型神经网络——将深度学习模型转移到微控制器和嵌入式系统
  • 【网络云SRE运维开发】2025第2周-每日【2025/01/11】小测-【第11章NAT理论和实操考试】解析和参考
  • 中国地面气候资料日值数据集(V3.0)格式和下载说明
  • 【深度学习】核心概念-数据驱动(Data-Driven)
  • 详解C#的文件写入和读取:从基础到高级应用
  • 初识JAVA-面向对象的三大特征之多态
  • DS1302模块学习笔记
  • 【gin】http方法了解,以及RESTful API与版本控制
  • [IGP]ospf ip frr 快速重路由技术
  • 认识微服务