flwr-nightly 1.11.0.dev20240823__py3-none-any.whl → 1.12.0.dev20240906__py3-none-any.whl

flwr/cli/new/templates/app/code/flwr_tune/{client.py.tpl → client_app.py.tpl}

@@ -1,20 +1,32 @@
 """$project_name: A Flower / FlowerTune app."""
 
+import os
+import warnings
 from collections import OrderedDict
-from typing import Callable, Dict, Tuple
+from typing import Dict, Tuple
 
 import torch
+from flwr.client import ClientApp, NumPyClient
+from flwr.common import Context
+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
 
-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.dataset import (
+    get_tokenizer_and_data_collator_and_propt_formatting,
+    load_data,
+    replace_keys,
+)
 from $import_name.models import cosine_annealing, get_model
 
+# Avoid warnings
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+os.environ["RAY_DISABLE_DOCKER_CPU_WARNING"] = "1"
+warnings.filterwarnings("ignore", category=UserWarning)
+
 
 # pylint: disable=too-many-arguments
 # pylint: disable=too-many-instance-attributes
@@ -29,7 +41,7 @@ class FlowerClient(NumPyClient):
         tokenizer,
         formatting_prompts_func,
         data_collator,
-        save_path,
+        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
@@ -37,13 +49,12 @@ class FlowerClient(NumPyClient):
         self.tokenizer = tokenizer
         self.formatting_prompts_func = formatting_prompts_func
         self.data_collator = data_collator
-        self.save_path = save_path
+        self.num_rounds = num_rounds
+        self.trainset = trainset
 
         # instantiate model
         self.model = get_model(model_cfg)
 
-        self.trainset = trainset
-
     def fit(
         self, parameters: NDArrays, config: Dict[str, Scalar]
     ) -> Tuple[NDArrays, int, Dict]:
@@ -52,13 +63,13 @@ class FlowerClient(NumPyClient):
 
         new_lr = cosine_annealing(
             int(config["current_round"]),
-            self.train_cfg.num_rounds,
+            self.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
+        self.training_argumnets.output_dir = config["save_path"]
 
         # Construct trainer
         trainer = SFTTrainer(
@@ -95,32 +106,31 @@ def get_parameters(model) -> NDArrays:
     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
+def client_fn(context: Context) -> FlowerClient:
+    """Create a Flower client representing a single organization."""
+    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)))
+
+    # Let's get the client partition
+    client_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()
+
+
+# Flower ClientApp
+app = ClientApp(client_fn)

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

@@ -1,8 +1,12 @@
 """$project_name: A Flower / FlowerTune app."""
 
+from flwr_datasets import FederatedDataset
+from flwr_datasets.partitioner import IidPartitioner
 from transformers import AutoTokenizer
 from trl import DataCollatorForCompletionOnlyLM
 
+FDS = None  # Cache FederatedDataset
+
 
 def formatting_prompts_func(example):
     """Construct prompts."""
@@ -24,7 +28,6 @@ def formatting_prompts_func(example):
 
 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"
     )
@@ -49,9 +52,36 @@ def formatting(dataset):
 def reformat(dataset, llm_task):
     """Reformat datasets."""
     dataset = dataset.rename_column("output", "response")
