Unsloth 将 LLM 训练速度提升 40% 以上

我们与 NVIDIA 合作，使 LLM 训练速度提升约 25%，本文将详细分解我们的优化方法。这些优化不会导致精度损失，是在 Unsloth 已有的 2-5 倍加速基础上的额外改进！这些新算法已在 RTX 笔记本、数据中心 GPU 和 DGX Spark 机器上自动启用，只需更新 Unsloth 即可获得最新优化。通过与 NVIDIA 的合作，我们演示了如何：

缓存打包序列元数据使训练快 14.3%。

使用双缓冲异步梯度检查点实现 8% 的加速。

gpt-oss 训练通过在 MoE 路由中使用 argsort 和 bincount 快 15%。

1. 缓存打包序列元数据

假设我们有几个短样例：

与其将它们全部填充到相同长度并浪费计算力在填充 token 上，我们将它们连接成一个更长的打包序列：

模型仍然需要知道每个原始序列从何处开始和结束。因此，除了打包 token 之外，我们还需要序列元数据，例如：

累积序列偏移量（cu_seqlens）

最大序列长度

从上述三项推导的注意力结构

关键点是：对于固定的打包批次，该元数据在每一层都是相同的。

如果我们将打包批次的边界信息写成：

B = { lengths, cu_seqlens, max_seqlen, mask structure }

那么该前向传播中的每个 transformer 层都使用相同的 B。

如果模型有 L 层，每层都重新构建或重新同步 B 不是新工作。这是相同信息被一次次重新构建。

换句话说，有用的工作是：

构建 B 一次，使用 L 次。

浪费的版本是：

构建 B + 构建 B + ⋯ + 构建 B（L 次）

这里的开销不主要是额外的 FLOPs。某些路径可能会强制设备到主机的同步，实际上创建了一个 GPU-CPU 同步点。一旦这发生在每层路径内，开销就会在每层重复出现。

这正是打包序列缓存改变所减少的。与其重复重建打包序列信息、SDPA 打包掩码和 xFormers 块掩码，我们缓存可重用的元数据和从中推导的注意力端结构，按设备，为当前打包批次。然后跨层重用这些缓存结构。

打包训练已经通过消除填充浪费来改善利用率。但如果元数据路径持续强制同步，某些收益会因与模型实际学习无关的开销而丧失。

缓存有帮助是因为它从热路径中移除了重复的协调工作。前向传播受益最多，因为这是相同打包元数据在多层中被重复使用的地方。

在 Qwen3-14B QLoRA SFT 上：

前向传播获得最大收益，因为重复的元数据和掩码准备在那里最直接地出现。反向传播也有改进，但效果较小。节省的时间相似，但反向传播，尤其是使用梯度检查点时，耗时较长，所以相对收益显得较小。

现在我们知道了测量的收益，可以问一个更简单的问题：这个数值是否合理？

如果我们假设每层大致相似，可以将打包注意力路径建模为：

T_uncached ≈ L · (A + s)

L 是层数，

A 是每层有用的注意力端工作，

s 是每层重复的元数据和掩码准备开销。

通过缓存，重复开销为批次付一次而非每层一次：

所以节省的时间约为：

T_saved ≈ (L − 1) · s

对于打包 SDPA 路径，我们在 NVIDIA Blackwell GPU 上的微基准测试显示，低级别的主机可见元数据调用是真实的但很小，约 0.2 ms 每次。主要的重复成本是打包 SDPA 掩码构造路径本身，对于包含 2048 个总打包 token 的综合打包批次测量约 13.7 ms。

对于 SDPA 后端，更好的心智模型是：

小流围栏 + 掩码重建 ≈ 掩码重建

这让我们可以做一个更清晰的一致性检查。如果一个打包掩码重建耗费 m 毫秒，那么在均匀层模型下：

T_saved ≈ (L − 1) · m

取 m ≈ 13.7 ms，这预测：

16 层：(16 - 1) x 13.7 ≈ 206 ms

