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/examples/speculative_decoding/scripts/ar_validate.py b/examples/speculative_decoding/scripts/ar_validate.py index 5699c480b7b..ccfb47dc085 100644 --- a/examples/speculative_decoding/scripts/ar_validate.py +++ b/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 diff --git a/examples/speculative_decoding/scripts/export_hf_checkpoint.py b/examples/speculative_decoding/scripts/export_hf_checkpoint.py index cee2e45d0eb..2b86a47004b 100644 --- a/examples/speculative_decoding/scripts/export_hf_checkpoint.py +++ b/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( diff --git a/examples/speculative_decoding/scripts/merge_lora.py b/examples/speculative_decoding/scripts/merge_lora.py index 25d311393b5..4580e166832 100644 --- a/examples/speculative_decoding/scripts/merge_lora.py +++ b/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}") 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..c419d24588f 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 @@ -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 '' in content %}{% set content = content.split('')[-1] %}{% endif %}" ) @@ -584,6 +590,39 @@ 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, @@ -591,6 +630,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,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