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

invarlock/__init__.py

@@ -6,13 +6,13 @@ Core runtime package — torch-independent utilities, configuration, and interfa
 
 This package provides the foundation for the InvarLock GuardChain without heavy dependencies.
 For torch-dependent functionality, see subpackages under `invarlock.*`:
-- `invarlock.adapters`: Model adapters (HF GPT-2/BERT/LLaMA, auto)
+- `invarlock.adapters`: Model adapters (HF causal/MLM/seq2seq + auto)
 - `invarlock.guards`: Safety mechanisms (invariants, spectral, RMT, variance)
 - `invarlock.edits`: Built-in quantization and edit interfaces
 - `invarlock.eval`: Metrics, guard-overhead checks, and certification
 """
 
-__version__ = "0.3.6"
+__version__ = "0.3.7"
 
 # Core exports - torch-independent
 from .config import CFG, Defaults, get_default_config

invarlock/adapters/__init__.py

@@ -29,13 +29,11 @@ from .capabilities import (
 )
 
 _LAZY_MAP = {
-    "HF_BERT_Adapter": ".hf_bert",
-    "HF_GPT2_Adapter": ".hf_gpt2",
-    "HF_LLaMA_Adapter": ".hf_llama",
-    "HF_T5_Adapter": ".hf_t5",
-    "HF_ORT_CausalLM_Adapter": ".hf_onnx",
-    "HF_Causal_Auto_Adapter": ".auto",
-    "HF_MLM_Auto_Adapter": ".auto",
+    "HF_Causal_Adapter": ".hf_causal",
+    "HF_MLM_Adapter": ".hf_mlm",
+    "HF_Seq2Seq_Adapter": ".hf_seq2seq",
+    "HF_Causal_ONNX_Adapter": ".hf_causal_onnx",
+    "HF_Auto_Adapter": ".auto",
 }
 
 
@@ -91,13 +89,11 @@ run_invarlock = _RemovedComponent("run_invarlock", "invarlock.cli.run")
 quick_prune_gpt2 = _RemovedComponent("quick_prune_gpt2")
 
 __all__ = [
-    "HF_GPT2_Adapter",
-    "HF_BERT_Adapter",
-    "HF_LLaMA_Adapter",
-    "HF_T5_Adapter",
-    "HF_ORT_CausalLM_Adapter",
-    "HF_Causal_Auto_Adapter",
-    "HF_MLM_Auto_Adapter",
+    "HF_Causal_Adapter",
+    "HF_MLM_Adapter",
+    "HF_Seq2Seq_Adapter",
+    "HF_Causal_ONNX_Adapter",
+    "HF_Auto_Adapter",
     "BaseAdapter",
     "AdapterConfig",
     "AdapterInterface",

invarlock/adapters/auto.py

@@ -110,21 +110,26 @@ class _DelegatingAdapter(ModelAdapter):
 
     def _load_adapter(self, adapter_name: str) -> ModelAdapter:
         """Load an adapter by name."""
-        if adapter_name == "hf_llama":
-            HF_LLaMA_Adapter = _importlib.import_module(
-                ".hf_llama", __package__
-            ).HF_LLaMA_Adapter
-            return HF_LLaMA_Adapter()
-        elif adapter_name == "hf_bert":
-            HF_BERT_Adapter = _importlib.import_module(
-                ".hf_bert", __package__
-            ).HF_BERT_Adapter
-            return HF_BERT_Adapter()
-        elif adapter_name == "hf_gpt2":
-            HF_GPT2_Adapter = _importlib.import_module(
-                ".hf_gpt2", __package__
-            ).HF_GPT2_Adapter
-            return HF_GPT2_Adapter()
+        if adapter_name == "hf_causal":
+            HF_Causal_Adapter = _importlib.import_module(
+                ".hf_causal", __package__
+            ).HF_Causal_Adapter
+            return HF_Causal_Adapter()
+        if adapter_name == "hf_mlm":
+            HF_MLM_Adapter = _importlib.import_module(
+                ".hf_mlm", __package__
+            ).HF_MLM_Adapter
+            return HF_MLM_Adapter()
+        if adapter_name == "hf_seq2seq":
+            HF_Seq2Seq_Adapter = _importlib.import_module(
+                ".hf_seq2seq", __package__
+            ).HF_Seq2Seq_Adapter
+            return HF_Seq2Seq_Adapter()
+        if adapter_name == "hf_causal_onnx":
+            HF_Causal_ONNX_Adapter = _importlib.import_module(
+                ".hf_causal_onnx", __package__
+            ).HF_Causal_ONNX_Adapter
+            return HF_Causal_ONNX_Adapter()
         elif adapter_name == "hf_bnb":
             HF_BNB_Adapter = _importlib.import_module(
                 "invarlock.plugins.hf_bnb_adapter"
@@ -141,11 +146,11 @@ class _DelegatingAdapter(ModelAdapter):
             ).HF_GPTQ_Adapter
             return HF_GPTQ_Adapter()
         else:
-            # Default to GPT2 adapter
-            HF_GPT2_Adapter = _importlib.import_module(
-                ".hf_gpt2", __package__
-            ).HF_GPT2_Adapter
-            return HF_GPT2_Adapter()
+            # Default to causal adapter
+            HF_Causal_Adapter = _importlib.import_module(
+                ".hf_causal", __package__
+            ).HF_Causal_Adapter
+            return HF_Causal_Adapter()
 
     def _ensure_delegate_from_id(self, model_id: str) -> ModelAdapter:
         if self._delegate is not None:
@@ -172,14 +177,16 @@ class _DelegatingAdapter(ModelAdapter):
             self._delegate = self._load_adapter(quant_adapter)
             return self._delegate
 
-        # Fall back to class name inspection
+        # Fall back to lightweight class-name inspection (no transformers import).
         cls_name = getattr(model, "__class__", type(model)).__name__.lower()
-        if any(k in cls_name for k in ["llama", "mistral", "qwen", "yi"]):
-            self._delegate = self._load_adapter("hf_llama")
-        elif any(k in cls_name for k in ["bert", "roberta", "albert", "deberta"]):
-            self._delegate = self._load_adapter("hf_bert")
+        if any(k in cls_name for k in ["bert", "roberta", "albert", "deberta"]):
+            self._delegate = self._load_adapter("hf_mlm")
         else:
-            self._delegate = self._load_adapter("hf_gpt2")
+            cfg = getattr(model, "config", None)
+            if getattr(cfg, "is_encoder_decoder", False):
+                self._delegate = self._load_adapter("hf_seq2seq")
+            else:
+                self._delegate = self._load_adapter("hf_causal")
         return self._delegate
 
     def can_handle(self, model: Any) -> bool:  # pragma: no cover - trivial
@@ -206,21 +213,9 @@ class _DelegatingAdapter(ModelAdapter):
         raise AttributeError(item)
 
 
-class HF_Causal_Auto_Adapter(_DelegatingAdapter):
-    name = "hf_causal_auto"
+class HF_Auto_Adapter(_DelegatingAdapter):
+    name = "hf_auto"
 
     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, **kwargs)
-
-
-class HF_MLM_Auto_Adapter(_DelegatingAdapter):
-    name = "hf_mlm_auto"
-
-    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, **kwargs)

invarlock/adapters/capabilities.py

@@ -359,7 +359,7 @@ def _detect_weight_tying(model: Any) -> dict[str, str]:
     tying: dict[str, str] = {}
 
     # Common weight tying patterns
-    # LLaMA/Mistral: lm_head.weight ↔ model.embed_tokens.weight
+    # Decoder embed_tokens style: lm_head.weight ↔ model.embed_tokens.weight
     if hasattr(model, "lm_head") and hasattr(model, "model"):
         inner = model.model
         if hasattr(inner, "embed_tokens"):
@@ -408,7 +408,7 @@ def _detect_primary_metric(model: Any) -> str:
             return "rouge"
         return "ppl_seq2seq"
 
-    # Decoder-only models (GPT-like, LLaMA-like)
+    # Decoder-only models (GPT-like, RoPE-style)
     return "ppl_causal"
 
 

invarlock/adapters/hf_causal.py

@@ -0,0 +1,418 @@
+"""
+HuggingFace causal LM adapter (decoder-only).
+=============================================
+
+Role-based adapter for HuggingFace decoder-only causal language models.
+
+This adapter intentionally avoids model-family naming. It selects a structural
+spec at runtime (dense FFN vs MoE vs GPT-2-like blocks) and exposes a stable
+`describe()` contract for InvarLock gates and reporting.
+"""
+
+from __future__ import annotations
+
+import os
+from types import SimpleNamespace
+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
+
+LIGHT_IMPORT = os.getenv("INVARLOCK_LIGHT_IMPORT", "").strip().lower() in {
+    "1",
+    "true",
+    "yes",
+}
+
+
+def _first_item(seq: Any) -> Any | None:
+    try:
+        if hasattr(seq, "__len__") and len(seq) > 0:  # type: ignore[arg-type]
+            return seq[0]  # type: ignore[index]
+    except Exception:
+        pass
+    try:
+        return next(iter(seq))
+    except Exception:
+        return None
+
+
+def _has_set_attr(obj: Any, name: str) -> bool:
+    d = getattr(obj, "__dict__", None)
+    if isinstance(d, dict) and name in d:
+        return True
+    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
+
+
+class _CausalSpec:
+    spec_name = "base"
+
+    def matches(self, model: Any, base: Any, layers: Any) -> bool:
+        raise NotImplementedError
+
+    def infer_mlp_dim(self, layer: Any, config: Any, hidden_size: int) -> int:
+        raise NotImplementedError
+
+    def layer_modules(self, model: Any, layer: Any) -> dict[str, Any]:
+        raise NotImplementedError
+
+    def tying_map(self, model: Any, base: Any) -> dict[str, str]:
+        return {}
+
+
+class _DenseDecoderSpec(_CausalSpec):
+    spec_name = "dense_decoder"
+
+    def matches(self, model: Any, base: Any, layers: Any) -> bool:
+        layer = _first_item(layers)
+        if layer is None:
+            return False
+        has_attn = (
+            hasattr(layer, "self_attn")
+            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")
+        )
+        has_mlp = (
+            hasattr(layer, "mlp")
+            and _has_set_attr(layer.mlp, "gate_proj")
+            and _has_set_attr(layer.mlp, "up_proj")
+            and _has_set_attr(layer.mlp, "down_proj")
+        )
+        has_norms = _has_set_attr(layer, "input_layernorm") and _has_set_attr(
+            layer, "post_attention_layernorm"
+        )
+        return bool(has_attn and has_mlp and has_norms)
+
+    def infer_mlp_dim(self, layer: Any, config: Any, hidden_size: int) -> int:
+        mlp_dim = int(getattr(config, "intermediate_size", hidden_size * 4) or 0)
+        try:
+            gate_proj = getattr(getattr(layer, "mlp", None), "gate_proj", None)
+            if gate_proj is not None and hasattr(gate_proj, "weight"):
+                mlp_dim = int(gate_proj.weight.shape[0])
+        except Exception:
+            pass
+        return int(mlp_dim)
+
+    def layer_modules(self, model: Any, layer: Any) -> dict[str, Any]:
+        mlp = layer.mlp
+        return {
+            "self_attn.q_proj": layer.self_attn.q_proj,
+            "self_attn.k_proj": layer.self_attn.k_proj,
+            "self_attn.v_proj": layer.self_attn.v_proj,
+            "self_attn.o_proj": layer.self_attn.o_proj,
+            "input_layernorm": layer.input_layernorm,
+            "post_attention_layernorm": layer.post_attention_layernorm,
+            "mlp.gate_proj": mlp.gate_proj,
+            "mlp.up_proj": mlp.up_proj,
+            "mlp.down_proj": mlp.down_proj,
+        }
+
+    def tying_map(self, model: Any, base: Any) -> dict[str, str]:
+        tying: dict[str, str] = {}
+        try:
+            if hasattr(model, "lm_head") and hasattr(base, "embed_tokens"):
+                if model.lm_head.weight is base.embed_tokens.weight:
+                    tying["lm_head.weight"] = "model.embed_tokens.weight"
+        except Exception:
+            pass
+        return tying
+
+
+class _MoEDecoderSpec(_CausalSpec):
+    spec_name = "moe_decoder"
+
+    def matches(self, model: Any, base: Any, layers: Any) -> bool:
+        layer = _first_item(layers)
+        if layer is None:
+            return False
+        has_attn = (
+            hasattr(layer, "self_attn")
+            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")
+        )
+        moe = getattr(layer, "block_sparse_moe", None)
+        experts = getattr(moe, "experts", None) if moe is not None else None
+        expert0 = _first_item(experts) if experts is not None else None
+        has_moe = bool(
+            expert0 is not None
+            and _has_set_attr(expert0, "w1")
+            and _has_set_attr(expert0, "w2")
+        )
+        has_norms = _has_set_attr(layer, "input_layernorm") and _has_set_attr(
+            layer, "post_attention_layernorm"
+        )
+        return bool(has_attn and has_moe and has_norms)
+
+    def infer_mlp_dim(self, layer: Any, config: Any, hidden_size: int) -> int:
+        mlp_dim = int(getattr(config, "intermediate_size", hidden_size * 4) or 0)
+        try:
+            moe = getattr(layer, "block_sparse_moe", None)
+            experts = getattr(moe, "experts", None) if moe is not None else None
+            expert0 = _first_item(experts) if experts is not None else None
+            if expert0 is not None:
+                w1 = getattr(expert0, "w1", None)
+                if w1 is not None and hasattr(w1, "weight"):
+                    mlp_dim = int(w1.weight.shape[0])
+        except Exception:
+            pass
+        return int(mlp_dim)
+
+    def layer_modules(self, model: Any, layer: Any) -> dict[str, Any]:
+        moe = layer.block_sparse_moe
+        expert0 = _first_item(moe.experts)
+        if expert0 is None:
+            raise AdapterError(
+                code="E202",
+                message="ADAPTER-STRUCTURE-INVALID: MoE layer missing experts",
+                details={"layer_class": layer.__class__.__name__},
+            )
+        return {
+            "self_attn.q_proj": layer.self_attn.q_proj,
+            "self_attn.k_proj": layer.self_attn.k_proj,
+            "self_attn.v_proj": layer.self_attn.v_proj,
+            "self_attn.o_proj": layer.self_attn.o_proj,
+            "input_layernorm": layer.input_layernorm,
+            "post_attention_layernorm": layer.post_attention_layernorm,
+            # Best-effort mapping to dense naming used elsewhere in the stack.
+            "mlp.gate_proj": expert0.w1,
+            "mlp.up_proj": getattr(expert0, "w3", expert0.w1),
+            "mlp.down_proj": expert0.w2,
+        }
+
+    def tying_map(self, model: Any, base: Any) -> dict[str, str]:
+        return _DenseDecoderSpec().tying_map(model, base)
+
+
+class _GPT2LikeDecoderSpec(_CausalSpec):
+    spec_name = "gpt2_like"
+
+    def matches(self, model: Any, base: Any, layers: Any) -> bool:
+        layer = _first_item(layers)
+        if layer is None:
+            return False
+        return bool(
+            hasattr(layer, "attn")
+            and hasattr(layer.attn, "c_proj")
+            and hasattr(layer, "mlp")
+            and hasattr(layer.mlp, "c_proj")
+        )
+
+    def infer_mlp_dim(self, layer: Any, config: Any, hidden_size: int) -> int:
+        try:
+            c_fc = getattr(getattr(layer, "mlp", None), "c_fc", None)
+            if c_fc is not None and hasattr(c_fc, "weight"):
+                # HF GPT-style uses Conv1D where nf is out_features.
+                if hasattr(c_fc, "nf"):
+                    return int(c_fc.nf)
+                return int(c_fc.weight.shape[0])
+        except Exception:
+            pass
+        return int(getattr(config, "n_inner", hidden_size * 4) or 0)
+
+    def layer_modules(self, model: Any, layer: Any) -> dict[str, Any]:
+        return {
+            "attn.c_attn": layer.attn.c_attn,
+            "attn.c_proj": layer.attn.c_proj,
+            "mlp.c_fc": layer.mlp.c_fc,
+            "mlp.c_proj": layer.mlp.c_proj,
+            "ln_1": layer.ln_1,
+            "ln_2": layer.ln_2,
+        }
+
+    def tying_map(self, model: Any, base: Any) -> dict[str, str]:
+        tying: dict[str, str] = {}
+        try:
+            if hasattr(model, "lm_head") and hasattr(base, "wte"):
+                if model.lm_head.weight is base.wte.weight:
+                    tying["lm_head.weight"] = "transformer.wte.weight"
+        except Exception:
+            pass
+        return tying
+
+
+_SPECS: list[_CausalSpec] = [
+    _MoEDecoderSpec(),
+    _DenseDecoderSpec(),
+    _GPT2LikeDecoderSpec(),
+]
+
+
+class HF_Causal_Adapter(HFAdapterMixin, ModelAdapter):
+    """Spec-driven adapter for decoder-only causal LMs."""
+
+    name = "hf_causal"
+
+    def load_model(
+        self, model_id: str, device: str = "auto", **kwargs: Any
+    ) -> ModuleType | Any:
+        try:
+            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)
+
+            return self._safe_to_device(model, device)
+        except DependencyError:
+            if LIGHT_IMPORT:
+                stub = SimpleNamespace(name="hf_causal_stub")
+                stub.to = lambda *_a, **_k: stub  # type: ignore[attr-defined]
+                return stub
+            raise
+
+    def _unwrap(self, model: Any) -> tuple[Any, Any, Any]:
+        config = getattr(model, "config", None)
+        if hasattr(model, "model") and hasattr(model.model, "layers"):
+            return model.model, model.model.layers, config
+        if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
+            return model.transformer, model.transformer.h, config
+        if hasattr(model, "layers"):
+            return model, model.layers, config
+        if hasattr(model, "h"):
+            return model, model.h, config
+        raise AdapterError(
+            code="E202",
+            message="ADAPTER-STRUCTURE-INVALID: unrecognized HF causal LM structure",
+            details={"model_class": model.__class__.__name__},
+        )
+
+    def _select_spec(self, model: Any, base: Any, layers: Any) -> _CausalSpec:
+        for spec in _SPECS:
+            try:
+                if spec.matches(model, base, layers):
+                    return spec
+            except Exception:
+                continue
+        return _DenseDecoderSpec()
+
+    def can_handle(self, model: ModuleType | Any) -> bool:
+        try:
+            base, layers, _cfg = self._unwrap(model)
+        except Exception:
+            return False
+        return any(spec.matches(model, base, layers) for spec in _SPECS)
+
+    def describe(self, model: ModuleType | Any) -> dict[str, Any]:
+        base, layers, config = self._unwrap(model)
+        if config is None:
+            raise AdapterError(
+                code="E202",
+                message="ADAPTER-STRUCTURE-INVALID: missing HuggingFace config on model",
+                details={"model_class": model.__class__.__name__},
+            )
+
+        try:
+            n_layers = len(layers)
+        except Exception:
+            n_layers = sum(1 for _ in iter(layers))
+
+        def _coerce_int(value: Any) -> int | None:
+            try:
+                if isinstance(value, bool):
+                    return None
+                if isinstance(value, int):
+                    return int(value)
+                if isinstance(value, float):
+                    return int(value)
+                if isinstance(value, str):
+                    stripped = value.strip()
+                    if stripped and stripped.isdigit():
+                        return int(stripped)
+            except Exception:
+                return None
+            return None
+
+        n_heads = _coerce_int(getattr(config, "num_attention_heads", None))
+        if n_heads is None:
+            n_heads = _coerce_int(getattr(config, "n_head", None))
+
+        hidden_size = _coerce_int(getattr(config, "hidden_size", None))
+        if hidden_size is None:
+            hidden_size = _coerce_int(getattr(config, "n_embd", None))
+
+        vocab_size = _coerce_int(getattr(config, "vocab_size", None))
+
+        if n_heads is None or hidden_size is None:
+            raise AdapterError(
+                code="E202",
+                message="ADAPTER-STRUCTURE-INVALID: missing head/hidden size metadata",
+                details={"model_class": model.__class__.__name__},
+            )
+
+        spec = self._select_spec(model, base, layers)
+
+        heads_per_layer = [int(n_heads)] * int(n_layers)
+        mlp_dims: list[int] = []
+        for idx in range(int(n_layers)):
+            layer = layers[idx]
+            mlp_dims.append(spec.infer_mlp_dim(layer, config, int(hidden_size)))
+
+        tying = spec.tying_map(model, base)
+
+        total_params = 0
+        try:
+            total_params = sum(p.numel() for p in model.parameters())
+        except Exception:
+            total_params = 0
+
+        try:
+            device = next(model.parameters()).device
+        except Exception:
+            device = torch.device("cpu")
+
+        return {
+            "n_layer": int(n_layers),
+            "heads_per_layer": heads_per_layer,
+            "mlp_dims": mlp_dims,
+            "tying": tying,
+            "model_type": str(getattr(config, "model_type", "") or "causal"),
+            "model_class": model.__class__.__name__,
+            "hf_model_type": str(getattr(config, "model_type", "") or ""),
+            "hf_config_class": config.__class__.__name__
+            if hasattr(config, "__class__")
+            else "unknown",
+            "n_heads": int(n_heads),
+            "hidden_size": int(hidden_size),
+            "vocab_size": int(vocab_size) if vocab_size is not None else None,
+            "total_params": int(total_params),
+            "device": str(device),
+            "spec": spec.spec_name,
+        }
+
+    def get_layer_modules(
+        self, model: ModuleType | Any, layer_idx: int
+    ) -> dict[str, Any]:
+        base, layers, _cfg = self._unwrap(model)
+        spec = self._select_spec(model, base, layers)
+        layer = layers[layer_idx]
+        return spec.layer_modules(model, layer)

invarlock/adapters/{hf_onnx.py → hf_causal_onnx.py}

@@ -22,7 +22,7 @@ from invarlock.core.error_utils import wrap_errors
 from invarlock.core.exceptions import DependencyError, ModelLoadError
 
 
-class HF_ORT_CausalLM_Adapter(ModelAdapter):
+class HF_Causal_ONNX_Adapter(ModelAdapter):
     """Optimum/ONNXRuntime causal LM adapter.
 
     Provides a lightweight bridge that loads an ORTModelForCausalLM and
