diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 478788c4f1c..cda6ada213c 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -352,11 +352,12 @@ def calibrate_with_adapters(model, args): print_rank_0("Disabling LoRA adapters during calibration...") model.disable_adapters() - yield - - if is_lora: - print_rank_0("Enabling LoRA adapters after calibration...") - model.enable_adapters() + try: + yield + finally: + if is_lora: + print_rank_0("Enabling LoRA adapters after calibration...") + model.enable_adapters() def disable_lora_quantizers_in_config(config, layers): @@ -381,9 +382,11 @@ def replace_function(package, name, new_func, og_func_cache_name=None): old_func = getattr(package, name) setattr(package, name, new_func) setattr(package, og_func_cache_name, old_func) - yield - setattr(package, name, old_func) - delattr(package, og_func_cache_name) + try: + yield + finally: + setattr(package, name, old_func) + delattr(package, og_func_cache_name) @contextmanager @@ -786,10 +789,12 @@ def enable_fake_quant(module): if hasattr(m, "weight_quantizer"): original_fake_quant.append(m.weight_quantizer._fake_quant) m.weight_quantizer._fake_quant = True - yield - for m in module.modules(): - if hasattr(m, "weight_quantizer"): - m.weight_quantizer._fake_quant = original_fake_quant.pop(0) + try: + yield + finally: + for m in module.modules(): + if hasattr(m, "weight_quantizer"): + m.weight_quantizer._fake_quant = original_fake_quant.pop(0) @contextmanager diff --git a/tests/unit/torch/quantization/test_utils.py b/tests/unit/torch/quantization/test_utils.py index 933c553ce12..0782b7a7ec6 100644 --- a/tests/unit/torch/quantization/test_utils.py +++ b/tests/unit/torch/quantization/test_utils.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import types + import pytest import torch @@ -21,6 +23,11 @@ reduce_amax, reduce_block_amax, ) +from modelopt.torch.quantization.utils.core_utils import ( + calibrate_with_adapters, + enable_fake_quant, + replace_function, +) from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -333,3 +340,61 @@ def _forward_loop(m): assert "L1" not in call_log_for_layer1 assert torch.allclose(inp2[0][0][0], torch.tensor([[10.0 + 1.0 + 2.0]])) + + +def test_replace_function_restores_on_exception(): + """replace_function must restore the original function even if the body raises.""" + + def _original(): + return "original" + + def _replacement(): + return "replacement" + + package = types.SimpleNamespace(func=_original) + + with pytest.raises(RuntimeError, match="boom"), replace_function(package, "func", _replacement): + assert package.func is _replacement + assert package._func is _original + raise RuntimeError("boom") + + assert package.func is _original + assert not hasattr(package, "_func") + + +def test_calibrate_with_adapters_reenables_on_exception(): + """calibrate_with_adapters must re-enable LoRA adapters even if calibration raises.""" + + class _Model: + def __init__(self): + self.adapters_enabled = True + + def disable_adapters(self): + self.adapters_enabled = False + + def enable_adapters(self): + self.adapters_enabled = True + + model = _Model() + args = types.SimpleNamespace(lora=True) + + with pytest.raises(RuntimeError, match="boom"), calibrate_with_adapters(model, args): + assert not model.adapters_enabled + raise RuntimeError("boom") + + assert model.adapters_enabled + + +def test_enable_fake_quant_restores_on_exception(): + """enable_fake_quant must restore per-module fake_quant flags even if the body raises.""" + model = torch.nn.Sequential(torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)) + model[0].weight_quantizer = types.SimpleNamespace(_fake_quant=False) + model[1].weight_quantizer = types.SimpleNamespace(_fake_quant=True) + + with pytest.raises(RuntimeError, match="boom"), enable_fake_quant(model): + assert model[0].weight_quantizer._fake_quant + assert model[1].weight_quantizer._fake_quant + raise RuntimeError("boom") + + assert not model[0].weight_quantizer._fake_quant + assert model[1].weight_quantizer._fake_quant