refactor(diffusion): make role residency one choice per role, not per phase

[nussejzz]

Fixes #429.

Summary

A colocated loop moved the whole train state across PCIe twice per rollout for no reason.
Generation ended by reloading the trainer, and the reward phase — the very next statement —
began by offloading it again, because the two were separate opt-ins (enable_fsdp_offload,
offload_train_during_reward) that each owned one phase and knew nothing about the other. On a
33B DiT that is tens of GB of host-to-device and device-to-host traffic per rollout, cancelling
out.

Residency becomes a property of a role rather than of a phase. train_resident,
rollout_resident and reward_resident each say whether that role keeps its weights on the GPU
while it is idle; only weights ever move, never a role's process. ResidencyPlanner translates
a phase's needs into transitions and issues only the ones that change something, so the reward
phase inherits an already-parked trainer and emits nothing. The trainer returns once,
immediately before the optimizer step, which also means it stays parked across a whole
accumulate_rollouts window instead of round-tripping per rollout.

Roles that do not share the slab are not tracked at all, so nothing moves memory the active role
could not have used: a reward with its own reward_fraction slab, and the trainer behind a
layout: separate or trainside rollout, whose generation reads the very weights that would be
parked. Trainside behaviour is therefore unchanged.

Two supporting changes fall out of parking the trainer before the rollout wakes:

• The LoRA push can no longer be a single call, because the read wants the trainer resident and
  the load wants the rollout awake. LocalLoraWeightSync gains the extract()/push() split
  that RemoteLoraWeightSync already had, and sync() keeps its exact previous behaviour as
  extract + push. The adapter only changes when the optimizer steps, so the cache survives
  the push and covers every later push until the next step: an accumulate window and an eval's
  later chunks reuse it instead of onloading the trainer to read weights that cannot have
  changed. That read goes through enter(), never set(), so it cannot land the trainer on the
  slab beside a reward the rollout phase has not parked yet.
• train_resident: false is rejected at startup wherever it cannot be honoured: with nothing to
  park (layout: separate, or a trainside rollout), with a full-weight sync (TensorWeightSync,
  NCCLWeightSync, IPCWeightSync, CheckpointWeightSync) whose single sync() reads the
  trainer and loads the rollout in one call, and with EMA/DiffusionNFT plus an external colocated
  rollout. The unparkable case is checked first, since reporting the sync would send the reader
  after the wrong knob. The sync check probes the built object for extract/push rather than
  matching a _target_ suffix, so a new implementation is classified by what it can do.
• AsyncDiffusionTrainer rejects both reward_resident: false (async scoring runs at reap time
  outside _reward_phase()) and rollout_resident: false. The second matters because the save
  hook lives on the base class: _boundary_evaluate deliberately leaves the engine awake
  (sleep_after=False) and the loop resubmits prompts straight after the checkpoint, so parking
  at that boundary would sleep an engine about to be asked to generate. The synchronous loop
  re-enters ROLLOUT before every generate and is unaffected.
• RewardService.offload()/onload() carried no dispatch decorator, so the role's Handle
  could not proxy them (AttributeError: 'Handle' object has no attribute 'onload'). Every
  scorer implements the pair and the managed-process backend has its own per_call
  choreography, but the training loop had no way to reach either — which is why reward residency
  had never been driven from a phase boundary. Broadcast rather than scatter: each reward worker
  holds its own copy of the scorer, so each has to move its own weights.

Two invariants are worth stating because getting them wrong is silent:

• Both the optimizer step and checkpointing read the trainer's weights, so it is made
  resident before each. train() runs train_step → evaluate → maybe_save_checkpoint, and
  evaluate parks the trainer to give the rollout the slab; without the second onload the save
  would gather CPU shards, which a full_state_dict all-gather on the NCCL group cannot do. The
  save-side onload is a _prepare_for_save() hook called after maybe_save_checkpoint's own
  due-or-not predicate, so the predicate stays in one place and a window that evaluates without
  saving does not pay a round trip for a save that will not happen.
• Parking in enter() is unconditional, including when the active role is untracked. A
  trainside rollout, or a trainer on a separate slab, still runs and still displaces whatever
  shares the slab it runs on — so a non-resident reward is parked for the optimizer step even
  when the trainer itself is not tracked.

Every transition goes through the planner, evaluate()'s rollout wake/sleep included, so its
state table is the only source of truth for what is resident. A no-sync multi-chunk evaluation
preserves the rollout while reward scoring runs between chunks; otherwise an unpinned SGLang
engine would discard the adapter that the caller deliberately did not repush.

Synchronous defaults reproduce the previous behaviour exactly (train_resident: true,
rollout_resident: false, reward_resident: true). Async diffusion intentionally overrides only
rollout_resident to true; its constructor and entry point now reference the same policy object,
so direct construction cannot silently fall back to the synchronous value. The three retired keys
are rejected at startup with the key that replaces them rather than being ignored. All three invert, so the
message says so instead of naming one direction.

The AR, PE and unified-model trainers have their own enable_fsdp_offload parameter; those
recipes are untouched, since the rejection lives in DiffusionTrainer.__init__.

Test Plan

Three checks: the planner's transition sequences per policy, the retired-key rejection, and two
hardware runs of the same recipe under different policies.

Hardware. BAGEL-7B-MoT t2i + PickScore on 1x8 H20, layout: colocate, external vLLM-Omni
rollout, two rollouts at weight_sync_interval: 1 so the run has to cross a LoRA push between
rollouts. Run twice against the same recipe:

• all three roles parked (train_resident/rollout_resident/reward_resident: false), so the
  planner has to sequence all of them — a recipe exercising one role would not test it. The
  reward is a managed vLLM child with its own CUDA context, which makes reward residency a real
  transition rather than bookkeeping.
• defaults (train_resident: true, reward_resident: true) — the policy every recipe that this
  PR does not touch runs under. This is the regression arm.
• all parked again, this time with eval_interval: 2 and save_interval: 2, so the run has to
  checkpoint right after an eval has parked the trainer. That ordering is the one that made the
  save gather CPU shards before the fix, and the gate's own eval_interval: 0 had hidden it.
• all parked with eval_interval: 1, eval_num_prompts: 24, eval_chunk_prompts: 8 and
  save_interval: 2 — three eval chunks per pass, an eval after every rollout, and one window
  that evaluates without being due to save. The later chunks resync, which is what reaches the
  adapter read on a parked trainer; a single-chunk eval never gets there.

Recipe added in this PR: examples/diffusion/bagel/bagel_t2i_residency_gate.yaml.

Planner. ResidencyPlanner driven directly through the loop's phase sequence with recording
callables in place of the engines, asserting the exact transition list for each policy. This is
what makes the claim engine-independent: the planner reaches an engine solely through the
(onload, offload) pair handed to it, so a policy's sequence is identical whether the rollout is
SGLang, vLLM-Omni, FastVideo or trainside. Per CLAUDE.md, this repository does not commit a
tests/ tree or one-off harnesses under another name; the planner harness was run externally and
its exact observed sequences are quoted below.

The follow-up review harness also checks: async direct construction defaults rollout residency to
true and rejects explicit false; a two-chunk no-sync eval emits no rollout offload; adapter
cache validity is owned and invalidated by the LoRA weight-sync object; and inherited distributed
cache methods bind on both local and remote LoRA sync handles.

ruff check and ruff format --check pass on every changed file.

Test Result

Defaults arm — nothing moves but the rollout. RES_GATE_RC=0, and over two rollouts the
loop emitted exactly:

2 x lifecycle residency: rollout onload
2 x lifecycle residency: rollout offload

No train and no reward transition at all, which is what the old code did with
enable_fsdp_offload: false + offload_train_during_reward: false: an unconditional
wake_up(), a conditional sleep(), and a _reward_phase() that was a no-op. Every recipe not
touched by this PR is on this path.

All-parked arm. RES_GATE_RC=0. Both rollouts emitted the same six transitions in the same
order, and nothing else:

lifecycle residency: train offload
lifecycle residency: rollout onload
lifecycle residency: rollout offload
lifecycle residency: reward onload
lifecycle residency: reward offload
lifecycle residency: train onload
rollout 1/2  reward=0.7880
... identical sequence ...
rollout 2/2  reward=0.6821

The rollout engine returned ~28 GiB per worker on each park
(Sleep Level 1: physically freed 28.60 GiB). Reward onload/offload appear exactly once per
rollout — the transition that was previously unreachable through the Handle. The trainer
offloads once before generation and returns once before the optimizer step; the redundant
reward-phase pair the old per-phase flags produced is gone.

Reward moved 0.7880 → 0.6821 across the two rollouts; at two points that is noise, and this gate
checks orchestration rather than a learning curve.

Eval + checkpoint arm. RES_GATE_RC=0, and the transition sequence ends the way the fix
requires — the trainer comes back after the final eval parked it:

EVAL step 0        train offload, reward offload / rollout onload, rollout offload
                   reward onload, reward offload
rollout 1/2        rollout onload, rollout offload, reward onload, reward offload, train onload
rollout 2/2        train offload, rollout onload, rollout offload,
                   reward onload, reward offload, train onload
