chore: import upstream snapshot with attribution

nemo/utils/__init__.py

@@ -0,0 +1,44 @@
+# Copyright (c) 2020, 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.app_state import AppState
+from nemo.utils.cast_utils import (
+    CastToFloat,
+    CastToFloatAll,
+    avoid_bfloat16_autocast_context,
+    avoid_float16_autocast_context,
+    cast_all,
+    cast_tensor,
+    monkeypatched,
+)
+from nemo.utils.dtype import str_to_dtype
+from nemo.utils.nemo_logging import Logger as _Logger
+from nemo.utils.nemo_logging import LogMode as logging_mode
+
+logging = _Logger()
+try:
+    from nemo.utils.lightning_logger_patch import add_memory_handlers_to_pl_logger
+
+    add_memory_handlers_to_pl_logger()
+except ModuleNotFoundError:
+    pass
+
+try:
+    import webdataset
+    from nemo.utils.data_utils import wds_url_opener
+
+    webdataset.tariterators.url_opener = wds_url_opener
+except ModuleNotFoundError:
+    pass

nemo/utils/app_state.py

@@ -0,0 +1,924 @@
+# Copyright (c) 2020, 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 dataclasses import dataclass
+from threading import Lock
+from typing import Optional
+
+from nemo.utils.metaclasses import Singleton
+
+
+@dataclass()
+class ModelMetadataRegistry:
+    """
+    Dataclass for model metadata registry.
+    """
+
+    guid: str
+    gidx: int
+    restoration_path: Optional[str] = None
+
+
+class AppState(metaclass=Singleton):
+    """
+    App state for the application.
+    """
+
+    def __init__(self):
+        # method call lock
+        self.__lock = Lock()
+
+        # TODO: should we store global config in hydra_runner?
+        self._app_cfg = None
+
+        # World info
+        self._device_id = None
+        self._local_rank = None
+        self._global_rank = None
+        self._tensor_model_parallel_rank = None
+        self._expert_model_parallel_rank = None
+        self._expert_tensor_parallel_rank = None
+        self._pipeline_model_parallel_rank = None
+        self._data_parallel_rank = None
+
+        self._world_size = None
+        self._model_parallel_size = None
+        self._tensor_model_parallel_size = None
+        self._tensor_model_parallel_group = None
+        self._expert_model_parallel_size = None
+        self._expert_tensor_parallel_size = None
+        self._pipeline_model_parallel_size = None
+        self._virtual_pipeline_model_parallel_size = None
+        self._encoder_tensor_model_parallel_size = None
+        self._encoder_pipeline_model_parallel_size = None
+        self._pipeline_model_parallel_group = None
+        self._pipeline_model_parallel_split_rank = None
+        self._pipeline_model_parallel_comm_backend = None
+        self._is_megatron_initialized = False
+        self._data_parallel_size = None
+        self._data_parallel_group = None
+        self._use_tp_pp_dp_mapping = False
+        self._num_distributed_optimizer_instances = 1
+        self._megatron_checkpoint_version = None
+        self._use_fp8 = False
+        self._context_parallel_size = None
+        self._init_mpi_proc_gruop = False
+        self._nccl_communicator_config_path = None
+        self._use_sharp = False
+        self._create_all_gather_group = False
+        self._use_gloo_process_groups = True
+
+        self._random_seed = None
+
+        # Logging info
+        self._log_dir = None
+        self._exp_dir = None
+        self._name = None
+        self._checkpoint_name = None
+        self._version = None
+        self._create_checkpoint_callback = None
+        self._checkpoint_callback_params = None
+
+        # Save and Restore (.nemo)
+        self._tmpdir_name = None
+        self._is_model_being_restored = False
+        self._nemo_file_folder = None
+        self._model_restore_path = None
+        self._all_model_restore_paths = []
+        self._model_guid_map = {}  # type: Dict[str, ModelMetadataRegistry]
+        self._restore = False  # TODO: are this and _is_model_being_restored both needed?
+
+        # files from a previous run to move into a new directory
+        self.files_to_move = []
+        # files to copy into log dir
+        self._files_to_copy = []
+        # command-ling arguments for run
+        self._cmd_args = None
+
+        # Insert NVTX ranges to categorize execution
+        self._nvtx_ranges = False
+
+    @property
+    def device_id(self):
+        """Property returns the device_id
+        Returns:
+            device_id
+        """
+        return self._device_id
+
+    @device_id.setter
+    def device_id(self, id):
+        """Property sets the device_id.
+        Args:
+            size (int): The device id.
+        """
+        self._device_id = id
+
+    @property
+    def world_size(self):
+        """Property returns the total number of GPUs.
+        Returns:
+            Total number of GPUs.
+        """
+        return self._world_size
+
+    @world_size.setter
+    def world_size(self, size):
+        """Property sets the total number of GPUs.
+        Args:
+            size (int):  Total number of GPUs.
+        """
+        self._world_size = size
+
+    @property
+    def model_parallel_size(self):
+        """Property returns the number of GPUs in each model parallel group.
+        Returns:
+            Number of GPUs in each model parallel group.
+        """
+        return self._model_parallel_size
+
+    @model_parallel_size.setter
+    def model_parallel_size(self, size):
+        """Property sets the number of GPUs in each model parallel group.
+        Args:
+            size (int):  Number of GPUs in each model parallel group.
+        """
+        self._model_parallel_size = size
+
+    @property
+    def tensor_model_parallel_size(self):
+        """Property returns the number of GPUs in each model parallel group.
+        Returns:
+            Number of GPUs in each model parallel group.
+        """
+        return self._tensor_model_parallel_size
+
+    @tensor_model_parallel_size.setter
+    def tensor_model_parallel_size(self, size):
+        """Property sets the number of GPUs in each model parallel group.
+        Args:
+            size (int):  Number of GPUs in each model parallel group.
+        """
+        self._tensor_model_parallel_size = size
+
+    @property
+    def expert_model_parallel_rank(self):
+        """Property returns the expert model parallel rank.
+        Returns:
+            Tensor model parallel rank.
+        """
+        return self._expert_model_parallel_rank
+
+    @expert_model_parallel_rank.setter
+    def expert_model_parallel_rank(self, rank):
+        """Property sets the expert model parallel rank.
+        Args:
+            rank (int):  Tensor model parallel rank.
+        """
+        self._expert_model_parallel_rank = rank
+
+    @property
+    def expert_model_parallel_size(self):
+        """Property returns the number of GPUs in each expert parallel group.
+        Returns:
+            Number of GPUs in each expert parallel group.
+        """
+        return self._expert_model_parallel_size
+
+    @expert_model_parallel_size.setter
+    def expert_model_parallel_size(self, size):
+        """Property returns the number of GPUs in each expert parallel group.
+        Returns:
+            Number of GPUs in each expert parallel group.
+        """
+        self._expert_model_parallel_size = size
+
+    @property
+    def expert_tensor_parallel_size(self):
+        """Property returns the number of GPUs in each expert tensor parallel group.
+        Returns:
+            Number of GPUs in each expert tensor parallel group.
+        """
+        return self._expert_tensor_parallel_size
+
+    @expert_tensor_parallel_size.setter
+    def expert_tensor_parallel_size(self, size):
+        """Property sets the number of GPUs in each expert tensor parallel group.
+        Args:
+            size (int):  Number of GPUs in each tensor expert parallel group.
+        """
+        self._expert_tensor_parallel_size = size
+
+    @property
+    def expert_tensor_parallel_rank(self):
+        """Property returns the expert tensor model parallel rank.
+        Returns:
+            Tensor model parallel rank.
+        """
+        return self._expert_tensor_parallel_rank
+
+    @expert_tensor_parallel_rank.setter
+    def expert_tensor_parallel_rank(self, rank):
+        """Property sets the expert tensor model parallel rank.
+        Args:
+            rank (int):  Tensor model parallel rank.
+        """
+        self._expert_tensor_parallel_rank = rank
+
+    @property
+    def pipeline_model_parallel_size(self):
+        """Property returns the number of GPUs in each model parallel group.
+        Returns:
+            Number of GPUs in each model parallel group.
+        """
+        return self._pipeline_model_parallel_size
+
+    @pipeline_model_parallel_size.setter
+    def pipeline_model_parallel_size(self, size):
+        """Property sets the number of GPUs in each model parallel group.
+        Args:
+            size (int):  Number of GPUs in each model parallel group.
+        """
+        self._pipeline_model_parallel_size = size
+
+    @property
+    def pipeline_model_parallel_comm_backend(self):
+        """Property returns the backend communication library of pipeline communication.
+        Returns:
+            Backend communication library of pipeline communication.
+        """
+        return self._pipeline_model_parallel_comm_backend
+
+    @pipeline_model_parallel_comm_backend.setter
+    def pipeline_model_parallel_comm_backend(self, backend):
+        """Property sets the backend communication library of pipeline communication.
+        Args:
+            backend (str): Backend communication library of pipeline communication.
+        """
+        self._pipeline_model_parallel_comm_backend = backend
+
+    @property
+    def encoder_tensor_model_parallel_size(self):
+        """Property returns the number of GPUs in each model parallel group.
+        Returns:
+            Number of GPUs in each model parallel group.
+        """
+        return self._encoder_tensor_model_parallel_size
+
+    @encoder_tensor_model_parallel_size.setter
+    def encoder_tensor_model_parallel_size(self, size):
+        """Property sets the number of GPUs in each model parallel group.
+        Args:
+            size (int):  Number of GPUs in each model parallel group.
+        """
+        self._encoder_tensor_model_parallel_size = size
+
+    @property
+    def encoder_pipeline_model_parallel_size(self):
+        """Property returns the number of GPUs in each model parallel group.
+        Returns:
+            Number of GPUs in each model parallel group.
+        """
+        return self._encoder_pipeline_model_parallel_size
+
+    @encoder_pipeline_model_parallel_size.setter
+    def encoder_pipeline_model_parallel_size(self, size):
+        """Property sets the number of GPUs in each model parallel group.
+        Args:
+            size (int):  Number of GPUs in each model parallel group.
+        """
+        self._encoder_pipeline_model_parallel_size = size
+
+    @property
+    def use_tp_pp_dp_mapping(self):
+        """Property returns whether to use TP-PP-DP mapping.
+        Returns:
+            Whether to use TP-PP-DP mapping.
+        """
+        return self._use_tp_pp_dp_mapping
+
+    @use_tp_pp_dp_mapping.setter
+    def use_tp_pp_dp_mapping(self, use_new_mapping):
+        """Property sets whether to use TP-PP-DP mapping.
+        Args:
+            use_new_mapping (bool):  Whether to use TP-PP-DP mapping.
+        """
+        self._use_tp_pp_dp_mapping = use_new_mapping
+
+    @property
+    def num_distributed_optimizer_instances(self):
+        """Property returns the factor by which the Partial DistOpt is sharded.
+        Returns:
+            The partial DistOpt shard factor
+        """
+        return self._num_distributed_optimizer_instances
+
+    @num_distributed_optimizer_instances.setter
+    def num_distributed_optimizer_instances(self, shard_factor):
+        """Property sets the factor by which the Partial DistOpt is sharded.
+        Args:
+            shard_factor (int):  The partial DistOpt shard factor.
+        """
+        self._num_distributed_optimizer_instances = shard_factor
+
+    @property
+    def virtual_pipeline_model_parallel_size(self):
+        """Property returns the number of GPUs in each model parallel group.
+        Returns:
+            Number of GPUs in each model parallel group.
+        """
+        return self._virtual_pipeline_model_parallel_size
+
+    @virtual_pipeline_model_parallel_size.setter
+    def virtual_pipeline_model_parallel_size(self, size):
+        """Property sets the size of the virtual pipeline parallel model.
+        Args:
+            size (int):  Number of modules in each pipeline parallel model.
+        """
+        self._virtual_pipeline_model_parallel_size = size
+
+    @property
+    def data_parallel_size(self):
+        """Property returns the number of GPUs in each data parallel group.
+        Returns:
+            Number of GPUs in each data parallel group.
+        """
+        return self._data_parallel_size
+
+    @data_parallel_size.setter
+    def data_parallel_size(self, size):
+        """Property sets the number of GPUs in each data parallel group.
+        Args:
+            size (int):  Number of GPUs in each data parallel group.
+        """
+        self._data_parallel_size = size
+
+    @property
+    def local_rank(self):
+        """Property returns the local rank.
+        Returns:
+            Local rank.
+        """
+        return self._local_rank
+
+    @local_rank.setter
+    def local_rank(self, rank):
+        """Property sets the local rank.
+        Args:
+            rank (int):  Local rank.
+        """
+        self._local_rank = rank
+
+    @property
+    def global_rank(self):
+        """Property returns the global rank.
+        Returns:
+            Global rank.
+        """
+        return self._global_rank
+
+    @global_rank.setter
+    def global_rank(self, rank):
+        """Property sets the global rank.
+        Args:
+            rank (int):  Global rank.
+        """
+        self._global_rank = rank
+
+    @property
+    def tensor_model_parallel_rank(self):
+        """Property returns the tensor model parallel rank.
+        Returns:
+            Tensor model parallel rank.
+        """
+        return self._tensor_model_parallel_rank
+
+    @tensor_model_parallel_rank.setter
+    def tensor_model_parallel_rank(self, rank):
+        """Property sets the tensor model parallel rank.
+        Args:
+            rank (int):  Tensor model parallel rank.
+        """
+        self._tensor_model_parallel_rank = rank
+
+    @property
+    def tensor_model_parallel_group(self):
+        """Property returns the tensor model parallel group.
+        Returns:
+            Tensor model parallel group.
+        """
+        return self._tensor_model_parallel_group
+
+    @tensor_model_parallel_group.setter
+    def tensor_model_parallel_group(self, group):
+        """Property sets the tensor model parallel group.
+        Args:
+            group:  Tensor model parallel group.
+        """
+        self._tensor_model_parallel_group = group
+
+    @property
+    def pipeline_model_parallel_rank(self):
+        """Property returns the pipeline model parallel rank.
+        Returns:
+            Pipeline model parallel rank.
+        """
+        return self._pipeline_model_parallel_rank
+
+    @pipeline_model_parallel_rank.setter
+    def pipeline_model_parallel_rank(self, rank):
+        """Property sets the pipeline model parallel rank.
+        Args:
+            rank (int):  Pipeline model parallel rank.
+        """
+        self._pipeline_model_parallel_rank = rank
+
+    @property
+    def virtual_pipeline_model_parallel_rank(self):
+        """Property returns the virtual pipeline parallel rank.
+        Returns:
+            Model parallel rank.
+        """
+        return self._virtual_pipeline_model_parallel_rank
+
+    @virtual_pipeline_model_parallel_rank.setter
+    def virtual_pipeline_model_parallel_rank(self, rank):
+        """Property sets the virtual pipeline parallel rank.
+        Args:
+            rank (int):  Virtual pipeline parallel rank.
+        """
+        self._virtual_pipeline_model_parallel_rank = rank
+
+    @property
+    def encoder_tensor_model_parallel_rank(self):
+        """Property returns the encoder tensor model parallel rank.
+        Returns:
+            Tensor model parallel rank.
+        """
+        return self._encoder_tensor_model_parallel_rank
+
+    @encoder_tensor_model_parallel_rank.setter
+    def encoder_tensor_model_parallel_rank(self, rank):
+        """Property sets the encoder tensor model parallel rank.
+        Args:
+            rank (int):  Tensor model parallel rank.
+        """
+        self._encoder_tensor_model_parallel_rank = rank
+
+    @property
+    def encoder_pipeline_model_parallel_rank(self):
+        """Property returns the encoder pipeline model parallel rank.
+        Returns:
+            Tensor model parallel rank.
+        """
+        return self._encoder_pipeline_model_parallel_rank
+
+    @encoder_pipeline_model_parallel_rank.setter
+    def encoder_pipeline_model_parallel_rank(self, rank):
+        """Property sets the encoder pipeline model parallel rank.
+        Args:
+            rank (int):  Tensor model parallel rank.
+        """
+        self._encoder_pipeline_model_parallel_rank = rank
+
+    @property
+    def pipeline_model_parallel_split_rank(self):
+        """Property returns the rank at which Encoder and Decoder are split into different pipelines for
+        Megatrron Encoder-Decoder models.
+        Returns:
+            Pipeline model parallel split rank.
+        """
+        return self._pipeline_model_parallel_split_rank
+
+    @pipeline_model_parallel_split_rank.setter
+    def pipeline_model_parallel_split_rank(self, rank):
+        """Property sets the rank at which Encoder and Decoder are split into different pipelines for
+        Megatron Encoder-Decoder models.
+        Args:
+            rank (int): Model parallel split rank.
+        """
+        self._pipeline_model_parallel_split_rank = rank
+
+    @property
+    def pipeline_model_parallel_group(self):
+        """Property returns the pipeline model parallel group.
+        Returns:
+            Pipeline model parallel group.
+        """
+        return self._pipeline_model_parallel_group
+
+    @pipeline_model_parallel_group.setter
+    def pipeline_model_parallel_group(self, group):
+        """Property sets the pipeline model parallel group.
+        Args:
+            group:  Pipeline model parallel group.
+        """
+        self._pipeline_model_parallel_group = group
+
+    @property
+    def data_parallel_rank(self):
+        """Property returns the data parallel rank.
+        Returns:
+            Data parallel rank.
+        """
+        return self._data_parallel_rank
+
+    @data_parallel_rank.setter
+    def data_parallel_rank(self, rank):
+        """Property sets the data parallel rank.
+        Args:
+            rank (int):  Data parallel rank.
+        """
+        self._data_parallel_rank = rank
+
+    @property
+    def data_parallel_group(self):
+        """Property returns the data parallel group.
+        Returns:
+            Data parallel group.
+        """
+        return self._data_parallel_group
+
+    @data_parallel_group.setter
+    def data_parallel_group(self, group):
+        """Property sets the data parallel group.
+        Args:
+            group:  Data parallel group.
+        """
+        self._data_parallel_group = group
+
+    @property
+    def use_fp8(self):
+        """Property returns the use of fp8 precision.
+        Returns:
+            Use of FP8.
+        """
+        return self._use_fp8
+
+    @use_fp8.setter
+    def use_fp8(self, use_fp8):
+        """Property sets the use of fp8 precision.
+        Args:
+            use_fp8:  Use of FP8.
+        """
+        self._use_fp8 = use_fp8
+
+    @property
+    def use_sharp(self):
+        """Property returns whether to use SHARP for all-reduce operations.
+        Returns:
+            Whether to use SHARP.
+        """
+        return self._use_sharp
+
+    @use_sharp.setter
+    def use_sharp(self, use_sharp):
+        """Property sets whether to use SHARP for all-reduce operations.
+        Args:
+            use_sharp (bool): Whether to use SHARP.
+        """
+        self._use_sharp = use_sharp
+
+    @property
+    def create_all_gather_group(self):
+        """Property returns whether to create a separate all-gather process group.
+        Returns:
+            Whether to create a separate all-gather process group.
+        """
+        return self._create_all_gather_group
+
+    @create_all_gather_group.setter
+    def create_all_gather_group(self, create_all_gather_group):
+        """Property sets whether to create a separate all-gather process group.
+        Args:
+            create_all_gather_group (bool): Whether to create a separate all-gather process group.
+        """
+        self._create_all_gather_group = create_all_gather_group
+
+    @property
+    def use_gloo_process_groups(self):
+        """Property returns whether to use Gloo process groups.
+        Returns:
+            Whether to use Gloo process groups.
+        """
+        return self._use_gloo_process_groups
+
+    @use_gloo_process_groups.setter
+    def use_gloo_process_groups(self, use_gloo_process_groups):
+        """Property sets whether to use Gloo process groups.
+        Args:
+            use_gloo_process_groups (bool): Whether to use Gloo process groups.
+        """
+        self._use_gloo_process_groups = use_gloo_process_groups
+
+    @property
+    def context_parallel_size(self):
+        """Property returns the number of GPUs in each context parallel group.
+        Returns:
+            Number of GPUs in each context parallel group.
+        """
+        return self._context_parallel_size
+
+    @context_parallel_size.setter
+    def context_parallel_size(self, size):
+        """Property sets the number of GPUs in each context parallel group.
+        Args:
+            size (int):  Number of GPUs in each context parallel group.
+        """
+        self._context_parallel_size = size
+
+    @property
+    def init_mpi_proc_group(self):
+        """Property sets the initialization of mpi process group.
+        Returns:
+            Initialize mpi process group.
+        """
+        return self._init_mpi_proc_group
+
+    @init_mpi_proc_group.setter
+    def init_mpi_proc_group(self, init_mpi_proc_group):
+        """Property sets the initialization of mpi process group.
+        Args:
+            init_mpi_proc_group:  Initialize mpi process group.
+        """
+        self._init_mpi_proc_group = init_mpi_proc_group
+
+    @property
+    def nccl_communicator_config_path(self):
+        """Property returns the path to the nccl communicator config.
+        Returns:
+            Path to the nccl communicator config.
+        """
+        return self._nccl_communicator_config_path
+
+    @nccl_communicator_config_path.setter
+    def nccl_communicator_config_path(self, path):
+        """Property sets the path to the nccl communicator config.
+        Args:
+            path (str):  Path to the nccl communicator config.
+        """
+        self._nccl_communicator_config_path = path
+
+    @property
+    def random_seed(self):
+        """Property returns the random seed.
+        Returns:
+            Random seed.
+        """
+        return self._random_seed
+
+    @random_seed.setter
+    def random_seed(self, seed):
+        """Property sets the random seed.
+        Args:
+            seed (int):  Random seed.
+        """
+        self._random_seed = seed
+
+    @property
+    def log_dir(self):
+        """Returns the log_dir set by exp_manager."""
+        return self._log_dir
+
+    @log_dir.setter
+    def log_dir(self, dir):
+        """Sets the log_dir property.
+
+        Args:
+            dir (str): Log_dir set by exp_manager.
+        """
+        self._log_dir = dir
+
+    @property
+    def exp_dir(self):
+        """Returns the exp_dir set by exp_manager."""
+        return self._exp_dir
+
+    @exp_dir.setter
+    def exp_dir(self, dir):
+        """Sets the log_dir property.
+
+        Args:
+            dir (str): Log_dir set by exp_manager.
+        """
+        self._exp_dir = dir
+
+    @property
+    def name(self):
+        """Returns the name set by exp_manager."""
+        return self._name
+
+    @name.setter
+    def name(self, name):
+        """Sets the name property.
+
+        Args:
+            dir (str): name set by exp_manager.
+        """
+        self._name = name
+
+    @property
+    def checkpoint_name(self):
+        """Returns the name set by exp_manager."""
+        return self._checkpoint_name
+
+    @checkpoint_name.setter
+    def checkpoint_name(self, name):
+        """Sets the name property.
+
+        Args:
+            dir (str): name set by exp_manager.
+        """
+        self._checkpoint_name = name
+
+    @property
+    def version(self):
+        """Returns the version set by exp_manager."""
+        return self._version
+
+    @version.setter
+    def version(self, version):
+        """Sets the version property.
+
+        Args:
+            dir (str): version set by exp_manager.
+        """
+        self._version = version
+
+    @property
+    def create_checkpoint_callback(self):
+        """Returns the create_checkpoint_callback set by exp_manager."""
+        return self._create_checkpoint_callback
+
+    @create_checkpoint_callback.setter
+    def create_checkpoint_callback(self, create_checkpoint_callback):
+        """Sets the create_checkpoint_callback property.
+
+        Args:
+            dir (bool): create_checkpoint_callback set by exp_manager.
+        """
+        self._create_checkpoint_callback = create_checkpoint_callback
+
+    @property
+    def checkpoint_callback_params(self):
+        """Returns the version set by exp_manager."""
+        return self._checkpoint_callback_params
+
+    @checkpoint_callback_params.setter
+    def checkpoint_callback_params(self, params):
+        """Sets the name property.
+
+        Args:
+            params (dict): checkpoint_callback_params set by exp_manager.
+        """
+        self._checkpoint_callback_params = params
+
+    @property
+    def files_to_move(self):
+        """Returns the list of files to move into a separate directory."""
+        return self._files_to_move
+
+    @files_to_move.setter
+    def files_to_move(self, files):
+        """Sets the files_to_move property.
+
+        Args:
+            files (list[str]): list of filenames to move.
+        """
+        self._files_to_move = files
+
+    @property
+    def files_to_copy(self):
+        """Returns the list of files to copy into the log dir."""
+        return self._files_to_copy
+
+    @files_to_copy.setter
+    def files_to_copy(self, files):
+        """Sets the files_to_copy property.
+
+        Args:
+            files (list[str]): list of filenames to copy.
+        """
+        self._files_to_copy = files
+
+    @property
+    def cmd_args(self):
+        """Returns the command line arguments for the current run."""
+        return self._cmd_args
+
+    @cmd_args.setter
+    def cmd_args(self, args):
+        """Sets the cmd_args property.
+
+        Args:
+            args (list[str]): list of the command line arguments
+                used to run the experiment.
+        """
+        self._cmd_args = args
+
+    @property
+    def model_restore_path(self):
+        """Property returns the model restore path.
+        Returns:
+            Model restore path.
+        """
+        restore_path = self._all_model_restore_paths[-1] if len(self._all_model_restore_paths) > 0 else None
+        return restore_path
+
+    @model_restore_path.setter
+    def model_restore_path(self, path):
+        """Property sets the model restore path.
+        Args:
+            path (str): Model restore path.
+        """
+        with self.__lock:
+            self._model_restore_path = path
+            self._all_model_restore_paths.append(path)
+
+    def register_model_guid(self, guid: str, restoration_path: Optional[str] = None):
+        """Maps a guid to its restore path (None or last absolute path).
+        Args:
+            guid (str): Guid.
+            restoration_path (Optional[str]): Restore path.
+        """
+        with self.__lock:
+            if guid in self._model_guid_map:
+                idx = self._model_guid_map[guid].gidx
+            else:
+                idx = len(self._model_guid_map)
+            self._model_guid_map[guid] = ModelMetadataRegistry(guid, idx, restoration_path=restoration_path)
+
+    def reset_model_guid_registry(self):
+        """Resets the guid mapping."""
+        with self.__lock:
+            self._model_guid_map.clear()
+
+    def get_model_metadata_from_guid(self, guid) -> ModelMetadataRegistry:
+        """Returns the global model idx and restoration path.
+        Args:
+            guid (str): Guid.
+        Returns:
+            Model metadata registry.
+        """
+        metadata = self._model_guid_map[guid]
+        return metadata
+
+    @property
+    def is_model_being_restored(self) -> bool:
+        """Property returns whether the model is being restored.
+        Returns:
+            Whether the model is being restored.
+        """
+        return self._is_model_being_restored
+
+    @is_model_being_restored.setter
+    def is_model_being_restored(self, is_restored: bool):
+        """Property sets whether the model is being restored.
+        Args:
+            is_restored (bool): Whether the model is being restored.
+        """
+        self._is_model_being_restored = is_restored
+
+    @property
+    def nemo_file_folder(self) -> str:
+        """Property returns the nemo file folder.
+        Returns:
+            Nemo file folder.
+        """
+        return self._nemo_file_folder
+
+    @nemo_file_folder.setter
+    def nemo_file_folder(self, path: str):
+        """Property sets the nemo file folder.
+        Args:
+            path (str): Nemo file folder.
+        """
+        self._nemo_file_folder = path
+
+    @property
+    def restore(self) -> bool:
+        """Property returns whether to restore the model.
+        Returns:
+            Whether to restore the model.
+        """
+        return self._restore
+
+    @restore.setter
+    def restore(self, restore: bool):
+        """Property sets whether to restore the model.
+        Args:
+            restore (bool): Whether to restore the model.
+        """
+        self._restore = restore

nemo/utils/arguments.py

@@ -0,0 +1,132 @@
+# Copyright (c) 2020, 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 argparse import ArgumentParser
+from typing import Any, Dict, List, Optional, Union
+
+
+def add_optimizer_args(
+    parent_parser: ArgumentParser,
+    optimizer: str = 'adam',
+    default_lr: float = None,
+    default_opt_args: Optional[Union[Dict[str, Any], List[str]]] = None,
+) -> ArgumentParser:
+    """Extends existing argparse with support for optimizers.
+
+    # Example of adding optimizer args to command line :
+    python train_script.py ... --optimizer "novograd" --lr 0.01 \
+        --opt_args betas=0.95,0.5 weight_decay=0.001
+
+    Args:
+        parent_parser (ArgumentParser): Custom CLI parser that will be extended.
+        optimizer (str): Default optimizer required.
+        default_lr (float): Default learning rate that should be overriden during training.
+        default_opt_args (list(str)): List of overriding arguments for the instantiated optimizer.
+
+    Returns:
+        ArgumentParser: Parser extended by Optimizers arguments.
+    """
+    if default_opt_args is None:
+        default_opt_args = []
+
+    parser = ArgumentParser(parents=[parent_parser], add_help=True, conflict_handler='resolve')
+
+    parser.add_argument('--optimizer', type=str, default=optimizer, help='Name of the optimizer. Defaults to Adam.')
+    parser.add_argument('--lr', type=float, default=default_lr, help='Learning rate of the optimizer.')
+    parser.add_argument(
+        '--opt_args',
+        default=default_opt_args,
+        nargs='+',
+        type=str,
+        help='Overriding arguments for the optimizer. \n Must follow the pattern : \n name=value separated by spaces.'
+        'Example: --opt_args weight_decay=0.001 eps=1e-8 betas=0.9,0.999',
+    )
+
+    return parser
+
+
+def add_scheduler_args(parent_parser: ArgumentParser) -> ArgumentParser:
+    """Extends existing argparse with default LR scheduler args.
+
+    Args:
+        parent_parser (ArgumentParser): Custom CLI parser that will be extended.
+
+    Returns:
+        ArgumentParser: Parser extended by LR Scheduler arguments.
+    """
+    parser = ArgumentParser(parents=[parent_parser], add_help=False, conflict_handler='resolve')
+    parser.add_argument("--warmup_steps", type=int, required=False, default=None, help="Number of warmup steps")
+    parser.add_argument(
+        "--warmup_ratio",
+        type=float,
+        required=False,
+        default=None,
+        help="Number of warmup steps as a percentage of total training steps",
+    )
+    parser.add_argument("--hold_steps", type=int, required=False, default=None, help="Number of hold LR steps")
+    parser.add_argument(
+        "--hold_ratio",
+        type=float,
+        required=False,
+        default=None,
+        help="Number of hold LR steps as a percentage of total training steps",
+    )
+    parser.add_argument("--min_lr", type=float, required=False, default=0.0, help="Minimum learning rate")
+    parser.add_argument(
+        "--last_epoch", type=int, required=False, default=-1, help="Last epoch id. -1 indicates training from scratch"
+    )
+    return parser
+
+
+def add_asr_args(parent_parser: ArgumentParser) -> ArgumentParser:
+    """Extends existing argparse with default ASR collection args.
+
+    Args:
+        parent_parser (ArgumentParser): Custom CLI parser that will be extended.
+
+    Returns:
+        ArgumentParser: Parser extended by NeMo ASR Collection arguments.
+    """
+    parser = ArgumentParser(parents=[parent_parser], add_help=False, conflict_handler='resolve')
+    parser.add_argument("--asr_model", type=str, required=True, default=None, help="Path to ASR model config yaml")
+    parser.add_argument("--train_dataset", type=str, required=True, default=None, help="training dataset path")
+    parser.add_argument("--eval_dataset", type=str, required=True, help="evaluation dataset path")
+    return parser
+
+
+def add_nlp_args(parent_parser: ArgumentParser) -> ArgumentParser:
+    """Extends existing argparse with default NLP collection args.
+
+    Args:
+        parent_parser (ArgumentParser): Custom CLI parser that will be extended.
+
+    Returns:
+        ArgumentParser: Parser extended by NeMo NLP Collection arguments.
+    """
+    parser = ArgumentParser(parents=[parent_parser], add_help=False, conflict_handler='resolve')
+    parser.add_argument(
+        "--data_dir", type=str, required=False, help="data directory to training or/and evaluation dataset"
+    )
+    parser.add_argument(
+        "--config_file", type=str, required=False, default=None, help="Huggingface model configuration file"
+    )
+    parser.add_argument(
+        "--pretrained_model_name", default='bert-base-uncased', type=str, required=False, help="pretrained model name"
+    )
+    parser.add_argument(
+        "--tokenizer_name", default='nemobert', type=str, choices=['sentencepiece', 'nemobert'], help="Tokenizer type"
+    )
+    parser.add_argument("--tokenizer_model", default=None, type=str, help="Tokenizer file for sentence piece")
+    parser.add_argument("--do_lower_case", action='store_true', required=False, help="lower case data")
+    return parser

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)

