FlashFFTConv: Эффективные свертки для длинных последовательностей с использованием тензорных ядер

FlashFFTConv: Efficient Convolutions for Long Sequences with Tensor Cores

November 10, 2023

Авторы: Daniel Y. Fu, Hermann Kumbong, Eric Nguyen, Christopher Ré

cs.AI

Аннотация

Сверточные модели с длинными фильтрами продемонстрировали передовые способности в решении задач с длинными последовательностями, но уступают наиболее оптимизированным трансформерам по времени выполнения. Основным узким местом является быстрое преобразование Фурье (FFT), которое позволяет выполнять длинные свертки за время O(N logN) для последовательности длины N, но имеет низкую эффективность использования аппаратного обеспечения. В данной работе мы исследуем, как оптимизировать свертку на основе FFT. Мы выявили два ключевых узких места: FFT неэффективно использует специализированные блоки для матричного умножения и вызывает дорогостоящие операции ввода-вывода между уровнями иерархии памяти. В ответ мы предлагаем FlashFFTConv. FlashFFTConv использует матричное разложение, которое вычисляет FFT с использованием блоков матричного умножения и позволяет объединять ядра для длинных последовательностей, сокращая операции ввода-вывода. Мы также представляем два алгоритма разреженной свертки: 1) частичные свертки и 2) частотно-разреженные свертки, которые могут быть реализованы простым пропуском блоков в матричном разложении, что открывает дополнительные возможности для экономии памяти и вычислений. FlashFFTConv ускоряет точные FFT-свертки до 7.93 раз по сравнению с PyTorch и обеспечивает до 4.4-кратного ускорения на всем конвейере. При одинаковом бюджете вычислений FlashFFTConv позволяет Hyena-GPT-s достичь на 2.3 пункта лучшей перплексии на наборе данных PILE, а M2-BERT-base — на 3.3 пункта выше оценки GLUE, что соответствует моделям с вдвое большим количеством параметров. FlashFFTConv также достигает 96.1% точности на задаче Path-512, высокоразрешающей задаче компьютерного зрения, где ранее ни одна модель не показывала точность выше 50%. Кроме того, частичные свертки позволяют создавать модели для более длинных последовательностей, что привело к появлению первой модели для обработки ДНК, способной обрабатывать самые длинные человеческие гены (2.3 млн пар оснований), а частотно-разреженные свертки ускоряют предобученные модели, сохраняя или улучшая их качество.

English

Convolution models with long filters have demonstrated state-of-the-art reasoning abilities in many long-sequence tasks but lag behind the most optimized Transformers in wall-clock time. A major bottleneck is the Fast Fourier Transform (FFT)--which allows long convolutions to run in O(N logN) time in sequence length N but has poor hardware utilization. In this paper, we study how to optimize the FFT convolution. We find two key bottlenecks: the FFT does not effectively use specialized matrix multiply units, and it incurs expensive I/O between layers of the memory hierarchy. In response, we propose FlashFFTConv. FlashFFTConv uses a matrix decomposition that computes the FFT using matrix multiply units and enables kernel fusion for long sequences, reducing I/O. We also present two sparse convolution algorithms--1) partial convolutions and 2) frequency-sparse convolutions--which can be implemented simply by skipping blocks in the matrix decomposition, enabling further opportunities for memory and compute savings. FlashFFTConv speeds up exact FFT convolutions by up to 7.93times over PyTorch and achieves up to 4.4times speedup end-to-end. Given the same compute budget, FlashFFTConv allows Hyena-GPT-s to achieve 2.3 points better perplexity on the PILE and M2-BERT-base to achieve 3.3 points higher GLUE score--matching models with twice the parameter count. FlashFFTConv also achieves 96.1% accuracy on Path-512, a high-resolution vision task where no model had previously achieved better than 50%. Furthermore, partial convolutions enable longer-sequence models--yielding the first DNA model that can process the longest human genes (2.3M base pairs)--and frequency-sparse convolutions speed up pretrained models while maintaining or improving model quality.