Add Transformer Encoder for ASR by nithinraok · Pull Request #15661 · NVIDIA-NeMo/NeMo

nithinraok · 2026-05-04T15:50:50Z

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 a new ASR transformer encoder with frame stacking, configurable subsampling, FlexAttention-based full attention, and optional QK Norm

Collection: ASR

Changelog

Add TransformerEncoder in nemo/collections/asr/modules/transformer_encoder.py as a lightweight pre-norm transformer encoder for ASR.
Add TransformerEncoderConfig dataclass to capture encoder hyperparameters such as d_model, n_heads, n_layers, qkv_bias, qk_norm, and
subsampling_factor.
Add FeatureStacking pre-encoder module to stack consecutive frames, reduce sequence length, and project stacked features into the model
dimension.
Add transformer building blocks: FeedForward, MultiHeadAttention, and TransformerBlock.
Implement full-attention masking with PyTorch FlexAttention and padding-aware block masks.
Add optional per-head QK normalization before attention score computation.
Restrict the initial PR scope to attn_mode="full" and raise an error for unsupported future modes.
Export TransformerEncoder from nemo/collections/asr/modules/__init__.py for config-based instantiation.
Add example config in examples/asr/conf/fastconformer/transformer_stacking_tdt_bpe.yaml showing how to train an RNNT/TDT model with the new
encoder.

Usage

  import torch

  from nemo.collections.asr.modules import TransformerEncoder

  encoder = TransformerEncoder(
      feat_in=128,
      d_model=512,
      n_heads=8,
      n_layers=12,
      drop_rate=0.1,
      qkv_bias=False,
      qk_norm=True,
      subsampling_factor=4,
      attn_mode="full",
  )

  audio_signal = torch.randn(2, 128, 400)   # (B, C, T)
  lengths = torch.tensor([400, 360])        # valid frame lengths

  encoded, encoded_lengths = encoder(audio_signal, lengths)

  print(encoded.shape)         # (B, D, T')
  print(encoded_lengths)       # ceil(length / subsampling_factor)

You can also instantiate it from Hydra config with:

  encoder:
    _target_: nemo.collections.asr.modules.transformer_encoder.TransformerEncoder
    feat_in: ${model.preprocessor.features}
    d_model: 1280
    n_heads: 16
    n_layers: 32
    drop_rate: 0.1
    qkv_bias: false
    qk_norm: true
    subsampling_factor: 8
    attn_mode: full

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?

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

Related to # (issue)

Signed-off-by: nithinraok <nithinrao.koluguri@gmail.com>

copy-pr-bot · 2026-05-04T15:50:55Z

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.

pzelasko · 2026-05-04T15:55:29Z

+        self.head_dim = cfg.d_model // cfg.n_heads
+        self.d_model = cfg.d_model
+
+        self.w_q = nn.Linear(cfg.d_model, cfg.d_model, bias=cfg.qkv_bias)


let's merge these 3 lines into w_qkv to batch 3 ops in one

This was done due to QK-Norm option to enable. But yes, I can unbind again for such ops.

pzelasko · 2026-05-04T15:56:25Z

+        self.qk_norm = cfg.qk_norm
+        if cfg.qk_norm:
+            self.q_norm = nn.LayerNorm(self.head_dim)
+            self.k_norm = nn.LayerNorm(self.head_dim)


what about RMSNorm?

I haven;t validated runs with RMSNorm yet, so will keep this for future work.

pzelasko · 2026-05-04T15:59:45Z

+            x = nn.functional.pad(x, (0, 0, 0, pad_size))
+        t_new = (t + pad_size) // self.subsampling_factor
+        x = x.reshape(b, t_new, c * self.subsampling_factor)
+        x = self.proj(x)


do we want to add some norm before projection to center the activations? rms norm could be reasonable although then we might need to change the zero-padding above to something that doesn't skew the norm too much

Thanks for suggestion. I will need to check by running experiments with this effect. So for now will defer to later PR.

pzelasko · 2026-05-04T16:03:04Z

