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]

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)

[qoderai]

👋 Review Summary

This adds an opt-in fast LoRA path for Qwen2/Qwen3 on NPU plus a fused SwiGLU implementation. The overall integration is thoughtfully scoped (LoRA SFT only, ENABLE_FAST_LORA gate, adapter compatibility checks), and the fallback behavior to the existing NPU path is clear when conditions are not met.

🛡️ Key Risks & Issues

• Fused SwiGLU on NPU is currently enabled purely based on Triton being importable. In swift/model/npu_patch/model.py:192-246, _get_FUSED_SWIGLU_KERNELS returns Triton-based kernels whenever the import succeeds, and npu_swiglu then unconditionally dispatches to FusedSwiGLUFunction for Qwen NPU patches. Because the underlying kernels in swift/model/npu_patch/fused_swiglu.py are Triton GPU-style ops and do not appear to be implemented for torch-npu tensors, this can cause runtime failures or undefined behavior on NPU hosts where Triton happens to be installed. We should guard this path by device type (or an explicit config flag) and keep NPU using torch_npu.npu_swiglu by default.
• The fast LoRA helpers (_get_active_adapter_name, _get_lora_dropout_p, _is_fast_lora_projection_compatible, can_use_fast_lora{mlp,qkv,o}) encode fairly nuanced eligibility rules (single active adapter, zero dropout, no trainable base bias, not merged/disabled). There are no explicit tests covering these rules, so regressions could silently change when the fast path triggers.

🧪 Verification Advice

• Add unit tests around the NPU fast LoRA enablement in tuner.prepare_model that simulate an NPU environment with ENABLE_FAST_LORA set, covering both the enabled and skipped branches and asserting the logging messages and enable_npu_fast_lora invocation behavior.
• Add targeted unit tests for _get_active_adapter_name, _get_lora_dropout_p, _has_trainable_base_bias, _is_fast_lora_projection_compatible, and _can_use_fast_lora_{mlp,qkv,o} using lightweight fake modules, to lock in the intended eligibility semantics (e.g., multiple active adapters, nonzero dropout, merged/disabled adapters).
• Add numerical regression tests for the fused SwiGLU kernels and LoRA fast paths: compare FusedSwiGLUFunction and LoRA_MLP / LoRA_QKV / LoRA_W against reference implementations across representative shapes/dtypes, checking both forward outputs and gradients.
• Add a test that simulates an NPU-only context where Triton is importable and asserts that the NPU path still uses torch_npu.npu_swiglu and does not attempt to execute the Triton kernels.

💡 Thoughts & Suggestions

• The gating on ENABLE_FAST_LORA, LoRA-only SFT, and adapter compatibility is a solid way to keep blast radius low while still offering a significant speedup. Once the device-guarding for fused SwiGLU is in place and tests cover the eligibility helpers and kernels, this should be a robust optimization path.
• Consider also exposing a simple configuration/flag to force-disable the fused SwiGLU path independently of fast LoRA, so that operators can bisect issues between the fused activation and the LoRA projection fusion if needed.

---

🤖 Generated by Qoder • View workflow run

[addsubmuldiv]

是否考虑一下把开关做成一个可配置选项比如--use_npu_fast_lora，而不是环境变量？或者给环境变量补个文档？