Source-linked AI summary
FedLab: A Flexible Federated Learning Framework
Dun Zeng, Siqi Liang, Xiangjing Hu, Hui Wang, Zenglin Xu
TL;DR
Federated learning seeks collaborative model training without direct data sharing, but research remains burdened by heterogeneous data and communication costs. FedLab addresses these needs as a lightweight, customizable framework for scalable simulation, with standard modules and deployment paradigms.
Problem
FL research must address heterogeneous client data and substantial communication costs while supporting privacy-preserving collaborative learning.
Method
FedLab provides a lightweight, open-source, modular FL simulation framework with customizable interfaces, standard implementations, data-partition tools, and scalable deployment support.
Results
FedLab offers common FL communication patterns, optimization modules, high-level APIs, open interfaces, and standalone, cross-machine, and hierarchical simulation support.
Takeaways & Limitations
FedLab is intended to relieve repetitive implementation work and support research on FL optimization effectiveness and communication efficiency.
Abstract
from arXiv · showhide
Federated learning (FL) is a machine learning field in which researchers try to facilitate model learning process among multiparty without violating privacy protection regulations. Considerable effort has been invested in FL optimization and communication related researches. In this work, we introduce \texttt{FedLab}, a lightweight open-source framework for FL simulation. The design of \texttt{FedLab} focuses on FL algorithm effectiveness and communication efficiency. Also, \texttt{FedLab} is scalable in different deployment scenario. We hope \texttt{FedLab} could provide flexible API as well as reliable baseline implementations, and relieve the burden of implementing novel approaches for researchers in FL community.
1 Introduction
FedLab is introduced to reduce the repetitive implementation burden in federated-learning research. It offers customizable, scalable, and standardized simulation components, data-partition tools, benchmarks, and open-source maintenance.
- Researchers often implement FL algorithms from scratch despite existing frameworks, hindering reproduction and rapid verification.
- FedLab provides customizable modules for communication, compression, optimization, data partitioning, and other FL simulation functions.Users can combine custom modules with defaults to construct simulation environments.
- FedLab supports standalone, cross-machine, and hierarchical simulation paradigms through customizable interfaces and scalable FL-system design.
- FedLab supplies data-partition functions and built-in schemes for comprehensive data-distribution scenarios across federations.
- FedLab presents standardized FL implementations, including synchronous and asynchronous systems, datasets, benchmarks, and functional simulation modules.
- An open-source GitHub group and documentation support FedLab’s continuous maintenance and public development.
2 Background
Federated learning research addresses client heterogeneity and communication costs while developing optimization and compression strategies. FedLab targets these research needs with a lightweight, modular framework for scalable simulation.
- Communication and Compression: Communication cost during model uploading and downloading is a distributed-learning bottleneck, especially in cross-device settings.
- Optimization Algorithms: Federated learning faces client-data heterogeneity that can reduce robustness and slow model convergence.
- Optimization Algorithms: Optimization research commonly changes client local training or server parameter aggregation to address non-I.I.D. data and personalization.
- Communication and Compression: Quantization lowers tensor precision, while sparsification transmits selected tensors and can achieve at least 100× compression.
- Communication and Compression: Because compression methods are model independent, a flexible FL framework should provide model-independent compression modules.
- Related Work: Existing frameworks variously emphasize industrial deployment, secure computation, single-machine simulation, broad research coverage, or communication, while some are heavy.
- Related Work: FedLab is designed as a lightweight simulator focused on optimization effectiveness and communication efficiency, with standard pipelines and custom interfaces.
3 Framework Overview
FedLab separates communication and optimization into customizable modules, supporting flexible communication patterns, scalable networking, standard optimization, data partitioning, and multiple deployment scenarios.
- Architecture: FedLab organizes each server and client around NetworkManager plus ParameterServerHandler or Trainer components.NetworkManager handles message processing, while handlers and trainers support optimization design.
- Communication: NetworkManager supports customizable communication agreements, tensor-based messages, compression interfaces, and scalable network topologies.Package and PackageProcessor hide tensor-packing details from users, while Scheduler extends scalability for larger client populations.
- Communication: FedLab implements synchronous and asynchronous communication patterns, covering initialization, client sampling, local training, uploads, and server aggregation.Asynchronous training updates the global model whenever the server receives a synchronization upload.
- Optimization: The optimization module separates local training from server aggregation and includes standard FedAvg and FedAsgd implementations.Trainer manages local datasets and PyTorch training, while ParameterServerHandler aggregates parameters.
- Deployment and data: FedLab provides data-partition functions and supports standalone, cross-process, and hierarchical deployment scenarios.Standalone simulation uses serial execution, cross-process simulation requires a shared network, and hierarchical deployment uses customizable Scheduler groups.
4 Pipeline and Examples
FedLab combines standard PyTorch model training with high-level communication and server-management APIs, enabling concise client and server implementations.
- Client pipeline: The client example uses standard PyTorch model, optimizer, loss, and dataset setup before invoking FedLab training and networking components.ClientSGDTrainer handles training, while DistNetwork and ClientPassiveManager provide communication and execution.
- Client pipeline: FedLab provides high-level APIs for defining network topology and standard network training processes.The example connects a client to a distributed network through DistNetwork and runs it with ClientPassiveManager.
- Server pipeline: The server example defines a SyncParameterServerHandler, connects it through DistNetwork, and runs ServerSynchronousManager.The handler is configured with the FedAvg algorithm, while the network identifies the server process with rank 0.
5 Development
FedLab is maintained through a public GitHub group with issue tracking, continuous integration, and published documentation.
- Maintenance: The GitHub group supports ongoing development through bug reports, feature requests, usage questions, and community contributions.Continuous integration is used to support package robustness.
- Documentation: Comprehensive documentation is generated with Sphinx and published on fedlab.readthedocs.io.
6 Summary and Future Work
FedLab is presented as a flexible, lightweight framework with reusable communication and optimization modules, and its future work expands implementations, research releases, and functional coverage.
- Summary: FedLab provides common FL communication patterns and optimization modules through high-level APIs and open interfaces for standardized simulation.The framework also maintains an open-source GitHub group for contributions and issue handling.
- Future work: Future work includes releasing optimization and communication research implementations through FedLab.
- Future work: The authors plan to re-implement additional FL works in FedLab to provide more standardized implementations.
- Future work: Planned functional additions include more complex network topologies, network configuration scripts, and modules for other learning techniques.The listed techniques include unsupervised, semi-supervised, and transfer learning.