feat(trainer): AsyncDiffusionTrainer for disaggregated async diffusion RL

[zzhuoxin1508]

Maintainer Update (2026-07-29) — migrated to sample-native main

Rebased onto current main. Two APIs this branch was written against no longer exist, so mergeable=MERGEABLE was textual only — the previous head raised ImportError on current main.

• The retired RolloutReq / RolloutResp / RolloutTrack triplet, deleted by the sample-native rollout boundary (feat(agentic): add Sample-native multi-turn rollout and training #214). The trainer 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 plus its _track_key bookkeeping are gone. The Hydra entry's stage_config became task_config.
• The duplicated async buffer / generate seam, which feat(agentic): add Sample-native multi-turn rollout and training #214 also lifted into unirl/rollout/async_runtime.py — the follow-up refactor this description anticipated. AsyncRolloutScheduler + RayGenerationDispatcher now own the loop, so _RolloutBuffer, _generate_async, _collect_resp, _is_ready, _launch, _reap_ready and _next_batch are gone and async_diffusion.py drops from 373 to 283 lines.

Adopting that runtime needed one addition to it: it launched before it reaped, which is exactly the ordering this path cannot use (see the key-mechanism section below). reap_before_launch selects the order and defaults to the existing launch-first behavior, so AsyncARTrainer is unchanged.

Current-head validation: ruff check / ruff format, full repository pre-commit, all recipe _target_ paths resolve, the trainer and entry point import against current main, BAGEL async Hydra compose passes, 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. The GPU numbers below were measured on the pre-migration head and have not been re-run.

Maintainer Update (2026-07-26)

Final review aligned the async-only path with the current synchronous and AsyncAR safety behavior:

• component reward tensors are hydrated before rollout metrics are computed;
• teardown drain failures no longer mask the original training exception, while standalone drain failures still propagate;
• no test file was added.

Current-head validation: non-persisted component-hydration and fault-injection smoke passed; BAGEL async Hydra compose passed; full repository pre-commit and git diff --check passed; synthetic merge with latest main is clean. The prior Tencent external-scan failure contained a literal null security count and an empty task ID rather than a reported defect; the new head has retriggered the checks.

Summary

AsyncDiffusionTrainer — the pure-DiT async RL trainer. Diffusion sibling of AsyncARTrainer: training and rollout run on disjoint GPU slabs, generation is overlapped with training, and weights are pushed cross-slab. Completed generations are synchronously localized and reward-scored at reap time before the next launch / training consume; reward itself is not overlapped.

The synchronous diffusion trainer's behavior is unchanged. The PR adds three async-specific files, one ordering option on the shared async rollout runtime, and a backward-compatible evaluation control seam in DiffusionTrainer.

Motivation: the synchronous diffusion path runs generate → reward → train in series each step. Async removes generation from the train critical path by running the next rollout concurrently with the current train step.

The key mechanism — reap-before-launch (what makes the overlap actually fast)

The buffer loop reaps (and cross-slab-transfers the completed generation's trajectory segment) BEFORE launching the next generation. That transfer runs on the rollout worker as an NCCL send; if a fresh generation were already queued on that worker (launch-first), the send blocks behind it — measured ~150s/rollout on BAGEL, even though the send itself is only ~3–8s. Reaping first gives the transfer an idle-worker window (~7–8s), then the next generation overlaps the caller's train step.

This requires max_inflight=1: a second in-flight generation co-tenants the same rollout workers and reintroduces the stall. On shared rollout workers generations serialize anyway, so max_inflight=1 costs no throughput while enabling the contention-free transfer + overlap. (This is also why plain vllm sync-separate never hit the stall — it never runs a generation concurrently with the transfer.)

This ordering now lives in the shared runtime as AsyncRolloutScheduler(reap_before_launch=True). It is load-bearing rather than cosmetic: launch-first at max_inflight=1 leaves nothing in flight when the step returns, so the trainer would still be correct and still pass every static check while silently running at vllm sync-separate speed.

What's changed

• unirl/trainer/async_diffusion.py — AsyncDiffusionTrainer(DiffusionTrainer). Reuses the layout="separate" two-slab build, cross-slab weight-sync wiring, _build_request_sample / _drop_decoded / checkpoint, and FlowGRPO stack.train_track. Drives the shared AsyncRolloutScheduler and supplies only the diffusion hooks: _build_async_sample (one data batch → one request Sample), _score_completed (reap-time reward, then Sample.split into tree-complete groups), _advantage_and_train, and _drain_all quiescence before weight sync / eval / checkpoint.
• unirl/rollout/async_runtime.py — reap_before_launch on AsyncRolloutScheduler, selecting whether each step launches or reaps first. Defaults to the existing launch-first order, so the AR path is byte-equivalent; the launch top-up is factored into _top_up so both orders share it.
• unirl/train_async_diffusion.py — Hydra entry (sibling of train_diffusion.py).
• examples/diffusion/bagel/bagel_vllmomni_async.yaml — BAGEL-7B-MoT async recipe (max_inflight=1, weight_sync_interval=4, buffer_max_staleness=2).
• unirl/trainer/diffusion.py — backward-compatible eval flags. Async eval uses the policy already resident in the rollout engine, does not sync train weights, and leaves the dedicated engine resident; synchronous callers retain the existing sync + sleep defaults.

Knobs: max_inflight (must be 1, see above); weight_sync_interval (>1 enables overlap across train steps); buffer_max_staleness (0 = buffered groups do not cross a regular rollout-weight sync; >0 = bounded policy-lag buffer that can survive sync boundaries).

Eval semantics

Periodic async eval measures the currently resident rollout policy. It deliberately skips train → rollout weight sync and does not sleep/offload the rollout engine, so evaluation neither changes _weight_version nor perturbs the resident async pipeline.

Validation (BAGEL-7B-MoT FlowGRPO, PickScore, 4 train + 4 rollout on 8×H20, batch=16 × 16 samples/prompt)

Measured on the pre-migration head; see the 2026-07-29 maintainer update.

• ratio = 1.0000 observed throughout the validation run; reward grows, no crash / OOM.
• reward/localize (cross-slab trajectory-segment transfer): 151s → ~7.5s — the reap-before-launch fix.
• Per-rollout ~148s on overlap rollouts.

config	reward/localize	per-rollout	vs vllm sync-separate
sync trainside colocate	~2s	~135s	2.04×
vllm colocate	—	~148s	1.86×
vllm sync-separate	2.6s	~276s	1.0×
async inflight=1 stale=0 interval=4	7.3s	~174s	1.59×
async inflight=1 stale=2 interval=4	7.6s	~148s	1.86×

• staleness=2 eliminates the per-window cold sync-boundary rollout (276s → 110s) → ~148s avg, matching vllm colocate at 1.85×.
• old_logp_source=rollout preserves the generating policy's emitted π_old for importance sampling. It does not mathematically force ratio=1; the near-1 ratio above is an empirical result of this run under bounded policy lag.
• Multi-node: the contention fix is node-agnostic (it only reorders rollout-worker scheduling). The cross-node segment transfer over IB was measured at 2.8 GB/s (~0.5s for the ~1–2GB segment) via the same ProcessGroupNCCL(TCPStore) mechanism localize uses, so localize stays ~8s cross-node — the fix holds multi-node.
• The trainer itself is model-agnostic (also validated on SD3.5-medium FlowGRPO).

Notes / risk

• Synchronous diffusion behavior is unchanged; its eval defaults still sync weights and sleep the rollout engine.
• The BAGEL-on-vllm_omni enablement fixes required by this recipe are included in merged PR perf(vllm-omni): batch BAGEL grouped t2i rollout into one packed generate_image #203.
• The async buffer / generate-seam machinery is no longer duplicated here: it is the shared unirl/rollout/async_runtime.py that AsyncARTrainer also drives. The one behavioral difference between the two callers is the reap_before_launch order.
• The migration is API-only and changes no async semantics (same launch ceiling, staleness eviction, freshness ordering and quiesce points), but it has not been re-validated on GPU — the reward-curve and localize-timing evidence predates it.

Test Plan

Post-migration, on current main:

• ruff check / ruff format clean; full repository pre-commit passes (including check-recipe-targets)
• unirl.trainer.async_diffusion and unirl.train_async_diffusion import against current main
• python -m unirl.train_async_diffusion --config-name=diffusion/bagel/bagel_vllmomni_async --cfg job --resolve composes
• layout != "separate" and max_inflight != 1 guards fire before any Ray construction
• scheduler ordering: reap-first at max_inflight=1 returns each step with one generation in flight and always reaps against idle rollout workers; launch-first at max_inflight=1 is serialized; launch-first at max_inflight=2 keeps AR's overlap
• BAGEL async GPU re-run on the migrated head

Pre-migration head (not re-run):

• BAGEL async GPU run: reward increases, ratio healthy, no crash
• reap-before-launch fix: localize 151s → ~7.5s, overlap rollout ~148s
• staleness=0 vs staleness=2 comparison; cross-node IB transfer validated
• longer-run reward curve for staleness=2
• periodic eval GPU smoke with resident rollout policy

AI-assisted: the sample-native / shared-runtime migration above was prepared with agent assistance and reviewed against the current main API surface; duplicate-work check — PR #272 covers the only other remaining RolloutReq consumer (sglang_diffusion/adapters/video.py) and does not overlap these files.

[CjhHa1]

Need a long run before merge to see the reward gain

[zzhuoxin1508]

the green line async (vLLM-Omni on a dedicated rollout slab + LoRA synced across slabs every 4 rollouts, SDPA) vs. the red line trainside (in-process sampling, fresh up-to-date weights every rollout, flash-attn) — aligned by rollout/step, the two rewards rise almost in lockstep, with async slightly lower because its sampling weights lag a bit behind.

[zzhuoxin1508]

[CjhHa1]

LGTM

[CjhHa1]

Post-merge validation of the async path on current main

This PR's numbers were measured on the pre-migration head; the sample-native / shared-runtime rewrite that landed with it was never re-validated on GPU. Two main bugs were blocking that (both now fixed: #276 for the vLLM-Omni engine boot, #277 for the BAGEL replay train/eval dispatch), so here is the completed experiment.

Tree: main @ 7c5a1260 (contains this PR as 8fe200e8, plus #276) + #277. No local workarounds — preflight asserted the migrated async trainer, reap_before_launch=True, both fixes present, and no escape hatch anywhere in unirl/.

Setup: 1x8 H20, diffusion/bagel/bagel_vllmomni_async via unirl.train_async_diffusion, num_rollouts=8, max_inflight=1, buffer_max_staleness=2, weight_sync_interval=4, 4 train + 4 rollout, transport_kind=colocate_store, BAGEL-7B-MoT + PickScore staged on local /data, CUDA_LAUNCH_BLOCKING unset.

Result — 8/8 rollouts, rc=0, 1271 s (21.2 min):

rollout 3/8  reward=0.6545  ratio=1.0000±0.0000  gn=0.0001  clip=0.20
rollout 4/8  reward=0.6725  ratio=1.0000±0.0000  gn=0.0001  clip=0.17
rollout 5/8  reward=0.7707  ratio=1.0000±0.0000  gn=0.0001  clip=0.28
rollout 6/8  reward=0.7552  ratio=1.0000±0.0000  gn=0.0000  clip=0.04
rollout 7/8  reward=0.7294  ratio=1.0000±0.0000  gn=0.0000  clip=0.05

per-rollout wall: 125, 125, 125, 130, 137, 138, 135 s
worker_deaths=0   packed_query_sequence=0

The reap-before-launch ordering survived the migration

That is the claim most at risk, because the ordering moved from this PR's hand-rolled _next_batch loop into AsyncRolloutScheduler(reap_before_launch=True), and the shared runtime's default is launch-first. Against the 2026-07-21 A/B over the same 8 rollouts (reap_first ~2010 s, launch_first ~2850 s), this run is 1271 s — firmly in the reap-first regime, with per-rollout ~131 s in the overlap band rather than the ~276 s sync-separate band. That separation is over 2x, so the conclusion holds despite the absolute numbers not being directly comparable to July 21 (this run reads locally staged weights, the A/B read them from CephFS, which is most of why it is faster than the 33.5 min reference).

buffer_max_staleness=2 also behaves as described

With weight_sync_interval=4 over 8 rollouts there are two weight-sync boundaries. The PR states that staleness=2 removes the per-window cold sync-boundary rollout (276 s -> 110 s). The per-rollout series above is uniform at 125-138 s with no spikes at the boundaries, which is the independent confirmation of that claim — if the staleness budget were not letting buffered groups cross a sync, two ~276 s outliers would be visible.

ratio=1.0000±0.0000 on every step reproduces the reported on-policy behaviour and is also evidence that the replay reconstructs the rollout's context faithfully.

Not covered

Only 8 rollouts, so this is a mechanism and throughput check, not the longer reward-curve claim. The reward/localize metric quoted in the description is not instrumented on main (it came from the transfer-A/B experiment branch), so wall clock and per-rollout deltas are used as the measurable equivalent. Multi-node was not re-tested.