Cut Your Losses in Large-Vocabulary Language Models
Authors: Erik Wijmans
Brody Huval
Alexander Hertzberg
Vladlen Koltun
Philipp Kraehenbuehl
Venue: ICLR 2025 Oral (2025)
Keywords: large language model
large vocabulary
efficient
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose Cut Cross-Entropy (CCE), a method that computes the cross-entropy loss with negligible memory consumption.

Abstract: 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.

Published as a conference paper at ICLR 2025
CUT YOUR LOSSES
IN
LANGUAGE MODELS
Erik Wijmans' Brody Huval Alexander Hertzberg Vladlen Koltun Philipp Krähenbühl Apple
ABSTRACT
As language models grow ever larger, so do their vocabularies. This has shifted the memory footprint of LLMs during training disproportionatly to one single layer: the cross-entropy in the loss computation. Cross-entropy builds up 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-Entropyy (CCE), 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 evalu- ates the log-sum-exy over all logits on the fly. We implement a custom kernel that performs the matrix multiplications and the log-sum-exp reduction over the vocab- ulary in flash memory, making global memory consumption for the cross-entropy computation negligible. This has dramatic effect. Taking the Gemma 2 (2B) model as an example. CCE reduces the memory footprint of the loss computa- tion from 24 GB 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 lever- age the inherent sparsity of softmax and propose to skip elements of the gradient computation that have 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. https:/ /gi thub. com/ apple/ml- www--ccooossoooooo
INTRODUCTION
Progress in large language models (LLMs) has been fueled in part by an increase in parameter count, context length, and vocabulary size (the number of tokens that can be used to represent the input). As LLMs grew, so did the associated infrastructure. Large mini-batch gradient descent (Goyal et al., 2017) combined with data-parall pata-parallelism (Hillis & Steele, 1986) enabled the harnessing of increasing computational power. ZeRO (Rajbhandari et al., 2020) broke the dependence between the number of GPUs and the memory used for model parameters, gradients, and optimizer state. Activation checkpointing (Chen et al., 2016) reduced the amount of memory used for activations, supporting the development of deeper models. FlashAttention (Dao et al., 2022) reduced the memory used in self- attention from O(N2) O(N), thereby supporting longer context windows. These improvements gradually shifted the memory consumption of LLM training to one single layer the cross-entropy loss, whose memory footprint grows with the product of vocabulary size and number of tokens per batch. The cross-entropy loss is responsible for to 90% of the memory footprint of modern LLM training (see Fig. la). The problem grows only more acute with time, since even the largest contemporary vocabularies (e.g., 256K tokens) may benefit from further expansion (Tao et al., 2024).
We propose a cross-entropy implementation, Cut Cross-Entropy (CCE), that has a negligible mem- ory footprint and scales to arbitrarily large vocabularies. Our key insight is that computation of the loss and its gradient only depends on a single proprbbbabllty, that the ground-truite label. With an arithmetic reformulation, we decompose the cross-entropy loss into index matrix multiplication over a single ground-truth label and log-sum-exp operation over all vocabulary entries for each token. Each operation has small and well-defined inputs the network embeddings and classifier
"Corresponding author; ewi jmans@apple..cmm

---

Published as a conference paper at ICLR 2025
Gemma 2 2B
Liemta 3.1
Ulrra
Demira
Liama S.
Weights Optimber Grodients Activatio Checkpoints
708
86
Jan 2024
Centrrm
278
Gomma 2 27E
Phi3
Phi 3 Metiun
Reiesse Date
Gemma 28
Medium
Jon 2024
1878
Mistral
9:78
1.5 75
Jul 2023
Uzma
130
7B
Limo
12B
Liaro 7B
Noo 2.78
OPT Neo
2021
GPT Net
2,75
GPT Nwo
1:38
OPT
Feb 2019
70
Max batch siw (M Takwne]
Mix batch sioe (M Tokkers))
(a) Regular cross-cntropy
(b) Cut cross-entropy (ours)
Figure Memory use and maximum attainable batch size (in millions of tokens) for a vari- ety of frontier models on a 16-GPU (80 GB each) fully-sharded data-parallel setup (Rajbhan- dari et al., 2020) with activation checkpointing (Chen et al., 2016) and a mrixed-precision 16-bit (fp16/bf16) AdamW optimizer (Kingma & Ba, 2015; Loshchilov & Hutter, 2019). For each model, we break its memory use down into weights and optimizer states, activation checkpoints, and the rog-pobbabilities computed by the cross-entropy loss layer. Our Cut Cross-Entropp Entropy (CCE) enables increasing the batch size by 1.5x (Llama 13B) to 10x (GPT 2, Gemma 2B), with no sacrifice in speed or convergence. Exact values in Table A4.
matrix and single scalar output per token. Both operations do, however, rely large interme- diate logit matrix that computes the score for each token and potential vocabulary entry. We show that there is no need materialize this logit matrix in GPU memory. Instead, compute logits as needed in SRAM a series of custom CUDA kernels. The result is a cross-entropy computation that has negligible memory footprint, with detrimental effect on latency convergence. See Fig. lb for a breakdown of memory savings and consequent batch size increases afforded by CCE.
2 RELATED WORK
Attention mechanisms. The effectiveness of transformers (Vaswaní al., 2017) in modeling lan- guage has drawn attention to their compute and memory requirements. Multiple works have pro- posed alternatives scaled dot-product attention that reduce transformers' computation and mem- ory (Kitaev et al., 2020; Wang et al., 2020; Choromanski et al., 2021). Other model classes, such as structured state-space models (Gu al., 2022; Gu & Dao, 2023), have also shown promising results. We study a different part of the model - its classifier head that is not considered in these works.
Attention implementations. In addition alternative attention mechanisms, the community has also tackled the daunting memory consumption of LLMs via efficient implementations. Rabe & Staats (2021) developed a self-attention implementation that makes use of chunking. Chen et al. (2023) proposed an implementation that broke the operation into two stages, reduction and matrix multiplication. This makes efficient use of GPU memory and registers but requires recomputation in the forward pass. FlashAttention (Dao et al., 2022) uses online softmax (Milakov & Gimelshein, 2018) and, like CCE, materializes blocks of the N2-sized self-attention matrix in on-chip SRAM rather than slower global DRAM. This is one of the key ideas that CCE builds on to develop memory-efficient cross-entropy formulation.
Vocabulary reduction. One way to minimize the amount of memory used by the log- pogaabbbabilities over the tokens is to reduce the number 'active" tokens in the vocabulary. Grave et al. (2017) proposed use vocabulary with hierarchical structure, thereby requiring the log-probabilities for only a subset of the vocabulary at any given time. Yu et al. (2023) explore tokenization free byte-level models that operate dramatically smaller vocabularies.
Sequence and model parallelism. Sequence parallelism (Jacobs et al., 2023; Li et al., 2023) enables training very large models (with large vocabularies) by splitting individual input sequence across
2

---

Published as a conference paper at ICLR 2025
multiple GPUs. Various model parallelism techniques (Huang et al., 2019; Narayanan et al., 2019; Shoeybi et al., 2019) achieve the same goal of training very large models (with large vocabularies) by distributing the computation and memory consumption of different pieces across multiple GPUs.
Efficient cross-entropy implementations. A number of recent implementations use chunking to reduce the memory usage of the cross-entropy layer. Yet chunking induces trade-off. Memory footprint is minimized when the number chunks high, but latency is minimized when the number of chunks is low. CCE utilizes only on-chip SRAM and minimizes both memory footprint and latency. Liger Kernels (Hsu et al., 2024) make efficient use of the GPU via chunking and by computing the losstgradient simultaneously. The latter requires that any transform applied the loss (such as masking) is implemented in the kernel itself. CCE has separate forward and backward stages, enabling user-defined transformations on the loss.
PRELIMINARIES
Let P(x)- II P(z; 2 ..... Ii-1 be Large Language Model (LLM) over a vocabulary V. The LLM parameterizes an autoregressive distribution over all possible tokens 2; E V given the preceding N 1 tokens. Specifically, this distribution is the combination of backbone network f: ...1...2-111......... ---1
RD and a linear classifier C IR D
P(aa .1... 1;-1) softmax, (CTf(a) .......2 f0 ..-1)),,,,,,,
exp(Uk)
softmax& (v) E, exp(w,,)
The backbone network f(z Ii E IRD encodes token sequence in the D-dimensional feature vector. The linear classifier C E PXIM projects the embedding into output space of the vocabulary V. The softmaxx (v) produces the probability over all vocabulary entries from the unnormalized log probabilities (logits) produced by C f(11...;;-1).
3.1 VOCABULARY
LLMs represent their input (and output) typically constructed by a method such as Byte Pair Encoding (BPE) (Gage, 1994), BPE initializes the vocabulary with all valid byte sequences from a standard text encoding, such utf-8. Then, over large corpus of text, BPE finds the most frequent pair of tokens and creates a new token that represents this pair. This continues iteratively until the maximum number of tokens reached.
of tokens in vocabulary V. The vocabulary is
Large vocabularies enable single token to represent multiple characters. This reduces the length of both input and output sequences, compresses larger and more diverse documents into shorter context windows, thus improving the model's comprehension while reducing computational demands.
3.2 INFERENCE AND TRAINING
Even with a large vocabulary, sampling from an LLM is memory-efficient at inference time. Specif- ically, the LLM produces one token a time, computing P( ............. ...I,1)) and sampling from this distribution (Kwon et al., 2023). Because the distribution over the vocabulary is only needed for single token at time, the memory footprint is independent of sequence length.
At training time, the LLM maximizes the log- ligelikel of the next token:
((X)
log
(3)
Due to the structure of most backbones (Vaswani et al., 2017: Gu et al., 2022; Gu Dao, 2023), f(x1),
.N) is efficiently computed in parallel. However, activations for non-linear layers have to be saved for the backward pass, consuming significant memory. Most LLM training frameworks make of aggressive activation checkpointing (Chen al., 2016), sharding (Rajbhandari et al., 2020), and specialized attention implementations (Dao et al., 2022) to keep this memory footprint manageable.
3

