Fix selected layers being skipped by post-init sparsity

[darkdi]

What does this PR do?

Type of change: Bug fix

do_sparsity() silently skips selected layers such as dense and gate: rstrip(".weight") treats its argument as a set of characters, turning dense.weight into dens and gate.weight into ga. Those names no longer match the selected modules. Remove the literal suffix instead.

Usage

No API changes. The regression test builds a small model and applies SparsityMethod2o4().do_sparsity(model) to each selected layer.

Testing

On main, the regression test fails for dense and gate because no pruning mask is attached; the proj control passes. With the fix, all three cases pass and verify the mask and resulting weights. All 13 tests in test_post_init_sparse.py and test_common.py pass on CPU. Pre-commit checks pass for both changed files.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: ✅
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: N/A
• Did you get Claude approval on this PR?: N/A

Additional Information

Found while inspecting the layer-name filtering in the post-initialization sparsity helper.

Summary by CodeRabbit

• Bug Fixes

  • Fixed sparsity processing for layers whose names include .weight, ensuring the correct model parameters are matched and updated.

• Tests

  • Added coverage for structured 2:4 sparsity across dense, gate, and proj MLP layers, including mask placement and resulting weight values.

[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: d0629883-72ae-4afa-bd0c-654ec6d055f0

📥 Commits

Reviewing files that changed from the base of the PR and between a7f339e and 9586109.

📒 Files selected for processing (2)

• modelopt/torch/puzzletron/tools/post_init_sparse.py
• tests/unit/torch/puzzletron/test_post_init_sparse.py

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

---

📝 Walkthrough

Walkthrough

The change fixes state-dict key normalization for sparsified layers. New parametrized tests verify 2:4 masks and masked weights for dense, gate, and proj MLP layers.

Changes

Post-initialization sparsity

Layer / File(s)	Summary
Exact weight-key normalization and sparsity validation modelopt/torch/puzzletron/tools/post_init_sparse.py, tests/unit/torch/puzzletron/test_post_init_sparse.py	The sparsity method removes the exact .weight suffix. Tests verify expected 2:4 masks and masked weights for supported MLP layer names.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to 95861

Post-initialization sparsity now preserves selected layer names when removing the .weight suffix, and regression coverage verifies correct 2:4 pruning for the affected MLP layer names. No current merge-blocking risk remains.

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 2 files.	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: fixing skipped layers during post-initialization sparsity.
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	PASS. The PR changes only two suffix-matching expressions in modelopt/torch/puzzletron/tools/post_init_sparse.py and adds a test file. The added and modified Python lines contain no `torch.load(...,…

• 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.}