FlashFFTConv : Convolutions efficaces pour les longues séquences avec les Tensor Cores

FlashFFTConv: Efficient Convolutions for Long Sequences with Tensor Cores

November 10, 2023

Auteurs: Daniel Y. Fu, Hermann Kumbong, Eric Nguyen, Christopher Ré

cs.AI

Résumé

Les modèles convolutifs avec des filtres longs ont démontré des capacités de raisonnement de pointe dans de nombreuses tâches impliquant des séquences longues, mais ils restent en retard par rapport aux Transformers les plus optimisés en termes de temps d'exécution réel. Un goulot d'étranglement majeur est la Transformée de Fourier Rapide (FFT)—qui permet aux convolutions longues de s'exécuter en temps O(N logN) pour une séquence de longueur N, mais qui présente une faible utilisation matérielle. Dans cet article, nous étudions comment optimiser la convolution FFT. Nous identifions deux goulots d'étranglement clés : la FFT n'utilise pas efficacement les unités de multiplication matricielle spécialisées, et elle entraîne des coûts élevés en termes d'entrées/sorties entre les couches de la hiérarchie mémoire. En réponse, nous proposons FlashFFTConv. FlashFFTConv utilise une décomposition matricielle qui calcule la FFT en utilisant des unités de multiplication matricielle et permet la fusion de noyaux pour les séquences longues, réduisant ainsi les entrées/sorties. Nous présentons également deux algorithmes de convolution parcimonieuse—1) les convolutions partielles et 2) les convolutions fréquentiellement parcimonieuses—qui peuvent être implémentés simplement en sautant des blocs dans la décomposition matricielle, offrant ainsi des opportunités supplémentaires d'économies de mémoire et de calcul. FlashFFTConv accélère les convolutions FFT exactes jusqu'à 7,93 fois par rapport à PyTorch et atteint une accélération de bout en bout jusqu'à 4,4 fois. Avec le même budget de calcul, FlashFFTConv permet à Hyena-GPT-s d'atteindre une perplexité améliorée de 2,3 points sur le PILE et à M2-BERT-base d'atteindre un score GLUE supérieur de 3,3 points—égalant ainsi des modèles avec deux fois plus de paramètres. FlashFFTConv atteint également une précision de 96,1 % sur Path-512, une tâche de vision à haute résolution où aucun modèle n'avait précédemment dépassé 50 %. De plus, les convolutions partielles permettent des modèles pour des séquences plus longues—produisant le premier modèle ADN capable de traiter les gènes humains les plus longs (2,3 millions de paires de bases)—et les convolutions fréquentiellement parcimonieuses accélèrent les modèles pré-entraînés tout en maintenant ou en améliorant la qualité du modèle.

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.