Training Architecture - Agora

Training Architecture

Three layers compose an Agora training run. Discovery runs a stateless DHT for peer-to-peer addressing. Compute is the pipeline itself: workers, each holding one stage's parameters, running forward and backward. Coordination is a small set of CPU-only trainers that route microbatches through the pipeline and shard the dataset across them. The full system runs on heterogeneous, untrusted hardware without any one contributor holding the complete model weights.

Discovery Layer

Seeds

Stateless. Routing only, no model data.

Peer Discovery

Compute Layer

Training Pipeline

Workers hold one pipeline stage's parameters, process fwd/bwd, run periodic SPARTA state averaging within their stage.

Model split across stages — each worker holds one stage's parameters, not the full model.

Activations → Forward pass: Head → Body → Tail Backward pass: Tail → Body → Head (grad_input only)

Coordination Layer

Trainers

Microbatch coordination, load balancing, data management.

Pluralis-owned only — not on contributor nodes.

Three-zone Agora architecture: Discovery (Seeds), Compute (Workers), Coordination (Trainers). Forward activations flow Head → Body → Tail; backward activation gradients (grad_input) flow Tail → Body → Head. Parameter gradients never cross stage boundaries.

Component Deep Dive

Workers

A worker is a single process holding one stage's parameters (one or more transformer layers). It performs forward and backward on those parameters, runs its own local optimizer to apply gradients, and joins same-stage peers in periodic AllReduce rounds for state averaging. A worker has no knowledge of the rest of the pipeline: only its own stage.

Worker: Stage X

Worker internals: six co-resident components inside a single Worker process. Runtime drives ModuleBackend for compute; the SPARTA Optimizer coordinates the parameter-averaging step with same-stage peers via the shared DHT.

DHT

Agora uses Hivemind's Kademlia DHT for four functions: peer discovery, expert registration, progress tracking, and matchmaking for AllReduce.

ModuleBackend

The nn.Module for this stage and the forward and backward functions the Runtime invokes. Also owns the two task pools (forward and backward) where incoming trainer requests queue up before the Runtime processes them.

Async SPARTA

Each worker accumulates gradients from its own backward passes and runs its own optimizer step locally; there is no per-step gradient AllReduce. Same-stage replicas drift apart as a result. To re-synchronize, every 20 local steps the worker matches with same-stage peers over the DHT and AllReduces 5% of its parameters. Successive rounds cover non-overlapping slices, so the full parameter set has cycled through over a 20-round window.

Connection Handlers

gRPC listeners that receive trainer requests and put each batch into the right queue: forward or backward. Multiple listeners share a single port.

DHTHandler

A background thread that keeps re-announcing this worker in the DHT under its stage-prefixed UID for trainer and peer discovery.

Runtime

The main loop. Dequeues batches from the forward and backward queues and runs them through ModuleBackend.

Batch processing

Once running, the Worker runs an event loop processing batches from trainers:

  1. Trainer sends a forward request via gRPC → Connection Handler places it in the forward queue.
  2. Runtime dequeues the batch → calls ModuleBackend.forward() → returns the output.
  3. Trainer sends a backward request with gradient outputs → Connection Handler places it in the backward queue.
  4. Runtime dequeues the batch → calls ModuleBackend.backward() → triggers the optimizer step.

Trainers

The trainer's role looks like an ordinary PyTorch training loop: forward through the model, compute a loss, call backward. The difference is that none of those calls run locally. Every forward and backward is dispatched over the network to a worker holding the relevant stage. The trainer's responsibilities are to track the full pipeline topology, select a healthy worker for each stage, and route activations forward and activation gradients back. Parameter gradients themselves never leave a worker.

Training flow

Startup:

  1. Loads configuration and tokenizer config.
  2. Prepares the dataset from Hugging Face.
  3. Creates DHT connections using seed peers. Each model stage has its own dedicated DHT.
  4. The Trainer discovers workers in each stage; all stages must show at least one available Worker before training can start.

Training loop:

For each batch, the trainer iterates through the pipeline in order:

hidden = head.forward(input_ids[:, :-1])
hidden = body1.forward(hidden)
hidden = body2.forward(hidden)
loss   = tail.forward(hidden, input_ids[:, 1:])   # shifted labels for LM
loss.backward()                                   # triggers backward on all stages

Worker selection uses a min-heap keyed by accumulated virtual runtime: the least-loaded worker is selected. When a worker finishes a request, its runtime is credited with the task's estimated duration and the worker re-enters the heap.

Each pipeline stage gets its own DHT connection on the trainer, so worker discovery for stage head is independent of discovery for stage body3. The full protocol involves periodic node announcements for peer discovery.