chore: import upstream snapshot with attribution

docs/source/core/adapters/adapter_bib.bib

@@ -0,0 +1,21 @@
+
+
+@inproceedings{houlsby2019adapter,
+  title={Parameter-efficient transfer learning for NLP},
+  author={Houlsby, Neil and Giurgiu, Andrei and Jastrzebski, Stanislaw and Morrone, Bruna and De Laroussilhe, Quentin and Gesmundo, Andrea and Attariyan, Mona and Gelly, Sylvain},
+  booktitle={International Conference on Machine Learning},
+  pages={2790--2799},
+  year={2019},
+  organization={PMLR}
+}
+
+@misc{Junxian2021unified,
+  doi = {10.48550/ARXIV.2110.04366},
+  url = {https://arxiv.org/abs/2110.04366},
+  author = {He, Junxian and Zhou, Chunting and Ma, Xuezhe and Berg-Kirkpatrick, Taylor and Neubig, Graham},
+  keywords = {Computation and Language (cs.CL), Machine Learning (cs.LG), FOS: Computer and information sciences, FOS: Computer and information sciences},
+  title = {Towards a Unified View of Parameter-Efficient Transfer Learning},
+  publisher = {arXiv},
+  year = {2021},
+  copyright = {arXiv.org perpetual, non-exclusive license}
+}

docs/source/core/adapters/api.rst

@@ -0,0 +1,66 @@
+Adapters API
+============
+
+Core
+----
+
+.. autoclass:: nemo.core.adapter_mixins.AdapterModuleMixin
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+    :undoc-members: adapter_module_names
+    :noindex:
+
+
+.. autoclass:: nemo.core.adapter_mixins.AdapterModelPTMixin
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+    :undoc-members: adapter_module_names
+    :noindex:
+
+
+Adapter Networks
+----------------
+
+
+.. autoclass:: nemo.collections.common.parts.adapter_modules.AdapterModuleUtil
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+    :noindex:
+
+
+.. autoclass:: nemo.collections.common.parts.adapter_modules.LinearAdapter
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+    :noindex:
+
+
+Adapter Strategies
+------------------
+
+
+.. autoclass:: nemo.core.classes.mixins.adapter_mixin_strategies.AbstractAdapterStrategy
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+    :undoc-members: adapter_module_names
+    :noindex:
+
+
+.. autoclass:: nemo.core.classes.mixins.adapter_mixin_strategies.ReturnResultAdapterStrategy
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+    :undoc-members: adapter_module_names
+    :noindex:
+
+
+.. autoclass:: nemo.core.classes.mixins.adapter_mixin_strategies.ResidualAddAdapterStrategy
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+    :undoc-members: adapter_module_names
+    :noindex:

docs/source/core/adapters/components.rst

@@ -0,0 +1,94 @@
+Adapter Components
+==================
+
+Adapters can be considered as any set of parameters that are added to a pre-existing module/model. In our case, we currently support the standard adapter in literature, more advanced adapter modules are being researched and can potentially be supported by NeMo.
+
+An adapter module can be any pytorch module, but it must follow certain straightforward requirements -
+
+1) The model accepts an input of some input dimension, and its output must match this dimension.
+2) Ideally, the module is initialized such that the output of the adapter when initialized is such that it does not modify the original input. This allows the model to produce the same output results, even when additional parameters have been added.
+
+According to Junxian et al :cite:`adapters-components-Junxian2021unified`, we can consider an adapter being represented as three components -
+
+1) Functional form - the trainable parameters that will modify the input
+2) Insertion form - Where the adapter outputs are integrated with the original input. The input to the adapters can be the last output of the layer, the input to some attention layer, or even the original input to the module itself (before even the modules forward pass).
+3) Composition function - How the adapters outputs are integrated with the inputs. It can be as simple as residual addition connection, or concatenation, or point-wise multiplication etc.
+
+Functional Form - Adapter Networks
+==================================
+
+Adapter modules represent the functional form of the adapter. We discuss an example of a most commonly used adapter module found in literature, titled the ``LinearAdapter`` (or Houlsby Adapter) :cite:`adapters-components-houlsby2019adapter`.
+
+.. note::
+
+    All adapter modules must extend :class:`~nemo.collections.common.parts.adapter_modules.AdapterModuleUtil` and should ideally have an equivalent DataClass config for easy instantiation !
+
+
+.. autoclass:: nemo.collections.common.parts.adapter_modules.AdapterModuleUtil
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+    :noindex:
+
+-----
+
+.. autoclass:: nemo.collections.common.parts.adapter_modules.LinearAdapter
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+    :noindex:
+
+
+Insertion Form - Module Adapters
+--------------------------------
+
+Adapter modules can be integrated into many different locations of a given module. For example, it is possible to have an adapter that affects only the outputs of the final layer in each module. We can also have a ``Parallel Adapter`` :cite:`adapters-components-Junxian2021unified` that operates at the input of the module itself, in parallel to the forward pass of the module. Yet another insertion location is inside the Multi Head Attention Layers.
+
+On top of this, while adapters are commonly used only in the layers containing the most parameters (say the Encoder of a network), some models can support adapters in multiple locations (Encoder-Decoder architecture for Language Models, Machine Translation, or even Encoder-Decoder-Joint for ASR with Transducer Loss). As such, NeMo utilizes the concept of ``Module Adapters``.
+
+``Module Adapters`` are very simply defined when adding an adapter - by specifying the module that the adapter should be inserted into.
+
+.. code-block:: python
+
+    # Get the list of supported modules / locations in a adapter compatible Model
+    print(model.adapter_module_names)  # assume ['', 'encoder', 'decoder']
+
+    # When calling add_adapter, specify the module name in the left of the colon symbol, and the adapter name afterwords.
+    # The adapter is then directed to the decoder module instead of the default / encoder module.
+    model.add_adapter("decoder:first_adapter", cfg=...)
+
+You might note that ``model.adapter_module_names`` can sometimes return ``''`` as one of the supported module names - this refers to the "default module". Generally we try to provide the default as the most commonly used adapter in literature - for example, Encoder adapters in NLP/NMT/ASR.
+
+Composition Function - Adapter Strategies
+-----------------------------------------
+
+Finally, we discuss how to compose the input and output of adapter modules. In order to generalize this step, we construct ``Adapter Strategies``.
+A strategy is any class (not torch.nn.Module!) that extends :class:`~nemo.core.classes.mixins.adapter_mixin_strategies.AbstractAdapterStrategy`, and provides a ``forward()`` method that accepts a specific signature of the inputs and produces an output tensor which combines the input and output with some specific method.
+
+We discuss a simple residual additional connection strategy below - that accepts an input to the adapter and an adapters output and simply adds them together. It also supports ``stochastic_depth`` which enables adapters to be dynamically switched off during training, making training more robust.
+
+.. autoclass:: nemo.core.classes.mixins.adapter_mixin_strategies.AbstractAdapterStrategy
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+    :undoc-members: adapter_module_names
+    :noindex:
+
+-----
+
+.. autoclass:: nemo.core.classes.mixins.adapter_mixin_strategies.ResidualAddAdapterStrategy
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+    :undoc-members: adapter_module_names
+    :noindex:
+
+-----
+
+
+References
+----------
+
+.. bibliography:: ./adapter_bib.bib
+    :style: plain
+    :keyprefix: adapters-components-

docs/source/core/adapters/intro.rst

