FLUTE：用于LUT量化LLM的灵活查找表引擎

更新

• 2024年8月5日 添加量化的LLaMA-3.1（8B/70B）模型。
• 2024年8月2日 添加对RTX4090的支持。
• 2024年7月27日 添加对LLaMA-3.1（405B）的支持，并优化了BF16性能。FP16仍然是推荐的数据类型，尤其是在3比特设置下。

安装

通过pip安装FLUTE或从源代码构建：

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

前往入门指南并尝试使用！

背景

均匀量化将全精度权重转换为等大小的低精度区间。查找表（LUT）量化是非均匀量化的一种灵活变体，可以通过查找表将区间映射到任意值。

其中$\mathbf{Q}$表示量化权重，$\mathbf{s}$表示（分组）缩放，$\widehat{\mathbf{W}}$表示反量化权重。以下是FLUTE支持的一些查找表示例。

FLUTE支持的新模型

内核的灵活性可能会带来新的量化算法。作为概念证明，我们发布了一些使用**学习正态浮点（NFL）**量化的模型 — 这是对QLoRA引入的nf4数据格式的简单扩展。NFL使用NF量化的查找表和缩放初始化查找表和缩放。然后，它使用校准数据通过直通估计来学习缩放，以获得关于缩放的梯度。

基准测试

有关其他基准测试、详细分析和相应的指令调优模型，请参阅论文和模型库。

LLaMA-3.1

Gemma-2

入门指南

FLUTE + vLLM

FLUTE量化模型（模型库）可以直接使用现有框架（如vLLM）进行服务。

- python -m vllm.entrypoints.openai.api_server \
+ python -m flute.integrations.vllm vllm.entrypoints.openai.api_server \
    --model [模型] \
    --revision [版本] \
    --tensor-parallel-size [TP大小] \
+   --quantization flute

例如，以下命令在单个GPU上运行FLUTE量化的LLaMA-3.1（8B）。

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

然后我们可以像往常一样查询vLLM服务器。

curl http://localhost:8000/v1/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "radi-cho/Meta-Llama-3.1-8B-FLUTE",
        "prompt": "San Francisco is a",
        "max_tokens": 7,
        "temperature": 0
    }'

FLUTE + HuggingFace

FLUTE也可以与HuggingFace及其accelerate扩展无缝配合。这种集成主要是实验性的，尚未优化。对性能考虑敏感的用户应该使用vLLM集成。

以下示例对密集模型执行简单量化。之后，模型可以正常使用。（即将支持加载预量化模型！）

import flute.integrations.base
flute.integrations.base.prepare_model_flute(
    name="model.model.layers",
    module=model.model.layers,  # 适用于LLaMA-3和Gemma-2
    num_bits=num_bits,
    group_size=group_size,
    fake=False,
    handle_hooks=True)  # 用于`accelerate`钩子

支持和兼容性

内核

[!警告] 在当前版本中，我们注意到torch.bfloat16比torch.float16慢。这可能是由于缺乏调优，以及Ampere GPU缺少对bfloat16矢量化原子加法的硬件加速。

[!警告] 我们注意到在使用位数=4，分组大小=256，GPU=A100时出现了一些数值不稳定的情况，尽管这种情况相对罕见（9360个测试用例中有8个失败）。我们还注意到在某些情况下使用位数=4，分组大小=256，数据类型=bfloat16，GPU=RTX4090时存在正确性问题（52个测试用例中有1个失败）。我们将对此进行调查，但目前建议避免使用这些特定配置（W4G256）。

模型

[!注意] 在当前版本中，由于历史原因，内核是形状专用的（即我们为每个矩阵形状调整tile大小等）。请参阅下表了解支持的用例，因为不同的平台和张量并行大小会改变矩阵形状。我们计划在不久的将来添加对广泛形状的支持。同时，如果您有任何特定的模型需求，请告诉我们，我们很乐意为其添加支持。

模型库

[!注意] 我们在此发布的模型是在更多数据上训练的，因此与论文中的模型不同。

[!提示] HuggingFace Hub链接默认使用NFL W4G64量化。要使用NFL W3G64量化，请添加--revision nfl_w3g64。

LLaMA-3.1 (8B)

LLaMA-3.1 (70B)

LLaMA-3.1 Instruct (8B)

LLaMA-3.1 Instruct (70B)

LLaMA-3 (8B)

LLaMA-3 (70B)

LLaMA-3 Instruct (8B)

LLaMA-3 Instruct (70B)

Gemma-2 (9B)

Gemma-2 (27B)

Gemma-2 Instruct (9B)

Gemma-2 Instruct (27B)

量化您自己的模型

我们提供了两个API来量化自定义模型。最简单的方法是使用命令行界面。

简单正态浮点量化

python -m flute.integrations.base \
    --pretrained_model_name_or_path meta-llama/Meta-Llama-3-70B-Instruct \
    --save_directory Meta-Llama-3-70B-Instruct-NF4 \
    --num_bits 4 \
    --group_size 128

CLI本质上包装了以下Python API，

from transformers import (
    LlamaForCausalLM,
    Gemma2ForCausalLM,
    AutoModelForCausalLM)
import flute.integrations.base

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

if isinstance(model, (LlamaForCausalLM, Gemma2ForCausalLM)):
    flute.integrations.base.prepare_model_flute(
        name="model.model.layers",
        module=model.model.layers,
        num_bits=num_bits,
        group_size=group_size,
        fake=False)
else:
    # 更多模型即将推出
    raise NotImplementedError

学习正态浮点量化

即将推出！

从源代码构建

1. 克隆CUTLASS库。

# 不幸的是，目前路径是硬编码的。如果您在不同的目录中安装CUTLASS，
# 请确保更新了`setup.py`中相应的路径。
cd /workspace

git clone https://github.com/NVIDIA/cutlass.git
cd cutlass
git checkout v3.4.1

2. 构建。

git clone https://github.com/HanGuo97/flute
cd flute
pip install -e .

**注意：**构建过程要求本地CUDA版本（nvcc --version）与PyTorch的CUDA匹配。在构建过程中出现与CUDA版本不匹配相关错误的情况下，请尝试添加--no-build-isolation。

致谢和引用

特别感谢Dmytro Ivchenko、Yijie Bei和Fireworks AI团队的有益讨论。如果您发现本仓库中的任何模型或代码有用，请随意引用：

@article{flute2024,
  title={Fast Matrix Multiplications for Lookup Table-Quantized LLMs},
  author={Guo, Han and Brandon, William and Cholakov, Radostin and Ragan-Kelley, Jonathan and Xing, Eric P and Kim, Yoon},
  journal={arXiv preprint arXiv:2407.10960},
  year={2024}
}