invarlock 0.2.0__py3-none-any.whl

invarlock/core/types.py

@@ -0,0 +1,154 @@
+"""
+InvarLock Core Types
+================
+
+Core type definitions and enums used throughout InvarLock.
+Torch-independent type system for cross-module compatibility.
+"""
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, NamedTuple
+
+
+class EditType(Enum):
+    """Types of model edits supported by InvarLock."""
+
+    QUANTIZATION = "quantization"
+    SPARSITY = "sparsity"
+    MIXED = "mixed"
+
+
+class GuardType(Enum):
+    """Types of safety guards available."""
+
+    INVARIANTS = "invariants"
+    SPECTRAL = "spectral"
+    VARIANCE = "variance"
+    RMT = "rmt"
+    NOOP = "noop"
+
+
+class RunStatus(Enum):
+    """Execution status for pipeline runs."""
+
+    PENDING = "pending"
+    RUNNING = "running"
+    SUCCESS = "success"
+    FAILED = "failed"
+    ROLLBACK = "rollback"
+    CANCELLED = "cancelled"
+
+
+class LogLevel(Enum):
+    """Logging levels for events."""
+
+    DEBUG = "DEBUG"
+    INFO = "INFO"
+    WARNING = "WARNING"
+    ERROR = "ERROR"
+    CRITICAL = "CRITICAL"
+
+
+@dataclass
+class ModelInfo:
+    """Basic model information."""
+
+    model_id: str
+    architecture: str
+    parameters: int
+    device: str
+    precision: str = "float32"
+
+
+@dataclass
+class EditInfo:
+    """Information about an applied edit."""
+
+    name: str
+    type: EditType
+    parameters: dict[str, Any]
+    compression_ratio: float | None = None
+    target_metrics: dict[str, float] | None = None
+
+
+@dataclass
+class GuardResult:
+    """Result from a guard validation."""
+
+    guard_name: str
+    passed: bool
+    score: float | None = None
+    threshold: float | None = None
+    message: str | None = None
+    details: dict[str, Any] | None = None
+
+
+class ValidationResult(NamedTuple):
+    """Result from validation operations."""
+
+    passed: bool
+    score: float
+    threshold: float
+    message: str = ""
+
+
+@dataclass
+class GuardOutcome:
+    """Result from a guard execution."""
+
+    name: str
+    passed: bool
+    action: str = "none"
+    violations: list[dict[str, Any]] | None = None
+    metrics: dict[str, Any] | None = None
+
+    def __post_init__(self):
+        if self.violations is None:
+            self.violations = []
+        if self.metrics is None:
+            self.metrics = {}
+
+
+@dataclass
+class PolicyConfig:
+    """Configuration for guard policies."""
+
+    on_violation: str = "warn"
+    guard_overrides: dict[str, str] | None = None
+    enable_auto_rollback: bool = False
+
+    def __post_init__(self):
+        if self.guard_overrides is None:
+            self.guard_overrides = {}
+
+    def get_action_for_guard(self, guard_name: str, requested_action: str) -> str:
+        """Get the action for a specific guard."""
+        # Check if there's an override for this guard
+        if self.guard_overrides and guard_name in self.guard_overrides:
+            return self.guard_overrides[guard_name]
+
+        # If requested action is not 'none', use it
+        if requested_action != "none":
+            return requested_action
+
+        # Fall back to global default
+        return self.on_violation
+
+
+def get_worst_action(actions: list[str]) -> str:
+    """Get the worst (most severe) action from a list."""
+    action_priority = {"none": 0, "warn": 1, "rollback": 2, "abort": 3}
+
+    if not actions:
+        return "none"
+
+    return max(actions, key=lambda action: action_priority.get(action, 0))
+
+
+# Type aliases for clarity
+DeviceSpec = str | Any  # Device specification
+ConfigDict = dict[str, Any]  # Configuration dictionary
+MetricsDict = dict[str, float | int | str | bool]  # Metrics
+LayerIndex = int  # Layer index
+HeadIndex = int  # Attention head index

invarlock/edits/__init__.py

@@ -0,0 +1,12 @@
+"""Edit namespace (`invarlock.edits`) re-exporting built-in edits."""
+
+from __future__ import annotations
+
+from invarlock.core.abi import INVARLOCK_CORE_ABI as INVARLOCK_CORE_ABI
+
+from .quant_rtn import RTNQuantEdit
+
+__all__ = [
+    "RTNQuantEdit",
+    "INVARLOCK_CORE_ABI",
+]

invarlock/edits/_edit_utils.py

