RNN-T and TDT inference: use CUDA graphs by default
State: closed | Pull request: yes
Labels: name: core

name: ASR

name: common

name: Run CICD

# What does this PR do ?

- enable CUDA graphs by default for inference (`transcribe_speech.py`, `speech_to_text_eval.py`)
- add fallback behavior with partial CUDA graphs for Label-Looping algorithm implementation (TDT, RNN-T)
  - Use Cuda graphs _with while loops_ when all requirements are met (driver 545+, Cuda 12.3+, cuda-python package)
  - Use Cuda graphs _without while loops_ (pure PyTorch functionality) in most other cases. The idea: use several separate Cuda graphs for parts of the decoding algorithm (`<graph before outer loop>` -> `while loop (python)` -> `<graph before inner loop>` -> `inner while loop (python)` -> `<graph for inner loop code>` etc.)

On my local machine, FastConformer L, LibriSpeech test-other decoding time, bfloat16, bs=16


|                | no graphs | partial graphs | full graph |
| -------------- | --------- | -------------- | ---------- |
| unsorted batch | 34s       | 21s            | 19s        |
| sorted batch   | 23s       | 15s            | 14s        |


**Collection**: [ASR]

# Changelog 
- Add specific line by line info of high level changes in this PR.

# Usage

```shell
python examples/asr/speech_to_text_eval.py <with rnnt/tdt model>
```

# 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**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [x] 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**:
- [x] 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](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Related to # (issue)


jenkins

jenkins

@galv, @titu1994 Please, review the PR

The only issue I see for now is that the Frame-Looping algorithm does not have fallback behavior, and setting  `loop_labels=false` without `use_cuda_graph_decoder=false` (default - `true`) can cause failures (if not all requirements are met for Frame-Looping+CUDA graphs).

On the other hand, this case is not important. Most users will use the Label-Looping algorithm since it is default and produces the same result as the Frame-Looping algorithm.

There is something I'm not quite understanding. Is there a reason why you did not set the default value of these values to true? 

https://github.com/NVIDIA/NeMo/blob/7e805bc8d0683cc76fa44b6a847611aaaedb9794/nemo/collections/asr/parts/submodules/rnnt_greedy_decoding.py#L592

and

https://github.com/NVIDIA/NeMo/blob/7e805bc8d0683cc76fa44b6a847611aaaedb9794/nemo/collections/asr/parts/submodules/rnnt_greedy_decoding.py#L2305

and

https://github.com/NVIDIA/NeMo/blob/7e805bc8d0683cc76fa44b6a847611aaaedb9794/nemo/collections/asr/parts/submodules/rnnt_greedy_decoding.py#L2655

It looks like you are turning cuda graphs on by default only when someone runs transcribe_speech.py or transcribe_speech_parallel.py right now.

> On the other hand, this case is not important. Most users will use the Label-Looping algorithm since it is default and produces the same result as the Frame-Looping algorithm.

I agree. That is was my first attempt. It was very educational, but I don't believe that we need to do the work to add a feature to it that won't be used very much.

I didn't hit approve yet. The changes seem good to me. But I want to give some time to defer to @titu1994 

> There is something I'm not quite understanding. Is there a reason why you did not set the default value of these values to true?
> 
> https://github.com/NVIDIA/NeMo/blob/7e805bc8d0683cc76fa44b6a847611aaaedb9794/nemo/collections/asr/parts/submodules/rnnt_greedy_decoding.py#L592
> 
> and
> 
> https://github.com/NVIDIA/NeMo/blob/7e805bc8d0683cc76fa44b6a847611aaaedb9794/nemo/collections/asr/parts/submodules/rnnt_greedy_decoding.py#L2305
> 
> and
> 
> https://github.com/NVIDIA/NeMo/blob/7e805bc8d0683cc76fa44b6a847611aaaedb9794/nemo/collections/asr/parts/submodules/rnnt_greedy_decoding.py#L2655
> 
> It looks like you are turning cuda graphs on by default only when someone runs transcribe_speech.py or transcribe_speech_parallel.py right now.

This is done intentionally. The problem is that using `True` as a default value here will result in using CUDA graphs when decoding in training mode, which is can result in issues with bucketing (dynamic batch sizes).
I think it is cleaner for now to use CUDA graphs only for transcribing speech.

See `training_step` for RNN-T models 
https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/asr/models/rnnt_models.py#L663
https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/asr/modules/rnnt.py#L1490
https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/asr/metrics/wer.py#L327



@galv, @titu1994 Please, review once again

I reworked the PR.
- CUDA graphs enabled by default now in all classes but turned off in the training pipeline
- I also added more tests and fixed some issues with confidence
- Now max_symbols is set to 10 if it is initially None (with an appropriate warning)