GRASS: Gradient-based Adaptive Layer-wise Importance Sampling for Memory-efficient Large Language Model Fine-tuning

Kaiyuan Tian; Yu Tang; Gongqingjian Jiang; Baihui Liu; Yifu Gao; Xialin Su; Linbo Qiao; Dongsheng Li

cs.CLApr 9, 2026

GRASS: Gradient-based Adaptive Layer-wise Importance Sampling for Memory-efficient Large Language Model Fine-tuning

Authors: Kaiyuan Tian, Yu Tang, Gongqingjian Jiang, Baihui Liu, Yifu Gao, Xialin Su, Linbo Qiao, Dongsheng Li

Organizations: National University of Defense Technology · Information Support Force Engineering University

Abstract

Full-parameter fine-tuning of large language models is constrained by substantial GPU memory requirements. Low-rank adaptation methods mitigate this challenge by updating only a subset of parameters. However, these approaches often limit model expressiveness and yield lower performance than full-parameter fine-tuning. Layer-wise fine-tuning methods have emerged as an alternative, enabling memory-efficient training through static layer importance sampling strategies. However, these methods overlook variations in layer importance across tasks and training stages, resulting in suboptimal performance on downstream tasks. To address these limitations, we propose GRASS, a gradient-based adaptive layer-wise importance sampling framework. GRASS utilizes mean gradient norms as a task-aware and training-stage-aware metric for estimating layer importance. Furthermore, GRASS adaptively adjusts layer sampling probabilities through an adaptive training strategy. We also introduce a layer-wise optimizer state offloading mechanism that overlaps computation and communication to further reduce memory usage while maintaining comparable training throughput. Extensive experiments across multiple models and benchmarks demonstrate that GRASS consistently outperforms state-of-the-art methods, achieving an average accuracy improvement of up to 4.38 points and reducing memory usage by up to 19.97%.

Figures & tables

Method	MultiA.	AddSub	GSM8K	AQuA	SingleEq	SVAMP	Avg. $\uparrow$
FFT	94.43	66.08	47.31	18.90	80.51	55.50	60.46
LISA	91.17	62.53	42.91	20.08	71.65	51.10	56.57

Table 1: Comparison of full-parameter fine-tuning (FFT) and LISA on arithmetic reasoning benchmarks. Although LISA approaches FFT performance on certain tasks, it exhibits notable degradation on others(e.g., GSM8K and SingleEq).

Figure 1: A comparison of normalized layer-wise mean gradient norm across different datasets on TinyLlama.

Figure 2: An overview of static layer-wise sampling (Left) and GRASS (Right). GRASS alternates between sampling a subset of model layers for tuning and updating sampling probabilities during training.

Figure 3: A comparison of vanilla and overlapped layer-wise optimizer state offloading. HtoD refers to host to device, while DtoH represents device to host. Each block with index

i

stands for a parameter in layer

i

Model	Method	MultiArith	AddSub	GSM8K	AQuA	SingleEq	SVAMP	Avg. $\uparrow$
TinyLlama	FFT	64.17	40.76	15.16	13.78	42.92	24.10	33.48
	LoRA r=128	61.17	25.32	15.16	12.20	38.19	27.00	29.84
	LoRA r=256	65.50	30.89	17.44	12.60	37.59	28.20	32.04
	DoRA	67.83	34.68	18.35	12.99	36.81	29.50	33.36
	LISA	65.00	38.04	17.74	12.99	43.11	24.90	33.63
	GRASS	68.00	35.19	17.13	15.16	42.52	27.30	34.22

Table 2: Accuracy comparison of multiple LLMs and PEFT methods on six math reasoning benchmarks. The best and second best results are marked in bold and underlined , respectively.

