refactor(async): budget staleness in rollout batches; version in optimizer updates

[CjhHa1]

Summary

• track async AR/diffusion policy freshness in committed optimizer updates, with the recipe knob (max_staleness) denominated in whole rollout batches so it means the same depth at any num_updates_per_batch
• split the two quantities that were both called lag: staleness is the off-policyness of the data being trained on, publish_lag is weight-sync debt
• replace group-level freshest-first buffering with atomic completion-order FIFO rollout batches
• drive weight sync from the staleness budget or a hard eval/checkpoint boundary instead of sync-window counters
• keep Agentic on its existing control path; fix(rollout): guard agentic sync and deterministic diffusion eval #294 supplies its independent live-drive sync guard and logging

Related Issue

Supersedes the batch AR/diffusion sync-counter and resident-policy eval portions removed from #294.

Test Plan

Static and unit:

• Local async validation suite: 26 passed (kept outside the PR)
• Ruff lint and format checks on all changed Python files
• Hydra compose/resolve (--cfg job) for all six async override sets exercised below

Control-plane equivalence and invariants (drives the real pre-rename control
object from fbefbfa and the current one through one model of the
single-threaded loop, then compares event traces):

• Old max_policy_lag vs new max_staleness produce byte-identical traces for all three shipped recipes. BAGEL at U=2: 123 identical events, syncs after batches 12/24/36, staleness reaching exactly 22 updates = 11 batches.
• A fixed max_staleness ∈ {0,1,3,11} gives sync period {1,2,4,12} batches at every num_updates_per_batch ∈ {1,2,3,4,8}, and staleness in batch units equals max_staleness exactly. The old unit could not express this: max_policy_lag=4 is depth 2 at U=4 but depth 3 at U=2 and depth 5 at U=1.
• Randomized fuzz, 4000 trials (1982 with partially-committed steps), 0 violations of: queue empty at every sync and hard boundary, consumed staleness never above budget, admission never starves the queue while refusing to launch.
• The eval/checkpoint clamp this run surfaced is now reported rather than silent: sync_period_batches was cross-checked against the simulated publication schedule over 8 (max_staleness, U, eval_interval, save_interval) combinations and matched every one.

GPU, 1x8 H20 per path (4 train + 4 rollout). This is the first GPU coverage of
the control path as it now stands — the earlier smokes ran the pre-dcf2e93
version whose harness cannot even import against this commit. Observable is the
weight publication itself (SGLang Cache flushed successfully!, vLLM-Omni
Activating adapter), since the async metrics only reach W&B:

arm	path	U	max_staleness	max_inflight	rollouts	predicted syncs	observed syncs	rc
A1	AR / SGLang	4	1	2	6	after 2, 4	after 2, 4	0
A2	AR / SGLang	2	1	2	6	after 2, 4	after 2, 4	0
A3	AR / SGLang	4	0	1	4	after 1, 2, 3	after 1, 2, 3	0
A4	AR / SGLang	4	3	2	8	after 4	after 4	0
B1	BAGEL / vLLM-Omni	2	1	1	6	after 2, 4	after 2, 4	0
B2	BAGEL / vLLM-Omni	2	3	1	8	after 4	after 4	0

• Every arm matched its predicted publication schedule and exited 0. A1 vs A2 is the differentiating case: same max_staleness=1 at U=4 and U=2, identical period. Spelling A1 with the old knob needed max_policy_lag=4; carrying that 4 into A2's U=2 gives 4//2+1 = 3, so the period would have moved.
• Preflight on both pods asserted AsyncBatchControl(max_staleness=11, num_updates_per_batch=2) yields staleness_budget=22 / admission_depth=12, reproducing what the BAGEL recipe previously spelled as max_policy_lag=22.
• Admission notes fired only where designed: A3 emitted the max_staleness=0 no-overlap warning; A4 and B2 (depth above max_inflight + 1) emitted the surplus note, each self-reporting a 4-batch weight-sync period that the observed schedule then confirmed; A1/A2/B1 emitted nothing.

What the budget buys in wall time (measured from the arm logs; the AR arms are a
clean comparison because their policy is static, so per-step compute is identical
and only the publication cadence moves):

• A publication costs ~8.4s on AR and ~5-6s on BAGEL (quiesce plus weight push). max_staleness=0 pays it every batch at 11.99s/step; a budget amortises it to 4.2s/batch at max_staleness=1 and 2.1s/batch at 3 (BAGEL 3.0 → 1.3s/batch). Mean AR step time is 78.9% higher at the synchronous floor than at any non-zero budget. A4/A5 reproduce to 0.03s on mean step. At batch_size 4 the fixed cost dominates, so 78.9% is an upper bound — a production step makes the same overhead a smaller fraction.

Limits of the GPU evidence, stated so it is not over-read. Two follow-up arms were
run specifically to remove these and both came back null:

