FLStore: Efficient Federated Learning Storage for Non-training Workloads

FLStore: Efficient Federated Learning Storage for Non-training Workloads

Speakers: Ahmad Faraz Khan, PhD Student, Virginia Tech; Samuel Fountain, PhD Student, Virginia Tech; Ahmed M. Abdelmoniem, Assistant Professor, University of Minnesota; Ali R. Butt, Professor, Virginia Tech; Ali Anwar, Assistant Professor, Queen Mary University of London

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

This article delves into FLStore, a novel architecture designed to enhance the efficiency of federated learning (FL) pipelines by optimizing the storage and processing of non-training workloads. Presented by Ahmad Faraz Khan and Samuel Fountain from Virginia Tech, this work addresses a critical bottleneck in federated learning systems: the high latency and cost associated with tasks beyond the core model training process. While FL has gained significant traction across diverse applications like healthcare, financial systems, and edge computing (e.g., Apple's voice recognition, Google's Gboard), existing cloud aggregators struggle to efficiently handle the increasing complexity and data requirements of tasks such as personalization, malicious client identification, and contribution calculation.

The core innovation of FLStore lies in its unification of the compute and data planes through a serverless memory cache, enabling data-local computation for non-training workloads. This approach fundamentally rethinks how data is accessed and processed in FL, moving away from the traditional model of data transfer between separate storage and compute services. By characterizing the unique access patterns of various non-training tasks and implementing specialized caching policies, FLStore achieves substantial reductions in both latency and operational costs, offering a more scalable and cost-effective solution for real-world federated learning deployments. The work highlights that non-training workloads, often overlooked, can account for a significant portion of the total cost and latency in an FL pipeline, making their optimization crucial for the broader adoption and efficiency of federated learning.

Background

▶ Watch: Introduction to FLStore and its purpose (0:00)

Federated learning (FL) represents a paradigm shift from traditional centralized machine learning. Instead of sending raw data from numerous client devices to a central server for training, FL distributes the model training process. A central aggregator sends a global model to multiple participating clients (e.g., mobile phones, IoT devices). Each client then trains the model locally on its private dataset, generating local model updates. These updates, rather than the raw data, are sent back to the aggregator, which then merges them to create an improved global model. This iterative process offers several advantages, including enhanced data privacy, reduced communication overhead compared to transmitting raw data, and the ability to leverage fresh data directly from edge devices. Prominent industry examples include Apple's use of FL for improving voice recognition and Google's application for Gboard word prediction, which reportedly yielded a 25% increase in accuracy.

However, the FL pipeline extends far beyond just model training. A myriad of non-training workloads are essential for robust, fair, and performant FL systems. These include, but are not limited to, model inference, aggregation (merging local updates), personalization (adapting global models to individual clients), model pruning, malicious client identification (detecting and mitigating poisoning attacks), accountable AI (ensuring transparency and fairness), contribution calculation (for incentivizing client participation), and dynamic client scheduling. These tasks can be executed either after the entire training process is complete or interleaved between training rounds, such as scheduling the next set of clients. The challenge arises because these non-training workloads are far from insignificant; the presenters highlight that they can contribute, on average, 89% of the cost and 38% of the latency to an FL pipeline.

The root cause of this inefficiency lies in the conventional cloud-based FL infrastructure. Typically, client data and model updates reside in a data plane (e.g., in-memory caches, object stores like Amazon S3). When a non-training workload needs to be executed, the necessary data is moved to a separate compute plane (e.g., AWS Sagemaker instances, virtual machines). This data movement incurs substantial communication latency and cost, especially when workloads require historical data for specific clients or multiple past rounds across all clients, rather than just the most recent model update. The existing frameworks often lack the ability to categorize the diverse access patterns of these iterative non-training workloads, leading to inefficient data retrieval and processing. The problem, therefore, is two-fold: the physical separation of data and compute, and the absence of intelligent, workload-aware data management strategies. FLStore aims to directly address these challenges by unifying data and compute and leveraging specialized caching.

Key Findings

▶ Watch: High cost and latency of traditional approaches (3:15)

FLStore introduces a paradigm shift in how federated learning handles non-training workloads, primarily by unifying the data and compute planes and implementing intelligent, workload-aware caching. The core findings and contributions are:

1. Unified Compute and Data Plane: FLStore’s most significant contribution is its architectural innovation that integrates data storage directly with computational resources. By leveraging serverless functions (e.g., AWS Lambda, OpenFaaS) as a serverless memory cache, FLStore enables data-local computation. This means that instead of moving data to compute, computation is brought to where the data resides, eliminating the substantial I/O costs and network latency associated with transferring data between separate storage (e.g., S3, Elastic Cache) and compute (e.g., Sagemaker) services. Each serverless function in the cache has attached CPU and GPU compute capabilities, allowing for efficient in-memory processing.

