mirror of
https://github.com/p-e-w/heretic.git
synced 2026-11-06 05:16:39 +00:00
Chris Leeandumran666
efb1be6e3f
* 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>
61 lines
1.9 KiB
Python
61 lines
1.9 KiB
Python
# 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()
|