feat(sft):Add optional fused linear ce for qwen2/qwen3 on npu

[yangguang-zhang]

PR type

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

PR information

Background

Training Large Language Models with massive vocabularies (e.g., Qwen2 with a 152K vocab size) on long-context tasks encounters a severe "Memory Wall" at the final layer. In the standard HuggingFace implementation, the LM-Head and CrossEntropyLoss are computed sequentially. This requires materializing a massive Logits tensor of shape [Batch * SeqLen, VocabSize] in HBM. For instance, with B=1, Seq=8192, this single tensor and its gradients can easily consume >4.5 GB of VRAM, often leading to Out-Of-Memory (OOM) errors on Ascend NPUs and severely limiting the maximum sequence length or batch size.

What this PR does

This PR introduces an optional Memory-Efficient Fused LM-Head & CrossEntropy path for Qwen series training on NPU.

When enabled, it patches the Qwen2ForCausalLM.forward (and its variants) to intercept the hidden_states before they enter the lm_head. It uses a Chunked Autograd strategy in the time dimension:

1. Slices the sequence into small chunks (e.g., CHUNK_SIZE = 2048).
2. Computes the local matrix multiplication (F.linear) and local cross-entropy loss.
3. Computes the gradients for the input features and weights on the fly (torch.autograd.grad) and accumulates them.

This ensures the full [Batch * SeqLen, VocabSize] Logits matrix is never materialized in memory.

This work is deeply inspired by Liger-Kernel and Unsloth, but is strictly implemented using PyTorch's native APIs (F.linear and F.cross_entropy) to ensure 100% compatibility with the Ascend NPU's Graph Engine (GE) and PTA dispatcher, completely avoiding compiler issues in triton-ascend.

Mathematical Equivalence

The chunking mechanism maintains strict mathematical equivalence to the standard global computation.
Given $Z = X \cdot W^T$ and $L = \text{CE}(Z, Y)$, the gradients are $\nabla X = \nabla Z \cdot W$ and $\nabla W = (\nabla Z)^T \cdot X$.
By chunking $X$ into $[X_1, X_2, ... X_C]$ along the sequence dimension:
This summation is implemented via grad_weight += torch.matmul(grad_logits.t(), x_chunk), guaranteeing exact mathematical fidelity (Gradient Cosine Similarity = 1.000000).

Scope

This PR is intentionally scoped to the following cases:

• Qwen2, Qwen3,architectures.
• NPU only.
• Pre-training / SFT workloads (where labels are provided).

Safety & Compatibility

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

• Opt-in Enablement: Users must explicitly set the environment variable FUSED_LINEAR_CE=1 to enable it.
• Graceful Fallback: In Inference/Generation mode (when labels=None), it automatically falls back to the original self.lm_head(hidden_states).

Limitations

The current implementation does not cover the following cases yet:

• Vision-Language Models (e.g., Qwen3-VL-MoE) due to complex image token routing.
• GRPO/PPO RLHF training (which relies on full logits for KL-divergence and advantage calculation, rather than standard CrossEntropy).

[gemini-code-assist]

Code Review

This pull request introduces a memory-efficient fused Linear and CrossEntropy Loss operator (NPUFusedLinearCrossEntropy) specifically optimized for Huawei Ascend NPUs to prevent Out-Of-Memory (OOM) errors during training. It chunks hidden states along the sequence dimension to compute local logits and gradients in-place, returning a zero-sized dummy tensor for logits. Additionally, it integrates this optimization into the Qwen2 and Qwen3 patches and adds a safeguard to SwiftMixin._compute_acc to handle the empty logits. The review feedback identifies two key issues: first, a critical PyTorch runtime error in fused_linear_ce.py caused by calling requires_grad_(True) on a non-leaf tensor slice, with a suggestion to set it on the logits chunk instead; second, a potential static analysis or linter warning in model.py due to the conditional import of fused_linear_ce, which can be resolved by defining it as None at the module level.

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]

Calling requires_grad_(True) on x_chunk will raise a RuntimeError: you can only change requires_grad flags of leaf variables because x_chunk is a slice (view) of x, which is not a leaf tensor.

To resolve this issue and completely avoid the error, we can compute logits_chunk inside the torch.no_grad() context (which is already active), and then call requires_grad_(True) on logits_chunk itself. Since logits_chunk is a newly created tensor with no history, it is a leaf tensor and can safely be marked as requiring gradients. We then perform the scaling and cross-entropy loss computation inside with torch.enable_grad():.

                # 1. Local Fused Linear (NPU Cube Engine full speed)
                logits_chunk = F.linear(x_chunk, weight)
                logits_chunk.requires_grad_(True)

                # Enable localized gradient tracking for the current chunk sandbox
                with torch.enable_grad():
                    # Apply model-specific scaling (e.g., Gemma-2 softcapping, Cohere scaling)
                    scaled_logits = logits_chunk
                    if logit_scaling != 0:
                        scaled_logits = scaled_logits * logit_scaling
                    if logit_softcapping != 0:
                        scaled_logits = logit_softcapping * torch.tanh(scaled_logits / logit_softcapping)

                    # 2. Local CrossEntropy Loss
                    loss_chunk = F.cross_entropy(scaled_logits.float(), y_chunk, ignore_index=-100, reduction='sum')
                    loss_chunk_mean = loss_chunk / denominator

                    total_loss += loss_chunk_mean.item()

                    # 3. Compute local gradients
                    grad_logits = torch.autograd.grad(loss_chunk_mean, logits_chunk)[0]
                    grad_logits = grad_logits.to(x.dtype)

[gemini-code-assist]

The name fused_linear_ce is conditionally imported inside an if block. If FUSED_LINEAR_CE_ENABLED is False, the name fused_linear_ce will not be defined in the module scope, which can trigger static analysis, linter, or IDE warnings when referencing it in the patch dictionaries.

To prevent this, define fused_linear_ce = None at the module level before checking the environment variable.

Suggested change

	FUSED_LINEAR_CE_ENABLED = os.getenv("FUSED_LINEAR_CE", "0").strip() == "1"
	fused_linear_ce = None
	FUSED_LINEAR_CE_ENABLED = os.getenv("FUSED_LINEAR_CE", "0").strip() == "1"