Decentralised simulations

Peer-to-peer decentralised learning simulation confronting Byzantine workers.

Each honest worker holds its own model. Each round runs two phases:

1. Local optimisation — each worker computes a gradient on its local batch and updates its own model.

2. Model mixing — each worker gathers models from other nodes and replaces its model with an aggregate of the received set.

Two abstract seams let protocols vary the local update rule (e.g. momentum-SGD) and the communication topology (which models each worker receives).

Available simulations

• MonnaSimulation

Base simulation class

class krum.simulations.decentralised.DecentralisedSimulation(*, model: Model, train_datasets: Sequence[Dataset[Any]], train_batch_size: int, test_set: Dataset[Any], test_batch_size: int, loss_fn: Callable[[Tensor, Tensor], Tensor], n: int, f: int, attack: type[Attack] | None = None, attack_kwargs: dict[str, Any] | None = None, aggregator: type[Aggregator], aggregator_kwargs: dict[str, Any] | None = None, seed: int | None = None)

Bases: ABC, Generic[StepResultT]

Base for decentralised simulations with per-worker model mixing.

Each honest worker holds its own flat parameter vector — one row of parameters. A round runs a local optimisation step, then a mixing phase in which every worker replaces its model with an aggregate of the models it received this round.

Two things vary between protocols, and each is an abstract seam:

• how a worker updates locally — local_update() maps the worker gradients to the post-local parameters theta_{t+1/2}. The base is agnostic to the rule and its state, so a momentum-free or optimiser-based protocol keeps none of MoNNA’s momentum machinery.

• which models a worker receives — gather_received_models() builds each worker’s received set (the communication topology).

The base owns everything else: gradient computation, the Byzantine generation hook, the mixing loop, the generic state commit, and the multi-round run() driver. Subclasses also implement build_step_result() so each protocol’s snapshot can carry its own fields (the base StepResult plus, e.g., momentum).

run() may be called repeatedly to continue training: all state (parameters, step_index, subclass optimiser state, and the worker dataloaders) lives on the instance and persists across calls. Each worker’s DataLoader is automatically re-iterated (a fresh epoch) once exhausted, so run() may be called for more rounds than one epoch provides.

Evaluation is defined by subclasses overriding evaluate() — it is never called from step(), so a caller decides when a test-set snapshot is worth its cost (e.g. every few rounds).

aggregate_over_received_nodes(local_parameters: Tensor, byzantine_parameters: Tensor) → Tensor

Run the model-mixing phase over post-local-update parameter vectors.

For each worker, builds its received set via the protocol-specific gather_received_models(), then aggregates it.

Parameters:

• local_parameters – Post-local-update honest models, one row per worker.

• byzantine_parameters – Byzantine models, shape (f, d).

Returns:

The mixed models, one row per honest worker.

aggregate_received_models(candidates: Tensor, *, pivot: Tensor) → Tensor

Aggregate the set of models one worker received.

NearestNeighborAverage anchors on the worker’s own model via pivot; pivot-free aggregators (e.g. Krum, Median) absorb it through **specialized.

Parameters:

• candidates – The received models for one worker.

• pivot – The worker’s own model, used as the distance reference.

Returns:

The single mixed model for the worker, shape `` (d,)

abstract build_step_result(*, honest_gradients: Tensor, local_parameters: Tensor, byzantine_parameters: Tensor, mixed_parameters: Tensor, losses: Tensor) → StepResultT

Build the round snapshot, called after the state has been committed.

Implementations return the protocol’s StepResult subtype, adding any optimiser-state fields (e.g. momentum). The committed step_index and parameters are available via self.

Parameters:

• honest_gradients – Stacked honest gradients this round.

• local_parameters – Post-local-update honest models.

• byzantine_parameters – Byzantine models injected this round.

• mixed_parameters – Mixed models (equal to the committed parameters).

• losses – Per-worker losses.

Returns:

The protocol snapshot, with a detached clone of each tensor.

collect_worker_batches() → list[tuple[Tensor, Tensor]]

Pull one local batch from every honest worker’s DataLoader.

When a worker’s loader is exhausted (its current epoch ends), a fresh iterator is created automatically and the first batch of the new epoch is used instead, so callers do not need to pre-cycle their loaders to run more rounds than one epoch provides.

Returns:

