Adding RoPE to TransformerEncoder

[ipmedenn]

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 (rotary position embedding) as a self_attention_model option to the ASR TransformerEncoder, alongside the existing "rel_pos" (default), "abs_pos", and "no_pos" modes. With self_attention_model="rope", position is injected by rotating Q and K inside attention — no additive positional embedding is added to the signal and attention stays plain scaled dot-product (FlexAttention with score_mod=None). This mirrors the RoPE support recently added to ConformerEncoder (#15714), keeping the two encoder surfaces aligned.

Collection: ASR

Changelog

• Added "rope" to the supported self_attention_model values of TransformerEncoder / TransformerEncoderConfig, plus rope_base (default 10000.0) and rotary_fraction (default 1.0) arguments and config fields.
• MultiHeadAttention now accepts an optional shared pos_enc; in rope mode it rotates Q/K in place via RotaryPositionalEncoding before FlexAttention (score_mod=None), and raises if rope is selected without a rotary module.
• TransformerEncoder builds a single RotaryPositionalEncoding (d_k = d_model // n_heads, rope_base, rotary_fraction, max_len = pos_emb_max_len) and passes it into every TransformerBlock so the cos/sin buffers are shared.
• In rope mode no additive positional embedding is added; xscaling (if set) and dropout_pre_encoder are applied to the pre-encoder output in forward_internal, and dropout_emb is unused.
• Reuses the existing RotaryPositionalEncoding from nemo.collections.asr.parts.submodules.multi_head_attention (no new positional-encoding module introduced).
• Added unit tests: extended the parametrized self_attention_model tests (test_valid_modes_are_accepted, test_non_rel_pos_modes_have_no_rel_params, test_forward_each_mode_cpu) to include "rope", asserted rope_base/rotary_fraction config defaults, and added test_rope_uses_shared_rotary_pos_enc (shared-module wiring) and test_rope_partial_rotation_forward_cpu (rotary_fraction=0.5).

Usage

Config:

model:
  encoder:
    _target_: nemo.collections.asr.modules.transformer_encoder.TransformerEncoder
    self_attention_model: rope
    rope_base: 10000.0
    rotary_fraction: 1.0

Python:

from nemo.collections.asr.modules.transformer_encoder import TransformerEncoder
encoder = TransformerEncoder(
    feat_in=128,
    d_model=512,
    n_heads=8,
    n_layers=17,
    subsampling="feature_stacking",
    subsampling_factor=4,
    self_attention_model="rope",  # one of "rel_pos" | "abs_pos" | "rope" | "no_pos" (or None)
    rope_base=10000.0,
    rotary_fraction=1.0,
)
encoded, encoded_len = encoder(audio_signal=features, length=feature_lengths)

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

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

Follow-up work after the initial TF encoder PRs (#15661, #15703). Adds the rope positional-encoding option mirroring the RoPE support added to ConformerEncoder (#15714).

[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.

[tango4j]

/ok to test f5962f4

[nithinraok]

pls add tests

[nithinraok]

/ok to test 25456fe

[nithinraok]

LGTM

[nithinraok]

/claude review

[claude]

LGTM

[github-actions]

[🤖]: Hi @ipmedenn 👋,

We wanted to let you know that a CICD pipeline for this PR just finished successfully.

So it might be time to merge this PR or get some approvals.

[nithinraok]

/ok to test 8c01d64

[github-actions]

[🤖]: Hi @ipmedenn 👋,

We wanted to let you know that a CICD pipeline for this PR just finished successfully.

So it might be time to merge this PR or get some approvals.