Repository navigation
DiLoCo: token-budget progress bar + rolling-rate ETA (webui + CLI) - #234
Merged
Merged
Conversation
Surface how far an open-ended (token-budgeted) run has got and roughly how long is left, so a run that's silently going much slower than it should is obvious at a glance instead of only after the fact. Server: each heartbeat that advances the aggregate appends a (monotonic, total_tokens) sample to an in-memory ring (only when a budget is set). _token_progress() reports completed/budget, a clamped fraction, a rolling tokens/sec over a recent window (~5 min — wide because heartbeats are ~30s apart, so a short window would see one or two samples and jitter), and an ETA. Exposed in /status as token_progress (a live gauge, not persisted). CLI (forgather diloco status): a TQDM-ish ASCII line under the existing budget field — bar + percent, then completed/total, rate, and ETA. webui (DiLoCo server view): a fill bar under the dashboard header when a budget is set, with completed/total + rate + ETA to the right. Always offered now; falls back to the plain budget count against a server without token_progress. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- _rolling_token_rate: after a long stall the only in-window sample could be the tail itself (baseline collapses onto now -> dt=0 -> None), blanking the rate/ETA exactly during the slowdown it exists to surface. Fall back to the previous sample so a stall reads as ~0 tok/s. (+ regression test) - Show 0/s for a genuine stall (rate is not None) instead of '—/s' (which reads as 'unknown') in both the CLI and the webui bar. - Reword the _record_token_sample thread-safety note: heartbeats race across workers here (not serial-per-worker); the dt<=0 / dtok<0 guards absorb any cross-worker append jitter. 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.
Summary
Better at-a-glance diagnostics for token-budgeted DiLoCo runs. The motivating
incident: a run was taking far longer than it should and nothing surfaced it —
a live progress/ETA would have caught it early.
server.py): each heartbeat that advances the aggregate appends a(monotonic, total_tokens)sample to an in-memory ring (only when a budget isset).
_token_progress()reportstokens_completed,token_budget, clampedfraction, a rollingtokens_per_secondover a recent window, andeta_seconds. Exposed in/statusastoken_progress(a live gauge — notpersisted, resets on restart). The window is ~5 min on purpose: heartbeats
are ~30 s apart, so a short window would see one or two samples and jitter; a
wide window stays stable while still reflecting a sustained slowdown within a
few minutes (rolling-recent-rate basis, so a slowdown shows as a falling rate /
rising ETA).
forgather diloco status): a TQDM-ish ASCII line under the budgetfield —
completed/total + rate + ETA to the right. Always offered when a budget is set
(an older server without
token_progressfalls back to the plain count).Tests
TestTokenProgressintest_token_budget.py: no-budget no-op, fields/fraction,clamp-at-1, rolling rate + ETA math, recent-window baseline selection. Full
tests/unit/ml/diloco/+tests/unit/cli/green (864).black/isortclean,webui
tsc --noEmitclean.Docs
docs/trainers/diloco-cli.md(status progress line),diloco-architecture.md(the sampling/rate gauge).
🤖 Generated with Claude Code