From 1d6823ed8b951bbf314b1ec8068b20b6346bb848 Mon Sep 17 00:00:00 2001 From: zhouzhuoxin Date: Wed, 8 Jul 2026 14:40:29 +0800 Subject: [PATCH 01/10] feat(trainer): add AsyncDiffusionTrainer for disaggregated async diffusion RL Diffusion sibling of AsyncARTrainer: subclasses DiffusionTrainer(layout=separate) to reuse the two-slab build + NCCLWeightSync handshake, and overlays the async rollout buffer loop (non-blocking generate, reap-time reward scoring off the train critical path, buffer of scored GRPO groups, train consumes the freshest batch). Knobs: max_inflight (overlap depth), buffer_max_staleness (0=on-policy). Adds unirl/trainer/async_diffusion.py, unirl/train_async_diffusion.py, and examples/diffusion/sd3/sd3_vllmomni_async.yaml. Purely additive. --- .../diffusion/sd3/sd3_vllmomni_async.yaml | 179 +++++++++ unirl/train_async_diffusion.py | 72 ++++ unirl/trainer/async_diffusion.py | 343 ++++++++++++++++++ 3 files changed, 594 insertions(+) create mode 100644 examples/diffusion/sd3/sd3_vllmomni_async.yaml create mode 100644 unirl/train_async_diffusion.py create mode 100644 unirl/trainer/async_diffusion.py diff --git a/examples/diffusion/sd3/sd3_vllmomni_async.yaml b/examples/diffusion/sd3/sd3_vllmomni_async.yaml new file mode 100644 index 000000000..f8c8a4ce8 --- /dev/null +++ b/examples/diffusion/sd3/sd3_vllmomni_async.yaml @@ -0,0 +1,179 @@ +# @package _global_ +# SD3 ASYNC trainer — for `python -m unirl.train_async_diffusion` (AsyncDiffusionTrainer). +# Async variant of sd3_vllmomni_full_nccl_separate.yaml: generation overlaps training, +# reward is scored off the train critical path. +# - layout: separate — train on `train_fraction` of the pool, rollout on the rest. +# - LoRA training (lora_cfg kept), pushed as a MERGED full model each sync +# (sync.lora_merged: true folds the LoRA delta into the base weights). +# - sync -> NCCLWeightSync (rank-0 broadcasts to all rollout GPUs). +# - rollout enable_sleep_mode: false (separate slabs don't time-share GPUs). +# - max_inflight / buffer_max_staleness: the async overlap knobs (see below). + +num_devices: 8 +batch_size: 48 # prompts_per_rollout +adv_use_global_std: true # advantage: divide by ONE batch-wide std (v1 parity), not per-group + +# Layout: two sibling device slabs. +layout: separate +train_fraction: 0.5 # 4 train GPUs + 4 rollout GPUs on an 8-GPU pool. + +weight_sync_interval: 1 +num_rollouts: 200 # struct key so `num_rollouts=N` overrides cleanly +save_interval: 0 # 0 = no checkpoints (validation run) + +# ---- async knobs (AsyncDiffusionTrainer) ----------------------------------- +# max_inflight: concurrent generations (overlap depth). 1 = one-step pipeline. +# buffer_max_staleness: weight-syncs a buffered group may cross. 0 = on-policy +# (a generation never crosses a weight sync → ratio≈1, sync-separate parity); +# >0 = off-policy continuous buffer. Start at 0 to match the synchronous recipe +# mathematically; the overlap alone hides gen+reward behind the train step. +max_inflight: 2 +buffer_max_staleness: 0 + +bundle: + _target_: unirl.models.sd3.bundle.SD3Bundle.from_config + config: + _target_: unirl.models.sd3.config.SD3PipelineConfig + pretrained_model_ckpt_path: ${oc.env:PRETRAINED_MODEL,stabilityai/stable-diffusion-3.5-medium} + model_precision: bf16 + shift: 3.0 + +pipeline: + _target_: unirl.models.sd3.pipeline.SD3Pipeline + shift: 3.0 + autocast_precision: bf16 + trajectory_precision: bf16 + logprob_precision: fp32 + strategy: + _target_: unirl.sde.kernels.FlowSDEStrategy + +backend: + _target_: unirl.train.backend.fsdp.FSDPBackend + block_class_names: ["JointTransformerBlock"] + trainable_attr: transformer + fsdp_cfg: + _target_: unirl.train.configs.FSDPConfig + param_dtype: bf16 + cpu_offload: false + mixed_precision: true + fsdp_mode: full + reshard_after_forward: false # SD3.5-medium is small (~5 GB bf16 full); keep params gathered, skip bwd re-gather + activation_checkpointing: false + use_torch_compile: false + optimizer_cfg: + _target_: unirl.train.backend.base.OptimizerConfig + learning_rate: 3.0e-4 + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_epsilon: 1.0e-8 + weight_decay: 0.0 + scheduler_cfg: + _target_: unirl.train.backend.base.LrSchedulerConfig + type: constant + warmup_steps: 0 + total_steps: 10000 + # LoRA training — synced as a MERGED full model (sync.lora_merged: true), so the + # rollout engine runs the merged weights without a separate adapter. + lora_cfg: + _target_: unirl.train.configs.LoraConfig + rank: 32 + alpha: 64 + dropout: 0.0 + bias: none + task_type: FEATURE_EXTRACTION + target_modules: + - attn.add_k_proj + - attn.add_q_proj + - attn.add_v_proj + - attn.to_add_out + - attn.to_k + - attn.to_out.0 + - attn.to_q + - attn.to_v + +rollout: + _target_: unirl.rollout.engine.vllm_omni.engine.VLLMOmniRolloutEngine + model_config: ${bundle.config} + config: + _target_: unirl.rollout.engine.vllm_omni.config.VLLMOmniEngineConfig + model_path: ${oc.env:PRETRAINED_MODEL,stabilityai/stable-diffusion-3.5-medium} + modality: sd3_t2i + # Separate slabs do not time-share GPUs, so sleep/wake is unnecessary. + enable_sleep_mode: false + +reward: + _target_: unirl.reward.service.RewardService + backend: + _target_: unirl.reward.local.pickscore.PickScoreRewardScorer + base_device: cuda + config: + _target_: unirl.reward.local.pickscore.PickScoreSpec + batch_size: 8 + device: auto + processor_id: laion/CLIP-ViT-H-14-laion2B-s32B-b79K + model_id: yuvalkirstain/PickScore_v1 + +algorithm: + _target_: unirl.algorithms.flowgrpo.FlowGRPO + stage_attr: diffusion + clip_range: 1.0e-4 + clip_schedule: constant + conditions_cls: + _target_: hydra.utils.get_class + path: unirl.models.sd3.conditions.SD3Conditions + params: ${sampling} + +stack: + _target_: unirl.train.stack.TrainStack + micro_batch_size: 1 + max_grad_norm: 1.0 + num_updates_per_batch: 2 # PPO mini-batches per rollout (π_old frozen once); v1 parity + +data_source: + # Real prompts loaded from file: MultimodalRLDataSource reads TXT/JSON/JSONL + # via data_path and exposes get_samples(bs). (DefaultDataSource returns only + # 8 hardcoded prompts — smoke-test use; training on it over-optimizes/collapses.) + _target_: unirl.data.data_source.MultimodalRLDataSource + args: + run: + data_path: datasets/pickscore/train.txt # relative to repo cwd (Hydra 1.3 no-chdir) + eval_data_path: datasets/pickscore/test.txt + seed: 42 + algorithm: + prompts_per_rollout: ${batch_size} + +sampling: + _target_: unirl.types.sampling.DiffusionSamplingParams + num_inference_steps: 10 + guidance_scale: 1.0 + height: 512 + width: 512 + eta: 0.7 + samples_per_prompt: 16 + seed: 42 + init_same_noise: false + scheduler: + _target_: unirl.utils.scheduler_utils.AllSDEScheduler + num_timesteps: ${..num_inference_steps} + num_sde_steps: 3 + # SDE-step window: [0,0.5] confines SDE noise to the early high-σ steps + # (the validated reproduce setting); [0,1] spreads it across all steps. + timestep_fraction: [0, 0.5] + +# Full base-weight sync → vLLM-Omni rollout (separate-slab NCCL broadcast). +# Wired to the rollout slab by DiffusionTrainer's one-time handshake +# (pick_master → set_rollout_targets → connect). NccclWeightSync takes the +# backend sibling only (rollout is cross-slab). +sync: + _target_: unirl.distributed.weight_sync.full.nccl.NCCLWeightSync + # Fold the trained LoRA delta into the base weights and push the merged full + # model (we LoRA-train but serve a merged model — no separate adapter). + lora_merged: true + group_name: weight_sync + bucket_size_mb: 512 + flush_cache: true + # Ordered first-match-wins name rewrite (see FullWeightSync._apply_name_remap): + # backend.model is the bare SD3 transformer, so the "*" catch-all nests every + # key under the receiver's transformer.* namespace. Applied on both the merged + # (lora_merged=true) and raw paths. + name_remap: {"*": "transformer.*"} diff --git a/unirl/train_async_diffusion.py b/unirl/train_async_diffusion.py new file mode 100644 index 000000000..c62be7720 --- /dev/null +++ b/unirl/train_async_diffusion.py @@ -0,0 +1,72 @@ +#!/usr/bin/env python +"""UniRL async diffusion training entry point (Hydra-native). + +Sibling of ``train_diffusion.py`` that drives +:class:`unirl.trainer.async_diffusion.AsyncDiffusionTrainer` — the disaggregated, +async variant of the diffusion path (training and rollout on DISJOINT GPU slabs, +generation overlapped with training, reward scored off the train critical path, +weights pushed cross-slab via ``NCCLWeightSync``). The synchronous diffusion +trainer is unchanged; this is purely additive. + +Launch (single node): + PRETRAINED_MODEL=/path/to/sd3.5 \ + python -m unirl.train_async_diffusion \ + --config-name=diffusion/sd3/sd3_vllmomni_async num_devices=8 + +Extra config knobs vs the synchronous separate recipe: + * ``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. +``layout`` is forced to ``separate`` (async needs disjoint train/rollout slabs). +""" + +from __future__ import annotations + +import hydra +from omegaconf import DictConfig + +from unirl.trainer.async_diffusion import AsyncDiffusionTrainer + + +@hydra.main(version_base=None, config_path="../examples", config_name="diffusion/sd3/sd3_vllmomni_async") +def main(cfg: DictConfig) -> None: + trainer = AsyncDiffusionTrainer( + cfg=cfg, + batch_size=cfg.batch_size, + bundle_cfg=cfg.bundle, + pipeline_cfg=cfg.pipeline, + backend_cfg=cfg.backend, + rollout_cfg=cfg.rollout, + reward_cfg=cfg.reward, + algorithm_cfg=cfg.algorithm, + stack_cfg=cfg.stack, + data_source_cfg=cfg.data_source, + sampling_cfg=cfg.sampling, + sync_cfg=cfg.get("sync"), + logging_cfg=cfg.get("logging"), + layout="separate", + train_fraction=cfg.get("train_fraction", 0.5), + reward_fraction=cfg.get("reward_fraction", 0.0), + adv_use_global_std=cfg.get("adv_use_global_std", False), + eval_interval=cfg.get("eval_interval", 0), + eval_num_prompts=cfg.get("eval_num_prompts", 64), + eval_samples_per_prompt=cfg.get("eval_samples_per_prompt", 4), + eval_chunk_prompts=cfg.get("eval_chunk_prompts", 16), + eval_cfg_text_scale=cfg.get("eval_cfg_text_scale", 4.0), + eval_eta=cfg.get("eval_eta", 0.0), + stage_config=cfg.get("stage_config"), + max_inflight=int(cfg.get("max_inflight", 1)), + buffer_max_staleness=cfg.get("buffer_max_staleness"), + ) + 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"), + save_mode=cfg.get("save_mode", "auto"), + ) + + +if __name__ == "__main__": + main() diff --git a/unirl/trainer/async_diffusion.py b/unirl/trainer/async_diffusion.py new file mode 100644 index 000000000..daad9cbc3 --- /dev/null +++ b/unirl/trainer/async_diffusion.py @@ -0,0 +1,343 @@ +"""Async diffusion RL trainer — disaggregated train/rollout slabs for DiT. + +Diffusion sibling of :class:`~unirl.trainer.async_ar.AsyncARTrainer`. It subclasses +:class:`~unirl.trainer.diffusion.DiffusionTrainer` with ``layout="separate"`` to +REUSE its two-slab build (train slab + dedicated rollout engine slab), the +``NCCLWeightSync`` cross-slab handshake (``_connect_separate``), and the diffusion +plumbing (``_build_req`` / ``_drop_decoded`` / ``evaluate`` / checkpoint / FlowGRPO +``stack.train_track``). On top of that it overlays the SAME single-threaded async +rollout buffer loop as ``AsyncARTrainer``: + +* Generation is launched as **non-blocking Ray futures** on the rollout slab + (``_generate_async``) and reaped on the driver thread (``_reap_ready``); no + producer thread, no locks. +* Reward is scored the moment a generation completes (``_score_into_buffer``), so + scoring runs on the rollout side and is OFF the train critical path — the buffer + holds already-scored GRPO groups. +* Training consumes the freshest ``batch_size`` groups per step + (``_advantage_and_train``: advantage + FlowGRPO optimizer step); it never calls + the reward. + +Two numeric knobs (identical semantics to AsyncARTrainer): + * ``max_inflight`` — concurrent generations (overlap depth). ``1`` ≈ one-step pipeline. + * ``buffer_max_staleness`` — weight-syncs a buffered group may cross. ``0`` (default) + = on-policy (the launch clamp never lets a generation cross a sync → ``ratio≈1``, + the sync-separate-parity regime). ``>0`` = off-policy continuous buffer. + +Draining all in-flight generations before each weight sync is MANDATORY (a +weight + KV update corrupts an in-flight generation); that is the single-threaded +``_drain_all`` quiesce. + +NOTE: the async buffer/generate-seam machinery below is intentionally a faithful +copy of ``AsyncARTrainer`` (it is engine- and modality-agnostic); a future refactor +could lift it into a shared mixin. Kept self-contained here to leave the validated +AR path untouched. +""" + +from __future__ import annotations + +import logging +import time +from typing import Any, Dict, List, Optional, Tuple + +import ray +import torch + +from unirl.distributed.group.dispatch import DISPATCH_MODE_REGISTRY, Dispatch +from unirl.distributed.tensor import WorkerLocalTransport, hydrate +from unirl.distributed.tensor.pytree import infer_batch_size +from unirl.train.stack import TrainStepResult +from unirl.trainer.diffusion import DiffusionTrainer +from unirl.types.rollout_req import RolloutReq +from unirl.types.rollout_resp import RolloutResp, RolloutTrack + +logger = logging.getLogger(__name__) + + +class _RolloutBuffer: + """Group-keyed rollout buffer (single-threaded; no lock needed). + + Each entry is one prompt's GRPO group — a ``RolloutTrack`` of + ``samples_per_prompt`` already-scored samples — stamped with the + ``weight_version`` it was generated under and a monotonic ``gen_id`` for + freshness ordering. Groups are always complete (the whole ``generate`` + finished before they are ``put``), so there is no partial-group bookkeeping. + """ + + def __init__(self) -> None: + self._items: List[Tuple[RolloutTrack, int, int]] = [] # (group, weight_version, gen_id) + + def put(self, track: RolloutTrack, *, weight_version: int, gen_id: int) -> None: + self._items.append((track, int(weight_version), int(gen_id))) + + def size(self) -> int: + return len(self._items) + + def drain_freshest( + self, + n: int, + *, + current_version: Optional[int] = None, + max_staleness: Optional[int] = None, + ) -> Optional[List[Tuple[RolloutTrack, int, int]]]: + """Pop the ``n`` freshest complete groups, carrying leftovers forward. + + Returns ``None`` if fewer than ``n`` groups remain after eviction. When + ``max_staleness`` is set, groups older than ``current_version - + max_staleness`` weight versions are evicted first (bounded off-policy). + """ + if max_staleness is not None and current_version is not None: + self._items = [it for it in self._items if current_version - it[1] <= max_staleness] + if len(self._items) < n: + return None + self._items.sort(key=lambda it: it[2], reverse=True) # freshest gen_id first + picked, self._items = self._items[:n], self._items[n:] + return picked + + +class AsyncDiffusionTrainer(DiffusionTrainer): + """Disaggregated async diffusion trainer (two slabs, resident engine, NCCL sync).""" + + def __init__( + self, + *, + max_inflight: int = 1, + buffer_max_staleness: Optional[int] = None, + **diffusion_kwargs: Any, + ) -> None: + # Async needs disjoint train/rollout slabs; force the separate layout. + layout = diffusion_kwargs.setdefault("layout", "separate") + if layout != "separate": + raise ValueError(f"AsyncDiffusionTrainer requires layout='separate', got {layout!r}.") + super().__init__(**diffusion_kwargs) + + if self.weight_sync is None: + raise ValueError( + "AsyncDiffusionTrainer requires a cross-slab weight sync (NCCLWeightSync) — " + "add a `sync:` block to the recipe." + ) + + # ---- async state ---- + self._max_inflight = max(1, int(max_inflight)) + self._buffer_max_staleness = buffer_max_staleness + self._weight_version = 0 # driver-tracked policy version (# of weight syncs issued) + # The rollout resp's single track key (e.g. "diffusion"), captured from the + # first reaped generation so the reassembled resp keeps the same key. + self._track_key: str = "diffusion" + + # ------------------------------------------------------------------ + # Non-blocking generate seam (split of the rollout Handle dispatch at ray.get; + # mirrors AsyncARTrainer._generate_async — engine-agnostic). + # ------------------------------------------------------------------ + + def _generate_async(self, req: RolloutReq): + """Launch ``generate`` non-blocking; return (refs, worker_local).""" + r = self.rollout + dispatch_fn = DISPATCH_MODE_REGISTRY[Dispatch.DP_SCATTER]["dispatch_fn"] + bs = infer_batch_size((req,), {}) + if bs is not None and bs % r.dp_size != 0: + raise ValueError(f"req batch_size={bs} not divisible by rollout dp_size={r.dp_size}") + shards = dispatch_fn(r, (req,), {}, bs) + worker_local = issubclass(r.pool.transport_cls, WorkerLocalTransport) + shards = r.pool.transport_cls.localize(shards, r.pool, r.device_ids, r.worker_ids) + refs = r._execute_all("generate", shards, grad_mode=False, call_id=None) + return refs, worker_local + + def _collect_resp(self, refs, worker_local) -> RolloutResp: + """Join a completed generate → full RolloutResp (blocks in ray.get).""" + r = self.rollout + collect_fn = DISPATCH_MODE_REGISTRY[Dispatch.DP_SCATTER]["collect_fn"] + results = ray.get(refs) + results = [r._rebind_tree(x, r.workers[i], worker_local=worker_local) for i, x in enumerate(results)] + return collect_fn(r, results) + + @staticmethod + def _is_ready(refs) -> bool: + """True iff every worker's generate ref is resolved (non-blocking reap).""" + ready, _ = ray.wait(refs, num_returns=len(refs), timeout=0) + return len(ready) == len(refs) + + # ------------------------------------------------------------------ + # In-flight bookkeeping + # ------------------------------------------------------------------ + + def _launch(self, gen_id: int) -> None: + """Build a request and launch one non-blocking generation.""" + req = self._build_req(self.data_source.get_samples(self.batch_size), gen_id) + refs, worker_local = self._generate_async(req) + self._inflight.append( + { + "refs": refs, + "worker_local": worker_local, + "req": req, + "gen_id": gen_id, + "weight_version": self._weight_version, + } + ) + + def _score_into_buffer(self, rec: Dict[str, Any], resp: RolloutResp) -> None: + """Score a completed generation and split its groups into the buffer. + + Scoring (``reward.score_and_attach``) runs here, on the rollout-completion + side — OFF the train critical path. Must precede ``_drop_decoded`` (the + reward reads ``decoded``). + """ + req = rec["req"] + for name, track in list(resp.tracks.items()): + if track.segment is not None: + resp.tracks[name] = self.reward.score_and_attach(req=req, track=track) + self._track_key = next(iter(resp.tracks)) + self._drop_decoded(req, resp, rollout_id=rec["gen_id"]) + (track,) = resp.tracks.values() + for group in track.split(): + self._buffer.put(group, weight_version=rec["weight_version"], gen_id=rec["gen_id"]) + + def _reap_ready(self) -> None: + """Move every completed in-flight generation into the buffer (scored).""" + still: List[Dict[str, Any]] = [] + for rec in self._inflight: + if self._is_ready(rec["refs"]): + self._score_into_buffer(rec, self._collect_resp(rec["refs"], rec["worker_local"])) + else: + still.append(rec) + self._inflight = still + + def _drain_all(self) -> None: + """Finish + buffer EVERY in-flight generation (single-threaded quiesce). + + Mandatory before a weight sync (a weight + KV update corrupts an in-flight + generate), before eval/checkpoint (shared engine), and in ``finally``. + """ + for rec in self._inflight: + self._score_into_buffer(rec, self._collect_resp(rec["refs"], rec["worker_local"])) + self._inflight = [] + + def _next_batch(self, rollout_id: int, interval: int, M: int, stale: int, num_rollouts: int): + """Top up launches, reap completed generations, and return the freshest + ``batch_size`` groups (blocking on the oldest in-flight generation if the + buffer is short). + + The launch clamp is the on-policy guarantee: ``stale=0`` ⇒ never launch + into a future sync-window ⇒ no generation crosses a sync ⇒ ``ratio≈1``. + """ + while True: + staleness_window = ((rollout_id // interval) + 1 + stale) * interval + ceiling = min(num_rollouts, staleness_window) + while self._launch_id < ceiling and len(self._inflight) < M: + self._launch(self._launch_id) + self._launch_id += 1 + + self._reap_ready() + picked = self._buffer.drain_freshest( + self.batch_size, current_version=self._weight_version, max_staleness=stale + ) + if picked is not None: + return picked + if self._inflight: + ray.get(self._inflight[0]["refs"]) # block on oldest; next _reap_ready harvests it + else: + raise RuntimeError("async-diffusion: buffer underflow with no in-flight generations") + + # ------------------------------------------------------------------ + # Train tail (mirrors DiffusionTrainer.train_step's post-generate half: + # advantage → FlowGRPO stack step; reward already attached at reap time). + # ------------------------------------------------------------------ + + def _advantage_and_train( + self, + track: RolloutTrack, + resp: RolloutResp, + *, + training_progress: float, + rollout_id: int, + t0: Optional[float] = None, + ) -> Tuple[TrainStepResult, float]: + """Advantage + optimizer step for a SCORED track (rewards already attached).""" + if t0 is None: + t0 = time.perf_counter() + mean_reward = 0.0 + if track.rewards is not None: + track.rewards = hydrate(track.rewards) + mean_reward = float(track.rewards.to(torch.float32).mean().item()) + track = track.compute_advantages(normalize=True, use_global_std=self._adv_use_global_std) + (name,) = resp.tracks.keys() # single-track diffusion + resp.tracks[name] = track + result = self.stack.train_track(track, training_progress=float(training_progress)) + self.wandb_logger.log_rollout_step(rollout_id, result, resp, step_time_s=time.perf_counter() - t0) + # train_step is bypassed, so BaseTrainer's per-step reset hook never fires; + # reclaim transport buffers here (no-op for colocate_store/gpu). + self._reset_transport_buffers() + return result, mean_reward + + # ------------------------------------------------------------------ + # Train loop + # ------------------------------------------------------------------ + + 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) + # Single-threaded: exactly one get_samples(batch_size) per launch and + # launches are 1:1 with gen_id, so replaying start_rollout times restores + # the exact stream position (deterministic resume). + 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, + "train_fraction": self._train_fraction if hasattr(self, "_train_fraction") else None, + }, + ) + + self._buffer = _RolloutBuffer() + self._inflight: List[Dict[str, Any]] = [] + self._launch_id = start_rollout + + if resumed and self.weight_sync is not None: + self.weight_sync.sync() # push restored weights into the fresh engine + if self.eval_interval > 0: + self.evaluate(start_rollout) # baseline; engine quiescent + + try: + for rollout_id in range(start_rollout, num_rollouts): + t0 = time.perf_counter() + picked = self._next_batch(rollout_id, interval, M, stale, num_rollouts) + track = RolloutTrack.concat([p[0] for p in picked]) + resp = RolloutResp(tracks={self._track_key: track}) + training_progress = rollout_id / max(1, num_rollouts - 1) + result, mean_reward = self._advantage_and_train( + track, resp, training_progress=training_progress, rollout_id=rollout_id, t0=t0 + ) + 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 shares the engine + self.evaluate(step) + if save_interval > 0 and (step % save_interval == 0 or step >= num_rollouts): + self._drain_all() # consistent engine + deterministic resume + self.maybe_save_checkpoint( + 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() # MANDATORY: weight/KV update corrupts in-flight generations + self.weight_sync.sync() + self._weight_version += 1 + finally: + self._drain_all() + self._finish_wandb() From e35bd9cb43bc37117f42c55f637bf22e70ec51ce Mon Sep 17 00:00:00 2001 From: zhouzhuoxin Date: Thu, 9 Jul 2026 16:41:12 +0800 Subject: [PATCH 02/10] Overlap async diffusion rollout with training via reap-before-launch segment transfer, add BAGEL async recipe, drop SD3 async recipe --- .../diffusion/bagel/bagel_vllmomni_async.yaml | 226 ++++++++++++++++++ .../diffusion/sd3/sd3_vllmomni_async.yaml | 179 -------------- unirl/trainer/async_diffusion.py | 10 +- 3 files changed, 235 insertions(+), 180 deletions(-) create mode 100644 examples/diffusion/bagel/bagel_vllmomni_async.yaml delete mode 100644 examples/diffusion/sd3/sd3_vllmomni_async.yaml diff --git a/examples/diffusion/bagel/bagel_vllmomni_async.yaml b/examples/diffusion/bagel/bagel_vllmomni_async.yaml new file mode 100644 index 000000000..375240b95 --- /dev/null +++ b/examples/diffusion/bagel/bagel_vllmomni_async.yaml @@ -0,0 +1,226 @@ +# @package _global_ +# BAGEL-7B-MoT x PickScore x LoRA — vLLM-Omni rollout GRPO recipe. +# +# This is the vllm_omni-rollout sibling of examples/diffusion/bagel/bagel_trainside_lora.yaml. +# It keeps that recipe's model / precision / algorithm / data setup verbatim and +# diverges ONLY on the rollout path: instead of the in-process +# TrainsideRolloutEngine, generation runs in a separate vLLM-Omni worker +# subprocess (the BAGEL single-stage DiT topology), and the freshly-trained LoRA +# is pushed into that worker each rollout via a LocalLoraWeightSync — exactly the +# shape examples/diffusion/sd3_vllmomni.yaml uses for SD3. +# +# What makes BAGEL's vllm_omni path different from SD3's (all handled in +# unirl/rollout/engine/vllm_omni/{adapters,pipelines}/bagel.*): +# - The worker uses RLBagelPipeline, which sets BAGEL's native generate_image +# scheduler to a BagelFlowSDEScheduler (SDE math == trainside FlowSDEStrategy) +# and injects the driver-authored x_T into bagel.prepare_vae_latent. +# - num_inference_steps is sent as steps+1 (BAGEL loops num_timesteps-1) and +# BAGEL builds its own sigma schedule (shift 3.0 == trainside) — the response +# sigma-echo verify asserts it matches the engine-pinned req.sigmas. +# - BAGEL conditioning (opaque KV caches) can't cross the IPC boundary, so the +# adapter ships the PROMPTS and BagelDiffusionStage rebuilds the KV contexts +# trainer-side at replay (the und/text path is frozen → identical contexts). +# +# 16 prompts/rollout, 16 samples/prompt, 2 optimizer updates/rollout, cfg=1. +# +# Launch: +# cd UniRL-main && PYTHONPATH=$PWD python -m unirl.train_diffusion \ +# --config-name diffusion/bagel/bagel_vllmomni --cfg job --resolve + +num_devices: 8 +batch_size: 16 # prompts_per_rollout (one UniRL rollout ~ one flow_grpo epoch) +adv_use_global_std: false # per-group GRPO normalization (mainline default) +num_rollouts: 10000 # intentionally large; stop manually +save_interval: 0 + +transport_kind: colocate_store +workers_per_device: 1 + +# ASYNC: disjoint train/rollout slabs (required for AsyncDiffusionTrainer). Async +# variant of bagel_vllmomni.yaml — same model/data/sampling/algorithm as +# bagel_trainside_lora, but a dedicated vLLM-Omni rollout engine on its own slab, +# generation overlapped with training, LoRA pushed cross-slab via RemoteLoraWeightSync. +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 keeps the +# on-policy ratio ≈ 1 (the rollout SDE math matches trainside, and +# old_logp_source=rollout anchors π_old to the emitted logp) while giving 3 of +# every 4 rollouts a generation overlapped with training. +weight_sync_interval: 4 + +# ---- async knobs (AsyncDiffusionTrainer) ---- +# max_inflight: concurrent generations. MUST be 1: the 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 reaps+transfers each generation in the idle window +# before launching the next, then overlaps that next generation with the train +# step (see AsyncDiffusionTrainer._next_batch reap-before-launch). Real overlap +# needs weight_sync_interval>1 (interval=1 drains every step). +# buffer_max_staleness: 0 = on-policy (a generation never crosses a weight sync, +# ratio≈1). +max_inflight: 1 +buffer_max_staleness: 0 + +logging: + report_to_wandb: false # flip to true to enable wandb (rank-0/driver only) + project_name: unirl + run_name: null + entity: null + tags: null + +bundle: + _target_: unirl.models.bagel.bundle.BagelBundle.from_config + config: + _target_: unirl.models.bagel.config.BagelPipelineConfig + pretrained_model_ckpt_path: ${oc.env:BAGEL_PATH,/root/hf_model/BAGEL-7B-MoT} + model_precision: bf16 + shift: 3.0 + use_lora: true # read by the engine's WeightSync (uses_lora) + +pipeline: + _target_: unirl.models.bagel.pipeline.BagelPipeline + autocast_precision: bf16 + trajectory_precision: fp32 + logprob_precision: fp32 + shift: 3.0 + strategy: + _target_: unirl.sde.kernels.FlowSDEStrategy + +backend: + _target_: unirl.train.backend.fsdp.FSDPBackend + block_class_names: ["Qwen2MoTDecoderLayer"] + trainable_attr: transformer + fsdp_cfg: + _target_: unirl.train.configs.FSDPConfig + param_dtype: bf16 + master_dtype: fp32 # v3: trainable LoRA master + Adam states in fp32 (reward-collapse fix) + cpu_offload: false + mixed_precision: true + fsdp_mode: full + reshard_after_forward: true + activation_checkpointing: true + use_torch_compile: false + root_wrap: false + optimizer_cfg: + _target_: unirl.train.backend.base.OptimizerConfig + learning_rate: 1.0e-4 # lowered from 1.0e-4 (reduce per-update drift / ratio swing) + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_epsilon: 1.0e-8 + weight_decay: 1.0e-4 + scheduler_cfg: + _target_: unirl.train.backend.base.LrSchedulerConfig + type: constant + warmup_steps: 0 + total_steps: 100000 + lora_cfg: + _target_: unirl.train.configs.LoraConfig + rank: 64 + alpha: 128 + dropout: 0.0 + bias: none + task_type: FEATURE_EXTRACTION + target_modules: # = BAGEL_MOE_GEN_LORA_TARGETS (gen experts only) + - self_attn.q_proj_moe_gen + - self_attn.k_proj_moe_gen + - self_attn.v_proj_moe_gen + - self_attn.o_proj_moe_gen + - mlp_moe_gen.gate_proj + - mlp_moe_gen.up_proj + - mlp_moe_gen.down_proj + +rollout: + _target_: unirl.rollout.engine.vllm_omni.engine.VLLMOmniRolloutEngine + # model_config carries the σ-schedule ``shift`` (read by the adapter's + # schedule_policy) and ``use_lora`` (read by WeightSync); point it at the + # bundle's config so train + rollout share one source. + model_config: ${bundle.config} + config: + _target_: unirl.rollout.engine.vllm_omni.config.VLLMOmniEngineConfig + # Required; same checkpoint the bundle loads. + model_path: ${oc.env:BAGEL_PATH,/root/hf_model/BAGEL-7B-MoT} + # BAGEL single-stage T2I diffusion modality (registers BagelT2iAdapter + + # boots stage_configs/bagel_t2i_rl.yaml with RLBagelPipeline). + modality: bagel_t2i + # Separate slabs don't time-share GPUs, so sleep/wake is unnecessary. + enable_sleep_mode: false + +reward: + _target_: unirl.reward.service.RewardService + backend: + _target_: unirl.reward.local.pickscore.PickScoreRewardScorer + base_device: cuda + config: + _target_: unirl.reward.local.pickscore.PickScoreSpec + batch_size: 8 + device: auto + processor_id: ${oc.env:PICKSCORE_PROCESSOR_ID,laion/CLIP-ViT-H-14-laion2B-s32B-b79K} + model_id: ${oc.env:PICKSCORE_MODEL_ID,yuvalkirstain/PickScore_v1} + +algorithm: + _target_: unirl.algorithms.flowgrpo.FlowGRPO + stage_attr: diffusion + clip_range: 1.0e-5 + clip_schedule: constant + # vllm_omni emits per-step log-probs from the worker SDE scheduler, so the + # PPO π_old anchor is the rollout's emitted sde_logp (default). The worker SDE + # math matches the trainside FlowSDEStrategy so the on-policy ratio stays ≈ 1. + old_logp_source: rollout + conditions_cls: + _target_: hydra.utils.get_class + path: unirl.models.bagel.conditions.BagelDiffusionConditions + params: ${sampling} + +stack: + _target_: unirl.train.stack.TrainStack + micro_batch_size: 1 + max_grad_norm: 1.0 + num_updates_per_batch: 2 # 2 optimizer updates/rollout (disjoint mini-batches) + +data_source: + _target_: unirl.data.data_source.MultimodalRLDataSource + args: + run: + data_path: datasets/pickscore/train.txt + eval_data_path: datasets/pickscore/test.txt + seed: 42 + algorithm: + prompts_per_rollout: ${batch_size} + +# Bagel diffusion sampling params (one object shared by trainer expand + diffusion stage). +sampling: + _target_: unirl.models.bagel.diffusion.BagelDiffusionParams + num_inference_steps: 14 # STEPS (σ schedule = steps+1 = 15 points); adapter sends +1 to the worker + guidance_scale: 1.0 + cfg_text_scale: 1.0 # cfg=1 → single-forward "No CFG" path + cfg_img_scale: 1.0 + cfg_interval: [0.0, 1.0] + cfg_renorm_min: 0.0 + cfg_renorm_type: global + eta: 1 # SDE noise scale (= flow_grpo noise_level) + samples_per_prompt: 16 + height: 512 + width: 512 + seed: 42 + init_same_noise: false # per-sample x_T; diverse GRPO groups + trajectory_precision: fp32 # forwarded to the worker scheduler's SDE log-prob storage round-trip + scheduler: + _target_: unirl.utils.scheduler_utils.AllSDEScheduler + num_timesteps: ${..num_inference_steps} + timestep_fraction: [0.0, 0.5] + num_sde_steps: 2 + +# 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. +sync: + _target_: unirl.distributed.weight_sync.lora.RemoteLoraWeightSync + verify: true # checksum read-back asserts the synced LoRA landed; catches a wrong prefix + # Mirrors BagelPipelineConfig.weight_sync_param_name_prefix; must match the + # engine-side LoRA key naming (the trainable module is model.language_model). + param_prefix: "language_model." + adapter_name: default diff --git a/examples/diffusion/sd3/sd3_vllmomni_async.yaml b/examples/diffusion/sd3/sd3_vllmomni_async.yaml deleted file mode 100644 index f8c8a4ce8..000000000 --- a/examples/diffusion/sd3/sd3_vllmomni_async.yaml +++ /dev/null @@ -1,179 +0,0 @@ -# @package _global_ -# SD3 ASYNC trainer — for `python -m unirl.train_async_diffusion` (AsyncDiffusionTrainer). -# Async variant of sd3_vllmomni_full_nccl_separate.yaml: generation overlaps training, -# reward is scored off the train critical path. -# - layout: separate — train on `train_fraction` of the pool, rollout on the rest. -# - LoRA training (lora_cfg kept), pushed as a MERGED full model each sync -# (sync.lora_merged: true folds the LoRA delta into the base weights). -# - sync -> NCCLWeightSync (rank-0 broadcasts to all rollout GPUs). -# - rollout enable_sleep_mode: false (separate slabs don't time-share GPUs). -# - max_inflight / buffer_max_staleness: the async overlap knobs (see below). - -num_devices: 8 -batch_size: 48 # prompts_per_rollout -adv_use_global_std: true # advantage: divide by ONE batch-wide std (v1 parity), not per-group - -# Layout: two sibling device slabs. -layout: separate -train_fraction: 0.5 # 4 train GPUs + 4 rollout GPUs on an 8-GPU pool. - -weight_sync_interval: 1 -num_rollouts: 200 # struct key so `num_rollouts=N` overrides cleanly -save_interval: 0 # 0 = no checkpoints (validation run) - -# ---- async knobs (AsyncDiffusionTrainer) ----------------------------------- -# max_inflight: concurrent generations (overlap depth). 1 = one-step pipeline. -# buffer_max_staleness: weight-syncs a buffered group may cross. 0 = on-policy -# (a generation never crosses a weight sync → ratio≈1, sync-separate parity); -# >0 = off-policy continuous buffer. Start at 0 to match the synchronous recipe -# mathematically; the overlap alone hides gen+reward behind the train step. -max_inflight: 2 -buffer_max_staleness: 0 - -bundle: - _target_: unirl.models.sd3.bundle.SD3Bundle.from_config - config: - _target_: unirl.models.sd3.config.SD3PipelineConfig - pretrained_model_ckpt_path: ${oc.env:PRETRAINED_MODEL,stabilityai/stable-diffusion-3.5-medium} - model_precision: bf16 - shift: 3.0 - -pipeline: - _target_: unirl.models.sd3.pipeline.SD3Pipeline - shift: 3.0 - autocast_precision: bf16 - trajectory_precision: bf16 - logprob_precision: fp32 - strategy: - _target_: unirl.sde.kernels.FlowSDEStrategy - -backend: - _target_: unirl.train.backend.fsdp.FSDPBackend - block_class_names: ["JointTransformerBlock"] - trainable_attr: transformer - fsdp_cfg: - _target_: unirl.train.configs.FSDPConfig - param_dtype: bf16 - cpu_offload: false - mixed_precision: true - fsdp_mode: full - reshard_after_forward: false # SD3.5-medium is small (~5 GB bf16 full); keep params gathered, skip bwd re-gather - activation_checkpointing: false - use_torch_compile: false - optimizer_cfg: - _target_: unirl.train.backend.base.OptimizerConfig - learning_rate: 3.0e-4 - adam_beta1: 0.9 - adam_beta2: 0.999 - adam_epsilon: 1.0e-8 - weight_decay: 0.0 - scheduler_cfg: - _target_: unirl.train.backend.base.LrSchedulerConfig - type: constant - warmup_steps: 0 - total_steps: 10000 - # LoRA training — synced as a MERGED full model (sync.lora_merged: true), so the - # rollout engine runs the merged weights without a separate adapter. - lora_cfg: - _target_: unirl.train.configs.LoraConfig - rank: 32 - alpha: 64 - dropout: 0.0 - bias: none - task_type: FEATURE_EXTRACTION - target_modules: - - attn.add_k_proj - - attn.add_q_proj - - attn.add_v_proj - - attn.to_add_out - - attn.to_k - - attn.to_out.0 - - attn.to_q - - attn.to_v - -rollout: - _target_: unirl.rollout.engine.vllm_omni.engine.VLLMOmniRolloutEngine - model_config: ${bundle.config} - config: - _target_: unirl.rollout.engine.vllm_omni.config.VLLMOmniEngineConfig - model_path: ${oc.env:PRETRAINED_MODEL,stabilityai/stable-diffusion-3.5-medium} - modality: sd3_t2i - # Separate slabs do not time-share GPUs, so sleep/wake is unnecessary. - enable_sleep_mode: false - -reward: - _target_: unirl.reward.service.RewardService - backend: - _target_: unirl.reward.local.pickscore.PickScoreRewardScorer - base_device: cuda - config: - _target_: unirl.reward.local.pickscore.PickScoreSpec - batch_size: 8 - device: auto - processor_id: laion/CLIP-ViT-H-14-laion2B-s32B-b79K - model_id: yuvalkirstain/PickScore_v1 - -algorithm: - _target_: unirl.algorithms.flowgrpo.FlowGRPO - stage_attr: diffusion - clip_range: 1.0e-4 - clip_schedule: constant - conditions_cls: - _target_: hydra.utils.get_class - path: unirl.models.sd3.conditions.SD3Conditions - params: ${sampling} - -stack: - _target_: unirl.train.stack.TrainStack - micro_batch_size: 1 - max_grad_norm: 1.0 - num_updates_per_batch: 2 # PPO mini-batches per rollout (π_old frozen once); v1 parity - -data_source: - # Real prompts loaded from file: MultimodalRLDataSource reads TXT/JSON/JSONL - # via data_path and exposes get_samples(bs). (DefaultDataSource returns only - # 8 hardcoded prompts — smoke-test use; training on it over-optimizes/collapses.) - _target_: unirl.data.data_source.MultimodalRLDataSource - args: - run: - data_path: datasets/pickscore/train.txt # relative to repo cwd (Hydra 1.3 no-chdir) - eval_data_path: datasets/pickscore/test.txt - seed: 42 - algorithm: - prompts_per_rollout: ${batch_size} - -sampling: - _target_: unirl.types.sampling.DiffusionSamplingParams - num_inference_steps: 10 - guidance_scale: 1.0 - height: 512 - width: 512 - eta: 0.7 - samples_per_prompt: 16 - seed: 42 - init_same_noise: false - scheduler: - _target_: unirl.utils.scheduler_utils.AllSDEScheduler - num_timesteps: ${..num_inference_steps} - num_sde_steps: 3 - # SDE-step window: [0,0.5] confines SDE noise to the early high-σ steps - # (the validated reproduce setting); [0,1] spreads it across all steps. - timestep_fraction: [0, 0.5] - -# Full base-weight sync → vLLM-Omni rollout (separate-slab NCCL broadcast). -# Wired to the rollout slab by DiffusionTrainer's one-time handshake -# (pick_master → set_rollout_targets → connect). NccclWeightSync takes the -# backend sibling only (rollout is cross-slab). -sync: - _target_: unirl.distributed.weight_sync.full.nccl.NCCLWeightSync - # Fold the trained LoRA delta into the base weights and push the merged full - # model (we LoRA-train but serve a merged model — no separate adapter). - lora_merged: true - group_name: weight_sync - bucket_size_mb: 512 - flush_cache: true - # Ordered first-match-wins name rewrite (see FullWeightSync._apply_name_remap): - # backend.model is the bare SD3 transformer, so the "*" catch-all nests every - # key under the receiver's transformer.* namespace. Applied on both the merged - # (lora_merged=true) and raw paths. - name_remap: {"*": "transformer.*"} diff --git a/unirl/trainer/async_diffusion.py b/unirl/trainer/async_diffusion.py index daad9cbc3..d2fe5608f 100644 --- a/unirl/trainer/async_diffusion.py +++ b/unirl/trainer/async_diffusion.py @@ -221,13 +221,21 @@ def _next_batch(self, rollout_id: int, interval: int, M: int, stale: int, num_ro into a future sync-window ⇒ no generation crosses a sync ⇒ ``ratio≈1``. """ while True: + # Reap (and cross-slab-transfer the completed generation's segment) + # BEFORE launching the next one. The transfer runs on the rollout + # worker as an NCCL send; if a fresh generation were already queued on + # that worker (launch-first), the send would block behind it (~150s). + # Reaping first gives the transfer an idle-worker window; the launch + # below then starts the NEXT generation, which overlaps the caller's + # train step. Contention-free as long as at most one generation is in + # flight at the transfer instant (max_inflight=1). + self._reap_ready() staleness_window = ((rollout_id // interval) + 1 + stale) * interval ceiling = min(num_rollouts, staleness_window) while self._launch_id < ceiling and len(self._inflight) < M: self._launch(self._launch_id) self._launch_id += 1 - self._reap_ready() picked = self._buffer.drain_freshest( self.batch_size, current_version=self._weight_version, max_staleness=stale ) From 9f42af4f7812f6d4a2668f8161fe434a15621a06 Mon Sep 17 00:00:00 2001 From: zhouzhuoxin Date: Thu, 9 Jul 2026 17:37:10 +0800 Subject: [PATCH 03/10] chore(trainer): make train_async_diffusion.py executable to match train_diffusion.py --- unirl/train_async_diffusion.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) mode change 100644 => 100755 unirl/train_async_diffusion.py diff --git a/unirl/train_async_diffusion.py b/unirl/train_async_diffusion.py old mode 100644 new mode 100755 From ae7b10384632bdcff62eed247f956c38368ec77a Mon Sep 17 00:00:00 2001 From: Jianghai <72591262+CjhHa1@users.noreply.github.com> Date: Mon, 20 Jul 2026 20:28:01 +0800 Subject: [PATCH 04/10] fix(trainer): point async diffusion entry at BAGEL recipe with stale=2 Default Hydra config still pointed at the dropped SD3 async yaml. Point train_async_diffusion at bagel_vllmomni_async and set buffer_max_staleness=2 (the throughput-optimal knob from the PR validation table). --- examples/diffusion/bagel/bagel_vllmomni_async.yaml | 14 +++++++++----- unirl/train_async_diffusion.py | 6 +++--- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/examples/diffusion/bagel/bagel_vllmomni_async.yaml b/examples/diffusion/bagel/bagel_vllmomni_async.yaml index 375240b95..c7e2ccc36 100644 --- a/examples/diffusion/bagel/bagel_vllmomni_async.yaml +++ b/examples/diffusion/bagel/bagel_vllmomni_async.yaml @@ -24,8 +24,9 @@ # 16 prompts/rollout, 16 samples/prompt, 2 optimizer updates/rollout, cfg=1. # # Launch: -# cd UniRL-main && PYTHONPATH=$PWD python -m unirl.train_diffusion \ -# --config-name diffusion/bagel/bagel_vllmomni --cfg job --resolve +# BAGEL_PATH=/path/to/BAGEL-7B-MoT PYTHONPATH=$PWD \ +# python -m unirl.train_async_diffusion \ +# --config-name diffusion/bagel/bagel_vllmomni_async num_devices=8 num_devices: 8 batch_size: 16 # prompts_per_rollout (one UniRL rollout ~ one flow_grpo epoch) @@ -59,10 +60,13 @@ weight_sync_interval: 4 # before launching the next, then overlaps that next generation with the train # step (see AsyncDiffusionTrainer._next_batch reap-before-launch). Real overlap # needs weight_sync_interval>1 (interval=1 drains every step). -# buffer_max_staleness: 0 = on-policy (a generation never crosses a weight sync, -# ratio≈1). +# buffer_max_staleness: how many weight syncs a buffered group may cross. +# 0 = on-policy (never crosses a sync; ~174s/rollout on BAGEL 4+4). +# 2 = throughput-optimal continuous buffer (~148s/rollout, matches vllm +# colocate); sync-boundary cold rollouts disappear. ratio stays ≈1 with +# old_logp_source=rollout (staleness moves gradient freshness, not ratio). max_inflight: 1 -buffer_max_staleness: 0 +buffer_max_staleness: 2 logging: report_to_wandb: false # flip to true to enable wandb (rank-0/driver only) diff --git a/unirl/train_async_diffusion.py b/unirl/train_async_diffusion.py index c62be7720..ef3db2b32 100755 --- a/unirl/train_async_diffusion.py +++ b/unirl/train_async_diffusion.py @@ -9,9 +9,9 @@ trainer is unchanged; this is purely additive. Launch (single node): - PRETRAINED_MODEL=/path/to/sd3.5 \ + BAGEL_PATH=/path/to/BAGEL-7B-MoT \ python -m unirl.train_async_diffusion \ - --config-name=diffusion/sd3/sd3_vllmomni_async num_devices=8 + --config-name=diffusion/bagel/bagel_vllmomni_async num_devices=8 Extra config knobs vs the synchronous separate recipe: * ``max_inflight`` — concurrent generations (overlap depth). ``1`` ≈ one-step pipeline. @@ -28,7 +28,7 @@ from unirl.trainer.async_diffusion import AsyncDiffusionTrainer -@hydra.main(version_base=None, config_path="../examples", config_name="diffusion/sd3/sd3_vllmomni_async") +@hydra.main(version_base=None, config_path="../examples", config_name="diffusion/bagel/bagel_vllmomni_async") def main(cfg: DictConfig) -> None: trainer = AsyncDiffusionTrainer( cfg=cfg, From 726662d068a6532816e23784d0480eb84cb58113 Mon Sep 17 00:00:00 2001 From: aimicahchen Date: Tue, 21 Jul 2026 20:11:04 +0800 Subject: [PATCH 05/10] fix(trainer): keep async diffusion evaluation policy-stable Evaluate the resident rollout policy without syncing or offloading the async engine, while preserving synchronous defaults and forwarding configured eval suites. --- unirl/train_async_diffusion.py | 1 + unirl/trainer/async_diffusion.py | 6 ++++-- unirl/trainer/diffusion.py | 18 +++++++++++++++--- 3 files changed, 20 insertions(+), 5 deletions(-) diff --git a/unirl/train_async_diffusion.py b/unirl/train_async_diffusion.py index ef3db2b32..eed1ded84 100755 --- a/unirl/train_async_diffusion.py +++ b/unirl/train_async_diffusion.py @@ -54,6 +54,7 @@ def main(cfg: DictConfig) -> None: eval_chunk_prompts=cfg.get("eval_chunk_prompts", 16), eval_cfg_text_scale=cfg.get("eval_cfg_text_scale", 4.0), eval_eta=cfg.get("eval_eta", 0.0), + eval_rewards_cfg=cfg.get("eval_rewards"), stage_config=cfg.get("stage_config"), max_inflight=int(cfg.get("max_inflight", 1)), buffer_max_staleness=cfg.get("buffer_max_staleness"), diff --git a/unirl/trainer/async_diffusion.py b/unirl/trainer/async_diffusion.py index d2fe5608f..8f6de550a 100644 --- a/unirl/trainer/async_diffusion.py +++ b/unirl/trainer/async_diffusion.py @@ -319,7 +319,9 @@ def train( if resumed and self.weight_sync is not None: self.weight_sync.sync() # push restored weights into the fresh engine if self.eval_interval > 0: - self.evaluate(start_rollout) # baseline; engine quiescent + # Evaluate the policy already resident on the rollout slab. Eval must + # neither advance the async weight version nor offload this engine. + self.evaluate(start_rollout, sync_weights=False, sleep_after=False) try: for rollout_id in range(start_rollout, num_rollouts): @@ -336,7 +338,7 @@ def train( step = rollout_id + 1 if self.eval_interval > 0 and step % self.eval_interval == 0: self._drain_all() # eval shares the engine - self.evaluate(step) + 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() # consistent engine + deterministic resume self.maybe_save_checkpoint( diff --git a/unirl/trainer/diffusion.py b/unirl/trainer/diffusion.py index cd808088f..ed199ec8d 100644 --- a/unirl/trainer/diffusion.py +++ b/unirl/trainer/diffusion.py @@ -503,7 +503,13 @@ def train_step( self.wandb_logger.log_rollout_step(rollout_id, result, resp, step_time_s=time.perf_counter() - t0) return result, mean_reward - def evaluate(self, step: int) -> float: + def evaluate( + self, + step: int, + *, + sync_weights: bool = True, + sleep_after: bool = True, + ) -> float: """Periodic eval on the eval set (no training); returns the mean reward. Mirrors :meth:`train_step`'s rollout+reward path but skips advantage/backward. @@ -515,6 +521,11 @@ def evaluate(self, step: int) -> float: own-set suite then gets its own generation pass over its own prompts. All means land in one ``eval/*`` row (``eval/reward`` + ``eval/``); returns ``eval/reward``. + + ``sync_weights=False`` evaluates the policy already resident in the + rollout engine without changing its weight version. ``sleep_after=False`` + leaves a dedicated rollout engine resident after evaluation. The defaults + preserve the synchronous trainer's existing behavior. """ # Override only the "diffusion" entry of the modality-keyed sampling dict # (mirrors the AR trainer's evaluate()). ``cfg_text_scale`` only exists @@ -533,7 +544,7 @@ def evaluate(self, step: int) -> float: eval_diffusion = dataclasses.replace(base_diffusion, **replace_kwargs) eval_sp = {**self.sampling_params, "diffusion": eval_diffusion} self.rollout.wake_up() - if self.weight_sync is not None: + if sync_weights and self.weight_sync is not None: self.weight_sync.sync() # Default pass: training reward + shared-set suites score the SAME images. scorers = [("reward", self.reward)] + [(s.name, s.reward) for s in self._eval_suites if s.data_source is None] @@ -542,7 +553,8 @@ def evaluate(self, step: int) -> float: if suite.data_source is not None: n = suite.num_prompts or self.eval_num_prompts metrics.update(self._eval_pass(suite.data_source, n, [(suite.name, suite.reward)], eval_sp, step)) - self.rollout.sleep() + if sleep_after: + self.rollout.sleep() logger.info( "EVAL step %d (%d samples/prompt, cfg=%.1f eta=%.1f) %s", step, From df9018b036dd9a145a4c10675f2221ce36deb804 Mon Sep 17 00:00:00 2001 From: aimicahchen Date: Tue, 21 Jul 2026 20:23:41 +0800 Subject: [PATCH 06/10] docs(trainer): clarify async reward scoring boundary State consistently that generation overlaps training while reap-time reward scoring remains synchronous. --- unirl/train_async_diffusion.py | 6 +++--- unirl/trainer/async_diffusion.py | 12 ++++++------ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/unirl/train_async_diffusion.py b/unirl/train_async_diffusion.py index eed1ded84..f6760bc01 100755 --- a/unirl/train_async_diffusion.py +++ b/unirl/train_async_diffusion.py @@ -4,9 +4,9 @@ Sibling of ``train_diffusion.py`` that drives :class:`unirl.trainer.async_diffusion.AsyncDiffusionTrainer` — the disaggregated, async variant of the diffusion path (training and rollout on DISJOINT GPU slabs, -generation overlapped with training, reward scored off the train critical path, -weights pushed cross-slab via ``NCCLWeightSync``). The synchronous diffusion -trainer is unchanged; this is purely additive. +generation overlapped with training, reward scored synchronously at reap time +rather than overlapped, weights pushed cross-slab via ``NCCLWeightSync``). The +synchronous diffusion trainer is unchanged. Launch (single node): BAGEL_PATH=/path/to/BAGEL-7B-MoT \ diff --git a/unirl/trainer/async_diffusion.py b/unirl/trainer/async_diffusion.py index 8f6de550a..0b1f4eb69 100644 --- a/unirl/trainer/async_diffusion.py +++ b/unirl/trainer/async_diffusion.py @@ -11,9 +11,9 @@ * Generation is launched as **non-blocking Ray futures** on the rollout slab (``_generate_async``) and reaped on the driver thread (``_reap_ready``); no producer thread, no locks. -* Reward is scored the moment a generation completes (``_score_into_buffer``), so - scoring runs on the rollout side and is OFF the train critical path — the buffer - holds already-scored GRPO groups. +* Reward is scored synchronously at reap time (``_score_into_buffer``) before + groups enter the buffer. Generation overlaps training; reward scoring itself + does not. * Training consumes the freshest ``batch_size`` groups per step (``_advantage_and_train``: advantage + FlowGRPO optimizer step); it never calls the reward. @@ -178,9 +178,9 @@ def _launch(self, gen_id: int) -> None: def _score_into_buffer(self, rec: Dict[str, Any], resp: RolloutResp) -> None: """Score a completed generation and split its groups into the buffer. - Scoring (``reward.score_and_attach``) runs here, on the rollout-completion - side — OFF the train critical path. Must precede ``_drop_decoded`` (the - reward reads ``decoded``). + Scoring (``reward.score_and_attach``) is synchronous at reap time, + before the next launch and training-batch consumption. It must precede + ``_drop_decoded`` because the reward reads ``decoded``. """ req = rec["req"] for name, track in list(resp.tracks.items()): From eb6f704e13e0935191eb6394fa1c935ce8583618 Mon Sep 17 00:00:00 2001 From: aimicahchen Date: Tue, 21 Jul 2026 20:24:50 +0800 Subject: [PATCH 07/10] fix(trainer): reject unsupported async diffusion depth Fail before worker construction unless max_inflight is exactly one, preserving the idle-worker window required by reap-time transfer. --- unirl/train_async_diffusion.py | 2 +- unirl/trainer/async_diffusion.py | 12 ++++++++++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/unirl/train_async_diffusion.py b/unirl/train_async_diffusion.py index f6760bc01..287e03bab 100755 --- a/unirl/train_async_diffusion.py +++ b/unirl/train_async_diffusion.py @@ -14,7 +14,7 @@ --config-name=diffusion/bagel/bagel_vllmomni_async num_devices=8 Extra config knobs vs the synchronous separate recipe: - * ``max_inflight`` — concurrent generations (overlap depth). ``1`` ≈ one-step pipeline. + * ``max_inflight`` — must be ``1``; other values fail during trainer initialization. * ``buffer_max_staleness`` — weight-syncs a buffered group may cross. ``0``/unset = on-policy (``ratio≈1``); ``>0`` = off-policy continuous buffer. ``layout`` is forced to ``separate`` (async needs disjoint train/rollout slabs). diff --git a/unirl/trainer/async_diffusion.py b/unirl/trainer/async_diffusion.py index 0b1f4eb69..42a4c583b 100644 --- a/unirl/trainer/async_diffusion.py +++ b/unirl/trainer/async_diffusion.py @@ -19,7 +19,8 @@ the reward. Two numeric knobs (identical semantics to AsyncARTrainer): - * ``max_inflight`` — concurrent generations (overlap depth). ``1`` ≈ one-step pipeline. + * ``max_inflight`` — must be ``1`` so reap-time transfer never competes with + a queued generation on the rollout workers. * ``buffer_max_staleness`` — weight-syncs a buffered group may cross. ``0`` (default) = on-policy (the launch clamp never lets a generation cross a sync → ``ratio≈1``, the sync-separate-parity regime). ``>0`` = off-policy continuous buffer. @@ -109,6 +110,13 @@ def __init__( layout = diffusion_kwargs.setdefault("layout", "separate") if layout != "separate": raise ValueError(f"AsyncDiffusionTrainer requires layout='separate', got {layout!r}.") + max_inflight = int(max_inflight) + if max_inflight != 1: + raise ValueError( + "AsyncDiffusionTrainer requires max_inflight=1: multiple queued generations " + "block the reap-time cross-slab transfer on the rollout workers; " + f"got {max_inflight}." + ) super().__init__(**diffusion_kwargs) if self.weight_sync is None: @@ -118,7 +126,7 @@ def __init__( ) # ---- async state ---- - self._max_inflight = max(1, int(max_inflight)) + self._max_inflight = max_inflight self._buffer_max_staleness = buffer_max_staleness self._weight_version = 0 # driver-tracked policy version (# of weight syncs issued) # The rollout resp's single track key (e.g. "diffusion"), captured from the From bea82dce41773a754f7119199722bffd86127d72 Mon Sep 17 00:00:00 2001 From: aimicahchen Date: Tue, 21 Jul 2026 21:21:24 +0800 Subject: [PATCH 08/10] fix(trainer): align async diffusion policy metadata Record the train slab fraction and describe the actual remote LoRA sync and bounded policy-lag ratio semantics without changing runtime behavior. --- .../diffusion/bagel/bagel_vllmomni_async.yaml | 20 +++++++------- unirl/train_async_diffusion.py | 8 +++--- unirl/trainer/async_diffusion.py | 27 ++++++++++--------- unirl/trainer/diffusion.py | 1 + 4 files changed, 29 insertions(+), 27 deletions(-) diff --git a/examples/diffusion/bagel/bagel_vllmomni_async.yaml b/examples/diffusion/bagel/bagel_vllmomni_async.yaml index c7e2ccc36..86109e83c 100644 --- a/examples/diffusion/bagel/bagel_vllmomni_async.yaml +++ b/examples/diffusion/bagel/bagel_vllmomni_async.yaml @@ -6,8 +6,7 @@ # diverges ONLY on the rollout path: instead of the in-process # TrainsideRolloutEngine, generation runs in a separate vLLM-Omni worker # subprocess (the BAGEL single-stage DiT topology), and the freshly-trained LoRA -# is pushed into that worker each rollout via a LocalLoraWeightSync — exactly the -# shape examples/diffusion/sd3_vllmomni.yaml uses for SD3. +# is pushed cross-slab on the configured cadence via RemoteLoraWeightSync. # # What makes BAGEL's vllm_omni path different from SD3's (all handled in # unirl/rollout/engine/vllm_omni/{adapters,pipelines}/bagel.*): @@ -46,10 +45,10 @@ 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 keeps the -# on-policy ratio ≈ 1 (the rollout SDE math matches trainside, and -# old_logp_source=rollout anchors π_old to the emitted logp) while giving 3 of -# every 4 rollouts a generation overlapped with training. +# 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) ---- @@ -61,10 +60,10 @@ weight_sync_interval: 4 # step (see AsyncDiffusionTrainer._next_batch reap-before-launch). 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 = on-policy (never crosses a sync; ~174s/rollout on BAGEL 4+4). +# 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. ratio stays ≈1 with -# old_logp_source=rollout (staleness moves gradient freshness, not ratio). +# colocate); sync-boundary cold rollouts disappear. The emitted π_old +# remains a valid anchor, while the measured ratio may reflect policy lag. max_inflight: 1 buffer_max_staleness: 2 @@ -171,7 +170,8 @@ algorithm: clip_schedule: constant # vllm_omni emits per-step log-probs from the worker SDE scheduler, so the # PPO π_old anchor is the rollout's emitted sde_logp (default). The worker SDE - # math matches the trainside FlowSDEStrategy so the on-policy ratio stays ≈ 1. + # math matches the trainside FlowSDEStrategy; the ratio then measures drift + # between that resident rollout policy and the current train policy. old_logp_source: rollout conditions_cls: _target_: hydra.utils.get_class diff --git a/unirl/train_async_diffusion.py b/unirl/train_async_diffusion.py index 287e03bab..c321093b8 100755 --- a/unirl/train_async_diffusion.py +++ b/unirl/train_async_diffusion.py @@ -5,8 +5,8 @@ :class:`unirl.trainer.async_diffusion.AsyncDiffusionTrainer` — the disaggregated, async variant of the diffusion path (training and rollout on DISJOINT GPU slabs, generation overlapped with training, reward scored synchronously at reap time -rather than overlapped, weights pushed cross-slab via ``NCCLWeightSync``). The -synchronous diffusion trainer is unchanged. +rather than overlapped, weights pushed via cross-slab weight sync). The synchronous +diffusion trainer is unchanged. Launch (single node): BAGEL_PATH=/path/to/BAGEL-7B-MoT \ @@ -15,8 +15,8 @@ Extra config knobs vs the synchronous separate recipe: * ``max_inflight`` — must be ``1``; other values fail during trainer initialization. - * ``buffer_max_staleness`` — weight-syncs a buffered group may cross. ``0``/unset = - on-policy (``ratio≈1``); ``>0`` = off-policy continuous buffer. + * ``buffer_max_staleness`` — regular rollout-weight syncs a buffered group may + cross. ``0``/unset never crosses a sync; ``>0`` enables bounded policy lag. ``layout`` is forced to ``separate`` (async needs disjoint train/rollout slabs). """ diff --git a/unirl/trainer/async_diffusion.py b/unirl/trainer/async_diffusion.py index 42a4c583b..a74d8f59e 100644 --- a/unirl/trainer/async_diffusion.py +++ b/unirl/trainer/async_diffusion.py @@ -3,10 +3,11 @@ Diffusion sibling of :class:`~unirl.trainer.async_ar.AsyncARTrainer`. It subclasses :class:`~unirl.trainer.diffusion.DiffusionTrainer` with ``layout="separate"`` to REUSE its two-slab build (train slab + dedicated rollout engine slab), the -``NCCLWeightSync`` cross-slab handshake (``_connect_separate``), and the diffusion -plumbing (``_build_req`` / ``_drop_decoded`` / ``evaluate`` / checkpoint / FlowGRPO -``stack.train_track``). On top of that it overlays the SAME single-threaded async -rollout buffer loop as ``AsyncARTrainer``: +cross-slab weight-sync wiring (``RemoteLoraWeightSync`` for the BAGEL recipe; +``NCCLWeightSync`` is also supported by ``_connect_separate``), and the diffusion +plumbing (``_build_req`` / ``_drop_decoded`` / ``evaluate`` / checkpoint / +FlowGRPO ``stack.train_track``). On top of that it overlays the SAME +single-threaded async rollout buffer loop as ``AsyncARTrainer``: * Generation is launched as **non-blocking Ray futures** on the rollout slab (``_generate_async``) and reaped on the driver thread (``_reap_ready``); no @@ -21,9 +22,9 @@ Two numeric knobs (identical semantics to AsyncARTrainer): * ``max_inflight`` — must be ``1`` so reap-time transfer never competes with a queued generation on the rollout workers. - * ``buffer_max_staleness`` — weight-syncs a buffered group may cross. ``0`` (default) - = on-policy (the launch clamp never lets a generation cross a sync → ``ratio≈1``, - the sync-separate-parity regime). ``>0`` = off-policy continuous buffer. + * ``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. Draining all in-flight generations before each weight sync is MANDATORY (a weight + KV update corrupts an in-flight generation); that is the single-threaded @@ -97,7 +98,7 @@ def drain_freshest( class AsyncDiffusionTrainer(DiffusionTrainer): - """Disaggregated async diffusion trainer (two slabs, resident engine, NCCL sync).""" + """Disaggregated async diffusion trainer (two slabs, resident engine, cross-slab sync).""" def __init__( self, @@ -121,8 +122,7 @@ def __init__( if self.weight_sync is None: raise ValueError( - "AsyncDiffusionTrainer requires a cross-slab weight sync (NCCLWeightSync) — " - "add a `sync:` block to the recipe." + "AsyncDiffusionTrainer requires a cross-slab weight sync; add a `sync:` block to the recipe." ) # ---- async state ---- @@ -225,8 +225,9 @@ def _next_batch(self, rollout_id: int, interval: int, M: int, stale: int, num_ro ``batch_size`` groups (blocking on the oldest in-flight generation if the buffer is short). - The launch clamp is the on-policy guarantee: ``stale=0`` ⇒ never launch - into a future sync-window ⇒ no generation crosses a sync ⇒ ``ratio≈1``. + The launch clamp guarantees that ``stale=0`` never launches into a + future sync window, so no generation crosses a regular rollout-weight + sync boundary. """ while True: # Reap (and cross-slab-transfer the completed generation's segment) @@ -316,7 +317,7 @@ def train( "max_inflight": M, "buffer_max_staleness": stale, "weight_sync_interval": interval, - "train_fraction": self._train_fraction if hasattr(self, "_train_fraction") else None, + "train_fraction": self._train_fraction, }, ) diff --git a/unirl/trainer/diffusion.py b/unirl/trainer/diffusion.py index ed199ec8d..22a6ca0ca 100644 --- a/unirl/trainer/diffusion.py +++ b/unirl/trainer/diffusion.py @@ -63,6 +63,7 @@ def __init__( super().__init__(cfg=cfg, logging_cfg=logging_cfg) self.batch_size = batch_size self._layout = str(layout) + self._train_fraction = float(train_fraction) # Colocate memory dance: offload the FSDP train state (params + grads + # optimizer) to CPU during the rollout's generate so a colocate # vLLM/SGLang engine fits, onload before the train backward. Off by From af1907166e2968b7669011fe409af78845c2e00a Mon Sep 17 00:00:00 2001 From: leviking98z-rgb Date: Sun, 26 Jul 2026 15:16:18 +0800 Subject: [PATCH 09/10] fix(trainer): harden async diffusion result handling --- unirl/trainer/async_diffusion.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/unirl/trainer/async_diffusion.py b/unirl/trainer/async_diffusion.py index a74d8f59e..b8c0a1736 100644 --- a/unirl/trainer/async_diffusion.py +++ b/unirl/trainer/async_diffusion.py @@ -39,6 +39,7 @@ from __future__ import annotations import logging +import sys import time from typing import Any, Dict, List, Optional, Tuple @@ -275,6 +276,8 @@ def _advantage_and_train( mean_reward = 0.0 if track.rewards is not None: track.rewards = hydrate(track.rewards) + if isinstance(track.component_rewards, dict): + track.component_rewards = {name: hydrate(value) for name, value in track.component_rewards.items()} mean_reward = float(track.rewards.to(torch.float32).mean().item()) track = track.compute_advantages(normalize=True, use_global_std=self._adv_use_global_std) (name,) = resp.tracks.keys() # single-track diffusion @@ -358,5 +361,13 @@ def train( self.weight_sync.sync() self._weight_version += 1 finally: - self._drain_all() - self._finish_wandb() + # Cleanup failures must not mask the exception that stopped training. + active_exception = sys.exc_info()[0] is not None + try: + self._drain_all() + except Exception: + if not active_exception: + raise + logger.exception("Failed to drain in-flight generations during async diffusion teardown") + finally: + self._finish_wandb() From 10e4c87abae9db8eafeb1f8862b40af78221579e Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Wed, 29 Jul 2026 18:29:41 +0800 Subject: [PATCH 10/10] refactor(trainer): async diffusion on the Sample API and the shared async runtime Two things broke this branch against current main, and both are fixed here. The trainer was written against the retired RolloutReq / RolloutResp / RolloutTrack triplet, deleted by the sample-native rollout boundary (#214). It is now sample-native: the request is the Sample from _build_request_sample, scoring is reward.score_and_attach(sample) on the self-contained filled Sample instead of the old (req=, track=) pair, groups reassemble with Sample.concat, and the RolloutResp(tracks=...) rebuild and its _track_key bookkeeping are gone. The entry point's stage_config was likewise renamed to task_config. The async buffer / generate seam this branch duplicated from AsyncARTrainer has since been lifted into unirl/rollout/async_runtime.py, the follow-up refactor this PR's description anticipated. _RolloutBuffer, _generate_async, _collect_resp, _is_ready, _launch, _reap_ready and the _next_batch loop are all replaced by AsyncRolloutScheduler + RayGenerationDispatcher, leaving only the diffusion hooks: build a request Sample, score-and-split at reap time, and advantage + FlowGRPO step. Adopting that runtime needs one addition to it, because it launched before it reaped and this path requires the opposite. Reaping pulls the trajectory segment off the rollout slab as an NCCL send issued on the rollout workers, so a generation launched ahead of that send blocks it -- the ~150s/rollout instead of ~8s that reap-before-launch was introduced to fix. Reap-first at max_inflight=1 hands the send idle workers while still launching before the step returns, so the next generation overlaps the caller's train step. The new reap_before_launch flag selects the order and defaults to the existing launch-first behavior, so the AR path is unchanged. Verified: ruff check and format clean, the trainer and entry point import against current main, Hydra compose of the BAGEL async recipe passes, all recipe _target_ paths resolve, both constructor guards fire before any Ray construction, and a fake-dispatcher check confirms reap-first at max_inflight=1 both keeps one generation in flight across every train step and always reaps against idle rollout workers. Not re-run: the GPU reward-curve and localize-timing validation in the PR description. --- .../diffusion/bagel/bagel_vllmomni_async.yaml | 16 +- unirl/rollout/async_runtime.py | 52 ++- unirl/train_async_diffusion.py | 2 +- unirl/trainer/README.md | 5 +- unirl/trainer/async_diffusion.py | 322 +++++++----------- 5 files changed, 172 insertions(+), 225 deletions(-) diff --git a/examples/diffusion/bagel/bagel_vllmomni_async.yaml b/examples/diffusion/bagel/bagel_vllmomni_async.yaml index 86109e83c..f567f83ce 100644 --- a/examples/diffusion/bagel/bagel_vllmomni_async.yaml +++ b/examples/diffusion/bagel/bagel_vllmomni_async.yaml @@ -15,7 +15,7 @@ # and injects the driver-authored x_T into bagel.prepare_vae_latent. # - num_inference_steps is sent as steps+1 (BAGEL loops num_timesteps-1) and # BAGEL builds its own sigma schedule (shift 3.0 == trainside) — the response -# sigma-echo verify asserts it matches the engine-pinned req.sigmas. +# sigma-echo verification asserts it matches the diffusion Part's pinned sigmas. # - BAGEL conditioning (opaque KV caches) can't cross the IPC boundary, so the # adapter ships the PROMPTS and BagelDiffusionStage rebuilds the KV contexts # trainer-side at replay (the und/text path is frozen → identical contexts). @@ -52,13 +52,13 @@ train_fraction: 0.5 # 4 train GPUs + 4 rollout GPUs on an 8-GPU pool weight_sync_interval: 4 # ---- async knobs (AsyncDiffusionTrainer) ---- -# max_inflight: concurrent generations. MUST be 1: the 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 reaps+transfers each generation in the idle window -# before launching the next, then overlaps that next generation with the train -# step (see AsyncDiffusionTrainer._next_batch reap-before-launch). Real overlap -# needs weight_sync_interval>1 (interval=1 drains every step). +# 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 shared async runtime (reap_before_launch) reaps and +# transfers each generation in the idle window before launching the next, then +# overlaps that next generation with 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 diff --git a/unirl/rollout/async_runtime.py b/unirl/rollout/async_runtime.py index bb6b9a603..f0eac5ccc 100644 --- a/unirl/rollout/async_runtime.py +++ b/unirl/rollout/async_runtime.py @@ -185,18 +185,31 @@ def collect(self, job: InflightGeneration) -> Sample: class AsyncRolloutScheduler: - """Single-threaded scheduler for complete, versioned rollout groups.""" + """Single-threaded scheduler for complete, versioned rollout groups. + + ``reap_before_launch`` picks the phase order inside :meth:`next_step`. + Launch-first (the default, used by the AR path) keeps the in-flight window as + full as possible. Reap-first instead guarantees that the reap-time work + (``collect`` plus the caller's ``on_complete``) runs while the rollout workers + hold no other queued generation — required when that work pulls a large payload + off those workers, because the transfer would otherwise queue behind a freshly + launched generation. It also keeps ``max_inflight=1`` overlapping: the post-reap + launch is still made before the step returns, so it runs during the caller's + train step. + """ def __init__( self, dispatcher: GenerationDispatcher, *, groups_per_step: int, + reap_before_launch: bool = False, ) -> None: if groups_per_step < 1: raise ValueError(f"groups_per_step must be >= 1, got {groups_per_step}") self._dispatcher = dispatcher self._groups_per_step = groups_per_step + self._reap_before_launch = bool(reap_before_launch) self._buffer = VersionedGroupBuffer() self._inflight: List[InflightGeneration] = [] self._launch_id = 0 @@ -226,6 +239,22 @@ def _launch_one( ) self._launch_id += 1 + def _top_up( + self, + *, + ceiling: int, + max_inflight: int, + build_sample: BuildSample, + current_version: int, + ) -> None: + """Launch generations until the launch ceiling or the in-flight cap binds.""" + + while self._launch_id < ceiling and len(self._inflight) < max_inflight: + self._launch_one( + build_sample=build_sample, + weight_version=current_version, + ) + def _complete( self, job: InflightGeneration, @@ -312,7 +341,9 @@ def next_step( A step is ``groups_per_step`` complete rollout groups. The launch ceiling is the load-bearing on-policy invariant: at ``max_staleness=0`` no - generation is launched into a future weight-sync window. + generation is launched into a future weight-sync window. Whether each + iteration launches or reaps first is fixed by ``reap_before_launch`` (see + the class docstring). ``sync_interval`` and ``max_inflight`` must already be ``>= 1``; callers (e.g. ``AsyncARTrainer``) clamp config before invoking this method. @@ -325,13 +356,16 @@ def next_step( while True: staleness_window = ((rollout_id // sync_interval) + 1 + max_staleness) * sync_interval ceiling = min(num_rollouts, staleness_window) - while self._launch_id < ceiling and len(self._inflight) < max_inflight: - self._launch_one( - build_sample=build_sample, - weight_version=current_version, - ) - - self.reap_ready(on_complete) + if self._reap_before_launch: + self.reap_ready(on_complete) + self._top_up( + ceiling=ceiling, + max_inflight=max_inflight, + build_sample=build_sample, + current_version=current_version, + ) + if not self._reap_before_launch: + self.reap_ready(on_complete) picked = self._buffer.drain_freshest( self._groups_per_step, current_version=current_version, diff --git a/unirl/train_async_diffusion.py b/unirl/train_async_diffusion.py index c321093b8..6b003f0cb 100755 --- a/unirl/train_async_diffusion.py +++ b/unirl/train_async_diffusion.py @@ -55,7 +55,7 @@ def main(cfg: DictConfig) -> None: eval_cfg_text_scale=cfg.get("eval_cfg_text_scale", 4.0), eval_eta=cfg.get("eval_eta", 0.0), eval_rewards_cfg=cfg.get("eval_rewards"), - stage_config=cfg.get("stage_config"), + task_config=cfg.get("task_config"), max_inflight=int(cfg.get("max_inflight", 1)), buffer_max_staleness=cfg.get("buffer_max_staleness"), ) diff --git a/unirl/trainer/README.md b/unirl/trainer/README.md index 59f9010e5..e7c107ef6 100644 --- a/unirl/trainer/README.md +++ b/unirl/trainer/README.md @@ -60,6 +60,7 @@ The current trainer surface is: | `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. | | `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. | | `RewardBackpropTrainer` | differentiable image reward → policy step | ReFL/DRaFT-K path; no rollout engine, `Sample`, GRPO advantage, replay ratio, or weight sync. | @@ -236,7 +237,9 @@ an evaluation and checkpoint fall on the same step, evaluation runs first. - `DiffusionTrainer`, `PETrainer`, `UnifiedModelTrainer`, and `RewardBackpropTrainer` 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. + 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. - 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/async_diffusion.py b/unirl/trainer/async_diffusion.py index b8c0a1736..06e442598 100644 --- a/unirl/trainer/async_diffusion.py +++ b/unirl/trainer/async_diffusion.py @@ -4,36 +4,39 @@ :class:`~unirl.trainer.diffusion.DiffusionTrainer` with ``layout="separate"`` to REUSE its two-slab build (train slab + dedicated rollout engine slab), the cross-slab weight-sync wiring (``RemoteLoraWeightSync`` for the BAGEL recipe; -``NCCLWeightSync`` is also supported by ``_connect_separate``), and the diffusion -plumbing (``_build_req`` / ``_drop_decoded`` / ``evaluate`` / checkpoint / -FlowGRPO ``stack.train_track``). On top of that it overlays the SAME -single-threaded async rollout buffer loop as ``AsyncARTrainer``: - -* Generation is launched as **non-blocking Ray futures** on the rollout slab - (``_generate_async``) and reaped on the driver thread (``_reap_ready``); no - producer thread, no locks. -* Reward is scored synchronously at reap time (``_score_into_buffer``) before - groups enter the buffer. Generation overlaps training; reward scoring itself - does not. -* Training consumes the freshest ``batch_size`` groups per step - (``_advantage_and_train``: advantage + FlowGRPO optimizer step); it never calls - the reward. +``NCCLWeightSync`` is also supported by ``_connect_separate``) and the diffusion +plumbing (``_build_request_sample`` / ``_drop_decoded`` / ``evaluate`` / +checkpoint / FlowGRPO ``stack.train_track``). + +The async loop itself is the shared +:class:`~unirl.rollout.async_runtime.AsyncRolloutScheduler` that ``AsyncARTrainer`` +drives — one single-threaded driver loop over non-blocking Ray dispatch, no +producer thread and no locks. This trainer supplies only the diffusion hooks: + +* ``_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. Two numeric knobs (identical semantics to AsyncARTrainer): - * ``max_inflight`` — must be ``1`` so reap-time transfer never competes with + * ``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 + * ``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. -Draining all in-flight generations before each weight sync is MANDATORY (a -weight + KV update corrupts an in-flight generation); that is the single-threaded -``_drain_all`` quiesce. +The scheduler runs in ``reap_before_launch`` mode, 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 — measured ~150s/rollout on +BAGEL instead of ~8s. Reaping first hands the send idle workers, and the launch +that follows still happens before the step returns, so the next generation +overlaps this step's training. -NOTE: the async buffer/generate-seam machinery below is intentionally a faithful -copy of ``AsyncARTrainer`` (it is engine- and modality-agnostic); a future refactor -could lift it into a shared mixin. Kept self-contained here to leave the validated -AR path untouched. +Draining all in-flight generations before each weight sync is MANDATORY (a +weight + KV update corrupts an in-flight generation); that is the +single-threaded ``_drain_all`` quiesce. """ from __future__ import annotations @@ -41,63 +44,24 @@ import logging import sys import time -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, List, Optional, Tuple -import ray import torch -from unirl.distributed.group.dispatch import DISPATCH_MODE_REGISTRY, Dispatch -from unirl.distributed.tensor import WorkerLocalTransport, hydrate -from unirl.distributed.tensor.pytree import infer_batch_size +from unirl.distributed.tensor import hydrate +from unirl.rollout.async_runtime import ( + AsyncRolloutScheduler, + BufferedRolloutGroup, + InflightGeneration, + RayGenerationDispatcher, +) from unirl.train.stack import TrainStepResult from unirl.trainer.diffusion import DiffusionTrainer -from unirl.types.rollout_req import RolloutReq -from unirl.types.rollout_resp import RolloutResp, RolloutTrack +from unirl.types.sample import Sample logger = logging.getLogger(__name__) -class _RolloutBuffer: - """Group-keyed rollout buffer (single-threaded; no lock needed). - - Each entry is one prompt's GRPO group — a ``RolloutTrack`` of - ``samples_per_prompt`` already-scored samples — stamped with the - ``weight_version`` it was generated under and a monotonic ``gen_id`` for - freshness ordering. Groups are always complete (the whole ``generate`` - finished before they are ``put``), so there is no partial-group bookkeeping. - """ - - def __init__(self) -> None: - self._items: List[Tuple[RolloutTrack, int, int]] = [] # (group, weight_version, gen_id) - - def put(self, track: RolloutTrack, *, weight_version: int, gen_id: int) -> None: - self._items.append((track, int(weight_version), int(gen_id))) - - def size(self) -> int: - return len(self._items) - - def drain_freshest( - self, - n: int, - *, - current_version: Optional[int] = None, - max_staleness: Optional[int] = None, - ) -> Optional[List[Tuple[RolloutTrack, int, int]]]: - """Pop the ``n`` freshest complete groups, carrying leftovers forward. - - Returns ``None`` if fewer than ``n`` groups remain after eviction. When - ``max_staleness`` is set, groups older than ``current_version - - max_staleness`` weight versions are evicted first (bounded off-policy). - """ - if max_staleness is not None and current_version is not None: - self._items = [it for it in self._items if current_version - it[1] <= max_staleness] - if len(self._items) < n: - return None - self._items.sort(key=lambda it: it[2], reverse=True) # freshest gen_id first - picked, self._items = self._items[:n], self._items[n:] - return picked - - class AsyncDiffusionTrainer(DiffusionTrainer): """Disaggregated async diffusion trainer (two slabs, resident engine, cross-slab sync).""" @@ -130,131 +94,41 @@ def __init__( self._max_inflight = max_inflight self._buffer_max_staleness = buffer_max_staleness self._weight_version = 0 # driver-tracked policy version (# of weight syncs issued) - # The rollout resp's single track key (e.g. "diffusion"), captured from the - # first reaped generation so the reassembled resp keeps the same key. - self._track_key: str = "diffusion" # ------------------------------------------------------------------ - # Non-blocking generate seam (split of the rollout Handle dispatch at ray.get; - # mirrors AsyncARTrainer._generate_async — engine-agnostic). + # Generic async-runtime hooks # ------------------------------------------------------------------ - def _generate_async(self, req: RolloutReq): - """Launch ``generate`` non-blocking; return (refs, worker_local).""" - r = self.rollout - dispatch_fn = DISPATCH_MODE_REGISTRY[Dispatch.DP_SCATTER]["dispatch_fn"] - bs = infer_batch_size((req,), {}) - if bs is not None and bs % r.dp_size != 0: - raise ValueError(f"req batch_size={bs} not divisible by rollout dp_size={r.dp_size}") - shards = dispatch_fn(r, (req,), {}, bs) - worker_local = issubclass(r.pool.transport_cls, WorkerLocalTransport) - shards = r.pool.transport_cls.localize(shards, r.pool, r.device_ids, r.worker_ids) - refs = r._execute_all("generate", shards, grad_mode=False, call_id=None) - return refs, worker_local - - def _collect_resp(self, refs, worker_local) -> RolloutResp: - """Join a completed generate → full RolloutResp (blocks in ray.get).""" - r = self.rollout - collect_fn = DISPATCH_MODE_REGISTRY[Dispatch.DP_SCATTER]["collect_fn"] - results = ray.get(refs) - results = [r._rebind_tree(x, r.workers[i], worker_local=worker_local) for i, x in enumerate(results)] - return collect_fn(r, results) - - @staticmethod - def _is_ready(refs) -> bool: - """True iff every worker's generate ref is resolved (non-blocking reap).""" - ready, _ = ray.wait(refs, num_returns=len(refs), timeout=0) - return len(ready) == len(refs) + def _build_async_sample(self, gen_id: int) -> Sample: + """Consume one data batch and build the request Sample for ``gen_id``.""" + return self._build_request_sample(self.data_source.get_samples(self.batch_size), gen_id) - # ------------------------------------------------------------------ - # In-flight bookkeeping - # ------------------------------------------------------------------ - - def _launch(self, gen_id: int) -> None: - """Build a request and launch one non-blocking generation.""" - req = self._build_req(self.data_source.get_samples(self.batch_size), gen_id) - refs, worker_local = self._generate_async(req) - self._inflight.append( - { - "refs": refs, - "worker_local": worker_local, - "req": req, - "gen_id": gen_id, - "weight_version": self._weight_version, - } - ) - - def _score_into_buffer(self, rec: Dict[str, Any], resp: RolloutResp) -> None: - """Score a completed generation and split its groups into the buffer. - - Scoring (``reward.score_and_attach``) is synchronous at reap time, - before the next launch and training-batch consumption. It must precede - ``_drop_decoded`` because the reward reads ``decoded``. + def _score_completed( + self, + job: InflightGeneration, + completed: Sample, + ) -> List[Sample]: + """Score a completed Sample and split it into tree-complete groups. + + Scoring is synchronous at reap time — before the next launch and before + training consumes the batch — and must precede ``_drop_decoded`` (the + reward reads the decoded primitive). Keyed by ``gen_id`` so media panels + behave like the synchronous path. The filled ``Sample`` is self-contained + (it carries its input Parts), so no request handle is kept on the + in-flight record. """ - req = rec["req"] - for name, track in list(resp.tracks.items()): - if track.segment is not None: - resp.tracks[name] = self.reward.score_and_attach(req=req, track=track) - self._track_key = next(iter(resp.tracks)) - self._drop_decoded(req, resp, rollout_id=rec["gen_id"]) - (track,) = resp.tracks.values() - for group in track.split(): - self._buffer.put(group, weight_version=rec["weight_version"], gen_id=rec["gen_id"]) - - def _reap_ready(self) -> None: - """Move every completed in-flight generation into the buffer (scored).""" - still: List[Dict[str, Any]] = [] - for rec in self._inflight: - if self._is_ready(rec["refs"]): - self._score_into_buffer(rec, self._collect_resp(rec["refs"], rec["worker_local"])) - else: - still.append(rec) - self._inflight = still + scored = self.reward.score_and_attach(completed) + self._drop_decoded(scored, rollout_id=job.gen_id) + return scored.split() def _drain_all(self) -> None: - """Finish + buffer EVERY in-flight generation (single-threaded quiesce). + """Finish + buffer EVERY in-flight generation (the single-threaded quiesce). Mandatory before a weight sync (a weight + KV update corrupts an in-flight - generate), before eval/checkpoint (shared engine), and in ``finally``. - """ - for rec in self._inflight: - self._score_into_buffer(rec, self._collect_resp(rec["refs"], rec["worker_local"])) - self._inflight = [] - - def _next_batch(self, rollout_id: int, interval: int, M: int, stale: int, num_rollouts: int): - """Top up launches, reap completed generations, and return the freshest - ``batch_size`` groups (blocking on the oldest in-flight generation if the - buffer is short). - - The launch clamp guarantees that ``stale=0`` never launches into a - future sync window, so no generation crosses a regular rollout-weight - sync boundary. + generate), before eval/checkpoint (shared engine), and in ``finally`` (no + leaked ObjectRefs). """ - while True: - # Reap (and cross-slab-transfer the completed generation's segment) - # BEFORE launching the next one. The transfer runs on the rollout - # worker as an NCCL send; if a fresh generation were already queued on - # that worker (launch-first), the send would block behind it (~150s). - # Reaping first gives the transfer an idle-worker window; the launch - # below then starts the NEXT generation, which overlaps the caller's - # train step. Contention-free as long as at most one generation is in - # flight at the transfer instant (max_inflight=1). - self._reap_ready() - staleness_window = ((rollout_id // interval) + 1 + stale) * interval - ceiling = min(num_rollouts, staleness_window) - while self._launch_id < ceiling and len(self._inflight) < M: - self._launch(self._launch_id) - self._launch_id += 1 - - picked = self._buffer.drain_freshest( - self.batch_size, current_version=self._weight_version, max_staleness=stale - ) - if picked is not None: - return picked - if self._inflight: - ray.get(self._inflight[0]["refs"]) # block on oldest; next _reap_ready harvests it - else: - raise RuntimeError("async-diffusion: buffer underflow with no in-flight generations") + self._async_scheduler.drain_all(self._score_completed) # ------------------------------------------------------------------ # Train tail (mirrors DiffusionTrainer.train_step's post-generate half: @@ -263,27 +137,28 @@ def _next_batch(self, rollout_id: int, interval: int, M: int, stale: int, num_ro def _advantage_and_train( self, - track: RolloutTrack, - resp: RolloutResp, + sample: Sample, *, training_progress: float, rollout_id: int, t0: Optional[float] = None, ) -> Tuple[TrainStepResult, float]: - """Advantage + optimizer step for a SCORED track (rewards already attached).""" + """Advantage + optimizer step for a SCORED ``Sample`` (rewards already attached).""" if t0 is None: t0 = time.perf_counter() + part = sample.parts[-1] mean_reward = 0.0 - if track.rewards is not None: - track.rewards = hydrate(track.rewards) - if isinstance(track.component_rewards, dict): - track.component_rewards = {name: hydrate(value) for name, value in track.component_rewards.items()} - mean_reward = float(track.rewards.to(torch.float32).mean().item()) - track = track.compute_advantages(normalize=True, use_global_std=self._adv_use_global_std) - (name,) = resp.tracks.keys() # single-track diffusion - resp.tracks[name] = track - result = self.stack.train_track(track, training_progress=float(training_progress)) - self.wandb_logger.log_rollout_step(rollout_id, result, resp, step_time_s=time.perf_counter() - t0) + if part.rewards is not None: + # Hydrate in place so the wandb reward/advantage stats reuse this fetch + # instead of re-pulling the TensorRef from the worker. + part.rewards = hydrate(part.rewards) + if isinstance(part.component_rewards, dict): + part.component_rewards = {name: hydrate(value) for name, value in part.component_rewards.items()} + mean_reward = float(part.rewards.to(torch.float32).mean().item()) + 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) # train_step is bypassed, so BaseTrainer's per-step reset hook never fires; # reclaim transport buffers here (no-op for colocate_store/gpu). self._reset_transport_buffers() @@ -324,9 +199,15 @@ def train( }, ) - self._buffer = _RolloutBuffer() - self._inflight: List[Dict[str, Any]] = [] - self._launch_id = start_rollout + # reap_before_launch: 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 (see the module docstring). + self._async_scheduler = AsyncRolloutScheduler( + RayGenerationDispatcher(self.rollout), + groups_per_step=self.batch_size, + reap_before_launch=True, + ) + self._async_scheduler.reset(start_rollout) if resumed and self.weight_sync is not None: self.weight_sync.sync() # push restored weights into the fresh engine @@ -338,12 +219,13 @@ def train( try: for rollout_id in range(start_rollout, num_rollouts): t0 = time.perf_counter() - picked = self._next_batch(rollout_id, interval, M, stale, num_rollouts) - track = RolloutTrack.concat([p[0] for p in picked]) - resp = RolloutResp(tracks={self._track_key: track}) + picked = self._next_step(rollout_id, interval, M, stale, num_rollouts) + # Reassemble the drained per-prompt group Samples into one batched + # Sample [input(P), gen(P*N)] — the inverse of Sample.split. + sample = Sample.concat([item.sample for item in picked]) training_progress = rollout_id / max(1, num_rollouts - 1) result, mean_reward = self._advantage_and_train( - track, resp, training_progress=training_progress, rollout_id=rollout_id, t0=t0 + sample, training_progress=training_progress, rollout_id=rollout_id, t0=t0 ) self.wandb_logger.log_progress(rollout_id, num_rollouts, result, mean_reward, logger=logger) @@ -371,3 +253,31 @@ def train( logger.exception("Failed to drain in-flight generations during async diffusion teardown") finally: self._finish_wandb() + + def _next_step( + self, + rollout_id: int, + interval: int, + M: int, + stale: int, + num_rollouts: int, + ) -> List[BufferedRolloutGroup]: + """Reap completed generations, top up launches, and return the freshest + ``batch_size`` groups 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 regular rollout-weight sync. + """ + return self._async_scheduler.next_step( + rollout_id=rollout_id, + sync_interval=interval, + max_inflight=M, + max_staleness=stale, + num_rollouts=num_rollouts, + current_version=self._weight_version, + build_sample=self._build_async_sample, + on_complete=self._score_completed, + )