flwr-nightly 1.10.0.dev20240612__py3-none-any.whl → 1.10.0.dev20240624__py3-none-any.whl

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

@@ -0,0 +1,124 @@
+"""$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.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[[str], FlowerClient]:  # pylint: disable=too-many-arguments
+    """Generate the client function that creates the Flower Clients."""
+
+    def client_fn(cid: str) -> FlowerClient:
+        """Create a Flower client representing a single organization."""
+        # Let's get the partition corresponding to the i-th client
+        client_trainset = fds.load_partition(int(cid), "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

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 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/pyproject.flowertune.toml.tpl

@@ -0,0 +1,42 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "$package_name"
+version = "1.0.0"
+description = ""
+authors = [
+    { name = "The Flower Authors", email = "hello@flower.ai" },
+]
+license = { text = "Apache License (2.0)" }
+dependencies = [
+    "flwr[simulation]>=1.9.0,<2.0",
+    "flwr-datasets>=0.1.0,<1.0.0",
+    "hydra-core==1.3.2",
+    "trl==0.8.1",
+    "bitsandbytes==0.43.0",
+    "scipy==1.13.0",
+    "peft==0.6.2",
+    "transformers==4.39.3",
+    "sentencepiece==0.2.0",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["."]
+
+[flower]
+publisher = "$username"
+
+[flower.components]
+serverapp = "$import_name.app:server"
+clientapp = "$import_name.app:client"
+
+[flower.engine]
+name = "simulation"
+
+[flower.engine.simulation.supernode]
+num = $num_clients
+
+[flower.engine.simulation]
+backend_config = { client_resources = { num_cpus = 8, num_gpus = 1.0 } }

flwr/cli/run/run.py

@@ -16,12 +16,18 @@
 
 import sys
 from enum import Enum
+from logging import DEBUG
 from typing import Optional
 
 import typer
 from typing_extensions import Annotated
 
 from flwr.cli import config_utils
+from flwr.common.constant import SUPEREXEC_DEFAULT_ADDRESS
+from flwr.common.grpc import GRPC_MAX_MESSAGE_LENGTH, create_channel
+from flwr.common.logger import log
+from flwr.proto.exec_pb2 import StartRunRequest  # pylint: disable=E0611
+from flwr.proto.exec_pb2_grpc import ExecStub
 from flwr.simulation.run_simulation import _run_simulation
 
 
@@ -31,20 +37,35 @@ class Engine(str, Enum):
     SIMULATION = "simulation"
 
 
+# pylint: disable-next=too-many-locals
 def run(
     engine: Annotated[
         Optional[Engine],
-        typer.Option(case_sensitive=False, help="The ML framework to use"),
+        typer.Option(
+            case_sensitive=False,
+            help="The engine to run FL with (currently only simulation is supported).",
+        ),
     ] = None,
+    use_superexec: Annotated[
+        bool,
+        typer.Option(
+            case_sensitive=False, help="Use this flag to use the new SuperExec API"
+        ),
+    ] = False,
 ) -> None:
     """Run Flower project."""
+    if use_superexec:
+        _start_superexec_run()
+        return
+
     typer.secho("Loading project configuration... ", fg=typer.colors.BLUE)
 
     config, errors, warnings = config_utils.load_and_validate()
 
     if config is None:
         typer.secho(
-            "Project configuration could not be loaded.\npyproject.toml is invalid:\n"
+            "Project configuration could not be loaded.\n"
+            "pyproject.toml is invalid:\n"
             + "\n".join([f"- {line}" for line in errors]),
             fg=typer.colors.RED,
             bold=True,
@@ -69,12 +90,16 @@ def run(
 
     if engine == Engine.SIMULATION:
         num_supernodes = config["flower"]["engine"]["simulation"]["supernode"]["num"]
+        backend_config = config["flower"]["engine"]["simulation"].get(
+            "backend_config", None
+        )
 
         typer.secho("Starting run... ", fg=typer.colors.BLUE)
         _run_simulation(
             server_app_attr=server_app_ref,
             client_app_attr=client_app_ref,
             num_supernodes=num_supernodes,
+            backend_config=backend_config,
         )
     else:
         typer.secho(
@@ -82,3 +107,22 @@ def run(
             fg=typer.colors.RED,
             bold=True,
         )
+
+
+def _start_superexec_run() -> None:
+    def on_channel_state_change(channel_connectivity: str) -> None:
+        """Log channel connectivity."""
+        log(DEBUG, channel_connectivity)
+
+    channel = create_channel(
+        server_address=SUPEREXEC_DEFAULT_ADDRESS,
+        insecure=True,
+        root_certificates=None,
+        max_message_length=GRPC_MAX_MESSAGE_LENGTH,
+        interceptors=None,
+    )
+    channel.subscribe(on_channel_state_change)
+    stub = ExecStub(channel)
+
+    req = StartRunRequest()
+    stub.StartRun(req)

flwr/client/__init__.py

@@ -28,8 +28,8 @@ __all__ = [
     "Client",
     "ClientApp",
     "ClientFn",
-    "mod",
     "NumPyClient",
+    "mod",
     "run_client_app",
     "run_supernode",
     "start_client",

flwr/client/app.py

@@ -19,7 +19,7 @@ import sys
 import time
 from dataclasses import dataclass
 from logging import DEBUG, ERROR, INFO, WARN
-from typing import Callable, ContextManager, Optional, Tuple, Type, Union
+from typing import Callable, ContextManager, Dict, Optional, Tuple, Type, Union
 
 from cryptography.hazmat.primitives.asymmetric import ec
 from grpc import RpcError
@@ -31,6 +31,7 @@ from flwr.common import GRPC_MAX_MESSAGE_LENGTH, EventType, Message, event
 from flwr.common.address import parse_address
 from flwr.common.constant import (
     MISSING_EXTRA_REST,
+    TRANSPORT_TYPE_GRPC_ADAPTER,
     TRANSPORT_TYPE_GRPC_BIDI,
     TRANSPORT_TYPE_GRPC_RERE,
     TRANSPORT_TYPE_REST,
@@ -41,6 +42,7 @@ from flwr.common.logger import log, warn_deprecated_feature
 from flwr.common.message import Error
 from flwr.common.retry_invoker import RetryInvoker, RetryState, exponential
 
+from .grpc_adapter_client.connection import grpc_adapter
 from .grpc_client.connection import grpc_connection
 from .grpc_rere_client.connection import grpc_request_response
 from .message_handler.message_handler import handle_control_message
@@ -177,7 +179,7 @@ def start_client(
 def _start_client_internal(
     *,
     server_address: str,
-    load_client_app_fn: Optional[Callable[[], ClientApp]] = None,
+    load_client_app_fn: Optional[Callable[[str, str], ClientApp]] = None,
     client_fn: Optional[ClientFn] = None,
     client: Optional[Client] = None,
     grpc_max_message_length: int = GRPC_MAX_MESSAGE_LENGTH,
@@ -252,7 +254,7 @@ def _start_client_internal(
 
             client_fn = single_client_factory
 
-        def _load_client_app() -> ClientApp:
+        def _load_client_app(_1: str, _2: str) -> ClientApp:
             return ClientApp(client_fn=client_fn)
 
         load_client_app_fn = _load_client_app
@@ -308,6 +310,8 @@ def _start_client_internal(
     )
 
     node_state = NodeState()
+    # run_id -> (fab_id, fab_version)
+    run_info: Dict[int, Tuple[str, str]] = {}
 
     while not app_state_tracker.interrupt:
         sleep_duration: int = 0
@@ -319,7 +323,6 @@ def _start_client_internal(
             root_certificates,
             authentication_keys,
         ) as conn:
-            # pylint: disable-next=W0612
             receive, send, create_node, delete_node, get_run = conn
 
             # Register node
@@ -356,13 +359,20 @@ def _start_client_internal(
                         send(out_message)
                         break
 
+                    # Get run info
+                    run_id = message.metadata.run_id
+                    if run_id not in run_info:
+                        if get_run is not None:
+                            run_info[run_id] = get_run(run_id)
+                        # If get_run is None, i.e., in grpc-bidi mode
+                        else:
+                            run_info[run_id] = ("", "")
+
                     # Register context for this run
-                    node_state.register_context(run_id=message.metadata.run_id)
+                    node_state.register_context(run_id=run_id)
 
                     # Retrieve context for this run
-                    context = node_state.retrieve_context(
-                        run_id=message.metadata.run_id
-                    )
+                    context = node_state.retrieve_context(run_id=run_id)
 
                     # Create an error reply message that will never be used to prevent
                     # the used-before-assignment linting error
@@ -373,7 +383,7 @@ def _start_client_internal(
                     # Handle app loading and task message
                     try:
                         # Load ClientApp instance
-                        client_app: ClientApp = load_client_app_fn()
+                        client_app: ClientApp = load_client_app_fn(*run_info[run_id])
 
                         # Execute ClientApp
                         reply_message = client_app(message=message, context=context)
@@ -411,7 +421,7 @@ def _start_client_internal(
                     else:
                         # No exception, update node state
                         node_state.update_context(
-                            run_id=message.metadata.run_id,
+                            run_id=run_id,
                             context=context,
                         )
 
@@ -592,6 +602,8 @@ def _init_connection(transport: Optional[str], server_address: str) -> Tuple[
         connection, error_type = http_request_response, RequestsConnectionError
     elif transport == TRANSPORT_TYPE_GRPC_RERE:
         connection, error_type = grpc_request_response, RpcError
+    elif transport == TRANSPORT_TYPE_GRPC_ADAPTER:
+        connection, error_type = grpc_adapter, RpcError
     elif transport == TRANSPORT_TYPE_GRPC_BIDI:
         connection, error_type = grpc_connection, RpcError
     else:

flwr/client/client_app.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Flower Labs GmbH. All Rights Reserved.
+# 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.

flwr/client/dpfedavg_numpy_client.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2022 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.

flwr/client/grpc_adapter_client/__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.
+# ==============================================================================
+"""Client-side part of the GrpcAdapter transport layer."""