remove deprecated get_default_load_sharded_strategy

[dimapihtar]

What does this PR do?

uses TorchDistLoadShardedStrategy instead of deprecated get_default_load_sharded_strategy.

Usage

# Add a code snippet demonstrating how to use this

Testing

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?: ✅ / ❌ / N/A
• 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

Summary by CodeRabbit

Release Notes

• Bug Fixes
  • Improved distributed checkpoint loading strategy for enhanced reliability and consistency when loading sharded checkpoint extra-state data.

[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: 0af2e769-9878-4fe5-b456-b02eb807f903

📥 Commits

Reviewing files that changed from the base of the PR and between 6b73e93 and 688be40.

📒 Files selected for processing (1)

• modelopt/torch/opt/plugins/mcore_dist_checkpointing.py

---

📝 Walkthrough

Walkthrough

This PR updates the sharded checkpoint loading strategy in the Megatron Core plugin. It imports TorchDistLoadShardedStrategy and replaces the factory-based get_default_load_sharded_strategy(checkpoint_name) call with explicit TorchDistLoadShardedStrategy() instantiation when loading extra-state from sharded checkpoints.

Changes

Sharded checkpoint loading strategy

Layer / File(s)	Summary
Sharded checkpoint loading strategy update modelopt/torch/opt/plugins/mcore_dist_checkpointing.py	Adds TorchDistLoadShardedStrategy import from Megatron Core and replaces the checkpoint extra-state loading strategy from the factory function get_default_load_sharded_strategy(checkpoint_name) to the explicit class constructor TorchDistLoadShardedStrategy().

---

🎯 1 (Trivial) | ⏱️ ~3 minutes

🚥 Pre-merge checks | ✅ 6

✅ Passed checks (6 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and specifically describes the main change: removing a deprecated function (get_default_load_sharded_strategy) and replacing it with TorchDistLoadShardedStrategy, which aligns with the actual code modifications.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
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-patterns detected in the PR. File refactors deprecated function call; no unsafe deserialization, eval/exec, or hardcoded trust settings present.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[codecov]

Codecov Report

❌ Patch coverage is 0% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 77.38%. Comparing base (6b73e93) to head (688be40).

Files with missing lines	Patch %	Lines
...lopt/torch/opt/plugins/mcore_dist_checkpointing.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1629   +/-   ##
=======================================
  Coverage   77.38%   77.38%           
=======================================
  Files         482      482           
  Lines       52960    52960           
=======================================
  Hits        40984    40984           
  Misses      11976    11976           

Flag	Coverage Δ
unit	53.93% <0.00%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.