From eb2cff3a704bb57a9c83ee8b42e9ffb3c2eb0884 Mon Sep 17 00:00:00 2001 From: Pablo Cabriada Sierra Date: Wed, 1 Jul 2026 21:03:29 +0200 Subject: [PATCH 01/11] Add Arcee squared-ReLU MLP adapter (ArceeForCausalLM) Adds TransformerBridge support for Arcee AI's AFM-4.5B (ArceeForCausalLM), a Llama-style dense decoder whose only architectural novelty is an ungated squared-ReLU (ReLU^2) MLP (up_proj -> ReLU^2 -> down_proj) instead of the gated SiLU/GeLU used by Llama. - Register the "relu2" activation (relu(x)**2) in SUPPORTED_ACTIVATIONS so the config guard and ActivationFunctionFactory accept HF's hidden_act="relu2". - New ArceeArchitectureAdapter: Llama attention (RoPE, RMSNorm, GQA, no QK-norm, no biases) + ungated MLPBridge. Post-activation ReLU^2 neurons are inspectable via blocks.{i}.mlp.hook_post. - Register at the four sites (factory, supported_architectures __init__, model_registry HF_SUPPORTED_ARCHITECTURES + CANONICAL_AUTHORS_BY_ARCH, generate_report descriptions) and add the two checkpoints to the registry. - Unit + integration parity tests (fp32 + eager); integration asserts bridge logits match HF and hook_post exposes non-negative d_mlp-width activations. Closes #1467 Co-Authored-By: Claude Opus 4.8 --- .../model_bridge/test_arcee_adapter.py | 119 ++++++++ .../test_arcee_adapter.py | 256 ++++++++++++++++++ .../factories/architecture_adapter_factory.py | 2 + .../supported_architectures/__init__.py | 4 + .../supported_architectures/arcee.py | 136 ++++++++++ .../tools/model_registry/__init__.py | 2 + .../model_registry/data/supported_models.json | 28 ++ .../tools/model_registry/generate_report.py | 1 + .../utilities/activation_functions.py | 14 + 9 files changed, 562 insertions(+) create mode 100644 tests/integration/model_bridge/test_arcee_adapter.py create mode 100644 tests/unit/model_bridge/supported_architectures/test_arcee_adapter.py create mode 100644 transformer_lens/model_bridge/supported_architectures/arcee.py diff --git a/tests/integration/model_bridge/test_arcee_adapter.py b/tests/integration/model_bridge/test_arcee_adapter.py new file mode 100644 index 000000000..4b85e1311 --- /dev/null +++ b/tests/integration/model_bridge/test_arcee_adapter.py @@ -0,0 +1,119 @@ +"""Integration tests for the Arcee architecture adapter. + +Loads the tiny-random ArceeForCausalLM checkpoint and asserts logit parity with +HF plus Arcee-specific structure: an ungated ReLU^2 MLP whose post-activation +neurons are inspectable via ``blocks.{i}.mlp.hook_post``. + +Runs at fp32 + eager attention (the adapter forces eager so output_attentions — +and therefore hook_attn_scores / hook_pattern — work). +""" +import pytest +import torch + +from transformer_lens.model_bridge.bridge import TransformerBridge + +MODEL = "optimum-intel-internal-testing/tiny-random-ArceeForCausalLM" + + +@pytest.fixture(scope="module") +def arcee_bridge(): + # The adapter forces eager attention via cfg.attn_implementation, so we do not + # (and cannot) pass attn_implementation to boot_transformers. + return TransformerBridge.boot_transformers(MODEL, device="cpu", dtype=torch.float32) + + +class TestArceeBridgeCreation: + def test_block_count(self, arcee_bridge): + assert len(arcee_bridge.blocks) == 2 + + def test_has_core_components(self, arcee_bridge): + assert hasattr(arcee_bridge, "embed") + assert hasattr(arcee_bridge, "unembed") + assert hasattr(arcee_bridge, "ln_final") + + def test_config_flags(self, arcee_bridge): + cfg = arcee_bridge.cfg + assert cfg.normalization_type == "RMS" + assert cfg.positional_embedding_type == "rotary" + assert cfg.gated_mlp is False # ungated ReLU^2 MLP + assert cfg.act_fn == "relu2" + + def test_gqa_config(self, arcee_bridge): + """tiny-random config: 4 attention heads, 2 KV heads.""" + assert arcee_bridge.cfg.n_key_value_heads == 2 + + +class TestArceeForwardEquivalence: + def test_forward_returns_logits(self, arcee_bridge): + tokens = torch.tensor([[1, 2, 3, 4]]) + with torch.no_grad(): + output = arcee_bridge(tokens) + assert output.shape[0] == 1 + assert output.shape[1] == 4 + assert not torch.isnan(output).any() + + def test_forward_matches_hf(self, arcee_bridge): + """Bridge delegates to HF native forward — output should match at fp32.""" + tokens = torch.tensor([[1, 2, 3, 4]]) + hf_model = arcee_bridge.original_model + with torch.no_grad(): + bridge_out = arcee_bridge(tokens) + hf_out = hf_model(tokens).logits + max_diff = (bridge_out - hf_out).abs().max().item() + assert max_diff < 1e-4, f"Bridge vs HF max diff = {max_diff}" + + +class TestArceeHFDelegation: + def test_mlp_delegates_to_hf_relu2_mlp(self, arcee_bridge): + """The ungated MLP bridge wraps HF's ArceeMLP, whose activation is HF's + ReLUSquaredActivation (ReLU^2) — i.e. the bridge delegates the MLP forward + to the real squared-ReLU implementation rather than a gated substitute.""" + hf_mlp = arcee_bridge.blocks[0].mlp.original_component + assert type(hf_mlp).__name__ == "ArceeMLP" + assert type(hf_mlp.act_fn).__name__ == "ReLUSquaredActivation" + + +class TestArceeHookShapes: + def test_attn_and_mlp_hooks_fire(self, arcee_bridge): + tokens = torch.tensor([[1, 2, 3, 4]]) + _, cache = arcee_bridge.run_with_cache(tokens) + for i in range(2): + assert f"blocks.{i}.attn.hook_in" in cache + assert f"blocks.{i}.attn.hook_out" in cache + assert f"blocks.{i}.mlp.hook_in" in cache + assert f"blocks.{i}.mlp.hook_out" in cache + + def test_residual_hooks_fire(self, arcee_bridge): + tokens = torch.tensor([[1, 2, 3, 4]]) + _, cache = arcee_bridge.run_with_cache(tokens) + for i in range(2): + assert f"blocks.{i}.hook_resid_pre" in cache + assert f"blocks.{i}.hook_resid_post" in cache + + +class TestArceeReLU2MLP: + """Arcee's distinguishing quirk: the post-activation MLP neurons (ReLU^2 + output, input to down_proj) are exposed via hook_post and are non-negative + (ReLU^2 >= 0), which is the sparse-activation structure the issue targets.""" + + def test_mlp_post_activation_hook_fires(self, arcee_bridge): + tokens = torch.tensor([[1, 2, 3, 4]]) + _, cache = arcee_bridge.run_with_cache(tokens) + for i in range(2): + assert f"blocks.{i}.mlp.hook_post" in cache + + def test_mlp_post_activation_shape(self, arcee_bridge): + """hook_post has the intermediate (d_mlp) width.""" + tokens = torch.tensor([[1, 2, 3, 4]]) + _, cache = arcee_bridge.run_with_cache(tokens) + post = cache["blocks.0.mlp.hook_post"] + assert post.shape[0] == 1 + assert post.shape[1] == 4 + assert post.shape[-1] == arcee_bridge.cfg.d_mlp + + def test_mlp_post_activation_nonnegative(self, arcee_bridge): + """ReLU^2 output is always >= 0.""" + tokens = torch.tensor([[1, 2, 3, 4]]) + _, cache = arcee_bridge.run_with_cache(tokens) + post = cache["blocks.0.mlp.hook_post"] + assert (post >= 0).all(), "ReLU^2 activations must be non-negative" diff --git a/tests/unit/model_bridge/supported_architectures/test_arcee_adapter.py b/tests/unit/model_bridge/supported_architectures/test_arcee_adapter.py new file mode 100644 index 000000000..b001d12a0 --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_arcee_adapter.py @@ -0,0 +1,256 @@ +"""Unit tests for ArceeArchitectureAdapter. + +Arcee (ArceeForCausalLM / AFM-4.5B) is Llama-shaped except for its MLP: an +*ungated* feed-forward block using the squared-ReLU (ReLU^2) activation. These +tests pin the architecture-specific quirks: + +- config flags (RMSNorm, rotary, eager attention, ``gated_mlp = False``, GQA) +- component mapping uses the ungated ``MLPBridge`` (no ``gate`` submodule) and + standard ``input_layernorm`` / ``post_attention_layernorm`` names with no QK-norm +- the ``"relu2"`` activation is registered and numerically correct +- ``TransformerBridgeConfig`` accepts ``act_fn = "relu2"`` (the config guard that + asserts ``act_fn in SUPPORTED_ACTIVATIONS`` — the reason the activation must be + registered) +- the ungated MLP exposes post-activation neurons via ``hook_post`` +""" +import pytest +import torch + +from transformer_lens.config import HookedTransformerConfig, TransformerBridgeConfig +from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion +from transformer_lens.conversion_utils.param_processing_conversion import ( + ParamProcessingConversion, +) +from transformer_lens.factories.activation_function_factory import ( + ActivationFunctionFactory, +) +from transformer_lens.model_bridge.generalized_components import ( + BlockBridge, + EmbeddingBridge, + GatedMLPBridge, + LinearBridge, + MLPBridge, + RMSNormalizationBridge, + RotaryEmbeddingBridge, + UnembeddingBridge, +) +from transformer_lens.model_bridge.generalized_components.position_embeddings_attention import ( + PositionEmbeddingsAttentionBridge, +) +from transformer_lens.model_bridge.supported_architectures.arcee import ( + ArceeArchitectureAdapter, +) +from transformer_lens.utilities.activation_functions import SUPPORTED_ACTIVATIONS + + +def _make_cfg( + n_heads: int = 8, + n_key_value_heads: int = 4, + d_model: int = 64, + n_layers: int = 2, + d_vocab: int = 100, + n_ctx: int = 128, +) -> TransformerBridgeConfig: + """Minimal TransformerBridgeConfig for Arcee adapter tests.""" + return TransformerBridgeConfig( + d_model=d_model, + d_head=d_model // n_heads, + n_layers=n_layers, + n_ctx=n_ctx, + n_heads=n_heads, + n_key_value_heads=n_key_value_heads, + d_vocab=d_vocab, + d_mlp=4 * d_model, + act_fn="relu2", + default_prepend_bos=False, + architecture="ArceeForCausalLM", + ) + + +@pytest.fixture +def cfg() -> TransformerBridgeConfig: + return _make_cfg() + + +@pytest.fixture +def adapter(cfg: TransformerBridgeConfig) -> ArceeArchitectureAdapter: + return ArceeArchitectureAdapter(cfg) + + +class TestArceeAdapterConfig: + """Adapter config defaults: RMSNorm, rotary, eager attention, ungated MLP, + and GQA propagation via n_key_value_heads.""" + + def test_normalization_is_rms(self, adapter: ArceeArchitectureAdapter) -> None: + assert adapter.cfg.normalization_type == "RMS" + assert adapter.cfg.final_rms is True + assert adapter.cfg.uses_rms_norm is True + + def test_positional_embedding_is_rotary(self, adapter: ArceeArchitectureAdapter) -> None: + assert adapter.cfg.positional_embedding_type == "rotary" + + def test_mlp_is_ungated(self, adapter: ArceeArchitectureAdapter) -> None: + """Arcee's defining feature: an ungated MLP (no gate projection).""" + assert adapter.cfg.gated_mlp is False + + def test_attn_not_only_and_eager(self, adapter: ArceeArchitectureAdapter) -> None: + assert adapter.cfg.attn_only is False + assert adapter.cfg.attn_implementation == "eager" + + def test_gqa_propagated(self, adapter: ArceeArchitectureAdapter) -> None: + assert adapter.cfg.n_key_value_heads == 4 + + +class TestArceeAdapterComponentMapping: + """Component-mapping structure and HF module names. Key contrasts with Llama: + the MLP is an ungated ``MLPBridge`` (no ``gate`` submodule), and attention has + no q_norm/k_norm.""" + + @staticmethod + def _mapping(adapter: ArceeArchitectureAdapter) -> dict: + mapping = adapter.component_mapping + assert mapping is not None + return mapping + + def test_embed(self, adapter: ArceeArchitectureAdapter) -> None: + mapping = self._mapping(adapter) + assert isinstance(mapping["embed"], EmbeddingBridge) + assert mapping["embed"].name == "model.embed_tokens" + + def test_rotary_emb(self, adapter: ArceeArchitectureAdapter) -> None: + mapping = self._mapping(adapter) + assert isinstance(mapping["rotary_emb"], RotaryEmbeddingBridge) + assert mapping["rotary_emb"].name == "model.rotary_emb" + + def test_blocks(self, adapter: ArceeArchitectureAdapter) -> None: + mapping = self._mapping(adapter) + assert isinstance(mapping["blocks"], BlockBridge) + assert mapping["blocks"].name == "model.layers" + + def test_ln_final(self, adapter: ArceeArchitectureAdapter) -> None: + mapping = self._mapping(adapter) + assert isinstance(mapping["ln_final"], RMSNormalizationBridge) + assert mapping["ln_final"].name == "model.norm" + + def test_unembed(self, adapter: ArceeArchitectureAdapter) -> None: + mapping = self._mapping(adapter) + assert isinstance(mapping["unembed"], UnembeddingBridge) + assert mapping["unembed"].name == "lm_head" + + def test_layernorms_use_standard_names(self, adapter: ArceeArchitectureAdapter) -> None: + blocks = self._mapping(adapter)["blocks"] + assert isinstance(blocks.submodules["ln1"], RMSNormalizationBridge) + assert blocks.submodules["ln1"].name == "input_layernorm" + assert isinstance(blocks.submodules["ln2"], RMSNormalizationBridge) + assert blocks.submodules["ln2"].name == "post_attention_layernorm" + + def test_attn_qkvo_names(self, adapter: ArceeArchitectureAdapter) -> None: + attn = self._mapping(adapter)["blocks"].submodules["attn"] + assert isinstance(attn, PositionEmbeddingsAttentionBridge) + assert attn.name == "self_attn" + assert attn.submodules["q"].name == "q_proj" + assert attn.submodules["k"].name == "k_proj" + assert attn.submodules["v"].name == "v_proj" + assert attn.submodules["o"].name == "o_proj" + + def test_attn_has_no_qk_norm(self, adapter: ArceeArchitectureAdapter) -> None: + """Unlike Qwen3/Apertus, Arcee uses plain Llama attention (no QK-norm).""" + attn = self._mapping(adapter)["blocks"].submodules["attn"] + assert "q_norm" not in attn.submodules + assert "k_norm" not in attn.submodules + + def test_mlp_is_ungated_bridge(self, adapter: ArceeArchitectureAdapter) -> None: + """The MLP must be the ungated MLPBridge, not a GatedMLPBridge, and must + expose only up_proj / down_proj (no gate_proj).""" + mlp = self._mapping(adapter)["blocks"].submodules["mlp"] + assert isinstance(mlp, MLPBridge) + assert not isinstance(mlp, GatedMLPBridge) + assert mlp.name == "mlp" + assert isinstance(mlp.submodules["in"], LinearBridge) + assert mlp.submodules["in"].name == "up_proj" + assert isinstance(mlp.submodules["out"], LinearBridge) + assert mlp.submodules["out"].name == "down_proj" + assert "gate" not in mlp.submodules + + def test_mlp_exposes_post_activation_hook(self, adapter: ArceeArchitectureAdapter) -> None: + """hook_post surfaces the post-activation MLP neurons (input to down_proj, + i.e. relu2(up_proj(x))) — the sparse-activation structure the issue targets.""" + mlp = self._mapping(adapter)["blocks"].submodules["mlp"] + assert mlp.hook_aliases["hook_post"] == "out.hook_in" + + +class TestArceeReLU2Activation: + """The squared-ReLU activation: registration, numerics, and the config guard + that requires it to be registered before an Arcee config can be constructed.""" + + def test_relu2_registered(self) -> None: + assert "relu2" in SUPPORTED_ACTIVATIONS + + def test_relu2_numerics(self) -> None: + """relu2(x) == relu(x) ** 2: zero for negatives, x^2 for positives. + + The activation is jaxtyping-annotated ``[batch, pos, d_mlp]`` (like the + other activations), so it is exercised with a 3-D tensor. + """ + fn = SUPPORTED_ACTIVATIONS["relu2"] + x = torch.tensor([-3.0, -0.5, 0.0, 0.5, 2.0]).reshape(1, 1, 5) + expected = torch.tensor([0.0, 0.0, 0.0, 0.25, 4.0]).reshape(1, 1, 5) + assert torch.allclose(fn(x), expected) + + def test_factory_picks_relu2(self) -> None: + """The activation factory resolves act_fn="relu2" to the squared-ReLU fn. + + Uses a HookedTransformerConfig (the factory's declared param type), which + also exercises the __post_init__ guard asserting act_fn in + SUPPORTED_ACTIVATIONS.""" + htc = HookedTransformerConfig( + n_layers=2, + d_model=64, + n_ctx=128, + d_head=16, + n_heads=4, + d_vocab=100, + d_mlp=256, + act_fn="relu2", + ) + fn = ActivationFunctionFactory.pick_activation_function(htc) + x = torch.tensor([-1.0, 3.0]).reshape(1, 1, 2) + assert torch.allclose(fn(x), torch.tensor([0.0, 9.0]).reshape(1, 1, 2)) + + def test_config_accepts_relu2(self) -> None: + """Constructing the config must not raise: the __post_init__ guard asserts + act_fn in SUPPORTED_ACTIVATIONS.""" + cfg = _make_cfg() + assert cfg.act_fn == "relu2" + + +class TestArceeAdapterWeightConversions: + """QKVO weight conversions with GQA-aware head counts: Q uses n_heads; + K and V use n_key_value_heads.""" + + def test_four_conversion_keys(self, adapter: ArceeArchitectureAdapter) -> None: + convs = adapter.weight_processing_conversions + assert convs is not None + assert len(convs) == 4 + + def test_q_uses_n_heads(self, adapter: ArceeArchitectureAdapter) -> None: + conv = adapter.weight_processing_conversions["blocks.{i}.attn.q.weight"] + assert isinstance(conv, ParamProcessingConversion) + assert isinstance(conv.tensor_conversion, RearrangeTensorConversion) + assert conv.tensor_conversion.axes_lengths["n"] == adapter.cfg.n_heads + + def test_k_uses_n_key_value_heads(self, adapter: ArceeArchitectureAdapter) -> None: + conv = adapter.weight_processing_conversions["blocks.{i}.attn.k.weight"] + assert isinstance(conv, ParamProcessingConversion) + assert conv.tensor_conversion.axes_lengths["n"] == adapter.cfg.n_key_value_heads + + def test_v_uses_n_key_value_heads(self, adapter: ArceeArchitectureAdapter) -> None: + conv = adapter.weight_processing_conversions["blocks.{i}.attn.v.weight"] + assert isinstance(conv, ParamProcessingConversion) + assert conv.tensor_conversion.axes_lengths["n"] == adapter.cfg.n_key_value_heads + + def test_o_pattern(self, adapter: ArceeArchitectureAdapter) -> None: + conv = adapter.weight_processing_conversions["blocks.{i}.attn.o.weight"] + assert isinstance(conv, ParamProcessingConversion) + assert conv.tensor_conversion.pattern == "m (n h) -> n h m" + assert conv.tensor_conversion.axes_lengths["n"] == adapter.cfg.n_heads diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py index 2342908d3..c7aa2acf2 100644 --- a/transformer_lens/factories/architecture_adapter_factory.py +++ b/transformer_lens/factories/architecture_adapter_factory.py @@ -11,6 +11,7 @@ from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter from transformer_lens.model_bridge.supported_architectures import ( ApertusArchitectureAdapter, + ArceeArchitectureAdapter, BaichuanArchitectureAdapter, BartArchitectureAdapter, BertArchitectureAdapter, @@ -82,6 +83,7 @@ # Export supported architectures SUPPORTED_ARCHITECTURES = { "ApertusForCausalLM": ApertusArchitectureAdapter, + "ArceeForCausalLM": ArceeArchitectureAdapter, "BaiChuanForCausalLM": BaichuanArchitectureAdapter, "BaichuanForCausalLM": BaichuanArchitectureAdapter, "BartForConditionalGeneration": BartArchitectureAdapter, diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py index 8f4cfea3e..efa25f5d6 100644 --- a/transformer_lens/model_bridge/supported_architectures/__init__.py +++ b/transformer_lens/model_bridge/supported_architectures/__init__.py @@ -6,6 +6,9 @@ from transformer_lens.model_bridge.supported_architectures.apertus import ( ApertusArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.arcee import ( + ArceeArchitectureAdapter, +) from transformer_lens.model_bridge.supported_architectures.baichuan import ( BaichuanArchitectureAdapter, ) @@ -207,6 +210,7 @@ __all__ = [ "ApertusArchitectureAdapter", + "ArceeArchitectureAdapter", "BaichuanArchitectureAdapter", "BartArchitectureAdapter", "BertArchitectureAdapter", diff --git a/transformer_lens/model_bridge/supported_architectures/arcee.py b/transformer_lens/model_bridge/supported_architectures/arcee.py new file mode 100644 index 000000000..803b13bba --- /dev/null +++ b/transformer_lens/model_bridge/supported_architectures/arcee.py @@ -0,0 +1,136 @@ +"""Arcee architecture adapter.""" + +from typing import Any + +from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter +from transformer_lens.model_bridge.generalized_components import ( + BlockBridge, + EmbeddingBridge, + LinearBridge, + MLPBridge, + PositionEmbeddingsAttentionBridge, + RMSNormalizationBridge, + RotaryEmbeddingBridge, + UnembeddingBridge, +) + + +class ArceeArchitectureAdapter(ArchitectureAdapter): + """Architecture adapter for Arcee models (ArceeForCausalLM / AFM-4.5B). + + Arcee is a Llama-style dense decoder: pre-norm RMSNorm, rotary position + embeddings (RoPE), grouped query attention (GQA), and no biases on any + projection. The single distinguishing feature is the MLP: an *ungated* + feed-forward block (``up_proj -> ReLU^2 -> down_proj``) using the squared-ReLU + activation (HF ``hidden_act = "relu2"``) instead of the gated SiLU/GeLU used by + Llama. The post-activation neurons are exposed via the MLP bridge's + ``hook_post`` (``mlp.out.hook_in``), which is useful for inspecting the sparse + activation structure ReLU^2 produces. + + Structurally identical to Llama except for the ungated ReLU^2 MLP; unlike + Apertus it uses standard ``input_layernorm`` / ``post_attention_layernorm`` + names and has no Q/K normalization. + + Optional Parameters (may not exist in state_dict): + ------------------------------------------------- + Arcee models do NOT have biases on attention or MLP projections + (``attention_bias = false``, ``mlp_bias = false``): + + - blocks.{i}.attn.b_Q / b_K / b_V / b_O - No bias on attention projections + - blocks.{i}.mlp.b_in - No bias on MLP input (up_proj) + - blocks.{i}.mlp.b_out - No bias on MLP output (down_proj) + - blocks.{i}.ln1.b / ln2.b / ln_final.b - RMSNorm has no bias + + Weight processing handles these missing biases gracefully via + ProcessWeights._safe_get_tensor(). + """ + + def __init__(self, cfg: Any) -> None: + """Initialize the Arcee architecture adapter.""" + super().__init__(cfg) + + # Set config variables for weight processing + self.cfg.normalization_type = "RMS" + self.cfg.positional_embedding_type = "rotary" + self.cfg.final_rms = True + self.cfg.gated_mlp = False # ungated ReLU^2 MLP (up_proj -> act -> down_proj) + self.cfg.attn_only = False + self.cfg.uses_rms_norm = True + + # Use eager attention so output_attentions works for hook_attn_scores / + # hook_pattern; SDPA does not support output_attentions. + self.cfg.attn_implementation = "eager" + + self.default_config = { + "d_model": cfg.d_model, + "d_head": cfg.d_model // cfg.n_heads, + "n_heads": cfg.n_heads, + "n_layers": cfg.n_layers, + "d_vocab": cfg.d_vocab, + } + + # GQA support (num_key_value_heads < num_attention_heads). Must set on cfg, + # not just default_config, so _qkvo_weight_conversions() picks it up. + if hasattr(cfg, "n_key_value_heads") and cfg.n_key_value_heads is not None: + self.default_config["n_key_value_heads"] = cfg.n_key_value_heads + self.cfg.n_key_value_heads = cfg.n_key_value_heads + + self.weight_processing_conversions = { + **self._qkvo_weight_conversions(), + } + + self.component_mapping = { + "embed": EmbeddingBridge(name="model.embed_tokens"), + "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"), + "blocks": BlockBridge( + name="model.layers", + submodules={ + "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), + "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg), + "attn": PositionEmbeddingsAttentionBridge( + name="self_attn", + config=self.cfg, + submodules={ + "q": LinearBridge(name="q_proj"), + "k": LinearBridge(name="k_proj"), + "v": LinearBridge(name="v_proj"), + "o": LinearBridge(name="o_proj"), + }, + requires_attention_mask=True, + requires_position_embeddings=True, + ), + "mlp": MLPBridge( + name="mlp", + submodules={ + "in": LinearBridge(name="up_proj"), + "out": LinearBridge(name="down_proj"), + }, + ), + }, + ), + "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), + "unembed": UnembeddingBridge(name="lm_head", config=self.cfg), + } + + def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None: + """Set up rotary embedding references for Arcee component testing. + + Arcee uses RoPE (Rotary Position Embeddings). We set the rotary_emb reference + on all attention bridge instances for component testing. + + Args: + hf_model: The HuggingFace Arcee model instance + bridge_model: The TransformerBridge model (if available, set rotary_emb on actual instances) + """ + # Get rotary embedding instance from the model + rotary_emb = hf_model.model.rotary_emb + + # Set rotary_emb on actual bridge instances in bridge_model if available + if bridge_model is not None and hasattr(bridge_model, "blocks"): + for block in bridge_model.blocks: + if hasattr(block, "attn"): + block.attn.set_rotary_emb(rotary_emb) + + # Also set on the template for get_generalized_component() calls + attn_bridge = self.get_generalized_component("blocks.0.attn") + attn_bridge.set_rotary_emb(rotary_emb) diff --git a/transformer_lens/tools/model_registry/__init__.py b/transformer_lens/tools/model_registry/__init__.py index 6d2a2927a..e4b2e93d6 100644 --- a/transformer_lens/tools/model_registry/__init__.py +++ b/transformer_lens/tools/model_registry/__init__.py @@ -46,6 +46,7 @@ # GPTNeoX) in config.architectures instead. HF_SUPPORTED_ARCHITECTURES: set[str] = { "ApertusForCausalLM", + "ArceeForCausalLM", "BaiChuanForCausalLM", "BaichuanForCausalLM", "BartForConditionalGeneration", @@ -117,6 +118,7 @@ # download-threshold bypass and the docs table's "Canonical only" toggle. CANONICAL_AUTHORS_BY_ARCH: dict[str, list[str]] = { "ApertusForCausalLM": ["swiss-ai"], + "ArceeForCausalLM": ["arcee-ai"], "BaiChuanForCausalLM": ["baichuan-inc"], "BaichuanForCausalLM": ["baichuan-inc"], "BartForConditionalGeneration": ["facebook"], diff --git a/transformer_lens/tools/model_registry/data/supported_models.json b/transformer_lens/tools/model_registry/data/supported_models.json index 6074ca7cd..f8b282aa1 100644 --- a/transformer_lens/tools/model_registry/data/supported_models.json +++ b/transformer_lens/tools/model_registry/data/supported_models.json @@ -179206,6 +179206,34 @@ "phase4_score": null, "phase7_score": null, "phase8_score": null + }, + { + "architecture_id": "ArceeForCausalLM", + "model_id": "optimum-intel-internal-testing/tiny-random-ArceeForCausalLM", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "ArceeForCausalLM", + "model_id": "arcee-ai/AFM-4.5B-Base", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null } ] } diff --git a/transformer_lens/tools/model_registry/generate_report.py b/transformer_lens/tools/model_registry/generate_report.py index 73f0aa8b1..5bfc96825 100644 --- a/transformer_lens/tools/model_registry/generate_report.py +++ b/transformer_lens/tools/model_registry/generate_report.py @@ -26,6 +26,7 @@ # Descriptions of common architectures (both supported and unsupported) ARCHITECTURE_DESCRIPTIONS: dict[str, str] = { # Supported architectures + "ArceeForCausalLM": "Arcee AI's AFM, a Llama-style decoder with an ungated squared-ReLU (ReLU^2) MLP", "GPT2LMHeadModel": "OpenAI's GPT-2 decoder-only transformer for causal language modeling", "GPTNeoForCausalLM": "EleutherAI's GPT-Neo, an open-source GPT-3-like model", "GPTNeoXForCausalLM": "EleutherAI's GPT-NeoX architecture used in Pythia models", diff --git a/transformer_lens/utilities/activation_functions.py b/transformer_lens/utilities/activation_functions.py index b0a280967..3594acae8 100644 --- a/transformer_lens/utilities/activation_functions.py +++ b/transformer_lens/utilities/activation_functions.py @@ -105,6 +105,19 @@ def xielu(input: Float[torch.Tensor, "batch pos d_mlp"]) -> Float[torch.Tensor, ) +def relu2(input: Float[torch.Tensor, "batch pos d_mlp"]) -> Float[torch.Tensor, "batch pos d_mlp"]: + """Squared-ReLU (ReLU^2) activation function: ``relu(x) ** 2``. + + Used by Arcee (ArceeForCausalLM / AFM-4.5B) in an ungated MLP. Matches + HuggingFace's ``ReLUSquaredActivation`` (registered as ``"relu2"`` in + ``transformers`` ACT2FN), so the key here mirrors the HF ``hidden_act`` string. + The squaring drives high activation sparsity, which is useful for neuron-level + interpretability (https://arxiv.org/abs/2109.08668). + """ + relu_applied = F.relu(input) + return relu_applied * relu_applied + + # Convenient type for the format of each activation function ActivationFunction = Callable[..., torch.Tensor] @@ -120,4 +133,5 @@ def xielu(input: Float[torch.Tensor, "batch pos d_mlp"]) -> Float[torch.Tensor, "gelu": F.gelu, "gelu_pytorch_tanh": gelu_pytorch_tanh, "xielu": xielu, + "relu2": relu2, } From 9f82ff8e833c3adcd51b627a4d407d8caf709ebd Mon Sep 17 00:00:00 2001 From: Jonah Larson Date: Thu, 2 Jul 2026 11:00:15 -0500 Subject: [PATCH 02/11] Redo of #1477: Qwen2-MoE bridge support (retarget to dev) (#1482) * Added minimal TransformerBridge support skeleton for Qwen2MoeForCausalLM * Fix return type annotation in Qwen2MoeRouterBridge.forward method * Enhance tests for Qwen2-MoE bridge to validate cache capturing of MoE hooks --------- Co-authored-by: Canonik --- .../model_bridge/test_qwen2_moe_bridge.py | 119 ++++++++++++++ .../test_qwen2_moe_adapter.py | 151 ++++++++++++++++++ .../factories/architecture_adapter_factory.py | 2 + .../model_bridge/sources/_bridge_builder.py | 1 + .../model_bridge/sources/transformers.py | 6 +- .../supported_architectures/__init__.py | 4 + .../supported_architectures/qwen2_moe.py | 67 ++++++++ .../tools/model_registry/__init__.py | 2 + .../model_registry/data/supported_models.json | 18 ++- .../tools/model_registry/generate_report.py | 1 + 10 files changed, 368 insertions(+), 3 deletions(-) create mode 100644 tests/integration/model_bridge/test_qwen2_moe_bridge.py create mode 100644 tests/unit/model_bridge/supported_architectures/test_qwen2_moe_adapter.py create mode 100644 transformer_lens/model_bridge/supported_architectures/qwen2_moe.py diff --git a/tests/integration/model_bridge/test_qwen2_moe_bridge.py b/tests/integration/model_bridge/test_qwen2_moe_bridge.py new file mode 100644 index 000000000..05ee6a351 --- /dev/null +++ b/tests/integration/model_bridge/test_qwen2_moe_bridge.py @@ -0,0 +1,119 @@ +"""Integration tests for the Qwen2-MoE TransformerBridge.""" + +import copy + +import torch +from transformers import Qwen2MoeConfig, Qwen2MoeForCausalLM + +from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.generalized_components import ( + GatedMLPBridge, + MoEBridge, + PositionEmbeddingsAttentionBridge, +) +from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_from_module, +) + + +def tiny_qwen2_moe_config() -> Qwen2MoeConfig: + return Qwen2MoeConfig( + vocab_size=128, + hidden_size=64, + intermediate_size=96, + moe_intermediate_size=32, + shared_expert_intermediate_size=96, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + num_experts=4, + num_experts_per_tok=2, + max_position_embeddings=64, + decoder_sparse_step=1, + mlp_only_layers=[], + ) + + +def tiny_qwen2_moe_bridge() -> tuple[TransformerBridge, Qwen2MoeForCausalLM]: + torch.manual_seed(0) + cfg = tiny_qwen2_moe_config() + cfg._attn_implementation = "eager" + hf_reference = Qwen2MoeForCausalLM(cfg).eval() + hf_reference.config._attn_implementation = "eager" + for layer in hf_reference.model.layers: + layer.self_attn.config._attn_implementation = "eager" + + hf_model = Qwen2MoeForCausalLM(copy.deepcopy(cfg)).eval() + hf_model.load_state_dict(hf_reference.state_dict()) + bridge = build_bridge_from_module( + hf_model, + "Qwen2MoeForCausalLM", + hf_config=copy.deepcopy(cfg), + tokenizer=None, + device="cpu", + ).eval() + bridge.adapter.setup_component_testing(hf_model, bridge_model=bridge) + return bridge, hf_reference + + +def _tokens() -> torch.Tensor: + return torch.tensor([[1, 2, 3, 4, 5]]) + + +class TestQwen2MoeBridge: + def test_bridge_structure(self) -> None: + bridge, _ = tiny_qwen2_moe_bridge() + + assert len(bridge.blocks) == 2 + assert isinstance(bridge.blocks[0].attn, PositionEmbeddingsAttentionBridge) + assert isinstance(bridge.blocks[0].mlp, MoEBridge) + assert isinstance(bridge.blocks[0].mlp.shared_expert, GatedMLPBridge) + assert hasattr(bridge.blocks[0].mlp, "shared_expert_gate") + + def test_forward_matches_hf(self) -> None: + bridge, hf_reference = tiny_qwen2_moe_bridge() + tokens = _tokens() + + with torch.no_grad(): + bridge_out = bridge(tokens) + hf_out = hf_reference(tokens).logits + + assert bridge_out.shape == (1, 5, 128) + assert not torch.isnan(bridge_out).any() + assert not torch.isinf(bridge_out).any() + max_diff = (bridge_out - hf_out).abs().max().item() + assert max_diff < 1e-5, f"Bridge vs HF max diff = {max_diff}" + + def test_run_with_cache_captures_moe_hooks(self) -> None: + bridge, _ = tiny_qwen2_moe_bridge() + tokens = _tokens() + batch, seq_len = tokens.shape + flat_tokens = batch * seq_len + d_model = bridge.cfg.d_model + num_experts = bridge.cfg.num_experts + + _, cache = bridge.run_with_cache(tokens) + + for layer_idx in range(len(bridge.blocks)): + gate_key = f"blocks.{layer_idx}.mlp.gate.hook_out" + assert gate_key in cache, f"Missing cache key: {gate_key}" + assert cache[gate_key].shape == (flat_tokens, num_experts) + + experts_key = f"blocks.{layer_idx}.mlp.experts.hook_out" + assert experts_key in cache, f"Missing cache key: {experts_key}" + assert cache[experts_key].shape == (flat_tokens, d_model) + + shared_expert_key = f"blocks.{layer_idx}.mlp.shared_expert.hook_out" + assert shared_expert_key in cache, f"Missing cache key: {shared_expert_key}" + assert cache[shared_expert_key].shape == (flat_tokens, d_model) + + shared_expert_gate_key = f"blocks.{layer_idx}.mlp.shared_expert_gate.hook_out" + assert shared_expert_gate_key in cache, f"Missing cache key: {shared_expert_gate_key}" + assert cache[shared_expert_gate_key].shape == (flat_tokens, 1) + + mlp_out_key = f"blocks.{layer_idx}.mlp.hook_out" + assert mlp_out_key in cache, f"Missing cache key: {mlp_out_key}" + assert cache[mlp_out_key].shape == (batch, seq_len, d_model) + + router_scores_key = f"blocks.{layer_idx}.mlp.hook_router_scores" + assert router_scores_key not in cache diff --git a/tests/unit/model_bridge/supported_architectures/test_qwen2_moe_adapter.py b/tests/unit/model_bridge/supported_architectures/test_qwen2_moe_adapter.py new file mode 100644 index 000000000..f405b2fd1 --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_qwen2_moe_adapter.py @@ -0,0 +1,151 @@ +"""Unit tests for the Qwen2MoeArchitectureAdapter.""" + +import pytest +import torch +from transformers import Qwen2MoeConfig + +from transformer_lens.conversion_utils.conversion_steps.rearrange_tensor_conversion import ( + RearrangeTensorConversion, +) +from transformer_lens.conversion_utils.param_processing_conversion import ( + ParamProcessingConversion, +) +from transformer_lens.factories.architecture_adapter_factory import ( + ArchitectureAdapterFactory, +) +from transformer_lens.model_bridge.generalized_components import ( + GatedMLPBridge, + LinearBridge, + MoEBridge, + PositionEmbeddingsAttentionBridge, +) +from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_config_from_hf, +) +from transformer_lens.model_bridge.sources.transformers import ( + determine_architecture_from_hf_config, +) +from transformer_lens.model_bridge.supported_architectures.qwen2_moe import ( + Qwen2MoeArchitectureAdapter, +) + + +def _tiny_config() -> Qwen2MoeConfig: + return Qwen2MoeConfig( + vocab_size=128, + hidden_size=64, + intermediate_size=96, + moe_intermediate_size=32, + shared_expert_intermediate_size=96, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + num_experts=4, + num_experts_per_tok=2, + max_position_embeddings=64, + decoder_sparse_step=1, + mlp_only_layers=[], + ) + + +@pytest.fixture +def bridge_cfg(): + return build_bridge_config_from_hf( + _tiny_config(), + "Qwen2MoeForCausalLM", + "qwen2-moe-test", + torch.float32, + ) + + +@pytest.fixture +def adapter(bridge_cfg) -> Qwen2MoeArchitectureAdapter: + return Qwen2MoeArchitectureAdapter(bridge_cfg) + + +class TestQwen2MoeDetection: + def test_model_type_detects_qwen2_moe(self) -> None: + assert determine_architecture_from_hf_config(_tiny_config()) == "Qwen2MoeForCausalLM" + + def test_factory_selects_adapter(self, bridge_cfg) -> None: + selected = ArchitectureAdapterFactory.select_architecture_adapter(bridge_cfg) + assert isinstance(selected, Qwen2MoeArchitectureAdapter) + + +class TestQwen2MoeConfigMapping: + def test_core_fields_map_from_hf_config(self, bridge_cfg) -> None: + hf_cfg = _tiny_config() + assert bridge_cfg.d_vocab == hf_cfg.vocab_size + assert bridge_cfg.d_model == hf_cfg.hidden_size + assert bridge_cfg.n_layers == hf_cfg.num_hidden_layers + assert bridge_cfg.n_heads == hf_cfg.num_attention_heads + assert bridge_cfg.n_key_value_heads == hf_cfg.num_key_value_heads + assert bridge_cfg.d_mlp == hf_cfg.intermediate_size + + def test_moe_fields_map_from_hf_config(self, bridge_cfg) -> None: + hf_cfg = _tiny_config() + assert bridge_cfg.num_experts == hf_cfg.num_experts + assert bridge_cfg.experts_per_token == hf_cfg.num_experts_per_tok + assert getattr(bridge_cfg, "moe_intermediate_size") == hf_cfg.moe_intermediate_size + assert ( + getattr(bridge_cfg, "shared_expert_intermediate_size") + == hf_cfg.shared_expert_intermediate_size + ) + + +class TestQwen2MoeComponentMapping: + def test_reuses_qwen2_attention_mapping(self, adapter: Qwen2MoeArchitectureAdapter) -> None: + blocks = adapter.component_mapping["blocks"] + attn = blocks.submodules["attn"] + assert isinstance(attn, PositionEmbeddingsAttentionBridge) + assert set(attn.submodules.keys()) == {"q", "k", "v", "o"} + assert attn.submodules["q"].name == "q_proj" + assert attn.submodules["k"].name == "k_proj" + assert attn.submodules["v"].name == "v_proj" + assert attn.submodules["o"].name == "o_proj" + + def test_mlp_is_moe_bridge(self, adapter: Qwen2MoeArchitectureAdapter) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + assert isinstance(mlp, MoEBridge) + assert mlp.name == "mlp" + + def test_moe_submodules(self, adapter: Qwen2MoeArchitectureAdapter) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + assert set(mlp.submodules.keys()) == { + "gate", + "experts", + "shared_expert", + "shared_expert_gate", + } + assert isinstance(mlp.submodules["gate"], LinearBridge) + assert isinstance(mlp.submodules["experts"], MoEBridge) + assert isinstance(mlp.submodules["shared_expert"], GatedMLPBridge) + assert isinstance(mlp.submodules["shared_expert_gate"], LinearBridge) + + def test_moe_hf_paths(self, adapter: Qwen2MoeArchitectureAdapter) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + assert mlp.submodules["gate"].name == "gate" + assert mlp.submodules["experts"].name == "experts" + assert mlp.submodules["shared_expert"].name == "shared_expert" + assert mlp.submodules["shared_expert_gate"].name == "shared_expert_gate" + + def test_shared_expert_submodules(self, adapter: Qwen2MoeArchitectureAdapter) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + shared_expert = mlp.submodules["shared_expert"] + assert set(shared_expert.submodules.keys()) == {"gate", "in", "out"} + assert shared_expert.submodules["gate"].name == "gate_proj" + assert shared_expert.submodules["in"].name == "up_proj" + assert shared_expert.submodules["out"].name == "down_proj" + + +class TestQwen2MoeWeightConversions: + def test_kv_rearrange_uses_n_key_value_heads( + self, adapter: Qwen2MoeArchitectureAdapter + ) -> None: + conversions = adapter.weight_processing_conversions + assert conversions is not None + for slot in ("k", "v"): + conv = conversions[f"blocks.{{i}}.attn.{slot}.weight"] + assert isinstance(conv, ParamProcessingConversion) + assert isinstance(conv.tensor_conversion, RearrangeTensorConversion) + assert conv.tensor_conversion.axes_lengths["n"] == 2 diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py index c7aa2acf2..cf18ffcdd 100644 --- a/transformer_lens/factories/architecture_adapter_factory.py +++ b/transformer_lens/factories/architecture_adapter_factory.py @@ -67,6 +67,7 @@ PhiArchitectureAdapter, PhiMoEArchitectureAdapter, Qwen2ArchitectureAdapter, + Qwen2MoeArchitectureAdapter, Qwen3_5ArchitectureAdapter, Qwen3_5MultimodalArchitectureAdapter, Qwen3ArchitectureAdapter, @@ -144,6 +145,7 @@ "PhiMoEForCausalLM": PhiMoEArchitectureAdapter, "QwenForCausalLM": QwenArchitectureAdapter, "Qwen2ForCausalLM": Qwen2ArchitectureAdapter, + "Qwen2MoeForCausalLM": Qwen2MoeArchitectureAdapter, "Qwen3ForCausalLM": Qwen3ArchitectureAdapter, "Qwen3MoeForCausalLM": Qwen3MoeArchitectureAdapter, "Qwen3NextForCausalLM": Qwen3NextArchitectureAdapter, diff --git a/transformer_lens/model_bridge/sources/_bridge_builder.py b/transformer_lens/model_bridge/sources/_bridge_builder.py index 55f53ada8..05e9fd5bf 100644 --- a/transformer_lens/model_bridge/sources/_bridge_builder.py +++ b/transformer_lens/model_bridge/sources/_bridge_builder.py @@ -52,6 +52,7 @@ # Hybrid/MoE architectures "layer_types", "moe_intermediate_size", + "shared_expert_intermediate_size", "norm_eps", "attention_bias", "lm_head_bias", diff --git a/transformer_lens/model_bridge/sources/transformers.py b/transformer_lens/model_bridge/sources/transformers.py index bce60d96f..1e6e6c90f 100644 --- a/transformer_lens/model_bridge/sources/transformers.py +++ b/transformer_lens/model_bridge/sources/transformers.py @@ -180,7 +180,9 @@ def map_default_transformer_lens_config(hf_config): tl_config.eps = source_config.layer_norm_epsilon elif hasattr(source_config, "norm_eps"): tl_config.eps = source_config.norm_eps - if hasattr(source_config, "num_local_experts"): + if hasattr(source_config, "num_experts"): + tl_config.num_experts = source_config.num_experts + elif hasattr(source_config, "num_local_experts"): tl_config.num_experts = source_config.num_local_experts if hasattr(source_config, "num_experts_per_tok"): tl_config.experts_per_token = source_config.num_experts_per_tok @@ -248,6 +250,7 @@ def determine_architecture_from_hf_config(hf_config): "phi3": "Phi3ForCausalLM", "qwen": "QwenForCausalLM", "qwen2": "Qwen2ForCausalLM", + "qwen2_moe": "Qwen2MoeForCausalLM", "qwen3": "Qwen3ForCausalLM", # qwen3_5 is the top-level multimodal config type; qwen3_5_text is # the text-only sub-config. Both map to the text-only adapter so @@ -544,6 +547,7 @@ def boot( # Hybrid/MoE architectures "layer_types", "moe_intermediate_size", + "shared_expert_intermediate_size", "norm_eps", "attention_bias", "lm_head_bias", diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py index efa25f5d6..ba5b56430 100644 --- a/transformer_lens/model_bridge/supported_architectures/__init__.py +++ b/transformer_lens/model_bridge/supported_architectures/__init__.py @@ -177,6 +177,9 @@ from transformer_lens.model_bridge.supported_architectures.qwen2 import ( Qwen2ArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.qwen2_moe import ( + Qwen2MoeArchitectureAdapter, +) from transformer_lens.model_bridge.supported_architectures.qwen3 import ( Qwen3ArchitectureAdapter, ) @@ -268,6 +271,7 @@ "PhiMoEArchitectureAdapter", "QwenArchitectureAdapter", "Qwen2ArchitectureAdapter", + "Qwen2MoeArchitectureAdapter", "Qwen3ArchitectureAdapter", "Qwen3MoeArchitectureAdapter", "Qwen3NextArchitectureAdapter", diff --git a/transformer_lens/model_bridge/supported_architectures/qwen2_moe.py b/transformer_lens/model_bridge/supported_architectures/qwen2_moe.py new file mode 100644 index 000000000..f3e6191a3 --- /dev/null +++ b/transformer_lens/model_bridge/supported_architectures/qwen2_moe.py @@ -0,0 +1,67 @@ +"""Qwen2-MoE architecture adapter.""" + +from typing import Any + +import torch + +from transformer_lens.model_bridge.generalized_components import ( + GatedMLPBridge, + LinearBridge, + MoEBridge, +) +from transformer_lens.model_bridge.supported_architectures.qwen2 import ( + Qwen2ArchitectureAdapter, +) + + +class Qwen2MoeRouterBridge(LinearBridge): + """Bridge Qwen2-MoE router logits while preserving HF's tuple return.""" + + def forward(self, input: torch.Tensor, *args: Any, **kwargs: Any) -> Any: + if self.original_component is None: + raise RuntimeError( + f"Original component not set for {self.name}. Call set_original_component() first." + ) + input = self.hook_in(input) + output = self.original_component(input, *args, **kwargs) + if not isinstance(output, tuple) or len(output) == 0: + return self.hook_out(output) + router_logits = self.hook_out(output[0]) + return (router_logits,) + output[1:] + + +class Qwen2MoeArchitectureAdapter(Qwen2ArchitectureAdapter): + """Architecture adapter for Qwen2-MoE models. + + Qwen2-MoE uses the Qwen2 attention stack plus a sparse MoE MLP with an + always-on shared expert path. + """ + + def __init__(self, cfg: Any) -> None: + """Initialize the Qwen2-MoE architecture adapter.""" + super().__init__(cfg) + + self.cfg.attn_implementation = "eager" + + if self.component_mapping is None: + raise ValueError("Qwen2 component mapping was not initialized") + + blocks = self.component_mapping["blocks"] + blocks.submodules["mlp"] = MoEBridge( + name="mlp", + config=self.cfg, + submodules={ + "gate": Qwen2MoeRouterBridge(name="gate"), + "experts": MoEBridge(name="experts", config=self.cfg), + "shared_expert": GatedMLPBridge( + name="shared_expert", + config=self.cfg, + submodules={ + "gate": LinearBridge(name="gate_proj"), + "in": LinearBridge(name="up_proj"), + "out": LinearBridge(name="down_proj"), + }, + ), + "shared_expert_gate": LinearBridge(name="shared_expert_gate"), + }, + ) diff --git a/transformer_lens/tools/model_registry/__init__.py b/transformer_lens/tools/model_registry/__init__.py index e4b2e93d6..b0be24999 100644 --- a/transformer_lens/tools/model_registry/__init__.py +++ b/transformer_lens/tools/model_registry/__init__.py @@ -101,6 +101,7 @@ "PhiMoEForCausalLM", "QwenForCausalLM", "Qwen2ForCausalLM", + "Qwen2MoeForCausalLM", "Qwen3ForCausalLM", "Qwen3MoeForCausalLM", "Qwen3NextForCausalLM", @@ -173,6 +174,7 @@ "PhiMoEForCausalLM": ["microsoft"], "PhiForCausalLM": ["microsoft"], "Qwen2ForCausalLM": ["Qwen", "nvidia"], + "Qwen2MoeForCausalLM": ["Qwen"], "Qwen3ForCausalLM": ["Qwen", "nvidia"], "Qwen3MoeForCausalLM": ["Qwen"], "Qwen3NextForCausalLM": ["Qwen"], diff --git a/transformer_lens/tools/model_registry/data/supported_models.json b/transformer_lens/tools/model_registry/data/supported_models.json index f8b282aa1..a01926111 100644 --- a/transformer_lens/tools/model_registry/data/supported_models.json +++ b/transformer_lens/tools/model_registry/data/supported_models.json @@ -6,10 +6,24 @@ "min_downloads": 500, "scan_duration_seconds": 8.1 }, - "total_architectures": 65, - "total_models": 12899, + "total_architectures": 66, + "total_models": 12900, "total_verified": 1017, "models": [ + { + "architecture_id": "Qwen2MoeForCausalLM", + "model_id": "hyper-accel/ci-random-qwen2-moe-a3b", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, { "architecture_id": "HunYuanDenseV1ForCausalLM", "model_id": "tencent/Hunyuan-0.5B-Instruct", diff --git a/transformer_lens/tools/model_registry/generate_report.py b/transformer_lens/tools/model_registry/generate_report.py index 5bfc96825..5e7f4f065 100644 --- a/transformer_lens/tools/model_registry/generate_report.py +++ b/transformer_lens/tools/model_registry/generate_report.py @@ -44,6 +44,7 @@ "GlmMoeDsaForCausalLM": "Z.ai's GLM-5 MoE model with DeepSeek Sparse Attention", "Glm4MoeForCausalLM": "Z.ai's GLM-4.5/4.6/4.7 sparse Mixture-of-Experts causal LM", "Qwen2ForCausalLM": "Alibaba's Qwen2 multilingual model", + "Qwen2MoeForCausalLM": "Alibaba's Qwen2 sparse Mixture-of-Experts model with shared experts", "Qwen3ForCausalLM": "Alibaba's Qwen3 latest generation", "Qwen3_5ForConditionalGeneration": "Alibaba's Qwen3.5 vision-language model", "BloomForCausalLM": "BigScience's BLOOM multilingual model", From d4560ca49c2ec309d604d8be5c66e685f251a964 Mon Sep 17 00:00:00 2001 From: Jonah Larson Date: Thu, 2 Jul 2026 13:17:35 -0500 Subject: [PATCH 03/11] SSM Interp Support (#1481) * Normalized SSM mixer across existing hybrids * GatedDeltaNet and compute_ssm_state improvements * Prepping Mamba1 mixer * Unify vocabulary * SSM2 interp support * Mamba1 read parity * Fixed Mamba1 bugs * Fix GPU device mismatch * Additional cleanup and testing * make format * Model verification + comment cleanup * documentation updates * Final SSM pass --- README.md | 12 +- docs/source/content/ssm_interpretability.md | 117 + docs/source/index.md | 1 + .../test_granite_moe_hybrid_adapter.py | 436 ++ .../model_bridge/test_mamba2_adapter.py | 307 +- .../model_bridge/test_mamba_adapter.py | 497 ++ .../model_bridge/test_nemotron_h_adapter.py | 8 +- .../model_bridge/test_nemotron_h_tiny.py | 135 + .../test_ssm_effective_attention_dispatch.py | 416 ++ .../model_bridge/test_ssm_real_weights.py | 156 + tests/mps/test_mps_ssm_eager_scan.py | 82 + .../test_granite_moe_hybrid_adapter.py | 36 +- .../test_nemotron_h_adapter.py | 5 +- .../test_qwen3_5_adapter.py | 471 ++ .../tools/model_registry/test_verify_gates.py | 100 + transformer_lens/ActivationCache.py | 182 + .../benchmarks/component_outputs.py | 15 + .../benchmarks/hook_registration.py | 6 + .../benchmarks/weight_processing.py | 24 +- .../generalized_components/__init__.py | 8 + .../generalized_components/gated_delta_net.py | 225 +- .../generalized_components/ssm2_mixer.py | 299 +- .../generalized_components/ssm_mixer.py | 288 +- .../generalized_components/ssm_protocol.py | 99 + .../granite_moe_hybrid.py | 28 +- .../supported_architectures/mamba.py | 8 +- .../supported_architectures/mamba2.py | 94 +- .../supported_architectures/nemotron_h.py | 23 +- .../tools/model_registry/AGENTS.md | 2 + .../data/architecture_gaps.json | 6294 +++++++++-------- .../model_registry/data/supported_models.json | 3580 +++++++++- .../data/verification_history.json | 472 +- .../tools/model_registry/verify_models.py | 11 +- 33 files changed, 10935 insertions(+), 3502 deletions(-) create mode 100644 docs/source/content/ssm_interpretability.md create mode 100644 tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py create mode 100644 tests/integration/model_bridge/test_nemotron_h_tiny.py create mode 100644 tests/integration/model_bridge/test_ssm_effective_attention_dispatch.py create mode 100644 tests/integration/model_bridge/test_ssm_real_weights.py create mode 100644 tests/mps/test_mps_ssm_eager_scan.py create mode 100644 tests/unit/tools/model_registry/test_verify_gates.py create mode 100644 transformer_lens/model_bridge/generalized_components/ssm_protocol.py diff --git a/README.md b/README.md index 74a4c3258..c76c8fe91 100644 --- a/README.md +++ b/README.md @@ -184,12 +184,12 @@ cover: Verification lives in the integration tests at `tests/integration/model_bridge/test_mamba_adapter.py` and -`tests/integration/model_bridge/test_mamba2_adapter.py` (81 tests total). The -`verify_models` benchmark suite does not currently cover SSM architectures — the -benchmark has transformer-shaped assumptions (hook path patterns, layer norm -folding, block submodule dispatch) that would need a dedicated refactor. SSM -benchmark coverage will be revisited if hybrid architectures like Jamba or -Falcon-H1 become priority or a user explicitly requests it. +`tests/integration/model_bridge/test_mamba2_adapter.py` (81 tests total), and the +`verify_models` benchmark suite now covers the SSM and hybrid families. Mamba-1, +Mamba-2, gated-delta-net (Qwen3.5 / Qwen3-Next), NemotronH, and GraniteMoeHybrid +all declare `applicable_phases = [1, 2, 3, 4]`, so their forward parity (P1, vs raw +HF), hook/cache coverage (P2/P3, which skip the HookedTransformer comparison SSMs +lack), and generation quality (P4) are benchmarked like any transformer. ## Credits diff --git a/docs/source/content/ssm_interpretability.md b/docs/source/content/ssm_interpretability.md new file mode 100644 index 000000000..f91d0bdaa --- /dev/null +++ b/docs/source/content/ssm_interpretability.md @@ -0,0 +1,117 @@ +# SSM / Recurrent-Model Interpretability + +TransformerBridge exposes a family-agnostic interpretability surface for state-space +and linear-attention models — Mamba-1, Mamba-2, gated-delta-net (Qwen3.5 / Qwen3-Next), +and the SSM hybrids (NemotronH, GraniteMoeHybrid). The same three `ActivationCache` +methods work across every family, dispatching to each block's SSM mixer regardless of +which slot it occupies (`.mixer` for Mamba, `.linear_attn` for gated-delta-net) and +skipping a hybrid's passthrough (attention / MLP / MoE) layers automatically. + +> These features are **bridge-only** — there is no `HookedTransformer` counterpart for +> SSM models. Load with `TransformerBridge.boot_transformers(...)`. + +--- + +## Discovery + +```python +bridge = TransformerBridge.boot_transformers("state-spaces/mamba-130m-hf") +_, cache = bridge.run_with_cache(tokens, use_cache=False) + +cache.ssm_layers() # -> [0, 1, 2, ...] block indices whose mixer is recurrent +``` + +`ssm_layers()` is purely structural (no dependence on `cfg.layers_block_type`): a +hybrid's attention/MLP layers are excluded, so on NemotronH you get only the Mamba +layers. + +## Read-only analysis (no forward re-run) + +Both methods reconstruct their quantities post-hoc from cached hooks. They return a +single tensor when *every* block is an SSM layer, else a `{layer_idx: tensor}` dict +over the SSM layers. + +```python +# Effective ("hidden") attention M = L ⊙ (C Bᵀ) [batch, heads, seq, seq] +M = cache.compute_ssm_effective_attention() # all SSM layers +M0 = cache.compute_ssm_effective_attention(layer=0) # one layer + +# Recurrent state trajectory S_t +S = cache.compute_ssm_state() # all SSM layers +S5 = cache.compute_ssm_state(layer=5, time_step=-1) # one layer, final step (memory-bounded) +``` + +State shape is family-specific (Mamba-1 is per-channel; Mamba-2 and gated-delta-net +are per-head), so `compute_ssm_state()` returns a dict except when one mixer type +covers every block. + +> **`use_cache=False` is required for gated-delta-net.** Its interior hooks +> (`hook_q/k/v`, `hook_beta`, `hook_log_decay`) fire only on the hooked prefill path; +> the default cached path exposes only `hook_in`/`hook_out`, and the reconstruction +> methods then raise. Mamba reconstructs from `in_proj`/`conv1d` hooks either way, but +> passing `use_cache=False` uniformly is safe. + +## Canonical hook vocabulary + +Every family exposes the same state-mutation names (defined once, in +`SSMStateHookMixin`), so interp tooling can address the same quantity across families: + +| Hook | Meaning | Kind | +|---|---|---| +| `hook_ssm_state` | post-scan recurrent-state trajectory `S_t` | real HookPoint (fires on the eager-scan path) | +| `hook_ssm_write` | per-step write influence | **real** for Mamba-1/2 (`dt·(x⊗B)`); **alias → `hook_beta`** for gated-delta-net (state-dependent write) | + +Additional canonical aliases resolve per family where the quantity exists: +`hook_ssm_out`, `hook_ssm_B`, `hook_ssm_C`, `hook_ssm_decay`, `hook_ssm_dt`. + +--- + +## Causal intervention: the `eager_scan` opt-in + +By default the mixer runs HF's fused recurrence kernel, which is opaque — the state +trajectory is never materialized, so it cannot be patched. Setting `eager_scan = True` +on a realized SSM mixer swaps the kernel for a readable Python scan that fires +`hook_ssm_state` (and, for the input-linear families, `hook_ssm_write`), so you can +read and edit the state mid-recurrence. + +```python +from transformer_lens.model_bridge.generalized_components.ssm_protocol import find_ssm_mixer + +# Enable on every realized SSM mixer (skips hybrid passthrough slots). +mixers = [m for b in bridge.blocks if (m := find_ssm_mixer(b)) is not None] +for m in mixers: + m.eager_scan = True + +# hook path uses the mixer's slot name: `.mixer` (Mamba) or `.linear_attn` (gated-delta-net) +STATE = "blocks.0.mixer.hook_ssm_state" + +def zero_state(state, hook): + return torch.zeros_like(state) # ablate the whole recurrent state + +logits = bridge.run_with_hooks(tokens, use_cache=False, fwd_hooks=[(STATE, zero_state)]) + +for m in mixers: + m.eager_scan = False # restore the default (bit-identical) fused path +``` + +**Two intervention semantics** (mirroring Mamba-Knockout and state-patching): + +- Patching **`hook_ssm_state`** changes only the *same-position* readout + (`y_t = C_t·S_t` or `S_t^T q_t`); it does not alter the forward recurrence. +- Patching **`hook_ssm_write`** re-runs the recurrence, so the edit **propagates** to + every later state. For gated-delta-net this is the write-strength gate + (`hook_beta`), so zeroing it suppresses all writes. + +### Caveats + +| | | +|---|---| +| **Prefill only** | The eager scan runs when `cache_params is None` (i.e. `use_cache=False`, no generation step). Autoregressive decode falls back to the fused kernel. | +| **Numerical** | The Python scan matches the fused kernel only to floating-point tolerance (≈1e-6 fp32), never bit-for-bit. | +| **Cost** | O(seq) Python with an O(batch·seq·…) state tensor — orders of magnitude slower/heavier than the kernel. Use short sequences. | +| **Default untouched** | With `eager_scan = False` (the default), `run_with_cache` is bit-identical to HF and `hook_ssm_state` does not fire. | + +> See [`ssm2_mixer.py`](../../../transformer_lens/model_bridge/generalized_components/ssm2_mixer.py), +> [`ssm_mixer.py`](../../../transformer_lens/model_bridge/generalized_components/ssm_mixer.py), and +> [`gated_delta_net.py`](../../../transformer_lens/model_bridge/generalized_components/gated_delta_net.py) +> for the per-family recurrences and reconstruction identities. diff --git a/docs/source/index.md b/docs/source/index.md index b10bd3a9d..72b90c4b0 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -56,6 +56,7 @@ content/citation content/contributing content/hook_system content/compatibility_mode +content/ssm_interpretability content/debugging_numerical_divergence generated/demos/Main_Demo generated/demos/Exploratory_Analysis_Demo diff --git a/tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py b/tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py new file mode 100644 index 000000000..f363e545b --- /dev/null +++ b/tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py @@ -0,0 +1,436 @@ +"""Integration tests for the GraniteMoeHybrid architecture adapter. + +GraniteMoeHybrid is a hybrid Mamba-2 + Attention + shared-MLP + sparse-MoE +model: each layer is *either* a Mamba SSM mixer *or* GQA attention, and every +layer additionally has a shared MLP and an optional MoE. Unlike NemotronH it is +a multi-slot transformer block (ln1/ln2/attn/mixer/shared_mlp/moe), not a +single-mixer SSMBlockBridge. + +Uses a tiny programmatic config with real (random) CPU weights — no network +access or weight downloads (mirrors tests/unit/model_bridge/test_gpt_oss_moe.py +and test_qwen3_moe_bridge.py for the from_config + direct-constructor pattern). + +Coverage focus (Phase 0, SSM mixer access normalization): +- The Mamba-2 mixer is reachable at the canonical ``.mixer`` slot on SSM layers. +- compute_effective_attention runs on a Granite SSM layer and reconstructs the + SSM output (proves dims are sourced from the wrapped HF mixer, not the shared + cfg whose n_heads/state_size hold the *attention* dims on a hybrid). +- The all-layers helper returns a per-SSM-layer dict on a heterogeneous hybrid. +- Attention / shared-MLP / MoE hooks still fire (rename did not disturb them). +""" + +import contextlib + +import pytest +import torch +from transformers import AutoModelForCausalLM +from transformers.models.granitemoehybrid import GraniteMoeHybridConfig + +from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.generalized_components import ( + GatedRMSNormBridge, + SSM2MixerBridge, +) +from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_config_from_hf, +) +from transformer_lens.model_bridge.supported_architectures.granite_moe_hybrid import ( + GraniteMoeHybridArchitectureAdapter, +) + +LAYER_TYPES = ["mamba", "attention", "mamba"] +MAMBA_LAYERS = [0, 2] +ATTN_LAYER = 1 + + +class _MockTokenizer: + """Stand-in to satisfy TransformerBridge(tokenizer=...) without a Hub call.""" + + pass + + +@pytest.fixture(scope="module") +def hf_model(): + cfg = GraniteMoeHybridConfig( + vocab_size=256, + hidden_size=64, + intermediate_size=32, + shared_intermediate_size=32, + num_hidden_layers=3, + num_attention_heads=8, + num_key_value_heads=4, + num_local_experts=4, + num_experts_per_tok=2, + max_position_embeddings=64, + layer_types=LAYER_TYPES, + mamba_n_heads=8, + mamba_n_groups=2, + mamba_d_state=16, + mamba_d_head=16, + mamba_d_conv=4, + mamba_expand=2, + mamba_chunk_size=16, + position_embedding_type="rope", + rope_parameters={"rope_type": "default", "rope_theta": 10000.0}, + ) + cfg.architectures = ["GraniteMoeHybridForCausalLM"] + torch.manual_seed(0) + return AutoModelForCausalLM.from_config(cfg).to(torch.float32).eval() + + +@pytest.fixture(scope="module") +def bridge(hf_model): + bridge_config = build_bridge_config_from_hf( + hf_model.config, "GraniteMoeHybridForCausalLM", "granitemoehybrid-tiny", torch.float32 + ) + adapter = GraniteMoeHybridArchitectureAdapter(bridge_config) + return TransformerBridge(model=hf_model, adapter=adapter, tokenizer=_MockTokenizer()) + + +@pytest.fixture(scope="module") +def tokens(): + return torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]]) + + +@pytest.fixture(scope="module") +def cache(bridge, tokens): + with torch.no_grad(): + _, c = bridge.run_with_cache(tokens) + return c + + +# --------------------------------------------------------------------------- +# Structure: Mamba mixer is reachable at the canonical .mixer slot +# --------------------------------------------------------------------------- + + +class TestGraniteMoeHybridStructure: + def test_block_count(self, bridge: TransformerBridge) -> None: + assert len(bridge.blocks) == len(LAYER_TYPES) + + def test_layers_block_type_surfaced(self, bridge: TransformerBridge) -> None: + assert getattr(bridge.cfg, "layers_block_type", None) == LAYER_TYPES + + def test_mamba_layers_expose_ssm2_mixer(self, bridge: TransformerBridge) -> None: + for i in MAMBA_LAYERS: + mixer = getattr(bridge.blocks[i], "mixer", None) + assert isinstance(mixer, SSM2MixerBridge), f"block {i} missing .mixer" + + def test_attention_layer_has_no_mixer_slot(self, bridge: TransformerBridge) -> None: + # The mamba bridge is optional and skipped on attention layers, so the + # structural `.mixer` slot must be absent there (blocks_with relies on this). + assert "mixer" not in bridge.blocks[ATTN_LAYER]._modules + assert "attn" in bridge.blocks[ATTN_LAYER]._modules + + def test_inner_norm_is_gated(self, bridge: TransformerBridge) -> None: + mixer = bridge.blocks[MAMBA_LAYERS[0]].mixer + assert isinstance(mixer.inner_norm, GatedRMSNormBridge) + + +# --------------------------------------------------------------------------- +# Forward parity: bridge delegates fully, so logits match HF exactly +# --------------------------------------------------------------------------- + + +class TestGraniteMoeHybridForwardPass: + def test_forward_matches_hf_exactly( + self, bridge: TransformerBridge, hf_model, tokens: torch.Tensor + ) -> None: + with torch.no_grad(): + bridge_out = bridge(tokens) + hf_out = hf_model(tokens).logits + max_diff = (bridge_out.float() - hf_out.float()).abs().max().item() + assert max_diff == 0.0, ( + f"Bridge vs HF forward max diff = {max_diff:.2e}. Expected 0 because " + "the bridge delegates the full forward to HF." + ) + + +# --------------------------------------------------------------------------- +# Hook coverage: rename to .mixer did not disturb attn / shared_mlp / moe hooks +# --------------------------------------------------------------------------- + + +class TestGraniteMoeHybridHookCoverage: + def test_mamba_mixer_submodule_hooks_fire(self, cache) -> None: + for i in MAMBA_LAYERS: + for submod in ("in_proj", "conv1d"): + assert f"blocks.{i}.mixer.{submod}.hook_out" in cache + assert f"blocks.{i}.mixer.inner_norm.hook_in" in cache + + def test_attention_hooks_fire_on_attention_layer(self, cache) -> None: + assert f"blocks.{ATTN_LAYER}.attn.hook_out" in cache + + def test_shared_mlp_and_moe_hooks_fire(self, cache) -> None: + # Every layer has a shared MLP and a sparse MoE; both must still hook. + assert f"blocks.{MAMBA_LAYERS[0]}.shared_mlp.hook_out" in cache + assert f"blocks.{MAMBA_LAYERS[0]}.moe.hook_out" in cache + + +# --------------------------------------------------------------------------- +# Effective attention: the Phase 0 acceptance surface +# --------------------------------------------------------------------------- + + +class TestGraniteMoeHybridEffectiveAttention: + def test_single_layer_shape(self, bridge: TransformerBridge, cache, tokens) -> None: + seq_len = tokens.shape[1] + mixer = bridge.blocks[MAMBA_LAYERS[0]].mixer + n_heads = mixer.original_component.num_heads + M = cache.compute_ssm_effective_attention(layer=MAMBA_LAYERS[0]) + assert isinstance(M, torch.Tensor) + assert M.shape == (1, n_heads, seq_len, seq_len) + assert torch.isfinite(M).all() + + def test_single_layer_is_causal(self, bridge: TransformerBridge, cache) -> None: + M = cache.compute_ssm_effective_attention(layer=MAMBA_LAYERS[0]) + seq_len = M.shape[-1] + upper = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool), diagonal=1) + assert torch.all(M[..., upper] == 0), "effective attention must be lower-triangular" + + def test_all_layers_returns_per_ssm_layer_dict(self, bridge: TransformerBridge, cache) -> None: + # Heterogeneous hybrid -> dict keyed by SSM layer index (attention layer skipped). + M_all = cache.compute_ssm_effective_attention() + assert isinstance(M_all, dict) + assert sorted(M_all.keys()) == MAMBA_LAYERS + for idx, M in M_all.items(): + assert torch.equal(M, cache.compute_ssm_effective_attention(layer=idx)) + + def test_attention_layer_index_raises_typeerror(self, bridge: TransformerBridge, cache) -> None: + with pytest.raises(TypeError): + cache.compute_ssm_effective_attention(layer=ATTN_LAYER) + + def test_reconstruction_matches_ssm_output( + self, bridge: TransformerBridge, cache, tokens + ) -> None: + """M @ x + D*x must reconstruct the SSM output (inner_norm.hook_in). + + Non-tautological: HF computes the SSM output via an independent chunked + path. Agreement to float32 precision proves compute_effective_attention + sourced the correct *Mamba* dims from the wrapped module — on a hybrid the + shared cfg holds attention dims, so a cfg-based read would produce garbage. + """ + layer = MAMBA_LAYERS[0] + seq_len = tokens.shape[1] + mixer = bridge.blocks[layer].mixer + oc = mixer.original_component + num_heads, head_dim = oc.num_heads, oc.head_dim + intermediate_size, n_groups, state_size = ( + oc.intermediate_size, + oc.n_groups, + oc.ssm_state_size, + ) + + M_full = cache.compute_ssm_effective_attention(layer=layer, include_dt_scaling=True) + + conv_out = cache[f"blocks.{layer}.mixer.conv1d.hook_out"][..., :seq_len].float() + conv_activated = torch.nn.functional.silu(conv_out).transpose(1, 2) + hidden_x, _, _ = conv_activated.split( + [intermediate_size, n_groups * state_size, n_groups * state_size], dim=-1 + ) + batch = hidden_x.shape[0] + x_per_head = hidden_x.view(batch, seq_len, num_heads, head_dim) + + D = mixer.D.float() + y_pred = torch.einsum("bhij,bjhd->bihd", M_full, x_per_head) + y_pred = y_pred + D[None, None, :, None] * x_per_head + y_pred_flat = y_pred.reshape(batch, seq_len, -1) + + y_actual = cache[f"blocks.{layer}.mixer.inner_norm.hook_in"].float() + max_diff = (y_actual - y_pred_flat).abs().max().item() + scale = max(y_actual.abs().max().item(), 1e-8) + assert max_diff / scale < 1e-4, ( + f"Reconstruction mismatch: rel {max_diff / scale:.2e}. The effective " + "attention dims are inconsistent with HF's SSM output (likely reading " + "attention dims off the shared cfg instead of the mamba module)." + ) + + +class TestGraniteMoeHybridSSMState: + """compute_ssm_state works across the hybrid: per-SSM-layer dict, correct dims.""" + + def test_single_layer_shape(self, bridge: TransformerBridge, cache, tokens) -> None: + seq_len = tokens.shape[1] + oc = bridge.blocks[MAMBA_LAYERS[0]].mixer.original_component + S = cache.compute_ssm_state(layer=MAMBA_LAYERS[0]) + assert isinstance(S, torch.Tensor) + assert S.shape == (1, oc.num_heads, seq_len, oc.head_dim, oc.ssm_state_size) + assert torch.isfinite(S).all() + + def test_all_layers_returns_per_ssm_layer_dict(self, bridge: TransformerBridge, cache) -> None: + S_all = cache.compute_ssm_state() + assert isinstance(S_all, dict) + assert sorted(S_all.keys()) == MAMBA_LAYERS + + def test_attention_layer_raises_typeerror(self, bridge: TransformerBridge, cache) -> None: + with pytest.raises(TypeError): + cache.compute_ssm_state(layer=ATTN_LAYER) + + def test_reconstructs_ssm_output(self, bridge: TransformerBridge, cache, tokens) -> None: + """y = C·S + D·x reconstructs HF's SSM output — proves hybrid dims are right.""" + layer = MAMBA_LAYERS[0] + seq_len = tokens.shape[1] + mixer = bridge.blocks[layer].mixer + oc = mixer.original_component + nh, hd, ns, ng = oc.num_heads, oc.head_dim, oc.ssm_state_size, oc.n_groups + + S = cache.compute_ssm_state(layer=layer) + conv = cache[f"blocks.{layer}.mixer.conv1d.hook_out"][..., :seq_len].float() + conv_act = torch.nn.functional.silu(conv).transpose(1, 2) + x_flat, _, C_flat = conv_act.split([oc.intermediate_size, ng * ns, ng * ns], dim=-1) + x = x_flat.view(1, seq_len, nh, hd) + C_h = C_flat.view(1, seq_len, ng, ns).repeat_interleave(nh // ng, dim=2) + D = mixer.D.float() + + y = torch.einsum("bthn,bhtpn->bthp", C_h, S) + D[None, None, :, None] * x + y_actual = cache[f"blocks.{layer}.mixer.inner_norm.hook_in"].float() + max_diff = (y_actual - y.reshape(1, seq_len, -1)).abs().max().item() + scale = max(y_actual.abs().max().item(), 1e-8) + assert max_diff / scale < 1e-4, f"state reconstruction rel diff {max_diff / scale:.2e}" + + +@contextlib.contextmanager +def _granite_eager(bridge): + """Enable the eager-scan intervention path on Granite's Mamba-2 mixer layers.""" + for i in MAMBA_LAYERS: + bridge.blocks[i].mixer.eager_scan = True + try: + yield bridge + finally: + for i in MAMBA_LAYERS: + bridge.blocks[i].mixer.eager_scan = False + + +def _build_granite_bridge(device="cpu"): + """Fresh tiny Granite bridge on the given device (for device-parametrized tests).""" + cfg = GraniteMoeHybridConfig( + vocab_size=256, + hidden_size=64, + intermediate_size=32, + shared_intermediate_size=32, + num_hidden_layers=3, + num_attention_heads=8, + num_key_value_heads=4, + num_local_experts=4, + num_experts_per_tok=2, + max_position_embeddings=64, + layer_types=LAYER_TYPES, + mamba_n_heads=8, + mamba_n_groups=2, + mamba_d_state=16, + mamba_d_head=16, + mamba_d_conv=4, + mamba_expand=2, + mamba_chunk_size=16, + position_embedding_type="rope", + rope_parameters={"rope_type": "default", "rope_theta": 10000.0}, + ) + cfg.architectures = ["GraniteMoeHybridForCausalLM"] + torch.manual_seed(0) + hf = AutoModelForCausalLM.from_config(cfg).to(torch.float32).eval() + bridge_config = build_bridge_config_from_hf( + hf.config, "GraniteMoeHybridForCausalLM", "granitemoehybrid-tiny", torch.float32 + ) + bridge = TransformerBridge( + model=hf, + adapter=GraniteMoeHybridArchitectureAdapter(bridge_config), + tokenizer=_MockTokenizer(), + ) + if device != "cpu": + hf.to(device) + return bridge + + +def _available_devices(): + devices = ["cpu"] + if torch.cuda.is_available(): + devices.append("cuda") + if torch.backends.mps.is_available(): + devices.append("mps") + return devices + + +class TestGraniteMoeHybridProcessedParity: + """Phase-3 parity guard: processed (compat-mode) bridge vs raw HF via log_softmax. + + The forward test only checks UNPROCESSED delegation (==0.0); this pins the PROCESSED + path so a compat-mode regression is caught in CI without the full-size checkpoint. + Granite is RMSNorm + fold_ln-off, so the centering rewrite must be output-invariant. + """ + + def test_compat_mode_logits_match_raw_hf(self): + bridge = _build_granite_bridge() + tokens = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]]) + with torch.no_grad(): + raw = bridge(tokens) + bridge.enable_compatibility_mode(disable_warnings=True) + with torch.no_grad(): + proc = bridge(tokens) + raw_lsm = torch.log_softmax(raw.float(), dim=-1) + proc_lsm = torch.log_softmax(proc.float(), dim=-1) + max_diff = (raw_lsm - proc_lsm).abs().max().item() + assert max_diff < 1e-4, ( + f"compat-mode processed vs raw-HF log_softmax max_diff {max_diff:.2e} — the " + "weight-centering rewrite must be output-invariant in fp32 (P3 logits_equivalence)." + ) + + +class TestGraniteMoeHybridEagerScanIntervention: + """Phase 4 eager-scan intervention on Granite's Mamba-2 mixer layers (hybrid): + hooks fire on the Mamba layers only, interventions propagate to logits, and the + default path is untouched. Eager scan needs use_cache=False (prefill).""" + + def test_default_path_has_no_eager_hooks(self, bridge, tokens): + with torch.no_grad(): + _, c = bridge.run_with_cache(tokens) + assert "blocks.0.mixer.hook_ssm_write" not in c + + def test_eager_hooks_fire_on_mamba_layers_only(self, bridge, tokens): + with _granite_eager(bridge), torch.no_grad(): + _, c = bridge.run_with_cache(tokens, use_cache=False) + for i in MAMBA_LAYERS: + assert f"blocks.{i}.mixer.hook_ssm_write" in c + assert f"blocks.{i}.mixer.hook_ssm_state" in c + assert f"blocks.{ATTN_LAYER}.mixer.hook_ssm_write" not in c # attention layer: no mixer + + def test_eager_scan_matches_fused(self, bridge, tokens): + with torch.no_grad(): + fused = bridge(tokens) + with _granite_eager(bridge), torch.no_grad(): + eager = bridge.run_with_hooks(tokens, use_cache=False, fwd_hooks=[]) + rel = (eager - fused).abs().max().item() / max(fused.abs().max().item(), 1e-8) + assert rel < 1e-4, f"Granite eager vs fused rel diff {rel:.2e}" + + def test_write_knockout_changes_logits(self, bridge, tokens): + def knockout(writes, hook): + writes = writes.clone() + writes[:, 2] = 0.0 + return writes + + with _granite_eager(bridge), torch.no_grad(): + base = bridge.run_with_hooks(tokens, use_cache=False, fwd_hooks=[]) + patched = bridge.run_with_hooks( + tokens, use_cache=False, fwd_hooks=[("blocks.0.mixer.hook_ssm_write", knockout)] + ) + assert (patched - base).abs().max().item() > 1e-6 + + +@pytest.mark.parametrize("device", _available_devices()) +def test_granite_eager_scan_device_correctness(device): + """Regression for the eager-scan device fix: ssm_state must be created on the + input's device. Parametrized over every available device (cpu + cuda/mps).""" + bridge = _build_granite_bridge(device) + toks = torch.tensor([[1, 2, 3, 4, 5]], device=device) + with torch.no_grad(): + fused = bridge(toks) + for i in MAMBA_LAYERS: + bridge.blocks[i].mixer.eager_scan = True + with torch.no_grad(): + _, cache = bridge.run_with_cache(toks, use_cache=False) + eager = bridge.run_with_hooks(toks, use_cache=False, fwd_hooks=[]) + + dev_type = torch.device(device).type + assert eager.device.type == dev_type + assert cache["blocks.0.mixer.hook_ssm_write"].device.type == dev_type + rel = (eager.cpu() - fused.cpu()).abs().max().item() / max(fused.abs().max().item(), 1e-8) + assert rel < 1e-4, f"eager vs fused rel diff on {device}: {rel:.2e}" diff --git a/tests/integration/model_bridge/test_mamba2_adapter.py b/tests/integration/model_bridge/test_mamba2_adapter.py index cb7f46dd0..7e5b1668f 100644 --- a/tests/integration/model_bridge/test_mamba2_adapter.py +++ b/tests/integration/model_bridge/test_mamba2_adapter.py @@ -16,6 +16,8 @@ mutated by the cache machinery. """ +import contextlib + import pytest import torch @@ -390,12 +392,9 @@ def test_different_layers_produce_different_attention(self, mamba2_bridge): assert not torch.allclose(M0, M5) -class TestMamba2EffectiveAttentionHelper: - """Tests for the mamba2.compute_effective_attention helper function. - - The helper wraps the mixer method so callers don't repeat the layer index - and can request all layers at once. - """ +class TestMamba2EffectiveAttentionDispatch: + """cache.compute_ssm_effective_attention wraps the mixer method so callers + don't repeat the layer index and can request all layers at once.""" @pytest.fixture(scope="class") def bridge_cache(self, mamba2_bridge): @@ -405,37 +404,22 @@ def bridge_cache(self, mamba2_bridge): return mamba2_bridge, cache def test_single_layer_matches_mixer_method(self, bridge_cache): - """helper(bridge, cache, layer=i) must match the underlying mixer call.""" - from transformer_lens.model_bridge.supported_architectures.mamba2 import ( - compute_effective_attention, - ) - bridge, cache = bridge_cache for layer in [0, 5, 23]: - M_helper = compute_effective_attention(bridge, cache, layer=layer) + M_helper = cache.compute_ssm_effective_attention(layer=layer) M_direct = bridge.blocks[layer].mixer.compute_effective_attention( cache, layer_idx=layer ) assert torch.equal(M_helper, M_direct) def test_all_layers_shape(self, bridge_cache, mamba2_bridge): - """helper(bridge, cache) with layer=None stacks all layers.""" - from transformer_lens.model_bridge.supported_architectures.mamba2 import ( - compute_effective_attention, - ) - - bridge, cache = bridge_cache - M_all = compute_effective_attention(bridge, cache) + _, cache = bridge_cache + M_all = cache.compute_ssm_effective_attention() assert M_all.shape == (mamba2_bridge.cfg.n_layers, 1, mamba2_bridge.cfg.n_heads, 5, 5) def test_all_layers_matches_individual_calls(self, bridge_cache): - """Stacked all-layers tensor must match per-layer calls along dim 0.""" - from transformer_lens.model_bridge.supported_architectures.mamba2 import ( - compute_effective_attention, - ) - bridge, cache = bridge_cache - M_all = compute_effective_attention(bridge, cache) + M_all = cache.compute_ssm_effective_attention() for layer in [0, 12, 23]: M_single = bridge.blocks[layer].mixer.compute_effective_attention( cache, layer_idx=layer @@ -443,15 +427,20 @@ def test_all_layers_matches_individual_calls(self, bridge_cache): assert torch.equal(M_all[layer], M_single) def test_include_dt_scaling_propagates(self, bridge_cache): - """The include_dt_scaling flag must flow through the helper to the mixer.""" + _, cache = bridge_cache + M_att = cache.compute_ssm_effective_attention(layer=0, include_dt_scaling=False) + M_full = cache.compute_ssm_effective_attention(layer=0, include_dt_scaling=True) + assert not torch.allclose(M_att, M_full) + + def test_deprecated_module_wrapper_warns_and_delegates(self, bridge_cache): from transformer_lens.model_bridge.supported_architectures.mamba2 import ( compute_effective_attention, ) bridge, cache = bridge_cache - M_att = compute_effective_attention(bridge, cache, layer=0, include_dt_scaling=False) - M_full = compute_effective_attention(bridge, cache, layer=0, include_dt_scaling=True) - assert not torch.allclose(M_att, M_full) + with pytest.warns(DeprecationWarning): + M_dep = compute_effective_attention(bridge, cache, layer=0) + assert torch.equal(M_dep, cache.compute_ssm_effective_attention(layer=0)) class TestMamba2ParameterAccess: @@ -562,3 +551,263 @@ def test_stop_returns_residual(self, mamba2_bridge): expected = cache["blocks.7.hook_in"] assert torch.allclose(stopped, expected) assert stopped.shape == (1, 4, mamba2_bridge.cfg.d_model) + + +class TestMamba2SSMState: + """compute_ssm_state reconstructs the recurrent state S read-only from cache. + + S_t = dA_t · S_{t-1} + dt_t · (x_t ⊗ B_t). Two independent checks: a naive + fp64 eager recurrence (validates the vectorized cumsum/einsum), and the + y = C·S + D·x reconstruction against HF's chunked SSD output (validates the + whole reconstruction against a numerically-independent path). + """ + + SEQ_LEN = 8 + + @pytest.fixture(scope="class") + def cache(self, mamba2_bridge): + tokens = torch.arange(1, self.SEQ_LEN + 1).unsqueeze(0) + with torch.no_grad(): + _, cache = mamba2_bridge.run_with_cache(tokens) + return cache + + @staticmethod + def _ssd_inputs(cache, mixer, layer, seq_len): + """Re-extract dt, per-head x/B/C, A, D from cache — independent of the impl.""" + oc = mixer.original_component + nh, hd, ns, ng = oc.num_heads, oc.head_dim, oc.ssm_state_size, oc.n_groups + in_proj = cache[f"blocks.{layer}.mixer.in_proj.hook_out"].float() + conv = cache[f"blocks.{layer}.mixer.conv1d.hook_out"][..., :seq_len].float() + dt = torch.nn.functional.softplus(in_proj[..., -nh:] + mixer.dt_bias.float()) + dt = torch.clamp(dt, float(oc.time_step_limit[0]), float(oc.time_step_limit[1])) + conv_act = torch.nn.functional.silu(conv).transpose(1, 2) + x_flat, B_flat, C_flat = conv_act.split([oc.intermediate_size, ng * ns, ng * ns], dim=-1) + x = x_flat.view(1, seq_len, nh, hd) + B_h = B_flat.view(1, seq_len, ng, ns).repeat_interleave(nh // ng, dim=2) + C_h = C_flat.view(1, seq_len, ng, ns).repeat_interleave(nh // ng, dim=2) + A = -torch.exp(mixer.A_log.float()) + return dt, x, B_h, C_h, A, mixer.D.float() + + def test_shape(self, cache, mamba2_bridge): + mixer = mamba2_bridge.blocks[0].mixer + oc = mixer.original_component + S = mixer.compute_ssm_state(cache, layer_idx=0) + assert S.shape == (1, oc.num_heads, self.SEQ_LEN, oc.head_dim, oc.ssm_state_size) + assert torch.isfinite(S).all() + + def test_matches_fp64_eager_recurrence(self, cache, mamba2_bridge): + """The vectorized state must match a naive fp64 step-by-step recurrence.""" + mixer = mamba2_bridge.blocks[0].mixer + S = mixer.compute_ssm_state(cache, layer_idx=0) + + dt, x, B_h, _, A, _ = self._ssd_inputs(cache, mixer, 0, self.SEQ_LEN) + dt, x, B_h, A = dt.double(), x.double(), B_h.double(), A.double() + b, nh = S.shape[0], S.shape[1] + hd, ns = S.shape[3], S.shape[4] + state = torch.zeros(b, nh, hd, ns, dtype=torch.float64) + ref = torch.zeros(b, nh, self.SEQ_LEN, hd, ns, dtype=torch.float64) + for t in range(self.SEQ_LEN): + dA = torch.exp(dt[:, t, :] * A[None, :]) # [b, nh] + write = dt[:, t, :, None, None] * x[:, t, :, :, None] * B_h[:, t, :, None, :] + state = dA[:, :, None, None] * state + write + ref[:, :, t] = state + + max_diff = (S.double() - ref).abs().max().item() + scale = max(ref.abs().max().item(), 1e-8) + assert ( + max_diff / scale < 1e-5 + ), f"S vs fp64 eager recurrence rel diff {max_diff / scale:.2e}" + + def test_reconstructs_ssm_output(self, cache, mamba2_bridge): + """y = C·S + D·x must reconstruct HF's SSM output (inner_norm.hook_in).""" + mixer = mamba2_bridge.blocks[0].mixer + S = mixer.compute_ssm_state(cache, layer_idx=0) + _, x, _, C_h, _, D = self._ssd_inputs(cache, mixer, 0, self.SEQ_LEN) + + y = torch.einsum("bthn,bhtpn->bthp", C_h, S) + D[None, None, :, None] * x + y_pred = y.reshape(1, self.SEQ_LEN, -1) + y_actual = cache["blocks.0.mixer.inner_norm.hook_in"].float() + max_diff = (y_actual - y_pred).abs().max().item() + scale = max(y_actual.abs().max().item(), 1e-8) + assert max_diff / scale < 1e-5, ( + f"y=C·S+D·x reconstruction rel diff {max_diff / scale:.2e}; the state " + "is inconsistent with HF's independently-computed SSM output." + ) + + def test_time_step_matches_full_slice(self, cache, mamba2_bridge): + mixer = mamba2_bridge.blocks[0].mixer + S = mixer.compute_ssm_state(cache, layer_idx=0) + S_t = mixer.compute_ssm_state(cache, layer_idx=0, time_step=3) + assert S_t.shape == (S.shape[0], S.shape[1], S.shape[3], S.shape[4]) + assert torch.equal(S_t, S[:, :, 3]) + + def test_cache_method_matches_mixer(self, cache, mamba2_bridge): + S_mixer = mamba2_bridge.blocks[0].mixer.compute_ssm_state(cache, layer_idx=0) + assert torch.equal(cache.compute_ssm_state(layer=0), S_mixer) + + def test_cache_all_layers_stacks(self, cache, mamba2_bridge): + # Pure Mamba-2 is homogeneous, so all-layers stacks along a new dim 0. + S_all = cache.compute_ssm_state() + assert torch.is_tensor(S_all) + assert S_all.shape[0] == mamba2_bridge.cfg.n_layers + + def test_deprecated_module_wrapper_warns_and_delegates(self, cache, mamba2_bridge): + from transformer_lens.model_bridge.supported_architectures.mamba2 import ( + compute_ssm_state, + ) + + S_mixer = mamba2_bridge.blocks[0].mixer.compute_ssm_state(cache, layer_idx=0) + with pytest.warns(DeprecationWarning): + S_dep = compute_ssm_state(mamba2_bridge, cache, layer=0) + assert torch.equal(S_dep, S_mixer) + + def test_raises_on_empty_cache(self, mamba2_bridge): + from transformer_lens.ActivationCache import ActivationCache + + empty = ActivationCache({}, model=mamba2_bridge) + with pytest.raises(RuntimeError, match="in cache"): + mamba2_bridge.blocks[0].mixer.compute_ssm_state(empty, layer_idx=0) + + +@contextlib.contextmanager +def _eager_scan(bridge): + """Enable the opt-in eager-scan intervention path on every mixer, then reset.""" + for blk in bridge.blocks: + blk.mixer.eager_scan = True + try: + yield bridge + finally: + for blk in bridge.blocks: + blk.mixer.eager_scan = False + + +class TestMamba2EagerScanIntervention: + """Phase 4: opt-in eager-scan path exposes hook_ssm_write / hook_ssm_state for + interventions (write-knockout, state-patch) that propagate to logits, while the + default run_with_cache path is untouched. Eager scan requires use_cache=False + (so cache_params is None — prefill only).""" + + TOKENS = torch.tensor([[1, 2, 3, 4, 5, 6]]) + + def test_default_path_has_no_eager_hooks(self, mamba2_bridge): + with torch.no_grad(): + _, cache = mamba2_bridge.run_with_cache(self.TOKENS) + assert "blocks.0.mixer.hook_ssm_write" not in cache + assert "blocks.0.mixer.hook_ssm_state" not in cache + + def test_eager_requires_use_cache_false(self, mamba2_bridge): + """With eager on but a default (stateful) forward, cache_params is present + so the eager path is skipped and the hooks do not fire.""" + with _eager_scan(mamba2_bridge), torch.no_grad(): + _, cache = mamba2_bridge.run_with_cache(self.TOKENS) # default use_cache + assert "blocks.0.mixer.hook_ssm_write" not in cache + + def test_eager_scan_matches_fused_logits(self, mamba2_bridge): + with torch.no_grad(): + fused = mamba2_bridge(self.TOKENS) + with _eager_scan(mamba2_bridge), torch.no_grad(): + eager = mamba2_bridge(self.TOKENS, use_cache=False) + rel = (eager - fused).abs().max().item() / max(fused.abs().max().item(), 1e-8) + assert rel < 1e-4, f"eager scan vs fused kernel rel diff {rel:.2e}" + + def test_eager_scan_matches_fp64_step_reference(self, mamba2_bridge): + """The eager scan's write/state must match an independent fp64 step recurrence. + + ``test_eager_scan_matches_fused_logits`` pins the eager scan to HF's fused + kernel — but a shared discretization error would pass both. This pins the + eager scan's own ``hook_ssm_write`` / ``hook_ssm_state`` to a naive fp64 + step-by-step recurrence built independently from the cached in_proj/conv1d + outputs (``S_t = dA_t·S_{t-1} + dt_t·(x_t⊗B_t)``), so it is a genuine + correctness gate on the eager scan, not just kernel agreement. + """ + seq = self.TOKENS.shape[1] + with _eager_scan(mamba2_bridge), torch.no_grad(): + _, cache = mamba2_bridge.run_with_cache(self.TOKENS, use_cache=False) + mixer = mamba2_bridge.blocks[0].mixer + write_hook = cache["blocks.0.mixer.hook_ssm_write"].double() + state_hook = cache["blocks.0.mixer.hook_ssm_state"].double() + + dt, x, B_h, _, A, _ = TestMamba2SSMState._ssd_inputs(cache, mixer, 0, seq) + dt, x, B_h, A = dt.double(), x.double(), B_h.double(), A.double() + b, nh, hd, ns = ( + state_hook.shape[0], + state_hook.shape[2], + state_hook.shape[3], + state_hook.shape[4], + ) + + state = torch.zeros(b, nh, hd, ns, dtype=torch.float64) + ref_write = torch.zeros(b, seq, nh, hd, ns, dtype=torch.float64) + ref_state = torch.zeros(b, seq, nh, hd, ns, dtype=torch.float64) + for t in range(seq): + dA = torch.exp(dt[:, t, :] * A[None, :]) # [b, nh] + write = dt[:, t, :, None, None] * x[:, t, :, :, None] * B_h[:, t, :, None, :] + ref_write[:, t] = write + state = dA[:, :, None, None] * state + write + ref_state[:, t] = state + + for name, got, ref in (("write", write_hook, ref_write), ("state", state_hook, ref_state)): + rel = (got - ref).abs().max().item() / max(ref.abs().max().item(), 1e-8) + assert rel < 1e-5, f"eager {name} vs fp64 step reference rel diff {rel:.2e}" + + def test_eager_scan_matches_fused_padded_batch(self, mamba2_bridge): + """Padded batch: the eager path must mirror HF's padding mask (applied both + before in_proj and after the conv) to stay at parity.""" + ids = torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 0]]) + mask = torch.tensor([[1, 1, 1, 1, 1], [1, 1, 1, 1, 0]]) + with torch.no_grad(): + fused = mamba2_bridge(ids, attention_mask=mask) + with _eager_scan(mamba2_bridge), torch.no_grad(): + eager = mamba2_bridge(ids, attention_mask=mask, use_cache=False) + rel = (eager - fused).abs().max().item() / max(fused.abs().max().item(), 1e-8) + assert rel < 1e-4, f"padded-batch eager vs fused rel diff {rel:.2e}" + + def test_eager_hooks_fire_with_use_cache_false(self, mamba2_bridge): + with _eager_scan(mamba2_bridge), torch.no_grad(): + _, cache = mamba2_bridge.run_with_cache(self.TOKENS, use_cache=False) + oc = mamba2_bridge.blocks[0].mixer.original_component + seq = self.TOKENS.shape[1] + expected = (1, seq, oc.num_heads, oc.head_dim, oc.ssm_state_size) + assert cache["blocks.0.mixer.hook_ssm_write"].shape == expected + assert cache["blocks.0.mixer.hook_ssm_state"].shape == expected + + def test_write_knockout_changes_logits(self, mamba2_bridge): + with _eager_scan(mamba2_bridge), torch.no_grad(): + base = mamba2_bridge(self.TOKENS, use_cache=False) + + def knockout(writes, hook): + writes = writes.clone() + writes[:, 2] = 0.0 # input position 2 writes nothing — column knockout + return writes + + patched = mamba2_bridge.run_with_hooks( + self.TOKENS, + use_cache=False, + fwd_hooks=[("blocks.0.mixer.hook_ssm_write", knockout)], + ) + assert (patched - base).abs().max().item() > 1e-6 + + def test_state_patch_changes_logits(self, mamba2_bridge): + with _eager_scan(mamba2_bridge), torch.no_grad(): + base = mamba2_bridge(self.TOKENS, use_cache=False) + + def patch(state, hook): + state = state.clone() + state[:, 3] = 0.0 + return state + + patched = mamba2_bridge.run_with_hooks( + self.TOKENS, + use_cache=False, + fwd_hooks=[("blocks.0.mixer.hook_ssm_state", patch)], + ) + assert (patched - base).abs().max().item() > 1e-6 + + def test_disabling_restores_fused_path(self, mamba2_bridge): + with torch.no_grad(): + fused_before = mamba2_bridge(self.TOKENS) + with _eager_scan(mamba2_bridge), torch.no_grad(): + mamba2_bridge(self.TOKENS, use_cache=False) + with torch.no_grad(): + fused_after = mamba2_bridge(self.TOKENS) + assert torch.equal(fused_before, fused_after) diff --git a/tests/integration/model_bridge/test_mamba_adapter.py b/tests/integration/model_bridge/test_mamba_adapter.py index 3f72449a4..32cae35dd 100644 --- a/tests/integration/model_bridge/test_mamba_adapter.py +++ b/tests/integration/model_bridge/test_mamba_adapter.py @@ -18,6 +18,8 @@ those hooks MUST `.clone()` captured state tensors. """ +import contextlib + import pytest import torch @@ -330,3 +332,498 @@ def perturb(t, hook): "Mid-layer mutation had no effect on generation — the stateful " "loop may be bypassing the bridge's forward() path." ) + + +class TestMamba1EffectiveAttention: + """Mamba-1's S6 scan as per-channel effective attention M. + + Two checks: M·x + D·x matches an fp64 eager recurrence (matrix-algebra + correctness), and the full mixer output reconstructed through gate + out_proj + matches HF's cached hook_out (ties M to HF's numerically-independent forward). + """ + + SEQ_LEN = 6 + + @pytest.fixture(scope="class") + def cache(self, mamba_bridge): + tokens = torch.arange(1, self.SEQ_LEN + 1).unsqueeze(0) + with torch.no_grad(): + _, cache = mamba_bridge.run_with_cache(tokens) + return cache + + @staticmethod + def _inputs(cache, mixer, layer, seq_len): + """Re-extract x (post-conv SiLU), B, C, dt, A, D, gate — independent of the impl.""" + import torch.nn.functional as F + + oc = mixer.original_component + d_inner, state, dt_rank = oc.intermediate_size, oc.ssm_state_size, oc.time_step_rank + conv = cache[f"blocks.{layer}.mixer.conv1d.hook_out"][..., :seq_len].float() + x_proj = cache[f"blocks.{layer}.mixer.x_proj.hook_out"].float() + dt_proj = cache[f"blocks.{layer}.mixer.dt_proj.hook_out"].float() + in_proj = cache[f"blocks.{layer}.mixer.in_proj.hook_out"].float() + x = oc.act(conv) # [b, d_inner, seq] + _ts, B, C = x_proj.split([dt_rank, state, state], dim=-1) + dt = F.softplus(dt_proj).transpose(1, 2) # [b, d_inner, seq] + A = -torch.exp(mixer.A_log.float()) + gate = in_proj.transpose(1, 2)[:, d_inner:, :] # [b, d_inner, seq] + return x, B, C, dt, A, mixer.D.float(), gate + + def test_shape(self, cache, mamba_bridge): + mixer = mamba_bridge.blocks[0].mixer + oc = mixer.original_component + M = mixer.compute_effective_attention(cache, layer_idx=0) + assert M.shape == (1, oc.intermediate_size, self.SEQ_LEN, self.SEQ_LEN) + assert torch.isfinite(M).all() + + def test_per_state_coord_sums_to_default(self, cache, mamba_bridge): + mixer = mamba_bridge.blocks[0].mixer + oc = mixer.original_component + M = mixer.compute_effective_attention(cache, layer_idx=0, include_dt_scaling=True) + M_coord = mixer.compute_effective_attention( + cache, layer_idx=0, include_dt_scaling=True, per_state_coord=True + ) + assert M_coord.shape == ( + 1, + oc.intermediate_size, + oc.ssm_state_size, + self.SEQ_LEN, + self.SEQ_LEN, + ) + assert torch.allclose(M_coord.sum(dim=2), M, atol=1e-6) + + def test_causal(self, cache, mamba_bridge): + M = mamba_bridge.blocks[0].mixer.compute_effective_attention(cache, layer_idx=0) + upper = torch.triu(torch.ones(self.SEQ_LEN, self.SEQ_LEN, dtype=torch.bool), diagonal=1) + assert torch.all(M[..., upper] == 0), "effective attention must be causal" + + def test_matches_fp64_eager_recurrence(self, cache, mamba_bridge): + """M·x + D·x must match a naive fp64 step-by-step S6 recurrence.""" + mixer = mamba_bridge.blocks[0].mixer + M = mixer.compute_effective_attention(cache, layer_idx=0, include_dt_scaling=True) + x, B, C, dt, A, D, _ = self._inputs(cache, mixer, 0, self.SEQ_LEN) + y_pred = torch.einsum("bcij,bcj->bci", M, x) + D[None, :, None] * x + + A, B, C, x, dt = A.double(), B.double(), C.double(), x.double(), dt.double() + b, d_inner = x.shape[0], x.shape[1] + ssm = torch.zeros(b, d_inner, A.shape[1], dtype=torch.float64) + y_ref = torch.zeros(b, d_inner, self.SEQ_LEN, dtype=torch.float64) + for i in range(self.SEQ_LEN): + dA = torch.exp(A[None] * dt[:, :, i, None]) # [b, d_inner, state] + dBu = dt[:, :, i, None] * B[:, i, None, :] * x[:, :, i, None] + ssm = dA * ssm + dBu + y_ref[:, :, i] = torch.einsum("bcn,bn->bc", ssm, C[:, i, :]) + y_ref = y_ref + D[None, :, None].double() * x + + rel = (y_pred.double() - y_ref).abs().max().item() / max(y_ref.abs().max().item(), 1e-8) + assert rel < 1e-5, f"M·x reconstruction vs fp64 eager scan rel diff {rel:.2e}" + + def test_reconstructs_mixer_output(self, cache, mamba_bridge): + """out_proj((M·x + D·x)·act(gate)) must reconstruct HF's mixer output.""" + mixer = mamba_bridge.blocks[0].mixer + M = mixer.compute_effective_attention(cache, layer_idx=0, include_dt_scaling=True) + x, _, _, _, _, D, gate = self._inputs(cache, mixer, 0, self.SEQ_LEN) + y = torch.einsum("bcij,bcj->bci", M, x) + D[None, :, None] * x + out = mixer.original_component.out_proj( + (y * mixer.original_component.act(gate)).transpose(1, 2) + ) + hook_out = cache["blocks.0.mixer.hook_out"].float() + rel = (out - hook_out).abs().max().item() / max(hook_out.abs().max().item(), 1e-8) + assert rel < 1e-5, ( + f"reconstructed mixer output vs hook_out rel diff {rel:.2e}; M is " + "inconsistent with HF's independently-computed forward." + ) + + def test_include_dt_scaling_changes_output(self, cache, mamba_bridge): + mixer = mamba_bridge.blocks[0].mixer + M_att = mixer.compute_effective_attention(cache, layer_idx=0, include_dt_scaling=False) + M_full = mixer.compute_effective_attention(cache, layer_idx=0, include_dt_scaling=True) + assert not torch.allclose(M_att, M_full) + + def test_raises_on_empty_cache(self, mamba_bridge): + from transformer_lens.ActivationCache import ActivationCache + + empty = ActivationCache({}, model=mamba_bridge) + with pytest.raises(RuntimeError, match="in cache"): + mamba_bridge.blocks[0].mixer.compute_effective_attention(empty, layer_idx=0) + + +class TestMamba1SSMState: + """Mamba-1 recurrent-state reconstruction (Phase 4.5): read-parity with Mamba-2. + + The vectorized state must match a naive fp64 step-by-step S6 scan, and + ``y = C·S + D·x`` reconstructed through gate + out_proj must match HF's cached + ``hook_out`` (ties S to HF's numerically-independent forward). + """ + + SEQ_LEN = 6 + + @pytest.fixture(scope="class") + def cache(self, mamba_bridge): + tokens = torch.arange(1, self.SEQ_LEN + 1).unsqueeze(0) + with torch.no_grad(): + _, cache = mamba_bridge.run_with_cache(tokens) + return cache + + @staticmethod + def _inputs(cache, mixer, layer, seq_len): + """Re-extract x (post-conv SiLU), B, C, dt, A, D, gate — independent of the impl.""" + import torch.nn.functional as F + + oc = mixer.original_component + d_inner, state, dt_rank = oc.intermediate_size, oc.ssm_state_size, oc.time_step_rank + conv = cache[f"blocks.{layer}.mixer.conv1d.hook_out"][..., :seq_len].float() + x_proj = cache[f"blocks.{layer}.mixer.x_proj.hook_out"].float() + dt_proj = cache[f"blocks.{layer}.mixer.dt_proj.hook_out"].float() + in_proj = cache[f"blocks.{layer}.mixer.in_proj.hook_out"].float() + x = oc.act(conv) # [b, d_inner, seq] + _ts, B, C = x_proj.split([dt_rank, state, state], dim=-1) + dt = F.softplus(dt_proj).transpose(1, 2) # [b, d_inner, seq] + A = -torch.exp(mixer.A_log.float()) + gate = in_proj.transpose(1, 2)[:, d_inner:, :] # [b, d_inner, seq] + return x, B, C, dt, A, mixer.D.float(), gate + + def test_shape(self, cache, mamba_bridge): + mixer = mamba_bridge.blocks[0].mixer + oc = mixer.original_component + S = mixer.compute_ssm_state(cache, layer_idx=0) + assert S.shape == (1, oc.intermediate_size, self.SEQ_LEN, oc.ssm_state_size) + assert torch.isfinite(S).all() + + def test_matches_fp64_eager_recurrence(self, cache, mamba_bridge): + """The vectorized state must match a naive fp64 step-by-step S6 scan.""" + mixer = mamba_bridge.blocks[0].mixer + S = mixer.compute_ssm_state(cache, layer_idx=0) + + x, B, _, dt, A, _, _ = self._inputs(cache, mixer, 0, self.SEQ_LEN) + A, B, x, dt = A.double(), B.double(), x.double(), dt.double() + b, d_inner = x.shape[0], x.shape[1] + ssm = torch.zeros(b, d_inner, A.shape[1], dtype=torch.float64) + ref = torch.zeros(b, d_inner, self.SEQ_LEN, A.shape[1], dtype=torch.float64) + for i in range(self.SEQ_LEN): + dA = torch.exp(A[None] * dt[:, :, i, None]) # [b, d_inner, state] + dBu = dt[:, :, i, None] * B[:, i, None, :] * x[:, :, i, None] + ssm = dA * ssm + dBu + ref[:, :, i] = ssm + + rel = (S.double() - ref).abs().max().item() / max(ref.abs().max().item(), 1e-8) + assert rel < 1e-5, f"S vs fp64 eager recurrence rel diff {rel:.2e}" + + def test_reconstructs_mixer_output(self, cache, mamba_bridge): + """out_proj((C·S + D·x)·act(gate)) must reconstruct HF's mixer output.""" + mixer = mamba_bridge.blocks[0].mixer + S = mixer.compute_ssm_state(cache, layer_idx=0) # [b, channels, seq, state] + x, _, C, _, _, D, gate = self._inputs(cache, mixer, 0, self.SEQ_LEN) + y = torch.einsum("bcts,bts->bct", S, C) + D[None, :, None] * x + out = mixer.original_component.out_proj( + (y * mixer.original_component.act(gate)).transpose(1, 2) + ) + hook_out = cache["blocks.0.mixer.hook_out"].float() + rel = (out - hook_out).abs().max().item() / max(hook_out.abs().max().item(), 1e-8) + assert rel < 1e-5, ( + f"C·S+D·x reconstruction vs hook_out rel diff {rel:.2e}; the state is " + "inconsistent with HF's independently-computed forward." + ) + + def test_time_step_matches_full_slice(self, cache, mamba_bridge): + mixer = mamba_bridge.blocks[0].mixer + S = mixer.compute_ssm_state(cache, layer_idx=0) + S_t = mixer.compute_ssm_state(cache, layer_idx=0, time_step=3) + assert S_t.shape == (S.shape[0], S.shape[1], S.shape[3]) + # Same math; the full ("bcsij") and single-step ("bcsj") einsums differ only + # in fp reduction order (measured ~1 fp32 ULP), so allclose, not equal. + assert torch.allclose(S[:, :, 3], S_t, atol=1e-6) + + def test_cache_method_matches_mixer(self, cache, mamba_bridge): + direct = mamba_bridge.blocks[0].mixer.compute_ssm_state(cache, layer_idx=0) + via_cache = cache.compute_ssm_state(layer=0) + assert torch.equal(direct, via_cache) + + def test_cache_all_layers_stacks(self, cache, mamba_bridge): + S_all = cache.compute_ssm_state() # pure Mamba-1 → every block is SSM → stacked + assert torch.is_tensor(S_all) + assert S_all.shape[0] == mamba_bridge.cfg.n_layers + assert torch.equal(S_all[0], cache.compute_ssm_state(layer=0)) + + def test_raises_on_empty_cache(self, mamba_bridge): + from transformer_lens.ActivationCache import ActivationCache + + empty = ActivationCache({}, model=mamba_bridge) + with pytest.raises(RuntimeError, match="in cache"): + mamba_bridge.blocks[0].mixer.compute_ssm_state(empty, layer_idx=0) + + +@contextlib.contextmanager +def _eager_scan(bridge): + """Enable the opt-in eager-scan intervention path on every Mamba-1 mixer.""" + for block in bridge.blocks: + block.mixer.eager_scan = True + try: + yield bridge + finally: + for block in bridge.blocks: + block.mixer.eager_scan = False + + +class TestMamba1EagerScanIntervention: + """Phase 4 (Mamba-1): opt-in eager S6 scan exposes hook_ssm_write / hook_ssm_state + for interventions that propagate to logits, while the default path is untouched. + Eager scan needs use_cache=False (prefill; cache_params is None).""" + + TOKENS = torch.tensor([[1, 2, 3, 4, 5, 6]]) + + def test_default_path_has_no_eager_hooks(self, mamba_bridge): + with torch.no_grad(): + _, cache = mamba_bridge.run_with_cache(self.TOKENS) + assert "blocks.0.mixer.hook_ssm_write" not in cache + assert "blocks.0.mixer.hook_ssm_state" not in cache + + def test_eager_requires_use_cache_false(self, mamba_bridge): + with _eager_scan(mamba_bridge), torch.no_grad(): + _, cache = mamba_bridge.run_with_cache(self.TOKENS) # default use_cache + assert "blocks.0.mixer.hook_ssm_write" not in cache + + def test_eager_scan_matches_fused_logits(self, mamba_bridge): + with torch.no_grad(): + fused = mamba_bridge(self.TOKENS) + with _eager_scan(mamba_bridge), torch.no_grad(): + eager = mamba_bridge.run_with_hooks(self.TOKENS, use_cache=False, fwd_hooks=[]) + rel = (eager - fused).abs().max().item() / max(fused.abs().max().item(), 1e-8) + assert rel < 1e-4, f"eager scan vs HF kernel rel diff {rel:.2e}" + + def test_eager_scan_matches_fused_padded_batch(self, mamba_bridge): + """Padded batch: eager path mirrors HF's channel-first padding mask (before + and after the conv).""" + ids = torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 0]]) + mask = torch.tensor([[1, 1, 1, 1, 1], [1, 1, 1, 1, 0]]) + with torch.no_grad(): + fused = mamba_bridge(ids, attention_mask=mask) + with _eager_scan(mamba_bridge), torch.no_grad(): + eager = mamba_bridge.run_with_hooks( + ids, use_cache=False, attention_mask=mask, fwd_hooks=[] + ) + rel = (eager - fused).abs().max().item() / max(fused.abs().max().item(), 1e-8) + assert rel < 1e-4, f"padded-batch eager vs fused rel diff {rel:.2e}" + + def test_eager_hooks_fire_with_use_cache_false(self, mamba_bridge): + with _eager_scan(mamba_bridge), torch.no_grad(): + _, cache = mamba_bridge.run_with_cache(self.TOKENS, use_cache=False) + oc = mamba_bridge.blocks[0].mixer.original_component + seq = self.TOKENS.shape[1] + expected = (1, oc.intermediate_size, seq, oc.ssm_state_size) + assert cache["blocks.0.mixer.hook_ssm_write"].shape == expected + assert cache["blocks.0.mixer.hook_ssm_state"].shape == expected + + def test_write_knockout_changes_logits(self, mamba_bridge): + def knockout(writes, hook): + writes = writes.clone() + writes[:, :, 2] = 0.0 # channel-first: zero all channels' write at position 2 + return writes + + with _eager_scan(mamba_bridge), torch.no_grad(): + base = mamba_bridge.run_with_hooks(self.TOKENS, use_cache=False, fwd_hooks=[]) + patched = mamba_bridge.run_with_hooks( + self.TOKENS, + use_cache=False, + fwd_hooks=[("blocks.0.mixer.hook_ssm_write", knockout)], + ) + assert (patched - base).abs().max().item() > 1e-6 + + def test_state_patch_changes_logits(self, mamba_bridge): + def patch(state, hook): + state = state.clone() + state[:, :, 3] = 0.0 + return state + + with _eager_scan(mamba_bridge), torch.no_grad(): + base = mamba_bridge.run_with_hooks(self.TOKENS, use_cache=False, fwd_hooks=[]) + patched = mamba_bridge.run_with_hooks( + self.TOKENS, use_cache=False, fwd_hooks=[("blocks.0.mixer.hook_ssm_state", patch)] + ) + assert (patched - base).abs().max().item() > 1e-6 + + def test_disabling_restores_fused_path(self, mamba_bridge): + with torch.no_grad(): + fused_before = mamba_bridge(self.TOKENS) + with _eager_scan(mamba_bridge), torch.no_grad(): + mamba_bridge.run_with_hooks(self.TOKENS, use_cache=False, fwd_hooks=[]) + with torch.no_grad(): + fused_after = mamba_bridge(self.TOKENS) + assert torch.equal(fused_before, fused_after) + + +class _DummyTok: + pass + + +def _build_synthetic_mamba_bridge(): + """Tiny random-init Mamba-1 bridge (from_config, no download) for CI-runnable tests.""" + from transformers import AutoModelForCausalLM + from transformers.models.mamba import MambaConfig + + from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_config_from_hf, + ) + from transformer_lens.model_bridge.supported_architectures.mamba import ( + MambaArchitectureAdapter, + ) + + torch.manual_seed(0) + cfg = MambaConfig( + vocab_size=128, + hidden_size=32, + state_size=16, + num_hidden_layers=2, + expand=2, + time_step_rank=4, + conv_kernel=4, + ) + cfg.architectures = ["MambaForCausalLM"] + hf = AutoModelForCausalLM.from_config(cfg).eval() + bridge_cfg = build_bridge_config_from_hf( + hf.config, "MambaForCausalLM", "m1-tiny", torch.float32 + ) + return TransformerBridge(hf, MambaArchitectureAdapter(bridge_cfg), tokenizer=_DummyTok()) + + +class TestMamba1EagerScanFp64Reference: + """Independent fp64 correctness gate on the Mamba-1 eager scan (_eager_scan_forward). + + Pins the eager scan's hook_ssm_write / hook_ssm_state to a naive fp64 step recurrence + built from the cached submodule outputs — distinct from both the fused-kernel parity + test and the read-path fp64 tests, so a shared discretization error can't pass. Uses + synthetic weights (from_config) so it runs in stock CI without a model download. + """ + + SEQ_LEN = 7 + LAYER = 0 + + @pytest.fixture(scope="class") + def bridge(self): + return _build_synthetic_mamba_bridge() + + @pytest.fixture(scope="class") + def eager_cache(self, bridge): + tokens = torch.randint(1, 128, (2, self.SEQ_LEN)) + with _eager_scan(bridge), torch.no_grad(): + _, cache = bridge.run_with_cache(tokens, use_cache=False) + return cache + + def _fp64_reference(self, bridge, cache): + """Reconstruct (writes, state) in fp64 from cached submodule hooks. + + Independent reimplementation of the S6 write term + recurrence; reads the same + HF submodule outputs the eager scan reuses (so their hooks fire), not the eager + code. Returns channel-first ``[b, d_inner, seq, state]`` tensors. + """ + import torch.nn.functional as F + + oc = bridge.blocks[self.LAYER].mixer.original_component + d_inner, state, dt_rank = oc.intermediate_size, oc.ssm_state_size, oc.time_step_rank + p = f"blocks.{self.LAYER}.mixer" + + conv = cache[f"{p}.conv1d.hook_out"][..., : self.SEQ_LEN].double() # [b, d_inner, seq] + x = oc.act(conv) # HF act on the trimmed conv output + xp = cache[f"{p}.x_proj.hook_out"].double() # [b, seq, dt_rank + 2*state] + _, B, _C = xp.split([dt_rank, state, state], dim=-1) # B: [b, seq, state] + dtp = cache[f"{p}.dt_proj.hook_out"].double() # [b, seq, d_inner] (pre-softplus) + dt = F.softplus(dtp).transpose(1, 2) # [b, d_inner, seq] + A = -torch.exp(bridge.blocks[self.LAYER].mixer.A_log.double()) # [d_inner, state] + + writes = (dt * x)[:, :, :, None] * B[:, None, :, :] # [b, d_inner, seq, state] + b = writes.shape[0] + ssm = torch.zeros(b, d_inner, state, dtype=torch.float64) + states = [] + for t in range(self.SEQ_LEN): + ssm = torch.exp(A[None] * dt[:, :, t, None]) * ssm + writes[:, :, t] + states.append(ssm) + return writes, torch.stack(states, dim=2) # [b, d_inner, seq, state] + + def test_eager_write_matches_fp64_reference(self, bridge, eager_cache): + writes_ref, _ = self._fp64_reference(bridge, eager_cache) + got = eager_cache[f"blocks.{self.LAYER}.mixer.hook_ssm_write"].double() + rel = (got - writes_ref).abs().max().item() / max(writes_ref.abs().max().item(), 1e-12) + assert rel < 1e-5, f"eager hook_ssm_write vs fp64 reference rel diff {rel:.2e}" + + def test_eager_state_matches_fp64_reference(self, bridge, eager_cache): + _, state_ref = self._fp64_reference(bridge, eager_cache) + got = eager_cache[f"blocks.{self.LAYER}.mixer.hook_ssm_state"].double() + rel = (got - state_ref).abs().max().item() / max(state_ref.abs().max().item(), 1e-12) + assert rel < 1e-5, f"eager hook_ssm_state vs fp64 reference rel diff {rel:.2e}" + + def test_write_knockout_propagates_forward(self, bridge): + """Zeroing hook_ssm_write at position k must change the state at every position + >= k (the recurrence carries it forward) and leave positions < k untouched.""" + tokens = torch.randint(1, 128, (1, self.SEQ_LEN)) + k = 3 + state_key = f"blocks.{self.LAYER}.mixer.hook_ssm_state" + write_key = f"blocks.{self.LAYER}.mixer.hook_ssm_write" + + def knock(writes, hook): + writes = writes.clone() + writes[:, :, k] = 0.0 + return writes + + captured = {} + + def grab(state, hook): + captured["s"] = state.detach().clone() + return state + + with _eager_scan(bridge), torch.no_grad(): + _, base_cache = bridge.run_with_cache(tokens, use_cache=False) + base_state = base_cache[state_key] + bridge.run_with_hooks( + tokens, use_cache=False, fwd_hooks=[(write_key, knock), (state_key, grab)] + ) + patched_state = captured["s"] + # positions before k unchanged; positions at/after k changed + assert torch.allclose(patched_state[:, :, :k], base_state[:, :, :k], atol=1e-6) + assert (patched_state[:, :, k:] - base_state[:, :, k:]).abs().max().item() > 1e-6 + + def test_state_patch_is_same_position_only(self, bridge): + """Patching hook_ssm_state at position k changes only the same-position mixer + output (y_t = C_t·S_t is a post-scan readout — it does not re-run the scan).""" + tokens = torch.randint(1, 128, (1, self.SEQ_LEN)) + k = 3 + state_key = f"blocks.{self.LAYER}.mixer.hook_ssm_state" + out_key = f"blocks.{self.LAYER}.mixer.hook_out" + + def patch(state, hook): + state = state.clone() + state[:, :, k] = 0.0 + return state + + base_out, patched_out = {}, {} + + def grab_base(o, hook): + base_out["o"] = o.detach().clone() + return o + + def grab_patched(o, hook): + patched_out["o"] = o.detach().clone() + return o + + with _eager_scan(bridge), torch.no_grad(): + bridge.run_with_hooks(tokens, use_cache=False, fwd_hooks=[(out_key, grab_base)]) + bridge.run_with_hooks( + tokens, use_cache=False, fwd_hooks=[(state_key, patch), (out_key, grab_patched)] + ) + delta = (patched_out["o"] - base_out["o"]).abs() + assert delta[:, k].max().item() > 1e-6, "same-position output must change" + other = torch.cat([delta[:, :k], delta[:, k + 1 :]], dim=1) + assert other.max().item() < 1e-6, "other positions must be untouched (no propagation)" + + def test_batch1_padded_matches_hf(self, bridge): + """Batch-1 with a padding mask must match HF: MambaMixer masks whenever + attention_mask is present (no batch>1 guard, unlike Mamba-2).""" + tokens = torch.tensor([[1, 2, 3, 4, 0]]) + mask = torch.tensor([[1, 1, 1, 1, 0]]) + with torch.no_grad(): + fused = bridge(tokens, attention_mask=mask) + with _eager_scan(bridge): + eager = bridge.run_with_hooks( + tokens, use_cache=False, attention_mask=mask, fwd_hooks=[] + ) + rel = (eager - fused).abs().max().item() / max(fused.abs().max().item(), 1e-8) + assert rel < 1e-4, f"batch-1 padded eager vs HF fused rel diff {rel:.2e}" diff --git a/tests/integration/model_bridge/test_nemotron_h_adapter.py b/tests/integration/model_bridge/test_nemotron_h_adapter.py index cab7d144d..5c56a8115 100644 --- a/tests/integration/model_bridge/test_nemotron_h_adapter.py +++ b/tests/integration/model_bridge/test_nemotron_h_adapter.py @@ -1,12 +1,12 @@ """Integration tests for the NemotronH architecture adapter. -Verifies forward-pass and generation parity against nvidia/Nemotron-H-8B-Base: +Verifies forward-pass and generation parity against nvidia/NVIDIA-Nemotron-Nano-9B-v2: - Forward-pass logits match HF exactly (bridge delegates the full forward to HF) - Greedy multi-token generation matches HF bit-for-bit (exercises DynamicCache state handling across attention, Mamba-2, MLP, and MoE layers) - Sanity checks: config flags, block count, hook coverage -Note: requires ~18 GB RAM (CPU) or ~16 GB VRAM (GPU) to load the 8B checkpoint. +Note: requires ~36 GB RAM (CPU, fp32) or ~18 GB VRAM (GPU, bf16) to load the 9B checkpoint. On a machine with less memory, skip with: pytest -m "not slow" tests/integration/model_bridge/test_nemotron_h_adapter.py @@ -28,7 +28,7 @@ pytestmark = pytest.mark.slow -MODEL = "nvidia/Nemotron-H-8B-Base" +MODEL = "nvidia/NVIDIA-Nemotron-Nano-9B-v2" # --------------------------------------------------------------------------- # Helpers @@ -82,7 +82,7 @@ def test_norm_bridges_are_rms(self, nemotron_bridge: TransformerBridge) -> None: assert isinstance(nemotron_bridge.ln_final, RMSNormalizationBridge) def test_block_count(self, nemotron_bridge: TransformerBridge) -> None: - # Nemotron-H-8B has 56 layers + # Nemotron-Nano-9B-v2 has 56 layers assert len(nemotron_bridge.blocks) == 56 def test_blocks_are_ssm_block_bridge(self, nemotron_bridge: TransformerBridge) -> None: diff --git a/tests/integration/model_bridge/test_nemotron_h_tiny.py b/tests/integration/model_bridge/test_nemotron_h_tiny.py new file mode 100644 index 000000000..7439d862f --- /dev/null +++ b/tests/integration/model_bridge/test_nemotron_h_tiny.py @@ -0,0 +1,135 @@ +"""Tiny from_config NemotronH integration tests — CI coverage without a real checkpoint. + +The full nvidia/NVIDIA-Nemotron-Nano-9B-v2 parity suite (test_nemotron_h_adapter.py) is +``@pytest.mark.slow`` and needs ~36 GB, so it does not run in normal CI. This file +builds a tiny synthetic NemotronH (real random CPU weights, no Hub download) that +exercises the same surfaces — the single passthrough ``.mixer`` slot on every +layer, forward parity, hook coverage, and the family-agnostic SSM dispatch. +""" +import pytest +import torch +from transformers import AutoModelForCausalLM +from transformers.models.nemotron_h import NemotronHConfig + +from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.generalized_components import ( + SSM2MixerBridge, + SSMBlockBridge, +) +from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_config_from_hf, +) +from transformer_lens.model_bridge.supported_architectures.nemotron_h import ( + NemotronHArchitectureAdapter, +) + +LAYERS = ["mamba", "attention", "mamba", "mlp"] +MAMBA_LAYERS = [0, 2] +ATTN_LAYER = 1 + + +class _Tok: + pass + + +@pytest.fixture(scope="module") +def bridge(): + torch.manual_seed(0) + cfg = NemotronHConfig( + vocab_size=256, + hidden_size=64, + layers_block_type=LAYERS, + num_attention_heads=4, + num_key_value_heads=2, + ssm_state_size=16, + mamba_num_heads=4, + mamba_head_dim=16, + n_groups=2, + conv_kernel=4, + expand=2, + intermediate_size=128, + chunk_size=8, + ) + cfg.architectures = ["NemotronHForCausalLM"] + hf = AutoModelForCausalLM.from_config(cfg).to(torch.float32).eval() + bridge_cfg = build_bridge_config_from_hf( + hf.config, "NemotronHForCausalLM", "nh-tiny", torch.float32 + ) + return TransformerBridge(hf, NemotronHArchitectureAdapter(bridge_cfg), tokenizer=_Tok()) + + +@pytest.fixture(scope="module") +def tokens(): + return torch.tensor([[1, 2, 3, 4, 5]]) + + +@pytest.fixture(scope="module") +def cache(bridge, tokens): + with torch.no_grad(): + _, c = bridge.run_with_cache(tokens) + return c + + +class TestNemotronHTinyStructure: + def test_block_count(self, bridge): + assert len(bridge.blocks) == len(LAYERS) + + def test_blocks_are_ssm_block_bridge(self, bridge): + assert isinstance(bridge.blocks[0], SSMBlockBridge) + + def test_every_block_wires_ssm2_mixer_passthrough(self, bridge): + # NemotronH wires one SSM2MixerBridge .mixer per layer, for all types. + for i in range(len(LAYERS)): + assert isinstance(bridge.blocks[i].mixer, SSM2MixerBridge) + + def test_layers_block_type_populated(self, bridge): + assert getattr(bridge.cfg, "layers_block_type", None) == LAYERS + + +class TestNemotronHTinyForwardPass: + def test_forward_matches_hf_exactly(self, bridge, tokens): + with torch.no_grad(): + bridge_out = bridge(tokens) + hf_out = bridge.original_model(tokens).logits + assert (bridge_out.float() - hf_out.float()).abs().max().item() == 0.0 + + def test_no_nan_longer_sequence(self, bridge): + with torch.no_grad(): + out = bridge(torch.arange(1, 17).unsqueeze(0)) + assert not torch.isnan(out).any() + + +class TestNemotronHTinyHookCoverage: + def test_block_hooks_fire(self, cache): + for i in range(len(LAYERS)): + assert f"blocks.{i}.hook_in" in cache + assert f"blocks.{i}.hook_out" in cache + + def test_mamba_submodule_hooks_fire(self, cache): + for i in MAMBA_LAYERS: + for submod in ("in_proj", "conv1d", "out_proj"): + assert f"blocks.{i}.mixer.{submod}.hook_out" in cache + + def test_no_transformer_specific_hooks(self, cache): + forbidden = ("hook_resid_mid", "hook_attn_out", "hook_mlp_out") + assert [k for k in cache if any(f in k for f in forbidden)] == [] + + +class TestNemotronHTinyDispatch: + def test_ssm_layers_excludes_passthrough(self, cache): + # Structural passthrough exclusion: only the real Mamba layers, no attn/mlp. + assert cache.ssm_layers() == MAMBA_LAYERS + + def test_effective_attention_per_ssm_layer_dict(self, cache): + M = cache.compute_ssm_effective_attention() + assert isinstance(M, dict) + assert sorted(M.keys()) == MAMBA_LAYERS + + def test_ssm_state_per_ssm_layer_dict(self, cache): + S = cache.compute_ssm_state() + assert isinstance(S, dict) + assert sorted(S.keys()) == MAMBA_LAYERS + + def test_attention_layer_raises_typeerror(self, cache): + with pytest.raises(TypeError): + cache.compute_ssm_effective_attention(layer=ATTN_LAYER) diff --git a/tests/integration/model_bridge/test_ssm_effective_attention_dispatch.py b/tests/integration/model_bridge/test_ssm_effective_attention_dispatch.py new file mode 100644 index 000000000..11ed87900 --- /dev/null +++ b/tests/integration/model_bridge/test_ssm_effective_attention_dispatch.py @@ -0,0 +1,416 @@ +"""Phase 3: family-agnostic SSM effective-attention dispatch + canonical hook vocabulary. + +Builds tiny synthetic models (no Hub access) for each SSM family and verifies: +- SSMMixerProtocol conformance (Mamba-1, Mamba-2, gated-delta-net; not attention), +- cache.ssm_layers / cache.compute_ssm_effective_attention dispatch across families, +- option forwarding (per_state_coord Mamba-1 only; include_dt_scaling not on GDN), +- the additive canonical hook vocabulary resolving cross-family via run_with_cache. +""" +from unittest.mock import MagicMock + +import pytest +import torch + +from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.generalized_components import ( + SSM2MixerBridge, + SSMMixerBridge, + SSMMixerProtocol, + find_ssm_mixer, +) +from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_config_from_hf, +) + +try: + from transformers import Qwen3_5ForCausalLM, Qwen3_5TextConfig + + _QWEN3_5_AVAILABLE = True +except ImportError: + _QWEN3_5_AVAILABLE = False + + +class _Tok: + pass + + +def _boot(hf_model, arch): + cfg = build_bridge_config_from_hf(hf_model.config, arch, "tiny", torch.float32) + from transformer_lens.factories.architecture_adapter_factory import ( + ArchitectureAdapterFactory, + ) + + adapter = ArchitectureAdapterFactory.select_architecture_adapter(cfg) + return TransformerBridge(model=hf_model, adapter=adapter, tokenizer=_Tok()) + + +@pytest.fixture(scope="module") +def mamba1_bridge(): + from transformers import AutoModelForCausalLM + from transformers.models.mamba import MambaConfig + + torch.manual_seed(0) + c = MambaConfig( + vocab_size=128, + hidden_size=32, + state_size=16, + num_hidden_layers=2, + expand=2, + time_step_rank=4, + conv_kernel=4, + ) + c.architectures = ["MambaForCausalLM"] + return _boot(AutoModelForCausalLM.from_config(c).eval(), "MambaForCausalLM") + + +@pytest.fixture(scope="module") +def mamba2_bridge(): + from transformers import AutoModelForCausalLM + from transformers.models.mamba2 import Mamba2Config + + torch.manual_seed(0) + c = Mamba2Config( + vocab_size=128, + hidden_size=32, + num_hidden_layers=2, + state_size=16, + expand=2, + head_dim=16, + num_heads=4, + n_groups=2, + conv_kernel=4, + chunk_size=8, + ) + c.architectures = ["Mamba2ForCausalLM"] + return _boot(AutoModelForCausalLM.from_config(c).eval(), "Mamba2ForCausalLM") + + +@pytest.fixture(scope="module") +def granite_bridge(): + from transformers import AutoModelForCausalLM + from transformers.models.granitemoehybrid import GraniteMoeHybridConfig + + torch.manual_seed(0) + c = GraniteMoeHybridConfig( + vocab_size=256, + hidden_size=64, + intermediate_size=32, + shared_intermediate_size=32, + num_hidden_layers=3, + num_attention_heads=8, + num_key_value_heads=4, + num_local_experts=4, + num_experts_per_tok=2, + max_position_embeddings=64, + layer_types=["mamba", "attention", "mamba"], + mamba_n_heads=8, + mamba_n_groups=2, + mamba_d_state=16, + mamba_d_head=16, + mamba_d_conv=4, + mamba_expand=2, + mamba_chunk_size=16, + position_embedding_type="rope", + rope_parameters={"rope_type": "default", "rope_theta": 10000.0}, + ) + c.architectures = ["GraniteMoeHybridForCausalLM"] + return _boot(AutoModelForCausalLM.from_config(c).eval(), "GraniteMoeHybridForCausalLM") + + +@pytest.fixture(scope="module") +def nemotronh_bridge(): + """Tiny NemotronH (Mamba-2 + attention + MLP hybrid) — exercises the single + passthrough ``.mixer`` slot wired on every layer.""" + from transformers import AutoModelForCausalLM + from transformers.models.nemotron_h import NemotronHConfig + + torch.manual_seed(0) + c = NemotronHConfig( + vocab_size=256, + hidden_size=64, + layers_block_type=["mamba", "attention", "mamba", "mlp"], + num_attention_heads=4, + num_key_value_heads=2, + ssm_state_size=16, + mamba_num_heads=4, + mamba_head_dim=16, + n_groups=2, + conv_kernel=4, + expand=2, + intermediate_size=128, + chunk_size=8, + ) + c.architectures = ["NemotronHForCausalLM"] + return _boot(AutoModelForCausalLM.from_config(c).eval(), "NemotronHForCausalLM") + + +@pytest.fixture(scope="module") +def qwen35_bridge(): + from transformer_lens.config.transformer_bridge_config import ( + TransformerBridgeConfig, + ) + from transformer_lens.model_bridge.supported_architectures.qwen3_5 import ( + Qwen3_5ArchitectureAdapter, + ) + + torch.manual_seed(0) + c = Qwen3_5TextConfig( + hidden_size=128, + num_hidden_layers=8, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=32, + intermediate_size=256, + vocab_size=512, + full_attention_interval=4, + linear_conv_kernel_dim=4, + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_num_key_heads=4, + linear_num_value_heads=4, + rope_parameters={ + "rope_theta": 10000.0, + "partial_rotary_factor": 0.25, + "rope_type": "default", + }, + ) + m = Qwen3_5ForCausalLM(c).eval() + bcfg = TransformerBridgeConfig( + d_model=128, + d_head=32, + n_heads=4, + n_layers=8, + n_ctx=2048, + d_vocab=512, + n_key_value_heads=2, + architecture="Qwen3_5ForCausalLM", + ) + return TransformerBridge(m, Qwen3_5ArchitectureAdapter(bcfg), tokenizer=MagicMock()) + + +TOKENS = torch.tensor([[1, 2, 3, 4, 5]]) + + +# --------------------------------------------------------------------------- +# SSMMixerProtocol conformance +# --------------------------------------------------------------------------- + + +class TestSSMMixerProtocol: + def test_mamba2_conforms(self, mamba2_bridge): + assert isinstance(mamba2_bridge.blocks[0].mixer, SSMMixerProtocol) + + def test_mamba1_conforms(self, mamba1_bridge): + assert isinstance(mamba1_bridge.blocks[0].mixer, SSMMixerProtocol) + + def test_attention_does_not_conform(self, granite_bridge): + # Block 1 is an attention layer — no SSM mixer. + assert not isinstance(granite_bridge.blocks[1].attn, SSMMixerProtocol) + + def test_find_ssm_mixer_on_attention_layer_is_none(self, granite_bridge): + assert find_ssm_mixer(granite_bridge.blocks[1]) is None + + def test_find_ssm_mixer_on_ssm_layer(self, granite_bridge): + assert isinstance(find_ssm_mixer(granite_bridge.blocks[0]), SSM2MixerBridge) + + def test_nemotronh_passthrough_mixer_is_not_found(self, nemotronh_bridge): + """NemotronH wires a passthrough SSM2MixerBridge .mixer on EVERY layer; it + conforms to the protocol by method presence, but find_ssm_mixer must reject + the no-op wrapper on attention / MLP layers (layer_types = mamba/attn/mamba/mlp).""" + assert isinstance(find_ssm_mixer(nemotronh_bridge.blocks[0]), SSM2MixerBridge) # mamba + assert find_ssm_mixer(nemotronh_bridge.blocks[1]) is None # attention passthrough + assert find_ssm_mixer(nemotronh_bridge.blocks[3]) is None # mlp passthrough + + @pytest.mark.skipif(not _QWEN3_5_AVAILABLE, reason="Qwen3_5 not available") + def test_gated_delta_net_conforms(self, qwen35_bridge): + assert isinstance(qwen35_bridge.blocks[0].linear_attn, SSMMixerProtocol) + + +# --------------------------------------------------------------------------- +# cache.ssm_layers +# --------------------------------------------------------------------------- + + +class TestSSMLayers: + def test_mamba2_all_layers(self, mamba2_bridge): + with torch.no_grad(): + _, cache = mamba2_bridge.run_with_cache(TOKENS) + assert cache.ssm_layers() == [0, 1] + + def test_granite_only_mamba_layers(self, granite_bridge): + with torch.no_grad(): + _, cache = granite_bridge.run_with_cache(TOKENS) + assert cache.ssm_layers() == [0, 2] + + def test_nemotronh_excludes_passthrough_layers(self, nemotronh_bridge): + # Structural detection must drop the passthrough .mixer on attn/mlp layers + # without relying on cfg.layers_block_type (which is empty in this path). + with torch.no_grad(): + _, cache = nemotronh_bridge.run_with_cache(TOKENS) + assert cache.ssm_layers() == [0, 2] + + def test_mixer_type_filter(self, granite_bridge): + with torch.no_grad(): + _, cache = granite_bridge.run_with_cache(TOKENS) + assert cache.ssm_layers(mixer_type=SSM2MixerBridge) == [0, 2] + assert cache.ssm_layers(mixer_type=SSMMixerBridge) == [] + + @pytest.mark.skipif(not _QWEN3_5_AVAILABLE, reason="Qwen3_5 not available") + def test_qwen35_linear_attn_layers(self, qwen35_bridge): + with torch.no_grad(): + _, cache = qwen35_bridge.run_with_cache(TOKENS, use_cache=False) + # full_attention_interval=4 -> full-attn at 3 and 7, linear-attn elsewhere + assert cache.ssm_layers() == [0, 1, 2, 4, 5, 6] + + +# --------------------------------------------------------------------------- +# cache.compute_ssm_effective_attention dispatch +# --------------------------------------------------------------------------- + + +class TestComputeSSMEffectiveAttention: + def test_mamba2_stacks_and_matches_direct(self, mamba2_bridge): + with torch.no_grad(): + _, cache = mamba2_bridge.run_with_cache(TOKENS) + M_all = cache.compute_ssm_effective_attention() + assert torch.is_tensor(M_all) + assert M_all.shape[0] == mamba2_bridge.cfg.n_layers + assert torch.equal(M_all[0], cache.compute_ssm_effective_attention(layer=0)) + + def test_granite_returns_dict(self, granite_bridge): + with torch.no_grad(): + _, cache = granite_bridge.run_with_cache(TOKENS) + M_all = cache.compute_ssm_effective_attention() + assert isinstance(M_all, dict) + assert sorted(M_all.keys()) == [0, 2] + direct = granite_bridge.blocks[0].mixer.compute_effective_attention(cache, layer_idx=0) + assert torch.equal(M_all[0], direct) + + def test_layer_without_mixer_raises_typeerror(self, granite_bridge): + with torch.no_grad(): + _, cache = granite_bridge.run_with_cache(TOKENS) + with pytest.raises(TypeError): + cache.compute_ssm_effective_attention(layer=1) # attention layer + + def test_nemotronh_dispatch_and_passthrough_contract(self, nemotronh_bridge): + """All-layers returns only the mamba layers; the single-layer path on a + passthrough (attention) layer raises the documented TypeError, not the + mixer's internal "needs ... in cache" RuntimeError.""" + with torch.no_grad(): + _, cache = nemotronh_bridge.run_with_cache(TOKENS) + M_all = cache.compute_ssm_effective_attention() + assert isinstance(M_all, dict) + assert sorted(M_all.keys()) == [0, 2] + direct = nemotronh_bridge.blocks[0].mixer.compute_effective_attention(cache, layer_idx=0) + assert torch.equal(M_all[0], direct) + with pytest.raises(TypeError): + cache.compute_ssm_effective_attention(layer=1) # passthrough attention layer + + def test_per_state_coord_mamba1(self, mamba1_bridge): + with torch.no_grad(): + _, cache = mamba1_bridge.run_with_cache(TOKENS) + M = cache.compute_ssm_effective_attention(layer=0, per_state_coord=True) + assert M.ndim == 5 # [batch, channels, state, seq, seq] + + def test_per_state_coord_unsupported_raises(self, mamba2_bridge): + with torch.no_grad(): + _, cache = mamba2_bridge.run_with_cache(TOKENS) + with pytest.raises(ValueError, match="per_state_coord"): + cache.compute_ssm_effective_attention(per_state_coord=True) + + @pytest.mark.skipif(not _QWEN3_5_AVAILABLE, reason="Qwen3_5 not available") + def test_qwen35_dict_and_include_dt_scaling_unsupported(self, qwen35_bridge): + with torch.no_grad(): + _, cache = qwen35_bridge.run_with_cache(TOKENS, use_cache=False) + M_all = cache.compute_ssm_effective_attention() + assert isinstance(M_all, dict) + assert sorted(M_all.keys()) == [0, 1, 2, 4, 5, 6] + with pytest.raises(ValueError, match="include_dt_scaling"): + cache.compute_ssm_effective_attention(include_dt_scaling=True) + + +# --------------------------------------------------------------------------- +# Canonical hook vocabulary (3b) — survives the hybrid alias resolution +# --------------------------------------------------------------------------- + + +class TestCanonicalHookVocabulary: + @pytest.mark.skipif(not _QWEN3_5_AVAILABLE, reason="Qwen3_5 not available") + def test_hook_ssm_out_resolves_cross_family(self, granite_bridge, qwen35_bridge): + """The critical test: one canonical name resolves on a Mamba-2 hybrid SSM + layer AND a gated-delta-net layer via run_with_cache (different slots).""" + with torch.no_grad(): + _, cg = granite_bridge.run_with_cache(TOKENS) + _, cq = qwen35_bridge.run_with_cache(TOKENS, use_cache=False) + assert "blocks.0.mixer.hook_ssm_out" in cg + assert "blocks.0.linear_attn.hook_ssm_out" in cq + + @pytest.mark.skipif(not _QWEN3_5_AVAILABLE, reason="Qwen3_5 not available") + def test_gated_delta_net_canonical_aliases(self, qwen35_bridge): + with torch.no_grad(): + _, cq = qwen35_bridge.run_with_cache(TOKENS, use_cache=False) + p = "blocks.0.linear_attn" + assert torch.equal(cq[f"{p}.hook_ssm_decay"], cq[f"{p}.hook_log_decay"]) + assert torch.equal(cq[f"{p}.hook_ssm_C"], cq[f"{p}.hook_q"]) + assert torch.equal(cq[f"{p}.hook_ssm_B"], cq[f"{p}.hook_k"]) + assert torch.equal(cq[f"{p}.hook_ssm_write"], cq[f"{p}.hook_beta"]) + + def test_mamba1_hook_ssm_dt_alias(self, mamba1_bridge): + with torch.no_grad(): + _, cache = mamba1_bridge.run_with_cache(TOKENS) + assert torch.equal( + cache["blocks.0.mixer.hook_ssm_dt"], cache["blocks.0.mixer.dt_proj.hook_out"] + ) + + +# --------------------------------------------------------------------------- +# cache.compute_ssm_state (Phase 4.5) — family-agnostic over Mamba-1 / Mamba-2 +# and gated-delta-net (each reconstructs its own recurrent state). +# --------------------------------------------------------------------------- + + +class TestComputeSsmStateDispatch: + def test_mamba1_reachable_via_cache(self, mamba1_bridge): + with torch.no_grad(): + _, cache = mamba1_bridge.run_with_cache(TOKENS) + S = cache.compute_ssm_state() # pure Mamba-1 → stacked + assert torch.is_tensor(S) + assert S.shape[0] == mamba1_bridge.cfg.n_layers + assert torch.equal(S[0], cache.compute_ssm_state(layer=0)) + + def test_mamba2_reachable_via_cache(self, mamba2_bridge): + with torch.no_grad(): + _, cache = mamba2_bridge.run_with_cache(TOKENS) + S = cache.compute_ssm_state() + assert torch.is_tensor(S) + assert S.shape[0] == mamba2_bridge.cfg.n_layers + + def test_granite_returns_per_ssm_layer_dict(self, granite_bridge): + with torch.no_grad(): + _, cache = granite_bridge.run_with_cache(TOKENS) + S = cache.compute_ssm_state() + assert isinstance(S, dict) + assert sorted(S.keys()) == [0, 2] + + @pytest.mark.skipif(not _QWEN3_5_AVAILABLE, reason="Qwen3_5 not available") + def test_gated_delta_net_reachable_via_cache(self, qwen35_bridge): + with torch.no_grad(): + _, cache = qwen35_bridge.run_with_cache(TOKENS, use_cache=False) + # Hybrid (full-attn layers interleaved) → dict over the linear-attn layers. + S = cache.compute_ssm_state() + assert isinstance(S, dict) + assert cache.ssm_layers() == sorted(S.keys()) + layer = cache.ssm_layers()[0] + # [batch, seq, n_v_heads, head_k_dim, head_v_dim] + assert S[layer].ndim == 5 + assert torch.equal(S[layer], cache.compute_ssm_state(layer=layer)) + + @pytest.mark.skipif(not _QWEN3_5_AVAILABLE, reason="Qwen3_5 not available") + def test_full_attention_layer_raises_typeerror(self, qwen35_bridge): + with torch.no_grad(): + _, cache = qwen35_bridge.run_with_cache(TOKENS, use_cache=False) + attn_layer = next( + i for i in range(qwen35_bridge.cfg.n_layers) if i not in cache.ssm_layers() + ) + with pytest.raises(TypeError, match="gated-delta-net"): + cache.compute_ssm_state(layer=attn_layer) diff --git a/tests/integration/model_bridge/test_ssm_real_weights.py b/tests/integration/model_bridge/test_ssm_real_weights.py new file mode 100644 index 000000000..f6c4710ab --- /dev/null +++ b/tests/integration/model_bridge/test_ssm_real_weights.py @@ -0,0 +1,156 @@ +"""Availability-gated real-weight integration tests for the SSM / recurrent families. + +The from_config tiny tests (test_mamba_adapter / test_mamba2_adapter / +test_nemotron_h_tiny / test_granite_moe_hybrid_adapter and the Qwen3_5 unit tests) +exercise the plumbing on random init. These load *real trained checkpoints* and run +the same interp surface — family-agnostic SSM-layer discovery, recurrent-state +reconstruction, effective attention, and the opt-in eager-scan intervention path — +where the numerics differ from random init. + +Every model is loaded through the ``bridge`` fixture, which SKIPS (never fails) when +the checkpoint or network is unavailable, the installed transformers lacks the +architecture, or there isn't enough memory. Point any family at a locally-cached +checkpoint with its env var (e.g. ``TL_MAMBA1_MODEL``); families without a small +official checkpoint (NemotronH 8B, Qwen3-Next 80B, …) skip in normal environments. + +All ``slow`` so ``make integration-test`` (``-m "not slow"``) never downloads. +""" +import gc +import os +from dataclasses import dataclass + +import pytest +import torch + +from transformer_lens.model_bridge.bridge import TransformerBridge + +pytestmark = pytest.mark.slow + + +@dataclass(frozen=True) +class _Case: + label: str + env: str + default: str # canonical HF id, or "" for env-only (no small official checkpoint) + + +# Ordered by how likely the checkpoint is small enough to actually run in CI. +CASES = [ + _Case("mamba1", "TL_MAMBA1_MODEL", "state-spaces/mamba-130m-hf"), + _Case("mamba2", "TL_MAMBA2_MODEL", "AntonV/mamba2-130m-hf"), + _Case("granite", "TL_GRANITE_MODEL", "ibm-granite/granite-4.0-tiny-preview"), + _Case("nemotron_h", "TL_NEMOTRON_H_MODEL", "nvidia/NVIDIA-Nemotron-Nano-9B-v2"), + _Case("qwen3_next", "TL_QWEN3_NEXT_MODEL", "Qwen/Qwen3-Next-80B-A3B-Instruct"), + _Case("qwen3_5", "TL_QWEN3_5_MODEL", ""), +] + + +def _device() -> str: + return "cuda" if torch.cuda.is_available() else "cpu" + + +@pytest.fixture(scope="module", params=CASES, ids=lambda c: c.label) +def bridge(request): + """Load one family's real checkpoint, or skip if it can't be obtained.""" + case: _Case = request.param + model_id = os.environ.get(case.env, case.default) + if not model_id: + pytest.skip(f"no default checkpoint for {case.label}; set {case.env} to run") + try: + # float32 (not bf16) so the reconstruction-faithfulness tolerances hold. + br = TransformerBridge.boot_transformers(model_id, device=_device(), dtype=torch.float32) + except pytest.skip.Exception: + raise + except Exception as e: # availability gate: network / missing arch / OOM / gated repo + pytest.skip(f"{case.label} checkpoint {model_id!r} unavailable: {type(e).__name__}: {e}") + yield br + del br + if torch.cuda.is_available(): + torch.cuda.empty_cache() + gc.collect() + + +@pytest.fixture(scope="module") +def tokens(): + # Small ids valid for every real vocab; short seq keeps the O(seq^2)/O(seq) work cheap. + return torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]]) + + +@pytest.fixture(scope="module") +def cache(bridge, tokens): + # use_cache=False forces the hooked prefill path so gated-delta-net interior hooks fire. + with torch.no_grad(): + _, c = bridge.run_with_cache(tokens.to(_device()), use_cache=False) + return c + + +def _realized_ssm_mixers(bridge): + """The realized SSM mixers (find_ssm_mixer excludes hybrid passthrough slots).""" + from transformer_lens.model_bridge.generalized_components.ssm_protocol import ( + find_ssm_mixer, + ) + + out = [] + for block in bridge.blocks: + mixer = find_ssm_mixer(block) + if mixer is not None: + out.append(mixer) + return out + + +class TestRealWeightSSMSurface: + def test_ssm_layers_nonempty(self, cache): + layers = cache.ssm_layers() + assert layers == sorted(layers) + assert len(layers) > 0, "a recurrent model must expose at least one SSM layer" + + def test_compute_ssm_state_finite(self, cache, tokens): + state = cache.compute_ssm_state() + mats = ( + list(state.values()) + if isinstance(state, dict) + else [state[i] for i in range(state.shape[0])] + ) + assert len(mats) == len(cache.ssm_layers()) + for S in mats: + assert torch.isfinite(S).all() + # batch dim leads; a seq axis of the right length is present somewhere. + assert S.shape[0] == tokens.shape[0] + assert tokens.shape[1] in tuple(S.shape) + + def test_effective_attention_finite_causal(self, cache, tokens): + M = cache.compute_ssm_effective_attention() + mats = list(M.values()) if isinstance(M, dict) else [M[i] for i in range(M.shape[0])] + seq = tokens.shape[1] + upper = torch.triu(torch.ones(seq, seq, dtype=torch.bool), diagonal=1) + for m in mats: + assert m.shape[-1] == m.shape[-2] == seq + assert torch.isfinite(m).all() + assert torch.all(m[..., upper] == 0), "effective attention must be causal" + + def test_eager_scan_matches_fused(self, bridge, tokens): + """Reimplemented eager recurrence must reproduce the fused-kernel logits. + + Enables eager_scan on every realized SSM mixer (Mamba-1/2, gated-delta-net, + and the SSM2-wired hybrid mixers) and compares to the default fused path on + real trained weights — the end-to-end faithfulness gate for the intervention + surface. + """ + mixers = [m for m in _realized_ssm_mixers(bridge) if hasattr(type(m), "eager_scan")] + if not mixers: + pytest.skip("no eager-scan-capable SSM mixers in this model") + + toks = tokens.to(_device()) + with torch.no_grad(): + fused = bridge(toks, use_cache=False) + for m in mixers: + m.eager_scan = True + try: + eager = bridge(toks, use_cache=False) + finally: + for m in mixers: + m.eager_scan = False + rel = (eager.float() - fused.float()).abs().max().item() / max( + fused.float().abs().max().item(), 1e-8 + ) + assert rel < 2e-2, f"eager vs fused logit parity rel diff {rel:.2e}" diff --git a/tests/mps/test_mps_ssm_eager_scan.py b/tests/mps/test_mps_ssm_eager_scan.py new file mode 100644 index 000000000..86074ce45 --- /dev/null +++ b/tests/mps/test_mps_ssm_eager_scan.py @@ -0,0 +1,82 @@ +"""MPS regression test for the SSM2 eager-scan intervention path (device correctness). + +The eager scan builds its recurrent ``ssm_state`` accumulator explicitly; if that +tensor is created on the default CPU device instead of the input's device, the +recurrence ``decay * ssm_state + writes`` crashes on any non-CPU model. This test +guards that on Metal. Skips on non-MPS runners; runs in the `mps-checks` CI job. +""" + +import gc + +import pytest +import torch + +pytestmark = pytest.mark.skipif( + not torch.backends.mps.is_available(), + reason="MPS not available on this runner — skipping Apple Silicon SSM tests", +) + + +class _Tok: + pass + + +def _tiny_mamba2_bridge(): + """Tiny synthetic Mamba-2 bridge (float32, no Hub download), built on CPU.""" + from transformers import AutoModelForCausalLM + from transformers.models.mamba2 import Mamba2Config + + from transformer_lens.model_bridge.bridge import TransformerBridge + from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_config_from_hf, + ) + from transformer_lens.model_bridge.supported_architectures.mamba2 import ( + Mamba2ArchitectureAdapter, + ) + + torch.manual_seed(0) + cfg = Mamba2Config( + vocab_size=128, + hidden_size=32, + num_hidden_layers=2, + state_size=16, + expand=2, + head_dim=16, + num_heads=4, + n_groups=2, + conv_kernel=4, + chunk_size=8, + ) + cfg.architectures = ["Mamba2ForCausalLM"] + hf = AutoModelForCausalLM.from_config(cfg).to(torch.float32).eval() + bridge_cfg = build_bridge_config_from_hf(hf.config, "Mamba2ForCausalLM", "x", torch.float32) + return TransformerBridge(hf, Mamba2ArchitectureAdapter(bridge_cfg), tokenizer=_Tok()) + + +def test_mps_eager_scan_runs_on_metal(): + """eager_scan=True must run on MPS (ssm_state created on the input's device) and + match the CPU fused kernel, with the intervention hooks firing on Metal.""" + bridge = _tiny_mamba2_bridge() + tokens = torch.tensor([[1, 2, 3, 4, 5]]) + with torch.no_grad(): + cpu_fused = bridge(tokens) + + bridge.original_model.to("mps") + for block in bridge.blocks: + block.mixer.eager_scan = True + tokens_mps = tokens.to("mps") + + with torch.no_grad(): + out = bridge(tokens_mps, use_cache=False) # would crash pre-fix (CPU ssm_state) + _, cache = bridge.run_with_cache(tokens_mps, use_cache=False) + + assert out.device.type == "mps", f"eager-scan output not on MPS: {out.device}" + assert cache["blocks.0.mixer.hook_ssm_write"].device.type == "mps" + assert cache["blocks.0.mixer.hook_ssm_state"].device.type == "mps" + + rel = (out.cpu() - cpu_fused).abs().max().item() / max(cpu_fused.abs().max().item(), 1e-8) + assert rel < 1e-3, f"MPS eager scan diverges from CPU fused kernel: rel {rel:.2e}" + + del bridge + torch.mps.empty_cache() + gc.collect() diff --git a/tests/unit/model_bridge/supported_architectures/test_granite_moe_hybrid_adapter.py b/tests/unit/model_bridge/supported_architectures/test_granite_moe_hybrid_adapter.py index 3d186e5a1..a9b6d2592 100644 --- a/tests/unit/model_bridge/supported_architectures/test_granite_moe_hybrid_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_granite_moe_hybrid_adapter.py @@ -7,6 +7,7 @@ BlockBridge, DepthwiseConv1DBridge, EmbeddingBridge, + GatedRMSNormBridge, LinearBridge, MLPBridge, MoEBridge, @@ -89,6 +90,11 @@ def test_common_granite_flags(self, adapter: GraniteMoeHybridArchitectureAdapter assert adapter.cfg.uses_rms_norm is True assert adapter.cfg.default_prepend_bos is False + def test_layers_block_type_surfaced(self, adapter: GraniteMoeHybridArchitectureAdapter) -> None: + # Sourced from HF's `layer_types`; exposed uniformly as `layers_block_type` + # (matching NemotronH) so analysis tools can find the Mamba layers. + assert getattr(adapter.cfg, "layers_block_type", None) == LAYER_TYPES + def test_non_rope_position_embedding_disables_rotary_mapping(self) -> None: adapter = GraniteMoeHybridArchitectureAdapter( _make_cfg(position_embedding_type="nope", num_experts=4) @@ -142,7 +148,7 @@ def test_block_submodule_mapping(self, adapter: GraniteMoeHybridArchitectureAdap "ln1", "ln2", "attn", - "mamba", + "mixer", "shared_mlp", "moe", } @@ -150,14 +156,15 @@ def test_block_submodule_mapping(self, adapter: GraniteMoeHybridArchitectureAdap assert isinstance(blocks.submodules["ln1"], RMSNormalizationBridge) assert isinstance(blocks.submodules["ln2"], RMSNormalizationBridge) assert isinstance(blocks.submodules["attn"], PositionEmbeddingsAttentionBridge) - assert isinstance(blocks.submodules["mamba"], SSM2MixerBridge) + assert isinstance(blocks.submodules["mixer"], SSM2MixerBridge) assert isinstance(blocks.submodules["shared_mlp"], MLPBridge) assert isinstance(blocks.submodules["moe"], MoEBridge) assert blocks.submodules["ln1"].name == "input_layernorm" assert blocks.submodules["ln2"].name == "post_attention_layernorm" assert blocks.submodules["attn"].name == "self_attn" - assert blocks.submodules["mamba"].name == "mamba" + # Canonical `.mixer` dict key; HF submodule path stays `.mamba`. + assert blocks.submodules["mixer"].name == "mamba" assert blocks.submodules["shared_mlp"].name == "shared_mlp" assert blocks.submodules["moe"].name == "block_sparse_moe" @@ -176,17 +183,18 @@ def test_attention_mapping(self, adapter: GraniteMoeHybridArchitectureAdapter) - assert isinstance(submodule, LinearBridge) def test_mamba_mapping(self, adapter: GraniteMoeHybridArchitectureAdapter) -> None: - mamba = adapter.component_mapping["blocks"].submodules["mamba"] - assert mamba.optional is True - assert set(mamba.submodules.keys()) == {"in_proj", "conv1d", "inner_norm"} - - assert isinstance(mamba.submodules["in_proj"], LinearBridge) - assert isinstance(mamba.submodules["conv1d"], DepthwiseConv1DBridge) - assert isinstance(mamba.submodules["inner_norm"], LinearBridge) - - assert mamba.submodules["in_proj"].name == "in_proj" - assert mamba.submodules["conv1d"].name == "conv1d" - assert mamba.submodules["inner_norm"].name == "norm" + mixer = adapter.component_mapping["blocks"].submodules["mixer"] + assert mixer.optional is True + assert set(mixer.submodules.keys()) == {"in_proj", "conv1d", "inner_norm"} + + assert isinstance(mixer.submodules["in_proj"], LinearBridge) + assert isinstance(mixer.submodules["conv1d"], DepthwiseConv1DBridge) + # HF's inner norm is the gated two-input MambaRMSNormGated, not a plain linear. + assert isinstance(mixer.submodules["inner_norm"], GatedRMSNormBridge) + + assert mixer.submodules["in_proj"].name == "in_proj" + assert mixer.submodules["conv1d"].name == "conv1d" + assert mixer.submodules["inner_norm"].name == "norm" def test_shared_mlp_mapping(self, adapter: GraniteMoeHybridArchitectureAdapter) -> None: shared_mlp = adapter.component_mapping["blocks"].submodules["shared_mlp"] diff --git a/tests/unit/model_bridge/supported_architectures/test_nemotron_h_adapter.py b/tests/unit/model_bridge/supported_architectures/test_nemotron_h_adapter.py index 8a3af1077..33790e4fd 100644 --- a/tests/unit/model_bridge/supported_architectures/test_nemotron_h_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_nemotron_h_adapter.py @@ -113,9 +113,8 @@ def test_conv_dim_propagated(self, adapter: NemotronHArchitectureAdapter) -> Non # intermediate=32, n_groups=2, ssm_state_size=4 → 32 + 2*2*4 = 48 assert getattr(adapter.cfg, "conv_dim", None) == 48 - def test_applicable_phases_empty(self) -> None: - # verify_models is transformer-shaped; SSM hybrids skip it. - assert NemotronHArchitectureAdapter.applicable_phases == [] + def test_applicable_phases_full(self) -> None: + assert NemotronHArchitectureAdapter.applicable_phases == [1, 2, 3, 4] def test_weight_processing_conversions_empty( self, adapter: NemotronHArchitectureAdapter diff --git a/tests/unit/model_bridge/supported_architectures/test_qwen3_5_adapter.py b/tests/unit/model_bridge/supported_architectures/test_qwen3_5_adapter.py index 75cc6b3b2..cced6757e 100644 --- a/tests/unit/model_bridge/supported_architectures/test_qwen3_5_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_qwen3_5_adapter.py @@ -680,3 +680,474 @@ def _capture(tensor: torch.Tensor, hook: object) -> torch.Tensor: f"Expected {hook_name} shape ({batch}, {seq}, {d_model}), " f"got {activations[0].shape}" ) + + +@pytest.mark.skipif( + not _QWEN3_5_AVAILABLE, + reason="Qwen3_5TextConfig / Qwen3_5ForCausalLM not available in installed transformers", +) +class TestQwen3_5GatedDeltaNetEffectiveAttention: + """Faithfulness gate for GatedDeltaNetBridge.compute_effective_attention. + + The reconstruction is a documented heuristic: (1) interior hooks fire only on + the hooked prefill path (use_cache=False); (2) the hooked Q/K are pre-norm + while the kernel L2-normalizes them; (3) the gated-linear-attention form omits + the delta rule's key-removal term. These tests pin the measured divergences. + """ + + LINEAR_LAYER = 0 # full-attn lands at 3 and 7 (interval=4); 0 is GatedDeltaNet + SEQ_LEN = 12 + + @pytest.fixture(scope="class") + def bridge(self): + br, _ = _make_tiny_bridge() + return br + + @pytest.fixture(scope="class") + def cache(self, bridge): + import torch + + torch.manual_seed(0) + tokens = torch.randint(0, 512, (1, self.SEQ_LEN)) + with torch.no_grad(): + # use_cache=False forces the hooked prefill path so interior hooks fire. + _, cache = bridge.run_with_cache(tokens, use_cache=False) + return cache + + @staticmethod + def _build_M(cache, layer, *, l2norm, dtype): + """Recompute M from cached hooks, independent of the bridge impl.""" + import torch + import torch.nn.functional as F + + prefix = f"blocks.{layer}.linear_attn" + q = cache[f"{prefix}.hook_q"].to(dtype) + k = cache[f"{prefix}.hook_k"].to(dtype) + beta = cache[f"{prefix}.hook_beta"].to(dtype) + g = cache[f"{prefix}.hook_log_decay"].to(dtype) + if l2norm: + q, k = F.normalize(q, p=2, dim=-1), F.normalize(k, p=2, dim=-1) + if q.shape[2] < beta.shape[-1]: + rep = beta.shape[-1] // q.shape[2] + q, k = q.repeat_interleave(rep, dim=2), k.repeat_interleave(rep, dim=2) + qk = torch.matmul(q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3).transpose(-2, -1)) + cum = torch.cumsum(g.permute(0, 2, 1), dim=-1) + L_log = cum[:, :, :, None] - cum[:, :, None, :] + seq = q.shape[1] + mask = torch.tril(torch.ones(seq, seq, dtype=torch.bool)) + L = torch.where(mask[None, None], torch.exp(L_log), torch.zeros_like(L_log)) + return qk * beta.permute(0, 2, 1)[:, :, None, :] * L + + def test_requires_use_cache_false(self, bridge): + """Default cache path omits interior hooks -> compute_effective_attention raises.""" + import torch + + tokens = torch.randint(0, 512, (1, self.SEQ_LEN)) + with torch.no_grad(): + _, cache = bridge.run_with_cache(tokens) # default: cache allocated + mixer = bridge.blocks[self.LINEAR_LAYER].linear_attn + with pytest.raises(RuntimeError, match="in cache"): + mixer.compute_effective_attention(cache, layer_idx=self.LINEAR_LAYER) + + def test_shape_causal_finite(self, bridge, cache): + import torch + + mixer = bridge.blocks[self.LINEAR_LAYER].linear_attn + M = mixer.compute_effective_attention(cache, layer_idx=self.LINEAR_LAYER) + assert M.ndim == 4 and M.shape[-1] == M.shape[-2] == self.SEQ_LEN + assert torch.isfinite(M).all() + upper = torch.triu(torch.ones(self.SEQ_LEN, self.SEQ_LEN, dtype=torch.bool), diagonal=1) + assert torch.all(M[..., upper] == 0), "effective attention must be causal" + + def test_impl_matches_fp64_reference(self, bridge, cache): + """The impl must match its own fp64 recomputation (no numerical/algorithm bug).""" + import torch + + mixer = bridge.blocks[self.LINEAR_LAYER].linear_attn + M = mixer.compute_effective_attention(cache, layer_idx=self.LINEAR_LAYER) + M_ref = self._build_M(cache, self.LINEAR_LAYER, l2norm=False, dtype=torch.float64) + rel = (M.double() - M_ref).abs().max().item() / max(M_ref.abs().max().item(), 1e-12) + assert rel < 1e-5, f"impl vs fp64 pre-norm reference rel diff {rel:.2e}" + + def test_l2norm_approximation_divergence_measured(self, bridge, cache): + """Pin the L2-norm approximation divergence (documented in the docstring). + + The hooked Q/K are pre-norm; the kernel L2-normalizes them. On the random- + init fixture (small, non-uniform Q/K norms) the relative divergence is ~1.0. + """ + import torch + + M_prenorm = self._build_M(cache, self.LINEAR_LAYER, l2norm=False, dtype=torch.float64) + M_l2 = self._build_M(cache, self.LINEAR_LAYER, l2norm=True, dtype=torch.float64) + rel = (M_prenorm - M_l2).abs().max().item() / max(M_l2.abs().max().item(), 1e-12) + assert torch.isfinite(torch.tensor(rel)) + assert 0.5 < rel < 1.5, ( + f"L2-norm approximation divergence {rel:.3f} outside the documented ~1.0 " + "band for the random-init fixture; update the docstring if the model changed." + ) + + +@pytest.mark.skipif( + not _QWEN3_5_AVAILABLE, + reason="Qwen3_5TextConfig / Qwen3_5ForCausalLM not available in installed transformers", +) +class TestQwen3_5GatedDeltaNetState: + """Faithful recurrent-state reconstruction + eager-scan interventions for GDN. + + Unlike compute_effective_attention (a gated-linear-attention heuristic that drops + key removal), compute_ssm_state replays the *full* delta rule, so ``S_t^T q_t`` + must reconstruct the fused kernel's hook_recurrence_out to fp tolerance. + """ + + LINEAR_LAYER = 0 # full-attn lands at 3 and 7 (interval=4); 0 is GatedDeltaNet + SEQ_LEN = 10 + + @pytest.fixture(scope="class") + def bridge(self): + br, _ = _make_tiny_bridge() + return br + + @pytest.fixture(scope="class") + def cache(self, bridge): + import torch + + torch.manual_seed(0) + tokens = torch.randint(0, 512, (2, self.SEQ_LEN)) + with torch.no_grad(): + _, cache = bridge.run_with_cache(tokens, use_cache=False) + return cache + + @staticmethod + def _ref_scan(cache, layer, dtype): + """Independent gated-delta-rule scan from cached hooks (not the impl). + + Uses the kernel's L2-norm convention (eps=1e-6) so the only differences vs + the impl are reduction order (fp32) and precision (fp64). + """ + import torch + + def l2norm(x): + return x / torch.sqrt((x * x).sum(-1, keepdim=True) + 1e-6) + + prefix = f"blocks.{layer}.linear_attn" + q = cache[f"{prefix}.hook_q"].to(dtype) + k = cache[f"{prefix}.hook_k"].to(dtype) + v = cache[f"{prefix}.hook_v"].to(dtype) + beta = cache[f"{prefix}.hook_beta"].to(dtype) + g = cache[f"{prefix}.hook_log_decay"].to(dtype) + n_v = v.shape[2] + if q.shape[2] < n_v: + rep = n_v // q.shape[2] + q, k = q.repeat_interleave(rep, dim=2), k.repeat_interleave(rep, dim=2) + q = l2norm(q) * (q.shape[-1] ** -0.5) + k = l2norm(k) + b, seq, _, kd = q.shape + vd = v.shape[-1] + state = torch.zeros(b, n_v, kd, vd, dtype=dtype) + states, outs = [], [] + for t in range(seq): + state = state * g[:, t].exp()[:, :, None, None] + kv = (state * k[:, t][:, :, :, None]).sum(-2) + delta = (v[:, t] - kv) * beta[:, t][:, :, None] + state = state + k[:, t][:, :, :, None] * delta[:, :, None, :] + states.append(state) + outs.append((state * q[:, t][:, :, :, None]).sum(-2)) + return torch.stack(states, 1), torch.stack(outs, 1) + + def test_requires_use_cache_false(self, bridge): + """Default cache path omits interior hooks -> compute_ssm_state raises.""" + import torch + + tokens = torch.randint(0, 512, (1, self.SEQ_LEN)) + with torch.no_grad(): + _, cache = bridge.run_with_cache(tokens) # default: cache allocated + mixer = bridge.blocks[self.LINEAR_LAYER].linear_attn + with pytest.raises(RuntimeError, match="in cache"): + mixer.compute_ssm_state(cache, layer_idx=self.LINEAR_LAYER) + + def test_state_shape(self, bridge, cache): + mixer = bridge.blocks[self.LINEAR_LAYER].linear_attn + S = mixer.compute_ssm_state(cache, layer_idx=self.LINEAR_LAYER) + # [batch, seq, n_v_heads, head_k_dim, head_v_dim] + assert S.ndim == 5 + assert S.shape[1] == self.SEQ_LEN + assert S.shape[2] == cache[f"blocks.{self.LINEAR_LAYER}.linear_attn.hook_v"].shape[2] + + def test_reconstructs_recurrence_out(self, bridge, cache): + """o_t = S_t^T q_t must match the fused kernel's hook_recurrence_out (faithful).""" + import torch + + mixer = bridge.blocks[self.LINEAR_LAYER].linear_attn + S = mixer.compute_ssm_state(cache, layer_idx=self.LINEAR_LAYER) + # o_t from the impl's own state, using the impl's normalized/scaled q. + prefix = f"blocks.{self.LINEAR_LAYER}.linear_attn" + q = cache[f"{prefix}.hook_q"].float() + n_v = cache[f"{prefix}.hook_v"].shape[2] + if q.shape[2] < n_v: + q = q.repeat_interleave(n_v // q.shape[2], dim=2) + q = mixer._l2norm(q) * (q.shape[-1] ** -0.5) + o = torch.einsum("bshkv,bshk->bshv", S, q) + rec = cache[f"{prefix}.hook_recurrence_out"].float() + rel = (o - rec).abs().max().item() / max(rec.abs().max().item(), 1e-8) + assert rel < 1e-4, f"delta-rule state reconstruction rel diff {rel:.2e}" + + def test_impl_matches_independent_reference(self, bridge, cache): + """The impl state must match an independent reimplementation of the scan. + + The impl works in fp32 by design, so it is compared to a fp32 reference at + tight tolerance (independent-algorithm check, only reduction-order noise). + The residual against a fp64 reference is then shown to be no larger than + fp32's own rounding gap — i.e. the impl adds no error beyond its working + precision (an actual algorithm bug would be O(1), not fp32-scale). + """ + import torch + + mixer = bridge.blocks[self.LINEAR_LAYER].linear_attn + S = mixer.compute_ssm_state(cache, layer_idx=self.LINEAR_LAYER).float() + S_ref32, _ = self._ref_scan(cache, self.LINEAR_LAYER, torch.float32) + S_ref64, _ = self._ref_scan(cache, self.LINEAR_LAYER, torch.float64) + + denom = max(S_ref64.abs().max().item(), 1e-12) + rel32 = (S - S_ref32).abs().max().item() / denom + assert rel32 < 1e-5, f"impl vs independent fp32 reference rel diff {rel32:.2e}" + + # fp64 gap of the impl vs fp64 gap of the fp32 reference: comparable ⇒ fp32 noise. + impl_vs_fp64 = (S.double() - S_ref64).abs().max().item() / denom + ref_fp32_vs_fp64 = (S_ref32.double() - S_ref64).abs().max().item() / denom + assert impl_vs_fp64 <= 4 * ref_fp32_vs_fp64 + 1e-6, ( + f"impl fp64 gap {impl_vs_fp64:.2e} exceeds fp32 rounding gap " + f"{ref_fp32_vs_fp64:.2e} — not attributable to fp32 accumulation" + ) + + def test_time_step_matches_full(self, bridge, cache): + import torch + + mixer = bridge.blocks[self.LINEAR_LAYER].linear_attn + S = mixer.compute_ssm_state(cache, layer_idx=self.LINEAR_LAYER) + for t in (0, self.SEQ_LEN // 2, self.SEQ_LEN - 1): + St = mixer.compute_ssm_state(cache, layer_idx=self.LINEAR_LAYER, time_step=t) + assert torch.equal(St, S[:, t]) + + +@pytest.mark.skipif( + not _QWEN3_5_AVAILABLE, + reason="Qwen3_5TextConfig / Qwen3_5ForCausalLM not available in installed transformers", +) +class TestQwen3_5GatedDeltaNetEagerScan: + """Opt-in eager_scan path: fused-kernel parity, hook_ssm_state interventions.""" + + LINEAR_LAYER = 0 + SEQ_LEN = 10 + + @pytest.fixture(scope="class") + def bridge(self): + br, _ = _make_tiny_bridge() + return br + + @pytest.fixture(scope="class") + def tokens(self): + import torch + + torch.manual_seed(0) + return torch.randint(0, 512, (2, self.SEQ_LEN)) + + @staticmethod + def _eager(bridge): + import contextlib + + from transformer_lens.model_bridge.generalized_components import ( + GatedDeltaNetBridge, + ) + + @contextlib.contextmanager + def _ctx(): + mixers = [ + b.linear_attn + for b in bridge.blocks + if isinstance(getattr(b, "linear_attn", None), GatedDeltaNetBridge) + ] + for m in mixers: + m.eager_scan = True + try: + yield + finally: + for m in mixers: + m.eager_scan = False + + return _ctx() + + def test_default_path_has_no_state_hook(self, bridge, tokens): + import torch + + with torch.no_grad(): + _, cache = bridge.run_with_cache(tokens, use_cache=False) + assert not any("hook_ssm_state" in k for k in cache) + + def test_eager_scan_matches_fused(self, bridge, tokens): + """Eager delta-rule scan reproduces the fused-kernel logits to fp tolerance.""" + import torch + + with torch.no_grad(): + base = bridge(tokens, use_cache=False) + with self._eager(bridge): + eager = bridge(tokens, use_cache=False) + rel = (eager.float() - base.float()).abs().max().item() / max( + base.float().abs().max().item(), 1e-8 + ) + assert rel < 1e-3, f"eager vs fused logit parity rel diff {rel:.2e}" + + def test_eager_scan_fires_state_hook(self, bridge, tokens): + import torch + + captured = {} + + def _grab(t, hook): + captured["shape"] = tuple(t.shape) + return t + + with self._eager(bridge): + with torch.no_grad(): + bridge.run_with_hooks( + tokens, + use_cache=False, + fwd_hooks=[(f"blocks.{self.LINEAR_LAYER}.linear_attn.hook_ssm_state", _grab)], + ) + assert captured.get("shape") is not None + assert captured["shape"][1] == self.SEQ_LEN and len(captured["shape"]) == 5 + + def test_state_patch_changes_logits(self, bridge, tokens): + """Zeroing the state trajectory (hook_ssm_state) must change the output.""" + import torch + + def _zero(t, hook): + return torch.zeros_like(t) + + with self._eager(bridge): + with torch.no_grad(): + base = bridge(tokens, use_cache=False) + patched = bridge.run_with_hooks( + tokens, + use_cache=False, + fwd_hooks=[(f"blocks.{self.LINEAR_LAYER}.linear_attn.hook_ssm_state", _zero)], + ) + assert (patched.float() - base.float()).abs().max().item() > 1e-4 + + def test_write_knockout_via_beta_alias(self, bridge, tokens): + """hook_ssm_write (alias -> hook_beta) knockout propagates through the scan.""" + import torch + + def _zero(t, hook): + return torch.zeros_like(t) + + with self._eager(bridge): + with torch.no_grad(): + base = bridge(tokens, use_cache=False) + ko = bridge.run_with_hooks( + tokens, + use_cache=False, + fwd_hooks=[(f"blocks.{self.LINEAR_LAYER}.linear_attn.hook_ssm_write", _zero)], + ) + assert (ko.float() - base.float()).abs().max().item() > 1e-4 + + def test_disabling_restores_default(self, bridge, tokens): + """Toggling eager_scan off must return bit-identical fused-path logits.""" + import torch + + with torch.no_grad(): + first = bridge(tokens, use_cache=False) + with self._eager(bridge): + _ = bridge(tokens, use_cache=False) + after = bridge(tokens, use_cache=False) + assert torch.equal(first, after) + + +def _make_tiny_gqa_bridge(): + """Qwen3_5 bridge with n_v_heads (4) > n_k_heads (2) to exercise the GDN GQA branch. + + The default _make_tiny_bridge uses n_k==n_v==4, so the repeat_interleave GQA + expansion in GatedDeltaNetBridge._hooked_forward never runs. Here Q/K carry 2 + heads and V carries 4, so the recurrence path must expand Q/K 2x. + """ + from unittest.mock import MagicMock + + from transformer_lens.config.transformer_bridge_config import ( + TransformerBridgeConfig, + ) + from transformer_lens.model_bridge import TransformerBridge + from transformer_lens.model_bridge.supported_architectures.qwen3_5 import ( + Qwen3_5ArchitectureAdapter, + ) + + cfg = Qwen3_5TextConfig( + hidden_size=128, + num_hidden_layers=8, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=32, + intermediate_size=256, + vocab_size=512, + full_attention_interval=4, + linear_conv_kernel_dim=4, + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_num_key_heads=2, # < value heads -> GQA expansion required + linear_num_value_heads=4, + rope_parameters={ + "rope_theta": 10000.0, + "partial_rotary_factor": 0.25, + "rope_type": "default", + }, + ) + hf_model = _Qwen3_5ForCausalLM(cfg).eval() + bridge_cfg = TransformerBridgeConfig( + d_model=128, + d_head=32, + n_heads=4, + n_layers=8, + n_ctx=2048, + d_vocab=512, + n_key_value_heads=2, + architecture="Qwen3_5ForCausalLM", + ) + return ( + TransformerBridge(hf_model, Qwen3_5ArchitectureAdapter(bridge_cfg), tokenizer=MagicMock()), + hf_model, + ) + + +@pytest.mark.skipif( + not _QWEN3_5_AVAILABLE, + reason="Qwen3_5TextConfig / Qwen3_5ForCausalLM not available in installed transformers", +) +class TestQwen3_5GatedDeltaNetGQA: + """Cover the GDN GQA expansion (n_v_heads > n_k_heads) in CI, not just env-gated + real-weight tests. Forward parity requires the Q/K repeat_interleave to be correct.""" + + LINEAR_LAYER = 0 + + def test_gqa_branch_engaged_and_parity(self): + import torch + + bridge, hf = _make_tiny_gqa_bridge() + assert hf.config.linear_num_value_heads > hf.config.linear_num_key_heads # GQA active + torch.manual_seed(0) + tokens = torch.randint(0, 512, (1, 10)) + with torch.no_grad(): + bridge_logits = bridge(tokens, use_cache=False) + hf_logits = hf(tokens).logits + rel = (bridge_logits.float() - hf_logits.float()).abs().max().item() / max( + hf_logits.float().abs().max().item(), 1e-8 + ) + assert rel < 1e-4, f"GQA GDN forward parity rel diff {rel:.2e}" + + def test_gqa_hook_head_counts(self): + """hook_q is pre-GQA (n_k_heads); hook_recurrence_out is post-expansion (n_v_heads).""" + import torch + + bridge, hf = _make_tiny_gqa_bridge() + tokens = torch.randint(0, 512, (1, 8)) + with torch.no_grad(): + _, cache = bridge.run_with_cache(tokens, use_cache=False) + p = f"blocks.{self.LINEAR_LAYER}.linear_attn" + assert cache[f"{p}.hook_q"].shape[2] == hf.config.linear_num_key_heads + assert cache[f"{p}.hook_recurrence_out"].shape[2] == hf.config.linear_num_value_heads diff --git a/tests/unit/tools/model_registry/test_verify_gates.py b/tests/unit/tools/model_registry/test_verify_gates.py new file mode 100644 index 000000000..be39abebd --- /dev/null +++ b/tests/unit/tools/model_registry/test_verify_gates.py @@ -0,0 +1,100 @@ +"""Cheap pure-function unit tests for the verify_models control-gates. + +These gates were previously exercised only by multi-hour real-model runs, so a slot-name +drift or a phase-filter regression would slip through CI. They are pure functions, so we +pin them directly: +- ``_phases_to_run`` — restricts requested phases to the adapter's ``applicable_phases`` + (SSM/recurrent families now run all four); phases 7/8 always pass through. +- ``_extract_phase_scores`` — per-phase pass-rate, with SKIPPED results excluded. +- ``_is_ssm_mixer_internal`` — the component-benchmark skip for SSM mixer internals. +""" +from types import SimpleNamespace + +from transformer_lens.benchmarks.component_outputs import _is_ssm_mixer_internal +from transformer_lens.benchmarks.utils import BenchmarkSeverity +from transformer_lens.tools.model_registry.verify_models import ( + _extract_phase_scores, + _phases_to_run, +) + +SSM_ARCHS = ["MambaForCausalLM", "Mamba2ForCausalLM", "NemotronHForCausalLM"] + + +class TestPhasesToRun: + def test_ssm_families_run_all_four_phases(self): + # Regression guard: SSM/hybrid families are no longer skipped or [4]-only. + for arch in SSM_ARCHS: + assert _phases_to_run(arch, [1, 2, 3, 4]) == [1, 2, 3, 4] + + def test_phases_7_and_8_always_pass_through(self): + # 7/8 are gated by is_multimodal/is_audio elsewhere, never filtered here. + assert _phases_to_run("MambaForCausalLM", [1, 7, 8]) == [1, 7, 8] + + def test_unknown_architecture_defaults_to_all_phases(self): + assert _phases_to_run("NotARealArchitecture", [1, 2, 3, 4]) == [1, 2, 3, 4] + + def test_restricted_architecture_is_filtered(self): + # Find any adapter with a restricted applicable_phases and confirm filtering. + from transformer_lens.factories.architecture_adapter_factory import ( + SUPPORTED_ARCHITECTURES, + ) + + restricted = None + for name, adapter in SUPPORTED_ARCHITECTURES.items(): + phases = getattr(adapter, "applicable_phases", [1, 2, 3, 4]) + if phases and set(phases) != {1, 2, 3, 4}: + restricted = (name, phases) + break + if restricted is None: + return # no restricted-phase adapter in the registry; nothing to assert + name, phases = restricted + result = _phases_to_run(name, [1, 2, 3, 4]) + assert result == [p for p in [1, 2, 3, 4] if p in phases] + assert set(result) <= set(phases) + + +def _result(phase, passed, severity, details=None): + return SimpleNamespace(phase=phase, passed=passed, severity=severity, details=details) + + +class TestExtractPhaseScores: + def test_skipped_results_excluded_from_score(self): + # A SKIPPED centering test (passed=False) must NOT drag the phase down — this is + # exactly the SSM P3=90 bug. One real PASS + one SKIPPED → 100. + results = [ + _result(3, True, BenchmarkSeverity.INFO), + _result(3, False, BenchmarkSeverity.SKIPPED), + ] + assert _extract_phase_scores(results)[3] == 100.0 + + def test_mixed_pass_fail_averaged(self): + results = [ + _result(1, True, BenchmarkSeverity.INFO), + _result(1, False, BenchmarkSeverity.DANGER), + ] + assert _extract_phase_scores(results)[1] == 50.0 + + def test_all_skipped_phase_omitted(self): + results = [_result(2, False, BenchmarkSeverity.SKIPPED)] + assert 2 not in _extract_phase_scores(results) + + def test_phase4_uses_quality_score_from_details(self): + results = [_result(4, True, BenchmarkSeverity.INFO, details={"score": 87.5})] + assert _extract_phase_scores(results)[4] == 87.5 + + +class TestIsSSMMixerInternal: + def test_mixer_and_linear_attn_internals_skipped(self): + assert _is_ssm_mixer_internal("blocks.0.mixer.conv1d") + assert _is_ssm_mixer_internal("blocks.5.mixer.in_proj") + assert _is_ssm_mixer_internal("blocks.0.linear_attn.conv1d") + assert _is_ssm_mixer_internal("blocks.0.mixer.conv1d.weight") + + def test_mixer_node_itself_is_tested(self): + # The mixer node (path ending in the slot) is still benchmarked end-to-end. + assert not _is_ssm_mixer_internal("blocks.0.mixer") + assert not _is_ssm_mixer_internal("blocks.0.linear_attn") + + def test_transformer_components_untouched(self): + for path in ("blocks.0.attn.q", "blocks.0.mlp.out", "embed", "unembed", "ln_final"): + assert not _is_ssm_mixer_internal(path) diff --git a/transformer_lens/ActivationCache.py b/transformer_lens/ActivationCache.py index 7d392b6da..52c4c11c7 100644 --- a/transformer_lens/ActivationCache.py +++ b/transformer_lens/ActivationCache.py @@ -17,6 +17,7 @@ class first, including the examples, and then skimming the available methods. Yo from typing import ( TYPE_CHECKING, Any, + Callable, Dict, Iterator, List, @@ -753,6 +754,187 @@ def compute_head_results( # Sum over d_head to get the contribution of each head to the residual stream self.cache_dict[f"blocks.{layer}.attn.hook_result"] = result.sum(dim=-2) + def ssm_layers(self, mixer_type: Optional[Union[type, Tuple[type, ...]]] = None) -> List[int]: + """Return the block indices whose mixer is an SSM / recurrent mixer. + + Family-agnostic and purely structural: finds each block's *realized* SSM + mixer in its variant slot (``.mixer`` / ``.linear_attn`` / …) via + ``find_ssm_mixer``, which excludes a hybrid's passthrough ``.mixer`` on + non-SSM layers (e.g. NemotronH attention/MLP/MoE) — no dependence on + ``cfg.layers_block_type``. + + Args: + mixer_type: Optional concrete bridge class to filter to (e.g. + ``SSM2MixerBridge`` for only Mamba-2 layers). + + Returns: + Ascending list of SSM block indices. + """ + from transformer_lens.model_bridge.generalized_components.ssm_protocol import ( + find_ssm_mixer, + ) + + bridge = self.model + layers: List[int] = [] + for i, block in enumerate(bridge.blocks): + mixer = find_ssm_mixer(block) + if mixer is None: + continue + if mixer_type is not None and not isinstance(mixer, mixer_type): + continue + layers.append(i) + return layers + + def _over_ssm_layers( + self, + fn: Callable[[Any, int], torch.Tensor], + mixer_type: Optional[Union[type, Tuple[type, ...]]] = None, + ) -> Union[torch.Tensor, Dict[int, torch.Tensor]]: + """Apply ``fn(mixer, layer_idx)`` over every SSM layer; stack or dict. + + The single SSM layer-enumeration used by both ``compute_ssm_state`` and + ``compute_ssm_effective_attention``. Returns a stacked tensor (dim 0 = + layer) when *every* block is an SSM layer, else a ``{layer_idx: result}`` + dict over the SSM layers (heterogeneous hybrids). + """ + from transformer_lens.model_bridge.generalized_components.ssm_protocol import ( + find_ssm_mixer, + ) + + bridge = self.model + indices = self.ssm_layers(mixer_type=mixer_type) + if not indices: + raise RuntimeError( + "No SSM mixer layers found. Use an SSM / hybrid bridge and " + "run_with_cache first (gated-delta-net needs use_cache=False)." + ) + results: Dict[int, torch.Tensor] = { + i: fn(cast(Any, find_ssm_mixer(bridge.blocks[i])), i) for i in indices + } + if indices == list(range(len(bridge.blocks))): + return torch.stack([results[i] for i in indices], dim=0) + return results + + def compute_ssm_effective_attention( + self, + layer: Optional[int] = None, + include_dt_scaling: bool = False, + per_state_coord: bool = False, + ) -> Union[torch.Tensor, Dict[int, torch.Tensor]]: + """Materialize SSM effective attention for one or all SSM layers, any family. + + The single discovery surface for effective attention across Mamba-1, + Mamba-2, and gated-delta-net layers; dispatches to each layer's mixer + ``compute_effective_attention`` regardless of family. Family-specific + options are forwarded only to mixers whose signature accepts them. + + Args: + layer: Specific block index, or None for every SSM layer. + include_dt_scaling: Forwarded to Mamba-1/Mamba-2 mixers (the + reconstruction form); gated-delta-net does not accept it. + per_state_coord: Forwarded to Mamba-1 only (per-state-coordinate + matrices). Setting it True for another family raises ValueError. + + Returns: + A per-layer matrix for a single ``layer``; for ``layer=None`` a + stacked tensor (dim 0 = layer) when every block is an SSM layer, else + a ``{layer_idx: matrix}`` dict over the SSM layers. + + Raises: + TypeError: If the requested ``layer`` has no SSM mixer. + ValueError: If an option is set that the target mixer does not support. + RuntimeError: If ``layer=None`` finds no SSM layers. + """ + import inspect + + from transformer_lens.model_bridge.generalized_components.ssm_protocol import ( + find_ssm_mixer, + ) + + def _call(mixer: Any, layer_idx: int) -> torch.Tensor: + fn = cast(Any, mixer).compute_effective_attention + params = inspect.signature(fn).parameters + kwargs: Dict[str, Any] = {} + for name, value in ( + ("include_dt_scaling", include_dt_scaling), + ("per_state_coord", per_state_coord), + ): + if name in params: + kwargs[name] = value + elif value: + raise ValueError( + f"{type(mixer).__name__}.compute_effective_attention does not " + f"support {name}=True." + ) + result: torch.Tensor = fn(cache=self, layer_idx=layer_idx, **kwargs) + return result + + if layer is not None: + single = find_ssm_mixer(self.model.blocks[layer]) + if single is None: + raise TypeError(f"Block {layer} has no SSM mixer (no compute_effective_attention).") + return _call(single, layer) + return self._over_ssm_layers(_call) + + def compute_ssm_state( + self, + layer: Optional[int] = None, + time_step: Optional[int] = None, + ) -> Union[torch.Tensor, Dict[int, torch.Tensor]]: + """Reconstruct the recurrent SSM state ``S`` from this cache. + + The single discovery surface for recurrent state across families — + ``SSMMixerBridge`` (Mamba-1), ``SSM2MixerBridge`` (Mamba-2) and + ``GatedDeltaNetBridge`` (gated delta rule) each reconstruct it — mirroring + ``compute_head_results``. Read-only post-hoc reconstruction from cached + hooks (no forward re-run); requires an SSM / SSM-hybrid bridge cached via + ``run_with_cache`` (gated-delta-net additionally needs ``use_cache=False`` + so its interior hooks fire). See the mixer's ``compute_ssm_state`` for the + recurrence, shapes, and the ``time_step`` memory bound; state shape is + family-specific, so ``layer=None`` returns a dict except when every block + shares one mixer type. + + Args: + layer: Specific block index, or None for every SSM-state layer. + time_step: Optional single position (memory-bounded); None for all. + + Returns: + A per-layer state tensor for a single ``layer``; for ``layer=None`` a + stacked tensor (dim 0 = layer) when every block is an SSM-state layer, + else a ``{layer_idx: state}`` dict over those layers. + + Raises: + TypeError: If the requested ``layer`` has no state-reconstructing mixer. + RuntimeError: If ``layer=None`` finds no such layers. + """ + from transformer_lens.model_bridge.generalized_components import ( + GatedDeltaNetBridge, + SSM2MixerBridge, + SSMMixerBridge, + ) + from transformer_lens.model_bridge.generalized_components.ssm_protocol import ( + find_ssm_mixer, + ) + + # Every recurrent family reconstructs S_t from its cached interior hooks. + state_mixers = (SSMMixerBridge, SSM2MixerBridge, GatedDeltaNetBridge) + + def _call(mixer: Any, layer_idx: int) -> torch.Tensor: + state: torch.Tensor = cast(Any, mixer).compute_ssm_state( + self, layer_idx=layer_idx, time_step=time_step + ) + return state + + if layer is not None: + single = find_ssm_mixer(self.model.blocks[layer]) + if not isinstance(single, state_mixers): + raise TypeError( + f"Block {layer} has no state-reconstructing SSM mixer; " + "compute_ssm_state supports Mamba-1 / Mamba-2 / gated-delta-net." + ) + return _call(single, layer) + return self._over_ssm_layers(_call, mixer_type=state_mixers) + def stack_head_results( self, layer: int = -1, diff --git a/transformer_lens/benchmarks/component_outputs.py b/transformer_lens/benchmarks/component_outputs.py index 91efc1130..f23315dee 100644 --- a/transformer_lens/benchmarks/component_outputs.py +++ b/transformer_lens/benchmarks/component_outputs.py @@ -196,6 +196,16 @@ def print_detailed_analysis(self) -> None: print("=" * 80 + "\n") +def _is_ssm_mixer_internal(component_path: str) -> bool: + """True for a submodule *inside* an SSM/recurrent mixer slot (``.mixer``/``.linear_attn``). + + Such submodules wrap the identical HF module (parity is covered by forward_pass_logits) + and take SSM-internal shapes, not the ``[b, seq, d_model]`` residual the isolated harness + feeds — so they are skipped. The mixer node itself (path ending in the slot) is not. + """ + return any(slot in component_path.split(".")[:-1] for slot in ("mixer", "linear_attn")) + + class ComponentBenchmarker: """Benchmarking utility for testing TransformerBridge components against HuggingFace.""" @@ -387,6 +397,11 @@ def _test_component_recursive( if component_path in skip_components: return + # SSM/recurrent mixer internal submodules can't be tested in isolation (see + # _is_ssm_mixer_internal); the mixer node itself is still tested against HF. + if _is_ssm_mixer_internal(component_path): + return + # Skip MLP components that don't exist as separate modules in HF (name=None) # These are virtual components where fc1/fc2 are directly on the layer # Component testing doesn't work for these because get_component returns the parent layer diff --git a/transformer_lens/benchmarks/hook_registration.py b/transformer_lens/benchmarks/hook_registration.py index d74932f08..77b819d25 100644 --- a/transformer_lens/benchmarks/hook_registration.py +++ b/transformer_lens/benchmarks/hook_registration.py @@ -394,6 +394,12 @@ def benchmark_gated_hooks_fire( name.rsplit(".", 1)[-1] == stem for name in target_hook_names ): failed.append((flag_name, stem)) + except NotImplementedError as e: + # Some architectures reject the split path only at forward time (e.g. + # gated q_proj: Qwen3.5 / Qwen3-Next). Treat like the setter-time gate — + # skipped (applicability gate is intentional), not failed. + tested_flags.remove(flag_name) + skipped.append((flag_name, str(e).split("\n", 1)[0][:120])) finally: setter(False) diff --git a/transformer_lens/benchmarks/weight_processing.py b/transformer_lens/benchmarks/weight_processing.py index 9ffe1c08f..d7368faff 100644 --- a/transformer_lens/benchmarks/weight_processing.py +++ b/transformer_lens/benchmarks/weight_processing.py @@ -559,11 +559,12 @@ def benchmark_attention_output_centering( attn_blocks = bridge.blocks_with("attn") if not attn_blocks: + # Attention-less (pure SSM) or hybrids whose attention lives inside a + # passthrough mixer (NemotronH): nothing to center — skip, don't fail. return BenchmarkResult( name="attention_output_centering", - severity=BenchmarkSeverity.WARNING, - message="No blocks have attention submodule", - passed=False, + severity=BenchmarkSeverity.SKIPPED, + message="No blocks expose an attention submodule (SSM / passthrough-mixer hybrid)", ) # Check W_O accessibility on first attention block @@ -642,17 +643,20 @@ def benchmark_mlp_output_centering( from transformer_lens.model_bridge.generalized_components.moe import MoEBridge mlp_module = None - block = bridge.blocks[0] - for name in ("mlp", "shared_mlp"): - if name in block._modules: - mlp_module = block._modules[name] + for block in bridge.blocks: + for name in ("mlp", "shared_mlp"): + if name in block._modules: + mlp_module = block._modules[name] + break + if mlp_module is not None: break if mlp_module is None: + # Pure SSM, or a hybrid whose MLP lives inside a passthrough mixer: + # no standalone MLP to center — skip, don't fail. return BenchmarkResult( name="mlp_output_centering", - severity=BenchmarkSeverity.WARNING, - message="No MLP submodule found on block 0", - passed=False, + severity=BenchmarkSeverity.SKIPPED, + message="No block exposes an MLP submodule (SSM / passthrough-mixer hybrid)", ) if isinstance(mlp_module, MoEBridge): diff --git a/transformer_lens/model_bridge/generalized_components/__init__.py b/transformer_lens/model_bridge/generalized_components/__init__.py index b272be36e..ac3dd5901 100644 --- a/transformer_lens/model_bridge/generalized_components/__init__.py +++ b/transformer_lens/model_bridge/generalized_components/__init__.py @@ -99,6 +99,11 @@ from transformer_lens.model_bridge.generalized_components.ssm2_mixer import SSM2MixerBridge from transformer_lens.model_bridge.generalized_components.ssm_block import SSMBlockBridge from transformer_lens.model_bridge.generalized_components.ssm_mixer import SSMMixerBridge +from transformer_lens.model_bridge.generalized_components.ssm_protocol import ( + SSMMixerProtocol, + SSMStateHookMixin, + find_ssm_mixer, +) from transformer_lens.model_bridge.generalized_components.symbolic import SymbolicBridge from transformer_lens.model_bridge.generalized_components.t5_block import T5BlockBridge from transformer_lens.model_bridge.generalized_components.t5gemma_decoder_block import ( @@ -155,6 +160,9 @@ "SiglipVisionEncoderBridge", "SiglipVisionEncoderLayerBridge", "SSM2MixerBridge", + "SSMMixerProtocol", + "SSMStateHookMixin", + "find_ssm_mixer", "SSMBlockBridge", "SSMMixerBridge", "VisionProjectionBridge", diff --git a/transformer_lens/model_bridge/generalized_components/gated_delta_net.py b/transformer_lens/model_bridge/generalized_components/gated_delta_net.py index 1e13fe4bf..f71d4320e 100644 --- a/transformer_lens/model_bridge/generalized_components/gated_delta_net.py +++ b/transformer_lens/model_bridge/generalized_components/gated_delta_net.py @@ -4,21 +4,22 @@ states. Falls back to HF native forward during autoregressive generation where cache state management is required. """ -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import Any, Dict, Optional import torch import torch.nn.functional as F +from transformer_lens.ActivationCache import ActivationCache from transformer_lens.hook_points import HookPoint from transformer_lens.model_bridge.generalized_components.base import ( GeneralizedComponent, ) - -if TYPE_CHECKING: - from transformer_lens.ActivationCache import ActivationCache +from transformer_lens.model_bridge.generalized_components.ssm_protocol import ( + SSMStateHookMixin, +) -class GatedDeltaNetBridge(GeneralizedComponent): +class GatedDeltaNetBridge(SSMStateHookMixin, GeneralizedComponent): """Bridge for GatedDeltaNet linear-attention with full hook decomposition. Hooks (prefill, in execution order): @@ -37,6 +38,11 @@ class GatedDeltaNetBridge(GeneralizedComponent): per v-head [batch, seq, n_v_heads] hook_recurrence_out: output of linear recurrence [batch, seq, n_v_heads, head_v_dim] hook_gate_input: z tensor (pre-silu) for GatedRMSNorm [batch, seq, n_v_heads, head_v_dim] + hook_ssm_state: recurrent state trajectory S_t [batch, seq, n_v_heads, head_k_dim, + head_v_dim] — fires ONLY on the opt-in eager-scan path (eager_scan=True), + which swaps the fused kernel for a Python delta-rule scan so S_t can be + read/patched. hook_ssm_write (alias -> hook_beta) is the write strength and + propagates through the scan. hook_out: final output to residual stream [batch, seq, d_model] During generation (cache_params present), only hook_in/hook_out fire. @@ -48,6 +54,15 @@ class GatedDeltaNetBridge(GeneralizedComponent): hook_aliases = { "hook_linear_attn_in": "hook_in", "hook_linear_attn_out": "hook_out", + # Canonical SSM vocabulary (additive) — maps onto GDN's existing hooks so + # interp tools can find these quantities by the same name across families. + # Semantic mapping: q reads the state (~C), k writes to it (~B), the gate + # g is the decay, beta is the per-step write strength. + "hook_ssm_out": "hook_out", + "hook_ssm_C": "hook_q", + "hook_ssm_B": "hook_k", + "hook_ssm_decay": "hook_log_decay", + "hook_ssm_write": "hook_beta", } property_aliases = { @@ -81,6 +96,8 @@ def __init__( # Recurrence output + gated norm input self.hook_recurrence_out = HookPoint() self.hook_gate_input = HookPoint() + # hook_ssm_state + eager_scan come from SSMStateHookMixin; hook_ssm_write is + # an alias onto hook_beta (the delta rule's write is state-dependent). def forward(self, *args: Any, **kwargs: Any) -> Any: if self.original_component is None: @@ -201,17 +218,23 @@ def _hooked_forward(self, *args: Any, **kwargs: Any) -> Any: query = query.repeat_interleave(repeat, dim=2) key = key.repeat_interleave(repeat, dim=2) - # --- Core linear recurrence (opaque fused kernel) --- - core_out, _ = hf.chunk_gated_delta_rule( - query, - key, - value, - g=g, - beta=beta, - initial_state=None, - output_final_state=False, - use_qk_l2norm_in_kernel=True, - ) + # --- Core linear recurrence --- + # Default: HF's opaque fused kernel. eager_scan: a Python delta-rule scan + # that fires hook_ssm_state so the state trajectory can be intervened on. + if self.eager_scan: + _, core_out = self._gated_delta_scan(query, key, value, g, beta, fire_hook=True) + core_out = core_out.to(value.dtype) + else: + core_out, _ = hf.chunk_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=True, + ) core_out = self.hook_recurrence_out(core_out) # --- Gated RMSNorm: norm(core_out) * silu(z) --- @@ -227,30 +250,176 @@ def _hooked_forward(self, *args: Any, **kwargs: Any) -> Any: output = hf.out_proj(core_out) return self.hook_out(output) + @staticmethod + def _l2norm(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: + """Match the kernel's internal Q/K L2 normalization (use_qk_l2norm_in_kernel).""" + return x / torch.sqrt((x * x).sum(-1, keepdim=True) + eps) + + def _gated_delta_scan( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + fire_hook: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Eager gated-delta-rule recurrence — the readable form of the fused kernel. + + Reproduces HF ``torch_recurrent_gated_delta_rule``: L2-normalize Q/K, scale + Q by ``1/sqrt(head_k_dim)``, then per step:: + + S_t = exp(g_t) · S_{t-1} (decay) + delta_t = (v_t - S_t^T k_t) · beta_t (delta rule: remove then write) + S_t = S_t + k_t ⊗ delta_t -> hook_ssm_state trajectory + o_t = S_t^T q_t + + Q/K are assumed already GQA-expanded to ``n_v_heads``. Matches the fused + kernel only to fp tolerance (≈1e-6 fp32), never bit-for-bit; O(seq) Python. + + Args: + query, key: ``[batch, seq, n_v_heads, head_k_dim]`` (post-conv, pre-norm). + value: ``[batch, seq, n_v_heads, head_v_dim]``. + g: log-space decay ``[batch, seq, n_v_heads]`` (NEGATIVE; decay = exp(g)). + beta: write strength ``[batch, seq, n_v_heads]``. + fire_hook: if True, fire ``hook_ssm_state`` on the trajectory and + recompute ``o_t`` from the (possibly patched) state. + + Returns: + ``(state_traj, core_out)`` — state ``[batch, seq, n_v_heads, head_k_dim, + head_v_dim]`` and output ``[batch, seq, n_v_heads, head_v_dim]``. + """ + q = self._l2norm(query.float()) + k = self._l2norm(key.float()) + v = value.float() + q = q * (q.shape[-1] ** -0.5) + g_f = g.float() + beta_f = beta.float() + + batch, seq_len, n_v, k_dim = q.shape + v_dim = v.shape[-1] + state = torch.zeros(batch, n_v, k_dim, v_dim, dtype=q.dtype, device=q.device) + states = [] + outs = [] + for t in range(seq_len): + q_t, k_t, v_t = q[:, t], k[:, t], v[:, t] # [batch, n_v, dim] + decay = g_f[:, t].exp()[:, :, None, None] # [batch, n_v, 1, 1] + beta_t = beta_f[:, t][:, :, None] # [batch, n_v, 1] + state = state * decay + kv_mem = (state * k_t[:, :, :, None]).sum(dim=-2) # [batch, n_v, v_dim] + delta = (v_t - kv_mem) * beta_t + state = state + k_t[:, :, :, None] * delta[:, :, None, :] + states.append(state) + outs.append((state * q_t[:, :, :, None]).sum(dim=-2)) + + state_traj = torch.stack(states, dim=1) # [batch, seq, n_v, k_dim, v_dim] + if fire_hook: + state_traj = self.hook_ssm_state(state_traj) + # Recompute o_t = S_t^T q_t from the (possibly patched) trajectory. + core_out = torch.einsum("bshkv,bshk->bshv", state_traj, q) + else: + core_out = torch.stack(outs, dim=1) # [batch, seq, n_v, v_dim] + return state_traj, core_out + + def compute_ssm_state( + self, + cache: ActivationCache, + layer_idx: int, + time_step: Optional[int] = None, + ) -> torch.Tensor: + """Reconstruct the recurrent state ``S`` of the gated delta rule from cache. + + Read-only: replays the eager delta-rule scan (``_gated_delta_scan``) on the + cached hook_q/k/v/beta/log_decay — no ``forward()`` re-run. Faithful (the + full delta rule, key-removal included), unlike ``compute_effective_attention`` + which is a gated-linear-attention heuristic that drops key removal. + + Requires the interior hooks, which fire only on the hooked prefill path: + call ``run_with_cache(tokens, use_cache=False)`` so ``cache_params`` is None. + + On padded batches the cached hooks are unmasked, so ``S`` is exact only at + non-pad positions; pad-position state is out of contract. + + Args: + cache: ActivationCache from ``run_with_cache(..., use_cache=False)``. + layer_idx: Block index for this linear_attn layer. + time_step: If given, return only ``S`` at that position + (``[batch, n_v_heads, head_k_dim, head_v_dim]``). None returns + every step. + + Returns: + ``[batch, seq, n_v_heads, head_k_dim, head_v_dim]`` for all steps, or + ``[batch, n_v_heads, head_k_dim, head_v_dim]`` for a single ``time_step``. + + Memory is O(batch · n_v_heads · seq · head_k_dim · head_v_dim); pass + ``time_step`` (or short sequences) when that is too large. + """ + prefix = f"blocks.{layer_idx}.linear_attn" + q_key = f"{prefix}.hook_q" + k_key = f"{prefix}.hook_k" + v_key = f"{prefix}.hook_v" + beta_key = f"{prefix}.hook_beta" + decay_key = f"{prefix}.hook_log_decay" + + for key in (q_key, k_key, v_key, beta_key, decay_key): + if key not in cache: + raise RuntimeError( + f"compute_ssm_state needs {key!r} in cache. Run " + "run_with_cache(tokens, use_cache=False) on the bridge first." + ) + + q = cache[q_key].float() # [batch, seq, n_k_heads, head_k_dim] + k = cache[k_key].float() + v = cache[v_key].float() # [batch, seq, n_v_heads, head_v_dim] + beta = cache[beta_key].float() # [batch, seq, n_v_heads] + g = cache[decay_key].float() # [batch, seq, n_v_heads] + + # GQA expansion to n_v_heads (Q/K carry n_k_heads). + n_v = v.shape[2] + if q.shape[2] < n_v: + repeat = n_v // q.shape[2] + q = q.repeat_interleave(repeat, dim=2) + k = k.repeat_interleave(repeat, dim=2) + + state_traj, _ = self._gated_delta_scan(q, k, v, g, beta, fire_hook=False) + if time_step is not None: + return state_traj[:, time_step] + return state_traj + def compute_effective_attention( self, - cache: "ActivationCache", + cache: ActivationCache, layer_idx: int, ) -> torch.Tensor: - """Materialize the effective attention matrix from cached hook values. + """Materialize a heuristic effective-attention matrix from cached hooks. - The gated delta rule recurrence is:: + Uses the gated-linear-attention form of the recurrence (the exact gated + *delta* rule additionally removes the key being written):: - S_t = exp(g_t) * S_{t-1} + beta_t * v_t @ k_t^T + S_t ≈ exp(g_t) * S_{t-1} + beta_t * v_t @ k_t^T o_t = S_t^T @ q_t + M[i,j] = (q_i^T @ k_j) * beta_j * prod_{t=j+1}^{i} exp(g_t) - The effective attention M[i,j] = contribution of input j to output i:: + so ``M`` is an interpretability heuristic, not a faithful output decomposition. - M[i,j] = (q_i^T @ k_j) * beta_j * prod_{t=j+1}^{i} exp(g_t) + Requires the interior hooks (hook_q/k/beta/log_decay), which fire only on + the hooked prefill path: call ``run_with_cache(tokens, use_cache=False)`` + so ``cache_params`` is None. The default cached path exposes only + hook_in/hook_out and this method then raises. + + **Measured divergence (tiny random-init test fixture, seed-stable):** - **Approximation note:** The fused kernel applies L2-normalization to Q - and K internally (``use_qk_l2norm_in_kernel=True``). The hooked Q/K are - pre-normalization, so this reconstruction diverges when Q/K norms vary - significantly across positions/heads. Accuracy is best when Q/K norms - are roughly uniform (common after training converges). + - *L2-norm gap.* The fused kernel L2-normalizes Q/K internally + (``use_qk_l2norm_in_kernel=True``) but the hooked Q/K are pre-norm, so + ``M`` differs from the normalized form by ≈1.0 relative when Q/K norms + are small/non-uniform (the random-init regime); the gap shrinks toward 0 + as norms equalize after training. + - *Delta-rule omission.* Even with normalized Q/K, ``M @ V`` reconstructs + the fused-kernel ``hook_recurrence_out`` only to O(1) relative error + because the key-removal term is dropped. Args: - cache: ActivationCache from ``run_with_cache``. + cache: ActivationCache from ``run_with_cache(..., use_cache=False)``. layer_idx: Block index for this linear_attn layer. Returns: diff --git a/transformer_lens/model_bridge/generalized_components/ssm2_mixer.py b/transformer_lens/model_bridge/generalized_components/ssm2_mixer.py index b642a0437..3f93f46bd 100644 --- a/transformer_lens/model_bridge/generalized_components/ssm2_mixer.py +++ b/transformer_lens/model_bridge/generalized_components/ssm2_mixer.py @@ -1,15 +1,29 @@ """Wrap-don't-reimplement bridge for HF's Mamba2Mixer, plus SSD effective attention.""" -from typing import Any +from typing import Any, NamedTuple, Optional import torch from transformer_lens.ActivationCache import ActivationCache +from transformer_lens.hook_points import HookPoint from transformer_lens.model_bridge.generalized_components.base import ( GeneralizedComponent, ) +from transformer_lens.model_bridge.generalized_components.ssm_protocol import ( + SSMStateHookMixin, +) + + +class _SSDTerms(NamedTuple): + """SSD intermediates reconstructed from cached Mamba-2 hooks (HF discretization).""" + dt: torch.Tensor # [batch, seq, num_heads] + L: torch.Tensor # [batch, num_heads, seq, seq] causal decay (upper triangle zero) + x: torch.Tensor # [batch, seq, num_heads, head_dim] SSM input per head + B: torch.Tensor # [batch, seq, num_heads, state] (group-expanded) + C: torch.Tensor # [batch, seq, num_heads, state] (group-expanded) -class SSM2MixerBridge(GeneralizedComponent): + +class SSM2MixerBridge(SSMStateHookMixin, GeneralizedComponent): """Opaque wrapper around Mamba-2's Mamba2Mixer. Structural differences from Mamba-1: @@ -29,24 +43,41 @@ class SSM2MixerBridge(GeneralizedComponent): "hook_in_proj": "in_proj.hook_out", "hook_conv": "conv1d.hook_out", "hook_inner_norm": "inner_norm.hook_out", + # Canonical SSM vocabulary (additive). Mamba-2 fuses B/C/dt into the + # in_proj/conv1d outputs, so only the mixer output has a granular hook by + # default; B/C/dt/decay are reconstructed via compute_effective_attention / + # compute_ssm_state. hook_ssm_write / hook_ssm_state are real HookPoints + # that fire only on the opt-in eager-scan intervention path. "hook_ssm_out": "hook_out", } + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) # mixin adds hook_ssm_state + eager_scan + # Real per-step write term dt·(x⊗B); fires only on the eager-scan path. + self.hook_ssm_write = HookPoint() + def forward(self, *args: Any, **kwargs: Any) -> Any: - """Hook the input, delegate to HF torch_forward, hook the output.""" + """Hook the input, run HF torch_forward (or the eager scan), hook the output.""" if self.original_component is None: raise RuntimeError( f"Original component not set for {self.name}. " "Call set_original_component() first." ) + hidden_states: Optional[torch.Tensor] = None if len(args) > 0 and isinstance(args[0], torch.Tensor): - hooked = self.hook_in(args[0]) - args = (hooked,) + args[1:] + hidden_states = self.hook_in(args[0]) + args = (hidden_states,) + args[1:] elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor): - kwargs["hidden_states"] = self.hook_in(kwargs["hidden_states"]) + hidden_states = self.hook_in(kwargs["hidden_states"]) + kwargs["hidden_states"] = hidden_states - output = self.original_component(*args, **kwargs) + # Eager-scan slow path: opt-in flag, prefill only (no cache_params). Keyed + # on explicit state, never on hook-registry introspection. + if self.eager_scan and hidden_states is not None and kwargs.get("cache_params") is None: + output: Any = self._eager_scan_forward(hidden_states, kwargs.get("attention_mask")) + else: + output = self.original_component(*args, **kwargs) if isinstance(output, tuple) and len(output) > 0: first = output[0] @@ -57,6 +88,94 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: return self.hook_out(output) return output + def _eager_scan_forward( + self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] + ) -> torch.Tensor: + """Reimplement the Mamba-2 mixer prefill forward with an eager Python scan. + + Reuses HF's in_proj / conv1d / inner-norm / out_proj submodules (so their + hooks still fire) and reimplements ONLY the SSD recurrence:: + + write_t = dt_t · (x_t ⊗ B_t) -> hook_ssm_write [b, seq, h, d, n] + S_t = exp(dt_t·A) · S_{t-1} + write_t + S = stack_t S_t -> hook_ssm_state [b, seq, h, d, n] + y_t = C_t · S_t (recomputed from the hooked state) + + Intervening on hook_ssm_write re-runs the recurrence (the change propagates + to every later state); hook_ssm_state is the post-scan trajectory, so a + patch there changes only the same-position output y_t = C_t·S_t, not the + forward recurrence — patch hook_ssm_write for a propagating state edit. + + Kernel-divergence caveat: this eager scan reproduces HF's fused chunk scan + only to floating-point tolerance (≈1e-6 fp32), never bit-for-bit, on any + family. It is O(seq) Python and materializes an O(b·seq·h·d·n) write/state + tensor — orders of magnitude slower/heavier than the kernel. Prefill only. + """ + oc: Any = self.original_component + batch, seq_len, _ = hidden_states.shape + intermediate, num_heads, head_dim = oc.intermediate_size, oc.num_heads, oc.head_dim + n_groups, state = oc.n_groups, oc.ssm_state_size + + def _mask_pad(states: torch.Tensor) -> torch.Tensor: + # Mirror HF apply_mask_to_padding_states (no-op for batch 1 / unpadded). + if attention_mask is not None and attention_mask.shape[1] > 1 and batch > 1: + return (states * attention_mask[:, :, None]).to(states.dtype) + return states + + # 1-2. Projection + conv — reuse the HF submodule bridges (their hooks fire). + # Match HF: mask padding before in_proj and after conv (before B/C split). + projected = oc.in_proj(_mask_pad(hidden_states)) + d_mlp = (projected.shape[-1] - 2 * intermediate - 2 * n_groups * state - num_heads) // 2 + _, _, gate, hidden_B_C, dt = projected.split( + [d_mlp, d_mlp, intermediate, oc.conv_dim, num_heads], dim=-1 + ) + hidden_B_C = oc.act(oc.conv1d(hidden_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2)) + hidden_B_C = _mask_pad(hidden_B_C) + x_flat, B_flat, C_flat = hidden_B_C.split( + [intermediate, n_groups * state, n_groups * state], dim=-1 + ) + + x = x_flat.reshape(batch, seq_len, num_heads, head_dim).float() + heads_per_group = num_heads // n_groups + B = ( + B_flat.reshape(batch, seq_len, n_groups, state) + .repeat_interleave(heads_per_group, dim=2) + .float() + ) + C = ( + C_flat.reshape(batch, seq_len, n_groups, state) + .repeat_interleave(heads_per_group, dim=2) + .float() + ) + + dt = torch.nn.functional.softplus(dt.float() + self.dt_bias.float()) + dt = torch.clamp(dt, float(oc.time_step_limit[0]), float(oc.time_step_limit[1])) + A = -torch.exp(self.A_log.float()) # [num_heads] + + # 3. Eager recurrence with intervention hooks. + writes = dt[:, :, :, None, None] * x[:, :, :, :, None] * B[:, :, :, None, :] + writes = self.hook_ssm_write(writes) # [batch, seq, heads, head_dim, state] + + ssm_state = torch.zeros( + batch, num_heads, head_dim, state, dtype=writes.dtype, device=writes.device + ) + states = [] + for t in range(seq_len): + decay = torch.exp(dt[:, t, :] * A[None, :]) # [batch, heads] + ssm_state = decay[:, :, None, None] * ssm_state + writes[:, t] + states.append(ssm_state) + state_traj = torch.stack(states, dim=1) # [batch, seq, heads, head_dim, state] + state_traj = self.hook_ssm_state(state_traj) + + y = torch.einsum("bthn,bthdn->bthd", C, state_traj) + y = y + self.D.float()[None, None, :, None] * x + y = y.reshape(batch, seq_len, intermediate) + + # 4. Gated norm + output projection — reuse the HF submodule bridges. + scan_output = oc.norm(y.to(hidden_states.dtype), gate) + contextualized: torch.Tensor = oc.out_proj(scan_output) + return contextualized + def compute_effective_attention( self, cache: ActivationCache, @@ -87,6 +206,69 @@ def compute_effective_attention( Cost is O(batch · num_heads · seq_len²); use on short sequences (≤2k). """ + terms = self._ssd_terms(cache, layer_idx) + + # CB[b, h, i, j] = + CB = torch.einsum("bihs,bjhs->bhij", terms.C, terms.B) + M = terms.L * CB # [batch, num_heads, seq, seq] + + if include_dt_scaling: + # Multiply column j by dt[j, h] to absorb the B discretization + M = M * terms.dt.permute(0, 2, 1)[:, :, None, :] + + return M + + def compute_ssm_state( + self, + cache: ActivationCache, + layer_idx: int, + time_step: Optional[int] = None, + ) -> torch.Tensor: + """Reconstruct the recurrent SSM state ``S`` from cached hook values. + + Mamba-2's recurrence is ``S_t = dA_t · S_{t-1} + dt_t · (x_t ⊗ B_t)`` with + ``dA_t[h] = exp(dt_t[h] · A[h])``, so post-hoc:: + + S_t[h, p, n] = sum_{j<=t} L[t, j] · dt_j[h] · x_j[h, p] · B_j[h, n] + + where ``L`` is the same causal decay matrix used by + ``compute_effective_attention``. Read-only: no ``forward()`` re-run. + Verify with ``y_t = C_t · S_t + D · x_t == inner_norm.hook_in``. + + On padded batches the cached hooks are unmasked, so ``S`` is exact only at + non-pad positions; pad-position state is out of contract. + + Args: + cache: ActivationCache from ``run_with_cache`` with this layer's + in_proj and conv1d hooks. + layer_idx: Block index for this mixer. + time_step: If given, return only ``S`` at that position + (``[batch, num_heads, head_dim, state]``); avoids materializing + the full tensor. None returns every step. + + Returns: + ``[batch, num_heads, seq, head_dim, state]`` for all steps, or + ``[batch, num_heads, head_dim, state]`` for a single ``time_step``. + + Memory is O(batch · num_heads · seq · head_dim · state) for the full + tensor; pass ``time_step`` (or short sequences) when that is too large. + """ + terms = self._ssd_terms(cache, layer_idx) + dtx = terms.dt[..., None] * terms.x # [batch, seq, num_heads, head_dim] + if time_step is not None: + # S_t = sum_{j<=t} L[t, j] · (dt_j x_j) ⊗ B_j + return torch.einsum("bhj,bjhp,bjhn->bhpn", terms.L[:, :, time_step, :], dtx, terms.B) + return torch.einsum("bhtj,bjhp,bjhn->bhtpn", terms.L, dtx, terms.B) + + def _ssd_terms(self, cache: ActivationCache, layer_idx: int) -> _SSDTerms: + """Reconstruct the shared SSD intermediates (dt, decay L, per-head x/B/C). + + Mirrors HF Mamba2Mixer's discretization from the cached in_proj/conv1d + outputs. Dims come from the wrapped HF mixer, not cfg: on a hybrid the + shared cfg holds the *attention* dims (cfg.n_heads etc.), while the HF + mixer always carries the true Mamba dims (as do A_log/dt_bias, already + read off the module via __getattr__). cfg is the fallback. + """ if self.config is None: raise RuntimeError("SSM2MixerBridge.config must be set") @@ -94,76 +276,59 @@ def compute_effective_attention( conv1d_key = f"blocks.{layer_idx}.mixer.conv1d.hook_out" if in_proj_key not in cache or conv1d_key not in cache: raise RuntimeError( - f"compute_effective_attention needs {in_proj_key!r} and " - f"{conv1d_key!r} in cache. Run `run_with_cache()` on the bridge " - "before calling this method." + f"SSD reconstruction needs {in_proj_key!r} and {conv1d_key!r} in " + "cache. Run `run_with_cache()` on the bridge before calling this." ) cfg = self.config - num_heads: int = cfg.n_heads - head_dim: int = cfg.d_head - intermediate_size: int = getattr(cfg, "intermediate_size", num_heads * head_dim) - state_size: int = getattr(cfg, "state_size", 128) - n_groups: int = getattr(cfg, "n_groups", 1) - - # Mirror HF's tuple convention so downstream equality checks stay consistent - time_step_limit = getattr(cfg, "time_step_limit", (0.0, float("inf"))) - time_step_min = float(time_step_limit[0]) - time_step_max = float(time_step_limit[1]) - - in_proj_out = cache[in_proj_key] # [batch, seq, proj_size] - conv1d_out = cache[conv1d_key] # [batch, conv_dim, seq + conv_kernel - 1] + oc = self.original_component + + def _mamba_dim(module_attr: str, cfg_attr: str, default: Any) -> Any: + if oc is not None: + val = getattr(oc, module_attr, None) + if val is not None: + return val + return getattr(cfg, cfg_attr, default) + + num_heads = int(_mamba_dim("num_heads", "n_heads", 0)) + head_dim = int(_mamba_dim("head_dim", "d_head", 0)) + intermediate_size = int( + _mamba_dim("intermediate_size", "intermediate_size", num_heads * head_dim) + ) + state_size = int(_mamba_dim("ssm_state_size", "state_size", 128)) + n_groups = int(_mamba_dim("n_groups", "n_groups", 1)) + time_step_limit = _mamba_dim("time_step_limit", "time_step_limit", (0.0, float("inf"))) + + in_proj_out = cache[in_proj_key].float() # [batch, seq, proj_size] + conv1d_out = cache[conv1d_key].float() # [batch, conv_dim, seq + kernel - 1] batch_size, seq_len = in_proj_out.shape[0], in_proj_out.shape[1] - # Match HF's SSM numerical precision - in_proj_out_f = in_proj_out.float() - conv1d_out_f = conv1d_out.float() - - # dt is the last num_heads features of in_proj output, post softplus+clamp - dt_raw = in_proj_out_f[..., -num_heads:] - dt_bias = self.dt_bias.float() - dt = torch.nn.functional.softplus(dt_raw + dt_bias) - dt = torch.clamp(dt, time_step_min, time_step_max) # [batch, seq, num_heads] - - # B, C come from the conv1d output after trimming to seq_len and applying SiLU - conv_trimmed = conv1d_out_f[..., :seq_len] - conv_activated = torch.nn.functional.silu(conv_trimmed).transpose(1, 2) - split_sizes = [intermediate_size, n_groups * state_size, n_groups * state_size] - _hidden_x, B_flat, C_flat = conv_activated.split(split_sizes, dim=-1) - B = B_flat.view(batch_size, seq_len, n_groups, state_size) - C = C_flat.view(batch_size, seq_len, n_groups, state_size) - - # GQA-style: each of n_groups B/C pairs is replicated to cover n_heads // n_groups heads + dt = torch.nn.functional.softplus(in_proj_out[..., -num_heads:] + self.dt_bias.float()) + dt = torch.clamp(dt, float(time_step_limit[0]), float(time_step_limit[1])) + + # x, B, C from conv1d output (trim to seq, SiLU, channel-last split) + conv_activated = torch.nn.functional.silu(conv1d_out[..., :seq_len]).transpose(1, 2) + x_flat, B_flat, C_flat = conv_activated.split( + [intermediate_size, n_groups * state_size, n_groups * state_size], dim=-1 + ) + x = x_flat.view(batch_size, seq_len, num_heads, head_dim) + # GQA-style: each of n_groups B/C pairs covers n_heads // n_groups heads heads_per_group = num_heads // n_groups - B_h = B.repeat_interleave(heads_per_group, dim=2) - C_h = C.repeat_interleave(heads_per_group, dim=2) + B = B_flat.view(batch_size, seq_len, n_groups, state_size).repeat_interleave( + heads_per_group, dim=2 + ) + C = C_flat.view(batch_size, seq_len, n_groups, state_size).repeat_interleave( + heads_per_group, dim=2 + ) + # L[i, j] = exp(sum_{k=j+1}^{i} A[h] · dt[k, h]) for i >= j, else 0. + # exp(cumsum[i] - cumsum[j]); cumsum[j] includes dt[j], so the sum is k>j. A = -torch.exp(self.A_log.float()) # [num_heads] - - # L[i, j] = exp(sum_{k=j+1}^{i} A[h] * dt[k, h]) for i >= j, else 0 - # Computed as exp(cumsum[i] - cumsum[j]) since cumsum[j] includes dt[j], - # so the remaining sum runs from k=j+1 to k=i. - log_a = dt * A[None, None, :] - cumsum_log_a = torch.cumsum(log_a, dim=1) - cs = cumsum_log_a.permute(0, 2, 1) # [batch, num_heads, seq] + cs = torch.cumsum(dt * A[None, None, :], dim=1).permute(0, 2, 1) # [batch, num_heads, seq] L_log = cs[:, :, :, None] - cs[:, :, None, :] causal_mask = torch.tril( torch.ones(seq_len, seq_len, dtype=torch.bool, device=L_log.device) ) - L = torch.where( - causal_mask[None, None, :, :], - torch.exp(L_log), - torch.zeros_like(L_log), - ) - - # CB[b, h, i, j] = - CB = torch.einsum("bihs,bjhs->bhij", C_h, B_h) - - M = L * CB # [batch, num_heads, seq, seq] + L = torch.where(causal_mask[None, None], torch.exp(L_log), torch.zeros_like(L_log)) - if include_dt_scaling: - # Multiply column j by dt[j, h] to absorb the B discretization - dt_col = dt.permute(0, 2, 1)[:, :, None, :] - M = M * dt_col - - return M + return _SSDTerms(dt=dt, L=L, x=x, B=B, C=C) diff --git a/transformer_lens/model_bridge/generalized_components/ssm_mixer.py b/transformer_lens/model_bridge/generalized_components/ssm_mixer.py index 58ab3ecfa..f0f3085e9 100644 --- a/transformer_lens/model_bridge/generalized_components/ssm_mixer.py +++ b/transformer_lens/model_bridge/generalized_components/ssm_mixer.py @@ -1,14 +1,28 @@ -"""Wrap-don't-reimplement bridge for HF's MambaMixer (Mamba-1).""" -from typing import Any +"""Wrap-don't-reimplement bridge for HF's MambaMixer (Mamba-1), plus S6 effective attention.""" +from typing import Any, NamedTuple, Optional import torch +from transformer_lens.ActivationCache import ActivationCache +from transformer_lens.hook_points import HookPoint from transformer_lens.model_bridge.generalized_components.base import ( GeneralizedComponent, ) +from transformer_lens.model_bridge.generalized_components.ssm_protocol import ( + SSMStateHookMixin, +) + + +class _S6Terms(NamedTuple): + """S6 intermediates reconstructed from cached Mamba-1 hooks (per-channel).""" + dt: torch.Tensor # [batch, channels, seq] + decay: torch.Tensor # [batch, channels, state, i, j] causal (upper triangle zero) + B: torch.Tensor # [batch, seq, state] (shared across channels) + C: torch.Tensor # [batch, seq, state] (shared across channels) -class SSMMixerBridge(GeneralizedComponent): + +class SSMMixerBridge(SSMStateHookMixin, GeneralizedComponent): """Opaque wrapper around Mamba-1's MambaMixer. Submodules (in_proj, conv1d, x_proj, dt_proj, out_proj) are swapped into @@ -26,24 +40,38 @@ class SSMMixerBridge(GeneralizedComponent): "hook_x_proj": "x_proj.hook_out", "hook_dt_proj": "dt_proj.hook_out", "hook_ssm_out": "hook_out", + # Canonical SSM vocabulary (additive): Mamba-1 exposes the discrete time + # step via dt_proj. B/C are bundled in x_proj (not separately hooked). + "hook_ssm_dt": "dt_proj.hook_out", } + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) # mixin adds hook_ssm_state + eager_scan + # Real per-step write term dt·x·B, per-channel [batch, channels, seq, state]. + self.hook_ssm_write = HookPoint() + def forward(self, *args: Any, **kwargs: Any) -> Any: - """Hook the input, delegate to HF slow_forward, hook the output.""" + """Hook the input, run HF slow_forward (or the eager scan), hook the output.""" if self.original_component is None: raise RuntimeError( f"Original component not set for {self.name}. " "Call set_original_component() first." ) - # Hook the hidden_states input (positional or keyword) + hidden_states: Optional[torch.Tensor] = None if len(args) > 0 and isinstance(args[0], torch.Tensor): - hooked = self.hook_in(args[0]) - args = (hooked,) + args[1:] + hidden_states = self.hook_in(args[0]) + args = (hidden_states,) + args[1:] elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor): - kwargs["hidden_states"] = self.hook_in(kwargs["hidden_states"]) + hidden_states = self.hook_in(kwargs["hidden_states"]) + kwargs["hidden_states"] = hidden_states - output = self.original_component(*args, **kwargs) + # Eager-scan slow path: opt-in, prefill only (no cache_params). Keyed on + # explicit state, never on hook-registry introspection. + if self.eager_scan and hidden_states is not None and kwargs.get("cache_params") is None: + output: Any = self._eager_scan_forward(hidden_states, kwargs.get("attention_mask")) + else: + output = self.original_component(*args, **kwargs) # Hook the primary output tensor, preserving tuple structure if isinstance(output, tuple) and len(output) > 0: @@ -54,3 +82,245 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: if isinstance(output, torch.Tensor): return self.hook_out(output) return output + + def _eager_scan_forward( + self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] + ) -> torch.Tensor: + """Reimplement the Mamba-1 mixer prefill forward with an eager Python S6 scan. + + Reuses HF's in_proj / conv1d / x_proj / dt_proj / out_proj submodules (so + their hooks still fire) and reimplements ONLY the S6 recurrence:: + + write_t[c, s] = dt_t[c] · x_t[c] · B_t[s] -> hook_ssm_write [b, c, seq, s] + S_t[c, s] = exp(A[c,s]·dt_t[c]) · S_{t-1} + write_t + S = stack_t S_t -> hook_ssm_state [b, c, seq, s] + y_t[c] = sum_s C_t[s] · S_t[c, s] + D[c] · x_t[c] + + Intervening on hook_ssm_write re-runs the recurrence (propagates to later + states); hook_ssm_state is the post-scan trajectory, so a patch there + changes only the same-position output y_t = C_t·S_t — patch hook_ssm_write + for a propagating state edit. + + Kernel-divergence caveat: reproduces HF's scan only to fp tolerance + (≈1e-6 fp32), never bit-for-bit. O(seq) Python; materializes an + O(b·channels·seq·state) write/state tensor. Prefill only. + """ + oc: Any = self.original_component + batch, seq_len, _ = hidden_states.shape + d_inner = oc.intermediate_size + state = oc.ssm_state_size + dt_rank = oc.time_step_rank + + def _mask_cf(states: torch.Tensor) -> torch.Tensor: + # HF MambaMixer masks whenever attention_mask is present (no batch guard, + # unlike Mamba-2), so match that — batch-1 padded inputs included. + if attention_mask is not None: + return states * attention_mask.unsqueeze(1).to(states.dtype) + return states + + # 1-2. in_proj + gate split, conv — reuse the HF submodule bridges (hooks fire). + # HF masks padding on the channel-first x both before and after the conv. + projected = oc.in_proj(hidden_states).transpose(1, 2) # [b, 2*d_inner, seq] + x_raw, gate = projected.chunk(2, dim=1) # [b, d_inner, seq] + x = oc.act(oc.conv1d(_mask_cf(x_raw))[..., :seq_len]) # [b, d_inner, seq] + x = _mask_cf(x) + + # 3. x_proj -> dt low-rank / B / C ; dt_proj -> dt (softplus, channel-first). + ssm_params = oc.x_proj(x.transpose(1, 2)) # [b, seq, dt_rank + 2*state] + time_step, B, C = ssm_params.split([dt_rank, state, state], dim=-1) # B, C: [b, seq, state] + dt = torch.nn.functional.softplus(oc.dt_proj(time_step)).transpose(1, 2).float() + A = -torch.exp(self.A_log.float()) # [d_inner, state] + x_f, B_f, C_f = x.float(), B.float(), C.float() + + # 4. Eager recurrence with intervention hooks. + writes = (dt * x_f)[:, :, :, None] * B_f[:, None, :, :] # [b, d_inner, seq, state] + writes = self.hook_ssm_write(writes) + + ssm_state = torch.zeros(batch, d_inner, state, dtype=writes.dtype, device=writes.device) + states = [] + for t in range(seq_len): + decay = torch.exp(A[None] * dt[:, :, t, None]) # [batch, d_inner, state] + ssm_state = decay * ssm_state + writes[:, :, t] + states.append(ssm_state) + state_traj = torch.stack(states, dim=2) # [batch, d_inner, seq, state] + state_traj = self.hook_ssm_state(state_traj) + + y = torch.einsum("bcts,bts->bct", state_traj, C_f) # [batch, d_inner, seq] + y = y + self.D.float()[None, :, None] * x_f + scan_output = y * oc.act(gate.float()) # HF gates with self.act, not hardcoded silu + + # 5. Output projection — reuse the HF submodule bridge. + contextualized: torch.Tensor = oc.out_proj( + scan_output.transpose(1, 2).to(hidden_states.dtype) + ) + return contextualized + + def compute_effective_attention( + self, + cache: ActivationCache, + layer_idx: int, + include_dt_scaling: bool = False, + per_state_coord: bool = False, + ) -> torch.Tensor: + """Materialize Mamba-1's per-channel effective attention from cached hooks. + + Mamba-1's S6 selective scan is equivalent to causal attention with a + per-channel, per-state-coordinate learned decay — see "The Hidden + Attention of Mamba" (Ali et al., ACL 2025). Unlike Mamba-2 there is no + head grouping: each of the ``intermediate_size`` channels has its own + ``A`` row, so the "head" axis here is the channel axis:: + + M[c, i, j] = sum_n C[i, n] · prod_{k=j+1..i} exp(A[c,n]·dt[c,k]) · B[j, n] + + Reads B/C from ``x_proj.hook_out`` and dt from ``dt_proj.hook_out`` + (dt = softplus of that output); A from the module via ``__getattr__``. + Read-only: no ``forward()`` re-run. + + Args: + cache: ActivationCache from ``run_with_cache`` with this layer's + x_proj and dt_proj hooks. + layer_idx: Block index for this mixer. + include_dt_scaling: False (default) returns the attention-like form; + True multiplies column j by dt[c, j], giving the reconstruction + form satisfying ``y[c,i] = sum_j M[c,i,j]·x[c,j] + D[c]·x[c,i]`` + (x is the post-conv SiLU input; y the pre-gate scan output). + per_state_coord: False (default) sums over the state coordinate and + returns ``[batch, channels, seq, seq]``. True returns the + unsummed ``[batch, channels, state, seq, seq]`` tensor (the + paper's D·N matrices) — OFF by default. + + Returns: + ``[batch, intermediate_size, seq, seq]``, or + ``[batch, intermediate_size, state_size, seq, seq]`` when + ``per_state_coord`` is True. + + Peak memory is O(batch · intermediate_size · state_size · seq²) — the + per-(channel, state) decay tensor — even for the summed default; use on + short sequences. + """ + t = self._s6_terms(cache, layer_idx) + + # M_coord[c, s, i, j] = C[i, s] · decay[c, s, i, j] · B[j, s] + C_col = t.C.permute(0, 2, 1)[:, None, :, :, None] # [batch, 1, state, i, 1] + B_col = t.B.permute(0, 2, 1)[:, None, :, None, :] # [batch, 1, state, 1, j] + M_coord = C_col * t.decay * B_col # [batch, channels, state, i, j] + + if include_dt_scaling: + M_coord = M_coord * t.dt[:, :, None, None, :] # × dt[c, j] + + if per_state_coord: + return M_coord + return M_coord.sum(dim=2) # [batch, channels, seq, seq] + + def compute_ssm_state( + self, + cache: ActivationCache, + layer_idx: int, + time_step: Optional[int] = None, + ) -> torch.Tensor: + """Reconstruct Mamba-1's recurrent state ``S`` from cached hook values. + + S6 recurrence ``h_t[c,s] = exp(A[c,s]·dt_t[c])·h_{t-1}[c,s] + dt_t[c]·x_t[c]·B_t[s]`` + unrolls to:: + + S_t[c, s] = sum_{j<=t} decay[c, s, t, j] · dt_j[c] · x_j[c] · B_j[s] + + with the same per-(channel, state) ``decay`` used by + ``compute_effective_attention``. ``x`` is the post-conv SiLU input + (``SiLU(conv1d.hook_out)``). Read-only: no ``forward()`` re-run. Verify + with ``y_t[c] = sum_s C_t[s]·S_t[c,s] + D[c]·x_t[c]``. + + On padded batches the cached hooks are unmasked, so ``S`` is exact only at + non-pad positions; pad-position state is out of contract. + + Args: + cache: ActivationCache from ``run_with_cache`` with this layer's + x_proj, dt_proj, and conv1d hooks. + layer_idx: Block index for this mixer. + time_step: If given, return only ``S`` at that position + (``[batch, channels, state]``); None returns every step. + + Returns: + ``[batch, channels, seq, state]`` for all steps, or + ``[batch, channels, state]`` for a single ``time_step``. + + Peak memory is O(batch · channels · state · seq²) — the per-(channel, + state) decay tensor, built in full by ``_s6_terms`` regardless of + ``time_step`` (which bounds only the returned tensor, not this peak); use + on short sequences. + """ + conv_key = f"blocks.{layer_idx}.mixer.conv1d.hook_out" + if conv_key not in cache: + raise RuntimeError( + f"compute_ssm_state needs {conv_key!r} in cache. Run " + "`run_with_cache()` on the bridge before calling this method." + ) + t = self._s6_terms(cache, layer_idx) + seq_len = t.dt.shape[-1] + oc: Any = self.original_component + # x = act(conv output), the SSM scan input; channel-first [batch, channels, seq]. + x = oc.act(cache[conv_key].float()[..., :seq_len]) + dtx = t.dt * x # dt_j[c]·x_j[c], [batch, channels, seq] + # S_t[c, s] = sum_{j<=t} decay[c, s, t, j] · dtx[c, j] · B_j[s] + if time_step is not None: + return torch.einsum("bcsj,bcj,bjs->bcs", t.decay[:, :, :, time_step, :], dtx, t.B) + return torch.einsum("bcsij,bcj,bjs->bcis", t.decay, dtx, t.B) + + def _s6_terms(self, cache: ActivationCache, layer_idx: int) -> _S6Terms: + """Reconstruct the shared S6 intermediates (dt, per-(channel,state) decay, B, C). + + Reused by ``compute_effective_attention`` and ``compute_ssm_state``. Reads + B/C from ``x_proj.hook_out`` (shared across channels), dt from + ``dt_proj.hook_out`` (softplus), and A via ``__getattr__``. Dims come from + the wrapped HF mixer, cfg is the fallback (mirrors the Mamba-2 bridge). + """ + if self.config is None: + raise RuntimeError("SSMMixerBridge.config must be set") + + x_proj_key = f"blocks.{layer_idx}.mixer.x_proj.hook_out" + dt_proj_key = f"blocks.{layer_idx}.mixer.dt_proj.hook_out" + for key in (x_proj_key, dt_proj_key): + if key not in cache: + raise RuntimeError( + f"S6 reconstruction needs {key!r} in cache. Run " + "`run_with_cache()` on the bridge before calling this." + ) + + cfg = self.config + oc = self.original_component + + def _mamba_dim(module_attr: str, cfg_attr: str, default: Any) -> Any: + if oc is not None: + val = getattr(oc, module_attr, None) + if val is not None: + return val + return getattr(cfg, cfg_attr, default) + + state_size = int(_mamba_dim("ssm_state_size", "state_size", 16)) + dt_rank = int(_mamba_dim("time_step_rank", "time_step_rank", 0)) + + x_proj_out = cache[x_proj_key].float() # [batch, seq, dt_rank + 2*state] + dt_proj_out = cache[dt_proj_key].float() # [batch, seq, channels] + seq_len = dt_proj_out.shape[1] + + # B, C from x_proj output (shared across channels); first dt_rank cols are dt low-rank. + _time_step, B, C = x_proj_out.split([dt_rank, state_size, state_size], dim=-1) + + dt = torch.nn.functional.softplus(dt_proj_out).transpose(1, 2) # [batch, channels, seq] + A = -torch.exp(self.A_log.float()) # [channels, state] + + # decay[c, s, i, j] = exp(A[c,s]·(cumdt[c,i]-cumdt[c,j])) for i >= j, else 0. + cumdt = torch.cumsum(dt, dim=-1) # [batch, channels, seq] + dcs = cumdt[:, :, :, None] - cumdt[:, :, None, :] # [batch, channels, i, j] + decay_exp = ( + A[None, :, :, None, None] * dcs[:, :, None, :, :] + ) # [batch, channels, state, i, j] + causal_mask = torch.tril( + torch.ones(seq_len, seq_len, dtype=torch.bool, device=decay_exp.device) + ) + decay = torch.where( + causal_mask[None, None, None, :, :], + torch.exp(decay_exp), + torch.zeros_like(decay_exp), + ) + return _S6Terms(dt=dt, decay=decay, B=B, C=C) diff --git a/transformer_lens/model_bridge/generalized_components/ssm_protocol.py b/transformer_lens/model_bridge/generalized_components/ssm_protocol.py new file mode 100644 index 000000000..75c346e10 --- /dev/null +++ b/transformer_lens/model_bridge/generalized_components/ssm_protocol.py @@ -0,0 +1,99 @@ +"""Family-agnostic discovery + shared state-mutation surface for SSM/recurrent mixers. + +Defines the structural contract (``SSMMixerProtocol``) that ``SSM2MixerBridge`` +(Mamba-2), ``SSMMixerBridge`` (Mamba-1), and ``GatedDeltaNetBridge`` all satisfy, +plus a lookup that finds a block's SSM mixer regardless of which variant slot it +occupies (``.mixer`` for Mamba, ``.linear_attn`` for gated-delta-net). A Protocol +(not a base class) keeps each family's own hook set — only the discovery surface +is shared. ``SSMStateHookMixin`` is the one place the canonical eager-scan +intervention hooks (``hook_ssm_state`` + the ``eager_scan`` opt-in) are defined, so +the three families share a single definition rather than repeating it ad hoc. +""" +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Optional, Protocol, runtime_checkable + +import torch +import torch.nn as nn + +from transformer_lens.hook_points import HookPoint +from transformer_lens.model_bridge.generalized_components.block import ( + VARIANT_SUBMODULE_NAMES, +) + +if TYPE_CHECKING: + from transformer_lens.ActivationCache import ActivationCache + + +@runtime_checkable +class SSMMixerProtocol(Protocol): + """An SSM/recurrent mixer that can materialize an effective-attention matrix. + + Only ``compute_effective_attention(cache, layer_idx)`` is required; families + add their own optional keyword options (e.g. ``include_dt_scaling``, + ``per_state_coord``) which callers discover by signature introspection. + """ + + def compute_effective_attention(self, cache: "ActivationCache", layer_idx: int) -> torch.Tensor: + ... + + +class SSMStateHookMixin: + """Canonical state-mutation hooks shared by the SSM/recurrent mixers. + + One definition point so Mamba-1/Mamba-2/gated-delta-net expose the same names: + + - ``hook_ssm_state`` — post-scan state trajectory ``S_t`` (created here for all); + patching it changes only the same-position readout, not the recurrence. + - ``hook_ssm_write`` — per-step write influence: a real HookPoint for the input- + linear families (Mamba-1/2 ``dt·(x⊗B)``, added by the subclass), an alias onto + ``hook_beta`` for the state-dependent delta rule; patching it re-runs the scan. + + ``eager_scan`` (opt-in, prefill only) swaps HF's fused kernel for a Python scan so + these hooks fire; default False leaves the cached path bit-for-bit untouched. Must + precede ``GeneralizedComponent`` in the bases so ``super().__init__`` reaches it. + """ + + eager_scan: bool = False + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.hook_ssm_state = HookPoint() + + +# Children present on *every* SSM2MixerBridge regardless of whether it wraps a +# real Mamba layer — universal I/O hooks, the wrapped module, and the eager-scan +# analysis hooks. None of these signal that the mixer is realized. +_PASSTHROUGH_CHILDREN = frozenset( + {"hook_in", "hook_out", "_original_component", "hook_ssm_write", "hook_ssm_state"} +) + + +def _is_realized_ssm_mixer(mixer: object) -> bool: + """True unless ``mixer`` is a no-op passthrough wrapper. + + A hybrid like NemotronH wires a single ``SSM2MixerBridge`` ``.mixer`` slot on + *every* layer; on attention / MLP / MoE layers its optional projection + submodules are skipped, leaving only the universal hooks. A realized mixer + always has something more — projection submodules (Mamba-1/2) or interior + hooks like ``hook_q`` / ``hook_log_decay`` (gated-delta-net). This structural + check needs no ``cfg.layers_block_type``. + """ + return any(name not in _PASSTHROUGH_CHILDREN for name in getattr(mixer, "_modules", {})) + + +def find_ssm_mixer(block: nn.Module) -> Optional[SSMMixerProtocol]: + """Return the block's SSM mixer submodule, or None if it has none. + + Scans the layer-type variant slots (``.mixer`` / ``.linear_attn`` / ``.mamba`` + / ``.ssm``) and returns the first that conforms to ``SSMMixerProtocol`` and is + a realized SSM mixer (not a passthrough wrapper). An attention layer (only + ``.attn``), or a hybrid's passthrough ``.mixer`` on a non-SSM layer, returns + None. + """ + modules = getattr(block, "_modules", {}) + for name in VARIANT_SUBMODULE_NAMES: + sub = modules.get(name) + if isinstance(sub, SSMMixerProtocol) and _is_realized_ssm_mixer(sub): + return sub + return None diff --git a/transformer_lens/model_bridge/supported_architectures/granite_moe_hybrid.py b/transformer_lens/model_bridge/supported_architectures/granite_moe_hybrid.py index 53229252e..636513182 100644 --- a/transformer_lens/model_bridge/supported_architectures/granite_moe_hybrid.py +++ b/transformer_lens/model_bridge/supported_architectures/granite_moe_hybrid.py @@ -5,7 +5,10 @@ has a shared MLP and optional sparse MoE. Both attention and Mamba are mapped as optional — each present only on its -respective layer type. Mamba hooks expose in_proj, conv1d, and inner_norm. +respective layer type. The Mamba mixer is wired under the canonical ``.mixer`` +slot (HF path is ``.mamba``) so SSM analyses (compute_effective_attention) +reach it the same way as on NemotronH / Mamba-2. Mamba hooks expose in_proj, +conv1d, and inner_norm (the gated two-input MambaRMSNormGated). """ from typing import Any @@ -14,6 +17,7 @@ from transformer_lens.model_bridge.generalized_components import ( BlockBridge, EmbeddingBridge, + GatedRMSNormBridge, LinearBridge, MLPBridge, MoEBridge, @@ -37,6 +41,9 @@ class GraniteMoeHybridArchitectureAdapter(GraniteArchitectureAdapter): universal. Inherits Granite config and attention bridge construction. """ + # Explicit for parity with the Mamba/NemotronH siblings. + applicable_phases: list[int] = [1, 2, 3, 4] + def __init__(self, cfg: Any) -> None: ArchitectureAdapter.__init__(self, cfg) self._setup_common_config(cfg) @@ -47,10 +54,23 @@ def __init__(self, cfg: Any) -> None: self.supports_fold_ln = False self.weight_processing_conversions = {} + + # Normalize the per-layer mixer-type list as cfg.layers_block_type (HF names + # it `layer_types`) so analysis tools can find the Mamba layers, as on NemotronH. + layers_block_type = ( + getattr(cfg, "layers_block_type", None) or getattr(cfg, "layer_types", None) or [] + ) + setattr(self.cfg, "layers_block_type", list(layers_block_type)) + self.component_mapping = self._build_component_mapping() def _build_mamba_bridge(self) -> SSM2MixerBridge: - """Mamba-2 mixer bridge with in_proj, conv1d, inner_norm hooks.""" + """Mamba-2 mixer bridge with in_proj, conv1d, inner_norm hooks. + + ``name="mamba"`` is the HF submodule path; the bridge exposes it under the + canonical ``.mixer`` dict key (see ``_build_component_mapping``). inner_norm + is the gated two-input MambaRMSNormGated, so it wraps GatedRMSNormBridge. + """ return SSM2MixerBridge( name="mamba", config=self.cfg, @@ -58,7 +78,7 @@ def _build_mamba_bridge(self) -> SSM2MixerBridge: submodules={ "in_proj": LinearBridge(name="in_proj"), "conv1d": DepthwiseConv1DBridge(name="conv1d"), - "inner_norm": LinearBridge(name="norm"), + "inner_norm": GatedRMSNormBridge(name="norm"), }, ) @@ -67,7 +87,7 @@ def _build_component_mapping(self) -> dict: "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg), "attn": self._build_attention_bridge(optional=True), - "mamba": self._build_mamba_bridge(), + "mixer": self._build_mamba_bridge(), "shared_mlp": MLPBridge( name="shared_mlp", config=self.cfg, diff --git a/transformer_lens/model_bridge/supported_architectures/mamba.py b/transformer_lens/model_bridge/supported_architectures/mamba.py index c45a936c3..06e3a2c05 100644 --- a/transformer_lens/model_bridge/supported_architectures/mamba.py +++ b/transformer_lens/model_bridge/supported_architectures/mamba.py @@ -23,11 +23,9 @@ class MambaArchitectureAdapter(ArchitectureAdapter): ``_HF_PASSTHROUGH_ATTRS`` in sources/transformers.py. """ - # Phases 1-3 are transformer-shaped (component/weight comparison) and don't - # fit SSMs; component-level coverage lives in integration tests: - # tests/integration/model_bridge/test_mamba_adapter.py. Phase 4 (generation - # + text-quality) needs no component comparison, so it applies. - applicable_phases: list[int] = [4] + # White-box forward: P1 is exact vs raw HF (mixer delegates to HF); P2/P3 skip + # without a HookedTransformer; P4 is generation. + applicable_phases: list[int] = [1, 2, 3, 4] def __init__(self, cfg: Any) -> None: super().__init__(cfg) diff --git a/transformer_lens/model_bridge/supported_architectures/mamba2.py b/transformer_lens/model_bridge/supported_architectures/mamba2.py index d09feb380..f9d716852 100644 --- a/transformer_lens/model_bridge/supported_architectures/mamba2.py +++ b/transformer_lens/model_bridge/supported_architectures/mamba2.py @@ -1,5 +1,6 @@ """Architecture adapter for HF's Mamba2ForCausalLM, plus the effective attention helper.""" -from typing import Any, Optional +import warnings +from typing import Any, Dict, Optional, Union import torch @@ -28,11 +29,9 @@ class Mamba2ArchitectureAdapter(ArchitectureAdapter): loop with Mamba-1. """ - # Phases 1-3 are transformer-shaped (component/weight comparison) and don't - # fit SSMs; component-level coverage lives in integration tests: - # tests/integration/model_bridge/test_mamba2_adapter.py. Phase 4 (generation - # + text-quality) needs no component comparison, so it applies. - applicable_phases: list[int] = [4] + # White-box forward: P1 is exact vs raw HF (mixer delegates to HF); P2/P3 skip + # without a HookedTransformer; P4 is generation. + applicable_phases: list[int] = [1, 2, 3, 4] def __init__(self, cfg: Any) -> None: super().__init__(cfg) @@ -115,60 +114,39 @@ def compute_effective_attention( cache: ActivationCache, layer: Optional[int] = None, include_dt_scaling: bool = False, -) -> torch.Tensor: - """Compute Mamba-2 effective attention M = L ⊙ (C B^T) for one or all layers. +) -> Union[torch.Tensor, Dict[int, torch.Tensor]]: + """Mamba-2 effective attention for one or all layers. - Wraps ``SSM2MixerBridge.compute_effective_attention`` so callers don't have - to repeat the layer index, and adds all-layers stacking when ``layer`` is - None. - - Args: - bridge: A loaded Mamba-2 ``TransformerBridge``. - cache: ActivationCache from ``run_with_cache`` with in_proj and conv1d - hooks populated for every requested layer. - layer: Specific block index, or None for all layers stacked. - include_dt_scaling: See ``SSM2MixerBridge.compute_effective_attention``. - - Returns: - Shape ``[batch, num_heads, seq, seq]`` for a single layer, or - ``[n_layers, batch, num_heads, seq, seq]`` when layer is None. - - Raises: - TypeError: If any targeted block's mixer isn't an ``SSM2MixerBridge``. + .. deprecated:: + Use the family-agnostic ``cache.compute_ssm_effective_attention(layer=...)`` + instead. This thin wrapper delegates to it and ignores ``bridge`` (the + cache already knows its model). + """ + warnings.warn( + "mamba2.compute_effective_attention is deprecated; use " + "cache.compute_ssm_effective_attention(layer=..., include_dt_scaling=...).", + DeprecationWarning, + stacklevel=2, + ) + return cache.compute_ssm_effective_attention(layer=layer, include_dt_scaling=include_dt_scaling) - Example:: - from transformer_lens.model_bridge.supported_architectures.mamba2 import ( - compute_effective_attention, - ) +def compute_ssm_state( + bridge: TransformerBridge, + cache: ActivationCache, + layer: Optional[int] = None, + time_step: Optional[int] = None, +) -> Union[torch.Tensor, Dict[int, torch.Tensor]]: + """Reconstruct the recurrent SSM state ``S`` for one or all Mamba-2 layers. - M5 = compute_effective_attention(bridge, cache, layer=5) - M_all = compute_effective_attention(bridge, cache) + .. deprecated:: + Use ``cache.compute_ssm_state(layer=..., time_step=...)`` instead. This + thin wrapper delegates to it and ignores ``bridge``. """ - if layer is not None: - mixer = bridge.blocks[layer].mixer - if not isinstance(mixer, SSM2MixerBridge): - raise TypeError( - f"Layer {layer} mixer is {type(mixer).__name__}, not " - "SSM2MixerBridge. compute_effective_attention requires a " - "Mamba-2 bridge." - ) - return mixer.compute_effective_attention( - cache, layer_idx=layer, include_dt_scaling=include_dt_scaling - ) - - matrices = [] - for layer_idx, block in enumerate(bridge.blocks): - mixer = block.mixer - if not isinstance(mixer, SSM2MixerBridge): - raise TypeError( - f"Layer {layer_idx} mixer is {type(mixer).__name__}, not " - "SSM2MixerBridge. compute_effective_attention requires a " - "Mamba-2 bridge." - ) - matrices.append( - mixer.compute_effective_attention( - cache, layer_idx=layer_idx, include_dt_scaling=include_dt_scaling - ) - ) - return torch.stack(matrices, dim=0) + warnings.warn( + "mamba2.compute_ssm_state is deprecated; use " + "cache.compute_ssm_state(layer=..., time_step=...).", + DeprecationWarning, + stacklevel=2, + ) + return cache.compute_ssm_state(layer=layer, time_step=time_step) diff --git a/transformer_lens/model_bridge/supported_architectures/nemotron_h.py b/transformer_lens/model_bridge/supported_architectures/nemotron_h.py index 2f322fa8d..216b334b4 100644 --- a/transformer_lens/model_bridge/supported_architectures/nemotron_h.py +++ b/transformer_lens/model_bridge/supported_architectures/nemotron_h.py @@ -1,6 +1,6 @@ """Nemotron-H hybrid Mamba2-Transformer architecture adapter. -Supports NemotronHForCausalLM (nvidia/Nemotron-H-8B-Base, Nemotron-H-47B-A13B). +Supports NemotronHForCausalLM (e.g. nvidia/NVIDIA-Nemotron-Nano-9B-v2, Nemotron-3 series). Architecture overview: - Heterogeneous layers defined by ``config.layers_block_type`` — each element is @@ -26,9 +26,8 @@ Mamba-specific inner submodules (in_proj, conv1d, inner_norm, out_proj) are declared ``optional=True`` so setup skips them gracefully on non-Mamba layers. - MLP layers use ``relu2`` activation (not SwiGLU); ``gated_mlp = False``. -- ``applicable_phases = []``: ``verify_models`` is transformer-shaped and would - require a dedicated refactor to cover SSM hybrids. Coverage lives in the - integration test instead. +- ``applicable_phases = [1, 2, 3, 4]``: P1 is exact vs raw HF (passthrough mixers); + P2/P3 skip without a HookedTransformer; P4 is generation. """ from typing import Any @@ -69,9 +68,9 @@ class NemotronHArchitectureAdapter(ArchitectureAdapter): is determined by ``config.layers_block_type[layer_idx]``. """ - # verify_models is transformer-shaped and requires a dedicated refactor to - # cover SSM hybrids. Integration tests cover forward-pass correctness instead. - applicable_phases: list[int] = [] + # White-box forward: P1 is exact vs raw HF (passthrough mixers); P2/P3 skip + # without a HookedTransformer; P4 is generation. + applicable_phases: list[int] = [1, 2, 3, 4] def __init__(self, cfg: Any) -> None: super().__init__(cfg) @@ -89,10 +88,12 @@ def __init__(self, cfg: Any) -> None: # Mamba layers require per-step SSM state; generation is stateful. self.cfg.is_stateful = True - # Expose the heterogeneous layer-type list so tests and analysis tools - # can inspect which layers are which without loading a full HF model. - layers_block_type = getattr(cfg, "layers_block_type", []) - setattr(self.cfg, "layers_block_type", layers_block_type) + # Normalize the per-layer type list as cfg.layers_block_type (HF names it + # `layer_types`) so analysis tools can find the Mamba layers, as on Granite. + layers_block_type = ( + getattr(cfg, "layers_block_type", None) or getattr(cfg, "layer_types", None) or [] + ) + setattr(self.cfg, "layers_block_type", list(layers_block_type)) # Mamba-2 dimensional config (mirrors Mamba2ArchitectureAdapter). mamba_num_heads = getattr(cfg, "mamba_num_heads", 128) diff --git a/transformer_lens/tools/model_registry/AGENTS.md b/transformer_lens/tools/model_registry/AGENTS.md index e3ac11134..e0bd42fa9 100644 --- a/transformer_lens/tools/model_registry/AGENTS.md +++ b/transformer_lens/tools/model_registry/AGENTS.md @@ -116,6 +116,8 @@ Never edit manually. | 7 | Multimodal (vision/text alignment) — only Llava / Gemma3-multimodal | | 8 | Audio — only Hubert | +SSM / recurrent families and the hybrids (Mamba-1/2, gated-delta-net, NemotronH, GraniteMoeHybrid, Qwen3.5/Qwen3-Next) declare `applicable_phases = [1, 2, 3, 4]` — all four apply. P2/P3 run but skip their HookedTransformer-comparison sub-tests (SSMs have no HT), which is scored as a pass. + ### Phase-score thresholds `verify_models` enforces hard pass/fail at the thresholds in `_MIN_PHASE_SCORES` ([`verify_models.py:508`](verify_models.py)). Below threshold OR a required-test failure → `STATUS_FAILED`. The contract: diff --git a/transformer_lens/tools/model_registry/data/architecture_gaps.json b/transformer_lens/tools/model_registry/data/architecture_gaps.json index 78c2fe754..0349cda8d 100644 --- a/transformer_lens/tools/model_registry/data/architecture_gaps.json +++ b/transformer_lens/tools/model_registry/data/architecture_gaps.json @@ -1,33 +1,14 @@ { - "generated_at": "2026-06-25", + "generated_at": "2026-07-01", "scan_info": { - "total_scanned": 5949, + "total_scanned": 5855, "task_filter": "text-generation", "min_downloads": 500, - "scan_duration_seconds": 8.1 + "scan_duration_seconds": 7.0 }, - "total_unsupported_architectures": 507, - "total_unsupported_models": 3862, + "total_unsupported_architectures": 514, + "total_unsupported_models": 4737, "gaps": [ - { - "architecture_id": "NemotronHForCausalLM", - "total_models": 153, - "total_downloads": 16859572, - "min_param_count": 4226684, - "sample_models": [ - "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16", - "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", - "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16", - "nvidia/NVIDIA-Nemotron-Nano-9B-v2", - "nvidia/NVIDIA-Nemotron-Nano-9B-v2-Japanese", - "trl-internal-testing/tiny-NemotronHForCausalLM-nano", - "nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16", - "trl-internal-testing/tiny-NemotronHForCausalLM-ultra", - "trl-internal-testing/tiny-NemotronHForCausalLM-super", - "nvidia/Nemotron-Cascade-2-30B-A3B" - ], - "relevancy_score": 99.7 - }, { "architecture_id": "MarianMTModel", "total_models": 154, @@ -47,10 +28,29 @@ ], "relevancy_score": 98.6 }, + { + "architecture_id": "Qwen3_5MoeForConditionalGeneration", + "total_models": 144, + "total_downloads": 1125570, + "min_param_count": 664944, + "sample_models": [ + "cyankiwi/Nex-N2-mini-AWQ-INT4", + "huihui-ai/Huihui-Qwen3.6-35B-A3B-Claude-4.7-Opus-abliterated", + "genevera/Qwen3.6-35B-A3B-Abliterated-Heretic-AWQ-4bit", + "nex-agi/Nex-N2-mini", + "Momix-44/Huihui-Qwen3.6-35B-A3B-Claude-4.7-Opus-abliterated", + "nex-agi/Nex-N2-Pro", + "atbender/Qwen3.6-VL-REAP-26B-A3B-W4A16", + "happypatrick/Qwen3.5-397B-A17B-heretic-int4-AutoRound", + "bullerwins/Nex-N2-Pro-4bit-W4A16", + "YuYu1015/Huihui-Qwen3.6-35B-A3B-abliterated-int4-AutoRound" + ], + "relevancy_score": 91.4 + }, { "architecture_id": "Lfm2ForCausalLM", - "total_models": 95, - "total_downloads": 1537174, + "total_models": 133, + "total_downloads": 2121790, "min_param_count": 4321856, "sample_models": [ "LiquidAI/LFM2.5-1.2B-Instruct", @@ -64,31 +64,31 @@ "LiquidAI/LFM2.5-1.2B-JP-202606", "LiquidAI/LFM2-2.6B" ], - "relevancy_score": 77.7 + "relevancy_score": 89.5 }, { - "architecture_id": "Qwen3_5MoeForConditionalGeneration", - "total_models": 100, - "total_downloads": 622389, - "min_param_count": 1034089088, + "architecture_id": "DFlashDraftModel", + "total_models": 88, + "total_downloads": 1674736, + "min_param_count": 385906176, "sample_models": [ - "cyankiwi/Nex-N2-mini-AWQ-INT4", - "huihui-ai/Huihui-Qwen3.6-35B-A3B-Claude-4.7-Opus-abliterated", - "genevera/Qwen3.6-35B-A3B-Abliterated-Heretic-AWQ-4bit", - "nex-agi/Nex-N2-mini", - "Momix-44/Huihui-Qwen3.6-35B-A3B-Claude-4.7-Opus-abliterated", - "nex-agi/Nex-N2-Pro", - "atbender/Qwen3.6-VL-REAP-26B-A3B-W4A16", - "happypatrick/Qwen3.5-397B-A17B-heretic-int4-AutoRound", - "bullerwins/Nex-N2-Pro-4bit-W4A16", - "YuYu1015/Huihui-Qwen3.6-35B-A3B-abliterated-int4-AutoRound" + "z-lab/Qwen3-8B-DFlash-b16", + "z-lab/Qwen3.6-35B-A3B-DFlash", + "z-lab/Qwen3.6-27B-DFlash", + "z-lab/Qwen3-4B-DFlash-b16", + "z-lab/LLaMA3.1-8B-Instruct-DFlash-UltraChat", + "z-lab/Qwen3.5-9B-DFlash", + "z-lab/gemma-4-31B-it-DFlash", + "z-lab/Qwen3.5-4B-DFlash", + "z-lab/gpt-oss-20b-DFlash", + "z-lab/Qwen3.5-35B-A3B-DFlash" ], - "relevancy_score": 73.3 + "relevancy_score": 75.9 }, { "architecture_id": "DeepseekV4ForCausalLM", - "total_models": 52, - "total_downloads": 13485551, + "total_models": 70, + "total_downloads": 16637388, "min_param_count": 110369621, "sample_models": [ "deepseek-ai/DeepSeek-V4-Pro", @@ -102,31 +102,12 @@ "Intel/DeepSeek-V4-Pro-W4A16-AutoRound", "JANGQ-AI/DeepSeek-V4-Flash-JANGTQ2" ], - "relevancy_score": 69.7 - }, - { - "architecture_id": "DFlashDraftModel", - "total_models": 64, - "total_downloads": 1200690, - "min_param_count": 385906176, - "sample_models": [ - "z-lab/Qwen3-8B-DFlash-b16", - "z-lab/Qwen3.6-35B-A3B-DFlash", - "z-lab/Qwen3.6-27B-DFlash", - "z-lab/Qwen3-4B-DFlash-b16", - "z-lab/LLaMA3.1-8B-Instruct-DFlash-UltraChat", - "z-lab/Qwen3.5-9B-DFlash", - "z-lab/gemma-4-31B-it-DFlash", - "z-lab/Qwen3.5-4B-DFlash", - "z-lab/gpt-oss-20b-DFlash", - "z-lab/Qwen3.5-35B-A3B-DFlash" - ], - "relevancy_score": 68.1 + "relevancy_score": 75.5 }, { "architecture_id": "MiniCPMForCausalLM", - "total_models": 57, - "total_downloads": 646259, + "total_models": 71, + "total_downloads": 927705, "min_param_count": 433873920, "sample_models": [ "openbmb/MiniCPM-2B-sft-bf16", @@ -140,12 +121,50 @@ "openbmb/BitCPM-CANN-0.5B-unquantized", "openbmb/MiniCPM-S-1B-sft" ], - "relevancy_score": 64.8 + "relevancy_score": 69.7 + }, + { + "architecture_id": "MiniMaxM2ForCausalLM", + "total_models": 99, + "total_downloads": 12619591, + "min_param_count": 15303312058, + "sample_models": [ + "MiniMaxAI/MiniMax-M2.7", + "MiniMaxAI/MiniMax-M2.5", + "MiniMaxAI/MiniMax-M2", + "mratsim/MiniMax-M2.5-BF16-INT4-AWQ", + "cyankiwi/MiniMax-M2.7-AWQ-4bit", + "cyankiwi/MiniMax-M2.5-AWQ-4bit", + "demon-zombie/MiniMax-M2.7-AWQ-4bit", + "MiniMaxAI/MiniMax-M2.1", + "QuantTrio/MiniMax-M2.7-AWQ", + "JANGQ-AI/MiniMax-M2.7-JANGTQ" + ], + "relevancy_score": 67.3 + }, + { + "architecture_id": "QWenLMHeadModel", + "total_models": 69, + "total_downloads": 1488558, + "min_param_count": 1836828672, + "sample_models": [ + "Qwen/Qwen-VL", + "cckevinn/SeeClick", + "Qwen/Qwen-VL-Chat", + "Qwen/Qwen-7B-Chat", + "Qwen/Qwen-7B", + "TheBloke/Qwen-7B-Chat-AWQ", + "Qwen/Qwen-Audio-Chat", + "Qwen/Qwen-14B", + "Qwen/Qwen-14B-Chat", + "Qwen/Qwen-1_8B" + ], + "relevancy_score": 66.1 }, { "architecture_id": "T5WithLMHeadModel", - "total_models": 42, - "total_downloads": 1032735, + "total_models": 56, + "total_downloads": 1336656, "min_param_count": 222903936, "sample_models": [ "google-t5/t5-3b", @@ -159,12 +178,31 @@ "unicamp-dl/ptt5-large-portuguese-vocab", "gagan3012/k2t" ], - "relevancy_score": 61.4 + "relevancy_score": 66.1 + }, + { + "architecture_id": "FalconH1ForCausalLM", + "total_models": 55, + "total_downloads": 736957, + "min_param_count": 91131072, + "sample_models": [ + "tiiuae/Falcon-H1-0.5B-Base", + "tiiuae/Falcon-H1-0.5B-Instruct", + "tiiuae/Falcon-H1-7B-Instruct", + "tiiuae/Falcon-H1-3B-Instruct", + "tiiuae/Falcon-H1R-7B", + "tiiuae/Falcon-H1-Tiny-90M-Instruct", + "tiiuae/Falcon-H1-1.5B-Instruct", + "tiiuae/Falcon-H1-1.5B-Deep-Instruct", + "tiiuae/Falcon-H1-34B-Instruct", + "tiiuae/Falcon-H1-34B-Base" + ], + "relevancy_score": 64.5 }, { "architecture_id": "Glm4MoeLiteForCausalLM", - "total_models": 27, - "total_downloads": 7000683, + "total_models": 35, + "total_downloads": 9799340, "min_param_count": 5768240, "sample_models": [ "zai-org/GLM-4.7-Flash", @@ -178,31 +216,12 @@ "cyankiwi/GLM-4.7-Flash-AWQ-8bit", "tiny-random/glm-4.7-flash" ], - "relevancy_score": 61.0 - }, - { - "architecture_id": "QWenLMHeadModel", - "total_models": 53, - "total_downloads": 1173335, - "min_param_count": 1836828672, - "sample_models": [ - "Qwen/Qwen-VL", - "cckevinn/SeeClick", - "Qwen/Qwen-VL-Chat", - "Qwen/Qwen-7B-Chat", - "Qwen/Qwen-7B", - "TheBloke/Qwen-7B-Chat-AWQ", - "Qwen/Qwen-Audio-Chat", - "Qwen/Qwen-14B", - "Qwen/Qwen-14B-Chat", - "Qwen/Qwen-1_8B" - ], - "relevancy_score": 60.9 + "relevancy_score": 64.1 }, { "architecture_id": "ParlerTTSForConditionalGeneration", - "total_models": 33, - "total_downloads": 2560642, + "total_models": 42, + "total_downloads": 3355473, "min_param_count": 647027186, "sample_models": [ "ai4bharat/indic-parler-tts", @@ -216,50 +235,12 @@ "atlithor/EmotiveIcelandic", "2121-8/japanese-parler-tts-mini" ], - "relevancy_score": 60.7 - }, - { - "architecture_id": "MiniMaxM2ForCausalLM", - "total_models": 76, - "total_downloads": 10035611, - "min_param_count": 15303312058, - "sample_models": [ - "MiniMaxAI/MiniMax-M2.7", - "MiniMaxAI/MiniMax-M2.5", - "MiniMaxAI/MiniMax-M2", - "mratsim/MiniMax-M2.5-BF16-INT4-AWQ", - "cyankiwi/MiniMax-M2.7-AWQ-4bit", - "cyankiwi/MiniMax-M2.5-AWQ-4bit", - "demon-zombie/MiniMax-M2.7-AWQ-4bit", - "MiniMaxAI/MiniMax-M2.1", - "QuantTrio/MiniMax-M2.7-AWQ", - "JANGQ-AI/MiniMax-M2.7-JANGTQ" - ], - "relevancy_score": 60.1 - }, - { - "architecture_id": "FalconH1ForCausalLM", - "total_models": 42, - "total_downloads": 550692, - "min_param_count": 91131072, - "sample_models": [ - "tiiuae/Falcon-H1-0.5B-Base", - "tiiuae/Falcon-H1-0.5B-Instruct", - "tiiuae/Falcon-H1-7B-Instruct", - "tiiuae/Falcon-H1-3B-Instruct", - "tiiuae/Falcon-H1R-7B", - "tiiuae/Falcon-H1-Tiny-90M-Instruct", - "tiiuae/Falcon-H1-1.5B-Instruct", - "tiiuae/Falcon-H1-1.5B-Deep-Instruct", - "tiiuae/Falcon-H1-34B-Instruct", - "tiiuae/Falcon-H1-34B-Base" - ], - "relevancy_score": 60.1 + "relevancy_score": 63.9 }, { "architecture_id": "Gemma4AssistantForCausalLM", - "total_models": 24, - "total_downloads": 3638749, + "total_models": 33, + "total_downloads": 5152069, "min_param_count": 77993476, "sample_models": [ "google/gemma-4-31B-it-assistant", @@ -270,33 +251,15 @@ "google/gemma-4-E4B-it-qat-q4_0-unquantized-assistant", "google/gemma-4-26B-A4B-it-qat-q4_0-unquantized-assistant", "guardiangate1775/gemma-4-26B-A4B-it-assistant-4bit", - "google/gemma-4-E2B-it-qat-q4_0-unquantized-assistant" - ], - "relevancy_score": 58.8 - }, - { - "architecture_id": "HunYuanDenseV1ForCausalLM", - "total_models": 41, - "total_downloads": 333772, - "min_param_count": 539010048, - "sample_models": [ - "tencent/Hy-MT2-1.8B", - "tencent/Hunyuan-7B-Instruct", - "tencent/Hy-MT2-7B", - "tencent/HY-MT1.5-1.8B", - "tencent/Hunyuan-MT-7B", - "shawnw3i/Hy-MT2-1.8B-AWQ", - "tencent/HY-MT1.5-7B", - "huihui-ai/Huihui-HY-MT1.5-7B-abliterated", - "tencent/Hunyuan-MT-Chimera-7B", - "tencent/HY-MT1.5-1.8B-GPTQ-Int4" + "google/gemma-4-E2B-it-qat-q4_0-unquantized-assistant", + "kunhunjon/gemma-4-26B-A4B-it-qat-assistant-w4a16-ct" ], - "relevancy_score": 58.7 + "relevancy_score": 62.2 }, { "architecture_id": "ExaoneForCausalLM", - "total_models": 30, - "total_downloads": 1208776, + "total_models": 40, + "total_downloads": 1649174, "min_param_count": 236983296, "sample_models": [ "LGAI-EXAONE/EXAONE-3.5-32B-Instruct-AWQ", @@ -310,46 +273,12 @@ "hyper-accel/tiny-random-exaone", "LGAI-EXAONE/EXAONE-Deep-2.4B" ], - "relevancy_score": 58.2 - }, - { - "architecture_id": "M2M100ForConditionalGeneration", - "total_models": 22, - "total_downloads": 2549811, - "min_param_count": 332735488, - "sample_models": [ - "AbteeXAILab/lumynax-translate-nllb-200-3b", - "raxtemur/SONAR_200_text_decoder", - "facebook/m2m100_1.2B", - "facebook/nllb-200-distilled-600M", - "facebook/m2m100_418M", - "facebook/nllb-200-distilled-1.3B", - "facebook/nllb-200-1.3B", - "alirezamsh/small100", - "stas/tiny-m2m_100", - "Xenova/nllb-200-distilled-600M" - ], - "relevancy_score": 57.5 - }, - { - "architecture_id": "Phi3VForCausalLM", - "total_models": 14, - "total_downloads": 6908063, - "min_param_count": 304612720, - "sample_models": [ - "microsoft/Phi-3.5-vision-instruct", - "TIGER-Lab/VLM2Vec-Full", - "microsoft/Phi-3-vision-128k-instruct", - "failspy/Phi-3-vision-128k-instruct-abliterated-alpha", - "Desm0nt/Phi-3-HornyVision-128k-instruct", - "yujiepan/phi-3-vision-tiny-random" - ], - "relevancy_score": 57.2 + "relevancy_score": 61.8 }, { "architecture_id": "Cohere2ForCausalLM", - "total_models": 27, - "total_downloads": 1085565, + "total_models": 36, + "total_downloads": 1515892, "min_param_count": 2099096, "sample_models": [ "trl-internal-testing/tiny-Cohere2ForCausalLM", @@ -362,42 +291,27 @@ "CohereLabs/tiny-aya-earth", "CohereLabs/tiny-aya-water" ], - "relevancy_score": 57.1 - }, - { - "architecture_id": "H2OVLChatModel", - "total_models": 6, - "total_downloads": 6524500, - "min_param_count": 826295808, - "sample_models": [ - "h2oai/h2ovl-mississippi-800m", - "h2oai/h2ovl-mississippi-2b" - ], - "relevancy_score": 54.8 + "relevancy_score": 60.5 }, { - "architecture_id": "MBartForConditionalGeneration", - "total_models": 20, - "total_downloads": 703568, - "min_param_count": 379691717, + "architecture_id": "Phi3VForCausalLM", + "total_models": 17, + "total_downloads": 9004672, + "min_param_count": 304612720, "sample_models": [ - "facebook/mbart-large-50-one-to-many-mmt", - "ai4bharat/IndicBART", - "moussaKam/mbarthez", - "facebook/mgenre-wiki", - "facebook/mbart-large-50-many-to-many-mmt", - "facebook/mbart-large-50", - "facebook/mbart-large-cc25", - "Babelscape/mrebel-large", - "bmd1905/vietnamese-correction-v2", - "moussaKam/barthez" + "microsoft/Phi-3.5-vision-instruct", + "TIGER-Lab/VLM2Vec-Full", + "microsoft/Phi-3-vision-128k-instruct", + "failspy/Phi-3-vision-128k-instruct-abliterated-alpha", + "Desm0nt/Phi-3-HornyVision-128k-instruct", + "yujiepan/phi-3-vision-tiny-random" ], - "relevancy_score": 54.2 + "relevancy_score": 58.7 }, { "architecture_id": "YatGPTForCausalLM", - "total_models": 39, - "total_downloads": 45146, + "total_models": 52, + "total_downloads": 53290, "min_param_count": 261096362, "sample_models": [ "mlnomad/yatnmn-softplus-sb-ca-d12-chinchilla-261M-seed2-pytorch", @@ -411,12 +325,31 @@ "mlnomad/yatnmn-softplus-d22-chinchilla-1B-pytorch", "mlnomad/yatnmn-softplus-sb-ca-d12-chinchilla-261M-pytorch" ], - "relevancy_score": 53.9 + "relevancy_score": 58.1 + }, + { + "architecture_id": "M2M100ForConditionalGeneration", + "total_models": 24, + "total_downloads": 2552801, + "min_param_count": 332735488, + "sample_models": [ + "AbteeXAILab/lumynax-translate-nllb-200-3b", + "raxtemur/SONAR_200_text_decoder", + "facebook/m2m100_1.2B", + "facebook/nllb-200-distilled-600M", + "facebook/m2m100_418M", + "facebook/nllb-200-distilled-1.3B", + "facebook/nllb-200-1.3B", + "alirezamsh/small100", + "stas/tiny-m2m_100", + "Xenova/nllb-200-distilled-600M" + ], + "relevancy_score": 58.1 }, { "architecture_id": "LlamaForCausalLMEagle3", - "total_models": 26, - "total_downloads": 236284, + "total_models": 35, + "total_downloads": 301662, "min_param_count": 215481600, "sample_models": [ "lightseekorg/kimi-k2.6-eagle3", @@ -429,38 +362,49 @@ "nvidia/gpt-oss-120b-Eagle3-throughput", "nvidia/Qwen3-235B-A22B-Eagle3" ], - "relevancy_score": 53.6 + "relevancy_score": 56.8 }, { - "architecture_id": "FalconMambaForCausalLM", - "total_models": 14, - "total_downloads": 1162249, - "min_param_count": 525400, - "sample_models": [ + "architecture_id": "H2OVLChatModel", + "total_models": 8, + "total_downloads": 8853218, + "min_param_count": 826295808, + "sample_models": [ + "h2oai/h2ovl-mississippi-800m", + "h2oai/h2ovl-mississippi-2b" + ], + "relevancy_score": 56.0 + }, + { + "architecture_id": "FalconMambaForCausalLM", + "total_models": 19, + "total_downloads": 1583076, + "min_param_count": 525400, + "sample_models": [ "tiiuae/falcon-mamba-7b", "tiiuae/falcon-mamba-tiny-dev", "trl-internal-testing/tiny-FalconMambaForCausalLM", "tiiuae/falcon-mamba-7b-instruct", "tiiuae/Falcon3-Mamba-7B-Instruct" ], - "relevancy_score": 53.5 + "relevancy_score": 55.6 }, { "architecture_id": "BambaForCausalLM", - "total_models": 9, - "total_downloads": 2166327, + "total_models": 12, + "total_downloads": 2698061, "min_param_count": 33110760, "sample_models": [ "hmellor/tiny-random-BambaForCausalLM", "ibm-ai-platform/Bamba-9B-v1", "ibm-ai-platform/Bamba-9B-v2" ], - "relevancy_score": 53.3 + "relevancy_score": 54.7 }, { "architecture_id": "LightOnOCRForConditionalGeneration", - "total_models": 21, - "total_downloads": 1407880, + "total_models": 28, + "total_downloads": 1835166, "min_param_count": 1005647872, "sample_models": [ "lightonai/LightOnOCR-2-1B", @@ -471,12 +415,50 @@ "lightonai/LightOnOCR-2-1B-ocr-soup", "lightonai/LightOnOCR-2-1B-bbox-base" ], - "relevancy_score": 51.9 + "relevancy_score": 54.5 + }, + { + "architecture_id": "LlavaLlamaForCausalLM", + "total_models": 55, + "total_downloads": 1690361, + "min_param_count": 7062902784, + "sample_models": [ + "liuhaotian/llava-v1.5-7b", + "ChantalPellegrini/RaDialog-interactive-radiology-report-generation", + "liuhaotian/llava-v1.6-vicuna-7b", + "liuhaotian/llava-v1.5-13b", + "liuhaotian/llava-v1.6-vicuna-13b", + "wisdomik/Quilt-Llava-v1.5-7b", + "lmms-lab/llama3-llava-next-8b", + "Vishal24/llava_1.5_image_classification_v3", + "ManishThota/Ollama_Video_llama_7B", + "etri-vilab/Ko-LLaVA-13b" + ], + "relevancy_score": 54.3 + }, + { + "architecture_id": "MBartForConditionalGeneration", + "total_models": 20, + "total_downloads": 703568, + "min_param_count": 379691717, + "sample_models": [ + "facebook/mbart-large-50-one-to-many-mmt", + "ai4bharat/IndicBART", + "moussaKam/mbarthez", + "facebook/mgenre-wiki", + "facebook/mbart-large-50-many-to-many-mmt", + "facebook/mbart-large-50", + "facebook/mbart-large-cc25", + "Babelscape/mrebel-large", + "bmd1905/vietnamese-correction-v2", + "moussaKam/barthez" + ], + "relevancy_score": 54.2 }, { "architecture_id": "JambaForCausalLM", - "total_models": 21, - "total_downloads": 128306, + "total_models": 28, + "total_downloads": 179399, "min_param_count": 127679344, "sample_models": [ "ai21labs/AI21-Jamba-Mini-1.5", @@ -487,12 +469,31 @@ "ai21labs/AI21-Jamba-Reasoning-3B", "ai21labs/AI21-Jamba2-3B" ], - "relevancy_score": 50.9 + "relevancy_score": 53.6 + }, + { + "architecture_id": "MultiScaleForCausalLM", + "total_models": 101, + "total_downloads": 71026, + "min_param_count": null, + "sample_models": [ + "KoinicLabs/AXL-Translate", + "KoinicLabs/AXL-Chat-10M", + "KoinicLabs/AXL-Secure-Lion", + "KoinicLabs/AXL-Code-1B", + "KoinicLabs/AXL-Secure-10M", + "KoinicLabs/AXL-Comment-5M", + "KoinicLabs/AXL-Vision-0.8M", + "KoinicLabs/AXL-Vision-v2", + "KoinicLabs/AXL-Chat-Pro", + "KoinicLabs/AXL-Refactor-Lion" + ], + "relevancy_score": 53.0 }, { "architecture_id": "Eagle3Speculator", - "total_models": 15, - "total_downloads": 296816, + "total_models": 20, + "total_downloads": 409203, "min_param_count": 950186496, "sample_models": [ "RedHatAI/Qwen3-8B-speculator.eagle3", @@ -501,12 +502,31 @@ "RedHatAI/Qwen3-32B-speculator.eagle3", "RedHatAI/Qwen3-14B-speculator.eagle3" ], - "relevancy_score": 50.9 + "relevancy_score": 53.0 + }, + { + "architecture_id": "Ovis", + "total_models": 24, + "total_downloads": 1176708, + "min_param_count": 1267417974, + "sample_models": [ + "AIDC-AI/Ovis2-1B", + "AIDC-AI/Ovis1.6-Llama3.2-3B", + "AIDC-AI/Ovis1.6-Gemma2-9B", + "AIDC-AI/Ovis2-4B", + "AIDC-AI/Ovis2-8B", + "AIDC-AI/Ovis2-34B", + "ATH-MaaS/Ovis2-1B", + "ATH-MaaS/Ovis1.6-Llama3.2-3B", + "ATH-MaaS/Ovis1.6-Gemma2-9B", + "ATH-MaaS/Ovis2-4B" + ], + "relevancy_score": 52.4 }, { "architecture_id": "LlavaQwenForCausalLM", - "total_models": 12, - "total_downloads": 350696, + "total_models": 16, + "total_downloads": 514384, "min_param_count": 893618208, "sample_models": [ "lmms-lab/llava-onevision-qwen2-7b-ov", @@ -514,23 +534,27 @@ "lmms-lab/llava-onevision-qwen2-0.5b-ov", "lmms-lab/LongVA-7B" ], - "relevancy_score": 50.4 + "relevancy_score": 52.4 }, { - "architecture_id": "HfMoondream", - "total_models": 6, - "total_downloads": 5626256, - "min_param_count": 1927237104, + "architecture_id": "Qwen2MoeForCausalLM", + "total_models": 24, + "total_downloads": 1075803, + "min_param_count": 1763452928, "sample_models": [ - "vikhyatk/moondream2", - "moondream/moondream3-preview" + "Qwen/Qwen1.5-MoE-A2.7B", + "Qwen/Qwen1.5-MoE-A2.7B-Chat", + "Qwen/Qwen2-57B-A14B-Instruct", + "Qwen/Qwen2-57B-A14B", + "Qwen/Qwen1.5-MoE-A2.7B-Chat-GPTQ-Int4", + "hyper-accel/ci-random-qwen2-moe-a3b" ], - "relevancy_score": 50.4 + "relevancy_score": 52.2 }, { "architecture_id": "ArceeForCausalLM", - "total_models": 12, - "total_downloads": 314521, + "total_models": 16, + "total_downloads": 442598, "min_param_count": 4129088, "sample_models": [ "optimum-intel-internal-testing/tiny-random-ArceeForCausalLM", @@ -538,22 +562,65 @@ "onnx-internal-testing/tiny-random-ArceeForCausalLM", "arcee-ai/AFM-4.5B" ], + "relevancy_score": 52.0 + }, + { + "architecture_id": "HfMoondream", + "total_models": 8, + "total_downloads": 7491036, + "min_param_count": 1927237104, + "sample_models": [ + "vikhyatk/moondream2", + "moondream/moondream3-preview" + ], + "relevancy_score": 51.7 + }, + { + "architecture_id": "NemotronLabsDiffusionModel", + "total_models": 24, + "total_downloads": 3704132, + "min_param_count": 3831659520, + "sample_models": [ + "nvidia/Nemotron-Labs-Diffusion-8B-Base", + "nvidia/Nemotron-Labs-Diffusion-3B-Base", + "nvidia/Nemotron-Labs-Diffusion-8B", + "nvidia/Nemotron-Labs-Diffusion-3B", + "nvidia/Nemotron-Labs-Diffusion-14B-Base", + "nvidia/Nemotron-Labs-Diffusion-14B" + ], + "relevancy_score": 50.9 + }, + { + "architecture_id": "HunYuanVLForConditionalGeneration", + "total_models": 4, + "total_downloads": 1140251, + "min_param_count": 996208112, + "sample_models": [ + "tencent/HunyuanOCR" + ], + "relevancy_score": 50.5 + }, + { + "architecture_id": "Gemma4UnifiedAssistantForCausalLM", + "total_models": 12, + "total_downloads": 299496, + "min_param_count": 422856964, + "sample_models": [ + "google/gemma-4-12B-it-assistant", + "google/gemma-4-12B-it-qat-q4_0-unquantized-assistant", + "kunhunjon/gemma-4-12B-it-qat-assistant-w4a16-ct" + ], "relevancy_score": 50.1 }, { - "architecture_id": "Ovis", - "total_models": 18, - "total_downloads": 875971, - "min_param_count": 1267417974, + "architecture_id": "OpenAIGPTLMHeadModel", + "total_models": 4, + "total_downloads": 945634, + "min_param_count": 119680512, "sample_models": [ - "AIDC-AI/Ovis2-1B", - "AIDC-AI/Ovis1.6-Llama3.2-3B", - "AIDC-AI/Ovis1.6-Gemma2-9B", - "AIDC-AI/Ovis2-4B", - "AIDC-AI/Ovis2-8B", - "AIDC-AI/Ovis2-34B" + "openai-community/openai-gpt" ], - "relevancy_score": 50.0 + "relevancy_score": 50.1 }, { "architecture_id": "PegasusForConditionalGeneration", @@ -574,148 +641,233 @@ "relevancy_score": 50.0 }, { - "architecture_id": "Qwen2MoeForCausalLM", - "total_models": 18, - "total_downloads": 790323, - "min_param_count": 1763452928, + "architecture_id": "Step3p7ForConditionalGeneration", + "total_models": 21, + "total_downloads": 546253, + "min_param_count": 2956129024, "sample_models": [ - "Qwen/Qwen1.5-MoE-A2.7B", - "Qwen/Qwen1.5-MoE-A2.7B-Chat", - "Qwen/Qwen2-57B-A14B-Instruct", - "Qwen/Qwen2-57B-A14B", - "Qwen/Qwen1.5-MoE-A2.7B-Chat-GPTQ-Int4", - "hyper-accel/ci-random-qwen2-moe-a3b" + "stepfun-ai/Step-3.7-Flash", + "cyankiwi/Step-3.7-Flash-AWQ-INT4", + "Hikari07jp/Step-3.7-Flash-MTP-draft", + "0xSero/Step-3.7-Flash-173B", + "0xSero/Step-3.7-Flash-148B", + "nameistoken/Step-3.7-Flash-Quark-W4A16" ], - "relevancy_score": 49.8 + "relevancy_score": 49.9 }, { - "architecture_id": "HunYuanVLForConditionalGeneration", - "total_models": 3, - "total_downloads": 865778, - "min_param_count": 996208112, + "architecture_id": "SeedOssForCausalLM", + "total_models": 15, + "total_downloads": 187264, + "min_param_count": 2497064, "sample_models": [ - "tencent/HunyuanOCR" + "ByteDance-Seed/Seed-OSS-36B-Instruct", + "cyankiwi/Hermes-4.3-36B-AWQ-4bit", + "ByteDance-Seed/Seed-OSS-36B-Base", + "tiny-random/seed-oss" ], - "relevancy_score": 49.6 + "relevancy_score": 49.9 }, { - "architecture_id": "LlavaLlamaForCausalLM", - "total_models": 40, - "total_downloads": 1414584, - "min_param_count": 7062902784, + "architecture_id": "NemotronForCausalLM", + "total_models": 23, + "total_downloads": 54883, + "min_param_count": 2150720, "sample_models": [ - "liuhaotian/llava-v1.5-7b", - "ChantalPellegrini/RaDialog-interactive-radiology-report-generation", - "liuhaotian/llava-v1.6-vicuna-7b", - "liuhaotian/llava-v1.5-13b", - "liuhaotian/llava-v1.6-vicuna-13b", - "wisdomik/Quilt-Llava-v1.5-7b", - "lmms-lab/llama3-llava-next-8b", - "Vishal24/llava_1.5_image_classification_v3", - "ManishThota/Ollama_Video_llama_7B", - "etri-vilab/Ko-LLaVA-13b" + "badaoui/tiny-random-NemotronForCausalLM", + "mgoin/Nemotron-4-340B-Instruct-hf", + "mgoin/nemotron-3-8b-chat-4k-sft-hf", + "thhaus/nemotron3-8b", + "domyn/Domyn-Small-v1.0", + "OpenLLM-France/Luciole-1B-Base" ], - "relevancy_score": 49.5 + "relevancy_score": 49.7 }, { - "architecture_id": "OpenAIGPTLMHeadModel", - "total_models": 3, - "total_downloads": 705713, - "min_param_count": 119680512, + "architecture_id": "Idefics3ForConditionalGeneration", + "total_models": 8, + "total_downloads": 443613, + "min_param_count": 257517120, "sample_models": [ - "openai-community/openai-gpt" + "ibm-granite/granite-docling-258M", + "docling-project/ScreenVLM" ], - "relevancy_score": 49.2 + "relevancy_score": 49.7 }, { - "architecture_id": "Qwen2AudioForConditionalGeneration", - "total_models": 3, - "total_downloads": 620941, - "min_param_count": 5033440, + "architecture_id": "RITAModelForCausalLM", + "total_models": 16, + "total_downloads": 137323, + "min_param_count": 85096320, "sample_models": [ - "Qwen/Qwen2-Audio-7B-Instruct", - "Qwen/Qwen2-Audio-7B", - "trl-internal-testing/tiny-Qwen2AudioForConditionalGeneration" + "lightonai/RITA_s", + "lightonai/RITA_m", + "lightonai/RITA_l", + "lightonai/RITA_xl" ], - "relevancy_score": 48.9 + "relevancy_score": 49.6 }, { - "architecture_id": "Idefics3ForConditionalGeneration", - "total_models": 6, - "total_downloads": 349241, - "min_param_count": 257517120, + "architecture_id": "Zamba2ForCausalLM", + "total_models": 12, + "total_downloads": 1547927, + "min_param_count": 1215064704, "sample_models": [ - "ibm-granite/granite-docling-258M", - "docling-project/ScreenVLM" + "Zyphra/Zamba2-1.2B-instruct", + "Zyphra/Zamba2-7B-Instruct", + "Zyphra/Zamba2-2.7B" ], - "relevancy_score": 48.6 + "relevancy_score": 49.5 }, { - "architecture_id": "NemotronLabsDiffusionModel", - "total_models": 18, - "total_downloads": 2679242, - "min_param_count": 3831659520, + "architecture_id": "LlavaQwen2ForCausalLM", + "total_models": 20, + "total_downloads": 71317, + "min_param_count": 758833760, "sample_models": [ - "nvidia/Nemotron-Labs-Diffusion-8B-Base", - "nvidia/Nemotron-Labs-Diffusion-3B-Base", - "nvidia/Nemotron-Labs-Diffusion-8B", - "nvidia/Nemotron-Labs-Diffusion-3B", - "nvidia/Nemotron-Labs-Diffusion-14B-Base", - "nvidia/Nemotron-Labs-Diffusion-14B" + "qnguyen3/nanoLLaVA", + "apple/FastVLM-0.5B", + "apple/FastVLM-1.5B", + "FreedomIntelligence/HuatuoGPT-Vision-7B", + "apple/FastVLM-7B" ], - "relevancy_score": 48.4 + "relevancy_score": 49.4 }, { - "architecture_id": "Gemma4UnifiedAssistantForCausalLM", - "total_models": 9, - "total_downloads": 192274, - "min_param_count": 422856964, + "architecture_id": "RWKV7ForCausalLM", + "total_models": 21, + "total_downloads": 59885, + "min_param_count": 27388416, "sample_models": [ - "google/gemma-4-12B-it-assistant", - "google/gemma-4-12B-it-qat-q4_0-unquantized-assistant", - "kunhunjon/gemma-4-12B-it-qat-assistant-w4a16-ct" + "yashmahe2018/vanilla-skip-rwkv7-strict-small-babylm2026", + "fla-hub/rwkv7-0.4B-g1", + "RWKV/RWKV7-Goose-World2.9-0.4B-HF", + "yashmahe2018/rwkv7-halfclm-strict-small", + "RWKV/RWKV7-Goose-World3-1.5B-HF", + "fla-hub/rwkv7-1.5B-world" ], - "relevancy_score": 48.2 + "relevancy_score": 49.3 + }, + { + "architecture_id": "HyenaDNAForCausalLM", + "total_models": 20, + "total_downloads": 69068, + "min_param_count": 450712, + "sample_models": [ + "LongSafari/hyenadna-small-32k-seqlen-hf", + "LongSafari/hyenadna-tiny-1k-seqlen-hf", + "LongSafari/hyenadna-medium-160k-seqlen-hf", + "LongSafari/hyenadna-large-1m-seqlen-hf", + "LongSafari/hyenadna-tiny-1k-seqlen-d256-hf", + "LongSafari/hyenadna-tiny-16k-seqlen-d128-hf" + ], + "relevancy_score": 49.3 + }, + { + "architecture_id": "DotsOCRForCausalLM", + "total_models": 19, + "total_downloads": 3591921, + "min_param_count": 3039179264, + "sample_models": [ + "rednote-hilab/dots.mocr", + "rednote-hilab/dots.ocr", + "rednote-hilab/dots.mocr-svg", + "kristaller486/dots.ocr-1.5", + "davanstrien/dots.ocr-1.5" + ], + "relevancy_score": 49.3 + }, + { + "architecture_id": "ProGen2ForPreTraining", + "total_models": 28, + "total_downloads": 21321, + "min_param_count": 151158821, + "sample_models": [ + "multimolecule/progen2-small", + "multimolecule/progen2-medium", + "multimolecule/progen2-base", + "multimolecule/progen2-xlarge", + "multimolecule/progen2-large", + "multimolecule/progen2-oas", + "multimolecule/progen2-bfd90" + ], + "relevancy_score": 49.2 + }, + { + "architecture_id": "DynColMaskMoELLaVAQwen2ForCausalLM", + "total_models": 28, + "total_downloads": 20383, + "min_param_count": 936779392, + "sample_models": [ + "KKHYA/llavaqwen2.5-0.5b-finetune-dyn-col-mask-moe-sparse-4e-2k-sp0.9-r128_20260422_053621", + "KKHYA/llavaqwen2.5-0.5b-finetune-dyn-col-mask-moe-sparse-4e-2k-sp0.9-r64_20260422_030205", + "KKHYA/llavaqwen2.5-0.5b-finetune-dyn-col-mask-moe-sparse-4e-2k-sp0.9-r32_20260422_002811", + "KKHYA/llavaqwen2.5-0.5b-finetune-dyn-col-mask-moe-sparse-4e-2k-sp0.9-r4_20260421_191943", + "KKHYA/llavaqwen2.5-0.5b-finetune-dyn-col-mask-moe-sparse-4e-2k-sp0.9-r8_20260421_215406", + "KKHYA/llavaqwen2.5-0.5b-finetune-dyn-col-mask-moe-sparse-4e-2k-sp0.9-r16_20260420_180515", + "KKHYA/llavaqwen2.5-0.5b-finetune-dyn-col-mask-moe-sparse-4e-2k-sp0.9-r16_20260418_062801" + ], + "relevancy_score": 49.1 }, { "architecture_id": "SundialForPrediction", - "total_models": 3, - "total_downloads": 443278, + "total_models": 4, + "total_downloads": 591580, "min_param_count": 128329680, "sample_models": [ "thuml/sundial-base-128m" ], - "relevancy_score": 48.2 + "relevancy_score": 49.1 }, { - "architecture_id": "SeedOssForCausalLM", - "total_models": 11, - "total_downloads": 136998, - "min_param_count": 2497064, + "architecture_id": "Qwen2AudioForConditionalGeneration", + "total_models": 3, + "total_downloads": 620941, + "min_param_count": 5033440, "sample_models": [ - "ByteDance-Seed/Seed-OSS-36B-Instruct", - "cyankiwi/Hermes-4.3-36B-AWQ-4bit", - "ByteDance-Seed/Seed-OSS-36B-Base", - "tiny-random/seed-oss" + "Qwen/Qwen2-Audio-7B-Instruct", + "Qwen/Qwen2-Audio-7B", + "trl-internal-testing/tiny-Qwen2AudioForConditionalGeneration" ], - "relevancy_score": 48.1 + "relevancy_score": 49.0 }, { - "architecture_id": "Zamba2ForCausalLM", - "total_models": 9, - "total_downloads": 1142055, - "min_param_count": 1215064704, + "architecture_id": "DeciLMForCausalLM", + "total_models": 32, + "total_downloads": 3300370, + "min_param_count": 7043551232, "sample_models": [ - "Zyphra/Zamba2-1.2B-instruct", - "Zyphra/Zamba2-7B-Instruct", - "Zyphra/Zamba2-2.7B" + "nvidia/Llama-3_3-Nemotron-Super-49B-v1_5", + "nvidia/Llama-3_3-Nemotron-Super-49B-v1", + "nvidia/Llama-3_1-Nemotron-Ultra-253B-v1", + "nvidia/Llama-3_1-Nemotron-51B-Instruct", + "Deci/DeciLM-7B-instruct", + "nvidia/Llama-3_1-Nemotron-Ultra-253B-CPT-v1", + "etsien/Llama-3_3-Nemotron-Super-49B-v1_5-GPTQ-w4a8", + "NewstaR/Porpoise-6b-instruct", + "Danielbrdz/Barcenas-6b", + "NewstaR/StableGalen-6b" + ], + "relevancy_score": 48.9 + }, + { + "architecture_id": "OuroForCausalLM", + "total_models": 20, + "total_downloads": 389795, + "min_param_count": 1434652673, + "sample_models": [ + "ByteDance/Ouro-1.4B", + "ByteDance/Ouro-2.6B-Thinking", + "ByteDance/Ouro-1.4B-Thinking", + "ByteDance/Ouro-2.6B", + "KristianS7/Ouro-1.4B" ], - "relevancy_score": 48.0 + "relevancy_score": 48.9 }, { "architecture_id": "BloomModel", - "total_models": 15, - "total_downloads": 70491, + "total_models": 17, + "total_downloads": 88286, "min_param_count": 16156544, "sample_models": [ "bigscience/bigscience-small-testing", @@ -726,25 +878,52 @@ "TurkuNLP/gpt3-finnish-13B", "BelleGroup/BELLE-7B-2M" ], - "relevancy_score": 47.9 + "relevancy_score": 48.9 }, { - "architecture_id": "RITAModelForCausalLM", + "architecture_id": "ProGenForCausalLM", + "total_models": 16, + "total_downloads": 89065, + "min_param_count": 151148576, + "sample_models": [ + "hugohrban/progen2-small", + "hugohrban/progen2-large", + "hugohrban/progen2-base", + "hugohrban/progen2-medium" + ], + "relevancy_score": 48.7 + }, + { + "architecture_id": "BertLMHeadModel", + "total_models": 24, + "total_downloads": 28854, + "min_param_count": 184474880, + "sample_models": [ + "sagawa/molscaletransfer-chemlm-0.06m", + "sagawa/molscaletransfer-chemlm-0.83m", + "sagawa/molscaletransfer-chemlm-2.30m", + "ielab/TILDE", + "dicta-il/BEREL_3.0", + "ENM/scibert_scivocab_cased-new-finetuned-leukaemia" + ], + "relevancy_score": 48.6 + }, + { + "architecture_id": "Eagle3DraftModel", "total_models": 12, - "total_downloads": 100793, - "min_param_count": 85096320, + "total_downloads": 139472, + "min_param_count": 522152832, "sample_models": [ - "lightonai/RITA_s", - "lightonai/RITA_m", - "lightonai/RITA_l", - "lightonai/RITA_xl" + "RedHatAI/gpt-oss-20b-speculator.eagle3", + "RedHatAI/gpt-oss-120b-speculator.eagle3", + "RedHatAI/Qwen3-30B-A3B-Instruct-2507-speculator.eagle3" ], - "relevancy_score": 47.7 + "relevancy_score": 48.4 }, { "architecture_id": "NandiForCausalLM", - "total_models": 14, - "total_downloads": 69602, + "total_models": 16, + "total_downloads": 71927, "min_param_count": 153412928, "sample_models": [ "FrontiersMind/Nandi-Mini-150M-Tool-Calling", @@ -754,116 +933,69 @@ "FrontiersMind/Nandi-Mini-600M-Early-Checkpoint", "FrontiersMind/Nandi-Mini-150M-Instruct" ], - "relevancy_score": 47.5 - }, - { - "architecture_id": "Step3p7ForConditionalGeneration", - "total_models": 15, - "total_downloads": 382404, - "min_param_count": 2956129024, - "sample_models": [ - "stepfun-ai/Step-3.7-Flash", - "cyankiwi/Step-3.7-Flash-AWQ-INT4", - "Hikari07jp/Step-3.7-Flash-MTP-draft", - "0xSero/Step-3.7-Flash-173B", - "0xSero/Step-3.7-Flash-148B" - ], - "relevancy_score": 47.4 - }, - { - "architecture_id": "NemotronForCausalLM", - "total_models": 17, - "total_downloads": 40787, - "min_param_count": 2150720, - "sample_models": [ - "badaoui/tiny-random-NemotronForCausalLM", - "mgoin/Nemotron-4-340B-Instruct-hf", - "mgoin/nemotron-3-8b-chat-4k-sft-hf", - "thhaus/nemotron3-8b", - "domyn/Domyn-Small-v1.0", - "OpenLLM-France/Luciole-1B-Base" - ], - "relevancy_score": 47.3 + "relevancy_score": 48.2 }, { - "architecture_id": "LlavaQwen2ForCausalLM", - "total_models": 15, - "total_downloads": 53603, - "min_param_count": 758833760, + "architecture_id": "OrionForCausalLM", + "total_models": 11, + "total_downloads": 144746, + "min_param_count": 93459456, "sample_models": [ - "qnguyen3/nanoLLaVA", - "apple/FastVLM-0.5B", - "apple/FastVLM-1.5B", - "FreedomIntelligence/HuatuoGPT-Vision-7B", - "apple/FastVLM-7B" + "OrionStarAI/Orion-14B-Chat", + "OrionStarAI/Orion-14B-Base", + "hyper-accel/tiny-random-orion" ], - "relevancy_score": 47.3 + "relevancy_score": 48.2 }, { - "architecture_id": "HyenaDNAForCausalLM", - "total_models": 15, - "total_downloads": 52395, - "min_param_count": 450712, + "architecture_id": "Eagle3DeepseekV2ForCausalLM", + "total_models": 16, + "total_downloads": 459368, + "min_param_count": 1841199104, "sample_models": [ - "LongSafari/hyenadna-small-32k-seqlen-hf", - "LongSafari/hyenadna-tiny-1k-seqlen-hf", - "LongSafari/hyenadna-medium-160k-seqlen-hf", - "LongSafari/hyenadna-large-1m-seqlen-hf", - "LongSafari/hyenadna-tiny-1k-seqlen-d256-hf", - "LongSafari/hyenadna-tiny-16k-seqlen-d128-hf" + "lightseekorg/kimi-k2.6-eagle3-mla", + "lightseekorg/kimi-k2.5-eagle3-mla", + "nvidia/Kimi-K2.5-Thinking-Eagle3", + "nvidia/Kimi-K2.6-Eagle3" ], - "relevancy_score": 47.2 + "relevancy_score": 48.1 }, { - "architecture_id": "DotsOCRForCausalLM", - "total_models": 14, - "total_downloads": 2702911, - "min_param_count": 3039179264, + "architecture_id": "PinyinCodeForCausalLM", + "total_models": 16, + "total_downloads": 55710, + "min_param_count": 33674240, "sample_models": [ - "rednote-hilab/dots.mocr", - "rednote-hilab/dots.ocr", - "rednote-hilab/dots.mocr-svg", - "kristaller486/dots.ocr-1.5", - "davanstrien/dots.ocr-1.5" + "CPSPX/babylm-zho-pinyin-code-33M", + "timorobrecht/full_chinese_gpu3.1", + "CPSPX/babylm-zho-pinyin-code-97M", + "timorobrecht/full_chinese_gpu3.2-dpo" ], - "relevancy_score": 47.2 + "relevancy_score": 47.7 }, { - "architecture_id": "OrionForCausalLM", + "architecture_id": "VibeVoiceForConditionalGeneration", "total_models": 9, - "total_downloads": 108635, - "min_param_count": 93459456, + "total_downloads": 969463, + "min_param_count": 2704021985, "sample_models": [ - "OrionStarAI/Orion-14B-Chat", - "OrionStarAI/Orion-14B-Base", - "hyper-accel/tiny-random-orion" + "microsoft/VibeVoice-1.5B", + "vibevoice/VibeVoice-1.5B", + "bezzam/VibeVoice-1.5B-hf" ], - "relevancy_score": 47.0 + "relevancy_score": 47.6 }, { - "architecture_id": "ProGenForCausalLM", + "architecture_id": "Fgclip2Model", "total_models": 12, - "total_downloads": 67740, - "min_param_count": 151148576, - "sample_models": [ - "hugohrban/progen2-small", - "hugohrban/progen2-large", - "hugohrban/progen2-base", - "hugohrban/progen2-medium" - ], - "relevancy_score": 46.9 - }, - { - "architecture_id": "Eagle3DraftModel", - "total_models": 9, - "total_downloads": 102420, - "min_param_count": 522152832, + "total_downloads": 76451, + "min_param_count": 383803394, "sample_models": [ - "RedHatAI/gpt-oss-20b-speculator.eagle3", - "RedHatAI/gpt-oss-120b-speculator.eagle3", - "RedHatAI/Qwen3-30B-A3B-Instruct-2507-speculator.eagle3" + "qihoo360/fg-clip2-base", + "qihoo360/fg-clip2-so400m", + "qihoo360/fg-clip2-large" ], - "relevancy_score": 46.9 + "relevancy_score": 47.2 }, { "architecture_id": "EncoderDecoderModel", @@ -880,81 +1012,49 @@ "relevancy_score": 46.9 }, { - "architecture_id": "OuroForCausalLM", - "total_models": 15, - "total_downloads": 288169, - "min_param_count": 1434652673, + "architecture_id": "Ernie4_5ForCausalLM", + "total_models": 8, + "total_downloads": 109570, + "min_param_count": 360748032, "sample_models": [ - "ByteDance/Ouro-1.4B", - "ByteDance/Ouro-2.6B-Thinking", - "ByteDance/Ouro-1.4B-Thinking", - "ByteDance/Ouro-2.6B", - "KristianS7/Ouro-1.4B" + "baidu/ERNIE-4.5-0.3B-PT", + "baidu/ERNIE-4.5-0.3B-Base-PT" ], "relevancy_score": 46.8 }, { - "architecture_id": "DeciLMForCausalLM", - "total_models": 26, - "total_downloads": 2687705, - "min_param_count": 7043551232, - "sample_models": [ - "nvidia/Llama-3_3-Nemotron-Super-49B-v1_5", - "nvidia/Llama-3_3-Nemotron-Super-49B-v1", - "nvidia/Llama-3_1-Nemotron-Ultra-253B-v1", - "nvidia/Llama-3_1-Nemotron-51B-Instruct", - "Deci/DeciLM-7B-instruct", - "nvidia/Llama-3_1-Nemotron-Ultra-253B-CPT-v1", - "etsien/Llama-3_3-Nemotron-Super-49B-v1_5-GPTQ-w4a8", - "NewstaR/Porpoise-6b-instruct", - "Danielbrdz/Barcenas-6b", - "NewstaR/StableGalen-6b" - ], - "relevancy_score": 46.7 - }, - { - "architecture_id": "RWKV7ForCausalLM", - "total_models": 15, - "total_downloads": 40337, - "min_param_count": 27388416, + "architecture_id": "RemoteForCausalLM", + "total_models": 4, + "total_downloads": 184810, + "min_param_count": 2054056, "sample_models": [ - "yashmahe2018/vanilla-skip-rwkv7-strict-small-babylm2026", - "fla-hub/rwkv7-0.4B-g1", - "RWKV/RWKV7-Goose-World2.9-0.4B-HF", - "yashmahe2018/rwkv7-halfclm-strict-small", - "RWKV/RWKV7-Goose-World3-1.5B-HF" + "trl-internal-testing/tiny-RemoteForCausalLM" ], "relevancy_score": 46.7 }, { - "architecture_id": "DynColMaskMoELLaVAQwen2ForCausalLM", - "total_models": 21, - "total_downloads": 16575, - "min_param_count": 936779392, + "architecture_id": "GlmForCausalLM", + "total_models": 20, + "total_downloads": 125616, + "min_param_count": 1593427968, "sample_models": [ - "KKHYA/llavaqwen2.5-0.5b-finetune-dyn-col-mask-moe-sparse-4e-2k-sp0.9-r128_20260422_053621", - "KKHYA/llavaqwen2.5-0.5b-finetune-dyn-col-mask-moe-sparse-4e-2k-sp0.9-r64_20260422_030205", - "KKHYA/llavaqwen2.5-0.5b-finetune-dyn-col-mask-moe-sparse-4e-2k-sp0.9-r32_20260422_002811", - "KKHYA/llavaqwen2.5-0.5b-finetune-dyn-col-mask-moe-sparse-4e-2k-sp0.9-r4_20260421_191943", - "KKHYA/llavaqwen2.5-0.5b-finetune-dyn-col-mask-moe-sparse-4e-2k-sp0.9-r8_20260421_215406", - "KKHYA/llavaqwen2.5-0.5b-finetune-dyn-col-mask-moe-sparse-4e-2k-sp0.9-r16_20260420_180515", - "KKHYA/llavaqwen2.5-0.5b-finetune-dyn-col-mask-moe-sparse-4e-2k-sp0.9-r16_20260418_062801" + "zai-org/glm-4-9b-chat-hf", + "zai-org/glm-4-9b-hf", + "zai-org/glm-edge-4b-chat", + "zai-org/glm-edge-1.5b-chat", + "zai-org/glm-4-9b-chat-1m-hf" ], "relevancy_score": 46.6 }, { - "architecture_id": "ProGen2ForPreTraining", - "total_models": 21, - "total_downloads": 16220, - "min_param_count": 151158821, + "architecture_id": "_A2DQwen3LMHeadModel", + "total_models": 11, + "total_downloads": 62966, + "min_param_count": 596049920, "sample_models": [ - "multimolecule/progen2-small", - "multimolecule/progen2-medium", - "multimolecule/progen2-base", - "multimolecule/progen2-xlarge", - "multimolecule/progen2-large", - "multimolecule/progen2-oas", - "multimolecule/progen2-bfd90" + "divelab/OPDLM-0.6B", + "divelab/OPDLM-4B", + "divelab/OPDLM-MATH-4B" ], "relevancy_score": 46.5 }, @@ -974,33 +1074,44 @@ "relevancy_score": 46.5 }, { - "architecture_id": "BertLMHeadModel", - "total_models": 18, - "total_downloads": 23524, - "min_param_count": 184474880, + "architecture_id": "IlamaForCausalLM", + "total_models": 4, + "total_downloads": 1126605, + "min_param_count": 1235814400, "sample_models": [ - "sagawa/molscaletransfer-chemlm-0.06m", - "sagawa/molscaletransfer-chemlm-0.83m", - "sagawa/molscaletransfer-chemlm-2.30m", - "ielab/TILDE", - "dicta-il/BEREL_3.0", - "ENM/scibert_scivocab_cased-new-finetuned-leukaemia" + "hmellor/Ilama-3.2-1B" ], - "relevancy_score": 46.4 + "relevancy_score": 46.5 }, { - "architecture_id": "Eagle3DeepseekV2ForCausalLM", - "total_models": 12, - "total_downloads": 349120, - "min_param_count": 1841199104, + "architecture_id": "Starcoder2ForCausalLM", + "total_models": 21, + "total_downloads": 675003, + "min_param_count": 3030371328, "sample_models": [ - "lightseekorg/kimi-k2.6-eagle3-mla", - "lightseekorg/kimi-k2.5-eagle3-mla", - "nvidia/Kimi-K2.5-Thinking-Eagle3", - "nvidia/Kimi-K2.6-Eagle3" + "bigcode/starcoder2-3b", + "bigcode/starcoder2-7b", + "bigcode/starcoder2-15b", + "bigcode/starcoder2-15b-instruct-v0.1", + "Johnblick187/starcoder2-15b", + "TechxGenus/starcoder2-7b-GPTQ" ], "relevancy_score": 46.4 }, + { + "architecture_id": "VisionEncoderDecoderModel", + "total_models": 17, + "total_downloads": 25186, + "min_param_count": 239195904, + "sample_models": [ + "AbteeXAILab/lumynax-ocr-trocr-large-printed", + "AbteeXAILab/lumynax-doc-nougat-base", + "AbteeXAILab/lumynax-ocr-trocr-large-handwritten", + "AbteeXAILab/lumynax-doc-donut-base", + "Molkaatb/ChestX" + ], + "relevancy_score": 46.3 + }, { "architecture_id": "FSMTForConditionalGeneration", "total_models": 5, @@ -1016,49 +1127,104 @@ "relevancy_score": 46.3 }, { - "architecture_id": "VibeVoiceForConditionalGeneration", - "total_models": 6, - "total_downloads": 721604, - "min_param_count": 2704021985, + "architecture_id": "SDARForCausalLM", + "total_models": 21, + "total_downloads": 89602, + "min_param_count": 2031739904, "sample_models": [ - "microsoft/VibeVoice-1.5B", - "vibevoice/VibeVoice-1.5B" + "JetLM/SDAR-8B-Chat", + "JetLM/SDAR-1.7B-Chat", + "JetLM/SDAR-8B-Chat-b32", + "yifanyu/I-DLM-8B", + "JetLM/SDAR-1.7B-Chat-b32", + "JetLM/SDAR-8B-Chat-b64" + ], + "relevancy_score": 46.1 + }, + { + "architecture_id": "Qwen3_5MoeForCausalLM", + "total_models": 37, + "total_downloads": 58152, + "min_param_count": 5642996864, + "sample_models": [ + "dawncr0w/Qwen3.6-35B-A3B-Uncensored-HauhauCS-Aggressive-text-oQ4", + "samuelfaj/Qwen3.6-35B-A3B-NSC-ACE-SABER-8bit-MTPLX-Optimized-Speed", + "FINAL-Bench/Darwin-36B-Opus", + "0xSero/Qwen3.6-28B", + "samuelfaj/Qwen3.6-35B-A3B-NSC-ACE-SABER-4bit-MTPLX-Optimized-Speed", + "dawncr0w/Qwen3.6-35B-A3B-Uncensored-HauhauCS-Aggressive-text-oQ8", + "dawncr0w/Qwen3.6-35B-A3B-Uncensored-HauhauCS-Aggressive-text-oQ6", + "samuelfaj/Qwen3.6-35B-A3B-NSC-ACE-SABER-6bit-MTPLX-Optimized-Speed", + "88plug/Qwen3.6-35B-A3B-W8A16", + "jenerallee78/Qwen3.6-35B-A3B-Abliterix-EGA-abliterated" + ], + "relevancy_score": 45.9 + }, + { + "architecture_id": "OLMoForCausalLM", + "total_models": 23, + "total_downloads": 58646, + "min_param_count": 1176764416, + "sample_models": [ + "allenai/OLMo-1B", + "allenai/OLMo-7B", + "allenai/OLMo-7B-Instruct", + "allenai/MolmoE-1B-0924", + "Nhoodie/omni-dna-ici-dc", + "allenai/OLMo-7B-Twin-2T" + ], + "relevancy_score": 45.8 + }, + { + "architecture_id": "CambrianQwenForCausalLM", + "total_models": 11, + "total_downloads": 46529, + "min_param_count": 893618208, + "sample_models": [ + "nyu-visionx/Scale-RAE-Qwen1.5B_DiT2.4B", + "nyu-visionx/Cambrian-S-7B", + "nyu-visionx/Cambrian-S-0.5B" ], - "relevancy_score": 46.1 + "relevancy_score": 45.8 }, { - "architecture_id": "Fgclip2Model", - "total_models": 9, - "total_downloads": 59355, - "min_param_count": 383803394, + "architecture_id": "OpenMythosForCausalLM", + "total_models": 4, + "total_downloads": 112460, + "min_param_count": 95907328, "sample_models": [ - "qihoo360/fg-clip2-base", - "qihoo360/fg-clip2-so400m", - "qihoo360/fg-clip2-large" + "maidacundo/open-mythos-140m" ], "relevancy_score": 45.7 }, { - "architecture_id": "PinyinCodeForCausalLM", - "total_models": 12, - "total_downloads": 35066, - "min_param_count": 33674240, + "architecture_id": "BitNetForCausalLM", + "total_models": 8, + "total_downloads": 62121, + "min_param_count": 849787090, "sample_models": [ - "CPSPX/babylm-zho-pinyin-code-33M", - "timorobrecht/full_chinese_gpu3.1", - "CPSPX/babylm-zho-pinyin-code-97M", - "timorobrecht/full_chinese_gpu3.2-dpo" + "microsoft/bitnet-b1.58-2B-4T", + "microsoft/bitnet-b1.58-2B-4T-bf16" ], - "relevancy_score": 45.5 + "relevancy_score": 45.6 }, { - "architecture_id": "Ernie4_5ForCausalLM", - "total_models": 6, - "total_downloads": 81731, - "min_param_count": 360748032, + "architecture_id": "Kanana2VecModel", + "total_models": 4, + "total_downloads": 738086, + "min_param_count": 2086979328, "sample_models": [ - "baidu/ERNIE-4.5-0.3B-PT", - "baidu/ERNIE-4.5-0.3B-Base-PT" + "kakaocorp/kanana-nano-2.1b-embedding" + ], + "relevancy_score": 45.6 + }, + { + "architecture_id": "Llama4ForCausalLM", + "total_models": 4, + "total_downloads": 105554, + "min_param_count": 3269144, + "sample_models": [ + "trl-internal-testing/tiny-Llama4ForCausalLM" ], "relevancy_score": 45.5 }, @@ -1075,173 +1241,122 @@ "relevancy_score": 45.5 }, { - "architecture_id": "IlamaForCausalLM", - "total_models": 3, - "total_downloads": 817683, - "min_param_count": 1235814400, + "architecture_id": "Exaone4ForCausalLM", + "total_models": 12, + "total_downloads": 220238, + "min_param_count": 1279391488, "sample_models": [ - "hmellor/Ilama-3.2-1B" + "LGAI-EXAONE/EXAONE-4.0-32B", + "LGAI-EXAONE/EXAONE-4.0-1.2B", + "LGAI-EXAONE/EXAONE-4.0.1-32B" ], - "relevancy_score": 45.5 + "relevancy_score": 45.4 }, { - "architecture_id": "OpenMythosForCausalLM", - "total_models": 3, - "total_downloads": 100124, - "min_param_count": 95907328, + "architecture_id": "AfmoeForCausalLM", + "total_models": 24, + "total_downloads": 237334, + "min_param_count": 6120003328, "sample_models": [ - "maidacundo/open-mythos-140m" + "arcee-ai/Trinity-Nano-Preview", + "arcee-ai/Trinity-Mini", + "arcee-ai/Trinity-Large-Thinking", + "arcee-ai/Trinity-Nano-Base", + "dphn/Dolphin-X1-Trinity-Nano", + "arcee-ai/Trinity-Large-Preview" ], "relevancy_score": 45.1 }, { - "architecture_id": "ProphetNetForConditionalGeneration", - "total_models": 2, - "total_downloads": 114477, - "min_param_count": 391321600, + "architecture_id": "HrmTextForCausalLM", + "total_models": 8, + "total_downloads": 334920, + "min_param_count": 1182795264, "sample_models": [ - "microsoft/prophetnet-large-uncased", - "hirotasoshu/tiny-random-prophetnet" + "sapientinc/HRM-Text-1B", + "LLM-OS-Models/KoHRM-Text-1.4B" ], "relevancy_score": 45.1 }, { - "architecture_id": "MultiScaleForCausalLM", - "total_models": 75, - "total_downloads": 55644, - "min_param_count": null, - "sample_models": [ - "KoinicLabs/AXL-Translate", - "KoinicLabs/AXL-Chat-10M", - "KoinicLabs/AXL-Secure-Lion", - "KoinicLabs/AXL-Code-1B", - "KoinicLabs/AXL-Secure-10M", - "KoinicLabs/AXL-Comment-5M", - "KoinicLabs/AXL-Vision-0.8M", - "KoinicLabs/AXL-Vision-v2", - "KoinicLabs/AXL-Chat-Pro", - "KoinicLabs/AXL-Refactor-Lion" - ], - "relevancy_score": 44.9 - }, - { - "architecture_id": "VisionEncoderDecoderModel", - "total_models": 13, - "total_downloads": 21676, - "min_param_count": 239195904, - "sample_models": [ - "AbteeXAILab/lumynax-ocr-trocr-large-printed", - "AbteeXAILab/lumynax-doc-nougat-base", - "AbteeXAILab/lumynax-ocr-trocr-large-handwritten", - "AbteeXAILab/lumynax-doc-donut-base", - "Molkaatb/ChestX" - ], - "relevancy_score": 44.8 - }, - { - "architecture_id": "RemoteForCausalLM", - "total_models": 3, - "total_downloads": 88262, - "min_param_count": 2054056, + "architecture_id": "FGCLIPModel", + "total_models": 4, + "total_downloads": 85495, + "min_param_count": 150136832, "sample_models": [ - "trl-internal-testing/tiny-RemoteForCausalLM" + "qihoo360/fg-clip-base" ], - "relevancy_score": 44.8 + "relevancy_score": 45.1 }, { - "architecture_id": "Llama4ForCausalLM", - "total_models": 3, - "total_downloads": 81342, - "min_param_count": 3269144, + "architecture_id": "ProphetNetForConditionalGeneration", + "total_models": 2, + "total_downloads": 114477, + "min_param_count": 391321600, "sample_models": [ - "trl-internal-testing/tiny-Llama4ForCausalLM" + "microsoft/prophetnet-large-uncased", + "hirotasoshu/tiny-random-prophetnet" ], - "relevancy_score": 44.7 + "relevancy_score": 45.1 }, { - "architecture_id": "Kanana2VecModel", - "total_models": 3, - "total_downloads": 565953, - "min_param_count": 2086979328, + "architecture_id": "HyenaDnaForCausalLM", + "total_models": 16, + "total_downloads": 15469, + "min_param_count": 436352, "sample_models": [ - "kakaocorp/kanana-nano-2.1b-embedding" + "multimolecule/hyenadna-tiny", + "multimolecule/hyenadna-small", + "multimolecule/hyenadna-large", + "multimolecule/hyenadna-medium" ], - "relevancy_score": 44.7 + "relevancy_score": 45.0 }, { - "architecture_id": "_A2DQwen3LMHeadModel", + "architecture_id": "Phi4MMForCausalLM", "total_models": 8, - "total_downloads": 38664, - "min_param_count": 596049920, - "sample_models": [ - "divelab/OPDLM-0.6B", - "divelab/OPDLM-4B", - "divelab/OPDLM-MATH-4B" - ], - "relevancy_score": 44.6 - }, - { - "architecture_id": "GlmForCausalLM", - "total_models": 15, - "total_downloads": 93391, - "min_param_count": 1593427968, + "total_downloads": 2008006, + "min_param_count": 5574460384, "sample_models": [ - "zai-org/glm-4-9b-chat-hf", - "zai-org/glm-4-9b-hf", - "zai-org/glm-edge-4b-chat", - "zai-org/glm-edge-1.5b-chat", - "zai-org/glm-4-9b-chat-1m-hf" + "microsoft/Phi-4-multimodal-instruct", + "Lexius/Phi-4-multimodal-instruct" ], - "relevancy_score": 44.5 + "relevancy_score": 44.9 }, { - "architecture_id": "BitNetForCausalLM", - "total_models": 6, - "total_downloads": 47719, - "min_param_count": 849787090, + "architecture_id": "TTTForCausalLM", + "total_models": 12, + "total_downloads": 20408, + "min_param_count": 733608192, "sample_models": [ - "microsoft/bitnet-b1.58-2B-4T", - "microsoft/bitnet-b1.58-2B-4T-bf16" + "RetentionLabs/TTT-Linear-760M-Base-Pile-8k", + "RetentionLabs/TTT-Linear-1.3B-Base-Pile-8k", + "RetentionLabs/TTT-MLP-760M-Base-Pile-8k" ], "relevancy_score": 44.4 }, { - "architecture_id": "Starcoder2ForCausalLM", - "total_models": 16, - "total_downloads": 496456, - "min_param_count": 3030371328, - "sample_models": [ - "bigcode/starcoder2-3b", - "bigcode/starcoder2-7b", - "bigcode/starcoder2-15b", - "bigcode/starcoder2-15b-instruct-v0.1", - "Johnblick187/starcoder2-15b", - "TechxGenus/starcoder2-7b-GPTQ" - ], - "relevancy_score": 44.3 - }, - { - "architecture_id": "CambrianQwenForCausalLM", - "total_models": 9, - "total_downloads": 29882, - "min_param_count": 893618208, + "architecture_id": "DbrxForCausalLM", + "total_models": 4, + "total_downloads": 58347, + "min_param_count": 4876152, "sample_models": [ - "nyu-visionx/Scale-RAE-Qwen1.5B_DiT2.4B", - "nyu-visionx/Cambrian-S-7B", - "nyu-visionx/Cambrian-S-0.5B" + "trl-internal-testing/tiny-DbrxForCausalLM" ], "relevancy_score": 44.3 }, { - "architecture_id": "HrmTextForCausalLM", - "total_models": 6, - "total_downloads": 297201, - "min_param_count": 1182795264, + "architecture_id": "Ernie4_5_MoeForCausalLM", + "total_models": 16, + "total_downloads": 432589, + "min_param_count": 4532627992, "sample_models": [ - "sapientinc/HRM-Text-1B", - "LLM-OS-Models/KoHRM-Text-1.4B" + "baidu/ERNIE-4.5-21B-A3B-PT", + "cyankiwi/ERNIE-4.5-21B-A3B-Thinking-AWQ-4bit", + "baidu/ERNIE-4.5-21B-A3B-Thinking", + "baidu/ERNIE-4.5-300B-A47B-PT" ], - "relevancy_score": 44.3 + "relevancy_score": 44.0 }, { "architecture_id": "LongT5ForConditionalGeneration", @@ -1257,31 +1372,55 @@ "relevancy_score": 44.0 }, { - "architecture_id": "SDARForCausalLM", + "architecture_id": "LLaDAModelLM", "total_models": 16, - "total_downloads": 63830, - "min_param_count": 2031739904, + "total_downloads": 2710535, + "min_param_count": 8015581184, "sample_models": [ - "JetLM/SDAR-8B-Chat", - "JetLM/SDAR-1.7B-Chat", - "JetLM/SDAR-8B-Chat-b32", - "yifanyu/I-DLM-8B", - "JetLM/SDAR-1.7B-Chat-b32", - "JetLM/SDAR-8B-Chat-b64" + "GSAI-ML/LLaDA-8B-Instruct", + "GSAI-ML/LLaDA-1.5", + "GSAI-ML/LLaDA-8B-Base", + "Zigeng/dParallel-LLaDA-8B-instruct" ], "relevancy_score": 43.9 }, { - "architecture_id": "Exaone4ForCausalLM", - "total_models": 9, - "total_downloads": 163194, - "min_param_count": 1279391488, + "architecture_id": "GPT", + "total_models": 12, + "total_downloads": 15138, + "min_param_count": 14259072, "sample_models": [ - "LGAI-EXAONE/EXAONE-4.0-32B", - "LGAI-EXAONE/EXAONE-4.0-1.2B", - "LGAI-EXAONE/EXAONE-4.0.1-32B" + "younissk/nanoBeard-sloop-14M", + "Imperius/mini-tron-50", + "ncncomplete/nanogpt-tinystories" ], - "relevancy_score": 43.9 + "relevancy_score": 43.8 + }, + { + "architecture_id": "Florence2ForConditionalGeneration", + "total_models": 11, + "total_downloads": 16883, + "min_param_count": 3549945, + "sample_models": [ + "onnx-community/Florence-2-base-ft", + "onnx-community/Florence-2-large-ft", + "Xenova/tiny-random-Florence2ForConditionalGeneration" + ], + "relevancy_score": 43.7 + }, + { + "architecture_id": "MiMoForCausalLM", + "total_models": 19, + "total_downloads": 1586785, + "min_param_count": 7833409536, + "sample_models": [ + "XiaomiMiMo/MiMo-7B-Base", + "XiaomiMiMo/MiMo-7B-RL", + "XiaomiMiMo/MiMo-7B-SFT", + "XiaomiMiMo/MiMo-7B-RL-Zero", + "XiaomiMiMo/MiMo-7B-RL-0530" + ], + "relevancy_score": 43.6 }, { "architecture_id": "PldrllmForCausalLM", @@ -1298,25 +1437,37 @@ "relevancy_score": 43.6 }, { - "architecture_id": "Phi4MMForCausalLM", - "total_models": 6, - "total_downloads": 1478049, - "min_param_count": 5574460384, + "architecture_id": "Plamo2ForCausalLM", + "total_models": 8, + "total_downloads": 161318, + "min_param_count": 1291441920, + "sample_models": [ + "pfnet/plamo-2-1b", + "pfnet/plamo-2-translate" + ], + "relevancy_score": 43.6 + }, + { + "architecture_id": "GeluGPTForCausalLM", + "total_models": 12, + "total_downloads": 12434, + "min_param_count": 261096338, "sample_models": [ - "microsoft/Phi-4-multimodal-instruct", - "Lexius/Phi-4-multimodal-instruct" + "mlnomad/gelu-d12-chinchilla-261M-seed2-pytorch", + "mlnomad/gelu-d12-chinchilla-261M-seed1-pytorch", + "mlnomad/gelu-d12-chinchilla-261M-pytorch" ], - "relevancy_score": 43.6 + "relevancy_score": 43.4 }, { - "architecture_id": "FGCLIPModel", - "total_models": 3, - "total_downloads": 49451, - "min_param_count": 150136832, + "architecture_id": "BD3LM", + "total_models": 4, + "total_downloads": 39185, + "min_param_count": 169627250, "sample_models": [ - "qihoo360/fg-clip-base" + "kuleshov-group/bd3lm-owt-block_size4" ], - "relevancy_score": 43.6 + "relevancy_score": 43.4 }, { "architecture_id": "T5Gemma2ForConditionalGeneration", @@ -1331,27 +1482,31 @@ "relevancy_score": 43.4 }, { - "architecture_id": "HyenaDnaForCausalLM", - "total_models": 12, - "total_downloads": 12116, - "min_param_count": 436352, + "architecture_id": "AraGPT2LMHeadModel", + "total_models": 11, + "total_downloads": 13612, + "min_param_count": 829369856, "sample_models": [ - "multimolecule/hyenadna-tiny", - "multimolecule/hyenadna-small", - "multimolecule/hyenadna-large", - "multimolecule/hyenadna-medium" + "aubmindlab/aragpt2-mega", + "aubmindlab/aragpt2-large", + "QCRI/Fanar-2-Diwan" ], "relevancy_score": 43.3 }, { - "architecture_id": "DbrxForCausalLM", - "total_models": 3, - "total_downloads": 43099, - "min_param_count": 4876152, + "architecture_id": "Glm4ForCausalLM", + "total_models": 23, + "total_downloads": 765649, + "min_param_count": 9400279040, "sample_models": [ - "trl-internal-testing/tiny-DbrxForCausalLM" + "mratsim/GLM-4-32B-0414.w4a16-gptq", + "zai-org/GLM-4-9B-0414", + "zai-org/GLM-4-32B-0414", + "zai-org/GLM-Z1-9B-0414", + "zai-org/GLM-Z1-32B-0414", + "unsloth/GLM-4-9B-0414-unsloth-bnb-4bit" ], - "relevancy_score": 43.3 + "relevancy_score": 43.2 }, { "architecture_id": "LEDForConditionalGeneration", @@ -1369,345 +1524,264 @@ "relevancy_score": 43.2 }, { - "architecture_id": "TTTForCausalLM", - "total_models": 9, - "total_downloads": 15207, - "min_param_count": 733608192, - "sample_models": [ - "RetentionLabs/TTT-Linear-760M-Base-Pile-8k", - "RetentionLabs/TTT-Linear-1.3B-Base-Pile-8k", - "RetentionLabs/TTT-MLP-760M-Base-Pile-8k" - ], - "relevancy_score": 42.9 - }, - { - "architecture_id": "AfmoeForCausalLM", - "total_models": 18, - "total_downloads": 175629, - "min_param_count": 6120003328, - "sample_models": [ - "arcee-ai/Trinity-Nano-Preview", - "arcee-ai/Trinity-Mini", - "arcee-ai/Trinity-Large-Thinking", - "arcee-ai/Trinity-Nano-Base", - "dphn/Dolphin-X1-Trinity-Nano", - "arcee-ai/Trinity-Large-Preview" - ], - "relevancy_score": 42.7 - }, - { - "architecture_id": "Qwen3_5MoeForCausalLM", + "architecture_id": "Mistral3ForConditionalGeneration", "total_models": 28, - "total_downloads": 42695, - "min_param_count": 5642996864, + "total_downloads": 51371, + "min_param_count": 3849090048, "sample_models": [ - "dawncr0w/Qwen3.6-35B-A3B-Uncensored-HauhauCS-Aggressive-text-oQ4", - "samuelfaj/Qwen3.6-35B-A3B-NSC-ACE-SABER-8bit-MTPLX-Optimized-Speed", - "FINAL-Bench/Darwin-36B-Opus", - "0xSero/Qwen3.6-28B", - "samuelfaj/Qwen3.6-35B-A3B-NSC-ACE-SABER-4bit-MTPLX-Optimized-Speed", - "dawncr0w/Qwen3.6-35B-A3B-Uncensored-HauhauCS-Aggressive-text-oQ8", - "dawncr0w/Qwen3.6-35B-A3B-Uncensored-HauhauCS-Aggressive-text-oQ6", - "samuelfaj/Qwen3.6-35B-A3B-NSC-ACE-SABER-6bit-MTPLX-Optimized-Speed", - "88plug/Qwen3.6-35B-A3B-W8A16", - "jenerallee78/Qwen3.6-35B-A3B-Abliterix-EGA-abliterated" + "ArmGPT/ArmenianGPT-1.0-3B", + "Surpem/Supertron2-24B", + "rdtand/Mistral-Medium-3.5-128B-PrismaQuant-4.75-vllm", + "JANGQ-AI/Mistral-Small-4-119B-A6B-JANG_2L", + "dgomes03/Ministral-3-14B-Instruct-2512-mixed-6-4-bit", + "aitf-komdigi/KomdigiITS-8B-DFK-TextClassification", + "TyKaoz/Mistral-Small-3.2-24B-Instruct-2506-4bit", + "squ11z1/LeChatonFat" ], - "relevancy_score": 42.6 + "relevancy_score": 43.0 }, { - "architecture_id": "OLMoForCausalLM", - "total_models": 18, - "total_downloads": 24900, - "min_param_count": 1176764416, + "architecture_id": "A2DQwen3LMHeadModel", + "total_models": 8, + "total_downloads": 18129, + "min_param_count": 751632384, "sample_models": [ - "allenai/OLMo-1B", - "allenai/OLMo-7B", - "allenai/OLMo-7B-Instruct", - "allenai/MolmoE-1B-0924", - "Nhoodie/omni-dna-ici-dc", - "allenai/OLMo-7B-Twin-2T" + "dllm-hub/Qwen3-0.6B-diffusion-bd3lm-v0.1", + "dllm-hub/Qwen3-0.6B-diffusion-mdlm-v0.1" ], - "relevancy_score": 42.6 + "relevancy_score": 43.0 }, { - "architecture_id": "GPT", - "total_models": 9, - "total_downloads": 13079, - "min_param_count": 14259072, + "architecture_id": "GLAForCausalLM", + "total_models": 10, + "total_downloads": 13317, + "min_param_count": 341707776, "sample_models": [ - "younissk/nanoBeard-sloop-14M", - "Imperius/mini-tron-50", - "ncncomplete/nanogpt-tinystories" + "fla-hub/gla-1.3B-100B", + "fla-hub/gla-340M-15B", + "fla-hub/gla-2.7B-100B" ], - "relevancy_score": 42.6 + "relevancy_score": 42.9 }, { - "architecture_id": "BD3LM", - "total_models": 3, - "total_downloads": 29349, - "min_param_count": 169627250, + "architecture_id": "HGRNBitForCausalLM", + "total_models": 8, + "total_downloads": 17154, + "min_param_count": 374108160, "sample_models": [ - "kuleshov-group/bd3lm-owt-block_size4" + "ridger/MMfreeLM-370M", + "ridger/MMfreeLM-2.7B" ], - "relevancy_score": 42.5 + "relevancy_score": 42.9 }, { - "architecture_id": "AraGPT2LMHeadModel", - "total_models": 9, - "total_downloads": 11813, - "min_param_count": 829369856, + "architecture_id": "Ovis2_5", + "total_models": 8, + "total_downloads": 118188, + "min_param_count": 2570429672, "sample_models": [ - "aubmindlab/aragpt2-mega", - "aubmindlab/aragpt2-large", - "QCRI/Fanar-2-Diwan" + "AIDC-AI/Ovis2.5-9B", + "AIDC-AI/Ovis2.5-2B", + "ATH-MaaS/Ovis2.5-9B", + "ATH-MaaS/Ovis2.5-2B" ], - "relevancy_score": 42.4 + "relevancy_score": 42.9 }, { - "architecture_id": "Plamo2ForCausalLM", - "total_models": 6, - "total_downloads": 120185, - "min_param_count": 1291441920, + "architecture_id": "Plamo3ForCausalLM", + "total_models": 7, + "total_downloads": 136743, + "min_param_count": 2603344384, "sample_models": [ - "pfnet/plamo-2-1b", - "pfnet/plamo-2-translate" + "pfnet/plamo-3-nict-2b-base", + "pfnet/plamo-3-nict-8b-base" ], - "relevancy_score": 42.4 + "relevancy_score": 42.9 }, { - "architecture_id": "LLaDAModelLM", + "architecture_id": "Qwen3AttnResForCausalLM", "total_models": 12, - "total_downloads": 2132115, - "min_param_count": 8015581184, + "total_downloads": 9743, + "min_param_count": 115578880, "sample_models": [ - "GSAI-ML/LLaDA-8B-Instruct", - "GSAI-ML/LLaDA-1.5", - "GSAI-ML/LLaDA-8B-Base", - "Zigeng/dParallel-LLaDA-8B-instruct" + "Ethangou/attention-residuals-100M-full", + "Ethangou/attention-residuals-0.6B-block", + "Ethangou/attention-residuals-100M-block" ], - "relevancy_score": 42.2 + "relevancy_score": 42.8 }, { - "architecture_id": "Ernie4_5_MoeForCausalLM", + "architecture_id": "ColMaskMoELLaVAQwen2ForCausalLM", "total_models": 12, - "total_downloads": 326101, - "min_param_count": 4532627992, - "sample_models": [ - "baidu/ERNIE-4.5-21B-A3B-PT", - "cyankiwi/ERNIE-4.5-21B-A3B-Thinking-AWQ-4bit", - "baidu/ERNIE-4.5-21B-A3B-Thinking", - "baidu/ERNIE-4.5-300B-A47B-PT" - ], - "relevancy_score": 42.2 - }, - { - "architecture_id": "Qwen3ASRForConditionalGeneration", - "total_models": 9, - "total_downloads": 10695, - "min_param_count": 782426112, + "total_downloads": 8536, + "min_param_count": 935673472, "sample_models": [ - "neosophie/Qwen3-ASR-1.7B-JA", - "bezzam/Qwen3-ASR-0.6B", - "NbAiLab/nb-asr-beta-qwen06b-lunde05", - "bezzam/Qwen3-ASR-0.6B-hf" + "KKHYA/llavaqwen2.5-0.5b-finetune-col-mask-moe-sparse-4e-2k-sp0.9_20260418_062427", + "KKHYA/llavaqwen2.5-0.5b-finetune-col-mask-moe-sparse-4e-2k-sp0.9_20260418_062617", + "KKHYA/llavaqwen2.5-0.5b-finetune-col-mask-moe-sparse-4e-2k-sp0.7_20260416_182923" ], - "relevancy_score": 42.1 + "relevancy_score": 42.6 }, { - "architecture_id": "GeluGPTForCausalLM", - "total_models": 9, - "total_downloads": 10541, - "min_param_count": 261096338, + "architecture_id": "MarlinForConditionalGeneration", + "total_models": 8, + "total_downloads": 97867, + "min_param_count": 2213241664, "sample_models": [ - "mlnomad/gelu-d12-chinchilla-261M-seed2-pytorch", - "mlnomad/gelu-d12-chinchilla-261M-seed1-pytorch", - "mlnomad/gelu-d12-chinchilla-261M-pytorch" + "NemoStation/Marlin-2B", + "lunahr/Marlin-2B-ungated" ], - "relevancy_score": 42.1 + "relevancy_score": 42.5 }, { "architecture_id": "HybridEchoForCausalLM", - "total_models": 4, - "total_downloads": 20829, + "total_models": 5, + "total_downloads": 21539, "min_param_count": 672190080, "sample_models": [ "mrs83/Kurtis-EON1-Hybrid-0.7B-v0.1.1", "mrs83/Kurtis-EON1-Hybrid-2B-v0.1.2" ], - "relevancy_score": 42.1 - }, - { - "architecture_id": "Plamo3ForCausalLM", - "total_models": 6, - "total_downloads": 102937, - "min_param_count": 2603344384, - "sample_models": [ - "pfnet/plamo-3-nict-2b-base", - "pfnet/plamo-3-nict-8b-base" - ], - "relevancy_score": 42.0 - }, - { - "architecture_id": "A2DQwen3LMHeadModel", - "total_models": 6, - "total_downloads": 14767, - "min_param_count": 751632384, - "sample_models": [ - "dllm-hub/Qwen3-0.6B-diffusion-bd3lm-v0.1", - "dllm-hub/Qwen3-0.6B-diffusion-mdlm-v0.1" - ], - "relevancy_score": 41.9 - }, - { - "architecture_id": "Ovis2_5", - "total_models": 6, - "total_downloads": 87234, - "min_param_count": 2570429672, - "sample_models": [ - "AIDC-AI/Ovis2.5-9B", - "AIDC-AI/Ovis2.5-2B" - ], - "relevancy_score": 41.7 + "relevancy_score": 42.5 }, { - "architecture_id": "MiMoForCausalLM", - "total_models": 15, - "total_downloads": 1061578, - "min_param_count": 7833409536, - "sample_models": [ - "XiaomiMiMo/MiMo-7B-Base", - "XiaomiMiMo/MiMo-7B-RL", - "XiaomiMiMo/MiMo-7B-SFT", - "XiaomiMiMo/MiMo-7B-RL-Zero", - "XiaomiMiMo/MiMo-7B-RL-0530" + "architecture_id": "FalconOCRForCausalLM", + "total_models": 4, + "total_downloads": 23305, + "min_param_count": 269944416, + "sample_models": [ + "tiiuae/Falcon-OCR" ], - "relevancy_score": 41.6 + "relevancy_score": 42.3 }, { - "architecture_id": "Qwen3AttnResForCausalLM", + "architecture_id": "Qwen3ASRForConditionalGeneration", "total_models": 9, - "total_downloads": 8065, - "min_param_count": 115578880, + "total_downloads": 10695, + "min_param_count": 782426112, "sample_models": [ - "Ethangou/attention-residuals-100M-full", - "Ethangou/attention-residuals-0.6B-block", - "Ethangou/attention-residuals-100M-block" + "neosophie/Qwen3-ASR-1.7B-JA", + "bezzam/Qwen3-ASR-0.6B", + "NbAiLab/nb-asr-beta-qwen06b-lunde05", + "bezzam/Qwen3-ASR-0.6B-hf" ], - "relevancy_score": 41.6 + "relevancy_score": 42.2 }, { - "architecture_id": "HGRNBitForCausalLM", - "total_models": 6, - "total_downloads": 12380, - "min_param_count": 374108160, + "architecture_id": "axiomForCausalLM", + "total_models": 8, + "total_downloads": 12388, + "min_param_count": 488664064, "sample_models": [ - "ridger/MMfreeLM-370M", - "ridger/MMfreeLM-2.7B" + "user-anto/Axiom-Dense-380M-Base", + "user-anto/Axiom-Dense-380M-Instruct" ], - "relevancy_score": 41.6 + "relevancy_score": 42.2 }, { - "architecture_id": "FalconOCRForCausalLM", - "total_models": 3, - "total_downloads": 18108, - "min_param_count": 269944416, + "architecture_id": "FalconPerceptionForSegmentation", + "total_models": 4, + "total_downloads": 21975, + "min_param_count": 632372288, "sample_models": [ - "tiiuae/Falcon-OCR" + "tiiuae/Falcon-Perception" ], - "relevancy_score": 41.5 + "relevancy_score": 42.2 }, { - "architecture_id": "GLAForCausalLM", - "total_models": 7, - "total_downloads": 9898, - "min_param_count": 341707776, + "architecture_id": "BharatGPTForCausalLM", + "total_models": 8, + "total_downloads": 11759, + "min_param_count": 353748992, "sample_models": [ - "fla-hub/gla-1.3B-100B", - "fla-hub/gla-340M-15B", - "fla-hub/gla-2.7B-100B" + "Harshhvm/bharat-minigpt-350m-pretrain-3b-tokens", + "anshullpal/ByteZira-Vaani-350M-pretrain-base-model" ], - "relevancy_score": 41.4 + "relevancy_score": 42.1 }, { - "architecture_id": "MarlinForConditionalGeneration", - "total_models": 6, - "total_downloads": 77941, - "min_param_count": 2213241664, + "architecture_id": "LanceAI", + "total_models": 8, + "total_downloads": 10551, + "min_param_count": 393241088, "sample_models": [ - "NemoStation/Marlin-2B", - "lunahr/Marlin-2B-ungated" + "NeuraCraft/Lance-AI-V2", + "NeuraCraft/Lance-AI" ], - "relevancy_score": 41.4 + "relevancy_score": 41.8 }, { - "architecture_id": "ColMaskMoELLaVAQwen2ForCausalLM", - "total_models": 9, - "total_downloads": 6908, - "min_param_count": 935673472, + "architecture_id": "SarvamMoEForCausalLM", + "total_models": 11, + "total_downloads": 224166, + "min_param_count": 6815594732, "sample_models": [ - "KKHYA/llavaqwen2.5-0.5b-finetune-col-mask-moe-sparse-4e-2k-sp0.9_20260418_062427", - "KKHYA/llavaqwen2.5-0.5b-finetune-col-mask-moe-sparse-4e-2k-sp0.9_20260418_062617", - "KKHYA/llavaqwen2.5-0.5b-finetune-col-mask-moe-sparse-4e-2k-sp0.7_20260416_182923" + "sarvamai/sarvam-30b", + "QuantTrio/sarvam-30b-AWQ", + "Orange/Sarvam-30b-GPTQ-w8a16" ], - "relevancy_score": 41.2 + "relevancy_score": 41.1 }, { - "architecture_id": "axiomForCausalLM", - "total_models": 6, - "total_downloads": 10221, - "min_param_count": 488664064, + "architecture_id": "CALIForCausalLM", + "total_models": 4, + "total_downloads": 13048, + "min_param_count": 123782400, "sample_models": [ - "user-anto/Axiom-Dense-380M-Base", - "user-anto/Axiom-Dense-380M-Instruct" + "Sandroeth/cali-0.1B" ], - "relevancy_score": 41.2 + "relevancy_score": 41.1 }, { - "architecture_id": "BharatGPTForCausalLM", - "total_models": 6, - "total_downloads": 10139, - "min_param_count": 353748992, + "architecture_id": "TransformerForCausalLM", + "total_models": 4, + "total_downloads": 80752, + "min_param_count": 1364297728, "sample_models": [ - "Harshhvm/bharat-minigpt-350m-pretrain-3b-tokens", - "anshullpal/ByteZira-Vaani-350M-pretrain-base-model" + "fla-hub/transformer-1.3B-100B" ], - "relevancy_score": 41.2 + "relevancy_score": 41.0 }, { - "architecture_id": "FalconPerceptionForSegmentation", - "total_models": 3, - "total_downloads": 16085, - "min_param_count": 632372288, + "architecture_id": "TinyMindForCausalLM", + "total_models": 8, + "total_downloads": 6526, + "min_param_count": 17731328, "sample_models": [ - "tiiuae/Falcon-Perception" + "HenrySentinel/tinyMind-SFT", + "HenrySentinel/tinyMind" ], - "relevancy_score": 41.2 + "relevancy_score": 40.8 }, { - "architecture_id": "Glm4ForCausalLM", - "total_models": 17, - "total_downloads": 576358, - "min_param_count": 9400279040, + "architecture_id": "LayerWiseMiniCPMForCausalLM", + "total_models": 4, + "total_downloads": 76865, + "min_param_count": 2724956928, "sample_models": [ - "mratsim/GLM-4-32B-0414.w4a16-gptq", - "zai-org/GLM-4-9B-0414", - "zai-org/GLM-4-32B-0414", - "zai-org/GLM-Z1-9B-0414", - "zai-org/GLM-Z1-32B-0414", - "unsloth/GLM-4-9B-0414-unsloth-bnb-4bit" + "BAAI/bge-reranker-v2-minicpm-layerwise" ], - "relevancy_score": 40.9 + "relevancy_score": 40.8 }, { - "architecture_id": "LanceAI", - "total_models": 6, - "total_downloads": 9030, - "min_param_count": 393241088, + "architecture_id": "HCXVisionForCausalLM", + "total_models": 4, + "total_downloads": 502121, + "min_param_count": 3721243520, "sample_models": [ - "NeuraCraft/Lance-AI-V2", - "NeuraCraft/Lance-AI" + "naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B" + ], + "relevancy_score": 40.8 + }, + { + "architecture_id": "SpatialLMQwenForCausalLM", + "total_models": 4, + "total_downloads": 10988, + "min_param_count": 603511168, + "sample_models": [ + "manycore-research/SpatialLM1.1-Qwen-0.5B" ], - "relevancy_score": 40.9 + "relevancy_score": 40.8 }, { "architecture_id": "GlmAsrForConditionalGeneration", @@ -1736,32 +1810,96 @@ "relevancy_score": 40.6 }, { - "architecture_id": "CALIForCausalLM", - "total_models": 3, - "total_downloads": 11266, - "min_param_count": 123782400, + "architecture_id": "Moondream", + "total_models": 4, + "total_downloads": 67776, + "min_param_count": 1857482608, "sample_models": [ - "Sandroeth/cali-0.1B" + "vikhyatk/moondream1" + ], + "relevancy_score": 40.6 + }, + { + "architecture_id": "MobilintQwen3ForCausalLM", + "total_models": 4, + "total_downloads": 10296, + "min_param_count": 388956160, + "sample_models": [ + "mobilint/Qwen3-4B" + ], + "relevancy_score": 40.6 + }, + { + "architecture_id": "DeepseekV32ForCausalLM", + "total_models": 20, + "total_downloads": 13965982, + "min_param_count": 685355329792, + "sample_models": [ + "deepseek-ai/DeepSeek-V3.2", + "QuantTrio/DeepSeek-V3.2-AWQ", + "deepseek-ai/DeepSeek-V3.2-Exp", + "deepseek-ai/DeepSeek-V3.2-Speciale", + "deepseek-ai/DeepSeek-Math-V2" ], "relevancy_score": 40.5 }, { - "architecture_id": "Mistral3ForConditionalGeneration", - "total_models": 21, - "total_downloads": 39261, - "min_param_count": 3849090048, + "architecture_id": "TRMTextISMForCausalLM", + "total_models": 7, + "total_downloads": 6571, + "min_param_count": 84312320, "sample_models": [ - "ArmGPT/ArmenianGPT-1.0-3B", - "Surpem/Supertron2-24B", - "rdtand/Mistral-Medium-3.5-128B-PrismaQuant-4.75-vllm", - "JANGQ-AI/Mistral-Small-4-119B-A6B-JANG_2L", - "dgomes03/Ministral-3-14B-Instruct-2512-mixed-6-4-bit", - "aitf-komdigi/KomdigiITS-8B-DFK-TextClassification", - "TyKaoz/Mistral-Small-3.2-24B-Instruct-2506-4bit", - "squ11z1/LeChatonFat" + "summerMC/TRM-textv3.5", + "summerMC/Trm-text-1B" + ], + "relevancy_score": 40.5 + }, + { + "architecture_id": "BitnetForCausalLM", + "total_models": 5, + "total_downloads": 8518, + "min_param_count": 728843904, + "sample_models": [ + "1bitLLM/bitnet_b1_58-large", + "1bitLLM/bitnet_b1_58-3B" + ], + "relevancy_score": 40.5 + }, + { + "architecture_id": "RecurrentGemmaForCausalLM", + "total_models": 10, + "total_downloads": 26908, + "min_param_count": 2682862080, + "sample_models": [ + "google/recurrentgemma-2b", + "google/recurrentgemma-2b-it", + "google/recurrentgemma-9b-it", + "google/recurrentgemma-9b" + ], + "relevancy_score": 40.4 + }, + { + "architecture_id": "DUO", + "total_models": 7, + "total_downloads": 6049, + "min_param_count": 167928865, + "sample_models": [ + "kekchpek/idlm-duo-tynigsm", + "s-sahoo/duo" ], "relevancy_score": 40.4 }, + { + "architecture_id": "Gemma4ForCausalLM", + "total_models": 7, + "total_downloads": 5918, + "min_param_count": 4496390, + "sample_models": [ + "veyra-ai/Kairo-5M-Gemma4-Base", + "zentropi-ai/cope-b-a4b" + ], + "relevancy_score": 40.3 + }, { "architecture_id": "EchoForCausalLM", "total_models": 3, @@ -1773,240 +1911,269 @@ "relevancy_score": 40.3 }, { - "architecture_id": "TransformerForCausalLM", - "total_models": 3, - "total_downloads": 59830, - "min_param_count": 1364297728, + "architecture_id": "DreamModel", + "total_models": 16, + "total_downloads": 482825, + "min_param_count": 7615616512, "sample_models": [ - "fla-hub/transformer-1.3B-100B" + "Dream-org/Dream-v0-Instruct-7B", + "Dream-org/Dream-v0-Base-7B", + "Dream-org/Dream-Coder-v0-Instruct-7B", + "neurovsa/Dream-Coder-v0-Instruct-7B-HM-Model" + ], + "relevancy_score": 40.2 + }, + { + "architecture_id": "Qwen3TTSForConditionalGeneration", + "total_models": 4, + "total_downloads": 8121, + "min_param_count": 914643008, + "sample_models": [ + "g-group-ai-lab/gwen-tts-0.6B" + ], + "relevancy_score": 40.1 + }, + { + "architecture_id": "SwarmForCausalLM", + "total_models": 4, + "total_downloads": 7686, + "min_param_count": 52729731, + "sample_models": [ + "reaperdoesntknow/SAGI" ], "relevancy_score": 40.0 }, { - "architecture_id": "SarvamMoEForCausalLM", - "total_models": 9, - "total_downloads": 168805, - "min_param_count": 6815594732, + "architecture_id": "DeepseekVLV2ForCausalLM", + "total_models": 4, + "total_downloads": 344873, + "min_param_count": 3370501440, "sample_models": [ - "sarvamai/sarvam-30b", - "QuantTrio/sarvam-30b-AWQ", - "Orange/Sarvam-30b-GPTQ-w8a16" + "Isotr0py/deepseek-vl2-tiny" ], - "relevancy_score": 39.9 + "relevancy_score": 40.0 }, { - "architecture_id": "LayerWiseMiniCPMForCausalLM", - "total_models": 3, - "total_downloads": 56751, - "min_param_count": 2724956928, + "architecture_id": "HiggsMultimodalQwen3ForConditionalGeneration", + "total_models": 4, + "total_downloads": 328453, + "min_param_count": 4654850537, "sample_models": [ - "BAAI/bge-reranker-v2-minicpm-layerwise" + "bosonai/higgs-audio-v3-tts-4b", + "bosonai/higgs-tts-v3-4b", + "bosonai/higgs-tts-3-4b" ], "relevancy_score": 39.9 }, { - "architecture_id": "SpatialLMQwenForCausalLM", + "architecture_id": "InductionForCausalLM", "total_models": 3, - "total_downloads": 8397, - "min_param_count": 603511168, + "total_downloads": 8595, + "min_param_count": 12288816, "sample_models": [ - "manycore-research/SpatialLM1.1-Qwen-0.5B" + "Lllllmmmmmm/conv-induction-babylm-strict-small" ], "relevancy_score": 39.9 }, { - "architecture_id": "TinyMindForCausalLM", + "architecture_id": "SquaredReLUQwen3ForCausalLM", "total_models": 6, - "total_downloads": 5382, - "min_param_count": 17731328, + "total_downloads": 5287, + "min_param_count": 524779008, "sample_models": [ - "HenrySentinel/tinyMind-SFT", - "HenrySentinel/tinyMind" + "Ba2han/qwen3_from_scratch", + "Ba2han/sft-test2" ], "relevancy_score": 39.8 }, { - "architecture_id": "HCXVisionForCausalLM", - "total_models": 3, - "total_downloads": 368595, - "min_param_count": 3721243520, + "architecture_id": "HybridMoRMoEForCausalLM", + "total_models": 4, + "total_downloads": 6963, + "min_param_count": 294194343, "sample_models": [ - "naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B" + "TorchLLM/HybridMoRMoE" ], "relevancy_score": 39.8 }, { - "architecture_id": "Gemma4ForCausalLM", - "total_models": 6, - "total_downloads": 5077, - "min_param_count": 4496390, + "architecture_id": "VexionLMForCausalLM", + "total_models": 4, + "total_downloads": 7140, + "min_param_count": 317774592, "sample_models": [ - "veyra-ai/Kairo-5M-Gemma4-Base", - "zentropi-ai/cope-b-a4b" + "DZER-Studios/Vexion-LM" ], - "relevancy_score": 39.7 + "relevancy_score": 39.8 }, { - "architecture_id": "BitnetForCausalLM", + "architecture_id": "GPTX3ForCausalLM", "total_models": 4, - "total_downloads": 6658, - "min_param_count": 728843904, + "total_downloads": 6919, + "min_param_count": 5158464, "sample_models": [ - "1bitLLM/bitnet_b1_58-large", - "1bitLLM/bitnet_b1_58-3B" + "AxiomicLabs/GPT-S-5M" ], - "relevancy_score": 39.7 + "relevancy_score": 39.8 }, { - "architecture_id": "Moondream", - "total_models": 3, - "total_downloads": 50863, - "min_param_count": 1857482608, + "architecture_id": "Llama4ForConditionalGeneration", + "total_models": 4, + "total_downloads": 6721, + "min_param_count": 6686880, "sample_models": [ - "vikhyatk/moondream1" + "yujiepan/llama-4-tiny-random" ], "relevancy_score": 39.7 }, { - "architecture_id": "MobilintQwen3ForCausalLM", - "total_models": 3, - "total_downloads": 7457, - "min_param_count": 388956160, + "architecture_id": "MetisMoRLMHeadModel", + "total_models": 4, + "total_downloads": 6353, + "min_param_count": 503772163, "sample_models": [ - "mobilint/Qwen3-4B" + "Lernex/Metis-1.4-base" ], "relevancy_score": 39.6 }, { - "architecture_id": "SwarmForCausalLM", - "total_models": 3, - "total_downloads": 6458, - "min_param_count": 52729731, + "architecture_id": "DebertaForSequenceClassification", + "total_models": 4, + "total_downloads": 6248, + "min_param_count": 139199753, "sample_models": [ - "reaperdoesntknow/SAGI" + "AbteeXAILab/lumynax-guard-text-moderation" ], - "relevancy_score": 39.3 + "relevancy_score": 39.6 }, { - "architecture_id": "T5Model", - "total_models": 1, - "total_downloads": 8315, - "min_param_count": 222903552, + "architecture_id": "TableTransformerForObjectDetection", + "total_models": 4, + "total_downloads": 6307, + "min_param_count": 28818631, "sample_models": [ - "sonoisa/t5-base-japanese" + "AbteeXAILab/lumynax-doc-table-transformer-detection" ], - "relevancy_score": 39.3 + "relevancy_score": 39.6 }, { - "architecture_id": "TRMTextISMForCausalLM", - "total_models": 5, - "total_downloads": 4538, - "min_param_count": 84312320, + "architecture_id": "MellumForCausalLM", + "total_models": 20, + "total_downloads": 198386, + "min_param_count": 12149923072, "sample_models": [ - "summerMC/TRM-textv3.5", - "summerMC/Trm-text-1B" + "JetBrains/Mellum2-12B-A2.5B-Thinking", + "JetBrains/Mellum2-12B-A2.5B-Base", + "JetBrains/Mellum2-12B-A2.5B-Instruct", + "cyankiwi/Mellum2-12B-A2.5B-Thinking-AWQ-INT4", + "JetBrains/Mellum2-12B-A2.5B-Thinking-SFT" ], - "relevancy_score": 39.2 + "relevancy_score": 39.5 }, { - "architecture_id": "SquaredReLUQwen3ForCausalLM", - "total_models": 5, - "total_downloads": 4681, - "min_param_count": 524779008, + "architecture_id": "Videollama3Qwen2ForCausalLM", + "total_models": 8, + "total_downloads": 23670, + "min_param_count": 1959993584, "sample_models": [ - "Ba2han/qwen3_from_scratch", - "Ba2han/sft-test2" + "DAMO-NLP-SG/VideoLLaMA3-7B", + "DAMO-NLP-SG/VideoLLaMA3-2B" ], - "relevancy_score": 39.2 + "relevancy_score": 39.5 }, { - "architecture_id": "Qwen3TTSForConditionalGeneration", - "total_models": 3, - "total_downloads": 6071, - "min_param_count": 914643008, + "architecture_id": "MolformerForCausalLM", + "total_models": 5, + "total_downloads": 5134, + "min_param_count": 46805760, "sample_models": [ - "g-group-ai-lab/gwen-tts-0.6B" + "ibm-research/GP-MoLFormer-Uniq", + "ralyn/NPComposer-v2" ], - "relevancy_score": 39.2 + "relevancy_score": 39.4 }, { - "architecture_id": "VexionLMForCausalLM", - "total_models": 3, - "total_downloads": 6165, - "min_param_count": 317774592, + "architecture_id": "DeltaNetForCausalLM", + "total_models": 8, + "total_downloads": 21528, + "min_param_count": 1365677056, "sample_models": [ - "DZER-Studios/Vexion-LM" + "fla-hub/delta_net-1.3B-8K-100B", + "fla-hub/delta_net-1.3B-100B" ], - "relevancy_score": 39.2 + "relevancy_score": 39.3 }, { - "architecture_id": "Llama4ForConditionalGeneration", - "total_models": 3, - "total_downloads": 5918, - "min_param_count": 6686880, + "architecture_id": "EnigmaForCausalLM", + "total_models": 4, + "total_downloads": 5390, + "min_param_count": 12882304, "sample_models": [ - "yujiepan/llama-4-tiny-random" + "dr-tkxx/Enigma" ], - "relevancy_score": 39.1 + "relevancy_score": 39.3 }, { - "architecture_id": "MetisMoRLMHeadModel", - "total_models": 3, - "total_downloads": 5691, - "min_param_count": 503772163, + "architecture_id": "T5Model", + "total_models": 1, + "total_downloads": 8315, + "min_param_count": 222903552, "sample_models": [ - "Lernex/Metis-1.4-base" + "sonoisa/t5-base-japanese" ], - "relevancy_score": 39.1 + "relevancy_score": 39.3 }, { - "architecture_id": "HybridMoRMoEForCausalLM", - "total_models": 3, - "total_downloads": 5751, - "min_param_count": 294194343, + "architecture_id": "GiddForDiffusionLM", + "total_models": 8, + "total_downloads": 18746, + "min_param_count": 2957629440, "sample_models": [ - "TorchLLM/HybridMoRMoE" + "dvruette/gidd-unif-10b", + "dvruette/gidd-unif-3b" ], - "relevancy_score": 39.1 + "relevancy_score": 39.0 }, { - "architecture_id": "DeepseekVLV2ForCausalLM", - "total_models": 3, - "total_downloads": 252416, - "min_param_count": 3370501440, + "architecture_id": "VrityaForCausalLM", + "total_models": 4, + "total_downloads": 4713, + "min_param_count": 163000320, "sample_models": [ - "Isotr0py/deepseek-vl2-tiny" + "curious-techie/Vritya-Tiny-163M-HF" ], "relevancy_score": 39.0 }, { - "architecture_id": "TableTransformerForObjectDetection", - "total_models": 3, - "total_downloads": 5425, - "min_param_count": 28818631, + "architecture_id": "NamerModel", + "total_models": 4, + "total_downloads": 4847, + "min_param_count": 883369, "sample_models": [ - "AbteeXAILab/lumynax-doc-table-transformer-detection" + "edwinhere/namer" ], "relevancy_score": 39.0 }, { - "architecture_id": "DUO", - "total_models": 5, - "total_downloads": 4073, - "min_param_count": 167928865, + "architecture_id": "SmallLLMForCausalLM", + "total_models": 4, + "total_downloads": 4874, + "min_param_count": 123920640, "sample_models": [ - "kekchpek/idlm-duo-tynigsm", - "s-sahoo/duo" + "snehangshu511/tinystories-gpt2-124M-scratch" ], - "relevancy_score": 38.9 + "relevancy_score": 39.0 }, { - "architecture_id": "DebertaForSequenceClassification", - "total_models": 3, - "total_downloads": 5374, - "min_param_count": 139199753, + "architecture_id": "OlmoHybridForCausalLM", + "total_models": 16, + "total_downloads": 256449, + "min_param_count": 7430870688, "sample_models": [ - "AbteeXAILab/lumynax-guard-text-moderation" + "allenai/Olmo-Hybrid-7B", + "allenai/Olmo-Hybrid-Instruct-SFT-7B", + "allenai/Olmo-Hybrid-Think-SFT-7B", + "allenai/Olmo-Hybrid-Instruct-DPO-7B" ], "relevancy_score": 38.9 }, @@ -2021,104 +2188,55 @@ "relevancy_score": 38.9 }, { - "architecture_id": "GPTX3ForCausalLM", - "total_models": 3, - "total_downloads": 5251, - "min_param_count": 5158464, - "sample_models": [ - "AxiomicLabs/GPT-S-5M" - ], - "relevancy_score": 38.9 - }, - { - "architecture_id": "HiggsMultimodalQwen3ForConditionalGeneration", - "total_models": 3, - "total_downloads": 229970, - "min_param_count": 4654850537, + "architecture_id": "YForCausalLM3", + "total_models": 4, + "total_downloads": 4241, + "min_param_count": 55064320, "sample_models": [ - "bosonai/higgs-audio-v3-tts-4b", - "bosonai/higgs-tts-v3-4b" + "SnifferCaptain/ymodel3-n1" ], "relevancy_score": 38.8 }, { - "architecture_id": "InductionForCausalLM", - "total_models": 2, - "total_downloads": 5588, - "min_param_count": 12288816, - "sample_models": [ - "Lllllmmmmmm/conv-induction-babylm-strict-small" - ], - "relevancy_score": 38.7 - }, - { - "architecture_id": "DreamModel", - "total_models": 12, - "total_downloads": 402170, - "min_param_count": 7615616512, - "sample_models": [ - "Dream-org/Dream-v0-Instruct-7B", - "Dream-org/Dream-v0-Base-7B", - "Dream-org/Dream-Coder-v0-Instruct-7B", - "neurovsa/Dream-Coder-v0-Instruct-7B-HM-Model" - ], - "relevancy_score": 38.6 - }, - { - "architecture_id": "MolformerForCausalLM", + "architecture_id": "VanFastForCausalLM", "total_models": 4, - "total_downloads": 4018, - "min_param_count": 46805760, - "sample_models": [ - "ibm-research/GP-MoLFormer-Uniq", - "ralyn/NPComposer-v2" - ], - "relevancy_score": 38.6 - }, - { - "architecture_id": "DeepseekV32ForCausalLM", - "total_models": 15, - "total_downloads": 10970298, - "min_param_count": 685355329792, + "total_downloads": 4132, + "min_param_count": 376644864, "sample_models": [ - "deepseek-ai/DeepSeek-V3.2", - "QuantTrio/DeepSeek-V3.2-AWQ", - "deepseek-ai/DeepSeek-V3.2-Exp", - "deepseek-ai/DeepSeek-V3.2-Speciale", - "deepseek-ai/DeepSeek-Math-V2" + "summerMC/summerV2" ], - "relevancy_score": 38.5 + "relevancy_score": 38.7 }, { - "architecture_id": "RecurrentGemmaForCausalLM", - "total_models": 6, - "total_downloads": 18182, - "min_param_count": 2682862080, + "architecture_id": "Qwen3VLForConditionalGeneration", + "total_models": 8, + "total_downloads": 97882, + "min_param_count": 4437815808, "sample_models": [ - "google/recurrentgemma-2b", - "google/recurrentgemma-2b-it" + "Goekdeniz-Guelmez/Josiefied-Qwen3-VL-4B-Instruct-abliterated-beta-v1", + "DreamFast/Qwen3-VL-8B-Heretic-1.3.0" ], - "relevancy_score": 38.4 + "relevancy_score": 38.5 }, { - "architecture_id": "NamerModel", - "total_models": 3, - "total_downloads": 4181, - "min_param_count": 883369, + "architecture_id": "LangFlow", + "total_models": 4, + "total_downloads": 3759, + "min_param_count": 170805361, "sample_models": [ - "edwinhere/namer" + "Continuous-Rivals-Discrete/langflow-owt" ], - "relevancy_score": 38.4 + "relevancy_score": 38.5 }, { - "architecture_id": "SmallLLMForCausalLM", - "total_models": 3, - "total_downloads": 4191, - "min_param_count": 123920640, + "architecture_id": "MobileLLMP1ForCausalLM", + "total_models": 4, + "total_downloads": 24933, + "min_param_count": 1084453120, "sample_models": [ - "snehangshu511/tinystories-gpt2-124M-scratch" + "facebook/MobileLLM-Pro-base" ], - "relevancy_score": 38.4 + "relevancy_score": 38.5 }, { "architecture_id": "YasinForCausalLM", @@ -2131,44 +2249,42 @@ "relevancy_score": 38.4 }, { - "architecture_id": "VrityaForCausalLM", - "total_models": 3, - "total_downloads": 4040, - "min_param_count": 163000320, + "architecture_id": "SLM", + "total_models": 4, + "total_downloads": 3445, + "min_param_count": 50550784, "sample_models": [ - "curious-techie/Vritya-Tiny-163M-HF" + "aman0419/Vitallm-50M" ], "relevancy_score": 38.3 }, { - "architecture_id": "EnigmaForCausalLM", + "architecture_id": "HanziForCausalLM", "total_models": 3, - "total_downloads": 3955, - "min_param_count": 12882304, + "total_downloads": 3868, + "min_param_count": 123535318, "sample_models": [ - "dr-tkxx/Enigma" + "EjZhou/chinese-babylm-2026-hanzi-a3large" ], "relevancy_score": 38.3 }, { - "architecture_id": "GiddForDiffusionLM", - "total_models": 6, - "total_downloads": 16470, - "min_param_count": 2957629440, + "architecture_id": "GoatVVVForCausalLM", + "total_models": 4, + "total_downloads": 3270, + "min_param_count": 88867190, "sample_models": [ - "dvruette/gidd-unif-10b", - "dvruette/gidd-unif-3b" + "mlnomad/goat-vvv-d12-fineweb-5x-pytorch" ], "relevancy_score": 38.2 }, { - "architecture_id": "Videollama3Qwen2ForCausalLM", - "total_models": 6, - "total_downloads": 16994, - "min_param_count": 1959993584, + "architecture_id": "RecursiveLanguageModel", + "total_models": 4, + "total_downloads": 3340, + "min_param_count": 198464806, "sample_models": [ - "DAMO-NLP-SG/VideoLLaMA3-7B", - "DAMO-NLP-SG/VideoLLaMA3-2B" + "Girinath11/recursive-language-model-198m" ], "relevancy_score": 38.2 }, @@ -2183,666 +2299,678 @@ "relevancy_score": 38.2 }, { - "architecture_id": "DeltaNetForCausalLM", - "total_models": 6, - "total_downloads": 15753, - "min_param_count": 1365677056, + "architecture_id": "LLaDA2MoeModelLM", + "total_models": 16, + "total_downloads": 1180876, + "min_param_count": 16255643392, "sample_models": [ - "fla-hub/delta_net-1.3B-8K-100B", - "fla-hub/delta_net-1.3B-100B" + "inclusionAI/LLaDA2.0-mini", + "inclusionAI/LLaDA2.1-flash", + "inclusionAI/LLaDA2.1-mini", + "inclusionAI/LLaDA2.0-mini-preview" ], "relevancy_score": 38.1 }, { - "architecture_id": "GPTX2ForCausalLM", - "total_models": 4, - "total_downloads": 3012, - "min_param_count": 126006398, + "architecture_id": "AV2TextForConditionalGeneration", + "total_models": 5, + "total_downloads": 2738, + "min_param_count": 480465000, "sample_models": [ - "AxiomicLabs/GPT-X2-125M", - "reaperdoesntknow/GPT-X2-125M-CIx-Long-Context" + "nguyenvulebinh/AV-HuBERT-MuAViC-en", + "nguyenvulebinh/AV-HuBERT" + ], + "relevancy_score": 38.1 + }, + { + "architecture_id": "KimiLinearForCausalLM", + "total_models": 12, + "total_downloads": 294022, + "min_param_count": 9373302656, + "sample_models": [ + "moonshotai/Kimi-Linear-48B-A3B-Instruct", + "moonshotai/Kimi-Linear-48B-A3B-Base", + "cyankiwi/Kimi-Linear-48B-A3B-Instruct-AWQ-4bit" ], "relevancy_score": 38.0 }, { - "architecture_id": "VanFastForCausalLM", - "total_models": 3, - "total_downloads": 3511, - "min_param_count": 376644864, + "architecture_id": "Rwkv7ForCausalLM", + "total_models": 4, + "total_downloads": 2912, + "min_param_count": 34158592, "sample_models": [ - "summerMC/summerV2" + "admijgjtjtjtjjg/dfdfdf" ], "relevancy_score": 38.0 }, { - "architecture_id": "YForCausalLM3", - "total_models": 3, - "total_downloads": 3339, - "min_param_count": 55064320, + "architecture_id": "TachBitModel", + "total_models": 4, + "total_downloads": 2919, + "min_param_count": 284873472, "sample_models": [ - "SnifferCaptain/ymodel3-n1" + "tachyonx-ai/TachBit-285" ], - "relevancy_score": 37.9 + "relevancy_score": 38.0 }, { - "architecture_id": "Qwen3VLForConditionalGeneration", - "total_models": 6, - "total_downloads": 93689, - "min_param_count": 4437815808, + "architecture_id": "DogeForCausalLM", + "total_models": 4, + "total_downloads": 2966, + "min_param_count": 13118728, "sample_models": [ - "Goekdeniz-Guelmez/Josiefied-Qwen3-VL-4B-Instruct-abliterated-beta-v1", - "DreamFast/Qwen3-VL-8B-Heretic-1.3.0" + "SmallDoge/Doge-20M" ], - "relevancy_score": 37.8 + "relevancy_score": 38.0 }, { - "architecture_id": "RecursiveLanguageModel", - "total_models": 3, - "total_downloads": 2711, - "min_param_count": 198464806, + "architecture_id": "SmoothieModel", + "total_models": 4, + "total_downloads": 2939, + "min_param_count": 155829504, "sample_models": [ - "Girinath11/recursive-language-model-198m" + "yasserrmd/smoothie-diffusion-qqp" ], - "relevancy_score": 37.5 + "relevancy_score": 38.0 }, { - "architecture_id": "GoatVVVForCausalLM", - "total_models": 3, - "total_downloads": 2705, - "min_param_count": 88867190, + "architecture_id": "ARAr381MForCausalLM", + "total_models": 4, + "total_downloads": 2927, + "min_param_count": 381480960, "sample_models": [ - "mlnomad/goat-vvv-d12-fineweb-5x-pytorch" + "AraFus/AraFusion-AR-381M-v2" ], - "relevancy_score": 37.5 + "relevancy_score": 38.0 }, { - "architecture_id": "MobileLLMP1ForCausalLM", - "total_models": 3, - "total_downloads": 18420, - "min_param_count": 1084453120, + "architecture_id": "GPTX2ForCausalLM", + "total_models": 4, + "total_downloads": 3012, + "min_param_count": 126006398, "sample_models": [ - "facebook/MobileLLM-Pro-base" + "AxiomicLabs/GPT-X2-125M", + "reaperdoesntknow/GPT-X2-125M-CIx-Long-Context" ], - "relevancy_score": 37.5 + "relevancy_score": 38.0 }, { - "architecture_id": "SLM", - "total_models": 3, - "total_downloads": 2764, - "min_param_count": 50550784, + "architecture_id": "YoutuVLForConditionalGeneration", + "total_models": 8, + "total_downloads": 11019, + "min_param_count": 2515984240, "sample_models": [ - "aman0419/Vitallm-50M" + "tencent/Youtu-Parsing", + "tencent/Youtu-VL-4B-Instruct" ], - "relevancy_score": 37.5 + "relevancy_score": 37.9 }, { - "architecture_id": "LangFlow", - "total_models": 3, - "total_downloads": 2688, - "min_param_count": 170805361, + "architecture_id": "HierColMaskMoELLaVAQwen2ForCausalLM", + "total_models": 4, + "total_downloads": 2840, + "min_param_count": 935731840, "sample_models": [ - "Continuous-Rivals-Discrete/langflow-owt" + "KKHYA/llavaqwen2.5-0.5b-finetune-hier-col-mask-moe-sparse-4e-2k-sp0.9_20260418_063208" ], - "relevancy_score": 37.5 + "relevancy_score": 37.9 }, { - "architecture_id": "MellumForCausalLM", - "total_models": 15, - "total_downloads": 139432, - "min_param_count": 12149923072, + "architecture_id": "LatentMoELLaVAQwen2ForCausalLM", + "total_models": 4, + "total_downloads": 2776, + "min_param_count": 935673472, "sample_models": [ - "JetBrains/Mellum2-12B-A2.5B-Thinking", - "JetBrains/Mellum2-12B-A2.5B-Base", - "JetBrains/Mellum2-12B-A2.5B-Instruct", - "cyankiwi/Mellum2-12B-A2.5B-Thinking-AWQ-INT4", - "JetBrains/Mellum2-12B-A2.5B-Thinking-SFT" + "KKHYA/llavaqwen2.5-0.5b-finetune-latent-sparse-moe-4e-1k_20260413_214436" ], - "relevancy_score": 37.3 + "relevancy_score": 37.9 }, { - "architecture_id": "OlmoHybridForCausalLM", - "total_models": 12, - "total_downloads": 209003, - "min_param_count": 7430870688, + "architecture_id": "RubiRLM", + "total_models": 4, + "total_downloads": 2794, + "min_param_count": 988446027, "sample_models": [ - "allenai/Olmo-Hybrid-7B", - "allenai/Olmo-Hybrid-Instruct-SFT-7B", - "allenai/Olmo-Hybrid-Think-SFT-7B", - "allenai/Olmo-Hybrid-Instruct-DPO-7B" + "DevHunterAI/RubiRLM-1B-Base" ], - "relevancy_score": 37.3 + "relevancy_score": 37.9 }, { - "architecture_id": "OpensciForCausalLM", - "total_models": 9, - "total_downloads": 7323, - "min_param_count": 1714377728, + "architecture_id": "KaniTTS2ForCausalLM", + "total_models": 4, + "total_downloads": 2871, + "min_param_count": 369977088, "sample_models": [ - "open-sci/open-sci-ref-v0.02-1.7b-fineweb-edu-1.4t-1T-4096-rope_theta-100k", - "open-sci/open-sci-ref-v0.02-1.7b-dclm-1T-4096-rope_theta-100k", - "open-sci/open-sci-ref-v0.02-1.7b-nemotron-hq-1T-4096-rope_theta-100k" + "nineninesix/kani-tts-2-en" ], - "relevancy_score": 37.3 + "relevancy_score": 37.9 }, { - "architecture_id": "SwitchTransformersForConditionalGeneration", + "architecture_id": "DynHierColMaskMoELLaVAQwen2ForCausalLM", "total_models": 4, - "total_downloads": 14420, - "min_param_count": 1978514688, + "total_downloads": 2822, + "min_param_count": 940155520, "sample_models": [ - "google/switch-base-8", - "google/switch-base-16", - "google/switch-base-128", - "google/switch-base-32" + "KKHYA/llavaqwen2.5-0.5b-finetune-dyn-hier-col-mask-moe-sparse-4e-2k-sp0.9-r16_20260418_063238" ], - "relevancy_score": 37.3 + "relevancy_score": 37.9 }, { - "architecture_id": "AV2TextForConditionalGeneration", + "architecture_id": "GlubLM", "total_models": 4, - "total_downloads": 2174, - "min_param_count": 480465000, + "total_downloads": 2855, + "min_param_count": 36055680, "sample_models": [ - "nguyenvulebinh/AV-HuBERT-MuAViC-en", - "nguyenvulebinh/AV-HuBERT" + "DenSec02/glublm-36m" ], - "relevancy_score": 37.3 + "relevancy_score": 37.9 }, { - "architecture_id": "TachBitModel", - "total_models": 3, - "total_downloads": 2371, - "min_param_count": 284873472, + "architecture_id": "HCXVisionV2ForCausalLM", + "total_models": 7, + "total_downloads": 526296, + "min_param_count": 10741664520, "sample_models": [ - "tachyonx-ai/TachBit-285" + "naver-hyperclovax/HyperCLOVAX-SEED-Think-32B", + "naver-hyperclovax/HyperCLOVAX-SEED-Omni-8B" ], - "relevancy_score": 37.2 + "relevancy_score": 37.8 }, { - "architecture_id": "GlubLM", - "total_models": 3, - "total_downloads": 2308, - "min_param_count": 36055680, + "architecture_id": "MiniMindForCausalLM", + "total_models": 4, + "total_downloads": 2684, + "min_param_count": 63912192, "sample_models": [ - "DenSec02/glublm-36m" + "chujiamo/baiheng_0405" ], - "relevancy_score": 37.2 + "relevancy_score": 37.8 }, { - "architecture_id": "Rwkv7ForCausalLM", - "total_models": 3, - "total_downloads": 2321, - "min_param_count": 34158592, + "architecture_id": "KsByteForCausalLM", + "total_models": 4, + "total_downloads": 2677, + "min_param_count": 15837440, "sample_models": [ - "admijgjtjtjtjjg/dfdfdf" + "Omarrran/ks-byte-lm-spacebyte-transformers" ], - "relevancy_score": 37.2 + "relevancy_score": 37.8 }, { - "architecture_id": "HierColMaskMoELLaVAQwen2ForCausalLM", - "total_models": 3, - "total_downloads": 2300, - "min_param_count": 935731840, + "architecture_id": "HeliumForCausalLM", + "total_models": 4, + "total_downloads": 18326, + "min_param_count": 2172643840, "sample_models": [ - "KKHYA/llavaqwen2.5-0.5b-finetune-hier-col-mask-moe-sparse-4e-2k-sp0.9_20260418_063208" + "kyutai/helium-1-preview-2b" ], - "relevancy_score": 37.2 + "relevancy_score": 37.8 }, { - "architecture_id": "ARAr381MForCausalLM", - "total_models": 3, - "total_downloads": 2373, - "min_param_count": 381480960, + "architecture_id": "HFHealthSLM", + "total_models": 4, + "total_downloads": 2731, + "min_param_count": 23468288, "sample_models": [ - "AraFus/AraFusion-AR-381M-v2" + "rohansingh2612/healthslm" ], - "relevancy_score": 37.2 + "relevancy_score": 37.8 }, { - "architecture_id": "DogeForCausalLM", - "total_models": 3, - "total_downloads": 2291, - "min_param_count": 13118728, + "architecture_id": "FusionInDecoderForConditionalGeneration", + "total_models": 4, + "total_downloads": 2759, + "min_param_count": 247577856, "sample_models": [ - "SmallDoge/Doge-20M" + "Intel/fid_flan_t5_base_nq" ], - "relevancy_score": 37.2 + "relevancy_score": 37.8 }, { - "architecture_id": "PegasusXForConditionalGeneration", - "total_models": 2, - "total_downloads": 2705, - "min_param_count": 272311296, + "architecture_id": "HybridQwen3ForCausalLM", + "total_models": 14, + "total_downloads": 181379, + "min_param_count": 8495712960, "sample_models": [ - "google/pegasus-x-base", - "BEE-spoke-data/pegasus-x-base-synthsumm_open-16k" + "amazon/GKA-primed-HQwen3-32B-Instruct", + "amazon/GKA-primed-HQwen3-8B-Reasoner", + "amazon/GKA-primed-HQwen3-32B-Reasoner", + "amazon/GDN-primed-HQwen3-8B-Instruct", + "amazon/BMOJOF-primed-HQwen3-8B-Instruct", + "amazon/Mamba2-primed-HQwen3-8B-Instruct" ], - "relevancy_score": 37.2 + "relevancy_score": 37.6 }, { - "architecture_id": "RubiRLM", - "total_models": 3, - "total_downloads": 2245, - "min_param_count": 988446027, + "architecture_id": "MoshiForConditionalGeneration", + "total_models": 6, + "total_downloads": 558308, + "min_param_count": 7783880545, "sample_models": [ - "DevHunterAI/RubiRLM-1B-Base" + "kmhf/hf-moshiko", + "kmhf/hf-moshika" ], - "relevancy_score": 37.1 + "relevancy_score": 37.6 }, { - "architecture_id": "LatentMoELLaVAQwen2ForCausalLM", - "total_models": 3, - "total_downloads": 2239, - "min_param_count": 935673472, + "architecture_id": "Qwen3VLMoeForConditionalGeneration", + "total_models": 16, + "total_downloads": 5664951, + "min_param_count": 31070754032, "sample_models": [ - "KKHYA/llavaqwen2.5-0.5b-finetune-latent-sparse-moe-4e-1k_20260413_214436" + "QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ", + "QuantTrio/Qwen3-VL-30B-A3B-Thinking-AWQ", + "QuantTrio/Qwen3-VL-235B-A22B-Instruct-AWQ", + "QuantTrio/Qwen3-VL-235B-A22B-Thinking-AWQ" ], - "relevancy_score": 37.1 + "relevancy_score": 37.4 }, { - "architecture_id": "HFHealthSLM", - "total_models": 3, - "total_downloads": 2192, - "min_param_count": 23468288, + "architecture_id": "OpensciForCausalLM", + "total_models": 9, + "total_downloads": 7323, + "min_param_count": 1714377728, "sample_models": [ - "rohansingh2612/healthslm" + "open-sci/open-sci-ref-v0.02-1.7b-fineweb-edu-1.4t-1T-4096-rope_theta-100k", + "open-sci/open-sci-ref-v0.02-1.7b-dclm-1T-4096-rope_theta-100k", + "open-sci/open-sci-ref-v0.02-1.7b-nemotron-hq-1T-4096-rope_theta-100k" ], - "relevancy_score": 37.1 + "relevancy_score": 37.4 }, { - "architecture_id": "DynHierColMaskMoELLaVAQwen2ForCausalLM", + "architecture_id": "SpikeWhaleLM", "total_models": 3, - "total_downloads": 2285, - "min_param_count": 940155520, + "total_downloads": 2534, + "min_param_count": 97272836, "sample_models": [ - "KKHYA/llavaqwen2.5-0.5b-finetune-dyn-hier-col-mask-moe-sparse-4e-2k-sp0.9-r16_20260418_063238" + "Quazim0t0/Escarda-86M-Base" ], - "relevancy_score": 37.1 + "relevancy_score": 37.4 }, { - "architecture_id": "KaniTTS2ForCausalLM", - "total_models": 3, - "total_downloads": 2170, - "min_param_count": 369977088, + "architecture_id": "RWForCausalLM", + "total_models": 16, + "total_downloads": 17669, + "min_param_count": 6854619456, "sample_models": [ - "nineninesix/kani-tts-2-en" + "h2oai/h2ogpt-gm-oasst1-en-2048-falcon-7b-v2", + "explosion-testing/refined-web-model-test", + "lightonai/alfred-40b-1023", + "projecte-aina/aguila-7b", + "h2oai/h2ogpt-gm-oasst1-en-2048-falcon-7b-v3", + "OpenAssistant/falcon-40b-sft-top1-560", + "nomic-ai/gpt4all-falcon", + "QuixiAI/WizardLM-Uncensored-Falcon-40b" ], - "relevancy_score": 37.0 + "relevancy_score": 37.3 }, { - "architecture_id": "FusionInDecoderForConditionalGeneration", - "total_models": 3, - "total_downloads": 2174, - "min_param_count": 247577856, + "architecture_id": "GraniteSwitchForCausalLM", + "total_models": 12, + "total_downloads": 32079, + "min_param_count": 4149977600, "sample_models": [ - "Intel/fid_flan_t5_base_nq" + "ibm-granite/granite-switch-4.1-3b-preview", + "ibm-granite/granite-switch-4.1-8b-preview", + "ibm-granite/granite-switch-4.1-30b-preview" ], - "relevancy_score": 37.0 + "relevancy_score": 37.3 }, { - "architecture_id": "MiniMindForCausalLM", - "total_models": 3, - "total_downloads": 2146, - "min_param_count": 63912192, + "architecture_id": "SwitchTransformersForConditionalGeneration", + "total_models": 4, + "total_downloads": 14420, + "min_param_count": 1978514688, "sample_models": [ - "chujiamo/baiheng_0405" + "google/switch-base-8", + "google/switch-base-16", + "google/switch-base-128", + "google/switch-base-32" ], - "relevancy_score": 37.0 + "relevancy_score": 37.3 }, { - "architecture_id": "SmoothieModel", + "architecture_id": "Veyra2ApricotForCausalLM", "total_models": 3, - "total_downloads": 2157, - "min_param_count": 155829504, + "total_downloads": 2442, + "min_param_count": 49303040, "sample_models": [ - "yasserrmd/smoothie-diffusion-qqp" + "veyra-ai/Veyra2-Apricot-50M-Base" ], - "relevancy_score": 37.0 + "relevancy_score": 37.3 }, { - "architecture_id": "KsByteForCausalLM", - "total_models": 3, - "total_downloads": 1999, - "min_param_count": 15837440, + "architecture_id": "TarsierForConditionalGeneration", + "total_models": 4, + "total_downloads": 598918, + "min_param_count": 7063427072, "sample_models": [ - "Omarrran/ks-byte-lm-spacebyte-transformers" + "omni-research/Tarsier-7b" ], - "relevancy_score": 36.9 + "relevancy_score": 37.2 }, { - "architecture_id": "HCXVisionV2ForCausalLM", - "total_models": 6, - "total_downloads": 388666, - "min_param_count": 10741664520, + "architecture_id": "CostWiseGemmaForCausalLM", + "total_models": 4, + "total_downloads": 614564, + "min_param_count": 9241831424, "sample_models": [ - "naver-hyperclovax/HyperCLOVAX-SEED-Think-32B", - "naver-hyperclovax/HyperCLOVAX-SEED-Omni-8B" + "BAAI/bge-reranker-v2.5-gemma2-lightweight" ], - "relevancy_score": 36.8 + "relevancy_score": 37.2 }, { - "architecture_id": "QEDForCausalLM", + "architecture_id": "MobilintLlamaForCausalLM", "total_models": 3, - "total_downloads": 1937, - "min_param_count": 94445952, + "total_downloads": 2353, + "min_param_count": 262668288, "sample_models": [ - "levossadtchi/QED-75M" + "mobilint/Llama-3.2-1B-Instruct" ], - "relevancy_score": 36.8 + "relevancy_score": 37.2 }, { - "architecture_id": "HanziForCausalLM", + "architecture_id": "PegasusXForConditionalGeneration", "total_models": 2, - "total_downloads": 2234, - "min_param_count": 123535318, + "total_downloads": 2705, + "min_param_count": 272311296, "sample_models": [ - "EjZhou/chinese-babylm-2026-hanzi-a3large" + "google/pegasus-x-base", + "BEE-spoke-data/pegasus-x-base-synthsumm_open-16k" ], - "relevancy_score": 36.8 + "relevancy_score": 37.2 }, { - "architecture_id": "MoshiForConditionalGeneration", - "total_models": 5, - "total_downloads": 425200, - "min_param_count": 7783880545, + "architecture_id": "VaultGemmaForCausalLM", + "total_models": 4, + "total_downloads": 13207, + "min_param_count": 1038741120, "sample_models": [ - "kmhf/hf-moshiko", - "kmhf/hf-moshika" + "google/vaultgemma-1b" ], - "relevancy_score": 36.7 + "relevancy_score": 37.1 }, { - "architecture_id": "NeuroBLASTForCausalLM", - "total_models": 3, - "total_downloads": 1859, - "min_param_count": 596728320, + "architecture_id": "Rwkv5ForCausalLM", + "total_models": 8, + "total_downloads": 7077, + "min_param_count": 1577754624, "sample_models": [ - "mkurman/NeuroBLAST-V3-SYNTH-EC-150000" + "RWKV/rwkv-5-world-1b5", + "RWKV/rwkv-5-world-3b" ], - "relevancy_score": 36.7 + "relevancy_score": 37.0 }, { - "architecture_id": "YoutuVLForConditionalGeneration", - "total_models": 6, - "total_downloads": 7973, - "min_param_count": 2515984240, + "architecture_id": "BolmoForCausalLM", + "total_models": 4, + "total_downloads": 12333, + "min_param_count": 1468911776, "sample_models": [ - "tencent/Youtu-Parsing", - "tencent/Youtu-VL-4B-Instruct" + "allenai/Bolmo-1B" ], - "relevancy_score": 36.6 + "relevancy_score": 37.0 }, { - "architecture_id": "KimiLinearForCausalLM", - "total_models": 9, - "total_downloads": 221546, - "min_param_count": 9373302656, + "architecture_id": "OrthrusLM", + "total_models": 8, + "total_downloads": 6914, + "min_param_count": 2072903680, "sample_models": [ - "moonshotai/Kimi-Linear-48B-A3B-Instruct", - "moonshotai/Kimi-Linear-48B-A3B-Base", - "cyankiwi/Kimi-Linear-48B-A3B-Instruct-AWQ-4bit" + "chiennv/Orthrus-Qwen3-8B", + "chiennv/Orthrus-Qwen3-1.7B" ], - "relevancy_score": 36.5 + "relevancy_score": 36.9 }, { - "architecture_id": "GomeForCausalLM", + "architecture_id": "FunAsrNanoForConditionalGeneration", "total_models": 3, - "total_downloads": 1703, - "min_param_count": 24761091, + "total_downloads": 1983, + "min_param_count": 985374304, "sample_models": [ - "Prositron/gome" + "FunAudioLLM/Fun-ASR-Nano-2512-hf" ], - "relevancy_score": 36.5 + "relevancy_score": 36.9 }, { - "architecture_id": "RWForCausalLM", - "total_models": 14, - "total_downloads": 15266, - "min_param_count": 6854619456, - "sample_models": [ - "h2oai/h2ogpt-gm-oasst1-en-2048-falcon-7b-v2", - "explosion-testing/refined-web-model-test", - "lightonai/alfred-40b-1023", - "projecte-aina/aguila-7b", - "h2oai/h2ogpt-gm-oasst1-en-2048-falcon-7b-v3", - "OpenAssistant/falcon-40b-sft-top1-560", - "nomic-ai/gpt4all-falcon", - "QuixiAI/WizardLM-Uncensored-Falcon-40b" + "architecture_id": "SLMForCausalLM", + "total_models": 3, + "total_downloads": 1998, + "min_param_count": 12065792, + "sample_models": [ + "liodon-ai/slm-10m" ], - "relevancy_score": 36.4 + "relevancy_score": 36.9 }, { - "architecture_id": "LLaDA2MoeModelLM", - "total_models": 12, - "total_downloads": 922374, - "min_param_count": 16255643392, + "architecture_id": "EshmunForCausalLM", + "total_models": 3, + "total_downloads": 2033, + "min_param_count": 341319680, "sample_models": [ - "inclusionAI/LLaDA2.0-mini", - "inclusionAI/LLaDA2.1-flash", - "inclusionAI/LLaDA2.1-mini", - "inclusionAI/LLaDA2.0-mini-preview" + "khairi/Eshmun-0.3B-CPT" ], - "relevancy_score": 36.4 + "relevancy_score": 36.9 }, { - "architecture_id": "HeliumForCausalLM", - "total_models": 3, - "total_downloads": 10984, - "min_param_count": 2172643840, + "architecture_id": "YoutuForCausalLM", + "total_models": 5, + "total_downloads": 9684, + "min_param_count": 1961560064, "sample_models": [ - "kyutai/helium-1-preview-2b" + "tencent/Youtu-LLM-2B-Base", + "tencent/Youtu-LLM-2B" ], - "relevancy_score": 36.4 + "relevancy_score": 36.8 }, { - "architecture_id": "CostWiseGemmaForCausalLM", + "architecture_id": "QEDForCausalLM", "total_models": 3, - "total_downloads": 457708, - "min_param_count": 9241831424, + "total_downloads": 1937, + "min_param_count": 94445952, "sample_models": [ - "BAAI/bge-reranker-v2.5-gemma2-lightweight" + "levossadtchi/QED-75M" ], - "relevancy_score": 36.3 + "relevancy_score": 36.8 }, { - "architecture_id": "BolmoForCausalLM", - "total_models": 3, - "total_downloads": 9669, - "min_param_count": 1468911776, + "architecture_id": "NanoChatForCausalLM", + "total_models": 5, + "total_downloads": 9328, + "min_param_count": 2217082880, "sample_models": [ - "allenai/Bolmo-1B" + "l0ulan/chinese-babylm-nanochat-d22-step10000", + "Twobombs/nanochat-d34-sft-hf" ], - "relevancy_score": 36.2 + "relevancy_score": 36.7 }, { - "architecture_id": "VaultGemmaForCausalLM", + "architecture_id": "NeuroBLASTForCausalLM", "total_models": 3, - "total_downloads": 9832, - "min_param_count": 1038741120, + "total_downloads": 1859, + "min_param_count": 596728320, "sample_models": [ - "google/vaultgemma-1b" + "mkurman/NeuroBLAST-V3-SYNTH-EC-150000" ], - "relevancy_score": 36.2 + "relevancy_score": 36.7 }, { - "architecture_id": "TarsierForConditionalGeneration", - "total_models": 3, - "total_downloads": 439836, - "min_param_count": 7063427072, + "architecture_id": "MolmoForCausalLM", + "total_models": 15, + "total_downloads": 99586, + "min_param_count": 7665032192, "sample_models": [ - "omni-research/Tarsier-7b" + "allenai/Molmo-7B-D-0924", + "allenai/Molmo-72B-0924", + "allenai/Molmo-7B-O-0924", + "cyan2k/molmo-7B-O-bnb-4bit" ], - "relevancy_score": 36.2 + "relevancy_score": 36.6 }, { - "architecture_id": "MobilintLlamaForCausalLM", - "total_models": 2, - "total_downloads": 1688, - "min_param_count": 262668288, + "architecture_id": "QuarkForCausalLM", + "total_models": 3, + "total_downloads": 1746, + "min_param_count": 302080768, "sample_models": [ - "mobilint/Llama-3.2-1B-Instruct" + "ThingAI/Quark-270m-Base" ], - "relevancy_score": 36.2 + "relevancy_score": 36.6 }, { - "architecture_id": "Veyra2ApricotForCausalLM", - "total_models": 2, - "total_downloads": 1620, - "min_param_count": 49303040, + "architecture_id": "GTLMForCausalLM", + "total_models": 8, + "total_downloads": 5514, + "min_param_count": 2095989760, "sample_models": [ - "veyra-ai/Veyra2-Apricot-50M-Base" + "Madras1/GTLM-1-2B-A350M", + "Madras1/GTLM-1-2B-A350M-fp16" ], - "relevancy_score": 36.1 + "relevancy_score": 36.5 }, { - "architecture_id": "SpikeWhaleLM", - "total_models": 2, - "total_downloads": 1614, - "min_param_count": 97272836, + "architecture_id": "GomeForCausalLM", + "total_models": 3, + "total_downloads": 1703, + "min_param_count": 24761091, "sample_models": [ - "Quazim0t0/Escarda-86M-Base" + "Prositron/gome" ], - "relevancy_score": 36.1 + "relevancy_score": 36.5 }, { - "architecture_id": "GraniteSwitchForCausalLM", - "total_models": 9, - "total_downloads": 25905, - "min_param_count": 4149977600, + "architecture_id": "XLNetLMHeadModel", + "total_models": 24, + "total_downloads": 1171367, + "min_param_count": null, "sample_models": [ - "ibm-granite/granite-switch-4.1-3b-preview", - "ibm-granite/granite-switch-4.1-8b-preview", - "ibm-granite/granite-switch-4.1-30b-preview" + "xlnet/xlnet-base-cased", + "xlnet/xlnet-large-cased", + "hfl/chinese-xlnet-base", + "sshleifer/tiny-xlnet-base-cased", + "textattack/xlnet-base-cased-imdb", + "textattack/xlnet-base-cased-CoLA" ], - "relevancy_score": 36.0 + "relevancy_score": 36.4 }, { - "architecture_id": "Rwkv5ForCausalLM", - "total_models": 6, - "total_downloads": 5719, - "min_param_count": 1577754624, + "architecture_id": "HawkForCausalLM", + "total_models": 1, + "total_downloads": 1834, + "min_param_count": 115270528, "sample_models": [ - "RWKV/rwkv-5-world-1b5", - "RWKV/rwkv-5-world-3b" + "NeTS-lab/babylm26_multiling_hawk_MoP16K_baseline" ], - "relevancy_score": 36.0 + "relevancy_score": 36.1 }, { - "architecture_id": "OrthrusLM", - "total_models": 6, - "total_downloads": 5429, - "min_param_count": 2072903680, + "architecture_id": "HulumedQwen3ForCausalLM", + "total_models": 4, + "total_downloads": 48149, + "min_param_count": 4833079536, "sample_models": [ - "chiennv/Orthrus-Qwen3-8B", - "chiennv/Orthrus-Qwen3-1.7B" + "ZJU-AI4H/Hulu-Med-4B" ], - "relevancy_score": 35.8 + "relevancy_score": 35.9 }, { - "architecture_id": "YoutuForCausalLM", + "architecture_id": "SMBUnstructuredForCausalLM", "total_models": 4, - "total_downloads": 7062, - "min_param_count": 1961560064, + "total_downloads": 6954, + "min_param_count": 1720060928, "sample_models": [ - "tencent/Youtu-LLM-2B-Base", - "tencent/Youtu-LLM-2B" + "standardmodelbio/SMB-v1_Qwen3-1.7b_multi-objective" ], "relevancy_score": 35.8 }, { - "architecture_id": "EshmunForCausalLM", - "total_models": 2, - "total_downloads": 1338, - "min_param_count": 341319680, + "architecture_id": "FlexOlmoForCausalLM", + "total_models": 12, + "total_downloads": 95858, + "min_param_count": 11627401216, "sample_models": [ - "khairi/Eshmun-0.3B-CPT" + "allenai/Flex-reddit-2x7B-1T", + "allenai/FlexOlmo-7x7B-1T", + "shanearora/Flex-reddit-2x7B-1T" ], "relevancy_score": 35.7 }, { - "architecture_id": "FunAsrNanoForConditionalGeneration", - "total_models": 2, - "total_downloads": 1296, - "min_param_count": 985374304, + "architecture_id": "InternLM3ForCausalLM", + "total_models": 4, + "total_downloads": 296531, + "min_param_count": 8804241408, "sample_models": [ - "FunAudioLLM/Fun-ASR-Nano-2512-hf" + "internlm/internlm3-8b-instruct" ], "relevancy_score": 35.7 }, { - "architecture_id": "SLMForCausalLM", - "total_models": 2, - "total_downloads": 1314, - "min_param_count": 12065792, + "architecture_id": "RavenForCausalLM", + "total_models": 4, + "total_downloads": 45315, + "min_param_count": 3911400096, "sample_models": [ - "liodon-ai/slm-10m" + "tomg-group-umd/huginn-0125" ], "relevancy_score": 35.7 }, { - "architecture_id": "Qwen3VLMoeForConditionalGeneration", - "total_models": 12, - "total_downloads": 4212596, - "min_param_count": 31070754032, + "architecture_id": "LlavaMistralForCausalLM", + "total_models": 8, + "total_downloads": 166198, + "min_param_count": 7566219264, "sample_models": [ - "QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ", - "QuantTrio/Qwen3-VL-30B-A3B-Thinking-AWQ", - "QuantTrio/Qwen3-VL-235B-A22B-Instruct-AWQ", - "QuantTrio/Qwen3-VL-235B-A22B-Thinking-AWQ" + "liuhaotian/llava-v1.6-mistral-7b", + "microsoft/llava-med-v1.5-mistral-7b" ], "relevancy_score": 35.6 }, { - "architecture_id": "NanoChatForCausalLM", - "total_models": 4, - "total_downloads": 6271, - "min_param_count": 2217082880, + "architecture_id": "ArkasrForConditionalGeneration", + "total_models": 5, + "total_downloads": 5424, + "min_param_count": 1299463168, "sample_models": [ - "l0ulan/chinese-babylm-nanochat-d22-step10000", - "Twobombs/nanochat-d34-sft-hf" + "AutoArk-AI/ARK-ASR-0.6B", + "AutoArk-AI/ARK-ASR-3B" ], "relevancy_score": 35.6 }, { - "architecture_id": "HybridQwen3ForCausalLM", - "total_models": 9, - "total_downloads": 132477, - "min_param_count": 8500245504, + "architecture_id": "RwkvForCausalLM", + "total_models": 17, + "total_downloads": 42973, + "min_param_count": 7518044160, "sample_models": [ - "amazon/GKA-primed-HQwen3-32B-Instruct", - "amazon/GKA-primed-HQwen3-8B-Reasoner", - "amazon/GKA-primed-HQwen3-32B-Reasoner" + "RWKV/v5-Eagle-7B-HF", + "RWKV/rwkv-4-169m-pile", + "RWKV/v5-EagleX-v2-7B-HF", + "RWKV/rwkv-4-430m-pile", + "RWKV/rwkv-raven-1b5", + "RWKV/rwkv-raven-14b", + "recursal/EagleX_1-7T_HF" ], "relevancy_score": 35.4 }, { - "architecture_id": "GTLMForCausalLM", - "total_models": 6, - "total_downloads": 4417, - "min_param_count": 2095989760, + "architecture_id": "IsaacForConditionalGeneration", + "total_models": 4, + "total_downloads": 5445, + "min_param_count": 2567073008, "sample_models": [ - "Madras1/GTLM-1-2B-A350M", - "Madras1/GTLM-1-2B-A350M-fp16" + "PerceptronAI/Isaac-0.1" ], - "relevancy_score": 35.4 + "relevancy_score": 35.3 }, { - "architecture_id": "QuarkForCausalLM", + "architecture_id": "NanochatForCausalLM", "total_models": 2, - "total_downloads": 1158, - "min_param_count": 302080768, - "sample_models": [ - "ThingAI/Quark-270m-Base" - ], - "relevancy_score": 35.4 - }, - { - "architecture_id": "HulumedQwen3ForCausalLM", - "total_models": 3, - "total_downloads": 42626, - "min_param_count": 4833079536, + "total_downloads": 1067, + "min_param_count": 524288824, "sample_models": [ - "ZJU-AI4H/Hulu-Med-4B" + "himalaya-ai/himalayagpt-0.5b-it" ], "relevancy_score": 35.3 }, @@ -2854,15 +2982,49 @@ "sample_models": [ "i3-lab/i3-tiny" ], - "relevancy_score": 35.3 + "relevancy_score": 35.3 + }, + { + "architecture_id": "LLaMAForCausalLM", + "total_models": 7, + "total_downloads": 23207, + "min_param_count": 6738425856, + "sample_models": [ + "maicomputer/alpaca-13b", + "AdaptLLM/finance-LLM" + ], + "relevancy_score": 35.2 + }, + { + "architecture_id": "IQuestCoderForCausalLM", + "total_models": 8, + "total_downloads": 130381, + "min_param_count": 7612810240, + "sample_models": [ + "IQuestLab/IQuest-Coder-V1-40B-Instruct", + "IQuestLab/IQuest-Coder-V1-7B-Instruct", + "IQuestLab/IQuest-Coder-V1-14B-Base" + ], + "relevancy_score": 35.1 + }, + { + "architecture_id": "Sarashina2VisionForCausalLM", + "total_models": 8, + "total_downloads": 19340, + "min_param_count": 3801475696, + "sample_models": [ + "sbintuitions/sarashina2.2-vision-3b", + "sbintuitions/sarashina2.2-ocr" + ], + "relevancy_score": 35.1 }, { - "architecture_id": "RavenForCausalLM", - "total_models": 3, - "total_downloads": 38693, - "min_param_count": 3911400096, + "architecture_id": "JinaVLMForConditionalGeneration", + "total_models": 4, + "total_downloads": 4948, + "min_param_count": 2481403760, "sample_models": [ - "tomg-group-umd/huginn-0125" + "jinaai/jina-vlm" ], "relevancy_score": 35.1 }, @@ -2877,119 +3039,131 @@ "relevancy_score": 35.1 }, { - "architecture_id": "MolmoForCausalLM", - "total_models": 12, - "total_downloads": 71197, - "min_param_count": 7665032192, + "architecture_id": "NeoLLMForCausalLM", + "total_models": 1, + "total_downloads": 1137, + "min_param_count": 117311018, "sample_models": [ - "allenai/Molmo-7B-D-0924", - "allenai/Molmo-72B-0924", - "allenai/Molmo-7B-O-0924", - "cyan2k/molmo-7B-O-bnb-4bit" + "KitsuVp/NeoLLM" + ], + "relevancy_score": 35.1 + }, + { + "architecture_id": "PenguinVLQwen3ForCausalLM", + "total_models": 4, + "total_downloads": 4881, + "min_param_count": 2167941120, + "sample_models": [ + "tencent/Penguin-VL-2B" ], "relevancy_score": 35.0 }, { - "architecture_id": "SMBUnstructuredForCausalLM", - "total_models": 3, - "total_downloads": 5288, - "min_param_count": 1720060928, + "architecture_id": "CheXagentForCausalLM", + "total_models": 4, + "total_downloads": 31982, + "min_param_count": 3140746752, "sample_models": [ - "standardmodelbio/SMB-v1_Qwen3-1.7b_multi-objective" + "StanfordAIMI/CheXagent-2-3b" ], - "relevancy_score": 34.9 + "relevancy_score": 35.0 }, { - "architecture_id": "InternLM3ForCausalLM", + "architecture_id": "BarbetForCausalLM", "total_models": 3, - "total_downloads": 217423, - "min_param_count": 8804241408, + "total_downloads": 5615, + "min_param_count": 1088124920, "sample_models": [ - "internlm/internlm3-8b-instruct" + "OpenFormosa/barbet-1b-base" ], - "relevancy_score": 34.7 + "relevancy_score": 35.0 }, { - "architecture_id": "MusicFlamingoForConditionalGeneration", - "total_models": 5, - "total_downloads": 158343, - "min_param_count": 8267215360, + "architecture_id": "JAISLMHeadModel", + "total_models": 15, + "total_downloads": 41502, + "min_param_count": 13462730280, "sample_models": [ - "nvidia/music-flamingo-2601-hf", - "nvidia/audio-flamingo-next-hf", - "nvidia/audio-flamingo-next-think-hf", - "nvidia/audio-flamingo-next-captioner-hf", - "nvidia/music-flamingo-think-2601-hf" + "inceptionai/jais-13b-chat", + "inceptionai/jais-13b", + "inceptionai/jais-family-13b-chat", + "inceptionai/jais-30b-chat-v3", + "inceptionai/jais-30b-v3" ], - "relevancy_score": 34.6 + "relevancy_score": 34.8 }, { - "architecture_id": "LlavaMistralForCausalLM", - "total_models": 6, - "total_downloads": 128525, - "min_param_count": 7566219264, + "architecture_id": "Jais2ForCausalLM", + "total_models": 8, + "total_downloads": 110000, + "min_param_count": 8090401280, "sample_models": [ - "liuhaotian/llava-v1.6-mistral-7b", - "microsoft/llava-med-v1.5-mistral-7b" + "inceptionai/Jais-2-8B-Chat", + "inceptionai/Jais-2-70B-Chat" ], - "relevancy_score": 34.5 + "relevancy_score": 34.8 }, { - "architecture_id": "LLaMAForCausalLM", - "total_models": 6, - "total_downloads": 17984, - "min_param_count": 6738425856, + "architecture_id": "Maira2ForConditionalGeneration", + "total_models": 4, + "total_downloads": 28821, + "min_param_count": 6880185600, "sample_models": [ - "maicomputer/alpaca-13b", - "AdaptLLM/finance-LLM" + "microsoft/maira-2" ], - "relevancy_score": 34.4 + "relevancy_score": 34.8 }, { - "architecture_id": "PenguinVLQwen3ForCausalLM", - "total_models": 3, - "total_downloads": 4239, - "min_param_count": 2167941120, + "architecture_id": "XCurOSForCausalLM", + "total_models": 4, + "total_downloads": 198170, + "min_param_count": 7615616512, "sample_models": [ - "tencent/Penguin-VL-2B" + "XCurOS/XCurOS0.1-8B-Instruct" ], - "relevancy_score": 34.4 + "relevancy_score": 34.8 }, { - "architecture_id": "RwkvForCausalLM", - "total_models": 15, - "total_downloads": 33087, - "min_param_count": 7518044160, + "architecture_id": "MusicFlamingoForConditionalGeneration", + "total_models": 5, + "total_downloads": 158343, + "min_param_count": 8267215360, "sample_models": [ - "RWKV/v5-Eagle-7B-HF", - "RWKV/rwkv-4-169m-pile", - "RWKV/v5-EagleX-v2-7B-HF", - "RWKV/rwkv-4-430m-pile", - "RWKV/rwkv-raven-1b5", - "RWKV/rwkv-raven-14b", - "recursal/EagleX_1-7T_HF" + "nvidia/music-flamingo-2601-hf", + "nvidia/audio-flamingo-next-hf", + "nvidia/audio-flamingo-next-think-hf", + "nvidia/audio-flamingo-next-captioner-hf", + "nvidia/music-flamingo-think-2601-hf" ], - "relevancy_score": 34.3 + "relevancy_score": 34.7 }, { - "architecture_id": "IQuestCoderForCausalLM", - "total_models": 7, - "total_downloads": 98324, - "min_param_count": 7612810240, + "architecture_id": "GritLM", + "total_models": 4, + "total_downloads": 181545, + "min_param_count": 7241732096, "sample_models": [ - "IQuestLab/IQuest-Coder-V1-40B-Instruct", - "IQuestLab/IQuest-Coder-V1-7B-Instruct", - "IQuestLab/IQuest-Coder-V1-14B-Base" + "parasail-ai/GritLM-7B-vllm" ], - "relevancy_score": 34.2 + "relevancy_score": 34.7 }, { - "architecture_id": "IsaacForConditionalGeneration", - "total_models": 3, - "total_downloads": 3780, - "min_param_count": 2567073008, + "architecture_id": "PanguEmbeddedForCausalLM", + "total_models": 4, + "total_downloads": 4014, + "min_param_count": 1391497728, "sample_models": [ - "PerceptronAI/Isaac-0.1" + "FreedomIntelligence/openPangu-Embedded-1B" + ], + "relevancy_score": 34.6 + }, + { + "architecture_id": "ParamBharatGenForCausalLM", + "total_models": 4, + "total_downloads": 3232, + "min_param_count": 2860693504, + "sample_models": [ + "bharatgenai/Param-1-2.9B-Instruct" ], "relevancy_score": 34.2 }, @@ -3005,75 +3179,71 @@ "relevancy_score": 34.2 }, { - "architecture_id": "FlexOlmoForCausalLM", - "total_models": 9, - "total_downloads": 71629, - "min_param_count": 11627401216, - "sample_models": [ - "allenai/Flex-reddit-2x7B-1T", - "allenai/FlexOlmo-7x7B-1T", - "shanearora/Flex-reddit-2x7B-1T" - ], - "relevancy_score": 34.1 - }, - { - "architecture_id": "JinaVLMForConditionalGeneration", - "total_models": 3, - "total_downloads": 3610, - "min_param_count": 2481403760, + "architecture_id": "MiMoV2ForCausalLM", + "total_models": 22, + "total_downloads": 512881, + "min_param_count": 52930771392, "sample_models": [ - "jinaai/jina-vlm" + "XiaomiMiMo/MiMo-V2.5-Pro", + "XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash", + "spectator2026/MiMo-V2.5-AWQ-int4", + "kernelpool/MiMo-V2.5-Pro-6bit", + "XiaomiMiMo/MiMo-V2.5-Pro-Base", + "nwzjk/MiMo-V2.5-AWQ-int4" ], "relevancy_score": 34.1 }, { - "architecture_id": "Sarashina2VisionForCausalLM", - "total_models": 6, - "total_downloads": 15158, - "min_param_count": 3801475696, + "architecture_id": "EveMoEForCausalLM", + "total_models": 1, + "total_downloads": 678, + "min_param_count": 271970816, "sample_models": [ - "sbintuitions/sarashina2.2-vision-3b", - "sbintuitions/sarashina2.2-ocr" + "anthonym21/Eve-2-MoE-IT-272M" ], "relevancy_score": 34.0 }, { - "architecture_id": "CheXagentForCausalLM", - "total_models": 3, - "total_downloads": 23029, - "min_param_count": 3140746752, + "architecture_id": "LagunaForCausalLM", + "total_models": 18, + "total_downloads": 801545, + "min_param_count": 33442617088, "sample_models": [ - "StanfordAIMI/CheXagent-2-3b" + "poolside/Laguna-XS.2", + "cyankiwi/Laguna-XS.2-AWQ-INT4", + "poolside/Laguna-XS.2-INT4", + "poolside/Laguna-M.1", + "poolside/Laguna-M.1-base" ], - "relevancy_score": 34.0 + "relevancy_score": 33.9 }, { - "architecture_id": "EveMoEForCausalLM", - "total_models": 1, - "total_downloads": 678, - "min_param_count": 271970816, + "architecture_id": "MaskMoELLaVAQwen2ForCausalLM", + "total_models": 4, + "total_downloads": 2805, + "min_param_count": 1563012736, "sample_models": [ - "anthonym21/Eve-2-MoE-IT-272M" + "KKHYA/llavaqwen2.5-0.5b-finetune-mask-moe-sparse-4e-2k-sp0.9_20260416_183544" ], - "relevancy_score": 34.0 + "relevancy_score": 33.9 }, { - "architecture_id": "XCurOSForCausalLM", - "total_models": 3, - "total_downloads": 145740, - "min_param_count": 7615616512, + "architecture_id": "CacaForCausalLM", + "total_models": 4, + "total_downloads": 2820, + "min_param_count": 1000020992, "sample_models": [ - "XCurOS/XCurOS0.1-8B-Instruct" + "Lyon28/caca-1B-untrained" ], "relevancy_score": 33.9 }, { - "architecture_id": "BarbetForCausalLM", - "total_models": 2, - "total_downloads": 3702, - "min_param_count": 1088124920, + "architecture_id": "Papagan", + "total_models": 4, + "total_downloads": 2780, + "min_param_count": 1298763264, "sample_models": [ - "OpenFormosa/barbet-1b-base" + "SutskeverFanBoy/papagan_1.3b" ], "relevancy_score": 33.9 }, @@ -3088,47 +3258,34 @@ "relevancy_score": 33.9 }, { - "architecture_id": "XLNetLMHeadModel", - "total_models": 18, - "total_downloads": 799167, - "min_param_count": null, - "sample_models": [ - "xlnet/xlnet-base-cased", - "xlnet/xlnet-large-cased", - "hfl/chinese-xlnet-base", - "sshleifer/tiny-xlnet-base-cased", - "textattack/xlnet-base-cased-imdb", - "textattack/xlnet-base-cased-CoLA" - ], - "relevancy_score": 33.8 - }, - { - "architecture_id": "ArkasrForConditionalGeneration", - "total_models": 3, - "total_downloads": 3070, - "min_param_count": 1299463168, + "architecture_id": "Phi3SmallForCausalLM", + "total_models": 8, + "total_downloads": 70709, + "min_param_count": 7392272384, "sample_models": [ - "AutoArk-AI/ARK-ASR-0.6B" + "microsoft/Phi-3-small-8k-instruct", + "microsoft/Phi-3-small-128k-instruct" ], "relevancy_score": 33.8 }, { - "architecture_id": "Maira2ForConditionalGeneration", - "total_models": 3, - "total_downloads": 21321, - "min_param_count": 6880185600, + "architecture_id": "InternS1ForConditionalGeneration", + "total_models": 8, + "total_downloads": 70025, + "min_param_count": 8538804224, "sample_models": [ - "microsoft/maira-2" + "internlm/Intern-S1", + "internlm/Intern-S1-mini" ], "relevancy_score": 33.8 }, { - "architecture_id": "PanguEmbeddedForCausalLM", - "total_models": 3, - "total_downloads": 3099, - "min_param_count": 1391497728, + "architecture_id": "MoELLaVAQwen2ForCausalLM", + "total_models": 4, + "total_downloads": 2662, + "min_param_count": 1406119552, "sample_models": [ - "FreedomIntelligence/openPangu-Embedded-1B" + "KKHYA/llavaqwen2.5-0.5b-finetune-moe-4e-2k_20260331_194516" ], "relevancy_score": 33.8 }, @@ -3143,12 +3300,12 @@ "relevancy_score": 33.8 }, { - "architecture_id": "GritLM", - "total_models": 3, - "total_downloads": 135264, - "min_param_count": 7241732096, + "architecture_id": "MotifForCausalLM", + "total_models": 4, + "total_downloads": 2574, + "min_param_count": 2597218432, "sample_models": [ - "parasail-ai/GritLM-7B-vllm" + "Motif-Technologies/Motif-2.6B" ], "relevancy_score": 33.7 }, @@ -3163,13 +3320,25 @@ "relevancy_score": 33.7 }, { - "architecture_id": "Jais2ForCausalLM", - "total_models": 6, - "total_downloads": 82006, - "min_param_count": 8090401280, + "architecture_id": "IdeficsForVisionText2Text", + "total_models": 12, + "total_downloads": 35400, + "min_param_count": 8929682192, "sample_models": [ - "inceptionai/Jais-2-8B-Chat", - "inceptionai/Jais-2-70B-Chat" + "HuggingFaceM4/idefics-80b-instruct", + "HuggingFaceM4/idefics-9b", + "HuggingFaceM4/idefics-9b-instruct" + ], + "relevancy_score": 33.6 + }, + { + "architecture_id": "ChatGLMModel", + "total_models": 8, + "total_downloads": 64320, + "min_param_count": 9399951392, + "sample_models": [ + "zai-org/glm-4-9b", + "zai-org/codegeex4-all-9b" ], "relevancy_score": 33.6 }, @@ -3184,6 +3353,16 @@ ], "relevancy_score": 33.6 }, + { + "architecture_id": "SMTModelForCausalLM", + "total_models": 1, + "total_downloads": 569, + "min_param_count": 16126425, + "sample_models": [ + "JuanCarlosMartinezSevilla/jazzmus-model" + ], + "relevancy_score": 33.6 + }, { "architecture_id": "XLMRobertaForCausalLM", "total_models": 1, @@ -3195,12 +3374,22 @@ "relevancy_score": 33.6 }, { - "architecture_id": "NanochatForCausalLM", + "architecture_id": "QuasarForCausalLM", + "total_models": 4, + "total_downloads": 106104, + "min_param_count": 8602037248, + "sample_models": [ + "silx-ai/Quasar-10B" + ], + "relevancy_score": 33.5 + }, + { + "architecture_id": "TinyQwen3NoveltyForCausalLM", "total_models": 1, - "total_downloads": 531, - "min_param_count": 524288824, + "total_downloads": 533, + "min_param_count": 8131968, "sample_models": [ - "himalaya-ai/himalayagpt-0.5b-it" + "User01110/tinyLM-8M-exp" ], "relevancy_score": 33.5 }, @@ -3224,16 +3413,6 @@ ], "relevancy_score": 33.4 }, - { - "architecture_id": "ParamBharatGenForCausalLM", - "total_models": 3, - "total_downloads": 2386, - "min_param_count": 2860693504, - "sample_models": [ - "bharatgenai/Param-1-2.9B-Instruct" - ], - "relevancy_score": 33.2 - }, { "architecture_id": "SpatialLMLlamaForCausalLM", "total_models": 3, @@ -3245,55 +3424,14 @@ "relevancy_score": 33.2 }, { - "architecture_id": "CacaForCausalLM", - "total_models": 3, - "total_downloads": 2269, - "min_param_count": 1000020992, - "sample_models": [ - "Lyon28/caca-1B-untrained" - ], - "relevancy_score": 33.1 - }, - { - "architecture_id": "MaskMoELLaVAQwen2ForCausalLM", - "total_models": 3, - "total_downloads": 2265, - "min_param_count": 1563012736, - "sample_models": [ - "KKHYA/llavaqwen2.5-0.5b-finetune-mask-moe-sparse-4e-2k-sp0.9_20260416_183544" - ], - "relevancy_score": 33.1 - }, - { - "architecture_id": "Papagan", - "total_models": 3, - "total_downloads": 2226, - "min_param_count": 1298763264, - "sample_models": [ - "SutskeverFanBoy/papagan_1.3b" - ], - "relevancy_score": 33.1 - }, - { - "architecture_id": "MoELLaVAQwen2ForCausalLM", - "total_models": 3, - "total_downloads": 2133, - "min_param_count": 1406119552, - "sample_models": [ - "KKHYA/llavaqwen2.5-0.5b-finetune-moe-4e-2k_20260331_194516" - ], - "relevancy_score": 33.0 - }, - { - "architecture_id": "JAISLMHeadModel", - "total_models": 10, - "total_downloads": 33463, - "min_param_count": 13462730280, + "architecture_id": "BailingMoeV2ForCausalLM", + "total_models": 12, + "total_downloads": 167726, + "min_param_count": 16255643392, "sample_models": [ - "inceptionai/jais-13b-chat", - "inceptionai/jais-13b", - "inceptionai/jais-family-13b-chat", - "inceptionai/jais-30b-chat-v3" + "inclusionAI/Ling-mini-2.0", + "inclusionAI/Ling-1T", + "inclusionAI/Ling-flash-2.0" ], "relevancy_score": 32.8 }, @@ -3308,55 +3446,60 @@ "relevancy_score": 32.8 }, { - "architecture_id": "MotifForCausalLM", - "total_models": 3, - "total_downloads": 1926, - "min_param_count": 2597218432, + "architecture_id": "EmoForCausalLM", + "total_models": 5, + "total_downloads": 9203, + "min_param_count": 3901818880, "sample_models": [ - "Motif-Technologies/Motif-2.6B" + "allenai/Emo_1b14b_1T", + "allenai/StdMoE_1b4b_130B" ], - "relevancy_score": 32.8 + "relevancy_score": 32.7 }, { - "architecture_id": "InternS1ForConditionalGeneration", - "total_models": 6, - "total_downloads": 52526, - "min_param_count": 8538804224, + "architecture_id": "AprielForCausalLM", + "total_models": 4, + "total_downloads": 10691, + "min_param_count": 4832071680, "sample_models": [ - "internlm/Intern-S1", - "internlm/Intern-S1-mini" + "ServiceNow-AI/Apriel-5B-Instruct" ], - "relevancy_score": 32.6 + "relevancy_score": 32.7 }, { - "architecture_id": "Phi3SmallForCausalLM", - "total_models": 6, - "total_downloads": 50758, - "min_param_count": 7392272384, + "architecture_id": "HulumedQwen2ForCausalLM", + "total_models": 7, + "total_downloads": 43346, + "min_param_count": 8044744944, "sample_models": [ - "microsoft/Phi-3-small-8k-instruct", - "microsoft/Phi-3-small-128k-instruct" + "ZJU-AI4H/Hulu-Med-7B", + "ZJU-AI4H/Hulu-Med-32B" ], "relevancy_score": 32.5 }, { - "architecture_id": "QuasarForCausalLM", - "total_models": 3, - "total_downloads": 75373, - "min_param_count": 8602037248, + "architecture_id": "Rwkv6ForCausalLM", + "total_models": 15, + "total_downloads": 13615, + "min_param_count": 7635746816, "sample_models": [ - "silx-ai/Quasar-10B" + "RWKV/rwkv-6-world-3b", + "RWKV/rwkv-6-world-1b6", + "RWKV/v6-Finch-1B6-HF", + "RWKV/v6-Finch-14B-HF", + "RWKV/v6-Finch-7B-HF" ], - "relevancy_score": 32.5 + "relevancy_score": 32.4 }, { - "architecture_id": "ChatGLMModel", - "total_models": 6, - "total_downloads": 48349, - "min_param_count": 9399951392, + "architecture_id": "BailingMoeForCausalLM", + "total_models": 10, + "total_downloads": 179673, + "min_param_count": 16801974272, "sample_models": [ - "zai-org/glm-4-9b", - "zai-org/codegeex4-all-9b" + "inclusionAI/Ling-lite-1.5", + "inclusionAI/Ling-lite", + "inclusionAI/Ling-plus" ], "relevancy_score": 32.4 }, @@ -3372,1223 +3515,1135 @@ "relevancy_score": 32.4 }, { - "architecture_id": "LagunaForCausalLM", - "total_models": 13, - "total_downloads": 705719, - "min_param_count": 33442617088, - "sample_models": [ - "poolside/Laguna-XS.2", - "cyankiwi/Laguna-XS.2-AWQ-INT4", - "poolside/Laguna-XS.2-INT4", - "poolside/Laguna-M.1", - "poolside/Laguna-M.1-base" - ], - "relevancy_score": 32.1 - }, - { - "architecture_id": "IdeficsForVisionText2Text", - "total_models": 9, - "total_downloads": 24198, - "min_param_count": 8929682192, - "sample_models": [ - "HuggingFaceM4/idefics-80b-instruct", - "HuggingFaceM4/idefics-9b", - "HuggingFaceM4/idefics-9b-instruct" - ], - "relevancy_score": 31.9 - }, - { - "architecture_id": "EmoForCausalLM", - "total_models": 4, - "total_downloads": 7549, - "min_param_count": 3901818880, - "sample_models": [ - "allenai/Emo_1b14b_1T", - "allenai/StdMoE_1b4b_130B" - ], - "relevancy_score": 31.9 - }, - { - "architecture_id": "MiMoV2ForCausalLM", - "total_models": 17, - "total_downloads": 354391, - "min_param_count": 52930771392, - "sample_models": [ - "XiaomiMiMo/MiMo-V2.5-Pro", - "XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash", - "spectator2026/MiMo-V2.5-AWQ-int4", - "kernelpool/MiMo-V2.5-Pro-6bit", - "XiaomiMiMo/MiMo-V2.5-Pro-Base", - "nwzjk/MiMo-V2.5-AWQ-int4" - ], - "relevancy_score": 31.8 - }, - { - "architecture_id": "AprielForCausalLM", - "total_models": 3, - "total_downloads": 8042, - "min_param_count": 4832071680, - "sample_models": [ - "ServiceNow-AI/Apriel-5B-Instruct" - ], - "relevancy_score": 31.8 - }, - { - "architecture_id": "EfficientDLM", - "total_models": 6, - "total_downloads": 4991, - "min_param_count": 4411424256, - "sample_models": [ - "nvidia/Efficient-DLM-8B", - "nvidia/Efficient-DLM-4B" - ], - "relevancy_score": 31.7 - }, - { - "architecture_id": "HulumedQwen2ForCausalLM", - "total_models": 5, - "total_downloads": 35622, - "min_param_count": 8044744944, - "sample_models": [ - "ZJU-AI4H/Hulu-Med-7B", - "ZJU-AI4H/Hulu-Med-32B" - ], - "relevancy_score": 31.5 - }, - { - "architecture_id": "PITForCausalLM", - "total_models": 4, - "total_downloads": 6019, - "min_param_count": 4232577024, - "sample_models": [ - "NurErtug/pit-finance-grpo-merged", - "Dapinsky/PIT-4B-FT-202212-math-reasoning-dpo" - ], - "relevancy_score": 31.5 - }, - { - "architecture_id": "Rwkv6ForCausalLM", - "total_models": 13, - "total_downloads": 11376, - "min_param_count": 7635746816, + "architecture_id": "InternLMForCausalLM", + "total_models": 20, + "total_downloads": 280342, + "min_param_count": null, "sample_models": [ - "RWKV/rwkv-6-world-3b", - "RWKV/rwkv-6-world-1b6", - "RWKV/v6-Finch-1B6-HF", - "RWKV/v6-Finch-14B-HF", - "RWKV/v6-Finch-7B-HF" + "internlm/internlm-chat-7b", + "internlm/internlm-7b", + "internlm/internlm-20b", + "TongjiFinLab/CFGPT1-pt-7B", + "internlm/internlm-chat-20b", + "internlm/Agent-FLAN-7b" ], - "relevancy_score": 31.4 + "relevancy_score": 32.2 }, - { - "architecture_id": "GPTSanJapaneseForConditionalGeneration", - "total_models": 2, - "total_downloads": 1158, - "min_param_count": 2779000992, + { + "architecture_id": "PITForCausalLM", + "total_models": 5, + "total_downloads": 6859, + "min_param_count": 4232577024, "sample_models": [ - "Tanrei/GPTSAN-japanese" + "NurErtug/pit-finance-grpo-merged", + "Dapinsky/PIT-4B-FT-202212-math-reasoning-dpo" ], - "relevancy_score": 31.4 + "relevancy_score": 32.1 }, { - "architecture_id": "BailingMoeV2ForCausalLM", - "total_models": 9, - "total_downloads": 126010, - "min_param_count": 16255643392, + "architecture_id": "DeepseekForCausalLM", + "total_models": 8, + "total_downloads": 192858, + "min_param_count": 16375728128, "sample_models": [ - "inclusionAI/Ling-mini-2.0", - "inclusionAI/Ling-1T", - "inclusionAI/Ling-flash-2.0" + "deepseek-ai/deepseek-moe-16b-base", + "deepseek-ai/deepseek-moe-16b-chat" ], - "relevancy_score": 31.3 + "relevancy_score": 32.0 }, { "architecture_id": "HunyuanImage3ForCausalMM", - "total_models": 3, - "total_downloads": 1789960, + "total_models": 4, + "total_downloads": 2091307, "min_param_count": 83009199459, "sample_models": [ "tencent/HunyuanImage-3.0" ], - "relevancy_score": 31.2 + "relevancy_score": 31.8 }, { - "architecture_id": "BailingMoeForCausalLM", - "total_models": 7, - "total_downloads": 133908, - "min_param_count": 16801974272, + "architecture_id": "EfficientDLM", + "total_models": 6, + "total_downloads": 4991, + "min_param_count": 4411424256, "sample_models": [ - "inclusionAI/Ling-lite-1.5", - "inclusionAI/Ling-lite", - "inclusionAI/Ling-plus" + "nvidia/Efficient-DLM-8B", + "nvidia/Efficient-DLM-4B" ], - "relevancy_score": 30.9 + "relevancy_score": 31.7 }, { - "architecture_id": "DeepseekForCausalLM", - "total_models": 6, - "total_downloads": 143523, - "min_param_count": 16375728128, + "architecture_id": "StepAudio2ForCausalLM", + "total_models": 4, + "total_downloads": 42507, + "min_param_count": 8315179264, "sample_models": [ - "deepseek-ai/deepseek-moe-16b-base", - "deepseek-ai/deepseek-moe-16b-chat" + "stepfun-ai/Step-Audio-2-mini" ], - "relevancy_score": 30.7 + "relevancy_score": 31.6 }, { - "architecture_id": "InternLMForCausalLM", - "total_models": 16, - "total_downloads": 221634, - "min_param_count": null, + "architecture_id": "GPTSanJapaneseForConditionalGeneration", + "total_models": 2, + "total_downloads": 1158, + "min_param_count": 2779000992, "sample_models": [ - "internlm/internlm-chat-7b", - "internlm/internlm-7b", - "internlm/internlm-20b", - "TongjiFinLab/CFGPT1-pt-7B", - "internlm/internlm-chat-20b", - "internlm/Agent-FLAN-7b" + "Tanrei/GPTSAN-japanese" ], - "relevancy_score": 30.6 + "relevancy_score": 31.4 }, { "architecture_id": "CXRMate2ForConditionalGeneration", - "total_models": 3, - "total_downloads": 4254, + "total_models": 4, + "total_downloads": 5501, "min_param_count": 3322260224, "sample_models": [ "aehrc/cxrmate-2" ], - "relevancy_score": 30.5 - }, - { - "architecture_id": "StepAudio2ForCausalLM", - "total_models": 3, - "total_downloads": 27607, - "min_param_count": 8315179264, - "sample_models": [ - "stepfun-ai/Step-Audio-2-mini" - ], - "relevancy_score": 30.4 + "relevancy_score": 31.3 }, { "architecture_id": "BioGptForCausalLM", - "total_models": 9, - "total_downloads": 446425, + "total_models": 11, + "total_downloads": 587762, "min_param_count": null, "sample_models": [ "microsoft/biogpt", "microsoft/BioGPT-Large", "microsoft/BioGPT-Large-PubMedQA" ], - "relevancy_score": 30.0 + "relevancy_score": 31.2 }, { "architecture_id": "DyninOmniForConditionalGeneration", - "total_models": 3, - "total_downloads": 22512, + "total_models": 4, + "total_downloads": 28962, "min_param_count": 8116244480, "sample_models": [ "snu-aidas/Dynin-Omni" ], - "relevancy_score": 30.0 - }, - { - "architecture_id": "AeroForConditionalGeneration", - "total_models": 1, - "total_downloads": 680, - "min_param_count": 2416221184, - "sample_models": [ - "lmms-lab/Aero-1-Audio" - ], - "relevancy_score": 30.0 + "relevancy_score": 30.8 }, { - "architecture_id": "StepVLForConditionalGeneration", - "total_models": 3, - "total_downloads": 20795, - "min_param_count": 10171750144, + "architecture_id": "CogVLMVideoForCausalLM", + "total_models": 8, + "total_downloads": 15909, + "min_param_count": 12507532544, "sample_models": [ - "QuantTrio/Step3-VL-10B-AWQ" + "zai-org/VisionReward-Video", + "zai-org/cogvlm2-llama3-caption" ], - "relevancy_score": 29.8 + "relevancy_score": 30.7 }, { - "architecture_id": "InstellaForCausalLM", - "total_models": 3, - "total_downloads": 2918, - "min_param_count": 3112675840, + "architecture_id": "GPTBERTForCausalLM", + "total_models": 32, + "total_downloads": 24615, + "min_param_count": null, "sample_models": [ - "amd/Instella-3B-Instruct" + "BabyLM-community/GPT-BERT-baseline-babylm-eng-nld-zho-equal", + "BabyLM-community/GPT-BERT-baseline-babylm-Strict", + "BabyLM-community/GPT-BERT-baseline-babylm-eng-nld-equal", + "BabyLM-community/GPT-BERT-baseline-babylm-zho", + "BabyLM-community/GPT-BERT-baseline-babylm-nld", + "BabyLM-community/GPT-BERT-baseline-babylm-Strict-Small", + "BabyLM-community/GPT-BERT-baseline-babylm-zho-nld-equal", + "BabyLM-community/GPT-BERT-baseline-babylm-eng-zho-equal" ], - "relevancy_score": 29.7 + "relevancy_score": 30.6 }, { "architecture_id": "Phi4FlashForCausalLM", - "total_models": 3, - "total_downloads": 3027, + "total_models": 4, + "total_downloads": 3949, "min_param_count": 3852562944, "sample_models": [ "microsoft/Phi-4-mini-flash-reasoning" ], - "relevancy_score": 29.7 + "relevancy_score": 30.6 }, { "architecture_id": "LamedPhi3ForCausalLM", - "total_models": 3, - "total_downloads": 2877, + "total_models": 4, + "total_downloads": 3814, "min_param_count": 4049101904, "sample_models": [ "GoodBaiBai88/M3D-LaMed-Phi-3-4B" ], - "relevancy_score": 29.6 - }, - { - "architecture_id": "TrainableM2MForConditionalGeneration", - "total_models": 1, - "total_downloads": 570, - "min_param_count": 1370638336, - "sample_models": [ - "Sunbird/translate-nllb-1.3b-salt" - ], - "relevancy_score": 29.6 + "relevancy_score": 30.5 }, { - "architecture_id": "CogVLMVideoForCausalLM", - "total_models": 6, - "total_downloads": 11923, - "min_param_count": 12507532544, + "architecture_id": "InstellaForCausalLM", + "total_models": 4, + "total_downloads": 3834, + "min_param_count": 3112675840, "sample_models": [ - "zai-org/VisionReward-Video", - "zai-org/cogvlm2-llama3-caption" + "amd/Instella-3B-Instruct" ], - "relevancy_score": 29.5 + "relevancy_score": 30.5 }, { - "architecture_id": "MoELLaVAQwen3ForCausalLM", - "total_models": 3, - "total_downloads": 2725, - "min_param_count": 3927104512, + "architecture_id": "StepVLForConditionalGeneration", + "total_models": 4, + "total_downloads": 24891, + "min_param_count": 10171750144, "sample_models": [ - "KKHYA/llavaqwen3-1.7b-finetune-moe-4e-2k_20260427_233320" + "QuantTrio/Step3-VL-10B-AWQ" ], - "relevancy_score": 29.5 + "relevancy_score": 30.5 }, { "architecture_id": "SolarForCausalLM", - "total_models": 3, - "total_downloads": 123183, + "total_models": 4, + "total_downloads": 165696, "min_param_count": 22140032000, "sample_models": [ "upstage/solar-pro-preview-instruct" ], - "relevancy_score": 29.5 + "relevancy_score": 30.5 }, { "architecture_id": "InternS1ProForConditionalGeneration", - "total_models": 3, - "total_downloads": 784387, + "total_models": 4, + "total_downloads": 1068937, "min_param_count": null, "sample_models": [ "internlm/Intern-S1-Pro" ], - "relevancy_score": 29.4 + "relevancy_score": 30.4 + }, + { + "architecture_id": "Emu3ForCausalLM", + "total_models": 8, + "total_downloads": 13179, + "min_param_count": 8492011520, + "sample_models": [ + "BAAI/Emu3-Gen", + "BAAI/Emu3.5-Image" + ], + "relevancy_score": 30.3 }, { "architecture_id": "Param2MoEForCausalLM", - "total_models": 3, - "total_downloads": 113565, + "total_models": 4, + "total_downloads": 151661, "min_param_count": 17151140480, "sample_models": [ "bharatgenai/Param2-17B-A2.4B-Thinking" ], - "relevancy_score": 29.4 + "relevancy_score": 30.3 }, { "architecture_id": "BunnyPhiForCausalLM", - "total_models": 3, - "total_downloads": 2559, + "total_models": 4, + "total_downloads": 3342, "min_param_count": 3182254624, "sample_models": [ "BAAI/Bunny-v1_0-3B" ], - "relevancy_score": 29.4 + "relevancy_score": 30.3 }, { "architecture_id": "LibraLlamaForCausalLM", - "total_models": 3, - "total_downloads": 2562, + "total_models": 4, + "total_downloads": 3408, "min_param_count": 6800327990, "sample_models": [ "X-iZhang/libra-v1.0-7b" ], - "relevancy_score": 29.4 + "relevancy_score": 30.3 + }, + { + "architecture_id": "MoELLaVAQwen3ForCausalLM", + "total_models": 4, + "total_downloads": 3283, + "min_param_count": 3927104512, + "sample_models": [ + "KKHYA/llavaqwen3-1.7b-finetune-moe-4e-2k_20260427_233320" + ], + "relevancy_score": 30.2 + }, + { + "architecture_id": "ReformerModelWithLMHead", + "total_models": 8, + "total_downloads": 539392, + "min_param_count": null, + "sample_models": [ + "google/reformer-crime-and-punishment", + "google/reformer-enwik8" + ], + "relevancy_score": 30.1 + }, + { + "architecture_id": "MiniCPMSALAForCausalLM", + "total_models": 4, + "total_downloads": 20389, + "min_param_count": 9477203968, + "sample_models": [ + "openbmb/MiniCPM-SALA" + ], + "relevancy_score": 30.1 }, { "architecture_id": "Step3p5ForCausalLM", - "total_models": 3, - "total_downloads": 737498, + "total_models": 4, + "total_downloads": 937972, "min_param_count": 199384301376, "sample_models": [ "stepfun-ai/Step-3.5-Flash" ], - "relevancy_score": 29.3 + "relevancy_score": 30.1 + }, + { + "architecture_id": "AeroForConditionalGeneration", + "total_models": 1, + "total_downloads": 680, + "min_param_count": 2416221184, + "sample_models": [ + "lmms-lab/Aero-1-Audio" + ], + "relevancy_score": 30.0 }, { "architecture_id": "VeridianForCausalLM", - "total_models": 3, - "total_downloads": 2295, + "total_models": 4, + "total_downloads": 2841, "min_param_count": 3938699264, "sample_models": [ "MagistrTheOne/veridian-beta" ], - "relevancy_score": 29.2 + "relevancy_score": 29.9 }, { - "architecture_id": "MiniCPMSALAForCausalLM", - "total_models": 3, - "total_downloads": 15650, - "min_param_count": 9477203968, + "architecture_id": "LaneformerForCausalLM", + "total_models": 1, + "total_downloads": 640, + "min_param_count": 2320069632, "sample_models": [ - "openbmb/MiniCPM-SALA" + "kogai/laneformer-2b-it" ], - "relevancy_score": 29.2 + "relevancy_score": 29.9 }, { - "architecture_id": "Emu3ForCausalLM", - "total_models": 6, - "total_downloads": 9683, - "min_param_count": 8492011520, + "architecture_id": "BailingMoeV2_5ForCausalLM", + "total_models": 15, + "total_downloads": 171333, + "min_param_count": 107494409216, "sample_models": [ - "BAAI/Emu3-Gen", - "BAAI/Emu3.5-Image" + "inclusionAI/Ring-2.5-1T", + "inclusionAI/Ling-2.6-flash", + "inclusionAI/Ring-2.6-1T", + "inclusionAI/Ling-2.6-1T" ], - "relevancy_score": 29.1 + "relevancy_score": 29.8 }, { - "architecture_id": "ReformerModelWithLMHead", + "architecture_id": "KORMoForCausalLM", + "total_models": 8, + "total_downloads": 10312, + "min_param_count": 10756624384, + "sample_models": [ + "KORMo-Team/KORMo-10B-base", + "KORMo-Team/KORMo-10B-sft" + ], + "relevancy_score": 29.8 + }, + { + "architecture_id": "Videollama2MistralForCausalLM", "total_models": 6, - "total_downloads": 413979, - "min_param_count": null, + "total_downloads": 13576, + "min_param_count": 8034428160, "sample_models": [ - "google/reformer-crime-and-punishment", - "google/reformer-enwik8" + "DAMO-NLP-SG/VideoLLaMA2-7B-16F", + "DAMO-NLP-SG/VideoLLaMA2-7B" ], - "relevancy_score": 29.0 + "relevancy_score": 29.8 }, { "architecture_id": "SmartCoderMoEForCausalLM", - "total_models": 3, - "total_downloads": 2180, + "total_models": 4, + "total_downloads": 2749, "min_param_count": 4717965312, "sample_models": [ "Johnblick187/SmartCoderMoE" ], - "relevancy_score": 29.0 + "relevancy_score": 29.8 }, { "architecture_id": "TalkieForCausalLM", - "total_models": 3, - "total_downloads": 13136, + "total_models": 4, + "total_downloads": 16874, "min_param_count": 13280257721, "sample_models": [ "lewtun/talkie-1930-13b-it-hf" ], - "relevancy_score": 28.8 + "relevancy_score": 29.7 }, { - "architecture_id": "EvafrillMoForCausalLM", - "total_models": 3, - "total_downloads": 1783, - "min_param_count": 3150567168, + "architecture_id": "TrainableM2MForConditionalGeneration", + "total_models": 1, + "total_downloads": 570, + "min_param_count": 1370638336, "sample_models": [ - "pathcosmos/EVAFRILL-Mo-3B" + "Sunbird/translate-nllb-1.3b-salt" ], - "relevancy_score": 28.6 + "relevancy_score": 29.7 }, { - "architecture_id": "KORMoForCausalLM", - "total_models": 6, - "total_downloads": 7588, - "min_param_count": 10756624384, + "architecture_id": "GENERatorForCausalLM", + "total_models": 1, + "total_downloads": 565, + "min_param_count": 1162061824, "sample_models": [ - "KORMo-Team/KORMo-10B-base", - "KORMo-Team/KORMo-10B-sft" + "GenerTeam/GENERator-v2-prokaryote-1.2b-base" ], - "relevancy_score": 28.5 + "relevancy_score": 29.6 }, { - "architecture_id": "BailingMoeV2_5ForCausalLM", - "total_models": 12, - "total_downloads": 131773, - "min_param_count": 107494409216, + "architecture_id": "HYV3ForCausalLM", + "total_models": 8, + "total_downloads": 367813, + "min_param_count": 30064725888, "sample_models": [ - "inclusionAI/Ring-2.5-1T", - "inclusionAI/Ling-2.6-flash", - "inclusionAI/Ring-2.6-1T", - "inclusionAI/Ling-2.6-1T" + "tencent/Hy3-preview", + "tencent/Hy-MT2-30B-A3B" ], - "relevancy_score": 28.3 + "relevancy_score": 29.3 }, { "architecture_id": "Step3VLForConditionalGeneration", - "total_models": 3, - "total_downloads": 459458, + "total_models": 4, + "total_downloads": 623319, "min_param_count": 320969912832, "sample_models": [ "stepfun-ai/step3" ], - "relevancy_score": 28.3 + "relevancy_score": 29.3 }, { - "architecture_id": "HYV3ForCausalLM", - "total_models": 6, - "total_downloads": 271805, - "min_param_count": 30064725888, + "architecture_id": "LlavaLlamaModel", + "total_models": 15, + "total_downloads": 118909, + "min_param_count": null, "sample_models": [ - "tencent/Hy3-preview", - "tencent/Hy-MT2-30B-A3B" + "Efficient-Large-Model/NVILA-Lite-8B", + "Efficient-Large-Model/VILA1.5-3b", + "Efficient-Large-Model/NVILA-8B", + "Efficient-Large-Model/Llama-3-VILA1.5-8B" ], - "relevancy_score": 28.1 + "relevancy_score": 29.0 }, { - "architecture_id": "JetMoEForCausalLM", + "architecture_id": "TinyLlavaForConditionalGeneration", "total_models": 3, - "total_downloads": 8333, - "min_param_count": 8522237952, + "total_downloads": 1986, + "min_param_count": 3217417280, "sample_models": [ - "jetmoe/jetmoe-8b" + "tinyllava/TinyLLaVA-Phi-2-SigLIP-3.1B" ], - "relevancy_score": 27.9 + "relevancy_score": 28.9 }, { - "architecture_id": "GPTBERTForCausalLM", - "total_models": 24, - "total_downloads": 19723, - "min_param_count": null, + "architecture_id": "JetMoEForCausalLM", + "total_models": 4, + "total_downloads": 11363, + "min_param_count": 8522237952, "sample_models": [ - "BabyLM-community/GPT-BERT-baseline-babylm-eng-nld-zho-equal", - "BabyLM-community/GPT-BERT-baseline-babylm-Strict", - "BabyLM-community/GPT-BERT-baseline-babylm-eng-nld-equal", - "BabyLM-community/GPT-BERT-baseline-babylm-zho", - "BabyLM-community/GPT-BERT-baseline-babylm-nld", - "BabyLM-community/GPT-BERT-baseline-babylm-Strict-Small", - "BabyLM-community/GPT-BERT-baseline-babylm-zho-nld-equal", - "BabyLM-community/GPT-BERT-baseline-babylm-eng-zho-equal" + "jetmoe/jetmoe-8b" ], - "relevancy_score": 27.8 + "relevancy_score": 28.8 }, { "architecture_id": "GPTNeoXJapaneseForCausalLM", - "total_models": 3, - "total_downloads": 356913, + "total_models": 4, + "total_downloads": 469476, "min_param_count": null, "sample_models": [ "abeja/gpt-neox-japanese-2.7b" ], - "relevancy_score": 27.8 + "relevancy_score": 28.7 }, { - "architecture_id": "HyperCLOVAXForCausalLM", + "architecture_id": "EvafrillMoForCausalLM", "total_models": 3, - "total_downloads": 52416, - "min_param_count": 14748112896, + "total_downloads": 1783, + "min_param_count": 3150567168, "sample_models": [ - "naver-hyperclovax/HyperCLOVAX-SEED-Think-14B" + "pathcosmos/EVAFRILL-Mo-3B" ], - "relevancy_score": 27.7 + "relevancy_score": 28.6 }, { - "architecture_id": "LlavaLlamaModel", - "total_models": 12, - "total_downloads": 93982, - "min_param_count": null, + "architecture_id": "DeepseekOCR2ForCausalLM", + "total_models": 3, + "total_downloads": 1688, + "min_param_count": 3389119360, "sample_models": [ - "Efficient-Large-Model/NVILA-Lite-8B", - "Efficient-Large-Model/VILA1.5-3b", - "Efficient-Large-Model/NVILA-8B", - "Efficient-Large-Model/Llama-3-VILA1.5-8B" + "bartsolutions/DeepSeek-OCR-2-GPTQ-INT8" ], - "relevancy_score": 27.6 + "relevancy_score": 28.5 }, { - "architecture_id": "TinyLlavaForConditionalGeneration", - "total_models": 2, - "total_downloads": 1264, - "min_param_count": 3217417280, + "architecture_id": "ChatGLMForConditionalGeneration", + "total_models": 6, + "total_downloads": 6869, + "min_param_count": 9399951392, "sample_models": [ - "tinyllava/TinyLLaVA-Phi-2-SigLIP-3.1B" + "qiuhuachuan/MeChat", + "IAAR-Shanghai/xVerify-9B-C", + "zai-org/LongWriter-glm4-9b" ], - "relevancy_score": 27.6 + "relevancy_score": 28.4 }, { - "architecture_id": "Videollama2MistralForCausalLM", + "architecture_id": "HyperCLOVAXForCausalLM", "total_models": 4, - "total_downloads": 5610, - "min_param_count": 8034428160, + "total_downloads": 61687, + "min_param_count": 14748112896, "sample_models": [ - "DAMO-NLP-SG/VideoLLaMA2-7B-16F", - "DAMO-NLP-SG/VideoLLaMA2-7B" + "naver-hyperclovax/HyperCLOVAX-SEED-Think-14B" ], - "relevancy_score": 27.3 + "relevancy_score": 28.4 }, { "architecture_id": "MiMoV2ASRForCausalLM", - "total_models": 3, - "total_downloads": 6322, + "total_models": 4, + "total_downloads": 8837, "min_param_count": 7622619136, "sample_models": [ "XiaomiMiMo/MiMo-V2.5-ASR" ], - "relevancy_score": 27.3 - }, - { - "architecture_id": "DeepseekOCR2ForCausalLM", - "total_models": 2, - "total_downloads": 1114, - "min_param_count": 3389119360, - "sample_models": [ - "bartsolutions/DeepSeek-OCR-2-GPTQ-INT8" - ], - "relevancy_score": 27.3 + "relevancy_score": 28.3 }, { "architecture_id": "MiniMaxM1ForCausalLM", - "total_models": 6, - "total_downloads": 162568, + "total_models": 8, + "total_downloads": 210297, "min_param_count": 456089655296, "sample_models": [ "MiniMaxAI/MiniMax-M1-40k", "MiniMaxAI/MiniMax-M1-80k" ], - "relevancy_score": 27.0 + "relevancy_score": 28.1 }, { - "architecture_id": "LongcatFlashForCausalLM", - "total_models": 3, - "total_downloads": 231734, - "min_param_count": 561862880256, + "architecture_id": "Cohere2MoeForCausalLM", + "total_models": 10, + "total_downloads": 134015, + "min_param_count": 30484303872, "sample_models": [ - "meituan-longcat/LongCat-Flash-Chat" + "CohereLabs/North-Mini-Code-1.0", + "CohereLabs/North-Mini-Code-1.0-w4a16", + "cyankiwi/North-Mini-Code-1.0-AWQ-INT4" ], - "relevancy_score": 26.9 + "relevancy_score": 27.8 }, { - "architecture_id": "LLaDAMoEModel", - "total_models": 3, - "total_downloads": 4948, - "min_param_count": 7356880896, + "architecture_id": "CogVLMForCausalLM", + "total_models": 8, + "total_downloads": 26881, + "min_param_count": 17639687424, "sample_models": [ - "inclusionAI/LLaDA-MoE-7B-A1B-Base" + "zai-org/cogvlm2-llama3-chat-19B", + "zai-org/cogvlm-chat-hf" ], - "relevancy_score": 26.8 + "relevancy_score": 27.8 }, { - "architecture_id": "ChatGLMForConditionalGeneration", + "architecture_id": "LongcatFlashForCausalLM", "total_models": 4, - "total_downloads": 4171, - "min_param_count": 9399951392, + "total_downloads": 310177, + "min_param_count": 561862880256, "sample_models": [ - "qiuhuachuan/MeChat", - "IAAR-Shanghai/xVerify-9B-C" + "meituan-longcat/LongCat-Flash-Chat" ], - "relevancy_score": 26.7 + "relevancy_score": 27.8 }, { "architecture_id": "Glm4vForConditionalGeneration", - "total_models": 3, - "total_downloads": 4685, + "total_models": 4, + "total_downloads": 6291, "min_param_count": 10357264896, "sample_models": [ "QuantTrio/GLM-4.1V-9B-Thinking-AWQ" ], - "relevancy_score": 26.7 - }, - { - "architecture_id": "CogVLMForCausalLM", - "total_models": 6, - "total_downloads": 20167, - "min_param_count": 17639687424, - "sample_models": [ - "zai-org/cogvlm2-llama3-chat-19B", - "zai-org/cogvlm-chat-hf" - ], - "relevancy_score": 26.6 + "relevancy_score": 27.6 }, { "architecture_id": "MiMoV2FlashForCausalLM", - "total_models": 3, - "total_downloads": 206430, + "total_models": 4, + "total_downloads": 271233, "min_param_count": 309785318400, "sample_models": [ "XiaomiMiMo/MiMo-V2-Flash" ], - "relevancy_score": 26.6 + "relevancy_score": 27.5 }, { - "architecture_id": "ChameleonXLLMXForConditionalGeneration", - "total_models": 3, - "total_downloads": 4313, - "min_param_count": 7013158912, + "architecture_id": "LLaDAMoEModel", + "total_models": 4, + "total_downloads": 6101, + "min_param_count": 7356880896, "sample_models": [ - "Alpha-VLLM/Lumina-mGPT-7B-768" + "inclusionAI/LLaDA-MoE-7B-A1B-Base" ], - "relevancy_score": 26.5 + "relevancy_score": 27.5 }, { - "architecture_id": "PhariaForCausalLM", - "total_models": 3, - "total_downloads": 4378, - "min_param_count": 7041544704, + "architecture_id": "DiffusionGemmaForBlockDiffusion", + "total_models": 8, + "total_downloads": 21747, + "min_param_count": 25823778864, "sample_models": [ - "Aleph-Alpha/Pharia-1-LLM-7B-control-hf" + "aidendle94/diffusiongemma-26B-A4B-it-INT8-dynamic", + "edwixx/diffusiongemma-26B-A4B-it-HERETIC-Uncensored" ], - "relevancy_score": 26.5 + "relevancy_score": 27.4 }, { "architecture_id": "ARCHunyuanVideoForConditionalGeneration", - "total_models": 3, - "total_downloads": 4380, + "total_models": 4, + "total_downloads": 5858, "min_param_count": 8635301872, "sample_models": [ "TencentARC/ARC-Hunyuan-Video-7B" ], - "relevancy_score": 26.5 + "relevancy_score": 27.4 }, { - "architecture_id": "Ministral3ForCausalLM", - "total_models": 1, - "total_downloads": 840, - "min_param_count": 3429006336, + "architecture_id": "PhariaForCausalLM", + "total_models": 4, + "total_downloads": 5660, + "min_param_count": 7041544704, "sample_models": [ - "Kezmark/ErniePEUnleashed" + "Aleph-Alpha/Pharia-1-LLM-7B-control-hf" ], - "relevancy_score": 26.5 + "relevancy_score": 27.4 }, { - "architecture_id": "Videollama2Qwen2ForCausalLM", - "total_models": 3, - "total_downloads": 4048, - "min_param_count": 8527095391, + "architecture_id": "MiniMaxM3SparseForConditionalGeneration", + "total_models": 12, + "total_downloads": 12139, + "min_param_count": 24949042838, "sample_models": [ - "DAMO-NLP-SG/VideoLLaMA2.1-7B-AV" + "JANGQ-AI/MiniMax-M3-REAP40-JANG_2L", + "JANGQ-AI/MiniMax-M3-Medium-JANG_2L", + "aquaman164/MiniMax-M3-AutoRound-3.2bit-longctx", + "JANGQ-AI/MiniMax-M3-REAP32-Coder", + "OsaurusAI/MiniMax-M3-Coder-Small", + "JANGQ-AI/MiniMax-M3-REAP22-Coder" ], - "relevancy_score": 26.3 + "relevancy_score": 27.3 }, { - "architecture_id": "Zaya1VLForConditionalGeneration", - "total_models": 3, - "total_downloads": 4034, - "min_param_count": 9722232312, + "architecture_id": "Videollama2Qwen2ForCausalLM", + "total_models": 4, + "total_downloads": 5445, + "min_param_count": 8527095391, "sample_models": [ - "Zyphra/ZAYA1-VL-8B" + "DAMO-NLP-SG/VideoLLaMA2.1-7B-AV" ], - "relevancy_score": 26.3 + "relevancy_score": 27.3 }, { - "architecture_id": "LizzyForCausalLM", - "total_models": 3, - "total_downloads": 3956, - "min_param_count": 7298011136, + "architecture_id": "ChameleonXLLMXForConditionalGeneration", + "total_models": 4, + "total_downloads": 5605, + "min_param_count": 7013158912, "sample_models": [ - "flwrlabs/Lizzy-7B" + "Alpha-VLLM/Lumina-mGPT-7B-768" ], - "relevancy_score": 26.3 + "relevancy_score": 27.3 }, { - "architecture_id": "DiffusionGemmaForBlockDiffusion", - "total_models": 6, - "total_downloads": 15675, - "min_param_count": 25823778864, + "architecture_id": "MiniCPM3ForCausalLM", + "total_models": 4, + "total_downloads": 224278, + "min_param_count": null, "sample_models": [ - "aidendle94/diffusiongemma-26B-A4B-it-INT8-dynamic", - "edwixx/diffusiongemma-26B-A4B-it-HERETIC-Uncensored" + "openbmb/MiniCPM3-4B" ], - "relevancy_score": 26.1 + "relevancy_score": 27.1 }, { "architecture_id": "ExaoneMoEForCausalLM", - "total_models": 3, - "total_downloads": 157963, + "total_models": 4, + "total_downloads": 209414, "min_param_count": 237099669632, "sample_models": [ "LGAI-EXAONE/K-EXAONE-236B-A23B" ], - "relevancy_score": 26.1 + "relevancy_score": 27.0 }, { - "architecture_id": "Cohere2MoeForCausalLM", - "total_models": 7, - "total_downloads": 87977, - "min_param_count": 30484303872, + "architecture_id": "CheXagentForConditionalGeneration", + "total_models": 4, + "total_downloads": 4538, + "min_param_count": 8362401664, "sample_models": [ - "CohereLabs/North-Mini-Code-1.0", - "CohereLabs/North-Mini-Code-1.0-w4a16", - "cyankiwi/North-Mini-Code-1.0-AWQ-INT4" + "StanfordAIMI/CheXagent-8b" ], - "relevancy_score": 26.0 + "relevancy_score": 26.9 }, { - "architecture_id": "MiniCPM3ForCausalLM", - "total_models": 3, - "total_downloads": 151062, - "min_param_count": null, + "architecture_id": "LizzyForCausalLM", + "total_models": 4, + "total_downloads": 4522, + "min_param_count": 7298011136, "sample_models": [ - "openbmb/MiniCPM3-4B" + "flwrlabs/Lizzy-7B" ], - "relevancy_score": 26.0 + "relevancy_score": 26.9 }, { - "architecture_id": "CheXagentForConditionalGeneration", - "total_models": 3, - "total_downloads": 3440, - "min_param_count": 8362401664, + "architecture_id": "Zaya1VLForConditionalGeneration", + "total_models": 4, + "total_downloads": 4620, + "min_param_count": 9722232312, "sample_models": [ - "StanfordAIMI/CheXagent-8b" + "Zyphra/ZAYA1-VL-8B" ], - "relevancy_score": 26.0 + "relevancy_score": 26.9 }, { "architecture_id": "AquilaForCausalLM", - "total_models": 6, - "total_downloads": 96975, + "total_models": 7, + "total_downloads": 126351, "min_param_count": null, "sample_models": [ "BAAI/AquilaChat2-7B", "BAAI/AquilaChat2-34B-16K", "BAAI/AquilaChat2-70B-Expr" ], - "relevancy_score": 25.9 + "relevancy_score": 26.8 }, { "architecture_id": "HunYuanMoEV1ForCausalLM", - "total_models": 3, - "total_downloads": 146663, + "total_models": 4, + "total_downloads": 196796, "min_param_count": 80393183232, "sample_models": [ "tencent/Hunyuan-A13B-Instruct" ], - "relevancy_score": 25.9 + "relevancy_score": 26.8 }, { - "architecture_id": "UMT5ForConditionalGeneration", + "architecture_id": "ZambaForCausalLM", "total_models": 4, - "total_downloads": 123848, - "min_param_count": null, + "total_downloads": 3926, + "min_param_count": 7232490496, "sample_models": [ - "google/umt5-xxl", - "google/umt5-base", - "google/umt5-small", - "google/umt5-xl" + "Zyphra/Zamba-7B-v1" ], - "relevancy_score": 25.8 + "relevancy_score": 26.6 }, { - "architecture_id": "ZambaForCausalLM", - "total_models": 3, - "total_downloads": 2980, - "min_param_count": 7232490496, + "architecture_id": "Ministral3ForCausalLM", + "total_models": 1, + "total_downloads": 840, + "min_param_count": 3429006336, "sample_models": [ - "Zyphra/Zamba-7B-v1" + "Kezmark/ErniePEUnleashed" ], - "relevancy_score": 25.7 + "relevancy_score": 26.5 }, { - "architecture_id": "Grok1ModelForCausalLM", - "total_models": 3, - "total_downloads": 123910, + "architecture_id": "MoAMetricLM", + "total_models": 16, + "total_downloads": 30359, "min_param_count": null, "sample_models": [ - "hpcai-tech/grok-1" + "reaperdoesntknow/MoA-400M", + "reaperdoesntknow/MoA-150M", + "reaperdoesntknow/MoA-155M", + "reaperdoesntknow/MoA-100M" ], - "relevancy_score": 25.5 + "relevancy_score": 26.4 }, { "architecture_id": "Neutrino", - "total_models": 3, - "total_downloads": 116318, + "total_models": 4, + "total_downloads": 161901, "min_param_count": null, "sample_models": [ "neuralcrew/neutrino-instruct" ], - "relevancy_score": 25.4 + "relevancy_score": 26.4 }, { - "architecture_id": "LatentMoELLaVAPhiForCausalLM", - "total_models": 1, - "total_downloads": 511, - "min_param_count": 3093139456, + "architecture_id": "Grok1ModelForCausalLM", + "total_models": 4, + "total_downloads": 162804, + "min_param_count": null, "sample_models": [ - "KKHYA/llavaphi2-2.7b-finetune-latent-sparse-moe-4e-2k-freeze-1.0_20260304_075653" + "hpcai-tech/grok-1" ], - "relevancy_score": 25.4 + "relevancy_score": 26.4 }, { "architecture_id": "VRT_Qwen2VLForConditionalGeneration", - "total_models": 2, - "total_downloads": 2792, + "total_models": 3, + "total_downloads": 4189, "min_param_count": 8317087232, "sample_models": [ "rp-yu/Qwen2-VL-7b-VPT-CLIP" ], - "relevancy_score": 25.3 - }, - { - "architecture_id": "SarvamMLAForCausalLM", - "total_models": 3, - "total_downloads": 105347, - "min_param_count": 106031767424, - "sample_models": [ - "sarvamai/sarvam-105b" - ], - "relevancy_score": 25.2 + "relevancy_score": 26.4 }, { - "architecture_id": "PLBartForConditionalGeneration", - "total_models": 2, - "total_downloads": 122903, - "min_param_count": null, + "architecture_id": "KimiK25ForConditionalGeneration", + "total_models": 19, + "total_downloads": 18969, + "min_param_count": 38802485184, "sample_models": [ - "uclanlp/plbart-base", - "uclanlp/plbart-java-cs" + "JANGQ-AI/Kimi-K2.6-Small-JANGTQ", + "bearzi/Kimi-K2.7-Code-JANGTQ_K", + "freakyskittle/kimi-k2.75-code", + "JANGQ-AI/Kimi-K2.6-Med-JANGTQ", + "sombra-x/Kimi-K2.6-REAP-Solidity", + "Ex0bit/Kimi-K2.5-PRISM-REAP-530B-A32B" ], - "relevancy_score": 25.2 + "relevancy_score": 26.3 }, { "architecture_id": "PersimmonForCausalLM", - "total_models": 6, - "total_downloads": 66659, + "total_models": 8, + "total_downloads": 89060, "min_param_count": null, "sample_models": [ "adept/persimmon-8b-chat", "adept/persimmon-8b-base" ], - "relevancy_score": 25.1 + "relevancy_score": 26.3 }, { - "architecture_id": "WeDLMForCausalLM", - "total_models": 3, - "total_downloads": 2271, - "min_param_count": 8190735360, + "architecture_id": "SarvamMLAForCausalLM", + "total_models": 4, + "total_downloads": 147076, + "min_param_count": 106031767424, "sample_models": [ - "tencent/WeDLM-8B-Instruct" + "sarvamai/sarvam-105b" ], - "relevancy_score": 25.1 + "relevancy_score": 26.2 }, { - "architecture_id": "FuxiTranyuForCausalLM", - "total_models": 3, - "total_downloads": 2130, - "min_param_count": 8094724096, + "architecture_id": "IQuestLoopCoderForCausalLM", + "total_models": 4, + "total_downloads": 129387, + "min_param_count": 39794696320, "sample_models": [ - "TJUNLP/FuxiTranyu-8B" + "IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct" ], - "relevancy_score": 25.0 + "relevancy_score": 25.9 }, { - "architecture_id": "ZayaForCausalLM", - "total_models": 3, - "total_downloads": 2159, - "min_param_count": 8840489464, + "architecture_id": "AXK1ForCausalLM", + "total_models": 4, + "total_downloads": 125199, + "min_param_count": 518983059456, "sample_models": [ - "Zyphra/ZAYA1-base" + "skt/A.X-K1" ], - "relevancy_score": 25.0 + "relevancy_score": 25.9 }, { - "architecture_id": "AXK1ForCausalLM", - "total_models": 3, - "total_downloads": 94019, - "min_param_count": 518983059456, + "architecture_id": "ArcticForCausalLM", + "total_models": 4, + "total_downloads": 129341, + "min_param_count": 478584608768, "sample_models": [ - "skt/A.X-K1" + "Snowflake/snowflake-arctic-instruct" ], - "relevancy_score": 25.0 + "relevancy_score": 25.9 }, { - "architecture_id": "IQuestLoopCoderForCausalLM", - "total_models": 3, - "total_downloads": 96837, - "min_param_count": 39794696320, + "architecture_id": "UMT5ForConditionalGeneration", + "total_models": 4, + "total_downloads": 123848, + "min_param_count": null, "sample_models": [ - "IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct" + "google/umt5-xxl", + "google/umt5-base", + "google/umt5-small", + "google/umt5-xl" ], - "relevancy_score": 25.0 + "relevancy_score": 25.9 }, { - "architecture_id": "ArcticForCausalLM", + "architecture_id": "KugelAudioForConditionalGeneration", "total_models": 3, - "total_downloads": 96604, - "min_param_count": 478584608768, + "total_downloads": 3188, + "min_param_count": 9343355361, "sample_models": [ - "Snowflake/snowflake-arctic-instruct" + "Roland-JAAI/klonaudio" ], - "relevancy_score": 25.0 + "relevancy_score": 25.9 }, { "architecture_id": "AX4VLForConditionalGeneration", - "total_models": 3, - "total_downloads": 2174, + "total_models": 4, + "total_downloads": 2737, "min_param_count": 7689660144, "sample_models": [ "skt/A.X-4.0-VL-Light" ], - "relevancy_score": 25.0 + "relevancy_score": 25.8 }, { - "architecture_id": "BlenderbotSmallForConditionalGeneration", - "total_models": 2, - "total_downloads": 112562, + "architecture_id": "CLIPT5ForConditionalGeneration", + "total_models": 8, + "total_downloads": 62316, "min_param_count": null, "sample_models": [ - "facebook/blenderbot_small-90M", - "facebook/blenderbot-90M" + "zhiqiulin/clip-flant5-xxl", + "zhiqiulin/clip-flant5-xl" ], - "relevancy_score": 25.0 + "relevancy_score": 25.6 }, { - "architecture_id": "MoAMetricLM", - "total_models": 12, - "total_downloads": 25598, - "min_param_count": null, + "architecture_id": "CambrianLlamaForCausalLM", + "total_models": 4, + "total_downloads": 2430, + "min_param_count": 8333022208, "sample_models": [ - "reaperdoesntknow/MoA-400M", - "reaperdoesntknow/MoA-150M", - "reaperdoesntknow/MoA-155M", - "reaperdoesntknow/MoA-100M" + "nyu-visionx/cambrian-8b" ], - "relevancy_score": 24.9 + "relevancy_score": 25.6 }, { - "architecture_id": "CambrianLlamaForCausalLM", + "architecture_id": "IQuestPLTCoderForCausalLM", "total_models": 3, - "total_downloads": 1916, - "min_param_count": 8333022208, + "total_downloads": 2881, + "min_param_count": 7615749680, "sample_models": [ - "nyu-visionx/cambrian-8b" + "Multilingual-Multimodal-NLP/LoopCoder-V2" ], - "relevancy_score": 24.8 + "relevancy_score": 25.6 }, { - "architecture_id": "CLIPT5ForConditionalGeneration", - "total_models": 6, - "total_downloads": 55162, + "architecture_id": "MiniMaxForCausalLM", + "total_models": 4, + "total_downloads": 106865, + "min_param_count": 456089655296, + "sample_models": [ + "MiniMaxAI/MiniMax-Text-01-hf" + ], + "relevancy_score": 25.5 + }, + { + "architecture_id": "OvisU1", + "total_models": 1, + "total_downloads": 538, + "min_param_count": 3644208367, + "sample_models": [ + "ATH-MaaS/Ovis-U1-3B" + ], + "relevancy_score": 25.5 + }, + { + "architecture_id": "LamedLlamaForCausalLM", + "total_models": 1, + "total_downloads": 523, + "min_param_count": 6983027792, + "sample_models": [ + "GoodBaiBai88/M3D-LaMed-Llama-2-7B" + ], + "relevancy_score": 25.5 + }, + { + "architecture_id": "BlockFFNForCausalLM", + "total_models": 16, + "total_downloads": 19263, "min_param_count": null, "sample_models": [ - "zhiqiulin/clip-flant5-xxl", - "zhiqiulin/clip-flant5-xl" + "SparseLLM/DECO-0.5B", + "SparseLLM/DECO-0.1B", + "SparseLLM/DECO-0.2B", + "SparseLLM/DECO-1.2B" ], - "relevancy_score": 24.7 + "relevancy_score": 25.4 }, { - "architecture_id": "MiniMaxForCausalLM", - "total_models": 3, - "total_downloads": 82647, - "min_param_count": 456089655296, + "architecture_id": "LatentMoELLaVAPhiForCausalLM", + "total_models": 1, + "total_downloads": 511, + "min_param_count": 3093139456, "sample_models": [ - "MiniMaxAI/MiniMax-Text-01-hf" + "KKHYA/llavaphi2-2.7b-finetune-latent-sparse-moe-4e-2k-freeze-1.0_20260304_075653" ], - "relevancy_score": 24.7 + "relevancy_score": 25.4 }, { - "architecture_id": "KugelAudioForConditionalGeneration", + "architecture_id": "PLBartForConditionalGeneration", "total_models": 2, - "total_downloads": 2096, - "min_param_count": 9343355361, + "total_downloads": 122903, + "min_param_count": null, "sample_models": [ - "Roland-JAAI/klonaudio" + "uclanlp/plbart-base", + "uclanlp/plbart-java-cs" ], - "relevancy_score": 24.7 + "relevancy_score": 25.3 }, { - "architecture_id": "KimiK25ForConditionalGeneration", - "total_models": 15, - "total_downloads": 14482, - "min_param_count": 38802485184, + "architecture_id": "WeDLMForCausalLM", + "total_models": 3, + "total_downloads": 2271, + "min_param_count": 8190735360, "sample_models": [ - "JANGQ-AI/Kimi-K2.6-Small-JANGTQ", - "bearzi/Kimi-K2.7-Code-JANGTQ_K", - "freakyskittle/kimi-k2.75-code", - "JANGQ-AI/Kimi-K2.6-Med-JANGTQ", - "sombra-x/Kimi-K2.6-REAP-Solidity", - "Ex0bit/Kimi-K2.5-PRISM-REAP-530B-A32B" + "tencent/WeDLM-8B-Instruct" ], - "relevancy_score": 24.5 + "relevancy_score": 25.1 }, { - "architecture_id": "MagmaForCausalLM", - "total_models": 3, - "total_downloads": 1707, - "min_param_count": 8903066496, + "architecture_id": "BlenderbotSmallForConditionalGeneration", + "total_models": 2, + "total_downloads": 112562, + "min_param_count": null, "sample_models": [ - "microsoft/Magma-8B" + "facebook/blenderbot_small-90M", + "facebook/blenderbot-90M" ], - "relevancy_score": 24.5 + "relevancy_score": 25.1 }, { - "architecture_id": "InternLM2ForRewardModel", - "total_models": 1, - "total_downloads": 94834, - "min_param_count": null, + "architecture_id": "FuxiTranyuForCausalLM", + "total_models": 3, + "total_downloads": 2130, + "min_param_count": 8094724096, "sample_models": [ - "internlm/internlm2_5-step-prover-critic" + "TJUNLP/FuxiTranyu-8B" ], - "relevancy_score": 24.4 + "relevancy_score": 25.0 }, { - "architecture_id": "IQuestPLTCoderForCausalLM", - "total_models": 2, - "total_downloads": 1750, - "min_param_count": 7615749680, + "architecture_id": "ZayaForCausalLM", + "total_models": 3, + "total_downloads": 2159, + "min_param_count": 8840489464, "sample_models": [ - "Multilingual-Multimodal-NLP/LoopCoder-V2" + "Zyphra/ZAYA1-base" ], - "relevancy_score": 24.3 + "relevancy_score": 25.0 }, { - "architecture_id": "MiniMaxM3SparseForConditionalGeneration", - "total_models": 6, - "total_downloads": 6040, - "min_param_count": 29658290213, + "architecture_id": "InternVideo3ForConditionalGeneration", + "total_models": 3, + "total_downloads": 1998, + "min_param_count": 9364294384, "sample_models": [ - "JANGQ-AI/MiniMax-M3-REAP40-JANG_2L", - "JANGQ-AI/MiniMax-M3-Medium-JANG_2L", - "aquaman164/MiniMax-M3-AutoRound-3.2bit-longctx" + "yanziang/InternVideo3-8B-Instruct" ], - "relevancy_score": 24.1 + "relevancy_score": 24.9 }, { - "architecture_id": "BlockFFNForCausalLM", + "architecture_id": "MossForCausalLM", "total_models": 12, - "total_downloads": 16500, + "total_downloads": 23596, "min_param_count": null, "sample_models": [ - "SparseLLM/DECO-0.5B", - "SparseLLM/DECO-0.1B", - "SparseLLM/DECO-0.2B", - "SparseLLM/DECO-1.2B" + "OpenMOSS-Team/moss-moon-003-sft", + "OpenMOSS-Team/moss-moon-003-base", + "OpenMOSS-Team/moss-moon-003-sft-plugin" ], - "relevancy_score": 23.9 + "relevancy_score": 24.7 }, { "architecture_id": "Dots1ForCausalLM", - "total_models": 3, - "total_downloads": 57077, + "total_models": 4, + "total_downloads": 72708, "min_param_count": 142774381696, "sample_models": [ "rednote-hilab/dots.llm1.inst" ], - "relevancy_score": 23.9 + "relevancy_score": 24.7 }, { - "architecture_id": "CogAgentForCausalLM", + "architecture_id": "VARGPTQwen2VLForConditionalGeneration", "total_models": 3, - "total_downloads": 7578, - "min_param_count": 18291854528, + "total_downloads": 1880, + "min_param_count": 10948528387, "sample_models": [ - "zai-org/cogagent-vqa-hf" + "VARGPT-family/VARGPT-v1.1" ], - "relevancy_score": 23.7 + "relevancy_score": 24.7 }, { "architecture_id": "MinistralForCausalLM", - "total_models": 2, - "total_downloads": 1246, + "total_models": 3, + "total_downloads": 1875, "min_param_count": 8237443260, "sample_models": [ "nassimjp/Ministral-8B-Instruct-2410-4bit" ], - "relevancy_score": 23.6 - }, - { - "architecture_id": "VARGPTQwen2VLForConditionalGeneration", - "total_models": 2, - "total_downloads": 1252, - "min_param_count": 10948528387, - "sample_models": [ - "VARGPT-family/VARGPT-v1.1" - ], - "relevancy_score": 23.6 - }, - { - "architecture_id": "SparseLlamaForCausalLM", - "total_models": 2, - "total_downloads": 1223, - "min_param_count": 8185270336, - "sample_models": [ - "openbmb/NOSA-8B", - "openbmb/NOSA-1B" - ], - "relevancy_score": 23.5 + "relevancy_score": 24.7 }, { - "architecture_id": "MossForCausalLM", - "total_models": 9, - "total_downloads": 19408, - "min_param_count": null, + "architecture_id": "MagmaForCausalLM", + "total_models": 3, + "total_downloads": 1707, + "min_param_count": 8903066496, "sample_models": [ - "OpenMOSS-Team/moss-moon-003-sft", - "OpenMOSS-Team/moss-moon-003-base", - "OpenMOSS-Team/moss-moon-003-sft-plugin" + "microsoft/Magma-8B" ], - "relevancy_score": 23.4 + "relevancy_score": 24.5 }, { - "architecture_id": "Qwen2TSForCausalLM", + "architecture_id": "CogAgentForCausalLM", "total_models": 4, - "total_downloads": 37377, - "min_param_count": null, + "total_downloads": 9499, + "min_param_count": 18291854528, "sample_models": [ - "bytedance-research/ChatTS-14B", - "kaluaim/ChatTS-14B-handler" + "zai-org/cogagent-vqa-hf" ], - "relevancy_score": 23.3 + "relevancy_score": 24.4 }, { - "architecture_id": "InternVideo3ForConditionalGeneration", - "total_models": 2, - "total_downloads": 1088, - "min_param_count": 9364294384, + "architecture_id": "InternLM2ForRewardModel", + "total_models": 1, + "total_downloads": 94834, + "min_param_count": null, "sample_models": [ - "yanziang/InternVideo3-8B-Instruct" + "internlm/internlm2_5-step-prover-critic" ], - "relevancy_score": 23.3 + "relevancy_score": 24.4 }, { "architecture_id": "HookedLlamaForCausalLM", - "total_models": 2, - "total_downloads": 1016, + "total_models": 3, + "total_downloads": 1529, "min_param_count": 8030261248, "sample_models": [ "FTK11558/Meta-Llama-3-8B-Instruct-DeepRefusal-Broken-APS" ], - "relevancy_score": 23.1 + "relevancy_score": 24.3 }, { - "architecture_id": "TransfoXLLMHeadModel", - "total_models": 3, - "total_downloads": 33474, + "architecture_id": "Qwen2TSForCausalLM", + "total_models": 6, + "total_downloads": 38480, "min_param_count": null, "sample_models": [ - "transfo-xl/transfo-xl-wt103" + "bytedance-research/ChatTS-14B", + "kaluaim/ChatTS-14B-handler" ], - "relevancy_score": 22.8 + "relevancy_score": 24.0 }, { "architecture_id": "QovaryxForCausalLM", - "total_models": 9, - "total_downloads": 12660, + "total_models": 12, + "total_downloads": 15365, "min_param_count": null, "sample_models": [ "tjarvis91/qovaryx-1b-scratch-base", "tjarvis91/qovaryx-50m-scratch-base", "tjarvis91/qovaryx-350m-scratch-base" ], - "relevancy_score": 22.5 + "relevancy_score": 23.8 }, { "architecture_id": "CpmBeeForCausalLM", - "total_models": 12, - "total_downloads": 8037, + "total_models": 15, + "total_downloads": 9615, "min_param_count": null, "sample_models": [ "openbmb/cpm-bee-10b", @@ -4596,192 +4651,225 @@ "openbmb/cpm-bee-5b", "openbmb/cpm-bee-2b" ], - "relevancy_score": 22.4 + "relevancy_score": 23.7 }, { - "architecture_id": "GPTRefactForCausalLM", + "architecture_id": "Qwen2_5_VLForConditionalGeneration", + "total_models": 12, + "total_downloads": 14087, + "min_param_count": null, + "sample_models": [ + "OmniSVG/OmniSVG1.1_8B", + "OmniSVG/OmniSVG", + "OmniSVG/OmniSVG1.1_4B" + ], + "relevancy_score": 23.6 + }, + { + "architecture_id": "TransfoXLLMHeadModel", "total_models": 4, - "total_downloads": 24366, + "total_downloads": 41546, "min_param_count": null, "sample_models": [ - "refactai/Refact-1_6B-fim", - "refactai/Refact-1_6-base" + "transfo-xl/transfo-xl-wt103" ], - "relevancy_score": 22.4 + "relevancy_score": 23.6 }, { - "architecture_id": "SentinelBrainForCausalLM", - "total_models": 3, - "total_downloads": 4068, - "min_param_count": 14814654656, + "architecture_id": "SparseLlamaForCausalLM", + "total_models": 2, + "total_downloads": 1223, + "min_param_count": 8185270336, "sample_models": [ - "lablab-ai-amd-developer-hackathon/SentinelBrain-14B-MoE-v0.1" + "openbmb/NOSA-8B", + "openbmb/NOSA-1B" ], - "relevancy_score": 22.4 + "relevancy_score": 23.6 }, { "architecture_id": "LISAForCausalLM", - "total_models": 9, - "total_downloads": 11365, + "total_models": 11, + "total_downloads": 14353, "min_param_count": null, "sample_models": [ "xinlai/LISA-7B-v1", "xinlai/LISA-13B-llama2-v1", "MBZUAI/GLaMM-GranD-Pretrained" ], - "relevancy_score": 22.3 + "relevancy_score": 23.4 }, { - "architecture_id": "Qwen2_5_VLForConditionalGeneration", - "total_models": 9, - "total_downloads": 10897, + "architecture_id": "QuasarLongForCausalLM", + "total_models": 4, + "total_downloads": 5749, + "min_param_count": 16553954240, + "sample_models": [ + "silx-ai/Quasar-Preview" + ], + "relevancy_score": 23.4 + }, + { + "architecture_id": "GPTRefactForCausalLM", + "total_models": 5, + "total_downloads": 32335, "min_param_count": null, "sample_models": [ - "OmniSVG/OmniSVG1.1_8B", - "OmniSVG/OmniSVG", - "OmniSVG/OmniSVG1.1_4B" + "refactai/Refact-1_6B-fim", + "refactai/Refact-1_6-base" ], - "relevancy_score": 22.2 + "relevancy_score": 23.3 }, { - "architecture_id": "Florence2ForConditionalGeneration", + "architecture_id": "StripedHyenaModelForCausalLM", "total_models": 8, - "total_downloads": 12586, + "total_downloads": 18895, "min_param_count": null, "sample_models": [ - "onnx-community/Florence-2-base-ft", - "onnx-community/Florence-2-large-ft" + "togethercomputer/evo-1-8k-base", + "togethercomputer/evo-1-131k-base" ], - "relevancy_score": 22.2 + "relevancy_score": 23.1 }, { - "architecture_id": "QuasarLongForCausalLM", - "total_models": 3, - "total_downloads": 3844, - "min_param_count": 16553954240, + "architecture_id": "MobileLlamaForCausalLM", + "total_models": 8, + "total_downloads": 19177, + "min_param_count": null, "sample_models": [ - "silx-ai/Quasar-Preview" + "mtgv/MobileVLM_V2-1.7B", + "mtgv/MobileVLM_V2-3B", + "mtgv/MobileVLM_V2-7B" ], - "relevancy_score": 22.2 + "relevancy_score": 23.1 }, { - "architecture_id": "InternLMXComposer2ForCausalLM", - "total_models": 3, - "total_downloads": 24081, + "architecture_id": "Share4VLlamaForCausalLM", + "total_models": 4, + "total_downloads": 31252, "min_param_count": null, "sample_models": [ - "internlm/internlm-xcomposer2-7b" + "Lin-Chen/ShareGPT4V-7B" ], - "relevancy_score": 22.1 + "relevancy_score": 23.0 }, { - "architecture_id": "InternLMXComposerForCausalLM", - "total_models": 3, - "total_downloads": 23774, - "min_param_count": null, + "architecture_id": "SentinelBrainForCausalLM", + "total_models": 4, + "total_downloads": 4725, + "min_param_count": 14814654656, "sample_models": [ - "internlm/internlm-xcomposer-7b" + "lablab-ai-amd-developer-hackathon/SentinelBrain-14B-MoE-v0.1" ], - "relevancy_score": 22.1 + "relevancy_score": 23.0 }, { - "architecture_id": "GravityMoEForCausalLM", - "total_models": 3, - "total_downloads": 3430, - "min_param_count": 16242181824, + "architecture_id": "InternLMXComposerForCausalLM", + "total_models": 4, + "total_downloads": 31328, + "min_param_count": null, "sample_models": [ - "trillionlabs/Gravity-16B-A3B-Base" + "internlm/internlm-xcomposer-7b" ], - "relevancy_score": 22.0 + "relevancy_score": 23.0 }, { - "architecture_id": "StripedHyenaModelForCausalLM", - "total_models": 6, - "total_downloads": 14327, + "architecture_id": "InternLMXComposer2ForCausalLM", + "total_models": 4, + "total_downloads": 31724, "min_param_count": null, "sample_models": [ - "togethercomputer/evo-1-8k-base", - "togethercomputer/evo-1-131k-base" + "internlm/internlm-xcomposer2-7b" ], - "relevancy_score": 21.9 + "relevancy_score": 23.0 }, { - "architecture_id": "MobileLlamaForCausalLM", - "total_models": 6, - "total_downloads": 14642, + "architecture_id": "MiniMindOmni", + "total_models": 8, + "total_downloads": 16807, "min_param_count": null, "sample_models": [ - "mtgv/MobileVLM_V2-1.7B", - "mtgv/MobileVLM_V2-3B" + "jingyaogong/minimind-3o-moe", + "jingyaogong/minimind-3o" ], - "relevancy_score": 21.9 + "relevancy_score": 22.8 }, { - "architecture_id": "Share4VLlamaForCausalLM", - "total_models": 3, - "total_downloads": 21820, - "min_param_count": null, + "architecture_id": "GravityMoEForCausalLM", + "total_models": 4, + "total_downloads": 4324, + "min_param_count": 16242181824, "sample_models": [ - "Lin-Chen/ShareGPT4V-7B" + "trillionlabs/Gravity-16B-A3B-Base" ], - "relevancy_score": 21.9 + "relevancy_score": 22.8 }, { - "architecture_id": "MiniMindOmni", - "total_models": 6, - "total_downloads": 13618, + "architecture_id": "AILOForCausalLM", + "total_models": 12, + "total_downloads": 9147, "min_param_count": null, "sample_models": [ - "jingyaogong/minimind-3o-moe", - "jingyaogong/minimind-3o" + "xxrickyxx/Ailo152m-v2", + "xxrickyxx/Ailo152m-events-en", + "xxrickyxx/Ailo152m-v2-ita" ], - "relevancy_score": 21.8 + "relevancy_score": 22.7 }, { "architecture_id": "LiquidForCausalLM", - "total_models": 6, - "total_downloads": 13394, + "total_models": 8, + "total_downloads": 15778, "min_param_count": null, "sample_models": [ "reaperdoesntknow/DNA-175M", "reaperdoesntknow/DNA-50M" ], - "relevancy_score": 21.7 + "relevancy_score": 22.7 + }, + { + "architecture_id": "ModernBertDecoderForCausalLM", + "total_models": 4, + "total_downloads": 26856, + "min_param_count": null, + "sample_models": [ + "jhu-clsp/ettin-decoder-400m" + ], + "relevancy_score": 22.6 }, { "architecture_id": "CTRLLMHeadModel", - "total_models": 3, - "total_downloads": 20376, + "total_models": 4, + "total_downloads": 25901, "min_param_count": null, "sample_models": [ "sshleifer/tiny-ctrl" ], - "relevancy_score": 21.7 + "relevancy_score": 22.6 }, { - "architecture_id": "ModernBertDecoderForCausalLM", - "total_models": 3, - "total_downloads": 19997, + "architecture_id": "GPT2LMHeadCustomModel", + "total_models": 4, + "total_downloads": 20128, "min_param_count": null, "sample_models": [ - "jhu-clsp/ettin-decoder-400m" + "bigcode/santacoder" ], - "relevancy_score": 21.7 + "relevancy_score": 22.0 }, { - "architecture_id": "LlavaLladaForMaskedDiffusion", - "total_models": 1, - "total_downloads": 558, - "min_param_count": 8434837024, + "architecture_id": "SolarOpenForCausalLM", + "total_models": 4, + "total_downloads": 19610, + "min_param_count": 102651793408, "sample_models": [ - "KonstantinosKK/lavida-llada-v1.0-instruct-hf-transformers" + "upstage/Solar-Open-100B" ], - "relevancy_score": 21.6 + "relevancy_score": 22.0 }, { "architecture_id": "XverseForCausalLM", - "total_models": 9, - "total_downloads": 7392, + "total_models": 10, + "total_downloads": 8396, "min_param_count": null, "sample_models": [ "xverse/XVERSE-7B-Chat", @@ -4792,680 +4880,668 @@ "xverse/XVERSE-65B", "xverse/XVERSE-MoE-A4.2B" ], - "relevancy_score": 21.4 + "relevancy_score": 21.9 }, { - "architecture_id": "AILOForCausalLM", - "total_models": 9, - "total_downloads": 6734, + "architecture_id": "VulavulaLlamaForCausalLM", + "total_models": 8, + "total_downloads": 10786, "min_param_count": null, "sample_models": [ - "xxrickyxx/Ailo152m-v2", - "xxrickyxx/Ailo152m-events-en", - "xxrickyxx/Ailo152m-v2-ita" + "Papizo/caracal-swahili-math-rigid", + "theophilusowiti/Caracal_instruct" ], - "relevancy_score": 21.2 + "relevancy_score": 21.9 }, { "architecture_id": "LongcatCausalLM", - "total_models": 3, - "total_downloads": 15706, + "total_models": 4, + "total_downloads": 18664, "min_param_count": 561862880256, "sample_models": [ "meituan-longcat/LongCat-Flash-Thinking-2601" ], - "relevancy_score": 21.2 - }, - { - "architecture_id": "SolarOpenForCausalLM", - "total_models": 3, - "total_downloads": 14835, - "min_param_count": 102651793408, - "sample_models": [ - "upstage/Solar-Open-100B" - ], - "relevancy_score": 21.1 - }, - { - "architecture_id": "GPT2LMHeadCustomModel", - "total_models": 3, - "total_downloads": 14230, - "min_param_count": null, - "sample_models": [ - "bigcode/santacoder" - ], - "relevancy_score": 21.0 + "relevancy_score": 21.9 }, { - "architecture_id": "VulavulaLlamaForCausalLM", - "total_models": 6, - "total_downloads": 8827, + "architecture_id": "GPT2Model", + "total_models": 8, + "total_downloads": 9832, "min_param_count": null, "sample_models": [ - "Papizo/caracal-swahili-math-rigid", - "theophilusowiti/Caracal_instruct" + "keshan/sinhala-gpt2", + "cerebras/Cerebras-GPT-13B" ], - "relevancy_score": 20.9 + "relevancy_score": 21.7 }, { - "architecture_id": "Int8OPTForCausalLM", - "total_models": 9, - "total_downloads": 5068, - "min_param_count": null, + "architecture_id": "LlavaLladaForMaskedDiffusion", + "total_models": 1, + "total_downloads": 558, + "min_param_count": 8434837024, "sample_models": [ - "mit-han-lab/opt-125m-smoothquant", - "mit-han-lab/opt-1.3b-smoothquant", - "mit-han-lab/opt-6.7b-smoothquant" + "KonstantinosKK/lavida-llada-v1.0-instruct-hf-transformers" ], - "relevancy_score": 20.6 + "relevancy_score": 21.6 }, { - "architecture_id": "GPT2Model", - "total_models": 6, - "total_downloads": 7698, - "min_param_count": null, + "architecture_id": "Olmo2RetrofitForCausalLM", + "total_models": 1, + "total_downloads": 516, + "min_param_count": 7298011136, "sample_models": [ - "keshan/sinhala-gpt2", - "cerebras/Cerebras-GPT-13B" + "allenai/Olmo-3-7B-RL-Zero-Mix" ], - "relevancy_score": 20.6 + "relevancy_score": 21.4 }, { "architecture_id": "RobertaForCausalLM", - "total_models": 6, - "total_downloads": 5729, + "total_models": 8, + "total_downloads": 7103, "min_param_count": null, "sample_models": [ "uf-aice-lab/math-roberta", "gokceuludogan/ChemBERTaLM" ], - "relevancy_score": 20.0 - }, - { - "architecture_id": "DenseLLM", - "total_models": 4, - "total_downloads": 7575, - "min_param_count": null, - "sample_models": [ - "AlgoDriveAI/Akkadian_English_DenseLLM_1B", - "AlgoDriveAI/Sanskrit_Akkadian_LLM_v1.0" - ], - "relevancy_score": 20.0 + "relevancy_score": 21.0 }, { "architecture_id": "TranceptionLMHeadModel", - "total_models": 6, - "total_downloads": 5328, + "total_models": 8, + "total_downloads": 6753, "min_param_count": null, "sample_models": [ "PascalNotin/Tranception_Large", "PascalNotin/Tranception_Small" ], - "relevancy_score": 19.8 + "relevancy_score": 20.9 }, { "architecture_id": "Qwen35ForCausalLM", - "total_models": 3, - "total_downloads": 7640, + "total_models": 4, + "total_downloads": 10762, "min_param_count": null, "sample_models": [ "JeffGreen311/Eve-V2-Unleashed-Qwen3.5-8B-Liberated-4K-4B-Merged" ], - "relevancy_score": 19.7 + "relevancy_score": 20.7 }, { "architecture_id": "MiniMaxText01ForCausalLM", - "total_models": 3, - "total_downloads": 7676, + "total_models": 4, + "total_downloads": 10676, "min_param_count": 456089655296, "sample_models": [ "MiniMaxAI/MiniMax-Text-01" ], - "relevancy_score": 19.7 + "relevancy_score": 20.7 }, { - "architecture_id": "BigBirdPegasusForConditionalGeneration", - "total_models": 2, - "total_downloads": 8916, + "architecture_id": "Int8OPTForCausalLM", + "total_models": 9, + "total_downloads": 5068, "min_param_count": null, "sample_models": [ - "google/bigbird-pegasus-large-arxiv", - "google/bigbird-pegasus-large-bigpatent" + "mit-han-lab/opt-125m-smoothquant", + "mit-han-lab/opt-1.3b-smoothquant", + "mit-han-lab/opt-6.7b-smoothquant" ], - "relevancy_score": 19.7 + "relevancy_score": 20.6 }, { - "architecture_id": "modeling_camelidae.LlamaForCausalLM", - "total_models": 7, - "total_downloads": 3781, + "architecture_id": "DenseLLM", + "total_models": 5, + "total_downloads": 8991, "min_param_count": null, "sample_models": [ - "hywu/Camelidae-8x7B", - "hywu/Camelidae-8x13B", - "hywu/Camelidae-8x34B" + "AlgoDriveAI/Akkadian_English_DenseLLM_1B", + "AlgoDriveAI/Sanskrit_Akkadian_LLM_v1.0" ], - "relevancy_score": 19.4 + "relevancy_score": 20.6 }, { "architecture_id": "GuppyLM", - "total_models": 6, - "total_downloads": 4503, + "total_models": 8, + "total_downloads": 5620, "min_param_count": null, "sample_models": [ "arman-bd/guppylm-9M", "tamsuiboy/guppylm-9M" ], - "relevancy_score": 19.4 + "relevancy_score": 20.5 }, { - "architecture_id": "ModelStarOLMhead", + "architecture_id": "BailingMoeLinearV2ForCausalLM", "total_models": 3, - "total_downloads": 6584, - "min_param_count": null, + "total_downloads": 1701, + "min_param_count": 16423448320, "sample_models": [ - "C-a-Star-Technology-Official/StarO-AI-2.69-Super" + "inclusionAI/Ring-mini-linear-2.0" ], - "relevancy_score": 19.4 + "relevancy_score": 20.5 }, { - "architecture_id": "BailingMoeLinearV2ForCausalLM", - "total_models": 2, - "total_downloads": 1128, - "min_param_count": 16423448320, + "architecture_id": "ModelStarOLMhead", + "total_models": 4, + "total_downloads": 8035, + "min_param_count": null, "sample_models": [ - "inclusionAI/Ring-mini-linear-2.0" + "C-a-Star-Technology-Official/StarO-AI-2.69-Super" ], - "relevancy_score": 19.4 + "relevancy_score": 20.1 }, { "architecture_id": "TAMELM", - "total_models": 3, - "total_downloads": 6345, + "total_models": 4, + "total_downloads": 7524, "min_param_count": null, "sample_models": [ "reaperdoesntknow/TameForCasualLM" ], - "relevancy_score": 19.3 + "relevancy_score": 20.0 + }, + { + "architecture_id": "MosaicGPT", + "total_models": 4, + "total_downloads": 7506, + "min_param_count": null, + "sample_models": [ + "anas-awadalla/mpt-1b-redpajama-200b" + ], + "relevancy_score": 20.0 + }, + { + "architecture_id": "MptForCausalLM", + "total_models": 4, + "total_downloads": 7296, + "min_param_count": null, + "sample_models": [ + "explosion-testing/mpt-test" + ], + "relevancy_score": 19.9 }, { "architecture_id": "GptOssPuzzleForCausalLM", - "total_models": 3, - "total_downloads": 5918, + "total_models": 4, + "total_downloads": 6851, "min_param_count": 90837823680, "sample_models": [ "nvidia/gpt-oss-puzzle-88B" ], - "relevancy_score": 19.1 + "relevancy_score": 19.8 }, { - "architecture_id": "Beit3LlavaLlamaForCausalLM", - "total_models": 6, - "total_downloads": 3667, - "min_param_count": null, + "architecture_id": "LongcatFlashNgramForCausalLM", + "total_models": 4, + "total_downloads": 6743, + "min_param_count": 69073335552, "sample_models": [ - "openbmb/RLHF-V", - "openbmb/RLHF-V-SFT" + "meituan-longcat/LongCat-Flash-Lite" ], - "relevancy_score": 19.0 + "relevancy_score": 19.7 }, { - "architecture_id": "MosaicGPT", - "total_models": 3, - "total_downloads": 5483, + "architecture_id": "CustomModel", + "total_models": 4, + "total_downloads": 6683, "min_param_count": null, "sample_models": [ - "anas-awadalla/mpt-1b-redpajama-200b" + "DZER-Studios/Create_Vexion-LM" ], - "relevancy_score": 19.0 + "relevancy_score": 19.7 }, - { - "architecture_id": "CustomModel", - "total_models": 3, - "total_downloads": 5512, + { + "architecture_id": "BigBirdPegasusForConditionalGeneration", + "total_models": 2, + "total_downloads": 8916, "min_param_count": null, "sample_models": [ - "DZER-Studios/Create_Vexion-LM" + "google/bigbird-pegasus-large-arxiv", + "google/bigbird-pegasus-large-bigpatent" ], - "relevancy_score": 19.0 + "relevancy_score": 19.7 }, { - "architecture_id": "MptForCausalLM", - "total_models": 3, - "total_downloads": 5334, + "architecture_id": "XMistralForCausalLM", + "total_models": 4, + "total_downloads": 6362, "min_param_count": null, "sample_models": [ - "explosion-testing/mpt-test" + "Hannibal046/xrag-7b" ], - "relevancy_score": 18.9 + "relevancy_score": 19.6 }, { - "architecture_id": "LongcatFlashNgramForCausalLM", - "total_models": 3, - "total_downloads": 5152, - "min_param_count": 69073335552, + "architecture_id": "modeling_camelidae.LlamaForCausalLM", + "total_models": 7, + "total_downloads": 3781, + "min_param_count": null, "sample_models": [ - "meituan-longcat/LongCat-Flash-Lite" + "hywu/Camelidae-8x7B", + "hywu/Camelidae-8x13B", + "hywu/Camelidae-8x34B" ], - "relevancy_score": 18.9 + "relevancy_score": 19.4 }, { - "architecture_id": "XMistralForCausalLM", - "total_models": 3, - "total_downloads": 4790, + "architecture_id": "URMForCausalLM", + "total_models": 4, + "total_downloads": 4910, "min_param_count": null, "sample_models": [ - "Hannibal046/xrag-7b" + "hudsongouge/rrt-40m-wiki-v1" ], - "relevancy_score": 18.7 + "relevancy_score": 19.1 }, { "architecture_id": "FiPhi-NeuralMark-V3", - "total_models": 3, - "total_downloads": 4331, + "total_models": 4, + "total_downloads": 5015, "min_param_count": null, "sample_models": [ "AccTB/FiPhi-NeuralMark-V3" ], - "relevancy_score": 18.5 + "relevancy_score": 19.1 }, { - "architecture_id": "YiForCausalLM", - "total_models": 5, - "total_downloads": 3207, - "min_param_count": 34388917248, + "architecture_id": "Beit3LlavaLlamaForCausalLM", + "total_models": 6, + "total_downloads": 3667, + "min_param_count": null, "sample_models": [ - "llmware/dragon-yi-6b-v0", - "OrionStarAI/OrionStar-Yi-34B-Chat", - "FreedomIntelligence/HuatuoGPT2-34B" + "openbmb/RLHF-V", + "openbmb/RLHF-V-SFT" ], - "relevancy_score": 18.4 + "relevancy_score": 19.0 }, { - "architecture_id": "URMForCausalLM", - "total_models": 3, - "total_downloads": 4236, + "architecture_id": "BTLMLMHeadModel", + "total_models": 4, + "total_downloads": 4768, "min_param_count": null, "sample_models": [ - "hudsongouge/rrt-40m-wiki-v1" + "cerebras/btlm-3b-8k-base" ], - "relevancy_score": 18.4 + "relevancy_score": 19.0 }, { "architecture_id": "NebulaXModel", - "total_models": 3, - "total_downloads": 4196, + "total_models": 4, + "total_downloads": 4864, "min_param_count": null, "sample_models": [ "Agnuxo/NEBULA-X-DEMO" ], - "relevancy_score": 18.4 + "relevancy_score": 19.0 }, { "architecture_id": "QHEARTForECGQA", - "total_models": 3, - "total_downloads": 4022, + "total_models": 4, + "total_downloads": 4752, "min_param_count": null, "sample_models": [ "Manhph2211/Q-HEART" ], - "relevancy_score": 18.3 - }, - { - "architecture_id": "LongcatNextForCausalLM", - "total_models": 3, - "total_downloads": 3963, - "min_param_count": 74257230752, - "sample_models": [ - "meituan-longcat/LongCat-Next" - ], - "relevancy_score": 18.3 + "relevancy_score": 19.0 }, { "architecture_id": "Qwen2ForSequenceClassification", - "total_models": 3, - "total_downloads": 3943, + "total_models": 4, + "total_downloads": 4606, "min_param_count": 71457234944, "sample_models": [ "nvidia/Qwen2.5-CascadeRL-RM-72B" ], - "relevancy_score": 18.3 + "relevancy_score": 18.9 }, { - "architecture_id": "CinnabarLMForCausalLM", - "total_models": 3, - "total_downloads": 3892, + "architecture_id": "YuanForCausalLM", + "total_models": 4, + "total_downloads": 4479, "min_param_count": null, "sample_models": [ - "MihaiPopa-1/CinnabarLM-4M-Base-Preview" + "IEITYuan/Yuan2-M32-hf" ], - "relevancy_score": 18.3 + "relevancy_score": 18.9 }, { - "architecture_id": "SimambaForCausalLM", - "total_models": 3, - "total_downloads": 3894, + "architecture_id": "LongcatNextForCausalLM", + "total_models": 4, + "total_downloads": 4591, + "min_param_count": 74257230752, + "sample_models": [ + "meituan-longcat/LongCat-Next" + ], + "relevancy_score": 18.9 + }, + { + "architecture_id": "Speech2TextTransformerForConditionalGeneration", + "total_models": 4, + "total_downloads": 4489, "min_param_count": null, "sample_models": [ - "soumil1/simamba-midpoint-10m-slimpajama-500m" + "valhalla/s2t_mustc_multilinguial_medium" ], - "relevancy_score": 18.3 + "relevancy_score": 18.9 }, { - "architecture_id": "StreamQwen3_5ForCausalLM", - "total_models": 1, - "total_downloads": 779, - "min_param_count": 26896131584, + "architecture_id": "SimambaForCausalLM", + "total_models": 4, + "total_downloads": 4552, + "min_param_count": null, "sample_models": [ - "JonasGeiping/stream-qwen3.5-27b" + "soumil1/simamba-midpoint-10m-slimpajama-500m" ], - "relevancy_score": 18.3 + "relevancy_score": 18.9 }, { "architecture_id": "GPT2CustomLMHeadModel", - "total_models": 3, - "total_downloads": 3811, + "total_models": 4, + "total_downloads": 4613, "min_param_count": null, "sample_models": [ "fxmarty/tiny-testing-gpt2-remote-code" ], - "relevancy_score": 18.2 + "relevancy_score": 18.9 }, { - "architecture_id": "BTLMLMHeadModel", - "total_models": 3, - "total_downloads": 3727, + "architecture_id": "CinnabarLMForCausalLM", + "total_models": 4, + "total_downloads": 4550, "min_param_count": null, "sample_models": [ - "cerebras/btlm-3b-8k-base" + "MihaiPopa-1/CinnabarLM-4M-Base-Preview" ], - "relevancy_score": 18.2 + "relevancy_score": 18.9 }, { - "architecture_id": "YuanForCausalLM", - "total_models": 3, - "total_downloads": 3827, + "architecture_id": "TFGPT2LMHeadModel", + "total_models": 4, + "total_downloads": 4305, "min_param_count": null, "sample_models": [ - "IEITYuan/Yuan2-M32-hf" + "mymusise/gpt2-medium-chinese" ], - "relevancy_score": 18.2 + "relevancy_score": 18.8 }, { "architecture_id": "CPMAntForCausalLM", - "total_models": 3, - "total_downloads": 3534, + "total_models": 4, + "total_downloads": 4345, "min_param_count": null, "sample_models": [ "openbmb/cpm-ant-10b" ], - "relevancy_score": 18.1 + "relevancy_score": 18.8 }, { - "architecture_id": "TFGPT2LMHeadModel", - "total_models": 3, - "total_downloads": 3487, + "architecture_id": "RoFormerForCausalLM", + "total_models": 4, + "total_downloads": 4260, "min_param_count": null, "sample_models": [ - "mymusise/gpt2-medium-chinese" + "junnyu/roformer_chinese_sim_char_base" ], - "relevancy_score": 18.0 + "relevancy_score": 18.8 }, { - "architecture_id": "Speech2TextTransformerForConditionalGeneration", - "total_models": 3, - "total_downloads": 3509, + "architecture_id": "CodeShell4bitForCausalLM", + "total_models": 4, + "total_downloads": 4065, "min_param_count": null, "sample_models": [ - "valhalla/s2t_mustc_multilinguial_medium" + "WisdomShell/CodeShell-7B-Chat-int4" ], - "relevancy_score": 18.0 + "relevancy_score": 18.7 }, { - "architecture_id": "RoFormerForCausalLM", + "architecture_id": "VStreamLlamaForCausalLM", "total_models": 3, - "total_downloads": 3446, + "total_downloads": 4487, "min_param_count": null, "sample_models": [ - "junnyu/roformer_chinese_sim_char_base" + "IVGSZ/Flash-VStream-7b" ], - "relevancy_score": 18.0 + "relevancy_score": 18.6 }, { - "architecture_id": "NorT5ForConditionalGeneration", - "total_models": 3, - "total_downloads": 3220, - "min_param_count": null, + "architecture_id": "YiForCausalLM", + "total_models": 5, + "total_downloads": 3207, + "min_param_count": 34388917248, "sample_models": [ - "ltg/nort5-base-en-no-translation", - "ltg/nort5-xs", - "ltg/nort5-base" + "llmware/dragon-yi-6b-v0", + "OrionStarAI/OrionStar-Yi-34B-Chat", + "FreedomIntelligence/HuatuoGPT2-34B" ], - "relevancy_score": 17.9 + "relevancy_score": 18.5 }, { - "architecture_id": "CodeShell4bitForCausalLM", - "total_models": 3, - "total_downloads": 3066, + "architecture_id": "ChatUniViLlamaForCausalLM", + "total_models": 4, + "total_downloads": 3843, "min_param_count": null, "sample_models": [ - "WisdomShell/CodeShell-7B-Chat-int4" + "Chat-UniVi/Chat-UniVi-7B-v1.5" ], - "relevancy_score": 17.8 + "relevancy_score": 18.5 }, { - "architecture_id": "BlueLMForCausalLM", - "total_models": 4, - "total_downloads": 2426, + "architecture_id": "A194LogitEnsembleForCausalLM", + "total_models": 3, + "total_downloads": 4359, "min_param_count": null, "sample_models": [ - "vivo-ai/BlueLM-7B-Chat", - "vivo-ai/BlueLM-7B-Base" + "LikC1606/chinesebabylm-a195-hanzi-best-20260611" ], - "relevancy_score": 17.6 + "relevancy_score": 18.5 }, { - "architecture_id": "BottleneckT5LMWithPerturb", + "architecture_id": "MegatronForCausalLM", "total_models": 4, - "total_downloads": 2454, + "total_downloads": 3509, "min_param_count": null, "sample_models": [ - "thesephist/contra-bottleneck-t5-large-wikipedia", - "thesephist/contra-bottleneck-t5-xl-wikipedia" + "hyunwoongko/megatron-11B" ], - "relevancy_score": 17.6 + "relevancy_score": 18.4 }, { - "architecture_id": "darkit-3-1m", + "architecture_id": "Qwen3OmniMoeForConditionalGeneration", "total_models": 3, - "total_downloads": 2897, - "min_param_count": null, + "total_downloads": 4136, + "min_param_count": 35273974321, "sample_models": [ - "darkai-1/darkit-v3-1M" + "88plug/Qwen3-Omni-30B-W4A16" ], - "relevancy_score": 17.6 + "relevancy_score": 18.4 }, { - "architecture_id": "ChatUniViLlamaForCausalLM", - "total_models": 3, - "total_downloads": 2709, + "architecture_id": "BottleneckT5LMWithPerturb", + "total_models": 5, + "total_downloads": 2955, "min_param_count": null, "sample_models": [ - "Chat-UniVi/Chat-UniVi-7B-v1.5" + "thesephist/contra-bottleneck-t5-large-wikipedia", + "thesephist/contra-bottleneck-t5-xl-wikipedia" ], - "relevancy_score": 17.5 + "relevancy_score": 18.3 }, { - "architecture_id": "ElectraForCausalLM", - "total_models": 3, - "total_downloads": 2696, + "architecture_id": "BartForCausalLM", + "total_models": 4, + "total_downloads": 3385, "min_param_count": null, "sample_models": [ - "smeoni/nbme-electra-large-generator" + "sanchit-gandhi/tiny-random-bart-fp16" ], - "relevancy_score": 17.5 + "relevancy_score": 18.3 }, { - "architecture_id": "HNetForCausalLM", - "total_models": 3, - "total_downloads": 2711, + "architecture_id": "ElectraForCausalLM", + "total_models": 4, + "total_downloads": 3369, "min_param_count": null, "sample_models": [ - "emarro/fineweb_hnet_78M_learned_N" + "smeoni/nbme-electra-large-generator" ], - "relevancy_score": 17.5 + "relevancy_score": 18.3 }, { - "architecture_id": "BartForCausalLM", - "total_models": 3, - "total_downloads": 2712, - "min_param_count": null, + "architecture_id": "StreamQwen3_5ForCausalLM", + "total_models": 1, + "total_downloads": 779, + "min_param_count": 26896131584, "sample_models": [ - "sanchit-gandhi/tiny-random-bart-fp16" + "JonasGeiping/stream-qwen3.5-27b" ], - "relevancy_score": 17.5 + "relevancy_score": 18.3 }, { - "architecture_id": "TransnormerForCausalLM", + "architecture_id": "HNetForCausalLM", "total_models": 4, - "total_downloads": 2289, + "total_downloads": 3267, "min_param_count": null, "sample_models": [ - "OpenNLPLab/TransNormerLLM-385M", - "OpenNLPLab/TransNormerLLM-1B" + "emarro/fineweb_hnet_78M_learned_N" ], - "relevancy_score": 17.4 + "relevancy_score": 18.2 }, { "architecture_id": "GeoChatLlamaForCausalLM", - "total_models": 3, - "total_downloads": 2558, + "total_models": 4, + "total_downloads": 3301, "min_param_count": null, "sample_models": [ "MBZUAI/geochat-7B" ], - "relevancy_score": 17.4 + "relevancy_score": 18.2 }, { - "architecture_id": "MegatronForCausalLM", - "total_models": 3, - "total_downloads": 2638, + "architecture_id": "OpenMoeForCausalLM", + "total_models": 4, + "total_downloads": 3250, "min_param_count": null, "sample_models": [ - "hyunwoongko/megatron-11B" + "OrionZheng/openmoe-base" ], - "relevancy_score": 17.4 + "relevancy_score": 18.2 }, { - "architecture_id": "VStreamLlamaForCausalLM", - "total_models": 2, - "total_downloads": 2934, + "architecture_id": "ErnieForCausalLM", + "total_models": 4, + "total_downloads": 3173, "min_param_count": null, "sample_models": [ - "IVGSZ/Flash-VStream-7b" + "mohitsha/tiny-ernie-random-remote-code" ], - "relevancy_score": 17.4 + "relevancy_score": 18.1 }, { - "architecture_id": "KonkanGPT", + "architecture_id": "NGen4OW10TForCausalLM", "total_models": 4, - "total_downloads": 2181, + "total_downloads": 2940, "min_param_count": null, "sample_models": [ - "omdeep22/Gonyai-teo2", - "omdeep22/Gonyai-v1" + "TNSA/NGen-4OW-10T-Expiremental" ], - "relevancy_score": 17.3 + "relevancy_score": 18.0 }, { - "architecture_id": "ErnieForCausalLM", - "total_models": 3, - "total_downloads": 2492, + "architecture_id": "SLMModel", + "total_models": 4, + "total_downloads": 2779, "min_param_count": null, "sample_models": [ - "mohitsha/tiny-ernie-random-remote-code" + "kiruthiga-99/slm-tinystories" ], - "relevancy_score": 17.3 + "relevancy_score": 17.9 }, { - "architecture_id": "OpenMoeForCausalLM", - "total_models": 3, - "total_downloads": 2454, + "architecture_id": "FinanceDecoder", + "total_models": 4, + "total_downloads": 2872, "min_param_count": null, "sample_models": [ - "OrionZheng/openmoe-base" + "tjarvis91/Q-Coder-50M-Sovereign" ], - "relevancy_score": 17.3 + "relevancy_score": 17.9 }, { - "architecture_id": "EmuForCausalLM", - "total_models": 4, - "total_downloads": 2037, + "architecture_id": "NorT5ForConditionalGeneration", + "total_models": 3, + "total_downloads": 3220, "min_param_count": null, "sample_models": [ - "BAAI/Emu2", - "BAAI/Emu2-Chat" + "ltg/nort5-base-en-no-translation", + "ltg/nort5-xs", + "ltg/nort5-base" ], - "relevancy_score": 17.2 + "relevancy_score": 17.9 }, { - "architecture_id": "Qwen3OmniMoeForConditionalGeneration", - "total_models": 2, - "total_downloads": 2744, - "min_param_count": 35273974321, + "architecture_id": "Olmo2ForSequenceClassification", + "total_models": 4, + "total_downloads": 2750, + "min_param_count": null, "sample_models": [ - "88plug/Qwen3-Omni-30B-W4A16" + "LifeWiki-ai/OLMo-2-1124-7B-RM" ], - "relevancy_score": 17.2 + "relevancy_score": 17.8 }, { - "architecture_id": "SLMModel", + "architecture_id": "darkit-3-1m", "total_models": 3, - "total_downloads": 2235, + "total_downloads": 2897, "min_param_count": null, "sample_models": [ - "kiruthiga-99/slm-tinystories" + "darkai-1/darkit-v3-1M" ], - "relevancy_score": 17.1 + "relevancy_score": 17.7 }, { - "architecture_id": "Olmo2ForSequenceClassification", - "total_models": 3, - "total_downloads": 2200, + "architecture_id": "BlueLMForCausalLM", + "total_models": 4, + "total_downloads": 2426, "min_param_count": null, "sample_models": [ - "LifeWiki-ai/OLMo-2-1124-7B-RM" + "vivo-ai/BlueLM-7B-Chat", + "vivo-ai/BlueLM-7B-Base" ], - "relevancy_score": 17.1 + "relevancy_score": 17.6 }, { - "architecture_id": "NGen4OW10TForCausalLM", - "total_models": 3, - "total_downloads": 2196, + "architecture_id": "TransnormerForCausalLM", + "total_models": 4, + "total_downloads": 2289, "min_param_count": null, "sample_models": [ - "TNSA/NGen-4OW-10T-Expiremental" + "OpenNLPLab/TransNormerLLM-385M", + "OpenNLPLab/TransNormerLLM-1B" ], - "relevancy_score": 17.1 + "relevancy_score": 17.5 }, { - "architecture_id": "A194LogitEnsembleForCausalLM", - "total_models": 2, - "total_downloads": 2564, + "architecture_id": "KonkanGPT", + "total_models": 4, + "total_downloads": 2181, "min_param_count": null, "sample_models": [ - "LikC1606/chinesebabylm-a195-hanzi-best-20260611" + "omdeep22/Gonyai-teo2", + "omdeep22/Gonyai-v1" ], - "relevancy_score": 17.1 + "relevancy_score": 17.4 }, { - "architecture_id": "FinanceDecoder", - "total_models": 3, - "total_downloads": 2133, + "architecture_id": "EmuForCausalLM", + "total_models": 4, + "total_downloads": 2037, "min_param_count": null, "sample_models": [ - "tjarvis91/Q-Coder-50M-Sovereign" + "BAAI/Emu2", + "BAAI/Emu2-Chat" ], - "relevancy_score": 17.0 + "relevancy_score": 17.2 }, { "architecture_id": "DarwinDuoOrchestrator", @@ -5478,22 +5554,32 @@ "relevancy_score": 16.9 }, { - "architecture_id": "GeoVForCausalLM", + "architecture_id": "GraphLlamaForCausalLM", "total_models": 3, - "total_downloads": 1736, + "total_downloads": 1741, "min_param_count": null, "sample_models": [ - "GeoV/GeoV-9b" + "Jiabin99/GraphGPT-7B-mix-all" ], "relevancy_score": 16.6 }, { - "architecture_id": "GraphLlamaForCausalLM", + "architecture_id": "RWKV", + "total_models": 3, + "total_downloads": 1773, + "min_param_count": null, + "sample_models": [ + "konpep/aether-rwkv-25m" + ], + "relevancy_score": 16.6 + }, + { + "architecture_id": "YUAZForCausalLM", "total_models": 3, "total_downloads": 1741, "min_param_count": null, "sample_models": [ - "Jiabin99/GraphGPT-7B-mix-all" + "Yuaz-Club/yuaz-mini-cheng-dx-1b" ], "relevancy_score": 16.6 }, @@ -5505,17 +5591,17 @@ "sample_models": [ "CofeAI/Tele-FLM-1T" ], - "relevancy_score": 16.5 + "relevancy_score": 16.6 }, { - "architecture_id": "LingoWhaleForCausalLM", + "architecture_id": "GeoVForCausalLM", "total_models": 3, - "total_downloads": 1698, + "total_downloads": 1736, "min_param_count": null, "sample_models": [ - "deeplang-ai/LingoWhale-8B" + "GeoV/GeoV-9b" ], - "relevancy_score": 16.5 + "relevancy_score": 16.6 }, { "architecture_id": "DuchifatCore", @@ -5525,37 +5611,37 @@ "sample_models": [ "Raziel1234/Duchifat-2" ], - "relevancy_score": 16.5 + "relevancy_score": 16.6 }, { - "architecture_id": "DotsVLMForCausalLM", - "total_models": 2, - "total_downloads": 1382, + "architecture_id": "LingoWhaleForCausalLM", + "total_models": 3, + "total_downloads": 1698, "min_param_count": null, "sample_models": [ - "rednote-hilab/dots.vlm1.inst" + "deeplang-ai/LingoWhale-8B" ], - "relevancy_score": 15.8 + "relevancy_score": 16.5 }, { - "architecture_id": "RWKV", - "total_models": 2, - "total_downloads": 1180, + "architecture_id": "HenlaConfedForCausalLM", + "total_models": 3, + "total_downloads": 1548, "min_param_count": null, "sample_models": [ - "konpep/aether-rwkv-25m" + "RthItalia/HENLA-CONFED-3B-PREFIX-V4-STEP300" ], - "relevancy_score": 15.5 + "relevancy_score": 16.3 }, { - "architecture_id": "YUAZForCausalLM", + "architecture_id": "DotsVLMForCausalLM", "total_models": 2, - "total_downloads": 1156, + "total_downloads": 1382, "min_param_count": null, "sample_models": [ - "Yuaz-Club/yuaz-mini-cheng-dx-1b" + "rednote-hilab/dots.vlm1.inst" ], - "relevancy_score": 15.4 + "relevancy_score": 15.8 }, { "architecture_id": "CrystalCoderLMHeadModel", @@ -5568,24 +5654,24 @@ "relevancy_score": 15.3 }, { - "architecture_id": "HenlaConfedForCausalLM", - "total_models": 2, - "total_downloads": 1028, + "architecture_id": "MoEGPT2", + "total_models": 1, + "total_downloads": 1069, "min_param_count": null, "sample_models": [ - "RthItalia/HENLA-CONFED-3B-PREFIX-V4-STEP300" + "NamrataThakur/Small_Language_Model_MOE_127M_Pretrained" ], - "relevancy_score": 15.2 + "relevancy_score": 15.0 }, { - "architecture_id": "MoEGPT2", + "architecture_id": "LlavaStableLMEpochForCausalLM", "total_models": 1, - "total_downloads": 1069, + "total_downloads": 822, "min_param_count": null, "sample_models": [ - "NamrataThakur/Small_Language_Model_MOE_127M_Pretrained" + "NousResearch/Obsidian-3B-V0.5" ], - "relevancy_score": 15.0 + "relevancy_score": 14.4 }, { "architecture_id": "JetNemotronForCausalLM", @@ -5598,14 +5684,14 @@ "relevancy_score": 14.4 }, { - "architecture_id": "LlavaStableLMEpochForCausalLM", + "architecture_id": "XModelForCausalLM", "total_models": 1, - "total_downloads": 822, + "total_downloads": 691, "min_param_count": null, "sample_models": [ - "NousResearch/Obsidian-3B-V0.5" + "XiaoduoAILab/Xmodel_LM" ], - "relevancy_score": 14.4 + "relevancy_score": 14.1 }, { "architecture_id": "JapaneseStableLMAlphaForCausalLM", @@ -5627,16 +5713,6 @@ ], "relevancy_score": 14.1 }, - { - "architecture_id": "XModelForCausalLM", - "total_models": 1, - "total_downloads": 691, - "min_param_count": null, - "sample_models": [ - "XiaoduoAILab/Xmodel_LM" - ], - "relevancy_score": 14.0 - }, { "architecture_id": "YayiForCausalLM", "total_models": 1, @@ -5668,32 +5744,32 @@ "relevancy_score": 13.8 }, { - "architecture_id": "OpenBAForConditionalGeneration", + "architecture_id": "D3PMSanskritModel", "total_models": 1, - "total_downloads": 594, + "total_downloads": 592, "min_param_count": null, "sample_models": [ - "OpenNLG/OpenBA-V1-Based" + "bhsinghgrid/sanskrit-translation" ], "relevancy_score": 13.7 }, { - "architecture_id": "D3PMSanskritModel", + "architecture_id": "OpenBAForConditionalGeneration", "total_models": 1, - "total_downloads": 592, + "total_downloads": 594, "min_param_count": null, "sample_models": [ - "bhsinghgrid/sanskrit-translation" + "OpenNLG/OpenBA-V1-Based" ], "relevancy_score": 13.7 }, { - "architecture_id": "GQAGPT2", + "architecture_id": "GPTBigCodeLMHeadModel", "total_models": 1, - "total_downloads": 560, + "total_downloads": 547, "min_param_count": null, "sample_models": [ - "NamrataThakur/Small_Language_Model_GQA_48M_Pretrained" + "bigcode/santacoderpack" ], "relevancy_score": 13.6 }, @@ -5708,52 +5784,52 @@ "relevancy_score": 13.6 }, { - "architecture_id": "GPTBigCodeLMHeadModel", + "architecture_id": "GQAGPT2", "total_models": 1, - "total_downloads": 547, + "total_downloads": 560, "min_param_count": null, "sample_models": [ - "bigcode/santacoderpack" + "NamrataThakur/Small_Language_Model_GQA_48M_Pretrained" ], "relevancy_score": 13.6 }, { - "architecture_id": "GPT2", + "architecture_id": "SeerAttnLlamaForCausalLM", "total_models": 1, - "total_downloads": 550, + "total_downloads": 559, "min_param_count": null, "sample_models": [ - "NamrataThakur/Small_Language_Model_MHA_53M_Pretrained" + "SeerAttention/SeerAttention-Llama-3.1-8B-AttnGates" ], "relevancy_score": 13.6 }, { - "architecture_id": "SeerAttnLlamaForCausalLM", + "architecture_id": "GPT2", "total_models": 1, - "total_downloads": 559, + "total_downloads": 550, "min_param_count": null, "sample_models": [ - "SeerAttention/SeerAttention-Llama-3.1-8B-AttnGates" + "NamrataThakur/Small_Language_Model_MHA_53M_Pretrained" ], "relevancy_score": 13.6 }, { - "architecture_id": "CoherenceMomentumModel", + "architecture_id": "LlaMAForCausalLM", "total_models": 1, - "total_downloads": 523, + "total_downloads": 545, "min_param_count": null, "sample_models": [ - "aisingapore/coherence-momentum" + "circulus/alpaca-7b" ], - "relevancy_score": 13.5 + "relevancy_score": 13.6 }, { - "architecture_id": "LlaMAForCausalLM", + "architecture_id": "CoherenceMomentumModel", "total_models": 1, - "total_downloads": 545, + "total_downloads": 523, "min_param_count": null, "sample_models": [ - "circulus/alpaca-7b" + "aisingapore/coherence-momentum" ], "relevancy_score": 13.5 }, @@ -5768,22 +5844,22 @@ "relevancy_score": 13.4 }, { - "architecture_id": "HunYuanForCausalLM", + "architecture_id": "GLaMMForCausalLM", "total_models": 1, - "total_downloads": 513, + "total_downloads": 512, "min_param_count": null, "sample_models": [ - "tencent/Hunyuan-7B-Pretrain-0124" + "MBZUAI/GLaMM-FullScope" ], "relevancy_score": 13.4 }, { - "architecture_id": "GLaMMForCausalLM", + "architecture_id": "HunYuanForCausalLM", "total_models": 1, - "total_downloads": 512, + "total_downloads": 513, "min_param_count": null, "sample_models": [ - "MBZUAI/GLaMM-FullScope" + "tencent/Hunyuan-7B-Pretrain-0124" ], "relevancy_score": 13.4 }, diff --git a/transformer_lens/tools/model_registry/data/supported_models.json b/transformer_lens/tools/model_registry/data/supported_models.json index a01926111..b82100137 100644 --- a/transformer_lens/tools/model_registry/data/supported_models.json +++ b/transformer_lens/tools/model_registry/data/supported_models.json @@ -1,29 +1,15 @@ { - "generated_at": "2026-06-25", + "generated_at": "2026-07-01", "scan_info": { - "total_scanned": 5949, + "total_scanned": 5855, "task_filter": "text-generation", "min_downloads": 500, - "scan_duration_seconds": 8.1 + "scan_duration_seconds": 7.0 }, - "total_architectures": 66, - "total_models": 12900, - "total_verified": 1017, + "total_architectures": 68, + "total_models": 13138, + "total_verified": 1028, "models": [ - { - "architecture_id": "Qwen2MoeForCausalLM", - "model_id": "hyper-accel/ci-random-qwen2-moe-a3b", - "status": 0, - "verified_date": null, - "metadata": null, - "note": null, - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": null, - "phase7_score": null, - "phase8_score": null - }, { "architecture_id": "HunYuanDenseV1ForCausalLM", "model_id": "tencent/Hunyuan-0.5B-Instruct", @@ -3918,14 +3904,14 @@ { "architecture_id": "Mamba2ForCausalLM", "model_id": "AntonV/mamba2-130m-hf", - "status": 2, - "verified_date": "2026-04-10", + "status": 1, + "verified_date": "2026-07-01", "metadata": null, - "note": "Architecture Mamba2ForCausalLM has applicable_phases=[]; verify_models coverage is deferred. Verification lives in integration tests.", - "phase1_score": 0.0, - "phase2_score": null, - "phase3_score": null, - "phase4_score": 91.6, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 93.3, "phase7_score": null, "phase8_score": null }, @@ -14289,28 +14275,28 @@ { "architecture_id": "MambaForCausalLM", "model_id": "EchoLabs33/mamba-130m-hxq", - "status": 1, - "verified_date": "2026-06-26", + "status": 3, + "verified_date": "2026-07-01", "metadata": null, - "note": "Full verification completed", - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": 51.5, + "note": "Below threshold: P1=0.0% < 100.0% (failed: all_components, forward_pass_logits) \u2014 24/51 components failed (24 critical)", + "phase1_score": 0.0, + "phase2_score": 100.0, + "phase3_score": 90.0, + "phase4_score": 51.0, "phase7_score": null, "phase8_score": null }, { "architecture_id": "Mamba2ForCausalLM", "model_id": "EchoLabs33/mamba2-1.3b-hxq", - "status": 1, - "verified_date": "2026-06-26", + "status": 3, + "verified_date": "2026-07-01", "metadata": null, - "note": "Full verification completed with issues, low text quality", - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": 31.3, + "note": "Below threshold: P1=0.0% < 100.0% (failed: all_components, forward_pass_logits) \u2014 50/99 components failed (50 critical)", + "phase1_score": 0.0, + "phase2_score": 100.0, + "phase3_score": 90.0, + "phase4_score": 28.2, "phase7_score": null, "phase8_score": null }, @@ -19442,14 +19428,14 @@ { "architecture_id": "Qwen3_5ForCausalLM", "model_id": "Evelyn67/Qwen3.5-2B-Her", - "status": 0, - "verified_date": null, + "status": 1, + "verified_date": "2026-07-01", "metadata": null, - "note": null, - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": null, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 96.8, "phase7_score": null, "phase8_score": null }, @@ -21620,28 +21606,28 @@ { "architecture_id": "Qwen3_5ForCausalLM", "model_id": "GoodStartLabs/gin-rummy-hbc-qwen3.5-0.8b", - "status": 0, - "verified_date": null, + "status": 1, + "verified_date": "2026-07-01", "metadata": null, - "note": null, - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": null, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 97.7, "phase7_score": null, "phase8_score": null }, { "architecture_id": "Qwen3_5ForCausalLM", "model_id": "GoodStartLabs/gin-rummy-hbc-qwen3.5-2b", - "status": 0, - "verified_date": null, + "status": 1, + "verified_date": "2026-07-01", "metadata": null, - "note": null, - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": null, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 98.4, "phase7_score": null, "phase8_score": null }, @@ -34584,13 +34570,13 @@ "architecture_id": "MambaForCausalLM", "model_id": "NYTK/PULI-HuBA-mamba-130M", "status": 1, - "verified_date": "2026-06-26", + "verified_date": "2026-07-01", "metadata": null, "note": "Full verification completed", - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": 62.3, + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 67.4, "phase7_score": null, "phase8_score": null }, @@ -59276,13 +59262,13 @@ "architecture_id": "GraniteMoeHybridForCausalLM", "model_id": "WithinUsAI/IBM-Grok4-Ultra.Fast.Coder-1B", "status": 1, - "verified_date": "2026-06-26", + "verified_date": "2026-07-01", "metadata": null, "note": "Full verification completed", "phase1_score": 100.0, "phase2_score": 100.0, "phase3_score": 100.0, - "phase4_score": 99.1, + "phase4_score": 92.3, "phase7_score": null, "phase8_score": null }, @@ -61130,14 +61116,14 @@ { "architecture_id": "Qwen3_5ForCausalLM", "model_id": "ZirTech/OmniMath-2B", - "status": 0, - "verified_date": null, + "status": 1, + "verified_date": "2026-07-01", "metadata": null, - "note": null, - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": null, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 95.0, "phase7_score": null, "phase8_score": null }, @@ -79515,14 +79501,14 @@ { "architecture_id": "Qwen3_5ForCausalLM", "model_id": "continuum-ai/qwen3.5-0.8b-general-forged", - "status": 0, - "verified_date": null, + "status": 1, + "verified_date": "2026-07-01", "metadata": null, - "note": null, - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": null, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 93.3, "phase7_score": null, "phase8_score": null }, @@ -82616,13 +82602,13 @@ "architecture_id": "Mamba2ForCausalLM", "model_id": "deqing/convergent-mamba2-300M-adamw-original", "status": 1, - "verified_date": "2026-06-26", + "verified_date": "2026-07-01", "metadata": null, "note": "Full verification completed with issues, low text quality", - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": 8.2, + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 12.3, "phase7_score": null, "phase8_score": null }, @@ -83075,13 +83061,13 @@ "architecture_id": "Mamba2ForCausalLM", "model_id": "deqing/mamba2-300M-v5-mamba2", "status": 1, - "verified_date": "2026-06-26", + "verified_date": "2026-07-01", "metadata": null, "note": "Full verification completed", - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": 91.0, + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 94.2, "phase7_score": null, "phase8_score": null }, @@ -84220,13 +84206,13 @@ "architecture_id": "MambaForCausalLM", "model_id": "dominguesm/mambarim-110m", "status": 1, - "verified_date": "2026-06-26", + "verified_date": "2026-07-01", "metadata": null, "note": "Full verification completed", - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": 73.2, + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 73.3, "phase7_score": null, "phase8_score": null }, @@ -102055,7 +102041,7 @@ "architecture_id": "GraniteMoeHybridForCausalLM", "model_id": "ibm-granite/granite-4.0-1b", "status": 1, - "verified_date": "2026-04-15", + "verified_date": "2026-07-01", "metadata": null, "note": "Full verification completed", "phase1_score": 100.0, @@ -102069,13 +102055,13 @@ "architecture_id": "GraniteMoeHybridForCausalLM", "model_id": "ibm-granite/granite-4.0-1b-base", "status": 1, - "verified_date": "2026-06-26", + "verified_date": "2026-07-01", "metadata": null, "note": "Full verification completed", "phase1_score": 100.0, "phase2_score": 100.0, "phase3_score": 100.0, - "phase4_score": 95.6, + "phase4_score": 97.7, "phase7_score": null, "phase8_score": null }, @@ -102083,11 +102069,11 @@ "architecture_id": "GraniteMoeHybridForCausalLM", "model_id": "ibm-granite/granite-4.0-350m", "status": 1, - "verified_date": "2026-04-15", + "verified_date": "2026-07-01", "metadata": null, - "note": "Full verification completed with issues: P2=91.7% (failed: generation)", + "note": "Full verification completed", "phase1_score": 100.0, - "phase2_score": 91.7, + "phase2_score": 100.0, "phase3_score": 100.0, "phase4_score": 94.7, "phase7_score": null, @@ -102292,11 +102278,11 @@ { "architecture_id": "GraniteMoeHybridForCausalLM", "model_id": "ibm-granite/granite-4.0-tiny-preview", - "status": 3, - "verified_date": "2026-04-15", + "status": 1, + "verified_date": "2026-07-01", "metadata": null, - "note": "Below threshold: P1=50.0% < 100.0% (failed: all_components) \u2014 72/347 components failed (72 critical)", - "phase1_score": 50.0, + "note": "Full verification completed", + "phase1_score": 100.0, "phase2_score": 100.0, "phase3_score": 100.0, "phase4_score": 98.7, @@ -112705,13 +112691,13 @@ "architecture_id": "Mamba2ForCausalLM", "model_id": "limloop/whiff-mamba2-50M-v0.1", "status": 1, - "verified_date": "2026-06-26", + "verified_date": "2026-07-01", "metadata": null, "note": "Full verification completed", - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": 80.5, + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 79.2, "phase7_score": null, "phase8_score": null }, @@ -127903,7 +127889,7 @@ "architecture_id": "Qwen3NextForCausalLM", "model_id": "optimum-intel-internal-testing/tiny-random-qwen3-next", "status": 1, - "verified_date": "2026-04-15", + "verified_date": "2026-07-01", "metadata": null, "note": "Full verification completed", "phase1_score": 100.0, @@ -138993,13 +138979,13 @@ { "architecture_id": "MambaForCausalLM", "model_id": "state-spaces/mamba-130m-hf", - "status": 2, - "verified_date": "2026-06-05", - "note": "Architecture MambaForCausalLM has applicable_phases=[]; verify_models coverage is deferred. Verification lives in integration tests.", - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": 91.6, + "status": 1, + "verified_date": "2026-07-01", + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 95.8, "phase7_score": null, "phase8_score": null }, @@ -139021,13 +139007,13 @@ "architecture_id": "MambaForCausalLM", "model_id": "state-spaces/mamba-370m-hf", "status": 1, - "verified_date": "2026-06-26", + "verified_date": "2026-07-01", "metadata": null, "note": "Full verification completed", - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": 89.0, + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 94.8, "phase7_score": null, "phase8_score": null }, @@ -139035,13 +139021,13 @@ "architecture_id": "MambaForCausalLM", "model_id": "state-spaces/mamba-790m-hf", "status": 1, - "verified_date": "2026-06-26", + "verified_date": "2026-07-01", "metadata": null, "note": "Full verification completed", - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": 97.6, + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 91.5, "phase7_score": null, "phase8_score": null }, @@ -141946,7 +141932,7 @@ "architecture_id": "Qwen3NextForCausalLM", "model_id": "tiny-random/qwen3-next-moe", "status": 1, - "verified_date": "2026-04-15", + "verified_date": "2026-07-01", "metadata": null, "note": "Full verification completed", "phase1_score": 100.0, @@ -154975,7 +154961,7 @@ "architecture_id": "Qwen3NextForCausalLM", "model_id": "yujiepan/qwen3-next-moe-tiny-random", "status": 1, - "verified_date": "2026-04-15", + "verified_date": "2026-07-01", "metadata": null, "note": "Full verification completed", "phase1_score": 100.0, @@ -159107,13 +159093,13 @@ "architecture_id": "Mamba2ForCausalLM", "model_id": "evolai/evolai_mamba_naive_kl", "status": 1, - "verified_date": "2026-06-26", + "verified_date": "2026-07-01", "metadata": null, "note": "Full verification completed", - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": 76.3, + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 77.6, "phase7_score": null, "phase8_score": null }, @@ -179221,6 +179207,3324 @@ "phase7_score": null, "phase8_score": null }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16", + "status": 1, + "verified_date": "2026-07-01", + "metadata": null, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 84.1, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-Base-BF16", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "nvidia/Nemotron-Cascade-2-30B-A3B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "nvidia/NVIDIA-Nemotron-Labs-3-Elastic-30B-A3B-BF16", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "nvidia/NVIDIA-Nemotron-Nano-9B-v2", + "status": 1, + "verified_date": "2026-07-02", + "metadata": null, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 90.1, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "nvidia/Nemotron-Elastic-12B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "nvidia/NVIDIA-Nemotron-Nano-9B-v2-Japanese", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-Base-BF16", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-GenRM", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Gemma4ForConditionalGeneration", + "model_id": "farbodtavakkoli/OTel-LLM-E4B-IT", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Lfm2MoeForCausalLM", + "model_id": "farbodtavakkoli/OTel-LLM-8B-A1B-IT", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "trl-internal-testing/tiny-NemotronHForCausalLM-nano", + "status": 1, + "verified_date": "2026-07-01", + "metadata": null, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 72.1, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "trl-internal-testing/tiny-NemotronHForCausalLM-ultra", + "status": 1, + "verified_date": "2026-07-01", + "metadata": null, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 71.9, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "trl-internal-testing/tiny-NemotronHForCausalLM-super", + "status": 1, + "verified_date": "2026-07-01", + "metadata": null, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 72.3, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3_5ForConditionalGeneration", + "model_id": "deepreinforce-ai/Ornith-1.0-9B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "tencent/HY-MT1.5-1.8B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "tencent/Hy-MT2-7B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "tencent/Hy-MT2-1.8B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "axana-labs/Llama-3.3-70B-Instruct-bnb-4bit", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "tencent/Hunyuan-7B-Instruct", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "unsloth/NVIDIA-Nemotron-3-Nano-4B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2ForCausalLM", + "model_id": "xinnn32/Qwen2.5-0.5B-Instruct-Gensyn-Swarm-amphibious_savage_mantis", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "c4tdr0ut/grok-oss-Revenant-70B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "shawnw3i/Hy-MT2-1.8B-AWQ", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2ForCausalLM", + "model_id": "ophirparwez/Qwen2.5-0.5B-Instruct-Gensyn-Swarm-invisible_mammalian_worm", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "Soofi-Project/Soofi-S-Isar-Preview", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "Soofi-Project/Soofi-S-Instruct-Preview", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "uzlm/alloma-3B-Instruct", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "shawnw3i/Hy-MT2-7B-AWQ", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "dealignai/Nemotron-3-Super-120B-A12B-UNCENSORED-JANG_2L", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "tencent/Hunyuan-MT-7B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3_5ForConditionalGeneration", + "model_id": "cyankiwi/Ornith-1.0-9B-AWQ-INT4", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "tencent/HY-MT1.5-7B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3ForCausalLM", + "model_id": "open-thoughts/OpenThinkerAgent-32B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "unsloth/NVIDIA-Nemotron-3-Super-120B-A12B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2ForCausalLM", + "model_id": "MihailBazarov/qwen2.5-coder-14b-bitrix-developer", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "RedHatAI/NVIDIA-Nemotron-3-Ultra-550B-A55B-quantized.w4a16", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "NeTS-lab/babylm26_multiling_gpt2_MoP16K_baseline", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3_5ForConditionalGeneration", + "model_id": "cpatonn/Qwopus3.6-27B-Coder-AWQ-INT4", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "tencent/HY-MT1.5-1.8B-GPTQ-Int4", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "DavidAU/NVIDIA-Nemotron-Labs-3-Elastic-12B-A2B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "sagnikM/hill_8k_300", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "nwzjk/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16-W4A16-G128", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2ForCausalLM", + "model_id": "jakiAJK/DeepSeek-R1-Distill-Qwen-1.5B_GPTQ-int4", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2ForCausalLM", + "model_id": "ATH-MaaS/Marco-LLM-ES", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "trohrbaugh/NVIDIA-Nemotron-3-Super-120B-A12B-BF16-heretic-v2", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "trohrbaugh/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16-heretic", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "OpenResearcher/OpenResearcher-30B-A3B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "manjunathshiva/Nemotron-3-Super-120B-A12B-tq3a-tq2e-g32", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3_5ForConditionalGeneration", + "model_id": "Merlin-Research/Qwen3.5-4B-Safety-Thinking", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "manjunathshiva/Nemotron-3-Super-120B-A12B-tq3", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Gemma4ForConditionalGeneration", + "model_id": "MobiusDevelopment/gemma4-E2B-it-qat-Q4-unsloth-heretic", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "RedHatAI/NVIDIA-Nemotron-3-Super-120B-A12B-BF16", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "HebArabNlpProject/Hebatron_base", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "locailabs/Jupiter-N-120B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2ForCausalLM", + "model_id": "BillyWang1/qwen2.5-7b-base-retool-slime-sft-v2", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "zai-org/LongWriter-llama3.1-8b", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3_5ForConditionalGeneration", + "model_id": "XReyRobert/Qwopus3.6-27B-Coder-GPTQ-Pro", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/swa-latn-10mb-100mb_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/swa-latn-100mb-after-ppt-shuff-dyck-100mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/swa-latn-100mb-after-ppt-Dp-100mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "OpenLLM-France/Luciole-8B-Base", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "tencent/Hunyuan-7B-Pretrain", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/ita-latn-10mb-10mb_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "ReadyArt/L3.3-The-Omega-Directive-70B-Unslop-v2.1", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Gemma4UnifiedForConditionalGeneration", + "model_id": "Rangle2/gemma-4-12B-it-uncensored-opus4.7-cot", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2ForCausalLM", + "model_id": "Kortix/FastApply-1.5B-v1.0", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/swa-latn-10mb-after-ppt-shuff-dyck-100mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2ForCausalLM", + "model_id": "Bilns/Qwen2.5-0.5B-Instruct-Gensyn-Swarm-mute_sedate_cobra", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/swa-latn-100mb-after-ppt-Dp-10mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "Recor2d/GPT2-ChineseDevBench", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Gemma2ForCausalLM", + "model_id": "Ftm23/cbd-gemma2-4pair", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/swa-latn-10mb-after-ppt-Dp-100mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Gemma3ForCausalLM", + "model_id": "TheKloakedSignal/GoblinGLMV02", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/ita-latn-100mb-after-ppt-Dp-100mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3ForCausalLM", + "model_id": "ermiaazarkhalili/Qwen3-4B-SFT-Fable5-Glint", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2ForCausalLM", + "model_id": "hayulalab/Averroes-Q-Instruct", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "RedHatAI/NVIDIA-Nemotron-Nano-9B-v2-quantized.w4a16", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/swa-latn-10mb-after-ppt-Dp-10mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3ForCausalLM", + "model_id": "clijo/qwen3-4b-instruct-2507-bf16-reco-grpo-b200-rapid-orange-quartz", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/zho-hans-10mb-after-ppt-shuff-dyck-100mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/swa-latn-100mb-after-ppt-shuff-dyck-10mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3MoeForCausalLM", + "model_id": "ThaiLLM/ThaiLLM-30B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "WestCode1357/gpt-sw3-126m-instruct", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/nld-latn-100mb-after-ppt-shuff-dyck-100mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Gemma3ForConditionalGeneration", + "model_id": "GalvinNguyen/vian_ai_shop_small", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "dnotitia/Llama-DNA-1.0-8B-Instruct", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "mii-llm/nesso-0.4B-agentic", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3ForCausalLM", + "model_id": "OpenOneRec/OneReason-0.8B-pretrain-competition", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "T5ForConditionalGeneration", + "model_id": "razent/SciFive-large-Pubmed_PMC", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2ForCausalLM", + "model_id": "onnx-community/Qwen2.5-Coder-3B-Instruct", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/ita-latn-100mb-after-ppt-shuff-dyck-100mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "ApertusForCausalLM", + "model_id": "EPFLiGHT/Apertus-8B-MeditronFO", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "HebArabNlpProject/Hebatron", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3_5ForConditionalGeneration", + "model_id": "CloudGoat/Qwen3.5-0.8B-JP-Tuned-v2.0", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GraniteForCausalLM", + "model_id": "treadon/granite-4.1-8b-Abliterated-AND-Disinhibited", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "Delentia/delentia-slm-jitna-v0.4", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "MistralForCausalLM", + "model_id": "TheBloke/OpenHermes-2.5-Mistral-7B-GPTQ", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "raulgdp/gpt-acredita-350m", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2ForCausalLM", + "model_id": "Multiverse4FM/Multiverse-32B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/nld-latn-100mb-after-ppt-Dp-100mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3_5ForConditionalGeneration", + "model_id": "shamsghi/Mistral-Le-Chaton-Fat", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "Raghav-Singhal/pretrain-normal-smollm-1p7b-100B-20n-2048sl-960gbsz", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3MoeForCausalLM", + "model_id": "unsloth/Qwen3-30B-A3B-Thinking-2507", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/ita-latn-10mb-after-ppt-Dp-100mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "hirundo-io/NVIDIA-Nemotron-3-Nano-4B-BF16", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3ForCausalLM", + "model_id": "m-a-p/OProver-32B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/swa-latn-100mb-10mb_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/swa-latn-10mb-10mb_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3ForCausalLM", + "model_id": "clijo/qwen3-4b-instruct-2507-bf16-reco-grpo-b200-rapid-lime-orbit", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "WestCode1357/gpt-sw3-1.3b-instruct", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Gemma4UnifiedForConditionalGeneration", + "model_id": "tpls/gemma-4-12B-coder-fable5-composer2.5-v1-sft-v3", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "empero-ai/openNemo-9B-Claude-Opus-4.6-distill", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "openbmb/MiniCPM5-1B-Base", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3ForCausalLM", + "model_id": "marin-community/delphi-1e23-25Bparams-628Btokens", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "trohrbaugh/NVIDIA-Nemotron-3-Super-120B-A12B-BF16-heretic", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GptOssForCausalLM", + "model_id": "nri-ai/gpt-oss-20b-Ja-Fin-Thinking", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/swa-latn-100mb-100mb_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/tur-latn-10mb-after-ppt-shuff-dyck-100mb-ckpt500_seed10", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "adamo1139/Danube3-4b-4chan-HESOYAM-2510", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GlmMoeDsaForCausalLM", + "model_id": "upmarking/kalki-2.5", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "WhirlwindAI/AtomZephyr", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "nmihtrug/NVIDIA-Nemotron-Nano-9B-v2-FinCoT-VN", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/ita-latn-10mb-after-ppt-shuff-dyck-100mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3_5ForConditionalGeneration", + "model_id": "huihui-ai/Huihui-Qwythos-9B-Claude-Mythos-5-1M-abliterated", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/zho-hans-10mb-after-ppt-Dp-100mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "MistralForCausalLM", + "model_id": "anthracite-org/magnum-v2.5-12b-kto", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/eng-latn-10mb-after-ppt-Dp-100mb-ckpt500_seed10", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "MistralForCausalLM", + "model_id": "unsloth/Mistral-Large-Instruct-2407-bnb-4bit", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/tur-latn-100mb-10mb_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "User01110/supralabs-50M-testing", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/eng-latn-100mb-after-ppt-shuff-dyck-100mb-ckpt500_seed10", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3_5ForCausalLM", + "model_id": "infgrad/Prism-Qwen3.5-Reranker-4B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "tencent/Hunyuan-0.5B-Instruct-AWQ-Int4", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "CohereForCausalLM", + "model_id": "AnishRacherla/aya-expanse-8b-pruned-4newlayers", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "MistralForCausalLM", + "model_id": "Ba2han/experimental3", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3ForCausalLM", + "model_id": "iamthecage/Qwen3-Embedding-4B-OptiQ-4bit", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/eng-latn-10mb-after-ppt-shuff-dyck-100mb-ckpt500_seed10", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GraniteForCausalLM", + "model_id": "SamsungSDS-Research/SGuard-JailbreakFilter-2B-v1", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "facebook/llm-compiler-7b", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/eng-latn-10mb-100mb_seed10", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3NextForCausalLM", + "model_id": "NexVeridian/Qwen3-Coder-Next-6bit", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Gemma4ForConditionalGeneration", + "model_id": "Nekochu/Gemma-4-31B-RedstoneDoor", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3ForCausalLM", + "model_id": "didula-wso2/Qwen3-8B-rl530_with_think_knowledge_merged", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "clowman/Llama-3.1-405B-Instruct-AWQ-Int4", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "Max-and-Omnis/Nemotron-3-Super-64B-A12B-Math-REAP-BF16", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3ForCausalLM", + "model_id": "RedHatAI/Qwen3-4B-Instruct-2507-quantized.w4a16", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "LinkSoul/Chinese-Llama-2-7b-4bit", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/eng-latn-10mb-ppt-shuff-dyck-10mb_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3ForCausalLM", + "model_id": "warshanks/Hermes-4-14B-AWQ", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2ForCausalLM", + "model_id": "huihui-ai/Qwen2.5-14B-Instruct-1M-abliterated", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3ForCausalLM", + "model_id": "Polygl0t/Tucano2-qwen-3.7B-Base", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Gemma4ForConditionalGeneration", + "model_id": "Gryphe/Gemma-4-26B-A4B-StyleTune-V2", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "Max-and-Omnis/Nemotron-3-Super-64B-A12B-Math-REAP-AWQ", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/eng-latn-100mb-after-ppt-Dp-100mb-ckpt500_seed10", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "vg1024/NVIDIA-Nemotron-Nano-9B-v2-Japanese", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3ForCausalLM", + "model_id": "rpreite/Qwen3-Embedding-8B-INT4-W4A16", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "SECWIKI/trohrbaugh-NVIDIA-Nemotron-3-Super-120B-A12B-BF16-heretic-v2", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "OsaurusAI/NVIDIA-Nemotron-3-Ultra-550B-A55B-JANGTQ_1L", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Gemma4UnifiedForConditionalGeneration", + "model_id": "adamm-hf/Gemma-4-12B-OBLITERATED", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "MistralForCausalLM", + "model_id": "SillyTilly/mistralai_Mistral-Nemo-Instruct-2407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Gemma3ForConditionalGeneration", + "model_id": "CYFRAGOVPL/PLLuM-4B-instruct-2512", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "MistralForCausalLM", + "model_id": "RedHatAI/Mistral-7B-Instruct-v0.3-quantized.w4a16", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "MuVeraAI/NVIDIA-Nemotron-3-Super-120B-A12B-BF16", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "RepublicOfKorokke/NVIDIA-Nemotron-3-Nano-4B-oQ3.5", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Gemma4UnifiedForConditionalGeneration", + "model_id": "tpls/gemma-4-12B-coder-fable5-composer2.5-v1-sft-v5-abliterated", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3_5ForConditionalGeneration", + "model_id": "shuhulx/Qwopus3.5-4B-Coder-Fable5-v1", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2ForCausalLM", + "model_id": "thirdeyeai/Qwen2.5-1.5B-Instruct-uncensored", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "Yadro13/NVIDIA-Nemotron-3-Super-120B-A12B-BF16", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Glm4MoeForCausalLM", + "model_id": "QuantTrio/GLM-4.6-AWQ", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3_5ForConditionalGeneration", + "model_id": "Avesed/Qwen3.6-27B-INT8-W8A8", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/rus-cyrl-100mb-10mb_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3_5ForConditionalGeneration", + "model_id": "GestaltLabs/Ornstein-3.5-9B-V2", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "huggingworld/NVIDIA-Nemotron-3-Nano-4B-BF16-ONNX", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "3DJ77/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/swe-latn-100mb-10mb_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2ForCausalLM", + "model_id": "zkxxxx/VibeThinker-3B-heretic-fc", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/rus-cyrl-10mb-after-ppt-shuff-dyck-100mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3_5ForConditionalGeneration", + "model_id": "apodex/Apodex-1.0-2B-SFT", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/dan-latn-10mb-after-ppt-Dp-10mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/swe-latn-100mb-after-ppt-Dp-10mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "GODELEV/Archaea-74M", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Gemma4UnifiedForConditionalGeneration", + "model_id": "ShaunGves/Gemma-4-12B-OBLITERATED", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Phi3ForCausalLM", + "model_id": "kaitchup/Phi-4-mini-instruct-AutoRoundGPTQ-4bit", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3MoeForCausalLM", + "model_id": "ATH-MaaS/Marco-Nano-Instruct", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/ind-latn-100mb-10mb_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/ind-latn-100mb-after-ppt-shuff-dyck-10mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3_5ForConditionalGeneration", + "model_id": "apodex/Apodex-1.0-0.8B-SFT", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "llm-jp/llm-jp-3-7.2b-instruct", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3ForCausalLM", + "model_id": "dicta-il/DictaLM-3.0-1.7B-Thinking", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/eus-latn-10mb-after-ppt-shuff-dyck-10mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3MoeForCausalLM", + "model_id": "ramblingpolymath/Qwen3-30B-A3B-Instruct-2507-W8A8", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3MoeForCausalLM", + "model_id": "ATH-MaaS/Marco-Mini-Instruct", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Gemma4UnifiedForConditionalGeneration", + "model_id": "Iambackup/Gemma-4-12B-OBLITERATED", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "TomerGamerTV/NVIDIA-Nemotron-3-Nano-4B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/eus-latn-10mb-after-ppt-Dp-10mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2ForCausalLM", + "model_id": "driaforall/Tiny-Agent-a-0.5B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "SupraLabs/Supra-Mini-v6-1M", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Gemma4ForConditionalGeneration", + "model_id": "smd20/socialengineering", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "NemotronHForCausalLM", + "model_id": "RepublicOfKorokke/NVIDIA-Nemotron-3-Nano-4B-oQ4", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/dan-latn-100mb-after-ppt-shuff-dyck-10mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Gemma2ForCausalLM", + "model_id": "Phani1479432/rocky-gemma-2b", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3_5ForConditionalGeneration", + "model_id": "DJLougen/Qwable-5-27B-Coder", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "OctoThinker/OctoThinker-3B-Long-Base", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/jpn-jpan-10mb-after-ppt-Dp-100mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/eus-latn-100mb-after-ppt-Dp-10mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "User01110/testing-50M", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2ForCausalLM", + "model_id": "iamzac/Qwen2.5-0.5B-Instruct-Gensyn-Swarm-graceful_reclusive_skunk", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/ind-latn-100mb-after-ppt-Dp-10mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Gemma3ForCausalLM", + "model_id": "alnnahwi/gemma-3-1b-arabic-gec-v1", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/rus-cyrl-10mb-after-ppt-Dp-100mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/ind-latn-10mb-10mb_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "c4tdr0ut/grok-oss-Revenant-8B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2ForCausalLM", + "model_id": "kaitchup/DeepSeek-R1-Distill-Qwen-7B-AutoRound-GPTQ-4bit", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3ForCausalLM", + "model_id": "zenlm/zen-agent-4b", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/dan-latn-10mb-10mb_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/hin-deva-10mb-10mb_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "open-unlearning/unlearn_tofu_Llama-3.2-1B-Instruct_forget10_RMU_lr5e-05_layer5_scoeff10_epoch10", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3_5ForConditionalGeneration", + "model_id": "YuYu1015/Huihui-Qwen3.5-27B-abliterated-int4-AutoRound", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/jpn-jpan-10mb-10mb_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "GPT2LMHeadModel", + "model_id": "fpadovani/dan-latn-10mb-after-ppt-shuff-dyck-10mb-ckpt500_seed3407", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2ForCausalLM", + "model_id": "RedHatAI/Qwen2.5-0.5B-quantized.w8a8", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen3ForCausalLM", + "model_id": "ceselder/nanonla-l24-av-qwen3-8b", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "sahilchachra/MiniCPM5-1B-Uncensored", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "LlamaForCausalLM", + "model_id": "HandH1998/QQQ-Llama-3-8b-g128", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "tencent/Hunyuan-7B-Instruct-0124", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "tencent/Hunyuan-4B-Pretrain", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "tencent/Hunyuan-7B-Instruct-GPTQ-Int4", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "tencent/Hunyuan-1.8B-Instruct-AWQ-Int4", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "tencent/Hunyuan-MT-Chimera-7B", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "tencent/HY-MT1.5-7B-GPTQ-Int4", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "tencent/Hy-MT1.5-1.8B-2bit", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "HunYuanDenseV1ForCausalLM", + "model_id": "tencent/Hy-MT1.5-1.8B-1.25bit", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "Qwen2MoeForCausalLM", + "model_id": "hyper-accel/ci-random-qwen2-moe-a3b", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, { "architecture_id": "ArceeForCausalLM", "model_id": "optimum-intel-internal-testing/tiny-random-ArceeForCausalLM", diff --git a/transformer_lens/tools/model_registry/data/verification_history.json b/transformer_lens/tools/model_registry/data/verification_history.json index 4946b5fd5..530f03e15 100644 --- a/transformer_lens/tools/model_registry/data/verification_history.json +++ b/transformer_lens/tools/model_registry/data/verification_history.json @@ -1,5 +1,5 @@ { - "last_updated": "2026-06-27T00:21:21.662427", + "last_updated": "2026-07-02T02:19:51.172022", "records": [ { "model_id": "Macropodus/macbert4mdcspell_v1", @@ -16390,6 +16390,476 @@ "notes": "Below threshold: P1=50.0% < 100.0% (failed: forward_pass_logits) \u2014 Tensors differ: max_diff=0.020484, mean_rel=0.006617", "invalidated": false, "invalidation_reason": null + }, + { + "model_id": "state-spaces/mamba-130m-hf", + "architecture_id": "MambaForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Below threshold: P1=50.0% < 100.0% (failed: all_components) \u2014 96/171 components failed (96 critical)", + "invalidated": true, + "invalidation_reason": "Superseded by the clean re-run (P1=100) after the component-benchmark fix that skips SSM mixer-internal submodules; the 96/171 component failures were the isolated harness feeding d_model-shaped inputs to SSM-internal projections, not a real divergence." + }, + { + "model_id": "state-spaces/mamba-130m-hf", + "architecture_id": "MambaForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed with issues: P3=90.0% (failed: attention_output_centering, mlp_output_centering)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "AntonV/mamba2-130m-hf", + "architecture_id": "Mamba2ForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed with issues: P3=90.0% (failed: attention_output_centering, mlp_output_centering)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "ibm-granite/granite-4.0-tiny-preview", + "architecture_id": "GraniteMoeHybridForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Below threshold: P3=94.7% but required tests failed: logits_equivalence \u2014 Tensors differ: max_diff=0.375000, mean_rel=0.002045", + "invalidated": true, + "invalidation_reason": "bf16 precision of compatibility-mode center_unembed (root-caused by toggle: off=0.000 both dtypes, fp32=4.2e-5); superseded by the clean fp32 run (P3=100). Not an algorithmic bug." + }, + { + "model_id": "NYTK/PULI-HuBA-mamba-130M", + "architecture_id": "MambaForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed with issues: P3=90.0% (failed: attention_output_centering, mlp_output_centering)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "deqing/mamba2-300M-v5-mamba2", + "architecture_id": "Mamba2ForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed with issues: P3=90.0% (failed: attention_output_centering, mlp_output_centering)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "ibm-granite/granite-4.0-tiny-preview", + "architecture_id": "GraniteMoeHybridForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "ibm-granite/granite-4.0-1b", + "architecture_id": "GraniteMoeHybridForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "yujiepan/qwen3-next-moe-tiny-random", + "architecture_id": "Qwen3NextForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed with issues: P2=92.3% (failed: gated_hooks_fire); P3=94.7% (failed: gated_hooks_fire)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "GoodStartLabs/gin-rummy-hbc-qwen3.5-0.8b", + "architecture_id": "Qwen3_5ForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed with issues: P2=92.3% (failed: gated_hooks_fire); P3=94.7% (failed: gated_hooks_fire)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "Evelyn67/Qwen3.5-2B-Her", + "architecture_id": "Qwen3_5ForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed with issues: P2=92.3% (failed: gated_hooks_fire); P3=94.7% (failed: gated_hooks_fire)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16", + "architecture_id": "NemotronHForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed with issues: P3=90.0% (failed: attention_output_centering, mlp_output_centering)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "nvidia/NVIDIA-Nemotron-Nano-9B-v2", + "architecture_id": "NemotronHForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed with issues: P3=90.0% (failed: attention_output_centering, mlp_output_centering)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "dominguesm/mambarim-110m", + "architecture_id": "MambaForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed with issues: P3=90.0% (failed: attention_output_centering, mlp_output_centering)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "state-spaces/mamba-370m-hf", + "architecture_id": "MambaForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed with issues: P3=90.0% (failed: attention_output_centering, mlp_output_centering)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "EchoLabs33/mamba-130m-hxq", + "architecture_id": "MambaForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Below threshold: P1=0.0% < 100.0% (failed: all_components, forward_pass_logits) \u2014 24/51 components failed (24 critical)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "deqing/convergent-mamba2-300M-adamw-original", + "architecture_id": "Mamba2ForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed with issues, low text quality: P3=90.0% (failed: attention_output_centering, mlp_output_centering)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "limloop/whiff-mamba2-50M-v0.1", + "architecture_id": "Mamba2ForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed with issues: P3=90.0% (failed: attention_output_centering, mlp_output_centering)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "EchoLabs33/mamba2-1.3b-hxq", + "architecture_id": "Mamba2ForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Below threshold: P1=0.0% < 100.0% (failed: all_components, forward_pass_logits) \u2014 50/99 components failed (50 critical)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "ibm-granite/granite-4.0-350m", + "architecture_id": "GraniteMoeHybridForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "ibm-granite/granite-4.0-1b-base", + "architecture_id": "GraniteMoeHybridForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "WithinUsAI/IBM-Grok4-Ultra.Fast.Coder-1B", + "architecture_id": "GraniteMoeHybridForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "GoodStartLabs/gin-rummy-hbc-qwen3.5-2b", + "architecture_id": "Qwen3_5ForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "ZirTech/OmniMath-2B", + "architecture_id": "Qwen3_5ForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "continuum-ai/qwen3.5-0.8b-general-forged", + "architecture_id": "Qwen3_5ForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "optimum-intel-internal-testing/tiny-random-qwen3-next", + "architecture_id": "Qwen3NextForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "tiny-random/qwen3-next-moe", + "architecture_id": "Qwen3NextForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "trl-internal-testing/tiny-NemotronHForCausalLM-nano", + "architecture_id": "NemotronHForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "trl-internal-testing/tiny-NemotronHForCausalLM-super", + "architecture_id": "NemotronHForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "trl-internal-testing/tiny-NemotronHForCausalLM-ultra", + "architecture_id": "NemotronHForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "state-spaces/mamba-790m-hf", + "architecture_id": "MambaForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed with issues: P3=90.0% (failed: attention_output_centering, mlp_output_centering)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "evolai/evolai_mamba_naive_kl", + "architecture_id": "Mamba2ForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed with issues: P3=90.0% (failed: attention_output_centering, mlp_output_centering)", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "state-spaces/mamba-130m-hf", + "architecture_id": "MambaForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "dominguesm/mambarim-110m", + "architecture_id": "MambaForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "NYTK/PULI-HuBA-mamba-130M", + "architecture_id": "MambaForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "state-spaces/mamba-370m-hf", + "architecture_id": "MambaForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "state-spaces/mamba-790m-hf", + "architecture_id": "MambaForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "AntonV/mamba2-130m-hf", + "architecture_id": "Mamba2ForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "deqing/convergent-mamba2-300M-adamw-original", + "architecture_id": "Mamba2ForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed with issues, low text quality", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "deqing/mamba2-300M-v5-mamba2", + "architecture_id": "Mamba2ForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "evolai/evolai_mamba_naive_kl", + "architecture_id": "Mamba2ForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "limloop/whiff-mamba2-50M-v0.1", + "architecture_id": "Mamba2ForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "yujiepan/qwen3-next-moe-tiny-random", + "architecture_id": "Qwen3NextForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "Evelyn67/Qwen3.5-2B-Her", + "architecture_id": "Qwen3_5ForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "GoodStartLabs/gin-rummy-hbc-qwen3.5-0.8b", + "architecture_id": "Qwen3_5ForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16", + "architecture_id": "NemotronHForCausalLM", + "verified_date": "2026-07-01", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "nvidia/NVIDIA-Nemotron-Nano-9B-v2", + "architecture_id": "NemotronHForCausalLM", + "verified_date": "2026-07-02", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null } ] } diff --git a/transformer_lens/tools/model_registry/verify_models.py b/transformer_lens/tools/model_registry/verify_models.py index 9b939fe7c..d84f72b36 100644 --- a/transformer_lens/tools/model_registry/verify_models.py +++ b/transformer_lens/tools/model_registry/verify_models.py @@ -114,7 +114,8 @@ def _phases_to_run(arch: str, phases: list[int]) -> list[int]: An adapter's ``applicable_phases`` declares which text phases (1-4) it covers. Phases 7/8 are gated separately by ``is_multimodal``/``is_audio`` in the benchmark, so they are never - filtered out here. An empty result means the architecture is skipped (e.g. SSMs). + filtered out here. An empty result means none of the requested phases apply to this + architecture (SSM / recurrent families run all four). """ from transformer_lens.factories.architecture_adapter_factory import ( SUPPORTED_ARCHITECTURES, @@ -837,11 +838,9 @@ def verify_models( _skip_model(model_id, arch, note, progress, quiet) continue - # Step 0b: Check adapter-level phase applicability. Architectures - # with applicable_phases=[] (e.g. SSMs) skip verify_models entirely - # because the benchmark suite has transformer-shaped assumptions that - # would need a dedicated refactor to cover them. Verification for - # these architectures lives in the integration test suite. + # Step 0b: Check adapter-level phase applicability. applicable_phases=[] + # means no phase applies (a genuinely unsupported architecture — rare); + # SSM / recurrent families run the full [1, 2, 3, 4]. from transformer_lens.factories.architecture_adapter_factory import ( SUPPORTED_ARCHITECTURES, ) From ef3df55eed5b8abb281404e6f6bcdd9336b3e62d Mon Sep 17 00:00:00 2001 From: Uday Phalak Date: Fri, 3 Jul 2026 20:39:24 +0530 Subject: [PATCH 04/11] Add Falcon-H1 parallel hybrid architecture adapter (FalconH1ForCausalLM) (#1480) * Add Falcon H1 Hybrid Parallel Mamba TransformerBridge Adapter * Add Tiny-90M parity tests and attention-branch ablation coverage for Falcon-H1 adapter * Drop redundant CI skipif from Falcon-H1 integration tests; rely on slow marker * Fix mypy attr-defined for Falcon-H1 eps_attr assignment --------- Co-authored-by: jlarson4 --- .../model_bridge/test_falcon_h1_adapter.py | 308 ++++++++++++++ .../test_falcon_h1_adapter.py | 381 ++++++++++++++++++ .../factories/architecture_adapter_factory.py | 2 + .../model_bridge/sources/_bridge_builder.py | 8 + .../model_bridge/sources/transformers.py | 8 + .../supported_architectures/__init__.py | 4 + .../supported_architectures/falcon_h1.py | 190 +++++++++ .../tools/model_registry/__init__.py | 2 + .../model_registry/data/supported_models.json | 178 +++++--- .../tools/model_registry/generate_report.py | 1 + 10 files changed, 1035 insertions(+), 47 deletions(-) create mode 100644 tests/integration/model_bridge/test_falcon_h1_adapter.py create mode 100644 tests/unit/model_bridge/supported_architectures/test_falcon_h1_adapter.py create mode 100644 transformer_lens/model_bridge/supported_architectures/falcon_h1.py diff --git a/tests/integration/model_bridge/test_falcon_h1_adapter.py b/tests/integration/model_bridge/test_falcon_h1_adapter.py new file mode 100644 index 000000000..cd38eecc0 --- /dev/null +++ b/tests/integration/model_bridge/test_falcon_h1_adapter.py @@ -0,0 +1,308 @@ +"""Integration tests for the Falcon-H1 architecture adapter. + +Verifies wrap-don't-reimplement behaviour against Falcon-H1 checkpoints: + +- ``tiiuae/Falcon-H1-Tiny-90M-Instruct`` — smallest variant (issue #1403 verify-first). +- ``tiiuae/Falcon-H1-0.5B-Base`` — primary parity suite. +- Forward-pass logits match HF **exactly** (the bridge delegates the whole block + to HF, which applies Falcon-H1's ~12 scalar multipliers natively). +- Both the attention branch AND the Mamba-2 branch are independently hookable in + every block — the defining value of a parallel-hybrid adapter. +- Greedy generation matches HF bit-for-bit through the standard unified + KV-cache path. +- Config propagation: SSM dims surface on ``bridge.cfg``. + +The Falcon-H1 checkpoints are not in the CI cached-model allowlist, so this +module is marked ``slow``. CI and ``make integration-test`` deselect it via +``-m "not slow"`` (same pattern as ``test_nemotron_h_adapter.py``). Run +locally with: + + pytest tests/integration/model_bridge/test_falcon_h1_adapter.py -v +""" + +import pytest +import torch + +from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.generalized_components import ( + BlockBridge, + DepthwiseConv1DBridge, + PositionEmbeddingsAttentionBridge, + SSM2MixerBridge, +) + +MODEL = "tiiuae/Falcon-H1-0.5B-Base" +TINY_MODEL = "tiiuae/Falcon-H1-Tiny-90M-Instruct" + +pytestmark = pytest.mark.slow + + +@pytest.fixture(scope="module") +def falcon_h1_bridge(): + return TransformerBridge.boot_transformers(MODEL, device="cpu", dtype=torch.float32) + + +# --------------------------------------------------------------------------- +# Bridge creation +# --------------------------------------------------------------------------- + + +class TestFalconH1BridgeCreation: + def test_block_count(self, falcon_h1_bridge): + assert len(falcon_h1_bridge.blocks) == 36 + + def test_config_flags(self, falcon_h1_bridge): + cfg = falcon_h1_bridge.cfg + assert cfg.normalization_type == "RMS" + assert cfg.uses_rms_norm is True + assert cfg.positional_embedding_type == "rotary" + assert cfg.gated_mlp is True + assert cfg.attn_only is False + assert cfg.final_rms is True + assert cfg.n_key_value_heads == 2 + # Standard KV-cache generation path, not the pure-Mamba stateful loop. + assert getattr(cfg, "is_stateful", False) is False + + def test_ssm_config_propagated(self, falcon_h1_bridge): + cfg = falcon_h1_bridge.cfg + assert getattr(cfg, "mamba_d_ssm", None) == 1536 + # conv_dim = d_ssm + 2 * n_groups * d_state = 1536 + 2*1*128 + assert getattr(cfg, "conv_dim", None) == 1536 + 2 * 1 * 128 + # in_proj fuses gate, hidden_BC, dt: 2*d_ssm + conv_dim + mamba_n_heads + assert getattr(cfg, "expected_in_proj_out_features", None) == 2 * 1536 + 1792 + 24 + + def test_both_branches_present(self, falcon_h1_bridge): + block = falcon_h1_bridge.blocks[0] + assert isinstance(block, BlockBridge) + assert isinstance(block.attn, PositionEmbeddingsAttentionBridge) + assert isinstance(block.mamba, SSM2MixerBridge) + + def test_has_core_components(self, falcon_h1_bridge): + for comp in ("embed", "unembed", "ln_final", "rotary_emb"): + assert hasattr(falcon_h1_bridge, comp) + + +# --------------------------------------------------------------------------- +# Forward equivalence +# --------------------------------------------------------------------------- + + +class TestFalconH1ForwardEquivalence: + def test_forward_returns_logits(self, falcon_h1_bridge): + tokens = torch.tensor([[1, 2, 3, 4]]) + with torch.no_grad(): + out = falcon_h1_bridge(tokens) + assert out.shape == (1, 4, falcon_h1_bridge.cfg.d_vocab) + assert not torch.isnan(out).any() + + def test_forward_matches_hf_exactly(self, falcon_h1_bridge): + """Passthrough: bridge logits must be bitwise-identical to HF. + + HF applies all of Falcon-H1's scalar multipliers in its own forward, so + with no weight folding the diff is exactly zero. + """ + tokens = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]]) + hf_model = falcon_h1_bridge.original_model + with torch.no_grad(): + bridge_out = falcon_h1_bridge(tokens) + hf_out = hf_model(tokens).logits + max_diff = (bridge_out - hf_out).abs().max().item() + assert max_diff == 0.0, f"Bridge vs HF max diff = {max_diff}" + + +# --------------------------------------------------------------------------- +# HF delegation (wrap-don't-reimplement) +# --------------------------------------------------------------------------- + + +class TestFalconH1HFDelegation: + """Bridge submodules wrap live HF torch modules rather than reimplementing.""" + + def test_attn_projections_wrap_linear(self, falcon_h1_bridge): + attn = falcon_h1_bridge.blocks[0].attn + for key in ("q", "k", "v", "o"): + assert isinstance(attn._modules[key].original_component, torch.nn.Linear) + + def test_mamba_submodules_wrap_real_modules(self, falcon_h1_bridge): + mamba = falcon_h1_bridge.blocks[0].mamba + assert isinstance(mamba.in_proj.original_component, torch.nn.Linear) + assert isinstance(mamba.out_proj.original_component, torch.nn.Linear) + assert isinstance( + falcon_h1_bridge.blocks[0].mamba.conv1d, + DepthwiseConv1DBridge, + ) + assert isinstance(mamba.conv1d.original_component, torch.nn.Conv1d) + + def test_mlp_projections_wrap_linear(self, falcon_h1_bridge): + mlp = falcon_h1_bridge.blocks[0].mlp + for key in ("gate", "in", "out"): + assert isinstance(mlp._modules[key].original_component, torch.nn.Linear) + + +# --------------------------------------------------------------------------- +# Parallel-branch hook coverage +# --------------------------------------------------------------------------- + + +class TestFalconH1ParallelHooks: + @pytest.fixture(scope="class") + def cache(self, falcon_h1_bridge): + tokens = torch.tensor([[1, 2, 3, 4]]) + with torch.no_grad(): + _, cache = falcon_h1_bridge.run_with_cache(tokens) + return cache + + def test_attention_branch_hooks_fire(self, cache): + for i in (0, 18, 35): + for proj in ("q", "k", "v", "o"): + assert f"blocks.{i}.attn.{proj}.hook_out" in cache + + def test_mamba_branch_hooks_fire(self, cache): + for i in (0, 18, 35): + for proj in ("in_proj", "conv1d", "out_proj"): + assert f"blocks.{i}.mamba.{proj}.hook_out" in cache + + def test_both_branches_fire_in_same_block(self, cache): + # The parallel-hybrid value proposition: in one block, both paths are + # observable so a researcher can ablate one and compare. + assert "blocks.0.attn.hook_out" in cache + assert "blocks.0.mamba.hook_out" in cache + + def test_residual_and_mlp_hooks_present(self, cache): + for i in (0, 35): + assert f"blocks.{i}.hook_in" in cache + assert f"blocks.{i}.hook_out" in cache + assert "blocks.0.mlp.hook_out" in cache + + def test_no_nan_in_hooked_activations(self, cache): + for key in ("blocks.0.attn.hook_out", "blocks.0.mamba.hook_out", "blocks.0.hook_out"): + assert not torch.isnan(cache[key]).any() + + +# --------------------------------------------------------------------------- +# Hook ablation +# --------------------------------------------------------------------------- + + +class TestFalconH1HookAblation: + def test_zero_ablation_changes_output(self, falcon_h1_bridge): + tokens = torch.tensor([[1, 2, 3, 4]]) + with torch.no_grad(): + baseline = falcon_h1_bridge(tokens) + + def zero(t, hook): + return torch.zeros_like(t) + + with torch.no_grad(): + ablated = falcon_h1_bridge.run_with_hooks( + tokens, fwd_hooks=[("blocks.18.hook_in", zero)] + ) + assert not torch.allclose(baseline, ablated) + + def test_mamba_branch_ablation_changes_output(self, falcon_h1_bridge): + # Ablating just the Mamba branch output must perturb logits, proving the + # branch is on the live forward path (not a dead hook). + tokens = torch.tensor([[1, 2, 3, 4]]) + with torch.no_grad(): + baseline = falcon_h1_bridge(tokens) + + def zero(t, hook): + return torch.zeros_like(t) + + with torch.no_grad(): + ablated = falcon_h1_bridge.run_with_hooks( + tokens, fwd_hooks=[("blocks.0.mamba.hook_out", zero)] + ) + assert not torch.allclose(baseline, ablated) + + def test_attn_branch_ablation_changes_output(self, falcon_h1_bridge): + # Symmetric to the Mamba ablation: zeroing the attention branch output + # must perturb logits, proving both parallel paths are live. + tokens = torch.tensor([[1, 2, 3, 4]]) + with torch.no_grad(): + baseline = falcon_h1_bridge(tokens) + + def zero(t, hook): + return torch.zeros_like(t) + + with torch.no_grad(): + ablated = falcon_h1_bridge.run_with_hooks( + tokens, fwd_hooks=[("blocks.0.attn.hook_out", zero)] + ) + assert not torch.allclose(baseline, ablated) + + +# --------------------------------------------------------------------------- +# Tiny-90M parity (issue #1403 verify-smallest-first) +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def falcon_h1_tiny_bridge(): + return TransformerBridge.boot_transformers(TINY_MODEL, device="cpu", dtype=torch.float32) + + +class TestFalconH1Tiny90MParity: + """Smoke parity on the smallest Falcon-H1 checkpoint.""" + + def test_block_count(self, falcon_h1_tiny_bridge): + assert len(falcon_h1_tiny_bridge.blocks) == 24 + + def test_both_branches_present(self, falcon_h1_tiny_bridge): + block = falcon_h1_tiny_bridge.blocks[0] + assert isinstance(block.attn, PositionEmbeddingsAttentionBridge) + assert isinstance(block.mamba, SSM2MixerBridge) + + def test_forward_matches_hf_exactly(self, falcon_h1_tiny_bridge): + tokens = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]]) + hf_model = falcon_h1_tiny_bridge.original_model + with torch.no_grad(): + bridge_out = falcon_h1_tiny_bridge(tokens) + hf_out = hf_model(tokens).logits + max_diff = (bridge_out - hf_out).abs().max().item() + assert max_diff == 0.0, f"Tiny-90M bridge vs HF max diff = {max_diff}" + + @pytest.mark.parametrize("prompt_len", [1, 4]) + def test_greedy_matches_hf(self, falcon_h1_tiny_bridge, prompt_len): + tokens = torch.arange(1, prompt_len + 1).unsqueeze(0) + with torch.no_grad(): + bridge_out = falcon_h1_tiny_bridge.generate(tokens, max_new_tokens=6, do_sample=False) + hf_out = falcon_h1_tiny_bridge.original_model.generate( + tokens, max_new_tokens=6, do_sample=False, pad_token_id=0 + ) + assert torch.equal( + bridge_out, hf_out + ), f"Tiny-90M prompt_len={prompt_len}: bridge={bridge_out.tolist()} vs HF={hf_out.tolist()}" + + +# --------------------------------------------------------------------------- +# Generation parity +# --------------------------------------------------------------------------- + + +class TestFalconH1Generation: + @pytest.mark.parametrize("prompt_len", [1, 4]) + def test_greedy_matches_hf(self, falcon_h1_bridge, prompt_len): + tokens = torch.arange(1, prompt_len + 1).unsqueeze(0) + with torch.no_grad(): + bridge_out = falcon_h1_bridge.generate(tokens, max_new_tokens=6, do_sample=False) + hf_out = falcon_h1_bridge.original_model.generate( + tokens, max_new_tokens=6, do_sample=False, pad_token_id=0 + ) + assert torch.equal( + bridge_out, hf_out + ), f"prompt_len={prompt_len}: bridge={bridge_out.tolist()} vs HF={hf_out.tolist()}" + + +# --------------------------------------------------------------------------- +# Config propagation vs HF +# --------------------------------------------------------------------------- + + +class TestFalconH1ConfigPropagation: + def test_mamba_dims_match_hf_config(self, falcon_h1_bridge): + hf_cfg = falcon_h1_bridge.original_model.config + assert falcon_h1_bridge.cfg.mamba_d_ssm == hf_cfg.mamba_d_ssm + assert falcon_h1_bridge.cfg.mamba_n_heads == hf_cfg.mamba_n_heads + assert falcon_h1_bridge.cfg.mamba_d_state == hf_cfg.mamba_d_state + assert falcon_h1_bridge.cfg.mamba_n_groups == hf_cfg.mamba_n_groups diff --git a/tests/unit/model_bridge/supported_architectures/test_falcon_h1_adapter.py b/tests/unit/model_bridge/supported_architectures/test_falcon_h1_adapter.py new file mode 100644 index 000000000..8567d7f83 --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_falcon_h1_adapter.py @@ -0,0 +1,381 @@ +"""Unit tests for FalconH1ArchitectureAdapter. + +Covers: config attribute propagation, component mapping bridge types and HF path +names, the parallel attn+mamba block shape, Mamba inner-norm optional flag, +applicable_phases, create_stateful_cache, factory + registry registration, and +guard tests against drift toward neighbouring (single-mixer) hybrid patterns. + +Structural-only: no weight load, no HF Hub access. +""" + +import pytest + +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.factories.architecture_adapter_factory import ( + ArchitectureAdapterFactory, +) +from transformer_lens.model_bridge.generalized_components import ( + BlockBridge, + DepthwiseConv1DBridge, + EmbeddingBridge, + GatedMLPBridge, + GatedRMSNormBridge, + LinearBridge, + PositionEmbeddingsAttentionBridge, + RMSNormalizationBridge, + RotaryEmbeddingBridge, + SSM2MixerBridge, + UnembeddingBridge, +) +from transformer_lens.model_bridge.supported_architectures.falcon_h1 import ( + FalconH1ArchitectureAdapter, +) + + +def _make_cfg( + n_layers: int = 2, + d_model: int = 64, + d_head: int = 16, + n_heads: int = 4, + n_key_value_heads: int = 2, + d_vocab: int = 100, + n_ctx: int = 128, + mamba_d_ssm: int = 32, + mamba_n_heads: int = 4, + mamba_d_head: int = 8, + mamba_d_state: int = 4, + mamba_n_groups: int = 1, +) -> TransformerBridgeConfig: + """Minimal TransformerBridgeConfig for Falcon-H1 adapter tests. + + Small Mamba-2 dimensions so tests run without loading any weights. The + Falcon-H1-specific fields are injected via ``setattr`` (the adapter reads + them with ``getattr``); they are not declared ``TransformerBridgeConfig`` + fields. + """ + cfg = TransformerBridgeConfig( + d_model=d_model, + d_head=d_head, + n_layers=n_layers, + n_ctx=n_ctx, + n_heads=n_heads, + n_key_value_heads=n_key_value_heads, + d_vocab=d_vocab, + default_prepend_bos=False, + architecture="FalconH1ForCausalLM", + ) + setattr(cfg, "mamba_d_ssm", mamba_d_ssm) + setattr(cfg, "mamba_n_heads", mamba_n_heads) + setattr(cfg, "mamba_d_head", mamba_d_head) + setattr(cfg, "mamba_d_state", mamba_d_state) + setattr(cfg, "mamba_n_groups", mamba_n_groups) + return cfg + + +@pytest.fixture(scope="class") +def cfg() -> TransformerBridgeConfig: + return _make_cfg() + + +@pytest.fixture(scope="class") +def adapter(cfg: TransformerBridgeConfig) -> FalconH1ArchitectureAdapter: + return FalconH1ArchitectureAdapter(cfg) + + +# --------------------------------------------------------------------------- +# Config attributes +# --------------------------------------------------------------------------- + + +class TestFalconH1AdapterConfig: + """Adapter propagates all required config attributes.""" + + def test_normalization_type_rms(self, adapter: FalconH1ArchitectureAdapter) -> None: + assert adapter.cfg.normalization_type == "RMS" + + def test_uses_rms_norm(self, adapter: FalconH1ArchitectureAdapter) -> None: + assert adapter.cfg.uses_rms_norm is True + + def test_positional_embedding_type_rotary(self, adapter: FalconH1ArchitectureAdapter) -> None: + # Falcon-H1 has a model-level rotary module (unlike NemotronH). + assert adapter.cfg.positional_embedding_type == "rotary" + + def test_gated_mlp_true(self, adapter: FalconH1ArchitectureAdapter) -> None: + # SwiGLU feed-forward. + assert adapter.cfg.gated_mlp is True + + def test_attn_only_false(self, adapter: FalconH1ArchitectureAdapter) -> None: + assert adapter.cfg.attn_only is False + + def test_final_rms_true(self, adapter: FalconH1ArchitectureAdapter) -> None: + assert adapter.cfg.final_rms is True + + def test_not_stateful(self, adapter: FalconH1ArchitectureAdapter) -> None: + # Falcon-H1 generates through the standard unified KV-cache path; the + # pure-Mamba stateful loop would diverge from HF after the first decode. + assert getattr(adapter.cfg, "is_stateful", False) is False + + def test_eps_attr_variance_epsilon(self, adapter: FalconH1ArchitectureAdapter) -> None: + assert adapter.cfg.eps_attr == "variance_epsilon" + + def test_n_key_value_heads_propagated(self, adapter: FalconH1ArchitectureAdapter) -> None: + assert adapter.cfg.n_key_value_heads == 2 + + def test_mamba_intermediate_size_propagated(self, adapter: FalconH1ArchitectureAdapter) -> None: + # mamba_d_ssm is the inner SSM width directly. + assert getattr(adapter.cfg, "mamba_intermediate_size", None) == 32 + + def test_conv_dim_propagated(self, adapter: FalconH1ArchitectureAdapter) -> None: + # d_ssm=32, n_groups=1, d_state=4 → 32 + 2*1*4 = 40 + assert getattr(adapter.cfg, "conv_dim", None) == 40 + + def test_expected_in_proj_out_features(self, adapter: FalconH1ArchitectureAdapter) -> None: + # 2*d_ssm + conv_dim + mamba_n_heads = 64 + 40 + 4 = 108 + assert getattr(adapter.cfg, "expected_in_proj_out_features", None) == 108 + + def test_applicable_phases_empty(self) -> None: + # verify_models is transformer-shaped; SSM hybrids skip it. + assert FalconH1ArchitectureAdapter.applicable_phases == [] + + +class TestFalconH1ConvDimFormula: + """conv_dim and in_proj derivations track the Mamba-2 layout.""" + + def test_conv_dim_formula(self) -> None: + a = FalconH1ArchitectureAdapter( + _make_cfg(mamba_d_ssm=48, mamba_n_groups=2, mamba_d_state=8) + ) + assert getattr(a.cfg, "conv_dim") == 48 + 2 * 2 * 8 + + def test_mamba_intermediate_equals_d_ssm(self) -> None: + a = FalconH1ArchitectureAdapter(_make_cfg(mamba_d_ssm=96)) + assert getattr(a.cfg, "mamba_intermediate_size") == 96 + + +# --------------------------------------------------------------------------- +# Top-level component mapping +# --------------------------------------------------------------------------- + + +class TestFalconH1TopLevelComponents: + """component_mapping has exactly the expected top-level keys.""" + + def test_required_keys(self, adapter: FalconH1ArchitectureAdapter) -> None: + assert set(adapter.component_mapping.keys()) == { + "embed", + "rotary_emb", + "blocks", + "ln_final", + "unembed", + } + + def test_rotary_emb_present(self, adapter: FalconH1ArchitectureAdapter) -> None: + # Falcon-H1 exposes a model-level rotary module. + assert isinstance(adapter.component_mapping["rotary_emb"], RotaryEmbeddingBridge) + assert adapter.component_mapping["rotary_emb"].name == "model.rotary_emb" + + def test_embed(self, adapter: FalconH1ArchitectureAdapter) -> None: + embed = adapter.component_mapping["embed"] + assert isinstance(embed, EmbeddingBridge) + assert embed.name == "model.embed_tokens" + + def test_blocks_is_block_bridge(self, adapter: FalconH1ArchitectureAdapter) -> None: + assert isinstance(adapter.component_mapping["blocks"], BlockBridge) + assert adapter.component_mapping["blocks"].name == "model.layers" + + def test_ln_final(self, adapter: FalconH1ArchitectureAdapter) -> None: + ln_final = adapter.component_mapping["ln_final"] + assert isinstance(ln_final, RMSNormalizationBridge) + assert ln_final.name == "model.final_layernorm" + + def test_unembed(self, adapter: FalconH1ArchitectureAdapter) -> None: + unembed = adapter.component_mapping["unembed"] + assert isinstance(unembed, UnembeddingBridge) + assert unembed.name == "lm_head" + + +# --------------------------------------------------------------------------- +# Block-level submodules (the parallel hybrid shape) +# --------------------------------------------------------------------------- + + +class TestFalconH1BlockSubmodules: + """The block carries two norms plus parallel attn + mamba and a SwiGLU MLP.""" + + @pytest.fixture(scope="class") + def blocks(self, adapter: FalconH1ArchitectureAdapter) -> BlockBridge: + return adapter.component_mapping["blocks"] + + def test_ln1_is_input_layernorm(self, blocks: BlockBridge) -> None: + ln1 = blocks.submodules["ln1"] + assert isinstance(ln1, RMSNormalizationBridge) + assert ln1.name == "input_layernorm" + + def test_ln2_is_pre_ff_layernorm(self, blocks: BlockBridge) -> None: + # Two-norm block: pre_ff_layernorm gives a real hook_resid_mid. + ln2 = blocks.submodules["ln2"] + assert isinstance(ln2, RMSNormalizationBridge) + assert ln2.name == "pre_ff_layernorm" + + def test_attn_branch_present(self, blocks: BlockBridge) -> None: + attn = blocks.submodules["attn"] + assert isinstance(attn, PositionEmbeddingsAttentionBridge) + assert attn.name == "self_attn" + + def test_mamba_branch_present(self, blocks: BlockBridge) -> None: + mamba = blocks.submodules["mamba"] + assert isinstance(mamba, SSM2MixerBridge) + assert mamba.name == "mamba" + + def test_both_branches_present_in_every_block(self, blocks: BlockBridge) -> None: + # The defining feature of Falcon-H1: attn AND mamba in parallel, not a + # single per-layer mixer slot. + assert "attn" in blocks.submodules + assert "mamba" in blocks.submodules + + def test_mlp_is_swiglu(self, blocks: BlockBridge) -> None: + mlp = blocks.submodules["mlp"] + assert isinstance(mlp, GatedMLPBridge) + assert mlp.name == "feed_forward" + + def test_block_submodule_keys(self, blocks: BlockBridge) -> None: + assert set(blocks.submodules.keys()) == {"ln1", "ln2", "attn", "mamba", "mlp"} + + +# --------------------------------------------------------------------------- +# Attention submodules +# --------------------------------------------------------------------------- + + +class TestFalconH1AttentionSubmodules: + """Split q/k/v/o projections (GQA).""" + + @pytest.fixture(scope="class") + def attn(self, adapter: FalconH1ArchitectureAdapter) -> PositionEmbeddingsAttentionBridge: + return adapter.component_mapping["blocks"].submodules["attn"] + + @pytest.mark.parametrize( + "key,hf_name", + [("q", "q_proj"), ("k", "k_proj"), ("v", "v_proj"), ("o", "o_proj")], + ) + def test_projection( + self, attn: PositionEmbeddingsAttentionBridge, key: str, hf_name: str + ) -> None: + proj = attn.submodules[key] + assert isinstance(proj, LinearBridge) + assert proj.name == hf_name + + +# --------------------------------------------------------------------------- +# Mamba mixer submodules +# --------------------------------------------------------------------------- + + +class TestFalconH1MambaSubmodules: + """SSM2MixerBridge submodules; inner_norm is optional (mamba_rms_norm=false).""" + + @pytest.fixture(scope="class") + def mamba(self, adapter: FalconH1ArchitectureAdapter) -> SSM2MixerBridge: + return adapter.component_mapping["blocks"].submodules["mamba"] + + def test_in_proj(self, mamba: SSM2MixerBridge) -> None: + assert isinstance(mamba.submodules["in_proj"], LinearBridge) + assert mamba.submodules["in_proj"].name == "in_proj" + + def test_conv1d(self, mamba: SSM2MixerBridge) -> None: + assert isinstance(mamba.submodules["conv1d"], DepthwiseConv1DBridge) + assert mamba.submodules["conv1d"].name == "conv1d" + + def test_out_proj(self, mamba: SSM2MixerBridge) -> None: + assert isinstance(mamba.submodules["out_proj"], LinearBridge) + assert mamba.submodules["out_proj"].name == "out_proj" + + def test_inner_norm_is_gated_and_optional(self, mamba: SSM2MixerBridge) -> None: + inner_norm = mamba.submodules["inner_norm"] + assert isinstance(inner_norm, GatedRMSNormBridge) + # HF calls it "norm" inside the mixer; absent when mamba_rms_norm=false. + assert inner_norm.name == "norm" + assert inner_norm.optional is True + + def test_mamba_submodule_keys(self, mamba: SSM2MixerBridge) -> None: + assert set(mamba.submodules.keys()) == {"in_proj", "conv1d", "inner_norm", "out_proj"} + + +# --------------------------------------------------------------------------- +# MLP submodules +# --------------------------------------------------------------------------- + + +class TestFalconH1MLPSubmodules: + """SwiGLU feed-forward: gate_proj / up_proj / down_proj.""" + + @pytest.fixture(scope="class") + def mlp(self, adapter: FalconH1ArchitectureAdapter) -> GatedMLPBridge: + return adapter.component_mapping["blocks"].submodules["mlp"] + + @pytest.mark.parametrize( + "key,hf_name", + [("gate", "gate_proj"), ("in", "up_proj"), ("out", "down_proj")], + ) + def test_projection(self, mlp: GatedMLPBridge, key: str, hf_name: str) -> None: + proj = mlp.submodules[key] + assert isinstance(proj, LinearBridge) + assert proj.name == hf_name + + +# --------------------------------------------------------------------------- +# Factory + registry registration +# --------------------------------------------------------------------------- + + +class TestFalconH1FactoryRegistration: + def test_factory_returns_falcon_h1_adapter(self) -> None: + adapter = ArchitectureAdapterFactory.select_architecture_adapter(_make_cfg()) + assert isinstance(adapter, FalconH1ArchitectureAdapter) + + def test_architecture_key_present(self) -> None: + from transformer_lens.factories.architecture_adapter_factory import ( + SUPPORTED_ARCHITECTURES, + ) + + assert SUPPORTED_ARCHITECTURES["FalconH1ForCausalLM"] is FalconH1ArchitectureAdapter + + +class TestFalconH1ModelRegistry: + def test_in_hf_supported_architectures(self) -> None: + from transformer_lens.tools.model_registry import HF_SUPPORTED_ARCHITECTURES + + assert "FalconH1ForCausalLM" in HF_SUPPORTED_ARCHITECTURES + + def test_canonical_author_is_tiiuae(self) -> None: + from transformer_lens.tools.model_registry import CANONICAL_AUTHORS_BY_ARCH + + assert CANONICAL_AUTHORS_BY_ARCH.get("FalconH1ForCausalLM") == ["tiiuae"] + + +# --------------------------------------------------------------------------- +# Guard tests +# --------------------------------------------------------------------------- + + +class TestFalconH1Guards: + """Guards against drift toward neighbouring adapter patterns.""" + + def test_uses_block_bridge_not_ssm_block_bridge( + self, adapter: FalconH1ArchitectureAdapter + ) -> None: + # Two-norm parallel hybrid → BlockBridge (real resid_mid), unlike the + # single-norm SSMBlockBridge used by NemotronH / Mamba2. + from transformer_lens.model_bridge.generalized_components import SSMBlockBridge + + blocks = adapter.component_mapping["blocks"] + assert isinstance(blocks, BlockBridge) + assert not isinstance(blocks, SSMBlockBridge) + + def test_block_has_ln2(self, adapter: FalconH1ArchitectureAdapter) -> None: + # ln2 (pre_ff_layernorm) must exist or BlockBridge raises at construction. + assert "ln2" in adapter.component_mapping["blocks"].submodules + + def test_no_weight_conversions_defined(self, adapter: FalconH1ArchitectureAdapter) -> None: + # Passthrough mode → HF applies the ~12 multipliers; no folding/rearrange. + assert len(adapter.weight_processing_conversions) == 0 diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py index cf18ffcdd..640ad42ca 100644 --- a/transformer_lens/factories/architecture_adapter_factory.py +++ b/transformer_lens/factories/architecture_adapter_factory.py @@ -21,6 +21,7 @@ DeepSeekV2ArchitectureAdapter, DeepSeekV3ArchitectureAdapter, FalconArchitectureAdapter, + FalconH1ArchitectureAdapter, Gemma1ArchitectureAdapter, Gemma2ArchitectureAdapter, Gemma3ArchitectureAdapter, @@ -95,6 +96,7 @@ "DeepseekV2ForCausalLM": DeepSeekV2ArchitectureAdapter, "DeepseekV3ForCausalLM": DeepSeekV3ArchitectureAdapter, "FalconForCausalLM": FalconArchitectureAdapter, + "FalconH1ForCausalLM": FalconH1ArchitectureAdapter, "GemmaForCausalLM": Gemma1ArchitectureAdapter, # Default to Gemma1 as it's the original version "Gemma1ForCausalLM": Gemma1ArchitectureAdapter, "Gemma2ForCausalLM": Gemma2ArchitectureAdapter, diff --git a/transformer_lens/model_bridge/sources/_bridge_builder.py b/transformer_lens/model_bridge/sources/_bridge_builder.py index 05e9fd5bf..1a2d516cc 100644 --- a/transformer_lens/model_bridge/sources/_bridge_builder.py +++ b/transformer_lens/model_bridge/sources/_bridge_builder.py @@ -44,6 +44,14 @@ # Mamba-2 (additional SSM config) "n_groups", "chunk_size", + # Falcon-H1 (parallel attn + Mamba-2 hybrid SSM config) + "mamba_d_ssm", + "mamba_n_heads", + "mamba_d_head", + "mamba_d_state", + "mamba_n_groups", + "mamba_d_conv", + "mamba_chunk_size", # Multimodal "vision_config", # Cohere diff --git a/transformer_lens/model_bridge/sources/transformers.py b/transformer_lens/model_bridge/sources/transformers.py index 1e6e6c90f..2f9f7c288 100644 --- a/transformer_lens/model_bridge/sources/transformers.py +++ b/transformer_lens/model_bridge/sources/transformers.py @@ -539,6 +539,14 @@ def boot( # Mamba-2 (additional SSM config) "n_groups", "chunk_size", + # Falcon-H1 (parallel attn + Mamba-2 hybrid SSM config) + "mamba_d_ssm", + "mamba_n_heads", + "mamba_d_head", + "mamba_d_state", + "mamba_n_groups", + "mamba_d_conv", + "mamba_chunk_size", # Multimodal "vision_config", # Cohere diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py index ba5b56430..5c499c07e 100644 --- a/transformer_lens/model_bridge/supported_architectures/__init__.py +++ b/transformer_lens/model_bridge/supported_architectures/__init__.py @@ -36,6 +36,9 @@ from transformer_lens.model_bridge.supported_architectures.falcon import ( FalconArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.falcon_h1 import ( + FalconH1ArchitectureAdapter, +) from transformer_lens.model_bridge.supported_architectures.gemma1 import ( Gemma1ArchitectureAdapter, ) @@ -223,6 +226,7 @@ "DeepSeekV2ArchitectureAdapter", "DeepSeekV3ArchitectureAdapter", "FalconArchitectureAdapter", + "FalconH1ArchitectureAdapter", "Gemma1ArchitectureAdapter", "Gemma2ArchitectureAdapter", "Gemma3ArchitectureAdapter", diff --git a/transformer_lens/model_bridge/supported_architectures/falcon_h1.py b/transformer_lens/model_bridge/supported_architectures/falcon_h1.py new file mode 100644 index 000000000..4154c1d97 --- /dev/null +++ b/transformer_lens/model_bridge/supported_architectures/falcon_h1.py @@ -0,0 +1,190 @@ +"""Falcon-H1 parallel-hybrid architecture adapter. + +Supports ``FalconH1ForCausalLM`` (e.g. ``tiiuae/Falcon-H1-0.5B-Base``). + +Architecture overview: +- Every block runs **GQA attention and a Mamba-2 mixer in parallel** on the same + ``input_layernorm`` output; their results are summed into the residual stream + alongside scalar multipliers. A second norm (``pre_ff_layernorm``) precedes a + SwiGLU feed-forward MLP. This is the parallel hybrid that distinguishes + Falcon-H1 from heterogeneous-layer hybrids like NemotronH (one mixer type per + layer) — here both branches are present in *every* block. +- RoPE via a model-level ``model.rotary_emb`` module (very large ``rope_theta``). +- ``mamba_rms_norm=false`` on the released checkpoints — the Mamba inner gated + RMSNorm is **absent**, so ``inner_norm`` is declared optional. +- ~12 scalar multipliers (``embedding_multiplier``, ``attention_out_multiplier``, + ``key_multiplier``, ``ssm_*_multiplier``, ``mlp_multipliers``, + ``lm_head_multiplier`` …). HF applies all of these in its own forward, so in + raw (passthrough) mode the bridge inherits them for free — no weight folding is + needed for forward parity. Folding would only matter for compatibility mode. + +Key adapter decisions: +- ``BlockBridge`` is the block container: the block has two norms + (``input_layernorm`` + ``pre_ff_layernorm``) and a real post-attention + residual, so transformer-shaped hooks (``hook_resid_mid`` via ``ln2``) are + meaningful — unlike single-norm SSM blocks that need ``SSMBlockBridge``. + ``BlockBridge.forward`` delegates the whole block to HF, so the parallel + attn+mamba combine happens natively. +- ``SSM2MixerBridge`` wraps the Mamba-2 mixer (passthrough forward). Its inner + ``in_proj`` / ``conv1d`` / ``out_proj`` are mapped for hookability; the gated + ``inner_norm`` is optional (absent when ``mamba_rms_norm=false``). +- ``applicable_phases = []``: ``verify_models`` is transformer-shaped and does + not meaningfully cover SSM hybrids. Correctness is gated by the integration + parity test instead (the standard set for Mamba/Mamba2/NemotronH). +- Generation uses the **standard transformer KV-cache path**, not the bridge's + stateful-Mamba loop. Falcon-H1's HF forward carries both attention KV and the + Mamba conv/recurrent state inside one unified ``past_key_values`` cache, so the + default path matches HF bit-for-bit. Setting ``is_stateful`` would instead + route generation through the pure-Mamba ``cache_params`` convention, which this + hybrid does not use, and diverges from HF after the first decode step. +""" + +from typing import Any + +from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter +from transformer_lens.model_bridge.generalized_components import ( + BlockBridge, + DepthwiseConv1DBridge, + EmbeddingBridge, + GatedMLPBridge, + GatedRMSNormBridge, + LinearBridge, + PositionEmbeddingsAttentionBridge, + RMSNormalizationBridge, + RotaryEmbeddingBridge, + SSM2MixerBridge, + UnembeddingBridge, +) +from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, +) + + +def _make_optional(component: GeneralizedComponent) -> GeneralizedComponent: + """Mark a GeneralizedComponent submodule as optional. + + Some bridge classes (e.g. ``GatedRMSNormBridge``) do not forward ``optional`` + through their own ``__init__`` even though ``GeneralizedComponent`` supports + it. Setting the attribute directly is safe because ``component_setup.py`` + reads ``getattr(submodule, "optional", False)`` at setup time. + """ + component.optional = True + return component + + +class FalconH1ArchitectureAdapter(ArchitectureAdapter): + """Architecture adapter for ``FalconH1ForCausalLM``. + + Parallel hybrid: every block runs GQA attention and a Mamba-2 mixer side by + side, then a SwiGLU MLP. Both branches are mapped on every block so each + sub-path is independently hookable for ablation studies. + """ + + # verify_models is transformer-shaped and would need a dedicated refactor to + # cover SSM hybrids. Forward-pass correctness lives in the integration test: + # tests/integration/model_bridge/test_falcon_h1_adapter.py + applicable_phases: list[int] = [] + + def __init__(self, cfg: Any) -> None: + super().__init__(cfg) + + self.cfg.normalization_type = "RMS" + self.cfg.uses_rms_norm = True + self.cfg.positional_embedding_type = "rotary" + self.cfg.final_rms = True + # SwiGLU feed-forward (gate_proj / up_proj / down_proj, silu activation). + self.cfg.gated_mlp = True + self.cfg.attn_only = False + # FalconH1RMSNorm stores its epsilon as `variance_epsilon` (like Llama). + setattr(self.cfg, "eps_attr", "variance_epsilon") + + if hasattr(cfg, "n_key_value_heads") and cfg.n_key_value_heads is not None: + self.cfg.n_key_value_heads = cfg.n_key_value_heads + + # Mamba-2 dimensional config. Falcon-H1 specifies the inner SSM width + # directly via `mamba_d_ssm` (= mamba_n_heads * mamba_d_head), unlike + # Mamba2's `expand`-derived intermediate size. conv_dim follows the + # Mamba-2 layout: inner width + 2 groups of state for B and C. + mamba_d_ssm = getattr(cfg, "mamba_d_ssm", self.cfg.d_model) + mamba_n_groups = getattr(cfg, "mamba_n_groups", 1) + mamba_d_state = getattr(cfg, "mamba_d_state", 128) + mamba_n_heads = getattr(cfg, "mamba_n_heads", 0) + conv_dim = mamba_d_ssm + 2 * mamba_n_groups * mamba_d_state + setattr(self.cfg, "mamba_intermediate_size", mamba_d_ssm) + setattr(self.cfg, "conv_dim", conv_dim) + # HF fuses gate, hidden_BC, and dt into one in_proj output; stored so a + # future HF layout change is caught by the integration test. + setattr( + self.cfg, "expected_in_proj_out_features", 2 * mamba_d_ssm + conv_dim + mamba_n_heads + ) + + # Raw (passthrough) mode delegates the full forward to HF, which applies + # all of Falcon-H1's scalar multipliers natively — no weight folding for + # forward parity. + self.weight_processing_conversions = {} + + self.component_mapping = { + "embed": EmbeddingBridge(name="model.embed_tokens"), + "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg), + "blocks": BlockBridge( + name="model.layers", + submodules={ + "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), + "ln2": RMSNormalizationBridge(name="pre_ff_layernorm", config=self.cfg), + "attn": PositionEmbeddingsAttentionBridge( + name="self_attn", + config=self.cfg, + submodules={ + "q": LinearBridge(name="q_proj"), + "k": LinearBridge(name="k_proj"), + "v": LinearBridge(name="v_proj"), + "o": LinearBridge(name="o_proj"), + }, + requires_attention_mask=True, + requires_position_embeddings=True, + ), + "mamba": SSM2MixerBridge( + name="mamba", + config=self.cfg, + submodules={ + "in_proj": LinearBridge(name="in_proj"), + "conv1d": DepthwiseConv1DBridge(name="conv1d"), + # HF names the inner gated norm "norm"; TL aliases it + # to "inner_norm" to disambiguate from the block norm. + # Absent when mamba_rms_norm=false, so optional. + "inner_norm": _make_optional(GatedRMSNormBridge(name="norm")), + "out_proj": LinearBridge(name="out_proj"), + }, + ), + "mlp": GatedMLPBridge( + name="feed_forward", + config=self.cfg, + submodules={ + "gate": LinearBridge(name="gate_proj"), + "in": LinearBridge(name="up_proj"), + "out": LinearBridge(name="down_proj"), + }, + ), + }, + ), + "ln_final": RMSNormalizationBridge(name="model.final_layernorm", config=self.cfg), + "unembed": UnembeddingBridge(name="lm_head", config=self.cfg), + } + + def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None: + """Wire the model-level rotary embedding onto the attention bridges.""" + if not hasattr(hf_model.model, "rotary_emb"): + return + + rotary_emb = hf_model.model.rotary_emb + + if bridge_model is not None and hasattr(bridge_model, "blocks"): + for block in bridge_model.blocks: + if hasattr(block, "attn"): + block.attn.set_rotary_emb(rotary_emb) + + try: + attn_bridge = self.get_generalized_component("blocks.0.attn") + attn_bridge.set_rotary_emb(rotary_emb) + except (AttributeError, KeyError, ValueError): + pass diff --git a/transformer_lens/tools/model_registry/__init__.py b/transformer_lens/tools/model_registry/__init__.py index b0be24999..50d3a4914 100644 --- a/transformer_lens/tools/model_registry/__init__.py +++ b/transformer_lens/tools/model_registry/__init__.py @@ -57,6 +57,7 @@ "DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM", "FalconForCausalLM", + "FalconH1ForCausalLM", "GemmaForCausalLM", "Gemma2ForCausalLM", "Gemma3ForCausalLM", @@ -130,6 +131,7 @@ "DeepseekV2ForCausalLM": ["deepseek-ai"], "DeepseekV3ForCausalLM": ["deepseek-ai"], "FalconForCausalLM": ["tiiuae"], + "FalconH1ForCausalLM": ["tiiuae"], "Gemma2ForCausalLM": ["google"], "Gemma3ForCausalLM": ["google"], "Gemma3ForConditionalGeneration": ["google"], diff --git a/transformer_lens/tools/model_registry/data/supported_models.json b/transformer_lens/tools/model_registry/data/supported_models.json index b82100137..399603b39 100644 --- a/transformer_lens/tools/model_registry/data/supported_models.json +++ b/transformer_lens/tools/model_registry/data/supported_models.json @@ -1,15 +1,99 @@ { - "generated_at": "2026-07-01", + "generated_at": "2026-06-25", "scan_info": { - "total_scanned": 5855, + "total_scanned": 5949, "task_filter": "text-generation", "min_downloads": 500, - "scan_duration_seconds": 7.0 + "scan_duration_seconds": 8.1 }, - "total_architectures": 68, - "total_models": 13138, + "total_architectures": 69, + "total_models": 13144, "total_verified": 1028, "models": [ + { + "architecture_id": "FalconH1ForCausalLM", + "model_id": "tiiuae/Falcon-H1-Tiny-90M-Instruct", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "FalconH1ForCausalLM", + "model_id": "tiiuae/Falcon-H1-0.5B-Base", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "FalconH1ForCausalLM", + "model_id": "tiiuae/Falcon-H1-0.5B-Instruct", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "FalconH1ForCausalLM", + "model_id": "tiiuae/Falcon-H1-1.5B-Base", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "FalconH1ForCausalLM", + "model_id": "tiiuae/Falcon-H1-3B-Base", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "FalconH1ForCausalLM", + "model_id": "tiiuae/Falcon-H1-7B-Base", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, { "architecture_id": "HunYuanDenseV1ForCausalLM", "model_id": "tencent/Hunyuan-0.5B-Instruct", @@ -179207,6 +179291,48 @@ "phase7_score": null, "phase8_score": null }, + { + "architecture_id": "Qwen2MoeForCausalLM", + "model_id": "hyper-accel/ci-random-qwen2-moe-a3b", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "ArceeForCausalLM", + "model_id": "optimum-intel-internal-testing/tiny-random-ArceeForCausalLM", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, + { + "architecture_id": "ArceeForCausalLM", + "model_id": "arcee-ai/AFM-4.5B-Base", + "status": 0, + "verified_date": null, + "metadata": null, + "note": null, + "phase1_score": null, + "phase2_score": null, + "phase3_score": null, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null + }, { "architecture_id": "NemotronHForCausalLM", "model_id": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", @@ -182510,48 +182636,6 @@ "phase4_score": null, "phase7_score": null, "phase8_score": null - }, - { - "architecture_id": "Qwen2MoeForCausalLM", - "model_id": "hyper-accel/ci-random-qwen2-moe-a3b", - "status": 0, - "verified_date": null, - "metadata": null, - "note": null, - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": null, - "phase7_score": null, - "phase8_score": null - }, - { - "architecture_id": "ArceeForCausalLM", - "model_id": "optimum-intel-internal-testing/tiny-random-ArceeForCausalLM", - "status": 0, - "verified_date": null, - "metadata": null, - "note": null, - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": null, - "phase7_score": null, - "phase8_score": null - }, - { - "architecture_id": "ArceeForCausalLM", - "model_id": "arcee-ai/AFM-4.5B-Base", - "status": 0, - "verified_date": null, - "metadata": null, - "note": null, - "phase1_score": null, - "phase2_score": null, - "phase3_score": null, - "phase4_score": null, - "phase7_score": null, - "phase8_score": null } ] } diff --git a/transformer_lens/tools/model_registry/generate_report.py b/transformer_lens/tools/model_registry/generate_report.py index 5e7f4f065..62b5e981d 100644 --- a/transformer_lens/tools/model_registry/generate_report.py +++ b/transformer_lens/tools/model_registry/generate_report.py @@ -53,6 +53,7 @@ "Phi3ForCausalLM": "Microsoft's Phi-3 improved small model", "PhiMoEForCausalLM": "Microsoft's Phi sparse Mixture of Experts model", "FalconForCausalLM": "TII's Falcon model series", + "FalconH1ForCausalLM": "TII's Falcon-H1 parallel attention + Mamba-2 hybrid series", "OlmoForCausalLM": "Allen AI's OLMo open language model", "Olmo2ForCausalLM": "Allen AI's OLMo 2 with improved training", "Olmo3ForCausalLM": "Allen AI's OLMo 3 latest generation", From 2fa5e5d9849382e2e115cad61880eb3fbb001112 Mon Sep 17 00:00:00 2001 From: Puranik Yashaswin Sharma Date: Mon, 6 Jul 2026 18:48:21 +0530 Subject: [PATCH 05/11] Add BD3LM architecture adapter (#1479) * Add BD3LM architecture adapter * Address review: fix dead hook_attn_out, elif ordering, add integration test and verify_models support --------- Co-authored-by: Jonah Larson --- .../model_bridge/test_bd3lm_adapter.py | 130 +++++++ .../test_bd3lm_adapter.py | 359 ++++++++++++++++++ .../factories/architecture_adapter_factory.py | 2 + .../model_bridge/sources/_bridge_builder.py | 6 + .../model_bridge/sources/transformers.py | 10 + .../supported_architectures/__init__.py | 4 + .../supported_architectures/bd3lm.py | 210 ++++++++++ .../tools/model_registry/__init__.py | 2 + .../tools/model_registry/generate_report.py | 1 + .../tools/model_registry/verify_models.py | 1 + transformer_lens/utilities/architectures.py | 1 + 11 files changed, 726 insertions(+) create mode 100644 tests/integration/model_bridge/test_bd3lm_adapter.py create mode 100644 tests/unit/model_bridge/supported_architectures/test_bd3lm_adapter.py create mode 100644 transformer_lens/model_bridge/supported_architectures/bd3lm.py diff --git a/tests/integration/model_bridge/test_bd3lm_adapter.py b/tests/integration/model_bridge/test_bd3lm_adapter.py new file mode 100644 index 000000000..53d389ff6 --- /dev/null +++ b/tests/integration/model_bridge/test_bd3lm_adapter.py @@ -0,0 +1,130 @@ +"""Integration tests for the BD3LM architecture adapter. + +Verifies forward-pass and hook parity against kuleshov-group/bd3lm-owt-block_size4: +- Forward-pass logits match HF exactly (bridge delegates the full forward to HF) +- Sanity checks: config flags, block count, hook coverage + +Note: runs on slow mark to avoid hitting HF on standard runs. +""" + +import gc + +import pytest +import torch + +from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.supported_architectures.bd3lm import ( + BD3LMArchitectureAdapter, +) + +pytestmark = pytest.mark.slow + +MODEL = "kuleshov-group/bd3lm-owt-block_size4" + + +def _device() -> str: + return "cuda" if torch.cuda.is_available() else "cpu" + + +@pytest.fixture(scope="module") +def bd3lm_bridge(): + device = _device() + # Override model_length to 4 so that we can run sequence length 8 (model_length * 2) forward passes + bridge = TransformerBridge.boot_transformers( + MODEL, + device=device, + trust_remote_code=True, + hf_config_overrides={"model_length": 4}, + ) + yield bridge + # Cleanup + del bridge + if torch.cuda.is_available(): + torch.cuda.empty_cache() + for _ in range(3): + gc.collect() + + +class TestBD3LMBridgeCreation: + """Smoke-test that the bridge loads with the right config flags.""" + + def test_block_count(self, bd3lm_bridge: TransformerBridge) -> None: + assert len(bd3lm_bridge.blocks) == 12 + + def test_adapter_type(self, bd3lm_bridge: TransformerBridge) -> None: + assert isinstance(bd3lm_bridge.adapter, BD3LMArchitectureAdapter) + + +class TestBD3LMForwardPass: + """Bridge logits must match HF logits exactly.""" + + @pytest.fixture(scope="class") + def tokens(self) -> torch.Tensor: + # Sequence length must be model_length * 2 = 8 + return torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]]) + + def test_forward_matches_hf_exactly( + self, bd3lm_bridge: TransformerBridge, tokens: torch.Tensor + ) -> None: + tokens = tokens.to(_device()) + hf_model = bd3lm_bridge.original_model + + timesteps = torch.tensor([10.0]).to(_device()) + + with torch.no_grad(): + bridge_out = bd3lm_bridge(tokens, timesteps=timesteps) + hf_out = hf_model(tokens, timesteps=timesteps) + + if hasattr(hf_out, "logits"): + hf_logits = hf_out.logits + else: + hf_logits = hf_out + + max_diff = (bridge_out.float() - hf_logits.float()).abs().max().item() + assert max_diff == 0.0, ( + f"Bridge vs HF forward max diff = {max_diff:.2e}. " + "Expected 0 because DelegatedAttentionBlockBridge.forward() is a pure passthrough." + ) + + +class TestBD3LMHookCoverage: + """run_with_cache captures residual stream, attention, MLP, and other hooks.""" + + @pytest.fixture(scope="class") + def cache(self, bd3lm_bridge: TransformerBridge): + # Sequence length must be model_length * 2 = 8 + tokens = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]]).to(_device()) + timesteps = torch.tensor([10.0]).to(_device()) + with torch.no_grad(): + _, cache = bd3lm_bridge.run_with_cache(tokens, timesteps=timesteps) + return cache + + def test_block_hooks_fire(self, cache, bd3lm_bridge: TransformerBridge) -> None: + for i in range(len(bd3lm_bridge.blocks)): + assert f"blocks.{i}.hook_in" in cache + assert f"blocks.{i}.hook_out" in cache + + def test_expected_hook_aliases_fire(self, cache, bd3lm_bridge: TransformerBridge) -> None: + """Verify that every registered hook alias in the block actually fires with real data.""" + for i in range(len(bd3lm_bridge.blocks)): + for alias in ( + "hook_resid_pre", + "hook_resid_mid", + "hook_resid_post", + "hook_attn_out", + "hook_mlp_out", + ): + key = f"blocks.{i}.{alias}" + assert key in cache, f"Missing hook alias {key} in cache" + val = cache[key] + assert val is not None + assert isinstance(val, torch.Tensor) + assert val.shape[0] == 1 # batch size + assert val.shape[1] == 8 # seq length (model_length * 2) + assert val.shape[2] == bd3lm_bridge.cfg.d_model # d_model + assert not torch.isnan(val).any(), f"NaN in cache['{key}']" + + def test_no_nan_in_cache(self, cache) -> None: + for key, val in cache.items(): + if isinstance(val, torch.Tensor) and val.is_floating_point(): + assert not torch.isnan(val).any(), f"NaN in cache['{key}']" diff --git a/tests/unit/model_bridge/supported_architectures/test_bd3lm_adapter.py b/tests/unit/model_bridge/supported_architectures/test_bd3lm_adapter.py new file mode 100644 index 000000000..01873d72f --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_bd3lm_adapter.py @@ -0,0 +1,359 @@ +"""Unit tests for BD3LMArchitectureAdapter. + +Covers: config attribute propagation, component mapping bridge types and HF +path names, applicable_phases, setup patching, and factory registration. +""" + +from unittest.mock import MagicMock + +import pytest +import torch + +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.factories.architecture_adapter_factory import ( + ArchitectureAdapterFactory, +) +from transformer_lens.model_bridge.generalized_components import ( + DelegatedAttentionBlockBridge, + EmbeddingBridge, + LinearBridge, + MLPBridge, + NormalizationBridge, + SymbolicBridge, + UnembeddingBridge, +) +from transformer_lens.model_bridge.supported_architectures.bd3lm import ( + BD3LMArchitectureAdapter, +) + + +def _make_cfg( + n_layers: int = 2, + d_model: int = 64, + d_head: int = 8, + n_heads: int = 8, + d_vocab: int = 100, + n_ctx: int = 128, + block_size: int = 4, + cond_dim: int = 128, + adaln: bool = True, + cross_attn: bool = True, +) -> TransformerBridgeConfig: + """Minimal TransformerBridgeConfig for BD3LM adapter tests.""" + cfg = TransformerBridgeConfig( + d_model=d_model, + d_head=d_head, + n_layers=n_layers, + n_ctx=n_ctx, + n_heads=n_heads, + d_vocab=d_vocab, + default_prepend_bos=False, + architecture="BD3LM", + ) + # Inject BD3LM-specific fields + cfg.block_size = block_size # type: ignore[attr-defined] + cfg.cond_dim = cond_dim # type: ignore[attr-defined] + cfg.adaln = adaln # type: ignore[attr-defined] + cfg.cross_attn = cross_attn # type: ignore[attr-defined] + return cfg + + +@pytest.fixture(scope="class") +def cfg() -> TransformerBridgeConfig: + return _make_cfg() + + +@pytest.fixture(scope="class") +def adapter(cfg: TransformerBridgeConfig) -> BD3LMArchitectureAdapter: + return BD3LMArchitectureAdapter(cfg) + + +# --------------------------------------------------------------------------- +# Config attributes +# --------------------------------------------------------------------------- + + +class TestBD3LMArchitectureAdapterConfig: + """Adapter propagates all required config attributes.""" + + def test_normalization_type(self, adapter: BD3LMArchitectureAdapter) -> None: + assert adapter.cfg.normalization_type == "LN" + + def test_uses_rms_norm(self, adapter: BD3LMArchitectureAdapter) -> None: + assert adapter.cfg.uses_rms_norm is False + + def test_positional_embedding_type(self, adapter: BD3LMArchitectureAdapter) -> None: + assert adapter.cfg.positional_embedding_type == "none" + + def test_gated_mlp_false(self, adapter: BD3LMArchitectureAdapter) -> None: + assert adapter.cfg.gated_mlp is False + + def test_final_rms_false(self, adapter: BD3LMArchitectureAdapter) -> None: + assert adapter.cfg.final_rms is False + + def test_applicable_phases(self, adapter: BD3LMArchitectureAdapter) -> None: + # BD3LM now supports all phases correctly + assert adapter.applicable_phases == [1, 2, 3] + + def test_supports_generation(self, adapter: BD3LMArchitectureAdapter) -> None: + assert adapter.supports_generation is False + + +# --------------------------------------------------------------------------- +# Component Mapping +# --------------------------------------------------------------------------- + + +class TestBD3LMArchitectureAdapterComponentMapping: + """Adapter maps all components to the correct HF submodules.""" + + def test_embed_bridge(self, adapter: BD3LMArchitectureAdapter) -> None: + assert adapter.component_mapping is not None + embed = adapter.component_mapping["embed"] + assert isinstance(embed, EmbeddingBridge) + assert embed.name == "backbone.vocab_embed" + + def test_blocks_bridge(self, adapter: BD3LMArchitectureAdapter) -> None: + assert adapter.component_mapping is not None + blocks = adapter.component_mapping["blocks"] + assert isinstance(blocks, DelegatedAttentionBlockBridge) + assert blocks.name == "backbone.blocks" + + def test_blocks_submodules(self, adapter: BD3LMArchitectureAdapter) -> None: + assert adapter.component_mapping is not None + blocks = adapter.component_mapping["blocks"] + assert isinstance(blocks, DelegatedAttentionBlockBridge) + sub = blocks.submodules + assert sub is not None + + assert isinstance(sub["ln1"], NormalizationBridge) + assert sub["ln1"].name == "norm1" + + assert isinstance(sub["ln2"], NormalizationBridge) + assert sub["ln2"].name == "norm2" + + assert isinstance(sub["adaln_modulation"], LinearBridge) + assert sub["adaln_modulation"].name == "adaLN_modulation" + + attn = sub["attn"] + assert isinstance(attn, SymbolicBridge) + assert attn.submodules is not None + assert isinstance(attn.submodules["qkv"], LinearBridge) + assert attn.submodules["qkv"].name == "attn_qkv" + assert isinstance(attn.submodules["o"], LinearBridge) + assert attn.submodules["o"].name == "attn_out" + + mlp = sub["mlp"] + assert isinstance(mlp, MLPBridge) + assert mlp.submodules is not None + assert isinstance(mlp.submodules["in"], LinearBridge) + assert mlp.submodules["in"].name == "0" + assert isinstance(mlp.submodules["out"], LinearBridge) + assert mlp.submodules["out"].name == "2" + + def test_sigma_map_bridge(self, adapter: BD3LMArchitectureAdapter) -> None: + assert adapter.component_mapping is not None + sigma_map = adapter.component_mapping["sigma_map"] + assert isinstance(sigma_map, MLPBridge) + assert sigma_map.name == "backbone.sigma_map.mlp" + assert sigma_map.submodules is not None + assert isinstance(sigma_map.submodules["in"], LinearBridge) + assert sigma_map.submodules["in"].name == "0" + assert isinstance(sigma_map.submodules["out"], LinearBridge) + assert sigma_map.submodules["out"].name == "2" + + def test_ln_final_bridge(self, adapter: BD3LMArchitectureAdapter) -> None: + assert adapter.component_mapping is not None + ln_final = adapter.component_mapping["ln_final"] + assert isinstance(ln_final, NormalizationBridge) + assert ln_final.name == "backbone.output_layer.norm_final" + + def test_unembed_bridge(self, adapter: BD3LMArchitectureAdapter) -> None: + assert adapter.component_mapping is not None + unembed = adapter.component_mapping["unembed"] + assert isinstance(unembed, UnembeddingBridge) + assert unembed.name == "backbone.output_layer.linear" + + def test_component_paths_resolve_to_real_submodules( + self, adapter: BD3LMArchitectureAdapter + ) -> None: + # Create a mock model matching the real BD3LM structure + import torch.nn as nn + + class MockBlock(nn.Module): + def __init__(self): + super().__init__() + self.norm1 = nn.LayerNorm(64) + self.attn_qkv = nn.Linear(64, 64) + self.attn_out = nn.Linear(64, 64) + self.adaLN_modulation = nn.Linear(128, 6 * 64) + self.norm2 = nn.LayerNorm(64) + self.mlp = nn.Sequential(nn.Linear(64, 128), nn.GELU(), nn.Linear(128, 64)) + + class MockSigmaMap(nn.Module): + def __init__(self): + super().__init__() + self.mlp = nn.Sequential(nn.Linear(128, 256), nn.GELU(), nn.Linear(256, 128)) + + class MockOutputLayer(nn.Module): + def __init__(self): + super().__init__() + self.norm_final = nn.LayerNorm(64) + self.linear = nn.Linear(64, 100) + + class MockBackbone(nn.Module): + def __init__(self): + super().__init__() + + class EmbedLayer(nn.Module): + def __init__(self): + super().__init__() + self.embedding = nn.Parameter(torch.randn(100, 64)) + + self.vocab_embed = EmbedLayer() + self.blocks = nn.ModuleList([MockBlock() for _ in range(2)]) + self.sigma_map = MockSigmaMap() + self.output_layer = MockOutputLayer() + + class MockBD3LM(nn.Module): + def __init__(self): + super().__init__() + self.backbone = MockBackbone() + + model = MockBD3LM() + + # Prepare the model with the adapter + adapter.prepare_model(model) + + assert adapter.component_mapping is not None + + # Verify get_remote_component works for every top-level mapping + # 1. embed + embed_bridge = adapter.component_mapping["embed"] + assert embed_bridge.name is not None + embed_comp = adapter.get_remote_component(model, embed_bridge.name) + assert hasattr(embed_comp, "weight") + + # 2. blocks + blocks_bridge = adapter.component_mapping["blocks"] + assert blocks_bridge.name is not None + blocks_comp = adapter.get_remote_component(model, blocks_bridge.name) + assert isinstance(blocks_comp, nn.ModuleList) + assert len(blocks_comp) == 2 + + # Verify block submodules resolve + block0 = blocks_comp[0] + assert isinstance(blocks_bridge, DelegatedAttentionBlockBridge) + assert blocks_bridge.submodules is not None + for sub_name, sub_bridge in blocks_bridge.submodules.items(): + if sub_name == "attn": + assert sub_bridge.submodules is not None + for attn_sub_name, attn_sub_bridge in sub_bridge.submodules.items(): + assert attn_sub_bridge.name is not None + comp = adapter.get_remote_component(block0, attn_sub_bridge.name) + assert isinstance(comp, nn.Linear) + elif sub_name == "mlp": + assert sub_bridge.submodules is not None + assert sub_bridge.name is not None + in_bridge = sub_bridge.submodules["in"] + out_bridge = sub_bridge.submodules["out"] + assert in_bridge.name is not None + assert out_bridge.name is not None + comp_in = adapter.get_remote_component( + block0, f"{sub_bridge.name}.{in_bridge.name}" + ) + comp_out = adapter.get_remote_component( + block0, f"{sub_bridge.name}.{out_bridge.name}" + ) + assert isinstance(comp_in, nn.Linear) + assert isinstance(comp_out, nn.Linear) + else: + assert sub_bridge.name is not None + comp = adapter.get_remote_component(block0, sub_bridge.name) + assert isinstance(comp, (nn.LayerNorm, nn.Linear)) + + # 3. sigma_map + sigma_map_bridge = adapter.component_mapping["sigma_map"] + assert isinstance(sigma_map_bridge, MLPBridge) + assert sigma_map_bridge.name is not None + assert sigma_map_bridge.submodules is not None + sigma_map_comp = adapter.get_remote_component(model, sigma_map_bridge.name) + in_bridge = sigma_map_bridge.submodules["in"] + out_bridge = sigma_map_bridge.submodules["out"] + assert in_bridge.name is not None + assert out_bridge.name is not None + sigma_in = adapter.get_remote_component(sigma_map_comp, in_bridge.name) + sigma_out = adapter.get_remote_component(sigma_map_comp, out_bridge.name) + assert isinstance(sigma_in, nn.Linear) + assert isinstance(sigma_out, nn.Linear) + + # 4. ln_final + ln_final_bridge = adapter.component_mapping["ln_final"] + assert ln_final_bridge.name is not None + ln_final = adapter.get_remote_component(model, ln_final_bridge.name) + assert isinstance(ln_final, nn.LayerNorm) + + # 5. unembed + unembed_bridge = adapter.component_mapping["unembed"] + assert unembed_bridge.name is not None + unembed = adapter.get_remote_component(model, unembed_bridge.name) + assert isinstance(unembed, nn.Linear) + + +# --------------------------------------------------------------------------- +# Setup and Patches +# --------------------------------------------------------------------------- + + +class TestBD3LMArchitectureAdapterSetup: + """Tests the architecture setup and monkeypatches.""" + + def test_setup_patches_vocab_embed(self, adapter: BD3LMArchitectureAdapter) -> None: + # Mock model structure + model = MagicMock() + model.backbone = MagicMock() + model.backbone.vocab_embed = MagicMock() + # Pretend vocab_embed has embedding but not weight + del model.backbone.vocab_embed.weight + model.backbone.vocab_embed.embedding = torch.randn(10, 10) + + adapter.prepare_model(model) + + # Should have aliased weight + assert hasattr(model.backbone.vocab_embed, "weight") + assert model.backbone.vocab_embed.weight is model.backbone.vocab_embed.embedding + + def test_setup_patches_attn_backend_cpu( + self, adapter: BD3LMArchitectureAdapter, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + adapter.cfg.attn_backend = "flex" # type: ignore[attr-defined] + + model = MagicMock() + model.backbone = MagicMock() + model.backbone.block_size = 4 + block = MagicMock() + block.attn_backend = "flex" + model.backbone.blocks = [block] + + adapter.prepare_model(model) + + assert adapter.cfg.attn_backend == "sdpa" # type: ignore[attr-defined] + assert block.attn_backend == "sdpa" + model.backbone.gen_mask.assert_called_once_with( + getattr(adapter.cfg, "model_length", getattr(adapter.cfg, "n_ctx", 2048)), + 4, + attn_backend="sdpa", + ) + + +# --------------------------------------------------------------------------- +# Registry Factory +# --------------------------------------------------------------------------- + + +class TestBD3LMArchitectureAdapterRegistry: + def test_factory_creates_adapter(self, cfg: TransformerBridgeConfig) -> None: + """The factory returns the correct adapter type for BD3LM.""" + adapter = ArchitectureAdapterFactory.select_architecture_adapter(cfg) + assert isinstance(adapter, BD3LMArchitectureAdapter) diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py index 640ad42ca..a839c179b 100644 --- a/transformer_lens/factories/architecture_adapter_factory.py +++ b/transformer_lens/factories/architecture_adapter_factory.py @@ -14,6 +14,7 @@ ArceeArchitectureAdapter, BaichuanArchitectureAdapter, BartArchitectureAdapter, + BD3LMArchitectureAdapter, BertArchitectureAdapter, BloomArchitectureAdapter, CodeGenArchitectureAdapter, @@ -89,6 +90,7 @@ "BaiChuanForCausalLM": BaichuanArchitectureAdapter, "BaichuanForCausalLM": BaichuanArchitectureAdapter, "BartForConditionalGeneration": BartArchitectureAdapter, + "BD3LM": BD3LMArchitectureAdapter, "BertForMaskedLM": BertArchitectureAdapter, "BloomForCausalLM": BloomArchitectureAdapter, "CodeGenForCausalLM": CodeGenArchitectureAdapter, diff --git a/transformer_lens/model_bridge/sources/_bridge_builder.py b/transformer_lens/model_bridge/sources/_bridge_builder.py index 1a2d516cc..a74d2ef9d 100644 --- a/transformer_lens/model_bridge/sources/_bridge_builder.py +++ b/transformer_lens/model_bridge/sources/_bridge_builder.py @@ -67,6 +67,12 @@ "router_jitter_noise", "input_jitter_noise", "eos_token_id", + # BD3LM + "model_length", + "block_size", + "cond_dim", + "adaln", + "cross_attn", ] diff --git a/transformer_lens/model_bridge/sources/transformers.py b/transformer_lens/model_bridge/sources/transformers.py index 2f9f7c288..7929ee7bd 100644 --- a/transformer_lens/model_bridge/sources/transformers.py +++ b/transformer_lens/model_bridge/sources/transformers.py @@ -71,6 +71,8 @@ def map_default_transformer_lens_config(hf_config): tl_config.d_model = source_config.model_dim elif hasattr(source_config, "d_model"): tl_config.d_model = source_config.d_model + elif hasattr(source_config, "hidden_dim"): + tl_config.d_model = source_config.hidden_dim if hasattr(source_config, "n_head"): tl_config.n_heads = source_config.n_head elif hasattr(source_config, "num_attention_heads"): @@ -128,6 +130,8 @@ def map_default_transformer_lens_config(hf_config): tl_config.n_layers = source_config.num_transformer_layers elif hasattr(source_config, "num_layers"): tl_config.n_layers = source_config.num_layers + elif hasattr(source_config, "n_blocks"): + tl_config.n_layers = source_config.n_blocks if hasattr(source_config, "vocab_size") and isinstance(source_config.vocab_size, int): tl_config.d_vocab = source_config.vocab_size if hasattr(source_config, "n_positions"): @@ -562,6 +566,12 @@ def boot( "router_jitter_noise", "input_jitter_noise", "eos_token_id", + # BD3LM + "model_length", + "block_size", + "cond_dim", + "adaln", + "cross_attn", ] for attr in _HF_PASSTHROUGH_ATTRS: val = getattr(hf_config, attr, None) diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py index 5c499c07e..bdb2d7311 100644 --- a/transformer_lens/model_bridge/supported_architectures/__init__.py +++ b/transformer_lens/model_bridge/supported_architectures/__init__.py @@ -12,6 +12,9 @@ from transformer_lens.model_bridge.supported_architectures.baichuan import ( BaichuanArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.bd3lm import ( + BD3LMArchitectureAdapter, +) from transformer_lens.model_bridge.supported_architectures.bart import ( BartArchitectureAdapter, ) @@ -217,6 +220,7 @@ __all__ = [ "ApertusArchitectureAdapter", "ArceeArchitectureAdapter", + "BD3LMArchitectureAdapter", "BaichuanArchitectureAdapter", "BartArchitectureAdapter", "BertArchitectureAdapter", diff --git a/transformer_lens/model_bridge/supported_architectures/bd3lm.py b/transformer_lens/model_bridge/supported_architectures/bd3lm.py new file mode 100644 index 000000000..9642d1b99 --- /dev/null +++ b/transformer_lens/model_bridge/supported_architectures/bd3lm.py @@ -0,0 +1,210 @@ +"""BD3LM (Block Diffusion Language Model) architecture adapter.""" + +from typing import Any + +import torch + +from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter +from transformer_lens.model_bridge.generalized_components import ( + EmbeddingBridge, + LinearBridge, + MLPBridge, + NormalizationBridge, + SymbolicBridge, + UnembeddingBridge, +) +from transformer_lens.model_bridge.generalized_components.block import ( + DelegatedAttentionBlockBridge, +) + + +class BD3LMArchitectureAdapter(ArchitectureAdapter): + """Architecture adapter for BD3LM (Block Diffusion LM, ICLR 2025). + + BD3LM uses adaLN conditioning on diffusion timesteps, a custom Rotary + embedding, joint QKV projections, and non-causal block-diffusion masking. + Because adaLN modulation varies per-timestep, it cannot be folded into + weights — the adapter uses ``DelegatedAttentionBlockBridge`` to delegate + each ``DDiTBlock.forward()`` wholesale to the original HF module. + Hooks fire at block boundaries and on mapped submodules. + """ + + # Phases 1–3 cover component mapping, weight conversion, and forward-pass + # parity. Phase 4 (autoregressive generation) is excluded because BD3LM + # uses iterative diffusion sampling. + applicable_phases: list[int] = [1, 2, 3] + + # BD3LM uses diffusion sampling, not autoregressive generation. + supports_generation: bool = False + + def __init__(self, cfg: Any) -> None: + super().__init__(cfg) + + # ── Config attributes ────────────────────────────────────────── + self.cfg.normalization_type = "LN" + self.cfg.uses_rms_norm = False + self.cfg.positional_embedding_type = "none" # custom Rotary, not HF + self.cfg.gated_mlp = False # standard GELU MLP, not gated + self.cfg.attn_only = False + self.cfg.final_rms = False # final norm is custom LayerNorm + self.cfg.tokenizer_name = "gpt2" + + # BD3LM-specific config fields. These live on the HF config and are + # forwarded via _HF_PASSTHROUGH_ATTRS; we also store them explicitly + # so unit tests can assert them without loading a real model. + block_size = getattr(self.cfg, "block_size", 4) + setattr(self.cfg, "block_size", block_size) + + cond_dim = getattr(self.cfg, "cond_dim", 128) + setattr(self.cfg, "cond_dim", cond_dim) + + adaln = getattr(self.cfg, "adaln", True) + setattr(self.cfg, "adaln", adaln) + + cross_attn = getattr(self.cfg, "cross_attn", True) + setattr(self.cfg, "cross_attn", cross_attn) + + # Compute d_mlp from the MLP ratio (hardcoded to 4 in DDiTBlock). + mlp_ratio = 4 + d_mlp = mlp_ratio * self.cfg.d_model + setattr(self.cfg, "d_mlp", d_mlp) + + # ── Weight processing conversions ────────────────────────────── + # Wrap-don't-reimplement: no weight rearrangement needed since we + # delegate forward to the original modules. + self.weight_processing_conversions = {} + + # ── Component mapping ────────────────────────────────────────── + # DelegatedAttentionBlockBridge delegates forward() wholesale to the + # original DDiTBlock (wrap-don't-reimplement), but exposes standard + # attn/mlp-shaped hook aliases (hook_resid_mid, hook_mlp_in) instead + # of the SSM-specific hook_mixer_in alias that SSMBlockBridge would + # have incorrectly implied for this attn+mlp architecture. Submodule + # bridges map to the HF module paths relative to each block. + # + # Module hierarchy (HF paths relative to backbone.blocks[i]): + # norm1 → pre-attention LayerNorm + # attn_qkv → joint Q/K/V projection (no bias) + # attn_out → output projection (no bias) + # adaLN_modulation → conditioning linear (cond_dim → 6*dim) + # norm2 → pre-MLP LayerNorm + # mlp.0 → MLP in-projection (bias) + # mlp.2 → MLP out-projection (bias) + self.component_mapping = { + "embed": EmbeddingBridge(name="backbone.vocab_embed"), + "blocks": DelegatedAttentionBlockBridge( + name="backbone.blocks", + submodules={ + "ln1": NormalizationBridge(name="norm1", config=self.cfg), + "attn": SymbolicBridge( + submodules={ + "qkv": LinearBridge(name="attn_qkv"), + "o": LinearBridge(name="attn_out"), + }, + ), + "adaln_modulation": LinearBridge(name="adaLN_modulation"), + "ln2": NormalizationBridge(name="norm2", config=self.cfg), + "mlp": MLPBridge( + name="mlp", + config=self.cfg, + submodules={ + "in": LinearBridge(name="0"), + "out": LinearBridge(name="2"), + }, + ), + }, + # hook_attn_out captures the raw projection, not the gate_msa-scaled + # value added to residual stream, as gating is fused inside a + # torch.jit.script function with no hookable module boundary. + hook_alias_overrides={ + "hook_attn_out": "attn.o.hook_out", + }, + ), + "sigma_map": MLPBridge( + name="backbone.sigma_map.mlp", + config=self.cfg, + submodules={ + "in": LinearBridge(name="0"), + "out": LinearBridge(name="2"), + }, + ), + "ln_final": NormalizationBridge( + name="backbone.output_layer.norm_final", config=self.cfg + ), + "unembed": UnembeddingBridge(name="backbone.output_layer.linear"), + } + + def prepare_loading(self, model_name: str, model_kwargs: dict) -> None: + """Patch BD3LM dynamic class before from_pretrained runs. + + Modeling code has a custom __getattr__ that fails to delegate back + to PreTrainedModel, raising AttributeError on all_tied_weights_keys. + """ + try: + from transformers.dynamic_module_utils import get_class_from_dynamic_module + + model_class = get_class_from_dynamic_module("modeling_bd3lm.BD3LM", model_name) + setattr(model_class, "all_tied_weights_keys", {}) + except Exception: + # Best-effort patch: should never block loading if dynamic module is missing/modified. + pass + + def prepare_model(self, hf_model: Any) -> None: + """Patch BD3LM quirks that prevent standard bridge construction. + + Three issues must be fixed before the bridge can wrap the model: + + 1. ``vocab_embed`` is an ``nn.Parameter``, not ``nn.Embedding``, so it + lacks a ``.weight`` attribute that ``EmbeddingBridge`` expects. + 2. The ``flex`` attention backend crashes on CPU; fall back to ``sdpa`` + and regenerate ``block_diff_mask`` for the new backend. + 3. The HF ``forward()`` does not accept ``output_attentions`` and other + kwargs the bridge unconditionally injects; patch at runtime because + no other hook point allows filtering them before HF's forward call. + """ + # Patch vocab_embed to have a weight attribute for EmbeddingBridge + if hasattr(hf_model, "backbone") and hasattr(hf_model.backbone, "vocab_embed"): + embed_mod = hf_model.backbone.vocab_embed + if not hasattr(embed_mod, "weight") and hasattr(embed_mod, "embedding"): + embed_mod.weight = embed_mod.embedding + + # Fix attention backend for CPU and ensure mask is on a real device + if hasattr(hf_model, "backbone"): + if not torch.cuda.is_available(): + backend = "sdpa" + setattr(self.cfg, "attn_backend", "sdpa") + for b in hf_model.backbone.blocks: + b.attn_backend = "sdpa" + else: + first_block = hf_model.backbone.blocks[0] + backend = getattr(first_block, "attn_backend", "sdpa") + + if hasattr(hf_model.backbone, "gen_mask"): + hf_model.backbone.gen_mask( + getattr(self.cfg, "model_length", getattr(self.cfg, "n_ctx", 2048)), + getattr(hf_model.backbone, "block_size", 4), + attn_backend=backend, + ) + + # Patch the model's forward method to filter out unsupported kwargs (like output_attentions) + original_forward = hf_model.forward + import inspect + + def patched_forward(*args, **kwargs): + sig = inspect.signature(original_forward) + valid_params = set(sig.parameters.keys()) + + # Check if VAR_KEYWORD (like **kwargs) is accepted + accepts_var_keyword = any( + p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values() + ) + + if not accepts_var_keyword: + kwargs = {k: v for k, v in kwargs.items() if k in valid_params} + return original_forward(*args, **kwargs) + + hf_model.forward = patched_forward + + def convert_weights(self) -> dict[str, torch.Tensor]: + """Return empty dict — delegation means no weight rearrangement.""" + return {} diff --git a/transformer_lens/tools/model_registry/__init__.py b/transformer_lens/tools/model_registry/__init__.py index 50d3a4914..8b0a42d87 100644 --- a/transformer_lens/tools/model_registry/__init__.py +++ b/transformer_lens/tools/model_registry/__init__.py @@ -50,6 +50,7 @@ "BaiChuanForCausalLM", "BaichuanForCausalLM", "BartForConditionalGeneration", + "BD3LM", "BertForMaskedLM", "BloomForCausalLM", "CodeGenForCausalLM", @@ -124,6 +125,7 @@ "BaiChuanForCausalLM": ["baichuan-inc"], "BaichuanForCausalLM": ["baichuan-inc"], "BartForConditionalGeneration": ["facebook"], + "BD3LM": ["kuleshov-group"], "BertForMaskedLM": ["google-bert"], "BloomForCausalLM": ["bigscience"], "CodeGenForCausalLM": ["Salesforce"], diff --git a/transformer_lens/tools/model_registry/generate_report.py b/transformer_lens/tools/model_registry/generate_report.py index 62b5e981d..eb3d1c923 100644 --- a/transformer_lens/tools/model_registry/generate_report.py +++ b/transformer_lens/tools/model_registry/generate_report.py @@ -63,6 +63,7 @@ "T5ForConditionalGeneration": "Google's T5 encoder-decoder model (partial support)", "T5GemmaForConditionalGeneration": "Google's T5Gemma encoder-decoder model with Gemma-style RoPE, GQA, and gated MLP", "BartForConditionalGeneration": "Facebook's BART encoder-decoder model", + "BD3LM": "Kuleshov Group's Block Diffusion Language Model (ICLR 2025) for masked text generation", "HunYuanDenseV1ForCausalLM": "Tencent's open source decoder models", # Unsupported architectures "BertModel": "Google's BERT bidirectional encoder for understanding tasks", diff --git a/transformer_lens/tools/model_registry/verify_models.py b/transformer_lens/tools/model_registry/verify_models.py index d84f72b36..407acc975 100644 --- a/transformer_lens/tools/model_registry/verify_models.py +++ b/transformer_lens/tools/model_registry/verify_models.py @@ -63,6 +63,7 @@ _BRIDGE_REMOTE_CODE_PREFIXES: tuple[str, ...] = ( "baichuan-inc/", # BaichuanForCausalLM — ships own modeling_baichuan.py "internlm/", # InternLM2ForCausalLM — ships own modeling_internlm2.py + "kuleshov-group/", # BD3LM — ships own custom modeling_d_dit.py ) # Data directory for registry files diff --git a/transformer_lens/utilities/architectures.py b/transformer_lens/utilities/architectures.py index ef216e596..53a982e96 100644 --- a/transformer_lens/utilities/architectures.py +++ b/transformer_lens/utilities/architectures.py @@ -26,6 +26,7 @@ "AlbertForMaskedLM", "DistilBertForMaskedLM", "ElectraForMaskedLM", + "BD3LM", } # Vision-language multimodal models From dfa54f7ab7098710dd1c850ecf786996ff4b503f Mon Sep 17 00:00:00 2001 From: pablocs116 <152398143+pablocs116@users.noreply.github.com> Date: Mon, 6 Jul 2026 15:18:44 +0200 Subject: [PATCH 06/11] Add RecurrentGemma (Griffin) architecture adapter (#1483) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add RecurrentGemma (Griffin) architecture adapter Adds a TransformerBridge adapter for RecurrentGemmaForCausalLM (Griffin): RG-LRU real-gated linear recurrence interleaved with local sliding-window attention (config.block_types, default recurrent/recurrent/attention). Because the temporal_block substructure varies per layer (recurrent vs. attention), the adapter wraps each decoder layer whole with residual-stream hooks only (RecurrentGemmaBlockBridge), mirroring Lfm2MoeArchitectureAdapter — rather than pretending every layer has a homogeneous attention/MLP substructure. applicable_phases=[4] (generation); finer-grained RG-LRU state hooks are a follow-up. Sets the Gemma-family numerics on the bridge config: RMSNorm (1.0 + weight) offset, gated MLP, final_rms, and the tanh final-logit soft cap (config.logits_soft_cap=30.0). ln_final maps to model.final_norm. Registers the adapter at the factory + supported_architectures __init__ sites and adds an offline structural unit test. Closes #1472 Co-Authored-By: Claude Opus 4.8 * Apply black formatting to RecurrentGemma adapter test Collapse two test method signatures onto single lines per black (line-length 100). No behavior change. Co-Authored-By: Claude Opus 4.8 * Register RecurrentGemma in the model-registry sync tables The RegistrySyncedWithFactory unit tests require every factory architecture to appear in HF_SUPPORTED_ARCHITECTURES and CANONICAL_AUTHORS_BY_ARCH (or INTENTIONAL_EXCLUDES). Add RecurrentGemmaForCausalLM to both (canonical author: google). Co-Authored-By: Claude Opus 4.8 * ci: re-trigger checks (flaky nbval notebook) Co-Authored-By: Claude Opus 4.8 --------- Co-authored-by: Claude Opus 4.8 --- .../test_recurrent_gemma_adapter.py | 130 ++++++++++++++++++ .../factories/architecture_adapter_factory.py | 2 + .../supported_architectures/__init__.py | 4 + .../recurrent_gemma.py | 126 +++++++++++++++++ .../tools/model_registry/__init__.py | 2 + 5 files changed, 264 insertions(+) create mode 100644 tests/unit/model_bridge/supported_architectures/test_recurrent_gemma_adapter.py create mode 100644 transformer_lens/model_bridge/supported_architectures/recurrent_gemma.py diff --git a/tests/unit/model_bridge/supported_architectures/test_recurrent_gemma_adapter.py b/tests/unit/model_bridge/supported_architectures/test_recurrent_gemma_adapter.py new file mode 100644 index 000000000..262cf8b50 --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_recurrent_gemma_adapter.py @@ -0,0 +1,130 @@ +"""Unit tests for the RecurrentGemmaArchitectureAdapter — no model downloads. + +RecurrentGemma (Griffin) is a hybrid of RG-LRU recurrent layers and local +sliding-window attention layers, wrapped whole-layer (residual hooks only) like +LFM2. These tests pin the architecture-specific quirks: + +- Gemma-family numerics: RMSNorm ``(1.0 + weight)`` offset, final-logit soft cap. +- Gated MLP, RMSNorm, ``final_rms``. +- The heterogeneous ``block_types`` pattern is exposed as ``layers_block_type``. +- Whole-layer block bridge advertising only residual-stream hooks (no attn/MLP + substructure, since recurrent layers have none). +""" + +import pytest + +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.model_bridge.generalized_components import ( + EmbeddingBridge, + RMSNormalizationBridge, + UnembeddingBridge, +) +from transformer_lens.model_bridge.supported_architectures.recurrent_gemma import ( + RecurrentGemmaArchitectureAdapter, + RecurrentGemmaBlockBridge, +) + + +@pytest.fixture(scope="class") +def cfg() -> TransformerBridgeConfig: + bridge_cfg = TransformerBridgeConfig( + d_model=64, + d_head=16, + n_layers=6, + n_ctx=128, + n_heads=4, + n_key_value_heads=1, + d_vocab=256, + d_mlp=128, + architecture="RecurrentGemmaForCausalLM", + ) + # HF RecurrentGemmaConfig fields the adapter reads off the raw config. + bridge_cfg.block_types = ["recurrent", "recurrent", "attention"] + bridge_cfg.rms_norm_eps = 1e-6 + bridge_cfg.logits_soft_cap = 30.0 + return bridge_cfg + + +@pytest.fixture(scope="class") +def adapter(cfg: TransformerBridgeConfig) -> RecurrentGemmaArchitectureAdapter: + return RecurrentGemmaArchitectureAdapter(cfg) + + +class TestRecurrentGemmaAdapterConfig: + def test_norm_config_is_gemma_rmsnorm(self, adapter: RecurrentGemmaArchitectureAdapter) -> None: + assert adapter.cfg.normalization_type == "RMS" + assert adapter.cfg.uses_rms_norm is True + assert adapter.cfg.final_rms is True + # Gemma RMSNorm applies (1.0 + weight); see HF PR #29402. + assert adapter.cfg.rmsnorm_uses_offset is True + assert adapter.cfg.eps == 1e-6 + + def test_gated_mlp_and_rotary(self, adapter: RecurrentGemmaArchitectureAdapter) -> None: + assert adapter.cfg.gated_mlp is True + assert adapter.cfg.attn_only is False + assert adapter.cfg.positional_embedding_type == "rotary" + + def test_final_logit_soft_cap(self, adapter: RecurrentGemmaArchitectureAdapter) -> None: + assert adapter.cfg.output_logits_soft_cap == 30.0 + + def test_default_prepend_bos_is_false(self, adapter: RecurrentGemmaArchitectureAdapter) -> None: + assert adapter.cfg.default_prepend_bos is False + + def test_block_types_exposed_for_analysis( + self, adapter: RecurrentGemmaArchitectureAdapter + ) -> None: + expected = ["recurrent", "recurrent", "attention"] + assert adapter.cfg.block_types == expected + # Normalized to the canonical name used by Nemotron-H / Granite tooling. + assert adapter.cfg.layers_block_type == expected + + def test_applicable_phases_is_generation_only( + self, adapter: RecurrentGemmaArchitectureAdapter + ) -> None: + assert adapter.applicable_phases == [4] + + +class TestRecurrentGemmaComponentMapping: + def test_has_residual_only_top_level_mapping( + self, adapter: RecurrentGemmaArchitectureAdapter + ) -> None: + mapping = adapter.component_mapping + assert mapping is not None + assert set(mapping) == {"embed", "blocks", "ln_final", "unembed"} + + def test_component_types(self, adapter: RecurrentGemmaArchitectureAdapter) -> None: + mapping = adapter.component_mapping + assert isinstance(mapping["embed"], EmbeddingBridge) + assert isinstance(mapping["blocks"], RecurrentGemmaBlockBridge) + assert isinstance(mapping["ln_final"], RMSNormalizationBridge) + assert isinstance(mapping["unembed"], UnembeddingBridge) + + def test_hf_module_paths(self, adapter: RecurrentGemmaArchitectureAdapter) -> None: + mapping = adapter.component_mapping + assert mapping["embed"].name == "model.embed_tokens" + assert mapping["blocks"].name == "model.layers" + # RecurrentGemma names the decoder-final norm `final_norm`, not `norm`. + assert mapping["ln_final"].name == "model.final_norm" + assert mapping["unembed"].name == "lm_head" + + def test_blocks_only_advertise_supported_residual_aliases( + self, adapter: RecurrentGemmaArchitectureAdapter + ) -> None: + blocks = adapter.component_mapping["blocks"] + assert blocks.hook_aliases == { + "hook_resid_pre": "hook_in", + "hook_resid_post": "hook_out", + } + assert blocks.submodules == {} + + +class TestRecurrentGemmaFactoryRegistration: + def test_registered_in_factory(self) -> None: + from transformer_lens.factories.architecture_adapter_factory import ( + SUPPORTED_ARCHITECTURES, + ) + + assert ( + SUPPORTED_ARCHITECTURES["RecurrentGemmaForCausalLM"] + is RecurrentGemmaArchitectureAdapter + ) diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py index a839c179b..5c73e1cf2 100644 --- a/transformer_lens/factories/architecture_adapter_factory.py +++ b/transformer_lens/factories/architecture_adapter_factory.py @@ -76,6 +76,7 @@ Qwen3MoeArchitectureAdapter, Qwen3NextArchitectureAdapter, QwenArchitectureAdapter, + RecurrentGemmaArchitectureAdapter, SmolLM3ArchitectureAdapter, StableLmArchitectureAdapter, T5ArchitectureAdapter, @@ -155,6 +156,7 @@ "Qwen3NextForCausalLM": Qwen3NextArchitectureAdapter, "Qwen3_5ForCausalLM": Qwen3_5ArchitectureAdapter, "Qwen3_5ForConditionalGeneration": Qwen3_5MultimodalArchitectureAdapter, + "RecurrentGemmaForCausalLM": RecurrentGemmaArchitectureAdapter, "SmolLM3ForCausalLM": SmolLM3ArchitectureAdapter, "StableLmForCausalLM": StableLmArchitectureAdapter, "T5ForConditionalGeneration": T5ArchitectureAdapter, diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py index bdb2d7311..ee7bc5b6a 100644 --- a/transformer_lens/model_bridge/supported_architectures/__init__.py +++ b/transformer_lens/model_bridge/supported_architectures/__init__.py @@ -201,6 +201,9 @@ from transformer_lens.model_bridge.supported_architectures.qwen3_next import ( Qwen3NextArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.recurrent_gemma import ( + RecurrentGemmaArchitectureAdapter, +) from transformer_lens.model_bridge.supported_architectures.smollm3 import ( SmolLM3ArchitectureAdapter, ) @@ -285,6 +288,7 @@ "Qwen3NextArchitectureAdapter", "Qwen3_5ArchitectureAdapter", "Qwen3_5MultimodalArchitectureAdapter", + "RecurrentGemmaArchitectureAdapter", "SmolLM3ArchitectureAdapter", "StableLmArchitectureAdapter", "T5ArchitectureAdapter", diff --git a/transformer_lens/model_bridge/supported_architectures/recurrent_gemma.py b/transformer_lens/model_bridge/supported_architectures/recurrent_gemma.py new file mode 100644 index 000000000..3f462ecd0 --- /dev/null +++ b/transformer_lens/model_bridge/supported_architectures/recurrent_gemma.py @@ -0,0 +1,126 @@ +"""RecurrentGemma (Griffin) architecture adapter. + +Supports ``RecurrentGemmaForCausalLM`` (e.g. ``google/recurrentgemma-2b``, +``google/recurrentgemma-9b``), the open instance of the **Griffin** architecture. + +Architecture overview: +- Heterogeneous layers defined by ``config.block_types`` (default pattern + ``("recurrent", "recurrent", "attention")`` repeated over ``num_hidden_layers``). + Each ``model.layers.{i}`` is a ``RecurrentGemmaDecoderLayer`` whose + ``temporal_block`` is *either*: + * ``RecurrentGemmaRecurrentBlock`` — the RG-LRU real-gated linear recurrence + (``linear_x`` / ``linear_y`` / ``conv_1d`` / ``rg_lru`` / ``linear_out``), or + * ``RecurrentGemmaSdpaAttention`` — local sliding-window GQA attention with + partial rotary (``q_proj`` / ``k_proj`` / ``v_proj`` / ``o_proj`` / ``rotary_emb``). +- Every layer additionally has ``temporal_pre_norm`` (pre-norm before the temporal + block), ``channel_pre_norm`` (pre-norm before the MLP), and a gated MLP + ``mlp_block`` (``gate_proj`` / ``up_proj`` / ``down_proj``, GELU-tanh). +- ``model.final_norm`` + ``lm_head``. + +Gemma-family numerics (shared with Gemma-2/3): RMSNorm applies ``(1.0 + weight)`` +(``rmsnorm_uses_offset``), token embeddings are scaled by ``sqrt(hidden_size)`` at +runtime inside the HF forward, and the final logits are tanh-soft-capped at +``config.logits_soft_cap`` (30.0). + +Key adapter decision (mirrors ``Lfm2MoeArchitectureAdapter``): because the +``temporal_block`` substructure varies per layer (recurrent vs. attention), we +wrap each decoder layer as a whole with residual-stream hooks only, rather than +pretending every layer has a homogeneous attention/MLP substructure. This keeps +execution correct on both layer types. Finer-grained RG-LRU state hooks +(``hook_ssm_write`` / ``hook_ssm_state`` style) are a natural follow-up. + +``applicable_phases = [4]``: the whole-layer bridge exposes only residual hooks, +so the component-level comparisons in phases 1-3 do not apply; phase 4 +(generation + text quality) does. +""" + +from typing import Any + +from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter +from transformer_lens.model_bridge.generalized_components import ( + BlockBridge, + EmbeddingBridge, + RMSNormalizationBridge, + UnembeddingBridge, +) + + +class RecurrentGemmaBlockBridge(BlockBridge): + """Whole-layer RecurrentGemma bridge exposing only residual-stream hooks. + + RecurrentGemma interleaves RG-LRU recurrent layers and local-attention layers. + Wrapping the HF decoder layer as a whole preserves correct execution while + avoiding unresolved standard attention/MLP aliases on the recurrent layers + (which have no q/k/v/o or gate/up/down substructure). + """ + + hook_aliases = { + "hook_resid_pre": "hook_in", + "hook_resid_post": "hook_out", + } + + +class RecurrentGemmaArchitectureAdapter(ArchitectureAdapter): + """Architecture adapter for ``RecurrentGemmaForCausalLM`` (Griffin). + + Hybrid RG-LRU recurrence + local sliding-window attention. The temporal-block + type per layer is determined by ``config.block_types[layer_idx % len(block_types)]``. + """ + + # Whole-layer residual hooks only; phases 1-3 compare component substructure + # this adapter intentionally does not expose. Phase 4 (generation) applies. + applicable_phases: list[int] = [4] + + def __init__(self, cfg: Any) -> None: + """Initialize the RecurrentGemma architecture adapter.""" + super().__init__(cfg) + + self.cfg.normalization_type = "RMS" + self.cfg.uses_rms_norm = True + self.cfg.final_rms = True + self.cfg.attn_only = False + # RG-LRU + local attention both use RoPE-style handling internally on the + # attention layers; there is no model-level rotary module (it lives inside + # each attention temporal_block), so we do not wire a rotary_emb component. + self.cfg.positional_embedding_type = "rotary" + # Gemma-family gated MLP (gate_proj -> GELU-tanh(up) -> down). + self.cfg.gated_mlp = True + # Gemma RMSNorm uses (1.0 + weight); see + # https://github.com/huggingface/transformers/pull/29402 + self.cfg.rmsnorm_uses_offset = True + + # Gemma models were not trained with BOS tokens prepended. + self.cfg.default_prepend_bos = False + + norm_eps = getattr(cfg, "rms_norm_eps", None) + if norm_eps is not None: + self.cfg.eps = norm_eps + + # Final logits are tanh-soft-capped at config.logits_soft_cap (30.0). + logits_soft_cap = getattr(cfg, "logits_soft_cap", None) + if logits_soft_cap is not None: + self.cfg.output_logits_soft_cap = logits_soft_cap + + # Expose the per-layer temporal-block pattern for analysis tools, using the + # canonical `layers_block_type` name as on Nemotron-H / Granite. + block_types = list(getattr(cfg, "block_types", None) or []) + setattr(self.cfg, "block_types", block_types) + setattr(self.cfg, "layers_block_type", block_types) + + self.component_mapping = { + "embed": EmbeddingBridge(name="model.embed_tokens"), + "blocks": RecurrentGemmaBlockBridge(name="model.layers", config=self.cfg), + "ln_final": RMSNormalizationBridge(name="model.final_norm", config=self.cfg), + "unembed": UnembeddingBridge(name="lm_head", config=self.cfg), + } + + def prepare_loading(self, model_name: str, model_kwargs: dict) -> None: + """Force eager attention when the HF config exposes the implementation knob.""" + config = model_kwargs.get("config") + if config is not None and hasattr(config, "_attn_implementation"): + config._attn_implementation = "eager" + + def prepare_model(self, hf_model: Any) -> None: + """Force eager attention on the loaded HF model when supported.""" + if hasattr(hf_model, "config") and hasattr(hf_model.config, "_attn_implementation"): + hf_model.config._attn_implementation = "eager" diff --git a/transformer_lens/tools/model_registry/__init__.py b/transformer_lens/tools/model_registry/__init__.py index 8b0a42d87..b01c8fe85 100644 --- a/transformer_lens/tools/model_registry/__init__.py +++ b/transformer_lens/tools/model_registry/__init__.py @@ -109,6 +109,7 @@ "Qwen3NextForCausalLM", "Qwen3_5ForCausalLM", "Qwen3_5ForConditionalGeneration", + "RecurrentGemmaForCausalLM", "SmolLM3ForCausalLM", "StableLmForCausalLM", "T5ForConditionalGeneration", @@ -185,6 +186,7 @@ "Qwen3_5ForCausalLM": ["Qwen"], "Qwen3_5ForConditionalGeneration": ["Qwen"], "QwenForCausalLM": ["Qwen"], + "RecurrentGemmaForCausalLM": ["google"], "SmolLM3ForCausalLM": ["HuggingFaceTB"], "StableLmForCausalLM": ["stabilityai"], "T5ForConditionalGeneration": ["google-t5", "google", "Salesforce", "MBZUAI"], From c689deb456ec38005e1ea40ce425a66da8013b15 Mon Sep 17 00:00:00 2001 From: Henry Su Date: Mon, 6 Jul 2026 10:45:04 -0500 Subject: [PATCH 07/11] Add Cohere2 architecture adapter (#1485) --- .../model_bridge/test_cohere2_adapter.py | 65 +++ .../test_cohere2_adapter.py | 373 ++++++++++++++++++ .../factories/architecture_adapter_factory.py | 2 + .../model_bridge/sources/_bridge_builder.py | 2 + .../model_bridge/sources/transformers.py | 3 + .../supported_architectures/__init__.py | 8 +- .../supported_architectures/cohere.py | 99 +++++ .../tools/model_registry/__init__.py | 2 + .../tools/model_registry/generate_report.py | 1 + 9 files changed, 552 insertions(+), 3 deletions(-) create mode 100644 tests/integration/model_bridge/test_cohere2_adapter.py create mode 100644 tests/unit/model_bridge/supported_architectures/test_cohere2_adapter.py diff --git a/tests/integration/model_bridge/test_cohere2_adapter.py b/tests/integration/model_bridge/test_cohere2_adapter.py new file mode 100644 index 000000000..2f140e1b2 --- /dev/null +++ b/tests/integration/model_bridge/test_cohere2_adapter.py @@ -0,0 +1,65 @@ +"""Integration tests for Cohere2 architecture adapter (Cohere2ForCausalLM). + +Model: trl-internal-testing/tiny-Cohere2ForCausalLM + - 2 layers, CPU-safe, no gated access. + - Current config has only sliding_attention layers, so full-attention NoPE is + covered by the synthetic unit tests in test_cohere2_adapter.py. +""" + +import pytest +import torch +from transformers import AutoModelForCausalLM + +from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.supported_architectures.cohere import ( + _Cohere2AttentionBridge, +) + +pytestmark = pytest.mark.slow + +MODEL = "trl-internal-testing/tiny-Cohere2ForCausalLM" + + +@pytest.fixture(scope="module") +def cohere2_bridge() -> TransformerBridge: + return TransformerBridge.boot_transformers(MODEL, device="cpu", dtype=torch.float32) + + +@pytest.fixture(scope="module") +def cohere2_hf() -> torch.nn.Module: + return AutoModelForCausalLM.from_pretrained( + MODEL, torch_dtype=torch.float32, attn_implementation="eager" + ).eval() + + +class TestCohere2BridgeCreation: + def test_boot_transformers_succeeds(self, cohere2_bridge: TransformerBridge) -> None: + assert cohere2_bridge is not None + + def test_block_count_matches_hf( + self, cohere2_bridge: TransformerBridge, cohere2_hf: torch.nn.Module + ) -> None: + assert len(cohere2_bridge.blocks) == cohere2_hf.config.num_hidden_layers + + def test_cfg_layer_types_match_hf( + self, cohere2_bridge: TransformerBridge, cohere2_hf: torch.nn.Module + ) -> None: + assert cohere2_bridge.cfg.layer_types == cohere2_hf.config.layer_types + + def test_attention_bridge_is_cohere2_nope_aware( + self, cohere2_bridge: TransformerBridge + ) -> None: + for block in cohere2_bridge.blocks: + assert type(block.attn) is _Cohere2AttentionBridge + + +class TestCohere2MatchesHuggingFace: + def test_forward_logits_match_hf( + self, cohere2_bridge: TransformerBridge, cohere2_hf: torch.nn.Module + ) -> None: + tokens = torch.tensor([[1, 2, 3, 4]]) + with torch.no_grad(): + bridge_logits = cohere2_bridge(tokens) + hf_logits = cohere2_hf(tokens).logits + max_diff = (bridge_logits - hf_logits).abs().max().item() + assert max_diff < 1e-5, f"Cohere2 bridge vs HF max diff = {max_diff:.6f}" diff --git a/tests/unit/model_bridge/supported_architectures/test_cohere2_adapter.py b/tests/unit/model_bridge/supported_architectures/test_cohere2_adapter.py new file mode 100644 index 000000000..acaa66602 --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_cohere2_adapter.py @@ -0,0 +1,373 @@ +"""Unit tests for Cohere2ArchitectureAdapter. + +Cohere2 reuses Cohere's parallel LayerNorm/GQA/logit-scale structure but mixes +sliding-window RoPE layers with full-attention NoPE layers. These tests stay +download-free and pin the adapter-specific layer-type and NoPE logic. +""" + +from types import SimpleNamespace +from typing import Any + +import pytest +import torch.nn as nn +from torch import allclose, float32, manual_seed, no_grad, ones, randn, zeros + +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.factories.architecture_adapter_factory import ( + SUPPORTED_ARCHITECTURES, + ArchitectureAdapterFactory, +) +from transformer_lens.model_bridge.generalized_components import LinearBridge +from transformer_lens.model_bridge.generalized_components.position_embeddings_attention import ( + PositionEmbeddingsAttentionBridge, +) +from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_config_from_hf, +) +from transformer_lens.model_bridge.sources.transformers import ( + determine_architecture_from_hf_config, +) +from transformer_lens.model_bridge.supported_architectures.cohere import ( + Cohere2ArchitectureAdapter, + CohereArchitectureAdapter, + _Cohere2AttentionBridge, +) +from transformer_lens.tools.model_registry import ( + CANONICAL_AUTHORS_BY_ARCH, + HF_SUPPORTED_ARCHITECTURES, +) + + +def _make_cfg( + n_heads: int = 4, + d_model: int = 64, + n_layers: int = 8, + d_mlp: int = 256, + d_vocab: int = 1000, + n_ctx: int = 512, + n_key_value_heads: int | None = 2, + layer_types: list[str] | None = None, + sliding_window_pattern: int | None = None, +) -> TransformerBridgeConfig: + cfg = TransformerBridgeConfig( + d_model=d_model, + d_head=d_model // n_heads, + n_layers=n_layers, + n_ctx=n_ctx, + n_heads=n_heads, + d_vocab=d_vocab, + d_mlp=d_mlp, + default_prepend_bos=True, + architecture="Cohere2ForCausalLM", + ) + if n_key_value_heads is not None: + cfg.n_key_value_heads = n_key_value_heads + setattr(cfg, "logit_scale", 1.0) + setattr(cfg, "rope_parameters", {"rope_theta": 50000.0, "rope_type": "default"}) + if layer_types is not None: + setattr(cfg, "layer_types", layer_types) + if sliding_window_pattern is not None: + setattr(cfg, "sliding_window_pattern", sliding_window_pattern) + return cfg + + +class FakeCohere2Attention(nn.Module): + """Minimal Cohere2-style attention module for bridge tests.""" + + def __init__( + self, + cfg: TransformerBridgeConfig, + sliding_window: int | None = 4096, + expose_sliding_window: bool = True, + layer_idx: int | None = None, + ) -> None: + super().__init__() + self.head_dim = cfg.d_head + self.num_key_value_groups = cfg.n_heads // (cfg.n_key_value_heads or cfg.n_heads) + self.scaling = cfg.d_head**-0.5 + self.attention_dropout = 0.0 + self.layer_idx = layer_idx + if expose_sliding_window: + self.sliding_window = sliding_window + + kv_width = (cfg.n_key_value_heads or cfg.n_heads) * cfg.d_head + self.q_proj = nn.Linear(cfg.d_model, cfg.n_heads * cfg.d_head, bias=False) + self.k_proj = nn.Linear(cfg.d_model, kv_width, bias=False) + self.v_proj = nn.Linear(cfg.d_model, kv_width, bias=False) + self.o_proj = nn.Linear(cfg.n_heads * cfg.d_head, cfg.d_model, bias=False) + + +def _cohere2_attention_bridge(cfg: TransformerBridgeConfig) -> _Cohere2AttentionBridge: + return _Cohere2AttentionBridge( + name="self_attn", + config=cfg, + submodules={ + "q": LinearBridge(name="q_proj"), + "k": LinearBridge(name="k_proj"), + "v": LinearBridge(name="v_proj"), + "o": LinearBridge(name="o_proj"), + }, + ) + + +class TestCohere2Config: + def test_explicit_layer_types_are_preserved(self) -> None: + layer_types = [ + "sliding_attention", + "full_attention", + "sliding_attention", + "full_attention", + ] + adapter = Cohere2ArchitectureAdapter(_make_cfg(n_layers=4, layer_types=layer_types)) + + assert adapter.cfg.layer_types == layer_types + + def test_sliding_window_pattern_generates_layer_types(self) -> None: + adapter = Cohere2ArchitectureAdapter(_make_cfg(n_layers=8, sliding_window_pattern=4)) + + assert adapter.cfg.layer_types == [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + ] + + def test_underscored_sliding_window_pattern_generates_layer_types(self) -> None: + cfg = _make_cfg(n_layers=3) + setattr(cfg, "_sliding_window_pattern", 2) + + adapter = Cohere2ArchitectureAdapter(cfg) + + assert adapter.cfg.layer_types == [ + "sliding_attention", + "full_attention", + "sliding_attention", + ] + + def test_layer_types_length_must_match_n_layers(self) -> None: + cfg = _make_cfg(n_layers=2, layer_types=["sliding_attention"]) + + with pytest.raises(ValueError, match="layer_types length"): + Cohere2ArchitectureAdapter(cfg) + + def test_sliding_window_pattern_must_be_positive(self) -> None: + cfg = _make_cfg(n_layers=2, sliding_window_pattern=0) + + with pytest.raises(ValueError, match="sliding_window_pattern must be positive"): + Cohere2ArchitectureAdapter(cfg) + + def test_inherits_cohere_gqa_and_logit_scale_behaviour(self) -> None: + adapter = Cohere2ArchitectureAdapter(_make_cfg(n_heads=8, n_key_value_heads=2)) + + assert isinstance(adapter, CohereArchitectureAdapter) + assert adapter.cfg.n_key_value_heads == 2 + assert getattr(adapter.cfg, "logit_scale") == 1.0 + assert set(adapter.weight_processing_conversions.keys()) == { + "blocks.{i}.attn.q.weight", + "blocks.{i}.attn.k.weight", + "blocks.{i}.attn.v.weight", + "blocks.{i}.attn.o.weight", + } + + +class TestCohere2ComponentMapping: + def test_attention_bridge_is_cohere2_nope_aware(self) -> None: + adapter = Cohere2ArchitectureAdapter(_make_cfg()) + + attn = adapter.component_mapping["blocks"].submodules["attn"] + assert type(attn) is _Cohere2AttentionBridge + assert issubclass(_Cohere2AttentionBridge, PositionEmbeddingsAttentionBridge) + + def test_parallel_cohere_block_shape_is_preserved(self) -> None: + adapter = Cohere2ArchitectureAdapter(_make_cfg()) + blocks = adapter.component_mapping["blocks"] + + assert "ln1" in blocks.submodules + assert "ln2" not in blocks.submodules + assert "mlp" in blocks.submodules + + +class TestCohere2NoPE: + def _record_super_forward(self, monkeypatch: pytest.MonkeyPatch) -> dict: + recorded: dict = {} + + def _fake_super_forward(self: Any, *args: Any, **kwargs: Any) -> str: + recorded["args"] = args + recorded["kwargs"] = kwargs + return "sentinel-output" + + monkeypatch.setattr( + PositionEmbeddingsAttentionBridge, "forward", _fake_super_forward, raising=True + ) + return recorded + + def test_full_attention_layer_suppresses_position_embeddings_kwarg( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + recorded = self._record_super_forward(monkeypatch) + cfg = _make_cfg() + bridge = _cohere2_attention_bridge(cfg) + bridge.set_original_component(FakeCohere2Attention(cfg, sliding_window=None)) + hidden = randn(2, 8, 64) + + result = bridge.forward(hidden, position_embeddings=(ones(1, 8, 16), zeros(1, 8, 16))) + + assert result == "sentinel-output" + assert recorded["kwargs"]["position_embeddings"] is None + + def test_sliding_attention_layer_passes_position_embeddings_through( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + recorded = self._record_super_forward(monkeypatch) + cfg = _make_cfg() + bridge = _cohere2_attention_bridge(cfg) + bridge.set_original_component(FakeCohere2Attention(cfg, sliding_window=4096)) + hidden = randn(2, 8, 64) + pos = (ones(1, 8, 16), zeros(1, 8, 16)) + + bridge.forward(hidden, position_embeddings=pos) + + assert recorded["kwargs"]["position_embeddings"] is pos + + def test_full_attention_layer_suppresses_positional_position_embeddings( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + recorded = self._record_super_forward(monkeypatch) + cfg = _make_cfg() + bridge = _cohere2_attention_bridge(cfg) + bridge.set_original_component(FakeCohere2Attention(cfg, sliding_window=None)) + hidden = randn(2, 8, 64) + pos = (ones(1, 8, 16), zeros(1, 8, 16)) + + bridge.forward(hidden, pos, None) + + assert recorded["args"][0] is hidden + assert recorded["args"][1] is None + + def test_layer_types_fallback_identifies_full_attention_layer( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + recorded = self._record_super_forward(monkeypatch) + cfg = _make_cfg( + n_layers=2, + layer_types=["sliding_attention", "full_attention"], + ) + adapter = Cohere2ArchitectureAdapter(cfg) + bridge = _cohere2_attention_bridge(adapter.cfg) + bridge.set_original_component( + FakeCohere2Attention( + adapter.cfg, + expose_sliding_window=False, + layer_idx=1, + ) + ) + hidden = randn(2, 8, 64) + + bridge.forward(hidden, position_embeddings=(ones(1, 8, 16), zeros(1, 8, 16))) + + assert recorded["kwargs"]["position_embeddings"] is None + + @staticmethod + def _wired_bridge( + cfg: TransformerBridgeConfig, + sliding_window: int | None, + ) -> _Cohere2AttentionBridge: + fake_attn = FakeCohere2Attention(cfg, sliding_window=sliding_window) + bridge = _cohere2_attention_bridge(cfg) + bridge.set_original_component(fake_attn) + for name, original in { + "q": fake_attn.q_proj, + "k": fake_attn.k_proj, + "v": fake_attn.v_proj, + "o": fake_attn.o_proj, + }.items(): + submodule = bridge.submodules[name] + submodule.set_original_component(original) + bridge.add_module(name, submodule) + bridge.setup_hook_compatibility() + return bridge + + @staticmethod + def _forward(bridge: _Cohere2AttentionBridge, hidden: Any, position_embeddings: Any) -> Any: + with no_grad(): + out = bridge(hidden, position_embeddings=position_embeddings, attention_mask=None) + return out[0] if isinstance(out, tuple) else out + + def test_full_attention_output_ignores_position_embeddings_end_to_end(self) -> None: + manual_seed(0) + cfg = _make_cfg(n_heads=4, n_key_value_heads=2, d_model=64) + bridge = self._wired_bridge(cfg, sliding_window=None) + hidden = randn(2, 8, 64) + cos = randn(1, 8, 16) + sin = randn(1, 8, 16) + + out_with_pos = self._forward(bridge, hidden, (cos, sin)) + out_without_pos = self._forward(bridge, hidden, None) + + assert allclose(out_with_pos, out_without_pos, atol=1e-6) + + def test_sliding_attention_output_depends_on_position_embeddings_end_to_end(self) -> None: + manual_seed(0) + cfg = _make_cfg(n_heads=4, n_key_value_heads=2, d_model=64) + bridge = self._wired_bridge(cfg, sliding_window=4096) + hidden = randn(2, 8, 64) + cos = randn(1, 8, 16) + sin = randn(1, 8, 16) + + out_with_pos = self._forward(bridge, hidden, (cos, sin)) + out_without_pos = self._forward(bridge, hidden, None) + + assert not allclose(out_with_pos, out_without_pos, atol=1e-6) + + +class TestCohere2Registration: + def test_factory_selects_cohere2_adapter(self) -> None: + cfg = _make_cfg() + + adapter = ArchitectureAdapterFactory.select_architecture_adapter(cfg) + + assert isinstance(adapter, Cohere2ArchitectureAdapter) + + def test_supported_architectures_contains_cohere2(self) -> None: + assert SUPPORTED_ARCHITECTURES["Cohere2ForCausalLM"] is Cohere2ArchitectureAdapter + + def test_registry_contains_cohere2(self) -> None: + assert "Cohere2ForCausalLM" in HF_SUPPORTED_ARCHITECTURES + assert CANONICAL_AUTHORS_BY_ARCH["Cohere2ForCausalLM"] == ["CohereLabs"] + + def test_model_type_routes_to_cohere2(self) -> None: + hf_config = SimpleNamespace(model_type="cohere2", architectures=None) + + assert determine_architecture_from_hf_config(hf_config) == "Cohere2ForCausalLM" + + def test_sliding_window_pattern_passthrough_survives_hf_config_translation(self) -> None: + hf_config = SimpleNamespace( + model_type="cohere2", + architectures=["Cohere2ForCausalLM"], + hidden_size=64, + head_dim=16, + num_attention_heads=4, + num_key_value_heads=2, + num_hidden_layers=4, + vocab_size=100, + max_position_embeddings=128, + intermediate_size=256, + hidden_act="silu", + layer_norm_eps=1e-5, + logit_scale=1.0, + rope_parameters={"rope_theta": 50000.0, "rope_type": "default"}, + sliding_window_pattern=4, + ) + + cfg = build_bridge_config_from_hf( + hf_config, + architecture="Cohere2ForCausalLM", + model_name="synthetic-cohere2", + dtype=float32, + ) + + assert cfg.sliding_window_pattern == 4 diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py index 5c73e1cf2..727e6eaaf 100644 --- a/transformer_lens/factories/architecture_adapter_factory.py +++ b/transformer_lens/factories/architecture_adapter_factory.py @@ -18,6 +18,7 @@ BertArchitectureAdapter, BloomArchitectureAdapter, CodeGenArchitectureAdapter, + Cohere2ArchitectureAdapter, CohereArchitectureAdapter, DeepSeekV2ArchitectureAdapter, DeepSeekV3ArchitectureAdapter, @@ -95,6 +96,7 @@ "BertForMaskedLM": BertArchitectureAdapter, "BloomForCausalLM": BloomArchitectureAdapter, "CodeGenForCausalLM": CodeGenArchitectureAdapter, + "Cohere2ForCausalLM": Cohere2ArchitectureAdapter, "CohereForCausalLM": CohereArchitectureAdapter, "DeepseekV2ForCausalLM": DeepSeekV2ArchitectureAdapter, "DeepseekV3ForCausalLM": DeepSeekV3ArchitectureAdapter, diff --git a/transformer_lens/model_bridge/sources/_bridge_builder.py b/transformer_lens/model_bridge/sources/_bridge_builder.py index a74d2ef9d..e5cf55027 100644 --- a/transformer_lens/model_bridge/sources/_bridge_builder.py +++ b/transformer_lens/model_bridge/sources/_bridge_builder.py @@ -57,6 +57,8 @@ # Cohere "logit_scale", "rope_parameters", + "sliding_window_pattern", + "_sliding_window_pattern", # Hybrid/MoE architectures "layer_types", "moe_intermediate_size", diff --git a/transformer_lens/model_bridge/sources/transformers.py b/transformer_lens/model_bridge/sources/transformers.py index 7929ee7bd..b79661ff5 100644 --- a/transformer_lens/model_bridge/sources/transformers.py +++ b/transformer_lens/model_bridge/sources/transformers.py @@ -246,6 +246,7 @@ def determine_architecture_from_hf_config(hf_config): "bert": "BertForMaskedLM", "bloom": "BloomForCausalLM", "codegen": "CodeGenForCausalLM", + "cohere2": "Cohere2ForCausalLM", "gptj": "GPTJForCausalLM", "gpt_neo": "GPTNeoForCausalLM", "gpt_neox": "GPTNeoXForCausalLM", @@ -556,6 +557,8 @@ def boot( # Cohere "logit_scale", "rope_parameters", + "sliding_window_pattern", + "_sliding_window_pattern", # Hybrid/MoE architectures "layer_types", "moe_intermediate_size", diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py index ee7bc5b6a..d783894c0 100644 --- a/transformer_lens/model_bridge/supported_architectures/__init__.py +++ b/transformer_lens/model_bridge/supported_architectures/__init__.py @@ -28,6 +28,7 @@ CodeGenArchitectureAdapter, ) from transformer_lens.model_bridge.supported_architectures.cohere import ( + Cohere2ArchitectureAdapter, CohereArchitectureAdapter, ) from transformer_lens.model_bridge.supported_architectures.deepseek_v2 import ( @@ -63,6 +64,9 @@ from transformer_lens.model_bridge.supported_architectures.glm4_moe import ( Glm4MoeArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.glm_moe_dsa import ( + GlmMoeDsaArchitectureAdapter, +) from transformer_lens.model_bridge.supported_architectures.gpt2 import ( GPT2ArchitectureAdapter, ) @@ -78,9 +82,6 @@ from transformer_lens.model_bridge.supported_architectures.gptj import ( GptjArchitectureAdapter, ) -from transformer_lens.model_bridge.supported_architectures.glm_moe_dsa import ( - GlmMoeDsaArchitectureAdapter, -) from transformer_lens.model_bridge.supported_architectures.granite import ( GraniteArchitectureAdapter, ) @@ -229,6 +230,7 @@ "BertArchitectureAdapter", "BloomArchitectureAdapter", "CodeGenArchitectureAdapter", + "Cohere2ArchitectureAdapter", "CohereArchitectureAdapter", "DeepSeekV2ArchitectureAdapter", "DeepSeekV3ArchitectureAdapter", diff --git a/transformer_lens/model_bridge/supported_architectures/cohere.py b/transformer_lens/model_bridge/supported_architectures/cohere.py index 97e8c3301..57e7731a1 100644 --- a/transformer_lens/model_bridge/supported_architectures/cohere.py +++ b/transformer_lens/model_bridge/supported_architectures/cohere.py @@ -193,3 +193,102 @@ def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> No # Also set on the template so get_generalized_component() calls work attn_bridge = self.get_generalized_component("blocks.0.attn") attn_bridge.set_rotary_emb(rotary_emb) + + +class _Cohere2AttentionBridge(PositionEmbeddingsAttentionBridge): + """Attention bridge that honours Cohere2's RoPE/NoPE interleaving. + + Cohere2 applies RoPE only on sliding-window layers. Full-attention global + layers receive the same position_embeddings tuple from the model loop but + intentionally skip apply_rotary_pos_emb inside HF's Cohere2Attention by + checking whether ``self.sliding_window`` is set. + + The base bridge rotates whenever position_embeddings is present, so full + layers must suppress that argument before delegating to the shared attention + reconstruction path. + """ + + def forward(self, *args: Any, **kwargs: Any) -> Any: + """Drop position_embeddings on Cohere2 full-attention NoPE layers.""" + if self._is_nope_layer(): + kwargs["position_embeddings"] = None + if len(args) >= 2 and not isinstance(args[1], torch.Tensor): + args = (args[0], None) + args[2:] + return super().forward(*args, **kwargs) + + def _is_nope_layer(self) -> bool: + """Return True when the wrapped Cohere2 attention is a full-attention layer.""" + hf_attn = self.original_component + if hf_attn is None: + return False + + if hasattr(hf_attn, "sliding_window"): + return getattr(hf_attn, "sliding_window") is None + + layer_idx = getattr(hf_attn, "layer_idx", None) + layer_types = getattr(self.config, "layer_types", None) + if layer_idx is None or layer_types is None: + return False + return layer_types[layer_idx] == "full_attention" + + +def _cohere2_layer_types(cfg: Any) -> list[str]: + """Resolve Cohere2 layer types from explicit config or sliding-window pattern.""" + n_layers = getattr(cfg, "n_layers", None) + if n_layers is None: + n_layers = getattr(cfg, "num_hidden_layers") + n_layers = int(n_layers) + layer_types = getattr(cfg, "layer_types", None) + if layer_types is not None: + resolved = list(layer_types) + if len(resolved) != n_layers: + raise ValueError( + f"Cohere2 layer_types length ({len(resolved)}) must match n_layers ({n_layers})." + ) + return resolved + + pattern = getattr(cfg, "sliding_window_pattern", None) + if pattern is None: + pattern = getattr(cfg, "_sliding_window_pattern", None) + if pattern is None: + pattern = 4 + pattern = int(pattern) + if pattern <= 0: + raise ValueError(f"Cohere2 sliding_window_pattern must be positive, got {pattern}.") + + return [ + "sliding_attention" if (layer_idx + 1) % pattern else "full_attention" + for layer_idx in range(n_layers) + ] + + +class Cohere2ArchitectureAdapter(CohereArchitectureAdapter): + """Architecture adapter for Cohere2 / Command-A models. + + Cohere2 keeps Cohere v1's parallel block, LayerNorm, GQA, gated MLP and + logit_scale behaviour, but interleaves sliding-window RoPE layers with + full-attention NoPE layers. HF represents that either as an explicit + ``layer_types`` list or as a legacy ``sliding_window_pattern`` integer. + """ + + def __init__(self, cfg: Any) -> None: + """Initialize the Cohere2 architecture adapter.""" + super().__init__(cfg) + + setattr(self.cfg, "layer_types", _cohere2_layer_types(cfg)) + + assert self.component_mapping is not None + blocks = self.component_mapping["blocks"] + assert blocks.submodules is not None + blocks.submodules["attn"] = _Cohere2AttentionBridge( + name="self_attn", + config=self.cfg, + submodules={ + "q": LinearBridge(name="q_proj"), + "k": LinearBridge(name="k_proj"), + "v": LinearBridge(name="v_proj"), + "o": LinearBridge(name="o_proj"), + }, + requires_attention_mask=True, + requires_position_embeddings=True, + ) diff --git a/transformer_lens/tools/model_registry/__init__.py b/transformer_lens/tools/model_registry/__init__.py index b01c8fe85..ad1effe43 100644 --- a/transformer_lens/tools/model_registry/__init__.py +++ b/transformer_lens/tools/model_registry/__init__.py @@ -54,6 +54,7 @@ "BertForMaskedLM", "BloomForCausalLM", "CodeGenForCausalLM", + "Cohere2ForCausalLM", "CohereForCausalLM", "DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM", @@ -130,6 +131,7 @@ "BertForMaskedLM": ["google-bert"], "BloomForCausalLM": ["bigscience"], "CodeGenForCausalLM": ["Salesforce"], + "Cohere2ForCausalLM": ["CohereLabs"], "CohereForCausalLM": ["CohereLabs"], "DeepseekV2ForCausalLM": ["deepseek-ai"], "DeepseekV3ForCausalLM": ["deepseek-ai"], diff --git a/transformer_lens/tools/model_registry/generate_report.py b/transformer_lens/tools/model_registry/generate_report.py index eb3d1c923..d62edda90 100644 --- a/transformer_lens/tools/model_registry/generate_report.py +++ b/transformer_lens/tools/model_registry/generate_report.py @@ -65,6 +65,7 @@ "BartForConditionalGeneration": "Facebook's BART encoder-decoder model", "BD3LM": "Kuleshov Group's Block Diffusion Language Model (ICLR 2025) for masked text generation", "HunYuanDenseV1ForCausalLM": "Tencent's open source decoder models", + "Cohere2ForCausalLM": "Cohere's Command-A architecture with interleaved sliding-window RoPE and full-attention NoPE layers", # Unsupported architectures "BertModel": "Google's BERT bidirectional encoder for understanding tasks", "BertForMaskedLM": "BERT with masked language modeling head", From ab46ac09432e0b7fad12065f4ec357d4279041cf Mon Sep 17 00:00:00 2001 From: Mukund Pandey Date: Mon, 6 Jul 2026 18:12:42 +0100 Subject: [PATCH 08/11] feat: add Zamba2ForCausalLM architecture adapter (#1486) * feat: add Zamba2ForCausalLM architecture adapter Closes #1463 Supports Zyphra/Zamba2-1.2B, Zamba2-7B, Zamba2-7B-Instruct via TransformerBridge. Zamba2 is a hybrid Mamba-2 / shared-attention model where layers_block_type alternates between "mamba" (pure SSM) and "hybrid" (SSM + shared global-attention block, weight-tied across num_mem_blocks slots). Key design decisions: - SSMBlockBridge for all layers: hook_in/hook_out fire regardless of layer type. Mamba layers expose mixer submodule hooks (in_proj, conv1d, inner_norm, out_proj); hybrid layers have norm and mixer marked optional=True so component_setup skips them gracefully. - is_stateful=False: Zamba2 threads its unified cache via past_key_values (standard KV path), not the Mamba cache_params path. Setting is_stateful=True collided with Zamba2's own cache call. - positional_embedding_type="none": RoPE is handled inside the HF attention block; no model-level rotary bridge needed. - Added mamba_expand, mamba_ngroups, num_mem_blocks, layers_block_type, use_shared_attention_adapter to _HF_PASSTHROUGH_ATTRS in both sources/transformers.py and sources/_bridge_builder.py so these config fields reach the adapter reliably. 21 integration tests pass against Zyphra/Zamba2-1.2B (38 layers: 32 mamba + 6 hybrid). Forward parity is exact (max_diff == 0.0); greedy generation matches HF bit-for-bit across both layer types. * fix: add Zamba2ForCausalLM to model registry (HF_SUPPORTED_ARCHITECTURES + CANONICAL_AUTHORS_BY_ARCH) --- .../model_bridge/test_zamba2_adapter.py | 302 ++++++++++++++++++ .../factories/architecture_adapter_factory.py | 2 + .../model_bridge/sources/_bridge_builder.py | 6 + .../model_bridge/sources/transformers.py | 6 + .../supported_architectures/__init__.py | 10 +- .../supported_architectures/zamba2.py | 166 ++++++++++ .../tools/model_registry/__init__.py | 2 + 7 files changed, 491 insertions(+), 3 deletions(-) create mode 100644 tests/integration/model_bridge/test_zamba2_adapter.py create mode 100644 transformer_lens/model_bridge/supported_architectures/zamba2.py diff --git a/tests/integration/model_bridge/test_zamba2_adapter.py b/tests/integration/model_bridge/test_zamba2_adapter.py new file mode 100644 index 000000000..b2ed82451 --- /dev/null +++ b/tests/integration/model_bridge/test_zamba2_adapter.py @@ -0,0 +1,302 @@ +"""Integration tests for the Zamba2 architecture adapter. + +Verifies forward-pass and generation parity against Zyphra/Zamba2-1.2B: +- Forward-pass logits match HF exactly (bridge delegates the full forward to HF) +- Greedy multi-token generation matches HF bit-for-bit (exercises the unified + Zamba2HybridDynamicCache threaded via past_key_values across Mamba-2 and + shared global-attention layers) +- Sanity checks: config flags, block count, hook coverage on both layer types + +Zamba2 has two layer types in ``config.layers_block_type``: + - ``"mamba"`` — pure Mamba-2 SSM (Zamba2MambaDecoderLayer) + - ``"hybrid"`` — Mamba-2 + shared global-attention (Zamba2HybridLayer) + +Run with GPU acceleration: + CUDA_VISIBLE_DEVICES=0 pytest tests/integration/model_bridge/test_zamba2_adapter.py -v -s + +On a CPU-only machine: + pytest tests/integration/model_bridge/test_zamba2_adapter.py -v -s +""" + +import gc + +import pytest +import torch + +from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.generalized_components import ( + RMSNormalizationBridge, + SSM2MixerBridge, + SSMBlockBridge, +) + +pytestmark = pytest.mark.slow + +MODEL = "Zyphra/Zamba2-1.2B" + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _device() -> str: + return "cuda" if torch.cuda.is_available() else "cpu" + + +def _dtype() -> torch.dtype: + # bfloat16 on GPU to match HF defaults; float32 on CPU for numerical safety + return torch.bfloat16 if torch.cuda.is_available() else torch.float32 + + +# --------------------------------------------------------------------------- +# Session fixture — load once, share across all test classes +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def zamba2_bridge(): + device = _device() + dtype = _dtype() + bridge = TransformerBridge.boot_transformers(MODEL, device=device, dtype=dtype) + yield bridge + del bridge + if torch.cuda.is_available(): + torch.cuda.empty_cache() + for _ in range(3): + gc.collect() + + +# --------------------------------------------------------------------------- +# Config and bridge structure +# --------------------------------------------------------------------------- + + +class TestZamba2BridgeCreation: + """Smoke-test that the bridge loads with the right config flags.""" + + def test_norm_bridges_are_rms(self, zamba2_bridge: TransformerBridge) -> None: + """normalization_type='RMS' must wire RMSNormalizationBridge, not LayerNorm. + + The block pre-norm (on Mamba layers) and the final norm should both be + RMS. Asserting the concrete bridge type catches a regression where the + adapter wires the wrong normalization component. + """ + lbt = getattr(zamba2_bridge.cfg, "layers_block_type", []) + first_mamba = next(i for i, t in enumerate(lbt) if t == "mamba") + assert isinstance(zamba2_bridge.blocks[first_mamba].norm, RMSNormalizationBridge) + assert isinstance(zamba2_bridge.ln_final, RMSNormalizationBridge) + + def test_block_count(self, zamba2_bridge: TransformerBridge) -> None: + # Zamba2-1.2B has 38 layers (32 "mamba" + 6 "hybrid"); the block count + # must match the per-layer type list rather than a hard-coded magic + # number so the test tracks the checkpoint's actual config. + lbt = getattr(zamba2_bridge.cfg, "layers_block_type", []) + assert len(zamba2_bridge.blocks) == len(lbt) == 38 + + def test_blocks_are_ssm_block_bridge(self, zamba2_bridge: TransformerBridge) -> None: + assert isinstance(zamba2_bridge.blocks[0], SSMBlockBridge) + + def test_mamba_layers_have_mixer(self, zamba2_bridge: TransformerBridge) -> None: + """Mamba (mamba) layers should expose SSM2MixerBridge as .mixer.""" + lbt = getattr(zamba2_bridge.cfg, "layers_block_type", []) + first_mamba = next(i for i, t in enumerate(lbt) if t == "mamba") + assert isinstance(zamba2_bridge.blocks[first_mamba].mixer, SSM2MixerBridge) + + def test_hybrid_layers_have_no_mixer(self, zamba2_bridge: TransformerBridge) -> None: + """Hybrid layers have no top-level .mamba; .mixer should be absent or None.""" + lbt = getattr(zamba2_bridge.cfg, "layers_block_type", []) + first_hybrid = next(i for i, t in enumerate(lbt) if t == "hybrid") + # The optional=True submodule should not be set up by component_setup + mixer = getattr(zamba2_bridge.blocks[first_hybrid], "mixer", None) + assert mixer is None or getattr( + mixer, "optional", False + ), "Hybrid layer should not have a wired .mixer (no top-level .mamba attribute)" + + def test_layers_block_type_populated(self, zamba2_bridge: TransformerBridge) -> None: + lbt = getattr(zamba2_bridge.cfg, "layers_block_type", []) + assert len(lbt) == len(zamba2_bridge.blocks) + # Both layer types must appear + assert "mamba" in lbt, "No mamba (Mamba) layers found" + assert "hybrid" in lbt, "No hybrid (Mamba + attention) layers found" + + def test_mamba_intermediate_size_positive(self, zamba2_bridge: TransformerBridge) -> None: + assert getattr(zamba2_bridge.cfg, "mamba_intermediate_size", 0) > 0 + + def test_conv_dim_positive(self, zamba2_bridge: TransformerBridge) -> None: + assert getattr(zamba2_bridge.cfg, "conv_dim", 0) > 0 + + def test_uses_standard_kv_cache_path(self, zamba2_bridge: TransformerBridge) -> None: + # Zamba2 threads its unified cache via past_key_values (standard KV + # path), NOT the Mamba cache_params path. is_stateful=True would select + # the Mamba path, whose cache_params kwarg collides with Zamba2's own + # forward. Keeping it False routes generation through the past_key_values + # path, which matches HF generate() bit-for-bit (see TestZamba2Generation). + assert zamba2_bridge.cfg.is_stateful is False + + def test_positional_embedding_none(self, zamba2_bridge: TransformerBridge) -> None: + # RoPE is handled inside the HF attention block; no model-level embedding + assert zamba2_bridge.cfg.positional_embedding_type == "none" + + +# --------------------------------------------------------------------------- +# Forward-pass parity +# --------------------------------------------------------------------------- + + +class TestZamba2ForwardPass: + """Bridge logits must match HF logits exactly. + + Zamba2ArchitectureAdapter uses SSMBlockBridge with a pure passthrough + forward, so the bridge never reimplements any computation. Parity with + HF should be exact (diff == 0), not just close. + """ + + @pytest.fixture(scope="class") + def tokens(self) -> torch.Tensor: + return torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]]) + + def test_forward_returns_logits( + self, zamba2_bridge: TransformerBridge, tokens: torch.Tensor + ) -> None: + tokens = tokens.to(_device()) + with torch.no_grad(): + out = zamba2_bridge(tokens) + assert out.shape == (1, 8, zamba2_bridge.cfg.d_vocab) + assert not torch.isnan(out).any(), "NaN in bridge logits" + assert not torch.isinf(out).any(), "Inf in bridge logits" + + def test_forward_matches_hf_exactly( + self, zamba2_bridge: TransformerBridge, tokens: torch.Tensor + ) -> None: + tokens = tokens.to(_device()) + hf_model = zamba2_bridge.original_model + with torch.no_grad(): + bridge_out = zamba2_bridge(tokens) + hf_out = hf_model(tokens).logits + max_diff = (bridge_out.float() - hf_out.float()).abs().max().item() + assert max_diff == 0.0, ( + f"Bridge vs HF forward max diff = {max_diff:.2e}. " + "Expected 0 because SSMBlockBridge.forward() is a pure passthrough." + ) + + def test_forward_no_nan_on_longer_sequence(self, zamba2_bridge: TransformerBridge) -> None: + # 32 tokens exercises multiple Mamba SSM steps and shared attention passes + tokens = torch.arange(1, 33).unsqueeze(0).to(_device()) + with torch.no_grad(): + out = zamba2_bridge(tokens) + assert not torch.isnan(out).any(), "NaN in logits for 32-token sequence" + + +# --------------------------------------------------------------------------- +# Multi-token generation parity (exercises DynamicCache state handling) +# --------------------------------------------------------------------------- + + +class TestZamba2Generation: + """Bridge greedy generation must match HF native generate() exactly. + + DynamicCache carries both KV-cache entries (from the shared attention + blocks in hybrid layers) and Mamba-2 conv/recurrent states. Token-level + equality with HF confirms the state threading is correct across both + layer types. + """ + + @pytest.fixture(scope="class") + def prompt(self) -> torch.Tensor: + return torch.tensor([[1, 2, 3, 4]]) + + def test_generation_produces_tokens( + self, zamba2_bridge: TransformerBridge, prompt: torch.Tensor + ) -> None: + prompt = prompt.to(_device()) + with torch.no_grad(): + result = zamba2_bridge.generate(prompt, max_new_tokens=5, do_sample=False) + assert isinstance(result, torch.Tensor) + assert result.shape == (1, 9) # 4 prompt + 5 new + + def test_greedy_matches_hf_exactly( + self, zamba2_bridge: TransformerBridge, prompt: torch.Tensor + ) -> None: + """Bit-for-bit equality with HF generate() over 8 new tokens.""" + prompt = prompt.to(_device()) + hf_model = zamba2_bridge.original_model + with torch.no_grad(): + bridge_out = zamba2_bridge.generate(prompt, max_new_tokens=8, do_sample=False) + hf_out = hf_model.generate(prompt, max_new_tokens=8, do_sample=False, pad_token_id=0) + assert torch.equal(bridge_out, hf_out), ( + f"Token mismatch between bridge and HF.\n" + f" bridge : {bridge_out.tolist()}\n" + f" hf : {hf_out.tolist()}\n" + "DynamicCache state threading across Mamba-2 and hybrid attention layers " + "is likely wrong." + ) + + def test_generation_is_deterministic( + self, zamba2_bridge: TransformerBridge, prompt: torch.Tensor + ) -> None: + prompt = prompt.to(_device()) + with torch.no_grad(): + out1 = zamba2_bridge.generate(prompt, max_new_tokens=4, do_sample=False) + out2 = zamba2_bridge.generate(prompt, max_new_tokens=4, do_sample=False) + assert torch.equal(out1, out2), "Greedy generation is not deterministic" + + +# --------------------------------------------------------------------------- +# Hook coverage: bridge hooks fire for both Mamba and hybrid layers +# --------------------------------------------------------------------------- + + +class TestZamba2HookCoverage: + """run_with_cache captures residual stream and mixer hooks.""" + + @pytest.fixture(scope="class") + def cache(self, zamba2_bridge: TransformerBridge): + tokens = torch.tensor([[1, 2, 3, 4, 5]]).to(_device()) + with torch.no_grad(): + _, cache = zamba2_bridge.run_with_cache(tokens) + return cache + + def test_block_hooks_fire_on_all_layers(self, cache, zamba2_bridge: TransformerBridge) -> None: + """hook_in and hook_out must fire on every layer regardless of type.""" + n_blocks = len(zamba2_bridge.blocks) + # Sample first, an early, the middle, and the last layer (indices derived + # from the actual block count, not a hard-coded magic number). + for i in sorted({0, 1, n_blocks // 2, n_blocks - 1}): + assert f"blocks.{i}.hook_in" in cache, f"Missing hook_in for block {i}" + assert f"blocks.{i}.hook_out" in cache, f"Missing hook_out for block {i}" + + def test_mamba_mixer_submodule_hooks_fire( + self, cache, zamba2_bridge: TransformerBridge + ) -> None: + """Mamba (mamba) layers must expose in_proj / conv1d / out_proj hooks.""" + lbt = getattr(zamba2_bridge.cfg, "layers_block_type", []) + mamba_indices = [i for i, t in enumerate(lbt) if t == "mamba"] + assert mamba_indices, "No mamba layers found in layers_block_type" + for i in mamba_indices[:3]: + for submod in ("in_proj", "conv1d", "out_proj"): + key_in = f"blocks.{i}.mixer.{submod}.hook_in" + key_out = f"blocks.{i}.mixer.{submod}.hook_out" + assert key_in in cache, f"Missing {key_in}" + assert key_out in cache, f"Missing {key_out}" + + def test_hybrid_layers_no_mixer_hooks(self, cache, zamba2_bridge: TransformerBridge) -> None: + """Hybrid layers have no top-level .mamba, so no mixer submodule hooks.""" + lbt = getattr(zamba2_bridge.cfg, "layers_block_type", []) + hybrid_indices = [i for i, t in enumerate(lbt) if t == "hybrid"] + assert hybrid_indices, "No hybrid layers found in layers_block_type" + for i in hybrid_indices[:3]: + # mixer hooks must not appear for hybrid layers + assert ( + f"blocks.{i}.mixer.in_proj.hook_in" not in cache + ), f"Unexpected mixer hook on hybrid layer {i}" + + def test_no_transformer_specific_hooks(self, cache) -> None: + """SSMBlockBridge must not inject transformer-shaped hook names.""" + forbidden = ("hook_resid_mid", "hook_attn_out", "hook_mlp_out") + bad = [k for k in cache if any(f in k for f in forbidden)] + assert bad == [], f"Unexpected transformer-shaped hooks: {bad[:5]}" + + def test_no_nan_in_cache(self, cache) -> None: + for key, val in cache.items(): + if isinstance(val, torch.Tensor) and val.is_floating_point(): + assert not torch.isnan(val).any(), f"NaN in cache['{key}']" diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py index 727e6eaaf..748c160ca 100644 --- a/transformer_lens/factories/architecture_adapter_factory.py +++ b/transformer_lens/factories/architecture_adapter_factory.py @@ -83,6 +83,7 @@ T5ArchitectureAdapter, T5GemmaArchitectureAdapter, XGLMArchitectureAdapter, + Zamba2ArchitectureAdapter, ) # Export supported architectures @@ -165,6 +166,7 @@ "MT5ForConditionalGeneration": T5ArchitectureAdapter, "T5GemmaForConditionalGeneration": T5GemmaArchitectureAdapter, "XGLMForCausalLM": XGLMArchitectureAdapter, + "Zamba2ForCausalLM": Zamba2ArchitectureAdapter, "NanoGPTForCausalLM": NanogptArchitectureAdapter, "TransformerLensNative": NativeArchitectureAdapter, "MinGPTForCausalLM": MingptArchitectureAdapter, diff --git a/transformer_lens/model_bridge/sources/_bridge_builder.py b/transformer_lens/model_bridge/sources/_bridge_builder.py index e5cf55027..4729f92cf 100644 --- a/transformer_lens/model_bridge/sources/_bridge_builder.py +++ b/transformer_lens/model_bridge/sources/_bridge_builder.py @@ -75,6 +75,12 @@ "cond_dim", "adaln", "cross_attn", + # Zamba2 (Mamba-2 + shared-attention hybrid) + "mamba_expand", + "mamba_ngroups", + "num_mem_blocks", + "layers_block_type", + "use_shared_attention_adapter", ] diff --git a/transformer_lens/model_bridge/sources/transformers.py b/transformer_lens/model_bridge/sources/transformers.py index b79661ff5..0ea8220af 100644 --- a/transformer_lens/model_bridge/sources/transformers.py +++ b/transformer_lens/model_bridge/sources/transformers.py @@ -575,6 +575,12 @@ def boot( "cond_dim", "adaln", "cross_attn", + # Zamba2 (Mamba-2 + shared-attention hybrid) + "mamba_expand", + "mamba_ngroups", + "num_mem_blocks", + "layers_block_type", + "use_shared_attention_adapter", ] for attr in _HF_PASSTHROUGH_ATTRS: val = getattr(hf_config, attr, None) diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py index d783894c0..0f0e2fd0b 100644 --- a/transformer_lens/model_bridge/supported_architectures/__init__.py +++ b/transformer_lens/model_bridge/supported_architectures/__init__.py @@ -12,12 +12,12 @@ from transformer_lens.model_bridge.supported_architectures.baichuan import ( BaichuanArchitectureAdapter, ) -from transformer_lens.model_bridge.supported_architectures.bd3lm import ( - BD3LMArchitectureAdapter, -) from transformer_lens.model_bridge.supported_architectures.bart import ( BartArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.bd3lm import ( + BD3LMArchitectureAdapter, +) from transformer_lens.model_bridge.supported_architectures.bert import ( BertArchitectureAdapter, ) @@ -220,6 +220,9 @@ from transformer_lens.model_bridge.supported_architectures.xglm import ( XGLMArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.zamba2 import ( + Zamba2ArchitectureAdapter, +) __all__ = [ "ApertusArchitectureAdapter", @@ -296,4 +299,5 @@ "T5ArchitectureAdapter", "T5GemmaArchitectureAdapter", "XGLMArchitectureAdapter", + "Zamba2ArchitectureAdapter", ] diff --git a/transformer_lens/model_bridge/supported_architectures/zamba2.py b/transformer_lens/model_bridge/supported_architectures/zamba2.py new file mode 100644 index 000000000..3d7f61f50 --- /dev/null +++ b/transformer_lens/model_bridge/supported_architectures/zamba2.py @@ -0,0 +1,166 @@ +"""Zamba2 hybrid Mamba2-Transformer architecture adapter. + +Supports Zamba2ForCausalLM (e.g. Zyphra/Zamba2-1.2B, Zamba2-7B, Zamba2-7B-Instruct). + +Architecture overview: +- Heterogeneous layers defined by ``config.layers_block_type`` — each element + is either ``"mamba"`` (pure Mamba-2 SSM) or ``"hybrid"`` + (Mamba-2 + shared global-attention block). +- Most layers are ``Zamba2MambaDecoderLayer``: a single pre-norm + (``input_layernorm``) followed by a Mamba-2 mixer (``.mamba``). +- A recurring subset are ``Zamba2HybridLayer``: each wraps a Mamba-2 decoder + layer plus a SHARED ``Zamba2AttentionDecoderLayer``. The shared attention + block's weights are tied across all hybrid layers, cycling through + ``config.num_mem_blocks`` unique blocks. When + ``config.use_shared_attention_adapter=True``, each hybrid layer carries + an independent per-layer LoRA adapter on top of the shared attention. +- No model-level rotary embedding module is wired by the bridge — the + attention block handles RoPE internally via ``position_ids``. +- Generation runs on the standard KV-cache path: HF threads a single unified + ``Zamba2HybridDynamicCache`` via ``past_key_values`` (carrying both KV-cache + entries for attention and SSM conv/recurrent states for Mamba-2), so the + bridge does not use the Mamba-specific ``cache_params`` stateful path. + +Key adapter decisions: +- ``SSMBlockBridge`` is used for all layers. Its forward delegates entirely to + the HF layer, giving ``hook_in`` / ``hook_out`` on every layer regardless + of type. +- For Mamba layers: ``norm`` (-> ``.input_layernorm``) and ``mixer`` + (-> ``.mamba``) are declared as submodules and expose inner hooks + (in_proj, conv1d, inner_norm, out_proj). +- For Hybrid layers: ``norm`` and ``mixer`` are marked ``optional=True`` so + component_setup skips them gracefully (Hybrid layers have no top-level + ``.input_layernorm`` or ``.mamba``). Block-level ``hook_in``/``hook_out`` + still fire on every layer. +- ``applicable_phases = [1, 2, 3, 4]``: P1 is exact vs raw HF (pure + passthrough); P2/P3 skip without a HookedTransformer; P4 exercises + ``past_key_values`` cache threading across Mamba-2 and attention layers. +""" + +from typing import Any + +from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter +from transformer_lens.model_bridge.generalized_components import ( + DepthwiseConv1DBridge, + EmbeddingBridge, + GatedRMSNormBridge, + LinearBridge, + RMSNormalizationBridge, + SSM2MixerBridge, + SSMBlockBridge, + UnembeddingBridge, +) +from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, +) + + +def _make_optional(component: "GeneralizedComponent") -> "GeneralizedComponent": + """Mark a GeneralizedComponent submodule as optional. + + ``component_setup.py`` reads ``getattr(submodule, 'optional', False)`` at + setup time, so setting the attribute directly is safe regardless of whether + the bridge class's ``__init__`` accepts an ``optional`` keyword argument. + """ + component.optional = True + return component + + +class Zamba2ArchitectureAdapter(ArchitectureAdapter): + """Architecture adapter for Zamba2ForCausalLM. + + Hybrid Mamba-2 + shared global-attention model. Most layers are pure + Mamba-2 SSM (``"mamba"``); a recurring subset are hybrid layers + (``"hybrid"``) that route through a shared attention block before the + Mamba-2 step. + """ + + # P1: exact passthrough vs raw HF; P2/P3: skip without HookedTransformer; + # P4: generation with past_key_values cache threading. + applicable_phases: list[int] = [1, 2, 3, 4] + + def __init__(self, cfg: Any) -> None: + super().__init__(cfg) + + self.cfg.normalization_type = "RMS" + self.cfg.uses_rms_norm = True + # RoPE is handled internally by the attention block (position_ids are + # threaded through HF layers); no model-level rotary bridge needed. + self.cfg.positional_embedding_type = "none" + # MLP inside the shared attention block uses GELU, not SwiGLU. + self.cfg.gated_mlp = False + self.cfg.attn_only = False + self.cfg.final_rms = True + # NOTE: is_stateful stays False even though Mamba-2 layers carry SSM + # state. In this bridge, ``is_stateful=True`` selects the *Mamba* cache + # path, which threads the cache as ``cache_params=`` and drives a + # conv-kernel ``cache_position``. Zamba2's HF forward instead threads a + # single unified ``Zamba2HybridDynamicCache`` via ``past_key_values=`` + # (both KV entries and SSM conv/recurrent states live in that one + # object). The standard KV-cache generation path already threads + # ``past_key_values`` correctly and matches HF ``generate`` bit-for-bit, + # so we use it rather than the Mamba path (whose ``cache_params`` kwarg + # collides with Zamba2's own ``cache_params=past_key_values`` call). + self.cfg.is_stateful = False + + # Expose the per-layer type list so analysis tools can identify which + # layers are Mamba-only vs hybrid. + # Values: "mamba" (Mamba-2) | "hybrid" (Mamba-2 + attention). + layers_block_type = list(getattr(cfg, "layers_block_type", [])) + setattr(self.cfg, "layers_block_type", layers_block_type) + + # Number of unique shared attention weight blocks (hybrid layers cycle + # through num_mem_blocks independent attention weight sets). + setattr(self.cfg, "num_mem_blocks", getattr(cfg, "num_mem_blocks", 1)) + + # Whether per-layer LoRA adapters are active on the shared attention. + setattr( + self.cfg, + "use_shared_attention_adapter", + getattr(cfg, "use_shared_attention_adapter", False), + ) + + # Mamba-2 dimensional config (mirrors Mamba2ArchitectureAdapter / + # NemotronHArchitectureAdapter patterns). + mamba_intermediate_size = int(getattr(cfg, "mamba_expand", 2) * self.cfg.d_model) + n_groups = getattr(cfg, "mamba_ngroups", 1) + ssm_state_size = getattr(cfg, "mamba_d_state", 64) + conv_dim = mamba_intermediate_size + 2 * n_groups * ssm_state_size + setattr(self.cfg, "mamba_intermediate_size", mamba_intermediate_size) + setattr(self.cfg, "conv_dim", conv_dim) + + self.weight_processing_conversions = {} + + self.component_mapping = { + "embed": EmbeddingBridge(name="model.embed_tokens"), + "blocks": SSMBlockBridge( + name="model.layers", + submodules={ + # Pre-norm: present on Zamba2MambaDecoderLayer (.input_layernorm), + # absent on Zamba2HybridLayer (no top-level pre-norm) -> optional. + "norm": _make_optional( + RMSNormalizationBridge(name="input_layernorm", config=self.cfg) + ), + # Mamba-2 mixer: present on Zamba2MambaDecoderLayer (.mamba), + # absent at the top level of Zamba2HybridLayer -> optional. + "mixer": _make_optional( + SSM2MixerBridge( + name="mamba", + config=self.cfg, + submodules={ + # -- Mamba-2 inner submodules (all optional) -- + "in_proj": LinearBridge(name="in_proj", optional=True), + "conv1d": DepthwiseConv1DBridge(name="conv1d", optional=True), + # HF names the gated RMS norm "norm" inside the + # mixer; TL uses "inner_norm" to avoid colliding + # with the block-level norm declared above. + "inner_norm": _make_optional(GatedRMSNormBridge(name="norm")), + "out_proj": LinearBridge(name="out_proj", optional=True), + }, + ) + ), + }, + ), + "ln_final": RMSNormalizationBridge(name="model.final_layernorm", config=self.cfg), + "unembed": UnembeddingBridge(name="lm_head"), + } diff --git a/transformer_lens/tools/model_registry/__init__.py b/transformer_lens/tools/model_registry/__init__.py index ad1effe43..9a38c22b5 100644 --- a/transformer_lens/tools/model_registry/__init__.py +++ b/transformer_lens/tools/model_registry/__init__.py @@ -117,6 +117,7 @@ "MT5ForConditionalGeneration", "T5GemmaForConditionalGeneration", "XGLMForCausalLM", + "Zamba2ForCausalLM", } # Foundation-trained orgs per architecture. Source of truth for the scraper's @@ -194,6 +195,7 @@ "T5ForConditionalGeneration": ["google-t5", "google", "Salesforce", "MBZUAI"], "T5GemmaForConditionalGeneration": ["google"], "XGLMForCausalLM": ["facebook"], + "Zamba2ForCausalLM": ["Zyphra"], } __all__ = [ From b34a37890936a77ea3a0fabe24a4b73d4e6ffefb Mon Sep 17 00:00:00 2001 From: Delaida Muminovic Date: Tue, 7 Jul 2026 16:15:58 +0200 Subject: [PATCH 09/11] Centralize soft-cap logic behind a shared helper and a named disabled sentinel (#1491) Consolidates the tanh soft-capping (Gemma-2 style) used for attention scores and output logits into apply_softcap()/softcap_enabled() in transformer_lens/utilities/activation_functions.py, replacing five independent cap * tanh(x / cap) implementations across HookedTransformer, abstract_attention, and the native TransformerBridge model, plus the center_unembed compatibility checks in both systems. Also replaces the bare -1.0 disabled sentinel with a named SOFTCAP_DISABLED constant in both HookedTransformerConfig and TransformerBridgeConfig. Pure refactor, no behavior change. A naive 'if cap:' truthiness check would treat -1.0 as enabled (computing -1.0 * tanh(x / -1.0) == tanh(x), saturating every large score to 1.0); softcap_enabled() makes that mistake structurally impossible instead of relying on every call site getting 'cap > 0' right independently. Verified: uv run mypy . clean (307 files), make unit-test clean (3204 passed, 32 skipped, 9 xfailed, 0 failed), docstring tests on changed files clean. Fixes #1489 --- .../utilities/test_activation_functions.py | 70 +++++++++++++++++++ transformer_lens/HookedTransformer.py | 10 ++- .../components/abstract_attention.py | 6 +- .../config/hooked_transformer_config.py | 9 ++- .../config/transformer_bridge_config.py | 6 +- transformer_lens/model_bridge/bridge.py | 3 +- .../model_bridge/sources/native/model.py | 9 +-- transformer_lens/utilities/__init__.py | 3 + .../utilities/activation_functions.py | 18 ++++- 9 files changed, 111 insertions(+), 23 deletions(-) create mode 100644 tests/unit/utilities/test_activation_functions.py diff --git a/tests/unit/utilities/test_activation_functions.py b/tests/unit/utilities/test_activation_functions.py new file mode 100644 index 000000000..3a0cf12b0 --- /dev/null +++ b/tests/unit/utilities/test_activation_functions.py @@ -0,0 +1,70 @@ +"""Unit tests for the shared soft-cap helpers. + +Covers ``softcap_enabled`` and ``apply_softcap`` in +``transformer_lens.utilities.activation_functions`` (see issue #1489). +""" + +import torch + +from transformer_lens.utilities.activation_functions import ( + SOFTCAP_DISABLED, + apply_softcap, + softcap_enabled, +) + + +class TestSoftcapEnabled: + """Test cases for softcap_enabled.""" + + def test_none_is_disabled(self): + assert softcap_enabled(None) is False + + def test_zero_is_disabled(self): + assert softcap_enabled(0.0) is False + + def test_disabled_sentinel_is_disabled(self): + assert softcap_enabled(SOFTCAP_DISABLED) is False + + def test_negative_is_disabled(self): + assert softcap_enabled(-5.0) is False + + def test_positive_is_enabled(self): + assert softcap_enabled(30.0) is True + + +class TestApplySoftcap: + """Test cases for apply_softcap.""" + + def test_disabled_is_identity(self): + x = torch.tensor([-100.0, -1.0, 0.0, 1.0, 100.0]) + result = apply_softcap(x, SOFTCAP_DISABLED) + assert torch.equal(result, x) + + def test_none_is_identity(self): + x = torch.tensor([-100.0, 0.0, 100.0]) + result = apply_softcap(x, None) + assert torch.equal(result, x) + + def test_enabled_matches_capped_tanh_formula(self): + x = torch.tensor([-50.0, -10.0, 0.0, 10.0, 50.0]) + cap = 30.0 + result = apply_softcap(x, cap) + expected = cap * torch.tanh(x / cap) + assert torch.allclose(result, expected) + + def test_enabled_output_stays_within_cap(self): + x = torch.tensor([-1000.0, 1000.0]) + cap = 30.0 + result = apply_softcap(x, cap) + assert torch.all(result.abs() <= cap) + + def test_disabled_sentinel_does_not_saturate_large_values(self): + """Regression guard for the truthiness trap the issue describes: + a naive ``if cap:`` check treats -1.0 as enabled, which computes + ``-1.0 * tanh(x / -1.0) == tanh(x)`` and saturates every large score + to 1.0. The disabled path must instead be a pure identity. + """ + x = torch.tensor([50.0, 100.0]) + result = apply_softcap(x, SOFTCAP_DISABLED) + assert torch.equal(result, x) + assert not torch.allclose(result, torch.ones_like(x)) diff --git a/transformer_lens/HookedTransformer.py b/transformer_lens/HookedTransformer.py index ee2134bd1..2b9810fe9 100644 --- a/transformer_lens/HookedTransformer.py +++ b/transformer_lens/HookedTransformer.py @@ -32,7 +32,6 @@ import numpy as np import torch import torch.nn as nn -import torch.nn.functional as F import tqdm.auto as tqdm from jaxtyping import Float, Int from transformers import AutoTokenizer, PreTrainedModel, PreTrainedTokenizerBase @@ -66,12 +65,14 @@ from transformer_lens.utilities import ( USE_DEFAULT_VALUE, TypedModuleList, + apply_softcap, get_best_available_device, get_device_for_block_index, init_kaiming_normal_, init_kaiming_uniform_, init_xavier_normal_, init_xavier_uniform_, + softcap_enabled, ) from transformer_lens.utilities.devices import move_to_and_update_config from transformer_lens.weight_processing import ProcessWeights @@ -669,10 +670,7 @@ def forward( return None else: logits = self.unembed(residual) # [batch, pos, d_vocab] - if self.cfg.output_logits_soft_cap > 0.0: - logits = self.cfg.output_logits_soft_cap * F.tanh( - logits / self.cfg.output_logits_soft_cap - ) + logits = apply_softcap(logits, self.cfg.output_logits_soft_cap) if return_type == "logits": return logits else: @@ -1420,7 +1418,7 @@ def from_pretrained( "architecture. Setting center_writing_weights=False." ) center_writing_weights = False - if center_unembed and cfg.output_logits_soft_cap > 0.0: + if center_unembed and softcap_enabled(cfg.output_logits_soft_cap): logging.warning( "You tried to specify center_unembed=True for a model using logit softcap, but this can't be done! Softcapping is not invariant upon adding a constant " "Setting center_unembed=False instead." diff --git a/transformer_lens/components/abstract_attention.py b/transformer_lens/components/abstract_attention.py index 01a0646eb..9f0e819ba 100644 --- a/transformer_lens/components/abstract_attention.py +++ b/transformer_lens/components/abstract_attention.py @@ -19,6 +19,7 @@ from transformer_lens.FactoredMatrix import FactoredMatrix from transformer_lens.hook_points import HookPoint from transformer_lens.utilities import get_offset_position_ids +from transformer_lens.utilities.activation_functions import apply_softcap from transformer_lens.utilities.attention import complex_attn_linear, simple_attn_linear if is_bitsandbytes_available(): @@ -522,10 +523,7 @@ def calculate_attention_scores( k, "batch key_pos head_index d_head -> batch head_index d_head key_pos" ) attn_scores = q_ @ k_ / self.attn_scale - if self.cfg.attn_scores_soft_cap > 0: - attn_scores = self.cfg.attn_scores_soft_cap * F.tanh( - attn_scores / self.cfg.attn_scores_soft_cap - ) + attn_scores = apply_softcap(attn_scores, self.cfg.attn_scores_soft_cap) return attn_scores def calculate_z_scores( diff --git a/transformer_lens/config/hooked_transformer_config.py b/transformer_lens/config/hooked_transformer_config.py index 57ed34745..0e4a757ef 100644 --- a/transformer_lens/config/hooked_transformer_config.py +++ b/transformer_lens/config/hooked_transformer_config.py @@ -14,7 +14,10 @@ import numpy as np import torch -from transformer_lens.utilities.activation_functions import SUPPORTED_ACTIVATIONS +from transformer_lens.utilities.activation_functions import ( + SOFTCAP_DISABLED, + SUPPORTED_ACTIVATIONS, +) from transformer_lens.utilities.devices import get_device from .transformer_lens_config import TransformerLensConfig @@ -280,8 +283,8 @@ class HookedTransformerConfig(TransformerLensConfig): decoder_start_token_id: Optional[int] = None tie_word_embeddings: bool = False use_normalization_before_and_after: bool = False - attn_scores_soft_cap: float = -1.0 - output_logits_soft_cap: float = -1.0 + attn_scores_soft_cap: float = SOFTCAP_DISABLED + output_logits_soft_cap: float = SOFTCAP_DISABLED use_NTK_by_parts_rope: bool = False NTK_by_parts_low_freq_factor: float = 1.0 NTK_by_parts_high_freq_factor: float = 4.0 diff --git a/transformer_lens/config/transformer_bridge_config.py b/transformer_lens/config/transformer_bridge_config.py index 593407451..9938c6313 100644 --- a/transformer_lens/config/transformer_bridge_config.py +++ b/transformer_lens/config/transformer_bridge_config.py @@ -4,6 +4,8 @@ import torch +from transformer_lens.utilities.activation_functions import SOFTCAP_DISABLED + from .transformer_lens_config import TransformerLensConfig @@ -77,8 +79,8 @@ def __init__( decoder_start_token_id: Optional[int] = None, tie_word_embeddings: bool = False, use_normalization_before_and_after: bool = False, - attn_scores_soft_cap: float = -1.0, - output_logits_soft_cap: float = -1.0, + attn_scores_soft_cap: float = SOFTCAP_DISABLED, + output_logits_soft_cap: float = SOFTCAP_DISABLED, use_NTK_by_parts_rope: bool = False, NTK_by_parts_low_freq_factor: float = 1.0, NTK_by_parts_high_freq_factor: float = 4.0, diff --git a/transformer_lens/model_bridge/bridge.py b/transformer_lens/model_bridge/bridge.py index 4800b0150..9d638b7b5 100644 --- a/transformer_lens/model_bridge/bridge.py +++ b/transformer_lens/model_bridge/bridge.py @@ -50,6 +50,7 @@ VARIANT_SUBMODULE_NAMES, ) from transformer_lens.model_bridge.get_params_util import get_bridge_params +from transformer_lens.utilities.activation_functions import softcap_enabled from transformer_lens.utilities.aliases import resolve_alias from transformer_lens.utilities.devices import move_to_and_update_config from transformer_lens.utilities.lm_utils import lm_cross_entropy_loss @@ -920,7 +921,7 @@ def process_weights( print(f"Processing weights for {self.cfg.model_name}...") # Soft capping (tanh) is not translation-invariant; centering would change output. - if center_unembed and getattr(self.cfg, "output_logits_soft_cap", -1.0) > 0.0: + if center_unembed and softcap_enabled(getattr(self.cfg, "output_logits_soft_cap", None)): import logging logging.warning( diff --git a/transformer_lens/model_bridge/sources/native/model.py b/transformer_lens/model_bridge/sources/native/model.py index 4cae95a2a..5617459e7 100644 --- a/transformer_lens/model_bridge/sources/native/model.py +++ b/transformer_lens/model_bridge/sources/native/model.py @@ -18,6 +18,7 @@ from transformer_lens.config import TransformerBridgeConfig from transformer_lens.utilities import TypedModuleList +from transformer_lens.utilities.activation_functions import apply_softcap # gelu_new = the tanh-approximation HF GPT-2 / HT use; F.gelu(approximate="tanh") # is the exact same formula. @@ -269,9 +270,7 @@ def forward( scores = torch.matmul(q, k.transpose(-2, -1)) / self.scale # Gemma2 soft-cap before the causal mask so masked positions stay -inf. - if self.attn_scores_soft_cap > 0: - c = self.attn_scores_soft_cap - scores = c * torch.tanh(scores / c) + scores = apply_softcap(scores, self.attn_scores_soft_cap) if self.causal: block_mask = self.causal_mask[:seq, :seq] @@ -468,7 +467,5 @@ def forward( ) hidden_states = self.ln_out(hidden_states) logits = self.head(hidden_states) - if self.output_logits_soft_cap > 0: - c = self.output_logits_soft_cap - logits = c * torch.tanh(logits / c) + logits = apply_softcap(logits, self.output_logits_soft_cap) return logits diff --git a/transformer_lens/utilities/__init__.py b/transformer_lens/utilities/__init__.py index e242e951c..319280256 100644 --- a/transformer_lens/utilities/__init__.py +++ b/transformer_lens/utilities/__init__.py @@ -1,8 +1,11 @@ from .activation_functions import ( + SOFTCAP_DISABLED, SUPPORTED_ACTIVATIONS, ActivationFunction, + apply_softcap, gelu_fast, gelu_new, + softcap_enabled, solu, ) from .attribute_utils import get_nested_attr, set_nested_attr diff --git a/transformer_lens/utilities/activation_functions.py b/transformer_lens/utilities/activation_functions.py index 3594acae8..0b650b360 100644 --- a/transformer_lens/utilities/activation_functions.py +++ b/transformer_lens/utilities/activation_functions.py @@ -3,7 +3,7 @@ Utilities for interacting with all supported activation functions. """ -from typing import Callable, Dict +from typing import Callable, Dict, Optional, cast import numpy as np import torch @@ -118,6 +118,22 @@ def relu2(input: Float[torch.Tensor, "batch pos d_mlp"]) -> Float[torch.Tensor, return relu_applied * relu_applied +SOFTCAP_DISABLED: float = -1.0 + + +def softcap_enabled(cap: Optional[float]) -> bool: + """A soft-cap is active only for a strictly positive value; -1.0 / 0 / None mean disabled.""" + return cap is not None and cap > 0 + + +def apply_softcap(x: torch.Tensor, cap: Optional[float]) -> torch.Tensor: + """Gemma-style ``cap * tanh(x / cap)`` when enabled, identity otherwise.""" + if not softcap_enabled(cap): + return x + cap = cast(float, cap) + return cap * torch.tanh(x / cap) + + # Convenient type for the format of each activation function ActivationFunction = Callable[..., torch.Tensor] From 8d01cbdbfd194327a22ecba66a393137bb026f79 Mon Sep 17 00:00:00 2001 From: Alessandro Canonico Date: Wed, 8 Jul 2026 20:28:46 +0200 Subject: [PATCH 10/11] Added Ouro architecture adapter and integration tests for OuroForCausalLM (#1488) * Add Ouro architecture adapter and integration tests for OuroForCausalLM * add missing quarantine row for the gated ouro integration test * feat: added Ouro architecture 'support_fold_ln = False' and test verification updates. Chore: formatting updates in .MD --- tests/QUARANTINES.md | 1 + .../model_bridge/test_ouro_adapter.py | 235 +++++++++++++++ .../test_ouro_adapter.py | 274 ++++++++++++++++++ .../factories/architecture_adapter_factory.py | 2 + .../model_bridge/sources/_bridge_builder.py | 3 + .../model_bridge/sources/transformers.py | 4 + .../supported_architectures/__init__.py | 4 + .../supported_architectures/ouro.py | 227 +++++++++++++++ .../tools/model_registry/__init__.py | 2 + .../model_registry/data/supported_models.json | 20 +- .../data/verification_history.json | 12 +- .../tools/model_registry/generate_report.py | 1 + .../tools/model_registry/verify_models.py | 1 + 13 files changed, 782 insertions(+), 4 deletions(-) create mode 100644 tests/integration/model_bridge/test_ouro_adapter.py create mode 100644 tests/unit/model_bridge/supported_architectures/test_ouro_adapter.py create mode 100644 transformer_lens/model_bridge/supported_architectures/ouro.py diff --git a/tests/QUARANTINES.md b/tests/QUARANTINES.md index 5a587ea38..0026b8867 100644 --- a/tests/QUARANTINES.md +++ b/tests/QUARANTINES.md @@ -46,6 +46,7 @@ Rule ([AGENTS.md §10](../AGENTS.md#10-hard-rules)): **never add `xfail` / `skip | [`unit/model_bridge/supported_architectures/test_gemma2_adapter.py`:49](unit/model_bridge/supported_architectures/test_gemma2_adapter.py) | "Network/disk fetch of tiny Gemma2" | | [`integration/model_bridge/test_bridge_integration.py`:801](integration/model_bridge/test_bridge_integration.py) | "Skip Gemma2 in CI to avoid timeout" | | [`acceptance/model_bridge/compatibility/test_hook_completeness.py`:156](acceptance/model_bridge/compatibility/test_hook_completeness.py) | "Gemma2 too large for CI" | +| [`integration/model_bridge/test_ouro_adapter.py`:35](integration/model_bridge/test_ouro_adapter.py) module-level (+ `slow` marker) | "ByteDance/Ouro-1.4B: 2.8GB download + ~11GB RAM, too large for CI" | **Un-skip:** locally with `HF_TOKEN` sourced. diff --git a/tests/integration/model_bridge/test_ouro_adapter.py b/tests/integration/model_bridge/test_ouro_adapter.py new file mode 100644 index 000000000..48f46b013 --- /dev/null +++ b/tests/integration/model_bridge/test_ouro_adapter.py @@ -0,0 +1,235 @@ +"""Integration tests for the Ouro architecture adapter (OuroForCausalLM). + +Model: ByteDance/Ouro-1.4B + - Remote-code (auto_map to modeling_ouro), so no tiny random checkpoint + exists; this loads the real 1.4B weights (~2.8GB download, ~11GB RAM for + the two fp32 copies). Gated out of CI on network/memory budget; run + locally with: + uv run pytest tests/integration/model_bridge/test_ouro_adapter.py -v -m slow + +Ouro is a looped-depth (Universal Transformer) model: the remote-code +OuroModel.forward applies the shared 24-layer stack total_ut_steps=4 times, +with model.norm after each pass. The bridge delegates the loop to HF's own +forward, so logit parity holds unchanged; the one observable difference from +a standard decoder is that each physical block's hooks fire once per UT step +(pinned by TestOuroLoopedDepth below) and a cache keeps the final step only. + +The bridge always reimplements attention in eager mode (so the score and +pattern hooks fire), so the reference HF model is loaded independently with +attn_implementation="eager" too. Comparing against an independent HF load +(never bridge.original_model, whose modules are hook-wrapped) keeps the +parity check honest. +""" + +import os +import platform + +import pytest +import torch +from transformers import AutoConfig, AutoModelForCausalLM + +from transformer_lens.model_bridge import TransformerBridge + +MODEL = "ByteDance/Ouro-1.4B" + +pytestmark = [ + pytest.mark.slow, + pytest.mark.skipif( + bool(os.getenv("CI")), + reason="ByteDance/Ouro-1.4B: 2.8GB download + ~11GB RAM, too large for CI", + ), +] + +# Sandwich-norm model with 96 effective layer applications at fp32; allow a +# wider op-order noise floor on GH Actions macOS-arm64 (same idiom as the +# SmolLM3 and general bridge-vs-HF parity tests). +_MACOS_ARM64 = platform.system() == "Darwin" and platform.machine() == "arm64" +FP32_NOISE_TOL = 1e-2 if _MACOS_ARM64 else 1e-3 + + +@pytest.fixture(scope="module") +def bridge() -> TransformerBridge: + return TransformerBridge.boot_transformers( + MODEL, device="cpu", dtype=torch.float32, trust_remote_code=True + ) + + +@pytest.fixture(scope="module") +def hf_eager(bridge: TransformerBridge) -> torch.nn.Module: + """HF model loaded independently of the bridge's wrapped instance. + + Depends on the bridge fixture only for ordering: booting the bridge runs + the adapter's prepare_loading patch (restoring the "default" RoPE init + that transformers v5 removed) on the cached modeling_ouro module, and + this plain from_pretrained load needs that patch too. The weights are + still a separate copy. + + The pad_token_id shim mirrors the bridge load path + (sources/transformers.py): transformers v5 no longer materializes a + default pad_token_id on configs, but Ouro's remote code reads + config.pad_token_id unconditionally. Mirroring the same eos fallback + keeps both copies' embedding padding_idx identical. + """ + config = AutoConfig.from_pretrained(MODEL, trust_remote_code=True) + if "pad_token_id" not in config.__dict__: + fallback_pad = getattr(config, "eos_token_id", None) + if isinstance(fallback_pad, list): + fallback_pad = fallback_pad[0] if fallback_pad else None + config.pad_token_id = fallback_pad + return AutoModelForCausalLM.from_pretrained( + MODEL, + config=config, + torch_dtype=torch.float32, + attn_implementation="eager", + trust_remote_code=True, + ).eval() + + +@pytest.fixture(scope="module") +def tokens(hf_eager: torch.nn.Module) -> torch.Tensor: + """A short, deterministic random token sequence. + + Random ids keep the test tokenizer-free and cover the vocab beyond the + few ids a short prompt would produce. + """ + generator = torch.Generator().manual_seed(0) + return torch.randint(0, hf_eager.config.vocab_size, (1, 16), generator=generator) + + +class TestOuroBridgeCreation: + """The bridge loads the remote-code model and wires the sandwich norms.""" + + def test_boot_transformers_succeeds(self, bridge: TransformerBridge) -> None: + assert bridge is not None + + def test_block_count_is_physical_layers( + self, bridge: TransformerBridge, hf_eager: torch.nn.Module + ) -> None: + """n_layers counts the 24 physical layers, not the 96 loop applications.""" + assert len(bridge.blocks) == hf_eager.config.num_hidden_layers + assert bridge.cfg.n_layers == hf_eager.config.num_hidden_layers + + def test_config_flags(self, bridge: TransformerBridge) -> None: + assert bridge.cfg.normalization_type == "RMS" + assert bridge.cfg.positional_embedding_type == "rotary" + assert bridge.cfg.gated_mlp is True + + def test_sandwich_norms_wired(self, bridge: TransformerBridge) -> None: + """All four per-layer RMSNorms are bridged.""" + block = bridge.blocks[0] + for ln_name in ("ln1", "ln1_post", "ln2", "ln2_post"): + assert hasattr(block, ln_name), f"missing {ln_name} on block 0" + + +class TestOuroHFDelegation: + """The bridge wraps the live remote-code modules in place (no copies). + + Wrapping is in-place: the HF tree slot holds the bridge component itself, + and the bridge's submodule IS the module the wrapped OuroDecoderLayer / + OuroAttention executes. Assert those identities directly. + """ + + def test_blocks_are_the_hf_tree_layers(self, bridge: TransformerBridge) -> None: + assert bridge.blocks[0] is bridge.original_model.model.layers[0] + assert type(bridge.blocks[0].original_component).__name__ == "OuroDecoderLayer" + + def test_attention_projections_wrap_live_hf_modules(self, bridge: TransformerBridge) -> None: + attn = bridge.blocks[0].attn + assert type(attn.original_component).__name__ == "OuroAttention" + assert attn.q is attn.original_component.q_proj + assert attn.k is attn.original_component.k_proj + assert attn.v is attn.original_component.v_proj + assert attn.o is attn.original_component.o_proj + + def test_sandwich_norms_wrap_live_hf_modules(self, bridge: TransformerBridge) -> None: + block = bridge.blocks[0] + layer = block.original_component + assert block.ln1 is layer.input_layernorm + assert block.ln1_post is layer.input_layernorm_2 + assert block.ln2 is layer.post_attention_layernorm + assert block.ln2_post is layer.post_attention_layernorm_2 + + +class TestOuroForwardEquivalence: + """Bridge output reproduces the HF eager reference through the full UT loop.""" + + def test_forward_logits_match_hf_eager( + self, bridge: TransformerBridge, hf_eager: torch.nn.Module, tokens: torch.Tensor + ) -> None: + with torch.inference_mode(): + bridge_logits = bridge(tokens) + hf_logits = hf_eager(tokens).logits + max_diff = (bridge_logits - hf_logits).abs().max().item() + assert max_diff < FP32_NOISE_TOL, ( + f"Ouro bridge vs HF eager logit drift={max_diff:.2e} exceeds the " + f"fp32-noise tolerance {FP32_NOISE_TOL:.0e}." + ) + + +class TestOuroLoopedDepth: + """Pin the Universal-Transformer loop semantics the adapter documents.""" + + def test_block_hook_fires_once_per_ut_step( + self, bridge: TransformerBridge, tokens: torch.Tensor + ) -> None: + """The same physical block executes total_ut_steps times per forward. + + This is the load-bearing behavioural difference from a standard + decoder: hooks on blocks.{i} fire once per loop step, and + run_with_cache keeps the final step's value. If HF's remote code ever + changes the loop (or the bridge stops delegating it), this fails. + """ + fired: list[bool] = [] + bridge.run_with_hooks( + tokens, + fwd_hooks=[("blocks.0.hook_out", lambda value, hook: fired.append(True))], + ) + assert len(fired) == bridge.cfg.total_ut_steps + + def test_hook_out_shape_matches_residual_stream( + self, bridge: TransformerBridge, tokens: torch.Tensor + ) -> None: + captured: list[torch.Tensor] = [] + + def _capture(value, hook): + captured.append(value[0] if isinstance(value, tuple) else value) + return value + + bridge.run_with_hooks(tokens, fwd_hooks=[("blocks.0.hook_out", _capture)]) + assert captured + assert captured[-1].shape == (1, tokens.shape[1], bridge.cfg.d_model) + + def test_bridge_runs_its_own_attention_reconstruction( + self, bridge: TransformerBridge, tokens: torch.Tensor + ) -> None: + """Anti-tautology guard: prove the bridge's custom attention path executes. + + If a future refactor made the bridge delegate to HF attention directly, + the parity test above would pass trivially. Assert a bridge-specific + hook fires during the forward pass instead. + """ + fired: list[bool] = [] + bridge.run_with_hooks( + tokens, + fwd_hooks=[ + ("blocks.0.attn.hook_attn_scores", lambda value, hook: fired.append(True)), + ], + ) + assert fired, ( + "blocks.0.attn.hook_attn_scores did not fire, so the bridge no longer " + "runs its own attention reconstruction and the parity test is tautological." + ) + + +class TestOuroConfigPropagation: + """Loop-related HF config attrs surface on bridge.cfg via _HF_PASSTHROUGH_ATTRS.""" + + def test_total_ut_steps_propagates( + self, bridge: TransformerBridge, hf_eager: torch.nn.Module + ) -> None: + assert bridge.cfg.total_ut_steps == hf_eager.config.total_ut_steps + + def test_early_exit_threshold_propagates( + self, bridge: TransformerBridge, hf_eager: torch.nn.Module + ) -> None: + assert bridge.cfg.early_exit_threshold == hf_eager.config.early_exit_threshold diff --git a/tests/unit/model_bridge/supported_architectures/test_ouro_adapter.py b/tests/unit/model_bridge/supported_architectures/test_ouro_adapter.py new file mode 100644 index 000000000..2031ca91e --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_ouro_adapter.py @@ -0,0 +1,274 @@ +"""Unit tests for OuroArchitectureAdapter. + +Tests cover: +- Config attributes set by the adapter +- Component mapping structure and HF module paths, including Ouro's + sandwich normalization (four RMSNorms per decoder layer) +- Standard Q/K/V/O weight conversion rules +- setup_component_testing rotary embedding wiring +- Factory registration + +Ouro's Universal-Transformer loop and early-exit gate live in the remote-code +HF forward and are deliberately NOT mapped by the adapter; the top-level-keys +test pins that scope (no "gate" / per-step components). +""" + +from types import SimpleNamespace + +import pytest + +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.model_bridge.generalized_components import ( + BlockBridge, + EmbeddingBridge, + GatedMLPBridge, + LinearBridge, + PositionEmbeddingsAttentionBridge, + RMSNormalizationBridge, + RotaryEmbeddingBridge, + UnembeddingBridge, +) +from transformer_lens.model_bridge.supported_architectures.ouro import ( + OuroArchitectureAdapter, +) + + +def _make_cfg( + n_heads: int = 4, + d_model: int = 64, + n_layers: int = 2, + d_mlp: int = 256, + d_vocab: int = 100, + n_ctx: int = 64, +) -> TransformerBridgeConfig: + # Keep dimensions tiny so adapter tests do not need HF downloads or real checkpoints. + # Ouro uses full MHA (num_key_value_heads == num_attention_heads). + return TransformerBridgeConfig( + d_model=d_model, + d_head=d_model // n_heads, + n_layers=n_layers, + n_ctx=n_ctx, + n_heads=n_heads, + n_key_value_heads=n_heads, + d_vocab=d_vocab, + d_mlp=d_mlp, + architecture="OuroForCausalLM", + ) + + +@pytest.fixture +def cfg() -> TransformerBridgeConfig: + return _make_cfg() + + +@pytest.fixture +def adapter(cfg: TransformerBridgeConfig) -> OuroArchitectureAdapter: + return OuroArchitectureAdapter(cfg) + + +def _fake_hf_model(rotary_emb: object) -> SimpleNamespace: + return SimpleNamespace(model=SimpleNamespace(rotary_emb=rotary_emb)) + + +class DummyAttention: + def __init__(self) -> None: + self.rotary_emb = None + + def set_rotary_emb(self, rotary_emb: object) -> None: + self.rotary_emb = rotary_emb + + +class DummyBlock: + def __init__(self, has_attention: bool = True) -> None: + if has_attention: + self.attn = DummyAttention() + + +class DummyBridgeModel: + def __init__(self, blocks: list[DummyBlock]) -> None: + self.blocks = blocks + + +class TestOuroAdapterConfig: + """Adapter-owned config defaults that downstream bridge code relies on.""" + + def test_normalization_flags(self, adapter: OuroArchitectureAdapter) -> None: + assert adapter.cfg.normalization_type == "RMS" + assert adapter.cfg.uses_rms_norm is True + assert adapter.cfg.final_rms is True + + def test_rotary_and_mlp_flags(self, adapter: OuroArchitectureAdapter) -> None: + assert adapter.cfg.positional_embedding_type == "rotary" + assert adapter.cfg.gated_mlp is True + assert adapter.cfg.attn_only is False + + def test_no_rmsnorm_offset(self, adapter: OuroArchitectureAdapter) -> None: + """Ouro's RMSNorm applies the weight directly (no Gemma-style +1.0 offset).""" + assert not getattr(adapter.cfg, "rmsnorm_uses_offset", False) + + def test_supports_fold_ln_is_false(self, adapter: OuroArchitectureAdapter) -> None: + """ln_final runs after every UT pass, feeding the next pass and the exit + gate; folding it into W_U would corrupt passes 1..N-1 in the live module.""" + assert adapter.supports_fold_ln is False + + +class TestOuroComponentMapping: + """The adapter contract: TL canonical names mapped to Ouro HF module paths.""" + + def test_top_level_keys(self, adapter: OuroArchitectureAdapter) -> None: + # No "gate" key: model.early_exit_gate is intentionally unmapped. + assert set(adapter.component_mapping.keys()) == { + "embed", + "rotary_emb", + "blocks", + "ln_final", + "unembed", + } + + def test_embed_path(self, adapter: OuroArchitectureAdapter) -> None: + assert adapter.component_mapping["embed"].name == "model.embed_tokens" + + def test_rotary_emb_path(self, adapter: OuroArchitectureAdapter) -> None: + assert adapter.component_mapping["rotary_emb"].name == "model.rotary_emb" + + def test_blocks_path(self, adapter: OuroArchitectureAdapter) -> None: + assert adapter.component_mapping["blocks"].name == "model.layers" + + def test_ln_final_path(self, adapter: OuroArchitectureAdapter) -> None: + assert adapter.component_mapping["ln_final"].name == "model.norm" + + def test_unembed_path(self, adapter: OuroArchitectureAdapter) -> None: + assert adapter.component_mapping["unembed"].name == "lm_head" + + def test_block_submodule_keys(self, adapter: OuroArchitectureAdapter) -> None: + """Sandwich norm: four RMSNorms per block, not the usual two.""" + blocks = adapter.component_mapping["blocks"] + assert set(blocks.submodules.keys()) == { + "ln1", + "ln1_post", + "ln2", + "ln2_post", + "attn", + "mlp", + } + + def test_sandwich_norm_hf_paths(self, adapter: OuroArchitectureAdapter) -> None: + """The extra norms map to Ouro's *_2 module names. + + Forward order in OuroDecoderLayer: input_layernorm (pre-attn) -> attn + -> input_layernorm_2 (post-attn, pre-residual); post_attention_layernorm + (pre-MLP) -> mlp -> post_attention_layernorm_2 (post-MLP, pre-residual). + """ + blocks = adapter.component_mapping["blocks"] + assert blocks.submodules["ln1"].name == "input_layernorm" + assert blocks.submodules["ln1_post"].name == "input_layernorm_2" + assert blocks.submodules["ln2"].name == "post_attention_layernorm" + assert blocks.submodules["ln2_post"].name == "post_attention_layernorm_2" + + def test_attention_submodule_keys(self, adapter: OuroArchitectureAdapter) -> None: + attn = adapter.component_mapping["blocks"].submodules["attn"] + assert set(attn.submodules.keys()) == {"q", "k", "v", "o"} + + def test_mlp_submodule_keys(self, adapter: OuroArchitectureAdapter) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + assert set(mlp.submodules.keys()) == {"gate", "in", "out"} + + def test_bridge_types(self, adapter: OuroArchitectureAdapter) -> None: + mapping = adapter.component_mapping + blocks = mapping["blocks"] + assert isinstance(mapping["embed"], EmbeddingBridge) + assert isinstance(mapping["rotary_emb"], RotaryEmbeddingBridge) + assert isinstance(blocks, BlockBridge) + assert isinstance(mapping["ln_final"], RMSNormalizationBridge) + assert isinstance(mapping["unembed"], UnembeddingBridge) + for ln_key in ("ln1", "ln1_post", "ln2", "ln2_post"): + assert isinstance(blocks.submodules[ln_key], RMSNormalizationBridge) + assert isinstance(blocks.submodules["attn"], PositionEmbeddingsAttentionBridge) + assert isinstance(blocks.submodules["mlp"], GatedMLPBridge) + + def test_attention_hf_paths(self, adapter: OuroArchitectureAdapter) -> None: + attn = adapter.component_mapping["blocks"].submodules["attn"] + assert attn.name == "self_attn" + assert attn.submodules["q"].name == "q_proj" + assert attn.submodules["k"].name == "k_proj" + assert attn.submodules["v"].name == "v_proj" + assert attn.submodules["o"].name == "o_proj" + + def test_mlp_hf_paths(self, adapter: OuroArchitectureAdapter) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + assert mlp.name == "mlp" + assert mlp.submodules["gate"].name == "gate_proj" + assert mlp.submodules["in"].name == "up_proj" + assert mlp.submodules["out"].name == "down_proj" + + def test_linear_submodule_bridge_types(self, adapter: OuroArchitectureAdapter) -> None: + blocks = adapter.component_mapping["blocks"] + attn = blocks.submodules["attn"] + mlp = blocks.submodules["mlp"] + for submodule in [*attn.submodules.values(), *mlp.submodules.values()]: + assert isinstance(submodule, LinearBridge) + + +class TestOuroWeightConversions: + """Standard split-QKV conversion rules from _qkvo_weight_conversions().""" + + def test_qkvo_conversion_keys_present(self, adapter: OuroArchitectureAdapter) -> None: + for key in ( + "blocks.{i}.attn.q.weight", + "blocks.{i}.attn.k.weight", + "blocks.{i}.attn.v.weight", + "blocks.{i}.attn.o.weight", + ): + assert key in adapter.weight_processing_conversions + + +class TestOuroSetupComponentTesting: + """setup_component_testing must wire Ouro's shared rotary embedding into attention bridges.""" + + def test_sets_rotary_emb_on_template_attention(self, adapter: OuroArchitectureAdapter) -> None: + rotary_emb = object() + attn_template = adapter.get_generalized_component("blocks.0.attn") + assert isinstance(attn_template, PositionEmbeddingsAttentionBridge) + assert attn_template._rotary_emb is None + + adapter.setup_component_testing(_fake_hf_model(rotary_emb)) + + assert attn_template._rotary_emb is rotary_emb + + def test_sets_rotary_emb_on_each_bridge_model_attention( + self, adapter: OuroArchitectureAdapter + ) -> None: + rotary_emb = object() + bridge_model = DummyBridgeModel([DummyBlock(), DummyBlock(), DummyBlock()]) + + adapter.setup_component_testing(_fake_hf_model(rotary_emb), bridge_model=bridge_model) + + for block in bridge_model.blocks: + assert block.attn.rotary_emb is rotary_emb + + def test_skips_bridge_blocks_without_attention(self, adapter: OuroArchitectureAdapter) -> None: + rotary_emb = object() + bridge_model = DummyBridgeModel([DummyBlock(), DummyBlock(has_attention=False)]) + + adapter.setup_component_testing(_fake_hf_model(rotary_emb), bridge_model=bridge_model) + + assert bridge_model.blocks[0].attn.rotary_emb is rotary_emb + + +class TestOuroFactoryRegistration: + """The factory resolves OuroForCausalLM to this adapter.""" + + def test_supported_architectures_entry(self) -> None: + from transformer_lens.factories.architecture_adapter_factory import ( + SUPPORTED_ARCHITECTURES, + ) + + assert SUPPORTED_ARCHITECTURES["OuroForCausalLM"] is OuroArchitectureAdapter + + def test_factory_selects_ouro_adapter(self, cfg: TransformerBridgeConfig) -> None: + from transformer_lens.factories.architecture_adapter_factory import ( + ArchitectureAdapterFactory, + ) + + adapter = ArchitectureAdapterFactory.select_architecture_adapter(cfg) + assert isinstance(adapter, OuroArchitectureAdapter) diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py index 748c160ca..cef0d22d7 100644 --- a/transformer_lens/factories/architecture_adapter_factory.py +++ b/transformer_lens/factories/architecture_adapter_factory.py @@ -66,6 +66,7 @@ OlmoeArchitectureAdapter, OpenElmArchitectureAdapter, OptArchitectureAdapter, + OuroArchitectureAdapter, Phi3ArchitectureAdapter, PhiArchitectureAdapter, PhiMoEArchitectureAdapter, @@ -148,6 +149,7 @@ "OlmoeForCausalLM": OlmoeArchitectureAdapter, "OpenELMForCausalLM": OpenElmArchitectureAdapter, "OPTForCausalLM": OptArchitectureAdapter, + "OuroForCausalLM": OuroArchitectureAdapter, "PhiForCausalLM": PhiArchitectureAdapter, "Phi3ForCausalLM": Phi3ArchitectureAdapter, "PhiMoEForCausalLM": PhiMoEArchitectureAdapter, diff --git a/transformer_lens/model_bridge/sources/_bridge_builder.py b/transformer_lens/model_bridge/sources/_bridge_builder.py index 4729f92cf..c64930a16 100644 --- a/transformer_lens/model_bridge/sources/_bridge_builder.py +++ b/transformer_lens/model_bridge/sources/_bridge_builder.py @@ -81,6 +81,9 @@ "num_mem_blocks", "layers_block_type", "use_shared_attention_adapter", + # Ouro (LoopLM) + "total_ut_steps", + "early_exit_threshold", ] diff --git a/transformer_lens/model_bridge/sources/transformers.py b/transformer_lens/model_bridge/sources/transformers.py index 0ea8220af..2d67a01b9 100644 --- a/transformer_lens/model_bridge/sources/transformers.py +++ b/transformer_lens/model_bridge/sources/transformers.py @@ -265,6 +265,7 @@ def determine_architecture_from_hf_config(hf_config): "qwen3_5_text": "Qwen3_5ForCausalLM", "smollm3": "SmolLM3ForCausalLM", "openelm": "OpenELMForCausalLM", + "ouro": "OuroForCausalLM", "stablelm": "StableLmForCausalLM", "t5": "T5ForConditionalGeneration", "mt5": "MT5ForConditionalGeneration", @@ -581,6 +582,9 @@ def boot( "num_mem_blocks", "layers_block_type", "use_shared_attention_adapter", + # Ouro (LoopLM) + "total_ut_steps", + "early_exit_threshold", ] for attr in _HF_PASSTHROUGH_ATTRS: val = getattr(hf_config, attr, None) diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py index 0f0e2fd0b..6ab2ad971 100644 --- a/transformer_lens/model_bridge/supported_architectures/__init__.py +++ b/transformer_lens/model_bridge/supported_architectures/__init__.py @@ -169,6 +169,9 @@ from transformer_lens.model_bridge.supported_architectures.opt import ( OptArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.ouro import ( + OuroArchitectureAdapter, +) from transformer_lens.model_bridge.supported_architectures.phi import ( PhiArchitectureAdapter, ) @@ -282,6 +285,7 @@ "Olmo3ArchitectureAdapter", "OlmoeArchitectureAdapter", "OptArchitectureAdapter", + "OuroArchitectureAdapter", "PhiArchitectureAdapter", "Phi3ArchitectureAdapter", "PhiMoEArchitectureAdapter", diff --git a/transformer_lens/model_bridge/supported_architectures/ouro.py b/transformer_lens/model_bridge/supported_architectures/ouro.py new file mode 100644 index 000000000..e443c49f9 --- /dev/null +++ b/transformer_lens/model_bridge/supported_architectures/ouro.py @@ -0,0 +1,227 @@ +"""Ouro architecture adapter.""" + +import sys +from typing import Any, Optional + +import torch + +from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter +from transformer_lens.model_bridge.generalized_components import ( + BlockBridge, + EmbeddingBridge, + GatedMLPBridge, + LinearBridge, + PositionEmbeddingsAttentionBridge, + RMSNormalizationBridge, + RotaryEmbeddingBridge, + UnembeddingBridge, +) + + +def _compute_default_rope_parameters( + config: Any, + device: Optional[torch.device] = None, + seq_len: Optional[int] = None, + **rope_kwargs: Any, +) -> tuple[torch.Tensor, float]: + """Standard (unscaled) RoPE inverse frequencies, as transformers v4 defined them. + + Transformers v5 removed the "default" entry from ROPE_INIT_FUNCTIONS (standard + RoPE moved to a per-model static method), but Ouro's remote code still looks it + up. Signature and return match the v4 contract the remote code calls with: + (config, device) -> (inv_freq, attention_scaling). + """ + base = config.rope_theta + partial_rotary_factor = getattr(config, "partial_rotary_factor", None) or 1.0 + head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + dim = int(head_dim * partial_rotary_factor) + inv_freq = 1.0 / ( + base + ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) + ) + return inv_freq, 1.0 + + +class OuroArchitectureAdapter(ArchitectureAdapter): + """Architecture adapter for ByteDance Ouro (LoopLM) models. + + Ouro is a looped-depth ("Universal Transformer") decoder: the remote-code + ``OuroModel.forward`` applies the same ``num_hidden_layers``-deep stack + ``total_ut_steps`` times (4 for the released checkpoints) within a single + forward pass, applying ``model.norm`` after every pass. The loop lives + entirely inside the HF forward, which the bridge delegates to, so logits + and generation are correct with no loop handling here. ``n_layers`` counts + the physical layers; each block's hooks fire once per loop step, and a + cache records the final step's value. The same holds for ``ln_final`` + (``model.norm``): it runs after EVERY UT pass, so its hooks fire + ``total_ut_steps`` times per forward and ``run_with_cache`` keeps only the + last pass. + + The backbone is Qwen2/Llama-shaped (RoPE, no-bias q/k/v/o projections, + SwiGLU gate/up/down MLP, untied lm_head) with one twist: sandwich + normalization. Each decoder layer has FOUR RMSNorms; the extra two + (``input_layernorm_2``, ``post_attention_layernorm_2``) apply to the + sublayer outputs before the residual add, exactly like Gemma2's + ``ln1_post``/``ln2_post`` but without Gemma's +1.0 RMSNorm offset. + + Deliberately not mapped by this adapter: + + - per-loop-step hooks (a cache holds the final UT step only) + - ``model.early_exit_gate``, the adaptive-exit halting head + - the ``UniversalTransformerCache`` slot layout (``step * n_layers + layer``) + + Loading requires ``trust_remote_code=True`` (``auto_map`` to + ``modeling_ouro``). + + Optional Parameters (may not exist in state_dict): + ------------------------------------------------- + Ouro models do NOT have biases on any mapped linear layers: + + - blocks.{i}.attn.b_Q / b_K / b_V / b_O - no attention biases + - blocks.{i}.mlp.b_gate / b_in / b_out - no MLP biases + - blocks.{i}.ln1.b / ln1_post.b / ln2.b / ln2_post.b - RMSNorm has no bias + - ln_final.b - RMSNorm has no bias + + Weight processing must handle these missing biases gracefully using + ProcessWeights._safe_get_tensor() or by checking for None values. + """ + + def __init__(self, cfg: Any) -> None: + """Initialize the Ouro architecture adapter.""" + super().__init__(cfg) + + # Set config variables for weight processing + self.cfg.normalization_type = "RMS" + self.cfg.positional_embedding_type = "rotary" + self.cfg.final_rms = True + self.cfg.gated_mlp = True + self.cfg.attn_only = False + self.cfg.uses_rms_norm = True + # default_prepend_bos stays at the framework default: the GPT2-style BPE + # tokenizer (bos == eos == <|endoftext|>) does not prepend BOS itself. + + # ln_final (model.norm) is applied after EVERY UT pass, feeding the next + # pass and the early-exit gate, so it is not a final-only norm. Folding + # it into W_U resets the live module's norm weight the loop reuses and + # corrupts UT passes 1..N-1. + self.supports_fold_ln = False + + self.weight_processing_conversions = { + **self._qkvo_weight_conversions(), + } + self.component_mapping = { + "embed": EmbeddingBridge(name="model.embed_tokens"), + "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"), + "blocks": BlockBridge( + name="model.layers", + config=self.cfg, + submodules={ + "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), + "ln1_post": RMSNormalizationBridge(name="input_layernorm_2", config=self.cfg), + "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg), + "ln2_post": RMSNormalizationBridge( + name="post_attention_layernorm_2", config=self.cfg + ), + "attn": PositionEmbeddingsAttentionBridge( + name="self_attn", + config=self.cfg, + submodules={ + "q": LinearBridge(name="q_proj"), + "k": LinearBridge(name="k_proj"), + "v": LinearBridge(name="v_proj"), + "o": LinearBridge(name="o_proj"), + }, + requires_attention_mask=True, + requires_position_embeddings=True, + ), + "mlp": GatedMLPBridge( + name="mlp", + config=self.cfg, + submodules={ + "gate": LinearBridge(name="gate_proj"), + "in": LinearBridge(name="up_proj"), + "out": LinearBridge(name="down_proj"), + }, + ), + }, + ), + "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), + "unembed": UnembeddingBridge(name="lm_head", config=self.cfg), + } + + def prepare_loading(self, model_name: str, model_kwargs: dict) -> None: + """Patch Ouro's remote code for compatibility with transformers v5. + + Ouro's modeling code was written against transformers 4.55, where + standard RoPE lived in ROPE_INIT_FUNCTIONS["default"]. Transformers v5 + removed that key and instead expects each *RotaryEmbedding class to + carry a compute_default_rope_parameters static method. Two call sites + break, so two patches: + + 1. OuroRotaryEmbedding.__init__ does ROPE_INIT_FUNCTIONS["default"] + (KeyError). Rebind the module-level name inside the imported + modeling_ouro module(s) to a copy with "default" restored; the + shared transformers dict is left untouched. + 2. v5's PreTrainedModel._init_weights re-initializes RotaryEmbedding + buffers via module.compute_default_rope_parameters(config) + (AttributeError). Attach the same function as a static method. + + Args: + model_name: The HuggingFace model name/path + model_kwargs: The kwargs dict for from_pretrained() + """ + # Force-import the modeling module so we can patch it + try: + from transformers.dynamic_module_utils import get_class_from_dynamic_module + + get_class_from_dynamic_module( + "modeling_ouro.OuroForCausalLM", + model_name, + ) + except Exception: + return + + # Each checkpoint revision gets its own module in sys.modules, so patch + # every imported Ouro modeling module (same idiom as openelm.py). + for key in list(sys.modules.keys()): + if "ouro" in key.lower() and "modeling" in key.lower(): + module = sys.modules[key] + rope_functions = getattr(module, "ROPE_INIT_FUNCTIONS", None) + if rope_functions is not None and "default" not in rope_functions: + setattr( + module, + "ROPE_INIT_FUNCTIONS", + {**rope_functions, "default": _compute_default_rope_parameters}, + ) + rope_class = getattr(module, "OuroRotaryEmbedding", None) + if rope_class is not None and not hasattr( + rope_class, "compute_default_rope_parameters" + ): + rope_class.compute_default_rope_parameters = staticmethod( + _compute_default_rope_parameters + ) + + def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None: + """Set up rotary embedding references for Ouro component testing. + + Ouro uses RoPE (Rotary Position Embeddings) with a single shared + ``model.rotary_emb``. We set the rotary_emb reference on all attention + bridge instances for component testing. + + Args: + hf_model: The HuggingFace Ouro model instance + bridge_model: The TransformerBridge model (if available, set rotary_emb on actual instances) + """ + # Get rotary embedding instance from the model + rotary_emb = hf_model.model.rotary_emb + + # Set rotary_emb on actual bridge instances in bridge_model if available + if bridge_model is not None and hasattr(bridge_model, "blocks"): + # Set on each layer's actual attention bridge instance + for block in bridge_model.blocks: + if hasattr(block, "attn"): + block.attn.set_rotary_emb(rotary_emb) + + # Also set on the template for get_generalized_component() calls + attn_bridge = self.get_generalized_component("blocks.0.attn") + attn_bridge.set_rotary_emb(rotary_emb) diff --git a/transformer_lens/tools/model_registry/__init__.py b/transformer_lens/tools/model_registry/__init__.py index 9a38c22b5..a304771ef 100644 --- a/transformer_lens/tools/model_registry/__init__.py +++ b/transformer_lens/tools/model_registry/__init__.py @@ -99,6 +99,7 @@ "OlmoForCausalLM", "OlmoeForCausalLM", "OPTForCausalLM", + "OuroForCausalLM", "PhiForCausalLM", "Phi3ForCausalLM", "PhiMoEForCausalLM", @@ -178,6 +179,7 @@ "OlmoForCausalLM": ["allenai"], "OpenELMForCausalLM": ["apple"], "OPTForCausalLM": ["facebook"], + "OuroForCausalLM": ["ByteDance"], "Phi3ForCausalLM": ["microsoft"], "PhiMoEForCausalLM": ["microsoft"], "PhiForCausalLM": ["microsoft"], diff --git a/transformer_lens/tools/model_registry/data/supported_models.json b/transformer_lens/tools/model_registry/data/supported_models.json index 399603b39..b42d76ff9 100644 --- a/transformer_lens/tools/model_registry/data/supported_models.json +++ b/transformer_lens/tools/model_registry/data/supported_models.json @@ -6,9 +6,9 @@ "min_downloads": 500, "scan_duration_seconds": 8.1 }, - "total_architectures": 69, - "total_models": 13144, - "total_verified": 1028, + "total_architectures": 70, + "total_models": 13145, + "total_verified": 1029, "models": [ { "architecture_id": "FalconH1ForCausalLM", @@ -156373,6 +156373,20 @@ "phase7_score": null, "phase8_score": null }, + { + "architecture_id": "OuroForCausalLM", + "model_id": "ByteDance/Ouro-1.4B", + "status": 1, + "verified_date": "2026-07-07", + "metadata": null, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": 95.9, + "phase7_score": null, + "phase8_score": null + }, { "architecture_id": "SmolLM3ForCausalLM", "model_id": "HuggingFaceTB/SmolLM3-3B", diff --git a/transformer_lens/tools/model_registry/data/verification_history.json b/transformer_lens/tools/model_registry/data/verification_history.json index 530f03e15..54eb14cfb 100644 --- a/transformer_lens/tools/model_registry/data/verification_history.json +++ b/transformer_lens/tools/model_registry/data/verification_history.json @@ -1,5 +1,5 @@ { - "last_updated": "2026-07-02T02:19:51.172022", + "last_updated": "2026-07-07T23:04:00.376842", "records": [ { "model_id": "Macropodus/macbert4mdcspell_v1", @@ -16860,6 +16860,16 @@ "notes": "Full verification completed", "invalidated": false, "invalidation_reason": null + }, + { + "model_id": "ByteDance/Ouro-1.4B", + "architecture_id": "OuroForCausalLM", + "verified_date": "2026-07-07", + "verified_by": "verify_models", + "transformerlens_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null } ] } diff --git a/transformer_lens/tools/model_registry/generate_report.py b/transformer_lens/tools/model_registry/generate_report.py index d62edda90..6f514dd6c 100644 --- a/transformer_lens/tools/model_registry/generate_report.py +++ b/transformer_lens/tools/model_registry/generate_report.py @@ -66,6 +66,7 @@ "BD3LM": "Kuleshov Group's Block Diffusion Language Model (ICLR 2025) for masked text generation", "HunYuanDenseV1ForCausalLM": "Tencent's open source decoder models", "Cohere2ForCausalLM": "Cohere's Command-A architecture with interleaved sliding-window RoPE and full-attention NoPE layers", + "OuroForCausalLM": "ByteDance's Ouro looped language model (LoopLM) with weight-shared iterated depth", # Unsupported architectures "BertModel": "Google's BERT bidirectional encoder for understanding tasks", "BertForMaskedLM": "BERT with masked language modeling head", diff --git a/transformer_lens/tools/model_registry/verify_models.py b/transformer_lens/tools/model_registry/verify_models.py index 407acc975..667daf4e6 100644 --- a/transformer_lens/tools/model_registry/verify_models.py +++ b/transformer_lens/tools/model_registry/verify_models.py @@ -62,6 +62,7 @@ # These are not in the legacy NEED_REMOTE_CODE_MODELS tuple (loading_from_pretrained.py). _BRIDGE_REMOTE_CODE_PREFIXES: tuple[str, ...] = ( "baichuan-inc/", # BaichuanForCausalLM — ships own modeling_baichuan.py + "ByteDance/Ouro-", # OuroForCausalLM — ships own modeling_ouro.py "internlm/", # InternLM2ForCausalLM — ships own modeling_internlm2.py "kuleshov-group/", # BD3LM — ships own custom modeling_d_dit.py ) From 9110cfd625701d8e906cd088f1814e3ec691efc1 Mon Sep 17 00:00:00 2001 From: "Katherine J. Ombrellaro" Date: Sat, 11 Jul 2026 04:59:17 +0000 Subject: [PATCH 11/11] Add weight converter for from-scratch pretrain.py foundation models Adds a converter and loader that bring checkpoints from an external decoder-only pretraining framework (RoPE with adjacent-pair rotation, RMSNorm, SwiGLU, optional dropless top-k MoE) into HookedTransformer so their residual streams can be read with the standard hook API. - convert_maritime_pretrain_weights: splits the fused QKV projection into per-head W_Q/W_K/W_V, maps SwiGLU gate/up/down onto W_gate/W_in/W_out, and emits router + per-expert weights under the MoE component's names. - maritime_pretrain_loader: builds the matching HookedTransformerConfig (rotary_adjacent_pairs=True, normalization_type=RMS, gated_mlp=True, norm_topk_prob=True for MoE), merges tensor-parallel shards, returns a ready HookedTransformer. - Tests assert bit-exact logit equivalence (max|delta| < 1e-4) for both dense and MoE configs plus converter key/shape correctness. --- .../test_maritime_pretrain.py | 299 ++++++++++++++++++ .../pretrained/maritime_pretrain_loader.py | 124 ++++++++ .../pretrained/weight_conversions/__init__.py | 7 +- .../weight_conversions/maritime_pretrain.py | 156 +++++++++ 4 files changed, 583 insertions(+), 3 deletions(-) create mode 100644 tests/unit/pretrained_weight_conversions/test_maritime_pretrain.py create mode 100644 transformer_lens/pretrained/maritime_pretrain_loader.py create mode 100644 transformer_lens/pretrained/weight_conversions/maritime_pretrain.py diff --git a/tests/unit/pretrained_weight_conversions/test_maritime_pretrain.py b/tests/unit/pretrained_weight_conversions/test_maritime_pretrain.py new file mode 100644 index 000000000..991eac08e --- /dev/null +++ b/tests/unit/pretrained_weight_conversions/test_maritime_pretrain.py @@ -0,0 +1,299 @@ +"""Tests for the maritime-pretrain (from-scratch foundation model) weight conversion. + +Coverage: + +* End-to-end numerical equivalence -- a source model and its converted + HookedTransformer produce the same logits -- across the axes a future change + is most likely to break: + - dense vs MoE MLPs, + - tied vs untied embeddings, + - several (n_heads, d_model) combinations. + This is the property downstream probing relies on: the residual stream read + through the hook API must be the model's real one. +* Converter key/shape correctness for the MoE branch. +* Friendly failure -- a config/checkpoint shape mismatch raises ValueError, not + an opaque reshape error. +* Tensor-parallel shard merge -- splitting a checkpoint into N shards the way + pretrain.py does and merging them back reconstructs the original weights and + still converts to an equivalent model, for N in {1, 2, 4}. + +The source model is a compact self-contained reimplementation of the +``pretrain.py`` architecture (RoPE with adjacent-pair rotation, RMSNorm, SwiGLU, +optional dropless top-k MoE), so the tests have no dependency on the external +training repo. +""" + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +from transformer_lens import HookedTransformer +from transformer_lens.pretrained.maritime_pretrain_loader import ( + _merge_tp_shards, + _shard_dim, + build_config, +) +from transformer_lens.pretrained.weight_conversions.maritime_pretrain import ( + convert_maritime_pretrain_weights, +) + +BASE_ARCH = dict( + d_model=64, + n_layers=2, + n_heads=4, + d_ff=192, + vocab_size=256, + max_seq_len=32, + rope_theta=10000.0, + rmsnorm_eps=1e-5, +) + + +# --- minimal reference model (mirrors pretrain.py exactly) -------------------- + + +def _rope_cache(seq_len, head_dim, theta): + inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim)) + freqs = torch.outer(torch.arange(seq_len).float(), inv) + return freqs.cos(), freqs.sin() + + +def _apply_rope(x, cos, sin): + x1, x2 = x[..., 0::2], x[..., 1::2] + T = x.shape[-2] + c, s = cos[:T].view(1, 1, T, -1), sin[:T].view(1, 1, T, -1) + out = torch.empty_like(x) + out[..., 0::2] = x1 * c - x2 * s + out[..., 1::2] = x1 * s + x2 * c + return out + + +class _RMSNorm(nn.Module): + def __init__(self, d, eps): + super().__init__() + self.weight = nn.Parameter(torch.ones(d)) + self.eps = eps + + def forward(self, x): + xf = x.float() + inv = torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) + return (xf * inv * self.weight.float()).to(x.dtype) + + +class _Attn(nn.Module): + def __init__(self, arch): + super().__init__() + d, h = arch["d_model"], arch["n_heads"] + self.h, self.hd = h, d // h + self.qkv = nn.Linear(d, 3 * d, bias=False) + self.proj = nn.Linear(d, d, bias=False) + + def forward(self, x, cos, sin): + B, T, _ = x.shape + qkv = self.qkv(x).view(B, T, 3, self.h, self.hd).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] + q, k = _apply_rope(q, cos, sin), _apply_rope(k, cos, sin) + o = F.scaled_dot_product_attention(q, k, v, is_causal=True) + return self.proj(o.transpose(1, 2).reshape(B, T, self.h * self.hd)) + + +class _SwiGLU(nn.Module): + def __init__(self, d, ff): + super().__init__() + self.gate = nn.Linear(d, ff, bias=False) + self.up = nn.Linear(d, ff, bias=False) + self.down = nn.Linear(ff, d, bias=False) + + def forward(self, x): + return self.down(F.silu(self.gate(x)) * self.up(x)) + + +class _MoE(nn.Module): + def __init__(self, arch): + super().__init__() + self.k = arch["top_k"] + self.router = nn.Linear(arch["d_model"], arch["n_experts"], bias=False) + self.experts = nn.ModuleList( + _SwiGLU(arch["d_model"], arch["d_ff"]) for _ in range(arch["n_experts"]) + ) + + def forward(self, x): + B, T, D = x.shape + flat = x.reshape(-1, D) + probs = self.router(flat.float()).softmax(-1) + top_p, top_i = probs.topk(self.k, dim=-1) + top_p = top_p / top_p.sum(-1, keepdim=True) + out = torch.zeros_like(flat) + for e in range(len(self.experts)): + rows, slot = (top_i == e).nonzero(as_tuple=True) + if rows.numel(): + y = self.experts[e](flat.index_select(0, rows)) + out.index_add_(0, rows, y * top_p[rows, slot].unsqueeze(-1).to(y.dtype)) + return out.reshape(B, T, D) + + +class _Block(nn.Module): + def __init__(self, arch, moe): + super().__init__() + self.norm1 = _RMSNorm(arch["d_model"], arch["rmsnorm_eps"]) + self.attn = _Attn(arch) + self.norm2 = _RMSNorm(arch["d_model"], arch["rmsnorm_eps"]) + self.mlp = _MoE(arch) if moe else _SwiGLU(arch["d_model"], arch["d_ff"]) + + def forward(self, x, cos, sin): + x = x + self.attn(self.norm1(x), cos, sin) + return x + self.mlp(self.norm2(x)) + + +class _GPT(nn.Module): + def __init__(self, arch, tie_embeddings=True): + super().__init__() + self.tie_embeddings = tie_embeddings + self.embed = nn.Embedding(arch["vocab_size"], arch["d_model"]) + self.blocks = nn.ModuleList( + _Block(arch, arch.get("moe", False)) for _ in range(arch["n_layers"]) + ) + self.norm_f = _RMSNorm(arch["d_model"], arch["rmsnorm_eps"]) + if not tie_embeddings: + self.lm_head = nn.Linear(arch["d_model"], arch["vocab_size"], bias=False) + cos, sin = _rope_cache( + arch["max_seq_len"], arch["d_model"] // arch["n_heads"], arch["rope_theta"] + ) + self.register_buffer("cos", cos) + self.register_buffer("sin", sin) + + def forward(self, toks): + x = self.embed(toks) + for blk in self.blocks: + x = blk(x, self.cos, self.sin) + x = self.norm_f(x) + w = self.embed.weight if self.tie_embeddings else self.lm_head.weight + return F.linear(x, w) + + def source_state_dict(self): + """State dict in pretrain.py's key layout: tied models omit lm_head.""" + sd = dict(self.state_dict()) + sd.pop("cos", None) + sd.pop("sin", None) + return sd + + +def _convert_to_tl(arch, src, dtype=torch.float32): + cfg = build_config(arch, dtype=dtype) + state_dict = convert_maritime_pretrain_weights(src.source_state_dict(), cfg) + tl = HookedTransformer(cfg) + tl.load_and_process_state_dict( + state_dict, + fold_ln=False, + center_writing_weights=False, + center_unembed=False, + fold_value_biases=False, + ) + return tl.eval() + + +def _max_logit_delta(arch, tie_embeddings=True): + torch.manual_seed(0) + src = _GPT(arch, tie_embeddings=tie_embeddings).eval() + tl = _convert_to_tl(arch, src) + toks = torch.randint(0, arch["vocab_size"], (2, 16)) + with torch.no_grad(): + ref = src(toks) + got = tl(toks) + return (ref - got).abs().max().item() + + +# --- equivalence across the axes most likely to regress ----------------------- + + +@pytest.mark.parametrize("tie_embeddings", [True, False], ids=["tied", "untied"]) +def test_dense_equivalence(tie_embeddings): + assert _max_logit_delta(dict(BASE_ARCH), tie_embeddings) < 1e-4 + + +@pytest.mark.parametrize("tie_embeddings", [True, False], ids=["tied", "untied"]) +def test_moe_equivalence(tie_embeddings): + arch = dict(BASE_ARCH, moe=True, n_experts=4, top_k=2, moe_every=1) + assert _max_logit_delta(arch, tie_embeddings) < 1e-4 + + +@pytest.mark.parametrize( + "d_model,n_heads", + [(64, 4), (64, 8), (128, 8), (96, 6)], + ids=["64d-4h", "64d-8h", "128d-8h", "96d-6h"], +) +def test_equivalence_varying_head_counts(d_model, n_heads): + arch = dict(BASE_ARCH, d_model=d_model, n_heads=n_heads, d_ff=3 * d_model) + assert _max_logit_delta(arch) < 1e-4 + + +# --- converter structure ------------------------------------------------------ + + +def test_moe_converter_emits_expected_keys(): + arch = dict(BASE_ARCH, moe=True, n_experts=4, top_k=2, moe_every=1) + torch.manual_seed(0) + src = _GPT(arch) + cfg = build_config(arch, dtype=torch.float32) + sd = convert_maritime_pretrain_weights(src.source_state_dict(), cfg) + assert "blocks.0.mlp.W_gate.weight" in sd + for e in range(arch["n_experts"]): + for w in ("W_gate", "W_in", "W_out"): + assert f"blocks.0.mlp.experts.{e}.{w}.weight" in sd + assert sd["blocks.0.attn.W_Q"].shape == ( + arch["n_heads"], + arch["d_model"], + arch["d_model"] // arch["n_heads"], + ) + + +# --- friendly failure on config/checkpoint mismatch --------------------------- + + +def test_shape_mismatch_raises_value_error(): + arch = dict(BASE_ARCH, n_layers=1) + cfg = build_config(arch, dtype=torch.float32) + torch.manual_seed(0) + src = _GPT(arch) + bad = dict(src.source_state_dict()) + # Corrupt one weight so it no longer matches the config the converter trusts. + bad["blocks.0.attn.qkv.weight"] = torch.randn(7, arch["d_model"]) + with pytest.raises(ValueError, match="expected shape"): + convert_maritime_pretrain_weights(bad, cfg) + + +# --- tensor-parallel shard merge regression ----------------------------------- + + +def _split_into_shards(state_dict, n_shards): + """Reproduce pretrain.py's tensor-parallel sharding: column-parallel weights + split along dim 0, row-parallel along dim 1, everything else replicated.""" + shards = [{} for _ in range(n_shards)] + for key, w in state_dict.items(): + dim = _shard_dim(key) + if dim is None: + for s in shards: + s[key] = w.clone() + else: + for s, piece in zip(shards, w.chunk(n_shards, dim=dim)): + s[key] = piece.clone() + return shards + + +@pytest.mark.parametrize("n_shards", [1, 2, 4]) +def test_tp_shard_merge_reconstructs_weights(tmp_path, n_shards): + arch = dict(BASE_ARCH, d_model=64, n_heads=8, d_ff=256) # divisible by 4 + torch.manual_seed(0) + src = _GPT(arch) + full = src.source_state_dict() + + run = tmp_path / "run" + (run / "best").mkdir(parents=True) + for rank, shard in enumerate(_split_into_shards(full, n_shards)): + torch.save({"model": shard}, run / "best" / f"model_tp{rank}.pt") + + merged = _merge_tp_shards(run, "best") + assert set(merged) == set(full) + for key in full: + assert torch.equal(merged[key], full[key]), key diff --git a/transformer_lens/pretrained/maritime_pretrain_loader.py b/transformer_lens/pretrained/maritime_pretrain_loader.py new file mode 100644 index 000000000..835840451 --- /dev/null +++ b/transformer_lens/pretrained/maritime_pretrain_loader.py @@ -0,0 +1,124 @@ +"""Load a from-scratch ``pretrain.py`` checkpoint as a HookedTransformer. + +``pretrain.py`` (the maritime-intent-probe foundation-model framework) writes +one weight shard per tensor-parallel rank plus a ``config.json`` describing the +architecture. This helper reassembles the (optionally TP-sharded) weights into a +single state dict, builds the matching :class:`HookedTransformerConfig`, and +returns a ready ``HookedTransformer`` whose residual stream can be probed with +the standard hook API. + +Example +------- +>>> from transformer_lens.pretrained.maritime_pretrain_loader import ( +... load_maritime_pretrain, +... ) +>>> model = load_maritime_pretrain("runs/base", tag="best") +>>> logits, cache = model.run_with_cache(tokens) +>>> resid = cache["resid_post", 6] # hand off to the linear probes in probe.py +""" + +from __future__ import annotations + +import json +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import torch + +from transformer_lens import HookedTransformer +from transformer_lens.config.hooked_transformer_config import HookedTransformerConfig +from transformer_lens.pretrained.weight_conversions.maritime_pretrain import ( + convert_maritime_pretrain_weights, +) + + +def _shard_dim(key: str) -> int | None: + """Which axis a pretrain.py tensor-parallel weight was split along -- or + None if the weight is replicated on every rank (router, norms, embedding, + lm_head), in which case any shard's copy is authoritative.""" + if key.endswith(("attn.qkv.weight", "gate.weight", "up.weight")): + return 0 # column-parallel: split along output features + if key.endswith(("attn.proj.weight", "down.weight")): + return 1 # row-parallel: split along input features + return None + + +def _merge_tp_shards(run_dir: Path, tag: str) -> dict[str, torch.Tensor]: + """Reassemble the full weights from per-rank shards.""" + shards = sorted(run_dir.joinpath(tag).glob("model_tp*.pt")) + if not shards: + raise FileNotFoundError(f"no model_tp*.pt under {run_dir / tag}") + states = [torch.load(s, map_location="cpu", weights_only=False)["model"] for s in shards] + if len(states) == 1: + return states[0] + return { + key: ( + states[0][key] + if (dim := _shard_dim(key)) is None + else torch.cat([st[key] for st in states], dim=dim) + ) + for key in states[0] + } + + +def build_config( + arch: Mapping[str, Any], dtype: torch.dtype = torch.float32 +) -> HookedTransformerConfig: + """Translate pretrain.py's ArchConfig dict into a HookedTransformerConfig.""" + d_model, n_heads = arch["d_model"], arch["n_heads"] + is_moe = arch.get("moe", False) + return HookedTransformerConfig( + n_layers=arch["n_layers"], + d_model=d_model, + n_ctx=arch["max_seq_len"], + d_head=d_model // n_heads, + n_heads=n_heads, + d_mlp=arch["d_ff"], + d_vocab=arch["vocab_size"], + act_fn="silu", + normalization_type="RMS", # RMSNorm, no bias, no mean-subtraction + eps=arch.get("rmsnorm_eps", 1e-5), + gated_mlp=True, # SwiGLU + positional_embedding_type="rotary", + rotary_dim=d_model // n_heads, + rotary_base=arch.get("rope_theta", 10000.0), + rotary_adjacent_pairs=True, # pretrain.py rotates adjacent pairs + final_rms=True, + num_experts=(arch["n_experts"] if is_moe else None), + experts_per_token=(arch["top_k"] if is_moe else None), + # pretrain.py renormalises the top-k routing weights (top_p / top_p.sum) + **({"norm_topk_prob": True} if is_moe else {}), + dtype=dtype, + ) + + +def load_maritime_pretrain( + run_dir: str | Path, + tag: str = "best", + device: str = "cpu", + dtype: torch.dtype = torch.float32, + fold_ln: bool = False, +) -> HookedTransformer: + """Return a HookedTransformer with the pretrain.py checkpoint loaded. + + ``fold_ln`` defaults to False because RMSNorm folding changes the residual + stream basis probes may depend on; enable it only if you specifically want + the folded form. + """ + run_dir = Path(run_dir) + arch = json.loads((run_dir / "config.json").read_text()) + cfg = build_config(arch, dtype=dtype) + + raw = _merge_tp_shards(run_dir, tag) + state_dict = convert_maritime_pretrain_weights(raw, cfg) + + model = HookedTransformer(cfg) + model.load_and_process_state_dict( + state_dict, + fold_ln=fold_ln, + center_writing_weights=False, # RMSNorm has no bias to center + center_unembed=False, + fold_value_biases=False, + ) + return model.to(device) diff --git a/transformer_lens/pretrained/weight_conversions/__init__.py b/transformer_lens/pretrained/weight_conversions/__init__.py index f2df6c917..202d6dbc8 100644 --- a/transformer_lens/pretrained/weight_conversions/__init__.py +++ b/transformer_lens/pretrained/weight_conversions/__init__.py @@ -1,10 +1,13 @@ +from .apertus import convert_apertus_weights from .bert import convert_bert_weights from .bloom import convert_bloom_weights from .coder import convert_coder_weights from .gemma import convert_gemma_weights from .gpt2 import convert_gpt2_weights from .gptj import convert_gptj_weights +from .hubert import convert_hubert_weights from .llama import convert_llama_weights +from .maritime_pretrain import convert_maritime_pretrain_weights from .mingpt import convert_mingpt_weights from .mistral import convert_mistral_weights from .mixtral import convert_mixtral_weights @@ -16,6 +19,7 @@ from .olmo2 import convert_olmo2_weights from .olmo3 import convert_olmo3_weights from .olmoe import convert_olmoe_weights +from .openai import convert_gpt_oss_weights from .opt import convert_opt_weights from .phi import convert_phi_weights from .phi3 import convert_phi3_weights @@ -23,6 +27,3 @@ from .qwen2 import convert_qwen2_weights from .qwen3 import convert_qwen3_weights from .t5 import convert_t5_weights -from .hubert import convert_hubert_weights -from .apertus import convert_apertus_weights -from .openai import convert_gpt_oss_weights diff --git a/transformer_lens/pretrained/weight_conversions/maritime_pretrain.py b/transformer_lens/pretrained/weight_conversions/maritime_pretrain.py new file mode 100644 index 000000000..fb43bcced --- /dev/null +++ b/transformer_lens/pretrained/weight_conversions/maritime_pretrain.py @@ -0,0 +1,156 @@ +from collections.abc import Mapping + +import einops +import torch + +from transformer_lens.config.hooked_transformer_config import HookedTransformerConfig + + +def _identity(tensor: torch.Tensor) -> torch.Tensor: + return tensor + + +def _check(tensor: torch.Tensor, expected: tuple[int, ...], name: str) -> torch.Tensor: + """Fail early and legibly when a source weight is not the shape the mapping + assumes. A clear ValueError here beats an opaque einops/reshape error three + frames deep, and turns a config/checkpoint mismatch into an actionable one.""" + if tuple(tensor.shape) != expected: + raise ValueError( + f"{name}: expected shape {expected}, got {tuple(tensor.shape)}. " + "This usually means the HookedTransformerConfig does not match the " + "checkpoint (wrong d_model / n_heads / d_mlp / n_experts)." + ) + return tensor + + +def convert_maritime_pretrain_weights( + old_state_dict: Mapping[str, torch.Tensor], + cfg: HookedTransformerConfig, +) -> dict[str, torch.Tensor]: + """Convert a checkpoint from the ``pretrain.py`` foundation-model framework + (https://github.com/kombrellaro/maritime-intent-probe) into TransformerLens + format, so its residual stream can be read with the standard hook API. + + The source is a decoder-only transformer with RoPE, RMSNorm, SwiGLU MLPs and + optional dropless top-k MoE. Three conventions must line up exactly; each is + named below so the mapping reads as a specification rather than a pile of + index gymnastics: + + * RoPE rotates *adjacent* dimension pairs ([x0, x1] -> [-x1, x0]), so the + matching config sets ``rotary_adjacent_pairs=True`` and applies rotation + inside attention. Q/K/V therefore need no reordering here. + * Attention fuses Q/K/V into one column-parallel projection stored as + ``blocks.{l}.attn.qkv.weight`` of shape [3*d_model, d_model]; it is split + into three per-head projections via :func:`split_heads`. + * SwiGLU maps onto gated-MLP naming as gate -> W_gate, up -> W_in, + down -> W_out. The MoE experts store these as ``nn.Linear`` ([out, in], + kept as-is under a ``.weight`` key); the dense gated MLP stores raw + [in, out] parameters (transposed, no suffix). Same mapping, two storage + conventions -- and that single fact is the whole dense/MoE asymmetry, + captured once in :func:`emit_gated_mlp`. + + Raises: + ValueError: if a source weight is not the shape ``cfg`` implies, so a + config/checkpoint mismatch fails with a clear message rather than an + opaque reshape error. + """ + n_heads, d_head, d_model = cfg.n_heads, cfg.d_head, cfg.d_model + d_mlp = cfg.d_mlp + + def split_heads(w: torch.Tensor) -> torch.Tensor: + # [feat, d_model] input projection -> [n_heads, d_model, d_head] + return einops.rearrange(w, "(h dh) m -> h m dh", h=n_heads, dh=d_head) + + def merge_heads(w: torch.Tensor) -> torch.Tensor: + # [d_model, feat] output projection -> [n_heads, d_head, d_model] + return einops.rearrange(w, "m (h dh) -> h dh m", h=n_heads, dh=d_head) + + def zeros(*shape: int) -> torch.Tensor: + return torch.zeros(*shape, dtype=cfg.dtype) + + def src(block: int, key: str) -> torch.Tensor: + # a block's source weight, addressed by suffix + return old_state_dict[f"blocks.{block}.{key}"] + + def emit_gated_mlp( + prefix: str, + gate: torch.Tensor, + up: torch.Tensor, + down: torch.Tensor, + *, + as_linear: bool, + ) -> dict[str, torch.Tensor]: + """gate/up/down -> W_gate/W_in/W_out under one of two storage styles: + nn.Linear weights ([out, in], ``.weight`` suffix) are kept as-is; raw + parameters ([in, out], no suffix) are transposed.""" + _check(gate, (d_mlp, d_model), f"{prefix}.W_gate") + _check(up, (d_mlp, d_model), f"{prefix}.W_in") + _check(down, (d_model, d_mlp), f"{prefix}.W_out") + suffix = ".weight" if as_linear else "" + orient = _identity if as_linear else torch.t + return { + f"{prefix}.W_gate{suffix}": orient(gate), + f"{prefix}.W_in{suffix}": orient(up), + f"{prefix}.W_out{suffix}": orient(down), + } + + tied_unembed = "lm_head.weight" not in old_state_dict + new_state_dict: dict[str, torch.Tensor] = { + "embed.W_E": old_state_dict["embed.weight"], + "ln_final.w": old_state_dict["norm_f.weight"], + "unembed.W_U": old_state_dict["embed.weight" if tied_unembed else "lm_head.weight"].T, + "unembed.b_U": zeros(cfg.d_vocab), + } + + for layer in range(cfg.n_layers): + p = f"blocks.{layer}" + + new_state_dict[f"{p}.ln1.w"] = src(layer, "norm1.weight") + new_state_dict[f"{p}.ln2.w"] = src(layer, "norm2.weight") + + # Attention: one fused QKV projection -> three per-head Q/K/V, plus out. + qkv = _check(src(layer, "attn.qkv.weight"), (3 * d_model, d_model), f"{p}.attn.qkv") + for name, w in zip("QKV", qkv.chunk(3, dim=0)): + new_state_dict[f"{p}.attn.W_{name}"] = split_heads(w) + new_state_dict[f"{p}.attn.b_{name}"] = zeros(n_heads, d_head) + w_o = _check(src(layer, "attn.proj.weight"), (d_model, d_model), f"{p}.attn.proj") + new_state_dict[f"{p}.attn.W_O"] = merge_heads(w_o) + new_state_dict[f"{p}.attn.b_O"] = zeros(d_model) + + # MLP: a dense SwiGLU, or a router plus a fleet of expert SwiGLUs. + # Both conditions are load-bearing: num_experts guards the *config* side + # (a dense config leaves it None and must take the dense branch even if a + # stray router key existed), and the key check guards the *checkpoint* + # side (an MoE config whose every-k-th-layer placement makes this + # particular block dense -- moe_every > 1 -- has no router here). + is_moe = cfg.num_experts is not None and f"{p}.mlp.router.weight" in old_state_dict + if is_moe: + router = _check( + src(layer, "mlp.router.weight"), (cfg.num_experts, d_model), f"{p}.mlp.router" + ) + new_state_dict[f"{p}.mlp.W_gate.weight"] = router + for e in range(cfg.num_experts): + ep = f"mlp.experts.{e}" + new_state_dict.update( + emit_gated_mlp( + f"{p}.{ep}", + src(layer, f"{ep}.gate.weight"), + src(layer, f"{ep}.up.weight"), + src(layer, f"{ep}.down.weight"), + as_linear=True, + ) + ) + else: + new_state_dict.update( + emit_gated_mlp( + f"{p}.mlp", + src(layer, "mlp.gate.weight"), + src(layer, "mlp.up.weight"), + src(layer, "mlp.down.weight"), + as_linear=False, + ) + ) + new_state_dict[f"{p}.mlp.b_in"] = zeros(d_mlp) + new_state_dict[f"{p}.mlp.b_out"] = zeros(d_model) + + return new_state_dict