---

Published as a conference paper at ICLR 2025
E
E
E
VE
Atorss Add C.S.
Indexed
Atonic
C
x
dot
c
Air
motmul
Bik.
Aup
S
S.E
VC
LS
Sar
log (exp(LSE))
exp(A
any(S by
If
>e) Sin VLSE.
S,,
#exp(LSE..)
LSE,
LSE,
V.SE.
(a) Indexed matmul (forward)
(b) Linear-log Linear-lg-ssumexp forward pass
Linear-log-suie caaa--loo-su--ppppp sum-exp,
backward pass
Figure 2: Access patterns and computation of blockwise (a) indexed matrix multiplication, (b) snrear-log-sum-e-e forward pass, and (c) Linear-lo--sum-ee backward pass. See Algorithms to for the corresponding algorithms.
With the aforementioned optimizations, the final (cross-entropp loss) layer of the LLM becomes by far the biggest memory hog. For large vocabularies, the final cross-entropy layer accounts for the majority of the model's memory footprint at training time (Fig. la). For example, the log- probabilities materialized by the cross-entropy layer account for 40% of the memory consumption of Phi 3.5 (Mini) (Abdin al., 2024) ([V] = 32, 82,064), 65% of the memory consumption of Llama 3 (8B) (Dubey et al., 2024) ([V] 128, 000), and 89% of the memory consumption of Gemma 2 (2B) (Rivière et al., 2024) (V) 256,1 256.
fact, the log- pogppobbabilities of Gemma (2B) for single sequence x with length N 80,000 use the entire available memory of an 80 GB H100 GPU. (The sequence length is factor due to the use of teacher forcing for parallelism.)
We show that a reformulation of the training objective leads to an implementation that has negligible memory consumption above what is required store the loss and the gradient.
CUT CROSS-ENTROPE
Consider the cross-entropy loss & over single prediction of the next token P(z 21-1):
(x() log softmax,,
(CT
E,
C E, log Σ cxp (C. E.).
Here the first term is a vector product over D-dimensional embeddings E; = f ... Fi-1) and a classifier C. The second term log-sum-exp operation and is independent of the next token a;- During training, we optimize all next-token predictions =e1... jointly using teacher forcing:
(c
- -log Σ e E),
=
where = F1.... Ex] and (CTE) [CT E1 ....I., EN]. The first term Equation (4) combination of an indexing operation and matrix multiplication. has efficient forward and backward passes, in terms of both compute and memory, as described in Section 4.1. The second term Equation (4) is joint log-sum-exp (LSE) and matrix multiplication operation. Section 4.2 describes how to compute the forward pass of this linear-log-sum-er joint matrix multiplication and reduction kernel. Section 4.3 describes how to compute its backward pass efficiently by taking advantage of the sparsity of the gradient over a large vocabulary. Putting all the pieces together yields memory-efficiee low-latency cross-entropy loss.
operation efficiently using
4.1 MEMORY- EFFICIENT INDEXED MATRIX MULTIPLICATION
naive computation of indexed matrix multiplication involves either explicit computation of the logits C. E with ((N]V)) memory cost, indexing into the classifier Cx [Ca, ... Cayl with
4

---

Published as a conference paper at ICLR 2025
Algorithm 1 Memory-efficient indexed matrix multiplication
Inputs:
E E R RDN,, C CER' RDx x E R Block sizes NB and DB-
Outputs:
o == (C
E)x
E RN
for blocks En, Xn. do
b Divide and x into blocks of size NB and NA, respectively
On ON for blocks Ed do
> Zero vector of size Ng in on-chip SRAM >> Divide E,, into blocks of size Dg x NB Indexed load into on-chip SRAM D. Column-wide dot product
C .n.d O, += En,d" C
end for
write On
> From on-chip SRAM to main GPU memory
end for
an O(ND) memory cost. Our implementation fuses the classifier indexing Cx with the consecutive dot product between columns Ca, and E3 in single CUDA/Triton kernel (Tillet et al., 2019). Our kernel retrieves the value T;, the æ;-thh column from C, and the i-th column from E, and stores them in on-chip shared memory (SRAM). It then performs dot product between Cs, and Eg and writes the result into global memory. The kernel uses only on-chip SRAM throughout and does not allocate any GPU memory. For efficiency, we perform all operations blockwise to make the best use of GPU cache structure. Algorithm and Fig. 2a summarize the computation and access patterns.
4.2 MEMORY-EFFICIENT LINEAR-LOG-SEE SUM-EXP, FORWARD PASS
Implementing a serial memory-efficient seear-log-sum-expp is fairly straightforward: use a triple for- loop. The innermost loop computes the dot product between C. and E,, for the v-th token and the n-th batch element. The middle loop iterates over the vocabulary, updating the log-sum-exp (LSE) along the way. Finally, the outermost loop iterates over all batch elements. Parallelizing over the outermost loop is trivial and would expose enough work to saturate the CPU due to the number of tokens in training batches (commonly in the thousands). Parallelization that exposes enough work to saturate the GPU is more challenging.
Let us first examine how efficient matrix multiplication between the batch of model output embed- dings E E DxN and the classifier C E DXIVI is implemented on modern GPUs (Kerr et al., 2017). A common method is to first divide the output O C E x into a set of blocks of size VB x NB. Independent CUDA blocks retrieve the corresponding parts E,, of with size D x NB and blocks Cm C with size x VB. and perform the inner product O Onm = C E, for-loop for large values of D. They loop over smaller size DB x NB and D x VB blocks and accumulate O On E. d C End in SRAM. Each CUDA block then writes O nm back into global memory. This method exposes enough work to the GPU and makes efficient of SRAM and 1.2 cache.
along the D dimension. Due to limited on-chip SRAM, most implementations
To produce log-sum-exp( C E),
the same blocking and parallelization strategy matrix multiplication. Each block first computes a matrix multiplication, then the log-sum-exp along the vocabulary dimension m for its block, and finally updates LSE with its result.
Note that multiple CUDA blocks are now all writing to the same location of LSF. This includes blocks in the same input range but different vocabulary ranges m. We use spin-lock an atomic operation in global memory to synchronize the updates by different CUDA blocks as this is simple to implement in our Triton framework and incurs little overhead. Alternative methods, such as an atomic compare-and-swee loop, may perform better when implementing in CUDA directly.
Algorithm and Fig. 2b summarize the computation and access patterns.
4.3 BEFIRNTEFFICIENT LINEAR-LOG-SUM-BSS BACKWARD PASS
The backward pass needs to efficiently compute two gradient updates:
0
πE log Σ cxp(CTE) and VC= =
log Σ ExppCC' TE)
E
ac
5

---

Published as a conference paper at ICLR 2025
Algorithm 2 Memory-efficient linear-log-sum-exp. forward pass
Inputs:
E E R PxD and C E IR Px|VI.
Block sizes NA, VB, and DB-
Outputs: LSE log E3 exp(C) B) RN
LSE = -O0N
00 vector of size in main GPU memory
for all pairs of blocks Ens C. do
Divide E and C into blocks of size D x Ng and D x VB > Zero matrix of size VB x NB in on-chip SRAM Divide E,, and Cp into blocks of DB NB and DB VB
Anu = OVgxNa Ny
for blocks En,d, Cr,d do
App += C Cv,d" Endd
> Blockwise matrix multiplication
end for
LSEEn = log Σ exp( TT)
Numerically stable implementation with max Locking thread-saa -afe log-add-exp
LSEn log( (exeppLSEEn) + exp(LSEnn))
end for
for a backpropagated gradient VLSE. Formally, the gradient is defined as
(S. VLSE) C and VC (S VLSE) T E
where S softmaxx C E) and refers to the row-by-row elementwise multiplication of the soft- max S and the gradient VLSE: $ S VLSE.
Computationally, the backward pass a double matrix multiplication C E and SC or S with intermediate matrices S and Å that do not fit into GPU memory and undergo a non-linear operation. We take a similar approach to the forward pass, recomputing the matrix C E implicitly in the GPU's shared memory. For the backward pass, we do not need to compute the normalization constant of the softmax, since S softmax( C E) exp(C" E LSE). This allows us reuse the global synchronization of the forward pass, and compute efficiently in parallel.
We implement the second matrix multiplication in the main memory of the GPU, as canonical blockwise implementation would require storing or synchronizing S. Algorithm 3 and Fig. 2c sum- marize the computation and access patterns. A naive implementation of this algorithm requires zero additional memory but is slow due to repeated global memory load and store operations. We use two techniques to improve the memory access pattern: gradient filtering and vocabulary sorting.
Gradient filtering. By definition, the softmax S sums to one over the vocabulary dimension. If stored in bfloat16 with 7-bit fraction, any value below & 2 2-12 will likely ignored due to truncation in the summation or rounding in the normalization." This has profound implications & 4096 entries have non-trivial values and
for the softmax matrix S: For any column, at most 1 contribute to the gradient computation. All other values are either rounded to zero truncated. In practice, the sparsity of the softmax matrix much higher: empirically, in frontier models we evaluate, less than 0.02% of elements are non-zero. Furthermore, the sparsity of the softmax matrix grows as vocabulary size increases.
Algorithm 3, we take advantage of this sparsity and skip gradient computation for any block whose corresponding softmax matrix Spm has only negligible elements. We chose the threshold = = 2-12 In practice, this leads to 3.5x speedup without loss of precision in any gradient computation. See Section 5 for a detailed analysis.
12 to be the smallest bfloat 16 value that not truncated.
The efficiency of gradient filtering is directly related to the block-level sparsity of the softmax matrix. We cannot control the overall sparsity pattern without changing the output. However, we can change the order of the vocabulary to create denser local blocks for more common tokens.
Vocabulary sorting. Ideally the vocabulary would be ordered such that all tokens with non-trivial gradients would be contiguously located. This reduces the amount of computation wasted by par- tially populated blocks - ideally blocks would either be entirely empty (and thus skipped) or entirely populated. We heuristically group the non-trivial gradients by ordering the tokens by their average logit. Specifically, during the forward pass (described in Section 4.2) we compute the average logit
'The 5 extra bits above the fractional size (7) account for rounding rules, and the consideration that small but tiny values will likely not get truncated due the blocking strategies used to compute sum.

