Revert "ci: add base_sha to codecov/codecov-action upload step (#4445)"

.github/workflows/cicd-main.yml

@@ -1006,11 +1006,6 @@ jobs:
       matrix:
         flag: [unit-test]
     steps:
-      - name: Get PR info
-        id: get-pr-info
-        if: startsWith(github.ref, 'refs/heads/pull-request/') && github.event_name == 'push'
-        uses: nv-gha-runners/get-pr-info@main
-
       - name: Checkout
         uses: actions/checkout@v6
 
@@ -1041,7 +1036,6 @@ jobs:
           token: ${{ secrets.CODECOV_TOKEN }}
           verbose: true
           flags: ${{ matrix.flag }}
-          base_sha: ${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').base.sha }}
 
       - name: Upload artifacts
         uses: actions/upload-artifact@v6

megatron/core/dist_checkpointing/serialization.py

@@ -148,7 +148,12 @@ def load(
         ckpt_sharded_metadata,
     )
 
-    async_strategy = getattr(common_state_dict.get("args"), "async_strategy", "nvrx")
+    ckpt_args = common_state_dict.get("args")
+    async_strategy = (
+        getattr(ckpt_args, "async_strategy", "mcore")
+        if getattr(ckpt_args, "async_save", False)
+        else "mcore"
+    )
     loaded_state_dict = sharded_strategy.load(sharded_state_dict, checkpoint_dir, async_strategy)
 
     merge(common_state_dict, loaded_state_dict)

megatron/core/dist_checkpointing/strategies/fully_parallel.py

@@ -189,7 +189,7 @@ def load(
         self,
         sharded_state_dict: ShardedStateDict,
         checkpoint_dir: Path,
-        async_strategy: str = "nvrx",
+        async_strategy: str = "mcore",
     ) -> StateDict:
         """Distributes the load and calls underlying strategy only for parts of the state dict.
 

megatron/core/dist_checkpointing/strategies/nvrx.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""Helpers for interacting with the experimental nvidia-resiliency-ext API."""
+
+from importlib import import_module
+from typing import Any, Callable, Dict
+
+
+def has_nvrx_async_support() -> bool:
+    """Checks whether the NVRx async checkpointing symbols Megatron uses are importable."""
+    try:
+        core = import_module("nvidia_resiliency_ext.checkpointing.async_ckpt.core")
+        cached_metadata_reader = import_module(
+            "nvidia_resiliency_ext.checkpointing.async_ckpt.cached_metadata_filesystem_reader"
+        )
+        filesystem_async = import_module(
+            "nvidia_resiliency_ext.checkpointing.async_ckpt.filesystem_async"
+        )
+        state_dict_saver = import_module(
+            "nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver"
+        )
+    except (ImportError, ModuleNotFoundError):
+        return False
+
+    required_symbols = (
+        getattr(core, "AsyncCallsQueue", None),
+        getattr(core, "AsyncRequest", None),
+        getattr(cached_metadata_reader, "CachedMetadataFileSystemReader", None),
+        getattr(filesystem_async, "FileSystemWriterAsync", None),
+        getattr(filesystem_async, "get_write_results_queue", None),
+        getattr(state_dict_saver, "CheckpointMetadataCache", None),
+        getattr(state_dict_saver, "save_state_dict_async_finalize", None),
+        getattr(state_dict_saver, "save_state_dict_async_plan", None),
+    )
+    return all(symbol is not None for symbol in required_symbols) and hasattr(
+        filesystem_async, "_results_queue"
+    )
+
+
+def make_nvrx_async_request(
+    async_request_cls: type,
+    async_fn: Callable[..., Any],
+    async_fn_args: Any,
+    finalize_fns: list[Callable[..., Any]],
+    async_fn_kwargs: Dict[str, Any] | None = None,
+    preload_fn: Callable[..., Any] | None = None,
+):
+    """Builds an AsyncRequest using the expected NVRx API."""
+    return async_request_cls(
+        async_fn,
+        async_fn_args,
+        finalize_fns,
+        async_fn_kwargs=async_fn_kwargs or {},
+        preload_fn=preload_fn,
+    )

megatron/core/dist_checkpointing/strategies/torch.py

@@ -6,13 +6,12 @@
 import os
 import pickle
 import warnings
-from abc import ABC
 from collections import defaultdict
 from contextlib import contextmanager
 from itertools import product
 from logging import getLogger
 from pathlib import Path
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union, cast
 
 import torch
 from packaging.version import Version as PkgVersion
@@ -51,19 +50,18 @@
 )
 from .async_utils import AsyncRequest
 from .checkpointable import CheckpointableShardedTensor, LocalShardsContainer
