Accelerate CTC greedy decoding by around 10%
State: closed | Pull request: yes
Labels: name: ASR

# What does this PR do ?

Accelerate CTC greedy decoding by around 10x, speeding up overall CTC runtime for parakeet-ctc-1.1b by close to 10%.

**Collection**: ASR

# Changelog 

There were a few problems with CTC greedy decoding before:

- It would copy GPU memory to pageable memory originally, which is very slow compared to copying to pinned memory. Unfortunately, cudaMallocHost() is synchronous and "slow", but fortunately pytorch has a free list of recent pinned memory allocations in its pinned memory allocator, allowing us to reduce the number of these calls. The cost of the memcpy before my change is excessive, taking ~20-40ms. After changing to use pinned memory, it takes only 2-4ms. This is the primary source of the 10x speedup. Note that doing the initial pinned memory allocations is slow (discussed later in this message).

<img width="789" alt="image" src="https://github.com/NVIDIA/NeMo/assets/4767568/11a1addd-b9a7-4060-84f7-712ea36062e3">

- A single scalar of logits_len/decoder_lengths was copied from GPU to CPU at a time, rather than all at once in a single call. This caused neddless overhead. We could also use a pinned memory allocation for copying logits_len as well, but using pinned memory is less important for small allocations. The original code's trace looks like this:
<img width="1061" alt="image" src="https://github.com/NVIDIA/NeMo/assets/4767568/9467224f-a1a0-492d-a4a5-9f3804e975ae"> Each red box followed by a green box is copying a single scalar value from GPU to CPU

- detatch() was called in function marked with @nograd

- For some reason, someone was using torch_tensor.numpy().tolist() instead of torch_tensor.tolist(). I don't believe numpy() makes a copy of the data, but it is unnecessary.

The main important improvements are the first two bullet points. Everything else is unimportant.

There are more opportunities for improvement. In particular, logits gets copied to cpu twice if trcfg.return_hypotheses is True. However, this overhead is now small since it will always just use the previous pinned host allocation from the free list.

Performance improvements:

I ran this code on an A100 GPU, on a machine with 16 CPU cores:

```
python examples/asr/speech_to_text_eval.py \
pretrained_name=nvidia/parakeet-ctc-1.1b \
dataset_manifest=/home/dgalvez/scratch/data/test_other.json \
batch_size=16  output_filename=test_other_decoded.jsonl  amp=true \
amp_dtype=bfloat16  use_cer=false num_workers=1
```

I applied this diff to run multiple times:

```
modified   examples/asr/transcribe_speech.py                                                                                                                                                              
@@ -367,14 +367,19 @@ def main(cfg: TranscriptionConfig) -> Union[TranscriptionConfig, List[Hypothesis                                                                                                    
                     decoder_type=cfg.decoder_type,                                                                                                                                                       
                 )                                                                                                                                                                                        
             else:                                                                                                                                                                                        
-                transcriptions = asr_model.transcribe(                                                                                                                                                   
-                    audio=filepaths,                                                                                                                                                                     
-                    batch_size=cfg.batch_size,                                                                                                                                                           
-                    num_workers=cfg.num_workers,                                                                                                                                                         
-                    return_hypotheses=cfg.return_hypotheses,                                                                                                                                             
-                    channel_selector=cfg.channel_selector,                                                                                                                                               
-                    augmentor=augmentor,                                                                                                                                                                 
-                )                                                                                                                                                                                        
+                for i in range(5):                                                                                                                                                                       
+                    if i == 1:                                                                                                                                                                           
+                        torch.cuda.cudart().cudaProfilerStart()                                                                                                                                          
+                    transcriptions = asr_model.transcribe(                                                                                                                                               
+                        audio=filepaths,                                                                                                                                                                 
+                        batch_size=cfg.batch_size,                                                                                                                                                       
+                        num_workers=cfg.num_workers,                                                                                                                                                     
+                        return_hypotheses=cfg.return_hypotheses,                                                                                                                                         
+                        channel_selector=cfg.channel_selector,                                                                                                                                           
+                        augmentor=augmentor,                                                                                                                                                             
+                    )                                                                                                                                                                                    
+                    if i == 1:                                                                                                                                                                           
+                        torch.cuda.cudart().cudaProfilerStop()                                                                                                                                           
                                                                                                                                                                                                          
     if cfg.dataset_manifest is not None:                                                                                                                                                                 
         logging.info(f"Finished transcribing from manifest file: {cfg.dataset_manifest}")                                                                                                                
```

Time to do each of 5 evaluations of librispeech test other before my changes: 

32 seconds
28
29
29
29

Average: 29.4 seconds
Average excluding first (warmup): 28.75 seconds

Time to do each of 5 evaluations of librispeech test other after my changes:

30 seconds
26
26
25
25
Average: 26.4 seconds
Average excluding first (warmup): 25.5 seconds

This corresponds to an almost 10% speedup. This meets expectations, since 10% of the time was originally spent on CTC greedy decoding before.

P99 latency can increase for the first few calls to the model with this change. This is because we are now calling cudaMallocHost(), and then doing a pinned memory copy from GPU to CPU, which is slower than doing a paged memory copy from GPU to CPU. However, the cudaMallocHost() calls will be cached over time, allowing us to avoid the overhead of them in later evaluations of the dataset. The expense of a cudaMallocHost() (or cudaHostAlloc()) call can be seen here:

<img width="921" alt="image" src="https://github.com/NVIDIA/NeMo/assets/4767568/fd7fd4bc-833d-43ba-9226-959ebb7c9ee6">



# Usage

See above

# 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?
- [X] 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.



jenkins

@titu1994 FYI, I made a fix. I tested by simply running running

```
python examples/asr/speech_to_text_eval.py  model_path=$HOME/.cache/huggingface/hub/models--nvidia--parakeet-ctc-1.1b/snapshots/085a3de63c7598065b072cd8f2182e6a5fa593eb/parakeet-ctc-1.1b.nemo  dataset_manifest=/home/dgalvez/scratch/data/test_other.json  batch_size=16  output_filename=test_other_decoded.jsonl  amp=true  amp_dtype=bfloat16  use_cer=false num_workers=1
```

with and without "CUDA_VISIBLE_DEVICES="

But the existing CTC test cases would have caught the error anyway if jenkins had run successfully, if jenkins had worked instead of getting stuck on fetching the git repo.

I also double checked my throughput results. They are a little bit better than originally reported because of a dumb error on my part. We are going from 28.75 seconds (RTFx 676) to 25.5 seconds (RTFx 762) after warmup on librispeech test other. Which is a nice speedup.

jenkins

@titu1994 are you good with merging this? It doesn't break any existing behaviors and I've been using it myself for a while. This change helps make comparisons of speed for other models against CTC more valid becuase CTC is no longer artificially slowed down.