perf(inference): skip the generation tower and video guardrail for reasoner-only runs

[lfengad]

Extends the reasoner-only hook added in #132 from the generation VAE to the rest of
the generation side. On Cosmos3-Edge this leaves 2.757 GiB/GPU (−36%) unallocated
for reasoner inference, with byte-identical output.

What is skipped

Module	Saved
MoT generation tower — all ten *_moe_gen modules (new include_gen_pathway)	1409.4M params, 2.62 GiB
VFM generation heads — time_embedder, proj_in/out, action_* (existing vision_gen/action_gen/sound_gen)	14.2M params, 0.03 GiB
Video guardrail runner — RetinaFace (new include_video)	0.10 GiB
Generation VAE	already handled by #132

Measured residency audit accounts for the full delta: 2.65 (net) + 0.10 (RetinaFace)
= 2.75 GiB, matching the 2.757 GiB observed end-to-end.

Why no weight-loading changes were needed

dcp.load is pull-based — it requests only the keys the live model exposes
(get_model_state_dict(model.model)), so a model built without the *_moe_gen
modules never reads their tensors off disk. Verified on a real Edge checkpoint: no
missing/unexpected key errors. This was an assumption at design time; it was
confirmed on hardware before relying on it.

Defaults are inert

All four flags default to the generation-enabled value (load_vision_tokenizer,
vision_gen/action_gen/sound_gen, include_gen_pathway, video_guardrail).
The only behavioural trigger is is_reasoner_only(...) in scripts/inference.py,
which requires a non-empty, all-reasoner sample list. Mixed batches, empty lists,
Ray serve, and the action policy servers are unaffected.

In the default path the only executed change is two always-true assertions (one in
_impl_forward, one in _run_video_guardrail) that turn a would-be
AttributeError on a misconfigured run into an explanatory failure.

The three model-config overrides are orthogonal and each degrades gracefully:
dropping any one entry from reasoner_only_overrides() restores the previous
behaviour for that module alone.

Verification (4× GB200, Cosmos3-Edge)

Path	Result
reasoner, text-only	byte-identical to main (md5 31165084…), 7.624 → 4.867 GiB
reasoner, image	byte-identical to main (md5 995e67f1…), 10.343 → 7.586 GiB
reasoner, video	model construction identical (4.735 GiB, same as text-only); decode step fails identically on main and this branch — the container lacks libnppicc.so.13, downstream of anything touched here
t2i generation, fixed seed	byte-identical to main (md5 48d22c11…)
Unit tests	13 new; full suite 192 passed / 6 failed, the 6 failing identically on main (GPU-less lint node)
just lint	passes

The t2i comparison was run only after confirming the path is bit-reproducible at a
fixed seed (two runs of the same code produced identical bytes). With the default
seed=None the path is not reproducible run-to-run, so a naive byte diff there is
meaningless — noted in case anyone else tries to regression-test image output.

Not covered

• No training smoke test. unified_mot.py is imported by the nano/edge/super SFT
  model configs. The change is default-inert there and the generative t2i run
  exercises the same construction path, but no training run was performed. The
  task="vlm" reasoner SFT builds VLMModel, which does not reference
  unified_mot at all; the recipe that would exercise it (vision_sft_edge) needs
  a dataset, VAE and 8 GPUs that were not available.
• Only Cosmos3-Edge was tested. Nano/Super route through the Qwen3-VL MoT
  classes, which the flag also threads through (default True, so unchanged), but
  no inference was run on them.

Note

include_gen_pathway does not leak into exported checkpoint configs: the LazyDict
only carries kwargs the experiment config passes explicitly, so the True default
stays in _MoTConfigBase.__init__. Confirmed by dumping load_model_config_dict()
on both branches — byte-identical, key absent in both.

🤖 Generated with Claude Code