让专家模型在单卡上高效全参数微调
RevFFN: Memory-Efficient Full-Parameter Fine-Tuning of Mixture-of-Experts LLMs with Reversible Blocks
- 用可逆变换块重构中间激活值，省去存储开销
- 单卡即可完成专家模型全参数微调，峰值内存大幅降低
- 适合资源有限但需精细调优的场景
全参数微调是适配大语言模型至下游任务的关键技术，但因反向传播需缓存大量中间激活值，导致内存开销巨大，难以在实际中应用。现有分布式训练框架如DeepSpeed虽通过ZeRO和FSDP缓解此问题，但依赖多卡或CPU内存溢出，常需额外硬件且降低训练速度。本文提出RevFFN，一种面向混合专家（MoE）大模型的内存高效微调范式。其采用精心设计的可逆Transformer块，在反向传播时可从输出重建层输入激活值，无需存储多数中间激活。该方法在保持MoE架构表达能力的同时，显著降低全参数微调的峰值内存消耗，使单张消费级或服务器级显卡即可实现高效微调。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。