进阶 unsloth.ai 2026-10-07 22:27:00 · 6 阅读
第3章 Unsloth：LLM 训练提速 30 倍
unsloth
Download
☰
下载
博客Introducing Unsloth：LLM 训练提速 30 倍2023 年 12 月 1 日 • 作者：Daniel Han我们很高兴地宣布，我们的 AI 创业公司正式成立，专注于打造有趣的 AI 产品！首个产品是 Unsloth，它能让 LLM 训练速度提升 30 倍！主要亮点：
提速 30 倍。Alpaca 训练时间从 85 小时缩短到 3 小时。
显存占用降低 60%，batch 大小可扩大 6 倍。
准确率零损失；Max 版本甚至可提升 20% 的准确率。
无需新硬件，只需更换软件。
Max 版本支持 NVIDIA、Intel 和 AMD GPU。
手写 autograd 及链式矩阵乘法优化。
使用 OpenAI 的 Triton 语言重写了所有 kernel。
通过 xformers 和 Tri Dao 的实现支持 Flash Attention。
免费开源版可将 finetuning 提速 2 倍、显存占用降低 50%。
快来试试我们的开源版本！pip install "unsloth[cu118] @ git+https://github.com/unslothai/unsloth.git"
pip install "unsloth[cu121] @ git+https://github.com/unslothai/unsloth.git"我们还提供了 Google Colab notebooks，包含 Alpaca 数据集和 SlimOrca 的训练示例。此外，我们还为 LAION 数据集提供了 Kaggle notebook。手动 Autograd
$Q
=
X
W
q
\tilde{}
=
X
(
W
q
+
A
q
B
q
)
$K
=
X
W
k
\tilde{}
=
X
(
W
k
+
A
k
B
k
)
$V
=
X
W
v
\tilde{}
=
X
(
W
v
+
A
v
B
v
)
A
(
X
)
=
\sigma
(
\frac{1}{d}
Q
K^T
+
M
)
V
PyTorch 的 Autograd 在大多数任务上效率尚可，但如果你追求极致性能，就必须手动推导矩阵微分。在臭名昭著的 Attention 机制与 LoRA 适配器结合的场景中，由于我们冻结了原始权重矩阵 $W$，只训练 $A$ 和 $B$，因此需要求解 6 个矩阵微分。$D(W)$ 仅作为反向传播过程中导数的记法。由于我们在 Llama 的所有 32 层上都应用了 LoRA，因此还必须推导 Attention Head 输入相对于输入的导数，以便将梯度反向传播直至第 1 层。
我们自行推导偏导数后发现，LoRA 权重矩阵极其瘦长，维度最大约 128，最小可能仅为 8；而 Llama 权重的维度通常是 1024 的整数倍，大多达到 4096 或更大。因此，在链式矩阵乘法中正确放置括号对于性能至关重要。通过简单的矩阵维度 FLOP 计算，我们发现将 LoRA 权重乘法放入括号内能带来出色的性能。同样，我们发现最终计算 dC/dX 时应使用 torch 的 inplace 操作以节省内存。我们对其他所有层（MLP、lm_head、layernorm、RoPE embeddings）也应用了这类手动微分优化。 性能分析 在赫赫有名的 Alpaca 数据集上，单张 Tesla T4 GPU 测试显示，Huggingface 原始实现耗时 23 小时 15 分钟，而我们的 Max 方案仅需 2 小时 34 分钟，速度提升 8.8 倍。在 SlimOrca 数据集上，耗时从 391 小时缩减至 51 小时，速度提升 7.6 倍。 使用 2 张 Tesla T4 GPU 进行 DDP 训练时，LAION Chip2 数据集耗时 164 小时，而我们的方案仅需 5 小时，速度提升 31 倍。在 SlimOrca 数据集上，耗时从 1301 小时（约 54 天）缩减至 54 小时，速度提升 24 倍。 关于 Open Assistant 数据集的内存占用，在单张支持 bfloat16 的 A10 GPU 上，峰值内存从 16.7GB 降至 6.9GB（减少 59%）；在 Tesla T4 上，峰值内存从 14.6GB 降至 7.5GB（减少 49%）。 OpenAI Triton kernels @triton.jit def _rope_embedding( Q, Q_row_stride, cos, cos_row_stride, sin, sin_row_stride, seqlen, head_dim, BACKWARD_PASS: tl.constexpr, BLOCK_SIZE : tl.constexpr, ): row_position = tl.program_id(0) head_position = tl.program_id(1) col_offsets = tl.arange(0, BLOCK_SIZE) half_head_dim = head_dim // 2 mask = col_offsets < half_head_dim rot_position = row_position % seqlen Q += row_position* Q_row_stride + head_position*head_dim cos += rot_position*cos_row_stride sin += rot_position*sin_row_stride Q1 = tl.load(Q + half_head_dim*0 + col_offsets, mask = mask, other = 0) sin1 = tl.load(sin + half_head_dim*0 + col_offsets, mask = mask, other = 0) cos1 = tl.load(cos + half_head_dim*0 + col_offsets, mask = mask, other = 0) Q2 = tl.load(Q + half_head_dim*1 + col_offsets, mask = mask, other = 0)
tl.store(Q + half_head_dim*0 + col_offsets,
Q1*cos1 - Q2*sin1, mask = mask)
tl.store(Q + half_head_dim*1 + col_offsets,
Q2*cos1 + Q1*sin1, mask = mask)
pass我们还用 OpenAI 的 Triton 语言重写了所有 kernel。比如上面这段就是带反向传播的 RoPE embedding kernel 实现，我们尽力让它简洁、易读、高效。更多手写 kernel 可以到我们的 GitHub 仓库查看。 我们的请求 快来试用我们的开源包 Unsloth，微调速度提升 2 倍！ 我们的 Pro 版本提供多 GPU 支持和更疯狂的加速等功能。Max 版本还提供从零完整训练 LLM 的 kernel，并支持把我们的代码移植到 AMD 和 Intel GPU 上。我们已经成功把 bitsandbytes 的 4bit 量化方法移植到了 Triton！详情请联系我们！ 我们打算推出一个训练和推理平台——如果有足够多的人感兴趣，我们就会做！目前我们正在用 ChatGPT / Llama 等 LLM 来搭建前端和后端，因为我（Daniel）完全没有给用户提供 Cloud GPU 的经验。如果你能帮上忙，欢迎发邮件到 [email protected]。 我们正在白手起家创业。如果有人愿意和我们建立商业合作伙伴关系，或者购买我们的代码（也许？），或者聊聊投资——欢迎发邮件！非常感谢！ 未来计划 在 NVIDIA、Intel 和 AMD GPU 上让推理更快。 Sqrt gradient checkpointing——进一步降低 25% 的内存占用。 参加 AI-MO 竞赛。 让 DPO 也更快。OpenAI 的 ChatGPT 用 PPO 做 RLHF，而 DPO 已被证明同样强大。 未来几个月我们还有多个产品发布计划（金融 LLM、经济衰退预测器、数据科学顾问 agent 等）！！！ 感谢阅读！🦥 Daniel Han
2023年11月30日 立即加速 AI 训练！ 免费开始 加入我们的 Discord
我们自行推导偏导数后发现，LoRA 权重矩阵极其瘦长，维度最大约 128，最小可能仅为 8；而 Llama 权重的维度通常是 1024 的整数倍，大多达到 4096 或更大。因此，在链式矩阵乘法中正确放置括号对于性能至关重要。通过简单的矩阵维度 FLOP 计算，我们发现将 LoRA 权重乘法放入括号内能带来出色的性能。同样，我们发现最终计算 dC/dX 时应使用 torch 的 inplace 操作以节省内存。我们对其他所有层（MLP、lm_head、layernorm、RoPE embeddings）也应用了这类手动微分优化。 性能分析 在赫赫有名的 Alpaca 数据集上，单张 Tesla T4 GPU 测试显示，Huggingface 原始实现耗时 23 小时 15 分钟，而我们的 Max 方案仅需 2 小时 34 分钟，速度提升 8.8 倍。在 SlimOrca 数据集上，耗时从 391 小时缩减至 51 小时，速度提升 7.6 倍。 使用 2 张 Tesla T4 GPU 进行 DDP 训练时，LAION Chip2 数据集耗时 164 小时，而我们的方案仅需 5 小时，速度提升 31 倍。在 SlimOrca 数据集上，耗时从 1301 小时（约 54 天）缩减至 54 小时，速度提升 24 倍。 关于 Open Assistant 数据集的内存占用，在单张支持 bfloat16 的 A10 GPU 上，峰值内存从 16.7GB 降至 6.9GB（减少 59%）；在 Tesla T4 上，峰值内存从 14.6GB 降至 7.5GB（减少 49%）。 OpenAI Triton kernels @triton.jit def _rope_embedding( Q, Q_row_stride, cos, cos_row_stride, sin, sin_row_stride, seqlen, head_dim, BACKWARD_PASS: tl.constexpr, BLOCK_SIZE : tl.constexpr, ): row_position = tl.program_id(0) head_position = tl.program_id(1) col_offsets = tl.arange(0, BLOCK_SIZE) half_head_dim = head_dim // 2 mask = col_offsets < half_head_dim rot_position = row_position % seqlen Q += row_position* Q_row_stride + head_position*head_dim cos += rot_position*cos_row_stride sin += rot_position*sin_row_stride Q1 = tl.load(Q + half_head_dim*0 + col_offsets, mask = mask, other = 0) sin1 = tl.load(sin + half_head_dim*0 + col_offsets, mask = mask, other = 0) cos1 = tl.load(cos + half_head_dim*0 + col_offsets, mask = mask, other = 0) Q2 = tl.load(Q + half_head_dim*1 + col_offsets, mask = mask, other = 0)
tl.store(Q + half_head_dim*0 + col_offsets,
Q1*cos1 - Q2*sin1, mask = mask)
tl.store(Q + half_head_dim*1 + col_offsets,
Q2*cos1 + Q1*sin1, mask = mask)
pass我们还用 OpenAI 的 Triton 语言重写了所有 kernel。比如上面这段就是带反向传播的 RoPE embedding kernel 实现，我们尽力让它简洁、易读、高效。更多手写 kernel 可以到我们的 GitHub 仓库查看。 我们的请求 快来试用我们的开源包 Unsloth，微调速度提升 2 倍！ 我们的 Pro 版本提供多 GPU 支持和更疯狂的加速等功能。Max 版本还提供从零完整训练 LLM 的 kernel，并支持把我们的代码移植到 AMD 和 Intel GPU 上。我们已经成功把 bitsandbytes 的 4bit 量化方法移植到了 Triton！详情请联系我们！ 我们打算推出一个训练和推理平台——如果有足够多的人感兴趣，我们就会做！目前我们正在用 ChatGPT / Llama 等 LLM 来搭建前端和后端，因为我（Daniel）完全没有给用户提供 Cloud GPU 的经验。如果你能帮上忙，欢迎发邮件到 [email protected]。 我们正在白手起家创业。如果有人愿意和我们建立商业合作伙伴关系，或者购买我们的代码（也许？），或者聊聊投资——欢迎发邮件！非常感谢！ 未来计划 在 NVIDIA、Intel 和 AMD GPU 上让推理更快。 Sqrt gradient checkpointing——进一步降低 25% 的内存占用。 参加 AI-MO 竞赛。 让 DPO 也更快。OpenAI 的 ChatGPT 用 PPO 做 RLHF，而 DPO 已被证明同样强大。 未来几个月我们还有多个产品发布计划（金融 LLM、经济衰退预测器、数据科学顾问 agent 等）！！！ 感谢阅读！🦥 Daniel Han
2023年11月30日 立即加速 AI 训练！ 免费开始 加入我们的 Discord
