量化 KV 缓存
FP8 KV 缓存¶
将 KV 缓存量化到 FP8 可以减少其内存占用。这增加了缓存中可以存储的 token 数量，从而提高了吞吐量。
FP8 格式¶
OCP (Open Compute Project) 定义了两种常见的 8 位浮点数据格式
- E5M2 (5 个指数位和 2 个尾数位)
- E4M3FN (4 个指数位和 3 个尾数位，通常简称为 E4M3)
与 E5M2 相比，E4M3 格式提供更高的精度。然而，由于其动态范围较小 (±240.0)，E4M3 通常需要在每个量化张量旁边附加一个更高精度 (FP32) 的缩放因子。
当前限制¶
目前，仅支持每张量（标量）的缩放因子。正在开发以支持更精细粒度（例如每通道）的缩放因子。
性能影响¶
当前的 FP8 KV 缓存实现主要通过允许约两倍的 KV 缓存分配空间来提升吞吐量。这使得可以：
- 处理单个请求更长的上下文长度，或
- 处理更多并发请求批次
然而，由于当前的实现尚未包含融合的去量化和注意力操作，目前没有延迟方面的改进。未来的版本将支持带有硬件加速的量化注意力，这将提供额外的性能优势。尽管最新的芯片产品（例如 AMD MI300、NVIDIA Hopper 或更新版本）支持 FP8 与其他格式（fp32、fp16、bf16）之间的原生硬件转换，但这一优势尚未完全实现。
研究表明，FP8 E4M3 量化通常只会对推理精度造成微小的损失，这使其成为吞吐量优化的实用选择。
使用示例¶
以下是启用 FP8 量化的示例
# To calculate kv cache scales on the fly enable the calculate_kv_scales
# parameter
from vllm import LLM, SamplingParams
sampling_params = SamplingParams(temperature=0.7, top_p=0.8)
llm = LLM(model="meta-llama/Llama-2-7b-chat-hf",
kv_cache_dtype="fp8",
calculate_kv_scales=True)
prompt = "London is the capital of"
out = llm.generate(prompt, sampling_params)[0].outputs[0].text
print(out)
kv_cache_dtype 参数指定 KV 缓存存储的数据类型： - "auto"：使用模型的默认“未量化”数据类型 - "fp8" 或 "fp8_e4m3"：支持 CUDA 11.8+ 和 ROCm (AMD GPU) - "fp8_e5m2"：支持 CUDA 11.8+
使用校准尺度以提高精度¶
在使用 FP8 KV 缓存时获得最佳模型质量，我们建议使用针对代表性推理数据调优的校准尺度。LLM Compressor 是推荐用于此过程的工具。
安装¶
首先，安装所需的依赖项
使用示例¶
以下是使用 meta-llama/Llama-3.1-8B-Instruct 的完整示例（大多数模型可以使用相同的模式）
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from llmcompressor.transformers import oneshot
# Select model and load it
MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto", torch_dtype="auto")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
# Select calibration dataset
DATASET_ID = "HuggingFaceH4/ultrachat_200k"
DATASET_SPLIT = "train_sft"
# Configure calibration parameters
NUM_CALIBRATION_SAMPLES = 512 # 512 samples is a good starting point
MAX_SEQUENCE_LENGTH = 2048
# Load and preprocess dataset
ds = load_dataset(DATASET_ID, split=DATASET_SPLIT)
ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
def process_and_tokenize(example):
text = tokenizer.apply_chat_template(example["messages"], tokenize=False)
return tokenizer(
text,
padding=False,
max_length=MAX_SEQUENCE_LENGTH,
truncation=True,
add_special_tokens=False,
)
ds = ds.map(process_and_tokenize, remove_columns=ds.column_names)
# Configure quantization settings
recipe = """
quant_stage:
quant_modifiers:
QuantizationModifier:
kv_cache_scheme:
num_bits: 8
type: float
strategy: tensor
dynamic: false
symmetric: true
"""
# Apply quantization
oneshot(
model=model,
dataset=ds,
recipe=recipe,
max_seq_length=MAX_SEQUENCE_LENGTH,
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
)
# Save quantized model: Llama-3.1-8B-Instruct-FP8-KV
SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8-KV"
model.save_pretrained(SAVE_DIR, save_compressed=True)
tokenizer.save_pretrained(SAVE_DIR)
上述脚本将在您当前目录中创建一个文件夹，其中包含带有校准尺度的量化模型（例如 Llama-3.1-8B-Instruct-FP8-KV）。
运行模型时，您必须指定 kv_cache_dtype="fp8" 以启用 KV 缓存量化并使用尺度。