correctness fixes across loading, quantization, configs, and processors

zeromodels/base/base_config.py

@@ -74,6 +74,30 @@ def _annotations(cls):
     def field_names(cls):
         return list(cls._annotations().keys())
 
+    @classmethod
+    def unknown_keys(cls, data):
+        """Keys in a repo spec ``data`` this config does not recognize.
+
+        A typed config's ``from_dict`` keeps only annotated fields, so a key a newer
+        zeromodels added (e.g. ``rope_scaling_factor``, nested under ``text_config``)
+        is dropped. This finds those by round-trip: a key present in ``data`` but
+        absent from ``to_dict(from_dict(data))`` was not consumed. It recurses into
+        the nested sub-config dicts the repo format wraps fields in, and never flags
+        a legitimate key (a recognized key always round-trips back).
+        """
+        known = cls.from_dict(data).to_dict()
+
+        def diff(spec, seen):
+            out = set()
+            for key, value in spec.items():
+                if key not in seen:
+                    out.add(key)
+                elif isinstance(value, dict) and isinstance(seen.get(key), dict):
+                    out |= {f"{key}.{inner}" for inner in diff(value, seen[key])}
+            return out
+
+        return diff(data, known)
+
     @classmethod
     def _defaults(cls):
         return {name: getattr(cls, name, None) for name in cls._annotations()}

zeromodels/base/base_generation.py

@@ -83,14 +83,21 @@ def call_with_cache(self, token_ids, cache, cache_update_index, key_padding=None
     FULL_CHECKPOINT_SOURCES = {}
 
     @classmethod
-    def _load_backbone_from_full(cls, full_cls, repo_id, load_weights=True, **kwargs):
+    def _load_backbone_from_full(
+        cls, full_cls, repo_id, load_weights=True, skip_mismatch=False, **kwargs
+    ):
         """Build ``full_cls`` from ``repo_id`` and copy this head's backbone out of it.
 
         The head is constructed from the constructor kwargs it shares with the built full
         model, so extra config the head does not take (vision dims, M-RoPE sections) is
         dropped. Weights are matched by keras path suffix (see :meth:`_head_from_full`).
+        ``skip_mismatch`` is forwarded to the full model's own load, so the caller's
+        skip-on-shape-mismatch request survives the indirection (e.g. a resized-head
+        fine-tune), mirroring the suffix branch.
         """
-        full = full_cls.from_weights(repo_id, load_weights=load_weights, **kwargs)
+        full = full_cls.from_weights(
+            repo_id, load_weights=load_weights, skip_mismatch=skip_mismatch, **kwargs
+        )
         return cls._head_from_full(full, copy_weights=load_weights)
 
     @classmethod

zeromodels/base/base_mixin.py

@@ -377,16 +377,40 @@ def from_weights(
                     quantization=quantization,
                     **kwargs,
                 )
-        # A no-float load already quantized in place (and recorded the config);
-        # only quantize here when that path didn't run (functional models, or the
-        # release-`.h5` / timm paths).
-        if (
-            quantization is not None
-            and getattr(model, "_quantization_config", None) is None
-        ):
-            from zeromodels.quantization import quantize_model
+        # Quantize here only when nothing already did. A model can arrive already
+        # quantized two ways: a no-float load applied the caller's own request, or a
+        # natively-quantized repo (e.g. mxfp4 GPT-OSS) packed itself from its
+        # zm_config. In the second case the requested format is baked into the
+        # weights and cannot be re-quantized on load, so a caller asking for a
+        # different format is refused rather than silently handed the repo's format.
+        if quantization is not None:
+            existing = getattr(model, "_quantization_config", None)
+            if existing is None:
+                from zeromodels.quantization import quantize_model
+
+                model = quantize_model(model, quantization)
+            else:
+                from zeromodels.quantization.quant_config import resolve_config
 
-            model = quantize_model(model, quantization)
+                requested = resolve_config(quantization).mode
+                current = (
+                    existing.get("quant_method")
+                    if isinstance(existing, dict)
+                    else getattr(existing, "mode", None)
+                )
+                if current != requested:
+                    raise ValueError(
+                        f"{cls.__name__}.from_weights({identifier!r}, "
+                        f"quantization={quantization!r}) cannot be honored: "
+                        f"{identifier!r} ships its own {current!r} quantization baked "
+                        f"into its weights, which cannot be re-quantized to "
+                        f"{requested!r} during load. Load it without quantization= to "
+                        f"keep {current!r}. To build {requested!r} instead, quantize an "
+                        f"unquantized build: load a float source (e.g. "
+                        f"from_weights('hf:<upstream repo>', quantization={requested!r}), "
+                        f"which dequantizes {current!r} first), or dequantize_model() "
+                        f"this model and quantize_model(..., {requested!r})."
+                    )
 
         # First-load cache write: store the converted result so a later identical
         # call rebuilds from it. Best-effort: never breaks the returned model.
@@ -638,6 +662,21 @@ def _config_from_zm_spec(cls, spec, variant):
             # flat format: hyperparameters at the top level (transformers style).
             fields = {k: v for k, v in spec.items() if k not in ZM_METADATA_KEYS}
             if cls.config_class is not None:
+                # A typed config silently drops keys it does not annotate, so a
+                # behavior-changing field a newer repo adds (e.g. rope_scaling_factor)
+                # would be ignored and the model built wrong with no signal. Surface
+                # it instead, matching the untyped branch below (which raises on an
+                # unaccepted constructor kwarg).
+                unknown = cls.config_class.unknown_keys(fields)
+                if unknown:
+                    raise ValueError(
+                        f"{cls.__name__}.from_weights: '{variant}' zm_config.json has "
+                        f"config key(s) {sorted(unknown)} that this zeromodels "
+                        f"version does not recognize. The repo is likely newer than "
+                        f"the installed zeromodels: upgrade it (ignoring these could "
+                        f"change the model, e.g. context length or rope scaling), or "
+                        f"remove the key(s) if you are hand-editing the config."
+                    )
                 return cls.config_class.from_dict(fields).constructor_kwargs()
             fields.pop("model_type", None)
             return retuple(fields)
@@ -774,7 +813,11 @@ def _load_from_checkpoint_source(
         if cs.match == "path":
             full_cls = getattr(importlib.import_module(cs.module), cs.source)
             return cls._load_backbone_from_full(
-                full_cls, repo_id, load_weights=load_weights, **kwargs
+                full_cls,
+                repo_id,
+                load_weights=load_weights,
+                skip_mismatch=skip_mismatch,
+                **kwargs,
             )
 
         from zeromodels.conversion import copy_weights_by_path_suffix