Escalando TransNormer a 175 Mil Millones de Parámetros

Scaling TransNormer to 175 Billion Parameters

July 27, 2023

Autores: 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

Resumen

Presentamos TransNormerLLM, el primer modelo de lenguaje grande (LLM) basado en atención lineal que supera a los modelos convencionales basados en atención softmax tanto en precisión como en eficiencia. TransNormerLLM evoluciona de la arquitectura de atención lineal anterior, TransNormer, mediante modificaciones avanzadas que incluyen incrustaciones posicionales, aceleración de atención lineal, mecanismos de compuerta, normalización de tensores, aceleración de inferencia y estabilización. Específicamente, utilizamos LRPE junto con un decaimiento exponencial para evitar problemas de dilución de la atención mientras permitimos que el modelo conserve interacciones globales entre tokens. Además, proponemos Lightning Attention, una técnica de vanguardia que acelera la atención lineal más del doble en tiempo de ejecución y reduce el uso de memoria en un notable cuatro veces. Para mejorar aún más el rendimiento de TransNormer, aprovechamos un mecanismo de compuerta para suavizar el entrenamiento y un nuevo esquema de normalización de tensores para acelerar el modelo, logrando una impresionante aceleración de más del 20%. Asimismo, hemos desarrollado un algoritmo de inferencia robusto que garantiza estabilidad numérica y velocidad de inferencia consistente, independientemente de la longitud de la secuencia, demostrando una eficiencia superior durante las etapas de entrenamiento e inferencia. La escalabilidad es el núcleo del diseño de nuestro modelo, permitiendo un despliegue fluido en clústeres a gran escala y facilitando la expansión a modelos aún más extensos, todo mientras se mantienen métricas de rendimiento excepcionales. La validación rigurosa de nuestro diseño de modelo se logra mediante una serie de experimentos exhaustivos en nuestro corpus autogenerado, que supera los 6TB de tamaño y contiene más de 2 billones de tokens. Para garantizar la calidad y relevancia de los datos, implementamos una nueva estrategia de autolimpieza para filtrar los datos recopilados. Nuestros modelos preentrenados serán liberados para fomentar avances comunitarios en LLMs eficientes.

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.