From d2943e00291a18f75ef8a77f0aecd3ce523cf4ee Mon Sep 17 00:00:00 2001
From: Ye Yu <yeyu@nvidia.com>
Date: Mon, 29 Jun 2026 11:51:29 -0700
Subject: [PATCH 01/10] feat: config_overrides + cosmos3_omni loading in
 load_vlm_or_llm

EAGLE3 base-model loading for NVIDIA Cosmos3 (Qwen3-VL text tower):
- ModelArguments.config_overrides: optional dict applied to model config +
  text_config before instantiation (some checkpoints don't propagate text_config
  dims; Cosmos3 needs intermediate_size/num_key_value_heads forced).
- load_vlm_or_llm: apply config_overrides; for model_type==cosmos3_omni use the
  transformers-cosmos3 Cosmos3ForConditionalGeneration (a Qwen3-VL subclass; the
  plugin registers only the config under Auto*), with ignore_mismatched_sizes for
  the unused vision tower.
- main.py: thread recipe.model.config_overrides into both load paths.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Ye Yu <yeyu@nvidia.com>
---
 examples/speculative_decoding/main.py         |  6 ++-
 .../speculative/plugins/hf_training_args.py   |  5 +++
 modelopt/torch/speculative/utils.py           | 37 +++++++++++++++++++
 3 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/examples/speculative_decoding/main.py b/examples/speculative_decoding/main.py
index 46848203606..8e5acd91544 100644
--- a/examples/speculative_decoding/main.py
+++ b/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,
diff --git a/modelopt/torch/speculative/plugins/hf_training_args.py b/modelopt/torch/speculative/plugins/hf_training_args.py
index 38d3f483e6b..b121d9f260d 100644
--- a/modelopt/torch/speculative/plugins/hf_training_args.py
+++ b/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):
diff --git a/modelopt/torch/speculative/utils.py b/modelopt/torch/speculative/utils.py
index 9fa8fde5e15..1f6250a3480 100644
--- a/modelopt/torch/speculative/utils.py
+++ b/modelopt/torch/speculative/utils.py
@@ -591,6 +591,7 @@ def load_vlm_or_llm(
     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,6 +606,9 @@ 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).
     """
     if use_offline_training and use_fake_base:
         from modelopt.torch.speculative.plugins.modeling_fakebase import FakeBaseModel
@@ -616,6 +620,16 @@ def load_vlm_or_llm(
         trust_remote_code=trust_remote_code,
     )
 
+    # Apply caller-supplied config corrections to both the parent config and its
+    # nested text_config (some checkpoints don't propagate text_config dims).
+    if config_overrides:
+        for cfg_obj in (model_config, getattr(model_config, "text_config", None)):
+            if cfg_obj is None:
+                continue
+            for key, value in config_overrides.items():
+                if hasattr(cfg_obj, key):
+                    setattr(cfg_obj, key, value)
+
     # 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(
@@ -646,8 +660,31 @@ def load_vlm_or_llm(
         if hasattr(model_config, "layer_types"):
             extra["layer_types"] = []
 
+    # Cosmos3 omni checkpoints: the transformers-cosmos3 plugin registers only the config
+    # (cosmos3_omni) with AutoConfig, not a model under the Auto* maps, so use 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
+
+        return Cosmos3ForConditionalGeneration.from_pretrained(
+            model_name_or_path,
+            config=model_config,
+            trust_remote_code=trust_remote_code,
+            torch_dtype=dtype,
+            device_map=device_map,
+            ignore_mismatched_sizes=True,
+            **extra,
+        )
+
+    if _is_vlm:
+        model_cls = transformers.AutoModelForVision2Seq
+    else:
+        model_cls = transformers.AutoModelForCausalLM
+
     model = model_cls.from_pretrained(
         model_name_or_path,
+        config=model_config if config_overrides else None,
         trust_remote_code=trust_remote_code,
         torch_dtype=dtype,
         device_map=device_map,

From ce564937b741ee4d536bc009f19f31f03ddf65ba Mon Sep 17 00:00:00 2001
From: Ye Yu <yeyu@nvidia.com>
Date: Mon, 29 Jun 2026 14:57:55 -0700
Subject: [PATCH 02/10] fix: import transformers_cosmos3 before AutoConfig so
 cosmos3_omni is registered

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Ye Yu <yeyu@nvidia.com>
---
 modelopt/torch/speculative/utils.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/modelopt/torch/speculative/utils.py b/modelopt/torch/speculative/utils.py
index 1f6250a3480..564f7f75ecc 100644
--- a/modelopt/torch/speculative/utils.py
+++ b/modelopt/torch/speculative/utils.py
@@ -615,6 +615,13 @@ def load_vlm_or_llm(
 
         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.
+    try:
+        import transformers_cosmos3  # noqa: F401
+    except ImportError:
+        pass
+
     model_config = transformers.AutoConfig.from_pretrained(
         model_name_or_path,
         trust_remote_code=trust_remote_code,

From 4ff57d5d5391d1c1d08bf10725d80470876120f0 Mon Sep 17 00:00:00 2001
From: Ye Yu <yeyu@nvidia.com>
Date: Tue, 18 Aug 2026 20:57:25 -0700
Subject: [PATCH 03/10] ar_validate: add --config_overrides passthrough to
 load_vlm_or_llm

load_vlm_or_llm already accepts config_overrides, but ar_validate.py had no way
to supply them, so it could not load checkpoints whose nested text_config dims
don't propagate to the parent config.

The Cosmos3-Nano EAGLE3 checkpoints are exactly that case: the saved config.json
carries num_hidden_layers/intermediate_size/num_key_value_heads/hidden_size
correctly under text_config but leaves the parent fields None, which is the same
propagation gap the cotrain works around at load time.

Optional and defaults to None, so existing callers are unaffected.

Signed-off-by: Ye Yu <yeyu@nvidia.com>
---
 .../scripts/ar_validate.py                     | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/examples/speculative_decoding/scripts/ar_validate.py b/examples/speculative_decoding/scripts/ar_validate.py
index 5699c480b7b..e689ea65d80 100644
--- a/examples/speculative_decoding/scripts/ar_validate.py
+++ b/examples/speculative_decoding/scripts/ar_validate.py
@@ -19,6 +19,7 @@
 """
 
 import argparse
