大模型推理加速：KV缓存、量化与模型剪枝实战

来利用 KV 缓存，避免重复计算已经处理过的键值对。每次生成一个新的 token 后，更新 KV 缓存并继续生成下一个 token。模型的第一个线性层作为示例，采用随机剪枝的方法，将权重的 20% 进行剪枝。最后移除剪枝的参数，使其永久生效。来加载量化后的模型。然后将输入文本转换为输入 ID 并进行推理，最后解码生成的 token。首先加载模型和分词器，然后将输入文本转换为输入 ID。在推理过程中

egzosn

639人浏览 · 2025-05-07 13:52:03

egzosn · 2025-05-07 13:52:03 发布

以下是大模型推理加速中KV缓存、量化与模型剪枝的一些实战方法：

KV缓存

优化缓存管理策略：采用弹性极速缓存EIC等技术，通过以存代算、GDR零拷贝等方式大幅降低推理GPU资源消耗，优化推理时延。EIC支持多种访问方式，能够支持不同的推理场景及性能需求，还能在内存和SSD之间通过数据沉降、回热及驱逐等智能数据流动策略，降低内存依赖，支持数据极速摄取，淘汰低频数据。
动态调整缓存大小：根据输入文本的长度和复杂度，以及模型的性能表现，动态调整KV缓存的大小。例如，对于较长的输入文本或复杂的推理任务，可以适当增加缓存大小，以减少重复计算；而对于简单的任务，则可以减小缓存大小，节省内存资源。
缓存复用与共享：在多个推理请求之间共享KV缓存，避免重复计算相同的键值对。例如，当处理多个相似的文本生成任务时，可以将前一个任务的KV缓存作为初始值，供后续任务使用，从而提高推理效率。
示例：

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载模型和分词器
model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 输入文本
input_text = "Once upon a time"
input_ids = tokenizer(input_text, return_tensors="pt").input_ids

# 初始化KV缓存
past_key_values = None

# 推理过程，模拟多步生成
for _ in range(5):
    outputs = model(input_ids, past_key_values=past_key_values)
    logits = outputs.logits
    next_token_id = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
    input_ids = next_token_id

    # 更新KV缓存
    past_key_values = outputs.past_key_values

    # 解码生成的token
    generated_text = tokenizer.decode(next_token_id[0], skip_special_tokens=True)
    print(generated_text)

在这个示例中，我们使用了transformers库中的GPT2模型。首先加载模型和分词器，然后将输入文本转换为输入 ID。在推理过程中，我们通过设置past_key_values来利用 KV 缓存，避免重复计算已经处理过的键值对。每次生成一个新的 token 后，更新 KV 缓存并继续生成下一个 token。

量化

选择合适的量化方案：根据模型的特点和硬件支持情况，选择合适的量化方法，如对称量化、非对称量化等。对于大语言模型，常见的量化方案包括W8A8量化算法等，即权重和激活值都采用8位整数表示。在实际应用中，可参考小红书团队的经验，对权重采取per - channel量化策略，以减少量化误差。
调整量化超参数：对于一些量化算法中的超参数，如SmoothQuant中的alpha值，需要根据模型的具体情况进行调整。不同的模型层可能需要不同的alpha值，通过超参搜索找到适合每个层的最佳alpha值，可以进一步降低量化误差，提高模型精度。
结合硬件进行优化：利用支持量化计算的硬件设备，如具有特定指令集的GPU或专用的AI芯片，来加速量化后的模型推理。这些硬件通常能够高效地处理低精度数据，从而提高计算速度，减少能耗。
示例：

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from bitsandbytes.nn import Int8Params

# 加载模型和分词器
model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 8位量化加载模型
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    load_in_8bit=True,
    device_map='auto',
    quantization_config=Int8Params()
)

# 输入文本
input_text = "Once upon a time"
input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(model.device)

# 推理
outputs = model(input_ids)
logits = outputs.logits
next_token_id = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
generated_text = tokenizer.decode(next_token_id[0], skip_special_tokens=True)
print(generated_text)

使用bitsandbytes库实现 8 位量化。通过设置load_in_8bit=True和quantization_config=Int8Params()来加载量化后的模型。然后将输入文本转换为输入 ID 并进行推理，最后解码生成的 token。

模型剪枝

确定剪枝策略：可以选择基于权重的剪枝策略，即根据权重的大小或重要性来决定是否剪除某些连接；也可以采用基于神经元的剪枝策略，直接删除不重要的神经元。此外，还有动态令牌剪枝等方法，如LazyLLM通过动态选择性地计算关键令牌的键值缓存，优化了模型的预填充和解码阶段。
渐进式剪枝：为了避免一次性剪枝过多导致模型性能下降，可以采用渐进式剪枝的方法。在模型的较浅层，倾向于保留更多的令牌，以保证信息的完整性和生成质量；随着层数的增加，逐渐增加剪枝的比例，因为深层的模型能够更加有效地从剩余的令牌中提取必要的信息。
结合辅助机制：引入辅助缓存机制来应对剪枝可能带来的信息损失。该缓存存储了在之前步骤中被剪枝但可能在后续生成中需要重新利用的令牌的状态，这样可以在不重新计算这些令牌的情况下，快速恢复它们的信息，从而提高推理的准确性。
示例：

import torch
import torch.nn.utils.prune as prune
from transformers import AutoModelForCausalLM

# 加载模型
model_name = "gpt2"
model = AutoModelForCausalLM.from_pretrained(model_name)

# 选择要剪枝的模块，这里以第一个线性层为例
module = model.transformer.h[0].attn.c_attn

# 对模块进行剪枝，这里采用随机剪枝，剪枝比例为0.2
prune.random_unstructured(module, name="weight", amount=0.2)

# 移除剪枝的参数，使其永久生效
prune.remove(module, 'weight')

# 打印剪枝后的权重形状
print(module.weight.shape)