EVAL step 2        train offload, rollout onload, rollout offload,
                   reward onload, reward offload, train onload   <- before the save

Saving checkpoint at rollout 2/2 -> .../checkpoint-2

The checkpoint is on disk (1.9 GB checkpoint.pt + trainer_state.json) with no all-gather
error after it. Note the first eval runs before any training, so it parks a trainer that then
returns for the first optimizer step — one round trip that the old code also paid.

Multi-chunk eval arm. RES_GATE_RC=0, 51 transitions over EVAL step 0, rollout 1/2,
EVAL step 1, rollout 2/2, EVAL step 2, then the save. Three properties are what this arm
is for:

Property	Observed
the trainer never returns beside a resident reward	all 3 train onloads are immediately preceded by reward offload
the trainer returns only where it is read	3 train onloads total, for 2 optimizer steps + 1 save, across 3 evals / 2 rollouts / 3 chunks per eval
the checkpoint is real	1.9 GB checkpoint.pt + trainer_state.json, no all-gather error

Before the fixes this arm would have shown a train onload with the reward still resident (the
peak violation), one train onload per rollout inside the window, and an extra onload for the
window that evaluates without saving.

Long arm with accumulation. RES_GATE_RC=0 over 10 rollouts at accumulate_rollouts: 2
(so five windows, one optimizer step each), eval_interval: 4, save_interval: 4. Three
checkpoints written (checkpoint-4, -8, -10, 5.5 GB total), no errors. The transition
counts are the point:

Transition	Count
rollout onload / rollout offload	16 / 16
reward onload / reward offload	16 / 17
train onload / train offload	7 / 7

Sixteen rollout phases against seven trainer onloads — five optimizer steps plus two
checkpoints (the third coincides with a step). One round trip per window, not per rollout,
which is the property the adapter cache buys and which a two-rollout run cannot show.

Planner and follow-up CPU checks. All pass, including the ones that encode the regression risk.
The follow-up output ended with planner_check_ok and followup_residency_regressions_ok:

Policy	Expected transitions per rollout
defaults	rollout onload, rollout offload
rollout_resident: true	rollout onload only — woken once, never slept
train_resident: false	train offload, rollout onload, rollout offload, train onload
all parked	the six above
train_resident: false, untracked train (trainside / separate)	rollout only; train never moves
reward_resident: false, reward on its own slab	rollout only; reward never moves

The all-parked row is computed independently of the hardware run and matches it transition for
transition. With accumulate_rollouts: 2 and a parked trainer, the planner issues one
train onload for the window rather than one per rollout. All three retired keys raise at
startup for either value with the replacing key named, and a config without them is accepted
unchanged.

Five of the checks exist because an earlier revision of this PR got them wrong, and each fails
against the revision it was written for:

Check	What it would have caught
untracked active role still parks the others	enter(TRAIN) returned early when the trainer was untracked, leaving a non-resident reward on the GPU for the whole optimizer step — on a layout: separate slab that is exactly when it should be gone
trainer is resident before the save	the trainer stayed parked from the last eval into maybe_save_checkpoint
the adapter read parks the reward first	the read used set(TRAIN, True), which moves only the trainer, so the peak held the trainer next to a reward the previous rollout had left resident — the 75GB-trainer-plus-27GB-judge shape that qwen_image_edit_plus_nft_managed_editscore.yaml documents as an H20 OOM
one train round trip per accumulate window	re-reading the adapter for every rollout in a window put the trainer back to one round trip per rollout, contradicting this PR's own claim
a window that will not save pays no onload	the pre-save onload ran unconditionally, so an eval_interval denser than save_interval paid a round trip on every eval
the save hook does park an unpinned rollout	pins the reason AsyncDiffusionTrainer has to reject rollout_resident: false: the hook is on the base class, and the async loop has no wake between the checkpoint and the next submission

Not run on hardware: SGLang. It is the most-used external engine for diffusion in examples/
(21 recipe references vs 15 for vLLM-Omni) and it is affected by this change, but no SGLang
build was available on the test host. Its coverage rests on the planner checks above plus the
interface being the one the planner uses — SGLangDiffusionRolloutEngine.sleep/wake_up are
the same BaseRolloutEngine contract vLLM-Omni implements, and no SGLang recipe sets any of the
retired keys, so all of them run the defaults arm verified above. Worth a maintainer sanity run
if you have a host with SGLang.

Not run: reward_resident: false and rollout_resident: false under AsyncDiffusionTrainer —
both rejected at startup by design, the first because async scoring happens at reap time
outside _reward_phase(), the second because the engine owns a dedicated slab and the loop
submits across the checkpoint boundary.