+import json
 from collections import defaultdict
 
 from accelerate import Accelerator
@@ -100,11 +101,26 @@ def main():
         default=None,
         help="Error if AR is below this threshold.",
     )
+    parser.add_argument(
+        "--config_overrides",
+        type=str,
+        default=None,
+        help=(
+            "JSON dict 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."
+        ),
+    )
     args = parser.parse_args()
 
+    config_overrides = json.loads(args.config_overrides) if args.config_overrides else None
+
     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

From b655c473e8f816ef4d95fc17209a0a0affb2479f Mon Sep 17 00:00:00 2001
From: Ye Yu <yeyu@nvidia.com>
Date: Tue, 25 Aug 2026 12:20:18 -0700
Subject: [PATCH 04/10] export_hf_checkpoint: add --config_overrides
 passthrough

Same gap ar_validate.py had: load_vlm_or_llm accepts config_overrides but the
export script had no way to supply them, so it cannot load checkpoints whose
nested text_config dims don't propagate to the parent config -- e.g. the
Cosmos3-Nano EAGLE3 checkpoints, whose saved config.json leaves the parent
num_hidden_layers/intermediate_size/num_key_value_heads/hidden_size as None.

Optional and defaults to None, so existing callers are unaffected.

Signed-off-by: Ye Yu <yeyu@nvidia.com>
---
 .../scripts/export_hf_checkpoint.py            | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/examples/speculative_decoding/scripts/export_hf_checkpoint.py b/examples/speculative_decoding/scripts/export_hf_checkpoint.py
index cee2e45d0eb..ab199d1d3e7 100644
--- a/examples/speculative_decoding/scripts/export_hf_checkpoint.py
+++ b/examples/speculative_decoding/scripts/export_hf_checkpoint.py
@@ -16,6 +16,7 @@
 """Export a HF checkpoint (with ModelOpt state) for deployment."""
 
 import argparse
