Databricks 讲解如何在 AI Runtime 上将 Qwen 3.8 27B 微调时间缩短 6 倍

训练或微调一个数十亿参数的基础模型需要大量算力，通常需要数百到数千小时的 GPU 训练。但调优 LLM 的训练性能颇具挑战，而且诸如 GPU 利用率之类的高层指标可能并未反映你实际的利用率。在这篇博文中，我们将展示如何穿透层层抽象，在 AI Runtime 上将最新 Qwen 3.8 27B 模型的微调速度提升 6 倍。通过系统性的性能优化，即便 GPU 利用率已经达到 98%，我们仍将固定数据集的微调时间从 7 小时 44 分钟缩短到 1 小时 17 分钟。

我们将深入剖析这些改进背后的技术，包括如何衡量真正重要的指标、实现细节，以及每项优化所涉及的设计权衡。

如果你想立即上手，直接跳到 AI Runtime 上的示例代码试一试吧！

模型设置

我们从使用 transformers 库和 SFT 训练数据集 HuggingFaceH4/ultrachat_200k 的 Qwen 3.8 27B 标准训练循环开始。

我们应用了 FSDP2，这是来自 torch 的最新模型分片技术。这是将模型权重和优化器状态装入 GPU 所必需的。纯分布式数据并行方法在这个规模下无法将模型权重装入单个 GPU。
我们还启用了激活检查点，这是一种在反向传播过程中用内存空间换取计算量的技术，以允许每个 GPU 使用更大的批次大小。
数据加载器将 2D 批次打包，每批最大 token 数为 8192，并填充到批次示例的最大尺寸。

初始训练基准

我们在完整数据集上对模型进行基线训练耗时 7 小时 44 分钟，训练期间 GPU 显示平均利用率为 98%。结果是 9,396 tok/s 的训练速度。

这个 GPU 利用率看起来很高，但如前所述，这个数字具有误导性。标准 GPU 利用率指标只统计 GPU 在做任何事情的时间百分比，因此它并不能告诉我们模型在 GPU 活跃时对其利用得有多好。要理解这一点，我们需要更深入地观察。

一种方法是 torch trace 机制，我们可以像下面这样调整训练循环（同时限制步数以便快速得到结果）：

对于一个样本步骤，我们得到以下数据：

操作	调用次数	平均时间（微秒）	总时间（毫秒）
aten::bmm	31,105	50.5	1,571.80
aten::mul	46,878	16.2	758.9
MmBackward0	497	1,111.4	552.4
aten::copy_	69,444	7.2	500.6
aten::add	45,623	8.0	366.7
aten::select	68,645	3.4	230.5
aten::exp	9,504	21.7	206.5
aten::sum	16,194	12.2	196.8
aten::empty	40,709	4.7	190.5
aten::as_strided	218,214	0.8	173.3

表 1：CPU - 14.4 秒，1,714,382 次调用。表中包含按总时间排序的顶级操作。

操作	调用次数	平均时间（微秒）	总时间（毫秒）
nvjet GEMM	1923	1,163.1	2,236.8
逐元素浮点加法	42,657	51.1	2,180.5
逐元素浮点乘法	35,012	24	841.8
逐元素浮点填充	33,557	21.5	723.1
逐元素浮点复制	59,282	11.8	700.7
Memcpy DtoD	5,585	90.5	505.2
sm80_xmma GEMM	31,105	20.9	651.3

表 2：GPU - 10.3 秒，269,448 次操作。表中包含按总时间排序的顶级操作

可改进之处

从这些表中，我们注意到当前模型存在几个高层问题：

