Source-linked AI summary
Full Parameter Fine-tuning for Large Language Models with Limited Resources
Kai Lv, Yuqing Yang, Tengxiao Liu, Qinghui Gao, Qipeng Guo, Xipeng Qiu
TL;DR
Full-parameter fine-tuning of large language models requires substantial GPU resources, while existing parameter-efficient methods do not address this setting directly. The paper introduces LOMO, which fuses gradient computation with parameter updates to reduce memory usage, and reports feasibility for fine-tuning a 65B model on consumer GPUs. The authors also evaluate memory, throughput, and downstream performance, while noting limitations in gradient normalization speed, benchmark coverage, and GPU evaluation.
Problem
Full-parameter fine-tuning of large language models remains difficult under limited resources, whereas existing parameter-efficient methods do not offer a practical solution for tuning all parameters.
Method
LOMO fuses gradient computation and parameter updates, while the training process also uses SGD-based optimization and mixed-precision stabilization techniques.
Results
LOMO enables successful full-parameter fine-tuning of a 65B model on a server equipped with 8 RTX 3090 GPUs and shows effectiveness through memory, throughput, and SuperGLUE experiments.
Takeaways & Limitations
LOMO facilitates full-parameter fine-tuning of large language models with limited resources while keeping the parameter update equivalent to SGD.
Takeaways & Limitations
Gradient normalization may require an additional backward pass that slows training, experiments covered only part of SuperGLUE, and throughput was not evaluated on A100 GPUs.
Abstract
from arXiv · showhide
Large Language Models (LLMs) have revolutionized Natural Language Processing (NLP) but demand massive GPU resources for training. Lowering the threshold for LLMs training would encourage greater participation from researchers, benefiting both academia and society. While existing approaches have focused on parameter-efficient fine-tuning, which tunes or adds a small number of parameters, few have addressed the challenge of tuning the full parameters of LLMs with limited resources. In this work, we propose a new optimizer, LOw-Memory Optimization (LOMO), which fuses the gradient computation and the parameter update in one step to reduce memory usage. By integrating LOMO with existing memory saving techniques, we reduce memory usage to 10.8% compared to the standard approach (DeepSpeed solution). Consequently, our approach enables the full parameter fine-tuning of a 65B model on a single machine with 8 RTX 3090, each with 24GB memory.Code and data are available at https://github.com/OpenLMLab/LOMO.
1 Introduction
The paper targets full-parameter LLM fine-tuning under resource constraints by introducing LOMO, which reduces memory use by fusing gradient computation and parameter updates. The approach supports large-model training while preserving an SGD-equivalent update process.
- Motivation: 30B to 175B-parameter LLMs often require expensive GPU resources, limiting participation by small labs and companies.The passage gives 8×80GB devices as an example of the required resources.
- Research gap: Parameter-efficient methods such as LoRA and Prefix-tuning do not provide a practical solution for full-parameter fine-tuning.The paper frames full-parameter fine-tuning as a more powerful approach and investigates resource-limited alternatives.
- LOMO approach: LOMO reduces gradient-tensor memory to O(1) by fusing gradient computation and parameter updates in one step.The optimizer is designed to minimize the size of stored gradient tensors during training.
- Memory reduction: The resulting memory usage equals parameters plus activations and the largest gradient tensor, approaching the memory required for inference.The parameter update remains equivalent to SGD, so the fine-tuning process is not changed by the memory optimization.
- Evaluation: 65B-model training succeeds on a server with only 8 RTX 3090 GPUs, while downstream experiments evaluate full-parameter tuning on SuperGLUE.The paper reports empirical assessments of memory, throughput, and downstream performance.
- Optimization rationale: The paper theoretically argues that SGD can successfully fine-tune full-parameter LLMs because prior obstacles may be less severe for this setting.This argument supports removing optimizer states as part of the memory-saving strategy.
2 Related Work
Related work reduces fine-tuning memory through activation checkpointing, mixed precision, heterogeneous memory, and memory-efficient optimizers. These approaches trade computation, precision management, data movement, or gradient approximation against memory consumption.
- Combination with LOMO: The related techniques are orthogonal to LOMO and can be combined with it to further reduce memory consumption.This combination is described as a way to extend LOMO’s memory savings.
- Activation checkpointing: Activation checkpointing saves memory by recomputing discarded activations, trading lower storage for additional computation.It is presented as a compromise between storing all activations and recomputing everything.
- Mixed-precision training: Mixed-precision training lowers memory use and accelerates computation through half-precision storage, with full-precision copies, loss scaling, and selected full-precision operations supporting stability.These techniques apply to parameters, activations, and gradients.
- Heterogeneous training: Heterogeneous training systems reduce GPU memory consumption by retaining tensors in CPU or NVMe memory and transferring only those needed for computation.L2L and ZeRO-Offload exemplify this strategy.
- Memory-efficient optimization: Recent optimizers save memory through zero-order gradient estimation, low-rank gradient approximation, or quantized optimizer states.The passage contrasts these approaches with LOMO, which does not approximate gradients in the described comparison.
3 Method
The method targets full-parameter LLM fine-tuning under limited resources by replacing state-heavy optimization with SGD and fusing gradient computation with parameter updates. LOMO reduces gradient-memory requirements while preserving an SGD-equivalent update, and combines with complementary memory-saving techniques.
- Using SGD: Optimizer states can dominate LLM training memory because Adam stores intermediate states twice the size of the parameters.
- Using SGD: SGD is presented as an acceptable fine-tuning optimizer when the scope is limited to LLM adaptation, despite no guarantee of superiority over modern optimizers.The rationale relies on assumptions about smooth loss surfaces, acceptable local optima, and distant saddle points.
- Implicit Batch Size: Smooth loss surfaces make sequential SGD updates approximately equivalent to larger-batch updates, supporting stable LLM fine-tuning.The analysis attributes this stability to smoothness and the stronger stability associated with larger batch sizes.
- LOMO: LOw-Memory Optimization: LOMO fuses gradient computation and parameter updates, storing at most one parameter gradient during backpropagation instead of all gradient tensors.The algorithm updates parameters layer by layer during backward propagation and then clears each gradient.
- LOMO: LOw-Memory Optimization: LOMO's memory usage approaches parameter, activation, and largest-gradient-tensor storage, making full fine-tuning close to inference memory without changing the SGD-equivalent update.Combining LOMO with parameter-efficient fine-tuning methods adds only a minor gradient-memory increase and enables tuning more parameters.
- Alternatives to Gradient Normalization and Clipping: Value-based gradient clipping avoids storing all gradients but can alter gradient direction and performs worse at high learning rates.The authors generally suggest this approach for learning rates below 1e-3; an additional backward pass can compute the gradient norm but sacrifices speed.
4 Experiment
Experiments evaluate LOMO's memory profile, throughput, downstream performance, scalability, and compatibility with LoRA. Across these settings, LOMO substantially reduces memory use, accelerates training, enables larger-model training, and generally maintains strong downstream results.
- 4.1 Memory Profile: 102.20GB to 14.58GB: LOMO substantially reduces LLaMA-7B memory usage versus AdamW, while also reducing it from 51.99GB versus SGD.The reduction is primarily attributed to lower gradient and optimizer-state requirements.
- 4.1 Memory Profile: 73.7% of AdamW's LLaMA-7B memory is allocated to optimizer states, which SGD reduces by avoiding full-precision momentums and variances.LOMO further eliminates optimizer-state memory by fusing parameter updates with backpropagation.
- 4.1 Memory Profile: 45.61GB to 1.79GB: activation checkpointing combined with LOMO reduces activation memory for 7B-model training with 512×8 tokens per batch.LOMO is compatible with activation memory reduction techniques such as activation checkpointing.
- 4.2 Throughput: About 11 times: LOMO surpasses AdamW and SGD in 7B throughput by training the model on a single GPU and reducing inter-GPU communication.Throughput is measured as tokens processed per GPU per second under sequence length 1024 and batch size 1.
- 4.2 Throughput: LOMO trains 30B models on 4 GPUs and successfully trains a 65B model on 8 RTX 3090 GPUs, achieving 4.93 TGS; 1,000 samples require approximately 3.6 hours.With 13B models, LOMO also outperforms SGD when model parallelism is necessary.
- 4.3 Downstream Performance: Across six SuperGLUE datasets and model sizes, LOMO consistently outperforms Zero-shot and generally outperforms LoRA, with average gains exceeding 20 points and reaching 2.8 points using LLaMA-13B.In some cases LOMO performs worse than LoRA, potentially because 1,000 training examples may be insufficient for full-parameter fine-tuning of large models.
5 Conclusion
The paper introduces LOMO to enable full-parameter fine-tuning of large language models with limited resources, demonstrating feasibility on a 65B model using consumer GPUs. It also identifies parameter quantization and broader applicability as future directions.
- LOMO is introduced as an optimizer designed to facilitate full-parameter fine-tuning of large language models with limited resources.
- A 65B model was fine-tuned on a server equipped with consumer RTX 3090 GPUs.
- The evaluation covered memory usage, throughput, and downstream performance on SuperGLUE.
- Future work targets parameter quantization to reduce memory usage and investigates additional applicable scenarios for LOMO.
Limitations
The paper identifies training-speed and evaluation-scope limitations: gradient normalization may require an additional backward pass, and experiments omitted some benchmark and hardware settings.
- An additional backward pass for gradient normalization can slow LOMO when gradient normalization is essential.
- Experiments covered only a subset of SuperGLUE and did not evaluate LOMO throughput on advanced GPUs such as A100.
Ethics statement
The paper uses open-source LLaMA models and publicly permitted datasets in accordance with their respective licenses.
- The study uses open-source LLaMA models and datasets whose licenses permit public use.
A Hyperparameters
The paper reports that experimental hyperparameters are listed in Table 4 and that results use experiments sharing the same random seed.
- Experimental hyperparameters are listed in Table 4.
- Reported results are the highest results from experiments conducted with the same random seed.
B Training Dynamics
LOMO training on LLaMA-7B BoolQ converges rapidly at first, then stabilizes and gradually declines in loss. Development-set accuracy generally rises with training steps.
- Training loss: LOMO’s training loss on LLaMA-7B BoolQ converges rapidly initially, then stabilizes and gradually declines.The corresponding curve is identified as Figure 4.
- Validation accuracy: Development-set accuracy generally increases as the number of training steps grows.The corresponding validation-accuracy plot is identified as Figure 5.
- Training Dynamics: The training-dynamics analysis uses loss and validation-accuracy curves for LLaMA-7B trained on BoolQ.These curves are presented in Figures 4 and 5, respectively.