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))以某一个训练步为例，可以得到下面这些数据：
|
|
|
|
|
|---|---|---|---|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
表1：CPU，总耗时14.4秒，共1,714,382次调用。表中列出了按总耗时排序的主要操作。
|
|
|
|
|
|---|---|---|---|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
表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分钟。下表中高亮标出的新内核，替代了前面表格里的逐元素函数。执行的操作总数也大幅下降，减少至原来的约十五分之一：
|
|
|
|
|
|---|---|---|---|
|
|
|
|
|
| ChunkGatedDeltaRuleFunction | 96 | 871.7 | 83.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次调用。表中列出了按总耗时排序的主要操作。
|
|
|
|
|
|---|---|---|---|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| chunk_cat flash-linear-attention | 65 | 1,402.0 | 91.2 |
|
|
|
|
|
| chunk_gated_delta_rule_fwd | 144 | 255.6 | 36.8 |
|
|
|
|
|
表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分钟。
原因有两个：
-
主要原因是，这种打包方式增加了每个批次中的序列数量，使总批次数减少了约36%。需要注意，这个比例取决于数据集中样本的长度及其分布。 -
次要原因是，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