Kati8/ICLR_2023_2025 huggingface dataset:
authors: Erik Wijmans
authors: Brody Huval
authors: Alexander Hertzberg
authors: Vladlen Koltun
authors: Philipp Kraehenbuehl
reviews: 
Comments: 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.
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.
- The tricks of gradient filtering and vocabulary sorting in the proposed method are enlightening.
- 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.
- Experiments on how the memory and latency of CCE kernel varies with the vocab size & model family can be added. 
  - 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:
  - 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)
  - 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).
Recommendation: 10
Reviewer confidence: 3
Soundness/correctness: 3
Clarity: 3
reviews: 
Comments: 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.
Strengths: 1. The problem this paper trying to solve is well-motivated.
2. The solution to avoid the materialization of the large logit tensor is clear and easy to understand
3. The CCE component is easy to deploy in realistic setting.
4. The performance do not degrade (since the algorithm is nearly lossless considering the high sparsity level)
5. 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?
2. 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.
3. 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.
4. 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.

Others:
What LSE stands for (Log-Sum-Exp) should be defined when it is on its first use.
Recommendation: 8
Reviewer confidence: 4
Soundness/correctness: 3
Clarity: 3
reviews: 
Comments: 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:
- A memory efficient indexed matrix multiplication method, which employs sparsity to accelerate the computation.
- A memory efficient linear-log-sum-exp method, which employs dynamic chunking to reduce memory requirements.
- Gradient filtering, which further improves sparsity of the gradient computation.
- Vocabulary sorting, which allows entire chunks to be skipped during computation.

The 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.
Strengths: - Well motivated problem. Reducing the memory footprint of LLMs during training is important.
- Method generalizes beyond transformer LLMs.
- Demonstrates convergence guarantees compared to cross entropy.
- Extensive benchmark results.
Weaknesses: - Preliminaries (section 3) does not adequately prepare the reader for the complexity of the notation in section 4. 

- 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. 
  - There is a lack of key insights, CCE seems like an arbitrary monolithic algorithm that came out of nowhere.

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

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

- A lack of ablation studies for the extensive modifications brought on by CCE
  - 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.
  - Unclear whether CCE's improvements is GPU dependent or not. Would it work in non-parallel cases such as on a single-threaded CPU?
Recommendation: 6
Reviewer confidence: 2
Soundness/correctness: 3
Clarity: 2
reviews: 
Comments: 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).

Putting 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).
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.
- The algorithm is clever and elegant, taking inspiration from FlashAttention, which avoids materializing full attention score matrix during attention computation.
- 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).
- 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.).
Recommendation: 10
Reviewer confidence: 4
Soundness/correctness: 4
Clarity: 3
reviews: 
Comments: The paper introduces Cut Cross-Entropy (CCE), a method designed to mitigate the substantial memory consumption associated with standard cross-entropy loss computation during LLM training. The approach employs tiling and fusion strategies for logits indexing and matrix multiplication operations between `lm_head` tiles and embeddings in the forward pass. During backpropagation, CCE leverages gradient filtering combined with vocabulary sorting to enhance memory access efficiency while maintaining negligible approximation error. The work includes empirical validation demonstrating significant memory footprint reduction and minimal performance degradation in fine-tuning scenarios.
Strengths: - The paper successfully identifies and addresses an important computational bottleneck in LLM training: the substantial memory overhead incurred by cross-entropy loss computation at the vocabulary scale.
- The proposed techniques of gradient filtering and vocabulary sorting provide valuable insights and appear novel in their application.
- The authors developed custom CUDA kernel implementations and conducted experiments that convincingly demonstrate both memory footprint and latency improvements.
Weaknesses: - The mathematical notation used throughout the derivation sections could benefit from enhanced clarity. Specifically, the notation introduced on page 4 (lines 186-215), including $C^T$, $C_{x_i}^T$, $C_X$, and $(C^T E)_X$, may be difficult to parse on initial reading. Including a comprehensive symbol reference guide in the appendix would improve accessibility.
- The experimental evaluation would be strengthened by analyzing how CCE's memory and latency performance scales across different vocabulary sizes and model architectures. Currently, Table 1 focuses exclusively on Gemma-2-2B with a vocabulary size of 256000, which is notably larger than many contemporary LLMs (e.g., Llama-3 variants at 128256, Mistral-7B-v0.3 at 32768, Phi-3.5 at 32064). Given that `lm_head` dimensions are `(model_hidden_size, vocab_size)`, the relative benefits of CCE may vary as hidden dimensions increase. The evaluation would gain substantially from investigating: 
  - Performance characteristics when applying CCE to Gemma-2-27B (maintaining the same vocabulary size as Table 1 but with increased hidden dimensions)
  - Performance characteristics when applying CCE to models with smaller vocabularies, such as Phi-3.5-mini (comparable model size but reduced vocabulary)
Recommendation: 10
Reviewer confidence: 3
Soundness/correctness: 3
Clarity: 3