chore: import upstream snapshot with attribution

nemo/utils/callbacks/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.utils.callbacks.cuda_graph import CUDAGraphCallback
+from nemo.utils.callbacks.nemo_model_checkpoint import NeMoModelCheckpoint
+from nemo.utils.callbacks.preemption import PreemptionCallback

nemo/utils/callbacks/cuda_graph.py

@@ -0,0 +1,451 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# CUDAGraphCallback is a full iteration CUDA graph callback designed for
+# models with PyTorch Lightning first, this has been tested with Stable
+# Diffusion right now.
+#
+# Prerequisites for this callback:
+# 1. Capturable: user has to make sure (almost) all the host & device
+#    synchronizations are removed, some of the syncs regarding logging
+#    of metrics introduced by PyTorch Lightning itself have been removed
+#    by this callback. This ensures the graph can be captured.
+# 2. Topology: user has to make sure there's no dynamic control flow
+#    within the iteration. Please use APEX alternatives for building
+#    blocks that contain dynamic control flow, e.g. gradient clipping.
+#    Otherwise the captured graph can run, but may raise silent failure,
+#    e.g. NaN loss.
+# 3. Parameters: user has to make sure pointers involved in the graph
+#    capturing range don't change across iterations. In this case users
+#    have to ensure that data is copied to static tensors. Otherwise this
+#    can also lead to silent failure.
+
+import os
+import time
+from dataclasses import dataclass
+from types import MethodType
+from typing import Any, Dict
+
+import lightning.pytorch as pl
+import torch
+from lightning.pytorch import LightningModule
+from lightning.pytorch.callbacks import Callback
+from lightning.pytorch.loops.optimization.automatic import ClosureResult
+from lightning.pytorch.trainer.connectors.logger_connector.result import _ResultCollection, _ResultMetric
+from lightning.pytorch.utilities import CombinedLoader, rank_zero_info
+from lightning.pytorch.utilities.signature_utils import is_param_in_hook_signature
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from torch.nn.parallel import DistributedDataParallel
+
+__all__ = ["CUDAGraphCallback"]
+
+
+def struct_copy_one(src):
+    if isinstance(src, tuple):
+        return tuple(struct_copy_one(i) for i in src)
+    elif isinstance(src, list):
+        return list(struct_copy_one(i) for i in src)
+    elif isinstance(src, dict):
+        return {k: struct_copy_one(src[k]) for k in src}
+    elif isinstance(src, torch.Tensor):
+        return src.clone().detach().cuda()
+    else:
+        return src
+
+
+def struct_copy_two(tgt, src):
+    if isinstance(src, tuple):
+        raise Exception(f"Unsupported copy for tuple yet: {type(src)}")
+    elif isinstance(src, list):
+        for i in range(len(src)):
+            if isinstance(src[i], (tuple, list, dict, torch.Tensor)):
+                struct_copy_two(tgt[i], src[i])
+            else:
+                tgt[i] = src[i]
+    elif isinstance(src, dict):
+        for k in src:
+            if isinstance(src[k], (tuple, list, dict, torch.Tensor)):
+                struct_copy_two(tgt[k], src[k])
+            else:
+                tgt[k] = src[k]
+    elif isinstance(src, torch.Tensor):
+        tgt.copy_(src, non_blocking=True)
+    else:
+        raise Exception(f"Expect top-level as container type but got: {type(src)}")
+
+
+class StaticBufferLoader:
+    """Load data to static buffers."""
+
+    def __init__(self, loader):
+        self.loader = loader
+        self.stream = torch.cuda.Stream()
+        self.static = None
+
+    def __iter__(self):
+        for inputs in self.loader:
+            if self.static is None:
+                with torch.cuda.stream(self.stream):
+                    self.static = struct_copy_one(inputs)
+
+            with torch.cuda.stream(self.stream):
+                struct_copy_two(self.static, inputs)
+            torch.cuda.current_stream().wait_stream(self.stream)
+            yield self.static
+
+    def __len__(self):
+        return len(self.loader)
+
+
+def get_lr(lr_scheduler):
+    lrs = lr_scheduler.__orig_get_lr__()
+    if not hasattr(lr_scheduler, "static_lrs"):
+        lr_scheduler.static_lrs = lrs
+    for i in range(len(lrs)):
+        lr_scheduler.static_lrs[i].copy_(lrs[i])
+    return lr_scheduler.static_lrs
+
+
+def zero_grad(optimizer, *args, **kwargs):
+    # We invoke zero_grad before graph capturing.
+    if torch.cuda.is_current_stream_capturing():
+        rank_zero_info("CUDAGraphCallback: set optimizer.zero_grad as nop during graph capturing.")
+    else:
+        optimizer.__orig_zero_grad__(*args, **kwargs)
+
+
+def to_tensor(self, value, name):
+    # Log metrics in PyTorch Lightning often invokes CPU & GPU synchronizations. Here
+    # we implement smart metrics to avoid those synchronizations.
+    # Refer to: https://github.com/Lightning-AI/pytorch-lightning/blob/2.0.7/src/lightning/pytorch/core/module.py#L615
+    value = value.clone().detach() if isinstance(value, torch.Tensor) else torch.tensor(value)
+    if not torch.numel(value) == 1:
+        raise ValueError(
+            f"`self.log({name}, {value})` was called, but the tensor must have a single element."
+            f" You can try doing `self.log({name}, {value}.mean())`"
+        )
+    value = value.squeeze()
+    return value
+
+
+def get_optimizer_step(state):
+    def optimizer_step(
+        self,
+        epoch,
+        batch_idx,
+        optimizer,
+        optimizer_closure=None,
+    ) -> None:
+        # Not all optimizer supports set_to_none.
+        if not hasattr(optimizer, "support_set_to_none"):
+            optimizer.support_set_to_none = is_param_in_hook_signature(
+                optimizer.zero_grad, "set_to_none", explicit=True
+            )
+        if optimizer.support_set_to_none:
+            zero_grad_kwargs = {"set_to_none": True}
+        else:
+            zero_grad_kwargs = {}
+
+        if 0 <= state.current_iteration < state.capture_iteration or state.capture_iteration < 0:
+            state.stream.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(state.stream):
+                optimizer.zero_grad(**zero_grad_kwargs)
+                self.__orig_optimizer_step__(
+                    epoch,
+                    batch_idx,
+                    optimizer,
+                    optimizer_closure=optimizer_closure,
+                )
+            torch.cuda.current_stream().wait_stream(state.stream)
+
+        if state.current_iteration == state.capture_iteration:
+            torch.cuda.synchronize()
+            # Sleep for one second to let environment stable
+            time.sleep(1)
+            rank_zero_info("CUDAGraphCallback: capturing CUDA graph for module %s.", self.__class__.__name__)
+            with torch.cuda.graph(state.graph, stream=state.stream, capture_error_mode="global"):
+                # PyTorch CUDA graph doc for whole-network capturing mentions:
+                #
+                #   Sets grads to None before capture, so backward() will create
+                #   .grad attributes with allocations from the graph's private pool
+                #
+                # But it's not necessary, and it can lead to CUDA kernels inside
+                # `zero_grad()` being not captured.
+                optimizer.zero_grad(**zero_grad_kwargs)
+                self.__orig_optimizer_step__(
+                    epoch,
+                    batch_idx,
+                    optimizer,
+                    optimizer_closure=optimizer_closure,
+                )
+            torch.cuda.synchronize()
+
+        # Graph replay and reconstruct missing result
+        if state.current_iteration >= state.capture_iteration >= 0:
+            state.graph.replay()
+            optimizer_closure._result = ClosureResult.from_training_step_output(state.output)
+
+        # If something is not capturable, try to put it there, e.g. `self.log()`.
+        if hasattr(self, "non_cuda_graph_capturable"):
+            self.non_cuda_graph_capturable()
+
+        state.current_iteration += 1
+
+    return optimizer_step
+
+
+def get_training_step(state):
+    def training_step(self, batch):
+        results = self.__orig_training_step__(batch)
+        if state.output is None:
+            state.output = struct_copy_one(results)
+
+        # Copy results to static buffer to rebuild states required by PL.
+        with torch.no_grad():
+            struct_copy_two(state.output, results)
+        return results
+
+    return training_step
+
+
+def get_amp_autocast_init(state):
+    def amp_autocast_init(self, *args, **kwargs):
+        if "cache_enabled" not in kwargs:
+            kwargs["cache_enabled"] = False
+        if state.current_iteration == 0:
+            rank_zero_info("CUDAGraphCallback: disable autocast cache.")
+        return self.__orig_init__(*args, **kwargs)
+
+    return amp_autocast_init
+
+
+def get_ddp_init(state):
+    def init(self, *args, **kwargs):
+        rank_zero_info("CUDAGraphCallback: init DDP on side stream.")
+        with torch.cuda.stream(state.stream):
+            self.__orig_init__(*args, **kwargs)
+
+    return init
+
+
+@dataclass
+class CUDAGraphState:
+    current_iteration: int = 0
+    capture_iteration: int = -1  # -1 to disable
+    stream: torch.cuda.Stream = None
+    graph: torch.cuda.CUDAGraph = None
+    output: Any = None  # static forward output
+
+
+class CUDAGraphCallback(Callback):
+    """Full iteration CUDA graph callback.
+
+    Dataloader and LR scheduler are not included in the CUDA graph with this callback.
+    """
+
+    def __init__(self, capture_iteration=-1):
+        super().__init__()
+
+        # Required by CUDA graph with DDP
+        # Ref: https://pytorch.org/docs/stable/notes/cuda.html#usage-with-distributeddataparallel
+        if 0 <= capture_iteration <= 11:
+            raise Exception("Warmup must run at least 11 DDP-enabled eager iterations before capture.")
+        if torch.distributed.is_initialized():
+            raise Exception("CUDAGraphCallback should be initialized before process group.")
+        os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"] = "0"
+
+        self.state = CUDAGraphState(capture_iteration=capture_iteration)
+
+    def setup(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", stage: str) -> None:
+        """Called when fit, validate, test, predict, or tune begins."""
+        if self.state.capture_iteration < 0:
+            return
+
+        # Hack to avoid CUDA graph issue with AMP, PyTorch Lightning doesn't support
+        # changing autocast arguments for now.
+        # https://github.com/pytorch/pytorch/blob/v1.13.1/torch/cuda/graphs.py#L234
+        torch.autocast.__orig_init__ = torch.autocast.__init__
+        torch.autocast.__init__ = get_amp_autocast_init(self.state)
+
+        # Before full-backward capture, DDP must be constructed in a side-stream context.
+        # We've merged the change that init DDP on side stream to PyTorch Lightning V2,
+        # but not all user defined strategy init DDP on side stream.
+        DistributedDataParallel.__orig_init__ = DistributedDataParallel.__init__
+        DistributedDataParallel.__init__ = get_ddp_init(self.state)
+
+    def teardown(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", stage: str) -> None:
+        """Called when fit, validate, test, predict, or tune ends."""
+        if self.state.capture_iteration < 0:
+            return
+
+        torch.autocast.__init__ = torch.autocast.__orig_init__
+        del torch.autocast.__orig_init__
+
+        DistributedDataParallel.__init__ = DistributedDataParallel.__orig_init__
+        del DistributedDataParallel.__orig_init__
+
+    def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        """Called when fit begins."""
+        if self.state.capture_iteration < 0:
+            return
+
+        if is_param_in_hook_signature(pl_module.training_step, "dataloader_iter", explicit=True):
+            raise Exception(
+                "Found `dataloader_iter` argument in the `training_step`. This is "
+                "not supported by full iteration CUDA graph capturing yet since "
+                "dataloader will be within the CUDA graph capturing range.\n"
+                "Try to change `dataloader_iter` to `batch` and remove "
+                "`next(dataloader_iter)` from `training_step`."
+            )
+
+        # Now that CUDA device has been set, we can init stream and graph now
+        self.state.stream = torch.cuda.Stream()
+        self.state.graph = torch.cuda.CUDAGraph()
+
+    def on_fit_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        """Called when fit ends."""
+        if self.state.capture_iteration < 0:
+            return
+
+    def on_train_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        """Called when the train begins."""
+        if self.state.capture_iteration < 0:
+            return
+
+        # Ensure training dataloader loads data to static buffer
+        dataloader = trainer.fit_loop._combined_loader._iterables
+        assert isinstance(
+            dataloader, torch.utils.data.dataloader.DataLoader
+        ), f"Expect Dataloader type but got {type(dataloader)}"
+        static_loader = StaticBufferLoader(dataloader)
+        _mode = trainer.fit_loop._combined_loader._mode
+        combined_loader = CombinedLoader(static_loader, mode=_mode)
+        trainer.fit_loop.__orig_combined_loader__ = trainer.fit_loop._combined_loader
+        trainer.fit_loop._combined_loader = combined_loader
+        trainer.fit_loop._data_fetcher.setup(trainer.fit_loop._combined_loader)
+        iter(trainer.fit_loop._data_fetcher)
+
+        # Warn if `optimizer.zero_grad()` invoked during graph capturing
+        for optimizer in trainer.optimizers:
+            assert isinstance(optimizer, torch.optim.Optimizer), f"Expect Optimizer type but got {type(optimizer)}"
+            optimizer.__orig_zero_grad__ = optimizer.zero_grad
+            optimizer.zero_grad = MethodType(zero_grad, optimizer)
+
+        # Ensure LR scheduler writes to static buffer
+        # We don't include LR scheduler in the full CUDA graph for now since
+        # its overhead is very small.
+        for config in trainer.lr_scheduler_configs:
+            assert isinstance(
+                config.scheduler, torch.optim.lr_scheduler._LRScheduler
+            ), f"Expect _LRScheduler type but got {type(config.scheduler)}"
+            config.scheduler.__orig_get_lr__ = config.scheduler.get_lr
+            config.scheduler.get_lr = MethodType(get_lr, config.scheduler)
+
+        # Use smart metrics to avoid syncs
+        LightningModule.__orig_to_tensor__ = LightningModule._LightningModule__to_tensor
+        LightningModule._LightningModule__to_tensor = to_tensor
+
+        # Save model outputs to static buffer for PL states reconstruct
+        pl_module.__orig_training_step__ = pl_module.training_step
+        training_step = get_training_step(self.state)
+        pl_module.training_step = MethodType(training_step, pl_module)
+
+        # Capture CUDA graph from model forward propagation to optimizer step
+        pl_module.__orig_optimizer_step__ = pl_module.optimizer_step
+        optimizer_step = get_optimizer_step(self.state)
+        pl_module.optimizer_step = MethodType(optimizer_step, pl_module)
+
+    def on_train_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        """Called when the train ends."""
+        if self.state.capture_iteration < 0:
+            return
+
+        trainer.fit_loop._combined_loader = trainer.fit_loop.__orig_combined_loader__
+        trainer.fit_loop._data_fetcher.setup(trainer.fit_loop._combined_loader)
+        iter(trainer.fit_loop._data_fetcher)
+        del trainer.fit_loop.__orig_combined_loader__
+
+        for optimizer in trainer.optimizers:
+            optimizer.zero_grad = optimizer.__orig_zero_grad__
+            del optimizer.__orig_zero_grad__
+
+        for config in trainer.lr_scheduler_configs:
+            config.scheduler.get_lr = config.scheduler.__orig_get_lr__
+            del config.scheduler.__orig_get_lr__
+
+        LightningModule._LightningModule__to_tensor = LightningModule.__orig_to_tensor__
+        del LightningModule.__orig_to_tensor__
+
+        pl_module.training_step = pl_module.__orig_training_step__
+        del pl_module.__orig_training_step__
+
+        pl_module.optimizer_step = pl_module.__orig_optimizer_step__
+        del pl_module.__orig_optimizer_step__
+
+    def on_train_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        """Called when the train epoch begins."""
+        pass
+
+    def on_train_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        """Called when the train epoch ends.
+
+        To access all batch outputs at the end of the epoch, either:
+
+        1. Implement `training_epoch_end` in the `LightningModule` and access outputs via the module OR
+        2. Cache data across train batch hooks inside the callback implementation to post-process in this hook.
+        """
+        pass
+
+    def on_train_batch_start(
+        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", batch: Any, batch_idx: int
+    ) -> None:
+        """Called when the train batch begins."""
+        pass
+
+    def on_train_batch_end(
+        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", outputs: STEP_OUTPUT, batch: Any, batch_idx: int
+    ) -> None:
+        """Called when the train batch ends.
+
+        Note:
+            The value ``outputs["loss"]`` here will be the normalized value w.r.t ``accumulate_grad_batches`` of the
+            loss returned from ``training_step``.
+        """
+        pass
+
+    def on_save_checkpoint(
+        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", checkpoint: Dict[str, Any]
+    ) -> None:
+        r"""
+        Called when saving a checkpoint to give you a chance to store anything else you might want to save.
+
+        Args:
+            trainer: the current :class:`~lightning.pytorch.trainer.Trainer` instance.
+            pl_module: the current :class:`~lightning.pytorch.core.module.LightningModule` instance.
+            checkpoint: the checkpoint dictionary that will be saved.
+        """
+        # Since we've add bound method to optimizer and lr_scheduler, it can lead to more
+        # CUDA tensors passed to consumer process unexpectedly.
+        if "optimizer_states" in checkpoint:
+            for optimizer_state in checkpoint["optimizer_states"]:
+                for k in list(optimizer_state.keys()):
+                    v = optimizer_state[k]
+                    if isinstance(v, MethodType) and hasattr(v, "__self__"):
+                        del optimizer_state[k]
+        if "lr_schedulers" in checkpoint:
+            for lr_scheduler in checkpoint["lr_schedulers"]:
+                for k in list(lr_scheduler.keys()):
+                    v = lr_scheduler[k]
+                    if isinstance(v, MethodType) and hasattr(v, "__self__"):
+                        del lr_scheduler[k]

nemo/utils/callbacks/dist_ckpt_io.py

@@ -0,0 +1,484 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import shutil
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from time import time
+from typing import Any, Dict, Optional, Union
+
+import lightning.pytorch as pl
+import torch
+from lightning.fabric.plugins import CheckpointIO
+from lightning.fabric.utilities.cloud_io import get_filesystem
+from lightning.fabric.utilities.types import _PATH
+from lightning.pytorch import Callback
+from lightning.pytorch.plugins.io.wrapper import _WrappingCheckpointIO
+
+from nemo.utils import logging
+
+try:
+    from megatron.core import dist_checkpointing
+    from megatron.core.dist_checkpointing.dict_utils import extract_matching_values
+    from megatron.core.dist_checkpointing.mapping import ShardedBase
+    from megatron.core.dist_checkpointing.serialization import (
+        get_default_load_sharded_strategy,
+        get_default_save_sharded_strategy,
+    )
+    from megatron.core.dist_checkpointing.strategies import tensorstore
+    from megatron.core.dist_checkpointing.strategies.async_utils import AsyncCallsQueue, AsyncRequest
+    from megatron.core.dist_checkpointing.strategies.base import SaveShardedStrategy
+    from megatron.core.dist_checkpointing.strategies.fully_parallel import (
+        FullyParallelLoadStrategyWrapper,
+        FullyParallelSaveStrategyWrapper,
+    )
+    from megatron.core.dist_checkpointing.strategies.torch import TorchDistSaveShardedStrategy
+    from megatron.core.dist_checkpointing.validation import StrictHandling
+    from megatron.core.parallel_state import get_data_parallel_group
+
+    HAVE_MEGATRON_CORE = True
+
+except (ImportError, ModuleNotFoundError) as e:
+
+    HAVE_MEGATRON_CORE = False
+    IMPORT_ERROR = (
+        "megatron-core was not found. "
+        "Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+        f" Exact error: {e}"
+    )
+
+
+@contextmanager
+def _debug_time(name: str):
+    """Simple context manager for timing functions/code blocks."""
+    start = time()
+    try:
+        yield
+    finally:
+        logging.debug(f'{name} took {time() - start:.3f}s')
+
+
+class AsyncCompatibleCheckpointIO(CheckpointIO, ABC):
+    """CheckpointIO that can be used together with async saving.
+
+    Differs from the regular CheckpointIO only by the `save_checkpoint`
+    return type. The `save_checkpoint` method itself is synchronous, but returns
+    callbacks that can be performed asynchronously.
+    """
+
+    @abstractmethod
+    def save_checkpoint(
+        self, checkpoint: Dict[str, Any], path: _PATH, storage_options: Optional[Any] = None
+    ) -> 'AsyncRequest':
+        """Interface to implement save_checkpoint and return an AsyncRequest"""
+        raise NotImplementedError
+
+
+class AsyncFinalizableCheckpointIO(_WrappingCheckpointIO):
+    """CheckpointIO wrapper for async checkpoint saving and synchronous finalization.
+
+    Runs main part of the checkpoint save in a separate process (not thread as the PTL
+    AsyncCheckpointIO does). Allows to perform a (synchronous) finalization
+    function after all ranks finish checkpoint saving.
+
+    NOTE: for correctness, this plugin must be used together with the
+    AsyncFinalizerCallback callback which performs the finalization checks.
+
+    Args:
+        checkpoint_io (CheckpointIO): wrapped checkpoint_io object. Must be
+            of type AsyncCompatibleCheckpointIO.
+    Requires the underlying checkpoint_io.save_checkpoint to return save_fn, save_args, finalize_fn.
+    """
+
+    def __init__(self, checkpoint_io: AsyncCompatibleCheckpointIO) -> None:
+        if not HAVE_MEGATRON_CORE:
+            raise ImportError(IMPORT_ERROR)
+        if not isinstance(checkpoint_io, AsyncCompatibleCheckpointIO):
+            raise ValueError(f'Incompatible wrapped checkpoint_io type: {type(checkpoint_io)}')
+
+        super().__init__(checkpoint_io)
+        self.async_calls_queue = AsyncCallsQueue()
+
+    def save_checkpoint(
+        self,
+        checkpoint: Dict[str, Any],
+        path: _PATH,
+        storage_options: Optional[Any] = None,
+    ) -> None:
+        """Executes async request returned from the underlying checkpoint_io asynchronously.
+
+        Requires the underlying checkpoint_io.save_checkpoint to return an AsyncRequest.
+        It is then applied with `self.async_calls_queue` asynchronously.
+
+        Args:
+            checkpoint (Dict[str, Any]): checkpoint to save. Passed to underlying
+                checkpoint_io without modifications.
+            path (_PATH): path to save the checkpoint. Passed to underlying
+                checkpoint_io without modifications.
+            storage_options (Any, optional): storage control modifiers. This class
+                consumed the `finalize_fn` parameter (if any), which is expected to be
+                a callback and is appended to async finalization functions.
+
+        Applies underlying checkpoint_io finalize callback first, then the external one (postfix order).
+        """
+        external_finalize_fn = (storage_options or {}).pop('finalize_fn', None)
+        assert isinstance(self.checkpoint_io, AsyncCompatibleCheckpointIO), type(self.checkpoint_io)
+        async_request = self.checkpoint_io.save_checkpoint(checkpoint, path, storage_options)
+        if external_finalize_fn is not None:
+            async_request.add_finalize_fn(external_finalize_fn)
+        call_idx = self.async_calls_queue.schedule_async_request(async_request)
+        logging.debug(f'Scheduled an async call #{call_idx}')
+
+    @_debug_time('AsyncFinalizableCheckpointIO.maybe_finalize_save_checkpoint')
+    def maybe_finalize_save_checkpoint(self, blocking: bool = False):
+        """Performs checkpoint finalization (if possible).
+
+        Args:
+            blocking (bool, optional): if True, waits until all async saves are
+                completed. Otherwise, finalizes only those async calls which are
+                already done on all ranks. Defaults to False.
+        """
+        if self.async_calls_queue.get_num_unfinalized_calls() == 0:
+            return False
+
+        start_time = time()
+        call_idx_finalized = self.async_calls_queue.maybe_finalize_async_calls(blocking)
+        if call_idx_finalized:
+            logging.debug(f'Finalized async calls: {[f"#{idx}" for idx in call_idx_finalized]}')
+        end_time = time()
+        logging.info(f"Async finalization time took {end_time - start_time:.3f} s")
+        return len(call_idx_finalized) > 0
+
+    def teardown(self) -> None:
+        """Warns if there are any pending checkpoint saves."""
+        super().teardown()
+        if self.async_calls_queue.get_num_unfinalized_calls() > 0:
+            # Can't do finalization now because some ranks might be lost
+            logging.warning('Some async checkpoint saves might be not finalized properly.')
+
+
+class AsyncFinalizerCallback(Callback):
+    """Callback which finalizes async saves initiated by the AsyncFinalizableCheckpointIO.
+
+    Tries to perform non-blocking finalization on train_batch_end and train_epoch_end.
+    On train_end performs a blocking finalization of all pending checkpoints.
+    """
+
+    def on_train_batch_end(self, trainer: "pl.Trainer", *args, **kwargs) -> None:
+        """Override hook to finalize pending checkpoint(s) if they exist."""
+        self._get_checkpoint_io(trainer).maybe_finalize_save_checkpoint(blocking=False)
+
+    def on_train_epoch_end(self, trainer: "pl.Trainer", *args, **kwargs) -> None:
+        """Override hook to finalize pending checkpoint(s) if they exist."""
+        self._get_checkpoint_io(trainer).maybe_finalize_save_checkpoint(blocking=False)
+
+    def on_train_end(self, trainer: "pl.Trainer", *args, **kwargs) -> None:
+        """Override hook to finalize pending checkpoint(s) if they exist."""
+        checkpoint_io = self._get_checkpoint_io(trainer)
+        if checkpoint_io.async_calls_queue.get_num_unfinalized_calls() > 0:
+            logging.info('Pending async checkpoint saves. Finalizing them synchronously now')
+        self._get_checkpoint_io(trainer).maybe_finalize_save_checkpoint(blocking=True)
+
+    def _get_checkpoint_io(self, trainer) -> AsyncFinalizableCheckpointIO:
+        checkpoint_io = trainer.strategy.checkpoint_io
+        if not isinstance(checkpoint_io, AsyncFinalizableCheckpointIO):
+            raise ValueError(
+                f'Async finalizer requires an async compatible CheckpointIO, got: {checkpoint_io.__class__}'
+            )
+        return checkpoint_io
+
+
+class DistributedCheckpointIO(AsyncCompatibleCheckpointIO):
+    """CheckpointIO for a distributed checkpoint format.
+
+    Args:
+        save_ckpt_format (str): Distributed checkpoint format to use for checkpoint saving.
+        load_directly_on_device (bool, optional): if True, loads the weights directly
+            on GPU. Has effect only for `zarr` based checkpoints (PyT Distributed
+            always loads on device). Defaults to True.
+        load_strictness (StrictHandling, optional): defines loading strictness.
+            If not None, overwrites the `strict` flag passed to `load_checkpoint`.
+            Defaults to None.
+        async_save (bool): whether to save asynchronously. Should be set to True if
+            this class will be wrapped with AsyncFinalizableCheckpointIO.
+        torch_dist_multiproc (int, optional): number of extra processes per rank
+            used during ckpt save with PyTorch distributed format. Defaults, to None
+            which means using an MCore default (2).
+        parallel_save (bool): parallelizes the save across ranks. Defaults to True
+        parallel_load (bool): parallelizes the load across ranks (followed by params all gather).
+            Defaults to False due to some extra memory usage requirement.
+    """
+
+    def __init__(
+        self,
+        save_ckpt_format: str,
+        load_directly_on_device: bool = True,
+        load_strictness: Optional['StrictHandling'] = None,
+        async_save: bool = False,
+        torch_dist_multiproc: Optional[int] = None,
+        assume_constant_structure: bool = False,
+        parallel_save: bool = False,
+        parallel_save_within_dp: bool = False,
+        parallel_load: bool = False,
+    ):
+        super().__init__()
+        if not HAVE_MEGATRON_CORE:
+            raise ImportError(IMPORT_ERROR)
+
+        self.save_ckpt_format = save_ckpt_format
+        self.load_directly_on_device = load_directly_on_device
+        self.load_strictness = load_strictness
+        self.async_save = async_save
+        self.torch_dist_multiproc = torch_dist_multiproc
+        self.assume_constant_structure = assume_constant_structure
+        self.parallel_save = parallel_save
+        self.parallel_save_within_dp = parallel_save_within_dp
+        self.parallel_load = parallel_load
+
+        self._save_sharded_strategy = None
+        self.validated_consistency = False
+
+    @classmethod
+    def from_config(cls, model_cfg: dict, async_save: bool = False):
+        """Instantiates a DistributedCheckpointIO from a config dict.
+
+        Args:
+            model_cfg (dict): model config dict. Most of the configuration
+                is extracted from this config.
+            async_save (bool, optional): async_save flag is not part of the model config,
+                it should be provided separately. Defaults to False.
+        """
+        return cls(
+            save_ckpt_format=model_cfg.get('dist_ckpt_format', 'torch_dist'),
+            load_directly_on_device=model_cfg.get('dist_ckpt_load_on_device', True),
+            load_strictness=model_cfg.get('dist_ckpt_load_strictness', None),
+            async_save=async_save,
+            torch_dist_multiproc=model_cfg.get('dist_ckpt_torch_dist_multiproc', None),
+            parallel_save=model_cfg.get('dist_ckpt_parallel_save', False),
+            parallel_save_within_dp=model_cfg.get('dist_ckpt_parallel_save_within_dp', False),
+            parallel_load=model_cfg.get('dist_ckpt_parallel_load', False),
+        )
+
+    @_debug_time('DistributedCheckpointIO.save_checkpoint')
+    def save_checkpoint(
+        self, checkpoint: Dict[str, Any], path: _PATH, storage_options: Optional[Any] = None
+    ) -> Optional['AsyncRequest']:
+        """Saves a distributed checkpoint. Creates the checkpoint root directory if doesn't exist.
+
+        Args:
+            checkpoint (Dict[str, Any]): sharded state dict to save
+            path (_PATH): checkpoint directory
+            storage_options (Any, optional): Optional parameters when saving the checkpoint
+        """
+        fs = get_filesystem(path)
+        fs.makedirs(path, exist_ok=True)
+
+        validate_sharding_integrity = not (self.validated_consistency and self.assume_constant_structure)
+        self.validated_consistency = True
+
+        rank = torch.distributed.get_rank()
+        iteration = _get_iteration_from_checkpoint(checkpoint)
+        start_time = time()
+        async_save_request = dist_checkpointing.save(
+            sharded_state_dict=checkpoint,
+            checkpoint_dir=path,
+            sharded_strategy=self.save_sharded_strategy,
+            validate_access_integrity=validate_sharding_integrity,
+            async_sharded_save=self.async_save,
+        )
+        end_time = time()
+        log_parts = (
+            "Global Checkpoint Save",
+            f"Rank: {rank}",
+            f"Iteration: {iteration}" if iteration is not None else None,
+            f"Start time: {start_time:.3f}s",
+            f"Save duration: {end_time - start_time:.3f}s",
+        )
+        log_message = " : ".join(part for part in log_parts if part is not None)
+        logging.info(log_message)
+
+        def iter_finalize_fn():
+            logging.info(f'Successfully saved checkpoint from iteration {int(iteration):7d} to {path}')
+
+        if self.async_save:
+            assert async_save_request is not None
+            async_save_request.add_finalize_fn(iter_finalize_fn)
+
+        return async_save_request
+
+    @_debug_time('DistributedCheckpointIO.load_checkpoint')
+    def load_checkpoint(
+        self,
+        path: _PATH,
+        map_location: Optional[Any] = None,
+        sharded_state_dict: Dict[str, Any] = None,
+        strict: Union[None, bool, 'StrictHandling'] = None,
+        validate_access_integrity: Optional[bool] = True,
+    ) -> Dict[str, Any]:
+        """Loads a distributed checkpoint.
+
+        Args:
+            path (_PATH): checkpoint directory
+            map_location (Any, optional): required to be None in this implementation
+            sharded_state_dict (Dict[str, Any], optional): state dict which
+                defines the loading procedure for the distributed checkpoint.
+                Defaults to None to comply with the CheckpointIO interface,
+                but it's a required argument.
+            strict (bool, StrictHandling, optional): adjust load strictness. bool value
+                is translated to StrictHandling instance. Gets overwritten by
+                `self.load_strictness`. Defaults to None. If `self.load_strictness`
+                is also None, strict becomes StrictHandling.ASSUME_OK_UNEXPECTED.
+
+        Returns:
+            Dist[str, Any]: loaded checkpoint.
+        """
+        if sharded_state_dict is None:
+            raise ValueError('DistributedCheckpointIO requires passing sharded_state_dict argument to load_checkpoint')
+        if map_location is not None:
+            raise ValueError('DistributedCheckpointIO doesnt handle map_location argument')
+
+        if self.save_ckpt_format == 'zarr' and self.load_directly_on_device:
+            sharded_strategy = tensorstore.TensorStoreLoadShardedStrategy(load_directly_on_device=True)
+        else:
+            sharded_strategy = None
+
+        if self.parallel_load:
+            if sharded_strategy is None:
+                sharded_strategy = get_default_load_sharded_strategy(path)
+            sharded_strategy = FullyParallelLoadStrategyWrapper(
+                sharded_strategy, get_data_parallel_group(with_context_parallel=True)
+            )
+
+        if sharded_strategy is not None:
+            logging.info(f'Using {sharded_strategy} dist-ckpt load strategy.')
+
+        if isinstance(strict, bool):
+            # For backward-compatibility reasons and a bug in MCore (strict check not applied to factories)
+            # we must apply a simple strict check here.
+            if not strict:
+                sharded_state_dict = self.adjust_non_strict_load(path, sharded_state_dict)
+            strict = StrictHandling.ASSUME_OK_UNEXPECTED if strict else StrictHandling.LOG_ALL
+        if self.load_strictness is not None:
+            # Overwrites function argument
+            strict = self.load_strictness
+        if strict is None:
+            # Default behavior
+            strict = StrictHandling.ASSUME_OK_UNEXPECTED
+
+        logging.debug(f'Dist ckpt load strictness: {strict}')
+
+        start_time = time()
+        ret = dist_checkpointing.load(
+            sharded_state_dict=sharded_state_dict,
+            checkpoint_dir=path,
+            sharded_strategy=sharded_strategy,
+            validate_access_integrity=validate_access_integrity,
+            strict=strict,
+        )
+        end_time = time()
+        duration = end_time - start_time
+        logging.info(
+            "Global Checkpoint Load : "
+            f"Rank : {torch.distributed.get_rank()} : "
+            f"Start time : {start_time:.3f}s : "
+            f"Time spent in load_checkpoint: {duration:.3f}s"
+        )
+        return ret
+
+    def adjust_non_strict_load(self, path: _PATH, sharded_state_dict: Dict[str, Any]):
+        """Remove unexpected keys from being loaded into the state dict."""
+        ckpt_sharded_metadata = dist_checkpointing.load_tensors_metadata(path)
+        loaded_keys = []
+        unexpected_keys = []
+
+        def should_remove_missing_sharded_base(x: Any):
+            if isinstance(x, ShardedBase):
+                if x.key in ckpt_sharded_metadata:
+                    loaded_keys.append(x.key)
+                    return False
+                else:
+                    unexpected_keys.append(x.key)
+                    return True
+            return False
+
+        _, sharded_state_dict = extract_matching_values(sharded_state_dict, should_remove_missing_sharded_base)
+        logging.info(f'The following keys are not in the checkpoint and will not be loaded: {unexpected_keys}')
+
+        # TODO: compute missing_keys by:
+        #  1. all_gather_object of loaded_keys
+        #  2. missing_keys = ckpt_sharded_metadata.keys() - loaded_keys
+        return sharded_state_dict
+
+    @_debug_time('DistributedCheckpointIO.remove_checkpoint')
+    def remove_checkpoint(self, path: _PATH) -> None:
+        """Remove a distributed checkpoint.
+
+        Due to potentially large number of files, the implementation remove the whole directory at once.
+        """
+        shutil.rmtree(path, ignore_errors=True)
+
+    @property
+    def save_sharded_strategy(self) -> 'SaveShardedStrategy':
+        """Conditionally initialize and get the sharded strategy to use for saving."""
+        if self._save_sharded_strategy is None:
+            self._save_sharded_strategy = self._determine_dist_ckpt_save_strategy()
+        return self._save_sharded_strategy
+
+    def _determine_dist_ckpt_save_strategy(self):
+        """Determine the saving strategy based on constructor args.
+
+        Relies on the default MCore strategy unless extra PyT Distributed format arguments
+        are passed in config or in case of a fully parallel save in which case
+        a parallelization wrapper is applied.
+        """
+        if self.save_ckpt_format == 'zarr':
+            logging.warning(
+                '`zarr` distributed checkpoint backend is deprecated.'
+                ' Distributed optimizer checkpoint saving might be extremely slow.'
+                ' Please switch to PyTorch Distributed format (model.dist_ckpt_format=torch_dist).'
+            )
+
+        if self.async_save and self.save_ckpt_format != 'torch_dist':
+            raise ValueError('Async dist-ckpt save supported only for torch_dist format')
+
+        torch_dist_kwargs = {} if self.torch_dist_multiproc is None else dict(thread_count=self.torch_dist_multiproc)
+        if self.save_ckpt_format == 'torch_dist' and torch_dist_kwargs:
+            save_strategy = TorchDistSaveShardedStrategy(self.save_ckpt_format, 1, **torch_dist_kwargs)
+        else:
+            save_strategy = get_default_save_sharded_strategy(self.save_ckpt_format, 1)
+
+        # MCore v0.8 introduces `use_cached_ckpt_structure` attribute
+        if hasattr(save_strategy, 'use_cached_ckpt_structure'):
+            save_strategy.use_cached_ckpt_structure = self.assume_constant_structure
+
+        if self.parallel_save:
+            parallelization_group = (
+                get_data_parallel_group(with_context_parallel=True) if self.parallel_save_within_dp else None
+            )
+            save_strategy = FullyParallelSaveStrategyWrapper(
+                save_strategy, parallelization_group, self.assume_constant_structure
+            )
+
+        logging.info(f'Using {save_strategy} dist-ckpt save strategy.')
+        return save_strategy
+
+
+def _get_iteration_from_checkpoint(checkpoint: Dict[str, Any]) -> Optional[int]:
+    return (
+        checkpoint.get("loops", {})
+        .get("fit_loop", {})
+        .get("epoch_loop.batch_progress", {})
+        .get("total", {})
+        .get("completed", None)
+    )

nemo/utils/callbacks/nemo_model_checkpoint.py

@@ -0,0 +1,740 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import re
+import shutil
+import time
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Union
+
+import torch
+from _weakref import proxy
+from lightning.fabric.utilities.cloud_io import get_filesystem
+from lightning.pytorch.callbacks.model_checkpoint import ModelCheckpoint, _is_local_file_protocol
+from lightning.pytorch.trainer import call
+from lightning.pytorch.utilities import rank_zero_info
+
+from nemo.collections.common.callbacks import EMA
+from nemo.utils import logging
+from nemo.utils.app_state import AppState
+from nemo.utils.callbacks.dist_ckpt_io import AsyncFinalizableCheckpointIO
+from nemo.utils.get_rank import is_global_rank_zero
+from nemo.utils.model_utils import ckpt_to_dir, inject_model_parallel_rank, uninject_model_parallel_rank
+from nemo.utils.msc_utils import import_multistorageclient, is_multistorageclient_url
+
+
+class NeMoModelCheckpoint(ModelCheckpoint):
+    """Light wrapper around Lightning's ModelCheckpoint to force a saved checkpoint on train_end.
+    Extends Lightning's on_save_checkpoint func to save the .nemo file. Saves the .nemo file based
+    on the best checkpoint saved (according to the monitor value).
+    Also contains func to save the EMA copy of the model.
+    """
+
+    UNFINISHED_CHECKPOINT_SUFFIX = "-unfinished"
+
+    def __init__(
+        self,
+        always_save_nemo: bool = False,
+        save_nemo_on_train_end: bool = True,
+        save_best_model: bool = False,
+        postfix: str = ".nemo",
+        n_resume: bool = False,
+        model_parallel_size: int = None,
+        async_save: bool = False,  # controls only finalize callbacks
+        save_last_n_optim_states: int = -1,
+        **kwargs,
+    ):
+        # Parse and store "extended" parameters: save_best model and postfix.
+        self.always_save_nemo = always_save_nemo
+        self.save_nemo_on_train_end = save_nemo_on_train_end
+        self.save_best_model = save_best_model
+        self.save_last_n_optim_states = save_last_n_optim_states
+        if self.save_best_model and not self.save_nemo_on_train_end:
+            logging.warning(
+                (
+                    "Found save_best_model is True and save_nemo_on_train_end is False. "
+                    "Set save_nemo_on_train_end to True to automatically save the best model."
+                )
+            )
+        self.postfix = postfix
+        self.previous_best_path = ""
+        self.model_parallel_size = model_parallel_size
+        self.async_save = async_save
+        self.async_finalize_cb = None
+        # Checkpoints which removal is deferred until async save is done.
+        # Each element of `deferred_ckpts_to_remove` is a growing list
+        # that `self._remove_checkpoint` adds to. Once `self._save_checkpoint`
+        # is called, the last element is frozen and a new element is added.
+        self.deferred_ckpts_to_remove: List[List[str]] = []
+
+        # `prefix` is deprecated
+        if 'prefix' in kwargs:
+            self.prefix = kwargs.pop('prefix')
+        else:
+            self.prefix = ""
+
+        # Call the parent class constructor with the remaining kwargs.
+        super().__init__(**kwargs)
+
+        if self.save_top_k != -1 and n_resume:
+            logging.debug("Checking previous runs")
+            self.nemo_topk_check_previous_run()
+
+    def nemo_topk_check_previous_run(self):
+        """
+        Check if there are previous runs.
+        """
+        try:
+            self.best_k_models
+            self.kth_best_model_path
+            self.best_model_score
+            self.best_model_path
+        except AttributeError:
+            raise AttributeError("Lightning's ModelCheckpoint was updated. NeMoModelCheckpoint will need an update.")
+        self.best_k_models = {}
+        self.kth_best_model_path = ""
+        self.best_model_score = None
+        self.best_model_path = ""
+
+        checkpoints = list(path for path in self._saved_checkpoint_paths if not self._is_ema_filepath(path))
+        for checkpoint in checkpoints:
+            if 'mp_rank' in str(checkpoint) or 'tp_rank' in str(checkpoint):
+                checkpoint = uninject_model_parallel_rank(checkpoint)
+            checkpoint = str(checkpoint)
+            # second case is for distributed checkpoints, since they are a directory there's no extension
+            if checkpoint[-10:] == '-last.ckpt' or checkpoint[-5:] == '-last':
+                continue
+            index = checkpoint.find(self.monitor) + len(self.monitor) + 1  # Find monitor in str + 1 for '='
+            if index != len(self.monitor):
+                match = re.search('[A-z]', checkpoint[index:])
+                if match:
+                    value = checkpoint[index : index + match.start() - 1]  # -1 due to separator hypen
+                    self.best_k_models[checkpoint] = float(value)
+        if len(self.best_k_models) < 1:
+            return  # No saved checkpoints yet
+
+        _reverse = False if self.mode == "min" else True
+
+        best_k_models = sorted(self.best_k_models, key=self.best_k_models.get, reverse=_reverse)
+
+        # This section should be ok as rank zero will delete all excess checkpoints, since all other ranks are
+        # instantiated after rank zero. models_to_delete should be 0 for all other ranks.
+        if self.model_parallel_size is not None:
+            # check for distributed checkpoint
+            if checkpoints[0].is_dir():
+                models_to_delete = len(best_k_models) - self.save_top_k
+            else:
+                models_to_delete = len(best_k_models) - self.model_parallel_size * self.save_top_k
+        else:
+            models_to_delete = len(best_k_models) - self.save_top_k
+
+        models_to_delete = max(0, models_to_delete)
+        logging.debug(f'Number of models to delete: {models_to_delete}')
+
+        # If EMA enabled, delete the additional EMA weights
+        ema_enabled = self._has_ema_ckpts(self._saved_checkpoint_paths)
+
+        for _ in range(models_to_delete):
+            model = best_k_models.pop(-1)
+            self.best_k_models.pop(model)
+            self._del_model_without_trainer(model)
+            if ema_enabled and self._fs.exists(self._ema_format_filepath(model)):
+                self._del_model_without_trainer(self._ema_format_filepath(model))
+            logging.debug(f"Removed checkpoint: {model}")
+
+        self.kth_best_model_path = best_k_models[-1]
+        self.best_model_path = best_k_models[0]
+        self.best_model_score = self.best_k_models[self.best_model_path]
+
+    def _remove_invalid_entries_from_topk(self):
+        # Removes invalid (incomplete or not existing) checkpoints from topk checkpoints.
+        # This might be needed if the checkpointing was abruptly terminated.
+        def __is_ckpt_ok(ckpt_path: str) -> bool:
+            exists = (
+                os.path.isfile(ckpt_path)
+                or os.path.isfile(inject_model_parallel_rank(ckpt_path))
+                or os.path.isdir(ckpt_path.removesuffix('.ckpt'))
+            )
+            return exists and not self.is_checkpoint_unfinished(ckpt_path)
+
+        self.best_k_models = {k: v for k, v in self.best_k_models.items() if __is_ckpt_ok(k)}
+        if len(self.best_k_models) > 0:
+            reverse_arr = self.mode != "min"
+            best_k_models_arr = sorted(self.best_k_models, key=self.best_k_models.get, reverse=reverse_arr)
+            self.kth_best_model_path = best_k_models_arr[-1]
+            self.kth_value = self.best_k_models[self.kth_best_model_path]
+            self.best_model_path = best_k_models_arr[0]
+            self.best_model_score = self.best_k_models[self.best_model_path]
+        else:
+            self.kth_best_model_path = ""
+            self.kth_value = None
+            self.best_model_path = ""
+            self.best_model_score = None
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        """
+        Load the state dict.
+        """
+        super().load_state_dict(state_dict)
+        self._remove_invalid_entries_from_topk()
+
+    def setup(self, trainer, pl_module, stage: str) -> None:
+        """
+        Setup the checkpoint.
+        """
+        if is_global_rank_zero():
+            logging.debug("Removing unfinished checkpoints if any...")
+            NeMoModelCheckpoint._remove_unfinished_checkpoints(self.dirpath)
+        # Ensure that all ranks continue with unfinished checkpoints removed
+        if torch.distributed.is_initialized():
+            torch.distributed.barrier()
+        super().setup(trainer, pl_module, stage)
+        # When using S3 checkpointing, only Rank 0 has the checkpoint and model path set in exp_manager.
+        # Sync the values across all ranks to ensure consistency.
+        path = trainer.strategy.broadcast(trainer.ckpt_path)
+        trainer.ckpt_path = path
+
+        self.last_model_path = trainer.strategy.broadcast(self.last_model_path)
+
+    def on_save_checkpoint(self, trainer, pl_module, checkpoint):
+        """
+        Save the checkpoint.
+        """
+        output = super().on_save_checkpoint(trainer, pl_module, checkpoint)
+        if not self.always_save_nemo:
+            return output
+        # Load the best model and then re-save it
+        app_state = AppState()
+        if app_state.model_parallel_size is not None and app_state.model_parallel_size > 1:
+            logging.warning('always_save_nemo will slow down training for model_parallel > 1.')
+        # since we are creating tarfile artifacts we need to update .nemo path
+        app_state.model_restore_path = self._format_nemo_checkpoint_name()
+        if app_state.model_parallel_size is not None and app_state.model_parallel_size > 1:
+            maybe_injected_best_model_path = inject_model_parallel_rank(self.best_model_path)
+        else:
+            maybe_injected_best_model_path = self.best_model_path
+
+        if self.save_best_model:
+            if not os.path.exists(maybe_injected_best_model_path):
+                return
+
+            if self.best_model_path == self.previous_best_path:
+                logging.debug('Best model has not changed, skipping save.')
+                return output
+
+            self.previous_best_path = self.best_model_path
+            if torch.distributed.is_initialized():
+                torch.distributed.barrier()
+            backup_path = self._backup_existing_nemo_ckpt(trainer)
+            pl_module.save_to(save_path=app_state.model_restore_path)
+            logging.info(f"New best .nemo model saved to: {app_state.model_restore_path}")
+        else:
+            if torch.distributed.is_initialized():
+                torch.distributed.barrier()
+            backup_path = self._backup_existing_nemo_ckpt(trainer)
+            pl_module.save_to(save_path=app_state.model_restore_path)
+            logging.info(f"New .nemo model saved to: {app_state.model_restore_path}")
+        if backup_path is not None and is_global_rank_zero():
+            logging.info(f'Removing old .nemo backup {backup_path}')
+            get_filesystem(backup_path).rm(backup_path)
+        return output
+
+    def on_train_end(self, trainer, pl_module):
+        """
+        Save the checkpoint on train end.
+        """
+        if trainer.fast_dev_run:
+            return None
+
+        # check if we need to save a last checkpoint manually as validation isn't always run based on the interval
+        if self.save_last and trainer.val_check_interval != 0:
+            should_save_last_checkpoint = False
+            if isinstance(trainer.val_check_interval, float) and trainer.val_check_interval % trainer.global_step != 0:
+                should_save_last_checkpoint = True
+            if isinstance(trainer.val_check_interval, int) and trainer.global_step % trainer.val_check_interval != 0:
+                should_save_last_checkpoint = True
+            if should_save_last_checkpoint:
+                monitor_candidates = self._monitor_candidates(trainer)
+                if self.last_model_path == self.format_checkpoint_name(monitor_candidates, self.CHECKPOINT_NAME_LAST):
+                    logging.debug(f'Last checkpoint {self.last_model_path} already saved')
+                else:
+                    super()._save_last_checkpoint(trainer, monitor_candidates)
+        # Call parent on_train_end() to save the -last checkpoint
+        super().on_train_end(trainer, pl_module)
+
+        # Load the best model and then re-save it
+        if self.save_best_model:
+            # wait for all processes
+            trainer.strategy.barrier("SaveBestCheckpointConnector.resume_end")
+            if self.best_model_path == "":
+                logging.warning(
+                    f"{self} was told to save the best checkpoint at the end of training, but no saved checkpoints "
+                    "were found. Saving latest model instead."
+                )
+            else:
+                if os.path.isdir(self.best_model_path.split('.ckpt')[0]):
+                    self.best_model_path = self.best_model_path.split('.ckpt')[0]
+                self.best_model_path = trainer.strategy.broadcast(self.best_model_path)
+                trainer._checkpoint_connector.restore(self.best_model_path)
+
+        if self.save_nemo_on_train_end:
+            save_to = getattr(pl_module, "save_to", None)
+            if not callable(save_to):
+                logging.warning(
+                    f"{type(pl_module).__name__} does not implement save_to(); "
+                    "skipping automatic .nemo export at train end."
+                )
+                return
+
+            backup_path = self._backup_existing_nemo_ckpt(trainer)
+            save_to(save_path=self._format_nemo_checkpoint_name())
+            if backup_path is not None and is_global_rank_zero():
+                logging.info(f'Removing old .nemo backup {backup_path}')
+                get_filesystem(backup_path).rm(backup_path)
+
+    def _backup_existing_nemo_ckpt(self, trainer) -> Optional[str]:
+        """Search for an available name with version infix and rename existing checkpoint.
+
+        NOTE: this behavior is slightly different from regular checkpoints.
+        PTL creates new regular checkpoint with the first available name.
+        Here, for backward compatibility, we create .nemo checkpoint as before
+        and create a backup under the first available name.
+
+        Args:
+            trainer (Trainer): trainer instance.
+
+        Returns:
+            Path to the backup checkpoint or None, if no backup was created
+        """
+        base_path = self._format_nemo_checkpoint_name()
+        available_path = base_path
+        if self._enable_version_counter:
+            version_cnt = self.STARTING_VERSION
+            while self.file_exists(available_path, trainer, check_dist_ckpt=False):
+                available_path = self._format_nemo_checkpoint_name(version_cnt)
+                version_cnt += 1
+        if available_path == base_path:
+            # no existing ckpt, no need to backup
+            return None
+        if trainer.is_global_zero:
+            logging.info(f'{base_path} already exists, moving existing checkpoint to {available_path}')
+            if is_multistorageclient_url(base_path):
+                # TODO: multistorageclient doesn't have "rename" function, therefore no-op but we should
+                # refactor this once multistorageclient have rename function supported.
+                pass
+            else:
+                shutil.move(base_path, available_path)
+        trainer.strategy.barrier()
+        return available_path
+
+    def _format_nemo_checkpoint_name(self, ver: Optional[int] = None) -> str:
+        version_infix = '' if ver is None else f'{self.CHECKPOINT_JOIN_CHAR}v{ver}'
+        if is_multistorageclient_url(self.dirpath):
+            return f"{self.dirpath}/{self.prefix + version_infix + self.postfix}"
+        return os.path.abspath(
+            os.path.expanduser(os.path.join(self.dirpath, self.prefix + version_infix + self.postfix))
+        )
+
+    def _del_model_without_trainer(self, filepath: str) -> None:
+
+        filepath = Path(filepath)
+
+        # check if filepath is a distributed a checkpoint
+        if ckpt_to_dir(filepath).is_dir():
+            if is_global_rank_zero():
+                try:
+                    dist_ckpt = ckpt_to_dir(filepath)
+                    shutil.rmtree(dist_ckpt, ignore_errors=True)
+                    logging.info(f"Removed distributed checkpoint: {dist_ckpt}")
+                except:
+                    logging.info(f"Tried to remove distributed checkpoint: {dist_ckpt} but failed.")
+
+        else:
+            app_state = AppState()
+
+            # legacy model parallel checkpoint
+            if app_state.model_parallel_size is not None and app_state.model_parallel_size > 1:
+                # filepath needs to be updated to include mp_rank
+                filepath = inject_model_parallel_rank(filepath)
+
+            # each model parallel rank needs to remove its model
+            if is_global_rank_zero() or (
+                app_state.model_parallel_size is not None and app_state.data_parallel_rank == 0
+            ):
+                try:
+                    self._fs.rm(filepath)
+                    logging.info(f"Removed checkpoint: {filepath}")
+                except:
+                    logging.info(f"Tried to remove checkpoint: {filepath} but failed.")
+
+    def _ema_callback(self, trainer: 'lightning.pytorch.Trainer') -> Optional[EMA]:  # noqa: F821
+        ema_callback = None
+        for callback in trainer.callbacks:
+            if isinstance(callback, EMA):
+                ema_callback = callback
+        return ema_callback
+
+    def _drop_optimizer_states(self, trainer, filepath: Union[str, Path], storage_options: Optional[Any]) -> None:
+        # Get list of saved checkpoints
+        checkpoints = self._get_checkpoints_list(filepath)
+        suffix = "-no-optim"
+
+        # Drop optimizer states
+        checkpoint_index = len(checkpoints) - self.save_last_n_optim_states - 1
+        if len(checkpoints) > self.save_last_n_optim_states:
+            checkpoint_path = checkpoints[checkpoint_index]
+
+            logging.info(f"Loading '{checkpoint_path}' checkpoint to drop optimizer states...")
+            checkpoint = trainer.strategy.load_checkpoint(checkpoint_path=checkpoint_path, load_optimizer_states=False)
+
+            # Load related state dict
+            self._load_current_state_dict(trainer, checkpoint)
+
+            # Save the checkpoint without optimizer states
+            if storage_options is None:
+                storage_options = dict(include_optimizer=False)
+            else:
+                storage_options["include_optimizer"] = False
+
+            trainer.save_checkpoint(
+                f"{checkpoint_path}{suffix}.ckpt", self.save_weights_only, storage_options=storage_options
+            )
+
+            # Remove the checkpoint version with optimizer states
+            if is_global_rank_zero():
+                trainer.strategy.remove_checkpoint(checkpoint_path)
+                shutil.move(f"{checkpoint_path}{suffix}", checkpoint_path)
+
+            if torch.distributed.is_initialized():
+                torch.distributed.barrier()
+
+            # Load the correct state_dict for current checkpoint.
+            # Temporary solution.
+            checkpoint = trainer.strategy.load_checkpoint(
+                checkpoint_path=ckpt_to_dir(filepath), load_optimizer_states=False
+            )
+            self._load_current_state_dict(trainer, checkpoint)
+
+            logging.info(f"Successfully dropped optimizer states for '{checkpoint_path}' checkpoint.")
+
+    def _get_checkpoints_list(self, filepath: Union[str, Path]) -> List[str]:
+        # Get a checkpoints directory
+        checkpoints_dir = os.path.dirname(filepath)
+
+        # Get a list of saved checkpoints
+        checkpoints = [
+            d
+            for d in os.listdir(checkpoints_dir)
+            if os.path.isdir(os.path.join(checkpoints_dir, d)) and '-last' not in d
+        ]
+        checkpoints = sorted(checkpoints, key=lambda x: int(x.split('-step=')[1].split('-')[0]))
+        checkpoints = [os.path.join(checkpoints_dir, checkpoint) for checkpoint in checkpoints]
+
+        return checkpoints
+
+    def _load_current_state_dict(self, trainer, checkpoint) -> None:
+        # Temporary solution for loading the correct state dict
+        # when dropping optimizer states "on the fly" during training.
+
+        # TODO @dimapihtar @mikolajblaz: provide a more elegant solution at the mcore level.
+
+        call._call_lightning_module_hook(trainer, "on_load_checkpoint", checkpoint)
+
+        # Load model state_dict
+        trainer.strategy.load_model_state_dict(
+            checkpoint,
+            strict=trainer.lightning_module.strict_loading,
+        )
+
+    @staticmethod
+    def format_checkpoint_unfinished_marker_path(checkpoint_path: Union[Path, str]) -> Path:
+        """Format the path to the unfinished checkpoint marker file.
+
+        If the marker file exists, corresponding checkpoint is considered unfinished/incomplete.
+        NOTE: Marker path for the EMA checkpoint part is the same as for the original checkpoint.
+
+        Args:
+            checkpoint_path: Path to the checkpoint file or dir.
+              Does not need to exist.
+
+        Returns:
+            Path to the unfinished checkpoint marker file.
+        """
+        marker_filepath = str(uninject_model_parallel_rank(checkpoint_path))
+        marker_filepath = marker_filepath.removesuffix(".nemo")
+        marker_filepath = marker_filepath.removesuffix(".ckpt")
+        marker_filepath = marker_filepath.removesuffix("-EMA")
+        return Path(marker_filepath + NeMoModelCheckpoint.UNFINISHED_CHECKPOINT_SUFFIX)
+
+    @staticmethod
+    def is_checkpoint_unfinished(checkpoint_path: Union[Path, str]) -> bool:
+        """Check if the checkpoint is unfinished.
+
+        Args:
+            checkpoint_path: Path to the checkpoint file or dir.
+              Does not need to exist.
+
+        Returns:
+            True if the checkpoint is unfinished, False otherwise.
+        """
+        return NeMoModelCheckpoint.format_checkpoint_unfinished_marker_path(checkpoint_path).exists()
+
+    @staticmethod
+    def set_checkpoint_unfinished_marker(checkpoint_path: Union[Path, str], barrier_after=False) -> None:
+        """Marks given checkpoint as unfinished.
+
+        Args:
+            checkpoint_filepath: Path to the checkpoint file or dir.
+              Does not need to exist.
+            barrier_after: Synchronize ranks after writing the marker file.
+              Defaults to False.
+        """
+        if is_global_rank_zero():
+            marker_path = NeMoModelCheckpoint.format_checkpoint_unfinished_marker_path(checkpoint_path)
+            marker_path.parent.mkdir(parents=True, exist_ok=True)
+            marker_path.touch()
+        if barrier_after and torch.distributed.is_initialized():
+            torch.distributed.barrier()
+
+    @staticmethod
+    def remove_checkpoint_unfinished_marker(checkpoint_path: Union[Path, str], barrier_before=False) -> None:
+        """Clear unfinished marker for given checkpoint.
+
+        Args:
+            checkpoint_path: Path to the checkpoint file or dir.
+              Does not need to exist.
+            barrier_before: Synchronize ranks before removing the marker file.
+              Defaults to False.
+        """
+        try:
+            if barrier_before and torch.distributed.is_initialized():
+                torch.distributed.barrier()
+            if is_global_rank_zero():
+                marker_path = NeMoModelCheckpoint.format_checkpoint_unfinished_marker_path(checkpoint_path)
+                if marker_path.exists():
+                    marker_path.unlink()
+        except:
+            return
+
+    def file_exists(
+        self, filepath: str, trainer: "lightning.pytorch.Trainer", check_dist_ckpt: bool = True  # noqa: F821
+    ) -> bool:
+        """Checks if a file or a file without a suffix (distributed checkpoint) exists."""
+        if is_multistorageclient_url(filepath):
+            exists = self._fs.exists(filepath)
+        else:
+            exists = self._fs.exists(filepath) or (check_dist_ckpt and self._fs.exists(ckpt_to_dir(filepath)))
+
+        return trainer.strategy.broadcast(exists)
+
+    def _save_checkpoint(self, trainer: 'lightning.pytorch.Trainer', filepath: str) -> None:  # noqa: F821
+        # barrier_after=True, so all ranks continue after the unfinished checkpoint marker is placed.
+        # if anything goes wrong during checkpointing, we should be able to detect that data is incomplete.
+        self.set_checkpoint_unfinished_marker(filepath, barrier_after=True)
+        ema_callback = self._ema_callback(trainer)
+        if ema_callback is not None:
+            if self.async_save:
+                raise ValueError('async_save with EMA not supported')
+            with ema_callback.save_original_optimizer_state(trainer):
+                super()._save_checkpoint(trainer, filepath)
+
+            # save EMA copy of the model as well.
+            with ema_callback.save_ema_model(trainer):
+                filepath = self._ema_format_filepath(filepath)
+                if self.verbose:
+                    rank_zero_info(f"Saving EMA weights to separate checkpoint {filepath}")
+                super()._save_checkpoint(trainer, filepath)
+            self.remove_checkpoint_unfinished_marker(filepath, barrier_before=True)
+        else:
+            # Async save passed the finalization function to checkpoint_io,
+            # sync save calls the finalization function immediately after save.
+            finalize_fn = self._get_finalize_save_checkpoint_callback(trainer, filepath, trainer.global_step)
+            if self.async_save:
+                checkpoint_io = trainer.strategy.checkpoint_io
+                if not isinstance(checkpoint_io, AsyncFinalizableCheckpointIO):
+                    raise ValueError('Async save requires async compatible CheckpointIO')
+                storage_options = dict(finalize_fn=finalize_fn)
+                # Each upcoming ckpt removal request will be executed as part of this save finalization
+                self.deferred_ckpts_to_remove.append([])
+            else:
+                storage_options = None
+            logging.info(f'Checkpoint save for step {trainer.global_step} started at {time.time()}.')
+            trainer.save_checkpoint(filepath, self.save_weights_only, storage_options=storage_options)
+            if self.async_save:
+                logging.info(f'Scheduled async checkpoint save for {filepath}')
+            else:
+                finalize_fn()
+
+        if self.save_last_n_optim_states >= 0 and '-last' in filepath:
+            self._drop_optimizer_states(trainer, filepath, storage_options)
+
+    def _get_finalize_save_checkpoint_callback(
+        self, trainer: 'lightning.pytorch.Trainer', filepath: str, global_step: int  # noqa: F821
+    ):
+        """Creates a callback that can be used to finalize async (and sync) ckpt saves."""
+
+        def _cb():
+            logging.debug(f'Finalize callback called for step {global_step}, filepath {filepath}')
+            self._last_global_step_saved = global_step
+            self._last_checkpoint_saved = filepath
+
+            # notify loggers
+            if trainer.is_global_zero:
+                for logger in trainer.loggers:
+                    logger.after_save_checkpoint(proxy(self))
+
+            # barrier_before=True, so all ranks synchronize before removing the unfinished checkpoint marker
+            # we don't want to remove the marker until all checkpointing is done.
+            self.remove_checkpoint_unfinished_marker(filepath, barrier_before=True)
+
+            if not self.async_save:
+                return
+
+            logging.info(
+                f'Async checkpoint save for step {global_step} ({filepath}) finalized successfully at {time.time()}.'
+            )
+
+            # Remove checkpoints marked for removal by `self._remove_checkpoint`
+            # For each finalization there is exactly one entry in self.deferred_ckpts_to_remove
+            assert self.deferred_ckpts_to_remove
+            ckpts_to_remove = self.deferred_ckpts_to_remove.pop(0)
+            logging.debug(f'Checkpoints to remove: {ckpts_to_remove}')
+            for ckpt_to_remove in ckpts_to_remove:
+                self._remove_checkpoint(trainer, ckpt_to_remove, override_async=True)
+
+        return _cb
+
+    def _remove_checkpoint(
+        self, trainer: "lightning.pytorch.Trainer", filepath: str, override_async=False  # noqa: F821
+    ) -> None:
+        """Performs checkpoint removal or deferred removal.
+
+        With async save, `self._remove_checkpoint` is called before the checkpoint
+        is actually finished so we can't remove it. Instead we add it to
+        `self.deferred_ckpts_to_remove` for future removal.
+        """
+        if self.async_save and not override_async:
+            # Register checkpoint removal in the last (active) checkpoint removal list
+            self.deferred_ckpts_to_remove[-1].append(filepath)
+            return
+        # barrier_after=True, so all ranks continue after the unfinished checkpoint marker is placed.
+        # if anything goes wrong during removal, we should be able to detect that data is incomplete.
+        self.set_checkpoint_unfinished_marker(filepath, barrier_after=True)
+        super()._remove_checkpoint(trainer, filepath)
+        ema_callback = self._ema_callback(trainer)
+        if ema_callback is not None:
+            # remove EMA copy of the state dict as well.
+            filepath = self._ema_format_filepath(filepath)
+            super()._remove_checkpoint(trainer, filepath)
+        # barrier_before=True, so all ranks synchronize before removing the unfinished checkpoint marker
+        # we don't want to remove the marker until the checkpoint is actually removed.
+        self.remove_checkpoint_unfinished_marker(filepath, barrier_before=True)
+
+    def _ema_format_filepath(self, filepath: str) -> str:
+        return filepath.replace(self.FILE_EXTENSION, f'-EMA{self.FILE_EXTENSION}')
+
+    def _has_ema_ckpts(self, checkpoints: Iterable[Path]) -> bool:
+        return any(self._is_ema_filepath(checkpoint_path) for checkpoint_path in checkpoints)
+
+    def _is_ema_filepath(self, filepath: Union[Path, str]) -> bool:
+        return str(filepath).endswith(f'-EMA{self.FILE_EXTENSION}')
+
+    @property
+    def _saved_checkpoint_paths(self) -> Iterable[Path]:
+        # distributed checkpoints are directories so we check for them here
+        # we filter out unfinished checkpoints, these should be deleted during next cleanup
+
+        if is_multistorageclient_url(self.dirpath):
+            msc = import_multistorageclient()
+            return msc.glob(f"{self.dirpath}/*.ckpt")
+        else:
+            dist_checkpoints = [d for d in Path(self.dirpath).glob("*") if d.is_dir()]
+        if dist_checkpoints:
+            return filter(lambda p: not self.is_checkpoint_unfinished(p), dist_checkpoints)
+        else:
+            checkpoint_files = [f for f in Path(self.dirpath).rglob("*.ckpt")]
+            return filter(lambda p: not self.is_checkpoint_unfinished(p), checkpoint_files)
+
+    @staticmethod
+    def _remove_unfinished_checkpoints(checkpoint_dir: Union[Path, str]) -> None:
+
+        # Delete unfinished checkpoints from the filesystems.
+        # "Unfinished marker" files are removed as well.
+
+        if not is_global_rank_zero():
+            raise AssertionError("_remove_unfinished_checkpoints should run only on rank 0")
+
+        if is_multistorageclient_url(checkpoint_dir):
+            msc = import_multistorageclient()
+            existing_marker_filepaths = msc.glob(
+                f"{checkpoint_dir}*{NeMoModelCheckpoint.UNFINISHED_CHECKPOINT_SUFFIX}"
+            )
+            fs = get_filesystem(checkpoint_dir)
+            for ckpt_filepath in existing_marker_filepaths:
+                fs.rm(ckpt_filepath)
+        else:
+            checkpoint_dir = Path(checkpoint_dir)
+
+            existing_marker_filepaths = {
+                f.resolve()
+                for f in checkpoint_dir.glob(f"*{NeMoModelCheckpoint.UNFINISHED_CHECKPOINT_SUFFIX}")
+                if f.is_file()
+            }
+
+            checkpoint_filepaths = {f.resolve() for f in checkpoint_dir.rglob("*.ckpt") if f.is_file()}
+            for ckpt_filepath in checkpoint_filepaths:
+                possible_marker_path = NeMoModelCheckpoint.format_checkpoint_unfinished_marker_path(ckpt_filepath)
+                if possible_marker_path in existing_marker_filepaths:
+                    logging.warning(f'Removing unfinished checkpoint: {ckpt_filepath}')
+                    os.remove(ckpt_filepath)
+
+            # some directories might be distributed checkpoints, we remove these if they have a unfinished marker
+            all_dirpaths = {d.resolve() for d in checkpoint_dir.glob("*") if d.is_dir()}
+            for ckpt_dirpath in all_dirpaths:
+                possible_marker_path = NeMoModelCheckpoint.format_checkpoint_unfinished_marker_path(ckpt_dirpath)
+                if possible_marker_path in existing_marker_filepaths:
+                    logging.warning(f'Removing unfinished dist checkpoint: {ckpt_dirpath}')
+                    shutil.rmtree(ckpt_dirpath)
+
+            # delete markers
+            for marker_path in existing_marker_filepaths:
+                os.remove(marker_path)
+
+    def _should_remove_checkpoint(self, trainer: "pl.Trainer", previous: str, current: str) -> bool:  # noqa: F821
+        """Checks if the previous checkpoint should be deleted.
+        A checkpoint won't be deleted if any of the cases apply:
+        - The previous checkpoint is the same as the current checkpoint (means the old was already overwritten by new)
+        - The previous checkpoint is not in the current checkpoint directory and the filesystem is local
+        - The previous checkpoint is the checkpoint the Trainer resumed from and the filesystem is local
+            and the resumed from checkpoint is not the last checkpoint
+        """
+        if previous == current:
+            return False
+        if not _is_local_file_protocol(previous):
+            return True
+        previous = Path(previous).absolute()
+        resume_path = Path(trainer.ckpt_path).absolute() if trainer.ckpt_path is not None else None
+
+        if resume_path is not None and previous == resume_path:
+            if str(current).endswith("-last.ckpt") and resume_path.name.endswith("-last.ckpt"):
+                # delete the previous `-last.ckpt` checkpoint when current saved checkpoint is also `-last.ckpt`,
+                # if they're in the same directory
+                pass
+            else:
+                return False
+        if self.dirpath is None:
+            raise ValueError(f"{self.__class__}.dirpath is None.")
+        dirpath = Path(self.dirpath).absolute()
+        return dirpath in previous.parents

nemo/utils/callbacks/preemption.py

@@ -0,0 +1,118 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import signal
+import sys
+
+import torch
+from lightning.pytorch.callbacks import Callback
+
+from nemo.utils import logging
+
+
+class PreemptionCallback(Callback):
+    """
+    PreemptionCallback class creates a callback that checks for preemption during training at the end of every step.
+    Upon preemption the callback provides a function to gracefully exit the training immediately and also saves the
+    current state in a checkpoint as *last.ckpt.
+    (to be able to start from the same step without wasting any compute while resuming the next time).
+
+    PreemptionCallback is always enabled by default via the arg create_preemption_callback under ExpManagerConfig.
+    To disable please pass create_preemption_callback: False in your config file.
+    """
+
+    def __init__(self, checkpoint_callback, sig=None):
+        """Store the checkpoint callback and the signal to listen for (defaults to SIGTERM)."""
+        self.sig = sig
+        if self.sig is None:
+            self.sig = signal.SIGTERM
+        self.checkpoint_callback = checkpoint_callback
+        self.preemption_enabled = False
+
+    @property
+    def interrupted(self):
+        """Return whether a preemption signal was received, broadcasting rank 0's state to all ranks."""
+        interrupted = torch.tensor(self._interrupted, device=torch.cuda.current_device(), dtype=torch.int32)
+        torch.distributed.broadcast(interrupted, 0)
+        interrupted = bool(interrupted.item())
+        return interrupted
+
+    def on_train_start(self, trainer, pl_module):
+        """
+        Defines custom handlers at the beginning of training to be executed when the
+        preemption signal is received.
+        """
+
+        # Check if torch distributed is initialised, required for broadcasting the preemption signal to all the ranks
+        if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
+            logging.info("Preemption requires torch distributed to be initialized, disabling preemption")
+        else:
+            self.preemption_enabled = True
+            # Bool var that's initialized to false and made True upon receving the preemption signal
+            self._interrupted = False
+            self.released = False
+            self.original_handler = signal.getsignal(self.sig)
+
+            # Master handler on rank 0 only upon preemption signal to avoid deadlock conditions
+            def master_handler(signum, frame):
+                self.release()
+                self._interrupted = True
+
+            # Handler executed by the non zero ranks
+            def ignoring_handler(signum, frame):
+                self.release()
+
+            self.private_rank = torch.distributed.get_rank()
+            if self.private_rank == 0:
+                signal.signal(self.sig, master_handler)
+            else:
+                signal.signal(self.sig, ignoring_handler)
+
+        return self
+
+    def on_train_end(self, trainer, pl_module):
+        """Restore the original signal handler when training finishes."""
+        if self.preemption_enabled:
+            self.release()
+
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx: int):
+        """Check for preemption after each batch and, if signaled, save a last checkpoint and exit."""
+        if self.preemption_enabled:
+            # check if the job was preempted at the end of every training step/iteration
+            # NOTE: "self.interrupted" is a property which triggers a
+            # distributed broadcast of "_interrupted" flag from rank 0 to all other
+            # ranks, to avoid performance overheads it's best to store the result in
+            # a regular local variable
+            interrupted = self.interrupted
+            if interrupted:
+                logging.info("Received SIGTERM, saving checkpoint and exiting")
+                # Same off-by-one as in StatelessTimer: on_train_batch_end fires before
+                # batch_progress.increment_completed(), but the batch's optim step has
+                # already advanced global_step. Flush the in-flight batch so resume
+                # doesn't replay it and double-count the optim step.
+                from nemo.utils.exp_manager import _flush_in_flight_batch_progress
+
+                _flush_in_flight_batch_progress(trainer)
+                monitor_candidates = self.checkpoint_callback._monitor_candidates(trainer)
+                self.checkpoint_callback._save_last_checkpoint(trainer, monitor_candidates)
+                sys.exit(0)
+
+    def release(self):
+        """Restore the original signal handler; returns False if already released, True otherwise."""
+        if self.released:
+            return False
+
+        signal.signal(self.sig, self.original_handler)
+        self.released = True
+        return True

nemo/utils/callbacks/s3_checkpoint_io.py

@@ -0,0 +1,289 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import time
+from concurrent.futures import ProcessPoolExecutor
+from io import BytesIO
+from multiprocessing import get_start_method
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+from typing import Any, Callable, Dict, Optional, Union
+
+import torch
+from lightning.fabric.plugins.io.checkpoint_io import CheckpointIO
+
+from nemo.utils import logging
+from nemo.utils.s3_utils import (
+    DEFAULT_CHUNK_SIZE_MB,
+    DEFAULT_MAX_READ_CONCURRENCY,
+    DEFAULT_MAX_WRITE_CONCURRENCY,
+    SHARED_MEM_DIR,
+    S3Utils,
+)
+
+
+class S3CheckpointIO(CheckpointIO):
+    """A custom S3CheckpointIO module that supports checkpoint reading/writing with s3 when filepath
+    is a s3 url.
+    """
+
+    def __init__(
+        self,
+        dirpath: str,
+        chunk_size_MB=DEFAULT_CHUNK_SIZE_MB,
+        max_read_concurrency=DEFAULT_MAX_READ_CONCURRENCY,
+        max_write_concurrency=DEFAULT_MAX_WRITE_CONCURRENCY,
+        async_checkpointing=False,
+    ):
+        """
+        Initialize the transfer configuration with custom values.
+
+        This method overrides the default TransferConfig values in boto3.
+        See https://boto3.amazonaws.com/v1/documentation/api/latest/_modules/boto3/s3/transfer.html#TransferConfig
+
+        Args:
+            chunk_size_MB (int, optional): The size of chunks to use when transferring files.
+                Default is 64 (MB).
+            max_read_concurrency (int, optional): The maximum number of threads that will be making
+                requests to perform a download. Default is 15.
+            max_write_concurrency (int, optional): The maximum number of threads that will be making
+                requests to perform an upload. Default is 10.
+            async_checkpointing (bool, optional): Uses a ProcessPoolExecutor to do the main saving logic.
+                This feature should be used with save_top_k as it's possible a previous checkpoint is removed while
+                the current checkpoint write fails.
+        """
+        if not S3Utils.is_s3_url(dirpath):
+            raise AssertionError(
+                f"Error attempting to initialize an S3CheckpointIO when {dirpath} is not an S3 url. Please use TorchCheckpointIO when using a non-S3 dirpath."
+            )
+
+        self.chunk_size_MB = chunk_size_MB
+        self.max_read_concurrency = max_read_concurrency
+        self.max_write_concurrency = max_write_concurrency
+        self._async_checkpointing = async_checkpointing
+        '''
+        When using shared memory, we create a temporary file to hold the checkpoint before uploading to S3. 
+        This list will track those temporary files, and clean up any leaked files that are still around during teardown. 
+        '''
+        self._temp_files = []
+
+        if self.async_checkpointing:
+            # create an executor that will asynchronously run functions
+            self._executor = ProcessPoolExecutor(max_workers=1) if self.async_checkpointing else None
+
+            # Eager creating a subprocess now so that forked subprocess does not inherit cuda context from parent
+            if get_start_method() == 'fork' and torch.cuda.is_initialized() is True:
+                raise Exception(
+                    f'torch.cuda should not be initialized when checkpointing subprocess is created by fork method'
+                )
+            logging.info(f'Creating asynchronous checkpointing subprocess')
+            future = self._executor.submit(dummy_func)
+            try:
+                future.result()
+                logging.info(f'Asynchronous heckpointing subprocess created successfully')
+            except Exception as e:
+                logging.error(f'Failed to create asynchronous checkpointing subprocess, exception: {e}')
+                raise e
+            self._futures = []
+
+        super().__init__()
+
+    @property
+    def async_checkpointing(self):
+        return self._async_checkpointing
+
+    def _serialize_checkpoint_to_shm(self, checkpoint: Dict, path: str) -> str:
+        """
+        Returns:
+            filename of the temporary file in shared memory.
+        """
+        start_time = time.perf_counter()
+        tempfile = NamedTemporaryFile(dir=SHARED_MEM_DIR, delete=False)
+        torch.save(checkpoint, tempfile)
+        logging.info(
+            f'Time elapsed saving checkpoint dict to {tempfile.name} for {path}: {(time.perf_counter() - start_time):.2f} seconds, rank {torch.distributed.get_rank()}'
+        )
+        del checkpoint
+        return tempfile.name
+
+    def _serialize_checkpoint_to_bytes(self, checkpoint: Dict, path: str) -> BytesIO:
+        """
+        Returns:
+            The bytestring of the checkpoint.
+        """
+        ss = time.perf_counter()
+        bytes = BytesIO()
+        torch.save(checkpoint, bytes)
+        tt = time.perf_counter() - ss
+        logging.info(
+            f'Time elapsed saving checkpoint dict to bytes for {path}: {tt:.2f} seconds, rank {torch.distributed.get_rank()}'
+        )
+        del checkpoint
+        return bytes
+
+    def _check_uploading_results_so_far(self):
+        """
+        self._future is a list of tuples of form (future, destination path, source path)
+        This function checks the result of all the futures, and updates the self._futures list appropriately.
+        It also updates the list of self._temp_files, which is used to clean up leaked temporary files in SHARED_MEM during teardown.
+        """
+        if not self._futures:
+            return
+        start_time = time.perf_counter()
+        done_futures = []
+        in_progress_futures = []
+        for item in self._futures:
+            if item[0].done():
+                done_futures.append(item)
+            else:
+                in_progress_futures.append(item)
+
+        for item in done_futures:
+            try:
+                item[0].result()
+            except Exception as e:
+                logging.error(f'Failed to upload {item[2]} to {item[1]}, exception: {e}')
+                raise e
+            # If the future is complete, we can remove the temp file since we choose to clear the temp file when uploading.
+            try:
+                self._temp_files.remove(item[2])
+            except:
+                pass  # When not using shared memory, we do not append anything to the temp_files list, so remove will do nothing.
+        self._futures = in_progress_futures
+        logging.debug(
+            f'Time elapsed checking uploading future results: {(time.perf_counter() - start_time):.2f} seconds'
+        )
+
+    def save_checkpoint(
+        self, checkpoint: Dict[str, Any], path: Union[str, Path], storage_options: Optional[Any] = None
+    ) -> None:
+        # if we have a shared memory directory, we can serialize as a file to shared memory instead of as bytes.
+        if os.path.exists(SHARED_MEM_DIR):
+            localfile = self._serialize_checkpoint_to_shm(checkpoint, path)
+            self._temp_files.append(localfile)
+            saved_as_file = True
+        else:
+            bytes = self._serialize_checkpoint_to_bytes(checkpoint, path)
+            saved_as_file = False
+
+        if self.async_checkpointing:
+            self._check_uploading_results_so_far()
+            logging.info(f'Uploading checkpoint to {path} in asynchronous mode, rank {torch.distributed.get_rank()}')
+            if saved_as_file:
+                future = self._executor.submit(
+                    _upload_file_to_s3, localfile, path, self.chunk_size_MB, self.max_write_concurrency, True
+                )
+                self._futures.append((future, path, localfile))
+            else:
+                future = self._executor.submit(
+                    _upload_bytes_to_s3, bytes, path, self.chunk_size_MB, self.max_write_concurrency
+                )
+                self._futures.append((future, path, 'bytes'))
+        else:
+            logging.info(f'Uploading checkpoint to {path} in synchronous mode, rank {torch.distributed.get_rank()}')
+            if saved_as_file:
+                _upload_file_to_s3(localfile, path, self.chunk_size_MB, self.max_write_concurrency, True)
+                self._temp_files.remove(localfile)
+            else:
+                _upload_bytes_to_s3(bytes, path, self.chunk_size_MB, self.max_write_concurrency)
+
+    def load_checkpoint(
+        self, path: Union[str, Path], map_location: Optional[Callable] = lambda storage, loc: storage
+    ) -> Dict[str, Any]:
+        if os.path.exists(SHARED_MEM_DIR):
+            with NamedTemporaryFile(dir=SHARED_MEM_DIR, delete=True) as tempfile:
+                logging.info(
+                    f'Loading checkpoint {path} into a temp file in shared memory {tempfile.name}, rank {torch.distributed.get_rank()}'
+                )
+                S3Utils.download_s3_file_to_path(
+                    s3_path=path,
+                    file_path=tempfile.name,
+                    chunk_size_MB=self.chunk_size_MB,
+                    max_concurrency=self.max_read_concurrency,
+                )
+                checkpoint = torch.load(tempfile.name)
+        else:
+            file_stream: BytesIO = S3Utils.download_s3_file_to_stream(
+                s3_path=path, chunk_size_MB=self.chunk_size_MB, max_concurrency=self.max_read_concurrency
+            )
+            checkpoint = torch.load(file_stream)
+        return checkpoint
+
+    def remove_checkpoint(self, path: Union[str, Path]) -> None:
+        if S3Utils.is_s3_url(path):
+            S3Utils.remove_object(path)
+        else:
+            super().remove_checkpoint(path)
+
+    def teardown(self) -> None:
+        # this ensure we wait for final checkpoint to finish uploading at train end.
+        rank = torch.distributed.get_rank()
+        if self.async_checkpointing:
+            logging.info(f'Entering teardown, waiting for all jobs to finish, rank {rank}')
+            start_time = time.perf_counter()
+            self._executor.shutdown(wait=True)
+            logging.info(f'executor shut down after {(time.perf_counter() - start_time):.2f} seconds, rank {rank}')
+
+        '''
+        this will be non-empty at the end of training if using asynchronous uploading since the futures are not processed with _check_uploading_results_so_far.
+        therefore, we check that the path exists first before trying to delete. 
+        '''
+        if self._temp_files:
+            for tfile in self._temp_files:
+                if os.path.exists(tfile):
+                    try:
+                        os.remove(tfile)
+                    except Exception as e:
+                        logging.info(f"Error occurred while deleting file {tfile}: {e}")
+
+
+def _clean_up_conflicting_checkpoint(filepath: str) -> None:
+    '''
+    before saving to s3, clean up any existing object with the same prefix megatron_gpt+step_count
+    e.g. before we save "megatron_gpt--step=1400-validation_loss=6.32-consumed_samples=55920.0-last.ckpt"
+    we need to clean up "megatron_gpt--step=1400-validation_loss=xxx-consumed_samples=yyy-last.ckpt"
+    so that in case later we need to resume from step 1400, it has a single checkpoint file at step 1400
+    '''
+
+    if S3Utils.is_s3_url(filepath):
+        prefix_with_step = S3Utils.parse_prefix_with_step(filepath)
+        logging.info(f'Looking for conflicting checkpoint under prefix {prefix_with_step}')
+
+        conflict_last_ckpts = S3Utils.find_files_with_suffix(
+            base_path=prefix_with_step, suffix='last.ckpt', return_key_only=False
+        )
+        for last_ckpt in conflict_last_ckpts:
+            logging.info(f'Cleaning up conflicting last ckpt {last_ckpt} before saving {filepath}')
+            S3Utils.remove_object(last_ckpt)
+
+
+def _upload_file_to_s3(localfile, path, chunk_size_MB, max_write_concurrency, remove_file):
+    try:
+        _clean_up_conflicting_checkpoint(path)
+        S3Utils.upload_file(localfile, path, chunk_size_MB, max_write_concurrency, remove_file)
+    except Exception as e:
+        raise e
+
+
+def _upload_bytes_to_s3(bytes, path, chunk_size_MB, max_write_concurrency):
+    try:
+        _clean_up_conflicting_checkpoint(path)
+        S3Utils.upload_file_stream_to_s3(bytes, path, chunk_size_MB, max_write_concurrency)
+    except Exception as e:
+        raise e
+
+
+def dummy_func():
+    time.sleep(0.01)