Memory & Communication#
This page covers two questions: which tensors reside on which hardware during training, and which of them cross the network for the swarm to make progress.
What Lives Where#
Under the default config (offload_optimizer: true, use_mixed_precision: true) the Adam moments are offloaded to host RAM and the forward pass runs in BF16. The full layout:
| Tensor | Dtype (mixed_precision=true) | Dtype (mixed_precision=false) | Location (default) | Location (offload=false) |
|---|---|---|---|---|
| Model parameters | FP32 | FP32 | GPU | GPU |
Param gradients (param.grad) |
FP32 | FP32 | GPU | GPU |
| Adam exp_avg | FP32 | FP32 | CPU | GPU |
| Adam exp_avg_sq | FP32 | FP32 | CPU | GPU |
| Saved activations | BF16 | FP32 | GPU | GPU |
| Forward / backward matmuls | BF16 (tensor cores) | FP32 | GPU | GPU |
| AllReduce buffers | FP32 | FP32 | CPU | CPU |
use_mixed_precision: true wraps the forward pass in torch.autocast(device_type="cuda", dtype=torch.bfloat16), so matmuls run on BF16 tensor cores and saved activations are stored in BF16. Each backward op runs in the dtype its matching forward op used. Parameters and param.grad remain FP32 throughout: PyTorch enforces matching dtypes between a parameter and its gradient.
GradScaler (the mechanism mixed precision normally uses to prevent FP16 underflow) is unnecessary here. BF16 has the same 8-bit exponent as FP32, so its dynamic range is identical and gradient underflow is not a risk.
Native BF16 tensor cores are required (compute capability ≥ 8.0): A100, H100, L40S, RTX 3090/4090, and similar. T4 and V100 do not have them; because they would emulate BF16 and run slower than FP32, those cards are excluded.
Activation Recomputation#
When a backward pass reaches a head or body worker, that worker re-runs its own forward pass before calling torch.autograd.backward(). (The tail is the exception: it computes the loss and its own backward inside the same call, so it never needs to recompute.) This is unusual and easy to mistake for gradient checkpointing, but the two mechanisms differ; this section explains why.
PyTorch's autograd graph is process-local. The graph holds references to in-process Python objects, intermediate tensors, and the grad metadata attached to them, none of which is serialisable across a network. In a pipeline-parallel run that spans multiple processes (typically multiple machines), the activations leaving one worker and arriving at the next stage cannot carry their autograd graph with them.
To send activations between stages over libp2p gRPC, Agora detaches them at every pipeline-stage boundary. The receiving worker sees a fresh leaf tensor with no history of how it was produced.
That detachment leaves a worker's local autograd graph severed at the input. When the backward pass arrives carrying a gradient with respect to that stage's output, there is no graph for autograd to traverse, and calling .backward() directly would do nothing.
The fix is to rebuild the graph on demand. At the start of the backward call, the worker re-runs the forward pass on its saved input, which constructs a fresh autograd graph local to this process. Backward then traverses that graph normally, populating param.grad for the worker's parameters.
Not gradient checkpointing
Gradient checkpointing (torch.utils.checkpoint) is a memory optimisation: skip saving activations during forward, recompute them lazily during backward, trade ~33% extra compute for ~60-80% less activation memory. Agora's recomputation looks superficially similar but is forced by the cross-process pipeline boundary, not chosen for memory. Activations are fully materialised in this worker's memory either way.
Memory Footprint#
Workers fit on a 24 GB consumer GPU under the default config; see Requirements for hardware bounds. Host RAM holds the offloaded optimizer state, which is why the requirements call for 70 GB system RAM per GPU.
Communication Patterns#
A training run has five distinct types of network traffic. Discovery is slow and periodic. Activation traffic is per-microbatch. Batch-size accumulation and matchmaking use the Metadata Store with minimal overhead. State averaging is rare, heavy, and synchronised between same-stage peers. Each is described in turn below.
Periodic node announcement (peer discovery)#
Workers re-announce themselves in the Metadata Store on a timer. A worker publishes under its stage-prefixed UID (head.0.0, body1.0.1, tail.0.0), tagged with its current sync_phase. The announcement thread refreshes every ~30s.
Each worker also maintains a live view of the next stage's active peers — the candidates it can route microbatches to — refreshing that roster on the same ~30s cadence. There is no leader election: a peer joining or leaving is a record write that other peers observe on their next refresh.
Forward / backward activation traffic#
Per microbatch, each worker pushes directly to a worker it selects in the next stage, in pipeline order:
- Forward payload: input tokens (originated on the head from its data shard) or hidden states from the previous stage (body / tail).
- Backward payload: activation gradients, pushed back along the same route in reverse, starting from the tail once it has computed the loss.
Activation gradients flow back stage-by-stage during the backward chain. Parameter gradients never traverse the network: they are accumulated locally and consumed by each worker's own optimizer.step(). The only cross-worker averaging is the sparse state AllReduce described below. See Activation Recomputation for how the backward chain works.
Batch-size accumulation sync (ProgressTracker)#
A worker cannot decide on its own when to step the optimizer; that decision is collective per stage. Each worker publishes its running sample count to the Metadata Store via ProgressTracker. When the per-stage sum crosses target_batch_size, every worker in the stage takes its optimizer step in lockstep.
Matchmaking (pre-cursor to SPARTA averaging)#
Before a SPARTA round runs, same-stage peers agree on a group: who is participating and which round it is. The matchmaking protocol does this through the Metadata Store:
- Peers register their willingness to average this round.
- Each peer polls until the expected participants have registered.
- Each peer sorts that set deterministically, so every member derives the same group on its own.
Matchmaking runs once every 5 local optimizer steps.
SPARTA state averaging (AllReduce within the matched group)#
Once a group is formed, its members run a butterfly AllReduce on a sparse parameter subset:
- Sparse: 5% of the parameter set is averaged per round.
- Rotating partitions: a different 5% slice each round, covering the full parameter set in 20 rounds.
- BF16 on the wire: parameter slices are cast to BF16 for transmission (2 bytes per parameter).
Optimizer state is not periodically averaged. AdamW's exp_avg and exp_avg_sq stay local during training; they are only transferred during the initial state download when a worker joins the run.
Bandwidth per averaging round#
For a body stage of the Pluralis-13B (~907M trainable parameters), each round transmits 5% × 907M × 2 bytes ≈ 91 MB per peer. The head is close behind at ~913M trainable parameters (~91 MB per round), carrying four layers plus a low-rank reparametrized token embedding. The tail is the largest stage at ~1.35B trainable parameters (~135 MB per round), carrying two layers plus a full-rank output projection, and its transformer layers use uncompressed projections rather than the low-rank reparametrizations used in the head and bodies. Spread across the group's butterfly all-reduce, this still fits within the WAN bandwidth budget.