Invalidate sharded sparsity masks after updating the source mask

CHANGELOG.rst

@@ -38,6 +38,7 @@ Changelog
 
 **Bug Fixes**
 
+- Fix updated sparsity masks being ignored when reading or exporting sharded FSDP2 weights after calling ``set_mask()``.
 - Fix ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` storing the MoE router at Megatron's ``moe_router_dtype``, which is a routing *compute* dtype, not a storage one. The router now exports at the export ``dtype`` like every other unquantized weight, matching what ``hf_ptq.py`` and the released NVFP4 checkpoints contain; pass ``moe_router_dtype`` to ``export_mcore_gpt_to_hf`` explicitly if you want the old fp32 storage.
 - Fix unified Megatron export writing a second, unreferenced copy of the vocab embedding when a model with MTP layers is exported with pipeline parallelism. The duplicate was never loaded but inflated the checkpoint by the size of the embedding (about 1 GB for Qwen3.6-35B-A3B); re-export to reclaim the space.
 - Fix ONNX INT8 entropy calibration failing or producing invalid quantization parameters for FP16 activations.

modelopt/torch/sparsity/weight_sparsity/module.py

@@ -87,11 +87,9 @@ def modify(self, *args, **kwargs):
 
     def set_mask(self, value: torch.BoolTensor | None):
         """Set the active sparse mask of the module weights."""
-        # invalidate the cached DTensor mask since the underlying mask is changing
-        self._weight_mask_dtensor = None
-
         if value is None:
             self._weight_mask = None
+            self._weight_mask_dtensor = None
             return
 
         # sanity checks on the mask
@@ -108,3 +106,7 @@ def set_mask(self, value: torch.BoolTensor | None):
                 self._weight_mask = value.detach().clone().to(self.weight.device)
             else:
                 self._weight_mask.copy_(value.to(self._weight_mask.device))
+
+        # Reading self.weight above can populate the cache with the old mask.
+        # Invalidate it only after the underlying mask has been updated.
+        self._weight_mask_dtensor = None

tests/gpu/torch/sparsity/weight_sparsity/test_sparse_fsdp.py

@@ -19,9 +19,11 @@
 import torch
 import torch.nn as nn
 from _test_utils.torch.distributed.fsdp_test import run_fsdp_test
+from torch.distributed.fsdp import fully_shard
 
 from modelopt.torch.nas.search_space import SearchSpace
 from modelopt.torch.opt.conversion import apply_mode
+from modelopt.torch.sparsity import export
 
 
 def _get_test_case():
@@ -48,3 +50,33 @@ def test_fsdp(dist_workers, use_orig_params):
             fsdp_kwargs={"use_orig_params": use_orig_params},
         ),
     )
+
+
+def _run_fsdp2_mask_updates(dtype, initial_mask, rank, world_size):
+    model, _ = _get_test_case()
+    model.to(dtype=dtype)
+    raw_weight = model[0]._parameters["weight"].detach().clone()
+    mask = torch.ones_like(raw_weight, dtype=torch.bool)
+    mask[:, ::2] = False
+    if initial_mask:
+        model[0].set_mask(mask)
+    model = fully_shard(model)
+
+    for new_mask in [~mask, mask, None, torch.ones_like(mask), ~mask]:
+        model[0].set_mask(new_mask)
+        expected = raw_weight if new_mask is None else raw_weight * new_mask
+        # Reading a sharded dynamic weight must reflect the most recent mask.
+        torch.testing.assert_close(model[0].weight.full_tensor(), expected, atol=0, rtol=0)
+
+    # Export materializes the dynamic weight, so a stale cache would bake the
+    # previous mask into the checkpoint even though the mask buffer was updated.
+    exported = export(model)
+    torch.testing.assert_close(
+        exported.state_dict()["0.weight"].full_tensor(), expected, atol=0, rtol=0
+    )
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("initial_mask", [False, True])
+def test_fsdp2_mask_updates(dist_workers, dtype, initial_mask):
+    dist_workers.run(partial(_run_fsdp2_mask_updates, dtype, initial_mask))