Beperk uw verliezen bij taalmodellen met een groot 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

Samenvatting

Naarmate taalmodellen steeds groter worden, groeit ook hun vocabulaire. Dit heeft de geheugenfootprint van LLM's tijdens training onevenredig verschoven naar één enkele laag: de kruis-entropie in de verliesberekening. Kruis-entropie bouwt een logit-matrix op met vermeldingen voor elk paar invoertokens en vocabulaire-items en verbruikt voor kleine modellen een orde van grootte meer geheugen dan de rest van de LLM bij elkaar. Wij stellen Cut Cross-Entropy (CCE) voor, een methode die de kruis-entropie verlies berekent zonder de logaritmes voor alle tokens in het globale geheugen te materialiseren. In plaats daarvan berekent CCE alleen de logaritme voor het juiste token en evalueert de log-sum-exp over alle logaritmes on-the-fly. We implementeren een aangepaste kernel die de matrixvermenigvuldigingen en de log-sum-exp reductie over het vocabulaire in flashgeheugen uitvoert, waardoor het globale geheugenverbruik voor de kruis-entropieberekening verwaarloosbaar wordt. Dit heeft een dramatisch effect. Als we het Gemma 2 (2B) model als voorbeeld nemen, vermindert CCE de geheugenfootprint van de verliesberekening van 24 GB naar 1 MB, en het totale geheugenverbruik tijdens de training van de classifier head van 28 GB naar 1 GB. Om de doorvoer van CCE te verbeteren, maken we gebruik van de inherente spaarzaamheid van softmax en stellen voor om elementen van de gradiëntberekening over te slaan die een verwaarloosbare (d.w.z. onder numerieke precisie) bijdrage aan de gradiënt hebben. Experimenten tonen aan dat de dramatische vermindering van geheugenverbruik wordt bereikt zonder snelheid of convergentie van de training op te offeren.

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.