解决异构设备内存不足时高效微调大模型的问题
Fed-pilot: Optimizing LoRA Allocation for Efficient Federated Fine-Tuning with Heterogeneous Clients
- 按背包问题优化选择训练哪些低秩模块，节省内存
- 在五组数据上表现优于现有方法，适配不同设备
- 首次将内存约束优化用于联邦大模型微调
联邦学习使分布式客户端能够针对特定任务微调基础模型（FMs），但其可扩展性受限于客户端内存容量的异质性。本文提出 Fed-pilot，一个面向内存受限客户端的高效联邦微调框架。通过仅在本地训练部分低秩适配（LoRA）模块，使资源受限设备也能参与微调。将最优可训练 LoRA 模块选择建模为背包优化问题，在内存约束下最大化模型性能。为缓解异构模块分配与非独立同分布（Non-IID）数据带来的不一致性，引入新型聚合规则，动态补偿更新不足层。在五个不同数据集、多种异构数据设置下的实验表明，相比当前最佳方法，Fed-pilot 在有效性与效率上均具优势。据我们所知，这是首个将内存约束优化集成到联邦大模型微调中的研究。代码将公开。
原文摘要 · Abstract (English)
Federated Learning enables the fine-tuning of foundation models (FMs) across distributed clients for specific tasks; however, its scalability is limited by the heterogeneity of client memory capacities. In this work, we propose Fed-pilot, a memory-efficient federated fine-tuning framework. It enables memory-constrained clients to participate in Low-Rank Adaptation (LoRA)-based fine-tuning by training only a subset of LoRA modules locally. Fed-pilot identifies the optimal selection of trainable LoRA modules as a knapsack optimization problem, maximizing model performance under memory constraints for each client. To mitigate inconsistencies arising from heterogeneous module allocations and Non-IID data, Fed-pilot employs a novel aggregation rule that dynamically compensates for under-updated layers. Extensive experiments on five diverse datasets across various heterogeneous data settings demonstrate Fed-pilot's effectiveness and efficiency compared to state-of-the-art methods. To the best of our knowledge, this is the first study on federated fine-tuning of FMs that integrates memory-constrained optimization. The code will be publicly available.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。