diff --git a/unirl/rollout/README.md b/unirl/rollout/README.md index 7ea2ea69..d06a0478 100644 --- a/unirl/rollout/README.md +++ b/unirl/rollout/README.md @@ -96,7 +96,9 @@ implements its weight-receive method and a matching `sync:` handler in `AsyncBatchRolloutEngine.quiesce()` drains every in-flight generation; a weight + KV update corrupts one mid-flight. The agentic quiesce is a turn-boundary `abort` + final poll, folded into - `AsyncAgenticRolloutEngine.quiesce()`. Reap-vs-launch ordering is trainer + `AsyncAgenticRolloutEngine.quiesce()`; its `sync_weights()` rejects a live + drive, then pairs the weight push with the version bump and logs the sync. + Reap-vs-launch ordering is trainer statement order (diffusion polls before topping up; see its `_next_step`). - **Reward/advantage methods are not engine code** — `Part.compute_advantages` and `Sample.propagate_rewards` are called by the trainer after scoring. An engine diff --git a/unirl/rollout/engine/asynchronous.py b/unirl/rollout/engine/asynchronous.py index 0e103dee..20018a8c 100644 --- a/unirl/rollout/engine/asynchronous.py +++ b/unirl/rollout/engine/asynchronous.py @@ -11,7 +11,7 @@ - :class:`VersionedBuffer` — payload-agnostic freshness/staleness buffer. - :class:`InflightPool` — non-blocking pool of distributed ``generate`` calls. -Engines, sharing one consumer surface (``poll`` / ``drain_freshest`` / +Engines share one consumer surface (``poll`` / ``drain_freshest`` / ``pop_evicted`` / ``quiesce`` + engine-owned ``weight_version``): - :class:`AsyncBatchRolloutEngine` — batch granularity over a single-turn @@ -336,8 +336,19 @@ def __init__(self, rollout: Any, *, group_size: int, start_gen_id: int = 0) -> N def weight_version(self) -> int: return self._weight_version - def bump_weight_version(self) -> int: + def sync_weights(self, weight_sync: Any) -> int: + """Push train weights via *weight_sync* and advance the version ledger. + + The only sanctioned weight-push path — pairing the push with the bump + is what keeps the ledger truthful. Raises while a drive is active (a + weight push must be decode-idle); a joined ``finalize_if_drained`` or + ``quiesce`` ends the drive. + """ + if self._drive_live: + raise RuntimeError("sync_weights with a drive active; finalize or quiesce() first") + weight_sync.sync() self._weight_version += 1 + logger.info("sync_weights: pushed train weights; weight_version -> %d", self._weight_version) return self._weight_version def submit(self, tasks: List["Sample"]) -> None: @@ -380,7 +391,7 @@ def pop_evicted(self) -> List[List["Sample"]]: def quiesce(self) -> List["Sample"]: """Turn-boundary stop: abort, then one final poll for trajectories that completed DURING the quiesce (before the next ``submit`` resets worker - buffers). Call before ``bump_weight_version`` so those groups carry the + buffers). Call before ``sync_weights`` so those groups carry the version they completed under.""" carried = self._rollout.abort()[0] self.poll() diff --git a/unirl/trainer/agentic_async.py b/unirl/trainer/agentic_async.py index 27302191..eb940da8 100644 --- a/unirl/trainer/agentic_async.py +++ b/unirl/trainer/agentic_async.py @@ -21,8 +21,8 @@ ``buffer_max_staleness``), reward + GRPO advantage + one optimizer step (reusing :class:`AgenticTrainer`'s helpers), then **quiesce + sync**: ``abort`` the in-flight tail at a turn boundary, apply the configured ``tail_policy`` (carry only when the - environment can resume from the ``Sample``; otherwise drop), ``weight_sync.sync()``, - bump the version. + environment can resume from the ``Sample``; otherwise drop), then + ``engine.sync_weights`` (one call: push + version bump). ONE single-threaded loop (the ``AsyncARTrainer`` shape): with disjoint slabs the rollout slab keeps generating in the background (the engine's per-worker drain) while @@ -356,7 +356,7 @@ def train( self._engine = AsyncAgenticRolloutEngine(self.rollout, group_size=self._n, start_gen_id=start_rollout) if start_rollout < num_rollouts and start_rollout and self.weight_sync is not None: - self.weight_sync.sync() + self._engine.sync_weights(self.weight_sync) # push restored weights into the fresh engine if start_rollout < num_rollouts: self._submit_drive(carried=[], rollout_id=start_rollout) @@ -387,8 +387,7 @@ def train( save_mode=save_mode, ) if need_sync: - self.weight_sync.sync() - self._engine.bump_weight_version() + self._engine.sync_weights(self.weight_sync) if step < num_rollouts: self._submit_drive(carried=carried, rollout_id=step) finally: diff --git a/unirl/trainer/agentic_partial.py b/unirl/trainer/agentic_partial.py index 4516b887..e2b5d003 100644 --- a/unirl/trainer/agentic_partial.py +++ b/unirl/trainer/agentic_partial.py @@ -155,8 +155,7 @@ def _collect_until(self, batch_size: int, rollout_id: int, stale: int) -> List[L def _drive_partial(self, rollout_id: int, sync_weights: bool, stale: int) -> List[List[Sample]]: self.rollout.wake_up() if sync_weights and self.weight_sync is not None: - self.weight_sync.sync() - self._engine.bump_weight_version() + self._engine.sync_weights(self.weight_sync) tasks = self._build_tasks(self._carried, rollout_id) self._carried = [] self._engine.submit(tasks) diff --git a/unirl/trainer/diffusion.py b/unirl/trainer/diffusion.py index aeba8dd1..2344bf07 100644 --- a/unirl/trainer/diffusion.py +++ b/unirl/trainer/diffusion.py @@ -388,8 +388,9 @@ def evaluate( Mirrors :meth:`train_step`'s rollout+reward path but skips advantage/backward. Generates at the deterministic best-quality setting (``cfg_text_scale= - eval_cfg_text_scale``, ``eta=eval_eta``; ``eval_samples_per_prompt`` x_T per - prompt) and scores. The training reward plus every shared-set + eval_cfg_text_scale``, ``eta=eval_eta`` — at ``eval_eta=0`` the SDE gate + is also cleared, so the request is pure ODE; ``eval_samples_per_prompt`` + x_T per prompt) and scores. The training reward plus every shared-set ``eval_rewards`` suite scores the SAME generated images over the default eval set (``run.eval_data_path``, ``eval_num_prompts`` prompts); each own-set suite then gets its own generation pass over its own prompts. @@ -407,6 +408,11 @@ def evaluate( samples_per_prompt=self.eval_samples_per_prompt, eta=self.eval_eta, ) + if self.eval_eta <= 0.0: + # Deterministic eval must also clear the SDE gate: eta=0 with gated + # steps is a contradictory request — the central kernel degrades such + # steps to ODE, but worker-resident schedulers (BAGEL) refuse the pair. + replace_kwargs.update(sde_indices=[], scheduler=None) if "cfg_text_scale" in {f.name for f in dataclasses.fields(base_diffusion)}: replace_kwargs["cfg_text_scale"] = self.eval_cfg_text_scale else: