Reduza suas perdas em Modelos de Linguagem de Grande Vocabulário.

Cut Your Losses in Large-Vocabulary Language Models

November 13, 2024

Autores: Erik Wijmans, Brody Huval, Alexander Hertzberg, Vladlen Koltun, Philipp Krähenbühl

cs.AI

Resumo

À medida que os modelos de linguagem crescem cada vez mais, também crescem seus vocabulários. Isso tem deslocado de forma desproporcional a pegada de memória dos LLMs durante o treinamento para uma única camada: a entropia cruzada no cálculo da perda. A entropia cruzada constrói uma matriz de logit com entradas para cada par de tokens de entrada e itens de vocabulário e, para modelos pequenos, consome uma ordem de magnitude mais memória do que o restante do LLM combinado. Propomos o Corte da Entropia Cruzada (CCE), um método que calcula a perda de entropia cruzada sem materializar os logit para todos os tokens na memória global. Em vez disso, o CCE calcula apenas o logit para o token correto e avalia a log-sum-exp sobre todos os logit dinamicamente. Implementamos um kernel personalizado que realiza as multiplicações de matriz e a redução de log-sum-exp sobre o vocabulário na memória flash, tornando o consumo de memória global para o cálculo da entropia cruzada negligenciável. Isso tem um efeito dramático. Tomando o modelo Gemma 2 (2B) como exemplo, o CCE reduz a pegada de memória do cálculo da perda de 24 GB para 1 MB, e o consumo total de memória durante o tempo de treinamento da cabeça classificadora de 28 GB para 1 GB. Para melhorar o rendimento do CCE, aproveitamos a esparsidade inerente do softmax e propomos pular elementos do cálculo do gradiente que têm uma contribuição negligenciável (ou seja, abaixo da precisão numérica) para o gradiente. Experimentos demonstram que a redução dramática no consumo de memória é alcançada sem sacrificar a velocidade ou a convergência do treinamento.

English

As language models grow ever larger, so do their vocabularies. This has shifted the memory footprint of LLMs during training disproportionately to one single layer: the cross-entropy in the loss computation. Cross-entropy builds up a logit matrix with entries for each pair of input tokens and vocabulary items and, for small models, consumes an order of magnitude more memory than the rest of the LLM combined. We propose Cut Cross-Entropy (CCE), a method that computes the cross-entropy loss without materializing the logits for all tokens into global memory. Rather, CCE only computes the logit for the correct token and evaluates the log-sum-exp over all logits on the fly. We implement a custom kernel that performs the matrix multiplications and the log-sum-exp reduction over the vocabulary in flash memory, making global memory consumption for the cross-entropy computation negligible. This has a dramatic effect. Taking the Gemma 2 (2B) model as an example, CCE reduces the memory footprint of the loss computation from 24 GB to 1 MB, and the total training-time memory consumption of the classifier head from 28 GB to 1 GB. To improve the throughput of CCE, we leverage the inherent sparsity of softmax and propose to skip elements of the gradient computation that have a negligible (i.e., below numerical precision) contribution to the gradient. Experiments demonstrate that the dramatic reduction in memory consumption is accomplished without sacrificing training speed or convergence.