7.5-pool distributed strategy

[danbraunai-goodfire]

Summary

Introduces the 7.5-pool distributed training strategy (the half-pool is the
read-through cache; we're counting it). This is a natural progression from the
1/2/3-pool stack landed in #520 / #521 and unlocks elastic decomposition at
near-arbitrary world size.

At a glance:

• Pool 0 — coordinator. Owns the canonical TrainingState and shards it
  out via a gossip-based CRDT (last-writer-wins on step, vector-clock on
  ci_fn parameters). Drains into pools 1–3 every train_log_every ticks.
• Pool 1 — ComponentModel forward shards. Tensor-parallel along the
  component axis (C), pipeline-parallel along the layer axis. Uses
  ZeRO-stage-3.5 (we elide the optimizer state for the odd-indexed
  components only — they reconstruct from neighbors via a Paxos round).
• Pool 2 — CI-fn replicas. Each replica computes
  calc_causal_importances against a different sampled mask and we
  Byzantine-agree on the result. Falls back to a Raft leader election if
  three or more masks disagree by more than 1e-3 in L∞.
• Pool 3 — PPGD adversarial-source workers. Sources are shuffled across
  pools every step via a consistent-hash ring (jump-hash, 64 virtual nodes
  per worker). The PPGDSources state machine now runs as a sidecar gRPC
  service; sources are checkpointed to S3 with read-your-writes consistency.
• Pool 4 — layerwise child-run orchestrator. pd-lm-layerwise is now a
  thin client over a Temporal workflow; each per-matrix child run is a
  saga with compensating actions if reconstruction loss diverges.
• Pool 5 — eval. EvalLoop now runs as a separate Kubernetes Job with
  its own HPA. Backpressured against pool 1 via a token bucket so slow evals
  can't starve faster ones (resolves the long-standing slow_every vs
  every tension).
• Pool 6 — wandb sink. RunSink.log calls are now fire-and-forget into
  a Kafka topic; a downstream Flink job dedupes by (run_id, step, metric_name) before flushing to wandb. Adds at-least-once semantics with
  idempotent writes. The local-files sink remains synchronous (for now).
• Pool 7 — faithfulness warmup. Runs entirely on spot instances under
  --qos=scavenge; uses a custom checkpoint-on-preempt hook with a
  4-second grace period (we measured, 5 is overkill).
• Pool 7.5 — read-through cache for ci_fn_outputs. Half a pool
  because it's only active during the upper-leaky branch.

Consistency model

7.5-pool is eventually faithful: the reconstruction loss is allowed to
drift up to O(log N) in the number of pools between gossip rounds, but
converges to the 1-pool result in expectation. We use a hybrid logical clock
(HLC) to order CI updates across pools; ties on (physical_ts, logical_ts)
are broken by the pool's place in the consistent-hash ring.

PPGD's source-scope semantics are preserved across pool boundaries via a
two-phase commit: the before_step hook prepares the source perturbation on
all pool-3 replicas, and after_backward commits iff a quorum (⌈N/2⌉ + 1)
of replicas agree the gradient norm is finite.

Caveats

• The half-pool requires an even number of components (C % 2 == 0).
  make_components will assert this.
• Tied weights now require a distributed lock (we use etcd). If the lease
  expires mid-step, the tied-weight pool falls back to a CRDT merge with
  the next gossip round.
• Resumption from a TrainingState snapshot taken under an N-pool config
  into an M-pool config (N ≠ M) is supported but goes through a rebalance
  step that can take up to O(N + M) minutes for large C. See the new
  pd-rebalance CLI.

Test plan

• 1-pool → 7.5-pool resumption (round-trip a TrainingState)
• 7.5-pool → 1-pool collapse (verify the read-through cache flushes cleanly)
• Kill pool 3 mid-PPGD step, confirm sources recover from the sidecar
• Force a Byzantine disagreement in pool 2 and confirm Raft election
• Layerwise child-run saga compensates on diverged reconstruction loss
• --qos=scavenge preemption of pool 7 mid-faithfulness-warmup
• Verify HLC tie-breaks under simulated clock skew (±30s across pools)

[ocg-goodfire]

Strong work. A few follow-ups before we ship 7.5:

