Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Flash RNN-T позволяет сократить память при обучении потокового распознавания речи, не урезая допустимые варианты сопоставления аудио и текста. На длинной тестовой нагрузке реализация поместилась в 80 ГБ памяти H100, тогда как все проверенные альтернативы с одним обратным проходом на пакет не запустились, хотя работа не рецензирована и числа в ней получили сами авторы. Команда может выбирать между более высокой скоростью и меньшим расходом памяти, не меняя целевую функцию RNN-T.
Почему функция потерь RNN-T занимает столько памяти
RNN-T подходит для потокового распознавания: модель обрабатывает звук последовательно и не ждёт окончания записи. При обучении ей нужно учесть разные способы сопоставить позиции аудио с символами расшифровки. Эти варианты образуют решётку выравниваний.
Совместная сеть RNN-T объединяет состояние звукового кодировщика с состоянием уже распознанного текста в каждой допустимой позиции решётки. Затем она вычисляет ненормализованные оценки для всего словаря. Объём промежуточных состояний и этих оценок растёт вместе с длиной записи, расшифровки и словаря.
Сама функция потерь использует гораздо меньше данных: вероятность пустого перехода и вероятность следующего правильного символа. Обычная реализация всё равно хранит оценки всего словаря, потому что они понадобятся при вычислении градиентов.
Flash RNN-T объединяет сложение состояний, функцию активации и dropout в специализированном ядре GPU. Другое ядро нормализует оценки, возвращает только две нужные вероятности, а при обратном проходе записывает градиенты поверх исходного буфера. Это сокращает как число временных массивов, так и обмен данными с памятью GPU.
Блоки ограничивают память внутри одной записи
Версия без разбиения обрабатывает все допустимые позиции пакета сразу. Она убирает лишние промежуточные массивы, но расход памяти всё ещё растёт с размером решётки. Такой режим подходит, пока пакет целиком помещается на GPU.
Блочная версия упаковывает позиции решётки в группы ограниченного размера. Группа может закончиться внутри одной записи или объединить позиции соседних записей. После прямого прохода реализация сохраняет только компактные вероятности переходов, а скрытые состояния и оценки словаря удаляет.
При обратном проходе Flash RNN-T вычисляет удалённые данные заново, обрабатывает градиенты и сразу освобождает память группы. Сохранённое начальное значение генератора случайных чисел восстанавливает ту же маску dropout, поэтому повторное вычисление не меняет результат.
Размер блока ограничивает крупнейшие промежуточные массивы, но не всю память процесса. Состояния кодировщиков, параметры модели, градиенты и сама решётка остаются в памяти. Чем меньше блок, тем ниже пик потребления, но тем больше запусков ядер и повторных вычислений.
Когда результат меняет план разработки
На средних и длинных нагрузках обе версии Flash RNN-T заняли всю измеренную границу компромисса: ни одна конфигурация конкурентов не была одновременно быстрее и экономнее по памяти. Без разбиения Flash RNN-T дал до 39% больше пропускной способности и потребовал до 56% меньше памяти, чем самый быстрый внешний вариант, когда помещался на GPU.
Блочная версия на средней нагрузке использовала в девять раз меньше памяти, чем Apex с сохранением графа всего пакета, но работала на 16% медленнее. Это прямой обмен вычислений на память: модель повторно строит оценки словаря во время обратного прохода.
Если совместная сеть и функция потерь уже упираются в память на средних или длинных записях, Flash RNN-T стоит включить в технический прототип. Сначала имеет смысл проверить режим без блоков: он сохраняет более высокую скорость. Разбиение нужно, когда полный пакет не помещается или важно оставить один обратный проход на пакет.
Работа сравнивает Flash RNN-T с Apex, NeMo Numba и Optimized Transducer. Для каждой реализации выбирали конфигурацию с минимальной памятью среди вариантов, которые сохраняли не менее 95% её максимальной измеренной скорости.
Проверка охватывает три пакета из речевых корпусов, две архитектуры GPU и одну конфигурацию модели. Измеряли только совместную сеть и функцию потерь, а не полное обучение с кодировщиком и загрузкой данных. Поэтому результаты обосновывают замену этого узла, но сами по себе не предсказывают ускорение всей системы.
На H100 и RTX 5090 преимущество по границе памяти и скорости сохранилось, однако влияние отдельных специализированных ядер различалось. Перед заменой реализации нужен замер на целевом GPU, с реальным распределением длительности записей и длины расшифровок.
Источники
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
