自回归生成模型中的推理加速详解

《DEEPSEEK原生应用与智能体开发实践 图书》【摘要 书评 试读】- 京东图书

在上一节（自回归生成模型中的资源计算-CSDN博客）中，我们详细讲解了自回归生成模型的原理与训练过程，揭示了训练环节的核心重要性。然而，在生成模型的全面实践中，更加关键的一步在于如何精妙地运用这些训练成熟的模型去执行实际的推理任务。而在此过程中，一个至关重要的考量因素，便是如何高效利用现有的设备和资源，以最优化的方式进行模型推理。

在本节中，我们将聚焦于生成模型中的推理加速内容，深入探讨如何通过技术手段提升模型推理的速度与效率。我们将介绍一系列策略和技巧，包括模型压缩、并行计算、硬件优化等，旨在帮助大家充分利用计算资源，减少推理时间，从而实现更快速、更高效的生成模型应用。通过本节的学习，你将能够掌握提升生成模型推理性能的关键方法，为你的实际项目注入更强大的动力。

8.2.1  模型推理中的“贪心生成”与“采样生成”

前面我们在最后讲解了模型生成时的参数temperature与采样个数Top_k，这实际上是我们在使用时默认使用了采样生成策略。在深度学习中，特别是在自然语言处理领域，文本生成是一个重要的任务。在这个过程中，贪心生成（Greedy Generation）和采样生成（Sampling Generation）是两种常用的策略。下面我们将详细讲解这两种策略及其特点。

1. 贪心生成

原理：贪心生成的核心思想是，在每一步生成过程中，都选择当前概率最大的token（即最可能的词或字符）作为预测值。这种方法简单直观，计算效率高。

优点：由于每次都选择最有可能的词，贪心生成往往能生成语法正确、语义通顺的文本。

缺点：然而，贪心策略也容易导致生成的文本缺乏多样性。因为每次都选择概率最大的词，所以生成的文本往往比较单一，缺乏创新和变化。

贪心生成的简单示例如下：

import torch
import torch.nn.functional as F
 
# 假设我们有一个训练好的模型model，和一个初始的输入序列input_seq
# model应该是一个PyTorch模型，其输出是词汇表大小的logits
# input_seq是一个张量，表示输入的序列
 
# 使用模型得到下一个token的预测分布
logits = model(input_seq)
# 使用Softmax函数得到概率分布
probs = F.softmax(logits, dim=-1)
# 选择概率最大的token作为下一个词
next_token = torch.argmax(probs, dim=-1)
 
print("Greedy generated token:", next_token.item())

2. 采样生成

为了增加生成的多样性，人们提出了采样生成的方法。这种方法通过采样的方式，使得非最大概率值的token也有机会被选中。

原理：在采样生成中，每个token被选中的概率与其在模型输出的概率分布中的值成正比。这样，即使某个token的概率不是最大的，也有可能被选中，从而增加了生成的多样性。

Top-k采样：在这种方法中，首先选取概率最高的k个token作为候选集，然后从这个候选集中随机选择一个token作为输出。这样做的好处是，既保证了生成的文本有一定的质量（因为候选集都是概率较高的token），又增加了多样性（因为是从k个候选词中随机选择的）。

采样生成的简单示例如下：

import torch
import torch.nn.functional as F

# 同样地，假设我们有一个训练好的model和一个初始的输入序列input_seq

# 使用模型得到下一个token的预测分布
logits = model(input_seq)
# 使用softmax函数得到概率分布
probs = F.softmax(logits, dim=-1)

# 采样生成：这里使用top-k采样作为示例
k = 5  # 设定top-k的值
# 获取概率最高的k个token的索引
top_k_probs, top_k_indices = torch.topk(probs, k, dim=-1)
# 从这k个token中随机选择一个
sampled_index = torch.multinomial(top_k_probs, 1).squeeze()
# 找到对应的token索引
next_token = top_k_indices[sampled_index]

print("Sampled token:", next_token.item())

除了前面的Top-k采样，我们还有一种称为Top-p采样的策略（核采样Nucleus Sampling）。它与Top-k采样不同，Top-p采样不是选择固定数量的候选词，而是选择一个概率阈值p。然后，从模型的预测分布中选择一个最小的词集合，使得这个集合中的单词的概率总和至少为p。这种方法更加灵活，因为它可以根据模型的输出动态调整候选词的数量。

