Full Parameter Fine-tuning for Large Language Models with Limited Resources (2306.09782v2)

Abstract: LLMs have revolutionized 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.

• The paper introduces LOMO, which fuses gradient computation and update to reduce memory usage while fine-tuning all parameters.
• LOMO leverages activation checkpointing and mixed-precision training to achieve substantial memory efficiency and competitive task performance.
• Empirical results reveal that LOMO enables full parameter tuning on large models using affordable hardware, broadening research accessibility.

Full Parameter Fine-tuning for LLMs with Limited Resources

The paper "Full Parameter Fine-tuning for LLMs with Limited Resources" (2306.09782) introduces LOw-Memory Optimization (LOMO) as a novel approach to facilitate full parameter fine-tuning of LLMs under constrained hardware resources. Here, we provide a detailed examination of the techniques proposed, the experimental results obtained, and discuss the broader implications and potential future directions in this research domain.

Introduction and Motivation

The paper addresses the prohibitive computational cost involved in full parameter fine-tuning of LLMs, which traditionally necessitates substantial GPU resources. The motivation stems from the need to democratize access to LLM fine-tuning capabilities for smaller laboratories and researchers without access to high-end hardware. Existing parameter-efficient fine-tuning approaches like LoRA and Prefix-tuning optimize only a subset of parameters, but do not fully realize the potential performance of fine-tuning all parameters in LLMs.

Methodology: Low-Memory Optimization (LOMO)

Optimizer Design

LOMO rethinks the optimizer functionality, proposing a fusion of gradient computation and parameter updates to reduce memory usage:

• Gradient and Parameter Update Fusion: By updating parameters immediately after computing their gradients, LOMO avoids the need to store complete gradient tensors, significantly reducing memory requirements.
• SGD as a Substitute: The paper suggests replacing complex optimizers like Adam with SGD for fine-tuning. It argues that the smooth loss surface of LLMs partially mitigates SGD’s known limitations.

  Figure 1: Comparison of SGD and LOMO in backpropagation and parameter update stages. Pi refers to the parameter of the model and Gi refers to the gradient corresponding to Pi. LOMO fused gradient computation and parameter update in one step to minimize the size of gradient tensors.

Memory Management Strategies

LOMO integrates additional strategies for memory efficiency:

• Activation Checkpointing: It reduces memory load by storing only essential activations during backpropagation.
• Mixed-Precision Training: By employing reduced precision for certain calculations, LOMO optimizes resource utilization without sacrificing model accuracy.

Experimental Evaluation

Memory and Throughput Performance

Empirical results demonstrate that LOMO significantly lowers memory usage compared to conventional solutions, enabling the training of a 65B LLM on a standard machine with multiple RTX 3090 GPUs.

Figure 2: The memory usage ratio of each part when using different optimizers to train LLaMA-7B. The sequence length and batch size are set to 512 and 8, respectively.

Downstream Task Performance

On the SuperGLUE benchmark, LOMO shows competitive or superior performance compared to parameter-efficient methods. It effectively scales to models as large as 65 billion parameters while maintaining high throughput and performance.

Figure 3: Results using LLaMA-13B on the BoolQ and MultiRC datasets (with 1,000 training examples). ``LoRA+LOMO" means injecting LoRA modules while fine-tuning the pre-trained model weights using LOMO.

Implications and Future Directions

LOMO presents a paradigm shift in the manner researchers might access and fine-tune LLMs, potentially broadening participation in NLP research. This work lays the groundwork for future explorations into reducing resource requirements for LLMs further, possibly by integrating quantization techniques.

In terms of future directions, exploring the combination of LOMO with other state-of-the-art memory optimization strategies and evaluating its theoretical robustness could provide additional insights. Moreover, expanding the application of LOMO to real-world scenarios could set a new standard in resource-efficient training workflows.

Conclusion

The proposed LOw-Memory Optimization strategy contributes significantly to the ongoing discourse on optimizing LLM resource usage. By enabling full parameter fine-tuning on less hardware, it paves the way for more inclusive AI research opportunities. Employing LOMO provides a practical balance between computational efficiency and model performance, warranting its consideration for future research initiatives in NLP.