Coverage for transformer_lens/model_bridge/supported_architectures/bert.py: 100%
30 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +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 BertPoolerBridge,
16 BlockBridge,
17 EmbeddingBridge,
18 LinearBridge,
19 MLPBridge,
20 NormalizationBridge,
21 PosEmbedBridge,
22 UnembeddingBridge,
23)
26class BertArchitectureAdapter(ArchitectureAdapter):
27 """Architecture adapter for BERT models."""
29 supports_generation: bool = False
31 def __init__(self, cfg: Any) -> None:
32 """Initialize the BERT architecture adapter.
34 Args:
35 cfg: The configuration object.
36 """
37 super().__init__(cfg)
39 # Set config variables for weight processing
40 self.cfg.normalization_type = "LN"
41 self.cfg.positional_embedding_type = "standard"
42 self.cfg.final_rms = False
43 self.cfg.gated_mlp = False
44 self.cfg.attn_only = False
45 # Encoder self-attention sees every position, so right padding must be masked.
46 self.cfg.attention_dir = "bidirectional"
48 # BERT uses post-LN (LayerNorm after residual, not before sublayer).
49 # fold_ln assumes pre-LN (LN before sublayer) and folds ln1 into attention
50 # QKV and ln2 into MLP. For post-LN, ln1 output feeds MLP (not attention)
51 # and ln2 output feeds next block's attention (not MLP), so folding into
52 # the wrong sublayer produces incorrect results.
53 self.supports_fold_ln = False
55 n_heads = self.cfg.n_heads
57 self.weight_processing_conversions = {
58 "blocks.{i}.attn.q.weight": ParamProcessingConversion(
59 tensor_conversion=RearrangeTensorConversion(
60 "(h d_head) d_model -> h d_model d_head", h=n_heads
61 ),
62 ),
63 "blocks.{i}.attn.k.weight": ParamProcessingConversion(
64 tensor_conversion=RearrangeTensorConversion(
65 "(h d_head) d_model -> h d_model d_head", h=n_heads
66 ),
67 ),
68 "blocks.{i}.attn.v.weight": ParamProcessingConversion(
69 tensor_conversion=RearrangeTensorConversion(
70 "(h d_head) d_model -> h d_model d_head", h=n_heads
71 ),
72 ),
73 "blocks.{i}.attn.q.bias": ParamProcessingConversion(
74 tensor_conversion=RearrangeTensorConversion("(h d_head) -> h d_head", h=n_heads),
75 ),
76 "blocks.{i}.attn.k.bias": ParamProcessingConversion(
77 tensor_conversion=RearrangeTensorConversion("(h d_head) -> h d_head", h=n_heads),
78 ),
79 "blocks.{i}.attn.v.bias": ParamProcessingConversion(
80 tensor_conversion=RearrangeTensorConversion("(h d_head) -> h d_head", h=n_heads),
81 ),
82 "blocks.{i}.attn.o.weight": ParamProcessingConversion(
83 tensor_conversion=RearrangeTensorConversion(
84 "d_model (h d_head) -> h d_head d_model", h=n_heads
85 ),
86 ),
87 }
89 # Set up component mapping
90 # MLM defaults; prepare_model() adjusts for other task heads (e.g., NSP).
91 self.component_mapping = {
92 "embed": EmbeddingBridge(name="bert.embeddings.word_embeddings"),
93 "token_type_embed": EmbeddingBridge(name="bert.embeddings.token_type_embeddings"),
94 "pos_embed": PosEmbedBridge(name="bert.embeddings.position_embeddings"),
95 "embed_ln": NormalizationBridge(
96 name="bert.embeddings.LayerNorm",
97 config=self.cfg,
98 use_native_layernorm_autograd=True,
99 ),
100 "blocks": BlockBridge(
101 name="bert.encoder.layer",
102 # BERT has no single MLP module (intermediate.dense and output.dense
103 # are siblings in BertLayer), so the MLPBridge forward is never called
104 # and mlp.hook_out never fires. Redirect hook_mlp_out to the actual
105 # MLP output hook (output of the "out" linear layer).
106 hook_alias_overrides={
107 "hook_mlp_out": "mlp.out.hook_out",
108 "hook_mlp_in": "mlp.in.hook_in",
109 },
110 submodules={
111 "ln1": NormalizationBridge(
112 name="attention.output.LayerNorm",
113 config=self.cfg,
114 use_native_layernorm_autograd=True,
115 ),
116 "ln2": NormalizationBridge(
117 name="output.LayerNorm",
118 config=self.cfg,
119 use_native_layernorm_autograd=True,
120 ),
121 "attn": AttentionBridge(
122 name="attention",
123 config=self.cfg,
124 submodules={
125 "q": LinearBridge(name="self.query"),
126 "k": LinearBridge(name="self.key"),
127 "v": LinearBridge(name="self.value"),
128 "o": LinearBridge(name="output.dense"),
129 },
130 ),
131 "mlp": MLPBridge(
132 name=None,
133 config=self.cfg,
134 submodules={
135 "in": LinearBridge(name="intermediate.dense"),
136 "out": LinearBridge(name="output.dense"),
137 },
138 ),
139 },
140 ),
141 "mlm_head": LinearBridge(name="cls.predictions.transform.dense"),
142 "unembed": UnembeddingBridge(name="cls.predictions.decoder"),
143 "ln_final": NormalizationBridge(
144 name="cls.predictions.transform.LayerNorm",
145 config=self.cfg,
146 use_native_layernorm_autograd=True,
147 ),
148 }
150 def prepare_model(self, hf_model: Any) -> None:
151 """Adjust component mapping based on the actual HF model variant.
153 BertForMaskedLM has cls.predictions (MLM head).
154 BertForNextSentencePrediction has cls.seq_relationship (NSP head)
155 and no MLM-specific LayerNorm.
156 """
157 if getattr(getattr(hf_model, "bert", None), "pooler", None) is not None:
158 # Wrap the pooler itself, not its inner Linear: HF applies tanh after
159 # the projection, so hook_out here is the pooled [CLS] that
160 # HookedEncoder's BertPooler exposes as hook_pooler_out. The dense
161 # stays hookable as a submodule for the pre-activation projection.
162 self.components["pooler"] = BertPoolerBridge(
163 name="bert.pooler",
164 submodules={"dense": LinearBridge(name="dense")},
165 )
167 has_predictions = hasattr(getattr(hf_model, "cls", None), "predictions")
168 has_nsp_head = hasattr(getattr(hf_model, "cls", None), "seq_relationship")
169 if has_nsp_head and has_predictions:
170 self.components["nsp_head"] = LinearBridge(name="cls.seq_relationship")
171 elif has_nsp_head:
172 # NSP-only model — swap head components.
173 self.components["unembed"] = UnembeddingBridge(name="cls.seq_relationship")
174 self.components.pop("mlm_head", None)
175 self.components.pop("ln_final", None)