+import json
 
 import torch
 
@@ -33,13 +34,28 @@ 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=(
+            "JSON dict 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."
+        ),
+    )
     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=json.loads(args.config_overrides) if args.config_overrides else None,
+)
 model.eval()
 with torch.inference_mode():
     export_speculative_decoding(

From 3308ffc843bb186fc7c3c3dead561ae05f5f63e1 Mon Sep 17 00:00:00 2001
From: Ye Yu <yeyu@nvidia.com>
Date: Wed, 26 Aug 2026 11:01:17 -0700
Subject: [PATCH 05/10] merge_lora: load the base via load_vlm_or_llm instead
 of AutoModelForCausalLM

AutoModelForCausalLM cannot load architectures that are absent from the Auto*
maps. Cosmos3 is the motivating case: the transformers-cosmos3 plugin registers
only the `cosmos3_omni` *config* with AutoConfig, never a model class, so merging
a LoRA into a Cosmos3-Nano base died with KeyError('cosmos3_omni') regardless of
what was imported.

load_vlm_or_llm already encapsulates this -- it imports the plugin, dispatches to
Cosmos3ForConditionalGeneration when model_type is cosmos3_omni, and applies
config_overrides for checkpoints whose nested text_config dims don't propagate.
For plain LLMs it falls back to AutoModelForCausalLM with the same dtype and
device_map, so behavior for existing models is unchanged; VLMs additionally route
to AutoModelForVision2Seq instead of failing.

Also adds --config_overrides, matching ar_validate.py and export_hf_checkpoint.py.

Signed-off-by: Ye Yu <yeyu@nvidia.com>
---
 .../scripts/merge_lora.py                     | 29 ++++++++++++++++---
 1 file changed, 25 insertions(+), 4 deletions(-)

diff --git a/examples/speculative_decoding/scripts/merge_lora.py b/examples/speculative_decoding/scripts/merge_lora.py
index 25d311393b5..52416062bcc 100644
--- a/examples/speculative_decoding/scripts/merge_lora.py
+++ b/examples/speculative_decoding/scripts/merge_lora.py
@@ -28,10 +28,13 @@
 """
 
 import argparse
+import json
 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 load_vlm_or_llm
 
 
 def parse_args():
@@ -61,6 +64,16 @@ 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=(
+            "JSON dict of config fields to override on the base 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."
+        ),
+    )
     return parser.parse_args()
 
 
@@ -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=json.loads(args.config_overrides) if args.config_overrides else None,
     )
     tokenizer = AutoTokenizer.from_pretrained(
         args.base_model_path, trust_remote_code=args.trust_remote_code

From 76860b174dc3c9855bb17f9b048986b095b8465a Mon Sep 17 00:00:00 2001
From: Ye Yu <yeyu@nvidia.com>
Date: Mon, 31 Aug 2026 12:02:50 -0700
Subject: [PATCH 06/10] use contextlib.suppress for the optional
 transformers_cosmos3 import

Satisfies ruff SIM105, which main's lint config now enforces.
---
 modelopt/torch/speculative/utils.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/modelopt/torch/speculative/utils.py b/modelopt/torch/speculative/utils.py
index 564f7f75ecc..ed0fe3e5a14 100644
--- a/modelopt/torch/speculative/utils.py
+++ b/modelopt/torch/speculative/utils.py
@@ -617,10 +617,8 @@ def load_vlm_or_llm(
 
     # Import the transformers-cosmos3 plugin if available: it registers the `cosmos3_omni`
     # architecture with AutoConfig on import, so the from_pretrained below recognizes it.
-    try:
+    with contextlib.suppress(ImportError):
         import transformers_cosmos3  # noqa: F401
-    except ImportError:
-        pass
 
     model_config = transformers.AutoConfig.from_pretrained(
         model_name_or_path,

From 8eed412f7aee126ddfc763a32f968bb2c22d405b Mon Sep 17 00:00:00 2001
From: Ye Yu <yeyu@nvidia.com>
Date: Mon, 31 Aug 2026 12:13:38 -0700
Subject: [PATCH 07/10] address review: fix VLM class regression, config/kwargs
 interaction, silent drops
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Review findings on this PR, all verified against the code:

1. CRITICAL — duplicate model_cls selection. A second `model_cls` block ran after
   the existing one and overwrote it with a bare
   `transformers.AutoModelForVision2Seq`, discarding main's Transformers-5
   fallback. Since pyproject allows transformers<5.15, every VLM load through
   load_vlm_or_llm would have raised AttributeError on a supported install. This
   was a bad conflict resolution on my side when rebasing; removed the duplicate.

2. CRITICAL — passing `config=` changed how `**extra` is interpreted.
   from_pretrained only forwards unrecognized kwargs into the config when it
   builds that config itself; given a PretrainedConfig instance it leaves them in
   model_kwargs, so `num_hidden_layers=0` silently stopped taking effect and the
   full base was materialized -- defeating the OOM-avoidance the offline path
   exists for. Offline fields are now set on the config when a config object is
   passed, and the original depth is captured before zeroing so
   num_orig_hidden_layers is still restored correctly.

3. The cosmos3 branch no longer returns early; it selects model_cls and falls
   through to the shared from_pretrained, so offline post-processing
   (num_orig_hidden_layers) applies to it too and the load call is not duplicated.

4. IMPORTANT — unmatched override keys were dropped silently, so a typo produced
   a wrong-shaped model while appearing to have applied. Now raises listing the
   offending keys.

5. FakeBaseModel paths rebuild the config from the checkpoint and cannot honor
   config_overrides; they now refuse instead of ignoring them. Scoped to those two
   paths -- offline training without fake-base handles overrides correctly.

6. merge_lora no longer copies the base config.json back over the saved one when
   config_overrides were applied: the weights were built from corrected dims, so
   the original config would disagree with them and force every downstream reader
   to re-supply the same overrides.

Signed-off-by: Ye Yu <yeyu@nvidia.com>
---
 .../scripts/merge_lora.py                     | 11 ++-
 modelopt/torch/speculative/utils.py           | 78 +++++++++++++------
 2 files changed, 64 insertions(+), 25 deletions(-)

diff --git a/examples/speculative_decoding/scripts/merge_lora.py b/examples/speculative_decoding/scripts/merge_lora.py
index 52416062bcc..80743eae0cd 100644
--- a/examples/speculative_decoding/scripts/merge_lora.py
+++ b/examples/speculative_decoding/scripts/merge_lora.py
@@ -156,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 args.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}")
 
diff --git a/modelopt/torch/speculative/utils.py b/modelopt/torch/speculative/utils.py
index ed0fe3e5a14..4d2236cfd10 100644
--- a/modelopt/torch/speculative/utils.py
+++ b/modelopt/torch/speculative/utils.py
@@ -610,7 +610,19 @@ def load_vlm_or_llm(
             its ``text_config`` before instantiation (e.g. to correct dims that don't
             propagate from a checkpoint's nested text config).
     """
+
+    def _reject_overrides_on_fake_base():
+        # FakeBaseModel.from_source re-reads the checkpoint config itself, so overrides applied
+        # here never reach it. Refuse rather than hand back a model built from the uncorrected
+        # dims the caller explicitly asked to fix.
+        if config_overrides:
+            raise NotImplementedError(
+                "config_overrides is not supported on the FakeBaseModel path: from_source "
+                "rebuilds the config from the checkpoint and would silently ignore them."
+            )
+
     if use_offline_training and use_fake_base:
+        _reject_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)
@@ -628,12 +640,24 @@ def load_vlm_or_llm(
     # Apply caller-supplied config corrections to both the parent config and its
     # nested text_config (some checkpoints don't propagate text_config dims).
     if config_overrides:
-        for cfg_obj in (model_config, getattr(model_config, "text_config", None)):
-            if cfg_obj is None:
-                continue
-            for key, value in config_overrides.items():
+        targets = [cfg for cfg in (model_config, getattr(model_config, "text_config", None)) if cfg]
+        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).
@@ -644,6 +668,7 @@ def load_vlm_or_llm(
     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").
+        _reject_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)
@@ -660,36 +685,41 @@ def load_vlm_or_llm(
         model_cls = transformers.AutoModelForCausalLM
 
     extra = {}
-    if use_offline_training:
-        extra["num_hidden_layers"] = 0
-        if hasattr(model_config, "layer_types"):
-            extra["layer_types"] = []
 
     # Cosmos3 omni checkpoints: the transformers-cosmos3 plugin registers only the config
-    # (cosmos3_omni) with AutoConfig, not a model under the Auto* maps, so use its
+    # (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
 
-        return Cosmos3ForConditionalGeneration.from_pretrained(
-            model_name_or_path,
-            config=model_config,
-            trust_remote_code=trust_remote_code,
-            torch_dtype=dtype,
-            device_map=device_map,
-            ignore_mismatched_sizes=True,
-            **extra,
-        )
+        model_cls = Cosmos3ForConditionalGeneration
+        extra["ignore_mismatched_sizes"] = True
 
-    if _is_vlm:
-        model_cls = transformers.AutoModelForVision2Seq
-    else:
-        model_cls = transformers.AutoModelForCausalLM
+    # 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:
+        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 config_overrides else None,
+        config=model_config if pass_config else None,
         trust_remote_code=trust_remote_code,
         torch_dtype=dtype,
         device_map=device_map,
@@ -698,7 +728,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
 

From efb2ca2ddf9e1d5627c7c269589e6164b10ed005 Mon Sep 17 00:00:00 2001
From: Ye Yu <yeyu@nvidia.com>
Date: Mon, 31 Aug 2026 12:24:39 -0700
Subject: [PATCH 08/10] address review: share one config_overrides parser
 across the three scripts

The json.loads expression and the identical help text were copy-pasted into
ar_validate.py, export_hf_checkpoint.py and merge_lora.py, and neither validated
the payload: '[1,2]' or '36' parsed fine and only failed later inside
load_vlm_or_llm with AttributeError, while malformed JSON surfaced as a raw
JSONDecodeError traceback.

Adds parse_config_overrides() and CONFIG_OVERRIDES_HELP next to load_vlm_or_llm.
Non-object payloads and bad JSON now fail immediately with an actionable message.

Signed-off-by: Ye Yu <yeyu@nvidia.com>
---
 .../scripts/ar_validate.py                    | 15 +++++-----
 .../scripts/export_hf_checkpoint.py           | 15 +++++-----
 .../scripts/merge_lora.py                     | 15 +++++-----
 modelopt/torch/speculative/utils.py           | 29 +++++++++++++++++++
 4 files changed, 50 insertions(+), 24 deletions(-)

diff --git a/examples/speculative_decoding/scripts/ar_validate.py b/examples/speculative_decoding/scripts/ar_validate.py
index e689ea65d80..ccfb47dc085 100644
--- a/examples/speculative_decoding/scripts/ar_validate.py
+++ b/examples/speculative_decoding/scripts/ar_validate.py
@@ -19,7 +19,6 @@
 """
 
 import argparse
-import json
 from collections import defaultdict
 
 from accelerate import Accelerator
@@ -29,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()
 
@@ -105,15 +108,11 @@ def main():
         "--config_overrides",
         type=str,
         default=None,
-        help=(
-            "JSON dict 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."
-        ),
+        help=CONFIG_OVERRIDES_HELP,
     )
     args = parser.parse_args()
 
-    config_overrides = json.loads(args.config_overrides) if args.config_overrides else None
+    config_overrides = parse_config_overrides(args.config_overrides)
 
     accelerator = Accelerator()
     model = load_vlm_or_llm(
diff --git a/examples/speculative_decoding/scripts/export_hf_checkpoint.py b/examples/speculative_decoding/scripts/export_hf_checkpoint.py
index ab199d1d3e7..2b86a47004b 100644
--- a/examples/speculative_decoding/scripts/export_hf_checkpoint.py
+++ b/examples/speculative_decoding/scripts/export_hf_checkpoint.py
@@ -16,13 +16,16 @@
 """Export a HF checkpoint (with ModelOpt state) for deployment."""
 
 import argparse
-import json
 
 import torch
 
 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():
@@ -38,11 +41,7 @@ def parse_args():
         "--config_overrides",
         type=str,
         default=None,
-        help=(
-            "JSON dict 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."
-        ),
+        help=CONFIG_OVERRIDES_HELP,
     )
     return parser.parse_args()
 
@@ -54,7 +53,7 @@ def parse_args():
     args.model_path,
     dtype="auto",
     trust_remote_code=args.trust_remote_code,
-    config_overrides=json.loads(args.config_overrides) if args.config_overrides else None,
+    config_overrides=parse_config_overrides(args.config_overrides),
 )
 model.eval()
 with torch.inference_mode():
diff --git a/examples/speculative_decoding/scripts/merge_lora.py b/examples/speculative_decoding/scripts/merge_lora.py
index 80743eae0cd..bee685ac729 100644
--- a/examples/speculative_decoding/scripts/merge_lora.py
+++ b/examples/speculative_decoding/scripts/merge_lora.py
@@ -28,13 +28,16 @@
 """
 
 import argparse
-import json
 from pathlib import Path
 
 from safetensors.torch import load_file
 from transformers import AutoTokenizer
 
-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():
@@ -68,11 +71,7 @@ def parse_args():
         "--config_overrides",
         type=str,
         default=None,
-        help=(
-            "JSON dict of config fields to override on the base 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."
-        ),
+        help=CONFIG_OVERRIDES_HELP,
     )
     return parser.parse_args()
 
