Added pFedGraph to Plato.

Author: baochunli

docs/docs/configurations/algorithm.md

@@ -6,6 +6,7 @@
     - `fedavg` the federated averaging algorithm
     - `split_learning` the Split Learning algorithm
     - `fedavg_personalized` the personalized federated learning algorithm
+    - `pfedgraph` the Personalized Federated Learning with Inferred Collaboration Graphs algorithm
 
 !!! example "cross_silo"
     Whether or not cross-silo training should be used.
@@ -26,3 +27,27 @@
         A float to show the proportion of clients participating in the federated training process. It is under `personalization`, which is a sub-config path that contains other personalized training parameters.
 
         Default value: `1.0`
+
+!!! example "pfedgraph"
+    Configuration for pFedGraph.
+
+    !!! example "pfedgraph_alpha"
+        Hyper-parameter controlling the collaboration graph update.
+
+        Default value: `0.8`
+
+    !!! example "pfedgraph_lambda"
+        Regularization strength for cosine similarity in the local objective.
+
+        Default value: `0.01`
+
+    !!! example "pfedgraph_similarity_metric"
+        Similarity metric scope for graph inference. Use `all` for all parameters
+        or `fc` to focus on the final fully-connected layers.
+
+        Default value: `all`
+
+    !!! example "pfedgraph_similarity_layers"
+        Optional list of layer name substrings to use when computing model
+        similarity for graph inference. Overrides `pfedgraph_similarity_metric`
+        when provided.

docs/docs/configurations/server.md

@@ -10,6 +10,7 @@
     - `fedavg_mpc_shamir` a Federated Averaging server that reconstructs Shamir MPC shares before aggregation. Requires clients of type `mpc` with the `mpc_model_encrypt_shamir` processor.
     - `split_learning` a Split Learning server that supports training different kinds of models in split learning framework. When this server is used, the `clients.per_round` in the configuration should be set to 1. Users should define the rules for updating models weights before cut from the clients to the server in the callback function `on_update_weights_before_cut`, depending on the specific model they use.
     - `fedavg_personalized` a personalized federated learning server that starts from a number of regular rounds of federated learning. In these regular rounds, only a subset of the total clients can be selected to perform the local update (the ratio of which is a configuration setting). After all regular rounds are completed, it starts a final round of personalization, where a selected subset of clients perform local training using their local dataset.
+    - `pfedgraph` a personalized federated learning server that aggregates models using an inferred collaboration graph and sends per-client aggregated weights.
 
 !!! example "address"
     The address of the central server, such as `127.0.0.1`.

plato/algorithms/pfedgraph.py

@@ -0,0 +1,12 @@
+"""
+pFedGraph algorithm wrapper.
+
+pFedGraph uses standard FedAvg weight handling with a personalized aggregation
+strategy on the server side.
+"""
+
+from plato.algorithms import fedavg
+
+
+class Algorithm(fedavg.Algorithm):
+    """pFedGraph reuses the FedAvg algorithm primitives."""

plato/algorithms/registry.py

@@ -16,6 +16,7 @@
     fedavg_personalized,
     lora,
     mlx_fedavg,
+    pfedgraph,
     split_learning,
 )
 from plato.algorithms.base import Algorithm as AlgorithmBase
@@ -27,6 +28,7 @@
     "fedavg_personalized": fedavg_personalized.Algorithm,
     "fedavg_lora": lora.Algorithm,
     "mlx_fedavg": mlx_fedavg.Algorithm,
+    "pfedgraph": pfedgraph.Algorithm,
     "split_learning": split_learning.Algorithm,
 }
 

plato/clients/pfedgraph.py

@@ -0,0 +1,11 @@
+"""
+pFedGraph client implementation.
+
+Uses the standard simple client pipeline.
+"""
+
+from plato.clients import simple
+
+
+class Client(simple.Client):
+    """A pFedGraph client using default composable strategies."""

plato/clients/registry.py

@@ -19,6 +19,7 @@
     edge,
     fedavg_personalized,
     mpc,
+    pfedgraph,
     self_supervised_learning,
     simple,
     split_learning,
@@ -107,6 +108,7 @@ def factory(**kwargs) -> Client:
     "simple": _simple_like_factory(simple.Client),
     "fedavg_personalized": _simple_like_factory(fedavg_personalized.Client),
     "mpc": _simple_like_factory(mpc.Client),
