[Quantization] Attention/ KV Cache Refactor

[kylesayrs]

Purpose

• Support fully-expressive attention and kv cache quantization
• Support running kv cache quantization evals with hf transformers
• Resolves [Bug]: KV cache quant - cannot be loaded in vllm #1949
• Resolves [Bug]: k_scale and v_scale is zero after kv cache fp8 quantization #1928

recipe = QuantizationModifier(
    config_groups={
        "attention": QuantizationScheme(
            targets=["LlamaAttention"],
            input_activations=QuantizationArgs(
                num_bits=8, type="float", strategy="tensor"
            ),
        )
    }
)

{
  "quantization_config": {
    "config_groups": {
      "group_0": {
        "format": null,
        "input_activations": {
          "dynamic": false,
          "num_bits": 8,
          "observer": "minmax",
          "strategy": "tensor",
          "symmetric": true,
          "type": "float"
        },
        "output_activations": null,
        "targets": [
          "LlamaAttention"
        ],
        "weights": null
      }
    },
    "format": "dense",
    "ignore": [],
    "kv_cache_scheme": {
      "dynamic": false,
      "group_size": null,
      "num_bits": 8,
      "observer": "minmax",
      "strategy": "tensor",
      "symmetric": true,
      "type": "float"
    },
    "quant_method": "compressed-tensors",
    "quantization_status": "frozen",
  },
}

Prerequisites

• Must be merged at the same time as [Transform] Attention/Cache transforms compressed-tensors#436

Changes

• Replace hooks

  • Remove calibrate_kv_cache_input_hook, calibrate_kv_cache_output_hook, initialize_quantized_kv_cache
  • Add calibrate_query_hook calibrate_key_hook, calibrate_value_hook
  • QuantizationMixin now initializes "q", "k", and "v" obsevers (depending on the attached submodules) and adds the appropriate hooks

• Miscellaneous

  • Fix minor shape bug in _flatten_attention
  • Add support for "attn_head" strategy in _flatten_attention

• Tests

  • Removed old QuantizationKVCache tests (these classes are now tested [here])(https://github.com/neuralmagic/compressed-tensors/pull/436/files#diff-6e33ff48047dc4f7c9d969293f87e32e4d5ec3f3e8b741ea757780c8c0aab775)
  • Updated scale names to avoid using enum
  • Avoid unnecessary tokenization to reduce runtime

Testing

• Kv cache regression tests pass
• Able to quantize attention with scripts (will add to examples once loadable in vllm)
  • kylesayrs/Llama-3.2-1B-Instruct-attention-fp8-head
  • kylesayrs/Llama-3.2-1B-Instruct-attention-nvfp4-head
• Nightly passes (in progress)

Evaluation

eval.py

import sys
import lm_eval

model_id = sys.argv[1]

print(model_id)
results = lm_eval.simple_evaluate(
    # 3) hf serialized
    model="hf",
    model_args={
        "pretrained": model_id,
        "add_bos_token": False,
        "dtype": "auto",
        "device_map": "cuda",
        #"max_length": 128000,
    },
    device="cuda",
    # 3/)

    #tasks=["gsm8k_platinum", "mmlu_llama", "longbench2_single"],
    tasks=["gsm8k_platinum"],
    batch_size=64,
    apply_chat_template=True,
    fewshot_as_multiturn=True,
)
print(model_id)
print(lm_eval.utils.make_table(results))

compress.py

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer

from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier
from llmcompressor.utils import dispatch_for_generation
from compressed_tensors.quantization import QuantizationScheme, QuantizationArgs

# Select model and load it.
#model_id = "Qwen/Qwen2.5-14B-Instruct-1M"
model_id = "meta-llama/Llama-3.1-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto")
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Select calibration dataset.
DATASET_ID = "ultrachat_200k"
DATASET_SPLIT = "train_sft"

# Select number of samples. 512 samples is a good place to start.
# Increasing the number of samples can improve accuracy.
NUM_CALIBRATION_SAMPLES = 512
MAX_SEQUENCE_LENGTH = 2048