Model	Method	BoolQ	PIQA	SIQA	HellaS.	WinoG.	ARC-e	ARC-c	OBQA	Avg. $\uparrow$
TinyLlama	FFT	61.90	51.25	36.64	29.72	50.43	27.95	25.85	31.00	39.34
	LoRA r=128	58.35	52.50	34.34	26.17	50.75	25.04	25.68	24.80	37.20
	LoRA r=256	58.69	50.76	33.52	25.59	51.70	26.64	27.82	26.20	37.62
	DoRA	57.71	51.74	34.39	25.98	50.83	26.01	27.05	26.40	37.51
	LISA	60.18	49.05	33.78	25.77	49.49	26.22	25.43	28.80	37.34
	GRASS	61.29	51.74	34.75	27.01	50.12	27.74	26.45	29.60	38.59

Table 3: Accuracy comparison of multiple LLMs and PEFT methods on eight commonsense reasoning benchmarks. The best and second best results are marked in bold and underlined , respectively.

	FFT	LoRA		DoRA	LISA		GRASS
Model	-	r=128	r=256	-	$\gamma$ =2	$\gamma$ =4	$\gamma$ =2	$\gamma$ =4
TinyLlama	8.76	4.71	5.03	5.71	4.93	5.31	4.49	4.55
Gemma-2B	21.19	11.26	11.81	11.64	13.33	14.27	12.45	12.46
LLaMA2-7B	51.32	19.97	20.86	23.23	20.68	22.31	19.08	19.22

Table 4: Comparison of peak GPU memory consumption (GB) for different models and fine-tuning methods. Batch size and sequence length are set to 1 and 1024, respectively.

Figure 4: (Left) Peak memory consumption of fine-tuning LLaMA2-7B with different methods. (Right) Memory consumption of fine-tuning LLaMA2-7B across different sequence lengths.

Model	Method	MultiArith	AddSub	GSM8K	AQuA	SingleEq	SVAMP	Avg. $\uparrow$
TinyLlama	GRASS*	54.50	39.24	13.12	15.75	40.75	20.80	30.69
TinyLlama	GRASS	68.00	35.19	17.13	15.16	42.52	27.30	34.22
Gemma-2B	GRASS*	91.00	63.80	41.40	31.10	76.77	49.90	59.00
Gemma-2B	GRASS	93.50	66.33	43.06	29.13	78.35	53.50	60.65
LLaMA2-7B	GRASS*	91.83	61.52	41.93	13.78	77.56	52.90	56.59
LLaMA2-7B	GRASS	92.00	67.09	42.15	23.62	77.56	55.10	59.59

Table 5: Comparison of GRASS and its variant without adaptive update of sampling probabilities (GRASS*) on arithmetic reasoning tasks.

$\gamma$	2					4					8
$T_{s}$	5	10	25	50	100	5	10	25	50	100	5	10	25	50	100
Avg.	57.68	57.12	59.59	58.40	58.85	58.13	59.64	60.34	60.37	59.32	62.22	61.87	62.32	62.46	60.30

Table 6: Comparison average accuracy of applying different GRASS configurations on arithmetic reasoning tasks.

Model	Seed 1	Seed 2	Seed 3
TinyLlama	34.36	34.22	34.12
Gemma-2B	60.64	60.65	60.85
LLaMA2-7B	59.45	59.59	59.37

Table 7: Average accuracy on arithmetic reasoning tasks with varying random seeds.

Figure 5: Training throughput of different methods on LLaMA2-7B.

\text{GRASS}^{{\dagger}}

represents a variant of GRASS that does not use overlap.

Model	Probing phase	MGN update	Total training time
Tinyllama	34 (3.24%)	32 (3.09%)	1050 (100%)
Gemma-2B	30 (0.99%)	42 (1.37%)	3036 (100%)
Llama2-7B	54 (0.95%)	60 (1.06%)	5708 (100%)

Table 8: Time breakdown of the additional runtime introduced by GRASS on the arithmetic reasoning dataset. Batch size and sequence length are set to 4 and 1024, respectively.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter		Arithmetic Reasoning			Commonsense Reasoning
Hyperparameter		TinyLlama	Gemma-2B	LLaMA2-7B	TinyLlama	Gemma-2B	LLaMA2-7B
FFT	Learning rate	1e-5			1e-5
FFT	Batch size	8	4	4	8	4	4
LoRA	Learning rate	1e-4			5e-5	2e-5	2e-5
	Batch size	8	4	4	8	4	4
	Rank	128/256			128/256

