fix(nemotron-vl): tighten model support scope and adapter placement

Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

docs/parallelisms.md

@@ -244,12 +244,6 @@ path performs a host scalar synchronization that is not capture-safe. They must
 provide equal per-rank dispatch shapes. Disable CUDA graphs when packed THD token
 counts can differ so the recipe can enable safe padding.
 
-Fixed-shape eager workloads may set
-`moe_hybridep_assume_equal_dispatch_inputs=true` to preserve an explicitly
-disabled padding setting and avoid the safety collective. This is a strict
-runtime contract: every rank in the HybridEP group must dispatch the same token
-count. Do not enable it for variable-length or dynamically packed batches.
-
 **GPU Architecture Requirements:**
 
 - **DeepEP**: Ampere (SM 8.x), Hopper (SM 9.x), B200, B300

examples/model_verification_cards/nemotron-3.5-super-vl-120b-a12b/card.yaml

@@ -9,6 +9,9 @@ summary: >
   The recorded results below used one MTP prediction depth. Verification is
   being refreshed for the corrected two-depth shared-block configuration;
   these historical results do not verify that new training objective.
+  They also predate the current FP32 optimizer-state and weight-decay defaults
+  and language-only PEFT adapter targets. Each verified status below applies
+  only to its recorded historical Bridge commit, not the current recipes.
   Real-data BF16 pretraining verification on 64 H100s and 64 GB200s
   uses DataComp image-caption samples, natural dropless expert routing, and a
   learned vision-pretrained checkpoint; both 120-step runs and their final
@@ -438,7 +441,9 @@ items:
         produced finite loss with zero skipped or NaN iterations. The
         step-100 adapter contains all 16 distributed shards and two Energon
         states and reloaded successfully. This dataset does not contain video
-        samples.
+        samples. This historical recipe used unqualified adapter targets that
+        also matched vision and projector modules; it does not verify the
+        current language-only adapter placement.
 
     GB200:
       status: verified
@@ -478,7 +483,9 @@ items:
         produced finite loss with zero skipped or NaN iterations. The
         step-100 adapter contains all 16 distributed shards and eight Energon
         states and reloaded successfully. This dataset does not contain video
-        samples.
+        samples. This historical recipe used unqualified adapter targets that
+        also matched vision and projector modules; it does not verify the
+        current language-only adapter placement.
 
   checkpoint_resume:
     H100:

examples/models/nemotron/nemotron_3/super_vl/README.md

@@ -66,15 +66,15 @@ These settings control whether the pretrained base parameters receive
 gradients. PEFT adapter placement is a separate concern: `peft.target_modules`
 determines which matching modules receive adapters. Freezing a base component
 does not, by itself, prevent an unscoped LoRA target name from adding a
-trainable adapter inside that component. Decoder-only PEFT recipes should use
-language-model-qualified target patterns and verify that no adapter is created
-under `vision_model` or `vision_projection`.
+trainable adapter inside that component. The H100 and GB200 PEFT recipes use
+language-model-qualified target patterns, covering the decoder and MTP while
+excluding `vision_model` and `vision_projection`. Custom PEFT configurations
+must preserve that qualification if language-only adapters are intended.
 
 ## Draft Completion Checklist
 
 - Add verified HF-to-Megatron and Megatron-to-HF conversion commands.
 - Add deterministic image and video inference examples.
 - Add SFT and PEFT dataset requirements and launch examples.
 - Document dynamic-resolution image and temporal-video preprocessing.
-- Document and test decoder-only PEFT target patterns.
 - Link the completed page from the Nemotron 3 example index.

scripts/inference/setup_inference.py

@@ -80,7 +80,7 @@ def _build_parser() -> argparse.ArgumentParser:
     execution.add_argument(
         "--tasks-per-node",
         type=int,
-        help="Inference tasks per node; defaults to one task per requested GPU.",
+        help="Inference tasks per node; defaults to one task per GPU.",
     )
     execution.add_argument("--cpus-per-task", type=int, help="CPUs allocated to each inference task.")
     execution.add_argument("--mem", help="Optional Slurm memory request, such as 64G.")
@@ -199,7 +199,6 @@ def _validate_task_args(task_name: str, inference_args: list[str]) -> None:
 def _build_executor(args: argparse.Namespace, env_names: list[str], mounts: list[str]) -> object:
     """Build the srun-native NeMo-Run Slurm executor."""
     gpu_kwargs = {} if args.no_gpu_resource_request else {"gpus_per_node": args.gpus_per_node}
-    tasks_per_node = args.tasks_per_node if args.tasks_per_node is not None else args.gpus_per_node
     # Slurm's --export=NIL removes the site PATH used to find scontrol and srun
     # on clusters where they are not installed in /usr/bin. Always inherit PATH
     # for the generated batch script, but expose it to the container only when
