add async save support for fsdp checkpoints

src/megatron/bridge/training/checkpointing.py

@@ -21,6 +21,7 @@
 import shutil
 import sys
 import threading
+from abc import ABC
 from dataclasses import dataclass, replace
 from enum import Enum, auto
 from logging import getLogger
@@ -114,6 +115,19 @@
 except ImportError:
     handle_gdn_in_state_dict = None
 
+try:
+    from nvidia_resiliency_ext.checkpointing.async_ckpt.core import AsyncRequest as NVRxAsyncRequest
+    from nvidia_resiliency_ext.checkpointing.async_ckpt.filesystem_async import FileSystemWriterAsync
+    from nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver import (
+        save_state_dict_async_finalize,
+        save_state_dict_async_plan,
+    )
+
+    HAVE_NVRX = True
+except (ImportError, ModuleNotFoundError):
+    NVRxAsyncRequest = ABC
+    HAVE_NVRX = False
+
 TRACKER_PREFIX = "latest"
 _CHECKPOINT_VERSION = None
 
@@ -389,6 +403,17 @@ def is_empty_async_queue(global_state: GlobalState) -> bool:
     return async_queue.get_num_unfinalized_calls() == 0
 
 
+def get_save_and_finalize_callbacks(writer, save_state_dict_ret) -> NVRxAsyncRequest:
+    """Creates an async save request for fsdp_dtensor & torch_dcp with a finalize function."""
+    save_fn, preload_fn, save_args = writer.get_save_function_and_args()
+
+    def finalize_fn():
+        """Finalizes async checkpointing and synchronizes processes."""
+        save_state_dict_async_finalize(*save_state_dict_ret)
+
+    return NVRxAsyncRequest(save_fn, save_args, [finalize_fn], async_fn_kwargs={}, preload_fn=preload_fn)
+
+
 def get_rng_state(
     data_parallel_random_init: bool,
     ckpt_format: str = "torch_dist",
@@ -875,7 +900,7 @@ def save_checkpoint(
 
     async_save_request = None
     if ckpt_cfg.async_save:
-        if ckpt_type == CheckpointType.GLOBAL and ckpt_cfg.ckpt_format != "torch_dist":
+        if ckpt_type == CheckpointType.GLOBAL and ckpt_cfg.ckpt_format not in ["torch_dist", "fsdp_dtensor"]:
             raise NotImplementedError(
                 f"Async checkpoint save not implemented for {ckpt_cfg.ckpt_format} distributed checkpoint format"
             )
@@ -962,16 +987,35 @@ def save_checkpoint(
             state_dict = preprocess_fsdp_dtensor_state_dict(cfg, state_dict, model[0])
 
             # FSDP DTensor checkpoint save path using PyTorch Distributed Checkpointing
-            if MultiStorageClientFeature.is_enabled():
-                from multistorageclient.contrib.torch.filesystem import MultiStorageFileSystemWriter
+            if ckpt_cfg.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=MultiStorageClientFeature.is_enabled(),
+                )
 
-                fs_storage_writer = MultiStorageFileSystemWriter(checkpoint_name)
+                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)
             else:
-                fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter(checkpoint_name)
-            torch.distributed.checkpoint.save(
-                state_dict=state_dict,
-                storage_writer=fs_storage_writer,
-            )
+                if MultiStorageClientFeature.is_enabled():
+                    from multistorageclient.contrib.torch.filesystem import MultiStorageFileSystemWriter
+
+                    fs_storage_writer = MultiStorageFileSystemWriter(checkpoint_name)
+                else:
+                    fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter(checkpoint_name)
+                torch.distributed.checkpoint.save(
+                    state_dict=state_dict,
+                    storage_writer=fs_storage_writer,
+                )
         else:
             # torch_dist and other formats using MCore distributed checkpointing
             if checkpointing_context is not None and "save_strategy" in checkpointing_context:

src/megatron/bridge/training/config.py

@@ -1199,8 +1199,8 @@ def validate(self) -> None:
 
         # Enforce async_save format restriction
         if self.checkpoint.async_save:
-            assert self.checkpoint.ckpt_format == "torch_dist", (
-                "async_save is only supported with ckpt_format='torch_dist'"
+            assert self.checkpoint.ckpt_format in ["torch_dist", "fsdp_dtensor"], (
+                "async_save is only supported with ckpt_format='torch_dist','fsdp_dtensor'"
             )
 
         # Set defaults for tensor inspect callback

tests/functional_tests/test_groups/recipes/test_llama_recipes_pretrain_1b.py

@@ -26,23 +26,59 @@
 
 
 LLAMA_PRETRAIN_RECIPES = [
-    # (config_func, name, parallelism_overrides, model_overrides)
-    (llama32_1b_config, "llama32_1b", {}, {"num_layers": 2}),
-    (llama32_3b_config, "llama32_3b", {}, {"num_layers": 2}),
+    # (config_func, name, parallelism_overrides, model_overrides, checkpoint_overrides, ddp_overrides)
+    (llama32_1b_config, "llama32_1b", {}, {"num_layers": 2}, {}, {}),
+    (llama32_3b_config, "llama32_3b", {}, {"num_layers": 2}, {}, {}),
+    # FSDP-async test case
+    (
+        llama32_1b_config,
+        "llama32_1b",
+        {},
+        {"num_layers": 2},
+        {
+            "ckpt_format": "fsdp_dtensor",
+            "strict_fsdp_dtensor_load": True,
+            "async_save": True,
+            "async_strategy": "nvrx",
+            "use_persistent_ckpt_worker": True,
+            "dist_ckpt_workers": 1,
+        },
+        {
+            "use_megatron_fsdp": True,
+            "use_distributed_optimizer": True,
+            "grad_reduce_in_fp32": True,
+            "average_in_collective": True,
+            "data_parallel_sharding_strategy": "optim_grads_params",
+        },
+    ),
 ]
 
 
 class TestLlamaRecipes:
     """Test class for LLaMA recipe functional tests."""
 
     @pytest.mark.run_only_on("GPU")
-    @pytest.mark.parametrize("config_func,recipe_name,parallelism_overrides,model_overrides", LLAMA_PRETRAIN_RECIPES)
-    def test_llama_pretrain_recipes(self, config_func, recipe_name, parallelism_overrides, model_overrides, tmp_path):
+    @pytest.mark.parametrize(
+        "config_func,recipe_name,parallelism_overrides,model_overrides,checkpoint_overrides,ddp_overrides",
+        LLAMA_PRETRAIN_RECIPES,
+    )
+    def test_llama_pretrain_recipes(
+        self,
+        config_func,
+        recipe_name,
+        parallelism_overrides,
+        model_overrides,
+        checkpoint_overrides,
+        ddp_overrides,
+        tmp_path,
+    ):
         """Functional test for LLaMA recipes with appropriate parallelism configurations."""
         run_pretrain_recipe_test(
             config_func,
             recipe_name,
             tmp_path,
             model_overrides=model_overrides,
+            checkpoint_overrides=checkpoint_overrides,
+            ddp_overrides=ddp_overrides,
             **parallelism_overrides,
         )

tests/functional_tests/test_groups/recipes/utils.py

@@ -36,6 +36,8 @@ def run_pretrain_recipe_test(
     pipeline_model_parallel_size: Optional[int] = None,
     expert_model_parallel_size: Optional[int] = None,
     model_overrides: Optional[dict] = None,
+    checkpoint_overrides: Optional[dict] = None,
+    ddp_overrides: Optional[dict] = None,
 ):
     """
     Common test implementation for pretrain recipe configurations.
@@ -116,6 +118,14 @@ def run_pretrain_recipe_test(
             for attribute_name, attribute_value in model_overrides.items():
                 setattr(config.model, attribute_name, attribute_value)
 
+        if checkpoint_overrides:
+            for attribute_name, attribute_value in checkpoint_overrides.items():
+                setattr(config.checkpoint, attribute_name, attribute_value)
+
+        if ddp_overrides:
+            for attribute_name, attribute_value in ddp_overrides.items():
+                setattr(config.ddp, attribute_name, attribute_value)
+
         pretrain(config, forward_step)
 
         # Basic verification that training completed successfully

tests/unit_tests/training/test_config.py

@@ -1985,9 +1985,7 @@ def test_async_save_format_validation_fsdp_dtensor_fails(self, monkeypatch):
             dist_config=dist_cfg,
         )
         try:
-            # Should raise error - async_save with fsdp_dtensor is not allowed
-            with pytest.raises(AssertionError, match="async_save is only supported with ckpt_format='torch_dist'"):
-                container.validate()
+            container.validate()
         finally:
             restore_get_world_size_safe(og_ws, cfg_mod)