Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Optimizing Federated Learning in the Era of LLMs: Message Quantization and Streaming

Ziyue Xu, Zhihong Zhang, Holger R. Roth, Chester Chen, Yan Cheng, Andrew Feng

TL;DR

FL faces significant communication and memory challenges when training large language models. This work extends NVIDIA FLARE with two core techniques—two-way message quantization and streaming-based object transmission (container and file streaming)—to alleviate bandwidth and memory bottlenecks in federated LLM scenarios. The approach preserves training and aggregation precision while substantially reducing message sizes (up to ~86% with 4-bit quantization) and peak memory usage in experiments with Llama-3.2-1B, demonstrating practicality for real-world FL deployments and outlining directions for broader multi-client evaluations and adaptive streaming.

Abstract

Federated Learning (FL) offers a promising solution for training machine learning models across distributed data sources while preserving data privacy. However, FL faces critical challenges related to communication overhead and local resource constraints, especially in the era of Large Language Models (LLMs) with billions of parameters. The sheer size of these models exacerbates both memory and communication constraints, making efficient transmission and processing essential for practical deployment. NVIDIA FLARE, an open-source SDK for federated learning, addresses these challenges by introducing advanced communication capabilities. Building upon existing solutions for large object streaming, we enhance FL workflows for LLMs through two key techniques: message quantization and container/file streaming. Quantization reduces message size, while streaming enables efficient memory management, improving scalability and integration with existing workflows. These advancements significantly enhance the robustness and efficiency of FL with LLMs, ensuring better performance in real-world federated learning scenarios.

Optimizing Federated Learning in the Era of LLMs: Message Quantization and Streaming

TL;DR

FL faces significant communication and memory challenges when training large language models. This work extends NVIDIA FLARE with two core techniques—two-way message quantization and streaming-based object transmission (container and file streaming)—to alleviate bandwidth and memory bottlenecks in federated LLM scenarios. The approach preserves training and aggregation precision while substantially reducing message sizes (up to ~86% with 4-bit quantization) and peak memory usage in experiments with Llama-3.2-1B, demonstrating practicality for real-world FL deployments and outlining directions for broader multi-client evaluations and adaptive streaming.

Abstract

Federated Learning (FL) offers a promising solution for training machine learning models across distributed data sources while preserving data privacy. However, FL faces critical challenges related to communication overhead and local resource constraints, especially in the era of Large Language Models (LLMs) with billions of parameters. The sheer size of these models exacerbates both memory and communication constraints, making efficient transmission and processing essential for practical deployment. NVIDIA FLARE, an open-source SDK for federated learning, addresses these challenges by introducing advanced communication capabilities. Building upon existing solutions for large object streaming, we enhance FL workflows for LLMs through two key techniques: message quantization and container/file streaming. Quantization reduces message size, while streaming enables efficient memory management, improving scalability and integration with existing workflows. These advancements significantly enhance the robustness and efficiency of FL with LLMs, ensuring better performance in real-world federated learning scenarios.

Paper Structure

This paper contains 11 sections, 5 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Message Quantization: Reducing Communication Overhead
  3. General FL Workflow of NVFlare
  4. Filter Mechanism
  5. Two-way Quantization Workflow
  6. Quantization Methods
  7. Streaming Functionality: Reducing Local Memory Usage
  8. Experiments and Results
  9. Quantization Results
  10. Memory-efficient Streaming
  11. Limitations and Future Work

Figures (5)

  • Figure 1: NVFlare's Data streaming API.
  • Figure 2: Federated Model Training with Message Quantization
  • Figure 3: Illustration of streaming under different settings, green region indicates the pre-allocated memory needed for communication.
  • Figure 4: Federated SFT comparison: centralized v.s. single-site FL.
  • Figure 5: Federated SFT comparison: single-site FL with message quantization.