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

TL;DR

The paper tackles the memory bottleneck of full-parameter fine-tuning for large language models by introducing RevFFN, a reversible Transformer framework tailored for Mixture-of-Experts models. RevFFN uses two-stream reversible blocks and lightweight projection adapters to preserve model capacity while eliminating the need to cache activations, enabling backpropagation through recomputation. A two-stage training schedule—adapter warm-up followed by joint fine-tuning with frozen MoE routing—stabilizes training and yields strong downstream performance. Empirical results on a single GPU show significant memory savings with competitive or superior task performance, offering a practical path to full fine-tuning of large MoE LLMs in modest hardware settings.

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.

Figures (1)

• Figure 1: RevFFN architecture: hidden states split, processed by Cross-Attention and MoE, projected, and concatenated as output.