Databricks 分享把 Qwen 3.8 27B 微调提速 6 倍的调优全过程

把 Qwen 微调提速 6 倍的实战经验分享

@databricks 分享了在自家 AI Runtime 上微调 Qwen 3.8 27B 的调优全过程，从 7 小时 44 分压到 1 小时 17 分，提速 6 倍。堪称教科书级的性能工程案例：用 torch profiler 找到真正的瓶颈（CPU 调度过载，而非 GPU 不够忙），再用两步常规手段解决它（融合内核 + 序列打包）！

# 起点：一个"看起来很健康"的慢基线

任务设定：用 HuggingFaceH4/ultrachat_200k 数据集对 Qwen 3.8 27B 做 SFT，采用 FSDP2 分片（27B 权重单卡放不下）+ 激活检查点，按最大 8192 token/batch 组 batch。

基线跑了 7h44m，GPU 平均利用率 98%；按常规监控逻辑这几乎无可挑剔。但有一个普遍误区：GPU 利用率只衡量 "GPU 是否在做事"，不衡量 "GPU 做得有多好"。

它无法区分高效执行和"把时间浪费在几万次微小 elementwise 操作上"。正确的诊断工具是 torch trace / profiler。

# 诊断：瓶颈在 CPU，不在 GPU

Profiler 结果显示一个反直觉的事实：CPU 每步耗时（14.4s）超过了 GPU（10.3s），CPU 忙于往 GPU 队列里塞操作，根本喂不满。单个训练 step 里 CPU 发起了约 171 万次调用（如 bmm 31,105 次、mul 46,878 次），GPU 侧也有 4 万多次琐碎的 elementwise add。

根源在于 Qwen 3.8 的门控 Delta 网络（gated delta net）层；这类线性注意力/卷积混合层在 transformers 的朴素实现里会被展开成海量小算子。也就是说，慢不是模型本身的问题，是参考实现没有针对训练做算子级融合。

Databricks 评估了三条路： · CUDA Graphs：否，与 FSDP2 的动态性冲突 · torch.compile：否，编译太慢，且会漏掉关键融合（如 Flash Attention） · 手动接入已知融合内核：是，选这条路

# 优化一：接入融合内核 → 3.7 倍

transformers 的实现其实标注了对两个融合内核的支持，安装后即生效： · flash-linear-attention（PyPI）覆盖 gated delta rule 层； · causal-conv1d（Dao-AILab）覆盖因果卷积。

效果立竿见影：7h44m → 2h05m（3.7x），每步算子数下降约 15 倍（CPU 调用从 171 万降到 11.8 万），trace 里最耗时的条目变成了 ChunkGatedDeltaRule、flash_attention_3 前反向和 causal-conv1d 内核。

两个值得记住的工程细节： · 不要依赖 HF Kernel Hub 动态加载内核：torch/CUDA 版本不匹配时它会静默回退到未优化路径，你以为什么都没变，其实优化已经失效。推荐直接安装，最好在 Docker 镜像构建时固化。 · 部分内核只实现了 forward，推理可用但训练不可用：选型时必须核对。

# 优化二：1D 序列打包去 padding → 再 1.6 倍

剩余瓶颈是数据侧的浪费：输入张量必须是矩形，导致每个 batch 内短序列被 padding 到最长序列，2D 矩形里大量格子是无效计算。

理想方案 NestedTensor 只兼容标准注意力；而 flash-linear-attention 要求 1D 打包（concat + position_ids） 格式。作者的做法是把 batch 内多个样本拼接成一维，并为每个样本传入正确的 position_ids，使注意力 mask 严格限制在样本内部，不发生跨样本泄漏（示例代码在 ai-runtime-cookbook 仓库）。

2h05m → 1h17m（1.6x），收益来自两头： · 序列打包使 batch 数减少约 36%（取决于长度分布）； · 输入 shape 从动态变为固定，flash-linear-attention 的 JIT 编译开销大幅下降。

也有客观代价：流水线复杂度上升，且 position_ids 一旦写错就是静默的质量 bug（训练照跑、模型变差）。备选方案"按长度排序组 batch"虽能减少 padding，但会引入样本顺序相关性带来的训练偏差。

# 方法论提炼

1. 测量要测对的东西：GPU 利用率会骗人，torch trace 不会； 2. 优先接入社区已验证的融合内核：flash-linear-attention、causal-conv1d、Flash Attention； 3. 消灭 padding 浪费：1D 打包 + 正确的 position_ids。

平台层面的配套：AI Runtime 打包了新版 CUDA/Flash Attention、保持 profiler 工具链更新、内置 GPU 指标，且 Genie Code 会在代码层主动寻找这类优化。