"Loop labels" greedy decoding: faster implementation #8286

artbataev · 2024-01-31T18:26:11Z

What does this PR do ?

Improves "loop labels" greedy decoding implementation, introduced in #7926, with using non-blocking operations (avoiding device-to-host synchronization).

On my local machine, the time for decoding unsorted test-other (LibriSpeech) with stt_en_fastconformer_transducer_large (2 repeated measurements, in seconds; after warmup the results are better):

bs	precision	RNN-T, loop_labels v1 (main)	RNN-T, loop_labels v2 (this)
16	fp32	50s / 50s	37s / 36s
128	fp32	34s / 33s	30s / 29s
16	bf16	48s / 46s	34s / 32s
128	bf16	25s / 24s	21s / 20s

This algorithm is planned for further improvements with torch.jit.script in #8203

Collection: [ASR]

Changelog

Add specific line by line info of high level changes in this PR.

Usage

python examples/asr/speech_to_text_eval.py \
   model_path=<nemo_model.nemo> \
   dataset_manifest=<manifest> \
   batch_size=128 \
   output_filename=<output_mainfest_path> \
   rnnt_decoding.greedy.loop_labels=true

Jenkins CI

To run Jenkins, a NeMo User with write access must comment jenkins on the PR.

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: Vladimir Bataev <vbataev@nvidia.com>

… rnnt_greedy_loop_lables_v2

artbataev · 2024-02-01T08:01:27Z

jenkins

GNroy

LGTM, see comments.

nemo/collections/asr/modules/rnnt.py

nemo/collections/asr/modules/rnnt_abstract.py

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

artbataev · 2024-02-01T17:44:36Z

jenkins

titu1994

Really incredible work ! Ready to merge

titu1994 · 2024-02-02T21:46:57Z

jenkins

* Loop labels greedy decoding v2 Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments. Clean up Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add tests for batched hypotheses Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add tests for batched alignments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Fix comment Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Fix test Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Method -> classmethod (self is not needed) Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> --------- Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> Co-authored-by: Somshubra Majumdar <titu1994@gmail.com>

* Loop labels greedy decoding v2 Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments. Clean up Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add tests for batched hypotheses Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add tests for batched alignments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Fix comment Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Fix test Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Method -> classmethod (self is not needed) Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> --------- Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> Co-authored-by: Somshubra Majumdar <titu1994@gmail.com> Signed-off-by: biscayan <skgudwn34@gmail.com>

* Loop labels greedy decoding v2 Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments. Clean up Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add tests for batched hypotheses Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add tests for batched alignments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Fix comment Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Fix test Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Method -> classmethod (self is not needed) Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> --------- Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> Co-authored-by: Somshubra Majumdar <titu1994@gmail.com> Signed-off-by: Sasha Meister <ameister@nvidia.com>

* Loop labels greedy decoding v2 Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments. Clean up Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add tests for batched hypotheses Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add tests for batched alignments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Fix comment Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Fix test Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Method -> classmethod (self is not needed) Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> --------- Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> Co-authored-by: Somshubra Majumdar <titu1994@gmail.com> Signed-off-by: Pablo Garay <pagaray@nvidia.com>

* Loop labels greedy decoding v2 Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments. Clean up Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add tests for batched hypotheses Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add tests for batched alignments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Add comments Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Fix comment Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Fix test Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> * Method -> classmethod (self is not needed) Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> --------- Signed-off-by: Vladimir Bataev <vbataev@nvidia.com> Co-authored-by: Somshubra Majumdar <titu1994@gmail.com>

Loop labels greedy decoding v2

abb998c

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

github-actions bot added the ASR label Jan 31, 2024

artbataev added 3 commits January 31, 2024 22:34

Add comments. Clean up

c4a5672

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Add comments

1ce53d3

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Add comments

eb1b425

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

artbataev changed the title ~~"Loop labels" greedy decoding: implementation without blocking operations~~ "Loop labels" greedy decoding: faster implementation Jan 31, 2024

artbataev added 7 commits January 31, 2024 23:44

Add tests for batched hypotheses

5982b4d

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Add tests for batched alignments

c2a8b1f

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Add comments

e6b6bb7

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Fix comment

fb6000f

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Merge branch 'main' into rnnt_greedy_loop_lables_v2

3445357

Fix test

00ffb75

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Merge remote-tracking branch 'origin/rnnt_greedy_loop_lables_v2' into…

712a78c

… rnnt_greedy_loop_lables_v2

artbataev marked this pull request as ready for review February 1, 2024 08:53

artbataev requested review from GNroy and titu1994 February 1, 2024 08:53

GNroy previously approved these changes Feb 1, 2024

View reviewed changes

nemo/collections/asr/modules/rnnt.py Show resolved Hide resolved

nemo/collections/asr/modules/rnnt.py Show resolved Hide resolved

nemo/collections/asr/modules/rnnt_abstract.py Show resolved Hide resolved

artbataev mentioned this pull request Feb 1, 2024

Add loop_labels algorithm for TDT greedy decoding #8215

Merged

8 tasks

Method -> classmethod (self is not needed)

088fc7d

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

artbataev dismissed GNroy’s stale review via 088fc7d February 1, 2024 17:44

titu1994 approved these changes Feb 2, 2024

View reviewed changes

Merge branch 'main' into rnnt_greedy_loop_lables_v2

3010cad

artbataev merged commit a5448f3 into main Feb 3, 2024
15 checks passed

artbataev deleted the rnnt_greedy_loop_lables_v2 branch February 3, 2024 09:08

artbataev mentioned this pull request Apr 5, 2024

Use Label-Looping algorithm for RNN-T decoding by default #8831

Merged

8 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

"Loop labels" greedy decoding: faster implementation #8286

"Loop labels" greedy decoding: faster implementation #8286

artbataev commented Jan 31, 2024 •

edited

Loading

artbataev commented Feb 1, 2024

GNroy left a comment

artbataev commented Feb 1, 2024

titu1994 left a comment

titu1994 commented Feb 2, 2024

"Loop labels" greedy decoding: faster implementation #8286

"Loop labels" greedy decoding: faster implementation #8286

Conversation

artbataev commented Jan 31, 2024 • edited Loading

What does this PR do ?

Changelog

Usage

Jenkins CI

Before your PR is "Ready for review"

Who can review?

Additional Information

artbataev commented Feb 1, 2024

GNroy left a comment

Choose a reason for hiding this comment

artbataev commented Feb 1, 2024

titu1994 left a comment

Choose a reason for hiding this comment

titu1994 commented Feb 2, 2024

artbataev commented Jan 31, 2024 •

edited

Loading