ml-cross-entropy：优化大型词汇表语言模型中的损失计算

ml-cross-entropy：优化大型词汇表语言模型中的损失计算

ml-cross-entropy 项目地址: https://gitcode.com/gh_mirrors/ml/ml-cross-entropy

项目介绍

ml-cross-entropy 是一个开源软件项目，旨在解决大型词汇表语言模型（LLM）训练中损失计算的高内存消耗问题。该项目的核心是提出了一种称为 Cut Cross-Entropy（CCE）的方法，该方法在不将所有 token 的 logits 存储到全局内存的情况下计算交叉熵损失。

项目技术分析

传统的交叉熵损失计算会在输入 token 和词汇表项之间构建一个 logits 矩阵，这在大型模型中会导致巨大的内存占用。ml-cross-entropy 通过仅计算正确 token 的 logits 并实时评估所有 logits 的 log-sum-exp 来优化这一过程。此外，项目实现了一个自定义内核，用于在闪存中执行矩阵乘法和 logits 的 log-sum-exp 累减，从而使得全局内存消耗变得微不足道。

项目及技术应用场景

ml-cross-entropy 的主要应用场景是在大型语言模型训练中，特别是那些具有数百万或数十亿词汇表项的模型。例如，对于 Gemma 2（2B）模型，CCE 将损失计算的内存占用从 24 GB 降低到了 1 MB，极大地减少了训练时间所需的内存。此外，该项目还提供了与 PyTorch 的集成，以及针对不同场景的优化实现，如使用 torch.compile 的快速版本。

项目特点

1. 内存高效：通过避免构建完整的 logits 矩阵，CCE 大幅减少了内存消耗，使得大型模型训练更加可行。

2. 训练速度：CCE 的实时计算和自定义内核优化使得损失计算更加迅速，而不会牺牲训练速度或收敛性。

3. 灵活性：除了语言模型，CCE 还可以应用于其他场景，如视觉模型、对比损失等，它非常适合任何可以用类似格式表示损失计算的场景。

4. 易于集成：项目提供了与 PyTorch 和 Transformers 库的集成，使得用户可以轻松地将 CCE 应用于自己的模型。

推荐理由

ml-cross-entropy 项目为处理大型语言模型训练中的内存瓶颈提供了一种高效且创新的解决方案。以下是几个推荐使用此开源项目的理由：

提高资源利用率

对于资源受限的环境，CCE 的高内存效率能够使得模型训练成为可能，尤其是在有限的 GPU 内存条件下。

加速模型训练

通过优化损失计算过程，CCE 不仅减少了内存消耗，还提高了计算效率，从而缩短了模型训练的时间。

灵活的部署选项

CCE 提供了多种使用方式，包括针对不同硬件和场景的优化版本，使得用户可以根据自己的需求灵活选择。

强大的社区支持

作为开源项目，ml-cross-entropy 拥有一个活跃的社区，不断提供改进和更新，确保项目的持续发展。

ml-cross-entropy 项目的引入无疑为大型语言模型训练带来了重要的优化，它的高效性和灵活性使其成为当前和未来模型训练中的一个宝贵工具。无论是研究还是工业应用，该项目都值得推荐给所有关注语言模型性能和效率的开发者和研究人员。

ml-cross-entropy 项目地址: https://gitcode.com/gh_mirrors/ml/ml-cross-entropy