add async save support for fsdp checkpoints by dimapihtar · Pull Request #3339 · NVIDIA-NeMo/Megatron-Bridge

dimapihtar · 2026-04-15T15:42:41Z

What does this PR do ?

Adds async save support for fsdp checkpoints.

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)

Summary by CodeRabbit

Improvements
- Enhanced asynchronous checkpoint saving with extended support for distributed checkpoint formats.
- Added graceful fallback behavior when optional dependencies are unavailable.

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

copy-pr-bot · 2026-04-15T15:42:45Z

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.

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

dimapihtar · 2026-04-15T15:44:16Z

/ok to test dba0517

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

dimapihtar · 2026-04-15T15:47:15Z

/ok to test a67f015

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

dimapihtar · 2026-04-15T15:48:00Z

/ok to test 2be7c04

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

dimapihtar · 2026-04-15T15:49:47Z

/ok to test 2e15c5c

coderabbitai · 2026-04-15T15:49:50Z

📝 Walkthrough

Walkthrough

Added runtime detection for the NVIDIA resiliency extension and introduced async checkpoint support for fsdp_dtensor format. Updated async checkpoint format gating to allow both torch_dist and fsdp_dtensor formats, with conditional fallback to synchronous saving when the extension is unavailable.

Changes

Cohort / File(s)	Summary
Async Checkpointing Enhancement `src/megatron/bridge/training/checkpointing.py`	Added `HAVE_NVRX` feature flag for nvidia_resiliency_ext.checkpointing detection. Introduced `get_save_and_finalize_callbacks()` to construct async request objects. Expanded async checkpoint format gating to support `fsdp_dtensor` alongside `torch_dist`. Implemented async save path using `FileSystemWriterAsync`, `DefaultSavePlanner`, and `save_state_dict_async_plan()` when async conditions are met; otherwise falls back to synchronous `torch.distributed.checkpoint.save()`.

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 introduces async save support for FSDP checkpoints but implementation is incomplete: checkpointing.py allows fsdp_dtensor while config.py still restricts it, causing feature to fail.	Update config.py to allow async_save with both torch_dist and fsdp_dtensor formats, update corresponding test, and document test results validating the end-to-end functionality.

✅ 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 'add async save support for fsdp checkpoints' directly and clearly summarizes the main change: adding asynchronous saving functionality for FSDP (Fully Sharded Data Parallel) format checkpoints, which matches the primary objective and the changes in the checkpointing.py file.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

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

✨ Finishing Touches

📝 Generate docstrings

Create stacked PR
Commit on current branch

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch dpykhtar/fsdp_async_save_support

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

coderabbitai

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

src/megatron/bridge/training/checkpointing.py (1)

950-965: ⚠️ Potential issue | 🔴 Critical

Fix mixed async+sync execution in the FSDP DTensor save path.

When async is enabled, Line 961-965 still performs a synchronous torch.distributed.checkpoint.save(...) unconditionally. This causes duplicate writes and defeats async behavior. Also, when ckpt_cfg.async_save=True but HAVE_NVRX=False, async_save_request stays None, and later async assertions fail instead of falling back.

💡 Suggested fix (use an effective async flag + true fallback)

@@
-    async_save_request = None
+    async_save_request = None
+    effective_async_save = ckpt_cfg.async_save
@@
-            if ckpt_cfg.async_save and HAVE_NVRX:
+            if effective_async_save and HAVE_NVRX:
                 planner = torch.distributed.checkpoint.DefaultSavePlanner()
                 coordinator_rank = 0
                 fs_storage_writer = FileSystemWriterAsync(
                     checkpoint_name, thread_count=ckpt_cfg.dist_ckpt_workers, use_msc=ckpt_cfg.enable_msc
                 )
