flwr-nightly 1.22.0.dev20250918__py3-none-any.whl → 1.22.0.dev20250919__py3-none-any.whl

flwr/clientapp/mod/__init__.py

@@ -17,9 +17,10 @@
 
 from flwr.client.mod.comms_mods import arrays_size_mod, message_size_mod
 
-from .centraldp_mods import fixedclipping_mod
+from .centraldp_mods import adaptiveclipping_mod, fixedclipping_mod
 
 __all__ = [
+    "adaptiveclipping_mod",
     "arrays_size_mod",
     "fixedclipping_mod",
     "message_size_mod",

flwr/clientapp/mod/centraldp_mods.py

@@ -16,13 +16,26 @@
 
 
 from collections import OrderedDict
-from logging import INFO, WARN
+from logging import ERROR, INFO
 from typing import cast
 
-from flwr.client.typing import ClientAppCallable
-from flwr.common import Array, ArrayRecord, Context, Message, MessageType, log
-from flwr.common.differential_privacy import compute_clip_model_update
-from flwr.common.differential_privacy_constants import KEY_CLIPPING_NORM
+from flwr.app import Error
+from flwr.clientapp.typing import ClientAppCallable
+from flwr.common import (
+    Array,
+    ArrayRecord,
+    ConfigRecord,
+    Context,
+    Message,
+    MetricRecord,
+    log,
+)
+from flwr.common.constant import ErrorCode
+from flwr.common.differential_privacy import (
+    compute_adaptive_clip_model_update,
+    compute_clip_model_update,
+)
+from flwr.common.differential_privacy_constants import KEY_CLIPPING_NORM, KEY_NORM_BIT
 
 
 # pylint: disable=too-many-return-statements
@@ -46,33 +59,15 @@ def fixedclipping_mod(
 
     Typically, fixedclipping_mod should be the last to operate on params.
     """
-    if msg.metadata.message_type != MessageType.TRAIN:
-        return call_next(msg, ctxt)
-
     if len(msg.content.array_records) != 1:
-        log(
-            WARN,
-            "fixedclipping_mod is designed to work with a single ArrayRecord. "
-            "Skipping.",
-        )
-        return call_next(msg, ctxt)
-
+        return _handle_multi_record_err("fixedclipping_mod", msg, ArrayRecord)
     if len(msg.content.config_records) != 1:
-        log(
-            WARN,
-            "fixedclipping_mod is designed to work with a single ConfigRecord. "
-            "Skipping.",
-        )
-        return call_next(msg, ctxt)
+        return _handle_multi_record_err("fixedclipping_mod", msg, ConfigRecord)
 
     # Get keys in the single ConfigRecord
     keys_in_config = set(next(iter(msg.content.config_records.values())).keys())
     if KEY_CLIPPING_NORM not in keys_in_config:
-        raise KeyError(
-            f"The {KEY_CLIPPING_NORM} value is not supplied by the "
-            f"`DifferentialPrivacyClientSideFixedClipping` wrapper at"
-            f" the server side."
-        )
+        return _handle_no_key_err("fixedclipping_mod", msg)
     # Record array record communicated to client and clipping norm
     original_array_record = next(iter(msg.content.array_records.values()))
     clipping_norm = cast(
@@ -86,26 +81,16 @@ def fixedclipping_mod(
     if out_msg.has_error():
         return out_msg
 
-    # Ensure there is a single ArrayRecord
+    # Ensure reply has a single ArrayRecord
     if len(out_msg.content.array_records) != 1:
-        log(
-            WARN,
-            "fixedclipping_mod is designed to work with a single ArrayRecord. "
-            "Skipping.",
-        )
-        return out_msg
+        return _handle_multi_record_err("fixedclipping_mod", out_msg, ArrayRecord)
 
     new_array_record_key, client_to_server_arrecord = next(
         iter(out_msg.content.array_records.items())
     )
     # Ensure keys in returned ArrayRecord match those in the one sent from server
     if set(original_array_record.keys()) != set(client_to_server_arrecord.keys()):
-        log(
-            WARN,
-            "fixedclipping_mod: Keys in ArrayRecord must match those from the model "
-            "that the ClientApp received. Skipping.",
-        )
-        return out_msg
+        return _handle_array_key_mismatch_err("fixedclipping_mod", out_msg)
 
     client_to_server_ndarrays = client_to_server_arrecord.to_numpy_ndarrays()
     # Clip the client update
@@ -130,3 +115,134 @@ def fixedclipping_mod(
         )
     )
     return out_msg
+
+
+def adaptiveclipping_mod(
+    msg: Message, ctxt: Context, call_next: ClientAppCallable
+) -> Message:
+    """Client-side adaptive clipping modifier.
+
+    This mod needs to be used with the DifferentialPrivacyClientSideAdaptiveClipping
+    server-side strategy wrapper.
+
+    The wrapper sends the clipping_norm value to the client.
+
+    This mod clips the client model updates before sending them to the server.
+
+    It also sends KEY_NORM_BIT to the server for computing the new clipping value.
+
+    It operates on messages of type `MessageType.TRAIN`.
+
+    Notes
+    -----
+    Consider the order of mods when using multiple.
+
+    Typically, adaptiveclipping_mod should be the last to operate on params.
+    """
+    if len(msg.content.array_records) != 1:
+        return _handle_multi_record_err("adaptiveclipping_mod", msg, ArrayRecord)
+    if len(msg.content.config_records) != 1:
+        return _handle_multi_record_err("adaptiveclipping_mod", msg, ConfigRecord)
+
+    # Get keys in the single ConfigRecord
+    keys_in_config = set(next(iter(msg.content.config_records.values())).keys())
+    if KEY_CLIPPING_NORM not in keys_in_config:
+        return _handle_no_key_err("adaptiveclipping_mod", msg)
+
+    # Record array record communicated to client and clipping norm
+    original_array_record = next(iter(msg.content.array_records.values()))
+    clipping_norm = cast(
+        float, next(iter(msg.content.config_records.values()))[KEY_CLIPPING_NORM]
+    )
+
+    # Call inner app
+    out_msg = call_next(msg, ctxt)
+
+    # Ensure reply has a single ArrayRecord
+    if len(out_msg.content.array_records) != 1:
+        return _handle_multi_record_err("adaptiveclipping_mod", out_msg, ArrayRecord)
+
+    # Ensure reply has a single MetricRecord
+    if len(out_msg.content.metric_records) != 1:
+        return _handle_multi_record_err("adaptiveclipping_mod", out_msg, MetricRecord)
+
+    # Check if the msg has error
+    if out_msg.has_error():
+        return out_msg
+
+    new_array_record_key, client_to_server_arrecord = next(
+        iter(out_msg.content.array_records.items())
+    )
+
+    # Ensure keys in returned ArrayRecord match those in the one sent from server
+    if set(original_array_record.keys()) != set(client_to_server_arrecord.keys()):
+        return _handle_array_key_mismatch_err("adaptiveclipping_mod", out_msg)
+
+    client_to_server_ndarrays = client_to_server_arrecord.to_numpy_ndarrays()
+    # Clip the client update
+    norm_bit = compute_adaptive_clip_model_update(
+        client_to_server_ndarrays,
+        original_array_record.to_numpy_ndarrays(),
+        clipping_norm,
+    )
+    log(
+        INFO,
+        "adaptiveclipping_mod: ndarrays are clipped by value: %.4f.",
+        clipping_norm,
+    )
+    # Replace outgoing ArrayRecord's Array while preserving their keys
+    out_msg.content.array_records[new_array_record_key] = ArrayRecord(
+        OrderedDict(
+            {
+                k: Array(v)
+                for k, v in zip(
+                    client_to_server_arrecord.keys(), client_to_server_ndarrays
+                )
+            }
+        )
+    )
+    # Add to the MetricRecords the norm bit (recall reply messages only contain
+    # one MetricRecord)
+    metric_record_key = list(out_msg.content.metric_records.keys())[0]
+    # We cast it to `int` because MetricRecord does not support `bool` values
+    out_msg.content.metric_records[metric_record_key][KEY_NORM_BIT] = int(norm_bit)
+    return out_msg
+
+
+def _handle_err(msg: Message, reason: str) -> Message:
+    """Log and return error message."""
+    log(ERROR, reason)
+    return Message(
+        Error(code=ErrorCode.MOD_FAILED_PRECONDITION, reason=reason),
+        reply_to=msg,
+    )
+
+
+def _handle_multi_record_err(mod_name: str, msg: Message, record_type: type) -> Message:
+    """Log and return multi-record error."""
+    cnt = sum(isinstance(_, record_type) for _ in msg.content.values())
+    return _handle_err(
+        msg,
+        f"{mod_name} expects exactly one {record_type.__name__}, "
+        f"but found {cnt} {record_type.__name__}(s).",
+    )
+
+
+def _handle_no_key_err(mod_name: str, msg: Message) -> Message:
+    """Log and return no-key error."""
+    return _handle_err(
+        msg,
+        f"{mod_name} requires the key '{KEY_CLIPPING_NORM}' to be present in the "
+        "ConfigRecord, but it was not found. "
+        "Please ensure the `DifferentialPrivacyClientSideFixedClipping` wrapper "
+        "is used in the ServerApp.",
+    )
+
+
+def _handle_array_key_mismatch_err(mod_name: str, msg: Message) -> Message:
+    """Create array-key-mismatch error reasons."""
+    return _handle_err(
+        msg,
+        f"{mod_name} expects the keys in the ArrayRecord of the reply message to match "
+        "those from the ArrayRecord that the ClientApp received, but they do not.",
+    )

flwr/clientapp/typing.py

@@ -0,0 +1,22 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Custom types for Flower clients."""
+
+
+from typing import Callable
+
+from flwr.common import Context, Message
+
+ClientAppCallable = Callable[[Message, Context], Message]

flwr/common/constant.py

@@ -202,6 +202,7 @@ class ErrorCode:
     MESSAGE_UNAVAILABLE = 3
     REPLY_MESSAGE_UNAVAILABLE = 4
     NODE_UNAVAILABLE = 5
+    MOD_FAILED_PRECONDITION = 6
 
     def __new__(cls) -> ErrorCode:
         """Prevent instantiation."""

flwr/common/record/typeddict.py

@@ -18,6 +18,8 @@
 from collections.abc import ItemsView, Iterator, KeysView, MutableMapping, ValuesView
 from typing import Callable, Generic, TypeVar, cast
 
+from typing_extensions import Self
+
 K = TypeVar("K")  # Key type
 V = TypeVar("V")  # Value type
 
@@ -86,3 +88,13 @@ class TypedDict(MutableMapping[K, V], Generic[K, V]):
     def items(self) -> ItemsView[K, V]:
         """D.items() -> a set-like object providing a view on D's items."""
         return cast(dict[K, V], self.__dict__["_data"]).items()
+
+    def copy(self) -> Self:
+        """Return a shallow copy of the dictionary."""
+        # Allocate instance without going through __init__
+        new = self.__class__.__new__(type(self))
+        # Copy internal state
+        new.__dict__["_check_key_fn"] = self.__dict__["_check_key_fn"]
+        new.__dict__["_check_value_fn"] = self.__dict__["_check_value_fn"]
+        new.__dict__["_data"] = cast(dict[K, V], self.__dict__["_data"]).copy()
+        return new

flwr/serverapp/strategy/__init__.py

@@ -15,6 +15,10 @@
 """ServerApp strategies."""
 
 
+from .dp_adaptive_clipping import (
+    DifferentialPrivacyClientSideAdaptiveClipping,
+    DifferentialPrivacyServerSideAdaptiveClipping,
+)
 from .dp_fixed_clipping import (
     DifferentialPrivacyClientSideFixedClipping,
     DifferentialPrivacyServerSideFixedClipping,
@@ -29,11 +33,15 @@ from .fedtrimmedavg import FedTrimmedAvg
 from .fedxgb_bagging import FedXgbBagging
 from .fedxgb_cyclic import FedXgbCyclic
 from .fedyogi import FedYogi
+from .krum import Krum
+from .qfedavg import QFedAvg
 from .result import Result
 from .strategy import Strategy
 
 __all__ = [
+    "DifferentialPrivacyClientSideAdaptiveClipping",
     "DifferentialPrivacyClientSideFixedClipping",
+    "DifferentialPrivacyServerSideAdaptiveClipping",
     "DifferentialPrivacyServerSideFixedClipping",
     "FedAdagrad",
     "FedAdam",
@@ -45,6 +53,8 @@ __all__ = [
     "FedXgbBagging",
     "FedXgbCyclic",
     "FedYogi",
+    "Krum",
+    "QFedAvg",
     "Result",
     "Strategy",
 ]

flwr/serverapp/strategy/dp_adaptive_clipping.py

@@ -0,0 +1,335 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Message-based Central differential privacy with adaptive clipping.
+
+Paper (Andrew et al.): https://arxiv.org/abs/1905.03871
+"""
+
+import math
+from abc import ABC
+from collections import OrderedDict
+from collections.abc import Iterable
+from logging import INFO
+from typing import Optional
+
+import numpy as np
+
+from flwr.common import Array, ArrayRecord, ConfigRecord, Message, MetricRecord, log
+from flwr.common.differential_privacy import (
+    adaptive_clip_inputs_inplace,
+    add_gaussian_noise_inplace,
+    compute_adaptive_noise_params,
+    compute_stdv,
+)
+from flwr.common.differential_privacy_constants import KEY_CLIPPING_NORM, KEY_NORM_BIT
+from flwr.server import Grid
+from flwr.serverapp.exception import AggregationError
+
+from .dp_fixed_clipping import validate_replies
+from .strategy import Strategy
+
+
+class DifferentialPrivacyAdaptiveBase(Strategy, ABC):
+    """Base class for DP strategies with adaptive clipping."""
+
+    # pylint: disable=too-many-arguments,too-many-instance-attributes,too-many-positional-arguments
+    def __init__(
+        self,
+        strategy: Strategy,
+        noise_multiplier: float,
+        num_sampled_clients: int,
+        initial_clipping_norm: float = 0.1,
+        target_clipped_quantile: float = 0.5,
+        clip_norm_lr: float = 0.2,
+        clipped_count_stddev: Optional[float] = None,
+    ) -> None:
+        super().__init__()
+
+        if strategy is None:
+            raise ValueError("The passed strategy is None.")
+        if noise_multiplier < 0:
+            raise ValueError("The noise multiplier should be a non-negative value.")
+        if num_sampled_clients <= 0:
+            raise ValueError(
+                "The number of sampled clients should be a positive value."
+            )
+        if initial_clipping_norm <= 0:
+            raise ValueError("The initial clipping norm should be a positive value.")
+        if not 0 <= target_clipped_quantile <= 1:
+            raise ValueError("The target clipped quantile must be in [0, 1].")
+        if clip_norm_lr <= 0:
+            raise ValueError("The learning rate must be positive.")
+        if clipped_count_stddev is not None and clipped_count_stddev < 0:
+            raise ValueError("The `clipped_count_stddev` must be non-negative.")
+
+        self.strategy = strategy
+        self.num_sampled_clients = num_sampled_clients
+        self.clipping_norm = initial_clipping_norm
+        self.target_clipped_quantile = target_clipped_quantile
+        self.clip_norm_lr = clip_norm_lr
+        (
+            self.clipped_count_stddev,
+            self.noise_multiplier,
+        ) = compute_adaptive_noise_params(
+            noise_multiplier,
+            num_sampled_clients,
+            clipped_count_stddev,
+        )
+
+    def _add_noise_to_aggregated_arrays(self, aggregated: ArrayRecord) -> ArrayRecord:
+        nds = aggregated.to_numpy_ndarrays()
+        stdv = compute_stdv(
+            self.noise_multiplier, self.clipping_norm, self.num_sampled_clients
+        )
+        add_gaussian_noise_inplace(nds, stdv)
+        log(INFO, "aggregate_fit: central DP noise with %.4f stdev added", stdv)
+        return ArrayRecord(
+            OrderedDict({k: Array(v) for k, v in zip(aggregated.keys(), nds)})
+        )
+
+    def _noisy_fraction(self, count: int, total: int) -> float:
+        return float(np.random.normal(count, self.clipped_count_stddev)) / float(total)
+
+    def _geometric_update(self, clipped_fraction: float) -> None:
+        self.clipping_norm *= math.exp(
+            -self.clip_norm_lr * (clipped_fraction - self.target_clipped_quantile)
+        )
+
+    def configure_evaluate(
+        self, server_round: int, arrays: ArrayRecord, config: ConfigRecord, grid: Grid
+    ) -> Iterable[Message]:
+        """Configure the next round of federated evaluation."""
+        return self.strategy.configure_evaluate(server_round, arrays, config, grid)
+
+    def aggregate_evaluate(
+        self, server_round: int, replies: Iterable[Message]
+    ) -> Optional[MetricRecord]:
+        """Aggregate MetricRecords in the received Messages."""
+        return self.strategy.aggregate_evaluate(server_round, replies)
+
+    def summary(self) -> None:
+        """Log summary configuration of the strategy."""
+        self.strategy.summary()
+
+
+class DifferentialPrivacyServerSideAdaptiveClipping(DifferentialPrivacyAdaptiveBase):
+    """Message-based central DP with server-side adaptive clipping."""
+
+    # pylint: disable=too-many-arguments,too-many-locals,too-many-positional-arguments
+    def __init__(
+        self,
+        strategy: Strategy,
+        noise_multiplier: float,
+        num_sampled_clients: int,
+        initial_clipping_norm: float = 0.1,
+        target_clipped_quantile: float = 0.5,
+        clip_norm_lr: float = 0.2,
+        clipped_count_stddev: Optional[float] = None,
+    ) -> None:
+        super().__init__(
+            strategy,
+            noise_multiplier,
+            num_sampled_clients,
+            initial_clipping_norm,
+            target_clipped_quantile,
+            clip_norm_lr,
+            clipped_count_stddev,
+        )
+        self.current_arrays: ArrayRecord = ArrayRecord()
+
+    def __repr__(self) -> str:
+        """Compute a string representation of the strategy."""
+        return "Differential Privacy Strategy Wrapper (Server-Side Adaptive Clipping)"
+
+    def summary(self) -> None:
+        """Log summary configuration of the strategy."""
+        log(INFO, "\t├──> DP settings:")
+        log(INFO, "\t│\t├── Noise multiplier: %s", self.noise_multiplier)
+        log(INFO, "\t│\t├── Clipping norm: %s", self.clipping_norm)
+        log(INFO, "\t│\t├── Target clipped quantile: %s", self.target_clipped_quantile)
+        log(INFO, "\t│\t└── Clip norm learning rate: %s", self.clip_norm_lr)
+        super().summary()
+
+    def configure_train(
+        self, server_round: int, arrays: ArrayRecord, config: ConfigRecord, grid: Grid
+    ) -> Iterable[Message]:
+        """Configure the next round of training."""
+        self.current_arrays = arrays
+        return self.strategy.configure_train(server_round, arrays, config, grid)
+
+    def aggregate_train(
+        self, server_round: int, replies: Iterable[Message]
+    ) -> tuple[Optional[ArrayRecord], Optional[MetricRecord]]:
+        """Aggregate ArrayRecords and MetricRecords in the received Messages."""
+        if not validate_replies(replies, self.num_sampled_clients):
+            return None, None
+
+        current_nd = self.current_arrays.to_numpy_ndarrays()
+        clipped_indicator_count = 0
+        replies_list = list(replies)
+
+        for reply in replies_list:
+            for arr_name, record in reply.content.array_records.items():
+                reply_nd = record.to_numpy_ndarrays()
+                model_update = [
+                    np.subtract(x, y) for (x, y) in zip(reply_nd, current_nd)
+                ]
+                norm_bit = adaptive_clip_inputs_inplace(
+                    model_update, self.clipping_norm
+                )
+                clipped_indicator_count += int(norm_bit)
+                # reconstruct array using clipped contribution from current round
+                restored = [c + u for c, u in zip(current_nd, model_update)]
+                reply.content[arr_name] = ArrayRecord(
+                    OrderedDict({k: Array(v) for k, v in zip(record.keys(), restored)})
+                )
+            log(
+                INFO,
+                "aggregate_train: arrays in `ArrayRecord` are clipped by value: %.4f.",
+                self.clipping_norm,
+            )
+
+        clipped_fraction = self._noisy_fraction(
+            clipped_indicator_count, len(replies_list)
+        )
+        self._geometric_update(clipped_fraction)
+
+        aggregated_arrays, aggregated_metrics = self.strategy.aggregate_train(
+            server_round, replies_list
+        )
+
+        if aggregated_arrays:
+            aggregated_arrays = self._add_noise_to_aggregated_arrays(aggregated_arrays)
+
+        return aggregated_arrays, aggregated_metrics
+
+
+class DifferentialPrivacyClientSideAdaptiveClipping(DifferentialPrivacyAdaptiveBase):
+    """Strategy wrapper for central DP with client-side adaptive clipping.
+
+    Use `adaptiveclipping_mod` modifier at the client side.
+
+    In comparison to `DifferentialPrivacyServerSideAdaptiveClipping`,
+    which performs clipping on the server-side,
+    `DifferentialPrivacyClientSideAdaptiveClipping`
+    expects clipping to happen on the client-side, usually by using the built-in
+    `adaptiveclipping_mod`.
+
+    Parameters
+    ----------
+    strategy : Strategy
+        The strategy to which DP functionalities will be added by this wrapper.
+    noise_multiplier : float
+        The noise multiplier for the Gaussian mechanism for model updates.
+    num_sampled_clients : int
+        The number of clients that are sampled on each round.
+    initial_clipping_norm : float
+        The initial value of clipping norm. Defaults to 0.1.
+        Andrew et al. recommends to set to 0.1.
+    target_clipped_quantile : float
+        The desired quantile of updates which should be clipped. Defaults to 0.5.
+    clip_norm_lr : float
+        The learning rate for the clipping norm adaptation. Defaults to 0.2.
+        Andrew et al. recommends to set to 0.2.
+    clipped_count_stddev : float
+        The stddev of the noise added to the count of
+        updates currently below the estimate.
+        Andrew et al. recommends to set to `expected_num_records/20`
+
+    Examples
+    --------
+    Create a strategy::
+
+        strategy = fl.serverapp.FedAvg(...)
+
+    Wrap the strategy with the `DifferentialPrivacyClientSideAdaptiveClipping` wrapper::
+
+        dp_strategy = DifferentialPrivacyClientSideAdaptiveClipping(
+            strategy, cfg.noise_multiplier, cfg.num_sampled_clients, ...
+        )
+
+    On the client, add the `adaptiveclipping_mod` to the client-side mods::
+
+        app = fl.client.ClientApp(mods=[adaptiveclipping_mod])
+    """
+
+    def __repr__(self) -> str:
+        """Compute a string representation of the strategy."""
+        return "Differential Privacy Strategy Wrapper (Client-Side Adaptive Clipping)"
+
+    def summary(self) -> None:
+        """Log summary configuration of the strategy."""
+        log(INFO, "\t├──> DP settings:")
+        log(INFO, "\t│\t├── Noise multiplier: %s", self.noise_multiplier)
+        log(INFO, "\t│\t├── Clipping norm: %s", self.clipping_norm)
+        log(INFO, "\t│\t├── Target clipped quantile: %s", self.target_clipped_quantile)
+        log(INFO, "\t│\t└── Clip norm learning rate: %s", self.clip_norm_lr)
+        super().summary()
+
+    def configure_train(
+        self, server_round: int, arrays: ArrayRecord, config: ConfigRecord, grid: Grid
+    ) -> Iterable[Message]:
+        """Configure the next round of training."""
+        config[KEY_CLIPPING_NORM] = self.clipping_norm
+        return self.strategy.configure_train(server_round, arrays, config, grid)
+
+    def aggregate_train(
+        self, server_round: int, replies: Iterable[Message]
+    ) -> tuple[Optional[ArrayRecord], Optional[MetricRecord]]:
+        """Aggregate ArrayRecords and MetricRecords in the received Messages."""
+        if not validate_replies(replies, self.num_sampled_clients):
+            return None, None
+
+        replies_list = list(replies)
+
+        # validate that KEY_NORM_BIT is present in all replies
+        for msg in replies_list:
+            for _, mrec in msg.content.metric_records.items():
+                if KEY_NORM_BIT not in mrec:
+                    raise AggregationError(
+                        f"KEY_NORM_BIT ('{KEY_NORM_BIT}') not found"
+                        f" in MetricRecord or metrics for reply."
+                    )
+
+        aggregated_arrays, aggregated_metrics = self.strategy.aggregate_train(
+            server_round, replies_list
+        )
+
+        self._update_clip_norm_from_replies(replies_list)
+
+        if aggregated_arrays:
+            aggregated_arrays = self._add_noise_to_aggregated_arrays(aggregated_arrays)
+
+        return aggregated_arrays, aggregated_metrics
+
+    def _update_clip_norm_from_replies(self, replies: list[Message]) -> None:
+        total = len(replies)
+        clipped_count = 0
+
+        for msg in replies:
+            # KEY_NORM_BIT is guaranteed to be present
+            for _, mrec in msg.content.metric_records.items():
+                if KEY_NORM_BIT in mrec:
+                    clipped_count += int(bool(mrec[KEY_NORM_BIT]))
+                    break
+            else:
+                # Check fallback location
+                if hasattr(msg.content, "metrics") and isinstance(
+                    msg.content.metrics, dict
+                ):
+                    clipped_count += int(bool(msg.content.metrics[KEY_NORM_BIT]))
+
+        clipped_fraction = self._noisy_fraction(clipped_count, total)
+        self._geometric_update(clipped_fraction)