CU15 — What a training run exports

for _ in range(rounds):
    for client in clients:
        client.fit(client.data)
    average = fedavg([client.export() for client in clients])
    for client in clients:
        client.load(average)

Status: closed. 410 tests in Python, 40 in the store.

The question CU11 put off, and how it changed shape on the way

CU11 asked is a node’s state a Value? and put it off. It came back as a better question — what does a training run export? — and the answer is the smallest one that is true: its weights, node by node, {node_id: {key: tensor}}.

By the same two ducks the rest of the project asks by: a state_dict by name, parameters() in order, and a node with neither has no weights and does not stop being a node. It has to be the same two state_digest and Graph._check_it_was_obeyed use, or a node could be told to settle and then have no way of being exported — one state, two questions, and a project that answers them differently in two places.

A `for` over a list, and it stays one

A federated round has no dependencies to declare: the clients do not read each other and the order they run in is nobody’s business. A graph earns its keep when there are dependencies, so this is not one — the original put training runs and graph slices in the same enum, and that was the mistake the three levels exist to avoid.

FedAvg, FedProx, FedYogi and SCAFFOLD differ in arithmetic. That is what a function is for, so fedavg is one. The day a topology stops being flat — hierarchical, gossip — that day it is a graph, and not before.

Three things it had to decide

What is trained where it runs cannot be exported from here. Those weights are on the other machine; what is here is the copy that was sent, and it never learnt anything. Handing it back would be silent, which is the only way this could go wrong, so it is refused with the node named. trains on its own is fine — running elsewhere is the half that matters.
The optimizer’s state is not in it. Momentum is a client’s own, and averaging it is not what averaging weights means.
What is not a number you can halve is not halved. A num_batches_tracked is a count and the mean of two counts is not a count; the first one’s stands. Every implementation of this does it and none of them says so out loud.

The test that was written twice

The demonstration has to be a control, not a loss going down: a net says that just as loudly when only a third of it is learning.

The first attempt gave each client one class. Cross-entropy then pushes its own logit up for ever, each client diverges alone, and the average of three diverged nets is a lesson about learning rates — measured, a loss of 2.7e8 against 4.5 for the client that stayed home. The task has to be the same for everyone; only where they draw their inputs from may differ. Three clients, one corner of the input space each, a fixed teacher they are all trying to learn, and a fourth trained alone on its corner for exactly as long. The average reads the union better.

Questionnaire

What it exports (soma-python/tests/test_federated.py)

the weights node by node, and a node with none is simply not in there
a node that says what its weights are called is asked by name, and one that does not is asked in order
what comes out is a snapshot and not a view
it goes back in where it came from
the optimizer’s state is not in it

What it refuses

loading something this graph does not have, by name
loading a weight of the wrong shape, with both shapes
and a refusal leaves the net as it was rather than half loaded
exporting or loading what is trained where it runs, which is not here
but one trained here is exported like any other

Putting several together

the average of one is that one, and of two is halfway between
sizes is what it weighs by, and ten times the data pulls ten times
what is not a number you can halve is not halved
two different networks, the same node with a different shape, nothing at all, and the wrong number of sizes: all refused before anything is computed
three clients that each see a corner average into one that reads the union better than the one that stayed home
and a round leaves every client at the same weights
three mutations — returning the first export, ignoring sizes, averaging the integers — every one caught

The store, opened by hand

The second piece, and what makes the first one reach another machine. Until now the store was a string you handed the engine — forward(store=...) — a place it kept things in and nobody else could open. That is enough while the only thing kept is what the engine decided to keep, and it stops being enough the moment a training run has weights of its own to write down.

store.keep("round/3", trainer.export())
trainer.load(store.recall("round/3"))

Store(directory) has the four the Rust trait has, one for one, dealing in bytes — put, get, bind, resolve, and bound() to look — and two more dealing in values, tensors included, by the codecs. The two are Keeper’s vocabulary and not new, and they are there because the thing this was opened for is a map of tensors: without them everybody writes their own torch.save, and whoever gets weights_only wrong writes a way into a shared directory.

Two things it found by being used:

An export’s keys were integers. The positional duck numbered them, and the one thing a map that crosses anything here may have for a key is text. Found by handing an export to a Store, which is what it was built for.
An edge and a store do not want the same rule. On an edge a bare tensor is refused and that is a feature — the mistake has two right answers, convert it or say Opaque and mean it, and refusing makes the cost of the first visible. In a store it is bytes either way, so the refusal defends nothing. One private Unknown::{Refused, Wrapped} is the whole difference.

