v1.3.0

CHANGELOG.md

@@ -1,5 +1,13 @@
 # Change log
 
+## [v1.3.0] 2026-MM-DD (WIP)
+
+(TODO: Add changelog for v1.3.0)
+
+### Bug Fix
+- Fixed VLM (e.g. Gemma3) quantized model save/load: safetensors key prefixes from `save_pretrained` (`model.language_model.model.layers.*`) did not match the `from_config` model (`model.language_model.layers.*`), so `load_state_dict(..., assign=True)` left `qweight` / `scales` / `qzeros` at zero and generation degraded (`onecomp/quantized_model_loader.py`)
+  - Added `_remap_state_dict_keys()` to normalize checkpoint keys before `_replace_quantized_layers`; `_apply_known_state_dict_key_rewrites()` rewrites `.language_model.model.` → `.language_model.` (pattern-based, without requiring keys to exist in the empty model yet)
+
 ## [v1.2.0] 2026-06-08
 
 ### Save/Load Support for JointQ, RTN, and OneBit Quantizers

onecomp/__version__.py

@@ -6,4 +6,4 @@
 
 """
 
-__version__ = "1.2.0"
+__version__ = "1.3.0"

onecomp/quantized_model_loader.py

@@ -82,6 +82,13 @@
         # Load state_dict from safetensors
         state_dict = cls._load_state_dict_from_dir(save_directory)
 
+        # Align checkpoint key prefixes with the empty model built from config.
+        # Gemma3 VLMs are a common case: weights saved from from_pretrained
+        # use model.language_model.model.layers. (language_model is a
+        # ForCausalLM wrapper) while from_config exposes
+        # model.language_model.layers.* directly.
+        state_dict = cls._remap_state_dict_keys(state_dict, model)
+
         # Replace quantized layers with empty modules
         cls._replace_quantized_layers(model, state_dict, quant_config)
 
@@ -425,6 +432,87 @@
         parent = name_to_module.get(parent_name, model)
         setattr(parent, child_name, module)
 
+    @classmethod
+    def _remap_state_dict_keys(cls, state_dict: dict, model: torch.nn.Module) -> dict:
+        """Rewrite checkpoint keys so they match model parameter paths.
+
+        Quantized models are saved from a from_pretrained instance whose
+        submodule naming can differ from the from_config model built at
+        load time.  Without remapping, load_state_dict(..., assign=True)
+        silently skips mismatched keys and leaves layers at their empty-model
+        initial values (often all zeros for quantized buffers).
+
+        Remapping runs before _replace_quantized_layers, so the empty
+        model still exposes nn.Linear.weight rather than GPTQ buffers
+        (``qweight``, ``scales``, …).  Known prefix rewrites are therefore
+        applied from checkpoint key patterns alone; they must not require the
+        destination key to already exist in model.named_parameters().
+
+        Args:
+            state_dict: Tensors loaded from *.safetensors.
+            model: Empty model returned by _build_empty_model_from_config.
+
+        Returns:
+            A new dict with keys renamed where a unique target exists in
+            model.  Unmatched keys are kept under their original names
+            so strict=False loading can still proceed.
+        """
+        model_keys = set(dict(model.named_parameters())) | set(dict(model.named_buffers()))
+        if not any(
+            cls._apply_known_state_dict_key_rewrites(key) is not None for key in state_dict
+        ) and all(key in model_keys for key in state_dict):
+            return state_dict
+
+        remapped: dict = {}
+        rewrite_count = 0
+        for ckpt_key, tensor in state_dict.items():
+            if ckpt_key in model_keys:
+                remapped[ckpt_key] = tensor
+                continue
+
+            target_key = cls._resolve_state_dict_key(ckpt_key, model_keys)
+            if target_key is not None and target_key != ckpt_key:
+                remapped[target_key] = tensor
+                rewrite_count += 1
+            else:
+                remapped[ckpt_key] = tensor
+
+        if rewrite_count:
+            logger.info(
+                "Remapped %d state_dict key(s) to match model module paths",
+                rewrite_count,
+            )
+        return remapped
+
+    @staticmethod
+    def _apply_known_state_dict_key_rewrites(ckpt_key: str) -> Optional[str]:
+        """Return a rewritten key for known save/load prefix drift, else None."""
+        if ".language_model.model." in ckpt_key:
+            return ckpt_key.replace(".language_model.model.", ".language_model.", 1)
+        if ckpt_key.startswith("language_model.model."):
+            return (
+                "model."
+                + ckpt_key.replace("language_model.model.", "language_model.", 1)
+            )
+        return None
+
+    @staticmethod
+    def _resolve_state_dict_key(ckpt_key: str, model_keys: set) -> Optional[str]:
+        """Return the remapped key for ckpt_key, or None if unknown."""
+        rewritten = QuantizedModelLoader._apply_known_state_dict_key_rewrites(ckpt_key)
+        if rewritten is not None:
+            return rewritten
+
+        # Suffix fallback for other prefix drift: only when the suffix maps
+        # uniquely onto the current model (non-quantized params/buffers).
+        match = re.search(r"(layers\.\d+(?:\..+)*)$", ckpt_key)
+        if match:
+            suffix = match.group(1)
+            hits = [name for name in model_keys if name.endswith(suffix)]
+            if len(hits) == 1:
+                return hits[0]
+        return None
+
     @staticmethod
     def _load_state_dict_from_dir(directory: str) -> dict:
         """Load all tensors from *.safetensors in *directory*.
@@ -521,12 +609,16 @@
                 suffix = m.group(1)
                 hits = [s for s in sd_prefix_map if s.endswith(suffix)]
                 if len(hits) > 1:
-                    logger.warning(
-                        "Ambiguous suffix %s for %s: %s",
-                        suffix,
-                        name,
-                        hits,
-                    )
+                    lang_hits = [h for h in hits if "language_model" in h]
+                    if len(lang_hits) == 1:
+                        hits = lang_hits
+                    else:
+                        logger.warning(
+                            "Ambiguous suffix %s for %s: %s",
+                            suffix,
+                            name,
+                            hits,
+                        )
                 if hits:
                     alt_prefix = hits[0] + "."
                     return {