[ckpt] feat: support MSC for fsdp_dtensors

[pavelgein]

What does this PR do ?

Support checkpoints for fsdp_dtensors

Changelog

• Add specific line by line info of high level changes in this PR.

GitHub Actions CI

See the CI sectionin the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (issue)
  [feature] Support MultiStorageClient (MSC) for FSDP checkpoints #3261

Summary by CodeRabbit

• New Features

  • Added support for flexible checkpoint storage backend selection via feature flag.

• Refactor

  • Centralized checkpoint reader and writer instantiation for improved consistency in checkpoint handling.

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

📝 Walkthrough

Walkthrough

Updated checkpoint I/O operations in fsdp_dtensor to use helper functions and conditional storage writer selection based on MultiStorageClientFeature.is_enabled(). Both checkpoint saving and loading paths now delegate to centralized reader/writer creation helpers instead of directly instantiating PyTorch checkpoint classes.

Changes

Cohort / File(s)

Summary

Checkpoint I/O Abstraction src/megatron/bridge/training/checkpointing.py

Modified fsdp_dtensor checkpoint saving to conditionally select MultiStorageFileSystemWriter or torch.distributed.checkpoint.FileSystemWriter based on feature flag. Updated checkpoint loading in fsdp_dtensor and _load_fsdp_dtensor_base_checkpoint to use centralized reader helpers _get_checkpoint_reader() and _get_filesystem_reader() instead of direct class instantiation. Adjusted imports accordingly.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Test Results For Major Changes	⚠️ Warning	PR implements major MSC feature for fsdp_dtensors but lacks documented test results, and contains critical bug: line 1917 calls undefined function _get_checkpoint_reader instead of _get_filesystem_reader.	Provide documented test results demonstrating MSC checkpoint functionality works; fix line 1917 bug by replacing _get_checkpoint_reader with _get_filesystem_reader; update PR description with test evidence.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: adding MultiStorageClient (MSC) support for fsdp_dtensor checkpoints.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

_{✏️ 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.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@src/megatron/bridge/training/checkpointing.py`:
- Line 1917: The code calls an undefined function _get_checkpoint_reader at
reader = _get_checkpoint_reader(checkpoint_name); replace that call with the
existing _get_filesystem_reader (the symbol imported earlier) so the code uses
_get_filesystem_reader(checkpoint_name) instead, and verify any surrounding
logic (e.g., variables used after reader is returned) still matches the reader
API used elsewhere (such as the other call at line ~2728).

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 0a4e457f-920d-4b46-9283-ebb0b40390c1

📥 Commits

Reviewing files that changed from the base of the PR and between 8dfe249 and e06901f.

📒 Files selected for processing (1)

• src/megatron/bridge/training/checkpointing.py

[yaoyu-33]

/ok to test 5654829

[yaoyu-33]

/ok to test 34e97a3