chore: import upstream snapshot with attribution

nemo/core/__init__.py

@@ -0,0 +1,16 @@
+# 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 nemo.core.neural_types
+from nemo.core.classes import *

nemo/core/classes/__init__.py

@@ -0,0 +1,35 @@
+# 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 hydra
+import lightning.pytorch
+import omegaconf
+
+from nemo.core.classes.common import (
+    FileIO,
+    Model,
+    PretrainedModelInfo,
+    Serialization,
+    Typing,
+    is_typecheck_enabled,
+    typecheck,
+)
+from nemo.core.classes.dataset import Dataset, IterableDataset
+from nemo.core.classes.exportable import Exportable, ExportFormat
+from nemo.core.classes.loss import Loss
+from nemo.core.classes.mixins import access_mixins, adapter_mixins, hf_io_mixin
+from nemo.core.classes.modelPT import ModelPT
+from nemo.core.classes.module import NeuralModule
+from nemo.utils import exceptions

nemo/core/classes/dataset.py

@@ -0,0 +1,108 @@
+# 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 typing import Optional
+
+from torch.utils import data
+
+from nemo.core.classes import Serialization, Typing, typecheck
+
+__all__ = ['Dataset', 'IterableDataset']
+
+
+class Dataset(data.Dataset, Typing, Serialization):
+    """Dataset with output ports
+
+    Please Note: Subclasses of IterableDataset should *not* implement input_types.
+    """
+
+    def _collate_fn(self, batch):
+        """
+        A default implementation of a collation function.
+        Users should override this method to define custom data loaders.
+        """
+        return data.dataloader.default_collate(batch)
+
+    @typecheck()
+    def collate_fn(self, batch):
+        """
+        This is the method that user pass as functor to DataLoader.
+        The method optionally performs neural type checking and add types to the outputs.
+
+        Please note, subclasses of Dataset should not implement `input_types`.
+
+        Usage:
+
+        .. code-block:: python
+
+            dataloader = torch.utils.data.DataLoader(
+                    ....,
+                    collate_fn=dataset.collate_fn,
+                    ....
+            )
+
+        Returns:
+            Collated batch, with or without types.
+        """
+        if self.input_types is not None:
+            raise TypeError("Datasets should not implement `input_types` as they are not checked")
+
+        # Simply forward the inner `_collate_fn`
+        return self._collate_fn(batch)
+
+
+class IterableDataset(data.IterableDataset, Typing, Serialization):
+    """Iterable Dataset with output ports
+
+    Please Note: Subclasses of IterableDataset should *not* implement input_types.
+    """
+
+    def _collate_fn(self, batch):
+        """
+        A default implementation of a collation function.
+        Users should override this method to define custom data loaders.
+        """
+        return data.dataloader.default_collate(batch)
+
+    @typecheck()
+    def collate_fn(self, batch):
+        """
+        This is the method that user pass as functor to DataLoader.
+        The method optionally performs neural type checking and add types to the outputs.
+
+        # Usage:
+        dataloader = torch.utils.data.DataLoader(
+                ....,
+                collate_fn=dataset.collate_fn,
+                ....
+        )
+
+        Returns:
+            Collated batch, with or without types.
+        """
+        if self.input_types is not None:
+            raise TypeError("Datasets should not implement `input_types` as they are not checked")
+
+        # Simply forward the inner `_collate_fn`
+        return self._collate_fn(batch)
+
+
+@dataclass
+class DatasetConfig:
+    # ...
+    batch_size: int = 32
+    drop_last: bool = False
+    shuffle: bool = False
+    num_workers: Optional[int] = 0
+    pin_memory: bool = True

nemo/core/classes/exportable.py

@@ -0,0 +1,370 @@
+# 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 abc import ABC
+from typing import Dict, List, Optional, Union
+
+import torch
+from lightning.pytorch.core.module import _jit_is_scripting
+
+from nemo.core.classes import typecheck
+from nemo.core.neural_types import NeuralType
+from nemo.core.utils.neural_type_utils import get_dynamic_axes, get_io_names
+from nemo.utils import logging, monkeypatched
+from nemo.utils.export_utils import (
+    ExportFormat,
+    augment_filename,
+    get_export_format,
+    parse_input_example,
+    rename_onnx_io,
+    replace_for_export,
+    verify_runtime,
+    verify_torchscript,
+    wrap_forward_method,
+)
+
+__all__ = ['ExportFormat', 'Exportable']
+
+
+class Exportable(ABC):
+    """
+    This Interface should be implemented by particular classes derived from nemo.core.NeuralModule or nemo.core.ModelPT.
+    It gives these entities ability to be exported for deployment to formats such as ONNX.
+
+    Usage:
+        # exporting pre-trained model to ONNX file for deployment.
+        model.eval()
+        model.to('cuda')  # or to('cpu') if you don't have GPU
+
+        model.export('mymodel.onnx', [options])  # all arguments apart from `output` are optional.
+    """
+
+    @property
+    def input_module(self):
+        return self
+
+    @property
+    def output_module(self):
+        return self
+
+    def export(
+        self,
+        output: str,
+        input_example=None,
+        verbose=False,
+        do_constant_folding=True,
+        onnx_opset_version=None,
+        check_trace: Union[bool, List[torch.Tensor]] = False,
+        dynamic_axes=None,
+        check_tolerance=0.01,
+        export_modules_as_functions=False,
+        keep_initializers_as_inputs=None,
+        use_dynamo=False,
+    ):
+        """
+        Exports the model to the specified format. The format is inferred from the file extension of the output file.
+
+        Args:
+            output (str): Output file name. File extension be .onnx, .pt, or .ts, and is used to select export
+                path of the model.
+            input_example (list or dict): Example input to the model's forward function. This is used to
+                trace the model and export it to ONNX/TorchScript. If the model takes multiple inputs, then input_example
+                should be a list of input examples. If the model takes named inputs, then input_example
+                should be a dictionary of input examples.
+            verbose (bool): If True, will print out a detailed description of the model's export steps, along with
+                the internal trace logs of the export process.
+            do_constant_folding (bool): If True, will execute constant folding optimization on the model's graph
+                before exporting. This is ONNX specific.
+            onnx_opset_version (int): The ONNX opset version to export the model to. If None, will use a reasonable
+                default version.
+            check_trace (bool): If True, will verify that the model's output matches the output of the traced
+                model, upto some tolerance.
+            dynamic_axes (dict): A dictionary mapping input and output names to their dynamic axes. This is
+                used to specify the dynamic axes of the model's inputs and outputs. If the model takes multiple inputs,
+                then dynamic_axes should be a list of dictionaries. If the model takes named inputs, then dynamic_axes
+                should be a dictionary of dictionaries. If None, will use the dynamic axes of the input_example
+                derived from the NeuralType of the input and output of the model.
+            check_tolerance (float): The tolerance to use when checking the model's output against the traced
+                model's output. This is only used if check_trace is True. Note the high tolerance is used because
+                the traced model is not guaranteed to be 100% accurate.
+            export_modules_as_functions (bool): If True, will export the model's submodules as functions. This is
+                ONNX specific.
+            keep_initializers_as_inputs (bool): If True, will keep the model's initializers as inputs in the onnx graph.
+                This is ONNX specific.
+            use_dynamo (bool): If True, use onnx.dynamo_export() instead of onnx.export(). This is ONNX specific.
+
+        Returns:
+            A tuple of two outputs.
+            Item 0 in the output is a list of outputs, the outputs of each subnet exported.
+            Item 1 in the output is a list of string descriptions. The description of each subnet exported can be
+            used for logging purposes.
+        """
+        all_out = []
+        all_descr = []
+        for subnet_name in self.list_export_subnets():
+            model = self.get_export_subnet(subnet_name)
+            out_name = augment_filename(output, subnet_name)
+            out, descr, out_example = model._export(
+                out_name,
+                input_example=input_example,
+                verbose=verbose,
+                do_constant_folding=do_constant_folding,
+                onnx_opset_version=onnx_opset_version,
+                check_trace=check_trace,
+                dynamic_axes=dynamic_axes,
+                check_tolerance=check_tolerance,
+                export_modules_as_functions=export_modules_as_functions,
+                keep_initializers_as_inputs=keep_initializers_as_inputs,
+                use_dynamo=use_dynamo,
+            )
+            # Propagate input example (default scenario, may need to be overriden)
+            if input_example is not None:
+                input_example = out_example
+            all_out.append(out)
+            all_descr.append(descr)
+            logging.info("Successfully exported {} to {}".format(model.__class__.__name__, out_name))
+        return (all_out, all_descr)
+
+    def _export(
+        self,
+        output: str,
+        input_example=None,
+        verbose=False,
+        do_constant_folding=True,
+        onnx_opset_version=None,
+        check_trace: Union[bool, List[torch.Tensor]] = False,
+        dynamic_axes=None,
+        check_tolerance=0.01,
+        export_modules_as_functions=False,
+        keep_initializers_as_inputs=None,
+        use_dynamo=False,
+    ):
+        my_args = locals().copy()
+        my_args.pop('self')
+
+        self.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+
+        exportables = []
+        for m in self.modules():
+            if isinstance(m, Exportable):
+                exportables.append(m)
+
+        qual_name = self.__module__ + '.' + self.__class__.__qualname__
+        format = get_export_format(output)
+        output_descr = f"{qual_name} exported to {format}"
+
+        # Pytorch's default opset version is too low, using reasonable latest one
+        if onnx_opset_version is None:
+            onnx_opset_version = 17
+
+        try:
+            # Disable typechecks
+            typecheck.set_typecheck_enabled(enabled=False)
+
+            # Allow user to completely override forward method to export
+            forward_method, old_forward_method = wrap_forward_method(self)
+
+            # Set module mode
+            with torch.inference_mode(), torch.no_grad(), torch.jit.optimized_execution(True), _jit_is_scripting():
+
+                if input_example is None:
+                    input_example = self.input_module.input_example()
+
+                # Remove i/o examples from args we propagate to enclosed Exportables
+                my_args.pop('output')
+                my_args.pop('input_example')
+
+                # Run (posibly overridden) prepare methods before calling forward()
+                for ex in exportables:
+                    ex._prepare_for_export(**my_args, noreplace=True)
+                self._prepare_for_export(output=output, input_example=input_example, **my_args)
+
+                input_list, input_dict = parse_input_example(input_example)
+                input_names = self.input_names
+                output_names = self.output_names
+                output_example = self.forward(*input_list, **input_dict)
+                if not isinstance(output_example, tuple):
+                    output_example = (output_example,)
+
+                if check_trace:
+                    if isinstance(check_trace, bool):
+                        check_trace_input = [input_example]
+                    else:
+                        check_trace_input = check_trace
+
+                if format == ExportFormat.TORCHSCRIPT:
+                    jitted_model = torch.jit.trace_module(
+                        self,
+                        {"forward": tuple(input_list) + tuple(input_dict.values())},
+                        strict=True,
+                        check_trace=check_trace,
+                        check_tolerance=check_tolerance,
+                    )
+                    jitted_model = torch.jit.freeze(jitted_model)
+                    if verbose:
+                        logging.info(f"JIT code:\n{jitted_model.code}")
+                    jitted_model.save(output)
+                    jitted_model = torch.jit.load(output)
+
+                    if check_trace:
+                        verify_torchscript(jitted_model, output, check_trace_input, check_tolerance)
+                elif format == ExportFormat.ONNX:
+                    # dynamic axis is a mapping from input/output_name => list of "dynamic" indices
+                    if dynamic_axes is None:
+                        dynamic_axes = self.dynamic_shapes_for_export(use_dynamo)
+                    # When using dynamo export, use dynamic_shapes instead of dynamic_axes
+                    if use_dynamo:
+                        dynamic_shapes = dynamic_axes
+                    if use_dynamo:
+                        typecheck.enable_wrapping(enabled=False)
+                        # https://github.com/pytorch/pytorch/issues/126339
+                        with monkeypatched(torch.nn.RNNBase, "flatten_parameters", lambda *args: None):
+                            logging.info(f"Running export.export, dynamic shapes:{dynamic_shapes}\n")
+
+                            # We have to use different types of arguments for dynamo_export to achieve
+                            # same external weights behaviour as onnx.export :
+                            # https://github.com/pytorch/pytorch/issues/126479
+                            # https://github.com/pytorch/pytorch/issues/126269
+                            mem_params = sum([param.nelement() * param.element_size() for param in self.parameters()])
+                            mem_bufs = sum([buf.nelement() * buf.element_size() for buf in self.buffers()])
+                            mem = mem_params + mem_bufs
+
+                            if mem > 2 * 1000 * 1000 * 1000:
+                                ex_model = torch.export.export(
+                                    self,
+                                    tuple(input_list),
+                                    kwargs=input_dict,
+                                    dynamic_shapes=dynamic_shapes,
+                                    strict=False,
+                                )
+                                ex_model = ex_model.run_decompositions()
+                                model_state = ex_model.state_dict
+                            else:
+                                model_state = None
+                                ex_model = self
+
+                            options = torch.onnx.ExportOptions(dynamic_shapes=True, op_level_debug=True)
+                            ex = torch.onnx.dynamo_export(ex_model, *input_list, **input_dict, export_options=options)
+                            ex.save(output, model_state=model_state)
+
+                            del ex
+                            del ex_model
+                            # Rename I/O after save - don't want to risk modifying ex._model_proto
+                            rename_onnx_io(output, input_names, output_names)
+                    else:
+                        torch.onnx.export(
+                            self,
+                            input_example,
+                            output,
+                            input_names=input_names,
+                            output_names=output_names,
+                            verbose=verbose,
+                            do_constant_folding=do_constant_folding,
+                            dynamic_axes=dynamic_axes,
+                            opset_version=onnx_opset_version,
+                            keep_initializers_as_inputs=keep_initializers_as_inputs,
+                            export_modules_as_functions=export_modules_as_functions,
+                            dynamo=False,  # Use legacy TorchScript-based exporter for LSTM compatibility
+                        )
+
+                    if check_trace:
+                        verify_runtime(self, output, check_trace_input, input_names, check_tolerance=check_tolerance)
+                else:
+                    raise ValueError(f'Encountered unknown export format {format}.')
+        finally:
+            typecheck.enable_wrapping(enabled=True)
+            typecheck.set_typecheck_enabled(enabled=True)
+            if forward_method:
+                type(self).forward = old_forward_method
+            self._export_teardown()
+        return (output, output_descr, output_example)
+
+    @property
+    def disabled_deployment_input_names(self) -> List[str]:
+        """Implement this method to return a set of input names disabled for export"""
+        return []
+
+    @property
+    def disabled_deployment_output_names(self) -> List[str]:
+        """Implement this method to return a set of output names disabled for export"""
+        return []
+
+    @property
+    def supported_export_formats(self) -> List[ExportFormat]:
+        """Implement this method to return a set of export formats supported. Default is all types."""
+        return [ExportFormat.ONNX, ExportFormat.TORCHSCRIPT]
+
+    def _prepare_for_export(self, **kwargs):
+        """
+        Override this method to prepare module for export. This is in-place operation.
+        Base version does common necessary module replacements (Apex etc)
+        """
+        if not 'noreplace' in kwargs:
+            replace_for_export(self)
+
+    def _export_teardown(self):
+        """
+        Override this method for any teardown code after export.
+        """
+        pass
+
+    @property
+    def input_names(self):
+        return get_io_names(self.input_module.input_types_for_export, self.disabled_deployment_input_names)
+
+    @property
+    def output_names(self):
+        return get_io_names(self.output_module.output_types_for_export, self.disabled_deployment_output_names)
+
+    @property
+    def input_types_for_export(self) -> Optional[Dict[str, NeuralType]]:
+        return self.input_types
+
+    @property
+    def output_types_for_export(self):
+        return self.output_types
+
+    def dynamic_shapes_for_export(self, use_dynamo=False):
+        return get_dynamic_axes(self.input_module.input_types_for_export, self.input_names, use_dynamo)
+
+    def get_export_subnet(self, subnet=None):
+        """
+        Returns Exportable subnet model/module to export
+        """
+        if subnet is None or subnet == 'self':
+            return self
+        else:
+            return getattr(self, subnet)
+
+    def list_export_subnets(self):
+        """
+        Returns default set of subnet names exported for this model
+        First goes the one receiving input (input_example)
+        """
+        return ['self']
+
+    def get_export_config(self):
+        """
+        Returns export_config dictionary
+        """
+        return getattr(self, 'export_config', {})
+
+    def set_export_config(self, args):
+        """
+        Sets/updates export_config dictionary
+        """
+        ex_config = self.get_export_config()
+        ex_config.update(args)
+        self.export_config = ex_config

nemo/core/classes/loss.py

@@ -0,0 +1,26 @@
+# 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.core.classes.common import Serialization, Typing
+
+__all__ = ['Loss']
+
+
+class Loss(torch.nn.modules.loss._Loss, Typing, Serialization):
+    """Inherit this class to implement custom loss."""
+
+    def __init__(self, **kwargs):
+        super(Loss, self).__init__(**kwargs)

nemo/core/classes/mixins/__init__.py

@@ -0,0 +1,28 @@
+# 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 nemo.core.classes.mixins.access_mixins import AccessMixin, set_access_cfg
+from nemo.core.classes.mixins.adapter_mixin_strategies import (
+    ResidualAddAdapterStrategy,
+    ResidualAddAdapterStrategyConfig,
+    ReturnResultAdapterStrategy,
+    ReturnResultAdapterStrategyConfig,
+)
+from nemo.core.classes.mixins.adapter_mixins import (
+    AdapterModelPTMixin,
+    AdapterModuleMixin,
+    get_registered_adapter,
+    register_adapter,
+)
+from nemo.core.classes.mixins.hf_io_mixin import HuggingFaceFileIO

nemo/core/classes/mixins/access_mixins.py