28 层：(28 - 1) x 13.7 ≈ 370 ms

较小的打包序列运行显示相同的模式：

Llama-3.2-1B，16 层：每步约 199 ms 节省，约是端到端步时间的 11.5% 降低

Qwen3-0.6B，28 层：每步约 319 ms 节省，约是端到端步时间的 14.8% 降低

这些百分比是相对于完整训练步时间，所以它们仍然包括打包注意力路径外的工作，例如嵌入、MLP、LM 头、损失和框架开销。这个估计只是关于打包 SDPA 路径端的块，而不是整个 transformer 层。它只是为了检查测量收益是否在打包 SDPA 路径的合理范围内。

2. 使用双缓冲检查点重载隐藏延迟

激活检查点是训练大型模型的标准技术。其思想是通过不在反向传播中保持每个中间激活活跃来节省内存。作为交换，我们在反向传播期间付出一些额外的工作。

这种权衡通常是值得的，特别是对于较大的模型。

但这引发了另一个系统问题：如果激活已被卸载，它如何在反向传播时回到 GPU？

在 Unsloth 的智能检查点路径中，激活可以在固定 CPU 内存中分阶段放置，并在需要时复制回来。这节省了 VRAM，但它可能引入瓶颈：

从 CPU 复制激活到 GPU。

等待复制完成。

对该激活运行反向计算。

这是一个序列化模式。如果一个缓冲区被同时用于复制和计算，复制流和计算流互相占用。

设 T_copy 为激活重载时间，T_compute 为当前层的反向计算时间。

使用单个缓冲区，这部分步的限制大约为：

T_single ≈ T_copy + T_compute

这是序列化情况。我们几乎完全为两者付费，一个接一个。

更好的处理方式是使用两个缓冲区。

当反向传播在缓冲区 A 上运行时，复制流可以将下一个激活预加载到缓冲区 B。然后角色交换。这创建了管道重叠，虽然不是完美重叠。

双缓冲不会减少数学计算量。它将复制延迟隐藏在有用的计算后面。

这种优化倾向于在模型足够大使得反向计算是实质性的，但不是非常占主导地位至掩盖所有复制开销为噪声的情况下变强。对于较大的模型，更高的隐藏维度意味着更多的数据移动，所以隐藏该移动有更大的影响。较大的模型还倾向于有更多层，这创建了更多隐藏复制在计算后面的机会。

这就是为什么较大的稠密模型适合这个改进。GPU 有足够真实的工作进行，复制可以与之重叠，第二个缓冲区需要的额外 VRAM 保持适度。

实现还保持了实用的护栏：

仅当有足够 VRAM 可用时使用额外缓冲区

内存预算紧张时干净地回退

保持正确性不变

在使用 NVIDIA B200 Blackwell GPU 进行的较大稠密模型运行上进行基准测试：

8B：0.3739 -> 0.4053 步/秒，+8.40%

14B：0.2245 -> 0.2395 步/秒，+6.70%

32B：0.1979 -> 0.2070 步/秒，+4.61%

内存开销保持适度：

在这些运行中，最终损失实际上没有改变。

加速在较大稠密模型中是一致的，额外 VRAM 成本保持相对较小。

一旦我们知道测量的收益，自然的后续问题是：这个数值是否合理？

如果我们假设有 L 个检查点层且每层大致相似：

每次重载耗费时间 c

每个反向计算块耗费时间 g

这也随批大小、序列长度和其他影响数据移动和计算的因素缩放。为简洁起见，我们省略这些术语。

T_single ≈ L · (c + g)

使用两个缓冲区，第一层仍然必须等待其激活到达，最后一层仍然必须完成计算。所以更好的近似是：

T_double ≈ c + (L − 1) · max(c, g) + g

所以节省的时间约为：

T_saved ≈ (L − 1) · min(c, g)

这是结果的有用解读：

第一个复制仍然暴露

最后的计算仍然暴露

但对于管道的中间部分，复制和计算可以重叠

