Centralised simulations#

Parameter-server distributed SGD simulation confronting Byzantine workers.

Each synchronous round follows the same pattern:

  1. Honest workers compute a gradient on their local data shard.

  2. Byzantine workers craft adversarial gradients.

  3. The aggregator combines all \(n\) gradients into a single update.

  4. The aggregated update is applied via an SGD step.

The CentralisedSimulation implements the full lifecycle (model initialisation, per-worker data loading, training loop). Evaluation is defined by subclasses overriding evaluate() — each protocol reports its own set of metrics.

Available simulations#

Base simulation class#

class krum.simulations.centralised.CentralisedSimulation(*, model_cls: type[~torch.nn.modules.module.Module], train_datasets: ~collections.abc.Sequence[~torch.utils.data.dataset.Dataset[~typing.Any]], test_set: ~torch.utils.data.dataset.Dataset[~typing.Any], aggregator: type[~krum.primitives.aggregators.Aggregator] | None = None, aggregator_kwargs: dict[str, ~typing.Any] | None = None, attack: type[~krum.primitives.attacks.Attack] | None = None, attack_kwargs: dict[str, ~typing.Any] | None = None, n: int, f: int, rounds: int, batch_size: int, lr: float, lr_schedule: ~typing.Literal['exponential', 'robbins_monro', 'none'] = 'exponential', lr_decay: float | None = 0.99, r_eta: float | None = None, weight_decay: float = 0.0, xavier_init: bool = False, stop_attack_at: int | None = None, loss_fn: ~typing.Callable[[...], ~torch.Tensor] = <function cross_entropy>, device: ~torch.device | None = None, seed: int = 42, eval_every: int = 10)[source]#

Bases: object

Parameter-server distributed SGD simulation with Byzantine workers.

One instance = one (aggregator, attack, dataset, model) configuration run over rounds synchronous rounds. Each of the \(n\) workers — of which \(f\) are Byzantine (up to the tolerance of the chosen aggregator) — brings its own training dataset, IID or not.

Evaluation is defined by subclasses overriding evaluate(). Each protocol reports its own set of metrics — e.g. KrumSimulation returns (test_loss, test_error) and HiddenVulnerabilitySimulation returns (test_loss, test_error, test_accuracy).

Parameters:

  • model_cls – nn.Module subclass to instantiate for training.

  • train_datasets – One dataset per worker (honest and Byzantine); len(train_datasets) must equal n. Only the first n - f (the honest workers) are ever wrapped into a DataLoader and trained on — Byzantine workers craft their gradients from the honest ones via attack, not from local data — but the full n-length sequence is required so a future data-consuming attack (e.g. label-flipping) has something to read. Splitting a full dataset into per-worker datasets (IID or not) is the caller’s responsibility, via a DataPartitioner.

  • test_set – Test dataset (evaluated via full-batch loader).

  • aggregator – Gradient aggregation rule class (e.g. Average, Krum). Pass the class itself — step() calls aggregator.aggregate(gradients, n=n, f=f, **aggregator_kwargs).

  • aggregator_kwargs – Extra keyword arguments forwarded to aggregator.aggregate (e.g. {"m": 12} for MultiKrum). n and f are automatically injected by the simulation. To override the f injected from f (the “real” Byzantine count), pass "f" explicitly here — for example, to run with f=0 honest workers but configure Bulyan as if it had to defend against 9 Byzantins, pass aggregator_kwargs={"f": 9}.

  • attack – Byzantine attack strategy class (e.g. GaussianAttack). Pass the class itself — step() calls attack.generate(honest_gradients, f=f, **attack_kwargs).

  • attack_kwargs – Extra keyword arguments forwarded to attack.generate (e.g. {"std": 200.0} for GaussianAttack).

  • n – Total number of workers.

  • f – Number of Byzantine workers (the real Byzantine count). Controls how many gradients the simulation produces per round: n - f honest workers always run, and the remaining f slots are filled with Byzantine gradients only when attack is provided. Must be within the aggregator’s Byzantine tolerance. To configure the aggregator with a different (declared) Byzantine budget, see aggregator_kwargs above.

  • rounds – Number of synchronous training rounds.

  • batch_size – Mini-batch size per honest worker.

  • lr – Initial learning rate for SGD. Used as \(η_0\) by every supported scheduler.

  • lr_schedule – Learning-rate schedule applied after each round. "exponential" uses multiplicative decay lr_decay per round (the default for the ICML 2018 protocol). "robbins_monro" uses the \(η(t) = r_η · η_0 / (t + r_η)\) schedule of El Mhamdi et al. (ICML 2018), with fading rate r_eta. "none" keeps a constant learning rate (the default for the NIPS 2017 protocol).

  • lr_decay – Multiplicative learning-rate decay per round, used when lr_schedule == "exponential". None disables the scheduler in that mode. Default: 0.99.

  • r_eta – Fading rate for the Robbins-Monro schedule \(η(t) = r_η · η_0 / (t + r_η)\). Required when lr_schedule == "robbins_monro".

  • weight_decay – \(ℓ_2\) regularization coefficient applied directly on the flat parameter tensor. Set to 0.0 (default) to disable. Section 5.1 of El Mhamdi et al. (ICML 2018) recommends 1e-4.

  • xavier_init – When True, apply Glorot/Xavier uniform initialization to every weight tensor of the model after construction, and zero-initialize biases. Matches Section 5.1 of El Mhamdi et al. (ICML 2018).

  • stop_attack_at – If set, the Byzantine attack is disabled at round t = stop_attack_at: the f Byzantine workers then send zero gradients. Used by Experiment 1 / Figure 2 of El Mhamdi et al. (ICML 2018), where the attack is maintained only up to round 50.

  • loss_fn – Per-sample loss function. Default: cross_entropy.

  • device – Device for training and evaluation. Auto-detected if None (CUDA → MPS → CPU).

  • seed – Random seed for reproducibility.

  • eval_every – Hint for how often evaluation should occur (used by experiment scripts, not enforced by the simulation).

