Table of Contents
Fetching ...

ParaBlock: Communication-Computation Parallel Block Coordinate Federated Learning for Large Language Models

Yujia Wang, Yuanpu Cao, Jinghui Chen

TL;DR

ParaBlock addresses the communication latency bottleneck in federated block-coordinate descent for large-language-model fine-tuning by employing a two-thread design that overlaps computation and communication on each client. It proves a non-convex convergence rate of $O(1/\sqrt{T})$ identical to standard FL-BCD while introducing a one-round staleness that is mitigated by a block-level correction mechanism. Empirically, ParaBlock achieves substantial wall-clock time reductions and maintains competitive performance across instruction-following and mathematical reasoning tasks on Llama-based models, outperforming many LoRA-based and cyclic baselines. This work demonstrates a scalable, latency-tolerant approach to FL-based LLM fine-tuning with resource-constrained clients, enabling faster deployment and broader participation.

Abstract

Federated learning (FL) has been extensively studied as a privacy-preserving training paradigm. Recently, federated block coordinate descent scheme has become a popular option in training large-scale models, as it allows clients to train only a subset of the model locally instead of the entire model. However, in the era of large language models (LLMs), even a single block can contain a significant number of parameters, posing substantial communication latency, particularly for resource-constrained clients. To address this challenge in federated training/fine-tuning LLMs, we propose ParaBlock, a novel approach that establishes two parallel threads for communication and computation to enhance communication efficiency. We theoretically prove that the proposed ParaBlock achieves the same convergence rate as the standard federated block coordinate descent methods. Empirical evaluations on fine-tuning LLMs on general instruction following and mathematical reasoning confirm that ParaBlock not only maintains strong performance but also significantly improves communication efficiency.

ParaBlock: Communication-Computation Parallel Block Coordinate Federated Learning for Large Language Models

TL;DR

ParaBlock addresses the communication latency bottleneck in federated block-coordinate descent for large-language-model fine-tuning by employing a two-thread design that overlaps computation and communication on each client. It proves a non-convex convergence rate of identical to standard FL-BCD while introducing a one-round staleness that is mitigated by a block-level correction mechanism. Empirically, ParaBlock achieves substantial wall-clock time reductions and maintains competitive performance across instruction-following and mathematical reasoning tasks on Llama-based models, outperforming many LoRA-based and cyclic baselines. This work demonstrates a scalable, latency-tolerant approach to FL-based LLM fine-tuning with resource-constrained clients, enabling faster deployment and broader participation.

Abstract

Federated learning (FL) has been extensively studied as a privacy-preserving training paradigm. Recently, federated block coordinate descent scheme has become a popular option in training large-scale models, as it allows clients to train only a subset of the model locally instead of the entire model. However, in the era of large language models (LLMs), even a single block can contain a significant number of parameters, posing substantial communication latency, particularly for resource-constrained clients. To address this challenge in federated training/fine-tuning LLMs, we propose ParaBlock, a novel approach that establishes two parallel threads for communication and computation to enhance communication efficiency. We theoretically prove that the proposed ParaBlock achieves the same convergence rate as the standard federated block coordinate descent methods. Empirical evaluations on fine-tuning LLMs on general instruction following and mathematical reasoning confirm that ParaBlock not only maintains strong performance but also significantly improves communication efficiency.

Paper Structure

This paper contains 32 sections, 6 theorems, 43 equations, 7 figures, 10 tables, 2 algorithms.

Key Result

Theorem 5.3

Under Assumptions as:smooth--as:bounded-v, let $T$ represent the total number of global rounds, $K$ be the number of local SGD training steps and $N$ be the number of the clients. If the learning rate $\eta$ and $\eta_l$ satisfy $\eta_l \leq \frac{1}{22KL}$ and $\eta \eta_l \leq \frac{1}{4KL}$ , the where $\mathcal{F} = f(\bm{\theta}_1) - f_*$ and $f_* =\min_{\bm{\theta}} f(\bm{\theta}) > -\infty$

Figures (7)

  • Figure 1: Comparison between the original federated BCD and the proposed ParaBlock. The original BCD's single-thread approach leads to higher runtime due to communication latency, while ParaBlock improves efficiency by overlapping communication and computation.
  • Figure 2: Time efficiency: wall-clock runtime for various network communication bandwidths and effective batch sizes.
  • Figure : LocalBlockTraining
  • Figure : Ablation for the block assignment the number of layers.
  • Figure : Ablation for block scheduling
  • ...and 2 more figures

Theorems & Definitions (11)

  • Theorem 5.3
  • Corollary 5.4
  • Remark 5.5
  • Lemma C.1
  • proof
  • Lemma C.2
  • proof
  • Lemma C.3
  • proof
  • Lemma C.4
  • ...and 1 more