如果重叠良好，中间的每层成本更接近：

T_middle ≈ max(T_copy, T_compute)

从测量的较大模型结果，每训练步节省的时间大约为：

这些主机缓冲区是固定分配，所以相关带宽是测量的固定内存主机到设备带宽，而不是可分页内存带宽。在我们的基于 NVIDIA B200 Blackwell 的系统上，该带宽约为 55.7 GB/s，64 GB/s 是有用的 PCIe 上限以供参考。

如果我们使用额外缓冲区大小作为一个激活重载的粗略代理，那么每次重载自然地在仅几毫秒的数量级：

8B，0.37 GB：在 55.7 GB/s 时约 6.6 ms，或在 64 GB/s 上限时 5.8 ms

14B，0.47 GB：在 55.7 GB/s 时约 8.4 ms，或在 64 GB/s 上限时 7.3 ms

32B，0.23 GB：在 55.7 GB/s 时约 4.1 ms，或在 64 GB/s 上限时 3.6 ms

要解释观察到的每步节省的时间，我们需要隐藏大约数十次这样的重载：

8B：在 55.7 GB/s 时约 31 次重载，或在 64 GB/s 时 36 次

14B：在 55.7 GB/s 时约 33 次重载，或在 64 GB/s 时 38 次

32B：在 55.7 GB/s 时约 54 次重载，或在 64 GB/s 时 62 次

在数十个检查点层中隐藏一次这样的重载落入了我们观察到的节省步时间的几百毫秒范围内。

同样，节省的时间是完整端到端训练步的一部分。它不应该解释嵌入、LM 头、损失、优化器工作或每个其他非检查点部分的步。关键是我们可以隐藏的通信足够大以合理解释测量的步时间收益。

3. 一个较小但有用的 MoE 优化

第三个改变更专门化，但它在 MoE 路由中展示了相同的模式。

在我们检查的基于 PyTorch 的 GPT-OSS MoE 路径中，路由的一个昂贵部分是找出哪些 token 去往哪个专家。一个天真的实现可能会做类似的事：

for expert_idx in range(num_experts):
    token_idx, _ = torch.where(router_indices == expert_idx)

乍一看，这看起来无害。但 torch.where 在这里是依赖数据的：路由到每个专家的 token 数量从批到批改变。这可能引入 CPU-GPU 同步或相关的运行时开销，因为输出大小依赖于路由模式。如果这对每个专家发生一次，动态查询的数量随 num_experts 缩放。

更好的方法是一次性组织所有内容：

展平所有专家分配。

按专家 ID 稳定排序。

使用一次 bincount 获取每个专家的 token 数。

从这些计数构建偏移量。

按专家切分分组的 token 列表。

在数学上，收益不是我们改变了路由逻辑。我们改变了我们询问运行时动态索引问题的频率。

动态查询开销 ∝ num_experts

因为我们对每个专家做一个动态查询，我们更接近：

动态查询开销 ∝ 1

加上之后的便宜记账工作。

这在更专门化的环境中是相同的主题：一次性分组，然后重用偏移量而不是重复询问动态 token 列表。

注意这些优化适用于任何使用 native_torch 后端的 MoE。

对于这个 GPT-OSS 特定的路由改进：

团队验证在 GPT-OSS 配置上显示约 10-15% 的加速

在目标路由路径中，我们看到 +23% 前向和 +13% 反向

这些改变的共同之处

虽然这三个优化存在于堆栈的不同部分，它们在解决相同的问题。

关键优化机会在于主内核周围的胶合代码中：

重建已存在的元数据

同步我们可能缓存的信息

让复制和计算序列化而不是重叠

这也是为什么这些改进在概念上组合。随着主内核变得更快，曾经不可见的开销开始成为总步时间的有意义部分。

这里有一个有用的工程课程。一旦数学内核被优化，"更快"通常意味着以下两件事之一：

做更少不必要的工作

使不可避免的工作并行发生

这正是这里发生的事。

使用 UI 训练和运行 LLM