FLib

FLib is a composable federated learning library for PyTorch. Nodes own models, optimizers, training, aggregation, evaluation, selection, and checkpointing. Architectures only coordinate communication, topology, update delivery, and round flow.

FLib provides:

• FedAvg, coordinate-wise Median, and Trimmed Mean strategies;
• centralized, synchronous decentralized, and hierarchical architectures;
• per-node aggregation functions;
• IID, Dirichlet Non-IID, and shard Non-IID partitioning;
• compact PyTorch reference models for common FL experiments;
• callbacks, typed results, and resumable checkpoints.

Installation

python -m pip install -e .
python -m pip install -e ".[dev,vision]"

FLib requires Python 3.10 or later. Torchvision is only required for vision examples such as MNIST.

Five-minute quickstart

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

from flib import (
    Centralized,
    Experiment,
    ExperimentConfig,
    FedAvg,
    TorchNode,
)


class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(4, 2)

    def forward(self, inputs):
        return self.linear(inputs)


features = torch.randn(120, 4)
labels = (features[:, 0] > 0).long()
client_datasets = [
    TensorDataset(features[index:index + 40], labels[index:index + 40])
    for index in range(0, 120, 40)
]
test_loader = DataLoader(TensorDataset(features, labels), batch_size=32)

server = TorchNode(
    node_id="server",
    model=MyModel(),
    aggregate_fn=FedAvg().aggregate,
)
clients = [
    TorchNode(node_id=f"client-{index}", model=MyModel())
    for index in range(len(client_datasets))
]

experiment = Experiment(
    architecture=Centralized(coordinator=server, clients=clients),
    client_datasets=client_datasets,
    test_loader=test_loader,
    config=ExperimentConfig(num_rounds=3, local_epochs=1),
)

result = experiment.run()
print(result.final_metrics)

Experiment can create training data loaders from datasets, but it does not create models or nodes. A sequence of datasets is matched to architecture.training_nodes by order. A mapping can bind datasets explicitly:

experiment = Experiment(
    architecture=architecture,
    client_datasets={
        "client-a": dataset_a,
        "client-b": dataset_b,
    },
    test_loader=test_loader,
)

Nodes with preconfigured train_loader and eval_loader can omit both data arguments.

Reference models

flib.models provides a deliberately small set of reference classifiers for tests and representative federated learning benchmarks:

from torch import nn

from flib.models import LeNet, LinearClassifier, MLP, resnet18

linear = LinearClassifier(
    input_dim=784,
    num_classes=10,
)
mlp = MLP(
    input_dim=784,
    hidden_dims=(256, 128),
    num_classes=10,
    dropout=0.1,
    activation=nn.ReLU,
)
lenet = LeNet(
    in_channels=1,
    num_classes=10,
)
cifar_resnet = resnet18(
    in_channels=3,
    num_classes=10,
    small_input=True,
    base_width=64,
)
standard_resnet = resnet18(
    in_channels=3,
    num_classes=1000,
    small_input=False,
)

LinearClassifier and MLP flatten inputs while preserving the batch dimension. LeNet uses adaptive pooling for both 28x28 and 32x32 images. resnet18 uses a 3x3 stem without max pooling for small inputs and a standard 7x7 stem with max pooling otherwise. Every model returns unnormalized logits, not probabilities, and can be paired with losses such as nn.CrossEntropyLoss.

The models do not download datasets or create data loaders. Dataset acquisition, transforms, partitioning, and batching remain separate user choices. FLib does not attempt to reproduce torchvision's full model collection or provide large general-purpose architectures.

BatchNorm running statistics can diverge across clients under Non-IID data and require an explicit federation policy. Experiments can keep statistics local, aggregate them separately, or replace BatchNorm with GroupNorm:

group_norm_resnet = resnet18(
    num_classes=10,
    norm_layer=lambda channels: nn.GroupNorm(8, channels),
)

Node aggregation contract

ModelUpdate contains the producing node ID, complete model state, sample count, base model version, and local metrics. AggregateFunction is a public protocol with this contract:

from collections.abc import Sequence

from flib import ModelState, ModelUpdate


def aggregate(updates: Sequence[ModelUpdate]) -> ModelState:
    ...

TorchNode.aggregate(updates) invokes its injected function, loads the returned state into its own model, increments model_version, and returns a detached CPU copy of the applied state. Calling it without aggregate_fn raises AggregationNotConfiguredError.

The same node class can therefore have different behavior per instance:

from flib import FedAvg, Median, TorchNode, TrimmedMean

server = TorchNode(
    node_id="server",
    model=MyModel(),
    aggregate_fn=FedAvg().aggregate,
)
robust_peer = TorchNode(
    node_id="peer-a",
    model=MyModel(),
    aggregate_fn=Median().aggregate,
)
trimmed_peer = TorchNode(
    node_id="peer-b",
    model=MyModel(),
    aggregate_fn=TrimmedMean(trim_ratio=0.2).aggregate,
)

