Speedup RNN-T greedy decoding

[artbataev]

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

# 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:

• 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)

[github-actions]

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.

[artbataev]

jenkins

[github-actions]

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.

[artbataev]

jenkins

[GNroy]

LGTM, but see comments.
I'd like to especially commend your tests, thanks for improving NeMo!

[GNroy]

Use the following commented print statements

not commented

[artbataev]

This was copied from the code nearby.
I reworked the test: instead of just printing the alignment, I use non-batched greedy decoding as a reference, and check if the batched version returns the same results.

[artbataev]

jenkins

[GNroy]

LGTM, thanks!

[titu1994]

Excellent work. Minor comments on inline documentation of the actual decoding loop and explain what is loop labels.

I also want to ask why the separation of joint into 3 functions - it seems ok but for example allows HAT to use less memory efficient path which can cause oom.

Finally, excellent tests, much better coverage of cases than before

[titu1994]

It would be better to have the similar API name across the RNNT Joints, is it necessary to change this ?

[artbataev]

The API is changed for all Joints, starting from AbstractRNNTJoint (see details in Slack)

[artbataev]

Now it is the following:

class AbstractRNNTJoint(NeuralModule, ABC):
    @abstractmethod
    def project_encoder(self, encoder_output):
        raise NotImplementedError()  # can be Linear or identity

    @abstractmethod
    def project_prednet(self, encoder_output):
        raise NotImplementedError()  # can be Linear or 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):
     """Full joint computation. Not abstract anymore!"""
    return self.joint_after_projection(self.project_encoder(f), self.project_prednet(g))

[titu1994]

Why remove the preemptive enc() pred() ? This is shown to be equivalent to RNNT and saves a ton of memory

[artbataev]

Inplace unsqueeze_ does not save memory.

Due to separating projections I needed to replace 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.

You can check it manually:

import torch

device = torch.device('cuda:0')

def print_allocated(device, prefix=""):
    allocated_mb = torch.cuda.max_memory_allocated() / 1024 / 1024
    print(f"{prefix}{allocated_mb:.0f}MB")


print_allocated(device, prefix="Before: ")  # Should be 0MB

# allocate memory ~projection result
data = torch.rand([128, 30 * 1000 // 10 // 8, 640], device=device)
print_allocated(device, prefix="After project encoder output: ")  # 118MB

# apply unsqueeze
data2 = data.unsqueeze(-1)  # unsqueeze returns a new tensor, but storage is the same (only metadata is new!)
print_allocated(device, prefix="After Unsqueeze: ")  # same, 118MB

[titu1994]

Revert name change

[artbataev]

It is essential to separate projections from other joint computations. It introduces no memory/computational overhead. See details in slack

[titu1994]

Revert name change. It's fine to keep joint

[artbataev]

See the comments above

[titu1994]

Explain in docstring what this isc

[artbataev]

Done, missed the class docstring before)

[titu1994]

Document line

[artbataev]

Added a comment

[titu1994]

Document

[artbataev]

Added a comment (above this line)

[titu1994]

Document

[artbataev]

Added a comment above

[titu1994]

Very neat, this is done so that jit compile is happy?

[artbataev]

Yep) There is also a test that torch.jit is fine with this structure :)

[titu1994]

Empty hys might be needed for beam search init and temp placeholders now that I remember

[artbataev]

I removed this function, this was used only when max_symbols=0 for the new decoding algorithm

[artbataev]

jenkins

[artbataev]

jenkins

[titu1994]

After detailed explanation, the changes make sense design wize. If you could put those explanation in the PR itself it will help for future discussion. Thanks again for the significant speedup !

[artbataev]

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) implementation 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

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)).