H i FT : A Hierarchical Full Parameter Fine-Tuning Strategy (2024)

Abstract: Full-parameter fine-tuning (FPFT) has become the go-to choice for adapting language models (LMs) to downstream tasks due to its excellent performance. As LMs grow in size, fine-tuning the full parameters of LMs requires a prohibitively large amount of GPU memory. Existing approaches utilize zeroth-order optimizer to conserve GPU memory, which potentially compromises the performance of LMs as non-zero order optimizers tend to converge more readily on most downstream tasks. We propose a novel, memory-efficient, optimizer-independent, end-to-end hierarchical fine-tuning strategy, HiFT, which only updates a subset of parameters at each training step. HiFT significantly reduces the amount of gradients and optimizer state parameters residing in GPU memory at the same time, thereby reducing GPU memory usage. Our results demonstrate that: (1) HiFT achieves comparable performance with parameter-efficient fine-tuning and standard FPFT. (2) Results on six models show that HiFT reduces the number of trainable parameters by about 89.18% on average compared to FPFT. (3) HiFT supports FPFT of 7B models for 24G GPU memory devices under mixed precision without using any memory saving techniques. (4) HiFT supports various optimizers including AdamW, AdaGrad, SGD, etc. The source code link is https://github.com/misonsky/HiFT.

