[Transform] Spinquant with R1 and R2 (#1615)

## Purpose ##
* Enable offline spinquant-style transforms

## Prerequisites ##
* neuralmagic/compressed-tensors#370
* neuralmagic/compressed-tensors#412
* neuralmagic/compressed-tensors#414

## Changes ##
* Added `spinquant_example.py` to examples folder
* Added `SpinQuantModifier` which handles the construction of a
spinquant-style transform config

## Testing ##
* Added modifier serialization and correctness tests

## Evaluation ##
Using this branch, and [the original SpinQuant
code](https://github.com/facebookresearch/SpinQuant), we see very
similar results for `meta-llama/Llama-3.2-1B-Instruct` with W4A16
quantization. Results are equivalent in hf (in-memory vs serialized and
re-loaded), and very similar in vllm. The symmetric scales calculation
in `llm-compressor` is slightly different than original SpinQuant paper,
which uses the original GPTQ implementation. When this is swapped in,
results are consistent, with hadamard improving results on `gsm8k_llama`
and `arc_challenge_llama`:

Scheme | Impl | gsm8k | gsm8k_llama | arc_challenge_llama
-- | -- | -- | -- | --
Hadamard+W4A16 | LC | 0.2403 | 0.2835 | 0.5262
W4A16 | LC | 0.1964 | 0.1933 | 0.4781
Hadamard+W4A16 | LC+SQscales | 0.1721 | 0.2183 | 0.485
W4A16 | LC+SQscales | 0.207 | 0.1706 | 0.4498
Hadamard+W4A16 | SQ | 0.1736 | 0.2282 | 0.4807
W4A16 | SQ | 0.1986 | 0.1774 | 0.4489

To run LC+SQScales, change [this line in
CT](https://github.com/neuralmagic/compressed-tensors/blob/b2df366797b00330ec765f5891dde14e4cc74c9d/src/compressed_tensors/quantization/utils/helpers.py#L111)
from

```python
scales = max_val_pos / (float(bit_range) / 2)
```
to
```python
scales = max_val_pos / (float(bit_max))
```

<details>
<summary>The following python script was used to generate these
results</summary>

Clone SpinQuant repo and paste this in the top-level directory:
```python
# coding=utf-8
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
from typing import Literal
import os

os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"

from torch import nn
import lm_eval

from transformers import LlamaForCausalLM, AutoTokenizer
import transformers
from train_utils.main import prepare_model
from train_utils.modeling_llama_quant import LlamaForCausalLM as LlamaForCausalLMQuant
from utils.hadamard_utils import random_hadamard_matrix, hadamard_matrix
from utils.process_args import process_args_ptq

# model_id = "meta-llama/Llama-3.1-8B-Instruct"
# model_id = "meta-llama/Llama-3.2-3B-Instruct"
model_id = "meta-llama/Llama-3.2-1B-Instruct"
dtype = torch.bfloat16


class RotateModule(nn.Module):
    def __init__(self, R_init):
        super(RotateModule, self).__init__()
        self.weight = nn.Parameter(R_init.to(torch.float32).to(torch.device("cuda")))

    def forward(self, x, transpose=False):
        if transpose:
            return x @ self.weight
        else:
            return self.weight @ x


def get_sq_model(
    r1r2=Literal["eye", "random-hadamard", "hadamard"],
    w_bits=Literal[4, 16],
    w_clip: bool = False,
) -> LlamaForCausalLMQuant:
    model_args, training_args, ptq_args = process_args_ptq()
    model_args.input_model = model_id
    if w_bits == 4:
        ptq_args.w_bits = 4
        ptq_args.w_groupsize = 128
        ptq_args.w_rtn = True  # if False, GPTQ is used
        ptq_args.w_clip = w_clip
    ptq_args.a_bits = 16
    ptq_args.k_bits = 16
    ptq_args.v_bits = 16

    print("=======ARGS=======", ptq_args)

    config = transformers.AutoConfig.from_pretrained(model_args.input_model)

    # Llama v3.2 specific: Spinquant is not compatiable with tie_word_embeddings, clone lm_head from embed_tokens
    process_word_embeddings = False
    if config.tie_word_embeddings:
        config.tie_word_embeddings = False
        process_word_embeddings = True

    model = LlamaForCausalLMQuant.from_pretrained(
        pretrained_model_name_or_path=model_args.input_model,
        config=config,
        torch_dtype=dtype,
        device_map="cuda",
    )

    if process_word_embeddings:
        model.lm_head.weight.data = model.model.embed_tokens.weight.data.clone()

    model = prepare_model(ptq_args, model)
    for param in model.parameters():
        param.requires_grad = False
    match r1r2:
        case "eye":
            R1 = torch.eye(model.config.hidden_size, device="cuda")
        case "random-hadamard":
            R1 = random_hadamard_matrix(model.config.hidden_size, "cuda")
        case _:
            R1 = hadamard_matrix(model.config.hidden_size, "cuda")
    model.R1 = RotateModule(R1)
    for i in range(model.config.num_hidden_layers):
        # Each head dim = 128 for Llama model
        match r1r2:
            case "eye":
                R2 = torch.eye(
                    model.config.hidden_size // model.config.num_attention_heads,
                    device="cuda",
                )
            case "random-hadamard":
                R2 = random_hadamard_matrix(
                    model.config.hidden_size // model.config.num_attention_heads, "cuda"
                )
            case _:
                R2 = hadamard_matrix(
                    model.config.hidden_size // model.config.num_attention_heads, "cuda"
                )
        model.model.layers[i].self_attn.R2 = RotateModule(R2)

    model.config.use_cache = False

    return model


def get_lc_model(
    r1r2=Literal["eye", "random-hadamard", "hadamard"], w_bits=Literal[4, 16]
) -> LlamaForCausalLM:
    from llmcompressor import oneshot
    from llmcompressor.modifiers.quantization import QuantizationModifier
    from llmcompressor.modifiers.transform import SpinQuantModifier

    model = LlamaForCausalLM.from_pretrained(
        pretrained_model_name_or_path=model_id,
        torch_dtype=dtype,
        device_map="cuda",
    )

    recipe = [
        SpinQuantModifier(
            rotations=[] if r1r2 == "eye" else ["R1", "R2"],
            transform_type="hadamard",
        )
    ]
    if w_bits == 4:
        recipe.append(
            QuantizationModifier(
                targets="Linear",
                scheme="W4A16",
                ignore=["lm_head"],
            )
        )

    oneshot(
        model=model,
        recipe=recipe,
        pipeline="datafree",
        log_dir=None,
    )

    return model


if __name__ == "__main__":
    for scales_impl in ["sq_min_hack", "lc_min_hack"]:
        for r1r2 in ["eye", "hadamard"]:
            for sq_lc in ["sq", "lc"]:
                w_bits = 4

                os.environ["SCALES_IMPL"] = scales_impl

                model = (
                    get_sq_model(r1r2=r1r2, w_bits=w_bits)
                    if sq_lc == "sq"
                    else get_lc_model(r1r2=r1r2, w_bits=w_bits)
                ).to("cuda")

                SAVE_DIR = model_id.split("/")[1] + f"-{scales_impl}-{r1r2}-w4a16"
                model.save_pretrained(SAVE_DIR, save_compressed=True)
                tokenizer = AutoTokenizer.from_pretrained(
                    model_id, trust_remote_code=True
                )
                tokenizer.save_pretrained(SAVE_DIR)

                del model
                del tokenizer
                torch.cuda.empty_cache()

                results = lm_eval.simple_evaluate(
                    # 1) hf in-memory
                    # model=lm_eval.models.huggingface.HFLM(
                    #     pretrained=model,
                    #     batch_size=32,
                    #     add_bos_token=False,
                    # ),
                    # 1/)
                    # 2) vllm serialized
                    model="vllm",
                    model_args={
                        "pretrained": SAVE_DIR,
                        "add_bos_token": False,
                        "dtype": "auto",
                        "max_model_len": 4096,
                        "gpu_memory_utilization": 0.5,
                        "enable_chunked_prefill": True,
                    },
                    # 2/)
                    # 3) hf serialized
                    # model="hf",
                    # model_args={
                    #     "pretrained": SAVE_DIR,
                    #     "add_bos_token": False,
                    #     "dtype": "auto",
                    # },
                    # device="cuda",
                    # 3/)
                    tasks=["gsm8k_llama", "gsm8k", "arc_challenge_llama"],
                    num_fewshot=8,
                    batch_size=32,
                    apply_chat_template=True,
                    fewshot_as_multiturn=True,
                )
                print(
                    f"RESULTS, {model_id} {sq_lc} R1R2 {r1r2} W_BITS {w_bits} SCALEIMPL {scales_impl}"
                )
                print(lm_eval.utils.make_table(results))
```
</details>


## Follow Ups ##
* Infer data free pipeline, even if a transform modifier is included
* Rotations R3 and R4
* Modify example to use GPTQ once basic evaluation has been performed

---------

Signed-off-by: Kyle Sayers <[email protected]>
Signed-off-by: Brian Dellabetta <[email protected]>
Co-authored-by: Kyle Sayers <[email protected]>

Author: brian-dellabetta

examples/transform/spinquant_example.py

@@ -0,0 +1,40 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from llmcompressor import oneshot
+from llmcompressor.modifiers.quantization import QuantizationModifier
+from llmcompressor.modifiers.transform import SpinQuantModifier
+from llmcompressor.utils import dispatch_for_generation
+
+# Select model and load it.
+MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"
+
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+# NOTE: currently only fused rotations (R1 & R2) are available
+# Learned rotations and online rotations (R3 & R4) will be added
+# in a future release.
+# Configure the quantization algorithm to run.
+#   * apply spinquant transforms to model to reduce quantization loss
+#   * quantize the weights to 4 bit with group size 128
+recipe = [
+    SpinQuantModifier(rotations=["R1", "R2"], transform_type="hadamard"),
+    QuantizationModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"]),
+]
+
+# Apply algorithms.
+oneshot(model=model, recipe=recipe, pipeline="datafree")
+
+# Confirm generations of the quantized model look sane.
+print("\n\n")
+print("========== SAMPLE GENERATION ==============")
+dispatch_for_generation(model)
+input_ids = tokenizer("Hello my name is", return_tensors="pt").input_ids.to("cuda")
+output = model.generate(input_ids, max_new_tokens=100)
+print(tokenizer.decode(output[0]))
+print("==========================================\n\n")
+
+# Save to disk compressed.
+SAVE_DIR = MODEL_ID.split("/")[1] + "-spinquantR1R2-w4a16"
+model.save_pretrained(SAVE_DIR, save_compressed=True)
+tokenizer.save_pretrained(SAVE_DIR)

src/llmcompressor/modeling/__init__.py

@@ -1,3 +1,4 @@
 # flake8: noqa
 
+from .fuse import *
 from .prepare import *

src/llmcompressor/modifiers/transform/__init__.py

@@ -0,0 +1,3 @@
+# flake8: noqa
+
+from .spinquant import SpinQuantModifier

src/llmcompressor/modifiers/transform/spinquant/__init__.py

@@ -0,0 +1,3 @@
+# flake8: noqa
+
+from .base import *

src/llmcompressor/modifiers/transform/spinquant/base.py

@@ -0,0 +1,246 @@
+from enum import Enum
+from typing import Iterable, List, Literal, Optional
+
+import torch
+from compressed_tensors import match_modules_set, match_named_modules
+from compressed_tensors.transform import (
+    TransformArgs,
+    TransformConfig,
+    TransformScheme,
+    apply_transform_config,
+)
+from compressed_tensors.utils import TorchDtype
+from pydantic import Field, ValidationInfo, field_validator
+from transformers import PreTrainedModel
+
+from llmcompressor.core import Event, EventType, State
+from llmcompressor.modeling import center_embeddings, fuse_norm_linears
+from llmcompressor.modifiers import Modifier
+
+from .mappings import SpinQuantMapping, infer_mapping_from_model
+from .norm_mappings import NormMapping, infer_norm_mapping_from_model
+
+
+class SpinquantRotation(str, Enum):
+    R1 = "R1"
+    R2 = "R2"
+    R3 = "R3"
+    R4 = "R4"
+
+
+class SpinQuantModifier(Modifier, use_enum_values=True):
+    """
+    Implements the transforms according to "SpinQuant: LLM quantization
+    with learned rotations" (https://arxiv.org/abs/2405.16406)
+
+    Transforms (rotations) are extra layers added to a model which reduce the accuracy
+    loss induced by quantization. This is achived through "rotating" weights and
+    activations into a space with a smaller dynamic range of values, thus decreasing
+    the range of scales required for quantization.
+
+    The SpinQuant authors describe four different rotations which can be applied to a
+    model. R1 and R2 are "offline" rotations, meaning that they can be fused into
+    existing weights and therefore do not induce runtime cost. R3 and R4 are "online"
+    rotations, meaning that they require additional computation at runtime.
+
+    Lifecycle:
+        - on_initialize
+            - infer SpinQuantMappings & NormMappings
+            - as needed, create transform schemes for R1, R2, R3, & R4
+        - on_start
+            - normalize embeddings
+            - fuse norm layers into subsequent Linear layers
+            - apply TransformConfig
+                - fuse transforms into weights for mergeable transforms
+                - add hooks for online transforms
+        - on sequential epoch end
+        - on_end
+        - on_finalize
+
+    :param rotations: A list containing the names of rotations to apply to the model.
+        Possible rotations include R1, R2, R3, and R4
+    :param transform_type: The type of transform to apply to the model.
+        `"hadamard"` has the least performance cost but only supports sizes which are
+        powers of power of two.
+        `"random-matrix"` has more performance cost, but supports a much larger set of
+            sizes.
+        `"random-matrix"` has the greatest performance cost, but supports any size
+    :param randomize: if True, create distinct transforms for each application
+    :param learnable: if True, attach gradients to transform weights for training
+    :param precision: Precision at which all transforms should be applied. This applies
+        to both weight fusing and online rotations
+    :param mappings: Specifies layers within a model to target for transforms.
+        A mapping will be inferred if None is provided
+    :param norm_mappings: Specifies layers within a model to target for norm fusing.
+        A mapping will be inferred if None is provided
+    :param transform_config: Optional transform config for overriding provided arguments
+    """
+
+    rotations: List[SpinquantRotation] = Field(default_factory=lambda: ["R1", "R2"])
+    transform_type: Literal["hadamard", "random-hadamard", "random-matrix"] = Field(
+        default="hadamard"
+    )
+    randomize: bool = Field(default=False)
+    learnable: bool = Field(default=False)
+    precision: TorchDtype = Field(default=torch.float64)
+
+    # norm mappings separate from spinquant mappings to allow users to
+    # override spinquant mappings with transform_config without overriding norms
+    mappings: Optional[SpinQuantMapping] = Field(
+        default=None,
+        repr=False,
+        exclude=True,
+    )
+    norm_mappings: Optional[List[NormMapping]] = Field(
+        default=None,
+        repr=False,
+        exclude=True,
+    )
+
+    # optional override for more fine-grained control
+    # also included in recipe serialization
+    transform_config: Optional[TransformConfig] = Field(default=None, repr=False)
+
+    @field_validator("randomize", "learnable", mode="before")
+    def validate_not_implemented(cls, value, info: ValidationInfo):
+        if value:
+            raise NotImplementedError(f"{info.field_name} is not supported right now")
+        return value
+
+    @field_validator("rotations", mode="before")
+    def validate_rotations(cls, value):
+        if isinstance(value, Iterable):
+            return tuple(v.upper() for v in value)
+        return value
+
+    def on_initialize(self, state: State, **kwargs) -> bool:
+        if self.transform_config is not None:
+            return True
+
+        self.mappings = infer_mapping_from_model(state.model)
+        self.norm_mappings = infer_norm_mapping_from_model(state.model)
+
+        config_groups = {}
+        if SpinquantRotation.R1 in self.rotations:
+            config_groups["R1"] = self._create_r1_scheme()
+
+        if SpinquantRotation.R2 in self.rotations:
+            config_groups["R2"] = self._create_r2_scheme(state.model)
+
+        if SpinquantRotation.R3 in self.rotations:
+            config_groups["R3"] = self._create_r3_scheme()
+
+        if SpinquantRotation.R4 in self.rotations:
+            config_groups["R4"] = self._create_r4_scheme()
+
+        self.transform_config = TransformConfig(config_groups=config_groups)
+
+        return True
+
+    def on_start(self, state: State, event: Event, **kwargs):
+        self.started_ = True
+
+        # needs to happen after the model has been hooked to execute on the GPU
+        # otherwise we're applying weight transforms on CPU
+        self._center_embeddings(state.model)
+        self._fuse_norms(state.model)
+        apply_transform_config(state.model, self.transform_config)
+
+    def on_event(self, state: State, event: Event, **kwargs):
+        if event.type_ == EventType.CALIBRATION_EPOCH_START:
+            if not self.started_:
+                self.on_start(state, None)
+
+        elif event.type_ == EventType.SEQUENTIAL_EPOCH_END:
+            pass
+
+        elif event.type_ == EventType.CALIBRATION_EPOCH_END:
+            if not self.ended_:
+                self.on_end(state, None)
+
+    def on_end(self, state: State, event: Event, **kwargs):
+        self.ended_ = True
+
+    def on_finalize(self, state: State, **kwargs) -> bool:
+        if not self.ended_:
+            self.on_end(state, None)
+
+        return True
+
+    def _center_embeddings(self, model: PreTrainedModel):
+        for _, embedding in match_named_modules(
+            model, [self.mappings.embedding], warn_on_fail=True
+        ):
+            center_embeddings(embedding)
+
+    def _fuse_norms(self, model: PreTrainedModel):
+        for mapping in self.norm_mappings:
+            for norm, *linears in match_modules_set(
+                model, (mapping.norm, *mapping.linears)
+            ):
+                fuse_norm_linears(norm, linears)
+
+    def _create_r1_scheme(self) -> TransformScheme:
+        return TransformScheme(
+            type=self.transform_type,
+            randomize=self.randomize,
+            requires_grad=self.learnable,
+            precision=self.precision,
+            apply=[
+                TransformArgs(
+                    targets=[
+                        self.mappings.embedding,
+                        self.mappings.attn_o,
+                        *self.mappings.mlp_out,
+                    ],
+                    location="weight_output",
+                ),
+                TransformArgs(
+                    targets=[
+                        self.mappings.attn_q,
+                        self.mappings.attn_k,
+                        self.mappings.attn_v,
+                        *self.mappings.mlp_in,
+                        self.mappings.lm_head,
+                    ],
+                    location="weight_input",
+                    inverse=True,
+                ),
+            ],
+        )
+
+    def _create_r2_scheme(self, model: PreTrainedModel) -> TransformScheme:
+        config = model.config
+
+        if hasattr(config, "head_dim"):
+            head_dim = config.head_dim
+        elif hasattr(config, "hidden_size") and hasattr(config, "num_attention_heads"):
+            head_dim = config.hidden_size // config.num_attention_heads
+        else:
+            raise NotImplementedError()
+
+        return TransformScheme(
+            type=self.transform_type,
+            randomize=self.randomize,
+            requires_grad=self.learnable,
+            precision=self.precision,
+            head_dim=head_dim,
+            apply=[
+                TransformArgs(targets=[self.mappings.attn_v], location="weight_output"),
+                TransformArgs(
+                    targets=[self.mappings.attn_o],
+                    location="weight_input",
+                    inverse=True,
+                ),
+            ],
+        )
+
+    def _create_r3_scheme(self) -> TransformScheme:
+        raise NotImplementedError(
+            "SpinQuant R3 and R4 rotations will be added in a future release"
+        )
+
+    def _create_r4_scheme(self) -> TransformScheme:
+        raise NotImplementedError(
+            "SpinQuant R3 and R4 rotations will be added in a future release"
+        )

src/llmcompressor/modifiers/transform/spinquant/mappings.py

@@ -0,0 +1,76 @@
+from typing import Dict, List, Optional
+
+from loguru import logger
+from pydantic import BaseModel, Field, field_validator
+from transformers import PreTrainedModel
+
+__all__ = ["SpinQuantMapping", "infer_mapping_from_model"]
+
+
+class SpinQuantMapping(BaseModel):
+    """
+    SpinQuant needs to know the entire architecture of the model,
+    as R1, R2, R3, and R4 rotations need to be applied to specific
+    layers (https://arxiv.org/pdf/2405.16406 Fig. 1).
+
+    :param embedding: name or regex of embedding layer
+    :param attn_q: name or regex of q_proj layer in attention block
+    :param attn_k: name or regex of k_proj layer in attention block
+    :param attn_v: name or regex of v_proj layer in attention block
+    :param attn_o: name or regex of o_proj layer in attention block
+    :param attn_head_dim: head_dim of the attention module, needed
+        because R2 needs to be applied "head-wisely" to v_proj and
+        o_proj
+    :param mlp_in: list of names or regexes for the mlp blocks that
+        receive the input to the MLP block, usually up_proj and gate_proj
+    :param mlp_out: list of names or regexes for the mlp blocks that
+        consitute the output of the MLP block, usually down_proj
+    """
+
+    embedding: str
+
+    attn_q: str
+    attn_k: str
+    attn_v: str
+    attn_o: str
+    attn_head_dim: Optional[int] = Field(default=None)
+
+    mlp_in: List[str]  # up_proj, gate_proj
+    mlp_out: List[str]  # down_proj
+
+    lm_head: str
+
+    @field_validator("mlp_in", "mlp_out", mode="before")
+    def cast_to_list(cls, value):
+        if isinstance(value, str):
+            return [value]
+
+        return value
+
+
+_default_mappings = SpinQuantMapping(
+    embedding="re:.*embed_tokens$",
+    attn_q="re:.*q_proj$",
+    attn_k="re:.*k_proj$",
+    attn_v="re:.*v_proj$",
+    attn_o="re:.*o_proj$",
+    mlp_in=["re:.*up_proj$", "re:.*gate_proj$"],
+    mlp_out="re:.*down_proj$",
+    lm_head="lm_head",
+)
+
+
+SPINQUANT_MAPPING_REGISTRY: Dict[str, SpinQuantMapping] = {
+    "LlamaForCausalLM": _default_mappings,
+}
+
+
+def infer_mapping_from_model(model: PreTrainedModel) -> SpinQuantMapping:
+    architecture = model.__class__.__name__
+    if architecture not in SPINQUANT_MAPPING_REGISTRY:
+        logger.info(
+            f"Unrecognized model architecture {architecture}. "
+            "Falling back to default mappings"
+        )
+
+    return SPINQUANT_MAPPING_REGISTRY.get(architecture, _default_mappings)