第9章 Unsloth 梯度检查点：支持 4 倍长上下文窗口的 LLM 微调算法

unsloth 下载 博客：Unsloth 梯度检查点 - 4 倍长的上下文窗口 2024 年 4 月 9 日 • 作者：Daniel & Michael 我们很高兴推出“Unsloth 梯度检查点”，一种新算法，可实现超长上下文窗口的 LLM 微调。在 NVIDIA H100 80GB GPU 上，它支持长达 228K tokens 的上下文——比 Hugging Face (HF) 搭配 Flash Attention 2 (FA2) 的 48K 提升了 4 倍。在 RTX 4090 24GB GPU 上，Unsloth 可实现 56K tokens 的上下文长度，是 HF+FA2（14K tokens）的 4 倍。 Unsloth 梯度检查点在额外增加 1.9% 时间开销的前提下，将内存占用进一步降低 30%，使基于 transformers 的长上下文 LLM 微调效率大幅提升。该算法适用于所有使用梯度检查点的模型架构（例如 stable diffusion、Mamba 等）。 我们提供了一个 Colab notebook，用于在 Tesla T4 上训练 16K 序列长度的模型，使用 Mistral 7b v2（基于长序列训练的新版 Mistral 7b 模型）并采用 ChatML 格式，点击这里即可访问！ 下表展示了在使用 4bit QLoRA、批量大小为 1、lora 秩为 32 以及 adamw_8bit 优化器的情况下，Mistral 7b 在各种流行 GPU 上的最大上下文长度计算结果。 这些数值是基于实验结果外推的，因此考虑到 VRAM 碎片化，建议将上下文窗口设置比数值低 10%。另一个好处是，现在你可以将批量大小提高 1.7 倍！ P.S. 别忘了在 Github 上 ⭐ 关注我们要，并加入我们的 Discord 服务器 ❤️ 最大上下文窗口基准测试 | GPU | VRAM | Unsloth
(New) | Unsloth
(Old) | Hugging Face+FA2 | | :--- | :--- | :--- | :--- | :--- | | RTX 4060 | 8 GB | 7,340 | 3,716 | 1,696 | | RTX 4070 | 12 GB | 19,610 | 11,055 | 4,797 | | RTX 4080 | 16 GB | 31,880 | 18,394 | 7,898 | | RTX 4090 | 24 GB | 56,420 | 33,073 | 14,099 | | A100 | 40 GB | 105,500 | 62,431 | 26,502 | | A6000 | 48 GB | 130,040 | 77,110 | 32,704 | | H100 | 80 GB | 228,199 | 135,826 | 57,510 | 💡 技术细节 我们是怎样实现的？我们利用纯 Pytorch 代码（仅需 20 行）自研的梯度检查点机制，异步地将激活值卸载到系统 RAM。令人惊讶的是，这仅带来了约 1.9% 的额外时间开销，主要得益于我们在前向和反向传播中巧妙地使用非阻塞调用，隐藏了 GPU 到 CPU 的通信延迟。 起初，我们担心卸载操作会导致性能严重下降，但实验结果证明并非如此。原生 Hugging Face 若不使用 Flash Attention 2，其 VRAM 使用量会呈二次方增长，在 L4 GPU 上运行 Mistral 7b 时仅能处理 5K 上下文窗口。使用 FA2 后，该数值可提升至约 14K（2.8 倍）。 由于 FA2 不会生成完整的注意力矩阵，VRAM 现在实现了线性增长。Unsloth 旧版支持 33K 上下文窗口（2.4 倍），而新版支持 56K 上下文窗口（1.7 倍）。这意味着在 Mistral 7b 上，Unsloth 支持的上下文窗口长度是 HF+FA2 的 4 倍，更是原生 HF 的 11.2 倍。 若要启用长上下文窗口微调，请先设置最大序列长度。你可以参考上文关于 Mistral 7b 的表格，并将该数值降低 10% 左右。注意在本地机器上更新 Unsloth，而在 Colab / Kaggle 上则无需更新！🖥️ 2.4 倍速度的 CodeGemma 7b + 节省 71% 内存 Google 今天发布了他们的 CodeGemma 新模型！CodeGemma 基于 Gemma 2b 和 7b 构建，是在 Gemma 7b 的 6 万亿词元数据集基础上，额外微调了 500B 代码数据。