@@ -31,7 +31,7 @@ class HF_ORT_CausalLM_Adapter(ModelAdapter):
     back to reload in the CLI runner.
     """
 
-    name = "hf_onnx"
+    name = "hf_causal_onnx"
 
     # --- Loading ---------------------------------------------------------
     def load_model(self, model_id: str, device: str = "cpu", **kwargs: Any):  # type: ignore[override]
@@ -109,4 +109,4 @@ class HF_ORT_CausalLM_Adapter(ModelAdapter):
         raise NotImplementedError("restore not supported for ONNXRuntime models")
 
 
-__all__ = ["HF_ORT_CausalLM_Adapter"]
+__all__ = ["HF_Causal_ONNX_Adapter"]

invarlock/adapters/hf_mixin.py

@@ -490,18 +490,39 @@ class HFAdapterMixin:
         """Return mapping of tied parameter names to source parameter names."""
 
         tying: dict[str, str] = {}
-        param_names = set(dict(model.named_parameters()).keys())
+        try:
+            named = model.named_parameters(remove_duplicate=False)  # type: ignore[call-arg]
+        except TypeError:  # pragma: no cover - torch version dependent
+            named = model.named_parameters()
+        params = dict(named)
+
+        def _is_tied(name_a: str, name_b: str) -> bool:
+            a = params.get(name_a)
+            b = params.get(name_b)
+            if a is None or b is None:
+                return False
+            try:
+                if a is b:
+                    return True
+                if hasattr(a, "data_ptr") and hasattr(b, "data_ptr"):
+                    return int(a.data_ptr()) == int(b.data_ptr())
+            except Exception:
+                return False
+            return False
 
-        if "lm_head.weight" in param_names and "transformer.wte.weight" in param_names:
+        if _is_tied("lm_head.weight", "transformer.wte.weight"):
             tying["lm_head.weight"] = "transformer.wte.weight"
 
+        if _is_tied("lm_head.weight", "model.embed_tokens.weight"):
+            tying["lm_head.weight"] = "model.embed_tokens.weight"
+
         decoder_name = "cls.predictions.decoder.weight"
-        if decoder_name in param_names:
+        if decoder_name in params:
             for candidate in (
                 "bert.embeddings.word_embeddings.weight",
                 "embeddings.word_embeddings.weight",
             ):
-                if candidate in param_names:
+                if _is_tied(decoder_name, candidate):
                     tying[decoder_name] = candidate
                     break