+from .nvrx import has_nvrx_async_support, make_nvrx_async_request
 
-try:
+if TYPE_CHECKING:
     from nvidia_resiliency_ext.checkpointing.async_ckpt.core import AsyncRequest as NVRxAsyncRequest
     from nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver import (
         CheckpointMetadataCache,
     )
+else:
+    CheckpointMetadataCache = Any
+    NVRxAsyncRequest = Any
 
-    HAVE_NVRX = True
-except (ImportError, ModuleNotFoundError):
-    CheckpointMetadataCache = ABC
-    NVRxAsyncRequest = ABC
-
-    HAVE_NVRX = False
+HAVE_NVRX = has_nvrx_async_support()
 
 try:
     if not torch.cuda.is_available():
@@ -103,6 +101,7 @@ class MCoreSavePlan:
 
 
 logger = getLogger(__name__)
+_logged_mcore_async_deprecation = False
 
 
 def flatten_state_dict(
@@ -651,9 +650,8 @@ def __init__(
         self.validated_loaded_metadata_reuse = False
 
     def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path):
-        """Each async strategy can be trivially used as a sync strategy."""
-        strategy = "nvrx" if HAVE_NVRX else "mcore"
-        async_request = self.async_save(sharded_state_dict, checkpoint_dir, async_strategy=strategy)
+        """Sync save always uses the built-in implementation."""
+        async_request = self.async_save(sharded_state_dict, checkpoint_dir, async_strategy="mcore")
         async_request.execute_sync()
         del async_request
 
@@ -671,11 +669,14 @@ def async_save(
 
         Returns: None
         """
+        global _logged_mcore_async_deprecation
         if async_strategy == "mcore":
-            logger.warning(
-                "MCore's async save is deprecated and will be removed in the future releases. "
-                "Please, use NVRx async solution by setting `async_strategy` to `nvrx`."
-            )
+            if not _logged_mcore_async_deprecation:
+                logger.warning(
+                    "MCore's async save is deprecated and will be removed in the future releases. "
+                    "Please, use NVRx async solution by setting `async_strategy` to `nvrx`."
+                )
+                _logged_mcore_async_deprecation = True
 
         # Translate the state dict
         (sharded_state_dict, flat_mapping, rename_mapping) = (
@@ -701,7 +702,9 @@ def async_save(
         if async_strategy == "nvrx":
             if self._metadata_cache is None:
                 self._metadata_cache = checkpointable_metadata_cache()
-                if self.cached_global_metadata is not None:
+                if self.cached_global_metadata is not None and hasattr(
+                    self._metadata_cache, "set_cached_global_metadata"
+                ):
                     self._metadata_cache.set_cached_global_metadata(self.cached_global_metadata)
             # Define additional arguments
             async_writer_kwargs["use_cached_data_structure"] = self.use_cached_ckpt_structure
@@ -818,11 +821,13 @@ def _get_save_and_finalize_callbacks(
         def finalize_fn():
             save_state_dict_async_finalize(*save_state_dict_ret)
 
-        return async_request(save_fn, save_args, [finalize_fn], preload_fn=preload_fn)
+        return make_nvrx_async_request(
+            async_request, save_fn, save_args, [finalize_fn], preload_fn=preload_fn
+        )
 
 
 def _get_filesystem_reader(
-    checkpoint_dir: Union[str, Path], cache_metadata: bool = False, async_strategy: str = "nvrx"
+    checkpoint_dir: Union[str, Path], cache_metadata: bool = False, async_strategy: str = "mcore"
 ) -> FileSystemReader:
     if MultiStorageClientFeature.is_enabled():
         msc = MultiStorageClientFeature.import_package()
@@ -846,7 +851,7 @@ def load(
         self,
         sharded_state_dict: ShardedStateDict,
         checkpoint_dir: Path,
-        async_strategy: str = "nvrx",
+        async_strategy: str = "mcore",
     ) -> StateDict:
         """Translates MCore ShardedTensors to PyT ShardedTensors & loads from PyT Distributed fmt.
 
@@ -1062,9 +1067,8 @@ def get_async_strategy(async_strategy: str = "nvrx", module: str = None) -> tupl
             async_strategy = "nvrx"
         except (ImportError, ModuleNotFoundError):
             raise ModuleNotFoundError(
-                "nvidia-resiliency-ext package is not installed. "
-                "Please, install nvidia-resiliency-ext package or set `async_strategy` to `mcore` "
-                "to enable async save strategy."
+                "A compatible `nvidia-resiliency-ext` installation is required for "
+                '`async_strategy="nvrx"`. Please install it or set `async_strategy` to `mcore`.'
             )
     elif async_strategy == "mcore":
         # do mcore async imports

megatron/training/arguments.py

@@ -1535,6 +1535,9 @@ def validate_args(args, defaults={}):
             )
             args.async_save = False
 
+    if not args.async_save:
+        args.async_strategy = "mcore"
+
     # Inference args
     if args.inference_batch_times_seqlen_threshold > -1:
         assert args.pipeline_model_parallel_size > 1, \

megatron/training/async_utils.py

@@ -8,14 +8,22 @@
 import logging
 import time
 from abc import ABC
+from typing import TYPE_CHECKING, Any
 
 from megatron.core.dist_checkpointing.strategies.async_utils import AsyncRequest
+from megatron.core.dist_checkpointing.strategies.nvrx import (
+    make_nvrx_async_request,
+)
 from megatron.core.dist_checkpointing.strategies.torch import get_async_strategy
 from megatron.training import get_args
 from megatron.training.utils import print_rank_0
 
-try:
+if TYPE_CHECKING:
     from nvidia_resiliency_ext.checkpointing.async_ckpt.core import AsyncRequest as NVRxAsyncRequest
+else:
+    NVRxAsyncRequest = Any
+
+try:
     from nvidia_resiliency_ext.checkpointing.async_ckpt.filesystem_async import _results_queue
     from nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver import (
         save_state_dict_async_finalize,
@@ -26,8 +34,6 @@
         save_state_dict_async_finalize,
     )
 
-    NVRxAsyncRequest = ABC
-
 logger = logging.getLogger(__name__)
 
 # Singleton manager of async calls
@@ -70,13 +76,13 @@ def init_persistent_async_worker(rank: int, mp_mode: str = 'spawn'):
         ),
     )
     # initialize the persistent caller with QoS priorities from args
