Federated Learning on a cluster of Raspberry Pis and Jetson Orin

Classic machine learning collects all data in one place and trains a model there. Federated Learning (FL) flips that around: the model travels to the data, each device trains locally, and only model updates (weights or gradients) are sent back to a coordinator. Raw data never leaves the device. In this post we'll build a real federated learning testbed from a handful of Raspberry Pi 5 boards and one or two NVIDIA Jetson Orin modules using the Flower framework and PyTorch. We'll cover the theory, the hardware, the network, the code, and the practical problems you hit once the devices are real and not simulated.

The Theory: How Federated Learning Works

A federated training run is organised in rounds. In each round:

1. The server holds the current global model w_t and selects a subset of available clients.
2. It sends w_t to the selected clients.
3. Each client k trains for E local epochs on its private dataset of n_k samples, producing w_t^k.
4. Clients send back the updated weights (plus the number of samples and metrics).
5. The server aggregates the updates into w_(t+1). In FedAvg this is a weighted mean: w_(t+1) = Σ_k (n_k / n) · w_t^k.

Federated vs Centralized vs Distributed Training

Aspect	Centralized	Distributed (data-parallel)	Federated
Where data lives	One server	Data center, shuffled across workers	On the devices, never moved
Data distribution	IID (shuffled)	IID	Usually non-IID, unbalanced
Network	None	Fast, reliable (InfiniBand, NVLink)	Slow, unreliable (Wi-Fi, LTE)
Sync frequency	-	Every batch	Every round (many local steps)
Workers	1	Homogeneous GPUs	Heterogeneous, may drop out
Privacy	Low	Low	Higher (plus DP / secure aggregation)

Common Aggregation Strategies

Strategy	Idea	When to use
FedAvg	Sample-weighted mean of client weights	Baseline, roughly IID data
FedProx	Adds proximal term μ/2 · ||w - w_t||² to local loss	Non-IID data, clients doing very different amounts of work
FedAdam / FedYogi	Server-side adaptive optimizer on the averaged update	Faster convergence, noisy updates
FedMedian / Krum	Robust statistics instead of mean	Untrusted clients, poisoning protection
SCAFFOLD	Control variates to correct client drift	Strongly non-IID data

Edge FL is interesting exactly because the devices differ. A Raspberry Pi trains on CPU only; a Jetson Orin has a CUDA GPU with Tensor Cores. Mixing them reproduces the system heterogeneity you see in real deployments.

Device	CPU	GPU / NPU	RAM	AI perf.	Power	Role
Raspberry Pi 4 (4/8 GB)	4× Cortex-A72 @ 1.8 GHz	-	4-8 GB	very low	3-7 W	Weak client
Raspberry Pi 5 (8/16 GB)	4× Cortex-A76 @ 2.4 GHz	optional Hailo-8L AI HAT (inference only)	8-16 GB	low (CPU training)	5-12 W	Client
Jetson Orin Nano Super 8 GB	6× Cortex-A78AE	1024 CUDA, 32 Tensor Cores	8 GB shared	~67 TOPS	7-25 W	Strong client
Jetson AGX Orin 64 GB	12× Cortex-A78AE	2048 CUDA, 64 Tensor Cores	64 GB shared	~275 TOPS	15-60 W	Strong client or server
x86 mini PC / laptop	any	optional	16+ GB	-	15-65 W	Server (SuperLink), monitoring

Other useful equipment:

• A gigabit switch (managed if you want VLANs) and short Cat6 cables - Wi-Fi works too and is more realistic, but makes debugging harder.
• A PoE+ switch with PoE HATs for the Pis to reduce cabling, or a multi-port USB-C PD charger (Pi 5 needs 5 V / 5 A).
• Active coolers - a Pi 5 training at 100% CPU will throttle without one.
• NVMe SSDs (Pi 5 via M.2 HAT, Jetson natively) - SD cards are slow for datasets and swap.
• A USB power meter or smart plug to measure energy per round.

Cluster Topology

Example Network Plan

Hostname	IP	Device	Data partition
fl-server	10.0.0.10	mini PC	test set only
rpi-01..04	10.0.0.21-24	Raspberry Pi 5 8 GB	partitions 0-3
orin-01	10.0.0.31	Jetson Orin Nano Super	partition 4

