feat(async): release rollout lanes during remote reward scoring

examples/diffusion/bagel/bagel_it2i_vllmomni_async_remote.yaml

@@ -0,0 +1,269 @@
+# @package _global_
+# BAGEL-7B-MoT image-EDIT async GRPO — vLLM-Omni + remote EditReward.
+#
+# The local 8-GPU pool is split into 4 train GPUs and 4 rollout GPUs. EditReward
+# runs as an external HTTP service on a separate node; the driver owns the thin
+# HTTP client, so no local reward GPU or reward worker is allocated.
+#
+# Each completed prompt group is scored asynchronously:
+#   generate -> release rollout lane -> start remote /score request
+#   -> enqueue scored group -> collect batch_size scored groups -> train.
+#
+# The it2i path is handled in unirl/{models,rollout}/bagel:
+#   - modality bagel_it2i — the adapter ships each sample's SOURCE image on the
+#     request (multi_modal_data["image"]) AND mirrors it onto the deferred
+#     conditions (input_images), because the trainer rebuilds the KV contexts at
+#     replay and must prefill the same source image the worker did.
+#   - RLBagelPipeline builds the three it2i KV contexts inside the worker with the
+#     vendored ImageTransform pair and injects them through vLLM-Omni's KV channel,
+#     so upstream's own img2img prefill never runs. That prefill would (a) override
+#     the output canvas with the resized SOURCE dims — breaking the driver-authored
+#     x_T shape — and (b) squash the source to a fixed 980x980 SigLIP square (~4x the
+#     ViT tokens, aspect ratio dropped) instead of the aspect-preserving navit
+#     transform the trainer replays with. A conditioning mismatch is NOT something
+#     replay can absorb: the reward would be earned under one conditioning while the
+#     gradient is taken under another, and no importance ratio corrects that.
+#   - bundle.enable_vit: true — the und ViT tower, needed on the TRAINER to rebuild
+#     the contexts at replay (the worker always builds its own).
+#   - packed rollout is off for it2i (upstream's grouped generate_image is cfg=1 t2i
+#     only), so every GRPO sibling is its own worker request.
+#
+# Deploy:
+#   reward node: cd /path/to/unirl-reward-service && \
+#     python -m reward_service --config configs/editreward_service.yaml
+#   train node:  export REWARD_SERVICE_URL=http://<reward_node_ip>:8080
+#                export BAGEL_PATH=/path/to/BAGEL-7B-MoT
+#                python -m unirl.train_async_diffusion \
+#                  --config-name diffusion/bagel/bagel_it2i_vllmomni_async_remote
+#
+# Compose check (CPU): python -m unirl.train_async_diffusion \
+#   --config-name diffusion/bagel/bagel_it2i_vllmomni_async_remote --cfg job --resolve
+# Smoke first — it2i is the tightest memory shape in the family (7B MoT + ViT on
+# the trainer plus source-image KV):
+#   override batch_size=8 num_rollouts=2 eval_interval=2 sampling.samples_per_prompt=2
+
+num_devices: 8
+batch_size: 8                  # prompts/rollout (x samples_per_prompt = images/rollout)
+adv_use_global_std: false      # per-group GRPO normalization
+num_rollouts: 10000            # intentionally large; stop manually
+
+# Periodic eval on the eval set (deterministic: eval_eta=0), logged under eval/*.
+# Eval prompts are chunked (eval_chunk_prompts) to stay within the it2i driver-memory
+# budget. Same values as the trainside editreward run.
+eval_interval: 25
+eval_num_prompts: 16
+eval_samples_per_prompt: 4
+eval_chunk_prompts: 8
+eval_eta: 0.0
+
+transport_kind: colocate_store
+workers_per_device: 1
+
+# AsyncDiffusionTrainer requires disjoint train/rollout slabs.
+layout: separate
+train_fraction: 0.5            # 4 train GPUs + 4 rollout GPUs
+reward_fraction: 0.0           # EditReward is external; allocate no local reward GPU
+
+max_inflight: 1                # protects cross-slab transfer and one SDE schedule per batch
+per_worker_inflight: 1         # concurrent prompt units per rollout-engine worker
+weight_sync_interval: 1        # publish the freshly-trained adapter every trained batch
+buffer_max_staleness: 2        # reject rollout groups more than two weight versions old
+
+# Score each completed prompt group immediately through the driver's HTTP client.
+async_reward: true
+reward_client_on_driver: true
+
+# Separate train/rollout slabs and the external reward service remain resident.
+enable_fsdp_offload: false
+offload_train_during_reward: false
+rollout_sleep_after_generate: false
+
+# NB: no `task_config.task: it2i` here. On the trainside recipe that pin is what
+# makes a source-image-less dataset fail loudly; on this path the modality does it —
+# BagelIt2iAdapter.validate_request rejects a Sample without image conditioning.
+
+logging:
+  report_to_wandb: true
+  project_name: ${oc.env:WANDB_PROJECT,bagel-editreward}
+  run_name: bagel_it2i_vllmomni_async_remote_editreward
+  # entity: set via the WANDB_ENTITY env var (wandb picks it up automatically)
+  tags: [bagel, it2i, edit, editreward, async, remote, lora, vllmomni, fp32master]
+  log_media: true              # logs the source|edited side-by-side preview for it2i
+  media_max_items: 8
+  media_log_interval: 5
+
+bundle:
+  _target_: unirl.models.bagel.bundle.BagelBundle.from_config
+  config:
+    _target_: unirl.models.bagel.config.BagelPipelineConfig
+    pretrained_model_ckpt_path: ${oc.env:BAGEL_PATH,ByteDance-Seed/BAGEL-7B-MoT}
+    model_precision: bf16
+    autocast_precision: bf16
+    trajectory_precision: fp32   # keep diffuse<->replay bit-exact; do NOT lower
+    logprob_precision: fp32
+    shift: 3.0
+    use_lora: true               # read by the engine's WeightSync (uses_lora)
+    enable_vit: true             # EDIT: the und ViT the replay context rebuild needs
+
+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                  # COMPUTE dtype (FSDP all-gather + forward/backward)
+    master_dtype: fp32                 # load-bearing: fp32 LoRA master + Adam (reward-collapse fix)
+    cpu_offload: false
+    mixed_precision: true
+    fsdp_mode: full
+    reshard_after_forward: true
+    activation_checkpointing: true     # 7B + multi-step replay: required
+    ac_wrap_order: inside              # recompute re-enters FSDP hooks; the order this recipe's smoke validated
+    use_torch_compile: false
+    root_wrap: false                   # vendored BAGEL calls embed_tokens/lm_head outside a root forward
+  optimizer_cfg:
+    _target_: unirl.train.backend.base.OptimizerConfig
+    learning_rate: 1.0e-4
+    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                           # must be <= the stage YAML's max_lora_rank (64)
+    alpha: 128
+    dropout: 0.0
+    bias: none
+    task_type: FEATURE_EXTRACTION
+    target_modules:                    # = BAGEL_MOE_GEN_LORA_TARGETS (gen experts only; und/ViT frozen)
+      - 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,ByteDance-Seed/BAGEL-7B-MoT}
+    # BAGEL single-stage editing modality (registers BagelIt2iAdapter + boots
+    # stage_configs/bagel_t2i_rl.yaml — one YAML serves both image-out modalities —
+    # with RLBagelPipeline).
+    modality: bagel_it2i
+    # Separate slabs don't time-share GPUs, so sleep/wake is unnecessary.
+    enable_sleep_mode: false
+
+# EditReward is hosted by an external reward-service process. The driver sends
+# bounded source/edit image batches over HTTP; no process, scorer, or lifecycle
+# configuration is managed by this training recipe.
+reward:
+  _target_: unirl.reward.service.RewardService
+  backend:
+    _target_: unirl.reward.remote.RemoteRewardBackend
+    base_device: cpu
+    config:
+      _target_: unirl.reward.remote.RemoteRewardSpec
+      base_url: ${oc.env:REWARD_SERVICE_URL,http://localhost:8080}
+      required_rewards: [editreward]
+      reward_weights: {editreward: 1.0}
+      batch_size: 8
+      request_batch_size: 8
+      timeout: 300.0
+      # EditReward's first column is reward; the second is log uncertainty, so do not average them.
+      sub_metric_reduce: first
+      aggregation_method: weighted_sum
+      input_kind: image
+
+algorithm:
+  _target_: unirl.algorithms.flowgrpo.FlowGRPO
+  stage_attr: diffusion
+  clip_range: 5.0e-3
+  clip_schedule: constant
+  # replay, NOT rollout (which the t2i vllmomni recipe can afford). The worker now
+  # conditions on the SAME source pixels and token geometry as the trainer, but it
+  # still builds the KV context with vLLM's ViT / VAE / MoT kernels while replay uses
+  # the vendored ones — so the emitted rollout log-probs carry a systematic offset.
+  # Anchoring pi_old in the trainer's own forward keeps the ratio well-defined
+  # (== 1 on update 1) instead of starting off it.
+  old_logp_source: replay
+  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                  # navit bs=1; BagelDiffusionConditions.single asserts it
+  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:
+      # Editing instruction + SOURCE image (role: condition -> the request's input image).
+      data_path: ${oc.env:EDIT_DATA_PATH,datasets/image_edit/train.jsonl}
+      eval_data_path: ${oc.env:EDIT_EVAL_DATA_PATH,datasets/image_edit/test.jsonl}
+      seed: 42
+    algorithm:
+      prompts_per_rollout: ${batch_size}
+
+# BAGEL edit sampling (BagelDiffusionParams). Dual-CFG: cfg_text amplifies the
+# INSTRUCTION, cfg_img amplifies the SOURCE (cfg_img active only when cfg_text_scale>1).
+sampling:
+  _target_: unirl.models.bagel.diffusion.BagelDiffusionParams
+  num_inference_steps: 14              # STEPS (σ schedule = steps+1 = 15 points); the adapter sends +1 to the worker
+  guidance_scale: 1.0
+  cfg_text_scale: 1                    # 1 = No-CFG single forward (validated). >1 = +1 forward/step -> more memory
+  cfg_img_scale: 1.0                   # source<->diversity knob; only active when cfg_text_scale>1
+  cfg_interval: [0.4, 1.0]             # CFG schedule; only active when cfg_text_scale>1
+  cfg_renorm_min: 0.0
+  cfg_renorm_type: text_channel        # CFG renorm; only active when cfg_text_scale>1
+  eta: 1.0                             # SDE noise scale (per-step stochasticity = GRPO exploration)
+  samples_per_prompt: 8                # GRPO group size; it2i sends one worker request per sample
+  height: 512                          # the OUTPUT canvas. Fixed square: the driver's x_T recipe carries ONE
+  width: 512                           # shape for the whole request, and the worker honors it via kv_metadata
+  seed: 42
+  init_same_noise: true                # siblings share x_T; in-group diversity comes from the SDE noise (eta>0)
+  trajectory_precision: fp32           # forwarded to the worker scheduler's SDE log-prob storage round-trip
+  scheduler:
+    _target_: unirl.sde.index_schedule.WindowScheduler
+    num_timesteps: 7                   # SDE window drawn from the first half of the 14 inference steps
+    config:
+      _target_: unirl.sde.index_schedule.WindowConfig
+      strategy: random
+      window_size: 3                   # a random contiguous 3-step SDE window per rollout
+
+# 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 publishes on weight_sync_interval or an earlier eval/save boundary.
+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/reward/async_dispatch.py

@@ -0,0 +1,115 @@
+"""Chain rollout completion to asynchronous reward scoring on the driver."""
+
+from __future__ import annotations
+
+import logging
+import threading
+from concurrent.futures import Future, ThreadPoolExecutor
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class _DriverFutureCall:
+    """Adapt a concurrent future to the async manager's call interface."""
+
+    def __init__(self, future: "Future") -> None:
+        self._future = future
+
+    def ready(self) -> bool:
+        return self._future.done()
+
+    def result(self) -> Any:
+        return self._future.result()
+
+
+class DriverRewardClient:
+    """Run a remote HTTP reward client on the driver without a GPU worker."""
+
+    # Handle-compatible shim: one scorer, so trainer DP-geometry sees dp_size=1.
+    dp_size = 1
+    world_size = 1
+
+    def __init__(self, service: Any, *, max_workers: int = 8) -> None:
+        self._service = service
+        self._pool = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="driver-reward")
+
+    def _score(self, sample: Any) -> Any:
+        # Materialize every nested TensorRef leaf on the driver before scoring.
+        from unirl.distributed.tensor.ref import TensorRef, map_tree
+
+        resolved = map_tree(sample, lambda o: o.materialize(backend=None) if isinstance(o, TensorRef) else o)
+        return self._service.score_and_attach(resolved)
+
+    def launch_nowait(self, method_name: str, *args: Any, **kwargs: Any) -> _DriverFutureCall:
+        if method_name != "score_and_attach":
+            raise AttributeError(f"DriverRewardClient only serves score_and_attach, got {method_name!r}")
+        return _DriverFutureCall(self._pool.submit(self._score, *args, **kwargs))
+
+    def score_and_attach(self, sample: Any) -> Any:
+        return self._score(sample)
+
+    def is_available(self) -> bool:
+        return self._service.is_available()
+
+    def offload(self) -> None:
+        pass
+
+    def onload(self) -> None:
+        pass
+
+    def shutdown(self) -> None:
+        # Every submitted score must finish before its backend/session is disposed.
+        self._pool.shutdown(wait=True, cancel_futures=False)
+        dispose = getattr(self._service, "dispose", None)
+        if callable(dispose):
+            dispose()
+
+
+class ChainedRewardCall:
+    """Release a rollout lane after generation while chained reward work continues."""
+
+    def __init__(self, rollout_call: Any, reward: Any) -> None:
+        self._rollout_call = rollout_call
+        self._reward = reward
+        self._reward_call: Optional[Any] = None
+        self._lock = threading.Lock()
+
+    def _start_if_ready(self, *, block: bool) -> bool:
+        with self._lock:
+            if self._reward_call is not None:
+                return True
+            if not block and not self._rollout_call.ready():
+                return False
+            sample = self._rollout_call.result()
+            self._reward_call = self._reward.launch_nowait("score_and_attach", sample)
+            return True
+
+    def is_capacity_released(self) -> bool:
+        """Release rollout capacity and start reward once generation completes."""
+        return self._start_if_ready(block=False)
+
+    def ready(self) -> bool:
+        """True only when the scored Sample can enter the completed queue."""
+        if not self.is_capacity_released():
+            return False
+        return self._reward_call.ready()
+
+    def result(self) -> Any:
+        self._start_if_ready(block=True)
+        return self._reward_call.result()
+
+    def discard_on_completion(self) -> None:
+        """Drain the chain before its reward client is shut down."""
+        try:
+            self.result()
+        except Exception:
+            logger.debug("discarded chained reward call failed during shutdown", exc_info=True)
+
+
+def chain_reward(rollout_call: Any, reward: Any) -> ChainedRewardCall:
+    """Return a future that starts reward as soon as ``rollout_call`` completes."""
+    return ChainedRewardCall(rollout_call, reward)
+
+
+__all__ = ["ChainedRewardCall", "DriverRewardClient", "chain_reward"]