From 5becc6405d20737e3ad832b4180047b01598624a Mon Sep 17 00:00:00 2001 From: harshal-96 Date: Tue, 1 Sep 2026 06:27:11 +0530 Subject: [PATCH 1/2] Fix fsdp2_aware_weight_update masking setup errors with UnboundLocalError fsdp2_aware_weight_update assigns root_module, fsdp_param_group and fsdp_param_mapping only after unshard() succeeds inside the try block, but the finally block referenced them unconditionally. When setup fails (typically a CUDA OOM in unshard() while exporting a large MoE model), the finally block raised UnboundLocalError, which replaces the original exception and hides the real failure. Initialize fsdp_param_mapping as a sentinel before the try and skip the finally-block update when setup never completed, so the original error propagates unchanged. Adds a CPU regression test that simulates a setup failure and asserts the original error type surfaces, plus a passthrough sanity test for non-FSDP roots. Fixes #1859 Signed-off-by: harshal-96 --- .../torch/quantization/utils/core_utils.py | 9 ++++- tests/unit/torch/quantization/test_utils.py | 37 +++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 3d1e017b3a5..ce7cc06eb3f 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -953,6 +953,10 @@ def fsdp2_aware_weight_update(root_model, modules_to_update, reshard=True): Returns: None """ + # Assigned only once setup below completes; the finally block keys off + # this so a setup failure (e.g. a CUDA OOM raised by unshard()) propagates + # instead of being replaced by an UnboundLocalError. See issue #1859. + fsdp_param_mapping = None try: if isinstance(root_model, FSDPModule): # Get FSDP root module, if none is returned, then the update is not made to a submodule of an FSDPModule @@ -995,7 +999,10 @@ def fsdp2_aware_weight_update(root_model, modules_to_update, reshard=True): finally: from modelopt.torch.quantization.qtensor.base_qtensor import QFSDPParam, QTensorWrapper - if isinstance(root_model, FSDPModule): + # fsdp_param_mapping is None when setup did not complete (nothing was + # unsharded or mapped) — skip the update so the original error from + # the try block is not masked by an UnboundLocalError here. + if isinstance(root_model, FSDPModule) and fsdp_param_mapping is not None: # Update FSDPParam list for module in modules_to_update: for param_name, param in module.named_parameters(): diff --git a/tests/unit/torch/quantization/test_utils.py b/tests/unit/torch/quantization/test_utils.py index 933c553ce12..7fdd51b667e 100644 --- a/tests/unit/torch/quantization/test_utils.py +++ b/tests/unit/torch/quantization/test_utils.py @@ -59,6 +59,43 @@ def test_reduce_block_amax(block_sizes, test_input, expected_scales): torch.allclose(scales, expected_scales) +def test_fsdp2_aware_weight_update_preserves_setup_error(monkeypatch): + """A failure during setup (e.g. a CUDA OOM raised by ``unshard()``) must + propagate to the caller, not be replaced by an ``UnboundLocalError`` from + the ``finally`` block referencing variables that were never assigned. + + Regression test for https://github.com/NVIDIA/Model-Optimizer/issues/1859. + """ + from unittest import mock + + from torch.distributed.fsdp import FSDPModule + + from modelopt.torch.quantization.utils import core_utils + + root_model = mock.MagicMock(spec=FSDPModule) + module = torch.nn.Linear(4, 4) + + def _raise_oom(*args, **kwargs): + raise RuntimeError("CUDA out of memory (simulated)") + + monkeypatch.setattr(core_utils, "_get_enclosing_fsdp_module", _raise_oom) + + with pytest.raises(RuntimeError, match="out of memory"): + with core_utils.fsdp2_aware_weight_update(root_model, module): + pass # setup fails before the body runs + + +def test_fsdp2_aware_weight_update_non_fsdp_body_error_passthrough(): + """For a non-FSDP root model the context manager is a no-op wrapper and + must transparently propagate errors raised in the body.""" + from modelopt.torch.quantization.utils.core_utils import fsdp2_aware_weight_update + + module = torch.nn.Linear(4, 4) + with pytest.raises(ValueError, match="body failure"): + with fsdp2_aware_weight_update(module, module): + raise ValueError("body failure") + + @pytest.mark.parametrize("fp8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) @pytest.mark.parametrize("axis", [None, 0, 1, (0, 1)]) def test_reduce_amax_fp8(fp8_dtype, axis): From 059460bf3317ff18c32437dcf6558b3ae85e012b Mon Sep 17 00:00:00 2001 From: harshal-96 Date: Tue, 1 Sep 2026 11:57:18 +0530 Subject: [PATCH 2/2] Address review: exercise unshard() failure path, module-scope imports - Parametrize the regression test over both setup-failure points: module discovery and the reported scenario where discovery succeeds and unshard() itself raises (verified to reproduce the masking before the fix and pass after). - Move test imports to module scope per test guidelines. Signed-off-by: harshal-96 --- tests/unit/torch/quantization/test_utils.py | 31 +++++++++++++-------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/tests/unit/torch/quantization/test_utils.py b/tests/unit/torch/quantization/test_utils.py index 7fdd51b667e..e593c52da8b 100644 --- a/tests/unit/torch/quantization/test_utils.py +++ b/tests/unit/torch/quantization/test_utils.py @@ -13,14 +13,19 @@ # See the License for the specific language governing permissions and # limitations under the License. +from contextlib import nullcontext +from unittest import mock + import pytest import torch +from torch.distributed.fsdp import FSDPModule from modelopt.torch.quantization.utils import ( convert_quantization_axis_to_reduce_axis, reduce_amax, reduce_block_amax, ) +from modelopt.torch.quantization.utils import core_utils from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -59,40 +64,44 @@ def test_reduce_block_amax(block_sizes, test_input, expected_scales): torch.allclose(scales, expected_scales) -def test_fsdp2_aware_weight_update_preserves_setup_error(monkeypatch): +@pytest.mark.parametrize("failure_point", ["module_discovery", "unshard"]) +def test_fsdp2_aware_weight_update_preserves_setup_error(monkeypatch, failure_point): """A failure during setup (e.g. a CUDA OOM raised by ``unshard()``) must propagate to the caller, not be replaced by an ``UnboundLocalError`` from the ``finally`` block referencing variables that were never assigned. Regression test for https://github.com/NVIDIA/Model-Optimizer/issues/1859. """ - from unittest import mock - - from torch.distributed.fsdp import FSDPModule - - from modelopt.torch.quantization.utils import core_utils - root_model = mock.MagicMock(spec=FSDPModule) module = torch.nn.Linear(4, 4) def _raise_oom(*args, **kwargs): raise RuntimeError("CUDA out of memory (simulated)") - monkeypatch.setattr(core_utils, "_get_enclosing_fsdp_module", _raise_oom) + if failure_point == "module_discovery": + monkeypatch.setattr(core_utils, "_get_enclosing_fsdp_module", _raise_oom) + else: # the reported scenario: discovery succeeds, unshard() OOMs + monkeypatch.setattr(core_utils, "_get_enclosing_fsdp_module", lambda m, r: root_model) + fake_fully_shard = mock.MagicMock() + fake_fully_shard.state.return_value._fsdp_param_group.is_sharded = True + monkeypatch.setattr(core_utils, "fully_shard", fake_fully_shard) + monkeypatch.setattr(core_utils, "enable_fake_quant", lambda m: nullcontext()) + root_model.unshard.side_effect = _raise_oom with pytest.raises(RuntimeError, match="out of memory"): with core_utils.fsdp2_aware_weight_update(root_model, module): pass # setup fails before the body runs + if failure_point == "unshard": + root_model.unshard.assert_called_once() + def test_fsdp2_aware_weight_update_non_fsdp_body_error_passthrough(): """For a non-FSDP root model the context manager is a no-op wrapper and must transparently propagate errors raised in the body.""" - from modelopt.torch.quantization.utils.core_utils import fsdp2_aware_weight_update - module = torch.nn.Linear(4, 4) with pytest.raises(ValueError, match="body failure"): - with fsdp2_aware_weight_update(module, module): + with core_utils.fsdp2_aware_weight_update(module, module): raise ValueError("body failure")