Cut Your Losses in Large-Vocabulary Language Models
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.
6736d4060cbcd485282f342c
Erik Wijmans
extracted_confirmed
2024-11-15T22:47:34.444Z
664ce8c217586a9634370552
/avatars/b8f96ed0dc321c833d654eab24a2443d.svg
Erik Wijmans
user
erikwijmans
6736d4060cbcd485282f342d
Brody Huval
admin_assigned
2024-11-15T13:59:30.106Z
6685b9b6eff6f8ec51725397
/avatars/2425f72d0d6be1dadc6f7797a94687ba.svg
Brody Huval
user
brodyhuval
6736d4060cbcd485282f342e
Alexander Hertzberg
6736d4060cbcd485282f342f
Vladlen Koltun
admin_assigned
2024-11-15T13:59:45.191Z
6707b1a9356f027ca986d6b1
/avatars/d32eae81b27dc4f7090196d410e96876.svg
Vladlen Koltun
user
vkoltun
6736d4060cbcd485282f3430
Philipp Krähenbühl
Cut Your Losses in Large-Vocabulary Language Models
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.
Cut Cross-Entropy reduces the memory footprint of large language models during training by efficiently computing cross-entropy loss without materializing logits for all tokens.
Cross-Entropy Cut Cross-Entropy (CCE) logit matrix flash memory vocabulary log-sum-exp matrix multiplications softmax gradient computation