flwr-nightly 1.8.0.dev20240315__py3-none-any.whl → 1.15.0.dev20250114__py3-none-any.whl

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

@@ -0,0 +1,87 @@
+"""$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."""
+    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."""
+    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 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="$llm_challenge_str")
+    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

@@ -0,0 +1,78 @@
+"""$project_name: A Flower / FlowerTune app."""
+
+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.utils import prepare_model_for_kbit_training
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig
+
+from flwr.common.typing import NDArrays
+
+
+def cosine_annealing(
+    current_round: int,
+    total_round: int,
+    lrate_max: float = 0.001,
+    lrate_min: float = 0.0,
+) -> float:
+    """Implement cosine annealing learning rate schedule."""
+    cos_inner = math.pi * current_round / total_round
+    return lrate_min + 0.5 * (lrate_max - lrate_min) * (1 + math.cos(cos_inner))
+
+
+def get_model(model_cfg: DictConfig):
+    """Load model with appropriate quantization config and other optimizations.
+    """
+    if model_cfg.quantization == 4:
+        quantization_config = BitsAndBytesConfig(load_in_4bit=True)
+    elif model_cfg.quantization == 8:
+        quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+    else:
+        raise ValueError(
+            f"Use 4-bit or 8-bit quantization. You passed: {model_cfg.quantization}/"
+        )
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_cfg.name,
+        quantization_config=quantization_config,
+        torch_dtype=torch.bfloat16,
+        low_cpu_mem_usage=True,
+    )
+
+    model = prepare_model_for_kbit_training(
+        model, use_gradient_checkpointing=model_cfg.gradient_checkpointing
+    )
+
+    peft_config = LoraConfig(
+        r=model_cfg.lora.peft_lora_r,
+        lora_alpha=model_cfg.lora.peft_lora_alpha,
+        lora_dropout=0.075,
+        task_type="CAUSAL_LM",
+    )
+
+    if model_cfg.gradient_checkpointing:
+        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

@@ -0,0 +1,94 @@
+"""$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.models import get_model, get_parameters, set_parameters
+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/server.huggingface.py.tpl

@@ -0,0 +1,38 @@
+"""$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
+from transformers import AutoModelForSequenceClassification
+
+from $import_name.task import get_weights
+
+
+def server_fn(context: Context):
+    # Read from config
+    num_rounds = context.run_config["num-server-rounds"]
+    fraction_fit = context.run_config["fraction-fit"]
+
+    # Initialize global model
+    model_name = context.run_config["model-name"]
+    num_labels = context.run_config["num-labels"]
+    net = AutoModelForSequenceClassification.from_pretrained(
+        model_name, num_labels=num_labels
+    )
+
+    weights = get_weights(net)
+    initial_parameters = ndarrays_to_parameters(weights)
+
+    # Define strategy
+    strategy = FedAvg(
+        fraction_fit=fraction_fit,
+        fraction_evaluate=1.0,
+        initial_parameters=initial_parameters,
+    )
+    config = ServerConfig(num_rounds=num_rounds)
+
+    return ServerAppComponents(strategy=strategy, config=config)
+
+
+# Create ServerApp
+app = ServerApp(server_fn=server_fn)

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

@@ -0,0 +1,26 @@
+"""$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
+from $import_name.task import get_params, load_model
+
+
+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)
+
+    return ServerAppComponents(strategy=strategy, config=config)
+
+
+# Create ServerApp
+app = ServerApp(server_fn=server_fn)

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

@@ -0,0 +1,31 @@
+"""$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
+from $import_name.task import MLP, get_params
+
+
+def server_fn(context: Context):
+    # Read from config
+    num_rounds = context.run_config["num-server-rounds"]
+
+    num_classes = 10
+    num_layers = context.run_config["num-layers"]
+    input_dim = context.run_config["input-dim"]
+    hidden_dim = context.run_config["hidden-dim"]
+
+    # Initialize global model
+    model = MLP(num_layers, input_dim, hidden_dim, num_classes)
+    params = get_params(model)
+    initial_parameters = ndarrays_to_parameters(params)
+
+    # Define strategy
+    strategy = FedAvg(initial_parameters=initial_parameters)
+    config = ServerConfig(num_rounds=num_rounds)
+
+    return ServerAppComponents(strategy=strategy, config=config)
+
+
+# Create ServerApp
+app = ServerApp(server_fn=server_fn)

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

@@ -1,12 +1,25 @@
-"""$project_name: A Flower / NumPy app."""
+"""$project_name: A Flower / $framework_str app."""
 
-import flwr as fl
+from flwr.common import Context, ndarrays_to_parameters
+from flwr.server import ServerApp, ServerAppComponents, ServerConfig
+from flwr.server.strategy import FedAvg
+from $import_name.task import get_dummy_model
 
-# Configure the strategy
-strategy = fl.server.strategy.FedAvg()
 
-# Flower ServerApp
-app = fl.server.ServerApp(
-    config=fl.server.ServerConfig(num_rounds=1),
-    strategy=strategy,
-)
+def server_fn(context: Context):
+    # Read from config
+    num_rounds = context.run_config["num-server-rounds"]
+
+    # Initial model
+    model = get_dummy_model()
+    dummy_parameters = ndarrays_to_parameters([model])
+
+    # Define strategy
+    strategy = FedAvg(initial_parameters=dummy_parameters)
+    config = ServerConfig(num_rounds=num_rounds)
+
+    return ServerAppComponents(strategy=strategy, config=config)
+
+
+# Create ServerApp
+app = ServerApp(server_fn=server_fn)

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

@@ -1,28 +1,31 @@
-"""$project_name: A Flower / PyTorch app."""
+"""$project_name: A Flower / $framework_str app."""
 
-from flwr.common import ndarrays_to_parameters
-from flwr.server import ServerApp, ServerConfig
+from flwr.common import Context, ndarrays_to_parameters
+from flwr.server import ServerApp, ServerAppComponents, ServerConfig
 from flwr.server.strategy import FedAvg
+from $import_name.task import Net, get_weights
 
-from $project_name.task import Net, get_weights
 
+def server_fn(context: Context):
+    # 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)
+    # Initialize model parameters
+    ndarrays = get_weights(Net())
+    parameters = ndarrays_to_parameters(ndarrays)
 
+    # Define strategy
+    strategy = FedAvg(
+        fraction_fit=fraction_fit,
+        fraction_evaluate=1.0,
+        min_available_clients=2,
+        initial_parameters=parameters,
+    )
+    config = ServerConfig(num_rounds=num_rounds)
 
-# Define strategy
-strategy = FedAvg(
-    fraction_fit=1.0,
-    fraction_evaluate=1.0,
-    min_available_clients=2,
-    initial_parameters=parameters,
-)
+    return ServerAppComponents(strategy=strategy, config=config)
 
 
 # Create ServerApp
-app = ServerApp(
-    config=ServerConfig(num_rounds=3),
-    strategy=strategy,
-)
+app = ServerApp(server_fn=server_fn)