First-Order LoRA Training on Mobile NPUs: No Need to Upload Private Data

Title

First-Order LoRA Training on Mobile NPUs: No Need to Upload Private Data

Meta Description

This article introduces a system for running first-order LoRA training on mobile NPUs, placing frozen weight product operations on the NPU while keeping the training graph and adapters on the CPU, enabling privacy-preserving on-device fine-tuning.

Article Content

On-device fine-tuning allows large language models (LLMs) to learn from personal data without uploading user privacy data. Low-rank adaptation (LoRA) reduces the number of trainable parameters, but each first-order training step still requires multiplying with the frozen base model weights during forward propagation, activation recomputation, and input gradient computation. Mobile neural processing units (NPUs) already hold these weights resident during inference; however, using them for training also requires computing input gradients, which the NPU's inference interface does not support.

System Architecture: NPU-CPU Collaborative Division of Work

The core idea of this system is to execute all three product operations involving frozen weights on the NPU, while keeping the training graph, adapters, and optimizer on the CPU. This leverages the NPU's advantage of weight residency while avoiding the overhead of duplicating weights separately for training.

Specifically, the input gradient kernel reuses the weight blocks resident on the NPU and the forward-pass dequantizer, thereby avoiding transposed weight copies and ensuring that forward and backward propagation use the same effective weights. This design resolves the key contradiction that the inference interface does not support input gradient computation.

The low-rank adaptation principle of LoRA: LoRA connects two small matrices $A$ (down-projection) and $B$ (up-projection) in parallel beside the frozen original weight matrix $W$. During training, only $A$ and $B$ are updated, and the forward computation becomes $h = Wx + BAx$. Since $W$ is frozen, there is no need to compute gradients for $W$ during backpropagation, but the gradient of the loss with respect to $x$ (the input gradient) must still be computed to continue backpropagating to shallower layers; this step necessarily uses $W$ itself. This is precisely why the article lists the three operations of "multiplication with frozen weights" separately and moves them to the NPU: although the number of parameters is small, the frozen weights remain the main body of computation and memory access.

Limitations of the NPU inference interface: Mobile NPUs typically execute inference as fixed operator graphs, and the interface only exposes a forward path of "given an input tensor, return an output tensor." It does not provide backward operators for differentiating with respect to inputs, nor does it guarantee that the weight layout can be directly transposed and reused. Therefore, to use the NPU to accelerate LoRA training, one must manually construct input gradient computation on the host side using equivalent kernels and reuse the weight blocks already resident and dequantized in the forward pass, avoiding storing a separate transposed copy of the weights for the backward pass.

Precision Selection and Gradient Stability

To maintain usable gradients, the system selects precision based on gradient sensitivity and applies per-row range protection to half-precision (FP16) outputs. This approach avoids gradient underflow or overflow caused by low precision while maintaining computational efficiency.

At runtime, the system fuses the remaining CPU operators and overlaps NPU execution with adapter computation to improve overall throughput. Recoverable execution and thermal throttling mechanisms allow training to continue as a background task, adapting to long-running mobile scenarios.

Gradient underflow and range protection in low-precision training: FP16 has only 5 exponent bits, and the smallest representable normal number is approximately $6\times10^{-5}$. Input gradients during backpropagation are often several orders of magnitude smaller than activations, making them prone to underflow to zero and causing shallow layers to learn nothing at all. The per-row range protection approach works as follows: before quantization or writing FP16, compute the maximum value of each row, dynamically adjust the scaling factor accordingly (similar to per-row loss scaling), and bring the values in that row into the safe range of FP16, then restore them upon reading. Compared with global uniform scaling, per-row scaling accommodates different channels within the same tensor that vary greatly in magnitude, avoiding the dilemma of "large values overflow, small values underflow."

Experimental Evaluation: Multi-Platform, Multi-Model Validation

The research team evaluated the system on two Qualcomm Snapdragon platforms and four LLMs from the Qwen and Llama families. The workloads cover overnight personalization tasks based on labels provided by the phone and the public LaMP-2 benchmark. Comparison baselines include an optimized CPU trainer, a mobile GPU trainer, forward-only training, and reference results from GPU training.

Experimental results show that the system can effectively complete first-order LoRA training under multiple configurations, providing a feasible path for privacy-preserving on-device fine-tuning.

LaMP-2 benchmark: LaMP (Language Model Personalization) is a set of benchmarks for evaluating the personalization capabilities of language models. LaMP-2 is the movie-tag personalization task within it: given a user's historical movie reviews and their tags (such as "suspense" or "heartwarming"), the model is required to predict personalized tags for the user's new reviews. Because the tag system varies from person to person and depends on the user's private history, it naturally fits the on-device fine-tuning scenario—the model must learn user preferences locally without uploading historical text, so it is often used to measure the effectiveness of personalized fine-tuning.

Key Takeaways

• All frozen-weight product operations are placed on the NPU, while the training graph, adapters, and optimizer remain on the CPU
• The input gradient kernel reuses resident weight blocks and the dequantizer, avoiding transposed copies
• Precision is selected based on gradient sensitivity, with per-row range protection for FP16 outputs
• Supports recoverable execution and thermal throttling, suitable for background training
• Effectiveness validated on Qualcomm Snapdragon platforms and Qwen/Llama models

According to source analysis, this system is the first to achieve first-order LoRA training on mobile NPUs, breaking through the limitation that inference interfaces do not support input gradient computation.

One-line Summary

All three frozen-weight products in LoRA are executed on the mobile NPU, achieving on-device first-order training for the first time.

Paragraph Summary

On-device fine-tuning allows large models to learn from personal data without uploading private data, but first-order LoRA training still requires repeated multiplication with frozen weights. Although mobile NPUs keep these weights resident, their inference interfaces do not support input gradient computation. This research proposes an NPU-CPU collaborative architecture: all three product operations involving frozen weights are completed on the NPU, while the training graph, adapters, and optimizer remain on the CPU; the input gradient kernel reuses the NPU-resident weight blocks and the forward dequantizer, avoiding transposed weight copies. The system selects precision based on gradient sensitivity and applies per-row range protection to FP16 outputs, while fusing CPU operators, overlapping NPU and adapter computation, and supporting recoverable execution and thermal throttling. Evaluations on two Qualcomm Snapdragon platforms and four Qwen/Llama models, covering overnight personalization and the LaMP-2 benchmark, validate its effectiveness and achieve first-order LoRA training on mobile NPUs for the first time.