第4章 Unsloth 更新：支持 Mistral 及多项性能提升

unsloth 下载 ☰ 下载 博客 用 Unsloth 微调 Mistral，速度提升 14 倍 2023年12月14日 • 作者 Daniel Han Unsloth 开源版： Mistral 7B（1xA100）快 2.2 倍；Code Llama 34B（1xA100）快 1.9 倍；Llama 7B（1xA100）快 2.2 倍；Llama 7B（1xT4）快 2 倍。 我们很高兴地宣布：QLoRA 现已支持 Mistral 7B、CodeLlama 34B 以及所有基于 Llama 架构的模型！我们还新增了滑动窗口注意力（sliding window attention）、初步的 Windows 和 DPO 支持，并将分享全部 59 个 notebook，方便你复现我们的数据。 现在你可以在单张 A100 上对 Mistral 7B 做 QLoRA 微调，速度提升 2.2 倍，显存节省 62%，峰值仅需 12.4GB。CodeLlama 34B 在单张 A100 上快 1.9 倍，显存节省 32%，峰值为 27GB——终于不再 OOM 了！ Unsloth Pro 版本： Mistral 7B（1xA100）快 14 倍；Code Llama 34B（1xA100）快 13 倍；Llama 7B（1xA100）快 21 倍；Llama 7B（2xT4）快 28 倍。 我们的 PRO 版本在单张 A100 上微调 Mistral 7B 快了整整 14 倍，峰值显存降低 70%；CodeLlama 34B 快 13 倍，显存节省 50%，峰值 20GB；Llama 7B 快 21 倍，峰值显存降低了惊人的 71%。借助 DDP 支持，Llama 7B 在 2 张 Tesla T4 上快 28 倍。 Unsloth Pro 版本： Mistral 7B（1xA100）峰值显存 -70%；Code Llama 34B（1xA100）峰值显存 -50%；Llama 7B（1xA100）峰值显存 -71%；Llama 7B（2xT4）峰值显存 -44%。 我们还新增了直接安装 Flash Attention v2 的方式，不再依赖 Xformers。如果你有 RTX 3060 或更高（A100、H100 等），请使用 "ampere" 安装路径： pip install "unsloth[cu118_ampere] @ git+https://github.com/unslothai/unsloth.git" pip install "unsloth[cu121_ampere] @ git+https://github.com/unslothai/unsloth.git" pip install "unsloth[colab_ampere] @ git+https://github.com/unslothai/unsloth.git" 应大家要求，我们提前分享 Unsloth 的加速原理及基准测试数据。 1 A100 40GB Hugging Face Flash Attention 2 Unsloth Open Unsloth Equal Unsloth Pro Unsloth Max Alpaca 1x 1.04x 1.98x 2.48x 5.32x 15.64x LAION Chip2 1x 0.92x 1.61x 1.84x 7.05x 20.73x OASST 1x 1.19x 2.17x 2.66x 5.04x 14.83x Slim Orca 1x 1.18x 2.22x 2.64x 5.04x 14.82x 我们在 1 张 A100（通过 Google Colab）上，将 Unsloth 与 Hugging Face 的原生实现及添加了 Flash Attention 2 支持的版本进行基准测试。Flash Attention 最多只能将训练速度提升 1.2 倍，而使用开源版本的 Unsloth 可使训练速度提升 2.2 倍。“Unsloth Equal” 是我们的 PRO 版本，前提是保持所有设置和损失曲线一致。在此场景下，我们进一步将训练速度提升至 2.7 倍。我们的 MAX 版本甚至能将 LAION 数据集上的训练速度提升至 21 倍！

