diff --git a/examples/diffusion/bagel/bagel_vllmomni_async.yaml b/examples/diffusion/bagel/bagel_vllmomni_async.yaml new file mode 100644 index 000000000..f567f83ce --- /dev/null +++ b/examples/diffusion/bagel/bagel_vllmomni_async.yaml @@ -0,0 +1,230 @@ +# @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 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.*): +# - 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 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). +# +# 16 prompts/rollout, 16 samples/prompt, 2 optimizer updates/rollout, cfg=1. +# +# Launch: +# 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) +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 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) ---- +# 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 +# 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 + +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; 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 + 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/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 new file mode 100755 index 000000000..6b003f0cb --- /dev/null +++ b/unirl/train_async_diffusion.py @@ -0,0 +1,73 @@ +#!/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 synchronously at reap time +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 \ + python -m unirl.train_async_diffusion \ + --config-name=diffusion/bagel/bagel_vllmomni_async num_devices=8 + +Extra config knobs vs the synchronous separate recipe: + * ``max_inflight`` — must be ``1``; other values fail during trainer initialization. + * ``buffer_max_staleness`` — regular rollout-weight syncs a buffered group may + cross. ``0``/unset never crosses a sync; ``>0`` enables bounded policy lag. +``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/bagel/bagel_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), + eval_rewards_cfg=cfg.get("eval_rewards"), + task_config=cfg.get("task_config"), + max_inflight=int(cfg.get("max_inflight", 1)), + buffer_max_staleness=cfg.get("buffer_max_staleness"), + ) + 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/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 new file mode 100644 index 000000000..06e442598 --- /dev/null +++ b/unirl/trainer/async_diffusion.py @@ -0,0 +1,283 @@ +"""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 +cross-slab weight-sync wiring (``RemoteLoraWeightSync`` for the BAGEL recipe; +``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 a reap-time transfer never competes with + a queued generation on the rollout workers. + * ``buffer_max_staleness`` — regular rollout-weight syncs a buffered group may + cross. ``0`` (default) never crosses a sync; ``>0`` enables a bounded + policy-lag buffer. + +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. + +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 + +import logging +import sys +import time +from typing import Any, List, Optional, Tuple + +import torch + +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.sample import Sample + +logger = logging.getLogger(__name__) + + +class AsyncDiffusionTrainer(DiffusionTrainer): + """Disaggregated async diffusion trainer (two slabs, resident engine, cross-slab 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}.") + 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: + raise ValueError( + "AsyncDiffusionTrainer requires a cross-slab weight sync; add a `sync:` block to the recipe." + ) + + # ---- async state ---- + self._max_inflight = max_inflight + self._buffer_max_staleness = buffer_max_staleness + self._weight_version = 0 # driver-tracked policy version (# of weight syncs issued) + + # ------------------------------------------------------------------ + # Generic async-runtime hooks + # ------------------------------------------------------------------ + + 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) + + 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. + """ + 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 (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`` (no + leaked ObjectRefs). + """ + self._async_scheduler.drain_all(self._score_completed) + + # ------------------------------------------------------------------ + # 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, + sample: Sample, + *, + training_progress: float, + rollout_id: int, + t0: Optional[float] = None, + ) -> Tuple[TrainStepResult, float]: + """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 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() + 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, + }, + ) + + # 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 + if self.eval_interval > 0: + # 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): + t0 = time.perf_counter() + 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( + 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) + + 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, 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( + 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: + # 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() + + 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, + ) diff --git a/unirl/trainer/diffusion.py b/unirl/trainer/diffusion.py index d5477f683..aaf57c5d2 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 @@ -519,7 +520,13 @@ def train_step( self.wandb_logger.log_rollout_step(rollout_id, result, sample, 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. @@ -531,6 +538,12 @@ 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, and ``sleep_after=False`` leaves + a dedicated engine resident afterwards — what the async trainer needs so + evaluation does not perturb its pipeline. 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 @@ -550,7 +563,7 @@ def evaluate(self, step: int) -> float: eval_sp = {**self.sampling_params, "diffusion": eval_diffusion} self.rollout.wake_up() try: - 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)] + [ @@ -562,7 +575,8 @@ def evaluate(self, step: int) -> float: 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)) finally: - self.rollout.sleep() + if sleep_after: + self.rollout.sleep() logger.info( "EVAL step %d (%d samples/prompt, cfg=%.1f eta=%.1f) %s", step,