Train, Freeze or Exit: Dynamic Module-wise Fine-Tuning under Memory Constraints

Yile Chen, Zeyi Wen, Jian Chen, Jin Huang

Keywords: fine-tuning, dynamic algorithm configuration

Abstract:

Full fine-tuning delivers strong performance on large language models (LLMs), but its high memory cost limits practical adoption. Although numerous methods have been proposed to reduce memory usage, they rarely model memory cost explicitly and lack the ability for dynamic adjustment. These limitations hinder fine-tuning performance under memory constraints and complicate the selection of suitable configurations.
To address this, we propose Three-State Module Scheduling (TriMS), a dynamic fine-tuning framework that assigns each module in the model to one of three states (trainable, frozen, or early exit). With this three-state formulation, TriMS can quantitatively estimate memory usage while clearly characterizing training configurations. During fine-tuning, TriMS constructs a performance–cost estimator, which is continuously updated by monitoring activation gradients and resource consumption, to evaluate candidate actions (e.g., shrinking or expanding trainable modules). By selecting actions with the best benefit–cost trade-off, TriMS achieves efficient fine-tuning under strict memory budgets.
Extensive experiments across diverse tasks and models demonstrate that TriMS effectively performs dynamic module scheduling under memory constraints. At moderate resource limit (i.e., 80\% of the peak memory required for full fine-tuning), TriMS matches or even outperforms the best baselines, consistently ranking among the top two methods. More importantly, under stricter constraints, where existing approaches often fail to adapt, TriMS maintains strong performance (e.g., achieving accuracy within 1.5\% of full fine-tuning at just 60\% of the memory cost).

Primary Area: optimization

Submission Number: 3549