Skip to content

Commit a9dcd8e

Browse files
jdinaltclaude
andauthored
DiLoCo: fix token-rate catch-up inflation + default Generate count to 1 (#235)
* diloco: stop the first-heartbeat catch-up from inflating the token rate A worker's first heartbeat reports its whole cumulative token count (the aggregator credits the full cumulative on a first report). With N workers those first reports land within ~1-2s, so total_tokens leaps from 0 to N-workers' worth of pre-heartbeat history in a tiny wall-clock window — and the rolling rate divided that jump by that interval, grossly overstating throughput (and underestimating the ETA) for the first few minutes until the window slid past the burst. Add StatsAggregator.steady_tokens: like total_tokens but excluding each worker's first-report catch-up (only true per-heartbeat increments), mirroring the existing flos_window treatment. The rolling-rate sampler reads steady_tokens instead of total_tokens, so the catch-up never enters the rate; total_tokens (with catch-up) still drives the budget percent/ETA-remaining, which is correct (those tokens are real). During warm-up — only catch-ups so far, steady_tokens still 0 — the rate reads 'unknown' (None) rather than a misleading 0 or an inflated burst. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * webui: default the DiLoCo 'Generate' worker count to 1 For the common small-worker-count case it's quicker to leave it at 1 and click Generate N times than to edit the number first (which is usually not the value you want anyway). 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 b6b44aa commit a9dcd8e

4 files changed

Lines changed: 62 additions & 10 deletions

File tree

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

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2881,32 +2881,42 @@ def _handle_heartbeat(self, handler: BaseHTTPRequestHandler):
28812881
self._maybe_log_stats()
28822882

28832883
def _record_token_sample(self):
2884-
"""Append a (monotonic ts, total_tokens) sample for the rolling rate.
2884+
"""Append a (monotonic ts, steady_tokens) sample for the rolling rate.
28852885
28862886
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.
2887+
runs). Samples ``steady_tokens`` (NOT ``total_tokens``) so the rate
2888+
excludes each worker's first-report cumulative catch-up — that burst
2889+
lands across all workers within ~1-2s and would otherwise grossly
2890+
overstate the rate on the first heartbeats. Lock-free on purpose:
2891+
``deque.append`` and the int read are each atomic under the GIL, and
2892+
different workers' heartbeats race here (this is NOT serial-per-worker
2893+
like ``pending_command``). A cross-worker read could append two samples a
2894+
hair out of order; the ``dt<=0`` / ``dtok<0`` guards in
2895+
``_rolling_token_rate`` absorb that, and a momentarily stale gauge is fine.
28932896
"""
28942897
if self.token_budget <= 0:
28952898
return
2896-
self._token_samples.append((time.monotonic(), self._stats.total_tokens))
2899+
self._token_samples.append((time.monotonic(), self._stats.steady_tokens))
28972900

28982901
def _rolling_token_rate(self) -> Optional[float]:
28992902
"""Tokens/second over the recent window, or None if not yet estimable.
29002903
29012904
Uses the oldest sample within ``_token_rate_window_s`` as the baseline
29022905
(or the oldest available if the buffer is younger than the window), so a
29032906
rate appears after the second heartbeat rather than waiting a full
2904-
window. Returns None until there are two spanning samples.
2907+
window. Returns None until there are two spanning samples, or while no
2908+
steady (non-catch-up) tokens have been counted yet — during that warm-up
2909+
the gauge reads "unknown" rather than a misleading 0 or an inflated burst.
29052910
"""
29062911
samples = list(self._token_samples)
29072912
if len(samples) < 2:
29082913
return None
29092914
now_t, now_tok = samples[-1]
2915+
if now_tok <= 0:
2916+
# No steady increment yet (only first-report catch-ups so far) —
2917+
# nothing real to rate. Distinct from a genuine stall, where
2918+
# steady_tokens is already > 0 and stays flat (reads as ~0 tok/s).
2919+
return None
29102920
cutoff = now_t - self._token_rate_window_s
29112921
base_t, base_tok = samples[0]
29122922
for t, tok in samples:

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,15 @@ def __init__(self, train_loss_decay: float = 0.7, eval_loss_decay: float = 0.5):
198198
self.total_flos: float = 0.0
199199
self.total_steps: int = 0
200200

201+
# Like total_tokens but EXCLUDING each worker's first-report catch-up
202+
# (the full cumulative a worker dumps on its first heartbeat). Used only
203+
# for the rolling throughput rate: that initial burst lands across all
204+
# workers within ~1-2s of wall-clock, so counting it would divide a huge
205+
# token jump by a tiny interval and grossly overstate the rate. Same
206+
# reasoning as the FLOPs `flos_window` below. Not persisted — it's a
207+
# live rate signal that resets cleanly on restart.
208+
self.steady_tokens: int = 0
209+
201210
# Per-worker last-seen cumulative counters, for delta accumulation
202211
# (persisted — must survive a restart so a resumed worker_id doesn't
203212
# re-add its history).
@@ -257,6 +266,10 @@ def update(self, worker_id: str, snap: Dict[str, Any]) -> None:
257266
live["flos_window"] = float(delta)
258267
elif attr == "total_tokens":
259268
self.total_tokens += int(delta)
269+
# Exclude the first-report catch-up from the rate counter
270+
# (only true per-heartbeat increments), mirroring flos_window.
271+
if prev is not None:
272+
self.steady_tokens += int(delta)
260273
else: # total_steps
261274
self.total_steps += int(delta)
262275

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,32 @@ def test_rate_window_uses_recent_baseline(self, tmp_path):
202202
# window [70,120]: baseline (100.0,100) -> (300-100)/(120-100)=10 tok/s
203203
assert abs(server._rolling_token_rate() - 10.0) < 1e-9
204204

205+
def test_steady_tokens_excludes_first_report_catchup(self, tmp_path):
206+
"""A worker's first heartbeat dumps its whole cumulative (catch-up).
207+
steady_tokens must count only true per-heartbeat increments, so the rate
208+
derived from it isn't inflated by that one-time burst."""
209+
agg = _server(tmp_path, token_budget=10_000)._stats
210+
# First reports from 4 workers (catch-up): total jumps, steady stays 0.
211+
for wid in ("w0", "w1", "w2", "w3"):
212+
agg.update(wid, {"tokens_total": 1000})
213+
assert agg.total_tokens == 4000
214+
assert agg.steady_tokens == 0
215+
# Second heartbeats (+100 each): only these count toward steady.
216+
for wid in ("w0", "w1", "w2", "w3"):
217+
agg.update(wid, {"tokens_total": 1100})
218+
assert agg.total_tokens == 4400
219+
assert agg.steady_tokens == 400
220+
221+
def test_rate_warmup_none_until_steady_tokens(self, tmp_path):
222+
"""During warm-up (only catch-up samples, steady_tokens still 0) the rate
223+
reads 'unknown' (None), not a misleading 0 or an inflated burst."""
224+
server = _server(tmp_path, token_budget=10_000)
225+
server._token_samples.extend([(1000.0, 0), (1010.0, 0)])
226+
assert server._rolling_token_rate() is None
227+
# First steady increment lands -> a real rate appears.
228+
server._token_samples.append((1020.0, 200))
229+
assert abs(server._rolling_token_rate() - 10.0) < 1e-9 # 200 / 20s
230+
205231
def test_stall_reports_zero_rate_not_none(self, tmp_path):
206232
"""After a long stall the only in-window sample is the tail; fall back to
207233
the previous sample so the gauge reads ~0 tok/s (visible slowdown), not a

‎tools/forgather_server/webui/src/components/SubmitModal.tsx‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1148,7 +1148,10 @@ function DiLoCoPicker(props: DiLoCoPickerProps) {
11481148
// Local UI state for the pool's add/generate controls. Hooks must run
11491149
// unconditionally, so they precede the cluster-fanout early return.
11501150
const [addName, setAddName] = useState<string>("");
1151-
const [genCount, setGenCount] = useState<string>("4");
1151+
// Default 1: for the common small-worker-count case it's quicker to leave it
1152+
// at 1 and click Generate N times than to edit the number first (which is
1153+
// usually not the value you want anyway).
1154+
const [genCount, setGenCount] = useState<string>("1");
11521155
const [generating, setGenerating] = useState<boolean>(false);
11531156
const [poolError, setPoolError] = useState<string | null>(null);
11541157

0 commit comments

Comments
 (0)