[Bug] Colocated diffusion loop moves the train state across PCIe twice per rollout: residency is per phase, not per role #429

System info

1x8 H20 96GB, single node. layout: colocate.

Rollout / inference engine

vLLM-Omni (colocated, external engine — not trainside).

Domain

Diffusion (DiffusionTrainer). Reproduced with BAGEL-7B-MoT; the cost scales with model size,
so it is worst on the 33B DiT recipes.

Recipe / config-name

examples/diffusion/bagel/bagel_it2i_managed_editscore.yaml — the one committed recipe that
turns the reward-phase offload on (offload_train_during_reward: true).

Reproduction

Run any colocated diffusion recipe with train offload enabled and a colocated reward:

layout: colocate
enable_fsdp_offload: false
offload_train_during_reward: true
rollout_sleep_after_generate: true

Then follow the trainer's weights across one rollout.

Expected behavior

The trainer's weights should leave the GPU once when they stop being needed and come back once
when they are needed again.

What happens instead

They make two extra full trips across PCIe per rollout, because generate-time offload and
reward-time offload are independent opt-ins that each own one phase and know nothing about the
other.

_generate_with_residency ends by reloading the trainer in its finally block:

elif train_offload_attempted:
    cleanup_steps.append(("generate train onload", self.backend.onload))

and _reward_phase — entered by the very next statement in _rollout_and_score — begins by
offloading it again:

@contextmanager
def _reward_phase(self):
    should_offload = self._offload_for_reward_phase()
    try:
        if should_offload:
            self.backend.offload()

So the sequence per rollout is offload → onload → offload → onload, where
offload → ... → onload would do. On a 33B DiT that is tens of GB of host-to-device plus the
same device-to-host, per rollout, with the two middle transfers cancelling out. With
accumulate_rollouts > 1 it repeats for every rollout in the window even though the optimizer
step only needs the trainer once at the end.

The shape of the problem is that residency is expressed per phase. Each flag can only answer
"should this phase offload?", which cannot express "stay parked, the next phase does not need
you either" — the information lives in the phase that has not run yet.

Two smaller consequences of the same framing:

• A colocated reward cannot be parked at all. Every scorer implements offload()/onload()
  and the managed-process backend has its own per_call choreography, but
  RewardService.offload/onload carry no dispatch decorator, so the role's Handle cannot
  proxy them (AttributeError: 'Handle' object has no attribute 'onload'). Nothing in the
  training loop had ever reached them, so the reward's weights sit on the slab while the rollout
  generates.
• rollout_sleep_after_generate is named after a phase too, which is why it reads as
  process lifetime rather than weight residency and has to be documented as "false keeps an
  external engine's weights resident".

Suggested direction

Make residency a property of a role rather than of a phase — one key per role saying whether it
keeps its weights on the GPU while idle — and have a single component turn a phase's needs into
transitions, issuing only the ones that change something. Then the reward phase inherits an
already-parked trainer and issues nothing, and the trainer can come back once per optimizer step
instead of once per rollout.

Note this only affects layout: colocate with an external engine. Trainside rollout reads the
very weights that would be parked, and a layout: separate trainer does not share the slab, so
neither should be touched.

PR #428 implements this.