Invalidate sharded sparsity masks after updating the source mask

[MrCapricornLiu]

What does this PR do?

Type of change: Bug fix

Calling set_mask() after FSDP2 sharding can leave the previous sparsity mask active when reading or exporting the sharded weight. The method invalidates its DTensor cache first, but then reads self.weight.shape for validation. That dynamic weight access rebuilds the cache from the old mask before the mask buffer is updated.

Invalidate the cache after updating the source mask instead, including the None path. Subsequent weight reads and sparse export then use the newly requested mask.

Usage

No API change. SparseModule.set_mask() updates the mask used by sharded weights and their exported state.

Testing

• Four new FSDP2 cases fail on the original implementation with incorrect masked weights.
• Six tests in tests/gpu/torch/sparsity/weight_sparsity/test_sparse_fsdp.py pass on two H800 GPUs, including the existing FSDP tests and new FP32/BF16 cases with and without a prior sparse mask.
• The new cases replace masks repeatedly, remove the mask, use an all-ones mask, and compare the exported state dict with the raw weights multiplied by the final requested mask.
• All 109 CPU tests in tests/unit/torch/sparsity/weight_sparsity/test_sparsify.py pass.
• Changed-file pre-commit checks pass, including Ruff, mypy, licenses and Bandit.

Multi-node execution, CPU offload and performance were not tested. The change addresses set_mask(); direct in-place mutation or checkpoint loading into an already populated cache is outside this regression.

Before your PR is "Ready for review"

• Backward compatible: yes; no configuration or checkpoint schema changes.
• Copied code or new PIP dependencies: none.
• Necessary tests: added to the existing FSDP sparsity test module.
• Changelog: updated.
• Upstream review: pending.

Additional Information

The DTensor mask cache was introduced in #1818. This corrects its invalidation ordering without changing the sharding strategy.

Summary by CodeRabbit

• Bug Fixes

  • Updated sparsity masks are now correctly respected when reading or exporting sharded FSDP2 weights after mask changes.
  • Removing or replacing a sparsity mask now reliably refreshes the affected weight data.

• Tests

  • Added coverage for dynamic mask updates across data types and initial-mask configurations, including sharded and exported weights.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: aa2a8271-de2c-4e1b-81af-fbf5f247a39e

📥 Commits

Reviewing files that changed from the base of the PR and between 079078d and 4a07aee.

📒 Files selected for processing (3)

• CHANGELOG.rst
• modelopt/torch/sparsity/weight_sparsity/module.py
• tests/gpu/torch/sparsity/weight_sparsity/test_sparse_fsdp.py

Included review availability: Your plan provides up to 12 included reviews per hour; 9 remain after this review.

---

📝 Walkthrough

Walkthrough

The change fixes DTensor mask cache handling in SparseModule.set_mask() and adds FSDP2 tests for dynamic masks across dtypes and initial-mask states. The changelog records the fix.

Changes

Sparsity mask cache updates

Layer / File(s)	Summary
Mask update and FSDP2 validation modelopt/torch/sparsity/weight_sparsity/module.py, tests/gpu/torch/sparsity/weight_sparsity/test_sparse_fsdp.py, CHANGELOG.rst	set_mask() clears the cached DTensor mask when removing a mask and invalidates it after non-null mask updates. Parameterized FSDP2 tests validate sharded weights and exports for float32 and bfloat16 models with and without initial masks. The changelog records the fix.

Estimated code review effort: 1 (Trivial) | ~5 minutes

Suggested reviewers: kaix-nv, kevalmorabia97

Merge Risk: ⚪ Minimal · up to 4a07a

This fixes stale sparsity masks for FSDP2 reads and exports without changing the public API or checkpoint format. The targeted regression coverage passes, so the change is ready to merge.

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 20.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 5 functions across 2 files. (1 skipped: 1…	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: invalidating cached sharded sparsity masks after the source mask changes.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	No security anti-pattern from the custom check was introduced. The authoritative diff changes one modelopt Python module and one GPU test; it adds no torch.load(..., weights_only=False), `numpy.lo…

Full details: Docstring Coverage

Explanation

Docstring coverage is 20.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 5 functions across 2 files. (1 skipped: 1 unsupported.)

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}