Databricks Cuts Qwen 3.8 27B Fine-Tuning Time 6x on AI Runtime

A high utilization number that hid the problem

Databricks starts from a standard training loop for Qwen 3.8 27B built with the transformers library and the SFT dataset HuggingFaceH4/ultrachat_200k. The model is sharded with FSDP2, the latest sharding technique from torch, which is needed to fit the weights and optimizer state onto the GPUs. Activation checkpointing is enabled so that a larger batch size per GPU is possible. The dataloader packs 2D batches with a maximum of 8192 tokens per batch and pads to the largest example in each batch.

This baseline needed 7 hours and 44 minutes to train over the complete dataset and reached 9,396 tok/s, while the GPU showed 98% average utilization. Databricks argues that the figure is misleading. The standard metric, the company says, "only counts the percentage of the time the GPU was doing anything", so it says little about how well the GPU is used while it is active.

What the torch trace showed

To look deeper, Databricks used the torch trace mechanism on a limited number of steps. For a sample step, the CPU side took 14.4 seconds with 1,714,382 calls, while the GPU side took 10.3 seconds with 269,448 operations. That is the reverse of what is expected, since the CPU normally runs well ahead of the GPU. The result is that the GPU stalls waiting for work.

The company identified the cause as too many small operations: the CPU cannot enqueue them fast enough to keep the GPU busy. The goal therefore became reducing the number of operations. Databricks considered several general techniques:

• CUDA graphs: they record kernels and replay them with a single CPU call, but have many limitations and conflict with the dynamism of FSDP2, so they were ruled out.
• torch.compile: it fuses kernels automatically with Triton, but Databricks chose manually fused kernels, which it says leave no room for improvement from torch.compile in the case of LLMs.
• Manually fused kernels: specialized fusions that torch's compiler cannot apply on its own, many of which already exist in the community.

A separate issue was excess padding, because each row of the rectangular input tensor is padded to the maximum token count in its batch. Databricks says it uses a trick for LLMs to minimize this.

Fused kernels deliver the first big gain

In the trace, the batch matrix multiplication kernel ran more than 30,000 times in one step. Databricks explains that a fused kernel combines several smaller GPU kernels into one unit, which avoids costly transfers between high bandwidth memory and the faster on-die memory and can allow more efficient scheduling of GPU cores.

By default, transformers uses torch's SDPA for standard attention layers. According to the post, many of the small operations came from the non-standard gated delta network layers in Qwen. For these layers, two kernels exist that fuse many operations: causal-conv1d and flash-linear-attention. The transformers implementation of Qwen itself lists which kernels it supports.

There are two ways to obtain them: install them yourself, for example with a pip install, a git clone and compile, or a precompiled package, or fetch precompiled kernels from Hugging Face's Kernel Hub. Databricks generally recommends the first, combined with its docker container support, so that kernels are installed at build time and jobs start faster. It warns that Kernel Hub may lack a kernel matching your torch and CUDA versions, in which case the transformers model silently falls back to an unoptimized kernel with low performance.

After installing flash-linear-attention and causal-conv1d, end-to-end time improved by a factor of 3.7x, down to 2 hours 5 minutes. The total number of operations fell by a factor of 15x. In the follow-up trace, the CPU side took 3.9 seconds with 118,316 calls, and the GPU side took 3.9 seconds with 18,060 calls, so the two sides were now balanced. New entries such as flash_attention_3 and chunk_gated_delta_rule_fwd appear among the top GPU operations.

What it means for teams fine-tuning large models

The main message of the post is that a high GPU utilization figure does not prove efficient training. Databricks says AI Runtime lets users diagnose performance issues with a standard torch trace of an LLM, and it points to example code that readers can try on the platform. The full path from 2 hours 5 minutes to the final 1 hour 17 minutes comes, per the company's summary, from further data pipeline improvements, which the excerpt of the post available here does not detail.

For teams that spend hundreds to thousands of GPU hours on fine-tuning, the practical lesson is to check which kernels the model library supports, make sure they are actually installed for the exact torch and CUDA versions in use, and watch for silent fallbacks. These figures are Databricks' own measurements on a fixed dataset with a specific setup, and results will vary with other models and configurations.