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.CL, cs.LG
发布日期: 2026-04-09
备注: Accepted by ACL 2026 Findings
💡 一句话要点
GRASS:基于梯度自适应层重要性采样,实现大语言模型高效微调。
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 大语言模型 微调 层重要性采样 梯度优化 内存优化
📋 核心要点
- 现有层级微调方法忽略了层重要性随任务和训练阶段的变化,导致下游任务性能受限。
- GRASS利用梯度范数动态评估层重要性,并自适应调整层采样概率,实现更有效的微调。
- 实验表明,GRASS在多个模型和基准测试中显著提升性能,并降低了内存使用。
📝 摘要(中文)
大语言模型的全参数微调受限于GPU内存需求。低秩适配方法通过仅更新部分参数来缓解这一问题,但通常限制了模型表达能力,性能低于全参数微调。层级微调方法通过静态层重要性采样策略实现内存高效训练。然而,这些方法忽略了层重要性在不同任务和训练阶段的变化,导致下游任务性能欠佳。为解决这些限制,我们提出了GRASS,一个基于梯度的自适应层级重要性采样框架。GRASS利用平均梯度范数作为任务感知和训练阶段感知的指标来估计层重要性,并通过自适应训练策略调整层采样概率。我们还引入了层级优化器状态卸载机制,通过计算和通信重叠来进一步降低内存使用,同时保持相当的训练吞吐量。在多个模型和基准测试上的实验表明,GRASS始终优于现有方法,平均准确率提高高达4.38个百分点,内存使用量减少高达19.97%。
🔬 方法详解
问题定义:大语言模型全参数微调需要大量GPU内存,而低秩适配方法虽然降低了内存需求,但牺牲了模型表达能力。现有的层级微调方法采用静态的层重要性采样策略,无法适应不同任务和训练阶段层重要性的变化,导致微调效果不佳。
核心思路:GRASS的核心思路是利用梯度信息动态地评估每一层的重要性,并根据重要性自适应地调整每一层被采样的概率。通过这种方式,模型可以在训练过程中更加关注重要的层,从而提高微调效率和性能。
技术框架:GRASS框架主要包含三个核心模块:1) 基于梯度范数的层重要性评估模块:计算每一层参数的平均梯度范数,作为该层重要性的度量。2) 自适应层采样概率调整模块:根据层重要性动态调整每一层被采样的概率,重要性高的层更容易被采样。3) 层级优化器状态卸载模块:将优化器的状态信息卸载到CPU或硬盘,以减少GPU内存占用,并通过计算和通信重叠来保证训练效率。
关键创新:GRASS的关键创新在于其自适应的层重要性采样策略。与传统的静态采样方法不同,GRASS能够根据任务和训练阶段动态地调整采样概率,从而更好地适应不同任务的需求。此外,层级优化器状态卸载机制进一步降低了内存占用,使得GRASS能够在资源有限的条件下进行高效微调。
关键设计:GRASS使用平均梯度范数作为层重要性的度量,并通过一个平滑因子来控制采样概率的更新速度。层级优化器状态卸载机制采用异步通信的方式,将优化器状态卸载到CPU或硬盘,并在需要时再将其加载回GPU。具体的损失函数与原始的全参数微调保持一致,旨在优化下游任务的性能。
📊 实验亮点
GRASS在多个模型和基准测试上取得了显著的性能提升。例如,在某些任务上,GRASS的平均准确率比现有方法提高了高达4.38个百分点,同时内存使用量减少了高达19.97%。这些实验结果表明,GRASS是一种高效且有效的微调方法,能够在降低内存需求的同时提高模型性能。
🎯 应用场景
GRASS适用于各种需要对大语言模型进行微调的场景,例如自然语言处理、文本生成、机器翻译等。该方法可以显著降低微调所需的GPU内存,使得在资源有限的条件下也能进行高效的模型微调。此外,GRASS的自适应采样策略可以提高微调后的模型性能,使其在下游任务中表现更佳。未来,GRASS有望被广泛应用于各种实际应用中,推动大语言模型在各个领域的应用。
📄 摘要(原文)
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\%.