Federated averaging

Federated averaging (FedAvg) is FlashRuntime's answer to a question the PyTorch guide and the JobSpec & isolation guide both run into: what do you do when the machines that want to help train a model cannot talk to each other? Volunteer nodes run their task containers with --network none (see the repo's docs/guides/donate-a-machine.md) — no LAN, no internet, no way for one container to find another. Coordinated multi-process training (DDP, FSDP) needs the opposite: every rank must rendezvous with every other rank over a process group before the first all_reduce. On a volunteer pool those two requirements are irreconcilable, so FlashRuntime does not attempt coordinated training there at all.

FedAvg sidesteps the rendezvous problem by never requiring it. Instead of ranks synchronizing gradients mid-step, rounds synchronize whole models between steps:

  1. The driver broadcasts the current weights as a plain artifact:// blob.
  2. Each participating node downloads the weights, trains independently on chunks of the data it was offered until its deadline, and uploads a weight delta (not the new weights — see below) plus the list of chunk ids it actually finished.
  3. Once the deltas that committed cover enough of the data between them, the driver averages them, applies the result to the broadcast weights, and starts the next round.

No node ever needs to see another node's IP address, let alone open a connection to it. Every cross-node interaction is a PUT/GET against the coordinator's artifact store, which is exactly the same shape of traffic a volunteer node already does to pull its task inputs and push its results. That is why this is a round loop implemented as a driver chaining ordinary lease jobs (flashml_workloads/fedavg_driver.py, flashml_workloads/fedavg_worker.py) — the same "pipelines are jobs chained by a driver, not a new execution mode" pattern as the sharded-k-means POC — rather than a new backend.

Why a delta, not the new weights

Each worker's task uploads delta.json (the change it made to the weights it started from) alongside metrics.json. The driver averages deltas, not raw weight snapshots, because a delta is a direction that stays meaningful even if the weights it was computed against are no longer the newest ones — the exact situation a straggling volunteer produces when it finally reports in after the round has moved on. Averaging final weights directly would require every worker to have started from the same snapshot; averaging deltas only requires knowing what each worker started from, which the driver already does.

Chunks, slots, and how a round closes

kmeans_driver (the other job-chaining driver in this codebase) requires every dispatched shard to report before it aggregates. FedAvg deliberately does not — and it does not count participants either. A round closes on data coverage.

A pass over the data is cut into total_chunks uniform chunks. Uniform is load-bearing rather than convenient: the reduce clips any contribution whose norm exceeds three times the round's median, and that cap is only silent on an honest round while every contribution covers the same amount of data — so heterogeneity is expressed as how many chunks a machine completes, never how big one is.

A round posts slots task offers. Every offer carries the same chunk budget and the same deadline; what differs is where in the sequence it starts, so two machines working at the same speed cover disjoint data and a machine that outruns its share simply wraps. A slot is not a partition and slots is not a machine count — it is a ceiling on offers, because how many machines are online is not knowable when the round is submitted. Unclaimed slots are never waited for. A machine claims an offer, trains as many chunks as it can before its deadline, and reports the ids it finished.

The round ends when the union of those ids meets its coverage target (sync_every of one pass), or when round_timeout_s fires — in which case it combines whatever did arrive rather than throwing the round away. The only way a round fails is having nothing creditable to combine at all. This is the correct policy for volunteer compute: machines that donate spare cycles are unequal and unreliable by construction — laptops close, Wi-Fi drops — and requiring a fixed number of them would let one closed laptop stall every other participant's round.

Consecutive rounds advance forward through the data. Every slot's start is advanced by the round index times the round's coverage target, so round 1 begins where round 0's target ended. Without that, slot geometry would repeat every round and — since a round stops the moment it hits a target below a full pass — the same prefix of the dataset would be retrained forever while the tail was never trained at all.

The advance wraps rather than running off the end: offsets are (round × target) mod total_chunks, so they cycle. One cycle is exactly one pass over the data — that is the guarantee, and it is a per-pass one, not a per-run one. Training four epochs means walking the same four offsets four times, which is what "four epochs" should mean.

