OmniFed: Configurable Federated Learning

• OmniFed is a federated learning framework that decouples configuration, orchestration, communication, and training to support heterogeneous, privacy-sensitive deployments.
• It enables arbitrary topology specification through a single YAML file and supports mixed protocols like gRPC, MPI, and MQTT along with plug-in privacy and compression methods.
• The architecture leverages HPC and edge device simulations with tools like Ray and Kafka to achieve near-linear scaling and efficient communication.

OmniFed is a federated learning (FL) framework designed for heterogeneous, privacy-sensitive deployments spanning edge devices and High Performance Computing (HPC) environments. It is presented as a modular system organized around decoupling configuration, orchestration, communication, and training logic, with support for configuration-driven prototyping, code-level override-what-you-need customization, arbitrary topology specification through a single YAML file, mixed communication protocols within a single deployment, and optional privacy and compression mechanisms including Differential Privacy (DP), Homomorphic Encryption (HE), Secure Aggregation (SA), quantization, and sparsification (Tyagi et al., 23 Sep 2025). The framework was published on 2025-09-23, and its repository is available at https://github.com/at-aaims/OmniFed.

1. Motivation and design target

OmniFed is motivated by the practical failure modes of monolithic FL systems as machine learning shifts from cloud-centralized settings to distributed environments such as edge devices, national-lab HPC, and multi-institution collaborations (Tyagi et al., 23 Sep 2025). The stated limitations of existing toolkits are rigid topologies, tight coupling of training logic with communication and orchestration, fragmented privacy and compression support, and a lack of real-time or streaming simulation combined with heterogeneous-resource orchestration. In this framing, most frameworks assume a star client-server pattern and offer poor support for peer-to-peer, ring, hierarchical, or mixed hub-and-spoke deployments; switching from FedAvg to FedProx, or from gRPC to MPI, often requires rewriting large swaths of boilerplate; and DP, HE, SA, and gradient-compression schemes tend to reside in separate libraries without unified extension points.

The framework is explicitly described as addressing those gaps by decoupling configuration, communication, orchestration, and training logic; supporting arbitrary topologies via a single YAML file; enabling mixed-protocol communication with gRPC, MPI, AMQP, HTTP, and MQTT in one run; exposing plug-in hooks for DP, HE, SA, and for quantization and sparsification compressors; and simulating streaming clients such as Kafka alongside HPC actor-based orchestration via Ray (Tyagi et al., 23 Sep 2025). This suggests that OmniFed is positioned not merely as an implementation of a specific FL algorithm, but as a systems substrate for configurable experimentation across deployment regimes that are usually treated separately.

2. Federated optimization model and algorithm life cycle

At its core, OmniFed targets the canonical federated optimization problem

$$\min_{w\in\mathbb{R}^d} \; F(w) \quad\text{where}\quad F(w)\;=\;\sum_{k=1}^K p_k\,F_k(w), \quad F_k(w)\;=\;\mathbb{E}_{\xi\sim\mathcal{D}_k}\big[L(w;\,\xi)\big].$$

Here, $K$ is the number of participants, $p_k \ge 0$, $\sum_k p_k = 1$ are aggregation weights, and a typical choice is $p_k=\frac{n_k}{\sum_j n_j}$ (Tyagi et al., 23 Sep 2025). Local optimization is described in terms of local SGD,

$w_k^{t+1} \;=\; w^t \;-\;\eta\,\nabla F_k(w^t),$

followed by submission of either a model $w_k$ or a gradient $g_k$, and then aggregation. For FedAvg, the recomputation rule is

$w^{t+1} \;=\;\sum_{k=1}^K p_k\,w_k^{t+1}.$

The framework’s Algorithm plug-in API is aligned to this training life cycle through hooks such as local_train(), aggregate(), and post_aggregate() (Tyagi et al., 23 Sep 2025). The paper describes the Algorithm layer as a pure-Python life-cycle class that supports FedAvg, FedProx, FedDyn, Moon, Ditto, DiLoCo, and 8+ other methods. Because the orchestration and communication layers are separated from the algorithm implementation, a change in optimization method is intended to affect only the relevant life-cycle stages rather than the surrounding system. A plausible implication is that OmniFed treats FL algorithms as interchangeable policies embedded in a stable execution substrate.

3. Layered architecture and separation of concerns

