HeteroTune: Efficient Federated Learning for Large Heterogeneous Models

TL;DR

HeteroTune is a federated learning framework that enables efficient fine-tuning of large heterogeneous models in resource-constrained, privacy-sensitive environments, achieving significant communication and memory savings while maintaining high performance.

Contribution

It introduces DeMA, a novel architecture for aggregating heterogeneous models, and CMGA, a gradient alignment mechanism, advancing federated learning for large, diverse models.

Findings

Reduces communication overhead by 99.5% on LLaMA models

Cuts peak memory usage by approximately 50%

Improves model performance by 4.61%

Abstract

While large pre-trained models have achieved impressive performance across AI tasks, their deployment in privacy-sensitive and distributed environments remains challenging. Federated learning (FL) offers a viable solution by enabling decentralized fine-tuning without data sharing, but real-world applications face significant obstacles due to heterogeneous client resources in compute and memory. To address this, we propose HeteroTune, a novel federated fine-tuning paradigm for large, heterogeneous models operating under limited communication and computation budgets. The core of our method lies in a novel architecture, DeMA (Dense Mixture of Adapters), which enables flexible and efficient aggregation of heterogeneous models by preserving their full representational capacity while facilitating seamless cross-model knowledge fusion. We further introduce CMGA (Cross-Model Gradient Alignment), a lightweight yet effective mechanism that enhances training stability by harmonizing gradient directions across heterogeneous client models during aggregation, mitigating update conflicts and promoting more consistent convergence in federated settings. We provide both theoretical analysis and empirical evidence showing that HeteroTune achieves state-of-the-art performance and efficiency across diverse tasks and model architectures. For example, on LLaMA models, it reduces communication overhead by 99.5%, cuts peak memory usage by ~50%, and improves performance by 4.61%.