fsdp2_aware_weight_update hides the real error with an UnboundLocalError

[dinhxuanvu]

Before submitting an issue, please make sure it hasn't been already addressed by searching through the existing and past issues.

Describe the bug

fsdp2_aware_weight_update (in modelopt/torch/quantization/utils/core_utils.py) is a context manager. Its finally block uses three variables — fsdp_param_mapping, fsdp_param_group, and root_module — that are only set after the unshard() call inside the try.

So if unshard() fails (typically a CUDA OOM when a large FSDP block can't fit on one GPU), the finally runs before those variables exist and raises:

UnboundLocalError: local variable 'fsdp_param_mapping' referenced before assignment

In Python, an error raised inside finally replaces the original error. So the real failure (the OOM) is thrown away, and you only see the misleading UnboundLocalError. The actual problem becomes very hard to debug.

This shows up during export of large (especially MoE) models, where unsharding a block is most likely to run out of memory.

Steps/Code to reproduce bug

The simplest deterministic repro: make the in-body unshard() raise and confirm the wrong error comes out.

1. Shard a model under FSDP2.
2. Patch the enclosing module's unshard to raise any error.
3. Enter fsdp2_aware_weight_update(model, sharded_module).
4. You get UnboundLocalError instead of the error you raised.

In the real world this happens on its own: export a large MoE model where one block's unshard() all-gather exceeds GPU memory, and the OOM is replaced by the UnboundLocalError.

Expected behavior

If something fails before those variables are set, the finally should be a no-op and the original error (e.g. the OOM) should propagate unchanged.

Who can help?

System information

• OS: Ubuntu 22.04
• CPU architecture: x86_64
• GPU name: H100
• Library versions (if applicable):
  • Python: 3.12
  • ModelOpt version or commit hash: present on current main
  • PyTorch: 2.x