The architecture is presented as a layered system with five principal components: Configuration, Engine, Topology, Node, and Communicator, with Algorithm as the customizable training abstraction (Tyagi et al., 23 Sep 2025). Configuration is handled through a single Hydra-managed YAML that specifies topology, algorithm, communicator, model, data, global rounds, and hyperparameters. The stated operational model is that users override only what they need; for example, switching from FedAvg to FedProx is described as a one-line change.

The Engine is a Ray-based driver process that reads the configuration, instantiates the chosen Topology, spawns Node actors, advances global rounds, collects metrics, and gracefully terminates. The Topology component encodes the graph of participants, including centralized star, ring, peer-to-peer, hierarchical tree, and custom arrangements, and instructs the Engine how Nodes should be wired together. Each Node holds a local model, data, and an Algorithm instance; invokes local training; encodes or compresses updates; and sends them through its Communicator. The Communicator exposes a uniform API for send/receive, all-reduce, and gather-and-scatter, backed by MPI via PyTorch Distributed, gRPC with protocol buffers, AMQP or MQTT, or other protocols (Tyagi et al., 23 Sep 2025).

This separation of concerns is central to the framework’s design. The paper emphasizes that custom topologies simply implement the Topology base class and register via Hydra’s plugin mechanism, with no changes to Engine or Node, and that new compressors can be attached through communicator configuration. The significance of this decomposition lies in its attempt to isolate systems concerns—deployment graph, transport substrate, orchestration, and training logic—so that each can vary independently without destabilizing the core runtime.

4. Topology specification and mixed-protocol communication

OmniFed ships with templates for four topology classes: centralized (star), decentralized ring, peer-to-peer (full mesh), and hierarchical tree (multi-level), while also allowing any arbitrary graph to be described in YAML (Tyagi et al., 23 Sep 2025). The distinctive claim is that each edge in that graph may choose its own Communicator. The paper gives a cross-facility hub-and-spoke example in which inner spokes use MPI all-reduce over high-bandwidth datacenter links, outer aggregation uses gRPC across geographically distributed higher-latency links, and gradient compression such as TopK is applied only on the slow outer link.

An excerpted mixed-protocol hierarchical configuration instantiates src.omnifed.topology.HierarchicalTopology, assigns src.omnifed.communicator.MpiCommunicator with backend: nccl to named sites such as us_east and us_west, designates a src.omnifed.communicator.GrpcCommunicator as the global_aggregator, and attaches src.omnifed.communicator.compression.TopK with k: 10000 (Tyagi et al., 23 Sep 2025). This is not merely a transport abstraction; it is a graph-level communication model in which protocol selection can vary across links inside a single deployment.

The framework also supports streaming-oriented and heterogeneous deployments. The motivation section specifically cites simulation of streaming clients such as Kafka alongside Ray-based orchestration, and the evaluation reports that Kafka-based clients sustain target rates up to 32 samples/s per stream across 16 clients (Tyagi et al., 23 Sep 2025). This suggests an attempt to collapse batch FL, streaming ingestion, and distributed systems experimentation into one configurable stack.

5. Privacy mechanisms and compression pipeline

Privacy is exposed as optional middleware in the Algorithm pipeline. OmniFed integrates Differential Privacy via the Gaussian mechanism, Homomorphic Encryption via TenSEAL/CKKS, and Secure Aggregation via pairwise HMAC-based masking (Tyagi et al., 23 Sep 2025). For DP, gradients are clipped to norm $C$ and perturbed:

$K$0

with

$K$1

The paper states that the resulting $K$2-DP guarantee composes over rounds. For HE, each client encrypts its clipped gradient vector to $K$3 and the aggregator performs ciphertext-space addition,

$K$4

For SA, client $K$5 submits

$K$6

so that masks cancel under summation and reveal only $K$7.

Compression is implemented as a Communicator-level plug-in facility. The built-in methods are TopK sparsification, Deep Gradient Compression (DGC), RandomK, SIDCo, RedSync, QSGD quantization, and PowerSGD low-rank compression (Tyagi et al., 23 Sep 2025). TopK keeps the $K$8 largest-magnitude components of gradient $K$9 and sets the rest to zero, with error-feedback accumulating dropped residuals for unbiasedness. The implementation steps are stated as: clip or scale gradient $p_k \ge 0$0; select indices $p_k \ge 0$1; form sparse vector $p_k \ge 0$2 with $p_k \ge 0$3 if $p_k \ge 0$4 and $p_k \ge 0$5 otherwise; and accumulate residual $p_k \ge 0$6 for the next round. QSGD is defined as

$p_k \ge 0$7

