Coverage for transformer_lens/model_bridge/supported_architectures/bert.py: 100%
29 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""BERT architecture adapter.
3This module provides the architecture adapter for BERT models.
4"""
6from typing import Any
8from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion
9from transformer_lens.conversion_utils.param_processing_conversion import (
10 ParamProcessingConversion,
11)
12from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
13from transformer_lens.model_bridge.generalized_components import (
14 AttentionBridge,
15 BlockBridge,
16 EmbeddingBridge,
17 LinearBridge,
18 MLPBridge,
19 NormalizationBridge,
20 PosEmbedBridge,
21 UnembeddingBridge,
22)
25class BertArchitectureAdapter(ArchitectureAdapter):
26 """Architecture adapter for BERT models."""
28 supports_generation: bool = False
30 def __init__(self, cfg: Any) -> None:
31 """Initialize the BERT architecture adapter.
33 Args:
34 cfg: The configuration object.
35 """
36 super().__init__(cfg)
38 # Set config variables for weight processing
39 self.cfg.normalization_type = "LN"
40 self.cfg.positional_embedding_type = "standard"
41 self.cfg.final_rms = False
42 self.cfg.gated_mlp = False
43 self.cfg.attn_only = False
45 # BERT uses post-LN (LayerNorm after residual, not before sublayer).
46 # fold_ln assumes pre-LN (LN before sublayer) and folds ln1 into attention
47 # QKV and ln2 into MLP. For post-LN, ln1 output feeds MLP (not attention)
48 # and ln2 output feeds next block's attention (not MLP), so folding into
49 # the wrong sublayer produces incorrect results.
50 self.supports_fold_ln = False
52 n_heads = self.cfg.n_heads
54 self.weight_processing_conversions = {
55 "blocks.{i}.attn.q.weight": ParamProcessingConversion(
56 tensor_conversion=RearrangeTensorConversion(
57 "(h d_head) d_model -> h d_model d_head", h=n_heads
58 ),
59 ),
60 "blocks.{i}.attn.k.weight": ParamProcessingConversion(
61 tensor_conversion=RearrangeTensorConversion(
62 "(h d_head) d_model -> h d_model d_head", h=n_heads
63 ),
64 ),
65 "blocks.{i}.attn.v.weight": ParamProcessingConversion(
66 tensor_conversion=RearrangeTensorConversion(
67 "(h d_head) d_model -> h d_model d_head", h=n_heads
68 ),
69 ),
70 "blocks.{i}.attn.q.bias": ParamProcessingConversion(
71 tensor_conversion=RearrangeTensorConversion("(h d_head) -> h d_head", h=n_heads),
72 ),
73 "blocks.{i}.attn.k.bias": ParamProcessingConversion(
74 tensor_conversion=RearrangeTensorConversion("(h d_head) -> h d_head", h=n_heads),
75 ),
76 "blocks.{i}.attn.v.bias": ParamProcessingConversion(
77 tensor_conversion=RearrangeTensorConversion("(h d_head) -> h d_head", h=n_heads),
78 ),
79 "blocks.{i}.attn.o.weight": ParamProcessingConversion(
80 tensor_conversion=RearrangeTensorConversion(
81 "d_model (h d_head) -> h d_head d_model", h=n_heads
82 ),
83 ),
84 }
86 # Set up component mapping
87 # MLM defaults; prepare_model() adjusts for other task heads (e.g., NSP).
88 self.component_mapping = {
89 "embed": EmbeddingBridge(name="bert.embeddings.word_embeddings"),
90 "token_type_embed": EmbeddingBridge(name="bert.embeddings.token_type_embeddings"),
91 "pos_embed": PosEmbedBridge(name="bert.embeddings.position_embeddings"),
92 "embed_ln": NormalizationBridge(
93 name="bert.embeddings.LayerNorm",
94 config=self.cfg,
95 use_native_layernorm_autograd=True,
96 ),
97 "blocks": BlockBridge(
98 name="bert.encoder.layer",
99 # BERT has no single MLP module (intermediate.dense and output.dense
100 # are siblings in BertLayer), so the MLPBridge forward is never called
101 # and mlp.hook_out never fires. Redirect hook_mlp_out to the actual
102 # MLP output hook (output of the "out" linear layer).
103 hook_alias_overrides={
104 "hook_mlp_out": "mlp.out.hook_out",
105 "hook_mlp_in": "mlp.in.hook_in",
106 },
107 submodules={
108 "ln1": NormalizationBridge(
109 name="attention.output.LayerNorm",
110 config=self.cfg,
111 use_native_layernorm_autograd=True,
112 ),
113 "ln2": NormalizationBridge(
114 name="output.LayerNorm",
115 config=self.cfg,
116 use_native_layernorm_autograd=True,
117 ),
118 "attn": AttentionBridge(
119 name="attention",
120 config=self.cfg,
121 submodules={
122 "q": LinearBridge(name="self.query"),
123 "k": LinearBridge(name="self.key"),
124 "v": LinearBridge(name="self.value"),
125 "o": LinearBridge(name="output.dense"),
126 },
127 ),
128 "mlp": MLPBridge(
129 name=None,
130 config=self.cfg,
131 submodules={
132 "in": LinearBridge(name="intermediate.dense"),
133 "out": LinearBridge(name="output.dense"),
134 },
135 ),
136 },
137 ),
138 "mlm_head": LinearBridge(name="cls.predictions.transform.dense"),
139 "unembed": UnembeddingBridge(name="cls.predictions.decoder"),
140 "ln_final": NormalizationBridge(
141 name="cls.predictions.transform.LayerNorm",
142 config=self.cfg,
143 use_native_layernorm_autograd=True,
144 ),
145 }
147 def prepare_model(self, hf_model: Any) -> None:
148 """Adjust component mapping based on the actual HF model variant.
150 BertForMaskedLM has cls.predictions (MLM head).
151 BertForNextSentencePrediction has cls.seq_relationship (NSP head)
152 and no MLM-specific LayerNorm.
153 """
154 if getattr(getattr(hf_model, "bert", None), "pooler", None) is not None:
155 self.components["pooler"] = LinearBridge(name="bert.pooler.dense")
157 has_predictions = hasattr(getattr(hf_model, "cls", None), "predictions")
158 has_nsp_head = hasattr(getattr(hf_model, "cls", None), "seq_relationship")
159 if has_nsp_head and has_predictions:
160 self.components["nsp_head"] = LinearBridge(name="cls.seq_relationship")
161 elif has_nsp_head:
162 # NSP-only model — swap head components.
163 self.components["unembed"] = UnembeddingBridge(name="cls.seq_relationship")
164 self.components.pop("mlm_head", None)
165 self.components.pop("ln_final", None)