[Background] [Benchmarks] [Getting Started] [Compatibility] [Model Zoo]

Update

Installation

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

Background

New Models Powered by FLUTE

Benchmarks

LLaMA-3.1

Gemma-2

Getting Started

FLUTE + vLLM

- python -m vllm.entrypoints.openai.api_server \+ python -m flute.integrations.vllm vllm.entrypoints.openai.api_server \    --model [MODEL] \    --revision [REVISION] \    --tensor-parallel-size [TP_SIZE] \+   --quantization flute

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

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

import flute.integrations.huggingface- model = AutoModelForCausalLM.from_pretrained(+ model = flute.integrations.huggingface.from_pretrained(    "radi-cho/Meta-Llama-3.1-8B-FLUTE",    # all of your favoriate HF flags will be forwarded    device_map="auto")

importflute.integrations.baseflute.integrations.base.prepare_model_flute(name="model.model.layers",module=model.model.layers,# for LLaMA-3 and Gemma-2num_bits=num_bits,group_size=group_size,fake=False,handle_hooks=True)# for `accelerate` hooks

Support and Compatibility

Kernel

Warning

In the current release, we noticedtorch.bfloat16 is slower thantorch.float16. This likely because of lack of tuning, and that Ampere GPUs lack a hardware acceleration forbfloat16vectorized atomic-add.

Warning

We noticed several numerically unstable situations usingbits=4, group-size=256, GPU=A100, though this is relatively rare (8 of 9360 test cases failed). We also noticed correctness issues in some situations withbits=4, group-size=256, dtype=bfloat16, GPU=RTX4090 (1 of 52 test cases failed). We will be looking into this, but we suggest avoiding these particular use cases (W4G256) for now.

Models

Note

As of the current release, the kernel is shape-specialized due to legacy reasons (i.e., we tune tile sizes etc for each matrix shape). Please see the below chart for the supported use cases, as different platform and tensor parallel size changes the matrix shapes. We plan to add supports for a broad range of shapes in the near future. In the meantime, please let us know if you have any specific models in mind and we are happy to add support for them.

Model Zoo

Note

The models we release here are trained on more data and hence different from those in the paper.

Tip

The HuggingFace Hub links are forNFL W4G64 quantization by default. To use theNFL W3G64 quantization, add--revision nfl_w3g64.

LLaMA-3.1 (8B)

LLaMA-3.1 (70B)

LLaMA-3.1 (405B)

LLaMA-3.1 Instruct (8B)

LLaMA-3.1 Instruct (70B)

LLaMA-3.1 Instruct (405B)

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)

Quantizing Your Own Models

Simple Normal Float Quantization

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

fromtransformersimport (LlamaForCausalLM,Gemma2ForCausalLM,AutoModelForCausalLM)importflute.integrations.basemodel=AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path,device_map="cpu",torch_dtype="auto")ifisinstance(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:# more models to comeraiseNotImplementedError

Convertingbitsandbytes Model into FLUTE Model

flute.integrations.base.prepare_model_flute(    name="model.model.layers",    module=model.model.layers,    num_bits=num_bits,    group_size=group_size,    fake=False,+   prepare_bnb_layers=True,+   default_bnb_dtype=torch.float16,)

Learned Normal Float Quantization (NFL)

importflute.integrations.learnableflute.integrations.learnable.learn_scales(model=model,tokenizer=tokenizer,num_bits=num_bits,group_size=group_size,custom_corpora=list_of_corpora,samples=num_samples,)

Extending to New Models (Experimental)

Step 1: Build theraw version of the library that exposes all templates.

cp flute/csrc/qgemm_kernel_generated.template.cu flute/csrc/qgemm_kernel_generated.cu

INSTANTIATE_TEMPLATE(NUM_SMs, DTYPE, cute::uint16_t, __half2, BITS, GROUP_SIZE);

Note

Although including other settings could still build, it could break the linking process and require re-compiling the library.

pip install -e. --no-build-isolation# `--no-build-isolation` is optional

Step 2: Tune FLUTE on the new matrix shapes.

importtorchfromflute.tuneimportTuneTask,tune_tasks_legacytasks= [TuneTask(M=1,# batch size (x sequence length, usually 1 for token-by-token generation)N=1024,# parameter dimension (note when using tensor-parallelism, this could change)K=4096,# parameter dimension (note when using tensor-parallelism, this could change)num_bits=4,# number of bitsgroup_size=64,# group sizenum_sms=108,# number of streaming multiprocessors of the GPUdtype=torch.float16,# data typedevice=torch.device("cuda:0")    ),]tune_tasks_legacy(tasks)

Step 3: Build the newly-tuned kernel

# remove changesgit checkout -- flute/csrc/# generating new dispatching logic based on tuning artifactsbash scripts/codegen_tuned.sh# remove changesgit checkout -- \    flute/ops.py \    flute/__init__.py# Buildpip install -e. --no-build-isolation

Build From Source

# Unfortunately, the path is hard-coded as of now. If you install CUTLASS# in a different directory, please make sure the corresponding path in# `setup.py` is updated.cd /workspacegit clone https://github.com/NVIDIA/cutlass.gitcd cutlassgit checkout v3.4.1

git clone https://github.com/HanGuo97/flutecd flutepip install -e.

Acknowledgement and Citation

@inproceedings{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 and Kim, Yoon},booktitle={Findings of the Association for Computational Linguistics: EMNLP 2024},pages={12419--12433},year={2024}}@article{higgs2024,title={Pushing the Limits of Large Language Model Quantization via the Linearity Theorem},author={Malinovskii, Vladimir and Panferov, Andrei and Ilin, Ivan and Guo, Han and Richt{\'a}rik, Peter and Alistarh, Dan},journal={arXiv preprint arXiv:2411.17525},year={2024}}