cut-cross-entropy

What it is and what it does

Cut Cross-Entropy (CCE) is a memory-efficient implementation of the cross-entropy loss computation for training large-vocabulary language models. During standard training, the cross-entropy layer materializes a full logit matrix covering all token-vocabulary pairs, consuming enormous amounts of GPU memory—often more than the rest of the model combined. CCE avoids this by computing logits on-the-fly in flash memory, computing only the logit for the correct token and evaluating log-sum-exp over all vocabulary items without materializing the full matrix.

The package provides a drop-in replacement function `linear_cross_entropy` that accepts embeddings, classifier weights, and labels, with optional support for token-shifting in causal language modeling. It includes custom Triton kernels for Ampere and newer GPUs, a torch.compile fallback for older hardware and non-Linux platforms, and direct integration patches for transformers models (Llama, Phi3, Mistral, Gemma2 families). According to the documentation, this reduces memory consumption from 24 GB to 1 MB for the loss computation on Gemma 2 (2B), with negligible impact on training speed.

Use it for:

• Training large language models on GPUs with limited memory by reducing classifier head memory footprint from tens of gigabytes to under a gigabyte.
• Fine-tuning transformer models (Llama, Phi3, Mistral, Gemma2) using the transformers library without modifying model code via cce_patch.
• Computing per-token loss and perplexity efficiently with reduction='none' for detailed loss analysis without materializing full logit matrices.
• Enabling training on older GPUs or non-Linux systems via torch.compile fallback when Triton kernels are unavailable.
• Reducing overall training-time memory consumption of the classifier head to fit larger models or batch sizes on fixed hardware.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Computes cross-entropy loss for large-vocabulary language models with minimal memory footprint by materializing logits selectively rather than globally, using custom kernels and torch.compile fallbacks.

Yes, with conditions. Install if you train large-vocabulary language models on memory-constrained GPUs and need dramatic memory reduction in the loss computation layer. The package is well-motivated by published research and offers both optimized Triton kernels and fallback implementations. However, verify the license terms before production use (license treatment is unclear), and be aware that maintenance is dormant—no active development is expected. Requires Python 3.10+, PyTorch 2.4+, and Triton 3.0+ on Ampere+ GPUs.

Install

cut-cross-entropy on PyPI

pip install cut-cross-entropy

uv add cut-cross-entropy

poetry add cut-cross-entropy

Installing cut-cross-entropy

Before you install

Low friction install with pure Python wheel distribution. Maintenance status is dormant (584 days since release), though the package received a recent update on 2025-01-07. Depends on torch and triton, both widely available.

License in practice

License treatment is unclear—no SPDX identifier or raw license text is available in the metadata. Verify the actual license terms before using in production or proprietary work.

Quickstart

pip install cut-cross-entropy

from cut_cross_entropy import linear_cross_entropy

embeddings = model.compute_embedding(inputs)
classifier = model.get_classifier_weights()
loss = linear_cross_entropy(embeddings, classifier, labels)

Requires Python 3.10+, PyTorch 2.4+, Triton 3.0+, and Ampere or newer GPU. Triton is not available on macOS; the package falls back to torch.compile on unsupported platforms.

Verify before relying

• Actual license terms and restrictions (license_treatment is unclear in metadata)
• Whether dormant maintenance status indicates ongoing support or abandonment
• Compatibility with PyTorch and Triton versions beyond those explicitly mentioned

Package facts

Evidence: cut_cross_entropy-25.1.1-py3-none-any.whl

Tags

More Artificial Intelligence packages

Further reading