nemo/utils/cast_utils.py

@@ -0,0 +1,119 @@
+# Copyright (c) 2022, 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 contextlib import contextmanager, nullcontext
+from typing import Any
+
+import torch
+
+
+def avoid_bfloat16_autocast_context():
+    """
+    If the current autocast context is bfloat16,
+    cast it to float32
+    """
+
+    if torch.is_autocast_enabled() and torch.get_autocast_gpu_dtype() == torch.bfloat16:
+        return torch.amp.autocast('cuda', dtype=torch.float32)
+    else:
+        return nullcontext()
+
+
+def avoid_float16_autocast_context():
+    """
+    If the current autocast context is float16, cast it to bfloat16
+    if available (unless we're in jit) or float32
+    """
+
+    if torch.is_autocast_enabled() and torch.get_autocast_gpu_dtype() == torch.float16:
+        if torch.jit.is_scripting() or torch.jit.is_tracing():
+            return torch.amp.autocast('cuda', dtype=torch.float32)
+
+        if torch.cuda.is_bf16_supported():
+            return torch.amp.autocast('cuda', dtype=torch.bfloat16)
+        else:
+            return torch.amp.autocast('cuda', dtype=torch.float32)
+    else:
+        return nullcontext()
+
+
+def cast_tensor(x, from_dtype=torch.float16, to_dtype=torch.float32):
+    return x.to(dtype=to_dtype) if x.dtype == from_dtype else x
+
+
+def cast_all(x, from_dtype=torch.float16, to_dtype=torch.float32):
+    if isinstance(x, torch.Tensor):
+        return cast_tensor(x, from_dtype=from_dtype, to_dtype=to_dtype)
+    else:
+        if isinstance(x, dict):
+            new_dict = {}
+            for k in x.keys():
+                new_dict[k] = cast_all(x[k], from_dtype=from_dtype, to_dtype=to_dtype)
+            return new_dict
+        elif isinstance(x, tuple):
+            return tuple(cast_all(y, from_dtype=from_dtype, to_dtype=to_dtype) for y in x)
+
+
+class CastToFloat(torch.nn.Module):
+    def __init__(self, mod):
+        super(CastToFloat, self).__init__()
+        self.mod = mod
+
+    def forward(self, x):
+        if torch.is_autocast_enabled() and x.dtype != torch.float32:
+            with torch.amp.autocast(x.device.type, enabled=False):
+                ret = self.mod.forward(x.to(torch.float32)).to(x.dtype)
+        else:
+            ret = self.mod.forward(x)
+        return ret
+
+
+class CastToFloatAll(torch.nn.Module):
+    def __init__(self, mod):
+        super(CastToFloatAll, self).__init__()
+        self.mod = mod
+
+    def forward(self, *args):
+        if torch.is_autocast_enabled():
+            from_dtype = args[0].dtype
+            with torch.amp.autocast(self.device.type, enabled=False):
+                ret = self.mod.forward(*cast_all(args, from_dtype=from_dtype, to_dtype=torch.float32))
+                return cast_all(ret, from_dtype=torch.float32, to_dtype=from_dtype)
+        else:
+            return self.mod.forward(*args)
+
+
+@contextmanager
+def monkeypatched(object, name, patch):
+    """Temporarily monkeypatches an object."""
+    pre_patched_value = getattr(object, name)
+    setattr(object, name, patch)
+    yield object
+    setattr(object, name, pre_patched_value)
+
+
+def maybe_cast_to_type(x: Any, type_: type) -> Any:
+    """Try to cast a value to int, if it fails, return the original value.
+
+    Args:
+        x (Any): The value to be casted.
+        type_ (type): The type to cast to, must be a callable.
+
+    Returns:
+        Any: The casted value or the original value if casting fails.
+    """
+    try:
+        return type_(x)
+    except Exception:
+        return x

nemo/utils/cloud.py