---

Published as a conference paper at ICLR 2025
Algorithm 3 Memory-efficient linear-lo--sum-exx, backward pass
Inputs:
E E RRNN, C E RDIVI S LSE E RR, and VLSE E R.V. Block sizes NB, VB, and DB-
Accuracy threshold
Outputs:
VE E RDN,
for all pairs of blocks En, Cp do
> Divide E and C into blocks of size D x NB and D x Vn Zero matrix of size VB x NB on-chip SRAM Divide E,, and C₂ into blocks of D3 NB and DB x VB
Ann = OVExNN
for blocks En,d. Cd,d do
Ane +- C,,d Endd end for
D Blockwise matrix multiplication
S,v exp(Ane (LSEm)
b> Compute the softmax
if (Sll(Snv < <E) then skip
> Skip computation below desired numerical precision
end if
for blocks E, Cr,d do
Divide E, and C. into blocks of DB NB and DB x VB > Locking thread-safe gradient update b Locking thread-safe gradient update
VE n.d += (Sno VLSEn) Cr,d
VCT udd += (Sne VLSE,) En,d
end for end for
per token using atomic addition. For the backward pass, we divide the vocabulary dimension [V] into blocks with similar average logit instead of arbitrarily. This requires a temporary buffer of size O(VV),, about 1 MB for the largest vocabularies in contemporary LLMs (Rivière et al., 2024).
Putting all the pieces together, we arrive at forward and backward implementations of cross-entropy that have negligible incremental memory footprint without sacrificing speed. Note that in practice, we found it to be easier and more memory-efficien to merge the indexed matrix-multilliation backward implementation with the backward pass of the linear-log- swear-log-sum-exp operator (Algorithm 3). The two operations share much of the computation and memory access pattern, see Algorithm 4.
ANALYSIS
5.1 RUNTIME AND MEMORY
First examine the runtime and memory of various implementations of the cross-entropy loss log softmax,, (C E). E). We consider batch of 8,192 tokens with vocabulary size of 256,000 and hidden dimension 2,304. This corresponds to Gemma 2 (2B) (Rivière et al., 2024). We use the Alpaca dataset (Taori et al., 2023) for inputs and labels and Gemma (2B) Instruct weights to compute E and for C. The analysis is summarized in Table 1.
The baseline implements the loss directly in PyTorch (Paszke et al., 2019). This is the default in popular frameworks such as Torch Tune (Torch Tune Team, 2024) and Transformers (Wolf et al., 2019). This method has reasonable throughput but peak memory usage of 28,000 MB of GPU memory to compute the loss+gradient (Table row 5). Due to memory fragmentation, just com- puting the loss+gradient for the classifier head requires an 80 GB GPU. torch. compile (Ansel et al., 2024) is able to reduce memory usage by 43% and computation time by 33%, demonstrat- ing the effectiveness of kernel fusion (Table row 4 vs. 5). Torch Tune (Torch Tune Team, 2024) includes method to compute the cross-entropy loss that divides the computation into chunks and uses torch. compile to save memory. This reduces memory consumption by 65% vs. Bascline and by 40% vs. torch. compile (to 9,631 MB, see Table row 3 vs. 4 and 5). Liger Kemnlss (Hsu al., 2024) provide a memory-efficiee
implementation of the cross-entropy loss that, like Torch Tune, makes uses of chunked computation reduce peak memory usage. While very effective reducing the memory footprint, using 95% less memory than Baseline, it has a detrimental effect on latency, more than doubling the wall-clock time for the computation (Table 1, row 2 vs. 4). The memory
"The gradient and loss are computed simultancously, in separate forward/hackwwrd passes.

---

Published as a conference paper at ICLR 2025
Loss
Gradient
Loss+Gradient
Method
Memory
Time
Memory
Time
Memory
Time
Lower bound
0.004 MB
1,161 MB
1,161 MB
1) CCE (Ours)
1 MB
46 ms
1,163 MB
100 mss
1,164MB
145 ms
2) Liger Kernels (Hsu et al., 2024y²
1,474 MB
304 ms
1,474 MB
304 ms
3) Torch Tune Team (2024) chunks)
8, 8,000 MB
55 ms
1,630 MB
115
9.631 MB
169
4) torch.come compile
4,000 MB
49 ms
12,000 MB
92
16,000 MB
143
5) Baseline
24,000 MB
ms
16,000 MB
122 ms
28,000 MB
208 ms
6) CCE (No Vocab Sorting)
0.09 MB
45 ms
1,162 MB
115 ms
1,162 MB
159 ms
7) CCE (No Grad. Filter)
0.09 MB
45 ms
1,163 MB
314 ms
1,162 MB
357 ms
8) CCE-Kahan
1 MB
47 ms
2,325 MB
114ms
2,326 MB
160
9) CCE- Kahan-FulIC
MB
ms
2,326 MB
268 ms
2,326 MB
313
10) CCE- SIE-Kahan-Fullll
MB
ms
M,3
247 ms
2,326 MB
292 ms
Table 1: Peak memory footprint and time to compute the loss, its gradient, and their combination. Note that intermediate buffers can often (but not always) be reused between the loss and gradient computation, resulting in lower peak memory consumption than the sum of the parts, Batch of 8, 192 tokens with vocabulary size of 256,000 and hidden dimension 2304. Embedding and classifier matrix taken during Gemma 2 (2B) training on Alpaca. Measured on an A100-SXM4 GPU with 80 GB of RAM, PyTorch 2.4.1, CUDA 12.4, rounded closest MB. Some numbers are multiples of 1,000 due to dimensions chosen and Py 's allocation strategy. 'Lower bound' is the amount of memory required for the output buffer(s), i.e., VE and VC, this is the lower bound for the memory footprint of any method. Results averaged over 5 seeds.
usage of CCE grows with O(N+IVI) as opposed O x [V] for Baseline, torch. compi le, and Torch Tune, and O(N x D) for Liger Kernels. In practice, CCE has a negligible memory footprint regardless of vocabulary size or sequence length.
Compared to the fastest method, torch. compile, CCE computes the loss slightly faster (5%, 4ms, Table row vs. 4). This is because CCE does not write all the logits global memory. CCE computes the loss+gradient slightly slower (6%, 2 ms). While CCE needs recompute C E, is able to save time in other parts of the computation. See Appendix C.I for breakdown of the backwards pass of CCE and Baseline. This increase is largely negligible as the backwrd pass for even a small LLM (2B parameters) is on the order of seconds.
The performance of CCE is enabled several fac- tors. Without vocabulary sorting CCE takes 15% ms) longer (Table row vs. 6) and without gradient filtering it is 3.4x (356 ms) longer (row vs. 7). CCE utilizes the final gradient floating point type (typically bf16) for summation in global memory. For increased numerical stability, we experiment with Kahan
Probability (log-scale)
10-2
10-5
10-8
10-11
summation (Kahan, 1965) with a higher time and memory cost (Table row vs. 8). We can further incraese the numerical stability by se- lectively applying gradient filtering to just VE and VC. When combined with Kahan sum- mation, removing gradient filtering from either VC or V E results in similar decrease of per- formance (Table row 9 or 10 vs. 8). The last variant
10-14
Token Probabilities
BF16 Cutoff
100
10
10a
10
10
10a
N-th most likely token (log-scale)
Figure Average probability for the ith most likely token, log-log plot. The probabilities very quickly vanish below numerical precision.
is particularly in- teresting for pretraining, where the numerical variants of CCE perform equivalently, shown in Section 5.3.
precision makes a difference. For fine-tuning all
In Appendix B, we demonstrate that CCE (and other methods) can be made up to 3 times faster by removing tokens that are ignored. In Appendix C we benchmark with more models. We find that as the vocabulary size (V)) to hidden size (D) ratio decreases, CCE's advantage in computation time for Loss+ Gradient decreases, but continues to save substantial amount of memory.
8