-    kwargs = {}
+    warmup_kwargs = {}
     if async_strategy == "mcore":
         # Note: nvidia-resiliency-ext uses is_daemon instead of mp_mode (always spawns)
-        kwargs["mp_mode"] = mp_mode
+        warmup_kwargs["mp_mode"] = mp_mode
     elif async_strategy == "nvrx":
         if "cpu_shm_mode" in inspect.signature(AsyncCallsQueue.warmup_persistent_caller).parameters:
-            kwargs["cpu_shm_mode"] = args.async_ckpt_use_cpu_shm
+            warmup_kwargs["cpu_shm_mode"] = args.async_ckpt_use_cpu_shm
         elif args.async_ckpt_use_cpu_shm:
             raise AssertionError(
                 "Installed nvidia-resiliency-ext does not support cpu_shm_mode. "
@@ -86,10 +92,16 @@ def init_persistent_async_worker(rank: int, mp_mode: str = 'spawn'):
         rank,
         cpu_priority=args.async_ckpt_cpu_priority,
         io_priority=args.async_ckpt_io_priority,
-        **kwargs,
+        **warmup_kwargs,
     )
     # initialize ckpt write results queue
-    get_write_results_queue('fork')
+    if async_strategy == "nvrx":
+        if "mp_mode" not in inspect.signature(get_write_results_queue).parameters:
+            raise AssertionError(
+                "Installed nvidia-resiliency-ext does not support "
+                "get_write_results_queue(mp_mode=...). Update nvidia-resiliency-ext."
+            )
+    get_write_results_queue(mp_mode="fork")
     if rank == 0:
         print(f"init_persistent_async_worker: rank {rank}, Async Caller Started in {time.time() - time_start} seconds", flush=True)
 
