fix: make replace_function / calibrate_with_adapters / enable_fake_quant exception-safe

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

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