Two simple techniques to ensure that your LLM training hits peak performance.
by Bruce Fontaine and Pardis Beikzadeh 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! We start with a standard training loop for Qwen 3.8 27B using the transformers library and the SFT training dataset HuggingFaceH4/ultrachat_200k.
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
From these tables we notice a few high level issues with the current model: 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.
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 SDPA for the standard attention layers. But, it turns out many of these small operations we see come from the non-standard gated delta network layers. We’re lucky that for these layers there are two kernels, causal-conv1d and flash-linear-attention which have been written to fuse many of these operations per layer. How do we discover this? The transformers library implementation of Qwen already tells us what kernels it supports:
There are two high level methods for getting these kernels:
We generally recommend the first, especially paired with our docker container support, build your container with all your kernels installed at build time to reduce the amount of time it takes for your job to start.
While kernelhub is convenient, since kernels often explicitly depend on the version torch and cuda installed, an available kernel for your specific configuration might not be available, and the transformers model will silently fallback to an unoptimized kernel with low performance. After installing the flash-linear-attention and causal-conv1d python packages, we improved our end to end time by a factor of 3.7x down to 2 hours 5 minutes. In the tables below, you can see the new kernels (highlighted), taking the place of the elementwise functions in the previous tables. As well the total number of operations executed has been greatly reduced, by a factor of 15x:
| Operation | Number of calls | Average time (micros) | Total time (millis) |
|---|---|---|---|
| cudaMemsetAsync | 1,854 | 283.9 | 526.3 | | ChunkGatedDeltaRuleFunction | 96 | 871.7 | 83.7 | | aten::mm | 1,987 | 39.6 | 78.6 | | MmBackward0 | 497 | 134.2 | 66.7 | | ChunkGatedDeltaRuleFunctionBackward | 48 | 1,159.9 | 55.7 | | DaoAILab::_causal_conv1d_fwd_cpp | 96 | 55.8 | 5.4 | | CausalConv1dFn | 96 | 51.4 | 4.9 |
Table 3: CPU - 3.9s with 118,316 calls. Table contains top operations by total time.
| Operation | Number of calls | Average time (micros) | Total time (millis) |
|---|---|---|---|
| nvjet GEMM | 1,923 | 1,183.3 | 2,275.5 | | Elementwise floating point multiply | 2184 | 128.1 | 279.9 | | Elementwise copy + cast | 1255 | 101.1 | 126.9 | | chunk_cat flash-linear-attention | 65 | 1,402.0 | 91.2 | | flash_attention_3 forwards | 16 | 2,693.8 | 43.1 | | chunk_gated_delta_rule_fwd | 144 | 255.6 | 36.8 | | flash_attention_3 backwards | 32 | 965.6 | 30.9 |
Table 4: GPU - 3.9s with 18,060 calls. Table contains top operations by total time.
Caveat: during training you will need both a kernel for the forwards and backwards pass and some kernels are forwards only, as they’ve only been implemented to speed up inference.
Takeaways
Our Qwen model is pretty fast now, but we can get even faster. By default our batching setup uses 2D rectangular tensors 8,192 tokens as this is the largest size we can fit without running out of memory. The downside of this approach is that we tend to waste a considerable amount of space in the tensor with padding.
If Qwen only used standard attention mechanisms, the torch’s NestedTensor is now well supported and would help reduce this cost by completely removing the padding. But flash-linear-attention only supports a more traditional 1D packed tensor approach: But we can’t just do this arbitrarily, simply passing the 1D tensor to the model may be acceptable for pretraining, but for finetuning, we need to ensure that in an example tokens only attend to tokens within the same example. To fix this we have to pass in the correct example position_ids to the model. In turn the model will pass these to the fused attention kernels and it will apply the correct attention mask. If you look at the example code here you can see how we compute the needed inputs to make this work.
With this simple looking change, our training time goes from 2 hours 5 minutes to 1 hour and 17 minutes.
There are two reasons why:
Trade off: This style of change does make the input pipeline more complex, and there’s a high chance of introducing errors that impact model quality. But in this case the 36% savings are worth it. There are other techniques for reducing padding such as keeping the original batching mechanism but sorting all inputs by length, so that the padding is minimized, but this order may introduce unexpected correlation in adjacent samples, biasing the model.
Take Aways: When training LLMs over a dataset with inputs shorter than your max trainable tokens (e.g. in this case 8192), efficiently packing data into batches can save a significant amount of time. For this dataset, we saw a 1.6x improvement in training speed, taking our training time from 2 hours 5 minutes to 1 hour and 17 minutes.
In this blog, we showed you how to 6x your training efficiency. It came down to (1) first measuring what matters, not just gpu utilization, but also where your model is spending time training and (2) optimizing and fusing kernels, using known libraries like Flash Linear Attention and (3) tailoring how batches of data are fed into the model.
To do this at scale, you need a platform that makes a performant path the default. In AIR, many of these performance improvements are built in. We bundle in the latest CUDA and flash attention with our base environments, automatically track core metrics like GPU utilization, temperature, and health for you, ensure torch versions (and the profilers are up to date). Our internal agent, Genie Code, is also tuned to look for opportunities like these to further optimize your training. Furthermore, reading data from Databricks is optimized through our UCVolumeDataset.
Please give AIR a try for your next training workload! Ready to try it? See the AI Runtime cookbook, with details on how to run this model yourself.
Subscribe to our blog and get the latest posts delivered to your inbox.