The corollary is what makes coverage-closing safe rather than merely convenient: once the driver has read the round's deltas and applied them, any delta that commits afterward for that round is discarded, never folded into a later round. run_fedavg freezes the contributor set at the moment the round closes and never re-reads that job's artifacts again (fedavg_driver.py, run_fedavg). A late delta was computed against weights that no longer exist by the time it arrives — the model has already moved past them — and applying it on top of a newer round's weights would not be "one more contribution," it would silently corrupt the average with a step that was never actually taken from the current state. Discarding is the honest behavior; a driver that tried to be more "inclusive" here would be quietly wrong instead.

tests/test_fedavg_convergence.py::test_round_completes_on_quorum_when_a_node_never_reports pins exactly this: three slots are offered but the test's agent pool is capped to exactly two successful claims and then stops claiming, so the third offer is never bound to any node and sits PENDING for the life of the test. The round still aggregates on the two that committed — with an exact participants == 2 assertion — rather than hanging until the deadline waiting for an offer nobody was ever going to take. (The cap on claims, not the node count, is what makes the third offer genuinely abandoned: either registered node can claim any offer, so without the cap both nodes could sequentially serve all three before the driver's poll notices coverage.)

Nothing sets a floor on the contributor count

There is no min_participants, and no other floor anywhere. A single fast machine that covers the target alone closes a round by itself, and the backstop can close a round on one contribution. This is deliberate — a floor is exactly the "one closed laptop stalls everyone" failure the coverage model exists to avoid — but it has a cost worth stating plainly: the median clip in fedavg_weights is robust statistics, and robust statistics need a majority of contributors to be honest. At two contributions the median is the mean, which an attacker moves directly; at one there is nothing to compare against. The influence cap is protection at a healthy round and is not protection at a thin one. Result verification is a separate, unbuilt concern.

What counts as a participant, and what the driver refuses

Everything a volunteer node produces — the delta, the reported chunk ids, the sample count, the metrics file, the filenames — is attacker-controlled input. Result verification (catching a node that lies about a delta it honestly computed) is a later milestone, but input validation and containment are not deferred:

Artifact PUT is now authenticated and lease-scoped when the coordinator sets FLASHML_NODE_TOKENS (the per-machine-token slice): a node token can only write under jobs/{job}/{task}/ for a task it currently holds a live lease on. The round-weights key (jobs/{job_id}/round-{round:03d}/weights.json) belongs to no task and no node's lease, so a plain node token cannot write it — the driver instead authenticates with an operator token (FLASHML_OPERATOR_TOKENS), which is attributable but deliberately not lease-scoped, exactly because drivers are legitimate writers outside any lease (see docs/guides/donate-a-machine.md). Result verification is still a separate, unbuilt concern: this scoping stops an unrelated node from overwriting the round weights, not from a participant lying about the delta it honestly computed.

The flashml.yaml shape

A federated-averaging round is submitted as an ordinary lease-mode job:

apiVersion: flashml.dev/v1alpha1
kind: Job
metadata:
  name: fedavg-r000
spec:
  execution:
    backend: leases
  image:
    repository: local/tier1
    tag: dev
  workload:
    type: federated_averaging
    parameters:
      round: 0
      # The round's chunk plan, stated once for every slot it expands into.
      slots: 2            # offers, a ceiling — not a machine count
      total_chunks: 4     # how finely one pass is cut
      chunk_budget: 2     # chunks each offer may walk
      chunk_offset: 0     # (round * coverage target) mod total_chunks
      deadline_s: 96.0    # when a worker must stop training and upload
      local_steps: 20
      lr: 0.1
      batch_size: 16
      seed: 0
      in_dim: 8
      hidden: 16
      out_dim: 2
      dataset_size: 256
      # weights: artifact://jobs/<prev-job>/round-000/weights.json
      # (omitted on round 0 — each worker seeds its own model from `seed`)