Raspberry Pi 5

Flash Raspberry Pi OS Lite (64-bit) - PyTorch only ships aarch64 wheels. Then:

sudo apt update && sudo apt full-upgrade -y
sudo apt install -y python3-venv python3-dev git
# more swap helps when installing / with larger models
sudo sed -i 's/^CONF_SWAPSIZE=.*/CONF_SWAPSIZE=2048/' /etc/dphys-swapfile
sudo systemctl restart dphys-swapfile

curl -LsSf https://astral.sh/uv/install.sh | sh
uv venv ~/fl-env && source ~/fl-env/bin/activate
uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
uv pip install "flwr[simulation]" flwr-datasets

Jetson Orin

Flash JetPack 6.x (Ubuntu 22.04, CUDA 12.x) with SDK Manager or the SD-card image. The generic PyPI torch wheel has no CUDA support for Jetson - use NVIDIA's wheel matching your JetPack, or simply run in a container from dustynv/jetson-containers:

sudo nvpmodel -m 0          # max performance mode
sudo jetson_clocks          # lock clocks for stable benchmarks
sudo apt install -y python3-pip && pip3 install jetson-stats   # jtop

# Option A: NVIDIA wheel for your JetPack version (see NVIDIA forums / pypi.jetson-ai-lab)
python3 -m venv ~/fl-env && source ~/fl-env/bin/activate
pip install torch torchvision --index-url https://pypi.jetson-ai-lab.dev/jp6/cu126
pip install flwr flwr-datasets

python -c "import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name(0))"

Software Stack per Node

Layer	Raspberry Pi 5	Jetson Orin	Server
OS	Raspberry Pi OS 64-bit (Bookworm)	JetPack 6 (Ubuntu 22.04)	Ubuntu 24.04
Python	3.11	3.10	3.11+
PyTorch	CPU wheel	CUDA wheel from NVIDIA	CPU is enough
Flower	flower-supernode	flower-supernode	flower-superlink
Monitoring	node_exporter, vcgencmd	jtop, tegrastats	Prometheus, Grafana

Important: install the same Flower version everywhere. Mismatched versions between SuperLink and SuperNodes are the most common cause of cryptic gRPC errors.

Since Flower 1.10+ the deployment engine separates long-running infrastructure from the user code:

• SuperLink - runs on the server, accepts SuperNode connections (Fleet API) and receives runs from flwr run (Exec API).
• SuperNode - long-running process on each Pi / Jetson that connects to the SuperLink.
• ServerApp - your strategy and round logic.
• ClientApp - your local training/evaluation code, executed by the SuperNode.

Create a project with flwr new edge-fl --framework PyTorch or by hand with this layout:

edge-fl/
├── pyproject.toml
└── edge_fl/
    ├── __init__.py
    ├── task.py        # model, data loading, train/test
    ├── client_app.py  # ClientApp
    └── server_app.py  # ServerApp + strategy

pyproject.toml

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[project]
name = "edge-fl"
version = "1.0.0"
dependencies = ["flwr>=1.20", "flwr-datasets[vision]>=0.5", "torch", "torchvision"]

[tool.hatch.build.targets.wheel]
packages = ["."]

[tool.flwr.app]
publisher = "scientific-dev"

[tool.flwr.app.components]
serverapp = "edge_fl.server_app:app"
clientapp = "edge_fl.client_app:app"

[tool.flwr.app.config]
num-server-rounds = 30
fraction-fit = 1.0
local-epochs = 1
batch-size = 32
learning-rate = 0.01
proximal-mu = 0.1

[tool.flwr.federations]
default = "local-sim"

[tool.flwr.federations.local-sim]
options.num-supernodes = 5

[tool.flwr.federations.edge-cluster]
address = "10.0.0.10:9093"
root-certificates = "./certificates/ca.crt"

task.py - model, data and training

A small CNN for CIFAR-10 is a good fit: it trains in reasonable time on a Pi 5 CPU and is fast on the Orin GPU. The data is split with flwr-datasets using a Dirichlet partitioner to simulate non-IID label distributions.

