Add EXAONE layer/attention/MLP name aliases; attn/trust flags

analyze_ov_circuits.py

@@ -93,13 +93,23 @@ def run(model_id: str, output: Path, token: str | None = None, trust_remote_code
     for layer_idx, layer in enumerate(layers):
         attn = getattr(layer, "self_attn", None)
         if attn is None:
-            print(f"[ov] layer {layer_idx}: no self_attn, skipping")
+            attn = getattr(layer, "attn", None)
+        if attn is None:
+            print(f"[ov] layer {layer_idx}: no self_attn/attn, skipping")
             continue
 
-        q_proj = attn.q_proj
-        k_proj = attn.k_proj
-        v_proj = attn.v_proj
-        o_proj = attn.o_proj
+        # EXAONE nests attention under attn.attention.
+        if hasattr(attn, "attention"):
+            attn = attn.attention
+
+        q_proj = getattr(attn, "q_proj", None)
+        k_proj = getattr(attn, "k_proj", None)
+        v_proj = getattr(attn, "v_proj", None)
+        o_proj = getattr(attn, "o_proj", None) or getattr(attn, "out_proj", None)
+
+        if q_proj is None or k_proj is None or v_proj is None or o_proj is None:
+            print(f"[ov] layer {layer_idx}: missing projections, skipping")
+            continue
 
         n_heads = getattr(attn, "num_heads", None)
         n_kv_heads = getattr(attn, "num_key_value_heads", None)

qwip_atlas/layers.py

@@ -20,22 +20,55 @@ def parse_layer_spec(spec: str) -> list[int]:
 
 
 def layers_container(model: Any):
-    """Find the module that owns the decoder `layers` ModuleList."""
+    """Find the module that owns the decoder layers ModuleList.
+
+    Handles both standard `.layers` (Llama-style) and GPT-NeoX-style `.h`.
+    """
     import torch.nn as nn
 
+    # Direct known aliases first.
+    transformer = getattr(model, "transformer", None)
+    if transformer is not None:
+        h = getattr(transformer, "h", None)
+        if isinstance(h, nn.ModuleList) and len(h) > 0:
+            return transformer
+        layers = getattr(transformer, "layers", None)
+        if isinstance(layers, nn.ModuleList) and len(layers) > 0:
+            return transformer
+
+    model_module = getattr(model, "model", None)
+    if model_module is not None:
+        h = getattr(model_module, "h", None)
+        if isinstance(h, nn.ModuleList) and len(h) > 0:
+            return model_module
+        layers = getattr(model_module, "layers", None)
+        if isinstance(layers, nn.ModuleList) and len(layers) > 0:
+            return model_module
+
+    # BFS fallback.
     queue = deque([model])
+    best = None
     while queue:
         module = queue.popleft()
-        layers = getattr(module, "layers", None)
-        if isinstance(layers, nn.ModuleList) and len(layers) > 0:
-            return module
+        for attr_name in ("layers", "h"):
+            layers = getattr(module, attr_name, None)
+            if isinstance(layers, nn.ModuleList) and len(layers) > 0:
+                # Prefer the deepest container (closest to actual decoder layers).
+                best = module
         for _, child in module.named_children():
             queue.append(child)
+    if best is not None:
+        return best
     raise RuntimeError(f"Cannot find decoder layers ModuleList on {type(model).__name__}")
 
 
 def resolve_layers(model: Any) -> list[Any]:
-    return list(layers_container(model).layers)
+    container = layers_container(model)
+    for attr_name in ("h", "layers"):
+        layers = getattr(container, attr_name, None)
+        if layers is not None:
+            return list(layers)
+    raise RuntimeError(f"Cannot enumerate layers from {type(container).__name__}")
 
 
 def inspect_layer(layer_mod: Any, text_cfg: Any) -> dict[str, Any]:
@@ -65,9 +98,27 @@ def inspect_layer(layer_mod: Any, text_cfg: Any) -> dict[str, Any]:
     mlp = getattr(layer_mod, "mlp", None) or getattr(layer_mod, "feed_forward", None)
     if mlp is not None:
         info["mlp"]["module"] = mlp
-        info["mlp"]["down_proj"] = getattr(mlp, "down_proj", None) or getattr(mlp, "wo", None)
-        info["mlp"]["gate_proj"] = getattr(mlp, "gate_proj", None) or getattr(mlp, "w1", None)
-        info["mlp"]["up_proj"] = getattr(mlp, "up_proj", None) or getattr(mlp, "w3", None)
+        # Llama/Gemma/Phi style
+        down_proj = getattr(mlp, "down_proj", None)
+        gate_proj = getattr(mlp, "gate_proj", None)
+        up_proj = getattr(mlp, "up_proj", None)
+        # GPT-NeoX / EXAONE style
+        if down_proj is None:
+            down_proj = getattr(mlp, "c_proj", None)
+        if gate_proj is None:
+            gate_proj = getattr(mlp, "c_fc_0", None) or getattr(mlp, "c_fc", None)
+        if up_proj is None:
+            up_proj = getattr(mlp, "c_fc_1", None)
+        # Mixtral/Mistral older names
+        if down_proj is None:
+            down_proj = getattr(mlp, "wo", None)
+        if gate_proj is None:
+            gate_proj = getattr(mlp, "w1", None)
+        if up_proj is None:
+            up_proj = getattr(mlp, "w3", None)
+        info["mlp"]["down_proj"] = down_proj
+        info["mlp"]["gate_proj"] = gate_proj
+        info["mlp"]["up_proj"] = up_proj
         info["mlp"]["act_fn"] = getattr(mlp, "act_fn", None)
         down_proj = info["mlp"]["down_proj"]
         if down_proj is not None and hasattr(down_proj, "in_features"):
@@ -81,6 +132,10 @@ def inspect_layer(layer_mod: Any, text_cfg: Any) -> dict[str, Any]:
     if attn is not None:
         info["attn"]["module"] = attn
         info["attn"]["class_name"] = type(attn).__name__
+        # EXAONE nests attention under attn.attention; most models use attn directly.
+        if hasattr(attn, "attention"):
+            attn = attn.attention
+
         info["attn"]["q_proj"] = getattr(attn, "q_proj", None)
         info["attn"]["k_proj"] = getattr(attn, "k_proj", None)
         info["attn"]["v_proj"] = getattr(attn, "v_proj", None)

qwip_atlas/sub_zero_surgery.py

@@ -326,36 +326,64 @@ def _model_device(model: torch.nn.Module) -> torch.device:
 
 
 def _resolve_layers(model: torch.nn.Module) -> List[torch.nn.Module]:
-    """Find the layers ModuleList in a model."""
+    """Find the decoder layers in a model. Handles Llama-style .layers and
+    GPT-NeoX/EXAONE-style transformer.h."""
     from collections import deque
+
+    transformer = getattr(model, "transformer", None)
+    if transformer is not None:
+        for attr in ("h", "layers"):
+            layers = getattr(transformer, attr, None)
+            if isinstance(layers, torch.nn.ModuleList) and len(layers) > 0:
+                return list(layers)
+
+    model_module = getattr(model, "model", None)
+    if model_module is not None:
+        for attr in ("h", "layers"):
+            layers = getattr(model_module, attr, None)
+            if isinstance(layers, torch.nn.ModuleList) and len(layers) > 0:
+                return list(layers)
+
     queue = deque([model])
     while queue:
         m = queue.popleft()
-        layers = getattr(m, "layers", None)
-        if isinstance(layers, torch.nn.ModuleList) and len(layers) > 0:
-            return list(layers)
+        for attr in ("h", "layers"):
+            layers = getattr(m, attr, None)
+            if isinstance(layers, torch.nn.ModuleList) and len(layers) > 0:
+                return list(layers)
         for _, child in m.named_children():
             queue.append(child)
-    raise RuntimeError(f"Cannot find layers ModuleList on {type(model).__name__}")
+    raise RuntimeError(f"Cannot find decoder layers on {type(model).__name__}")
 
 
 def _get_projection_map(layer: torch.nn.Module) -> Dict[str, torch.nn.Module]:
     """Get projection modules for a layer."""
     result = {}
-    # MLP projections
+    # MLP projections (Llama and GPT-NeoX/EXAONE aliases)
     mlp = getattr(layer, "mlp", None)
     if mlp is not None:
-        for name in ["gate_proj", "up_proj", "down_proj"]:
-            mod = getattr(mlp, name, None)
-            if mod is not None and hasattr(mod, "weight"):
-                result[name] = mod
+        aliases = {
+            "gate_proj": ["gate_proj", "c_fc_0", "c_fc"],
+            "up_proj":   ["up_proj", "c_fc_1"],
+            "down_proj": ["down_proj", "c_proj", "wo"],
+        }
+        for canon, names in aliases.items():
+            for name in names:
+                mod = getattr(mlp, name, None)
+                if mod is not None and hasattr(mod, "weight"):
+                    result[canon] = mod
+                    break
     # Attention projections
     attn = getattr(layer, "self_attn", None) or getattr(layer, "attention", None) or getattr(layer, "attn", None)
+    if attn is not None and hasattr(attn, "attention"):
+        attn = attn.attention
     if attn is not None:
-        for name in ["q_proj", "k_proj", "v_proj", "o_proj"]:
+        for name in ["q_proj", "k_proj", "v_proj", "o_proj", "out_proj"]:
             mod = getattr(attn, name, None)
             if mod is not None and hasattr(mod, "weight"):
-                result[name] = mod
+                # Normalize out_proj to o_proj for downstream naming.
+                key = "o_proj" if name == "out_proj" else name
+                result[key] = mod
     return result