@@ -0,0 +1,178 @@
+# Copyright (c) 2020, 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 urllib.request
+from pathlib import Path
+from time import sleep
+
+from lightning.pytorch.plugins.environments import LightningEnvironment
+from lightning.pytorch.strategies import DDPStrategy, StrategyRegistry
+
+from nemo.utils import logging
+
+
+def maybe_download_from_cloud(url, filename, subfolder=None, cache_dir=None, refresh_cache=False) -> str:
+    """
+    Helper function to download pre-trained weights from the cloud
+    Args:
+        url: (str) URL of storage
+        filename: (str) what to download. The request will be issued to url/filename
+        subfolder: (str) subfolder within cache_dir. The file will be stored in cache_dir/subfolder. Subfolder can
+            be empty
+        cache_dir: (str) a cache directory where to download. If not present, this function will attempt to create it.
+            If None (default), then it will be $HOME/.cache/torch/NeMo
+        refresh_cache: (bool) if True and cached file is present, it will delete it and re-fetch
+
+    Returns:
+        If successful - absolute local path to the downloaded file
+        else - empty string
+    """
+    # try:
+    if cache_dir is None:
+        cache_location = Path.joinpath(Path.home(), ".cache/torch/NeMo")
+    else:
+        cache_location = cache_dir
+    if subfolder is not None:
+        destination = Path.joinpath(cache_location, subfolder)
+    else:
+        destination = cache_location
+
+    if not os.path.exists(destination):
+        os.makedirs(destination, exist_ok=True)
+
+    destination_file = Path.joinpath(destination, filename)
+
+    if os.path.exists(destination_file):
+        logging.info(f"Found existing object {destination_file}.")
+        if refresh_cache:
+            logging.info("Asked to refresh the cache.")
+            logging.info(f"Deleting file: {destination_file}")
+            os.remove(destination_file)
+        else:
+            logging.info(f"Re-using file from: {destination_file}")
+            return str(destination_file)
+    # download file
+    wget_uri = url + filename
+    logging.info(f"Downloading from: {wget_uri} to {str(destination_file)}")
+    # NGC links do not work everytime so we try and wait
+    i = 0
+    max_attempts = 3
+    while i < max_attempts:
+        i += 1
+        try:
+            urllib.request.urlretrieve(wget_uri, str(destination_file))
+            if os.path.exists(destination_file):
+                return destination_file
+            else:
+                return ""
+        except:
+            logging.info(f"Download from cloud failed. Attempt {i} of {max_attempts}")
+            sleep(0.05)
+            continue
+    raise ValueError("Not able to download url right now, please try again.")
+
+
+class SageMakerDDPStrategy(DDPStrategy):
+    """DDP strategy configured for AWS SageMaker distributed training."""
+
+    @property
+    def cluster_environment(self):
+        """Return a LightningEnvironment configured from SageMaker environment variables."""
+        env = LightningEnvironment()
+        env.world_size = lambda: int(os.environ["WORLD_SIZE"])
+        env.global_rank = lambda: int(os.environ["RANK"])
+        return env
+
+    @cluster_environment.setter
+    def cluster_environment(self, env):
+        """No-op setter to prevent Lightning from overriding the SageMaker environment."""
+        pass
+
+
+def initialize_sagemaker() -> None:
+    """
+    Helper function to initiate sagemaker with NeMo.
+    This function installs libraries that NeMo requires for the ASR toolkit + initializes sagemaker ddp.
+    """
+
+    logging.info("Registering SageMaker DDP strategy 'smddp'.")
+    StrategyRegistry.register(
+        name='smddp',
+        strategy=SageMakerDDPStrategy,
+        process_group_backend="smddp",
+        find_unused_parameters=False,
+    )
+
+    def _install_system_libraries() -> None:
+        import subprocess
+
+        logging.info("Installing system libraries: libsndfile1, ffmpeg")
+        try:
+            logging.info("Running apt-get update")
+            subprocess.run(
+                ["apt-get", "update"],
+                check=True,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+            logging.info("Running apt-get install for libsndfile1 and ffmpeg")
+            subprocess.run(
+                ["apt-get", "install", "-y", "libsndfile1", "ffmpeg"],
+                check=True,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+        except subprocess.CalledProcessError as e:
+            logging.error(
+                "Failed to install system libraries via apt-get (command=%s, returncode=%s): %s",
+                getattr(e, "cmd", None),
+                getattr(e, "returncode", None),
+                e,
+            )
+            logging.info("System libraries installed successfully.")
+        except Exception as e:
+            logging.error(f"Failed to install system libraries: {e}")
+
+    def _patch_torch_metrics() -> None:
+        """
+        Patches torchmetrics to not rely on internal state.
+        This is because sagemaker DDP overrides the `__init__` function of the modules to do automatic-partitioning.
+        """
+        from torchmetrics import Metric
+
+        def __new_hash__(self):
+            hash_vals = [self.__class__.__name__, id(self)]
+            return hash(tuple(hash_vals))
+
+        Metric.__hash__ = __new_hash__
+
+    logging.info("Patching torchmetrics hash function for SageMaker compatibility.")
+    _patch_torch_metrics()
+
+    if os.environ.get("RANK") and os.environ.get("WORLD_SIZE"):
+        import smdistributed.dataparallel.torch.distributed as dist
+
+        # has to be imported, as it overrides torch modules and such when DDP is enabled.
+        import smdistributed.dataparallel.torch.torch_smddp  # noqa: F401
+
+        logging.info("Initializing SageMaker distributed process group.")
+        dist.init_process_group()
+
+        if dist.get_local_rank():
+            _install_system_libraries()
+        logging.info("Waiting at barrier for all processes.")
+        return dist.barrier()  # wait for main process
+    _install_system_libraries()
+    return

nemo/utils/config_utils.py

@@ -0,0 +1,266 @@
+# Copyright (c) 2020, 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 copy
+import inspect
+from dataclasses import is_dataclass
+from typing import Dict, List, Optional
+
+from omegaconf import DictConfig, OmegaConf, open_dict
+
+from nemo.utils import logging
+
+
+def update_model_config(
+    model_cls: 'nemo.core.config.modelPT.NemoConfig', update_cfg: 'DictConfig', drop_missing_subconfigs: bool = True
+):
+    """
+    Helper class that updates the default values of a ModelPT config class with the values
+    in a DictConfig that mirrors the structure of the config class.
+
+    Assumes the `update_cfg` is a DictConfig (either generated manually, via hydra or instantiated via yaml/model.cfg).
+    This update_cfg is then used to override the default values preset inside the ModelPT config class.
+
+    If `drop_missing_subconfigs` is set, the certain sub-configs of the ModelPT config class will be removed, if
+    they are not found in the mirrored `update_cfg`. The following sub-configs are subject to potential removal:
+        -   `train_ds`
+        -   `validation_ds`
+        -   `test_ds`
+        -   `optim` + nested `sched`.
+
+    Args:
+        model_cls: A subclass of NemoConfig, that details in entirety all of the parameters that constitute
+            the NeMo Model.
+
+        update_cfg: A DictConfig that mirrors the structure of the NemoConfig data class. Used to update the
+            default values of the config class.
+
+        drop_missing_subconfigs: Bool which determins whether to drop certain sub-configs from the NemoConfig
+            class, if the corresponding sub-config is missing from `update_cfg`.
+
+    Returns:
+        A DictConfig with updated values that can be used to instantiate the NeMo Model along with supporting
+        infrastructure.
+    """
+    if not (is_dataclass(model_cls) or isinstance(model_cls, DictConfig)):
+        raise ValueError("`model_cfg` must be a dataclass or a structured OmegaConf object")
+
+    if not isinstance(update_cfg, DictConfig):
+        update_cfg = OmegaConf.create(update_cfg)
+
+    if is_dataclass(model_cls):
+        model_cls = OmegaConf.structured(model_cls)
+
+    # Update optional configs
+    model_cls = _update_subconfig(
+        model_cls, update_cfg, subconfig_key='train_ds', drop_missing_subconfigs=drop_missing_subconfigs
+    )
+    model_cls = _update_subconfig(
+        model_cls, update_cfg, subconfig_key='validation_ds', drop_missing_subconfigs=drop_missing_subconfigs
+    )
+    model_cls = _update_subconfig(
+        model_cls, update_cfg, subconfig_key='test_ds', drop_missing_subconfigs=drop_missing_subconfigs
+    )
+    model_cls = _update_subconfig(
+        model_cls, update_cfg, subconfig_key='optim', drop_missing_subconfigs=drop_missing_subconfigs
+    )
+
+    # Add optim and sched additional keys to model cls
+    model_cls = _add_subconfig_keys(model_cls, update_cfg, subconfig_key='optim')
+
+    # Perform full merge of model config class and update config
+    # Remove ModelPT artifact `target`
+    if 'target' in update_cfg.model:
+        # Assume artifact from ModelPT and pop
+        if 'target' not in model_cls.model:
+            with open_dict(update_cfg.model):
+                update_cfg.model.pop('target')
+
+    # Remove ModelPT artifact `nemo_version`
+    if 'nemo_version' in update_cfg.model:
+        # Assume artifact from ModelPT and pop
+        if 'nemo_version' not in model_cls.model:
+            with open_dict(update_cfg.model):
+                update_cfg.model.pop('nemo_version')
+
+    model_cfg = OmegaConf.merge(model_cls, update_cfg)
+
+    return model_cfg
+
+
+def _update_subconfig(
+    model_cfg: 'DictConfig', update_cfg: 'DictConfig', subconfig_key: str, drop_missing_subconfigs: bool
+):
+    """
+    Updates the NemoConfig DictConfig such that:
+    1)  If the sub-config key exists in the `update_cfg`, but does not exist in ModelPT config:
+        - Add the sub-config from update_cfg to ModelPT config
+
+    2) If the sub-config key does not exist in `update_cfg`, but exists in ModelPT config:
+        - Remove the sub-config from the ModelPT config; iff the `drop_missing_subconfigs` flag is set.
+
+    Args:
+        model_cfg: A DictConfig instantiated from the NemoConfig subclass.
+        update_cfg: A DictConfig that mirrors the structure of `model_cfg`, used to update its default values.
+        subconfig_key: A str key used to check and update the sub-config.
+        drop_missing_subconfigs: A bool flag, whether to allow deletion of the NemoConfig sub-config,
+            if its mirror sub-config does not exist in the `update_cfg`.
+
+    Returns:
+        The updated DictConfig for the NemoConfig
+    """
+    with open_dict(model_cfg.model):
+        # If update config has the key, but model cfg doesnt have the key
+        # Add the update cfg subconfig to the model cfg
+        if subconfig_key in update_cfg.model and subconfig_key not in model_cfg.model:
+            model_cfg.model[subconfig_key] = update_cfg.model[subconfig_key]
+
+        # If update config does not the key, but model cfg has the key
+        # Remove the model cfg subconfig in order to match layout of update cfg
+        if subconfig_key not in update_cfg.model and subconfig_key in model_cfg.model:
+            if drop_missing_subconfigs:
+                model_cfg.model.pop(subconfig_key)
+
+    return model_cfg
+
+
+def _add_subconfig_keys(model_cfg: 'DictConfig', update_cfg: 'DictConfig', subconfig_key: str):
+    """
+    For certain sub-configs, the default values specified by the NemoConfig class is insufficient.
+    In order to support every potential value in the merge between the `update_cfg`, it would require
+    explicit definition of all possible cases.
+
+    An example of such a case is Optimizers, and their equivalent Schedulers. All optimizers share a few basic
+    details - such as name and lr, but almost all require additional parameters - such as weight decay.
+    It is impractical to create a config for every single optimizer + every single scheduler combination.
+
+    In such a case, we perform a dual merge. The Optim and Sched Dataclass contain the bare minimum essential
+    components. The extra values are provided via update_cfg.
+
+    In order to enable the merge, we first need to update the update sub-config to incorporate the keys,
+    with dummy temporary values (merge update config with model config). This is done on a copy of the
+    update sub-config, as the actual override values might be overriden by the NemoConfig defaults.
+
+    Then we perform a merge of this temporary sub-config with the actual override config in a later step
+    (merge model_cfg with original update_cfg, done outside this function).
+
+    Args:
+        model_cfg: A DictConfig instantiated from the NemoConfig subclass.
+        update_cfg: A DictConfig that mirrors the structure of `model_cfg`, used to update its default values.
+        subconfig_key: A str key used to check and update the sub-config.
+
+    Returns:
+        A ModelPT DictConfig with additional keys added to the sub-config.
+    """
+    with open_dict(model_cfg.model):
+        # Create copy of original model sub config
+        if subconfig_key in update_cfg.model:
+            if subconfig_key not in model_cfg.model:
+                # create the key as a placeholder
+                model_cfg.model[subconfig_key] = None
+
+            subconfig = copy.deepcopy(model_cfg.model[subconfig_key])
+            update_subconfig = copy.deepcopy(update_cfg.model[subconfig_key])
+
+            # Add the keys and update temporary values, will be updated during full merge
+            subconfig = OmegaConf.merge(update_subconfig, subconfig)
+            # Update sub config
+            model_cfg.model[subconfig_key] = subconfig
+
+    return model_cfg
+
+
+def assert_dataclass_signature_match(
+    cls: 'class_type',
+    datacls: 'dataclass',
+    ignore_args: Optional[List[str]] = None,
+    remap_args: Optional[Dict[str, str]] = None,
+):
+    """
+    Analyses the signature of a provided class and its respective data class,
+    asserting that the dataclass signature matches the class __init__ signature.
+
+    Note:
+        This is not a value based check. This function only checks if all argument
+        names exist on both class and dataclass and logs mismatches.
+
+    Args:
+        cls: Any class type - but not an instance of a class. Pass type(x) where x is an instance
+            if class type is not easily available.
+        datacls: A corresponding dataclass for the above class.
+        ignore_args: (Optional) A list of string argument names which are forcibly ignored,
+            even if mismatched in the signature. Useful when a dataclass is a superset of the
+            arguments of a class.
+        remap_args: (Optional) A dictionary, mapping an argument name that exists (in either the
+            class or its dataclass), to another name. Useful when argument names are mismatched between
+            a class and its dataclass due to indirect instantiation via a helper method.
+
+    Returns:
+        A tuple containing information about the analysis:
+        1) A bool value which is True if the signatures matched exactly / after ignoring values.
+            False otherwise.
+        2) A set of arguments names that exist in the class, but *do not* exist in the dataclass.
+            If exact signature match occurs, this will be None instead.
+        3) A set of argument names that exist in the data class, but *do not* exist in the class itself.
+            If exact signature match occurs, this will be None instead.
+    """
+    class_sig = inspect.signature(cls.__init__)
+
+    class_params = dict(**class_sig.parameters)
+    class_params.pop('self')
+
+    dataclass_sig = inspect.signature(datacls)
+
+    dataclass_params = dict(**dataclass_sig.parameters)
+    dataclass_params.pop("_target_", None)
+
+    class_params = set(class_params.keys())
+    dataclass_params = set(dataclass_params.keys())
+
+    if remap_args is not None:
+        for original_arg, new_arg in remap_args.items():
+            if original_arg in class_params:
+                class_params.remove(original_arg)
+                class_params.add(new_arg)
+                logging.info(f"Remapped {original_arg} -> {new_arg} in {cls.__name__}")
+
+            if original_arg in dataclass_params:
+                dataclass_params.remove(original_arg)
+                dataclass_params.add(new_arg)
+                logging.info(f"Remapped {original_arg} -> {new_arg} in {datacls.__name__}")
+
+    if ignore_args is not None:
+        ignore_args = set(ignore_args)
+
+        class_params = class_params - ignore_args
+        dataclass_params = dataclass_params - ignore_args
+        logging.info(f"Removing ignored arguments - {ignore_args}")
+
+    intersection = set.intersection(class_params, dataclass_params)
+    subset_cls = class_params - intersection
+    subset_datacls = dataclass_params - intersection
+
+    if (len(class_params) != len(dataclass_params)) or len(subset_cls) > 0 or len(subset_datacls) > 0:
+        logging.error(f"Class {cls.__name__} arguments do not match " f"Dataclass {datacls.__name__}!")
+
+        if len(subset_cls) > 0:
+            logging.error(f"Class {cls.__name__} has additional arguments :\n" f"{subset_cls}")
+
+        if len(subset_datacls):
+            logging.error(f"Dataclass {datacls.__name__} has additional arguments :\n{subset_datacls}")
+
+        return False, subset_cls, subset_datacls
+
+    else:
+        return True, None, None

nemo/utils/data_utils.py

@@ -0,0 +1,390 @@
+# Copyright (c) 2022, 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.
+
+"""Utility functions for handling data operations, including datastore access and caching."""
+
+import os
+import pathlib
+import shutil
+import subprocess
+from functools import lru_cache
+from typing import Any, Callable, Dict, Iterable, Tuple
+from urllib.parse import urlparse
+
+try:
+    from nemo import __version__ as NEMO_VERSION
+except ImportError:
+    NEMO_VERSION = 'git'
+
+from nemo import constants
+from nemo.utils import logging
+from nemo.utils.nemo_logging import LogMode
+
+try:
+    from lhotse.serialization import open_best as lhotse_open_best
+
+    LHOTSE_AVAILABLE = True
+except ImportError:
+    LHOTSE_AVAILABLE = False
+
+
+def resolve_cache_dir() -> pathlib.Path:
+    """
+    Utility method to resolve a cache directory for NeMo that can be overriden by an environment variable.
+
+    Example:
+        NEMO_CACHE_DIR="~/nemo_cache_dir/" python nemo_example_script.py
+
+    Returns:
+        A Path object, resolved to the absolute path of the cache directory. If no override is provided,
+        uses an inbuilt default which adapts to nemo versions strings.
+    """
+    override_dir = os.environ.get(constants.NEMO_ENV_CACHE_DIR, "")
+    if override_dir == "":
+        path = pathlib.Path.joinpath(pathlib.Path.home(), f'.cache/torch/NeMo/NeMo_{NEMO_VERSION}')
+    else:
+        path = pathlib.Path(override_dir).resolve()
+    return path
+
+
+def is_datastore_path(path) -> bool:
+    """Check if a path is from a data object store."""
+    try:
+        result = urlparse(path)
+        return bool(result.scheme) and bool(result.netloc)
+    except AttributeError:
+        return False
+
+
+def is_tarred_path(path) -> bool:
+    """Check if a path is for a tarred file."""
+    return path.endswith('.tar')
+
+
+def is_datastore_cache_shared() -> bool:
+    """Check if store cache is shared."""
+    # Assume cache is shared by default, e.g., as in resolve_cache_dir (~/.cache)
+    cache_shared = int(os.environ.get(constants.NEMO_ENV_DATA_STORE_CACHE_SHARED, 1))
+
+    if cache_shared == 0:
+        return False
+    if cache_shared == 1:
+        return True
+    raise ValueError(f'Unexpected value of env {constants.NEMO_ENV_DATA_STORE_CACHE_SHARED}')
+
+
+def ais_cache_base() -> str:
+    """Return path to local cache for AIS."""
+    override_dir = os.environ.get(constants.NEMO_ENV_DATA_STORE_CACHE_DIR, "")
+    if override_dir == "":
+        cache_dir = resolve_cache_dir().as_posix()
+    else:
+        cache_dir = pathlib.Path(override_dir).resolve().as_posix()
+
+    if cache_dir.endswith(NEMO_VERSION):
+        # Prevent re-caching dataset after upgrading NeMo
+        cache_dir = os.path.dirname(cache_dir)
+    return os.path.join(cache_dir, 'ais')
+
+
+def ais_endpoint() -> str:
+    """Get configured AIS endpoint."""
+    return os.getenv('AIS_ENDPOINT')
+
+
+def bucket_and_object_from_uri(uri: str) -> Tuple[str, str]:
+    """Parse a path to determine bucket and object path.
+
+    Args:
+        uri: Full path to an object on an object store
+
+    Returns:
+        Tuple of strings (bucket_name, object_path)
+    """
+    if not is_datastore_path(uri):
+        raise ValueError(f'Provided URI is not a valid store path: {uri}')
+    uri_parts = pathlib.PurePath(uri).parts
+    bucket = uri_parts[1]
+    object_path = pathlib.PurePath(*uri_parts[2:])
+
+    return str(bucket), str(object_path)
+
+
+def ais_endpoint_to_dir(endpoint: str) -> str:
+    """Convert AIS endpoint to a valid dir name.
+    Used to build cache location.
+
+    Args:
+        endpoint: AIStore endpoint in format https://host:port
+
+    Returns:
+        Directory formed as `host/port`.
+    """
+    result = urlparse(endpoint)
+    if not result.hostname or not result.port:
+        raise ValueError(f"Unexpected format for ais endpoint: {endpoint}")
+    return os.path.join(result.hostname, str(result.port))
+
+
+@lru_cache(maxsize=1)
+def ais_binary() -> str:
+    """Return location of `ais` binary if available."""
+    path = shutil.which('ais')
+
+    if path is not None:
+        logging.debug('Found AIS binary at %s', path)
+        return path
+
+    # Double-check if it exists at the default path
+    default_path = '/usr/local/bin/ais'
+    if os.path.isfile(default_path):
+        logging.info('ais available at the default path: %s', default_path, mode=LogMode.ONCE)
+        return default_path
+    logging.warning(
+        f'AIS binary not found with `which ais` and at the default path {default_path}.', mode=LogMode.ONCE
+    )
+    return None
+
+
+def datastore_path_to_local_path(store_path: str) -> str:
+    """Convert a data store path to a path in a local cache.
+
+    Args:
+        store_path: a path to an object on an object store
+
+    Returns:
+        Path to the same object in local cache.
+    """
+    if is_datastore_path(store_path):
+        endpoint = ais_endpoint()
+        if not endpoint:
+            raise RuntimeError(f'AIS endpoint not set, cannot resolve {store_path}')
+
+        local_ais_cache = os.path.join(ais_cache_base(), ais_endpoint_to_dir(endpoint))
+        store_bucket, store_object = bucket_and_object_from_uri(store_path)
+        local_path = os.path.join(local_ais_cache, store_bucket, store_object)
+    else:
+        raise ValueError(f'Unexpected store path format: {store_path}')
+
+    return local_path
+
+
+def open_datastore_object_with_binary(path: str, num_retries: int = 5):
+    """Open a datastore object and return a file-like object.
+
+    Args:
+        path: path to an object
+        num_retries: number of retries if the get command fails with ais binary,
+            as AIS Python SDK has its own retry mechanism
+
+    Returns:
+        File-like object that supports read()
+    """
+
+    if is_datastore_path(path):
+        endpoint = ais_endpoint()
+        if endpoint is None:
+            raise RuntimeError(f'AIS endpoint not set, cannot resolve {path}')
+
+        binary = ais_binary()
+
+        if not binary:
+            raise RuntimeError(
+                f"AIS binary is not found, cannot resolve {path}. "
+                "Please either install it or install Lhotse with `pip install lhotse`.\n"
+                "Lhotse's native open_best supports AIS Python SDK, "
+                "which is the recommended way to operate with the data from AIStore.\n"
+                "See AIS binary installation instructions at "
+                "https://github.com/NVIDIA/aistore?tab=readme-ov-file#install-from-release-binaries.\n"
+            )
+
+        cmd = [binary, 'get', path, '-']
+
+        done = False
+
+        for _ in range(num_retries):
+            with subprocess.Popen(
+                cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=False  # bytes mode
+            ) as proc:
+                stream = proc.stdout
+                if stream.peek(1):
+                    done = True
+                    return stream
+
+        if not done:
+            with subprocess.Popen(
+                cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=False
+            ) as proc:
+                error = proc.stderr.read().decode("utf-8", errors="ignore").strip()
+            raise ValueError(
+                f"{path} couldn't be opened with AIS binary "
+                f"after {num_retries} attempts because of the following exception: {error}"
+            )
+    return None
+
+
+def open_best(path: str, mode: str = "rb"):
+    """Open a file using the best available method (Lhotse, datastore binary, or standard open).
+
+    Args:
+        path: path to the file or datastore object
+        mode: file opening mode (default: "rb")
+
+    Returns:
+        File-like object
+    """
+    if LHOTSE_AVAILABLE:
+        return lhotse_open_best(path, mode=mode)
+    if is_datastore_path(path):
+        return open_datastore_object_with_binary(path)
+    return open(path, mode=mode, encoding='utf-8' if 'b' not in mode else None)
+
+
+def get_datastore_object(path: str, force: bool = False, num_retries: int = 5) -> str:
+    """Download an object from a store path and return the local path.
+    If the input `path` is a local path, then nothing will be done, and
+    the original path will be returned.
+
+    Args:
+        path: path to an object
+        force: force download, even if a local file exists
+        num_retries: number of retries if the get command fails with ais binary,
+            as AIS Python SDK has its own retry mechanism
+
+    Returns:
+        Local path of the object.
+    """
+    if is_datastore_path(path):
+
+        local_path = datastore_path_to_local_path(store_path=path)
+
+        if not os.path.isfile(local_path) or force:
+            # Either we don't have the file in cache or we force download it
+            # Enhancement: if local file is present, check some tag and compare against remote
+            local_dir = os.path.dirname(local_path)
+            if not os.path.isdir(local_dir):
+                os.makedirs(local_dir, exist_ok=True)
+
+            with open(local_path, 'wb') as f:
+                f.write(open_best(path).read(), num_retries=num_retries)
+
+        return local_path
+
+    # Assume the file is local
+    return path
+
+
+class DataStoreObject:
+    """A simple class for handling objects in a data store.
+    Currently, this class supports objects on AIStore.
+
+    Args:
+        store_path: path to a store object
+        local_path: path to a local object, may be used to upload local object to store
+        get: get the object from a store
+    """
+
+    def __init__(self, store_path: str, local_path: str = None, get: bool = False):
+        if local_path is not None:
+            raise NotImplementedError('Specifying a local path is currently not supported.')
+
+        self._store_path = store_path
+        self._local_path = local_path
+
+        if get:
+            self.get()
+
+    @property
+    def store_path(self) -> str:
+        """Return store path of the object."""
+        return self._store_path
+
+    @property
+    def local_path(self) -> str:
+        """Return local path of the object."""
+        return self._local_path
+
+    def get(self, force: bool = False) -> str:
+        """Get an object from the store to local cache and return the local path.
+
+        Args:
+            force: force download, even if a local file exists
+
+        Returns:
+            Path to a local object.
+        """
+        if not self.local_path:
+            # Assume the object needs to be downloaded
+            self._local_path = get_datastore_object(self.store_path, force=force)
+        return self.local_path
+
+    def put(self, force: bool = False) -> str:
+        """Push to remote and return the store path
+
+        Args:
+            force: force download, even if a local file exists
+
+        Returns:
+            Path to a (remote) object object on the object store.
+        """
+        raise NotImplementedError()
+
+    def __str__(self):
+        """Return a human-readable description of the object."""
+        description = f'{type(self)}: store_path={self.store_path}, local_path={self.local_path}'
+        return description
+
+
+def datastore_object_get(store_object: DataStoreObject) -> bool:
+    """A convenience wrapper for multiprocessing.imap.
+
+    Args:
+        store_object: An instance of DataStoreObject
+
+    Returns:
+        True if get() returned a path.
+    """
+    return store_object.get() is not None
+
+
+def wds_url_opener(  # pylint: disable=unused-argument
+    data: Iterable[Dict[str, Any]],
+    handler: Callable[[Exception], bool],
+    **kw: Dict[str, Any],
+):
+    """
+    Open URLs and yield a stream of url+stream pairs.
+    This is a workaround to use lhotse's open_best instead of webdataset's default url_opener.
+    webdataset's default url_opener uses gopen, which does not support opening datastore paths.
+
+    Args:
+        data: Iterator over dict(url=...).
+        handler: Exception handler.
+        **kw: Keyword arguments for gopen.gopen (unused, kept for API compatibility).
+
+    Yields:
+        A stream of url+stream pairs.
+    """
+    for sample in data:
+        assert isinstance(sample, dict), sample
+        assert "url" in sample
+        url = sample["url"]
+        try:
+            stream = open_best(url, mode="rb")
+            sample.update(stream=stream)
+            yield sample
+        except Exception as exn:  # pylint: disable=broad-exception-caught
+            if handler(exn):
+                continue
+            break

nemo/utils/debug_hook.py

@@ -0,0 +1,202 @@
+# Copyright (c) 2022, 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 torch
+
+
+def get_forward_hook(name, trainer, rank, logger, dump_to_file=False):
+    """
+    A forward hook to dump all of the module input and output norms.
+    It is called every time after forward() has computed an output.
+    Only float type input/output tensor norms are computed.
+    For more details about the forward hook, check
+    https://pytorch.org/docs/stable/generated/torch.nn.modules.module.register_module_forward_hook.html
+
+    Args:
+        name: tensor name
+        trainer: PTL trainer
+        rank: worker rank
+        logger: PTL log function
+        dump_to_file:  wether dump the csv file to the disk
+    """
+    if dump_to_file:
+        os.makedirs('debug_info', exist_ok=True)
+        fp = open(f'debug_info/forward_{name}_rank{rank}.txt', 'w')
+        header = False
+
+    def forward_hook(module, inputs, outputs):
+        nonlocal header
+        nonlocal fp
+        if trainer.training:
+            values = []
+            headers = []
+            for n, i in enumerate(inputs):
+                if isinstance(i, torch.Tensor) and (
+                    i.dtype == torch.float or i.dtype == torch.half or i.dtype == torch.bfloat16
+                ):
+                    if not header:
+                        headers.append('input')
+                    input_norm = i.data.norm()
+                    values.append(f'{input_norm}')
+                    logger(f'debug_info_forward/{name}_rank{rank}_input{n}', input_norm)
+            if isinstance(outputs, tuple):
+                for n, i in enumerate(outputs):
+                    if isinstance(i, torch.Tensor) and (
+                        i.dtype == torch.float or i.dtype == torch.half or i.dtype == torch.bfloat16
+                    ):
+                        if not header:
+                            headers.append('output')
+                        output_norm = i.data.norm()
+                        values.append(f'{output_norm}')
+                        logger(f'debug_info_forward/{name}_rank{rank}_output{n}', output_norm)
+            else:
+                headers.append('output')
+                values.append(f'{outputs.data.norm()}')
+            values.append(f'{trainer.global_step}')
+            if not header:
+                headers.append('step')
+                fp.write(','.join(headers) + '\n')
+                header = True
+            fp.write(','.join(values) + '\n')
+        fp.flush()
+
+    return forward_hook
+
+
+def get_backward_hook(name, trainer, rank, logger, dump_to_file=False):
+    """
+    A backward hook to dump all of the module input and output grad norms.
+    The hook is called every time gradients with respect to module inputs are computed.
+    Only float type input/output grad tensor norms are computed.
+    For more details about the backward hook, check
+    https://pytorch.org/docs/stable/generated/torch.nn.modules.module.register_module_full_backward_hook.html
+
+    Args:
+        name: tensor name
+        trainer: PTL trainer
+        rank: worker rank
+        logger: PTL log function
+        dump_to_file:  wether dump the csv file to the disk
+    """
+    if dump_to_file:
+        os.makedirs('debug_info', exist_ok=True)
+        fp = open(f'debug_info/backward_{name}_rank{rank}.txt', 'w')
+        header = False
+
+    def backward_hook(module, inputs, outputs):
+        nonlocal header
+        nonlocal fp
+        if trainer.training:
+            values = []
+            headers = []
+            for n, i in enumerate(inputs):
+                if isinstance(i, torch.Tensor) and (
+                    i.dtype == torch.float or i.dtype == torch.half or i.dtype == torch.bfloat16
+                ):
+                    if not header:
+                        headers.append('input')
+                    input_norm = i.data.norm()
+                    values.append(f'{input_norm}')
+                    logger(f'debug_info_backward/{name}_rank{rank}_input{n}', input_norm)
+            if isinstance(outputs, tuple):
+                for n, i in enumerate(outputs):
+                    if isinstance(i, torch.Tensor) and (
+                        i.dtype == torch.float or i.dtype == torch.half or i.dtype == torch.bfloat16
+                    ):
+                        if not header:
+                            headers.append('output')
+                        output_norm = i.data.norm()
+                        values.append(f'{output_norm}')
+                        logger(f'debug_info_backward/{name}_rank{rank}_output{n}', output_norm)
+            else:
+                headers.append('output')
+                values.append(f'{outputs.data.norm()}')
+            values.append(f'{trainer.global_step}')
+            if not header:
+                headers.append('step')
+                fp.write(','.join(headers) + '\n')
+                header = True
+            fp.write(','.join(values) + '\n')
+        fp.flush()
+
+    return backward_hook
+
+
+def get_tensor_hook(module, name, trainer, rank, logger, dump_to_file=False):
+    """
+    A tensor hook to dump tensor weight norms and grad norms at the end of each backward step.
+    For more details about the tensor hook, check
+    https://pytorch.org/docs/stable/generated/torch.Tensor.register_hook.html
+
+    Args:
+        module: the model module
+        name: tensor name
+        trainer: PTL trainer
+        rank: worker rank
+        logger: PTL log function
+        dump_to_file:  wether dump the csv file to the disk
+    """
+    if dump_to_file:
+        os.makedirs('debug_info', exist_ok=True)
+        fp = open(f'debug_info/tensor_{name}_rank{rank}.csv', 'w')
+        header = False
+
+    def tensor_hook(grad):
+        nonlocal header
+        nonlocal fp
+        values = []
+        headers = []
+
+        weight = module.get_parameter(name)
+        weight_norm = weight.data.norm()
+        grad_norm = grad.data.norm()
+        logger(f'debug_info_tensors/{name}_rank{rank}_grad_norm', grad_norm)
+        logger(f'debug_info_tensors/{name}_rank{rank}_weight_norm', weight_norm)
+        values.append(f'{weight_norm}')
+        values.append(f'{grad_norm}')
+        values.append(f'{trainer.global_step}')
+        if dump_to_file:
+            if not header:
+                headers.append('weight')
+                headers.append('grad')
+                headers.append('step')
+                fp.write(','.join(headers) + '\n')
+                header = True
+            fp.write(','.join(values) + '\n')
+            fp.flush()
+        return grad
+
+    return tensor_hook
+
+
+def register_debug_hooks(module, trainer, logger, dump_to_file=False):
+    """
+    Register debug hooks. It can
+    1. track the module forward step input/ouput norm
+    2. track the module backward step input/output grad norm
+    3. track the parameter weight norm and grad norm.
+    """
+    # default rank 0
+    rank = 0
+    if torch.distributed.is_initialized():
+        rank = torch.distributed.get_rank()
+    for name, tensor in module.named_parameters():
+        if name != '':
+            tensor.register_hook(get_tensor_hook(module, name, trainer, rank, logger, dump_to_file))
+    for name, layer in module.named_modules():
+        if name != '':
+            layer.register_forward_hook(get_forward_hook(name, trainer, rank, logger, dump_to_file))
+            layer.register_full_backward_hook(get_backward_hook(name, trainer, rank, logger, dump_to_file))

nemo/utils/decorators/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2020, 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.decorators.deprecated import deprecated, deprecated_warning
+from nemo.utils.decorators.experimental import experimental
+from nemo.utils.decorators.port_docs import add_port_docs

nemo/utils/decorators/deprecated.py

@@ -0,0 +1,96 @@
+# Copyright (c) 2020, 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.
+
+
+__all__ = [
+    'deprecated',
+]
+
+import functools
+import inspect
+import time
+import wrapt
+
+from nemo.utils import logging
+
+# Remember which deprecation warnings have been printed already.
+_PRINTED_WARNING = {}
+
+
+def deprecated(wrapped=None, version=None, explanation=None, wait_seconds=0):
+    """
+    Decorator which can be used for indicating that a function/class is deprecated and going to be removed.
+    Tracks down which function/class printed the warning and will print it only once per call.
+
+    Args:
+      version: Version in which the function/class will be removed (optional).
+      explanation: Additional explanation, e.g. "Please, ``use another_function`` instead." (optional).
+      wait_seconds: Sleep for a few seconds after the deprecation message appears in case it gets drowned
+      with subsequent logging messages.
+    """
+
+    if wrapped is None:
+        return functools.partial(deprecated, version=version, explanation=explanation, wait_seconds=wait_seconds)
+
+    @wrapt.decorator
+    def wrapper(wrapped, instance, args, kwargs):
+        # Check if we already warned about that function.
+        if wrapped.__name__ not in _PRINTED_WARNING.keys():
+            # Add to list so we won't print it again.
+            _PRINTED_WARNING[wrapped.__name__] = True
+
+            # Prepare the warning message.
+            entity_name = "Class" if inspect.isclass(wrapped) else "Function"
+            msg = f"{entity_name} ``{wrapped.__name__}`` is deprecated."
+
+            # Optionally, add version and explanation.
+            if version is not None:
+                msg = f"{msg} It is going to be removed in the {version} version."
+
+            if explanation is not None:
+                msg = f"{msg} {explanation}"
+
+            # Display the deprecated warning.
+            logging.warning(msg)
+            if wait_seconds > 0:
+                logging.warning(f'Waiting for {wait_seconds} seconds before this message disappears')
+                time.sleep(wait_seconds)
+
+        # Call the function.
+        return wrapped(*args, **kwargs)
+
+    return wrapper(wrapped)
+
+
+def deprecated_warning(old_method=None, new_method=None, wait_seconds=2):
+    """
+    Function which can be used for indicating that a function/class is deprecated and going to be removed.
+
+    Args:
+      old_method: Name of deprecated class/function.
+      new_method: Name of new class/function to use.
+      wait_seconds: Sleep for a few seconds after the deprecation message appears in case it gets drowned
+      with subsequent logging messages.
+    """
+
+    # Create a banner
+    if new_method is not None:
+        msg = f"*****  {old_method} is deprecated. Please, use {new_method} instead.  *****"
+    else:
+        msg = f"*****  {old_method} is deprecated and will be removed soon.  *****"
+    banner = '\n'.join(['*' * len(msg)] * 2 + [msg] + ['*' * len(msg)] * 2)
+
+    logging.warning(f"\n\n{banner}\n")
+    logging.warning(f"Waiting for {wait_seconds} seconds before this message disappears.")
+    time.sleep(wait_seconds)

nemo/utils/decorators/experimental.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2020, 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.
+
+
+__all__ = ['experimental']
+
+
+import wrapt
+
+from nemo.utils import logging
+
+
+@wrapt.decorator
+def experimental(wrapped, instance, args, kwargs):
+    logging.warning(f"`{wrapped}` is experimental and not ready for production yet. Use at your own risk.")
+    return wrapped(*args, **kwargs)

nemo/utils/decorators/port_docs.py

@@ -0,0 +1,90 @@
+# Copyright (c) 2020, 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.
+
+
+# The "add_port_docs" decorator is needed to nicely generate neural types in Sphynx for input and output ports
+
+__all__ = [
+    'add_port_docs',
+]
+
+import functools
+import sys
+
+import wrapt
+
+
+def _normalize_docstring(docstring):
+    """Normalizes the docstring.
+    Replaces tabs with spaces, removes leading and trailing blanks lines, and
+    removes any indentation.
+    Copied from PEP-257:
+    https://www.python.org/dev/peps/pep-0257/#handling-docstring-indentation
+    Args:
+        docstring: the docstring to normalize
+    Returns:
+        The normalized docstring
+    """
+    if not docstring:
+        return ''
+    # Convert tabs to spaces (following the normal Python rules)
+    # and split into a list of lines:
+    lines = docstring.expandtabs().splitlines()
+    # Determine minimum indentation (first line doesn't count):
+    # (we use sys.maxsize because sys.maxint doesn't exist in Python 3)
+    indent = sys.maxsize
+    for line in lines[1:]:
+        stripped = line.lstrip()
+        if stripped:
+            indent = min(indent, len(line) - len(stripped))
+    # Remove indentation (first line is special):
+    trimmed = [lines[0].strip()]
+    if indent < sys.maxsize:
+        for line in lines[1:]:
+            trimmed.append(line[indent:].rstrip())
+    # Strip off trailing and leading blank lines:
+    while trimmed and not trimmed[-1]:
+        trimmed.pop()
+    while trimmed and not trimmed[0]:
+        trimmed.pop(0)
+    # Return a single string:
+    return '\n'.join(trimmed)
+
+
+def add_port_docs(wrapped=None, instance=None, value=''):
+    if wrapped is None:
+        return functools.partial(add_port_docs, value=value)
+
+    @wrapt.decorator
+    def wrapper(wrapped, instance=None, args=None, kwargs=None):
+        return wrapped(*args, **kwargs)
+
+    decorated = wrapper(wrapped)
+    try:
+        port_2_ntype = decorated(instance)
+    except:
+        port_2_ntype = None
+
+    port_description = ""
+    if port_2_ntype is not None:
+        for port, ntype in port_2_ntype.items():
+            port_description += "* *" + port + "* : " + str(ntype)
+            port_description += "\n\n"
+
+    __doc__ = _normalize_docstring(wrapped.__doc__) + '\n\n' + str(port_description)
+    __doc__ = _normalize_docstring(__doc__)
+
+    wrapt.FunctionWrapper.__setattr__(decorated, "__doc__", __doc__)
+
+    return decorated

nemo/utils/dependency.py

@@ -0,0 +1,82 @@
+# Copyright (c) 2026, 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 importlib
+import importlib.util
+import types
+
+
+def is_module_available(*modules: str) -> bool:
+    """Check whether the given modules are installed without importing them.
+
+    This is safer than a ``try: import ...`` block because some packages have
+    side effects at import time (e.g., changing the multiprocessing start
+    method).  Use this for lightweight availability checks such as test-skip
+    decorators or conditional registration.
+
+    Args:
+        *modules: One or more top-level module names to check.
+
+    Returns:
+        ``True`` if **all** listed modules are found, ``False`` otherwise.
+    """
+    return all(importlib.util.find_spec(m) is not None for m in modules)
+
+
+def assert_optional_dependency_available(module_name: str, *, pip_name: str | None = None) -> None:
+    """Raise an ``ImportError`` if *module_name* is not installed.
+
+    Unlike :func:`import_optional_dependency` this does **not** import the
+    module — it only checks availability via :func:`is_module_available` and
+    raises with an actionable install hint on failure.  Use this for early
+    fail-fast checks (e.g., at the top of a CLI entry-point or ``__init__``).
+
+    Args:
+        module_name: The module to check (e.g. ``"lhotse"``).
+        pip_name: The pip install name if it differs from *module_name*.
+            When *None*, the top-level package name is used.
+
+    Raises:
+        ImportError: If the module is not found.
+    """
+    if not is_module_available(module_name):
+        install_name = pip_name if pip_name is not None else module_name.split(".")[0]
+        raise ImportError(
+            f"Optional dependency '{module_name}' is not installed. " f"Install it with:  pip install {install_name}"
+        )
+
+
+def import_optional_dependency(module_name: str, *, pip_name: str | None = None) -> types.ModuleType:
+    """Import an optional dependency, raising a clear error if it is not installed.
+
+    Args:
+        module_name: The module to import (e.g. ``"lhotse"`` or ``"torchaudio.transforms"``).
+        pip_name: The pip install name if it differs from *module_name*
+            (e.g. ``pip_name="Cython"`` for ``module_name="cython"``).
+            When *None*, the top-level package name is used.
+
+    Returns:
+        The imported module.
+
+    Raises:
+        ImportError: If the module cannot be imported, with an actionable
+            install hint.
+    """
+    try:
+        return importlib.import_module(module_name)
+    except ImportError:
+        install_name = pip_name if pip_name is not None else module_name.split(".")[0]
+        raise ImportError(
+            f"Optional dependency '{module_name}' is not installed. " f"Install it with:  pip install {install_name}"
+        ) from None

nemo/utils/distributed.py

@@ -0,0 +1,147 @@
+# Copyright (c) 2022, 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 contextlib
+import os
+import tempfile
+
+import torch
+import torch.distributed as dist
+
+from nemo.utils import logging
+from nemo.utils.get_rank import is_global_rank_zero
+
+try:
+    from megatron.core import parallel_state
+
+    HAVE_MEGATRON_CORE = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_MEGATRON_CORE = False
+
+
+def initialize_distributed(args, backend='nccl'):
+    """Initialize torch.distributed."""
+    # Get local rank in case it is provided.
+    local_rank = args.local_rank
+
+    # Get rank and world size.
+    rank = int(os.getenv('RANK', '0'))
+    world_size = int(os.getenv("WORLD_SIZE", '1'))
+
+    logging.info(
+        f'Initializing torch.distributed with local_rank: {local_rank}, rank: {rank}, world_size: {world_size}'
+    )
+
+    # Set the device id.
+    device = rank % torch.cuda.device_count()
+    if local_rank is not None:
+        device = local_rank
+    torch.cuda.set_device(device)
+
+    # Call the init process.
+    init_method = 'tcp://'
+    master_ip = os.getenv('MASTER_ADDR', 'localhost')
+    master_port = os.getenv('MASTER_PORT', '6000')
+    init_method += master_ip + ':' + master_port
+    torch.distributed.init_process_group(backend=backend, world_size=world_size, rank=rank, init_method=init_method)
+    return local_rank, rank, world_size
+
+
+def gather_objects(partial_results_list, main_rank=None):
+    """
+    Collect objects (e.g., results) from all GPUs.
+    Useful for inference over multiple GPUs with DDP.
+
+    Use main_rank to specify which rank will be used to gather results.
+    This allows to continue execution on the main_rank only after the gather.
+
+    Args:
+        partial_results_list: list of partial results from each GPU
+        main_rank: rank of the main process to collect results from all GPUs (useful for collecting results in a target rank)
+
+
+    Example:
+        predictions = gather_objects(predictions,main_rank=0)
+        # all but rank 0 will return None
+        if predictions is None:
+            return
+
+        # from here only rank 0 should contiue
+        pickle.dump(predictions, open(output_fname, "wb"))
+    """
+    # do not fail when DDP is not initialized
+    if not parallel_state.is_initialized():
+        return partial_results_list
+
+    rank = parallel_state.get_data_parallel_rank()
+    world_size = parallel_state.get_data_parallel_world_size()
+    # return input when no DDP is used
+    if world_size == 1:
+        return partial_results_list
+
+    gathered_results = [None for _ in range(world_size)]
+    torch.distributed.all_gather_object(gathered_results, partial_results_list)
+
+    # return None to non-main ranks
+    if main_rank is not None:
+        if rank != main_rank:
+            return None
+
+    # return collected results
+    results_list = []
+    for r in gathered_results:
+        results_list.extend(r)
+
+    return results_list
+
+
+@contextlib.contextmanager
+def temporary_directory():
+    """Create a shared temporary directory across ranks in distributed setup.
+
+    This function assumes that the distributed setup has been already
+    correctly initialized. It is intended to be used only in single-node
+    setup so that all ranks can access the directory created."""
+
+    if is_global_rank_zero():
+        tmp_dir = [tempfile.TemporaryDirectory()]
+    else:
+        tmp_dir = [None]
+    dist.broadcast_object_list(tmp_dir)
+    yield tmp_dir[0].name
+    # We use barrier below to make sure that rank zero won't exit
+    # and delete tmp_dir while other ranks may still use it
+    dist.barrier()
+    if is_global_rank_zero():
+        tmp_dir[0].cleanup()
+
+
+def webdataset_split_by_workers(src):
+    """
+    This is for latest webdataset>=0.2.6
+    This function will make sure that each worker gets a different subset of the dataset.
+    """
+    # group = torch.distributed.group.WORLD
+    # rank = torch.distributed.get_rank(group=group)
+    # world_size = torch.distributed.get_world_size(group=group)
+    worker_info = torch.utils.data.get_worker_info()
+    num_workers = 1
+    if worker_info is not None:
+        worker = worker_info.id
+        num_workers = worker_info.num_workers
+
+    if num_workers > 1:
+        yield from list(src)[worker::num_workers]
+    else:
+        yield from src

nemo/utils/dtype.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  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 typing import Dict, Union
+
+import torch
+
+_str_to_dtype: Dict[str, torch.dtype] = dict(
+    float32=torch.float32,
+    float=torch.float32,
+    float64=torch.float64,
+    double=torch.float64,
+    float16=torch.float16,
+    half=torch.float16,
+    bfloat16=torch.bfloat16,
+    bf16=torch.bfloat16,
+    uint8=torch.uint8,
+    byte=torch.uint8,
+    int8=torch.int8,
+    char=torch.int8,
+    int16=torch.int16,
+    short=torch.int16,
+    int32=torch.int32,
+    int=torch.int32,
+    int64=torch.int64,
+    long=torch.int64,
+    bool=torch.bool,
+)
+
+
+def str_to_dtype(dtype: Union[str, torch.dtype]) -> torch.dtype:
+    """Convert a data type name to a PyTorch data type"""
+    if isinstance(dtype, torch.dtype):
+        return dtype
+    name = str(dtype).strip().lower()
+    if name.startswith("torch."):
+        name = name.replace("torch.", "", 1)
+    if name.startswith("fp"):
+        name = name.replace("fp", "float", 1)
+    if name not in _str_to_dtype:
+        raise ValueError(f"Unrecognized dtype ({name})")
+    return _str_to_dtype[name]

nemo/utils/enum.py

@@ -0,0 +1,40 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  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 enum import Enum
+
+
+class PrettyStrEnum(Enum):
+    """
+    Pretty enum to work with string values for config options with choices
+    Provides en automatic error message with possible values, if the value is not in the enum
+    Converting to string will show the actual string value, which makes serialization/deserialization straightforward
+
+    Example:
+        class ASRModelType(PrettyStrEnum):
+            CTC = "ctc"
+            RNNT = "rnnt"
+        ...
+        model_type = ModelType(model_type_string)  # automatically validated
+        if model_type == ModelType.CTC:  # more error-prone (to typos) compared to pure string literals
+            ...  # do something specific to CTC model
+    """
+
+    def __str__(self):
+        return self.value
+
+    @classmethod
+    def _missing_(cls, value: object):
+        choices = ', '.join(map(str, cls))
+        raise ValueError(f"{value} is not a valid {cls.__name__}. Possible choices: {choices}")

nemo/utils/env_var_parsing.py

@@ -0,0 +1,255 @@
+# The MIT Licence (MIT)
+#
+# Copyright (c) 2016 YunoJuno Ltd
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+# Vendored dependency from : https://github.com/yunojuno/python-env-utils/blob/master/env_utils/utils.py
+#
+#
+# Modified by NVIDIA
+#
+# Copyright (c) 2020, 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 decimal
+import json
+import os
+from datetime import datetime
+
+__all__ = [
+    "get_env",
+    "get_envbool",
+    "get_envint",
+    "get_envfloat",
+    "get_envdecimal",
+    "get_envdate",
+    "get_envdatetime",
+    "get_envlist",
+    "get_envdict",
+    "CoercionError",
+    "RequiredSettingMissingError",
+]
+
+
+class CoercionError(Exception):
+    """Custom error raised when a value cannot be coerced."""
+
+    def __init__(self, key, value, func):
+        msg = "Unable to coerce '{}={}' using {}.".format(key, value, func.__name__)
+        super(CoercionError, self).__init__(msg)
+
+
+class RequiredSettingMissingError(Exception):
+    """Custom error raised when a required env var is missing."""
+
+    def __init__(self, key):
+        msg = "Required env var '{}' is missing.".format(key)
+        super(RequiredSettingMissingError, self).__init__(msg)
+
+
+def _get_env(key, default=None, coerce=lambda x: x, required=False):
+    """
+    Return env var coerced into a type other than string.
+    This function extends the standard os.getenv function to enable
+    the coercion of values into data types other than string (all env
+    vars are strings by default).
+    Args:
+        key: string, the name of the env var to look up
+    Kwargs:
+        default: the default value to return if the env var does not exist. NB the
+            default value is **not** coerced, and is assumed to be of the correct type.
+        coerce: a function that is used to coerce the value returned into
+            another type
+        required: bool, if True, then a RequiredSettingMissingError error is raised
+            if the env var does not exist.
+    Returns the env var, passed through the coerce function
+    """
+    try:
+        value = os.environ[key]
+    except KeyError:
+        if required is True:
+            raise RequiredSettingMissingError(key)
+        else:
+            return default
+
+    try:
+        return coerce(value)
+    except Exception:
+        raise CoercionError(key, value, coerce)
+
+
+# standard type coercion functions
+def _bool(value):
+    if isinstance(value, bool):
+        return value
+
+    return not (value is None or value.lower() in ("false", "0", "no", "n", "f", "none"))
+
+
+def _int(value):
+    return int(value)
+
+
+def _float(value):
+    return float(value)
+
+
+def _decimal(value):
+    return decimal.Decimal(value)
+
+
+def _dict(value):
+    return json.loads(value)
+
+
+_DATE_FORMATS = (
+    "%Y-%m-%d",
+    "%d-%m-%Y",
+    "%m/%d/%Y",
+    "%d %B %Y",
+    "%B %d, %Y",
+)
+
+_DATETIME_FORMATS = (
+    "%Y-%m-%dT%H:%M:%S%z",
+    "%Y-%m-%dT%H:%M:%S",
+    "%Y-%m-%d %H:%M:%S%z",
+    "%Y-%m-%d %H:%M:%S",
+    "%d-%m-%Y %H:%M:%S%z",
+    "%d-%m-%Y %H:%M:%S",
+    "%m/%d/%Y %H:%M:%S%z",
+    "%m/%d/%Y %H:%M:%S",
+    "%d %B %Y %H:%M:%S%z",
+    "%d %B %Y %H:%M:%S",
+    "%B %d, %Y %H:%M:%S%z",
+    "%B %d, %Y %H:%M:%S",
+)
+
+
+def _parse_datetime(value):
+    value = value.strip()
+    if value.endswith("Z"):
+        value = value[:-1] + "+00:00"
+
+    try:
+        return datetime.fromisoformat(value)
+    except ValueError:
+        pass
+
+    for fmt in _DATETIME_FORMATS:
+        try:
+            return datetime.strptime(value, fmt)
+        except ValueError:
+            pass
+
+    for fmt in _DATE_FORMATS:
+        try:
+            return datetime.strptime(value, fmt)
+        except ValueError:
+            pass
+
+    raise ValueError(value)
+
+
+def _datetime(value):
+    return _parse_datetime(value)
+
+
+def _date(value):
+    return _parse_datetime(value).date()
+
+
+def get_env(key, *default, **kwargs):
+    """
+    Return env var.
+    This is the parent function of all other get_foo functions,
+    and is responsible for unpacking args/kwargs into the values
+    that _get_env expects (it is the root function that actually
+    interacts with environ).
+    Args:
+        key: string, the env var name to look up.
+        default: (optional) the value to use if the env var does not
+            exist. If this value is not supplied, then the env var is
+            considered to be required, and a RequiredSettingMissingError
+            error will be raised if it does not exist.
+    Kwargs:
+        coerce: a func that may be supplied to coerce the value into
+            something else. This is used by the default get_foo functions
+            to cast strings to builtin types, but could be a function that
+            returns a custom class.
+    Returns the env var, coerced if required, and a default if supplied.
+    """
+    assert len(default) in (0, 1), "Too many args supplied."
+    func = kwargs.get('coerce', lambda x: x)
+    required = len(default) == 0
+    default = default[0] if not required else None
+    return _get_env(key, default=default, coerce=func, required=required)
+
+
+def get_envbool(key, *default):
+    """Return env var cast as boolean."""
+    return get_env(key, *default, coerce=_bool)
+
+
+def get_envint(key, *default):
+    """Return env var cast as integer."""
+    return get_env(key, *default, coerce=_int)
+
+
+def get_envfloat(key, *default):
+    """Return env var cast as float."""
+    return get_env(key, *default, coerce=_float)
+
+
+def get_envdecimal(key, *default):
+    """Return env var cast as Decimal."""
+    return get_env(key, *default, coerce=_decimal)
+
+
+def get_envdate(key, *default):
+    """Return env var as a date."""
+    return get_env(key, *default, coerce=_date)
+
+
+def get_envdatetime(key, *default):
+    """Return env var as a datetime."""
+    return get_env(key, *default, coerce=_datetime)
+
+
+def get_envlist(key, *default, **kwargs):
+    """Return env var as a list."""
+    separator = kwargs.get('separator', ' ')
+    return get_env(key, *default, coerce=lambda x: x.split(separator))
+
+
+def get_envdict(key, *default):
+    """Return env var as a dict."""
+    return get_env(key, *default, coerce=_dict)

nemo/utils/exceptions.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2020, 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.
+
+
+class NeMoBaseException(Exception):
+    """NeMo Base Exception. All exceptions created in NeMo should inherit from this class"""
+
+
+class LightningNotInstalledException(NeMoBaseException):
+    """Raised when optional Lightning dependencies are required but unavailable."""
+
+    def __init__(self, obj):
+        message = (
+            f" You are trying to use {obj} without installing all of pytorch_lightning, hydra, and "
+            f"omegaconf. Please install those packages before trying to access {obj}."
+        )
+        super().__init__(message)
+
+
+class CheckInstall:
+    """Placeholder that raises when optional Lightning dependencies are missing."""
+
+    def __init__(self, *args, **kwargs):
+        raise LightningNotInstalledException(self)
+
+    def __call__(self, *args, **kwargs):
+        raise LightningNotInstalledException(self)
+
+    def __getattr__(self, *args, **kwargs):
+        raise LightningNotInstalledException(self)

nemo/utils/export_utils.py

@@ -0,0 +1,564 @@
+# Copyright (c) 2020, 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.
+
+# flake8: noqa
+# pylint: skip-file
+
+import os
+from enum import Enum
+from typing import Callable, Dict, Optional, Type
+
+import onnx
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from nemo.utils import CastToFloat, CastToFloatAll, logging
+from nemo.utils.megatron_utils import ApexGuardDefaults
+
+try:
+    import onnxruntime
+
+    ort_available = True
+except (ImportError, ModuleNotFoundError):
+    ort_available = False
+
+try:
+    from apex.transformer.functional.fused_softmax import FusedScaleMaskSoftmax
+
+    HAVE_APEX = True
+
+except (ImportError, ModuleNotFoundError):
+    HAVE_APEX = False
+
+
+if HAVE_APEX:
+
+    class MatchedScaleMaskSoftmax(FusedScaleMaskSoftmax):
+        """
+        fused operation: scaling + mask + softmax
+        match the behavior of fused softmax and torch softmax.
+        This is a workaround for https://github.com/NVIDIA/apex/issues/1493.
+
+        Arguments:
+            input_in_fp16: flag to indicate if input in fp16 data format.
+            input_in_bf16: flag to indicate if input in bf16 data format.
+            attn_mask_type: attention mask type (pad or causal)
+            scaled_masked_softmax_fusion: flag to indicate user want to use softmax fusion
+            mask_func: mask function to be applied.
+            softmax_in_fp32: if true, softmax in performed at fp32 precision.
+            scale: scaling factor used in input tensor scaling.
+        """
+
+        def forward_torch_softmax(self, input, mask):
+            if self.input_in_float16 and self.softmax_in_fp32:
+                input = input.float()
+
+            if self.scale is not None:
+                input = input * self.scale
+            mask_output = self.mask_func(input, mask) if mask is not None else input
+            probs = torch.nn.Softmax(dim=-1)(mask_output)
+            if mask is not None:
+                all_k_masked = mask.all(axis=-1)
+                zero_attention_mask = (1.0 - all_k_masked.type(probs.type()))[:, :, :, None]
+                probs = probs * zero_attention_mask
+
+            if self.input_in_float16 and self.softmax_in_fp32:
+                if self.input_in_fp16:
+                    probs = probs.half()
+                else:
+                    probs = probs.bfloat16()
+            return probs
+
+else:
+
+    class MatchedScaleMaskSoftmax(ApexGuardDefaults):
+        def __init__(self):
+            super().__init__()
+            logging.warning(
+                "Apex was not found. ColumnLinear will not work. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+            )
+
+
+class ExportFormat(Enum):
+    """Which format to use when exporting a Neural Module for deployment"""
+
+    ONNX = 1
+    TORCHSCRIPT = 2
+
+
+_EXT_DICT = {
+    ".pt": ExportFormat.TORCHSCRIPT,
+    ".ts": ExportFormat.TORCHSCRIPT,
+    ".onnx": ExportFormat.ONNX,
+}
+
+
+class TorchRMSNorm(nn.Module):
+    def __init__(self, weight, eps=1e-6):
+        """
+        LayerNorm without bias
+        """
+        super().__init__()
+        self.weight = weight
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        # can be only calculated with precision=32
+        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        if self.weight.dtype in [torch.float16, torch.bfloat16]:
+            hidden_states = hidden_states.to(self.weight.dtype)
+
+        return self.weight * hidden_states
+
+
+class LinearWithBiasSkip(nn.Module):
+    def __init__(self, weight, bias, skip_bias_add):
+        super(LinearWithBiasSkip, self).__init__()
+        self.bias = bias
+        self.weight = weight
+        self.skip_bias_add = skip_bias_add
+
+    def forward(self, x, weight=None):
+        if weight is None:
+            weight = self.weight
+        if self.skip_bias_add:
+            return F.linear(x, weight), self.bias
+        return F.linear(x, weight, self.bias), None
+
+
+def get_export_format(filename: str):
+    _, ext = os.path.splitext(filename)
+    try:
+        return _EXT_DICT[ext.lower()]
+    except KeyError:
+        raise ValueError(f"Export file {filename} extension does not correspond to any export format!")
+
+
+def augment_filename(output: str, prepend: str):
+    if prepend == 'self':
+        return output
+
+    path, filename = os.path.split(output)
+    filename = f"{prepend}-{filename}"
+    return os.path.join(path, filename)
+
+
+def forward_method(self):
+    if hasattr(self, "forward_for_export"):
+        return self.forward_for_export
+    else:
+        return self.forward
+
+
+def wrap_forward_method(self):
+    tp = type(self)
+    old_forward_method = None
+    if hasattr(tp, "forward_for_export"):
+        forward_method = tp.forward_for_export
+        old_forward_method = tp.forward
+        tp.forward = forward_method
+    else:
+        forward_method = None
+    return forward_method, old_forward_method
+
+
+def parse_input_example(input_example):
+    input_list = list(input_example)
+    input_dict = {}
+    # process possible kwargs
+    if isinstance(input_list[-1], dict):
+        input_dict = input_list[-1]
+        input_list = input_list[:-1]
+    return input_list, input_dict
+
+
+def to_onnxrt_input(ort_input_names, input_names, input_dict, input_list):
+    odict = {}
+    if not input_names:
+        input_list.extend(input_dict.values())
+        for k, v in zip(ort_input_names, input_list):
+            odict[k] = v.cpu().numpy()
+        return odict
+    for k in reversed(input_names):
+        val = None
+        if k in input_dict:
+            val = input_dict[k].cpu().numpy()
+        elif len(input_list) > 0:
+            val = input_list.pop().cpu().numpy()
+        if k in ort_input_names and val is not None:
+            odict[k] = val
+    return odict
+
+
+def verify_torchscript(model, output, input_examples, check_tolerance=0.01):
+    all_good = True
+    for input_example in input_examples:
+        input_list, input_dict = parse_input_example(input_example)
+        # We disable autocast here to make sure exported TS will run under Triton or other C++ env
+        with torch.amp.autocast('cuda', enabled=False):
+            output_example = model.forward(*input_list, **input_dict)
+            ts_model = torch.jit.load(output)
+            all_good = all_good and run_ts_and_compare(
+                ts_model, input_list, input_dict, output_example, check_tolerance
+            )
+    status = "SUCCESS" if all_good else "FAIL"
+    logging.info(f"Torchscript generated at {output} verified with torchscript forward : " + status)
+    return all_good
+
+
+def verify_runtime(model, output, input_examples, input_names, check_tolerance=0.01):
+    onnx_model = onnx.load(output)
+    ort_input_names = [node.name for node in onnx_model.graph.input]
+
+    global ort_available
+    if not ort_available:
+        logging.warning(f"ONNX generated at {output}, not verified - please install onnxruntime_gpu package.\n")
+        onnx.checker.check_model(onnx_model, full_check=True)
+        return
+    onnx_session_opt = onnxruntime.SessionOptions()
+    onnx_session_opt.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_BASIC
+    sess = onnxruntime.InferenceSession(
+        onnx_model.SerializeToString(), sess_options=onnx_session_opt, providers=['CUDAExecutionProvider']
+    )
+    del onnx_model
+    all_good = True
+    for input_example in input_examples:
+        input_list, input_dict = parse_input_example(input_example)
+        output_example = model.forward(*input_list, **input_dict)
+        if not isinstance(output_example, tuple):
+            output_example = (output_example,)
+        ort_input = to_onnxrt_input(ort_input_names, input_names, input_dict, input_list)
+        all_good = all_good and run_ort_and_compare(sess, ort_input, output_example, check_tolerance)
+    status = "SUCCESS" if all_good else "FAIL"
+    logging.info(f"ONNX generated at {output} verified with onnxruntime : " + status)
+    return all_good
+
+
+def run_ts_and_compare(ts_model, ts_input_list, ts_input_dict, output_example, check_tolerance=0.01):
+    # Verify the model can be read, and is valid
+    ts_out = ts_model(*ts_input_list, **ts_input_dict)
+
+    all_good = True
+    for i, out in enumerate(ts_out):
+        expected = output_example[i]
+
+        if torch.is_tensor(expected):
+            tout = out.to('cpu')
+            logging.debug(f"Checking output {i}, shape: {expected.shape}:\n")
+            this_good = True
+            try:
+                if not torch.allclose(tout, expected.cpu(), rtol=check_tolerance, atol=check_tolerance):
+                    this_good = False
+            except Exception:  # there may ne size mismatch and it may be OK
+                this_good = False
+            if not this_good:
+                logging.info(f"Results mismatch! PyTorch(expected):\n{expected}\nTorchScript:\n{tout}")
+                all_good = False
+    return all_good
+
+
+def run_ort_and_compare(sess, ort_input, output_example, check_tolerance=0.01):
+    # Verify the model can be read, and is valid
+    ort_out = sess.run(None, ort_input)
+    all_good = True
+    for i, out in enumerate(ort_out):
+        expected = output_example[i]
+
+        if torch.is_tensor(expected):
+            tout = torch.from_numpy(out)
+            logging.debug(f"Checking output {i}, shape: {expected.shape}:\n")
+            this_good = True
+            try:
+                if not torch.allclose(tout, expected.cpu(), rtol=check_tolerance, atol=100 * check_tolerance):
+                    this_good = False
+            except Exception:  # there may be size mismatch and it may be OK
+                this_good = False
+            if not this_good:
+                logging.info(
+                    f"onnxruntime results mismatch! PyTorch(expected, {expected.shape}):\n{expected}\nONNXruntime, {tout.shape}:\n{tout}"
+                )
+                all_good = False
+    return all_good
+
+
+apex_available = True
+
+try:
+    from apex.contrib.layer_norm.layer_norm import FastLayerNorm
+    from apex.normalization import MixedFusedRMSNorm
+    from apex.normalization.fused_layer_norm import FusedLayerNorm, MixedFusedLayerNorm
+    from megatron.core.fusions.fused_layer_norm import FusedLayerNorm as MCoreFusedLayerNorm
+    from megatron.core.fusions.fused_softmax import FusedScaleMaskSoftmax
+    from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+
+    def replace_FusedLayerNorm(n: nn.Module) -> Optional[nn.LayerNorm]:
+        """
+        Replaces Apex's FusedLayerNorm with nn.LayerNorm. This is required for ONNX export.
+        Args:
+           n: the FusedLayerNorm pytorch module to replace
+        Returns:
+           Equivalent LayerNorm module
+        """
+
+        p = next(n.parameters())
+
+        if isinstance(n, FusedLayerNorm) or isinstance(n, MixedFusedLayerNorm):
+            shape, eps, affine = n.normalized_shape, n.eps, n.elementwise_affine
+        elif isinstance(n, MCoreFusedLayerNorm):
+            shape, eps, affine = n.weight.shape, n.eps, True
+        elif isinstance(n, FastLayerNorm):
+            shape, eps, affine = n.weight.shape, n.epsilon, True
+        else:
+            return None
+
+        n_state = n.state_dict()
+        mod = nn.LayerNorm(shape, eps=eps, elementwise_affine=affine, device=p.device, dtype=p.dtype)
+
+        mod.load_state_dict(n_state, strict=True)
+
+        return mod
+
+    def replace_MixedFusedRMSNorm(n: nn.Module):
+        """
+        Replaces Apex's MixedFusedRMSNorm with equivalent Pytorch layer. This is required for ONNX export.
+        Args:
+           n: the MixedFusedRMSNorm pytorch module to replace
+        Returns:
+           Equivalent module
+        """
+
+        p = next(n.parameters())
+
+        if isinstance(n, MixedFusedRMSNorm):
+            mod = TorchRMSNorm(n.state_dict()['weight'], n.eps).to(p.device)
+        else:
+            return None
+
+        return mod
+
+    def replace_ParallelLinear(n: nn.Module) -> Optional[nn.Linear]:
+        """
+        Replaces Apex's ColumnParallelLinear or RowParallelLinear with nn.Linear
+        Args:
+           n: the nn.Module pytorch module to replace
+        Returns:
+           Equivalent Linear module
+        """
+        if not (isinstance(n, ColumnParallelLinear) or isinstance(n, RowParallelLinear)):
+            raise ValueError("This function can only change the ColumnParallelLinear or RowParallelLinear module.")
+
+        dev = next(n.parameters()).device
+        mod = LinearWithBiasSkip(n.weight, n.bias, n.skip_bias_add).to(dev)
+
+        n_state = n.state_dict()
+        mod.load_state_dict(n_state, strict=False)
+        return mod
+
+    def replace_FusedScaleMaskSoftmax(n: nn.Module) -> Optional[nn.Linear]:
+        """
+        Replaces Apex's FusedScaleMaskSoftmax with nn.LayerNorm. This is required for ONNX export.
+        Args:
+           n: the FusedScaleMaskSoftmax module to replace
+        Returns:
+           Equivalent LayerNorm module
+        """
+        if not isinstance(n, FusedScaleMaskSoftmax):
+            logging.warning(f"This function can only change the FusedScaleMaskSoftmax module, got: {n.__class__}")
+            return n
+
+        # disable the fusion only
+        mod = FusedScaleMaskSoftmax(
+            n.input_in_fp16, n.input_in_bf16, n.attn_mask_type, False, n.mask_func, n.softmax_in_fp32, n.scale
+        )
+
+        return mod
+
+    default_Apex_replacements = {
+        "FusedLayerNorm": replace_FusedLayerNorm,
+        "MixedFusedLayerNorm": replace_FusedLayerNorm,
+        "MCoreFusedLayerNorm": replace_FusedLayerNorm,
+        "FastLayerNorm": replace_FusedLayerNorm,
+        "RowParallelLinear": replace_ParallelLinear,
+        "ColumnParallelLinear": replace_ParallelLinear,
+        "FusedScaleMaskSoftmax": replace_FusedScaleMaskSoftmax,
+        "MixedFusedRMSNorm": replace_MixedFusedRMSNorm,
+    }
+
+except Exception:
+    default_Apex_replacements = {}
+    apex_available = False
+
+
+def simple_replace(BaseT: Type[nn.Module], DestT: Type[nn.Module]) -> Callable[[nn.Module], Optional[nn.Module]]:
+    """
+    Generic function generator to replace BaseT module with DestT. BaseT and DestT should have same atrributes. No weights are copied.
+    Args:
+        BaseT : module type to replace
+        DestT : destination module type
+    Returns:
+        swap function to replace BaseT module with DestT
+    """
+
+    def expansion_fn(mod: nn.Module) -> Optional[nn.Module]:
+        if not isinstance(mod, BaseT):
+            return None
+        args = [getattr(mod, name, None) for name in mod.__constants__]
+        out = DestT(*args)
+        return out
+
+    return expansion_fn
+
+
+def replace_MatchedScaleMaskSoftmax(n: nn.Module) -> Optional[nn.Linear]:
+    """
+    Replaces MatchedScaleMaskSoftmax with exportable softmax layer
+    Args:
+        n: module to replace
+    Returns:
+        exportable module
+    """
+
+    # disabling fusion for the MatchedScaleMaskSoftmax
+    mod = MatchedScaleMaskSoftmax(
+        n.input_in_fp16, n.input_in_bf16, n.attn_mask_type, False, n.mask_func, n.softmax_in_fp32, n.scale
+    )
+    return mod
+
+
+def wrap_module(BaseT: Type[nn.Module], DestT: Type[nn.Module]) -> Callable[[nn.Module], Optional[nn.Module]]:
+    """
+    Generic function generator to replace BaseT module with DestT wrapper.
+    Args:
+        BaseT : module type to replace
+        DestT : destination module type
+    Returns:
+        swap function to replace BaseT module with DestT
+    """
+
+    def expansion_fn(mod: nn.Module) -> Optional[nn.Module]:
+        out = DestT(mod)
+        return out
+
+    return expansion_fn
+
+
+def swap_modules(model: nn.Module, mapping: Dict[str, nn.Module]):
+    """
+    This function swaps nested modules as specified by "dot paths" in mod with a desired replacement. This allows
+    for swapping nested modules through arbitrary levels if children
+
+    NOTE: This occurs in place, if you want to preserve model then make sure to copy it first.
+
+    """
+    for path, new_mod in mapping.items():
+        expanded_path = path.split(".")
+        parent_mod = model
+        for sub_path in expanded_path[:-1]:
+            parent_mod = parent_mod._modules[sub_path]  # noqa
+        parent_mod._modules[expanded_path[-1]] = new_mod  # noqa
+
+    return model
+
+
+def replace_modules(
+    model: nn.Module, expansions: Dict[str, Callable[[nn.Module], Optional[nn.Module]]] = None
+) -> nn.Module:
+    """
+    Top-level function to replace modules in model, specified by class name with a desired replacement.
+    NOTE: This occurs in place, if you want to preserve model then make sure to copy it first.
+    Args:
+        model : top level module
+        expansions : replacement dictionary: module class name -> replacement function generator
+    Returns:
+        model, possibly modified in-place
+    """
+    mapping: Dict[str, nn.Module] = {}
+    for name, m in model.named_modules():
+        m_type = type(m).__name__
+        if m_type in expansions:
+            swapped = expansions[m_type](m)
+            if swapped:
+                mapping[name] = swapped
+    if len(mapping) > 0:
+        logging.info(f"Swapped {len(mapping)} modules")
+    swap_modules(model, mapping)
+    return model
+
+
+def script_module(m: nn.Module):
+    return torch.jit.script(m)
+
+
+script_replacements = {}
+
+
+def replace_for_export(model: nn.Module) -> nn.Module:
+    """
+    Top-level function to replace 'default set' of modules in model, called from _prepare_for_export.
+    NOTE: This occurs in place, if you want to preserve model then make sure to copy it first.
+    Args:
+        model : top level module
+    Returns:
+        model, possibly modified in-place
+    """
+    default_replacements = {
+        "MatchedScaleMaskSoftmax": wrap_module(None, replace_MatchedScaleMaskSoftmax),
+    }
+
+    replace_modules(model, default_Apex_replacements)
+    replace_modules(model, default_replacements)
+    # This one has to be the last
+    replace_modules(model, script_replacements)
+
+
+def add_casts_around_norms(model: nn.Module):
+    """
+    Function to put additional to/from float32 casts around operations known to require full precision.
+    It was used with an extra post-parse script to have TRT preserve extra precision when --fp16 needed.
+    Should not be needed with TRT 8.6.1 or later.
+    """
+    from nemo.collections.tts.modules.submodules import MaskedInstanceNorm1d
+
+    default_cast_replacements = {
+        "BatchNorm1d": wrap_module(nn.BatchNorm1d, CastToFloat),
+        "BatchNorm2d": wrap_module(nn.BatchNorm2d, CastToFloat),
+        "LayerNorm": wrap_module(nn.LayerNorm, CastToFloat),
+        "InstanceNorm1d": wrap_module(nn.InstanceNorm1d, CastToFloat),
+        "MaskedInstanceNorm1d": wrap_module(MaskedInstanceNorm1d, CastToFloatAll),
+    }
+    replace_modules(model, default_cast_replacements)
+
+
+def rename_onnx_io(output, input_names, output_names):
+    onnx_model = onnx.load(output)
+    rename_map = {}
+    for inp, name in zip(onnx_model.graph.input, input_names):
+        rename_map[inp.name] = name
+    for out, name in zip(onnx_model.graph.output, output_names):
+        rename_map[out.name] = name
+    for n in onnx_model.graph.node:
+        for inp in range(len(n.input)):
+            if n.input[inp] in rename_map:
+                n.input[inp] = rename_map[n.input[inp]]
+        for out in range(len(n.output)):
+            if n.output[out] in rename_map:
+                n.output[out] = rename_map[n.output[out]]
+
+    for i in range(len(input_names)):
+        onnx_model.graph.input[i].name = input_names[i]
+    for i in range(len(output_names)):
+        onnx_model.graph.output[i].name = output_names[i]
+    onnx.save(onnx_model, output)

nemo/utils/file_utils.py

@@ -0,0 +1,31 @@
+# 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 pathlib import Path
+from typing import Union
+
+
+def robust_copy(src: Union[Path, str], dst: Union[Path, str]) -> str:
+    """
+    Copy file from src to dst, falling back to shutil.copy if shutil.copy2 fails.
+    shutil.copy2 preserves metadata, but can fail on some filesystems.
+    """
+    try:
+        return shutil.copy2(src, dst)
+    except PermissionError:
+        # copy2 can fail on some filesystems due to metadata copy errors
+        # (e.g., permission errors on setting timestamps).
+        # In such cases, we fallback to a plain copy.
+        return shutil.copy(src, dst)

nemo/utils/formatters/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2020, 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.

nemo/utils/formatters/base.py

@@ -0,0 +1,135 @@
+# Copyright (c) 2020, 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 logging
+
+from nemo.utils.formatters.colors import Fore as ForegroundColors
+from nemo.utils.formatters.utils import check_color_support, to_unicode
+
+__all__ = ["BaseNeMoFormatter"]
+
+
+class BaseFormatter(logging.Formatter):
+    """
+    Log formatter used in Tornado. Key features of this formatter are:
+    * Color support when logging to a terminal that supports it.
+    * Timestamps on every log line.
+    * Robust against str/bytes encoding problems.
+    """
+
+    DEFAULT_FORMAT = "%(color)s[%(levelname)1.1s %(asctime)s %(module)s:%(lineno)d]%(end_color)s %(message)s"
+
+    DEFAULT_DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
+
+    DEFAULT_COLORS = {
+        logging.DEBUG: ForegroundColors.CYAN,
+        logging.INFO: ForegroundColors.GREEN,
+        logging.WARNING: ForegroundColors.YELLOW,
+        logging.ERROR: ForegroundColors.MAGENTA,
+        logging.CRITICAL: ForegroundColors.RED,
+    }
+
+    def __init__(self, color=True, fmt=None, datefmt=None, colors=None):
+        r"""
+        :arg bool color: Enables color support.
+        :arg string fmt: Log message format.
+          It will be applied to the attributes dict of log records. The
+          text between ``%(color)s`` and ``%(end_color)s`` will be colored
+          depending on the level if color support is on.
+        :arg dict colors: color mappings from logging level to terminal color
+          code
+        :arg string datefmt: Datetime format.
+          Used for formatting ``(asctime)`` placeholder in ``prefix_fmt``.
+        .. versionchanged:: 3.2
+           Added ``fmt`` and ``datefmt`` arguments.
+        """
+
+        if fmt is None:
+            fmt = self.DEFAULT_FORMAT
+
+        if datefmt is None:
+            datefmt = self.DEFAULT_DATE_FORMAT
+
+        if colors is None:
+            colors = self.DEFAULT_COLORS
+
+        logging.Formatter.__init__(self, datefmt=datefmt)
+
+        self._fmt = fmt
+        self._colors = {}
+        self._normal = ""
+
+        if color and check_color_support():
+            self._colors = colors
+            self._normal = ForegroundColors.RESET
+
+    def format(self, record):
+        try:
+            message = record.getMessage()
+            assert isinstance(message, str)  # guaranteed by logging
+            # Encoding notes:  The logging module prefers to work with character
+            # strings, but only enforces that log messages are instances of
+            # basestring.  In python 2, non-ascii bytestrings will make
+            # their way through the logging framework until they blow up with
+            # an unhelpful decoding error (with this formatter it happens
+            # when we attach the prefix, but there are other opportunities for
+            # exceptions further along in the framework).
+            #
+            # If a byte string makes it this far, convert it to unicode to
+            # ensure it will make it out to the logs.  Use repr() as a fallback
+            # to ensure that all byte strings can be converted successfully,
+            # but don't do it by default so we don't add extra quotes to ascii
+            # bytestrings.  This is a bit of a hacky place to do this, but
+            # it's worth it since the encoding errors that would otherwise
+            # result are so useless (and tornado is fond of using utf8-encoded
+            # byte strings wherever possible).
+            record.message = to_unicode(message)
+
+        except Exception as e:
+            record.message = "Bad message (%r): %r" % (e, record.__dict__)
+
+        record.asctime = self.formatTime(record, self.datefmt)
+
+        if record.levelno in self._colors:
+            record.color = self._colors[record.levelno]
+            record.end_color = self._normal
+        else:
+            record.color = record.end_color = ""
+
+        formatted = self._fmt % record.__dict__
+
+        if record.exc_info:
+            if not record.exc_text:
+                record.exc_text = self.formatException(record.exc_info)
+
+        if record.exc_text:
+            # exc_text contains multiple lines.  We need to _safe_unicode
+            # each line separately so that non-utf8 bytes don't cause
+            # all the newlines to turn into '\n'.
+            lines = [formatted.rstrip()]
+            lines.extend(to_unicode(ln) for ln in record.exc_text.split("\n"))
+
+            formatted = "\n".join(lines)
+        return formatted.replace("\n", "\n    ")
+
+
+class BaseNeMoFormatter(BaseFormatter):
+    DEFAULT_FORMAT = "%(color)s[NeMo %(levelname)1.1s %(asctime)s %(module)s:%(lineno)d]%(end_color)s %(message)s"
+
+
+class DebugNeMoFormatter(BaseFormatter):
+    DEFAULT_FORMAT = (
+        "%(color)s[NeMo %(levelname)1.1s %(asctime)s %(module)s:%(lineno)d rank:%(rank)s]%(end_color)s %(message)s"
+    )

nemo/utils/formatters/colors.py

@@ -0,0 +1,121 @@
+# Source: https://github.com/tartley/colorama/blob/master/colorama/ansi.py
+# Copyright: Jonathan Hartley 2013. BSD 3-Clause license.
+#
+# Copyright (c) 2020, 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.
+
+CSI = "\033["
+OSC = "\033]"
+BEL = "\007"
+
+
+def code_to_chars(code):
+    return CSI + str(code) + "m"
+
+
+def set_title(title):
+    return OSC + "2;" + title + BEL
+
+
+def clear_screen(mode=2):
+    return CSI + str(mode) + "J"
+
+
+def clear_line(mode=2):
+    return CSI + str(mode) + "K"
+
+
+class AnsiCodes(object):
+    def __init__(self):
+        # the subclasses declare class attributes which are numbers.
+        # Upon instantiation we define instance attributes, which are the same
+        # as the class attributes but wrapped with the ANSI escape sequence
+        for name in dir(self):
+            if not name.startswith("_"):
+                value = getattr(self, name)
+                setattr(self, name, code_to_chars(value))
+
+
+class AnsiCursor(object):
+    def UP(self, n=1):
+        return CSI + str(n) + "A"
+
+    def DOWN(self, n=1):
+        return CSI + str(n) + "B"
+
+    def FORWARD(self, n=1):
+        return CSI + str(n) + "C"
+
+    def BACK(self, n=1):
+        return CSI + str(n) + "D"
+
+    def POS(self, x=1, y=1):
+        return CSI + str(y) + ";" + str(x) + "H"
+
+
+class AnsiFore(AnsiCodes):
+    BLACK = 30
+    RED = 31
+    GREEN = 32
+    YELLOW = 33
+    BLUE = 34
+    MAGENTA = 35
+    CYAN = 36
+    WHITE = 37
+    RESET = 39
+
+    # These are fairly well supported, but not part of the standard.
+    LIGHTBLACK_EX = 90
+    LIGHTRED_EX = 91
+    LIGHTGREEN_EX = 92
+    LIGHTYELLOW_EX = 93
+    LIGHTBLUE_EX = 94
+    LIGHTMAGENTA_EX = 95
+    LIGHTCYAN_EX = 96
+    LIGHTWHITE_EX = 97
+
+
+class AnsiBack(AnsiCodes):
+    BLACK = 40
+    RED = 41
+    GREEN = 42
+    YELLOW = 43
+    BLUE = 44
+    MAGENTA = 45
+    CYAN = 46
+    WHITE = 47
+    RESET = 49
+
+    # These are fairly well supported, but not part of the standard.
+    LIGHTBLACK_EX = 100
+    LIGHTRED_EX = 101
+    LIGHTGREEN_EX = 102
+    LIGHTYELLOW_EX = 103
+    LIGHTBLUE_EX = 104
+    LIGHTMAGENTA_EX = 105
+    LIGHTCYAN_EX = 106
+    LIGHTWHITE_EX = 107
+
+
+class AnsiStyle(AnsiCodes):
+    BRIGHT = 1
+    DIM = 2
+    NORMAL = 22
+    RESET_ALL = 0
+
+
+Fore = AnsiFore()
+Back = AnsiBack()
+Style = AnsiStyle()
+Cursor = AnsiCursor()

nemo/utils/formatters/utils.py

@@ -0,0 +1,46 @@
+# Copyright (c) 2020, 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 sys
+
+from nemo.constants import NEMO_ENV_VARNAME_ENABLE_COLORING
+from nemo.utils.env_var_parsing import get_envbool
+
+__all__ = ["check_color_support", "to_unicode"]
+
+
+def check_color_support():
+    # Colors can be forced with an env variable
+    if not sys.platform.lower().startswith("win") and get_envbool(NEMO_ENV_VARNAME_ENABLE_COLORING, False):
+        return True
+
+
+def to_unicode(value):
+    """
+    Converts a string argument to a unicode string.
+    If the argument is already a unicode string or None, it is returned
+    unchanged.  Otherwise it must be a byte string and is decoded as utf8.
+    """
+    try:
+        if isinstance(value, (str, type(None))):
+            return value
+
+        if not isinstance(value, bytes):
+            raise TypeError("Expected bytes, unicode, or None; got %r" % type(value))
+
+        return value.decode("utf-8")
+
+    except UnicodeDecodeError:
+        return repr(value)

nemo/utils/get_rank.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2020, 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 torch
+
+from nemo.utils.env_var_parsing import get_envint
+
+
+def is_global_rank_zero():
+    """Helper function to determine if the current process is global_rank 0 (the main process)"""
+    # Try to get the pytorch RANK env var
+    # RANK is set by torch.distributed.launch
+    rank = get_envint("RANK", None)
+    if rank is not None:
+        return rank == 0
+
+    # Try to get the SLURM global rank env var
+    # SLURM_PROCID is set by SLURM
+    slurm_rank = get_envint("SLURM_PROCID", None)
+    if slurm_rank is not None:
+        return slurm_rank == 0
+
+    # Try to get the MPI global rank env var
+    mpi_rank = get_envint("OMPI_COMM_WORLD_RANK", None)
+    if mpi_rank is not None:
+        return mpi_rank == 0
+
+    # if neither pytorch, SLURM nor MPI env vars are set
+    # check NODE_RANK/GROUP_RANK and LOCAL_RANK env vars
+    # assume global_rank is zero if undefined
+    node_rank = get_envint("NODE_RANK", get_envint("GROUP_RANK", 0))
+    local_rank = get_envint("LOCAL_RANK", 0)
+    return node_rank == 0 and local_rank == 0
+
+
+def get_rank():
+    """Helper function that returns torch.distributed.get_rank() if DDP has been initialized otherwise it returns 0."""
+
+    if is_global_rank_zero():
+        return 0
+    else:
+        return torch.distributed.get_rank()
+
+
+def get_last_rank() -> int:
+    """Get the last rank in the distributed group"""
+    if not torch.distributed.is_initialized():
+        return 0
+    return torch.distributed.get_world_size() - 1

nemo/utils/import_utils.py

@@ -0,0 +1,402 @@
+# 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.
+
+# This file is taken from https://github.com/NVIDIA-NeMo/Curator/blob/dask/nemo_curator/utils/import_utils.py,
+# which is adapted from cuML's safe_imports module:
+# https://github.com/rapidsai/cuml/blob/e93166ea0dddfa8ef2f68c6335012af4420bc8ac/python/cuml/internals/safe_imports.py
+
+
+import importlib
+import logging
+import traceback
+from contextlib import contextmanager
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+logger.addHandler(logging.StreamHandler())
+
+GPU_INSTALL_STRING = (
+    """Install GPU packages via `pip install --extra-index-url """
+    """https://pypi.nvidia.com nemo-curator[cuda12x]`
+or use `pip install --extra-index-url https://pypi.nvidia.com ".[cuda12x]"` if installing from source"""
+)
+
+
+class UnavailableError(Exception):
+    """Error thrown if a symbol is unavailable due to an issue importing it"""
+
+
+@contextmanager
+def null_decorator(*args, **kwargs):
+    """null_decorator"""
+    if len(kwargs) == 0 and len(args) == 1 and callable(args[0]):
+        return args[0]
+    else:
+
+        def inner(func):
+            return func
+
+        return inner
+
+
+class UnavailableMeta(type):
+    """A metaclass for generating placeholder objects for unavailable symbols
+
+    This metaclass allows errors to be deferred from import time to the time
+    that a symbol is actually used in order to streamline the usage of optional
+    dependencies. This is particularly useful for attempted imports of GPU-only
+    modules which will only be invoked if GPU-only functionality is
+    specifically used.
+
+    If an attempt to import a symbol fails, this metaclass is used to generate
+    a class which stands in for that symbol. Any attempt to call the symbol
+    (instantiate the class) or access its attributes will throw an
+    UnavailableError exception. Furthermore, this class can be used in
+    e.g. isinstance checks, since it will (correctly) fail to match any
+    instance it is compared against.
+
+    In addition to calls and attribute access, a number of dunder methods are
+    implemented so that other common usages of imported symbols (e.g.
+    arithmetic) throw an UnavailableError, but this is not guaranteed for
+    all possible uses. In such cases, other exception types (typically
+    TypeErrors) will be thrown instead.
+    """
+
+    def __new__(meta, name, bases, dct):
+        if dct.get("_msg", None) is None:
+            dct["_msg"] = f"{name} could not be imported"
+        name = f"MISSING{name}"
+        return super(UnavailableMeta, meta).__new__(meta, name, bases, dct)
+
+    def __call__(cls, *args, **kwargs):
+        raise UnavailableError(cls._msg)
+
+    def __getattr__(cls, name):
+        raise UnavailableError(cls._msg)
+
+    def __eq__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __lt__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __gt__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __le__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __ge__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __ne__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __abs__(cls):
+        raise UnavailableError(cls._msg)
+
+    def __add__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __radd__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __iadd__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __floordiv__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __rfloordiv__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __ifloordiv__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __lshift__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __rlshift__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __mul__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __rmul__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __imul__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __ilshift__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __pow__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __rpow__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __ipow__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __rshift__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __rrshift__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __irshift__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __sub__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __rsub__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __isub__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __truediv__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __rtruediv__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __itruediv__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __divmod__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __rdivmod__(cls, other):
+        raise UnavailableError(cls._msg)
+
+    def __neg__(cls):
+        raise UnavailableError(cls._msg)
+
+    def __invert__(cls):
+        raise UnavailableError(cls._msg)
+
+    def __hash__(cls):
+        raise UnavailableError(cls._msg)
+
+    def __index__(cls):
+        raise UnavailableError(cls._msg)
+
+    def __iter__(cls):
+        raise UnavailableError(cls._msg)
+
+    def __delitem__(cls, name):
+        raise UnavailableError(cls._msg)
+
+    def __setitem__(cls, name, value):
+        raise UnavailableError(cls._msg)
+
+    def __enter__(cls, *args, **kwargs):
+        raise UnavailableError(cls._msg)
+
+    def __get__(cls, *args, **kwargs):
+        raise UnavailableError(cls._msg)
+
+    def __delete__(cls, *args, **kwargs):
+        raise UnavailableError(cls._msg)
+
+    def __len__(cls):
+        raise UnavailableError(cls._msg)
+
+
+def is_unavailable(obj):
+    """Helper to check if given symbol is actually a placeholder"""
+    return type(obj) is UnavailableMeta
+
+
+class UnavailableNullContext:
+    """A placeholder class for unavailable context managers
+
+    This context manager will return a value which will throw an
+    UnavailableError if used in any way, but the context manager itself can be
+    safely invoked.
+    """
+
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def __enter__(self):
+        return UnavailableMeta(
+            "MissingContextValue",
+            (),
+            {"_msg": "Attempted to make use of placeholder context return value."},
+        )
+
+    def __exit__(self, *args, **kwargs):
+        pass
+
+
+def safe_import(module, *, msg=None, alt=None):
+    """A function used to import modules that may not be available
+
+    This function will attempt to import a module with the given name, but it
+    will not throw an ImportError if the module is not found. Instead, it will
+    return a placeholder object which will raise an exception only if used.
+
+    Parameters
+    ----------
+    module: str
+        The name of the module to import.
+    msg: str or None
+        An optional error message to be displayed if this module is used
+        after a failed import.
+    alt: object
+        An optional module to be used in place of the given module if it
+        fails to import
+
+    Returns
+    -------
+    Tuple(object, bool)
+        The imported module, the given alternate, or a class derived from
+        UnavailableMeta, and a boolean indicating whether the intended import was successful.
+    """
+    try:
+        return importlib.import_module(module), True
+    except ImportError:
+        exception_text = traceback.format_exc()
+        logger.debug(f"Import of {module} failed with: {exception_text}")
+    except Exception:
+        exception_text = traceback.format_exc()
+        raise
+    if msg is None:
+        msg = f"{module} could not be imported"
+    if alt is None:
+        return UnavailableMeta(module.rsplit(".")[-1], (), {"_msg": msg}), False
+    else:
+        return alt, False
+
+
+def safe_import_from(module, symbol, *, msg=None, alt=None, fallback_module=None):
+    """A function used to import symbols from modules that may not be available
+
+    This function will attempt to import a symbol with the given name from
+    the given module, but it will not throw an ImportError if the symbol is not
+    found. Instead, it will return a placeholder object which will raise an
+    exception only if used.
+
+    Parameters
+    ----------
+    module: str
+        The name of the module in which the symbol is defined.
+    symbol: str
+        The name of the symbol to import.
+    msg: str or None
+        An optional error message to be displayed if this symbol is used
+        after a failed import.
+    alt: object
+        An optional object to be used in place of the given symbol if it fails
+        to import
+    fallback_module: str
+        Alternative name of the model in which the symbol is defined. The function will first to
+        import using the `module` value and if that fails will also try the `fallback_module`.
+
+    Returns
+    -------
+    Tuple(object, bool)
+        The imported symbol, the given alternate, or a class derived from
+        UnavailableMeta, and a boolean indicating whether the intended import was successful.
+    """
+    try:
+        imported_module = importlib.import_module(module)
+        return getattr(imported_module, symbol), True
+    except ImportError:
+        exception_text = traceback.format_exc()
+        logger.debug(f"Import of {module} failed with: {exception_text}")
+    except AttributeError:
+        # if there is a fallback module try it.
+        if fallback_module is not None:
+            return safe_import_from(fallback_module, symbol, msg=msg, alt=alt, fallback_module=None)
+        exception_text = traceback.format_exc()
+        logger.info(f"Import of {symbol} from {module} failed with: {exception_text}")
+    except Exception:
+        exception_text = traceback.format_exc()
+        raise
+    if msg is None:
+        msg = f"{module}.{symbol} could not be imported"
+    if alt is None:
+        return UnavailableMeta(symbol, (), {"_msg": msg}), False
+    else:
+        return alt, False
+
+
+def gpu_only_import(module, *, alt=None):
+    """A function used to import modules required only in GPU installs
+
+    This function will attempt to import a module with the given name.
+    This function will attempt to import a symbol with the given name from
+    the given module, but it will not throw an ImportError if the symbol is not
+    found. Instead, it will return a placeholder object which will raise an
+    exception only if used with instructions on installing a GPU build.
+
+    Parameters
+    ----------
+    module: str
+        The name of the module to import.
+    alt: object
+        An optional module to be used in place of the given module if it
+        fails to import in a non-GPU-enabled install
+
+    Returns
+    -------
+    object
+        The imported module, the given alternate, or a class derived from
+        UnavailableMeta.
+    """
+
+    return safe_import(
+        module,
+        msg=f"{module} is not enabled in non GPU-enabled installations or environemnts. {GPU_INSTALL_STRING}",
+        alt=alt,
+    )
+
+
+def gpu_only_import_from(module, symbol, *, alt=None):
+    """A function used to import symbols required only in GPU installs
+
+    This function will attempt to import a module with the given name.
+    This function will attempt to import a symbol with the given name from
+    the given module, but it will not throw an ImportError if the symbol is not
+    found. Instead, it will return a placeholder object which will raise an
+    exception only if used with instructions on installing a GPU build.
+
+    Parameters
+    ----------
+    module: str
+        The name of the module to import.
+    symbol: str
+        The name of the symbol to import.
+    alt: object
+        An optional object to be used in place of the given symbol if it fails
+        to import in a non-GPU-enabled install
+
+    Returns
+    -------
+    object
+        The imported symbol, the given alternate, or a class derived from
+        UnavailableMeta.
+    """
+    return safe_import_from(
+        module,
+        symbol,
+        msg=f"{module}.{symbol} is not enabled in non GPU-enabled installations or environments. {GPU_INSTALL_STRING}",
+        alt=alt,
+    )

nemo/utils/lightning_logger_patch.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2020, 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 logging as _logging
+from logging.handlers import MemoryHandler
+
+import lightning.pytorch as pl
+
+HANDLERS = {}
+PATCHED = False
+
+
+def add_memory_handlers_to_pl_logger():
+    """
+    Adds two MemoryHandlers to pytorch_lightning's logger. These two handlers are essentially message buffers. This
+    function is called in nemo.utils.__init__.py. These handlers are used in add_filehandlers_to_pl_logger to flush
+    buffered messages to files.
+    """
+    if not HANDLERS:
+        HANDLERS["memory_err"] = MemoryHandler(-1)
+        HANDLERS["memory_err"].addFilter(lambda record: record.levelno > _logging.INFO)
+        HANDLERS["memory_all"] = MemoryHandler(-1)
+        pl._logger.addHandler(HANDLERS["memory_err"])
+        pl._logger.addHandler(HANDLERS["memory_all"])
+
+
+def add_filehandlers_to_pl_logger(all_log_file, err_log_file):
+    """
+    Adds two filehandlers to pytorch_lightning's logger. Called in nemo.utils.exp_manager(). The first filehandler
+    logs all messages to all_log_file while the second filehandler logs all WARNING and higher messages to err_log_file.
+    If "memory_err" and "memory_all" exist in HANDLERS, then those buffers are flushed to err_log_file and all_log_file
+    respectively, and then closed.
+    """
+    HANDLERS["file"] = _logging.FileHandler(all_log_file)
+    pl._logger.addHandler(HANDLERS["file"])
+    HANDLERS["file_err"] = _logging.FileHandler(err_log_file)
+    HANDLERS["file_err"].addFilter(lambda record: record.levelno > _logging.INFO)
+    pl._logger.addHandler(HANDLERS["file_err"])
+
+    if HANDLERS.get("memory_all", None):
+        HANDLERS["memory_all"].setTarget(HANDLERS["file"])
+        HANDLERS["memory_all"].close()
+        del HANDLERS["memory_all"]
+    if HANDLERS.get("memory_err", None):
+        HANDLERS["memory_err"].setTarget(HANDLERS["file_err"])
+        HANDLERS["memory_err"].close()
+        del HANDLERS["memory_err"]

nemo/utils/loggers/__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.loggers.clearml_logger import ClearMLLogger, ClearMLParams
+from nemo.utils.loggers.dllogger import DLLogger, DLLoggerParams
+from nemo.utils.loggers.mlflow_logger import MLFlowParams

nemo/utils/loggers/clearml_logger.py

@@ -0,0 +1,191 @@
+# 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
+from argparse import Namespace
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, List, Literal, Mapping, Optional, Union
+
+import pandas as pd
+from lightning.pytorch.callbacks import Checkpoint
+from lightning.pytorch.loggers import Logger
+from lightning.pytorch.utilities.parsing import AttributeDict
+from lightning_utilities.core.apply_func import apply_to_collection
+from omegaconf import DictConfig, ListConfig, OmegaConf
+from torch import Tensor
+
+from nemo.utils import logging
+
+try:
+    from clearml import OutputModel, Task
+
+    HAVE_CLEARML_LOGGER = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_CLEARML_LOGGER = False
+
+
+@dataclass
+class ClearMLParams:  # pylint: disable=C0115
+    project: Optional[str] = None
+    task: Optional[str] = None
+    connect_pytorch: Optional[bool] = False
+    model_name: Optional[str] = None
+    tags: Optional[List[str]] = None
+    log_model: Optional[bool] = False
+    log_cfg: Optional[bool] = False
+    log_metrics: Optional[bool] = False
+
+
+class ClearMLLogger(Logger):  # pylint: disable=C0115
+    @property
+    def name(self) -> str:  # pylint: disable=C0116
+        return self.clearml_task.name
+
+    @property
+    def version(self) -> str:  # pylint: disable=C0116
+        return self.clearml_task.id
+
+    def __init__(
+        self, clearml_cfg: DictConfig, log_dir: str, prefix: str, save_best_model: bool, postfix: str = ".nemo"
+    ) -> None:  # pylint: disable=C0116
+        if not HAVE_CLEARML_LOGGER:
+            raise ImportError(
+                "Found create_clearml_logger is True."
+                "But ClearML not found. Please see the README for installation instructions:"
+                "https://github.com/clearml/clearml"
+            )
+
+        self.clearml_task = None
+        self.clearml_model = None
+        self.clearml_cfg = clearml_cfg
+        self.path_nemo_model = os.path.abspath(
+            os.path.expanduser(os.path.join(log_dir, "checkpoints", prefix + postfix))
+        )
+        self.save_best_model = save_best_model
+        self.prefix = prefix
+        self.previos_best_model_path = None
+        self.last_metrics = None
+        self.save_blocked = True
+
+        self.project_name = os.getenv("CLEARML_PROJECT", clearml_cfg.project if clearml_cfg.project else "NeMo")
+        self.task_name = os.getenv("CLEARML_TASK", clearml_cfg.task if clearml_cfg.task else f"Trainer {self.prefix}")
+
+        tags = ["NeMo"]
+        if clearml_cfg.tags:
+            tags.extend(clearml_cfg.tags)
+
+        self.clearml_task: Task = Task.init(
+            project_name=self.project_name,
+            task_name=self.task_name,
+            auto_connect_frameworks={"pytorch": clearml_cfg.connect_pytorch},
+            output_uri=True,
+            tags=tags,
+        )
+
+        if clearml_cfg.model_name:
+            model_name = clearml_cfg.model_name
+        elif self.prefix:
+            model_name = self.prefix
+        else:
+            model_name = self.task_name
+
+        if clearml_cfg.log_model:
+            self.clearml_model: OutputModel = OutputModel(
+                name=model_name, task=self.clearml_task, tags=tags, framework="NeMo"
+            )
+
+    def log_hyperparams(self, params, *args, **kwargs) -> None:  # pylint: disable=C0116
+        if self.clearml_model and self.clearml_cfg.log_cfg:
+            if isinstance(params, Namespace):
+                params = vars(params)
+            elif isinstance(params, AttributeDict):
+                params = dict(params)
+            params = apply_to_collection(params, (DictConfig, ListConfig), OmegaConf.to_container, resolve=True)
+            params = apply_to_collection(params, Path, str)
+            params = OmegaConf.to_yaml(params)
+            self.clearml_model.update_design(config_text=params)
+
+    def log_metrics(self, metrics: Mapping[str, float], step: Optional[int] = None) -> None:  # pylint: disable=C0116
+        if self.clearml_model and self.clearml_cfg.log_metrics:
+            metrics = {
+                k: {
+                    "value": str(v.item() if type(v) == Tensor else v),
+                    "type": str(type(v.item() if type(v) == Tensor else v)),
+                }
+                for k, v in metrics.items()
+            }
+            self.last_metrics = metrics
+
+    # pylint: disable=C0116
+    def log_table(
+        self,
+        key: str,
+        columns: List[str] = None,
+        data: List[List[Any]] = None,
+        dataframe: Any = None,
+        step: Optional[int] = None,
+    ) -> None:
+        table: Optional[Union[pd.DataFrame, List[List[Any]]]] = None
+
+        if dataframe is not None:
+            table = dataframe
+            if columns is not None:
+                table.columns = columns
+
+        if data is not None:
+            table = data
+            assert len(columns) == len(table[0]), "number of column names should match the total number of columns"
+            table.insert(0, columns)
+
+        if table is not None:
+            self.clearml_task.logger.report_table(title=key, series=key, iteration=step, table_plot=table)
+
+    def after_save_checkpoint(self, checkpoint_callback: Checkpoint) -> None:  # pylint: disable=C0116
+        if self.clearml_model:
+            if self.save_best_model:
+                if self.save_blocked:
+                    self.save_blocked = False
+                    return None
+                if not os.path.exists(checkpoint_callback.best_model_path):
+                    return None
+                if self.previos_best_model_path == checkpoint_callback.best_model_path:
+                    return None
+                self.previos_best_model_path = checkpoint_callback.best_model_path
+            self._log_model(self.path_nemo_model)
+
+    def finalize(self, status: Literal["success", "failed", "aborted"] = "success") -> None:  # pylint: disable=C0116
+        if status == "success":
+            self.clearml_task.mark_completed()
+        elif status == "failed":
+            self.clearml_task.mark_failed()
+        elif status == "aborted":
+            self.clearml_task.mark_stopped()
+
+    def _log_model(self, save_path: str) -> None:  # pylint: disable=C0116
+        if self.clearml_model:
+            if os.path.exists(save_path):
+                self.clearml_model.update_weights(
+                    weights_filename=save_path,
+                    upload_uri=self.clearml_task.storage_uri or self.clearml_task._get_default_report_storage_uri(),
+                    auto_delete_file=False,
+                    is_package=True,
+                )
+
+                if self.clearml_cfg.log_metrics and self.last_metrics:
+                    self.clearml_model.set_all_metadata(self.last_metrics)
+
+                self.save_blocked = True
+            else:
+                logging.warning((f"Logging model enabled, but cant find .nemo file!" f" Path: {save_path}"))

nemo/utils/loggers/dllogger.py

@@ -0,0 +1,104 @@
+# Copyright (c) 2020, 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 argparse import Namespace
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+from lightning.pytorch.loggers import Logger
+from lightning.pytorch.utilities import rank_zero_only
+from lightning.pytorch.utilities.parsing import AttributeDict
+from lightning_utilities.core.apply_func import apply_to_collection
+from omegaconf import DictConfig, ListConfig, OmegaConf
+
+from nemo.utils import logging
+
+try:
+    import dllogger
+    from dllogger import Verbosity
+
+    HAVE_DLLOGGER = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_DLLOGGER = False
+
+try:
+    from lightning.fabric.utilities.logger import _convert_params, _flatten_dict, _sanitize_callable_params
+
+    PL_LOGGER_UTILITIES = True
+except (ImportError, ModuleNotFoundError):
+    PL_LOGGER_UTILITIES = False
+
+
+@dataclass
+class DLLoggerParams:
+    verbose: Optional[bool] = False
+    stdout: Optional[bool] = False
+    json_file: Optional[str] = "./dllogger.json"
+
+
+class DLLogger(Logger):
+    @property
+    def name(self):
+        return self.__class__.__name__
+
+    @property
+    def version(self):
+        return None
+
+    def __init__(self, stdout: bool, verbose: bool, json_file: str):
+        if not HAVE_DLLOGGER:
+            raise ImportError(
+                "DLLogger was not found. Please see the README for installation instructions: "
+                "https://github.com/NVIDIA/dllogger"
+            )
+        if not PL_LOGGER_UTILITIES:
+            raise ImportError(
+                "DLLogger utilities were not found. You probably need to update PyTorch Lightning>=1.9.0. "
+                "pip install pytorch-lightning -U"
+            )
+        verbosity = Verbosity.VERBOSE if verbose else Verbosity.DEFAULT
+        backends = []
+        if json_file:
+            Path(json_file).parent.mkdir(parents=True, exist_ok=True)
+            backends.append(dllogger.JSONStreamBackend(verbosity, json_file))
+        if stdout:
+            backends.append(dllogger.StdOutBackend(verbosity))
+
+        if not backends:
+            logging.warning(
+                "Neither stdout nor json_file DLLogger parameters were specified." "DLLogger will not log anything."
+            )
+        dllogger.init(backends=backends)
+
+    @rank_zero_only
+    def log_hyperparams(self, params, *args, **kwargs):
+        if isinstance(params, Namespace):
+            params = vars(params)
+        elif isinstance(params, AttributeDict):
+            params = dict(params)
+        params = apply_to_collection(params, (DictConfig, ListConfig), OmegaConf.to_container, resolve=True)
+        params = apply_to_collection(params, Path, str)
+        params = _sanitize_callable_params(_flatten_dict(_convert_params(params)))
+        dllogger.log(step="PARAMETER", data=params)
+
+    @rank_zero_only
+    def log_metrics(self, metrics, step=None):
+        if step is None:
+            step = tuple()
+
+        dllogger.log(step=step, data=metrics)
+
+    def save(self):
+        dllogger.flush()

nemo/utils/loggers/mlflow_logger.py

@@ -0,0 +1,36 @@
+# 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 dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+
+@dataclass
+class MLFlowParams:
+    """ML Flow Configuration Dataclass."""
+
+    # name of experiment, if none, defaults to the globally set experiment name
+    experiment_name: Optional[str] = None
+    run_name: Optional[str] = None
+    # if no run_name is set, it's set by version
+    # local or remote tracking seerver. If tracking_uri is not set, it defaults to save_dir
+    tracking_uri: Optional[str] = None
+    tags: Optional[Dict[str, Any]] = None
+    save_dir: Optional[str] = "./mlruns"
+    prefix: str = ""
+    artifact_location: Optional[str] = None
+    # provide run_id if resuming a previously started run
+    run_id: Optional[str] = None
+    # Log checkpoints created by ModelCheckpoint as MLFlow artifacts.
+    log_model: bool = False

nemo/utils/mcore_logger.py

@@ -0,0 +1,31 @@
+# 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 logging as _logging
+
+from nemo.utils import logging as nemo_logger
+
+
+def add_handlers_to_mcore_logger():
+    """Add NeMo handlers to MCore loggers.
+
+    MCore doesn't have and handlers for loggers (see
+    https://docs.python.org/3/howto/logging-cookbook.html#adding-handlers-other-than-nullhandler-to-a-logger-in-a-library
+    for a rationale). We have to add handlers explicitly.
+    """
+    mcore_logger = _logging.getLogger('megatron.core')
+    for handler in nemo_logger._handlers.values():
+        mcore_logger.addHandler(handler)
+    mcore_logger.propagate = False
+    mcore_logger.setLevel(nemo_logger._logger.level)

nemo/utils/megatron_utils.py

@@ -0,0 +1,328 @@
+# 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.
+
+# flake8: noqa
+# pylint: skip-file
+
+"""Utilities for models."""
+import itertools
+from typing import Dict, Iterator, List, Optional, Union
+
+import torch
+from torch import Tensor
+
+from nemo.utils import logging, logging_mode
+
+try:
+    from apex.transformer.enums import AttnMaskType
+
+    HAVE_APEX = True
+
+except (ImportError, ModuleNotFoundError):
+
+    HAVE_APEX = False
+
+try:
+    from megatron.core import parallel_state
+
+    HAVE_MEGATRON_CORE = True
+
+except (ImportError, ModuleNotFoundError):
+
+    HAVE_MEGATRON_CORE = False
+
+
+def ApproxGELUActivation(input: Tensor):
+    """
+    Applies GELU approximation that is fast but somewhat inaccurate. See: https://github.com/hendrycks/GELUs
+    """
+    return input * torch.sigmoid(1.702 * input)
+
+
+class ApexGuardDefaults(object):
+    """
+    This class can be used to replace missing classes when apex is missing.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def __getattr__(self, item):
+        return None
+
+
+def init_method_kaiming_uniform(val):
+    def init_(tensor):
+        return torch.nn.init.kaiming_uniform_(tensor, a=val)
+
+    return init_
+
+
+def init_method_const(val):
+    def init_(tensor):
+        return torch.nn.init.constant_(tensor, val)
+
+    return init_
+
+
+def init_method_normal(sigma):
+    """Init method based on N(0, sigma)."""
+
+    def init_(tensor):
+        return torch.nn.init.normal_(tensor, mean=0.0, std=sigma)
+
+    return init_
+
+
+def average_losses_across_data_parallel_group(losses):
+    """Reduce a tensor of losses across all GPUs."""
+    averaged_losses = torch.cat([loss.clone().detach().view(1) for loss in losses])
+    torch.distributed.all_reduce(averaged_losses, group=parallel_state.get_data_parallel_group())
+    averaged_losses = averaged_losses / torch.distributed.get_world_size(
+        group=parallel_state.get_data_parallel_group()
+    )
+
+    return averaged_losses
+
+
+def get_ltor_masks_and_position_ids(
+    data, eod_token, reset_position_ids, reset_attention_mask, eod_mask_loss, compute_attention_mask=True
+):
+    """Build masks and position id for left to right model."""
+
+    # Extract batch size and sequence length.
+    micro_batch_size, seq_length = data.size()
+
+    # Attention mask (lower triangular).
+    if reset_attention_mask:
+        att_mask_batch = micro_batch_size
+    else:
+        att_mask_batch = 1
+
+    attention_mask = None
+    if compute_attention_mask:
+        attention_mask = torch.tril(torch.ones((att_mask_batch, seq_length, seq_length), device=data.device)).view(
+            att_mask_batch, 1, seq_length, seq_length
+        )
+
+    # Loss mask.
+    loss_mask = torch.ones(data.size(), dtype=torch.float, device=data.device)
+    if eod_mask_loss:
+        loss_mask[data == eod_token] = 0.0
+
+    # Position ids.
+    position_ids = torch.arange(seq_length, dtype=torch.long, device=data.device)
+    position_ids = position_ids.unsqueeze(0).repeat(micro_batch_size, 1)
+    # We need to clone as the ids will be modifed based on batch index.
+    if reset_position_ids:
+        position_ids = position_ids.clone()
+
+    if reset_position_ids or reset_attention_mask:
+        # Loop through the batches:
+        for b in range(micro_batch_size):
+
+            # Find indecies where EOD token is.
+            eod_index = position_ids[b, data[b] == eod_token]
+            # Detach indecies from positions if going to modify positions.
+            if reset_position_ids:
+                eod_index = eod_index.clone()
+
+            # Loop through EOD indicies:
+            prev_index = 0
+            for j in range(eod_index.size()[0]):
+                i = eod_index[j]
+                # Mask attention loss.
+                if reset_attention_mask:
+                    attention_mask[b, 0, (i + 1) :, : (i + 1)] = 0
+                # Reset positions.
+                if reset_position_ids:
+                    position_ids[b, (i + 1) :] -= i + 1 - prev_index
+                    prev_index = i + 1
+
+    if compute_attention_mask:
+        # Convert attention mask to binary:
+        attention_mask = attention_mask < 0.5
+
+    return attention_mask, loss_mask, position_ids
+
+
+def build_position_ids(token_ids):
+    # Create position ids
+    seq_length = token_ids.size(1)
+    position_ids = torch.arange(seq_length, dtype=torch.long, device=token_ids.device)
+    position_ids = position_ids.unsqueeze(0).expand_as(token_ids).clone()
+
+    return position_ids
+
+
+def make_attention_mask_3d(source_mask, target_mask):
+    """
+    Returns a 3-dimensional (3-D) attention mask
+    :param source_block: 2-D array
+    :param target_block: 2-D array
+    """
+    mask = target_mask[:, None, :] * source_mask[:, :, None]
+    return mask
+
+
+def make_inference_attention_mask_3d(source_block, target_block, pad_id):
+    """
+    Returns a 3-dimensional (3-D) attention mask
+    :param source_block: 2-D array
+    :param target_block: 2-D array
+    """
+    # mask = (target_block[:, None, :] != pad_id) * (source_block[:, :, None] != pad_id)
+    return make_attention_mask_3d(source_block != pad_id, target_block != pad_id)
+
+
+def make_inference_history_mask_3d(block):
+    batch, length = block.shape
+    arange = torch.arange(length, device=block.device)
+    history_mask = (arange[None,] <= arange[:, None])[None,]
+    history_mask = history_mask.expand(batch, length, length)
+    return history_mask
+
+
+def build_attention_mask_3d_padding(source_mask, target_mask):
+    """
+    Returns a 3D joint attention mask for Megatron given two 2D masks
+    :param source_mask - True for non-masked, else masked [batch, src length]
+    :param target_mask - True for non-masked, else masked [batch, tgt length]
+    """
+    mask = make_attention_mask_3d(source_mask, target_mask)
+    # invert mask for Megatron
+    return mask < 0.5
+
+
+def build_attention_mask_3d_causal(source_mask, target_mask):
+    """
+    Returns a 3D joint attention mask for Megatron given two 2D masks
+    :param source_mask - True for non-masked, else masked [batch, src length]
+    :param target_mask - True for non-masked, else masked [batch, tgt length]
+    """
+    causal_mask = make_inference_history_mask_3d(target_mask)
+    mask = make_attention_mask_3d(source_mask, target_mask)
+    mask = mask * causal_mask
+    # invert mask for Megatron
+    return mask < 0.5
+
+
+def build_attention_mask_3d(source_mask, target_mask, attn_mask_type):
+    """
+    Returns a 3D attention mask for Megatron given two 2D masks
+    :param source_mask - < 0.5 for non-masked, else masked [batch, src length]
+    :param target_mask - < 0.5 for non-masked, else masked [batch, tgt length]
+    :param attn_mask_type - AttnMaskType enum
+    """
+    if attn_mask_type == AttnMaskType.padding:
+        mask = build_attention_mask_3d_padding(source_mask, target_mask)
+    elif attn_mask_type == AttnMaskType.causal:
+        mask = build_attention_mask_3d_causal(source_mask, target_mask)
+    else:
+        raise ValueError(f"Unsupported attention mask attn_mask_type = {attn_mask_type}")
+
+    return mask
+
+
+def split_list(inputs, num_chunks, enforce_divisible_batch: Optional[bool] = True):
+    """
+    Split a list into equal sized chunks
+    """
+    chunk_size = len(inputs) // num_chunks
+    if enforce_divisible_batch:
+        assert len(inputs) % chunk_size == 0, "Issue with batch size configuration!"
+    return [inputs[i : i + chunk_size] for i in range(0, len(inputs), chunk_size)]
+
+
+def get_iterator_k_split(
+    batch: Union[Dict, List[torch.Tensor]], num_microbatches: int, enforce_divisible_batch: Optional[bool] = True
+) -> Iterator:
+    """
+    Split a batch into k microbatches, where the batch size is divisible by k. Batch could be
+    a dictionary of tensors or a list of tensors. A dictionary batch could also have items of List type,
+    as long as the length of that list is the same as the batch size.
+    """
+    if isinstance(batch, dict):
+        discard_items = [k for k, v in batch.items() if not isinstance(v, (torch.Tensor, list))]
+        if len(discard_items) > 0:
+            logging.warning(
+                f"Only support splitting torch.Tensor and List[torch.Tensor]. Discarding the following keys from the batch: {discard_items}",
+                mode=logging_mode.ONCE,
+            )
+
+        batch = {k: v for k, v in batch.items() if isinstance(v, (torch.Tensor, list))}
+        tensor_items = {k: v for k, v in batch.items() if isinstance(v, torch.Tensor)}
+        list_items = {k: v for k, v in batch.items() if isinstance(v, list)}
+
+        # Split tensor items
+        items = list(tensor_items.items())
+
+        if enforce_divisible_batch:
+            if items[0][1].shape[0] % num_microbatches != 0:
+                raise ValueError(
+                    f"Issue with batch size configuration: batch size {items[0][1].shape[0]} is not divisible by {num_microbatches}!"
+                )
+
+        split_batch = [torch.tensor_split(item[1], num_microbatches, dim=0) for item in items]
+        # handle the case where the batch size from dynamic bucketting is not divisible
+        if items[0][1].shape[0] % num_microbatches != 0:
+            chunk_size = split_batch[0][-1].shape[0]
+            split_batch = [[j[:chunk_size] for j in i] for i in split_batch]
+
+        if len(list_items) == 0:
+            # Only have tensor items
+            microbatches = [
+                [(items[i][0], split_batch[i][j]) for i in range(len(items))] for j in range(num_microbatches)
+            ]
+        else:
+            # Split list items
+            list_items = list(list_items.items())
+            split_list_batch = [
+                split_list(item[1], num_microbatches, enforce_divisible_batch=enforce_divisible_batch)
+                for item in list_items
+            ]
+            # Merge tensor and list items
+            all_keys = [item[0] for item in items] + [item[0] for item in list_items]
+            all_split_batch = split_batch + split_list_batch
+            microbatches = [
+                [(all_keys[i], all_split_batch[i][j]) for i in range(len(all_keys))] for j in range(num_microbatches)
+            ]
+        microbatches = [dict(elem) for elem in microbatches]
+    else:
+        # Split a list of torch tensors
+        assert batch[0].shape[0] % num_microbatches == 0, "Issue with batch size configuration!"
+        split_batch = []
+        for item in batch:
+            if torch.is_tensor(item):
+                split_batch.append(torch.tensor_split(item, num_microbatches, dim=0))
+            elif isinstance(item, list):
+                if isinstance(item[0], torch.Tensor):
+                    split_tensors = [torch.tensor_split(elem, num_microbatches, dim=0) for elem in item]
+                    split_tuple = []
+                    for mbi in range(num_microbatches):
+                        split_tuple.append([split_tensors[i][mbi] for i in range(len(split_tensors))])
+                    split_tuple = tuple(split_tuple)
+                    split_batch.append(split_tuple)
+                else:
+                    split_batch.append(split_list(item, num_microbatches))
+            elif item is None:
+                split_batch.append(item)
+            else:
+                raise ValueError(f"Unsupported item type: {type(item)}")
+
+        microbatches = [
+            [elem[i] if elem is not None else elem for elem in split_batch] for i in range(num_microbatches)
+        ]
+
+    return itertools.chain(microbatches)

nemo/utils/metaclasses.py

@@ -0,0 +1,39 @@
+# Copyright (c) 2020, 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 threading
+
+
+class Singleton(type):
+    """Implementation of a generic, tread-safe singleton meta-class.
+    Can be used as meta-class, i.e. will create
+    """
+
+    # List of instances - one per class.
+    __instances = {}
+    # Lock used for accessing the instance.
+    __lock = threading.Lock()
+
+    def __call__(cls, *args, **kwargs):
+        """Returns singleton instance. A thread safe implementation."""
+        if cls not in cls.__instances:
+            # Enter critical section.
+            with cls.__lock:
+                # Check once again.
+                if cls not in cls.__instances:
+                    # Create a new object instance - one per class.
+                    cls.__instances[cls] = super(Singleton, cls).__call__(*args, **kwargs)
+        # Return the instance.
+        return cls.__instances[cls]

nemo/utils/model_utils.py

@@ -0,0 +1,768 @@
+# Copyright (c) 2020, 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 contextlib
+import copy
+import fnmatch
+import importlib
+import os
+import shutil
+import tarfile
+import tempfile
+from dataclasses import dataclass, is_dataclass
+from enum import Enum
+from functools import lru_cache
+from pathlib import Path, PurePosixPath
+from typing import TYPE_CHECKING, List, Optional, Tuple, Type, Union
+
+import wrapt
+from omegaconf import DictConfig, ListConfig, OmegaConf
+from omegaconf import errors as omegaconf_errors
+from packaging import version
+
+from nemo.utils import AppState, logging
+from nemo.utils.data_utils import (  # imported for compatibility: model_utils.resolve_cache_dir()  # noqa: F401  # pylint: disable=unused-import,line-too-long
+    is_datastore_path,
+    resolve_cache_dir,
+)
+from nemo.utils.tar_utils import TarPathTraversalError, safe_extract
+
+if TYPE_CHECKING:
+    import lightning.pytorch as pl
+
+    from nemo.core.classes import ModelPT, PretrainedModelInfo
+    from nemo.core.config.modelPT import NemoConfig
+
+
+MODEL_CONFIG = "model_config.yaml"
+_VAL_TEST_FASTPATH_KEY = 'ds_item'
+
+
+class ArtifactPathType(Enum):
+    """
+    ArtifactPathType refers to the type of the path that the artifact is located at.
+
+    LOCAL_PATH: A user local filepath that exists on the file system.
+    TAR_PATH: A (generally flattened) filepath that exists inside of an archive (that may have its own full path).
+    """
+
+    LOCAL_PATH = 0
+    TAR_PATH = 1
+
+
+@dataclass
+class ArtifactItem:
+    path: str = ""
+    path_type: ArtifactPathType = ArtifactPathType.LOCAL_PATH
+    hashed_path: Optional[str] = None
+
+
+def detect_prefix(names: List[str]) -> str:
+    """Detect model config prefix for a list of file names.
+
+    Useful to identify prefix used within .nemo tarball checkpoint."""
+    model_config = fnmatch.filter(names, f"*{MODEL_CONFIG}")
+    assert len(model_config) == 1, f"Exactly one model config path expected, found: {model_config}."
+    prefix = model_config[0].removesuffix(MODEL_CONFIG)
+    return prefix
+
+
+def load_config(model_file: str) -> DictConfig:
+    """Load model config from extracted directory or '.nemo' tarball."""
+    if os.path.isfile(model_file):
+        with tempfile.TemporaryDirectory() as tmp, tarfile.open(model_file, "r:") as tar:
+            prefix = detect_prefix(tar.getnames())
+            safe_extract(tar, tmp, members=[f"{prefix}{MODEL_CONFIG}"])
+            model_config = OmegaConf.load(os.path.join(tmp, MODEL_CONFIG))
+    elif os.path.isdir(model_file):
+        model_config = OmegaConf.load(os.path.join(model_file, MODEL_CONFIG))
+    else:
+        raise FileNotFoundError(model_file)
+
+    return model_config
+
+
+def _validate_artifact_path(path: str) -> PurePosixPath:
+    artifact_path = PurePosixPath(path)
+    if artifact_path.is_absolute() or ".." in artifact_path.parts:
+        raise TarPathTraversalError(f"Unsafe artifact path: {path}")
+    return artifact_path
+
+
+def unwrap_model(model, module_instances: Optional[Union[Type, Tuple[Type]]] = None):
+    """Unwrap model from wrapper classes like Float16Module, for example."""
+
+    # TODO: Import this from megatron.core once moved there from megatron.training.
+    return_list = True
+    if not isinstance(model, list):
+        model = [model]
+        return_list = False
+    unwrapped_model = []
+    for model_module in model:
+        if module_instances:
+            while isinstance(model_module, module_instances):
+                model_module = model_module.module
+        else:  # remove any wrappers that have a '.module' attribute
+            while hasattr(model_module, "module"):
+                model_module = model_module.module
+        unwrapped_model.append(model_module)
+    if not return_list:
+        return unwrapped_model[0]
+    return unwrapped_model
+
+
+def param_is_not_shared(param):
+    return not hasattr(param, 'shared') or not param.shared
+
+
+def resolve_dataset_name_from_cfg(cfg: 'DictConfig') -> Optional[str]:
+    """
+    Parses items of the provided sub-config to find the first potential key that
+    resolves to an existing file or directory.
+
+    # Fast-path Resolution
+    In order to handle cases where we need to resolve items that are not paths, a fastpath
+    key can be provided as defined in the global `_VAL_TEST_FASTPATH_KEY`.
+
+    This key can be used in two ways :
+
+    ## _VAL_TEST_FASTPATH_KEY points to another key in the config
+
+    If this _VAL_TEST_FASTPATH_KEY points to another key in this config itself,
+    then we assume we want to loop through the values of that key.
+
+    This allows for any key in the config to become a fastpath key.
+
+    Example:
+    validation_ds:
+        splits: "val"
+        ...
+        <_VAL_TEST_FASTPATH_KEY>: "splits"  <-- this points to the key name "splits"
+
+    Then we can write the following when overriding in hydra:
+    ```python
+    python train_file.py ... \
+        model.validation_ds.splits=[val1, val2, dev1, dev2] ...
+    ```
+
+    ## _VAL_TEST_FASTPATH_KEY itself acts as the resolved key
+
+    If this _VAL_TEST_FASTPATH_KEY does not point to another key in the config, then
+    it is assumed that the items of this key itself are used for resolution.
+
+    Example:
+    validation_ds:
+        ...
+        <_VAL_TEST_FASTPATH_KEY>: "val"  <-- this points to the key name "splits"
+
+    Then we can write the following when overriding in hydra:
+    ```python
+    python train_file.py ... \
+        model.validation_ds.<_VAL_TEST_FASTPATH_KEY>=[val1, val2, dev1, dev2] ...
+    ```
+
+    # IMPORTANT NOTE:
+    It <can> potentially mismatch if there exist more than 2 valid paths, and the
+    first path does *not* resolve the the path of the data file (but does resolve to
+    some other valid path).
+
+    To avoid this side-effect, place the data path as the first item on the config file.
+
+    Args:
+        cfg: DictConfig (Sub-config) that should be parsed.
+
+    Returns:
+        A str representing the `key` of the config which hosts the filepath(s),
+        or None in case path could not be resolved.
+    """
+    if _VAL_TEST_FASTPATH_KEY in cfg and cfg[_VAL_TEST_FASTPATH_KEY] is not None:
+        fastpath_key = cfg[_VAL_TEST_FASTPATH_KEY]
+
+        if isinstance(fastpath_key, str) and fastpath_key in cfg:
+            return cfg[fastpath_key]
+        else:
+            return _VAL_TEST_FASTPATH_KEY
+
+    for key, value in cfg.items():
+        if type(value) in [list, tuple, ListConfig]:
+            # Count the number of valid paths in the list
+            values_are_paths = 0
+            for val_i in value:
+                val_i = str(val_i)
+                if os.path.exists(val_i) or os.path.isdir(val_i) or is_datastore_path(val_i):
+                    values_are_paths += 1
+                else:
+                    # reset counter and break inner loop
+                    break
+
+            if values_are_paths == len(value):
+                return key
+
+        else:
+            if os.path.exists(str(value)) or os.path.isdir(str(value)) or is_datastore_path(str(value)):
+                return key
+
+    return None
+
+
+def parse_dataset_as_name(name: str) -> str:
+    """
+    Constructs a valid prefix-name from a provided file path.
+
+    Args:
+        name: str path to some valid data/manifest file or a python object that
+            will be used as a name for the data loader (via str() cast).
+
+    Returns:
+        str prefix used to identify uniquely this data/manifest file.
+    """
+    if os.path.exists(str(name)) or os.path.isdir(str(name)) or is_datastore_path(str(name)):
+        name = Path(name).stem
+    else:
+        name = str(name)
+
+    # cleanup name
+    name = name.replace('-', '_')
+
+    if 'manifest' in name:
+        name = name.replace('manifest', '')
+
+    if 'dataset' in name:
+        name = name.replace('dataset', '')
+
+    # Test if the manifes/dataset name was simply `manifest.yaml` or `dataset.yaml`: Invalid names.
+    if name == '':
+        raise ValueError(
+            "Provided dataset / manifest filename was `manifest.json` or `dataset.json`.\n"
+            "Such a name is invalid, since multiple datasets/manifests can share the same name,\n"
+            "thereby overriding their results during logging. Please pick a more discriptive filename \n"
+            "for the provided dataset / manifest file."
+        )
+
+    if '_' != name[-1]:
+        name = name + '_'
+
+    return name
+
+
+def unique_names_check(name_list: Optional[List[str]]):
+    """
+    Performs a uniqueness check on the name list resolved, so that it can warn users
+    about non-unique keys.
+
+    Args:
+        name_list: List of strings resolved for data loaders.
+    """
+    if name_list is None:
+        return
+
+    # Name uniqueness checks
+    names = set()
+    for name in name_list:
+        if name in names:
+            logging.warning(
+                "Name resolution has found more than one data loader having the same name !\n"
+                "In such cases, logs will nor be properly generated. "
+                "Please rename the item to have unique names.\n"
+                f"Resolved name : {name}"
+            )
+        else:
+            names.add(name)  # we need just hash key check, value is just a placeholder
+
+
+def resolve_validation_dataloaders(model: 'ModelPT'):
+    """
+    Helper method that operates on the ModelPT class to automatically support
+    multiple dataloaders for the validation set.
+
+    It does so by first resolving the path to one/more data files via `resolve_dataset_name_from_cfg()`.
+    If this resolution fails, it assumes the data loader is prepared to manually support / not support
+    multiple data loaders and simply calls the appropriate setup method.
+
+    If resolution succeeds:
+        Checks if provided path is to a single file or a list of files.
+        If a single file is provided, simply tags that file as such and loads it via the setup method.
+        If multiple files are provided:
+            Inject a new manifest path at index "i" into the resolved key.
+            Calls the appropriate setup method to set the data loader.
+            Collects the initialized data loader in a list and preserves it.
+            Once all data loaders are processed, assigns the list of loaded loaders to the ModelPT.
+            Finally assigns a list of unique names resolved from the file paths to the ModelPT.
+
+    Args:
+        model: ModelPT subclass, which requires >=1 Validation Dataloaders to be setup.
+    """
+    cfg = copy.deepcopy(model._cfg)
+    dataloaders = []
+
+    # process val_loss_idx
+    if 'val_dl_idx' in cfg.validation_ds:
+        cfg = OmegaConf.to_container(cfg)
+        val_dl_idx = cfg['validation_ds'].pop('val_dl_idx')
+        cfg = OmegaConf.create(cfg)
+    else:
+        val_dl_idx = 0
+
+    # Set val_loss_idx
+    model._val_dl_idx = val_dl_idx
+
+    ds_key = resolve_dataset_name_from_cfg(cfg.validation_ds)
+
+    if ds_key is None or val_dl_idx < 0:
+        logging.debug(
+            "Could not resolve file path from provided config - {}. "
+            "Disabling support for multi-dataloaders.".format(cfg.validation_ds)
+        )
+
+        model.setup_validation_data(cfg.validation_ds)
+        return
+
+    ds_values = cfg.validation_ds[ds_key]
+
+    if isinstance(ds_values, (list, tuple, ListConfig)):
+
+        for ds_value in ds_values:
+            if isinstance(ds_value, (dict, DictConfig)):
+                # this is a nested dataset
+                cfg.validation_ds = ds_value
+            else:
+                cfg.validation_ds[ds_key] = ds_value
+
+            model.setup_validation_data(cfg.validation_ds)
+            dataloaders.append(model._validation_dl)
+
+        model._validation_dl = dataloaders
+        if len(ds_values) > 0 and isinstance(ds_values[0], (dict, DictConfig)):
+            # using the name of each of the nested dataset
+            model._validation_names = [ds.name for ds in ds_values]
+        else:
+            ds_names = cfg.validation_ds.get('name', [])
+            if len(ds_names) > 0:
+                if len(ds_names) != len(ds_values):
+                    raise ValueError(
+                        f"Number of names ({len(ds_names)}) does not match number of "
+                        f"datasets ({len(ds_values)}). Got {ds_names} and {ds_values}"
+                    )
+                model._validation_names = [parse_dataset_as_name(n) for n in ds_names]
+            else:
+                model._validation_names = [parse_dataset_as_name(ds) for ds in ds_values]
+        unique_names_check(name_list=model._validation_names)
+
+        return
+
+    else:
+        model.setup_validation_data(cfg.validation_ds)
+        ds_names = cfg.validation_ds.get('name', None)
+        if ds_names is not None:
+            if not isinstance(ds_names, str):
+                raise ValueError(f"`name` must be a string for single manifest, got {ds_names}")
+            model._validation_names = [parse_dataset_as_name(ds_names)]
+        else:
+            model._validation_names = [parse_dataset_as_name(ds_values)]
+        unique_names_check(name_list=model._validation_names)
+
+
+def resolve_test_dataloaders(model: 'ModelPT'):
+    """
+    Helper method that operates on the ModelPT class to automatically support
+    multiple dataloaders for the test set.
+
+    It does so by first resolving the path to one/more data files via `resolve_dataset_name_from_cfg()`.
+    If this resolution fails, it assumes the data loader is prepared to manually support / not support
+    multiple data loaders and simply calls the appropriate setup method.
+
+    If resolution succeeds:
+        Checks if provided path is to a single file or a list of files.
+        If a single file is provided, simply tags that file as such and loads it via the setup method.
+        If multiple files are provided:
+            Inject a new manifest path at index "i" into the resolved key.
+            Calls the appropriate setup method to set the data loader.
+            Collects the initialized data loader in a list and preserves it.
+            Once all data loaders are processed, assigns the list of loaded loaders to the ModelPT.
+            Finally assigns a list of unique names resolved from the file paths to the ModelPT.
+
+    Args:
+        model: ModelPT subclass, which requires >=1 Test Dataloaders to be setup.
+    """
+    cfg = copy.deepcopy(model._cfg)
+    dataloaders = []
+
+    # process test_loss_idx
+    if 'test_dl_idx' in cfg.test_ds:
+        cfg = OmegaConf.to_container(cfg)
+        test_dl_idx = cfg['test_ds'].pop('test_dl_idx')
+        cfg = OmegaConf.create(cfg)
+    else:
+        test_dl_idx = 0
+
+    # Set val_loss_idx
+    model._test_dl_idx = test_dl_idx
+
+    ds_key = resolve_dataset_name_from_cfg(cfg.test_ds)
+
+    if ds_key is None:
+        logging.debug(
+            "Could not resolve file path from provided config - {}. "
+            "Disabling support for multi-dataloaders.".format(cfg.test_ds)
+        )
+
+        model.setup_test_data(cfg.test_ds)
+        return
+
+    ds_values = cfg.test_ds[ds_key]
+
+    if isinstance(ds_values, (list, tuple, ListConfig)):
+
+        for ds_value in ds_values:
+            if isinstance(ds_value, (dict, DictConfig)):
+                # this is a nested dataset
+                cfg.test_ds = ds_value
+            else:
+                cfg.test_ds[ds_key] = ds_value
+
+            model.setup_test_data(cfg.test_ds)
+            dataloaders.append(model._test_dl)
+
+        model._test_dl = dataloaders
+        if len(ds_values) > 0 and isinstance(ds_values[0], (dict, DictConfig)):
+            # using the name of each of the nested dataset
+            model._test_names = [ds.name for ds in ds_values]
+        else:
+            ds_names = cfg.test_ds.get('name', [])
+            if len(ds_names) > 0:
+                if len(ds_names) != len(ds_values):
+                    raise ValueError(
+                        f"Number of names ({len(ds_names)}) does not match number of "
+                        f"datasets ({len(ds_values)}). Got {ds_names} and {ds_values}"
+                    )
+                model._test_names = [parse_dataset_as_name(n) for n in ds_names]
+            else:
+                model._test_names = [parse_dataset_as_name(ds) for ds in ds_values]
+
+        unique_names_check(name_list=model._test_names)
+        return
+
+    else:
+        model.setup_test_data(cfg.test_ds)
+        ds_names = cfg.test_ds.get('name', None)
+        if ds_names is not None:
+            if not isinstance(ds_names, str):
+                raise ValueError(f"`name` must be a string for single manifest, got {ds_names}")
+            model._test_names = [parse_dataset_as_name(ds_names)]
+        else:
+            model._test_names = [parse_dataset_as_name(ds_values)]
+
+        unique_names_check(name_list=model._test_names)
+
+
+@wrapt.decorator
+def wrap_training_step(wrapped, instance: 'pl.LightningModule', args, kwargs):
+    output_dict = wrapped(*args, **kwargs)
+
+    if isinstance(output_dict, dict) and output_dict is not None and 'log' in output_dict:
+        log_dict = output_dict.pop('log')
+        instance.log_dict(log_dict, on_step=True)
+
+    return output_dict
+
+
+def convert_model_config_to_dict_config(cfg: Union['DictConfig', 'NemoConfig']) -> 'DictConfig':
+    """
+    Converts its input into a standard DictConfig.
+    Possible input values are:
+    -   DictConfig
+    -   A dataclass which is a subclass of NemoConfig
+
+    Args:
+        cfg: A dict-like object.
+
+    Returns:
+        The equivalent DictConfig
+    """
+    if not isinstance(cfg, (OmegaConf, DictConfig)) and is_dataclass(cfg):
+        cfg = OmegaConf.structured(cfg)
+
+    if not isinstance(cfg, DictConfig):
+        raise ValueError(f"cfg constructor argument must be of type DictConfig/dict but got {type(cfg)} instead.")
+
+    config = OmegaConf.to_container(cfg, resolve=True)
+    config = OmegaConf.create(config)
+
+    return config
+
+
+def _convert_config(cfg: 'OmegaConf'):
+    """Recursive function convertint the configuration from old hydra format to the new one."""
+    # Get rid of cls -> _target_.
+    if 'cls' in cfg and '_target_' not in cfg:
+        cfg._target_ = cfg.pop('cls')
+
+    # Get rid of params.
+    if 'params' in cfg:
+        params = cfg.pop('params')
+        for param_key, param_val in params.items():
+            cfg[param_key] = param_val
+
+    # Recursion.
+    try:
+        for _, sub_cfg in cfg.items():
+            if isinstance(sub_cfg, (dict, DictConfig)):
+                _convert_config(sub_cfg)
+    except omegaconf_errors.OmegaConfBaseException as e:
+        logging.warning(f"Skipped conversion for config/subconfig:\n{cfg}\n Reason: {e}.")
+
+
+def maybe_update_config_version(cfg: 'DictConfig', make_copy: bool = True):
+    """
+    Recursively convert Hydra 0.x configs to Hydra 1.x configs.
+
+    Changes include:
+    -   `cls` -> `_target_`.
+    -   `params` -> drop params and shift all arguments to parent.
+    -   `target` -> `_target_` cannot be performed due to ModelPT injecting `target` inside class.
+
+    Args:
+        cfg: Any Hydra compatible DictConfig
+        make_copy: bool to indicating if the config should be copied before updating
+
+    Returns:
+        An updated DictConfig that conforms to Hydra 1.x format.
+    """
+    if cfg is not None and not isinstance(cfg, DictConfig):
+        try:
+            temp_cfg = OmegaConf.create(cfg)
+            cfg = temp_cfg
+        except omegaconf_errors.OmegaConfBaseException:
+            # Cannot be cast to DictConfig, skip updating.
+            return cfg
+
+    # Make a copy if requested
+    if make_copy:
+        cfg = copy.deepcopy(cfg)
+
+    OmegaConf.set_struct(cfg, False)
+
+    # Convert config
+    _convert_config(cfg)
+
+    OmegaConf.set_struct(cfg, True)
+
+    return cfg
+
+
+@lru_cache(maxsize=1024)
+def import_class_by_path(path: str):
+    """
+    Recursive import of class by path string.
+    """
+    paths = path.split('.')
+    path = ".".join(paths[:-1])
+    class_name = paths[-1]
+    mod = __import__(path, fromlist=[class_name])
+    mod = getattr(mod, class_name)
+    return mod
+
+
+def resolve_subclass_pretrained_model_info(base_class) -> List['PretrainedModelInfo']:
+    """
+    Recursively traverses the inheritance graph of subclasses to extract all pretrained model info.
+    First constructs a set of unique pretrained model info by performing DFS over the inheritance graph.
+    All model info belonging to the same class is added together.
+
+    Args:
+        base_class: The root class, whose subclass graph will be traversed.
+
+    Returns:
+        A list of unique pretrained model infos belonging to all of the inherited subclasses of
+        this baseclass.
+    """
+    list_of_models = set()
+
+    def recursive_subclass_walk(cls):
+        for subclass in cls.__subclasses__():
+            # step into its immediate subclass
+            recursive_subclass_walk(subclass)
+
+            subclass_models = subclass.list_available_models()
+
+            if subclass_models is not None and len(subclass_models) > 0:
+                # Inject subclass info into pretrained model info
+                # if not already overriden by subclass
+                for model_info in subclass_models:
+                    # If subclass manually injects class_, dont override.
+                    if model_info.class_ is None:
+                        model_info.class_ = subclass
+
+                for model_info in subclass_models:
+                    list_of_models.add(model_info)
+
+    recursive_subclass_walk(base_class)
+
+    list_of_models = list(sorted(list_of_models))
+    return list_of_models
+
+
+def check_lib_version(lib_name: str, checked_version: str, operator) -> Tuple[Optional[bool], str]:
+    """
+    Checks if a library is installed, and if it is, checks the operator(lib.__version__, checked_version) as a result.
+    This bool result along with a string analysis of result is returned.
+
+    If the library is not installed at all, then returns None instead, along with a string explaining
+    that the library is not installed
+
+    Args:
+        lib_name: lower case str name of the library that must be imported.
+        checked_version: semver string that is compared against lib.__version__.
+        operator: binary callable function func(a, b) -> bool; that compares lib.__version__ against version in
+            some manner. Must return a boolean.
+
+    Returns:
+        A tuple of results:
+        -   Bool or None. Bool if the library could be imported, and the result of
+            operator(lib.__version__, checked_version) or False if __version__ is not implemented in lib.
+            None is passed if the library is not installed at all.
+        -   A string analysis of the check.
+    """
+    try:
+        if '.' in lib_name:
+            mod = import_class_by_path(lib_name)
+        else:
+            mod = importlib.import_module(lib_name)
+
+        if hasattr(mod, '__version__'):
+            lib_ver = version.Version(mod.__version__)
+            match_ver = version.Version(checked_version)
+
+            if operator(lib_ver, match_ver):
+                msg = f"Lib {lib_name} version is satisfied !"
+                return True, msg
+            else:
+                msg = (
+                    f"Lib {lib_name} version ({lib_ver}) is not {operator.__name__} "
+                    f"than required version {checked_version}.\n"
+                    f"Please upgrade the lib using either pip or conda to the latest version."
+                )
+                return False, msg
+        else:
+            msg = (
+                f"Lib {lib_name} does not implement __version__ in its init file. "
+                f"Could not check version compatibility."
+            )
+            return False, msg
+    except (AttributeError, ImportError, ModuleNotFoundError):
+        pass
+
+    msg = f"Lib {lib_name} has not been installed. Please use pip or conda to install this package."
+    return None, msg
+
+
+def uninject_model_parallel_rank(filepath):
+    filepath = str(filepath)
+    if any([s for s in ['mp_rank', 'tp_rank', 'fsdp_shard'] if s in filepath]):
+        dirname = os.path.dirname(os.path.dirname(filepath))
+        basename = os.path.basename(filepath)
+        filepath = os.path.join(dirname, basename)
+        return filepath
+    else:
+        return filepath
+
+
+def inject_model_parallel_rank(filepath, fsdp_sharded_ckpt=False):
+    """
+    Injects tensor/pipeline model parallel ranks into the filepath.
+    Does nothing if not using model parallelism.
+    """
+    # first make sure filepath does not have rank
+    filepath = uninject_model_parallel_rank(filepath)
+
+    app_state = AppState()
+    dirname = os.path.dirname(filepath)
+    basename = os.path.basename(filepath)
+    if app_state.model_parallel_size is not None and app_state.model_parallel_size > 1:
+        fsdp_shard = f'_fsdp_shard_{app_state.data_parallel_rank:05d}' if fsdp_sharded_ckpt else ''
+        if app_state.pipeline_model_parallel_size is None or app_state.pipeline_model_parallel_size == 1:
+            filepath = f'{dirname}/mp_rank_{app_state.tensor_model_parallel_rank:02d}{fsdp_shard}/{basename}'
+        else:
+            filepath = f'{dirname}/tp_rank_{app_state.tensor_model_parallel_rank:02d}_pp_rank_{app_state.pipeline_model_parallel_rank:03d}/{basename}'  # pylint: disable=line-too-long
+        return filepath
+    else:
+        fsdp_shard = f'/fsdp_shard_{app_state.data_parallel_rank:05d}' if fsdp_sharded_ckpt else ''
+        return f'{dirname}{fsdp_shard}/{basename}'
+
+
+def ckpt_to_dir(filepath: Union[str, Path]) -> Path:
+    """PTL considers checkpoints as .ckpt files.
+    This method removes the extension and returns a path
+    to be used as a directory for distributed checkpoints
+    """
+
+    filepath = Path(filepath)
+    # if it is already a distributed checkpoint, then return
+    if filepath.suffix != ".ckpt" and filepath.is_dir():
+        return filepath
+
+    # adding this assert because we will later remove directories based on the return value of this method
+    assert filepath.suffix == ".ckpt", f"filepath: {filepath} must have .ckpt extension"
+
+    # create a new path whose name is the original filepath without the .ckpt extension
+    checkpoint_dir = filepath.with_name(filepath.stem)
+
+    return checkpoint_dir
+
+
+def save_artifacts(model, output_dir: str, use_abspath: bool = False) -> None:
+    """Save all model artifacts and tokenizer config to a given output directory."""
+    app_state = AppState()
+    model_file = app_state.model_restore_path
+    model_cfg = copy.deepcopy(model.cfg)
+
+    if model_cfg.tokenizer.library == "huggingface":
+        model.tokenizer.save_pretrained(output_dir)
+
+    if not hasattr(model, "artifacts"):
+        if hasattr(model_cfg, "tokenizer"):
+            OmegaConf.save(model_cfg.tokenizer, os.path.join(output_dir, "tokenizer_config.yaml"))
+        return
+
+    # Setup model file handling context: directory or tarball
+    if os.path.isfile(model_file):
+        model_file_handler = tarfile.open
+        kwargs = {"name": model_file, "mode": "r:"}
+    elif os.path.isdir(model_file):
+        model_file_handler = contextlib.nullcontext
+        kwargs = {}
+    else:
+        raise FileNotFoundError(model_file)
+
+    # Copy or extract artifacts depending on the context
+    with model_file_handler(**kwargs) as maybe_tar:
+        if maybe_tar is not None:
+            prefix = detect_prefix(maybe_tar.getnames())
+        for arti_name, arti_item in model.artifacts.items():
+            _, arti_file = arti_item.path.split("nemo:")
+            artifact_path = _validate_artifact_path(arti_file)
+            arti_path = os.path.join(output_dir, arti_name)
+            if maybe_tar is not None:
+                member_name = f"{prefix}{artifact_path.as_posix()}"
+                safe_extract(maybe_tar, output_dir, members=[member_name])
+                os.rename(os.path.join(output_dir, *PurePosixPath(member_name).parts), arti_path)
+            else:
+                shutil.copy(os.path.join(model_file, *artifact_path.parts), arti_path)
+            # Store artifact path as basename by default. Otherwise save absolute path but bear in mind
+            # that in this case output directory should be permanent for correct artifact recovery later
+            arti_path = os.path.abspath(arti_path) if use_abspath else os.path.basename(arti_path)
+            OmegaConf.update(model_cfg, arti_name, arti_path)
+
+    if hasattr(model_cfg, "tokenizer"):
+        OmegaConf.save(model_cfg.tokenizer, os.path.join(output_dir, "tokenizer_config.yaml"))

nemo/utils/msc_utils.py

@@ -0,0 +1,65 @@
+# 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.
+
+
+from pathlib import Path
+from typing import Union
+
+try:
+    import multistorageclient as msc
+
+    HAVE_MSC = True
+except (ImportError, ModuleNotFoundError):
+    msc = None
+
+    HAVE_MSC = False
+
+
+MSC_PROTOCOL = "msc://"
+
+
+def is_multistorageclient_url(path: Union[str, Path]):
+    """
+    Check if the path is a multistorageclient URL (e.g. msc://<profile>/<path>).
+
+    Args:
+        path: str, the path to check.
+
+    Returns:
+        bool, True if the path is a multistorageclient URL, False otherwise.
+    """
+    if isinstance(path, Path):
+        return False
+
+    has_msc_prefix = path and str(path).startswith(MSC_PROTOCOL)
+
+    if HAVE_MSC:
+        return has_msc_prefix
+
+    if not HAVE_MSC and has_msc_prefix:
+        raise ValueError(
+            "Multi-Storage Client is not installed. Please install it with "
+            '"pip install multi-storage-client" to handle msc:// URLs.'
+        )
+
+    return False
+
+
+def import_multistorageclient():
+    """Import multistorageclient if it is installed."""
+    if not HAVE_MSC:
+        raise ValueError(
+            "Multi-Storage Client is not installed. Please install it with " '"pip install multi-storage-client".'
+        )
+    return msc

nemo/utils/nemo_logging.py

@@ -0,0 +1,436 @@
+# Copyright (c) 2020, 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 enum
+import logging as _logging
+import sys
+import threading
+import warnings
+from contextlib import contextmanager
+from logging.handlers import MemoryHandler
+
+from nemo.constants import NEMO_ENV_VARNAME_REDIRECT_LOGS_TO_STDERR, NEMO_ENV_VARNAME_TESTING
+from nemo.utils.env_var_parsing import get_envbool
+from nemo.utils.formatters.base import BaseNeMoFormatter, DebugNeMoFormatter
+from nemo.utils.get_rank import is_global_rank_zero
+from nemo.utils.metaclasses import Singleton
+
+__all__ = ["Logger", "LogMode"]
+
+
+class LogMode(enum.IntEnum):
+    """Enum to control how many times to log messages in NeMo logging"""
+
+    EACH = 0  # Log the message each time
+    ONCE = 1  # Log the message only once. The same message will not be logged again.
+
+
+class Logger(metaclass=Singleton):
+    """NeMo's logging class. Makes some changes on top of python's logging module to aid model devs."""
+
+    # Level 0
+    NOTSET = _logging.NOTSET
+
+    # Level 10
+    DEBUG = _logging.DEBUG
+
+    # Level 20
+    INFO = _logging.INFO
+
+    # Level 30
+    WARNING = _logging.WARNING
+
+    # Level 40
+    ERROR = _logging.ERROR
+
+    # Level 50
+    CRITICAL = _logging.CRITICAL
+
+    _level_names = {
+        0: "NOTSET",
+        10: "DEBUG",
+        20: "INFO",
+        30: "WARNING",
+        40: "ERROR",
+        50: "CRITICAL",
+    }
+
+    def __init__(self, capture_warnings=True):
+
+        self._logger = None
+        # Multi-GPU runs run in separate processes, thread locks shouldn't be needed
+        self._logger_lock = threading.Lock()
+        self._handlers = dict()
+        self.old_warnings_showwarning = None
+        self._define_logger(capture_warnings)
+        self.once_logged = set()
+        self.rank = 0 if is_global_rank_zero() else "UNK"
+
+    def _define_logger(self, capture_warnings=True):
+        """Creates the logger if not already created. Called in init"""
+
+        # Use double-checked locking to avoid taking lock unnecessarily.
+        if self._logger is not None:
+            return self._logger
+
+        with self._logger_lock:
+            try:
+                self._logger = _logging.getLogger("nemo_logger")
+                # By default, silence all loggers except the logger for rank 0
+                self.remove_stream_handlers()
+                # If NEMO_TESTING is set, add a streamhandler to all ranks
+                if get_envbool(NEMO_ENV_VARNAME_TESTING, False):
+                    old_factory = _logging.getLogRecordFactory()
+
+                    def record_factory(*args, **kwargs):
+                        record = old_factory(*args, **kwargs)
+                        record.rank = self.rank
+                        return record
+
+                    _logging.setLogRecordFactory(record_factory)
+                    self.add_stream_handlers(formatter=DebugNeMoFormatter)
+                elif is_global_rank_zero():
+                    self.add_stream_handlers()
+
+                # Add memoryhandlers, essentially buffers. They are used to save messages that we will flush to file
+                # once the appropriate file handlers are added.
+                if is_global_rank_zero():
+                    # Add a memoryhandler for error messages. Only logged on rank 0
+                    self._handlers["memory_err"] = MemoryHandler(-1)
+                    self._handlers["memory_err"].addFilter(lambda record: record.levelno > _logging.INFO)
+                    formatter = BaseNeMoFormatter
+                    self._handlers["memory_err"].setFormatter(formatter())
+                    self._logger.addHandler(self._handlers["memory_err"])
+                # Add a memoryhandler for all messages on all ranks
+                self._handlers["memory_all"] = MemoryHandler(-1)
+                formatter = BaseNeMoFormatter
+                self._handlers["memory_all"].setFormatter(formatter())
+                self._logger.addHandler(self._handlers["memory_all"])
+
+            finally:
+                level = Logger.INFO
+                if get_envbool(NEMO_ENV_VARNAME_TESTING, False):
+                    level = Logger.DEBUG
+                self.set_verbosity(verbosity_level=level)
+                self.captureWarnings(capture_warnings)
+
+        self._logger.propagate = False
+
+    def remove_stream_handlers(self):
+        """Removes StreamHandler that log to stdout and stderr from the logger."""
+        if self._logger is None:
+            raise RuntimeError("Impossible to set handlers if the Logger is not predefined")
+
+        # ======== Remove Handler if already existing ========
+
+        try:
+            self._logger.removeHandler(self._handlers["stream_stdout"])
+            del self._handlers["stream_stdout"]
+        except KeyError:
+            pass
+
+        try:
+            self._logger.removeHandler(self._handlers["stream_stderr"])
+            del self._handlers["stream_stderr"]
+        except KeyError:
+            pass
+
+    def add_stream_handlers(self, formatter=BaseNeMoFormatter):
+        """Add StreamHandler that log to stdout and stderr to the logger. INFO and lower logs are streamed to stdout
+        while WARNING and higher are streamed to stderr. If the NEMO_ENV_VARNAME_REDIRECT_LOGS_TO_STDERR environment
+        variable is set, all logs are sent to stderr instead.
+        """
+        if self._logger is None:
+            raise RuntimeError("Impossible to set handlers if the Logger is not predefined")
+
+        # Add the output handler.
+        if get_envbool(NEMO_ENV_VARNAME_REDIRECT_LOGS_TO_STDERR, False):
+            self._handlers["stream_stdout"] = _logging.StreamHandler(sys.stderr)
+
+        else:
+            self._handlers["stream_stdout"] = _logging.StreamHandler(sys.stdout)
+            self._handlers["stream_stdout"].addFilter(lambda record: record.levelno <= _logging.INFO)
+
+            self._handlers["stream_stderr"] = _logging.StreamHandler(sys.stderr)
+            self._handlers["stream_stderr"].addFilter(lambda record: record.levelno > _logging.INFO)
+
+        self._handlers["stream_stdout"].setFormatter(formatter())
+        self._logger.addHandler(self._handlers["stream_stdout"])
+
+        try:
+            self._handlers["stream_stderr"].setFormatter(formatter())
+            self._logger.addHandler(self._handlers["stream_stderr"])
+        except KeyError:
+            pass
+
+    def reset_stream_handler(self, formatter=BaseNeMoFormatter):
+        """Removes then adds stream handlers."""
+        self.remove_stream_handlers()
+        self.add_stream_handlers(formatter=formatter)
+
+    def add_file_handler(self, log_file):
+        """Add a FileHandler to logger that logs all messages to a file. If the logger had a MemoryHandler at
+        self._handlers["memory_all"], those buffered messages are flushed to the new file, and the MemoryHandler is
+        closed."""
+        if self._logger is None:
+            raise RuntimeError("Impossible to set handlers if the Logger is not predefined")
+
+        self._handlers["file"] = _logging.FileHandler(log_file)
+        formatter = BaseNeMoFormatter
+        self._handlers["file"].setFormatter(formatter())
+        self._logger.addHandler(self._handlers["file"])
+
+        if self._handlers.get("memory_all", None):
+            self._handlers["memory_all"].setTarget(self._handlers["file"])
+            self._handlers["memory_all"].close()  # flush and remove
+            del self._handlers["memory_all"]
+
+    def add_err_file_handler(self, log_file):
+        """Add a FileHandler to logger that logs all WARNING and higher messages to a file. If the logger had a
+        MemoryHandler at self._handlers["memory_err"], those buffered messages are flushed to the new file, and the
+        MemoryHandler is closed."""
+        if self._logger is None:
+            raise RuntimeError("Impossible to set handlers if the Logger is not predefined")
+
+        self._handlers["file_err"] = _logging.FileHandler(log_file)
+        self._handlers["file_err"].addFilter(lambda record: record.levelno > _logging.INFO)
+
+        formatter = BaseNeMoFormatter
+        self._handlers["file_err"].setFormatter(formatter())
+        self._logger.addHandler(self._handlers["file_err"])
+
+        if self._handlers.get("memory_err", None):
+            self._handlers["memory_err"].setTarget(self._handlers["file_err"])
+            self._handlers["memory_err"].close()  # flush and remove
+            del self._handlers["memory_err"]
+
+    def getEffectiveLevel(self):
+        """Return how much logging output will be produced."""
+        if self._logger is not None:
+            return self._logger.getEffectiveLevel()
+
+    def get_verbosity(self):
+        """See getEffectiveLevel"""
+        return self.getEffectiveLevel()
+
+    def setLevel(self, verbosity_level):
+        """Sets the threshold for what messages will be logged."""
+        if self._logger is not None:
+            self._logger.setLevel(verbosity_level)
+
+            for handler in self._logger.handlers:
+                handler.setLevel(verbosity_level)
+
+    def set_verbosity(self, verbosity_level):
+        """See setLevel"""
+        self.setLevel(verbosity_level)
+
+    @contextmanager
+    def patch_stderr_handler(self, stream):
+        """Sends messages that should log to stderr to stream instead. Useful for unittests"""
+        if self._logger is not None:
+            try:
+                old_stream = self._handlers["stream_stderr"].stream
+                if old_stream is None:
+                    raise ValueError
+
+                # Port backwards set_stream() from python 3.7
+                self._handlers["stream_stderr"].acquire()
+                try:
+                    self._handlers["stream_stderr"].flush()
+                    self._handlers["stream_stderr"].stream = stream
+                finally:
+                    self._handlers["stream_stderr"].release()
+
+                yield stream
+            except (KeyError, ValueError):
+                raise RuntimeError("Impossible to patch logging handlers if handler does not exist")
+            finally:
+                # Port backwards set_stream() from python 3.7
+                self._handlers["stream_stderr"].acquire()
+                try:
+                    self._handlers["stream_stderr"].flush()
+                    self._handlers["stream_stderr"].stream = old_stream
+                finally:
+                    self._handlers["stream_stderr"].release()
+
+        else:
+            raise RuntimeError("Impossible to patch logging handlers if handler does not exist")
+
+    @contextmanager
+    def patch_stdout_handler(self, stream):
+        """Sends messages that should log to stdout to stream instead. Useful for unittests"""
+        if self._logger is not None:
+            try:
+                old_stream = self._handlers["stream_stdout"].stream
+                if old_stream is None:
+                    raise ValueError
+
+                # Port backwards set_stream() from python 3.7
+                self._handlers["stream_stdout"].acquire()
+                try:
+                    self._handlers["stream_stdout"].flush()
+                    self._handlers["stream_stdout"].stream = stream
+                finally:
+                    self._handlers["stream_stdout"].release()
+
+                yield stream
+            except (KeyError, ValueError):
+                raise RuntimeError("Impossible to patch logging handlers if handler does not exist")
+            finally:
+                # Port backwards set_stream() from python 3.7
+                self._handlers["stream_stdout"].acquire()
+                try:
+                    self._handlers["stream_stdout"].flush()
+                    self._handlers["stream_stdout"].stream = old_stream
+                finally:
+                    self._handlers["stream_stdout"].release()
+
+        else:
+            raise RuntimeError("Impossible to patch logging handlers if handler does not exist")
+
+    @contextmanager
+    def temp_verbosity(self, verbosity_level):
+        """Sets the a temporary threshold for what messages will be logged."""
+
+        if self._logger is not None:
+
+            old_verbosity = self.get_verbosity()
+
+            try:
+                self.set_verbosity(verbosity_level)
+                yield
+
+            finally:
+                self.set_verbosity(old_verbosity)
+
+        else:
+            try:
+                yield
+
+            finally:
+                pass
+
+    def captureWarnings(self, capture):
+        """
+        If capture is true, redirect all warnings to the logging package.
+        If capture is False, ensure that warnings are not redirected to logging
+        but to their original destinations.
+        """
+
+        if self._logger is not None:
+
+            if capture and self.old_warnings_showwarning is None:
+                # Backup Method
+                self.old_warnings_showwarning = warnings.showwarning
+                warnings.showwarning = self._showwarning
+
+            elif not capture and self.old_warnings_showwarning is not None:
+                # Restore Method
+                warnings.showwarning = self.old_warnings_showwarning
+                self.old_warnings_showwarning = None
+
+    def _warning_is_ignored(self, category):
+        from warnings import filters
+
+        # Search the filters
+        for action, msg, cat, mod, ln in filters:
+            # least-common demoninator if multiple filters for the same class.
+            if cat == category and action == 'ignore':
+                return True
+        return False
+
+    def _showwarning(self, message, category, filename, lineno, file=None, line=None):
+        """
+        Implementation of showwarnings which redirects to logging.
+        It will call warnings.formatwarning and will log the resulting string
+        with level logging.WARNING.
+        """
+        s = warnings.formatwarning(message, category, filename, lineno, line)
+        if self._warning_is_ignored(category):
+            return
+        self.warning("%s", s)
+
+    def _logged_once(self, msg, mode):
+        PREFIX_LEN = 12
+        if mode == LogMode.ONCE:
+            if msg[PREFIX_LEN:] in self.once_logged:
+                return True
+            self.once_logged.add(msg[PREFIX_LEN:])
+        return False
+
+    def debug(self, msg, *args, mode=LogMode.EACH, **kwargs):
+        """
+        Log 'msg % args' with severity 'DEBUG'.
+
+        To pass exception information, use the keyword argument exc_info with
+        a true value, e.g.
+
+        logger.debug("Houston, we have a %s", "thorny problem", exc_info=1)
+        """
+        if self._logger is not None and self._logger.isEnabledFor(Logger.DEBUG) and not self._logged_once(msg, mode):
+            self._logger._log(Logger.DEBUG, msg, args, **kwargs, stacklevel=2)
+
+    def info(self, msg, *args, mode=LogMode.EACH, **kwargs):
+        """
+        Log 'msg % args' with severity 'INFO'.
+
+        To pass exception information, use the keyword argument exc_info with
+        a true value, e.g.
+
+        logger.info("Houston, we have a %s", "interesting problem", exc_info=1)
+        """
+        if self._logger is not None and self._logger.isEnabledFor(Logger.INFO) and not self._logged_once(msg, mode):
+            self._logger._log(Logger.INFO, msg, args, **kwargs, stacklevel=2)
+
+    def warning(self, msg, *args, mode=LogMode.EACH, **kwargs):
+        """
+        Log 'msg % args' with severity 'WARNING'.
+
+        To pass exception information, use the keyword argument exc_info with
+        a true value, e.g.
+
+        logger.warning("Houston, we have a %s", "bit of a problem", exc_info=1)
+        """
+        if self._logger is not None and self._logger.isEnabledFor(Logger.WARNING) and not self._logged_once(msg, mode):
+            self._logger._log(Logger.WARNING, msg, args, **kwargs, stacklevel=2)
+
+    def error(self, msg, *args, mode=LogMode.EACH, **kwargs):
+        """
+        Log 'msg % args' with severity 'ERROR'.
+
+        To pass exception information, use the keyword argument exc_info with
+        a true value, e.g.
+
+        logger.error("Houston, we have a %s", "major problem", exc_info=1)
+        """
+        if self._logger is not None and self._logger.isEnabledFor(Logger.ERROR) and not self._logged_once(msg, mode):
+            self._logger._log(Logger.ERROR, msg, args, **kwargs, stacklevel=2)
+
+    def critical(self, msg, *args, mode=LogMode.EACH, **kwargs):
+        """
+        Log 'msg % args' with severity 'CRITICAL'.
+
+        To pass exception information, use the keyword argument exc_info with
+        a true value, e.g.
+
+        logger.critical("Houston, we have a %s", "major disaster", exc_info=1)
+        """
+        if (
+            self._logger is not None
+            and self._logger.isEnabledFor(Logger.CRITICAL)
+            and not self._logged_once(msg, mode)
+        ):
+            self._logger._log(Logger.CRITICAL, msg, args, **kwargs, stacklevel=2)

nemo/utils/notebook_utils.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2022, 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 glob
+import json
+import os
+import os.path
+import subprocess
+import tarfile
+import urllib.request
+from typing import Optional
+
+from nemo.utils.dependency import import_optional_dependency
+from nemo.utils.tar_utils import safe_extract
+
+
+def build_manifest(transcripts_path, manifest_path, data_dir, mount_dir, wav_path):
+    """Build an AN4 manifest with local or mounted audio paths."""
+    # create manifest with reference to this directory. This is useful when mounting the dataset.
+    mount_dir = mount_dir if mount_dir else data_dir
+    sox = import_optional_dependency("sox")
+    with open(transcripts_path, 'r') as fin:
+        with open(manifest_path, 'w') as fout:
+            for line in fin:
+                # Lines look like this:
+                # <s> transcript </s> (fileID)
+                transcript = line[: line.find('(') - 1].lower()
+                transcript = transcript.replace('<s>', '').replace('</s>', '')
+                transcript = transcript.strip()
+
+                file_id = line[line.find('(') + 1 : -2]  # e.g. "cen4-fash-b"
+                audio_path = os.path.join(
+                    data_dir, wav_path, file_id[file_id.find('-') + 1 : file_id.rfind('-')], file_id + '.wav'
+                )
+
+                mounted_audio_path = os.path.join(
+                    mount_dir, wav_path, file_id[file_id.find('-') + 1 : file_id.rfind('-')], file_id + '.wav'
+                )
+                duration = sox.file_info.duration(audio_path)
+
+                # Write the metadata to the manifest
+                metadata = {"audio_filepath": mounted_audio_path, "duration": duration, "text": transcript}
+                json.dump(metadata, fout)
+                fout.write('\n')
+
+
+def download_an4(data_dir: str = "./", train_mount_dir: Optional[str] = None, test_mount_dir: Optional[str] = None):
+    """
+    Function to download the AN4 dataset. This hides pre-processing boilerplate for notebook ASR examples.
+
+    Args:
+        data_dir: Path to store the data.
+        train_mount_dir: If you plan to mount the dataset, use this to prepend the mount directory to the
+            audio filepath in the train manifest.
+        test_mount_dir: If you plan to mount the dataset, use this to prepend the mount directory to the
+            audio filepath in the test manifest.
+    """
+    print("******")
+    os.makedirs(data_dir, exist_ok=True)
+    if not os.path.exists(data_dir + '/an4_sphere.tar.gz'):
+        an4_url = 'https://dldata-public.s3.us-east-2.amazonaws.com/an4_sphere.tar.gz'
+        an4_path = os.path.join(data_dir, 'an4_sphere.tar.gz')
+        urllib.request.urlretrieve(an4_url, an4_path)
+        print(f"Dataset downloaded at: {an4_path}")
+    else:
+        print("Tarfile already exists.")
+        an4_path = data_dir + '/an4_sphere.tar.gz'
+
+    if not os.path.exists(data_dir + '/an4/'):
+        with tarfile.open(an4_path) as tar:
+            safe_extract(tar, data_dir)
+
+        print("Converting .sph to .wav...")
+        sph_list = glob.glob(data_dir + '/an4/**/*.sph', recursive=True)
+        for sph_path in sph_list:
+            wav_path = sph_path[:-4] + '.wav'
+            cmd = ["sox", sph_path, wav_path]
+            subprocess.run(cmd)
+    print("Finished conversion.\n******")
+
+    # Building Manifests
+    print("******")
+    train_transcripts = data_dir + '/an4/etc/an4_train.transcription'
+    train_manifest = data_dir + '/an4/train_manifest.json'
+
+    if not os.path.isfile(train_manifest):
+        build_manifest(train_transcripts, train_manifest, data_dir, train_mount_dir, 'an4/wav/an4_clstk')
+        print("Training manifest created.")
+
+    test_transcripts = data_dir + '/an4/etc/an4_test.transcription'
+    test_manifest = data_dir + '/an4/test_manifest.json'
+    if not os.path.isfile(test_manifest):
+        build_manifest(test_transcripts, test_manifest, data_dir, test_mount_dir, 'an4/wav/an4test_clstk')
+        print("Test manifest created.")
+    print("***Done***")

nemo/utils/nvtx.py

@@ -0,0 +1,63 @@
+# 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 functools
+from typing import Optional
+
+import torch
+
+from nemo.utils.app_state import AppState
+
+# pylint: disable=C0116
+
+
+@functools.lru_cache(maxsize=None)
+def _nvtx_enabled() -> bool:
+    """Check if NVTX range profiling is enabled"""
+    return AppState()._nvtx_ranges
+
+
+# Messages associated with active NVTX ranges
+_nvtx_range_messages: list[str] = []
+
+
+def nvtx_range_push(msg: str) -> None:
+    # Return immediately if NVTX range profiling is not enabled
+    if not _nvtx_enabled():
+        return
+
+    # Push NVTX range to stack
+    _nvtx_range_messages.append(msg)
+    torch.cuda.nvtx.range_push(msg)
+
+
+def nvtx_range_pop(msg: Optional[str] = None) -> None:
+    # Return immediately if NVTX range profiling is not enabled
+    if not _nvtx_enabled():
+        return
+
+    # Update list of NVTX range messages and check for consistency
+    if not _nvtx_range_messages:
+        raise RuntimeError("Attempted to pop NVTX range from empty stack")
+    last_msg = _nvtx_range_messages.pop()
+    if msg is not None and msg != last_msg:
+        raise ValueError(
+            f"Attempted to pop NVTX range from stack with msg={msg}, " f"but last range has msg={last_msg}"
+        )
+
+    # Pop NVTX range
+    torch.cuda.nvtx.range_pop()
+
+
+# pylint: enable=C0116

nemo/utils/oomptimizer.py

@@ -0,0 +1,119 @@
+#!/usr/bin/env python
+# Copyright (c) 2024, 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 math
+from dataclasses import dataclass
+from numbers import Number
+from typing import Literal
+
+from lhotse import compute_num_samples
+from omegaconf import OmegaConf
+
+from nemo.core.neural_types import LabelsType, NeuralType
+
+
+def is_2d_bucketing(buckets) -> bool:
+    """Return whether the bucket list contains input/output sequence-length pairs."""
+    return all(
+        isinstance(item, (list, tuple)) and len(item) == 2 and all(isinstance(v, Number) for v in item)
+        for item in buckets
+    )
+
+
+@dataclass
+class SequenceLengthResolver:
+    """Resolve OOMptimizer bucket values into synthetic input and output sequence lengths."""
+
+    cfg: object
+    ratio: float
+    salm_audio_token_ratio: float
+    module_name: str | None = None
+    model: object | None = None
+    schema: dict | None = None
+
+    def resolve_many(self, buckets) -> list[tuple[int, int]]:
+        """Resolve a list of OOMptimizer buckets into input and output sequence lengths."""
+        return [self.resolve_one(bucket) for bucket in buckets]
+
+    def resolve_one(self, bucket) -> tuple[int, int]:
+        """Resolve one OOMptimizer bucket into input and output sequence lengths."""
+        if self._uses_audio_locator_expansion():
+            return self._audio_locator_lens(bucket)
+
+        if is_2d_bucketing([bucket]):
+            input_len, output_len = bucket
+            return int(input_len), int(output_len)
+
+        input_len = bucket
+        output_len = int(math.ceil(self.ratio * input_len))
+        if self.schema is None:
+            return compute_num_samples(input_len, sampling_rate=16000), output_len
+
+        sampling_rate = self._sampling_rate()
+        match self._modalities():
+            case ("audio", "audio"):
+                return (
+                    compute_num_samples(input_len, sampling_rate=sampling_rate),
+                    compute_num_samples(output_len, sampling_rate=sampling_rate),
+                )
+            case ("audio", "text"):
+                return compute_num_samples(input_len, sampling_rate=sampling_rate), output_len
+            case ("text", "audio"):
+                return int(input_len), compute_num_samples(output_len, sampling_rate=sampling_rate)
+            case ("text", "text"):
+                return int(input_len), output_len
+            case unexpected:
+                raise RuntimeError(f"Unexpected modality combination: {unexpected}")
+
+    def _matches_model_name(self, *suffixes: str) -> bool:
+        return self.module_name is not None and any(self.module_name.endswith(suffix) for suffix in suffixes)
+
+    def _matches_model_class_name(self, *names: str) -> bool:
+        return self.model is not None and type(self.model).__name__ in names
+
+    def _uses_audio_locator_expansion(self) -> bool:
+        return self._matches_model_name(
+            "SALMAutomodel", "SALM", "SALMWithAsrDecoder"
+        ) or self._matches_model_class_name("SALMAutomodel", "SALM", "SALMWithAsrDecoder")
+
+    def _modalities(self) -> tuple[str, str]:
+        if self.schema is None:
+            return "audio", "text"
+
+        def _modality(direction: Literal["input", "output"]) -> str:
+            for item in self.schema["inputs"]:
+                nt = item["type"]
+                if nt == "dummy":
+                    continue
+                if (
+                    isinstance(nt, NeuralType)
+                    and isinstance(nt.elements_type, LabelsType)
+                    and item["seq_length"] == direction
+                ):
+                    return "text"
+            return "audio"
+
+        return _modality("input"), _modality("output")
+
+    def _sampling_rate(self) -> int:
+        return int(getattr(self.model, "sample_rate", 16000))
+
+    def _audio_locator_lens(self, bucket) -> tuple[int, int]:
+        sampling_rate = OmegaConf.select(self.cfg, "data.train_ds.sample_rate", default=16000)
+        token_equivalent_duration = OmegaConf.select(self.cfg, "data.train_ds.token_equivalent_duration", default=0.08)
+        audio_tokens = max(1, int(math.ceil(self.salm_audio_token_ratio * bucket)))
+        text_tokens = max(2, int(math.ceil((1.0 - self.salm_audio_token_ratio) * bucket)))
+        audio_len = int(math.ceil(audio_tokens * token_equivalent_duration * sampling_rate))
+        return audio_len, text_tokens

nemo/utils/s3_dirpath_utils.py

@@ -0,0 +1,36 @@
+# 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.
+
+from pathlib import Path
+from typing import Optional
+
+S3_PATH_PREFIX = 's3://'
+
+
+def build_s3_url(bucket, key) -> str:
+    """
+    This function constructs an s3 address given a bucket and key.
+    It has no reliance on any S3-related dependencies as the file pre-defines the S3 path prefix.
+    """
+    return f'{S3_PATH_PREFIX}{bucket}/{key}'
+
+
+def is_s3_url(path: Optional[str]) -> bool:
+    """
+    This function checks if a path is an S3 url.
+    It has no reliance on any S3-related dependencies as the file pre-defines the S3 path prefix.
+    """
+    if isinstance(path, Path):
+        path = str(path)
+    return path is not None and path.strip().startswith(S3_PATH_PREFIX)

nemo/utils/s3_utils.py

@@ -0,0 +1,356 @@
+# 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 re
+import time
+from io import BytesIO
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+import boto3
+import botocore
+from boto3.s3.transfer import TransferConfig
+from botocore.exceptions import ClientError
+from tenacity import before_sleep_log, retry, retry_if_exception, stop_after_delay, wait_exponential
+
+from nemo.utils import logging
+from nemo.utils.s3_dirpath_utils import build_s3_url, is_s3_url
+
+try:
+    import awscrt
+    import s3transfer.crt
+
+    crt_available = True
+except ImportError as e:
+    crt_available = False
+
+MB = 1024**2
+GB = 1024**3
+
+SHARED_MEM_DIR = '/dev/shm'
+DEFAULT_CHUNK_SIZE_MB = 64
+DEFAULT_MAX_READ_CONCURRENCY = 15
+DEFAULT_MAX_WRITE_CONCURRENCY = 10
+
+
+class S3Utils:
+    """
+    Utility class for interacting with S3. Handles downloading and uploading to S3, and parsing/formatting S3 urls.
+    """
+
+    '''
+    Avoid caching boto3 client or resource as a class variable as it gets executed once during class construction.
+    When the security token expires, the client or resouece will be no longer valid.
+    Create a new resource as needed. To avoid multithreading errors, use different session for each thread.
+    '''
+
+    @staticmethod
+    def s3_path_exists(s3_path: str, match_directory: bool = False) -> bool:
+        """
+        :s3_path: the path
+        :match_directory: if the content is known to be a directory then set it to `True`. Since s3 isn't a file system, paths are funky and the concept of folders doesn't really exist.
+        """
+        bucket_name, prefix = S3Utils.parse_s3_url(s3_path)
+        if not prefix:
+            return False
+
+        s3 = S3Utils._get_s3_resource()
+        # bucket = s3.Bucket(bucket_name)
+        s3_client = s3.meta.client
+
+        try:
+            objs = s3_client.list_objects_v2(Bucket=bucket_name, MaxKeys=1, Prefix=prefix).get('Contents', [])
+        except s3_client.exceptions.NoSuchBucket:
+            return False
+
+        if prefix == '':  # bucket only
+            return True
+
+        return len(objs) > 0 and (match_directory or objs[0]['Key'].startswith(prefix))
+
+    @staticmethod
+    def remove_object(s3_path: str) -> None:
+        s3_client = S3Utils._get_s3_resource(get_client=True)
+        bucket, key = S3Utils.parse_s3_url(s3_path)
+        s3_client.delete_object(Bucket=bucket, Key=key)
+
+    @staticmethod
+    def download_s3_file_to_stream(
+        s3_path: str, chunk_size_MB: int = DEFAULT_CHUNK_SIZE_MB, max_concurrency: int = DEFAULT_MAX_READ_CONCURRENCY
+    ) -> BytesIO:
+        bytes_buffer = BytesIO()
+
+        s3_client = S3Utils._get_s3_resource(get_client=True)
+        bucket, key = S3Utils.parse_s3_url(s3_path)
+        chunk_size = chunk_size_MB * MB
+        config = TransferConfig(multipart_chunksize=chunk_size, max_concurrency=max_concurrency)
+
+        start_time = time.perf_counter()
+        _download_fileobj_with_retry(s3_client, bucket, key, bytes_buffer, config)
+        logging.info(
+            f'Time elapsed downloading {s3_path} to file stream with chunk_size={chunk_size_MB}MB '
+            f'and max_concurrency={max_concurrency}: {(time.perf_counter() - start_time):.2f} seconds'
+        )
+
+        bytes_buffer.seek(0)
+        return bytes_buffer
+
+    @staticmethod
+    def download_s3_file_to_path(
+        s3_path: str,
+        file_path: str,
+        chunk_size_MB: int = DEFAULT_CHUNK_SIZE_MB,
+        max_concurrency: int = DEFAULT_MAX_READ_CONCURRENCY,
+    ) -> None:
+        s3_client = S3Utils._get_s3_resource(get_client=True)
+        bucket, key = S3Utils.parse_s3_url(s3_path)
+        chunk_size = chunk_size_MB * MB
+        config = TransferConfig(multipart_chunksize=chunk_size, max_concurrency=max_concurrency)
+
+        logging.info(
+            f'Downloading {s3_path} to {file_path} with chunk_size={chunk_size_MB}MB and max_threads={max_concurrency}'
+        )
+        start_time = time.perf_counter()
+        _download_file_with_retry(s3_client, bucket, key, file_path, config)
+        logging.info(
+            f'Time elapsed downloading {s3_path} to {file_path} with chunk_size={chunk_size_MB}MB '
+            f'and max_concurrency={max_concurrency}: {(time.perf_counter() - start_time):.2f} seconds'
+        )
+
+    @staticmethod
+    def upload_file_stream_to_s3(
+        bytes_buffer: BytesIO,
+        s3_path: str,
+        chunk_size_MB: int = DEFAULT_CHUNK_SIZE_MB,
+        max_concurrency: int = DEFAULT_MAX_WRITE_CONCURRENCY,
+    ) -> None:
+        s3_client = S3Utils._get_s3_resource(get_client=True)
+        bucket, key = S3Utils.parse_s3_url(s3_path)
+        chunk_size = chunk_size_MB * MB
+        config = TransferConfig(multipart_chunksize=chunk_size, max_concurrency=max_concurrency)
+        bytes_buffer.seek(0)
+
+        start_time = time.perf_counter()
+        _upload_fileobj_with_retry(s3_client, bytes_buffer, bucket, key, config)
+        logging.info(
+            f'Time elapsed uploading bytes buffer to {s3_path} with chunk_size={chunk_size_MB}MB '
+            f'and max_concurrency={max_concurrency}: {(time.perf_counter() - start_time):.2f} seconds'
+        )
+
+    @staticmethod
+    def upload_file(
+        file_path: str,
+        s3_path: str,
+        chunk_size_MB=DEFAULT_CHUNK_SIZE_MB,
+        max_concurrency=DEFAULT_MAX_WRITE_CONCURRENCY,
+        remove_file=False,
+    ):
+        total_size = os.path.getsize(file_path)
+        assert total_size > 0, f"file size is zero, {file_path}"
+
+        s3_client = S3Utils._get_s3_resource(get_client=True)
+        bucket, key = S3Utils.parse_s3_url(s3_path)
+
+        chunk_size = chunk_size_MB * MB
+        config = TransferConfig(
+            multipart_threshold=chunk_size, multipart_chunksize=chunk_size, max_concurrency=max_concurrency
+        )
+
+        start_time = time.perf_counter()
+        _upload_file_with_retry(s3_client, file_path, bucket, key, config)
+        if remove_file and os.path.exists(file_path):
+            os.remove(file_path)
+        logging.info(
+            f'Time elapsed uploading file {file_path} of size {(total_size/GB):.1f}GB to {s3_path} with chunk_size={chunk_size_MB}MB '
+            f'and max_concurrency={max_concurrency}: {(time.perf_counter() - start_time):.2f} seconds'
+        )
+
+    @staticmethod
+    def find_files_with_suffix(
+        base_path: str,
+        suffix: str = None,
+        return_key_only: bool = True,
+        profile: Optional[str] = None,
+        creds: botocore.credentials.Credentials = None,
+    ) -> List[str]:
+        """
+        Returns a list of keys that have the specified suffix
+        :param base_path: the root of search
+        :param suffix: the suffix to match, case sensitive
+        :return: list of keys matching the suffix, relative to the base_path
+        """
+        s3 = S3Utils._get_s3_resource(profile, creds)
+        bucket_name, prefix = S3Utils.parse_s3_url(base_path)
+
+        start_time = time.perf_counter()
+        bucket = s3.Bucket(bucket_name)
+        objects_list = _scan_objects_with_retry(s3_bucket=bucket, s3_prefix=prefix)
+        logging.info(
+            f'Time elapsed reading all objects under path {base_path}: {(time.perf_counter() - start_time):.2f} seconds'
+        )
+
+        if suffix:
+            objects_list = list(filter(lambda o: o.key.endswith(suffix), objects_list))
+
+        if return_key_only:
+            return [o.key for o in objects_list]
+        else:
+            return [S3Utils.build_s3_url(o.bucket_name, o.key) for o in objects_list]
+
+    @staticmethod
+    def _get_s3_resource(
+        profile: str = None,
+        creds: botocore.credentials.Credentials = None,
+        get_client: bool = False,
+        session=None,
+        config={},
+    ):
+        config = botocore.config.Config(max_pool_connections=30, **config)
+
+        if profile is not None and creds is not None:
+            raise ValueError('Please provide profile or creds or neither, not both.')
+
+        if profile is not None:
+            s3 = boto3.Session(profile_name=profile).resource('s3', config=config)
+        elif creds is not None:
+            s3 = boto3.Session().resource(
+                's3',
+                aws_access_key_id=creds["AccessKeyId"],
+                aws_secret_access_key=creds["SecretAccessKey"],
+                aws_session_token=creds["SessionToken"],
+                config=config,
+            )
+        else:
+            s3 = (
+                boto3.Session().resource('s3', config=config) if not session else session.resource('s3', config=config)
+            )
+
+        if get_client:
+            return s3.meta.client
+        else:
+            return s3
+
+    @staticmethod
+    def parse_s3_url(s3_url: str) -> Optional[Tuple[str, str]]:
+        match = re.match(r"s3://([^/]+)/(.*)", s3_url, flags=re.UNICODE)
+
+        if match is None:
+            return None, None
+
+        return match.groups()[0], match.groups()[1]
+
+    @staticmethod
+    def build_s3_url(bucket, key) -> str:
+        return build_s3_url(bucket, key)
+
+    @staticmethod
+    def is_s3_url(path: Optional[str]) -> bool:
+        return is_s3_url(path)
+
+    @staticmethod
+    def parse_prefix_with_step(path: str) -> str:
+        """
+        Use regex to find the pattern up to "-step=900-"
+        s3://path/to/checkpoints/tp_rank_00_pp_rank_000/megatron_gpt--step=900-validation_loss=6.47-consumed_samples=35960.0-last.ckpt
+        should return s3://path/to/checkpoints/tp_rank_00_pp_rank_000/megatron_gpt--step=900-
+        """
+        match = re.search(r'(.*step=\d+-)', path)
+
+        if match:
+            return match.group(1)
+
+        return path
+
+
+def _scan_objects_with_retry(s3_bucket, s3_prefix):
+    # this returns a collection https://boto3.amazonaws.com/v1/documentation/api/latest/guide/collections.html
+    # This collection acts as an iterable that automatically makes additional requests to retrieve more objects from S3 as needed
+    objects = s3_bucket.objects.filter(Prefix=s3_prefix)
+    return list(objects)
+
+
+def is_slow_down_error(exception):
+    """
+    This function checks if the error is due to slowdown or is throttling related.
+    If so, returns true to allow tenacity to retry the upload/download to S3.
+    """
+    class_name = exception.__class__.__name__
+    module_name = exception.__class__.__module__
+    full_class_name = f"{module_name}.{class_name}"
+    logging.error(f'Caught exception of type {full_class_name}: {exception}')
+
+    # 2023-12-07T05:59:25.913721576Z stdout F 2023-12-07 05:59:25,913 [ERROR] - s3_utils.py:354 - Caught exception:
+    # AWS_ERROR_S3_INVALID_RESPONSE_STATUS: Invalid response status from request. Body from error request is: b'<?xml version="1.0" encoding="UTF-8"?>\n<Error><Code>RequestTimeout</Code><Message>Your socket connection to the server was not read from or written to within the timeout period. Idle connections will be closed.</Message><RequestId>XPHS9896G3RJE364</RequestId><HostId>ZAiF3HPpUD5IgSr/mfkP2QPs7ttuvY+uTRG9MET/jZZ45MJ6bVbnvSBQLggICvPCROPP/1k85p4=</HostId></Error>'
+    message = str(exception)
+    if (
+        "<Code>SlowDown</Code>" in message
+        or "<Code>RequestTimeout</Code>" in message
+        or "<Code>InternalError</Code>" in message
+    ):
+        logging.info("Identified the Retriable Error retrying the job")
+        return True
+
+    if crt_available and isinstance(exception, awscrt.exceptions.AwsCrtError):
+        logging.error(f'Caught awscrt.exceptions.AwsCrtError: {exception.__repr__()}')
+        return True
+
+    if isinstance(exception, ClientError):
+        logging.error(f'Caught ClientError, response is: {exception.response}')
+        error_code = exception.response['Error']['Code'] if exception.response else None
+        return error_code in ['SlowDown', 'RequestTimeout', 'InternalError']
+    logging.info("Non Retriable Error - Terminating the job")
+    return False
+
+
+@retry(
+    wait=wait_exponential(multiplier=1, min=1, max=16),
+    stop=stop_after_delay(2 * 60),
+    retry=retry_if_exception(is_slow_down_error),
+    before_sleep=before_sleep_log(logging, logging.ERROR),
+)
+def _download_fileobj_with_retry(
+    s3_client, bucket: str, key: str, bytes_buffer: BytesIO, config: TransferConfig = None
+):
+    s3_client.download_fileobj(bucket, key, bytes_buffer, Config=config)
+
+
+@retry(
+    wait=wait_exponential(multiplier=1, min=1, max=16),
+    stop=stop_after_delay(2 * 60),
+    retry=retry_if_exception(is_slow_down_error),
+    before_sleep=before_sleep_log(logging, logging.ERROR),
+)
+def _download_file_with_retry(s3_client, bucket: str, key: str, file_path: str, config: TransferConfig = None):
+    s3_client.download_file(bucket, key, file_path, Config=config)
+
+
+@retry(
+    wait=wait_exponential(multiplier=1, min=1, max=16),
+    stop=stop_after_delay(2 * 60),
+    retry=retry_if_exception(is_slow_down_error),
+    before_sleep=before_sleep_log(logging, logging.ERROR),
+)
+def _upload_fileobj_with_retry(s3_client, bytes_buffer: BytesIO, bucket: str, key: str, config: TransferConfig = None):
+    s3_client.upload_fileobj(bytes_buffer, bucket, key, Config=config)
+
+
+@retry(
+    wait=wait_exponential(multiplier=1, min=1, max=16),
+    stop=stop_after_delay(2 * 60),
+    retry=retry_if_exception(is_slow_down_error),
+    before_sleep=before_sleep_log(logging, logging.ERROR),
+)
+def _upload_file_with_retry(s3_client, file_path: str, bucket: str, key: str, config: TransferConfig = None):
+    s3_client.upload_file(file_path, bucket, key, Config=config)

nemo/utils/sequence_packing_utils.py

@@ -0,0 +1,276 @@
+# 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 collections
+from typing import Dict, List, Tuple
+
+import numpy as np
+from tqdm import tqdm
+
+from nemo.utils import logging
+
+PACKING_ALGOS = ["first_fit_decreasing", "first_fit_shuffle"]
+
+
+def find_first_bin_that_fits(bins: List[List[int]], s: int, bin_size: int) -> int:
+    """
+    Finds the first bin in a list of bins that has enough space to fit a sequence of size 's'.
+
+    Args:
+      bins: A list of lists, where each inner list represents a bin and contains the current elements in that bin.
+      s: The size of the sequence to be placed in a bin.
+      bin_size: The maximum capacity of each bin.
+
+    Returns:
+      The index of the first bin that can fit the sequence 's', or -1 if no such bin exists.
+    """
+    for i, abin in enumerate(bins):
+        if sum(abin) + s <= bin_size:
+            return i
+    return -1
+
+
+def first_fit(seqlens: List[int], pack_size: int) -> List[List[int]]:
+    """
+    Packs sequences of varying lengths into bins using the First-Fit algorithm.
+
+    Args:
+      seqlens: A list of integers, representing the lengths of the sequences to be packed.
+      pack_size: The maximum capacity of each bin.
+
+    Returns:
+      A list of lists, where each inner list represents a bin and contains the indices
+        of the sequences assigned to that bin.
+    """
+    res = []
+    for s in seqlens:
+        first_bin = find_first_bin_that_fits(res, s, pack_size)
+        if first_bin == -1:  # open a new bin
+            res.append([s])
+        else:
+            res[first_bin].append(s)
+    return res
+
+
+def first_fit_decreasing(seqlens: List[int], pack_size: int) -> List[List[int]]:
+    """
+    Packs sequences of varying lengths into bins using the First-Fit Decreasing algorithm.
+
+    This is a variation of the First-Fit algorithm where the sequences are sorted by decreasing length before packing.
+
+    Args:
+      seqlens: A list of integers, representing the lengths of the sequences to be packed.
+      pack_size: The maximum capacity of each bin.
+
+    Returns:
+      A list of lists, similar to the output of the 'first_fit' function.
+    """
+    sorted_seqlens = sorted(seqlens, reverse=True)
+    return first_fit(sorted_seqlens, pack_size)
+
+
+def first_fit_shuffle(seqlens: List[int], pack_size: int) -> List[List[int]]:
+    """
+    Packs sequences of varying lengths into bins using the First-Fit with Shuffling algorithm.
+
+    This variation shuffles the order of the sequences before applying the First-Fit algorithm.
+
+    Args:
+      seqlens: A list of integers, representing the lengths of the sequences to be packed.
+      pack_size: The maximum capacity of each bin.
+
+    Returns:
+      A list of lists, similar to the output of the 'first_fit' function.
+    """
+    shuffled_seqlens = seqlens[:]
+    np.random.shuffle(shuffled_seqlens)
+    return first_fit(shuffled_seqlens, pack_size)
+
+
+def create_hist(dataset: np.array, truncate_seq_len: int):
+    """
+    Creates a histogram of sequence lengths from a tokenized dataset.
+
+    This function analyzes the tokenized dataset and creates a histogram showing the distribution of sequence lengths.
+
+    Args:
+      dataset: A NumPy array containing the tokenized sequences. Each element is a dictionary that contains at minimum
+               the key `input_ids`.
+      truncate_seq_len: The maximum sequence length to consider in the histogram.
+
+    Returns:
+      sequences: A dictionary where keys are sequence lengths and values are lists
+                 of corresponding sequences from the dataset.
+      histogram: A list representing the histogram data (number of sequences for each length).
+    """
+    logging.info("Creating histogram from tokenized dataset...")
+
+    sequences = collections.defaultdict(list)
+    counts = [0] * (truncate_seq_len + 1)
+
+    for item_dict in dataset:
+        # Minus 1 here to account for the fact that transformer input and label
+        # have one less token than the full sequence.
+        # Input is missing the last token and label is missing the first token
+        # (this way the tokens are aligned for next token prediction).
+        # We want pack size to be the length of the actual input and label, hence minus 1.
+        seq_len = len(item_dict["input_ids"]) - 1
+        sequences[seq_len].append(item_dict)
+        counts[seq_len] += 1
+
+    logging.debug("Histogram of sequence lengths")
+    logging.debug(counts)
+
+    histogram = []
+    for seq_len in range(truncate_seq_len + 1):
+        histogram.append(len(sequences[seq_len]))
+
+    return sequences, histogram
+
+
+def create_packing_strategy(
+    histogram: List[int], pack_size: int, packing_algorithm: str = "first_fit"
+) -> Tuple[List[List[int]], dict]:
+    """
+    Packs sequences into bins using the specified packing algorithm.
+
+    This function takes the histogram of sequence lengths, desired pack size, and a string representing the packing
+    algorithm to use. It then calls the corresponding function (e.g., 'first_fit_decreasing') and performs the
+    packing process using only sequence lengths as input (without the actual sequences).
+
+    Args:
+          histogram: A list representing the histogram data (number of sequences for each length).
+          pack_size: The maximum capacity of each bin.
+          packing_algorithm: One of the supported packing algorithms from ['first_fit_decreasing', 'first_fit_shuffle']
+
+    Returns:
+          assignments: A list of lists, where each inner list represents a bin and contains the indices of the
+                        sequence lengths assigned to that bin.
+          pack_metadata: A dict that records packing metadata, for instance the max number of samples per bin.
+    """
+
+    logging.info(f"Packing sequences to length {pack_size}...")
+
+    all_seq_lens = []
+    for i, count in enumerate(histogram):
+        all_seq_lens.extend([i] * count)
+
+    packing_fn = globals()[packing_algorithm]
+    assignments: List[List[int]] = packing_fn(all_seq_lens, pack_size)
+    packed_seq_lens = [sum(x) for x in assignments]
+    packing_factor = len(all_seq_lens) / len(packed_seq_lens)
+
+    max_seqlen = max(all_seq_lens)
+    max_samples_per_bin = max([len(b) for b in assignments])
+    min_packed_seqlen = min(packed_seq_lens)
+    packing_metadata = {
+        "dataset_max_seqlen": max_seqlen,
+        "max_samples_per_bin": max_samples_per_bin,
+        "packing_factor": round(packing_factor, 2),
+        "packing_efficiency": round(sum(packed_seq_lens) / len(packed_seq_lens) / pack_size * 100, 2),
+        "pack_size": pack_size,
+        'min_packed_seqlen': min_packed_seqlen,
+    }
+    logging.debug("Packed sequence lengths:")
+    logging.debug(packed_seq_lens)
+    logging.info(f"Packing is {sum(packed_seq_lens) / len(packed_seq_lens) / pack_size * 100:.2f}% efficient")
+    logging.info(
+        f">>>>> For pack size {pack_size}, average number of sequences per pack is n = {packing_factor:.3f} <<<<<"
+    )
+    return assignments, packing_metadata
+
+
+def fill_packing_strategy(
+    assignments: List[List[int]],
+    sequences: Dict[int, List[Dict]],
+    pack_size: int,
+    pad_id: int,
+) -> List[Dict]:
+    """
+    Fills the packing strategy with actual sequence data based on assignments and sequence information.
+
+    This function takes the assignments generated by the packing algorithm (containing sequence length indices),
+    the original sequences data, and the pack size. It iterates through the assignments, retrieves the corresponding
+    sequences from the sequences dictionary, and constructs the final output data structure with input IDs, loss masks
+    (if available), and starting indices for each sequence in a packed sequence.
+
+    Args:
+          assignments: A list of lists, where each inner list represents a bin and contains the indices of the
+                        sequence lengths assigned to that bin (output of 'create_packing_strategy').
+          sequences: A dictionary where keys are sequence lengths and values are lists of corresponding sequences
+                      from the dataset (output of 'create_hist').
+          pack_size: The maximum capacity of each bin.
+          pad_id: The tokenizer's padding token.
+
+    Returns:
+          output_data: A list of dictionaries, where each dictionary represents a packed sequence with its input IDs,
+                        loss mask (if available), and starting indices.
+    """
+    ifile_handles = dict()
+    for seq_len in tqdm(range(pack_size + 1)):
+        per_seq_data = sequences[seq_len]
+        if len(per_seq_data) > 0:
+            perm = np.random.permutation(len(per_seq_data))
+            input_ids = np.array([x["input_ids"] for x in per_seq_data])[perm].tolist()
+            try:
+                loss_mask = np.array([x["loss_mask"] for x in per_seq_data])[perm].tolist()
+                # roll loss mask by 1 to align with labels. We want to train on the output after the last context token
+                loss_mask = [x[1:] + [False] for x in loss_mask]
+            except KeyError:
+                try:
+                    loss_mask = np.array(
+                        [
+                            [
+                                # (x['answer_start_idx'] - 1) because we want to train on the output
+                                # after the last context token
+                                idx >= (x["answer_start_idx"] - 1)
+                                for idx in range(len(x["input_ids"]))
+                            ]
+                            for x in per_seq_data
+                        ]
+                    )[perm].tolist()
+                except KeyError as err:
+                    err_msg = "Key errors loss_mask and answer_start_idx missing in example - "
+                    err_msg += f"{err} {per_seq_data[0]}"
+                    logging.error(err_msg)
+                    raise ValueError(err_msg)
+
+            ifile_handles[seq_len] = (input_ids, loss_mask)
+
+    input_ids, loss_mask, seq_start_id = {}, {}, {}
+
+    for oindex, assignment in tqdm(enumerate(assignments), total=len(assignments)):
+        _input_ids, _loss_mask, _seq_start_id = [], [], [0]
+
+        for seq_length in assignment:
+            _input_ids.extend(ifile_handles[seq_length][0].pop())
+            _loss_mask.extend(ifile_handles[seq_length][1].pop())
+            _seq_start_id.append(len(_input_ids))
+
+        input_ids[oindex] = _input_ids
+        loss_mask[oindex] = _loss_mask
+        seq_start_id[oindex] = _seq_start_id[:-1]
+
+    output_data = []
+    for i in range(len(input_ids)):
+        item_dict = {
+            "input_ids": input_ids[i],
+            "loss_mask": loss_mask[i],
+            "seq_start_id": seq_start_id[i],
+        }
+        output_data.append(item_dict)
+
+    assert all(not seq[0] for seq in ifile_handles.values()), "Error: There are items left over from the assignment"
+    assert all(not seq[1] for seq in ifile_handles.values()), "Error: There are items left over from the assignment"
+    return output_data

nemo/utils/tar_utils.py

@@ -0,0 +1,64 @@
+# Copyright (c) 2026, 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 logging
+import os
+import tarfile
+from pathlib import PurePosixPath
+from typing import Iterable, Optional, Union
+
+
+class TarPathTraversalError(ValueError):
+    """Raised when a tar member would extract outside the target directory."""
+
+
+def is_safe_tar_member(member: tarfile.TarInfo, extract_to: str) -> bool:
+    member_path = PurePosixPath(member.name)
+    if member_path.is_absolute() or ".." in member_path.parts:
+        return False
+
+    if member.issym() or member.islnk():
+        return False
+
+    destination = os.path.realpath(os.path.join(extract_to, *member_path.parts))
+    return os.path.commonpath([destination, extract_to]) == extract_to
+
+
+def safe_extract(
+    tar: tarfile.TarFile,
+    path: str,
+    members: Optional[Iterable[Union[tarfile.TarInfo, str]]] = None,
+    *,
+    skip_unsafe: bool = False,
+) -> list[tarfile.TarInfo]:
+    extract_to = os.path.realpath(path)
+    os.makedirs(extract_to, exist_ok=True)
+    if members is None:
+        members = tar.getmembers()
+
+    extracted_members = []
+    for member in members:
+        member = tar.getmember(member) if isinstance(member, str) else member
+        if is_safe_tar_member(member, extract_to):
+            tar.extract(member, extract_to, filter="data")
+            extracted_members.append(member)
+            continue
+
+        message = f"Skipping potentially unsafe tar member: {member.name}"
+        if skip_unsafe:
+            logging.warning(message)
+            continue
+        raise TarPathTraversalError(message)
+
+    return extracted_members

nemo/utils/te_utils.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2020, 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 functools
+import importlib.metadata
+from typing import Tuple
+
+import packaging
+import torch
+from nemo.utils.import_utils import safe_import_from
+
+# Check if Transformer Engine has quantized tensor classes
+Float8Tensor, HAVE_TE_FLOAT8TENSOR = safe_import_from("transformer_engine.pytorch.float8_tensor", "Float8Tensor")
+MXFP8Tensor, HAVE_TE_MXFP8TENSOR = safe_import_from("transformer_engine.pytorch.mxfp8_tensor", "MXFP8Tensor")
+
+
+def is_float8tensor(tensor: torch.Tensor) -> bool:
+    """Check if a tensor is a Transformer Engine Float8Tensor"""
+    return HAVE_TE_FLOAT8TENSOR and isinstance(tensor, Float8Tensor)
+
+
+def is_mxfp8tensor(tensor: torch.Tensor) -> bool:
+    """Check if a tensor is a Transformer Engine MXFP8Tensor"""
+    return HAVE_TE_MXFP8TENSOR and isinstance(tensor, MXFP8)
+
+
+@functools.lru_cache(maxsize=None)
+def te_version() -> Tuple[int, ...]:
+    """Transformer Engine version"""
+    return packaging.version.Version(importlib.metadata.version("transformer-engine")).release

nemo/utils/timers.py

@@ -0,0 +1,217 @@
+"""
+This module support timing of code blocks.
+"""
+
+# Copyright (c) 2021, 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 time
+from typing import Optional
+
+import numpy as np
+import torch
+
+__all__ = ["NamedTimer", "SimpleTimer"]
+
+
+class NamedTimer(object):
+    """
+    A timer class that supports multiple named timers.
+    A named timer can be used multiple times, in which case the average
+    dt will be returned.
+    A named timer cannot be started if it is already currently running.
+    Use case: measuring execution of multiple code blocks.
+    """
+
+    _REDUCTION_TYPE = ["mean", "sum", "min", "max", "none"]
+
+    def __init__(self, reduction="mean", sync_cuda=False, buffer_size=-1):
+        """
+        Args:
+            reduction (str): reduction over multiple timings of the same timer
+                             (none - returns the list instead of a scalar)
+            sync_cuda (bool): if True torch.cuda.synchronize() is called for start/stop
+            buffer_size (int): if positive, limits the number of stored measures per name
+        """
+        if reduction not in self._REDUCTION_TYPE:
+            raise ValueError(f"Unknown reduction={reduction} please use one of {self._REDUCTION_TYPE}")
+
+        self._reduction = reduction
+        self._sync_cuda = sync_cuda
+        self._buffer_size = buffer_size
+
+        self.reset()
+
+    def __getitem__(self, k):
+        return self.get(k)
+
+    @property
+    def buffer_size(self):
+        """Return the number of recent timings retained per timer."""
+        return self._buffer_size
+
+    @property
+    def _reduction_fn(self):
+        if self._reduction == "none":
+            fn = lambda x: x
+        else:
+            fn = getattr(np, self._reduction)
+
+        return fn
+
+    def reset(self, name=None):
+        """
+        Resents all / specific timer
+
+        Args:
+            name (str): timer name to reset (if None all timers are reset)
+        """
+        if name is None:
+            self.timers = {}
+        else:
+            self.timers[name] = {}
+
+    def start(self, name=""):
+        """
+        Starts measuring a named timer.
+
+        Args:
+            name (str): timer name to start
+        """
+        timer_data = self.timers.get(name, {})
+
+        if "start" in timer_data:
+            raise RuntimeError(f"Cannot start timer = '{name}' since it is already active")
+
+        # synchronize pytorch cuda execution if supported
+        if self._sync_cuda and torch.cuda.is_initialized():
+            torch.cuda.synchronize()
+
+        timer_data["start"] = time.time()
+
+        self.timers[name] = timer_data
+
+    def stop(self, name=""):
+        """
+        Stops measuring a named timer.
+
+        Args:
+            name (str): timer name to stop
+        """
+        timer_data = self.timers.get(name, None)
+        if (timer_data is None) or ("start" not in timer_data):
+            raise RuntimeError(f"Cannot end timer = '{name}' since it is not active")
+
+        # synchronize pytorch cuda execution if supported
+        if self._sync_cuda and torch.cuda.is_initialized():
+            torch.cuda.synchronize()
+
+        # compute dt and make timer inactive
+        dt = time.time() - timer_data.pop("start")
+
+        # store dt
+        timer_data["dt"] = timer_data.get("dt", []) + [dt]
+
+        # enforce buffer_size if positive
+        if self._buffer_size > 0:
+            timer_data["dt"] = timer_data["dt"][-self._buffer_size :]
+
+        self.timers[name] = timer_data
+
+    def is_active(self, name=""):
+        """Return whether the named timer is currently active."""
+        timer_data = self.timers.get(name, {})
+        if "start" in timer_data:
+            return True
+        return False
+
+    def active_timers(self):
+        """
+        Return list of all active named timers
+        """
+        return [k for k, v in self.timers.items() if ("start" in v)]
+
+    def get(self, name=""):
+        """
+        Returns the value of a named timer
+
+        Args:
+            name (str): timer name to return
+        """
+        dt_list = self.timers[name].get("dt", [])
+
+        return self._reduction_fn(dt_list)
+
+    def export(self):
+        """
+        Exports a dictionary with average/all dt per named timer
+        """
+        fn = self._reduction_fn
+
+        data = {k: fn(v["dt"]) for k, v in self.timers.items() if ("dt" in v)}
+
+        return data
+
+
+class SimpleTimer:
+    """
+    Simple Timer with maximum possible resolution, uses `time.perf_counter_ns`.
+    """
+
+    def __init__(self, sync_cuda=True):
+        """
+
+        Args:
+            sync_cuda: synchronize CUDA device.
+                The synchronization is done only if the device for start/stop is None or CUDA device.
+        """
+        self.total_time = 0
+        self._start_time: Optional[int] = None
+        self.sync_cuda = sync_cuda
+
+    def reset(self):
+        """Reset timer"""
+        self.total_time = 0
+        self._start_time = None
+
+    def start(self, device: Optional[torch.device] = None):
+        """
+        Start timer.
+
+        Args:
+            device: CUDA device to synchronize (optional).
+        """
+        if self.sync_cuda and torch.cuda.is_initialized() and (device is None or device.type == "cuda"):
+            torch.cuda.synchronize(device=device)
+        if self._start_time is not None:
+            raise RuntimeError("Timer already started")
+        self._start_time = time.perf_counter_ns()
+
+    def stop(self, device: Optional[torch.device] = None):
+        """
+        Stop device.
+
+        Args:
+            device: CUDA device to synchronize (optional).
+        """
+        if self.sync_cuda and torch.cuda.is_initialized() and (device is None or device.type == "cuda"):
+            torch.cuda.synchronize(device=device)
+        if self._start_time is None:
+            raise RuntimeError("Timer not started")
+        self.total_time += time.perf_counter_ns() - self._start_time
+        self._start_time = None
+
+    def total_sec(self) -> float:
+        """Return total time in seconds"""
+        return self.total_time / 1e9

nemo/utils/trainer_utils.py

@@ -0,0 +1,251 @@
+# 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.
+# pylint: disable=C0116
+
+from contextlib import nullcontext
+from typing import Any, ContextManager, Mapping, Sequence
+
+import torch
+from lightning.fabric.plugins.precision.utils import _convert_fp_tensor
+from lightning.pytorch.plugins import HalfPrecision
+from lightning.pytorch.plugins.precision.precision import Precision
+from lightning_utilities import apply_to_collection
+from omegaconf import DictConfig, OmegaConf
+from torch import Tensor
+from typing_extensions import override
+
+from nemo.core.classes.common import safe_instantiate
+
+
+_FLASH_PRECISION_ALIASES = {
+    "fp16-flash": "fp16-flash",
+    "bf16-flash": "bf16-flash",
+    # Temporary backward-compatible aliases retained during migration.
+    "fp16-automodel": "fp16-flash",
+    "bf16-automodel": "bf16-flash",
+}
+
+
+def resolve_trainer_cfg(trainer_cfg: DictConfig) -> DictConfig:
+    """
+    Resolves and processes a trainer configuration.
+
+    This function handles specific trainer configuration details:
+    - For half precision setups, replaces precision settings with custom plugins
+    - Instantiates strategy objects from mapping configurations
+    - Instantiates custom callbacks from sequences
+
+    Args:
+        trainer_cfg: A DictConfig containing trainer configuration parameters
+
+    Returns:
+        A processed DictConfig with resolved configuration values
+    """
+    trainer_cfg = OmegaConf.to_container(trainer_cfg, resolve=True)
+
+    # Avoids downcasting 'audio' tensors in half precision setups and enables
+    # the specialized flash precision plugin without mutating global dtype state.
+    precision = trainer_cfg.get("precision")
+    if precision in ("fp16-true", "bf16-true"):
+        trainer_cfg.pop("precision", None)
+        trainer_cfg["plugins"] = [HalfPrecisionForAudio(precision)]
+    elif (flash_precision := _normalize_flash_precision(precision)) is not None:
+        trainer_cfg.pop("precision", None)
+        trainer_cfg["plugins"] = [FlashPrecision(flash_precision)]
+
+    # Allows customizable strategies (eg ModelParallelStrategy) in YAML configs.
+    if (strategy := trainer_cfg.get("strategy", None)) is not None and isinstance(strategy, Mapping):
+        trainer_cfg["strategy"] = safe_instantiate(strategy)
+        # Convert dict-valued nemo_automodel configs to proper dataclass instances.
+        # This must happen AFTER Hydra instantiation because Hydra's recursive
+        # processing chokes on dataclass fields with Union types (e.g. MoEParallelizerConfig).
+        _resolve_automodel_configs(trainer_cfg["strategy"])
+
+    # Allows to add custom callbacks (e.g. NsysCallback) from YAML config.
+    if (cbs := trainer_cfg.get("callbacks", None)) is not None and isinstance(cbs, Sequence):
+        resolved = []
+        for cb in cbs:
+            resolved.append(safe_instantiate(cb))
+        trainer_cfg["callbacks"] = resolved
+
+    return trainer_cfg
+
+
+def _resolve_automodel_configs(strategy) -> None:
+    """Convert plain dicts for ``distributed_config`` and ``moe_config`` to nemo_automodel objects.
+
+    When :class:`AutomodelParallelStrategy` is specified in YAML, ``distributed_config``
+    and ``moe_config`` arrive as plain dicts (Hydra passes them through as-is).
+    This function converts them to proper dataclass instances on the
+    already-instantiated strategy object.
+
+    Does nothing if the strategy doesn't have these attributes or if they are
+    already proper objects (not dicts).
+    """
+    if isinstance(getattr(strategy, '_distributed_config', None), Mapping):
+        from nemo_automodel.components.distributed.config import FSDP2Config
+
+        cfg = strategy._distributed_config
+        # Instantiate any nested _target_ dicts (e.g. a custom mp_policy)
+        resolved = {}
+        for k, v in cfg.items():
+            if isinstance(v, Mapping) and "_target_" in v:
+                resolved[k] = safe_instantiate(v)
+            else:
+                resolved[k] = v
+        strategy._distributed_config = FSDP2Config(**resolved)
+
+    if isinstance(getattr(strategy, '_moe_config', None), Mapping):
+        from nemo_automodel.components.moe.config import MoEParallelizerConfig
+
+        strategy._moe_config = MoEParallelizerConfig(**strategy._moe_config)
+
+
+class HalfPrecisionForAudio(HalfPrecision):
+    """
+    Adjusted Pytorch Lightning plugin for training with half precision.
+    It avoids downcasting audio to bfloat16 when the mini-batch is a dict
+    with 'audio' string in the keys corresponding to audio tensors.
+    """
+
+    @override
+    def convert_input(self, data: Any) -> Any:
+        """
+        Converts input data to the appropriate precision format, preserving audio tensor precision.
+
+        This method overrides the parent class implementation to avoid downcasting tensors
+        with 'audio' in their dictionary keys. It processes input data recursively when
+        encountering nested dictionaries.
+
+        Args:
+            data: The input data to convert (can be tensor, dict, or other types)
+
+        Returns:
+            The converted data with appropriate precision for each element
+        """
+        if not isinstance(data, dict):
+            return super().convert_input(data)
+
+        return _convert_audio_preserving(data, self._desired_input_dtype)
+
+
+class FlashPrecision(Precision):
+    """Precision plugin for flash optimizer training.
+
+    Unlike Lightning's :class:`HalfPrecision`, this does **not** call
+    :func:`torch.set_default_dtype` and does **not** use :func:`torch.autocast`.
+    It's recommended to use this class together with ``flashoptim`` optimizers.
+
+    This ensures that model-specific fp32 escapes (for example custom norms or
+    gating layers) and FlashOptim's master-weight correction terms are never
+    silently downcast by a global dtype override.
+
+    Note: it won't downcast your model's weights to half precision if some of them
+    have already been downcast (manual downcasting) or if the model is using DTensor
+    (in that case you have to downcast them yourself, typically in configure_model()).
+
+    Opt in by setting ``trainer.precision: bf16-flash`` in the YAML config.
+    """
+
+    precision: str = "bf16-flash"
+
+    def __init__(self, precision: str = "bf16-flash") -> None:
+        normalized = _normalize_flash_precision(precision) or precision
+        self.precision = normalized
+        self._desired_input_dtype = torch.bfloat16 if "bf16" in normalized else torch.float16
+
+    @override
+    def convert_module(self, module: torch.nn.Module) -> torch.nn.Module:
+        # Some models manage dtype explicitly inside configure_model() and may
+        # intentionally keep select parameters in fp32. Only cast modules that
+        # are still entirely plain fp32 tensors.
+        if _should_skip_flash_module_conversion(module):
+            return module
+
+        from flashoptim import cast_model
+
+        cast_model(module, dtype=self._desired_input_dtype)
+        return module
+
+    @override
+    def forward_context(self) -> ContextManager:
+        return nullcontext()
+
+    @override
+    def convert_input(self, data: Any) -> Any:
+        if not isinstance(data, dict):
+            return apply_to_collection(
+                data, function=_convert_fp_tensor, dtype=Tensor, dst_type=self._desired_input_dtype
+            )
+
+        return _convert_audio_preserving(data, self._desired_input_dtype)
+
+
+def _convert_audio_preserving(data: dict, dtype: torch.dtype) -> dict:
+    """Convert dict batch to *dtype*, keeping tensors whose key contains ``'audio'`` in fp32."""
+
+    def _convert(v):
+        if isinstance(v, dict):
+            ans = {}
+            for k, v in v.items():
+                if "audio" not in k or not torch.is_tensor(v):
+                    v = _convert(v)
+                ans[k] = v
+            return ans
+        if isinstance(v, torch.Tensor) and torch.is_floating_point(v):
+            return v.to(dtype)
+        return v
+
+    return _convert(data)
+
+
+def _normalize_flash_precision(precision: str | None) -> str | None:
+    if precision is None:
+        return None
+
+    return _FLASH_PRECISION_ALIASES.get(precision)
+
+
+def _should_skip_flash_module_conversion(module: torch.nn.Module) -> bool:
+    """Return True when a module should keep its existing parameter dtype policy."""
+
+    saw_fp_tensor = False
+    for tensor in _iter_module_tensors(module):
+        if not torch.is_floating_point(tensor):
+            continue
+
+        saw_fp_tensor = True
+        if _is_distributed_tensor(tensor):
+            return True
+        if tensor.dtype != torch.float32:
+            return True
+
+    return not saw_fp_tensor
+
+
+def _iter_module_tensors(module: torch.nn.Module):
+    yield from module.parameters()
+    yield from module.buffers()
+
+
+def _is_distributed_tensor(tensor: Tensor) -> bool:
+    if hasattr(tensor, "device_mesh") or hasattr(tensor, "placements"):
+        return True
+
+    try:
+        from torch.distributed.tensor import DTensor
+    except (ImportError, ModuleNotFoundError):
+        return False
+
+    return isinstance(tensor, DTensor)

nemo/utils/trt_utils.py

@@ -0,0 +1,64 @@
+# 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.
+
+try:
+    import tensorrt as trt
+    from polygraphy.backend.trt import CreateConfig, Profile, engine_from_network, network_from_onnx_path, save_engine
+
+    HAVE_TRT = True
+
+except (ImportError, ModuleNotFoundError):
+
+    HAVE_TRT = False
+
+
+def build_engine(
+    onnx_path,
+    output_path,
+    fp16,
+    input_profile=None,
+    enable_refit=False,
+    enable_preview=False,
+    timing_cache=None,
+    workspace_size=0,
+    int8=False,
+    builder_optimization_level=None,
+):
+    print(f"Building TensorRT engine for {onnx_path}: {output_path}")
+    p = Profile()
+    if input_profile:
+        for name, dims in input_profile.items():
+            assert len(dims) == 3
+            p.add(name, min=dims[0], opt=dims[1], max=dims[2])
+
+    preview_features = None
+
+    config_kwargs = {}
+    if workspace_size > 0:
+        config_kwargs["memory_pool_limits"] = {trt.MemoryPoolType.WORKSPACE: workspace_size}
+    engine = engine_from_network(
+        network_from_onnx_path(onnx_path),
+        config=CreateConfig(
+            fp16=fp16,
+            refittable=enable_refit,
+            profiles=[p],
+            preview_features=preview_features,
+            load_timing_cache=timing_cache,
+            int8=int8,
+            builder_optimization_level=builder_optimization_level,
+            **config_kwargs,
+        ),
+        save_timing_cache=timing_cache,
+    )
+    save_engine(engine, path=output_path)