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

invarlock/adapters/hf_llama.py

@@ -1,487 +0,0 @@
-"""
-HuggingFace LLaMA Model Adapter
-===============================
-
-ModelAdapter implementation for HuggingFace LLaMA architecture models.
-
-This adapter provides LLaMA-specific integration including:
-- Support for LLaMA, LLaMA-2, Code Llama, and other LLaMA variants
-- Proper handling of RMSNorm layers and SwiGLU activation
-- RoPE (Rotary Position Embedding) support
-- Group Query Attention (GQA) handling for LLaMA-2
-- Proper device-aware state serialization
-"""
-
-from typing import Any
-
-import torch
-import torch.nn as nn
-
-from invarlock.core.api import ModelAdapter
-from invarlock.core.error_utils import wrap_errors
-from invarlock.core.exceptions import AdapterError, DependencyError, ModelLoadError
-
-from .hf_mixin import HFAdapterMixin
-
-TensorType = torch.Tensor
-ModuleType = nn.Module
-
-
-class HF_LLaMA_Adapter(HFAdapterMixin, ModelAdapter):
-    """
-    HuggingFace-specific ModelAdapter implementation for LLaMA models.
-
-    Supports LLaMA, LLaMA-2, Code Llama, and other LLaMA variants with:
-    - Enhanced LLaMA model detection and validation
-    - Support for Group Query Attention (GQA) in LLaMA-2
-    - RMSNorm layer handling
-    - RoPE position embedding support
-    - Device-aware state serialization
-    """
-
-    name = "hf_llama"
-
-    def load_model(
-        self, model_id: str, device: str = "auto", **kwargs: Any
-    ) -> ModuleType | Any:
-        """
-        Load a HuggingFace LLaMA model.
-
-        Args:
-            model_id: Model identifier (e.g. "meta-llama/Llama-2-7b-hf")
-            device: Target device ("auto", "cuda", "mps", "cpu")
-
-        Returns:
-            Loaded LLaMA model
-        """
-        # Lazy import to map missing dependency
-        with wrap_errors(
-            DependencyError,
-            "E203",
-            "DEPENDENCY-MISSING: transformers",
-            lambda e: {"dependency": "transformers"},
-        ):
-            from transformers import AutoModelForCausalLM  # type: ignore
-
-        with wrap_errors(
-            ModelLoadError,
-            "E201",
-            "MODEL-LOAD-FAILED: transformers AutoModelForCausalLM",
-            lambda e: {"model_id": model_id},
-        ):
-            model = AutoModelForCausalLM.from_pretrained(model_id, **kwargs)
-
-        # Use safe device movement that respects quantization constraints
-        return self._safe_to_device(model, device)
-
-    def can_handle(self, model: ModuleType | Any) -> bool:
-        """
-        Check if this adapter can handle the given model.
-
-        Enhanced detection for HuggingFace LLaMA models with validation
-        of expected structure and configuration.
-
-        Args:
-            model: The model to check
-
-        Returns:
-            True if this is a HuggingFace LLaMA compatible model
-        """
-
-        # Helper to detect explicitly set attributes (avoid unittest.mock auto-creation)
-        def _has_set_attr(obj, name: str) -> bool:
-            # Only treat attributes as present if explicitly set to avoid Mock auto-creation
-            d = getattr(obj, "__dict__", None)
-            if isinstance(d, dict) and name in d:
-                return True
-            # For nn.Module, also consider registered submodules/params/buffers
-            if isinstance(obj, nn.Module):
-                if hasattr(obj, "_modules") and name in obj._modules:
-                    return True
-                if hasattr(obj, "_parameters") and name in obj._parameters:
-                    return True
-                if hasattr(obj, "_buffers") and name in obj._buffers:
-                    return True
-            return False
-
-        # Check for HuggingFace LLaMA class names
-        model_name = model.__class__.__name__
-        if model_name in ["LlamaModel", "LlamaForCausalLM"]:
-            # Verify it has HF config
-            if hasattr(model, "config") and hasattr(model.config, "model_type"):
-                return model.config.model_type == "llama"
-
-        # Early bare-structure acceptance (no wrapper), minimal checks for tests
-        if hasattr(model, "layers"):
-            layers_obj = model.layers
-            # Obtain first layer via index or iterator
-            first_layer = None
-            try:
-                if hasattr(layers_obj, "__len__") and len(layers_obj) > 0:
-                    first_layer = layers_obj[0]
-            except Exception:
-                first_layer = None
-            if first_layer is None:
-                try:
-                    first_layer = next(iter(layers_obj))
-                except Exception:
-                    first_layer = None
-            if first_layer is not None:
-                candidate_layer = first_layer
-                # Minimal structural check for bare models (satisfies test expectations)
-                if hasattr(candidate_layer, "self_attn") and hasattr(
-                    candidate_layer, "mlp"
-                ):
-                    return True
-
-        # Structural validation for LLaMA-like models
-        if hasattr(model, "config") and hasattr(model, "model"):
-            config = model.config
-            llama_model = model.model
-
-            # Check for LLaMA configuration attributes
-            if (
-                hasattr(config, "num_hidden_layers")
-                and hasattr(config, "num_attention_heads")
-                and hasattr(config, "hidden_size")
-                and hasattr(llama_model, "layers")
-            ):
-                # Validate LLaMA structure
-                try:
-                    layers = llama_model.layers
-                    layer = None
-                    # Length-based path with robust exception handling
-                    try:
-                        if hasattr(layers, "__len__") and len(layers) > 0:
-                            layer = layers[0]
-                    except Exception:
-                        layer = None
-                    # Iterator fallback
-                    if layer is None and hasattr(layers, "__iter__"):
-                        try:
-                            # Call mocked __iter__ directly to support unittest.mock patterns
-                            layer = next(layers.__iter__())
-                        except (StopIteration, TypeError, AttributeError):
-                            return False
-                    if layer is None:
-                        return False
-
-                    # Check for LLaMA layer structure (strict: only count explicitly set attributes)
-                    if (
-                        hasattr(layer, "self_attn")
-                        and hasattr(layer, "mlp")
-                        and _has_set_attr(layer.self_attn, "q_proj")
-                        and _has_set_attr(layer.self_attn, "k_proj")
-                        and _has_set_attr(layer.self_attn, "v_proj")
-                        and _has_set_attr(layer.self_attn, "o_proj")
-                        and _has_set_attr(layer.mlp, "gate_proj")
-                        and _has_set_attr(layer.mlp, "up_proj")
-                        and _has_set_attr(layer.mlp, "down_proj")
-                    ):
-                        # Check for RMSNorm (characteristic of LLaMA)
-                        if _has_set_attr(layer, "input_layernorm") and _has_set_attr(
-                            layer, "post_attention_layernorm"
-                        ):
-                            return True
-                        else:
-                            return False
-                    else:
-                        return False
-
-                except (AttributeError, TypeError):
-                    return False
-
-        # Check for bare LLaMA model structure (less common but possible)
-        # Accept list/tuple/ModuleList and iterator-only mocks
-        if hasattr(model, "layers") and hasattr(model, "config"):
-            try:
-                layers = model.layers
-                first_layer = None
-                # Length-based access
-                try:
-                    if hasattr(layers, "__len__") and len(layers) > 0:
-                        first_layer = layers[0]
-                except Exception:
-                    first_layer = None
-                # Iterator-based access
-                if first_layer is None and hasattr(layers, "__iter__"):
-                    try:
-                        # Call __iter__ directly to support unittest.mock patterns
-                        first_layer = (
-                            next(layers.__iter__())
-                            if hasattr(layers, "__iter__")
-                            else next(iter(layers))
-                        )
-                    except Exception:
-                        first_layer = None
-                if first_layer is not None:
-                    candidate_layer = first_layer
-                    if (
-                        hasattr(candidate_layer, "self_attn")
-                        and hasattr(candidate_layer, "mlp")
-                        and hasattr(candidate_layer.self_attn, "q_proj")
-                        and hasattr(candidate_layer.mlp, "gate_proj")
-                    ):
-                        return True
-            except Exception:
-                pass
-
-        return False
-
-    def describe(self, model: ModuleType | Any) -> dict[str, Any]:
-        """
-        Get structural description of the HuggingFace LLaMA model.
-
-        Returns the required format for validation gates:
-        - n_layer: int
-        - heads_per_layer: List[int]
-        - mlp_dims: List[int]
-        - tying: Dict[str, str] (weight tying map)
-
-        Args:
-            model: The HuggingFace LLaMA model to describe
-
-        Returns:
-            Dictionary with model structure info in required format
-        """
-        # Determine model structure
-        if hasattr(model, "model"):
-            # LlamaForCausalLM structure
-            llama_model = model.model
-            layers = llama_model.layers
-            config = model.config
-        elif hasattr(model, "layers"):
-            # Direct LlamaModel structure
-            layers = model.layers
-            config = model.config
-            llama_model = model
-        else:
-            raise AdapterError(
-                code="E202",
-                message=(
-                    "ADAPTER-STRUCTURE-INVALID: unrecognized HuggingFace LLaMA model structure"
-                ),
-                details={"model_class": model.__class__.__name__},
-            )
-
-        # Extract basic configuration
-        # Robust layer count with Mock/iterator support; allow empty layers
-        try:
-            n_layers = len(layers)
-        except Exception:
-            try:
-                # Fallback: count via iteration
-                n_layers = sum(1 for _ in iter(layers))
-            except Exception as err:
-                raise AdapterError(
-                    code="E202",
-                    message=(
-                        "ADAPTER-STRUCTURE-INVALID: unrecognized HuggingFace LLaMA model structure"
-                    ),
-                    details={"error": str(err)},
-                ) from err
-        n_heads = getattr(config, "num_attention_heads", None)
-        hidden_size = getattr(config, "hidden_size", None)
-        vocab_size = getattr(config, "vocab_size", None)
-
-        # LLaMA-2 specific: Group Query Attention support
-        num_key_value_heads = getattr(config, "num_key_value_heads", n_heads)
-
-        if n_heads is None or hidden_size is None:
-            raise AdapterError(
-                code="E202",
-                message=(
-                    "ADAPTER-STRUCTURE-INVALID: missing num_attention_heads or hidden_size"
-                ),
-                details={"model_class": model.__class__.__name__},
-            )
-
-        # Get device info
-        try:
-            device = next(model.parameters()).device
-        except StopIteration:
-            device = torch.device("cpu")
-
-        # Calculate total parameters
-        total_params = sum(p.numel() for p in model.parameters())
-
-        # Get MLP dimensions for each layer
-        mlp_dims = []
-        heads_per_layer = []
-
-        for layer_idx in range(n_layers):
-            layer = layers[layer_idx]
-
-            # For LLaMA, all layers have the same head count
-            heads_per_layer.append(n_heads)
-
-            # Get MLP intermediate dimension (gate_proj/up_proj output size)
-            if hasattr(layer.mlp.gate_proj, "weight"):
-                # Linear layer: (out_features, in_features)
-                mlp_dim = layer.mlp.gate_proj.weight.shape[0]
-            else:
-                # Fallback to config
-                mlp_dim = getattr(config, "intermediate_size", hidden_size * 4)
-
-            mlp_dims.append(mlp_dim)
-
-        # Detect weight tying (lm_head ↔ embed_tokens)
-        tying_map = {}
-        if hasattr(model, "lm_head") and hasattr(llama_model, "embed_tokens"):
-            # Check if the weights are the same tensor (tied)
-            if model.lm_head.weight is llama_model.embed_tokens.weight:
-                tying_map["lm_head.weight"] = "model.embed_tokens.weight"
-
-        # Build the required description format
-        description = {
-            # Required fields for validation gates
-            "n_layer": n_layers,
-            "heads_per_layer": heads_per_layer,
-            "mlp_dims": mlp_dims,
-            "tying": tying_map,
-            # Additional useful information
-            "model_type": "llama",
-            "model_class": model.__class__.__name__,
-            "n_heads": n_heads,
-            "num_key_value_heads": num_key_value_heads,  # GQA support
-            "hidden_size": hidden_size,
-            "vocab_size": vocab_size,
-            "total_params": total_params,
-            "device": str(device),
-            # HuggingFace specific info
-            "hf_model_type": getattr(config, "model_type", "llama"),
-            "hf_config_class": config.__class__.__name__
-            if hasattr(config, "__class__")
-            else "unknown",
-            # LLaMA specific architecture details
-            "architecture": {
-                "has_lm_head": hasattr(model, "lm_head"),
-                "has_model_wrapper": hasattr(model, "model"),
-                "layer_norm_type": "rms",  # LLaMA uses RMSNorm
-                "activation": "silu",  # LLaMA uses SwiGLU (SiLU activation)
-                "positional_encoding": "rope",  # LLaMA uses RoPE
-                "use_bias": getattr(
-                    config, "use_bias", False
-                ),  # LLaMA typically no bias
-                "rope_theta": getattr(config, "rope_theta", 10000.0),
-                "max_position_embeddings": getattr(
-                    config, "max_position_embeddings", 2048
-                ),
-                "is_gqa": num_key_value_heads != n_heads,  # Group Query Attention
-                "gqa_ratio": n_heads // num_key_value_heads
-                if num_key_value_heads != n_heads
-                else 1,
-                "pretraining_tp": getattr(
-                    config, "pretraining_tp", 1
-                ),  # Tensor parallelism
-                "rms_norm_eps": getattr(config, "rms_norm_eps", 1e-6),
-            },
-        }
-
-        return description
-
-    def _extract_weight_tying_info(self, model: ModuleType | Any) -> dict[str, str]:
-        """
-        Extract weight tying relationships from the model.
-
-        Args:
-            model: The model to analyze
-
-        Returns:
-            Dictionary mapping tied parameter names to their source parameter names
-        """
-        tying_info = {}
-
-        # Check for lm_head ↔ embed_tokens tying (common in LLaMA)
-        if hasattr(model, "lm_head") and hasattr(model, "model"):
-            if hasattr(model.model, "embed_tokens"):
-                if model.lm_head.weight is model.model.embed_tokens.weight:
-                    tying_info["lm_head.weight"] = "model.embed_tokens.weight"
-
-        return tying_info
-
-    def _restore_weight_tying(
-        self, model: ModuleType | Any, tied_param: str, source_param: str
-    ) -> None:
-        """
-        Restore a weight tying relationship between parameters.
-
-        Args:
-            model: The model to modify
-            tied_param: Name of the parameter that should be tied
-            source_param: Name of the source parameter to tie to
-        """
-        # This is a placeholder for weight tying restoration logic
-        print(
-            f"Warning: Weight tying relationship {tied_param} -> {source_param} may have been broken during restore"
-        )
-
-    def get_layer_modules(
-        self, model: ModuleType | Any, layer_idx: int
-    ) -> dict[str, ModuleType | Any]:
-        """
-        Get the modules for a specific layer (utility method).
-
-        Args:
-            model: The HuggingFace LLaMA model
-            layer_idx: Index of the layer to get modules for
-
-        Returns:
-            Dictionary mapping module names to modules
-        """
-        if hasattr(model, "model"):
-            layer = model.model.layers[layer_idx]
-        else:
-            layer = model.layers[layer_idx]
-
-        modules = {
-            "self_attn.q_proj": layer.self_attn.q_proj,  # Query projection
-            "self_attn.k_proj": layer.self_attn.k_proj,  # Key projection
-            "self_attn.v_proj": layer.self_attn.v_proj,  # Value projection
-            "self_attn.o_proj": layer.self_attn.o_proj,  # Output projection
-            "mlp.gate_proj": layer.mlp.gate_proj,  # Gate projection (SwiGLU)
-            "mlp.up_proj": layer.mlp.up_proj,  # Up projection (SwiGLU)
-            "mlp.down_proj": layer.mlp.down_proj,  # Down projection
-            "input_layernorm": layer.input_layernorm,  # RMSNorm before attention
-            "post_attention_layernorm": layer.post_attention_layernorm,  # RMSNorm before MLP
-        }
-
-        return modules
-
-    def get_attention_info(self, model: ModuleType | Any) -> dict[str, Any]:
-        """
-        Get attention-specific information for LLaMA models.
-
-        Args:
-            model: The HuggingFace LLaMA model
-
-        Returns:
-            Dictionary with attention configuration details
-        """
-        config = model.config
-
-        def _safe_int(val):
-            return val if isinstance(val, int) else None
-
-        num_heads = _safe_int(getattr(config, "num_attention_heads", None))
-        hidden_size = _safe_int(getattr(config, "hidden_size", None))
-        num_key_value_heads = (
-            _safe_int(getattr(config, "num_key_value_heads", None)) or num_heads
-        )
-
-        head_dim = None
-        if isinstance(hidden_size, int) and isinstance(num_heads, int) and num_heads:
-            head_dim = hidden_size // num_heads
-
-        return {
-            "num_attention_heads": num_heads,
-            "num_key_value_heads": num_key_value_heads,
-            "head_dim": head_dim,
-            "is_group_query_attention": num_key_value_heads != num_heads,
-            "gqa_groups": num_heads // num_key_value_heads
-            if num_key_value_heads != num_heads
-            else 1,
-            "rope_theta": getattr(config, "rope_theta", 10000.0),
-            "max_position_embeddings": getattr(config, "max_position_embeddings", 2048),
-            "attention_dropout": getattr(config, "attention_dropout", 0.0),
-        }