FedML: A Research Library and Benchmark for Federated Machine Learning

In plain words

Federated learning trains one shared model across many separate data holders, such as phones or organisations, without collecting all their data in one place. Writing such a system is hard: each participant may run different code, send different kinds of messages, and see a different slice of the data. Results from different research groups are also hard to compare, because each group often uses its own datasets, models and data splits.

FedML is a research library that tries to fix both problems. You describe what each participant does and which messages it sends, and the library handles the communication underneath. The same code can run on one computer, on several machines, or on real phones and small boards. The library also ships common algorithms as starting points, and a benchmark that fixes which datasets, models and data splits to use, so that a new method can be compared fairly.

TL;DR

• FedML is an open library plus a benchmark for federated learning (FL) research §1.
• Its core is a worker-oriented API: each worker defines the messages it sends and handles the ones it receives, and can exchange payloads beyond gradients and models §3.
• It provides reference implementations of FedAvg, decentralized FL, vertical FL, split learning, FedNAS and Turbo-Aggregate §4.1 app. B.1.
• The benchmark fixes datasets, models and non-IID partitions §4.2 tab. 2 tab. 3 tab. 4. The author-reported FedAvg drop from IID to non-IID is 1.6 to 9.1 points across the reported runs tab. 5.
• The paper states no training hyperparameters. Any exact reproduction has to come from the repository, and its example configs do not match the paper's 10-worker setup (see Pitfalls).

Problem

FL trains models from scattered, isolated data and differs from data-centre training in statistical heterogeneity, system constraints and trustworthiness §1. Existing libraries are either data-centre training frameworks that cannot express arbitrary message flows and topologies, or simulation-only libraries that support centralised FedAvg on one machine, or production systems that are not built for algorithm prototyping §1 tab. 1. The paper also argues that benchmark settings (datasets, models, partitions, number of clients per round) differ so much between published papers that results are hard to compare §1 tab. 8.

Method

Architecture §2 fig. 1. The library has two layers. FedML-core separates communication from training. Its communication module uses MPI. A TopologyManager inside that module holds the network topology, and the module also supports security and privacy functions §2. The model-training module is built on PyTorch, and users implement their own trainers and coordinators §2. FedML-API is built on FedML-core and contains the algorithms, models and datasets, kept separate from one another §2. Two on-device testbeds are described: one for Android smartphones and one for Raspberry Pi 4 and NVIDIA Jetson Nano §2 app. C.

Worker-oriented programming §3 fig. 2. A user writes a worker by subclassing a worker manager class and calling its predefined functions to register receive handlers and to send messages. The communication mechanism stays hidden §3, "Worker/client-oriented programming". The paper contrasts this with a training-procedure style, where the whole distributed training loop is written once and the messaging is fixed to gradient aggregation through an all-reduce call fig. 2.

Messages beyond gradients and models §3. Each worker defines the message types it sends. Any message type and its parameters can be sent from the training function. Sending normally happens after the worker has handled an incoming message §3 fig. 2.

Topology management §3 fig. 3. After setup, each trainer can query its neighbour IDs from the TopologyManager and send messages to any of them. The paper lists vertical FL, split learning, decentralized FL and hierarchical FL as topologies that this supports §3.

Coordinator and trainer. These are left entirely to the developer. The paper points to the reference implementations for examples §3.

Security, privacy and robustness §3. The paper describes low-level primitives for secret sharing, key agreement, digital signatures and public-key infrastructure. Lagrange Coded Computing and a secure-aggregation sample are stated as plans, not as delivered features §3. The paper also lists robust aggregators (norm-difference clipping, weak differential privacy, RFA, Krum and Multi-Krum) and two backdoor attacks (model replacement and edge-case) as supported §3.

Federated optimizer, FedAvg §4.1 app. B.1. Each selected client trains its local model for several epochs and uploads it. The server aggregates the uploads by weighted coordinate-wise averaging, where each client's weight is its local number of data points. The server then sends the new global model to all clients app. B.1. The official code implements this as follows.

