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).

Experimental Results

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.