@@ -27,6 +27,7 @@
 import shlex
 import sys
 from collections.abc import Iterator
+from functools import cache
 from pathlib import Path
 from typing import overload
 
@@ -41,9 +42,11 @@
 CONVERSION_ARGUMENTS = REPO_ROOT / "scripts" / "conversion" / "arguments.py"
 SETUP_CONVERSION = REPO_ROOT / "scripts" / "conversion" / "setup_conversion.py"
 SETUP_INFERENCE = REPO_ROOT / "scripts" / "inference" / "setup_inference.py"
+PublishedCommands = tuple[tuple[str, tuple[str, ...]], ...]
 
 
-def _declared_options(source_path: Path, *, append_only: bool = False, function: str | None = None) -> set[str]:
+@cache
+def _declared_options(source_path: Path, *, append_only: bool = False, function: str | None = None) -> frozenset[str]:
     """Read parser options, including called argument helpers in Bridge source."""
     options: set[str] = set()
     tree = ast.parse(source_path.read_text(encoding="utf-8"))
@@ -77,7 +80,15 @@ def _declared_options(source_path: Path, *, append_only: bool = False, function:
             if isinstance(argument, ast.Constant) and isinstance(argument.value, str):
                 if argument.value.startswith("-"):
                     options.add(argument.value)
-    return options
+    return frozenset(options)
+
+
+@pytest.fixture(autouse=True)
+def _clear_declared_options_cache() -> Iterator[None]:
+    # Reuse parser results within a check, but never across tests that edit source.
+    _declared_options.cache_clear()
+    yield
+    _declared_options.cache_clear()
 
 
 def _inference_tasks() -> dict[str, Path]:
@@ -128,15 +139,21 @@ def _published_commands() -> Iterator[tuple[str, list[str]]]:
                 yield f"{card.parent.name}{location}", tokens
 
 
+@pytest.fixture(scope="module")
+def published_commands() -> PublishedCommands:
+    """Parse every shipped command once, with immutable argv shared by the checks."""
+    return tuple((leaf, tuple(tokens)) for leaf, tokens in _published_commands())
+
+
 @overload
-def _flag_value(tokens: list[str], flag: str, default: int) -> int: ...
+def _flag_value(tokens: tuple[str, ...], flag: str, default: int) -> int: ...
 
 
 @overload
-def _flag_value(tokens: list[str], flag: str, default: None) -> int | None: ...
+def _flag_value(tokens: tuple[str, ...], flag: str, default: None) -> int | None: ...
 
 
-def _flag_value(tokens: list[str], flag: str, default: int | None) -> int | None:
+def _flag_value(tokens: tuple[str, ...], flag: str, default: int | None) -> int | None:
     """Return the last integer value passed to a flag, or the launcher default."""
     value = default
     for index, token in enumerate(tokens[:-1]):
@@ -145,17 +162,19 @@ def _flag_value(tokens: list[str], flag: str, default: int | None) -> int | None
     return value
 
 
-def _selected_task(tokens: list[str]) -> str:
+def _selected_task(tokens: tuple[str, ...]) -> str:
     task = "text-generation"
     for index, token in enumerate(tokens[:-1]):
         if token == "--task":
             task = tokens[index + 1]
     return task
 
 
-def test_published_gpu_conversion_commands_decompose_over_the_published_world_size() -> None:
+def test_published_gpu_conversion_commands_decompose_over_the_published_world_size(
+    published_commands: PublishedCommands,
+) -> None:
     offenders = []
-    for leaf, tokens in _published_commands():
+    for leaf, tokens in published_commands:
         if "convert.sh" not in tokens[0]:
             continue
         device = "gpu"
@@ -178,12 +197,12 @@ def test_published_gpu_conversion_commands_decompose_over_the_published_world_si
     assert not offenders, "setup_conversion.py refuses these published commands: " + "; ".join(offenders)
 
 
-def test_published_inference_commands_carry_no_conversion_launcher_flag() -> None:
+def test_published_inference_commands_carry_no_conversion_launcher_flag(published_commands: PublishedCommands) -> None:
     launcher_options = _declared_options(SETUP_INFERENCE)
     conversion_options = _declared_options(CONVERSION_ARGUMENTS) | _declared_options(SETUP_CONVERSION)
     tasks = _inference_tasks()
     offenders = []
-    for leaf, tokens in _published_commands():
+    for leaf, tokens in published_commands:
         if "infer.sh" not in tokens[0]:
             continue
         accepted = launcher_options | _declared_options(tasks[_selected_task(tokens)])
@@ -196,12 +215,12 @@ def test_published_inference_commands_carry_no_conversion_launcher_flag() -> Non
     assert not offenders, "these published inference commands die in the container: " + "; ".join(offenders)
 
 
-def test_published_card_commands_do_not_repeat_a_flag() -> None:
+def test_published_card_commands_do_not_repeat_a_flag(published_commands: PublishedCommands) -> None:
     repeatable = _declared_options(CONVERSION_ARGUMENTS, append_only=True) | _declared_options(
         SETUP_INFERENCE, append_only=True
     )
     offenders = []
-    for leaf, tokens in _published_commands():
+    for leaf, tokens in published_commands:
         flags = [token.split("=", 1)[0] for token in tokens[1:] if token.startswith("--")]
         repeated = sorted({flag for flag in flags if flags.count(flag) > 1 and flag not in repeatable})
         if repeated:
@@ -252,10 +271,7 @@ def test_declared_options_follow_only_called_helpers(tmp_path: Path, monkeypatch
         ),
     ],
 )