---

Published as a conference paper at ICLR 2025
Confidence Interval (p=0.95)
122
torch.compile
Training Loss
Training Loss
1.1-
Cut Cross-Entropy (Ours)
0.9
0.9
500
Gradient Sleps
Grudient Steps
(a) Gemma 2 2B
(b) Phi 3,5 Mini
1.3
1.3
Training Loss
1.1-
Training
0.93
0.9
0.8
0.5
0.7
100
200
300
400
500
600
700
0.7
100
200
300
400
500
600
700
Gradient Steps
Gradient Steps
(c) Qwen 2.5 7B
(d) Mistral Nemo
Figure 4: Training loss curves for four models on the Alpaca dataset (Taori et al., 2023). The loss curves for CCE and torch. compile are nearly indistinguishable, showing that the gradient filtering in CCE does not impair convergence. Results averaged over 5 seeds.
5.2 GRADIENT FILTERING
Fig. 3 shows the sorted softmax probability of vocabulary entries. Note that the probabilities vanish very quickly and, for the top 10t most likely tokens, there linear relationship between log rank and log probability. Second, by the ~50tth most likely token, the probability has fallen bellow our threshold for gradient filtering.
This explains why we are able to filter so many values from the gradient computation without af- fecting the result. these sparsity levels, most blocks of the softmax matrix S are empty.
5.3 TRAINING STABILITY
Fine-tuning. We fine-tune Qwen 2.5 7B Instruct (Qwen Team, 2024), Phi 3.5 Mini Instruct (Abdin al., 2024), Gemma 2B Instruct (Rivière al., 2024), and Mistral NeMo (Mistral AI Team, 2024) on the Alpaca Dataset (Taori et al., 2023) using CCE and torch. compile ll as the control. CCE and torch. compile have indistinguishable loss curves, demonstrating that the gradient filtering in CCE does not impair convergence (Fig. 4).
Prefraining. In our initial experiments using CCE for pretraining, we found that validation per- plexity suffered due to two sources of eтror. First, gradient filtering when applied to VC' causes no gradient to be propagated to tokens that have little no support the training set. This does not cause issues when fine-tuning but does when pretraining. Second, CCE performs a summation in most efficient to perform this reduction in the desired final floating point type. In pretraining, the resulting loss of precision reduces performance. We use Kahan summation (Ka- han, 1965) to recover this loss of precision. This changes correspond to
global memory.
FulIC.
We pretrain Qwen 2.5 7B Instruct (Qwen Team, 2024), Phi 3.5 Mini Instruct (Abdin et al., 2024), Gemma 2B Instruct (Rivière et al., 2024), and Mistral NeMo (Mistral AI Team, 2024) on the 5% of the Open WebText Dataset (Gokaslan et al., 2019) using CCE-Kahan- FullC and torch. compile il. We report validation perplexity on a held-out 0.25% of Open WebText and find that CCE-Kahaan FullC produces identical curves torch. compile (Fig. 5).
We make two notes about CCE- Krea-Kahm-- KulIC.FFlIC. First, the increased memory usage of CCE-Kahan FulIC CCE is due temporary buffers used in the backward pass. The size of these buffers

---

Published as a conference paper at ICLR 2025
Validatrom Perplexily
Validation Perplexity
Confidence Interval (p-0.95)
1803
100
torch. torch.compile
1403
CCE-Kahan-Fullee (Ours)
1203
100
100
80
40
1000
1.00
500
1000
1500
Gradient Steps
Gradient Steps
(a) Gemma 2 2B
(b) Phi 3.5 Mini
200
200-
Validation Perpiexity
1802
Validatiion Perplexity
180
1603
160g
140-
140.
120
120
1003
100
803
60
60
40
1500
1000
Gradient Steps
Gradient Steps
(c) Qwen 2.5 7B
(d) Mistral Nemo
Figure 5: Validation perplexity curves for four models on trained using 5% of the Open WebText dataset (Gokaslan et al., 2019). The validation set is 0.25% subset of Open WebText that does not overlap with the train set. We find that CCE-Kahan- Kahan-Full FullC matches torch. compile. Results averaged over 5 seeds.
is typically less than the amount of free memory needed to rematerialize activations when using activation/rradient checkpoint (Chen al., 2016). Thus CCE-Kahan-Fai FullC often shares the same memory saving benefits as CCE. Second, the increased computation time of CCE-Kahan-Fulle torch. compile often offset by the larger batch sizes CCE- Kahan-FullC enables. In our experi- ments with Mistral NeMo, CCE-Kahan-Fulle training time by hours (16%) compared to torch. compile.
vs.
enabled doubling the batch size, thereby decreasing
6 DISCUSSION
As vocabulary size [V] has grown in language models, so has the memory footprint of the loss layer. The memory used by this one layer dominates the training-time memory footprint of many recent language models. We described CCE, an algorithm to compute & = log softmax,
.;-1) and its gradient with negligible memory footprint.
Beyond the immediate impact on compact large-vocabulary LLMs, as illustrated in Fig. 1, we expect that CCE may prove beneficial for training very large models. Specifically, very large models are trained with techniques such as pipeline parallelism (Huang et al., 2019; Narayanan et al., 2019). Pipeline parallelism works best when all stages are equally balanced in computation load. Achiev- ing this balance is easiest when all blocks in the network have similar nemory-to-computarii ratios. The classification head is currently an outlier, with disproportionately high memory-to- computation ratio. CCE may enable better pipeline balancing or reducing the number of stages.
We implemented CCE using Triton (Tillet et al., 2019). Triton creates efficient GPU kernels and enables rapid experimentation but has some limitations in control flow. Specifically, the control flow must be specified at the block level and therefore our thread-safe log-add-exp and gradient filtering are constrained to operate at the block level as well. We expect that implementing CCE in CUDA may bring further performance gains because control flow could be performed at finer-grained levels.
could also be interesting to extend CCE to other classification problems where the number of classes is large, such as image classification and contrastive learning.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Marah Abdin, Sam Ade Jacobs, Ammar Ahmad Awan, Jyoti Aneja, Ahmed Awadallah, Hany Awadalla, Nguyen Bach, Amit Bahree, Arash Bakhtiari, Harkirat S. Behl, et al. Phi-3 technical report: A highly capable language model locally on your phone, 2024. URL https://arxiv org/abs/2404 14219.
Jason Ansel, Edward Z. Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky, Bin Bao, Peter Bell, David Berard, Evgeni Burovski, et al. Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation. In ACM International Conference on Architectural Support for Programming Languages and Operating Systems, 2024.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost, 2016. URL htttt://aaaxiv http:://wwwxvvv...bbb/o//
06174.
Yu-Hui Chen, Raman Sarokin, Juhyun Lee, Jiuqiang Tang, Chuo-Ling Chang, Andrei Kulik, and Matthias Grundmann. Speed is all you need: On-device acceleration of large diffusion mod- els via GPU-aware optimizations. In Conference on Computer Vision and Pattern Recognition, Workshops, 2023.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane. Tamás Sarlós. Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking attention with per- formers. In International Conference on Learning Representations, 2021.
Tri Dao, Daniel Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness, In Neural Information Processing Systems, 2022.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Llama 3 herd of models, 2024. URL httpp:://arki wwwp:/ggg///////eee/eooes 21783.
Philip Gage. A new algorithm for data compression. The C Users Journal, 12(2):23-38, 1994.
Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Openwebtext corpus, 2019. URL http://Skyliono00 github. ht/pp/eeeeeextorppu.....
Priya Goyal, Piotr Dollár, Ross B. Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training Ima- geNet in hour, 2017. URL http://arxiv, wwwperiwvw.ooraeeeeeeeeee orggabs///0/0 02677.
Edouard Grave, Armand Joulin, Moustapha Cissé, David Grangier, and Hervé Jégou. Efficient softmax approximation for gpus. In International Conference on Machine Learning, 2017.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2023.
URL https://aaaiiii orgg http///bs//3.....eeee 00752.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2022.
W. Daniel Hillis and Guy L. Steele. Data parallel algorithms. Commun. ACM, 29012):1170-1188 1986.
Pin-Lun Hsu, Yun Dai, Vignesh Kothapalli, Qingquan Song, Shao Tang, and Siyu Zhu. Liger- Kerne:: Efficient Triton kernels for LLM training, 2024. URL https::////t thub. com/linkedin/ Liger-Kernel.
Yanping Huang, Youlong Cheng. Ankur Bapna, Orhan Firat, Dehao Chen, Mia Xu Chen, Hy- oukJoong Lee, Jiquan Ngiam, Quoc V. Le, Yonghui Wu, and Zhifeng Chen. GPipe: Efficient training of giant neural networks using pipeline parallelism. Neural Information Processing Systems, 2019.
11

---

