Skip to content

Commit b6b44aa

Browse files
jdinaltclaude
andauthored
DiLoCo: token-budget progress bar + rolling-rate ETA (webui + CLI) (#234)
* diloco: token-budget progress + rolling-rate ETA (webui bar + CLI line) 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> * diloco: address review — keep the rate gauge live through a stall - _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> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 035dda3 commit b6b44aa

7 files changed

Lines changed: 367 additions & 8 deletions

File tree

‎docs/trainers/diloco-architecture.md‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -627,6 +627,14 @@ can race the workers: if the process exits before they finish stopping and
627627
deregistering, they are orphaned — and a reused port (the next arm's server) then
628628
404s their heartbeats, the failure mode the worker orphan self-exit also guards.
629629

630+
Each heartbeat that advances the aggregate appends a `(monotonic, total_tokens)`
631+
sample to an in-memory ring (`_record_token_sample`, only when a budget is set).
632+
`_token_progress()` reports completed/budget, the clamped fraction, a rolling
633+
tokens/sec over a recent window (`_rolling_token_rate`, ~5 min — wide because
634+
heartbeats are ~30 s apart) and an ETA, surfaced in `/status` as `token_progress`
635+
for the webui progress bar and the `forgather diloco status` line. It's a live
636+
gauge (in-memory, resets on restart), not persisted state.
637+
630638
**Critical locking:**
631639

632640
| Lock | Protects | Used by |

‎docs/trainers/diloco-cli.md‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,21 @@ status, and DyLU configuration. The known-worker line expands into a
206206
**resumable roster** — the not-running worker_ids (with last-seen times) you can
207207
relaunch under to resume from their checkpoints.
208208

209+
When a `--token-budget` is set, the status renders a **progress line** —
210+
completed/total tokens, percent, the rolling token rate, and an ETA:
211+
212+
```
213+
Token budget: [##########--------------] 41.0%
214+
820.00M / 2.00B · 1.20M/s · ETA 16m23s
215+
```
216+
217+
The rate is averaged over a recent window (~5 min — wide on purpose, since
218+
heartbeats are ~30 s apart), so a sustained slowdown shows up as a falling rate
219+
and a climbing ETA, a quick at-a-glance health check for an open-ended run. The
220+
webui's DiLoCo server view shows the same as a fill bar under the header. (Both
221+
fall back to the plain budget count against an older server that predates the
222+
progress field.)
223+
209224
Add `--queues` to also show **work-unit dispatch** per `(dataset_id,
210225
shuffle_seed)` queue: a human-readable dataset label (from the worker's
211226
load-args hint, e.g. `wikitext@train`) with the raw `dataset_id` hash kept as a

‎src/forgather/cli/diloco_orch.py‎

Lines changed: 51 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -321,6 +321,38 @@ def _render_aggregate_stats(agg):
321321
print(line)
322322

323323

324+
def _fmt_tokens(n):
325+
"""Compact K/M/B token count, e.g. 1_500_000 -> '1.50M'."""
326+
if n is None:
327+
return "—"
328+
n = float(n)
329+
for unit, div in (("B", 1e9), ("M", 1e6), ("K", 1e3)):
330+
if abs(n) >= div:
331+
return f"{n / div:.2f}{unit}"
332+
return f"{n:.0f}"
333+
334+
335+
def _fmt_duration(secs):
336+
"""Human ETA, e.g. 3725 -> '1h02m'; None/negative -> '—'."""
337+
if secs is None or secs < 0:
338+
return "—"
339+
secs = int(secs)
340+
h, rem = divmod(secs, 3600)
341+
m, s = divmod(rem, 60)
342+
if h:
343+
return f"{h}h{m:02d}m"
344+
if m:
345+
return f"{m}m{s:02d}s"
346+
return f"{s}s"
347+
348+
349+
def _ascii_bar(frac, width=24):
350+
"""A TQDM-ish ASCII bar: '[######------------------]'."""
351+
frac = max(0.0, min(1.0, float(frac)))
352+
filled = int(round(frac * width))
353+
return "[" + "#" * filled + "-" * (width - filled) + "]"
354+
355+
324356
def render_status(merged, *, want_queues):
325357
"""Human-readable rendering of an :func:`assemble_status` result."""
326358
import datetime
@@ -383,7 +415,25 @@ def render_status(merged, *, want_queues):
383415

384416
if status.get("token_budget", 0):
385417
stopped = " (stop sent)" if status.get("budget_stop_sent") else ""
386-
print(f" Token budget: {status['token_budget']:,}{stopped}")
418+
prog = status.get("token_progress") or {}
419+
done = prog.get("tokens_completed")
420+
budget = status["token_budget"]
421+
if done is not None:
422+
frac = prog.get("fraction") or 0.0
423+
rate = prog.get("tokens_per_second")
424+
eta = prog.get("eta_seconds")
425+
# A real stall reports rate 0.0 (not None) — show "0/s" (reads as
426+
# "stalled") rather than "—/s" (reads as "unknown / not yet known").
427+
rate_s = f"{_fmt_tokens(rate)}/s" if rate is not None else "—/s"
428+
eta_s = "done" if status.get("budget_stop_sent") else _fmt_duration(eta)
429+
print(f" Token budget: {_ascii_bar(frac)} {frac * 100:5.1f}%{stopped}")
430+
print(
431+
f" {_fmt_tokens(done)} / {_fmt_tokens(budget)}"
432+
f" · {rate_s} · ETA {eta_s}"
433+
)
434+
else:
435+
# Older server without token_progress.
436+
print(f" Token budget: {budget:,}{stopped}")
387437

388438
deaths = status.get("total_worker_deaths", 0)
389439
if deaths:

‎src/forgather/ml/diloco/server.py‎

Lines changed: 86 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@
3737
import struct
3838
import threading
3939
import time
40-
from collections import defaultdict
40+
from collections import defaultdict, deque
4141
from dataclasses import dataclass, field
4242
from http.server import BaseHTTPRequestHandler, HTTPServer, ThreadingHTTPServer
4343
from typing import Any, Callable, Dict, List, Optional, Tuple
@@ -619,6 +619,16 @@ def __init__(
619619
# persisted, so a restart still over budget re-sends the stop.
620620
self.token_budget = max(0, int(token_budget))
621621
self._budget_stop_sent = False
622+
# Rolling-rate samples for the token-budget progress / ETA — (monotonic
623+
# ts, total_tokens) appended whenever a heartbeat advances the aggregate.
624+
# In-memory only (a live gauge; resets on restart). The window is wide on
625+
# purpose: heartbeats are ~30 s apart, so a short window would see only a
626+
# sample or two and give a jittery rate. 300 s spans 6+ heartbeats per
627+
# worker for a stable estimate while still reflecting a sustained
628+
# slowdown within a few minutes. Not a CLI knob — tune here if ever
629+
# needed. maxlen bounds memory if heartbeats are unusually dense.
630+
self._token_samples: deque = deque(maxlen=512)
631+
self._token_rate_window_s: float = 300.0
622632
# Verbose per-round sync logging (off by default). Server-authoritative:
623633
# advertised in /info so each worker gates its own per-round sync log to
624634
# match, and the server gates its per-round outer-step log. A targeted
@@ -2837,6 +2847,7 @@ def _handle_heartbeat(self, handler: BaseHTTPRequestHandler):
28372847
# the EMA math).
28382848
if worker_stats:
28392849
self._stats.update(worker_id, worker_stats)
2850+
self._record_token_sample()
28402851
self._maybe_trigger_token_budget()
28412852

28422853
# Compute DyLU recommendation if enabled
@@ -2869,6 +2880,76 @@ def _handle_heartbeat(self, handler: BaseHTTPRequestHandler):
28692880
if worker_stats:
28702881
self._maybe_log_stats()
28712882

2883+
def _record_token_sample(self):
2884+
"""Append a (monotonic ts, total_tokens) sample for the rolling rate.
2885+
2886+
No-op without a token budget (the rate/ETA is only surfaced for budgeted
2887+
runs). Lock-free on purpose: ``deque.append`` and the ``total_tokens``
2888+
int read are each atomic under the GIL, and different workers' heartbeats
2889+
race here (this is NOT serial-per-worker like ``pending_command``). A
2890+
cross-worker read could append two samples a hair out of order; the
2891+
``dt<=0`` / ``dtok<0`` guards in ``_rolling_token_rate`` absorb that, and
2892+
a momentarily stale gauge is fine.
2893+
"""
2894+
if self.token_budget <= 0:
2895+
return
2896+
self._token_samples.append((time.monotonic(), self._stats.total_tokens))
2897+
2898+
def _rolling_token_rate(self) -> Optional[float]:
2899+
"""Tokens/second over the recent window, or None if not yet estimable.
2900+
2901+
Uses the oldest sample within ``_token_rate_window_s`` as the baseline
2902+
(or the oldest available if the buffer is younger than the window), so a
2903+
rate appears after the second heartbeat rather than waiting a full
2904+
window. Returns None until there are two spanning samples.
2905+
"""
2906+
samples = list(self._token_samples)
2907+
if len(samples) < 2:
2908+
return None
2909+
now_t, now_tok = samples[-1]
2910+
cutoff = now_t - self._token_rate_window_s
2911+
base_t, base_tok = samples[0]
2912+
for t, tok in samples:
2913+
if t >= cutoff:
2914+
base_t, base_tok = t, tok
2915+
break
2916+
# After a long stall the only in-window sample can be the tail itself
2917+
# (base collapses onto ``now``). Fall back to the previous sample so we
2918+
# still produce a rate — a stall then reads as ~0 tok/s (the slowdown we
2919+
# most want visible), not a blank gauge. len>=2 guaranteed above.
2920+
if base_t >= now_t:
2921+
base_t, base_tok = samples[-2]
2922+
dt = now_t - base_t
2923+
dtok = now_tok - base_tok
2924+
if dt <= 0 or dtok < 0:
2925+
return None
2926+
return dtok / dt
2927+
2928+
def _token_progress(self) -> Optional[Dict[str, Any]]:
2929+
"""Token-budget progress + rolling-rate ETA, or None when no budget.
2930+
2931+
Shape (all keys present when a budget is set):
2932+
tokens_completed, token_budget, fraction (0..1, clamped),
2933+
tokens_per_second (None until estimable), eta_seconds (None when the
2934+
rate is unknown or zero), rate_window_seconds, budget_stop_sent.
2935+
"""
2936+
if self.token_budget <= 0:
2937+
return None
2938+
done = int(self._stats.total_tokens)
2939+
budget = int(self.token_budget)
2940+
rate = self._rolling_token_rate()
2941+
remaining = max(0, budget - done)
2942+
eta = remaining / rate if (rate and rate > 0) else None
2943+
return {
2944+
"tokens_completed": done,
2945+
"token_budget": budget,
2946+
"fraction": min(1.0, done / budget) if budget > 0 else 0.0,
2947+
"tokens_per_second": rate,
2948+
"eta_seconds": eta,
2949+
"rate_window_seconds": self._token_rate_window_s,
2950+
"budget_stop_sent": self._budget_stop_sent,
2951+
}
2952+
28722953
def _maybe_trigger_token_budget(self):
28732954
"""When the aggregated global token count reaches ``token_budget``,
28742955
relay ``save_and_stop`` to every worker (once). The command rides each
@@ -3176,6 +3257,10 @@ def _handle_status(self, handler: BaseHTTPRequestHandler):
31763257
if self.token_budget > 0:
31773258
response["token_budget"] = self.token_budget
31783259
response["budget_stop_sent"] = self._budget_stop_sent
3260+
# Progress + rolling-rate ETA (tokens_completed, fraction,
3261+
# tokens_per_second, eta_seconds, rate_window_seconds) for the
3262+
# webui progress bar and the CLI status line.
3263+
response["token_progress"] = self._token_progress()
31793264

31803265
if self.async_mode:
31813266
response["total_submissions"] = self._total_submissions

‎tests/unit/ml/diloco/test_token_budget.py‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,3 +150,74 @@ def test_raise_clears_prior_stop(self, tmp_path):
150150
assert _pending(server, "w0") == [None]
151151
finally:
152152
server.stop()
153+
154+
155+
class TestTokenProgress:
156+
"""The progress/ETA gauge the webui bar and CLI status line consume."""
157+
158+
def test_none_without_budget(self, tmp_path):
159+
server = _server(tmp_path, token_budget=0)
160+
assert server._token_progress() is None
161+
# No samples are recorded when there's no budget.
162+
server._stats.total_tokens = 5
163+
server._record_token_sample()
164+
assert len(server._token_samples) == 0
165+
166+
def test_progress_fields_and_fraction(self, tmp_path):
167+
server = _server(tmp_path, token_budget=1000)
168+
server._stats.total_tokens = 250
169+
prog = server._token_progress()
170+
assert prog["tokens_completed"] == 250
171+
assert prog["token_budget"] == 1000
172+
assert prog["fraction"] == 0.25
173+
assert prog["rate_window_seconds"] == server._token_rate_window_s
174+
# No rate from a single sample => no ETA.
175+
assert prog["tokens_per_second"] is None
176+
assert prog["eta_seconds"] is None
177+
178+
def test_fraction_clamps_at_one(self, tmp_path):
179+
server = _server(tmp_path, token_budget=1000)
180+
server._stats.total_tokens = 1500 # overshoot past the budget
181+
assert server._token_progress()["fraction"] == 1.0
182+
183+
def test_rolling_rate_and_eta(self, tmp_path):
184+
server = _server(tmp_path, token_budget=2000)
185+
# Seed (monotonic ts, total_tokens) samples directly: 100 tokens/10s.
186+
server._token_samples.extend((1000.0 + i * 10.0, i * 100) for i in range(6))
187+
server._stats.total_tokens = 500 # = last sample
188+
rate = server._rolling_token_rate()
189+
assert abs(rate - 10.0) < 1e-9 # 100 tokens / 10 s
190+
prog = server._token_progress()
191+
assert abs(prog["tokens_per_second"] - 10.0) < 1e-9
192+
# remaining 1500 tokens / 10 tok/s = 150 s
193+
assert abs(prog["eta_seconds"] - 150.0) < 1e-6
194+
195+
def test_rate_window_uses_recent_baseline(self, tmp_path):
196+
server = _server(tmp_path, token_budget=10_000)
197+
server._token_rate_window_s = 50.0
198+
# Old slow segment then a recent fast one; only the last 50s counts.
199+
samples = [(0.0, 0), (10.0, 10), (100.0, 100), (110.0, 200), (120.0, 300)]
200+
server._token_samples.extend(samples)
201+
server._stats.total_tokens = 300
202+
# window [70,120]: baseline (100.0,100) -> (300-100)/(120-100)=10 tok/s
203+
assert abs(server._rolling_token_rate() - 10.0) < 1e-9
204+
205+
def test_stall_reports_zero_rate_not_none(self, tmp_path):
206+
"""After a long stall the only in-window sample is the tail; fall back to
207+
the previous sample so the gauge reads ~0 tok/s (visible slowdown), not a
208+
blank None."""
209+
server = _server(tmp_path, token_budget=10_000)
210+
server._token_rate_window_s = 50.0
211+
# All progress is old; the last two samples are flat (a stall) and only
212+
# the tail falls inside the 50s window.
213+
server._token_samples.extend(
214+
[(0.0, 1000), (10.0, 1000), (200.0, 1000), (260.0, 1000)]
215+
)
216+
server._stats.total_tokens = 1000
217+
rate = server._rolling_token_rate()
218+
assert rate is not None
219+
assert rate == 0.0 # baseline = samples[-2]=(200,1000), now=(260,1000)
220+
# A 0 rate yields no ETA (can't divide), but progress still renders.
221+
prog = server._token_progress()
222+
assert prog["tokens_per_second"] == 0.0
223+
assert prog["eta_seconds"] is None

‎tools/forgather_server/webui/src/api.ts‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1196,6 +1196,20 @@ export interface DiLoCoWorkerStatus {
11961196
}
11971197

11981198
/** Upstream ``/status`` response. Field set tracks DiLoCoServer._handle_status. */
1199+
/** Token-budget progress for the dashboard bar / CLI status line. */
1200+
export interface DiLoCoTokenProgress {
1201+
tokens_completed: number;
1202+
token_budget: number;
1203+
/** 0..1, clamped. */
1204+
fraction: number;
1205+
/** Rolling tokens/sec over the recent window; null until estimable. */
1206+
tokens_per_second: number | null;
1207+
/** Seconds remaining at the rolling rate; null when rate unknown/zero. */
1208+
eta_seconds: number | null;
1209+
rate_window_seconds: number;
1210+
budget_stop_sent: boolean;
1211+
}
1212+
11991213
export interface DiLoCoStatus {
12001214
status?: string;
12011215
mode?: "sync" | "async";
@@ -1217,6 +1231,8 @@ export interface DiLoCoStatus {
12171231
min_workers?: number;
12181232
token_budget?: number;
12191233
budget_stop_sent?: boolean;
1234+
/** Token-budget progress + rolling-rate ETA (present when token_budget>0). */
1235+
token_progress?: DiLoCoTokenProgress;
12201236
outer_lr?: number;
12211237
outer_momentum?: number;
12221238
save_dir?: string;

0 commit comments

Comments
 (0)