From 5e6d97fffbed6fb14cee937dea08797eadca6b1d Mon Sep 17 00:00:00 2001 From: Canonik Date: Sun, 11 Oct 2026 13:05:33 +0200 Subject: [PATCH 1/6] feat(bridge): RoBERTa masked-LM adapter RobertaForMaskedLM is BertForMaskedLM with the trunk under roberta.* and the MLM head under lm_head.*, so the adapter subclasses BertArchitectureAdapter, calls super().__init__ and rebuilds component_mapping with the new paths. token_type_embed stays (a real [1, d_model] table even at type_vocab_size 1) and so do the hook_mlp_in/out redirects (RobertaLayer keeps BertLayer's sibling intermediate/output layout). The inherited prepare_model finds no bert.pooler or cls.* heads and is a no-op. n_ctx is max_position_embeddings minus the two rows HF reserves for position ids that start at padding_idx + 1, which is what makes the W_pos accessor slice them off. Subtracting instead of pinning 512 keeps longer-table checkpoints right too. Wired into the factory, the registry sets and the report descriptions. No tokenizer flags on the adapter: / handling comes from each checkpoint's tokenizer at boot. --- .../factories/architecture_adapter_factory.py | 2 + .../supported_architectures/__init__.py | 4 + .../supported_architectures/roberta.py | 119 ++++++++++++++++++ .../tools/model_registry/__init__.py | 2 + .../tools/model_registry/generate_report.py | 2 +- 5 files changed, 128 insertions(+), 1 deletion(-) create mode 100644 transformer_lens/model_bridge/supported_architectures/roberta.py diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py index 5c0f0a080..02a44f9dd 100644 --- a/transformer_lens/factories/architecture_adapter_factory.py +++ b/transformer_lens/factories/architecture_adapter_factory.py @@ -138,6 +138,7 @@ QwenArchitectureAdapter, RavenArchitectureAdapter, RecurrentGemmaArchitectureAdapter, + RobertaArchitectureAdapter, RWKV7ArchitectureAdapter, RwkvArchitectureAdapter, SeedOssArchitectureAdapter, @@ -189,6 +190,7 @@ "SwitchTransformersForConditionalGeneration": SwitchTransformersArchitectureAdapter, "BertForMaskedLM": BertArchitectureAdapter, "BertLMHeadModel": BertArchitectureAdapter, + "RobertaForMaskedLM": RobertaArchitectureAdapter, "BitNetForCausalLM": BitNetArchitectureAdapter, "BlenderbotForConditionalGeneration": BlenderbotArchitectureAdapter, "BloomForCausalLM": BloomArchitectureAdapter, diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py index 0d952dd1f..95ea53873 100644 --- a/transformer_lens/model_bridge/supported_architectures/__init__.py +++ b/transformer_lens/model_bridge/supported_architectures/__init__.py @@ -273,6 +273,9 @@ from transformer_lens.model_bridge.supported_architectures.recurrent_gemma import ( RecurrentGemmaArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.roberta import ( + RobertaArchitectureAdapter, +) from transformer_lens.model_bridge.supported_architectures.rwkv7 import ( RWKV7ArchitectureAdapter, ) @@ -449,6 +452,7 @@ "Qwen3_5MultimodalArchitectureAdapter", "RavenArchitectureAdapter", "RecurrentGemmaArchitectureAdapter", + "RobertaArchitectureAdapter", "RWKV7ArchitectureAdapter", "SeedOssArchitectureAdapter", "SmolLM3ArchitectureAdapter", diff --git a/transformer_lens/model_bridge/supported_architectures/roberta.py b/transformer_lens/model_bridge/supported_architectures/roberta.py new file mode 100644 index 000000000..380c300f5 --- /dev/null +++ b/transformer_lens/model_bridge/supported_architectures/roberta.py @@ -0,0 +1,119 @@ +"""RoBERTa architecture adapter. + +RoBERTa reuses BERT's encoder layer and only renames what surrounds it: the +trunk lives under ``roberta.*`` instead of ``bert.*`` and the masked-LM head +under ``lm_head.*`` instead of ``cls.predictions.*``. +""" + +from typing import Any + +from transformer_lens.model_bridge.generalized_components import ( + AttentionBridge, + BlockBridge, + EmbeddingBridge, + LinearBridge, + MLPBridge, + NormalizationBridge, + PosEmbedBridge, + UnembeddingBridge, +) +from transformer_lens.model_bridge.supported_architectures.bert import ( + BertArchitectureAdapter, +) + + +class RobertaArchitectureAdapter(BertArchitectureAdapter): + """Architecture adapter for ``RobertaForMaskedLM`` checkpoints. + + Everything BERT sets up carries over unchanged: the post-LN config flags + with ``supports_fold_ln = False``, the q/k/v/o rearrange conversions and + ``supports_generation = False``. Only the HF module paths and the context + length differ. Tokenizer policy is not set here: the boot path reads the + ```` / ```` behaviour from each checkpoint's tokenizer. + + HF builds RoBERTa position ids inside ``RobertaEmbeddings`` starting at + ``padding_idx + 1 = 2`` (padding keeps id 1), so ``pos_embed.hook_out`` + holds rows ``2 .. seq_len + 1`` of the table and is not row-aligned with + BERT's. ``W_pos`` drops the two reserved rows and is aligned. + + ``BertArchitectureAdapter.prepare_model`` only reacts to ``bert.pooler`` + and ``cls.*`` task heads, which ``RobertaForMaskedLM`` lacks, so the + inherited version is a no-op here. + """ + + def __init__(self, cfg: Any) -> None: + """Initialize the RoBERTa architecture adapter. + + Args: + cfg: The configuration object. ``n_ctx`` is expected to arrive as + HF's ``max_position_embeddings`` (the full table size). + """ + super().__init__(cfg) + + # Components are keyed by HF module path, so BERT's mapping cannot be + # renamed in place: rebuild it with RoBERTa's paths. + self.component_mapping = { + "embed": EmbeddingBridge(name="roberta.embeddings.word_embeddings"), + # A real [type_vocab_size, d_model] table that HF adds on every forward, + # even though type_vocab_size is 1 and the row is a constant. + "token_type_embed": EmbeddingBridge(name="roberta.embeddings.token_type_embeddings"), + "pos_embed": PosEmbedBridge(name="roberta.embeddings.position_embeddings"), + "embed_ln": NormalizationBridge( + name="roberta.embeddings.LayerNorm", + config=self.cfg, + use_native_layernorm_autograd=True, + ), + "blocks": BlockBridge( + name="roberta.encoder.layer", + # RobertaLayer keeps BERT's layout: intermediate.dense and output.dense + # are siblings with no MLP module, so the MLP aliases point at the two + # Linear hooks exactly as in bert.py. + hook_alias_overrides={ + "hook_mlp_out": "mlp.out.hook_out", + "hook_mlp_in": "mlp.in.hook_in", + }, + submodules={ + "ln1": NormalizationBridge( + name="attention.output.LayerNorm", + config=self.cfg, + use_native_layernorm_autograd=True, + ), + "ln2": NormalizationBridge( + name="output.LayerNorm", + config=self.cfg, + use_native_layernorm_autograd=True, + ), + "attn": AttentionBridge( + name="attention", + config=self.cfg, + submodules={ + "q": LinearBridge(name="self.query"), + "k": LinearBridge(name="self.key"), + "v": LinearBridge(name="self.value"), + "o": LinearBridge(name="output.dense"), + }, + ), + "mlp": MLPBridge( + name=None, + config=self.cfg, + submodules={ + "in": LinearBridge(name="intermediate.dense"), + "out": LinearBridge(name="output.dense"), + }, + ), + }, + ), + "mlm_head": LinearBridge(name="lm_head.dense"), + "unembed": UnembeddingBridge(name="lm_head.decoder"), + "ln_final": NormalizationBridge( + name="lm_head.layer_norm", + config=self.cfg, + use_native_layernorm_autograd=True, + ), + } + + # HF's max_position_embeddings (514 for roberta-base) counts the two + # reserved rows, so the usable context is two shorter. Subtracting rather + # than pinning 512 keeps longer-table checkpoints right too, and W_pos + # only slices the reserved rows off when the table has n_ctx + 2 of them. + self.cfg.n_ctx -= 2 diff --git a/transformer_lens/tools/model_registry/__init__.py b/transformer_lens/tools/model_registry/__init__.py index d3808b017..5388076be 100644 --- a/transformer_lens/tools/model_registry/__init__.py +++ b/transformer_lens/tools/model_registry/__init__.py @@ -181,6 +181,7 @@ "Qwen3_5MoeForConditionalGeneration", "RavenForCausalLM", "RecurrentGemmaForCausalLM", + "RobertaForMaskedLM", "SeedOssForCausalLM", "RWKV7ForCausalLM", "SmolLM3ForCausalLM", @@ -347,6 +348,7 @@ "QwenForCausalLM": ["Qwen"], "RavenForCausalLM": ["tomg-group-umd"], "RecurrentGemmaForCausalLM": ["google"], + "RobertaForMaskedLM": ["FacebookAI", "distilbert"], "RWKV7ForCausalLM": ["fla-hub"], "SeedOssForCausalLM": ["ByteDance-Seed"], "SmolLM3ForCausalLM": ["HuggingFaceTB"], diff --git a/transformer_lens/tools/model_registry/generate_report.py b/transformer_lens/tools/model_registry/generate_report.py index 3a95db69f..fa25690f6 100644 --- a/transformer_lens/tools/model_registry/generate_report.py +++ b/transformer_lens/tools/model_registry/generate_report.py @@ -93,6 +93,7 @@ "NemotronForCausalLM": "NVIDIA's dense Nemotron/Minitron decoder (LayerNorm1P, squared-ReLU MLP)", "BloomModel": "Headless Bloom checkpoints (loads as BloomForCausalLM with tied embeddings)", "BertLMHeadModel": "BERT with a causal LM head (decoder-style BERT; same module tree as BertForMaskedLM)", + "RobertaForMaskedLM": "RoBERTa with masked language modeling head", "Idefics3ForConditionalGeneration": "HuggingFace's Idefics3/SmolVLM vision-language model (granite-docling)", "Qwen2AudioForConditionalGeneration": "Alibaba's Qwen2-Audio (Whisper-style encoder + Qwen2 decoder)", "Ernie4_5ForCausalLM": "Baidu's dense ERNIE 4.5 Llama-layout decoder", @@ -145,7 +146,6 @@ "BertForMaskedLM": "BERT with masked language modeling head", "BertForSequenceClassification": "BERT fine-tuned for classification", "RobertaModel": "Facebook's RoBERTa, optimized BERT training", - "RobertaForMaskedLM": "RoBERTa with masked language modeling head", "DistilBertModel": "Distilled version of BERT, 40% smaller", "AlbertModel": "A Lite BERT with parameter sharing", "XLNetLMHeadModel": "Google/CMU's XLNet with permutation language modeling", From 09a4be7f9d63921d23c6152649494a36073f59bb Mon Sep 17 00:00:00 2001 From: Canonik Date: Fri, 16 Oct 2026 01:27:43 +0200 Subject: [PATCH 2/6] trim adapter comments per review Keep the one-liners the issue notes asked to be visible (inherited prepare_model, position offset and which W_pos slices it, token_type_embed, alias redirects, n_ctx) and drop the rest. 40 comment lines down to 14. --- .../supported_architectures/roberta.py | 50 +++++-------------- 1 file changed, 12 insertions(+), 38 deletions(-) diff --git a/transformer_lens/model_bridge/supported_architectures/roberta.py b/transformer_lens/model_bridge/supported_architectures/roberta.py index 380c300f5..e2b635c91 100644 --- a/transformer_lens/model_bridge/supported_architectures/roberta.py +++ b/transformer_lens/model_bridge/supported_architectures/roberta.py @@ -1,9 +1,4 @@ -"""RoBERTa architecture adapter. - -RoBERTa reuses BERT's encoder layer and only renames what surrounds it: the -trunk lives under ``roberta.*`` instead of ``bert.*`` and the masked-LM head -under ``lm_head.*`` instead of ``cls.predictions.*``. -""" +"""RoBERTa architecture adapter.""" from typing import Any @@ -23,39 +18,22 @@ class RobertaArchitectureAdapter(BertArchitectureAdapter): - """Architecture adapter for ``RobertaForMaskedLM`` checkpoints. - - Everything BERT sets up carries over unchanged: the post-LN config flags - with ``supports_fold_ln = False``, the q/k/v/o rearrange conversions and - ``supports_generation = False``. Only the HF module paths and the context - length differ. Tokenizer policy is not set here: the boot path reads the - ```` / ```` behaviour from each checkpoint's tokenizer. + """Architecture adapter for ``RobertaForMaskedLM``. - HF builds RoBERTa position ids inside ``RobertaEmbeddings`` starting at - ``padding_idx + 1 = 2`` (padding keeps id 1), so ``pos_embed.hook_out`` - holds rows ``2 .. seq_len + 1`` of the table and is not row-aligned with - BERT's. ``W_pos`` drops the two reserved rows and is aligned. - - ``BertArchitectureAdapter.prepare_model`` only reacts to ``bert.pooler`` - and ``cls.*`` task heads, which ``RobertaForMaskedLM`` lacks, so the - inherited version is a no-op here. + BERT's encoder layer under ``roberta.*`` with the MLM head under ``lm_head.*``; + flags, weight conversions and ``prepare_model`` (a no-op here: no pooler, no + ``cls.*`` heads) are inherited. Position ids start at ``padding_idx + 1 = 2``, + so ``pos_embed.hook_out`` is offset by two rows from BERT's; + ``TransformerBridge.W_pos`` slices those rows off. """ def __init__(self, cfg: Any) -> None: - """Initialize the RoBERTa architecture adapter. - - Args: - cfg: The configuration object. ``n_ctx`` is expected to arrive as - HF's ``max_position_embeddings`` (the full table size). - """ + """Initialize the RoBERTa architecture adapter.""" super().__init__(cfg) - # Components are keyed by HF module path, so BERT's mapping cannot be - # renamed in place: rebuild it with RoBERTa's paths. self.component_mapping = { "embed": EmbeddingBridge(name="roberta.embeddings.word_embeddings"), - # A real [type_vocab_size, d_model] table that HF adds on every forward, - # even though type_vocab_size is 1 and the row is a constant. + # Real [1, d_model] table at type_vocab_size = 1; HF adds it on every forward. "token_type_embed": EmbeddingBridge(name="roberta.embeddings.token_type_embeddings"), "pos_embed": PosEmbedBridge(name="roberta.embeddings.position_embeddings"), "embed_ln": NormalizationBridge( @@ -65,9 +43,7 @@ def __init__(self, cfg: Any) -> None: ), "blocks": BlockBridge( name="roberta.encoder.layer", - # RobertaLayer keeps BERT's layout: intermediate.dense and output.dense - # are siblings with no MLP module, so the MLP aliases point at the two - # Linear hooks exactly as in bert.py. + # Same split intermediate.dense / output.dense layout as BertLayer, see bert.py. hook_alias_overrides={ "hook_mlp_out": "mlp.out.hook_out", "hook_mlp_in": "mlp.in.hook_in", @@ -112,8 +88,6 @@ def __init__(self, cfg: Any) -> None: ), } - # HF's max_position_embeddings (514 for roberta-base) counts the two - # reserved rows, so the usable context is two shorter. Subtracting rather - # than pinning 512 keeps longer-table checkpoints right too, and W_pos - # only slices the reserved rows off when the table has n_ctx + 2 of them. + # max_position_embeddings counts the two reserved rows; a table of n_ctx + 2 + # rows is what makes TransformerBridge.W_pos drop them. self.cfg.n_ctx -= 2 From 9446af409d4d73dfbcc5771638309ae0e188d4bb Mon Sep 17 00:00:00 2001 From: Canonik Date: Sun, 11 Oct 2026 13:05:33 +0200 Subject: [PATCH 3/6] registry: add roberta-base and distilroberta-base, both verified verify_models on FacebookAI/roberta-base and distilbert/distilroberta-base: P1/P2/P3 at 100%, P4 not applicable to a masked LM, same as bert-base-cased. Entries, header counts and the verification_history records are what the tool wrote. --- .../model_registry/data/supported_models.json | 36 +++++++++++++++++-- .../data/verification_history.json | 26 +++++++++++++- 2 files changed, 58 insertions(+), 4 deletions(-) diff --git a/transformer_lens/tools/model_registry/data/supported_models.json b/transformer_lens/tools/model_registry/data/supported_models.json index 55962e92b..2f26a8f2d 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": 0.1 }, - "total_architectures": 145, - "total_models": 15734, - "total_verified": 1196, + "total_architectures": 146, + "total_models": 15736, + "total_verified": 1198, "total_provisional": 7, "models": [ { @@ -221545,6 +221545,36 @@ "phase4_score": null, "phase7_score": null, "phase8_score": null + }, + { + "architecture_id": "RobertaForMaskedLM", + "model_id": "FacebookAI/roberta-base", + "status": 1, + "verified_date": "2026-10-10", + "metadata": null, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null + }, + { + "architecture_id": "RobertaForMaskedLM", + "model_id": "distilbert/distilroberta-base", + "status": 1, + "verified_date": "2026-10-10", + "metadata": null, + "note": "Full verification completed", + "phase1_score": 100.0, + "phase2_score": 100.0, + "phase3_score": 100.0, + "phase4_score": null, + "phase7_score": null, + "phase8_score": null, + "phase9_score": null } ] } diff --git a/transformer_lens/tools/model_registry/data/verification_history.json b/transformer_lens/tools/model_registry/data/verification_history.json index bec8a8227..8d0c82b8c 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-08-31T10:41:22.140312", + "last_updated": "2026-10-10T23:22:28.618265", "records": [ { "model_id": "Macropodus/macbert4mdcspell_v1", @@ -41424,6 +41424,30 @@ "notes": null, "invalidated": false, "invalidation_reason": null + }, + { + "model_id": "distilbert/distilroberta-base", + "architecture_id": "RobertaForMaskedLM", + "verified_date": "2026-10-10", + "verified_by": "verify_models", + "transformerlens_version": "0.0.0", + "prompt_profile": null, + "p4_scoring_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null + }, + { + "model_id": "FacebookAI/roberta-base", + "architecture_id": "RobertaForMaskedLM", + "verified_date": "2026-10-10", + "verified_by": "verify_models", + "transformerlens_version": "0.0.0", + "prompt_profile": null, + "p4_scoring_version": null, + "notes": "Full verification completed", + "invalidated": false, + "invalidation_reason": null } ] } From c019f51f43e0bf5020dcafddac3502397cfe9906 Mon Sep 17 00:00:00 2001 From: Canonik Date: Sun, 11 Oct 2026 13:05:34 +0200 Subject: [PATCH 4/6] test(roberta): adapter unit test and tiny from_config coverage The unit test pins the mapping against BertArchitectureAdapter (keys, bridge types, block layout, weight conversions) instead of restating it, so a BERT-side change surfaces here too, plus the RoBERTa paths, the n_ctx adjustment, the no-op prepare_model, untouched tokenizer flags and the factory routing. test_roberta_tiny.py is the CI tier: a random-weight two-layer RobertaForMaskedLM, live wiring and tied decoder weight, W_pos / W_E_pos shapes, hook capture against the HF modules including the redirected hook_mlp_in/out aliases and ln_final, the position rows HF actually looks up (2.. for content, 1 for padding), padded and unpadded parity, and a block-output intervention. --- .../model_bridge/test_roberta_tiny.py | 245 ++++++++++++++++++ .../test_roberta_adapter.py | 158 +++++++++++ 2 files changed, 403 insertions(+) create mode 100644 tests/integration/model_bridge/test_roberta_tiny.py create mode 100644 tests/unit/model_bridge/supported_architectures/test_roberta_adapter.py diff --git a/tests/integration/model_bridge/test_roberta_tiny.py b/tests/integration/model_bridge/test_roberta_tiny.py new file mode 100644 index 000000000..8f2d0615c --- /dev/null +++ b/tests/integration/model_bridge/test_roberta_tiny.py @@ -0,0 +1,245 @@ +"""Tiny from_config RoBERTa tests: CI stand-in for the slow real-weights parity file.""" + +import copy +from dataclasses import dataclass + +import pytest +import torch +import torch.nn.functional as F +from transformers import RobertaConfig, RobertaForMaskedLM + +from transformer_lens.model_bridge import TransformerBridge +from transformer_lens.model_bridge.sources import build_bridge_from_module + +D_MODEL = 16 +D_HEAD = 4 +N_HEADS = D_MODEL // D_HEAD +D_MLP = 32 +N_LAYERS = 2 +N_CTX = 8 +D_VOCAB = 37 +PAD_ID = 1 +PARITY_ATOL = 1e-5 +CACHE_NAMES = ( + "embed.hook_out", + "token_type_embed.hook_out", + "pos_embed.hook_out", + "blocks.0.hook_in", + "blocks.0.attn.q.hook_out", + "blocks.0.hook_mlp_in", + "blocks.0.hook_mlp_out", + "ln_final.hook_normalized", + "unembed.hook_out", +) + + +@dataclass +class TinyRobertaCase: + bridge: TransformerBridge + reference: RobertaForMaskedLM + + +@pytest.fixture(scope="module") +def case() -> TinyRobertaCase: + torch.manual_seed(0) + config = RobertaConfig( + vocab_size=D_VOCAB, + hidden_size=D_MODEL, + num_hidden_layers=N_LAYERS, + num_attention_heads=N_HEADS, + intermediate_size=D_MLP, + # HF counts the two reserved rows in the table size. + max_position_embeddings=N_CTX + 2, + type_vocab_size=1, + pad_token_id=PAD_ID, + bos_token_id=0, + eos_token_id=2, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + ) + reference = RobertaForMaskedLM(config).eval() + source = copy.deepcopy(reference).eval() + bridge = build_bridge_from_module( + source, + "RobertaForMaskedLM", + hf_config=copy.deepcopy(config), + tokenizer=None, + dtype=torch.float32, + device="cpu", + ) + bridge.eval() + return TinyRobertaCase(bridge=bridge, reference=reference) + + +@pytest.fixture(scope="module") +def tokens() -> torch.Tensor: + return torch.tensor([[0, 5, 6, 7, 2]]) + + +@pytest.fixture(scope="module") +def padded_tokens() -> torch.Tensor: + return torch.tensor( + [ + [0, 5, 6, 2, PAD_ID, PAD_ID], + [0, 5, 6, 7, 8, 2], + ] + ) + + +def test_tiny_roberta_forward_matches_independent_reference( + case: TinyRobertaCase, + tokens: torch.Tensor, +) -> None: + with torch.inference_mode(): + bridge_logits = case.bridge(tokens) + reference_logits = case.reference(tokens).logits + + torch.testing.assert_close(bridge_logits, reference_logits, rtol=0.0, atol=PARITY_ATOL) + + +def test_tiny_roberta_live_wiring_and_tied_weights(case: TinyRobertaCase) -> None: + bridge = case.bridge + hf_model = bridge.original_model + + assert bridge.cfg.architecture == "RobertaForMaskedLM" + assert hf_model.roberta.embeddings.word_embeddings is bridge.embed + assert hf_model.roberta.embeddings.token_type_embeddings is bridge.token_type_embed + assert hf_model.roberta.embeddings.position_embeddings is bridge.pos_embed + assert hf_model.roberta.encoder.layer[0].attention.self.query is bridge.blocks[0].attn.q + assert hf_model.lm_head.decoder is bridge.unembed + assert bridge.token_type_embed.original_component.weight.shape == (1, D_MODEL) + assert bridge.unembed.original_component.weight is bridge.embed.original_component.weight + + +def test_tiny_roberta_reports_usable_context_and_aligned_w_pos(case: TinyRobertaCase) -> None: + bridge = case.bridge + table = bridge.pos_embed.original_component.weight + + assert bridge.original_model.config.max_position_embeddings == N_CTX + 2 + assert bridge.cfg.n_ctx == N_CTX + assert bridge.pos_embed.W_pos.shape == (N_CTX + 2, D_MODEL) + assert bridge.W_pos.shape == (N_CTX, D_MODEL) + torch.testing.assert_close(bridge.W_pos, table[2:], rtol=0.0, atol=0.0) + assert bridge.W_E_pos.shape == (D_VOCAB + N_CTX, D_MODEL) + + +def test_tiny_roberta_cache_matches_reference_activations( + case: TinyRobertaCase, + tokens: torch.Tensor, +) -> None: + bridge = case.bridge + reference = case.reference + + with torch.inference_mode(): + bridge_logits, cache = bridge.run_with_cache(tokens, names_filter=list(CACHE_NAMES)) + + embeddings = reference.roberta.embeddings + token_type_ids = torch.zeros_like(tokens) + position_ids = embeddings.create_position_ids_from_input_ids( + tokens, + embeddings.padding_idx, + ) + word_embeddings = embeddings.word_embeddings(tokens) + token_type_embeddings = embeddings.token_type_embeddings(token_type_ids) + position_embeddings = embeddings.position_embeddings(position_ids) + block_input = embeddings(input_ids=tokens) + query = ( + reference.roberta.encoder.layer[0] + .attention.self.query(block_input) + .reshape(tokens.shape[0], tokens.shape[1], N_HEADS, D_HEAD) + ) + layer0 = reference.roberta.encoder.layer[0] + attn_out = layer0.attention(block_input) + mlp_in = attn_out[0] if isinstance(attn_out, tuple) else attn_out + mlp_out = layer0.output.dense(layer0.intermediate(mlp_in)) + hidden = reference.roberta(input_ids=tokens).last_hidden_state + pre_layer_norm = F.gelu(reference.lm_head.dense(hidden)) + normalized = F.layer_norm( + pre_layer_norm, + (pre_layer_norm.shape[-1],), + eps=reference.lm_head.layer_norm.eps, + ) + reference_logits = reference.lm_head(hidden) + + # Filtering by an alias caches it next to the canonical hook it redirects to. + assert set(cache.keys()) - set(CACHE_NAMES) == { + "blocks.0.mlp.in.hook_in", + "blocks.0.mlp.out.hook_out", + } + expected = { + "embed.hook_out": word_embeddings, + "token_type_embed.hook_out": token_type_embeddings, + "pos_embed.hook_out": position_embeddings, + "blocks.0.hook_in": block_input, + "blocks.0.attn.q.hook_out": query, + "blocks.0.hook_mlp_in": mlp_in, + "blocks.0.hook_mlp_out": mlp_out, + "ln_final.hook_normalized": normalized, + "unembed.hook_out": reference_logits, + } + for name, activation in expected.items(): + torch.testing.assert_close(cache[name], activation, rtol=0.0, atol=PARITY_ATOL) + torch.testing.assert_close(bridge_logits, reference_logits, rtol=0.0, atol=PARITY_ATOL) + + +def test_tiny_roberta_pos_embed_hook_starts_at_row_two( + case: TinyRobertaCase, + tokens: torch.Tensor, +) -> None: + table = case.bridge.pos_embed.original_component.weight + seq_len = tokens.shape[1] + + with torch.inference_mode(): + _, cache = case.bridge.run_with_cache(tokens, names_filter=["pos_embed.hook_out"]) + + torch.testing.assert_close( + cache["pos_embed.hook_out"][0], + table[2 : 2 + seq_len], + rtol=0.0, + atol=0.0, + ) + + +def test_tiny_roberta_padded_batch_matches_reference( + case: TinyRobertaCase, + padded_tokens: torch.Tensor, +) -> None: + bridge = case.bridge + reference = case.reference + attention_mask = padded_tokens.ne(PAD_ID).long() + table = bridge.pos_embed.original_component.weight + + with torch.inference_mode(): + bridge_logits, cache = bridge.run_with_cache( + padded_tokens, + attention_mask=attention_mask, + names_filter=["pos_embed.hook_out"], + ) + unmasked_logits = bridge(padded_tokens) + reference_logits = reference(padded_tokens, attention_mask=attention_mask).logits + + torch.testing.assert_close(bridge_logits, reference_logits, rtol=0.0, atol=PARITY_ATOL) + assert (bridge_logits - unmasked_logits).abs().max().item() > PARITY_ATOL + pad_rows = cache["pos_embed.hook_out"][0, 4:] + torch.testing.assert_close(pad_rows, table[PAD_ID].expand(2, -1), rtol=0.0, atol=0.0) + torch.testing.assert_close(cache["pos_embed.hook_out"][0, :4], table[2:6], rtol=0.0, atol=0.0) + + +def test_tiny_roberta_block_hook_intervenes(case: TinyRobertaCase, tokens: torch.Tensor) -> None: + fired: list[str] = [] + + def zero_block_output(value: torch.Tensor, hook) -> torch.Tensor: + fired.append(hook.name) + return torch.zeros_like(value) + + with torch.inference_mode(): + baseline = case.bridge(tokens) + intervened = case.bridge.run_with_hooks( + tokens, + fwd_hooks=[("blocks.0.hook_out", zero_block_output)], + ) + after = case.bridge(tokens) + + assert fired == ["blocks.0.hook_out"] + assert not torch.allclose(intervened, baseline, rtol=0.0, atol=PARITY_ATOL) + torch.testing.assert_close(after, baseline, rtol=0.0, atol=0.0) diff --git a/tests/unit/model_bridge/supported_architectures/test_roberta_adapter.py b/tests/unit/model_bridge/supported_architectures/test_roberta_adapter.py new file mode 100644 index 000000000..d9a1987ea --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_roberta_adapter.py @@ -0,0 +1,158 @@ +"""Unit tests for RobertaArchitectureAdapter.""" + +from types import SimpleNamespace +from typing import Any + +import pytest + +from transformer_lens.config.transformer_bridge_config import TransformerBridgeConfig +from transformer_lens.model_bridge.generalized_components import ( + BlockBridge, + EmbeddingBridge, +) +from transformer_lens.model_bridge.supported_architectures.bert import ( + BertArchitectureAdapter, +) +from transformer_lens.model_bridge.supported_architectures.roberta import ( + RobertaArchitectureAdapter, +) + + +def _make_cfg( + architecture: str = "RobertaForMaskedLM", + n_ctx: int = 514, + **overrides: Any, +) -> TransformerBridgeConfig: + """n_ctx mimics HF's max_position_embeddings.""" + cfg = TransformerBridgeConfig( + d_model=768, + d_head=64, + n_heads=12, + n_layers=12, + n_ctx=n_ctx, + d_vocab=50265, + architecture=architecture, + ) + for name, value in overrides.items(): + setattr(cfg, name, value) + return cfg + + +def _layout(component: Any) -> dict[str, Any]: + """Nested (bridge type, HF path, native-autograd flag) layout, compared against BERT's.""" + return { + key: ( + type(sub).__name__, + sub.name, + getattr(sub, "use_native_layernorm_autograd", None), + _layout(sub), + ) + for key, sub in component.submodules.items() + } + + +@pytest.fixture(scope="module") +def adapter() -> RobertaArchitectureAdapter: + return RobertaArchitectureAdapter(_make_cfg()) + + +@pytest.fixture(scope="module") +def bert_adapter() -> BertArchitectureAdapter: + return BertArchitectureAdapter(_make_cfg(architecture="BertForMaskedLM", n_ctx=512)) + + +class TestRobertaComponentMapping: + def test_top_level_keys_match_bert( + self, adapter: RobertaArchitectureAdapter, bert_adapter: BertArchitectureAdapter + ) -> None: + assert set(adapter.component_mapping) == set(bert_adapter.component_mapping) + + def test_top_level_bridge_types_match_bert( + self, adapter: RobertaArchitectureAdapter, bert_adapter: BertArchitectureAdapter + ) -> None: + for key, component in adapter.component_mapping.items(): + bert_component = bert_adapter.component_mapping[key] + assert type(component) is type(bert_component), key + assert getattr(component, "use_native_layernorm_autograd", None) == getattr( + bert_component, "use_native_layernorm_autograd", None + ), key + + def test_top_level_hf_paths(self, adapter: RobertaArchitectureAdapter) -> None: + mapping = adapter.component_mapping + assert mapping["embed"].name == "roberta.embeddings.word_embeddings" + assert mapping["token_type_embed"].name == "roberta.embeddings.token_type_embeddings" + assert mapping["pos_embed"].name == "roberta.embeddings.position_embeddings" + assert mapping["embed_ln"].name == "roberta.embeddings.LayerNorm" + assert mapping["blocks"].name == "roberta.encoder.layer" + assert mapping["mlm_head"].name == "lm_head.dense" + assert mapping["ln_final"].name == "lm_head.layer_norm" + assert mapping["unembed"].name == "lm_head.decoder" + + def test_keeps_token_type_embedding(self, adapter: RobertaArchitectureAdapter) -> None: + token_type_embed = adapter.component_mapping["token_type_embed"] + assert isinstance(token_type_embed, EmbeddingBridge) + + def test_block_layout_matches_bert( + self, adapter: RobertaArchitectureAdapter, bert_adapter: BertArchitectureAdapter + ) -> None: + blocks = adapter.component_mapping["blocks"] + assert isinstance(blocks, BlockBridge) + assert _layout(blocks) == _layout(bert_adapter.component_mapping["blocks"]) + + def test_preserves_bert_block_hook_aliases(self, adapter: RobertaArchitectureAdapter) -> None: + blocks = adapter.component_mapping["blocks"] + assert blocks.hook_aliases["hook_mlp_in"] == "mlp.in.hook_in" + assert blocks.hook_aliases["hook_mlp_out"] == "mlp.out.hook_out" + + +class TestRobertaAdapterConfig: + def test_context_length_drops_the_two_reserved_rows(self) -> None: + assert RobertaArchitectureAdapter(_make_cfg(n_ctx=514)).cfg.n_ctx == 512 + assert RobertaArchitectureAdapter(_make_cfg(n_ctx=1026)).cfg.n_ctx == 1024 + + def test_inherits_bert_post_ln_behavior(self, adapter: RobertaArchitectureAdapter) -> None: + assert adapter.cfg.normalization_type == "LN" + assert adapter.cfg.positional_embedding_type == "standard" + assert adapter.cfg.final_rms is False + assert adapter.supports_fold_ln is False + + def test_inherits_bert_weight_conversions( + self, adapter: RobertaArchitectureAdapter, bert_adapter: BertArchitectureAdapter + ) -> None: + assert set(adapter.weight_processing_conversions) == set( + bert_adapter.weight_processing_conversions + ) + + @pytest.mark.parametrize("prepend_bos", [True, False]) + def test_does_not_override_tokenizer_policy(self, prepend_bos: bool) -> None: + adapter = RobertaArchitectureAdapter( + _make_cfg(default_prepend_bos=prepend_bos, default_padding_side="left") + ) + assert adapter.cfg.default_prepend_bos is prepend_bos + assert adapter.cfg.default_padding_side == "left" + + def test_generation_is_not_supported(self) -> None: + assert RobertaArchitectureAdapter.supports_generation is False + + +class TestRobertaInheritedPrepareModel: + def test_prepare_model_is_a_no_op(self) -> None: + adapter = RobertaArchitectureAdapter(_make_cfg()) + mapping_before = dict(adapter.components) + hf_model = SimpleNamespace( + roberta=SimpleNamespace(pooler=None), + lm_head=SimpleNamespace(), + ) + + adapter.prepare_model(hf_model) + + assert adapter.components == mapping_before + + +class TestRobertaFactoryRouting: + def test_roberta_for_masked_lm_maps_to_roberta_adapter(self) -> None: + from transformer_lens.factories.architecture_adapter_factory import ( + SUPPORTED_ARCHITECTURES, + ) + + assert SUPPORTED_ARCHITECTURES["RobertaForMaskedLM"] is RobertaArchitectureAdapter From eb97db31eb2a8670f5b62b3ae8a688938b0001f6 Mon Sep 17 00:00:00 2001 From: Canonik Date: Sun, 11 Oct 2026 13:05:34 +0200 Subject: [PATCH 5/6] test(roberta): real-weights parity on distilroberta-base Follows test_bridge_vs_hf_bert_parity.py against an independent eager HF reference, with the padding-mask case on pad id 1 and the boot-time / detection pinned. Slow-marked since the checkpoint is not in the CI cache, quarantine row added. 6/6 locally. --- tests/QUARANTINES.md | 1 + .../test_bridge_vs_hf_roberta_parity.py | 180 ++++++++++++++++++ 2 files changed, 181 insertions(+) create mode 100644 tests/integration/model_bridge/test_bridge_vs_hf_roberta_parity.py diff --git a/tests/QUARANTINES.md b/tests/QUARANTINES.md index c49bc3ffd..b257424e0 100644 --- a/tests/QUARANTINES.md +++ b/tests/QUARANTINES.md @@ -98,6 +98,7 @@ Big-model adapter tests use `@pytest.mark.slow`, CI tier filters `-m "not slow"` | [`integration/model_bridge/test_ouro_adapter.py`:35](integration/model_bridge/test_ouro_adapter.py) (skipif + `slow`) | ByteDance/Ouro-1.4B: 2.8GB download + ~11GB RAM | | [`integration/model_bridge/test_raven_adapter.py`:42](integration/model_bridge/test_raven_adapter.py) (skipif + `slow`) | huginn-0125: ~14GB download + ~28GB RAM | | [`integration/model_bridge/test_rwkv7_adapter.py`:36](integration/model_bridge/test_rwkv7_adapter.py) (skipif + `slow`; `importorskip("fla")` at :30) | rwkv7-0.1B-g1: remote code + flash-linear-attention dep | +| [`integration/model_bridge/test_bridge_vs_hf_roberta_parity.py`:23](integration/model_bridge/test_bridge_vs_hf_roberta_parity.py) | distilroberta-base 82M; not in the CI cache allowlist (test_roberta_tiny.py keeps the same surfaces in CI on random weights) | | [`integration/test_jacobian_lens_kurtosis.py`](integration/test_jacobian_lens_kurtosis.py) | Qwen3.5-0.8B (~1.6GB) + lens artifacts + wikitext validation split; CPU-feasible (#1539 Tier-1) | **Un-skip:** locally with `HF_TOKEN` sourced (slow-marked files: run the file directly or `-m slow`). diff --git a/tests/integration/model_bridge/test_bridge_vs_hf_roberta_parity.py b/tests/integration/model_bridge/test_bridge_vs_hf_roberta_parity.py new file mode 100644 index 000000000..12482c16a --- /dev/null +++ b/tests/integration/model_bridge/test_bridge_vs_hf_roberta_parity.py @@ -0,0 +1,180 @@ +"""RoBERTa Bridge parity against an independent eager HF reference. Slow: not in the CI cache.""" + +import pytest +import torch +import torch.nn.functional as F +from transformers import AutoModelForMaskedLM + +from transformer_lens.model_bridge import TransformerBridge + +MODEL_NAME = "distilbert/distilroberta-base" +PAD_ID = 1 +PARITY_ATOL = 1e-5 +CACHE_NAMES = ( + "embed.hook_out", + "token_type_embed.hook_out", + "pos_embed.hook_out", + "blocks.0.hook_in", + "blocks.0.attn.q.hook_out", + "ln_final.hook_normalized", + "unembed.hook_out", +) + +pytestmark = pytest.mark.slow + + +@pytest.fixture(scope="module") +def tokens() -> torch.Tensor: + return torch.tensor([[0, 133, 812, 16, 50264, 2]]) + + +@pytest.fixture(scope="module") +def bridge(): + return TransformerBridge.boot_transformers( + MODEL_NAME, + device="cpu", + dtype=torch.float32, + ) + + +@pytest.fixture(scope="module") +def hf_eager(): + return AutoModelForMaskedLM.from_pretrained( + MODEL_NAME, + dtype=torch.float32, + attn_implementation="eager", + ).eval() + + +def test_bridge_boots_with_roberta_context_metadata(bridge) -> None: + assert bridge.cfg.architecture == "RobertaForMaskedLM" + assert bridge.cfg.tokenizer_prepends_bos is True + assert bridge.cfg.tokenizer_appends_eos is True + assert bridge.tokenizer.pad_token_id == PAD_ID + assert bridge.cfg.n_ctx == 512 + assert bridge.original_model.config.max_position_embeddings == 514 + assert bridge.W_pos.shape == (512, bridge.cfg.d_model) + torch.testing.assert_close( + bridge.W_pos, + bridge.pos_embed.original_component.weight[2:], + rtol=0.0, + atol=0.0, + ) + + +def test_bridge_logits_match_independent_hf_eager(bridge, hf_eager, tokens) -> None: + assert bridge.original_model is not hf_eager + assert next(bridge.parameters()).dtype == torch.float32 + assert next(hf_eager.parameters()).dtype == torch.float32 + + with torch.inference_mode(): + bridge_logits = bridge(tokens) + hf_logits = hf_eager(tokens).logits + + torch.testing.assert_close( + bridge_logits, + hf_logits, + rtol=0.0, + atol=PARITY_ATOL, + ) + + +def test_run_with_cache_matches_hf_activations(bridge, hf_eager, tokens) -> None: + with torch.inference_mode(): + bridge_logits, cache = bridge.run_with_cache(tokens, names_filter=list(CACHE_NAMES)) + + embeddings = hf_eager.roberta.embeddings + token_type_ids = torch.zeros_like(tokens) + position_ids = embeddings.create_position_ids_from_input_ids( + tokens, + embeddings.padding_idx, + ) + word_embeddings = embeddings.word_embeddings(tokens) + token_type_embeddings = embeddings.token_type_embeddings(token_type_ids) + position_embeddings = embeddings.position_embeddings(position_ids) + block_input = embeddings(input_ids=tokens) + query = ( + hf_eager.roberta.encoder.layer[0] + .attention.self.query(block_input) + .reshape( + tokens.shape[0], + tokens.shape[1], + bridge.cfg.n_heads, + bridge.cfg.d_head, + ) + ) + hidden = hf_eager.roberta(input_ids=tokens).last_hidden_state + pre_layer_norm = F.gelu(hf_eager.lm_head.dense(hidden)) + normalized = F.layer_norm( + pre_layer_norm, + (pre_layer_norm.shape[-1],), + eps=hf_eager.lm_head.layer_norm.eps, + ) + hf_logits = hf_eager.lm_head.decoder(hf_eager.lm_head.layer_norm(pre_layer_norm)) + + assert set(cache.keys()) == set(CACHE_NAMES) + expected = { + "embed.hook_out": word_embeddings, + "token_type_embed.hook_out": token_type_embeddings, + "pos_embed.hook_out": position_embeddings, + "blocks.0.hook_in": block_input, + "blocks.0.attn.q.hook_out": query, + "ln_final.hook_normalized": normalized, + "unembed.hook_out": hf_logits, + } + for name, activation in expected.items(): + torch.testing.assert_close(cache[name], activation, rtol=0.0, atol=PARITY_ATOL) + torch.testing.assert_close(bridge_logits, hf_logits, rtol=0.0, atol=PARITY_ATOL) + + +def test_padding_mask_logits_match_hf_eager(bridge, hf_eager) -> None: + padded_tokens = torch.tensor( + [ + [0, 133, 812, 2, PAD_ID, PAD_ID], + [0, 133, 812, 16, 50264, 2], + ] + ) + attention_mask = padded_tokens.ne(PAD_ID).long() + + with torch.inference_mode(): + bridge_logits = bridge(padded_tokens, attention_mask=attention_mask) + unmasked_logits = bridge(padded_tokens) + hf_logits = hf_eager( + input_ids=padded_tokens, + attention_mask=attention_mask, + ).logits + + torch.testing.assert_close(bridge_logits, hf_logits, rtol=0.0, atol=PARITY_ATOL) + assert (bridge_logits - unmasked_logits).abs().max().item() > PARITY_ATOL + + +def test_live_hf_tree_uses_bridge_wrappers_and_preserves_tied_weights(bridge) -> None: + hf_model = bridge.original_model + + assert hf_model.roberta.embeddings.word_embeddings is bridge.embed + assert hf_model.roberta.embeddings.token_type_embeddings is bridge.token_type_embed + assert hf_model.roberta.embeddings.position_embeddings is bridge.pos_embed + assert hf_model.roberta.encoder.layer[0].attention.self.query is bridge.blocks[0].attn.q + assert hf_model.lm_head.decoder is bridge.unembed + assert bridge.unembed.original_component.weight is bridge.embed.original_component.weight + assert bridge.token_type_embed.original_component.weight.shape == (1, bridge.cfg.d_model) + + +def test_block_output_hook_can_intervene(bridge, tokens) -> None: + fired: list[str] = [] + + def zero_block_output(value: torch.Tensor, hook) -> torch.Tensor: + fired.append(hook.name) + return torch.zeros_like(value) + + with torch.inference_mode(): + baseline = bridge(tokens) + intervened = bridge.run_with_hooks( + tokens, + fwd_hooks=[("blocks.0.hook_out", zero_block_output)], + ) + after = bridge(tokens) + + assert fired == ["blocks.0.hook_out"] + assert not torch.allclose(intervened, baseline, rtol=0.0, atol=PARITY_ATOL) + torch.testing.assert_close(after, baseline, rtol=0.0, atol=0.0) From 4c3d97ec9daeaca7199db25db98dc3289f45bc9b Mon Sep 17 00:00:00 2001 From: Canonik Date: Fri, 16 Oct 2026 02:09:37 +0200 Subject: [PATCH 6/6] fix(bridge): leave position ids to models that derive them from input ids RoBERTa builds its own pad-aware position ids inside RobertaEmbeddings (padding_idx + 1 onwards), but accepts_mask_derived_position_ids only recognised OPT-style embeddings that take the mask, so a left-padded batch got the bridge's 0-based cumsum instead and content tokens landed on the two reserved rows (max logit diff 38 vs HF on distilroberta-base; right padding was unaffected). A module exposing create_position_ids_from_input_ids now counts as owning its derivation. Regression cases on the gate, the tiny RoBERTa and the real-weights parity file. --- .../test_bridge_vs_hf_roberta_parity.py | 16 ++++++++++++++ .../model_bridge/test_roberta_tiny.py | 21 +++++++++++++++++++ .../test_position_ids_injection_gate.py | 17 +++++++++++++++ transformer_lens/utilities/position_ids.py | 6 ++++++ 4 files changed, 60 insertions(+) diff --git a/tests/integration/model_bridge/test_bridge_vs_hf_roberta_parity.py b/tests/integration/model_bridge/test_bridge_vs_hf_roberta_parity.py index 12482c16a..a2ecb300a 100644 --- a/tests/integration/model_bridge/test_bridge_vs_hf_roberta_parity.py +++ b/tests/integration/model_bridge/test_bridge_vs_hf_roberta_parity.py @@ -148,6 +148,22 @@ def test_padding_mask_logits_match_hf_eager(bridge, hf_eager) -> None: assert (bridge_logits - unmasked_logits).abs().max().item() > PARITY_ATOL +def test_left_padding_mask_logits_match_hf_eager(bridge, hf_eager) -> None: + padded_tokens = torch.tensor( + [ + [PAD_ID, PAD_ID, 0, 133, 812, 2], + [0, 133, 812, 16, 50264, 2], + ] + ) + attention_mask = padded_tokens.ne(PAD_ID).long() + + with torch.inference_mode(): + bridge_logits = bridge(padded_tokens, attention_mask=attention_mask) + hf_logits = hf_eager(input_ids=padded_tokens, attention_mask=attention_mask).logits + + torch.testing.assert_close(bridge_logits, hf_logits, rtol=0.0, atol=PARITY_ATOL) + + def test_live_hf_tree_uses_bridge_wrappers_and_preserves_tied_weights(bridge) -> None: hf_model = bridge.original_model diff --git a/tests/integration/model_bridge/test_roberta_tiny.py b/tests/integration/model_bridge/test_roberta_tiny.py index 8f2d0615c..14ca0913b 100644 --- a/tests/integration/model_bridge/test_roberta_tiny.py +++ b/tests/integration/model_bridge/test_roberta_tiny.py @@ -10,6 +10,7 @@ from transformer_lens.model_bridge import TransformerBridge from transformer_lens.model_bridge.sources import build_bridge_from_module +from transformer_lens.utilities.position_ids import accepts_mask_derived_position_ids D_MODEL = 16 D_HEAD = 4 @@ -225,6 +226,26 @@ def test_tiny_roberta_padded_batch_matches_reference( torch.testing.assert_close(cache["pos_embed.hook_out"][0, :4], table[2:6], rtol=0.0, atol=0.0) +def test_tiny_roberta_left_padded_batch_matches_reference(case: TinyRobertaCase) -> None: + left_padded = torch.tensor( + [ + [PAD_ID, PAD_ID, 0, 5, 6, 2], + [0, 5, 6, 7, 8, 2], + ] + ) + attention_mask = left_padded.ne(PAD_ID).long() + # HF derives pad-aware positions itself, so the bridge must not inject its own. + assert accepts_mask_derived_position_ids(case.bridge.original_model) is False + + with torch.inference_mode(): + bridge_logits = case.bridge(left_padded, attention_mask=attention_mask) + reference_logits = case.reference(left_padded, attention_mask=attention_mask).logits + unpadded_logits = case.reference(left_padded[:1, 2:]).logits + + torch.testing.assert_close(bridge_logits, reference_logits, rtol=0.0, atol=PARITY_ATOL) + torch.testing.assert_close(bridge_logits[0, 2:], unpadded_logits[0], rtol=0.0, atol=PARITY_ATOL) + + def test_tiny_roberta_block_hook_intervenes(case: TinyRobertaCase, tokens: torch.Tensor) -> None: fired: list[str] = [] diff --git a/tests/unit/model_bridge/test_position_ids_injection_gate.py b/tests/unit/model_bridge/test_position_ids_injection_gate.py index 57414facb..3a44e8bc4 100644 --- a/tests/unit/model_bridge/test_position_ids_injection_gate.py +++ b/tests/unit/model_bridge/test_position_ids_injection_gate.py @@ -74,6 +74,17 @@ def forward( # type: ignore[override] return super().forward(position_ids) +class _DerivesFromInputIds(nn.Module): + """Mirrors ``RobertaEmbeddings``: positions come from the input ids, pad-aware.""" + + @staticmethod + def create_position_ids_from_input_ids( + input_ids: torch.Tensor, padding_idx: int + ) -> torch.Tensor: + mask = input_ids.ne(padding_idx).int() + return (torch.cumsum(mask, dim=1) * mask).long() + padding_idx + + class _OwnsPositions(_AcceptsPositionIds): """mRoPE models compute a 3-D index here, but only while position_ids is None.""" @@ -111,6 +122,12 @@ def test_refuses_wrapper_whose_language_model_owns_positions(self) -> None: wrapper.language_model = _OwnsPositions() assert gate(_bridge_over(wrapper)) is False + def test_refuses_model_deriving_positions_from_input_ids(self) -> None: + """RoBERTa-style embeddings derive only while position_ids is None.""" + wrapper = _AcceptsPositionIds() + wrapper.embeddings = _DerivesFromInputIds() + assert gate(_bridge_over(wrapper)) is False + def test_refuses_mrope_section_in_config(self) -> None: """Config-level backstop: the section list is what makes positions 3-D.""" model = _AcceptsPositionIds() diff --git a/transformer_lens/utilities/position_ids.py b/transformer_lens/utilities/position_ids.py index a41389864..1d767ad19 100644 --- a/transformer_lens/utilities/position_ids.py +++ b/transformer_lens/utilities/position_ids.py @@ -23,6 +23,12 @@ def accepts_mask_derived_position_ids(model: nn.Module) -> bool: scaling = getattr(candidate, "rope_scaling", None) if isinstance(scaling, dict) and "mrope_section" in scaling: return False + # fairseq-style embeddings (RoBERTa family) build pad-aware positions from the + # input ids at padding_idx + 1 onwards; an injected 0-based arange would land + # content tokens on the reserved rows. + for module in model.modules(): + if callable(getattr(module, "create_position_ids_from_input_ids", None)): + return False # A positional embedding that takes the mask derives positions for itself; only # embeddings that override nn.Embedding.forward are worth inspecting. for module in model.modules():