Published as a conference paper at ICLR 2025
Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Leon Song, Samyam Rajbhandari, and Yuxiong He. Deepspeed ulysses: System optimizations for en- abling training of extreme long sequence transformer models, 2023. URL https://doi.oo org/ 10. 20.48550/arX/5 2309. 14589.
William Kahan. Pracniques: further remarks reducing truncation emros. Communications of the
ACM, 1965.
Andrew Kerr, Duane Merrill, Julien Demouth, and John Tran. CUTLASS: Fast lin- algebra in CUDA C++, 2017. URL https://developer. cutlass- Linear-algebra-cuue
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International
Conference on Learning Representations, 2015.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2020.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Symposium on Operating Systems Principles, 2023.
Shenggui Li, Fuzhao Xue, Chaitanya Baranwal, Yongbin Li, and Yang You. Sequence parallelism: Long sequence training from system perspective. In Association for Computational, 2023.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019.
Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax, 2018. URL
http://arxxv.org aww.aa////.. 02867.
Mistral AI Team. Mistral NeMo, 2024. URL htt.://t.s//mstrraaaaaaa///iiirr
Deepak Narayanan, Aaron Harlap. Amar Phanishayee, Vivek Seshadri, Nikhil R. Devanur, Gre- gory R. Ganger, Phillip B. Gibbons, and Matei Zaharia. Pipedream: Generalized pipeline paral- lelism for DNN training. In ACM Symposium on Operating Systems Principles, 2019.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. PyTorch: An imperative style, high-performance deep learning library. In Neural Information Processing Systems, 2019.
Qwen Team. Qwen2.5: A party of foundation models, September 2024. URL https:///wwnt: github. tttblo////////oo 5/
Markus N. Rabe and Charles Staats. Self-attention does not need O(n2) memory, 2021. URL
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. ZeRO: Memory optimizations toward training trillion parameter models. In International Conference for High Performance Computing, Networking, Storage and Analysis, 2020.
Morgane Rivière, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhupatiraju, Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, Johan Ferret, et al. Gemma 2: Improving open language models at a practical size, 2024. URL https://arxiv.orr/altt 2408. 2408.00118. 00118.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-LM: Training multi-billion parameter language models using model par- allclism, 2019. URL http: http:///arxti /atxivaxiivv.v///////////////i petrxiv.orgge orgo hrrg/bss/1n abbs/1909 1909. 08053.
Chaofan Tao, Qian Liu, Longxu Dou, Niklas Muennighoff, Zhongwei Wan, Ping Luo, Min Lin, and Ngai Wong. Scaling laws with vocabulary: Larger models deserve larger vocabularies, 2024. URL https://aarae wtw.aaiii...o/ab/////////// 13623.
12

---

Published as a conference paper at ICLR 2025
Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford Alpaca: An instruction-followinn LLaMA model, 2023. URL https://gitt thub. com/tatsu- stanfoo alpaca.
Philippe Tillet, Hsiang-Tsung Kung, and David D. Cox. Triton: An intermediate language and compiler for tiled neural network computations. In ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, 2019.
Torch Tune Team. torchtune, 2024. URL htt.ttttt://gthhhee http:/ihthh......oo.eeeeeeeee httpla/ccoopppppppeee com/pyyorrch/torchtuus
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez,
Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In Neural Information Processing Systems, 2017.
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity, 2020. URL https://arxiv orgg ha//ab///200 2006. 04768.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowic7, and Jamie Brew. Huggingface's trans- formers: State-of-the-arr natural language processing. 2019.
Lili Yu, Daniel Simig, Colin Flaherty, Armen Aghajanyan, Luke Zettlemoyer, and Mike Lewis. MEGABYTE: Predicting million-byte sequences with multiscale transformers. In Neural Infor- mation Processing Systems, 2023.
13

---

Published as a conference paper at ICLR 2025
A NOTATION
Throughout the paper, we use the following notation conventions. Matrices are bold, capital letters, e.g., A. Indexed matrices are capital letters and are indexed by column and then, optionally, row. For example, given A E R N then e.g., Aj is the length N vector that is the jth column for A, AJ, then the ith value in the vector Ay. When we combine indexing and transposing, we always index and then transpose.
Vectors are bold lower-case letters, e.g., x. with the exception of LSE which the vector containing the log-sum-exp (LSE). Indexed vectors are lower-case letters,
In addition to scalar indexing, we also block index matrices when describing how our algorithms are implemented. In these cases, the matrix and vector will maintain their bold to indicate that the indexing refers to block and thus are still a matrix or vector.
Notation
Description
N matrix containing batch of inputs.
E,
D D-dimensional vector containing the embedding for the ith input.
C
A D) x classifier matrix used compute the logit for each token.
C
A D D-dimensional vector used to create the logit for the ith token.
X
A length N vector containing the inputs.
Z
A scalar that is the ith input.
Cz
length containing the vector used to create the logit for the s7th token.
C E
A [V] N matrix containing the logits over the vocabulary for each input.
(CTE)
length N vector where the ith entry is the logit for the zęth token.
LSE
length vector containing the log- log-sum-exp (LSE) for each input over the vocabulary.
E,
The vth D Na block of
Ennd
The dth DB NE block of En. An indicator matrix where value the ith column and jth is a, and otherwise.
b b']]
B REMOVING IGNORED TOKENS
It is common to have tokens that have loss computation when training LLMs in practice. Ex- amples include padding, the system prompt, user input, etc.. While these tokens must be processed by the backbone rect context for its prediction in the case of system prompts and use inputs they do not contribute directly the loss.
enable efficient batching in the of padding to give the model the cor-
In all implementations we are aware of, the logits and loss for these ignored tokens is first com- puted and then set to zero. We notice that this is unnecessary. These tokens can be removed before logits+loss computation with no change to the loss/gradient and save significant amount of com- putation.
Table Al shows the performance of all methods in Table I with filter that removes ignored tokens before logits+loss computation. This represents significant speed up for all methods but Liger Kernels. Due to heavy chunking in Liger Kernels to save memory, it is bound by kernel launch overhead, not computation, and therefore reducing the amount of computation does not increase speed. Filtering ignored tokens is also a significant memory saving for most all but CCE (because CCE already uses the minimum amount of memory possible).
C ADDITIONAL RESULTS
C.I FURTHER PERFORMANCE ANALYSIS
Table A2 shows breakdown of the time spent for different components of in the backward pass of CCE and Baseline. For CCE, we selectively disabled/cnabled portions of the kernel and measured the time saved to determine the amount of time taken by that component. For Baseline, we manually implemented each operation of the backward pass and timed them seperately.
'The gradient and loss are computed simultancously, in separate forward/hackwwrd passes.
14

---

