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

cosmos_framework/inference/args.py

@@ -202,6 +202,34 @@ def is_reasoner_only(sample_overrides: Sequence["OmniSampleOverrides"]) -> bool:
     return bool(sample_overrides) and all(sample.sample_meta.model_mode.is_reasoner for sample in sample_overrides)
 
 
+def reasoner_only_overrides() -> list[str]:
+    """Config overrides that drop generation-side modules for reasoner-only runs.
+
+    Reasoner decoding runs entirely in the understanding tower, so none of these
+    modules are reachable. ``dcp.load`` is pull-based, so not building them also
+    means their tensors are never read from the checkpoint.
+
+    The three flags are deliberately orthogonal and each defaults to the
+    generation-enabled value: dropping any one entry from this list degrades to
+    the previous behaviour for that module alone, leaving the others in effect.
+    """
+    return [
+        # The generation VAE (Wan2.2 for Edge): ~1.3 GiB resident, and its weights
+        # are downloaded lazily at tokenizer construction, so skipping it also
+        # avoids a ~2.6 GiB download on a cold cache.
+        "model.config.load_vision_tokenizer=false",
+        # VFM-level generation heads: time_embedder, proj_in/proj_out, action_*, sound_*.
+        # All three modality flags must fall together: ``Cosmos3VFMNetworkConfig``
+        # asserts that action and sound generation each imply vision generation, and
+        # Edge checkpoints ship with ``action_gen=true``.
+        "model.config.vision_gen=false",
+        "model.config.action_gen=false",
+        "model.config.sound_gen=false",
+        # The MoT generation tower (every ``*_moe_gen`` module): ~2.0 GiB resident.
+        "model.config.vlm_config.model_instance.config.include_gen_pathway=false",
+    ]
+
+
 class VisionMode(StrEnum):
     IMAGE = "image"
     VIDEO = "video"

cosmos_framework/inference/args_test.py

@@ -20,6 +20,7 @@
     SoundDataOverrides,
     _get_nvml_device_memory_info,
     is_reasoner_only,
+    reasoner_only_overrides,
 )
 from cosmos_framework.inference.common.config import structure_config
 