One batch per honest worker, in worker order.

commit_state(parameters: Tensor) → None

Persist the next parameters and advance the round counter.

Subclasses commit their own optimiser state inside local_update(); this only handles the parameter state shared by every protocol.

Parameters:

parameters – The mixed models to persist as the next parameters.

compute_honest_worker_gradients(batches: Sequence[tuple[Tensor, Tensor]]) → tuple[Tensor, Tensor]

Compute gradients at each honest worker’s current parameters.

Batches are moved to the model’s device (wherever parameters already lives) before the forward pass, so callers may freely mix a CPU-only DataLoader with a model placed on an accelerator.

Parameters:

batches – One batch per honest worker, in worker order.

Returns:

• A tuple `` (gradients, losses)

• per honest worker.

copy_parameters_to_model(parameters: Tensor) → None

Copy one flat parameter vector into the shared model wrapper.

Parameters:

parameters – Flat parameter vector of shape (d,) to load.

abstract evaluate() → tuple[float, ...]

Evaluate every honest worker’s current local model on the test set.

This is the evaluation seam: each protocol reports its own set of metrics, the way CentralisedSimulation subclasses do. Never called automatically — invoke it explicitly whenever a test-set snapshot is wanted (e.g. every few rounds), since it is typically far more expensive than one step().

Implementations typically loop over parameters (one row per honest worker), load each into model via copy_parameters_to_model(), evaluate against test_loader, and average the per-worker metrics.

Returns:

The protocol’s test metrics, averaged across honest workers.

abstract gather_received_models(honest_vectors: Tensor, byzantine_parameters: Tensor, *, worker_index: int) → Tensor

Build the set of models received by one honest worker this round.

This is the communication-topology seam: each decentralised protocol decides which models (honest and Byzantine) land in a worker’s received set. Implementations should lead the set with the worker’s own model so a pivot-anchored aggregator can rely on its position.

Parameters:

• honest_vectors – Post-local-update honest models, one row per worker.

• byzantine_parameters – Byzantine models, shape (f, d), as produced by generate_byzantine_models() (may be empty or ignored by protocols that generate Byzantine replies per recipient).

• worker_index – Index of the receiving honest worker.

Returns:

The received models for the worker, with its own model first.

generate_byzantine_models(local_parameters: Tensor) → Tensor

Generate the Byzantine model vectors injected into the mixing phase.

Called once per round before the per-worker mixing loop. The same Byzantine models are then placed into received sets by gather_received_models(). Protocols whose Byzantine replies depend on the recipient (e.g. recipient-specific attacks) override this to produce them inside gather_received_models instead.

Parameters:

local_parameters – Post-local-update honest models, one row per worker, passed to the attack.

Returns:

The Byzantine models, shape `` (f, d)

abstract local_update(gradients: Tensor) → Tensor

Compute the post-local-update parameters theta_{t+1/2}.

This is the optimisation seam: each protocol defines how a worker turns its gradient into its pre-mixing parameters, and owns any optimiser state (momentum, moments, …). Implementations should update that state in place so build_step_result() can snapshot it.

Parameters:

gradients – Stacked honest gradients, one row per worker.

Returns:

The post-local-update parameters, one row per honest worker.

run(rounds: int) → list[StepResultT]

Execute several rounds and collect their snapshots.

State persists on the instance, so successive calls continue training from where the previous call left off.

Parameters:

rounds – Number of rounds to run; must be non-negative.

Returns:

One snapshot per round, in execution order.

Raises:

ValueError – If rounds is negative.

step() → StepResultT

Execute one decentralised training round.

Runs one local optimisation phase (local_update()) followed by one model-mixing phase over the per-worker received sets, then commits the resulting state and builds the snapshot.

Returns:

A snapshot dict of the round, as built by :meth:`build_step_result`.

class krum.simulations.decentralised.StepResult

Bases: TypedDict

Snapshot returned by DecentralisedSimulation.step() for one round.

Holds only the fields common to every decentralised protocol. Subclasses whose local step keeps extra state (e.g. momentum) extend this TypedDict and bind it as the simulation’s StepResultT.

byzantine_parameters: Tensor

honest_gradients: Tensor

local_parameters: Tensor

losses: Tensor

mixed_parameters: Tensor

parameters: Tensor

step: int