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

flwr/cli/app.py

@@ -15,6 +15,7 @@
 """Flower command line interface."""
 
 import typer
+from typer.main import get_command
 
 from .build import build
 from .example import example
@@ -37,5 +38,7 @@ app.command()(run)
 app.command()(build)
 app.command()(install)
 
+typer_click_object = get_command(app)
+
 if __name__ == "__main__":
     app()

flwr/cli/build.py

@@ -36,13 +36,9 @@ def build(
 ) -> str:
     """Build a Flower project into a Flower App Bundle (FAB).
 
-    You can run `flwr build` without any argument to bundle the current directory:
-
-        `flwr build`
-
-    You can also build a specific directory:
-
-        `flwr build --directory ./projects/flower-hello-world`
+    You can run ``flwr build`` without any arguments to bundle the current directory,
+    or you can use ``--directory`` to build a specific directory:
+    ``flwr build --directory ./projects/flower-hello-world``.
     """
     if directory is None:
         directory = Path.cwd()

flwr/cli/new/new.py

@@ -41,6 +41,16 @@ class MlFramework(str, Enum):
     HUGGINGFACE = "HF"
     MLX = "MLX"
     SKLEARN = "sklearn"
+    FLOWERTUNE = "FlowerTune"
+
+
+class LlmChallengeName(str, Enum):
+    """Available LLM challenges."""
+
+    GENERALNLP = "GeneralNLP"
+    FINANCE = "Finance"
+    MEDICAL = "Medical"
+    CODE = "Code"
 
 
 class TemplateNotFound(Exception):
@@ -81,6 +91,7 @@ def render_and_create(file_path: str, template: str, context: Dict[str, str]) ->
     create_file(file_path, content)
 
 