@@ -209,7 +208,7 @@ def _build_executor(args: argparse.Namespace, env_names: list[str], mounts: list
         account=args.account,
         partition=args.partition,
         nodes=args.nodes,
-        ntasks_per_node=tasks_per_node,
+        ntasks_per_node=args.tasks_per_node or args.gpus_per_node,
         cpus_per_task=args.cpus_per_task,
         mem=args.mem,
         exclusive=True if args.exclusive else None,

skills/create-model-verification-card/scripts/verify_hf_inference.py

@@ -210,10 +210,17 @@ def _load_runtime(args: argparse.Namespace) -> tuple[Any, Any, Any]:
 
     dtype = getattr(torch, args.dtype)
     if args.image:
-        from transformers import AutoModelForImageTextToText, AutoProcessor
+        from transformers import AutoConfig, AutoModelForImageTextToText, AutoModelForMultimodalLM, AutoProcessor
 
+        config = AutoConfig.from_pretrained(args.hf_model, trust_remote_code=args.trust_remote_code)
         processor = AutoProcessor.from_pretrained(args.hf_model, trust_remote_code=args.trust_remote_code)
-        model_cls = AutoModelForImageTextToText
+        auto_map = getattr(config, "auto_map", None) or {}
+        # Keep the broader multimodal loader for built-in models, including Omni.
+        # Some remote-code VLMs register only the image-text auto class.
+        if "AutoModelForImageTextToText" in auto_map and "AutoModelForMultimodalLM" not in auto_map:
+            model_cls = AutoModelForImageTextToText
+        else:
+            model_cls = AutoModelForMultimodalLM
     else:
         from transformers import AutoModelForCausalLM, AutoTokenizer
 
@@ -224,6 +231,8 @@ def _load_runtime(args: argparse.Namespace) -> tuple[Any, Any, Any]:
         "trust_remote_code": args.trust_remote_code,
         "output_loading_info": True,
     }
+    if args.image:
+        model_kwargs["config"] = config
     if args.device_map:
         model_kwargs["device_map"] = args.device_map
     model, loading_info = model_cls.from_pretrained(args.hf_model, **model_kwargs)

src/megatron/bridge/models/transformer_config.py

@@ -123,7 +123,6 @@ def _uses_legacy_full_iteration(value: object) -> bool:
         or getattr(config, "moe_token_dispatcher_type", None) != "flex"
         or getattr(config, "moe_flex_dispatcher_backend", None) != "hybridep"
         or cuda_graphs_enabled
-        or getattr(config, "moe_hybridep_assume_equal_dispatch_inputs", False)
     ):
         return
 
@@ -156,9 +155,6 @@ class TransformerConfig(MCoreTransformerConfig):
         config.finalize()
     """
 
-    moe_hybridep_assume_equal_dispatch_inputs: bool = False
-    """Keep HybridEP uneven-input padding disabled for fixed-shape eager workloads."""
-
     _NO_COPY_KEYS = {"_pg_collection"}
 
     # Generalized tensor-parallel metadata was added after the frozen MCore dev pin.

src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_35_super_vl.py

@@ -40,6 +40,7 @@ def _nemotron_35_super_vl_base() -> ConfigContainer:
     cfg = _sft_common_vlm()
     cfg.model = AutoBridge.from_hf_pretrained(
         NEMOTRON_35_SUPER_VL_HF_MODEL_ID,
+        revision=NEMOTRON_35_SUPER_VL_HF_REVISION,
         trust_remote_code=True,
     ).to_megatron_provider(load_weights=False)
     cfg.model.seq_length = 4096
@@ -163,7 +164,6 @@ def nemotron_35_super_vl_pretrain_64gpu_h100_bf16_config() -> ConfigContainer:
     cfg.model.moe_expert_capacity_factor = None
     cfg.model.moe_pad_expert_input_to_capacity = False
     cfg.model.moe_hybridep_pad_uneven_dispatch_inputs = False
-    cfg.model.moe_hybridep_assume_equal_dispatch_inputs = True
     cfg.model.moe_flex_dispatcher_num_sms = 32
     cfg.ddp.overlap_param_gather = False
     cfg.model.moe_hybridep_num_sms = None
@@ -253,7 +253,12 @@ def nemotron_35_super_vl_peft_16gpu_h100_bf16_config(
         The Super-VL PEFT configuration.
     """
     cfg = _nemotron_35_super_vl_base()