## introduction full-parameter fine-tuning (fpft) language models (lms) have been a successful paradigm in various downstream tasks @xcite . however, as the size of lms becomes larger, fpft lms require immense memory, which has become an obstacle to conducting research. one line of research to reduce memory is to use heterogeneous memory @xcite (e.g., gpu, cpu, * corresponding author and nvme memory) or distributed techniques (e.g., tensor parallelism @xcite @xcite @xcite ). these strategies require parameter sharing across diverse devices and thus usually introduce a significant communication burden. parameter-efficient fine-tuning (peft) is another line of strategies for memory reduction, categorized into addition-based, selection-based, and reparametrization-based methods @xcite . the addition-based methods (e.g., @xcite , attentionfusion @xcite ) reduce the number of trainable parameters by only updating newly added parameters and freezing the weights of lms. although these methods reduce the number of parameters for fine-tuning, they expand the number of model parameters and increase the burden on forward propagation. the selection-based methods (e.g, bitfit @xcite , lt-sft @xcite , far @xcite ), on the other hand, fine-tune a subset of model parameters, resulting in a performance gap with fpft. the reparametrization-based methods (e.g., lora @xcite , krona @xcite , s4model @xcite ) leverage low-rank decomposition to minimize the number of trainable parameters. using low-rank representations inevitably leads to information loss and performance degradation. peft involves a trade-off between serving efficiency and quality. according to existing works @xcite @xcite , fpft still maintains advantages in performance on most benchmarks. some works have reduced memory usage for fpft by removing the momentum state of the optimizer. lomo @xcite reduces the memory usage of the optimizer momentum and gradients by integrating gradient calculation and update. nevertheless, lomo requires forward propagation twice. in addition, lomo forces the model to be 16-bit layer-by-layer fine-tuning quantized and uses the gradient checkpointing technique @xcite to reduce memory usage while lomo has limited memory savings in realworld scenarios. mezo @xcite designs a zeroth-order optimizer to reduce memory usage. however, mezo is unstable and performs poorly without prompts. these methods make the momentum optimizers unusable, while the momentum optimizers such as adamw @xcite have been proven to be superior in improving performance. in this paper, we propose a novel memoryefficient hierarchical fine-tuning (hift) strategy, adopting the idea of block-by-block training. hift divides the layers of the model into different groups (a group is a block.). at each training step, hift updates the parameters of one group while freezing the others. compared to standard fpft, hift leads to different groups of parameters being updated with different learning rates. this causes the model parameters to be updated in an inconsistent amplitude, which leads to a decrease in model performance. to solve this problem, we adopt to delay the learning rate update, which only updates the learning rate once when all layers of the model are updated. hift is also different from layer-wise training @xcite , where the layer-wise training incrementally adds new layers to a pre-trained shallow model, only updating the newly added parameters at each training stage until all layers are updated. as a result, the layer-wise strategy produces accumulated errors at different training stages due to its pipeline training. hift can significantly reduce the number of trainable parameters per training step. we only keep the momentum and gradients of the parameters that need to be updated on the gpu device due to only a portion of the parameters are updated at each training step. this helps to reduce the gpu memory usage of the optimizer states and gradients. hift supports full-parameter fine-tuning of a 7b model on devices with 24g memory. our contributions are summarized as follows: ## related work full-parameter fine-tuning fpft fine-tunes the pre-trained lms on specific tasks by updating all parameters @xcite @xcite , which requires massive computing power as the parameters of lms increase. mixedprecision training enables high-throughput computations by employing half-precision storage for parameters, activations, and gradients @xcite . staged training incrementally increases the amount of compute and reuse the compute from prior stages @xcite . these methods increase the parameter consumption when training precision or operators. lomo @xcite identifies the memory saving of sgd @xcite , fuses the gradient computation and the parameter update in one step. mezo @xcite designs a gradient-free method to update the model. although it can reduce memory usage, its performance has a big gap than fpft, especially when there is no prompt. these methods waste the superiority of momentum optimizers. parameter-efficient fine-tuning peft minimizes resource utilization from the perspective of parameters with additon, selection or decomposition methods @xcite . the additionbased methods add and update new parameters with the weights of lms frozen, such as prefix-tuning (li and liang, 2021), attentionfusion @xcite , while the added parameters increase the burden on forward propagation. the selectionbased methods fine-tune a subset of the parameters of lms, such as bitfit @xcite , lt-sft @xcite , far @xcite , but has a performance gap with fpft. the reparametrization-based methods leverage lowrank decomposition to minimize the number of trainable parameters, such as lora @xcite , phm @xcite , krona @xcite , s4-model @xcite , while using low-rank representations inevitably leads to information loss and performance degradation. peft involves a trade-off between serving efficiency and quality. memory-efficient fine-tuning meft minimizes memory usage with heterogeneous memory (e.g., gpu, cpu and nvme) or parallel methods (e.g., tensor and pipeline parallelism). in a layer-to-layer strategy @xcite , only the tensors necessary for the computation of a particular layer are transferred to gpu, while the remaining tensors are retained in cpu. zero-infinity @xcite enables the partitioned states and tensors to cpu and nvme. tensor parallelism accelerates training by parallelizing tensor computations across different gpus, but requires multiple global communications during each propagation @xcite . pipeline parallelism accelerates training by breaking the model into segments or layers and processing them sequentially in a pipeline fashion @xcite @xcite . these methods transfer massive memory to heterogeneous devices, although temporarily saving memory, still requires a large number of devices. different from existing works @xcite , hift adopts the idea of blockby-block training to save memory of fpft, and can be seamlessly integrated with any optimizer. ## approach the training strategy of our hift is shown in figure 1 . we first present some necessary notations. notation given the training dataset @xmath0, y i )} n @xmath1, the goal of the training is to learn a model m with n layers, where n is the number of the training samples, (x i , y i ) is the labeled data pair. we use p to represent the optimizer, and η t to represent the learning rate schedule. the number of layers in each group is represented by m and the number of groups is represented by k. if m is divisible by n, then @xmath2, otherwise @xmath3. queue q is used to store special identifiers that uniquely identify different layers. @xmath4,"top2down","random"} represents the adopted update strategy. consider a pre-trained lm f θpre parameterized by θ pre . let θ f pf t and θ hif t denote parameters after full fine-tuning and hierarchical full-parameter finetuning after one training step, respectively. let l τ (d; θ) be the objective to minimize during finetuning, with d being the input, θ being updated parameters, and τ being the task in fine-tuning. in the process of full fine-tuning, we optimize the model by adjusting its full parameters: in the process of hift, only a subset of parameters are updated at one training step. more formally, with optimizing group @xmath5, ..., k}, we have: ## experiments please refer to appendix for baselines(c), datasets (d) and implementation details (f).