where $p_k \ge 0$8 is a randomized rounding bit, and the paper gives the mean-squared error bound

$p_k \ge 0$9

PowerSGD is described as decomposing a gradient matrix $\sum_k p_k = 1$0 as $\sum_k p_k = 1$1 with rank $\sum_k p_k = 1$2 and transmitting $\sum_k p_k = 1$3 and $\sum_k p_k = 1$4 (Tyagi et al., 23 Sep 2025).

The architectural significance is that privacy and compression are not external wrappers but first-class extension points in the same execution pipeline. This suggests a deliberate unification of statistical optimization, privacy middleware, and communication efficiency within one framework interface.

6. Experimental evaluation, scalability characteristics, and stated limitations

The reported experiments run on an NVIDIA DGX with 8 × H100 GPUs and 16 logical clients (Tyagi et al., 23 Sep 2025). The evaluated models and datasets are ResNet-18 / CIFAR-10, VGG-11 / CIFAR-100, AlexNet / CalTech-101, and MobileNetV3 / CalTech-256. The compared FL algorithms are FedAvg, FedProx, FedMom, FedNova, Moon, FedDyn, FedBN, FedPer, Ditto, and DiLoCo. Compression settings include TopK at 10× and 1000×, DGC, QSGD at 8-bit and 16-bit, and PowerSGD at rank 32 and 64. Privacy settings include DP with $\sum_k p_k = 1$5, HE with CKKS, and SA with HMAC masks. The reported metrics are test accuracy, epoch wall-clock time, communication payload, and privacy-compute overhead.

The convergence results reported in Table 1 are: ResNet-18, FedAvg 99.32% and Moon best 99.46%; VGG-11, FedAvg 86.6% and FedProx 86.3%; AlexNet, FedDyn 88.8% and FedAvg 87.9%; and MobileNetV3, FedProx 82.96% (Tyagi et al., 23 Sep 2025). Figure 1 reports that Moon and FedAvg are fastest per epoch, while FedNova and FedMom incur extra prox/momentum overhead. Compression results in Table 2 and Figure 2 state that TopK 10× reduces payload ≈10× with accuracy drop <0.3%; TopK 1000× loses up to 1.5% accuracy on AlexNet; QSGD 8-bit/16-bit trades modest accuracy (<1%) for 2×–4× reduction in byte-transfer cost; and PowerSGD rank-32 can collapse accuracy on VGG-11 to 6.7% unless higher rank is used. Privacy results in Table 3 state that DP($\sum_k p_k = 1$6) yields ResNet-18 at 97.98% versus 99.32% without noise, and $\sum_k p_k = 1$7 recovers 98.06%; HE and SA incur 10×–100× compute overhead compared to DP, with VGG-11 SA at approximately 2000 s per round (Tyagi et al., 23 Sep 2025).

The systems evaluation further reports that Kafka-based clients sustain target rates up to 32 samples/s per stream across 16 clients; that in a cross-facility setting inner MPI All-Reduce is approximately 3× faster than global gRPC, while outer gRPC incurs 2× latency but can be compressed selectively; that HPC runs exploit Ray’s fine-grained actor scheduling and NCCL-based MPI to achieve near-linear strong scaling up to 64 GPUs in early SLURM tests; that edge settings on Raspberry Pi and Jetson Nano run TensorFlow Lite or ExecuTorch with smaller batch sizes and per-round wall-clock times dominated by network latency through gRPC or MQTT; and that large-scale stress tests up to 256 clients in simulation show less than 5% degradation in orchestration throughput due to Hydra’s lazy-init and Ray’s object-store pipelining (Tyagi et al., 23 Sep 2025).

The paper also states several limitations and future directions. Peer-to-peering protocol, native without MPI, is under development; MQTT integration is prototyped but requires validation on lossy wireless links; asynchronous FL algorithms such as FedAsync and asynchronous gossip are not yet integrated; and containerization with Docker/Kubernetes and SLURM deployment scripts are forthcoming (Tyagi et al., 23 Sep 2025). The near-term roadmap is to finish peer-to-peer and MQTT Communicators, containerize the framework, extend to asynchronous training and heterogeneity-aware scheduling, integrate with large-model backends such as DeepSpeed and TorchTitan and edge runtimes such as ExecuTorch, and pursue full SLURM/Flux integration for exascale. These stated limitations are important because they delimit the scope of the current implementation: OmniFed is presented as a unified configurable FL stack, but not yet as a complete solution for asynchronous FL or fully productionized multi-environment deployment.