Optimal Hydra: Model Heterogeneous Federated Learning for LLMs using Data-Informed Hydra-Nets
Organizations: Concordia University · Mila · Toyota Motor Europe
Abstract
Large language models (LLMs) trained on vast amounts of data have demonstrated impressive performance over a range of tasks, but future performance gains will require increasing amounts of data and compute. Federated Learning (FL) facilitates training on decentralized private data, but client resource limitations can make training impractical, particularly as model sizes continue to expand. These resource constraints make the assumption of a shared server-client architecture impractical. Model heterogeneous FL offers a promising alternative by adapting model size to client capabilities; however, previous work is limited to the pre-training setting, and conclusions on sub-network selection do not necessarily hold for fine-tuning. Furthermore, prior approaches create client sub-networks without considering their unique data distributions, thereby overlooking the opportunity to tailor sub-network architecture to local data. Motivated by evidence that attention heads exhibit task dependency within layer defined feature spaces, we propose Optimal Hydra, a transformer based heterogeneous FL framework for fine-tuning that combines global layer selection with client specific attention head selection, aligned with local data distributions. Optimal Hydra makes explicit use of client data to produce client-specific sub-networks we call Hydra-nets. To our knowledge, this is the first model heterogeneous FL approach tailored to LLM fine-tuning. Experiments on three datasets show significant improvements over baseline methods while also reducing local resource demands.
Figures & tables
| Config | Comp. | % Impr. | Mem. | % Impr. |
|---|---|---|---|---|
| (GFLOPS) | w.r.t. full | (GB) | w.r.t. full | |
| f(20) | 56.6 | – | 61.9 | – |
| 3q(20) | 50.6 | 10.7 | 55.8 | 9.7 |
| f(5)3q(5)h(5)q(5) | 47.6 | 16.0 | 52.8 | 14.6 |
| h(20) | 44.5 | 21.3 | 49.8 | 19.5 |
| 3q(6)h(7)q(7) | 44.2 | 21.9 | 49.5 | 20.0 |
| Alpaca | Natural Instructions | Dolly | |||||||
|---|---|---|---|---|---|---|---|---|---|
| Hydra | HeteroFL | FedRolex | Hydra | HeteroFL | FedRolex | Hydra | HeteroFL | FedRolex | |
| F(20) ZS | |||||||||
| F(20) FedAvg | |||||||||
| 3Q(20) | |||||||||
| H(20) | |||||||||
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
| Name | Var. | DataJuicer-1.3B |
|---|---|---|
| No. trf blocks 1 | 24 | |
| Hidden dim. | 2048 | |
| Tot. No. heads | 16 | |
| Reduced No. heads | variable | |
| Head dim. | 128 | |
| FF dim. | 5504 |
| Dataset | Config | (OH) | (OH) | (HFL) | (Rolex) |
|---|---|---|---|---|---|
| Alpaca | 3Q(20) | ||||
| H(20) | |||||
| Q(20) | |||||
| F(5)3Q(5)H(5)Q(5) | 25 | ||||
| H(10)Q(10) | |||||
| 3Q(6)H(7)Q(7) |
Federated Learning
Federated learning trains shared models across data holders without requiring them to pool their raw data.
See all 951 papersFedTreeLoRA: Reconciling Statistical and Functional Heterogeneity in Federated LoRA Fine-Tuning
trunks' while progressively specializing on deep branches'. Experiments on NLU and NLG benchmarks demonstrate that FedTreeLoRA significantly outperforms state-of-the-art methods by effectively reconciling generalization and personalization.