feat(trainer): AsyncDiffusionTrainer for disaggregated async diffusion RL

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

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,

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()