@@ -0,0 +1,149 @@
+Adapters
+========
+
+In NeMo, we often train models and fine-tune them for a specific task. This is a reasonable approach when the models are just a few million parameters. However, this approach quickly becomes infeasible when approaching hundreds of millions or even billions of parameters. As a potential solution to such a scenario, where fine-tuning a massive model is no longer feasible, we look to `Adapters <https://arxiv.org/abs/1902.00751>`_ :cite:`adapters-houlsby2019adapter` to specialize our model on a specific domain or task. Adapters require a fraction of the total number of parameters as the original model and are much more efficient to fine-tune.
+
+.. note::
+
+    For a detailed tutorial on adding ``Adapter`` support to any PyTorch module, please refer to the NeMo Adapter :doc:`Tutorials <../../starthere/tutorials>`.
+
+
+What are Adapters?
+------------------
+
+Adapters are a straightforward concept - one formulation can be shown by the diagram below. At their simplest, they are residual Feedforward layers that compress the input dimension (:math:`D`) to a small bottleneck dimension (:math:`H`), such that :math:`R^D \text{->} R^H`, compute an activation (such as ReLU), finally mapping :math:`R^H \text{->} R^D` with another Feedforward layer. This output is then added to the input via a simple residual connection.
+
+.. raw:: html
+
+    <div align="center">
+      <img src="https://mermaid.ink/img/pako:eNptkLFqwzAQhl9F3ORAPDSjA4EUx6RgXEjbycpwWOdG1JaMfEoakrx7ZcfpUKrlxH_fz4d0gcoqggTqxp6qAzoW76k0Ipx1-WI6z3sRxyuRF1GOZ3KisK6d3YG8GFdZ9hRJeLbMDRmqvkRGpDLrTuiUiEWUigBtlyIVqzBnEqZ66I39dcX6iKytKXeUf-wn-286QoFeBMvmu0PTD-EfyXaQpP9JFmP_1XN4S3kfD8W4ue6o18pjc52gYQlzaMm1qFX4msuQSOADtSQhCdfaOupZgjS3QPpOIdNGabYOkhqbnuaAnu3b2VSQsPP0gFKNnw7bibr9AJkZdXU" height=100% />
+    </div>
+
+Adapter modules such as this are usually initialized such that the initial output of the adapter will always be zeros so as to prevent degradation of the original model's performance due to addition of such modules.
+
+``torch.nn.Module`` with Adapters
+---------------------------------
+
+In NeMo, Adapters are supported via a ``Mixin`` class that can be attached to any ``torch.nn.Module``. Such a module will
+have multiple additional methods which will enable adapter capabilities in that module.
+
+.. code-block:: python
+
+    # Import the adapter mixin from NeMo
+    from nemo.core import adapter_mixins
+
+    # NOTE: See the *two* classes being inherited here !
+    class MyModule(torch.nn.Module, adapter_mixins.AdapterModuleMixin):
+        pass
+
+
+AdapterModuleMixin
+------------------
+Let's look into what :class:`~nemo.core.adapter_mixins.AdapterModuleMixin` adds to the general PyTorch module. Some of the most important methods that are required are listed below :
+
+1) ``add_adapter``: Used to add an adapter with a unique name to the module.
+2) ``get_enabled_adapters``: Returns a list of names of all enabled adapter modules.
+3) ``set_enabled_adapters``: Sets whether a single (or all) adapters are enabled or disabled.
+4) ``is_adapter_available``: Check if any adapter is available and enabled or not.
+
+Modules that extend this mixin usually can directly use these methods without extending them, but we will cover a case below
+where you may wish to extend these methods.
+
+.. autoclass:: nemo.core.adapter_mixins.AdapterModuleMixin
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+    :undoc-members: adapter_module_names
+
+
+Using the Adapter Module
+------------------------
+
+Now that ``MyModule`` supports adapters, we can easily add adapters, set their state, check if they are available and
+perform their forward pass. Note that if multiple adapters are enabled, they are called in a chain, the output of the previous adapter is passed as input to the next adapter and so on.
+
+.. code-block:: python
+
+    # Import the adapter mixin and modules from NeMo
+    import torch
+    from nemo.core import adapter_mixins
+    from nemo.collections.common.parts import adapter_modules
+
+    class MyModule(torch.nn.Module, adapter_mixins.AdapterModuleMixin):
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            output = self.layers(x)  # assume self.layers is some Sequential() module
+
+            if self.is_adapter_available():  # check if adapters were added or not
+                output = self.forward_enabled_adapters()  # perform the forward of all enabled adapters in a chain
+
+            return output
+
+    # Now let us create a module, add an adapter and do a forward pass with some random inputs
+    module = MyModule(dim)  # assume dim is some input and output dimension of the module.
+
+    # Add an adapter
+    module.add_adapter("first_adapter", cfg=adapter_modules.LinearAdapter(in_features=dim, dim=5))
+
+    # Check if adapter is available
+    module.is_adapter_available()  # returns True
+
+    # Check the name(s) of the enabled adapters
+    module.get_enabled_adapters()  # returns ['first_adapter']
+
+    # Set the state of the adapter (by name)
+    module.set_enabled_adapters(name="first_adapter", enabled=True)
+
+    # Freeze all the weights of the original module  (equivalent to calling module.freeze() for a NeuralModule)
+    for param in module.parameters():
+        param.requires_grad = False
+
+    # Unfreeze only the adapter weights (so that we finetune only the adapters and not the original weights !)
+    module.unfreeze_enabled_adapters()
+
+    # Now you can train this model's adapters !
+    input_data = torch.randn(4, dim)  # assume dim is the input-output dim of the module
+    outputs_with_adapter = module(input_data)
+
+    # Compute loss and backward ...
+
+
+Adapter Compatible Models
+-------------------------
+
+If the goal was to support adapters in a single module, then the goal has been accomplished. In the real world however, we
+build large composite models out of multiple modules and combine them to build a final model that we then train. We do this using the
+:class:`~nemo.core.adapter_mixins.AdapterModelPTMixin`.
+
+.. note::
+
+    For an in-depth guide to supporting hierarchical adapter modules, please refer to the NeMo Adapter :doc:`Tutorials <../../starthere/tutorials>`.
+
+
+.. autoclass:: nemo.core.adapter_mixins.AdapterModelPTMixin
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+    :undoc-members: adapter_module_names
+
+Below, we will discuss some useful functionality of Adapter Compatible Models.
+
+1) ``Save and restore a Model with adapter capability``: Any NeMo model that implements this class correctly can save and restore NeMo models with adapter capabilities, thereby allowing sharing of adapters.
+2) ``save_adapters`` and ``load_adapters``: Adapters are usually a very small number of parameters, there is no need for the entire model to be duplicated for each adapter. This method allows storing just the adapter module(s) separately from the Model, so that you can use the same "base" model, and share just the Adapter modules.
+
+
+.. toctree::
+   :maxdepth: 8
+   :caption: Adapters
+
+   components
+   api
+
+
+References
+----------
+
+.. bibliography:: ./adapter_bib.bib
+    :style: plain
+    :labelprefix: adapters
+    :keyprefix: adapters-

docs/source/core/api.rst

@@ -0,0 +1,130 @@
+
+NeMo Core APIs
+==============
+
+Base class for all NeMo models
+------------------------------
+
+.. autoclass:: nemo.core.ModelPT
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+    :undoc-members: cfg, num_weights
+    :exclude-members: set_eff_save, use_eff_save, teardown
+
+Base Neural Module class
+------------------------
+
+.. autoclass:: nemo.core.NeuralModule
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+
+Base Mixin classes
+------------------
+
+.. autoclass:: nemo.core.Typing
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+    :private-members:
+    :exclude-members: _abc_impl
+    :noindex:
+
+
+.. autoclass:: nemo.core.Serialization
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+    :noindex:
+
+
+.. autoclass:: nemo.core.FileIO
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+    :noindex:
+
+
+Base Connector classes
+----------------------
+
+.. autoclass:: nemo.core.connectors.save_restore_connector.SaveRestoreConnector
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+
+
+Base Mixin Classes
+------------------
+
+.. autoclass:: nemo.core.classes.mixins.access_mixins.AccessMixin
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+
+
+.. autoclass:: nemo.core.classes.mixins.hf_io_mixin.HuggingFaceFileIO
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+
+
+Neural Type checking
+--------------------
+
+.. autoclass:: nemo.core.classes.common.typecheck
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+
+    .. automethod:: __call__
+
+Neural Type classes
+-------------------
+
+.. autoclass:: nemo.core.neural_types.NeuralType
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+
+
+.. autoclass:: nemo.core.neural_types.axes.AxisType
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+
+
+.. autoclass:: nemo.core.neural_types.elements.ElementType
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+
+
+.. autoclass:: nemo.core.neural_types.comparison.NeuralTypeComparisonResult
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+
+Experiment manager
+------------------
+
+.. autoclass:: nemo.utils.exp_manager.exp_manager
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+
+.. autoclass:: nemo.utils.exp_manager.ExpManagerConfig
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+
+
+Exportable
+----------
+
+.. autoclass:: nemo.core.classes.exportable.Exportable
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+

docs/source/core/core.rst