import torch
import torch.nn as nn
import torch.nn.functional as F
from collections import OrderedDict
from flwr_datasets import FederatedDataset
from flwr_datasets.partitioner import DirichletPartitioner
from torch.utils.data import DataLoader
from torchvision.transforms import Compose, Normalize, ToTensor


class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 8 * 8, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = torch.flatten(x, 1)
        return self.fc2(F.relu(self.fc1(x)))


def get_device():
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


_fds = None
_transforms = Compose([ToTensor(), Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])


def _apply(batch):
    batch["img"] = [_transforms(img) for img in batch["img"]]
    return batch


def load_data(partition_id, num_partitions, batch_size):
    global _fds
    if _fds is None:
        partitioner = DirichletPartitioner(
            num_partitions=num_partitions, partition_by="label", alpha=0.5, seed=42
        )
        _fds = FederatedDataset(dataset="uoft-cs/cifar10", partitioners={"train": partitioner})
    partition = _fds.load_partition(partition_id).train_test_split(test_size=0.2, seed=42)
    partition = partition.with_transform(_apply)
    train = DataLoader(partition["train"], batch_size=batch_size, shuffle=True)
    test = DataLoader(partition["test"], batch_size=batch_size)
    return train, test


def train(net, loader, epochs, lr, device, global_params=None, mu=0.0):
    net.to(device).train()
    opt = torch.optim.SGD(net.parameters(), lr=lr, momentum=0.9)
    criterion = nn.CrossEntropyLoss()
    total_loss = 0.0
    for _ in range(epochs):
        for batch in loader:
            x, y = batch["img"].to(device), batch["label"].to(device)
            opt.zero_grad()
            loss = criterion(net(x), y)
            if global_params is not None and mu > 0:  # FedProx proximal term
                prox = sum((w - g).norm(2) ** 2 for w, g in zip(net.parameters(), global_params))
                loss = loss + mu / 2 * prox
            loss.backward()
            opt.step()
            total_loss += loss.item()
    return total_loss / (len(loader) * epochs)


def test(net, loader, device):
    net.to(device).eval()
    criterion = nn.CrossEntropyLoss()
    correct, loss = 0, 0.0
    with torch.no_grad():
        for batch in loader:
            x, y = batch["img"].to(device), batch["label"].to(device)
            out = net(x)
            loss += criterion(out, y).item()
            correct += (out.argmax(1) == y).sum().item()
    return loss / len(loader), correct / len(loader.dataset)


def get_weights(net):
    return [v.cpu().numpy() for v in net.state_dict().values()]


def set_weights(net, weights):
    state = OrderedDict((k, torch.tensor(v)) for k, v in zip(net.state_dict().keys(), weights))
    net.load_state_dict(state, strict=True)

client_app.py

The same ClientApp runs on Pis and Jetsons. The partition-id comes from the SuperNode's --node-config, so each device knows its slice of data. The client also reports its training time and device type - useful for analysing stragglers.

import platform
import time

import torch
from flwr.client import ClientApp, NumPyClient
from flwr.common import Context

from edge_fl.task import Net, get_device, get_weights, load_data, set_weights, test, train


class EdgeClient(NumPyClient):
    def __init__(self, trainloader, valloader, epochs, lr, mu):
        self.net = Net()
        self.trainloader, self.valloader = trainloader, valloader
        self.epochs, self.lr, self.mu = epochs, lr, mu
        self.device = get_device()

    def fit(self, parameters, config):
        set_weights(self.net, parameters)
        global_params = [p.detach().clone().to(self.device) for p in self.net.to(self.device).parameters()]
        start = time.perf_counter()
        loss = train(self.net, self.trainloader, self.epochs, self.lr, self.device, global_params, self.mu)
        metrics = {
            "train_loss": loss,
            "train_time": time.perf_counter() - start,
            "device": "cuda" if self.device.type == "cuda" else platform.machine(),
        }
        return get_weights(self.net), len(self.trainloader.dataset), metrics

    def evaluate(self, parameters, config):
        set_weights(self.net, parameters)
        loss, acc = test(self.net, self.valloader, self.device)
        return loss, len(self.valloader.dataset), {"accuracy": acc}