# Configure the quantization algorithm to run.
args = QuantizationArgs(
    num_bits=8,
    type="float",
    strategy="attn_head",
    symmetric=True,
    observer="static_minmax",
)
recipe = QuantizationModifier(
    # config_groups={
    #     "attention": QuantizationScheme(
    #         #targets=["Qwen2Attention"],
    #         targets=["LlamaAttention"],
    #         input_activations=args,
    #     )
    # }
    kv_cache_scheme=args,
)

# Apply algorithms.
oneshot(
    model=model,
    dataset=DATASET_ID,
    splits={"calibration": f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]"},
    recipe=recipe,
    max_seq_length=MAX_SEQUENCE_LENGTH,
    num_calibration_samples=NUM_CALIBRATION_SAMPLES,
)

# Confirm generations of the quantized model look sane.
print("\n\n")
print("========== SAMPLE GENERATION ==============")
dispatch_for_generation(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] + f"-KV-FP8-{args.strategy}-{args.observer}"
model.save_pretrained(SAVE_DIR, save_compressed=True)
tokenizer.save_pretrained(SAVE_DIR)

Model	GSM8K
nm-testing/Llama-3.1-8B-Instruct	0.8337
nm-testing/Llama-3.1-8B-Instruct-KV-FP8-Tensor	0.8271
nm-testing/Llama-3.1-8B-Instruct-KV-FP8-Head	0.8354
nm-testing/Llama-3.1-8B-Instruct-QKV-FP8-Tensor	0.8321
nm-testing/Llama-3.1-8B-Instruct-QKV-FP8-Head	0.8238

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[gemini-code-assist]

Summary of Changes

Hello @kylesayrs, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request integrates the foundational components for applying 'online rotations' (specifically R1 and R2 from the SpinQuant paper) into the llmcompressor framework. It primarily introduces a new SpinQuantModifier that leverages novel model transformation utilities, such as embedding normalization and norm-linear fusion, to prepare models for more effective quantization. Additionally, it refines the handling of tied word embeddings, ensuring compatibility and robustness across various model configurations.

Highlights

• New Feature: SpinQuantModifier: Introduced a new SpinQuantModifier to apply 'offline' rotations (R1 and R2) from the SpinQuant paper. These rotations transform model weights and activations to improve quantization accuracy without introducing runtime overhead.
• Model Transformation Utilities: Added new utilities for normalizing embedding layers and fusing norm layers into subsequent linear layers. These are crucial preprocessing steps for applying SpinQuant rotations, ensuring transform invariance.
• Improved Tied Word Embedding Handling: Refactored and enhanced the utility for untying word embeddings. The updated implementation is more robust, correctly handling cases where embeddings are tied, especially with offloaded parameters, and centralizes the untying logic.
• Example Usage and Integration: Provided new example scripts (compress_model.py, spinquant_example.py) demonstrating how to use the SpinQuantModifier for model compression. The modifier is also integrated into the data-free pipeline for seamless application.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces support for SpinQuant online rotations, a technique for improving quantization performance. It adds a new SpinQuantModifier, along with utilities for model transformation like layer fusion and embedding normalization. The changes also include updates to the data-free pipeline, improvements to handling tied word embeddings, and new example scripts and tests.

My review identified a critical bug in a Pydantic validator within the new SpinQuantModifier that prevents it from being used. I've also pointed out a few medium-severity issues, including a required argument missing in a script, brittle directory name construction, a documentation typo, and a maintainability concern with a hardcoded pipeline selection. Addressing these points will improve the correctness and robustness of the new features.

[kylesayrs]

https://github.com/neuralmagic/llm-compressor-testing/actions/runs/17246609866

[kylesayrs]

https://github.com/neuralmagic/llm-compressor-testing/actions/runs/18477528039

[kylesayrs]

Last nightly worked, but e2e failed due to model storage issues
https://github.com/neuralmagic/llm-compressor-testing/actions/runs/18483826999

[HDCharles]

looks a lot cleaner

[rahul-tuli]

Nice!

The base branch was changed.

[kylesayrs]

https://github.com/neuralmagic/llm-compressor-testing/actions/runs/18663440870