我们上传了 CodeGemma 的 4bit 预量化模型，下载速度快 4 倍，还因减少 GPU 内存碎片而节省 1GB 显存。模型已上传到我们的 HF 页面，也可以使用我们的 CodeGemma 7b + ChatML 模板 Colab notebook。 Gemma 7b / 1x A100：速度提升 243%，显存节省 71% Gemma 2b / 1x A100：速度提升 200%，显存节省 68% ❤️‍🩹 自修复 Tokenizer 我们实现了把慢速 HF tokenizer 智能地、实时地转换成快速版本，并在加载 tokenizer 时自动修复一些遗留的错误 token。这有什么用呢？ 1. 像 Starling 或 CodeLlama 那样有问题的 tokenizer 可以"自愈"后正常工作。不修复的话可能导致越界内存访问。 2. 无需手动编辑 tokenizer 文件来支持 ChatML 格式，Unsloth 会自动修改 sentencepiece 的 tokenizer.model 及其他文件。 3. 有时模型上传者会要求使用慢速 tokenizer，因为快速 tokenizer（HF 的 Rust 版本）结果不正确。我们会尝试转换成快速版本，并确认其分词结果正确。 提醒一下，Unsloth 无缝支持所有 Chat 模板（Vicuna、ChatML、Zephyr 等），详见这里，或使用我们的 ChatML notebook。 🧶 RoPE Embeddings 提速 28% HuyNguyen-hust 把 Unsloth 的 RoPE Embeddings 提速了约 28%！这对长上下文窗口尤其有用。通过 torch profiler 分析，Unsloth 原本的 kernel 中 RoPE 只占总运行时间的不到 2%，所以实际提速大约在 0.5% 到 1% 之间，对大规模训练尤其明显。任何提速都值得欢迎！ ⚡ 更多 4 倍速下载的 4bit 预量化模型 我们又上传了更多 4bit 模型（下载快 4 倍 + 少用 1GB 显存）。如果你有其他需求，欢迎在 Discord 上提出，或在 Github 上提交 issue！ CodeGemma 2b 和 7b 新发布的 Gemma 1.1 instruct 2b 和 7b Mistral 新重训的 7b v2，支持长上下文 Teknium 的 Open Hermes 2.5 Starling Beta Nous Research 的 Hermes Pro 🐛 重要 bug 修复 修复了 Gemma 因 tied weights 导致无法正确转换为 GGUF 的问题。 Kaggle 磁盘空间只有 20GB，导致合并到 16bit 时会失败——现在我们会自动删除 4GB 的 model.safetensors 文件，让你顺利合并到 16bit。 修复了批量生成时的推理问题，之前我们遗漏了对 attention mask 和 position ids 的处理。提醒一下，我们的原生推理速度快 2 倍！ 对 lm_head 和 embed_tokens 的微调现在可以正常工作了！记得设置 modules_to_save。 🔮 未来计划 我们正在开发自动模型优化器。社区呼声较高的需求包括 Mixtral、Command R、DBRX 等，目标是用我们的 kernel 自动优化模型，类似 torch.compile，但用的是我们自己的 kernel。 我们还在开发 Colab 上的一键式增强微调系统，希望能让大家更简单地使用 Unsloth。 💕 支持我们！欢迎通过我们的 Ko-fi 捐赠页面支持我们。特别感谢新支持者们：Roman、Henrik、Rajesh、007ok、Netrve、Goblin、pacozaa、Datta Nimmaturi、Hamel Husain、Ratish、Chris、Steffen、Remek、Anthony、Richard、Chrismcmaster、Trelis Research、preemware 和 Nam！🙏

一如既往，欢迎加入我们的 Discord 社区获取帮助或表达支持！你也可以关注我们的 Twitter 和 Substack。感谢阅读！Daniel & Michael Han 🦥
2024 年 4 月 9 日Unsloth Studio 加载中……免费开始 加入我们的 Discord