TLDR: OmniFed is a modular, open-source framework built on PyTorch that simplifies Federated Learning (FL) and Collaborative Learning (CL) deployments from edge devices to High-Performance Computing (HPC) systems. It features a configurable architecture that decouples concerns for topology, communication, and training logic, supporting various FL algorithms, communication compression, real-time streaming data, and privacy mechanisms like Differential Privacy, Homomorphic Encryption, and Secure Aggregation. Its design allows for rapid prototyping and customization, enabling researchers to easily experiment with different FL strategies and topologies.
In today’s data-driven world, where information is increasingly distributed, sensitive, and vast, traditional centralized Artificial Intelligence (AI) methods are becoming less practical. This is where Federated Learning (FL) and Collaborative Learning (CL) step in, offering a way to train powerful Deep Neural Network (DNN) models using private and sensitive data located at various endpoints without moving the data itself. This approach is crucial for fields like scientific discovery, healthcare, and finance, where data privacy and locality are paramount.
Researchers at Oak Ridge National Laboratory have introduced a new, open-source framework called OmniFed. Built on PyTorch, OmniFed aims to simplify and streamline the deployment of FL/CL across a wide range of environments, from edge devices to high-performance computing (HPC) systems. Its core strength lies in its modular and extensible design, which allows users to easily configure and customize various aspects of federated learning without getting bogged down in complex infrastructure details.
A Flexible and Configurable Architecture
OmniFed’s design philosophy emphasizes modularity, flexibility, and extensibility. It achieves this by clearly separating concerns for configuration, orchestration, communication, and training logic. This means that components like local computation, communication protocols, and algorithmic control strategies can be easily swapped out. The framework uses YAML-based configuration management with Hydra, enabling a plug-and-play approach where users can add or remove privacy features and compression techniques with minimal code changes.
One of the key innovations of OmniFed is its support for diverse training topologies. While many existing frameworks assume a centralized setup, OmniFed allows for the implementation of complex, custom topologies such as centralized, decentralized, ring, peer-to-peer, and hierarchical structures. This flexibility is vital for real-world scenarios where federated nodes might be connected in various ways, each with its own trade-offs in terms of communication and coordination.
Core Components for Seamless Operation
The framework is composed of several key modules that work together to manage FL/CL experiments:
- Engine: The central orchestrator, responsible for launching and coordinating experiments, managing node lifecycles, allocating resources, and collecting metrics.
- Topology: Defines the structure and coordination patterns for distributed nodes, supporting templates for common setups and allowing for custom graph-based representations.
- Node: Represents a participant in the federation, managing local model state, data, and resources. Nodes can act as trainers (performing local training) or aggregators (collecting and merging model updates).
- Communicator: Provides a unified API for data exchange between nodes, abstracting away the underlying communication protocols like gRPC, MPI, or MQTT.
- Algorithm: Encapsulates the federated learning logic through configurable lifecycle hooks, allowing users to choose from a suite of built-in algorithms or implement their own.
Also Read:
- Enhancing Edge Federated Learning Through Client Overlap and Relays
- Optimizing Large Language Model Training in Mobile Edge Networks with CollaPipe
Rich Feature Set for Practical FL
OmniFed comes equipped with a comprehensive set of features designed to address practical challenges in federated learning:
- Diverse FL Algorithms: It supports over 10 popular FL algorithms out-of-the-box, including FedAvg, FedProx, FedMom, and Scaffold. Users can easily switch between algorithms and define hyperparameters through the configuration file, making it simple to compare their performance.
- Communication Compression: To reduce communication overhead, OmniFed integrates various gradient compression techniques like TopK, Deep Gradient Compression (DGC), QSGD quantization, and PowerSGD low-rank approximation. These can be configured within the communication module to optimize data exchange.
- Real-time Learning with Streaming Data: The framework can simulate real-time learning scenarios by integrating with Apache Kafka to handle streaming datasets. This allows FL clients to subscribe to topics and collect data as it arrives, enabling adaptive, real-time modeling.
- Privacy-Preserving Mechanisms: Privacy is a first-class citizen in OmniFed, offering optional mechanisms such as Differential Privacy (DP), Homomorphic Encryption (HE), and Secure Aggregation (SA). Users can specify privacy budgets and compare the computational overhead of different techniques to meet their specific security requirements.
- Mixed-Protocol Communication: OmniFed supports complex, cross-facility training scenarios where different communication protocols can be used within a single deployment. For example, high-bandwidth MPI collectives can be used for intra-site aggregation, while gRPC handles inter-site communication over slower networks, with optional compression applied only to the slower links.
By unifying topology configuration, mixed-protocol communication, and pluggable modules, OmniFed significantly streamlines FL deployment in heterogeneous environments. It empowers researchers to focus on algorithmic design and innovation rather than infrastructure complexities, paving the way for the next generation of secure and distributed AI. For more details, you can refer to the full research paper: OmniFed: A Modular Framework for Configurable Federated Learning from Edge to HPC.
