进阶 unsloth.ai 2026-10-07 22:27:00 · 6 阅读
第13章 使用 Unsloth 微调 Gemma 2：实现 2 倍速度与显存优化
unsloth
下载
☰
下载
博客 2倍速度的 Gemma 2 微调 + 节省 63% 显存 2024 年 7 月 3 日 • 作者 Daniel & Michael 2024 年 7 月 3 日 • 作者 Daniel & Michael Gemma 2 (9B) 1xL4 24GB 速度快 200% Gemma 2 (9B) 1xL4 24GB 显存减 63% Gemma 2 (27B) 1xA100 80GB 速度快 190% Gemma 2 (27B) 1xA100 80GB 显存减 53% 现在，你可以比 Flash Attention 2 (FA2) + Hugging Face (HF) 快 2 倍微调 Google 最新的 Gemma 2 (9B) 模型，同时减少 63.2% 的内存占用。对于 Gemma 2 (27B)，速度提升 1.9 倍，显存减少 51%。我们还与出色的 Gemma 和 HF 团队合作修复了一些小 bug，详情见下文。
使用 Unsloth，你可以在 40GB GPU 上通过 QLoRA 微调 Gemma 2 (27B)，支持 9.7K 上下文长度（最大支持 8K），而 HF+FA2 仅支持 3K 长度。对于 9B 模型，Unsloth 在 24GB 显卡上支持 11K 上下文长度，而 HF+FA2 仅支持 2.6K。这意味着使用 Unsloth 的上下文长度可以是 4-5 倍！
我们上传了 Gemma 2 (9B) 的 Colab Notebook，以及预量化为 4bit 的模型，支持 4 倍速下载，其中包含 Instruct 和 Base 版本的 Gemma 2 4bit 模型。
我们还升级了对微软最新更新 Phi-3 mini 的支持。请确保使用我们最新的 Phi-3 mini notebook 和模型。 Gemma 2 基准测试 模型 显存 🦥Unsloth 速度 🦥 VRAM 减少幅度 🦥 更长上下文 🤗Hugging Face+FA2 Gemma 2 (9B) 80GB 2x 63.2% 长 4-5 倍 1x Gemma 2 (27B) 80GB 1.9x 51% 长 3 倍 1x 我们使用 Alpaca 数据集进行测试，批次大小为 2，梯度累积步数为 4，秩为 32，并在所有线性层（q, k, v, o, gate, up, down）上应用 QLoRA。🌠 上下文长度提升 4-5 倍 通过采用 Unsloth 的 offloaded gradient checkpointing（卸载梯度检查点），训练速度仅慢约 2%，但显存占用降低 30%！现在，Unsloth 默认启用此功能，设置 `use_gradient_checkpointing = "unsloth"`。
有趣的是，由于注意力软截断（attention softcapping），Flash Attention 目前不支持该功能，导致注意力矩阵实际上使用 O(N^2) 内存，而 FA2 允许在创建 QK^T 矩阵时按需使用内存。这意味着在 Flash Attention 支持软截断之前，Gemma 的内存使用量将随序列长度呈二次方增长。
Unsloth 允许 27B 模型的上下文长度增加 3 倍，9B 模型增加 4-5 倍！ Gemma 2 (27B) 最大上下文长度 GPU 显存 Unsloth （新增）Hugging Face+FA216 GB285OOM24 GB3,43667540 GB9,7373,11648 GB12,8884,33780 GB25,4919,221在我们的所有实验中，我们都使用 rank 为 32 的 QLoRA，并将 LoRA 适配器应用到所有线性层（q、k、v、o、gate、up、down）上。batch size 设为 1，并通过重复数据使其填满最大上下文窗口。 Gemma 2 (9B) 最大上下文长度GPU 显存UnslothHugging Face+FA28 GB74OOM12 GB2,84228416 GB5,6091,07024 GB11,1452,64240 GB22,2155,78748 GB27,7507,35980 GB49,89113,649我们同样用重复数据的方式测试 9b 模型，以填满最大上下文长度。 👨💻 Softcapping 函数的导数 我们需要在 Gemma 2 中实现 softcapping 机制。我们利用 Desmos 的数值微分和积分功能，验证了导数的正确性，并借助一些三角恒等式推导出 softcapping 函数的梯度。 我们还通过将 softcapping 机制融合到交叉熵损失计算中，减少了 500MB 甚至更多的显存占用，这里同样需要用到导数。通过融合处理，我们无需再保留 softcapping 操作之前的 logits 副本，从而降低显存用量。我们通过确认损失值是否一致，验证了梯度的准确性。 📈 Softcapping 调研 对于 9b 模型，我们发现必须为 lm head 的 logits 开启 softcapping，attention 的 softcapping 也普遍建议开启。如果 lm head 不开启，训练损失就会不正确。
对于 27b 模型，我们测试了全部 4 种组合（attention softcapping、logit softcapping、两者都开、只开一个、都不开）。结果显示 27b 模型对 softcapping 敏感得多，attention 和 lm head logits 的 softcapping 都必须开启。这不像 9b 模型那样有可选择性——在 9b 上关闭 attention softcapping 只会对损失造成轻微影响。
总体而言，我们建议对 9b 和 27b 模型开启注意力机制和 lm head logits 的 softcapping。这意味着将无法再使用 Flash Attention。对于 Unsloth，我们利用 torch.compile 将 softcapping 与注意力矩阵计算进行融合。🐛 Gemma 2 Pytorch 修复 我们还为官方 Gemma 2 Pytorch 仓库提供了一些修复！见 Pull Request 67。我们演示了必须遵循之前的 Gemma 错误修复方法，即需要仔细对代码的特定区域进行降精度（downcast）和升精度（upcast）操作。我们提供了 2 个修复！这主要是因为 GPU 混合精度训练与 TPU 混合精度训练存在一定差异。
我们还发现 Gemma 团队将 tokenizer 词表填充了 128 个 token，这相当酷！这很可能是为了让训练速度稍快一些。🪟 Phi-3 mini 更新 得益于 Microsoft，Phi-3 mini 迎来了出色的新更新，因此我们已更新所有基础设施以支持该版本。如果你仍希望使用旧版本，我们也上传了 Phi-3 的原始模型。
在新标签页中打开下图可查看所有新的基准测试：🌎 2024 AI 世界博览会 非常感谢上周出席我们在 AI Engineer World's Fair 上两场活动的各位。我们结识了许多杰出的人才，并衷心感谢 Swyx 及其出色的团队组织此次盛会。我们举办了一场 3 小时的工作坊，观众极具互动性，令人印象深刻，希望部分参与者拿走了我们的贴纸。
在 YouTube 上观看我们 20 分钟的闪电演讲，主题是修复开源模型中的 bug。💕 感谢！ 欢迎通过我们的 Ko-fi 捐赠页面支持我们。特别感谢以下新支持者：Creivailty, kearm, MrDragonFox, Sebastien, Fimbul, Jeff, Steffen, Andrew & Shailendra！🙏
一如既往地，欢迎加入我们的 Discord 服务器寻求帮助或表达支持！你也可以在 Twitter 和 Substack 上关注我们要。 谢谢阅读！ Daniel & Michael Han 🦥
2024 年 7 月 3 日 MultiGPU 即将上线，免费开始使用 加入我们的 Discord
使用 Unsloth，你可以在 40GB GPU 上通过 QLoRA 微调 Gemma 2 (27B)，支持 9.7K 上下文长度（最大支持 8K），而 HF+FA2 仅支持 3K 长度。对于 9B 模型，Unsloth 在 24GB 显卡上支持 11K 上下文长度，而 HF+FA2 仅支持 2.6K。这意味着使用 Unsloth 的上下文长度可以是 4-5 倍！
我们上传了 Gemma 2 (9B) 的 Colab Notebook，以及预量化为 4bit 的模型，支持 4 倍速下载，其中包含 Instruct 和 Base 版本的 Gemma 2 4bit 模型。
我们还升级了对微软最新更新 Phi-3 mini 的支持。请确保使用我们最新的 Phi-3 mini notebook 和模型。 Gemma 2 基准测试 模型 显存 🦥Unsloth 速度 🦥 VRAM 减少幅度 🦥 更长上下文 🤗Hugging Face+FA2 Gemma 2 (9B) 80GB 2x 63.2% 长 4-5 倍 1x Gemma 2 (27B) 80GB 1.9x 51% 长 3 倍 1x 我们使用 Alpaca 数据集进行测试，批次大小为 2，梯度累积步数为 4，秩为 32，并在所有线性层（q, k, v, o, gate, up, down）上应用 QLoRA。🌠 上下文长度提升 4-5 倍 通过采用 Unsloth 的 offloaded gradient checkpointing（卸载梯度检查点），训练速度仅慢约 2%，但显存占用降低 30%！现在，Unsloth 默认启用此功能，设置 `use_gradient_checkpointing = "unsloth"`。
有趣的是，由于注意力软截断（attention softcapping），Flash Attention 目前不支持该功能，导致注意力矩阵实际上使用 O(N^2) 内存，而 FA2 允许在创建 QK^T 矩阵时按需使用内存。这意味着在 Flash Attention 支持软截断之前，Gemma 的内存使用量将随序列长度呈二次方增长。
Unsloth 允许 27B 模型的上下文长度增加 3 倍，9B 模型增加 4-5 倍！ Gemma 2 (27B) 最大上下文长度 GPU 显存 Unsloth （新增）Hugging Face+FA216 GB285OOM24 GB3,43667540 GB9,7373,11648 GB12,8884,33780 GB25,4919,221在我们的所有实验中，我们都使用 rank 为 32 的 QLoRA，并将 LoRA 适配器应用到所有线性层（q、k、v、o、gate、up、down）上。batch size 设为 1，并通过重复数据使其填满最大上下文窗口。 Gemma 2 (9B) 最大上下文长度GPU 显存UnslothHugging Face+FA28 GB74OOM12 GB2,84228416 GB5,6091,07024 GB11,1452,64240 GB22,2155,78748 GB27,7507,35980 GB49,89113,649我们同样用重复数据的方式测试 9b 模型，以填满最大上下文长度。 👨💻 Softcapping 函数的导数 我们需要在 Gemma 2 中实现 softcapping 机制。我们利用 Desmos 的数值微分和积分功能，验证了导数的正确性，并借助一些三角恒等式推导出 softcapping 函数的梯度。 我们还通过将 softcapping 机制融合到交叉熵损失计算中，减少了 500MB 甚至更多的显存占用，这里同样需要用到导数。通过融合处理，我们无需再保留 softcapping 操作之前的 logits 副本，从而降低显存用量。我们通过确认损失值是否一致，验证了梯度的准确性。 📈 Softcapping 调研 对于 9b 模型，我们发现必须为 lm head 的 logits 开启 softcapping，attention 的 softcapping 也普遍建议开启。如果 lm head 不开启，训练损失就会不正确。
对于 27b 模型，我们测试了全部 4 种组合（attention softcapping、logit softcapping、两者都开、只开一个、都不开）。结果显示 27b 模型对 softcapping 敏感得多，attention 和 lm head logits 的 softcapping 都必须开启。这不像 9b 模型那样有可选择性——在 9b 上关闭 attention softcapping 只会对损失造成轻微影响。
总体而言，我们建议对 9b 和 27b 模型开启注意力机制和 lm head logits 的 softcapping。这意味着将无法再使用 Flash Attention。对于 Unsloth，我们利用 torch.compile 将 softcapping 与注意力矩阵计算进行融合。🐛 Gemma 2 Pytorch 修复 我们还为官方 Gemma 2 Pytorch 仓库提供了一些修复！见 Pull Request 67。我们演示了必须遵循之前的 Gemma 错误修复方法，即需要仔细对代码的特定区域进行降精度（downcast）和升精度（upcast）操作。我们提供了 2 个修复！这主要是因为 GPU 混合精度训练与 TPU 混合精度训练存在一定差异。
我们还发现 Gemma 团队将 tokenizer 词表填充了 128 个 token，这相当酷！这很可能是为了让训练速度稍快一些。🪟 Phi-3 mini 更新 得益于 Microsoft，Phi-3 mini 迎来了出色的新更新，因此我们已更新所有基础设施以支持该版本。如果你仍希望使用旧版本，我们也上传了 Phi-3 的原始模型。
在新标签页中打开下图可查看所有新的基准测试：🌎 2024 AI 世界博览会 非常感谢上周出席我们在 AI Engineer World's Fair 上两场活动的各位。我们结识了许多杰出的人才，并衷心感谢 Swyx 及其出色的团队组织此次盛会。我们举办了一场 3 小时的工作坊，观众极具互动性，令人印象深刻，希望部分参与者拿走了我们的贴纸。
在 YouTube 上观看我们 20 分钟的闪电演讲，主题是修复开源模型中的 bug。💕 感谢！ 欢迎通过我们的 Ko-fi 捐赠页面支持我们。特别感谢以下新支持者：Creivailty, kearm, MrDragonFox, Sebastien, Fimbul, Jeff, Steffen, Andrew & Shailendra！🙏
一如既往地，欢迎加入我们的 Discord 服务器寻求帮助或表达支持！你也可以在 Twitter 和 Substack 上关注我们要。 谢谢阅读！ Daniel & Michael Han 🦥
2024 年 7 月 3 日 MultiGPU 即将上线，免费开始使用 加入我们的 Discord