+            length: (B,) — output lengths after subsampling.
+        """
+        x, length = self.pre_encode(audio_signal, length)
+        x = x * (self.d_model**0.5)


why re-scale just before norm?

Added to prevent NaNs during training, similar to the fix we applied in FastConformer during scaling.

I need to check again to remove this if this has no effect after I added LayerNorm post Encoder for additional stability. Will revist

Where would you get NaN from if you have embed_norm just below? I don't get it.

So couple of things:

I was getting NaNs in training loss so applied bunch of techniques to see which could fix like embedding scaling, layer norm after last layer and padding after this model converged well.

Then I didn't went back and change each to see exactly which one contributed heavily for smoother training. I think Post LayerNorm but I haven't validated.

I don't think this might have effect but I didn't check without it

Good point.

So I disabled it and re-ran for already trained ckpt and here are results:

with:

librispeech_test_clean_manifest: WER=0.0177 librispeech_test_other_manifest: WER=0.0359

commenting it out:

librispeech_test_clean_manifest: WER=0.0177 librispeech_test_other_manifest: WER=0.0359

You are correct, looks like pre-norm from attention is removing the added scale.

so I believe its the post Layer Norm that helped NaNs issue.

pzelasko

nice work!

tango4j · 2026-05-05T06:38:41Z

+    def __init__(self, cfg: TransformerEncoderConfig):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(cfg.d_model, 4 * cfg.d_model),


This should never be hardcoded since this is very frequently customized variable.
cfg.ff_expansion should be added, replacing the hardcoded 4.

Also, I think we should add this function to TransformerEncoderConfig

@property def ff_hidden_size(self) -> int: return int(self.ff_expansion * self.d_model)

and make this line
nn.Linear(cfg.d_model, cfg.ff_hidden_size)

I donlt think adding property to dataclass is a good idea. however added the ff_exansion parameter

Don't need to add property function. But wrapping the variable with int() is necessary.
And ff_expansion should be float.
Maybe one line of a variable sanity check on the "ffn_hidden_size" would make the code easy to track errors.

Yes added now, pls check

Checked the change and ff_hidden variable. Looks good.

Please approve if all looks good from your end

tango4j · 2026-05-05T07:06:32Z

+            length: (B,) — output lengths after subsampling.
+        """
+        x, length = self.pre_encode(audio_signal, length)
+        x = x * (self.d_model**0.5)


Since the layer norm is after this, this line has no effect.
If this is needed, needs to be after x = self.embed_norm(x)..?

Added to prevent NaNs during training, similar to the fix we applied in FastConformer during scaling.

I need to check again to remove this if this has no effect after I added LayerNorm post Encoder for additional stability. Will revist

Good point.

So I disabled it and re-ran for already trained ckpt and here are results:

with:

librispeech_test_clean_manifest: WER=0.0177 librispeech_test_other_manifest: WER=0.0359

commenting it out:

librispeech_test_clean_manifest: WER=0.0177 librispeech_test_other_manifest: WER=0.0359

You are correct, looks like pre-norm from attention is removing the added scale.

Very nice. This was a concern for me for breaking the symmetry with other TF encoder implementations.

Signed-off-by: nithinraok <nithinrao.koluguri@gmail.com>

nithinraok · 2026-05-05T12:33:41Z

/ok to test ef9547a

github-actions · 2026-05-05T13:54:30Z

[🤖]: Hi @nithinraok 👋,

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.

Signed-off-by: nithinraok <nithinrao.koluguri@gmail.com>

pzelasko

Thanks @nithinraok ! It looks good from my side, let's wait for @tango4j to approve and LGTM

tango4j · 2026-05-06T03:53:32Z

This doesn't need to be happening in this PR but to ensure standardization, but later on, we probably need to verify this implementation is 100% compatible (transferable weights and identical results) with standard baselines like the Hugging Face Transformer.

If there are any custom scaling or structural tweaks, it would be great to make them optional flags. This keeps our baseline universally compatible while still allowing for optimizations when needed.

I think we should merge this after getting approval from @KunalDhawan and @stevehuang52 too.

nithinraok added 2 commits May 4, 2026 08:37

flex attention code asr

9dc987f

Signed-off-by: nithinraok <nithinrao.koluguri@gmail.com>

update lengths to match for shortter durations

0d257c6

Signed-off-by: nithinraok <nithinrao.koluguri@gmail.com>

github-actions Bot added the ASR label May 4, 2026

nithinraok requested a review from pzelasko May 4, 2026 15:51

pzelasko reviewed May 4, 2026

View reviewed changes

tango4j requested changes May 5, 2026

View reviewed changes