[Paper Review] FedLab: A Flexible Federated Learning Framework

Dun Zeng, Siqi Liang+3 more|arXiv (Cornell University)|Jul 24, 2021

Privacy-Preserving Technologies in DataComputer Science35 references64 citations

TL;DR

FedLab is a lightweight, open-source framework for flexible federated learning simulation that emphasizes modularity, communication efficiency, and algorithmic flexibility across deployment scenarios.

ABSTRACT

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 exttt{FedLab}, a lightweight open-source framework for FL simulation. The design of exttt{FedLab} focuses on FL algorithm effectiveness and communication efficiency. Also, exttt{FedLab} is scalable in different deployment scenario. We hope exttt{FedLab} could provide flexible API as well as reliable baseline implementations, and relieve the burden of implementing novel approaches for researchers in FL community.

Motivation & Objective

• Introduce FedLab as a flexible, lightweight FL framework for simulation and research.
• Provide highly customizable interfaces to explore FL algorithms and communication strategies.
• Offer data partition tools and standard FL pipelines to reduce duplication of effort in research.

Proposed method

• Define two main roles (Server and Client) with modular components (NetworkManager and Trainer/ParameterServerHandler).
• Implement flexible communication backends using torch.distributed with synchronous and asynchronous patterns.
• Provide data partition utilities and NIID/baseline datasets to simulate non-IID settings.
• Support multiple deployment scenarios (Standalone, Cross-process, Hierarchical) to scale experiments.
• Offer standard FL implementations (e.g., FedAvg, FedAsgd) and example code snippets for rapid prototyping.
• Maintain an open-source workflow with CI and documentation for community collaboration.

Experimental results

Research questions

• RQ1 How can a flexible framework accelerate FL research by decoupling communication from optimization components?
• RQ2 What deployment scenarios and data partition strategies best support experimentation with heterogeneous, non-IID FL data?
• RQ3 Can standard FL baselines and modular APIs reduce the burden of implementing new FL algorithms for researchers?

Key findings

• FedLab provides a modular architecture with NetworkManager and Trainer/ParameterServerHandler to separate communication from optimization.
• The framework supports synchronous and asynchronous communication patterns and scalable scheduling for many clients.
• FedLab includes data partition utilities and NIID benchmarks to simulate diverse data distributions across clients.
• Three deployment scenarios (Standalone, Cross-process, Hierarchical) enable flexible experimentation across resource and network configurations.
• FedLab ships with standard FL baselines (e.g., FedAvg, FedAsgd) and example code to facilitate rapid prototyping.