Flash RNN-T Loss: 2x Faster Loss and up to 50x VRAM Reduction

[MahmoudAshraf97]

What does this PR do ?

This is an exact implementation of RNNT loss using time-tiling and kernel fusion to reduce memory usage

We can safely delete all intermediate tensors since we recalculate the activations in the backward step

Results

All tests were using H100 SXM
Below are the VRAM consumption and Throughput in samples/s, the semi-transparent part of the bars is the result of tuning fused batch size between 8 to 64 in a batch of 64 across different buckets, batch sizes 1-4 were barely different in memory consumption than 8 but much slower, so they are excluded from the chart
Note that the current main branch does not support Numba BF16 and #15996 is used to enable it

The tunable memory-throughput knob is max_joint_rows, it decides how much lattice elements per tile, 200k was choosen as the default as it is the sweet spot IMO

The loss is almost identical and the difference can be attributed to different GEMM reduction order and different dropout implementation, the mask is drawn per call and is not what F.dropout produces for the same torch seed.

Flash run finished 20k steps in 166 minutes and Numba finished in 207 minutes, so 20% e2e time saved, Flash also used 22% less total memory when using the same batch sizes for all buckets. The model used was a fast conformer medium with 1023+1 Vocab size, 8xH100 were used.

Usage

Add this to your rnnt training config

  model:
    joint:
      fuse_loss_wer: true
      fused_batch_size: 4  # required by the fused joint step; this loss does not read it
    loss:
      loss_name: "flash_rnnt"
      flash_rnnt_kwargs:
        fastemit_lambda: 0.0
        clamp: -1.0
        max_joint_rows: 200000  # rows of one tile

Notes for Reviewers

1. the numpy implementation of warp loss still imports numba, this prevents some triton tests from running in any environment with missing numba although it is not needed and only needs a reference implementation.
2. wer computation was moved into a separate util for usage in both regular fused path and flash path, note that this causes a slight difference in WER calculations since flash computes the WER for the whole batch while the other losses compute the mean of the subbatches.
3. This is completely opt-in, the default pathes are not touched

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

Who can review?

@artbataev @GNroy @titu1994 @hainan-xv

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[MahmoudAshraf97]

This code path is not reachable since we raise an error in the joint if the required activation isn't implemented, so the return type is never None

[MahmoudAshraf97]

This code path is not reachable since we raise an error in the joint if the required activation isn't implemented, so the return type is never None

[nithinraok]

@artbataev could you help review this PR