flwr 1.21.0__py3-none-any.whl → 1.22.0__py3-none-any.whl

flwr/cli/new/templates/app/code/flwr_tune/client_app.py.tpl

@@ -2,15 +2,12 @@
 
 import os
 import warnings
-from typing import Dict, Tuple
 
-import torch
-from flwr.client import ClientApp, NumPyClient
-from flwr.common import Context
+from flwr.app import ArrayRecord, Context, Message, MetricRecord, RecordDict
+from flwr.clientapp import ClientApp
 from flwr.common.config import unflatten_dict
-from flwr.common.typing import NDArrays, Scalar
 from omegaconf import DictConfig
-
+from peft import get_peft_model_state_dict, set_peft_model_state_dict
 from transformers import TrainingArguments
 from trl import SFTTrainer
 
@@ -19,12 +16,7 @@ from $import_name.dataset import (
     load_data,
     replace_keys,
 )
-from $import_name.models import (
-    cosine_annealing,
-    get_model,
-    set_parameters,
-    get_parameters,
-)
+from $import_name.models import cosine_annealing, get_model
 
 # Avoid warnings
 os.environ["TOKENIZERS_PARALLELISM"] = "true"
@@ -32,95 +24,69 @@ os.environ["RAY_DISABLE_DOCKER_CPU_WARNING"] = "1"
 warnings.filterwarnings("ignore", category=UserWarning)
 
 
-# pylint: disable=too-many-arguments
-# pylint: disable=too-many-instance-attributes
-class FlowerClient(NumPyClient):
-    """Flower client for LLM fine-tuning."""
+# Avoid warnings
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+os.environ["RAY_DISABLE_DOCKER_CPU_WARNING"] = "1"
+warnings.filterwarnings("ignore", category=UserWarning)
 
-    def __init__(
-        self,
-        model_cfg: DictConfig,
-        train_cfg: DictConfig,
-        trainset,
-        tokenizer,
-        formatting_prompts_func,
-        data_collator,
-        num_rounds,
-    ):  # pylint: disable=too-many-arguments
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        self.train_cfg = train_cfg
-        self.training_arguments = TrainingArguments(**train_cfg.training_arguments)
-        self.tokenizer = tokenizer
-        self.formatting_prompts_func = formatting_prompts_func
-        self.data_collator = data_collator
-        self.num_rounds = num_rounds
-        self.trainset = trainset
-
-        # instantiate model
-        self.model = get_model(model_cfg)
-
-    def fit(
-        self, parameters: NDArrays, config: Dict[str, Scalar]
-    ) -> Tuple[NDArrays, int, Dict]:
-        """Implement distributed fit function for a given client."""
-        set_parameters(self.model, parameters)
-
-        new_lr = cosine_annealing(
-            int(config["current_round"]),
-            self.num_rounds,
-            self.train_cfg.learning_rate_max,
-            self.train_cfg.learning_rate_min,
-        )
-
-        self.training_arguments.learning_rate = new_lr
-        self.training_arguments.output_dir = config["save_path"]
-
-        # Construct trainer
-        trainer = SFTTrainer(
-            model=self.model,
-            tokenizer=self.tokenizer,
-            args=self.training_arguments,
-            max_seq_length=self.train_cfg.seq_length,
-            train_dataset=self.trainset,
-            formatting_func=self.formatting_prompts_func,
-            data_collator=self.data_collator,
-        )
-
-        # Do local training
-        results = trainer.train()
-
-        return (
-            get_parameters(self.model),
-            len(self.trainset),
-            {"train_loss": results.training_loss},
-        )
-
-
-def client_fn(context: Context) -> FlowerClient:
-    """Create a Flower client representing a single organization."""
+
+# Flower ClientApp
+app = ClientApp()
+
+
+@app.train()
+def train(msg: Message, context: Context):
+    """Train the model on local data."""
+    # Parse config
     partition_id = context.node_config["partition-id"]
     num_partitions = context.node_config["num-partitions"]
     num_rounds = context.run_config["num-server-rounds"]
     cfg = DictConfig(replace_keys(unflatten_dict(context.run_config)))
