第44章 如何利用 Unsloth 与 NVIDIA 技术加速 LLM 训练

# 如何用 Unsloth 和 NVIDIA 让 LLM 训练更快 2026年5月6日 作者：Daniel、Michael、Mathew 和 Datta，感谢 NVIDIA 的协助 我们与 NVIDIA 合作，将 LLM 训练速度提升了约 25%。本文将详细拆解我们的实现方法。这些优化不影响精度，且是在 Unsloth 原有 2-5 倍加速之上的额外提升。新算法会在 RTX 笔记本、数据中心 GPU 和 DGX Spark 上自动启用，只需更新 Unsloth 即可获得最新改进。通过与 NVIDIA 合作，我们实现了： - 缓存 packed sequence 元数据，训练提速 14.3% - 双缓冲异步梯度检查点，提速 8% - 在 MoE 路由中使用 argsort 和 bincount，gpt-oss 训练提速 15% ## 1. 缓存 Packed-Sequence 元数据 假设我们有几条较短的样本。与其把它们全部 padding 到相同长度、在 padding token 上浪费算力，不如把它们拼接成一条更长的 packed sequence： 模型仍需要知道每个原始序列的起止位置。因此除了拼接后的 token，我们还需要携带序列元数据，包括： - 序列长度 - 累积偏移量（cu_seqlens） - 最大序列长度 - 由以上三项推导出的 attention 结构 关键在于：对固定的 packed batch 来说，这些元数据在每一层都是一样的。 如果把一个 packed batch 的边界信息记为： $$B = \{\text{lengths}, \text{cu\_seqlens}, \text{max\_seqlen}, \text{mask 结构}\}$$ 那么该次前向传播中的每个 transformer 层消费的都是同一个 B。 如果模型有 L 层，每层都重新构建或重新同步一次 B，并不是什么新工作，只是同样的信息被反复重建。 换句话说，正确的做法是：构建 B 一次，使用 L 次。 而浪费的做法是：构建 B + 构建 B + ⋯ + 构建 B（共 L 次）。 这里的开销主要不是额外的 FLOPs。某些路径会强制 device-to-host 同步，实际上形成一个 GPU-CPU 同步点。一旦这种同步出现在逐层执行的路径中，开销就会在每一层重复出现。 这正是 packed-sequence 缓存改动所消除的问题。它不再反复重建 packed sequence 信息、SDPA packed mask 和 xFormers block mask，而是把可复用的元数据以及从中推导出的 attention 侧结构按设备缓存下来，供当前 packed batch 使用。这些缓存结构随后在各层之间复用。 ### 为什么有效 Packed 训练本身已经通过消除 padding 浪费提升了利用率。但如果元数据路径不断强制同步，部分收益就会被与模型实际学习无关的开销吃掉。 缓存的好处在于把重复的协调工作从热路径中移除。前向传播受益最大，因为同一个 packed 元数据会在多层中被反复消费。 ### 基准测试 在 Qwen3-14B QLoRA SFT 上： - forward：+43.3% - backward：+5.8% - 每 batch：+14.3% 前向传播收益最大，因为重复的元数据和 mask 准备工作在这里体现得最直接。backward 也有提升，但幅度较小。节省的时间差不多，只是 backward（尤其是开启梯度检查点时）耗时更长，所以相对收益显得更小。 知道了实测收益之后，自然会问：这个量级合理吗？ ### 快速合理性验证 假设各层大致相似，可以把 packed-attention 路径建模为： $$T_{uncached} \approx L \cdot (A + s)$$ 其中： - L 是层数 - A 是每层有用的 attention 侧工作 - s 是每层重复的元数据和 mask 准备开销 启用缓存后，重复开销每个 batch 只付一次，而不是每层一次： $$T_{cached} \approx L \cdot A + s$$ 因此节省的时间约为： $$T_{saved} \approx (L - 1) \cdot s$$ 在 packed SDPA 路径上，我们在 NVIDIA Blackwell GPU 上的微基准测试显示，底层 host 可见的元数据调用开销真实但很小，每次约 0.2 ms。主要的重复开销来自 packed SDPA mask 构建路径本身，对一个总共 2048 个 token 的合成 packed batch 来说约 13.7 ms。 对 SDPA 后端来说，更合适的心智模型是： $$\text{小的 stream fence + mask 重建} \approx \text{mask 重建}$$ 这样就能做一次更干净的一致性检查。如果一次 packed-mask 重建耗时 m 毫秒，在均匀层模型下： $$T_{saved} \approx (L - 1) \cdot m$$ 取 m ≈ 13.7 ms，预测结果为： - 16 层：(16 - 1) × 13.7 ≈ 206 ms - 28 层：(28 - 1) × 13.7 ≈ 370 ms 较小的 packed-sequence 实验也呈现同样的规律： - Llama-3.2-1B（16 层）：每步约节省 199 ms，端到端步时降低约 11.5% - Qwen3-0.6B（28 层）：每步约节省 319 ms，端到端步时降低约 14.8% 这些百分比是相对完整训练步时间的，其中还包含 packed-attention 路径之外的工作，比如 embeddings、MLP、LM head、loss 和框架开销。这个估算刻意只针对 packed-attention 部分，而不是整个 transformer 层，目的只是验证实测收益处于合理范围。 ## 2. 用双缓冲检查点重载隐藏延迟 激活检查点是训练大模型的常用技术。思路是不在反向传播过程中保留所有中间激活，从而节省显存，代价是 backward 阶段要做一些额外工作。 这个权衡通常是值得的，尤其对更大的模型。 但它引出一个系统层面的问题：如果激活已被 offload，backward 时怎么把它送回 GPU？ 在 Unsloth 的智能检查点路径中，激活可以先存放在 pinned CPU 内存中，需要时再拷回。这节省了 VRAM，但可能引入瓶颈： 1. 把激活从 CPU 拷到 GPU 2. 等待拷贝完成 3. 对该激活执行 backward 计算 4. 开始下一次拷贝 这是一种串行模式。如果同一个 buffer 同时用于拷贝和计算，拷贝流和计算流就只能轮流执行。 设 $T_{copy}$ 为激活重载时间，$T_{compute}$ 为当前层的 backward 计算时间。 单缓冲时，这一步大致受限于： $$T_{single} \approx T_{copy} + T_{compute}$$ 这是串行情况，两笔开销几乎全额、先后支付。 更优雅的做法是用两个 buffer。当 backward 在 buffer A 上运行时，拷贝流可以提前把下一个激活加载到 buffer B，然后两者角色互换。这样就形成了流水线重叠——虽然不是完全重叠。 双缓冲并不减少计算量，它只是把拷贝延迟隐藏在有效计算背后。 ### 为什么有效 这种优化在模型足够大、backward 计算量可观、但又不至于大到让拷贝开销完全被淹没时效果最明显。模型越大，隐藏维度越高，数据搬运量越大，隐藏这些搬运的影响也越大。大模型层数通常也更多，隐藏拷贝的机会随之增加。 这就是为什么更大的 dense 模型更适合这项改进。GPU 上有足够的真实计算可供拷贝重叠，而第二个 buffer 的额外 VRAM 开销也很有限。 实现上还保留了务实的保护措施： - 仅在 VRAM 充足时使用额外 buffer - 内存紧张时干净地回退 - 正确性完全不变 ### 基准测试 在更大的 dense 模型上、使用 NVIDIA B200 Blackwell GPU 基准测试： - 8B：0.3739 → 0.4053 steps/s，+8.40% - 14B：0.2245 → 0.2395 steps/s，+6.70% - 32B：0.1979 → 0.2070 steps/s，+4.61% 显存开销保持在较低水平： - 8B：+0.37 GB - 14B：+0.47 GB - 32B：+0.23 GB 这些实验中，最终 loss 基本没有变化。 加速在较大 dense 模型上表现一致，额外 VRAM 成本也相对很小。 知道实测收益之后，自然的追问是：这个量级合理吗？ ### 快速合理性验证 假设有 L 个做了检查点的层，各层大致相似： - 每次重载耗时 c - 每段 backward 计算耗时 g 这也会随 batch size、序列长度以及其他影响数据搬运和计算的因素而变化，为简洁起见省略这些项。 单缓冲时： $$T_{single} \approx L \cdot (c + g)$$ 双缓冲时，第一层仍要等待激活到达，最后一层仍要完成计算，因此更好的近似是： $$T_{double} \approx c + (L - 1) \cdot \max(c, g) + g$$ 所以节省的时间约为： $$T_{saved} \approx (L - 1) \cdot \min(c, g)$$ 对这个结果的正确解读是： - 第一次拷贝仍然暴露在外 - 最后一次计算仍然暴露在外 - 但流水线中段的拷贝和计算可以重叠 如果重叠良好，中段每层的成本会接近： $$T_{middle} \approx \max(T_{copy}, T_{compute})$$ 从大模型的实测结果看，每个训练步节省的时间大约是： - 8B：约 207 ms - 14B：约 279 ms - 32B：约 222 ms 这些 host buffer 是 pinned 分配，因此相关带宽是 pinned 内存的 host-to-device 带宽，而非 pageable 内存带宽。在我们的 NVIDIA B200 Blackwell 系统上，该带宽约 55.7 GB/s，PCIe 理论上限 64 GB/s 可作为对比参考。 如果把额外 buffer 的大小粗略视为一次激活重载的数据量，那么每次重载本身就只有几毫秒量级： - 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 次 在几十个做了检查点的层中隐藏这种重载，正好对应几百毫秒量级的步时节省，与我们的观测完全吻合。 再次强调，这个节省时间是完整端到端训练步的一部分，并不打算解释 embeddings、LM head、loss、优化器工作或检查点之外的其他环节。这里的要点只是：可隐藏的数据传输规模足以合理解释实测的步时收益。 ## 3. 一个更小但实用的 MoE 优化 第三项改动更加专门化，但在 MoE 路由中体现了同样的模式。 在我们分析的基于 PyTorch 的 GPT-OSS MoE 路径中，路由中一个昂贵的环节是确定哪些 token 分给哪个 expert。朴素的实现可能写成： ```python for expert_idx in range(num_experts): token_idx, _ = torch.where(router_indices == expert_idx) ``` 乍看之下没什么问题。但这里的 torch.where 是数据依赖的：每个 expert 分到的 token 数量随 batch 变化。由于输出尺寸取决于路由模式，这可能引入 CPU-GPU 同步或相关运行时开销。如果每个 expert 都查询一次，动态查询的次数就随 num_experts 增长。 更好的做法是一次性分组： 1. 展平所有 expert 分配 2. 按 expert ID 稳定排序 3. 用一次 bincount 得到每个 expert 的 token 数 4. 根据这些计数构建偏移量 5. 按 expert 切分分组后的 token 列表 从数学角度看，收益并不在于改变了路由逻辑，而在于减少了向运行时询问动态索引问题的次数。 原来是： $$\text{动态查询开销} \propto \text{num\_experts}$$ （每个 expert 一次动态查询） 现在更接近： $$\text{动态查询开销} \propto 1$$ 之后只需廉价的簿记工作。 这是同一主题在更专门场景下的体现：一次分组，然后复用偏移量，而不是反复请求动态 token 列表。 ### 基准测试 注意，这些优化适用于任何使用 native_torch 后端的 MoE。 针对 GPT-OSS 的路由改进： - 团队验证显示 GPT-OSS 配置上有约 10-15% 的加速 - 在目标路由路径上，forward +23%，backward +13% ## 这些改动的共同点 尽管这三项优化位于技术栈的不同位置，它们解决的是同一个问题。 关键的优化机会都在主 kernel 周边的胶水代码中： - 重建已经存在的元数据 - 对本可以缓存的信息反复同步 - 让拷贝和计算串行执行而非并行重叠 这也是这些改进在概念上可以叠加的原因。随着主 kernel 越来越快，过去不可见的开销开始占据步时中可观的比重。 这里有一个实用的工程经验：一旦数学 kernel 已经优化到位，"更快"往往意味着以下二者之一： - 减少不必要的工作 - 让不可避免的工作并行发生 这正是我们在这次合作中做到的。 ## 总结 | 优化 | 消除的主要瓶颈 | 实测收益 | |---|---|---| | Packed-sequence 元数据缓存（PR link） | 各层重复的元数据重建与同步 | Qwen3-14B QLoRA SFT：forward +43.3%，backward +5.8%，每 batch +14.3% | | 双缓冲检查点重载（PR link） | 拷贝与 backward 计算在单 buffer 上串行 | 8B +8.4%，14B +6.7%，32B +4.6% | | 基于 bincount 的 GPT-OSS MoE 路由（PR link） | 逐 expert 动态索引带来的重复同步 | 团队验证约 10-15%，目标路径 forward +23%、backward +13% | 💕 谢谢大家！衷心感谢 NVIDIA 在开源工作中给予社区的帮助，以及对我们撰写本文的支持。也感谢大家阅读和使用 Unsloth，我们不胜感激。🙏

老规矩，欢迎加入我们的 Reddit 页面和 Discord 服务器，获取帮助或表达支持！你也可以在 Twitter 上关注我们，并订阅我们的 Substack 通讯。感谢阅读！Daniel & Michael Han 🦥
2026年5月6日 通过 UI 训练和运行 LLM，免费开始 加入我们的 Discord