Escalonando o TransNormer para 175 Bilhões 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

Resumo

Apresentamos o TransNormerLLM, o primeiro Modelo de Linguagem de Grande Escala (LLM) baseado em atenção linear que supera os modelos convencionais baseados em atenção softmax tanto em precisão quanto em eficiência. O TransNormerLLM evolui da arquitetura de atenção linear anterior, o TransNormer, por meio de modificações avançadas que incluem incorporação posicional, aceleração de atenção linear, mecanismo de portão, normalização de tensores, aceleração e estabilização de inferência. Especificamente, utilizamos o LRPE juntamente com um decaimento exponencial para evitar problemas de diluição de atenção, permitindo que o modelo mantenha interações globais entre os tokens. Além disso, propomos o Lightning Attention, uma técnica de ponta que acelera a atenção linear em mais de duas vezes no tempo de execução e reduz o uso de memória em impressionantes quatro vezes. Para aprimorar ainda mais o desempenho do TransNormer, empregamos um mecanismo de portão para suavizar o treinamento e um novo esquema de normalização de tensores para acelerar o modelo, resultando em uma aceleração impressionante de mais de 20%. Adicionalmente, desenvolvemos um algoritmo robusto de inferência que garante estabilidade numérica e velocidade de inferência consistente, independentemente do comprimento da sequência, demonstrando eficiência superior tanto nas etapas de treinamento quanto de inferência. A escalabilidade está no cerne do design do nosso modelo, permitindo implantação contínua em clusters de grande escala e facilitando a expansão para modelos ainda mais extensos, tudo isso mantendo métricas de desempenho excepcionais. A validação rigorosa do design do nosso modelo é alcançada por meio de uma série de experimentos abrangentes em nosso corpus auto-coletado, que possui um tamanho superior a 6TB e contém mais de 2 trilhões de tokens. Para garantir a qualidade e relevância dos dados, implementamos uma nova estratégia de autolimpeza para filtrar os dados coletados. Nossos modelos pré-treinados serão disponibilizados para promover avanços da comunidade em 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.