diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 01a0fbe1b47..d633433693a 100755 --- a/CHANGELOG.rst +++ b/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. diff --git a/modelopt/torch/sparsity/weight_sparsity/module.py b/modelopt/torch/sparsity/weight_sparsity/module.py index 533976fa9ca..9d9d0330bfe 100644 --- a/modelopt/torch/sparsity/weight_sparsity/module.py +++ b/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 diff --git a/tests/gpu/torch/sparsity/weight_sparsity/test_sparse_fsdp.py b/tests/gpu/torch/sparsity/weight_sparsity/test_sparse_fsdp.py index a087b59a20d..53e911e7f47 100644 --- a/tests/gpu/torch/sparsity/weight_sparsity/test_sparse_fsdp.py +++ b/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))