invarlock 0.2.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

invarlock/adapters/hf_mixin.py

@@ -4,10 +4,12 @@ Shared HuggingFace adapter mixin.
 
 Provides reusable functionality for InvarLock's HuggingFace adapters:
 - Device resolution helpers
+- Safe device movement for quantized models
 - Snapshot/restore with device awareness
 - Chunked snapshot helpers to reduce peak memory usage
 - Lightweight config serialization
 - Weight-tying detection plumbing
+- Quantization detection and capabilities
 """
 
 from __future__ import annotations
@@ -17,12 +19,15 @@ import json
 import os
 import tempfile
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 
 from invarlock.security import is_secure_path
 
+if TYPE_CHECKING:
+    from .capabilities import ModelCapabilities, QuantizationConfig
+
 SCALAR_TYPES = (int, float, str, bool)
 
 
@@ -91,6 +96,122 @@ class HFAdapterMixin:
 
         return torch.device(device_str)
 
+    def _safe_to_device(
+        self,
+        model: torch.nn.Module,
+        device: str | torch.device | None = "auto",
+        capabilities: ModelCapabilities | None = None,
+    ) -> torch.nn.Module:
+        """
+        Safely move model to device, respecting quantization constraints.
+
+        For quantized models (BNB, AWQ, GPTQ), device movement may be
+        impossible or already handled by the loading mechanism. This
+        method checks the model's capabilities before attempting .to().
+
+        Args:
+            model: The model to move.
+            device: Target device ("auto", "cuda", "mps", "cpu").
+            capabilities: Pre-computed capabilities, or None to auto-detect.
+
+        Returns:
+            The model (possibly on the new device, or unchanged if not movable).
+        """
+        target_device = self._resolve_device(device)
+
+        # Auto-detect capabilities if not provided
+        if capabilities is None:
+            capabilities = self._detect_capabilities(model)
+
+        # Check if model can be moved
+        if capabilities is not None and not capabilities.device_movable:
+            # Model handles its own device placement (e.g., BNB, AWQ, GPTQ)
+            # Log this decision for debugging but don't attempt .to()
+            return model
+
+        # Safe to move
+        return model.to(target_device)
+
+    def _detect_capabilities(self, model: torch.nn.Module) -> ModelCapabilities | None:
+        """
+        Detect model capabilities from a loaded model instance.
+
+        Args:
+            model: Loaded model instance.
+
+        Returns:
+            ModelCapabilities if detection succeeds, None otherwise.
+        """
+        try:
+            from .capabilities import detect_capabilities_from_model
+
+            return detect_capabilities_from_model(model)
+        except ImportError:
+            return None
+
+    def _is_quantized_model(self, model: torch.nn.Module) -> bool:
+        """
+        Check if a model is quantized (BNB, AWQ, GPTQ).
+
+        This is a quick heuristic check that doesn't require full
+        capability detection.
+
+        Args:
+            model: Model to check.
+
+        Returns:
+            True if the model appears to be quantized.
+        """
+        config = getattr(model, "config", None)
+        if config is None:
+            return False
+
+        # Check for quantization_config attribute
+        quant_cfg = getattr(config, "quantization_config", None)
+        if quant_cfg is not None:
+            return True
+
+        # Check for BNB-specific attributes on the model
+        if hasattr(model, "is_loaded_in_8bit") and model.is_loaded_in_8bit:
+            return True
+        if hasattr(model, "is_loaded_in_4bit") and model.is_loaded_in_4bit:
+            return True
+
+        # Check for quantized module types in the model
+        for module in model.modules():
+            module_name = module.__class__.__name__.lower()
+            if any(
+                q in module_name
+                for q in ["linear8bit", "linear4bit", "quantlinear", "awqlinear"]
+            ):
+                return True
+
+        return False
+
+    def _detect_quantization_config(
+        self, model: torch.nn.Module
+    ) -> QuantizationConfig | None:
+        """
+        Detect quantization configuration from a model.
+
+        Args:
+            model: Model to inspect.
+
+        Returns:
+            QuantizationConfig if quantization detected, None otherwise.
+        """
+        try:
+            from .capabilities import detect_quantization_from_config
+
+            config = getattr(model, "config", None)
+            if config is not None:
+                quant_cfg = detect_quantization_from_config(config)
+                if quant_cfg.is_quantized():
+                    return quant_cfg
+        except ImportError:
+            pass
+        return None
+
     # ------------------------------------------------------------------
     # HF save/export helpers
     # ------------------------------------------------------------------

invarlock/cli/commands/doctor.py

@@ -326,8 +326,14 @@ def doctor_command(
     try:
         import torch
 
+        torch_version = getattr(torch, "__version__", None)
         if not json_out:
-            console.print(f"[green]✅ PyTorch {torch.__version__}[/green]")
+            if torch_version:
+                console.print(f"[green]✅ PyTorch {torch_version}[/green]")
+            else:
+                console.print(
+                    "[yellow]⚠️  PyTorch present but version unavailable[/yellow]"
+                )
 
         # Device information
         from ..device import get_device_info

invarlock/cli/commands/run.py

@@ -81,6 +81,137 @@ GUARD_OVERHEAD_THRESHOLD = 0.01
 SPLIT_ALIASES: tuple[str, ...] = ("validation", "val", "dev", "eval", "test")
 
 
+def _coerce_mapping(obj: object) -> dict[str, Any]:
+    """Best-effort conversion of config-like objects to plain dicts."""
+
+    if isinstance(obj, dict):
+        return obj
+    try:
+        raw = getattr(obj, "_data", None)
+        if isinstance(raw, dict):
+            return raw
+    except Exception:
+        pass
+    try:
+        dumped = obj.model_dump()  # type: ignore[attr-defined]
+        if isinstance(dumped, dict):
+            return dumped
+    except Exception:
+        pass
+    try:
+        data = vars(obj)
+        if isinstance(data, dict):
+            return data
+    except Exception:
+        pass
+    return {}
+
+
+def _resolve_pm_acceptance_range(
+    cfg: InvarLockConfig | dict[str, Any] | None,
+) -> dict[str, float]:
+    """Resolve primary-metric acceptance bounds from config/env with safe defaults."""
+
+    base_min = 0.95
+    base_max = 1.10
+
+    cfg_min = None
+    cfg_max = None
+    try:
+        cfg_map = _coerce_mapping(cfg) if cfg is not None else {}
+        pm_section = cfg_map.get("primary_metric") if isinstance(cfg_map, dict) else {}
+        pm_map = _coerce_mapping(pm_section)
+        acceptance = (
+            pm_map.get("acceptance_range") if isinstance(pm_map, dict) else None
+        )
+        if isinstance(acceptance, dict):
+            if acceptance.get("min") is not None:
+                try:
+                    cfg_min = float(acceptance["min"])
+                except (TypeError, ValueError):
+                    cfg_min = None
+            if acceptance.get("max") is not None:
+                try:
+                    cfg_max = float(acceptance["max"])
+                except (TypeError, ValueError):
+                    cfg_max = None
+    except Exception:
+        cfg_min = None
+        cfg_max = None
+
+    def _parse_env(name: str) -> float | None:
+        try:
+            raw = os.environ.get(name, "")
+            if raw is None or str(raw).strip() == "":
+                return None
+            return float(raw)
+        except Exception:
+            return None
+
+    env_min = _parse_env("INVARLOCK_PM_ACCEPTANCE_MIN")
+    env_max = _parse_env("INVARLOCK_PM_ACCEPTANCE_MAX")
+
+    has_explicit = any(v is not None for v in (cfg_min, cfg_max, env_min, env_max))
+    if not has_explicit:
+        return {}
+
+    min_val = (
+        env_min if env_min is not None else cfg_min if cfg_min is not None else base_min
+    )
+    max_val = (
+        env_max if env_max is not None else cfg_max if cfg_max is not None else base_max
+    )
+
+    try:
+        if min_val is not None and min_val <= 0:
+            min_val = base_min
+    except Exception:
+        min_val = base_min
+    try:
+        if max_val is not None and max_val <= 0:
+            max_val = base_max
+    except Exception:
+        max_val = base_max
+
+    try:
+        if max_val is not None and min_val is not None and max_val < min_val:
+            max_val = min_val
+    except Exception:
+        max_val = base_max
+
+    return {"min": float(min_val), "max": float(max_val)}
+
+
+def _free_model_memory(model: object | None) -> None:
+    """Best-effort cleanup to release GPU memory for a model object."""
+    if model is None:
+        return
+    try:
+        import gc
+
+        del model
+        gc.collect()
+        if torch is not None and torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+    except Exception:
+        # Cleanup should never raise; fallback is to proceed without cache purge
+        pass
+
+
+def _should_measure_overhead(profile_normalized: str) -> tuple[bool, bool]:
+    """Return (measure_guard_overhead, skip_overhead) derived from env/profile."""
+
+    skip_overhead_env = (
+        os.environ.get("INVARLOCK_SKIP_OVERHEAD_CHECK", "").strip().lower()
+    )
+    skip_overhead = skip_overhead_env in {"1", "true", "yes"}
+    measure_guard_overhead = (
+        profile_normalized in {"ci", "release"} and not skip_overhead
+    )
+    return measure_guard_overhead, skip_overhead
+
+
 def _choose_dataset_split(
     *, requested: str | None, available: list[str] | None
 ) -> tuple[str, bool]:
@@ -1671,6 +1802,7 @@ def run_command(
             "edit": edit_meta,
             "guards": guard_metadata,
         }
+        pm_acceptance_range = _resolve_pm_acceptance_range(cfg)
 
         console.print(f"🔌 Adapter: {adapter.name}")
 
@@ -1746,6 +1878,10 @@ def run_command(
             "plugins": plugin_provenance,
             "run_id": run_id,
         }
+        run_context.setdefault("primary_metric", {})["acceptance_range"] = (
+            pm_acceptance_range
+        )
+        run_context["pm_acceptance_range"] = pm_acceptance_range
         run_context["model_profile"] = {
             "family": model_profile.family,
             "default_loss": model_profile.default_loss,
@@ -2756,18 +2892,26 @@ def run_command(
 
                 restore_fn = _restore2
             else:
-                # reload path
+                # reload path - properly free GPU memory before setting to None
+                _free_model_memory(model)
                 model = None
                 restore_fn = None
         except Exception:
             # On any failure, fall back to reload-per-attempt path
+            _free_model_memory(model)
             model = None
             restore_fn = None
 
         # RETRY LOOP - All report processing inside loop
         attempt = 1
         profile_normalized = (profile or "").lower()
-        measure_guard_overhead = profile_normalized in {"ci", "release"}
+        measure_guard_overhead, skip_overhead = _should_measure_overhead(
+            profile_normalized
+        )
+        if skip_overhead and profile_normalized in {"ci", "release"}:
+            console.print(
+                "[yellow]⚠️  Overhead check skipped via INVARLOCK_SKIP_OVERHEAD_CHECK[/yellow]"
+            )
 
         while True:
             # Reset RNG streams each attempt to guarantee determinism across retries
@@ -2933,6 +3077,8 @@ def run_command(
             if env_flags:
                 meta_payload["env_flags"] = env_flags
             report["meta"].update(meta_payload)
+            if pm_acceptance_range:
+                report["meta"]["pm_acceptance_range"] = pm_acceptance_range
             report["meta"]["model_profile"] = {
                 "family": model_profile.family,
                 "default_loss": model_profile.default_loss,

invarlock/core/registry.py

@@ -117,14 +117,24 @@ class CoreRegistry:
             module: str,
             class_name: str,
             status: str = "Available (fallback)",
+            required_deps: list[str] | None = None,
         ) -> None:
             if name not in registry:
+                # Check runtime dependencies for optional plugins
+                actual_available = True
+                actual_status = status
+                if required_deps:
+                    missing = self._check_runtime_dependencies(required_deps)
+                    if missing:
+                        actual_available = False
+                        actual_status = f"Needs extra: {', '.join(missing)}"
+
                 registry[name] = PluginInfo(
                     name=name,
                     module=module,
                     class_name=class_name,
-                    available=True,
-                    status=status,
+                    available=actual_available,
+                    status=actual_status,
                     package="invarlock",
                     version=INVARLOCK_VERSION,
                 )
@@ -147,27 +157,30 @@ class CoreRegistry:
         _fallback(
             self._adapters, "hf_mlm_auto", "invarlock.adapters", "HF_MLM_Auto_Adapter"
         )
-        # Optional plugin adapters (available when modules present)
+        # Optional plugin adapters (verify runtime dependencies)
         _fallback(
             self._adapters,
             "hf_gptq",
             "invarlock.plugins.hf_gptq_adapter",
             "HF_GPTQ_Adapter",
-            status="Available (fallback plugin)",
+            status="Available (plugin)",
+            required_deps=["auto_gptq"],
         )
         _fallback(
             self._adapters,
             "hf_awq",
             "invarlock.plugins.hf_awq_adapter",
             "HF_AWQ_Adapter",
-            status="Available (fallback plugin)",
+            status="Available (plugin)",
+            required_deps=["autoawq"],
         )
         _fallback(
             self._adapters,
             "hf_bnb",
             "invarlock.plugins.hf_bnb_adapter",
             "HF_BNB_Adapter",
-            status="Available (fallback plugin)",
+            status="Available (plugin)",
+            required_deps=["bitsandbytes"],
         )
 
         # Register built-in edits (quant-only core) and internal no-op
@@ -181,6 +194,21 @@ class CoreRegistry:
         _fallback(self._guards, "rmt", "invarlock.guards", "RMTGuard")
         _fallback(self._guards, "hello_guard", "invarlock.plugins", "HelloGuard")
 
+    def _check_runtime_dependencies(self, deps: list[str]) -> list[str]:
+        """
+        Check if runtime dependencies are actually importable.
+
+        Returns:
+            List of missing dependency names.
+        """
+        missing = []
+        for dep in deps:
+            try:
+                importlib.import_module(dep)
+            except ImportError:
+                missing.append(dep)
+        return missing
+
     def _create_plugin_info(
         self, entry_point: EntryPoint, plugin_type: str
     ) -> PluginInfo:

invarlock/guards/variance.py

@@ -39,6 +39,30 @@ from .policies import VariancePolicyDict
 __all__ = ["equalise_residual_variance", "equalise_branch_variance", "VarianceGuard"]
 
 
+def _safe_mean(
+    samples: list[float] | np.ndarray, default: float | None = None
+) -> float | None:
+    """
+    Compute mean of samples, returning default if empty.
+
+    Avoids numpy RuntimeWarning "Mean of empty slice" when samples is empty
+    or contains no valid values.
+
+    Args:
+        samples: List or array of float values.
+        default: Value to return if samples is empty.
+
+    Returns:
+        Mean value or default if samples is empty.
+    """
+    if samples is None:
+        return default
+    arr = np.asarray(samples)
+    if arr.size == 0:
+        return default
+    return float(np.nanmean(arr))
+
+
 try:  # Optional dependency: tqdm (progress bars)
     from tqdm.auto import tqdm as _tqdm
 except Exception:  # pragma: no cover - exercised only when tqdm is absent
@@ -1472,7 +1496,14 @@ class VarianceGuard(Guard):
 
         if coverage >= min_coverage and not self._scales:
             ppl_no_ve_samples = ppl_no_ve_samples[:coverage]
-            ppl_no_ve_mean = float(np.mean(ppl_no_ve_samples))
+            ppl_no_ve_mean = _safe_mean(ppl_no_ve_samples)
+            if ppl_no_ve_mean is None:
+                # No valid samples - cannot compute mean
+                self._ratio_ci = None
+                predictive_state["reason"] = "no_valid_samples"
+                self._predictive_gate_state = predictive_state
+                self._stats["predictive_gate"] = predictive_state.copy()
+                return
             self.set_ab_results(
                 ppl_no_ve=ppl_no_ve_mean,
                 ppl_with_ve=ppl_no_ve_mean,
@@ -1527,8 +1558,12 @@ class VarianceGuard(Guard):
                     n_bootstrap=500,
                     seed=calib_seed,
                 )
-                ppl_no_ve_mean = float(np.mean(ppl_no_ve_samples))
-                ppl_with_ve_mean = float(np.mean(ppl_with_ve_samples))
+                ppl_no_ve_mean = _safe_mean(ppl_no_ve_samples)
+                ppl_with_ve_mean = _safe_mean(ppl_with_ve_samples)
+                if ppl_no_ve_mean is None or ppl_with_ve_mean is None:
+                    # Fallback if means couldn't be computed
+                    ppl_no_ve_mean = ppl_no_ve_mean or 0.0
+                    ppl_with_ve_mean = ppl_with_ve_mean or 0.0
                 self.set_ab_results(
                     ppl_no_ve=ppl_no_ve_mean,
                     ppl_with_ve=ppl_with_ve_mean,
@@ -2118,7 +2153,7 @@ class VarianceGuard(Guard):
 
                 if coverage >= min_coverage and not self._scales:
                     ppl_no_ve_samples = ppl_no_ve_samples[:coverage]
-                    ppl_no_ve_mean = float(np.mean(ppl_no_ve_samples))
+                    ppl_no_ve_mean = _safe_mean(ppl_no_ve_samples, default=0.0)
                     self.set_ab_results(
                         ppl_no_ve=ppl_no_ve_mean,
                         ppl_with_ve=ppl_no_ve_mean,
@@ -2158,8 +2193,8 @@ class VarianceGuard(Guard):
                             n_bootstrap=500,
                             seed=calib_seed,
                         )
-                        ppl_no_ve_mean = float(np.mean(ppl_no_ve_samples))
-                        ppl_with_ve_mean = float(np.mean(ppl_with_ve_samples))
+                        ppl_no_ve_mean = _safe_mean(ppl_no_ve_samples, default=0.0)
+                        ppl_with_ve_mean = _safe_mean(ppl_with_ve_samples, default=0.0)
                         self.set_ab_results(
                             ppl_no_ve=ppl_no_ve_mean,
                             ppl_with_ve=ppl_with_ve_mean,

invarlock/plugins/hf_awq_adapter.py

@@ -4,12 +4,16 @@ HuggingFace AWQ Adapter (plugin)
 
 Optional adapter for loading AWQ-quantized causal LMs from the Hub.
 Requires the `autoawq` extra on supported platforms (typically Linux/CUDA).
+
+AWQ models are pre-quantized and typically handle device placement internally
+during loading. This adapter does NOT call .to() on the loaded model.
 """
 
 from __future__ import annotations
 
 from typing import Any
 
