RASMUS/Finnish-ASR-Canary-v2
02.2k
1{2 "cells": [3 {4 "cell_type": "code",5 "execution_count": null,6 "metadata": {7 "id": "yS2xVGrLrphl"8 },9 "outputs": [],10 "source": [11 "\"\"\"\n",12 "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",13 "\n",14 "Instructions for setting up Colab are as follows:\n",15 "1. Open a new Python 3 notebook.\n",16 "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",17 "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",18 "4. Run this cell to set up dependencies.\n",19 "\"\"\"\n",20 "# If you're using Google Colab and not running locally, run this cell.\n",21 "\n",22 "## Install dependencies\n",23 "!apt-get install sox libsndfile1 ffmpeg\n",24 "!pip install wget\n",25 "!pip install text-unidecode\n",26 "\n",27 "# ## Install NeMo\n",28 "BRANCH = 'main'\n",29 "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]\n",30 "\n",31 "## Grab the config we'll use in this example\n",32 "!mkdir configs"33 ]34 },35 {36 "cell_type": "markdown",37 "metadata": {38 "id": "ivKMObsjy9Om"39 },40 "source": [41 "# Adapters Support in NeMo Models\n",42 "\n",43 "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",44 "\n",45 "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",46 "\n",47 "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"48 ]49 },50 {51 "cell_type": "markdown",52 "metadata": {53 "id": "lZ7mPouNEXjB"54 },55 "source": [56 "## What are Adapters?\n",57 "\n",58 "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",59 "\n",60 "<div align=\"center\">\n",61 " <img src=\"https://mermaid.ink/img/pako:eNptkLFqwzAQhl9F3ORAPDSjA4EUx6RgXEjbycpwWOdG1JaMfEoakrx7ZcfpUKrlxH_fz4d0gcoqggTqxp6qAzoW76k0Ipx1-WI6z3sRxyuRF1GOZ3KisK6d3YG8GFdZ9hRJeLbMDRmqvkRGpDLrTuiUiEWUigBtlyIVqzBnEqZ66I39dcX6iKytKXeUf-wn-286QoFeBMvmu0PTD-EfyXaQpP9JFmP_1XN4S3kfD8W4ue6o18pjc52gYQlzaMm1qFX4msuQSOADtSQhCdfaOupZgjS3QPpOIdNGabYOkhqbnuaAnu3b2VSQsPP0gFKNnw7bibr9AJkZdXU\" height=100% />\n",62 "</div>\n",63 "\n",64 "-----\n",65 "\n",66 "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."67 ]68 },69 {70 "cell_type": "markdown",71 "metadata": {72 "id": "_kE1oh1_IdLW"73 },74 "source": [75 "## Emulating a standard architecture\n",76 "\n",77 "For this tutorial, the focus will be on demonstrating how to modify an existing architecture to add Adapter support.\n",78 "\n",79 "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",80 "\n",81 "We will also skip the implementation of datasets, data loaders, losses, metrics, and the Pytorch Lightning \"steps\" (trainer, validation, test). "82 ]83 },84 {85 "cell_type": "code",86 "execution_count": null,87 "metadata": {88 "id": "3iYvsUFpIISX"89 },90 "outputs": [],91 "source": [92 "import os\n",93 "import torch\n",94 "import torch.nn as nn\n",95 "from nemo.core import NeuralModule, ModelPT\n",96 "\n",97 "from hydra.utils import instantiate\n",98 "from omegaconf import DictConfig, OmegaConf\n",99 "\n",100 "# As of PyTorch 2.6, torch.load defaults to weights_only=True. Adapter checkpoints\n",101 "# contain non-tensor objects (OmegaConf DictConfig), so we must allow full loading.\n",102 "# Only do this with trusted checkpoint files.\n",103 "os.environ[\"TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD\"] = \"1\""104 ]105 },106 {107 "cell_type": "code",108 "execution_count": null,109 "metadata": {110 "id": "5qchb0kHZ6xV"111 },112 "outputs": [],113 "source": [114 "class MLP(torch.nn.Module):\n",115 " def __init__(self, dim: int = 50):\n",116 " super().__init__()\n",117 "\n",118 " self.fc = torch.nn.Linear(dim, dim)\n",119 " self.ln = torch.nn.LayerNorm(dim)\n",120 "\n",121 " def forward(self, x):\n",122 " x = self.fc(x)\n",123 " x = self.ln(x)\n",124 " return x\n",125 "\n",126 "class ResidualMLP(torch.nn.Module):\n",127 " def __init__(self, dim: int, num_layers: int):\n",128 " super().__init__()\n",129 "\n",130 " self.dim = dim\n",131 " self.num_layers = num_layers\n",132 " self.layers = nn.ModuleList([MLP(dim) for _ in range(num_layers)])\n",133 " \n",134 " def forward(self, x):\n",135 " input = x\n",136 " for layer in self.layers:\n",137 " x = layer(x)\n",138 " x = x + input\n",139 " input = x\n",140 " return x"141 ]142 },143 {144 "cell_type": "markdown",145 "metadata": {146 "id": "NgBYZMyFcJiO"147 },148 "source": [149 "-----\n",150 "Next we implement a simple model that has two \"modules\""151 ]152 },153 {154 "cell_type": "code",155 "execution_count": null,156 "metadata": {157 "id": "P4lo4E-Abfm_"158 },159 "outputs": [],160 "source": [161 "class SimpleModel(ModelPT):\n",162 " def __init__(self, cfg, trainer=None):\n",163 " super().__init__(cfg, trainer=trainer)\n",164 "\n",165 " self.encoder = instantiate(cfg.encoder) # type: ResidualMLP\n",166 " self.decoder = instantiate(cfg.decoder) # type: ResidualMLP\n",167 " self.projection = nn.Linear(self.decoder.dim, cfg.out_features)\n",168 "\n",169 " def forward(self, x):\n",170 " y = self.encoder(x)\n",171 " z = self.decoder(y)\n",172 " out = self.projection(z)\n",173 " return out\n",174 "\n",175 " def list_available_models(cls):\n",176 " return []\n",177 "\n",178 " def setup_training_data(self, train_data_config):\n",179 " pass\n",180 "\n",181 " def setup_validation_data(self, val_data_config):\n",182 " pass"183 ]184 },185 {186 "cell_type": "markdown",187 "metadata": {188 "id": "oE6401-Bdj-K"189 },190 "source": [191 "## Initialize the basic model\n",192 "\n",193 "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",194 "\n",195 "Next, let's create a helper to generate a config for this model and create a new model using that config!"196 ]197 },198 {199 "cell_type": "code",200 "execution_count": null,201 "metadata": {202 "id": "6xgMDKDvdbKw"203 },204 "outputs": [],205 "source": [206 "def get_classpath(cls):\n",207 " return f'{cls.__module__}.{cls.__name__}'\n",208 "\n",209 "def get_model_config(dim=512):\n",210 " config = OmegaConf.create(\n",211 " {\n",212 " 'in_features': dim,\n",213 " 'out_features': 10,\n",214 " 'encoder': {'_target_': get_classpath(ResidualMLP), 'dim': dim, 'num_layers': 4},\n",215 " 'decoder': {'_target_': get_classpath(ResidualMLP), 'dim': dim, 'num_layers': 2},\n",216 " }\n",217 " )\n",218 " return config"219 ]220 },221 {222 "cell_type": "code",223 "execution_count": null,224 "metadata": {225 "id": "HuCHbKM6eXgs"226 },227 "outputs": [],228 "source": [229 "dim = 512\n",230 "model_cfg = get_model_config(dim)\n",231 "model = SimpleModel(model_cfg)\n",232 "model.summarize()"233 ]234 },235 {236 "cell_type": "code",237 "execution_count": null,238 "metadata": {239 "id": "Ecba1X-6egs8"240 },241 "outputs": [],242 "source": [243 "# Check if the forward pass works !\n",244 "with torch.no_grad():\n",245 " input_data = torch.randn(8, dim)\n",246 " out = model(input_data)\n",247 " print(out.shape)"248 ]249 },250 {251 "cell_type": "markdown",252 "metadata": {253 "id": "bQUitVFafW6q"254 },255 "source": [256 "# Incorporating Adapters - Module by Module\n",257 "\n",258 "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",259 "\n",260 "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)."261 ]262 },263 {264 "cell_type": "markdown",265 "metadata": {266 "id": "-YcePlyLgUD4"267 },268 "source": [269 "# Adapter support in the lowest level module\n",270 "\n",271 "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",272 "\n",273 "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",274 "\n",275 "-----\n"276 ]277 },278 {279 "cell_type": "markdown",280 "metadata": {281 "id": "daSapfg0lpUj"282 },283 "source": [284 "## What is a `mixin`? \n",285 "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. "286 ]287 },288 {289 "cell_type": "code",290 "execution_count": null,291 "metadata": {292 "id": "WXE2cJrre9SN"293 },294 "outputs": [],295 "source": [296 "from nemo.core import adapter_mixins"297 ]298 },299 {300 "cell_type": "code",301 "execution_count": null,302 "metadata": {303 "id": "U7m8MScqkw8_"304 },305 "outputs": [],306 "source": [307 "help(adapter_mixins.AdapterModuleMixin)"308 ]309 },310 {311 "cell_type": "code",312 "execution_count": null,313 "metadata": {314 "id": "ML_Uaig8iLKR"315 },316 "outputs": [],317 "source": [318 "# NOTE: See the *two* classes being inherited here !\n",319 "class MLP(torch.nn.Module, adapter_mixins.AdapterModuleMixin):\n",320 " def __init__(self, dim: int = 50):\n",321 " super().__init__()\n",322 "\n",323 " self.fc = torch.nn.Linear(dim, dim)\n",324 " self.ln = torch.nn.LayerNorm(dim)\n",325 "\n",326 " def forward(self, x):\n",327 " x = self.fc(x)\n",328 " x = self.ln(x)\n",329 "\n",330 " # The only necessary change to the module code !\n",331 " if self.is_adapter_available():\n",332 " x = self.forward_enabled_adapters(x)\n",333 " return x\n",334 "\n",335 " # add a utility method to calculate number of parameters (or we could simple extend nemo.core.NeuralModule instead)\n",336 " @property\n",337 " def num_weights(self):\n",338 " num: int = 0\n",339 " for p in self.parameters():\n",340 " if p.requires_grad:\n",341 " num += p.numel()\n",342 " return num"343 ]344 },345 {346 "cell_type": "markdown",347 "metadata": {348 "id": "cUsZL7MJlKly"349 },350 "source": [351 "-----\n",352 "\n",353 "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"354 ]355 },356 {357 "cell_type": "markdown",358 "metadata": {359 "id": "fV4Ww-kHlc2_"360 },361 "source": [362 "## Experimenting with a module level adapter\n",363 "\n",364 "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!"365 ]366 },367 {368 "cell_type": "markdown",369 "metadata": {370 "id": "OJBBbWFTmgni"371 },372 "source": [373 "-----\n",374 "\n",375 "First, let's create a `MLP` module and print the number of trainable parameters (before adding any adapters)"376 ]377 },378 {379 "cell_type": "code",380 "execution_count": null,381 "metadata": {382 "id": "bUczqEM4lJYb"383 },384 "outputs": [],385 "source": [386 "mlp = MLP(dim)\n",387 "\n",388 "print(mlp)\n",389 "print(\"Num trainable parameters (without adapters):\", mlp.num_weights)"390 ]391 },392 {393 "cell_type": "markdown",394 "metadata": {395 "id": "VY6rJroGmwEg"396 },397 "source": [398 "## Adapter Modules\n",399 "\n",400 "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!"401 ]402 },403 {404 "cell_type": "code",405 "execution_count": null,406 "metadata": {407 "id": "dol4-4vmmenZ"408 },409 "outputs": [],410 "source": [411 "from nemo.collections.common.parts import adapter_modules"412 ]413 },414 {415 "cell_type": "code",416 "execution_count": null,417 "metadata": {418 "id": "rtZtXNSHo1LB"419 },420 "outputs": [],421 "source": [422 "# Next we look at one of the adapter modules - the LinearAdapter\n",423 "linear_adapter = adapter_modules.LinearAdapter(in_features=dim, dim=5)\n",424 "print(linear_adapter)"425 ]426 },427 {428 "cell_type": "markdown",429 "metadata": {430 "id": "0E2877IlIVoM"431 },432 "source": [433 "-----\n",434 "You will often not directly with this module, instead of passing the Config Dataclass to the `AdapterModuleMixin` methods. We see an example below - "435 ]436 },437 {438 "cell_type": "markdown",439 "metadata": {440 "id": "8eH6mW792lkY"441 },442 "source": [443 "## [Optional] Constructing Adapter Components\n",444 "\n",445 "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",446 "\n",447 "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"448 ]449 },450 {451 "cell_type": "markdown",452 "metadata": {453 "id": "lgsyaQHI3w5X"454 },455 "source": [456 "------\n",457 "First, let us start with a simple PyTorch module."458 ]459 },460 {461 "cell_type": "code",462 "execution_count": null,463 "metadata": {464 "id": "wAfA3r0b3fpi"465 },466 "outputs": [],467 "source": [468 "class SimpleModule(torch.nn.Module):\n",469 " def __init__(self, size: int):\n",470 " super().__init__()\n",471 " self.size = size\n",472 " self.model = torch.nn.Sequential(\n",473 " torch.nn.Linear(size, size, bias=False),\n",474 " torch.nn.Identity(),\n",475 " )\n",476 " \n",477 " def forward(self, x):\n",478 " return self.model(x)"479 ]480 },481 {482 "cell_type": "markdown",483 "metadata": {484 "id": "hMKoxc0e5c14"485 },486 "source": [487 "### Adapter Strategy\n",488 "\n",489 "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",490 "\n",491 "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."492 ]493 },494 {495 "cell_type": "code",496 "execution_count": null,497 "metadata": {498 "id": "1DVeRqH65IN8"499 },500 "outputs": [],501 "source": [502 "# The earlier LinearAdapter has a simple ResidualAddStrategy\n",503 "# Uncomment below to see the ResidualAddAdapterStrategy definition\n",504 "# help(linear_adapter.adapter_strategy)"505 ]506 },507 {508 "cell_type": "markdown",509 "metadata": {510 "id": "5mWowiS269h8"511 },512 "source": [513 "### Creating a custom Adapter Strategy\n",514 "\n",515 "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",516 "\n",517 "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",518 "\n",519 "-----\n",520 "\n",521 "Below, we will create a Multiplication adapter strategy simply as a demonstration."522 ]523 },524 {525 "cell_type": "code",526 "execution_count": null,527 "metadata": {528 "id": "teiTVBMq687x"529 },530 "outputs": [],531 "source": [532 "from nemo.core.classes.mixins import adapter_mixin_strategies"533 ]534 },535 {536 "cell_type": "markdown",537 "metadata": {538 "id": "BS2W1a919KDr"539 },540 "source": [541 "We will implement a special `forward` method of adapters as follows"542 ]543 },544 {545 "cell_type": "code",546 "execution_count": null,547 "metadata": {548 "id": "G9_bq05P9Izh"549 },550 "outputs": [],551 "source": [552 "# Uncomment to see the definition of the AbstractAdapterStrategy\n",553 "# help(adapter_mixin_strategies.AbstractAdapterStrategy)"554 ]555 },556 {557 "cell_type": "code",558 "execution_count": null,559 "metadata": {560 "id": "T3agPtjA3fst"561 },562 "outputs": [],563 "source": [564 "class MultiplicationAdapterStrategy(adapter_mixin_strategies.AbstractAdapterStrategy):\n",565 "\n",566 " def __init__(self, scaling_factor: float = 1.0):\n",567 " super().__init__()\n",568 " self.scale = scaling_factor\n",569 "\n",570 " def forward(self, input: torch.Tensor, adapter: torch.nn.Module, *, module: 'AdapterModuleMixin'):\n",571 " # 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",572 " # The second argument is the adapter that is currently being applied to this input\n",573 " # The final argument is the entire nn.Module that supports adapters.\n",574 " # In this case, the final argument would be the entire `MLP` module\n",575 " \n",576 " # Equivalent to f(x) = x * adapter(x)\n",577 " adapter_out = adapter(input) # compute the adapter output from the input(s)\n",578 " result = input * adapter_out\n",579 "\n",580 " # Apply scaling factor. Equivalent to f(x) = scale * (x * adapter(x))\n",581 " result = self.scale * result\n",582 " return result\n"583 ]584 },585 {586 "cell_type": "markdown",587 "metadata": {588 "id": "ZOS9Z3KVAGH2"589 },590 "source": [591 "### Design a corresponding dataclass for the Adapter Strategy\n",592 "\n",593 "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:"594 ]595 },596 {597 "cell_type": "code",598 "execution_count": null,599 "metadata": {600 "id": "MI4oYRYDAeqb"601 },602 "outputs": [],603 "source": [604 "from dataclasses import dataclass\n",605 "\n",606 "@dataclass\n",607 "class MultiplicationAdapterStrategyConfig:\n",608 " scaling_factor: float = 1.0\n",609 "\n",610 " # mandatory field\n",611 " _target_: str = \"{0}.{1}\".format(\n",612 " MultiplicationAdapterStrategy.__module__, MultiplicationAdapterStrategy.__name__\n",613 " ) "614 ]615 },616 {617 "cell_type": "markdown",618 "metadata": {619 "id": "mZ22m_ZK_ifY"620 },621 "source": [622 "### Creating a Custom Adapter Component\n",623 "\n",624 "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",625 "\n",626 "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. "627 ]628 },629 {630 "cell_type": "code",631 "execution_count": null,632 "metadata": {633 "id": "wCcdXIix__Yq"634 },635 "outputs": [],636 "source": [637 "class SimpleModuleAdapter(SimpleModule, adapter_modules.AdapterModuleUtil):\n",638 "\n",639 " def __init__(self, size: int, adapter_strategy: MultiplicationAdapterStrategy = None):\n",640 " \"\"\"\n",641 " The input arguments should match the original module so you can pass the inputs to the module.\n",642 " It should also accept an adapter strategy.\n",643 "\n",644 " We will then use the method `setup_adapter_strategy()` to prepare the component to be used as an adapter.\n",645 " Note: Passing None to the strategy will let it pick a default strategy provided by the method\n",646 " `get_default_strategy_config()`.\n",647 " \"\"\"\n",648 " super().__init__(size=size)\n",649 "\n",650 " # Prepare the adapter strategy\n",651 " self.setup_adapter_strategy(adapter_strategy)\n",652 "\n",653 " # Initialize the weights to be 0 at init\n",654 " self.reset_parameters()\n",655 "\n",656 " # Note: In this case, because we didn't add new modules, nor change how the original forward works\n",657 " # We dont need to subclass and override forward() !\n",658 " \n",659 " def reset_parameters(self):\n",660 " # We normally want an adapter at initialization to have no effect on the output\n",661 " # Therefore we replace the random uniform with a simple identity matrix, which will cause\n",662 " # the output of the adapter to match the input\n",663 " with torch.no_grad():\n",664 " self.model[0].weight = torch.nn.Parameter(torch.eye(self.size))\n",665 " \n",666 "\n",667 " def get_default_strategy_config(self) -> 'dataclass':\n",668 " \"\"\"\n",669 " Make the default adapter strategy of this component be the `MultiplicationAdapterStrategy()` \n",670 " \"\"\"\n",671 " return MultiplicationAdapterStrategyConfig()"672 ]673 },674 {675 "cell_type": "markdown",676 "metadata": {677 "id": "lYfjhWBtEBYj"678 },679 "source": [680 "-----\n",681 "Let's quickly test whether the adapter behaves as expected"682 ]683 },684 {685 "cell_type": "code",686 "execution_count": null,687 "metadata": {688 "id": "tVkikcmCEJun"689 },690 "outputs": [],691 "source": [692 "simple_adapter = SimpleModuleAdapter(size=5)\n",693 "multiplication_strategy = simple_adapter.adapter_strategy\n",694 "x = torch.randn(1, 5)\n",695 "adapter_x = simple_adapter(x)\n",696 "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",697 "print(\"Original input :\", x)\n",698 "print(\"Adapter output :\", adapter_x)\n",699 "print(\"Strategy output:\", output)"700 ]701 },702 {703 "cell_type": "markdown",704 "metadata": {705 "id": "mi7WpFgxHmGQ"706 },707 "source": [708 "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",709 "\n",710 "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."711 ]712 },713 {714 "cell_type": "markdown",715 "metadata": {716 "id": "OrmCQCX6ImKs"717 },718 "source": [719 "### Design a corresponding dataclass for the Adapter Component\n",720 "\n",721 "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:"722 ]723 },724 {725 "cell_type": "code",726 "execution_count": null,727 "metadata": {728 "id": "Ml17OoOVJOwR"729 },730 "outputs": [],731 "source": [732 "from typing import Optional\n",733 "\n",734 "@dataclass\n",735 "class SimpleModuleAdapterConfig:\n",736 " size: int\n",737 " adapter_strategy: Optional[MultiplicationAdapterStrategyConfig] = None\n",738 "\n",739 " # mandatory field\n",740 " _target_: str = \"{0}.{1}\".format(\n",741 " SimpleModuleAdapter.__module__, SimpleModuleAdapter.__name__\n",742 " ) "743 ]744 },745 {746 "cell_type": "markdown",747 "metadata": {748 "id": "sB9qPX1qoqJU"749 },750 "source": [751 "## Adding an Adapter Module\n",752 "\n",753 "Since `MLP` inherits `AdapterModuleMixin`, it also inherits a set of methods that perform adapter module manipulations - such as adding a new adapter.\n",754 "\n",755 "When users want to add an adapter, they can call `add_adapter()` with two specific arguments - `name` and `cfg`.\n",756 "\n",757 "Arguments - \n",758 "- `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",759 "- `cfg`: A dataclass / OmegaConf config that contains the `_target_` attribute pointing to the classpath of an Adapter Module, along with any additional required attributes."760 ]761 },762 {763 "cell_type": "code",764 "execution_count": null,765 "metadata": {766 "id": "Du2uPUSdni9D"767 },768 "outputs": [],769 "source": [770 "mlp.add_adapter(name='adapter_1', cfg=adapter_modules.LinearAdapterConfig(in_features=dim, dim=5))"771 ]772 },773 {774 "cell_type": "code",775 "execution_count": null,776 "metadata": {777 "id": "92vSc2k_xuHk"778 },779 "outputs": [],780 "source": [781 "# Now check the new parameter count of this MLP module, it should be higher than the previous count\n",782 "print(\"New param count :\", mlp.num_weights)"783 ]784 },785 {786 "cell_type": "markdown",787 "metadata": {788 "id": "eh4qTdTUyWTF"789 },790 "source": [791 "-----\n",792 "\n",793 "**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",794 "\n",795 "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."796 ]797 },798 {799 "cell_type": "markdown",800 "metadata": {801 "id": "w0CUr4m2sVRH"802 },803 "source": [804 "-----\n",805 "**Note**: If you try to add the same adapter multiple times, you will see the below error message! \n",806 "\n",807 "Remember, adapter names must be **locally** unique at the module level and **globally** unique at the model level!"808 ]809 },810 {811 "cell_type": "code",812 "execution_count": null,813 "metadata": {814 "id": "p497PTCBsoIM"815 },816 "outputs": [],817 "source": [818 "# Uncomment to see the error message - \n",819 "# mlp.add_adapter(name='adapter_1', cfg=adapter_modules.LinearAdapterConfig(in_features=dim, dim=10))"820 ]821 },822 {823 "cell_type": "markdown",824 "metadata": {825 "id": "EWgsqsCns4Yr"826 },827 "source": [828 "## Get all enabled Adapter Modules\n",829 "\n",830 "Next, we use `get_enabled_adapters()` to return a list of names of all the enabled adapters currently available to this module."831 ]832 },833 {834 "cell_type": "code",835 "execution_count": null,836 "metadata": {837 "id": "jXuGQoFus3qd"838 },839 "outputs": [],840 "source": [841 "mlp.get_enabled_adapters()"842 ]843 },844 {845 "cell_type": "markdown",846 "metadata": {847 "id": "H-UeYsFXtQqj"848 },849 "source": [850 "## Set the state of Adapter Modules\n",851 "\n",852 "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",853 "\n",854 "For that, we use the `set_enabled_adapter()` method. It has a few arguments - \n",855 "- `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",856 "- `enabled`: A bool, whether the adapter should be enabled or not.\n",857 "\n",858 "-----\n",859 "\n",860 "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",861 "\n",862 "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."863 ]864 },865 {866 "cell_type": "code",867 "execution_count": null,868 "metadata": {869 "id": "0__7RFaWtPD6"870 },871 "outputs": [],872 "source": [873 "# Disable all adapters\n",874 "mlp.set_enabled_adapters(enabled=False)\n",875 "print(\"Enabled adapters :\", mlp.get_enabled_adapters())\n",876 "\n",877 "# Enable just one adapter\n",878 "mlp.set_enabled_adapters(name=\"adapter_1\", enabled=True)\n",879 "print(\"Enabled adapters :\", mlp.get_enabled_adapters())"880 ]881 },882 {883 "cell_type": "markdown",884 "metadata": {885 "id": "KzT3kwvRubFT"886 },887 "source": [888 "## Check if Adapter Module(s) are available / enabled\n",889 "\n",890 "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"891 ]892 },893 {894 "cell_type": "code",895 "execution_count": null,896 "metadata": {897 "id": "5quE42kAusiX"898 },899 "outputs": [],900 "source": [901 "mlp.is_adapter_available()"902 ]903 },904 {905 "cell_type": "markdown",906 "metadata": {907 "id": "xRCobs9hvEeU"908 },909 "source": [910 "## Adapter functionality methods\n",911 "\n",912 "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",913 "\n",914 "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)."915 ]916 },917 {918 "cell_type": "markdown",919 "metadata": {920 "id": "6ybsJ9zRx6W7"921 },922 "source": [923 "### `forward_enabled_adapters()`\n",924 "To use these adapters, we utilize the `forward_adapter_modules()` method.\n",925 "\n",926 "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"927 ]928 },929 {930 "cell_type": "code",931 "execution_count": null,932 "metadata": {933 "id": "1F9IibFKv9va"934 },935 "outputs": [],936 "source": [937 "# Check `forward_enabled_adapters()`\n",938 "out = mlp.forward_enabled_adapters(input_data)\n",939 "print(out.shape)"940 ]941 },942 {943 "cell_type": "markdown",944 "metadata": {945 "id": "zzN0RlZI5M-l"946 },947 "source": [948 "### `forward_single_enabled_adapter_()`\n",949 "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",950 "\n",951 "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."952 ]953 },954 {955 "cell_type": "code",956 "execution_count": null,957 "metadata": {958 "id": "YO2dXe7T5PIQ"959 },960 "outputs": [],961 "source": [962 "# Check `forward_single_enabled_adapter_()`\n",963 "adapter_name = mlp.get_enabled_adapters()[0] # we have enabled just one adapter\n",964 "adapter_module = mlp.adapter_layer[adapter_name] # get the adapter module with this name\n",965 "adapter_strategy = adapter_module.adapter_strategy # get the adapter strategy for this adapter\n",966 "\n",967 "out = mlp.forward_single_enabled_adapter_(input_data, adapter_module, adapter_name=adapter_name, adapter_strategy=adapter_strategy)\n",968 "print(out.shape)"969 ]970 },971 {972 "cell_type": "markdown",973 "metadata": {974 "id": "bU-OS3wk7FW9"975 },976 "source": [977 "-----\n",978 "For further information about adapter forward pass, adapter strategy please refer to the documentation section for adapters."979 ]980 },981 {982 "cell_type": "markdown",983 "metadata": {984 "id": "b1mPF9qVx-HA"985 },986 "source": [987 "### `unfreeze_enabled_adapters()`\n",988 "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",989 "\n",990 "We can do this in two steps - \n",991 "- Call model.freeze() (at the highest level)\n",992 "- Call `unfreeze_enabled_adapters()` that will recursively unfreeze just the adapter modules that are enabled."993 ]994 },995 {996 "cell_type": "code",997 "execution_count": null,998 "metadata": {999 "id": "t6Qdk8YHw5W_"1000 },1001 "outputs": [],1002 "source": [1003 "# First setup some utility functions (this is part of NeuralModule)\n",1004 "def freeze(m):\n",1005 " for param in m.parameters():\n",1006 " param.requires_grad = False\n",1007 " m.eval()\n",1008 "\n",1009 "def unfreeze(m):\n",1010 " for param in m.parameters():\n",1011 " param.requires_grad = True\n",1012 " m.train()"1013 ]1014 },1015 {1016 "cell_type": "code",1017 "execution_count": null,1018 "metadata": {1019 "id": "95KGRNwxxlia"1020 },1021 "outputs": [],1022 "source": [1023 "freeze(mlp)\n",1024 "print(\"MLP frozen params :\", mlp.num_weights)"1025 ]1026 },1027 {1028 "cell_type": "code",1029 "execution_count": null,1030 "metadata": {1031 "id": "Y3xjntCgxhee"1032 },1033 "outputs": [],1034 "source": [1035 "# Check `unfreeze_enabled_adapters()` - param count should be lower than the previous total (original + adapter)\n",1036 "mlp.unfreeze_enabled_adapters()\n",1037 "print(\"MLP unfrozen adapter params :\", mlp.num_weights)"1038 ]1039 },1040 {1041 "cell_type": "markdown",1042 "metadata": {1043 "id": "vQjz_RIilg1L"1044 },1045 "source": [1046 "# Adapter support in intermediate level modules\n",1047 "\n",1048 "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",1049 "\n",1050 "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."1051 ]1052 },1053 {1054 "cell_type": "markdown",1055 "metadata": {1056 "id": "DsHudP-bz3bj"1057 },1058 "source": [1059 "## Intermediate modules that are instantiated via config\n",1060 "\n",1061 "Currently, we have a 3 level model -- \n",1062 "\n",1063 "`Top level Model (SimpleModel) -> Intermediate level Module (ResidualMLP) -> Bottom level Module (MLP)`. \n",1064 "\n",1065 "-----\n",1066 "\n",1067 "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",1068 "\n",1069 "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",1070 "\n",1071 "For this guide, we will show the recommended setup so that best practices can be followed."1072 ]1073 },1074 {1075 "cell_type": "markdown",1076 "metadata": {1077 "id": "_f-Y-Oxm1tjH"1078 },1079 "source": [1080 "## Creating an Adapter-compatible \"Penultimate\" module\n",1081 "\n",1082 "First, we create the new Adapter compatible module as a separate class."1083 ]1084 },1085 {1086 "cell_type": "code",1087 "execution_count": null,1088 "metadata": {1089 "id": "DC-L1SsVllLy"1090 },1091 "outputs": [],1092 "source": [1093 "# NOTE: We subclass the original ResidualMLP, and add in the AdapterModuleMixin too\n",1094 "class ResidualMLPAdapter(ResidualMLP, adapter_mixins.AdapterModuleMixin):\n",1095 " pass"1096 ]1097 },1098 {1099 "cell_type": "markdown",1100 "metadata": {1101 "id": "CuIlnrSW2m9t"1102 },1103 "source": [1104 "## Overriding the adapter methods\n",1105 "\n",1106 "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",1107 "\n",1108 "Therefore, this will create/update the state / forward an adapter module inside the `MLP` modules!"1109 ]1110 },1111 {1112 "cell_type": "code",1113 "execution_count": null,1114 "metadata": {1115 "id": "N3bcnXIy2mY9"1116 },1117 "outputs": [],1118 "source": [1119 "from typing import List, Optional\n",1120 "\n",1121 "class ResidualMLPAdapter(ResidualMLP, adapter_mixins.AdapterModuleMixin):\n",1122 " def add_adapter(self, name: str, cfg: DictConfig):\n",1123 " # call the same method on each `MLP` layer, collecting results\n",1124 " for layer in self.layers:\n",1125 " layer.add_adapter(name, cfg)\n",1126 " \n",1127 " def get_enabled_adapters(self) -> List[str]:\n",1128 " # call the same method on each `MLP` layer, collecting results\n",1129 " enabled_adapters = set([])\n",1130 " for layer in self.layers:\n",1131 " names = layer.get_enabled_adapters()\n",1132 " enabled_adapters.update(names)\n",1133 " return list(enabled_adapters)\n",1134 " \n",1135 " def set_enabled_adapters(self, name: Optional[str], enabled: bool):\n",1136 " # call the same method on each `MLP` layer, collecting results\n",1137 " for layer in self.layers:\n",1138 " layer.set_enabled_adapters(name, enabled)\n",1139 " \n",1140 " def is_adapter_available(self) -> bool:\n",1141 " # call the same method on each `MLP` layer, collecting results\n",1142 " is_available = any([layer.is_adapter_available() for layer in self.layers])\n",1143 " return is_available"1144 ]1145 },1146 {1147 "cell_type": "markdown",1148 "metadata": {1149 "id": "_-RHMG4W4pSY"1150 },1151 "source": [1152 "## Register the new adapter\n",1153 "\n",1154 "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",1155 "\n",1156 "The steps below are : \n",1157 "- Check if the registry has the base class via `get_registered_adapter()`.\n",1158 "- If it returns None, then register the base class and its compatible adapter class via `register_adapter()`.\n",1159 "\n",1160 "-----\n",1161 "\n",1162 "**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."1163 ]1164 },1165 {1166 "cell_type": "code",1167 "execution_count": null,1168 "metadata": {1169 "id": "sOUo-b042S9m"1170 },1171 "outputs": [],1172 "source": [1173 "if adapter_mixins.get_registered_adapter(ResidualMLP) is None:\n",1174 " adapter_mixins.register_adapter(ResidualMLP, ResidualMLPAdapter)"1175 ]1176 },1177 {1178 "cell_type": "markdown",1179 "metadata": {1180 "id": "BSaAnSxC6A6f"1181 },1182 "source": [1183 "-----\n",1184 "\n",1185 "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",1186 "\n",1187 "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."1188 ]1189 },1190 {1191 "cell_type": "markdown",1192 "metadata": {1193 "id": "_yeL62Js6sGb"1194 },1195 "source": [1196 "# Adapters support at top level Model\n",1197 "\n",1198 "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",1199 "\n",1200 "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",