fix(bagel): rebuild replay KV contexts in eval() so navit keeps inference dispatch

[CjhHa1]

Summary

BAGEL FlowGRPO died in the training forward with:

TypeError: Qwen2Model.forward_train() got an unexpected keyword argument 'packed_query_sequence'

Every vendored navit module dispatches forward_train vs forward_inference on self.training. The vllm_omni rollout ships only the prompt text (KV caches cannot cross the worker->driver IPC boundary), so replay rebuilds the three KV contexts trainer-side through the vendored inferencer's inference (packed-query) signature. TrainStack.train_track puts the model in train() before the gradient-bearing replay so HF gradient checkpointing engages, so the rebuild arrived in train() mode and the packed kwargs were routed into forward_train.

rl_ops.forward_flow already forces eval() for exactly this reason, and require_inference_dispatch documents the contract ("replay runs in eval() with grads enabled") — the context rebuild was the one vendored inference-signature call in the replay path still missing the guard. Restoring the caller's mode afterwards is safe: the und/text path is frozen and the block already runs under no_grad, so no activation-checkpoint recompute revisits it during .backward().

Only the deferred-prompt path is affected; a trainside/colocate rollout carries opaque contexts and has_contexts() short-circuits before the rebuild, which is why this surfaced only on the vllm_omni BAGEL recipe.

Related Issue

N/A

Test Plan

Ran on one 8xH20 pod with main's own recipe forced onto two slabs:

python -m unirl.train_diffusion --config-name=diffusion/bagel/bagel_vllmomni \
  +layout=separate +train_fraction=0.5 num_devices=8 +devices_per_node=8 \
  num_rollouts=1 weight_sync_interval=4 \
  sync._target_=unirl.distributed.weight_sync.lora.RemoteLoraWeightSync \
  rollout.config.enable_sleep_mode=false transport_kind=colocate_store

Result: a complete generate -> reward -> train step, navit_dispatch_errors=0, optimizer step finishing.

Reaching the training forward at all on this recipe additionally requires #276, which fixes an unrelated PR_SET_PDEATHSIG bug that SIGKILLs healthy DiffusionWorkers before the first generate. The acceptance run above was executed with that fix also applied. This PR is independent of it and touches a disjoint file.

ruff check / ruff format clean, and the full pre-commit suite (including check-recipe-targets) passes via the pre-push hook.

Compatibility / Risk

No config, checkpoint, data-format or API change. One behavioural change confined to one function: the eval() guard is scoped to the frozen, no_grad context rebuild and restores the caller's mode, so gradient checkpointing for the trained surface is untouched.

Reviewer Notes

This PR was originally filed bundled with a PR_SET_PDEATHSIG fix in unirl/rollout/engine/vllm_omni/patches/runtime.py. That half has been dropped and the branch rewritten (force-push 1c20d989 -> 81ca0c23), because #276 fixes the same bug, was filed first, and carries the correct root-cause account. See the discussion on #276 for the kernel probe that settles it.

Checklist

• I reviewed the changed code and removed unrelated/generated artifacts.
• I updated tests, docs, and configs where needed, or explained why not.

[CjhHa1]

Closing in favour of #276, which fixes the same bug and was filed first (10:09 UTC vs 10:46 UTC here). Same file, same function — they cannot both land.

#276 is also the more accurate account of the failure. The rationale in this PR said the spawning process "legitimately exits after handing off", but that cannot be what happened: install_fate_sharing's watchdog still carries a process-level reparent check (original_ppid != 1 and os.getppid() != original_ppid -> os._exit(1)), so a genuinely exited parent would have had the worker reaped within one 5s poll and dropping prctl would not have helped. What actually exits is the short-lived stage-init thread, which is what PR_SET_PDEATHSIG binds to — exactly the framing in #275 and #276.

The A/B table here also overstates its isolation: the two arms toggled both prctl and the poll at once rather than one variable, and the "the poll never fired in either arm" claim is not assertable for the arm where the poll was disabled. The conclusion happened to be right, but #276 reaches it with a cleaner argument.

The empirical data from this branch is still useful and has been carried over to #276 as a review comment: the 4-replica layout="separate" bagel_vllmomni repro shows the bug is neither single-GPU- nor SD3-specific and reproduces under the plain synchronous DiffusionTrainer.

The one behavioural difference: this PR removed PDEATHSIG entirely, whereas #276 keeps it for children created on the parent's main thread. #276's narrower scope is the safer change.

[CjhHa1]

GPU end-to-end: PASS on clean main

Now that #276 has landed, this could be verified without any local workaround.

Tree under test: main @ 7c5a1260 (the #276 PDEATHSIG boot fix) + this PR's two commits, and nothing else. Preflight asserted all three conditions before launching: the mot.train(was_training) guard present, arm_pdeathsig present, and no UNIRL_DISABLE_FATE_SHARING escape hatch anywhere.

Setup: 1x8 H20, diffusion/bagel/bagel_vllmomni_async driven by unirl.train_diffusion, num_rollouts=1, layout=separate / train_fraction=0.5 (4 train + 4 rollout), weight_sync_interval=4, transport_kind=colocate_store, BAGEL-7B-MoT + PickScore staged on local /data, CUDA_LAUNCH_BLOCKING unset.

Result:

RESULT rc=0 wall_s=388 engines_booted=4 worker_deaths=0 executor_closed=0 packed_query_sequence=0
rollout 1/1  reward=0.7502  loss=0.0000 gn=0.0001 lr=1.00e-04 ratio=1.0000±0.0000 clip=0.03

A complete generate -> reward -> advantage -> optimizer step finished in 388 s. All four vLLM-Omni replicas booted and none lost its DiffusionWorker (that is #276 doing its job), and the packed_query_sequence TypeError this PR removes did not occur.

For contrast, the same recipe and topology on the pre-fix tree failed in two distinct ways, both now gone: with the watchdog active every replica lost DiffusionWorker-0 right after its engine init and the run died at the first generate with DiffusionExecutor is closed; with the watchdog worked around, generate and reward completed and the run then raised TypeError: Qwen2Model.forward_train() got an unexpected keyword argument 'packed_query_sequence' in TrainStack._run_update -> flowgrpo.compute_loss_and_backward -> BagelDiffusionStage.replay.

ratio=1.0000 on the first step is the expected on-policy signature and is incidental evidence that the rebuilt KV contexts match what the rollout worker used -- a mis-dispatched or differently-built context would show up as a ratio away from 1.

Also folded in a simplification of the guard itself: nn.Module.train() takes the mode, so train(was_training) replaces both if was_training branches (the idiom already used at unirl/train/stack/base.py:462), and the comment is trimmed to the constraint the code cannot express. Net effect on the diff: 24+/4- becomes 15+/4-.

Duplicate note: #278 proposed the same fix behind an rl_ops.inference_dispatch context manager and has been closed in favour of this PR -- the inline guard keeps the mode change local and does not need that helper's grad-conditional restore rule.

[haonan3]

LGTM

[haonan3]

Post-merge review note (evidence hygiene, the fix itself checks out): the Test Plan line citing navit_dispatch_errors=0 references a metric that doesn't exist — git grep navit_dispatch_errors is empty across unirl/, recipes/, and the pinned vllm-omni checkout, so no code emits that counter. The correctness of this fix stands on the code path (verified independently: eval scope covers all dispatching modules, rebuilds are stateless per call, no rollout regression), just not on that line — flagging so future readers don't go hunting for the metric.

One real follow-up it surfaced: BagelPipeline._build_contexts (trainside/SFT) runs the same inference-signature prefills with no eval guard and only avoids the TypeError by call-order luck today. Guarded in #290 with the same pattern as this PR.