Improve model loading for compressed tensor models
Author: rahul-tuli
Pull request: yes

This PR improves the `from_pretrained` model-loading pipeline in Hugging Face `transformers` by **suppressing unnecessary warnings** for models saved using **compressed tensors** (e.g., sparse and quantized models). 

Currently, models that store **compression metadata** (e.g., `bitmask`, `compressed`, `row_offsets`) instead of direct weight tensors trigger misleading warnings, Due to our compression method, certain extra and missing keys are expected:

- **Extra keys**: Represent compression metadata stored on disk but not explicitly expected by the model graph.
- **Missing keys**: Expected in the model graph but omitted from storage since they can be reconstructed from compression metadata.

This PR fixes these issues by **introducing key filtering mechanisms** through Hugging Face’s `hf_quantizer` extension points.

---

## **Problem: Unnecessary Warnings for Compressed Models**
### 🚨 **Current Behavior (Before Fix)**
When loading compressed models, users encounter:
- ❌ **"Missing key" warnings** for weight tensors that will be reconstructed.
- ❌ **"Unexpected key" warnings** for compression metadata in the checkpoint.

### ✅ **Expected Behavior (After Fix)**
- ✅ No unnecessary **missing key** warnings for compressed weights.
- ✅ No unnecessary **unexpected key** warnings for compression metadata.
- ✅ **Seamless integration** with Hugging Face’s model-loading process.

---

## **Solution: Key Filtering for Compressed Models**
### **🔹 Step 1: Suppress Warnings for Expected Missing Keys**
🔧 **Fix:**  
- `update_missing_keys_after_loading()` removes weight keys (e.g., `.*weight`) from `missing_keys`.  
- Since these weights are **reconstructed later**, the warning is unnecessary.

<p><strong>💡 Impact:</strong> No more misleading "missing key" warnings for compressed models.</p>
<hr>
<h3><strong>🔹 Step 2: Suppress Warnings for Expected Unexpected Keys</strong></h3>
<p>🔧 <strong>Fix:</strong></p>
<ul>
<li><code inline="">update_unexpected_keys()</code> removes <strong>compression parameters</strong> (e.g., <code inline="">bitmask</code>, <code inline="">compressed</code>, <code inline="">row_offsets</code>) from <code inline="">unexpected_keys</code>.</li>
<li>These are <strong>metadata, not actual model weights</strong>, and should not raise warnings.</li>
</ul>
<p><strong>💡 Impact:</strong> No more misleading "unexpected key" warnings from compression metadata.</p>
<hr>
<h3><strong>🔹 Step 3: Seamless Integration with Model Loading</strong></h3>
<ul>
<li>✅ Uses <strong>existing Hugging Face extension points</strong> (<code inline="">hf_quantizer</code>)—no changes to core <code inline="">transformers</code> code.</li>
<li>✅ <strong>Standard models remain unaffected</strong>—only compressed models benefit.</li>
<li>✅ <strong>Ensures genuine issues (e.g., truly missing parameters) still raise warnings.</strong></li>
</ul>
<hr>
<h2><strong>Lifecycle Overview: Before and After the Fix</strong></h2>

Step | Before (Current Behavior) ❌ | After (With Fix) ✅
-- | -- | --
6 | Load checkpoint (compressed weights missing!) | Load checkpoint (compressed weights missing!)
7 | Identify missing keys ⚠ (Triggers Warning!) | Identify missing keys (Before Filtering) ✅
8 | Identify unexpected keys ⚠ (Triggers Warning!) | Filter missing keys (Removes .*weight) ✅
9 | Raise unnecessary warnings ❌ | Identify unexpected keys (Before Filtering) ✅
10 | Assign weights (weights missing!) | Filter unexpected keys (Removes compression params) ✅
14 | Set model to evaluation mode (with unnecessary warnings) | Apply quantization postprocessing (Reconstructs weights) ✅


<hr>
<h2><strong>Testing</strong></h2>
<p>This PR has been tested across multiple model configurations to ensure correctness:</p>
<ul>
<li>✅ <strong>Quantized-only models</strong> – No unexpected warnings.</li>
<li>✅ <strong>Sparse-only models</strong> – No missing or unexpected key warnings.</li>
<li>✅ <strong>Stacked cases (both sparse and quantized)</strong> – Loads correctly without unnecessary warnings.</li>
</ul>
<hr>
<h2><strong>Why This Should Be Merged</strong></h2>
<ul>
<li>🚀 <strong>Fixes a real issue</strong> affecting compressed models.</li>
<li>✅ <strong>No impact on standard model-loading workflows</strong>.</li>
<li>🔧 <strong>Uses Hugging Face’s extension points</strong>—no core code changes.</li>
<li>🔄 <strong>Maintains extensibility</strong> while improving compressed model support.</li>
</ul>
<hr>
<h2><strong>Dependencies</strong></h2>
<p>This PR depends on <a href="https://github.com/neuralmagic/compressed-tensors/pull/250">neuralmagic/compressed-tensors#250</a>, which introduces the necessary filtering mechanisms at the compression framework level.</p>
<hr>
<h2><strong>Conclusion</strong></h2>
<p>With this fix:
✅ Expected missing keys (e.g., <code inline="">.*weight</code>) <strong>do not trigger warnings</strong>.<br>
✅ Known compression parameters <strong>do not raise unexpected key warnings</strong>.<br>
✅ Weights <strong>are reconstructed properly in postprocessing</strong>.<br>
✅ The <strong>Hugging Face <code inline="">from_pretrained</code> API remains modular and extensible</strong>.</p>
<p>This significantly improves the developer experience when working with compressed models. <strong>Looking forward to feedback and merging this in!</strong> 🚀</p>


Comments:
cc @SunMarc @MekkCyber for quantization!

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_36152). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Also make sure to rebase the PR, this should solve the issue with the tests in the CI