98% 占用并没有说明已经跑满

Databricks 10 月 14 日的工程博文写,他们把 Qwen 3.8 27B 的微调加快大约 6 倍。站点 RSS 的时间是世界时 11:30,北京时间 19:30。基线在完整数据集上要 7 小时 44 分,训练时 GPU 平均占用 98%,吞吐每秒 9396 个 token。

博文的判断是:占用率只表示 GPU 有没有在做事,不表示算得满。他们看到 CPU 排布一步的时间比 GPU 真正计算还长,小操作有几十万次,CPU 来不及把工作交给 GPU。

[1]
一块炭黑色的扁方块。
一块炭黑色的扁方块,侧面能看见厚度。, AI 生成插画,不是新闻照片

先换现成的融合核

训练循环用 transformers,监督微调数据是 HuggingFaceH4/ultrachat_200k。权重和优化器状态用 FSDP2 摊到多张卡上,并打开激活检查点,用计算换显存,好让每张卡的批次大一点。批次是二维的,每批最多 8192 个 token,短样本要补到这一批里最长的那条。

默认注意力走 torch 的 SDPA。博文说,大量小操作来自非标准的 gated delta 层。他们装上 flash-linear-attention 和 causal-conv1d 之后,端到端时间降到 2 小时 5 分,大约 3.7 倍。一步里的操作次数大约少到原来的十五分之一。博文还写,只做推理的核不够,训练要前向和反向都有。

[1]

再把补齐的空白拿掉

二维矩形批次会把短样本补满。flash-linear-attention 不吃嵌套张量,他们改成一维打包,并传入每条样本自己的 position_ids,让 token 只看同一条样本。时间从 2 小时 5 分降到 1 小时 17 分。

主因是每批能放进更多序列,总批次数大约少 36%,少多少取决于样本有多长。次因是融合核以前要按变化的形状即时编译,打包后形状稳定,少了几分钟编译。两步合在一起,就是博文说的大约 6 倍。

[1]

6 倍只覆盖这一次微调

博文自己写,这种改法让输入管线更复杂,很有可能引入影响模型质量的错误;他们认为这次少掉的 36% 时间值得换。质量有没有掉,这篇没有给损失或评测对照。

也没有写这套倍数会不会出现在别的模型、别的数据集或预训练上。基线占用已经 98% 还能再快,说明那个 98% 不是“余量已经没了”。

[1]