The two tests that matter run a real second interpreter: it trains, writes what it learnt, and this one reads it back to the same weights — a federated round’s client half with nothing between the two but a folder both can see. And twice over: the same weights written by two processes are one blob, which is what makes a round that changed nothing free.

Claiming, which is how work gets handed out

The third piece. bind replaces, and that is right for what it is for — a name is a question and its answer can be refreshed. It is exactly wrong for handing out work: two processes given the same round would both bind it, both do it, and nobody would do the next one.

me = store.put(f"{gethostname()}/{getpid()}".encode())
if store.claim(f"round/{r}/client/{k}", me):
    ...

One operation that either takes the name or finds it taken, on the trait and with no default: one written out of resolve and bind would be a race with a doc comment on it. link and not rename is the whole difference — a rename replaces, and link fails when the name is taken. It is also the one that has always been trusted over NFS, where O_EXCL has not, and a network folder is where this is going to live.

The tests really race. Eight threads in Rust, eight processes in Python, because that is the case this exists for: Slurm tasks on a folder they all mounted. And not just that one wins — that the one told it won is the one written down, or the winner does the work with somebody else’s name on it. Against the mutant written as resolve then bind, every other test still passes and these say seven of the eight racers were told they had it.

The barrier, with nobody in charge

The last piece, and the half the study’s design never had: trials are independent and rounds are not.

for r in range(rounds):
    trainer.fit(my_data)
    trainer.load(gather(store, trainer.export(), run="cifar", round=r,
                        clients=4, mine=int(os.environ["SLURM_PROCID"])))

The same script on every machine. The obvious answer is a coordinator, and a coordinator is a process that has to stay alive — and a run that hangs over a weekend when it does not. Instead: whoever finds the round complete claims the averaging, and exactly one can win that, because that is what a claim is. One client does a little more work than the rest, once a round, and nobody babysits anything.

A round, on disk:

<run>/round/<r>/client/<k>    what client k learnt   (its size in the record)
<run>/round/<r>/averaging     who is doing the mean  (claimed, so exactly one)
<run>/round/<r>/average       the mean               (what everybody leaves with)

The deadline says who is missing by name, and in two messages rather than one, because they are not the same thing: somebody missing is a client that never came; nobody missing is worse — the round is complete, so whoever claimed the averaging died holding it, and no one else will try. That is what a claim is, said from the other side.

Waiting longer is a number. Going on without them is a policy, and not this function’s to make: fedavg takes whatever list it is handed.

Questionnaire, the distributed half

The store, opened by hand (soma-python/tests/test_store.py)

bytes by what they are, names that point at them, and both directions of each
a value with tensors in it goes in and comes out alive, and a bare tensor is kept although it would not cross an edge
something nobody registered a codec for says which type it was
a training run written down by another interpreter is read back here to the same weights, and two processes that wrote the same weights wrote them once

Claiming (soma-store/tests/unit/local.rs, test_store.py)

a name nobody has can be claimed and one somebody has cannot
what bind replaces, claim refuses
eight threads and eight processes on one name: exactly one wins, and the one told it won is the one written down
a claim leaves nothing behind in the temporaries
against the mutant written as resolve then bind, seven of the eight racers were told they had it

The round (soma-python/tests/test_round.py)

the only client of a round averages it itself, and publishes it
a client that arrives after the average finds it and does not wait
two runs sharing a directory are two runs, and so are two rounds of one
the deadline says which clients are missing — and says a different thing when the round is complete and the average never came
the size travels in the record; all of them saying it or none
four processes, one folder, two rounds: all four leave with the same number, and exactly two averagings happen, which nobody arranged
a client that never starts stops the others by name
four mutations — bind for claim, never waiting, never publishing, not weighing — every one caught

What is NOT in it yet

FedProx, FedYogi, SCAFFOLD, the same shape with different arithmetic · secure aggregation · partial rounds, which is a policy and belongs to whoever is running one · and what a round is worth remembering by, which is where the digest of a state stops being only a cache key.

None of it went through the graph’s transport, and that was the point: no Plan::Remote, no port, no protocol. A folder, and Slurm hands the work out.