Fix alignment logits when using fusion models in TDT greedy decoder

[Garvys]

What does this PR do ?

Fixes an inconsistency in alignment logits when using preserve_alignments=True with TDT models and fusion models (e.g., LM, boosting tree).

Problem

When fusion models are enabled, the label selection logic correctly applies:

• If blank is the most likely token before fusion → emit blank
• Otherwise → emit the non-blank token with the highest logit after fusion

However, the logits stored in alignments were always the pre-fusion logits. This caused a mismatch where the stored label (correctly chosen with the above logic) did not correspond to the argmax of the stored logits.

Solution

Update the logits sent to alignments to be consistent with the label selection:

• If blank was selected (best before fusion) → store pre-fusion logits (unchanged)
• If a non-blank token was selected → store post-fusion logits with blank set to -inf

This ensures argmax(logits) == label for every alignment entry, making the alignments reliable for downstream analysis.

Scope

Implementation in torch_impl method only. CUDA graphs implementation is not affected.

Collection: ASR

Changelog

• tdt_label_looping.py: Updated torch_impl to store alignment-consistent logits when fusion models are used (both outer loop and inner loop)
• test_rnnt_decoding.py: Added test_tdt_greedy_decoding_preserve_alignments_with_fusion to verify label/logits consistency with fusion models

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

None

[Garvys]

If you have time for a review @nithinraok. Thanks a lot!!

[artbataev]

First, thanks for this contribution.
Unfortunately, I'm working on a concurrent PR with modifications to the decoding algorithm, so asking to wait before it is merged.
But I would be happy to fix bugs and inconsistencies in this implementation.

Also, it would be great if you could clarify your use-case of alignments: alignments data structure (which stores all logits) requires an enormous amount of memory and makes decoding slower.

Regarding the PR itself:

• I agree about using consistent logits with LM scores in alignments (currently inconsistent, but for both TDT and RNN-T)
• Also, we need to maintain consistency of tdt_label_looping and rnnt_label_looping
• I disagree with using -inf for blank, since it will break some of our use cases of alignments.

If you clarify the core use-case, we could discuss how to implement it.

[artbataev]

Concurrent PR with decoding changes: #15315

[Garvys]

Hello @artbataev !

Thanks a lot for taking the time to review my PR!

No worries, I can wait.

Regarding my use-case for the alignments data structure: I use it to have access to the top tokens at each decoding step in order to troubleshoot why the model is not behaving as expected. I attached a screenshot of the tool.

It is very interesting to see, at each prediction step, which tokens are the most likely according to the model and to compare that with how the Fusion model is trained. This really allow to have a deep understanding of the behaviour of the models and to see which changes could improve the results (fine tuning, changing how the fusion model is trained ...)

What i noticed is that the token that was predicted at each prediction step was not necessarily the one with the highest score in the logits of the relevant prediction step. My proposed change allows to fix that but indeed I agree that setting -inf for the blank logit is not very satisfying.

When the blank is not the most likely, keeping the blank logit to its value and updating the logits with the fusion scores of all the other tokens won't guarantee that the blank token is not the most likely token.
It does seem to me that the good way of implementing this would be to store the blank logits in another tensor when shallow fusion is applied.

WDYT ? How would you fix this issue with another implementation ?

Also, thanks a lot for linking the other PR. I am having a look !

Thanks a lot for all the PRs you merged related to TDT, very impressive work 💪