fix(bagel): force eval() for replay's packed-query KV-context rebuild

[CjhHa1]

Summary

BAGEL FlowGRPO training against a deferred-prompt rollout engine (the vllm_omni / sglang_diffusion recipes) dies in the first optimizer step with:

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

Three things line up to produce it:

1. The vendored navit dispatches on module mode — Qwen2Model.forward routes to forward_train() when self.training and forward_inference() otherwise (unirl/models/bagel/vendor/modeling/bagel/qwen2_navit.py:974-978).
2. Since feat(agentic): add Sample-native multi-turn rollout and training #214, TrainStack.train_track switches the model to train() before the gradient-bearing replay so HF gradient checkpointing engages (unirl/train/stack/base.py:508). Before that commit the replay inherited the eval() set a few lines above for the π_old anchor freeze, so the whole replay ran in eval().
3. BagelDiffusionStage._build_contexts_from_prompt issues packed-query LM forwards (update_context_text → forward_cache_update_text → forward_inference) with no eval() guard. forward_flow already forces eval() for its own call — this is the one replay entry point reached before any forward_flow, so it is the only unguarded one.

It stayed latent because it needs both halves. Only deferred-prompt engines reach the rebuild at all: trainside/colocate carry the opaque KV contexts, so _resolve_single returns early and never issues these forwards. And only post-#214 is replay in train() mode.

The fix adds rl_ops.inference_dispatch, which names the rule forward_flow already applies: force eval() for a packed-query region, and restore the previous mode only when no backward follows. Activation-checkpointing recompute re-enters these forwards during .backward(), so a restored train() would mis-dispatch them again — the same asymmetry _forward_flow_train_safe documents. Guarding the rebuild itself (rather than one call site) covers all three of its callers: diffuse, replay, and build_forward_kwargs.

Related Issue

N/A.

Complements #277 (fix(vllm_omni): stop PR_SET_PDEATHSIG from killing booted DiffusionWorkers), which explicitly names this as needing its own fix. The two are independent and both are required to run a separate-slab BAGEL FlowGRPO job on current main: #277 unblocks the four-replica engine boot, this unblocks the train step immediately after it. No file overlap — #277 touches only patches/runtime.py.

Test Plan

Static + unit, on this branch rebased onto main:

• ruff check and ruff format --check on both changed files: clean. Full repository pre-commit passed via the pre-push hook.
• inference_dispatch semantics, exercised against a stub module that mimics the navit if self.training dispatch:
  • caller in train() under no_grad → region dispatches inference, previous train() restored;
  • caller in train() with grad enabled → region dispatches inference and is deliberately left in eval(), so the backward's AC recompute also dispatches inference;
  • caller already in eval() → no-op.
• PR_SET_PDEATHSIG / mode probes used during diagnosis are not part of this change.

GPU end-to-end: not run; reason: reaching the train step on a separate-slab BAGEL job first requires #277 — without it every DiffusionWorker is SIGKILLed during engine boot and the job dies at the first generate, before any of this code runs.

What I did confirm on GPU (1×8 H20, diffusion/bagel/bagel_vllmomni_async via unirl.train_diffusion, num_rollouts=1, 4 train + 4 rollout, transport_kind=colocate_store, with the watchdog temporarily gated off locally): all four engines boot, generate and reward.score_and_attach both complete, and the run then raises exactly the packed_query_sequence TypeError above in TrainStack._run_update → flowgrpo.compute_loss_and_backward → BagelDiffusionStage.replay. That pins the failure this PR removes; verifying the fix itself needs one run stacked on #277, which I will post here.

Compatibility / Risk

• No config, checkpoint, data-format or public-API change. One new helper in unirl/models/bagel/rl_ops.py, exported via __all__.
• The only behavior change is the MoT's mode during the context rebuild. No-grad callers (rollout) get their previous mode restored; the grad path is left in eval(), which is already where forward_flow leaves it for the remainder of the step, so the mode at the end of a replay is unchanged either way.
• Not a memory regression: forward_flow already forces eval() for the heavy velocity forward, so on main the LM is effectively already in eval() for the expensive part of replay.
• Worth a maintainer's opinion, deliberately not changed here: the self.fsdp_backend.model.train() at unirl/train/stack/base.py:508 buys BAGEL nothing, because the vendored navit has no self.training-gated gradient checkpointing and BAGEL relies on FSDP activation checkpointing instead. Narrowing that switch would affect every other model, so this PR fixes the BAGEL side only.

Reviewer Notes

• Start with the restore rule in inference_dispatch (the "only when no backward follows" branch). That asymmetry is the subtle part, and it intentionally mirrors _forward_flow_train_safe rather than inventing a second convention.
• Duplicate-work check: scanned open PRs for bagel / fate / replay / dispatch / forward_train keywords. Only fix(bagel): rebuild replay KV contexts in eval() so navit keeps inference dispatch #277 is adjacent, and it does not touch these files.
• AI assistance: the diagnosis and patch were prepared with agent assistance. The root cause is backed by the full runtime traceback and the #214 diff rather than inference, and every changed line is accounted for above.
• Follow-up: post the GPU end-to-end result once fix(bagel): rebuild replay KV contexts in eval() so navit keeps inference dispatch #277 lands.

Checklist

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

No test file added: the repository ships no unit-test suite for this area (pytest is a dev dependency but there is no test tree), so the guard's three semantics were verified with the inline stub described in the Test Plan. The stale contract in BagelDiffusionStage.replay's docstring ("caller owns .train() mode") is corrected in this PR, since that is precisely the assumption that does not hold for BAGEL.

[CjhHa1]

Closing as a duplicate of #277, which fixes the same bug in the same function and is the one we are keeping.

Both PRs diagnosed the identical root cause (navit dispatches forward_train vs forward_inference on self.training; _build_contexts_from_prompt uses the packed-query inference signature; TrainStack.train_track leaves the model in train() before the gradient-bearing replay since #214). #277's inline guard is the better shape:

• it keeps the mode change local and restores the caller's mode unconditionally, which is sound here because the whole block is non-gradient-bearing (no_grad, frozen und/text path), so no activation-checkpointing recompute revisits these forwards during .backward();
• that makes the rl_ops.inference_dispatch helper this PR added — and in particular its "leave it in eval() when grads are enabled" asymmetry — unnecessary complexity for a single call site. The subsequent forward_flow handles its own mode either way.

I have pushed the simplification to #277: nn.Module.train() takes the mode, so train(was_training) replaces both if was_training branches (the idiom already used in unirl/train/stack/base.py), and the comment is trimmed to the constraint the code cannot express.

Verification carried over to #277: the failure was reproduced on 1x8 H20 (diffusion/bagel/bagel_vllmomni_async via unirl.train_diffusion, 4 train + 4 rollout, num_rollouts=1) with the engine-boot bug worked around locally -- generate and reward both completed, then the packed_query_sequence TypeError was raised in TrainStack._run_update -> flowgrpo.compute_loss_and_backward -> BagelDiffusionStage.replay. Now that #276 has landed, that workaround is no longer needed and the end-to-end run can be done on plain main.