Repository: transformers
Issue: Quantization support for heads and embeddings
Labels: Feature request,Quantization

### Feature request

Hi! I’ve been researching LLM quantization recently ([this paper](https://arxiv.org/abs/2405.14852)), and noticed a potentially improtant issue that arises when using LLMs with 1-2 bit quantization.

### Problem description :mag:

Transformers supports several great ways for quantizing transformer ‘body’, but it seems that there is no built-in way
to quantize embeddings and/or lm head.

The reason why this is important is that some of the recent LLMs have very large vocabularies, and as a result,
their embeddings and heads can get massive. For instance, [Llama 3](https://huggingface.co/meta-llama/Meta-Llama-3-8B)
has 128K token vocabulary, [Qwen 2](https://huggingface.co/Qwen/Qwen2-72B-Instruct) has over 150K, [Gemma 2b](https://huggingface.co/google/gemma-2b) has 256K

As a result, if you load NF4 or AQLM quantized models, their **embeddings can take up 50% or more of the model footprint**. 
This is even more critical for lower bitwidth quantization:

![https://galqiwi.ru/persistent/2024-06-18/embed-1.png](https://galqiwi.ru/persistent/2024-06-18/embed-1.png)

### Feature Request :rocket:

It would be great if transformers had a flag to quantize embeddings and heads using some of the existing quantization methods. One simple way would be to use LLM.int8 or NF4 by Tim Dettmers since transformers already supports this.

I’ve investigated how quantizing embeddings with these methods affects common models. Below is model perplexity for [Llama 3 8B using AQLM+PV 2-bit quantization](https://huggingface.co/ISTA-DASLab/Meta-Llama-3-8B-AQLM-PV-2Bit-1x16). I measured three configurations: fp16 embeddings, int8 embeddings and NF4 embeddings with the same parameters that transformers uses for linear layers.

![https://galqiwi.ru/persistent/2024-06-18/emb_v3.png](https://galqiwi.ru/persistent/2024-06-18/emb_v3.png)
![https://galqiwi.ru/persistent/2024-06-18/head_v3.png](https://galqiwi.ru/persistent/2024-06-18/head_v3.png)

The values represent perplexity on [WikiText-2 test set](https://huggingface.co/datasets/Salesforce/wikitext/viewer/wikitext-2-v1)
measured with the same protocol used in [GPTQ](https://arxiv.org/abs/2210.17323) / [AQLM](https://arxiv.org/abs/2401.06118)
/ [QuIP#](https://arxiv.org/pdf/2402.04396.pdf) papers.
The code for these measurements can be found [here](https://gist.github.com/galqiwi/cb896f39052d1f4f718cb772040f3088).

Overall, 8-bit compression looks nearly lossless, the increase in perplexity does not exceed the error you get when 
quantizing the transformer with the same LLM int8 codec. In turn, NF4 introduces some error (within 0.05 for Llama 3),
but I would argue that this trade-off makes sense for low memory applications. Also, embeddings appear
easier to quantize than heads.

### Implementation details :gear:

There are multiple obstacles on the way to implementing this feature:
#### No support for mixed quantization
Currently, transformers does not support quantizing with multiple `HfQuantizer`s. IMO this is a good behaviour, as interactions between different quantizators can be messy. The problem is that this feature requires for transformers library to use different compression methods for body and heads/embeddings. I think that can be solved by extending `HfQuantizer` interface by adding embedding/head quantization methods and adding new `[embed,head]_quantization_config` arguments to `QuantizationConfigMixin` or something in this area.
#### No support for embedding quantization in bitsandbytes
As far as I know, no quantization method supports `nn.Embedding`-like interface. I can ask bitsandbytes maintainers if they would accept a PR that fixes that.

Also, there is a caveat that some models use tied embeddings/heads, while implementing, one need to be mindful of them.

### Cool things that this can enable :trophy:

If we can implement 4-bit embeddings, it will be possible to write a colab notebook that runs [Llama 3 70B model](https://huggingface.co/meta-llama/Meta-Llama-3-70B)
on the free tier T4 GPU without offoading, by combining embedding/heads quantization and the  PV-tuned model 
https://huggingface.co/ISTA-DASLab/Meta-Llama-3-70B-AQLM-PV-1Bit-1x16 .

Another use case is running quantized LLMs on smartphones or embedded devices: for instance, the [gemma-2b](https://huggingface.co/google/gemma-2b) can fit into 1GB RAM, but only if you quantize embeddings/heads in addition to transformer weights.

If you’re interested in making a demo out of this, I’d be excited to implement this with your review / recommendations if you prefer, or wait for you to implement it your way.


What do you think?

### Motivation

We are faced with a new bottleneck in model quantization. I think we can manage to fix it

### Your contribution

I can allocate my time to submitting PR, but we need to figure out what to do first