@@ -0,0 +1,777 @@
+NeMo Models
+===========
+
+Basics
+------
+
+NeMo models contain everything needed to train and reproduce conversational AI models:
+
+- neural network architectures 
+- datasets/data loaders
+- data preprocessing/postprocessing
+- data augmentors
+- optimizers and schedulers
+- tokenizers
+- language models
+
+NeMo uses `Hydra <https://hydra.cc/>`_ for configuring both NeMo models and the PyTorch Lightning Trainer.
+
+.. note::
+    Every NeMo model has an example configuration file and training script that can be found `here <https://github.com/NVIDIA/NeMo/tree/stable/examples>`__.
+
+The end result of using NeMo, `Pytorch Lightning <https://github.com/PyTorchLightning/pytorch-lightning>`__, and Hydra is that NeMo models all have the same look and feel and are also fully compatible with the PyTorch ecosystem.
+
+Pretrained
+----------
+
+NeMo comes with many pretrained models for each of our collections: ASR, TTS, Audio, and SpeechLM2. Every pretrained NeMo model can be downloaded 
+and used with the ``from_pretrained()`` method.
+
+As an example, we can instantiate a Parakeet model with the following:
+
+.. code-block:: Python
+
+    import nemo.collections.asr as nemo_asr
+
+    model = nemo_asr.models.ASRModel.from_pretrained(model_name="nvidia/parakeet-tdt-0.6b-v2")
+
+To see all available pretrained models for a specific NeMo model, use the ``list_available_models()`` method:
+
+.. code-block:: Python
+
+    nemo_asr.models.EncDecCTCModel.list_available_models()
+
+For detailed information on the available pretrained models, refer to the collections documentation: 
+
+- :doc:`Automatic Speech Recognition (ASR) <../asr/intro>`
+- :doc:`Text-to-Speech Synthesis (TTS) <../tts/intro>`
+
+Training
+--------
+
+NeMo leverages `PyTorch Lightning <https://www.pytorchlightning.ai/>`__ for model training. PyTorch Lightning lets NeMo decouple the
+conversational AI code from the PyTorch training code. This means that NeMo users can focus on their domain (ASR, NLP, TTS) and 
+build complex AI applications without having to rewrite boilerplate code for PyTorch training.
+
+When using PyTorch Lightning, NeMo users can automatically train with:
+
+- multi-GPU/multi-node
+- mixed precision
+- model checkpointing
+- logging
+- early stopping
+- and more
+
+The two main aspects of the Lightning API are the `LightningModule <https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#>`_ 
+and the `Trainer <https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html>`_.
+
+PyTorch Lightning ``LightningModule``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Every NeMo model is a ``LightningModule`` which is an ``nn.module``. This means that NeMo models are compatible with the PyTorch 
+ecosystem and can be plugged into existing PyTorch workflows.
+
+Creating a NeMo model is similar to any other PyTorch workflow. We start by initializing our model architecture, then define the forward pass:
+
+.. code-block:: python
+
+    class TextClassificationModel(NLPModel, Exportable):
+        ...
+        def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+            """Initializes the BERTTextClassifier model."""
+            ...
+            super().__init__(cfg=cfg, trainer=trainer)
+
+            # instantiate a BERT based encoder
+            self.bert_model = get_lm_model(
+                config_file=cfg.language_model.config_file,
+                config_dict=cfg.language_model.config,
+                vocab_file=cfg.tokenizer.vocab_file,
+                trainer=trainer,
+                cfg=cfg,
+            )
+
+            # instantiate the FFN for classification
+            self.classifier = SequenceClassifier(
+                hidden_size=self.bert_model.config.hidden_size,
+                num_classes=cfg.dataset.num_classes,
+                num_layers=cfg.classifier_head.num_output_layers,
+                activation='relu',
+                log_softmax=False,
+                dropout=cfg.classifier_head.fc_dropout,
+                use_transformer_init=True,
+                idx_conditioned_on=0,
+            )
+
+.. code-block:: python
+
+        def forward(self, input_ids, token_type_ids, attention_mask):
+            """
+            No special modification required for Lightning, define it as you normally would
+            in the `nn.Module` in vanilla PyTorch.
+            """
+            hidden_states = self.bert_model(
+                input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask
+            )
+            logits = self.classifier(hidden_states=hidden_states)
+            return logits
+
+The ``LightningModule`` organizes PyTorch code so that across all NeMo models we have a similar look and feel.
+For example, the training logic can be found in ``training_step``:
+
+.. code-block:: python
+
+    def training_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the training loop with the data from the training dataloader
+        passed in as `batch`.
+        """
+        # forward pass
+        input_ids, input_type_ids, input_mask, labels = batch
+        logits = self.forward(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask)
+
+        train_loss = self.loss(logits=logits, labels=labels)
+
+        lr = self._optimizer.param_groups[0]['lr']
+
+        self.log('train_loss', train_loss)
+        self.log('lr', lr, prog_bar=True)
+
+        return {
+            'loss': train_loss,
+            'lr': lr,
+        }
+
+While validation logic can be found in ``validation_step``:
+
+.. code-block:: python
+
+    def validation_step(self, batch, batch_idx):
+        """
+        Lightning calls this inside the validation loop with the data from the validation dataloader
+        passed in as `batch`.
+        """
+        if self.testing:
+            prefix = 'test'
+        else:
+            prefix = 'val'
+
+        input_ids, input_type_ids, input_mask, labels = batch
+        logits = self.forward(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask)
+
+        val_loss = self.loss(logits=logits, labels=labels)
+
+        preds = torch.argmax(logits, axis=-1)
+
+        tp, fn, fp, _ = self.classification_report(preds, labels)
+
+        return {'val_loss': val_loss, 'tp': tp, 'fn': fn, 'fp': fp}
+
+PyTorch Lightning then handles all of the boilerplate code needed for training. Virtually any aspect of training can be customized
+via PyTorch Lightning `hooks <https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#hooks>`_, 
+`Plugins <https://pytorch-lightning.readthedocs.io/en/stable/extensions/plugins.html>`_, 
+`callbacks <https://pytorch-lightning.readthedocs.io/en/stable/extensions/callbacks.html>`_, or by overriding `methods <https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#methods>`_. 
+
+For more domain-specific information, see:
+
+- :doc:`Automatic Speech Recognition (ASR) <../asr/intro>`
+- :doc:`Text-to-Speech Synthesis (TTS) <../tts/intro>`
+
+PyTorch Lightning Trainer
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Since every NeMo model is a ``LightningModule``, we can automatically take advantage of the PyTorch Lightning ``Trainer``. Every NeMo 
+`example <https://github.com/NVIDIA/NeMo/tree/v1.0.2/examples>`_ training script uses the ``Trainer`` object to fit the model.
+
+First, instantiate the model and trainer, then call ``.fit``:
+
+.. code-block:: python
+    
+    # We first instantiate the trainer based on the model configuration.
+    # See the model configuration documentation for details.    
+    trainer = pl.Trainer(**cfg.trainer)
+
+    # Then pass the model configuration and trainer object into the NeMo model
+    model = TextClassificationModel(cfg.model, trainer=trainer)
+
+    # Now we can train with by calling .fit
+    trainer.fit(model)
+
+    # Or we can run the test loop on test data by calling
+    trainer.test(model=model)
+
+All `trainer flags <https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html#trainer-flags>`_ can be set from from the NeMo configuration. 
+    
+
+Configuration
+-------------
+
+Hydra is an open-source Python framework that simplifies configuration for complex applications that must bring together many different 
+software libraries. Conversational AI model training is a great example of such an application. To train a conversational AI model, we 
+must be able to configure:
+
+- neural network architectures
+- training and optimization algorithms 
+- data pre/post processing
+- data augmentation
+- experiment logging/visualization
+- model checkpointing   
+
+For an introduction to using Hydra, refer to the `Hydra Tutorials <https://hydra.cc/docs/tutorials/intro>`_.
+
+With Hydra, we can configure everything needed for NeMo with three interfaces:
+
+- Command Line (CLI) 
+- Configuration Files (YAML)
+- Dataclasses (Python)
+
+YAML
+~~~~
+
+NeMo provides YAML configuration files for all of our `example <https://github.com/NVIDIA/NeMo/tree/v1.0.2/examples>`_ training scripts.
+YAML files make it easy to experiment with different model and training configurations.
+
+Every NeMo example YAML has the same underlying configuration structure:
+
+- trainer
+- exp_manager
+- model
+
+The model configuration always contains ``train_ds``, ``validation_ds``, ``test_ds``, and ``optim``.  Model architectures, however, can vary across domains.
+Refer to the documentation of specific collections (LLM, ASR etc.) for detailed information on model architecture configuration.
+
+A NeMo configuration file should look similar to the following:
+
+.. code-block:: yaml
+
+    # PyTorch Lightning Trainer configuration
+    # any argument of the Trainer object can be set here
+    trainer:
+        devices: 1 # number of gpus per node
+        accelerator: gpu
+        num_nodes: 1 # number of nodes
+        max_epochs: 10 # how many training epochs to run
+        val_check_interval: 1.0 # run validation after every epoch
+
+    # Experiment logging configuration
+    exp_manager:
+        exp_dir: /path/to/my/nemo/experiments
+        name: name_of_my_experiment
+        create_tensorboard_logger: True
+        create_wandb_logger: True
+
+    # Model configuration
+    # model network architecture, train/val/test datasets, data augmentation, and optimization
+    model:
+        train_ds:
+            manifest_filepath: /path/to/my/train/manifest.json
+            batch_size: 256
+            shuffle: True
+        validation_ds:
+            manifest_filepath: /path/to/my/validation/manifest.json
+            batch_size: 32
+            shuffle: False
+        test_ds:
+            manifest_filepath: /path/to/my/test/manifest.json
+            batch_size: 32
+            shuffle: False
+        optim:
+            name: novograd
+            lr: .01
+            betas: [0.8, 0.5]
+            weight_decay: 0.001
+        # network architecture can vary greatly depending on the domain
+        encoder:
+            ...
+        decoder:
+            ...
+
+CLI
+~~~
+
+With NeMo and Hydra, every aspect of model training can be modified from the command-line. This is extremely helpful for running lots 
+of experiments on compute clusters or for quickly testing parameters during development.
+
+All NeMo `examples <https://github.com/NVIDIA/NeMo/tree/stable/examples>`_ come with instructions on how to
+run the training/inference script from the command-line (e.g. see `here <https://github.com/NVIDIA/NeMo/blob/stable/examples/asr/asr_ctc/speech_to_text_ctc.py>`__
+for an example).
+
+With Hydra, arguments are set using the ``=`` operator:
+
+.. code-block:: bash
+
+    python examples/asr/asr_ctc/speech_to_text_ctc.py \
+        model.train_ds.manifest_filepath=/path/to/my/train/manifest.json \
+        model.validation_ds.manifest_filepath=/path/to/my/validation/manifest.json \
+        trainer.devices=2 \
+        trainer.accelerator='gpu' \
+        trainer.max_epochs=50
+
+We can use the ``+`` operator to add arguments from the CLI:
+
+.. code-block:: bash
+
+    python examples/asr/asr_ctc/speech_to_text_ctc.py \
+        model.train_ds.manifest_filepath=/path/to/my/train/manifest.json \
+        model.validation_ds.manifest_filepath=/path/to/my/validation/manifest.json \
+        trainer.devices=2 \
+        trainer.accelerator='gpu' \
+        trainer.max_epochs=50 \
+        +trainer.fast_dev_run=true
+
+We can use the ``~`` operator to remove configurations:
+
+.. code-block:: bash
+
+    python examples/asr/asr_ctc/speech_to_text_ctc.py \
+        model.train_ds.manifest_filepath=/path/to/my/train/manifest.json \
+        model.validation_ds.manifest_filepath=/path/to/my/validation/manifest.json \
+        ~model.test_ds \
+        trainer.devices=2 \
+        trainer.accelerator='gpu' \
+        trainer.max_epochs=50 \
+        +trainer.fast_dev_run=true
+
+We can specify configuration files using the ``--config-path`` and ``--config-name`` flags:
+
+.. code-block:: bash
+
+    python examples/asr/asr_ctc/speech_to_text_ctc.py \
+        --config-path=conf/conformer \
+        --config-name=conformer_ctc_bpe \
+        model.train_ds.manifest_filepath=/path/to/my/train/manifest.json \
+        model.validation_ds.manifest_filepath=/path/to/my/validation/manifest.json \
+        ~model.test_ds \
+        trainer.devices=2 \
+        trainer.accelerator='gpu' \
+        trainer.max_epochs=50 \
+        +trainer.fast_dev_run=true
+
+Dataclasses
+~~~~~~~~~~~
+
+Dataclasses allow NeMo to ship model configurations as part of the NeMo library and also enables pure Python configuration of NeMo models. 
+With Hydra, dataclasses can be used to create `structured configs <https://hydra.cc/docs/tutorials/structured_config/intro>`_ for the conversational AI application. 
+
+As an example, refer to the code block below for an *Attenion is All You Need* machine translation model. The model configuration can 
+be instantiated and modified like any Python `Dataclass <https://docs.python.org/3/library/dataclasses.html>`_.
+
+.. code-block:: Python
+
+    from nemo.collections.nlp.models.machine_translation.mt_enc_dec_config import AAYNBaseConfig
+
+    cfg = AAYNBaseConfig()
+
+    # modify the number of layers in the encoder
+    cfg.encoder.num_layers = 8
+
+    # modify the training batch size
+    cfg.train_ds.tokens_in_batch = 8192
+
+.. note:: Configuration with Hydra always has the following precedence CLI > YAML > Dataclass.
+
+.. _optimization-label:
+
+Optimization
+------------
+
+Optimizers and learning rate schedules are configurable across all NeMo models and have their own namespace. Here is a sample YAML 
+configuration for a Novograd optimizer with a Cosine Annealing learning rate schedule.
+
+.. code-block:: yaml
+
+    optim:
+        name: novograd
+        lr: 0.01
+    
+        # optimizer arguments
+        betas: [0.8, 0.25]
+        weight_decay: 0.001
+    
+        # scheduler setup
+        sched:
+            name: CosineAnnealing
+    
+            # Optional arguments
+            max_steps: -1 # computed at runtime or explicitly set here
+            monitor: val_loss
+            reduce_on_plateau: false
+    
+            # scheduler config override
+            warmup_steps: 1000
+            warmup_ratio: null
+            min_lr: 1e-9:
+
+.. note:: `NeMo Examples <https://github.com/NVIDIA/NeMo/tree/stable/examples>`_ has optimizer and scheduler configurations for every NeMo model.
+
+Optimizers can be configured from the CLI as well:
+
+.. code-block:: bash
+
+    python examples/asr/asr_ctc/speech_to_text_ctc.py \
+        --config-path=conf/conformer \
+        --config-name=conformer_ctc_bpe \
+        ...
+        # train with the adam optimizer
+        model.optim=adam \
+        # change the learning rate
+        model.optim.lr=.0004 \
+        # modify betas 
+        model.optim.betas=[.8, .5]
+
+.. _optimizers-label:
+
+Optimizers
+~~~~~~~~~~
+
+``name`` corresponds to the lowercase name of the optimizer. To view a list of available optimizers, run:
+
+.. code-block:: Python
+
+    from nemo.core.optim.optimizers import AVAILABLE_OPTIMIZERS
+
+    for name, opt in AVAILABLE_OPTIMIZERS.items():
+        print(f'name: {name}, opt: {opt}')
+
+.. code-block:: bash
+
+    name: sgd opt: <class 'torch.optim.sgd.SGD'>
+    name: adam opt: <class 'torch.optim.adam.Adam'>
+    name: adamw opt: <class 'torch.optim.adamw.AdamW'>
+    name: adadelta opt: <class 'torch.optim.adadelta.Adadelta'>
+    name: adamax opt: <class 'torch.optim.adamax.Adamax'>
+    name: adagrad opt: <class 'torch.optim.adagrad.Adagrad'>
+    name: rmsprop opt: <class 'torch.optim.rmsprop.RMSprop'>
+    name: rprop opt: <class 'torch.optim.rprop.Rprop'>
+    name: novograd opt: <class 'nemo.core.optim.novograd.Novograd'>
+
+Optimizer Params
+~~~~~~~~~~~~~~~~
+
+Optimizer params can vary between optimizers but the ``lr`` param is required for all optimizers. To see the available params for an 
+optimizer, we can look at its corresponding dataclass.
+
+.. code-block:: python
+
+    from nemo.core.config.optimizers import NovogradParams
+
+    print(NovogradParams())
+
+.. code-block:: bash
+
+    NovogradParams(lr='???', betas=(0.95, 0.98), eps=1e-08, weight_decay=0, grad_averaging=False, amsgrad=False, luc=False, luc_trust=0.001, luc_eps=1e-08)
+
+``'???'`` indicates that the lr argument is required.
+
+Register Optimizer
+~~~~~~~~~~~~~~~~~~
+
+To register a new optimizer to be used with NeMo, run:
+
+.. autofunction:: nemo.core.optim.optimizers.register_optimizer
+
+.. _learning-rate-schedulers-label:
+
+Learning Rate Schedulers
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+Learning rate schedulers can be optionally configured under the ``optim.sched`` namespace.
+
+``name`` corresponds to the name of the learning rate schedule. To view a list of available schedulers, run: 
+    
+.. code-block:: Python
+
+    from nemo.core.optim.lr_scheduler import AVAILABLE_SCHEDULERS
+
+    for name, opt in AVAILABLE_SCHEDULERS.items():
+        print(f'name: {name}, schedule: {opt}')
+
+.. code-block:: bash
+
+    name: WarmupPolicy, schedule: <class 'nemo.core.optim.lr_scheduler.WarmupPolicy'>
+    name: WarmupHoldPolicy, schedule: <class 'nemo.core.optim.lr_scheduler.WarmupHoldPolicy'>
+    name: SquareAnnealing, schedule: <class 'nemo.core.optim.lr_scheduler.SquareAnnealing'>
+    name: CosineAnnealing, schedule: <class 'nemo.core.optim.lr_scheduler.CosineAnnealing'>
+    name: NoamAnnealing, schedule: <class 'nemo.core.optim.lr_scheduler.NoamAnnealing'>
+    name: WarmupAnnealing, schedule: <class 'nemo.core.optim.lr_scheduler.WarmupAnnealing'>
+    name: InverseSquareRootAnnealing, schedule: <class 'nemo.core.optim.lr_scheduler.InverseSquareRootAnnealing'>
+    name: SquareRootAnnealing, schedule: <class 'nemo.core.optim.lr_scheduler.SquareRootAnnealing'>
+    name: PolynomialDecayAnnealing, schedule: <class 'nemo.core.optim.lr_scheduler.PolynomialDecayAnnealing'>
+    name: PolynomialHoldDecayAnnealing, schedule: <class 'nemo.core.optim.lr_scheduler.PolynomialHoldDecayAnnealing'>
+    name: StepLR, schedule: <class 'torch.optim.lr_scheduler.StepLR'>
+    name: ExponentialLR, schedule: <class 'torch.optim.lr_scheduler.ExponentialLR'>
+    name: ReduceLROnPlateau, schedule: <class 'torch.optim.lr_scheduler.ReduceLROnPlateau'>
+    name: CyclicLR, schedule: <class 'torch.optim.lr_scheduler.CyclicLR'>
+
+Scheduler Params
+~~~~~~~~~~~~~~~~
+
+To see the available params for a scheduler, we can look at its corresponding dataclass:
+
+.. code-block:: Python
+
+    from nemo.core.config.schedulers import CosineAnnealingParams
+
+    print(CosineAnnealingParams())
+
+.. code-block:: bash
+
+    CosineAnnealingParams(last_epoch=-1, warmup_steps=None, warmup_ratio=None, min_lr=0.0)
+
+Register scheduler
+~~~~~~~~~~~~~~~~~~
+
+To register a new scheduler to be used with NeMo, run:
+
+.. autofunction:: nemo.core.optim.lr_scheduler.register_scheduler
+
+Save and Restore
+----------------
+
+NeMo models all come with ``.save_to`` and ``.restore_from`` methods.  
+
+Save
+~~~~
+
+To save a NeMo model, run:
+
+.. code-block:: Python
+
+    model.save_to('/path/to/model.nemo')
+
+Everything needed to use the trained model is packaged and saved in the ``.nemo`` file. For example, in the NLP domain, ``.nemo`` files 
+include the necessary tokenizer models and/or vocabulary files, etc.
+
+.. note:: A ``.nemo`` file is simply an archive like any other ``.tar`` file.
+
+Restore
+~~~~~~~
+
+To restore a NeMo model, run:
+
+.. code-block:: Python
+
+    # Here, you should usually use the class of the model, or simply use ModelPT.restore_from() for simplicity.
+    model.restore_from('/path/to/model.nemo')
+
+When using the PyTorch Lightning Trainer, a PyTorch Lightning checkpoint is created. These are mainly used within NeMo to auto-resume 
+training. Since NeMo models are ``LightningModules``, the PyTorch Lightning method ``load_from_checkpoint`` is available. Note that 
+``load_from_checkpoint`` won't necessarily work out-of-the-box for all models as some models require more artifacts than just the 
+checkpoint to be restored. For these models, the user will have to override ``load_from_checkpoint`` if they want to use it.
+
+It's highly recommended to use ``restore_from`` to load NeMo models.
+
+Restore with Modified Config
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Sometimes, there may be a need to modify the model (or it's sub-components) prior to restoring a model. A common case is when
+the model's internal config must be updated due to various reasons (such as deprecation, newer versioning, support a new feature).
+As long as the model has the same parameters as compared to the original config, the parameters can once again be restored safely.
+
+In NeMo, as part of the .nemo file, the model's internal config will be preserved. This config is used during restoration, and
+as shown below we can update this config prior to restoring the model.
+
+.. code-block::
+
+    # When restoring a model, you should generally use the class of the model
+    # Obtain the config (as an OmegaConf object)
+    config = model_class.restore_from('/path/to/model.nemo', return_config=True)
+    # OR
+    config = model_class.from_pretrained('name_of_the_model', return_config=True)
+
+    # Modify the config as needed
+    config.x.y = z
+
+    # Restore the model from the updated config
+    model = model_class.restore_from('/path/to/model.nemo', override_config_path=config)
+    # OR
+    model = model_class.from_pretrained('name_of_the_model', override_config_path=config)
+
+Register Artifacts
+------------------
+
+Restoring conversational AI models can be complicated because it requires more than just the checkpoint weights; additional information is also needed to use the model.
+NeMo models can save additional artifacts in the .nemo file by calling ``.register_artifact``.
+When restoring NeMo models using ``.restore_from`` or ``.from_pretrained``, any artifacts that were registered will be available automatically.
+
+As an example, consider an NLP model that requires a trained tokenizer model. 
+The tokenizer model file can be automatically added to the .nemo file with the following:
+
+.. code-block:: python
+
+    self.encoder_tokenizer = get_nmt_tokenizer(
+        ...
+        tokenizer_model=self.register_artifact(config_path='encoder_tokenizer.tokenizer_model',
+                                               src='/path/to/tokenizer.model',
+                                               verify_src_exists=True),
+    )
+
+By default, ``.register_artifact`` will always return a path. If the model is being restored from a .nemo file, 
+then that path will be to the artifact in the .nemo file. Otherwise, ``.register_artifact`` will return the local path specified by the user.
+
+``config_path`` is the artifact key. It usually corresponds to a model configuration but does not have to.
+The model config that is packaged with the .nemo file will be updated according to the ``config_path`` key.
+In the above example, the model config will have 
+
+.. code-block:: YAML
+
+    encoder_tokenizer:
+        ...
+        tokenizer_model: nemo:4978b28103264263a03439aaa6560e5e_tokenizer.model
+
+``src`` is the path to the artifact and the base-name of the path will be used when packaging the artifact in the .nemo file.
+Each artifact will have a hash prepended to the basename of ``src`` in the .nemo file. This is to prevent collisions with basenames 
+base-names that are identical (say when there are two or more tokenizers, both called `tokenizer.model`).
+The resulting .nemo file will then have the following file:
+
+.. code-block:: bash
+
+    4978b28103264263a03439aaa6560e5e_tokenizer.model
+
+If ``verify_src_exists`` is set to ``False``, then the artifact is optional. This means that ``.register_artifact`` will return ``None`` 
+if the ``src`` cannot be found.
+
+Push to Hugging Face Hub
+------------------------
+
+NeMo models can be pushed to the `Hugging Face Hub <https://huggingface.co/>`_ with the :meth:`~nemo.core.classes.mixins.hf_io_mixin.HuggingFaceFileIO.push_to_hf_hub` method. This method performs the same actions as ``save_to()`` and then uploads the model to the HuggingFace Hub. It offers an additional ``pack_nemo_file`` argument that allows the user to upload the entire NeMo file or just the ``.nemo`` file. This is useful for large language models that have a massive number of parameters, and a single NeMo file could exceed the max upload size of Hugging Face Hub.
+
+
+Upload a model to the Hub
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: python
+
+    token = "<HF TOKEN>" or None
+    pack_nemo_file = True  # False will upload multiple files that comprise the NeMo file onto HF Hub; Generally useful for LLMs
+
+    model.push_to_hf_hub(
+       repo_id=repo_id, pack_nemo_file=pack_nemo_file, token=token,
+    )
+
+Use a Custom Model Card Template for the Hub
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: python
+
+    # Override the default model card
+    template = """ <Your own custom template>
+    # {model_name}
+    """
+    kwargs = {"model_name": "ABC", "repo_id": "nvidia/ABC_XYZ"}
+    model_card = model.generate_model_card(template=template, template_kwargs=kwargs, type="hf")
+
+    model.push_to_hf_hub(
+        repo_id=repo_id, token=token, model_card=model_card
+    )
+
+    # Write your own model card class
+    class MyModelCard:
+      def __init__(self, model_name):
+        self.model_name = model_name
+
+      def __repr__(self):
+        template = """This is the {model_name} model""".format(model_name=self.model_name)
+        return template
+
+    model.push_to_hf_hub(
+        repo_id=repo_id, token=token, model_card=MyModelCard("ABC")
+    )
+
+
+Nested NeMo Models
+------------------
+
+In some cases, it may be helpful to use NeMo models inside other NeMo models. For example, we can incorporate language models into ASR models to use in a decoding process to improve accuracy.
+
+There are three ways to instantiate child models inside parent models:
+
+- use subconfig directly
+- use the ``.nemo`` checkpoint path to load the child model
+- use a pretrained NeMo model
+
+To register a child model, use the ``register_nemo_submodule`` method of the parent model. This method will add the child model to a specified model attribute. During serialization, it will correctly handle child artifacts and store the child model’s configuration in the parent model’s ``config_field``.
+
+.. code-block:: python
+
+    from nemo.core.classes import ModelPT
+
+    class ChildModel(ModelPT):
+        ...  # implement necessary methods
+
+    class ParentModel(ModelPT):
+        def __init__(self, cfg, trainer=None):
+            super().__init__(cfg=cfg, trainer=trainer)
+
+            # optionally annotate type for IDE autocompletion and type checking
+            self.child_model: Optional[ChildModel]
+            if cfg.get("child_model") is not None:
+                # load directly from config
+                # either if config provided initially, or automatically
+                # after model restoration
+                self.register_nemo_submodule(
+                    name="child_model",
+                    config_field="child_model",
+                    model=ChildModel(self.cfg.child_model, trainer=trainer),
+                )
+            elif cfg.get('child_model_path') is not None:
+                # load from .nemo model checkpoint
+                # while saving, config will be automatically assigned/updated
+                # in cfg.child_model
+                self.register_nemo_submodule(
+                    name="child_model",
+                    config_field="child_model",
+                    model=ChildModel.restore_from(self.cfg.child_model_path, trainer=trainer),
+                )
+            elif cfg.get('child_model_name') is not None:
+                # load from pretrained model
+                # while saving, config will be automatically assigned/updated
+                # in cfg.child_model
+                self.register_nemo_submodule(
+                    name="child_model",
+                    config_field="child_model",
+                    model=ChildModel.from_pretrained(self.cfg.child_model_name, trainer=trainer),
+                )
+            else:
+                self.child_model = None
+
+
+
+Profiling 
+---------
+
+NeMo offers users two options for profiling: Nsys and CUDA memory profiling. These two options allow users
+to debug performance issues as well as memory issues such as memory leaks.
+
+To enable Nsys profiling, add the following options to the model config:
+
+.. code-block:: yaml
+
+    nsys_profile: False
+        start_step: 10  # Global batch to start profiling
+        end_step: 10 # Global batch to end profiling
+        ranks: [0] # Global rank IDs to profile
+        gen_shape: False # Generate model and kernel details including input shapes
+
+Finally, run the model training script with:
+
+.. code-block:: bash
+
+    nsys profile -s none -o <profile filepath> -t cuda,nvtx --force-overwrite true --capture-range=cudaProfilerApi --capture-range-end=stop python ./examples/...
+
+See more options at `nsight user guide <https://docs.nvidia.com/nsight-systems/UserGuide/index.html#cli-profiling>`_.
+
+
+
+To enable CUDA memory profiling, add the following options to the model config:
+
+.. code-block:: yaml
+
+    memory_profile:
+        enabled: True
+        start_step: 10  # Global batch to start profiling
+        end_step: 10 # Global batch to end profiling
+        rank: 0 # Global rank ID to profile
+        output_path: None # Path to store the profile output file
+
+Then invoke your NeMo script without any changes in the invocation command.

docs/source/core/exp_manager.rst

@@ -0,0 +1,502 @@
+
+.. _exp-manager-label:
+
+Experiment Manager
+==================
+
+The NeMo Toolkit Experiment Manager leverages PyTorch Lightning for model checkpointing, TensorBoard Logging, Weights and Biases, DLLogger and MLFlow logging. The
+Experiment Manager is included by default in all NeMo example scripts.
+
+To use the Experiment Manager, call :class:`~nemo.utils.exp_manager.exp_manager` and pass in the PyTorch Lightning ``Trainer``.
+
+.. code-block:: python
+
+    exp_dir = exp_manager(trainer, cfg.get("exp_manager", None))
+
+The Experiment Manager is configurable using YAML with Hydra.
+
+.. code-block:: bash
+
+    exp_manager:
+        exp_dir: /path/to/my/experiments
+        name: my_experiment_name
+        create_tensorboard_logger: True
+        create_checkpoint_callback: True
+
+Optionally, launch TensorBoard to view the training results in ``exp_dir``, which by default is set to ``./nemo_experiments``.
+
+.. code-block:: bash
+
+    tensorboard --bind_all --logdir nemo_experiments
+
+..
+
+If ``create_checkpoint_callback`` is set to ``True``, then NeMo automatically creates checkpoints during training
+using PyTorch Lightning's `ModelCheckpoint <https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html>`_.
+We can configure the ``ModelCheckpoint`` via YAML or CLI:
+
+.. code-block:: yaml
+
+    exp_manager:
+        ...
+        # configure the PyTorch Lightning ModelCheckpoint using checkpoint_call_back_params
+        # any ModelCheckpoint argument can be set here
+
+        # save the best checkpoints based on this metric
+        checkpoint_callback_params.monitor=val_loss
+
+        # choose how many total checkpoints to save
+        checkpoint_callback_params.save_top_k=5
+
+Resume Training
+---------------
+
+To auto-resume training, configure the ``exp_manager``. This feature is important for long training runs that might be interrupted or
+shut down before the procedure has completed. To auto-resume training, set the following parameters via YAML or CLI:
+
+.. code-block:: yaml
+
+    exp_manager:
+        ...
+        # resume training if checkpoints already exist
+        resume_if_exists: True
+
+        # to start training with no existing checkpoints
+        resume_ignore_no_checkpoint: True
+
+        # by default experiments will be versioned by datetime
+        # we can set our own version with
+        exp_manager.version: my_experiment_version
+
+
+Experiment Loggers
+------------------
+
+Alongside Tensorboard, NeMo also supports Weights and Biases, MLFlow, DLLogger, ClearML and NeptuneLogger. To use these loggers,  set the following
+via YAML or :class:`~nemo.utils.exp_manager.ExpManagerConfig`.
+
+
+Weights and Biases (WandB)
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. _exp_manager_weights_biases-label:
+
+.. code-block:: yaml
+
+    exp_manager:
+        ...
+        create_checkpoint_callback: True
+        create_wandb_logger: True
+        wandb_logger_kwargs:
+            name: ${name}
+            project: ${project}
+            entity: ${entity}
+            <Add any other arguments supported by WandB logger here>
+
+
+MLFlow
+~~~~~~
+
+.. _exp_manager_mlflow-label:
+
+.. code-block:: yaml
+
+    exp_manager:
+        ...
+        create_checkpoint_callback: True
+        create_mlflow_logger: True
+        mlflow_logger_kwargs:
+            experiment_name: ${name}
+            tags:
+                <Any key:value pairs>
+            save_dir: './mlruns'
+            prefix: ''
+            artifact_location: None
+            # provide run_id if resuming a previously started run
+            run_id: Optional[str] = None
+
+DLLogger
+~~~~~~~~
+
+.. _exp_manager_dllogger-label:
+
+.. code-block:: yaml
+
+    exp_manager:
+        ...
+        create_checkpoint_callback: True
+        create_dllogger_logger: True
+        dllogger_logger_kwargs:
+            verbose: False
+            stdout: False
+            json_file: "./dllogger.json"
+
+ClearML
+~~~~~~~
+
+.. _exp_manager_clearml-label:
+
+.. code-block:: yaml
+
+    exp_manager:
+        ...
+        create_checkpoint_callback: True
+        create_clearml_logger: True
+        clearml_logger_kwargs:
+            project: None  # name of the project
+            task: None  # optional name of task
+            connect_pytorch: False
+            model_name: None  # optional name of model
+            tags: None  # Should be a list of str
+            log_model: False  # log model to clearml server
+            log_cfg: False  # log config to clearml server
+            log_metrics: False  # log metrics to clearml server
+
+Neptune
+~~~~~~~
+
+.. _exp_manager_neptune-label:
+
+.. code-block:: yaml
+
+    exp_manager:
+        ...
+        create_checkpoint_callback: True
+        create_neptune_logger: false
+        neptune_logger_kwargs:
+            project: ${project}
+            name: ${name}
+            prefix: train
+            log_model_checkpoints: false # set to True if checkpoints need to be pushed to Neptune
+            tags: null # can specify as an array of strings in yaml array format
+            description: null
+            <Add any other arguments supported by Neptune logger here>
+
+Exponential Moving Average
+--------------------------
+
+.. _exp_manager_ema-label:
+
+NeMo supports using exponential moving average (EMA) for model parameters. This can be useful for improving model generalization
+and stability. To use EMA, set the following parameters via YAML or :class:`~nemo.utils.exp_manager.ExpManagerConfig`.
+
+.. code-block:: yaml
+
+    exp_manager:
+        ...
+        # use exponential moving average for model parameters
+        ema:
+            enabled: True  # False by default
+            decay: 0.999  # decay rate
+            cpu_offload: False  # If EMA parameters should be offloaded to CPU to save GPU memory
+            every_n_steps: 1  # How often to update EMA weights
+            validate_original_weights: False  # Whether to use original weights for validation calculation or EMA weights
+
+..  Support for Preemption
+    ----------------------
+
+    .. _exp_manager_preemption_support-label:
+
+    NeMo adds support for a callback upon preemption while running the models on clusters. The callback takes care of saving the current state of training via the ``.ckpt``
+    file followed by a graceful exit from the run. The checkpoint saved upon preemption has the ``*last.ckpt`` suffix and replaces the previously saved last checkpoints.
+    This feature is useful to increase utilization on clusters.
+    The ``PreemptionCallback`` is enabled by default. To disable it, add ``create_preemption_callback: False`` under exp_manager in the config YAML file.
+
+    Stragglers Detection
+    ----------------------
+
+    .. _exp_manager_straggler_det_support-label:
+
+    .. note::
+        Stragglers Detection feature is included in the optional NeMo resiliency package.
+
+    Distributed training can be affected by stragglers, which are workers that slow down the overall training process.
+    NeMo provides a straggler detection feature that can identify slower GPUs.
+
+    This feature is implemented in the ``StragglerDetectionCallback``, which is disabled by default.
+
+    The callback computes normalized GPU performance scores, which are scalar values ranging from 0.0 (worst) to 1.0 (best). 
+    A performance score can be interpreted as the ratio of current performance to reference performance.
+
+    There are two types of performance scores provided by the callback:
+        * Relative GPU performance score: The best-performing GPU in the workload is used as a reference.
+        * Individual GPU performance score: The best historical performance of the GPU is used as a reference.
+
+    Examples:
+        * If the relative performance score is 0.5, it means that a GPU is twice slower than the fastest GPU.
+        * If the individual performance score is 0.5, it means that a GPU is twice slower than its best observed performance.
+
+    If a GPU performance score drops below the specified threshold, it is identified as a straggler.
+
+    To enable straggler detection, add ``create_straggler_detection_callback: True`` under exp_manager in the config YAML file. 
+    You might also want to adjust the callback parameters:
+
+    .. code-block:: yaml
+
+        exp_manager:
+            ...
+            create_straggler_detection_callback: True
+            straggler_detection_callback_params:
+                report_time_interval: 300      # Interval [seconds] of the straggler check
+                calc_relative_gpu_perf: True   # Calculate relative GPU performance
+                calc_individual_gpu_perf: True # Calculate individual GPU performance
+                num_gpu_perf_scores_to_log: 5       # Log 5 best and 5 worst GPU performance scores, even if no stragglers are detected
+                gpu_relative_perf_threshold: 0.7    # Threshold for relative GPU performance scores
+                gpu_individual_perf_threshold: 0.7  # Threshold for individual GPU performance scores
+                stop_if_detected: True              # Terminate the workload if stragglers are detected
+
+    Straggler detection may require inter-rank synchronization and should be performed at regular intervals, such as every few minutes.
+
+
+.. Fault Tolerance
+    ---------------
+
+    .. _exp_manager_fault_tolerance_support-label:
+
+    .. note::
+        Fault Tolerance feature is included in the optional NeMo resiliency package.
+
+    When training Deep Neural Network (DNN models), faults may occur, hindering the progress of the entire training process.
+    This is particularly common in distributed, multi-node training scenarios, with many nodes and GPUs involved. 
+
+    NeMo incorporates a fault tolerance mechanism to detect training halts. 
+    In response, it can terminate a hung workload and, if requested, restart it from the last checkpoint.
+
+    Fault tolerance ("FT") relies on a special launcher (``ft_launcher``), which is a modified ``torchrun``. 
+    The FT launcher runs background processes called rank monitors. **You need to use ft_launcher to start 
+    your workload if you are using FT**. I.e., `NeMo-Framework-Launcher <https://github.com/NVIDIA/NeMo-Framework-Launcher>`_  
+    can be used to generate SLURM batch scripts with FT support. 
+
+    Each training process (rank) sends `heartbeats` to its monitor during training and validation steps.
+    If a rank monitor stops receiving `heartbeats`, a training failure is detected.
+
+    Fault detection is implemented in the ``FaultToleranceCallback`` and is disabled by default. 
+    To enable it, add a ``create_fault_tolerance_callback: True`` option under ``exp_manager`` in the 
+    config YAML file. Additionally, you can customize FT parameters by adding ``fault_tolerance`` section:
+
+    .. code-block:: yaml
+
+        exp_manager:
+            ...
+            create_fault_tolerance_callback: True
+            fault_tolerance:
+                initial_rank_heartbeat_timeout: 600  # wait for 10 minutes for the initial heartbeat
+                rank_heartbeat_timeout: 300  # wait for 5 minutes for subsequent heartbeats
+                calculate_timeouts: True # estimate more accurate timeouts based on observed intervals
+
+    Timeouts for fault detection need to be adjusted for a given workload:
+        * ``initial_rank_heartbeat_timeout`` should be long enough to allow for workload initialization.
+        * ``rank_heartbeat_timeout`` should be at least as long as the longest possible interval between steps. 
+
+    **Importantly, `heartbeats` are not sent during checkpoint loading and saving**, so time for 
+    checkpointing related operations should be taken into account.
+
+    If ``calculate_timeouts: True``, timeouts will be automatically estimated based on observed intervals.
+    Estimated timeouts take precedence over timeouts defined in the config file. **Timeouts are estimated
+    at the end of a training run when checkpoint loading and saving were observed.** Hence, in a multi-part
+    training started from scratch, estimated timeouts won't be available during the initial two runs.
+    Estimated timeouts are stored in a separate JSON file.
+
+    ``max_subsequent_job_failures`` allows for the automatic continuation of training on a SLURM cluster. 
+    This feature requires SLURM job to be scheduled with ``NeMo-Framework-Launcher``. If ``max_subsequent_job_failures`` 
+    value is `>0` continuation job is prescheduled. It will continue  the work until ``max_subsequent_job_failures`` 
+    subsequent jobs failed (SLURM job exit code is `!= 0`) or the training is completed successfully 
+    ("end of training" marker file is produced by the ``FaultToleranceCallback``, i.e. due to iters or time limit reached).
+
+    All FT configuration items summary:
+        * ``workload_check_interval`` (float, default=5.0) Periodic workload check interval [seconds] in the workload monitor.
+        * ``initial_rank_heartbeat_timeout`` (Optional[float], default=60.0 * 60.0) Timeout [seconds] for the first heartbeat from a rank. 
+        * ``rank_heartbeat_timeout`` (Optional[float], default=45.0 * 60.0) Timeout [seconds] for subsequent heartbeats from a rank. 
+        * ``calculate_timeouts`` (bool, default=True) Try to calculate ``rank_heartbeat_timeout`` and ``initial_rank_heartbeat_timeout`` 
+        based on the observed heartbeat intervals.
+        * ``safety_factor``: (float, default=5.0) When calculating the timeouts, multiply the maximum observed heartbeat interval 
+        by this factor to obtain the timeout estimate. Can be made smaller for stable environments and larger for unstable ones.  
+        * ``rank_termination_signal`` (signal.Signals, default=signal.SIGKILL) Signal used to terminate the rank when failure is detected.
+        * ``log_level`` (str, default='INFO') Log level for the FT client and server(rank monitor).
+        * ``max_rank_restarts`` (int, default=0) Used by FT launcher. Max number of restarts for a rank. 
+        If ``>0`` ranks will be restarted on existing nodes in case of a failure.
+        * ``max_subsequent_job_failures`` (int, default=0) Used by FT launcher. How many subsequent job failures are allowed until stopping autoresuming. 
+        ``0`` means do not auto-resume.
+        * ``additional_ft_launcher_args`` (str, default='') Additional FT launcher params (for advanced use).
+
+
+.. _nemo_multirun-label:
+
+Hydra Multi-Run with NeMo
+-------------------------
+
+When training neural networks, it is common to perform a hyperparameter search to improve the model’s performance on validation data.
+However, manually preparing a grid of experiments and managing all checkpoints and their metrics can be tedious.
+To simplify these tasks, NeMo integrates with `Hydra Multi-Run support <https://hydra.cc/docs/tutorials/basic/running_your_app/multi-run/>`_,
+providing a unified way to run a set of experiments directly from the configuration.
+
+There are certain limitations to this framework, which we list below:
+
+* All experiments are assumed to be run on a single GPU, and multi GPU for single run (model parallel models are not supported as of now).
+* NeMo Multi-Run currently supports only grid search over a set of hyperparameters. Support for advanced hyperparameter search strategies will be added in the future.
+* **NeMo Multi-Run requires one or more GPUs** to function and will not work without GPU devices.
+
+Config Setup
+~~~~~~~~~~~~
+
+In order to enable NeMo Multi-Run, we first update our YAML configs with some information to let Hydra know we expect to run multiple experiments from this one config -
+
+.. code-block:: yaml
+
+    # Required for Hydra launch of hyperparameter search via multirun
+    defaults:
+      - override hydra/launcher: nemo_launcher
+
+    # Hydra arguments necessary for hyperparameter optimization
+    hydra:
+      # Helper arguments to ensure all hyper parameter runs are from the directory that launches the script.
+      sweep:
+        dir: "."
+        subdir: "."
+
+      # Define all the hyper parameters here
+      sweeper:
+        params:
+          # Place all the parameters you wish to search over here (corresponding to the rest of the config)
+          # NOTE: Make sure that there are no spaces between the commas that separate the config params !
+          model.optim.lr: 0.001,0.0001
+          model.encoder.dim: 32,64,96,128
+          model.decoder.dropout: 0.0,0.1,0.2
+
+      # Arguments to the process launcher
+      launcher:
+        num_gpus: -1  # Number of gpus to use. Each run works on a single GPU.
+        jobs_per_gpu: 1  # If each GPU has large memory, you can run multiple jobs on the same GPU for faster results (until OOM).
+
+
+Next, we will setup the config for ``Experiment Manager``. When we perform hyper parameter search, each run may take some time to complete.
+We want to therefore avoid the case where a run ends (say due to OOM or timeout on the machine) and we need to redo all experiments.
+We therefore setup the experiment manager config such that every experiment has a unique "key", whose value corresponds to a single
+resumable experiment.
+
+Let us see how to setup such a unique "key" via the experiment name. Simply attach all the hyper parameter arguments to the experiment
+name as shown below -
+
+.. code-block:: yaml
+
+    exp_manager:
+      exp_dir: null  # Can be set by the user.
+
+      # Add a unique name for all hyper parameter arguments to allow continued training.
+      # NOTE: It is necessary to add all hyperparameter arguments to the name !
+      # This ensures successful restoration of model runs in case HP search crashes.
+      name: ${name}-lr-${model.optim.lr}-adim-${model.adapter.dim}-sd-${model.adapter.adapter_strategy.stochastic_depth}
+
+      ...
+      checkpoint_callback_params:
+        ...
+        save_top_k: 1  # Dont save too many .ckpt files during HP search
+        always_save_nemo: True # saves the checkpoints as nemo files for fast checking of results later
+      ...
+
+      # We highly recommend use of any experiment tracking took to gather all the experiments in one location
+      create_wandb_logger: True
+      wandb_logger_kwargs:
+        project: "<Add some project name here>"
+
+      # HP Search may crash due to various reasons, best to attempt continuation in order to
+      # resume from where the last failure case occurred.
+      resume_if_exists: true
+      resume_ignore_no_checkpoint: true
+
+
+Run a NeMo Multi-Run Configuration
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Once the config has been updated, we can now run it just like any normal Hydra script, with one special flag (``-m``).
+
+.. code-block:: bash
+
+    python script.py --config-path=ABC --config-name=XYZ -m \
+        trainer.max_steps=5000 \  # Any additional arg after -m will be passed to all the runs generated from the config !
+        ...
+
+Tips and Tricks
+---------------
+
+This section provides recommendations for using the Experiment Manager.
+
+Preserving disk space for a large number of experiments
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Some models may have a large number of parameters, making it very expensive to save numerous checkpoints on physical storage drives.
+For example, if you use the Adam optimizer, each PyTorch Lightning ".ckpt" file will be three times the size of just the model
+parameters. This can become exorbitant if you have multiple runs.
+
+In the above configuration, we explicitly set ``save_top_k: 1`` and ``always_save_nemo: True``. This limits the number of ".ckpt"
+files to just one and also saves a NeMo file, which contains only the model parameters without the optimizer state.
+This NeMo file can be restored immediately for further work.
+
+We can further save storage space by using NeMo's utility functions to automatically delete either ".ckpt" or NeMo files
+after a training run has finished. This is sufficient if you are collecting results in an experiment tracking tool and can
+simply rerun the best configuration after the search is completed.
+
+.. code-block:: python
+
+    # Import `clean_exp_ckpt` along with exp_manager
+    from nemo.utils.exp_manager import clean_exp_ckpt, exp_manager
+
+    @hydra_runner(...)
+    def main(cfg):
+        ...
+
+        # Keep track of the experiment directory
+        exp_log_dir = exp_manager(trainer, cfg.get("exp_manager", None))
+
+        ... add any training code here as needed ...
+
+        # Add following line to end of the training script
+        # Remove PTL ckpt file, and potentially also remove .nemo file to conserve storage space.
+        clean_exp_ckpt(exp_log_dir, remove_ckpt=True, remove_nemo=False)
+
+
+Debugging Multi-Run Scripts
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+When running Hydra scripts, you may encounter configuration issues that crash the program. In NeMo Multi-Run, a crash in
+any single run will not crash the entire program. Instead, we will note the error and proceed to the next job. Once all
+jobs are completed, we will raise the errors in the order they occurred, crashing the program with the first error’s stack trace.
+
+
+To debug NeMo Multi-Run, we recommend commenting out the entire hyperparameter configuration set inside ``sweep.params``.
+Instead, run a single experiment with the configuration, which will immediately raise the error.
+
+
+Experiment name cannot be parsed by Hydra
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Sometimes our hyperparameters include PyTorch Lightning ``trainer`` arguments, such as the number of steps, number of epochs,
+and whether to use gradient accumulation. When we attempt to add these as keys to the experiment manager's ``name``,
+Hydra may complain that ``trainer.xyz`` cannot be resolved.
+
+A simple solution is to finalize the Hydra config before you call ``exp_manager()`` as follows:
+
+.. code-block:: python
+
+    @hydra_runner(...)
+    def main(cfg):
+        # Make any changes as necessary to the config
+        cfg.xyz.abc = uvw
+
+        # Finalize the config
+        cfg = OmegaConf.resolve(cfg)
+
+        # Carry on as normal by calling trainer and exp_manager
+        trainer = pl.Trainer(**cfg.trainer)
+        exp_log_dir = exp_manager(trainer, cfg.get("exp_manager", None))
+        ...
+
+
+ExpManagerConfig
+----------------
+
+.. autoclass:: nemo.utils.exp_manager.ExpManagerConfig
+    :show-inheritance:
+    :members:
+    :member-order: bysource
+    :noindex:

docs/source/core/neural_modules.rst

@@ -0,0 +1,88 @@
+Neural Modules
+==============
+
+NeMo is built around Neural Modules, conceptual blocks of neural networks that take typed inputs and produce typed outputs. Such 
+modules typically represent data layers, encoders, decoders, language models, loss functions, or methods of combining activations.
+NeMo makes it easy to combine and re-use these building blocks while providing a level of semantic correctness checking via its neural 
+type system.
+
+.. note:: *All Neural Modules inherit from ``torch.nn.Module`` and are therefore compatible with the PyTorch ecosystem.*
+
+There are 3 types on Neural Modules:
+
+    - Regular modules
+    - Dataset/IterableDataset
+    - Losses
+
+Every Neural Module in NeMo must inherit from `nemo.core.classes.module.NeuralModule` class.
+
+.. autoclass:: nemo.core.classes.module.NeuralModule
+
+Every Neural Modules inherits the ``nemo.core.classes.common.Typing`` interface and needs to define neural types for its inputs and outputs.
+This is done by defining two properties: ``input_types`` and ``output_types``. Each property should return an ordered dictionary of 
+"port name"->"port neural type" pairs. Here is the example from :class:`~nemo.collections.asr.modules.ConvASREncoder` class:
+
+.. code-block:: python
+
+    @property
+    def input_types(self):
+        return OrderedDict(
+            {
+                "audio_signal": NeuralType(('B', 'D', 'T'), SpectrogramType()),
+                "length": NeuralType(tuple('B'), LengthsType()),
+            }
+        )
+
+    @property
+    def output_types(self):
+        return OrderedDict(
+            {
+                "outputs": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
+                "encoded_lengths": NeuralType(tuple('B'), LengthsType()),
+            }
+        )
+
+    @typecheck()
+    def forward(self, audio_signal, length=None):
+        ...
+
+The code snippet above means that ``nemo.collections.asr.modules.conv_asr.ConvASREncoder`` expects two arguments:
+    * First one, named ``audio_signal`` of shape ``[batch, dimension, time]`` with elements representing spectrogram values.
+    * Second one, named ``length`` of shape ``[batch]`` with elements representing lengths of corresponding signals.
+
+It also means that ``.forward(...)`` and ``__call__(...)`` methods each produce two outputs:
+    * First one, of shape ``[batch, dimension, time]`` but with elements representing encoded representation (``AcousticEncodedRepresentation`` class).
+    * Second one, of shape ``[batch]``, corresponding to their lengths.
+
+.. tip:: It is a good practice to define types and add ``@typecheck()`` decorator to your ``.forward()`` method after your module is ready for use by others.
+
+.. note:: The outputs of ``.forward(...)`` method will always be of type ``torch.Tensor`` or container of tensors and will work with any other Pytorch code. The type information is attached to every output tensor. If tensors without types is passed to your module, it will not fail, however the types will not be checked. Thus, it is recommended to define input/output types for all your modules, starting with data layers and add ``@typecheck()`` decorator to them.
+
+.. note:: To temporarily disable typechecking, you can enclose your code in ```with typecheck.disable_checks():``` statement.
+
+
+Dynamic Layer Freezing
+----------------------
+
+You can selectively freeze any modules inside a Nemo model by specifying a freezing schedule in the config yaml. Freezing stops any gradient updates
+to that module, so that its weights are not changed for that step. This can be useful for combatting catastrophic forgetting, for example
+when finetuning a large pretrained model on a small dataset.
+
+The default approach is to freeze a module for the first N training steps, but you can also enable freezing for a specific range of steps,
+for example, from step 20 - 100, or even activate freezing from some N until the end of training. You can also freeze a module for the entire training run.
+Dynamic freezing is specified in training steps, not epochs.
+
+To enable freezing, add the following to your config:
+
+.. code-block:: yaml
+
+  model:
+    ...
+    freeze_updates:
+      enabled: true  # set to false if you want to disable freezing
+      
+      modules:   # list all of the modules you want to have freezing logic for
+        encoder: 200       # module will be frozen for the first 200 training steps
+        decoder: [50, -1]  # module will be frozen at step 50 and will remain frozen until training ends
+        joint: [10, 100]   # module will be frozen between step 10 and step 100 (step >= 10 and step <= 100)
+        transcoder: -1     # module will be frozen for the entire training run

docs/source/core/neural_types.rst

@@ -0,0 +1,181 @@
+
+Neural Types
+============
+
+Motivation
+----------
+
+Neural Types describe the semantics, axis order, and dimensions of a tensor. The purpose of this type system is to catch semantic and
+dimensionality errors during model creation and facilitate module re-use.
+
+.. image:: whyntypes.gif
+  :width: 900
+  :alt: Neural Types Motivation
+
+``NeuralType`` class
+--------------------
+
+Neural Types perform semantic checks for modules and models inputs/outputs. They contain information about:
+
+    - Semantics of what is stored in the tensors. For example, logits, logprobs, audiosignal, embeddings, etc.
+    - Axes layout, semantic and (optionally) dimensionality. For example: ``[Batch, Time, Channel]``
+
+Types are implemented in ``nemo.core.neural_types.NeuralType`` class. When you instantiate an instance of this class, you
+are expected to include both *axes* information and *element type* information.
+
+.. autoclass:: nemo.core.neural_types.NeuralType
+    :noindex:
+
+Type Comparison Results
+-----------------------
+
+When comparing two neural types, the following comparison results are generated.
+
+.. autoclass:: nemo.core.neural_types.NeuralTypeComparisonResult
+    :noindex:
+
+Examples
+--------
+
+Long vs short notation
+~~~~~~~~~~~~~~~~~~~~~~
+
+NeMo's ``NeuralType`` class allows you to express axis semantics information in long and short form. Consider these two equivalent types. Both encoder 3 dimensional tensors and both contain elements of type ``AcousticEncodedRepresentation`` (this type is a typical output of ASR encoders).
+
+.. code-block:: python
+
+    long_version = NeuralType(
+            axes=(AxisType(AxisKind.Batch, None), AxisType(AxisKind.Dimension, None), AxisType(AxisKind.Time, None)),
+            elements_type=AcousticEncodedRepresentation(),
+        )
+    short_version = NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation())
+    assert long_version.compare(short_version) == NeuralTypeComparisonResult.SAME
+
+Transpose same
+~~~~~~~~~~~~~~
+
+Often it is useful to know if a simple transposition will solve type incompatibility. This is the case if the comparison result of two types equals ``nemo.core.neural_types.NeuralTypeComparisonResult.TRANSPOSE_SAME``.
+
+.. code-block:: python
+
+    type1 = NeuralType(axes=('B', 'T', 'C'))
+    type2 = NeuralType(axes=('T', 'B', 'C'))
+    assert type1.compare(type2) == NeuralTypeComparisonResult.TRANSPOSE_SAME
+    assert type2.compare(type1) == NeuralTypeComparisonResult.TRANSPOSE_SAME
+
+Note that in this example, we dropped ``elements_type`` argument of ``NeuralType`` constructor. If not supplied, the element type is ``VoidType``.
+
+``VoidType`` for elements
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Sometimes it is useful to express that elements' types don't matter but axes layout do. ``VoidType`` for elements can be used to express this.
+
+.. note:: ``VoidType`` is compatible with every other elements' type but not the other way around. See the following code snippet below for details.
+
+.. code-block:: python
+
+        btc_spctr = NeuralType(('B', 'T', 'C'), SpectrogramType())
+        btc_spct_bad = NeuralType(('B', 'T'), SpectrogramType())
+        # Note the VoidType for elements here
+        btc_void = NeuralType(('B', 'T', 'C'), VoidType())
+
+        # This is true because VoidType is compatible with every other element type (SpectrogramType in this case)
+        # And axes layout between btc_void and btc_spctr is the same
+        assert btc_void.compare(btc_spctr) == NeuralTypeComparisonResult.SAME
+        # These two types are incompatible because even though VoidType is used for elements on one side,
+        # the axes layout is different
+        assert btc_void.compare(btc_spct_bad) == NeuralTypeComparisonResult.INCOMPATIBLE
+        # Note that even though VoidType is compatible with every other type, other types are not compatible with VoidType!
+        # It is one-way compatibility
+        assert btc_spctr.compare(btc_void) == NeuralTypeComparisonResult.INCOMPATIBLE
+
+Element type inheritance
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+Neural types in NeMo support Python inheritance between element types. Consider an example where you want to develop a Neural Module which performs data augmentation for all kinds of spectrograms.
+In ASR, two types of spectrograms are frequently used: mel and mfcc. To express this, we will create 3 classes to express
+element's types: ``SpectrogramType``, ``MelSpectrogramType(SpectrogramType)``, ``MFCCSpectrogramType(SpectrogramType)``.
+
+.. code-block:: python
+
+        input = NeuralType(('B', 'D', 'T'), SpectrogramType())
+        out1 = NeuralType(('B', 'D', 'T'), MelSpectrogramType())
+        out2 = NeuralType(('B', 'D', 'T'), MFCCSpectrogramType())
+
+        # MelSpectrogram and MFCCSpectrogram are not interchangeable.
+        assert out1.compare(out2) == NeuralTypeComparisonResult.INCOMPATIBLE
+        assert out2.compare(out1) == NeuralTypeComparisonResult.INCOMPATIBLE
+        # Type comparison detects that MFCC/MelSpectrogramType is a kind of SpectrogramType and can be accepted.
+        assert input.compare(out1) == NeuralTypeComparisonResult.GREATER
+        assert input.compare(out2) == NeuralTypeComparisonResult.GREATER
+
+Custom element types
+~~~~~~~~~~~~~~~~~~~~
+
+It is possible to create user-defined element types to express the semantics of elements in your tensors. To do so, the user will need to inherit and implement abstract methods of the ``nemo.core.neural_types.elements.ElementType`` class
+
+.. autoclass:: nemo.core.neural_types.elements.ElementType
+    :noindex:
+
+Note that element types can be parametrized. Consider this example where it distinguishes between audio sampled at 8Khz and 16Khz.
+
+.. code-block:: python
+
+    audio16K = NeuralType(axes=('B', 'T'), elements_type=AudioSignal(16000))
+    audio8K = NeuralType(axes=('B', 'T'), elements_type=AudioSignal(8000))
+
+    assert audio8K.compare(audio16K) == NeuralTypeComparisonResult.SAME_TYPE_INCOMPATIBLE_PARAMS
+    assert audio16K.compare(audio8K) == NeuralTypeComparisonResult.SAME_TYPE_INCOMPATIBLE_PARAMS
+
+Enforcing dimensions
+~~~~~~~~~~~~~~~~~~~~
+
+In addition to specifying tensor layout and elements' semantics, neural types also allow you to enforce tensor dimensions.
+The user will have to use long notations to specify dimensions. Short notations only allows you to specify axes semantics and assumes
+arbitrary dimensions.
+
+.. code-block:: python
+
+        type1 = NeuralType(
+        (AxisType(AxisKind.Batch, 64), AxisType(AxisKind.Time, 10), AxisType(AxisKind.Dimension, 128)),
+        SpectrogramType(),
+        )
+        type2 = NeuralType(('B', 'T', 'C'), SpectrogramType())
+
+        # type2 will accept elements of type1 because their axes semantics match and type2 does not care about dimensions
+        assert type2.compare(type1), NeuralTypeComparisonResult.SAME
+        # type1 will not accept elements of type2 because it need dimensions to match strictly.
+        assert type1.compare(type2), NeuralTypeComparisonResult.DIM_INCOMPATIBLE
+
+Generic Axis kind
+~~~~~~~~~~~~~~~~~
+
+Sometimes (especially in the case of loss modules) it is useful to be able to specify a "generic" axis kind which will make it
+compatible with any other kind of axis. This is easy to express with Neural Types by using ``nemo.core.neural_types.axes.AxisKind.Any`` for axes.
+
+.. code-block:: python
+
+        type1 = NeuralType(('B', 'Any', 'Any'), SpectrogramType())
+        type2 = NeuralType(('B', 'T', 'C'), SpectrogramType())
+        type3 = NeuralType(('B', 'C', 'T'), SpectrogramType())
+
+        # type1 will accept elements of type2 and type3 because it only cares about element kind (SpectrogramType)
+        # number of axes (3) and that first one corresponds to batch
+        assert type1.compare(type2) == NeuralTypeComparisonResult.SAME
+        assert type1.compare(type3) == NeuralTypeComparisonResult.INCOMPATIBLE
+
+Container types
+~~~~~~~~~~~~~~~
+
+The NeMo-type system understands Python containers (lists). If your module returns a nested list of typed tensors, the way to express it is by
+using Python list notation and Neural Types together when defining your input/output types.
+
+The example below shows how to express that your module returns single output ("out") which is list of lists of two dimensional tensors of shape ``[batch, dimension]`` containing logits.
+
+.. code-block:: python
+
+    @property
+    def output_types(self):
+        return {
+            "out": [[NeuralType(('B', 'D'), LogitsType())]],
+        }