Centralised simulation walkthrough#

Problem: You need to run a parameter-server simulation with multiple workers, a gradient aggregator, and Byzantine attacks, but you are not sure how to configure the training loop.

Krum ships with ready-to-use centralised simulations that handle the worker loop, gradient computation, and evaluation for you.

See also

Centralised simulations

Reference for KrumSimulation and HiddenVulnerabilitySimulation.

Data preparation#

Each of the n workers (honest and Byzantine) brings its own training dataset. The caller is responsible for splitting a full dataset into per-worker datasets — this can be IID or non-IID. The built-in data_partitioners strategies handle this:

from torchvision import datasets, transforms

from krum.primitives.data_partitioners.iid import IidPartitioner

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,)),
])
train_set = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
test_set = datasets.MNIST(root="./data", train=False, download=True, transform=transform)

# One dataset per worker: IidPartitioner gives each worker an equal,
# shuffled portion of the training set
train_datasets = IidPartitioner.partition(train_set, n=10, seed=42)

For non-IID data, DirichletPartitioner produces per-class label skew (controlled by an alpha parameter). See Working with Data Partitioners for the full partitioner family.

len(train_datasets) must equal n, including Byzantine workers: only the first n - f are ever trained on (Byzantine workers craft their gradients from the honest ones via the attack), but the full n-length sequence is required so a future data-consuming attack (e.g. label-flipping) has something to read.

Minimal example#

Instantiation#

Aggregator and attack are passed as classes, not instances. The simulation calls aggregator.aggregate() and attack.generate() each round. Extra parameters go through aggregator_kwargs and attack_kwargs:

from torchvision import datasets, transforms

from krum.primitives.aggregators.multikrum import MultiKrum
from krum.primitives.attacks.sign_flip import SignFlipAttack
from krum.primitives.data_partitioners.iid import IidPartitioner
from krum.primitives.models.mlp import Krum2017MLPMnist
from krum.simulations.centralised.krum_nips_2017 import KrumSimulation

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,)),
])
train_set = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
test_set = datasets.MNIST(root="./data", train=False, download=True, transform=transform)

# Partition the training set into one dataset per worker
train_datasets = IidPartitioner.partition(train_set, n=10, seed=42)

sim = KrumSimulation(
     model_cls=Krum2017MLPMnist,
     train_datasets=train_datasets,
     test_set=test_set,
     aggregator=MultiKrum,
     attack=SignFlipAttack,
     attack_kwargs={"scale": 1.5},
     n=10,
     f=2,
     rounds=50,
     batch_size=64,
     lr=0.01,
     seed=42,
 )

Note that CentralisedSimulation does not provide a run(rounds) method. The training loop is always manual so you can evaluate when you want (compare with the Decentralised simulation walkthrough, which does have run).

Setup#

setup() initialises the model parameters, wraps each honest worker’s dataset into a dedicated DataLoader (batch size, shuffling), and seeds all RNG. The result is deterministic for a given seed:

sim.setup()

Step#

Each call to step() runs one synchronous round:

  • Broadcast the model to all \(n\) workers.

  • Honest workers compute gradients on their local data shard.

  • Byzantine workers generate attack gradients.

  • Aggregator combines all \(n\) gradients into one.

  • SGD update is applied.

for round_idx in range(50):
    sim.step()
    if round_idx % 10 == 0:
        loss, accuracy = sim.evaluate()
        print(f"round {round_idx}: loss={loss:.4f}  accuracy={accuracy:.4f}")

Evaluate#

evaluate() returns the metrics specific to the protocol ((test_loss, test_accuracy) for KrumSimulation, (test_loss, test_error, test_accuracy) for HiddenVulnerabilitySimulation). You can also read the training loss with evaluate_train().

Using the ICML 2018 simulation#

Switching to the other built-in simulation changes only the import and one extra parameter. The ICML 2018 variant adds Xavier weight initialisation, L2 regularisation, and the Robbins-Monro learning-rate schedule:

from krum.simulations.centralised.hidden_vulnerability_icml_2018 import (
    HiddenVulnerabilitySimulation,
)

# Same data preparation as above
train_datasets = IidPartitioner.partition(train_set, n=39, seed=42)

sim = HiddenVulnerabilitySimulation(
     model_cls=Krum2017MLPMnist,
     train_datasets=train_datasets,
     test_set=test_set,
     aggregator=Bulyan,
     attack=SmallPerturbationAttack,
     n=39,
     f=9,
     r_eta=10_000,
     rounds=100,
     batch_size=32,
     lr=0.1,
     seed=42,
)
sim.setup()
for round_idx in range(100):
    sim.step()
    if round_idx % 10 == 0:
        loss, error, accuracy = sim.evaluate()
        print(f"round {round_idx}: loss={loss:.4f}  error={error:.4f}  accuracy={accuracy:.4f}")

train_loss = sim.evaluate_train()
print(f"final training loss: {train_loss:.4f}")

HiddenVulnerabilitySimulation also accepts stop_attack_at. This is an optional round index after which the Byzantine attack is disabled (used by the ICML 2018 paper, Experiment 1). Pass stop_attack_at=50 to stop the attack after round 50 while continuing training.

Note that unlike the decentralised simulation, the centralised one stores the round count at construction time (rounds=50 above) but you drive the loop yourself. This lets you evaluate, log, or even change hyperparameters between rounds.

Next steps#