@@
                 save_state_dict_ret = save_state_dict_async_plan(
                     state_dict, fs_storage_writer, None, coordinator_rank, planner=planner, enable_cache=ckpt_cfg.ckpt_assume_constant_structure
                 )
                 async_save_request = get_save_and_finalize_callbacks(fs_storage_writer, save_state_dict_ret)
-            fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter(checkpoint_name)
-            torch.distributed.checkpoint.save(
-                state_dict=state_dict,
-                storage_writer=fs_storage_writer,
-            )
+            else:
+                if effective_async_save and not HAVE_NVRX:
+                    print_rank_0(
+                        "WARNING: async_save=True for fsdp_dtensor but nvidia_resiliency_ext async_ckpt is unavailable; "
+                        "falling back to synchronous save."
+                    )
+                effective_async_save = False
+                fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter(checkpoint_name)
+                torch.distributed.checkpoint.save(
+                    state_dict=state_dict,
+                    storage_writer=fs_storage_writer,
+                )
@@
-    if ckpt_type != CheckpointType.LOCAL:
-        if not ckpt_cfg.async_save:
+    if ckpt_type != CheckpointType.LOCAL:
+        if not effective_async_save:
@@
-        if ckpt_cfg.async_save:
+        if effective_async_save:
@@
-    if ckpt_cfg.async_save:
+    if effective_async_save:
@@
-            if ckpt_cfg.async_save:
+            if effective_async_save:
@@
-                if cfg.logger.log_progress and ckpt_cfg.async_save:
+                if cfg.logger.log_progress and effective_async_save:
@@
-        if ckpt_cfg.async_save:
+        if effective_async_save:
@@
-    if ckpt_cfg.async_save:
+    if effective_async_save:
@@
-        if ckpt_cfg.async_save:
+        if effective_async_save:
@@
-        if ckpt_cfg.async_save:
+        if effective_async_save:
@@
-    if ckpt_cfg.async_save:
+    if effective_async_save:
         schedule_async_save(state, async_save_request)
@@
-            save_dir, leave_ckpt_num=ckpt_cfg.most_recent_k, do_async=ckpt_cfg.async_save
+            save_dir, leave_ckpt_num=ckpt_cfg.most_recent_k, do_async=effective_async_save
         )

🤖 Prompt for AI Agents

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

In `@src/megatron/bridge/training/checkpointing.py` around lines 950 - 965, The
code unconditionally calls torch.distributed.checkpoint.save causing duplicate
sync writes and doesn't fallback when async is requested but HAVE_NVRX is false;
fix by computing an effective_async flag = ckpt_cfg.async_save and HAVE_NVRX,
only create FileSystemWriterAsync / call save_state_dict_async_plan /
get_save_and_finalize_callbacks when effective_async is true, otherwise create
the synchronous torch.distributed.checkpoint.FileSystemWriter and call
torch.distributed.checkpoint.save; ensure async_save_request is set/left None
predictably so later assertions can detect fallback. Reference symbols:
ckpt_cfg.async_save, HAVE_NVRX, FileSystemWriterAsync,
save_state_dict_async_plan, get_save_and_finalize_callbacks,
torch.distributed.checkpoint.FileSystemWriter,
torch.distributed.checkpoint.save, async_save_request.

