Flash RNN-T: calcolo della rete congiunta e della funzione di perdita del trasduttore con un uso efficiente della memoria

Abstract degli autori

Il trasduttore basato su rete neurale ricorrente (RNN-T) resta una scelta pratica per il riconoscimento vocale in streaming, ma l’addestramento richiede grandi tensori intermedi. Presentiamo Flash RNN-T, un’implementazione della rete congiunta e della funzione di perdita che combina kernel specializzati con la possibilità di ricalcolare a blocchi i valori sull’intero spazio degli allineamenti e sull’intero vocabolario. Flash, con e senza suddivisione in blocchi, occupa l’intera frontiera misurata tra throughput e uso della memoria per i carichi di lavoro medi e lunghi. Su H100, Flash senza suddivisione in blocchi raggiunge un throughput superiore del 24-39% usando il 26-56% di memoria in meno rispetto alla baseline più veloce per questi carichi, nei casi in cui può essere eseguito entro la memoria disponibile. La suddivisione in blocchi riduce ulteriormente l’uso della memoria mediante il ricalcolo. Sui carichi medi, Flash con suddivisione in blocchi usa nove volte meno memoria della stessa baseline, con un throughput inferiore del 16%, quando entrambi eseguono una chiamata backward per batch. Sui carichi lunghi, può essere eseguito entro 80 GB di memoria, mentre tutte le alternative testate con la stessa cadenza di chiamate backward esauriscono la memoria. Per ciascuna implementazione, questi confronti usano la configurazione che richiede meno memoria tra quelle che raggiungono almeno il 95% del suo throughput massimo.