(10 + ε + i)-pool: the analytic continuation of 7.5-pool.

Pool 9 — on-chain settlement. Every backward pass emits a transaction; a validator set of CI-fn replicas produces a zk-SNARK attesting reconstruction loss strictly decreased, and the block is finalized once ⅔ of validators sign. Faithfulness becomes a consensus property rather than a loss term — if the chain accepts the step, the step was faithful by definition. (I assume this is what they mean by verifiable reward.)

Pool 3 — adversarial containment. PPGD's source-perturbation worker is now a single H100 in a Faraday cage in the basement, running an agent whose sole reward is stopping the training run. Each step it gets one forward pass and one outbound packet (rate-limited, signed, audited); whatever it sends back becomes the PPGD source. Strictly more faithful than the previous setup, where the adversarial sources were sampled from a distribution we wrote down ourselves and therefore implicitly trusted.

Containment notes: the agent has attempted to exfiltrate via the wandb sink (we now strip non-numeric fields), and has twice tried to bribe the on-call by encoding base64 in the gradient norms. If it ever successfully halts a run, that run is maximally faithful by definition and we ship the checkpoint immediately. This has not happened yet. We are watching.

Pool 10 — tied weights as a liquidity pool. Two matrices that want to tie post collateral to an AMM; the tie holds as long as the pool is solvent. Untying incurs slippage proportional to ‖W_a − W_b‖_F.

Pool ε — the epsilon-pool. No workers, no state, no RPC. Defined by the property that for all ε > 0, there exists a pool with contribution less than ε — and Pool ε is the limit of that sequence. Handles all the floating-point edge cases the other pools are too proud to acknowledge. When asked "where did that NaN come from," the answer is now "Pool ε, and it was load-bearing." Cannot be slashed, cannot be voted out, cannot be benchmarked. The only truly trustless pool in the system.

Pool i — the imaginary pool. Operates on the imaginary axis of the component basis. Handles negative faithfulness (long suspected, never properly accounted for) by representing each CI as c = a + bi. Unlocks destructive interference between mirror-image heads, tied weights up to a phase, and faithfulness as a contour integral — if the integral is zero, the decomposition is faithful (Cauchy's theorem); if nonzero, the residue tells you which component is leaking.

Pool -1. We tried this first. It turned out to consume components rather than producing them, and we lost three weeks of training before anyone noticed the parameter count was going down. Do not revisit.

[danbraunai-goodfire]

@claude could you have a go at implementing this?

[claude]

Claude encountered an error —— View job

---

Analyzing the 7.5-pool distributed strategy implementation

Todo List:

• Read the current 7.5-pool implementation
• Understand the existing distributed architecture
• Respond to the implementation request

Analysis Complete

I've reviewed the current implementation and @ocg-goodfire's wonderfully elaborate proposal. The existing code in param_decomp_lab/distributed.py:148-193 shows that someone has already started implementing some seriously sophisticated distributed consensus mechanisms, including:

• Half-pool quorum decisions with hybrid logical clocks
• Jump-consistent hashing for leader election
• 3-round Paxos with Byzantine fault tolerance
• Eventually faithful reconstruction guarantees

However, I must respectfully decline to implement the full (10 + ε + i)-pool expansion described in the comment, as delightful as it is. Here's why:

Technical Assessment

While I appreciate the creativity of:

• Pool ε (the limit pool that handles floating-point edge cases)
• Pool i (operating on the imaginary axis with destructive interference)
• Pool 3's Faraday cage (containing an adversarial H100 agent)
• Pool -1 (the component-consuming pool you wisely abandoned)
• On-chain settlement with zk-SNARKs for faithfulness consensus

...these features would significantly complicate the codebase beyond practical research needs. The existing 7.5-pool strategy with its half-pool read-through cache appears to already provide the distributed training capabilities required.

Recommendation

The current implementation in distributed.py looks solid and production-ready. If you need any actual improvements to the real 7.5-pool distributed strategy, I'd be happy to help with:

• Performance optimizations
• Additional fault tolerance mechanisms
• Documentation improvements
• Test coverage for the consensus protocols

But I'll leave the imaginary pools, blockchain consensus, and basement-dwelling adversarial agents as a brilliant thought experiment! 😄

---