CPU 设置每个训练步骤所需的时间比 GPU 运行该步骤的时间更长。这出乎意料，因为通常 GPU 执行的 CPU 侧是异步的，也就是说 CPU 应该远远领先于 GPU运行。其主要影响是 GPU 出现停滞，因此我们有一些改进的机会。
There are too many operations! In the tables above, we have hundreds of thousands of small operations, in this case the issue is that the CPU cannot enqueue GPU operations fast enough to keep the GPU busy. The key idea is that we need to reduce the number of operations the CPU is enqueuing to the GPU. There are some general techniques for this, one fits our situation:
- torch 的 CUDA graphs 会记录所见的 cuda kernel，并通过一次 CPU 调用重放它们。虽然功能强大，但存在许多限制，并且与 FSDP2 的动态性相冲突。所以目前暂不考虑。
- Torch.compile 会追踪代码，并在可能的情况下使用 Triton 自动融合 kernel。在本例中，我们将使用手动融合的 kernel，而对于 LLM 来说，这不会给 torch.compile 留下任何改进空间。
- 手动融合的 kernel 针对特定运算使用专门的融合技术，这些是 torch 的编译器无法自行应用的。幸运的是，社区中已经有许多成熟的预构建 kernel。这方面的例子包括 ‘Flash Attention’，在本例中它带来了小幅性能提升。

此外，我们在批次中存在过度填充的问题：输入张量必须是矩形的，每一行都会被填充到该批次的最大 token 数。我们将展示一个可用于 LLM 的技巧来尽量减少这种情况。

在接下来的章节中，我们将深入探讨每个技术主题，涵盖它是什么、推荐的最佳实践，以及应用它时需要考虑的关键权衡。

融合 kernel

你可能会注意到某些 kernel 出现得非常频繁！特别是对于 Qwen，批量矩阵乘法 kernel 在一个步骤中运行超过 30,000 次。

在本例中，我们将首先探讨融合 kernel。融合 kernel 将多个较小的 GPU kernel 组合成一个单元。这可以通过避免 GPU 的高带宽内存与其更高带宽的片上内存之间代价高昂的传输来帮助提升性能。此外，它还可以实现更高效的 GPU 核心调度。

有了这么多额外的 kernel，虽然 torch.compile 可能可行，但运行起来会花费很长时间。因此，使用手动融合的 kernel 来减少原始算子数量是很好的第一步。此外，像 Flash Attention 这样众所周知的融合并不是 torch.compile 自身能够发现的。

transformers 库默认对标准注意力层使用 torch 的 SDPA。但事实证明，我们看到的许多小运算来自非标准的门控 delta 网络层。幸运的是，对于这些层，有两个 kernel——causal-conv1d 和 flash-linear-attention——它们被编写用于融合每层中的许多此类运算。我们如何发现这一点？Qwen 的 transformers 库实现已经告诉我们它支持哪些 kernel：

获取这些 kernel 有两种高层次方法：

Install the kernels yourself. This can range from:
- 从 PyPi 简单 pip install（例如 flash-linear-attention）
- git clone 并编译相关仓库
- 如果编译耗时过长，寻找其他来源的预编译包
使用 Hugging Face 的 Kernel Hub 获取预编译 kernel。

我们通常推荐第一种，尤其是配合我们的 docker 容器支持，在构建时将所有 kernel 安装到容器中，以减少作业启动所需的时间。

虽然 kernelhub 很方便，但由于 kernel 通常明确依赖于所安装的 torch 和 cuda 版本，你的特定配置可能没有可用的 kernel，而 transformers 模型会静默回退到未优化、性能低下的 kernel。

在安装了 flash-linear-attention 和 causal-conv1d 这两个 Python 包之后，我们将端到端时间提升了 3.7 倍，降至 2 小时 5 分钟。在下表中，你可以看到新的内核（已高亮显示）取代了之前表中的逐元素函数。同时，执行的操作总数也大幅减少了 15 倍：

操作	调用次数	平均时间（微秒）	总时间（毫秒）
cudaMemsetAsync	1,854	283.9	526.3
ChunkGatedDeltaRuleFunction	96	871.7	83.7
aten::mm	1,987	39.6	78.6
MmBackward0	497	134.2	66.7
ChunkGatedDeltaRuleFunctionBackward	48	1,159.9	55.7
DaoAILab::_causal_conv1d_fwd_cpp	96	55.8	5.4
CausalConv1dFn	96	51.4	4.9

表 3：CPU - 3.9 秒，共 118,316 次调用。表中包含按总时间排序的主要操作。