def client_fn(context: Context):
    pid = int(context.node_config["partition-id"])
    num = int(context.node_config["num-partitions"])
    cfg = context.run_config
    trainloader, valloader = load_data(pid, num, cfg["batch-size"])
    torch.set_num_threads(4)  # Pi 5 has 4 cores
    return EdgeClient(trainloader, valloader, cfg["local-epochs"], cfg["learning-rate"], cfg["proximal-mu"]).to_client()


app = ClientApp(client_fn=client_fn)

server_app.py

from flwr.common import Context, Metrics, ndarrays_to_parameters
from flwr.server import ServerApp, ServerAppComponents, ServerConfig
from flwr.server.strategy import FedProx

from edge_fl.task import Net, get_weights


def weighted_average(metrics: list[tuple[int, Metrics]]) -> Metrics:
    total = sum(n for n, _ in metrics)
    return {"accuracy": sum(n * m["accuracy"] for n, m in metrics) / total}


def fit_metrics(metrics: list[tuple[int, Metrics]]) -> Metrics:
    times = [m["train_time"] for _, m in metrics]
    return {"max_train_time": max(times), "min_train_time": min(times)}


def server_fn(context: Context):
    cfg = context.run_config
    strategy = FedProx(
        fraction_fit=cfg["fraction-fit"],
        fraction_evaluate=1.0,
        min_available_clients=3,
        min_fit_clients=3,
        proximal_mu=cfg["proximal-mu"],
        initial_parameters=ndarrays_to_parameters(get_weights(Net())),
        evaluate_metrics_aggregation_fn=weighted_average,
        fit_metrics_aggregation_fn=fit_metrics,
    )
    config = ServerConfig(num_rounds=cfg["num-server-rounds"], round_timeout=600)
    return ServerAppComponents(strategy=strategy, config=config)


app = ServerApp(server_fn=server_fn)

Tip: before touching any hardware, run flwr run . - it uses the local-sim federation and simulates 5 clients on your laptop. Fix bugs there, not over SSH on five boards.

1. TLS certificates

Create a small CA and a server certificate whose SAN contains the SuperLink IP, then copy ca.crt to every device:

openssl req -x509 -newkey rsa:4096 -days 365 -nodes -keyout ca.key -out ca.crt -subj "/CN=fl-ca"
openssl req -newkey rsa:4096 -nodes -keyout server.key -out server.csr -subj "/CN=fl-server"
printf "subjectAltName=IP:10.0.0.10,DNS:fl-server\n" > san.ext
openssl x509 -req -in server.csr -CA ca.crt -CAkey ca.key -CAcreateserial -out server.pem -days 365 -extfile san.ext
for h in rpi-01 rpi-02 rpi-03 rpi-04 orin-01; do scp ca.crt $h:~/certificates/; done

2. Start the SuperLink (server)

flower-superlink \
  --ssl-ca-certfile certificates/ca.crt \
  --ssl-certfile certificates/server.pem \
  --ssl-keyfile certificates/server.key

3. Start SuperNodes (each device)

# on rpi-01 (partition 0); rpi-02 uses partition-id=1, ... orin-01 uses 4
flower-supernode \
  --superlink 10.0.0.10:9092 \
  --root-certificates ~/certificates/ca.crt \
  --node-config "partition-id=0 num-partitions=5"

To survive reboots, wrap it in a systemd unit:

# /etc/systemd/system/flower-supernode.service
[Unit]
Description=Flower SuperNode
After=network-online.target
Wants=network-online.target

[Service]
User=pi
WorkingDirectory=/home/pi
ExecStart=/home/pi/fl-env/bin/flower-supernode --superlink 10.0.0.10:9092 \
  --root-certificates /home/pi/certificates/ca.crt \
  --node-config "partition-id=0 num-partitions=5"
Restart=always
RestartSec=10

[Install]
WantedBy=multi-user.target

sudo systemctl daemon-reload && sudo systemctl enable --now flower-supernode

Managing five boards by hand gets old fast - an Ansible playbook (or parallel-ssh) that installs the venv, copies the cert and templates the partition-id per host is worth the 20 minutes.

4. Launch the run

flwr run . edge-cluster --stream
flwr ls . edge-cluster          # list runs
flwr log <run-id> . edge-cluster

Lifecycle of a Node

Stragglers and heterogeneity

