[kevalmorabia97]

What does this PR do?

Type of change: Bug fix (tests)

Fixes the two unit tests that fail on main since #3091 merged:
test_fsdp2_owns_its_parameters_in_every_shard_state and test_fsdp2_weight_access_leaves_parameters_fsdp2_does_not_own in tests/unit/torch/quantization/test_utils.py.

#3088 added these tests, which count fully_shard.state calls as a stand-in for "how many times the window scanned the FSDP2 unit". #3091 then switched _get_fsdp2_mesh from the private _get_module_state to the public fully_shard.state, and the weight-access window calls _get_fsdp2_mesh on every entry, so the counts went up by one per entry. The two PRs did not conflict textually. The tests now count _fsdp2_param_groups, the scan itself.

Usage

N/A

Testing

tests/unit/torch/quantization/test_utils.py on torch 2.10: 27 passed (both tests fail on main before this change, as in its linux CI run).

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?: N/A
• Did you update Changelog?: N/A
• Did you get Claude approval on this PR?: N/A

Additional Information

Unblocks linux on #3092 and #3093.

🤖 Generated with Claude Code

[coderabbitai]

Caution

Review failed

The pull request is closed.

Finishing Touches

📝 Generate docstrings

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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

[kevalmorabia97]

Closing: #2745 already landed the same fix on main (the tests now count _fsdp2_param_groups).