Reviewer Notes

Follow-up review removed the duplicate trainer-side adapter-cache flag, the unused
ResidencyPlanner.on_gpu() method, repeated literal residency defaults and redundant internal bool
coercions. Review fixes were AI-assisted and then checked against the full diff and focused harnesses.

[nussejzz]

Rebased onto main at 1e87b0d and re-ran the full hardware matrix against the rebased head (3569918), so the results below are for the tree as it would merge. Same recipe (bagel_t2i_residency_gate.yaml, now an overlay on bagel_vllmomni.yaml), same 1x8 H20 nodes, BAGEL-7B-MoT t2i + local PickScore, external vLLM-Omni rollout. The arms were spread over three nodes so they could run at once; each arm's pass condition is asserted by the harness (exact per-role onload counts, plus the invariant that every train onload is immediately preceded by a reward offload), not eyeballed.

Arm	Overrides	train / rollout / reward onloads	Verdict
all parked	(recipe defaults: all three *_resident: false)	2 / 2 / 2	PASS — the six-transition cycle, twice, nothing else
defaults	train_resident=true reward_resident=true	0 / 2 / 0	PASS — only the rollout pair; the path every untouched recipe runs
eval + checkpoint	eval_interval=2 save_interval=2	3 / 10 / 10	PASS — 2 steps + 1 post-eval save; checkpoint-2 on disk (1.9 GB checkpoint.pt + trainer_state.json), no all-gather error
multi-chunk eval	eval_interval=1 eval_num_prompts=24 eval_chunk_prompts=8 save_interval=2	3 / 11 / 11	PASS — 3 evals x 3 chunks + 2 rollouts; all 3 train onloads preceded by reward offload
long, accumulate	accumulate_rollouts=2 stack.num_updates_per_batch=1 eval_interval=4 save_interval=4 eval_num_prompts=16 eval_chunk_prompts=8, 10 rollouts	7 / 16 / 16	PASS — five windows (one step each) + two post-eval saves = 7 train onloads against 16 rollout phases (10 rollouts + 3 evals x 2 chunks); checkpoint-4/-8/-10 on disk

The all-parked sequence per rollout, verbatim from the log:

train offload, reward offload
rollout onload
rollout offload
reward onload
reward offload
train onload

Two things changed in the branch besides the rebase:

• The tip commit (gate recipe as an overlay) was authored under a stale local identity; re-committed as Ding Zuhao. Its message also claimed to drop cfg_scale_of(), which was already gone from the tree; that paragraph is removed.
• pre-commit run --all-files (the CI lint job, including the one-line-docstring and recipe-target hooks) passes on the rebased head.

One thing the long arm surfaced that the recipe should say out loud: accumulate_rollouts > 1 is rejected at startup unless stack.num_updates_per_batch == 1, and bagel_vllmomni.yaml sets 2. The override is in the table above; nothing in the trainer changed.

Still current against today's main (c8af5d6)

main has moved 15 commits since that base. The branch still merges cleanly, and none of those commits reach the paths this PR orchestrates:

• Only two touch files this PR also edits, both in disjoint regions: fix(reward): make prompt semantics explicit for composed rollouts #477 rewrites _build_reward_request in reward/service.py (this PR only adds dispatch decorators to offload/onload at the end of the same file), and feat(trainer): teacher-anchored capability negotiation in ARTrainer #457 edits a different section of trainer/README.md.
• perf(train): release consumed grads at the optimizer-step boundary #460 adds optimizer.zero_grad(set_to_none=True) at the end of optimizer_step(); its own message cites grads staying alive across offload()/onload(), so it is complementary — it shrinks the parked train state. Backward still happens entirely inside one train_track call, so accumulation windows are unaffected.
• fix(checkpoint): preserve cold AdamW state during checkpoint export #440 changes optimizer-state export for cold AdamW. That is the path _prepare_for_save() exists to protect, and the two are independent: this PR guarantees the trainer is resident before the gather, fix(checkpoint): preserve cold AdamW state during checkpoint export #440 changes what the gather emits.
• unirl/trainer/rollout_sleep.py, new on main in feat(vLLM): support direct vLLM TP rollout engine and vLLM native IPC weight sync #480, is an AR-only helper, not a second residency mechanism.

I have not re-run the matrix on the newer base. The five arms assert transition sequences, and nothing above changes a transition sequence, so the results stand; say the word if you would rather see the all-parked and defaults arms re-run against the merge result before approving.

Not run on hardware, unchanged from the description: SGLang. No SGLang build on the test host. Every SGLang recipe runs the defaults arm's policy, which is verified above.

[leviking98z-rgb]

LGTM