specdec: config_overrides for nested text_config checkpoints + load VLM-capable bases in merge_lora

examples/speculative_decoding/main.py

@@ -217,7 +217,10 @@ def train():
         assert checkpoint is not None  # guaranteed by checkpoint_is_hf
         with patch_transformers5_params_loading():
             model = load_vlm_or_llm(
-                checkpoint, dtype="auto", trust_remote_code=recipe.model.trust_remote_code
+                checkpoint,
+                dtype="auto",
+                trust_remote_code=recipe.model.trust_remote_code,
+                config_overrides=recipe.model.config_overrides,
             )
         tokenizer = transformers.AutoTokenizer.from_pretrained(
             checkpoint, trust_remote_code=recipe.model.trust_remote_code
@@ -242,6 +245,7 @@ def train():
             dtype="auto",
             device_map="cpu",
             trust_remote_code=recipe.model.trust_remote_code,
+            config_overrides=recipe.model.config_overrides,
         )
         tokenizer = transformers.AutoTokenizer.from_pretrained(
             model_name_or_path,

examples/speculative_decoding/scripts/ar_validate.py

@@ -28,7 +28,11 @@
 
 import modelopt.torch.opt as mto
 from modelopt.torch.speculative.plugins.hf_eagle import HFARValidation
-from modelopt.torch.speculative.utils import load_vlm_or_llm
+from modelopt.torch.speculative.utils import (
+    CONFIG_OVERRIDES_HELP,
+    load_vlm_or_llm,
+    parse_config_overrides,
+)
 
 mto.enable_huggingface_checkpointing()
 
@@ -100,11 +104,22 @@ def main():
         default=None,
         help="Error if AR is below this threshold.",
     )
+    parser.add_argument(
+        "--config_overrides",
+        type=str,
+        default=None,
+        help=CONFIG_OVERRIDES_HELP,
+    )
     args = parser.parse_args()
 
