RevFFN: Memory-Efficient Full-Parameter Fine-Tuning of Mixture-of-Experts LLMs with Reversible Blocks

TL;DR

RevFFN introduces reversible Transformer blocks for mixture-of-experts LLMs, enabling memory-efficient full parameter fine-tuning on standard GPUs by reconstructing activations during backpropagation.

Contribution

It presents a novel reversible Transformer design that reduces memory usage in fine-tuning large language models with mixture-of-experts architecture.

Findings

Significantly reduces peak memory consumption during fine-tuning.

Enables full parameter fine-tuning on single GPU hardware.

Maintains model expressive capacity while improving efficiency.

Abstract

Full parameter fine tuning is a key technique for adapting large language models (LLMs) to downstream tasks, but it incurs substantial memory overhead due to the need to cache extensive intermediate activations for backpropagation. This bottleneck makes full fine tuning of contemporary large scale LLMs challenging in practice. Existing distributed training frameworks such as DeepSpeed alleviate this issue using techniques like ZeRO and FSDP, which rely on multi GPU memory or CPU offloading, but often require additional hardware resources and reduce training speed. We introduce RevFFN, a memory efficient fine tuning paradigm for mixture of experts (MoE) LLMs. RevFFN employs carefully designed reversible Transformer blocks that allow reconstruction of layer input activations from outputs during backpropagation, eliminating the need to store most intermediate activations in memory. While preserving the expressive capacity of MoE architectures, this approach significantly reduces peak memory consumption for full parameter fine tuning. As a result, RevFFN enables efficient full fine tuning on a single consumer grade or server grade GPU.