# FedAvg, one communication round  [app. B.1]
# Server side: aggregate the uploaded client models by sample count.
def fedavg_round(global_weights, client_ids, client_data, client_sizes):
    # Each selected client starts from the broadcast global model  [app. B.1]
    # Local training uses plain SGD; see the next pseudocode block  [app. B.1]
    updates = []
    for client_id in client_ids:
        local_weights = local_sgd(global_weights, client_data[client_id])
        updates.append((client_sizes[client_id], local_weights))
    # Weight of each client = its share of all samples in this round  [app. B.1]
    # [code: python/fedml/simulation/mpi/fedavg/FedAVGAggregator.py#L99-L116]
    total = sum(size for size, _ in updates)
    new_weights = {}
    for name in global_weights:
        new_weights[name] = sum(size * w[name] / total for size, w in updates)
    return new_weights

# Local update on one client
# [code: python/fedml/ml/trainer/my_model_trainer_classification.py#L21-L53]
def local_sgd(weights, batches, epochs, learning_rate):
    # Start from the broadcast global weights
    model = load_model(weights)
    # Plain SGD: no momentum, and no weight decay in this branch
    # [code: python/fedml/ml/trainer/my_model_trainer_classification.py#L29-L33]
    optimizer = sgd(model.parameters(), lr=learning_rate)
    for epoch in range(epochs):
        for inputs, labels in batches:
            # Cross-entropy on one minibatch; one optimizer step per batch
            loss = cross_entropy(model(inputs), labels)
            loss.backward()
            optimizer.step()
    return model.state_dict()

Client sampling [code]. If the number of clients per round equals the total, all clients take part. Otherwise a random subset is drawn with np.random.seed(round_idx), so every run selects the same clients in each round FedAVGAggregator.py. The paper does not state the participation rule.

Non-IID partition, latent Dirichlet allocation app. B.2 tab. 4. For CIFAR-10 the official code draws, for each class, Dirichlet proportions over the clients and splits that class accordingly. It zeroes proportions for clients already above the average size and retries until every client has at least 10 samples data_loader.py.

# Heterogeneous partition of CIFAR-10 by latent Dirichlet allocation
# [code: python/fedml/data/cifar10/data_loader.py#L122-L160]
def partition_hetero(labels, num_clients, alpha):
    # Retry until every client holds at least 10 samples
    while smallest_client_size < 10:
        for each class k:
            indices = shuffled(indices_of_class(k))
            # Dirichlet weights over clients, concentration alpha
            proportions = dirichlet(alpha, num_clients)
            # Balance: clients already at or above average get no more
            proportions = proportions * (size < total / num_clients)
            proportions = proportions / proportions.sum()
            split_points = cumsum(proportions) * len(indices)
            assign(indices, split_points)
    return client_index_lists

Datasets and models §4.2 tab. 2 tab. 3 tab. 4. The benchmark has three categories. The linear category uses convex models on MNIST, Federated EMNIST and the synthetic (α, β) data from FedProx, with logistic regression tab. 2. The shallow-network category uses Federated EMNIST, CIFAR-100, Shakespeare and StackOverflow with small CNNs, ResNet-18 with group normalization, and LSTMs tab. 3. The deep-network category uses CIFAR-10, CIFAR-100 and CINIC-10 with ResNet-56 and MobileNet, plus StackOverflow with an LSTM tab. 4. The paper describes the partitions in words app. B.2.

Experiments §5 fig. 5 tab. 5 tab. 6. The authors train ResNet-56 and MobileNet with FedAvg on CIFAR-10, CIFAR-100 and CINIC-10, with 10 workers, under an IID and an LDA non-IID split. They also compare wall-clock time of a single-GPU standalone simulation against multi-GPU distributed computing. The text states that a single machine with 8 GPUs can run 112 ResNet workers on CIFAR-10 through multi-processing on each GPU §5.

Hyperparameters

Only the values in the table below are stated. Everything else is "not stated" in the paper. Code values are shown only where they exist in the repository, and they belong to example configs, not to the paper's runs.

Results

Author-reported. Setup: FedAvg, 10 workers, ResNet-56 and MobileNet, LDA partition, hardware in tab. 6. The table lists IID and non-IID test accuracy per model; the unit is not stated in the table, and the values look like percentages. "Drop" is the paper's own difference tab. 5.

Wall-clock time for training the two CNNs with 10 workers tab. 6:

The paper reads this as the standalone simulation being about 8 times slower than distributed computing with 10 parallel workers §5. The two rows do not use the same number of GPUs (see Pitfalls). Training curves for ResNet-56 on the three datasets are shown in fig. 5 and are not reproduced here.

