Speedup RNN-T greedy decoding
State: closed | Pull request: yes
Labels: name: ASR

# What does this PR do ?

New algorithm for greedy batched decoding for RNN-Transducer.
With large batch sizes (e.g., 128) the expected speedup for large Fast Conformer-Transducer (full evaluation time including Encoder) is 1.7x-1.9x (when using `speech_to_text_eval.py`). For small batch sizes, e.g., 16, the observed speedup is ~1.3x.
The original algorithm is preserved and can be enabled by using `loop_labels=False`

E.g., on my local machine, with bf16, bs=128, Fast Conformer-Transducer Large, full `test-other` decoding

| Algorithm | Greedy | Greedy + Alignments |
| ---- | ---- | ---- |
| Current NeMo | 45 sec | 1 min 38 sec |
| Proposed | 24 sec | 30 sec |

**Collection**: [ASR]

# Changelog 
- Add specific line by line info of high level changes in this PR.

# Usage

```shell
# default - new decoding algorithm
python examples/asr/speech_to_text_eval.py \
   model_path=<nemo_model.nemo> \
   dataset_manifest=<manifest> \
   batch_size=128 \
   output_filename=<output_mainfest_path> 

# previous algorithm is preserved and can be used with `loop_labels=false`
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=false
```

# 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)


This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

jenkins

This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

Waiting for torch.jit support #8027 

jenkins

jenkins

jenkins

jenkins

jenkins

jenkins

jenkins

jenkins

Pasting here the discussion from Slack about Joint refactoring, `joint_after_projection`, non-inplace `unsqueeze`.

## Main points:

- (1) **separating projections in Joint** from other operations is not only needed for the "memory vs speed" tradeoff. It is also helpful for speed optimization without additional memory usage (this is done not only for (2))
- (2) **The immediate projection of encoder output** is a tiny overhead, and I'm sure it is negligible compared to other operations for the RNN-T system
- (3) **inplace `unsqueeze_` does not save memory**, memory consumption is the same with unsqueeze, and there is no overhead for the change
- (4) i**mplementation of separation of projections** – I tried to preserve compatibility, readability, and usability for inheritance.


### 1) separating projections in Joint from other operations is helpful in many cases.
Even in the original encoder algorithm, when we loop over encoder frames, we can project the frame immediately (one-by-one => no memory overhead), but this will save computations: for each encoder frame, multiple evaluations for Joint are used => we waste time when recalculating the encoder vector's projection.
The new algorithm is even more sensitive to operations in Joint, and I see a substantial speedup for separating projections

### 2) The immediate projection of encoder output is a tiny overhead

I see the speedup from projecting the encoder output immediately.
So, what's the overhead, and is it significant?
* This could be considered a significant overhead when we used tiny encoders with linear memory/time complexity. For modern encoders with quadratic complexity (due to attention)
* for bs 128, 30 sec, subsampling 8, joint_hidden=640, fp32, the size of tensor will be ~118MB, for bf16 – ~59MB
    * To compare with the memory consumption of one piece of the encoder, I tried a MultiHeadAttention block (used by Conformer). It uses ~2129MB memory (one block!) due to quadratic complexity. I'm sure that 118MB or even 59MB is a tiny piece of memory compared to modern encoders (it's from Conformer large, not x+large!)
* From a practical point of view, I can easily fit bs 256,  fp32 to my desktop GPU (LibriSpeech test-other, Fast Conformer Large), and we are targeting bf16, bs 128, and smaller.
* Comparison with CTC system: We project to the final output with vocabulary size dimension (not one-by-one!), which is larger than RNN-T projection, and do not optimize this for better memory usage, sacrificing the speed

Given all these facts, it is acceptable to project the encoder output immediately. If we need a robust memory consumption optimization, we can use a separate flag (`preserve_memory`), but I don't think it is now required.

### 3) in-place `unsqueeze_` does not save memory (no overhead after separating projections)

Due to separating projections, I needed to replace the in-place `unsqueeze_` operation with `unsqueeze`. There is no overhead in memory.
According to the documentation https://pytorch.org/docs/stable/generated/torch.unsqueeze.html
> The returned tensor shares the same underlying data with this tensor.

### 4) Implementation of separation of projections.
I think the acceptable solution should:
* expose projections as the public API (`AbstractRNNTJoint`)
* should be developer-friendly
* should not lead to unnecessary code duplication
* should not break checkpoints
* should not introduce any significant overhead

I considered several possibilities.
a) we can duplicate the code for `joint` in `joint_after_projection`, but there I do not think it is a good practice to maintain the same code in 2 places (it must be the same except applying projections)

b) using `enc()` and `pred()` as functions: undesirable, since it will break the checkpoints.

c) use `enc` and `pred` with type annotations in public API
```python
class AbstractRNNTJoint(NeuralModule, ABC):
   enc: Callable # can be Identity
   pred: Callable # can be Identity
   
  @abstractmethod
  def joint_after_projection(self, f, g):
     """This is the main method that one should implement for Joint""" 
    raise NotImplementedError()
   
  def joint(self, f, g):
     # not abstract anymore!
    return self.joint_after_projection(self.enc(f), self.pred(g))
```

d) separate abstract `project_prednet` and `project_encoder` methods – the current solution
I think it is better since if `project_prednet` and `project_encoder` are not implemented, a clear error will indicate this.

I prefer the last one because there is no overhead for the current implementation (see (3)).