feat: add iMatrix weighted MSE observer and IMatrixGatherer

- imatrix_mse observer with E[x²] importance weighting
- IMatrixGatherer transform using match_named_modules + CPU offload
- Unit tests for observer and gatherer
- E2E integration tests

RFC #2456

Signed-off-by: Gilles Turpin <turpingilles15@gmail.com>

Author: Yatimai

examples/quantization_w4a16/README.md

@@ -138,6 +138,105 @@ We can see the resulting scores look good!
 |     |       |strict-match    |     5|exact_match|↑  |0.720|±  |0.0285|
 ```
 
+---
+
+## iMatrix Importance-Weighted Quantization
+
+`imatrix_mse` is an observer that uses per-channel activation importance (E[x²]) to weight quantization error during range selection. Channels that carry more signal get more careful range optimization.
+
+Two components work together:
+- **`IMatrixGatherer`**: collects E[x²] per input channel via forward pre-hooks during calibration
+- **`imatrix_mse` observer**: extends the MSE grid search with importance weighting: `err = sum(importance * |Q(w) - w|^p)`
+
+> See [RFC #2456](https://github.com/vllm-project/llm-compressor/discussions/2456) for the full design discussion.
+
+### Usage
+
+```bash
+python3 llama3_imatrix_example.py
+```
+
+The simplest setup uses `preset_name_to_scheme` to configure W4A16 and swaps in the `imatrix_mse` observer:
+
+```python
+from compressed_tensors.quantization import preset_name_to_scheme
+from llmcompressor.modifiers.quantization import QuantizationModifier
+from llmcompressor.modifiers.transform.imatrix import IMatrixGatherer
+
+scheme = preset_name_to_scheme("W4A16", ["Linear"])
+scheme.weights.observer = "imatrix_mse"
+
+recipe = [
+    IMatrixGatherer(ignore=["lm_head"]),
+    QuantizationModifier(
+        config_groups={"group_0": scheme},
+        ignore=["lm_head"],
+    ),
+]
+```
+
+### Composing with GPTQ
+
+iMatrix composes with GPTQ by providing importance-weighted ranges for the Hessian-based rounding:
+
+```python
+from llmcompressor.modifiers.gptq import GPTQModifier
+
+scheme = preset_name_to_scheme("W4A16", ["Linear"])
+scheme.weights.observer = "imatrix_mse"
+
+recipe = [
+    IMatrixGatherer(ignore=["lm_head"]),
+    GPTQModifier(
+        config_groups={"group_0": scheme},
+        ignore=["lm_head"],
+    ),
+]
+```
+
+### Results
+
+W4A16, Llama-3.1-8B, WikiText-2 token-level perplexity (141 chunks x 2048):
+
+**group_size=128:**
+
+| Config | PPL |
+|---|---|
+| FP16 baseline | 6.24 |
+| RTN `memoryless_minmax` | 6.96 |
+| RTN `imatrix_mse` | 6.97 |
+| GPTQ | 6.89 |
+| GPTQ + `imatrix_mse` | 6.82 |
+
+**group_size=32:**
+
+| Config | PPL |
+|---|---|
+| RTN `memoryless_minmax` | 6.74 |
+| RTN `imatrix_mse` | 6.73 |
+| GPTQ | 6.70 |
+| GPTQ + `imatrix_mse` | 6.66 |
+
+GPTQ + `imatrix_mse` is the best result at both group sizes with default observer settings. iMatrix never degrades quality.
+
+### Observer Parameters
+
+The observer accepts optional `observer_kwargs` for fine-tuning:
+
+| Parameter | Default | Description |
+|---|---|---|
+| `norm` | 2.4 | Error exponent (`\|Q(w) - w\|^norm`) |
+| `maxshrink` | 0.20 | Max fraction to shrink the range |
+| `grid` | 20 | Number of grid search steps |
+| `patience` | 5 | Early stopping after N steps without improvement |
+| `maxgrow` | 0.0 | Max fraction to grow the range beyond observed min/max |
+
+The defaults work well for GPTQ composition. For RTN, increasing `maxshrink` (e.g. 0.80) allows the observer to optimize ranges more aggressively:
+
+```python
+scheme.weights.observer_kwargs = {"maxshrink": 0.80}
+```
+
 ### Questions or Feature Request?
 
 Please open up an issue on `vllm-project/llm-compressor`

examples/quantization_w4a16/llama3_imatrix_example.py

@@ -0,0 +1,61 @@
+from compressed_tensors.offload import dispatch_model
+from compressed_tensors.quantization import preset_name_to_scheme
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from llmcompressor import oneshot
+from llmcompressor.modifiers.quantization import QuantizationModifier
+from llmcompressor.modifiers.transform.imatrix import IMatrixGatherer
+
+# Select model and load it.
+model_id = "meta-llama/Meta-Llama-3.1-8B"
+
+model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto")
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+# Select calibration dataset.
+DATASET_ID = "open_platypus"
+
+# 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.
+#   * collect per-channel importance statistics (E[x²]) with IMatrixGatherer
+#   * quantize the weights to 4 bit with group size 128
+#   * use imatrix_mse observer to weight quantization error by channel importance
+scheme = preset_name_to_scheme("W4A16", ["Linear"])
+scheme.weights.observer = "imatrix_mse"
+
+recipe = [
+    IMatrixGatherer(ignore=["lm_head"]),
+    QuantizationModifier(
+        config_groups={"group_0": scheme},
+        ignore=["lm_head"],
+    ),
+]
+
+# Apply algorithms.
+oneshot(
+    model=model,
+    dataset=DATASET_ID,
+    splits={"calibration": "train[:5%]"},
+    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_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] + "-W4A16-G128-imatrix"
+model.save_pretrained(SAVE_DIR, save_compressed=True)
+tokenizer.save_pretrained(SAVE_DIR)

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

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