Pré-treinamento Eficiente com Superposição de Tokens

Resumo

O pré-treinamento de Grandes Modelos de Linguagem é frequentemente proibitivamente caro e ineficiente em escala, exigindo modificações complexas e invasivas para alcançar alta taxa de transferência de dados. Neste trabalho, apresentamos o Treinamento por Superposição de Tokens (TST), um método simples de substituição direta que melhora significativamente a taxa de transferência de dados por FLOPs durante o pré-treinamento, sem modificar o paralelismo, o otimizador, o tokenizador, os dados ou a arquitetura do modelo. O TST é realizado em duas fases: (i) uma fase de superposição altamente eficiente, onde combinamos muitos tokens contíguos em um único pacote e treinamos usando um objetivo de entropia cruzada multi-hot (MCE), e (ii) uma fase de recuperação, onde retornamos ao treinamento padrão. Avaliamos extensivamente o TST nas escalas de 270M e 600M parâmetros e validamos em modelos de 3B e um modelo de mistura de especialistas A1B de 10B, demonstrando que ele é altamente robusto em diferentes configurações. Em última análise, o TST supera consistentemente a perda de base e as avaliações downstream, e em configurações de perda igual, o TST proporciona uma redução de até 2,5x no tempo total de pré-treinamento na escala de 10B A1B.

English

Pre-training of Large Language Models is often prohibitively expensive and inefficient at scale, requiring complex and invasive modifications in order to achieve high data throughput. In this work, we present Token-Superposition Training (TST), a simple drop-in method that significantly improves the data throughput per FLOPs during pre-training without modifying the parallelism, optimizer, tokenizer, data, or model architecture. TST is done in two phases: (i) A highly efficient superposition phase where we combine many contiguous tokens into one bag and train using a multi-hot cross-entropy (MCE) objective, and (ii) a recovery phase where we revert back to standard training. We extensively evaluate TST on the scale of 270M and 600M parameters and validate on 3B and a 10B A1B mixture of experts model, demonstrating that it is highly robust in different settings. Ultimately, TST consistently outperforms baseline loss and downstream evaluations, and under equal-loss settings, TST yields up to a 2.5x reduction in total pre-training time at the 10B A1B scale.