Enable encoder adapters for Canary and MultiTaskAED models
State: closed | Pull request: yes
Labels: name: core

name: ASR

name: NLP

name: Run CICD

# What does this PR do ?

Enable multi task AED models to add adapter support to its encoders (speech encoder and transformer encoder) and decoder (transformer decoder). Also add new API to automatically swap out original modules with adapter compatible modules.

**Collection**: [ASR]

# Changelog 
- Add support for adapter module mixin to AED multitask model
- [CORE] Add support for `replace_adapter_compatible_modules()` which automatically swaps out registered original module class with the adapter class. It also updates config so save_to() and restore_from() will automatically pick current class.
- Support mha adapter and linear adapter for both speech encoder and transformer encoder + decoder modules
- Refactor AdapterModuleMixin to separate adapter type assertion from `add_adapter`
- Refactor ASRAdapterModelMixin to separate the logic of what model is parsed  - now its determined by a simple string list of module names
- Add tests for aed multi task model with adapter support


# Usage

```python
from nemo.collections.asr.models import ASRModel
from nemo.collections.asr.parts.submodules.adapters import TransformerEncoderMultiHeadAttentionAdapterConfig
from nemo.collections.common.parts.adapter_modules import LinearAdapterConfig
 
model = ASRModel.from_pretrained("nvidia/canary-1b")

# New API to automatically swap out all original modules in model with adapter compatible modules via registry
model.replace_adapter_compatible_modules()

# Tiny Attention Adapters
adapter_cfg = TransformerEncoderMultiHeadAttentionAdapterConfig(
  hidden_size=1024,
  num_attention_heads=1,
  proj_dim=8,
)

# OR

# Linear Adapters
adapter_cfg = LinearAdapterConfig(
  in_features=1024,
  dim=8,
)

# Load adapter weights and freeze all other weights
# Adapter targets can be "", "encoder", "transf_encoder" and "transf_decoder". "" refers to encoder by default.
# Verify supported adapter target modules using `model.adapter_module_names`.
model.add_adapter("encoder:my_adapter", adapter_cfg)
model.freeze()
model.unfreeze_enabled_adapters()

# Train model ...
```


# GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

# Before your PR is "Ready for review"
**Pre checks**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [x] Did you write any new necessary tests?
- [x] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [x] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.