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

📄 arXiv: 2512.20920v1 📥 PDF

作者: Ningyuan Liu, Jing Yang, Kaitong Cai, Keze Wang

分类: cs.LG, cs.AI

发布日期: 2025-12-24

备注: Under submission

💡 一句话要点

RevFFN：利用可逆块实现MoE LLM全参数高效微调

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 全参数微调 混合专家模型 可逆神经网络 内存优化

📋 核心要点

  1. 现有全参数微调方法面临内存瓶颈，需要缓存大量中间激活值，限制了大型语言模型在消费级硬件上的应用。
  2. RevFFN通过引入可逆Transformer块，在反向传播时从输出重建输入激活，避免存储大量中间激活值，降低内存消耗。
  3. RevFFN在保持模型表达能力的同时，实现了在单GPU上对MoE LLM进行全参数微调，提升了训练效率。

📝 摘要（中文）

全参数微调是使大型语言模型(LLMs)适应下游任务的关键技术，但由于需要缓存大量的中间激活值以进行反向传播，因此会产生巨大的内存开销。这种瓶颈使得对当前大规模LLM进行全参数微调在实践中具有挑战性。现有的分布式训练框架，如DeepSpeed，使用ZeRO和FSDP等技术来缓解这个问题，这些技术依赖于多GPU内存或CPU卸载，但通常需要额外的硬件资源并降低训练速度。我们引入了RevFFN，一种用于混合专家(MoE) LLM的内存高效微调范例。RevFFN采用精心设计的可逆Transformer块，允许在反向传播期间从输出重建层输入激活，从而消除了在内存中存储大多数中间激活的需要。在保留MoE架构的表达能力的同时，这种方法显著降低了全参数微调的峰值内存消耗。因此，RevFFN能够在单个消费级或服务器级GPU上实现高效的全参数微调。

🔬 方法详解

问题定义：全参数微调能够有效提升LLM在特定任务上的性能，但其巨大的内存需求限制了在资源受限环境下的应用，尤其是在MoE模型中，参数量巨大，中间激活值的存储成为瓶颈。现有方法如ZeRO和FSDP虽然能缓解内存问题，但依赖多GPU或CPU卸载，增加了硬件成本和通信开销。

核心思路：RevFFN的核心在于利用可逆Transformer块，使得每一层的输入激活可以从输出激活中精确重建。这样，在反向传播过程中，无需存储大量的中间激活值，从而显著降低了内存占用。这种方法在不牺牲模型表达能力的前提下，实现了内存效率的提升。

技术框架：RevFFN采用改进的Transformer架构，核心是可逆Transformer块。在正向传播过程中，计算每一层的输出，但只保留最后一层的输出。在反向传播过程中，从最后一层的输出开始，利用可逆块的特性，逐层重建每一层的输入激活，然后进行梯度计算。整体流程与标准Transformer类似，但内存占用显著降低。

关键创新：RevFFN最重要的创新在于可逆Transformer块的设计，它允许从输出精确重建输入，从而避免了存储中间激活值的需求。与传统的Transformer块相比，RevFFN在内存效率上具有显著优势，尤其是在大规模模型和全参数微调场景下。

关键设计：RevFFN的关键设计包括：1) 可逆Transformer块的具体结构，需要保证输入和输出之间的可逆性；2) 损失函数的设计，需要考虑可逆性对训练的影响；3) 参数初始化策略，需要保证模型在训练初期的稳定性。具体的可逆块实现细节（例如，如何进行输入输出的拆分和重组）以及损失函数的选择（例如，是否需要额外的正则化项来保证可逆性）需要在论文中详细描述。

🖼️ 关键图片

fig_0

📊 实验亮点

RevFFN通过可逆块设计，显著降低了MoE LLM全参数微调的内存消耗，使得在单GPU上进行全参数微调成为可能。实验结果表明，RevFFN在保持模型性能的同时，能够大幅降低峰值内存占用，提升训练效率。具体的性能数据和对比基线需要在论文中给出。

🎯 应用场景

RevFFN可应用于各种需要对大型语言模型进行全参数微调的场景，例如自然语言处理、机器翻译、文本生成等。它降低了微调的硬件门槛，使得研究人员和开发者能够在消费级或服务器级GPU上高效地进行模型定制和优化，加速了LLM在各个领域的应用。

📄 摘要（原文）

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.