• The AR arms validate the control plane only, not staleness effects. grad_norm prints 0.00000 in every AR rollout: at 4 prompts x 4 samples on DAPO math a 4B base model gets whole groups wrong, GRPO's within-group advantage vanishes, and the policy does not move — so all arms are bit-identical and the ratio spread (0.9988-1.0329) is the SGLang-vs-FSDP numerical gap, not drift. Re-running at 8 prompts x 8 samples with 2048 new tokens (max_staleness=0 vs 3) produced within-group reward variance but still byte-identical arms: mean |ratio-1| 0.00493 both, clip 0.0000 both. This is the base model's pass rate, not a batch-size problem. Precisely: grad_norm = 0.0 is TrainStack's no-backward fallback and that path also sets optimizer_updates = 0, which would freeze publish_lag and prevent any sync — the observed schedules require U updates per batch, so backward did run and the norm is genuinely below the %.4f print precision.
• Comparing two budgets by their means is not sensitive enough. Three seeds per budget: max_staleness=1 mean clip 0.3478 ±0.0439, max_staleness=3 mean clip 0.4011 ±0.0735, reward 0.6252 vs 0.6253 — the between-budget gap is 0.91x the within-budget spread, so not separable from run-to-run variation. The single-run gap that first looked like an AReaL-style effect was seed noise.

Measured against the shipped budget, 24 rollouts per arm at 8 prompts x 8 samples x
8 SDE steps, same seed, so max_staleness=11 runs two complete publication cycles:

• Drift grows with the budget, significantly. The period is 12 batches, so the batch consumed at rollout i is (i-1) mod 12 batches stale by construction — one run contains its own staleness sweep, paired against a single policy trajectory and free of the seed noise above. Clip fraction rises 0.140 → 0.420 across the cycle; grouped, staleness ≤ 2 (n=6) clips 0.1500 against staleness ≥ 7 (n=10) at 0.3170, a factor of 2.11, difference +0.167 ± 0.050, t = 3.33. It cross-checks: the entire max_staleness=1 run sits at clip 0.1550, i.e. at this run's fresh level, so the deep budget's higher overall clip (0.2675) is the same response curve sampled mostly at its stale end. Clip is the instrument rather than the ratio because clip_range is 1e-5 while the console prints ratio to four decimals, so |ratio-1| reads 0.000000 at every level.
• Whether that drift costs reward is unresolved. Neither arm has a significant reward trend over 24 rollouts (slope +0.000405 ± 0.000905, t=0.45 at max_staleness=1; −0.000795 ± 0.000833, t=−0.95 at 11), and the deep budget's −0.0136 mean gap is 0.47 of one rollout's standard deviation. Direction disfavours the deep budget in both the mean and the slope sign, but establishing it needs a production-length run — enough steps to read a reward trajectory, and ideally the decoupled objective AReaL and verl rely on at this depth, which UniRL does not implement. That is a training question rather than a control-plane one and should not gate this PR.

Compatibility / Risk

• Breaking async AR/diffusion config change: weight_sync_interval and buffer_max_staleness are no longer read on this path; use max_staleness, counted in whole rollout batches. Both keys stay live elsewhere — weight_sync_interval in 69 synchronous recipes, buffer_max_staleness on the agentic path — so a stale key here is silently ignored rather than rejected; check async recipes by hand when upgrading.
• max_staleness is enforced at batch entry. With num_updates_per_batch = U the anchor stays frozen while the weights move, so update i trains at staleness + i - 1 and the worst case any gradient step sees is staleness_budget + U - 1. That extra span is the in-batch off-policyness PPO already assumes and clip_range covers.
• An eval or checkpoint interval below max_staleness + 1 becomes the effective sync period, because next_hard_boundary clamps admission too: at max_staleness=11 with eval_interval=10 the run publishes every 10 batches and data never gets past 9 of the configured 11. This used to be silent — the startup note reported the budget-derived 12 — and is now reported by sync_period_batches, quoted in the surplus note, warned about when a budget can never be spent, and logged to W&B as sync_period_batches. No shipped recipe changes behaviour: the AR recipes are max_staleness=0 (period 1 either way) and BAGEL sets no eval_interval, so its 12 stands.
• BAGEL ships max_staleness=11, i.e. 11 batches between publications. That is past AReaL's recommended eta <= 8, and FlowGRPO's rollout-anchored ratio is their "naive PPO" arm with no decoupled objective to absorb the drift. The recipe comment says so; lower it first if the ratio starts to swing.
• Async eval/checkpoint boundaries now drain admitted work and sync the current train version before evaluating/saving.
• FIFO consumption changes rollout-data order relative to freshest-first selection.
• Agentic async/partial behavior is intentionally unchanged by this PR.

Checklist

• I reviewed the changed code and removed unrelated/generated artifacts.
• Repository-level tests are not included; the local validation suite is retained outside the PR.
• Documentation and example recipes are updated for the new control semantics.