feat(trainer): AsyncDiffusionTrainer for disaggregated async diffusion RL

[zzhuoxin1508]

Summary

AsyncDiffusionTrainer — the pure-DiT async RL trainer. Diffusion sibling of AsyncARTrainer: training and rollout on disjoint GPU slabs, generation overlapped with training, reward scored off the train critical path, weights pushed cross-slab. The synchronous diffusion trainer is unchanged — purely additive (3 new files).

Motivation: the synchronous diffusion path runs generate → reward → train in series each step, so generation + reward sit on the train critical path. Async overlaps them behind the 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.)

What's added (3 files, purely additive)

• unirl/trainer/async_diffusion.py — AsyncDiffusionTrainer(DiffusionTrainer). Reuses layout="separate" two-slab build + _connect_separate NCCL handshake + _build_req / _drop_decoded / evaluate / checkpoint + FlowGRPO stack.train_track. Overlays the async buffer loop (mirrors AsyncARTrainer, engine/modality-agnostic): non-blocking _generate_async, reap-before-launch _next_batch, _reap_ready → _score_into_buffer (reward at reap time), _RolloutBuffer of scored GRPO groups with freshness/staleness eviction, _drain_all quiesce before every weight sync.
• 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=0).

Knobs: max_inflight (must be 1, see above); weight_sync_interval (>1 required for overlap — interval=1 drains every step); buffer_max_staleness (0 = on-policy, ratio≈1; >0 = buffer survives weight syncs, off-policy continuous buffer).

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

• ratio = 1.0000 throughout, 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×; ratio stays 1.0 (old_logp_source=rollout anchors π_old to the emitted logp, so staleness moves gradient freshness, not the ratio).
• 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

• Purely additive: 3 new files, no changes to existing trainer / model code.
• Running the BAGEL recipe additionally requires two small BAGEL-on-vllm_omni **enablement fixes not included
• The async buffer / generate-seam machinery is a faithful copy of AsyncARTrainer (engine/modality-agnostic); a shared mixin is a natural follow-up refactor.

Test Plan

• --cfg job compose passes
• 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 (off-policy drift check) — follow-up