invarlock 0.3.3__py3-none-any.whl → 0.3.5__py3-none-any.whl

invarlock/__init__.py

@@ -12,7 +12,7 @@ For torch-dependent functionality, see subpackages under `invarlock.*`:
 - `invarlock.eval`: Metrics, guard-overhead checks, and certification
 """
 
-__version__ = "0.3.3"
+__version__ = "0.3.5"
 
 # Core exports - torch-independent
 from .config import CFG, Defaults, get_default_config

invarlock/adapters/auto.py

@@ -209,18 +209,18 @@ class _DelegatingAdapter(ModelAdapter):
 class HF_Causal_Auto_Adapter(_DelegatingAdapter):
     name = "hf_causal_auto"
 
-    def load_model(self, model_id: str, device: str = "auto") -> Any:
+    def load_model(self, model_id: str, device: str = "auto", **kwargs: Any) -> Any:
         delegate = self._ensure_delegate_from_id(model_id)
-        return delegate.load_model(model_id, device=device)
+        return delegate.load_model(model_id, device=device, **kwargs)
 
 
 class HF_MLM_Auto_Adapter(_DelegatingAdapter):
     name = "hf_mlm_auto"
 
-    def load_model(self, model_id: str, device: str = "auto") -> Any:
+    def load_model(self, model_id: str, device: str = "auto", **kwargs: Any) -> Any:
         # Force BERT-like adapter for MLM families
         HF_BERT_Adapter = _importlib.import_module(
             ".hf_bert", __package__
         ).HF_BERT_Adapter
         self._delegate = HF_BERT_Adapter()
-        return self._delegate.load_model(model_id, device=device)
+        return self._delegate.load_model(model_id, device=device, **kwargs)

invarlock/adapters/hf_bert.py

@@ -41,7 +41,9 @@ class HF_BERT_Adapter(HFAdapterMixin, ModelAdapter):
 
     name = "hf_bert"
 
-    def load_model(self, model_id: str, device: str = "auto") -> ModuleType | Any:
+    def load_model(
+        self, model_id: str, device: str = "auto", **kwargs: Any
+    ) -> ModuleType | Any:
         """
         Load a HuggingFace BERT model.
 
@@ -68,7 +70,7 @@ class HF_BERT_Adapter(HFAdapterMixin, ModelAdapter):
                 "MODEL-LOAD-FAILED: transformers AutoModelForMaskedLM",
                 lambda e: {"model_id": model_id},
             ):
-                model = AutoModelForMaskedLM.from_pretrained(model_id)
+                model = AutoModelForMaskedLM.from_pretrained(model_id, **kwargs)
         except Exception:
             with wrap_errors(
                 ModelLoadError,
@@ -76,10 +78,9 @@ class HF_BERT_Adapter(HFAdapterMixin, ModelAdapter):
                 "MODEL-LOAD-FAILED: transformers AutoModel",
                 lambda e: {"model_id": model_id},
             ):
-                model = AutoModel.from_pretrained(model_id)
+                model = AutoModel.from_pretrained(model_id, **kwargs)
 
-        target_device = self._resolve_device(device)
-        return model.to(target_device)
+        return self._safe_to_device(model, device)
 
     def can_handle(self, model: ModuleType | Any) -> bool:
         """

invarlock/adapters/hf_gpt2.py

@@ -48,7 +48,9 @@ class HF_GPT2_Adapter(HFAdapterMixin, ModelAdapter):
 
     name = "hf_gpt2"
 
-    def load_model(self, model_id: str, device: str = "auto") -> ModuleType | Any:
+    def load_model(
+        self, model_id: str, device: str = "auto", **kwargs: Any
+    ) -> ModuleType | Any:
         """
         Load a HuggingFace GPT-2 model.
 
@@ -75,10 +77,9 @@ class HF_GPT2_Adapter(HFAdapterMixin, ModelAdapter):
                 "MODEL-LOAD-FAILED: transformers AutoModelForCausalLM",
                 lambda e: {"model_id": model_id},
             ):
-                model = AutoModelForCausalLM.from_pretrained(model_id)
+                model = AutoModelForCausalLM.from_pretrained(model_id, **kwargs)
 
-            target_device = self._resolve_device(device)
-            return model.to(target_device)
+            return self._safe_to_device(model, device)
         except DependencyError:
             if LIGHT_IMPORT:
                 # Minimal stand-in that satisfies downstream interface requirements

