Chunked and Fused Cross-Entropy: How Online Logit Tiling Slashes Large-Vocabulary VRAM Bottlenecks in LLM Training

As frontier large language models have scaled, tokenizer vocabularies have expanded substantially. Where early architectures such as LLaMA and Mistral relied on 32,000 subword tokens, contemporary models routinely employ vocabularies of 128,256 tokens (Llama 3), 152,064 tokens (Qwen 2.5), and 256,000 tokens (Gemma 2). Larger vocabularies compress text more densely, improve multilingual representation efficiency, and lower the generation step count per prompt.

However, this scaling creates a severe memory bottleneck in the final layer of the network during pre-training and fine-tuning. The unembedding projection and cross-entropy loss computation scale linearly with vocabulary size $V$, sequence length $S$, and batch size $B$. In standard training pipelines, computing cross-entropy loss requires materializing an intermediate logits tensor of shape $[B \times S, V]$ directly in GPU High Bandwidth Memory (HBM). For models with 128k to 256k tokens, this single tensor can consume more VRAM than all preceding transformer layers combined.

Recent advances in kernel design—most notably LinkedIn's Liger Kernel and Apple's Cut Cross-Entropy (CCE)—resolve this bottleneck through fused operations and online logit tiling. By computing matrix projections, log-sum-exp reductions, and backpropagation gradients within localized SRAM tiles, chunked cross-entropy eliminates global logit materialization and reduces peak loss memory by up to 85%.

The Vocabulary Memory Wall in Modern LLMs

