Fix #8891 by supported GPU-side batched CTC Greedy Decoding

[galv]

What does this PR do ?

I add a new flag "batched_inference" to the CTC greedy decoder.

Collection: ASR

Changelog

• Support batched inference of greedy CTC decoding.
• Add support for batched inference for label inputs as well.

Usage

Here is an example turning this on for transcribe_speech.py

batch_size=32
amp=true

python examples/asr/speech_to_text_eval.py  pretrained_name=nvidia/parakeet-ctc-1.1b \ 
dataset_manifest=/home/dgalvez/scratch/data/test_other_sorted_downward.json  \
batch_size=$batch_size  output_filename=test_clean_decoded.jsonl  \
amp=$amp  amp_dtype=bfloat16  use_cer=false num_workers=1 \
return_hypotheses=false ctc_decoding.greedy.batched_inference=true

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

Additional Information

• Related to Consider refactoring CTC greedy decoding #8891

Here are the speedups in terms of throughput. First, I applied this diff to transcribe_speech.py in order to run multiple times:

modified   examples/asr/transcribe_speech.py                                                                                                                                                                
@@ -407,7 +407,33 @@ def main(cfg: TranscriptionConfig) -> Union[TranscriptionConfig, List[Hypothesis                                                                                                       
                 override_cfg.augmentor = augmentor                                                                                                                                                         
                 override_cfg.text_field = cfg.gt_text_attr_name                                                                                                                                            
                 override_cfg.lang_field = cfg.gt_lang_attr_name                                                                                                                                            
-                transcriptions = asr_model.transcribe(audio=filepaths, override_config=override_cfg,)                                                                                                      
+                                                                                                                                                                                                           
+                transcriptions = None                                                                                                                                                                      
+                                                                                                                                                                                                           
+                torch.cuda.cudart().cudaProfilerStart()                                                                                                                                                    
+                torch.cuda.nvtx.range_push("GALVEZ_START")                                                                                                                                                 
+                for i in range(5):                                                                                                                                                                         
+                    if i == 1:                                                                                                                                                                             
+                        # import nvtx                                                                                                                                                                      
+                        # pr = nvtx.Profile()                                                                                                                                                              
+                        # pr.enable()  # begin annotating function calls                                                                                                                                   
+                        # ctx = torch.autograd.profiler.emit_nvtx()                                                                                                                                        
+                        # ctx.__enter__()                                                                                                                                                                  
+                        # torch.cuda.cudart().cudaProfilerStart()                                                                                                                                          
+                        pass                                                                                                                                                                               
+                    import time                                                                                                                                                                            
+                    del transcriptions                                                                                                                                                                     
+                    start_time = time.time()                                                                                                                                                               
+                    transcriptions = asr_model.transcribe(audio=filepaths, override_config=override_cfg,)                                                                                                  
+                    end_time = time.time()                                                                                                                                                                 
+                    print("RTFx=", 5.1 * 60 * 60 / (end_time - start_time))                                                                                                                                
+                    if i == 1:                                                                                                                                                                             
+                        # pr.disable()                                                                                                                                                                     
+                        # ctx.__exit__(None, None, None)                                                                                                                                                   
+                        # torch.cuda.cudart().cudaProfilerStop()                                                                                                                                           
+                        pass                                                                                                                                                                               
+                torch.cuda.nvtx.range_pop() # "GALVEZ_START"                                                                                                                                               
+                torch.cuda.cudart().cudaProfilerStop()                                                                                                                                                     
                                                                                                                                                                                                            
     if cfg.dataset_manifest is not None:                                                                                                                                                                   
         logging.info(f"Finished transcribing from manifest file: {cfg.dataset_manifest}")                                                                                                                  

Then I ran with and without this option turned on:

Without batched inference:

set -euo pipefail

batch_size=32
amp=true

echo "GALVEZ: Conformer CTC"
python examples/asr/speech_to_text_eval.py  pretrained_name=nvidia/parakeet-ctc-1.1b dataset_manifest=/home/dgalvez/scratch/data/test_other_sorted_downward.json  batch_size=$batch_size  output_filename\
=test_clean_decoded.jsonl  amp=$amp  amp_dtype=bfloat16  use_cer=false num_workers=1 return_hypotheses=false ctc_decoding.greedy.batched_inference=false
RTFx= 1128.651983273993
RTFx= 1248.7464311388146
RTFx= 1327.39890551856
RTFx= 1344.5921792715505
RTFx= 1346.41703283495
[NeMo I 2024-05-02 15:07:14 speech_to_text_eval:210] Dataset WER/CER 3.76%/1.41%

With batched inference:

set -euo pipefail

batch_size=32
amp=true

echo "GALVEZ: Conformer CTC"
python examples/asr/speech_to_text_eval.py  pretrained_name=nvidia/parakeet-ctc-1.1b dataset_manifest=/home/dgalvez/scratch/data/test_other_sorted_downward.json  batch_size=$batch_size  output_filename=test_clean_decoded.jsonl  amp=$amp  amp_dtype=bfloat16  use_cer=false num_workers=1 return_hypotheses=false ctc_decoding.greedy.batched_inference=true
RTFx= 1157.725053779039
RTFx= 1265.9191489003545
RTFx= 1361.67639098725
RTFx= 1354.9981484942944
RTFx= 1350.2222749634843
[NeMo I 2024-05-02 15:25:17 speech_to_text_eval:210] Dataset WER/CER 3.76%/1.41%

You can see that RTFx throughput improvements are only modest. However, in terms of inference latency, there is a large speed up at the P99 latency levels, because the sometimes long calls to max() on the CPU no longer happen. We don't have a great way to measure this right now with NeMo's current tooling, so you will just have to believe me for now. See #8891 for details about this problem.

[galv]

To verify that latency does indeed fall, I ran the above prefixed by nsys profile --env-var=NSYS_NVTX_PROFILER_REGISTER_ONLY=0 -t nvtx -c nvtx -p GALVEZ_START with batch_inference=false and batch_inference=true. I added an nvtx range just for the decoder.

batch_inference=true case:

Time	Total Time	Instances	Avg	Med	Min	Max	StdDev	Style	Range
93.9%	71.801 s	1	71.801 s	71.801 s	71.801 s	71.801 s	0 ns	PushPop	GALVEZ_START
6.1%	4.683 s	460	10.180 ms	1.173 ms	796.420 μs	117.883 ms	19.312 ms	PushPop	decoder

batched_inference=false case:

Time	Total Time	Instances	Avg	Med	Min	Max	StdDev	Style	Range
86.7%	79.410 s	1	79.410 s	79.410 s	79.410 s	79.410 s	0 ns	PushPop	GALVEZ_START
13.3%	12.217 s	460	26.558 ms	3.800 ms	2.125 ms	171.248 ms	38.294 ms	PushPop	decoder

You can see that the decoder takes half the time it used to, and that its maximum runtime is 118ms, down from 171 ms.

[titu1994]

The new algorithm part looks great, but shouldn't be fused into CTC decoding but he treated as a strategy similar to how RNNT does it.

@nithinraok upto you if you want to fuse the two together like this with a flag or use rnnt strategy pattern here in CTC

[titu1994]

Looks far cleaner than before, minor comments but overall good enough to merge.

1. Can we make this the default algo ? If users want inference features that don't work in this mode they can switch to slower code path. NeMo 2.0 is the time for switching defaults if we want to.

2. can we name this greedy_batch similar to RNNT

[titu1994]

Hmm can't we call it greedy_batched similar to RNNT? I understand both are technically batched, and vectorized is more appropriate, however conformity is nice for both us and users

[galv]

Good observation. I changed to "greedy_batched", and generally speaking changed "vectorized" to "batched" based on your observation about conformity.

[galv]

Will need to get back to you about what it takes to make this the default algo. But the main concern is that many existing ".nemo" files will have the "strategy" field of the relevant config set to "greedy". It's not really clear to me how we can convert those users using pretrained models to use the "greedy_batched" strategy. For what it's worth, my tests check that the interfaces are effectively the same, so I could just use GreedyBatchedCTCInfer whenever "greedy" strategy is specified, unless that distresses you.

[pablo-garay]

I see CICD pipeline passed: https://github.com/NVIDIA/NeMo/actions/runs/8976148083
Let me know if you need/ready to merge

[titu1994]

Hmm I see your point. Lets do this -

1. Set default in config to greedy_batch
2. Add logging message to greedy (onetime log) that informs users that greedy_batch strategy is faster.
3. All newer models will automatically be using the greedy batch path then

I'm not average to replacing "greedy" with greedy_batch, but then we need a different flag name to allow users to fall back to older codepath. Maybe we can call it "greedy_legacy" then?

[galv]

@titu1994 done. To clarify, I went with the warn-once path rather than making "greedy" pointed to the new implementation.

[titu1994]

Looks great ! Thanks for the refactor !