tutorials/02_NeMo_Adapters.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yS2xVGrLrphl"
      },
      "outputs": [],
      "source": [
        "\"\"\"\n",
        "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
        "\n",
        "Instructions for setting up Colab are as follows:\n",
        "1. Open a new Python 3 notebook.\n",
        "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
        "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
        "4. Run this cell to set up dependencies.\n",
        "\"\"\"\n",
        "# If you're using Google Colab and not running locally, run this cell.\n",
        "\n",
        "## Install dependencies\n",
        "!apt-get install sox libsndfile1 ffmpeg\n",
        "!pip install wget\n",
        "!pip install text-unidecode\n",
        "\n",
        "# ## Install NeMo\n",
        "BRANCH = 'r1.17.0'\n",
        "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]\n",
        "\n",
        "## Grab the config we'll use in this example\n",
        "!mkdir configs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ivKMObsjy9Om"
      },
      "source": [
        "# Adapters Support in NeMo Models\n",
        "\n",
        "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. \n",
        "\n",
        "As a potential solution to such a scenario, where fine-tuning a massive model is no longer feasible, we look to specialized [Adapters](https://arxiv.org/abs/1902.00751) 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.\n",
        "\n",
        "In this tutorial, we will discuss how to update any torch.nn.Module to support Adapters, and going further, how to enable NeMo models with Adapter support for their components.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lZ7mPouNEXjB"
      },
      "source": [
        "## What are Adapters?\n",
        "\n",
        "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 ($D$) to a small bottleneck dimension ($H$), such that $R^D \\text{->} R^H$, compute an activation (such as ReLU), finally mapping $R^H \\text{->} R^D$ with another Feedforward layer. This output is then added to the input via a simple residual connection.\n",
        "\n",
        "<div align=\"center\">\n",
        "  <img src=\"https://mermaid.ink/img/pako:eNptkLFqwzAQhl9F3ORAPDSjA4EUx6RgXEjbycpwWOdG1JaMfEoakrx7ZcfpUKrlxH_fz4d0gcoqggTqxp6qAzoW76k0Ipx1-WI6z3sRxyuRF1GOZ3KisK6d3YG8GFdZ9hRJeLbMDRmqvkRGpDLrTuiUiEWUigBtlyIVqzBnEqZ66I39dcX6iKytKXeUf-wn-286QoFeBMvmu0PTD-EfyXaQpP9JFmP_1XN4S3kfD8W4ue6o18pjc52gYQlzaMm1qFX4msuQSOADtSQhCdfaOupZgjS3QPpOIdNGabYOkhqbnuaAnu3b2VSQsPP0gFKNnw7bibr9AJkZdXU\" height=100% />\n",
        "</div>\n",
        "\n",
        "-----\n",
        "\n",
        "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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_kE1oh1_IdLW"
      },
      "source": [
        "## Emulating a standard architecture\n",
        "\n",
        "For this tutorial, the focus will be on demonstrating how to modify an existing architecture to add Adapter support.\n",
        "\n",
        "We will focus on a trivial model implemented using simple Multi-Layer Perceptrons. Still, the model itself will emulate a standard Encoder-Decoder architecture (commonly used in multiple domains, such as ASR, NLP, NMT etc). \n",
        "\n",
        "We will also skip the implementation of datasets, data loaders, losses, metrics, and the Pytorch Lightning \"steps\" (trainer, validation, test). "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3iYvsUFpIISX"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from nemo.core import NeuralModule, ModelPT\n",
        "\n",
        "from hydra.utils import instantiate\n",
        "from omegaconf import DictConfig, OmegaConf"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5qchb0kHZ6xV"
      },
      "outputs": [],
      "source": [
        "class MLP(torch.nn.Module):\n",
        "    def __init__(self, dim: int = 50):\n",
        "        super().__init__()\n",
        "\n",
        "        self.fc = torch.nn.Linear(dim, dim)\n",
        "        self.ln = torch.nn.LayerNorm(dim)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.fc(x)\n",
        "        x = self.ln(x)\n",
        "        return x\n",
        "\n",
        "class ResidualMLP(torch.nn.Module):\n",
        "  def __init__(self, dim: int, num_layers: int):\n",
        "    super().__init__()\n",
        "\n",
        "    self.dim = dim\n",
        "    self.num_layers = num_layers\n",
        "    self.layers = nn.ModuleList([MLP(dim) for _ in range(num_layers)])\n",
        "  \n",
        "  def forward(self, x):\n",
        "    input = x\n",
        "    for layer in self.layers:\n",
        "      x = layer(x)\n",
        "      x = x + input\n",
        "      input = x\n",
        "    return x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NgBYZMyFcJiO"
      },
      "source": [
        "-----\n",
        "Next we implement a simple model that has two \"modules\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P4lo4E-Abfm_"
      },
      "outputs": [],
      "source": [
        "class SimpleModel(ModelPT):\n",
        "    def __init__(self, cfg, trainer=None):\n",
        "        super().__init__(cfg, trainer=trainer)\n",
        "\n",
        "        self.encoder = instantiate(cfg.encoder)  # type: ResidualMLP\n",
        "        self.decoder = instantiate(cfg.decoder)  # type: ResidualMLP\n",
        "        self.projection = nn.Linear(self.decoder.dim, cfg.out_features)\n",
        "\n",
        "    def forward(self, x):\n",
        "        y = self.encoder(x)\n",
        "        z = self.decoder(y)\n",
        "        out = self.projection(z)\n",
        "        return out\n",
        "\n",
        "    def list_available_models(cls):\n",
        "        return []\n",
        "\n",
        "    def setup_training_data(self, train_data_config):\n",
        "        pass\n",
        "\n",
        "    def setup_validation_data(self, val_data_config):\n",
        "        pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oE6401-Bdj-K"
      },
      "source": [
        "## Initialize the basic model\n",
        "\n",
        "The above model is a simple residual MLP network with two components, an encoder, and a decoder block. It may not do so well on real-world tasks, but it is sufficient for this demonstration.\n",
        "\n",
        "Next, let's create a helper to generate a config for this model and create a new model using that config!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6xgMDKDvdbKw"
      },
      "outputs": [],
      "source": [
        "def get_classpath(cls):\n",
        "    return f'{cls.__module__}.{cls.__name__}'\n",
        "\n",
        "def get_model_config(dim=512):\n",
        "    config = OmegaConf.create(\n",
        "        {\n",
        "            'in_features': dim,\n",
        "            'out_features': 10,\n",
        "            'encoder': {'_target_': get_classpath(ResidualMLP), 'dim': dim, 'num_layers': 4},\n",
        "            'decoder': {'_target_': get_classpath(ResidualMLP), 'dim': dim, 'num_layers': 2},\n",
        "        }\n",
        "    )\n",
        "    return config"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HuCHbKM6eXgs"
      },
      "outputs": [],
      "source": [
        "dim = 512\n",
        "model_cfg = get_model_config(dim)\n",
        "model = SimpleModel(model_cfg)\n",
        "model.summarize()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ecba1X-6egs8"
      },
      "outputs": [],
      "source": [
        "# Check if the forward pass works !\n",
        "with torch.no_grad():\n",
        "  input_data = torch.randn(8, dim)\n",
        "  out = model(input_data)\n",
        "  print(out.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bQUitVFafW6q"
      },
      "source": [
        "# Incorporating Adapters - Module by Module\n",
        "\n",
        "Now that we have a basic Model that we can successfully perform a forward pass on, we can add adapter support to the Model and its modules - layer by layer.\n",
        "\n",
        "When considering the addition of adapter support, we work backward, going from the lowest level module used, and build a chain that forwards the methods of the adapters from the top-level Model to the bottom-level module(s) / layer(s)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-YcePlyLgUD4"
      },
      "source": [
        "# Adapter support in the lowest level module\n",
        "\n",
        "As we work backward in the model chain, we look at the `MLP` module that creates a `Linear` and `LayerNorm` layer. We now extend this MLP module with the `AdapterModuleMixin` that is available inside `nemo.core.adapter_mixins`.\n",
        "\n",
        "It is generally advised to directly update the code of the module, though there are other ways to implement this (shown later in the tutorial).\n",
        "\n",
        "-----\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "daSapfg0lpUj"
      },
      "source": [
        "## What is a `mixin`? \n",
        "A `mixin` is generally a term used to refer to a class that is **inherited by another class**, **adds some functionality to another class**, _but cannot be used on its own_. A mixin can be loosely considered a relatively safe way to incorporate additional functionality into a class via multiple inheritances. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WXE2cJrre9SN"
      },
      "outputs": [],
      "source": [
        "from nemo.core import adapter_mixins"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U7m8MScqkw8_"
      },
      "outputs": [],
      "source": [
        "help(adapter_mixins.AdapterModuleMixin)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ML_Uaig8iLKR"
      },
      "outputs": [],
      "source": [
        "# NOTE: See the *two* classes being inherited here !\n",
        "class MLP(torch.nn.Module, adapter_mixins.AdapterModuleMixin):\n",
        "    def __init__(self, dim: int = 50):\n",
        "        super().__init__()\n",
        "\n",
        "        self.fc = torch.nn.Linear(dim, dim)\n",
        "        self.ln = torch.nn.LayerNorm(dim)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.fc(x)\n",
        "        x = self.ln(x)\n",
        "\n",
        "        # The only necessary change to the module code !\n",
        "        if self.is_adapter_available():\n",
        "          x = self.forward_enabled_adapters(x)\n",
        "        return x\n",
        "\n",
        "    # add a utility method to calculate number of parameters (or we could simple extend nemo.core.NeuralModule instead)\n",
        "    @property\n",
        "    def num_weights(self):\n",
        "      num: int = 0\n",
        "      for p in self.parameters():\n",
        "          if p.requires_grad:\n",
        "              num += p.numel()\n",
        "      return num"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cUsZL7MJlKly"
      },
      "source": [
        "-----\n",
        "\n",
        "That's it! We now have an MLP layer that has nearly full adapter support! We will try out a few of the adapter functionalities below to get a teaser of what we can expect as we go further into this tutorial"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fV4Ww-kHlc2_"
      },
      "source": [
        "## Experimenting with a module level adapter\n",
        "\n",
        "We will now instantiate the newly augmented `MLP` model above and explore all the functionality that has been added via the `AdapterModuleMixin` class - without having to write too much supporting code!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OJBBbWFTmgni"
      },
      "source": [
        "-----\n",
        "\n",
        "First, let's create a `MLP` module and print the number of trainable parameters (before adding any adapters)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bUczqEM4lJYb"
      },
      "outputs": [],
      "source": [
        "mlp = MLP(dim)\n",
        "\n",
        "print(mlp)\n",
        "print(\"Num trainable parameters (without adapters):\", mlp.num_weights)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VY6rJroGmwEg"
      },
      "source": [
        "## Adapter Modules\n",
        "\n",
        "Next, let us import and add an adapter or two to this module! We first import `adapter_modules` from the NeMo `common` collections. This module contains pre-defined Adapter modules that can be attached to other torch.nn.Modules!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dol4-4vmmenZ"
      },
      "outputs": [],
      "source": [
        "from nemo.collections.common.parts import adapter_modules"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rtZtXNSHo1LB"
      },
      "outputs": [],
      "source": [
        "# Next we look at one of the adapter modules - the LinearAdapter\n",
        "linear_adapter = adapter_modules.LinearAdapter(in_features=dim, dim=5)\n",
        "print(linear_adapter)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "-----\n",
        "You will often not directly with this module, instead of passing the Config Dataclass to the `AdapterModuleMixin` methods. We see an example below - "
      ],
      "metadata": {
        "id": "0E2877IlIVoM"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## [Optional] Constructing Adapter Components\n",
        "\n",
        "Linear Adapter Modules are not the only type of adapters you can create ! In PyTorch, any torch.nn.Module can be made into an Adapter component.\n",
        "\n",
        "For example, you can potentially convert a pre-existing pytorch module into an adapter component. The below section is **optional**, but is recommended if you wish to create your own adapters.\n"
      ],
      "metadata": {
        "id": "8eH6mW792lkY"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "------\n",
        "First, let us start with a simple PyTorch module."
      ],
      "metadata": {
        "id": "lgsyaQHI3w5X"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class SimpleModule(torch.nn.Module):\n",
        "  def __init__(self, size: int):\n",
        "    super().__init__()\n",
        "    self.size = size\n",
        "    self.model = torch.nn.Sequential(\n",
        "        torch.nn.Linear(size, size, bias=False),\n",
        "        torch.nn.Identity(),\n",
        "    )\n",
        "  \n",
        "  def forward(self, x):\n",
        "    return self.model(x)"
      ],
      "metadata": {
        "id": "wAfA3r0b3fpi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Adapter Strategy\n",
        "\n",
        "Adapter modules are, at the end of the day, simply PyTorch modules. Just as PyTorch modules, they take some input tensors, perform some operation and then return some result.\n",
        "\n",
        "There are many ways to integrate adapters - add them as a residual, multiply pointwise, concatenate with the input (at the end or the beginning). The Adapter Strategy class determines how an adapter integrates with its input."
      ],
      "metadata": {
        "id": "hMKoxc0e5c14"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# The earlier LinearAdapter has a simple ResidualAddStrategy\n",
        "# Uncomment below to see the ResidualAddAdapterStrategy definition\n",
        "# help(linear_adapter.adapter_strategy)"
      ],
      "metadata": {
        "id": "1DVeRqH65IN8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Creating a custom Adapter Strategy\n",
        "\n",
        "Residual Add strategy can be considered the simple operation $f(x) = x + adapter(x)$ such that $adapter$'s initial outputs without training should be 0. \n",
        "\n",
        "In doing so, the output of the adapter augmented model is originally just $f(x) = x$, and so the model retains the exact performance of the original model (without any adapters).\n",
        "\n",
        "-----\n",
        "\n",
        "Below, we will create a Multiplication adapter strategy simply as a demonstration."
      ],
      "metadata": {
        "id": "5mWowiS269h8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from nemo.core.classes.mixins import adapter_mixin_strategies"
      ],
      "metadata": {
        "id": "teiTVBMq687x"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will implement a special `forward` method of adapters as follows"
      ],
      "metadata": {
        "id": "BS2W1a919KDr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Uncomment to see the definition of the AbstractAdapterStrategy\n",
        "# help(adapter_mixin_strategies.AbstractAdapterStrategy)"
      ],
      "metadata": {
        "id": "G9_bq05P9Izh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class MultiplicationAdapterStrategy(adapter_mixin_strategies.AbstractAdapterStrategy):\n",
        "\n",
        "  def __init__(self, scaling_factor: float = 1.0):\n",
        "    super().__init__()\n",
        "    self.scale = scaling_factor\n",
        "\n",
        "  def forward(self, input: torch.Tensor, adapter: torch.nn.Module, *, module: 'AdapterModuleMixin'):\n",
        "     # This is the forward method that takes in the previous input (here, its a tensor, but it can be a dictionary, a tuple, a class, anything really).\n",
        "     # The second argument is the adapter that is currently being applied to this input\n",
        "     # The final argument is the entire nn.Module that supports adapters.\n",
        "     # In this case, the final argument would be the entire `MLP` module\n",
        "     \n",
        "     # Equivalent to f(x) = x * adapter(x)\n",
        "     adapter_out = adapter(input)  # compute the adapter output from the input(s)\n",
        "     result = input * adapter_out\n",
        "\n",
        "     # Apply scaling factor. Equivalent to f(x) = scale * (x * adapter(x))\n",
        "     result = self.scale * result\n",
        "     return result\n"
      ],
      "metadata": {
        "id": "T3agPtjA3fst"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Design a corresponding dataclass for the Adapter Strategy\n",
        "\n",
        "In order to make usage of this class easier, you should create a Dataclass that can be used to create the strategy easily. We show an example below:"
      ],
      "metadata": {
        "id": "ZOS9Z3KVAGH2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from dataclasses import dataclass\n",
        "\n",
        "@dataclass\n",
        "class MultiplicationAdapterStrategyConfig:\n",
        "    scaling_factor: float = 1.0\n",
        "\n",
        "    # mandatory field\n",
        "    _target_: str = \"{0}.{1}\".format(\n",
        "        MultiplicationAdapterStrategy.__module__, MultiplicationAdapterStrategy.__name__\n",
        "    )  "
      ],
      "metadata": {
        "id": "MI4oYRYDAeqb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Creating a Custom Adapter Component\n",
        "\n",
        "Now that we have both the basic PyTorch module (`SimpleModule`) as well as the Adapter Strategy (`MultiplicationAdapterStrategy`), we can now construct a new adapter component.\n",
        "\n",
        "The prime difference between a basic PyTorch module and an adapter component is the `adapter_strategy` - it defines how an adapter integrates with the original input. "
      ],
      "metadata": {
        "id": "mZ22m_ZK_ifY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class SimpleModuleAdapter(SimpleModule, adapter_modules.AdapterModuleUtil):\n",
        "\n",
        "  def __init__(self, size: int, adapter_strategy: MultiplicationAdapterStrategy = None):\n",
        "    \"\"\"\n",
        "    The input arguments should match the original module so you can pass the inputs to the module.\n",
        "    It should also accept an adapter strategy.\n",
        "\n",
        "    We will then use the method `setup_adapter_strategy()` to prepare the component to be used as an adapter.\n",
        "    Note: Passing None to the strategy will let it pick a default strategy provided by the method\n",
        "    `get_default_strategy_config()`.\n",
        "    \"\"\"\n",
        "    super().__init__(size=size)\n",
        "\n",
        "    # Prepare the adapter strategy\n",
        "    self.setup_adapter_strategy(adapter_strategy)\n",
        "\n",
        "    # Initialize the weights to be 0 at init\n",
        "    self.reset_parameters()\n",
        "\n",
        "  # Note: In this case, because we didn't add new modules, nor change how the original forward works\n",
        "  # We dont need to subclass and override forward() !\n",
        "  \n",
        "  def reset_parameters(self):\n",
        "    # We normally want an adapter at initialization to have no effect on the output\n",
        "    # Therefore we replace the random uniform with a simple identity matrix, which will cause\n",
        "    # the output of the adapter to match the input\n",
        "    with torch.no_grad():\n",
        "      self.model[0].weight = torch.nn.Parameter(torch.eye(self.size))\n",
        "  \n",
        "\n",
        "  def get_default_strategy_config(self) -> 'dataclass':\n",
        "    \"\"\"\n",
        "    Make the default adapter strategy of this component be the `MultiplicationAdapterStrategy()`  \n",
        "    \"\"\"\n",
        "    return MultiplicationAdapterStrategyConfig()"
      ],
      "metadata": {
        "id": "wCcdXIix__Yq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "-----\n",
        "Let's quickly test whether the adapter behaves as expected"
      ],
      "metadata": {
        "id": "lYfjhWBtEBYj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "simple_adapter = SimpleModuleAdapter(size=5)\n",
        "multiplication_strategy = simple_adapter.adapter_strategy\n",
        "x = torch.randn(1, 5)\n",
        "adapter_x = simple_adapter(x)\n",
        "output = multiplication_strategy(input=x, adapter=simple_adapter, module=None)  # Normally you would pass the module here, but in this example can be skipped.\n",
        "print(\"Original input :\", x)\n",
        "print(\"Adapter output :\", adapter_x)\n",
        "print(\"Strategy output:\", output)"
      ],
      "metadata": {
        "id": "tVkikcmCEJun"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We see that the original input passes through the adapter, which results in the original values being returned successfully, and then the adapter strategy multiplies the two values (effectively computing the square of the input).\n",
        "\n",
        "This is a sufficient demonstration of creating custom adapters, and we would normally not perform elementwise multiplication as an adapter strategy. Normally we would prefer the output of the strategy to be equal to the original init, at least at initialization."
      ],
      "metadata": {
        "id": "mi7WpFgxHmGQ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Design a corresponding dataclass for the Adapter Component\n",
        "\n",
        "In order to make usage of this Adapter component easier, you should create a Dataclass that can be used to create the component easily. We show an example below:"
      ],
      "metadata": {
        "id": "OrmCQCX6ImKs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from typing import Optional\n",
        "\n",
        "@dataclass\n",
        "class SimpleModuleAdapterConfig:\n",
        "    size: int\n",
        "    adapter_strategy: Optional[MultiplicationAdapterStrategyConfig] = None\n",
        "\n",
        "    # mandatory field\n",
        "    _target_: str = \"{0}.{1}\".format(\n",
        "        SimpleModuleAdapter.__module__, SimpleModuleAdapter.__name__\n",
        "    )  "
      ],
      "metadata": {
        "id": "Ml17OoOVJOwR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sB9qPX1qoqJU"
      },
      "source": [
        "## Adding an Adapter Module\n",
        "\n",
        "Since `MLP` inherits `AdapterModuleMixin`, it also inherits a set of methods that perform adapter module manipulations - such as adding a new adapter.\n",
        "\n",
        "When users want to add an adapter, they can call `add_adapter()` with two specific arguments - `name` and `cfg`.\n",
        "\n",
        "Arguments - \n",
        "- `name`: A string name that must be **locally unique** (for modules) and **globally unique** (for models). The name may also support \":\" to delegate that the adapter belongs to specific modules only (this is discussed towards the end of the tutorial).\n",
        "- `cfg`: A dataclass / OmegaConf config that contains the `_target_` attribute pointing to the classpath of an Adapter Module, along with any additional required attributes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Du2uPUSdni9D"
      },
      "outputs": [],
      "source": [
        "mlp.add_adapter(name='adapter_1', cfg=adapter_modules.LinearAdapterConfig(in_features=dim, dim=5))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "92vSc2k_xuHk"
      },
      "outputs": [],
      "source": [
        "# Now check the new parameter count of this MLP module, it should be higher than the previous count\n",
        "print(\"New param count :\", mlp.num_weights)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eh4qTdTUyWTF"
      },
      "source": [
        "-----\n",
        "\n",
        "**Note**: You can add as many adapters as are needed! While in this tutorial, we will only add one, we usually recommend adding one adapter for every task you want to specialize in. \n",
        "\n",
        "Also, note that while it is possible to train multiple adapters at once (add many adapters, enable them all, then unfreeze them), we recommend training just one adapter per task."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w0CUr4m2sVRH"
      },
      "source": [
        "-----\n",
        "**Note**: If you try to add the same adapter multiple times, you will see the below error message! \n",
        "\n",
        "Remember, adapter names must be **locally** unique at the module level and **globally** unique at the model level!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p497PTCBsoIM"
      },
      "outputs": [],
      "source": [
        "# Uncomment to see the error message - \n",
        "# mlp.add_adapter(name='adapter_1', cfg=adapter_modules.LinearAdapterConfig(in_features=dim, dim=10))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EWgsqsCns4Yr"
      },
      "source": [
        "## Get all enabled Adapter Modules\n",
        "\n",
        "Next, we use `get_enabled_adapters()` to return a list of names of all the enabled adapters currently available to this module."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jXuGQoFus3qd"
      },
      "outputs": [],
      "source": [
        "mlp.get_enabled_adapters()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H-UeYsFXtQqj"
      },
      "source": [
        "## Set the state of Adapter Modules\n",
        "\n",
        "We can get the enabled adapter names with the above method, but how do we set whether an adapter module should be enabled or disabled? \n",
        "\n",
        "For that, we use the `set_enabled_adapter()` method. It has a few arguments - \n",
        "- `name`: An optional string name of an adapter, which will specifically enable or disable only that adapter. If no `name` is provided, all adapter modules will have their state set to the new value.\n",
        "- `enabled`: A bool, whether the adapter should be enabled or not.\n",
        "\n",
        "-----\n",
        "\n",
        "Enabling an adapter simply enables the forward pass of that adapter and nothing more. It does not freeze / unfreeze the weights of the adapter itself, allowing more complex interactions to occur in combination with other adapters.\n",
        "\n",
        "For example, one can add an adapter to a model, train it and then save the model. The restored model can then add yet another adapter. Prior to training this second adapter, the user can decide to utilize the outputs of the first adapter instead of the original model's outputs. To accomplish this, we can enable both adapters, but freezing the weights of the first adapter, and train just the second adapter."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0__7RFaWtPD6"
      },
      "outputs": [],
      "source": [
        "# Disable all adapters\n",
        "mlp.set_enabled_adapters(enabled=False)\n",
        "print(\"Enabled adapters :\", mlp.get_enabled_adapters())\n",
        "\n",
        "# Enable just one adapter\n",
        "mlp.set_enabled_adapters(name=\"adapter_1\", enabled=True)\n",
        "print(\"Enabled adapters :\", mlp.get_enabled_adapters())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KzT3kwvRubFT"
      },
      "source": [
        "## Check if Adapter Module(s) are available / enabled\n",
        "\n",
        "An extension of the above two methods is to check if the current module has any active adapter module or not. To do so, you can use `is_adapter_available()`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5quE42kAusiX"
      },
      "outputs": [],
      "source": [
        "mlp.is_adapter_available()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xRCobs9hvEeU"
      },
      "source": [
        "## Adapter functionality methods\n",
        "\n",
        "The above few methods form the core of the functionality to enable adapters to be added and modified to a module, but they don't use the added adapter modules!\n",
        "\n",
        "Therefore, the following functionality methods are used to leverage adapters properly and need not be overridden by the user (unless required for some special case)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6ybsJ9zRx6W7"
      },
      "source": [
        "### `forward_enabled_adapters()`\n",
        "To use these adapters, we utilize the `forward_adapter_modules()` method.\n",
        "\n",
        "To utilize any enabled adapters, the module that inherits `AdapterModuleMixin` should first check if any adapters are enabled and then call this method to forward the adapter modules on the input data. \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1F9IibFKv9va"
      },
      "outputs": [],
      "source": [
        "# Check `forward_enabled_adapters()`\n",
        "out = mlp.forward_enabled_adapters(input_data)\n",
        "print(out.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zzN0RlZI5M-l"
      },
      "source": [
        "### `forward_single_enabled_adapter_()`\n",
        "A method that can be sub-classed in order to provide custom logic for the forward pass of the adapters. For example, we may wish to provide some adapters with different set of inputs, or check whether we support an adapter type or not before we perform forward pass.\n",
        "\n",
        "It can be useful to check the type of the adapter, and then use the additional information prior to forwarding the input to any specific adapter."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YO2dXe7T5PIQ"
      },
      "outputs": [],
      "source": [
        "# Check `forward_single_enabled_adapter_()`\n",
        "adapter_name = mlp.get_enabled_adapters()[0]  # we have enabled just one adapter\n",
        "adapter_module = mlp.adapter_layer[adapter_name]  # get the adapter module with this name\n",
        "adapter_strategy = adapter_module.adapter_strategy  # get the adapter strategy for this adapter\n",
        "\n",
        "out = mlp.forward_single_enabled_adapter_(input_data, adapter_module, adapter_name=adapter_name, adapter_strategy=adapter_strategy)\n",
        "print(out.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bU-OS3wk7FW9"
      },
      "source": [
        "-----\n",
        "For further information about adapter forward pass, adapter strategy please refer to the documentation section for adapters."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b1mPF9qVx-HA"
      },
      "source": [
        "### `unfreeze_enabled_adapters()`\n",
        "One of the core benefits of adapters is that they do not need the entire model to be trained. We can freeze the rest of the original model/modules and train the adapter modules themselves. \n",
        "\n",
        "We can do this in two steps - \n",
        "- Call model.freeze() (at the highest level)\n",
        "- Call `unfreeze_enabled_adapters()` that will recursively unfreeze just the adapter modules that are enabled."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t6Qdk8YHw5W_"
      },
      "outputs": [],
      "source": [
        "# First setup some utility functions (this is part of NeuralModule)\n",
        "def freeze(m):\n",
        "    for param in m.parameters():\n",
        "      param.requires_grad = False\n",
        "    m.eval()\n",
        "\n",
        "def unfreeze(m):\n",
        "    for param in m.parameters():\n",
        "      param.requires_grad = True\n",
        "    m.train()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "95KGRNwxxlia"
      },
      "outputs": [],
      "source": [
        "freeze(mlp)\n",
        "print(\"MLP frozen params :\", mlp.num_weights)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y3xjntCgxhee"
      },
      "outputs": [],
      "source": [
        "# Check `unfreeze_enabled_adapters()` - param count should be lower than the previous total (original + adapter)\n",
        "mlp.unfreeze_enabled_adapters()\n",
        "print(\"MLP unfrozen adapter params :\", mlp.num_weights)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vQjz_RIilg1L"
      },
      "source": [
        "# Adapter support in intermediate level modules\n",
        "\n",
        "Above, we discussed many of the methods and capabilities added to a simple nn.Module via the `AdapterModuleMixin`. However, this module was the lowest building block in the model. Next, we will look into how to \"dispatch\" the calls from the intermediate module to the lower modules.\n",
        "\n",
        "We will aim for simplicity in this tutorial, modifying the minimal amount of code as possible. However, it is entirely possible to add much more sophisticated handling of intermediate layer dispatches to lower level modules."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DsHudP-bz3bj"
      },
      "source": [
        "## Intermediate modules that are instantiated via config\n",
        "\n",
        "Currently, we have a 3 level model -- \n",
        "\n",
        "`Top level Model (SimpleModel) -> Intermediate level Module (ResidualMLP) -> Bottom level Module (MLP)`. \n",
        "\n",
        "-----\n",
        "\n",
        "As you may have noticed, in earlier primer tutorials (NeMo Model Primer), we recommend the Model utilize configs to instantiate its intermediate modules. This allows users to swap in equivalent modules via the config and enjoy the rest of the utility of the Model itself without too many code changes.\n",
        "\n",
        "For such \"penultimate\" modules, we recommend creating a separate Adapter supported module that extends the original module rather than modifying the original module itself. This is merely a preference to avoid cluttering the original module code and can be ignored if the user wishes.\n",
        "\n",
        "For this guide, we will show the recommended setup so that best practices can be followed."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_f-Y-Oxm1tjH"
      },
      "source": [
        "## Creating an Adapter-compatible \"Penultimate\" module\n",
        "\n",
        "First, we create the new Adapter compatible module as a separate class."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DC-L1SsVllLy"
      },
      "outputs": [],
      "source": [
        "# NOTE: We subclass the original ResidualMLP, and add in the AdapterModuleMixin too\n",
        "class ResidualMLPAdapter(ResidualMLP, adapter_mixins.AdapterModuleMixin):\n",
        "  pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CuIlnrSW2m9t"
      },
      "source": [
        "## Overriding the adapter methods\n",
        "\n",
        "Next, we override a few adapter methods, such that we dispatch these methods to all the blocks of `MLP` inside of the `ResidualMLP` module.\n",
        "\n",
        "Therefore, this will create/update the state / forward an adapter module inside the `MLP` modules!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N3bcnXIy2mY9"
      },
      "outputs": [],
      "source": [
        "from typing import List, Optional\n",
        "\n",
        "class ResidualMLPAdapter(ResidualMLP, adapter_mixins.AdapterModuleMixin):\n",
        "  def add_adapter(self, name: str, cfg: DictConfig):\n",
        "      # call the same method on each `MLP` layer, collecting results\n",
        "      for layer in self.layers:\n",
        "        layer.add_adapter(name, cfg)\n",
        "      \n",
        "  def get_enabled_adapters(self) -> List[str]:\n",
        "      # call the same method on each `MLP` layer, collecting results\n",
        "      enabled_adapters = set([])\n",
        "      for layer in self.layers:\n",
        "        names = layer.get_enabled_adapters()\n",
        "        enabled_adapters.update(names)\n",
        "      return list(enabled_adapters)\n",
        "  \n",
        "  def set_enabled_adapters(self, name: Optional[str], enabled: bool):\n",
        "      # call the same method on each `MLP` layer, collecting results\n",
        "      for layer in self.layers:\n",
        "        layer.set_enabled_adapters(name, enabled)\n",
        "  \n",
        "  def is_adapter_available(self) -> bool:\n",
        "      # call the same method on each `MLP` layer, collecting results\n",
        "      is_available = any([layer.is_adapter_available() for layer in self.layers])\n",
        "      return is_available"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_-RHMG4W4pSY"
      },
      "source": [
        "## Register the new adapter\n",
        "\n",
        "When we subclass a module to add Adapter functionality, it is essential to register such modules with the Adapter registry so that many convenient functions can be used later on. The adapter registry is a global collection of the base class and adapter compatible class that can be used later to update model configs more easily.\n",
        "\n",
        "The steps below are : \n",
        "- Check if the registry has the base class via `get_registered_adapter()`.\n",
        "- If it returns None, then register the base class and its compatible adapter class via `register_adapter()`.\n",
        "\n",
        "-----\n",
        "\n",
        "**Note**: that while in this trivial case, our penultimate module is, in fact the intermediate module, there may be real-world models with many more intermediate modules. In such a case, you may update such intermediate modules by directly extending `AdapterModuleMixin` and following the above steps without creating a new subclass. In such cases, you can also skip registering for these modules."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sOUo-b042S9m"
      },
      "outputs": [],
      "source": [
        "if adapter_mixins.get_registered_adapter(ResidualMLP) is None:\n",
        "  adapter_mixins.register_adapter(ResidualMLP, ResidualMLPAdapter)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BSaAnSxC6A6f"
      },
      "source": [
        "-----\n",
        "\n",
        "That's all it takes to add support for intermediate modules! While adding the same (or similar) code for all intermediate modules may seem a little redundant, that is only because we are implementing the most naive dispatching. \n",
        "\n",
        "There are many interesting approaches to building adapters, such as adapters for only attention layers (before or after) or only for the final feed-forward (in conventional attention-based blocks). As such, intermediate layers have the total flexibility to dispatch these functions to lower layers."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_yeL62Js6sGb"
      },
      "source": [
        "# Adapters support at top level Model\n",
        "\n",
        "Finally, after dispatching the above methods from the intermediate modules to the bottom module, we need to perform the final dispatch from the Model itself to the first (or penultimate if moving backward) module.\n",
        "\n",
        "In this case, we will subclass a different mixin class than the one we have been using till now. Instead of `AdapterModuleMixin`, we will instead subclass `AdapterModelPTMixin` - which has some functionality built into it to manage model level config (including saving and restoring adapter compatible models !)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aQpl_ML78Pw2"
      },
      "source": [
        "## Extending `AdapterModelPTMixin`\n",
        "\n",
        "There are two ways in which one can inherit the top level mixin - \n",
        "\n",
        "(1)  directly extend it in your current Model class\n",
        "\n",
        "(2)  create a class that implements the additional functionality, then inherit that class.\n",
        "\n",
        "It might seem that option (2) is a more roundabout way of achieving (1). Still, it is done to keep the logic of adapter management outside of the complicated Model codebase since the Model itself is involved with many important details such as the setup of the modules, data loaders, optimizer/schedulers, losses, metrics, and the Pytorch Lightning \"steps\" - training, validation, and testing steps.\n",
        "\n",
        "-----\n",
        "We will follow option (2) in this tutorial for clarity. Also, note that we will create new subclasses for each step for transparency (and to avoid burdening you with too much information at once). It is preferred to do all these steps within just one new class."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "85eprLBu55a0"
      },
      "outputs": [],
      "source": [
        "class SimpleModelAdapter(adapter_mixins.AdapterModelPTMixin):\n",
        "  pass"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j38nVuRoCiMb"
      },
      "outputs": [],
      "source": [
        "help(adapter_mixins.AdapterModelPTMixin)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5ILtH3PT8PQs"
      },
      "source": [
        "## Overriding methods for selective dispatch\n",
        "\n",
        "There was an apparent distinction between how to dispatch adapter calls to the subsequent module. This was because the modules were homogeneous and shared typical behavior.\n",
        "\n",
        "However, there is no reason for the subsequent layers to share standard behavior at the model level. Think in terms of `encoder` vs. `decoder` Transformer layers - they are significantly different modules! Therefore why should their adapters be similar?\n",
        "\n",
        "At the top level, we can utilize input from the user to determine how to construct adapters for such logically heterogeneous components. The following sections will describe how we can utilize **global** and **module** level adapters to separate the behavior and construction of adapters for each **component** of the Model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5EwR-Onc-xbr"
      },
      "source": [
        "## Overriding setup_adapters()\n",
        "\n",
        "When a model must be restored, it must carefully load the parameters of all the modules within it. Till now, we have been able to augment torch.nn.Module(s) with Adapter information and capabilities, but we have not preserved this information anywhere.\n",
        "\n",
        "Therefore if we were to turn off the notebook and try to restore a saved checkpoint, restoration would fail because the new model does not have the information of the adapters that it previously added, and so state dict matching will fail.\n",
        "\n",
        "This issue is resolved by overriding `setup_adapters()` and calling it inside the Model constructor."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rsSAb99a9l86"
      },
      "outputs": [],
      "source": [
        "# Import the class explicitly to make instance checks easier\n",
        "from nemo.core.classes.mixins.adapter_mixins import AdapterModuleMixin\n",
        "\n",
        "class SimpleModelAdapterSetupAdapters(SimpleModelAdapter):\n",
        "  def setup_adapters(self):\n",
        "    # First check that any of the modules support adapters or not\n",
        "    supports_adapters = False\n",
        "\n",
        "    # Check the inheriting class' modules supports adapters or not\n",
        "    if hasattr(self, 'encoder') and isinstance(self.encoder, AdapterModuleMixin):\n",
        "        supports_adapters |= True\n",
        "\n",
        "    if hasattr(self, 'decoder') and isinstance(self.decoder, AdapterModuleMixin):\n",
        "        supports_adapters |= True\n",
        "\n",
        "    # If any class supports it, try to restore adapters\n",
        "    if supports_adapters:\n",
        "        super().setup_adapters()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kvmY6XJtApK4"
      },
      "source": [
        "-----\n",
        "\n",
        "In this step, we merely check if any of the modules we have created support adapters or not. If any of them do, we call the super() method to try to restore any adapters if needed."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z2eTu1pFA2j-"
      },
      "source": [
        "## Overriding add_adapter()\n",
        "\n",
        "Next up, we override `add_adapter`. Before jumping to the code, we must first discuss the types of adapters supported in NeMo.\n",
        "\n",
        "- `Global Adapters`: These adapters share their name and functionality with all supported modules. They are helpful when you are sure that one adapter can be shared between multiple model components. For example, the encoder and decoder share the same adapter.\n",
        "- `Module Adapters`: These adapters are specific to each module they are designated to and, therefore, cannot share their name across multiple components of a model. They are denoted via the adapter name of the format `{module_name}:{adapter_name}`.\n",
        "\n",
        "**Note**: After the module adapter is added, it can be referred to simply by the `adapter_name` part of its name. There is no need to provide `module_name` again since it is guaranteed that all adapter names are globally unique at the Model level."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W47wBr-eB6IR"
      },
      "source": [
        "-----\n",
        "\n",
        "It is entirely up to the user whether they should support just `Global Adapters`, `Module Adapters` or both. For the purpose of this tutorial, we will support both, and also add support for a `Default Module Adapter` for the `encoder`.\n",
        "\n",
        "**Note**: In order to easily distinguish between a `Global` and `Module` adapter, use the convenient method `resolve_adapter_module_name_(name)`. We encourage use of the property `adapter_module_names` to determine the valid adapter modules that can be used."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yCFZt7GCCQ1K"
      },
      "outputs": [],
      "source": [
        "class SimpleModelAdapterAddAdapter(SimpleModelAdapterSetupAdapters):\n",
        "\n",
        "  def add_adapter(self, name: str, cfg: DictConfig):\n",
        "      # Setup the config first. At the model level, super does not automatically call any of the subclass methods\n",
        "      # It just sets up the model.cfg for users\n",
        "      super().add_adapter(name, cfg)\n",
        "\n",
        "      # Resolve module name and adapter name\n",
        "      module_name, adapter_name = self.resolve_adapter_module_name_(name)\n",
        "\n",
        "      # Try to retrieve global adapter config\n",
        "      global_config = self._get_global_cfg()\n",
        "\n",
        "      # forward the method call to the individual modules\n",
        "      # If module name is empty, it is a default and global adapter, otherwise it is a module adapter\n",
        "      if (module_name == '' and global_config.get('encoder_adapter', True)) or (module_name == 'encoder'):\n",
        "          self.encoder.add_adapter(name, cfg)\n",
        "\n",
        "      if (module_name == '' and global_config.get('decoder_adapter', False)) or (module_name == 'decoder'):\n",
        "          self.decoder.add_adapter(name, cfg)\n",
        "    \n",
        "  def resolve_adapter_module_name_(self, name: str) -> (str, str):\n",
        "      # resolve name and module\n",
        "      module_name, adapter_name = super().resolve_adapter_module_name_(name)\n",
        "\n",
        "      # '' as module name means \"default module\"\n",
        "      # assert that the module name (if provided) is valid - default, encoder or decoder\n",
        "      valid_module_names = self.adapter_module_names  # Get the list of supported adapter modules from property\n",
        "      if module_name not in valid_module_names:\n",
        "          raise ValueError(f\"Provided module name `{module_name}` is not in valid list : {valid_module_names}\")\n",
        "\n",
        "      return (module_name, adapter_name)\n",
        "\n",
        "  def _get_global_cfg(self):\n",
        "      # Utility method to get a default \"global\" adapter config (can be given any value by the user in this config)\n",
        "      global_config = DictConfig({})\n",
        "      if 'adapters' in self.cfg and self.adapter_global_cfg_key in self.cfg.adapters:\n",
        "          global_config = self.adapter_cfg[self.adapter_global_cfg_key]\n",
        "      return global_config\n",
        "\n",
        "  @property\n",
        "  def adapter_module_names(self) -> List[str]:\n",
        "      module_names = super().adapter_module_names  # \"Default\" adapter module: ''\n",
        "      module_names.extend(['encoder', 'decoder'])  # Add support for `encoder` and `decoder` modules\n",
        "      return module_names\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HHAZFqo7DrCV"
      },
      "source": [
        "-----\n",
        "\n",
        "Since that was a lot of code, let's break it down. First we define the utility method `_get_global_cfg()` that tries to get the \"model.cfg.adapters.global_cfg\" sub-config from the model config. This config is user-defined and can be used by the user to configure any logic as necessary. If it is not found, a default dict is created instead.\n",
        "\n",
        "Next, we override the `resolve_adapter_module_name_(name)` method. This method in the base class takes a string name and tries to split it into `module_name` and `adapter_name`. We override this method to assert some valid `module_name's.\n",
        "\n",
        "-----\n",
        "\n",
        "Finally, we override `add_adapter` - first, we call super() to update the config. Next, we call our overridden `resolve_adapter_module_name_(name)` to check if the provided adapter name is valid or not. Then we gather the Adapter `global_cfg` if it exists in the model config.\n",
        "\n",
        "With all this information, we are now ready to add adapters as needed. We have some \"user-defined\" logic, such that we add an encoder adapter if either. \n",
        "- The user provides a \"Global\" adapter with a default module name, OR has set the value of `global_cfg.encoder_adapter` to True (True by default). This means that at least the encoder adapter will always be added by default.\n",
        "- The user provides a `Module` adapter with the `decoder` module name, OR has set the `global_cfg.decoder_adapter` flag to True explicitly (False by default)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4fVSt9wPFK7U"
      },
      "source": [
        "## Overriding get_enabled_adapters()\n",
        "\n",
        "Next, we will override the `get_enabled_adapters()` method. This is often simple enough, where we need only to check if the Model components support adapters or not, and then if they do, gather and collate the results of those nodules."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RuwE_mLXAko6"
      },
      "outputs": [],
      "source": [
        "class SimpleModelAdapterGetEnabledAdapters(SimpleModelAdapterAddAdapter):\n",
        "\n",
        "  def get_enabled_adapters(self) -> List[str]:\n",
        "      enabled_adapters = super().get_enabled_adapters()\n",
        "\n",
        "      # Forward the method call to the individual modules\n",
        "      if isinstance(self.encoder, AdapterModuleMixin):\n",
        "          encoder_adapters = self.encoder.get_enabled_adapters()\n",
        "          enabled_adapters.extend(encoder_adapters)\n",
        "\n",
        "      if isinstance(self.decoder, AdapterModuleMixin):\n",
        "          decoder_adapters = self.decoder.get_enabled_adapters()\n",
        "          enabled_adapters.extend(decoder_adapters)\n",
        "\n",
        "      return enabled_adapters"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O0fgtkVqFvFS"
      },
      "source": [
        "## Overriding set_enabled_adapters()\n",
        "\n",
        "Similar to the above, we only need to check if the components support adapters or not and then dispatch the call to those components if they are supported.\n",
        "\n",
        "**Note**: We will perform a logic check here instead of the usual inheritance check. An inheritance check alone does not mean the component has added an adapter or not - remember, we have two use cases of `default/global/module encoder` and `module decoder` adapters. So we need to check for those conditions.\n",
        "\n",
        "**Note 2**: As you may remember, set_enabled_adapters() does take the value `None` for the name to set the state of all adapters. However, `resolve_adapter_module_name(name)_` must always accept a valid string name. So care must be taken not to pass `None` to that method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nuuzuedEFt4Y"
      },
      "outputs": [],
      "source": [
        "class SimpleModelAdapterSetEnabledAdapters(SimpleModelAdapterGetEnabledAdapters):\n",
        "\n",
        "  def set_enabled_adapters(self, name: Optional[str] = None, enabled: bool = True):\n",
        "      # check if valid model with some adapter support\n",
        "      super().set_enabled_adapters(name, enabled)\n",
        "\n",
        "      # Resolve module name and adapter name\n",
        "      if name is not None:\n",
        "          module_name, _ = self.resolve_adapter_module_name_(name)\n",
        "      else:\n",
        "          module_name = None\n",
        "\n",
        "      # Try to retrieve global adapter config\n",
        "      global_config = self._get_global_cfg()\n",
        "\n",
        "      # Forward the method call to the individual modules\n",
        "      # Note the OR checks - \n",
        "      # if module_name is None - ie explicitly None was passed, set the state for all modules\n",
        "      # if module name was '' or 'encoder, or if `global_cfg.encoder_adapter` was true, or module_name was '' or 'encoder', forward to encoder.\n",
        "      # if `global_cfg.decoder_adapter` was true, or module_name was 'decoder', forward to decoder.\n",
        "      # The user can chose to simplify this logic, or add more complex logic as required.\n",
        "      if name is None or global_config.get('encoder_adapter', True) or module_name in ('', 'encoder'):\n",
        "        if self.encoder.is_adapter_available():\n",
        "          self.encoder.set_enabled_adapters(name, enabled)\n",
        "\n",
        "      if name is None or global_config.get('decoder_adapter', False) or module_name == 'decoder':\n",
        "        if self.decoder.is_adapter_available():\n",
        "          self.decoder.set_enabled_adapters(name, enabled)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wPf5ga4YH0lb"
      },
      "source": [
        "## Overriding `check_valid_model_with_adapter_support_()`\n",
        "\n",
        "In the above implementation, we implicitly check that the components of the model support adapters or not and carry forward. This is fine, but we may want to perform more strict checks so as to give meaningful warnings or errors about invalid combinations of actions. \n",
        "\n",
        "For this purpose, we provide `check_valid_model_with_adapter_support()_`. This method is called before nearly all adapter operations, and attempts to assert some truths. Users can raise any error or warning here to notify the user about invalid operations / configurations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l4vM2u3FHvaN"
      },
      "outputs": [],
      "source": [
        "from nemo.utils import logging, logging_mode\n",
        "\n",
        "class SimpleModelAdapterFinal(SimpleModelAdapterSetEnabledAdapters):\n",
        "\n",
        "  def check_valid_model_with_adapter_support_(self):\n",
        "      global_cfg = DictConfig({})\n",
        "      if self.adapter_global_cfg_key in self.adapter_cfg:\n",
        "          global_cfg = self.adapter_cfg[self.adapter_global_cfg_key]\n",
        "\n",
        "      encoder_adapter = global_cfg.get('encoder_adapter', True)\n",
        "      decoder_adapter = global_cfg.get('decoder_adapter', False)\n",
        "\n",
        "      if encoder_adapter and not hasattr(self, 'encoder'):\n",
        "          logging.warning(\"Encoder not available\", mode=logging_mode.ONCE)\n",
        "      elif encoder_adapter and not isinstance(self.encoder, AdapterModuleMixin):\n",
        "          logging.warning(\"Encoder does not support adapters !\", mode=logging_mode.ONCE)\n",
        "\n",
        "      if decoder_adapter and not hasattr(self, 'decoder'):\n",
        "          logging.warning(\"Decoder is not available\", mode=logging_mode.ONCE)\n",
        "      elif decoder_adapter and not isinstance(self.decoder, AdapterModuleMixin):\n",
        "          logging.warning(\"Decoder does not support adapters !\", mode=logging_mode.ONCE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eTVGaXpqIzWH"
      },
      "source": [
        "## Updating the Model\n",
        "\n",
        "After the top-level Model mixin class has been implemented separately, we can easily attach it to our original Model. For the tutorial's sake, we will duplicate the code here, but know that you can also subclass the Model and override its `__init__` method for similar functionality."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Tl4aJr9zIx_v"
      },
      "outputs": [],
      "source": [
        "# Note how we added `SimpleModelAdapterFinal` to the class inheritance scheme.\n",
        "# The only other change is the addition of `self.setup_adapters()` to the __init__ method.\n",
        "class SimpleModel(ModelPT, SimpleModelAdapterFinal):\n",
        "    def __init__(self, cfg, trainer=None):\n",
        "        super().__init__(cfg, trainer=trainer)\n",
        "\n",
        "        self.encoder = instantiate(cfg.encoder)  # type: ResidualMLP\n",
        "        self.decoder = instantiate(cfg.decoder)  # type: ResidualMLP\n",
        "        self.projection = nn.Linear(self.decoder.dim, cfg.out_features)\n",
        "\n",
        "        # NOTE: The only important change - calling `setup_adapters()` !\n",
        "        self.setup_adapters()\n",
        "\n",
        "    def forward(self, x):\n",
        "        y = self.encoder(x)\n",
        "        z = self.decoder(y)\n",
        "        out = self.projection(z)\n",
        "        return out\n",
        "\n",
        "    def list_available_models(cls):\n",
        "        return []\n",
        "\n",
        "    def setup_training_data(self, train_data_config):\n",
        "        pass\n",
        "\n",
        "    def setup_validation_data(self, val_data_config):\n",
        "        pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9HqU-8pBJjlP"
      },
      "source": [
        "-----\n",
        "\n",
        "And there we go ! Just subclassing the new mixin, as well as calling `setup_adapters()` is the only modification necessary !"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFlt2sCPJgvU"
      },
      "outputs": [],
      "source": [
        "old_config = get_model_config(dim)\n",
        "model = SimpleModel(old_config)\n",
        "model.summarize()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6LioRW_TJ6Qs"
      },
      "source": [
        "-----\n",
        "\n",
        "Now, let us try adding a `decoder` Module adapter to this Model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9W5B_0wiJ0V5"
      },
      "outputs": [],
      "source": [
        "# This cell will error out if uncommented\n",
        "# model.add_adapter(\"decoder:adapter_1\", cfg=adapter_modules.LinearAdapterConfig(in_features=dim, dim=5))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gW71Jhq_KhD7"
      },
      "source": [
        "-----\n",
        "\n",
        "It fails with an error `Encoder does not support adapters !`. This is because, if you recall, the original config of this model (old_config) has the classpath to the `ResidualMLP` class, but not the `ResidualMLPAdapter` class!\n",
        "\n",
        "This can be easily rectified because we have already registered this class properly (refer to `Register the new adapter` sub-section)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oyvHTCjpKgKi"
      },
      "outputs": [],
      "source": [
        "def get_adapter_model_config() -> DictConfig:\n",
        "  config = get_model_config()\n",
        "\n",
        "  # Find the metadata in the registry, and get the correct adapter capable class path\n",
        "  enc_adapter_metadata = adapter_mixins.get_registered_adapter(config.encoder._target_)\n",
        "  if enc_adapter_metadata is not None:\n",
        "      print(\"Updated encoder to support adapters !\")\n",
        "      config.encoder._target_ = enc_adapter_metadata.adapter_class_path\n",
        "\n",
        "  # Find the metadata in the registry, and get the correct adapter capable class path\n",
        "  dec_adapter_metadata = adapter_mixins.get_registered_adapter(config.decoder._target_)\n",
        "  if dec_adapter_metadata is not None:\n",
        "      print(\"Updated decoder to support adapters !\")\n",
        "      config.decoder._target_ = dec_adapter_metadata.adapter_class_path\n",
        "\n",
        "  return config"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "K5HD91NwKMTp"
      },
      "outputs": [],
      "source": [
        "new_config = get_adapter_model_config()\n",
        "model = SimpleModel(new_config)\n",
        "model.summarize()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YvnZW4Y8LnEH"
      },
      "source": [
        "-----\n",
        "Now let us again try to add a `decoder` Module Adapter -"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TCYE7afTLcYY"
      },
      "outputs": [],
      "source": [
        "model.add_adapter('decoder:adapter_1', cfg=adapter_modules.LinearAdapterConfig(in_features=dim, dim=5))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y4enwIFsL3BE"
      },
      "source": [
        "-----\n",
        "You will see multiple log messages stating `adapter_1` was added (one for each block). Now lets check if the adapter is in the correct module - "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lXBD5l1PL0zm"
      },
      "outputs": [],
      "source": [
        "print(\"Encoder adapter available :\", model.encoder.is_adapter_available())\n",
        "print(\"Decoder adapter available :\", model.decoder.is_adapter_available())\n",
        "print(\"Decoder adapter(s) :\", model.decoder.get_enabled_adapters())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iKRWH0ZSMBpR"
      },
      "outputs": [],
      "source": [
        "model.summarize()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sw8ZXBBYMnVH"
      },
      "source": [
        "## Preparing to train an adapter\n",
        "\n",
        "Finally, now that our model has adapter capabilities, we can train just the adapter while freezing the rest of the model.\n",
        "\n",
        "In the following section, we show how to perform this setup."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fyNB7Jh0Mp2q"
      },
      "outputs": [],
      "source": [
        "# disable all adapters, enable just one adapter that we want to train\n",
        "model.set_enabled_adapters(enabled=False)\n",
        "model.set_enabled_adapters('adapter_1', enabled=True)  # note : we directly use the adapter_name of adapter_1\n",
        "\n",
        "# freeze all the weights, unfreeze just the enabled adapters\n",
        "model.freeze()\n",
        "model.unfreeze_enabled_adapters()\n",
        "\n",
        "print()\n",
        "model.summarize()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u1aA0GxI_kKc"
      },
      "source": [
        "## Save and Restore Adapters\n",
        "\n",
        "Now that we have implemented this Model, we can always use `model.save_to()` to save a NeMo model with full adapter support. Let us try to save and then restore this adapter model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AH_v_InH_xhO"
      },
      "outputs": [],
      "source": [
        "model.save_to('full_model.nemo')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LH0FNfVHAs4u"
      },
      "outputs": [],
      "source": [
        "!ls -d -- *.nemo"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0Z_kajXeBDyI"
      },
      "outputs": [],
      "source": [
        "new_model = ModelPT.restore_from('full_model.nemo')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XvTufIQWBNgi"
      },
      "outputs": [],
      "source": [
        "new_model.decoder.get_enabled_adapters()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9OzpPgfxBda1"
      },
      "source": [
        "-----\n",
        "\n",
        "While we can save and restore the entire model, we don't need to. Think about adapters for a moment - they are additional modules on top of a base model. For every new adapter we add, it would not be feasible to save and restore the entire model to file (especially if the models are multiple billions of parameters !).\n",
        "\n",
        "Next, we discuss how to save and restore just the modules themselves into separate .pt files using `save_adapters()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jbf23dgFB8NU"
      },
      "outputs": [],
      "source": [
        "model.save_adapters('adapters.pt', name=None)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n-dkPsxNCEGv"
      },
      "outputs": [],
      "source": [
        "!du -sh adapters.pt full_model.nemo"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wJc-dA7aCH5I"
      },
      "source": [
        "-----\n",
        "As you can see, the whole model is much larger than just the adapter modules themselves. This can be used to share just the adapter modules with others without using ample disk space for the entire model.\n",
        "\n",
        "Next, we show how to restore such adapter checkpoint to a new model using `load_adapters()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3w2ob6_ECgxf"
      },
      "outputs": [],
      "source": [
        "new_config = get_adapter_model_config()\n",
        "model_2 = SimpleModel(new_config)\n",
        "model_2.summarize()  # no adapters in basic model with adapter support"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZmceRIylCnry"
      },
      "outputs": [],
      "source": [
        "model_2.load_adapters('adapters.pt', name=None, map_location='cpu')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2qEjDjcLEfC2"
      },
      "outputs": [],
      "source": [
        "model_2.freeze()\n",
        "model_2.unfreeze_enabled_adapters()\n",
        "model_2.summarize()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a80Pt-_5E_kN"
      },
      "outputs": [],
      "source": [
        "model_2.get_enabled_adapters()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lzk6hXGuEdqm"
      },
      "source": [
        "-----\n",
        "You will note that we passed None to the above model, restoring all the adapters, but what if we want to restore just one adapter from the checkpoint? Then we can pass the `name` argument to the `load_adapters()` method. \n",
        "\n",
        "Remember: The name passed here must be the name used to create the adapter itself, so for module-level adapters, the `module_name` must also be provided."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qg48swJpMiS5"
      },
      "source": [
        "# Further reading\n",
        "\n",
        "For further information about Adapters, please refer to the NeMo Documentation page for Adapters. It also includes sections that detail how to create your Adapter modules.\n",
        "\n",
        "For further details regarding how adapters can be used for, please refer to the following articles - \n",
        "- [Parameter-Efficient Transfer Learning for NLP](https://arxiv.org/abs/1902.00751)\n",
        "- [Exploiting Adapters for Cross-lingual Low-resource Speech Recognition](https://arxiv.org/abs/2105.11905)\n",
        "- [Efficient Adapter Transfer of Self-Supervised Speech Models for Automatic Speech Recognition](https://arxiv.org/abs/2202.03218)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iz2wF3cd-6MF"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "8eH6mW792lkY"
      ],
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}