--- title: 'RevFFN: Memory-Efficient Fine-Tuning for MoE LLMs' url: https://www.emergentmind.com/topics/reversible-ffn-for-moe-llms type: topic --- # RevFFN: Memory-Efficient Fine-Tuning for MoE LLMs RevFFN is a memory-efficient paradigm for full-parameter fine-tuning of Mixture-of-Experts (MoE) large language models (LLMs) utilizing reversible Transformer block architectures. It addresses the activation memory bottleneck inherent in conventional fine-tuning approaches by enabling input reconstruction from outputs during the backward pass, thereby eliminating the need to store intermediate activations. This mechanism significantly reduces peak VRAM requirements and enables single-GPU training for large-scale MoE LLMs without sacrificing expressive capacity or downstream performance [2512.20920]. ## 1. Architectural Principles ### 1.1 Standard MoE Transformer Layer Traditional Transformer decoder layers consist of two sublayers: multi-head self-attention with residual connections and a feed-forward network (FFN). The residual formulation is: - Self-attention: $H' = H + \mathrm{Attn}(\mathrm{LN}(H), \mathrm{LN}(H), \mathrm{LN}(H))$ - Feed-forward: $H_{out} = H' + \mathrm{FFN}(\mathrm{LN}(H'))$, with $\mathrm{FFN}(x) = W_2 \sigma(W_1 x)$ The MoE variant replaces the FFN with a sparsely-gated expert layer. A gating network $g(x) = \mathrm{softmax}(W_g x) \in \mathbb{R}^E$ assigns each token to the top-$k$ experts, with individual two-layer MLPs per expert $E_e(x) = W_2^{(e)} \sigma(W_1^{(e)} x)$, aggregated as $F_{MoE}(x) = \sum_{e=1}^E g_e(x) E_e(x)$. All computation occurs in the full model dimension $d_{model}$. ### 1.2 Reversibility in Residual Blocks Conventional residual blocks require caching input activations for backpropagation, incurring $O(LBSd_{model})$ memory for $L$ layers, batch size $B$, sequence length $S$, and model dimension $d_{model}$. A reversible block implements a bijective mapping $y \leftrightarrow x$, enabling input reconstruction during the backward pass: - Forward: $y = x + F(x)$ - Inverse: $x = y - F(x)$ RevFFN employs a two-stream coupling method, splitting activations into halves and ensuring exact invertibility. ## 2. Reversible MoE Block Construction ### 2.1 Formulation The hidden tensor $H \in \mathbb{R}^{B \times S \times d_{model}}$ is partitioned as $X = [X_1, X_2]$ with $X_i \in \mathbb{R}^{B \times S \times (d_{model}/2)}$. The reversible update equations for a decoder layer are: 1. $Y_1 = X_1 + \mathrm{Attn}(\mathrm{Norm}(X_1), \mathrm{Norm}(X_2), \mathrm{Norm}(X_2))$ 2. $Y_2 = X_2 + F_{MoE}(\mathrm{Norm}(Y_1))$ 3. $H_{out} = [Y_1, Y_2]$ Inverse mapping is defined as: 1. $\hat X_2 = Y_2 - F_{MoE}(\mathrm{Norm}(Y_1))$ 2. $\hat X_1 = Y_1 - \mathrm{Attn}(\mathrm{Norm}(\hat X_1), \mathrm{Norm}(\hat X_2), \mathrm{Norm}(\hat X_2))$ A single fixed-point iteration initialized at $\hat X_1^{(0)} = Y_1$ achieves machine-precision convergence. ### 2.2 MoE Feed-Forward Layer Structure For $Z = \mathrm{Norm}(Y_1)$: - Routing: $\ell = W_g Z$, $g = \mathrm{softmax}(\ell) \in \mathbb{R}^E$ - Experts: $E_e(Z) = W_2^{(e)} \sigma(W_1^{(e)} Z)$ - Aggregation: $F_{MoE}(Z) = \sum_{e=1}^E g_e \cdot E_e(Z)$ To maintain compatibility with pre-trained MoE modules, inputs are projected via adapter matrices $P_\uparrow \in \mathbb{R}^{(d/2) \times d}$ and $P_\downarrow \in \mathbb{R}^{d \times (d/2)}$, yielding $F_{MoE}(Z) = P_\downarrow\ \cdot\ \mathrm{MoE}_{pretrained}(P_\uparrow Z)$. ## 3. Memory Savings and Activation Reconstruction ### 3.1 Back-Propagation Strategy Standard layers require storing $X_1, X_2$, and LayerNorm inputs for gradient calculations. In RevFFN, the backward pass proceeds as: 1. Reconstruct $X_1, X_2$ from $Y_1, Y_2$ using the inverse mapping. 2. Re-execute LayerNorm, Attention, and MoE blocks to recreate required intermediates. 3. Compute gradients with respect to parameters and inputs using chain-rule. Each sublayer is executed twice per step (forward and backward), trading memory savings for compute overhead. ### 3.2 Memory Complexity | Method | Memory Complexity | |---------------------------------|-------------------------------| | Standard fine-tuning | $O(L \cdot B \cdot S \cdot d)$| | RevFFN reversible architecture | $O(B \cdot S \cdot d) + O(1)$ | RevFFN eliminates the dependence on layer count $L$—activation memory scales only with batch size, sequence length, and model dimension. ## 4. Training Modifications and Performance ### 4.1 Backward Hook Implementation RevFFN requires a custom backward hook: - At each reversible block, output activations $Y_1, Y_2$ are popped. - Inverse mapping reconstructs $X_1, X_2$. - Forward sublayers are re-executed to materialize intermediates (LayerNorm, Attn, MoE). - Autograd applies the chain-rule for gradients with respect to model parameters and inputs. Only expert parameters $W_1^{(e)}, W_2^{(e)}$ and adapters $P_\uparrow, P_\downarrow$ are updated; the gating network remains frozen during fine-tuning. ### 4.2 Computational Overhead Each reversible layer incurs roughly 2$\times$ FLOPs compared to the standard layer (due to re-execution in backward). In practice, MoE computation is dominant, resulting in $\approx$ 20–30% training overhead: - Throughput drops from 31.0 to 24.6 samples/s on NVIDIA H800. | Method | Peak VRAM (GB) | Throughput (samples/s) | |---------------------|:--------------:|:----------------------:| | SFT + Checkpointing | 65.4 | 19.7 | | GaLore | 45.1 | 35.2 | | RevFFN | **39.5** | **24.6** | ## 5. Empirical Validation Downstream task performance is evaluated on MMLU, GSM8K, MT-Bench, and a multilingual benchmark: | Method | MMLU | GSM8K | Multilingual | MT-Bench | |---------------------|:-----:|:-----:|:------------:|:--------:| | SFT + Checkpointing | 66.1% | 74.8% | 39.5% | 7.52 | | RevFFN |**66.7%**|**75.1%**|38.8%| **7.65** | RevFFN provides a $\approx$ 49% reduction in peak memory (vs. SFT+Checkpointing), with task accuracy matching or slightly exceeding baseline methods. Ablation studies confirm that both stages of the two-stage schedule are essential for training stability and optimal performance. ## 6. Usage Scenarios and Implementation Recommendations ### 6.1 Application Context RevFFN is indicated in settings with VRAM constraints under $\sim$80GB and model scales ranging from several to tens of billions of parameters. It is suited for scenarios where full-parameter adaptation is required and multi-GPU or CPU offloading is infeasible or suboptimal. ### 6.2 PyTorch Implementation The reversible block is constructed as follows: ```python class RevFFNBlock(torch.nn.Module): def __init__(self, pre_attn, pre_moe, d): super().__init__() self.p_up = nn.Linear(d//2, d) self.p_down = nn.Linear(d, d//2) self.attn = pre_attn self.moe = pre_moe self.norm1 = nn.LayerNorm(d//2) self.norm2 = nn.LayerNorm(d//2) def forward(self, H): X1, X2 = H.chunk(2, dim=-1) q = self.p_up(self.norm1(X1)) k = self.p_up(self.norm1(X2)) v = self.p_up(self.norm1(X2)) A = self.attn(q, k, v) Y1 = X1 + self.p_down(A) z = self.p_up(self.norm1(Y1)) M = self.moe(z) Y2 = X2 + self.p_down(M) return torch.cat([Y1, Y2], dim=-1) def backward_hook(self, grad_Y): # Custom backward logic (omitted) pass ``` Backward hooks must be registered to: 1. Pop output activations, 2. Run inverse mapping, 3. Re-execute submodules to recover intermediates, 4. Apply autograd for parameter gradients. ## 7. Summary and Implications RevFFN refactors Transformer decoder layers into reversible two-stream blocks, maintaining full MoE routing and expert computation across the model dimension. It achieves comparable downstream task performance to standard full fine-tuning while providing significant reductions in peak memory usage—enabling practical, single-GPU training of billion-parameter MoE LLMs by trading a moderate computational overhead ($\sim$2$\times$ per block) for an effective 2$\times$ reduction in activation memory [2512.20920]. This suggests broader applicability of reversible computation techniques for efficient adaptation of modern LLM architectures lacking distributed infrastructure. Source: https://www.emergentmind.com/topics/reversible-ffn-for-moe-llms