FedSubMuon: Communication-Efficient Federated LLM Fine-Tuning via Structured Subspace Muon
Organizations: King Abdullah University of Science and Technology (KAUST) · Provable Responsible AI and Data Analytics (PRADA) Lab · Imperial College London · Technische Universität Berlin (TU Berlin) ·
Abstract
Federated fine-tuning adapts large language models (LLMs) to decentralized client data, but its scalability in cross-device training is often limited by the high communication cost. Muon is an optimizer that improves optimization performance by orthogonalizing momentum for matrix-valued parameters. Existing federated Muon methods demonstrate the benefit of matrix-aware optimization in federated learning, but still require transmitting full layer-size updates and optimizer state. A natural way to reduce communication is to directly apply Muon to LoRA factors, but this changes the optimized object and weakens Muon's matrix-aware update geometry. We propose FedSubMuon, a communication-efficient federated Muon fine-tuning method that optimizes compact coefficient matrices within shared structured subspaces. This design keeps Muon on a single matrix-valued trainable object, while reducing the client upload to compact coefficient matrices. We further introduce FedSubMuon-GT, an accuracy-oriented extension that uses projected gradients to adapt tracked subspace bases toward task-relevant gradient directions. Experiments on instruction tuning and mathematical reasoning show that FedSubMuon-GT achieves the best overall accuracy on four of five dataset-model pairs, while FedSubMuon performs best under all matched communication budgets. On Dolly-15K, the closest communication baseline requires 5.5 times and 1.4 times more total communication on Llama-1B and Qwen-4B, respectively.
Figures & tables
| Method | AdamW | Muon | SGD |
|---|---|---|---|
| FedIT | |||
| FLoRA |
| Method | Natural Instructions | Dolly-15K | GSM8K | MATH | |
|---|---|---|---|---|---|
| Llama 1B | Llama 1B | Qwen 4B | Qwen 8B | Qwen 8B | |
| FedIT | 30.14 0.82 | 28.32 0.41 | 36.33 0.17 | 79.59 0.33 | 45.24 1.22 |
| FeDeRA | 31.07 0.43 | 29.63 0.53 | 36.78 0.37 | 80.29 0.26 | 47.52 0.65 |
| FLoRA | 31.12 0.28 | 29.30 0.45 | 36.39 0.36 | 76.58 0.82 | 43.48 0.42 |
| FedEx-LoRA | 30.29 0.51 | 28.83 0.10 | 36.57 0.76 | 79.51 0.25 | 45.82 1.26 |
| FLoRG | 29.83 0.42 | 28.53 0.54 | 36.79 0.70 | 80.12 0.51 | 43.48 0.42 |
| Method | Model / Dataset | ||
|---|---|---|---|
| Llama/NI | Llama/Dolly | Qwen/Dolly | |
| FedSubMuon | 30.59 | 30.09 | 36.87 |
| FedSubMuon-GT | 29.19 | 29.32 | 36.83 |
| FedIT | 30.12 | 24.70 | 34.62 |
| FeDeRA | 30.10 | 26.68 | 36.49 |
| FLoRA | 10.28 | 16.95 | 28.77 |
| Model | Metric | FedSubMuon | FedSubMuon-GT | FedIT | FeDeRA | FLoRA | FedEx-LoRA | FLoRG | FedKRSO |
|---|---|---|---|---|---|---|---|---|---|
| Llama 1B | Uplink (MB) | 105 (1 ) | 923 (8.8 ) | 920 (8.8 ) | 1159 (11 ) | 2045 (19 ) | 920 (8.8 ) | 721 (6.9 ) | 1311 (12 ) |
| Total Comm. (MB) | 264 (1 ) | 11848 (45 ) | 1840 (7.0 ) | 2317 (8.8 ) | 22492 (85 ) | 11042 (42 ) | 1442 (5.5 ) | 85679 (325 ) | |
| GPU Mem. (MB) | 5530 | 5555 | 5761 | 5762 | 5764 | 5762 | 6567 | 24136 | |
| Wall time (s) | 77.19 | 75.33 | 83.47 | 81.25 | 82.65 | 84.61 | 114.18 | 77.61 | |
| Qwen 4B | Uplink (MB) | 1410 (1 ) | 7078 (5.0 ) | 4720 (3.3 ) | 5782 (4.1 ) | 7078 (5.0 ) | 6960 (4.9 ) | 2490 (1.8 ) | 13920 (9.9 ) |
| Total Comm. (MB) | 3457 (1 ) | 83802 (24 ) | 9440 (2.7 ) | 11564 (3.3 ) | 77880 (23 ) | 83519 (24 ) | 4980 (1.4 ) | 1127508 (326 ) |
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
| Model | Method group | Rank | Trainable params (ratio) |
|---|---|---|---|
| Llama 1B | FedSubMuon/ FedSubMuon-GT | ( ) | |
| FedIT/FeDeRA/FLoRA/ FedEx-LoRA | ( ) | ||
| FLoRG | ( ) | ||
| FedKRSO | ( ) | ||
| Qwen 4B | FedSubMuon/ FedSubMuon-GT | ( ) | |
| FedIT/FeDeRA/FLoRA/ FedEx-LoRA | ( ) |