diloco_features: paper-form Async + Streaming DiLoCo reproduction (small scale)#239
Open
jdinalt wants to merge 20 commits into
Open
diloco_features: paper-form Async + Streaming DiLoCo reproduction (small scale)#239jdinalt wants to merge 20 commits into
jdinalt wants to merge 20 commits into
Conversation
…t config
Restructure the project into a pre-registered, paper-form reproduction study of
async (DN/DyLU/grace, arXiv:2401.09135) and Streaming DiLoCo (arXiv:2501.18512),
measured against synchronous DiLoCo at small scale (4x4090, 34.4M Llama, 10 blocks).
Honest scope: on identical 4090s synchronous DiLoCo is fine; async/grace/DyLU solve
problems this homogeneous rig lacks. The study validates that the mechanisms function
and reproduce the papers' per-token *convergence* trends — it does not demonstrate the
*wall-clock* benefit (the async paper's Fig-2 win), which needs the two-population
(4090+3090) / WAN setup and is deferred to Future Directions.
Config (templates/configs/default.yaml):
- Override the inherited max_steps default (ns.total_steps, a ~16k self-stop) to -1
with num_train_epochs=1, so the server's --token-budget is the SOLE stop authority
for these open-ended runs. Overridable via the --max-steps dynamic arg (validate).
End-to-end smoked: server relays save_and_stop at the budget; workers stop cleanly
well before any step cap.
Run matrix (experiment.sh): 10 single-run arms, token-budget global stop (no per-worker
max_steps), arms tagged [trend] vs [mech]. Grace is cut from the matrix to a
validate-only mechanism check (coalesces near-simultaneous finishers AND proceeds when
alone) — its payoff is wall-clock tail-reduction under heterogeneity, not convergence.
Capture /status live DURING the run (per-worker staleness + the grace histogram vanish
once workers hit the budget and deregister); unwrap the orchestrator-routed
`status --json` (snapshot nested under "status").
analysis/: + perplexity; new streaming.py (strided/sequential at N in {2,5}),
grace_batches.py (mechanism check + all-k guardrail), staleness.py (gate split into
should-be-stale ~ k-1 vs the DyLU-on reducer-vs-control); read status.json
(total_tokens, derived staleness, grace); dn_sweep.py removed.
Related: #232 (CLI --token-budget units inconsistent with the webui's millions).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The CLI `--token-budget` (the `diloco server` flag and the runtime `diloco token-budget` subcommand) was raw tokens with no unit stated in help, while the webui field is millions — ambiguous (#232). The human CLI flag and the webui/orchestrator path share one parser (`build_diloco_server_command` emits the same `--token-budget <raw>`), so redefining a bare number as "millions" would break the webui/orchestrator and is the larger #232 change. Instead, accept an optional case-insensitive K/M/B (or G) suffix with decimals (`2.08B`, `2080M`, `500K`) while a bare number stays raw tokens — non-breaking (the webui/orchestrator emit bare raw counts, unchanged) and unambiguous. Applied to both `--token-budget` sites (metavar TOKENS, units in help). - experiment.sh: budgets now `2.08B` / `1.04B` / `3M`; the start_one budget guard is a string check (was a numeric `-gt 0`, which a suffixed value would break). - harness.sh: add a `token-budget` functional smoke recipe (workers submit with no `--max-steps`, so the config's max_steps=-1 lets the server's small 2M budget be the sole stop authority) and include it in `all`. Addresses the #232 ambiguity non-breakingly; bare-number=millions parity with the webui remains a separate, breaking option. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ference cmd + hyperparam doc Tighten the study to the two papers' communication mechanisms vs the sync baseline. Drop two arms that were not experiments: - #9 `baseline_2w` (2w-vs-4w token-efficiency scaling) — DiLoCo's own batch-scaling is covered in the sibling diloco / pretrain/small-llm projects; doesn't belong here. - #10 `wire_http_pk` (wire/transport lossless) — never an experiment, a one-off validation that the re-implemented gRPC+safetensors matches HTTP+pickle; already done. Retire analysis/worker_scaling.py + analysis/verify_baseline.py with them. Matrix is now 8 arms, all 4-worker. Token budget 2.08B -> 1B total: the 2.08B was an unprincipled "520M/worker x 4" artifact (inherited per-worker length x 4 workers). 1B matches the sibling diloco long reference run and is ~2x Chinchilla — the model's Chinchilla-optimal is 525M tokens (20 x its 26.2M non-embedding params). Single BUDGET (all 4w); BUDGET_2W gone. README: restore the shared **reference command** (the base every arm runs, with per-arm deltas only in the diloco server flags) and add a **"Hyperparameters (and why)"** table documenting every magic number with its rationale — H=100, k=4, outer LR 0.7/mom 0.9, inner AdamW lr 2.07e-4 (base 1.5e-4 x sqrt batch scaling), jitter 0.15, the ~2x DyLU spread, DN N=4, fragments {2,5}, the 1B budget, etc. Remove the stale pre-#225 bug-artifact assets (curves.csv + PNGs; several referenced the removed arms or a deleted script) — Results are pending, regenerated by the run. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Document the actual resolved config values in the §3.1 hyperparameter table: the inner batch is per_device_train_batch_size=8 (set in small.yaml) x 4096-token packed sequences = 32,768 tok/step/worker — 2x the 16,384-token LR reference, which is what upscales the inner AdamW lr from base_lr 1.5e-4 by sqrt(2) to the effective 2.07e-4 (WSD schedule, 1606 warmup steps, min_lr 2.07e-5). Removes a stale "shared with the sibling project" note on the batch (it is in fact what differs). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Two review agents audited the harness + analysis pipeline before the GPU runs; no blockers. Applying the cheap, non-run-invalidating fixes they found: - plot_experiment.py: stale "520M tokens" suptitle -> "1B total tokens". - plot_experiment.py / streaming.py / dylu_control.py: x-axis "step" -> "local step (∝ total tokens)" (the design's primary axis; faithful at fixed batch). - experiment.sh submit_sync: add --heartbeat-interval 5 (parity with async/dylu) so the heartbeat-driven token budget ticks at the same cadence on every arm — tighter equal-total-tokens alignment across the matrix. - experiment.sh: fix the backwards v_budget comment; add a final wait_no_servers so the last arm's server is cleaned up. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Each arm now `rm -rf "$RESULTS/$name"` at the start of do_one (a named path, internal to the already-running script — no glob, no manual sweep needed). This makes re-runs idempotent (a fresh capture dir per arm) and removes any need to manually clear runs/ with a wildcard `rm` between phases (a glob rm trips a confirmation prompt that would stall an unattended sweep). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…obes) JITTER and BUDGET join GRACE_S as env-overridable knobs so the staleness-gate calibration / short probes can be driven without editing the file (e.g. BUDGET=80M ./experiment.sh run async_dn4 for a ~5-min staleness probe). Defaults unchanged (0.15 / 1B). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…/2, not k-1)
The live async_dn4 run exposed a round-robin conflation in the gate. For k
equal-average-speed jitter workers the sync cycle is round-robin, so at any instant
the workers occupy staleness {0,1,...,k-1} (one per cycle phase). The SNAPSHOT MEAN
of (sync_round - last_sync_server_round) is therefore (k-1)/2 = 1.5 at k=4 — the
signature of *full* decorrelation — while the async paper's "staleness ~ k-1" is the
PER-SUBMISSION staleness (how stale a gradient is when applied) = the MAX of the
cycle = 3. The old gate expected mean ~ k-1, which full decorrelation can never
reach (more jitter can't push the snapshot mean above (k-1)/2 at equal speed), so it
would have falsely FAILed the (correctly calculated) jitter 0.15 and triggered a
pointless recalibration loop.
staleness.py now gates the jitter arms (async_nodn, async_dn4) on snapshot
mean ~ (k-1)/2 and reports the per-submission max ~ k-1. The DyLU arms use a delay
spread (not jitter, not round-robin): dylu_off is the un-gated control, dylu_on
succeeds by landing at lower mean staleness. Verified on the live async_dn4 capture:
mean 1.50 = exp, max 3, PASS. README §3.3 + harvest comment updated to match.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
After the de-risk async_dn4 run showed the runtime is affordable, bump the matrix budget to 2B total (~4x Chinchilla; the model's Chinchilla-optimal is 525M from its 26.2M non-embedding params) to capture async's longer-term dynamics, where DiLoCo-family behavior tends to separate. Distinct rationale from the old 2.08B (which was an unprincipled 520M/worker x 4 artifact) — this is a deliberate long-budget choice. ~500M tok/worker, ~150 sync rounds at H=100. Updates the hyperparameter table + plot caption to match. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
A baseline timing probe (80M, compile on) gave a steady-state step time of ~0.18-0.20 s/step. The DyLU speed spread is set so the max per-worker delay (~0.18 s) approximately equals the step time, making the slowest worker ~2x the fastest (a realistic 4090+3090 mix): DYLU_SPREAD = (0 .06 .12 .18). README updated. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… arm Make each arm fail loud rather than limp if a worker dies below the launch count, and stop the long compile/eval phase from starving the heartbeat thread into a false eviction. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Run the 10-arm from-scratch matrix to completion and fill §3.7 Results with real numbers (was _TBD_ placeholders). Adds a DN-buffer depth sweep (async_dn8/dn16 alongside the paper-default async_dn4) — the headline new finding. Findings (single seed, 4 workers, 2B tokens, H=100, from scratch): - Streaming trades per-token quality for wall-clock; strided N=5 closest per token (+0.060) and fastest — nearly matches the baseline on relative time. - DN buffer depth is a non-monotonic lever: N=8 optimum (+0.249 vs baseline), N=4 under-buffered (+0.802), N=16 over-buffered (+0.514). The paper's N=k=4 default is under-tuned for the from-scratch + staleness regime. - async-from-scratch does not reach the sync baseline even at the DN optimum; the warm-start sweep (§6) is the open question. - DyLU neutral here (within seed noise) — the rig's induced spread is too mild. - async_nodn diverges (control). Harness/analysis: - experiment.sh: add async_dn8 / async_dn16 arms (DN-buffer sweep). - harvest.py: register the two new arms (table + staleness set). - analysis/dn_sweep.py (new): DN-depth eval-loss figure (dn_sweep.png). - analysis/plot_walltime.py (new): eval-loss vs relative wall-clock (walltime_comparison.png) — the fair axis for the streaming overlap trade. - analysis/regen_tb.py: per-worker TB regen from captured logs (was untracked). - README: §3.7 table + commentary; abstract TL;DR and §5 hypotheses updated from predictions to measured outcomes; run matrix 8->10 arms. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The central from-scratch finding was that async DiLoCo trails the sync baseline even at the DN-buffer optimum (§3.7 finding 3). This adds the pre-registered (§2.6) test of whether that is the async path or the from-scratch *regime*: pretrain the same arch with plain 4xDDP to ~500M tokens, then re-run the async arms + a warm SYNC baseline from that checkpoint (2B further tokens each). Result — async DiLoCo ≈ sync once warm-started: - Every warm async arm lands +0.03..+0.06 eval of the warm sync baseline, vs +0.25..+0.80 from scratch. The from-scratch regime, not async, was the obstacle — exactly why the source papers warm-start their async runs. - DN-buffer depth stops mattering warm: warm dn4 (+0.064) vs warm dn8 (+0.062), within 0.002 (vs a 0.55 gap from scratch). Small, well-aligned warm pseudo-gradients => staleness barely bites => the buffer is barely needed. - DyLU still neutral warm (within noise). Tooling (warm path, verified end-to-end — server loads the assembled master, workers train from the pretrained loss ~3.2 not ~9): - templates/configs/warm_pretrain.yaml: plain 4xDDP pretrain of the same arch (enable_diloco off, save_safetensors, --total-tokens => global budget). - make_warm_master.py: assemble the trained checkpoint into a root-safetensors DiLoCo server master (copies safetensors, or converts .bin; keys match pristine). - experiment.sh: start_one honors PRISTINE_WARM; WARM_MASTER arm set (warm baseline + warm async dn4/dn8 + dylu off/on), PRISTINE_WARM set per-call so the scratch arms stay on the pristine master. - analysis/harvest.py: register the 5 warm arms (+ warm async in the staleness set). - analysis/warm_compare.py -> assets/warm_compare.png: the gap-collapse figure. - README §3.8 + abstract/§5 hypotheses/§6 updated to the warm-start result. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- Inline the 7 result figures at their findings in the README (loss_comparison, walltime_comparison, streaming, dn_sweep, dylu_control, training_health, warm_compare) instead of referencing them by filename; update Appendix D file list to include the new analysis scripts + plots. - Remove two stale PNGs not produced by this study (baseline_vs_h100.png from an earlier comparison, worker_scaling.png whose analysis script was retired). - Mirror the 7 figures as per-file symlinks under docs/examples/.../assets/ so they render in the mkdocs build (matching the repo's example-mirror convention). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The streaming, async+DN, and DyLU arms are not directly comparable — they test different axes under different conditions — so reporting them in one combined table/plot implied a comparability they don't have. Split §3.7 Results into three self-contained clusters, each with its own table, figures, and appropriate reference: - §3.7.1 Async + DN (equal-speed jitter): baseline + async-no-DN + DN sweep N=4/8/16, vs the sync baseline. loss_comparison/training_health refocused to these arms only (was a cross-axis "headline" mix). - §3.7.2 Streaming (orthogonal — synchronous DiLoCo, fragmented comm): baseline + str2/seq2/str5, its own per-token + wall-clock comparison. plot_walltime.py is now streaming-only (async's wall-clock win is off-rig, §6). - §3.7.3 DyLU (A/B under a ~2x speed spread): dylu_off control vs dylu_on, judged against the control — not the equal-speed baseline. No data changes (same curves.csv); regrouped presentation + regenerated the two refocused figures. Cross-refs (§3.4, §3.8) updated to the new subsections. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
We asserted a ~2x worker speed spread but never showed it. Add analysis/worker_speeds.py, which recovers each worker's median ms/step from the captured logs and confirms the spread is real and correctly ordered: step time rises monotonically with the injected DILOCO_DEBUG_STEP_DELAY (0/.06/.12/.18s), w0 fastest (156 ms/step, 17.7k steps) -> w3 slowest (250 ms/step, 11.9k steps), consistent across both dylu_off and dylu_on (DyLU changes sync cadence, not compute speed). Correction: the realized spread is ~1.6x, not the ~2x the delays were calibrated to target. The per-step CPU sleep partially overlaps asynchronous GPU compute, so a 0.18s nominal delay adds only ~0.09s of wall-clock. Updated every ~2x claim to distinguish the calibration target from the measured ~1.6x (§3.7.3, §3.3, the hyperparameter table, run matrix, captions, experiment.sh). This further weakens the heterogeneity DyLU had to work with, reinforcing the DyLU-neutral caveat. Adds the worker_speeds table + figure to §3.7.3, the analysis script + docs mirror. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Per review: the warm-start runs read better alongside their from-scratch counterparts for joint analysis than as a separate §3.8 coda. Dissolve §3.8: - Warm async (baseline + DN N=4/8) now lives at the end of §3.7.1, with a JOINT scratch-vs-warm table (per arm, gap to each group's own baseline) — the comparison the warm experiment is for. - Warm DyLU A/B (off vs on) now lives at the end of §3.7.3. Also: never overlay the two groups on one loss axis (their y-ranges differ vastly, ~9->2.85 scratch vs ~3.2->2.83 warm, which squashes the warm detail). warm_compare.py's left panel is now WARM-ONLY trajectories on their own scale; the cross-comparison is the right panel's gap bars (deltas to each group's baseline, so comparable). Updated all §3.8 cross-refs -> §3.7.1, and the stale Related-Work arm numbers (DN arms 5-8, DyLU arms 9-10). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…7.3) The ~1.6x spread was too mild to exercise DyLU (no signal). Re-run with a real 4:1 spread, and warm-only (the scratch-vs-warm story is settled by the async arms; DyLU is no exception). Smoke-verified the calibration before the long runs: DYLU_SPREAD=(0 0.24 0.40 0.56) lands a clean 4.00x slowest/fastest step time (the per-step CPU sleep partially overlaps async GPU compute, landing ~delay-0.09s). New result — the mechanism fires, but the convergence payoff is unmeasurable here: - DyLU cuts mean staleness 2.33 -> 1.50 and ~doubles sync rounds (1220 vs 615) — the adaptive sync_every responding to the spread as designed (the staleness gate now PASSES, vs the inverted noise at 1.6x). - But eval is unmoved: warm_dylu_on 2.851 vs warm_dylu_off 2.855 (Δ 0.004, within noise); both reach the warm baseline. The DN buffer already absorbs the staleness, so reducing it further doesn't move per-token loss. DyLU's real payoff is wall-clock tail reduction (off-rig). Harness/analysis: - experiment.sh: DYLU_SPREAD -> 4:1; drop the from-scratch dylu_off/dylu_on arms (DyLU warm-only); smoke-verify note. - harvest.py: drop scratch dylu; worker_speeds/dylu_control/staleness point at the warm_dylu_* arms; warm_compare drops the (now warm-only) DyLU rows from the gap chart. Regenerated dylu_control/worker_speeds/warm_compare. - README: §3.7.3 rewritten (warm-only, verified 4:1, mechanism-fires-but-neutral); matrix is 8 from-scratch arms + a warm set; abstract/§3.1/§3.3/§3.4/§3.5/§4/§5/§6 and the reproduce commands updated. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Per review: comparing raw eval loss across the scratch (2B from random init) and warm (2B on top of 500M) groups is meaningless. Use relative perplexity vs each group's OWN sync baseline — "% worse perplexity than the matched baseline" — which normalizes out the different starting points (and perplexity is the DiLoCo paper's reporting unit). - §3.7.1 warm table: drop the raw scratch-vs-warm loss juxtaposition; show warm eval/ppl + % worse ppl vs warm base, alongside the scratch arms' % worse ppl vs scratch base (+123%/+28% scratch -> +6.6%/+6.4% warm). - warm_compare.py bar chart: relative perplexity (% worse vs matched baseline), not raw eval-loss gap; with value labels. - Perplexity column now in every results table (added to DyLU); all comparison columns expressed as % worse ppl for consistency (diverged no-DN reads "diverged"). Prose/abstract/§5 updated to match. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Two axis corrections to the loss-vs-x plots. X-axis: was worker0's local step count, which is NOT comparable across arms — worker0 is a 1/k-share worker on the equal-speed arms but the FASTEST worker on the DyLU speed-spread arms, so the same global progress logs ~2x the steps (the DyLU baseline appeared to "end early"). harvest.py now maps worker0 step -> AGGREGATE tokens across all workers (step * aggregate_total / worker0_final_step, exact since speeds are constant); curves.csv x column is now `mtokens`. Every arm spans the same 0..2B axis. Y-axis: perplexity (= exp(eval loss), the DiLoCo papers' axis) everywhere, with scale chosen by dynamic range — log for the from-scratch trajectories (early ppl ~1000s would drown the converged tail), linear for the warm trajectories and the tight endgame-zoom panels (plain ticks, best tail contrast). Grad-norm panel unchanged (log). - harvest.py: emit aggregate mtokens as the x; rename CSV column step -> mtokens. - plot_experiment / dn_sweep / streaming / dylu_control / warm_compare / plot_walltime: read mtokens; y = perplexity; log full + linear endgame/warm; x label "total tokens (M, all workers)"; final-perplexity bar in streaming. - README §3.2 + figure captions updated (perplexity, log/linear rationale). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Paper-form small-scale reproduction study of Async DiLoCo (arXiv:2401.09135 — Delayed-Nesterov, DyLU) and Streaming DiLoCo (arXiv:2501.18512 — fragmented overlapped sync), measured against synchronous DiLoCo on a single small Llama (34.4M) trained from scratch on Fineweb-Edu across 4 workers (4×4090). Lives in
examples/tiny_experiments/diloco_features/.The design was fixed as a pre-registration before any GPU time was spent (methodology, controls, and pass/fail gates predate the numbers). One run per arm at a single small scale — suggestive, not a benchmark.
Results (README §3.7–3.8)
1. Streaming trades per-token quality for wall-clock. All streaming arms are slightly worse per token (+0.06–0.14 eval) but finish faster (44.5–45.1 min vs 50.5); strided N=5 is closest per token and fastest — ≈ the baseline trajectory in real time.
2. DN-buffer depth is a non-monotonic lever. Sweeping N ∈ {4, 8, 16} from scratch: N=8 optimal (eval 3.107, ~half the gap to baseline closed), N=4 under-buffered (3.661), N=16 regresses (3.373). The paper's
N=k=4default is under-tuned for the from-scratch + staleness regime. (First clean post-#225 evidence on DN depth.)3. Warm-start closes the async gap (headline). Pretrained 4×DDP to ~500M tokens, then re-ran the async arms + a warm sync baseline from that checkpoint: every warm async arm lands +0.03–0.06 of the warm sync baseline (vs +0.25–0.80 from scratch), and the DN-depth dependence collapses to within 0.002. The from-scratch regime, not async, was the obstacle — exactly why the source papers warm-start. (async_nodn diverges as the no-DN control; DyLU neutral on this homogeneous rig.)
What's in here
experiment.sh(10-arm from-scratch matrix + DN sweep,WARM_MASTERwarm arms, token-budget global stop),templates/configs/default.yaml(server-token-budget DiLoCo worker),templates/configs/warm_pretrain.yaml(plain 4×DDP pretrain),make_warm_master.py(assemble a warm server master).harvest.py,plot_experiment.py,streaming.py,dn_sweep.py,plot_walltime.py,warm_compare.py,dylu_control.py,staleness.py(gate),grace_batches.py(validate-only),regen_tb.py.forgather diloco --token-budgetaccepts K/M/B suffixes (non-breaking units fix).docs/examples/), and pre-registered hypotheses with their outcomes.Caveats
Single seed, single small scale. The DyLU-neutral and N=16-regression results are the two most worth a second seed if firming up; the headline streaming / DN-optimum / warm-start effects are large.
🤖 Generated with Claude Code