flwr-nightly 1.8.0.dev20240314__py3-none-any.whl → 1.11.0.dev20240813__py3-none-any.whl

flwr/cli/new/templates/app/code/client.mlx.py.tpl

@@ -0,0 +1,93 @@
+"""$project_name: A Flower / $framework_str app."""
+
+import mlx.core as mx
+import mlx.nn as nn
+import mlx.optimizers as optim
+from flwr.client import NumPyClient, ClientApp
+from flwr.common import Context
+
+from $import_name.task import (
+    batch_iterate,
+    eval_fn,
+    get_params,
+    load_data,
+    loss_fn,
+    set_params,
+    MLP,
+)
+
+
+# Define Flower Client and client_fn
+class FlowerClient(NumPyClient):
+    def __init__(
+        self,
+        data,
+        num_layers,
+        hidden_dim,
+        num_classes,
+        batch_size,
+        learning_rate,
+        num_epochs,
+    ):
+        self.num_layers = num_layers
+        self.hidden_dim = hidden_dim
+        self.num_classes = num_classes
+        self.batch_size = batch_size
+        self.learning_rate = learning_rate
+        self.num_epochs = num_epochs
+
+        self.train_images, self.train_labels, self.test_images, self.test_labels = data
+        self.model = MLP(
+            num_layers, self.train_images.shape[-1], hidden_dim, num_classes
+        )
+        self.optimizer = optim.SGD(learning_rate=learning_rate)
+        self.loss_and_grad_fn = nn.value_and_grad(self.model, loss_fn)
+        self.num_epochs = num_epochs
+        self.batch_size = batch_size
+
+    def get_parameters(self, config):
+        return get_params(self.model)
+
+    def set_parameters(self, parameters):
+        set_params(self.model, parameters)
+
+    def fit(self, parameters, config):
+        self.set_parameters(parameters)
+        for _ in range(self.num_epochs):
+            for X, y in batch_iterate(
+                self.batch_size, self.train_images, self.train_labels
+            ):
+                _, grads = self.loss_and_grad_fn(self.model, X, y)
+                self.optimizer.update(self.model, grads)
+                mx.eval(self.model.parameters(), self.optimizer.state)
+        return self.get_parameters(config={}), len(self.train_images), {}
+
+    def evaluate(self, parameters, config):
+        self.set_parameters(parameters)
+        accuracy = eval_fn(self.model, self.test_images, self.test_labels)
+        loss = loss_fn(self.model, self.test_images, self.test_labels)
+        return loss.item(), len(self.test_images), {"accuracy": accuracy.item()}
+
+
+def client_fn(context: Context):
+    partition_id = context.node_config["partition-id"]
+    num_partitions = context.node_config["num-partitions"]
+    data = load_data(partition_id, num_partitions)
+
+    num_layers = context.run_config["num-layers"]
+    hidden_dim = context.run_config["hidden-dim"]
+    num_classes = 10
+    batch_size = context.run_config["batch-size"]
+    learning_rate = context.run_config["lr"]
+    num_epochs = context.run_config["local-epochs"]
+
+    # Return Client instance
+    return FlowerClient(
+        data, num_layers, hidden_dim, num_classes, batch_size, learning_rate, num_epochs
+    ).to_client()
+
+
+# Flower ClientApp
+app = ClientApp(
+    client_fn,
+)

flwr/cli/new/templates/app/code/client.numpy.py.tpl

@@ -1,6 +1,7 @@
-"""$project_name: A Flower / NumPy app."""
+"""$project_name: A Flower / $framework_str app."""
 
 from flwr.client import NumPyClient, ClientApp
+from flwr.common import Context
 import numpy as np
 
 
@@ -15,7 +16,7 @@ class FlowerClient(NumPyClient):
         return float(0.0), 1, {"accuracy": float(1.0)}
 
 
-def client_fn(cid: str):
+def client_fn(context: Context):
     return FlowerClient().to_client()
 
 

flwr/cli/new/templates/app/code/client.pytorch.py.tpl

@@ -1,10 +1,11 @@
-"""$project_name: A Flower / PyTorch app."""
+"""$project_name: A Flower / $framework_str app."""
 
