--- id: cut-cross-entropy version: "25.1.1" license: unclear license_treatment: unclear maintenance: dormant --- # cut-cross-entropy — Code for cut cross entropy, a memory efficient implementation of linear-cross-entropy loss. License: unclear · Maintenance: dormant · Downloads: 609.7K/mo ## 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 above — 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 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 - License: not declared (unclear) - Python support: supports_current - Install friction: low - Maintenance: dormant - Downloads: 609.7K/month (top 15,000 on PyPI) - Known vulnerabilities: none known ## Tags memory efficient cross entropy loss, language model training optimization, large vocabulary loss computation, GPU memory reduction training, linear cross entropy implementation, memory-optimization, language-models, gpu-training [View on SkillFed](https://skillfed.io/packages/cut-cross-entropy) · [View on PyPI](https://pypi.org/project/cut-cross-entropy/)