@@ -157,14 +169,24 @@ def reset_persistent_async_worker(async_strategy):
     module.clear_metadata_cache()
 
 
-def get_save_and_finalize_callbacks(writer, save_state_dict_ret) -> NVRxAsyncRequest:
+def get_save_and_finalize_callbacks(
+    writer, save_state_dict_ret, async_strategy: str = "nvrx"
+) -> AsyncRequest | 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()
+    _, async_modules = get_async_strategy(async_strategy)
+    async_request_cls = async_modules["AsyncRequest"]
+    save_state_dict_async_finalize = async_modules["save_state_dict_async_finalize"]
 
     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
+    return make_nvrx_async_request(
+        async_request_cls,
+        save_fn,
+        save_args,
+        [finalize_fn],
+        async_fn_kwargs={},
+        preload_fn=preload_fn,
     )

megatron/training/checkpointing.py

@@ -33,6 +33,7 @@
 from megatron.core.dist_checkpointing.strategies.torch import (
     TorchDistLoadShardedStrategy,
     TorchDistSaveShardedStrategy,
+    get_async_strategy,
 )
 from megatron.core.msc_utils import MultiStorageClientFeature, open_file
 from megatron.core.num_microbatches_calculator import update_num_microbatches
@@ -72,19 +73,6 @@
     has_nvidia_modelopt = False
 
 
-try:
-    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_plan,
-    )
-
-    HAVE_NVRX = True
-except (ImportError, ModuleNotFoundError):
-
-    HAVE_NVRX = False
-
 _CHECKPOINT_VERSION = None
 _LOADED_ITERATION = None
 
@@ -694,6 +682,9 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati
             if args.async_save:
                 planner = torch.distributed.checkpoint.DefaultSavePlanner()
                 coordinator_rank = 0
+                _, async_modules = get_async_strategy(args.async_strategy)
+                FileSystemWriterAsync = async_modules["FileSystemWriterAsync"]
+                save_state_dict_async_plan = async_modules["save_state_dict_async_plan"]
                 _cpu_shm = getattr(args, 'async_ckpt_use_cpu_shm', False)
                 _writer_kwargs = {}
                 if _cpu_shm:
@@ -718,7 +709,9 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati
                 save_state_dict_ret = save_state_dict_async_plan(
                     state_dict, fs_storage_writer, None, coordinator_rank, planner=planner, enable_cache=args.ckpt_assume_constant_structure
                 )
-                async_save_request = get_save_and_finalize_callbacks(fs_storage_writer, save_state_dict_ret)
+                async_save_request = get_save_and_finalize_callbacks(
+                    fs_storage_writer, save_state_dict_ret, args.async_strategy
+                )
             else:
                 fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter(checkpoint_name)
                 torch.distributed.checkpoint.save(