1. Categorization of Non-Training Workload Access Patterns: The research systematically categorizes the diverse data access patterns exhibited by various non-training workloads. Recognizing that FL is iterative and many non-training tasks have specific, recurring data needs (e.g., historical data for a specific client, multi-round data for all clients), this categorization forms the basis for developing optimized data management strategies. This is a crucial step that existing, workload-agnostic FL aggregators typically overlook.

1. Specialized Caching Policies: Building upon the workload categorization, FLStore develops and implements specialized caching policies. These policies are designed to intelligently pre-fetch or retain data based on the predicted future needs of specific non-training workloads, significantly improving cache hit rates. This contrasts sharply with traditional, generic caching policies like LRU (Least Recently Used) or FIFO (First-In, First-Out), which are less effective in the context of FL's iterative and varied non-training tasks.

1. Dramatic Latency and Cost Reductions: Through its architectural innovations and intelligent caching, FLStore demonstrates unprecedented efficiency gains.

• Compared to a baseline of AWS Sagemaker coupled with Amazon S3 (a cloud object store), FLStore reduces average latency by 71% and average costs by 92% per request.
• Against a faster but more expensive baseline of AWS Sagemaker with Elastic Cache (an in-memory caching service), FLStore still achieves average latency reductions of 65% and an astounding 99% reduction in costs.
• When comparing FLStore's specialized caching policies against traditional policies (LRU, FIFO), the specialized approach leads to a substantial drop in both latency and cost.

1. Significantly Increased Cache Hit Rate: The effectiveness of FLStore's specialized caching policies is quantitatively demonstrated by a 99% increase in the cache hit rate compared to using workload-agnostic policies. This high hit rate is a direct consequence of the system's ability to predict and prepare for data access patterns specific to non-training workloads.

These findings collectively position FLStore as a highly efficient and cost-effective solution for managing the often-overlooked but critical non-training aspects of federated learning pipelines, making advanced FL applications more feasible and economically viable.

Technical Deep Dive

▶ Watch: Introducing FLStore's architecture and components (4:20)

FLStore's architecture is meticulously designed to overcome the limitations of traditional federated learning infrastructure by integrating data storage and computation. It comprises three primary components: the Request Tracker, the Cache Engine, and the Serverless Memory Cache.

