FedLab

FedLab is a flexible federated learning (FL) framework designed to standardize FL simulation procedures. It provides modular tools and standard implementations for:

• Synchronous and Asynchronous Algorithms: Standardized implementations for various FL training modes.
• Communication Compression: Tools to handle package transmission and reduce communication overhead.
• Modular Customization: Ability to define custom communication strategies and federated optimization processes.
• Data Management: Tools for handling federated datasets and implementing IID (Independent and Identically Distributed) or non-IID data partitioning.

FedLab provides a PyTorch-based implementation of the LEAF benchmark, migrating it from the original TensorFlow version. It includes a unified dataloader interface located in fedlab_benchmarks/leaf/dataloader.py.

Supported datasets include:

• FEMNIST: Image classification (62 classes, 28x28 pixels, ~3500 users).
• Sentiment140: Sentiment analysis using Tweets (~660k users).
• Shakespeare: Next-character prediction using dialogues (~1129 users).
• Celeba: Image classification (Smiling vs. Not smiling, ~9343 users).
• Synthetic: Customizable classification task for challenging federated settings.
• Reddit: Next-word prediction using Reddit comments (~1.6M users).

FedLab provides two standard communication patterns out of the box:

1. Synchronous Mode: Uses SynchronousServerManager and PassiveClientManager. In this mode, clients typically wait for instructions from the server before performing local training.
2. Asynchronous Mode: Uses ServerAsynchronousManager and ClientActiveManager. This allows for more flexible communication where clients may not be strictly synchronized with the server's global state.

When building custom algorithms, you should choose or implement a pair that maintains a consistent communication protocol.

When preprocessing CelebA data, you can choose between two sampling strategies via the -s flag:

• IID (Independent and Identically Distributed): Each datapoint is equally likely to be sampled. This results in all users having the same underlying data distribution.
• Non-IID: The underlying distribution of data for each user is consistent with the raw data. This simulates real-world scenarios where data distributions vary between users.

FedLab provides utilities to prepare data for federated learning. This includes:

• Dataset Acquisition: Accessing supported federated datasets.
• Data Partitioning: Using partitioners to split data into client-side subsets, supporting both IID and non-IID settings to simulate realistic federated environments.

Standard Federated Learning (FL) optimization in FedLab is split into two distinct parts that you must implement to customize the training process:

1. Local Training (Client side): Handled by subclassing ClientTrainer. This defines how individual clients train on their local data.
2. Global Aggregation (Server side): Handled by subclassing ServerHandler. This defines how the central server aggregates updates from clients to update the global model.

Both ClientTrainer and ServerHandler inherit from ModelMaintainer, which provides the necessary infrastructure to manage PyTorch models, including GPU assignment and parameter serialization.

The output of the preprocessing steps is stored in .json files. Each file is a JSON object containing three keys:

1. 'users': A list of user identifiers.
2. 'num_samples': A list containing the number of samples associated with each user.
3. 'user_data': An object where keys are user names and values are lists of images. Each image is represented as a size-784 integer list (a flattened 28x28 image).

FedLab supports two main communication patterns for Federated Learning:

• Synchronous FL: The server controls the training rounds. The server first performs client sampling and then broadcasts the global model parameters to the selected clients.
• Asynchronous FL: The clients drive the training rounds. Clients request the current global model parameters from the server and then proceed with their local training independently.

FCUBE supports two partitioning methods via the FCUBEPartitioner class:

1. synthetic: Divides the 3D cube into 8 parts using planes $x_1=0$, $x_2=0$, and $x_3=0$. Symmetric pairs of parts are allocated to the same client. This results in exactly 4 clients with feature distribution skew but balanced labels.
2. iid: Standard Independent and Identically Distributed partitioning.

Use FCUBEPartitioner(data, partition='method_name') to perform the partition. The resulting object can be indexed by client ID to retrieve data indices.

# Synthetic partition (results in 4 clients)
synthetic_part = FCUBEPartitioner(trainset.data, partition="synthetic")

# IID partition
iid_part = FCUBEPartitioner(trainset.data, partition="iid")

FedLab supports various simulation modes to balance speed and resource usage. You can run simulations in Standalone mode, or use Cross-process modes to distribute workloads across multiple processes or machines.

Note that while cross-process simulation (e.g., 1M-10P on 1 machine with 10 processes) significantly reduces time compared to standalone execution, using multiple machines (e.g., 2M-10P) may introduce communication bottlenecks that can slow down the simulation depending on your hardware interconnects.

FedLab provides implementations for communication-efficient federated learning, allowing you to compare standard algorithms against compressed versions.

Key baselines and compression techniques available for testing include:

• QSGD (Quantized SGD) with various bit-widths: QSGD-4bit, QSGD-8bit, and QSGD-16bit.
• Top-k sparsification methods: Top-5%, Top-10%, and Top-20%.

These can be evaluated by following the communication tutorials provided in the repository.