Fix broken remove_sharded_tensors public API and re-enable its unit test

[huthvincent]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Repairs megatron.core.dist_checkpointing.remove_sharded_tensors, which is currently broken in three independent ways, and re-enables the unit test that would have caught this.

1. serialization.py: the wrapper called the method unbound on TorchDistSaveShardedStrategy, but it is defined only on TorchDistLoadShardedStrategy → every call raised AttributeError. Now instantiates the load strategy, restoring the pre-2beb593b4 behavior (sharded_strategy.remove_sharded_tensors(...) on an instance).
2. strategies/torch.py: the method read fs_writer.metadata_path, which no longer exists on FileSystemWriter in current PyTorch. It now uses the metadata_filename computed two lines above (identical path, works across torch versions).
3. strategies/torch.py: fs_writer.fs.rm_file(checkpoint_dir / f) raised TypeError for the documented str argument type; now Path(checkpoint_dir) / f.

Test changes (tests/unit_tests/dist_checkpointing/test_serialization.py::test_remove_sharded_tensors):

• Removed the @pytest.mark.flaky / @pytest.mark.flaky_in_dev markers. With both markers the test runs in no CI environment, which is precisely how the API could break unnoticed (timeline in the linked issue).
• Updated the post-removal metadata assertion to include common_state/shard_0_1: commit 40b1fd3 ("deprecate common strategy") updated the pre-removal assertion for the new common-state entry but could not update the post-removal one, since the test no longer reached it.

Issue tracking

Linked issue: Fixes #5758

Contribution process

Pre-checks

• I have added relevant unit tests (re-enabled and corrected the existing one)
• I have added relevant functional tests (not applicable)
• I have added proper typing to my code (no signature changes)
• I have added relevant documentation (behavior unchanged; restores documented behavior)
• I have run the autoformatter on the changed files (black 24.10, isort)

Local verification (the CI unit test needs 8 GPUs; verified end-to-end at world_size=1 on a single GPU with the same save/remove/inspect sequence as the unit test, passing str paths to exercise fix 3):

files before: ['.metadata', '__0_0.distcp', 'metadata.json', 'prefix__0_1.distcp']
files after: ['.metadata', '__0_0.distcp', 'metadata.json']
metadata keys after: ['common_state/shard_0_1', 'keyA']
PASS: remove_sharded_tensors works end-to-end

On unpatched main the same script fails with AttributeError: type object 'TorchDistSaveShardedStrategy' has no attribute 'remove_sharded_tensors'; with only fix 1 applied it fails with AttributeError: 'FileSystemWriter' object has no attribute 'metadata_path'.

[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.

[cspades]

/ok to test 5218f35

[huthvincent]

@cspades Same situation as #5756. After your /ok to test, I pushed a small follow-up commit (bb92537c5) that bumps the NVIDIA copyright headers on the changed files to the current year so the copyright-check job passes. That moved the PR head, so the earlier approval is tied to the previous commit and CI did not start on the new one. Could you re-run /ok to test bb92537c5 when you get a chance? Thanks a lot, and sorry for the extra round.

[cspades]

/ok to test bb92537

[dimapihtar]

/ok to test c656b65

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29677835075