-    if llm_task == "finance" or llm_task == "code":
+    if llm_task in ["finance", "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
+
+
+def load_data(partition_id: int, num_partitions: int, dataset_name: str):
+    """Load partition data."""
+    # Only initialize `FederatedDataset` once
+    global FDS
+    if FDS is None:
+        partitioner = IidPartitioner(num_partitions=num_partitions)
+        FDS = FederatedDataset(
+            dataset=dataset_name,
+            partitioners={"train": partitioner},
+        )
+    client_trainset = FDS.load_partition(partition_id, "train")
+    client_trainset = reformat(client_trainset, llm_task="generalnlp")
+    return client_trainset
+
+
+def replace_keys(input_dict, match="-", target="_"):
+    """Recursively replace match string with target string in dictionary keys."""
+    new_dict = {}
+    for key, value in input_dict.items():
+        new_key = key.replace(match, target)
+        if isinstance(value, dict):
+            new_dict[new_key] = replace_keys(value, match, target)
+        else:
+            new_dict[new_key] = value
+    return new_dict

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

@@ -22,9 +22,6 @@ def cosine_annealing(
 
 def get_model(model_cfg: DictConfig):
     """Load model with appropriate quantization config and other optimizations.
-
-    Please refer to this example for `peft + BitsAndBytes`:
-    https://github.com/huggingface/peft/blob/main/examples/fp4_finetuning/finetune_fp4_opt_bnb_peft.py
     """
     if model_cfg.quantization == 4:
         quantization_config = BitsAndBytesConfig(load_in_4bit=True)

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

@@ -0,0 +1,95 @@
+"""$project_name: A Flower / FlowerTune app."""
+
+import os
+from datetime import datetime
+
+from flwr.common import Context, ndarrays_to_parameters
+from flwr.common.config import unflatten_dict
+from flwr.server import ServerApp, ServerAppComponents, ServerConfig
+from omegaconf import DictConfig
+
+from $import_name.client_app import get_parameters, set_parameters
+from $import_name.models import get_model
+from $import_name.dataset import replace_keys
+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
+
+
+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."""
+    # 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)
+
+    # Read from config
+    num_rounds = context.run_config["num-server-rounds"]
+    cfg = DictConfig(replace_keys(unflatten_dict(context.run_config)))
+
+    # 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)
+
+    # Define strategy
+    strategy = FlowerTuneLlm(
+        fraction_fit=cfg.strategy.fraction_fit,
+        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,
+        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)
+
+
+# Flower ServerApp
+app = ServerApp(server_fn=server_fn)

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

@@ -0,0 +1,83 @@
+"""$project_name: A Flower / FlowerTune app."""
+
+from io import BytesIO
+from logging import INFO, WARN
+from typing import List, Tuple, Union
+
+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
+
+
+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.
+    """
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.comm_tracker = CommunicationTracker()
+
+    def configure_fit(
+        self, server_round: int, parameters: Parameters, client_manager: ClientManager
+    ):
+        """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
+        )
+
+        return parameters_aggregated, metrics_aggregated
+
+
+class CommunicationTracker:
+    """Communication costs tracker over FL rounds."""
+    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
+
+        self.curr_comm_cost += comm_cost
+        log(
+            INFO,
+            "Communication budget: used %.2f MB (+%.2f MB this round) / 200,000 MB",
+            self.curr_comm_cost,
+            comm_cost,
+        )
+
+        if self.curr_comm_cost > 2e5:
+            log(
+                WARN,
+                "The accumulated communication cost has exceeded 200,000 MB. "
+                "Please consider reducing it if you plan to participate "
+                "FlowerTune LLM Leaderboard.",
+            )

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

@@ -0,0 +1,80 @@
+"""$project_name: A Flower Baseline."""
+
+from collections import OrderedDict
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class Net(nn.Module):
+    """Model (simple CNN adapted from 'PyTorch: A 60 Minute Blitz')."""
+
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(3, 6, 5)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(6, 16, 5)
+        self.fc1 = nn.Linear(16 * 5 * 5, 120)
+        self.fc2 = nn.Linear(120, 84)
+        self.fc3 = nn.Linear(84, 10)
+
+    def forward(self, x):
+        """Do forward."""
+        x = self.pool(F.relu(self.conv1(x)))
+        x = self.pool(F.relu(self.conv2(x)))
+        x = x.view(-1, 16 * 5 * 5)
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        return self.fc3(x)
+
+
+def train(net, trainloader, epochs, device):
+    """Train the model on the training set."""
+    net.to(device)  # move model to GPU if available
+    criterion = torch.nn.CrossEntropyLoss()
+    criterion.to(device)
+    optimizer = torch.optim.SGD(net.parameters(), lr=0.1, momentum=0.9)
+    net.train()
+    running_loss = 0.0
+    for _ in range(epochs):
+        for batch in trainloader:
+            images = batch["img"]
+            labels = batch["label"]
+            optimizer.zero_grad()
+            loss = criterion(net(images.to(device)), labels.to(device))
+            loss.backward()
+            optimizer.step()
+            running_loss += loss.item()
+
+    avg_trainloss = running_loss / len(trainloader)
+    return avg_trainloss
+
+
+def test(net, testloader, device):
+    """Validate the model on the test set."""
+    net.to(device)
+    criterion = torch.nn.CrossEntropyLoss()
+    correct, loss = 0, 0.0
+    with torch.no_grad():
+        for batch in testloader:
+            images = batch["img"].to(device)
+            labels = batch["label"].to(device)
+            outputs = net(images)
+            loss += criterion(outputs, labels).item()
+            correct += (torch.max(outputs.data, 1)[1] == labels).sum().item()
+    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.tensor(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

@@ -0,0 +1,46 @@
+"""$project_name: A Flower Baseline."""
+
+from typing import List, Tuple
+
+from flwr.common import Context, Metrics, ndarrays_to_parameters
+from flwr.server import ServerApp, ServerAppComponents, ServerConfig
+from flwr.server.strategy import FedAvg
+from $import_name.model import Net, get_weights
+
+
+# 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)}
+
+
+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"]
+
+    # Initialize model parameters
+    ndarrays = get_weights(Net())
+    parameters = ndarrays_to_parameters(ndarrays)
+
+    # Define strategy
+    strategy = FedAvg(
+        fraction_fit=float(fraction_fit),
+        fraction_evaluate=1.0,
+        min_available_clients=2,
+        initial_parameters=parameters,
+        evaluate_metrics_aggregation_fn=weighted_average,
+    )
+    config = ServerConfig(num_rounds=int(num_rounds))
+
+    return ServerAppComponents(strategy=strategy, config=config)
+
+
+# Create ServerApp
+app = ServerApp(server_fn=server_fn)

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

@@ -0,0 +1 @@
+"""$project_name: A Flower Baseline."""

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

@@ -0,0 +1 @@
+"""$project_name: A Flower Baseline."""

flwr/cli/new/templates/app/pyproject.baseline.toml.tpl

@@ -0,0 +1,138 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "$package_name"
+version = "1.0.0"
+description = ""
+license = "Apache-2.0"
+dependencies = [
+    "flwr[simulation]>=1.11.0",
+    "flwr-datasets[vision]>=0.3.0",
+    "torch==2.2.1",
+    "torchvision==0.17.1",
+]
+
+[tool.hatch.metadata]
+allow-direct-references = true
+
+[project.optional-dependencies]
+dev = [
+    "isort==5.13.2",
+    "black==24.2.0",
+    "docformatter==1.7.5",
+    "mypy==1.8.0",
+    "pylint==3.2.6",
+    "flake8==5.0.4",
+    "pytest==6.2.4",
+    "pytest-watch==4.2.0",
+    "ruff==0.1.9",
+    "types-requests==2.31.0.20240125",
+]
+
+[tool.isort]
+profile = "black"
+known_first_party = ["flwr"]
+
+[tool.black]
+line-length = 88
+target-version = ["py38", "py39", "py310", "py311"]
+
+[tool.pytest.ini_options]
+minversion = "6.2"
+addopts = "-qq"
+testpaths = [
+    "flwr_baselines",
+]
+
+[tool.mypy]
+ignore_missing_imports = true
+strict = false
+plugins = "numpy.typing.mypy_plugin"
+
+[tool.pylint."MESSAGES CONTROL"]
+disable = "duplicate-code,too-few-public-methods,useless-import-alias"
+good-names = "i,j,k,_,x,y,X,Y,K,N"
+max-args = 10
+max-attributes = 15
+max-locals = 36
+max-branches = 20
+max-statements = 55
+
+[tool.pylint.typecheck]
+generated-members = "numpy.*, torch.*, tensorflow.*"
+
+[[tool.mypy.overrides]]
+module = [
+    "importlib.metadata.*",
+    "importlib_metadata.*",
+]
+follow_imports = "skip"
+follow_imports_for_stubs = true
+disallow_untyped_calls = false
+
+[[tool.mypy.overrides]]
+module = "torch.*"
+follow_imports = "skip"
+follow_imports_for_stubs = true
+
+[tool.docformatter]
+wrap-summaries = 88
+wrap-descriptions = 88
+
+[tool.ruff]
+target-version = "py38"
+line-length = 88
+select = ["D", "E", "F", "W", "B", "ISC", "C4"]
+fixable = ["D", "E", "F", "W", "B", "ISC", "C4"]
+ignore = ["B024", "B027"]
+exclude = [
+    ".bzr",
+    ".direnv",
+    ".eggs",
+    ".git",
+    ".hg",
+    ".mypy_cache",
+    ".nox",
+    ".pants.d",
+    ".pytype",
+    ".ruff_cache",
+    ".svn",
+    ".tox",
+    ".venv",
+    "__pypackages__",
+    "_build",
+    "buck-out",
+    "build",
+    "dist",
+    "node_modules",
+    "venv",
+    "proto",
+]
+
+[tool.ruff.pydocstyle]
+convention = "numpy"
+
+[tool.hatch.build.targets.wheel]
+packages = ["."]
+
+[tool.flwr.app]
+publisher = "$username"
+
+[tool.flwr.app.components]
+serverapp = "$import_name.server_app:app"
+clientapp = "$import_name.client_app:app"
+
+[tool.flwr.app.config]
+num-server-rounds = 3
+fraction-fit = 0.5
+local-epochs = 1
+
+[tool.flwr.federations]
+default = "local-simulation"
+
+[tool.flwr.federations.local-simulation]
+options.num-supernodes = 10
+options.backend.client-resources.num-cpus = 2
+options.backend.client-resources.num-gpus = 0.0