DropBP: Accelerating Fine-Tuning of Large Language Models by Dropping Backward Propagation

TL;DR

DropBP addresses the high computational and activation-memory cost of fine-tuning LLMs by randomly dropping layers during backward propagation, effectively training shallow submodules while keeping forward paths intact. Layer sensitivity S_l guides a greedy allocation of per-layer drop rates to satisfy a FLOPs budget, enabling stable training. Empirical results show DropBP maintains near-baseline accuracy on MMLU, commonsense, and MT-Bench tasks across LLaMA2 and LLaMA3-8B with substantial speedups (up to $57\%$) and memory reductions (down to $32$GB) on a single NVIDIA A100, plus up to $6.2\times$ longer context and significant throughput gains. The method is orthogonal to PEFT, easy to integrate as a PyTorch extension, and enhances usability for long-context fine-tuning on resource-constrained hardware. Overall, DropBP offers a practical, scalable approach to speeding up fine-tuning of large language models without sacrificing accuracy.

Abstract

Large language models (LLMs) have achieved significant success across various domains. However, training these LLMs typically involves substantial memory and computational costs during both forward and backward propagation. While parameter-efficient fine-tuning (PEFT) considerably reduces the training memory associated with parameters, it does not address the significant computational costs and activation memory. In this paper, we propose Dropping Backward Propagation (DropBP), a novel approach designed to reduce computational costs and activation memory while maintaining accuracy. DropBP randomly drops layers during backward propagation, which is essentially equivalent to training shallow submodules generated by undropped layers and residual connections. Additionally, DropBP calculates the sensitivity of each layer to assign an appropriate drop rate, thereby stabilizing the training process. DropBP is not only applicable to full fine-tuning but can also be orthogonally integrated with all types of PEFT by dropping layers during backward propagation. Specifically, DropBP can reduce training time by 44% with comparable accuracy to the baseline, accelerate convergence to the same perplexity by 1.5x, and enable training with a sequence length 6.2x larger on a single NVIDIA-A100 GPU. Furthermore, our DropBP enabled a throughput increase of 79% on a NVIDIA A100 GPU and 117% on an Intel Gaudi2 HPU. The code is available at https://github.com/WooSunghyeon/dropbp.

Figures (14)

• Figure 1: Performance enhancements in fine-tuning large language models using DropBP when the $p$ represents the target average drop rate for backward propagation: (a) Training time per sample for fine-tuning LLaMA2-7B with DropBP, at a sequence length of 512 and a micro batch size of 2. (b) Available max sequence length for fine-tuning LLaMA2-70B with DropBP, at a micro batch size of 1 on an NVIDIA-A100 GPU.
• Figure 2: Interpreting the model with residual connections as a combination of multiple submodules.
• Figure 3: The overveiw of DropBP.
• Figure 4: Code implementation for integrating DropBP.
• Figure 5: Validation perplexity (PPL) for fine-tuning LLaMA2-70B through QLoRA (baseline) with DropBP on the Alpaca dataset. The $p$ represents the target average drop rate for backward propagation.