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

Abstract: Full parameter fine tuning is a key technique for adapting 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.

• The paper introduces a novel reversible transformer module that reconstructs intermediate activations to reduce memory usage during full-parameter fine-tuning of MoE LLMs.
• It employs a two-stage training curriculum, combining adapter warm-up and joint fine-tuning, to align representation spaces and maximize model performance.
• Empirical findings show that RevFFN achieves competitive task performance on benchmarks such as MMLU while reducing VRAM usage by approximately 49% compared to activation checkpointing.

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

Introduction and Motivation

The ever-increasing scale of LLMs—driven by architectures such as GPT-4, LLaMA, and Qwen—has elevated their generalization capacity but also exacerbated the memory bottlenecks confronted during full fine-tuning. Conventional techniques, notably Adam-based backpropagation, necessitate the retention of layer-wise activations whose footprint grows with both model parameter count and batch size. Distributed optimization frameworks like DeepSpeed (ZeRO) and FSDP mitigate per-GPU VRAM requirements via sharding/offloading, but at the expense of higher hardware complexity and decreased throughput due to network overhead. PEFT methods—LoRA, DoRA, (IA)$^3$—reduce trainable parameter count but inherently restrict task-adaptive model capacity.

RevFFN addresses this central limitation by fundamentally redesigning the Transformer module through reversible architectures, enabling reconstruction rather than caching of intermediate activations. This paradigm shift facilitates full-parameter fine-tuning of contemporary Mixture-of-Experts models on a single commodity GPU, with significant memory savings.

RevFFN Architecture

RevFFN is predicated on the principles of reversible neural networks and adapts them for transformers and MoE LLMs. Each hidden state $$H \in \mathbb{R}^{B \times S \times d_\text{model}}$$ is split into two streams, updated in a coupled and asymmetrical fashion:

• $X_1$ is updated using cross-branch attention, where queries are sourced from $X_1$ and keys/values from $X_2$.
• $X_2$ is updated using an MoE-driven feed-forward sub-layer that leverages the normalized output of $Y_1$.

The concatenated output $[Y_1, Y_2]$ enables exact reconstruction via a bijective update rule, supporting on-the-fly backward pass recomputation.

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

A key challenge arises as pre-trained Qwen1.5-MoE weights expect $d_\text{model}$-wide inputs, while each stream only provides $d_\text{model}/2$ features. Lightweight projection adapters ($P_\uparrow$ and $P_\downarrow$) address this dimensionality mismatch, preserving model capacity and requiring negligible new parameters.

Two-Stage Training Curriculum

RevFFN employs a curriculum designed to align representation spaces and exploit backbone capacity without catastrophic forgetting:

1. Adapter Warm-Up: All pre-trained MoE weights are frozen while adapters $P_\uparrow$, $P_\downarrow$ are trained with a low learning rate. This phase ensures stable bridging between reversible and backbone spaces.
2. Joint Fine-Tuning: The entire model is unfrozen (except MoE routers), enabling gradient flow through all parameters for maximal expressivity.

Ablation experiments conclusively demonstrate that omitting either stage results in substantial performance degradation, highlighting the necessity of this schedule.

Empirical Evaluation

Benchmarks and Baselines: RevFFN is evaluated on Qwen1.5-MoE-A2.7B with the databricks-dolly-15k dataset. Competing approaches span PEFT (LoRA, DoRA, (IA)$^3$) and memory-efficient full tuning (Activation Checkpointing, LoMo, GaLore).

Memory and Throughput: RevFFN exhibits a ~49% reduction in peak VRAM relative to activation checkpointing (39.5GB vs. 65.4GB) and is more efficient than GaLore or LoMo, though PEFT methods remain fastest due to architectural minimalist updates. The recomputation overhead intrinsic to reversible networks yields lower throughput than PEFT but remains competitive versus other full tuning schemes.

Downstream Task Results: On MMLU, GSM8K, Multilingual benchmarks, and MT-Bench, RevFFN marginally surpasses all baselines, indicating that memory-efficient reversibility does not compromise, and may even enhance, task adaptation. Performance gaps compared with PEFT highlight the criticality of updating all parameters for high-complexity tasks.

Ablation and Component Analysis

The study isolates the contribution of each training stage. Adapter-only, projection-based fine-tuning is insufficient, yielding far lower benchmark scores (e.g., 54.5% on MMLU). Joint training from scratch destabilizes optimization and also underperforms (57.1%). The full two-stage method attains 66.7% on MMLU—establishing the necessity of sequential adapter alignment and global parameter updates.

Implications and Future Directions

RevFFN's architecture makes single-GPU fine-tuning of MoE LLMs viable, drastically lowering the operational barrier for full-parameter updates. This has broad implications for democratizing large-scale LLM customization, especially in resource-constrained contexts. The reversibility principle may synergize with quantization, offloading, or cross-modal transfer to further optimize the memory-performance frontier. Future research could extend RevFFN to trillion-parameter models, integrate advanced curriculum strategies, or combine with knowledge distillation for training acceleration and model compression.

Conclusion

RevFFN demonstrates that memory-efficient full fine-tuning of Mixture-of-Experts LLMs is achievable via reversible transformer blocks and adaptive projection layers. The approach attains strong benchmark performance, substantially reduces memory requirements, and removes the multi-GPU constraint for large model adaptation. This establishes a promising direction for scalable, cost-effective LLM deployment and further underlines the utility of reversible design for resource-bound optimization.

(2512.20920)