FLUTE: 一种灵活高效的查找表量化引擎

FLUTE简介

FLUTE (Flexible Lookup Table Engine)是一个为查找表(LUT)量化大语言模型设计的灵活引擎。它可以实现快速矩阵乘法,支持多种量化方案,在保持模型性能的同时大大降低内存和计算需求。

FLUTE logo

FLUTE的主要特点包括:

支持多种量化方案,包括uniform、lookup table等
兼容多种主流模型,如LLaMA、Gemma等
与vLLM、Hugging Face等框架无缝集成
提供简单易用的API和命令行工具
在多个基准测试中表现出色

本文将详细介绍FLUTE的背景、特性、使用方法以及在各种模型上的表现。

背景

随着大语言模型的规模不断增长,如何在有限的计算资源下高效运行这些模型成为一个重要问题。量化是一种有效的模型压缩方法,可以将模型参数从高精度(如FP32)压缩到低精度(如INT8、INT4等),从而减少内存占用和计算量。

传统的uniform量化将全精度权重映射到等间隔的低精度区间。而lookup table(LUT)量化则是一种更灵活的非均匀量化方案,可以通过查找表将量化后的值映射到任意值。FLUTE就是基于LUT量化设计的一个高效引擎。

FLUTE支持多种量化方案:

int4, int3, int2: 对应uniform/integer量化
fp4, fp3, fp2: 浮点量化
nf4, nf3, nf2: 扩展了QLoRA中引入的nf4数据格式
任意自定义查找表: 甚至可以学习得到最优表

这种灵活性使FLUTE能够适应不同模型和任务的需求,在压缩率和性能之间取得更好的平衡。

FLUTE的主要特性

1. 高效的矩阵乘法实现

FLUTE针对LUT量化模型优化了矩阵乘法实现,可以显著提高计算速度。下图展示了FLUTE在不同设置下的性能表现:

Performance benchmark

可以看到,FLUTE在各种量化位数和分组大小下都能保持较高的吞吐量。

2. 广泛的模型支持

FLUTE目前支持以下模型:

LLaMA-3/3.1 (8B, 70B, 405B)
Gemma-2 (9B, 27B)

这些模型覆盖了目前主流的大语言模型架构,使FLUTE能够广泛应用于各种场景。

3. 灵活的量化选项

FLUTE支持多种量化设置:

位数: 4-bit, 3-bit, 2-bit
分组大小: 32, 64, 128, 256
输入数据类型: torch.float16, torch.bfloat16

用户可以根据具体需求选择合适的量化参数。

4. 与主流框架集成

FLUTE可以无缝集成到vLLM和Hugging Face等主流框架中。例如,使用FLUTE量化的模型可以直接通过vLLM部署服务:

python -m flute.integrations.vllm vllm.entrypoints.openai.api_server \
    --model radi-cho/Meta-Llama-3.1-8B-FLUTE \
    --quantization flute

这种集成使得FLUTE可以轻松融入现有的机器学习工作流程。

使用FLUTE

安装

FLUTE可以通过pip安装:

# For CUDA 12.1
pip install flute-kernel
# For CUDA 11.8
pip install flute-kernel -i https://flute-ai.github.io/whl/cu118

量化自定义模型

FLUTE提供了简单的API来量化自定义模型。以下是一个基本示例:

import flute.integrations.base

model = AutoModelForCausalLM.from_pretrained(
    pretrained_model_name_or_path,
    device_map="cpu",
    torch_dtype="auto")

flute.integrations.base.prepare_model_flute(
    name="model.model.layers",
    module=model.model.layers,
    num_bits=4,
    group_size=128,
    fake=False)