Published as a conference paper at ICLR 2025
Loss
Gradient
Loss Grasiena
Method
Memory
Time
Memory
Time
Memory
Time
Lower bound
0.004 MB
1,161 MB
1,161 MB
1) CCE (Ours)
245 MB
17 ms
1,163 MB
37 ms
1,164 MB
54 ms
2) Liger Kemnels (Hsu al., 2024)²
1,316 MB
301
1,314 MB
303
3) Torch Tune Team (2024) (8 chunks)
3,688 MB
23
2,789 MB
6,157 MB
77
4) torch. compile
1,847 MB
19 ms
5,490 MB
34 ms
7.337
5) Baseline
10,997 MB
30 ms
7,320 MB
ms
12,826 MB
75 ms
6) CCE (No Vocab Sorting)
0.06 MB
17
1,162 MB
43 ms
1,163
60 ms
7) CCE (No Grad. Filter)
0.06 MB
17 ms
1,163 MB
110ms
1,163
8) CCE--ahan
MB
18 ms
2,325 MB
ms
2,327 MB
59 ms
9) CCE- TAEE-Kaann-Fulll
MB
18 ms
2,326 MB
98 ms
2,327 MB
114 ms
10) CCE- Hat--Kahn--FulE
MB
18ms
2,325 MB
92 ms
2,327 MB
109 ms
Table Al: Table I where all methods include a filter that removes tokens that are ignored in loss computation. This simple change represents large improvements in practice. Results averaged over 5 seeds.
Component
Baseline
CCE
logits = softcap
C
recomputation
45 ms (43.2 %)
V log soft (logits)
35 ms (28.5 %)
4.7 ms (4.4 %)
Gradient Filter
1.3 ms (1.2 %)
Vsoftcap (c
17 ms (13.7 %)
4.7 ms (4.4 %)
VE
37 ms (30.0 %)
31 ms (29.6 %)
VC
34 ms (27.7 %)
18 ms (17.3 %)
Table A2: Performance breakdown for the backward pass of CCE and Baseline. Gemma 2 (2 B) model. Batch of 8192 tokens. Alpaca dataset used to generate inputs.
CCE spends considerably less time the cross-entropy loss and softcap portions of the gradient computation. For Baseline, these are very memory intensive operations as there relatively very little computation done compared the amount of reading/writing. For CCE, the logits already in SRAM (they were just recomputed) and CCE does not write the result of this computation to main memory, saving significant amount of time.
Coincidentally, CCE spends a very similar amount of time computing the gradient wrt. the embed- dings. CCE spends less time computing the gradient wrt. the classifier. This is because the axis reduce along for the classifier, N, is shorter than the axis for the embeddings, -V-, and thus leads less contention global memory.
Compared to Baseline, CCE saves 30 ms on the gradient of the logits wrt. cross-entropy loss, 12 ms on the gradient wrt. softcapping, 5 ms the gradient wrt. E, and 15 the gradient wrt. C. This saving of 62 ms more than offsets the 45 ms spent re-computing and applying the gradient filter.
C.2 ADDITIONAL RUNTIME AND MEMORY
Table A3 shows additional results for Gemma (9 B), Gemma 2 (27 B), Qwen 2.5 (7 B) (Qwen Team, 2024), Qwen 2.5 (32 B), PHI 3.5 Mini (Abdin et al., 2024), and Mistral NeMo (Mistral AI Team, 2024) in the same setting Table 1. For each model CCE is able to reduce the total memory consumed by the loss by an order of magnitude from the baseline. For forward (Loss) and backward (Gradient) passes combined, CCE is within 3 MB of the lowest possible memory consumption. Compared to Gemma 2 (2 B) all these models have a smaller ratio of the vocabulary size to hidden dimension. This has impacts.
15

---

Published as a conference paper at ICLR 2025
First, the number of tokens that have significant gradient is largely constant (it is dependent on the data type). Therefore proportionally less of the gradient will be filtered out.
Second, for all other methods increasing the hidden dimension increase the amount of parallclism that can be achieved. Liger Kernels (Hsu et al., 2024) sets its chunk size based on IVI/D the lower that ratio, the bigger the chunk size. As WI/D continues decrease, Liger Kernels is able to make better use of the GPU. All other methods use two matrix multiplications to compute the gradient. The amount of work that can be performed in parallel compute V E and VC is B D and [V] x D, respectively'. The amount of parallel work for CCE is B x IV,, thus increasing D increases the amount of work but not the amount of parallelism. It may be possible leverage ideas from split-k matrix multiplication kernels to expose more parallelism to CCE for large values of D.
For the smallest IV/D considered, Phi 3.5 Mini (|V]=22,064, D=3,072) ours is approximately 50% slower (12 ms) than torch. compile (although it uses substantially less memory). In our exper- iments, this increase in linear-cross-e enttee loss computation time is largely negligible and only increases training time by one to two percent.
We also consider how changing the number of tokens changes performance (Figs. Al and A2). We find that CCE behaves very similarly to Baseline and torch. compile. Further, because CCE does not utilize chunking, it does not reach a point where the overhead of dispatching all the kernels becomes the dominating factor. We also find that while Kernel and Torch Tune baselines with large number of tokens, it becomes more performant than those baselines as the number of tokens reduces.
is slower than the Liger
D MEMORY USE METHOD DETAILS
Table A4 contains the numbers used to create Fig. 1. The maximum batch size for 16 GPUs was calculated by assuming that the total amount of memory available is 75 x 16 (i.e., each 80 GB GPU will be fully occupied expect for 5 GB buffer for various libraries), then subtracting the memory used for weights + optimizer + gradients and then diving by the memory used per token.
The numbers in Table A4 are computed using the following methods. When present, the number of tokens is assumed to be 65,536.
We compute the amount of memory used for intermediate activations as the number of layers times the hidden size times number of tokens times 2 bytes per bfloat 6.. This assumes the use of activa- tion/gradient checkpointing (Chen al., 2016) for transformer layer.
The amount of memory used by the logits is the number of tokens times the vocabulary size times bytes per float32. This likely undercounts the amount of memory used for computing the probability distribution, as its common to also keep a copy of the logits in 6 and, for models like Gemma 2 (Rivière al., 2024) that use logit softcapping, an additional copy of the logits after softcapping may be needed. However, this method can be uniformly applied to all models.
The amount of memory used by Weights+Opt+Grad is the number of parameters times 4 (parame- ters, gradient, and Adam first and second moments) times 2 bytes per bfloat16.
E FLOATING POINT ADDITION
Here we provide a brief explanation of floating point addition and how it relates to our proposed gradient filtering.
Given two numbers and b represented using floating point, such that [a] < [b], the following steps are performed
1. Separate the mantissa (the fractional part) and the exponent from both numbers a and b.
2. Re-write the mantissa of the smaller number (a in our case) such that it shares the same exponent as the b.
3. Add the re-written mantissa of a to the mantissa of b.
Ignoring split-k matrix multiplication kemness for simplicity.
16

---

Published as a conference paper at ICLR 2025
Loss
Gradient
Loss-Gradien
Method
Memury
Memory
Memoty
Time
Gemma By (Rivibre et al., 2024) (V]=256,000 D=3,584)
Lower bound
0.004 MB
1,806 MB
1,806MB
CCE (Ours)
MB
68
1,808 MB
141
1,809 MB
205
Liger Kernels (Hsu et al., 2024)
2,119 MB
418 ms
11,264 2,119 MB
m19
Torch Tune Team (2024) chunks)
8,000
3.264 MB
168 mss
mis
torch. compile
4,000 MB
70 ms
12,000 MR
134 ms
16,000 MB
207 ms
Baseline
24,000
102 ms
16,000 MB
ms
28, 19,000
271 ms
CCE-Kahan-FallCC
MMB
3,558 MB
384
3,559 MB
450
Gemma (27 B) (Rivière et al., 2024) (V1-)((()) D=1,605)
Lower bound
2,322
2,322
CCE (Ours)
MB
83
2,324 MB
200
2,325 MB
281 ms
Liger Kernels (Hsu et al., 2024)
2,948 MB
361 ms
2,948 MB
363 ms
Torch Tune Team (2024) chunks)
8,000 MB
4,768 MB
204
12.768 MB
296ms
worch.oopile
1,000 MB
86 ms
12,000 MB
168 ms
16,000 MB
256 ms
Baseline
24,000
119 16,000 MB
28,000
ms
CCF-Kahan Full-
1 MB
83 ms
4.
513 ms
4,575 MB
593 ms
Mistral NeMo (Mistral Al Team, 2024) [V|==331,072 1=5,120)
Lower bound
0.004 MB
1,360
1,360
CCE (Ours)
0.6 MB
52 ms
1,361 MB
129 ms
1,362 MB
180 ms
Liger Kernels (Hsu et al., 2024)
1,872
166 ms
1,872
ms
Torch Tune Team (2024) (8 chunks)
2,048 MB
19
3,318
113
5,396 MB
161 ms
torch. compile Baseline
2,048
48 ms
6.144 MB
ms
8,192
143 ms
10,240
58
8,192 MB
100
12,288 MB
161 ms
CCE-Kahan-Fulle
0.6
52 ms
2,641 MB
291 ms
2,6-12 MB
312 ms
Phi 3.5 (Abdin et al., 2024) (|V|=32,064, D=3,072)
Lower bound
0.001 MB
236 MB
236 MB
CCE (Ours)
0.2 MB
8 ms
236 MB
26 ms
236 MB
31 ms
Liger Kernels (Hsu al., 2024)
26 ms
Torch Tune Team (2024) (8 chunks)
502 MB
ms
451 MB
18ms
953 MB
30
torch. compile
1,504 MB
ms
Baseline
2,506
11 ms
2,004 MB
16 ms
3,006
27 ms
CCE Kahun-FullC
0.2
8 ms
424 MB
46 ms
424
54 mss
Qwen By (Qwen Team, 2024) (V15152666, D=3,584)
Lower bound
0.004 MB
1,096 MB
1,096 MB
CCE (Ours)
m3
1,098 MB
93 ms
1,097 MB
136 ms
Liger Kernels (Hsu al., 2024)
1,391 MB
171 ms
1,391 MB
171 ms
Torch Tune Team (2024) (8 chunks)
2,379 MB
41 m2
2.540
96 ms
4,921 MB
138 ms
tarch. compi le
2,376 MB
7,128
79
9.504 MB
121
Baseline
11,880
53 ms
0,504MB
86 ms
14,256 MB
142 ms
CCE-Kahan- Kahan-FulIC
0.6 MB
ms
2,138 MB
225 ms
2,138
ms
Qwen 2.5 (32 (Qwen Team, 2024) (VV]=152,064, 225, Lower bound
0.004 MB
1,565 MB
1,565
CCE (Ours)
0.6
1,566 MB
1,567
mss
Liger Kernels (Hsu et al., 2024) Torch Tune Team (2024) chunks) torch.compile
2,159 MB
192 ms
2,161 MB
192 ms
2,376
3,882 MB
ms
6,259
ms
2,376 MB
56 ms
7.128 MB
108 ms
9,501 MB
165 ms
Baseline
11,880 MB
68 ms
9,504 MB
115 ms
14,256 MB
186
CCE-Kahan-FulIC
0B6
3,052
326 3,033
381ms
Table A3: Memory usage and time of CCE, Liger Kernels, Torch Tune, torch. compile, and Base- line for additional models. Batch of 8, 192 tokens. Results averaged over seeds.
4. Combine the resulting mantissa and exponent of and then convert them into normalized form.
Step 2 is where truncation happens and the intuition of gradient filtering comes from. In 6 if the exponent of b is more than 27 times larger than that of a, the 7-bit mantissa no longer has enough precision to represent any of a's mantissa and in the process of re- re-writing, will be, in effect, set to zero. For gradient filtering, we are only concerned with values in the range [0, 1]. so the threshold of 2 means that only keep values that don't get rounded to zero when b 2-5
17