@@ -108,7 +107,7 @@ def main():
         dtype="auto",
         device_map="cpu",
         trust_remote_code=args.trust_remote_code,
-        config_overrides=json.loads(args.config_overrides) if args.config_overrides else None,
+        config_overrides=parse_config_overrides(args.config_overrides),
     )
     tokenizer = AutoTokenizer.from_pretrained(
         args.base_model_path, trust_remote_code=args.trust_remote_code
diff --git a/modelopt/torch/speculative/utils.py b/modelopt/torch/speculative/utils.py
index 4d2236cfd10..4d03a34d471 100644
--- a/modelopt/torch/speculative/utils.py
+++ b/modelopt/torch/speculative/utils.py
@@ -18,6 +18,7 @@
 import contextlib
 import copy
 import importlib.util
+import json
 import os
 import sys
 import warnings
@@ -584,6 +585,34 @@ 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:
+        parsed = json.loads(raw)
+    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,

From 4322e5ec99024acd84600792d573a8558809c6c5 Mon Sep 17 00:00:00 2001
From: Ye Yu <yeyu@nvidia.com>
Date: Mon, 31 Aug 2026 12:36:41 -0700
Subject: [PATCH 09/10] address second review round: llm_config targeting,
 fake-base warn, parse once
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

1. IMPORTANT — overrides only reached text_config, but VLM detection also accepts
   llm_config. For a checkpoint nesting its text tower under llm_config with the
   fields mirrored as None on the parent, the override landed on the parent,
   counted as applied, and never reached the real text tower -- silently
   reproducing the wrong-shaped model this flag exists to prevent. Overrides now
   target every nested config, via a NESTED_CONFIG_ATTRS constant that VLM
   detection also uses so the two lists cannot drift.

2. Downgraded the FakeBaseModel guard from NotImplementedError to a warning.
   Verified the reviewer's point: FakeBaseModel.from_source resolves dims from the
   nested text_config/llm_config first (modeling_fakebase.py:185-192), so that
   path is already correct without overrides -- the case I rejected would have
   worked. Since main.py forwards config_overrides unconditionally, raising made
   a single recipe unable to run offline at all.

3. merge_lora parses the overrides once and branches on the parsed dict, so
   --config_overrides '{}' (a no-op) keeps the config.json restore instead of
   taking the 'overrides were applied' path.

Signed-off-by: Ye Yu <yeyu@nvidia.com>
---
 .../scripts/merge_lora.py                     |  5 ++-
 modelopt/torch/speculative/utils.py           | 44 ++++++++++++++-----
 2 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/examples/speculative_decoding/scripts/merge_lora.py b/examples/speculative_decoding/scripts/merge_lora.py
index bee685ac729..4580e166832 100644
--- a/examples/speculative_decoding/scripts/merge_lora.py
+++ b/examples/speculative_decoding/scripts/merge_lora.py
@@ -78,6 +78,7 @@ def 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)
@@ -107,7 +108,7 @@ def main():
         dtype="auto",
         device_map="cpu",
         trust_remote_code=args.trust_remote_code,