所有基准测试均采用以下配置（若出现 OOM 错误，我们会相应降低所有测试的 batch size）： QLoRA nf4 layers = [
"q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj",
]
QLoRA rank = 16, alpha = 16, dropout = 0
max_seq_length = 2048
learning_rate = 2e-4
weight_decay = 0.01
max_steps = 240
warmup_steps = 10
batch_size = 4
gradient_accumulation_steps = 4
lr_scheduler_type = "linear"
optimizer = "adamw_8bit", bfloat16
use_gradient_checkpointing = True
random_state = 3407 1 A100 40GB Hugging Face Flash Attention 2 Unsloth Open Unsloth Equal Unsloth Pro Unsloth Max Mistral 7B Slim Orca 1x 1.15x 2.15x 2.53x 4.61x 13.69x code Code Code Code Code seconds 1813 1571 842 718 393 132 memory MB 32853 19385 12465 10271 memory saved % 40.99 62.06 68.74 在 1 张 A100 上训练 Mistral 7B 时，Flash Attention v2 将训练速度提升至 1.15 倍，Unsloth Open 则提升至 2.15 倍，同时内存占用减少 62%。同样地，“Unsloth Equal” 执行均衡的训练流程，可将速度提升至 2.53 倍，且峰值 VRAM 使用量减少 69%。Unsloth MAX 版本更是将训练速度提升了 13.7 倍。

点击上方 “Code” 按钮即可查看我们的共享 Notebook，以确保实验可复现。“Unsloth Equal” 仅展示训练损失，未显示其他代码路径。 1 A100 40GB Hugging Face Flash Attention 2 Unsloth Open Unsloth Equal Unsloth Pro Unsloth Max Code Llama 34B OOM ❌ 0.99x 1.87x 2.61x 4.27x 12.82x code Code Code Code Code seconds 1953 1982 1043 748 458 152 memory MB 40000 33217 27413 22161 memory saved % 16.96 31.47 44.60 在 CodeLlama 34B 上，我们采用 batch size 为 1、序列长度为 4096 的设置。遗憾的是，Hugging Face 的原始实现在 batch size 为 2 时会出现显存溢出（OOM），因此基准测试只能针对 bsz = 1 进行。Flash Attention v2 对运行时间的影响微乎其微，而 Unsloth Open 的速度提升了 1.87 倍，且峰值 VRAM 占用减少了 32%。“Unsloth Equal” 速度提升 2.6 倍，内存占用降低 45%，而 MAX 版速度更是快了 12.8 倍。 2 T4 DDP Hugging Face Flash Attention Unsloth Open Unsloth Equal Unsloth Pro Unsloth Max LAION Chip2 1x 1.12x 5.28x 4.21x 10.01x 28.32x code Code Code Code seconds 5418 4854 1027 1286 541 191 memory MB 7316 7316 5732 5934 memory saved % 0.00 21.65 18.89 通过 Kaggle 的 2 个 Tesla T4 实例，我们发现仅使用 1 个 GPU 时，Unsloth Open 的训练速度就快了 5.3 倍。为公平起见，我们将梯度累积步数乘以 2，因为开源版本不支持多 GPU。“Unsloth Equal” 通过 DDP 加速了 4.21 倍。由于 DDP 需要在每步同步梯度，因此存在一定的开销。Unsloth MAX 的训练速度提升了 28 倍！

这篇博客末尾我们附上了完整的基准测试表格和全部 59 个 notebook 链接，方便复现结果。逐项性能拆解

No.	Method	Time (s)	Peak VRAM (GB)	Time saved (%)	VRAM saved (%)	Final error
1	Huggingface Original PEFT QLoRA	594	16.7			1.0202
2	Reduce data upcasting	465	15.5	21.7%	7.2%	1.0203
3	Bitsandbytes bfloat16	424	15.3	8.9%	1.3%	1.0208
4	SDPA	418	14.9	1.4%	2.6%	1.0214
5	SDPA causal = True	384	14.9	8.1%	0.0%	1.0219
6	Xformers	353	9.1	8.1%	38.9%	1.021
7	Flash Attention 2	353	9.1	0.0%	0.0%	1.0215
8	Fast RoPE Embeddings	326	9	7.6%	1.1%	1.0211
9	Fast RMS Layernorm	316	9	3.1%	0.0%	1.021
10	Fast Cross Entropy Loss	315	7.4	0.4%	17.8%	1.021
11	Manual Autograd MLP	302	6.8	4.0%	8.1%	1.0222
12	Manual Autograd QKV	297	6.8	1.7%	0.0%	1.0217

