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

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

@@ -0,0 +1,59 @@
+"""$project_name: A Flower / FlowerTune app."""
+
+import math
+
+import torch
+from omegaconf import DictConfig
+from peft import LoraConfig, get_peft_model
+from peft.utils import prepare_model_for_kbit_training
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig
+
+
+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.
+
+    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)
+    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)

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

@@ -0,0 +1,48 @@
+"""$project_name: A Flower / FlowerTune app."""
+
+from $import_name.client_app import set_parameters
+from $import_name.models import get_model
+
+
+# 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():
+    """
+    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 = {"current_round": server_round}
+        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)}

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

@@ -0,0 +1,11 @@
+# Federated Instruction Tuning (static)
+---
+dataset:
+  name: $dataset_name
+
+# FL experimental settings
+num_clients: $num_clients # total number of clients
+num_rounds: 200
+partitioner:
+  _target_: flwr_datasets.partitioner.IidPartitioner
+  num_partitions: $num_clients

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

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

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

@@ -0,0 +1,24 @@
+"""$project_name: A Flower / $framework_str app."""
+
+from flwr.common import Context
+from flwr.server import ServerApp, ServerAppComponents, ServerConfig
+from flwr.server.strategy import FedAvg
+
+
+def server_fn(context: Context):
+    # Read from config
+    num_rounds = context.run_config["num-server-rounds"]
+
+    # Define strategy
+    strategy = FedAvg(
+        fraction_fit=1.0,
+        fraction_evaluate=1.0,
+        min_available_clients=2,
+    )
+    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.tensorflow.py.tpl

@@ -1 +1,29 @@
-"""$project_name: A Flower / TensorFlow app."""
+"""$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 load_model
+
+
+def server_fn(context: Context):
+    # Read from config
+    num_rounds = context.run_config["num-server-rounds"]
+
+    # Get parameters to initialize global model
+    parameters = ndarrays_to_parameters(load_model().get_weights())
+
+    # Define strategy
+    strategy = strategy = FedAvg(
+        fraction_fit=1.0,
+        fraction_evaluate=1.0,
+        min_available_clients=2,
+        initial_parameters=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/task.huggingface.py.tpl

@@ -0,0 +1,99 @@
+"""$project_name: A Flower / $framework_str app."""
+
+import warnings
+from collections import OrderedDict
+
+import torch
+from evaluate import load as load_metric
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer, DataCollatorWithPadding
+
+from flwr_datasets import FederatedDataset
+from flwr_datasets.partitioner import IidPartitioner
+
+
+warnings.filterwarnings("ignore", category=UserWarning)
+DEVICE = torch.device("cpu")
+CHECKPOINT = "distilbert-base-uncased"  # transformer model checkpoint
+
+
+fds = None  # Cache FederatedDataset
+
+
+def load_data(partition_id: int, num_partitions: int):
+    """Load IMDB data (training and eval)"""
+    # Only initialize `FederatedDataset` once
+    global fds
+    if fds is None:
+        partitioner = IidPartitioner(num_partitions=num_partitions)
+        fds = FederatedDataset(
+            dataset="stanfordnlp/imdb",
+            partitioners={"train": partitioner},
+        )
+    partition = fds.load_partition(partition_id)
+    # Divide data: 80% train, 20% test
+    partition_train_test = partition.train_test_split(test_size=0.2, seed=42)
+
+    tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
+
+    def tokenize_function(examples):
+        return tokenizer(examples["text"], truncation=True)
+
+    partition_train_test = partition_train_test.map(tokenize_function, batched=True)
+    partition_train_test = partition_train_test.remove_columns("text")
+    partition_train_test = partition_train_test.rename_column("label", "labels")
+
+    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+    trainloader = DataLoader(
+        partition_train_test["train"],
+        shuffle=True,
+        batch_size=32,
+        collate_fn=data_collator,
+    )
+
+    testloader = DataLoader(
+        partition_train_test["test"], batch_size=32, collate_fn=data_collator
+    )
+
+    return trainloader, testloader
+
+
+def train(net, trainloader, epochs):
+    optimizer = AdamW(net.parameters(), lr=5e-5)
+    net.train()
+    for _ in range(epochs):
+        for batch in trainloader:
+            batch = {k: v.to(DEVICE) for k, v in batch.items()}
+            outputs = net(**batch)
+            loss = outputs.loss
+            loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+
+
+def test(net, testloader):
+    metric = load_metric("accuracy")
+    loss = 0
+    net.eval()
+    for batch in testloader:
+        batch = {k: v.to(DEVICE) for k, v in batch.items()}
+        with torch.no_grad():
+            outputs = net(**batch)
+        logits = outputs.logits
+        loss += outputs.loss.item()
+        predictions = torch.argmax(logits, dim=-1)
+        metric.add_batch(predictions=predictions, references=batch["labels"])
+    loss /= len(testloader.dataset)
+    accuracy = metric.compute()["accuracy"]
+    return loss, accuracy
+
+
+def get_weights(net):
+    return [val.cpu().numpy() for _, val in net.state_dict().items()]
+
+
+def set_weights(net, parameters):
+    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/task.jax.py.tpl

@@ -0,0 +1,57 @@
+"""$project_name: A Flower / $framework_str app."""
+
+import jax
+import jax.numpy as jnp
+from sklearn.datasets import make_regression
+from sklearn.model_selection import train_test_split
+import numpy as np
+
+key = jax.random.PRNGKey(0)
+
+
+def load_data():
+    # Load dataset
+    X, y = make_regression(n_features=3, random_state=0)
+    X, X_test, y, y_test = train_test_split(X, y)
+    return X, y, X_test, y_test
+
+
+def load_model(model_shape):
+    # Extract model parameters
+    params = {"b": jax.random.uniform(key), "w": jax.random.uniform(key, model_shape)}
+    return params
+
+
+def loss_fn(params, X, y):
+    # Return MSE as loss
+    err = jnp.dot(X, params["w"]) + params["b"] - y
+    return jnp.mean(jnp.square(err))
+
+
+def train(params, grad_fn, X, y):
+    loss = 1_000_000
+    num_examples = X.shape[0]
+    for epochs in range(50):
+        grads = grad_fn(params, X, y)
+        params = jax.tree_map(lambda p, g: p - 0.05 * g, params, grads)
+        loss = loss_fn(params, X, y)
+    return params, loss, num_examples
+
+
+def evaluation(params, grad_fn, X_test, y_test):
+    num_examples = X_test.shape[0]
+    err_test = loss_fn(params, X_test, y_test)
+    loss_test = jnp.mean(jnp.square(err_test))
+    return loss_test, num_examples
+
+
+def get_params(params):
+    parameters = []
+    for _, val in params.items():
+        parameters.append(np.array(val))
+    return parameters
+
+
+def set_params(local_params, global_params):
+    for key, value in list(zip(local_params.keys(), global_params)):
+        local_params[key] = value

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

@@ -0,0 +1,102 @@
+"""$project_name: A Flower / $framework_str app."""
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+from datasets.utils.logging import disable_progress_bar
+from flwr_datasets import FederatedDataset
+from flwr_datasets.partitioner import IidPartitioner
+
+
+disable_progress_bar()
+
+
+class MLP(nn.Module):
+    """A simple MLP."""
+
+    def __init__(
+        self, num_layers: int, input_dim: int, hidden_dim: int, output_dim: int
+    ):
+        super().__init__()
+        layer_sizes = [input_dim] + [hidden_dim] * num_layers + [output_dim]
+        self.layers = [
+            nn.Linear(idim, odim)
+            for idim, odim in zip(layer_sizes[:-1], layer_sizes[1:])
+        ]
+
+    def __call__(self, x):
+        for l in self.layers[:-1]:
+            x = mx.maximum(l(x), 0.0)
+        return self.layers[-1](x)
+
+
+def loss_fn(model, X, y):
+    return mx.mean(nn.losses.cross_entropy(model(X), y))
+
+
+def eval_fn(model, X, y):
+    return mx.mean(mx.argmax(model(X), axis=1) == y)
+
+
+def batch_iterate(batch_size, X, y):
+    perm = mx.array(np.random.permutation(y.size))
+    for s in range(0, y.size, batch_size):
+        ids = perm[s : s + batch_size]
+        yield X[ids], y[ids]
+
+
+fds = None  # Cache FederatedDataset
+
+
+def load_data(partition_id: int, num_partitions: int):
+    # Only initialize `FederatedDataset` once
+    global fds
+    if fds is None:
+        partitioner = IidPartitioner(num_partitions=num_partitions)
+        fds = FederatedDataset(
+            dataset="ylecun/mnist",
+            partitioners={"train": partitioner},
+            trust_remote_code=True,
+        )
+    partition = fds.load_partition(partition_id)
+    partition_splits = partition.train_test_split(test_size=0.2, seed=42)
+
+    partition_splits["train"].set_format("numpy")
+    partition_splits["test"].set_format("numpy")
+
+    train_partition = partition_splits["train"].map(
+        lambda img: {
+            "img": img.reshape(-1, 28 * 28).squeeze().astype(np.float32) / 255.0
+        },
+        input_columns="image",
+    )
+    test_partition = partition_splits["test"].map(
+        lambda img: {
+            "img": img.reshape(-1, 28 * 28).squeeze().astype(np.float32) / 255.0
+        },
+        input_columns="image",
+    )
+
+    data = (
+        train_partition["img"],
+        train_partition["label"].astype(np.uint32),
+        test_partition["img"],
+        test_partition["label"].astype(np.uint32),
+    )
+
+    train_images, train_labels, test_images, test_labels = map(mx.array, data)
+    return train_images, train_labels, test_images, test_labels
+
+
+def get_params(model):
+    layers = model.parameters()["layers"]
+    return [np.array(val) for layer in layers for _, val in layer.items()]
+
+
+def set_params(model, parameters):
+    new_params = {}
+    new_params["layers"] = [
+        {"weight": mx.array(parameters[i]), "bias": mx.array(parameters[i + 1])}
+        for i in range(0, len(parameters), 2)
+    ]
+    model.update(new_params)

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

@@ -1,4 +1,4 @@
-"""$project_name: A Flower / PyTorch app."""
+"""$project_name: A Flower / $framework_str app."""
 
 from collections import OrderedDict
 
@@ -6,11 +6,9 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
-from torchvision.datasets import CIFAR10
 from torchvision.transforms import Compose, Normalize, ToTensor
 from flwr_datasets import FederatedDataset
-
-DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+from flwr_datasets.partitioner import IidPartitioner
 
 
 class Net(nn.Module):
@@ -34,12 +32,22 @@ class Net(nn.Module):
         return self.fc3(x)
 
 
-def load_data(partition_id, num_partitions):
+fds = None  # Cache FederatedDataset
+
+
+def load_data(partition_id: int, num_partitions: int):
     """Load partition CIFAR10 data."""
-    fds = FederatedDataset(dataset="cifar10", partitioners={"train": num_partitions})
+    # Only initialize `FederatedDataset` once
+    global fds
+    if fds is None:
+        partitioner = IidPartitioner(num_partitions=num_partitions)
+        fds = FederatedDataset(
+            dataset="uoft-cs/cifar10",
+            partitioners={"train": partitioner},
+        )
     partition = fds.load_partition(partition_id)
     # Divide data on each node: 80% train, 20% test
-    partition_train_test = partition.train_test_split(test_size=0.2)
+    partition_train_test = partition.train_test_split(test_size=0.2, seed=42)
     pytorch_transforms = Compose(
         [ToTensor(), Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
     )
@@ -55,44 +63,41 @@ def load_data(partition_id, num_partitions):
     return trainloader, testloader
 
 
-def train(net, trainloader, valloader, epochs, device):
+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().to(device)
-    optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
+    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()
-            criterion(net(images.to(DEVICE)), labels.to(DEVICE)).backward()
+            loss = criterion(net(images.to(device)), labels.to(device))
+            loss.backward()
             optimizer.step()
+            running_loss += loss.item()
 
-    train_loss, train_acc = test(net, trainloader)
-    val_loss, val_acc = test(net, valloader)
-
-    results = {
-        "train_loss": train_loss,
-        "train_accuracy": train_acc,
-        "val_loss": val_loss,
-        "val_accuracy": val_acc,
-    }
-    return results
+    avg_trainloss = running_loss / len(trainloader)
+    return avg_trainloss
 
 
-def test(net, testloader):
+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)
+            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