The expansion (service/modea._expand_fedavg) turns that into one task per slot, each carrying the same budget and deadline and its own start_chunk — chunk_offset is consumed there and never reaches the worker. It fails the submission outright on anything the worker would only discover inside a container on a volunteer's machine: a missing worker parameter, a non-positive chunk_budget, a missing or non-finite deadline_s, or a total_chunks larger than dataset_size (which would leave chunks with no rows in them at all). One bad body must not burn every attempt of every slot across the fleet and read as a fleet-wide node fault.

Almost nobody writes that block by hand. run_fedavg derives all of it — you give it epochs and sync_every (both in passes over the data), total_chunks, slots, and expected_machines, and the round count, the per-round coverage target, the chunk budget, the offset and the training deadline all fall out:

run_fedavg(
    HttpCoordinator(coordinator_url),
    epochs=4,              # how much training you want
    sync_every=1.0,        # how often to combine, in passes
    total_chunks=4,        # how finely one pass is cut
    slots=2,               # offers per round (a ceiling)
    expected_machines=4,   # who is expected to CLAIM — sizes the budget
    worker_params=WORKER_PARAMS,
    initial_weights=initial_weights,
)

expected_machines is separate from slots because the two pull in opposite directions: the offer ceiling should be generous, and dividing the round's work by a generous ceiling would starve every machine that actually turned up. Nobody types a round count and nobody types a machine count.

isolation.tier is left at its default, "standard", deliberately: unlike the argv runner tier for arbitrary bring-your-code jobs, a federated_averaging task's payload is a fixed, trusted module execution (flashml_workloads.fedavg_worker), so it does not need the sandboxed argv path and its argv_capable gate. A node only needs module_capable (fail-open — absent counts as capable) to be eligible. run_fedavg builds this JobSpec once per round and submits it as a new job (flashml_workloads/fedavg_driver.py:_round_body) — the round number, its chunk_offset and the weights URI are the only things that change between the driver's own resume points. The image and isolation tier are run_fedavg parameters (image=, isolation_tier=); the defaults above are this repo's e2e fixture image, which only works because SubprocessRunner ignores image entirely — a docker-tier volunteer needs a real, pullable reference.

What this proves — and what it does not

tests/test_fedavg_convergence.py runs this loop against a real coordinator over real HTTP: real job expansion, real leases, real local artifact storage (FLASHML_LOCAL_ARTIFACTS_DIR), and real commit-time sha256 validation on every uploaded artifact. Two independent worker "agents" (a few lines of urllib, standing in for flashnode work — see the test file's docstring for why an in-repo test cannot import flashnode directly) pull leases, train, and commit without ever talking to each other. The measured per-round mean loss across four rounds with two participating nodes:

round 0  participants 2/2  covered 100%  mean_loss 0.3501
round 1  participants 2/2  covered 100%  mean_loss 0.1615
round 2  participants 2/2  covered 100%  mean_loss 0.1032
round 3  participants 2/2  covered 100%  mean_loss 0.0797
converged: 0.3501 -> 0.0797 over 4 rounds

covered is the number the round actually closed on, printed rather than assumed: 100% here because four chunks over two slots at expected_machines=4 gives each offer half the pass, so the two nodes cover it between them. A round that had to fall back on its backstop would print a lower figure, and that is exactly when you want to see it.

(scripts/fedavg_local_demo.py reproduces this and exits non-zero if the final round's loss is not below the first — a demo that prints numbers nobody checks is not evidence. The run above is that script's own output.)

Read that number correctly: this proves collaborative training, not faster training. Two nodes did not finish training in half the wall-clock time of one — they trained sequentially through four rounds, each doing its own local steps, and the loss came down because their independently computed updates were combined. Nothing here claims a throughput or speed-up result; DDP/FSDP make that claim, on a coordinated pool that can rendezvous, and that claim is out of scope for volunteer nodes entirely (the repo's docs/guides/donate-a-machine.md has the full list of things the volunteer pool does not attempt, including "no coordinated multi-process training"). What FedAvg proves is that machines which cannot see or trust each other — and in the volunteer case, cannot even reach each other over the network — can still jointly move one model's loss in the right direction, coordinated entirely through the coordinator's leases and artifact store.