Reducing Qwen training time by 6x on AI Runtime

Databricks demonstrated a sixfold reduction in fine-tuning time for the Qwen 3.8 27B model on AI Runtime, lowering the process from 7 hours 44 minutes to 1 hour 17 minutes. The optimization targeted GPU utilization metrics that masked inefficiencies, revealing that high GPU activity did not equate to effective compute usage. The team identified over 30,000 batch matrix multiplication kernel calls per step, indicating substantial overhead in standard training loops.

The first optimization involved manually fusing kernels to reduce the raw number of GPU operations, avoiding costly memory transfers and improving scheduling efficiency. This approach bypassed limitations of automated tools like torch.compile, which struggled with the scale of operations. Fused kernels combined smaller GPU tasks into single units, directly addressing the inefficiencies uncovered in the baseline training loop.

Further gains came from replacing default attention mechanisms with specialized kernels such as flash-linear-attention and causal-conv1d, which fused multiple operations per layer. Installing these kernels reduced end-to-end training time by 3.7x, from 7 hours 44 minutes to 2 hours 5 minutes, while cutting total operations by 15x. The improvements were validated through detailed profiling of CPU and GPU operations.

To eliminate padding waste in batch tensors, the team adopted a 1D packed tensor approach with corrected position IDs to maintain example-level attention masking. This adjustment, combined with the fused kernels, enabled the final sixfold speedup. The techniques and example code are available on AI Runtime for immediate implementation by users.