@@ -0,0 +1,249 @@
+"""
+InvarLock Edit Utilities
+====================
+
+Shared helper functions for edit implementations.
+Common functionality used across multiple edit types.
+"""
+
+from __future__ import annotations
+
+import warnings
+from typing import Any
+
+try:
+    import torch.nn as nn
+
+    TORCH_AVAILABLE = True
+except ImportError:
+    TORCH_AVAILABLE = False
+
+__all__ = [
+    "validate_model_structure",
+    "calculate_compression_ratio",
+    "get_layer_info",
+    "validate_edit_parameters",
+    "create_edit_metadata",
+]
+
+
+def validate_model_structure(
+    model_desc: dict[str, Any], required_fields: list[str]
+) -> bool:
+    """
+    Validate that a model description has required fields for an edit.
+
+    Args:
+        model_desc: Model description from adapter
+        required_fields: List of required field names
+
+    Returns:
+        True if all required fields are present
+    """
+    missing_fields = []
+
+    for field in required_fields:
+        if field not in model_desc:
+            missing_fields.append(field)
+
+    if missing_fields:
+        warnings.warn(
+            f"Model missing required fields for edit: {missing_fields}", stacklevel=2
+        )
+        return False
+
+    return True
+
+
+def calculate_compression_ratio(original_params: int, final_params: int) -> float:
+    """
+    Calculate parameter compression ratio.
+
+    Args:
+        original_params: Original parameter count
+        final_params: Final parameter count after edit
+
+    Returns:
+        Compression ratio (final/original)
+    """
+    if original_params == 0:
+        return 1.0
+
+    return final_params / original_params
+
+
+def get_layer_info(model_desc: dict[str, Any], layer_idx: int) -> dict[str, Any]:
+    """
+    Extract information for a specific layer.
+
+    Args:
+        model_desc: Model description from adapter
+        layer_idx: Index of the layer
+
+    Returns:
+        Dictionary with layer-specific information
+    """
+    n_layers = model_desc.get("n_layer", 0)
+
+    if layer_idx >= n_layers:
+        raise IndexError(f"Layer index {layer_idx} >= {n_layers}")
+
+    heads_per_layer = model_desc.get("heads_per_layer", [])
+    mlp_dims = model_desc.get("mlp_dims", [])
+
+    layer_info = {
+        "layer_idx": layer_idx,
+        "n_heads": heads_per_layer[layer_idx]
+        if layer_idx < len(heads_per_layer)
+        else None,
+        "mlp_dim": mlp_dims[layer_idx] if layer_idx < len(mlp_dims) else None,
+        "hidden_size": model_desc.get("hidden_size"),
+    }
+
+    return layer_info
+
+
+def validate_edit_parameters(
+    params: dict[str, Any],
+    required_params: list[str],
+    optional_params: dict[str, Any] | None = None,
+) -> tuple[bool, str]:
+    """
+    Validate edit parameters.
+
+    Args:
+        params: Parameters to validate
+        required_params: List of required parameter names
+        optional_params: Dict of optional parameters with default values
+
+    Returns:
+        (is_valid, error_message)
+    """
+    # Check required parameters
+    missing_required = []
+    for param in required_params:
+        if param not in params:
+            missing_required.append(param)
+
+    if missing_required:
+        return False, f"Missing required parameters: {missing_required}"
+
+    # Add default values for optional parameters
+    if optional_params:
+        for param, default_value in optional_params.items():
+            if param not in params:
+                params[param] = default_value
+
+    return True, "Parameters valid"
+
+
+def create_edit_metadata(
+    edit_name: str,
+    model_desc: dict[str, Any],
+    parameters: dict[str, Any],
+    result: dict[str, Any],
+) -> dict[str, Any]:
+    """
+    Create standardized edit metadata.
+
+    Args:
+        edit_name: Name of the edit operation
+        model_desc: Original model description
+        parameters: Edit parameters used
+        result: Edit operation results
+
+    Returns:
+        Standardized metadata dictionary
+    """
+    original_params = model_desc.get("total_params", 0)
+    final_params = result.get("final_params", original_params)
+
+    metadata: dict[str, Any] = {
+        "name": edit_name,
+        "parameters": parameters.copy(),
+        "original_model": {
+            "n_layer": model_desc.get("n_layer", 0),
+            "total_params": original_params,
+            "model_type": model_desc.get("model_type", "unknown"),
+        },
+        "results": {
+            "final_params": final_params,
+            "compression_ratio": calculate_compression_ratio(
+                original_params, final_params
+            ),
+            "layers_modified": result.get("layers_modified", []),
+            "success": result.get("success", True),
+        },
+    }
+
+    # Add edit-specific results
+    results_dict = metadata["results"]
+    assert isinstance(results_dict, dict)
+    for key, value in result.items():
+        if key not in results_dict:
+            results_dict[key] = value
+
+    return metadata
+
+
+def get_module_parameter_count(module) -> int:
+    """
+    Count parameters in a module (if torch is available).
+
+    Args:
+        module: PyTorch module
+
+    Returns:
+        Number of parameters
+    """
+    if not TORCH_AVAILABLE:
+        return 0
+
+    if not isinstance(module, nn.Module):
+        return 0
+
+    return sum(p.numel() for p in module.parameters())
+
+
+def validate_compression_target(
+    target_ratio: Any, min_ratio: float = 0.1, max_ratio: float = 1.0
+) -> bool:
+    """
+    Validate compression target ratio.
+
+    Args:
+        target_ratio: Desired compression ratio
+        min_ratio: Minimum allowed ratio
+        max_ratio: Maximum allowed ratio
+
+    Returns:
+        True if ratio is valid
+    """
+    if not isinstance(target_ratio, int | float):
+        return False
+
+    return min_ratio <= target_ratio <= max_ratio
+
+
+def create_layer_mask(
+    n_layers: int, layers_to_edit: list[int] | None = None
+) -> list[bool]:
+    """
+    Create a boolean mask for layers to edit.
+
+    Args:
+        n_layers: Total number of layers
+        layers_to_edit: List of layer indices to edit (None = all layers)
+
+    Returns:
+        Boolean mask where True indicates layer should be edited
+    """
+    if layers_to_edit is None:
+        return [True] * n_layers
+
+    mask = [False] * n_layers
+    for layer_idx in layers_to_edit:
+        if 0 <= layer_idx < n_layers:
+            mask[layer_idx] = True
+
+    return mask