Appendix

Table 9: Hyperparameter configurations of different fine-tuning methods for arithmetic reasoning and common reasoning tasks.

Figure 6: Comparison of layer-wise mean gradient norm across different datasets on Gemma-2B (Left) and LLaMA2-7B (Right).

Figure 7: Loss curves of fine-tuning TinyLlama (Left), Gemma-2B (Middle), and LLaMA2-7B (Right) with different methods on Alpaca-GPT4 dataset.

$T_{p}$	25	50	100	150	200
Avg.	58.95	59.31	59.33	59.59	58.20

Appendix

Table 10: Average accuracy on arithmetic reasoning tasks with varying lengths of probing phase.

LLM Fine-Tuning

LLM fine-tuning adapts pretrained large language models to tasks, domains, languages, and behaviors through additional training.

See all 755 papers

ChunkFT: Byte-Streamed Optimization for Memory-Efficient Full Fine-Tuning
May 20, 2026Yongkang Liu, Zijing Wang, Mengjie Zhao +7Fine-Tuning Memory-Efficient Fine-Tuning
Northeastern University, China · Shanghai Jiao Tong University, China · CIS, LMU Munich, Germany +2
Activation-Efficient Fine-Tuning via Sequence-Compressed Low-Rank Subspaces
Apr 3, 2026Irene Tenison, Stella Ahn, Miriam Kim +2Memory-Efficient Fine-Tuning LLM Fine-Tuning
MIT CSAIL · Harvard SEAS
Beyond LoRA vs. Full Fine-Tuning: Gradient-Guided Optimizer Routing for LLM Adaptation
May 8, 2026Haozhan Tang, Xiuqi Zhu, Xinyin Zhang +3LLM Fine-Tuning Low-Rank Adaptation
Carnegie Mellon University · Tsinghua University · Infinigence AI

May 20, 2026J/K move · Enter open · S save

ChunkFT: Byte-Streamed Optimization for Memory-Efficient Full Fine-Tuning

Yongkang Liu, Zijing Wang, Mengjie Zhao, Ercong Nie+6

Northeastern University, China · Shanghai Jiao Tong University, China · CIS, LMU Munich, Germany+2

Fine-Tuning Memory-Efficient Fine-Tuning LLM Fine-Tuning Memory-Efficient Optimization Memory-Efficient Training

Open paper AskPDF arXiv

This work presents ChunkFT, a memory-efficient fine-tuning framework that reformulates full-parameter fine-tuning around a dynamically activated working set. ChunkFT enables gradient computation for arbitrary sub-tensors without modifying the network architecture, providing an algorithmic foundation for optimizing arbitrary sub-networks while avoiding standard dense gradient computation. We provide a theoretical convergence analysis of ChunkFT in the deterministic setting. Empirically, we apply ChunkFT to fine-tune Llama 3-8B and Llama 3-70B using a single RTX 4090-24GB GPU and 2

\times

H800-80GB GPUs, respectively. Full-parameter fine-tuning of a 7B model with a 1K input length requires only 13.72GB of GPU memory. The results demonstrate the effectiveness of ChunkFT in memory usage, running time, and optimization quality. Moreover, downstream evaluations on language understanding, mathematical reasoning, and MT-Bench show that ChunkFT consistently outperforms existing memory-efficient baselines. Notably, ChunkFT achieves performance comparable to, and in some cases exceeding, full-parameter fine-tuning. Our repository is on https://github.com/misonsky/chunk.

Abstract

LLM Fine-Tuning

Explore similar work

ChunkFT: Byte-Streamed Optimization for Memory-Efficient Full Fine-Tuning

Activation-Efficient Fine-Tuning via Sequence-Compressed Low-Rank Subspaces

Beyond LoRA vs. Full Fine-Tuning: Gradient-Guided Optimizer Routing for LLM Adaptation