diff --git a/tests/test_strategies.py b/tests/test_strategies.py index c65a3f9..362eac1 100644 --- a/tests/test_strategies.py +++ b/tests/test_strategies.py @@ -2,8 +2,11 @@ from __future__ import annotations +from types import SimpleNamespace + import pytest import torch +import torch.nn as nn from obliteratus.strategies.base import AblationSpec from obliteratus.strategies.registry import STRATEGY_REGISTRY, get_strategy @@ -51,6 +54,61 @@ def handle(): return _make_dummy_handle() +# --------------------------------------------------------------------------- +# Architecture navigation +# --------------------------------------------------------------------------- + +class _Qwen35MoeLayer(nn.Module): + def __init__(self, *, linear_attention: bool = False): + super().__init__() + if linear_attention: + self.linear_attn = nn.Module() + else: + self.self_attn = nn.Module() + self.mlp = nn.Module() + + +class _Qwen35MoeModel(nn.Module): + def __init__(self): + super().__init__() + self.model = nn.Module() + self.model.layers = nn.ModuleList( + [_Qwen35MoeLayer(), _Qwen35MoeLayer(linear_attention=True)] + ) + + +def test_qwen35_moe_resolves_layer_attention_and_ffn_paths(): + from obliteratus.models.loader import ModelHandle + from obliteratus.strategies.utils import ( + get_attention_module, + get_ffn_module, + get_layer_modules, + ) + + model = _Qwen35MoeModel() + handle = ModelHandle( + model=model, + tokenizer=SimpleNamespace(pad_token="", eos_token=""), + config=SimpleNamespace( + model_type="qwen3_5_moe", + num_hidden_layers=2, + num_attention_heads=4, + hidden_size=8, + intermediate_size=32, + ), + model_name="Qwen/Qwen3.5-35B-A3B", + task="causal_lm", + ) + + layers = get_layer_modules(handle) + + assert layers is model.model.layers + assert get_attention_module(layers[0], handle.architecture) is layers[0].self_attn + assert get_attention_module(layers[1], handle.architecture) is layers[1].linear_attn + assert get_ffn_module(layers[0], handle.architecture) is layers[0].mlp + assert get_ffn_module(layers[1], handle.architecture) is layers[1].mlp + + # --------------------------------------------------------------------------- # Registry tests # ---------------------------------------------------------------------------