在 AI Runtime 上将 Qwen 训练时间缩短 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，其中包含如何自行运行该模型的详细信息。