During the forward pass of a decoder-only transformer, hidden representations $X \in \mathbb{R}^{N \times H}$ (where $N = B \times S$ denotes the total number of flattened tokens across the batch, and $H$ is the model's hidden dimension) are projected onto the output vocabulary space via the unembedding weight matrix $W_{\text{vocab}} \in \mathbb{R}^{V \times H}$:

$$Z = X W_{\text{vocab}}^T \in \mathbb{R}^{N \times V}$$

Once logits $Z$ are calculated, standard cross-entropy loss computes the negative log-likelihood of the ground-truth target tokens $Y \in \{1, \dots, V\}^N$:

$$\mathcal{L} = \frac{1}{N} \sum_{i=1}^N \ell_i, \quad \text{where} \quad \ell_i = -\log P(y_i \mid x_i) = -z_{i, y_i} + \log \left( \sum_{j=1}^V \exp(z_{i, j}) \right)$$

The memory footprint of this operation in standard PyTorch (torch.nn.CrossEntropyLoss) is governed by three separate global HBM allocations:

1. Forward Logits Tensor ($Z$): Size $N \times V \times \text{sizeof}(\text{dtype})$.
2. Softmax Normalization / Log-Sum-Exp Tensor: Storing intermediate probabilities or reduction buffers.
3. Backward Logit Gradients ($\nabla_Z$): Size $N \times V \times \text{sizeof}(\text{dtype})$, where $\nabla_{z_{i, j}} = P(j \mid x_i) - \mathbb{I}(j = y_i)$.

Standard PyTorch Memory Footprint:
Total HBM Allocated = 2 × (B × S × V × bytes_per_element) + Overhead

For a training setup processing a batch of $B = 4$, sequence length $S = 8,192$ ($N = 32,768$ tokens), and a 16-bit floating point format (bfloat16, 2 bytes per element):

• 32k Vocabulary (LLaMA 2): $32,768 \times 32,000 \times 2 \approx 2.10\text{ GB}$
• 128k Vocabulary (Llama 3): $32,768 \times 128,256 \times 2 \approx 8.41\text{ GB}$
• 256k Vocabulary (Gemma 2): $32,768 \times 256,000 \times 2 \approx 16.78\text{ GB}$

When factoring in both the forward logits tensor and the backward gradient tensor, Gemma 2 allocates over 33.5 GB of GPU memory solely to compute cross-entropy on a single GPU. As documented by Apple ML Research, the log-probabilities materialized by standard cross-entropy account for roughly 40% of peak memory in Phi-3.5-mini, 65% in Llama 3-8B, and up to 89% in Gemma 2-2B during standard training runs.

The Standard Execution Pipeline vs. Fused Tiling

In standard unfused pipelines, the linear projection and cross-entropy loss are distinct computational layers executed as separate CUDA kernel launches:

Unfused Execution Flow:
[Hidden States X] ---> [GEMM Kernel] ---> [Write Logits Z to HBM (8-16 GB)]
                                                      |
                                                      v
[Loss Value L]    <--- [CrossEntropy Kernel] <--------+
                                                      |
                                                      v
[Write Grad dZ to HBM (8-16 GB)] <--------------------+
       |
       v
[GEMM Backward]  ---> [Compute dX and dW]

This separation requires writing the entire $N \times V$ tensor to HBM only for the subsequent cross-entropy kernel to read it back immediately, creating an intense memory bandwidth tax in addition to capacity exhaustion.

Chunked and fused cross-entropy refactors this execution graph by merging the final linear projection, the log-sum-exp reduction, and the backward gradient derivation into a single streaming workflow.

Mathematical Mechanics of Chunked Online Cross-Entropy

Chunked cross-entropy decomposes the token sequence dimension $N$ into small micro-chunks of size $C$ (typically $C \in [256, 1024]$ tokens), such that $N = K \times C$ across $K$ chunks.

For each micro-chunk $k \in \{1, \dots, K\}$:

1. Chunked Linear Projection in SRAM

The model takes the chunked hidden state slice $X_k \in \mathbb{R}^{C \times H}$ and computes the corresponding logit slice $Z_k \in \mathbb{R}^{C \times V}$ using on-chip matrix multiplication:

$$Z_k = X_k W_{\text{vocab}}^T$$

Crucially, $Z_k$ is held inside the GPU's streaming multiprocessor (SM) SRAM and L2 cache hierarchies or allocated in a small temporary scratchpad buffer of size $C \times V$ rather than allocating the full $N \times V$ matrix.

2. Numerically Stable Online Log-Sum-Exp

For each token $i$ within chunk $k$, the kernel computes the partition function using a running online max and sum to prevent floating-point overflow:

$$m_i = \max_{j \in [1, V]} z_{i, j}$$

$$d_i = \sum_{j=1}^V \exp(z_{i, j} - m_i)$$

$$\text{LSE}(z_i) = m_i + \log(d_i)$$

The scalar loss for token $i$ is evaluated directly:

$$\ell_i = -z_{i, y_i} + m_i + \log(d_i)$$

The running total loss $\mathcal{L} = \sum_{i=1}^N \ell_i$ is accumulated into a global scalar.

3. Immediate In-Place Gradient Generation

Because the gradient of the cross-entropy loss with respect to logit $z_{i, j}$ is simply the softmax probability minus the indicator function:

$$\frac{\partial \mathcal{L}}{\partial z_{i, j}} = \frac{1}{N} \left( \frac{\exp(z_{i, j} - m_i)}{d_i} - \mathbb{I}(j = y_i) \right)$$

the kernel computes $\nabla_{Z_k} \in \mathbb{R}^{C \times V}$ immediately while $Z_k$, $m_i$, and $d_i$ are still active in fast memory.

4. Backward Hidden and Weight Accumulation

The gradients with respect to the input hidden states $X_k$ and the unembedding weights $W_{\text{vocab}}$ are computed before moving to the next chunk:

$$\nabla_{X_k} = \nabla_{Z_k} W_{\text{vocab}} \in \mathbb{R}^{C \times H}$$

$$\nabla_{W_{\text{vocab}}} \mathrel{+}= \nabla_{Z_k}^T X_k \in \mathbb{R}^{V \times H}$$

Once $\nabla_{X_k}$ is written to the output buffer and $\nabla_{W_{\text{vocab}}}$ is updated, the intermediate logit slice $Z_k$ and its gradient $\nabla_{Z_k}$ are discarded. The memory allocated for the next chunk simply reuses the same $C \times V$ scratchpad buffer.

Peak Logit Memory Complexity:
Standard Cross-Entropy: O(B × S × V)
Chunked Cross-Entropy:  O(C × V) where C << (B × S)

When $B \times S = 32,768$ and $C = 512$, the peak memory consumed by logits and logit gradients is reduced by a factor of $64\times$.

Architectural Comparison: Liger Kernel vs. Cut Cross-Entropy

Two prominent open-source paradigms have emerged for fused cross-entropy optimization:

+--------------------------+-----------------------------------+-----------------------------------+
| Feature                  | Liger Kernel (FLCE)               | Cut Cross-Entropy (CCE)           |
+--------------------------+-----------------------------------+-----------------------------------+
| Primary Strategy         | Input Token Chunking              | Gradient Filtering & Target-Only  |
| Triton Implementation    | Fused Linear + Cross-Entropy      | Tiled LogSumExp & Selective Grad  |
| Memory Reduction         | 40% - 60% Overall Training Peak   | Up to 85% Loss Layer Reduction    |
| Vocabulary Parallelism   | Chunk-by-chunk full vocab GEMM    | Tiled vocabulary dot products     |
| Numerical Precision      | Exact mathematical equivalent     | Optional dynamic thresholding     |
| Framework Integration    | Hugging Face, torchtune, Axolotl  | PyTorch native, transformers hook |
+--------------------------+-----------------------------------+-----------------------------------+

LinkedIn Liger Kernel

The Liger Kernel suite implements LigerFusedLinearCrossEntropyLoss in OpenAI Triton. It flattens 3D hidden states into 2D tensors, chunks the token sequence along the batch/time dimension, and performs an in-place overwrite of logits with gradients during the backward pass. This design minimizes kernel launch overhead while fitting cleanly into standard distributed training frameworks such as PyTorch Fully Sharded Data Parallel (FSDP) and DeepSpeed ZeRO.

Apple Cut Cross-Entropy (CCE)

Apple's CCE takes optimization further by leveraging an information-theoretic property of softmax distributions: for large vocabularies, the vast majority of non-target tokens receive negligible softmax probability ($P(j \mid x_i) < \epsilon$). In 16-bit precision (bfloat16), values below $2^{-12}$ round to zero during normalization or contribute negligibly to gradient updates.

CCE computes only the exact target logits $z_{i, y_i}$ and performs a streaming reduction for the denominator $\sum \exp(z_i)$ in fast SRAM tiles. During the backward pass, CCE applies gradient filtering to skip calculating and materializing gradients for token dimensions that do not meaningfully impact weight updates.

VRAM Footprint Across Model Families

The quantitative impact of chunked and fused loss kernels becomes more pronounced as vocabulary sizes scale:

+-----------------------+------------+------------+--------------------+-------------------+
| Model Architecture    | Vocab Size | Hidden Dim | Standard CE Peak   | Fused CE Peak     |
| (Context: 8k, B: 4)   | (V)        | (H)        | Logit VRAM (GB)    | Logit VRAM (GB)   |
+-----------------------+------------+------------+--------------------+-------------------+
| LLaMA 2 (7B)          | 32,000     | 4,096      | 4.19 GB            | 0.07 GB (512-blk) |
| Mistral NeMo (12B)    | 128,000    | 5,120      | 16.78 GB           | 0.26 GB (512-blk) |
| Llama 3.1 (8B)        | 128,256    | 4,096      | 16.81 GB           | 0.26 GB (512-blk) |
| Qwen 2.5 (7B)         | 152,064    | 3,584      | 19.93 GB           | 0.31 GB (512-blk) |
| Gemma 2 (9B)          | 256,000    | 3,584      | 33.55 GB           | 0.52 GB (512-blk) |
+-----------------------+------------+------------+--------------------+-------------------+

Note: Peak VRAM figures reflect the combined allocation of forward logits and backward gradient buffers at float16/bfloat16 precision.

By eliminating tens of gigabytes of ephemeral logit allocations, engineering teams can reallocate GPU memory toward:

1. Longer Context Windows: Expanding fine-tuning contexts from 8k to 32k or 128k tokens without encountering out-of-memory (OOM) errors.
2. Larger Micro-Batch Sizes: Increasing per-GPU batch sizes to saturate tensor core compute capacity and accelerate wall-clock training throughput by 20% to 35%.
3. Hardware Accessibility: Enabling full-parameter fine-tuning of 7B-8B models on single 24GB or 48GB workstation GPUs (such as RTX 4090 or A6000 Ada) that previously failed during the final loss step.

Numerical Considerations and Practical Trade-Offs

While chunked cross-entropy produces mathematical results identical to standard implementations, practitioners must account for several implementation details:

1. Floating-Point Accumulation Order: Summing loss values across sequential micro-chunks alters the accumulation order of floating-point numbers compared to a monolithic reduction. Small numerical discrepancies (typically between the 5th and 7th decimal places) are expected and do not degrade model convergence.
2. Loss Masking and Padding: In instruction tuning and multi-turn conversations, user prompt tokens are masked out from loss computation (label = -100). Optimized chunked kernels inspect label buffers to skip forward projections and backward computations for entire chunks composed exclusively of padding tokens.
3. Regularization Integration: Auxiliary objectives such as label smoothing or $z$-loss regularization ($\mathcal{L}_z = \alpha \log^2 \sum \exp(z_i)$) can be computed directly within the SRAM tile during the log-sum-exp step without incurring additional memory passes.

As modern foundation models continue expanding tokenizer vocabularies to support global languages and specialized modalities, chunked and fused cross-entropy has shifted from an optional optimization to an essential standard in LLM training infrastructure.

Sources

• Liger Kernel: Efficient Triton Kernels for LLM Training (arXiv:2410.10989)
• Cut Your Losses in Large-Vocabulary Language Models - Apple Machine Learning Research (arXiv:2411.09009)
• GitHub: linkedin/Liger-Kernel
• GitHub: apple/ml-cross-entropy
• PyTorch Blog: Optimizing torchtune Performance with Liger Kernel