Mise à l'échelle de TransNormer à 175 milliards de paramètres

Scaling TransNormer to 175 Billion Parameters

July 27, 2023

Auteurs: Zhen Qin, Dong Li, Weigao Sun, Weixuan Sun, Xuyang Shen, Xiaodong Han, Yunshen Wei, Baohong Lv, Fei Yuan, Xiao Luo, Yu Qiao, Yiran Zhong

cs.AI

Résumé

Nous présentons TransNormerLLM, le premier modèle de langage de grande taille (LLM) basé sur une attention linéaire qui surpasse les modèles conventionnels utilisant l'attention softmax, tant en termes de précision que d'efficacité. TransNormerLLM évolue à partir de l'architecture d'attention linéaire précédente, TransNormer, en apportant des modifications avancées incluant l'incorporation positionnelle, l'accélération de l'attention linéaire, un mécanisme de gating, la normalisation tensorielle, l'accélération et la stabilisation de l'inférence. Plus précisément, nous utilisons LRPE conjointement avec une décroissance exponentielle pour éviter les problèmes de dilution de l'attention tout en permettant au modèle de conserver les interactions globales entre les tokens. De plus, nous proposons Lightning Attention, une technique de pointe qui accélère l'attention linéaire de plus de deux fois en temps d'exécution et réduit l'utilisation de la mémoire par un facteur remarquable de quatre. Pour améliorer davantage les performances de TransNormer, nous exploitons un mécanisme de gating pour fluidifier l'entraînement et un nouveau schéma de normalisation tensorielle pour accélérer le modèle, résultant en une accélération impressionnante de plus de 20%. Par ailleurs, nous avons développé un algorithme d'inférence robuste qui assure la stabilité numérique et une vitesse d'inférence constante, indépendamment de la longueur de la séquence, démontrant une efficacité supérieure tant pendant l'entraînement que lors de l'inférence. La scalabilité est au cœur de la conception de notre modèle, permettant un déploiement fluide sur des clusters à grande échelle et facilitant l'expansion vers des modèles encore plus vastes, tout en maintenant des performances exceptionnelles. Une validation rigoureuse de notre conception de modèle est réalisée à travers une série d'expériences exhaustives sur notre corpus auto-collectionné, qui dépasse 6TB de taille et contient plus de 2 trillions de tokens. Pour garantir la qualité et la pertinence des données, nous mettons en œuvre une nouvelle stratégie d'auto-nettoyage pour filtrer les données collectées. Nos modèles pré-entraînés seront publiés pour favoriser les avancées communautaires dans le domaine des LLM efficaces.

English

We present TransNormerLLM, the first linear attention-based Large Language Model (LLM) that outperforms conventional softmax attention-based models in terms of both accuracy and efficiency. TransNormerLLM evolves from the previous linear attention architecture TransNormer by making advanced modifications that include positional embedding, linear attention acceleration, gating mechanism, tensor normalization, inference acceleration and stabilization. Specifically, we use LRPE together with an exponential decay to avoid attention dilution issues while allowing the model to retain global interactions between tokens. Additionally, we propose Lightning Attention, a cutting-edge technique that accelerates linear attention by more than twice in runtime and reduces memory usage by a remarkable four times. To further enhance the performance of TransNormer, we leverage a gating mechanism to smooth training and a new tensor normalization scheme to accelerate the model, resulting in an impressive acceleration of over 20%. Furthermore, we have developed a robust inference algorithm that ensures numerical stability and consistent inference speed, regardless of the sequence length, showcasing superior efficiency during both training and inference stages. Scalability is at the heart of our model's design, enabling seamless deployment on large-scale clusters and facilitating expansion to even more extensive models, all while maintaining outstanding performance metrics. Rigorous validation of our model design is achieved through a series of comprehensive experiments on our self-collected corpus, boasting a size exceeding 6TB and containing over 2 trillion tokens. To ensure data quality and relevance, we implement a new self-cleaning strategy to filter our collected data. Our pre-trained models will be released to foster community advancements in efficient LLMs.