fix: support flattened hyper-connection residual streams (#525)

* fix: support flattened hyper-connection residual streams

Qwen4-Exp flattens its hc_count residual streams into the component
dimension, so its hidden states have shape (prompt, position, streams *
hidden_size) instead of the 4D layout handled by #519. Collapse the flattened
streams to their mean before stacking residuals.

* fix: collect Qwen4-Exp shared experts as ablation targets

Qwen4-Exp uses the singular mlp.shared_expert spelling, while DeepSeek and
GLM use the plural mlp.shared_experts spelling. Collect both.

Co-authored-by: umran666 <shaikumran666@gmail.com>

* test: update qwen3.5-moe hashes for shared experts

Collecting mlp.shared_expert.down_proj adds it to the abliterable modules
for Qwen3.5-MoE, so the merged model changes. Add the new Linux and
Python 3.10 CI hashes.

---------

Co-authored-by: umran666 <shaikumran666@gmail.com>

src/heretic/model.py

@@ -44,6 +44,32 @@ def get_model_class(
         return AutoModelForCausalLM
 
 
+def _collapse_flattened_residual_streams(
+    hidden_states: tuple[FloatTensor, ...],
+    config: Any,
+) -> tuple[FloatTensor, ...]:
+    """
+    Average flattened multi-stream hidden states to hidden_size.
+
+    Some models (e.g. Qwen4-Exp) flatten their residual streams into the
+    component dimension instead of exposing a separate stream dimension, so
+    their hidden states have shape (prompt, position, streams * hidden_size).
+    The final entry is the merged model output (hidden_size) and is unchanged.
+    """
+    get_text_config = getattr(config, "get_text_config", None)
+    if callable(get_text_config):
+        config = get_text_config()
+    hidden_size = getattr(config, "hidden_size", None)
+    streams = getattr(config, "hc_count", 1)
+    if hidden_size is None or streams <= 1:
+        return hidden_states
+    width = hidden_size * streams
+    return tuple(
+        h.unflatten(-1, (streams, hidden_size)).mean(-2) if h.shape[-1] == width else h
+        for h in hidden_states
+    )
+
+
 # Some chat templates (e.g. gpt-oss) insert the current date into the prompt,
 # which makes the residuals, and therefore the modified model, depend on the day
 # Heretic is run. Rendering templates with a fixed date keeps results reproducible.
@@ -454,6 +480,10 @@ class Model:
         with suppress(Exception):
             try_add("mlp.down_proj", layer.mlp.shared_experts.down_proj)  # ty: ignore[unresolved-attribute]
 
+        # Shared expert of MoE models (e.g. Qwen4-Exp).
+        with suppress(Exception):
+            try_add("mlp.down_proj", layer.mlp.shared_expert.down_proj)  # ty: ignore[unresolved-attribute]
+
         # Phi-3.5-MoE (and possibly others).
         with suppress(Exception):
             for expert in layer.block_sparse_moe.experts:  # ty:ignore[not-iterable, unresolved-attribute]
@@ -654,6 +684,11 @@ class Model:
             for layer_hidden_states in hidden_states
         )
 
+        hidden_states = _collapse_flattened_residual_streams(
+            cast(tuple[FloatTensor, ...], hidden_states),
+            self.model.config,
+        )
+
         # The returned tensor has shape (prompt, layer, component).
         residuals = torch.stack(
             # layer_hidden_states has shape (prompt, position, component),

tests/qwen3.5-moe/SHA256SUMS.ci

@@ -1,7 +1,7 @@
 a4aee8afcf2e0711942cf848899be66016f8d14a889ff9ede07bca099c28f715 *chat_template.jinja
 4b501bef0727793de90d446fcfdc7757a9f8d19efc42ceaa3d26e90db4f876d5 *config.json
 532190619b0b243e0558a846257cc62883c26df812d456f8133e7b1043c69ad8 *generation_config.json
-aee5ba8012fe581a2905a72e278d80b6e65cc1ac08dfb328772c18825724deb3 *model.safetensors
+d72dad62b90e40eb726baf0f2f1867db1b8289c183a60b75cd147e68e674d7b1 *model.safetensors
 3814737f2d4bfdd153a82364b3850f1be5dbc6f074a55c8d5fbf4d40127c3e87 *processor_config.json
 a5cd9732badce41de57e6efce8302930ded1c1188c5f81feb2bd6c24c4a1941f *tokenizer.json
 5e55dc6b9d9d28d49b6b2b80d2a2f2558a0b6e0e7b194bbca8664fc269ab3ebc *tokenizer_config.json

tests/qwen3.5-moe/SHA256SUMS.ci2

@@ -1,7 +1,7 @@
 a4aee8afcf2e0711942cf848899be66016f8d14a889ff9ede07bca099c28f715 *chat_template.jinja
 4b501bef0727793de90d446fcfdc7757a9f8d19efc42ceaa3d26e90db4f876d5 *config.json
 532190619b0b243e0558a846257cc62883c26df812d456f8133e7b1043c69ad8 *generation_config.json
-845fda6e6bfa82f7e28930606816f7f0962d053996a912e8d18d5b95ce5d154e *model.safetensors
+327e02f00fccfb8dd580c1b3ca1620438dd6d8c8224b3b38b70cabf893a8539c *model.safetensors
 3814737f2d4bfdd153a82364b3850f1be5dbc6f074a55c8d5fbf4d40127c3e87 *processor_config.json
 a5cd9732badce41de57e6efce8302930ded1c1188c5f81feb2bd6c24c4a1941f *tokenizer.json
 5e55dc6b9d9d28d49b6b2b80d2a2f2558a0b6e0e7b194bbca8664fc269ab3ebc *tokenizer_config.json

tests/qwen3.5-moe/SHA256SUMS.linux

@@ -1,7 +1,7 @@
 a4aee8afcf2e0711942cf848899be66016f8d14a889ff9ede07bca099c28f715 *chat_template.jinja
 4b501bef0727793de90d446fcfdc7757a9f8d19efc42ceaa3d26e90db4f876d5 *config.json
 532190619b0b243e0558a846257cc62883c26df812d456f8133e7b1043c69ad8 *generation_config.json
-4c29b8ce99c59c90eca59aa6453029bda87165efa57cabc2c668599c02afc31a *model.safetensors
+327e02f00fccfb8dd580c1b3ca1620438dd6d8c8224b3b38b70cabf893a8539c *model.safetensors
 3814737f2d4bfdd153a82364b3850f1be5dbc6f074a55c8d5fbf4d40127c3e87 *processor_config.json
 a5cd9732badce41de57e6efce8302930ded1c1188c5f81feb2bd6c24c4a1941f *tokenizer.json
 5e55dc6b9d9d28d49b6b2b80d2a2f2558a0b6e0e7b194bbca8664fc269ab3ebc *tokenizer_config.json

tests/test_residual_streams.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: AGPL-3.0-or-later
+# Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
+
+import unittest
+from types import SimpleNamespace
+from typing import cast
+
+import torch
+from torch import FloatTensor
+
+from heretic.model import _collapse_flattened_residual_streams
+
+
+class FlattenedResidualStreamTests(unittest.TestCase):
+    def test_collapses_flattened_streams(self) -> None:
+        torch.manual_seed(0)
+        hidden_states = cast(
+            tuple[FloatTensor, ...],
+            (
+                torch.randn(2, 5, 12),
+                torch.randn(2, 5, 12),
+                torch.randn(2, 5, 4),
+            ),
+        )
+        config = SimpleNamespace(hidden_size=4, hc_count=3)
+
+        collapsed = _collapse_flattened_residual_streams(hidden_states, config)
+
+        self.assertEqual([h.shape[-1] for h in collapsed], [4, 4, 4])
+        self.assertTrue(
+            torch.allclose(
+                collapsed[0],
+                hidden_states[0].unflatten(-1, (3, 4)).mean(-2),
+            )
+        )
+        self.assertIs(collapsed[2], hidden_states[2])
+
+    def test_uses_text_config(self) -> None:
+        hidden_states = cast(tuple[FloatTensor, ...], (torch.randn(1, 2, 16),))
+        text_config = SimpleNamespace(hidden_size=4, hc_count=4)
+
+        class Config:
+            def get_text_config(self) -> SimpleNamespace:
+                return text_config
+
+        collapsed = _collapse_flattened_residual_streams(hidden_states, Config())
+
+        self.assertEqual(collapsed[0].shape, (1, 2, 4))
+
+    def test_leaves_single_stream_models_unchanged(self) -> None:
+        hidden_states = cast(tuple[FloatTensor, ...], (torch.randn(1, 2, 4),))
+        config = SimpleNamespace(hidden_size=4, hc_count=1)
+
+        collapsed = _collapse_flattened_residual_streams(hidden_states, config)
+
+        self.assertIs(collapsed, hidden_states)
+
+
+if __name__ == "__main__":
+    unittest.main()