From 7196d83735129ea998b68aaf8af2d1c11c36f340 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Mon, 3 Aug 2026 22:57:12 +0800 Subject: [PATCH 1/9] refactor(async): measure policy lag in optimizer updates Replace sync-window/LIFO control with FIFO rollout batches and demand-driven weight sync so freshness has one explicit unit. --- .../qwen3_drpo_4b_base_dapo_sglang_async.yaml | 14 +- .../qwen3_grpo_4b_base_dapo_sglang_async.yaml | 17 +- .../diffusion/bagel/bagel_vllmomni_async.yaml | 25 +- pyproject.toml | 3 + tests/rollout/test_async_batch_engine.py | 74 ++++++ tests/trainer/test_async_policy.py | 138 +++++++++++ .../trainer/test_async_trainer_boundaries.py | 57 +++++ unirl/rollout/README.md | 18 +- unirl/rollout/engine/asynchronous.py | 169 ++++++++------ unirl/rollout/engine/synchronous.py | 14 ++ unirl/train/backend/base_backend.py | 6 + unirl/train/stack/base.py | 16 +- unirl/train/unified_model_stack.py | 11 +- unirl/train_async_ar.py | 8 +- unirl/train_async_diffusion.py | 6 +- unirl/trainer/README.md | 20 +- unirl/trainer/agentic.py | 2 +- unirl/trainer/agentic_async.py | 2 +- unirl/trainer/async_ar.py | 214 +++++++++++++----- unirl/trainer/async_diffusion.py | 201 +++++++++++----- unirl/trainer/async_policy.py | 181 +++++++++++++++ 21 files changed, 940 insertions(+), 256 deletions(-) create mode 100644 tests/rollout/test_async_batch_engine.py create mode 100644 tests/trainer/test_async_policy.py create mode 100644 tests/trainer/test_async_trainer_boundaries.py create mode 100644 unirl/trainer/async_policy.py diff --git a/examples/ar/qwen3_drpo_4b_base_dapo_sglang_async.yaml b/examples/ar/qwen3_drpo_4b_base_dapo_sglang_async.yaml index 0679bcf50..4fac94496 100644 --- a/examples/ar/qwen3_drpo_4b_base_dapo_sglang_async.yaml +++ b/examples/ar/qwen3_drpo_4b_base_dapo_sglang_async.yaml @@ -23,17 +23,11 @@ num_rollouts: 800 # Train slab share (rollout gets the rest). 0.5 -> equal split. train_fraction: 0.5 -# Async knobs (one unified loop): max_inflight = concurrent generations (overlap -# depth; 1 ≈ one-step pipeline). buffer_max_staleness = weight-syncs a buffered -# group may cross: 0 = on-policy (ratio≈1), >0 = off-policy continuous buffer. +# Async control. Policy lag is measured in committed optimizer updates. +# max_policy_lag=0 keeps the behavior/train versions aligned at batch +# admission; max_inflight remains a resource cap. max_inflight: 1 -buffer_max_staleness: 0 - -# =1 is the on-policy reward-parity baseline vs the colocate ARTrainer (a sync -# sits between every gen/train pair -> NO overlap). Set >1 to overlap -# generate(N+1) with train(N): data is then off-policy by <= interval weight -# versions, absorbed by DRPO's rollout-anchored ratio (old_logp_source: rollout). -weight_sync_interval: 1 +max_policy_lag: 0 eval_interval: 10 adv_normalization_scope: group diff --git a/examples/ar/qwen3_grpo_4b_base_dapo_sglang_async.yaml b/examples/ar/qwen3_grpo_4b_base_dapo_sglang_async.yaml index 6ee77cbde..2e79b0808 100644 --- a/examples/ar/qwen3_grpo_4b_base_dapo_sglang_async.yaml +++ b/examples/ar/qwen3_grpo_4b_base_dapo_sglang_async.yaml @@ -26,19 +26,12 @@ num_rollouts: 800 # train-bound, lower for generation-bound workloads (mind DIVISIBILITY above). train_fraction: 0.5 -# Async knobs (one unified loop): max_inflight = concurrent generations (overlap -# depth; 1 ≈ one-step pipeline). buffer_max_staleness = weight-syncs a buffered -# group may cross: 0 = on-policy (ratio≈1), >0 = off-policy continuous buffer. +# Async control. Policy lag is measured in committed optimizer updates. +# max_policy_lag=0 means every batch starts training against the same +# train version that generated it (rollout/replay numerical differences may +# still make ratio only approximately 1). max_inflight remains a resource cap. max_inflight: 1 -buffer_max_staleness: 0 - -# Push the full dense weights into SGLang every N rollouts. =1 is the on-policy -# reward-parity baseline vs the colocate ARTrainer (a sync sits between every -# gen/train pair, so there is NO overlap). Set >1 to actually overlap -# generate(N+1) with train(N): the data is then off-policy by <= interval -# weight versions, absorbed by the rollout-anchored PPO ratio (old_logp = -# segment.log_probs), exactly the slime train_async trade-off. -weight_sync_interval: 1 +max_policy_lag: 0 # AIME eval every 10 rollouts + a pre-train baseline (rollout -1). eval_interval: 10 diff --git a/examples/diffusion/bagel/bagel_vllmomni_async.yaml b/examples/diffusion/bagel/bagel_vllmomni_async.yaml index e3b766d2f..df4ef7fcd 100644 --- a/examples/diffusion/bagel/bagel_vllmomni_async.yaml +++ b/examples/diffusion/bagel/bagel_vllmomni_async.yaml @@ -43,29 +43,18 @@ workers_per_device: 1 layout: separate train_fraction: 0.5 # 4 train GPUs + 4 rollout GPUs on an 8-GPU pool -# LoRA weight-sync cadence (loop concern, read by the trainer). >1 is required -# for async overlap: interval=1 drains every step (every rollout is a cold sync -# boundary, no generation can overlap a train step). interval=4 gives 3 of every -# 4 rollouts a generation overlapped with training. old_logp_source=rollout keeps -# the emitted π_old as a valid importance-sampling anchor; the ratio may move as -# the resident rollout policy and current train policy diverge. -weight_sync_interval: 4 - -# ---- async knobs (AsyncDiffusionTrainer) ---- +# ---- async control (optimizer-update units) ---- # max_inflight: concurrent generations. MUST be 1: the trajectory-segment # cross-slab transfer (NCCL send) runs on the rollout worker; a second in-flight # generation co-tenanting that worker blocks the send behind it (~150s/rollout). # With max_inflight=1 the trainer polls (reaps) before topping up launches, so # each generation transfers in the idle window before the next launch, then -# that next generation overlaps the train step. Real overlap needs -# weight_sync_interval>1 (interval=1 drains every step). -# buffer_max_staleness: how many weight syncs a buffered group may cross. -# 0 = never crosses a regular rollout-weight sync (~174s/rollout on BAGEL 4+4). -# 2 = throughput-optimal continuous buffer (~148s/rollout, matches vllm -# colocate); sync-boundary cold rollouts disappear. The emitted π_old -# remains a valid anchor, while the measured ratio may reflect policy lag. +# that next generation overlaps the train step. Real overlap needs admission +# depth > 1. +# With num_updates_per_batch=2, max_policy_lag=22 admits a 12-batch +# behavior-policy cycle: floor(22 / 2) + 1 = 12. max_inflight: 1 -buffer_max_staleness: 2 +max_policy_lag: 22 logging: report_to_wandb: false # flip to true to enable wandb (rank-0/driver only) @@ -220,7 +209,7 @@ sampling: # LoRA weight sync → vLLM-Omni rollout. Separate slabs ⇒ RemoteLoraWeightSync: # rank 0 ships the freshly-trained LoRA adapter to each cross-slab rollout Worker # by Ray RPC (no NCCL rendezvous, no name_remap — pushes the adapter directly). -# The train loop drives cadence via weight_sync_interval. +# The train loop publishes only when the optimizer-update lag budget binds. sync: _target_: unirl.distributed.weight_sync.lora.RemoteLoraWeightSync verify: true # checksum read-back asserts the synced LoRA landed; catches a wrong prefix diff --git a/pyproject.toml b/pyproject.toml index 97254ada7..fbed401a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -196,3 +196,6 @@ known-first-party = ["unirl"] [tool.ruff.format] quote-style = "double" + +[tool.pytest.ini_options] +testpaths = ["tests"] diff --git a/tests/rollout/test_async_batch_engine.py b/tests/rollout/test_async_batch_engine.py new file mode 100644 index 000000000..eb4520669 --- /dev/null +++ b/tests/rollout/test_async_batch_engine.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import pytest + +from unirl.rollout.engine.asynchronous import AsyncBatchRolloutEngine + + +@dataclass +class _Pending: + value: str + + def ready(self) -> bool: + return True + + def result(self) -> str: + return self.value + + def wait(self) -> None: + return None + + +class _RolloutHandle: + def __init__(self) -> None: + self.launched: list[str] = [] + + def launch_nowait(self, method: str, sample: str) -> _Pending: + assert method == "generate" + self.launched.append(sample) + return _Pending(sample) + + +def test_batch_engine_stamps_train_version_and_consumes_generation_fifo() -> None: + rollout = _RolloutHandle() + engine = AsyncBatchRolloutEngine( + rollout, + process_completion=lambda _gen_id, completed: [f"{completed}/0", f"{completed}/1"], + groups_per_batch=2, + ) + + engine.submit("batch-0", behavior_version=0) + engine.submit("batch-1", behavior_version=4) + assert engine.poll() == 2 + + first = engine.pop_next_batch(train_version=4, max_policy_lag=4) + second = engine.pop_next_batch(train_version=4, max_policy_lag=4) + assert first is not None and first.groups == ["batch-0/0", "batch-0/1"] + assert first.behavior_version == 0 + assert second is not None and second.groups == ["batch-1/0", "batch-1/1"] + assert second.behavior_version == 4 + + +def test_batch_engine_rejects_non_atomic_generation() -> None: + engine = AsyncBatchRolloutEngine( + _RolloutHandle(), + process_completion=lambda _gen_id, completed: [completed], + groups_per_batch=2, + ) + engine.submit("short", behavior_version=0) + with pytest.raises(RuntimeError, match="expected groups_per_batch=2"): + engine.poll() + + +def test_batch_engine_fails_closed_on_stale_ready_batch() -> None: + engine = AsyncBatchRolloutEngine( + _RolloutHandle(), + process_completion=lambda _gen_id, completed: [completed], + groups_per_batch=1, + ) + engine.submit("old", behavior_version=0) + engine.poll() + with pytest.raises(RuntimeError, match="exceeded policy lag budget"): + engine.pop_next_batch(train_version=2, max_policy_lag=1) diff --git a/tests/trainer/test_async_policy.py b/tests/trainer/test_async_policy.py new file mode 100644 index 000000000..e4571a480 --- /dev/null +++ b/tests/trainer/test_async_policy.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +import pytest + +from unirl.train.stack.base import TrainStepResult, _aggregate_update_results +from unirl.trainer.async_policy import PolicyVersionState, launch_slots, next_hard_boundary, unwrap_replicated_int + + +def _slots( + *, + train: int = 0, + rollout: int = 0, + updates: int = 4, + max_lag: int = 0, + inflight: int = 0, + ready: int = 0, + max_inflight: int = 8, + trained: int = 0, + total: int = 20, + boundary: int = 20, +) -> int: + return launch_slots( + train_version=train, + rollout_version=rollout, + num_updates_per_batch=updates, + max_policy_lag=max_lag, + inflight_count=inflight, + ready_count=ready, + max_inflight=max_inflight, + trained_batches=trained, + num_rollouts=total, + hard_boundary=boundary, + ) + + +def test_launch_slots_uses_optimizer_update_lag() -> None: + assert _slots(max_lag=0) == 1 + assert _slots(max_lag=4) == 2 + assert _slots(train=4, max_lag=4) == 1 + assert _slots(train=8, max_lag=4) == 0 + + +def test_launch_slots_subtracts_all_outstanding_batches() -> None: + assert _slots(max_lag=8, inflight=1, ready=1) == 1 + assert _slots(max_lag=8, inflight=2, ready=1) == 0 + + +def test_launch_slots_respects_concurrency_target_and_hard_boundary() -> None: + assert _slots(max_lag=100, max_inflight=2) == 2 + assert _slots(max_lag=100, inflight=1, max_inflight=2) == 1 + assert _slots(max_lag=100, trained=3, boundary=4) == 1 + assert _slots(max_lag=100, trained=4, boundary=4) == 0 + + +@pytest.mark.parametrize( + ("updates", "max_lag", "expected_batches"), + [ + (4, 0, 1), + (4, 12, 4), + (2, 22, 12), + ], +) +def test_single_thread_cycle_exhausts_queue_before_sync( + updates: int, + max_lag: int, + expected_batches: int, +) -> None: + train = 0 + rollout = 0 + ready = 0 + consumed = 0 + while True: + ready += _slots( + train=train, + rollout=rollout, + updates=updates, + max_lag=max_lag, + ready=ready, + max_inflight=64, + trained=consumed, + total=100, + boundary=100, + ) + if ready == 0: + break + ready -= 1 + consumed += 1 + train += updates + + assert consumed == expected_batches + assert ready == 0 + + +def test_launch_slots_rejects_future_rollout_version() -> None: + with pytest.raises(ValueError, match="ahead of train_version"): + _slots(train=2, rollout=3) + + +def test_policy_versions_track_train_updates_and_rollout_sync_separately() -> None: + state = PolicyVersionState() + state.record_optimizer_updates(4) + assert state.train_version == 4 + assert state.rollout_version == 0 + assert state.rollout_lag == 4 + + state.mark_rollout_synced(4) + assert state.rollout_version == 4 + state.mark_rollout_synced(4) + assert state.train_version == 4 + assert state.rollout_version == 4 + + +def test_policy_versions_fail_closed_on_future_batch_or_sync() -> None: + state = PolicyVersionState(train_version=3, rollout_version=2) + with pytest.raises(ValueError, match="future behavior version"): + state.behavior_lag(4) + with pytest.raises(ValueError, match="future train version"): + state.mark_rollout_synced(4) + + +def test_next_hard_boundary_uses_nearest_eval_save_or_final() -> None: + assert next_hard_boundary(0, num_rollouts=20, eval_interval=4, save_interval=3) == 3 + assert next_hard_boundary(3, num_rollouts=20, eval_interval=4, save_interval=3) == 4 + assert next_hard_boundary(19, num_rollouts=20, eval_interval=4, save_interval=3) == 20 + + +def test_train_step_result_aggregates_only_committed_optimizer_updates() -> None: + committed = TrainStepResult(1.0, 1.0, 1e-6, True, [], {}, optimizer_updates=1) + skipped = TrainStepResult(float("nan"), float("nan"), 1e-6, True, [], {}, optimizer_updates=0) + result = _aggregate_update_results([committed, skipped]) + assert result.has_backward is True + assert result.optimizer_updates == 1 + + +def test_replicated_optimizer_count_must_agree() -> None: + assert unwrap_replicated_int([7, 7], name="optimizer count") == 7 + with pytest.raises(RuntimeError, match="disagrees across workers"): + unwrap_replicated_int([7, 8], name="optimizer count") diff --git a/tests/trainer/test_async_trainer_boundaries.py b/tests/trainer/test_async_trainer_boundaries.py new file mode 100644 index 000000000..abe081eff --- /dev/null +++ b/tests/trainer/test_async_trainer_boundaries.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +from unirl.trainer.async_ar import AsyncARTrainer +from unirl.trainer.async_policy import PolicyVersionState + + +class _Engine: + def __init__(self) -> None: + self.ready_count = 0 + self.quiesces = 0 + + def quiesce(self) -> None: + self.quiesces += 1 + + +class _WeightSync: + def __init__(self) -> None: + self.calls = 0 + + def sync(self) -> None: + self.calls += 1 + + +class _Rollout: + def __init__(self) -> None: + self.versions: list[int] = [] + + def set_policy_version(self, train_version: int) -> None: + self.versions.append(train_version) + + +def test_async_ar_baseline_eval_syncs_once_even_when_version_is_zero() -> None: + trainer = object.__new__(AsyncARTrainer) + trainer._policy_versions = PolicyVersionState() + trainer._rollout_initialized = False + trainer._async_engine = _Engine() + trainer.weight_sync = _WeightSync() + trainer.rollout = _Rollout() + + trainer._prepare_rollout(sync_weights=True) + trainer._prepare_rollout(sync_weights=True) + + assert trainer.weight_sync.calls == 1 + assert trainer.rollout.versions == [0] + + +def test_async_ar_sync_assigns_exact_train_version() -> None: + trainer = object.__new__(AsyncARTrainer) + trainer._policy_versions = PolicyVersionState(train_version=7, rollout_version=0) + trainer._rollout_initialized = False + trainer._async_engine = _Engine() + trainer.weight_sync = _WeightSync() + trainer.rollout = _Rollout() + + assert trainer._sync_rollout_weights() is True + assert trainer._policy_versions.rollout_version == 7 + assert trainer.rollout.versions == [7] diff --git a/unirl/rollout/README.md b/unirl/rollout/README.md index 7ea2ea69a..8671fabd3 100644 --- a/unirl/rollout/README.md +++ b/unirl/rollout/README.md @@ -64,14 +64,14 @@ wrong objective. plus a `sync:` block; *colocate* — a dedicated engine sharing GPUs with train, plus offload/onload and `sync:`. - **Driver-side async engines** (`engine/asynchronous.py`, the driver-side half next - to `engine/synchronous.py`'s worker-side sync contracts). Both engines expose the - same consumer verbs the async trainers program against: `poll` / `drain_freshest` / - `pop_evicted` / `quiesce` + engine-owned `weight_version`. `AsyncBatchRolloutEngine` - (batch granularity; non-blocking `Handle.launch_nowait` generations, stamps - versions at launch, used by `AsyncARTrainer`/`AsyncDiffusionTrainer`) and - `AsyncAgenticRolloutEngine` (trajectory granularity over the agentic rank-0 + to `engine/synchronous.py`'s worker-side sync contracts). `AsyncBatchRolloutEngine` + (batch granularity; non-blocking `Handle.launch_nowait` generations, stamps the + synced train version at launch, and exposes completion-order FIFO train batches + to `AsyncARTrainer`/`AsyncDiffusionTrainer`) and `AsyncAgenticRolloutEngine` + (trajectory granularity over the agentic rank-0 coordinator; normalizes the `[0]` unwraps, assembles n-sibling GRPO groups, - stamps versions at completion, used by the partial/async agentic trainers). + stamps sync-generation versions at completion, and retains the + `drain_freshest`/`pop_evicted` surface used by partial/async agentic trainers). **Extending it:** a new single-turn engine adds `engine//config.py` (a `BaseEngineConfig` whose `make_engine(**deps)` lazily imports and builds it) and @@ -94,7 +94,9 @@ implements its weight-receive method and a matching `sync:` handler in engine also can't live on a `layout: separate` slab — `_build_rollout` raises. - **Quiesce before weight sync / eval / checkpoint on the batch async path** — `AsyncBatchRolloutEngine.quiesce()` drains every in-flight generation; a - weight + KV update corrupts one mid-flight. The agentic quiesce is a + weight + KV update corrupts one mid-flight. A batch weight publication also + requires the completed FIFO to be empty; hard-boundary admission guarantees + this rather than silently discarding data. The agentic quiesce is a turn-boundary `abort` + final poll, folded into `AsyncAgenticRolloutEngine.quiesce()`. Reap-vs-launch ordering is trainer statement order (diffusion polls before topping up; see its `_next_step`). diff --git a/unirl/rollout/engine/asynchronous.py b/unirl/rollout/engine/asynchronous.py index 0e103dee8..5a6d1e273 100644 --- a/unirl/rollout/engine/asynchronous.py +++ b/unirl/rollout/engine/asynchronous.py @@ -5,18 +5,18 @@ the driver side — everything is single-threaded, lock-free, and ray-free (non-blocking dispatch is ``Handle.launch_nowait``). -Mechanisms (policy-free — launch ceilings, reap/launch ordering, and step +Mechanisms (policy-free — launch admission, reap/launch ordering, and step loops live in the trainers): - :class:`VersionedBuffer` — payload-agnostic freshness/staleness buffer. - :class:`InflightPool` — non-blocking pool of distributed ``generate`` calls. -Engines, sharing one consumer surface (``poll`` / ``drain_freshest`` / -``pop_evicted`` / ``quiesce`` + engine-owned ``weight_version``): +Engines share launch/poll/quiesce mechanisms but own different queue semantics: - :class:`AsyncBatchRolloutEngine` — batch granularity over a single-turn engine slab; one ``submit`` is one non-blocking distributed ``generate``. - ``(weight_version, gen_id)`` are stamped at LAUNCH. + ``(behavior_version, gen_id)`` are stamped at LAUNCH and one complete + generation is consumed atomically in completion-order FIFO. - :class:`AsyncAgenticRolloutEngine` — trajectory granularity over the ``AgenticRolloutEngine`` rank-0 coordinator; ``submit`` fires a task-pool drive and completions stream in via ``poll``. ``(weight_version, gen_id)`` @@ -32,11 +32,13 @@ from __future__ import annotations import logging +from collections import deque from dataclasses import dataclass from typing import ( TYPE_CHECKING, Any, Callable, + Deque, Dict, Generic, Iterable, @@ -106,22 +108,63 @@ def pop_evicted(self) -> List[T]: return evicted -Complete = Callable[[int, int, Any], None] +@dataclass(frozen=True) +class RolloutBatch: + """One atomic train batch produced by one batch-generation launch.""" + + groups: List["Sample"] + behavior_version: int + gen_id: int + + +class RolloutBatchQueue: + """Completion-order FIFO for single-turn batch generations.""" + + def __init__(self) -> None: + self._items: Deque[RolloutBatch] = deque() + + def put(self, batch: RolloutBatch) -> None: + self._items.append(batch) + + def __len__(self) -> int: + return len(self._items) + + def pop_next( + self, + *, + train_version: int, + max_policy_lag: int, + ) -> Optional[RolloutBatch]: + if not self._items: + return None + item = self._items[0] + lag = int(train_version) - item.behavior_version + if lag < 0: + raise RuntimeError( + f"generation {item.gen_id} has future behavior version " + f"{item.behavior_version} > train version {train_version}" + ) + if lag > int(max_policy_lag): + raise RuntimeError(f"generation {item.gen_id} exceeded policy lag budget: lag={lag} > max={max_policy_lag}") + return self._items.popleft() + + +CompleteGeneration = Callable[[int, int, Any], None] @dataclass(frozen=True) class _InflightJob: gen_id: int - weight_version: int + behavior_version: int pending: Any class InflightPool: """Non-blocking pool of distributed ``generate`` launches on a rollout Handle. - Mechanism only: launch ceilings, reap/launch ordering, and step loops are + Mechanism only: launch admission, reap/launch ordering, and step loops are caller policy. Jobs are launched via ``Handle.launch_nowait`` and completed - through ``complete(gen_id, weight_version, payload)`` — all of ``complete``'s + through ``complete(gen_id, behavior_version, payload)`` — all of ``complete``'s fallible work must happen before it mutates caller state, because a job whose completion raises stays in flight for retry. """ @@ -138,14 +181,14 @@ def next_gen_id(self) -> int: def __len__(self) -> int: return len(self._jobs) - def launch(self, sample: Any, *, weight_version: int) -> int: + def launch(self, sample: Any, *, behavior_version: int) -> int: gen_id = self._next_gen_id pending = self._rollout.launch_nowait("generate", sample) - self._jobs.append(_InflightJob(gen_id, int(weight_version), pending)) + self._jobs.append(_InflightJob(gen_id, int(behavior_version), pending)) self._next_gen_id += 1 return gen_id - def reap_ready(self, complete: Complete) -> int: + def reap_ready(self, complete: CompleteGeneration) -> int: """Complete every ready job; leave unresolved and failed jobs in flight. A job whose ``result()``/``complete`` raises stays in flight for retry; @@ -161,7 +204,7 @@ def reap_ready(self, complete: Complete) -> int: still.append(job) continue try: - complete(job.gen_id, job.weight_version, job.pending.result()) + complete(job.gen_id, job.behavior_version, job.pending.result()) completed += 1 except Exception as exc: still.append(job) @@ -174,7 +217,7 @@ def reap_ready(self, complete: Complete) -> int: raise first_error return completed - def drain_all(self, complete: Complete) -> int: + def drain_all(self, complete: CompleteGeneration) -> int: """Quiesce: complete every job, blocking as needed. Same error contract as :meth:`reap_ready`.""" jobs, self._jobs = self._jobs, [] @@ -182,7 +225,7 @@ def drain_all(self, complete: Complete) -> int: completed = 0 for job in jobs: try: - complete(job.gen_id, job.weight_version, job.pending.result()) + complete(job.gen_id, job.behavior_version, job.pending.result()) completed += 1 except Exception as exc: self._jobs.append(job) @@ -201,12 +244,11 @@ def wait_oldest(self) -> None: class AsyncBatchRolloutEngine: - """Batch-granular async engine over a ``SyncRolloutEngine`` slab Handle; buffers ``Sample`` groups. + """Batch-granular async engine over a ``SyncRolloutEngine`` slab Handle. - ``complete(gen_id, completed) -> groups`` runs at reap time — scoring must - precede training, and on transfer-sensitive backends the next launch. All - of ``complete``'s fallible work happens before any buffer mutation, so a - failed job stays in flight for retry without double-inserting groups. + One generation must produce exactly ``groups_per_batch`` groups and enters a + completion-order FIFO atomically. ``behavior_version`` is captured at launch + from the exact train snapshot currently synced to the rollout engine. ``quiesce()`` (drain everything) is MANDATORY before a weight sync, eval, or checkpoint: a weight + KV update corrupts an in-flight generation. @@ -216,21 +258,16 @@ def __init__( self, rollout: Any, *, - complete: Callable[[int, "Sample"], List["Sample"]], + process_completion: Callable[[int, "Sample"], List["Sample"]], + groups_per_batch: int, start_gen_id: int = 0, ) -> None: - self._complete = complete + if int(groups_per_batch) < 1: + raise ValueError(f"groups_per_batch must be >= 1, got {groups_per_batch}") + self._process_completion = process_completion + self._groups_per_batch = int(groups_per_batch) self._pool = InflightPool(rollout, start_gen_id=start_gen_id) - self._buffer: VersionedBuffer["Sample"] = VersionedBuffer() - self._weight_version = 0 - - @property - def weight_version(self) -> int: - return self._weight_version - - def bump_weight_version(self) -> int: - self._weight_version += 1 - return self._weight_version + self._ready = RolloutBatchQueue() @property def next_gen_id(self) -> int: @@ -238,34 +275,52 @@ def next_gen_id(self) -> int: return self._pool.next_gen_id @property - def inflight(self) -> int: + def inflight_count(self) -> int: return len(self._pool) - def submit(self, sample: "Sample") -> int: - """Launch one non-blocking distributed ``generate``; stamps the CURRENT version.""" - return self._pool.launch(sample, weight_version=self._weight_version) + @property + def ready_count(self) -> int: + return len(self._ready) + + def submit(self, sample: "Sample", *, behavior_version: int) -> int: + """Launch one generation under the supplied rollout policy version.""" + + return self._pool.launch(sample, behavior_version=behavior_version) def poll(self) -> int: return self._pool.reap_ready(self._on_complete) - def drain_freshest(self, n: int, *, max_staleness: int) -> Optional[List["Sample"]]: - return self._buffer.drain_freshest(n, current_version=self._weight_version, max_staleness=max_staleness) - - def pop_evicted(self) -> List["Sample"]: - return self._buffer.pop_evicted() - - def quiesce(self) -> List["Sample"]: + def pop_next_batch( + self, + *, + train_version: int, + max_policy_lag: int, + ) -> Optional[RolloutBatch]: + return self._ready.pop_next( + train_version=train_version, + max_policy_lag=max_policy_lag, + ) + + def quiesce(self) -> None: self._pool.drain_all(self._on_complete) - return [] def wait_oldest(self) -> None: """Block until the oldest in-flight generation resolves (reap via ``poll``).""" self._pool.wait_oldest() - def _on_complete(self, gen_id: int, weight_version: int, completed: "Sample") -> None: - groups = self._complete(gen_id, completed) # fallible (scoring) before any buffer put - for group in groups: - self._buffer.put(group, weight_version=weight_version, gen_id=gen_id) + def _on_complete(self, gen_id: int, behavior_version: int, completed: "Sample") -> None: + groups = self._process_completion(gen_id, completed) # fallible before any queue mutation + if len(groups) != self._groups_per_batch: + raise RuntimeError( + f"generation {gen_id} produced {len(groups)} groups; expected groups_per_batch={self._groups_per_batch}" + ) + self._ready.put( + RolloutBatch( + groups=list(groups), + behavior_version=int(behavior_version), + gen_id=int(gen_id), + ) + ) def root_of(traj: "Sample") -> str: @@ -407,27 +462,9 @@ def _ingest(self, completed: List["Sample"]) -> int: return len(completed) -def launch_ceiling(rollout_id: int, *, sync_interval: int, max_staleness: int, num_rollouts: int) -> int: - """The batch trainers' on-policy launch clamp — trainer POLICY, defined once. - - A generation launched now is consumed later, so how far ahead the gen_id - allocator may run is bounded to ``max_staleness`` weight-sync windows: - ``max_staleness=0`` ⇒ never launch into a future sync-window ⇒ no - generation crosses a sync ⇒ ``ratio≈1`` (on-policy). - - OWNERSHIP: this is trainer-side POLICY, not engine surface — its vocabulary - (``rollout_id`` / ``sync_interval`` / ``num_rollouts``) is the trainers', - the engine classes never call it, and it must never become an engine - method. It is hosted in this module only because it is the two batch - trainers' one shared torch-free home; the step loops that use it stay in - the trainers as visible statement order. - """ - return min(num_rollouts, ((rollout_id // sync_interval) + 1 + max_staleness) * sync_interval) - - __all__ = [ "AsyncAgenticRolloutEngine", "AsyncBatchRolloutEngine", - "launch_ceiling", + "RolloutBatch", "root_of", ] diff --git a/unirl/rollout/engine/synchronous.py b/unirl/rollout/engine/synchronous.py index 87966a9a0..74496f41e 100644 --- a/unirl/rollout/engine/synchronous.py +++ b/unirl/rollout/engine/synchronous.py @@ -157,6 +157,20 @@ class SyncRolloutEngine(BaseRolloutEngine, ABC): def generate(self, sample: Sample) -> Sample: """Synchronously fill and return one request ``Sample``.""" + @distributed(dispatch_mode=Dispatch.BROADCAST) + def set_policy_version(self, train_version: int) -> None: + """Assign the optimizer-update version of the currently loaded weights. + + Weight transports may receive one model through several buckets, so + receiver-local API-call counters are not valid policy provenance. The + trainer calls this only after a full weight publication succeeds. + """ + + train_version = int(train_version) + if train_version < 0: + raise ValueError(f"train_version must be >= 0, got {train_version}") + self._weight_version = train_version + def _stamp_weight_version(self, sample: Sample) -> Sample: """Stamp ``self._weight_version`` onto the frontier (last) gen Part.""" v = getattr(self, "_weight_version", None) diff --git a/unirl/train/backend/base_backend.py b/unirl/train/backend/base_backend.py index 7fe51d10a..1d560bee3 100644 --- a/unirl/train/backend/base_backend.py +++ b/unirl/train/backend/base_backend.py @@ -370,6 +370,12 @@ def on_rollout_end(self) -> None: if self.ema is not None: self.ema.on_rollout_end(self._optimizer_step_count) + @distributed(dispatch_mode=Dispatch.BROADCAST) + def get_optimizer_step_count(self) -> int: + """Return the authoritative number of committed optimizer updates.""" + + return int(self._optimizer_step_count) + @property def rollout_adapter_name(self) -> str: """Adapter the rollout must sample under (single source of truth). diff --git a/unirl/train/stack/base.py b/unirl/train/stack/base.py index 1af3b0c47..058d61f0e 100644 --- a/unirl/train/stack/base.py +++ b/unirl/train/stack/base.py @@ -41,6 +41,7 @@ from __future__ import annotations import logging +import math from contextlib import nullcontext from dataclasses import dataclass, replace from typing import Dict, List, Mapping, Optional, Tuple @@ -61,7 +62,7 @@ @dataclass(frozen=True) class TrainStepResult: - """Result of one full optimizer step on this stage.""" + """Result of one ``train_track`` call, possibly spanning multiple optimizer updates.""" loss: float grad_norm: float @@ -69,6 +70,9 @@ class TrainStepResult: has_backward: bool micros: List[AlgorithmStepResult] metrics: Mapping[str, object] + # Number of optimizer steps that actually committed. This is zero when an + # update had no backward or the backend rejected a non-finite grad norm. + optimizer_updates: int per_update: Tuple[Mapping[str, object], ...] = () @@ -92,6 +96,7 @@ def _aggregate_update_results(results: List["TrainStepResult"]) -> "TrainStepRes has_backward=any(r.has_backward for r in results), micros=micros, metrics=metrics, + optimizer_updates=sum(r.optimizer_updates for r in results), ) @@ -302,6 +307,7 @@ def _run_update( has_backward=has_backward, micros=micro_results, metrics=aggregated_metrics, + optimizer_updates=int(has_backward and math.isfinite(grad_norm)), ) def on_rollout_end(self) -> None: @@ -511,7 +517,13 @@ def _run_updates( return results[0] aggregated = _aggregate_update_results(results) per_update = tuple( - {**dict(r.metrics), "loss": float(r.loss), "grad_norm": float(r.grad_norm), "lr": float(r.lr)} + { + **dict(r.metrics), + "loss": float(r.loss), + "grad_norm": float(r.grad_norm), + "lr": float(r.lr), + "optimizer_updates": float(r.optimizer_updates), + } for r in results ) return replace(aggregated, per_update=per_update) diff --git a/unirl/train/unified_model_stack.py b/unirl/train/unified_model_stack.py index a43ad5eed..24e0e9463 100644 --- a/unirl/train/unified_model_stack.py +++ b/unirl/train/unified_model_stack.py @@ -34,6 +34,7 @@ from __future__ import annotations import logging +import math from contextlib import nullcontext from dataclasses import replace from typing import Dict, List, Mapping, Tuple @@ -215,6 +216,7 @@ def _backward_part( has_backward=has_backward, micros=micros, metrics=aggregated, + optimizer_updates=0, ) return partial, has_backward @@ -250,9 +252,16 @@ def _train_one_step( logger.warning("UnifiedModelTrainStack._train_one_step: no algorithm reported backward; skipping step.") lr = self._current_lr() + optimizer_updates = int(any_backward and math.isfinite(grad_norm)) for name, r in list(results.items()): results[name] = TrainStepResult( - loss=r.loss, grad_norm=grad_norm, lr=lr, has_backward=r.has_backward, micros=r.micros, metrics=r.metrics + loss=r.loss, + grad_norm=grad_norm, + lr=lr, + has_backward=r.has_backward, + micros=r.micros, + metrics=r.metrics, + optimizer_updates=optimizer_updates, ) return results diff --git a/unirl/train_async_ar.py b/unirl/train_async_ar.py index 37c359ab9..9480f3dfd 100755 --- a/unirl/train_async_ar.py +++ b/unirl/train_async_ar.py @@ -16,9 +16,8 @@ Constraint: ``train_fraction * num_devices`` and ``(1-train_fraction) * num_devices`` must both be integers, AND ``batch_size * samples_per_prompt`` must be divisible by each slab size (DP_SCATTER divisibility). - * ``max_inflight`` — concurrent generations (overlap depth). ``1`` ≈ one-step pipeline. - * ``buffer_max_staleness`` — weight-syncs a buffered group may cross. ``0``/unset = - on-policy (``ratio≈1``); ``>0`` = off-policy continuous buffer. + * ``max_inflight`` — concurrent generations (resource/overlap limit). + * ``max_policy_lag`` — maximum optimizer-update lag at batch admission. """ from __future__ import annotations @@ -56,11 +55,10 @@ def main(cfg: DictConfig) -> None: eval_temperature=float(cfg.get("eval_temperature", 1.0)), train_fraction=float(cfg.get("train_fraction", 0.5)), max_inflight=int(cfg.get("max_inflight", 1)), - buffer_max_staleness=cfg.get("buffer_max_staleness"), + max_policy_lag=int(cfg.get("max_policy_lag", 0)), ) trainer.train( num_rollouts=int(cfg.get("num_rollouts", 100)), - weight_sync_interval=int(cfg.get("weight_sync_interval", 1)), save_interval=int(cfg.get("save_interval", 0)), save_dir=cfg.get("save_dir"), load_dir=cfg.get("load_dir"), diff --git a/unirl/train_async_diffusion.py b/unirl/train_async_diffusion.py index 6b003f0cb..37b8879ad 100755 --- a/unirl/train_async_diffusion.py +++ b/unirl/train_async_diffusion.py @@ -15,8 +15,7 @@ Extra config knobs vs the synchronous separate recipe: * ``max_inflight`` — must be ``1``; other values fail during trainer initialization. - * ``buffer_max_staleness`` — regular rollout-weight syncs a buffered group may - cross. ``0``/unset never crosses a sync; ``>0`` enables bounded policy lag. + * ``max_policy_lag`` — maximum optimizer-update lag at batch admission. ``layout`` is forced to ``separate`` (async needs disjoint train/rollout slabs). """ @@ -57,11 +56,10 @@ def main(cfg: DictConfig) -> None: eval_rewards_cfg=cfg.get("eval_rewards"), task_config=cfg.get("task_config"), max_inflight=int(cfg.get("max_inflight", 1)), - buffer_max_staleness=cfg.get("buffer_max_staleness"), + max_policy_lag=int(cfg.get("max_policy_lag", 0)), ) trainer.train( num_rollouts=cfg.get("num_rollouts", 100), - weight_sync_interval=cfg.get("weight_sync_interval", 1), save_interval=cfg.get("save_interval", 0), save_dir=cfg.get("save_dir"), load_dir=cfg.get("load_dir"), diff --git a/unirl/trainer/README.md b/unirl/trainer/README.md index 842ff1914..aec0c5a5a 100644 --- a/unirl/trainer/README.md +++ b/unirl/trainer/README.md @@ -59,8 +59,8 @@ The current trainer surface is: | `DiffusionTrainer` | one diffusion `Part` → one `TrainStack` | Reference diffusion loop; supports trainside or dedicated rollout, optional separate reward GPUs, FSDP offload, and DiffusionNFT's EMA-adapter rollout. | | `ARTrainer` | one AR `Part` → one `TrainStack` | Text or multimodal AR rollout with group/global advantage normalization and optional token-balanced DP shards. | | `SFTTrainer` | dataset records → one standalone training `Part` | Reuses the RL TrainStack without rollout, reward, or advantages; owns exact epoch/cursor resume and full-set evaluation. | -| `AsyncARTrainer` | buffered AR `Sample` groups → one `TrainStack` | Separate train/rollout slabs with resident generation, bounded staleness, and quiescence before sync, eval, or checkpoint. | -| `AsyncDiffusionTrainer` | buffered diffusion `Sample` groups → one `TrainStack` | The same separate-slab async loop for DiT. Requires `max_inflight=1` and reaps each generation before launching the next, so the cross-slab trajectory transfer never queues behind a fresh generation. | +| `AsyncARTrainer` | FIFO AR generation batch → one `TrainStack` | Separate train/rollout slabs with resident generation, optimizer-update policy-lag admission, and quiescence before weight sync, eval, or checkpoint. | +| `AsyncDiffusionTrainer` | FIFO diffusion generation batch → one `TrainStack` | The same update-versioned separate-slab loop for DiT. Requires `max_inflight=1` and reaps each generation before launching the next, so the cross-slab trajectory transfer never queues behind a fresh generation. | | `PETrainer` | `ar` + `diffusion` Parts → two `TrainStack`s | Composed prompt-rewrite/image rollout; image rewards propagate to AR rewrites. `freeze_llm=true` trains and checkpoints diffusion only. | | `UnifiedModelTrainer` | whole `Sample` → one `UnifiedModelTrainStack` | AR and image losses accumulate into shared-backbone optimizer steps while prompt-tree lineage remains intact during DP scatter. | | `AgenticTrainer` / `AgenticEnvTrainer` | variable-depth `List[Sample]` → concatenated turn `Part` | Barrier multi-turn tool use. The base variant scores terminal answers; the env variant consumes per-trajectory environment returns. | @@ -69,9 +69,9 @@ The current trainer surface is: The async variants program against the driver-side async engines in `unirl/rollout/engine/asynchronous.py`: `AsyncBatchRolloutEngine` (AR/diffusion — non-blocking -batched generations, launch-time version stamps) and `AsyncAgenticRolloutEngine` +batched generations, launch-time train-version stamps, completion-order FIFO batches) and `AsyncAgenticRolloutEngine` (partial/async agentic — trajectory drives, group assembly, completion-time stamps). -The trainers keep the policy: launch ceilings, reap-vs-launch order, quiesce points, +The trainers keep the policy: optimizer-update launch slots, reap-vs-launch order, quiesce points, and tail carry/drop. **Extending it:** a new domain is a new `Trainer(BaseTrainer)` that builds its @@ -180,8 +180,10 @@ so EMA decay schedules continue) and resumes the loop from the saved step. Synchronous Sample-based trainers continue `training_progress` and driver-authored x_T scheduling, fast-forward a deterministically seeded data stream, and force the restored weights into a freshly started rollout engine -when needed. `AsyncARTrainer` also fast-forwards its deterministic input stream -but rebuilds its rollout buffer. Partial-agentic resume can consume a different +when needed. Async AR/diffusion resume reads the backend optimizer count as the +train version, fast-forwards the deterministic input stream to the saved +rollout step, and syncs those restored weights into the fresh engine. +Partial-agentic resume can consume a different input sequence when an earlier over-sampled drive required refills, and ReFL does not currently fast-forward its data source. @@ -243,9 +245,9 @@ an evaluation and checkpoint fall on the same step, evaluation runs first. - `DiffusionTrainer`, `PETrainer`, and `UnifiedModelTrainer` report image reward; optional `eval_rewards` suites can score the same generated samples or their own prompt sets. PE scores only the - diffusion/image frontier. `AsyncDiffusionTrainer` quiesces first and then scores - the policy already resident in its rollout engine, without a weight sync and - without offloading that engine afterwards. + diffusion/image frontier. `AsyncDiffusionTrainer` reaches an empty hard + boundary, syncs the current train version when needed, and then scores + the resident rollout engine without offloading it afterwards. - Agentic evaluation is not implemented. Barrier and partial variants raise if evaluation is enabled; async variants currently force it off. diff --git a/unirl/trainer/agentic.py b/unirl/trainer/agentic.py index 9013c6465..ee4855c26 100644 --- a/unirl/trainer/agentic.py +++ b/unirl/trainer/agentic.py @@ -216,7 +216,7 @@ def _advantage_train_and_log( ) if not train_parts: # pathological: every sampled trajectory failed to generate logger.warning("AgenticTrainer rollout %d produced no trainable turns.", rollout_id) - return TrainStepResult(0.0, 0.0, 0.0, False, [], {}), mean_reward + return TrainStepResult(0.0, 0.0, 0.0, False, [], {}, optimizer_updates=0), mean_reward train_part = Part.concat(train_parts) train_part = self._pad_to_dp_multiple(train_part) diff --git a/unirl/trainer/agentic_async.py b/unirl/trainer/agentic_async.py index 27302191c..0728b6712 100644 --- a/unirl/trainer/agentic_async.py +++ b/unirl/trainer/agentic_async.py @@ -297,7 +297,7 @@ def _train_on_groups( depths = [len(tr.gen_parts()) for tr in trajs] if not train_parts: logger.warning("AsyncAgenticTrainer rollout %d produced no trainable turns.", rollout_id) - return TrainStepResult(0.0, 0.0, 0.0, False, [], {}), mean_reward + return TrainStepResult(0.0, 0.0, 0.0, False, [], {}, optimizer_updates=0), mean_reward train_part = self._pad_to_dp_multiple(Part.concat(train_parts)) result = self.stack.train_track(train_part, training_progress=float(training_progress)) diff --git a/unirl/trainer/async_ar.py b/unirl/trainer/async_ar.py index 85b72d33d..1f3029e22 100644 --- a/unirl/trainer/async_ar.py +++ b/unirl/trainer/async_ar.py @@ -9,20 +9,19 @@ ONE single-threaded loop (slime's "one trainer loop; async-depth is a knob" principle, implemented with UniRL-native non-blocking Ray dispatch instead of -slime's thread+asyncio). The async behavior is set by **two numeric knobs**: +slime's thread+asyncio). Async freshness is measured directly in committed +optimizer updates: * ``max_inflight`` — how many generations run concurrently (overlap/parallelism depth). ``1`` ≈ the classic one-step pipeline; higher fans out more. -* ``buffer_max_staleness`` — how many weight-syncs a buffered group may cross - before it is evicted. ``0`` (default) = **on-policy**: the launch clamp never - lets a generation cross a weight sync, so ``ratio≈1`` (the colocate-parity - regime). ``>0`` = **off-policy continuous buffer**: generations may run ahead - across syncs, bounded by eviction; the rollout-anchored DRPO ratio absorbs it. +* ``max_policy_lag`` — inclusive train-minus-behavior optimizer-update + lag at batch admission. ``0`` aligns policy versions; the rollout-anchored + PPO ratio remains the numerical source of truth. Generation runs through :class:`~unirl.rollout.engine.asynchronous.AsyncBatchRolloutEngine` (non-blocking Ray futures over the rollout Handle) on the single driver thread — -no producer thread, no locks; the trainer's ``_next_step`` loop owns the policy -(launch ceiling, launch-then-reap order). Draining all in-flight generations +no producer thread, no locks; the trainer's ``_next_rollout_batch`` loop owns the policy +(optimizer-update launch admission, launch-then-reap order). Draining all in-flight generations before each weight sync is **mandatory** (the engine corrupts an in-flight generation when weights + KV cache update mid-flight); this is the single-threaded ``_drain_all`` quiesce. @@ -45,9 +44,10 @@ from unirl.distributed.group.placement import placement, remote from unirl.distributed.tensor import hydrate from unirl.models.qwen3_5.validation import validate_qwen3_5_training_contract -from unirl.rollout.engine.asynchronous import AsyncBatchRolloutEngine, launch_ceiling +from unirl.rollout.engine.asynchronous import AsyncBatchRolloutEngine, RolloutBatch from unirl.train.stack import TrainStepResult from unirl.trainer.ar import ARTrainer +from unirl.trainer.async_policy import PolicyVersionState, launch_slots, next_hard_boundary, unwrap_replicated_int from unirl.trainer.base import BaseTrainer, build_sampling_dict from unirl.types.sample import Sample from unirl.types.sampling import BaseSamplingParams, total_samples_per_prompt @@ -101,7 +101,7 @@ def __init__( eval_temperature: float = 1.0, train_fraction: float = 0.5, max_inflight: int = 1, - buffer_max_staleness: Optional[int] = None, + max_policy_lag: int = 0, ) -> None: validate_qwen3_5_training_contract( pipeline_cfg=pipeline_cfg, @@ -138,7 +138,26 @@ def __init__( self._train_fraction = float(train_fraction) self._max_inflight = max(1, int(max_inflight)) - self._buffer_max_staleness = buffer_max_staleness + self._max_policy_lag = int(max_policy_lag) + self._num_updates_per_batch = int(stack_cfg.get("num_updates_per_batch", 1)) + if self._max_policy_lag < 0: + raise ValueError(f"max_policy_lag must be >= 0, got {self._max_policy_lag}") + if self._num_updates_per_batch < 1: + raise ValueError(f"stack.num_updates_per_batch must be >= 1, got {self._num_updates_per_batch}") + freshness_depth = self._max_policy_lag // self._num_updates_per_batch + 1 + if self._max_inflight > freshness_depth: + logger.warning( + "max_inflight=%d exceeds the policy-lag admission depth %d; the extra concurrency cannot be used", + self._max_inflight, + freshness_depth, + ) + if freshness_depth == 1: + logger.warning( + "async policy-lag settings admit one generation at a time; " + "generation cannot overlap the preceding train batch" + ) + self._policy_versions = PolicyVersionState() + self._rollout_initialized = False self._train_devices = int(round(self.num_devices * self._train_fraction)) if self._train_devices <= 0 or self._train_devices >= self.num_devices: raise ValueError( @@ -177,13 +196,14 @@ def __init__( with placement(self.pool, fraction=1.0 - self._train_fraction, shared_workers=True): self.rollout = remote(**rollout_parsed) - if self.weight_sync is not None: - self._connect_separate(sync_cfg) + if self.weight_sync is None: + raise ValueError("AsyncARTrainer requires a cross-slab weight sync; add a `sync:` block.") + self._connect_separate(sync_cfg) def _prepare_rollout(self, *, sync_weights: bool) -> bool: """Sync a resident separate-slab engine without colocate handoffs.""" - if sync_weights and self.weight_sync is not None: - self.weight_sync.sync() + if sync_weights: + self._sync_rollout_weights() return False def _finish_rollout(self, *, train_state_offloaded: bool) -> None: @@ -244,6 +264,32 @@ def _drain_all(self) -> None: """ self._async_engine.quiesce() + def _sync_rollout_weights(self, *, force: bool = False) -> bool: + """Load the current train weights into an empty rollout engine.""" + + versions = self._policy_versions + if not force and self._rollout_initialized and versions.rollout_version == versions.train_version: + return False + self._drain_all() + if self._async_engine.ready_count != 0: + raise RuntimeError( + f"cannot sync rollout weights with completed batches queued: " + f"ready_count={self._async_engine.ready_count}" + ) + target = versions.train_version + self.weight_sync.sync() + self.rollout.set_policy_version(target) + versions.mark_rollout_synced(target) + self._rollout_initialized = True + return True + + def _policy_metrics(self, batch: RolloutBatch) -> Dict[str, float]: + versions = self._policy_versions + return { + "async/behavior_version": float(batch.behavior_version), + "async/behavior_lag": float(versions.behavior_lag(batch.behavior_version)), + } + def _advantage_and_train( self, sample: Sample, @@ -251,8 +297,9 @@ def _advantage_and_train( training_progress: float, rollout_id: int, t0: Optional[float] = None, + extra_metrics: Optional[Dict[str, float]] = None, ) -> Tuple[TrainStepResult, float]: - """Advantage + optimizer step for a SCORED ``Sample`` (rewards already attached).""" + """Advantage + optimizer updates for a scored ``Sample`` (rewards already attached).""" if t0 is None: t0 = time.perf_counter() part = sample.parts[-1] @@ -270,12 +317,22 @@ def _advantage_and_train( if self.balance_shards: train_part = part.balance_shards(self._train_devices) result = self.stack.train_track(train_part, training_progress=float(training_progress)) + self._policy_versions.record_optimizer_updates(int(result.optimizer_updates)) + if extra_metrics is not None: + extra_metrics.update( + { + "async/train_version": float(self._policy_versions.train_version), + "async/rollout_lag": float(self._policy_versions.rollout_lag), + "async/optimizer_updates": float(result.optimizer_updates), + } + ) self.wandb_logger.log_rollout_step( rollout_id, result, sample, step_time_s=time.perf_counter() - t0, trunc_len=getattr(self.sampling_params.get("ar"), "max_new_tokens", None), + extra_metrics=extra_metrics, ) self._reset_transport_buffers() return result, mean_reward @@ -284,65 +341,90 @@ def train( self, *, num_rollouts: int, - weight_sync_interval: int = 1, save_interval: int = 0, save_dir: Optional[str] = None, load_dir: Optional[str] = None, save_mode: str = "full", ) -> None: - interval = max(1, weight_sync_interval) - stale = self._buffer_max_staleness if self._buffer_max_staleness is not None else 0 - M = self._max_inflight - start_rollout = self.maybe_load_checkpoint(load_dir, num_rollouts=num_rollouts) resumed = bool(load_dir) + train_version = unwrap_replicated_int( + self.backend.get_optimizer_step_count(), + name="backend optimizer step count", + ) + self._policy_versions = PolicyVersionState(train_version=train_version) for _ in range(start_rollout): self.data_source.get_samples(self.batch_size) self._init_wandb( num_rollouts=num_rollouts, extra={ "adv_normalization_scope": self.adv_normalization_scope, - "max_inflight": M, - "buffer_max_staleness": stale, - "weight_sync_interval": interval, + "max_inflight": self._max_inflight, + "max_policy_lag": self._max_policy_lag, + "num_updates_per_batch": self._num_updates_per_batch, }, ) self._async_engine = AsyncBatchRolloutEngine( self.rollout, - complete=self._score_completed, + process_completion=self._score_completed, + groups_per_batch=self.batch_size, start_gen_id=start_rollout, ) - if resumed and self.weight_sync is not None: - self.weight_sync.sync() + if resumed: + self._sync_rollout_weights(force=True) if self.eval_interval > 0: self.evaluate(rollout_id=-1) try: for rollout_id in range(start_rollout, num_rollouts): t0 = time.perf_counter() - picked = self._next_step(rollout_id, interval, M, stale, num_rollouts) - sample = Sample.concat(picked) + hard_boundary = next_hard_boundary( + rollout_id, + num_rollouts=num_rollouts, + eval_interval=self.eval_interval, + save_interval=save_interval, + ) + batch = self._next_rollout_batch( + rollout_id, + num_rollouts=num_rollouts, + hard_boundary=hard_boundary, + ) + sample = Sample.concat(batch.groups) training_progress = rollout_id / max(1, num_rollouts - 1) result, mean_reward = self._advantage_and_train( - sample, training_progress=training_progress, rollout_id=rollout_id, t0=t0 + sample, + training_progress=training_progress, + rollout_id=rollout_id, + t0=t0, + extra_metrics=self._policy_metrics(batch), ) self.wandb_logger.log_progress(rollout_id, num_rollouts, result, mean_reward, logger=logger) step = rollout_id + 1 - if self.eval_interval > 0 and step % self.eval_interval == 0: - self._drain_all() + eval_due = self.eval_interval > 0 and step % self.eval_interval == 0 + save_due = save_interval > 0 and (step % save_interval == 0 or step >= num_rollouts) + sync_due = step < num_rollouts and self._policy_versions.rollout_lag > self._max_policy_lag + if eval_due or save_due or sync_due: + if self._async_engine.inflight_count + self._async_engine.ready_count != 0: + raise RuntimeError( + "async sync boundary retained rollout work: " + f"inflight_count={self._async_engine.inflight_count}, " + f"ready_count={self._async_engine.ready_count}" + ) + self._sync_rollout_weights() + + if eval_due: self.evaluate(rollout_id=rollout_id) - if save_interval > 0 and (step % save_interval == 0 or step >= num_rollouts): - self._drain_all() + if save_due: self.maybe_save_checkpoint( - rollout_id, num_rollouts, save_interval=save_interval, save_dir=save_dir, save_mode=save_mode + rollout_id, + num_rollouts, + save_interval=save_interval, + save_dir=save_dir, + save_mode=save_mode, ) - if step % interval == 0 and self.weight_sync is not None: - self._drain_all() - self.weight_sync.sync() - self._async_engine.bump_weight_version() finally: active_exception = sys.exc_info()[0] is not None try: @@ -354,34 +436,42 @@ def train( finally: self._finish_wandb() - def _next_step( + def _next_rollout_batch( self, rollout_id: int, - interval: int, - M: int, - stale: int, + *, num_rollouts: int, - ) -> List[Sample]: - """Top up launches, reap completed generations, and return the freshest - ``batch_size`` scored group Samples for ``rollout_id`` (blocking on the - oldest in-flight generation if the buffer is short). - - The launch clamp is the load-bearing on-policy guarantee: a generation - launched now is consumed later, so bound how far ahead we launch to - ``stale`` weight-syncs. ``stale=0`` ⇒ never launch into a future - sync-window ⇒ no generation crosses a sync ⇒ ``ratio≈1`` (on-policy). - """ + hard_boundary: int, + ) -> RolloutBatch: + """Launch, reap, and consume one completion-order FIFO train batch.""" + engine = self._async_engine while True: - ceiling = launch_ceiling(rollout_id, sync_interval=interval, max_staleness=stale, num_rollouts=num_rollouts) - while engine.next_gen_id < ceiling and engine.inflight < M: - engine.submit(self._build_async_sample(engine.next_gen_id)) + slots = launch_slots( + train_version=self._policy_versions.train_version, + rollout_version=self._policy_versions.rollout_version, + num_updates_per_batch=self._num_updates_per_batch, + max_policy_lag=self._max_policy_lag, + inflight_count=engine.inflight_count, + ready_count=engine.ready_count, + max_inflight=self._max_inflight, + trained_batches=rollout_id, + num_rollouts=num_rollouts, + hard_boundary=hard_boundary, + ) + for _ in range(slots): + engine.submit( + self._build_async_sample(engine.next_gen_id), + behavior_version=self._policy_versions.rollout_version, + ) engine.poll() - picked = engine.drain_freshest(self.batch_size, max_staleness=stale) - engine.pop_evicted() - if picked is not None: - return picked - if engine.inflight: + batch = engine.pop_next_batch( + train_version=self._policy_versions.train_version, + max_policy_lag=self._max_policy_lag, + ) + if batch is not None: + return batch + if engine.inflight_count: engine.wait_oldest() else: - raise RuntimeError("async rollout buffer underflow with no in-flight generations") + raise RuntimeError("async rollout queue is empty and policy lag admits no new generation") diff --git a/unirl/trainer/async_diffusion.py b/unirl/trainer/async_diffusion.py index 6b41db9f8..2b7e51709 100644 --- a/unirl/trainer/async_diffusion.py +++ b/unirl/trainer/async_diffusion.py @@ -17,17 +17,16 @@ * ``_build_async_sample`` — one data batch → one request ``Sample``. * ``_score_completed`` — reward at reap time, then split into tree-complete groups. Generation overlaps training; reward scoring itself does not. -* ``_advantage_and_train`` — advantage + FlowGRPO optimizer step over the - freshest ``batch_size`` groups; it never calls the reward. +* ``_advantage_and_train`` — advantage + FlowGRPO optimizer step over the next + FIFO rollout batch; it never calls the reward. -Two numeric knobs (identical semantics to AsyncARTrainer): +Async control uses the same optimizer-update clock as AsyncARTrainer: * ``max_inflight`` — must be ``1`` so a reap-time transfer never competes with a queued generation on the rollout workers. - * ``buffer_max_staleness`` — regular rollout-weight syncs a buffered group may - cross. ``0`` (default) never crosses a sync; ``>0`` enables a bounded - policy-lag buffer. + * ``max_policy_lag`` — inclusive train-minus-behavior lag at batch + admission, measured in committed optimizer updates. -``_next_step`` polls (reaps) BEFORE topping up launches, which is what makes the +``_next_rollout_batch`` polls (reaps) BEFORE topping up launches, which is what makes the overlap fast here: reaping a generation pulls its trajectory segment off the rollout slab (the reward's cross-slab localize, an NCCL send issued on the rollout workers), so a generation launched ahead of that send blocks it — @@ -50,8 +49,9 @@ import torch from unirl.distributed.tensor import hydrate -from unirl.rollout.engine.asynchronous import AsyncBatchRolloutEngine, launch_ceiling +from unirl.rollout.engine.asynchronous import AsyncBatchRolloutEngine, RolloutBatch from unirl.train.stack import TrainStepResult +from unirl.trainer.async_policy import PolicyVersionState, launch_slots, next_hard_boundary, unwrap_replicated_int from unirl.trainer.diffusion import DiffusionTrainer from unirl.types.sample import Sample @@ -65,7 +65,7 @@ def __init__( self, *, max_inflight: int = 1, - buffer_max_staleness: Optional[int] = None, + max_policy_lag: int = 0, **diffusion_kwargs: Any, ) -> None: layout = diffusion_kwargs.setdefault("layout", "separate") @@ -86,7 +86,20 @@ def __init__( ) self._max_inflight = max_inflight - self._buffer_max_staleness = buffer_max_staleness + self._max_policy_lag = int(max_policy_lag) + stack_cfg = diffusion_kwargs["stack_cfg"] + self._num_updates_per_batch = int(stack_cfg.get("num_updates_per_batch", 1)) + if self._max_policy_lag < 0: + raise ValueError(f"max_policy_lag must be >= 0, got {self._max_policy_lag}") + if self._num_updates_per_batch < 1: + raise ValueError(f"stack.num_updates_per_batch must be >= 1, got {self._num_updates_per_batch}") + freshness_depth = self._max_policy_lag // self._num_updates_per_batch + 1 + if freshness_depth == 1: + logger.warning( + "async policy-lag settings admit one generation at a time; " + "generation cannot overlap the preceding train batch" + ) + self._policy_versions = PolicyVersionState() def _build_async_sample(self, gen_id: int) -> Sample: """Consume one data batch and build the request Sample for ``gen_id``.""" @@ -114,6 +127,31 @@ def _drain_all(self) -> None: """ self._async_engine.quiesce() + def _sync_rollout_weights(self, *, force: bool = False) -> bool: + """Load the current train weights into an empty rollout engine.""" + + versions = self._policy_versions + if not force and versions.rollout_version == versions.train_version: + return False + self._drain_all() + if self._async_engine.ready_count != 0: + raise RuntimeError( + f"cannot sync rollout weights with completed batches queued: " + f"ready_count={self._async_engine.ready_count}" + ) + target = versions.train_version + self.weight_sync.sync() + self.rollout.set_policy_version(target) + versions.mark_rollout_synced(target) + return True + + def _policy_metrics(self, batch: RolloutBatch) -> dict[str, float]: + versions = self._policy_versions + return { + "async/behavior_version": float(batch.behavior_version), + "async/behavior_lag": float(versions.behavior_lag(batch.behavior_version)), + } + def _advantage_and_train( self, sample: Sample, @@ -121,8 +159,9 @@ def _advantage_and_train( training_progress: float, rollout_id: int, t0: Optional[float] = None, + extra_metrics: Optional[dict[str, float]] = None, ) -> Tuple[TrainStepResult, float]: - """Advantage + optimizer step for a SCORED ``Sample`` (rewards already attached).""" + """Advantage + optimizer updates for a scored ``Sample`` (rewards already attached).""" if t0 is None: t0 = time.perf_counter() part = sample.parts[-1] @@ -135,7 +174,22 @@ def _advantage_and_train( part = part.compute_advantages(normalize=True, use_global_std=self._adv_use_global_std) sample = sample.replace_frontier(part) result = self.stack.train_track(sample.parts[-1], training_progress=float(training_progress)) - self.wandb_logger.log_rollout_step(rollout_id, result, sample, step_time_s=time.perf_counter() - t0) + self._policy_versions.record_optimizer_updates(int(result.optimizer_updates)) + if extra_metrics is not None: + extra_metrics.update( + { + "async/train_version": float(self._policy_versions.train_version), + "async/rollout_lag": float(self._policy_versions.rollout_lag), + "async/optimizer_updates": float(result.optimizer_updates), + } + ) + self.wandb_logger.log_rollout_step( + rollout_id, + result, + sample, + step_time_s=time.perf_counter() - t0, + extra_metrics=extra_metrics, + ) self._reset_transport_buffers() return result, mean_reward @@ -143,65 +197,90 @@ def train( self, *, num_rollouts: int, - weight_sync_interval: int = 1, save_interval: int = 0, save_dir: Optional[str] = None, load_dir: Optional[str] = None, save_mode: str = "auto", ) -> None: - interval = max(1, weight_sync_interval) - stale = self._buffer_max_staleness if self._buffer_max_staleness is not None else 0 - M = self._max_inflight - start_rollout = self.maybe_load_checkpoint(load_dir, num_rollouts=num_rollouts) resumed = bool(load_dir) + train_version = unwrap_replicated_int( + self.backend.get_optimizer_step_count(), + name="backend optimizer step count", + ) + self._policy_versions = PolicyVersionState(train_version=train_version) for _ in range(start_rollout): self.data_source.get_samples(self.batch_size) self._init_wandb( num_rollouts=num_rollouts, extra={ - "max_inflight": M, - "buffer_max_staleness": stale, - "weight_sync_interval": interval, + "max_inflight": self._max_inflight, + "max_policy_lag": self._max_policy_lag, + "num_updates_per_batch": self._num_updates_per_batch, "train_fraction": self._train_fraction, }, ) self._async_engine = AsyncBatchRolloutEngine( self.rollout, - complete=self._score_completed, + process_completion=self._score_completed, + groups_per_batch=self.batch_size, start_gen_id=start_rollout, ) - if resumed and self.weight_sync is not None: - self.weight_sync.sync() + if resumed: + self._sync_rollout_weights(force=True) if self.eval_interval > 0: self.evaluate(start_rollout, sync_weights=False, sleep_after=False) try: for rollout_id in range(start_rollout, num_rollouts): t0 = time.perf_counter() - picked = self._next_step(rollout_id, interval, M, stale, num_rollouts) - sample = Sample.concat(picked) + hard_boundary = next_hard_boundary( + rollout_id, + num_rollouts=num_rollouts, + eval_interval=self.eval_interval, + save_interval=save_interval, + ) + batch = self._next_rollout_batch( + rollout_id, + num_rollouts=num_rollouts, + hard_boundary=hard_boundary, + ) + sample = Sample.concat(batch.groups) training_progress = rollout_id / max(1, num_rollouts - 1) result, mean_reward = self._advantage_and_train( - sample, training_progress=training_progress, rollout_id=rollout_id, t0=t0 + sample, + training_progress=training_progress, + rollout_id=rollout_id, + t0=t0, + extra_metrics=self._policy_metrics(batch), ) self.wandb_logger.log_progress(rollout_id, num_rollouts, result, mean_reward, logger=logger) step = rollout_id + 1 - if self.eval_interval > 0 and step % self.eval_interval == 0: - self._drain_all() + eval_due = self.eval_interval > 0 and step % self.eval_interval == 0 + save_due = save_interval > 0 and (step % save_interval == 0 or step >= num_rollouts) + sync_due = step < num_rollouts and self._policy_versions.rollout_lag > self._max_policy_lag + if eval_due or save_due or sync_due: + if self._async_engine.inflight_count + self._async_engine.ready_count != 0: + raise RuntimeError( + "async sync boundary retained rollout work: " + f"inflight_count={self._async_engine.inflight_count}, " + f"ready_count={self._async_engine.ready_count}" + ) + self._sync_rollout_weights() + + if eval_due: self.evaluate(step, sync_weights=False, sleep_after=False) - if save_interval > 0 and (step % save_interval == 0 or step >= num_rollouts): - self._drain_all() + if save_due: self.maybe_save_checkpoint( - rollout_id, num_rollouts, save_interval=save_interval, save_dir=save_dir, save_mode=save_mode + rollout_id, + num_rollouts, + save_interval=save_interval, + save_dir=save_dir, + save_mode=save_mode, ) - if step % interval == 0 and self.weight_sync is not None: - self._drain_all() - self.weight_sync.sync() - self._async_engine.bump_weight_version() finally: # Cleanup failures must not mask the exception that stopped training. active_exception = sys.exc_info()[0] is not None @@ -214,38 +293,46 @@ def train( finally: self._finish_wandb() - def _next_step( + def _next_rollout_batch( self, rollout_id: int, - interval: int, - M: int, - stale: int, + *, num_rollouts: int, - ) -> List[Sample]: - """Reap completed generations, top up launches, and return the freshest - ``batch_size`` scored group Samples for ``rollout_id`` (blocking on the - oldest in-flight generation if the buffer is short). + hard_boundary: int, + ) -> RolloutBatch: + """Reap, launch, and consume one completion-order FIFO train batch. Polls BEFORE topping up: reaping pulls the trajectory segment off the rollout slab, so it must not queue behind a freshly launched generation, and the post-reap launch is what overlaps this step (module docstring). - - The launch clamp is the load-bearing on-policy guarantee: a generation - launched now is consumed later, so bound how far ahead we launch to - ``stale`` weight-syncs. ``stale=0`` ⇒ never launch into a future - sync-window ⇒ no generation crosses a regular rollout-weight sync. """ engine = self._async_engine while True: - ceiling = launch_ceiling(rollout_id, sync_interval=interval, max_staleness=stale, num_rollouts=num_rollouts) engine.poll() - while engine.next_gen_id < ceiling and engine.inflight < M: - engine.submit(self._build_async_sample(engine.next_gen_id)) - picked = engine.drain_freshest(self.batch_size, max_staleness=stale) - engine.pop_evicted() - if picked is not None: - return picked - if engine.inflight: + slots = launch_slots( + train_version=self._policy_versions.train_version, + rollout_version=self._policy_versions.rollout_version, + num_updates_per_batch=self._num_updates_per_batch, + max_policy_lag=self._max_policy_lag, + inflight_count=engine.inflight_count, + ready_count=engine.ready_count, + max_inflight=self._max_inflight, + trained_batches=rollout_id, + num_rollouts=num_rollouts, + hard_boundary=hard_boundary, + ) + for _ in range(slots): + engine.submit( + self._build_async_sample(engine.next_gen_id), + behavior_version=self._policy_versions.rollout_version, + ) + batch = engine.pop_next_batch( + train_version=self._policy_versions.train_version, + max_policy_lag=self._max_policy_lag, + ) + if batch is not None: + return batch + if engine.inflight_count: engine.wait_oldest() else: - raise RuntimeError("async rollout buffer underflow with no in-flight generations") + raise RuntimeError("async rollout queue is empty and policy lag admits no new generation") diff --git a/unirl/trainer/async_policy.py b/unirl/trainer/async_policy.py new file mode 100644 index 000000000..a3ec83be6 --- /dev/null +++ b/unirl/trainer/async_policy.py @@ -0,0 +1,181 @@ +"""Shared policy-version accounting for disaggregated async trainers. + +Versions in this module count committed optimizer updates, never weight-sync +calls or consumed rollout batches. A rollout batch records the train version +whose weights were resident in the rollout engine when generation started. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class PolicyVersionState: + """Driver-owned train/rollout policy clocks. + + ``train_version`` advances by the number of optimizer steps that actually + committed. ``rollout_version`` is assigned to the exact train snapshot most + recently synced to the rollout engine. + """ + + train_version: int = 0 + rollout_version: int = 0 + + def __post_init__(self) -> None: + self.train_version = int(self.train_version) + self.rollout_version = int(self.rollout_version) + if self.train_version < 0 or self.rollout_version < 0: + raise ValueError("policy versions must be non-negative") + if self.rollout_version > self.train_version: + raise ValueError( + f"rollout_version cannot be ahead of train_version: {self.rollout_version} > {self.train_version}" + ) + + @property + def rollout_lag(self) -> int: + return self.train_version - self.rollout_version + + def record_optimizer_updates(self, committed_updates: int) -> int: + """Advance by optimizer steps that successfully changed train weights.""" + + committed_updates = int(committed_updates) + if committed_updates < 0: + raise ValueError(f"committed optimizer updates must be >= 0, got {committed_updates}") + self.train_version += committed_updates + return self.train_version + + def mark_rollout_synced(self, train_version: int) -> int: + """Record the train version loaded by a successful rollout weight sync.""" + + train_version = int(train_version) + if train_version < self.rollout_version: + raise ValueError(f"synced train version cannot move backwards: {train_version} < {self.rollout_version}") + if train_version > self.train_version: + raise ValueError(f"cannot sync future train version {train_version} > {self.train_version}") + self.rollout_version = train_version + return self.rollout_version + + def behavior_lag(self, behavior_version: int) -> int: + """Optimizer-update lag between train and a batch's behavior policy.""" + + behavior_version = int(behavior_version) + lag = self.train_version - behavior_version + if lag < 0: + raise ValueError( + f"rollout batch has future behavior version {behavior_version} > train version {self.train_version}" + ) + return lag + + +def launch_slots( + *, + train_version: int, + rollout_version: int, + num_updates_per_batch: int, + max_policy_lag: int, + inflight_count: int, + ready_count: int, + max_inflight: int, + trained_batches: int, + num_rollouts: int, + hard_boundary: int, +) -> int: + """How many batch generations may be launched by the single-threaded loop. + + Freshness is measured in committed optimizer updates. The first outstanding + batch would train at the current lag; every additional outstanding batch is + conservatively reserved ``num_updates_per_batch`` future updates. Capacity and + durable-boundary clamps are applied in the same generation-batch unit. + """ + + train = _non_negative("train_version", train_version) + rollout = _non_negative("rollout_version", rollout_version) + updates = _positive("num_updates_per_batch", num_updates_per_batch) + max_lag = _non_negative("max_policy_lag", max_policy_lag) + active = _non_negative("inflight_count", inflight_count) + queued = _non_negative("ready_count", ready_count) + max_active = _positive("max_inflight", max_inflight) + trained = _non_negative("trained_batches", trained_batches) + total = _non_negative("num_rollouts", num_rollouts) + boundary = _non_negative("hard_boundary", hard_boundary) + if rollout > train: + raise ValueError(f"rollout_version cannot be ahead of train_version: {rollout} > {train}") + if active > max_active: + raise ValueError(f"inflight_count={active} exceeds max_inflight={max_active}") + if boundary < trained: + raise ValueError(f"hard_boundary={boundary} is behind trained_batches={trained}") + if trained >= total: + return 0 + + current_lag = train - rollout + if current_lag > max_lag: + return 0 + + freshness_slots = (max_lag - current_lag) // updates + 1 + allowed_outstanding = min( + freshness_slots, + total - trained, + boundary - trained, + ) + outstanding = active + queued + return max( + 0, + min( + max_active - active, + allowed_outstanding - outstanding, + ), + ) + + +def next_hard_boundary( + trained_batches: int, + *, + num_rollouts: int, + eval_interval: int = 0, + save_interval: int = 0, +) -> int: + """Nearest eval/checkpoint/final boundary for launch admission.""" + + trained = _non_negative("trained_batches", trained_batches) + total = _non_negative("num_rollouts", num_rollouts) + boundary = total + for interval in (int(eval_interval), int(save_interval)): + if interval > 0 and trained < total: + boundary = min(boundary, ((trained // interval) + 1) * interval) + return boundary + + +def unwrap_replicated_int(value: object, *, name: str) -> int: + """Normalize a BROADCAST return and verify all worker replicas agree.""" + + if isinstance(value, (list, tuple)): + if not value: + raise ValueError(f"{name} returned no worker values") + first = int(value[0]) + if any(int(item) != first for item in value[1:]): + raise RuntimeError(f"{name} disagrees across workers: {value!r}") + return first + return int(value) + + +def _non_negative(name: str, value: int) -> int: + value = int(value) + if value < 0: + raise ValueError(f"{name} must be >= 0, got {value}") + return value + + +def _positive(name: str, value: int) -> int: + value = int(value) + if value < 1: + raise ValueError(f"{name} must be >= 1, got {value}") + return value + + +__all__ = [ + "PolicyVersionState", + "launch_slots", + "next_hard_boundary", + "unwrap_replicated_int", +] From 314f8dd317afcb9268cdc3dec08018926da98b3f Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Mon, 3 Aug 2026 23:02:46 +0800 Subject: [PATCH 2/9] refactor(async): remove redundant numeric casts Keep policy-version values in their native integer types so the control path stays explicit and readable. --- unirl/rollout/engine/asynchronous.py | 16 ++++++++-------- unirl/rollout/engine/synchronous.py | 1 - unirl/train/backend/base_backend.py | 2 +- unirl/train/stack/base.py | 12 ++++++------ unirl/train/unified_model_stack.py | 2 +- unirl/train_async_ar.py | 2 +- unirl/train_async_diffusion.py | 2 +- unirl/trainer/async_ar.py | 16 ++++++++-------- unirl/trainer/async_diffusion.py | 16 ++++++++-------- unirl/trainer/async_policy.py | 19 ++++++++----------- 10 files changed, 42 insertions(+), 46 deletions(-) diff --git a/unirl/rollout/engine/asynchronous.py b/unirl/rollout/engine/asynchronous.py index 5a6d1e273..82799ef49 100644 --- a/unirl/rollout/engine/asynchronous.py +++ b/unirl/rollout/engine/asynchronous.py @@ -138,13 +138,13 @@ def pop_next( if not self._items: return None item = self._items[0] - lag = int(train_version) - item.behavior_version + lag = train_version - item.behavior_version if lag < 0: raise RuntimeError( f"generation {item.gen_id} has future behavior version " f"{item.behavior_version} > train version {train_version}" ) - if lag > int(max_policy_lag): + if lag > max_policy_lag: raise RuntimeError(f"generation {item.gen_id} exceeded policy lag budget: lag={lag} > max={max_policy_lag}") return self._items.popleft() @@ -184,7 +184,7 @@ def __len__(self) -> int: def launch(self, sample: Any, *, behavior_version: int) -> int: gen_id = self._next_gen_id pending = self._rollout.launch_nowait("generate", sample) - self._jobs.append(_InflightJob(gen_id, int(behavior_version), pending)) + self._jobs.append(_InflightJob(gen_id, behavior_version, pending)) self._next_gen_id += 1 return gen_id @@ -262,10 +262,10 @@ def __init__( groups_per_batch: int, start_gen_id: int = 0, ) -> None: - if int(groups_per_batch) < 1: + if groups_per_batch < 1: raise ValueError(f"groups_per_batch must be >= 1, got {groups_per_batch}") self._process_completion = process_completion - self._groups_per_batch = int(groups_per_batch) + self._groups_per_batch = groups_per_batch self._pool = InflightPool(rollout, start_gen_id=start_gen_id) self._ready = RolloutBatchQueue() @@ -316,9 +316,9 @@ def _on_complete(self, gen_id: int, behavior_version: int, completed: "Sample") ) self._ready.put( RolloutBatch( - groups=list(groups), - behavior_version=int(behavior_version), - gen_id=int(gen_id), + groups=groups, + behavior_version=behavior_version, + gen_id=gen_id, ) ) diff --git a/unirl/rollout/engine/synchronous.py b/unirl/rollout/engine/synchronous.py index 74496f41e..43d6ea1eb 100644 --- a/unirl/rollout/engine/synchronous.py +++ b/unirl/rollout/engine/synchronous.py @@ -166,7 +166,6 @@ def set_policy_version(self, train_version: int) -> None: trainer calls this only after a full weight publication succeeds. """ - train_version = int(train_version) if train_version < 0: raise ValueError(f"train_version must be >= 0, got {train_version}") self._weight_version = train_version diff --git a/unirl/train/backend/base_backend.py b/unirl/train/backend/base_backend.py index 1d560bee3..e1476fb12 100644 --- a/unirl/train/backend/base_backend.py +++ b/unirl/train/backend/base_backend.py @@ -374,7 +374,7 @@ def on_rollout_end(self) -> None: def get_optimizer_step_count(self) -> int: """Return the authoritative number of committed optimizer updates.""" - return int(self._optimizer_step_count) + return self._optimizer_step_count @property def rollout_adapter_name(self) -> str: diff --git a/unirl/train/stack/base.py b/unirl/train/stack/base.py index 058d61f0e..065826e99 100644 --- a/unirl/train/stack/base.py +++ b/unirl/train/stack/base.py @@ -307,7 +307,7 @@ def _run_update( has_backward=has_backward, micros=micro_results, metrics=aggregated_metrics, - optimizer_updates=int(has_backward and math.isfinite(grad_norm)), + optimizer_updates=1 if has_backward and math.isfinite(grad_norm) else 0, ) def on_rollout_end(self) -> None: @@ -518,11 +518,11 @@ def _run_updates( aggregated = _aggregate_update_results(results) per_update = tuple( { - **dict(r.metrics), - "loss": float(r.loss), - "grad_norm": float(r.grad_norm), - "lr": float(r.lr), - "optimizer_updates": float(r.optimizer_updates), + **r.metrics, + "loss": r.loss, + "grad_norm": r.grad_norm, + "lr": r.lr, + "optimizer_updates": r.optimizer_updates, } for r in results ) diff --git a/unirl/train/unified_model_stack.py b/unirl/train/unified_model_stack.py index 24e0e9463..0592eb77b 100644 --- a/unirl/train/unified_model_stack.py +++ b/unirl/train/unified_model_stack.py @@ -252,7 +252,7 @@ def _train_one_step( logger.warning("UnifiedModelTrainStack._train_one_step: no algorithm reported backward; skipping step.") lr = self._current_lr() - optimizer_updates = int(any_backward and math.isfinite(grad_norm)) + optimizer_updates = 1 if any_backward and math.isfinite(grad_norm) else 0 for name, r in list(results.items()): results[name] = TrainStepResult( loss=r.loss, diff --git a/unirl/train_async_ar.py b/unirl/train_async_ar.py index 9480f3dfd..ad4f90b17 100755 --- a/unirl/train_async_ar.py +++ b/unirl/train_async_ar.py @@ -55,7 +55,7 @@ def main(cfg: DictConfig) -> None: eval_temperature=float(cfg.get("eval_temperature", 1.0)), train_fraction=float(cfg.get("train_fraction", 0.5)), max_inflight=int(cfg.get("max_inflight", 1)), - max_policy_lag=int(cfg.get("max_policy_lag", 0)), + max_policy_lag=cfg.get("max_policy_lag", 0), ) trainer.train( num_rollouts=int(cfg.get("num_rollouts", 100)), diff --git a/unirl/train_async_diffusion.py b/unirl/train_async_diffusion.py index 37b8879ad..6ec0f3848 100755 --- a/unirl/train_async_diffusion.py +++ b/unirl/train_async_diffusion.py @@ -56,7 +56,7 @@ def main(cfg: DictConfig) -> None: eval_rewards_cfg=cfg.get("eval_rewards"), task_config=cfg.get("task_config"), max_inflight=int(cfg.get("max_inflight", 1)), - max_policy_lag=int(cfg.get("max_policy_lag", 0)), + max_policy_lag=cfg.get("max_policy_lag", 0), ) trainer.train( num_rollouts=cfg.get("num_rollouts", 100), diff --git a/unirl/trainer/async_ar.py b/unirl/trainer/async_ar.py index 1f3029e22..ffedbbebe 100644 --- a/unirl/trainer/async_ar.py +++ b/unirl/trainer/async_ar.py @@ -138,8 +138,8 @@ def __init__( self._train_fraction = float(train_fraction) self._max_inflight = max(1, int(max_inflight)) - self._max_policy_lag = int(max_policy_lag) - self._num_updates_per_batch = int(stack_cfg.get("num_updates_per_batch", 1)) + self._max_policy_lag = max_policy_lag + self._num_updates_per_batch = stack_cfg.get("num_updates_per_batch", 1) if self._max_policy_lag < 0: raise ValueError(f"max_policy_lag must be >= 0, got {self._max_policy_lag}") if self._num_updates_per_batch < 1: @@ -286,8 +286,8 @@ def _sync_rollout_weights(self, *, force: bool = False) -> bool: def _policy_metrics(self, batch: RolloutBatch) -> Dict[str, float]: versions = self._policy_versions return { - "async/behavior_version": float(batch.behavior_version), - "async/behavior_lag": float(versions.behavior_lag(batch.behavior_version)), + "async/behavior_version": batch.behavior_version, + "async/behavior_lag": versions.behavior_lag(batch.behavior_version), } def _advantage_and_train( @@ -317,13 +317,13 @@ def _advantage_and_train( if self.balance_shards: train_part = part.balance_shards(self._train_devices) result = self.stack.train_track(train_part, training_progress=float(training_progress)) - self._policy_versions.record_optimizer_updates(int(result.optimizer_updates)) + self._policy_versions.record_optimizer_updates(result.optimizer_updates) if extra_metrics is not None: extra_metrics.update( { - "async/train_version": float(self._policy_versions.train_version), - "async/rollout_lag": float(self._policy_versions.rollout_lag), - "async/optimizer_updates": float(result.optimizer_updates), + "async/train_version": self._policy_versions.train_version, + "async/rollout_lag": self._policy_versions.rollout_lag, + "async/optimizer_updates": result.optimizer_updates, } ) self.wandb_logger.log_rollout_step( diff --git a/unirl/trainer/async_diffusion.py b/unirl/trainer/async_diffusion.py index 2b7e51709..b8051bd19 100644 --- a/unirl/trainer/async_diffusion.py +++ b/unirl/trainer/async_diffusion.py @@ -86,9 +86,9 @@ def __init__( ) self._max_inflight = max_inflight - self._max_policy_lag = int(max_policy_lag) + self._max_policy_lag = max_policy_lag stack_cfg = diffusion_kwargs["stack_cfg"] - self._num_updates_per_batch = int(stack_cfg.get("num_updates_per_batch", 1)) + self._num_updates_per_batch = stack_cfg.get("num_updates_per_batch", 1) if self._max_policy_lag < 0: raise ValueError(f"max_policy_lag must be >= 0, got {self._max_policy_lag}") if self._num_updates_per_batch < 1: @@ -148,8 +148,8 @@ def _sync_rollout_weights(self, *, force: bool = False) -> bool: def _policy_metrics(self, batch: RolloutBatch) -> dict[str, float]: versions = self._policy_versions return { - "async/behavior_version": float(batch.behavior_version), - "async/behavior_lag": float(versions.behavior_lag(batch.behavior_version)), + "async/behavior_version": batch.behavior_version, + "async/behavior_lag": versions.behavior_lag(batch.behavior_version), } def _advantage_and_train( @@ -174,13 +174,13 @@ def _advantage_and_train( part = part.compute_advantages(normalize=True, use_global_std=self._adv_use_global_std) sample = sample.replace_frontier(part) result = self.stack.train_track(sample.parts[-1], training_progress=float(training_progress)) - self._policy_versions.record_optimizer_updates(int(result.optimizer_updates)) + self._policy_versions.record_optimizer_updates(result.optimizer_updates) if extra_metrics is not None: extra_metrics.update( { - "async/train_version": float(self._policy_versions.train_version), - "async/rollout_lag": float(self._policy_versions.rollout_lag), - "async/optimizer_updates": float(result.optimizer_updates), + "async/train_version": self._policy_versions.train_version, + "async/rollout_lag": self._policy_versions.rollout_lag, + "async/optimizer_updates": result.optimizer_updates, } ) self.wandb_logger.log_rollout_step( diff --git a/unirl/trainer/async_policy.py b/unirl/trainer/async_policy.py index a3ec83be6..9783d9444 100644 --- a/unirl/trainer/async_policy.py +++ b/unirl/trainer/async_policy.py @@ -23,8 +23,6 @@ class PolicyVersionState: rollout_version: int = 0 def __post_init__(self) -> None: - self.train_version = int(self.train_version) - self.rollout_version = int(self.rollout_version) if self.train_version < 0 or self.rollout_version < 0: raise ValueError("policy versions must be non-negative") if self.rollout_version > self.train_version: @@ -39,7 +37,6 @@ def rollout_lag(self) -> int: def record_optimizer_updates(self, committed_updates: int) -> int: """Advance by optimizer steps that successfully changed train weights.""" - committed_updates = int(committed_updates) if committed_updates < 0: raise ValueError(f"committed optimizer updates must be >= 0, got {committed_updates}") self.train_version += committed_updates @@ -48,7 +45,6 @@ def record_optimizer_updates(self, committed_updates: int) -> int: def mark_rollout_synced(self, train_version: int) -> int: """Record the train version loaded by a successful rollout weight sync.""" - train_version = int(train_version) if train_version < self.rollout_version: raise ValueError(f"synced train version cannot move backwards: {train_version} < {self.rollout_version}") if train_version > self.train_version: @@ -59,7 +55,6 @@ def mark_rollout_synced(self, train_version: int) -> int: def behavior_lag(self, behavior_version: int) -> int: """Optimizer-update lag between train and a batch's behavior policy.""" - behavior_version = int(behavior_version) lag = self.train_version - behavior_version if lag < 0: raise ValueError( @@ -140,7 +135,7 @@ def next_hard_boundary( trained = _non_negative("trained_batches", trained_batches) total = _non_negative("num_rollouts", num_rollouts) boundary = total - for interval in (int(eval_interval), int(save_interval)): + for interval in (eval_interval, save_interval): if interval > 0 and trained < total: boundary = min(boundary, ((trained // interval) + 1) * interval) return boundary @@ -152,22 +147,24 @@ def unwrap_replicated_int(value: object, *, name: str) -> int: if isinstance(value, (list, tuple)): if not value: raise ValueError(f"{name} returned no worker values") - first = int(value[0]) - if any(int(item) != first for item in value[1:]): + if any(not isinstance(item, int) for item in value): + raise TypeError(f"{name} returned non-integer worker values: {value!r}") + first = value[0] + if any(item != first for item in value[1:]): raise RuntimeError(f"{name} disagrees across workers: {value!r}") return first - return int(value) + if not isinstance(value, int): + raise TypeError(f"{name} returned {type(value).__name__}, expected int") + return value def _non_negative(name: str, value: int) -> int: - value = int(value) if value < 0: raise ValueError(f"{name} must be >= 0, got {value}") return value def _positive(name: str, value: int) -> int: - value = int(value) if value < 1: raise ValueError(f"{name} must be >= 1, got {value}") return value From 1fd3363c370b1b4b1c3b27c54b8c278c7b9d41f0 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Mon, 3 Aug 2026 23:22:28 +0800 Subject: [PATCH 3/9] chore(test): keep async validation local Remove repository-level async test scaffolding from the pull request while retaining the files in the local worktree for validation. --- pyproject.toml | 3 - tests/rollout/test_async_batch_engine.py | 74 ---------- tests/trainer/test_async_policy.py | 138 ------------------ .../trainer/test_async_trainer_boundaries.py | 57 -------- 4 files changed, 272 deletions(-) delete mode 100644 tests/rollout/test_async_batch_engine.py delete mode 100644 tests/trainer/test_async_policy.py delete mode 100644 tests/trainer/test_async_trainer_boundaries.py diff --git a/pyproject.toml b/pyproject.toml index fbed401a3..97254ada7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -196,6 +196,3 @@ known-first-party = ["unirl"] [tool.ruff.format] quote-style = "double" - -[tool.pytest.ini_options] -testpaths = ["tests"] diff --git a/tests/rollout/test_async_batch_engine.py b/tests/rollout/test_async_batch_engine.py deleted file mode 100644 index eb4520669..000000000 --- a/tests/rollout/test_async_batch_engine.py +++ /dev/null @@ -1,74 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass - -import pytest - -from unirl.rollout.engine.asynchronous import AsyncBatchRolloutEngine - - -@dataclass -class _Pending: - value: str - - def ready(self) -> bool: - return True - - def result(self) -> str: - return self.value - - def wait(self) -> None: - return None - - -class _RolloutHandle: - def __init__(self) -> None: - self.launched: list[str] = [] - - def launch_nowait(self, method: str, sample: str) -> _Pending: - assert method == "generate" - self.launched.append(sample) - return _Pending(sample) - - -def test_batch_engine_stamps_train_version_and_consumes_generation_fifo() -> None: - rollout = _RolloutHandle() - engine = AsyncBatchRolloutEngine( - rollout, - process_completion=lambda _gen_id, completed: [f"{completed}/0", f"{completed}/1"], - groups_per_batch=2, - ) - - engine.submit("batch-0", behavior_version=0) - engine.submit("batch-1", behavior_version=4) - assert engine.poll() == 2 - - first = engine.pop_next_batch(train_version=4, max_policy_lag=4) - second = engine.pop_next_batch(train_version=4, max_policy_lag=4) - assert first is not None and first.groups == ["batch-0/0", "batch-0/1"] - assert first.behavior_version == 0 - assert second is not None and second.groups == ["batch-1/0", "batch-1/1"] - assert second.behavior_version == 4 - - -def test_batch_engine_rejects_non_atomic_generation() -> None: - engine = AsyncBatchRolloutEngine( - _RolloutHandle(), - process_completion=lambda _gen_id, completed: [completed], - groups_per_batch=2, - ) - engine.submit("short", behavior_version=0) - with pytest.raises(RuntimeError, match="expected groups_per_batch=2"): - engine.poll() - - -def test_batch_engine_fails_closed_on_stale_ready_batch() -> None: - engine = AsyncBatchRolloutEngine( - _RolloutHandle(), - process_completion=lambda _gen_id, completed: [completed], - groups_per_batch=1, - ) - engine.submit("old", behavior_version=0) - engine.poll() - with pytest.raises(RuntimeError, match="exceeded policy lag budget"): - engine.pop_next_batch(train_version=2, max_policy_lag=1) diff --git a/tests/trainer/test_async_policy.py b/tests/trainer/test_async_policy.py deleted file mode 100644 index e4571a480..000000000 --- a/tests/trainer/test_async_policy.py +++ /dev/null @@ -1,138 +0,0 @@ -from __future__ import annotations - -import pytest - -from unirl.train.stack.base import TrainStepResult, _aggregate_update_results -from unirl.trainer.async_policy import PolicyVersionState, launch_slots, next_hard_boundary, unwrap_replicated_int - - -def _slots( - *, - train: int = 0, - rollout: int = 0, - updates: int = 4, - max_lag: int = 0, - inflight: int = 0, - ready: int = 0, - max_inflight: int = 8, - trained: int = 0, - total: int = 20, - boundary: int = 20, -) -> int: - return launch_slots( - train_version=train, - rollout_version=rollout, - num_updates_per_batch=updates, - max_policy_lag=max_lag, - inflight_count=inflight, - ready_count=ready, - max_inflight=max_inflight, - trained_batches=trained, - num_rollouts=total, - hard_boundary=boundary, - ) - - -def test_launch_slots_uses_optimizer_update_lag() -> None: - assert _slots(max_lag=0) == 1 - assert _slots(max_lag=4) == 2 - assert _slots(train=4, max_lag=4) == 1 - assert _slots(train=8, max_lag=4) == 0 - - -def test_launch_slots_subtracts_all_outstanding_batches() -> None: - assert _slots(max_lag=8, inflight=1, ready=1) == 1 - assert _slots(max_lag=8, inflight=2, ready=1) == 0 - - -def test_launch_slots_respects_concurrency_target_and_hard_boundary() -> None: - assert _slots(max_lag=100, max_inflight=2) == 2 - assert _slots(max_lag=100, inflight=1, max_inflight=2) == 1 - assert _slots(max_lag=100, trained=3, boundary=4) == 1 - assert _slots(max_lag=100, trained=4, boundary=4) == 0 - - -@pytest.mark.parametrize( - ("updates", "max_lag", "expected_batches"), - [ - (4, 0, 1), - (4, 12, 4), - (2, 22, 12), - ], -) -def test_single_thread_cycle_exhausts_queue_before_sync( - updates: int, - max_lag: int, - expected_batches: int, -) -> None: - train = 0 - rollout = 0 - ready = 0 - consumed = 0 - while True: - ready += _slots( - train=train, - rollout=rollout, - updates=updates, - max_lag=max_lag, - ready=ready, - max_inflight=64, - trained=consumed, - total=100, - boundary=100, - ) - if ready == 0: - break - ready -= 1 - consumed += 1 - train += updates - - assert consumed == expected_batches - assert ready == 0 - - -def test_launch_slots_rejects_future_rollout_version() -> None: - with pytest.raises(ValueError, match="ahead of train_version"): - _slots(train=2, rollout=3) - - -def test_policy_versions_track_train_updates_and_rollout_sync_separately() -> None: - state = PolicyVersionState() - state.record_optimizer_updates(4) - assert state.train_version == 4 - assert state.rollout_version == 0 - assert state.rollout_lag == 4 - - state.mark_rollout_synced(4) - assert state.rollout_version == 4 - state.mark_rollout_synced(4) - assert state.train_version == 4 - assert state.rollout_version == 4 - - -def test_policy_versions_fail_closed_on_future_batch_or_sync() -> None: - state = PolicyVersionState(train_version=3, rollout_version=2) - with pytest.raises(ValueError, match="future behavior version"): - state.behavior_lag(4) - with pytest.raises(ValueError, match="future train version"): - state.mark_rollout_synced(4) - - -def test_next_hard_boundary_uses_nearest_eval_save_or_final() -> None: - assert next_hard_boundary(0, num_rollouts=20, eval_interval=4, save_interval=3) == 3 - assert next_hard_boundary(3, num_rollouts=20, eval_interval=4, save_interval=3) == 4 - assert next_hard_boundary(19, num_rollouts=20, eval_interval=4, save_interval=3) == 20 - - -def test_train_step_result_aggregates_only_committed_optimizer_updates() -> None: - committed = TrainStepResult(1.0, 1.0, 1e-6, True, [], {}, optimizer_updates=1) - skipped = TrainStepResult(float("nan"), float("nan"), 1e-6, True, [], {}, optimizer_updates=0) - result = _aggregate_update_results([committed, skipped]) - assert result.has_backward is True - assert result.optimizer_updates == 1 - - -def test_replicated_optimizer_count_must_agree() -> None: - assert unwrap_replicated_int([7, 7], name="optimizer count") == 7 - with pytest.raises(RuntimeError, match="disagrees across workers"): - unwrap_replicated_int([7, 8], name="optimizer count") diff --git a/tests/trainer/test_async_trainer_boundaries.py b/tests/trainer/test_async_trainer_boundaries.py deleted file mode 100644 index abe081eff..000000000 --- a/tests/trainer/test_async_trainer_boundaries.py +++ /dev/null @@ -1,57 +0,0 @@ -from __future__ import annotations - -from unirl.trainer.async_ar import AsyncARTrainer -from unirl.trainer.async_policy import PolicyVersionState - - -class _Engine: - def __init__(self) -> None: - self.ready_count = 0 - self.quiesces = 0 - - def quiesce(self) -> None: - self.quiesces += 1 - - -class _WeightSync: - def __init__(self) -> None: - self.calls = 0 - - def sync(self) -> None: - self.calls += 1 - - -class _Rollout: - def __init__(self) -> None: - self.versions: list[int] = [] - - def set_policy_version(self, train_version: int) -> None: - self.versions.append(train_version) - - -def test_async_ar_baseline_eval_syncs_once_even_when_version_is_zero() -> None: - trainer = object.__new__(AsyncARTrainer) - trainer._policy_versions = PolicyVersionState() - trainer._rollout_initialized = False - trainer._async_engine = _Engine() - trainer.weight_sync = _WeightSync() - trainer.rollout = _Rollout() - - trainer._prepare_rollout(sync_weights=True) - trainer._prepare_rollout(sync_weights=True) - - assert trainer.weight_sync.calls == 1 - assert trainer.rollout.versions == [0] - - -def test_async_ar_sync_assigns_exact_train_version() -> None: - trainer = object.__new__(AsyncARTrainer) - trainer._policy_versions = PolicyVersionState(train_version=7, rollout_version=0) - trainer._rollout_initialized = False - trainer._async_engine = _Engine() - trainer.weight_sync = _WeightSync() - trainer.rollout = _Rollout() - - assert trainer._sync_rollout_weights() is True - assert trainer._policy_versions.rollout_version == 7 - assert trainer.rollout.versions == [7] From dcf2e93a03cd83a4c530523dee1bf9d34f6846c5 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 4 Aug 2026 11:57:55 +0800 Subject: [PATCH 4/9] refactor(async): centralize batch control state Share policy versions, launch admission, sync boundaries, and metrics so AR and diffusion keep their distinct loop order without duplicating control logic. --- unirl/trainer/async_ar.py | 91 ++++----------- unirl/trainer/async_diffusion.py | 84 +++----------- unirl/trainer/async_policy.py | 193 ++++++++++++------------------- 3 files changed, 112 insertions(+), 256 deletions(-) diff --git a/unirl/trainer/async_ar.py b/unirl/trainer/async_ar.py index ffedbbebe..2af04dd67 100644 --- a/unirl/trainer/async_ar.py +++ b/unirl/trainer/async_ar.py @@ -47,7 +47,7 @@ from unirl.rollout.engine.asynchronous import AsyncBatchRolloutEngine, RolloutBatch from unirl.train.stack import TrainStepResult from unirl.trainer.ar import ARTrainer -from unirl.trainer.async_policy import PolicyVersionState, launch_slots, next_hard_boundary, unwrap_replicated_int +from unirl.trainer.async_policy import AsyncBatchControl, next_hard_boundary, unwrap_replicated_int from unirl.trainer.base import BaseTrainer, build_sampling_dict from unirl.types.sample import Sample from unirl.types.sampling import BaseSamplingParams, total_samples_per_prompt @@ -138,26 +138,21 @@ def __init__( self._train_fraction = float(train_fraction) self._max_inflight = max(1, int(max_inflight)) - self._max_policy_lag = max_policy_lag - self._num_updates_per_batch = stack_cfg.get("num_updates_per_batch", 1) - if self._max_policy_lag < 0: - raise ValueError(f"max_policy_lag must be >= 0, got {self._max_policy_lag}") - if self._num_updates_per_batch < 1: - raise ValueError(f"stack.num_updates_per_batch must be >= 1, got {self._num_updates_per_batch}") - freshness_depth = self._max_policy_lag // self._num_updates_per_batch + 1 - if self._max_inflight > freshness_depth: + self._control = AsyncBatchControl( + max_policy_lag=max_policy_lag, + num_updates_per_batch=stack_cfg.get("num_updates_per_batch", 1), + ) + if self._max_inflight > self._control.freshness_depth: logger.warning( "max_inflight=%d exceeds the policy-lag admission depth %d; the extra concurrency cannot be used", self._max_inflight, - freshness_depth, + self._control.freshness_depth, ) - if freshness_depth == 1: + if self._control.freshness_depth == 1: logger.warning( "async policy-lag settings admit one generation at a time; " "generation cannot overlap the preceding train batch" ) - self._policy_versions = PolicyVersionState() - self._rollout_initialized = False self._train_devices = int(round(self.num_devices * self._train_fraction)) if self._train_devices <= 0 or self._train_devices >= self.num_devices: raise ValueError( @@ -203,7 +198,7 @@ def __init__( def _prepare_rollout(self, *, sync_weights: bool) -> bool: """Sync a resident separate-slab engine without colocate handoffs.""" if sync_weights: - self._sync_rollout_weights() + self._control.sync_rollout(self._async_engine, self.rollout, self.weight_sync) return False def _finish_rollout(self, *, train_state_offloaded: bool) -> None: @@ -264,32 +259,6 @@ def _drain_all(self) -> None: """ self._async_engine.quiesce() - def _sync_rollout_weights(self, *, force: bool = False) -> bool: - """Load the current train weights into an empty rollout engine.""" - - versions = self._policy_versions - if not force and self._rollout_initialized and versions.rollout_version == versions.train_version: - return False - self._drain_all() - if self._async_engine.ready_count != 0: - raise RuntimeError( - f"cannot sync rollout weights with completed batches queued: " - f"ready_count={self._async_engine.ready_count}" - ) - target = versions.train_version - self.weight_sync.sync() - self.rollout.set_policy_version(target) - versions.mark_rollout_synced(target) - self._rollout_initialized = True - return True - - def _policy_metrics(self, batch: RolloutBatch) -> Dict[str, float]: - versions = self._policy_versions - return { - "async/behavior_version": batch.behavior_version, - "async/behavior_lag": versions.behavior_lag(batch.behavior_version), - } - def _advantage_and_train( self, sample: Sample, @@ -317,15 +286,9 @@ def _advantage_and_train( if self.balance_shards: train_part = part.balance_shards(self._train_devices) result = self.stack.train_track(train_part, training_progress=float(training_progress)) - self._policy_versions.record_optimizer_updates(result.optimizer_updates) + self._control.record_optimizer_updates(result.optimizer_updates) if extra_metrics is not None: - extra_metrics.update( - { - "async/train_version": self._policy_versions.train_version, - "async/rollout_lag": self._policy_versions.rollout_lag, - "async/optimizer_updates": result.optimizer_updates, - } - ) + extra_metrics.update(self._control.train_metrics(result.optimizer_updates)) self.wandb_logger.log_rollout_step( rollout_id, result, @@ -352,7 +315,7 @@ def train( self.backend.get_optimizer_step_count(), name="backend optimizer step count", ) - self._policy_versions = PolicyVersionState(train_version=train_version) + self._control.restore(train_version) for _ in range(start_rollout): self.data_source.get_samples(self.batch_size) self._init_wandb( @@ -360,8 +323,8 @@ def train( extra={ "adv_normalization_scope": self.adv_normalization_scope, "max_inflight": self._max_inflight, - "max_policy_lag": self._max_policy_lag, - "num_updates_per_batch": self._num_updates_per_batch, + "max_policy_lag": self._control.max_policy_lag, + "num_updates_per_batch": self._control.num_updates_per_batch, }, ) @@ -373,7 +336,7 @@ def train( ) if resumed: - self._sync_rollout_weights(force=True) + self._control.sync_rollout(self._async_engine, self.rollout, self.weight_sync) if self.eval_interval > 0: self.evaluate(rollout_id=-1) @@ -398,22 +361,16 @@ def train( training_progress=training_progress, rollout_id=rollout_id, t0=t0, - extra_metrics=self._policy_metrics(batch), + extra_metrics=self._control.behavior_metrics(batch.behavior_version), ) self.wandb_logger.log_progress(rollout_id, num_rollouts, result, mean_reward, logger=logger) step = rollout_id + 1 eval_due = self.eval_interval > 0 and step % self.eval_interval == 0 save_due = save_interval > 0 and (step % save_interval == 0 or step >= num_rollouts) - sync_due = step < num_rollouts and self._policy_versions.rollout_lag > self._max_policy_lag + sync_due = step < num_rollouts and self._control.rollout_lag > self._control.max_policy_lag if eval_due or save_due or sync_due: - if self._async_engine.inflight_count + self._async_engine.ready_count != 0: - raise RuntimeError( - "async sync boundary retained rollout work: " - f"inflight_count={self._async_engine.inflight_count}, " - f"ready_count={self._async_engine.ready_count}" - ) - self._sync_rollout_weights() + self._control.sync_rollout(self._async_engine, self.rollout, self.weight_sync) if eval_due: self.evaluate(rollout_id=rollout_id) @@ -447,11 +404,7 @@ def _next_rollout_batch( engine = self._async_engine while True: - slots = launch_slots( - train_version=self._policy_versions.train_version, - rollout_version=self._policy_versions.rollout_version, - num_updates_per_batch=self._num_updates_per_batch, - max_policy_lag=self._max_policy_lag, + slots = self._control.launch_slots( inflight_count=engine.inflight_count, ready_count=engine.ready_count, max_inflight=self._max_inflight, @@ -462,12 +415,12 @@ def _next_rollout_batch( for _ in range(slots): engine.submit( self._build_async_sample(engine.next_gen_id), - behavior_version=self._policy_versions.rollout_version, + behavior_version=self._control.rollout_version, ) engine.poll() batch = engine.pop_next_batch( - train_version=self._policy_versions.train_version, - max_policy_lag=self._max_policy_lag, + train_version=self._control.train_version, + max_policy_lag=self._control.max_policy_lag, ) if batch is not None: return batch diff --git a/unirl/trainer/async_diffusion.py b/unirl/trainer/async_diffusion.py index b8051bd19..912c8cf4b 100644 --- a/unirl/trainer/async_diffusion.py +++ b/unirl/trainer/async_diffusion.py @@ -51,7 +51,7 @@ from unirl.distributed.tensor import hydrate from unirl.rollout.engine.asynchronous import AsyncBatchRolloutEngine, RolloutBatch from unirl.train.stack import TrainStepResult -from unirl.trainer.async_policy import PolicyVersionState, launch_slots, next_hard_boundary, unwrap_replicated_int +from unirl.trainer.async_policy import AsyncBatchControl, next_hard_boundary, unwrap_replicated_int from unirl.trainer.diffusion import DiffusionTrainer from unirl.types.sample import Sample @@ -86,20 +86,15 @@ def __init__( ) self._max_inflight = max_inflight - self._max_policy_lag = max_policy_lag - stack_cfg = diffusion_kwargs["stack_cfg"] - self._num_updates_per_batch = stack_cfg.get("num_updates_per_batch", 1) - if self._max_policy_lag < 0: - raise ValueError(f"max_policy_lag must be >= 0, got {self._max_policy_lag}") - if self._num_updates_per_batch < 1: - raise ValueError(f"stack.num_updates_per_batch must be >= 1, got {self._num_updates_per_batch}") - freshness_depth = self._max_policy_lag // self._num_updates_per_batch + 1 - if freshness_depth == 1: + self._control = AsyncBatchControl( + max_policy_lag=max_policy_lag, + num_updates_per_batch=diffusion_kwargs["stack_cfg"].get("num_updates_per_batch", 1), + ) + if self._control.freshness_depth == 1: logger.warning( "async policy-lag settings admit one generation at a time; " "generation cannot overlap the preceding train batch" ) - self._policy_versions = PolicyVersionState() def _build_async_sample(self, gen_id: int) -> Sample: """Consume one data batch and build the request Sample for ``gen_id``.""" @@ -127,31 +122,6 @@ def _drain_all(self) -> None: """ self._async_engine.quiesce() - def _sync_rollout_weights(self, *, force: bool = False) -> bool: - """Load the current train weights into an empty rollout engine.""" - - versions = self._policy_versions - if not force and versions.rollout_version == versions.train_version: - return False - self._drain_all() - if self._async_engine.ready_count != 0: - raise RuntimeError( - f"cannot sync rollout weights with completed batches queued: " - f"ready_count={self._async_engine.ready_count}" - ) - target = versions.train_version - self.weight_sync.sync() - self.rollout.set_policy_version(target) - versions.mark_rollout_synced(target) - return True - - def _policy_metrics(self, batch: RolloutBatch) -> dict[str, float]: - versions = self._policy_versions - return { - "async/behavior_version": batch.behavior_version, - "async/behavior_lag": versions.behavior_lag(batch.behavior_version), - } - def _advantage_and_train( self, sample: Sample, @@ -174,15 +144,9 @@ def _advantage_and_train( part = part.compute_advantages(normalize=True, use_global_std=self._adv_use_global_std) sample = sample.replace_frontier(part) result = self.stack.train_track(sample.parts[-1], training_progress=float(training_progress)) - self._policy_versions.record_optimizer_updates(result.optimizer_updates) + self._control.record_optimizer_updates(result.optimizer_updates) if extra_metrics is not None: - extra_metrics.update( - { - "async/train_version": self._policy_versions.train_version, - "async/rollout_lag": self._policy_versions.rollout_lag, - "async/optimizer_updates": result.optimizer_updates, - } - ) + extra_metrics.update(self._control.train_metrics(result.optimizer_updates)) self.wandb_logger.log_rollout_step( rollout_id, result, @@ -208,15 +172,15 @@ def train( self.backend.get_optimizer_step_count(), name="backend optimizer step count", ) - self._policy_versions = PolicyVersionState(train_version=train_version) + self._control.restore(train_version) for _ in range(start_rollout): self.data_source.get_samples(self.batch_size) self._init_wandb( num_rollouts=num_rollouts, extra={ "max_inflight": self._max_inflight, - "max_policy_lag": self._max_policy_lag, - "num_updates_per_batch": self._num_updates_per_batch, + "max_policy_lag": self._control.max_policy_lag, + "num_updates_per_batch": self._control.num_updates_per_batch, "train_fraction": self._train_fraction, }, ) @@ -229,7 +193,7 @@ def train( ) if resumed: - self._sync_rollout_weights(force=True) + self._control.sync_rollout(self._async_engine, self.rollout, self.weight_sync) if self.eval_interval > 0: self.evaluate(start_rollout, sync_weights=False, sleep_after=False) @@ -254,22 +218,16 @@ def train( training_progress=training_progress, rollout_id=rollout_id, t0=t0, - extra_metrics=self._policy_metrics(batch), + extra_metrics=self._control.behavior_metrics(batch.behavior_version), ) self.wandb_logger.log_progress(rollout_id, num_rollouts, result, mean_reward, logger=logger) step = rollout_id + 1 eval_due = self.eval_interval > 0 and step % self.eval_interval == 0 save_due = save_interval > 0 and (step % save_interval == 0 or step >= num_rollouts) - sync_due = step < num_rollouts and self._policy_versions.rollout_lag > self._max_policy_lag + sync_due = step < num_rollouts and self._control.rollout_lag > self._control.max_policy_lag if eval_due or save_due or sync_due: - if self._async_engine.inflight_count + self._async_engine.ready_count != 0: - raise RuntimeError( - "async sync boundary retained rollout work: " - f"inflight_count={self._async_engine.inflight_count}, " - f"ready_count={self._async_engine.ready_count}" - ) - self._sync_rollout_weights() + self._control.sync_rollout(self._async_engine, self.rollout, self.weight_sync) if eval_due: self.evaluate(step, sync_weights=False, sleep_after=False) @@ -309,11 +267,7 @@ def _next_rollout_batch( engine = self._async_engine while True: engine.poll() - slots = launch_slots( - train_version=self._policy_versions.train_version, - rollout_version=self._policy_versions.rollout_version, - num_updates_per_batch=self._num_updates_per_batch, - max_policy_lag=self._max_policy_lag, + slots = self._control.launch_slots( inflight_count=engine.inflight_count, ready_count=engine.ready_count, max_inflight=self._max_inflight, @@ -324,11 +278,11 @@ def _next_rollout_batch( for _ in range(slots): engine.submit( self._build_async_sample(engine.next_gen_id), - behavior_version=self._policy_versions.rollout_version, + behavior_version=self._control.rollout_version, ) batch = engine.pop_next_batch( - train_version=self._policy_versions.train_version, - max_policy_lag=self._max_policy_lag, + train_version=self._control.train_version, + max_policy_lag=self._control.max_policy_lag, ) if batch is not None: return batch diff --git a/unirl/trainer/async_policy.py b/unirl/trainer/async_policy.py index 9783d9444..053574f50 100644 --- a/unirl/trainer/async_policy.py +++ b/unirl/trainer/async_policy.py @@ -1,30 +1,26 @@ -"""Shared policy-version accounting for disaggregated async trainers. - -Versions in this module count committed optimizer updates, never weight-sync -calls or consumed rollout batches. A rollout batch records the train version -whose weights were resident in the rollout engine when generation started. -""" +"""Optimizer-update policy control shared by async AR and diffusion trainers.""" from __future__ import annotations from dataclasses import dataclass +from typing import Any @dataclass -class PolicyVersionState: - """Driver-owned train/rollout policy clocks. - - ``train_version`` advances by the number of optimizer steps that actually - committed. ``rollout_version`` is assigned to the exact train snapshot most - recently synced to the rollout engine. - """ +class AsyncBatchControl: + """Track train/rollout versions and gate batch generation.""" + max_policy_lag: int + num_updates_per_batch: int train_version: int = 0 rollout_version: int = 0 + rollout_initialized: bool = False def __post_init__(self) -> None: - if self.train_version < 0 or self.rollout_version < 0: - raise ValueError("policy versions must be non-negative") + if self.max_policy_lag < 0: + raise ValueError(f"max_policy_lag must be >= 0, got {self.max_policy_lag}") + if self.num_updates_per_batch < 1: + raise ValueError(f"num_updates_per_batch must be >= 1, got {self.num_updates_per_batch}") if self.rollout_version > self.train_version: raise ValueError( f"rollout_version cannot be ahead of train_version: {self.rollout_version} > {self.train_version}" @@ -34,27 +30,16 @@ def __post_init__(self) -> None: def rollout_lag(self) -> int: return self.train_version - self.rollout_version - def record_optimizer_updates(self, committed_updates: int) -> int: - """Advance by optimizer steps that successfully changed train weights.""" - - if committed_updates < 0: - raise ValueError(f"committed optimizer updates must be >= 0, got {committed_updates}") - self.train_version += committed_updates - return self.train_version - - def mark_rollout_synced(self, train_version: int) -> int: - """Record the train version loaded by a successful rollout weight sync.""" + @property + def freshness_depth(self) -> int: + return self.max_policy_lag // self.num_updates_per_batch + 1 - if train_version < self.rollout_version: - raise ValueError(f"synced train version cannot move backwards: {train_version} < {self.rollout_version}") - if train_version > self.train_version: - raise ValueError(f"cannot sync future train version {train_version} > {self.train_version}") - self.rollout_version = train_version - return self.rollout_version + def restore(self, train_version: int) -> None: + self.train_version = train_version + self.rollout_version = 0 + self.rollout_initialized = False def behavior_lag(self, behavior_version: int) -> int: - """Optimizer-update lag between train and a batch's behavior policy.""" - lag = self.train_version - behavior_version if lag < 0: raise ValueError( @@ -62,65 +47,51 @@ def behavior_lag(self, behavior_version: int) -> int: ) return lag - -def launch_slots( - *, - train_version: int, - rollout_version: int, - num_updates_per_batch: int, - max_policy_lag: int, - inflight_count: int, - ready_count: int, - max_inflight: int, - trained_batches: int, - num_rollouts: int, - hard_boundary: int, -) -> int: - """How many batch generations may be launched by the single-threaded loop. - - Freshness is measured in committed optimizer updates. The first outstanding - batch would train at the current lag; every additional outstanding batch is - conservatively reserved ``num_updates_per_batch`` future updates. Capacity and - durable-boundary clamps are applied in the same generation-batch unit. - """ - - train = _non_negative("train_version", train_version) - rollout = _non_negative("rollout_version", rollout_version) - updates = _positive("num_updates_per_batch", num_updates_per_batch) - max_lag = _non_negative("max_policy_lag", max_policy_lag) - active = _non_negative("inflight_count", inflight_count) - queued = _non_negative("ready_count", ready_count) - max_active = _positive("max_inflight", max_inflight) - trained = _non_negative("trained_batches", trained_batches) - total = _non_negative("num_rollouts", num_rollouts) - boundary = _non_negative("hard_boundary", hard_boundary) - if rollout > train: - raise ValueError(f"rollout_version cannot be ahead of train_version: {rollout} > {train}") - if active > max_active: - raise ValueError(f"inflight_count={active} exceeds max_inflight={max_active}") - if boundary < trained: - raise ValueError(f"hard_boundary={boundary} is behind trained_batches={trained}") - if trained >= total: - return 0 - - current_lag = train - rollout - if current_lag > max_lag: - return 0 - - freshness_slots = (max_lag - current_lag) // updates + 1 - allowed_outstanding = min( - freshness_slots, - total - trained, - boundary - trained, - ) - outstanding = active + queued - return max( - 0, - min( - max_active - active, - allowed_outstanding - outstanding, - ), - ) + def record_optimizer_updates(self, optimizer_updates: int) -> None: + self.train_version += optimizer_updates + + def launch_slots( + self, + *, + inflight_count: int, + ready_count: int, + max_inflight: int, + trained_batches: int, + num_rollouts: int, + hard_boundary: int, + ) -> int: + if self.rollout_lag > self.max_policy_lag: + return 0 + freshness = (self.max_policy_lag - self.rollout_lag) // self.num_updates_per_batch + 1 + allowed = min(freshness, num_rollouts - trained_batches, hard_boundary - trained_batches) + return max(0, min(max_inflight - inflight_count, allowed - inflight_count - ready_count)) + + def sync_rollout(self, engine: Any, rollout: Any, weight_sync: Any) -> bool: + if self.rollout_initialized and self.rollout_lag == 0: + return False + engine.quiesce() + if engine.ready_count: + raise RuntimeError( + f"cannot sync rollout weights with completed batches queued: ready_count={engine.ready_count}" + ) + weight_sync.sync() + rollout.set_policy_version(self.train_version) + self.rollout_version = self.train_version + self.rollout_initialized = True + return True + + def behavior_metrics(self, behavior_version: int) -> dict[str, int]: + return { + "async/behavior_version": behavior_version, + "async/behavior_lag": self.behavior_lag(behavior_version), + } + + def train_metrics(self, optimizer_updates: int) -> dict[str, int]: + return { + "async/train_version": self.train_version, + "async/rollout_lag": self.rollout_lag, + "async/optimizer_updates": optimizer_updates, + } def next_hard_boundary( @@ -132,12 +103,10 @@ def next_hard_boundary( ) -> int: """Nearest eval/checkpoint/final boundary for launch admission.""" - trained = _non_negative("trained_batches", trained_batches) - total = _non_negative("num_rollouts", num_rollouts) - boundary = total + boundary = num_rollouts for interval in (eval_interval, save_interval): - if interval > 0 and trained < total: - boundary = min(boundary, ((trained // interval) + 1) * interval) + if interval > 0 and trained_batches < num_rollouts: + boundary = min(boundary, ((trained_batches // interval) + 1) * interval) return boundary @@ -145,34 +114,14 @@ def unwrap_replicated_int(value: object, *, name: str) -> int: """Normalize a BROADCAST return and verify all worker replicas agree.""" if isinstance(value, (list, tuple)): - if not value: - raise ValueError(f"{name} returned no worker values") - if any(not isinstance(item, int) for item in value): - raise TypeError(f"{name} returned non-integer worker values: {value!r}") - first = value[0] - if any(item != first for item in value[1:]): + if not value or any(not isinstance(item, int) for item in value): + raise TypeError(f"{name} returned invalid worker values: {value!r}") + if any(item != value[0] for item in value[1:]): raise RuntimeError(f"{name} disagrees across workers: {value!r}") - return first + return value[0] if not isinstance(value, int): raise TypeError(f"{name} returned {type(value).__name__}, expected int") return value -def _non_negative(name: str, value: int) -> int: - if value < 0: - raise ValueError(f"{name} must be >= 0, got {value}") - return value - - -def _positive(name: str, value: int) -> int: - if value < 1: - raise ValueError(f"{name} must be >= 1, got {value}") - return value - - -__all__ = [ - "PolicyVersionState", - "launch_slots", - "next_hard_boundary", - "unwrap_replicated_int", -] +__all__ = ["AsyncBatchControl", "next_hard_boundary", "unwrap_replicated_int"] From fbefbfa2ed327a3990de3de781bff72155160852 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 4 Aug 2026 13:44:39 +0800 Subject: [PATCH 5/9] refactor(async): make bootstrap sync explicit Remove the hidden rollout initialization flag and force only bootstrap/resume syncs at their call sites. --- unirl/trainer/async_ar.py | 4 ++-- unirl/trainer/async_diffusion.py | 2 +- unirl/trainer/async_policy.py | 7 ++----- 3 files changed, 5 insertions(+), 8 deletions(-) diff --git a/unirl/trainer/async_ar.py b/unirl/trainer/async_ar.py index 2af04dd67..e8d527b3b 100644 --- a/unirl/trainer/async_ar.py +++ b/unirl/trainer/async_ar.py @@ -335,8 +335,8 @@ def train( start_gen_id=start_rollout, ) - if resumed: - self._control.sync_rollout(self._async_engine, self.rollout, self.weight_sync) + if resumed or self.eval_interval > 0: + self._control.sync_rollout(self._async_engine, self.rollout, self.weight_sync, force=True) if self.eval_interval > 0: self.evaluate(rollout_id=-1) diff --git a/unirl/trainer/async_diffusion.py b/unirl/trainer/async_diffusion.py index 912c8cf4b..a7d556214 100644 --- a/unirl/trainer/async_diffusion.py +++ b/unirl/trainer/async_diffusion.py @@ -193,7 +193,7 @@ def train( ) if resumed: - self._control.sync_rollout(self._async_engine, self.rollout, self.weight_sync) + self._control.sync_rollout(self._async_engine, self.rollout, self.weight_sync, force=True) if self.eval_interval > 0: self.evaluate(start_rollout, sync_weights=False, sleep_after=False) diff --git a/unirl/trainer/async_policy.py b/unirl/trainer/async_policy.py index 053574f50..ac3960f34 100644 --- a/unirl/trainer/async_policy.py +++ b/unirl/trainer/async_policy.py @@ -14,7 +14,6 @@ class AsyncBatchControl: num_updates_per_batch: int train_version: int = 0 rollout_version: int = 0 - rollout_initialized: bool = False def __post_init__(self) -> None: if self.max_policy_lag < 0: @@ -37,7 +36,6 @@ def freshness_depth(self) -> int: def restore(self, train_version: int) -> None: self.train_version = train_version self.rollout_version = 0 - self.rollout_initialized = False def behavior_lag(self, behavior_version: int) -> int: lag = self.train_version - behavior_version @@ -66,8 +64,8 @@ def launch_slots( allowed = min(freshness, num_rollouts - trained_batches, hard_boundary - trained_batches) return max(0, min(max_inflight - inflight_count, allowed - inflight_count - ready_count)) - def sync_rollout(self, engine: Any, rollout: Any, weight_sync: Any) -> bool: - if self.rollout_initialized and self.rollout_lag == 0: + def sync_rollout(self, engine: Any, rollout: Any, weight_sync: Any, *, force: bool = False) -> bool: + if not force and self.rollout_lag == 0: return False engine.quiesce() if engine.ready_count: @@ -77,7 +75,6 @@ def sync_rollout(self, engine: Any, rollout: Any, weight_sync: Any) -> bool: weight_sync.sync() rollout.set_policy_version(self.train_version) self.rollout_version = self.train_version - self.rollout_initialized = True return True def behavior_metrics(self, behavior_version: int) -> dict[str, int]: From 8ed6bb3d2918b23c10ad02fbcb41866ecbf1a97d Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 4 Aug 2026 20:11:53 +0800 Subject: [PATCH 6/9] refactor(async): denominate the staleness budget in rollout batches max_policy_lag counted committed optimizer updates, but admission and consumption only ever run at batch boundaries, so the knob quantized to num_updates_per_batch and the same value meant different depths per recipe. max_staleness now counts whole rollout batches (AReaL's unit for eta) and staleness_budget converts once into the update clock, which keeps a partially-committed step honest. Also 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. Shipped recipes are behavior-identical. --- .../qwen3_drpo_4b_base_dapo_sglang_async.yaml | 9 +- .../qwen3_grpo_4b_base_dapo_sglang_async.yaml | 13 ++- .../diffusion/bagel/bagel_vllmomni_async.yaml | 15 ++- unirl/rollout/engine/asynchronous.py | 21 ++-- unirl/train_async_ar.py | 7 +- unirl/train_async_diffusion.py | 7 +- unirl/trainer/README.md | 2 +- unirl/trainer/async_ar.py | 37 +++--- unirl/trainer/async_diffusion.py | 27 +++-- unirl/trainer/async_policy.py | 107 +++++++++++++++--- 10 files changed, 172 insertions(+), 73 deletions(-) diff --git a/examples/ar/qwen3_drpo_4b_base_dapo_sglang_async.yaml b/examples/ar/qwen3_drpo_4b_base_dapo_sglang_async.yaml index 4fac94496..ae6fa4984 100644 --- a/examples/ar/qwen3_drpo_4b_base_dapo_sglang_async.yaml +++ b/examples/ar/qwen3_drpo_4b_base_dapo_sglang_async.yaml @@ -23,11 +23,12 @@ num_rollouts: 800 # Train slab share (rollout gets the rest). 0.5 -> equal split. train_fraction: 0.5 -# Async control. Policy lag is measured in committed optimizer updates. -# max_policy_lag=0 keeps the behavior/train versions aligned at batch -# admission; max_inflight remains a resource cap. +# Async control. max_staleness counts whole rollout batches the behavior policy +# may trail the train policy by; admission depth is max_staleness + 1. +# max_staleness=0 keeps the behavior/train versions aligned at batch admission; +# max_inflight remains a resource cap. max_inflight: 1 -max_policy_lag: 0 +max_staleness: 0 eval_interval: 10 adv_normalization_scope: group diff --git a/examples/ar/qwen3_grpo_4b_base_dapo_sglang_async.yaml b/examples/ar/qwen3_grpo_4b_base_dapo_sglang_async.yaml index 2e79b0808..e29450b9b 100644 --- a/examples/ar/qwen3_grpo_4b_base_dapo_sglang_async.yaml +++ b/examples/ar/qwen3_grpo_4b_base_dapo_sglang_async.yaml @@ -26,12 +26,15 @@ num_rollouts: 800 # train-bound, lower for generation-bound workloads (mind DIVISIBILITY above). train_fraction: 0.5 -# Async control. Policy lag is measured in committed optimizer updates. -# max_policy_lag=0 means every batch starts training against the same -# train version that generated it (rollout/replay numerical differences may -# still make ratio only approximately 1). max_inflight remains a resource cap. +# Async control. max_staleness counts whole rollout batches the behavior policy +# may trail the train policy by; admission depth is max_staleness + 1. +# max_staleness=0 means every batch starts training against the same train +# version that generated it (rollout/replay numerical differences may still make +# ratio only approximately 1) — note that with num_updates_per_batch=4 below, +# updates 2-4 of a batch are still 1-3 updates off-policy by design, which the +# frozen pi_old anchor and clip_range cover. max_inflight remains a resource cap. max_inflight: 1 -max_policy_lag: 0 +max_staleness: 0 # AIME eval every 10 rollouts + a pre-train baseline (rollout -1). eval_interval: 10 diff --git a/examples/diffusion/bagel/bagel_vllmomni_async.yaml b/examples/diffusion/bagel/bagel_vllmomni_async.yaml index df4ef7fcd..5b38b2017 100644 --- a/examples/diffusion/bagel/bagel_vllmomni_async.yaml +++ b/examples/diffusion/bagel/bagel_vllmomni_async.yaml @@ -43,18 +43,23 @@ workers_per_device: 1 layout: separate train_fraction: 0.5 # 4 train GPUs + 4 rollout GPUs on an 8-GPU pool -# ---- async control (optimizer-update units) ---- +# ---- async control ---- # max_inflight: concurrent generations. MUST be 1: the trajectory-segment # cross-slab transfer (NCCL send) runs on the rollout worker; a second in-flight # generation co-tenanting that worker blocks the send behind it (~150s/rollout). # With max_inflight=1 the trainer polls (reaps) before topping up launches, so # each generation transfers in the idle window before the next launch, then # that next generation overlaps the train step. Real overlap needs admission -# depth > 1. -# With num_updates_per_batch=2, max_policy_lag=22 admits a 12-batch -# behavior-policy cycle: floor(22 / 2) + 1 = 12. +# depth > 1, i.e. max_staleness >= 1. +# +# max_staleness counts whole rollout batches, so admission depth is +# max_staleness + 1 = 12 regardless of num_updates_per_batch. Note this is also +# the weight-sync period: 11 batches (22 optimizer updates) elapse between +# publications, which is a long way 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. Lower this first if the ratio starts to swing. max_inflight: 1 -max_policy_lag: 22 +max_staleness: 11 logging: report_to_wandb: false # flip to true to enable wandb (rank-0/driver only) diff --git a/unirl/rollout/engine/asynchronous.py b/unirl/rollout/engine/asynchronous.py index 6d0c6bf06..8ce47aff7 100644 --- a/unirl/rollout/engine/asynchronous.py +++ b/unirl/rollout/engine/asynchronous.py @@ -133,19 +133,22 @@ def pop_next( self, *, train_version: int, - max_policy_lag: int, + staleness_budget: int, ) -> Optional[RolloutBatch]: if not self._items: return None item = self._items[0] - lag = train_version - item.behavior_version - if lag < 0: + staleness = train_version - item.behavior_version + if staleness < 0: raise RuntimeError( f"generation {item.gen_id} has future behavior version " f"{item.behavior_version} > train version {train_version}" ) - if lag > max_policy_lag: - raise RuntimeError(f"generation {item.gen_id} exceeded policy lag budget: lag={lag} > max={max_policy_lag}") + if staleness > staleness_budget: + raise RuntimeError( + f"generation {item.gen_id} exceeded its staleness budget: " + f"staleness={staleness} > budget={staleness_budget} optimizer updates" + ) return self._items.popleft() @@ -248,7 +251,9 @@ class AsyncBatchRolloutEngine: One generation must produce exactly ``groups_per_batch`` groups and enters a completion-order FIFO atomically. ``behavior_version`` is captured at launch - from the exact train snapshot currently synced to the rollout engine. + from the exact train snapshot currently synced to the rollout engine, so + ``staleness_budget`` is compared in committed optimizer updates — the caller + resolves it from a batch count before passing it down. ``quiesce()`` (drain everything) is MANDATORY before a weight sync, eval, or checkpoint: a weight + KV update corrupts an in-flight generation. @@ -294,11 +299,11 @@ def pop_next_batch( self, *, train_version: int, - max_policy_lag: int, + staleness_budget: int, ) -> Optional[RolloutBatch]: return self._ready.pop_next( train_version=train_version, - max_policy_lag=max_policy_lag, + staleness_budget=staleness_budget, ) def quiesce(self) -> None: diff --git a/unirl/train_async_ar.py b/unirl/train_async_ar.py index ad4f90b17..6e7e4139c 100755 --- a/unirl/train_async_ar.py +++ b/unirl/train_async_ar.py @@ -17,7 +17,8 @@ must both be integers, AND ``batch_size * samples_per_prompt`` must be divisible by each slab size (DP_SCATTER divisibility). * ``max_inflight`` — concurrent generations (resource/overlap limit). - * ``max_policy_lag`` — maximum optimizer-update lag at batch admission. + * ``max_staleness`` — how many whole rollout batches the behavior policy may + trail the train policy by at batch admission. """ from __future__ import annotations @@ -26,10 +27,12 @@ from omegaconf import DictConfig from unirl.trainer.async_ar import AsyncARTrainer +from unirl.trainer.async_policy import reject_removed_async_keys @hydra.main(version_base=None, config_path="../examples", config_name="ar/qwen3_grpo_4b_base_dapo_sglang_async") def main(cfg: DictConfig) -> None: + reject_removed_async_keys(cfg) trainer = AsyncARTrainer( cfg=cfg, batch_size=cfg.batch_size, @@ -55,7 +58,7 @@ def main(cfg: DictConfig) -> None: eval_temperature=float(cfg.get("eval_temperature", 1.0)), train_fraction=float(cfg.get("train_fraction", 0.5)), max_inflight=int(cfg.get("max_inflight", 1)), - max_policy_lag=cfg.get("max_policy_lag", 0), + max_staleness=cfg.get("max_staleness", 0), ) trainer.train( num_rollouts=int(cfg.get("num_rollouts", 100)), diff --git a/unirl/train_async_diffusion.py b/unirl/train_async_diffusion.py index 6ec0f3848..9c4c534d0 100755 --- a/unirl/train_async_diffusion.py +++ b/unirl/train_async_diffusion.py @@ -15,7 +15,8 @@ Extra config knobs vs the synchronous separate recipe: * ``max_inflight`` — must be ``1``; other values fail during trainer initialization. - * ``max_policy_lag`` — maximum optimizer-update lag at batch admission. + * ``max_staleness`` — how many whole rollout batches the behavior policy may + trail the train policy by at batch admission. ``layout`` is forced to ``separate`` (async needs disjoint train/rollout slabs). """ @@ -25,10 +26,12 @@ from omegaconf import DictConfig from unirl.trainer.async_diffusion import AsyncDiffusionTrainer +from unirl.trainer.async_policy import reject_removed_async_keys @hydra.main(version_base=None, config_path="../examples", config_name="diffusion/bagel/bagel_vllmomni_async") def main(cfg: DictConfig) -> None: + reject_removed_async_keys(cfg) trainer = AsyncDiffusionTrainer( cfg=cfg, batch_size=cfg.batch_size, @@ -56,7 +59,7 @@ def main(cfg: DictConfig) -> None: eval_rewards_cfg=cfg.get("eval_rewards"), task_config=cfg.get("task_config"), max_inflight=int(cfg.get("max_inflight", 1)), - max_policy_lag=cfg.get("max_policy_lag", 0), + max_staleness=cfg.get("max_staleness", 0), ) trainer.train( num_rollouts=cfg.get("num_rollouts", 100), diff --git a/unirl/trainer/README.md b/unirl/trainer/README.md index aec0c5a5a..27d209650 100644 --- a/unirl/trainer/README.md +++ b/unirl/trainer/README.md @@ -59,7 +59,7 @@ The current trainer surface is: | `DiffusionTrainer` | one diffusion `Part` → one `TrainStack` | Reference diffusion loop; supports trainside or dedicated rollout, optional separate reward GPUs, FSDP offload, and DiffusionNFT's EMA-adapter rollout. | | `ARTrainer` | one AR `Part` → one `TrainStack` | Text or multimodal AR rollout with group/global advantage normalization and optional token-balanced DP shards. | | `SFTTrainer` | dataset records → one standalone training `Part` | Reuses the RL TrainStack without rollout, reward, or advantages; owns exact epoch/cursor resume and full-set evaluation. | -| `AsyncARTrainer` | FIFO AR generation batch → one `TrainStack` | Separate train/rollout slabs with resident generation, optimizer-update policy-lag admission, and quiescence before weight sync, eval, or checkpoint. | +| `AsyncARTrainer` | FIFO AR generation batch → one `TrainStack` | Separate train/rollout slabs with resident generation, batch-denominated `max_staleness` admission over an optimizer-update clock, and quiescence before weight sync, eval, or checkpoint. | | `AsyncDiffusionTrainer` | FIFO diffusion generation batch → one `TrainStack` | The same update-versioned separate-slab loop for DiT. Requires `max_inflight=1` and reaps each generation before launching the next, so the cross-slab trajectory transfer never queues behind a fresh generation. | | `PETrainer` | `ar` + `diffusion` Parts → two `TrainStack`s | Composed prompt-rewrite/image rollout; image rewards propagate to AR rewrites. `freeze_llm=true` trains and checkpoints diffusion only. | | `UnifiedModelTrainer` | whole `Sample` → one `UnifiedModelTrainStack` | AR and image losses accumulate into shared-backbone optimizer steps while prompt-tree lineage remains intact during DP scatter. | diff --git a/unirl/trainer/async_ar.py b/unirl/trainer/async_ar.py index e8d527b3b..2dab92236 100644 --- a/unirl/trainer/async_ar.py +++ b/unirl/trainer/async_ar.py @@ -9,14 +9,17 @@ ONE single-threaded loop (slime's "one trainer loop; async-depth is a knob" principle, implemented with UniRL-native non-blocking Ray dispatch instead of -slime's thread+asyncio). Async freshness is measured directly in committed -optimizer updates: +slime's thread+asyncio). Async freshness rides a clock of committed optimizer +updates, but the budget is stated in whole rollout batches: * ``max_inflight`` — how many generations run concurrently (overlap/parallelism depth). ``1`` ≈ the classic one-step pipeline; higher fans out more. -* ``max_policy_lag`` — inclusive train-minus-behavior optimizer-update - lag at batch admission. ``0`` aligns policy versions; the rollout-anchored - PPO ratio remains the numerical source of truth. +* ``max_staleness`` — how many whole rollout batches the behavior policy may + trail the train policy by, checked at batch admission and consumption. ``0`` + aligns policy versions; the rollout-anchored PPO ratio remains the numerical + source of truth. Batches are the unit because admission only ever runs at a + batch boundary; with ``num_updates_per_batch > 1`` the updates after the first + in a batch additionally drift by up to ``num_updates_per_batch - 1``. Generation runs through :class:`~unirl.rollout.engine.asynchronous.AsyncBatchRolloutEngine` (non-blocking Ray futures over the rollout Handle) on the single driver thread — @@ -101,7 +104,7 @@ def __init__( eval_temperature: float = 1.0, train_fraction: float = 0.5, max_inflight: int = 1, - max_policy_lag: int = 0, + max_staleness: int = 0, ) -> None: validate_qwen3_5_training_contract( pipeline_cfg=pipeline_cfg, @@ -139,19 +142,18 @@ def __init__( self._train_fraction = float(train_fraction) self._max_inflight = max(1, int(max_inflight)) self._control = AsyncBatchControl( - max_policy_lag=max_policy_lag, + max_staleness=max_staleness, num_updates_per_batch=stack_cfg.get("num_updates_per_batch", 1), ) - if self._max_inflight > self._control.freshness_depth: + if self._max_inflight > self._control.admission_depth: logger.warning( - "max_inflight=%d exceeds the policy-lag admission depth %d; the extra concurrency cannot be used", + "max_inflight=%d exceeds the staleness admission depth %d; the extra concurrency cannot be used", self._max_inflight, - self._control.freshness_depth, + self._control.admission_depth, ) - if self._control.freshness_depth == 1: + if self._control.max_staleness == 0: logger.warning( - "async policy-lag settings admit one generation at a time; " - "generation cannot overlap the preceding train batch" + "max_staleness=0 admits one generation at a time; generation cannot overlap the preceding train batch" ) self._train_devices = int(round(self.num_devices * self._train_fraction)) if self._train_devices <= 0 or self._train_devices >= self.num_devices: @@ -323,7 +325,8 @@ def train( extra={ "adv_normalization_scope": self.adv_normalization_scope, "max_inflight": self._max_inflight, - "max_policy_lag": self._control.max_policy_lag, + "max_staleness": self._control.max_staleness, + "staleness_budget": self._control.staleness_budget, "num_updates_per_batch": self._control.num_updates_per_batch, }, ) @@ -368,7 +371,7 @@ def train( step = rollout_id + 1 eval_due = self.eval_interval > 0 and step % self.eval_interval == 0 save_due = save_interval > 0 and (step % save_interval == 0 or step >= num_rollouts) - sync_due = step < num_rollouts and self._control.rollout_lag > self._control.max_policy_lag + sync_due = step < num_rollouts and self._control.publish_lag > self._control.staleness_budget if eval_due or save_due or sync_due: self._control.sync_rollout(self._async_engine, self.rollout, self.weight_sync) @@ -420,11 +423,11 @@ def _next_rollout_batch( engine.poll() batch = engine.pop_next_batch( train_version=self._control.train_version, - max_policy_lag=self._control.max_policy_lag, + staleness_budget=self._control.staleness_budget, ) if batch is not None: return batch if engine.inflight_count: engine.wait_oldest() else: - raise RuntimeError("async rollout queue is empty and policy lag admits no new generation") + raise RuntimeError("async rollout queue is empty and the staleness budget admits no new generation") diff --git a/unirl/trainer/async_diffusion.py b/unirl/trainer/async_diffusion.py index a7d556214..e77b7320f 100644 --- a/unirl/trainer/async_diffusion.py +++ b/unirl/trainer/async_diffusion.py @@ -23,8 +23,13 @@ Async control uses the same optimizer-update clock as AsyncARTrainer: * ``max_inflight`` — must be ``1`` so a reap-time transfer never competes with a queued generation on the rollout workers. - * ``max_policy_lag`` — inclusive train-minus-behavior lag at batch - admission, measured in committed optimizer updates. + * ``max_staleness`` — how many whole rollout batches the behavior policy may + trail the train policy by, checked at batch admission and consumption. The + clock underneath counts committed optimizer updates; the budget is stated in + batches because those are the only points where admission can react. + ``num_updates_per_batch > 1`` therefore also means the updates after the + first in a batch run at up to ``num_updates_per_batch - 1`` more updates of + drift than the admitted figure. ``_next_rollout_batch`` polls (reaps) BEFORE topping up launches, which is what makes the overlap fast here: reaping a generation pulls its trajectory segment off the @@ -65,7 +70,7 @@ def __init__( self, *, max_inflight: int = 1, - max_policy_lag: int = 0, + max_staleness: int = 0, **diffusion_kwargs: Any, ) -> None: layout = diffusion_kwargs.setdefault("layout", "separate") @@ -87,13 +92,12 @@ def __init__( self._max_inflight = max_inflight self._control = AsyncBatchControl( - max_policy_lag=max_policy_lag, + max_staleness=max_staleness, num_updates_per_batch=diffusion_kwargs["stack_cfg"].get("num_updates_per_batch", 1), ) - if self._control.freshness_depth == 1: + if self._control.max_staleness == 0: logger.warning( - "async policy-lag settings admit one generation at a time; " - "generation cannot overlap the preceding train batch" + "max_staleness=0 admits one generation at a time; generation cannot overlap the preceding train batch" ) def _build_async_sample(self, gen_id: int) -> Sample: @@ -179,7 +183,8 @@ def train( num_rollouts=num_rollouts, extra={ "max_inflight": self._max_inflight, - "max_policy_lag": self._control.max_policy_lag, + "max_staleness": self._control.max_staleness, + "staleness_budget": self._control.staleness_budget, "num_updates_per_batch": self._control.num_updates_per_batch, "train_fraction": self._train_fraction, }, @@ -225,7 +230,7 @@ def train( step = rollout_id + 1 eval_due = self.eval_interval > 0 and step % self.eval_interval == 0 save_due = save_interval > 0 and (step % save_interval == 0 or step >= num_rollouts) - sync_due = step < num_rollouts and self._control.rollout_lag > self._control.max_policy_lag + sync_due = step < num_rollouts and self._control.publish_lag > self._control.staleness_budget if eval_due or save_due or sync_due: self._control.sync_rollout(self._async_engine, self.rollout, self.weight_sync) @@ -282,11 +287,11 @@ def _next_rollout_batch( ) batch = engine.pop_next_batch( train_version=self._control.train_version, - max_policy_lag=self._control.max_policy_lag, + staleness_budget=self._control.staleness_budget, ) if batch is not None: return batch if engine.inflight_count: engine.wait_oldest() else: - raise RuntimeError("async rollout queue is empty and policy lag admits no new generation") + raise RuntimeError("async rollout queue is empty and the staleness budget admits no new generation") diff --git a/unirl/trainer/async_policy.py b/unirl/trainer/async_policy.py index ac3960f34..7c61de97a 100644 --- a/unirl/trainer/async_policy.py +++ b/unirl/trainer/async_policy.py @@ -1,23 +1,72 @@ -"""Optimizer-update policy control shared by async AR and diffusion trainers.""" +"""Optimizer-update policy control shared by async AR and diffusion trainers. + +Two quantities ride the optimizer-update clock and only one of them is staleness: + +* ``staleness`` — updates between the current train weights and the behavior + policy that generated a batch, i.e. the off-policyness of the data actually + being trained on. This is AReaL's ``eta`` / ``max_head_offpolicyness``. +* ``publish_lag`` — updates between the current train weights and the snapshot + last published to the rollout engine. This is sync debt: no batch is that + stale, but it is what makes a weight sync due. + +Recipes state the budget in whole rollout batches (``max_staleness``) because +admission and consumption only ever run at a batch boundary. Stating it in raw +updates instead would quantize it to ``num_updates_per_batch`` — 22 and 23 would +both mean a 12-batch depth — and silently change meaning whenever that count +changes. ``staleness_budget`` converts once into the clock the versions count in. +""" from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass from typing import Any +_REMOVED_ASYNC_KEYS = { + "max_policy_lag": ( + "max_policy_lag was replaced by max_staleness: the budget now counts whole rollout " + "batches instead of committed optimizer updates. Divide the old value by " + "stack.num_updates_per_batch and rename the key." + ), + "buffer_max_staleness": ( + "buffer_max_staleness (weight-syncs a buffered group could cross, under " + "freshest-first selection) was replaced by max_staleness, which counts whole " + "rollout batches under completion-order FIFO. The units do not correspond; pick a " + "new value rather than carrying the old one over." + ), + "weight_sync_interval": ( + "weight_sync_interval was removed: weight sync is now driven by max_staleness and " + "by eval/checkpoint boundaries, not by a fixed counter." + ), +} + + +def reject_removed_async_keys(cfg: Mapping[str, Any]) -> None: + """Fail a recipe still carrying a removed or re-denominated async knob. + + Every one of these is a silent-wrong hazard rather than a crash: an ignored + key runs at the default, and reading ``max_policy_lag`` as ``max_staleness`` + yields a perfectly valid budget that is ``num_updates_per_batch`` times the + intended one. Nothing downstream can tell either case from a deliberate one. + """ + + for key, message in _REMOVED_ASYNC_KEYS.items(): + if key in cfg: + raise ValueError(message) + @dataclass class AsyncBatchControl: """Track train/rollout versions and gate batch generation.""" - max_policy_lag: int + max_staleness: int num_updates_per_batch: int train_version: int = 0 rollout_version: int = 0 def __post_init__(self) -> None: - if self.max_policy_lag < 0: - raise ValueError(f"max_policy_lag must be >= 0, got {self.max_policy_lag}") + if self.max_staleness < 0: + raise ValueError(f"max_staleness must be >= 0, got {self.max_staleness}") if self.num_updates_per_batch < 1: raise ValueError(f"num_updates_per_batch must be >= 1, got {self.num_updates_per_batch}") if self.rollout_version > self.train_version: @@ -26,24 +75,36 @@ def __post_init__(self) -> None: ) @property - def rollout_lag(self) -> int: + def staleness_budget(self) -> int: + """``max_staleness`` batches expressed in committed optimizer updates.""" + + return self.max_staleness * self.num_updates_per_batch + + @property + def publish_lag(self) -> int: + """Sync debt: updates the published rollout snapshot trails train by.""" + return self.train_version - self.rollout_version @property - def freshness_depth(self) -> int: - return self.max_policy_lag // self.num_updates_per_batch + 1 + def admission_depth(self) -> int: + """Outstanding generations the budget allows, in whole batches.""" + + return self.max_staleness + 1 def restore(self, train_version: int) -> None: self.train_version = train_version self.rollout_version = 0 - def behavior_lag(self, behavior_version: int) -> int: - lag = self.train_version - behavior_version - if lag < 0: + def staleness(self, behavior_version: int) -> int: + """Optimizer updates between train and a batch's behavior policy.""" + + stale = self.train_version - behavior_version + if stale < 0: raise ValueError( f"rollout batch has future behavior version {behavior_version} > train version {self.train_version}" ) - return lag + return stale def record_optimizer_updates(self, optimizer_updates: int) -> None: self.train_version += optimizer_updates @@ -58,14 +119,17 @@ def launch_slots( num_rollouts: int, hard_boundary: int, ) -> int: - if self.rollout_lag > self.max_policy_lag: + if self.publish_lag > self.staleness_budget: return 0 - freshness = (self.max_policy_lag - self.rollout_lag) // self.num_updates_per_batch + 1 + # Floor division is what keeps a partially-committed step honest: a step + # that skipped updates moved the clock by less than a whole batch, and + # rounding down never admits a generation the budget cannot cover. + freshness = (self.staleness_budget - self.publish_lag) // self.num_updates_per_batch + 1 allowed = min(freshness, num_rollouts - trained_batches, hard_boundary - trained_batches) return max(0, min(max_inflight - inflight_count, allowed - inflight_count - ready_count)) def sync_rollout(self, engine: Any, rollout: Any, weight_sync: Any, *, force: bool = False) -> bool: - if not force and self.rollout_lag == 0: + if not force and self.publish_lag == 0: return False engine.quiesce() if engine.ready_count: @@ -77,16 +141,18 @@ def sync_rollout(self, engine: Any, rollout: Any, weight_sync: Any, *, force: bo self.rollout_version = self.train_version return True - def behavior_metrics(self, behavior_version: int) -> dict[str, int]: + def behavior_metrics(self, behavior_version: int) -> dict[str, float]: + staleness = self.staleness(behavior_version) return { "async/behavior_version": behavior_version, - "async/behavior_lag": self.behavior_lag(behavior_version), + "async/staleness_updates": staleness, + "async/staleness_batches": staleness / self.num_updates_per_batch, } def train_metrics(self, optimizer_updates: int) -> dict[str, int]: return { "async/train_version": self.train_version, - "async/rollout_lag": self.rollout_lag, + "async/publish_lag": self.publish_lag, "async/optimizer_updates": optimizer_updates, } @@ -121,4 +187,9 @@ def unwrap_replicated_int(value: object, *, name: str) -> int: return value -__all__ = ["AsyncBatchControl", "next_hard_boundary", "unwrap_replicated_int"] +__all__ = [ + "AsyncBatchControl", + "next_hard_boundary", + "reject_removed_async_keys", + "unwrap_replicated_int", +] From 512dbfda0d799abbf32f5ea91efacaeaa9b47b54 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 4 Aug 2026 21:15:54 +0800 Subject: [PATCH 7/9] refactor(async): report what the admission numbers actually buy The startup notes on async admission lived in both trainers and covered only two of the three ways a value can mislead. Move them next to the admission math and add the missing case: depth past max_inflight + 1 can never become concurrency, because the loop reaps before it launches and so holds at most one completed batch behind the in-flight ones. The surplus only sets the weight-sync period, which is what bagel's max_staleness=11 is really buying. Drop reject_removed_async_keys with it. A shared blacklist misdescribes two of its three keys: buffer_max_staleness is still the agentic consumer-buffer bound, and weight_sync_interval is still live in 69 recipes on the synchronous path. Only max_policy_lag is genuinely gone, and it never shipped outside this branch. --- unirl/train_async_ar.py | 2 - unirl/train_async_diffusion.py | 2 - unirl/trainer/async_ar.py | 18 ++++---- unirl/trainer/async_diffusion.py | 12 +++--- unirl/trainer/async_policy.py | 71 +++++++++++++++++--------------- 5 files changed, 52 insertions(+), 53 deletions(-) diff --git a/unirl/train_async_ar.py b/unirl/train_async_ar.py index 6e7e4139c..32526a080 100755 --- a/unirl/train_async_ar.py +++ b/unirl/train_async_ar.py @@ -27,12 +27,10 @@ from omegaconf import DictConfig from unirl.trainer.async_ar import AsyncARTrainer -from unirl.trainer.async_policy import reject_removed_async_keys @hydra.main(version_base=None, config_path="../examples", config_name="ar/qwen3_grpo_4b_base_dapo_sglang_async") def main(cfg: DictConfig) -> None: - reject_removed_async_keys(cfg) trainer = AsyncARTrainer( cfg=cfg, batch_size=cfg.batch_size, diff --git a/unirl/train_async_diffusion.py b/unirl/train_async_diffusion.py index 9c4c534d0..a837f6d2c 100755 --- a/unirl/train_async_diffusion.py +++ b/unirl/train_async_diffusion.py @@ -26,12 +26,10 @@ from omegaconf import DictConfig from unirl.trainer.async_diffusion import AsyncDiffusionTrainer -from unirl.trainer.async_policy import reject_removed_async_keys @hydra.main(version_base=None, config_path="../examples", config_name="diffusion/bagel/bagel_vllmomni_async") def main(cfg: DictConfig) -> None: - reject_removed_async_keys(cfg) trainer = AsyncDiffusionTrainer( cfg=cfg, batch_size=cfg.batch_size, diff --git a/unirl/trainer/async_ar.py b/unirl/trainer/async_ar.py index 2dab92236..6dd185c12 100644 --- a/unirl/trainer/async_ar.py +++ b/unirl/trainer/async_ar.py @@ -50,7 +50,12 @@ from unirl.rollout.engine.asynchronous import AsyncBatchRolloutEngine, RolloutBatch from unirl.train.stack import TrainStepResult from unirl.trainer.ar import ARTrainer -from unirl.trainer.async_policy import AsyncBatchControl, next_hard_boundary, unwrap_replicated_int +from unirl.trainer.async_policy import ( + AsyncBatchControl, + log_admission_notes, + next_hard_boundary, + unwrap_replicated_int, +) from unirl.trainer.base import BaseTrainer, build_sampling_dict from unirl.types.sample import Sample from unirl.types.sampling import BaseSamplingParams, total_samples_per_prompt @@ -145,16 +150,7 @@ def __init__( max_staleness=max_staleness, num_updates_per_batch=stack_cfg.get("num_updates_per_batch", 1), ) - if self._max_inflight > self._control.admission_depth: - logger.warning( - "max_inflight=%d exceeds the staleness admission depth %d; the extra concurrency cannot be used", - self._max_inflight, - self._control.admission_depth, - ) - if self._control.max_staleness == 0: - logger.warning( - "max_staleness=0 admits one generation at a time; generation cannot overlap the preceding train batch" - ) + log_admission_notes(self._control, max_inflight=self._max_inflight) self._train_devices = int(round(self.num_devices * self._train_fraction)) if self._train_devices <= 0 or self._train_devices >= self.num_devices: raise ValueError( diff --git a/unirl/trainer/async_diffusion.py b/unirl/trainer/async_diffusion.py index e77b7320f..970666d1a 100644 --- a/unirl/trainer/async_diffusion.py +++ b/unirl/trainer/async_diffusion.py @@ -56,7 +56,12 @@ from unirl.distributed.tensor import hydrate from unirl.rollout.engine.asynchronous import AsyncBatchRolloutEngine, RolloutBatch from unirl.train.stack import TrainStepResult -from unirl.trainer.async_policy import AsyncBatchControl, next_hard_boundary, unwrap_replicated_int +from unirl.trainer.async_policy import ( + AsyncBatchControl, + log_admission_notes, + next_hard_boundary, + unwrap_replicated_int, +) from unirl.trainer.diffusion import DiffusionTrainer from unirl.types.sample import Sample @@ -95,10 +100,7 @@ def __init__( max_staleness=max_staleness, num_updates_per_batch=diffusion_kwargs["stack_cfg"].get("num_updates_per_batch", 1), ) - if self._control.max_staleness == 0: - logger.warning( - "max_staleness=0 admits one generation at a time; generation cannot overlap the preceding train batch" - ) + log_admission_notes(self._control, max_inflight=self._max_inflight) def _build_async_sample(self, gen_id: int) -> Sample: """Consume one data batch and build the request Sample for ``gen_id``.""" diff --git a/unirl/trainer/async_policy.py b/unirl/trainer/async_policy.py index 7c61de97a..fe1fd55fd 100644 --- a/unirl/trainer/async_policy.py +++ b/unirl/trainer/async_policy.py @@ -18,41 +18,11 @@ from __future__ import annotations -from collections.abc import Mapping +import logging from dataclasses import dataclass from typing import Any -_REMOVED_ASYNC_KEYS = { - "max_policy_lag": ( - "max_policy_lag was replaced by max_staleness: the budget now counts whole rollout " - "batches instead of committed optimizer updates. Divide the old value by " - "stack.num_updates_per_batch and rename the key." - ), - "buffer_max_staleness": ( - "buffer_max_staleness (weight-syncs a buffered group could cross, under " - "freshest-first selection) was replaced by max_staleness, which counts whole " - "rollout batches under completion-order FIFO. The units do not correspond; pick a " - "new value rather than carrying the old one over." - ), - "weight_sync_interval": ( - "weight_sync_interval was removed: weight sync is now driven by max_staleness and " - "by eval/checkpoint boundaries, not by a fixed counter." - ), -} - - -def reject_removed_async_keys(cfg: Mapping[str, Any]) -> None: - """Fail a recipe still carrying a removed or re-denominated async knob. - - Every one of these is a silent-wrong hazard rather than a crash: an ignored - key runs at the default, and reading ``max_policy_lag`` as ``max_staleness`` - yields a perfectly valid budget that is ``num_updates_per_batch`` times the - intended one. Nothing downstream can tell either case from a deliberate one. - """ - - for key, message in _REMOVED_ASYNC_KEYS.items(): - if key in cfg: - raise ValueError(message) +logger = logging.getLogger(__name__) @dataclass @@ -157,6 +127,41 @@ def train_metrics(self, optimizer_updates: int) -> dict[str, int]: } +def log_admission_notes(control: AsyncBatchControl, *, max_inflight: int) -> None: + """Report admission settings whose effect differs from what the value suggests. + + All three are legitimate configurations, so none of them is an error; each is + a case where the recipe's number does not buy what its name implies. + """ + + if control.max_staleness == 0: + logger.warning( + "max_staleness=0 admits one generation at a time; generation cannot overlap the preceding train batch" + ) + if max_inflight > control.admission_depth: + logger.warning( + "max_inflight=%d exceeds the staleness admission depth %d; the extra concurrency cannot be used", + max_inflight, + control.admission_depth, + ) + # The loop reaps before it launches, so one completed batch can sit in the + # ready queue behind the in-flight ones; anything past that never becomes + # concurrency, it only defers the sync (which fires after admission_depth + # batches, when publish_lag first exceeds the budget). + usable_depth = max_inflight + 1 + if control.admission_depth > usable_depth: + logger.info( + "max_staleness=%d admits %d outstanding batches but the loop holds at most %d " + "(max_inflight=%d plus one reaped); the surplus does not deepen the pipeline, it " + "sets the weight-sync period to %d batches", + control.max_staleness, + control.admission_depth, + usable_depth, + max_inflight, + control.admission_depth, + ) + + def next_hard_boundary( trained_batches: int, *, @@ -189,7 +194,7 @@ def unwrap_replicated_int(value: object, *, name: str) -> int: __all__ = [ "AsyncBatchControl", + "log_admission_notes", "next_hard_boundary", - "reject_removed_async_keys", "unwrap_replicated_int", ] From 700da05f3c90e2c0c93b66a6976855eec0267805 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 4 Aug 2026 22:49:12 +0800 Subject: [PATCH 8/9] docs(async): scope the staleness budget to batch entry max_staleness is checked when a batch enters training, but the anchor stays frozen while the weights move, so update i of num_updates_per_batch trains at staleness + i - 1 and the worst case any gradient step sees is one short of a full extra batch beyond the budget. That span is the in-batch off-policyness PPO already assumes and the clip range covers. Say so where the two clock quantities are defined, so the budget is not read as a bound on every optimizer step. --- unirl/trainer/async_policy.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/unirl/trainer/async_policy.py b/unirl/trainer/async_policy.py index fe1fd55fd..729a5e9a7 100644 --- a/unirl/trainer/async_policy.py +++ b/unirl/trainer/async_policy.py @@ -2,9 +2,9 @@ Two quantities ride the optimizer-update clock and only one of them is staleness: -* ``staleness`` — updates between the current train weights and the behavior - policy that generated a batch, i.e. the off-policyness of the data actually - being trained on. This is AReaL's ``eta`` / ``max_head_offpolicyness``. +* ``staleness`` — updates between the behavior policy that generated a batch and + the train weights that batch starts training against, i.e. the off-policyness + of the data. This is AReaL's ``eta`` / ``max_head_offpolicyness``. * ``publish_lag`` — updates between the current train weights and the snapshot last published to the rollout engine. This is sync debt: no batch is that stale, but it is what makes a weight sync due. @@ -14,6 +14,13 @@ updates instead would quantize it to ``num_updates_per_batch`` — 22 and 23 would both mean a 12-batch depth — and silently change meaning whenever that count changes. ``staleness_budget`` converts once into the clock the versions count in. + +Batch entry is also the only point the budget is enforced at, which matters once +``num_updates_per_batch > 1``: the anchor is frozen for the whole batch while the +weights keep moving, so update ``i`` trains at ``staleness + i - 1`` and the worst +case any gradient step sees is ``staleness_budget + num_updates_per_batch - 1``. +That extra span is the in-batch off-policyness PPO already assumes — the frozen +anchor and ``clip_range`` cover it — so it is deliberately outside the budget. """ from __future__ import annotations From 1665525ad5be26f57e1a12591034e9c59e97ef92 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 4 Aug 2026 23:25:24 +0800 Subject: [PATCH 9/9] fix(async): report the publication period every admission limit allows log_admission_notes derived the weight-sync period from the staleness budget alone, but next_hard_boundary clamps admission too, so an eval or checkpoint interval below max_staleness + 1 becomes the real period. Measured at max_staleness=11 with eval_interval=10: the run publishes every 10 batches and data never gets past 9 of the configured 11, while the note still reported 12. sync_period_batches now takes the tightest of the three limits, the surplus note quotes it instead of the budget-derived depth, and a budget the boundaries never let you spend gets its own warning. Both trainers report from train() rather than __init__ because save_interval only arrives with the train call; the period also goes to W&B as async config so a run records the cadence it actually ran at. 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. --- unirl/trainer/async_ar.py | 16 +++++++++- unirl/trainer/async_diffusion.py | 16 +++++++++- unirl/trainer/async_policy.py | 50 ++++++++++++++++++++++++++++---- 3 files changed, 75 insertions(+), 7 deletions(-) diff --git a/unirl/trainer/async_ar.py b/unirl/trainer/async_ar.py index 6dd185c12..c98bb4caf 100644 --- a/unirl/trainer/async_ar.py +++ b/unirl/trainer/async_ar.py @@ -54,6 +54,7 @@ AsyncBatchControl, log_admission_notes, next_hard_boundary, + sync_period_batches, unwrap_replicated_int, ) from unirl.trainer.base import BaseTrainer, build_sampling_dict @@ -150,7 +151,6 @@ def __init__( max_staleness=max_staleness, num_updates_per_batch=stack_cfg.get("num_updates_per_batch", 1), ) - log_admission_notes(self._control, max_inflight=self._max_inflight) self._train_devices = int(round(self.num_devices * self._train_fraction)) if self._train_devices <= 0 or self._train_devices >= self.num_devices: raise ValueError( @@ -324,8 +324,22 @@ def train( "max_staleness": self._control.max_staleness, "staleness_budget": self._control.staleness_budget, "num_updates_per_batch": self._control.num_updates_per_batch, + "sync_period_batches": sync_period_batches( + self._control, + eval_interval=self.eval_interval, + save_interval=save_interval, + ), }, ) + # Reported here rather than in __init__: save_interval only arrives with + # the train call, and it clamps the publication period just as the + # staleness budget does. + log_admission_notes( + self._control, + max_inflight=self._max_inflight, + eval_interval=self.eval_interval, + save_interval=save_interval, + ) self._async_engine = AsyncBatchRolloutEngine( self.rollout, diff --git a/unirl/trainer/async_diffusion.py b/unirl/trainer/async_diffusion.py index 970666d1a..f8662b565 100644 --- a/unirl/trainer/async_diffusion.py +++ b/unirl/trainer/async_diffusion.py @@ -60,6 +60,7 @@ AsyncBatchControl, log_admission_notes, next_hard_boundary, + sync_period_batches, unwrap_replicated_int, ) from unirl.trainer.diffusion import DiffusionTrainer @@ -100,7 +101,6 @@ def __init__( max_staleness=max_staleness, num_updates_per_batch=diffusion_kwargs["stack_cfg"].get("num_updates_per_batch", 1), ) - log_admission_notes(self._control, max_inflight=self._max_inflight) def _build_async_sample(self, gen_id: int) -> Sample: """Consume one data batch and build the request Sample for ``gen_id``.""" @@ -188,9 +188,23 @@ def train( "max_staleness": self._control.max_staleness, "staleness_budget": self._control.staleness_budget, "num_updates_per_batch": self._control.num_updates_per_batch, + "sync_period_batches": sync_period_batches( + self._control, + eval_interval=self.eval_interval, + save_interval=save_interval, + ), "train_fraction": self._train_fraction, }, ) + # Reported here rather than in __init__: save_interval only arrives with + # the train call, and it clamps the publication period just as the + # staleness budget does. + log_admission_notes( + self._control, + max_inflight=self._max_inflight, + eval_interval=self.eval_interval, + save_interval=save_interval, + ) self._async_engine = AsyncBatchRolloutEngine( self.rollout, diff --git a/unirl/trainer/async_policy.py b/unirl/trainer/async_policy.py index 729a5e9a7..b846ed21f 100644 --- a/unirl/trainer/async_policy.py +++ b/unirl/trainer/async_policy.py @@ -134,13 +134,41 @@ def train_metrics(self, optimizer_updates: int) -> dict[str, int]: } -def log_admission_notes(control: AsyncBatchControl, *, max_inflight: int) -> None: +def sync_period_batches( + control: AsyncBatchControl, + *, + eval_interval: int = 0, + save_interval: int = 0, +) -> int: + """Batches between weight publications, once every admission limit is applied. + + The staleness budget alone would publish every ``admission_depth`` batches, + but :func:`next_hard_boundary` clamps admission as well, so an eval or + checkpoint interval below that depth becomes the period instead. + """ + + period = control.admission_depth + for interval in (eval_interval, save_interval): + if interval > 0: + period = min(period, interval) + return period + + +def log_admission_notes( + control: AsyncBatchControl, + *, + max_inflight: int, + eval_interval: int = 0, + save_interval: int = 0, +) -> None: """Report admission settings whose effect differs from what the value suggests. - All three are legitimate configurations, so none of them is an error; each is - a case where the recipe's number does not buy what its name implies. + All of these are legitimate configurations, so none is an error; each is a + case where the recipe's number does not buy what its name implies. """ + period = sync_period_batches(control, eval_interval=eval_interval, save_interval=save_interval) + if control.max_staleness == 0: logger.warning( "max_staleness=0 admits one generation at a time; generation cannot overlap the preceding train batch" @@ -153,8 +181,7 @@ def log_admission_notes(control: AsyncBatchControl, *, max_inflight: int) -> Non ) # The loop reaps before it launches, so one completed batch can sit in the # ready queue behind the in-flight ones; anything past that never becomes - # concurrency, it only defers the sync (which fires after admission_depth - # batches, when publish_lag first exceeds the budget). + # concurrency, it only defers the sync. usable_depth = max_inflight + 1 if control.admission_depth > usable_depth: logger.info( @@ -165,7 +192,19 @@ def log_admission_notes(control: AsyncBatchControl, *, max_inflight: int) -> Non control.admission_depth, usable_depth, max_inflight, + period, + ) + if period < control.admission_depth: + logger.warning( + "eval/checkpoint boundaries publish every %d batches, below the %d the staleness " + "budget allows (eval_interval=%d, save_interval=%d), so max_staleness=%d is never " + "fully spent — data tops out at %d batches stale", + period, control.admission_depth, + eval_interval, + save_interval, + control.max_staleness, + period - 1, ) @@ -203,5 +242,6 @@ def unwrap_replicated_int(value: object, *, name: str) -> int: "AsyncBatchControl", "log_admission_notes", "next_hard_boundary", + "sync_period_batches", "unwrap_replicated_int", ]