---

Published as a conference paper at ICLR 2025
Gradient
Lows Gradient
80
203
250
CCE (Dars)
Runtine (mz)
Runtine (es)
Runtime (mx)
Liger Kemels
160
Turch Tume (8 chunis)
1000
forch cample
Baalins
200
Kahan-FulC
70
50
50
2'
2/4
2/2
2'
213
212
Num Takens (Log Base 2 Stale)
Num Tokens (Log Bose Scale)
Num Tokens (Log Base Scale)
(a) Gemma 2 B
Loes
Grodient
Loss+ Gradient
100
CCE (Gurs)
Runbime (ms)
000
Figh Keetels
61)
Runtime
2000
Runtime
Torch Tune chunks)
200
sorch comple
411
Basefrw
100
CCE Kulln-FulIC
21
207
2
24
300
Num Tokens (Log Base Scale)
Num Tokens (Log Base Scale)
Num Tokens (Log Bose 2 Scale)
(b) Gemma 2 9 B
Loss
Gradient
Losse Gradient
600
1000
500
400
CCF (Ours]
Runtine (me)
Runtime (res)
Liger Kernels
300
Torch chunks)
60
300
torch.comple
200
Raseire
CCE Full-
100
100
g1z
y014
213
Mam Tokens (log Scale)
Num Takems (Log Rasa Scale)
Num Tokems (log Rase Scale)
(c) Gemma 2 27 B
Loss
Grodient
Loss+-Grodient
60
300
350
30
300
200
CCF (Ours)
Buntime (ms)
Buntime (ma2)
Buntime (ms)
Liger Kumels
200
Torch chunks)
20
150
.00
torch.omple
Rawdine
20
200
FuE-Kalun-FuIC
30
Num Tokens (Log Base 2 Scale)
Num Tokens (Log Base Scale)
Num Tokens (Log Base 2 Scale)
(d) Mistral NeMo
Figure A1: Performance of CCE and baselines for all models with a varying batch sizes. Results averaged over 5 seeds. Continued in Fig. A2.
18

---

Published as a conference paper at ICLR 2025
Loss
Grodient
Loss+G Grodient
50
80
CCF (Oun)
Runtine (mo)
Runtime (ms)
Uiger Kemels
Tarc Tune (8 chunks)
20
larsh
Bazdire
Sahanahn-Full
10
10
Num Tokens (Log Base 2 Scale)
Num Tokens (Log Buse Scale)
Num Tokens (Log Base 2 Scale)
(a) Phi 3,5 Mini
Gradient
Loss Gradient
250
40
200
CCE (Ours)
Runtine (mz)
Rantine (mo)
150
Runtime (mc)
Liger Kamels
30
200
Torch Tune chunks)
furch cumple
700
100
Baudine
Kahan FulC
2
212
y1
2'
213
Nam Tokees (lon Rase 2 Stale)
Num Takens (log Rese Scale]
Nur Tokens (log Basa Scale)
(b) Qwen 2.5 7 B
Loss
Grodient
Loss+GGaddent
201
250
CCE (Ouns
Bundims (mc)
201
Burdims (mc)
Liger Kemels
Runine
Torch chunks)
30
150
torch.comple
Baseline
20
200
CCE-Kolun-FuIC
10
21n
21
21a
21
211
2.
Num Tokers (Log Base 2 Scale)
Tokens (Log Scale)
Num Tokons (Log Base Scale)
(c) Qwen 2.5 32 B
Figure A2: Performance of CCE and baselines for all models with varying batch sizes. Results averaged over 5 seeds.
Model
Logits
Activations
Weights+Opt--ruu
Max Batch Size (Before)
Max Butch Size (After)
Increase
GPT
12,564 MB
1,152 MB
1,045 MB
3,866,190
69,815,595
11.9x
GPT Neo (1.3B)
12,564
MB
10,421 MB
4,268,
12,996.042
3.0x
GPT Neo (2.7 B)
12.561 MB
10,240 MB
20,740 MB
3,171,7881
7,731,585
2.2x
Gemma (2
64.000 MB
4,608 MB
19,121 MB
1,155,51 5,5
17,204.330
14.9%
Gemma (27
61.000 MB
266 MB
207,727 MB
739.
2,525,551
3.4x
Gemma 2 (2 B)
64,000 MB
7,488 MB
19,946 MB
1,108,206
10,580,057
9.5x
Llama 2 B)
8.000 MB
25.600 MB
99,303 MB
2,203,057
2,891,512
1.3x
Llama 2 17
8,000 MB
16,384 MB
51,410 MB
3,164,129
4,709.560
1.5%
Llama 3 (70 B)
32,064 MB
81.920 MB
338,282 MB
397 009
552,414
1.4x
Llama (8
32,061
16.384 MB
61,266 MB
1,579,333
1,670,136
3.0X
Mistral 7B
8,000 MB
16.384 MB
55,250 MB
3,154,108
4,694,200
1.5x
Mixtral 8x7B
8,000 MB
16.384 MB
356,314 MB
2,3111991
3,489,911
1.5x
Phi 1.5
12.574 MB
6,144 MB
10.821 MB
4,26-4,482
12,991,781
3.0x
Phi 3 Medium
8,003 MB
25.600
106, MB
2,158 821
2,873,067
1.3x
Qwen (.5 B)
37,912
16,384
58,909 MB
1,412,087
4,679,564
3.3x
Table A4: Raw data for Fig. 1. Memory usage calculated using a global batch size of 65,536.
19

---

Published as a conference paper at ICLR 2025
Algorithm 4 Memory-efficient linear-cross-entropy loss, backward pass
Inputs:
E R CE CxIVI, ISE E R.,, VCEL R., and x E R.V. Block sizes NB, VB, and DB. Accuracy threshold
= [1...., [V]].
Outputs:
VE E E VC E DXIIV]
for all pairs of blocks En, Cp do
b Divide E and C into blocks of size D x NB and D x VB Zero matrix of size Vn Nn in on-chip SRAM Divide E, and C into blocks of DB x NB and DB x VB > Blockwise matrix multiplication
Anu = for blocks En,d, Cy,d do
AN += C.T v.d En,d end for
S.v exp(Anu - LSEE)
Compute the softmax Gradient of cross-entropy loss wrt. logits
G.D. [[v X, Sn
if aal[ < E) then
skip
b Skip computation below desired numerical precision
end if
for blocks En,d, Cr.d do
Divide E,, and C. into blocks DR NB and Dr VB b Locking thread-safe gradient update > Locking thread-safe gradient update
VET n. + (Gnv VCEL.) Cr.d
VC += (Gnr VCEL.) T En,d
end for end for
20


Reviews:
[{"review_id": "hJIYTO7Sho", "reviewer": "Reviewer_XSRM", "summary": "This paper proposes \"Cut cross-entropy\" (CCE), a method that reduces the memory footprint of computing the cross-entropy loss during LM training dramatically, by never materializing the full matrix of logits/probabilities (which can be huge: batch * sequence_length * vocab_size).  To accomplish this, it realizes that the cross-entropy loss can be broken down into two components: (1) the logit for the correct next token, and (2) the log-sum-exp (log of softmax denominator) --- both of these terms are scalars, and can be computed without materializing the full logit tensor. In particular, (1) is computed via simple vector dot-products (Algorithm 1), while (2) can be computed by accumulating partial sums of the log-sum-exp, without every materializing all elements of the sum at once (Algorithm 2).  For the backward pass (Algorithm 3), the paper proposes two methods --- gradient filtering and vocabulary sorting --- that reduce the backward pass time by skipping gradient computations for blocks of the softmax matrix where all values are < 2^(-12).\n\nPutting all these pieces together, CCE is able to match the speed and quality of existing implementations of the cross-entropy loss, while only requiring a very small percentage of HBM memory (e.g., 1 MB instead of 1GB->24 GB).", "questions": "- Are there regimes where CCE is meaningfully slower than the torch.compile method?\n- There were a few elements of the backward computation that I think could be explained more clearly:\n  - What is the \"v\" index in the lines 339-341 (Algorithm 3)?\n  - Why doesn't recomputing the large $C^T E$ matrix multiplication in the backward pass (Algorithm 3) lead to slow-downs? If I understand correctly, this is because although much extra time is spent on this recomputation, less time is spent on the subsequent matrix multiplications, due to the gradient filtering/vocab sorting? Can you break down more granularly how much time each component of CCE (especially the backward pass) takes, and compare this to the naive implementations, so that it becomes clear what is happening here?\n  - Can you explain this paragraph in more detail please: \"We implement the second matrix multiplication in the main memory of the GPU, as a blockwise implementation would require storing or synchronizing S...\"? Here, is \"main memory\" HBM?\n  - I think there may be mistakes in the backward pass equations at the bottom of page 5 (lines 266-269).  Letting V be vocab size, L be sequence length, and D be hidden dimension, we can see that C is [D,V], E is [D,L], and S is [V,L].  Then for the matrix shapes to be correct, shouldn't it be:\n    - $d/dE = C (S \\cdot LSE_{grad})$ --- which is a [D,V] * [V,L] multiplication, which gives [D,L], which is the correct shape of E,\n    - $d/dC = E (S \\cdot LSE_{grad})^T$ --- which is a [D,L] * [L,V] multiplication, which gives [D,V], which is the correct shape of C.\n  - Can you include, at least in the appendix, a version of algorithm 3 that also includes the backward pass of the indexed matrix multiplication? \n  - NIT: I think it could be clearer to update the notation to be something like the following, to make Algorithms 2 and 3 easier to follow. For example, you could use $V$, $L$, $D$ to denote vocab size, sequence length, and hidden dim, and correspondingly $V_B$, $L_B$, and $D_B$ to denote the dimensions of the blocks, and $B_V = V/V_B$, $B_L = L/L_B$, $B_D = D/D_B$ to denote the number of blocks, and v, l, d to index into these blocks.\n- Can you add a discussion around sequence parallelism approaches, which can also reduce logit memory per GPU by splitting logits along sequence dimensions?", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- Reducing the memory requirement for computing the CE loss in LLMs is a strong contribution, especially as the vocabulary sizes, batch sizes, and sequence lengths of LLMs continue to grow.  This custom kernel could save many people lots of time trying to get around OOM errors during training, and make it easier to train models with larger sequence lengths/batch sizes/vocab sizes.\n- The algorithm is clever and elegant, taking inspiration from FlashAttention, which avoids materializing full attention score matrix during attention computation.\n- The experiments demonstrate that CCE can reduce training memory requirements without impacting quality/convergence during training, or training speed, relative to strong baselines (e.g., torch.compile).", "weaknesses": "- I think the section about the backward pass could be explained more clearly (see my questions below for points of confusion that could be clarified).\n- I think there could have been additional experiments to explore how CCE performs relative to baselines as different hyperparameters vary (e.g., relative size of vocabulary vs sequence length vs. hidden dim, sparsity of S, etc.).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tQmemldjqm", "reviewer": "Reviewer_2Zay", "summary": "The paper proposes a novel method of skipping most of the unneeded computation inside LM heads during training when using cross-entropy loss. It's key contributions are:\n- A memory efficient indexed matrix multiplication method, which employs sparsity to accelerate the computation.\n- A memory efficient linear-log-sum-exp method, which employs dynamic chunking to reduce memory requirements.\n- Gradient filtering, which further improves sparsity of the gradient computation.\n- Vocabulary sorting, which allows entire chunks to be skipped during computation.\n\nThe method is evaulated on both speed, memory usage and convergence, where it shows massive improvements in memory usage for computing losses compared to alternative methods, marginal improvements in speed and negligible degradation in convergence and training quality.", "questions": "- What are the theoretical justifications on why CCE might be much more computationally and memory efficient compared to traditional CE? For example, why can't you apply the same chunking strategies used in CCE for traditional CE? \n- What are the key insights that make CCE work? It seems to me currently that all of the contributions are mixed together, where CCE is an all or nothing monolithic algorithm.\n- How does CCE compare to CE on the CPU, or non-parallel cases? Are the improvements algorithmic or does it need to take advantage of GPU parallelization strategies?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- Well motivated problem. Reducing the memory footprint of LLMs during training is important.\n- Method generalizes beyond transformer LLMs.\n- Demonstrates convergence guarantees compared to cross entropy.\n- Extensive benchmark results.", "weaknesses": "- Preliminaries (section 3) does not adequately prepare the reader for the complexity of the notation in section 4. \n\n- Section 4 is particularly hard to understand if the reader does not have a deep understanding of GPU kernels and the architecture of modern LLMs. \n  - There is a lack of key insights, CCE seems like an arbitrary monolithic algorithm that came out of nowhere.\n\n  - Prehaps decoupling the theoretical reasoning from the actual GPU implementation could make the explanation clearer. For example, in line 201, it says \"section 4.2 describes how to compute the [...] operation efficiently\", but it is initially unclear to the reader why that operation might be efficient unless the reader can fully understand the intricacies of creating an efficient GPU kernel as described in section 4.2. Same goes for section 4.1 and 4.3.\n\n  - Otherwise, starting from an already efficient GPU implementation of standard CE and focusing on the steps needed to modify it into the CCE method could further improve readability and clarity.\n\n- A lack of ablation studies for the extensive modifications brought on by CCE\n  - Section 4.1, 4.2 and 4.3 makes a large number of significant assumptions, modifications and improvements to the traditional CE algorithm, it is not clear whether each modification is actually necessary or which are the most important ones.\n  - Unclear whether CCE's improvements is GPU dependent or not. Would it work in non-parallel cases such as on a single-threaded CPU?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iJeAtZ0mvs", "reviewer": "Reviewer_dN8N", "summary": "This paper proposes Cut Cross-Entropy(CCE) to reduce the memory consumption of Classifier Head and CrossEntropy Loss. They find that the vocabulary of LLMs continuously to grows, and under the gradient checkpointing setting, these part takes more than 50% of the memory consumption. CCE reduce the memory overhead by fusing the classifier head and the calculation of cross entropy loss into 1 kernel, and not materializing the intermediate logits in the forward process. In the backward pass they re-compute the intermediate values to avoid this additional memory overhead (which is quite similar to FlashAttention's design). They further propose to leverage the sparsity pattern in the gradient of classifier head to reduce the amount of computation. CCE reduce the memory overhead by 20x for \"Loss+Gradient\" part and their loss curve matches the BF16 training baseline.", "questions": "Please see weakness.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The problem this paper trying to solve is well-motivated.\n2. The solution to avoid the materialization of the large logit tensor is clear and easy to understand\n3. The CCE component is easy to deploy in realistic setting.\n4. The performance do not degrade (since the algorithm is nearly lossless considering the high sparsity level)\n5. The paper writing is very clear and easy to follow (represent C, E, and LSE in different colors)", "weaknesses": "1. \"The probabilities materialized by the cross-entropy layer account for 89% of the memory consumption of Gemma 2 for single sequence x with length N = 80000\". Can you provide details about how the 89% number is calculated and include a brief calculation or breakdown of the memory usage in the paper or appendix?\n2. Does this assumption still hold true when gradient checkpointing = False? I think most of the analysis in this paper is based on the assumption that gradient checkpointing = True. Include a subsection to discuss or analysis of how your method performs when gradient checkpointing is disabled is appreciated.\n3. Similar to 2, In Table 1, can you explain where 1477MB, 8000MB, 4000MB, and 24000MB come from? If I understand correctly, the logits.shape is (8192, 256000) in float32, which should take 8000MB memory in total.\n4. In Section 4.3, the Gradient filtering paragraph, \"If stored in bfloat16 with a 7-bit fraction, any value below 2^{-12} will likely be ignored due to truncation in the summation or rounding in the normalization.\" Can you explain this in detail? Providing a brief explanation of the numerical precision issues in bfloat16 and how they relate to the gradient filtering threshold is appreciated.\n\nOthers:\nWhat LSE stands for (Log-Sum-Exp) should be defined when it is on its first use.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hRqTGWmZgJ", "reviewer": "Reviewer_HR95", "summary": "This work proposes Cut Cross-Entropy (CCE) to address the massive memory footprint of the standard cross-entropy loss calculation in LLM training. CCE tiles and fuses the logits indexing and the matmul between tiles of `lm_head` and embeddings in the forward pass. In backward propagation, CCE introduces gradient filtering and vocabulary sorting to optimize memory access pattern with negligible approximation errors. CCE presents experiments showing the effectively reduced memory footprint and indistinguishable influence on fine-tuning.", "questions": "- Current Fig. 4 verifies that CCE has negligible influence on LLM fine-tuning. I am also curious about the impact of CCE on LLM pretraining. If an LLM is trained from scratch using CCE, how will the randomly initialized weights influence the gradient filtering and vocab sorting? \n\nFor other questions, please refer to my concerns in the Weakness section.", "limitations": "", "rating": 10, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- This paper identifies a new challenge brought by the large size of vocabulary of language models, especially LLMs, i.e., the massive memory footprint consumed by the cross-entropy loss computation.\n- The tricks of gradient filtering and vocabulary sorting in the proposed method are enlightening.\n- The author implemented CUDA kernels to support the algorithm and provides experiments to verify the reduced memory footprint and latency.", "weaknesses": "- The symbols in the derivation of CCE can be clearer. For example, the symbols in page 4 between line 186 and 215, such as $C^T$, $C_{x_i}^T$, $C_X$, and $(C^T E)_X$, look confusing at first glance. It may be helpful to have a table of symbol definition in the appendix.\n- Experiments on how the memory and latency of CCE kernel varies with the vocab size & model family can be added. \n  - Current Tab 1 presents the memory and latency results of Gemma-2-2B.  The vocab size of Gemma-2-2B is 256000, which is larger than other LLMs. For example, the vocab size is 128256 for Llama-3-8B/70B/405B, 32768 for mistral-7B-v0.3, and 32064 for Phi-3.5. If the size of `lm_head` is `(model_hidden_size, vocab_size)`. When the `model_hidden_size` increases, do we expect a diminishing benefit of CCE? The evaluation will be more comprehensive if the author could discuss:\n  - Compared to the baselines, how the memory and latency change if CCE is applied to Gemma-2-27B training (same vocab size as Tab 1, but larger model hidden size)\n  - Compared to the baselines, how the memory and latency change if CCE is applied to training of the models like Phi3.5-mini (smaller vocab size, similar model size).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)