以下拆解了我们在 AWS 的 A10G GPU 上对 Unsloth 开源版本所做的每一项优化。全部落地后，训练速度提升 2 倍，峰值显存占用减少 61%。本节数学内容偏多，提前提醒！1. 减少数据类型提升 通过减少 QLoRA 过程中权重的类型提升，就能轻松省下 7.2% 显存，训练时间缩短 21.7%。2. Bitsandbytes bfloat16 Bitsandbytes 内部使用 float16，所以需要额外一次内存拷贝才能转换成 bfloat16。我们在内部修复了这个问题，省下 9% 时间。3. Scaled Dot Product Attention 我们改用 Pytorch 的高性能 attention 实现，节省 1.4% 时间。4、5、6. 因果掩码、Xformers、Flash Attention 2 用因果掩码代替单独的 attention mask 后，由于无需再读取 attention 矩阵，速度快了 8.1%。接着切换到 Xformers，又提速 8.1%，显存占用大幅下降 39%。换用 Flash Attention v2 则没有明显效果，因为 Xformers 内部本来就是调用 FA2。7. 快速 RoPE Embeddings 用 OpenAI 的 Triton 实现 RoPE Embeddings，又省下 7.6% 时间。但为此需要手动推导 RoPE 函数的导数。注意 RoPE 可以改写成旋转矩阵 R 与原矩阵 Q 的矩阵乘法，这样导数就是 R 的转置。8. 快速 RMS Layernorm 遗憾的是，RMS Layernorm 的导数复杂得多。仔细运用链式法则推导后，会得到一个相当丑陋的导数表达式。我们同样用 OpenAI 的 Triton 实现它，训练再提速 3.1%。9. 快速交叉熵损失 \begin{align} \text{loss} = \text{CrossEntropyLoss}(\text{logits}, \text{labels}) \\ CE_{loss} = \frac{1}{n} \sum{ - y_i \log{p_i}} \\ CE_{loss} = \frac{1}{n} \sum{ - y_i \log{\frac{\exp(x_i)}{\sum{\exp{x_i}}}}} \\ CE_i = - y_i \log{\frac{\exp(x_i)}{\sum{\exp{x_i}}}} \\ = -y_i (x_i - \log{\sum{\exp{x_i}}}) \\ = y_i ( \text{logsumexp}(x) - x ) \\ = \left\{ \\ \begin{array}{ c l } 0 & \quad \textrm{if } y = 0 \\ \text{logsumexp}(x)-x & \quad \textrm{otherwise} \\ \end{array} \\ \right. \\ \frac{dC}{dx_i} = y_i \cdot \text{exp} \big( x - \text{logsumexp}(x) \big) - \frac{d}{dx_i} x_k \cdot y_i \end{align} 交叉熵损失也要复杂一些。我们用 log 技巧（即 x = exp(log(x))）来推导导数，还借助维基百科确认了臭名昭著的 logsumexp 函数的导数其实就是 softmax 函数！显存占用因此砍掉了 17%。10、11. 手动 Autograd 只要正确地加括号，就能大幅减少 LoRA 微调时实际的 FLOPs 数量！通常 Pytorch 的 autograd 引擎会从计算图末端向起点反向传播。我们发现，把多个算子融合为一个，并通过链式矩阵乘法正确加括号，实际的 FLOPs 数量就能降下来。 (X_{(m,d)}^T \times dW_{(m,h)}) \times B_{(h,r)}^T 如果括号加得不对——Pytorch 的 autograd 目前正是如此——就会先计算 X.T 和 dW 的乘积。设 X 的形状为 (bsz, seq_len, d)，先把它 reshape 成 (m, d)，其中 m 就是 bsz * seq_len，d 是 attention 维度：Llama 7b 中是 4096，Llama 70b 中是 8192。

dW 的维度为 (m, h)，其中 h 表示 MLP 的中间层维度。在 Llama 7b 中该值为 11,008，而在 Llama 70b 中则为 28,672。B.T 是 LoRA 权重，维度为 (h, r)，其中 r 为 LoRA 矩阵的秩，通常可以很小，例如 16 或 64。

我们发现，采用慢速计算路径时，所需的 FLOPs 约为 (h * d)(m + r)。

而在快速路径中，我们对第二项施加括号优先运算，此时所需的 FLOPs 为 (m * r)(h + d)。将慢速路径的 FLOPs 除以快速路径，即可得到加速比：

为了简化上述表达式，我们注意到 r 通常很小，例如 16 或 64。而 m 可以很大，比如批量大小为 4、序列长度为 4096 时，m 就会达到 4 * 4096 = 16,384。因此 (m + r) 中的加数影响微乎其微，可以忽略不计。但对于 (m * r) 中的乘法，由于 16 的乘积远大于 16 的加数，故不能省略。

经过这样简化后，我们得到了一个简化的表达式：加速比主要取决于 MLP 中间层维度 h、注意力维度 d 以及 LoRA 秩 r。

对于 Llama 7b（h = 11,008，d = 4,096，r = 16），加速比达到 186.58。
对于 Llama 70b（h = 28,672，d = 8,192，r = 16），加速比达到 398.22。其他新特性152334H 成功让 Unsloth 支持 DPO！目前还是初步支持，但通过 TRL 似乎已经可以正常运行。RandomInternetPreson 成功让 Unsloth 在 WSL 上运行！也就是说 Windows 支持目前处于初步阶段。其他 bug 修复——支持最大到 2^16（65536）的所有词表大小，group query attention 现在也能正常工作了。旧款 GPU 上的 GQA 现在通过 Xformers 得到了完整支持——我们手动对 K 和 V 进行 reshape，让 Xformers 以普通 attention 的方式完成计算。遗憾的是，Xformers 不支持 GQA 的反向传播。常见问题Q：支持 Mixtral 吗？
正在开发中！Q：如何购买 PRO 或 MAX 版本？
我们正在搭建一个平台，敬请期待！Q：会减少 FLOPs 吗？
会。Q：开源版支持全量微调吗？
不支持，参见 Issue。开源版关闭了所有优化，除了 Flash Attention 和部分 Triton kernel 之外，速度不会有明显提升。Q：支持 LoRA（而非 QLoRA）吗？
支持。把 load_in_4bit 设为 False 即可。Q：PRO / MAX 版本支持全量微调和预训练吗？
支持。感谢阅读！🦥Daniel Han
2023 年 12 月 13 日 完整基准测试表 1 T4 16GB Hugging Face Flash Attention Unsloth Open Unsloth Pro Equal Unsloth Pro Unsloth Max Alpaca 1x 1.09x 1.69x 1.79x 2.93x 8.3x 代码 Code Code Code Code 秒 1599 1468 942 894 545 193 内存 MB 7199 7059 6459 5443 内存节省百分比 1.94 10.28 24.39 1 T4 16GB Hugging Face Flash Attention Unsloth Open Unsloth Pro Equal Unsloth Pro Unsloth Max LAION Chip2 1x 0.99x 1.80x 1.75x 4.15x 11.75x 代码 Code Code Code Code 秒 952 955 529 543 229 81 内存 MB 6037 6033 5797 4855 内存节省百分比 0.07 3.98 19.58 1 T4 16GB Hugging Face Flash Attention Unsloth Open Unsloth Pro Equal Unsloth Pro Unsloth Max OASST 1x 1.19x 1.95x 1.86x 2.58x 7.3x 代码 Code Code Code Code 秒 2640 2222 1355 1421 1024 362 内存 MB 14827 10391 8413 7031 内存节省百分比 29.92 43.26 52.58 1 T4 16GB Hugging Face Flash Attention Unsloth Open Unsloth Pro Equal Unsloth Pro Unsloth Max Slim Orca 1x 1.21x 1.77x 1.85x 2.71x 7.67x 代码 Code Code Code Code 秒 2735 2262 1545 1478 1009 356 内存 MB 13933 10489 7661 6563 内存节省百分比 24.72 45.02 52.90 2 T4 DDP Hugging Face Flash Attention Unsloth Open Unsloth Equal Unsloth Pro Unsloth Max Alpaca 1x 0.99x 4.95x 4.44x 7.28x 20.61x 代码 Code Code Code 秒 9882 9946 1996 2227 1357 480 内存 MB 9176 9128 6904 6782 内存节省百分比 0.52 24.76 26.09 2 T4 DDP Hugging Face Flash Attention Unsloth Open Unsloth Equal Unsloth Pro Unsloth Max LAION Chip2 1x 1.12x 5.28x 4.21x 10.01x 28.32x 代码 Code Code Code 秒 5418 4854 1027 1286 541 191 内存 MB 7316 7316 5732 5934 内存节省百分比 0.00 21.65 18.89 2 T4 DDP Hugging Face Flash Attention Unsloth Open Unsloth Equal Unsloth Pro Unsloth Max OASST (bsz=1) 1x 1.14x 5.56x 5.09x 5.64x 15.97x 代码 Code Code Code 秒 4503 3955 811 885 798 282 内存 MB 11896 11628 6616 7105 内存节省百分比 2.25 44.38 40.27 2 T4 DDP Hugging Face Flash Attention Unsloth Open Unsloth Equal Unsloth Pro Unsloth Max Slim Orca (bsz=1) 1x 0.97x 5.54x 4.68x 6.88x 19.46x 代码 Code Code Code 秒 4042 4158 729 863 588 208 内存 MB 11010 11042 6492 7410 内存节省百分比 -0.29 41.04 32.70 2 T4 DDP Hugging Face Flash Attention Unsloth Open Unsloth Equal Unsloth Pro Unsloth Max OASST (bsz=2) OOM ❌ OOM ❌ ✓ ✓ ✓ ✓ 代码 Code Code Code 秒 OOM OOM 2719 3391 2794 987 内存 MB OOM OOM 8134 9600 内存节省百分比 OOM OOM 2 T4 DDP Hugging Face Flash Attention Unsloth Open Unsloth Equal Unsloth Pro Unsloth Max Slim Orca (bsz=2) OOM ❌ OOM ❌ ✓ ✓ ✓ ✓ 代码 Code Code Code 秒 OOM OOM 2990 3444 2351 831 内存 MB OOM OOM 7594 8881 内存节省百分比 OOM OOM 1 A100 40GB Hugging Face Flash Attention 2 Unsloth Open Unsloth Equal Unsloth Pro Unsloth Max Alpaca 1x 1.04x 1.98x 2.48x 5.32x 15.64x 代码 Code Code Code Code 秒 1040 1001 525 419 196 67 内存 MB 18235 15365 9631 8525 内存节省百分比 15.74 47.18 53.25 1 A100 40GB Hugging Face Flash Attention 2 Unsloth Open Unsloth Equal Unsloth Pro Unsloth Max LAION Chip2 1x 0.92x 1.61x 1.84x 7.05x 20.73x 代码 Code Code Code Code 秒 581 631 361 315 82 28 内存 MB 7763 8047 7763 6441 内存节省百分比 -3.66 0.00 17.03 1 A100 40GB Hugging Face Flash Attention 2 Unsloth Open Unsloth Equal Unsloth Pro Unsloth Max OASST 1x 1.19x 2.17x 2.66x 5.04x 14.83x 代码 Code Code Code Code 秒 1852 1558 852 696 367 125 内存 MB 26431 16565 12267 11223 内存节省百分比 37.33 53.59 57.54 1 A100 40GB Hugging Face Flash Attention 2 Unsloth Open Unsloth Equal Unsloth Pro Unsloth Max Slim Orca 1x 1.18x 2.22x 2.64x 5.04x 14.82x 代码 Code Code Code Code 秒 1824 1545 821 691 362 123 内存 MB 24557 15681 10595 9007 内存节省百分比 36.14 56.86 63.32 1 A100 40GB Hugging Face Flash Attention 2 Unsloth Open Unsloth Equal Unsloth Pro Unsloth Max Mistral 7B Slim Orca 1x 1.15x 2.15x 2.53x 4.61x 13.69x 代码 Code Code Code Code 秒 1813 1571 842 718 393 132 内存 MB 32853 19385 12465 10271 内存节省百分比 40.99 62.06 68.74 1 A100 40GB Hugging Face Flash Attention 2 Unsloth Open Unsloth Equal Unsloth Pro Unsloth Max Code Llama 34B OOM ❌ 0.99x 1.87x 2.61x 4.27x 12.82x 代码 Code Code Code Code 秒 1953 1982 1043 748 458 152 内存 MB 40000 33217 27413 22161 内存节省百分比 16.96 31.47 44.60 现在 AI 训练不再慢了！免费开始 加入我们的 Discord