invarlock/edits/_external_utils.py

@@ -0,0 +1,268 @@
+"""
+External Utilities for Edit Operations
+=====================================
+
+Utilities for integrating with external edit backends and guard chains.
+Provides common functionality for model snapshots, validation, and guard policies.
+"""
+
+from __future__ import annotations
+
+import warnings
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+
+class ExternalBackend:
+    """Base class for external edit backends."""
+
+    def apply_edit(self, model: nn.Module, config: dict, calib: Any) -> dict:
+        """Apply edit to model."""
+        raise NotImplementedError()
+
+    def get_edit_info(self, model: nn.Module, config: dict) -> dict:
+        """Get edit information without applying."""
+        raise NotImplementedError()
+
+
+class ExternalEditWrapper:
+    """Wrapper for external edit implementations."""
+
+    def __init__(self, backend: ExternalBackend):
+        self.backend = backend
+
+    def preview(self, model: nn.Module, adapter: Any, calib: Any) -> dict:
+        """Preview edit operation."""
+        return {"plan": {}, "estimated_sparsity": {}, "preview_metrics": {}}
+
+    def apply(self, model: nn.Module, adapter: Any, plan: dict) -> dict:
+        """Apply edit operation."""
+        return {"actual_sparsity": {}, "modified_layers": [], "metrics": {}}
+
+
+def check_dependencies(edit_type: str) -> bool:
+    """
+    Check if dependencies for edit type are available.
+
+    Args:
+        edit_type: Type of edit (e.g., 'quant')
+
+    Returns:
+        True if dependencies are available
+    """
+    if edit_type == "svd":
+        return False
+    elif edit_type == "quant":
+        # Quantization uses standard PyTorch
+        return True
+    else:
+        return False
+
+
+def validate_edit_config(config: dict, edit_type: str) -> dict:
+    """
+    Validate and normalize edit configuration.
+
+    Args:
+        config: Configuration dictionary
+        edit_type: Type of edit operation
+
+    Returns:
+        Validated configuration dictionary
+    """
+    validated_config = config.copy()
+
+    if edit_type == "svd":
+        raise ValueError("Unsupported edit type: svd")
+
+    elif edit_type == "quant":
+        # Validate quantization specific config
+        bits = validated_config.get("bits", 8)
+        if bits not in [4, 8, 16]:
+            warnings.warn(
+                f"Unusual bit width {bits}, common values are 4, 8, or 16", stacklevel=2
+            )
+
+    return validated_config
+
+
+def compute_edit_metrics(
+    model_before: nn.Module, model_after: nn.Module, config: dict
+) -> dict:
+    """
+    Compute metrics comparing model before and after edit.
+
+    Args:
+        model_before: Model state before edit
+        model_after: Model state after edit
+        config: Edit configuration
+
+    Returns:
+        Dictionary of computed metrics
+    """
+    # Count parameters
+    params_before = sum(p.numel() for p in model_before.parameters())
+    params_after = sum(p.numel() for p in model_after.parameters())
+
+    # Count non-zero parameters (for sparsity)
+    nonzero_before = sum((p != 0).sum().item() for p in model_before.parameters())
+    nonzero_after = sum((p != 0).sum().item() for p in model_after.parameters())
+
+    # Calculate compression ratios
+    param_ratio = params_after / params_before if params_before > 0 else 1.0
+    sparsity_ratio = (
+        1.0 - (nonzero_after / nonzero_before) if nonzero_before > 0 else 0.0
+    )
+
+    return {
+        "params_before": params_before,
+        "params_after": params_after,
+        "nonzero_before": nonzero_before,
+        "nonzero_after": nonzero_after,
+        "param_compression_ratio": param_ratio,
+        "sparsity_achieved": sparsity_ratio,
+        "memory_saved_mb": (params_before - params_after) * 4 / (1024 * 1024),
+    }
+
+
+def prepare_calibration_data(calib_data: Any, config: dict) -> Any:
+    """
+    Prepare calibration data for edit operations.
+
+    Args:
+        calib_data: Raw calibration data
+        config: Edit configuration
+
+    Returns:
+        Prepared calibration data
+    """
+    if calib_data is None:
+        return None
+
+    # If it's already prepared (has __iter__ and __len__), return as-is
+    if hasattr(calib_data, "__iter__") and hasattr(calib_data, "__len__"):
+        return calib_data
+
+    # If it's a list, return as-is
+    if isinstance(calib_data, list):
+        return calib_data
+
+    # If it's a tensor, wrap in a list
+    if isinstance(calib_data, torch.Tensor):
+        return [calib_data]
+
+    return calib_data
+
+
+def safe_model_snapshot(model: nn.Module, adapter: Any) -> dict:
+    """
+    Create a safe snapshot of model state for potential rollback.
+
+    Args:
+        model: Model to snapshot
+        adapter: Model adapter for model-specific operations
+
+    Returns:
+        Model snapshot dictionary
+    """
+    try:
+        # Create deep copy of state dict for safety
+        state_dict = {}
+        for name, param in model.named_parameters():
+            if param.requires_grad:
+                state_dict[name] = param.data.clone()
+
+        # Store model configuration if available through adapter
+        model_config = {}
+        if hasattr(adapter, "get_config"):
+            try:
+                model_config = adapter.get_config(model)
+            except Exception:
+                pass
+
+        return {
+            "state_dict": state_dict,
+            "config": model_config,
+            "snapshot_time": torch.cuda.Event(enable_timing=True)
+            if torch.cuda.is_available()
+            else None,
+        }
+    except Exception as e:
+        warnings.warn(f"Failed to create model snapshot: {e}", stacklevel=2)
+        return {}
+
+
+def safe_model_restore(model: nn.Module, snapshot: dict, adapter: Any) -> bool:
+    """
+    Safely restore model from snapshot.
+
+    Args:
+        model: Model to restore
+        snapshot: Snapshot from safe_model_snapshot()
+        adapter: Model adapter for model-specific operations
+
+    Returns:
+        True if restore was successful
+    """
+    try:
+        if "state_dict" not in snapshot:
+            warnings.warn("Snapshot missing state_dict, cannot restore", stacklevel=2)
+            return False
+
+        # Restore parameters
+        state_dict = snapshot["state_dict"]
+        for name, param in model.named_parameters():
+            if name in state_dict:
+                param.data.copy_(state_dict[name])
+
+        return True
+    except Exception as e:
+        warnings.warn(f"Failed to restore model from snapshot: {e}", stacklevel=2)
+        return False
+
+
+def create_baseline_guard_policy(edit_type: str) -> dict:
+    """
+    Create baseline guard policy for edit type.
+
+    Args:
+        edit_type: Type of edit operation
+
+    Returns:
+        Guard policy configuration
+    """
+    baseline_policies = {
+        "quant": {
+            "spectral_monitoring": False,  # Quantization may change spectral properties
+            "rmt_detection": True,
+            "weight_change_threshold": 0.5,  # Higher tolerance for quantization
+            "activation_change_threshold": 0.3,
+            "max_spectral_norm_increase": 3.0,
+        },
+    }
+
+    return baseline_policies.get(
+        edit_type,
+        {
+            "spectral_monitoring": True,
+            "rmt_detection": True,
+            "weight_change_threshold": 0.1,
+            "activation_change_threshold": 0.1,
+            "max_spectral_norm_increase": 2.0,
+        },
+    )
+
+
+__all__ = [
+    "ExternalBackend",
+    "ExternalEditWrapper",
+    "check_dependencies",
+    "validate_edit_config",
+    "compute_edit_metrics",
+    "prepare_calibration_data",
+    "safe_model_snapshot",
+    "safe_model_restore",
+    "create_baseline_guard_policy",
+]