优点：采样生成方法通过引入随机性，增加了生成的多样性，使得生成的文本更加丰富多彩。

缺点：采样生成也可能导致生成的文本质量下降。因为非最大概率值的token被选中的机会增加，所以可能会生成一些语法错误或语义不通的文本。

对于Top-p采样，实现会稍微复杂一些，因为你需要计算累积概率并找到一个阈值，使得累积概率之和达到或超过p。这通常涉及对概率分布进行排序和累积求和，直到累积和达到或超过p为止。然后，从这个累积和达到p的子集中随机选择一个token。

在实际应用中，文本生成通常涉及更复杂的处理，如处理序列长度、处理特殊标记（如<EOS>表示序列结束）以及可能的批处理操作等。此外，model和input_seq需要根据你具体的模型和数据进行替换。在实际使用时，还需要确保模型处于评估模式（model.eval()），并处理任何可能的设备（CPU/GPU）兼容性问题。

8.2.2  模型推理过程中的冗余计算问题解析

在探讨基于自回归架构的推理过程时，我们可以清晰地看到其自回归生成机制是如何运作的。以用户输入“上海的天气”为例，模型逐步生成了“是晴天”这一完整输出。整个生成过程既精妙又直观，如图8-4所示。

首先，模型接收输入“上海的天气”，并计算出每个token的注意力表示（这些注意力表示在图中以绿色部分呈现）。利用“天气”这一token的注意力表示，模型预测并生成了下一个token“是”。

随后，模型将新生成的“是”拼接到原始输入后，形成新的输入序列“上海的天气是”。再次通过模型处理，利用“是”的注意力表示，模型预测并生成了下一个token“晴”。这一过程不断重复，模型依次将新生成的token拼接到输入序列中，并基于最新的输入序列预测下一个token，直至生成完整的输出“上海的天气是晴天”。

图8-4  自回归生成示例

我们将这个推理过程用矩阵Embedding的形式进行表示，如图8-5所示。

图8-5  矩阵表示的推理过程

用公式表示则为：

根据上面的图解和公式可以看到，在自回归生成过程中，推理中的每个token是逐个生成的。以预测token3这个token为例，模型只需要考虑当前token的Query向量（Q3）与所有之前token的Key和Value向量的关系，而不需要重新计算之前token的注意力表示。这种不必要的重复计算不仅增加了计算量，还延长了生成过程的时间。

具体来说，我们希望有一种方法。在预填充阶段，模型会计算输入文本的Key和Value向量，并将它们缓存起来。然后，在解码阶段，模型会逐一生成token。对于每个新生成的token，模型会利用其对应的Query向量和缓存的Key、Value向量来计算注意力权重，从而预测出下一个token。带有缓存的推理过程如图8-6所示。

通过这种技术显著减少了计算量，提高了推理性能。

图8-6  带有缓存的推理过程

8.2.3  初识模型推理中的KV Cache与代码实现

在深度学习领域，特别是在自然语言处理任务中，自回归模型由于其生成特性而备受关注。然而，这类模型在生成长文本时往往面临计算量巨大的挑战。为了提升计算效率，KV Cache技术应运而生，成为优化自回归生成模型推理速度的重要手段。

KV Cache是一种通过缓存Attention机制中的键（K）和值（V）向量来减少重复计算的技术。在自回归生成模型中，每个token的生成都依赖于其之前的所有token，这意味着在生成长文本时，需要反复计算每个token的键（K）和值（V）向量。KV Cache通过预先计算并缓存这些向量，避免了在每次生成新token时的重复计算，从而显著提高了推理速度。

下面示例演示了我们在没有经过KV Cache缓存的条件下进行生成的过程：

import torch
import torch.nn as nn

batch_size = 1
seq_length = 48
hidden_size = 384
vocab_size = 1024

Wq = torch.randn(hidden_size, hidden_size)
Wk = torch.randn(hidden_size, hidden_size)
Wv = torch.randn(hidden_size, hidden_size)

embedding_layer = torch.nn.Embedding(vocab_size,hidden_size)
lm_head = torch.nn.Linear(hidden_size,vocab_size)