+from invarlock.adapters.capabilities import ModelCapabilities
 from invarlock.adapters.hf_mixin import HFAdapterMixin
 from invarlock.core.api import ModelAdapter
 from invarlock.core.error_utils import wrap_errors
@@ -56,7 +60,24 @@ class HF_AWQ_Adapter(HFAdapterMixin, ModelAdapter):
                 trust_remote_code=True,
                 **{k: v for k, v in kwargs.items() if k != "device"},
             )
-        return model.to(self._resolve_device(device))
+
+        # AWQ models are pre-quantized; use safe device movement
+        # which respects the model's device constraints
+        return self._safe_to_device(
+            model, device, capabilities=ModelCapabilities.for_awq()
+        )
+
+    def get_capabilities(self, model: Any) -> ModelCapabilities:
+        """Return capabilities for an AWQ-quantized model."""
+        config = getattr(model, "config", None)
+        group_size = 128  # Default AWQ group size
+        if config is not None:
+            quant_cfg = getattr(config, "quantization_config", None)
+            if isinstance(quant_cfg, dict):
+                group_size = quant_cfg.get("group_size", 128)
+            elif quant_cfg is not None:
+                group_size = getattr(quant_cfg, "group_size", 128)
+        return ModelCapabilities.for_awq(group_size=group_size)
 
     def can_handle(self, model: Any) -> bool:
         cfg = getattr(model, "config", None)