Add support for MHA adapters to ASR
State: closed | Pull request: yes
Labels: name: core

name: TTS

name: ASR

name: NLP

name: CI

name: common

# What does this PR do ?

Adds Multi Head Attention Adapter support, along with support for Tiny-Attention adapter from https://arxiv.org/abs/2211.01979 

**Collection**: [Core, ASR]

# Changelog 
- Add modules for tiny attention based adapter
  - Add `MHAResidualAddAdapterStrategy` and config
  - Add `RelPositionMultiHeadAttentionAdapter` and config
  - Add `PositionalEncodingAdapter` and config (unused for now)
  - Add `RelPositionalEncodingAdapter` and config (unused for now)
- Add support to Conformer and Squeezeformer for attention adapters
- Add new support for ASR adapter training script to support multiple types of adapters (linear, mha, relpos mha, tiny attention etc)
- Add exp manager utility function to clean up ckpt and nemo files during grid search.
- Add `ReturnResultAdapterStrategy` to simply return the result of the adapter without residual connection - used for PositionalEncodingAdapter and `RelPositionalEncodingAdapter` 
- Extend `ResidualAddAdapterStrategy` to  `MHAResidualAddAdapterStrategy` for attention adapters.
- Update docstrings

# Usage

## Single adapter training

``` bash

python train_asr_adapter.py \
    --config-path="../conf/asr_adapters" \
    --config-name="asr_adaptation.yaml" \
    model.pretrained_model=null \
    model.nemo_model=null \
    model.adapter.adapter_name=<Unique adapter name> \
    model.adapter.adapter_type="tiny_attn" \
    model.adapter.adapter_module_name="encoder" \
    model.adapter.tiny_attn.n_feat=<dimension of the layer outputs of the model> \
    model.adapter.tiny_attn.n_head=1 \
    model.adapter.tiny_attn.proj_dim=-1 \
    model.adapter.linear.dropout_rate=0.0 \
    model.train_ds.manifest_filepath=<Path to manifest> \
    model.train_ds.batch_size=8 \
    model.validation_ds.manifest_filepath=<Path to manifest> \
    model.validation_ds.batch_size=8 \
    model.optim.lr=0.001 \
    model.optim.weight_decay=0.0 \
    model.optim.sched.warmup_steps=100 \
    trainer.max_steps=300 \
    trainer.devices=1 \
    trainer.precision=32 \
    exp_manager.exp_dir=<Some directory for experiment manager>

```

## Hyper parameter search

First, modify hydra.sweep section in the `asr_adaptation_hp.yaml` file - 

``` yaml

hydra:
  sweep:
    dir: "."
    subdir: "."

  sweeper:
    params:  # place all the parameters you wish to search over here (corresponding to the rest of the config)
      model.optim.lr: 0.001  # ,0.0001
      model.adapter.tiny_attn.n_head: 1  # 1,2,4,8
      model.adapter.tiny_attn.proj_dim: 1,4,8,16,32,64  # 1,4,8,16,32,64
      model.adapter.tiny_attn.adapter_strategy.stochastic_depth: 0.0,0.5,0.6,0.7,0.8,0.9

  # Arguments to the hyperparameter runner
  launcher:
    num_gpus: -1  # Number of gpus to use. Each run works on a single GPU.
    jobs_per_gpu: 1  # If each GPU has large memory, you can run multiple jobs on the same GPU for faster results (until OOM).

```

Then run the adapter train script as below - 

``` bash

python train_asr_adapter.py \
    --config-path="../conf/asr_adapters" \
    --config-name="asr_adaptation_hp.yaml" \
    -m \
    model.pretrained_model=null \
    model.nemo_model=null \
    model.adapter.adapter_name=<Unique adapter name> \
    model.adapter.adapter_type="tiny_attn" \
    model.adapter.adapter_module_name="encoder" \
    model.adapter.tiny_attn.n_feat=<dimension of the layer outputs of the model> \
    model.train_ds.manifest_filepath=<Path to manifest> \
    model.train_ds.batch_size=8 \
    model.validation_ds.manifest_filepath=<Path to manifest> \
    model.validation_ds.batch_size=8 \
    exp_manager.exp_dir="<some directory>" \
    exp_manager.create_wandb_logger=true \
    exp_manager.wandb_logger_kwargs.project="<Project Name>" \
    ++delete_ckpt_after_train=True

```

# 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?
- [x] 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