+# pylint: disable=too-many-locals,too-many-branches,too-many-statements
 def new(
     project_name: Annotated[
         Optional[str],
@@ -125,6 +136,19 @@ def new(
 
     framework_str = framework_str.lower()
 
+    if framework_str == "flowertune":
+        llm_challenge_value = prompt_options(
+            "Please select LLM challenge by typing in the number",
+            sorted([challenge.value for challenge in LlmChallengeName]),
+        )
+        selected_value = [
+            name
+            for name, value in vars(LlmChallengeName).items()
+            if value == llm_challenge_value
+        ]
+        llm_challenge_str = selected_value[0]
+        llm_challenge_str = llm_challenge_str.lower()
+
     print(
         typer.style(
             f"\n🔨 Creating Flower project {project_name}...",
@@ -139,33 +163,6 @@ def new(
     import_name = package_name.replace("-", "_")
     project_dir = os.path.join(cwd, package_name)
 
-    # List of files to render
-    files = {
-        ".gitignore": {"template": "app/.gitignore.tpl"},
-        "README.md": {"template": "app/README.md.tpl"},
-        "pyproject.toml": {"template": f"app/pyproject.{framework_str}.toml.tpl"},
-        f"{import_name}/__init__.py": {"template": "app/code/__init__.py.tpl"},
-        f"{import_name}/server.py": {
-            "template": f"app/code/server.{framework_str}.py.tpl"
-        },
-        f"{import_name}/client.py": {
-            "template": f"app/code/client.{framework_str}.py.tpl"
-        },
-    }
-
-    # Depending on the framework, generate task.py file
-    frameworks_with_tasks = [
-        MlFramework.PYTORCH.value.lower(),
-        MlFramework.JAX.value.lower(),
-        MlFramework.HUGGINGFACE.value.lower(),
-        MlFramework.MLX.value.lower(),
-        MlFramework.TENSORFLOW.value.lower(),
-    ]
-    if framework_str in frameworks_with_tasks:
-        files[f"{import_name}/task.py"] = {
-            "template": f"app/code/task.{framework_str}.py.tpl"
-        }
-
     context = {
         "project_name": project_name,
         "package_name": package_name,
@@ -173,6 +170,85 @@ def new(
         "username": username,
     }
 
+    # List of files to render
+    if framework_str == "flowertune":
+        files = {
+            ".gitignore": {"template": "app/.gitignore.tpl"},
+            "pyproject.toml": {"template": f"app/pyproject.{framework_str}.toml.tpl"},
+            "README.md": {"template": f"app/README.{framework_str}.md.tpl"},
+            f"{import_name}/__init__.py": {"template": "app/code/__init__.py.tpl"},
+            f"{import_name}/server.py": {
+                "template": "app/code/flwr_tune/server.py.tpl"
+            },
+            f"{import_name}/client.py": {
+                "template": "app/code/flwr_tune/client.py.tpl"
+            },
+            f"{import_name}/app.py": {"template": "app/code/flwr_tune/app.py.tpl"},
+            f"{import_name}/models.py": {
+                "template": "app/code/flwr_tune/models.py.tpl"
+            },
+            f"{import_name}/dataset.py": {
+                "template": "app/code/flwr_tune/dataset.py.tpl"
+            },
+            f"{import_name}/conf/config.yaml": {
+                "template": "app/code/flwr_tune/config.yaml.tpl"
+            },
+            f"{import_name}/conf/static_config.yaml": {
+                "template": "app/code/flwr_tune/static_config.yaml.tpl"
+            },
+        }
+
+        # Challenge specific context
+        fraction_fit = "0.2" if llm_challenge_str == "code" else "0.1"
+        if llm_challenge_str == "generalnlp":
+            challenge_name = "General NLP"
+            num_clients = "20"
+            dataset_name = "vicgalle/alpaca-gpt4"
+        elif llm_challenge_str == "finance":
+            challenge_name = "Finance"
+            num_clients = "50"
+            dataset_name = "FinGPT/fingpt-sentiment-train"
+        elif llm_challenge_str == "medical":
+            challenge_name = "Medical"
+            num_clients = "20"
+            dataset_name = "medalpaca/medical_meadow_medical_flashcards"
+        else:
+            challenge_name = "Code"
+            num_clients = "10"
+            dataset_name = "lucasmccabe-lmi/CodeAlpaca-20k"
+
+        context["llm_challenge_str"] = llm_challenge_str
+        context["fraction_fit"] = fraction_fit
+        context["challenge_name"] = challenge_name
+        context["num_clients"] = num_clients
+        context["dataset_name"] = dataset_name
+    else:
+        files = {
+            ".gitignore": {"template": "app/.gitignore.tpl"},
+            "README.md": {"template": "app/README.md.tpl"},
+            "pyproject.toml": {"template": f"app/pyproject.{framework_str}.toml.tpl"},
+            f"{import_name}/__init__.py": {"template": "app/code/__init__.py.tpl"},
+            f"{import_name}/server.py": {
+                "template": f"app/code/server.{framework_str}.py.tpl"
+            },
+            f"{import_name}/client.py": {
+                "template": f"app/code/client.{framework_str}.py.tpl"
+            },
+        }
+
+        # Depending on the framework, generate task.py file
+        frameworks_with_tasks = [
+            MlFramework.PYTORCH.value.lower(),
+            MlFramework.JAX.value.lower(),
+            MlFramework.HUGGINGFACE.value.lower(),
+            MlFramework.MLX.value.lower(),
+            MlFramework.TENSORFLOW.value.lower(),
+        ]
+        if framework_str in frameworks_with_tasks:
+            files[f"{import_name}/task.py"] = {
+                "template": f"app/code/task.{framework_str}.py.tpl"
+            }
+
     for file_path, value in files.items():
         render_and_create(
             file_path=os.path.join(project_dir, file_path),
@@ -190,7 +266,7 @@ def new(
     )
     print(
         typer.style(
-            f"	cd {project_name}\n" + "	pip install -e .\n	flwr run\n",
+            f"	cd {package_name}\n" + "	pip install -e .\n	flwr run\n",
             fg=typer.colors.BRIGHT_CYAN,
             bold=True,
         )

flwr/cli/new/templates/app/README.flowertune.md.tpl

@@ -0,0 +1,56 @@
+# FlowerTune LLM on $challenge_name Dataset
+
+This directory conducts federated instruction tuning with a pretrained [Mistral-7B](https://huggingface.co/mistralai/Mistral-7B-v0.3) model on a [$challenge_name dataset](https://huggingface.co/datasets/$dataset_name).
+We use [Flower Datasets](https://flower.dev/docs/datasets/) to download, partition and preprocess the dataset.
+Flower's Simulation Engine is used to simulate the LLM fine-tuning process in federated way,
+which allows users to perform the training on a single GPU.
+
+
+## Methodology
+
+This baseline performs federated LLM fine-tuning with [LoRA](https://arxiv.org/pdf/2106.09685) using the [🤗PEFT](https://huggingface.co/docs/peft/en/index) library.
+The clients' models are aggregated with FedAvg strategy.
+This provides a baseline performance for the leaderboard of $challenge_name challenge.
+
+
+## Environments setup
+
+Project dependencies are defined in `pyproject.toml`. Install them in an activated Python environment with:
+
+```shell
+pip install -e .
+```
+
+## Experimental setup
+
+The dataset is partitioned into $num_clients shards with IID fashion serving as clients.
+We randomly sample $fraction_fit clients to be available for each round,
+and the federated fine-tuning lasts for `200` rounds.
+All settings are defined in `$project_name/conf/static_config.yaml`, which is not allowed to be modified for fair competition if you plan to participated in the [LLM leaderboard](https://flower.ai/benchmarks/llm-leaderboard).
+
+
+## Running the challenge
+
+First make sure that you have got the access to [Mistral-7B](https://huggingface.co/mistralai/Mistral-7B-v0.3) model with your Hugging-Face account. You can request access directly from the Hugging-Face website.
+Then, follow the instruction [here](https://huggingface.co/docs/huggingface_hub/en/quick-start#login-command) to log in your account. Note you only need to complete this stage once in your development machine:
+
+```bash
+huggingface-cli login
+```
+
+Run the challenge with default config values.
+The configs are in `$project_name/conf/config.yaml` and `$project_name/conf/static_config.yaml`, and are loaded automatically.
+
+```bash
+flwr run
+```
+
+## VRAM consumption
+
+We use Mistral-7B model with 4-bit quantization as default. The estimated VRAM consumption per client for each challenge is shown below:
+
+| Challenges | GeneralNLP |   Finance  |   Medical  |    Code    |
+| :--------: | :--------: | :--------: | :--------: | :--------: |
+|    VRAM    | ~25.50 GB  | ~17.30 GB  | ~22.80 GB  | ~17.40 GB  |
+
+You can adjust the CPU/GPU resources you assign to each of the clients based on your device, which is specified with `flower.engine.simulation` in `pyproject.toml`.

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,86 @@
+"""$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 ndarrays_to_parameters
+from flwr.server import ServerApp, ServerConfig
+
+from $import_name.client 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 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)
+
+# 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
+    ),
+)
+
+# ServerApp for Flower Next
+server = ServerApp(
+    config=ServerConfig(num_rounds=cfg_static.num_rounds),
+    strategy=strategy,
+)

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)