Source-linked AI summary
FedML: A Research Library and Benchmark for Federated Machine Learning
Chaoyang He, Songze Li, Jinhyun So, Xiao Zeng, Mi Zhang, Hongyi Wang, Xiaoyang Wang, Praneeth Vepakomma, Abhishek Singh, Hang Qiu, Xinghua Zhu, Jianzong Wang, Li Shen, Peilin Zhao, Yan Kang, Yang Liu, Ramesh Raskar, Qiang Yang, Murali Annavaram, Salman Avestimehr
TL;DR
Existing FL libraries do not adequately support diverse computing configurations and make fair algorithm comparison difficult because implementations, datasets, and benchmarks vary. FedML addresses this gap with an open toolkit combining flexible APIs, multiple computing paradigms, standardized algorithms and benchmarks, and reusable models and datasets. It is intended to support reproducible FL development and evaluation, while standard FL settings still assume no central authority verifies training data or hardware.
Problem
Existing FL libraries lack support for diverse computing configurations, while inconsistent implementations, datasets, and benchmarks hinder fair performance comparison.
Method
FedML provides an open end-to-end toolkit with flexible APIs, three computing paradigms, standardized algorithms and benchmarks, and separated model, dataset, and algorithm implementations.
Results
FedML provides reference benchmark results and supports distributed computing that is about 8 times faster than standalone simulation for large CNN training with 10 parallel workers.
Takeaways & Limitations
FedML offers researchers and engineers a toolkit for developing FL algorithms and fairly comparing them across diverse computing paradigms and configurations.
Takeaways & Limitations
Standard FL settings assume no single central authority owns or verifies training data or user hardware, creating exposure to adversarial attacks during decentralized training.
Abstract
from arXiv · showhide
Federated learning (FL) is a rapidly growing research field in machine learning. However, existing FL libraries cannot adequately support diverse algorithmic development; inconsistent dataset and model usage make fair algorithm comparison challenging. In this work, we introduce FedML, an open research library and benchmark to facilitate FL algorithm development and fair performance comparison. FedML supports three computing paradigms: on-device training for edge devices, distributed computing, and single-machine simulation. FedML also promotes diverse algorithmic research with flexible and generic API design and comprehensive reference baseline implementations (optimizer, models, and datasets). We hope FedML could provide an efficient and reproducible means for developing and evaluating FL algorithms that would benefit the FL research community. We maintain the source code, documents, and user community at https://fedml.ai.
1 Introduction
Federated learning faces challenges from heterogeneous data, system constraints, and trustworthiness, while existing libraries lack broad configuration support and standardized comparisons. FedML addresses these gaps with multiple computing paradigms, flexible programming interfaces, standardized implementations and benchmarks, and an open development model.
- Motivation: Federated learning must address statistical heterogeneity, system constraints, and trustworthiness across multiple disciplines.These challenges involve machine learning, wireless communication, mobile computing, distributed systems, and information security.
- Limitations of Existing Tools: Existing libraries inadequately support diverse FL computing paradigms and algorithmic procedures.Simulation-oriented libraries mainly support centralized topologies and are unsuitable for complex auxiliary information exchange or customized training procedures.
- Limitations of Existing Tools: Existing FL work lacks standardized implementations and benchmarks, making fair performance comparison difficult under non-I.I.D. data distributions.Different libraries, datasets, and non-I.I.D. partitions can produce varying accuracies even for the same model and dataset.
- FedML Contributions: FedML supports on-device training, distributed computing, and single-machine simulation for different deployment scenarios.The paradigms target edge devices such as smartphones and IoT devices, distributed systems, and standalone simulation.
- FedML Contributions: FedML uses a worker/client-oriented interface to support diverse network topologies, information exchange, and training procedures.The interface is designed for flexible algorithm development across varied FL configurations.
- FedML Contributions: FedML provides standardized algorithm implementations and benchmarks with defined metrics, diverse non-I.I.D. datasets, and verified baseline results.The implementations serve both as API learning resources and comparison baselines, while the benchmarks facilitate fair evaluation.
- FedML Contributions: FedML is fully open and intended to evolve through expanded computing paradigms, algorithms, benchmarks, models, datasets, and community contributions.The paper identifies applications including computer vision, natural language processing, finance, transportation, digital health, recommendation, robotics, and smart cities.
2 FedML Library: Architecture Design
FedML separates high-level APIs from low-level communication and training components, enabling reusable implementations and flexible distributed experimentation. Its architecture also supports real-world hardware testbeds for measuring practical system performance.
- Architecture: FedML consists of FedML-API and FedML-core, representing high-level and low-level APIs.The architecture organizes the library into distinct abstraction levels.
- FedML-core: FedML-core separates distributed communication from model training and uses MPI-based communication with topology management.TopologyManager supports multiple network topologies used by FL algorithms.
- FedML-API: FedML-API builds on FedML-core and uses a client-oriented interface for implementing distributed algorithms.This design supports large DNN training scenarios that standalone simulation cannot handle because of GPU memory and training-time constraints.
- FedML-API: Separating models, datasets, and algorithms enables code reuse and fair comparison while reducing implementation differences across algorithms.Applications can contribute models and realistic datasets without understanding distributed optimization details.
- Hardware Testbeds: FedML-Mobile and FedML-IoT provide hardware-based testbeds for Android smartphones, Raspberry Pi 4, and NVIDIA Jetson Nano.These platforms allow evaluation of training time, communication, and computation costs under realistic system conditions.
3 FedML Library: Programming Interface
FedML uses flexible worker- and client-oriented interfaces to support customized message flows, topologies, training procedures, and distributed execution. Its APIs also address real-world FL concerns including hardware evaluation, privacy, security, and robustness.
- Worker/client-oriented programming: FedML lets users customize worker behavior, message handlers, and communication without managing the underlying communication mechanism.Workers inherit WorkerManager and use register_message_receive_handler and send_message to define message behavior.
- Message definition beyond gradient and model: FedML supports arbitrary message types beyond gradients and models, enabling auxiliary information exchange through the training procedure.Workers can send message types and parameters using train().
- Topology management: TopologyManager supports diverse FL topologies by allowing each trainer to query neighborhood worker IDs and send messages to arbitrary neighbors.The supported examples include vertical, split, decentralized, and hierarchical FL.
- Trainer and coordinator: FedML separates trainer and coordinator implementation from the framework, leaving their behavior to developers for greater flexibility.In FedAvg, the central worker acts as coordinator while other workers act as trainers.
- Privacy, security, and robustness: FedML includes hardware testbeds and security APIs for evaluating realistic system costs and implementing cryptographic, privacy, and robust-aggregation techniques.The supplied passages describe real-world mobile and IoT platforms, cryptographic primitives, robust aggregation, and reference attacks.
4 FedML Benchmark: Algorithms, Models, and Datasets
FedML provides diverse algorithm implementations and a standardized benchmark spanning multiple optimization settings, datasets, and model scales. These resources are intended to support reproducible development and fairer comparison of FL algorithms.
- Algorithms: FedML implements FL algorithms spanning diverse network topologies, exchanged information, and training procedures as reference examples and baselines.Examples include FedAvg, decentralized FL, vertical FL, split learning, FedNAS, and Turbo-Aggregate.
- Benchmark design: FedML benchmark fixes combinations of datasets, models, and non-I.I.D. partition methods to improve fairness in algorithm comparisons.Inconsistent choices across existing work make performance comparisons difficult.
- Benchmark design: The benchmark covers linear convex optimization, lightweight shallow neural networks, and deep neural networks for non-convex optimization.These categories organize the benchmark by model class and optimization setting.
- Models and datasets: The linear-model category includes MNIST, Federated EMNIST, and Synthetic (α, β), using logistic regression as the baseline model.This category targets convex optimization experiments.
- Models and datasets: The shallow-network category includes Federated EMNIST, CIFAR-100, Shakespeare, and StackOverflow under edge-device resource constraints.This category targets lightweight non-convex optimization experiments.
- Models and datasets: The deep-network category targets large DNN models, whose accuracy advantages motivate benchmarking them despite the challenges of edge-device FL.The paper also notes that large DNNs are commonly trained in cross-silo FL because of edge-device resource constraints.
5 Experiments
FedML reports benchmark results and compares distributed computing with standalone simulation. The experiments show a substantial speed advantage for distributed training on large CNNs and support many workers on limited GPU hardware.
- Benchmark results: FedML maintains benchmark experimental results as references for newly developed FL algorithms and updates them through Weight and Bias.The results are available through the project’s GitHub repository.
- Experimental evaluation: The experiments include modern CNN training results and a ResNet-56 test-accuracy trajectory.The supplied figures and tables identify these as experimental outputs.
- Training-time comparison: About 8 times slower: standalone simulation versus distributed computing with 10 parallel workers when training large CNNs.This comparison is reported in Table 6.
- Hardware scalability: FedML can run 112 workers on a server with 8 GPUs when training ResNet on CIFAR-10.The system supports multiprocessing within a single GPU card, allowing more workers than GPUs.
6 Conclusion
FedML is presented as a research-oriented federated learning library and benchmark intended to support algorithm development and fair comparison. The project is maintained as an evolving resource for the federated learning community.
- FedML provides an end-to-end toolkit for developing FL algorithms and fairly comparing them with existing algorithms.
- The paper frames FedML as both a federated learning library and benchmark.
- The project includes maintained source code, documentation, and a user community for continued research support.
B.1 Details of Supported Algorithms
FedML supports diverse federated learning algorithms and configurations through examples spanning centralized, decentralized, vertical, split, and architecture-search settings. These examples demonstrate support for customized communication and training procedures.
- Federated Averaging (FedAvg): FedAvg is implemented as a baseline in which workers train locally and the server aggregates client models by weighted coordinate-wise averaging.
- Decentralized FL: Decentralized FL is demonstrated with directed worker communication and no central server.
- Vertical Federated Learning (VFL): VFL supports shared samples with different features through privacy-preserving feature aggregation, logistic regression, customizable feature extractors, and two datasets.
- Split Learning: Split learning communicates activations and gradients at one model layer instead of weights for the entire model.
- Federated Neural Architecture Search (FedNAS): FedNAS enables scattered clients to collaboratively search for neural architectures while exchanging information beyond gradients.
B.2 Details of Datasets
FedML’s dataset coverage includes federated benchmarks with natural, synthetic, and cross-source heterogeneity. The section documents datasets, partitioning schemes, and the challenge of comparing results across diverse non-I.I.D. data and models.
- Federated EMNIST contains 62 digit and English-character classes partitioned by author, producing heterogeneity from individual writing styles.
- Federated CIFAR-100 randomly partitions training data among 500 clients with 100 examples per client using PAM.
- Shakespeare treats each speaking role in each play as a different device in a dataset drawn from Shakespeare’s complete works.
- The benchmark incorporates the federated StackOverflow dataset maintained by the TensorFlow Federated team.
- CIFAR-10 and CIFAR-100 use LDA to partition data according to the workers involved in each training round.
- CINIC-10 combines ImageNet and CIFAR-10 and supports evaluation on samples from similar but non-identical distributions.
C IoT Devices
FedML supports on-device federated learning experiments on two IoT platforms: Raspberry Pi 4 and NVIDIA Jetson Nano.
- FedML supports Raspberry Pi 4 for edge CPU computing.
- FedML supports NVIDIA Jetson Nano for edge GPU computing.
- The documented IoT-device coverage includes both an edge CPU and an edge GPU platform.
C.1 Raspberry Pi 4 (Edge CPU Computing - ARMv7l)
The Raspberry Pi 4 Desktop kit combines a Raspberry Pi 4 Model B with peripherals, power, and Raspberry Pi OS storage. The NVIDIA Jetson Nano Developer Kit is a compact, low-power computer for parallel neural-network applications.
- The kit includes a Raspberry Pi 4 Model B available with 2GB, 4GB, or 8GB of memory.
- It supplies a Raspberry Pi keyboard, mouse, and 15.3W USB-C power supply.
- The package includes a 16GB NOOBS microSD card with Raspberry Pi OS.
- The NVIDIA Jetson Nano Developer Kit runs multiple neural networks in parallel for image classification, object detection, segmentation, and speech processing, using as little as 5 watts.