-    target_modules = ["linear_qkv", "linear_proj", "in_proj", "out_proj", "linear_fc1", "linear_fc2"]
+    # Base freeze flags do not prevent PEFT from inserting new adapters. Keep
+    # targets language-qualified so the vision tower and projector stay frozen.
+    target_modules = [
+        f"*language_model.*.{name}"
+        for name in ("linear_qkv", "linear_proj", "in_proj", "out_proj", "linear_fc1", "linear_fc2")
+    ]
     cfg.peft = default_peft_config(
         peft_scheme,
         target_modules=target_modules,

tests/unit_tests/models/test_transformer_config.py

@@ -92,19 +92,6 @@ def test_cuda_graph_config_does_not_require_padding_field(self):
 
         _enable_safe_hybridep_dispatch(cfg, uses_thd=True)
 
-    def test_equal_dispatch_input_contract_preserves_disabled_padding(self):
-        cfg = SimpleNamespace(
-            moe_token_dispatcher_type="flex",
-            moe_flex_dispatcher_backend="hybridep",
-            moe_hybridep_pad_variable_tokens=False,
-            moe_hybridep_assume_equal_dispatch_inputs=True,
-            cuda_graph_impl="none",
-        )
-
-        _enable_safe_hybridep_dispatch(cfg)
-
-        assert cfg.moe_hybridep_pad_variable_tokens is False
-
 
 # ---------------------------------------------------------------------------
 # _resolve_string_fields

tests/unit_tests/recipes/nemotron_omni/test_nemotron_omni_recipes.py

@@ -29,6 +29,7 @@
 )
 from megatron.bridge.data.collators.registry import resolve_model_collate
 from megatron.bridge.models.nemotron_omni.data.collate_fn import nemotron_omni_expanded_collate_fn
+from megatron.bridge.peft.lora_layers import LoRALinear
 from megatron.bridge.training.config import ConfigContainer
 from megatron.bridge.training.mixed_precision import get_mixed_precision_config
 from megatron.bridge.training.optim import _get_scheduler
