AI 服務暫時不可用，以下為來源正文，待恢復後補全翻譯。
Reducing Qwen training time by 6x on AI Runtime | Databricks Blog Skip to main content 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: OperationNumber of callsAverage time (micros)Total time (millis) aten::bmm31,10550.51,571.80 aten::mul46,87816.2758.9 MmBackward04971,111.4552.4 aten::copy_69,4447.2500.6 aten::add45,6238.0366.7 aten::select68,6453.4230.5 aten::exp9,50421.7206.5 aten::sum16,19412.2196.8 aten::empty40,7094.7190.5 aten::as_strided218,2140.8173.3 Table 1: CPU - 14.4s with 1,714,382 calls. Table contains top operations by total time. OperationNumber of callsAverage time (micros)Total time (millis) nvjet GEMM19231,163.12,236.8 Elementwise floating point add42,65751.12,180.5 Elementwise floating point multiply35,01224841.8 Elementwise floating point fill33,55721.5723.1 Elementwise floating point copy59,28211.8700.7 Memcpy DtoD5,58590.5505.2 sm80_xmma GEMM31,10520.9651.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 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: flash-linear-attention causal-conv1d There are two high level methods for getting these kernels: Install the kernels yourself. This can range from: A simple pip install (e.g. for flash-linear-attention) from PyPi A git clone and compile of the relevant repository Find another source for a precompiled package in case the compilation takes too long Use Hugging Face’s Kernel Hub to fetch precompiled 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: OperationNumber of callsAverage time (micros)Total time (millis) cudaMemsetAsync1,854283.9526.3 ChunkGatedDeltaRuleFunction96871.783.7 aten::mm1,98739.678.6 MmBackward0497134.266.7 ChunkGatedDeltaRuleFunctionBackward481,159.955.7 DaoAILab::_causal_conv1d_fwd_cpp9655.85.4 CausalConv1dFn9651.44.9 Table 3: CPU - 3.9s with 118,316 calls. Table contains top operations by total time. OperationNumber of callsAverage time (micros)Total time (millis) nvjet GEMM1,9231,183.32,275.5 Elementwise floating point multiply2184128.1279.9 Elementwise copy + cast1255101.1126.9 chunk_cat flash-linear-attention651,402.091.2 flash_attention_3 forwards162,693.843.1 chunk_gated_delta_rule_fwd144255.636.8 flash_attention_3 backwards32965.630.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 Consider using well-established, highly optimized kernels that are readily available for common models. Leveraging these kernels can deliver significant performance improvements with relatively little effort. In this specific case we saw a 3.7x improvement in training speed, decreasing our training time from 7 hours 44 minutes to 2 hours 5 minutes. Batching, tokenization and padding 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: The primary reason is that this packing increases the number of sequences per batch and reduces the total number of batches by approximately 36%. Note that this depends on the distribution and length of examples in your dataset. A secondary reason is that the flash linear attention kernels need to be just-in-time compiled for each input shape and the shape was previously dynamic, leading to a few extra minutes of compilation. 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 [truncated for AI cost control]