🤖 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`:
- Around line 898-901: The config validation currently only allows async_save
for 'torch_dist', which blocks the new fsdp_dtensor async path; update the
validation that checks async_save (the assertion/conditional referencing
'torch_dist') to also permit 'fsdp_dtensor' (i.e., include 'fsdp_dtensor'
alongside 'torch_dist' in the allowed ckpt_format check for async_save) and
adjust any related error/assert messages to mention both formats (refer to
async_save, ckpt_format, and the 'fsdp_dtensor' and 'torch_dist' format
strings).

---

Outside diff comments:
In `@src/megatron/bridge/training/checkpointing.py`:
- Around line 950-965: The code unconditionally calls
torch.distributed.checkpoint.save causing duplicate sync writes and doesn't
fallback when async is requested but HAVE_NVRX is false; fix by computing an
effective_async flag = ckpt_cfg.async_save and HAVE_NVRX, only create
FileSystemWriterAsync / call save_state_dict_async_plan /
get_save_and_finalize_callbacks when effective_async is true, otherwise create
the synchronous torch.distributed.checkpoint.FileSystemWriter and call
torch.distributed.checkpoint.save; ensure async_save_request is set/left None
predictably so later assertions can detect fallback. Reference symbols:
ckpt_cfg.async_save, HAVE_NVRX, FileSystemWriterAsync,
save_state_dict_async_plan, get_save_and_finalize_callbacks,
torch.distributed.checkpoint.FileSystemWriter,
torch.distributed.checkpoint.save, async_save_request.

🪄 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 Plus

Run ID: 4dc37e54-6e6a-44ec-a288-38ffca6e5ddb

📥 Commits

Reviewing files that changed from the base of the PR and between 44797c1 and dba0517.

📒 Files selected for processing (1)

src/megatron/bridge/training/checkpointing.py

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

dimapihtar · 2026-04-15T16:00:26Z

/ok to test 04be225

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

dimapihtar · 2026-04-16T12:29:21Z

/ok to test 94ce030

yaoyu-33 · 2026-04-16T18:34:44Z

@dimapihtar needs tests

cspades · 2026-04-17T22:13:19Z

New checkpoint configs should be added to the checkpoint config: dist_ckpt_workers and enable_msc

[rank4]:   File "/opt/Megatron-Bridge/src/megatron/bridge/training/checkpointing.py", line 956, in save_checkpoint
[rank4]:     checkpoint_name, thread_count=ckpt_cfg.dist_ckpt_workers, use_msc=ckpt_cfg.enable_msc
[rank4]:                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: AttributeError: 'CheckpointConfig' object has no attribute 'dist_ckpt_workers'

[rank4]:   File "/opt/Megatron-Bridge/src/megatron/bridge/training/checkpointing.py", line 956, in save_checkpoint
[rank4]:     checkpoint_name, thread_count=ckpt_cfg.dist_ckpt_workers, use_msc=ckpt_cfg.enable_msc
[rank4]:                                                                       ^^^^^^^^^^^^^^^^^^^
[rank4]: AttributeError: 'CheckpointConfig' object has no attribute 'enable_msc'

So far looking good, just saved a checkpoint async. Doing stop-and-go tests...

In addition to this, use_precision_aware_optimizer=True requires my FusedAdam fix as mentioned before: NVIDIA/TransformerEngine#2795 and NVRx installed.

cspades

CheckpointConfig needs to be updated and then I'll approve!

dimapihtar · 2026-04-22T17:58:14Z

/ok to test 9f90f6c

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

dimapihtar · 2026-04-23T18:18:13Z

/ok to test ae4eacf

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

dimapihtar · 2026-04-23T18:22:05Z

/ok to test 558f799

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

dimapihtar · 2026-04-23T18:24:09Z

/ok to test 0fccf12

dimapihtar · 2026-04-24T16:21:53Z

/ok to test 59e3fe4

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

dimapihtar · 2026-04-24T17:54:26Z

/ok to test c60768d

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

dimapihtar · 2026-04-24T18:14:54Z

/ok to test dcfbfea

Signed-off-by: dimapihtar <dpykhtar@nvidia.com> Signed-off-by: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com> Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>

add async support for fsdp

806226b

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

revert MLM version

dba0517

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

dimapihtar requested a review from cspades April 15, 2026 15:43

dimapihtar marked this pull request as ready for review April 15, 2026 15:43

fix code style

a67f015

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

fix if statement

2be7c04

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

fix code style

2e15c5c

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

coderabbitai Bot reviewed Apr 15, 2026

View reviewed changes

Comment thread src/megatron/bridge/training/checkpointing.py Outdated

dimapihtar added 2 commits April 15, 2026 08:58

fix code style

83dce36

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

fix aassertion

04be225

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

dimapihtar and others added 2 commits April 16, 2026 05:28

fix code style

020cbb3

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

Merge branch 'main' into dpykhtar/fsdp_async_save_support

94ce030

copy-pr-bot Bot temporarily deployed to test April 16, 2026 12:30 Inactive

yaoyu-33 added feature New capabilities, enhancements, or enablement work area:ckpt Checkpoint conversion, loading, export, and save paths needs-review PR is ready for code review and waiting on a reviewer labels Apr 16, 2026

cspades mentioned this pull request Apr 16, 2026

[feature] Support Asynchronous FSDP Checkpointing #3130

Closed

yaoyu-33 added needs-author and removed needs-review PR is ready for code review and waiting on a reviewer labels Apr 16, 2026

cspades requested changes Apr 17, 2026

View reviewed changes

copy-pr-bot Bot temporarily deployed to test April 22, 2026 17:59 Inactive

dimapihtar requested a review from cspades April 22, 2026 18:00

cspades previously approved these changes Apr 22, 2026

View reviewed changes

dimapihtar enabled auto-merge (squash) April 22, 2026 20:39

dimapihtar added 2 commits April 23, 2026 08:29

add fsdp-async functional test

a9a0961

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

remove extra file

c668bf3

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

dimapihtar dismissed cspades’s stale review via c668bf3 April 23, 2026 15:29

dimapihtar added 3 commits April 23, 2026 08:30

fix typo

28539cd

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

add unit test

40dd348

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

fix functional test

ae4eacf

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

fix code style

558f799

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

fix code style

0fccf12

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

copy-pr-bot Bot temporarily deployed to test April 23, 2026 18:24 Inactive

Merge branch 'main' into dpykhtar/fsdp_async_save_support

59e3fe4

copy-pr-bot Bot temporarily deployed to test April 24, 2026 16:22 Inactive

fix unit test

c60768d

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

fix style

dcfbfea

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

copy-pr-bot Bot temporarily deployed to test April 24, 2026 18:15 Inactive

dimapihtar requested a review from cspades April 24, 2026 18:53

cspades approved these changes Apr 24, 2026

View reviewed changes

dimapihtar merged commit 4963de0 into main Apr 24, 2026
80 checks passed

dimapihtar deleted the dpykhtar/fsdp_async_save_support branch April 24, 2026 19:06

Uh oh!

Conversation

dimapihtar commented Apr 15, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Changelog

GitHub Actions CI

Before your PR is "Ready for review"

Additional Information

Summary by CodeRabbit

Uh oh!

copy-pr-bot Bot commented Apr 15, 2026

Uh oh!

dimapihtar commented Apr 15, 2026

Uh oh!

dimapihtar commented Apr 15, 2026

Uh oh!

dimapihtar commented Apr 15, 2026

Uh oh!

dimapihtar commented Apr 15, 2026

Uh oh!

coderabbitai Bot commented Apr 15, 2026

Walkthrough

Changes

Estimated code review effort

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

dimapihtar commented Apr 15, 2026

Uh oh!

dimapihtar commented Apr 16, 2026

Uh oh!

yaoyu-33 commented Apr 16, 2026

Uh oh!

cspades commented Apr 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cspades left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dimapihtar commented Apr 22, 2026

Uh oh!

dimapihtar commented Apr 23, 2026

Uh oh!

dimapihtar commented Apr 23, 2026

Uh oh!

dimapihtar commented Apr 23, 2026

Uh oh!

dimapihtar commented Apr 24, 2026

Uh oh!

dimapihtar commented Apr 24, 2026

Uh oh!

dimapihtar commented Apr 24, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

dimapihtar commented Apr 15, 2026 •

edited by coderabbitai Bot

Loading

cspades commented Apr 17, 2026 •

edited

Loading

cspades left a comment •

edited

Loading