Raises:

ValueError – If lr_schedule is invalid, required hyperparameters (\(r_eta\)) are missing, the number of train datasets does not equal n, or an honest worker’s dataset is empty.

property model: Model#

The encapsulated Model, available after setup() or run().

Returns:

The wrapped ``nn.Module`` with zero-copy flat parameter/gradient views.

Raises:

RuntimeError – If the simulation has not been set up yet.

setup() → None[source]#

Initialise the model, learning rate, and per-worker dataloaders.

Each honest worker’s train_datasets entry is wrapped into a dedicated DataLoader with its own RNG generator, so mini-batch sampling is reproducible across runs. The datasets are used as-is — splitting a full dataset into per-worker datasets (IID or not) is the caller’s responsibility, via a DataPartitioner.

The learning rate is initialised to self.lr. The Robbins-Monro schedule updates it each round inside step(); the exponential schedule decays it after every round via self._current_lr *= lr_decay.

When xavier_init is enabled, every weight tensor of the instantiated model is re-initialized with the Glorot/Xavier uniform rule, and every bias is reset to zero. The Xavier draws use a per-parameter local torch.Generator (matched to the parameter device), so they do not perturb the global PyTorch RNG and work correctly on MPS/CUDA.

Determinism: setup() is fully deterministic for a given seed. Local torch.Generator instances are used for the per-worker dataloaders and (when enabled) the Xavier re-initialization, so the global RNG state is left untouched and re-running setup() reproduces the exact same model, weights, and dataloaders.

Safe to call multiple times — each call resets all internal state.

step() → None[source]#

Advance the simulation by one synchronous round.

  1. The learning rate is updated to the value of the current round (only for the Robbins-Monro schedule; the exponential schedule updates the rate after the optimizer step instead).

  2. Each of the \(n - f\) honest workers computes a gradient on its local data shard via _train_one_worker().

  3. If \(f > 0\) and the attack has not been stopped, Byzantine workers generate attack gradients. For FullGradientNegationAttack, the full-dataset honest gradient is computed first.

  4. The aggregator combines all \(n\) gradients into a single update via self.aggregator.aggregate(...).

  5. The aggregated gradient is written to self._model.gradients and applied via an in-place SGD update on the flat parameter tensor.

  6. The learning rate (if scheduled) is decayed.

Raises:

RuntimeError – If setup() has not been called.