-def test_published_command_checks_reject_original_errors(
-    command: str, check_name: str, message: str, monkeypatch: pytest.MonkeyPatch
-) -> None:
+def test_published_command_checks_reject_original_errors(command: str, check_name: str, message: str) -> None:
     module = sys.modules[__name__]
-    monkeypatch.setattr(module, "_published_commands", lambda: iter([("regression", shlex.split(command))]))
     with pytest.raises(AssertionError, match=message):
-        getattr(module, check_name)()
+        getattr(module, check_name)((("regression", tuple(shlex.split(command))),))
@@ -581,7 +581,9 @@ def test_nemotron_omni_vision_projection_uses_squared_relu():
     values = torch.tensor([-2.0, 0.0, 3.0])
 
     assert vision_projection_config.activation_func is squared_relu
-    assert torch.equal(vision_projection_config.activation_func(values), torch.tensor([0.0, 0.0, 9.0]))
+    # This checks configuration and arithmetic; the model test covers compiled execution.
+    activation = torch.compiler.disable(vision_projection_config.activation_func)
+    assert torch.equal(activation(values), torch.tensor([0.0, 0.0, 9.0]))
 
 
 def test_nemotron_omni_providers_use_contract_specific_cpe_defaults():
 
@@ -152,10 +152,32 @@ def fake_all_gather(outputs, _local_counts, group):
 
 @pytest.fixture(scope="module")
 def hf_config():
-    """Create a local HuggingFace config once for all tests."""
-    config = Qwen3VLMoeConfig()
-    config.vision_config.out_hidden_size = config.text_config.hidden_size
-    return config
+    """Keep real vision, deepstack and language paths with small, compatible shapes."""
+    return Qwen3VLMoeConfig(
+        text_config={
+            "num_hidden_layers": 2,
+            "hidden_size": 128,
+            "intermediate_size": 256,
+            "num_attention_heads": 4,
+            "num_key_value_heads": 2,
+            "head_dim": 32,
+            "vocab_size": 256,
+            "max_position_embeddings": 256,
+            "rope_parameters": {"rope_type": "default", "rope_theta": 5000000.0, "mrope_section": [4, 6, 6]},
+        },
+        vision_config={
+            "depth": 2,
+            "hidden_size": 128,
+            "intermediate_size": 256,
+            "num_heads": 4,
+            "out_hidden_size": 128,
+            "deepstack_visual_indexes": [0, 1],
+        },
+        image_token_id=240,
+        video_token_id=241,
+        vision_start_token_id=242,
+        vision_end_token_id=243,
+    )
 
 
 class TestQwen3VLModel:
@@ -229,8 +251,7 @@ def get_vision_transformer_config(hf_config):
     def get_language_transformer_config(hf_config):
         """Create a language transformer config for testing.
 
-        Uses actual Qwen3-VL-30B-A3B model sizes to ensure compatibility
-        with the vision model output (2048 hidden size).
+        Match the tiny HuggingFace fixture's language and vision dimensions.
 
         Args:
             hf_config: HuggingFace config object.
@@ -239,16 +260,18 @@ def get_language_transformer_config(hf_config):
             Qwen3VLTransformerConfig: Configuration for the language model.
         """
         return Qwen3VLTransformerConfig(
-            # Use actual model dimensions from HF config
-            num_layers=4,  # Reduced for testing (actual: hf_config.text_config.num_hidden_layers)
-            hidden_size=hf_config.text_config.hidden_size,  # Must match vision output: 2048
+            num_layers=hf_config.text_config.num_hidden_layers,
+            hidden_size=hf_config.text_config.hidden_size,
             num_attention_heads=hf_config.text_config.num_attention_heads,
             num_query_groups=hf_config.text_config.num_key_value_heads,
             kv_channels=hf_config.text_config.hidden_size // hf_config.text_config.num_attention_heads,
             ffn_hidden_size=hf_config.text_config.intermediate_size,
             # Qwen3-VL specific
             vocab_size=hf_config.text_config.vocab_size,
             language_max_sequence_length=hf_config.text_config.max_position_embeddings,
+            image_token_id=hf_config.image_token_id,
+            video_token_id=hf_config.video_token_id,
+            vision_start_token_id=hf_config.vision_start_token_id,
             # Vision parameters
             patch_size=hf_config.vision_config.patch_size,
             temporal_patch_size=hf_config.vision_config.temporal_patch_size,
@@ -513,17 +536,17 @@ def test_set_input_tensor(self, hf_config):
 
         if torch.cuda.is_available():
             model_pre.to("cuda")
-            test_tensor = torch.randn(2, 4, 2048).cuda()
+            test_tensor = torch.randn(2, 4, language_transformer_config.hidden_size).cuda()
         else:
-            test_tensor = torch.randn(2, 4, 2048)
+            test_tensor = torch.randn(2, 4, language_transformer_config.hidden_size)
 
         # Test with single tensor (not a list)
         model_pre.set_input_tensor(test_tensor)
-        assert model_pre.encoder_hidden_state is not None
+        assert model_pre.encoder_hidden_state is test_tensor
 
         # Test with list of tensors
         model_pre.set_input_tensor([test_tensor])
-        assert model_pre.encoder_hidden_state is not None
+        assert model_pre.encoder_hidden_state is test_tensor
 
         # Test with pre_process=False
         model_no_pre = Qwen3VLModel(
@@ -543,7 +566,7 @@ def test_set_input_tensor(self, hf_config):
 
         # This should set the input tensor on the language model instead
         model_no_pre.set_input_tensor([test_tensor])
-        # No assertion here as it sets internal state
+        assert model_no_pre.language_model.decoder.input_tensor is test_tensor
 
     @staticmethod
     def _attach_dist_train(language_transformer_config, vision_to_llm_dp_ratio: int = 1) -> None:
@@ -651,6 +674,10 @@ def test_forward_dist_train_decoder_only(self, hf_config):
             )
             vision_payload = enc_out["vision_module"].detach()
             decoder.set_input_tensor([{"vision_module": vision_payload}])
+            assert decoder.vision_embeds is not None
+            assert len(decoder.deepstack_feature_lists) == len(vision_transformer_config.deepstack_visual_indexes)
+            received_payload = torch.cat((*decoder.deepstack_feature_lists, decoder.vision_embeds), dim=0)
+            assert torch.equal(received_payload, vision_payload.flatten(0, 1))
             out = decoder(
                 input_ids=batch["input_ids"],
                 attention_mask=batch["attention_mask"],
 
@@ -14,7 +14,6 @@
 
 """Unit tests for PEFT-specific recompute helpers."""
 
-import gc
 import weakref
 from types import SimpleNamespace
 
@@ -116,7 +115,7 @@ def test_recompute_patch_registry_tracks_model_lifetime(monkeypatch):
     maybe_enable_recompute_inputs_grad(first_model)
     first_model_ref = weakref.ref(first_model)
     del first_model
-    gc.collect()
+    # The model has no reference cycle; the registry must not keep it alive.
     assert first_model_ref() is None
 
     second_model = DummyModel(DummyHybridStack)
 
@@ -1158,7 +1158,10 @@ def test_no_checkpoint_when_disabled(self, mock_check_nvrx, mock_barrier_log, mo
     @patch("megatron.bridge.training.train.save_checkpoint_and_time")
     @patch("megatron.bridge.training.train.barrier_and_log")
     @patch("megatron.bridge.training.train.check_nvrx_straggler_detection")
-    def test_multiple_exit_conditions_signal_wins(self, mock_check_nvrx, mock_barrier_log, mock_save_checkpoint):
+    @patch("megatron.bridge.training.train.end_active_mlflow_run")
+    def test_multiple_exit_conditions_signal_wins(
+        self, mock_end_mlflow, mock_check_nvrx, mock_barrier_log, mock_save_checkpoint
+    ):
         """Test that signal handler exit takes precedence over other exit conditions."""
         mock_check_nvrx.return_value = True  # Straggler detection would also trigger
 
@@ -1181,6 +1184,10 @@ def test_multiple_exit_conditions_signal_wins(self, mock_check_nvrx, mock_barrie
         # Verify signal handler message was logged, not straggler message
         mock_barrier_log.assert_called_once_with("exiting program after receiving SIGTERM.")
 
+        mock_end_mlflow.assert_called_once_with("KILLED")
+        mock_check_nvrx.assert_not_called()
+        mock_save_checkpoint.assert_called_once()
+
     @patch("megatron.bridge.training.train.save_checkpoint_and_time")
     @patch("megatron.bridge.training.train.barrier_and_log")
     @patch("megatron.bridge.training.train.check_nvrx_straggler_detection")