diff --git a/examples/quantization_embedding/README.md b/examples/quantization_embedding/README.md new file mode 100644 index 0000000000..96227ba7a3 --- /dev/null +++ b/examples/quantization_embedding/README.md @@ -0,0 +1,111 @@ +# Quantizing the Input Embedding + +`llm-compressor` can quantize a model's **input embedding table** (the vocab +lookup) to weight-only `intN` (`WNA16`), in addition to the linear layers. vLLM +loads these checkpoints and runs an efficient fused gather + dequant for the +looked-up rows, so the packed table is never densified. + +Embedding quantization is most useful for **large-vocabulary models**, where the +embedding table is a meaningful fraction of memory (e.g. a 150k-vocab model has a +~1 GB embedding table in fp16). It is **weight-only and data-free** -- no +calibration dataset is required -- and its accuracy impact is typically +negligible. + +## Installation + +```bash +git clone https://github.com/vllm-project/llm-compressor.git +cd llm-compressor +pip install -e . +``` + +## Quickstart + +```bash +python3 llama3_example.py +``` + +The resulting model `Meta-Llama-3-8B-Instruct-embedding-W4A16-G64` is ready to be +loaded into vLLM. + +## Code Walkthrough + +### 1) Load the model + +```python +from transformers import AutoModelForCausalLM, AutoTokenizer + +model_id = "meta-llama/Meta-Llama-3-8B-Instruct" +model = AutoModelForCausalLM.from_pretrained(model_id) +tokenizer = AutoTokenizer.from_pretrained(model_id) +``` + +> Note: `meta-llama/Meta-Llama-3-8B-Instruct` is gated on Hugging Face. Make sure +> your account has access and that you are authenticated locally (e.g. via +> `huggingface-cli login`) before running, or substitute an open model. + +### 2) Apply embedding quantization (data-free) + +Target the `Embedding` module **by class name** rather than by a layer-name +regex. Class-based matching is portable across architectures and does not depend +on the module's prefix, so the same recipe works for `model.embed_tokens`, +`gpt_neox.embed_in`, and so on: + +```python +from llmcompressor import oneshot +from llmcompressor.modifiers.quantization import QuantizationModifier + +recipe = QuantizationModifier( + config_groups={ + "embedding": { + "targets": ["Embedding"], + "weights": { + "num_bits": 4, + "type": "int", + "symmetric": True, + "strategy": "group", + "group_size": 64, + }, + } + } +) + +# Weight-only -> data-free; no calibration dataset needed. +oneshot(model=model, recipe=recipe) +``` + +Use `"strategy": "channel"` (and drop `group_size`) for per-row channel scales, +or `"num_bits": 8` for a more conservative, effectively lossless setting. + +### 3) Save the compressed model + +```python +SAVE_DIR = "Meta-Llama-3-8B-Instruct-embedding-W4A16-G64" +model.save_pretrained(SAVE_DIR, save_compressed=True) +tokenizer.save_pretrained(SAVE_DIR) +``` + +## Composing with weight quantization + +Embedding quantization can be combined with linear-weight quantization by adding +the linear modifier alongside it, e.g.: + +```python +from llmcompressor.modifiers.quantization import GPTQModifier + +recipe = [ + GPTQModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"]), + QuantizationModifier(config_groups={"embedding": embedding_config}), # where embedding_config is defined as above +] +``` + +## Accuracy + +Quantizing only the input embedding is near-lossless. On `pythia-1.4b` +(`lm-eval`, `wikitext` word perplexity / `arc_easy` acc): + +| scheme | wikitext ppl | arc_easy acc | +| --- | --- | --- | +| baseline (fp16) | 14.733 | 0.6048 | +| embedding W8 channel | 14.732 | 0.6052 | +| embedding W4 group-64 | 14.752 | 0.6061 | diff --git a/examples/quantization_embedding/llama3_example.py b/examples/quantization_embedding/llama3_example.py new file mode 100644 index 0000000000..c66f70db7c --- /dev/null +++ b/examples/quantization_embedding/llama3_example.py @@ -0,0 +1,52 @@ +from compressed_tensors.offload import dispatch_model +from transformers import AutoModelForCausalLM, AutoTokenizer + +from llmcompressor import oneshot +from llmcompressor.modifiers.quantization import QuantizationModifier + +# Select model and load it. +model_id = "meta-llama/Meta-Llama-3-8B-Instruct" +model = AutoModelForCausalLM.from_pretrained(model_id) +tokenizer = AutoTokenizer.from_pretrained(model_id) + +# Configure the quantization algorithm to run. +# * quantize the input embedding table to int4 with a group size of 64 +# * weight-only: there are no input/output activations to quantize for a lookup +# +# The target is the `Embedding` module *class* (not a name regex). Matching by +# class name keeps the recipe portable across architectures and independent of a +# model's module prefix, so the same recipe works whether the embedding is named +# `model.embed_tokens`, `gpt_neox.embed_in`, etc. +recipe = QuantizationModifier( + config_groups={ + "embedding": { + "targets": ["Embedding"], + "weights": { + "num_bits": 4, + "type": "int", + "symmetric": True, + "strategy": "group", + "group_size": 64, + }, + } + } +) + +# Apply the quantization algorithm. Embedding quantization is weight-only, so it +# is data-free -- no calibration dataset is required. +oneshot(model=model, recipe=recipe) + +# Confirm generations of the quantized model look sane. +print("\n\n") +print("========== SAMPLE GENERATION ==============") +dispatch_model(model) +sample = tokenizer("Hello my name is", return_tensors="pt") +sample = {key: value.to(model.device) for key, value in sample.items()} +output = model.generate(**sample, max_new_tokens=100) +print(tokenizer.decode(output[0])) +print("==========================================\n\n") + +# Save to disk compressed. +SAVE_DIR = model_id.rstrip("/").split("/")[-1] + "-embedding-W4A16-G64" +model.save_pretrained(SAVE_DIR, save_compressed=True) +tokenizer.save_pretrained(SAVE_DIR)