Optimal Hydra: Model Heterogeneous Federated Learning for LLMs using Data-Informed Hydra-Nets

Gwen Legate; Rahaf Aljundi; Eugene Belilovsky

cs.LGOct 14, 2026

Optimal Hydra: Model Heterogeneous Federated Learning for LLMs using Data-Informed Hydra-Nets

Authors: Gwen Legate, Rahaf Aljundi, Eugene Belilovsky

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

Figure 1: Each model is represented as a set of horizontal rectangles denoting transformer layers. The segregation of each layer represents attention heads present in the layer. Note: for simplicity, attention heads are the only architectural components depicted. Hydra-Net selection algorithm: Clients first collaboratively select a subset of common layers from which no attention heads may be removed; these layers are highlighted in orange to indicate the protected status of their attention heads during the individual head selection process. Next, clients individually select a subset of heads in each unimportant layer via a gradient based scoring mechanism. Heads that were not selected by clients in the individual selection process are grey to indicate their removal from the sub-network. The end result is a unique, data aligned Hydra-Net for each client that is subsequently used in the model heterogeneous FL framework, Optimal Hydra.

Figure 2: Each row consists of results for one dataset (1. Alpaca 2. Natural Instructions 3. Dolly 15K) and each column corresponds to one model heterogeneous FL configuration (1. F(5)3Q(5)H(5)Q(5) 2. 3Q(6)H(7)Q(7) 3. 3Q(2)H(2)Q(16) 4. H(10)Q(10) 5. H(5)Q(15) ). Each sub-figure contains the mean Rouge-L

\pm

std. dev. of four seeds for Optimal Hydra (blue), and our transformer adapted HeteroFL (orange) and FedRolex (green) baselines. The black dotted line indicates the FedAvg baseline whose values are reported in Tab. 2 . Optimal Hydra performs better than the baselines in all cases, it does particularly well with the very data heterogeneous NI dataset and the very resource constrained 3Q(2)H(2)Q(16) model heterogeneous configuration.