TorchNode also accepts an optimizer instance or an optimizer_factory, custom loss and metric functions, custom train and evaluation functions, gradient clipping, and standalone node checkpoints.

Architectures

Architectures never receive or store an Aggregator. The node responsible for a stage performs that stage's aggregation.

Centralized

architecture = Centralized(
    coordinator=server,
    clients=clients,
)

The coordinator selects clients, sends its current model, receives local updates, calls its own aggregate(), and distributes its resulting model to all clients.

Decentralized

from flib import Decentralized, FullyConnected

architecture = Decentralized(
    nodes=[peer_a, peer_b, peer_c],
    topology=FullyConnected(),
)

Each selected peer trains locally and receives updates from topology neighbors. Each peer then calls its own aggregate(). All local updates are captured and independently cloned before any aggregation call, so an earlier peer's result cannot affect a later peer's inputs in the same round. Network evaluation is the sample-weighted mean of peer evaluation metrics; it does not create or own a global aggregation strategy.

Hierarchical

from flib import Hierarchical

architecture = Hierarchical(
    cloud=cloud,
    groups={
        edge_a: clients_a,
        edge_b: clients_b,
    },
)

The mapping explicitly represents edge-to-client relationships. Each edge calls its own aggregate() for its clients. The cloud receives typed edge updates and calls its own aggregate() in the second stage. No node IDs or locations are used to infer groups.

Data partitioning

from flib import partition_dirichlet, partition_iid, partition_shards

iid_clients = partition_iid(dataset, num_clients=10, seed=42)

dirichlet_clients = partition_dirichlet(
    dataset,
    num_clients=10,
    alpha=0.3,
    min_size=20,
    seed=42,
)

shard_clients = partition_shards(
    dataset,
    num_clients=10,
    shards_per_client=2,
    seed=42,
)

Smaller Dirichlet alpha values create stronger label skew. Labels are inferred from TensorDataset, Subset, or a dataset's targets or labels attribute. Custom datasets can pass labels= explicitly.

Custom aggregate strategies

An aggregate function does not need to inherit from Aggregator; any callable matching AggregateFunction is accepted:

from collections.abc import Sequence

import torch

from flib import ModelState, ModelUpdate


def mean(updates: Sequence[ModelUpdate]) -> ModelState:
    if not updates:
        raise ValueError("At least one update is required.")
    return {
        name: torch.stack(
            [update.state_dict[name] for update in updates]
        ).mean(dim=0)
        for name in updates[0].state_dict
    }


node = TorchNode(
    node_id="custom",
    model=MyModel(),
    aggregate_fn=mean,
)

Aggregator remains available as an optional reusable strategy base class. Built-in strategies expose bound aggregate methods that can be passed directly to nodes.

Checkpointing and resume

experiment.run(max_rounds=3)
experiment.save_checkpoint("checkpoints/round-3.pt")

restored = Experiment(
    architecture=Centralized(
        coordinator=new_server,
        clients=new_clients,
    ),
    client_datasets=client_datasets,
    test_loader=test_loader,
)
restored.load_checkpoint("checkpoints/round-3.pt")
result = restored.resume()

Recreate the same architecture roles, node IDs, models, optimizers, and aggregate functions before loading. Checkpoints contain node model and optimizer states, architecture role state, completed rounds, configuration, random states, and the FLib version. Python callables are configuration and are not serialized.

Breaking changes in the node-owned aggregation API

This refactor intentionally changes construction and checkpoint contracts:

• Centralized(aggregator=...) is replaced by Centralized(coordinator=..., clients=...).
• Decentralized(aggregator=..., topology=...) is replaced by Decentralized(nodes=..., topology=...).
• Hierarchical(aggregator=..., client_groups=...) is replaced by Hierarchical(cloud=..., groups={edge_node: client_nodes}).
• Experiment no longer accepts model_factory, node_factory, or node_kwargs; construct nodes first and pass them through an architecture.
• Architecture checkpoints no longer contain an aggregator field. Aggregation functions must be configured on recreated nodes before loading a checkpoint.
• Checkpoints created with the previous architecture-owned aggregator layout are not compatible with this version.

The breaking API is released as FLib 0.2.0. No compatibility constructor or automatic migration shim is provided for architecture-owned aggregators.

MNIST example

python examples/mnist/quickstart.py

The example downloads MNIST through torchvision, creates five IID client partitions, trains flib.models.LeNet with centralized FedAvg, and saves a final checkpoint. Torchvision is used only by the example for the dataset and transforms; the reference model package depends only on PyTorch.

Development

python -m pytest
ruff check src tests examples
mypy src/flib