examples/secure_aggregation and examples/split_learning passed ty checks.

Author: baochunli

docs/docs/index.md

@@ -9,7 +9,6 @@ Welcome to *Plato*, a software framework to facilitate scalable, reproducible, a
 
     - **[Installation](install.md)** - Installing Plato and setting up your development environment
     - **[Quick Start](quickstart.md)** - Getting started with Plato
-    - Plato supports both PyTorch and MLX backends (MLX for Apple Silicon devices)
 
 ## Examples
 

docs/docs/quickstart.md

@@ -157,7 +157,7 @@ To fix all linter errors automatically, run:
 uvx ruff check --fix
 ```
 
-# Type Checking
+## Type Checking
 
 It is also strongly recommended that new additions and revisions of the code base to pass Astral's [ty](https://docs.astral.sh/ty/) type checker cleanly. To install `ty` globally using `uv`, run:
 

examples/async/fedbuff/fedbuff_cifar10.toml

@@ -17,7 +17,7 @@ speed_simulation = true
 max_sleep_time = 30
 
 # Should clients really go to sleep, or should we just simulate the sleep times?
-sleep_simulation = false
+sleep_simulation = true 
 
 # If we are simulating client training times, what is the average training time?
 avg_training_time = 20

examples/model_pruning/fedscr/fedscr_server.py

@@ -4,7 +4,7 @@
 """
 
 import os
-from typing import Dict, Optional, TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Dict, Optional, cast
 
 import numpy as np
 
@@ -58,11 +58,7 @@ def __init__(
     def customize_server_response(self, server_response: dict, client_id) -> dict:
         """Wraps up generating the server response with any additional information."""
         trainer = cast(Optional["FedSCRTrainer"], self.trainer)
-        if (
-            trainer is not None
-            and trainer.use_adaptive
-            and self.current_round > 1
-        ):
+        if trainer is not None and trainer.use_adaptive and self.current_round > 1:
             self.calc_threshold()
             server_response["update_thresholds"] = self.update_thresholds
         return server_response
@@ -101,9 +97,7 @@ def weights_aggregated(self, updates):
                 float(np.var([update.report.loss for update in updates]))
             )
             if self.current_round > 3:
-                self.mean_variance = sum(self.loss_variances) / (
-                    self.current_round - 2
-                )
+                self.mean_variance = sum(self.loss_variances) / (self.current_round - 2)
             else:
                 self.mean_variance = 0.0
 

examples/secure_aggregation/maskcrypt/maskcrypt_client.py

@@ -50,15 +50,18 @@ def setup(self, context: ClientContext) -> None:
     def load_payload(self, context: ClientContext, server_payload: Any) -> None:
         """Store inbound payload or delegate weight loading based on the round."""
         state = self._state(context)
+        owner = getattr(context, "owner", None)
 
         if context.current_round % 2 != 0:
             state["final_mask"] = None
-            context.owner.final_mask = None
+            if owner is not None and hasattr(owner, "final_mask"):
+                setattr(owner, "final_mask", None)
             super().load_payload(context, server_payload)
             return
 
         state["final_mask"] = server_payload
-        context.owner.final_mask = server_payload
+        if owner is not None and hasattr(owner, "final_mask"):
+            setattr(owner, "final_mask", server_payload)
 
     async def train(self, context: ClientContext) -> tuple[Any, Any]:
         """Alternate between mask proposal computation and weight submission."""
@@ -136,12 +139,17 @@ def _compute_mask(
 class MaskCryptClientProxy(simple.Client):
     """Client variant exposing MaskCrypt state via a convenient property."""
 
+    encrypt_ratio: float
+    random_mask: bool
+    attack_prep_dir: str
+    checkpoint_path: str
+
     @property
-    def final_mask(self):
+    def final_mask(self) -> Any | None:
         return self._context.state.get("maskcrypt", {}).get("final_mask")
 
     @final_mask.setter
-    def final_mask(self, value):
+    def final_mask(self, value: Any | None) -> None:
         self._context.state.setdefault("maskcrypt", {})["final_mask"] = value
 
 

examples/secure_aggregation/maskcrypt/maskcrypt_server.py

@@ -2,6 +2,8 @@
 A MaskCrypt server with selective homomorphic encryption support.
 """
 
+from typing import cast
+
 from maskcrypt_algorithm import Algorithm as MaskCryptAlgorithm
 
 from plato.servers import fedavg_he
@@ -50,14 +52,21 @@ async def aggregate_weights(self, updates, baseline_weights, weights_received):
 
             return baseline_weights
         else:
-            # Clients send model updates in even rounds, conduct aggregation
-            aggregated_weights = await super().aggregate_weights(
-                updates, baseline_weights, weights_received
+            strategy = getattr(self, "aggregation_strategy", None)
+            if strategy is None or not hasattr(strategy, "aggregate_weights"):
+                raise AttributeError(
+                    "Aggregation strategy must expose an 'aggregate_weights' coroutine."
+                )
+            aggregated_weights = await strategy.aggregate_weights(
+                updates, baseline_weights, weights_received, self.context
             )
+            if aggregated_weights is None:
+                raise RuntimeError("Aggregation strategy failed to produce weights.")
 
             return aggregated_weights
 
     def _mask_consensus(self, updates):
         """Conduct mask consensus on the reported mask proposals."""
         proposals = [update.payload for update in updates]
-        self.final_mask = self.algorithm.build_consensus_mask(proposals)
+        algorithm = cast(MaskCryptAlgorithm, self.require_algorithm())
+        self.final_mask = algorithm.build_consensus_mask(proposals)

examples/split_learning/llm_split_learning/split_learning_llm_model.py

@@ -1,7 +1,9 @@
 """
-Obtain LLM models from HuggingFace, specifically designed for split learning
+Obtain LLM models from HuggingFace, specifically designed for split learning.
 """
 
+from typing import Any, Union, cast
+
 import torch
 from peft import LoraConfig, get_peft_model
 from transformers import AutoConfig, AutoModelForCausalLM
@@ -17,7 +19,7 @@ def get_lora_model(model):
     return model
 
 
-def get_module(start_module: torch.nn.Module, module_names):
+def get_module(start_module: torch.nn.Module, module_names) -> torch.nn.Module:
     """
     Recursively get a PyTorch module starting from the start module with
     a given list of module names.
@@ -28,6 +30,9 @@ def get_module(start_module: torch.nn.Module, module_names):
     return module
 
 
+TransformerSequence = Union[torch.nn.Sequential, torch.nn.ModuleList]
+
+
 class BaseModel(torch.nn.Module):
     """
     The basic model loading HuggingFace model used for the server model and the client model
@@ -47,14 +52,16 @@ def __init__(self, *args, **kwargs) -> None:
 
         self.config = AutoConfig.from_pretrained(self.model_name, **config_kwargs)
 
-        self.base_model = AutoModelForCausalLM.from_pretrained(
+        base_model = AutoModelForCausalLM.from_pretrained(
             self.model_name,
             config=self.config,
             cache_dir=Config().params["model_path"] + "/huggingface",
             token=use_auth_token,
         )
-        if hasattr(self.base_model, "loss_type"):
-            self.base_model.loss_type = "ForCausalLM"
+        base_model_for_loss = cast(Any, base_model)
+        if hasattr(base_model_for_loss, "loss_type"):
+            base_model_for_loss.loss_type = "ForCausalLM"
+        self.base_model = base_model
         self.cut_layer = Config().parameters.model.cut_layer
 
     def get_input_embeddings(self):
@@ -79,9 +86,11 @@ def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         # replace the layers in the base model
         # which should be on the cloud with Identity layers()
-        transformer_module = self.base_model
-        for module_name in Config().parameters.model.transformer_module_name.split("."):
-            transformer_module = getattr(transformer_module, module_name)
+        transformer_module_raw = get_module(
+            self.base_model,
+            Config().parameters.model.transformer_module_name.split("."),
+        )
+        transformer_module = cast(TransformerSequence, transformer_module_raw)
         client_layers = transformer_module[: self.cut_layer]
         client_module_names = Config().parameters.model.transformer_module_name.split(
             "."
@@ -126,18 +135,21 @@ def __init__(self, *args, **kwargs) -> None:
         # The first copy of the model is the whole model which is used for test.
         # The second copy of the model only contains the layers on the server
         # used for training.
-        self.server_model = AutoModelForCausalLM.from_pretrained(
+        server_model = AutoModelForCausalLM.from_pretrained(
             self.model_name,
             config=self.config,
             cache_dir=Config().params["model_path"] + "/huggingface",
         )
-        if hasattr(self.server_model, "loss_type"):
-            self.server_model.loss_type = "ForCausalLM"
+        server_model_for_loss = cast(Any, server_model)
+        if hasattr(server_model_for_loss, "loss_type"):
+            server_model_for_loss.loss_type = "ForCausalLM"
+        self.server_model = server_model
         transformer_module = get_module(
             self.base_model,
             Config().parameters.model.transformer_module_name.split("."),
         )
-        server_layers = transformer_module[self.cut_layer :]
+        transformer_sequence = cast(TransformerSequence, transformer_module)
+        server_layers = transformer_sequence[self.cut_layer :]
         server_module_names = Config().parameters.model.transformer_module_name.split(
             "."
         )
@@ -159,9 +171,8 @@ def copy_weight(self):
         base_model_weights = self.base_model.state_dict()
         server_model_weights = self.server_model.state_dict()
 
-        transformer_module = self.base_model
-        for module_name in basic_name.split("."):
-            transformer_module = getattr(transformer_module, module_name)
+        transformer_module_raw = get_module(self.base_model, basic_name.split("."))
+        transformer_module = cast(TransformerSequence, transformer_module_raw)
         layer_names = [
             basic_name + "." + str(index)
             for index in range(

examples/split_learning/llm_split_learning/split_learning_lora_algorithm.py

@@ -2,10 +2,14 @@
 A split learning algorithm supporting LoRA fine-tuning LLMs.
 """
 
+from typing import Dict, cast
+
 from peft import (
     get_peft_model_state_dict,
     set_peft_model_state_dict,
 )
+from torch import Tensor
+from torch.nn import Module
 
 from plato.algorithms import split_learning
 
@@ -15,13 +19,22 @@ class Algorithm(split_learning.Algorithm):
     Extract and load only the LoRA weights.
     """
 
-    def extract_weights(self, model=None):
-        # Extract LoRA wegiths
-        return {
-            k: v.cpu()
-            for k, v in get_peft_model_state_dict(self.model.base_model).items()
-        }
+    def _get_base_model(self, model: object | None = None) -> Module:
+        """Return the wrapped HuggingFace base model."""
+        model_obj = model if model is not None else self.model
+        if model_obj is None or not hasattr(model_obj, "base_model"):
+            raise AttributeError(
+                "LoRA split learning requires a model with a `base_model` attribute."
+            )
+        base_model = getattr(model_obj, "base_model")
+        return cast(Module, base_model)
+
+    def extract_weights(self, model=None) -> Dict[str, Tensor]:
+        """Extract LoRA weights from the underlying base model."""
+        base_model = self._get_base_model(model)
+        return {k: v.cpu() for k, v in get_peft_model_state_dict(base_model).items()}
 
-    def load_weights(self, weights):
-        # Load LoRA weights
-        return set_peft_model_state_dict(self.model.base_model, weights)
+    def load_weights(self, weights: Dict[str, Tensor]):
+        """Load LoRA weights into the underlying base model."""
+        base_model = self._get_base_model()
+        return set_peft_model_state_dict(base_model, weights)

examples/split_learning/llm_split_learning/split_learning_server_attack.py

@@ -42,7 +42,15 @@ def attack(self, update):
         """
         self.attack_started = True
         intermediate_features, labels = update[0]
-        evaluation_metrics = self.trainer.attack(intermediate_features, labels)
+        trainer = self.trainer
+        if trainer is None or not hasattr(trainer, "attack"):
+            raise AttributeError(
+                "Trainer must define an `attack` method for curious server attacks."
+            )
+        attack_fn = getattr(trainer, "attack")
+        if not callable(attack_fn):
+            raise TypeError("Trainer attack must be callable.")
+        evaluation_metrics = attack_fn(intermediate_features, labels)
         rouge_metrics = evaluation_metrics["ROUGE"]
         self.rouge["rouge1_fm"] = rouge_metrics["rouge1_fmeasure"].item()
         self.rouge["rouge1_p"] = rouge_metrics["rouge1_precision"].item()

examples/split_learning/llm_split_learning/split_learning_trainer.py

@@ -6,7 +6,7 @@
 """
 
 from collections import OrderedDict
-from typing import Optional
+from typing import Any, Optional, Sized, cast
 
 import evaluate
 from torch import Tensor, reshape
@@ -40,13 +40,14 @@ def preprocess_logits_for_metrics(logits, labels):
 
 def compute_metrics(eval_preds):
     """Calculate the accuracy for evaluation stage."""
-    metric = evaluate.load("accuracy")
+    metric: Any = evaluate.load("accuracy")
     preds, labels = eval_preds
     # preds have the same shape as the labels, after the argmax(-1) has been calculated
     # by preprocess_logits_for_metrics but we need to shift the labels
     labels = labels.reshape(-1)
     preds = preds.reshape(-1)
-    return metric.compute(predictions=preds, references=labels)
+    compute_fn = getattr(metric, "compute")
+    return compute_fn(predictions=preds, references=labels)
 
 
 # ============================================================================
@@ -88,7 +89,10 @@ def __init__(
     def _get_train_sampler(self) -> Sampler | None:
         """Get training sampler."""
         if self.sampler is None:
-            return RandomSampler(self.train_dataset)
+            if self.train_dataset is None:
+                raise ValueError("Training dataset is not initialized.")
+            dataset = cast(Sized, self.train_dataset)
+            return RandomSampler(dataset)
         return self.sampler
 
     def _get_eval_sampler(self, eval_dataset) -> Sampler | None:
@@ -334,6 +338,8 @@ def server_forward_from(self, batch, config):
         inputs, labels = batch
         batch_size = inputs.size(0)
         inputs = inputs.detach().requires_grad_(True)
+        if self.model is None or not hasattr(self.model, "forward_from"):
+            raise AttributeError("Model must provide a `forward_from` method.")
         outputs = self.model.forward_from(inputs, labels)
         loss = outputs.loss
         loss.backward()