1. Request Tracker: This component acts as the central orchestrator for all non-training workloads. When a non-training task is initiated, the request is first directed to the Request Tracker. Its role is to understand the nature of the workload, identify the specific data requirements (e.g., which client's data, which rounds, what metadata), and then coordinate with the Cache Engine to locate or prepare the necessary data. It essentially serves as the "bookkeeper" for ongoing and pending non-training requests, ensuring they are routed efficiently.

1. Cache Engine: The Cache Engine is the "bookkeeper" for all data within FLStore. It maintains an up-to-date registry of where specific data segments are stored across the Serverless Memory Cache. During the FL training phase, when clients send their local model updates or other data, the Cache Engine intelligently decides where to place this data. It can allocate data to existing serverless functions that have available space or, leveraging the elasticity of serverless platforms, spin up new serverless functions as needed. This dynamic scaling ensures that data storage capacity can adapt to the fluctuating demands of the FL pipeline. Furthermore, for long-term persistence or disaster recovery, the Cache Engine can also push data to a traditional cloud object store (e.g., MinIO in the implementation). Conversely, if historical data not currently in the Serverless Memory Cache is required, the Cache Engine can pull it from the persistent object store into the serverless functions.

1. Serverless Memory Cache: This is the core innovation where data and compute converge. The Serverless Memory Cache is not a single entity but rather a group of serverless functions (e.g., Amazon Lambda, OpenFaaS, Whisk). Each serverless function in this group is provisioned with not only memory for data storage but also attached CPU and GPU compute resources. This is the crucial element that enables data-local computation. When the Request Tracker identifies a non-training workload and the Cache Engine locates the relevant data spread across one or more serverless functions, the request itself (the computation logic) is sent directly to these functions. Instead of moving potentially large datasets over the network to a separate compute instance, the computation is performed directly within the memory and on the compute resources of the serverless function holding the data. This eliminates the significant I/O overhead and network latency that plague traditional architectures.

Workflow Integration:

During the federated training process, client devices send their data (e.g., local model updates, metadata) to the Cache Engine. The Cache Engine, based on its internal logic and resource availability, places this data into one or more serverless functions within the Serverless Memory Cache, potentially creating replicas for fault tolerance or concurrent access. Data can also be offloaded to a persistent object store like MinIO for durability.

When a non-training workload needs to be executed, the request goes to the Request Tracker. The Request Tracker collaborates with the Cache Engine to pinpoint the exact locations of all necessary data within the Serverless Memory Cache. Rather than initiating data transfer, the Request Tracker dispatches the computational task (the non-training workload's logic) directly to the serverless functions that hold the required data. These functions then execute the workload using their integrated CPU/GPU resources, performing in-memory computation. This "compute-to-data" rather than "data-to-compute" paradigm is central to FLStore's efficiency.

Workload Categorization and Specialized Caching:

A critical technical aspect is the categorization of non-training workloads based on their data access patterns. The iterative nature of FL means that many tasks exhibit predictable patterns (e.g., accessing data from the last 'k' rounds, or data belonging to a specific cluster of clients). By analyzing these patterns, FLStore can develop specialized caching policies. Unlike generic policies such as LRU or FIFO, these specialized policies are "workload-aware." They proactively inform the Cache Engine about what data might be needed in the near future, allowing for pre-fetching or intelligent retention of data, significantly boosting the cache hit rate and further reducing latency. The implementation used OpenFaaS for the serverless function environment and MinIO to simulate the persistent object store, demonstrating a practical and extensible deployment.

Experimental Setup & Results

▶ Watch: Achieving efficiency with data-local compute (6:00)

To validate the efficacy of FLStore, the research team conducted a comprehensive experimental evaluation comparing its performance against state-of-the-art cloud solutions. The setup was designed to simulate realistic federated learning scenarios and measure key metrics: latency and cost.

Models and Workloads:

The evaluation incorporated a diverse set of deep learning models commonly used in cross-device federated learning, which are typically smaller in size but represent prevalent use cases. These included EfficientNet, ResNet, MobileNet, and even a Transformer model. While these models are relatively smaller, the presenters explicitly stated that FLStore was tested with models up to 1 billion parameters, specifically mentioning Llama 312, demonstrating its scalability to moderately large models.

For non-training workloads, the team identified and tested up to 10 distinct types. These included critical tasks such as:

• Personalization: Customizing models for individual client preferences.
• Clustering: Grouping clients with similar data characteristics.
• Debugging: Identifying issues in model updates or client behavior.
• Malicious Filtering: Detecting and isolating clients submitting harmful updates.
• Other tasks like contribution calculation and scheduling.

Implementation and Baselines:

FLStore's serverless memory cache was implemented using OpenFaaS, an open-source framework for building serverless functions, while MinIO was employed to simulate a persistent object store.

For comparison, two primary baselines were established, representing typical cloud-based FL aggregation setups:

1. AWS Sagemaker + S3: This configuration pairs AWS Sagemaker (a managed service for machine learning compute) with Amazon S3 (a highly durable and inexpensive cloud object storage service). This represents a common, cost-effective baseline, though S3 is known for higher latency due to its object storage nature.
2. AWS Sagemaker + Elastic Cache: This configuration uses AWS Sagemaker for compute but couples it with Elastic Cache (an in-memory caching service, e.g., Redis or Memcached). Elastic Cache offers significantly lower latency than S3 but comes at a higher cost. This baseline represents a performance-optimized, albeit more expensive, traditional cloud setup.

Evaluation Metrics:

The primary metrics of interest were:

• Latency: The time taken to complete a non-training workload request.
• Cost: The financial expenditure associated with executing the workload.

Both metrics were analyzed on a "per request" basis, acknowledging that real-world FL pipelines can involve thousands or millions of such requests.

Headline Results:

The experimental results demonstrated compelling improvements across all evaluated scenarios:

• FLStore vs. AWS Sagemaker + S3:
• Latency Reduction: FLStore achieved an average reduction in latency of 71%. This significant drop is attributed to avoiding data transfer between S3 and Sagemaker by performing computation directly within the serverless memory cache.
• Cost Reduction: The cost per request was reduced by an average of 92%. This massive saving stems from minimizing data egress charges, efficient resource utilization of serverless functions, and the elimination of dedicated, always-on compute instances often associated with Sagemaker.

• FLStore vs. AWS Sagemaker + Elastic Cache:
• Latency Reduction: Even against the faster Elastic Cache baseline, FLStore reduced latency by 65%. While Elastic Cache is fast, data still needs to be moved over the network to Sagemaker. FLStore's in-memory, data-local computation within the serverless functions still provides a decisive advantage.
• Cost Reduction: The cost reduction in this scenario was even more dramatic, reaching a staggering 99%. This highlights the high cost of managed in-memory caching services like Elastic Cache compared to the pay-per-execution model and optimized resource usage of FLStore's serverless architecture.

Impact of Specialized Caching Policies:

Beyond the architectural advantages, the study also quantified the benefits of FLStore's workload-aware caching:

• Latency and Cost: Comparing FLStore variants utilizing specialized caching policies against those employing traditional policies (LRU, FIFO) showed a "significant drop" in both latency and cost. This confirms the value of intelligently predicting and managing data access patterns.
• Cache Hit Rate: The specialized caching policies led to an impressive 99% increase in the cache hit rate. This means that data required for non-training workloads was almost always found in the fast, in-memory serverless cache, further minimizing trips to slower persistent storage.

These results unequivocally establish FLStore as a superior solution for handling non-training workloads in federated learning, offering substantial improvements in both performance and cost efficiency.

Practical Implications

▶ Watch: Transition to evaluation and experimental setup (6:50)

FLStore presents a transformative approach to managing the often-overlooked but critical non-training workloads in federated learning, with significant practical implications for various stakeholders.

For practitioners and infrastructure teams deploying federated learning systems, FLStore offers a path to dramatically reduce operational costs and improve the responsiveness of their FL pipelines. The reported 92% to 99% cost reductions are a game-changer, making advanced FL applications that require frequent non-training tasks (like continuous personalization or real-time malicious client detection) economically viable. Infrastructure teams can leverage serverless platforms, which are inherently scalable and managed, to build FL aggregators that automatically adapt to varying workload demands without constant manual provisioning or scaling. This shift can free up resources currently spent on managing complex distributed storage and compute clusters.

Model builders and deployers will find FLStore particularly beneficial for developing more sophisticated and robust FL models. The ability to efficiently run diverse non-training workloads – from detailed debugging and model analysis to advanced personalization and accountability checks – means that the entire FL lifecycle can be more thoroughly managed. For instance, quickly identifying and filtering malicious clients or calculating client contributions for incentivization becomes a low-latency, low-cost operation, enhancing the security and fairness of the FL ecosystem. The architecture simplifies the deployment of complex FL pipelines by unifying data and compute, reducing the need for intricate data movement logic between disparate services. Furthermore, FLStore provides an on-demand solution, allowing non-training workloads to be executed efficiently even after the main training phase has concluded, facilitating post-training analysis or ad-hoc queries.

However, FLStore also introduces certain tradeoffs and limitations. While the concept of specialized caching policies is powerful, their initial development requires a deep understanding of the specific data access patterns of each non-training workload. This might demand upfront analysis and engineering effort. The presenters acknowledge this, outlining future work to automate the development of these policies. Another limitation is the current scale of tested models; while FLStore was tested with models up to 1 billion parameters (e.g., Llama 312), modern large language models often exceed 10 billion or even 100 billion parameters. Scaling FLStore to efficiently handle such colossal models, especially their intermediate representations and updates, remains a future challenge. The reliance on serverless functions means that the performance and cost benefits are tied to the specific serverless platform's capabilities and pricing model, which can vary.

Despite these considerations, FLStore's core contribution – unifying data and compute via serverless functions and leveraging workload-aware caching – addresses fundamental inefficiencies in FL. It enables a future where federated learning pipelines are not only private and distributed but also highly performant and cost-effective across their entire operational spectrum, fostering the development and deployment of more advanced and responsible AI systems at the edge.

Key Takeaways

• Unification of Compute and Data Planes: FLStore fundamentally redesigns federated learning infrastructure by unifying the compute and data planes using a serverless memory cache composed of serverless functions with attached CPU/GPU compute, enabling data-local computation.
• Dramatic Efficiency Gains: The architecture achieves significant reductions in both latency (65-71% average) and cost (92-99% average) for non-training workloads compared to state-of-the-art cloud solutions like AWS Sagemaker paired with S3 or Elastic Cache.
• Workload-Aware Caching: FLStore categorizes diverse non-training workload access patterns and implements specialized caching policies that intelligently pre-fetch and retain data, leading to a 99% increase in cache hit rate over traditional, workload-agnostic policies (e.g., LRU, FIFO).
• Addressing Overlooked Bottlenecks: The work highlights that non-training workloads can account for a substantial portion (up to 89% of cost, 38% of latency) of an FL pipeline, making their optimization critical for overall efficiency and scalability.
• Practical Applicability: FLStore provides an on-demand, scalable, and cost-effective solution for executing complex non-training tasks such as personalization, malicious client identification, and debugging, making advanced FL deployments more viable for practitioners.
• Future Directions: Future work aims to automate the development of specialized caching policies and extend FLStore's capabilities to efficiently handle models exceeding 10 billion parameters, addressing the growing scale of modern AI.

About the Speaker(s)

The talk was presented by Ahmad Faraz Khan and Samuel Fountain, both PhD students from Virginia Tech, who are credited as co-authors and equal contributors to the FLStore project. Their research focuses on optimizing federated learning systems. The work itself was a collaborative effort, involving contributions from researchers at Virginia Tech, the University of Minnesota, and Queen Mary University of London, bringing together expertise from multiple institutions in the field of distributed systems and machine learning.