@@ -526,7 +527,7 @@ def test_super_vl_sft_recipe_reuses_omni_data_and_super_training_stack(fake_proc
 
     assert isinstance(cfg, ConfigContainer)
     assert _FakeAutoBridge.hf_path == _TEST_SUPER_VL_HF_ID
-    assert _FakeAutoBridge.kwargs == {"trust_remote_code": True}
+    assert _FakeAutoBridge.kwargs == {"revision": _SUPER_VL_HF_REVISION, "trust_remote_code": True}
     assert _FakeAutoBridge.load_weights is False
 
     assert isinstance(cfg.dataset, EnergonDatasetConfig)
@@ -593,12 +594,12 @@ def test_super_vl_peft_recipe_uses_native_lora_targets_and_frozen_vision(fake_pr
     assert cfg.dataset.pad_to_max_length is True
     assert cfg.dataset.do_validation is False
     assert cfg.peft.target_modules == [
-        "linear_qkv",
-        "linear_proj",
-        "in_proj",
-        "out_proj",
-        "linear_fc1",
-        "linear_fc2",
+        "*language_model.*.linear_qkv",
+        "*language_model.*.linear_proj",
+        "*language_model.*.in_proj",
+        "*language_model.*.out_proj",
+        "*language_model.*.linear_fc1",
+        "*language_model.*.linear_fc2",
     ]
     assert cfg.peft.dim == 32
     assert cfg.peft.alpha == 32
@@ -648,6 +649,46 @@ def test_super_vl_peft_recipe_uses_native_lora_targets_and_frozen_vision(fake_pr
     assert cfg.ddp.grad_reduce_in_fp32 is True
 
 
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "recipe_func",
+    [
+        _super_vl_h100_recipe_module.nemotron_35_super_vl_peft_16gpu_h100_bf16_config,
+        _super_vl_gb200_recipe_module.nemotron_35_super_vl_peft_16gpu_gb200_bf16_config,
+    ],
+    ids=["h100", "gb200"],
+)
+@pytest.mark.parametrize("wrapped", [False, True], ids=["bare", "wrapped"])
+def test_super_vl_peft_adapters_only_modify_language_modules(recipe_func, wrapped, fake_processor):
+    cfg = _build_config(recipe_func, fake_processor)
+    target_names = ("linear_qkv", "linear_proj", "in_proj", "out_proj", "linear_fc1", "linear_fc2")
+
+    def linear_stack():
+        return torch.nn.ModuleDict({name: torch.nn.Linear(4, 4) for name in target_names})
+
+    # Identical leaf names expose accidental unqualified matches in both media
+    # components, while the language MTP path must still receive adapters.
+    model = torch.nn.Module()
+    model.language_model = torch.nn.ModuleDict({"decoder": linear_stack(), "mtp": linear_stack()})
+    model.vision_model = torch.nn.ModuleDict({"decoder": linear_stack()})
+    model.vision_projection = torch.nn.ModuleDict({"encoder": linear_stack()})
+    if wrapped:
+        wrapper = torch.nn.Module()
+        wrapper.module = model
+        model = wrapper
+    base_parameters = list(model.parameters())
+    model = cfg.peft(model)
+
+    prefix = "module." if wrapped else ""
+    expected = {f"{prefix}language_model.{stack}.{name}" for stack in ("decoder", "mtp") for name in target_names}
+    adapted = {name for name, module in model.named_modules() if isinstance(module, LoRALinear)}
+    assert adapted == expected
+    assert all(not parameter.requires_grad for parameter in base_parameters)
+    trainable = {name for name, parameter in model.named_parameters() if parameter.requires_grad}
+    assert trainable
+    assert all(name.startswith(f"{prefix}language_model.") and ".adapter." in name for name in trainable)
+
+
 def test_super_vl_pretrain_recipe_uses_tuned_h100_training_policy(fake_processor):
     cfg = _build_config(
         _super_vl_h100_recipe_module.nemotron_35_super_vl_pretrain_64gpu_h100_bf16_config,
@@ -682,7 +723,6 @@ def test_super_vl_pretrain_recipe_uses_tuned_h100_training_policy(fake_processor
     assert cfg.model.moe_expert_capacity_factor is None
     assert cfg.model.moe_pad_expert_input_to_capacity is False
     assert cfg.model.moe_hybridep_pad_uneven_dispatch_inputs is False
-    assert cfg.model.moe_hybridep_assume_equal_dispatch_inputs is True
     assert cfg.model.moe_flex_dispatcher_num_sms == 32
     assert cfg.model.moe_hybridep_num_sms is None
     assert cfg.model.moe_hybridep_num_sms_preprocessing == 108
@@ -774,7 +814,6 @@ def test_super_vl_pretrain_recipe_uses_gb200_nvl72_policy(fake_processor):
     assert cfg.model.moe_expert_capacity_factor is None
     assert cfg.model.moe_pad_expert_input_to_capacity is False
     assert cfg.model.moe_hybridep_pad_uneven_dispatch_inputs is False
-    assert cfg.model.moe_hybridep_assume_equal_dispatch_inputs is True
     assert cfg.model.moe_flex_dispatcher_num_sms == 32
     assert cfg.model.moe_hybridep_num_sms == 32
     assert cfg.model.moe_permute_fusion_into_hybridep is False
@@ -926,12 +965,12 @@ def test_super_vl_peft_recipe_uses_gb200_support_topology(fake_processor):
     assert isinstance(cfg.dataset, EnergonDatasetConfig)
     assert cfg.dataset.task_encoder.hf_processor_path == _TEST_SUPER_VL_HF_ID
     assert cfg.peft.target_modules == [
-        "linear_qkv",
-        "linear_proj",
-        "in_proj",
-        "out_proj",
-        "linear_fc1",
-        "linear_fc2",
+        "*language_model.*.linear_qkv",
+        "*language_model.*.linear_proj",
+        "*language_model.*.in_proj",
+        "*language_model.*.out_proj",
+        "*language_model.*.linear_fc1",
+        "*language_model.*.linear_fc2",
     ]
     assert cfg.peft.dim == 32
     assert cfg.peft.alpha == 32

tests/unit_tests/scripts/inference/test_setup_inference.py

@@ -332,35 +332,6 @@ def __init__(self, **kwargs):
     assert executor.kwargs["gpus_per_node"] == 4
 
 
-def test_slurm_executor_can_assign_multiple_gpus_to_one_task(tmp_path, monkeypatch):
-    module = _load_setup_inference_module()
-
-    class _SlurmExecutor:
-        def __init__(self, **kwargs):
-            self.kwargs = kwargs
-
-    module.run.Packager = lambda: "packager"
-    module.run.LocalTunnel = lambda **kwargs: types.SimpleNamespace(**kwargs)
-    module.run.SlurmExecutor = _SlurmExecutor
-    monkeypatch.setattr(module, "get_nemorun_home", lambda: str(tmp_path))
-    args, inference_args = module.parse_args(
-        _launcher_args(
-            "--gpus-per-node",
-            "4",
-            "--tasks-per-node",
-            "1",
-            "--device-map",
-            "balanced_low_0",
-        )
-    )
-
-    executor = module._build_executor(args, [], [])
-
-    assert executor.kwargs["ntasks_per_node"] == 1
-    assert executor.kwargs["gpus_per_node"] == 4
-    assert inference_args == ["--device-map", "balanced_low_0"]
-
-
 def test_slurm_executor_can_request_exclusive_nodes(tmp_path, monkeypatch):
     module = _load_setup_inference_module()