Repository: huggingface/transformers
Pull request: Fix: Unexpected Keys, Improve `run_compressed`, Rename Test Folder

In the latest release, the removal of `unexpected_keys` filtering for compressed tensors models reintroduced warnings that were previously resolved in [#36152](https://github.com/huggingface/transformers/pull/36152). This PR addresses that regression, enhances the user experience for the `run_compressed` flag, and updates the test folder naming to avoid conflicts and align with conventions.

#### Changes and Objectives
This pull request accomplishes three key improvements:

1. **Restores Filtering of Unexpected Keys for Compressed Tensors Models**  
   - The removal of `unexpected_keys` filtering caused warnings to reappear when loading compressed tensors models. This PR reintroduces the necessary logic by adding `unexpected_keys = hf_quantizer.update_unexpected_keys(model, unexpected_keys, prefix)` in `modeling_utils.py`. This ensures unexpected keys are properly managed during model loading, eliminating warnings and restoring the behavior from [#36152](https://github.com/huggingface/transformers/pull/36152).

2. **Enhances User Experience for `run_compressed` Misconfiguration**  
   - Previously, setting `run_compressed=True` in unsupported cases (e.g., sparsified models or non-compressed quantized models) triggered a `ValueError` and halted execution. This PR improves this by:
     - Adding checks in `quantizer_compressed_tensors.py` to identify unsupported scenarios (`is_sparsification_compressed` or `is_quantized` without `is_quantization_compressed`).
     - Issuing a `logger.warn` message instead of raising an error, notifying users that `run_compressed` is unsupported for the given model type.
     - Automatically setting `run_compressed=False` in these cases, allowing the process to proceed gracefully.
   - This change enhances usability by replacing hard failures with warnings and safe fallbacks.

3. **Renames Test Folder to Avoid Name Collisions**  
   - The test folder `tests/quantization/compressed_tensors` has been renamed to `tests/quantization/compressed_tensors_integration`. This prevents potential name collisions when running `pytest`, ensuring smoother test execution. The new name also aligns with the naming conventions of other integration tests in the repository, improving consistency.

#### Impact
- **Fixes Regression**: Eliminates unexpected keys warnings for compressed tensors models, ensuring a seamless loading experience.
- **Better UX**: Replaces abrupt failures with warnings and automatic corrections for `run_compressed`, making the library more robust and user-friendly.
- **Improved Testing**: Avoids test execution issues.

#### Files Modified
- `src/transformers/modeling_utils.py`: Added `update_unexpected_keys` call to restore filtering.
- `src/transformers/quantizers/quantizer_compressed_tensors.py`: Updated `run_compressed` logic with warnings and overrides.
- `tests/quantization/compressed_tensors/*`: Renamed folder to `compressed_tensors_integration` (including `__init__.py`, `test_compressed_models.py`, and `test_compressed_tensors.py`).


Absolutely! Here's the full markdown with both the **script** and the **output** wrapped in collapsible `<details>` blocks for a clean and structured PR description:

---

### Local Testing

This test verifies that compressed and uncompressed models can be loaded using `AutoModelForCausalLM` with various `run_compressed` settings. It also surfaces any warnings, decompression events, or fallbacks.

<details>
<summary><strong>Loading Script</strong></summary>

```python
from transformers import AutoModelForCausalLM
from transformers.utils.quantization_config import CompressedTensorsConfig
import traceback

# List of model stubs to test
model_stubs = [
    "nm-testing/llama2.c-stories42M-gsm8k-quantized-only-compressed",
    "nm-testing/llama2.c-stories42M-gsm8k-quantized-only-uncompressed",
    "nm-testing/llama2.c-stories42M-gsm8k-sparse-only-compressed",
    "nm-testing/llama2.c-stories42M-gsm8k-sparse-only-uncompressed",
    "nm-testing/llama2.c-stories42M-gsm8k-stacked-compressed",
    "nm-testing/llama2.c-stories42M-gsm8k-stacked-uncompressed",
]

print("\n=== Starting Model Load Tests ===\n")

for stub in model_stubs:
    print("=" * 40)
    print(f"Testing model stub: {stub}")

    # Infer model style (check 'uncompressed' before 'compressed')
    style = "uncompressed" if "uncompressed" in stub else "compressed"

    for run_compressed in [True, False]:
        print(f"\n→ Attempting load with run_compressed={run_compressed} (model style: {style})")

        try:
            model = AutoModelForCausalLM.from_pretrained(
                stub,
                torch_dtype="auto",
                device_map="auto",
                quantization_config=CompressedTensorsConfig(run_compressed=run_compressed),
            )
            print(f"✓ Successfully loaded ({style}, run_compressed={run_compressed})")
        except Exception:
            print(f"✗ Failed to load ({style}, run_compressed={run_compressed})")
            print("Traceback:")
            traceback.print_exc()

    print(f"\nFinished testing: {stub}")
    print("=" * 40 + "\n")

print("=== ✅ All Tests Completed ===")
```

</details>

---

<details>
<summary><strong>Test Output</strong></summary>

```

=== Starting Model Load Tests ===

========================================
Testing model stub: nm-testing/llama2.c-stories42M-gsm8k-quantized-only-compressed

→ Attempting load with run_compressed=True (model style: compressed)
2025-04-02 02:31:01.841947: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2025-04-02 02:31:01.882386: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2025-04-02 02:31:01.882424: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2025-04-02 02:31:01.883778: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2025-04-02 02:31:01.890147: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2025-04-02 02:31:02.655687: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
/home/rahul/upstream-transformers/src/transformers/quantizers/auto.py:212: UserWarning: You passed `quantization_config` or equivalent parameters to `from_pretrained` but the model you're loading already has a `quantization_config` attribute. The `quantization_config` from the model will be used.However, loading attributes (e.g. ['run_compressed']) will be overwritten with the one you passed to `from_pretrained`. The rest will be ignored.
  warnings.warn(warning_msg)
✓ Successfully loaded (compressed, run_compressed=True)

→ Attempting load with run_compressed=False (model style: compressed)
Decompressing model: 56it [00:00, 996.06it/s]
✓ Successfully loaded (compressed, run_compressed=False)

Finished testing: nm-testing/llama2.c-stories42M-gsm8k-quantized-only-compressed
========================================

========================================
Testing model stub: nm-testing/llama2.c-stories42M-gsm8k-quantized-only-uncompressed

→ Attempting load with run_compressed=True (model style: uncompressed)
`run_compressed` is only supported for compressed models.Setting `run_compressed=False`
✓ Successfully loaded (uncompressed, run_compressed=True)

→ Attempting load with run_compressed=False (model style: uncompressed)
✓ Successfully loaded (uncompressed, run_compressed=False)

Finished testing: nm-testing/llama2.c-stories42M-gsm8k-quantized-only-uncompressed
========================================

========================================
Testing model stub: nm-testing/llama2.c-stories42M-gsm8k-sparse-only-compressed

→ Attempting load with run_compressed=True (model style: compressed)
`run_compressed` is only supported for quantized_compressed models and not for sparsified models. Setting `run_compressed=False`
Decompressing model: 75it [00:00, 631.24it/s]
✓ Successfully loaded (compressed, run_compressed=True)

→ Attempting load with run_compressed=False (model style: compressed)
Decompressing model: 75it [00:00, 618.48it/s]
✓ Successfully loaded (compressed, run_compressed=False)

Finished testing: nm-testing/llama2.c-stories42M-gsm8k-sparse-only-compressed
========================================

========================================
Testing model stub: nm-testing/llama2.c-stories42M-gsm8k-sparse-only-uncompressed

→ Attempting load with run_compressed=True (model style: uncompressed)
✓ Successfully loaded (uncompressed, run_compressed=True)

→ Attempting load with run_compressed=False (model style: uncompressed)
✓ Successfully loaded (uncompressed, run_compressed=False)

Finished testing: nm-testing/llama2.c-stories42M-gsm8k-sparse-only-uncompressed
========================================

========================================
Testing model stub: nm-testing/llama2.c-stories42M-gsm8k-stacked-compressed

→ Attempting load with run_compressed=True (model style: compressed)
`run_compressed` is only supported for quantized_compressed models and not for sparsified models. Setting `run_compressed=False`
Decompressing model: 131it [00:00, 278.01it/s]
Decompressing model: 56it [00:00, 5845.43it/s]
✓ Successfully loaded (compressed, run_compressed=True)

→ Attempting load with run_compressed=False (model style: compressed)
Decompressing model: 131it [00:00, 878.16it/s]
Decompressing model: 56it [00:00, 12393.47it/s]
✓ Successfully loaded (compressed, run_compressed=False)

Finished testing: nm-testing/llama2.c-stories42M-gsm8k-stacked-compressed
========================================

========================================
Testing model stub: nm-testing/llama2.c-stories42M-gsm8k-stacked-uncompressed

→ Attempting load with run_compressed=True (model style: uncompressed)
`run_compressed` is only supported for compressed models.Setting `run_compressed=False`
✓ Successfully loaded (uncompressed, run_compressed=True)

→ Attempting load with run_compressed=False (model style: uncompressed)
✓ Successfully loaded (uncompressed, run_compressed=False)

Finished testing: nm-testing/llama2.c-stories42M-gsm8k-stacked-uncompressed
========================================

=== ✅ All Tests Completed ===
```

</details>

### Cases when `run_compressed=True` is not supported and overridden to `False`

- Model is uncompressed
- Model is sparse (both sparse and sparse quantized)

| Model ID                                           | Run Compressed | Overridden Columns                  |
|----------------------------------------------------|----------------|-------------------------------------|
| nm-testing/llama2.c-stories42M-gsm8k-quantized-only-compressed | True           | None                                |
| nm-testing/llama2.c-stories42M-gsm8k-quantized-only-compressed | False          | None                                |
| nm-testing/llama2.c-stories42M-gsm8k-quantized-only-uncompressed | True           | `run_compressed` set to False       |
| nm-testing/llama2.c-stories42M-gsm8k-quantized-only-uncompressed | False          | None                                |
| nm-testing/llama2.c-stories42M-gsm8k-sparse-only-compressed | True           | `run_compressed` set to False       |
| nm-testing/llama2.c-stories42M-gsm8k-sparse-only-compressed | False          | None                                |
| nm-testing/llama2.c-stories42M-gsm8k-sparse-only-uncompressed | True           | `run_compressed` set to False       |
| nm-testing/llama2.c-stories42M-gsm8k-sparse-only-uncompressed | False          | None                                |
| nm-testing/llama2.c-stories42M-gsm8k-stacked-compressed | True           | `run_compressed` set to False       |
| nm-testing/llama2.c-stories42M-gsm8k-stacked-compressed | False          | None                                |
| nm-testing/llama2.c-stories42M-gsm8k-stacked-uncompressed | True           | `run_compressed` set to False       |
| nm-testing/llama2.c-stories42M-gsm8k-stacked-uncompressed | False          | None                                |


Changed files:
filename: src/transformers/modeling_utils.py
patch: @@ -1352,6 +1352,7 @@ def _find_missing_and_unexpected_keys(
 
     if hf_quantizer is not None:
         missing_keys = hf_quantizer.update_missing_keys(model, missing_keys, prefix)
+        unexpected_keys = hf_quantizer.update_unexpected_keys(model, unexpected_keys, prefix)
 
     # Model-specific exceptions for missing and unexpected keys (e.g. if the modeling change over time, or any other reason...)
     if cls._keys_to_ignore_on_load_missing is not None:

filename: src/transformers/quantizers/quantizer_compressed_tensors.py
patch: @@ -46,6 +46,10 @@ def __init__(self, quantization_config: CompressedTensorsConfig, **kwargs):
                 "`pip install compressed-tensors`"
             )
 
+        # Call post_init here to ensure proper config setup when `run_compressed`
+        # is provided directly via CompressedTensorsConfig, and to avoid duplicate logging.
+
+        quantization_config.post_init()
         from compressed_tensors.compressors import ModelCompressor
 
         self.compressor = ModelCompressor.from_compression_config(quantization_config)
@@ -117,16 +121,16 @@ def _process_model_before_weight_loading(self, model, **kwargs):
         ct_quantization_config = self.compressor.quantization_config
 
         if self.run_compressed:
-            if not self.is_quantization_compressed:
-                raise ValueError("`run_compressed` is only supported for quantized_compressed models")
             apply_quantization_config(model, ct_quantization_config, run_compressed=True)
-        elif self.is_quantized and not self.is_quantization_compressed:
+        elif not self.quantization_config.is_quantization_compressed:
             apply_quantization_config(model, ct_quantization_config)
 
     def _process_model_after_weight_loading(self, model, **kwargs):
         """Decompress loaded model if necessary - need for qat"""
 
-        if (self.is_quantization_compressed and not self.run_compressed) or self.is_sparsification_compressed:
+        if (
+            self.quantization_config.is_quantization_compressed and not self.run_compressed
+        ) or self.quantization_config.is_sparsification_compressed:
             config = kwargs.get("config", None)
             cache_path = config._name_or_path
 
@@ -136,44 +140,20 @@ def _process_model_after_weight_loading(self, model, **kwargs):
                 config_file_path = cached_file(cache_path, "config.json")
                 cache_path = os.path.sep.join(config_file_path.split(os.path.sep)[:-1])
 
-            if self.is_quantization_compressed and not self.run_compressed:
+            if self.quantization_config.is_quantization_compressed and not self.run_compressed:
                 from compressed_tensors.quantization import QuantizationStatus
 
                 self.compressor.quantization_config.quantization_status = QuantizationStatus.FROZEN
             self.compressor.decompress(model_path=cache_path, model=model)
 
-    @property
-    def is_quantized(self):
-        return self.quantization_config.quantization_config is not None and bool(
-            self.quantization_config.quantization_config.config_groups
-        )
-
-    @property
-    def is_quantization_compressed(self):
-        from compressed_tensors.quantization import QuantizationStatus
-
-        return (
-            self.quantization_config.quantization_config is not None
-            and self.quantization_config.quantization_config.quantization_status == QuantizationStatus.COMPRESSED
-        )
-
-    @property
-    def is_sparsification_compressed(self):
-        from compressed_tensors.config.base import CompressionFormat
-
-        return (
-            self.quantization_config.sparsity_config is not None
-            and self.quantization_config.sparsity_config.format != CompressionFormat.dense.value
-        )
-
     @property
     def is_trainable(self):
         return True
 
     def is_qat_trainable(self) -> bool:
         """Loaded Models can carry out quantization aware training"""
         # models need to be decompressed carry out qat
-        return not self.run_compressed or not self.is_quantization_compressed
+        return not self.run_compressed or not self.quantization_config.is_quantization_compressed
 
     def is_serializable(self, safe_serialization=None) -> bool:
         """Models quantized using compressed tensors can be saved to disk"""

filename: src/transformers/utils/quantization_config.py
patch: @@ -1263,7 +1263,7 @@ def __init__(
 
         # parse from dict to load nested QuantizationScheme objects
         if config_groups or kv_cache_scheme:
-            self.quantization_config = QuantizationConfig.parse_obj(
+            self.quantization_config = QuantizationConfig.model_validate(
                 {
                     "config_groups": config_groups,
                     "quant_method": quant_method,
@@ -1282,7 +1282,19 @@ def __init__(
                 sparsity_config.get("format"), **sparsity_config
             )
 
-        super().__init__(quant_method=QuantizationMethod.COMPRESSED_TENSORS)
+        self.quant_method = QuantizationMethod.COMPRESSED_TENSORS
+
+    def post_init(self):
+        if self.run_compressed:
+            if self.is_sparsification_compressed:
+                logger.warn(
+                    "`run_compressed` is only supported for quantized_compressed models"
+                    " and not for sparsified models. Setting `run_compressed=False`"
+                )
+                self.run_compressed = False
+            elif not self.is_quantization_compressed:
+                logger.warn("`run_compressed` is only supported for compressed models. Setting `run_compressed=False`")
+                self.run_compressed = False
 
     @classmethod
     def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs):
@@ -1356,6 +1368,28 @@ def to_diff_dict(self) -> Dict[str, Any]:
     def get_loading_attributes(self):
         return {"run_compressed": self.run_compressed}
 
+    @property
+    def is_quantized(self):
+        return bool(self.quantization_config) and bool(self.quantization_config.config_groups)
+
+    @property
+    def is_quantization_compressed(self):
+        from compressed_tensors.quantization import QuantizationStatus
+
+        return self.is_quantized and self.quantization_config.quantization_status == QuantizationStatus.COMPRESSED
+
+    @property
+    def is_sparsification_compressed(self):
+        from compressed_tensors.config import (
+            CompressionFormat,
+            SparsityCompressionConfig,
+        )
+
+        return (
+            isinstance(self.sparsity_config, SparsityCompressionConfig)
+            and self.sparsity_config.format != CompressionFormat.dense.value
+        )
+
 
 @dataclass
 class FbgemmFp8Config(QuantizationConfigMixin):

filename: tests/quantization/compressed_tensors_integration/test_compressed_models.py
patch: @@ -185,6 +185,7 @@ def test_default_run_compressed__True(self):
     def test_default_run_compressed__False(self):
         from compressed_tensors.linear.compressed_linear import CompressedLinear
         from compressed_tensors.quantization.utils import iter_named_leaf_modules
+
         from transformers.utils.quantization_config import CompressedTensorsConfig
 
         quantization_config = CompressedTensorsConfig(run_compressed=False)