+import torch
 from flwr.client import NumPyClient, ClientApp
+from flwr.common import Context
 
-from $project_name.task import (
+from $import_name.task import (
     Net,
-    DEVICE,
     load_data,
     get_weights,
     set_weights,
@@ -15,29 +16,40 @@ from $project_name.task import (
 
 # Define Flower Client and client_fn
 class FlowerClient(NumPyClient):
-    def __init__(self, net, trainloader, valloader):
+    def __init__(self, net, trainloader, valloader, local_epochs):
         self.net = net
         self.trainloader = trainloader
         self.valloader = valloader
+        self.local_epochs = local_epochs
+        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        self.net.to(self.device)
 
     def fit(self, parameters, config):
         set_weights(self.net, parameters)
-        results = train(self.net, self.trainloader, self.valloader, 1, DEVICE)
-        return get_weights(self.net), len(self.trainloader.dataset), results
+        train_loss = train(
+            self.net,
+            self.trainloader,
+            self.local_epochs,
+            self.device,
+        )
+        return get_weights(self.net), len(self.trainloader.dataset), {"train_loss": train_loss}
 
     def evaluate(self, parameters, config):
         set_weights(self.net, parameters)
-        loss, accuracy = test(self.net, self.valloader)
+        loss, accuracy = test(self.net, self.valloader, self.device)
         return loss, len(self.valloader.dataset), {"accuracy": accuracy}
 
 
-def client_fn(cid):
+def client_fn(context: Context):
     # Load model and data
-    net = Net().to(DEVICE)
-    trainloader, valloader = load_data(int(cid), 2)
+    net = Net()
+    partition_id = context.node_config["partition-id"]
+    num_partitions = context.node_config["num-partitions"]
+    trainloader, valloader = load_data(partition_id, num_partitions)
+    local_epochs = context.run_config["local-epochs"]
 
     # Return Client instance
-    return FlowerClient(net, trainloader, valloader).to_client()
+    return FlowerClient(net, trainloader, valloader, local_epochs).to_client()
 
 
 # Flower ClientApp

flwr/cli/new/templates/app/code/client.sklearn.py.tpl

@@ -0,0 +1,97 @@
+"""$project_name: A Flower / $framework_str app."""
+
+import warnings
+
+import numpy as np
+from flwr.client import NumPyClient, ClientApp
+from flwr.common import Context
+from flwr_datasets import FederatedDataset
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import log_loss
+
+
+def get_model_parameters(model):
+    if model.fit_intercept:
+        params = [
+            model.coef_,
+            model.intercept_,
+        ]
+    else:
+        params = [model.coef_]
+    return params
+
+
+def set_model_params(model, params):
+    model.coef_ = params[0]
+    if model.fit_intercept:
+        model.intercept_ = params[1]
+    return model
+
+
+def set_initial_params(model):
+    n_classes = 10  # MNIST has 10 classes
+    n_features = 784  # Number of features in dataset
+    model.classes_ = np.array([i for i in range(10)])
+
+    model.coef_ = np.zeros((n_classes, n_features))
+    if model.fit_intercept:
+        model.intercept_ = np.zeros((n_classes,))
+
+
+class FlowerClient(NumPyClient):
+    def __init__(self, model, X_train, X_test, y_train, y_test):
+        self.model = model
+        self.X_train = X_train
+        self.X_test = X_test
+        self.y_train = y_train
+        self.y_test = y_test
+
+    def get_parameters(self, config):
+        return get_model_parameters(self.model)
+
+    def fit(self, parameters, config):
+        set_model_params(self.model, parameters)
+
+        # Ignore convergence failure due to low local epochs
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            self.model.fit(self.X_train, self.y_train)
+
+        return get_model_parameters(self.model), len(self.X_train), {}
+
+    def evaluate(self, parameters, config):
+        set_model_params(self.model, parameters)
+
+        loss = log_loss(self.y_test, self.model.predict_proba(self.X_test))
+        accuracy = self.model.score(self.X_test, self.y_test)
+
+        return loss, len(self.X_test), {"accuracy": accuracy}
+
+
+def client_fn(context: Context):
+    partition_id = context.node_config["partition-id"]
+    num_partitions = context.node_config["num-partitions"]
+    fds = FederatedDataset(dataset="mnist", partitioners={"train": num_partitions})
+    dataset = fds.load_partition(partition_id, "train").with_format("numpy")
+
+    X, y = dataset["image"].reshape((len(dataset), -1)), dataset["label"]
+
+    # Split the on edge data: 80% train, 20% test
+    X_train, X_test = X[: int(0.8 * len(X))], X[int(0.8 * len(X)) :]
+    y_train, y_test = y[: int(0.8 * len(y))], y[int(0.8 * len(y)) :]
+
+    # Create LogisticRegression Model
+    model = LogisticRegression(
+        penalty="l2",
+        max_iter=1,  # local epoch
+        warm_start=True,  # prevent refreshing weights when fitting
+    )
+
+    # Setting initial parameters, akin to model.compile for keras models
+    set_initial_params(model)
+
+    return FlowerClient(model, X_train, X_test, y_train, y_test).to_client()
+
+
+# Flower ClientApp
+app = ClientApp(client_fn=client_fn)

flwr/cli/new/templates/app/code/client.tensorflow.py.tpl

@@ -1 +1,60 @@
-"""$project_name: A Flower / TensorFlow app."""
+"""$project_name: A Flower / $framework_str app."""
+
+from flwr.client import NumPyClient, ClientApp
+from flwr.common import Context
+
+from $import_name.task import load_data, load_model
+
+
+# Define Flower Client and client_fn
+class FlowerClient(NumPyClient):
+    def __init__(
+        self, model, data, epochs, batch_size, verbose
+    ):
+        self.model = model
+        self.x_train, self.y_train, self.x_test, self.y_test = data
+        self.epochs = epochs
+        self.batch_size = batch_size
+        self.verbose = verbose
+
+    def get_parameters(self, config):
+        return self.model.get_weights()
+
+    def fit(self, parameters, config):
+        self.model.set_weights(parameters)
+        self.model.fit(
+            self.x_train,
+            self.y_train,
+            epochs=self.epochs,
+            batch_size=self.batch_size,
+            verbose=self.verbose,
+        )
+        return self.model.get_weights(), len(self.x_train), {}
+
+    def evaluate(self, parameters, config):
+        self.model.set_weights(parameters)
+        loss, accuracy = self.model.evaluate(self.x_test, self.y_test, verbose=0)
+        return loss, len(self.x_test), {"accuracy": accuracy}
+
+
+def client_fn(context: Context):
+    # Load model and data
+    net = load_model()
+
+    partition_id = context.node_config["partition-id"]
+    num_partitions = context.node_config["num-partitions"]
+    data = load_data(partition_id, num_partitions)
+    epochs = context.run_config["local-epochs"]
+    batch_size = context.run_config["batch-size"]
+    verbose = context.run_config.get("verbose")
+
+    # Return Client instance
+    return FlowerClient(
+        net, data, epochs, batch_size, verbose
+    ).to_client()
+
+
+# Flower ClientApp
+app = ClientApp(
+    client_fn=client_fn,
+)

flwr/cli/new/templates/app/code/flwr_tune/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Flower CLI `new` command app / code / flwr_tune templates."""

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

@@ -0,0 +1,89 @@
+"""$project_name: A Flower / FlowerTune app."""
+
+import os
+import warnings
+from datetime import datetime
+
+from flwr_datasets import FederatedDataset
+from hydra import compose, initialize
+from hydra.utils import instantiate
+
+from flwr.client import ClientApp
+from flwr.common import Context, ndarrays_to_parameters
+from flwr.server import ServerApp, ServerAppComponents, ServerConfig
+
+from $import_name.client_app import gen_client_fn, get_parameters
+from $import_name.dataset import get_tokenizer_and_data_collator_and_propt_formatting
+from $import_name.models import get_model
+from $import_name.server_app import fit_weighted_average, get_evaluate_fn, get_on_fit_config
+
+# Avoid warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+os.environ["RAY_DISABLE_DOCKER_CPU_WARNING"] = "1"
+
+# Initialise regular config
+with initialize(config_path="conf", version_base="1.1"):
+    cfg = compose(config_name="config")
+
+# Initialise static config
+with initialize(config_path="conf", version_base="1.1"):
+    cfg_static = compose(config_name="static_config")
+
+cfg.train.num_rounds = cfg_static.num_rounds
+
+# Create output directory given current timestamp
+current_time = datetime.now()
+folder_name = current_time.strftime("%Y-%m-%d_%H-%M-%S")
+save_path = os.path.join(os.getcwd(), f"results/{folder_name}")
+os.makedirs(save_path, exist_ok=True)
+
+# Partition dataset and get dataloaders
+partitioner = instantiate(cfg_static.partitioner)
+fds = FederatedDataset(
+    dataset=cfg_static.dataset.name, partitioners={"train": partitioner}
+)
+(
+    tokenizer,
+    data_collator,
+    formatting_prompts_func,
+) = get_tokenizer_and_data_collator_and_propt_formatting(cfg.model.name)
+
+# ClientApp for Flower Next
+client = ClientApp(
+    client_fn=gen_client_fn(
+        fds,
+        tokenizer,
+        formatting_prompts_func,
+        data_collator,
+        cfg.model,
+        cfg.train,
+        save_path,
+    ),
+)
+
+# 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)
+
+def server_fn(context: Context):
+    # Instantiate strategy according to config. Here we pass other arguments
+    # that are only defined at runtime.
+    strategy = instantiate(
+        cfg.strategy,
+        on_fit_config_fn=get_on_fit_config(),
+        fit_metrics_aggregation_fn=fit_weighted_average,
+        initial_parameters=init_model_parameters,
+        evaluate_fn=get_evaluate_fn(
+            cfg.model, cfg.train.save_every_round, cfg_static.num_rounds, save_path
+        ),
+    )
+
+    config = ServerConfig(num_rounds=cfg_static.num_rounds)
+
+    return ServerAppComponents(strategy=strategy, config=config)
+
+
+# ServerApp for Flower Next
+server = ServerApp(server_fn=server_fn)

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

@@ -0,0 +1,126 @@
+"""$project_name: A Flower / FlowerTune app."""
+
+from collections import OrderedDict
+from typing import Callable, Dict, Tuple
+
+import torch
+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
+
+from flwr.client import NumPyClient
+from flwr.common import Context
+from flwr.common.typing import NDArrays, Scalar
+from $import_name.dataset import reformat
+from $import_name.models import cosine_annealing, get_model
+
+
+# pylint: disable=too-many-arguments
+# pylint: disable=too-many-instance-attributes
+class FlowerClient(NumPyClient):
+    """Standard Flower client for CNN training."""
+
+    def __init__(
+        self,
+        model_cfg: DictConfig,
+        train_cfg: DictConfig,
+        trainset,
+        tokenizer,
+        formatting_prompts_func,
+        data_collator,
+        save_path,
+    ):  # 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_argumnets = TrainingArguments(**train_cfg.training_arguments)
+        self.tokenizer = tokenizer
+        self.formatting_prompts_func = formatting_prompts_func
+        self.data_collator = data_collator
+        self.save_path = save_path
+
+        # instantiate model
+        self.model = get_model(model_cfg)
+
+        self.trainset = trainset
+
+    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.train_cfg.num_rounds,
+            self.train_cfg.learning_rate_max,
+            self.train_cfg.learning_rate_min,
+        )
+
+        self.training_argumnets.learning_rate = new_lr
+        self.training_argumnets.output_dir = self.save_path
+
+        # Construct trainer
+        trainer = SFTTrainer(
+            model=self.model,
+            tokenizer=self.tokenizer,
+            args=self.training_argumnets,
+            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 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()]
+
+
+def gen_client_fn(
+    fds,
+    tokenizer,
+    formatting_prompts_func,
+    data_collator,
+    model_cfg: DictConfig,
+    train_cfg: DictConfig,
+    save_path: str,
+) -> Callable[[Context], FlowerClient]:  # pylint: disable=too-many-arguments
+    """Generate the client function that creates the Flower Clients."""
+
+    def client_fn(context: Context) -> FlowerClient:
+        """Create a Flower client representing a single organization."""
+        # Let's get the partition corresponding to the i-th client
+        partition_id = context.node_config["partition-id"]
+        client_trainset = fds.load_partition(partition_id, "train")
+        client_trainset = reformat(client_trainset, llm_task="$llm_challenge_str")
+
+        return FlowerClient(
+            model_cfg,
+            train_cfg,
+            client_trainset,
+            tokenizer,
+            formatting_prompts_func,
+            data_collator,
+            save_path,
+        ).to_client()
+
+    return client_fn

flwr/cli/new/templates/app/code/flwr_tune/config.yaml.tpl

@@ -0,0 +1,34 @@
+# Federated Instruction Tuning
+---
+model:
+  name: "mistralai/Mistral-7B-v0.3"
+  quantization: 4 # 8 or 4 if you want to do quantization with BitsAndBytes
+  gradient_checkpointing: True
+  lora:
+    peft_lora_r: 32
+    peft_lora_alpha: 64
+
+train:
+  num_rounds: null
+  save_every_round: 5
+  learning_rate_max: 5e-5
+  learning_rate_min: 1e-6
+  seq_length: 512
+  training_arguments:
+    output_dir: null # to be set by hydra
+    learning_rate: null # to be set by the client
+    per_device_train_batch_size: 16
+    gradient_accumulation_steps: 1
+    logging_steps: 10
+    num_train_epochs: 3
+    max_steps: 10
+    report_to: null
+    save_steps: 1000
+    save_total_limit: 10
+    gradient_checkpointing: True
+    lr_scheduler_type: "constant"
+
+strategy:
+  _target_: flwr.server.strategy.FedAvg
+  fraction_fit: $fraction_fit
+  fraction_evaluate: 0.0 # no client evaluation

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

@@ -0,0 +1,57 @@
+"""$project_name: A Flower / FlowerTune app."""
+
+from transformers import AutoTokenizer
+from trl import DataCollatorForCompletionOnlyLM
+
+
+def formatting_prompts_func(example):
+    """Construct prompts."""
+    output_texts = []
+    # Constructing a standard Alpaca
+    # (https://github.com/tatsu-lab/stanford_alpaca#data-release) prompt
+    mssg = (
+        "Below is an instruction that describes a task. "
+        "Write a response that appropriately completes the request."
+    )
+    for i in range(len(example["instruction"])):
+        text = (
+            f"{mssg}\n### Instruction:\n{example['instruction'][i]}\n"
+            f"### Response: {example['response'][i]}"
+        )
+        output_texts.append(text)
+    return output_texts
+
+
+def get_tokenizer_and_data_collator_and_propt_formatting(model_name: str):
+    """Get tokenizer, data_collator and prompt formatting."""
+    # From: https://huggingface.co/docs/trl/en/sft_trainer
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name, use_fast=True, padding_side="right"
+    )
+    tokenizer.pad_token = tokenizer.eos_token
+    response_template_with_context = "\n### Response:"  # alpaca response tag
+    response_template_ids = tokenizer.encode(
+        response_template_with_context, add_special_tokens=False
+    )[2:]
+    data_collator = DataCollatorForCompletionOnlyLM(
+        response_template_ids, tokenizer=tokenizer
+    )
+
+    return tokenizer, data_collator, formatting_prompts_func
+
+
+def formatting(dataset):
+    """Format dataset."""
+    dataset["instruction"] = dataset["instruction"] + " " + dataset["input"]
+    return dataset
+
+
+def reformat(dataset, llm_task):
+    """Reformat datasets."""
+    dataset = dataset.rename_column("output", "response")
+    if llm_task == "finance" or llm_task == "code":
+        dataset = dataset.map(formatting, remove_columns=["input"])
+    if llm_task == "medical":
+        dataset = dataset.remove_columns(["instruction"])
+        dataset = dataset.rename_column("input", "instruction")
+    return dataset