invarlock/adapters/hf_llama.py

@@ -41,7 +41,9 @@ class HF_LLaMA_Adapter(HFAdapterMixin, ModelAdapter):
 
     name = "hf_llama"
 
-    def load_model(self, model_id: str, device: str = "auto") -> ModuleType | Any:
+    def load_model(
+        self, model_id: str, device: str = "auto", **kwargs: Any
+    ) -> ModuleType | Any:
         """
         Load a HuggingFace LLaMA model.
 
@@ -67,7 +69,7 @@ class HF_LLaMA_Adapter(HFAdapterMixin, ModelAdapter):
             "MODEL-LOAD-FAILED: transformers AutoModelForCausalLM",
             lambda e: {"model_id": model_id},
         ):
-            model = AutoModelForCausalLM.from_pretrained(model_id)
+            model = AutoModelForCausalLM.from_pretrained(model_id, **kwargs)
 
         # Use safe device movement that respects quantization constraints
         return self._safe_to_device(model, device)

invarlock/adapters/hf_mixin.py

@@ -119,6 +119,10 @@ class HFAdapterMixin:
         """
         target_device = self._resolve_device(device)
 
+        # If transformers already sharded/placed the model, skip explicit .to().
+        if getattr(model, "hf_device_map", None):
+            return model
+
         # Auto-detect capabilities if not provided
         if capabilities is None:
             capabilities = self._detect_capabilities(model)
@@ -334,7 +338,9 @@ class HFAdapterMixin:
             if hasattr(model, "config")
             else {},
             "params": {},
+            "params_meta": {},
             "buffers": {},
+            "buffers_meta": {},
             "device_map": {},
             "weight_tying": self._extract_weight_tying_info(model),
         }
@@ -344,6 +350,10 @@ class HFAdapterMixin:
             file_path = snapshot_dir / filename
             torch.save(param.detach().cpu(), file_path)
             manifest["params"][name] = filename
+            manifest["params_meta"][name] = {
+                "shape": [int(x) for x in param.shape],
+                "dtype": str(param.dtype),
+            }
             manifest["device_map"][name] = str(param.device)
 
         for name, buffer in model.named_buffers():
@@ -351,6 +361,10 @@ class HFAdapterMixin:
             file_path = snapshot_dir / filename
             torch.save(buffer.detach().cpu(), file_path)
             manifest["buffers"][name] = filename
+            manifest["buffers_meta"][name] = {
+                "shape": [int(x) for x in buffer.shape],
+                "dtype": str(buffer.dtype),
+            }
             manifest["device_map"][f"buffer::{name}"] = str(buffer.device)
 
         manifest_path = snapshot_dir / "manifest.json"
@@ -377,24 +391,89 @@ class HFAdapterMixin:
 
         device_map = manifest.get("device_map", {})
 
-        for name, filename in manifest.get("params", {}).items():
+        params_manifest = manifest.get("params", {})
+        if not isinstance(params_manifest, dict):
+            raise TypeError("Invalid snapshot manifest: params must be a mapping")
+        buffers_manifest = manifest.get("buffers", {})
+        if not isinstance(buffers_manifest, dict):
+            raise TypeError("Invalid snapshot manifest: buffers must be a mapping")
+        params_meta = manifest.get("params_meta", {})
+        buffers_meta = manifest.get("buffers_meta", {})
+
+        # Preflight: ensure manifest/model agreement and tensor readability before copying.
+        for name, filename in params_manifest.items():
+            if name not in param_map:
+                raise KeyError(f"Snapshot parameter missing in target model: {name}")
+            if not isinstance(filename, str) or not filename:
+                raise TypeError(f"Invalid snapshot manifest filename for param: {name}")
             file_path = snapshot_dir / filename
