详细内容或原文请订阅后点击阅览
将 LLM 内存减少 84%：深入研究融合内核
为什么你的最终 LLM 层出现 OOMing 以及如何使用自定义 Triton 内核修复它。将 LLM 内存削减 84%：深入研究融合内核的帖子首先出现在 Towards Data Science 上。
来源:走向数据科学或微调 LLM，您可能在最后一步遇到了障碍：交叉熵损失。
罪魁祸首是 Logit 瓶颈。为了预测下一个标记，我们将隐藏状态投影到巨大的词汇空间中。对于 Llama 3（128,256 个标记），仅权重矩阵就超过 5.25 亿个参数。虽然 bfloat16 中只有 ~1GB，但中间 Logit 张量才是真正的问题。对于大批量，仅计算单个标量损失就可以轻松超过 80GB VRAM。
优化这一层是 Unsloth 和 Liger-Kernel 等库实现如此大量内存减少的方式。在本文中，我们将在 Triton 中从头开始构建一个融合的线性 + 交叉熵内核。我们将推导数学并实现平铺的前向和后向传递，将峰值内存使用量削减 84%。
关于性能的说明：此实现主要是为了教育目的。我们通过使用全局原子操作来优先考虑数学清晰度和可读的 Triton 代码。虽然它解决了内存瓶颈，但匹配生产级速度将需要更加复杂的实现，这超出了本文的范围。
这篇文章是我的 Triton 系列的一部分。我们将使用之前介绍过的平铺和在线 softmax 等概念。如果这些听起来不熟悉，我建议先了解一下！
Logit 瓶颈
首先，让我们在 Logit 瓶颈上添加更多数字。我们考虑形状为 [NxD] 的输入矩阵 X、形状为 [DxV] 的权重矩阵 W 和形状为 [NxV] 的 Logit 矩阵 Y=X@W。在 LLM 的上下文中，N 是序列长度乘以批次大小（即批次中的标记总数），D 是隐藏状态的大小，V 是词汇表大小。
对于 Llama3 8B 模型，我们将拥有一个包含 8192 个标记的上下文窗口、一个具有 4096 个维度的隐藏状态以及一个包含 128,256 个标记的词汇表大小。使用适度的批量大小 8，我们得到 N = 8192x8 = 65,536。
一般解决这个问题的方法是使用：