Databricks：如何让 Qwen 3.8 27B 微调快 6 倍

GPU 平均利用率已经达到98%，训练却还能快6倍。Databricks 这篇博文写下了排查过程、算子变化和数据打包细节。以下为全文翻译，文中的“我们”指原作者 Bruce Fontaine 和 Pardis Beikzadeh。

让大模型训练发挥最佳性能的两个简单方法。

训练或微调一个数十亿参数的基础模型，需要大量计算资源，往往要耗费数百乃至数千个 GPU 小时。但优化大模型的训练性能并不容易，GPU 利用率这样的高层指标，也未必能反映 GPU 实际用得有多充分。本文将介绍如何穿透这些抽象层，在 AI Runtime 上将最新 Qwen 3.8 27B 模型的微调提速6倍。通过系统性的性能优化，即使原来的 GPU 利用率已经达到98%，我们仍将固定数据集上的微调耗时从7小时44分钟缩短到了1小时17分钟。

我们会深入介绍这些改进背后的方法：应该测量什么、如何实现，以及每项优化涉及哪些设计取舍。

如果你想马上动手，可以直接查看 AI Runtime 上的示例代码，自己试一试！

模型设置

我们从一个标准训练循环开始，使用 transformers 库和监督微调数据集 HuggingFaceH4/ultrachat_200k，训练 Qwen 3.8 27B。

# Create the Dataloader
data_loader = make_dataloader("HuggingFaceH4/ultrachat_200k")
model = AutoModelForImageTextToText.from_pretrained(
    "Qwen/Qwen3.8-27B", torch_dtype=torch.bfloat16)
# Enabled gradient checkpointing to reduce memory requirements
model.gradient_checkpointing_enable(
    gradient_checkpointing_kwargs={"use_reentrant": False})
# Apply FSDP2
mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)
layers = model.model.language_model.layers
for layer in layers:
  fully_shard(layer, mp_policy=mp)
fully_shard(model, mp_policy=mp)
model.cuda()
opt = torch.optim.AdamW(
    model.parameters(), lr=1e-5, betas=(0.9, 0.95), weight_decay=0.0)
model.train()
for batch in data_loader:
  out = model(input_ids=batch["inputs"], labels=batch["labels"])
  out.loss.backward()
  opt.step()
  opt.zero_grad(set_to_none=True)

• 我们使用了 FSDP2，这是 torch 最新的模型分片技术。要让模型权重和优化器状态装进 GPU，就需要这种技术。在这个模型规模下，单纯使用分布式数据并行，无法把模型权重放进单张 GPU。

• 我们还开启了激活检查点。这种技术会在反向传播时用额外计算换取内存空间，从而增大每张 GPU 可以承载的 batch size。

• 数据加载器将数据组成二维批次，每批最大 token 数为8192，并将批内样本补齐到最长样本的长度。

初始训练基准

基线模型在完整数据集上训练需要7小时44分钟，训练期间 GPU 的平均利用率显示为98%。训练吞吐量为每秒9,396个 token。

这个 GPU 利用率看起来很高，但正如前面提到的，它具有误导性。标准的 GPU 利用率指标，只统计 GPU 处于工作状态的时间占比，因此无法告诉我们：GPU 活跃时，模型究竟把它用得有多充分。要弄清楚这一点，需要继续往下看。

一种办法是使用 torch 的追踪机制。我们可以像下面这样调整训练循环，同时限制步数，以便尽快得到结果：

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,  # Remove if running only on CPU
    ],
    schedule=torch.profiler.schedule(
        wait=1,      # Step 1: Skip tracing to let things warm up
        warmup=0,    # Skip warmup for this specific single-step trace
        active=1,    # Step 2: Trace exactly 1 step
        repeat=1
    )) as prof:
  <training code here>
print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=10))

以某一个训练步为例，可以得到下面这些数据：

操作	调用次数	平均耗时（微秒）	总耗时（毫秒）
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 会停下来等待，因此还有改进空间。

