feat(trainer): AsyncDiffusionTrainer for disaggregated async diffusion RL

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

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