From 4d3c535b638b631a5ad20d6ee4e32336a3dfc922 Mon Sep 17 00:00:00 2001 From: hsusul Date: Sun, 12 Jul 2026 18:42:47 -0500 Subject: [PATCH] Add forward-only LLaDA architecture adapter --- .../model_bridge/test_llada_adapter.py | 841 ++++++++++++++++++ .../test_llada_adapter.py | 324 +++++++ .../model_bridge/test_architecture_adapter.py | 17 + .../unit/model_bridge/test_component_setup.py | 28 + .../unit/model_bridge/test_native_features.py | 67 +- .../factories/architecture_adapter_factory.py | 2 + .../model_bridge/architecture_adapter.py | 26 +- transformer_lens/model_bridge/bridge.py | 42 +- .../model_bridge/component_setup.py | 10 +- .../generalized_components/attention.py | 6 +- .../joint_qkv_attention.py | 2 + .../position_embeddings_attention.py | 2 + .../model_bridge/sources/_bridge_builder.py | 15 + .../model_bridge/sources/transformers.py | 44 +- .../supported_architectures/__init__.py | 4 + .../supported_architectures/llada.py | 371 ++++++++ .../tools/model_registry/__init__.py | 2 + .../model_registry/data/supported_models.json | 18 +- .../tools/model_registry/generate_report.py | 1 + .../tools/model_registry/verify_models.py | 1 + 20 files changed, 1795 insertions(+), 28 deletions(-) create mode 100644 tests/integration/model_bridge/test_llada_adapter.py create mode 100644 tests/unit/model_bridge/supported_architectures/test_llada_adapter.py create mode 100644 transformer_lens/model_bridge/supported_architectures/llada.py diff --git a/tests/integration/model_bridge/test_llada_adapter.py b/tests/integration/model_bridge/test_llada_adapter.py new file mode 100644 index 0000000000..972da19a5e --- /dev/null +++ b/tests/integration/model_bridge/test_llada_adapter.py @@ -0,0 +1,841 @@ +"""Offline integration tests for one dense LLaDA transformer forward pass.""" + +from __future__ import annotations + +import copy +import gc +import math +import weakref +from dataclasses import dataclass, field +from typing import NamedTuple, Optional, cast +from unittest.mock import patch + +import pytest +import torch +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace +from torch import nn +from transformers import PreTrainedTokenizerFast + +from transformer_lens.model_bridge import TransformerBridge +from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_from_module, +) + + +@dataclass +class TinyLLaDAConfig: + """Small structural copy of the released dense LLaDA config contract.""" + + activation_type: str = "silu" + alibi: bool = False + architectures: list[str] = field(default_factory=lambda: ["LLaDAModelLM"]) + attention_dropout: float = 0.0 + attention_layer_norm: bool = False + block_group_size: int = 1 + block_type: str = "llama" + d_model: int = 32 + embedding_dropout: float = 0.0 + embedding_size: int = 64 + eos_token_id: int = 1 + include_bias: bool = False + include_qkv_bias: bool = False + input_emb_norm: bool = False + layer_norm_type: str = "rms" + mask_token_id: int = 63 + max_sequence_length: int = 16 + mlp_hidden_size: int = 48 + model_type: str = "llada" + n_heads: int = 4 + n_kv_heads: int = 2 + n_layers: int = 2 + output_attentions: bool = True + pad_token_id: int = 0 + residual_dropout: float = 0.0 + rms_norm_eps: float = 1e-5 + rope: bool = True + rope_full_precision: bool = True + rope_theta: float = 500_000.0 + scale_logits: bool = False + use_cache: bool = False + vocab_size: int = 64 + weight_tying: bool = False + + @property + def effective_n_kv_heads(self) -> int: + return self.n_kv_heads + + +class TinyLLaDAOutput(NamedTuple): + logits: torch.Tensor + + +class TinyRMSNorm(nn.Module): + def __init__(self, config: TinyLLaDAConfig) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(config.d_model)) + self.eps = config.rms_norm_eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + original_dtype = x.dtype + normalized = x.float() + variance = normalized.pow(2).mean(dim=-1, keepdim=True) + normalized = normalized * torch.rsqrt(variance + self.eps) + return self.weight * normalized.to(original_dtype) + + +class TinyRotaryEmbedding(nn.Module): + """LLaDA's half-vector RoPE pairing, computed in fp32.""" + + def __init__(self, config: TinyLLaDAConfig) -> None: + super().__init__() + self.config = config + + @staticmethod + def _rotate_half(x: torch.Tensor) -> torch.Tensor: + batch, heads, positions, head_dim = x.shape + paired = x.view(batch, heads, positions, 2, head_dim // 2) + first, second = paired.unbind(dim=-2) + return torch.cat((-second, first), dim=-1) + + def _apply_rotary( + self, tensor: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor + ) -> torch.Tensor: + return (tensor * cos + self._rotate_half(tensor) * sin).to(tensor.dtype) + + def forward(self, query: torch.Tensor, key: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + original_query_dtype = query.dtype + original_key_dtype = key.dtype + if self.config.rope_full_precision: + query = query.float() + key = key.float() + + head_dim = query.shape[-1] + key_len = key.shape[-2] + query_len = query.shape[-2] + inverse_frequency = 1.0 / ( + self.config.rope_theta + ** (torch.arange(0, head_dim, 2, device=query.device).float() / head_dim) + ) + frequency = torch.outer( + torch.arange(key_len, device=query.device).float(), inverse_frequency + ) + position = torch.cat((frequency, frequency), dim=-1) + sin = position.sin()[None, None, :, :] + cos = position.cos()[None, None, :, :] + query = self._apply_rotary( + query, sin[:, :, key_len - query_len :], cos[:, :, key_len - query_len :] + ) + key = self._apply_rotary(key, sin, cos) + return query.to(original_query_dtype), key.to(original_key_dtype) + + +class TinyLLaDALlamaBlock(nn.Module): + """Faithful tiny copy of the remote dense ``LLaDALlamaBlock`` math.""" + + def __init__(self, config: TinyLLaDAConfig) -> None: + super().__init__() + self.config = config + self.dropout = nn.Dropout(config.residual_dropout) + self.act = nn.SiLU() + self.rotary_emb = TinyRotaryEmbedding(config) + self.attn_norm = TinyRMSNorm(config) + self.ff_norm = TinyRMSNorm(config) + + head_dim = config.d_model // config.n_heads + kv_width = config.effective_n_kv_heads * head_dim + self.q_proj = nn.Linear(config.d_model, config.d_model, bias=False) + self.k_proj = nn.Linear(config.d_model, kv_width, bias=False) + self.v_proj = nn.Linear(config.d_model, kv_width, bias=False) + self.attn_out = nn.Linear(config.d_model, config.d_model, bias=False) + self.ff_proj = nn.Linear(config.d_model, config.mlp_hidden_size, bias=False) + self.up_proj = nn.Linear(config.d_model, config.mlp_hidden_size, bias=False) + self.ff_out = nn.Linear(config.mlp_hidden_size, config.d_model, bias=False) + + self.last_attn_norm: Optional[torch.Tensor] = None + self.last_q: Optional[torch.Tensor] = None + self.last_k: Optional[torch.Tensor] = None + self.last_v: Optional[torch.Tensor] = None + self.last_rot_q: Optional[torch.Tensor] = None + self.last_rot_k: Optional[torch.Tensor] = None + self.last_scores: Optional[torch.Tensor] = None + self.last_pattern: Optional[torch.Tensor] = None + self.last_attn_output: Optional[torch.Tensor] = None + self.last_resid_mid: Optional[torch.Tensor] = None + self.last_mlp_norm: Optional[torch.Tensor] = None + self.last_gate: Optional[torch.Tensor] = None + self.last_up: Optional[torch.Tensor] = None + self.last_activation: Optional[torch.Tensor] = None + self.last_product: Optional[torch.Tensor] = None + self.last_mlp_output: Optional[torch.Tensor] = None + self.last_output: Optional[torch.Tensor] = None + + @staticmethod + def _cast_attn_bias(bias: torch.Tensor, input_dtype: torch.dtype) -> torch.Tensor: + return bias.to(dtype=input_dtype) + + def attention( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attention_bias: Optional[torch.Tensor] = None, + layer_past: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + use_cache: bool = False, + ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]: + batch, query_len, q_width = q.shape + head_dim = q_width // self.config.n_heads + q = q.view(batch, query_len, self.config.n_heads, head_dim).transpose(1, 2) + k = k.view(batch, query_len, self.config.effective_n_kv_heads, head_dim).transpose(1, 2) + v = v.view(batch, query_len, self.config.effective_n_kv_heads, head_dim).transpose(1, 2) + + self.last_q = q + self.last_k = k + self.last_v = v + + if layer_past is not None: + past_key, past_value = layer_past + k = torch.cat((past_key, k), dim=-2) + v = torch.cat((past_value, v), dim=-2) + present = (k, v) if use_cache else None + + q, k = self.rotary_emb(q, k) + self.last_rot_q = q + self.last_rot_k = k + key_len = k.shape[-2] + if attention_bias is not None: + attention_bias = self._cast_attn_bias( + attention_bias[:, :, key_len - query_len : key_len, :key_len], + k.dtype, + ) + + groups = self.config.n_heads // self.config.effective_n_kv_heads + k = k.repeat_interleave(groups, dim=1, output_size=self.config.n_heads) + v = v.repeat_interleave(groups, dim=1, output_size=self.config.n_heads) + scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(head_dim) + if attention_bias is not None: + scores = scores + attention_bias + self.last_scores = scores + pattern = torch.nn.functional.softmax(scores, dim=-1, dtype=torch.float32).to(q.dtype) + z = torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attention_bias, + dropout_p=self.config.attention_dropout if self.training else 0.0, + is_causal=False, + ) + z = z.transpose(1, 2).contiguous() + projected = self.attn_out(z.view(batch, query_len, self.config.d_model)) + + self.last_pattern = pattern + self.last_attn_output = projected + return projected, present + + def forward( + self, + x: torch.Tensor, + attention_bias: Optional[torch.Tensor] = None, + layer_past: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + use_cache: bool = False, + ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]: + normalized = self.attn_norm(x) + self.last_attn_norm = normalized + q = self.q_proj(normalized) + k = self.k_proj(normalized) + v = self.v_proj(normalized) + attention_output, cache = self.attention( + q, + k, + v, + attention_bias, + layer_past=layer_past, + use_cache=use_cache, + ) + resid_mid = x + self.dropout(attention_output) + + mlp_norm = self.ff_norm(resid_mid) + gate = self.ff_proj(mlp_norm) + up = self.up_proj(mlp_norm) + activation = self.act(gate) + product = activation * up + mlp_output = self.dropout(self.ff_out(product)) + output = resid_mid + mlp_output + + self.last_resid_mid = resid_mid + self.last_mlp_norm = mlp_norm + self.last_gate = gate + self.last_up = up + self.last_activation = activation + self.last_product = product + self.last_mlp_output = mlp_output + self.last_output = output + return output, cache + + +class TinyLLaDAModel(nn.Module): + def __init__(self, config: TinyLLaDAConfig) -> None: + super().__init__() + self.config = config + self.transformer = nn.ModuleDict( + { + "wte": nn.Embedding(config.embedding_size, config.d_model), + "emb_drop": nn.Dropout(config.embedding_dropout), + "blocks": nn.ModuleList( + [TinyLLaDALlamaBlock(config) for _ in range(config.n_layers)] + ), + "ln_f": TinyRMSNorm(config), + "ff_out": nn.Linear(config.d_model, config.embedding_size, bias=False), + } + ) + self.last_final_norm: Optional[torch.Tensor] = None + + @staticmethod + def _padding_bias( + attention_mask: Optional[torch.Tensor], + batch_size: int, + ) -> Optional[torch.Tensor]: + if attention_mask is None or not bool((attention_mask == 0).any()): + return None + mask = attention_mask.float().view(batch_size, -1)[:, None, None, :] + return (1.0 - mask) * torch.finfo(mask.dtype).min + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + attention_bias: Optional[torch.Tensor] = None, + use_cache: bool = False, + output_hidden_states: Optional[bool] = None, + ) -> TinyLLaDAOutput: + del output_hidden_states + if use_cache: + raise AssertionError("The KV cache is not supported for LLaDA") + + batch_size, seq_len = input_ids.shape + x = self.transformer["wte"](input_ids) + x = self.transformer["emb_drop"](x) + padding_bias = self._padding_bias(attention_mask, batch_size) + if attention_bias is None and padding_bias is not None: + attention_bias = torch.zeros(1, 1, seq_len, seq_len, device=x.device) + if attention_bias is not None and attention_bias.dtype in (torch.bool, torch.int8): + allowed = attention_bias.bool() + attention_bias = torch.where( + allowed, + torch.zeros((), device=x.device), + torch.full((), torch.finfo(torch.float32).min, device=x.device), + ) + if attention_bias is not None and padding_bias is not None: + attention_bias = attention_bias.float() + padding_bias + + blocks = cast(nn.ModuleList, self.transformer["blocks"]) + for block_module in blocks: + x, _ = block_module(x, attention_bias=attention_bias, use_cache=False) + x = self.transformer["ln_f"](x) + self.last_final_norm = x + logits = self.transformer["ff_out"](x) + return TinyLLaDAOutput(logits=logits) + + +class TinyLLaDAModelLM(nn.Module): + def __init__(self, config: TinyLLaDAConfig) -> None: + super().__init__() + self.config = config + self.model = TinyLLaDAModel(config) + self.output_attentions_requests: list[Optional[bool]] = [] + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + attention_bias: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + ) -> TinyLLaDAOutput: + self.output_attentions_requests.append(output_attentions) + if output_attentions: + raise ValueError("output_attentions is not yet supported in LLaDA") + if use_cache is None: + use_cache = self.config.use_cache + return self.model( + input_ids, + attention_mask=attention_mask, + attention_bias=attention_bias, + use_cache=use_cache, + output_hidden_states=output_hidden_states, + ) + + +class TinyModels(NamedTuple): + bridge: TransformerBridge + reference: TinyLLaDAModelLM + + +def _build_models( + config: Optional[TinyLLaDAConfig] = None, + *, + dtype: torch.dtype = torch.float32, + eval_bridge: bool = True, +) -> TinyModels: + torch.manual_seed(0) + config = config or TinyLLaDAConfig() + reference = TinyLLaDAModelLM(config).to(dtype=dtype).eval() + candidate = copy.deepcopy(reference).eval() + bridge = build_bridge_from_module( + candidate, + architecture="LLaDAModelLM", + hf_config=copy.deepcopy(candidate.config), + tokenizer=None, + device="cpu", + dtype=dtype, + model_name="offline-tiny-llada", + ) + if eval_bridge: + bridge.eval() + return TinyModels(bridge=bridge, reference=reference) + + +@pytest.fixture(scope="module") +def models() -> TinyModels: + return _build_models() + + +def _offline_tokenizer() -> PreTrainedTokenizerFast: + vocabulary = { + "": 0, + "": 1, + "": 2, + "": 3, + "": 63, + } + vocabulary.update( + {f"token_{index}": index for index in range(64) if index not in {0, 1, 2, 3, 63}} + ) + backend = Tokenizer(WordLevel(vocabulary, unk_token="")) + backend.pre_tokenizer = Whitespace() + return PreTrainedTokenizerFast( + tokenizer_object=backend, + eos_token="", + pad_token="", + unk_token="", + mask_token="", + ) + + +def _reference_block(reference: TinyLLaDAModelLM, index: int = 0) -> TinyLLaDALlamaBlock: + blocks = cast(nn.ModuleList, reference.model.transformer["blocks"]) + block = blocks[index] + assert isinstance(block, TinyLLaDALlamaBlock) + return block + + +def _require_tensor(value: Optional[torch.Tensor]) -> torch.Tensor: + assert value is not None + return value + + +@pytest.mark.parametrize("trust_remote_code", [False, True]) +def test_remote_code_and_revision_reach_all_hf_loaders(trust_remote_code: bool) -> None: + import transformer_lens.model_bridge.sources.transformers as bridge_source + + config = TinyLLaDAConfig() + model = TinyLLaDAModelLM(TinyLLaDAConfig()).eval() + tokenizer = _offline_tokenizer() + with ( + patch.object( + bridge_source.AutoConfig, + "from_pretrained", + return_value=copy.deepcopy(config), + ) as config_load, + patch.object( + bridge_source.AutoModelForCausalLM, + "from_pretrained", + return_value=model, + ) as model_load, + patch.object( + bridge_source.AutoTokenizer, + "from_pretrained", + return_value=tokenizer, + ) as tokenizer_load, + ): + bridge = TransformerBridge.boot_transformers( + "offline-tiny-llada", + device="cpu", + trust_remote_code=trust_remote_code, + revision="reviewed-revision", + ) + + assert bridge.cfg.trust_remote_code is trust_remote_code + config_kwargs = config_load.call_args.kwargs + model_kwargs = model_load.call_args.kwargs + tokenizer_kwargs = tokenizer_load.call_args.kwargs + assert config_kwargs["trust_remote_code"] is trust_remote_code + assert config_kwargs["revision"] == "reviewed-revision" + assert model_kwargs.get("trust_remote_code", False) is trust_remote_code + assert model_kwargs["revision"] == "reviewed-revision" + assert tokenizer_kwargs["trust_remote_code"] is trust_remote_code + assert tokenizer_kwargs["revision"] == "reviewed-revision" + + +def test_masked_forward_matches_reference_logits_and_intermediates( + models: TinyModels, +) -> None: + tokens = torch.tensor([[5, 63, 7, 9]]) + with torch.inference_mode(): + reference_logits = models.reference(tokens).logits + bridge_logits, cache = models.bridge.run_with_cache(tokens) + + reference_block = _reference_block(models.reference) + torch.testing.assert_close(bridge_logits, reference_logits, rtol=1e-5, atol=1e-6) + attn_norm = _require_tensor(reference_block.last_attn_norm) + torch.testing.assert_close(cache["blocks.0.ln1.hook_out"], attn_norm, rtol=1e-5, atol=1e-6) + torch.testing.assert_close(cache["blocks.0.attn.hook_in"], attn_norm, rtol=1e-5, atol=1e-6) + + reference_q = _require_tensor(reference_block.last_q) + reference_k = _require_tensor(reference_block.last_k) + reference_v = _require_tensor(reference_block.last_v) + torch.testing.assert_close( + cache["blocks.0.attn.q.hook_out"], + reference_q.transpose(1, 2), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + cache["blocks.0.attn.k.hook_out"], + reference_k.transpose(1, 2), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + cache["blocks.0.attn.v.hook_out"], + reference_v.transpose(1, 2), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + cache["blocks.0.attn.hook_rot_q"], + _require_tensor(reference_block.last_rot_q).transpose(1, 2), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + cache["blocks.0.attn.hook_rot_k"], + _require_tensor(reference_block.last_rot_k).transpose(1, 2), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + cache["blocks.0.attn.hook_attn_scores"], + _require_tensor(reference_block.last_scores), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + cache["blocks.0.hook_out"], + _require_tensor(reference_block.last_output), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + cache["blocks.0.ln2.hook_in"], + _require_tensor(reference_block.last_resid_mid), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + cache["blocks.0.ln2.hook_out"], + _require_tensor(reference_block.last_mlp_norm), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + cache["blocks.0.mlp.hook_in"], + _require_tensor(reference_block.last_mlp_norm), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + cache["blocks.0.attn.hook_pattern"], + _require_tensor(reference_block.last_pattern), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + cache["blocks.0.attn.hook_out"], + _require_tensor(reference_block.last_attn_output), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + cache["blocks.0.mlp.gate.hook_out"], + _require_tensor(reference_block.last_gate), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + cache["blocks.0.mlp.in.hook_out"], + _require_tensor(reference_block.last_up), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + cache["blocks.0.mlp.act.hook_out"], + _require_tensor(reference_block.last_activation), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + cache["blocks.0.mlp.out.hook_in"], + _require_tensor(reference_block.last_product), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + cache["blocks.0.mlp.hook_out"], + _require_tensor(reference_block.last_mlp_output), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close( + cache["ln_final.hook_out"], + _require_tensor(models.reference.model.last_final_norm), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close(cache["unembed.hook_out"], reference_logits, rtol=1e-5, atol=1e-6) + + pattern = cache["blocks.0.attn.hook_pattern"] + assert torch.all(pattern[0, :, 0, -1] > 0) + + changed_future = tokens.clone() + changed_future[0, -1] = 10 + with torch.inference_mode(): + changed_logits = models.bridge(changed_future) + assert (bridge_logits[:, 0] - changed_logits[:, 0]).abs().max() > 1e-8 + + +def test_padding_mask_blocks_keys_without_becoming_causal(models: TinyModels) -> None: + tokens = torch.tensor([[5, 63, 7, 9]]) + attention_mask = torch.tensor([[1, 1, 1, 0]]) + with torch.inference_mode(): + reference_logits = models.reference(tokens, attention_mask=attention_mask).logits + bridge_logits, cache = models.bridge.run_with_cache( + tokens, + attention_mask=attention_mask, + ) + + torch.testing.assert_close(bridge_logits, reference_logits, rtol=1e-5, atol=1e-6) + pattern = cache["blocks.0.attn.hook_pattern"] + torch.testing.assert_close( + pattern, + _require_tensor(_reference_block(models.reference).last_pattern), + rtol=1e-5, + atol=1e-6, + ) + assert torch.count_nonzero(pattern[..., -1]) == 0 + assert torch.all(pattern[0, :, 0, 1:3] > 0) + + changed_padding = tokens.clone() + changed_padding[0, -1] = 11 + with torch.inference_mode(): + changed_logits = models.bridge(changed_padding, attention_mask=attention_mask) + torch.testing.assert_close( + bridge_logits[:, :3], + changed_logits[:, :3], + rtol=1e-5, + atol=1e-6, + ) + + +def test_run_with_cache_exposes_hooks_without_hf_output_attentions( + models: TinyModels, +) -> None: + models.bridge.original_model.output_attentions_requests.clear() + tokens = torch.tensor([[5, 63, 7, 9]]) + _, cache = models.bridge.run_with_cache(tokens) + + expected_hooks = { + "embed.hook_out", + "blocks.0.hook_in", + "blocks.0.hook_out", + "blocks.0.ln1.hook_out", + "blocks.0.attn.hook_in", + "blocks.0.attn.q.hook_out", + "blocks.0.attn.k.hook_out", + "blocks.0.attn.v.hook_out", + "blocks.0.attn.o.hook_in", + "blocks.0.attn.o.hook_out", + "blocks.0.attn.hook_rot_q", + "blocks.0.attn.hook_rot_k", + "blocks.0.attn.hook_attn_scores", + "blocks.0.attn.hook_pattern", + "blocks.0.attn.hook_out", + "blocks.0.ln2.hook_out", + "blocks.0.mlp.hook_in", + "blocks.0.mlp.gate.hook_out", + "blocks.0.mlp.in.hook_out", + "blocks.0.mlp.act.hook_out", + "blocks.0.mlp.out.hook_in", + "blocks.0.mlp.hook_out", + "ln_final.hook_out", + "unembed.hook_out", + } + assert expected_hooks <= set(cache.keys()) + assert models.bridge.original_model.output_attentions_requests == [None] + + +def test_explicit_hf_output_attentions_request_is_not_silently_suppressed( + models: TinyModels, +) -> None: + models.bridge.original_model.output_attentions_requests.clear() + tokens = torch.tensor([[5, 63, 7, 9]]) + + with pytest.raises(ValueError, match="output_attentions is not yet supported"): + models.bridge(tokens, output_attentions=True) + + assert models.bridge.original_model.output_attentions_requests == [True] + + +def test_per_head_result_hook_matches_projected_attention(models: TinyModels) -> None: + tokens = torch.tensor([[5, 63, 7, 9]]) + previous = models.bridge.cfg.use_attn_result + models.bridge.set_use_attn_result(True) + try: + _, cache = models.bridge.run_with_cache( + tokens, + names_filter=[ + "blocks.0.attn.hook_result", + "blocks.0.attn.hook_out", + ], + ) + finally: + models.bridge.set_use_attn_result(previous) + + result = cache["blocks.0.attn.hook_result"] + projected = cache["blocks.0.attn.hook_out"] + assert result.shape == (1, 4, 4, 32) + torch.testing.assert_close(result.sum(dim=2), projected, rtol=1e-5, atol=1e-6) + + +def test_bridge_wraps_live_direct_block_components(models: TinyModels) -> None: + bridge_block = models.bridge.blocks[0] + original_block = bridge_block.original_component + assert isinstance(original_block, TinyLLaDALlamaBlock) + assert original_block.q_proj is bridge_block.attn.q + assert original_block.k_proj is bridge_block.attn.k + assert original_block.v_proj is bridge_block.attn.v + assert original_block.attn_out is bridge_block.attn.o + assert original_block.ff_proj is bridge_block.mlp.gate + assert original_block.up_proj is getattr(bridge_block.mlp, "in") + assert original_block.act is bridge_block.mlp.act + assert original_block.ff_out is bridge_block.mlp.out + + +def test_complete_reference_state_dict_mapping(models: TinyModels) -> None: + key_mapping = { + "model.transformer.wte.weight": "embed.weight", + "model.transformer.ln_f.weight": "ln_final.weight", + "model.transformer.ff_out.weight": "unembed.weight", + } + for layer in range(2): + reference_prefix = f"model.transformer.blocks.{layer}" + bridge_prefix = f"blocks.{layer}" + key_mapping.update( + { + f"{reference_prefix}.attn_norm.weight": f"{bridge_prefix}.ln1.weight", + f"{reference_prefix}.ff_norm.weight": f"{bridge_prefix}.ln2.weight", + f"{reference_prefix}.q_proj.weight": f"{bridge_prefix}.attn.q.weight", + f"{reference_prefix}.k_proj.weight": f"{bridge_prefix}.attn.k.weight", + f"{reference_prefix}.v_proj.weight": f"{bridge_prefix}.attn.v.weight", + f"{reference_prefix}.attn_out.weight": f"{bridge_prefix}.attn.o.weight", + f"{reference_prefix}.ff_proj.weight": f"{bridge_prefix}.mlp.gate.weight", + f"{reference_prefix}.up_proj.weight": f"{bridge_prefix}.mlp.in.weight", + f"{reference_prefix}.ff_out.weight": f"{bridge_prefix}.mlp.out.weight", + } + ) + + reference_state = models.reference.state_dict() + bridge_state = models.bridge.state_dict() + assert set(reference_state) == set(key_mapping) + assert set(bridge_state) == {*key_mapping.values(), "unembed.bias"} + for reference_key, bridge_key in key_mapping.items(): + torch.testing.assert_close(bridge_state[bridge_key], reference_state[reference_key]) + torch.testing.assert_close( + bridge_state["unembed.bias"], + torch.zeros_like(bridge_state["unembed.bias"]), + ) + + +def test_deepcopy_survives_original_bridge_collection() -> None: + local = _build_models() + tokens = torch.tensor([[5, 63, 7, 9]]) + with torch.inference_mode(): + expected_logits = local.bridge(tokens) + + original_block = local.bridge.blocks[0].original_component + original_block_ref = weakref.ref(original_block) + clone = copy.deepcopy(local.bridge) + assert clone.blocks[0].attn.original_component is clone.blocks[0].original_component + + del original_block + del local + gc.collect() + assert original_block_ref() is None + + with torch.inference_mode(): + clone_logits, clone_cache = clone.run_with_cache(tokens) + torch.testing.assert_close(clone_logits, expected_logits, rtol=1e-5, atol=1e-6) + assert { + "blocks.0.attn.hook_in", + "blocks.0.mlp.hook_in", + "blocks.0.mlp.hook_out", + } <= set(clone_cache) + + +def test_attention_dropout_follows_native_block_mode() -> None: + local = _build_models( + TinyLLaDAConfig(attention_dropout=1.0), + eval_bridge=False, + ) + assert local.bridge.training is True + assert local.bridge.blocks[0].original_component.training is False + tokens = torch.tensor([[5, 63, 7, 9]]) + + with torch.inference_mode(): + reference_logits = local.reference(tokens).logits + bridge_logits = local.bridge(tokens) + + torch.testing.assert_close(bridge_logits, reference_logits, rtol=1e-5, atol=1e-6) + + +def test_bfloat16_forward_matches_native_sdpa_with_dtype_appropriate_bound() -> None: + local = _build_models(dtype=torch.bfloat16) + tokens = torch.tensor([[5, 63, 7, 9]]) + with torch.inference_mode(): + reference_logits = local.reference(tokens).logits + bridge_logits = local.bridge(tokens) + + max_absolute_drift = (bridge_logits.float() - reference_logits.float()).abs().max().item() + assert max_absolute_drift < 0.05 + assert torch.equal(bridge_logits.argmax(dim=-1), reference_logits.argmax(dim=-1)) + + +@pytest.mark.parametrize("api_name", ["generate", "generate_stream", "hf_generate"]) +def test_autoregressive_generation_apis_fail_clearly(models: TinyModels, api_name: str) -> None: + tokens = torch.tensor([[5, 63, 7, 9]]) + with pytest.raises(NotImplementedError, match="generation is not supported"): + if api_name == "generate_stream": + next(models.bridge.generate_stream(tokens, max_new_tokens=1)) + else: + getattr(models.bridge, api_name)(tokens, max_new_tokens=1) + + +@pytest.mark.parametrize("return_type", ["loss", "both"]) +def test_causal_loss_modes_fail_clearly(models: TinyModels, return_type: str) -> None: + tokens = torch.tensor([[5, 63, 7, 9]]) + calls_before = len(models.bridge.original_model.output_attentions_requests) + with pytest.raises(NotImplementedError, match="causal.*loss|shifted.*loss"): + models.bridge(tokens, return_type=return_type) + assert len(models.bridge.original_model.output_attentions_requests) == calls_before diff --git a/tests/unit/model_bridge/supported_architectures/test_llada_adapter.py b/tests/unit/model_bridge/supported_architectures/test_llada_adapter.py new file mode 100644 index 0000000000..ae0fc981ce --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_llada_adapter.py @@ -0,0 +1,324 @@ +"""Offline unit tests for the dense LLaDA architecture adapter.""" + +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.conversion_utils.conversion_steps 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 import TransformerBridge +from transformer_lens.model_bridge.generalized_components import ( + AttentionBridge, + BlockBridge, + EmbeddingBridge, + GatedMLPBridge, + LinearBridge, + RMSNormalizationBridge, + UnembeddingBridge, +) +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.llada import ( + LLaDAArchitectureAdapter, +) + + +def _hf_config(**overrides: Any) -> SimpleNamespace: + values: dict[str, Any] = { + "activation_type": "silu", + "alibi": False, + "architectures": ["LLaDAModelLM"], + "attention_dropout": 0.0, + "attention_layer_norm": False, + "block_group_size": 1, + "block_type": "llama", + "d_model": 32, + "embedding_dropout": 0.0, + "embedding_size": 64, + "eos_token_id": 1, + "include_bias": False, + "include_qkv_bias": False, + "input_emb_norm": False, + "layer_norm_type": "rms", + "mask_token_id": 63, + "max_sequence_length": 16, + "mlp_hidden_size": 48, + "model_type": "llada", + "n_heads": 4, + "n_kv_heads": 2, + "n_layers": 2, + "pad_token_id": 0, + "residual_dropout": 0.0, + "rms_norm_eps": 1e-5, + "rope": True, + "rope_full_precision": True, + "rope_theta": 500_000.0, + "scale_logits": False, + "use_cache": False, + "vocab_size": 64, + "weight_tying": False, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _bridge_config(**overrides: Any) -> TransformerBridgeConfig: + config = build_bridge_config_from_hf( + _hf_config(), + architecture="LLaDAModelLM", + model_name="tiny-llada", + dtype=torch.float32, + ) + for name, value in overrides.items(): + setattr(config, name, value) + return config + + +@pytest.fixture(scope="module") +def adapter() -> LLaDAArchitectureAdapter: + return LLaDAArchitectureAdapter(_bridge_config()) + + +def _rearrange(adapter: LLaDAArchitectureAdapter, key: str) -> RearrangeTensorConversion: + conversions = adapter.weight_processing_conversions + assert conversions is not None + conversion = conversions[key] + assert isinstance(conversion, ParamProcessingConversion) + assert isinstance(conversion.tensor_conversion, RearrangeTensorConversion) + return conversion.tensor_conversion + + +def test_config_conversion_preserves_llada_fields() -> None: + config = _bridge_config() + + expected = { + "d_model": 32, + "d_head": 8, + "n_layers": 2, + "n_ctx": 16, + "n_heads": 4, + "n_key_value_heads": 2, + "d_mlp": 48, + "d_vocab": 64, + "act_fn": "silu", + "eps": 1e-5, + "rotary_base": 500_000, + "tie_word_embeddings": False, + "mask_token_id": 63, + "pad_token_id": 0, + "eos_token_id": 1, + "block_type": "llama", + "block_group_size": 1, + "rope": True, + "rope_full_precision": True, + "attention_layer_norm": False, + "include_bias": False, + "include_qkv_bias": False, + "scale_logits": False, + "input_emb_norm": False, + "layer_norm_type": "rms", + } + for name, value in expected.items(): + assert getattr(config, name) == value + + +def test_factory_and_model_type_dispatch_to_llada() -> None: + selected = ArchitectureAdapterFactory.select_architecture_adapter(_bridge_config()) + assert isinstance(selected, LLaDAArchitectureAdapter) + + config = SimpleNamespace(architectures=[], model_type="llada") + assert determine_architecture_from_hf_config(config) == "LLaDAModelLM" + + +def test_adapter_sets_bidirectional_forward_only_semantics( + adapter: LLaDAArchitectureAdapter, +) -> None: + assert adapter.cfg.attention_dir == "bidirectional" + assert adapter.cfg.normalization_type == "RMS" + assert adapter.cfg.uses_rms_norm is True + assert adapter.cfg.positional_embedding_type == "rotary" + assert adapter.cfg.rotary_adjacent_pairs is False + assert adapter.cfg.gated_mlp is True + assert adapter.cfg.final_rms is True + assert adapter.cfg.d_vocab_out == adapter.cfg.d_vocab == 64 + assert adapter.cfg.default_prepend_bos is False + assert adapter.cfg.default_padding_side == "right" + assert adapter.supports_generation is False + assert adapter.supports_hf_output_attentions is False + assert adapter.supports_causal_loss is False + + attention = adapter.component_mapping["blocks"].submodules["attn"] + assert isinstance(attention, AttentionBridge) + assert attention.is_causal is False + + +def test_component_mapping_matches_block_local_module_tree( + adapter: LLaDAArchitectureAdapter, +) -> None: + mapping = adapter.component_mapping + assert isinstance(mapping["embed"], EmbeddingBridge) + assert isinstance(mapping["blocks"], BlockBridge) + assert isinstance(mapping["ln_final"], RMSNormalizationBridge) + assert isinstance(mapping["unembed"], UnembeddingBridge) + assert mapping["embed"].name == "model.transformer.wte" + assert mapping["blocks"].name == "model.transformer.blocks" + assert mapping["ln_final"].name == "model.transformer.ln_f" + assert mapping["unembed"].name == "model.transformer.ff_out" + + block = mapping["blocks"] + assert block.submodules["ln1"].name == "attn_norm" + assert block.submodules["ln2"].name == "ff_norm" + assert isinstance(block.submodules["ln1"], RMSNormalizationBridge) + assert isinstance(block.submodules["ln2"], RMSNormalizationBridge) + + attention = block.submodules["attn"] + assert isinstance(attention, AttentionBridge) + assert {name: component.name for name, component in attention.submodules.items()} == { + "q": "q_proj", + "k": "k_proj", + "v": "v_proj", + "o": "attn_out", + } + assert all(isinstance(component, LinearBridge) for component in attention.submodules.values()) + + mlp = block.submodules["mlp"] + assert isinstance(mlp, GatedMLPBridge) + assert {name: component.name for name, component in mlp.submodules.items()} == { + "gate": "ff_proj", + "in": "up_proj", + "act": "act", + "out": "ff_out", + } + + +@pytest.mark.parametrize( + "tl_path,hf_path", + [ + ("blocks.0.attn.W_Q", "model.transformer.blocks.0.q_proj.weight"), + ("blocks.0.attn.W_K", "model.transformer.blocks.0.k_proj.weight"), + ("blocks.0.attn.W_V", "model.transformer.blocks.0.v_proj.weight"), + ("blocks.0.attn.W_O", "model.transformer.blocks.0.attn_out.weight"), + ("blocks.0.mlp.W_gate", "model.transformer.blocks.0.ff_proj.weight"), + ("blocks.0.mlp.W_in", "model.transformer.blocks.0.up_proj.weight"), + ("blocks.0.mlp.W_out", "model.transformer.blocks.0.ff_out.weight"), + ], +) +def test_containerless_paths_translate_to_direct_block_fields( + adapter: LLaDAArchitectureAdapter, tl_path: str, hf_path: str +) -> None: + assert adapter.translate_transformer_lens_path(tl_path) == hf_path + + +@pytest.mark.parametrize( + "hf_key,tl_key", + [ + ("model.transformer.blocks.0.q_proj.weight", "blocks.0.attn.q.weight"), + ("model.transformer.blocks.0.k_proj.weight", "blocks.0.attn.k.weight"), + ("model.transformer.blocks.0.v_proj.weight", "blocks.0.attn.v.weight"), + ("model.transformer.blocks.0.attn_out.weight", "blocks.0.attn.o.weight"), + ("model.transformer.blocks.0.ff_proj.weight", "blocks.0.mlp.gate.weight"), + ("model.transformer.blocks.0.up_proj.weight", "blocks.0.mlp.in.weight"), + ("model.transformer.blocks.0.ff_out.weight", "blocks.0.mlp.out.weight"), + ], +) +def test_hf_keys_translate_back_to_containerless_components( + adapter: LLaDAArchitectureAdapter, hf_key: str, tl_key: str +) -> None: + assert adapter.convert_hf_key_to_tl_key(hf_key) == tl_key + + +def test_qkvo_weight_conversions_preserve_gqa_head_counts( + adapter: LLaDAArchitectureAdapter, +) -> None: + conversions = adapter.weight_processing_conversions + assert conversions is not None + assert set(conversions) == { + "blocks.{i}.attn.q.weight", + "blocks.{i}.attn.k.weight", + "blocks.{i}.attn.v.weight", + "blocks.{i}.attn.o.weight", + } + assert _rearrange(adapter, "blocks.{i}.attn.q.weight").axes_lengths["n"] == 4 + assert _rearrange(adapter, "blocks.{i}.attn.k.weight").axes_lengths["n"] == 2 + assert _rearrange(adapter, "blocks.{i}.attn.v.weight").axes_lengths["n"] == 2 + assert _rearrange(adapter, "blocks.{i}.attn.o.weight").axes_lengths["n"] == 4 + + +def test_equal_kv_and_query_head_counts_use_mha_default() -> None: + config = build_bridge_config_from_hf( + _hf_config(n_kv_heads=4), + architecture="LLaDAModelLM", + model_name="tiny-llada", + dtype=torch.float32, + ) + + assert config.n_key_value_heads is None + + +@pytest.mark.parametrize( + "field,value,error_fragment", + [ + ("block_type", "sequential", "block_type"), + ("block_group_size", 2, "block_group_size"), + ("rope", False, "rope"), + ("rope_full_precision", False, "rope_full_precision"), + ("alibi", True, "alibi"), + ("attention_layer_norm", True, "attention_layer_norm"), + ("include_bias", True, "include_bias"), + ("include_qkv_bias", True, "include_qkv_bias"), + ("scale_logits", True, "scale_logits"), + ("input_emb_norm", True, "input_emb_norm"), + ("layer_norm_type", "default", "layer_norm_type"), + ("act_fn", "relu", "activation_type"), + ("tie_word_embeddings", True, "tied embeddings"), + ("embedding_size", 128, "embedding_size"), + ], +) +def test_unsupported_dense_variants_fail_clearly( + field: str, value: Any, error_fragment: str +) -> None: + with pytest.raises(ValueError, match=error_fragment): + LLaDAArchitectureAdapter(_bridge_config(**{field: value})) + + +def test_existing_causal_attention_default_remains_triangular() -> None: + config = TransformerBridgeConfig( + d_model=16, + d_head=4, + n_heads=4, + n_layers=1, + n_ctx=5, + d_vocab=32, + d_mlp=32, + act_fn="silu", + normalization_type="RMS", + final_rms=True, + gated_mlp=True, + architecture="TransformerLensNative", + attention_dir="causal", + default_prepend_bos=False, + seed=0, + ) + bridge = TransformerBridge.boot_native(config, device="cpu") + tokens = torch.tensor([[1, 2, 3, 4, 5]]) + + _, cache = bridge.run_with_cache( + tokens, + names_filter="blocks.0.attn.hook_pattern", + ) + + pattern = cache["blocks.0.attn.hook_pattern"] + assert pattern.shape == (1, 4, 5, 5) + assert torch.count_nonzero(torch.triu(pattern, diagonal=1)) == 0 diff --git a/tests/unit/model_bridge/test_architecture_adapter.py b/tests/unit/model_bridge/test_architecture_adapter.py index d42a0cf34d..2609c5b52b 100644 --- a/tests/unit/model_bridge/test_architecture_adapter.py +++ b/tests/unit/model_bridge/test_architecture_adapter.py @@ -206,6 +206,23 @@ def test_translate_transformer_lens_path_last_component(adapter: Gemma3Architect ) +def test_translate_containerless_descendant_without_inventing_container_path( + adapter: Gemma3ArchitectureAdapter, +) -> None: + adapter.component_mapping["blocks"].submodules["synthetic"] = MLPBridge( + name=None, + submodules={"in": LinearBridge(name="fc1")}, + ) + + assert adapter.translate_transformer_lens_path("blocks.0.synthetic.in") == "model.layers.0.fc1" + assert ( + adapter.translate_transformer_lens_path("blocks.0.synthetic.in", last_component_only=True) + == "fc1" + ) + with pytest.raises(ValueError, match="Synthetic component synthetic has no remote path"): + adapter.translate_transformer_lens_path("blocks.0.synthetic") + + def test_component_mapping_structure(adapter: Gemma3ArchitectureAdapter) -> None: """Test that the component mapping has the expected structure.""" mapping = adapter.get_component_mapping() diff --git a/tests/unit/model_bridge/test_component_setup.py b/tests/unit/model_bridge/test_component_setup.py index 09cc578833..98639e2c03 100644 --- a/tests/unit/model_bridge/test_component_setup.py +++ b/tests/unit/model_bridge/test_component_setup.py @@ -19,9 +19,13 @@ AttentionBridge, BlockBridge, EmbeddingBridge, + LinearBridge, MLPBridge, NormalizationBridge, ) +from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, +) from transformer_lens.model_bridge.generalized_components.joint_qkv_attention import ( JointQKVAttentionBridge, ) @@ -121,6 +125,30 @@ def test_setup_submodules_empty(self): # Should not raise any errors setup_submodules(component, adapter, original_ln) + def test_containerless_executable_promotes_children_for_processed_weights(self): + adapter = MockArchitectureAdapter() + original_block = nn.Module() + original_block.fc_in = nn.Linear(4, 8, bias=False) + original_block.fc_out = nn.Linear(8, 4, bias=False) + mlp = MLPBridge( + name=None, + submodules={ + "in": LinearBridge(name="fc_in"), + "out": LinearBridge(name="fc_out"), + }, + ) + block = GeneralizedComponent(name="block", submodules={"mlp": mlp}) + block.set_original_component(original_block) + + setup_submodules(block, adapter, original_block) + + mlp_in = getattr(mlp, "in") + assert set(block.real_components) == {"mlp.in", "mlp.out"} + assert block.real_components["mlp.in"][1] is mlp_in + new_weight = torch.randn_like(mlp_in.original_component.weight) + block.set_processed_weights({"mlp.in.weight": new_weight}) + torch.testing.assert_close(mlp_in.original_component.weight, new_weight) + def test_setup_components_regular_component(self): """Test setting up regular (non-list) components.""" adapter = MockArchitectureAdapter() diff --git a/tests/unit/model_bridge/test_native_features.py b/tests/unit/model_bridge/test_native_features.py index abc5aebf45..85f6b5ad46 100644 --- a/tests/unit/model_bridge/test_native_features.py +++ b/tests/unit/model_bridge/test_native_features.py @@ -9,11 +9,15 @@ from __future__ import annotations +from unittest.mock import patch + +import pytest import torch from transformer_lens.config import TransformerBridgeConfig from transformer_lens.model_bridge import TransformerBridge from transformer_lens.model_bridge.generalized_components import ( + AttentionBridge, GatedMLPBridge, MLPBridge, NormalizationBridge, @@ -29,6 +33,11 @@ ) +class _ResultOnlyAttentionBridge(AttentionBridge): + supports_split_qkv_fork = False + supports_attn_result = True + + def _cfg(**overrides) -> TransformerBridgeConfig: base = dict( d_model=32, @@ -51,6 +60,61 @@ def _forward(bridge: TransformerBridge) -> torch.Tensor: return bridge(inputs, return_type="logits") +def test_result_hook_capability_is_independent_of_qkv_input_forks() -> None: + bridge = TransformerBridge.boot_native(_cfg()) + bridge.blocks[0]._modules["attn"] = _ResultOnlyAttentionBridge( + name="result_only", + config=bridge.cfg, + ) + + bridge.set_use_attn_result(True) + assert bridge.cfg.use_attn_result is True + with pytest.raises(NotImplementedError, match="use_split_qkv_input"): + bridge.set_use_split_qkv_input(True) + + +def test_forward_only_capabilities_reject_generation_and_causal_loss() -> None: + bridge = TransformerBridge.boot_native(_cfg()) + bridge.adapter.supports_generation = False + bridge.adapter.supports_causal_loss = False + tokens = torch.randint(0, bridge.cfg.d_vocab, (1, 4)) + + with pytest.raises(NotImplementedError, match="generation is not supported"): + bridge.generate(tokens, max_new_tokens=1) + with pytest.raises(NotImplementedError, match="shifted causal loss"): + bridge(tokens, return_type="loss") + + +def test_output_attention_capability_controls_only_the_implicit_request() -> None: + bridge = TransformerBridge.boot_native(_cfg()) + tokens = torch.randint(0, bridge.cfg.d_vocab, (1, 4)) + + with patch.object( + bridge.original_model, + "forward", + wraps=bridge.original_model.forward, + ) as forward: + bridge.run_with_cache(tokens) + assert forward.call_args.kwargs["output_attentions"] is True + + bridge.adapter.supports_hf_output_attentions = False + with patch.object( + bridge.original_model, + "forward", + wraps=bridge.original_model.forward, + ) as forward: + bridge.run_with_cache(tokens) + assert "output_attentions" not in forward.call_args.kwargs + + with patch.object( + bridge.original_model, + "forward", + wraps=bridge.original_model.forward, + ) as forward: + bridge.run_with_cache(tokens, output_attentions=True) + assert forward.call_args.kwargs["output_attentions"] is True + + # -- soft-cap ----------------------------------------------------------------- @@ -534,9 +598,6 @@ def test_rotary_pattern_differs_from_absolute(): # -- init modes --------------------------------------------------------------- -import pytest - - @pytest.mark.parametrize( "init_mode", ["gpt2", "xavier_uniform", "xavier_normal", "kaiming_uniform", "kaiming_normal"], diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py index 2342908d37..3b0b4aa6bd 100644 --- a/transformer_lens/factories/architecture_adapter_factory.py +++ b/transformer_lens/factories/architecture_adapter_factory.py @@ -40,6 +40,7 @@ HunYuanDenseV1ArchitectureAdapter, InternLM2ArchitectureAdapter, Lfm2MoeArchitectureAdapter, + LLaDAArchitectureAdapter, LlamaArchitectureAdapter, LlavaArchitectureAdapter, LlavaNextArchitectureAdapter, @@ -117,6 +118,7 @@ "HubertModel": HubertArchitectureAdapter, "HunYuanDenseV1ForCausalLM": HunYuanDenseV1ArchitectureAdapter, "InternLM2ForCausalLM": InternLM2ArchitectureAdapter, + "LLaDAModelLM": LLaDAArchitectureAdapter, "LlamaForCausalLM": LlamaArchitectureAdapter, "LlavaForConditionalGeneration": LlavaArchitectureAdapter, "LlavaNextForConditionalGeneration": LlavaNextArchitectureAdapter, diff --git a/transformer_lens/model_bridge/architecture_adapter.py b/transformer_lens/model_bridge/architecture_adapter.py index 39e3ce1df4..d1945ac2c2 100644 --- a/transformer_lens/model_bridge/architecture_adapter.py +++ b/transformer_lens/model_bridge/architecture_adapter.py @@ -48,6 +48,14 @@ class ArchitectureAdapter: # Encoder-only models (e.g. BERT, HuBERT) should set this to False. supports_generation: bool = True + # Whether run_with_cache should request attention tensors from Hugging Face. + # Set False when an adapter reconstructs those tensors but the HF wrapper + # rejects output_attentions=True. + supports_hf_output_attentions: bool = True + + # Whether Bridge's shifted next-token cross-entropy is meaningful. + supports_causal_loss: bool = True + # Optional libraries this adapter needs at load time (e.g. the multimodal group's timm). # Checked at construction so a missing one raises a clear error, not a deep HF failure. required_libraries: list[str] = [] @@ -453,22 +461,16 @@ def translate_transformer_lens_path( subcomponent_bridge = bridge_component.submodules[subcomponent_name] if len(parts) > 3: current_bridge = subcomponent_bridge - subcomponent_name_str = subcomponent_bridge.name - if subcomponent_name_str is None: - raise ValueError( - f"Subcomponent {subcomponent_name} must have a name for path translation" - ) - remote_path_parts = [items_path, item_index, subcomponent_name_str] + remote_path_parts = [items_path, item_index] + if subcomponent_bridge.name is not None: + remote_path_parts.append(subcomponent_bridge.name) for i in range(3, len(parts)): deeper_component_name = parts[i] if deeper_component_name in current_bridge.submodules: current_bridge = current_bridge.submodules[deeper_component_name] deeper_name = current_bridge.name - if deeper_name is None: - raise ValueError( - f"Component {deeper_component_name} must have a name for path translation" - ) - remote_path_parts.append(deeper_name) + if deeper_name is not None: + remote_path_parts.append(deeper_name) else: raise ValueError( f"Component {deeper_component_name} not found in {'.'.join(parts[:i])} components" @@ -483,7 +485,7 @@ def translate_transformer_lens_path( subcomponent_name_str = subcomponent_bridge.name if subcomponent_name_str is None: raise ValueError( - f"Subcomponent {subcomponent_name} must have a name for path translation" # type: ignore[assignment] + f"Synthetic component {subcomponent_name} has no remote path" ) remote_path = f"{items_path}.{item_index}.{subcomponent_name_str}" if param_suffix: diff --git a/transformer_lens/model_bridge/bridge.py b/transformer_lens/model_bridge/bridge.py index 4800b0150d..b000b76c96 100644 --- a/transformer_lens/model_bridge/bridge.py +++ b/transformer_lens/model_bridge/bridge.py @@ -1746,6 +1746,14 @@ def forward( Model output based on return_type """ + if return_type in ("loss", "both") and not self.adapter.supports_causal_loss: + architecture = self.cfg.architecture or type(self.adapter).__name__ + raise NotImplementedError( + f"{architecture} does not support TransformerBridge's shifted causal " + "loss. Request return_type='logits' and compute the architecture-specific " + "masked-token objective explicitly." + ) + if start_at_layer is not None: raise NotImplementedError( "start_at_layer is not supported in TransformerBridge. " @@ -2189,7 +2197,10 @@ def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor: stacklevel=2, ) try: - if "output_attentions" not in filtered_kwargs: + if ( + "output_attentions" not in filtered_kwargs + and self.adapter.supports_hf_output_attentions + ): filtered_kwargs["output_attentions"] = True if processed_args: output = self.forward(processed_args[0], **filtered_kwargs) @@ -2671,6 +2682,15 @@ def _generate_tokens( if all_finished: return + def _ensure_generation_supported(self, api_name: str) -> None: + """Reject autoregressive generation for forward-only architectures.""" + if not self.adapter.supports_generation: + architecture = self.cfg.architecture or type(self.adapter).__name__ + raise NotImplementedError( + f"TransformerBridge.{api_name}() generation is not supported by " + f"the {architecture} architecture." + ) + def generate( self, input: Union[str, List[str], torch.Tensor] = "", @@ -2790,6 +2810,7 @@ def generate( the tokens that were fed to the model, useful for verifying BOS handling with chat templates. """ + self._ensure_generation_supported("generate") # padding_side is handled internally: for batched list inputs, left-padding # is forced to ensure correct generation. See _is_batched_list logic below. @@ -3206,6 +3227,7 @@ def generate_stream( max_tokens_per_yield tokens between yields. First yield includes the input tokens; subsequent yields contain only new tokens. """ + self._ensure_generation_supported("generate_stream") # --- Input parsing (mirrors generate()) --- _is_batched_list = isinstance(input, list) and len(input) > 1 @@ -3437,6 +3459,7 @@ def hf_generate( print(result.logits) # Logits for each generation step print(result.attentions) # Attention weights """ + self._ensure_generation_supported("hf_generate") # Handle string input by tokenizing it if isinstance(input, str): inputs = self.tokenizer(input, return_tensors="pt", padding=False, truncation=False).to( @@ -3942,18 +3965,25 @@ def _validate_attention_fork_supported(self, flag_name: str) -> None: continue total_with_attn += 1 attn_classes.add(type(attn).__name__) - if isinstance(attn, supported_classes): + supports_flag = ( + bool(getattr(attn, "supports_attn_result", False)) + if flag_name == "use_attn_result" + else isinstance(attn, supported_classes) + ) + if supports_flag: supporting_layers.append(idx) if total_with_attn == 0: raise NotImplementedError(f"{flag_name}: no attention bridges found on self.blocks.") if not supporting_layers: + if flag_name == "use_attn_result": + capability_detail = "Per-head result computation is unavailable." + else: + capability_detail = "No hook point is available before the Q/K/V projections." raise NotImplementedError( f"{flag_name}: none of this model's attention bridges support " - "the fine-grained Q/K/V hook fork. Found attention classes: " + "the requested fine-grained attention hook. Found attention classes: " f"{sorted(attn_classes)}. Supported classes: " - f"{[c.__name__ for c in supported_classes]}. Plain " - "AttentionBridge delegates to HuggingFace and exposes no hook " - "point before the Q/K/V projection." + f"{[c.__name__ for c in supported_classes]}. {capability_detail}" ) if len(supporting_layers) < total_with_attn: skipped = total_with_attn - len(supporting_layers) diff --git a/transformer_lens/model_bridge/component_setup.py b/transformer_lens/model_bridge/component_setup.py index 7821d03542..fc21f99eac 100644 --- a/transformer_lens/model_bridge/component_setup.py +++ b/transformer_lens/model_bridge/component_setup.py @@ -137,8 +137,14 @@ def setup_submodules( if module_name not in component._modules: component.add_module(module_name, submodule) - # Add to real_components mapping (for non-list components) - if not submodule.is_list_item and submodule.name is not None: + # Containerless executable bridges still own real child components. + # Promote them just like SymbolicBridge so processed weights reach + # projections that live directly on the parent block. + if not submodule.is_list_item and submodule.name is None: + for sub_name, (sub_path, sub_comp) in submodule.real_components.items(): + prefixed_key = f"{module_name}.{sub_name}" + component.real_components[prefixed_key] = (sub_path, sub_comp) + elif not submodule.is_list_item: component.real_components[module_name] = (submodule.name, submodule) # Clean up so architecture_adapter traversal won't find stale entries diff --git a/transformer_lens/model_bridge/generalized_components/attention.py b/transformer_lens/model_bridge/generalized_components/attention.py index cd79d2386b..d3498f314c 100644 --- a/transformer_lens/model_bridge/generalized_components/attention.py +++ b/transformer_lens/model_bridge/generalized_components/attention.py @@ -40,6 +40,8 @@ class AttentionBridge(GeneralizedComponent): # Override to False on variants without a pre-LN fork (e.g. MLA); skips # the split-qkv HookPoints and the BlockBridge pre-ln1 capture. supports_split_qkv_fork: bool = True + # Reconstructed variants can opt in independently of Q/K/V input forks. + supports_attn_result: bool = False property_aliases = { "W_Q": "q.weight", "W_K": "k.weight", @@ -53,7 +55,7 @@ class AttentionBridge(GeneralizedComponent): def __init__( self, - name: str, + name: Optional[str], config: Any, submodules: Optional[Dict[str, GeneralizedComponent]] = None, conversion_rule: Optional[BaseTensorConversion] = None, @@ -70,7 +72,7 @@ def __init__( """Initialize the attention bridge. Args: - name: The name of this component + name: The name of this component, or None when projections live on the parent config: Model configuration (required for auto-conversion detection) submodules: Dictionary of submodules to register (e.g., q_proj, k_proj, etc.) conversion_rule: Optional conversion rule. If None, AttentionAutoConversion will be used diff --git a/transformer_lens/model_bridge/generalized_components/joint_qkv_attention.py b/transformer_lens/model_bridge/generalized_components/joint_qkv_attention.py index 6b54871509..2490c3ee38 100644 --- a/transformer_lens/model_bridge/generalized_components/joint_qkv_attention.py +++ b/transformer_lens/model_bridge/generalized_components/joint_qkv_attention.py @@ -27,6 +27,8 @@ class JointQKVAttentionBridge(AttentionBridge): the individual activations from the separated q, k, and v matrices are hooked and accessible. """ + supports_attn_result: bool = True + # property_aliases inherited from AttentionBridge (W_Q, W_K, W_V, W_O, b_Q, b_K, b_V, b_O) def __init__( diff --git a/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py b/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py index 430bf32afa..908012d2cf 100644 --- a/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py +++ b/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py @@ -112,6 +112,8 @@ class PositionEmbeddingsAttentionBridge(PositionEmbeddingHooksMixin, AttentionBr component with dummy Q/K tensors and position_ids. """ + supports_attn_result: bool = True + def __init__( self, name: str, diff --git a/transformer_lens/model_bridge/sources/_bridge_builder.py b/transformer_lens/model_bridge/sources/_bridge_builder.py index 55f53ada85..9f40ef054a 100644 --- a/transformer_lens/model_bridge/sources/_bridge_builder.py +++ b/transformer_lens/model_bridge/sources/_bridge_builder.py @@ -58,6 +58,21 @@ "router_jitter_noise", "input_jitter_noise", "eos_token_id", + # LLaDA remote-code model contract and tokenizer metadata + "block_type", + "block_group_size", + "rope", + "rope_full_precision", + "attention_layer_norm", + "include_bias", + "include_qkv_bias", + "scale_logits", + "input_emb_norm", + "layer_norm_type", + "embedding_size", + "mask_token_id", + "pad_token_id", + "bos_token_id", ] diff --git a/transformer_lens/model_bridge/sources/transformers.py b/transformer_lens/model_bridge/sources/transformers.py index bce60d96f7..8509af5bbf 100644 --- a/transformer_lens/model_bridge/sources/transformers.py +++ b/transformer_lens/model_bridge/sources/transformers.py @@ -80,6 +80,8 @@ def map_default_transformer_lens_config(hf_config): tl_config.n_heads = n_heads elif hasattr(source_config, "num_heads"): tl_config.n_heads = source_config.num_heads + elif hasattr(source_config, "n_heads"): + tl_config.n_heads = source_config.n_heads elif hasattr(source_config, "num_query_heads") and isinstance( source_config.num_query_heads, list ): @@ -120,6 +122,13 @@ def map_default_transformer_lens_config(hf_config): tl_config.n_key_value_heads = num_kv_heads except (TypeError, ValueError, AttributeError): pass + elif hasattr(source_config, "n_kv_heads") and source_config.n_kv_heads is not None: + try: + num_kv_heads = int(source_config.n_kv_heads) + if num_kv_heads != getattr(tl_config, "n_heads", None): + tl_config.n_key_value_heads = num_kv_heads + except (TypeError, ValueError, AttributeError): + pass if hasattr(source_config, "n_layer"): tl_config.n_layers = source_config.n_layer elif hasattr(source_config, "num_hidden_layers"): @@ -128,6 +137,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_layers"): + tl_config.n_layers = source_config.n_layers 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"): @@ -140,6 +151,8 @@ def map_default_transformer_lens_config(hf_config): tl_config.n_ctx = source_config.max_length elif hasattr(source_config, "seq_length"): tl_config.n_ctx = source_config.seq_length + elif hasattr(source_config, "max_sequence_length"): + tl_config.n_ctx = source_config.max_sequence_length else: # Models like Bloom use ALiBi (no positional embeddings) and have no # context length field. Default to 2048 as a reasonable fallback. @@ -155,6 +168,8 @@ def map_default_transformer_lens_config(hf_config): if isinstance(intermediate_size, (list, tuple)): intermediate_size = max(intermediate_size) if intermediate_size else None tl_config.d_mlp = intermediate_size + elif hasattr(source_config, "mlp_hidden_size"): + tl_config.d_mlp = source_config.mlp_hidden_size elif hasattr(tl_config, "d_model"): tl_config.d_mlp = getattr(source_config, "n_inner", 4 * tl_config.d_model) if hasattr(source_config, "head_dim") and source_config.head_dim is not None: @@ -171,6 +186,13 @@ def map_default_transformer_lens_config(hf_config): tl_config.act_fn = source_config.activation_function elif hasattr(source_config, "hidden_act"): tl_config.act_fn = source_config.hidden_act + elif hasattr(source_config, "activation_type"): + activation_type = source_config.activation_type + tl_config.act_fn = getattr(activation_type, "value", activation_type) + if hasattr(source_config, "rope_theta"): + tl_config.rotary_base = source_config.rope_theta + if hasattr(source_config, "weight_tying"): + tl_config.tie_word_embeddings = bool(source_config.weight_tying) # Layer norm / RMS norm epsilon — HF uses 3 different field names if hasattr(source_config, "rms_norm_eps"): tl_config.eps = source_config.rms_norm_eps @@ -221,6 +243,7 @@ def determine_architecture_from_hf_config(hf_config): "hubert": "HubertModel", "bart": "BartForConditionalGeneration", "llama": "LlamaForCausalLM", + "llada": "LLaDAModelLM", "mamba": "MambaForCausalLM", "mamba2": "Mamba2ForCausalLM", "mistral": "MistralForCausalLM", @@ -403,7 +426,7 @@ def boot( the field name. If larger than the model's default, a warning is emitted — quality may degrade past the trained length for rotary models. revision: Optional HF revision string (branch, tag, or commit). Forwarded to - ``AutoConfig.from_pretrained`` and ``AutoModelForCausalLM.from_pretrained``. + config, model, and tokenizer loading. Mutually exclusive with ``checkpoint_index`` and ``checkpoint_value``. checkpoint_index: Index into the available training checkpoints for the model family. Convenience over ``revision`` for checkpointed models like EleutherAI/pythia* and @@ -458,6 +481,7 @@ def boot( "max_context_length", "max_length", "seq_length", + "max_sequence_length", ): if hasattr(hf_config, _field): _n_ctx_field = _field @@ -504,6 +528,7 @@ def boot( bridge_config.architecture = architecture bridge_config.model_name = model_name bridge_config.dtype = dtype + bridge_config.trust_remote_code = trust_remote_code # Propagate HF-specific config attributes that adapters may need. # Any attribute present on the HF config and not None is copied to bridge_config. # This is architecture-agnostic — new architectures don't need changes here. @@ -550,6 +575,21 @@ def boot( "router_jitter_noise", "input_jitter_noise", "eos_token_id", + # LLaDA remote-code model contract and tokenizer metadata + "block_type", + "block_group_size", + "rope", + "rope_full_precision", + "attention_layer_norm", + "include_bias", + "include_qkv_bias", + "scale_logits", + "input_emb_norm", + "layer_norm_type", + "embedding_size", + "mask_token_id", + "pad_token_id", + "bos_token_id", ] for attr in _HF_PASSTHROUGH_ATTRS: val = getattr(hf_config, attr, None) @@ -746,6 +786,7 @@ def boot( use_fast=use_fast, token=token_arg, trust_remote_code=trust_remote_code, + revision=revision, ) except ValueError: # Model doesn't have a BOS token, load without add_bos_token @@ -754,6 +795,7 @@ def boot( use_fast=use_fast, token=token_arg, trust_remote_code=trust_remote_code, + revision=revision, ) tokenizer = setup_tokenizer( base_tokenizer, diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py index 8f4cfea3ed..ce39aa7fea 100644 --- a/transformer_lens/model_bridge/supported_architectures/__init__.py +++ b/transformer_lens/model_bridge/supported_architectures/__init__.py @@ -93,6 +93,9 @@ from transformer_lens.model_bridge.supported_architectures.lfm2_moe import ( Lfm2MoeArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.llada import ( + LLaDAArchitectureAdapter, +) from transformer_lens.model_bridge.supported_architectures.llama import ( LlamaArchitectureAdapter, ) @@ -235,6 +238,7 @@ "HubertArchitectureAdapter", "HunYuanDenseV1ArchitectureAdapter", "InternLM2ArchitectureAdapter", + "LLaDAArchitectureAdapter", "LlamaArchitectureAdapter", "LlavaArchitectureAdapter", "LlavaNextArchitectureAdapter", diff --git a/transformer_lens/model_bridge/supported_architectures/llada.py b/transformer_lens/model_bridge/supported_architectures/llada.py new file mode 100644 index 0000000000..2ef605679f --- /dev/null +++ b/transformer_lens/model_bridge/supported_architectures/llada.py @@ -0,0 +1,371 @@ +"""LLaDA architecture support for one masked-token transformer forward pass. + +LLaDA is a masked discrete-diffusion model. Its transformer uses bidirectional +self-attention and returns logits for every input position; the iterative +denoising sampler lives outside the model. This adapter intentionally supports +only the transformer forward pass. TransformerBridge's autoregressive +generation APIs are disabled for this architecture. + +The remote ``LLaDALlamaBlock`` owns its Q/K/V, output, and gated-MLP projections +directly instead of grouping them in attention and MLP modules. The adapter +preserves the reviewed remote block forward and replaces only its +``attention(...)`` method with a hook-aware reconstruction. This keeps the +native pre-norm residual order, dropout, RoPE implementation, and MLP math while +exposing the standard attention scores and pattern hooks. +""" + +from __future__ import annotations + +import math +from typing import Any, Dict, Optional, cast + +import torch + +from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter +from transformer_lens.model_bridge.generalized_components import ( + AttentionBridge, + BlockBridge, + EmbeddingBridge, + GatedMLPBridge, + LinearBridge, + RMSNormalizationBridge, + UnembeddingBridge, +) +from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, +) + + +class _LLaDAAttentionBridge(AttentionBridge): + """Reconstruct LLaDA's block-local, bidirectional attention method.""" + + supports_split_qkv_fork: bool = False + supports_attn_result: bool = True + + def set_original_component(self, original_component: torch.nn.Module) -> None: + """Patch the owning block without registering it as a child module.""" + object.__setattr__(self, "_original_block", original_component) + object.__setattr__(original_component, "attention", self) + + @property + def original_component(self) -> Optional[torch.nn.Module]: + """Return the owning LLaDA block.""" + return self.__dict__.get("_original_block") + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attention_bias: Optional[torch.Tensor] = None, + layer_past: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + use_cache: bool = False, + ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]: + """Run LLaDA attention with observable pre/post-softmax tensors.""" + block = self.original_component + if block is None: + raise RuntimeError("LLaDA attention is not attached to an owning block") + config = self.config + if config is None: + raise RuntimeError("LLaDA attention is not attached to a bridge config") + + batch_size, query_len, q_width = q.shape + n_heads = int(config.n_heads) + n_kv_heads = int(getattr(config, "n_key_value_heads", None) or n_heads) + head_dim = q_width // n_heads + input_dtype = k.dtype + + q_norm = getattr(block, "q_norm", None) + k_norm = getattr(block, "k_norm", None) + if q_norm is not None and k_norm is not None: + q = q_norm(q).to(dtype=input_dtype) + k = k_norm(k).to(dtype=input_dtype) + + q = q.view(batch_size, query_len, n_heads, head_dim).transpose(1, 2) + k = k.view(batch_size, query_len, n_kv_heads, head_dim).transpose(1, 2) + v = v.view(batch_size, query_len, n_kv_heads, head_dim).transpose(1, 2) + + if layer_past is not None: + past_key, past_value = layer_past + k = torch.cat((past_key, k), dim=-2) + v = torch.cat((past_value, v), dim=-2) + present = (k, v) if use_cache else None + + if bool(getattr(block.config, "rope", False)): + q, k = cast(Any, block.rotary_emb)(q, k) + if hasattr(self, "hook_rot_q"): + q = self.hook_rot_q(q) + if hasattr(self, "hook_rot_k"): + k = self.hook_rot_k(k) + + key_len = k.shape[-2] + if attention_bias is not None: + attention_bias = cast(Any, block._cast_attn_bias)( + attention_bias[:, :, key_len - query_len : key_len, :key_len], + input_dtype, + ) + + if n_heads != n_kv_heads: + if n_heads % n_kv_heads != 0: + raise ValueError( + f"n_heads ({n_heads}) must be divisible by n_key_value_heads " f"({n_kv_heads})" + ) + groups = n_heads // n_kv_heads + k = k.repeat_interleave(groups, dim=1, output_size=n_heads) + v = v.repeat_interleave(groups, dim=1, output_size=n_heads) + + attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(head_dim) + attn_scores = self._apply_reconstruct_attention_mask( + attn_scores=attn_scores, + attention_mask=attention_bias, + seq_len=key_len, + q_seq_len=query_len, + ) + attn_scores = self.hook_attn_scores(attn_scores) + + pattern = torch.nn.functional.softmax(attn_scores, dim=-1, dtype=torch.float32).to(q.dtype) + dropout = float(getattr(block.config, "attention_dropout", 0.0)) + if block.training and dropout > 0.0: + pattern = torch.nn.functional.dropout(pattern, p=dropout, training=True) + pattern = self.hook_pattern(pattern) + + z = torch.matmul(pattern, v).transpose(1, 2).contiguous() + flat_z = z.view(batch_size, query_len, n_heads * head_dim) + if bool(getattr(self.config, "use_attn_result", False)): + flat_z = self.o.hook_in(flat_z) + z = flat_z.view(batch_size, query_len, n_heads, head_dim) + output = self._compute_per_head_result(z, n_heads, head_dim) + else: + output = self.o(flat_z) + output = self.hook_out(output) + return output, present + + +class _LLaDABlockBridge(BlockBridge): + """Preserve native block math while routing container-level hooks.""" + + def __init__( + self, + name: str, + config: Optional[Any] = None, + submodules: Optional[Dict[str, GeneralizedComponent]] = None, + ) -> None: + super().__init__( + name, + config=config, + submodules=submodules, + hook_alias_overrides={ + "hook_attn_in": "attn.hook_in", + "hook_q_input": "attn.q.hook_in", + "hook_k_input": "attn.k.hook_in", + "hook_v_input": "attn.v.hook_in", + }, + ) + self._llada_container_hooks_wired = False + self._llada_container_hook_handles: list[torch.utils.hooks.RemovableHandle] = [] + + def _route_attn_input( + self, _module: torch.nn.Module, _args: tuple[Any, ...], output: torch.Tensor + ) -> torch.Tensor: + return self.attn.hook_in(output) + + def _route_mlp_input( + self, _module: torch.nn.Module, _args: tuple[Any, ...], output: torch.Tensor + ) -> torch.Tensor: + return self.mlp.hook_in(output) + + def _route_mlp_output( + self, _module: torch.nn.Module, _args: tuple[Any, ...], output: torch.Tensor + ) -> torch.Tensor: + if bool(getattr(self.mlp, "_executing_container", False)): + return output + return self.mlp.hook_out(output) + + def _route_pre_mlp_norm( + self, _module: torch.nn.Module, args: tuple[Any, ...] + ) -> Optional[tuple[Any, ...]]: + if not self._read_use_hook_mlp_in(): + return None + if args and isinstance(args[0], torch.Tensor): + return (self.hook_mlp_in(args[0]),) + args[1:] + return None + + def _maybe_wire_pre_ln_capture(self) -> None: + """Use deepcopy-safe bound hooks for LLaDA's pre-MLP residual hook.""" + if self._pre_ln_capture_wired: + return + if self.ln2.original_component is not None: + self._pre_ln_capture_handles.append( + self.ln2.register_forward_pre_hook(self._route_pre_mlp_norm) + ) + self._pre_ln_capture_wired = True + + def _wire_llada_container_hooks(self) -> None: + if self._llada_container_hooks_wired: + return + + self._llada_container_hook_handles.extend( + ( + self.ln1.register_forward_hook(self._route_attn_input), + self.ln2.register_forward_hook(self._route_mlp_input), + self.mlp.out.register_forward_hook(self._route_mlp_output), + ) + ) + self._llada_container_hooks_wired = True + + def forward(self, *args: Any, **kwargs: Any) -> Any: + """Wire LLaDA's symbolic containers, then delegate to the native block.""" + self._wire_llada_container_hooks() + return super().forward(*args, **kwargs) + + +class _LLaDAGatedMLPBridge(GatedMLPBridge): + """Executable view over LLaDA's block-local gated MLP projections.""" + + def set_original_component(self, original_component: torch.nn.Module) -> None: + """The executable view needs the block's children, not ownership of the block.""" + del original_component + + def forward(self, hidden_states: torch.Tensor, **kwargs: Any) -> torch.Tensor: + """Run SiLU-gated MLP math using the live wrapped projections.""" + del kwargs + hidden_states = self.hook_in(hidden_states) + self._executing_container = True + try: + up_projection = getattr(self, "in") + gate = self.gate(hidden_states) + up = up_projection(hidden_states) + activated = self.act(gate) + gated = activated * up + output = self.out(gated) + finally: + self._executing_container = False + return self.hook_out(output) + + +class LLaDAArchitectureAdapter(ArchitectureAdapter): + """Adapter for the dense ``LLaDAModelLM`` architecture. + + Support is deliberately limited to the released dense LLaDA block contract: + Llama-style blocks, RMSNorm, separate bias-free projections, RoPE, + bidirectional attention, an untied LM head, and no KV cache. The external + iterative denoising/remasking loop is not a TransformerBridge generation API. + Loading the Hugging Face checkpoint requires the caller to opt in with + ``trust_remote_code=True``. + """ + + applicable_phases: list[int] = [] + supports_generation: bool = False + supports_hf_output_attentions: bool = False + supports_causal_loss: bool = False + + def __init__(self, cfg: Any) -> None: + super().__init__(cfg) + self._validate_config() + + self.cfg.d_vocab_out = self.cfg.d_vocab + self.cfg.normalization_type = "RMS" + self.cfg.uses_rms_norm = True + self.cfg.positional_embedding_type = "rotary" + self.cfg.rotary_adjacent_pairs = False + self.cfg.attention_dir = "bidirectional" + self.cfg.final_rms = True + self.cfg.gated_mlp = True + self.cfg.attn_only = False + self.cfg.default_prepend_bos = False + self.cfg.default_padding_side = "right" + + self.weight_processing_conversions = { + **self._qkvo_weight_conversions(), + } + + self.component_mapping = { + "embed": EmbeddingBridge(name="model.transformer.wte", config=self.cfg), + "blocks": _LLaDABlockBridge( + name="model.transformer.blocks", + config=self.cfg, + submodules={ + "ln1": RMSNormalizationBridge(name="attn_norm", config=self.cfg), + "attn": _LLaDAAttentionBridge( + name=None, + config=self.cfg, + is_causal=False, + submodules={ + "q": LinearBridge(name="q_proj"), + "k": LinearBridge(name="k_proj"), + "v": LinearBridge(name="v_proj"), + "o": LinearBridge(name="attn_out"), + }, + ), + "ln2": RMSNormalizationBridge(name="ff_norm", config=self.cfg), + "mlp": _LLaDAGatedMLPBridge( + name=None, + config=self.cfg, + submodules={ + "gate": LinearBridge(name="ff_proj"), + "in": LinearBridge(name="up_proj"), + "act": GeneralizedComponent(name="act"), + "out": LinearBridge(name="ff_out"), + }, + ), + }, + ), + "ln_final": RMSNormalizationBridge(name="model.transformer.ln_f", config=self.cfg), + "unembed": UnembeddingBridge(name="model.transformer.ff_out", config=self.cfg), + } + + @staticmethod + def _enum_value(value: Any) -> Any: + return getattr(value, "value", value) + + def _require_value(self, name: str, expected: Any) -> None: + if not hasattr(self.cfg, name): + raise ValueError(f"LLaDAModelLM config is missing required field '{name}'") + actual = self._enum_value(getattr(self.cfg, name)) + if actual != expected: + raise ValueError(f"LLaDAModelLM requires {name}={expected!r}; got {actual!r}") + + def _validate_config(self) -> None: + self._require_value("block_type", "llama") + self._require_value("block_group_size", 1) + self._require_value("rope", True) + self._require_value("rope_full_precision", True) + self._require_value("alibi", False) + self._require_value("attention_layer_norm", False) + self._require_value("include_bias", False) + self._require_value("include_qkv_bias", False) + self._require_value("scale_logits", False) + self._require_value("input_emb_norm", False) + self._require_value("layer_norm_type", "rms") + if self.cfg.act_fn != "silu": + raise ValueError( + f"LLaDAModelLM requires activation_type='silu'; got {self.cfg.act_fn!r}" + ) + if bool(self.cfg.tie_word_embeddings): + raise ValueError( + "LLaDAModelLM tied embeddings are not supported by the initial dense adapter" + ) + embedding_size = int(getattr(self.cfg, "embedding_size", self.cfg.d_vocab)) + if embedding_size != self.cfg.d_vocab: + raise ValueError( + "LLaDAModelLM requires embedding_size == vocab_size in the initial " + f"dense adapter; got {embedding_size} and {self.cfg.d_vocab}" + ) + + def prepare_loading(self, model_name: str, model_kwargs: dict) -> None: + """Disable remote branches incompatible with single-pass hook support.""" + del model_name + config = model_kwargs.get("config") + if config is not None: + config.output_attentions = False + config.use_cache = False + + def prepare_model(self, hf_model: Any) -> None: + """Keep the wrapper and underlying model on the no-cache path.""" + for config in ( + getattr(hf_model, "config", None), + getattr(getattr(hf_model, "model", None), "config", None), + ): + if config is not None: + config.output_attentions = False + config.use_cache = False diff --git a/transformer_lens/tools/model_registry/__init__.py b/transformer_lens/tools/model_registry/__init__.py index 6d2a2927ac..e3b0aa787c 100644 --- a/transformer_lens/tools/model_registry/__init__.py +++ b/transformer_lens/tools/model_registry/__init__.py @@ -79,6 +79,7 @@ "HubertModel", "HunYuanDenseV1ForCausalLM", "InternLM2ForCausalLM", + "LLaDAModelLM", "LlamaForCausalLM", "LlavaForConditionalGeneration", "LlavaNextForConditionalGeneration", @@ -149,6 +150,7 @@ "HubertModel": ["facebook"], "HunYuanDenseV1ForCausalLM": ["tencent"], "InternLM2ForCausalLM": ["internlm"], + "LLaDAModelLM": ["GSAI-ML"], "LlamaForCausalLM": ["meta-llama", "huggyllama", "codellama", "SimpleStories"], "LlavaForConditionalGeneration": ["llava-hf"], "LlavaNextForConditionalGeneration": ["llava-hf"], diff --git a/transformer_lens/tools/model_registry/data/supported_models.json b/transformer_lens/tools/model_registry/data/supported_models.json index 6074ca7cd3..1c0facecac 100644 --- a/transformer_lens/tools/model_registry/data/supported_models.json +++ b/transformer_lens/tools/model_registry/data/supported_models.json @@ -6,8 +6,8 @@ "min_downloads": 500, "scan_duration_seconds": 8.1 }, - "total_architectures": 65, - "total_models": 12899, + "total_architectures": 66, + "total_models": 12900, "total_verified": 1017, "models": [ { @@ -179193,6 +179193,20 @@ "phase7_score": null, "phase8_score": null }, + { + "architecture_id": "LLaDAModelLM", + "model_id": "GSAI-ML/LLaDA-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": "Glm4MoeForCausalLM", "model_id": "zai-org/GLM-4.5-Base", diff --git a/transformer_lens/tools/model_registry/generate_report.py b/transformer_lens/tools/model_registry/generate_report.py index 73f0aa8b18..44a928d63b 100644 --- a/transformer_lens/tools/model_registry/generate_report.py +++ b/transformer_lens/tools/model_registry/generate_report.py @@ -31,6 +31,7 @@ "GPTNeoXForCausalLM": "EleutherAI's GPT-NeoX architecture used in Pythia models", "GPTJForCausalLM": "EleutherAI's GPT-J 6B parameter model", "LlamaForCausalLM": "Meta's LLaMA architecture, basis for many open models", + "LLaDAModelLM": "GSAI's bidirectional masked-diffusion language model (single-forward support)", "Lfm2MoeForCausalLM": "LiquidAI's LFM2 hybrid convolution/attention Mixture of Experts model", "MistralForCausalLM": "Mistral AI's efficient 7B parameter model with sliding window attention", "MixtralForCausalLM": "Mistral AI's Mixture of Experts model", diff --git a/transformer_lens/tools/model_registry/verify_models.py b/transformer_lens/tools/model_registry/verify_models.py index 9b939fe7cf..f0e0c53db8 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 + "GSAI-ML/LLaDA", # LLaDAModelLM — ships configuration_llada.py/modeling_llada.py ) # Data directory for registry files