From 8b1beabb2332d8a4666cc1f5e3a7a8498ca272c1 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Thu, 30 Jul 2026 00:05:09 -0700 Subject: [PATCH 1/4] feat(mot): add include_gen_pathway to skip the generation tower Reasoner-only inference never runs the MoT generation tower, so it can leave the ten *_moe_gen modules unbuilt. dcp.load is pull-based, so a model built without them also never reads their tensors off disk. The flag defaults to True at every level, leaving all existing callers and the training path unchanged. _impl_forward asserts the pathway is present so a misconfigured generative run fails loudly instead of on a missing attribute. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: liang.feng --- .../generator/mot/include_gen_pathway_test.py | 123 ++++++++++++++++++ .../model/generator/mot/unified_mot.py | 106 ++++++++++----- 2 files changed, 196 insertions(+), 33 deletions(-) create mode 100644 cosmos_framework/model/generator/mot/include_gen_pathway_test.py diff --git a/cosmos_framework/model/generator/mot/include_gen_pathway_test.py b/cosmos_framework/model/generator/mot/include_gen_pathway_test.py new file mode 100644 index 00000000..2e767bba --- /dev/null +++ b/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 diff --git a/cosmos_framework/model/generator/mot/unified_mot.py b/cosmos_framework/model/generator/mot/unified_mot.py index cd539702..22a7c3cd 100644 --- a/cosmos_framework/model/generator/mot/unified_mot.py +++ b/cosmos_framework/model/generator/mot/unified_mot.py @@ -236,6 +236,7 @@ def __init__( qk_norm_for_text: bool = True, qk_norm_for_diffusion: bool = True, include_visual: bool = False, + include_gen_pathway: bool = True, gen_noisy_gating: bool = False, gen_cosine_router_config: CosineRouterConfig | None = None, gen_aux_loss_free_load_balancing_config: AuxLossFreeLoadBalancingConfig | None = None, @@ -247,6 +248,9 @@ def __init__( self.qk_norm_for_text = qk_norm_for_text self.qk_norm_for_diffusion = qk_norm_for_diffusion self.include_visual = include_visual + # Build the MoT generation tower (the ``*_moe_gen`` duplicates). Reasoner-only + # inference disables this; every other caller keeps the default and is unchanged. + self.include_gen_pathway = include_gen_pathway # Noisy top-k gating on the generation-tower MoE blocks (Shazeer 2017). # Gen-tower only; the understanding tower never receives this flag. self.gen_noisy_gating = gen_noisy_gating @@ -485,9 +489,11 @@ def __init__( qk_norm_for_text: bool, qk_norm_for_diffusion: bool, use_und_k_norm_for_gen: bool = False, + include_gen_pathway: bool = True, ): super().__init__() self.config = config + self.include_gen_pathway = include_gen_pathway self.layer_idx = layer_idx self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) self.hidden_size = config.hidden_size @@ -513,13 +519,15 @@ def __init__( self.q_norm = nn.Identity() self.k_norm = nn.Identity() - # Generation pathway QK norm - if qk_norm_for_diffusion: - self.q_norm_moe_gen = layer_types.rms_norm(self.head_dim, eps=eps) - self.k_norm_moe_gen = layer_types.rms_norm(self.head_dim, eps=eps) - else: - self.q_norm_moe_gen = nn.Identity() - self.k_norm_moe_gen = nn.Identity() + # Generation pathway QK norm. Everything below this point belongs to the + # generation tower and is skipped wholesale when it is not built. + if include_gen_pathway: + if qk_norm_for_diffusion: + self.q_norm_moe_gen = layer_types.rms_norm(self.head_dim, eps=eps) + self.k_norm_moe_gen = layer_types.rms_norm(self.head_dim, eps=eps) + else: + self.q_norm_moe_gen = nn.Identity() + self.k_norm_moe_gen = nn.Identity() # Cross-attention K norm: normalises und K tokens seen by the generator in the # gen→und cross-attention path. Only needed when the generation pathway has QK @@ -529,24 +537,26 @@ def __init__( # uncontrolled magnitude and dominates attention over the gen self-attention path. # When both pathways share the same QK norm (or neither has one) k_norm_und_for_gen # is None and the standard packed K tensor is used for all paths unchanged. - if use_und_k_norm_for_gen and qk_norm_for_diffusion and not qk_norm_for_text: + # It serves the generation pathway only, so it is None whenever that tower is absent. + if include_gen_pathway and use_und_k_norm_for_gen and qk_norm_for_diffusion and not qk_norm_for_text: self.k_norm_und_for_gen: nn.Module | None = layer_types.rms_norm(self.head_dim, eps=eps) else: self.k_norm_und_for_gen = None # Generation pathway linear projections - self.q_proj_moe_gen = nn.Linear( - self.hidden_size, self.num_attention_heads * self.head_dim, bias=config.attention_bias - ) - self.k_proj_moe_gen = nn.Linear( - self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias - ) - self.v_proj_moe_gen = nn.Linear( - self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias - ) - self.o_proj_moe_gen = nn.Linear( - self.num_attention_heads * self.head_dim, self.hidden_size, bias=config.attention_bias - ) + if include_gen_pathway: + self.q_proj_moe_gen = nn.Linear( + self.hidden_size, self.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj_moe_gen = nn.Linear( + self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj_moe_gen = nn.Linear( + self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj_moe_gen = nn.Linear( + self.num_attention_heads * self.head_dim, self.hidden_size, bias=config.attention_bias + ) self._apply_rotary_pos_emb = layer_types.apply_rotary_pos_emb self.dispatch_attention_fn = dispatch_attention @@ -823,6 +833,7 @@ def _impl_init( gen_noisy_gating: bool = False, gen_cosine_router_config: CosineRouterConfig | None = None, gen_aux_loss_free_load_balancing_config: AuxLossFreeLoadBalancingConfig | None = None, + include_gen_pathway: bool = True, ) -> None: """Shared ``__init__`` body for the three MoT text-model variants. @@ -832,6 +843,9 @@ def _impl_init( """ self.padding_idx = getattr(config, "pad_token_id", None) self.vocab_size = config.vocab_size + # Read back by ``_impl_forward``'s guard: the joint generation forward cannot + # run on a model built without the generation tower. + self.include_gen_pathway = include_gen_pathway self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) @@ -848,13 +862,15 @@ def _impl_init( gen_noisy_gating=gen_noisy_gating, gen_cosine_router_config=gen_cosine_router_config, gen_aux_loss_free_load_balancing_config=gen_aux_loss_free_load_balancing_config, + include_gen_pathway=include_gen_pathway, ) ) # Reasoner-pathway final norm. self.norm = layer_types.rms_norm(config.hidden_size, eps=config.rms_norm_eps) - # Generation-pathway final norm (parallel to ``self.norm``). - self.norm_moe_gen = layer_types.rms_norm(config.hidden_size, eps=config.rms_norm_eps) + if include_gen_pathway: + # Generation-pathway final norm (parallel to ``self.norm``). + self.norm_moe_gen = layer_types.rms_norm(config.hidden_size, eps=config.rms_norm_eps) # Rotary embedding (text-only optimized) self.rotary_emb = layer_types.rotary_embedding(config) @@ -907,6 +923,15 @@ def _impl_forward( forward passes. """ + # The joint forward drives both towers, so it cannot run on a model built + # without the generation pathway. Reasoner-only decoding goes through + # ``_impl_reasoner_forward`` instead and is unaffected. + assert getattr(self, "include_gen_pathway", True), ( + "The joint generation forward requires the MoT generation pathway, but this model was " + "built with include_gen_pathway=False (reasoner-only). Drop the " + "model.config.vlm_config.model_instance.config.include_gen_pathway=false override." + ) + # Create position embeddings (Qwen3 style) - squeeze once at model level # tensor below is only used for its dtype and device device, dtype = get_device_and_dtype(pack) @@ -1030,9 +1055,11 @@ def __init__( gen_noisy_gating: bool = False, gen_cosine_router_config: CosineRouterConfig | None = None, gen_aux_loss_free_load_balancing_config: AuxLossFreeLoadBalancingConfig | None = None, + include_gen_pathway: bool = True, ) -> None: super().__init__() self.hidden_size = config.hidden_size + self.include_gen_pathway = include_gen_pathway self.self_attn = PackedAttentionMoT( config, layer_types=layer_types, @@ -1040,6 +1067,7 @@ def __init__( qk_norm_for_text=qk_norm_for_text, qk_norm_for_diffusion=qk_norm_for_diffusion, use_und_k_norm_for_gen=use_und_k_norm_for_gen, + include_gen_pathway=include_gen_pathway, ) if ( @@ -1048,22 +1076,25 @@ def __init__( and (config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0) ): self.mlp = Qwen3VLMoeTextSparseMoeBlock(config) - # Noisy gating, the cosine router, and aux-loss-free load balancing - # are gen-tower only. - self.mlp_moe_gen = Qwen3VLMoeTextSparseMoeBlock( - config, - noisy_gating=gen_noisy_gating, - cosine_router_config=gen_cosine_router_config, - aux_loss_free_load_balancing_config=gen_aux_loss_free_load_balancing_config, - ) + if include_gen_pathway: + # Noisy gating, the cosine router, and aux-loss-free load balancing + # are gen-tower only. + self.mlp_moe_gen = Qwen3VLMoeTextSparseMoeBlock( + config, + noisy_gating=gen_noisy_gating, + cosine_router_config=gen_cosine_router_config, + aux_loss_free_load_balancing_config=gen_aux_loss_free_load_balancing_config, + ) else: self.mlp = layer_types.mlp(config) - self.mlp_moe_gen = layer_types.mlp(config) + if include_gen_pathway: + self.mlp_moe_gen = layer_types.mlp(config) self.input_layernorm = layer_types.rms_norm(config.hidden_size, eps=config.rms_norm_eps) - self.input_layernorm_moe_gen = layer_types.rms_norm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = layer_types.rms_norm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm_moe_gen = layer_types.rms_norm(config.hidden_size, eps=config.rms_norm_eps) + if include_gen_pathway: + self.input_layernorm_moe_gen = layer_types.rms_norm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm_moe_gen = layer_types.rms_norm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, @@ -1272,6 +1303,7 @@ def __init__( qk_norm_for_text: bool, qk_norm_for_diffusion: bool, use_und_k_norm_for_gen: bool, + include_gen_pathway: bool = True, ): super().__init__(config) _impl_init( @@ -1281,6 +1313,7 @@ def __init__( qk_norm_for_text=qk_norm_for_text, qk_norm_for_diffusion=qk_norm_for_diffusion, use_und_k_norm_for_gen=use_und_k_norm_for_gen, + include_gen_pathway=include_gen_pathway, ) def init_taylorseer(self, cache_dic=None, current=None): @@ -1310,6 +1343,7 @@ def __init__( gen_noisy_gating: bool = False, gen_cosine_router_config: CosineRouterConfig | None = None, gen_aux_loss_free_load_balancing_config: AuxLossFreeLoadBalancingConfig | None = None, + include_gen_pathway: bool = True, ) -> None: super().__init__(config) _impl_init( @@ -1322,6 +1356,7 @@ def __init__( gen_noisy_gating=gen_noisy_gating, gen_cosine_router_config=gen_cosine_router_config, gen_aux_loss_free_load_balancing_config=gen_aux_loss_free_load_balancing_config, + include_gen_pathway=include_gen_pathway, ) def init_taylorseer(self, cache_dic=None, current=None): @@ -1348,6 +1383,7 @@ def __init__( qk_norm_for_text: bool, qk_norm_for_diffusion: bool, use_und_k_norm_for_gen: bool, + include_gen_pathway: bool = True, ): super().__init__(config) _impl_init( @@ -1357,6 +1393,7 @@ def __init__( qk_norm_for_text=qk_norm_for_text, qk_norm_for_diffusion=qk_norm_for_diffusion, use_und_k_norm_for_gen=use_und_k_norm_for_gen, + include_gen_pathway=include_gen_pathway, ) def init_taylorseer(self, cache_dic=None, current=None) -> None: @@ -2037,6 +2074,7 @@ def __init__(self, config: Qwen3VLMoTConfig): qk_norm_for_text=config.qk_norm_for_text, qk_norm_for_diffusion=config.qk_norm_for_diffusion, use_und_k_norm_for_gen=getattr(config, "use_und_k_norm_for_gen", False), + include_gen_pathway=getattr(config, "include_gen_pathway", True), ) self.vocab_size = text_config.vocab_size self.lm_head = nn.Linear(text_config.hidden_size, text_config.vocab_size, bias=False) @@ -2181,6 +2219,7 @@ def __init__(self, config: Qwen3VLMoeMoTConfig): qk_norm_for_text=config.qk_norm_for_text, qk_norm_for_diffusion=config.qk_norm_for_diffusion, use_und_k_norm_for_gen=getattr(config, "use_und_k_norm_for_gen", False), + include_gen_pathway=getattr(config, "include_gen_pathway", True), gen_noisy_gating=config.gen_noisy_gating, gen_cosine_router_config=getattr(config, "gen_cosine_router_config", None), gen_aux_loss_free_load_balancing_config=config.gen_aux_loss_free_load_balancing_config, @@ -2372,6 +2411,7 @@ def __init__(self, config: Nemotron3DenseVLMoTConfig) -> None: qk_norm_for_text=config.qk_norm_for_text, qk_norm_for_diffusion=config.qk_norm_for_diffusion, use_und_k_norm_for_gen=getattr(config, "use_und_k_norm_for_gen", False), + include_gen_pathway=getattr(config, "include_gen_pathway", True), ) self.vocab_size = text_config.vocab_size self.lm_head = nn.Linear(text_config.hidden_size, text_config.vocab_size, bias=False) From cbc41f3133937a69d227ab16345b9222f850be7d Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Thu, 30 Jul 2026 01:02:51 -0700 Subject: [PATCH 2/4] feat(guardrail): allow skipping the video guardrail runner The video runner only holds RetinaFace (~109 MB, fp32) and is only reached by generative samples, so reasoner-only runs can skip building it entirely. video_guardrail lives on GuardrailArgs rather than GuardrailOverrides so it never becomes a CLI flag: a user-facing switch that disables face blurring on generative runs would be a safety footgun. It is excluded from model_dump() because SetupArgs dumps must round-trip through SetupOverrides, which forbids extra keys. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: liang.feng --- cosmos_framework/inference/common/args.py | 12 +++++++ .../inference/common/inference.py | 26 +++++++++++--- .../inference/common/inference_test.py | 35 +++++++++++++++++++ 3 files changed, 69 insertions(+), 4 deletions(-) diff --git a/cosmos_framework/inference/common/args.py b/cosmos_framework/inference/common/args.py index 83d2fe01..8b081ef6 100644 --- a/cosmos_framework/inference/common/args.py +++ b/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): diff --git a/cosmos_framework/inference/common/inference.py b/cosmos_framework/inference/common/inference.py index 9755f945..e747bcdb 100644 --- a/cosmos_framework/inference/common/inference.py +++ b/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}") diff --git a/cosmos_framework/inference/common/inference_test.py b/cosmos_framework/inference/common/inference_test.py index d9872256..df68c2a6 100644 --- a/cosmos_framework/inference/common/inference_test.py +++ b/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 From df96316a62082a622263387626655a22d1fce237 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Thu, 30 Jul 2026 01:02:52 -0700 Subject: [PATCH 3/4] feat(inference): drop every generation-side module for reasoner-only runs Extends the reasoner-only hook added in #132 from the generation VAE to the whole generation side: the MoT generation tower (include_gen_pathway), the VFM-level generation heads (vision_gen), and the video guardrail runner. On Cosmos3-Edge this leaves ~3.4 GiB/GPU unallocated (VAE 1.31 + MoT tower 1.97 + VFM heads 0.03 + RetinaFace 0.10) and avoids a ~2.6 GiB cold-cache download of Wan2.2_VAE.pth. The three config flags are orthogonal and each defaults to the generation-enabled value, so dropping any one entry from reasoner_only_overrides degrades to the previous behaviour for that module alone. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: liang.feng --- cosmos_framework/inference/args.py | 23 ++++++++++++++++++++ cosmos_framework/inference/args_test.py | 29 +++++++++++++++++++++++++ cosmos_framework/scripts/inference.py | 8 ++++--- 3 files changed, 57 insertions(+), 3 deletions(-) diff --git a/cosmos_framework/inference/args.py b/cosmos_framework/inference/args.py index 755a6b26..70bec903 100644 --- a/cosmos_framework/inference/args.py +++ b/cosmos_framework/inference/args.py @@ -202,6 +202,29 @@ 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_*. + "model.config.vision_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" diff --git a/cosmos_framework/inference/args_test.py b/cosmos_framework/inference/args_test.py index defc8b7c..b10cdd57 100644 --- a/cosmos_framework/inference/args_test.py +++ b/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,34 @@ 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 + + def test_build_parallelism(monkeypatch: pytest.MonkeyPatch): parallelism_args = OmniSetupOverrides( checkpoint_path=DEFAULT_CHECKPOINT_NAME, diff --git a/cosmos_framework/scripts/inference.py b/cosmos_framework/scripts/inference.py index d363c47c..59e70b23 100644 --- a/cosmos_framework/scripts/inference.py +++ b/cosmos_framework/scripts/inference.py @@ -12,7 +12,7 @@ import pydantic import tyro -from cosmos_framework.inference.args import OmniSetupOverrides, is_reasoner_only +from cosmos_framework.inference.args import OmniSetupOverrides, is_reasoner_only, reasoner_only_overrides from cosmos_framework.inference.common.args import SampleOutputs, SetupOverrides, tyro_cli from cosmos_framework.inference.common.init import init_output_dir from cosmos_framework.utils import log @@ -45,8 +45,10 @@ def inference(args: InferenceArgs): ) log.info(f"Loaded {len(sample_overrides_list)} samples") if is_reasoner_only(sample_overrides_list): - setup_args.experiment_overrides.append("model.config.load_vision_tokenizer=false") - log.info("Reasoner-only inputs detected; generation vision tokenizer will not be loaded") + setup_args.experiment_overrides.extend(reasoner_only_overrides()) + # Reasoner output is text, so no frames ever reach the video guardrail. + setup_args.video_guardrail = False + log.info("Reasoner-only inputs detected; generation-side modules will not be loaded") for sample_overrides in sample_overrides_list: assert sample_overrides.name sample_overrides.output_dir = setup_args.output_dir / sample_overrides.name From 877c29d650f670b1302e403cabd301461ce4392a Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Thu, 30 Jul 2026 02:41:46 -0700 Subject: [PATCH 4/4] fix(inference): switch off action/sound generation with vision generation Cosmos3VFMNetworkConfig asserts that action and sound generation each imply vision generation. Edge checkpoints ship with action_gen=true, so overriding vision_gen alone aborted model construction: AssertionError('Action generation requires visual generation! We do NOT support action only training!') Found by the first real reasoner run on GB200; unit tests could not catch it because they assert on module construction, not on config-invariant coherence. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: liang.feng --- cosmos_framework/inference/args.py | 7 ++++++- cosmos_framework/inference/args_test.py | 4 ++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/cosmos_framework/inference/args.py b/cosmos_framework/inference/args.py index 70bec903..ba636817 100644 --- a/cosmos_framework/inference/args.py +++ b/cosmos_framework/inference/args.py @@ -218,8 +218,13 @@ def reasoner_only_overrides() -> list[str]: # 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_*. + # 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", ] diff --git a/cosmos_framework/inference/args_test.py b/cosmos_framework/inference/args_test.py index b10cdd57..17c33891 100644 --- a/cosmos_framework/inference/args_test.py +++ b/cosmos_framework/inference/args_test.py @@ -84,6 +84,10 @@ def test_reasoner_only_overrides_disable_every_generation_side_module(tmp_path: 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):