Add RoPE to ConformerEncoder

[MahmoudAshraf97]

Important

The Update branch button must only be pressed in very rare occassions.
An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.

What does this PR do ?

Adds RoPE as a positional encoding method to ConformerEncoder, it enables the use of torch SDPA while maintaining accuracy
This is a least-friction implementation, there is still path for more improvements (such as flex attention and fused qkv proj) but I preferred to not diverge from other conformer submodules
Collection: asr

Motivation

Since #9590, SDPA caused loss explosion in training when Relative Positional Embeddings were used, investigation showed that SDPA had issues with gradients when attention additive biases were used

as per this paper, RoPe has the same baseline or better performance than RelPos and trains faster because of its ability to use SDPA

Results

I tested two Fastconformer medium models, everything was constant except the positional encoding method (RoPE vs RelPos) and attention implementation (SDPA vs NeMo), both models were randomly initialized, although I forgot to fix the seed.

One takeaway is that when the two models were initialized from a pretrained checkpoint except for the positional encoding parameters, RelPos converged faster, implying that RoPE attention geometry is different than RelPos and not necessarily transferable

Improvements

The Rope was around 6% faster E2E, and used less memory which allowed for higher fused batch size for the RNNT loss and better compute utilization. The training time was dominated by RNNT loss calculation, so it diminishes any benefits from SDPA or torch.compile

RNNT loss compute time is data-dependent, that's why the step time is much lower at the beginning of the run when the model hasn't learned anything yet, once it starts to generate tokens, the time goes up

Changelog

• Add RoPE as Positional Encoding Method for ConformerEncoder

Usage

• pass self_attention_model: rope to encoder config

model:
  encoder:
    self_attention_model: rope

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:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• 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:

• New Feature
• Bugfix
• Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

cc @pzelasko

Additional Information

• Related to # (issue)

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.