Config	Comp.	% Impr.	Mem.	% Impr.
Config	(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

Table 1: Memory and compute requirements for the different resource allocation configurations we study. We provide both the absolute value and a percentage improvement over full FedAvg fine-tuning. Configurations are ordered from most to least resource intensive.

	Alpaca			Natural Instructions			Dolly
	Hydra	HeteroFL	FedRolex	Hydra	HeteroFL	FedRolex	Hydra	HeteroFL	FedRolex
F(20) ZS	$17.157$			$5.698$			$26.694$
F(20) FedAvg	$31.668\ (\pm.002)$			$26.915\ (\pm.002)$			$32.465\ (\pm.001)$
3Q(20)	$\mathbf{31.267}$	$26.183$	$30.186$	$25.593$	$22.991$	$\mathbf{25.630}$	$34.295$	$29.535$	$\mathbf{34.714}$
	$\bm{(\pm.001)}$	$(\pm.002)$	$(\pm.001)$	$(\pm.005)$	$(\pm.005)$	$\bm{(\pm.004)}$	$(\pm.005)$	$(\pm.002)$	$\bm{(\pm.003)}$
H(20)	$\mathbf{30.621}$	$24.077$	$29.258$	$\mathbf{22.738}$	$17.821$	$22.157$	$\mathbf{35.489}$	$31.516$	$35.431$

Table 2: Rouge-L (%)

\pm

the std. dev. for four seeds, is reported for The full model zero-shot and FedAvg baselines as well as the setting in which all clients have homogeneous, but limited resource capabilities.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Name	Var.	DataJuicer-1.3B
No. trf blocks 1	$L$	24
Hidden dim.	$d$	2048
Tot. No. heads	$h$	16
Reduced No. heads	$h^{\prime}$	variable
Head dim.	$d_{h}$	128
FF dim.	$i$	5504

Appendix

Table 3: Provides a summary of variables and their values for the DataJuicer-1.3B model architectures

Dataset	Config	$\eta_{c}$ (OH)	$\lambda$ (OH)	$\eta_{c}$ (HFL)	$\eta_{c}$ (Rolex)
Alpaca	3Q(20)	$5\times 10^{3}$	$100$	$5\times 10^{4}$	$1\times 10^{3}$
	H(20)	$5\times 10^{3}$	$10$	$5\times 10^{4}$	$5\times 10^{4}$
	Q(20)	$1\times 10^{3}$	$100$	$5\times 10^{4}$	$5\times 10^{4}$
	F(5)3Q(5)H(5)Q(5)	$5\times 10^{3}$	25	$1\times 10^{3}$	$1\times 10^{3}$
	H(10)Q(10)	$1\times 10^{3}$	$25$	$1\times 10^{3}$	$1\times 10^{4}$
	3Q(6)H(7)Q(7)	$5\times 10^{3}$	$50$	$5\times 10^{4}$	$5\times 10^{3}$

Appendix

Table 4: Provides selected hyper-parameters for Optimal Hydra each dataset, heterogeneity combination evaluated in the main body of the paper. OH stands for Optimal Hydra, HFL and Rolex for our modified HeteroFL and FedRolex baselines, respectively.

Panel of Figure 3 — Figure 3: For Client 1 at round 20 and transformer layer 1, we see head 7 is unimportant to both $\nabla\mathbf{(W_{v})^{2}}$ and $(\nabla\mathbf{W_{q}})^{2}+(\nabla\mathbf{W_{k}})^{2}+(\nabla\mathbf{W_{v}})^{2}$ and important to $(\nabla\mathbf{W_{q}})^{2}$ and $(\nabla\mathbf{W_{k}})^{2}$ showing the dominance of $\nabla\mathbf{(W_{v})^{2}}$ in the scoring mechanism proposed in Eq. 11

Panel of Figure 4 — Figure 4: For Client 3 at round 5 and transformer layer 13, we see head 2 is important to both $\nabla\mathbf{(W_{v})^{2}}$ and $(\nabla\mathbf{W_{q}})^{2}+(\nabla\mathbf{W_{k}})^{2}+(\nabla\mathbf{W_{v}})^{2}$ and unimportant to $(\nabla\mathbf{W_{q}})^{2}$ and $(\nabla\mathbf{W_{k}})^{2}$ again that $\mathbf{(W_{v})^{2}}$ is dominant.

Federated Learning

Federated learning trains shared models across data holders without requiring them to pool their raw data.

See all 951 papers

SplitFT: An Adaptive Federated Split Learning System For LLMs Fine-Tuning
Apr 29, 2026Yimeng Shan, Zhaorui Zhang, Sheng Di +3Communication-Efficient Distributed Training LLM Fine-Tuning
Hong Kong Polytechnic University Hong Kong · Argonne National Labratory USA · University of California, Merced USA +1
FedTreeLoRA: Reconciling Statistical and Functional Heterogeneity in Federated LoRA Fine-Tuning
Feb 27, 2026Jieming Bian, Lei Wang, Letian Zhang +1Federated Fine-Tuning Federated Learning
University of Florida, Gainesville, FL 32611 · Middle Tennessee State University Murfreesboro, TN 37132
Federated Foundation Models Fine-Tuning with Heterogeneous Compressed Clients
Jul 31, 2026Shengkun Zhu, Jinshan Zeng, Zhihua Allen-Zhao +5Federated Learning Aggregation Federated Fine-Tuning
The Hong Kong Polytechnic University · OceanBase, Ant Group · School of Management, Xi’an Jiaotong University +3

Apr 29, 2026J/K move · Enter open · S save

SplitFT: An Adaptive Federated Split Learning System For LLMs Fine-Tuning

Yimeng Shan, Zhaorui Zhang, Sheng Di, Yu Liu+2

Hong Kong Polytechnic University Hong Kong · Argonne National Labratory USA · University of California, Merced USA+1

Communication-Efficient Distributed Training LLM Fine-Tuning Split Federated Learning Federated Fine-Tuning Federated Learning Heterogeneous Federated Learning

Open paper AskPDF arXiv

Federated Split Learning has been identified as an efficient approach to address the computational resource constraints of clients in classical federated learning, while guaranteeing data privacy for distributed model training across data owners. However, it faces some critical challenges when such a training strategy meets large language models (LLMs) for fine-tuning. Such challenges include setting the cutlayer adaptively across different clients to address the data and device heterogeneity issues, which affect the system performance significantly. In addition, efficiently reducing the communication overhead during the fine-tuning procedure is also another challenge. No work tries to address these challenges. To bridge this gap, we propose SplitTF, an adaptive federated split learning system for LLMs fine-tuning. SplitFT enables different clients to set different cut layers according to their computation resources and trained model performance. SplitFT also proposes to reduce the LoRA rank in cutlayer to reduce the communication overhead. In addition to simulating the heterogeneous data in real-world applications for our proposed split federated learning system, we propose a length-based Dirichlet approach to divide the training data into different clients. Extensive experimental results show that our proposed approach outperforms the state-of-the-art approach for fine-tuning time efficiency and model performance based on various popular benchmarks.

Abstract

Federated Learning

Explore similar work

SplitFT: An Adaptive Federated Split Learning System For LLMs Fine-Tuning

FedTreeLoRA: Reconciling Statistical and Functional Heterogeneity in Federated LoRA Fine-Tuning

Federated Foundation Models Fine-Tuning with Heterogeneous Compressed Clients