+    training_arguments = TrainingArguments(**cfg.train.training_arguments)
 
     # Let's get the client partition
-    client_trainset = load_data(partition_id, num_partitions, cfg.static.dataset.name)
+    trainset = load_data(partition_id, num_partitions, cfg.static.dataset.name)
     (
         tokenizer,
         data_collator,
         formatting_prompts_func,
     ) = get_tokenizer_and_data_collator_and_propt_formatting(cfg.model.name)
 
-    return FlowerClient(
-        cfg.model,
-        cfg.train,
-        client_trainset,
-        tokenizer,
-        formatting_prompts_func,
-        data_collator,
-        num_rounds,
-    ).to_client()
-
+    # Load the model and initialize it with the received weights
+    model = get_model(cfg.model)
+    set_peft_model_state_dict(model, msg.content["arrays"].to_torch_state_dict())
 
-# Flower ClientApp
-app = ClientApp(client_fn)
+    # Set learning rate for current round
+    new_lr = cosine_annealing(
+        msg.content["config"]["server-round"],
+        num_rounds,
+        cfg.train.learning_rate_max,
+        cfg.train.learning_rate_min,
+    )
+
+    training_arguments.learning_rate = new_lr
+    training_arguments.output_dir = msg.content["config"]["save_path"]
+
+    # Construct trainer
+    trainer = SFTTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        args=training_arguments,
+        max_seq_length=cfg.train.seq_length,
+        train_dataset=trainset,
+        formatting_func=formatting_prompts_func,
+        data_collator=data_collator,
+    )
+
+    # Do local training
+    results = trainer.train()
+
+    # Construct and return reply Message
+    model_record = ArrayRecord(get_peft_model_state_dict(model))
+    metrics = {
+        "train_loss": results.training_loss,
+        "num-examples": len(trainset),
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"arrays": model_record, "metrics": metric_record})
+    return Message(content=content, reply_to=msg)

flwr/cli/new/templates/app/code/flwr_tune/models.py.tpl

@@ -4,18 +4,10 @@ import math
 
 import torch
 from omegaconf import DictConfig
-from collections import OrderedDict
-from peft import (
-    LoraConfig,
-    get_peft_model,
-    get_peft_model_state_dict,
-    set_peft_model_state_dict,
-)
+from peft import LoraConfig, get_peft_model
 from peft.utils import prepare_model_for_kbit_training
 from transformers import AutoModelForCausalLM, BitsAndBytesConfig
 
