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..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,6 +64,47 @@ def test_reduce_block_amax(block_sizes, test_input, expected_scales): torch.allclose(scales, expected_scales) +@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. + """ + root_model = mock.MagicMock(spec=FSDPModule) + module = torch.nn.Linear(4, 4) + + def _raise_oom(*args, **kwargs): + raise RuntimeError("CUDA out of memory (simulated)") + + 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.""" + module = torch.nn.Linear(4, 4) + with pytest.raises(ValueError, match="body failure"): + with core_utils.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):