Codeseys's picture
Wave 11: cross-model adversarial review + honest down-revision
f16fa23
|
Raw
History Blame Contribute Delete
4.97 kB

Spike 008 — VERDICT

Status: ⚠️ PARTIAL (acceptance criterion redefined; see § "Honest re-statement" below) Date: 2026-05-26 Wave: 9 Cross-model review: BLOCKER 2 of docs/research/WAVE_7_10_FINAL_REVIEW.md

Headline

make_diloco_outer_loop() wraps torchft.local_sgd.DiLoCo (BSD-3, Meta-maintained) and the framework's outer-loop sign convention is pinned by an explicit unit test. Cross-replica convergence is NOT verified — the BACKLOG acceptance criterion required two replicas; what shipped is one replica + passthrough no-op allreduce.

Honest re-statement

The BACKLOG required:

Smoke test: 2 replicas × 4 inner steps × 2 outer rounds on the toy model from Spike 005, both replicas converge toward the same solution within tolerance.

What was attempted: the recon doc (DILOCO_RECONNAISSANCE.md) provided a "ready-to-paste" 2-replica pattern with a shared-buffer mock allreduce that averages tensors across replicas. That pattern does not work in single-process: each inner.step()'s post-hook runs prepare_sync + perform_sync to completion (including outer optimizer step) before yielding back to the caller. By the time replica B's post-hook starts, replica A has already finished its outer step using A's un-averaged pseudo-gradient. The mock allreduce can compute the cross-replica mean, but it can't write that mean back into A's _grads[name] buffer in time for A's outer step.

The fix would require either:

  • A real torch.distributed barrier (NCCL or Gloo) — out of scope for a CPU-only single-process smoke.
  • A multi-process test using torch.multiprocessing.spawn with two real processes — feasible but ~200 LOC of additional test infrastructure that would need its own review.

What ships instead: a single-replica machinery test (allreduce is a no-op passthrough). Verifies that outer optimizer fires, Nesterov state populates, sign convention is correct. Does NOT verify cross-replica convergence. This is a redefinition of the BACKLOG acceptance criterion. Documented explicitly in this verdict + the test file's _make_passthrough_manager docstring + composer_diloco.py.

What the test suite DOES verify (5/5 pass)

Criterion Status
allreduce/start_quorum/should_commit fire at the right step boundaries ✅ test 1
Nesterov momentum state populated for every parameter ✅ test 1
Pseudo-gradient sign convention verified (θ_initial − θ_local) with explicit arithmetic prediction ✅ test 2
No regression in Spike 005 imports ✅ test 3
make_diloco_outer_loop() factory wraps the right object ✅ test 4
Streaming DiLoCo with 2 fragments + nonzero fragment_sync_delay accepts the config ✅ test 5

Sign convention pinned (the most important result here)

Per torchft's _save_grads() (line 324 of torchft/local_sgd.py):

pseudograd = θ_initial − θ_local

DiLoCo defines pseudo-gradient as θ_initial - θ_local. This is the negative of the local update direction. Standard SGD subtracts gradients (p ← p - lr * grad), so the outer step moves in the local-update direction. No negation needed in our outer optimizer wrapper.

The test exercises this exact math with local_param_after_nudge = θ_initial + 0.5 and asserts final ≈ θ_local. A sign flip in either _save_grads or the outer optimizer would land at θ_initial - 0.5 (movement in the wrong direction); the test reports both values in the failure message so a future flip is immediately diagnosable. This is the single best test in Wave 7-10 per the cross-model reviewer.

What this CLAIMS to close

  • V2 (DiLoCo "deferred to v0.2") in docs/VISION_VALIDATION.md — re-scored as ⚠️ partial, not ✅, in the 2026-05-26 update at the bottom of § 3 of that doc.

What this does NOT close

  • True multi-replica convergence — see § "Honest re-statement" above. Either needs a real torch.distributed test on multiple processes, or a redesigned single-process pattern that overrides _DummyWork.wait() to do lazy averaging.
  • Trainer integrationComposerReplicationTrainer does NOT yet use make_diloco_outer_loop. The DiLoCo wrapper is an independent context manager. Wiring it into the trainer's lifecycle is a separate spike.
  • Streaming DiLoCo with fragment_sync_delay > 0 (overlapped sync, CUDA streams). The framework's make_diloco_outer_loop() accepts the parameter; tests only exercise delay=0 (vanilla DiLoCo).

Files

  • composer_diloco.pymake_diloco_outer_loop() wrapper. Documents the sign convention.
  • tests/test_diloco_smoke.py — 5 acceptance tests. Test 2 (sign convention) is the highest-value test.

Dependencies added

  • torchft-nightly (BSD-3, Meta-maintained, pip install torchft-nightly)

Cost / time

  • Pure CPU, single process, no GPU.
  • Test suite: ~5 seconds for 5 tests.