Flame: Simplifying Topology Extension in Federated Learning

TL;DR

Flame introduces Topology Abstraction Graphs (TAGs) to decouple federated learning topology design from deployment infrastructure, enabling flexible, extensible, and per-channel communication backends. Its management plane, resource annotations, and dual programming models (user and developer) support rapid composition and extension of FL topologies, including Coordinated FL (CO-FL) and Hybrid FL, with TAG expansion mapping abstract roles to physical workers. Evaluations show TAG expansion overhead is low and scalable, and Flame enables easier topology transformations with modest code changes compared to FedML. The work demonstrates practical benefits in efficiency and resilience across diverse deployment contexts, and provides open-source tooling (Flame and Flame-In-A-Box) to accelerate FL methodology development.

Abstract

Distributed machine learning approaches, including a broad class of federated learning (FL) techniques, present a number of benefits when deploying machine learning applications over widely distributed infrastructures. The benefits are highly dependent on the details of the underlying machine learning topology, which specifies the functionality executed by the participating nodes, their dependencies and interconnections. Current systems lack the flexibility and extensibility necessary to customize the topology of a machine learning deployment. We present Flame, a new system that provides flexibility of the topology configuration of distributed FL applications around the specifics of a particular deployment context, and is easily extensible to support new FL architectures. Flame achieves this via a new high-level abstraction Topology Abstraction Graphs (TAGs). TAGs decouple the ML application logic from the underlying deployment details, making it possible to specialize the application deployment with reduced development effort. Flame is released as an open source project, and its flexibility and extensibility support a variety of topologies and mechanisms, and can facilitate the development of new FL methodologies.

Figures (16)

• Figure 1: Topologies that can be used in federated learning: (a) distributed, (b) classical FL, (c) hierarchical, (d) hierarchical with replicas and coordinator, and (e) hybrid. : training node, : aggregation node, : global aggregation node, and : coordinator. In (d), dotted lines denote connections between coordinator and other components.
• Figure 2: Federated learning experiments where each topologies have different time-to-accuracy performance and communication cost: Exp1) straggler node, Exp2) client node failures, and Exp3) aggregator & leader node failures.
• Figure 3: Conceptual overview of the Flame system.
• Figure 4: (a) Building blocks of TAG. TAG representation of topologies: (b) distributed, (c) classical FL, (d) hierarchical FL, (e) hybrid. is a trainer node with isDataConsumer set. The groupBy attribute is used to create groups of the same role.
• Figure 5: Expansion of a TAG into a physical topology. (a) TAG representation of hierarchical FL (H-FL); each dataset belongs to a group; (b) dataset is expanded; (c) one training worker is allocated per dataset and the worker's group is determined by the dataset's group; and (d) each entry in groupAssociation corresponds to one aggregation worker; hence, two aggregators are created; each belongs to a different group for param-channel while both have one "default" group for agg-channel. Since there is one entry for groupAssociation in the global aggregation role, expansion finishes by creating one global aggregation worker. Note that the expansion of roles can be done in an arbitrary order since groupAssociation has all the necessary information for expansion. A channel's attributes (e.g., groupBy) are used to validate the expansion.