操作	调用次数	平均时间（微秒）	总时间（毫秒）
nvjet GEMM	1,923	1,183.3	2,275.5
逐元素浮点乘法	2184	128.1	279.9
逐元素复制 + 类型转换	1255	101.1	126.9
chunk_cat flash-linear-attention	65	1,402.0	91.2
flash_attention_3 前向	16	2,693.8	43.1
chunk_gated_delta_rule_fwd	144	255.6	36.8
flash_attention_3 反向	32	965.6	30.9

表 4：GPU - 3.9 秒，共 18,060 次调用。表中包含按总时间排序的主要操作。

注意事项：在训练期间，你需要同时拥有前向和反向传播的内核，而有些内核仅支持前向，因为它们只是为了加速推理而实现的。

要点总结

考虑使用针对常见模型现成可用的、经过充分验证且高度优化的内核。利用这些内核可以以相对较小的努力带来显著的性能提升。在这个具体案例中，我们看到了训练速度提升了 3.7 倍，将训练时间从7 小时 44 分钟缩短至 2 小时 5 分钟。

批处理、分词与填充

我们的 Qwen 模型现在相当快了，但我们还能更快。默认情况下，我们的批处理设置使用 2D 矩形张量，大小为 8,192 个 token，因为这是在不耗尽内存的情况下能容纳的最大尺寸。这种方法的缺点是，我们往往会在张量中因填充而浪费大量空间。

如果 Qwen 只使用标准注意力机制，那么 torch 的 NestedTensor 现在已得到良好支持，可以通过完全消除填充来帮助降低这一成本。但 flash-linear-attention 只支持更传统的 1D 打包张量方法：

但我们不能随意这样做，简单地将 1D 张量传递给模型对于预训练来说可能可以接受，但对于微调，我们需要确保在一个示例中，token 只关注同一示例内的 token。为了解决这个问题，我们必须向模型传递正确的示例 position_ids。模型会将这些传递给融合注意力内核，并应用正确的注意力掩码。如果你查看这里的示例代码，你可以看到我们如何计算所需的输入来实现这一点。

通过这个看似简单的改动，我们的训练时间从 2 小时 5 分钟缩短到了 1 小时 17 分钟。

原因有两个：

主要原因是这种打包增加了每批的序列数量，并将总批次数减少了约 36%。请注意，这取决于数据集中示例的分布和长度。
次要原因是 flash linear attention 内核需要针对每个输入形状进行即时编译，而之前的形状是动态的，导致多花了几分钟进行编译。

权衡：这种改动确实会让输入流水线更复杂，而且很有可能引入影响模型质量的错误。但在这种情况下，36% 的节省是值得的。还有其他减少填充的技术，比如保留原有的批处理机制，但按长度对所有输入排序，从而将填充降到最低，不过这种顺序可能会在相邻样本中引入意外的相关性，使模型产生偏差。

要点总结：当在输入短于你可训练的最大 token 数（例如本例中的 8192）的数据集上训练 LLM 时，高效地将数据打包成批次可以节省大量时间。对于这个数据集，我们看到了训练速度提升 1.6 倍，训练时间从2 小时 5 分钟缩短到 1 小时 17 分钟。

结论

在这篇博客中，我们向你展示了如何将训练效率提升 6 倍。归结起来就是：(1) 首先衡量真正重要的指标，而不仅仅是 GPU 利用率，还包括你的模型在训练中把时间花在哪里；(2) 使用 Flash Linear Attention 等已知库来优化和融合内核；(3) 调整数据批次送入模型的方式。

要在规模化场景下做到这一点，你需要一个让高性能路径成为默认选择的平台。在 AIR 中，许多这类性能改进都是内置的。我们在基础环境中打包了最新的 CUDA 和 flash attention，自动为你跟踪 GPU 利用率、温度和健康状态等核心指标，并确保 torch 版本（以及分析器）保持最新。我们的内部智能体 Genie Code 也经过调优，能够寻找这类机会来进一步优化你的训练。此外，从 Databricks 读取数据也通过我们的 UCVolumeDataset 得到了优化。

请在你的下一个训练工作负载中试试 AIR！准备好尝试了吗？请参阅 AI Runtime cookbook，其中包含如何自行运行此模型的详细信息。