Training speech recognition models can quickly exhaust GPU memory. For RNN-T, an architecture used in streaming speech recognition, a major source of that memory usage is the calculation that connects the audio to its transcript. As both grow longer, the temporary data produced by this calculation grows rapidly.
We built Flash RNN-T (referred to as ‘Flash’ in this post) to reduce that memory requirement. It combines more efficient GPU operations with a way to process the calculation in small pieces, keeping less temporary data and recreating it when needed.
Why the memory grows
Consider a recording of someone saying “hello there.” The transcript tells us the words and their order, but it does not tell us exactly when each word was spoken. A different speaker might stretch the first word or pause before the second. Training needs to account for these possible timings rather than assume one fixed match between the audio and the text.
RNN-T handles this by considering all valid alignments between the recording and its transcript. Picture a grid: one axis follows the audio through time, while the other follows progress through the transcript. A path through the grid describes one possible alignment. The training calculation combines the probabilities of all valid paths to measure how well the model explains the transcript. That measure is called the loss.
The same transcript, different possible timings.
At each position, a part of the model called the joint network combines what the model has heard with the text it has already accounted for. It produces scores for every possible next token - a word or word fragment in the model's vocabulary. But speech unfolds over many audio steps, and each step does not necessarily produce a new token. The model therefore also scores a special blank token: “advance to the next audio step without adding text.” Blank can occur during speech, not just silence.
This is where the storage grows. Twice as many audio steps means roughly twice as many grid positions. If the transcript is also twice as long, the grid is roughly four times larger. Each position carries an entire vocabulary's worth of scores, alongside the joint network's internal values. Training on a batch of recordings multiplies that storage again.
Calculating the loss is only the first part of training. The model then uses backpropagation to work backward through the calculation and determine how its parameters should change. This requires intermediate values from the earlier calculation, so they are normally kept in memory until then. These saved values can consume much of the available GPU memory, even when the recordings are only around 20 seconds long.
Compute it, use It, make Room
Keeping those intermediate values avoids repeating work, but the loss calculation needs much less information than the joint network produces. At each grid position, an alignment matching the transcript can either advance through the audio using blank or produce the next transcript token. Calculating these two probabilities requires scores across the vocabulary; once calculated, however, the two probabilities are enough to continue the alignment calculation.
This gives us a way to work through the grid without keeping all its vocabulary scores in memory. Flash groups grid positions into tiles. For each tile, it computes the scores, saves the two probabilities per position, and discards the scores and internal joint network values. The next tile reuses that working space. Once all tiles are processed, their saved probabilities provide what we need to calculate the loss.
To use that loss to train the model, we still need the values we discarded. Flash therefore revisits each tile during backpropagation, recreates the values needed for that calculation, and frees them before moving to the next tile. This is called activation checkpointing: we repeat part of the computation so that the large intermediate tensors only need to exist for one tile at a time.
That exchange gives us control over the memory requirement. Smaller tiles use less working space, while larger tiles give the GPU more work at once and reduce the overhead of processing separate pieces. The saved probabilities still occupy space across the grid, but the much larger joint network intermediates now depend on the tile size rather than the length of the recording.
Making each operation more efficient
Recomputing those values saves memory, but it also means doing more work. To make these calculations more efficient, Flash combines consecutive steps so the GPU can pass results directly between them, instead of writing them to memory and reading them back. This also removes the need to store some temporary results.
Flash implements these combined steps in specialized GPU functions, called kernels. Its joint kernel combines the audio and text representations and processes the result without saving a separate intermediate result between steps.
The same approach helps when converting vocabulary scores into the two probabilities needed by the loss. Rather than storing a probability for every vocabulary entry and then selecting two, the kernel calculates and writes those two directly. During backpropagation, it also overwrites the scores with their gradients, avoiding another equally large allocation.
These improvements are useful with or without recomputation. Flash therefore also offers an untiled mode that keeps the joint intermediates rather than recreating them. Both modes use the specialized kernels; tiled execution spends additional computation to save more memory, while untiled execution avoids that cost when the available memory allows it.
What changed in practice
Together, these changes make the full model cheaper to train. Switching from NeMo’s Numba-based RNN-T loss to tiled Flash reduced training-step time by 22% and peak GPU memory by 46%. We measured this with a 114-million-parameter speech recognition model on eight H100 GPUs, keeping the model, data and batch size the same.
Projecting those short timing runs to 30,000 training steps gives about 4 hours 11 minutes with Flash instead of 5 hours 23 minutes—roughly 72 minutes saved. At the same hourly GPU rate, that means 22% less compute cost for those steps. Across repeated training runs, the same GPU budget can support more experiments.
What this unlocks
Memory limits affect which experiments a team can run. Reducing the working space required by the joint network and loss leaves more room for the rest of the model, larger batches, or longer training examples. Which option is useful depends on the training setup; the important change is that fewer of those choices are dictated by temporary tensors that do not need to coexist.
The paper explains the implementation and experiments in detail. The repository includes the benchmark code and measurement records for reproducing the comparisons and exploring the tradeoff on other hardware.