-        config_overrides=parse_config_overrides(args.config_overrides),
+        config_overrides=config_overrides,
     )
     tokenizer = AutoTokenizer.from_pretrained(
         args.base_model_path, trust_remote_code=args.trust_remote_code
@@ -161,7 +162,7 @@ def main():
     # re-supply the same overrides.
     base_config = Path(args.base_model_path) / "config.json"
     output_config = Path(args.output_path) / "config.json"
-    if args.config_overrides:
+    if config_overrides:
         print(
             "  Keeping the saved config.json (config_overrides were applied, so the original "
             "base config would not match the merged weights)"
diff --git a/modelopt/torch/speculative/utils.py b/modelopt/torch/speculative/utils.py
index 4d03a34d471..d757b160970 100644
--- a/modelopt/torch/speculative/utils.py
+++ b/modelopt/torch/speculative/utils.py
@@ -37,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 %}"
 )
@@ -640,18 +645,23 @@ def load_vlm_or_llm(
             propagate from a checkpoint's nested text config).
     """
 
-    def _reject_overrides_on_fake_base():
+    def _warn_overrides_on_fake_base():
         # FakeBaseModel.from_source re-reads the checkpoint config itself, so overrides applied
-        # here never reach it. Refuse rather than hand back a model built from the uncorrected
-        # dims the caller explicitly asked to fix.
+        # 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:
-            raise NotImplementedError(
-                "config_overrides is not supported on the FakeBaseModel path: from_source "
-                "rebuilds the config from the checkpoint and would silently ignore them."
+            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:
-        _reject_overrides_on_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)
@@ -666,10 +676,20 @@ def _reject_overrides_on_fake_base():
         trust_remote_code=trust_remote_code,
     )
 
-    # Apply caller-supplied config corrections to both the parent config and its
-    # nested text_config (some checkpoints don't propagate text_config dims).
+    # 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:
-        targets = [cfg for cfg in (model_config, getattr(model_config, "text_config", None)) if cfg]
+        # 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
@@ -691,13 +711,13 @@ def _reject_overrides_on_fake_base():
     # 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").
-        _reject_overrides_on_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)

From 91402faa8deccaafe96dc2b3d8bcb714ebc332d3 Mon Sep 17 00:00:00 2001
From: Ye Yu <yeyu@nvidia.com>
Date: Mon, 31 Aug 2026 12:37:55 -0700
Subject: [PATCH 10/10] address review: reject NaN/Infinity in
 --config_overrides

Python's json.loads accepts the JSON5-ish constants NaN, Infinity and -Infinity
by default, so they would parse cleanly and land on a config field as a model
dimension. parse_constant now rejects them with the same actionable error as the
other malformed inputs.

Signed-off-by: Ye Yu <yeyu@nvidia.com>
---
 modelopt/torch/speculative/utils.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/speculative/utils.py b/modelopt/torch/speculative/utils.py
index d757b160970..c419d24588f 100644
--- a/modelopt/torch/speculative/utils.py
+++ b/modelopt/torch/speculative/utils.py
@@ -607,7 +607,12 @@ def parse_config_overrides(raw: str | None) -> dict | None:
     if not raw:
         return None
     try:
-        parsed = json.loads(raw)
+        # 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):