## What does this PR do?
**Type of change:** New feature, new tests
**Overview:** Replace hardcoded per-model MoE class registrations
(Mixtral, Qwen2Moe, Qwen3Moe, Qwen3Next, Llama4TextMoe, Qwen3VLMoe,
MiniMaxM2, etc.) with a single generic auto-detection mechanism
(`register_sparse_moe_on_the_fly`) that walks the model tree and
identifies MoE blocks by their structural attributes (`gate` + `experts`
with `top_k`/`num_experts`). This makes MoE quantization
forward-compatible with new HuggingFace MoE architectures without
requiring explicit registration for each model family.
Additionally, this PR:
- Tracks per-expert token routing counts during calibration via a gate
forward hook, enabling visibility into expert utilization.
- Saves an HTML report of expert token counts during export
(`save_expert_token_count_table`), highlighting under-utilized experts.
- Fixes the `topk` -> `top_k` attribute name for transformers >= 5.0
compatibility.
- Also move the ptq summary prints to a file in hf_ptq.py to reduce the
prints
## Usage
Auto-detection is transparent -- no user-facing API changes are needed.
Any HuggingFace MoE model with the standard `gate`/`experts` pattern is
automatically detected and quantized:
import modelopt.torch.quantization as mtq
# Any HuggingFace MoE model (Mixtral, Qwen3Moe, DeepSeek, etc.)
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-30B-A3B")
mtq.quantize(model, mtq.INT8_DEFAULT_CFG, forward_loop)
# During export, an .moe.html report with per-expert token counts is
saved automatically
## Testing
unittest, also test exporting qwen MOE
## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->
- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->
## Additional Information
<!-- E.g. related issue. -->
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* Added expert token count visualization for Mixture of Experts models,
exported as HTML reports during model export.
* Enhanced sparse MoE quantization with improved calibration-aware
routing and automatic model block detection.
* **Tests**
* Added comprehensive test suite for sparse MoE quantization validation.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
## What does this PR do?
**Type of change:** Bug fix
**Overview:** Fix `AttributeError: Can't get local object
'add_hook_to_module.<locals>.new_forward'` when saving a quantized model
a second time after restoring it with `device_map="auto"`.
When a model is loaded with `device_map="auto"`, accelerate's
`add_hook_to_module` patches every submodule (including
`TensorQuantizer` instances) and injects three instance attributes:
`_hf_hook`, `_old_forward`, and `forward` (a `functools.partial`
wrapping a local function). These are not picklable and were leaking
into the modelopt state dict collected by `get_modelopt_state()`,
causing `torch.save` to fail.
This PR adds the three accelerate-injected attributes to
`TensorQuantizer._skip_properties_for_save_restore` so they are excluded
from the serialized state, matching the existing pattern used for
HuggingFace and DeepSpeed attributes.
## Usage
```python
mto.enable_huggingface_checkpointing()
# Quantize and save
model = AutoModelForCausalLM.from_pretrained(name, device_map="auto")
model = mtq.quantize(model, mtq.FP8_DEFAULT_CFG, forward_loop=forward_loop)
model.save_pretrained(save_dir)
# Restore and save again (this previously failed)
model2 = AutoModelForCausalLM.from_pretrained(save_dir, device_map="auto")
model2.save_pretrained(save_dir_round2) # now works
```
## Testing
- Added unit test
`test_tensor_quantizer_modelopt_state_with_accelerate_hook` in
`tests/unit/torch/quantization/plugins/test_accelerate.py` that verifies
accelerate hook attributes are excluded from modelopt state and the
state dict remains picklable.
## Before your PR is "*Ready for review*"
- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes — only adds entries to a
skip set; existing saved checkpoints are unaffected.
- **Did you write any new necessary tests?**: Yes
- **Did you add or update any necessary documentation?**: No (internal
fix, no API change)
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
No
## Additional Information
The root cause is in accelerate's `add_hook_to_module`, which defines
`new_forward` as a local function and binds it via `functools.partial`
onto `module.forward`. Since local functions cannot be pickled, any
`TensorQuantizer` that has been hooked by accelerate becomes
unserializable unless these attributes are excluded.
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **Bug Fixes**
* Enhanced compatibility with accelerate library by excluding
framework-specific hooks and attributes from model state serialization,
preventing issues during save/restore operations.
* **Tests**
* Added test to validate that accelerate-related attributes are properly
excluded from model state and that the state remains picklable.
* **Public API**
* TensorQuantizer is now publicly exported.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
Signed-off-by: realAsma <akuriparambi@nvidia.com>
## What does this PR do? new feature
**Overview:** use kitchen FA in huggingface plugin
## Usage
<!-- You can potentially add a usage example below. -->
```python
# Add a code snippet demonstrating how to use this
```
## Testing
<!-- Mention how have you tested your change if applicable. -->
## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->
- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->
## Additional Information
<!-- E.g. related issue. -->
---------
Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
## What does this PR do?
**Type of change:** Bug fix
**Overview:**
`BertSDPASelfAttention` quantization was failing
Issue: Decorators (e.g., @deprecate_kwarg in BERT) wrap methods and
change their free variables. The modified AST needs __class__ for
super(), but the decorated wrapper has different freevars (e.g.,
'additional_message', 'func').
Fix: Search the decorator's closure to find the actual undecorated
method with matching free variables, then use its globals and closure
for bytecode reconstruction.
## Usage
```
cd examples/chained_optimizations
bash scripts/1_prune.sh
bash scripts/2_int8_quantize.sh
```
## Testing
Added unit test: test_kv_quant_bert
## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->
- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes
- **Did you write any new necessary tests?**: Yes
- **Did you add or update any necessary documentation?**: Yes
- **Did you update
[Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**:
NA
## Additional Information
<!-- E.g. related issue. -->
Signed-off-by: Kinjal Patel <kinjalpravin@nvidia.com>
## What does this PR do?
**Type of change:**
Bug fix
[aq_Qwen3-235B-A22B-Thinking-2507_scores.html](https://github.com/user-attachments/files/23915284/aq_Qwen3-235B-A22B-Thinking-2507_scores.html)
**Overview:**
For Qwen3 MoE, sensitivity should be measured at `layer.x.mlp` not
`layer.x.mlp.experts` (`layer.x.mlp.experts` module forward is never
called. Hence sensitivity was not correctly estimated previously).
After this PR, Qwen3 MoE sensitivity is correctly estimated. (See
attached file)
## Usage
<!-- You can potentially add a usage example below. -->
```python
# Add a code snippet demonstrating how to use this
```
## Testing
<!-- Mention how have you tested your change if applicable. -->
## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->
- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->
## Additional Information
<!-- E.g. related issue. -->
---------
Signed-off-by: realAsma <akuriparambi@nvidia.com>
## What does this PR do?
**Type of change:** Refator; Minor new feature
**Overview:** ?
1. Refactored AutoQuantizeSearcher to _AutoQuantizeBaseSearcher &
AutoQuantizeGradientSearcher - Prepares architecture for additional
search methods.
2. seperated quant modules and score modules - separate quantization
modules from scoring modules, enabling auto-quantization to measure
sensitivity at parent layers (e.g., MLP output for MoE experts) rather
than individual ops.
3. Also see https://github.com/NVIDIA/TensorRT-Model-Optimizer/pull/592
and https://github.com/NVIDIA/TensorRT-Model-Optimizer/pull/588
## Testing
See unittests; `tests/unit/torch/quantization/test_autoquant.py` and
`tests/unit/torch/quantization/plugins/test_huggingface.py`
## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->
- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes
- **Did you write any new necessary tests?**: Yes
- **Did you add or update any necessary documentation?**: Yes
- **Did you update
[Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Not Required
## Additional Information
<!-- E.g. related issue. -->
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* Added support for score modules in quantization workflows.
* Added optional naming for quantization recipes.
* **Bug Fixes**
* Improved quantization grouping rules documentation with clearer
configuration examples.
* **Refactor**
* Renamed quantization module parameters for improved clarity.
* Enhanced quantization search architecture for better scalability.
<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: realAsma <akuriparambi@nvidia.com>
Co-authored-by: Asma Kuriparambil Thekkumpate <akuriparambi@akuriparambi-mlt.client.nvidia.com>
## What does this PR do?
**Type of change:** New feature <!-- Use one of the following: Bug fix,
new feature, new example, new tests, documentation. -->
**Overview:** Support quantization of PyTorch Geometric
```python
# Add a code snippet demonstrating how to use this
```
## Testing
`python -m pytest
tests/unit/torch/quantization/plugins/test_pytorch_geometric_plugin.py
-v`
## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->
- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes <!--- If No, explain why.
-->
- **Did you write any new necessary tests?**: Yes
- **Did you add or update any necessary documentation?**: Yes
- **Did you update
[Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes<!--- Only for new features, API changes, critical bug fixes or bw
breaking changes. -->
## Additional Information
- https://github.com/NVIDIA/TensorRT-Model-Optimizer/issues/298
---------
Signed-off-by: Riyad Islam <rislam@nvidia.com>
Chenjie Luo