Paper vs code

Checked against the official repository at code_commit (FedML-AI/FedML, Apache-2.0) LICENSE. The paper links this repository in its IoT section and in its experiments section §2 §5.

• Structure. The paper's FedML-core and FedML-API appear to correspond to python/fedml/core/ and python/fedml/simulation/mpi/, and the topology managers are in python/fedml/core/distributed/topology/ [code]. The repository also contains cross-silo, cross-device and serving packages . This commit is therefore not necessarily the code the paper's runs used.
• FedAvg aggregation. Matches the paper: weights are local sample counts over the round total app. B.1 FedAVGAggregator.py.
• Local optimizer. The paper does not name it. The code uses plain SGD with no momentum and no weight decay in the sgd branch, although the example configs set weight_decay my_model_trainer_classification.py. The Adam branch does use weight decay my_model_trainer_classification.py.
• Gradient clipping. Not in the paper. The code has it commented out my_model_trainer_classification.py.
• Client sampling. Not in the paper. The code selects all clients when the per-round number equals the total, and otherwise seeds NumPy with the round index FedAVGAggregator.py.
• Partition. The paper says LDA tab. 4. The code's heterogeneous partition matches that description data_loader.py. The paper's appendix says CIFAR-100 uses Pachinko Allocation app. B.2, while Table 4 lists LDA for CIFAR-100. This is a conflict inside the paper.
• Worker count. Table 5 uses 10 workers. The example configs use 20 clients with 10 per round cifar10_resnet56_config.yaml, and the MPI CINIC-10 example uses 5 clients with 3 per round fedml_config.yaml. No config in the repository matches the paper's experiment.
• Topology defaults. The paper gives no numbers. The symmetric topology defaults to 2 neighbours symmetric_topology_manager.py, and the asymmetric one to 3 undirected and 3 out-directed neighbours asymmetric_topology_manager.py.

Pitfalls

• No training recipe in the paper. Rounds, epochs, batch size, learning rate, optimizer, client participation and α are not stated for tab. 5. Reproducing those numbers means finding the exact configuration that was run, which the paper does not give. The repository's examples are not that configuration (see Paper vs code).
• Weight decay is silently dropped with SGD. A config that sets weight_decay with client_optimizer: sgd has no effect in the local trainer my_model_trainer_classification.py.
• Fixed client subsets. Sampling seeds NumPy with the round index, so repeated runs with the same configuration select identical client subsets FedAVGAggregator.py. The paper does not say whether Table 5 used this rule.
• Partition retries. The LDA split loops until every client has at least 10 samples, with no iteration cap. The balancing step also changes which clients receive a class's samples data_loader.py. The paper gives no α, so the effect of the chosen α cannot be checked.
• StackOverflow numbers conflict. Table 3 gives 16,586,035 test samples and an RNN with 1 LSTM and 2 FC layers. Table 4 gives 10,586,035 test samples and 2 LSTM and 1 FC layers tab. 3 tab. 4. Table 4 also writes "342477 (10)" for the client count, and the parenthesis is not explained.
• CIFAR-100 partition conflict. Table 3 uses Pachinko Allocation with ResNet-18 and group normalization; Table 4 uses LDA with ResNet-56 tab. 3 tab. 4 app. B.2.
• Wall-clock comparison is not like for like. The standalone row is described as single-GPU and the distributed row as multi-GPU on a machine with 8 GPUs tab. 6. The 8× speed-up therefore mixes parallelism with the choice of paradigm.
• Claims to check. The paper says the benchmark has "verified baseline results" §1 (iv), but Table 5 covers six configurations, all with FedAvg. Table 1 is the authors' own feature comparison and is not independently checked.
• Security features are partly planned. LCC and the secure-aggregation sample are described as plans §3. Do not assume they exist in the code version you use.
• The 112-worker claim is hardware-specific. It is for ResNet on CIFAR-10 on one server with 8 GPUs §5 tab. 6.

Lineage

• Builds on: FedAvg (McMahan et al., 2017), which the paper cites as its baseline algorithm §4.1. Not compiled here.
• Compared against: TensorFlow Federated, PySyft, LEAF, FATE and PaddleFL §1 tab. 1. These are named, not compiled here.
• Superseded by: not determined from this card's content.