• 操作太多了！上面的表格里有数十万个小操作。这里的问题在于，CPU 无法以足够快的速度将 GPU 操作加入队列，让 GPU 持续保持忙碌。关键思路是减少 CPU 向 GPU 提交的操作数量。有几种通用方法可以考虑，其中一种适合我们这次的情况：

  - torch 的 CUDA graphs 可以记录执行过的 CUDA 内核，再通过一次 CPU 调用重放它们。它很强大，但也有不少限制，而且与 FSDP2 的动态行为存在冲突，所以我们暂时不采用。   

  - torch.compile 会追踪代码，并在可能的地方使用 Triton 自动融合内核。本例中我们将使用手动融合内核；对于这些大模型操作，手动融合后就没有给 torch.compile 留下进一步优化的空间。   

  - 手动融合内核针对特定操作使用专门的融合技术，这些技术是 torch 编译器无法自行应用的。好在社区里已经有不少成熟的预构建内核，例如 Flash Attention。它在本例中带来了一点性能提升。

此外，批次中还有 padding 过多的问题：输入张量必须是矩形，每一行都要补齐到该批次的最大 token 数。接下来我们会介绍一种适用于大模型的技巧，尽量减少这部分浪费。

下面将逐一深入这些技术，介绍它们是什么、推荐的实践方式，以及应用时需要考虑的主要取舍。

融合内核

你可能注意到了，有些内核出现的次数特别多。以 Qwen 为例，批量矩阵乘法内核在一个训练步中就运行了超过30,000次。

我们先研究融合内核。融合内核会把多个较小的 GPU 内核合成一个执行单元。这样可以避免 GPU 高带宽显存与带宽更高的片上存储之间代价高昂的数据搬运，从而提高性能；同时，也能让 GPU 核心的调度更加高效。

面对这么多额外内核，即使 torch.compile 能起作用，运行它也要花很长时间。因此，先用手动融合内核减少操作总数，是一个不错的起点。而且，Flash Attention 这类成熟的融合方案，也不是 torch.compile 自己能够发现的。

对于标准注意力层，transformers 库默认使用 torch 的缩放点积注意力（SDPA）。不过，我们看到的许多小操作，其实来自非标准的门控 Delta 网络层。幸运的是，针对这些层，已经有 causal-conv1d 和 flash-linear-attention 两种内核，可以融合每一层里的许多操作。我们是怎么发现它们的？transformers 库中 Qwen 的实现，已经写明了支持哪些内核：

• flash-linear-attention

• causal-conv1d

获取这些内核，主要有两种方式：

• 自行安装。具体方式包括：

  - 直接通过 pip 从 PyPI 安装，例如 flash-linear-attention。   - 使用 git clone 克隆对应仓库，然后编译。   

  - 如果编译耗时太长，也可以寻找其他提供预编译包的来源。

• 使用 Hugging Face 的 Kernel Hub 获取预编译内核。

我们通常推荐第一种方式，尤其是结合我们的 Docker 容器支持：在构建容器时就装好所有内核，以减少训练任务的启动时间。

Kernel Hub 虽然方便，但内核往往明确依赖所安装的 torch 和 CUDA 版本，因此未必有适配你当前配置的可用内核。这时，transformers 模型会静默回退到未经优化、性能较低的内核。

安装 flash-linear-attention 和 causal-conv1d 两个 Python 包之后，我们的端到端训练提速了3.7倍，耗时降至2小时5分钟。下表中高亮标出的新内核，替代了前面表格里的逐元素函数。执行的操作总数也大幅下降，减少至原来的约十五分之一：

操作	调用次数	平均耗时（微秒）	总耗时（毫秒）
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 forwards	16	2,693.8	43.1
chunk_gated_delta_rule_fwd	144	255.6	36.8
flash_attention_3 backwards	32	965.6	30.9

表4：GPU，总耗时3.9秒，共18,060次调用。表中列出了按总耗时排序的主要操作。

注意：训练时，前向和反向传播都需要对应内核。有些内核只支持前向计算，因为它们的实现仅用于加速推理。

要点

• 可以考虑使用常见模型已有的成熟、高度优化的内核。借助这些现成内核，通常只需较少工作，就能获得显著的性能提升。此次训练速度提升了3.7倍，训练耗时从7小时44分钟降至2小时5分钟。

批处理、分词与填充

现在 Qwen 模型已经相当快了，但我们还能再快一些。默认的批处理方式采用二维矩形张量，长度为8,192个 token，因为这是不发生显存不足时能够容纳的最大长度。缺点是，padding 往往会浪费张量中相当多的空间。

如果 Qwen 只使用标准注意力机制，torch 现在已经提供了较完善的 NestedTensor 支持，可以完全移除 padding，从而降低这部分开销。但 flash-linear-attention 只支持更传统的一维打包张量方式：