+    "pfedgraph": _simple_like_factory(pfedgraph.Client),
     "self_supervised_learning": _simple_like_factory(self_supervised_learning.Client),
     "split_learning": _simple_like_factory(split_learning.Client),
     "edge": _edge_factory(),

plato/servers/pfedgraph.py

@@ -0,0 +1,81 @@
+"""
+pFedGraph server implementation.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+from plato.config import Config
+from plato.servers import fedavg
+from plato.servers.strategies.aggregation.pfedgraph import (
+    PFedGraphAggregationStrategy,
+)
+
+
+class Server(fedavg.Server):
+    """Federated learning server implementing pFedGraph."""
+
+    def __init__(
+        self,
+        model=None,
+        datasource=None,
+        algorithm=None,
+        trainer=None,
+        callbacks=None,
+        aggregation_strategy=None,
+        client_selection_strategy=None,
+    ):
+        if aggregation_strategy is None:
+            similarity_layers = None
+            similarity_metric = "all"
+            alpha = 0.8
+
+            if hasattr(Config(), "algorithm"):
+                if hasattr(Config().algorithm, "pfedgraph_similarity_metric"):
+                    similarity_metric = Config().algorithm.pfedgraph_similarity_metric
+                elif hasattr(Config().algorithm, "pfedgraph_similarity"):
+                    similarity_metric = Config().algorithm.pfedgraph_similarity
+
+                if hasattr(Config().algorithm, "pfedgraph_similarity_layers"):
+                    similarity_layers = Config().algorithm.pfedgraph_similarity_layers
+
+                if hasattr(Config().algorithm, "pfedgraph_alpha"):
+                    alpha = Config().algorithm.pfedgraph_alpha
+
+            aggregation_strategy = PFedGraphAggregationStrategy(
+                alpha=alpha,
+                similarity_metric=similarity_metric,
+                similarity_layers=similarity_layers,
+            )
+
+        super().__init__(
+            model=model,
+            datasource=datasource,
+            algorithm=algorithm,
+            trainer=trainer,
+            callbacks=callbacks,
+            aggregation_strategy=aggregation_strategy,
+            client_selection_strategy=client_selection_strategy,
+        )
+
+        self.client_models: dict[int, dict[str, Any]] = {}
+
+    def update_client_model(
+        self,
+        aggregated_clients_models: Sequence[dict[str, Any]],
+        updates: Sequence[Any],
+    ) -> None:
+        """Update the stored model for each client."""
+        for client_model, update in zip(aggregated_clients_models, updates):
+            client_id = getattr(update, "client_id", None)
+            if client_id is None:
+                continue
+            self.client_models[client_id] = client_model
+
+    def customize_server_payload(self, payload: Any) -> Any:
+        """Send per-client aggregated weights when available."""
+        client_id = self.selected_client_id
+        if client_id in self.client_models:
+            return self.client_models[client_id]
+        return payload

plato/servers/registry.py

@@ -16,6 +16,7 @@
     fedavg_mpc_additive,
     fedavg_mpc_shamir,
     fedavg_personalized,
+    pfedgraph,
     split_learning,
 )
 
@@ -34,6 +35,7 @@
         "fedavg_personalized": fedavg_personalized.Server,
         "fedavg_mpc_additive": fedavg_mpc_additive.Server,
         "fedavg_mpc_shamir": fedavg_mpc_shamir.Server,
+        "pfedgraph": pfedgraph.Server,
         "split_learning": split_learning.Server,
     }
 

plato/servers/strategies/aggregation/__init__.py

@@ -11,6 +11,7 @@
 from plato.servers.strategies.aggregation.gan import FedAvgGanAggregationStrategy
 from plato.servers.strategies.aggregation.he import FedAvgHEAggregationStrategy
 from plato.servers.strategies.aggregation.hermes import HermesAggregationStrategy
+from plato.servers.strategies.aggregation.pfedgraph import PFedGraphAggregationStrategy
 from plato.servers.strategies.aggregation.port import PortAggregationStrategy
 
 __all__ = [
@@ -20,6 +21,7 @@
     "FedAsyncAggregationStrategy",
     "PortAggregationStrategy",
     "HermesAggregationStrategy",
+    "PFedGraphAggregationStrategy",
     "FedAvgGanAggregationStrategy",
     "FedAvgHEAggregationStrategy",
 ]