Techniques for Peak Memory Reduction for LoRA Fine-tuning of LLMs on Edge Devices

2026-06-17

AI Summary

This work addresses the challenge of excessive peak memory consumption during LoRA fine-tuning of large language models on edge devices, which often exceeds the capacity of consumer-grade hardware. The authors propose a synergistic optimization framework that integrates dynamic dequantization, selective activation caching, and offloading to disk, complemented by a novel semantic-correlation-based token subset approximation of softmax, logits masking, and an efficient checkpointing mechanism. This approach substantially reduces memory footprint with negligible degradation in model performance. Experimental results demonstrate up to 26× and 28× reductions in peak memory usage on Llama-3.2-3B and Qwen-2.5-3B, respectively, enabling billion-parameter-scale models to undergo full LoRA fine-tuning for the first time on resource-constrained devices.

Abstract

Fine-tuning of Large Language Models (LLMs) using Low-Rank Adaptation (LoRA) on an end-user's data offers personalized experiences while keeping data private, but faces severe memory constraints on consumer hardware. Peak memory during fine-tuning often exceeds device limits, especially for models with billions of parameters and long-context training data. This paper introduces a suite of complementary techniques to reduce memory footprint without sacrificing model quality: (1) base model quantization with on-the-fly dequantization, (2) memory-efficient checkpointing combining selective activation caching and disk offloading, (3) softmax approximation using semantically relevant token subsets, and (4) logits masking. Experiments on Llama-3.2 3B and Qwen-2.5 3B demonstrate up to $26\times$ and $28\times$ reduction in peak memory, enabling fine-tuning on resource-constrained devices.

Problem

• Peak Memory
• LoRA Fine-tuning
• Large Language Models
• Edge Devices
• Memory Constraints

Innovation

• LoRA fine-tuning
• peak memory reduction
• model quantization
• memory-efficient checkpointing
• softmax approximation