Réduisez vos pertes dans les modèles de langage à large vocabulaire.

Cut Your Losses in Large-Vocabulary Language Models

November 13, 2024

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

cs.AI

Résumé

À mesure que les modèles linguistiques deviennent de plus en plus grands, leurs vocabulaires s'élargissent également. Cela a déplacé de manière disproportionnée l'empreinte mémoire des LLMs pendant l'entraînement vers une seule couche : l'entropie croisée dans le calcul de la perte. L'entropie croisée construit une matrice de logit avec des entrées pour chaque paire de jetons d'entrée et d'éléments de vocabulaire et, pour les petits modèles, consomme un ordre de grandeur de mémoire plus important que le reste du LLM combiné. Nous proposons Cut Cross-Entropy (CCE), une méthode qui calcule la perte d'entropie croisée sans matérialiser les logit pour tous les jetons dans la mémoire globale. Au lieu de cela, CCE ne calcule que le logit pour le jeton correct et évalue le log-sum-exp sur tous les logit à la volée. Nous mettons en œuvre un noyau personnalisé qui effectue les multiplications de matrices et la réduction log-sum-exp sur le vocabulaire en mémoire flash, rendant la consommation de mémoire globale pour le calcul de l'entropie croisée négligeable. Cela a un effet spectaculaire. En prenant le modèle Gemma 2 (2B) comme exemple, CCE réduit l'empreinte mémoire du calcul de la perte de 24 Go à 1 Mo, et la consommation totale de mémoire pendant le temps d'entraînement de la tête de classification de 28 Go à 1 Go. Pour améliorer le débit de CCE, nous exploitons la sparsité inhérente de softmax et proposons de sauter les éléments du calcul du gradient qui ont une contribution négligeable (c'est-à-dire inférieure à la précision numérique) au gradient. Les expériences démontrent que la réduction spectaculaire de la consommation de mémoire est accomplie sans sacrifier la vitesse d'entraînement ou la convergence.

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.