In a synchronous round the server waits for the slowest client. With a Pi 5 and an Orin in the same round, the Orin sits idle most of the time. Typical figures for the CNN above (1 local epoch, ~8000 samples per client):

Device	Time / local epoch	Avg power	Energy / round
Raspberry Pi 4 (4 threads)	~110 s	~6 W	~0.18 Wh
Raspberry Pi 5 (4 threads)	~45 s	~9 W	~0.11 Wh
Jetson Orin Nano Super (GPU)	~4 s	~15 W	~0.02 Wh
Jetson AGX Orin (GPU)	~2 s	~30 W	~0.02 Wh

These are indicative values from a small CNN; measure on your own setup - cooling, power mode and data loading change them a lot.

Mitigations:

• Time budget instead of fixed epochs - send config["time_budget"] and let each client train as many batches as fit; FedProx keeps divergent amounts of work under control.
• Adaptive local work - give the Orin more epochs or a larger batch size.
• round_timeout and min_fit_clients - don't block a round on a node that dropped out.
• Asynchronous / buffered aggregation (FedBuff) - aggregate as soon as K updates arrive.

Non-IID data

The Dirichlet alpha controls how skewed each client's label distribution is:

alpha	Distribution	Effect on FedAvg
100	Almost IID	Close to centralized accuracy
1.0	Moderately skewed	Slower convergence
0.5	Skewed (our default)	Noticeable drift, FedProx helps
0.1	Each client sees 1-2 classes	Oscillating accuracy; use SCAFFOLD / FedProx, fewer local epochs

Communication cost

Our CNN has ~545k parameters → ~2.2 MB in float32 per direction per client per round. For 5 clients and 30 rounds that's ~660 MB of traffic in total - fine on Ethernet, noticeable on LTE. Options: send float16 weights (half the size), quantize to 8 bits, sparsify updates (top-k), or train more locally and use fewer rounds.

Thermals and memory

• Watch vcgencmd measure_temp and vcgencmd get_throttled on the Pi - throttling starts at 80-85 °C.
• On Jetson, CPU and GPU share RAM: use jtop, keep num_workers of the DataLoader low, and consider mixed precision (torch.autocast("cuda")).
• Pre-download and cache the dataset partition on each node so the first round doesn't spend minutes on downloads.

FL keeps raw data local, but weight updates can still leak information (membership inference, gradient inversion). Layer the defences:

Threat	Defence	Flower support
Eavesdropping on the network	TLS	--ssl-* flags on SuperLink
Rogue devices joining	Node authentication (public key allow-list)	--auth-list-public-keys / --auth-supernode-private-key
Server inspecting individual updates	Secure aggregation	SecAgg+ workflow
Reconstruction from the final model	Differential privacy (clip + noise)	DifferentialPrivacyClientSideFixedClipping wrappers
Model poisoning by a client	Robust aggregation	FedMedian, FedTrimmedAvg, Krum

• Training metrics - log aggregated accuracy, loss and max_train_time per round to MLflow or Weights & Biases from the ServerApp.
• System metrics - node_exporter on every Pi, and a jtop-based exporter on Jetson, scraped by Prometheus and shown in Grafana (CPU, GPU, temperature, power, network).
• Logs - journalctl -u flower-supernode -f on each node.

Ideas for Experiments

1. Compare FedAvg vs FedProx vs FedAdam at alpha = 0.1, 0.5 and 100.
2. Measure accuracy vs energy (Wh) instead of vs rounds.
3. Unplug a Pi mid-round and check that the run continues (min_fit_clients).
4. Replace CIFAR-10 with real sensor data: each Pi with a camera module or a temperature/IMU sensor collecting its own dataset.
5. Fine-tune only the last layers of a pretrained MobileNetV3 on the Pis while the Orin trains the full network (partial / split training).

Conclusion

A cluster of Raspberry Pis plus one Jetson Orin is a cheap and surprisingly realistic federated learning lab: real network latency, real hardware heterogeneity, real thermal limits and real dropouts - all the things a simulation hides. With Flower, the code you write for a laptop simulation runs unchanged on the cluster; the difference is only in the federation address and a systemd unit on each device. Start in simulation, move to two or three boards, add TLS and node authentication, and then use the testbed to study the real questions of FL: non-IID data, stragglers, communication cost and privacy.