Reducing Qwen training time by 6x on AI Runtime

Reducing Qwen training time by 6x on AI Runtime | Databricks Blog Skip to main content

Summary

In this blog post we show how we speed up Qwen fine tuning by 6x by installing the correct kernels and making improvements to the data pipeline.

We get that performance increase despite the fact that the baseline model already showed 98% GPU utilization.

AIR helps you diagnose performance issues with a standard torch trace of an LLM.

Training or finetuning a multi-billion parameter foundation model requires significant compute, often spanning hundreds to thousands of hours of GPU training. But tuning training performance for an LLM is challenging, and high level metrics, such as GPU utilization may not be showing your actual utilization. In this blog post, we show how to cut through the layers of abstraction, and speed up fine-tuning of the latest Qwen 3.8 27B model on AI Runtime by 6x. Through systematic performance optimization, even with GPU utilization already at 98%, we go from 7 hours 44 minutes to 1 hour 17 minutes fine tuning a fixed dataset. We will provide a deep dive into the techniques behind these improvements, from how to measure what matters, implementation details, and design trade-offs involved in each optimization. If you want to get started right away, jump directly to the example code on AI Runtime and give it a try! Model setup We start with a standard training loop for Qwen 3.8 27B using the transformers library and the SFT training dataset HuggingFaceH4/ultrachat_200k .

We apply FSDP2 , the latest model sharding technique from torch. This is required to fit the model weights and optimizer state onto the GPUs. A pure distributed data parallel approach would fail to fit the model weights, at this size, on a single GPU. We also enable activation checkpointing, a technique which trades memory space for compute in the backwards pass to allow for a larger batch size per GPU. The dataloader is packing 2D batches with a maximum tokens per batch of 8192, padding to the maximum size of the batches example.

Initial training benchmarks Our baseline training of the model over the complete dataset takes 7 hours and 44 minutes and during the training the GPU shows 98% average utilization. The result is 9,396 tok/s training. This GPU utilization looks high, but as mentioned earlier that number is misleading. The standard GPU utilization metric only counts the percentage of the time the GPU was doing anything , so this does not tell us how well the model was using the GPU while it was active. To understand this we need to look deeper. One way of doing this is the torch trace mechanism, we can adjust our training loop like the following (while also limiting the number of steps so we get a quick response):

For a sample step this gives us the following data: Operation Number of calls Average time (micros) Total time (millis) aten::bmm 31,105 50.5 1,571.80 aten::mul 46,878 16.2 758.9 MmBackward0 497 1,111.4 552.4 aten::copy_ 69,444 7.2 500.6 aten::add 45,623 8.0 366.7 aten::select 68,645 3.4 230.5 aten::exp 9,504 21.7 206.5 aten::sum 16,194 12.2 196.8 aten::empty 40,709 4.7 190.5 aten::as_strided 218,214 0.8 173.3

Table 1: CPU - 14.4s with 1,714,382 calls. Table contains top operations by total time. Operation Number of calls Average time (micros) Total time (millis) nvjet GEMM 1923 1,163.1 2,236.8 Elementwise floating point add 42,657 51.1 2,180.5 Elementwise floating point multiply 35,012 24 841.8 Elementwise floating point fill 33,557 21.5 723.1 Elementwise floating point copy 59,282 11.8 700.7 Memcpy DtoD 5,585 90.5 505.2 sm80_xmma GEMM 31,105 20.9 651.3

Table 2: GPU - 10.3s with 269,448 operations. Table contains top operations by total time Areas of improvement From these tables we notice a few high level issues with the current model: The CPU takes longer to set up each training step than the GPU takes to run it. This is unexpected as generally the CPU side of GPU execution is asynchronous, i.e. the CPU should run well ahead of the GPU . The main effect of this is the GPU is stalling, so we have some opportunity to improve this. There are too many operations! In the tables above, we have hundreds of thousands of small operations, in this case the issue is that the CPU cannot enqueue GPU operations fast enough to keep the GPU busy. The key idea is that we need to reduce the number of operations the CPU is enqueuing to the GPU. There are some general techniques for this, one fits our situation: torch’s CUDA graphs record the cuda kernels seen and replay them with a single CPU call. While powerful, there are many limitations and it conflicts with the dynamism of FSDP2. So for now, this is off the table. Torch.compile traces the code and uses Triton to automatically fuse kernels where possible. In this example, we will use manually fused kernels and in the case of LLMs, this doesn’t leave any room for improvements with torch.compile. Manually fused kernels use specialized fusion techniques for particular operations that torch’s compiler can’t apply on its own. Fortunately there are many established pre-built kernels already in the community. Examples of this include ‘Flash Attention ’, which gives a small improvement in performance in this case.

As well, we have an issue with excess padding in the batch: the input tensor must be rectangular and each row gets padded to the maximum number of tokens for that batch. We will show a trick for LLMs that can be used to minimize it. In the following sections, we will dive deeper into each technical topic, covering what it is, recommended best practices, and the key trade-offs to consider when applying it. Fused kernels You may notice that some kernels occur a lot! In particular, for Qwen, the batch matrix multiplication kernel runs over 30,000 times in one step. In this case we’ll explore fused kernels first. A fused kernel takes multiple smaller GPU kernels and combines them into one unit. This can help performance by avoiding costly transfers between a GPU’s high bandwidth memory and its even higher bandwidth on die memory. As well it can also allow for more efficient scheduling of GPU cores. With this many extra kernels, while torch.compile might work, it will take a long time to run. So reducing the raw number of ops with manually fused kernels is a good first step. Also, well-known fusions such as Flash Attention aren’t something that torch.compile by itself would find. The transformers library by default uses torch’s...