Invalidate sharded sparsity masks after updating the source mask

CHANGELOG.rst

@@ -25,6 +25,7 @@ Changelog
 
 **Bug Fixes**
 
+- Fix updated sparsity masks being ignored when reading or exporting sharded FSDP2 weights after calling ``set_mask()``.
 - Fix ONNX INT8 entropy calibration failing or producing invalid quantization parameters for FP16 activations.
 - Fix ``--use_fsdp2`` HuggingFace checkpoint export gathering the whole model onto rank 0, which made export the dominant phase of a PTQ run and could exhaust host memory on large models. The model is now split into per-decoder-layer units dealt round-robin across ranks; each rank gathers every unit but keeps, packs, and writes only the ones it owns, so a rank buffers roughly ``model / world_size`` instead of the whole checkpoint, and rank 0 writes the combined index. Export configurations that cannot be split this way now raise instead of producing a mismatched checkpoint: FSDP2 combined with another DTensor parallelism (for example FSDP2 + tensor parallel on a 2-D mesh; HSDP is supported), models whose decoder layers cannot be discovered, a decoder layer object reused across layers, and a module that holds the decoder layers while owning parameters of its own.
 - Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration.

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))