interp-engine 1.6.0__tar.gz → 1.7.0__tar.gz

{interp_engine-1.6.0 → interp_engine-1.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: interp-engine
-Version: 1.6.0
+Version: 1.7.0
 Summary: A fast, standardized, and easy to use interpretability engine.
 Project-URL: Homepage, https://github.com/decoderesearch/interp-engine
 Project-URL: Repository, https://github.com/decoderesearch/interp-engine

{interp_engine-1.6.0 → interp_engine-1.7.0}/interp_engine/facts.py

@@ -478,6 +478,7 @@ class GateUpLayout(StrEnum):
 
     #: ``[all_gate | all_up]``, two contiguous halves, gate first. Phi-3's ``gate_up_proj``, whose
     #: forward is ``gate, up = self.gate_up_proj(x).chunk(2, dim=-1); down(up * act(gate))``.
+    #: GLM-4's ``Glm4MLP`` performs the identical chunk.
     GATE_FIRST = "gate_first"
     #: ``[gate_0 up_0 | gate_1 up_1 | ...]``, the two branches interleaved per neuron. gpt-oss's
     #: MXFP4 expert weights, where the kernel reads ``[..., ::2]`` and ``[..., 1::2]``.
@@ -488,8 +489,10 @@ class GateUpLayout(StrEnum):
 #: packed. Absent means the refusal in ``ArchSpec.mlp_projection`` stands: a family may fuse and pack
 #: either way, and "probably gate first" is exactly the guess that produces silent garbage.
 #: Verified by the identity the branches exist to satisfy -- ``act(mlp_pre) * mlp_pre_linear`` is the
-#: down projection's input, which is captured independently as ``mlp_act`` (tests/test_fused_mlp.py).
+#: down projection's input, which is captured independently as ``mlp_act``
+#: (tests/test_mlp_internals.py).
 FUSED_GATE_UP_LAYOUTS: dict[str, GateUpLayout] = {
+    "Glm4ForCausalLM": GateUpLayout.GATE_FIRST,
     "Phi3ForCausalLM": GateUpLayout.GATE_FIRST,
 }
 

{interp_engine-1.6.0 → interp_engine-1.7.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "interp-engine"
-version = "1.6.0"
+version = "1.7.0"
 description = "A fast, standardized, and easy to use interpretability engine."
 readme = "README.md"
 license = "Apache-2.0"

{interp_engine-1.6.0 → interp_engine-1.7.0}/tests/synthetic_families.py

@@ -437,6 +437,44 @@ def shrunk_phi3(*, seed: int = 0) -> Any:
         return EagerModel("Phi3ForCausalLM", hf_model=model, tokenizer=NoTokenizer(), device=None, dtype="float32")
 
 
+#: Dims for :func:`shrunk_glm4`. Same subject as :data:`_PHI3_SHRUNK` -- the MLP's single
+#: ``gate_up_proj`` -- so the only constraint is an even ``intermediate_size``. GLM-4 rotates half
+#: its head dimension, so ``head_dim`` stays a multiple of four.
+_GLM4_SHRUNK: dict[str, Any] = {
+    "hidden_size": 32,
+    "intermediate_size": 64,
+    "num_hidden_layers": 2,
+    "num_attention_heads": 4,
+    "num_key_value_heads": 2,
+    "head_dim": 8,
+    "vocab_size": 128,
+    "max_position_embeddings": 64,
+    "pad_token_id": 0,
+    "eos_token_id": 1,
+}
+
+
+def shrunk_glm4(*, seed: int = 0) -> Any:
+    """A tiny real GLM-4 with float32 weights on CPU, as an ``EagerModel``.
+
+    The second family to fuse its two pre-activation projections into one ``gate_up_proj``, and the
+    reason the packing is a table rather than a default: ``Glm4MLP`` chunks gate-first exactly as
+    Phi-3 does, but that is a fact about the family, established by running it, not one the shared
+    attribute name implies.
+    """
+    from transformers import Glm4Config, Glm4ForCausalLM
+
+    from interp_engine import EagerModel
+
+    config = Glm4Config(**_GLM4_SHRUNK)
+    config.architectures = ["Glm4ForCausalLM"]
+    torch.manual_seed(seed)
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        model = Glm4ForCausalLM(config).to(torch.float32).eval()
+        return EagerModel("Glm4ForCausalLM", hf_model=model, tokenizer=NoTokenizer(), device=None, dtype="float32")
+
+
 def shrunk_lfm2_moe(*, seed: int = 0) -> Any:
     """A tiny LFM2-MoE with real float32 weights on CPU, as an ``EagerModel``.
 

{interp_engine-1.6.0 → interp_engine-1.7.0}/tests/test_mlp_internals.py

@@ -29,7 +29,7 @@ from types import SimpleNamespace
 import pytest
 import torch
 from harness import GEMMA_IT, GPT2, QWEN_THINKING, ModelSpec, load_model, require_hf_token
-from synthetic_families import shrunk_opt, shrunk_phi3
+from synthetic_families import shrunk_glm4, shrunk_opt, shrunk_phi3
 
 from interp_engine import run_with_cache
 from interp_engine.capture import AddressLike
@@ -176,8 +176,9 @@ def test_a_plain_mlp_refuses_the_multiplied_branch():
         model.resolve_point("mlp_pre_linear", 0)
 
 
-def test_a_fused_gate_up_is_sliced_into_its_two_branches():
-    """Phi-3's shape: one projection holding both branches, so neither is a module output.
+@pytest.mark.parametrize("build", [shrunk_phi3, shrunk_glm4], ids=["phi3", "glm4"])
+def test_a_fused_gate_up_is_sliced_into_its_two_branches(build):
+    """Phi-3's and GLM-4's shape: one projection holding both branches, so neither is a module output.
 
     Served rather than refused, because a dense MLP's neuron basis exists whether or not the
     checkpoint stores the two matrices concatenated -- and this is a read plus a last-axis slice, the
@@ -186,9 +187,10 @@ def test_a_fused_gate_up_is_sliced_into_its_two_branches():
     Checked by the identity the branches exist to satisfy: `act(mlp_pre) * mlp_pre_linear` is the
     down projection's input, which arrives independently as `mlp_act`. That is what says the halves
     were cut the right way round, and it is exact -- swapping them changes the answer on any
-    non-symmetric activation, which SiLU is.
+    non-symmetric activation, which SiLU is. Run per family rather than once, because the packing is
+    a property of the family: both of these chunk gate-first, and nothing but running them says so.
     """
-    model = shrunk_phi3()
+    model = build()
     ids = torch.arange(7).unsqueeze(0) % 128
     cache = run_with_cache(model, ids, [("mlp_pre", 1), ("mlp_pre_linear", 1), ("mlp_act", 1)])
     pre, linear, act = (cache.get(name, 1) for name in ("mlp_pre", "mlp_pre_linear", "mlp_act"))