@@ -0,0 +1,146 @@
+# 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 abc import ABC
+from typing import Optional
+
+import torch
+from omegaconf import DictConfig
+
+_DEFAULT_ACCESS_GUID = "default"
+_ACCESS_CFG = DictConfig({_DEFAULT_ACCESS_GUID: {"detach": False, "convert_to_cpu": False}})
+_ACCESS_ENABLED = {_DEFAULT_ACCESS_GUID: False}
+
+
+def set_access_cfg(cfg: DictConfig, guid: Optional[str] = None):
+    if cfg is None or not isinstance(cfg, DictConfig):
+        raise TypeError(f"cfg must be a DictConfig")
+    global _ACCESS_CFG
+    global _DEFAULT_ACCESS_GUID
+    if guid is not None:
+        _ACCESS_CFG[guid] = cfg
+    else:
+        _ACCESS_CFG[_DEFAULT_ACCESS_GUID] = cfg
+
+
+class AccessMixin(ABC):
+    """
+    Allows access to output of intermediate layers of a model
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._registry = {}  # dictionary of lists
+
+    def register_accessible_tensor(self, name, tensor):
+        """
+        Register tensor for later use.
+        """
+        if self.access_cfg.get('convert_to_cpu', False):
+            tensor = tensor.cpu()
+
+        if self.access_cfg.get('detach', False):
+            tensor = tensor.detach()
+
+        if not hasattr(self, '_registry'):
+            self._registry = {}
+
+        if name not in self._registry:
+            self._registry[name] = []
+
+        self._registry[name].append(tensor)
+
+    @classmethod
+    def get_module_registry(cls, module: torch.nn.Module):
+        """
+        Extract all registries from named submodules, return dictionary where
+        the keys are the flattened module names, the values are the internal registry
+        of each such module.
+        """
+        module_registry = {}
+        for name, m in module.named_modules():
+            if hasattr(m, '_registry') and len(m._registry) > 0:
+                module_registry[name] = m._registry
+        return module_registry
+
+    def reset_registry(self: torch.nn.Module, registry_key: Optional[str] = None):
+        """
+        Reset the registries of all named sub-modules
+        """
+        if hasattr(self, "_registry"):
+            if registry_key is None:
+                self._registry.clear()
+            else:
+                if registry_key in self._registry:
+                    self._registry.pop(registry_key)
+                else:
+                    raise KeyError(
+                        f"Registry key `{registry_key}` provided, but registry does not have this key.\n"
+                        f"Available keys in registry : {list(self._registry.keys())}"
+                    )
+
+        for _, m in self.named_modules():
+            if hasattr(m, "_registry"):
+                if registry_key is None:
+                    m._registry.clear()
+                else:
+                    if registry_key in self._registry:
+                        self._registry.pop(registry_key)
+                    else:
+                        raise KeyError(
+                            f"Registry key `{registry_key}` provided, but registry does not have this key.\n"
+                            f"Available keys in registry : {list(self._registry.keys())}"
+                        )
+
+        # Explicitly disable registry cache after reset
+        AccessMixin.set_access_enabled(access_enabled=False, guid=getattr(self, "model_guid", None))
+
+    @property
+    def access_cfg(self):
+        """
+        Returns:
+            The global access config shared across all access mixin modules.
+        """
+        global _ACCESS_CFG
+        global _DEFAULT_ACCESS_GUID
+        guid = self.model_guid if getattr(self, "model_guid", None) else _DEFAULT_ACCESS_GUID
+        if hasattr(self, "propagate_model_guid"):
+            self.propagate_model_guid()
+        if guid not in _ACCESS_CFG:
+            _ACCESS_CFG[guid] = DictConfig({})
+        return _ACCESS_CFG[guid]
+
+    @classmethod
+    def update_access_cfg(cls, cfg: dict, guid: Optional[str] = None):
+        global _ACCESS_CFG
+        global _DEFAULT_ACCESS_GUID
+        guid = guid if guid is not None else _DEFAULT_ACCESS_GUID
+        if guid not in _ACCESS_CFG:
+            _ACCESS_CFG[guid] = cfg
+        else:
+            _ACCESS_CFG[guid].update(cfg)
+
+    @classmethod
+    def is_access_enabled(cls, guid: Optional[str] = None):
+        global _ACCESS_ENABLED
+        global _DEFAULT_ACCESS_GUID
+        guid = guid if guid is not None else _DEFAULT_ACCESS_GUID
+        return _ACCESS_ENABLED.get(guid, False)
+
+    @classmethod
+    def set_access_enabled(cls, access_enabled: bool, guid: Optional[str] = None):
+        global _ACCESS_ENABLED
+        global _DEFAULT_ACCESS_GUID
+        guid = guid if guid is not None else _DEFAULT_ACCESS_GUID
+        _ACCESS_ENABLED[guid] = access_enabled

nemo/core/classes/mixins/adapter_mixin_strategies.py

@@ -0,0 +1,259 @@
+# 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 abc import ABC
+from dataclasses import dataclass
+from typing import Any, Dict, List, Tuple, Union
+
+import torch
+
+from nemo.core.classes.mixins import AccessMixin
+
+
+class AbstractAdapterStrategy(ABC):
+    def forward(self, input: torch.Tensor, adapter: torch.nn.Module, *, module: 'AdapterModuleMixin'):
+        """
+        Forward method that defines how the output of the adapter should be merged with the input, or if it
+        should be merged at all.
+
+        Also provides the module that called this strategy - thereby allowing access to all other
+        adapters in the calling module. This can be useful if one adapter is a meta adapter, that
+        combines the outputs of various adapters. In such a case, the input can be forwarded across
+        all other adapters, collecting their outputs, and those outputs can then be merged via some
+        strategy. For example, refer to :
+
+        - [AdapterFusion: Non-Destructive Task Composition for Transfer Learning](https://arxiv.org/abs/2005.00247)
+        - [Exploiting Adapters for Cross-lingual Low-resource Speech Recognition](https://arxiv.org/abs/2105.11905)
+
+        Args:
+            input: Original output tensor of the module, or the output of the previous adapter (if more than
+                one adapters are enabled).
+            adapter: The adapter module that is currently required to perform the forward pass.
+            module: The calling module, in its entirety. It is a module that implements `AdapterModuleMixin`,
+                therefore the strategy can access all other adapters in this module via `module.adapter_layer`.
+
+        Returns:
+            The result tensor, after one of the active adapters has finished its forward passes.
+        """
+        raise NotImplementedError()
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
+
+
+class ReturnResultAdapterStrategy(AbstractAdapterStrategy):
+    """
+    An implementation of an adapter strategy that simply returns the result of the adapter.
+    Supports stochastic
+    """
+
+    def forward(self, input: torch.Tensor, adapter: torch.nn.Module, *, module: 'AdapterModuleMixin'):
+        """
+        A basic strategy, which simply returns the result of the adapter's calculation as the output.
+
+        Args:
+            input: Original output tensor of the module, or the output of the previous adapter (if more than
+                one adapters are enabled).
+            adapter: The adapter module that is currently required to perform the forward pass.
+            module: The calling module, in its entirety. It is a module that implements `AdapterModuleMixin`,
+                therefore the strategy can access all other adapters in this module via `module.adapter_layer`.
+
+        Returns:
+            The result tensor, after one of the active adapters has finished its forward passes.
+        """
+        result = self.compute_output(input, adapter, module=module)
+
+        return result
+
+    def compute_output(
+        self,
+        input: Union[torch.Tensor, List[torch.Tensor], Tuple[torch.Tensor], Dict[str, Any]],
+        adapter: torch.nn.Module,
+        *,
+        module: 'AdapterModuleMixin',
+    ) -> torch.Tensor:
+        """
+        Compute the output of a single adapter to some input.
+
+        Args:
+            input: Original output tensor of the module, or the output of the previous adapter (if more than
+                one adapters are enabled).
+            adapter: The adapter module that is currently required to perform the forward pass.
+            module: The calling module, in its entirety. It is a module that implements `AdapterModuleMixin`,
+                therefore the strategy can access all other adapters in this module via `module.adapter_layer`.
+
+        Returns:
+            The result tensor, after one of the active adapters has finished its forward passes.
+        """
+        if isinstance(input, (list, tuple)):
+            out = adapter(*input)
+        elif isinstance(input, dict):
+            out = adapter(**input)
+        else:
+            out = adapter(input)
+        return out
+
+
+@dataclass
+class ReturnResultAdapterStrategyConfig:
+    _target_: str = "{0}.{1}".format(
+        ReturnResultAdapterStrategy.__module__, ReturnResultAdapterStrategy.__name__
+    )  # mandatory field
+
+
+class ResidualAddAdapterStrategy(AbstractAdapterStrategy):
+    """
+    An implementation of residual addition of an adapter module with its input.
+    Supports stochastic depth regularization.
+    """
+
+    def __init__(self, stochastic_depth: float = 0.0, l2_lambda: float = 0.0):
+        """
+        An implementation of residual addition of an adapter module with its input.
+        Performs output = input + adapter(input).
+
+        Args:
+            stochastic_depth: float, when greater than one, can optionally dropout the output of
+                the adapter's forward pass.
+            l2_lambda: L2 norm of the difference between the original input to the function, and the adapter's
+                output result. Disabled if set to 0.0.
+        """
+        super().__init__()
+        self.stochastic_depth = stochastic_depth
+        self.l2_lambda = l2_lambda
+
+    def forward(self, input: torch.Tensor, adapter: torch.nn.Module, *, module: 'AdapterModuleMixin'):
+        """
+        A basic strategy, comprising of a residual connection over the input, after forward pass by
+        the underlying adapter.
+
+        Args:
+            input: Original output tensor of the module, or the output of the previous adapter (if more than
+                one adapters are enabled).
+            adapter: The adapter module that is currently required to perform the forward pass.
+            module: The calling module, in its entirety. It is a module that implements `AdapterModuleMixin`,
+                therefore the strategy can access all other adapters in this module via `module.adapter_layer`.
+
+        Returns:
+            The result tensor, after one of the active adapters has finished its forward passes.
+        """
+        out = self.compute_output(input, adapter, module=module)
+
+        # If not in training mode, or probability of stochastic depth is 0, skip step.
+        p = self.stochastic_depth
+        if not module.training or p == 0.0:
+            pass
+        else:
+            out = self.apply_stochastic_depth(out, input, adapter, module=module)
+
+        # Return the residual connection output = input + adapter(input)
+        result = input + out
+
+        # If l2_lambda is activated, register the loss value
+        self.compute_auxiliary_losses(result, input, adapter, module=module)
+
+        return result
+
+    def compute_output(
+        self, input: torch.Tensor, adapter: torch.nn.Module, *, module: 'AdapterModuleMixin'
+    ) -> torch.Tensor:
+        """
+        Compute the output of a single adapter to some input.
+
+        Args:
+            input: Original output tensor of the module, or the output of the previous adapter (if more than
+                one adapters are enabled).
+            adapter: The adapter module that is currently required to perform the forward pass.
+            module: The calling module, in its entirety. It is a module that implements `AdapterModuleMixin`,
+                therefore the strategy can access all other adapters in this module via `module.adapter_layer`.
+
+        Returns:
+            The result tensor, after one of the active adapters has finished its forward passes.
+        """
+        out = adapter(input)
+        return out
+
+    def apply_stochastic_depth(
+        self, output: torch.Tensor, input: torch.Tensor, adapter: torch.nn.Module, *, module: 'AdapterModuleMixin'
+    ):
+        """
+        Compute and apply stochastic depth if probability is greater than 0.
+
+        Args:
+            output: The result tensor, after one of the active adapters has finished its forward passes.
+            input: Original output tensor of the module, or the output of the previous adapter (if more than
+                one adapters are enabled).
+            adapter: The adapter module that is currently required to perform the forward pass.
+            module: The calling module, in its entirety. It is a module that implements `AdapterModuleMixin`,
+                therefore the strategy can access all other adapters in this module via `module.adapter_layer`.
+
+        Returns:
+            The result tensor, after stochastic depth has been potentially applied to it.
+        """
+        # Perform stochastic depth if needed.
+        p = self.stochastic_depth
+        if p < 0.0 or p > 1.0:
+            raise ValueError(f"Stochastic depth probability has to be between 0 and 1, but got {p}")
+
+        # Apply stochastic depth to the output of adapter.
+        keep_prob = 1.0 - p
+        shape = [1] * output.ndim
+        noise = torch.empty(shape, dtype=output.dtype, device=output.device)
+        noise = noise.bernoulli_(keep_prob)
+        if keep_prob > 0.0:  # Done to normalize activation for inference mode
+            noise.div_(keep_prob)
+
+        output = noise * output
+
+        return output
+
+    def compute_auxiliary_losses(
+        self, output: torch.Tensor, input: torch.Tensor, adapter: torch.nn.Module, *, module: 'AdapterModuleMixin'
+    ):
+        """
+        Compute any auxiliary losses and preserve it in the tensor registry.
+
+        Args:
+            output: The result tensor, after one of the active adapters has finished its forward passes.
+            input: Original output tensor of the module, or the output of the previous adapter (if more than
+                one adapters are enabled).
+            adapter: The adapter module that is currently required to perform the forward pass.
+            module: The calling module, in its entirety. It is a module that implements `AdapterModuleMixin`,
+                therefore the strategy can access all other adapters in this module via `module.adapter_layer`.
+        """
+        if module.training and self.l2_lambda > 0.0:
+            if not isinstance(adapter, AccessMixin):
+                raise ValueError(f"Module {adapter.__class__.__name__} does not implement AccessMixin !")
+
+            # Only add auxiliary loss if adapter has trainable parameters that require gradients
+            if next(adapter.parameters()).requires_grad is True:
+                # Check if globally allowed to compute aux loss
+                compute_aux_loss = adapter.access_cfg.get('compute_adapter_loss', True)
+
+                if compute_aux_loss:
+                    # if l2 lambda is enabled, also enable AccessMixin
+                    adapter.set_access_enabled(access_enabled=True, guid=getattr(self, "model_guid", None))
+
+                    l2_loss = self.l2_lambda * (input - output).square().reshape(input.size(0), -1).sum(dim=-1).mean()
+                    adapter.register_accessible_tensor(name='adapter_loss', tensor=l2_loss)
+
+
+@dataclass
+class ResidualAddAdapterStrategyConfig:
+    stochastic_depth: float = 0.0
+    l2_lambda: float = 0.0
+
+    _target_: str = "{0}.{1}".format(
+        ResidualAddAdapterStrategy.__module__, ResidualAddAdapterStrategy.__name__
+    )  # mandatory field

nemo/core/classes/mixins/hf_io_mixin.py

@@ -0,0 +1,264 @@
+# 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 abc import ABC
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Union
+
+from huggingface_hub import HfApi, ModelCard, ModelCardData
+from huggingface_hub import get_token as get_hf_token
+from huggingface_hub.hf_api import ModelInfo
+from huggingface_hub.utils import SoftTemporaryDirectory
+
+
+class HuggingFaceFileIO(ABC):
+    """
+    Mixin that provides Hugging Face file IO functionality for NeMo models.
+    It is usually implemented as a mixin to `ModelPT`.
+
+    This mixin provides the following functionality:
+    - `search_huggingface_models()`: Search the hub programmatically via some model filter.
+    - `push_to_hf_hub()`: Push a model to the hub.
+    """
+
+    @classmethod
+    def get_hf_model_filter(cls) -> Dict[str, Any]:
+        """
+        Generates a filter for HuggingFace models.
+
+        Additionaly includes default values of some metadata about results returned by the Hub.
+
+        Metadata:
+            resolve_card_info: Bool flag, if set, returns the model card metadata. Default: False.
+            limit_results: Optional int, limits the number of results returned.
+
+        Returns:
+            A dict representing the arguments passable to huggingface list_models().
+        """
+        model_filter = dict(
+            author=None,
+            filter=['nemo'],
+            model_name=None,
+            limit=None,
+            full=None,
+            cardData=False,
+        )
+
+        return model_filter
+
+    @classmethod
+    def search_huggingface_models(cls, model_filter: Optional[Dict[str, Any]] = None) -> Iterable['ModelInfo']:
+        """
+        Should list all pre-trained models available via Hugging Face Hub.
+
+        The following metadata can be passed via the `model_filter` for additional results.
+        Metadata:
+
+            resolve_card_info: Bool flag, if set, returns the model card metadata. Default: False.
+
+            limit_results: Optional int, limits the number of results returned.
+
+        .. code-block:: python
+
+            # You can replace <DomainSubclass> with any subclass of ModelPT.
+            from nemo.core import ModelPT
+
+            # Get default filter dict
+            filt = <DomainSubclass>.get_hf_model_filter()
+
+            # Make any modifications to the filter as necessary
+            filt['filter'].append('en')  # Add language filter
+            filt['filter'].append('automatic-speech-recognition')  # Add task filter
+
+            # Add any metadata to the filter as needed (kwargs to list_models)
+            filt['limit'] = 5
+
+            # Obtain model info
+            model_infos = <DomainSubclass>.search_huggingface_models(model_filter=filt)
+
+            # Browse through cards and select an appropriate one
+            card = model_infos[0]
+
+            # Restore model using `modelId` of the card.
+            model = ModelPT.from_pretrained(card.modelId)
+
+        Args:
+            model_filter: Optional Dictionary (for Hugging Face Hub kwargs)
+                that filters the returned list of compatible model cards, and selects all results from each filter.
+                Users can then use `model_card.modelId` in `from_pretrained()` to restore a NeMo Model.
+
+        Returns:
+            A list of ModelInfo entries.
+        """
+        # Resolve model filter if not provided as argument
+        if model_filter is None:
+            model_filter = cls.get_hf_model_filter()
+
+        # Check if api token exists, use if it does
+        hf_token = get_hf_token()
+
+        # Search for all valid models after filtering
+        api = HfApi()
+
+        results = api.list_models(token=hf_token, sort="lastModified", **model_filter)  # type: Iterable[ModelInfo]
+
+        return results
+
+    def push_to_hf_hub(
+        self,
+        repo_id: str,
+        *,
+        pack_nemo_file: bool = True,
+        model_card: Optional['ModelCard'] | object | str = None,
+        commit_message: str = "Push model using huggingface_hub.",
+        private: bool = False,
+        api_endpoint: Optional[str] = None,
+        token: Optional[str] = None,
+        branch: Optional[str] = None,
+        allow_patterns: Optional[Union[List[str], str]] = None,
+        ignore_patterns: Optional[Union[List[str], str]] = None,
+        delete_patterns: Optional[Union[List[str], str]] = None,
+    ):
+        """
+        Upload model checkpoint to the Hub.
+
+        Use `allow_patterns` and `ignore_patterns` to precisely filter which files should be pushed to the hub. Use
+        `delete_patterns` to delete existing remote files in the same commit. See [`upload_folder`] reference for more
+        details.
+
+        Args:
+            repo_id (`str`):
+                ID of the repository to push to (example: `"username/my-model"`).
+            pack_nemo_file (`bool`, *optional*, defaults to `True`): Whether to pack the model checkpoint and
+                configuration into a single `.nemo` file. If set to false, uploads the contents of the directory
+                containing the model checkpoint and configuration plus additional artifacts.
+            model_card (`ModelCard`, *optional*): Model card to upload with the model. If None, will use the model
+                card template provided by the class itself via `generate_model_card()`. Any object that implements
+                str(obj) can be passed here. Two keyword replacements are passed to `generate_model_card()`:
+                `model_name` and `repo_id`. If the model card generates a string, and it contains `{model_name}` or
+                `{repo_id}`, they will be replaced with the actual values.
+            commit_message (`str`, *optional*):
+                Message to commit while pushing.
+            private (`bool`, *optional*, defaults to `False`):
+                Whether the repository created should be private.
+            api_endpoint (`str`, *optional*):
+                The API endpoint to use when pushing the model to the hub.
+            token (`str`, *optional*):
+                The token to use as HTTP bearer authorization for remote files. By default, it will use the token
+                cached when running `huggingface-cli login`.
+            branch (`str`, *optional*):
+                The git branch on which to push the model. This defaults to `"main"`.
+            allow_patterns (`List[str]` or `str`, *optional*):
+                If provided, only files matching at least one pattern are pushed.
+            ignore_patterns (`List[str]` or `str`, *optional*):
+                If provided, files matching any of the patterns are not pushed.
+            delete_patterns (`List[str]` or `str`, *optional*):
+                If provided, remote files matching any of the patterns will be deleted from the repo.
+
+        Returns:
+            The url of the uploaded HF repo.
+        """
+        if "/" not in repo_id or len(repo_id.split("/")) != 2:
+            raise ValueError("Invalid repo_id provided. Please provide a repo_id of the form `username/repo-name`.")
+
+        domain_name, model_name = repo_id.split("/")
+
+        if token is None:
+            token = get_hf_token()
+
+        api = HfApi(endpoint=api_endpoint, token=token)
+        repo_id = api.create_repo(repo_id=repo_id, private=private, exist_ok=True).repo_id
+
+        # Push the files to the repo in a single commit
+        with SoftTemporaryDirectory() as tmp:
+            saved_path = Path(tmp) / repo_id
+            saved_path.mkdir(parents=True, exist_ok=True)
+
+            # Save nemo file in temp dir
+            # Get SaveRestoreConnector from subclass implementation
+            if not hasattr(self, '_save_restore_connector'):
+                raise NotImplementedError(
+                    "Model must implement a `_save_restore_connector` property to push to the HuggingFace Hub."
+                )
+
+            # We want to save a NeMo file, but not pack its contents into a tarfile by default
+            save_restore_connector = self._save_restore_connector
+            save_restore_connector.pack_nemo_file = pack_nemo_file
+
+            nemo_filepath = saved_path / f"{model_name}.nemo"
+            self.save_to(nemo_filepath)
+
+            # Save model card in temp dir
+            if model_card is None:
+                card_model_name = model_name.replace("_", " ").split(" ")
+                card_model_name = " ".join([word.capitalize() for word in card_model_name])
+                template_kwargs = {
+                    'model_name': card_model_name,
+                    'repo_id': repo_id,
+                }
+                # Generate model card from subclass that implements this method
+                model_card = self.generate_model_card(type='hf', template_kwargs=template_kwargs)
+
+            # Convert model card to str
+            model_card = str(model_card)
+
+            # Write model card to temp dir
+            model_card_filepath = saved_path / "README.md"
+            model_card_filepath.write_text(str(model_card), encoding='utf-8', errors='ignore')
+
+            api.upload_folder(
+                repo_id=repo_id,
+                repo_type="model",
+                folder_path=saved_path,
+                commit_message=commit_message,
+                revision=branch,
+                allow_patterns=allow_patterns,
+                ignore_patterns=ignore_patterns,
+                delete_patterns=delete_patterns,
+            )
+
+            if branch is None:
+                branch = "main"
+
+            return f"https://huggingface.co/{repo_id}/tree/{branch}"
+
+    def _get_hf_model_card(self, template: str, template_kwargs: Optional[Dict[str, str]] = None):
+        """
+        Generate a HuggingFace ModelCard from a str template. The template may have markers with `{key}` that will be
+        populated by values from `template_kwargs` if provided.
+
+        Args:
+            template: Str template for the model card.
+            template_kwargs (optional): Dict of key-value pairs to populate the template with.
+
+        Returns:
+            A HuggingFace ModelCard object that can be converted to a model card string.
+        """
+        card_data = ModelCardData(
+            library_name='nemo',
+            tags=['pytorch', 'NeMo'],
+            license='cc-by-4.0',
+            ignore_metadata_errors=True,
+        )
+
+        if 'card_data' not in template_kwargs:
+            template_kwargs['card_data'] = card_data.to_yaml()
+
+        # Update template with kwargs
+        # We need to do a manual replace because not all keys may be provided in the kwargs
+        for key, val in template_kwargs.items():
+            template = template.replace("{" + key.strip() + "}", val)
+
+        hf_model_card = ModelCard(template)
+        return hf_model_card

nemo/core/classes/module.py

@@ -0,0 +1,141 @@
+# 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 contextlib import contextmanager
+
+import torch
+from torch.nn import Module
+
+from nemo.core.classes.common import FileIO, Serialization, Typing
+from nemo.utils import logging
+
+__all__ = ['NeuralModule', 'freeze', 'unfreeze']
+
+
+def freeze(module: Module) -> None:
+    """Freeze all parameters of ``module`` and snapshot their prior ``requires_grad`` state.
+
+    The snapshot is stored on ``module._frozen_grad_map`` so a later call to ``unfreeze(..., partial=True)``
+    can restore the pre-freeze state instead of unconditionally enabling gradients.
+    """
+    grad_map = {pname: param.requires_grad for pname, param in module.named_parameters()}
+    for param in module.parameters():
+        param.requires_grad = False
+    if not hasattr(module, '_frozen_grad_map'):
+        module._frozen_grad_map = grad_map
+    else:
+        module._frozen_grad_map.update(grad_map)
+    module.eval()
+
+
+def unfreeze(module: Module, partial: bool = False) -> None:
+    """Unfreeze parameters of ``module``.
+
+    If ``partial=True``, restore each parameter's ``requires_grad`` from the snapshot recorded by
+    ``freeze(module)``; otherwise enable gradients on every parameter. The snapshot is cleared in
+    both cases and ``module.train()`` is called.
+    """
+    if partial and not hasattr(module, '_frozen_grad_map'):
+        raise ValueError("Cannot unfreeze partially without first freezing the module with `freeze()`")
+
+    for pname, param in module.named_parameters():
+        if not partial:
+            param.requires_grad = True
+        elif pname in module._frozen_grad_map:
+            param.requires_grad = module._frozen_grad_map[pname]
+        else:
+            logging.warning(
+                f"Parameter {pname} not found in list of previously frozen parameters. Unfreezing this parameter."
+            )
+            param.requires_grad = True
+
+    if hasattr(module, '_frozen_grad_map'):
+        delattr(module, '_frozen_grad_map')
+
+    module.train()
+
+
+class NeuralModule(Module, Typing, Serialization, FileIO):
+    """
+    Abstract class offering interface shared between all PyTorch Neural Modules.
+    """
+
+    @property
+    def num_weights(self):
+        """
+        Utility property that returns the total number of parameters of NeuralModule.
+        """
+        return self._num_weights()
+
+    @torch.jit.ignore
+    def _num_weights(self):
+        num: int = 0
+        for p in self.parameters():
+            if p.requires_grad:
+                num += p.numel()
+        return num
+
+    def input_example(self, max_batch=None, max_dim=None):
+        """
+        Override this method if random inputs won't work
+        Returns:
+            A tuple sample of valid input data.
+        """
+
+        return None
+
+    def freeze(self) -> None:
+        r"""Freeze all params for inference. See :func:`freeze` for details."""
+        freeze(self)
+
+    def unfreeze(self, partial: bool = False) -> None:
+        """Unfreeze parameters for training. See :func:`unfreeze` for details.
+
+        Example:
+            ```python
+            model.encoder.freeze()              # caller freezes encoder
+            model.freeze()                      # freezes everything; encoder snapshot preserved
+            model.unfreeze(partial=True)        # decoder unfrozen, encoder stays frozen
+            ```
+        """
+        unfreeze(self, partial=partial)
+
+    @contextmanager
+    def as_frozen(self):
+        """
+        Context manager which temporarily freezes a module, yields control and finally unfreezes the module partially
+        to return to original state.
+
+        Allows for either total unfreeze or partial unfreeze (if the module was explicitly frozen
+        previously with `freeze()`). The `partial` argument is used to determine whether to unfreeze
+        all parameters or only the parameters that were previously unfrozen prior `freeze()`.
+
+        Example:
+            with model.as_frozen():  # by default, partial = True
+                # Do something with the model
+                pass
+
+            # Model's parameters are now back to original state of requires_grad
+        """
+        training_mode = self.training
+        self.freeze()
+        try:
+            yield
+        finally:
+            self.unfreeze(partial=True)
+
+            if training_mode:
+                self.train()
+            else:
+                self.eval()

nemo/core/config/__init__.py

@@ -0,0 +1,48 @@
+# 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.core.config.base_config import Config
+from nemo.core.config.hydra_runner import hydra_runner
+from nemo.core.config.optimizers import (
+    AdadeltaParams,
+    AdagradParams,
+    AdamaxParams,
+    AdamParams,
+    AdamWParams,
+    NovogradParams,
+    OptimizerParams,
+    RMSpropParams,
+    RpropParams,
+    SGDParams,
+    get_optimizer_config,
+    register_optimizer_params,
+)
+from nemo.core.config.pytorch import DataLoaderConfig
+from nemo.core.config.pytorch_lightning import TrainerConfig
+from nemo.core.config.schedulers import (
+    CosineAnnealingParams,
+    InverseSquareRootAnnealingParams,
+    NoamAnnealingParams,
+    PolynomialDecayAnnealingParams,
+    PolynomialHoldDecayAnnealingParams,
+    SchedulerParams,
+    SquareAnnealingParams,
+    SquareRootAnnealingParams,
+    SquareRootConstantSchedulerParams,
+    WarmupAnnealingParams,
+    WarmupHoldSchedulerParams,
+    WarmupSchedulerParams,
+    get_scheduler_config,
+    register_scheduler_params,
+)

nemo/core/config/base_config.py

@@ -0,0 +1,30 @@
+# 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 typing import Optional
+
+__all__ = ['Config']
+
+
+@dataclass
+class Config:
+    """
+    Abstract NeMo Configuration class.
+
+    Args:
+        name: name of the module/dataset/loss/model object (used in serialization, DEFAULT: None)
+    """
+
+    name: Optional[str] = None

nemo/core/config/hydra_runner.py

@@ -0,0 +1,145 @@
+# 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 os
+import sys
+from typing import Any, Callable, Optional
+
+from hydra._internal.utils import _run_hydra, get_args_parser
+from hydra.core.config_store import ConfigStore
+from hydra.types import TaskFunction
+from omegaconf import DictConfig, OmegaConf
+
+
+def _get_gpu_name():
+    try:
+        import pynvml
+    except (ImportError, ModuleNotFoundError):
+        return None
+
+    pynvml.nvmlInit()
+    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+    cuda_capability, _ = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
+    pynvml.nvmlShutdown()
+    if cuda_capability == 8:
+        return "a100"
+    elif cuda_capability == 9:
+        return "h100"
+    else:
+        return None
+
+
+OmegaConf.register_new_resolver("gpu_name", _get_gpu_name)
+
+# multiple interpolated values in the config
+OmegaConf.register_new_resolver("multiply", lambda x, y: x * y, replace=True)
+
+# sum interpolated values in the config
+OmegaConf.register_new_resolver("sum", lambda x, y: x + y, replace=True)
+
+
+def hydra_runner(
+    config_path: Optional[str] = ".", config_name: Optional[str] = None, schema: Optional[Any] = None
+) -> Callable[[TaskFunction], Any]:
+    """
+    Decorator used for passing the Config paths to main function.
+    Optionally registers a schema used for validation/providing default values.
+
+    Args:
+        config_path: Optional path that will be added to config search directory.
+            NOTE: The default value of `config_path` has changed between Hydra 1.0 and Hydra 1.1+.
+            Please refer to https://hydra.cc/docs/next/upgrades/1.0_to_1.1/changes_to_hydra_main_config_path/
+            for details.
+        config_name: Pathname of the config file.
+        schema: Structured config  type representing the schema used for validation/providing default values.
+    """
+
+    def decorator(task_function: TaskFunction) -> Callable[[], None]:
+        @functools.wraps(task_function)
+        def wrapper(cfg_passthrough: Optional[DictConfig] = None) -> Any:
+            # Check it config was passed.
+            if cfg_passthrough is not None:
+                return task_function(cfg_passthrough)
+            else:
+                args = get_args_parser()
+
+                # Parse arguments in order to retrieve overrides
+                parsed_args = args.parse_args()  # type: argparse.Namespace
+
+                # Get overriding args in dot string format
+                overrides = parsed_args.overrides  # type: list
+
+                # Disable the creation of .hydra subdir
+                # https://hydra.cc/docs/tutorials/basic/running_your_app/working_directory
+                overrides.append("hydra.output_subdir=null")
+                # Hydra logging outputs only to stdout (no log file).
+                # https://hydra.cc/docs/configure_hydra/logging
+                overrides.append("hydra/job_logging=stdout")
+
+                # Set run.dir ONLY for ExpManager "compatibility" - to be removed.
+                overrides.append("hydra.run.dir=.")
+
+                # Check if user set the schema.
+                if schema is not None:
+                    # Create config store.
+                    cs = ConfigStore.instance()
+
+                    # Get the correct ConfigStore "path name" to "inject" the schema.
+                    if parsed_args.config_name is not None:
+                        path, name = os.path.split(parsed_args.config_name)
+                        # Make sure the path is not set - as this will disable validation scheme.
+                        if path != '':
+                            sys.stderr.write(
+                                "ERROR Cannot set config file path using `--config-name` when "
+                                "using schema. Please set path using `--config-path` and file name using "
+                                "`--config-name` separately.\n"
+                            )
+                            sys.exit(1)
+                    else:
+                        name = config_name
+
+                    # Register the configuration as a node under the name in the group.
+                    cs.store(name=name, node=schema)  # group=group,
+
+                # Wrap a callable object with name `parse_args`
+                # This is to mimic the ArgParser.parse_args() API.
+                def parse_args(self, args=None, namespace=None):
+                    return parsed_args
+
+                parsed_args.parse_args = parse_args
+
+                # no return value from run_hydra() as it may sometime actually run the task_function
+                # multiple times (--multirun)
+                # argparse_wrapper = _argparse_wrapper(args)
+                argparse_wrapper = parsed_args
+
+                try:
+                    _run_hydra(
+                        args=argparse_wrapper,
+                        args_parser=args,
+                        task_function=task_function,
+                        config_path=config_path,
+                        config_name=config_name,
+                    )
+                finally:
+                    # Import here to avoid circular import
+                    from nemo.lightning.callback_group import CallbackGroup
+
+                    # Ensure on_app_end is called even if _run_hydra raises an exception
+                    CallbackGroup.get_instance().on_app_end()
+
+        return wrapper
+
+    return decorator

nemo/core/config/modelPT.py

@@ -0,0 +1,183 @@
+# 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, field
+from typing import Any, Dict, Optional
+
+from omegaconf import MISSING
+
+from nemo.core import config
+from nemo.core.classes.dataset import DatasetConfig
+from nemo.utils import exp_manager
+
+
+@dataclass
+class SchedConfig:
+    name: str = MISSING
+    min_lr: float = 0.0
+    last_epoch: int = -1
+
+
+@dataclass
+class OptimConfig:
+    name: str = MISSING
+    sched: Optional[SchedConfig] = None
+
+
+@dataclass
+class ModelConfig:
+    """
+    Model component inside ModelPT
+    """
+
+    # ...
+    train_ds: Optional[DatasetConfig] = None
+    validation_ds: Optional[DatasetConfig] = None
+    test_ds: Optional[DatasetConfig] = None
+    optim: Optional[OptimConfig] = None
+
+
+@dataclass
+class HydraConfig:
+    run: Dict[str, Any] = field(default_factory=lambda: {"dir": "."})
+    job_logging: Dict[str, Any] = field(default_factory=lambda: {"root": {"handlers": None}})
+
+
+@dataclass
+class NemoConfig:
+    name: str = MISSING
+    model: ModelConfig = MISSING
+    trainer: config.TrainerConfig = field(
+        default_factory=lambda: config.TrainerConfig(
+            strategy="ddp", enable_checkpointing=False, logger=False, log_every_n_steps=1, accelerator='gpu'
+        )
+    )
+    exp_manager: Optional[Any] = field(default_factory=lambda: exp_manager.ExpManagerConfig())
+    hydra: HydraConfig = field(default_factory=lambda: HydraConfig())
+
+
+class ModelConfigBuilder:
+    def __init__(self, model_cfg: ModelConfig):
+        """
+        Base class for any Model Config Builder.
+
+        A Model Config Builder is a utility class that accepts a ModelConfig dataclass,
+        and via a set of utility methods (that are implemented by the subclassed ModelConfigBuilder),
+        builds a finalized ModelConfig that can be supplied to a NemoModel dataclass as
+        the `model` component.
+
+        Subclasses *must* implement the private method `_finalize_cfg`.
+            Inside this method, they must update `self.model_cfg` with all interdependent config
+            options that need to be set (either updated by user explicitly or with their default value).
+
+            The updated model config must then be preserved in `self.model_cfg`.
+
+        Example:
+            # Create the config builder
+            config_builder = <subclass>ModelConfigBuilder()
+
+            # Update the components of the config that are modifiable
+            config_builder.set_X(X)
+            config_builder.set_Y(Y)
+
+            # Create a "finalized" config dataclass that will contain all the updates
+            # that were specified by the builder
+            model_config = config_builder.build()
+
+            # Use model config as is (or further update values), then create a new Model
+            model = nemo.<domain>.models.<ModelName>Model(cfg=model_config, trainer=Trainer())
+
+        Supported build methods:
+        -   set_train_ds: All model configs can accept a subclass of `DatasetConfig` as their
+                training config. Subclasses can override this method to enable auto-complete
+                by replacing `Optional[DatasetConfig]` with `Optional[<subclass of DatasetConfig>]`.
+
+        -   set_validation_ds: All model configs can accept a subclass of `DatasetConfig` as their
+                validation config. Subclasses can override this method to enable auto-complete
+                by replacing `Optional[DatasetConfig]` with `Optional[<subclass of DatasetConfig>]`.
+
+        -   set_test_ds: All model configs can accept a subclass of `DatasetConfig` as their
+                test config. Subclasses can override this method to enable auto-complete
+                by replacing `Optional[DatasetConfig]` with `Optional[<subclass of DatasetConfig>]`.
+
+        -   set_optim: A build method that supports changes to the Optimizer (and optionally,
+                the Scheduler) used for training the model. The function accepts two inputs -
+
+                `cfg`: A subclass of `OptimizerParams` - any OptimizerParams subclass can be used,
+                    in order to select an appropriate Optimizer. Examples: AdamParams.
+
+                `sched_cfg`: A subclass of `SchedulerParams` - any SchedulerParams subclass can be used,
+                    in order to select an appropriate Scheduler. Examples: CosineAnnealingParams.
+                    Note that this argument is optional.
+
+        -   build(): The method which should return a "finalized" ModelConfig dataclass.
+                Subclasses *should* always override this method, and update the signature
+                of this method with the return type of the Dataclass, so that it enables
+                autocomplete for the user.
+
+                Example:
+                    def build(self) -> EncDecCTCConfig:
+                        return super().build()
+
+        Any additional build methods must be added by subclasses of ModelConfigBuilder.
+
+        Args:
+            model_cfg:
+        """
+        self.model_cfg = model_cfg
+        self.train_ds_cfg = None
+        self.validation_ds_cfg = None
+        self.test_ds_cfg = None
+        self.optim_cfg = None
+
+    def set_train_ds(self, cfg: Optional[DatasetConfig] = None):
+        self.model_cfg.train_ds = cfg
+
+    def set_validation_ds(self, cfg: Optional[DatasetConfig] = None):
+        self.model_cfg.validation_ds = cfg
+
+    def set_test_ds(self, cfg: Optional[DatasetConfig] = None):
+        self.model_cfg.test_ds = cfg
+
+    def set_optim(self, cfg: config.OptimizerParams, sched_cfg: Optional[config.SchedulerParams] = None):
+        @dataclass
+        class WrappedOptimConfig(OptimConfig, cfg.__class__):
+            pass
+
+        # Setup optim
+        optim_name = cfg.__class__.__name__.replace("Params", "").lower()
+        wrapped_cfg = WrappedOptimConfig(name=optim_name, sched=None, **vars(cfg))
+
+        if sched_cfg is not None:
+
+            @dataclass
+            class WrappedSchedConfig(SchedConfig, sched_cfg.__class__):
+                pass
+
+            # Setup scheduler
+            sched_name = sched_cfg.__class__.__name__.replace("Params", "")
+            wrapped_sched_cfg = WrappedSchedConfig(name=sched_name, **vars(sched_cfg))
+
+            wrapped_cfg.sched = wrapped_sched_cfg
+
+        self.model_cfg.optim = wrapped_cfg
+
+    def _finalize_cfg(self):
+        raise NotImplementedError()
+
+    def build(self) -> ModelConfig:
+        # validate config
+        self._finalize_cfg()
+
+        return self.model_cfg

nemo/core/config/optimizers.py

@@ -0,0 +1,295 @@
+# 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 functools import partial
+from typing import Any, Dict, Optional, Tuple
+
+from omegaconf import MISSING, OmegaConf
+
+__all__ = [
+    'OptimizerParams',
+    'AdamParams',
+    'NovogradParams',
+    'SGDParams',
+    'AdadeltaParams',
+    'AdamaxParams',
+    'AdagradParams',
+    'AdamWParams',
+    'RMSpropParams',
+    'RpropParams',
+]
+
+
+@dataclass
+class OptimizerParams:
+    """
+    Base Optimizer params with no values. User can chose it to explicitly override via
+    command line arguments
+    """
+
+    lr: Optional[float] = MISSING
+
+
+@dataclass
+class SGDParams(OptimizerParams):
+    """
+    Default configuration for Adam optimizer.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch.org/docs/stable/optim.html?highlight=sgd#torch.optim.SGD
+    """
+
+    momentum: float = 0
+    dampening: float = 0
+    weight_decay: float = 0
+    nesterov: bool = False
+
+
+@dataclass
+class AdamParams(OptimizerParams):
+    """
+    Default configuration for Adam optimizer.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch.org/docs/stable/optim.html?highlight=adam#torch.optim.Adam
+    """
+
+    # betas: Tuple[float, float] = (0.9, 0.999)
+    eps: float = 1e-08
+    weight_decay: float = 0
+    amsgrad: bool = False
+
+
+@dataclass
+class AdamWParams(OptimizerParams):
+    """
+    Default configuration for AdamW optimizer.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch.org/docs/stable/optim.html#torch.optim.AdamW
+    """
+
+    betas: Tuple[float, float] = (0.9, 0.999)
+    eps: float = 1e-08
+    weight_decay: float = 0
+    amsgrad: bool = False
+
+
+@dataclass
+class AdadeltaParams(OptimizerParams):
+    """
+    Default configuration for Adadelta optimizer.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch.org/docs/stable/optim.html#torch.optim.Adadelta
+    """
+
+    rho: float = 0.9
+    eps: float = 1e-6
+    weight_decay: float = 0
+
+
+@dataclass
+class AdamaxParams(OptimizerParams):
+    """
+    Default configuration for Adamax optimizer.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch.org/docs/stable/optim.html#torch.optim.Adamax
+    """
+
+    betas: Tuple[float, float] = (0.9, 0.999)
+    eps: float = 1e-8
+    weight_decay: float = 0
+
+
+@dataclass
+class AdagradParams(OptimizerParams):
+    """
+    Default configuration for Adagrad optimizer.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch.org/docs/stable/optim.html#torch.optim.Adagrad
+    """
+
+    lr_decay: float = 0
+    weight_decay: float = 0
+    initial_accumulator_value: float = 0
+    eps: float = 1e-10
+
+
+@dataclass
+class RMSpropParams(OptimizerParams):
+    """
+    Default configuration for RMSprop optimizer.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch.org/docs/stable/optim.html#torch.optim.RMSprop
+    """
+
+    alpha: float = 0.99
+    eps: float = 1e-8
+    weight_decay: float = 0
+    momentum: float = 0
+    centered: bool = False
+
+
+@dataclass
+class RpropParams(OptimizerParams):
+    """
+    Default configuration for RpropParams optimizer.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch.org/docs/stable/optim.html#torch.optim.Rprop
+    """
+
+    etas: Tuple[float, float] = (0.5, 1.2)
+    step_sizes: Tuple[float, float] = (1e-6, 50)
+
+
+@dataclass
+class NovogradParams(OptimizerParams):
+    """
+    Configuration of the Novograd optimizer.
+
+    It has been proposed  in "Stochastic Gradient Methods with Layer-wise
+    Adaptive Moments for Training of Deep Networks"
+    (https://arxiv.org/abs/1905.11286)
+
+    Args:
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper "On the Convergence of Adam and Beyond"
+    """
+
+    betas: Tuple[float, float] = (0.95, 0.98)
+    eps: float = 1e-8
+    weight_decay: float = 0
+    grad_averaging: bool = False
+    amsgrad: bool = False
+    luc: bool = False
+    luc_trust: float = 1e-3
+    luc_eps: float = 1e-8
+
+
+@dataclass
+class AdafactorParams(OptimizerParams):
+    """
+    Configuration of the Adafactor optimizer.
+
+    It has been proposed  in "Adafactor: Adaptive Learning Rates with Sublinear Memory Cost"
+    (https://arxiv.org/abs/1804.04235)
+
+    Args:
+        lr (float, optional): learning rate (default: 1e-3)
+        beta1 (float, optional): coefficients used for computing
+            running averages of gradient and its square (default: None)
+        eps (Tuple [float, float] optional)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        scale_parameter (float, optional): scale parameter (default: False)
+        relative_step (bool, optional): whether to use relative step sizes (default: False)
+        warmup_init (bool, optional): whether to warmup the learning rate linearly (default: False)
+    """
+
+    beta1: float = None
+    eps: Tuple[float, float] = (1e-30, 1e-3)
+    clip_threshold: float = 1.0
+    decay_rate: float = 0.8
+    weight_decay: float = 0
+    scale_parameter: bool = True
+    relative_step: bool = False
+    warmup_init: bool = False
+
+
+def register_optimizer_params(name: str, optimizer_params: OptimizerParams):
+    """
+    Checks if the optimizer param name exists in the registry, and if it doesnt, adds it.
+
+    This allows custom optimizer params to be added and called by name during instantiation.
+
+    Args:
+        name: Name of the optimizer. Will be used as key to retrieve the optimizer.
+        optimizer_params: Optimizer class
+    """
+    if name in AVAILABLE_OPTIMIZER_PARAMS:
+        raise ValueError(f"Cannot override pre-existing optimizers. Conflicting optimizer name = {name}")
+
+    AVAILABLE_OPTIMIZER_PARAMS[name] = optimizer_params
+
+
+def get_optimizer_config(name: str, **kwargs: Optional[Dict[str, Any]]) -> OptimizerParams:
+    """
+    Convenience method to obtain a OptimizerParams class and partially instantiate it with optimizer kwargs.
+
+    Args:
+        name: Name of the OptimizerParams in the registry.
+        kwargs: Optional kwargs of the optimizer used during instantiation.
+
+    Returns:
+        a partially instantiated OptimizerParams
+    """
+    if name is None:
+        return kwargs
+
+    if name not in AVAILABLE_OPTIMIZER_PARAMS:
+        raise ValueError(
+            f"Cannot resolve optimizer parameters '{name}'. Available optimizer parameters are : "
+            f"{AVAILABLE_OPTIMIZER_PARAMS.keys()}"
+        )
+
+    scheduler_params = AVAILABLE_OPTIMIZER_PARAMS[name]
+
+    if kwargs is not None and len(kwargs) != 0:
+        kwargs = OmegaConf.create(kwargs)
+        OmegaConf.merge(scheduler_params(), kwargs)
+
+    scheduler_params = partial(scheduler_params, **kwargs)
+    return scheduler_params
+
+
+AVAILABLE_OPTIMIZER_PARAMS = {
+    'optim_params': OptimizerParams,
+    'adam_params': AdamParams,
+    'novograd_params': NovogradParams,
+    'sgd_params': SGDParams,
+    'adadelta_params': AdadeltaParams,
+    'adamax_params': AdamaxParams,
+    'adagrad_params': AdagradParams,
+    'adamw_params': AdamWParams,
+    'rmsprop_params': RMSpropParams,
+    'rprop_params': RpropParams,
+    'adafactor_params': AdafactorParams,
+}

nemo/core/config/pytorch.py

@@ -0,0 +1,45 @@
+# 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 typing import Any, Optional
+
+from omegaconf import MISSING
+
+__all__ = ['DataLoaderConfig']
+
+
+@dataclass
+class DataLoaderConfig:
+    """
+    Configuration of PyTorch DataLoader.
+
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader
+    """
+
+    batch_size: int = MISSING
+    shuffle: bool = False
+    sampler: Optional[Any] = None
+    batch_sampler: Optional[Any] = None
+    num_workers: int = 0
+    collate_fn: Optional[Any] = None
+    pin_memory: bool = False
+    drop_last: bool = False
+    timeout: int = 0
+    worker_init_fn: Optional[Any] = None
+    multiprocessing_context: Optional[Any] = None

nemo/core/config/pytorch_lightning.py

@@ -0,0 +1,84 @@
+# 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 typing import Any, Optional
+
+from hydra.core.config_store import ConfigStore
+
+__all__ = ['TrainerConfig']
+
+
+cs = ConfigStore.instance()
+
+
+@dataclass
+class TrainerConfig:
+    """
+    Configuration of PyTorch Lightning Trainer.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    ..warning:
+        Picked just few params of the PTL trainer for now. This needs to be discussed.
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch-lightning.readthedocs.io/en/latest/common/trainer.html
+    """
+
+    logger: Any = True
+    callbacks: Optional[Any] = None
+    default_root_dir: Optional[str] = None
+    gradient_clip_val: float = 0
+    num_nodes: int = 1
+    enable_progress_bar: bool = True
+    overfit_batches: Any = 0.0
+    check_val_every_n_epoch: int = 1
+    fast_dev_run: bool = False
+    accumulate_grad_batches: Any = 1
+    max_epochs: int = 1000
+    min_epochs: int = 1
+    max_steps: Optional[int] = -1
+    min_steps: Optional[int] = None
+    limit_train_batches: Any = 1.0
+    limit_val_batches: Any = 1.0
+    limit_test_batches: Any = 1.0
+    val_check_interval: Any = 1.0
+    log_every_n_steps: int = 50
+    accelerator: Optional[str] = 'auto'
+    sync_batchnorm: bool = False
+    precision: Any = 32
+    num_sanity_val_steps: int = 2
+    profiler: Optional[Any] = None
+    benchmark: bool = False
+    deterministic: bool = False
+    use_distributed_sampler: bool = True
+    detect_anomaly: bool = False
+    plugins: Optional[Any] = None  # Optional[Union[str, list]]
+    limit_predict_batches: float = 1.0
+    gradient_clip_algorithm: str = 'norm'
+    max_time: Optional[Any] = None  # can be one of Union[str, timedelta, Dict[str, int], None]
+    reload_dataloaders_every_n_epochs: int = 0
+    devices: Any = 'auto'
+    strategy: Any = 'auto'
+    enable_checkpointing: bool = False
+    enable_model_summary: bool = True
+    inference_mode: bool = True
+    barebones: bool = False
+
+
+# Register the trainer config.
+cs.store(
+    group="trainer",
+    name="trainer",
+    node=TrainerConfig,
+)

nemo/core/config/schedulers.py

@@ -0,0 +1,288 @@
+# 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 functools import partial
+from typing import Any, Dict, Optional
+
+
+@dataclass
+class SchedulerParams:
+    """
+    Base configuration for all schedulers.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    last_epoch: int = -1
+
+
+@dataclass
+class SquareRootConstantSchedulerParams(SchedulerParams):
+    """
+    Base configuration for all schedulers.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    constant_steps: Optional[float] = None
+    constant_ratio: Optional[float] = None
+
+
+@dataclass
+class WarmupSchedulerParams(SchedulerParams):
+    """
+    Base configuration for all schedulers.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    max_steps: int = 0
+    warmup_steps: Optional[float] = None
+    warmup_ratio: Optional[float] = None
+
+
+@dataclass
+class WarmupHoldSchedulerParams(WarmupSchedulerParams):
+    """
+    Base configuration for all schedulers.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    hold_steps: Optional[float] = None
+    hold_ratio: Optional[float] = None
+    min_lr: float = 0.0
+
+
+@dataclass
+class WarmupAnnealingHoldSchedulerParams(WarmupSchedulerParams):
+    """
+    Base configuration for all schedulers.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    constant_steps: Optional[float] = None
+    constant_ratio: Optional[float] = None
+    min_lr: float = 0.0
+
+
+@dataclass
+class SquareAnnealingParams(WarmupSchedulerParams):
+    """
+    Square Annealing parameter config
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    min_lr: float = 1e-5
+
+
+@dataclass
+class SquareRootAnnealingParams(WarmupSchedulerParams):
+    """
+    Square Root Annealing parameter config
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    min_lr: float = 0.0
+
+
+@dataclass
+class CosineAnnealingParams(WarmupAnnealingHoldSchedulerParams):
+    """
+    Cosine Annealing parameter config
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    min_lr: float = 0.0
+
+
+@dataclass
+class NoamAnnealingParams(WarmupSchedulerParams):
+    """
+    Cosine Annealing parameter config
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    min_lr: float = 0.0
+
+
+@dataclass
+class NoamHoldAnnealingParams(WarmupHoldSchedulerParams):
+    """
+    Polynomial Hold Decay Annealing parameter config
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    decay_rate: float = 0.5
+
+
+@dataclass
+class WarmupAnnealingParams(WarmupSchedulerParams):
+    """
+    Warmup Annealing parameter config
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    warmup_ratio: Optional[float] = None
+
+
+@dataclass
+class InverseSquareRootAnnealingParams(WarmupSchedulerParams):
+    """
+    Inverse Square Root Annealing parameter config
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+
+@dataclass
+class PolynomialDecayAnnealingParams(WarmupSchedulerParams):
+    """
+    Polynomial Decay Annealing parameter config
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    power: float = 1.0
+    cycle: bool = False
+
+
+@dataclass
+class PolynomialHoldDecayAnnealingParams(WarmupSchedulerParams):
+    """
+    Polynomial Hold Decay Annealing parameter config
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    power: float = 1.0
+    cycle: bool = False
+
+
+"""
+Pytorch Optimizers
+"""
+
+
+@dataclass
+class StepLRParams(SchedulerParams):
+    """
+    Config for StepLR.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    step_size: float = 0.1
+    gamma: float = 0.1
+
+
+@dataclass
+class ExponentialLRParams(SchedulerParams):
+    """
+    Config for ExponentialLR.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    gamma: float = 0.9
+
+
+@dataclass
+class ReduceLROnPlateauParams:
+    """
+    Config for ReduceLROnPlateau.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    mode: str = 'min'
+    factor: float = 0.1
+    patience: int = 10
+    verbose: bool = False
+    threshold: float = 1e-4
+    threshold_mode: str = 'rel'
+    cooldown: int = 0
+    min_lr: float = 0
+    eps: float = 1e-8
+
+
+@dataclass
+class CyclicLRParams(SchedulerParams):
+    """
+    Config for CyclicLR.
+    NOTE:
+    # `scale_fn` is not supported
+
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    base_lr: float = 0.001
+    max_lr: float = 0.1
+    step_size_up: int = 2000
+    step_size_down: Optional[int] = None
+    mode: str = 'triangular'
+    gamma: float = 1.0
+    scale_mode: str = 'cycle'
+    # scale_fn is not supported
+    cycle_momentum: bool = True
+    base_momentum: float = 0.8
+    max_momentum: float = 0.9
+
+
+def register_scheduler_params(name: str, scheduler_params: SchedulerParams):
+    """
+    Checks if the schduler config name exists in the registry, and if it doesnt, adds it.
+
+    This allows custom schedulers to be added and called by name during instantiation.
+
+    Args:
+        name: Name of the optimizer. Will be used as key to retrieve the optimizer.
+        scheduler_params: SchedulerParams class
+    """
+    if name in AVAILABLE_SCHEDULER_PARAMS:
+        raise ValueError(f"Cannot override pre-existing optimizers. Conflicting optimizer name = {name}")
+
+    AVAILABLE_SCHEDULER_PARAMS[name] = scheduler_params
+
+
+def get_scheduler_config(name: str, **kwargs: Optional[Dict[str, Any]]) -> SchedulerParams:
+    """
+    Convenience method to obtain a SchedulerParams class and partially instantiate it with optimizer kwargs.
+
+    Args:
+        name: Name of the SchedulerParams in the registry.
+        kwargs: Optional kwargs of the optimizer used during instantiation.
+
+    Returns:
+        a partially instantiated SchedulerParams
+    """
+    if name not in AVAILABLE_SCHEDULER_PARAMS:
+        raise ValueError(
+            f"Cannot resolve scheduler parameters '{name}'. Available scheduler parameters are : "
+            f"{AVAILABLE_SCHEDULER_PARAMS.keys()}"
+        )
+
+    scheduler_params = AVAILABLE_SCHEDULER_PARAMS[name]
+    scheduler_params = partial(scheduler_params, **kwargs)
+    return scheduler_params
+
+
+AVAILABLE_SCHEDULER_PARAMS = {
+    'SchedulerParams': SchedulerParams,
+    'WarmupPolicyParams': WarmupSchedulerParams,
+    'WarmupHoldPolicyParams': WarmupHoldSchedulerParams,
+    'WarmupAnnealingHoldSchedulerParams': WarmupAnnealingHoldSchedulerParams,
+    'SquareAnnealingParams': SquareAnnealingParams,
+    'SquareRootAnnealingParams': SquareRootAnnealingParams,
+    'InverseSquareRootAnnealingParams': InverseSquareRootAnnealingParams,
+    'SquareRootConstantSchedulerParams': SquareRootConstantSchedulerParams,
+    'CosineAnnealingParams': CosineAnnealingParams,
+    'NoamAnnealingParams': NoamAnnealingParams,
+    'NoamHoldAnnealingParams': NoamHoldAnnealingParams,
+    'WarmupAnnealingParams': WarmupAnnealingParams,
+    'PolynomialDecayAnnealingParams': PolynomialDecayAnnealingParams,
+    'PolynomialHoldDecayAnnealingParams': PolynomialHoldDecayAnnealingParams,
+    'ReduceLROnPlateauParams': ReduceLROnPlateauParams,
+}

nemo/core/config/templates/__init__.py

@@ -0,0 +1,13 @@
+# 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.

nemo/core/config/templates/model_card.py

@@ -0,0 +1,210 @@
+# 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.
+
+
+NEMO_DEFAULT_MODEL_CARD_TEMPLATE = """---
+{card_data}
+---
+
+# {model_name}
+
+<style>
+img {
+ display: inline;
+}
+</style>
+
+[![Model architecture](https://img.shields.io/badge/Model_Arch-PUT-YOUR-ARCHITECTURE-HERE-lightgrey#model-badge)](#model-architecture)
+| [![Model size](https://img.shields.io/badge/Params-PUT-YOUR-MODEL-SIZE-HERE-lightgrey#model-badge)](#model-architecture)
+| [![Language](https://img.shields.io/badge/Language-PUT-YOUR-LANGUAGE-HERE-lightgrey#model-badge)](#datasets)
+
+**Put a short model description here.**
+
+See the [model architecture](#model-architecture) section and [NeMo documentation](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/index.html) for complete architecture details.
+
+
+## NVIDIA NeMo: Training
+
+To train, fine-tune, or experiment with the model, install the PyTorch build for your platform first, then install [NVIDIA NeMo](https://docs.nvidia.com/nemo/speech/nightly/starthere/install.html) with the extras you need.
+```
+pip install 'nemo-toolkit[all]'
+``` 
+
+## How to Use this Model
+
+The model is available for use in the NeMo toolkit [3], and can be used as a pre-trained checkpoint for inference or for fine-tuning on another dataset.
+
+### Automatically instantiate the model
+
+**NOTE**: Please update the model class below to match the class of the model being uploaded.
+
+```python
+import nemo.core import ModelPT
+model = ModelPT.from_pretrained("{repo_id}")
+```
+
+### NOTE
+
+    Add some information about how to use the model here. An example is provided for ASR inference below.
+
+    ### Transcribing using Python
+    First, let's get a sample
+    ```
+    wget https://dldata-public.s3.us-east-2.amazonaws.com/2086-149220-0033.wav
+    ```
+    Then simply do:
+    ```
+    asr_model.transcribe(['2086-149220-0033.wav'])
+    ```
+
+    ### Transcribing many audio files
+
+    ```shell
+    python [NEMO_GIT_FOLDER]/examples/asr/transcribe_speech.py \
+     pretrained_name="{repo_id}" \
+     audio_dir=""
+    ```
+
+### Input
+
+**Add some information about what are the inputs to this model**
+
+### Output
+
+**Add some information about what are the outputs of this model**
+
+## Model Architecture
+
+**Add information here discussing architectural details of the model or any comments to users about the model.**
+
+## Training
+
+**Add information here about how the model was trained. It should be as detailed as possible, potentially including the the link to the script used to train as well as the base config used to train the model. If extraneous scripts are used to prepare the components of the model, please include them here.**
+
+### NOTE
+
+    An example is provided below for ASR
+
+    The NeMo toolkit [3] was used for training the models for over several hundred epochs. These model are trained with this [example script](https://github.com/NVIDIA/NeMo/blob/main/examples/asr/asr_transducer/speech_to_text_rnnt_bpe.py) and this [base config](https://github.com/NVIDIA/NeMo/blob/main/examples/asr/conf/fastconformer/fast-conformer_transducer_bpe.yaml).
+
+    The tokenizers for these models were built using the text transcripts of the train set with this [script](https://github.com/NVIDIA/NeMo/blob/main/scripts/tokenizers/process_asr_text_tokenizer.py).
+
+
+### Datasets
+
+**Try to provide as detailed a list of datasets as possible. If possible, provide links to the datasets on HF by adding it to the manifest section at the top of the README (marked by ---).**
+
+### NOTE
+
+    An example for the manifest section is provided below for ASR datasets
+
+    datasets:
+    - librispeech_asr
+    - fisher_corpus
+    - Switchboard-1
+    - WSJ-0
+    - WSJ-1
+    - National-Singapore-Corpus-Part-1
+    - National-Singapore-Corpus-Part-6
+    - vctk
+    - voxpopuli
+    - europarl
+    - multilingual_librispeech
+    - mozilla-foundation/common_voice_8_0
+    - MLCommons/peoples_speech
+
+    The corresponding text in this section for those datasets is stated below -
+
+    The model was trained on 64K hours of English speech collected and prepared by NVIDIA NeMo and Suno teams.
+
+    The training dataset consists of private subset with 40K hours of English speech plus 24K hours from the following public datasets:
+
+    - Librispeech 960 hours of English speech
+    - Fisher Corpus
+    - Switchboard-1 Dataset
+    - WSJ-0 and WSJ-1
+    - National Speech Corpus (Part 1, Part 6)
+    - VCTK
+    - VoxPopuli (EN)
+    - Europarl-ASR (EN)
+    - Multilingual Librispeech (MLS EN) - 2,000 hour subset
+    - Mozilla Common Voice (v7.0)
+    - People's Speech  - 12,000 hour subset
+
+
+## Performance
+
+**Add information here about the performance of the model. Discuss what is the metric that is being used to evaluate the model and if there are external links explaning the custom metric, please link to it.
+
+### NOTE
+
+    An example is provided below for ASR metrics list that can be added to the top of the README
+    
+    model-index:
+    - name: PUT_MODEL_NAME
+      results:
+      - task:
+          name: Automatic Speech Recognition
+          type: automatic-speech-recognition
+        dataset:
+          name: AMI (Meetings test)
+          type: edinburghcstr/ami
+          config: ihm
+          split: test
+          args:
+            language: en
+        metrics:
+        - name: Test WER
+          type: wer
+          value: 17.10
+      - task:
+          name: Automatic Speech Recognition
+          type: automatic-speech-recognition
+        dataset:
+          name: Earnings-22
+          type: revdotcom/earnings22
+          split: test
+          args:
+            language: en
+        metrics:
+        - name: Test WER
+          type: wer
+          value: 14.11
+
+Provide any caveats about the results presented in the top of the discussion so that nuance is not lost. 
+
+It should ideally be in a tabular format (you can use the following website to make your tables in markdown format - https://www.tablesgenerator.com/markdown_tables)**
+
+## Limitations
+
+**Discuss any practical limitations to the model when being used in real world cases. They can also be legal disclaimers, or discussion regarding the safety of the model (particularly in the case of LLMs).**
+
+
+### Note
+
+    An example is provided below 
+
+    Since this model was trained on publicly available speech datasets, the performance of this model might degrade for speech which includes technical terms, or vernacular that the model has not been trained on. The model might also perform worse for accented speech.
+
+
+## License
+
+License to use this model is covered by the [CC-BY-4.0](https://creativecommons.org/licenses/by/4.0/). By downloading the public and release version of the model, you accept the terms and conditions of the [CC-BY-4.0](https://creativecommons.org/licenses/by/4.0/) license.
+
+## References
+
+**Provide appropriate references in the markdown link format below. Please order them numerically.**
+
+[1] [NVIDIA NeMo Toolkit](https://github.com/NVIDIA/NeMo)
+"""

nemo/core/connectors/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+# Copyright 2015 and onwards Google, Inc.
+#
+# 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/core/connectors/save_restore_connector.py

@@ -0,0 +1,809 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+# Copyright 2015 and onwards Google, Inc.
+#
+# 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 __future__ import annotations  # necessary for lazy types evaluation
+
+import os
+import shutil
+import tarfile
+import tempfile
+import time
+import uuid
+from contextlib import contextmanager, nullcontext
+from pathlib import PurePosixPath
+from typing import Callable, Generator, Optional, Set, Union
+
+import torch
+from lightning.pytorch.trainer.trainer import Trainer
+from omegaconf import DictConfig, OmegaConf
+from omegaconf.omegaconf import open_dict
+
+from nemo.core import classes as nemo_classes  # to avoid circular import do not import ModelPT directly
+from nemo.utils import logging, model_utils
+from nemo.utils.app_state import AppState
+from nemo.utils.file_utils import robust_copy
+from nemo.utils.get_rank import is_global_rank_zero
+from nemo.utils.model_utils import inject_model_parallel_rank
+from nemo.utils.msc_utils import import_multistorageclient, is_multistorageclient_url
+from nemo.utils.tar_utils import is_safe_tar_member, safe_extract
+
+
+class SaveRestoreConnector:
+    """
+    Connector for saving and restoring models.
+    """
+
+    def __init__(self) -> None:
+        self._model_config_yaml = "model_config.yaml"
+        self._model_weights_ckpt = "model_weights.ckpt"
+        self._model_extracted_dir = None
+        self._pack_nemo_file = True
+
+    def save_to(self, model: "nemo_classes.ModelPT", save_path: str):
+        """
+        Saves model instance (weights and configuration) into .nemo file.
+        You can use "restore_from" method to fully restore instance from .nemo file.
+
+        .nemo file is an archive (tar.gz) with the following:
+
+        - model_config.yaml - model configuration in .yaml format.
+          You can deserialize this into cfg argument for model's constructor
+        - model_wights.ckpt - model checkpoint
+
+        Args:
+            model: ModelPT object to be saved.
+            save_path: Path to .nemo file where model instance should be saved
+
+        Returns:
+            str: Path to .nemo file where model instance was saved (same as save_path argument) or None if not rank 0
+                The path can be a directory if the flag `pack_nemo_file` is set to False.
+        """
+
+        if is_global_rank_zero():
+            with tempfile.TemporaryDirectory() as tmpdir:
+                config_yaml = os.path.join(tmpdir, self.model_config_yaml)
+                model_weights = os.path.join(tmpdir, self.model_weights_ckpt)
+                model.to_config_file(path2yaml_file=config_yaml)
+                # update subconfigs, if there are child model, since child model can change its config
+                self._update_subconfigs(model, path2yaml_file=config_yaml)
+                if model.has_native_or_submodules_artifacts():
+                    self._handle_artifacts(model, nemo_file_folder=tmpdir)
+                    # We should not update self._cfg here - the model can still be in use
+                    self._update_artifact_paths(model, path2yaml_file=config_yaml)
+                self._save_state_dict_to_disk(model.state_dict(), model_weights)
+
+                # Check if we are packing the folder into a nemo file
+                if self.pack_nemo_file:
+                    self._make_nemo_file_from_folder(filename=save_path, source_dir=tmpdir)
+                else:
+                    # Get the folder path from the save_path and move all values inside the tmpdir to the folder
+                    folder_path = os.path.dirname(save_path)
+
+                    for file in os.listdir(tmpdir):
+                        shutil.move(os.path.join(tmpdir, file), folder_path)
+        else:
+            return
+
+    def load_config_and_state_dict(
+        self,
+        calling_cls,
+        restore_path: str,
+        override_config_path: Optional[Union[OmegaConf, str]] = None,
+        map_location: Optional[torch.device] = None,
+        strict: bool = True,
+        return_config: bool = False,
+        trainer: Trainer = None,
+        validate_access_integrity: bool = True,
+    ):
+        """
+        Restores model instance (weights and configuration) into .nemo file
+
+        Args:
+            restore_path: path to .nemo file from which model should be instantiated
+            override_config_path: path to a yaml config that will override the internal
+                config file or an OmegaConf / DictConfig object representing the model config.
+            map_location: Optional torch.device() to map the instantiated model to a device.
+                By default (None), it will select a GPU if available, falling back to CPU otherwise.
+            strict: Passed to load_state_dict. By default True
+            return_config: If set to true, will return just the underlying config of the restored
+                model as an OmegaConf DictConfig object without instantiating the model.
+
+        Example:
+            ```
+            model = nemo.collections.asr.models.EncDecCTCModel.restore_from('asr.nemo')
+            assert isinstance(model, nemo.collections.asr.models.EncDecCTCModel)
+            ```
+
+        Returns:
+            An instance of type cls or its underlying config (if return_config is set).
+        """
+        # Get path where the command is executed - the artifacts will be "retrieved" there
+        # (original .nemo behavior)
+        cwd = os.getcwd()
+
+        if map_location is None:
+            if torch.cuda.is_available():
+                map_location = torch.device('cuda')
+            else:
+                map_location = torch.device('cpu')
+
+        app_state = AppState()
+
+        # Determine if we should use a pre-extracted directory
+        use_extracted_dir = self.model_extracted_dir is not None and os.path.isdir(self.model_extracted_dir)
+
+        if use_extracted_dir:
+            logging.info(f"Restoration will occur within pre-extracted directory : " f"`{self.model_extracted_dir}`.")
+
+        # Use nullcontext if we have an extracted dir, otherwise create a temp directory
+        dir_context = nullcontext(self.model_extracted_dir) if use_extracted_dir else tempfile.TemporaryDirectory()
+
+        with dir_context as tmpdir:
+            try:
+                if not use_extracted_dir:
+                    # Extract the nemo file into the temporary directory
+                    filter_fn = None
+                    if return_config:
+                        filter_fn = lambda name: '.yaml' in name
+                    members = self._filtered_tar_info(restore_path, filter_fn=filter_fn)
+                    self._unpack_nemo_file(path2file=restore_path, out_folder=tmpdir, members=members)
+
+                # Change current working directory to
+                os.chdir(tmpdir)
+                if override_config_path is None:
+                    config_yaml = self.model_config_yaml
+                else:
+                    # can be str path or OmegaConf / DictConfig object
+                    config_yaml = override_config_path
+                if not isinstance(config_yaml, (OmegaConf, DictConfig)):
+                    conf = OmegaConf.load(config_yaml)
+                else:
+                    conf = config_yaml
+                    if override_config_path is not None:
+                        # Resolve the override config
+                        conf = OmegaConf.to_container(conf, resolve=True)
+                        conf = OmegaConf.create(conf)
+                # If override is top level config, extract just `model` from it
+                if 'model' in conf:
+                    conf = conf.model
+
+                if return_config:
+                    instance = conf
+                    return instance
+                else:
+                    if app_state.model_parallel_size is not None and app_state.model_parallel_size > 1:
+                        model_weights = self._inject_model_parallel_rank_for_ckpt(tmpdir, self.model_weights_ckpt)
+                    else:
+                        model_weights = os.path.join(tmpdir, self.model_weights_ckpt)
+                OmegaConf.set_struct(conf, True)
+                os.chdir(cwd)
+                # get the class
+                calling_cls._set_model_restore_state(is_being_restored=True, folder=tmpdir)
+                instance = calling_cls.from_config_dict(config=conf, trainer=trainer)
+                instance = instance.to(map_location)
+                # add load_state_dict override
+                if app_state.model_parallel_size is not None and app_state.model_parallel_size > 1:
+                    model_weights = self._inject_model_parallel_rank_for_ckpt(tmpdir, self.model_weights_ckpt)
+                state_dict = self._load_state_dict_from_disk(model_weights, map_location=map_location)
+            finally:
+                os.chdir(cwd)
+
+        return (conf, instance, state_dict)
+
+    def modify_state_dict(self, conf, state_dict):
+        """
+        Utility method that allows to modify the state dict before loading parameters into a model.
+        Args:
+            conf: A model level OmegaConf object.
+            state_dict: The state dict restored from the checkpoint.
+        Returns:
+            A potentially modified state dict.
+        """
+
+        # NOTE and TODO (sandeepsub) This is duplicated across save_restore_connector and nlp_save_restore_connector.
+        # This shouldn't be here.
+        if conf.get('megatron_amp_O2', False):
+            new_state_dict = {}
+            for key in state_dict.keys():
+                new_key = key.replace('model.', 'model.module.', 1)
+                new_state_dict[new_key] = state_dict[key]
+            state_dict = new_state_dict
+        return state_dict
+
+    def load_instance_with_state_dict(self, instance, state_dict, strict):
+        """
+        Utility method that loads a model instance with the (potentially modified) state dict.
+
+        Args:
+            instance: ModelPT subclass instance.
+            state_dict: The state dict (which may have been modified)
+            strict: Bool, whether to perform strict checks when loading the state dict.
+        """
+        instance.load_state_dict(state_dict, strict=strict)
+        instance._set_model_restore_state(is_being_restored=False)
+
+    def restore_from(
+        self,
+        calling_cls,
+        restore_path: str,
+        override_config_path: Optional[Union[OmegaConf, str]] = None,
+        map_location: Optional[torch.device] = None,
+        strict: bool = True,
+        return_config: bool = False,
+        trainer: Trainer = None,
+        validate_access_integrity: bool = True,
+    ):
+        """
+        Restores model instance (weights and configuration) into .nemo file
+
+        Args:
+            restore_path: path to .nemo file from which model should be instantiated
+            override_config_path: path to a yaml config that will override the internal
+                config file or an OmegaConf / DictConfig object representing the model config.
+            map_location: Optional torch.device() to map the instantiated model to a device.
+                By default (None), it will select a GPU if available, falling back to CPU otherwise.
+            strict: Passed to load_state_dict. By default True
+            return_config: If set to true, will return just the underlying config of the restored
+                model as an OmegaConf DictConfig object without instantiating the model.
+            trainer: An optional Trainer object, passed to the model constructor.
+
+        Example:
+            ```
+            model = nemo.collections.asr.models.EncDecCTCModel.restore_from('asr.nemo')
+            assert isinstance(model, nemo.collections.asr.models.EncDecCTCModel)
+            ```
+
+        Returns:
+            An instance of type cls or its underlying config (if return_config is set).
+        """
+        # Get path where the command is executed - the artifacts will be "retrieved" there
+        # (original .nemo behavior)
+        loaded_params = self.load_config_and_state_dict(
+            calling_cls,
+            restore_path,
+            override_config_path,
+            map_location,
+            strict,
+            return_config,
+            trainer,
+            validate_access_integrity,
+        )
+        if not isinstance(loaded_params, tuple) or return_config is True:
+            return loaded_params
+        conf, instance, state_dict = loaded_params
+        state_dict = self.modify_state_dict(conf, state_dict)
+        self.load_instance_with_state_dict(instance, state_dict, strict)
+        logging.info(f'Model {instance.__class__.__name__} was successfully restored from {restore_path}.')
+        return instance
+
+    def extract_state_dict_from(self, restore_path: str, save_dir: str, split_by_module: bool = False):
+        """
+        Extract the state dict(s) from a provided .nemo tarfile and save it to a directory.
+
+        Args:
+            restore_path: path to .nemo file from which state dict(s) should be extracted
+            save_dir: directory in which the saved state dict(s) should be stored
+            split_by_module: bool flag, which determins whether the output checkpoint should
+                be for the entire Model, or the individual module's that comprise the Model
+
+        Example:
+            To convert the .nemo tarfile into a single Model level PyTorch checkpoint::
+
+                state_dict = nemo.collections.asr.models.EncDecCTCModel.extract_state_dict_from('asr.nemo', './ckpts')
+
+            To restore a model from a Model level checkpoint::
+
+                model = nemo.collections.asr.models.EncDecCTCModel(cfg)  # or any other method of restoration
+                model.load_state_dict(torch.load("./ckpts/model_weights.ckpt"))
+
+            To convert the .nemo tarfile into multiple Module level PyTorch checkpoints::
+
+                state_dict = nemo.collections.asr.models.EncDecCTCModel.extract_state_dict_from(
+                    'asr.nemo', './ckpts', split_by_module=True
+                )
+
+            To restore a module from a Module level checkpoint::
+
+                model = nemo.collections.asr.models.EncDecCTCModel(cfg)  # or any other method of restoration
+
+                # load the individual components
+                model.preprocessor.load_state_dict(torch.load("./ckpts/preprocessor.ckpt"))
+                model.encoder.load_state_dict(torch.load("./ckpts/encoder.ckpt"))
+                model.decoder.load_state_dict(torch.load("./ckpts/decoder.ckpt"))
+
+        Returns:
+            The state dict that was loaded from the original .nemo checkpoint
+        """
+
+        cwd = os.getcwd()
+
+        save_dir = os.path.abspath(save_dir)
+        if not os.path.exists(save_dir):
+            os.makedirs(save_dir, exist_ok=True)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            try:
+                self._unpack_nemo_file(path2file=restore_path, out_folder=tmpdir)
+                os.chdir(tmpdir)
+                model_weights = os.path.join(tmpdir, self.model_weights_ckpt)
+                state_dict = self._load_state_dict_from_disk(model_weights)
+
+                if not split_by_module:
+                    filepath = os.path.join(save_dir, self.model_weights_ckpt)
+                    self._save_state_dict_to_disk(state_dict, filepath)
+
+                else:
+                    key_set = set([key.split(".")[0] for key in state_dict.keys()])
+                    for primary_key in key_set:
+                        inner_keys = [key for key in state_dict.keys() if key.split(".")[0] == primary_key]
+                        state_dict_subset = {
+                            ".".join(inner_key.split(".")[1:]): state_dict[inner_key] for inner_key in inner_keys
+                        }
+                        filepath = os.path.join(save_dir, f"{primary_key}.ckpt")
+                        self._save_state_dict_to_disk(state_dict_subset, filepath)
+
+                logging.info(f'Checkpoints from {restore_path} were successfully extracted into {save_dir}.')
+            finally:
+                os.chdir(cwd)
+
+        return state_dict
+
+    def register_artifact(self, model, config_path: str, src: str, verify_src_exists: bool = True):
+        """
+        Register model artifacts with this function. These artifacts (files) will be included inside .nemo file
+        when model.save_to("mymodel.nemo") is called.
+
+        How it works:
+
+        1. It always returns existing absolute path which can be used during Model constructor call
+            EXCEPTION: src is None or "" in which case nothing will be done and src will be returned
+        2. It will add (config_path, model_utils.ArtifactItem()) pair to self.artifacts
+
+            .. code-block::
+
+              If "src" is local existing path:
+                  then it will be returned in absolute path form
+              elif "src" starts with "nemo_file:unique_artifact_name":
+                  .nemo will be untarred to a temporary folder location and an actual existing path will be returned
+              else:
+                  an error will be raised.
+
+        WARNING: use .register_artifact calls in your models' constructors.
+        The returned path is not guaranteed to exist after you have exited your model's constructor.
+
+        Args:
+            model: ModelPT object to register artifact for.
+            config_path (str): Artifact key. Usually corresponds to the model config.
+            src (str): Path to artifact.
+            verify_src_exists (bool): If set to False, then the artifact is optional and register_artifact will return
+                None even if src is not found. Defaults to True.
+
+        Returns:
+            str: If src is not None or empty it always returns absolute path which is guaranteed to exists during model
+                instance life
+        """
+        app_state = AppState()
+
+        artifact_item = model_utils.ArtifactItem()
+
+        # This is for backward compatibility, if the src objects exists simply inside of the tarfile
+        # without its key having been overriden, this pathway will be used.
+        src_obj_name = os.path.basename(src)
+        if app_state.nemo_file_folder is not None:
+            src_obj_path = os.path.abspath(os.path.join(app_state.nemo_file_folder, src_obj_name))
+        else:
+            src_obj_path = src_obj_name
+
+        # src is a local existing path - register artifact and return exact same path for usage by the model
+        if os.path.exists(os.path.abspath(src)):
+            return_path = os.path.abspath(src)
+            artifact_item.path_type = model_utils.ArtifactPathType.LOCAL_PATH
+
+        # this is the case when artifact must be retried from the nemo file
+        # we are assuming that the location of the right nemo file is available from _MODEL_RESTORE_PATH
+        elif src.startswith("nemo:"):
+            return_path = os.path.abspath(os.path.join(app_state.nemo_file_folder, src[5:]))
+            artifact_item.path_type = model_utils.ArtifactPathType.TAR_PATH
+
+        # backward compatibility implementation
+        elif os.path.exists(src_obj_path):
+            return_path = src_obj_path
+            artifact_item.path_type = model_utils.ArtifactPathType.TAR_PATH
+        else:
+            if verify_src_exists:
+                raise FileNotFoundError(
+                    f"src path does not exist or it is not a path in nemo file. "
+                    f"src value I got was: {src}. Absolute: {os.path.abspath(src)}"
+                )
+            else:
+                # artifact is optional and we simply return None
+                logging.warning(
+                    f"src path does not exist or it is not a path in nemo file. "
+                    f"src value I got was: {src}. Absolute: {os.path.abspath(src)}"
+                )
+                return None
+
+        if not os.path.exists(return_path):
+            nemo_folder = app_state.nemo_file_folder
+            existing_files = os.listdir(nemo_folder) if nemo_folder and os.path.isdir(nemo_folder) else []
+            raise FileNotFoundError(
+                f"Artifact not found at expected path: {return_path}\n"
+                f"  src: {src}\n"
+                f"  nemo_file_folder: {nemo_folder}\n"
+                f"  Files in nemo_file_folder: {existing_files}"
+            )
+
+        artifact_item.path = os.path.abspath(src)
+        model.artifacts[config_path] = artifact_item
+        # we were called by ModelPT
+        if hasattr(model, "cfg"):
+            with open_dict(model._cfg):
+                OmegaConf.update(model.cfg, config_path, return_path)
+        return return_path
+
+    def _handle_artifacts(self, model, nemo_file_folder):
+        tarfile_artifacts = []
+        app_state = AppState()
+
+        # aggregate artifacts from self and all children recursively
+        artifacts_containers = []
+        for _, config_path, module in model.named_nemo_modules():
+            if module.has_artifacts():  # NeMo model with artifacts
+                artifacts_containers.append((config_path, module.artifacts))
+
+        if len(artifacts_containers) > 0 and (not hasattr(model, "artifacts") or model.artifacts is None):
+            # model has no artifacts, but submodules have some
+            model.artifacts = dict()
+        for config_path, artifacts in artifacts_containers:
+            for subconf_path, artiitem in artifacts.items():
+                conf_path = f"{config_path}.{subconf_path}" if config_path else f"{subconf_path}"
+                if artiitem.path_type == model_utils.ArtifactPathType.LOCAL_PATH:
+                    if not os.path.exists(artiitem.path):
+                        raise FileNotFoundError(f"Artifact {conf_path} not found at location: {artiitem.path}")
+
+                    # Generate new uniq artifact name and copy it to nemo_file_folder
+                    # Note uuid.uuid4().hex is guaranteed to be 32 character long
+                    artifact_base_name = os.path.basename(artiitem.path)
+                    artifact_uniq_name = f"{uuid.uuid4().hex}_{artifact_base_name}"
+                    robust_copy(artiitem.path, os.path.join(nemo_file_folder, artifact_uniq_name))
+
+                    # Update artifacts registry
+                    artiitem.hashed_path = "nemo:" + artifact_uniq_name
+                    model.artifacts[conf_path] = artiitem
+
+                elif artiitem.path_type == model_utils.ArtifactPathType.TAR_PATH:
+                    # process all tarfile artifacts in one go, so preserve key-value pair
+                    tarfile_artifacts.append((conf_path, artiitem))
+                    if subconf_path:  # artifact from submodule
+                        model.artifacts[conf_path] = artiitem
+
+                else:
+                    raise ValueError("Directly referencing artifacts from other nemo files isn't supported yet")
+
+        # Process current tarfile artifacts by unpacking the previous tarfile and extract the artifacts
+        # that are currently required.
+        # artifacts can be native (from the model itself) and from submodules
+        restoration_paths: Set[str] = set()  # model + submodules restoration paths, handle only unique paths
+        model_metadata = app_state.get_model_metadata_from_guid(model.model_guid)
+        if model_metadata.restoration_path is not None:
+            restoration_paths.add(model_metadata.restoration_path)
+        # aggregate restoration paths for all submodules recursively
+        for module in model.modules():
+            if isinstance(module, nemo_classes.ModelPT):  # if NeMo model
+                submodule_restoration_path = app_state.get_model_metadata_from_guid(module.model_guid).restoration_path
+                if submodule_restoration_path is not None:
+                    restoration_paths.add(submodule_restoration_path)
+        if len(tarfile_artifacts) > 0 and len(restoration_paths) == 0:
+            # TODO: see cases when this can occur, and if we can fix them
+            logging.warning("Model contains registered artifacts, but no restoration paths found")
+        if len(tarfile_artifacts) > 0 and len(restoration_paths) > 0:
+
+            def check_artifact_and_query_basename_match(query_path: str) -> bool:
+                for _, artiitem in tarfile_artifacts:
+                    # Get basename and copy it to nemo_file_folder
+                    if 'nemo:' in artiitem.path:
+                        artifact_base_name = artiitem.path.split('nemo:')[1]
+                    else:
+                        artifact_base_name = os.path.basename(artiitem.path)
+
+                    if artifact_base_name == os.path.basename(query_path):
+                        return True
+                return False
+
+            artifact_rel_paths = {}
+            for path in restoration_paths:
+                if self.model_extracted_dir:
+                    artifact_rel_paths[path] = self._filtered_recursive_walk(
+                        path, filter_fn=check_artifact_and_query_basename_match
+                    )
+                else:
+                    artifact_rel_paths[path] = self._filtered_tar_info(
+                        path, filter_fn=check_artifact_and_query_basename_match
+                    )
+            # Need to step into nemo archive to extract file
+            # Get path where the command is executed - the artifacts will be "retrieved" there
+            # (original .nemo behavior)
+            cwd = os.getcwd()
+            # Step into the nemo archive to try and find the file
+            # TemporaryDirectory context must always be outer to try-catch chdir otherwise it crashes on Windows
+            with tempfile.TemporaryDirectory() as archive_dir:
+                try:
+                    # unpack artifacts from all restorations paths (nemo checkpoints)
+                    # in nemo checkpoints all resources contain hash in name, so there should be no collisions
+                    for path in restoration_paths:
+                        if self.model_extracted_dir:
+                            for rel_path in artifact_rel_paths[path]:
+                                robust_copy(rel_path, archive_dir)
+                        else:
+                            self._unpack_nemo_file(
+                                path2file=path, out_folder=archive_dir, members=artifact_rel_paths[path]
+                            )
+                    os.chdir(archive_dir)
+                    for conf_path, artiitem in tarfile_artifacts:
+                        # Get basename and copy it to nemo_file_folder
+                        if 'nemo:' in artiitem.path:
+                            artifact_base_name = artiitem.path.split('nemo:')[1]
+                        else:
+                            artifact_base_name = os.path.basename(artiitem.path)
+                        # no need to hash here as we are in tarfile_artifacts which are already hashed
+                        artifact_uniq_name = artifact_base_name
+                        robust_copy(artifact_base_name, os.path.join(nemo_file_folder, artifact_base_name))
+
+                        # Update artifacts registry
+                        new_artiitem = model_utils.ArtifactItem()
+                        new_artiitem.path = "nemo:" + artifact_uniq_name
+                        new_artiitem.path_type = model_utils.ArtifactPathType.TAR_PATH
+                        model.artifacts[conf_path] = new_artiitem
+                finally:
+                    # change back working directory
+                    os.chdir(cwd)
+
+    @staticmethod
+    def _update_subconfigs(model: "nemo_classes.ModelPT", path2yaml_file):
+        """
+        Update subconfigs of the model if ModelPT has submodules
+        Should be called before updating artifacts paths
+        """
+        if not model.has_nemo_submodules():
+            # no submodules => nothing to update
+            return
+        conf = OmegaConf.load(path2yaml_file)
+        # update subconfigs for all children recoursively
+        # parent configs updated before children
+        for _, conf_path, submodule in model.named_nemo_modules():
+            if not conf_path:  # self
+                continue
+            OmegaConf.update(conf, conf_path, submodule.cfg)
+        with open(path2yaml_file, 'w', encoding='utf-8') as fout:
+            OmegaConf.save(config=conf, f=fout, resolve=True)
+
+    def _update_artifact_paths(self, model, path2yaml_file):
+        if hasattr(model, "artifacts") and model.artifacts is not None and len(model.artifacts) > 0:
+            conf = OmegaConf.load(path2yaml_file)
+            for conf_path, item in model.artifacts.items():
+                if item.hashed_path is None:
+                    OmegaConf.update(conf, conf_path, item.path)
+                else:
+                    OmegaConf.update(conf, conf_path, item.hashed_path)
+            with open(path2yaml_file, 'w', encoding='utf-8') as fout:
+                OmegaConf.save(config=conf, f=fout, resolve=True)
+
+    def _inject_model_parallel_rank_for_ckpt(self, dirname, basename):
+        model_weights = os.path.join(dirname, basename)
+        model_weights = inject_model_parallel_rank(model_weights)
+        return model_weights
+
+    @staticmethod
+    def _make_nemo_file_from_folder(filename, source_dir):
+        if is_multistorageclient_url(filename):
+            SaveRestoreConnector._make_nemo_file_from_folder_with_multistorageclient(filename, source_dir)
+        else:
+            dirname = os.path.dirname(filename)
+            os.makedirs(dirname, exist_ok=True)
+            with tarfile.open(filename, "w:") as tar:
+                tar.add(source_dir, arcname=".")
+
+    @staticmethod
+    def _make_nemo_file_from_folder_with_multistorageclient(filename, source_dir):
+        msc = import_multistorageclient()
+        # Use PurePosixPath for cloud storage paths which always use forward slashes
+        filename_with_extension = PurePosixPath(filename).name
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tar_file = os.path.join(tmpdir, filename_with_extension)
+            with tarfile.open(tar_file, "w:") as tar:
+                tar.add(source_dir, arcname=".")
+                start_time = time.time()
+                msc.upload_file(filename, tar_file)
+                logging.debug(
+                    f"Time spent for msc.upload_file from {tar_file} to {filename}: {time.time() - start_time:.4f}"
+                )
+
+    @staticmethod
+    def _is_safe_path(member, extract_to):
+        return is_safe_tar_member(member, os.path.realpath(extract_to))
+
+    @staticmethod
+    def _safe_extract(tar, out_folder: str, members=None):
+        return safe_extract(tar, out_folder, members=members, skip_unsafe=True)
+
+    @staticmethod
+    def _filtered_tar_info(tar_path: str, filter_fn: Optional[Callable[[str], bool]] = None) -> list[tarfile.TarInfo]:
+        """
+        Returns the members of the tarball filtered by a function
+        """
+        with SaveRestoreConnector._tar_open(tar_path) as tar:
+            members = tar.getmembers()
+            if filter_fn is None:
+                return members
+
+            return [x for x in members if filter_fn(x.name)]
+
+    @staticmethod
+    def _filtered_recursive_walk(path: str, filter_fn: Optional[Callable[[str], bool]] = None) -> list[str]:
+        """
+        Returns the result of recursive walking a path and filtering each element
+        """
+        if not os.path.isdir(path):
+            raise NotADirectoryError(f"Expected {path=} to be a directory")
+
+        filtered_rel_paths = []
+        for root, _, files in os.walk(path):
+            for f in files:
+                full_rel_path = os.path.join(root, f)
+                if filter_fn is None or filter_fn(full_rel_path):
+                    filtered_rel_paths.append(full_rel_path)
+        return filtered_rel_paths
+
+    @staticmethod
+    @contextmanager
+    def _tar_open(path2file: str) -> Generator[tarfile.TarFile, None, None]:
+        if not os.path.exists(path2file):
+            raise FileNotFoundError(f"{path2file} does not exist")
+
+        # we start with an assumption of uncompressed tar,
+        # which should be true for versions 1.7.0 and above
+        tar_header = "r:"
+        try:
+            tar_test = tarfile.open(path2file, tar_header)
+            tar_test.close()
+        except tarfile.ReadError:
+            # can be older checkpoint => try compressed tar
+            tar_header = "r:gz"
+
+        tar = tarfile.open(path2file, tar_header)
+        try:
+            yield tar
+        finally:
+            tar.close()
+
+    @staticmethod
+    def _unpack_nemo_file(path2file: str, out_folder: str, members: Optional[list[str]] = None) -> str:
+        """
+        Unpack a nemo file.
+        """
+        if is_multistorageclient_url(path2file):
+            out_folder = SaveRestoreConnector._unpack_nemo_file_with_multistorageclient(path2file, out_folder, members)
+        else:
+            with SaveRestoreConnector._tar_open(path2file) as tar:
+                if members is None:
+                    SaveRestoreConnector._safe_extract(tar, out_folder)
+                else:
+                    SaveRestoreConnector._safe_extract(tar, out_folder, members)
+        return out_folder
+
+    @staticmethod
+    def _unpack_nemo_file_with_multistorageclient(
+        path2file: str, out_folder: str, members: Optional[list[str]] = None
+    ) -> str:
+        """
+        Unpack a nemo file with multistorageclient.
+        """
+        msc = import_multistorageclient()
+        if not msc.os.path.exists(path2file):
+            raise FileNotFoundError(f"{path2file} does not exist")
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Use PurePosixPath for cloud storage paths which always use forward slashes
+            filename_with_extension = PurePosixPath(path2file).name
+            downloaded_file_path = os.path.join(tmpdir, filename_with_extension)
+            start_time = time.time()
+            msc.download_file(path2file, downloaded_file_path)
+            logging.info(
+                f"Time spent for msc.download_file from {downloaded_file_path}: {time.time() - start_time:.4f}"
+            )
+
+            # we start with an assumption of uncompressed tar,
+            # which should be true for versions 1.7.0 and above
+            tar_header = "r:"
+            try:
+                tar_test = tarfile.open(downloaded_file_path, tar_header)
+                tar_test.close()
+            except tarfile.ReadError:
+                # can be older checkpoint => try compressed tar
+                tar_header = "r:gz"
+            tar = tarfile.open(downloaded_file_path, tar_header)
+            if members is None:
+                SaveRestoreConnector._safe_extract(tar, out_folder)
+            else:
+                SaveRestoreConnector._safe_extract(tar, out_folder, members)
+            tar.close()
+        return out_folder
+
+    @staticmethod
+    def _save_state_dict_to_disk(state_dict, filepath):
+        torch.save(state_dict, filepath)
+
+    @staticmethod
+    def _load_state_dict_from_disk(model_weights, map_location='cpu'):
+        """
+        Load model state dict from disk.
+
+        Args:
+            model_weights: Path to the checkpoint file
+            map_location: Device to map tensors to
+
+        Returns:
+            State dict loaded from checkpoint
+
+        """
+        try:
+            # Use torch's default weights_only handling so TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD is honored.
+            return torch.load(model_weights, map_location=map_location)
+        except Exception as e:
+            logging.error(f"Failed to load checkpoint: {e}")
+            raise e
+
+    @property
+    def model_config_yaml(self) -> str:
+        """
+        Get the path to the model config yaml file.
+        """
+        return self._model_config_yaml
+
+    @model_config_yaml.setter
+    def model_config_yaml(self, path: str):
+        self._model_config_yaml = path
+
+    @property
+    def model_weights_ckpt(self) -> str:
+        """
+        Get the path to the model weights checkpoint file.
+        """
+        return self._model_weights_ckpt
+
+    @model_weights_ckpt.setter
+    def model_weights_ckpt(self, path: str):
+        self._model_weights_ckpt = path
+
+    @property
+    def model_extracted_dir(self) -> Optional[str]:
+        """
+        Get the path to the model extracted directory.
+        """
+        return self._model_extracted_dir
+
+    @model_extracted_dir.setter
+    def model_extracted_dir(self, path: Optional[str]):
+        self._model_extracted_dir = path
+
+    @property
+    def pack_nemo_file(self) -> bool:
+        """
+        Get the flag for packing a nemo file.
+        """
+        return self._pack_nemo_file
+
+    @pack_nemo_file.setter
+    def pack_nemo_file(self, save_nemo_file: bool):
+        self._pack_nemo_file = save_nemo_file

nemo/core/neural_types/__init__.py

@@ -0,0 +1,19 @@
+# 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.core.neural_types.axes import *
+from nemo.core.neural_types.comparison import *
+from nemo.core.neural_types.elements import *
+from nemo.core.neural_types.neural_type import *

nemo/core/neural_types/axes.py

@@ -0,0 +1,108 @@
+# 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 enum import Enum
+from typing import Optional
+
+__all__ = ['AxisKindAbstract', 'AxisKind', 'AxisType']
+
+
+class AxisKindAbstract(Enum):
+    """This is an abstract Enum to represents what does varying axis dimension mean.
+    In practice, you will almost always use AxisKind Enum. This Enum should be inherited by
+    your OWN Enum if you aren't satisfied with AxisKind. Then your own Enum can be used
+    instead of AxisKind."""
+
+    pass
+
+
+class AxisKind(AxisKindAbstract):
+    """This Enum represents what does varying axis dimension mean.
+    For example, does this dimension correspond to width, batch, time, etc.
+    The "Dimension" and "Channel" kinds are the same and used to represent
+    a general axis. "Any" axis will accept any axis kind fed to it.
+    """
+
+    Batch = 0
+    Time = 1
+    Dimension = 2
+    Channel = 2
+    Width = 3
+    Height = 4
+    Any = 5
+    Sequence = 6
+    FlowGroup = 7
+    Singleton = 8  # Used to represent a axis that has size 1
+
+    def __repr__(self):
+        return self.__str__()
+
+    def __str__(self):
+        return str(self.name).lower()
+
+    def t_with_string(self, text):
+        """Check whether a dynamic time-axis label matches this axis kind."""
+        # it checks if text is "t_<any string>"
+        return text.startswith("t_") and text.endswith("_") and text[2:-1] == self.__str__()
+
+    @staticmethod
+    def from_str(label):
+        """Returns AxisKind instance based on short string representation"""
+        _label = label.lower().strip()
+        if _label == "b" or _label == "n" or _label == "batch":
+            return AxisKind.Batch
+        elif _label == "t" or _label == "time" or (len(_label) > 2 and _label.startswith("t_")):
+            return AxisKind.Time
+        elif _label == "d" or _label == "c" or _label == "channel":
+            return AxisKind.Dimension
+        elif _label == "w" or _label == "width":
+            return AxisKind.Width
+        elif _label == "h" or _label == "height":
+            return AxisKind.Height
+        elif _label == "s" or _label == "singleton":
+            return AxisKind.Singleton
+        elif _label == "seq" or _label == "sequence":
+            return AxisKind.Sequence
+        elif _label == "flowgroup":
+            return AxisKind.FlowGroup
+        elif _label == "any":
+            return AxisKind.Any
+        else:
+            raise ValueError(f"Can't create AxisKind from {label}")
+
+
+class AxisType(object):
+    """This class represents axis semantics and (optionally) it's dimensionality
+    Args:
+        kind (AxisKindAbstract): what kind of axis it is? For example Batch, Height, etc.
+        size (int, optional): specify if the axis should have a fixed size. By default it is set to None and you
+        typically do not want to set it for Batch and Time
+        is_list (bool, default=False): whether this is a list or a tensor axis
+    """
+
+    def __init__(self, kind: AxisKindAbstract, size: Optional[int] = None, is_list=False):
+        if size is not None and is_list:
+            raise ValueError("The axis can't be list and have a fixed size")
+        self.kind = kind
+        self.size = size
+        self.is_list = is_list
+
+    def __repr__(self):
+        if self.size is None:
+            representation = str(self.kind)
+        else:
+            representation = f"{str(self.kind)}:{self.size}"
+        if self.is_list:
+            representation += "_listdim"
+        return representation

nemo/core/neural_types/comparison.py

@@ -0,0 +1,32 @@
+# 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 enum import Enum
+
+__all__ = ['NeuralTypeComparisonResult']
+
+
+class NeuralTypeComparisonResult(Enum):
+    """The result of comparing two neural type objects for compatibility.
+    When comparing A.compare_to(B):"""
+
+    SAME = 0
+    LESS = 1  # A is B
+    GREATER = 2  # B is A
+    DIM_INCOMPATIBLE = 3  # Resize connector might fix incompatibility
+    TRANSPOSE_SAME = 4  # A transpose and/or converting between lists and tensors will make them same
+    CONTAINER_SIZE_MISMATCH = 5  # A and B contain different number of elements
+    INCOMPATIBLE = 6  # A and B are incompatible
+    SAME_TYPE_INCOMPATIBLE_PARAMS = 7  # A and B are of the same type but parametrized differently
+    UNCHECKED = 8  # type comparison wasn't done

nemo/core/neural_types/elements.py

@@ -0,0 +1,589 @@
+# 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 abc import ABC
+from typing import Any, Dict, Optional
+
+import torch
+
+from nemo.core.neural_types.comparison import NeuralTypeComparisonResult
+
+__all__ = [
+    'ElementType',
+    'VoidType',
+    'BoolType',
+    'ChannelType',
+    'AcousticEncodedRepresentation',
+    'AudioSignal',
+    'VideoSignal',
+    'SpectrogramType',
+    'MelSpectrogramType',
+    'MFCCSpectrogramType',
+    'LogitsType',
+    'LabelsType',
+    'HypothesisType',
+    'LossType',
+    'RegressionValuesType',
+    'CategoricalValuesType',
+    'PredictionsType',
+    'LogprobsType',
+    'ProbsType',
+    'LengthsType',
+    'EmbeddedTextType',
+    'EncodedRepresentation',
+    'MaskType',
+    'Target',
+    'ClassificationTarget',
+    'ImageFeatureValue',
+    'Index',
+    'ImageValue',
+    'NormalizedImageValue',
+    'StringLabel',
+    'StringType',
+    'TokenIndex',
+    'Length',
+    'IntType',
+    'FloatType',
+    'BoolType',
+    'NormalDistributionSamplesType',
+    'NormalDistributionMeanType',
+    'NormalDistributionLogVarianceType',
+    'TokenDurationType',
+    'TokenLogDurationType',
+    'LogDeterminantType',
+    'SequenceToSequenceAlignmentType',
+]
+
+
+class ElementType(ABC):
+    """Abstract class defining semantics of the tensor elements.
+    We are relying on Python for inheritance checking"""
+
+    def __str__(self):
+        if torch.jit.is_scripting():
+            return "SuppressedForTorchScript"
+        return self.__doc__
+
+    def __repr__(self):
+        if torch.jit.is_scripting():
+            return "SuppressedForTorchScript"
+        return self.__class__.__name__
+
+    @property
+    def type_parameters(self) -> Dict[str, Any]:
+        """Override this property to parametrize your type. For example, you can specify 'storage' type such as
+        float, int, bool with 'dtype' keyword. Another example, is if you want to represent a signal with a
+        particular property (say, sample frequency), then you can put sample_freq->value in there.
+        When two types are compared their type_parameters must match."""
+        return {}
+
+    @property
+    def fields(self):
+        """This should be used to logically represent tuples/structures. For example, if you want to represent a
+        bounding box (x, y, width, height) you can put a tuple with names ('x', y', 'w', 'h') in here.
+        Under the hood this should be converted to the last tesnor dimension of fixed size = len(fields).
+        When two types are compared their fields must match."""
+        return None
+
+    def compare(self, second) -> NeuralTypeComparisonResult:
+        if torch.jit.is_scripting():
+            # Neural types for TorchScript are suppressed
+            # This is a stub to make TorchScript happy
+            return NeuralTypeComparisonResult.SAME
+        # First, check general compatibility
+        first_t = type(self)
+        second_t = type(second)
+
+        if first_t == second_t:
+            result = NeuralTypeComparisonResult.SAME
+        elif issubclass(first_t, second_t):
+            result = NeuralTypeComparisonResult.LESS
+        elif issubclass(second_t, first_t):
+            result = NeuralTypeComparisonResult.GREATER
+        else:
+            result = NeuralTypeComparisonResult.INCOMPATIBLE
+
+        if result != NeuralTypeComparisonResult.SAME:
+            return result
+        else:
+            # now check that all parameters match
+            check_params = set(self.type_parameters.keys()) == set(second.type_parameters.keys())
+            if check_params is False:
+                return NeuralTypeComparisonResult.SAME_TYPE_INCOMPATIBLE_PARAMS
+            else:
+                for k1, v1 in self.type_parameters.items():
+                    if v1 is None or second.type_parameters[k1] is None:
+                        # Treat None as Void
+                        continue
+                    if v1 != second.type_parameters[k1]:
+                        return NeuralTypeComparisonResult.SAME_TYPE_INCOMPATIBLE_PARAMS
+            # check that all fields match
+            if self.fields == second.fields:
+                return NeuralTypeComparisonResult.SAME
+            else:
+                return NeuralTypeComparisonResult.INCOMPATIBLE
+
+
+class VoidType(ElementType):
+    """Void-like type which is compatible with everything.
+    It is a good practice to use this type only as necessary.
+    For example, when you need template-like functionality.
+    """
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            # see details: https://github.com/pytorch/pytorch/issues/42885
+            super().__init__()
+
+    def compare(self, second) -> NeuralTypeComparisonResult:
+        return NeuralTypeComparisonResult.SAME
+
+
+# TODO: Consider moving these files elsewhere
+class ChannelType(ElementType):
+    """Element to represent convolutional input/output channel."""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class EmbeddedTextType(ChannelType):
+    """Element to represent output on word/text embedding layers"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class LogitsType(ElementType):
+    """Element type to represent logits"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class ProbsType(ElementType):
+    """Element type to represent probabilities. For example, outputs of softmax layers."""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class LogprobsType(ElementType):
+    """Element type to represent log-probabilities. For example, outputs of log softmax layers."""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class LabelsType(ElementType):
+    """Element type to represent some sort of labels. This is often used as a base class to create
+    a more concrete types such as RegressionValuesType, etc."""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class HypothesisType(LabelsType):
+    """Element type to represent some decoded hypothesis, which may further be processed to obtain
+    a concrete label."""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class LengthsType(ElementType):
+    """Element type representing lengths of something"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class LossType(ElementType):
+    """Element type to represent outputs of Loss modules"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class EncodedRepresentation(ChannelType):
+    """Element type to represent encoded representation, for example, encoder's output"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class AcousticEncodedRepresentation(EncodedRepresentation):
+    """Element type to represent encoded representation returned by the acoustic encoder model"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class AudioSignal(ElementType):
+    """Element type to represent encoded representation returned by the acoustic encoder model
+    Args:
+        freq (int): sampling frequency of a signal. Note that two signals will only be the same if their
+        freq is the same.
+    """
+
+    def __init__(self, freq: Optional[int] = None):
+        self._params: Dict[str, Any] = {}
+        self._params['freq'] = freq
+
+    @property
+    def type_parameters(self):
+        return self._params
+
+
+class VideoSignal(ElementType):
+    """Element type to represent encoded representation returned by the visual encoder model
+    Args:
+        fps (int): frames per second.
+    """
+
+    def __init__(self, fps: Optional[int] = None):
+        self._params: dict[str, Any] = {}
+        self._params['fps'] = fps
+
+    @property
+    def type_parameters(self):
+        return self._params
+
+
+class SpectrogramType(ChannelType):
+    """Element type to represent generic spectrogram signal"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class MelSpectrogramType(SpectrogramType):
+    """Element type to represent mel spectrogram signal"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class MFCCSpectrogramType(SpectrogramType):
+    """Element type to represent MFCC spectrogram signal"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class PredictionsType(LabelsType):
+    """Element type to represent some sort of predictions returned by model"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class RegressionValuesType(PredictionsType):
+    """Element type to represent labels for regression task"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class CategoricalValuesType(PredictionsType):
+    """Element type to represent labels for categorical classification task"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class MaskType(PredictionsType):
+    """Element type to represent a boolean mask"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class Index(ElementType):
+    """Type representing an element being an index of the sample."""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class Target(ElementType):
+    """
+    Type representing an element being a target value.
+    """
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class ClassificationTarget(Target):
+    """
+    Type representing an element being target value in the classification task, i.e. identifier of a desired class.
+    """
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class ImageValue(ElementType):
+    """
+    Type representing an element/value of a single image channel,
+    e.g. a single element (R) of RGB image.
+    """
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class NormalizedImageValue(ImageValue):
+    """
+    Type representing an element/value of a single image channel normalized to <0-1> range,
+    e.g. a single element (R) of normalized RGB image.
+    """
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class ImageFeatureValue(ImageValue):
+    """Type representing an element (single value) of a (image) feature maps."""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class StringType(ElementType):
+    """Element type representing a single string"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class StringLabel(StringType):
+    """
+    Type representing an label being a string with class name (e.g. the "hamster" class in CIFAR100).
+    """
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class BoolType(ElementType):
+    """Element type representing a single integer"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class IntType(ElementType):
+    """Element type representing a single integer"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class FloatType(ElementType):
+    """Element type representing a single float"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class TokenIndex(IntType):
+    """Type representing an element being index of a token in some kind of a vocabulary."""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class Length(IntType):
+    """Type representing an element storing a "length" (e.g. length of a list)."""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class ProbabilityDistributionSamplesType(ElementType):
+    """Element to represent tensors that meant to be sampled from a valid probability distribution"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class NormalDistributionSamplesType(ProbabilityDistributionSamplesType):
+    """Element to represent tensors that meant to be sampled from a valid normal distribution"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class SequenceToSequenceAlignmentType(ElementType):
+    """Class to represent the alignment from seq-to-seq attention outputs. Generally a mapping from endcoder time steps
+    to decoder time steps."""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class NormalDistributionMeanType(ElementType):
+    """Element to represent the mean of a normal distribution"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class NormalDistributionLogVarianceType(ElementType):
+    """Element to represent the log variance of a normal distribution"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class TokenDurationType(ElementType):
+    """Element for representing the duration of a token"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class TokenLogDurationType(ElementType):
+    """Element for representing the log-duration of a token"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()
+
+
+class LogDeterminantType(ElementType):
+    """Element for representing log determinants usually used in flow models"""
+
+    def __init__(self):
+        """Dummy init for TorchScript compatibility"""
+        if not torch.jit.is_scripting():
+            # torch.jit does not support super() call
+            super().__init__()

nemo/core/neural_types/neural_type.py

@@ -0,0 +1,272 @@
+# 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 typing import Any, Optional
+
+import torch
+
+from nemo.core.neural_types.axes import AxisKind, AxisType
+from nemo.core.neural_types.comparison import NeuralTypeComparisonResult
+from nemo.core.neural_types.elements import ElementType, VoidType
+
+__all__ = [
+    'NeuralType',
+    'NeuralTypeError',
+    'NeuralPortNameMismatchError',
+    'NeuralPortNmTensorMismatchError',
+]
+
+
+class NeuralType:
+    """This is the main class which would represent neural type concept.
+    It is used to represent *the types* of inputs and outputs.
+
+    Args:
+        axes (Optional[Tuple]): a tuple of AxisTypes objects representing the semantics of what varying each axis means
+            You can use a short, string-based form here. For example: ('B', 'C', 'H', 'W') would correspond to an NCHW
+            format frequently used in computer vision. ('B', 'T', 'D') is frequently used for signal processing and
+            means [batch, time, dimension/channel].
+        elements_type (ElementType): an instance of ElementType class representing the semantics of what is stored
+            inside the tensor. For example: logits (LogitsType), log probabilities (LogprobType), etc.
+        optional (bool): By default, this is false. If set to True, it would means that input to the port of this
+            type can be optional.
+    """
+
+    def __str__(self):
+        if torch.jit.is_scripting():
+            return "SuppressedForTorchScript"
+        if self.axes is not None:
+            return f"axes: {self.axes}; elements_type: {self.elements_type.__class__.__name__}"
+        else:
+            return f"axes: None; elements_type: {self.elements_type.__class__.__name__}"
+
+    def __init__(self, axes: Optional[Any] = None, elements_type: Optional[Any] = None, optional: bool = False):
+        """
+        Args:
+            axes: a tuple of AxisTypes objects representing the semantics of what varying each axis means
+            elements_type: None or ElementType; we need Any annotation here to avoid problems with TorchScript
+                (it is checked in _init_internal)
+            optional: If input to the port of this type can be optional (False by default).
+        """
+        if not torch.jit.is_scripting():
+            self._init_internal(axes=axes, elements_type=elements_type, optional=optional)
+
+    @torch.jit.unused
+    def _init_internal(
+        self, axes: Optional[Any] = None, elements_type: Optional[ElementType] = None, optional: bool = False
+    ):
+        """Internals of __init__, separated to make TorchScript and autodoc work"""
+        if elements_type is None:
+            elements_type = VoidType()
+        if not isinstance(elements_type, ElementType):
+            raise ValueError(
+                "elements_type of NeuralType must be an instance of a class derived from ElementType. "
+                "Did you pass a class instead?"
+            )
+        self.elements_type = elements_type
+        if axes is not None:
+            NeuralType.__check_sanity(axes)
+            axes_list = []
+            for axis in axes:
+                if isinstance(axis, str):
+                    axes_list.append(AxisType(AxisKind.from_str(axis), None))
+                elif isinstance(axis, AxisType):
+                    axes_list.append(axis)
+                else:
+                    raise ValueError("axis type must be either str or AxisType instance")
+            self.axes = tuple(axes_list)
+        else:
+            self.axes = None
+        self.optional = optional
+
+    def compare(self, second) -> NeuralTypeComparisonResult:
+        """Performs neural type comparison of self with second. When you chain two modules' inputs/outputs via
+        __call__ method, this comparison will be called to ensure neural type compatibility."""
+        if torch.jit.is_scripting():
+            # suppress check for TorchScript
+            return NeuralTypeComparisonResult.SAME
+        # First, handle dimensionality
+        axes_a = self.axes
+        axes_b = second.axes
+
+        # "Big void" type
+        if isinstance(self.elements_type, VoidType) and self.axes is None:
+            return NeuralTypeComparisonResult.SAME
+
+        if self.axes is None:
+            if second.axes is None:
+                return self.elements_type.compare(second.elements_type)
+            else:
+                return NeuralTypeComparisonResult.INCOMPATIBLE
+
+        dimensions_pass = NeuralType.__compare_axes(axes_a, axes_b)
+        element_comparison_result = self.elements_type.compare(second.elements_type)
+
+        # SAME DIMS
+        if dimensions_pass == 0:
+            return element_comparison_result
+        # TRANSPOSE_SAME DIMS
+        elif dimensions_pass == 1:
+            if element_comparison_result == NeuralTypeComparisonResult.SAME:
+                return NeuralTypeComparisonResult.TRANSPOSE_SAME
+            else:
+                return NeuralTypeComparisonResult.INCOMPATIBLE
+        # DIM_INCOMPATIBLE DIMS
+        elif dimensions_pass == 2:
+            if element_comparison_result == NeuralTypeComparisonResult.SAME:
+                return NeuralTypeComparisonResult.DIM_INCOMPATIBLE
+            else:
+                return NeuralTypeComparisonResult.INCOMPATIBLE
+        else:
+            return NeuralTypeComparisonResult.INCOMPATIBLE
+
+    def compare_and_raise_error(self, parent_type_name, port_name, second_object):
+        """Method compares definition of one type with another and raises an error if not compatible."""
+        if torch.jit.is_scripting():
+            # suppress for TorchScript
+            return
+        type_comatibility = self.compare(second_object)
+        if (
+            type_comatibility != NeuralTypeComparisonResult.SAME
+            and type_comatibility != NeuralTypeComparisonResult.GREATER
+        ):
+            raise NeuralPortNmTensorMismatchError(
+                parent_type_name, port_name, str(self), str(second_object.ntype), type_comatibility
+            )
+
+    def __eq__(self, other):
+        if isinstance(other, NeuralType):
+            return self.compare(other)
+
+        return False
+
+    @staticmethod
+    def __check_sanity(axes):
+        # check that list come before any tensor dimension
+        are_strings = True
+        for axis in axes:
+            if not isinstance(axis, str):
+                are_strings = False
+            if isinstance(axis, str) and not are_strings:
+                raise ValueError("Either use full class names or all strings")
+        if are_strings:
+            return
+        checks_passed = True
+        saw_tensor_dim = False
+        for axis in axes:
+            if not axis.is_list:
+                saw_tensor_dim = True
+            else:  # current axis is a list
+                if saw_tensor_dim:  # which is preceded by tensor dim
+                    checks_passed = False
+        if not checks_passed:
+            raise ValueError(
+                "You have list dimension after Tensor dimension. All list dimensions must preceed Tensor dimensions"
+            )
+
+    @staticmethod
+    def __compare_axes(axes_a, axes_b) -> int:
+        """
+        Compares axes_a and axes_b
+        Args:
+            axes_a: first axes tuple
+            axes_b: second axes tuple
+
+        Returns:
+            0 - if they are exactly the same
+            1 - if they are "TRANSPOSE_SAME"
+            2 - if the are "DIM_INCOMPATIBLE"
+            3 - if they are different
+        """
+        if axes_a is None and axes_b is None:
+            return 0
+        elif axes_a is None and axes_b is not None:
+            return 3
+        elif axes_a is not None and axes_b is None:
+            return 3
+        elif len(axes_a) != len(axes_b):
+            return 3
+        # After these ifs we know that len(axes_a) == len(axes_b)
+
+        same = True
+        kinds_a = dict()
+        kinds_b = dict()
+        for axis_a, axis_b in zip(axes_a, axes_b):
+            kinds_a[axis_a.kind] = axis_a.size
+            kinds_b[axis_b.kind] = axis_b.size
+            if axis_a.kind == AxisKind.Any:
+                same = True
+            elif (
+                axis_a.kind != axis_b.kind
+                or axis_a.is_list != axis_b.is_list
+                or (axis_a.size != axis_b.size and axis_a.size is not None)
+            ):
+                same = False
+        if same:
+            return 0
+        else:
+            # can be TRANSPOSE_SAME, DIM_INCOMPATIBLE
+            if kinds_a.keys() == kinds_b.keys():
+                for key, value in kinds_a.items():
+                    if kinds_b[key] != value:
+                        return 2
+                return 1
+            else:
+                return 3
+
+    def __repr__(self):
+        if torch.jit.is_scripting():
+            return "SuppressedForTorchScript"
+
+        if self.axes is not None:
+            axes = str(self.axes)
+        else:
+            axes = "None"
+
+        if self.elements_type is not None:
+            element_type = repr(self.elements_type)
+        else:
+            element_type = "None"
+
+        data = f"axis={axes}, element_type={element_type}"
+
+        if self.optional:
+            data = f"{data}, optional={self.optional}"
+
+        final = f"{self.__class__.__name__}({data})"
+        return final
+
+
+class NeuralTypeError(Exception):
+    """Base class for neural type related exceptions."""
+
+
+class NeuralPortNameMismatchError(NeuralTypeError):
+    """Exception raised when neural module is called with incorrect port
+    names."""
+
+    def __init__(self, input_port_name):
+        super().__init__()
+        self.message = "Wrong input port name: {0}".format(input_port_name)
+
+
+class NeuralPortNmTensorMismatchError(NeuralTypeError):
+    """Exception raised when a port is fed with a NmTensor of incompatible
+    type."""
+
+    def __init__(self, class_name, port_name, first_type, second_type, type_comatibility):
+        super().__init__()
+        self.message = "\nIn {}. \nPort: {} and a NmTensor it was fed are \n".format(class_name, port_name)
+        self.message += "of incompatible neural types:\n\n{} \n\n and \n\n{}".format(first_type, second_type)
+        self.message += "\n\nType comparison result: {}".format(type_comatibility)

nemo/core/optim/__init__.py

@@ -0,0 +1,38 @@
+# 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
+
+from nemo.core.optim.adafactor import Adafactor
+from nemo.core.optim.adan import Adan
+from nemo.core.optim.flash_optim import patch_flashoptim_uneven_shard_support
+from nemo.core.optim.lr_scheduler import (
+    CosineAnnealing,
+    InverseSquareRootAnnealing,
+    NoamAnnealing,
+    PolynomialDecayAnnealing,
+    PolynomialHoldDecayAnnealing,
+    SquareAnnealing,
+    SquareRootAnnealing,
+    T5InverseSquareRootAnnealing,
+    WarmupAnnealing,
+    WarmupHoldAnnealLinear,
+    WarmupHoldAnnealOneMinusSquareRoot,
+    WarmupHoldPolicy,
+    WarmupPolicy,
+    prepare_lr_scheduler,
+)
+from nemo.core.optim.novograd import Novograd
+from nemo.core.optim.optimizer_with_main_params import MainParamsOptimizerWrapper
+from nemo.core.optim.optimizers import get_optimizer, parse_optimizer_args, register_optimizer

nemo/core/optim/adafactor.py

@@ -0,0 +1,218 @@
+# 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.
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Most of the code here has been copied from:
+# https://github.com/pytorch/fairseq/blob/main/fairseq/optim/adafactor.py
+
+import math
+
+import torch
+from torch.optim.optimizer import Optimizer
+
+__all__ = ['Adafactor']
+
+
+class Adafactor(Optimizer):
+    """Implements Adafactor algorithm.
+    This implementation is based on:
+    `Adafactor: Adaptive Learning Rates with Sublinear Memory Cost`
+    (see https://arxiv.org/abs/1804.04235)
+    Note that this optimizer internally adjusts the learning rate
+    depending on the *scale_parameter*, *relative_step* and
+    *warmup_init* options. To use a manual (external) learning rate
+    schedule you should set `scale_parameter=False` and
+    `relative_step=False`.
+    Args:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): external learning rate (default: None)
+        eps (tuple[float, float]): regularization constans for square gradient
+            and parameter scale respectively (default: (1e-30, 1e-3))
+        clip_threshold (float): threshold of root mean square of
+            final gradient update (default: 1.0)
+        decay_rate (float): coefficient used to compute running averages of square
+            gradient (default: -0.8)
+        beta1 (float): coefficient used for computing running averages of gradient
+            (default: None)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        scale_parameter (bool): if True, learning rate is scaled by root mean square of
+            parameter (default: True)
+        relative_step (bool): if True, time-dependent learning rate is computed
+            instead of external learning rate (default: True)
+        warmup_init (bool): time-dependent learning rate computation depends on
+            whether warm-up initialization is being used (default: False)
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=None,
+        eps=(1e-30, 1e-3),
+        clip_threshold=1.0,
+        decay_rate=-0.8,
+        beta1=None,
+        weight_decay=0.0,
+        scale_parameter=True,
+        relative_step=True,
+        warmup_init=False,
+        min_step=1e-2,
+    ):
+        if lr is not None and relative_step:
+            raise ValueError("Cannot combine manual lr and relative_step options")
+        if warmup_init and not relative_step:
+            raise ValueError("warmup_init requires relative_step=True")
+        self.min_step = min_step
+
+        defaults = dict(
+            lr=lr,
+            eps=eps,
+            clip_threshold=clip_threshold,
+            decay_rate=decay_rate,
+            beta1=beta1,
+            weight_decay=weight_decay,
+            scale_parameter=scale_parameter,
+            relative_step=relative_step,
+            warmup_init=warmup_init,
+            min_step=min_step,
+        )
+        super(Adafactor, self).__init__(params, defaults)
+
+    @property
+    def supports_memory_efficient_fp16(self):
+        """Whether this optimizer supports memory-efficient fp16 training."""
+        return True
+
+    @property
+    def supports_flat_params(self):
+        """Whether this optimizer supports flattened parameters."""
+        return False
+
+    def _get_lr(self, param_group, param_state):
+        rel_step_sz = param_group["lr"]
+        if param_group["relative_step"]:
+            min_step = 1e-6 * param_state["step"] if param_group["warmup_init"] else self.min_step
+            rel_step_sz = min(min_step, 1.0 / math.sqrt(param_state["step"]))
+        param_scale = 1.0
+        if param_group["scale_parameter"]:
+            param_scale = max(param_group["eps"][1], param_state["RMS"])
+        return param_scale * rel_step_sz
+
+    def _get_options(self, param_group, param_shape):
+        factored = len(param_shape) >= 2
+        use_first_moment = param_group["beta1"] is not None
+        return factored, use_first_moment
+
+    def _rms(self, tensor):
+        return tensor.norm(2) / (tensor.numel() ** 0.5)
+
+    def _approx_sq_grad(self, exp_avg_sq_row, exp_avg_sq_col):
+        r_factor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1)
+        c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
+        return torch.mul(r_factor, c_factor)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        Args:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.dtype in {torch.float16, torch.bfloat16}:
+                    grad = grad.float()
+                if grad.is_sparse:
+                    raise RuntimeError("Adafactor does not support sparse gradients.")
+
+                state = self.state[p]
+                grad_shape = grad.shape
+
+                factored, use_first_moment = self._get_options(group, grad_shape)
+                # State Initialization
+                if len(state) == 0:
+                    state["step"] = 0
+
+                    if use_first_moment:
+                        # Exponential moving average of gradient values
+                        state["exp_avg"] = torch.zeros_like(grad)
+                    if factored:
+                        state["exp_avg_sq_row"] = torch.zeros(grad_shape[:-1]).to(grad)
+                        state["exp_avg_sq_col"] = torch.zeros(grad_shape[:-2] + grad_shape[-1:]).to(grad)
+                    else:
+                        state["exp_avg_sq"] = torch.zeros_like(grad)
+
+                    state["RMS"] = 0
+                else:
+                    if use_first_moment:
+                        state["exp_avg"] = state["exp_avg"].to(grad)
+                    if factored:
+                        state["exp_avg_sq_row"] = state["exp_avg_sq_row"].to(grad)
+                        state["exp_avg_sq_col"] = state["exp_avg_sq_col"].to(grad)
+                    else:
+                        state["exp_avg_sq"] = state["exp_avg_sq"].to(grad)
+
+                p_data_fp32 = p.data
+                if p.data.dtype in {torch.float16, torch.bfloat16}:
+                    p_data_fp32 = p_data_fp32.float()
+
+                state["step"] += 1
+                state["RMS"] = self._rms(p_data_fp32)
+                group["lr"] = self._get_lr(group, state)
+
+                beta2t = 1.0 - math.pow(state["step"], group["decay_rate"])
+                update = (grad**2) + group["eps"][0]
+                if factored:
+                    exp_avg_sq_row = state["exp_avg_sq_row"]
+                    exp_avg_sq_col = state["exp_avg_sq_col"]
+
+                    exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=1.0 - beta2t)
+                    exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=1.0 - beta2t)
+
+                    # Approximation of exponential moving average of square of gradient
+                    update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)
+                    update.mul_(grad)
+                else:
+                    exp_avg_sq = state["exp_avg_sq"]
+
+                    exp_avg_sq.mul_(beta2t).add_(update, alpha=1.0 - beta2t)
+                    update = exp_avg_sq.rsqrt().mul_(grad)
+
+                update.div_((self._rms(update) / group["clip_threshold"]).clamp_(min=1.0))
+                update.mul_(group["lr"])
+
+                if use_first_moment:
+                    exp_avg = state["exp_avg"]
+                    exp_avg.mul_(group["beta1"]).add_(update, alpha=1 - group["beta1"])
+                    update = exp_avg
+
+                if group["weight_decay"] != 0:
+                    p_data_fp32.add_(p_data_fp32, alpha=-group["weight_decay"] * group["lr"])
+
+                p_data_fp32.add_(-update)
+
+                if p.data.dtype in {torch.float16, torch.bfloat16}:
+                    p.data.copy_(p_data_fp32)
+
+        return loss

nemo/core/optim/adan.py

@@ -0,0 +1,453 @@
+# 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.
+
+# Copyright 2022 Garena Online Private Limited
+#
+# 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 typing import List
+
+import torch
+from torch import Tensor
+from torch.optim.optimizer import Optimizer
+
+
+class MultiTensorApply(object):
+    available = False
+    warned = False
+
+    def __init__(self, chunk_size):
+        try:
+            MultiTensorApply.available = True
+            self.chunk_size = chunk_size
+        except ImportError as err:
+            MultiTensorApply.available = False
+            MultiTensorApply.import_err = err
+
+    def __call__(self, op, noop_flag_buffer, tensor_lists, *args):
+        return op(self.chunk_size, noop_flag_buffer, tensor_lists, *args)
+
+
+class Adan(Optimizer):
+    """
+    Implements a pytorch variant of Adan
+    Adan was proposed in
+    Adan: Adaptive Nesterov Momentum Algorithm for
+        Faster Optimizing Deep Models[J].arXiv preprint arXiv:2208.06677, 2022.
+    https://arxiv.org/abs/2208.06677
+    Arguments:
+        params (iterable): iterable of parameters to optimize or
+            dicts defining parameter groups.
+        lr (float, optional): learning rate. (default: 1e-3)
+        betas (Tuple[float, float, flot], optional): coefficients used for
+            first- and second-order moments. (default: (0.98, 0.92, 0.99))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability. (default: 1e-8)
+        weight_decay (float, optional): decoupled weight decay
+            (L2 penalty) (default: 0)
+        max_grad_norm (float, optional): value used to clip
+            global grad norm (default: 0.0 no clip)
+        no_prox (bool): how to perform the decoupled weight decay
+            (default: False)
+        foreach (bool): if True would use torch._foreach implementation.
+            It's faster but uses slightly more memory. (default: True)
+        fused (bool, optional): whether fused implementation is used.
+            (default: False)
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=1e-3,
+        betas=(0.98, 0.92, 0.99),
+        eps=1e-8,
+        weight_decay=0.0,
+        max_grad_norm=0.0,
+        no_prox=False,
+        foreach: bool = True,
+        fused: bool = False,
+    ):
+        if not 0.0 <= max_grad_norm:
+            raise ValueError('Invalid Max grad norm: {}'.format(max_grad_norm))
+        if not 0.0 <= lr:
+            raise ValueError('Invalid learning rate: {}'.format(lr))
+        if not 0.0 <= eps:
+            raise ValueError('Invalid epsilon value: {}'.format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError('Invalid beta parameter at index 0: {}'.format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError('Invalid beta parameter at index 1: {}'.format(betas[1]))
+        if not 0.0 <= betas[2] < 1.0:
+            raise ValueError('Invalid beta parameter at index 2: {}'.format(betas[2]))
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            max_grad_norm=max_grad_norm,
+            no_prox=no_prox,
+            foreach=foreach,
+            fused=fused,
+        )
+        super().__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(Adan, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault('no_prox', False)
+
+    @torch.no_grad()
+    def restart_opt(self):
+        for group in self.param_groups:
+            group['step'] = 0
+            for p in group['params']:
+                if p.requires_grad:
+                    state = self.state[p]
+                    # State initialization
+
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p)
+                    # Exponential moving average of gradient difference
+                    state['exp_avg_diff'] = torch.zeros_like(p)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step."""
+
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        if self.defaults['max_grad_norm'] > 0:
+            device = self.param_groups[0]['params'][0].device
+            global_grad_norm = torch.zeros(1, device=device)
+
+            max_grad_norm = torch.tensor(self.defaults['max_grad_norm'], device=device)
+            for group in self.param_groups:
+
+                for p in group['params']:
+                    if p.grad is not None:
+                        grad = p.grad
+                        global_grad_norm.add_(grad.pow(2).sum())
+
+            global_grad_norm = torch.sqrt(global_grad_norm)
+
+            clip_global_grad_norm = torch.clamp(max_grad_norm / (global_grad_norm + group['eps']), max=1.0).item()
+        else:
+            clip_global_grad_norm = 1.0
+
+        for group in self.param_groups:
+            params_with_grad = []
+            grads = []
+            exp_avgs = []
+            exp_avg_sqs = []
+            exp_avg_diffs = []
+            neg_pre_grads = []
+
+            beta1, beta2, beta3 = group['betas']
+            # assume same step across group now to simplify things
+            # per parameter step can be easily support
+            # by making it tensor, or pass list into kernel
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            bias_correction1 = 1.0 - beta1 ** group['step']
+            bias_correction2 = 1.0 - beta2 ** group['step']
+            bias_correction3 = 1.0 - beta3 ** group['step']
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                params_with_grad.append(p)
+                grads.append(p.grad)
+
+                state = self.state[p]
+                if len(state) == 0:
+                    state['exp_avg'] = torch.zeros_like(p)
+                    state['exp_avg_sq'] = torch.zeros_like(p)
+                    state['exp_avg_diff'] = torch.zeros_like(p)
+
+                if 'neg_pre_grad' not in state or group['step'] == 1:
+                    state['neg_pre_grad'] = p.grad.clone().mul_(-clip_global_grad_norm)
+
+                exp_avgs.append(state['exp_avg'])
+                exp_avg_sqs.append(state['exp_avg_sq'])
+                exp_avg_diffs.append(state['exp_avg_diff'])
+                neg_pre_grads.append(state['neg_pre_grad'])
+
+            kwargs = dict(
+                params=params_with_grad,
+                grads=grads,
+                exp_avgs=exp_avgs,
+                exp_avg_sqs=exp_avg_sqs,
+                exp_avg_diffs=exp_avg_diffs,
+                neg_pre_grads=neg_pre_grads,
+                beta1=beta1,
+                beta2=beta2,
+                beta3=beta3,
+                bias_correction1=bias_correction1,
+                bias_correction2=bias_correction2,
+                bias_correction3_sqrt=math.sqrt(bias_correction3),
+                lr=group['lr'],
+                weight_decay=group['weight_decay'],
+                eps=group['eps'],
+                no_prox=group['no_prox'],
+                clip_global_grad_norm=clip_global_grad_norm,
+            )
+
+            if group['foreach']:
+                if group['fused']:
+                    if torch.cuda.is_available():
+                        _fused_adan_multi_tensor(**kwargs)
+                    else:
+                        raise ValueError('Fused Adan does not support CPU')
+                else:
+                    _multi_tensor_adan(**kwargs)
+            elif group['fused']:
+                if torch.cuda.is_available():
+                    _fused_adan_single_tensor(**kwargs)
+                else:
+                    raise ValueError('Fused Adan does not support CPU')
+            else:
+                _single_tensor_adan(**kwargs)
+
+        return loss
+
+
+def _single_tensor_adan(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    exp_avg_diffs: List[Tensor],
+    neg_pre_grads: List[Tensor],
+    *,
+    beta1: float,
+    beta2: float,
+    beta3: float,
+    bias_correction1: float,
+    bias_correction2: float,
+    bias_correction3_sqrt: float,
+    lr: float,
+    weight_decay: float,
+    eps: float,
+    no_prox: bool,
+    clip_global_grad_norm: Tensor,
+):
+    for i, param in enumerate(params):
+        grad = grads[i]
+        exp_avg = exp_avgs[i]
+        exp_avg_sq = exp_avg_sqs[i]
+        exp_avg_diff = exp_avg_diffs[i]
+        neg_grad_or_diff = neg_pre_grads[i]
+
+        grad.mul_(clip_global_grad_norm)
+
+        # for memory saving, we use `neg_grad_or_diff`
+        # to get some temp variable in a inplace way
+        neg_grad_or_diff.add_(grad)
+
+        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)  # m_t
+        exp_avg_diff.mul_(beta2).add_(neg_grad_or_diff, alpha=1 - beta2)  # diff_t
+
+        neg_grad_or_diff.mul_(beta2).add_(grad)
+        exp_avg_sq.mul_(beta3).addcmul_(neg_grad_or_diff, neg_grad_or_diff, value=1 - beta3)  # n_t
+
+        denom = ((exp_avg_sq).sqrt() / bias_correction3_sqrt).add_(eps)
+        step_size_diff = lr * beta2 / bias_correction2
+        step_size = lr / bias_correction1
+
+        if no_prox:
+            param.mul_(1 - lr * weight_decay)
+            param.addcdiv_(exp_avg, denom, value=-step_size)
+            param.addcdiv_(exp_avg_diff, denom, value=-step_size_diff)
+        else:
+            param.addcdiv_(exp_avg, denom, value=-step_size)
+            param.addcdiv_(exp_avg_diff, denom, value=-step_size_diff)
+            param.div_(1 + lr * weight_decay)
+
+        neg_grad_or_diff.zero_().add_(grad, alpha=-1.0)
+
+
+def _multi_tensor_adan(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    exp_avg_diffs: List[Tensor],
+    neg_pre_grads: List[Tensor],
+    *,
+    beta1: float,
+    beta2: float,
+    beta3: float,
+    bias_correction1: float,
+    bias_correction2: float,
+    bias_correction3_sqrt: float,
+    lr: float,
+    weight_decay: float,
+    eps: float,
+    no_prox: bool,
+    clip_global_grad_norm: Tensor,
+):
+    if len(params) == 0:
+        return
+
+    torch._foreach_mul_(grads, clip_global_grad_norm)
+
+    # for memory saving, we use `neg_pre_grads`
+    # to get some temp variable in a inplace way
+    torch._foreach_add_(neg_pre_grads, grads)
+
+    torch._foreach_mul_(exp_avgs, beta1)
+    torch._foreach_add_(exp_avgs, grads, alpha=1 - beta1)  # m_t
+
+    torch._foreach_mul_(exp_avg_diffs, beta2)
+    torch._foreach_add_(exp_avg_diffs, neg_pre_grads, alpha=1 - beta2)  # diff_t
+
+    torch._foreach_mul_(neg_pre_grads, beta2)
+    torch._foreach_add_(neg_pre_grads, grads)
+    torch._foreach_mul_(exp_avg_sqs, beta3)
+    torch._foreach_addcmul_(exp_avg_sqs, neg_pre_grads, neg_pre_grads, value=1 - beta3)  # n_t
+
+    denom = torch._foreach_sqrt(exp_avg_sqs)
+    torch._foreach_div_(denom, bias_correction3_sqrt)
+    torch._foreach_add_(denom, eps)
+
+    step_size_diff = lr * beta2 / bias_correction2
+    step_size = lr / bias_correction1
+
+    if no_prox:
+        torch._foreach_mul_(params, 1 - lr * weight_decay)
+        torch._foreach_addcdiv_(params, exp_avgs, denom, value=-step_size)
+        torch._foreach_addcdiv_(params, exp_avg_diffs, denom, value=-step_size_diff)
+    else:
+        torch._foreach_addcdiv_(params, exp_avgs, denom, value=-step_size)
+        torch._foreach_addcdiv_(params, exp_avg_diffs, denom, value=-step_size_diff)
+        torch._foreach_div_(params, 1 + lr * weight_decay)
+    torch._foreach_zero_(neg_pre_grads)
+    torch._foreach_add_(neg_pre_grads, grads, alpha=-1.0)
+
+
+def _fused_adan_multi_tensor(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    exp_avg_diffs: List[Tensor],
+    neg_pre_grads: List[Tensor],
+    *,
+    beta1: float,
+    beta2: float,
+    beta3: float,
+    bias_correction1: float,
+    bias_correction2: float,
+    bias_correction3_sqrt: float,
+    lr: float,
+    weight_decay: float,
+    eps: float,
+    no_prox: bool,
+    clip_global_grad_norm: Tensor,
+):
+    import fused_adan
+
+    multi_tensor_applier = MultiTensorApply(2048 * 32)
+    _dummy_overflow_buf = torch.cuda.IntTensor([0])
+    multi_tensor_applier(
+        fused_adan.adan_multi_tensor,
+        _dummy_overflow_buf,
+        [params, grads, exp_avgs, exp_avg_sqs, exp_avg_diffs, neg_pre_grads],
+        beta1,
+        beta2,
+        beta3,
+        bias_correction1,
+        bias_correction2,
+        bias_correction3_sqrt,
+        lr,
+        weight_decay,
+        eps,
+        no_prox,
+        clip_global_grad_norm,
+    )
+    torch._foreach_zero_(neg_pre_grads)
+    torch._foreach_add_(neg_pre_grads, grads, alpha=-1.0)
+
+
+def _fused_adan_single_tensor(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    exp_avg_diffs: List[Tensor],
+    neg_pre_grads: List[Tensor],
+    *,
+    beta1: float,
+    beta2: float,
+    beta3: float,
+    bias_correction1: float,
+    bias_correction2: float,
+    bias_correction3_sqrt: float,
+    lr: float,
+    weight_decay: float,
+    eps: float,
+    no_prox: bool,
+    clip_global_grad_norm: Tensor,
+):
+    for i, param in enumerate(params):
+        p_data_fp32 = param.data.float()
+        out_p = param.data
+        grad = grads[i]
+        exp_avg = exp_avgs[i]
+        exp_avg_sq = exp_avg_sqs[i]
+        exp_avg_diff = exp_avg_diffs[i]
+        neg_grad = neg_pre_grads[i]
+        with torch.cuda.device(param.device):
+            import fused_adan
+
+            fused_adan.adan_single_tensor(
+                p_data_fp32,
+                out_p,
+                grad,
+                exp_avg,
+                exp_avg_sq,
+                exp_avg_diff,
+                neg_grad,
+                beta1,
+                beta2,
+                beta3,
+                bias_correction1,
+                bias_correction2,
+                bias_correction3_sqrt,
+                lr,
+                weight_decay,
+                eps,
+                no_prox,
+                clip_global_grad_norm,
+            )
+        neg_grad.zero_().add_(grad, alpha=-1.0)

nemo/core/optim/distributed_adam.py

@@ -0,0 +1,853 @@
+# 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 collections
+import contextlib
+import itertools
+from typing import Callable, Dict, Iterable, Optional, Tuple, Union
+
+import torch
+from apex.contrib.optimizers.distributed_fused_adam import (
+    DistributedFusedAdam,
+    _disable_pre_forward_hook,
+    _multi_tensor_copy,
+)
+
+try:
+    import apex.contrib.nccl_allocator as nccl_allocator
+except ImportError:
+    nccl_allocator = None
+
+from megatron.core import parallel_state
+from megatron.core.dist_checkpointing.dict_utils import dict_list_map_inplace
+from megatron.core.dist_checkpointing.mapping import ShardedTensor
+from megatron.core.dist_checkpointing.optimizer import get_param_id_to_sharded_param_map, optim_state_to_sharding_state
+
+from nemo.utils import logging, str_to_dtype
+from nemo.utils.te_utils import is_float8tensor, is_mxfp8tensor, te_version
+
+if te_version() >= (2, 0):
+
+    # TE quantization logic using quantizer API
+    # Supported TE versions: 2.0+
+
+    from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor
+
+    def _quantize_param_fragment_impl(
+        input_: torch.Tensor,
+        *,
+        out: torch.Tensor,
+        param: torch.nn.Parameter,
+    ) -> None:
+        quantizer = param._quantizer
+        out = Float8Tensor(
+            shape=input_.size(),
+            dtype=param.dtype,
+            requires_grad=False,
+            data=out,
+            fp8_scale_inv=param._scale_inv,
+            fp8_dtype=param._fp8_dtype,
+            quantizer=quantizer,
+        )
+        quantizer.update_quantized(input_, out)
+
+    def _get_fp8_scale_and_amax_impl(tensor: Float8Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        quantizer = tensor._quantizer
+        return quantizer.scale, quantizer.amax
+
+elif te_version() >= (1, 0):
+
+    # TE quantization logic with fp8_meta dicts
+    # Supported TE versions: 1.0 - 1.14
+
+    from transformer_engine.pytorch.cpp_extensions import cast_to_fp8
+
+    def _quantize_param_fragment_impl(
+        input_: torch.Tensor,
+        *,
+        out: torch.Tensor,
+        param: torch.nn.Parameter,
+    ) -> None:
+        cast_to_fp8(
+            src.view(1, -1),
+            param._fp8_meta["scaling_fwd"],
+            param._fp8_meta_index,
+            param._fp8_dtype,
+            out=dst.view(1, -1),
+        )
+
+    def _get_fp8_scale_and_amax_impl(tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        fp8_meta = tensor._fp8_meta["scaling_fwd"]
+        fp8_meta_index = tensor._fp8_meta_index
+        return fp8_meta.scale[fp8_meta_index], fp8_meta.amax_history[0][fp8_meta_index]
+
+else:
+
+    # Fallback impl if TE version is invalid
+    def _quantize_param_fragment_impl(*args, **kwargs) -> None:
+        raise RuntimeError("Invalid Transformer Engine version for FP8 distributed optimizer")
+
+    def _get_fp8_scale_and_amax_impl(*args, **kwargs):
+        raise RuntimeError("Invalid Transformer Engine version for FP8 distributed optimizer")
+
+
+def quantize_param_fragment(
+    input_: torch.Tensor,
+    *,
+    out: torch.Tensor,
+    param: torch.nn.Parameter,
+) -> None:
+    """Cast values in parameter fragment to FP8
+
+    Arguments:
+      input_ (torch.Tensor): Values to quantize.
+      out (torch.Tensor): Raw UINT8 buffer to fill with FP8 values.
+          Dimensions should match input_.
+      param (torch.nn.Parameter): Parameter containing this parameter
+          fragment. Must be a Float8Tensor.
+
+    """
+    _quantize_param_fragment_impl(input_, out=out, param=param)
+
+
+def get_fp8_scale_and_amax(tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Get FP8 scale and amax from Float8Tensor"""
+    return _get_fp8_scale_and_amax_impl(tensor)
+
+
+_distributed_pgs = {}
+
+
+def create_distributed_pgs(*, distributed_size: int) -> Dict:
+    """Create process groups for distributing within multiple devices.
+
+    User can reuse this function to reorder communicators for SHArP.
+
+    Arguments:
+        distributed_size (int): the number of devices to distribute optimizer
+            state over.
+
+    """
+    global _distributed_pgs
+    assert torch.distributed.is_initialized()
+    if _distributed_pgs:
+        return _distributed_pgs
+
+    world_size = torch.distributed.get_world_size()
+    rank = torch.distributed.get_rank()
+    devices = distributed_size
+    nodes = world_size // devices
+
+    if nodes * devices != world_size:
+        logging.warning("Expected all nodes have the same amout of devices, disable distribute_within_nodes.")
+        return {}
+
+    node_id = rank // devices
+    device_id = rank % devices
+
+    distributed_pgs = []
+    for i in range(nodes):
+        ranks = [i * devices + j for j in range(devices)]
+        pg = torch.distributed.new_group(ranks=ranks)
+        distributed_pgs.append(pg)
+
+    redundant_pgs = []
+    for i in range(devices):
+        ranks = [i + j * devices for j in range(nodes)]
+        pg = torch.distributed.new_group(ranks=ranks)
+        redundant_pgs.append(pg)
+
+    # To re-order SHArP communicator right after distributed init,
+    # we have to expose redundant_process_group to user.
+    # User has too invoke allreduce through redundant_process_group
+    # before all other communicators to lock SHArP tree.
+    _distributed_pgs = {
+        'world_size': world_size,
+        'rank': rank,
+        'devices': devices,
+        'nodes': nodes,
+        'node_id': node_id,
+        'device_id': device_id,
+        'distributed_process_group': distributed_pgs[node_id],
+        'redundant_process_group': redundant_pgs[device_id],
+    }
+    return _distributed_pgs
+
+
+def create_distribute_within_nodes_pgs():
+    """Create process groups for distributing within nodes.
+
+    User can reuse this function to reorder communicators for SHArP.
+    This funcion is kept for backward compatibility.
+    """
+    return create_distributed_pgs(distributed_size=torch.cuda.device_count())
+
+
+class MegatronDistributedFusedAdam(DistributedFusedAdam):
+    """Adam optimizer with ZeRO algorithm
+
+    Child class of Apex DistributedFusedAdam, with optimizations for
+    NeMo-Megatron.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts
+            defining parameter groups.
+        disable_distributed_parameters (bool, optional): use standard
+            data-parallel communication instead of ZeRO.
+            (default: False)
+        distribute_within_nodes (bool, optional): distribute states
+            within the same node, e.g. DGX. This can improve performance
+            but requires larger memory than distributing within all
+            ranks, especially for pure data parallel models.
+            (default: False).
+        distributed_size (int, optional): the number of devices to
+            distribute optimizer state over.
+        lock_timeout (float, optional): timeout for callback mutex in
+            seconds.
+        **kwargs: keyword arguments to pass to Apex
+            DistributedFusedAdam.
+
+    """
+
+    def __init__(
+        self,
+        params: Union[Iterable[torch.nn.Parameter], Iterable[dict]],
+        disable_distributed_parameters: bool = False,
+        distribute_within_nodes: bool = False,
+        distributed_size: Optional[int] = None,
+        lock_timeout: Optional[float] = None,
+        **kwargs,
+    ):
+
+        # Update distributed_size settings
+        if distribute_within_nodes:
+            if distributed_size is not None and distributed_size != torch.cuda.device_count():
+                raise ValueError("Inconsistent distributed_size value")
+            distributed_size = torch.cuda.device_count()
+
+        # Initialize process groups
+        if 'process_group' not in kwargs and parallel_state.is_initialized():
+            kwargs['process_group'] = parallel_state.get_data_parallel_group(with_context_parallel=True)
+        if disable_distributed_parameters:
+            world_size = torch.distributed.get_world_size()
+            rank = torch.distributed.get_rank()
+            self_groups = [torch.distributed.new_group(ranks=[i]) for i in range(world_size)]
+            kwargs['distributed_process_group'] = self_groups[rank]
+            kwargs['redundant_process_group'] = kwargs['process_group']
+        elif distributed_size is not None:
+            dist_pg_infos = create_distributed_pgs(distributed_size=distributed_size)
+            if dist_pg_infos:
+                kwargs['distributed_process_group'] = dist_pg_infos['distributed_process_group']
+                kwargs['redundant_process_group'] = dist_pg_infos['redundant_process_group']
+                global _distributed_pgs
+                _distributed_pgs = {}
+
+        # Make sure dtypes are in right type
+        for keyword in ('dtype', 'grad_sync_dtype', 'param_sync_dtype'):
+            if keyword in kwargs:
+                kwargs[keyword] = str_to_dtype(kwargs[keyword])
+
+        # Make sure params are in consistent format (list of param group dicts)
+        param_groups = list(params)
+        assert param_groups
+        if not isinstance(param_groups[0], dict):
+            param_groups = [{'params': param_groups}]
+
+        # Construct distributed optimizer
+        super().__init__(param_groups, **kwargs)
+
+        # Create mutex with timeout
+        self._lock_with_timeout = None
+        if lock_timeout is not None:
+
+            @contextlib.contextmanager
+            def lock_with_timeout():
+                result = self._lock.acquire(timeout=lock_timeout)
+                try:
+                    yield result
+                finally:
+                    if result:
+                        # Acquired lock before timeout
+                        self._lock.release()
+                    else:
+                        # Failed to acquire lock before timeout
+                        print(f'MegatronDistributedFusedAdam: Failed to acquire lock within {lock_timeout} seconds.')
+
+            self._lock_with_timeout = lock_with_timeout
+
+        # Check for MXFP8 parameters
+        if any(is_mxfp8tensor(param) for param in self.parameters()):
+            raise ValueError("Distributed optimizer currently does not support MXFP8 parameters")
+
+    def _broadcast_params(self) -> None:
+        # Assume params have already been synchronized
+        pass
+
+    def _make_post_backward_hook(self, param: torch.nn.Parameter, param_group_id: int, param_id: int) -> Callable:
+        def hook(*unused):
+            if getattr(param, '_pre_forward_hook_is_enabled', False):
+                raise RuntimeError(
+                    'A parameter called its post-backward hook '
+                    'before its pre-forward hook. '
+                    'Please manually interact with the parameter '
+                    'before the forward pass (e.g. by calling data_ptr) '
+                    'or run DistributedFusedAdam with overlap_param_sync=False.'
+                )
+            lock = self._lock
+            if self._lock_with_timeout is not None:
+                lock = self._lock_with_timeout()
+            with lock:
+                need_to_initialize = 'fragments' not in self.state[param]
+                if need_to_initialize:
+                    self._init_param_state(param, param_group_id, param_id)
+                if self.greedy_grad_copy and not getattr(param, '_disable_greedy_grad_copy', False):
+                    self._grad_copy(param)
+                    if self.overlap_grad_sync and not getattr(param, '_disable_overlap_grad_sync', False):
+                        self._try_start_bucket_grad_sync(
+                            params=[param],
+                            ignore_last_bucket=need_to_initialize,
+                        )
+
+        return hook
+
+    def init_params(
+        self,
+        params: Optional[Iterable[torch.nn.Parameter]] = None,
+        param_sync_dtype: Optional[torch.dtype] = None,
+        **kwargs,
+    ) -> None:
+        """Initialize optimizer state for parameters
+
+        Initializes FP8 and non-FP8 params separately.
+
+        """
+
+        # Default cases
+        if params is None:
+            params = self.parameters()
+        elif isinstance(params, torch.Tensor):
+            params = [params]
+
+        # Ignore parameters that have already been initialized
+        params = [param for param in params if "fragments" not in self.state[param]]
+        if not params:
+            return
+
+        # Initialize FP8 and non-FP8 tensors separately
+        if any(is_float8tensor(param) for param in params):
+            super().init_params(
+                filter(is_float8tensor, params),
+                param_sync_dtype=torch.uint8,
+                **kwargs,
+            )
+        super().init_params(
+            params,
+            param_sync_dtype=param_sync_dtype,
+            **kwargs,
+        )
+
+    def init_params_bucket(
+        self,
+        params: Iterable[torch.nn.Parameter],
+        grad_sync_dtype: Optional[torch.dtype] = None,
+        param_sync_dtype: Optional[torch.dtype] = None,
+        **kwargs,
+    ) -> None:
+        """Initialize optimizer state for parameters in one effective bucket"""
+
+        # Ignore parameters that have already been initialized
+        if isinstance(params, torch.Tensor):
+            params = [params]
+        params = [param for param in params if "fragments" not in self.state[param]]
+        if not params:
+            return
+
+        # Initialize parameters with FP32 grads
+        fp32_params = []
+        remaining_params = []
+        for param in params:
+            if getattr(param, '_with_fp32_optimizer', False):
+                fp32_params.append(param)
+            else:
+                remaining_params.append(param)
+        params = remaining_params
+        start_bucket_id = len(self.state["buckets"])
+        super().init_params_bucket(
+            fp32_params,
+            grad_sync_dtype=torch.float32,
+            param_sync_dtype=param_sync_dtype,
+            **kwargs,
+        )
+        end_bucket_id = len(self.state["buckets"])
+        fp32_buckets = self.state["buckets"][start_bucket_id:end_bucket_id]
+
+        # Initialize FP8 parameters
+        fp8_params = []
+        remaining_params = []
+        for param in params:
+            if is_float8tensor(param):
+                fp8_params.append(param)
+            else:
+                remaining_params.append(param)
+        params = remaining_params
+        start_bucket_id = len(self.state["buckets"])
+        super().init_params_bucket(
+            fp8_params,
+            grad_sync_dtype=grad_sync_dtype,
+            param_sync_dtype=torch.uint8,
+            **kwargs,
+        )
+        end_bucket_id = len(self.state["buckets"])
+        fp8_buckets = self.state["buckets"][start_bucket_id:end_bucket_id]
+
+        # Initialize remaining parameters as usual
+        normal_buckets = []
+        start_bucket_id = len(self.state["buckets"])
+        super().init_params_bucket(
+            params,
+            grad_sync_dtype=grad_sync_dtype,
+            param_sync_dtype=param_sync_dtype,
+            **kwargs,
+        )
+        end_bucket_id = len(self.state["buckets"])
+        normal_buckets = self.state["buckets"][start_bucket_id:end_bucket_id]
+
+        def add_param_to_bucket(
+            param: torch.nn.Parameter,
+            bucket: self.StateBucket,
+        ) -> None:
+            """Add trivial param fragment to bucket"""
+            param_fragments = self.state[param]["fragments"]
+            param_group_id = param_fragments[0].param_group_id
+            param_id = param_fragments[0].param_id
+            bucket_id = bucket.fragments[0].bucket_id
+            param_size = param.numel()
+            bucket_size = bucket.bucket_size
+            fragment = self.ParameterFragment(
+                param_group_id=param_group_id,
+                param_id=param_id,
+                bucket_id=bucket_id,
+                param_range=(param_size, param_size),
+                bucket_range=(bucket_size, bucket_size),
+                in_local_shard=False,
+                shard_range=None,
+                shard_bucket_range=None,
+                shard_param_range=None,
+            )
+            param_fragments.append(fragment)
+            bucket.fragments.append(fragment)
+
+        # Make sure all added buckets depend on provided params
+        for bucket in fp32_buckets:
+            for param in itertools.chain(fp8_params, params):
+                add_param_to_bucket(param, bucket)
+        for bucket in fp8_buckets:
+            for param in itertools.chain(fp32_params, params):
+                add_param_to_bucket(param, bucket)
+        for bucket in normal_buckets:
+            for param in itertools.chain(fp32_params, fp8_params):
+                add_param_to_bucket(param, bucket)
+
+    def _init_param_state(
+        self,
+        param: torch.nn.Parameter,
+        param_group_id: int,
+        param_id: int,
+        param_sync_dtype: Optional[torch.dtype] = None,
+        **kwargs,
+    ) -> None:
+        """Initialize optimizer state for a parameter
+
+        Initializing the master weights requires slicing a flattened
+        view of the param. FP8 tensors do not handle these operations
+        gracefully, so we hack around it by explicitly casting to
+        FP32.
+
+        """
+
+        # Initialize non-FP8 params as usual
+        if not is_float8tensor(param):
+            super()._init_param_state(
+                param,
+                param_group_id,
+                param_id,
+                param_sync_dtype=param_sync_dtype,
+                **kwargs,
+            )
+
+        # Return immediately if already initialized
+        if "fragments" in self.state[param]:
+            return
+
+        # Initialize with FP32 copy of param
+        fp32_param = param.float()
+        super()._init_param_state(
+            fp32_param,
+            param_group_id,
+            param_id,
+            param_sync_dtype=torch.uint8,
+            **kwargs,
+        )
+        self.state[param].update(self.state[fp32_param])
+        del self.state[fp32_param]
+
+    @torch.no_grad()
+    def init_param_buffer(self) -> None:
+        """Allocate contiguous buffers for param buckets
+
+        For FP8 params, the FP8 data buffer is made a view into a
+        contiguous buffer.
+
+        """
+
+        # Make sure all params are initialized
+        self.contiguous_param_buffer = True
+        self.init_params()
+
+        # Construct param buffers
+        buffer_sizes = collections.defaultdict(lambda: 0)
+        for bucket in self.state["buckets"]:
+            dtypes = bucket.dtypes()
+            buffer_sizes[dtypes] = max(bucket.contiguous_buffer_offset + bucket.bucket_size, buffer_sizes[dtypes])
+        for dtypes, buffer_size in buffer_sizes.items():
+            _, _, param_sync_dtype = dtypes
+            if getattr(self, "nccl_ub", False):
+                if not nccl_allocator:
+                    raise RuntimeError("NCCL allocator importing failed but nccl ub is still requested")
+                with nccl_allocator.nccl_mem():
+                    self._param_buffers[dtypes] = torch.zeros(
+                        [buffer_size], dtype=param_sync_dtype, device=self.device
+                    )
+            else:
+                self._param_buffers[dtypes] = torch.zeros([buffer_size], dtype=param_sync_dtype, device=self.device)
+        # Figure out corresponding positions in params and param buffer
+        params = list(self.parameters())
+        param_flat_views = []
+        param_buffer_views = []
+        for i, param in enumerate(params):
+            fragment = self.state[param]["fragments"][0]
+            bucket_id = fragment.bucket_id
+            bucket = self.state["buckets"][bucket_id]
+            param_size = param.numel()
+            bucket_start, _ = fragment.bucket_range
+            buffer_offset = bucket.contiguous_buffer_offset
+            buffer_start = buffer_offset + bucket_start
+            buffer_end = buffer_start + param_size
+            param_buffer = self._param_buffers[bucket.dtypes()]
+            param_buffer_view = param_buffer[buffer_start:buffer_end].detach()
+            if param_buffer_view.device != param.device:
+                raise RuntimeError(
+                    "Attempted to change a parameter with device={param.device} "
+                    f"into a buffer view with device={param_buffer_view.device}"
+                )
+            if is_float8tensor(param):
+                param_flat_views.append(param._data.detach().view(-1))
+            else:
+                if param_buffer_view.dtype != param.dtype:
+                    raise RuntimeError(
+                        f"Attempted to change a parameter with dtype={param.dtype} "
+                        f"into a buffer view with dtype={param_buffer_view.dtype}"
+                    )
+                if param.is_contiguous(memory_format=torch.channels_last):
+                    param = param.permute(0, 2, 3, 1)
+                param_flat_views.append(param.detach().view(-1))
+            param_buffer_views.append(param_buffer_view)
+
+        # Copy values into param buffer
+        _multi_tensor_copy(
+            param_flat_views,
+            param_buffer_views,
+            dummy_overflow_buf=self._dummy_overflow_buf,
+        )
+
+        # Make all params a view into the param buffer
+        for param, buffer_view in zip(params, param_buffer_views):
+            if is_float8tensor(param):
+                param._data = buffer_view.view(param.size())
+            else:
+                # Preserve memory format for param here, i.e. NHWC tensors
+                # `param.data.set_()` failed to change storage.
+                # `param.set_()` invalidates bprop hook.
+                param.data = torch.as_strided(
+                    buffer_view,
+                    param.size(),
+                    param.stride(),
+                    storage_offset=buffer_view.storage_offset(),
+                )
+
+    def try_grad_sync(self, params: Iterable[torch.nn.Parameter]) -> None:
+        """Attempt to launch gradient synchronization"""
+
+        def is_grad_copy_enabled(param: torch.nn.Parameter) -> bool:
+            return not getattr(param, '_disable_greedy_grad_copy', False) and not getattr(
+                param, '_disable_overlap_grad_sync', False
+            )
+
+        params = list(filter(is_grad_copy_enabled, params))
+        for p in params:
+            self._grad_copy(p)
+        self._try_start_bucket_grad_sync(params=params)
+
+    def zero_grad(self, *args, **kwargs) -> None:
+        """Clear parameter gradients"""
+        super().zero_grad(*args, **kwargs)
+
+        # Reset main grads
+        if self.contiguous_grad_buffer:
+            for param in self.parameters():
+                with _disable_pre_forward_hook(param):
+                    param.main_grad = self.grad_buffer_view(param)
+
+    def grad_norm(
+        self,
+        parameters: Optional[Iterable[torch.nn.Parameter]] = None,
+        norm_type: float = 2.0,
+        force: bool = False,
+    ) -> torch.Tensor:
+        """L2 norm of parameter gradients"""
+        assert norm_type == 2
+
+        if parameters is not None:
+            # Make sure we can access iterable multiple times
+            parameters = list(parameters)
+
+        # Compute grad norm
+        if force or self._grad_norm is None:
+
+            # Compute norm of local gradients for distributed optimizer
+            grad_norm_sq = self._local_grad_norm(parameters=parameters, norm_type=norm_type)
+            if self.redundant_size > 1:
+                grad_norm_sq /= self.redundant_size
+
+            # Sum over all procs to get grad norm
+            torch.distributed.all_reduce(
+                grad_norm_sq,
+                op=torch.distributed.ReduceOp.SUM,
+            )
+            self._grad_norm = grad_norm_sq.sqrt()
+
+        # Use cached grad norm
+        return super().grad_norm()
+
+    @torch.no_grad()
+    def _param_copy_fragments(self, fragments: Iterable[DistributedFusedAdam.ParameterFragment]) -> None:
+        """Update parameter fragments with values from parameter buckets
+
+        For FP8 params, values are copied directly into the FP8 data
+        buffer.
+
+        """
+
+        # Figure out corresponding positions in param buckets and params
+        buffers_in = []
+        buffers_out = []
+        fragments = list(fragments)
+        for fragment in fragments:
+
+            # Check if fragment needs to be updated
+            bucket_id = fragment.bucket_id
+            bucket_start, bucket_end = fragment.bucket_range
+            param_start, param_end = fragment.param_range
+            if param_end <= param_start or bucket_id not in self._params_buckets:
+                continue
+
+            # Corresponding positions in bucket and param
+            param_bucket = self._params_buckets[bucket_id]
+            param = self.parameter(fragment)
+            buffer_in = param_bucket.params_bucket[bucket_start:bucket_end]
+            if is_float8tensor(param):
+                # Copy into FP8 params's data buffer
+                assert (
+                    param_bucket.params_bucket.dtype == torch.uint8
+                ), "Expected FP8 params to perform param sync in UINT8"
+                buffer_out = param._data.view(-1)[param_start:param_end]
+                buffers_in.append(buffer_in)
+                buffers_out.append(buffer_out)
+            elif torch.is_floating_point(buffer_in) and torch.is_floating_point(param):
+                # Conv with NHWC layout, i.e. shape (N, C, H, W) and stride
+                # (HWC, 1, WC, C), can't `.view(-1)`. Here to turn it to
+                # tensor with shape (N, H, W, C) and stride (HWC, WC, C, 1).
+                # Note: https://github.com/NVIDIA/apex/pull/1794
+                if param.is_contiguous(memory_format=torch.channels_last):
+                    param = param.permute(0, 2, 3, 1)
+
+                # Cast between floating-point dtypes
+                buffer_out = param.detach().view(-1)[param_start:param_end]
+                buffers_in.append(buffer_in)
+                buffers_out.append(buffer_out)
+            else:
+                # Copy most significant bytes for non-floating-point
+                # dtypes
+                # Note: Assume dtypes are little-endian
+                buffer_out = param.detach().view(-1)[param_start:param_end]
+                in_bytes = buffer_in.unsqueeze(-1).view(torch.uint8)
+                out_bytes = buffer_out.unsqueeze(-1).view(torch.uint8)
+                copy_size = min(in_bytes.size(-1), out_bytes.size(-1))
+                buffers_in.append(in_bytes[..., -copy_size:])
+                buffers_out.append(out_bytes[..., -copy_size:])
+                if copy_size < out_bytes.size(-1):
+                    out_bytes[..., :-copy_size].zero_()
+
+        # Copy data from parameter buckets to parameters
+        _multi_tensor_copy(
+            buffers_in,
+            buffers_out,
+            dummy_overflow_buf=self._dummy_overflow_buf,
+        )
+
+        # Update transpose caches
+        params = set(self.parameter(fragment) for fragment in fragments)
+        for param in params:
+            if is_float8tensor(param):
+                param._reset_caches()
+
+    @torch.no_grad()
+    def _check_params_shard_dtypes(self, params_buckets: Dict[int, DistributedFusedAdam.ParameterBucket]) -> None:
+        """Make sure local shards of parameters are in expected datatypes
+
+        For FP8 params, FP32 values are cast into FP8 using per-param
+        scaling factors and per-param amaxes are computed and reduced.
+
+        """
+
+        # Just call base class function if there are no FP8 tensors
+        num_fp8_params = sum(1 for param in self.parameters() if is_float8tensor(param))
+        if num_fp8_params == 0:
+            super()._check_params_shard_dtypes(params_buckets)
+            return
+
+        # Cast local data to FP8
+        fp8_params_shards = dict()
+        for bucket_id, param_bucket in params_buckets.items():
+            state_bucket = self.state["buckets"][bucket_id]
+            if state_bucket.param_sync_dtype != torch.uint8:
+                continue
+
+            # Initialize FP8 buffer for param sync
+            params_shard = param_bucket.params_shard
+            if self.contiguous_param_buffer:
+                shard_size = state_bucket.shard_size
+                buffer_offset = state_bucket.contiguous_buffer_offset
+                buffer_start = buffer_offset + self.distributed_rank * shard_size
+                buffer_end = buffer_start + shard_size
+                param_buffer = self._param_buffers[state_bucket.dtypes()]
+                fp8_params_shard = param_buffer[buffer_start:buffer_end]
+            else:
+                fp8_params_shard = torch.empty_like(params_shard, dtype=torch.uint8)
+            param_bucket.params_shard = fp8_params_shard
+
+            # Cast param fragments to FP8
+            for fragment in self.state["buckets"][bucket_id].fragments:
+                param = self.parameter(fragment)
+                if not is_float8tensor(param):
+                    continue
+                if not fragment.in_local_shard:
+                    continue
+                shard_start, shard_end = fragment.shard_range
+                if shard_end <= shard_start:
+                    continue
+                shard_range = slice(shard_start, shard_end)
+                quantize_param_fragment(
+                    params_shard[shard_range],
+                    out=fp8_params_shard[shard_range],
+                    param=param,
+                )
+
+        # Update FP8 scaling factors when all buckets have processed
+        if getattr(self, "_check_params_shard_dtypes_progress", None) is None:
+            self._check_params_shard_dtypes_progress = []
+        self._check_params_shard_dtypes_progress.extend(params_buckets.keys())
+        if len(self._check_params_shard_dtypes_progress) == len(self.state["buckets"]):
+            assert len(set(self._check_params_shard_dtypes_progress)) == len(self.state["buckets"])
+
+            # FP8 scaling factors
+            amaxes = []
+            scales = []
+            scale_invs = []
+            i = -1
+            for param in self.parameters():
+                if not is_float8tensor(param):
+                    continue
+                i += 1
+                scale, amax = get_fp8_scale_and_amax(param)
+                amaxes.append(amax.view(1))
+                scales.append(scale.view(1))
+                scale_invs.append(param._scale_inv.view(1))
+
+            # Update cached scale-inverses
+            packed_scales = torch.empty(num_fp8_params, dtype=torch.float32, device=self.device)
+            packed_scale_views = [packed_scales[i].view(1) for i in range(num_fp8_params)]
+            _multi_tensor_copy(
+                scales,
+                packed_scale_views,
+                dummy_overflow_buf=self._dummy_overflow_buf,
+            )
+            torch.reciprocal(packed_scales, out=packed_scales)
+            _multi_tensor_copy(
+                packed_scale_views,
+                scale_invs,
+                dummy_overflow_buf=self._dummy_overflow_buf,
+            )
+
+            # Reduce amaxes
+            # Note: Assume each param has a separate amax
+            packed_amaxes = torch.empty(num_fp8_params, dtype=torch.float32, device=self.device)
+            packed_amax_views = [packed_amaxes[i].view(1) for i in range(num_fp8_params)]
+            _multi_tensor_copy(
+                amaxes,
+                packed_amax_views,
+                dummy_overflow_buf=self._dummy_overflow_buf,
+            )
+            torch.distributed.all_reduce(
+                packed_amaxes,
+                op=torch.distributed.ReduceOp.MAX,
+                group=self.distributed_process_group,
+            )
+            _multi_tensor_copy(
+                packed_amax_views,
+                amaxes,
+                dummy_overflow_buf=self._dummy_overflow_buf,
+            )
+
+            # Reset
+            self._check_params_shard_dtypes_progress = None
+
+        # Handle any remaining dtype conversions
+        super()._check_params_shard_dtypes(params_buckets)
+
+    def sharded_state_dict(self, model_sharded_state_dict, optimizer_state_dict=None):
+        """Create sharded state dict"""
+        if optimizer_state_dict is None:
+            optimizer_state_dict = self.state_dict()
+
+        id_to_sharded_param_map = get_param_id_to_sharded_param_map(
+            model_sharded_state_dict=model_sharded_state_dict,
+            optim_params_iter=self.parameters(),
+        )
+        # Convert state
+        step = optimizer_state_dict['state'].pop('step')
+        state_dict_format = optimizer_state_dict.pop('format', None)
+        optim_state_to_sharding_state(optimizer_state_dict, id_to_sharded_param_map)
+        optimizer_state_dict['state']['step'] = step
+        if state_dict_format is not None:
+            optimizer_state_dict['format'] = state_dict_format
+
+        def rename_fp32_params(x):
+            if isinstance(x, ShardedTensor) and x.key.startswith('optimizer.state.param'):
+                x.key = x.key.replace('optimizer.state.param', 'optimizer.state.fp32_param')
+            return x
+
+        dict_list_map_inplace(rename_fp32_params, optimizer_state_dict)
+
+        return optimizer_state_dict

nemo/core/optim/flash_optim.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.
+
+from typing import Any
+
+import torch
+
+from nemo.utils import logging
+
+
+def patch_flashoptim_uneven_shard_support(optimizer) -> None:
+    """Patch flashoptim to handle FSDP2 unevenly-sharded parameters in DCP.
+
+    FlashOptim <= 0.1.3 raises ``ValueError`` when saving optimizer state for
+    parameters whose shard dimension is not evenly divisible by the FSDP mesh
+    size. The root cause is that ``DTensor.from_local()`` is called without an
+    explicit ``shape``, so it infers ``global = local * mesh_size`` which is
+    wrong for padded (uneven) shards.
+
+    This patch replaces ``_wrap_state_as_dtensor`` on the optimizer class so
+    that ``shape=param.shape`` and ``stride=param.stride()`` are always passed,
+    which is correct for both even and uneven shards.
+    """
+
+    klass = type(optimizer)
+    if not hasattr(klass, "_wrap_state_as_dtensor"):
+        return
+    if getattr(klass, "_nemo_patched_uneven_shard", False):
+        return
+
+    @staticmethod
+    def _fixed_wrap_state_as_dtensor(state: dict[str, Any], param: torch.Tensor) -> None:  # noqa: UP006
+        if not hasattr(param, "device_mesh"):
+            return
+
+        from torch.distributed.tensor import DTensor
+
+        mesh = param.device_mesh
+        placements = param.placements
+
+        for key, val in state.items():
+            if isinstance(val, torch.Tensor) and not isinstance(val, DTensor) and val.dim() > 0:
+                state[key] = DTensor.from_local(
+                    val,
+                    mesh,
+                    placements,
+                    shape=param.shape,
+                    stride=param.stride(),
+                )
+
+    klass._wrap_state_as_dtensor = _fixed_wrap_state_as_dtensor
+    klass._nemo_patched_uneven_shard = True
+    logging.info("Patched flashoptim %s to support unevenly-sharded FSDP2 parameters in DCP.", klass.__name__)

nemo/core/optim/novograd.py

@@ -0,0 +1,150 @@
+# 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 torch.optim.optimizer import Optimizer
+
+__all__ = ['Novograd']
+
+
+def _check_valid_opt_params(lr, eps, betas):
+    if lr < 0:
+        raise ValueError(f"Invalid learning rate: {lr}")
+    if eps < 0:
+        raise ValueError(f"Invalid epsilon value: {eps}")
+    if not (0.0 <= betas[0] < 1.0 and 0.0 <= betas[1] < 1.0):
+        raise ValueError(f"Betas have to be between 0 and 1: {betas}")
+
+
+class Novograd(Optimizer):
+    """Implements Novograd algorithm.
+    It has been proposed  in "Stochastic Gradient Methods with Layer-wise
+    Adaptive Moments for Training of Deep Networks"
+    (https://arxiv.org/abs/1905.11286)
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper "On the Convergence of Adam and Beyond"
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=1e-3,
+        betas=(0.95, 0.98),
+        eps=1e-8,
+        weight_decay=0,
+        grad_averaging=False,
+        amsgrad=False,
+        luc=False,
+        luc_trust=1e-3,
+        luc_eps=1e-8,
+    ):
+        _check_valid_opt_params(lr, eps, betas)
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            grad_averaging=grad_averaging,
+            amsgrad=amsgrad,
+        )
+        self.luc = luc
+        self.luc_trust = luc_trust
+        self.luc_eps = luc_eps
+        super(Novograd, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(Novograd, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault("amsgrad", False)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError("Sparse gradients are not supported.")
+                amsgrad = group["amsgrad"]
+                state = self.state[p]
+
+                # State initialization
+                if not state:
+                    state["step"] = 0
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros([]).to(state["exp_avg"].device)
+                    if amsgrad:
+                        # Maintains max of all exp moving avg of squared grad
+                        state["max_exp_avg_sq"] = torch.zeros([]).to(state["exp_avg"].device)
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                if amsgrad:
+                    max_exp_avg_sq = state["max_exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+
+                state["step"] += 1
+
+                norm = grad.norm().pow(2)
+
+                if exp_avg_sq == 0:
+                    exp_avg_sq.copy_(norm)
+                else:
+                    exp_avg_sq.mul_(beta2).add_(norm, alpha=1.0 - beta2)
+
+                if amsgrad:
+                    # Maintains max of all 2nd moment running avg till now
+                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+                    # Use the max for normalizing running avg. of gradient
+                    denom = max_exp_avg_sq.sqrt().add_(group["eps"])
+                else:
+                    denom = exp_avg_sq.sqrt().add_(group["eps"])
+
+                grad.div_(denom)
+                if group["weight_decay"] != 0:
+                    grad.add_(p.data, alpha=group["weight_decay"])
+                if group["grad_averaging"]:
+                    grad.mul_(1 - beta1)
+                exp_avg.mul_(beta1).add_(grad)
+
+                if self.luc:
+                    # Clip update so that updates are less than eta*weights
+                    data_norm = torch.norm(p.data)
+                    grad_norm = torch.norm(exp_avg.data)
+                    luc_factor = self.luc_trust * data_norm / (grad_norm + self.luc_eps)
+                    luc_factor = min(luc_factor, group["lr"])
+                    p.data.add_(exp_avg, alpha=-luc_factor)
+                else:
+                    p.data.add_(exp_avg, alpha=-group["lr"])
+
+        return loss

nemo/core/optim/optimizer_with_main_params.py

@@ -0,0 +1,566 @@
+# 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.
+
+from contextlib import contextmanager
+
+import torch
+
+from nemo.utils import logging
+
+try:
+    import amp_C
+    from apex.multi_tensor_apply import multi_tensor_applier
+
+    HAVE_APEX = True
+
+except (ImportError, ModuleNotFoundError):
+
+    HAVE_APEX = False
+
+try:
+    from megatron.core.parallel_state import (
+        get_data_parallel_group,
+        get_data_parallel_world_size,
+        get_expert_data_parallel_group,
+    )
+    from megatron.core.tensor_parallel import copy_tensor_model_parallel_attributes
+
+    HAVE_MEGATRON_CORE = True
+
+except (ImportError, ModuleNotFoundError):
+
+    HAVE_MEGATRON_CORE = False
+
+
+def _zero_grad_group_helper(group, set_to_none):
+    """Zero out the gradient for a group of parameters.
+    Note: copied from torch.optim.optimizer."""
+    for param in group:
+        if param.grad is not None:
+            if set_to_none:
+                param.grad = None
+            else:
+                if param.grad.grad_fn is not None:
+                    param.grad.detach_()
+                else:
+                    param.grad.requires_grad_(False)
+                param.grad.zero_()
+
+
+def _multi_tensor_copy_this_to_that(this, that, overflow_buf):
+    """Use multi-tensor-applier to copy values from one list to another.
+    We don't have a blfoat16 implementation so for now if the overflow_buf
+    is not provided, we default back to simple loop copy to be compatible
+    with bfloat16."""
+    if overflow_buf:
+        # Scaling with factor `1.0` is equivalent to copy.
+        multi_tensor_applier(amp_C.multi_tensor_scale, overflow_buf, [this, that], 1.0)
+    else:
+        # FIXME: use multi-tensor applier for bf16
+        for this_, that_ in zip(this, that):
+            that_.copy_(this_)
+
+
+def _get_grad_data_group(is_expert_group):
+    if is_expert_group:
+        data_group = get_expert_data_parallel_group()
+    else:
+        data_group = get_data_parallel_group(with_context_parallel=True)
+    return data_group
+
+
+class GradBucket(object):
+    """
+    Persistent buffer for main gradients that remains allocated between training iterations
+    """
+
+    def __init__(self, numel, chunk_size_mb, data_group):
+        if not HAVE_APEX:
+            raise ImportError(
+                "Apex was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+            )
+
+        if not HAVE_MEGATRON_CORE:
+            raise ImportError(
+                "megatron-core was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+            )
+
+        self.numel = numel
+        self.data = torch.zeros(self.numel, dtype=torch.float, device=torch.cuda.current_device(), requires_grad=False)
+
+        self._data_group = data_group
+        self.chunk_size_mb = chunk_size_mb
+        if self.chunk_size_mb > 0:
+            chunk_size_bytes = chunk_size_mb * 1024 * 1024
+            self.chunk_size_numel = chunk_size_bytes // 4
+            self.num_chunks = self.numel // self.chunk_size_numel
+            self.numel_per_chunk = [self.chunk_size_numel] * self.num_chunks
+            if self.numel % self.chunk_size_numel != 0:
+                self.num_chunks += 1
+                self.numel_per_chunk.append(self.numel % self.chunk_size_numel)
+
+            self.start_index_per_chunk = torch.cumsum(torch.tensor([0] + self.numel_per_chunk[:-1]), dim=0)
+            self.current_chunk = 0
+            self.computed_numel_per_chunk = [0] * self.num_chunks
+
+    def zero(self):
+        """Reset the buffer to zero."""
+        self.data.zero_()
+
+    def allreduce_buffer(self):
+        """Synchronous buffer data allreduce"""
+        self.data.div_(get_data_parallel_world_size())
+        torch.distributed.all_reduce(self.data, group=self._data_group)
+
+    def get(self, shape, start_index):
+        """Return a tensor with the input `shape` as a view into the
+        1-D data starting at `start_index`."""
+        end_index = start_index + shape.numel()
+        assert end_index <= self.numel, 'requested tensor is out of the buffer range.'
+        buffer_tensor = self.data[start_index:end_index]
+        buffer_tensor = buffer_tensor.view(shape)
+
+        grad_chunk_info = None
+        if self.chunk_size_mb > 0:
+            grad_chunk_info = {}
+            chunk = start_index // self.chunk_size_numel
+            chunk_start_index = self.start_index_per_chunk[chunk]
+            chunk_end_index = chunk_start_index + self.numel_per_chunk[chunk]
+            grad_chunk_info[chunk] = min(chunk_end_index, end_index) - start_index
+            while chunk_end_index < end_index:
+                chunk += 1
+                chunk_start_index = self.start_index_per_chunk[chunk]
+                chunk_end_index = chunk_start_index + self.numel_per_chunk[chunk]
+                grad_chunk_info[chunk] = min(chunk_end_index, end_index) - chunk_start_index
+
+        return buffer_tensor, grad_chunk_info
+
+    def update_chunk_info(self, grad_chunk_info):
+        for chunk in grad_chunk_info.keys():
+            self.computed_numel_per_chunk[chunk] += grad_chunk_info[chunk]
+
+    def get_allreduce_tensor(self):
+        if self.computed_numel_per_chunk[self.current_chunk] == self.numel_per_chunk[self.current_chunk]:
+            chunk_start_index = self.start_index_per_chunk[self.current_chunk]
+            chunk_end_index = chunk_start_index + self.numel_per_chunk[self.current_chunk]
+            allreduce_tensor = self.data[chunk_start_index:chunk_end_index]
+
+            self.computed_numel_per_chunk[self.current_chunk] = 0
+            self.current_chunk += 1
+            if self.current_chunk == self.num_chunks:
+                self.current_chunk = 0
+
+            return allreduce_tensor
+
+        return None
+
+
+class MainParamsOptimizerWrapper(torch.optim.Optimizer):
+    """
+    Float16 optimizer wrapper for half precision (fp16 and bf16) data types.
+    This optimizer wrapper holds main parameters and gradients in fp32 to support
+    stable convergence.
+
+    Arguments:
+        optimizer: base optimizer such as Adam or SGD.
+        fp32_grad_accum: to enable the use of fp32 in gradient accumulation and allreduce.
+        contiguous_grad_bucket: to enable allocating the master gradients in the
+            contiguous memory space to reduce memory fragmentation.
+        async_grad_allreduce: enable asynchronous gradient allreduce that is executed
+            along with the training step backprop.
+    """
+
+    def __init__(
+        self,
+        optimizer,
+        fp32_grad_accum=False,
+        contiguous_grad_bucket=False,
+        async_grad_allreduce=False,
+        grad_div_ar_fusion=True,
+        grad_allreduce_chunk_size_mb=0,
+    ):
+        if not HAVE_APEX:
+            raise ImportError(
+                "Apex was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+            )
+
+        if not HAVE_MEGATRON_CORE:
+            raise ImportError(
+                "megatron-core was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+            )
+
+        self.optimizer = optimizer
+        assert self.optimizer, 'no optimizer is provided.'
+        if contiguous_grad_bucket:
+            assert fp32_grad_accum, 'contiguous gradient buffer assumes using fp32 grad.'
+        if async_grad_allreduce:
+            assert fp32_grad_accum, (
+                'async allreduce applies to master gradients only, '
+                'which is supposed to be accumulated after grad op.'
+            )
+            assert contiguous_grad_bucket, (
+                'currently async_grad_allreduce is supported only ' 'with contiguous_grad_bucket.'
+            )
+
+        self._fp32_grad_accum = fp32_grad_accum
+        self._contiguous_grad_bucket = contiguous_grad_bucket
+
+        # used with tensor parallel only (no pipeline parallelism)
+        # be careful, weight update cannot start until all async grad AR works are done
+        self._async_grad_allreduce = (
+            async_grad_allreduce and get_data_parallel_world_size(with_context_parallel=True) > 1
+        )
+
+        if self._async_grad_allreduce:
+            # use @no_sync to disable backward grad sync during gradient accumulation
+            self._require_backward_grad_sync = True
+            self._grad_div_ar_fusion = grad_div_ar_fusion
+            self._grad_allreduce_chunk_size_mb = grad_allreduce_chunk_size_mb
+        else:
+            self._require_backward_grad_sync = False
+            self._grad_div_ar_fusion = False
+            self._grad_allreduce_chunk_size_mb = 0
+
+        # Dummy tensor needed for apex multi-apply tensor.
+        self._dummy_overflow_buf = None
+
+        # Create persistent buffers for main gradients in contiguous memory space
+        # - Chunked element-wise and allreduce ops without creating a temporary buffer for merged operation
+        # - Low memory fragmentation
+        self._main_grad_buffers = None
+        if self._contiguous_grad_bucket:
+            self._main_grad_buffers = {}
+            # get the size of buffers
+            num_elements = {}
+            for i, param_group in enumerate(self.optimizer.param_groups):
+                num_elements[i] = sum(
+                    map(lambda x: x.data.nelement(), filter(lambda p: p.requires_grad, param_group['params']))
+                )
+
+                # Allocate gradient memory buffers for each data type
+                if num_elements[i] > 0:
+                    self._main_grad_buffers[i] = GradBucket(
+                        num_elements[i],
+                        self._grad_allreduce_chunk_size_mb,
+                        _get_grad_data_group(param_group.get('is_expert', False)),
+                    )
+
+        # Three groups of parameters:
+        self.float16_groups = []  # original float16 parameters
+        self.fp32_from_float16_groups = []  # fp32 copy of float16 parameters
+        self.fp32_from_fp32_groups = []  # original fp32 parameters
+
+        # gradient function hooks
+        if self._fp32_grad_accum:
+            self.grad_accs = []
+
+        # For all the groups in the original optimizer:
+        for i, param_group in enumerate(self.optimizer.param_groups):
+            float16_params_this_group = []
+            fp32_params_this_group = []
+            fp32_from_float16_params_this_group = []
+            # For all the parameters in this group:
+            is_expert_group = param_group.get('is_expert', False)
+            for j, param in enumerate(param_group['params']):
+                main_param = None
+                if param.requires_grad:
+                    # float16 params:
+                    if param.type() in ['torch.cuda.HalfTensor', 'torch.cuda.BFloat16Tensor']:
+                        float16_params_this_group.append(param)
+
+                        # Allocate the main parameter
+                        main_param = param.detach().clone().float()
+
+                        # Copy tensor model parallel attributes.
+                        copy_tensor_model_parallel_attributes(main_param, param)
+                        if hasattr(param, 'shared'):
+                            main_param.shared = param.shared
+                        if hasattr(param, 'allreduce'):
+                            main_param.allreduce = param.allreduce
+
+                        # Assign the grad buffer offset to main parameters
+                        if self._contiguous_grad_bucket:
+                            num_elements[i] -= param.data.nelement()
+                            main_param.grad, grad_chunk_info = self._main_grad_buffers[i].get(
+                                param.data.shape, num_elements[i]
+                            )
+                            # Add a pointer to main_grad in model param for first-last stage embedding param reduction
+                            param.main_grad = main_param.grad
+
+                        # Replace the optimizer params with the new fp32 copy.
+                        param_group['params'][j] = main_param
+                        fp32_from_float16_params_this_group.append(main_param)
+                        # Reset existing state dict key to the new main param.
+                        if param in self.optimizer.state:
+                            self.optimizer.state[main_param] = self.optimizer.state.pop(param)
+                    # fp32 params.
+                    elif param.type() == 'torch.cuda.FloatTensor':
+                        fp32_params_this_group.append(param)
+                        param_group['params'][j] = param
+
+                    else:
+                        raise TypeError(
+                            'Wrapped parameters must be one of '
+                            'torch.cuda.FloatTensor,  '
+                            'torch.cuda.HalfTensor, or '
+                            'torch.cuda.BFloat16Tensor. '
+                            'Received {}'.format(param.type())
+                        )
+
+                # Add gradient accumulation hook for fp32 grad accumulation
+                if main_param is not None and self._fp32_grad_accum and param.requires_grad:
+                    # Expand so we get access to grad_fn
+                    param_tmp = param.expand_as(param)
+                    # Get the gradient accumulator function.
+                    grad_acc = param_tmp.grad_fn.next_functions[0][0]
+                    grad_acc.register_hook(
+                        self._make_param_hook(param, main_param, i, grad_chunk_info, is_expert_group)
+                    )
+                    self.grad_accs.append(grad_acc)
+
+            self.float16_groups.append(float16_params_this_group)
+            self.fp32_from_float16_groups.append(fp32_from_float16_params_this_group)
+            self.fp32_from_fp32_groups.append(fp32_params_this_group)
+
+        # Leverage state_dict() and load_state_dict() to
+        # recast preexisting per-param state tensors
+        self.optimizer.load_state_dict(self.optimizer.state_dict())
+
+    def _make_param_hook(self, param, main_param, i, grad_chunk_info, is_expert_group):
+        """Create the grad accumulation and all-reduce hook for backprop."""
+
+        # Hook used for back-prop.
+        def param_hook(*unused):
+            # Accumulates gradients on main gradients
+            if param.grad is not None:
+                if main_param.grad is None:
+                    main_param.grad = param.grad.float()
+                else:
+                    main_param.grad.add_(param.grad.data)
+                # Deallocate grad memory.
+                param.grad = None
+
+            def allreduce_grads(use_fused_div, tensor, data_group, grad_mult):
+                if use_fused_div:
+                    torch.distributed.all_reduce(
+                        tensor,
+                        group=data_group,
+                        async_op=True,
+                        op=torch.distributed._make_nccl_premul_sum(1 / grad_mult),
+                    )
+                else:
+                    tensor.div_(grad_mult)
+                    torch.distributed.all_reduce(
+                        tensor,
+                        group=data_group,
+                        async_op=True,
+                    )
+
+            # Asynchronous gradients allreduce accross data_parallel ranks
+            grad_mult = get_data_parallel_world_size()
+            if self._require_backward_grad_sync:
+                data_group = _get_grad_data_group(is_expert_group)
+                if self._grad_allreduce_chunk_size_mb > 0:
+                    self._main_grad_buffers[i].update_chunk_info(grad_chunk_info)
+                    while True:
+                        allreduce_tensor = self._main_grad_buffers[i].get_allreduce_tensor()
+                        if allreduce_tensor is None:
+                            break
+                        allreduce_grads(self._grad_div_ar_fusion, allreduce_tensor, data_group, grad_mult)
+                else:
+                    allreduce_grads(self._grad_div_ar_fusion, main_param.grad, data_group, grad_mult)
+
+        return param_hook
+
+    def zero_grad(self, set_to_none=True):
+        """We only need to zero the model related parameters, i.e.,
+        float16_groups & fp32_from_fp32_groups. We additionally zero
+        fp32_from_float16_groups as a memory optimization to reduce
+        fragmentation; in the case of set_to_none==True, the space
+        used by this field can be safely deallocated at this point."""
+        for group in self.float16_groups:
+            _zero_grad_group_helper(group, set_to_none)
+        if self._contiguous_grad_bucket:
+            for i in self._main_grad_buffers:
+                self._main_grad_buffers[i].zero()
+        else:
+            for group in self.fp32_from_float16_groups:
+                _zero_grad_group_helper(group, set_to_none)
+        for group in self.fp32_from_fp32_groups:
+            _zero_grad_group_helper(group, set_to_none)
+
+    def copy_model_grads_to_main_grads(self):
+        # This only needs to be done for the float16 group.
+        for model_group, main_group in zip(self.float16_groups, self.fp32_from_float16_groups):
+            for model_param, main_param in zip(model_group, main_group):
+                if model_param.grad is not None:
+                    main_param.grad = model_param.grad.float()
+
+                # Safe to deallocate model's grad after copying.
+                # (If using contiguous buffers, main_grad's memory should
+                # persist and therefore should not be deallocated.)
+                model_param.grad = None
+
+    def _get_model_and_main_params_data_float16(self):
+        model_data = []
+        main_data = []
+        half_dtype = None
+        for model_group, main_group in zip(self.float16_groups, self.fp32_from_float16_groups):
+            for model_param, main_param in zip(model_group, main_group):
+                if half_dtype is None:
+                    half_dtype = model_param.data.dtype
+                model_data.append(model_param.data)
+                main_data.append(main_param.data)
+        return model_data, main_data, half_dtype
+
+    def _set_overflow_buffer(self, half_dtype):
+        if half_dtype == torch.float16:
+            if self._dummy_overflow_buf is None:
+                self._dummy_overflow_buf = torch.cuda.IntTensor([0])
+            else:
+                self._dummy_overflow_buf.fill_(0)
+
+    def _copy_main_params_to_model_params(self):
+        # Only needed for the float16 params.
+        model_data, main_data, half_dtype = self._get_model_and_main_params_data_float16()
+        self._set_overflow_buffer(half_dtype)
+        _multi_tensor_copy_this_to_that(this=main_data, that=model_data, overflow_buf=self._dummy_overflow_buf)
+
+    def _copy_model_params_to_main_params(self):
+        # Only needed for the float16 params.
+        model_data, main_data, half_dtype = self._get_model_and_main_params_data_float16()
+        self._set_overflow_buffer(half_dtype)
+        _multi_tensor_copy_this_to_that(this=model_data, that=main_data, overflow_buf=self._dummy_overflow_buf)
+
+    def reload_model_params(self):
+        self._copy_model_params_to_main_params()
+
+    @torch.no_grad()
+    def step(self, **kwargs):
+        # while async grad allreduce is enabled, bprop will keep moving forward without waiting for
+        # the finish of async grad AR works. Hence, to guarantee the correctness of grads reduction,
+        # we cannot start weight update until all async grad AR works are done.
+        if self._async_grad_allreduce:
+            torch.cuda.synchronize()
+
+        closure = kwargs.pop('closure', None)
+        # Allows applications to specify closure as None without erroring due to duplicate specification
+        assert closure is None, f"closure should be None but was passed {closure}"
+
+        # Step the optimizer.
+        self.optimizer.step(closure=None, **kwargs)
+
+        # Update params from main params.
+        with torch.no_grad():
+            self._copy_main_params_to_model_params()
+
+        # Successful update.
+        return True
+
+    def state_dict(self):
+        state_dict = {}
+        state_dict['optimizer'] = self.optimizer.state_dict()
+        state_dict['fp32_from_fp16_params'] = self.fp32_from_float16_groups
+        return state_dict
+
+    def load_state_dict(self, state_dict):
+        # Optimizer.
+        optimizer_key = 'optimizer'
+        if optimizer_key not in state_dict:
+            optimizer_key = 'optimizer_state_dict'
+            logging.info('***WARNING*** loading optimizer from ' 'an old checkpoint ...')
+        if 'state' not in state_dict[optimizer_key]:
+            state_dict[optimizer_key]['state'] = {}
+        self.optimizer.load_state_dict(state_dict[optimizer_key])
+
+        # Copy data for the main params.
+        fp32_from_float16_params_key = 'fp32_from_fp16_params'
+        if fp32_from_float16_params_key not in state_dict:
+            fp32_from_float16_params_key = 'fp32_from_fp16'
+        if fp32_from_float16_params_key not in state_dict:
+            state_dict[fp32_from_float16_params_key] = []
+        for current_group, saved_group in zip(self.fp32_from_float16_groups, state_dict[fp32_from_float16_params_key]):
+            for current_param, saved_param in zip(current_group, saved_group):
+                current_param.data.copy_(saved_param.data)
+
+    def allreduce_main_grads(self):
+        for i in self._main_grad_buffers:
+            self._main_grad_buffers[i].allreduce_buffer()
+
+    @contextmanager
+    def no_sync(self):
+        """A context manager to disable gradient synchronizations across
+        data-parallel ranks."""
+        old_require_backward_grad_sync = self._require_backward_grad_sync
+        self._require_backward_grad_sync = False
+        try:
+            yield
+        finally:
+            self._require_backward_grad_sync = old_require_backward_grad_sync
+
+    @property
+    def async_master_grads_allreudce(self):
+        return self._async_grad_allreduce
+
+    @property
+    def fp32_grad_accumulation(self):
+        return self._fp32_grad_accum
+
+    def get_parameters_with_grad(self):
+        params = []
+        for param_group in self.optimizer.param_groups:
+            for param in param_group['params']:
+                if param.grad is not None:  # (@adithyare) added to enable pp>1 training for adapters
+                    params.append(param)
+        return params
+
+    # Promote state so it can be retrieved or set via
+    # "optimizer_instance.state"
+    def _get_state(self):
+        if hasattr(self, 'optimizer'):
+            return self.optimizer.state
+        else:
+            return []
+
+    def _set_state(self, value):
+        self.optimizer.state = value
+
+    state = property(_get_state, _set_state)
+
+    # Promote param_groups so it can be retrieved or set via
+    # "optimizer_instance.param_groups"
+    # (for example, to adjust the learning rate)
+    def _get_param_groups(self):
+        if hasattr(self, 'optimizer'):
+            return self.optimizer.param_groups
+        else:
+            return []
+
+    def _set_param_groups(self, value):
+        self.optimizer.param_groups = value
+
+    param_groups = property(_get_param_groups, _set_param_groups)
+
+    # Promote defaults so it can be retrieved or set via
+    # "optimizer_instance.defaults
+    def _get_defaults(self):
+        if hasattr(self, 'optimizer'):
+            return self.optimizer.defaults
+        else:
+            return []
+
+    def _set_defaults(self, value):
+        self.optimizer.defaults = value
+
+    defaults = property(_get_defaults, _set_defaults)

nemo/core/optim/optimizers.py

@@ -0,0 +1,226 @@
+# 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
+from functools import partial
+from typing import Any, Dict, Optional, Union
+
+import torch
+import torch.optim as optim
+from omegaconf import DictConfig, OmegaConf
+from torch.optim import adadelta, adagrad, adamax, rmsprop, rprop
+from torch.optim.optimizer import Optimizer
+
+from nemo.core.classes.common import safe_instantiate
+from nemo.core.config.optimizers import OptimizerParams, get_optimizer_config, register_optimizer_params
+from nemo.core.optim.adafactor import Adafactor
+from nemo.core.optim.adan import Adan
+from nemo.core.optim.novograd import Novograd
+
+from nemo.utils.model_utils import maybe_update_config_version
+
+AVAILABLE_OPTIMIZERS = {
+    'sgd': optim.SGD,
+    'adam': optim.Adam,
+    'adamw': optim.AdamW,
+    'adadelta': adadelta.Adadelta,
+    'adamax': adamax.Adamax,
+    'adagrad': adagrad.Adagrad,
+    'rmsprop': rmsprop.RMSprop,
+    'rprop': rprop.Rprop,
+    'novograd': Novograd,
+    'adafactor': Adafactor,
+    'adan': Adan,
+}
+
+try:
+    from apex.optimizers import FusedAdam, FusedLAMB
+
+    HAVE_APEX = True
+
+    AVAILABLE_OPTIMIZERS['lamb'] = FusedLAMB
+    AVAILABLE_OPTIMIZERS['fused_adam'] = FusedAdam
+except ModuleNotFoundError:
+    HAVE_APEX = False
+
+HAVE_APEX_DISTRIBUTED_ADAM = False
+if HAVE_APEX:
+    try:
+        # Try importing wrapper for Apex distributed Adam optimizer
+        from nemo.core.optim.distributed_adam import MegatronDistributedFusedAdam
+
+        HAVE_APEX_DISTRIBUTED_ADAM = True
+
+        AVAILABLE_OPTIMIZERS['distributed_fused_adam'] = MegatronDistributedFusedAdam
+    except (ImportError, ModuleNotFoundError):
+        HAVE_APEX_DISTRIBUTED_ADAM = False
+
+    try:
+        # Try importing wrapper for Apex FusedAdam optimizer
+        from nemo.core.optim.megatron_fused_adam import MegatronFusedAdam
+
+        AVAILABLE_OPTIMIZERS['megatron_fused_adam'] = MegatronFusedAdam
+    except (ImportError, ModuleNotFoundError):
+        pass
+
+__all__ = ['get_optimizer', 'register_optimizer', 'parse_optimizer_args']
+
+
+def parse_optimizer_args(
+    optimizer_name: str, optimizer_kwargs: Union[DictConfig, Dict[str, Any]]
+) -> Union[Dict[str, Any], DictConfig]:
+    """
+    Parses a list of strings, of the format "key=value" or "key2=val1,val2,..."
+    into a dictionary of type {key=value, key2=[val1, val2], ...}
+
+    This dictionary is then used to instantiate the chosen Optimizer.
+
+    Args:
+        optimizer_name: string name of the optimizer, used for auto resolution of params
+        optimizer_kwargs: Either a list of strings in a specified format,
+            or a dictionary. If a dictionary is provided, it is assumed the dictionary
+            is the final parsed value, and simply returned.
+            If a list of strings is provided, each item in the list is parsed into a
+            new dictionary.
+
+    Returns:
+        A dictionary
+    """
+    kwargs = {}
+
+    if optimizer_kwargs is None:
+        return kwargs
+
+    optimizer_kwargs = copy.deepcopy(optimizer_kwargs)
+    optimizer_kwargs = maybe_update_config_version(optimizer_kwargs)
+
+    if isinstance(optimizer_kwargs, DictConfig):
+        optimizer_kwargs = OmegaConf.to_container(optimizer_kwargs, resolve=True)
+
+    # If it is a dictionary, perform stepwise resolution
+    if hasattr(optimizer_kwargs, 'keys'):
+        # Attempt class path resolution
+        if '_target_' in optimizer_kwargs:  # captures (target, _target_)
+            optimizer_kwargs_config = OmegaConf.create(optimizer_kwargs)
+            optimizer_instance = safe_instantiate(optimizer_kwargs_config)  # type: DictConfig
+            optimizer_instance = vars(optimizer_instance)
+            return optimizer_instance
+
+        # If class path was not provided, perhaps `name` is provided for resolution
+        if 'name' in optimizer_kwargs:
+            # If `auto` is passed as name for resolution of optimizer name,
+            # then lookup optimizer name and resolve its parameter config
+            if optimizer_kwargs['name'] == 'auto':
+                optimizer_params_name = "{}_params".format(optimizer_name)
+                optimizer_kwargs.pop('name')
+            else:
+                optimizer_params_name = optimizer_kwargs.pop('name')
+
+            # Override arguments provided in the config yaml file
+            if 'params' in optimizer_kwargs:
+                # If optimizer kwarg overrides are wrapped in yaml `params`
+                optimizer_params_override = optimizer_kwargs.get('params')
+            else:
+                # If the kwargs themselves are a DictConfig
+                optimizer_params_override = optimizer_kwargs
+
+            if isinstance(optimizer_params_override, DictConfig):
+                optimizer_params_override = OmegaConf.to_container(optimizer_params_override, resolve=True)
+
+            optimizer_params_cls = get_optimizer_config(optimizer_params_name, **optimizer_params_override)
+
+            # If we are provided just a Config object, simply return the dictionary of that object
+            if optimizer_params_name is None:
+                optimizer_params = vars(optimizer_params_cls)
+                return optimizer_params
+
+            else:
+                # If we are provided a partial class instantiation of a Config,
+                # Instantiate it and retrieve its vars as a dictionary
+                optimizer_params = optimizer_params_cls()  # instantiate the parameters object
+                optimizer_params = vars(optimizer_params)
+                return optimizer_params
+
+        # simply return the dictionary that was provided
+        return optimizer_kwargs
+
+    return kwargs
+
+
+def register_optimizer(name: str, optimizer: Optimizer, optimizer_params: OptimizerParams):
+    """
+    Checks if the optimizer name exists in the registry, and if it doesnt, adds it.
+
+    This allows custom optimizers to be added and called by name during instantiation.
+
+    Args:
+        name: Name of the optimizer. Will be used as key to retrieve the optimizer.
+        optimizer: Optimizer class
+        optimizer_params: The parameters as a dataclass of the optimizer
+    """
+    if name in AVAILABLE_OPTIMIZERS:
+        raise ValueError(f"Cannot override pre-existing optimizers. Conflicting optimizer name = {name}")
+
+    AVAILABLE_OPTIMIZERS[name] = optimizer
+
+    optim_name = "{}_params".format(optimizer.__name__)
+    register_optimizer_params(name=optim_name, optimizer_params=optimizer_params)
+
+
+def get_optimizer(name: str, **kwargs: Optional[Dict[str, Any]]) -> Optimizer:
+    """
+    Convenience method to obtain an Optimizer class and partially instantiate it with optimizer kwargs.
+
+    Args:
+        name: Name of the Optimizer in the registry.
+        kwargs: Optional kwargs of the optimizer used during instantiation.
+
+    Returns:
+        a partially instantiated Optimizer
+    """
+    if name not in AVAILABLE_OPTIMIZERS:
+        raise ValueError(
+            f"Cannot resolve optimizer '{name}'. Available optimizers are : " f"{AVAILABLE_OPTIMIZERS.keys()}"
+        )
+    if name == 'fused_adam':
+        if not torch.cuda.is_available():
+            raise ValueError('CUDA must be available to use fused_adam.')
+
+    optimizer = AVAILABLE_OPTIMIZERS[name]
+    optimizer = partial(optimizer, **kwargs)
+    return optimizer
+
+
+def init_optimizer_states(optimizer: Optimizer):
+    """
+    Initialize optimizer states for Adam-based optimizers.
+
+    This function initializes the exponential moving averages (exp_avg and exp_avg_sq)
+    for Adam, AdamW, and FusedAdam optimizers if they haven't been initialized yet.
+
+    Args:
+        optimizer: The optimizer instance to initialize states for
+    """
+    adam_nondist_optims = (optim.Adam, optim.AdamW)
+    if HAVE_APEX:
+        adam_nondist_optims += (FusedAdam,)
+    if isinstance(optimizer, adam_nondist_optims):
+        for group in optimizer.param_groups:
+            for p in group['params']:
+                state = optimizer.state[p]
+                if len(state) == 0:
+                    state['exp_avg'] = torch.zeros_like(p.data, memory_format=torch.preserve_format)
+                    state['exp_avg_sq'] = torch.zeros_like(p.data, memory_format=torch.preserve_format)
+                    if group.get('amsgrad'):
+                        state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)

nemo/core/optim/radam.py

@@ -0,0 +1,129 @@
+# 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.
+"""RAdam
+Original source taken from https://github.com/LiyuanLucasLiu/RAdam
+
+Copyright 2019 Liyuan Liu
+
+   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
+
+import torch
+from torch.optim.optimizer import Optimizer
+
+
+class RAdam(Optimizer):
+    """RAdam optimizer"""
+
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
+        """
+        Init
+
+        :param params: parameters to optimize
+        :param lr: learning rate
+        :param betas: beta
+        :param eps: numerical precision
+        :param weight_decay: weight decay weight
+        """
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+        self.buffer = [[None, None, None] for _ in range(10)]
+        super().__init__(params, defaults)
+
+    def step(self, closure=None):
+
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data.float()
+                if grad.is_sparse:
+                    raise RuntimeError('RAdam does not support sparse gradients')
+
+                p_data_fp32 = p.data.float()
+
+                state = self.state[p]
+
+                if len(state) == 0:
+                    state['step'] = 0
+                    state['exp_avg'] = torch.zeros_like(p_data_fp32)
+                    state['exp_avg_sq'] = torch.zeros_like(p_data_fp32)
+                else:
+                    state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)
+                    state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                beta1, beta2 = group['betas']
+
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=(1.0 - beta2))
+                exp_avg.mul_(beta1).add_(grad, alpha=(1.0 - beta1))
+
+                state['step'] += 1
+                buffered = self.buffer[int(state['step'] % 10)]
+                if state['step'] == buffered[0]:
+                    N_sma, step_size = buffered[1], buffered[2]
+                else:
+                    buffered[0] = state['step']
+                    beta2_t = beta2 ** state['step']
+                    N_sma_max = 2 / (1 - beta2) - 1
+                    N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t)
+                    buffered[1] = N_sma
+
+                    # more conservative since it's an approximated value
+                    if N_sma >= 5:
+                        step_size = (
+                            group['lr']
+                            * math.sqrt(
+                                (1 - beta2_t)
+                                * (N_sma - 4)
+                                / (N_sma_max - 4)
+                                * (N_sma - 2)
+                                / N_sma
+                                * N_sma_max
+                                / (N_sma_max - 2)
+                            )
+                            / (1 - beta1 ** state['step'])
+                        )
+                    else:
+                        step_size = group['lr'] / (1 - beta1 ** state['step'])
+                    buffered[2] = step_size
+
+                if group['weight_decay'] != 0:
+                    p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32)
+
+                # more conservative since it's an approximated value
+                if N_sma >= 5:
+                    denom = exp_avg_sq.sqrt().add_(group['eps'])
+                    p_data_fp32.addcdiv_(-step_size, exp_avg, denom)
+                else:
+                    p_data_fp32.add_(-step_size, exp_avg)
+
+                p.data.copy_(p_data_fp32)
+
+        return loss

nemo/core/utils/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+
+from nemo.core.utils import numba_utils, process_launcher

nemo/core/utils/cuda_python_utils.py

@@ -0,0 +1,275 @@
+# 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 contextlib
+
+import numpy as np
+import torch
+from packaging.version import Version
+
+from nemo.core.utils.optional_libs import CUDA_PYTHON_AVAILABLE, cuda_python_required
+from nemo.utils.exceptions import NeMoBaseException
+
+if CUDA_PYTHON_AVAILABLE:
+    from cuda.bindings import __version__ as cuda_python_version
+    from cuda.bindings import driver as cuda
+    from cuda.bindings import nvrtc
+    from cuda.bindings import runtime as cudart
+
+__CUDA_PYTHON_MINIMUM_VERSION_CUDA_GRAPH_CONDITIONAL_NODES_SUPPORTED__ = (12, 6)  # 12060
+
+
+class NeMoCUDAPythonException(NeMoBaseException):
+    """Exception caused by python-cuda in NeMo"""
+
+    pass
+
+
+CUDA_GRAPH_COMPILE_ERROR_TYPES = (NeMoCUDAPythonException,)
+if hasattr(torch, "AcceleratorError"):
+    CUDA_GRAPH_COMPILE_ERROR_TYPES += (torch.AcceleratorError,)
+
+
+def check_cuda_python_cuda_graphs_conditional_nodes_supported():
+    """Check if CUDA and CUDA-Python are available with CUDA Graphs with conditional nodes support"""
+    # for CPU-only environment we need to raise an exception, otherwise cuda-python library will fail
+    if not torch.cuda.is_available():
+        raise EnvironmentError("CUDA is not available")
+
+    try:
+        from cuda.bindings import driver as cuda
+    except ImportError:
+        raise ModuleNotFoundError("No `cuda-python` module. Please do `pip install cuda-python>=12.3`")
+
+    from cuda.bindings import __version__ as cuda_python_version
+
+    if Version(cuda_python_version) < Version("12.3.0"):
+        raise ImportError(f"Found cuda-python {cuda_python_version}, but at least version 12.3.0 is needed.")
+
+    error, driver_version = cuda.cuDriverGetVersion()
+    if error != cuda.CUresult.CUDA_SUCCESS:
+        raise ImportError(f"cuDriverGetVersion() returned {cuda.cuGetErrorString(error)}")
+
+    driver_version_major = driver_version // 1000
+    driver_version_minor = (driver_version % 1000) // 10
+
+    driver_version = (driver_version_major, driver_version_minor)
+    if driver_version < __CUDA_PYTHON_MINIMUM_VERSION_CUDA_GRAPH_CONDITIONAL_NODES_SUPPORTED__:
+        required_version = __CUDA_PYTHON_MINIMUM_VERSION_CUDA_GRAPH_CONDITIONAL_NODES_SUPPORTED__
+        raise ImportError(
+            f"""Driver supports cuda toolkit version \
+{driver_version_major}.{driver_version_minor}, but the driver needs to support \
+at least {required_version[0]},{required_version[1]}. Please update your cuda driver."""
+        )
+
+
+def skip_cuda_python_test_if_cuda_graphs_conditional_nodes_not_supported():
+    """
+    Helper method to skip pytest test case if cuda graph conditionals nodes are not supported.
+    """
+    try:
+        check_cuda_python_cuda_graphs_conditional_nodes_supported()
+    except (ImportError, ModuleNotFoundError, EnvironmentError) as e:
+        import pytest
+
+        pytest.skip(
+            "Test using cuda graphs with conditional nodes is being skipped because "
+            f"cuda graphs with conditional nodes aren't supported. Error message: {e}"
+        )
+
+
+@cuda_python_required
+def assert_drv(err):
+    """
+    Throws an exception if the return value of a cuda-python call is not success.
+    """
+    if isinstance(err, cuda.CUresult):
+        if err != cuda.CUresult.CUDA_SUCCESS:
+            raise NeMoCUDAPythonException("Cuda Error: {}".format(err))
+    elif isinstance(err, nvrtc.nvrtcResult):
+        if err != nvrtc.nvrtcResult.NVRTC_SUCCESS:
+            raise NeMoCUDAPythonException("Nvrtc Error: {}".format(err))
+    elif isinstance(err, cudart.cudaError_t):
+        if err != cudart.cudaError_t.cudaSuccess:
+            raise NeMoCUDAPythonException("Cuda Runtime Error: {}".format(err))
+    else:
+        raise NeMoCUDAPythonException("Unknown error type: {}".format(err))
+
+
+@cuda_python_required
+def cu_call(f_call_out):
+    """
+    Makes calls to cuda-python's functions inside cuda.cuda more python by throwing an exception
+    if they return a status which is not cudaSuccess
+    """
+    error, *others = f_call_out
+    if error != cudart.cudaError_t.cudaSuccess:
+        raise NeMoCUDAPythonException(f"CUDA failure! {error}")
+    else:
+        return tuple(others)
+
+
+@contextlib.contextmanager
+@cuda_python_required
+def with_conditional_node(while_loop_kernel, while_loop_args, while_loop_conditional_handle, device):
+    """
+    Even though we add a conditional node only once, we need to
+    capture the kernel that calls cudaGraphSetConditional() both
+    before in the parent graph containing the while loop body graph
+    and after the rest of the while loop body graph (because we need
+    to decide both whether to enter the loop, and also whether to
+    execute the next iteration of the loop).
+    """
+    # NB: depending on cuda-python version, cudaStreamGetCaptureInfo can return either 5 or 6 elements
+    capture_status, _, graph, *_ = cu_call(
+        cudart.cudaStreamGetCaptureInfo(torch.cuda.current_stream(device=device).cuda_stream)
+    )
+    assert capture_status == cudart.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive
+
+    cuda.cuLaunchKernel(
+        while_loop_kernel,
+        1,
+        1,
+        1,
+        1,
+        1,
+        1,
+        0,
+        torch.cuda.current_stream(device=device).cuda_stream,
+        while_loop_args.ctypes.data,
+        0,
+    )
+
+    # NB: depending on cuda-python version, cudaStreamGetCaptureInfo can return either 5 or 6 elements
+    capture_status, _, graph, dependencies, *_ = cu_call(
+        cudart.cudaStreamGetCaptureInfo(torch.cuda.current_stream(device=device).cuda_stream)
+    )
+    assert capture_status == cudart.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive
+
+    driver_params = cuda.CUgraphNodeParams()
+    driver_params.type = cuda.CUgraphNodeType.CU_GRAPH_NODE_TYPE_CONDITIONAL
+    driver_params.conditional.handle = while_loop_conditional_handle
+    driver_params.conditional.type = cuda.CUgraphConditionalNodeType.CU_GRAPH_COND_TYPE_WHILE
+    driver_params.conditional.size = 1
+    if Version(cuda_python_version) == Version("12.3.0"):
+        # Work around for https://github.com/NVIDIA/cuda-python/issues/55
+        # Originally, cuda-python version 12.3.0 failed to allocate phGraph_out
+        # on its own.
+        # This bug is fixed in cuda-python version 12.4.0. In fact, we can
+        # no longer write to phGraph_out in cuda-python 12.4.0, so we must
+        # condition on the version number.
+        driver_params.conditional.phGraph_out = [cuda.CUgraph()]
+    (ctx,) = cu_call(cuda.cuCtxGetCurrent())
+    driver_params.conditional.ctx = ctx
+
+    # Use driver API here because of bug in cuda-python runtime API: https://github.com/NVIDIA/cuda-python/issues/55
+    # TODO: Change call to this after fix goes in (and we bump minimum cuda-python version to 12.4.0):
+    # node, = cu_call(cudart.cudaGraphAddNode(graph, dependencies, len(dependencies), driver_params))
+    # CUDA 13 (cuda-python >= 13.0.0) adds an edgeData parameter to cuGraphAddNode and
+    # cudaStreamUpdateCaptureDependencies; CUDA 12 does not accept it.
+    _cuda13 = Version(cuda_python_version) >= Version("13.0.0")
+    if _cuda13:
+        (node,) = cu_call(cuda.cuGraphAddNode(graph, dependencies, None, len(dependencies), driver_params))
+    else:
+        (node,) = cu_call(cuda.cuGraphAddNode(graph, dependencies, len(dependencies), driver_params))
+    body_graph = driver_params.conditional.phGraph_out[0]
+
+    if _cuda13:
+        cu_call(
+            cudart.cudaStreamUpdateCaptureDependencies(
+                torch.cuda.current_stream(device=device).cuda_stream,
+                [node],
+                None,
+                1,
+                cudart.cudaStreamUpdateCaptureDependenciesFlags.cudaStreamSetCaptureDependencies,
+            )
+        )
+    else:
+        cu_call(
+            cudart.cudaStreamUpdateCaptureDependencies(
+                torch.cuda.current_stream(device=device).cuda_stream,
+                [node],
+                1,
+                cudart.cudaStreamUpdateCaptureDependenciesFlags.cudaStreamSetCaptureDependencies,
+            )
+        )
+    body_stream = torch.cuda.Stream(device)
+    previous_stream = torch.cuda.current_stream(device=device)
+    body_capture_active = False
+    cu_call(
+        cudart.cudaStreamBeginCaptureToGraph(
+            body_stream.cuda_stream,
+            body_graph,
+            None,
+            None,
+            0,
+            cudart.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal,
+        )
+    )
+    body_capture_active = True
+
+    try:
+        torch.cuda.set_stream(body_stream)
+        yield body_stream, body_graph
+
+        cuda.cuLaunchKernel(
+            while_loop_kernel, 1, 1, 1, 1, 1, 1, 0, body_stream.cuda_stream, while_loop_args.ctypes.data, 0
+        )
+
+        end_capture_out = cudart.cudaStreamEndCapture(body_stream.cuda_stream)
+        body_capture_active = False
+        cu_call(end_capture_out)
+    finally:
+        if body_capture_active:
+            try:
+                end_capture_out = cudart.cudaStreamEndCapture(body_stream.cuda_stream)
+                body_capture_active = False
+                cu_call(end_capture_out)
+            except Exception:
+                pass
+        torch.cuda.set_stream(previous_stream)
+
+
+@cuda_python_required
+def run_nvrtc(kernel_string: str, kernel_name: bytes, program_name: bytes):
+    """Run CUDA kernel using CUDA-Python"""
+    err, prog = nvrtc.nvrtcCreateProgram(str.encode(kernel_string), program_name, 0, [], [])
+    assert_drv(err)
+    # Compile program
+    # Not specifying --gpu-architecture will default us to a fairly low compute capability, which is a safe bet.
+    # Otherwise, there are ways to query the current device's compute capability.
+    # https://stackoverflow.com/questions/48283009/nvcc-get-device-compute-capability-in-runtime
+    opts = []
+    (err,) = nvrtc.nvrtcCompileProgram(prog, len(opts), opts)
+    assert_drv(err)
+    err, size = nvrtc.nvrtcGetProgramLogSize(prog)
+    assert_drv(err)
+    buf = b" " * size
+    (err,) = nvrtc.nvrtcGetProgramLog(prog, buf)
+    assert_drv(err)
+
+    # Get PTX from compilation
+    err, ptxSize = nvrtc.nvrtcGetPTXSize(prog)
+    assert_drv(err)
+    ptx = b" " * ptxSize
+    (err,) = nvrtc.nvrtcGetPTX(prog, ptx)
+    assert_drv(err)
+
+    ptx = np.char.array(ptx)
+    err, module = cuda.cuModuleLoadData(ptx.ctypes.data)
+    assert_drv(err)
+    err, kernel = cuda.cuModuleGetFunction(module, kernel_name)
+    assert_drv(err)
+
+    return kernel

nemo/core/utils/k2_guard.py

@@ -0,0 +1,53 @@
+# 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.
+
+"""
+Guard for importing optional NeMo dependency k2.
+Contains checks for k2 availability and version.
+Use `from nemo.core.utils.k2_guard import k2` to import k2 instead of direct import.
+If there is an error, the module will raise an exception with a helpful message.
+"""
+
+import textwrap
+
+from lightning.pytorch.utilities.imports import package_available
+from packaging.version import Version
+
+from nemo.core.utils.k2_utils import K2_INSTALLATION_MESSAGE
+
+__K2_MINIMUM_MAJOR_VERSION = 1
+__K2_MINIMUM_MINOR_VERSION = 14
+
+__K2_MINIMUM_VERSION = Version(f"{__K2_MINIMUM_MAJOR_VERSION}.{__K2_MINIMUM_MINOR_VERSION}")
+
+
+if not package_available("k2"):
+    raise ModuleNotFoundError("Module k2 is not available.\n" + K2_INSTALLATION_MESSAGE)
+
+import k2  # noqa: E402
+
+try:
+    __k2_version = Version(k2.__dev_version__)
+except AttributeError:
+    raise ImportError("Module k2 is corrupted.\n" + K2_INSTALLATION_MESSAGE)
+
+if __k2_version < __K2_MINIMUM_VERSION:
+    raise ImportError(
+        textwrap.dedent(
+            f"""
+            Minimum required k2 version: {__K2_MINIMUM_VERSION};
+            Installed k2 version: {__k2_version}
+            """
+        )
+    )

nemo/core/utils/k2_utils.py

@@ -0,0 +1,24 @@
+# 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.
+
+
+K2_INSTALLATION_MESSAGE = (
+    "Could not import `k2`.\n"
+    "Please install k2 in one of the following ways:\n"
+    "1) (recommended) Run `bash scripts/installers/install_k2.sh`\n"
+    "2) Use any approach from https://k2-fsa.github.io/k2/installation/index.html "
+    "if your your cuda and pytorch versions are supported.\n"
+    "It is advised to always install k2 using setup.sh only, "
+    "as different versions of k2 may not interact with the NeMo code as expected."
+)

nemo/core/utils/neural_type_utils.py

@@ -0,0 +1,82 @@
+# 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 collections import defaultdict
+from typing import Dict, List, Optional
+import torch
+from nemo.core.neural_types import AxisKind, NeuralType
+
+
+def get_io_names(types: Optional[Dict[str, NeuralType]], disabled_names: List[str]) -> List[str]:
+    if types is None:
+        return []
+    names = list(types.keys())
+    for name in disabled_names:
+        if name in names:
+            names.remove(name)
+    return names
+
+
+def extract_dynamic_axes(name: str, ntype: NeuralType):
+    """
+    This method will extract BATCH and TIME dimension ids from each provided input/output name argument.
+
+    For example, if module/model accepts argument named "input_signal" with type corresponding to [Batch, Time, Dim]
+    shape, then the returned result should contain "input_signal" -> [0, 1] because Batch and Time are dynamic axes
+    as they can change from call to call during inference.
+
+    Args:
+        name: Name of input or output parameter
+        ntype: Corresponding Neural Type
+
+    Returns:
+
+    """
+
+    def unpack_nested_neural_type(neural_type):
+        if type(neural_type) in (list, tuple):
+            return unpack_nested_neural_type(neural_type[0])
+        return neural_type
+
+    dynamic_axes = defaultdict(list)
+    if type(ntype) in (list, tuple):
+        ntype = unpack_nested_neural_type(ntype)
+
+    if ntype.axes:
+        for ind, axis in enumerate(ntype.axes):
+            if axis.kind in [AxisKind.Batch, AxisKind.Time, AxisKind.Width, AxisKind.Height]:
+                dynamic_axes[name].append(ind)
+    return dynamic_axes
+
+
+def get_dynamic_axes(types, names, use_dynamo=False):
+    dynamic_axes = defaultdict(list)
+    if names is not None:
+        for name in names:
+            if name in types:
+                dynamic_axes.update(extract_dynamic_axes(name, types[name]))
+    if use_dynamo:
+        dynamic_shapes = {}
+        batch = torch.export.Dim("batch")
+        for name, dims in dynamic_axes.items():
+            ds = {}
+            for d in dims:
+                if d == 0:
+                    ds[d] = batch
+                # this currently has issues: https://github.com/pytorch/pytorch/issues/126127
+                else:
+                    ds[d] = torch.export.Dim(name + '__' + str(d))
+            dynamic_shapes[name] = ds
+        dynamic_axes = dynamic_shapes
+    return dynamic_axes

nemo/core/utils/numba_utils.py

@@ -0,0 +1,195 @@
+# 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 contextlib
+import logging as pylogger
+import operator
+import os
+
+from typing import Tuple, Union
+
+from nemo.utils import model_utils
+
+# Prevent Numba CUDA logs from showing at info level
+cuda_logger = pylogger.getLogger('numba.cuda.cudadrv.driver')
+cuda_logger.setLevel(pylogger.ERROR)  # only show error
+
+__NUMBA_DEFAULT_MINIMUM_VERSION__ = "0.53.0"
+__NUMBA_MINIMUM_VERSION__ = os.environ.get("NEMO_NUMBA_MINVER", __NUMBA_DEFAULT_MINIMUM_VERSION__)
+
+__NUMBA_MINIMUM_VERSION_FP16_SUPPORTED__ = "0.57.0"
+
+
+NUMBA_INSTALLATION_MESSAGE = (
+    "Could not import `numba`.\n"
+    "Please install numba in one of the following ways."
+    "1) If using conda, simply install it with conda using `conda install -c numba numba`\n"
+    "2) If using pip (not recommended), `pip install --upgrade numba`\n"
+    "followed by `export NUMBAPRO_LIBDEVICE='/usr/local/cuda/nvvm/libdevice/'` and \n"
+    "`export NUMBAPRO_NVVM='/usr/local/cuda/nvvm/lib64/libnvvm.so'`.\n"
+    "It is advised to always install numba using conda only, "
+    "as pip installations might interfere with other libraries such as llvmlite.\n"
+    "If pip install does not work, you can also try adding `--ignore-installed` to the pip command,\n"
+    "but this is not advised."
+)
+
+STRICT_NUMBA_COMPAT_CHECK = True
+
+# Get environment key if available
+if 'STRICT_NUMBA_COMPAT_CHECK' in os.environ:
+    check_str = os.environ.get('STRICT_NUMBA_COMPAT_CHECK')
+    check_bool = str(check_str).lower() in ("yes", "true", "t", "1")
+    STRICT_NUMBA_COMPAT_CHECK = check_bool
+
+
+def is_numba_compat_strict() -> bool:
+    """
+    Returns strictness level of numba cuda compatibility checks.
+
+    If value is true, numba cuda compatibility matrix must be satisfied.
+    If value is false, only cuda availability is checked, not compatibility.
+    Numba Cuda may still compile and run without issues in such a case, or it may fail.
+    """
+    return STRICT_NUMBA_COMPAT_CHECK
+
+
+def set_numba_compat_strictness(strict: bool):
+    """
+    Sets the strictness level of numba cuda compatibility checks.
+
+    If value is true, numba cuda compatibility matrix must be satisfied.
+    If value is false, only cuda availability is checked, not compatibility.
+    Numba Cuda may still compile and run without issues in such a case, or it may fail.
+
+    Args:
+        strict: bool value, whether to enforce strict compatibility checks or relax them.
+    """
+    global STRICT_NUMBA_COMPAT_CHECK
+    STRICT_NUMBA_COMPAT_CHECK = strict
+
+
+@contextlib.contextmanager
+def with_numba_compat_strictness(strict: bool):
+    """Context manager for setting numba compatibility checks temporary"""
+    initial_strictness = is_numba_compat_strict()
+    set_numba_compat_strictness(strict=strict)
+    yield
+    set_numba_compat_strictness(strict=initial_strictness)
+
+
+def numba_cpu_is_supported(min_version: str) -> bool:
+    """
+    Tests if an appropriate version of numba is installed.
+
+    Args:
+        min_version: The minimum version of numba that is required.
+
+    Returns:
+        bool, whether numba CPU supported with this current installation or not.
+    """
+    module_available, msg = model_utils.check_lib_version('numba', checked_version=min_version, operator=operator.ge)
+
+    # If numba is not installed
+    if module_available is None:
+        return False
+    else:
+        return True
+
+
+def numba_cuda_is_supported(min_version: str) -> bool:
+    """
+    Tests if an appropriate version of numba is installed, and if it is,
+    if cuda is supported properly within it.
+
+    Args:
+        min_version: The minimum version of numba that is required.
+
+    Returns:
+        bool, whether cuda is supported with this current installation or not.
+    """
+    module_available = numba_cpu_is_supported(min_version)
+
+    # If numba is not installed
+    if module_available is None:
+        return False
+
+    # If numba version is installed and available
+    if module_available is True:
+        from numba import cuda
+
+        try:
+            cuda_available = cuda.is_available()
+            if cuda_available:
+                cuda_compatible = cuda.cudadrv.runtime.get_version()[0] in (12, 13)
+            else:
+                cuda_compatible = False
+
+            if is_numba_compat_strict():
+                return cuda_available and cuda_compatible
+            else:
+                return cuda_available
+
+        except Exception:
+            # dlopen(libcudart.dylib) might fail if CUDA was never installed in the first place.
+            return False
+
+    else:
+        return False
+
+
+def is_numba_cuda_fp16_supported(return_reason: bool = False) -> Union[bool, Tuple[bool, str]]:
+    """
+    Utility method that returns a bool, stating if FP16 is supported for numba cuda kernels or not.
+
+    Returns:
+        bool, whether Numba CUDA will support fp16 or not.
+    """
+    reason = ""
+    use_nvidia_binding = os.environ.get('NUMBA_CUDA_USE_NVIDIA_BINDING', None)
+    if use_nvidia_binding is not None:
+        use_nvidia_binding = use_nvidia_binding.lower() == "1"
+        reason += "Env variable `NUMBA_CUDA_USE_NVIDIA_BINDING` is available and set to `1`. "
+    else:
+        use_nvidia_binding = False
+        reason += "Env variable `NUMBA_CUDA_USE_NVIDIA_BINDING` is not available or has not set to `1`."
+
+    numba_fp16_version_correct = model_utils.check_lib_version(
+        'numba', __NUMBA_MINIMUM_VERSION_FP16_SUPPORTED__, operator=operator.ge
+    )[0]
+
+    if numba_fp16_version_correct:
+        reason += "Numba CUDA FP16 is supported in installed numba version."
+    else:
+        reason += "Numba CUDA FP16 is not supported in installed numba version."
+
+    result = use_nvidia_binding and numba_fp16_version_correct
+
+    if return_reason:
+        return result, reason
+    else:
+        return result
+
+
+def skip_numba_cuda_test_if_unsupported(min_version: str):
+    """
+    Helper method to skip pytest test case if numba cuda is not supported.
+
+    Args:
+        min_version: The minimum version of numba that is required.
+    """
+    numba_cuda_support = numba_cuda_is_supported(min_version)
+    if not numba_cuda_support:
+        import pytest
+
+        pytest.skip(f"Numba cuda test is being skipped. Minimum version required : {min_version}")

nemo/core/utils/optional_libs.py

@@ -0,0 +1,137 @@
+# 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.
+
+"""
+Module with guards for optional libraries, that cannot be listed in `requirements.txt`.
+Provides helper constants and decorators to check if the library is available in the system.
+"""
+
+__all__ = [
+    # kenlm
+    "KENLM_AVAILABLE",
+    "kenlm_required",
+    # k2
+    "K2_AVAILABLE",
+    "k2_required",
+    # triton
+    "TRITON_AVAILABLE",
+    "triton_required",
+    # cuda-python
+    "CUDA_PYTHON_AVAILABLE",
+    "cuda_python_required",
+    # numba
+    "NUMBA_AVAILABLE",
+    "numba_required",
+    # numba-cuda
+    "NUMBA_CUDA_AVAILABLE",
+    "numba_cuda_required",
+]
+
+import importlib.util
+from functools import wraps
+
+from packaging.version import Version
+
+from nemo.core.utils.k2_utils import K2_INSTALLATION_MESSAGE
+from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__, numba_cpu_is_supported, numba_cuda_is_supported
+
+
+def is_lib_available(name: str) -> bool:
+    """
+    Checks if the library/package with `name` is available in the system
+    NB: try/catch with importlib.import_module(name) requires importing the library, which can be slow.
+    So, `find_spec` should be preferred
+    """
+    return importlib.util.find_spec(name) is not None
+
+
+KENLM_AVAILABLE = is_lib_available("kenlm")
+KENLM_INSTALLATION_MESSAGE = "Try installing kenlm with `pip install kenlm`"
+
+TRITON_AVAILABLE = is_lib_available("triton")
+TRITON_INSTALLATION_MESSAGE = "Try installing triton with `pip install triton`"
+
+NUMBA_AVAILABLE = numba_cpu_is_supported(__NUMBA_MINIMUM_VERSION__)
+NUMBA_INSTALLATION_MESSAGE = (
+    "Numba is not found. Install with `pip install numba`. "
+    "For GPU support install with `pip install numba-cuda[cu12]` or `pip install numba-cuda[cu13]`"
+)
+
+NUMBA_CUDA_AVAILABLE = numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__)
+NUMBA_CUDA_INSTALLATION_MESSAGE = (
+    "Numba with GPU support is not available. "
+    "For GPU support install with `pip install numba-cuda[cu12]` or `pip install numba-cuda[cu13]`"
+)
+
+try:
+    from nemo.core.utils.k2_guard import k2 as _  # noqa: F401
+
+    K2_AVAILABLE = True
+except (ImportError, ModuleNotFoundError):
+    K2_AVAILABLE = False
+
+try:
+    from cuda.bindings import __version__ as cuda_python_version
+
+    if Version(cuda_python_version) >= Version("12.6.0"):
+        CUDA_PYTHON_AVAILABLE = True
+    else:
+        CUDA_PYTHON_AVAILABLE = False
+except (ImportError, ModuleNotFoundError):
+    CUDA_PYTHON_AVAILABLE = False
+
+CUDA_PYTHON_INSTALLATION_MESSAGE = "Try installing cuda-python with `pip install cuda-python>=12.6.0`"
+
+
+def identity_decorator(f):
+    """Identity decorator for further using in conditional decorators"""
+    return f
+
+
+def _lib_required(is_available: bool, name: str, message: str | None = None):
+    """
+    Decorator factory. Returns identity decorator if lib `is_available`,
+    otherwise returns a decorator which returns a function that raises an error when called.
+    Such decorator can be used for conditional checks for optional libraries in functions and methods
+    with zero computational overhead.
+    """
+    if is_available:
+        return identity_decorator
+
+    # return wrapper that will raise an error when the function is called
+    def function_stub_with_error_decorator(f):
+        """Decorator that replaces the function and raises an error when called"""
+
+        @wraps(f)
+        def wrapper(*args, **kwargs):
+            error_msg = f"Module {name} required for the function {f.__name__} is not found."
+            if message:
+                error_msg += f" {message}"
+            raise ModuleNotFoundError(error_msg)
+
+        return wrapper
+
+    return function_stub_with_error_decorator
+
+
+kenlm_required = _lib_required(is_available=KENLM_AVAILABLE, name="kenlm", message=KENLM_INSTALLATION_MESSAGE)
+triton_required = _lib_required(is_available=TRITON_AVAILABLE, name="triton", message=TRITON_INSTALLATION_MESSAGE)
+k2_required = _lib_required(is_available=K2_AVAILABLE, name="k2", message=K2_INSTALLATION_MESSAGE)
+cuda_python_required = _lib_required(
+    is_available=CUDA_PYTHON_AVAILABLE, name="cuda_python", message=CUDA_PYTHON_INSTALLATION_MESSAGE
+)
+numba_required = _lib_required(is_available=NUMBA_AVAILABLE, name="numba", message=NUMBA_INSTALLATION_MESSAGE)
+numba_cuda_required = _lib_required(
+    is_available=NUMBA_CUDA_AVAILABLE, name="numba-cuda", message=NUMBA_CUDA_INSTALLATION_MESSAGE
+)

nemo/core/utils/process_launcher/__init__.py

@@ -0,0 +1,15 @@
+# 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 nemo.core.utils.process_launcher.launcher import ProcessLauncher, ProcessLauncherConfig

nemo/core/utils/process_launcher/launcher.py

@@ -0,0 +1,386 @@
+# 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 copy
+import hashlib
+import os
+import subprocess
+import sys
+import threading
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, Optional, Sequence
+
+import torch
+from hydra.core.config_store import ConfigStore
+from hydra.core.hydra_config import HydraConfig
+from hydra.core.plugins import Plugins
+from hydra.core.singleton import Singleton
+from hydra.core.utils import JobReturn, JobStatus, configure_log, filter_overrides, setup_globals
+from hydra.plugins.launcher import Launcher
+from hydra.types import HydraContext, TaskFunction
+from omegaconf import DictConfig, OmegaConf, open_dict
+
+from nemo.utils import logging
+
+
+# monkey-patch hydra func
+def is_in_toplevel_plugins_module(*args, **kwargs) -> bool:
+    """Treat NeMo launcher plugins as top-level Hydra plugins."""
+    return True
+
+
+# Monkey-patch Hydra
+Plugins.instance().is_in_toplevel_plugins_module = is_in_toplevel_plugins_module
+
+
+@dataclass
+class ProcessLauncherConfig:
+    """Configuration for the NeMo process launcher."""
+
+    _target_: str = "nemo.core.utils.process_launcher.launcher.ProcessLauncher"
+    num_gpus: int = -1
+    jobs_per_gpu: int = 1
+
+
+def execute_job(
+    idx: int,
+    overrides: Sequence[str],
+    hydra_context: HydraContext,
+    config: DictConfig,
+    singleton_state: Dict[Any, Any],
+    gpu_idx: int,
+):
+    """
+    Creates a process that launches a "single" job that is identical in config + updated with sweep hyperparams.
+    Since a different process is being used, CUDA can work in non-ddp mode without issue.
+    Attempting ddp when using this script will not work as ddp cannot be used in shared contexts.
+
+    Args:
+        idx: Global index of the job.
+        overrides: List of str overrides that correspond to this job
+        hydra_context: Hydra Context used to load the sweep params into the global config
+        config: Global config that will be updated with sweep hyper parameters.
+        singleton_state: Hydra state.
+        gpu_idx: The GPU ID on which this process will be run.
+
+    Returns:
+        - The Process object that corresponds to this sweep
+        - The JobReturn object holding some metadata about this run
+    """
+    # Required by Hydra (lookup other Hydra Launchers for details)
+    setup_globals()
+    Singleton.set_state(singleton_state)
+
+    # Update base config with overrides to create sweep config
+    sweep_config = hydra_context.config_loader.load_sweep_config(config, list(overrides))
+    with open_dict(sweep_config):
+        sweep_config.hydra.job.id = "{}_{}".format(sweep_config.hydra.job.name, idx)
+        sweep_config.hydra.job.num = idx
+    HydraConfig.instance().set_config(sweep_config)
+
+    # Setup a directory where the config will temporarily be stored.
+    script_path = os.path.join(os.getcwd(), sys.argv[0])
+    script_path = os.path.abspath(script_path)
+
+    hash_salt = "|".join([script_path, str(OmegaConf.to_yaml(config))]).encode('utf-8')
+    hash_val = hashlib.sha256(hash_salt).hexdigest()
+
+    config_dir = os.path.join(os.getcwd(), "hydra_cfg", str(hash_val))
+    if not os.path.exists(config_dir):
+        os.makedirs(config_dir, exist_ok=True)
+
+    task_cfg = copy.deepcopy(sweep_config)
+
+    # Remove hydra from sweep config
+    # This is done to prevent recursive call to multirun launcher in Hydra.
+    with open_dict(task_cfg):
+        task_cfg.pop('hydra', '')
+
+    # Save the current jobs config to directory
+    temp_config_name = f"config_{idx}.yaml"
+    temp_config = os.path.join(config_dir, temp_config_name)
+    OmegaConf.save(task_cfg, temp_config)
+
+    # Compute the overides as a dict
+    overrides = OmegaConf.to_container(config.hydra.overrides.task)
+
+    # Check and replace trainer.devices in config with gpu_idx
+    found_devices = False
+    gpu_override = f'trainer.devices=[{gpu_idx}]'
+    for oidx, val in enumerate(overrides):
+        if 'trainer.devices' in val:
+            overrides[oidx] = gpu_override
+            found_devices = True
+
+    if not found_devices:
+        overrides.append(gpu_override)
+
+    # Build launch command
+    # Note: We depend on PTL doing the right thing since this command has global visibility of all CUDA_VISIBLE_DEVICES
+    cmd = [
+        'python',
+        script_path,
+        "--config-path",
+        config_dir,
+        "--config-name",
+        temp_config_name,
+        *overrides,
+    ]
+
+    # Launch the subprocess; pipe the stderr
+    # NOTE: If this hangs due to some reason after prolonged training, it means that the stderr pipe buffer
+    # has become full at the OS level and we need to explicitly empty it (either parallel thread or manually
+    # call proc.communicate(). It should not happen in general case as stderr is filled only in case retcode != 0
+    # If it does happen though, implement the code here
+    # https://stackoverflow.com/questions/39607172/python-subprocess-popen-poll-seems-to-hang-but-communicate-works
+    proc = subprocess.Popen(cmd, stderr=subprocess.PIPE)
+
+    # Setup data thread for stderr
+    std_error_buffer = []
+    # Trivial thread just reads lines from stdout into the list
+    drainerthread = threading.Thread(target=std_error_buffer.extend, args=(proc.stderr,))
+    drainerthread.daemon = True
+    drainerthread.start()
+
+    # Construct JobReturn object for Hydra
+    res = JobReturn()
+    res.cfg = task_cfg
+    res.overrides = overrides
+    res.hydra_cfg = config
+    res.working_dir = os.getcwd()
+    res.return_value = None
+
+    return proc, res, (std_error_buffer, drainerthread)
+
+
+def launch(
+    launcher,
+    job_overrides: Sequence[Sequence[str]],
+    initial_job_idx: int,
+) -> Sequence[JobReturn]:
+    """
+    Args:
+        launcher: Reference to the Launched subclass
+        job_overrides: A List of List<String>, where each inner list is the arguments for one job run
+        initial_job_idx: Initial job idx in batch
+
+    Returns:
+        A list of JobReturn objects.
+    """
+    # Needed for Hydra, lookup JoblibLauncher in Hydra
+    setup_globals()
+    assert launcher.config is not None
+    assert launcher.task_function is not None
+    assert launcher.hydra_context is not None
+
+    configure_log(launcher.config.hydra.hydra_logging, launcher.config.hydra.verbose)
+    sweep_dir = Path(str(launcher.config.hydra.sweep.dir))
+    sweep_dir.mkdir(parents=True, exist_ok=True)
+
+    # Extraact the runner's config (its actually a DictConfig, but type is used for autocomplete)
+    runner_cfg = launcher.runner  # type: ProcessLauncherConfig
+
+    logging.info(
+        "ProcessLauncher({}) is launching {} jobs".format(
+            ",".join([f"{k}={v}" for k, v in runner_cfg.items()]),
+            len(job_overrides),
+        )
+    )
+    logging.info("Launching jobs, sweep output dir : {}".format(sweep_dir))
+    for idx, overrides in enumerate(job_overrides):
+        logging.info("\t#{} : {}".format(idx, " ".join(filter_overrides(overrides))))
+
+    # Needed by Hydra
+    singleton_state = Singleton.get_state()
+
+    # Process the runner's config to build up the multiplex config
+    num_gpus = runner_cfg.get('num_gpus', -1)
+    jobs_per_gpu = runner_cfg.get('jobs_per_gpu', 1)
+
+    # Only GPUs are supported for now.
+    if num_gpus <= 0:
+        if torch.cuda.is_available():
+            num_gpus = torch.cuda.device_count()
+        else:
+            raise ValueError(f"{launcher.__class__.__name__} only supports GPU operations.")
+
+    # Setup arguments for multiplex runner
+    overrides = list(job_overrides)
+    num_overrides = len(overrides)
+
+    job_idx = 0
+    batch_size = num_gpus * jobs_per_gpu
+    gpu_idx = 0
+
+    ret = []  # List of returned JobResult
+    subprocess_list = []  # Buffer of subprocess
+    results = []  # Buffer of JobResult
+
+    # STD ERROR cache
+    std_error_buffers = []  # type: List[List[str]]
+    std_error_threads = []  # type: threading.Thread
+
+    # Run over all job combinations
+    while job_idx < num_overrides:
+        # Fill up subprocess buffer while its size is smaller than multiplex batch size
+        while len(subprocess_list) < batch_size:
+            # If we run out of jobs, stop trying to submit more jobs
+            if job_idx >= num_overrides:
+                break
+
+            # Submit a job as a new process
+            process, res, error_tup = execute_job(
+                initial_job_idx + job_idx,
+                overrides[job_idx],
+                launcher.hydra_context,
+                launcher.config,
+                singleton_state,
+                gpu_idx % num_gpus,  # This will evenly distribute GPU load
+            )
+
+            # Store the subprocesses and JobResults
+            subprocess_list.append(process)
+            results.append(res)
+
+            # Manage stderror thread data
+            std_error_buffers.append(error_tup[0])
+            std_error_threads.append(error_tup[1])
+
+            job_idx += 1
+            gpu_idx += 1
+
+        # Poll for samples in batch to finish.
+        if len(subprocess_list) > 0:
+            finished_processes = [0] * len(subprocess_list)
+
+            # Check if all processes are completed or not
+            # This is busy waiting, this is actually quite necessary
+            # Turns out that when you do proc.communicate(), you block all other threads immediately.
+            # IE they may fill up their buffers entirely, and hang while they wait for the first thread
+            # who called communicate() to finish its work or crash.
+            # Effectively it entirely stops multiprocessing jobs or multiplexed runs.
+            # Must poll and busy wait to keep threads alive, along with drain the pipes with thread buffers.
+            while sum(finished_processes) < len(subprocess_list):
+                # Check all processes to make sure they have a retcode (doesnt matter yet if 0 or not)
+                for proc_idx, proc in enumerate(subprocess_list):
+                    # poll() is cheaper op than communicate()
+                    retcode = proc.poll()
+
+                    if retcode is not None:
+                        # Log that the process with some ID has finished
+                        if finished_processes[proc_idx] == 0:
+                            logging.info(f"Processed job : {len(ret) + proc_idx} :: Ret code = {retcode}")
+
+                        finished_processes[proc_idx] = 1
+
+                        # Join this thread and merge its stderror buffer
+                        proc.wait()
+                        std_error_threads[proc_idx].join()
+                        error_data = std_error_buffers[proc_idx]
+                        error_data = [
+                            str(data, encoding='utf-8').encode('utf-8').decode('utf-8').encode('utf-8')
+                            for data in error_data
+                        ]
+
+                        std_error_buffers[proc_idx] = error_data
+
+                time.sleep(1.0)
+
+            # Process all the subprocess results
+            for proc_idx, (proc, res) in enumerate(zip(subprocess_list, results)):
+                # Wait until completion of process
+                output, error = proc.communicate()
+
+                # 0 is for successful run
+                if proc.returncode == 0:
+                    res.status = JobStatus.COMPLETED
+                else:
+                    # > 0 is for error, log the error.
+                    # Note: For the sake of efficiency while we log the error and raise an exception,
+                    # It will only raise the 1st wrong job in all the jobs.
+                    # If multiple jobs fail, it will still try to execute every job first before
+                    # raising the error for the first one.
+                    # This is done so that even if some jobs fail (say OOM or something),
+                    # other jobs can still run.
+                    err_buffer = std_error_buffers[proc_idx]
+                    if isinstance(err_buffer, (list, tuple)):
+                        err_string = ""
+                        for err_line in err_buffer:
+                            err_string = (
+                                err_string + f"{str(err_line, encoding='utf-8').encode('utf-8').decode('utf-8')}"
+                            )
+                    else:
+                        err_string = err_buffer
+
+                    error_msg = (
+                        f"\nHyperparameter Arguments : {proc.args}\n"
+                        f"Process Return code : {proc.returncode}\n"
+                        f"Error Trace :\n"
+                        f"{err_string}"
+                    )
+                    res.return_value = Exception(error_msg)
+                    res.status = JobStatus.FAILED
+
+                logging.info(f"Finished executing job : {len(ret)}. Return Code = {proc.returncode}")
+                ret.append(res)
+
+            # Reset for next batch
+            subprocess_list.clear()
+            results.clear()
+
+    return ret
+
+
+class ProcessLauncher(Launcher):
+    """Hydra launcher that multiplexes jobs across local GPU processes."""
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Process Launcher
+        Based on the JoblibLauncher, but uses processes to scatter jobs in a multiplexed manner across
+        some number of GPUs on a single machine.
+        """
+        self.config: Optional[DictConfig] = None
+        self.task_function: Optional[TaskFunction] = None
+        self.hydra_context: Optional[HydraContext] = None
+
+        self.runner = kwargs  # type: ProcessLauncherConfig
+
+    def setup(
+        self,
+        *,
+        hydra_context: HydraContext,
+        task_function: TaskFunction,
+        config: DictConfig,
+    ) -> None:
+        """Store Hydra launch context and task function."""
+        self.config = config
+        self.task_function = task_function
+        self.hydra_context = hydra_context
+
+    def launch(self, job_overrides: Sequence[Sequence[str]], initial_job_idx: int) -> Sequence[JobReturn]:
+        """Launch jobs with the configured process launcher."""
+
+        return launch(launcher=self, job_overrides=job_overrides, initial_job_idx=initial_job_idx)
+
+
+ConfigStore.instance().store(
+    group="hydra/launcher",
+    name="nemo_launcher",
+    node=ProcessLauncherConfig,
+    provider="nemo_process_launcher",
+)
+
+Plugins.instance().register(ProcessLauncher)