+    config_overrides = parse_config_overrides(args.config_overrides)
+
     accelerator = Accelerator()
     model = load_vlm_or_llm(
-        args.model_path, device_map="auto", trust_remote_code=args.trust_remote_code
+        args.model_path,
+        device_map="auto",
+        trust_remote_code=args.trust_remote_code,
+        config_overrides=config_overrides,
     )
     tokenizer = AutoTokenizer.from_pretrained(
         args.model_path, trust_remote_code=args.trust_remote_code

examples/speculative_decoding/scripts/export_hf_checkpoint.py

@@ -21,7 +21,11 @@
 
 import modelopt.torch.opt as mto
 from modelopt.torch.export import export_speculative_decoding
-from modelopt.torch.speculative.utils import load_vlm_or_llm
+from modelopt.torch.speculative.utils import (
+    CONFIG_OVERRIDES_HELP,
+    load_vlm_or_llm,
+    parse_config_overrides,
+)
 
 
 def parse_args():
@@ -33,13 +37,24 @@ def parse_args():
     parser.add_argument(
         "--export_path", type=str, default="Destination directory for exported files."
     )
+    parser.add_argument(
+        "--config_overrides",
+        type=str,
+        default=None,
+        help=CONFIG_OVERRIDES_HELP,
+    )
     return parser.parse_args()
 
 
 mto.enable_huggingface_checkpointing()
 
 args = parse_args()
-model = load_vlm_or_llm(args.model_path, dtype="auto", trust_remote_code=args.trust_remote_code)
+model = load_vlm_or_llm(
+    args.model_path,
+    dtype="auto",
+    trust_remote_code=args.trust_remote_code,
+    config_overrides=parse_config_overrides(args.config_overrides),
+)
 model.eval()
 with torch.inference_mode():
     export_speculative_decoding(

examples/speculative_decoding/scripts/merge_lora.py

@@ -31,7 +31,13 @@
 from pathlib import Path
 
 from safetensors.torch import load_file
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer
+
+from modelopt.torch.speculative.utils import (
+    CONFIG_OVERRIDES_HELP,
+    load_vlm_or_llm,
+    parse_config_overrides,
+)
 
 
 def parse_args():
@@ -61,11 +67,18 @@ def parse_args():
         action="store_true",
         help="Allow loading models that define custom code on the HF Hub. Off by default.",
     )
+    parser.add_argument(
+        "--config_overrides",
+        type=str,
+        default=None,
+        help=CONFIG_OVERRIDES_HELP,
+    )
     return parser.parse_args()
 
 
 def main():
     args = parse_args()
+    config_overrides = parse_config_overrides(args.config_overrides)
     lora_dir = Path(args.exported_lora_dir)
 
     # Verify exported files exist (standard peft naming)
@@ -81,13 +94,21 @@ def main():
     print(f"Loaded {len(lora_sd)} LoRA tensors from {lora_dir}")
     print(f"  Sample keys: {list(lora_sd.keys())[:4]}")
 
-    # Load the original base model
+    # Load the original base model.
+    #
+    # Use load_vlm_or_llm rather than AutoModelForCausalLM directly: it falls back to
+    # AutoModelForCausalLM for plain LLMs (same dtype/device_map, so unchanged behavior), but also
+    # handles VLMs and registers/loads architectures the Auto* maps don't cover. Cosmos3 is the
+    # motivating case -- the transformers-cosmos3 plugin registers only the `cosmos3_omni` config,
+    # never a model under Auto*, so AutoModelForCausalLM raises KeyError('cosmos3_omni') no matter
+    # what is imported.
     print(f"Loading base model from {args.base_model_path}...")
-    model = AutoModelForCausalLM.from_pretrained(
+    model = load_vlm_or_llm(
         args.base_model_path,
-        torch_dtype="auto",
+        dtype="auto",
         device_map="cpu",
         trust_remote_code=args.trust_remote_code,
+        config_overrides=config_overrides,
     )
     tokenizer = AutoTokenizer.from_pretrained(
         args.base_model_path, trust_remote_code=args.trust_remote_code
@@ -135,9 +156,18 @@ def main():
     # Since LoRA only changes weights — not architecture — the original config is correct.
     import shutil
 
+    # ...but only when the loaded config still matches the base's. With --config_overrides the
+    # weights were built from corrected dims, so copying the uncorrected base config back would
+    # leave config.json disagreeing with model.safetensors and force every downstream reader to
+    # re-supply the same overrides.
     base_config = Path(args.base_model_path) / "config.json"
     output_config = Path(args.output_path) / "config.json"
-    if base_config.exists():
+    if config_overrides:
+        print(
+            "  Keeping the saved config.json (config_overrides were applied, so the original "
+            "base config would not match the merged weights)"
+        )
+    elif base_config.exists():
         shutil.copy2(str(base_config), str(output_config))
         print(f"  Restored original config.json from {base_config}")
 

modelopt/torch/speculative/plugins/hf_training_args.py

@@ -44,6 +44,11 @@ class ModelArguments(BaseModel):
     model_name_or_path: str | None = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
     use_fake_base_for_offline: bool = False
     trust_remote_code: bool = False
+    # Optional config field overrides applied to the loaded model config (and its
+    # text_config) before instantiation. Needed for checkpoints whose config doesn't
+    # round-trip cleanly through transformers (e.g. Cosmos3's Qwen3-VL text tower,
+    # where intermediate_size/num_key_value_heads don't propagate from text_config).
+    config_overrides: dict | None = None
 
 
 class DataArguments(BaseModel):

modelopt/torch/speculative/utils.py

@@ -18,6 +18,7 @@
 import contextlib
 import copy
 import importlib.util
+import json
 import os
 import sys
 import warnings
@@ -36,6 +37,11 @@
 KIMI_K2_PACKAGE_NAME = "kimi_k2_temp"
 
 
+# Attributes under which a checkpoint may nest its text-tower config. Mirrors
+# modelopt.torch.speculative.plugins.modeling_fakebase._VLM_CONFIG_ATTRS.
+NESTED_CONFIG_ATTRS = ["text_config", "llm_config"]
+
+
 REMOVE_THINK_CHAT_TEMPLATE = (
     "{% if '</think>' in content %}{% set content = content.split('</think>')[-1] %}{% endif %}"
 )
@@ -584,13 +590,47 @@ def enable_cp_ttt_patch(cp_size: int = 1):
             modelopt.torch.speculative.plugins.hf_eagle.ENABLE_CP_TTT_PATCH = False
 
 
+CONFIG_OVERRIDES_HELP = (
+    "JSON object of config fields to override on the model config and its text_config before "
+    "instantiation, e.g. '{\"num_hidden_layers\": 36}'. Needed for checkpoints whose nested "
+    "text_config dims don't propagate to the parent config."
+)
+
+
+def parse_config_overrides(raw: str | None) -> dict | None:
+    """Parse a ``--config_overrides`` CLI value into a dict, or ``None`` if not supplied.
+
+    Rejects malformed JSON and non-object payloads here, with an actionable message, rather than
+    letting them surface later as a raw ``JSONDecodeError`` or an ``AttributeError`` from deep
+    inside model loading.
+    """
+    if not raw:
+        return None
+    try:
+        # Reject the JSON5-ish constants Python's json accepts by default: NaN/Infinity
+        # would sail through as floats and land on a config field as a dimension.
+        def _reject(const):
+            raise ValueError(f"--config_overrides contains non-finite value {const!r}")
+
+        parsed = json.loads(raw, parse_constant=_reject)
+    except json.JSONDecodeError as e:
+        raise ValueError(f"--config_overrides is not valid JSON: {e}") from e
+    if not isinstance(parsed, dict):
+        raise ValueError(
+            f"--config_overrides must be a JSON object mapping field names to values, got "
+            f"{type(parsed).__name__}: {raw!r}"
+        )
+    return parsed
+
+
 def load_vlm_or_llm(
     model_name_or_path: str,
     use_fake_base: bool = False,
     use_offline_training: bool = False,
     dtype: str | torch.dtype | None = None,
     device_map: str | None = None,
     trust_remote_code: bool = False,
+    config_overrides: dict | None = None,
 ):
     """Load a VLM or LLM. Returns the model.
 
@@ -605,26 +645,84 @@ def load_vlm_or_llm(
         dtype: dtype to use when loading the model.
         device_map: Device map passed to ``from_pretrained``.
         trust_remote_code: Whether to trust remote code.
+        config_overrides: Optional config field overrides applied to the model config and
+            its ``text_config`` before instantiation (e.g. to correct dims that don't
+            propagate from a checkpoint's nested text config).
     """
+
+    def _warn_overrides_on_fake_base():
+        # FakeBaseModel.from_source re-reads the checkpoint config itself, so overrides applied
+        # here never reach it -- but it is not silently wrong: from_source resolves dims from the
+        # nested text_config/llm_config first (modeling_fakebase._VLM_CONFIG_ATTRS), which is the
+        # very problem config_overrides exists to work around, so this path is already correct
+        # without them. Warn rather than raise: main.py forwards config_overrides unconditionally,
+        # so hard-failing would leave a single recipe unable to run offline at all.
+        if config_overrides:
+            warnings.warn(
+                "config_overrides is ignored on the FakeBaseModel path: from_source rebuilds the "
+                "config from the checkpoint, reading dims from the nested text_config/llm_config "
+                "directly, so the overrides are not needed there.",
+                stacklevel=2,
+            )
+
     if use_offline_training and use_fake_base:
+        _warn_overrides_on_fake_base()
         from modelopt.torch.speculative.plugins.modeling_fakebase import FakeBaseModel
 
         return FakeBaseModel.from_source(model_name_or_path, trust_remote_code=trust_remote_code)
 
+    # Import the transformers-cosmos3 plugin if available: it registers the `cosmos3_omni`
+    # architecture with AutoConfig on import, so the from_pretrained below recognizes it.
+    with contextlib.suppress(ImportError):
+        import transformers_cosmos3  # noqa: F401
+
     model_config = transformers.AutoConfig.from_pretrained(
         model_name_or_path,
         trust_remote_code=trust_remote_code,
     )
 
+    # Apply caller-supplied config corrections to the parent config and every nested config
+    # (some checkpoints don't propagate the nested dims up to the parent).
+    if config_overrides:
+        # Cover every nested attribute VLM detection accepts, not just text_config: an
+        # llm_config-nesting checkpoint mirrors the fields on the parent as None, so an override
+        # would land on the parent, count as "applied", and never reach the real text tower.
+        targets = [
+            cfg
+            for cfg in (
+                model_config,
+                *(getattr(model_config, a, None) for a in NESTED_CONFIG_ATTRS),
+            )
+            if cfg is not None
+        ]
+        unmatched = []
+        for key, value in config_overrides.items():
+            applied = False
+            for cfg_obj in targets:
+                if hasattr(cfg_obj, key):
+                    setattr(cfg_obj, key, value)
+                    applied = True
+            if not applied:
+                unmatched.append(key)
+        if unmatched:
+            # Silently skipping a key would hand back a wrong-shaped model while appearing to
+            # have applied the override -- the exact failure this option exists to correct.
+            raise ValueError(
+                f"config_overrides key(s) {sorted(unmatched)} matched no field on the model "
+                f"config (model_type={getattr(model_config, 'model_type', None)!r}) or its "
+                "text_config. Check for typos."
+            )
+
     # Detect VLMs: either "vl" in model_type (e.g. "llava") or has a nested text config
     # (e.g. Mistral3Config with model_type="mistral3" and text_config attribute).
     _is_vlm = "vl" in model_config.model_type.lower() or any(
-        getattr(model_config, attr, None) is not None for attr in ["text_config", "llm_config"]
+        getattr(model_config, attr, None) is not None for attr in NESTED_CONFIG_ATTRS
     )
 
     if _is_vlm and use_offline_training:
         # For VLMs in offline training, FakeBaseModel loads only embed_tokens + lm_head
         # and auto-detects VLM weight key layouts (e.g. "language_model.model.embed_tokens").
+        _warn_overrides_on_fake_base()
         from modelopt.torch.speculative.plugins.modeling_fakebase import FakeBaseModel
 
         return FakeBaseModel.from_source(model_name_or_path, trust_remote_code=trust_remote_code)
@@ -641,13 +739,41 @@ def load_vlm_or_llm(
         model_cls = transformers.AutoModelForCausalLM
 
     extra = {}
+
+    # Cosmos3 omni checkpoints: the transformers-cosmos3 plugin registers only the config
+    # (cosmos3_omni) with AutoConfig, never a model under the Auto* maps, so dispatch to its
+    # Cosmos3ForConditionalGeneration (a Qwen3-VL subclass) directly. The unused vision
+    # tower has mismatched dims vs the text-only use, so ignore those on load.
+    if getattr(model_config, "model_type", None) == "cosmos3_omni":
+        from transformers_cosmos3 import Cosmos3ForConditionalGeneration
+
+        model_cls = Cosmos3ForConditionalGeneration
+        extra["ignore_mismatched_sizes"] = True
+
+    # Pass our config object only when we had to modify it (overrides) or when the model class
+    # needs the plugin-built config; otherwise let from_pretrained build its own, exactly as before.
+    pass_config = bool(config_overrides) or "ignore_mismatched_sizes" in extra
+
+    # Capture the true depth before any zeroing below, since it is restored after load.
+    orig_num_hidden_layers = getattr(model_config, "num_hidden_layers", None)
+
     if use_offline_training:
-        extra["num_hidden_layers"] = 0
-        if hasattr(model_config, "layer_types"):
-            extra["layer_types"] = []
+        if pass_config:
+            # from_pretrained only forwards unrecognized kwargs into the config when it builds
+            # that config itself. Given a PretrainedConfig instance it deep-copies it and leaves
+            # the rest in model_kwargs, so num_hidden_layers=0 would never reach the config and
+            # the full model would be materialized. Set the fields on the config directly.
+            model_config.num_hidden_layers = 0
+            if hasattr(model_config, "layer_types"):
+                model_config.layer_types = []
+        else:
+            extra["num_hidden_layers"] = 0
+            if hasattr(model_config, "layer_types"):
+                extra["layer_types"] = []
 
     model = model_cls.from_pretrained(
         model_name_or_path,
+        config=model_config if pass_config else None,
         trust_remote_code=trust_remote_code,
         torch_dtype=dtype,
         device_map=device_map,
@@ -656,7 +782,7 @@ def load_vlm_or_llm(
 
     if use_offline_training:
         # Preserve the original layer count since we loaded with num_hidden_layers=0
-        model.config.num_orig_hidden_layers = model_config.num_hidden_layers
+        model.config.num_orig_hidden_layers = orig_num_hidden_layers
 
     return model