Fed-pilot: Optimizing LoRA Allocation for Efficient Federated Fine-Tuning with Heterogeneous Clients

📄 arXiv: 2410.10200v2 📥 PDF

作者: Zikai Zhang, Rui Hu, Ping Liu, Jiahao Xu

分类: cs.LG, cs.DC

发布日期: 2024-10-14 (更新: 2025-06-20)

💡 一句话要点

Fed-pilot：优化LoRA分配，实现异构客户端高效联邦微调

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

关键词: 联邦学习 低秩适应 LoRA 异构客户端 内存优化

📋 核心要点

  1. 现有联邦学习微调方法难以应对客户端内存异构性，限制了基础模型在资源受限设备上的应用。
  2. Fed-pilot通过优化LoRA模块的选择，在满足客户端内存约束的同时，最大化模型性能。
  3. 实验结果表明，Fed-pilot在多种数据集和异构数据设置下，显著优于现有联邦微调方法。

📝 摘要（中文）

联邦学习允许在分布式客户端上对基础模型（FMs）进行微调以适应特定任务；然而，其可扩展性受到客户端内存容量异构性的限制。本文提出了Fed-pilot，一个内存高效的联邦微调框架。它通过仅在本地训练LoRA模块的子集，使内存受限的客户端能够参与基于低秩适应（LoRA）的微调。Fed-pilot将可训练LoRA模块的最佳选择确定为一个背包优化问题，从而在每个客户端的内存约束下最大化模型性能。为了减轻由异构模块分配和非独立同分布（Non-IID）数据引起的不一致性，Fed-pilot采用了一种新颖的聚合规则，动态补偿欠更新的层。在各种异构数据设置下的五个不同数据集上的大量实验表明，与最先进的方法相比，Fed-pilot具有有效性和效率。据我们所知，这是第一个整合内存约束优化的FM联邦微调研究。代码将公开。

🔬 方法详解

问题定义：论文旨在解决联邦学习中，由于客户端内存容量的异构性，导致无法充分利用所有客户端进行基础模型微调的问题。现有方法要么需要所有客户端都具备足够的内存来容纳整个模型，要么只能进行部分参数的微调，导致模型性能下降。

核心思路：核心思路是将LoRA模块的选择问题转化为一个背包优化问题。每个客户端根据自身的内存限制，选择一部分LoRA模块进行训练，从而保证所有客户端都能参与到联邦学习过程中。同时，为了解决由于不同客户端训练的LoRA模块不同导致的模型不一致性问题，提出了动态补偿的聚合规则。

技术框架：Fed-pilot框架主要包含以下几个阶段：1) 客户端LoRA模块选择：每个客户端根据自身内存限制，使用背包算法选择一部分LoRA模块进行训练。2) 本地训练：客户端使用本地数据训练所选择的LoRA模块。3) 参数聚合：服务器接收来自客户端的LoRA模块参数，并使用动态补偿的聚合规则进行聚合。4) 模型更新：服务器将聚合后的LoRA模块参数更新到全局模型中。

关键创新：主要创新点在于：1) 将LoRA模块选择问题建模为背包优化问题，实现了在内存约束下的最优模块选择。2) 提出了动态补偿的聚合规则，有效缓解了由于异构模块分配和Non-IID数据导致的模型不一致性问题。

关键设计：1) 背包优化：使用动态规划算法求解背包问题，目标是在内存约束下最大化模型性能。性能的评估可以通过预先评估每个LoRA模块对模型性能的贡献来实现。2) 动态补偿聚合：对于每个LoRA模块，服务器根据参与训练该模块的客户端数量，动态调整聚合权重。参与训练的客户端越多，权重越高，反之亦然。这种方法可以有效地补偿欠更新的层，提高模型性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Fed-pilot在五个不同的数据集上，相比于现有的联邦微调方法，取得了显著的性能提升。例如，在某些数据集上，Fed-pilot的性能提升超过了5%。此外，Fed-pilot还能够有效地降低客户端的内存占用，使得更多的客户端能够参与到联邦学习过程中，提高了联邦学习的效率。

🎯 应用场景

Fed-pilot可应用于各种需要联邦学习的场景，尤其是在客户端设备资源受限的情况下，例如移动设备、物联网设备等。它可以帮助企业或组织利用大量分散在不同设备上的数据，训练出高性能的基础模型，从而提升各种应用的用户体验，例如个性化推荐、智能助手等。该研究对于推动联邦学习在资源受限环境下的应用具有重要意义。

📄 摘要（原文）

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.