不过，我们不能随意这样做。预训练时，直接把一维张量交给模型或许可以接受；但微调时，必须确保每个样本中的 token 只关注同一个样本内的 token。为此，需要将正确的样本 position_ids 传给模型。模型再把这些信息传给融合注意力内核，由内核应用正确的注意力掩码。在示例代码中，可以看到我们如何计算所需输入，让这套方式正常工作。

经过这个看起来很简单的改动，训练耗时从2小时5分钟进一步缩短到了1小时17分钟。

原因有两个：

1. 主要原因是，这种打包方式增加了每个批次中的序列数量，使总批次数减少了约36%。需要注意，这个比例取决于数据集中样本的长度及其分布。

2. 次要原因是，Flash Linear Attention 内核需要针对每一种输入形状进行即时编译。此前输入形状是动态变化的，因此额外增加了几分钟编译时间。

取舍在于：这种改动会让输入流水线更加复杂，而且很容易引入影响模型质量的错误。但在本例中，节省36%的批次值得这样做。还有其他减少 padding 的办法，例如保留原来的批处理方式，先按长度对所有输入排序，尽量减少填充。不过，这种排序可能让相邻样本之间出现意料之外的相关性，进而使模型产生偏差。

要点：当训练数据集中的输入短于模型可训练的最大 token 长度时，例如本例中的8192，高效地将数据打包成批次，可以节省大量时间。在这个数据集上，我们获得了1.6倍的训练提速，耗时从2小时5分钟降至1小时17分钟。

结论

本文介绍了如何将训练效率提升至6倍。归结起来有三点：（1）先测量真正重要的东西，不只看 GPU 利用率，还要看训练时间究竟花在了哪里；（2）借助 Flash Linear Attention 等成熟库优化并融合内核；（3）针对模型调整数据批次的输入方式。

要规模化地做到这些，需要一个能把高性能路径变成默认选择的平台。AI Runtime（AIR）已经内置了其中许多改进：基础环境打包了最新的 CUDA 和 Flash Attention；自动跟踪 GPU 利用率、温度及健康状况等核心指标；确保 torch 和性能分析工具保持更新。我们的内部 Agent——Genie Code，也专门调优过，会寻找这类机会，进一步优化训练。此外，UCVolumeDataset 还优化了从 Databricks 读取数据的过程。

欢迎在下一次训练任务中试试 AIR！想开始使用，可以查看 AI Runtime cookbook，其中详细介绍了如何自行运行这个模型。

参考链接

• 原文：https://www.databricks.com/blog/reducing-qwen-training-time-6x-ai-runtime

• AI Runtime 示例代码：https://github.com/databricks-solutions/ai-runtime-cookbook/tree/fd81d9866bf9b35522afbbd2387b68e91958958f/post-training/sft-qwen3.8-27b-fsdp2-multinode

• UltraChat 200k 数据集：https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k

• FSDP2 文档：https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html

• CUDA graphs：https://pytorch.org/blog/accelerating-pytorch-with-cuda-graphs/

• torch.compile 文档：https://docs.pytorch.org/docs/2.14/user_guide/torch_compiler/compile/programming_model.html

• Flash Attention 3：https://pytorch.org/blog/flashattention-3/

• 缩放点积注意力文档：https://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html

• transformers 中的 Qwen 实现：https://github.com/huggingface/transformers/blob/a10365f95226614b2e9b9501e24bba36088c6b4f/src/transformers/models/qwen3_5/modular_qwen3_5.py

• Flash Linear Attention：https://github.com/fla-org/flash-linear-attention

• causal-conv1d：https://github.com/Dao-AILab/causal-conv1d

• AI Runtime Docker 容器文档：https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/docker-images

• NestedTensor 文档：https://docs.pytorch.org/docs/2.14/nested.html

• 样本打包与训练代码：https://github.com/databricks-solutions/ai-runtime-cookbook/blob/b27c10cfda94263588214f8824e0dd8bb393371e/training/sft-qwen3.8-27b-fsdp2-multinode/train.py

• Databricks 数据加载优化：https://www.databricks.com/blog/fast-fault-tolerant-pytorch-training-ai-runtime

• AI Runtime 示例代码：https://github.com/databricks-solutions/ai-runtime-cookbook/tree/fd81d9866bf9b35522afbbd2387b68e91958958f/post-training/sft-qwen3.8-27b-fsdp2-multinode