prompt_len = 2
#step1: 预填充阶段，处理输入prompt，生成第一个token
inputs = torch.randint(0,vocab_size,[2]) #[prompt_len]，设置输入的prompt长度
print("step1:输入prompt的token:",inputs)
inputs_emb = embedding_layer(inputs) #[prompt_len, hidden_size]
Q,K,V=inputs_emb@Wq,inputs_emb@Wk ,inputs_emb@Wv#[prompt_len, hidden_size]
att_weight=Q@(K.T)#[prompt_len, prompt_len]
print("step1:att_weight:",att_weight.shape)
att_output=att_weight@V #[prompt_len, hidden_size]
output = lm_head(att_output) #[prompt_len, vocab_size]
output_token = torch.argmax(output,dim=-1)[-1:]
print("step1:输出生成的token:",output_token)

print("##########################")
#step2：解码阶段，下一个token生成
inputs = torch.cat((inputs,output_token),dim=-1)
print("step2:输入的token:",inputs)
inputs_emb = embedding_layer(inputs) #[prompt_len+1, hidden_size]
Q,K,V=inputs_emb@Wq,inputs_emb@Wk ,inputs_emb@Wv#[prompt_len+1, hidden_size]
att_weight=Q@(K.T)#[prompt_len+1, prompt_len+1]
print("step2:att_weight:",att_weight.shape)
att_output=att_weight@V #[prompt_len+1, hidden_size]
output = lm_head(att_output) #[prompt_len+1, vocab_size]
output_token = torch.argmax(output,dim=-1)[-1:]
print("step2:生成的token:",output_token)

打印结果如下所示：

step1:输入prompt的token: tensor([130, 830])
step1:att_weight: torch.Size([2, 2])
step1:输出生成的token: tensor([1010])
##########################
step2:输入的token: tensor([ 130,  830, 1010])
step2:att_weight: torch.Size([3, 3])
step2:生成的token: tensor([90])

在这个基础上我们调整一下代码，完成带有KV Cache缓存的输入，代码如下所示：

import torch

batch_size = 1
seq_length = 48
hidden_size = 384
vocab_size = 1024

Wq = torch.randn(hidden_size, hidden_size)
Wk = torch.randn(hidden_size, hidden_size)
Wv = torch.randn(hidden_size, hidden_size)

embedding_layer = torch.nn.Embedding(vocab_size,hidden_size)
lm_head = torch.nn.Linear(hidden_size,vocab_size)

prompt_len = 2
#step1: 预填充阶段，处理输入prompt，生成第一个token
inputs = torch.randint(0,vocab_size,[2]) #[prompt_len]，设置输入的prompt长度
print("step1:输入prompt的token:",inputs)

inputs_emb = embedding_layer(inputs) #[prompt_len, hidden_size]
Q,K,V=inputs_emb@Wq,inputs_emb@Wk ,inputs_emb@Wv#[prompt_len, hidden_size]

cache_k = K #将prompt计算embedding作为cache_k
cache_v = V #将prompt计算embedding作为cache_v

att_weight=Q@(K.T)#[prompt_len, prompt_len]
print("step1:att_weight:",att_weight.shape)
att_output=att_weight@V #[prompt_len, hidden_size]
output = lm_head(att_output) #[prompt_len, vocab_size]
output_token = torch.argmax(output,dim=-1)[-1:]
print("step1:输出生成的token:",output_token)

print("##########################")
#step2：解码阶段，一个个token逐个生成
inputs = output_token   #这里相对于前面，没有进行concat操作
print("step2:输入的token:",inputs)

#相对于前面的#[prompt_len+1, hidden_size]，这里只有[1, hidden_size]
inputs_emb = embedding_layer(inputs) #[1, hidden_size]
#QKV计算量减少
Q,K,V=inputs_emb@Wq,inputs_emb@Wk ,inputs_emb@Wv#[1, hidden_size]
#kv cache显存占用
cache_k=torch.cat((cache_k,K),dim=0)#[prompt_len+1, hidden_size]
cache_v=torch.cat((cache_v,V),dim=0)#[prompt_len+1, hidden_size]

att_weight=Q@(cache_k.T)#[1, prompt_len+1]
att_output=att_weight@cache_v #[1, hidden_size]
output = lm_head(att_output) #[1, vocab_size]
output_token = torch.argmax(output,dim=-1)
print("step2:生成的token:",output_token)

打印结果如下：

step1:输入prompt的token: tensor([508, 621])
step1:att_weight: torch.Size([2, 2])
step1:输出生成的token: tensor([8])
##########################
step2:输入的token: tensor([8])
step2:生成的token: tensor([529])

从代码运行结果上来看，我们可以看到随着内容的输入，也可以获得对应的next token的输出。