Add optional fast lora path for qwen2/qwen3 on npu

[burm95]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Background

LoRA training on NPU still has room for optimization, especially for common Qwen dense model SFT workloads. In these scenarios, the default LoRA execution path works correctly, but there is still an opportunity to reduce overhead in the MLP and attention projection path.

What this PR does

This PR introduces an optional fast_lora path for Qwen2/Qwen3 training on NPU.

When enabled, it patches the Qwen2/Qwen3 MLP and Attention forward path to use fused SwiGLU and fast LoRA projection logic. The original implementation is kept as the fallback path, so the optimization is only applied when the required conditions are satisfied.

This path relies on triton-ascend for the fused SwiGLU kernel implementation.

Based on current internal experiments, this optimization can reduce SFT training time by approximately 25% to 40% in supported NPU LoRA workloads.

This work is inspired by Unsloth which cannot not be used on npu.

Scope

This PR is intentionally scoped to the following cases:

• Qwen2 / Qwen3 only
• LoRA SFT only
• NPU only
• triton-ascend is required

Safety

To avoid changing the default training behavior, this feature is disabled by default.

When the fast path is not applicable, the code automatically falls back to the original implementation.

Compatibility checks

The fast path is only enabled when all of the following conditions are satisfied:

• dropout == 0
• single active adapter
• no trainable base bias
• not merged

Benchmark

Benchmark results should compare the original path and the fast_lora path in terms of:

• training throughput
• memory usage
• end-to-end SFT training time

Limitations

The current implementation does not cover the following cases yet:

• Qwen3.5
• MoE models
• RLHF training
• environments without triton-ascend

[gemini-code-assist]

Code Review

This pull request introduces NPU-specific optimizations, including fast LoRA and fused SwiGLU implementations for Qwen2 and Qwen3 models, along with integration into the model preparation pipeline. Key feedback highlights that fused_swiglu.py is an accidental duplicate of fast_lora.py resulting in circular imports and missing Triton kernels. Additionally, the reviewer points out potential gradient corruption from in-place modifications of saved tensor X, redundant double transpositions and del statements in fast_lora.py, inefficient .clone() calls on large tensors in model.py, a missing Callable import, and unidiomatic type checking in tuner.py.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The file fused_swiglu.py appears to be an accidental duplicate of fast_lora.py. It contains the exact same 663 lines of code and attempts to import swiglu_fg_kernel and swiglu_DWf_DW_dfg_kernel from itself (from .fused_swiglu import ...), which will result in a circular import error or AttributeError at runtime. Furthermore, the actual Triton implementations for these SwiGLU kernels are completely missing from the pull request. Please replace this file with the correct Triton kernel implementations.

[gemini-code-assist]

Overwriting the saved tensor X in-place during the backward pass (out = X if ctx.inplace else None) is extremely dangerous. X is a saved tensor used for gradient computation of other parameters (e.g., d_upA, d_gateA) and is also part of the wider autograd graph (such as residual connections). Modifying it in-place can lead to silent gradient corruption or trigger PyTorch's runtime error: RuntimeError: one of the variables needed for gradient computation has been modified by an in-place operation. It is highly recommended to disable this in-place optimization or default inplace to False to ensure correctness.

[gemini-code-assist]

There is a redundant double transposition of weights during the backward pass. For example, upW = fast_dequantize(upW.t(), upW_quant) transposes upW, and then torch.matmul(df, upW.t(), ...) transposes it again. Since fast_dequantize currently just returns the tensor as-is, this results in upW.t().t(), which is equivalent to upW. This is redundant and confusing. You should pass upW directly without transposing it twice.

Suggested change

	upW = fast_dequantize(upW.t(), upW_quant)
	dX = torch.matmul(df, upW.t(), out = X if ctx.inplace else None)
	upW = fast_dequantize(upW, upW_quant)
	dX = torch.matmul(df, upW, out = X if ctx.inplace else None)

[gemini-code-assist]

The del W statement (and other similar del statements like del upW, del gateW) only deletes the local reference within the function scope. It does not free the underlying tensor memory because the tensors are still referenced elsewhere (e.g., by the module or autograd context). Since Python automatically cleans up local references when the function exits, these del statements are redundant and can be removed to improve readability.

[gemini-code-assist]

Calling .clone() on large activation and gradient tensors (grad_output, gate, up) during the backward pass is highly inefficient and significantly increases memory usage during training. Since Triton kernels do not modify these inputs in-place, making them contiguous with .contiguous() is sufficient. Removing the .clone() calls will reduce memory overhead.

Suggested change

	grad_output_2d = grad_output.reshape(-1, grad_output.shape[-1]).contiguous().clone()
	gate_2d = gate.reshape(-1, gate.shape[-1]).contiguous().clone()
	up_2d = up.reshape(-1, up.shape[-1]).contiguous().clone()
	grad_output_2d = grad_output.reshape(-1, grad_output.shape[-1]).contiguous()
	gate_2d = gate.reshape(-1, gate.shape[-1]).contiguous()
	up_2d = up.reshape(-1, up.shape[-1]).contiguous()

[gemini-code-assist]

The type hint Callable is used here, but Callable is not imported from typing in this file. Although from __future__ import annotations prevents a runtime NameError during module import, it will still fail static analysis or runtime type inspection. Please add from typing import Callable to the imports.

[gemini-code-assist]

Checking the type of args using type(args).__name__ == 'SftArguments' is unidiomatic and fragile (e.g., it doesn't support subclasses). Since SftArguments is already imported on line 11, you should use isinstance(args, SftArguments) instead.

Suggested change

	is_sft = type(args).__name__ == 'SftArguments'
	is_sft = isinstance(args, SftArguments)