mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-11-10 09:34:01 +08:00
jingyu-ml
87237e7dd1
## 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:**
This MR improves robustness when `forward()` is monkey‑patched (replaced
at runtime) on modules that later get wrapped/converted by ModelOpt
(DynamicModule + quant wrappers).
It addresses two concrete failure modes introduced/exposed by supporting
“patched forward” modules:
1. Forward “leakage” after export: a dynamic wrapper forward could
remain bound on an instance even after export() restores the original
(non‑dynamic) class, causing runtime errors in unrelated codepaths (e.g.
KD export/save/restore chains).
1. Infinite recursion in quant wrappers: _forward_pre_dm can sometimes
point to a wrapper forward that already participates in the class chain,
causing a recursion loop when quant wrappers call _forward_pre_dm
directly.
## Usage
<!-- You can potentially add a usage example below. -->
```python
lin = torch.nn.Linear(4, 4)
def upcast_forward(x):
# external closure: NOT part of any class MRO
return torch.nn.functional.linear(x, lin.weight.to(x.dtype), lin.bias.to(x.dtype))
lin.forward = upcast_forward # framework/user patches forward
# Later, ModelOpt converts/wraps the module.
# It stashes the patched function as `_forward_pre_dm` and binds the wrapper forward on the class.
# During quantization, QuantInputBase.forward sees `_forward_pre_dm` is NOT in MRO -> calls it.
```
```
# Imagine a module already wrapped by quant classes:
# QuantLinearConvBase.forward -> super().forward -> QuantInputBase.forward -> ...
# If `_forward_pre_dm` accidentally points to QuantLinearConvBase.forward (which IS in MRO),
# and QuantInputBase.forward calls it directly, you get:
# QuantInputBase.forward -> _forward_pre_dm (QuantLinearConvBase.forward)
# -> super().forward -> QuantInputBase.forward -> ...
# infinite recursion
# The fix: if `_forward_pre_dm` is a forward already in MRO, ignore it and use super().forward.
```
## 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?**: No <!--- If No, explain why.
-->
- **Did you write any new necessary tests?**: Yes
- **Did you add or update any necessary documentation?**: No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**: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
## Release Notes
* **Bug Fixes**
* Improved forward method restoration during module export to prevent
state leakage
* Enhanced quantization behavior when using chained optimization modes
* **Tests**
* Added regression tests for quantization with runtime forward patching
* Added validation tests for sparse quantization combined with
distillation workflows
<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: Jingyu Xin <jingyux@nvidia.com>
94 lines
3.6 KiB
Python
94 lines
3.6 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import types
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch import nn
|
|
|
|
import modelopt.torch.quantization as mtq
|
|
from modelopt.torch.quantization import QuantModuleRegistry
|
|
from modelopt.torch.quantization.nn.modules.quant_module import QuantLinearConvBase
|
|
|
|
|
|
def test_quant_input_base_ignores_forward_pre_dm_in_mro():
|
|
"""Regression test for recursion when `_forward_pre_dm` points to a wrapper forward in the MRO.
|
|
|
|
In complex wrapper stacks, `_forward_pre_dm` may accidentally end up referencing a `forward`
|
|
method already present in the quant wrapper MRO (e.g. QuantLinearConvBase.forward). If
|
|
QuantInputBase.forward calls that directly, it can recurse indefinitely:
|
|
|
|
QuantLinearConvBase.forward -> super().forward (QuantInputBase.forward)
|
|
-> _forward_pre_dm (QuantLinearConvBase.forward) -> ...
|
|
|
|
The fix is to detect this case and fall back to `super().forward` instead.
|
|
"""
|
|
lin = nn.Linear(8, 8, bias=False)
|
|
QuantModuleRegistry.convert(lin)
|
|
|
|
# Force the problematic state: `_forward_pre_dm` points to a wrapper forward already in MRO.
|
|
lin._forward_pre_dm = types.MethodType(QuantLinearConvBase.forward, lin)
|
|
|
|
x = torch.randn(2, 8)
|
|
y = lin(x)
|
|
assert isinstance(y, torch.Tensor)
|
|
assert y.shape == (2, 8)
|
|
|
|
|
|
def test_quantize_calibration_calls_quantizers_with_runtime_forward_patch():
|
|
"""Regression test for on-the-fly forward patching during mtq.quantize calibration.
|
|
|
|
Some frameworks replace `module.forward` on-the-fly with a closure just before a forward pass.
|
|
During mtq.quantize calibration, quantizers must still run (input + weight at minimum).
|
|
"""
|
|
lin = nn.Linear(8, 8, bias=True).to(torch.float32)
|
|
|
|
called = {"patched_forward": 0, "input_q": 0, "weight_q": 0}
|
|
|
|
# Monkey patch instance-level forward (closure-style, no `self` argument).
|
|
def patched_forward(x):
|
|
called["patched_forward"] += 1
|
|
# Use module parameters directly; if quantization wrappers are active, weight access
|
|
# should still be routed through the quantized path.
|
|
w = lin.weight.to(dtype=x.dtype)
|
|
b = lin.bias.to(dtype=x.dtype) if lin.bias is not None else None
|
|
return F.linear(x, w, b)
|
|
|
|
def _count_input(_m, _inp, _out):
|
|
called["input_q"] += 1
|
|
|
|
def _count_weight(_m, _inp, _out):
|
|
called["weight_q"] += 1
|
|
|
|
lin.forward = patched_forward
|
|
x = torch.randn(2, 8, dtype=torch.float16)
|
|
|
|
def forward_loop(model):
|
|
# Patch forward on-the-fly (after conversion, right before calibration forward).
|
|
|
|
# Count quantizer executions during calibration.
|
|
model.input_quantizer.register_forward_hook(_count_input)
|
|
model.weight_quantizer.register_forward_hook(_count_weight)
|
|
|
|
model(x)
|
|
|
|
mtq.quantize(lin, mtq.INT8_DEFAULT_CFG, forward_loop)
|
|
lin(x)
|
|
|
|
assert called["patched_forward"] == 2
|
|
assert called["input_q"] == 2
|
|
assert called["weight_q"] == 2
|