-from flwr.common.typing import NDArrays
-
 
 def cosine_annealing(
     current_round: int,
@@ -62,17 +54,3 @@ def get_model(model_cfg: DictConfig):
         model.config.use_cache = False
 
     return get_peft_model(model, peft_config)
-
-
-def set_parameters(model, parameters: NDArrays) -> None:
-    """Change the parameters of the model using the given ones."""
-    peft_state_dict_keys = get_peft_model_state_dict(model).keys()
-    params_dict = zip(peft_state_dict_keys, parameters)
-    state_dict = OrderedDict({k: torch.Tensor(v) for k, v in params_dict})
-    set_peft_model_state_dict(model, state_dict)
-
-
-def get_parameters(model) -> NDArrays:
-    """Return the parameters of the current net."""
-    state_dict = get_peft_model_state_dict(model)
-    return [val.cpu().numpy() for _, val in state_dict.items()]

flwr/cli/new/templates/app/code/flwr_tune/server_app.py.tpl

@@ -3,62 +3,23 @@
 import os
 from datetime import datetime
 
-from flwr.common import Context, ndarrays_to_parameters
+from flwr.app import ArrayRecord, ConfigRecord, Context, MetricRecord
 from flwr.common.config import unflatten_dict
-from flwr.server import ServerApp, ServerAppComponents, ServerConfig
+from flwr.serverapp import Grid, ServerApp
 from omegaconf import DictConfig
+from peft import get_peft_model_state_dict, set_peft_model_state_dict
 
-from $import_name.models import get_model, get_parameters, set_parameters
 from $import_name.dataset import replace_keys
+from $import_name.models import get_model
 from $import_name.strategy import FlowerTuneLlm
 
-
-# Get function that will be executed by the strategy's evaluate() method
-# Here we use it to save global model checkpoints
-def get_evaluate_fn(model_cfg, save_every_round, total_round, save_path):
-    """Return an evaluation function for saving global model."""
-
-    def evaluate(server_round: int, parameters, config):
-        # Save model
-        if server_round != 0 and (
-            server_round == total_round or server_round % save_every_round == 0
-        ):
-            # Init model
-            model = get_model(model_cfg)
-            set_parameters(model, parameters)
-
-            model.save_pretrained(f"{save_path}/peft_{server_round}")
-
-        return 0.0, {}
-
-    return evaluate
-
-
-def get_on_fit_config(save_path):
-    """Return a function that will be used to construct the config that the
-    client's fit() method will receive."""
-
-    def fit_config_fn(server_round: int):
-        fit_config = {}
-        fit_config["current_round"] = server_round
-        fit_config["save_path"] = save_path
-        return fit_config
-
-    return fit_config_fn
+# Create ServerApp
+app = ServerApp()
 
 
-def fit_weighted_average(metrics):
-    """Aggregate (federated) evaluation metrics."""
-    # Multiply accuracy of each client by number of examples used
-    losses = [num_examples * m["train_loss"] for num_examples, m in metrics]
-    examples = [num_examples for num_examples, _ in metrics]
-
-    # Aggregate and return custom metric (weighted average)
-    return {"train_loss": sum(losses) / sum(examples)}
-
-
-def server_fn(context: Context):
-    """Construct components that set the ServerApp behaviour."""
+@app.main()
+def main(grid: Grid, context: Context) -> None:
+    """Main entry point for the ServerApp."""
     # Create output directory given current timestamp
     current_time = datetime.now()
     folder_name = current_time.strftime("%Y-%m-%d_%H-%M-%S")
@@ -71,24 +32,42 @@ def server_fn(context: Context):
 
     # Get initial model weights
     init_model = get_model(cfg.model)
-    init_model_parameters = get_parameters(init_model)
-    init_model_parameters = ndarrays_to_parameters(init_model_parameters)
+    arrays = ArrayRecord(get_peft_model_state_dict(init_model))
 
     # Define strategy
     strategy = FlowerTuneLlm(
-        fraction_fit=cfg.strategy.fraction_fit,
+        fraction_train=cfg.strategy.fraction_train,
         fraction_evaluate=cfg.strategy.fraction_evaluate,
-        on_fit_config_fn=get_on_fit_config(save_path),
-        fit_metrics_aggregation_fn=fit_weighted_average,
-        initial_parameters=init_model_parameters,
+    )
+
+    # Start strategy, run FedAvg for `num_rounds`
+    strategy.start(
+        grid=grid,
+        initial_arrays=arrays,
+        train_config=ConfigRecord({"save_path": save_path}),
+        num_rounds=num_rounds,
         evaluate_fn=get_evaluate_fn(
             cfg.model, cfg.train.save_every_round, num_rounds, save_path
         ),
     )
-    config = ServerConfig(num_rounds=num_rounds)
 
-    return ServerAppComponents(strategy=strategy, config=config)
 
+# Get function that will be executed by the strategy
+# Here we use it to save global model checkpoints
+def get_evaluate_fn(model_cfg, save_every_round, total_round, save_path):
+    """Return an evaluation function for saving global model."""
 
-# Flower ServerApp
-app = ServerApp(server_fn=server_fn)
+    def evaluate(server_round: int, arrays: ArrayRecord) -> MetricRecord:
+        # Save model
+        if server_round != 0 and (
+            server_round == total_round or server_round % save_every_round == 0
+        ):
+            # Init model
+            model = get_model(model_cfg)
+            set_peft_model_state_dict(model, arrays.to_torch_state_dict())
+
+            model.save_pretrained(f"{save_path}/peft_{server_round}")
+
+        return MetricRecord()
+
+    return evaluate

flwr/cli/new/templates/app/code/flwr_tune/strategy.py.tpl

@@ -1,53 +1,48 @@
 """$project_name: A Flower / FlowerTune app."""
 
-from io import BytesIO
+from collections.abc import Iterable
 from logging import INFO, WARN
-from typing import List, Tuple, Union
+from typing import Optional
 
-from flwr.common import FitIns, FitRes, Parameters, log, parameters_to_ndarrays
-from flwr.server.client_manager import ClientManager
-from flwr.server.client_proxy import ClientProxy
-from flwr.server.strategy import FedAvg
+from flwr.app import ArrayRecord, ConfigRecord, Message, MetricRecord
+from flwr.common import log
+from flwr.serverapp import Grid
+from flwr.serverapp.strategy import FedAvg
 
 
 class FlowerTuneLlm(FedAvg):
     """Customised FedAvg strategy implementation.
-    
+
     This class behaves just like FedAvg but also tracks the communication
-    costs associated with `fit` over FL rounds.
+    costs associated with `train` over FL rounds.
     """
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
         self.comm_tracker = CommunicationTracker()
 
-    def configure_fit(
-        self, server_round: int, parameters: Parameters, client_manager: ClientManager
-    ):
+    def configure_train(
+            self, server_round: int, arrays: ArrayRecord, config: ConfigRecord, grid: Grid
+    ) -> Iterable[Message]:
         """Configure the next round of training."""
-        return_clients = super().configure_fit(server_round, parameters, client_manager)
-
-        # Test communication costs
-        fit_ins_list = [fit_ins for _, fit_ins in return_clients]
-        self.comm_tracker.track(fit_ins_list)
-
-        return return_clients
-
-    def aggregate_fit(
-        self,
-        server_round: int,
-        results: List[Tuple[ClientProxy, FitRes]],
-        failures: List[Union[Tuple[ClientProxy, FitRes], BaseException]],
-    ):
-        """Aggregate fit results using weighted average."""
-        # Test communication costs
-        fit_res_list = [fit_res for _, fit_res in results]
-        self.comm_tracker.track(fit_res_list)
-
-        parameters_aggregated, metrics_aggregated = super().aggregate_fit(
-            server_round, results, failures
-        )
+        messages = super().configure_train(server_round, arrays, config, grid)
+
+        # Track communication costs
+        self.comm_tracker.track(messages)
+
+        return messages
+
+    def aggregate_train(
+            self,
+            server_round: int,
+            replies: Iterable[Message],
+    ) -> tuple[Optional[ArrayRecord], Optional[MetricRecord]]:
+        """Aggregate ArrayRecords and MetricRecords in the received Messages."""
+        # Track communication costs
+        self.comm_tracker.track(replies)
 
-        return parameters_aggregated, metrics_aggregated
+        arrays, metrics = super().aggregate_train(server_round, replies)
+
+        return arrays, metrics
 
 
 class CommunicationTracker:
@@ -55,16 +50,16 @@ class CommunicationTracker:
     def __init__(self):
         self.curr_comm_cost = 0.0
 
-    @staticmethod
-    def _compute_bytes(parameters):
-        return sum([BytesIO(t).getbuffer().nbytes for t in parameters.tensors])
-
-    def track(self, fit_list: List[Union[FitIns, FitRes]]):
-        size_bytes_list = [
-            self._compute_bytes(fit_ele.parameters)
-            for fit_ele in fit_list
-        ]
-        comm_cost = sum(size_bytes_list) / 1024**2
+    def track(self, messages: Iterable[Message]):
+        comm_cost = (
+            sum(
+                record.count_bytes()
+                for msg in messages
+                if msg.has_content()
+                for record in msg.content.array_records.values()
+            )
+            / 1024**2
+        )
 
         self.curr_comm_cost += comm_cost
         log(

flwr/cli/new/templates/app/code/model.baseline.py.tpl

@@ -1,7 +1,5 @@
 """$project_name: A Flower Baseline."""
 
-from collections import OrderedDict
-
 import torch
 import torch.nn.functional as F
 from torch import nn
@@ -66,15 +64,3 @@ def test(net, testloader, device):
     accuracy = correct / len(testloader.dataset)
     loss = loss / len(testloader)
     return loss, accuracy
-
-
-def get_weights(net):
-    """Extract model parameters as numpy arrays from state_dict."""
-    return [val.cpu().numpy() for _, val in net.state_dict().items()]
-
-
-def set_weights(net, parameters):
-    """Apply parameters to an existing model."""
-    params_dict = zip(net.state_dict().keys(), parameters)
-    state_dict = OrderedDict({k: torch.from_numpy(v) for k, v in params_dict})
-    net.load_state_dict(state_dict, strict=True)

flwr/cli/new/templates/app/code/server.baseline.py.tpl

@@ -1,45 +1,43 @@
 """$project_name: A Flower Baseline."""
 
-from flwr.common import Context, Metrics, ndarrays_to_parameters
-from flwr.server import ServerApp, ServerAppComponents, ServerConfig
-from flwr.server.strategy import FedAvg
+import torch
+from flwr.app import ArrayRecord, Context
+from flwr.serverapp import Grid, ServerApp
+from flwr.serverapp.strategy import FedAvg
 
-from $import_name.model import Net, get_weights
+from $import_name.model import Net
 
+# Create ServerApp
+app = ServerApp()
 
-# Define metric aggregation function
-def weighted_average(metrics: list[tuple[int, Metrics]]) -> Metrics:
-    """Do weighted average of accuracy metric."""
-    # Multiply accuracy of each client by number of examples used
-    accuracies = [num_examples * float(m["accuracy"]) for num_examples, m in metrics]
-    examples = [num_examples for num_examples, _ in metrics]
-
-    # Aggregate and return custom metric (weighted average)
-    return {"accuracy": sum(accuracies) / sum(examples)}
 
+@app.main()
+def main(grid: Grid, context: Context) -> None:
+    """Main entry point for the ServerApp."""
 
-def server_fn(context: Context):
-    """Construct components that set the ServerApp behaviour."""
     # Read from config
     num_rounds = context.run_config["num-server-rounds"]
-    fraction_fit = context.run_config["fraction-fit"]
+    fraction_train = context.run_config["fraction-train"]
 
-    # Initialize model parameters
-    ndarrays = get_weights(Net())
-    parameters = ndarrays_to_parameters(ndarrays)
+    # Load global model
+    global_model = Net()
+    arrays = ArrayRecord(global_model.state_dict())
 
-    # Define strategy
+    # Initialize FedAvg strategy
     strategy = FedAvg(
-        fraction_fit=float(fraction_fit),
+        fraction_train=fraction_train,
         fraction_evaluate=1.0,
-        min_available_clients=2,
-        initial_parameters=parameters,
-        evaluate_metrics_aggregation_fn=weighted_average,
+        min_available_nodes=2,
     )
-    config = ServerConfig(num_rounds=int(num_rounds))
-
-    return ServerAppComponents(strategy=strategy, config=config)
 
+    # Start strategy, run FedAvg for `num_rounds`
+    result = strategy.start(
+        grid=grid,
+        initial_arrays=arrays,
+        num_rounds=num_rounds,
+    )
 
-# Create ServerApp
-app = ServerApp(server_fn=server_fn)
+    # Save final model to disk
+    print("\nSaving final model to disk...")
+    state_dict = result.arrays.to_torch_state_dict()
+    torch.save(state_dict, "final_model.pt")

flwr/cli/new/templates/app/code/server.huggingface.py.tpl

@@ -1,17 +1,22 @@
 """$project_name: A Flower / $framework_str app."""
 
-from flwr.common import Context, ndarrays_to_parameters
-from flwr.server import ServerApp, ServerAppComponents, ServerConfig
-from flwr.server.strategy import FedAvg
+import torch
+from flwr.app import ArrayRecord, Context
+from flwr.serverapp import Grid, ServerApp
+from flwr.serverapp.strategy import FedAvg
 from transformers import AutoModelForSequenceClassification
 
-from $import_name.task import get_weights
+# Create ServerApp
+app = ServerApp()
+
 
+@app.main()
+def main(grid: Grid, context: Context) -> None:
+    """Main entry point for the ServerApp."""
 
-def server_fn(context: Context):
     # Read from config
     num_rounds = context.run_config["num-server-rounds"]
-    fraction_fit = context.run_config["fraction-fit"]
+    fraction_train = context.run_config["fraction-train"]
 
     # Initialize global model
     model_name = context.run_config["model-name"]
@@ -19,20 +24,19 @@ def server_fn(context: Context):
     net = AutoModelForSequenceClassification.from_pretrained(
         model_name, num_labels=num_labels
     )
+    arrays = ArrayRecord(net.state_dict())
 
-    weights = get_weights(net)
-    initial_parameters = ndarrays_to_parameters(weights)
+    # Initialize FedAvg strategy
+    strategy = FedAvg(fraction_train=fraction_train)
 
-    # Define strategy
-    strategy = FedAvg(
-        fraction_fit=fraction_fit,
-        fraction_evaluate=1.0,
-        initial_parameters=initial_parameters,
+    # Start strategy, run FedAvg for `num_rounds`
+    result = strategy.start(
+        grid=grid,
+        initial_arrays=arrays,
+        num_rounds=num_rounds,
     )
-    config = ServerConfig(num_rounds=num_rounds)
-
-    return ServerAppComponents(strategy=strategy, config=config)
 
-
-# Create ServerApp
-app = ServerApp(server_fn=server_fn)
+    # Save final model to disk
+    print("\nSaving final model to disk...")
+    state_dict = result.arrays.to_torch_state_dict()
+    torch.save(state_dict, "final_model.pt")

flwr/cli/new/templates/app/code/server.jax.py.tpl

@@ -1,26 +1,39 @@
 """$project_name: A Flower / $framework_str app."""
 
-from flwr.common import Context, ndarrays_to_parameters
-from flwr.server import ServerApp, ServerAppComponents, ServerConfig
-from flwr.server.strategy import FedAvg
+import numpy as np
+from flwr.app import ArrayRecord, Context
+from flwr.serverapp import Grid, ServerApp
+from flwr.serverapp.strategy import FedAvg
+
 from $import_name.task import get_params, load_model
 
+# Create ServerApp
+app = ServerApp()
+
+
+@app.main()
+def main(grid: Grid, context: Context) -> None:
+    """Main entry point for the ServerApp."""
 
-def server_fn(context: Context):
     # Read from config
     num_rounds = context.run_config["num-server-rounds"]
     input_dim = context.run_config["input-dim"]
 
-    # Initialize global model
-    params = get_params(load_model((input_dim,)))
-    initial_parameters = ndarrays_to_parameters(params)
-
-    # Define strategy
-    strategy = FedAvg(initial_parameters=initial_parameters)
-    config = ServerConfig(num_rounds=num_rounds)
+    # Load global model
+    model = load_model((input_dim,))
+    arrays = ArrayRecord(get_params(model))
 
-    return ServerAppComponents(strategy=strategy, config=config)
+    # Initialize FedAvg strategy
+    strategy = FedAvg()
 
+    # Start strategy, run FedAvg for `num_rounds`
+    result = strategy.start(
+        grid=grid,
+        initial_arrays=arrays,
+        num_rounds=num_rounds,
+    )
 
-# Create ServerApp
-app = ServerApp(server_fn=server_fn)
+    # Save final model to disk
+    print("\nSaving final model to disk...")
+    ndarrays = result.arrays.to_numpy_ndarrays()
+    np.savez("final_model.npz", *ndarrays)