-            if name not in param_map or not file_path.exists():
-                continue
+            if not file_path.exists():
+                raise FileNotFoundError(
+                    f"Missing snapshot tensor for param: {file_path}"
+                )
+            tensor = torch.load(file_path, map_location="cpu")
+            if not isinstance(tensor, torch.Tensor):
+                raise TypeError(f"Invalid snapshot tensor payload for param: {name}")
+            meta = params_meta.get(name) if isinstance(params_meta, dict) else None
+            if isinstance(meta, dict):
+                expected_shape = meta.get("shape")
+                expected_dtype = meta.get("dtype")
+                if isinstance(expected_shape, list) and list(tensor.shape) != list(
+                    expected_shape
+                ):
+                    raise ValueError(
+                        f"Snapshot tensor shape mismatch for param: {name}"
+                    )
+                if isinstance(expected_dtype, str) and expected_dtype:
+                    if str(tensor.dtype) != expected_dtype:
+                        raise ValueError(
+                            f"Snapshot tensor dtype mismatch for param: {name}"
+                        )
+
+        for name, filename in buffers_manifest.items():
+            if name not in buffer_map:
+                raise KeyError(f"Snapshot buffer missing in target model: {name}")
+            if not isinstance(filename, str) or not filename:
+                raise TypeError(
+                    f"Invalid snapshot manifest filename for buffer: {name}"
+                )
+            file_path = snapshot_dir / filename
+            if not file_path.exists():
+                raise FileNotFoundError(
+                    f"Missing snapshot tensor for buffer: {file_path}"
+                )
+            tensor = torch.load(file_path, map_location="cpu")
+            if not isinstance(tensor, torch.Tensor):
+                raise TypeError(f"Invalid snapshot tensor payload for buffer: {name}")
+            meta = buffers_meta.get(name) if isinstance(buffers_meta, dict) else None
+            if isinstance(meta, dict):
+                expected_shape = meta.get("shape")
+                expected_dtype = meta.get("dtype")
+                if isinstance(expected_shape, list) and list(tensor.shape) != list(
+                    expected_shape
+                ):
+                    raise ValueError(
+                        f"Snapshot tensor shape mismatch for buffer: {name}"
+                    )
+                if isinstance(expected_dtype, str) and expected_dtype:
+                    if str(tensor.dtype) != expected_dtype:
+                        raise ValueError(
+                            f"Snapshot tensor dtype mismatch for buffer: {name}"
+                        )
+
+        # Restore parameters/buffers (second pass) after successful preflight.
+        for name, filename in params_manifest.items():
             target = param_map[name]
             target_device = torch.device(device_map.get(name, str(target.device)))
-            tensor = torch.load(file_path, map_location="cpu")
+            tensor = torch.load(snapshot_dir / filename, map_location="cpu")
             with torch.no_grad():
                 target.copy_(tensor.to(target_device))
 
-        for name, filename in manifest.get("buffers", {}).items():
-            file_path = snapshot_dir / filename
-            if name not in buffer_map or not file_path.exists():
-                continue
+        for name, filename in buffers_manifest.items():
             target = buffer_map[name]
             key = f"buffer::{name}"
             target_device = torch.device(device_map.get(key, str(target.device)))
-            tensor = torch.load(file_path, map_location="cpu")
+            tensor = torch.load(snapshot_dir / filename, map_location="cpu")
             target.copy_(tensor.to(target_device))
 
         original_tying = manifest.get("weight_tying", {})

invarlock/adapters/hf_t5.py

@@ -30,7 +30,9 @@ class HF_T5_Adapter(HFAdapterMixin, ModelAdapter):
 
     name = "hf_t5"
 
-    def load_model(self, model_id: str, device: str = "auto") -> ModuleType | Any:  # type: ignore[override]
+    def load_model(  # type: ignore[override]
+        self, model_id: str, device: str = "auto", **kwargs: Any
+    ) -> ModuleType | Any:
         with wrap_errors(
             DependencyError,
             "E203",
@@ -45,8 +47,8 @@ class HF_T5_Adapter(HFAdapterMixin, ModelAdapter):
             "MODEL-LOAD-FAILED: transformers AutoModelForSeq2SeqLM",
             lambda e: {"model_id": model_id},
         ):
-            model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
-        return model.to(self._resolve_device(device))
+            model = AutoModelForSeq2SeqLM.from_pretrained(model_id, **kwargs)
+        return self._safe_to_device(model, device)
 
     def can_handle(self, model: ModuleType | Any) -> bool:  # type: ignore[override]
         cfg = getattr(model, "config", None)

invarlock/cli/commands/certify.py

@@ -378,7 +378,7 @@ def certify_command(
                         f"Context: device={device}, adapter={adapter_name}, edit={edit_name}. "
                         "Baseline ok; edited failed to compute ppl. "
                         "Try: use an accelerator (mps/cuda), force float32, reduce max_modules, "
-                        "or lower batch size (INVARLOCK_SCORES_BATCH_SIZE)."
+                        "or lower the evaluation batch size."
                     ),
                     details={
                         "device": device,