@@ -57,6 +58,38 @@ def test_reasoner_only_override_disables_vision_tokenizer_in_model_config(tmp_pa
     assert model_dict.config.load_vision_tokenizer is False
 
 
+def test_reasoner_only_overrides_disable_every_generation_side_module(tmp_path: Path) -> None:
+    """All three overrides must land on the live model config, not just parse.
+
+    They are orthogonal on purpose: dropping any one of them degrades to the
+    previous behaviour for that module alone.
+    """
+    setup_args = OmniSetupOverrides(
+        checkpoint_path=DEFAULT_CHECKPOINT_NAME,
+        output_dir=tmp_path / "outputs",
+    ).build_setup(world_size=1, local_world_size=1, device_memory_bytes=_H100_MEMORY_BYTES)
+
+    model_dict = structure_config(setup_args.load_model_config_dict(), omegaconf.DictConfig)
+    assert model_dict.config.load_vision_tokenizer is True
+    assert model_dict.config.vision_gen is True
+    # ``include_gen_pathway`` is absent until overridden: the LazyDict only carries
+    # kwargs the experiment config passed explicitly, so the True default lives in
+    # ``_MoTConfigBase.__init__``. That is also why enabling this flag cannot leak
+    # into an exported checkpoint config.
+    assert "include_gen_pathway" not in model_dict.config.vlm_config.model_instance.config
+
+    setup_args.experiment_overrides.extend(reasoner_only_overrides())
+
+    model_dict = structure_config(setup_args.load_model_config_dict(), omegaconf.DictConfig)
+    assert model_dict.config.load_vision_tokenizer is False
+    assert model_dict.config.vision_gen is False
+    assert model_dict.config.vlm_config.model_instance.config.include_gen_pathway is False
+    # Cosmos3VFMNetworkConfig asserts action/sound generation each imply vision
+    # generation, so the three modality flags must be switched off together.
+    assert model_dict.config.action_gen is False
+    assert model_dict.config.sound_gen is False
+
+
 def test_build_parallelism(monkeypatch: pytest.MonkeyPatch):
     parallelism_args = OmniSetupOverrides(
         checkpoint_path=DEFAULT_CHECKPOINT_NAME,

cosmos_framework/inference/common/args.py

@@ -748,6 +748,18 @@ class GuardrailArgs(ArgsBase):
 
     guardrails: bool
     offload_guardrail_models: bool
+    video_guardrail: bool = pydantic.Field(default=True, exclude=True)
+    """Build the video guardrail runner.
+
+    Runtime-only policy field. It has no ``GuardrailOverrides`` counterpart on
+    purpose: exposing a CLI flag that disables face blurring for generative runs
+    would be a safety footgun, so only the reasoner-only path sets it False.
+
+    ``exclude=True`` keeps it out of ``model_dump()`` because ``SetupArgs`` dumps
+    must round-trip back through ``SetupOverrides``, which forbids extra keys
+    (see ``args_test.test_setup_args``). Anything that rebuilds args from a dump
+    therefore gets the safe default of building the runner.
+    """
 
 
 class GuardrailOverrides(OverridesBase):

cosmos_framework/inference/common/inference.py

@@ -78,15 +78,25 @@ def sync_distributed_errors():
 @dataclass
 class GuardrailRunners:
     text: "GuardrailRunner"
-    video: "GuardrailRunner"
+    video: "GuardrailRunner | None"
 
     @classmethod
-    def create(cls, args: GuardrailArgs, /) -> Self:
+    def create(cls, args: GuardrailArgs, /, *, include_video: bool = True) -> Self:
+        """Build the guardrail runners.
+
+        ``include_video=False`` skips the video runner entirely (its RetinaFace
+        weights are ~109 MB of otherwise idle GPU memory). Reasoner-only runs use
+        it because they never produce frames; every other caller keeps the default.
+        """
         from cosmos_framework.auxiliary.guardrail.common import presets
 
         return cls(
             text=presets.create_text_guardrail_runner(offload_model_to_cpu=args.offload_guardrail_models),
-            video=presets.create_video_guardrail_runner(offload_model_to_cpu=args.offload_guardrail_models),
+            video=(
+                presets.create_video_guardrail_runner(offload_model_to_cpu=args.offload_guardrail_models)
+                if include_video
+                else None
+            ),
         )
 
 
@@ -132,7 +142,11 @@ def generate_batch(
     def create(cls, setup_args: SetupArgs, /) -> Self:
         """Create instance."""
         timer = TrainingTimer() if setup_args.benchmark else None
-        guardrails = GuardrailRunners.create(setup_args) if setup_args.guardrails else None
+        guardrails = (
+            GuardrailRunners.create(setup_args, include_video=setup_args.video_guardrail)
+            if setup_args.guardrails
+            else None
+        )
         return cls._create(setup_args, guardrails=guardrails, _timer=timer)
 
     @torch.no_grad()
@@ -284,6 +298,10 @@ def _run_video_guardrail(self, name: str, video_cthw: torch.Tensor) -> torch.Ten
         """Run guardrail checks on the video and apply face blur."""
         if self.guardrails is None:
             return video_cthw
+        assert self.guardrails.video is not None, (
+            "The video guardrail runner was not built (include_video=False), but a sample "
+            "produced frames. Only reasoner-only runs may skip it — this is a policy bug."
+        )
         processed_video_cthw, message = _run_video_guardrail(self.guardrails.video, video_cthw)
         if processed_video_cthw is None:
             raise ValueError(f"Guardrail blocked video '{name}': {message}")

cosmos_framework/inference/common/inference_test.py

@@ -213,6 +213,41 @@ def broadcast(payload: list[str | None], *, src: int) -> None:
     download.assert_not_called()
 
 
+def test_guardrail_runners_skip_video_construction(monkeypatch: pytest.MonkeyPatch) -> None:
+    """``include_video=False`` must not even construct the video runner.
+
+    Reasoner-only runs never reach the video guardrail, and the point of the flag
+    is to avoid paying for RetinaFace at all — asserting on ``video is None`` alone
+    would still pass if the runner were built and then discarded.
+    """
+    from cosmos_framework.auxiliary.guardrail.common import presets
+
+    text_runner = cast(Any, object())
+    video_builds: list[bool] = []
+    monkeypatch.setattr(presets, "create_text_guardrail_runner", lambda **_: text_runner)
+    monkeypatch.setattr(presets, "create_video_guardrail_runner", lambda **_: video_builds.append(True))
+
+    guardrail_args = GuardrailArgs(guardrails=True, offload_guardrail_models=False)
+    runners = GuardrailRunners.create(guardrail_args, include_video=False)
+
+    assert runners.text is text_runner
+    assert runners.video is None
+    assert video_builds == []
+
+
+def test_guardrail_runners_build_video_by_default(monkeypatch: pytest.MonkeyPatch) -> None:
+    from cosmos_framework.auxiliary.guardrail.common import presets
+
+    video_runner = cast(Any, object())
+    monkeypatch.setattr(presets, "create_text_guardrail_runner", lambda **_: cast(Any, object()))
+    monkeypatch.setattr(presets, "create_video_guardrail_runner", lambda **_: video_runner)
+
+    guardrail_args = GuardrailArgs(guardrails=True, offload_guardrail_models=False)
+    runners = GuardrailRunners.create(guardrail_args)
+
+    assert runners.video is video_runner
+
+
 def test_guardrail_runners() -> None:
     from cosmos_framework.auxiliary.guardrail.common import presets
 

cosmos_framework/model/generator/mot/include_gen_pathway_test.py

@@ -0,0 +1,123 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+"""Unit tests for ``include_gen_pathway``.
+
+Reasoner-only inference never runs the MoT generation tower, so it can leave the
+``*_moe_gen`` duplicates unbuilt.  ``dcp.load`` is pull-based (it requests only
+the keys the live model exposes), so a model built without those modules also
+never reads their tensors off disk.
+
+The flag defaults to ``True`` everywhere; these tests pin both the disabled
+behaviour and the unchanged default.
+"""
+
+import torch.nn as nn
+
+from cosmos_framework.model.generator.mot.unified_mot import (
+    LayerTypes,
+    MoTDecoderLayer,
+    Nemotron3DenseVLMoTConfig,
+    PackedAttentionMoT,
+)
+from cosmos_framework.model.generator.reasoner.nemotron_3_dense_vl.configuration_nemotron_3_dense_vl import (
+    Nemotron3DenseVLTextConfig,
+)
+
+NUM_Q_HEADS = 4
+NUM_KV_HEADS = 2
+HEAD_DIM = 16
+
+# Every generation-pathway module gated by the flag, by owning class.
+_ATTN_GEN_MODULES = (
+    "q_proj_moe_gen",
+    "k_proj_moe_gen",
+    "v_proj_moe_gen",
+    "o_proj_moe_gen",
+    "q_norm_moe_gen",
+    "k_norm_moe_gen",
+)
+_LAYER_GEN_MODULES = (
+    "mlp_moe_gen",
+    "input_layernorm_moe_gen",
+    "post_attention_layernorm_moe_gen",
+)
+
+
+def _tiny_config() -> Nemotron3DenseVLTextConfig:
+    return Nemotron3DenseVLTextConfig(
+        hidden_size=NUM_Q_HEADS * HEAD_DIM,
+        num_attention_heads=NUM_Q_HEADS,
+        num_key_value_heads=NUM_KV_HEADS,
+        num_hidden_layers=1,
+        attention_bias=False,
+    )
+
+
+def _make_attention(*, include_gen_pathway: bool | None = None) -> PackedAttentionMoT:
+    kwargs = {} if include_gen_pathway is None else {"include_gen_pathway": include_gen_pathway}
+    return PackedAttentionMoT(
+        _tiny_config(),
+        layer_idx=0,
+        layer_types=LayerTypes("nemotron_dense"),
+        qk_norm_for_text=False,
+        qk_norm_for_diffusion=True,
+        **kwargs,
+    )
+
+
+def _make_layer(*, include_gen_pathway: bool | None = None) -> MoTDecoderLayer:
+    kwargs = {} if include_gen_pathway is None else {"include_gen_pathway": include_gen_pathway}
+    return MoTDecoderLayer(
+        config=_tiny_config(),
+        layer_idx=0,
+        layer_types=LayerTypes("nemotron_dense"),
+        qk_norm_for_text=False,
+        qk_norm_for_diffusion=True,
+        **kwargs,
+    )
+
+
+def test_attention_omits_gen_modules_when_disabled() -> None:
+    attn = _make_attention(include_gen_pathway=False)
+
+    for name in _ATTN_GEN_MODULES:
+        assert not hasattr(attn, name), f"{name} should not be built when include_gen_pathway=False"
+    # The cross-attention K norm only exists to serve the generation pathway.
+    assert attn.k_norm_und_for_gen is None
+    # The understanding pathway is untouched.
+    assert isinstance(attn.q_proj, nn.Linear)
+    assert isinstance(attn.o_proj, nn.Linear)
+
+
+def test_attention_builds_gen_modules_by_default() -> None:
+    attn = _make_attention()
+
+    for name in _ATTN_GEN_MODULES:
+        assert hasattr(attn, name), f"{name} must still be built by default"
+
+
+def test_decoder_layer_omits_gen_modules_when_disabled() -> None:
+    layer = _make_layer(include_gen_pathway=False)
+
+    for name in _LAYER_GEN_MODULES:
+        assert not hasattr(layer, name), f"{name} should not be built when include_gen_pathway=False"
+    # Nothing anywhere in the layer (including the nested attention) carries gen weights.
+    assert not [name for name, _ in layer.named_parameters() if "moe_gen" in name]
+    # The understanding pathway still has its full parameter set.
+    assert [name for name, _ in layer.named_parameters() if name.startswith("mlp.")]
+
+
+def test_decoder_layer_builds_gen_modules_by_default() -> None:
+    layer = _make_layer()
+
+    for name in _LAYER_GEN_MODULES:
+        assert hasattr(layer, name), f"{name} must still be built by default"
+    assert [name for name, _ in layer.named_parameters() if "moe_gen" in name]
+
+
+def test_mot_config_includes_gen_pathway_by_default() -> None:
+    assert Nemotron3DenseVLMoTConfig({}).include_gen_pathway is True
+
+
+def test_mot_config_forwards_disabled_flag() -> None:
+    assert Nemotron3DenseVLMoTConfig({}, include_gen_pathway=False).include_gen_pathway is False