Flash RNN-T cuts memory use and speeds speech transcription
Flash RNN-T: Memory-Efficient Joint Network and Transducer Loss Computation
Machine LearningDistributed, Parallel, and Cluster ComputingSound
Summary
Speech recognition models called RNN-Ts can be slow and need a lot of memory when learning to understand language. The authors created a new way called Flash RNN-T that uses less memory and processes speech faster by optimizing how computations are done. This method works especially well on modern processors like the H100 and can handle longer speech inputs where other methods run out of memory. They show a balance between speed and memory use, making streaming speech recognition more efficient.
What this means in practice
- •For speech recognition engineers: Build efficient streaming speech models that use less memory and run faster on GPUs like H100 by adopting Flash RNN-T implementations.
- •For gpu system architects: Design GPU kernels and memory scheduling that optimize throughput and resource use for neural network training involving large vocabularies and alignments.
Authors
Mahmoud Ashraf
Abstract
The recurrent neural network transducer (RNN-T) remains a practical choice for streaming speech recognition, but training requires large intermediate tensors. We present Flash RNN-T, a joint and loss implementation combining specialized kernels with optional tiled recomputation over the full alignment space and vocabulary. Tiled and untiled Flash form the entire measured throughput-memory frontier on medium and long workloads. On H100, untiled Flash achieves 24-39% higher throughput with 26-56% less memory than the fastest baseline on these workloads wherever it fits. Tiling reduces memory further through recomputation. On medium, tiled Flash uses ninefold less memory than that baseline at 16% lower throughput when both use one backward call per batch. On long, it fits where all tested alternatives with that backward schedule run out of memory with 80 GB. These comparisons use each implementation's lowest-memory configuration reaching at least 95% of its peak throughput.