Coverage for transformer_lens/model_bridge/supported_architectures/vit.py: 100%
38 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-08-11 18:50 +0000
1"""ViT / DeiT architecture adapter.
3Supports HF `ViTModel`, `ViTForImageClassification`, `DeiTModel`,
4`DeiTForImageClassification` (single CLS-token classifier head). Encoder blocks
5are structurally near-identical between ViT and DeiT — same field names
6(layernorm_before/after, attention.{q,k,v,o}_proj, mlp.{fc1,fc2}) — differing
7only in the embeddings module used (DeiT's carries an extra distillation token,
8which is invisible to this adapter — see vision_embeddings.py).
10NOT covered: `DeiTForImageClassificationWithTeacher` (dual cls+distillation head,
11averaged). See vision_classifier_head.py's docstring for why, and prepare_model()
12below raises loudly if you load one anyway rather than silently producing wrong
13logits.
15ViT/DeiT blocks are pre-LN (LayerNorm applied *before* attention/MLP, residual
16added after — same shape as Llama/GPT2), unlike BERT's post-LN. That's why
17`supports_fold_ln = True` here where BertArchitectureAdapter sets it False.
19NOTE (transformers >= the ViT/DeiT flattening refactor): as of this version of
20`modeling_vit.py`, HF removed the separate `ViTEncoder` wrapper — the blocks
21now live directly at `<prefix>.layers` on the model, not `<prefix>.encoder.layer`.
22Attention was flattened too: `ViTAttention` now owns `q_proj`/`k_proj`/`v_proj`/
23`o_proj` directly (no more nested `attention.attention.{query,key,value}` +
24`output.dense`). And `ViTLayer` already exposes a flat `.mlp` submodule
25(`ViTMLP` with `.fc1`/`.fc2`) instead of the old `intermediate`/`output` split.
26Because of this, the old `ViTMLPWrapper` shim, the block-forward tuple-unwrapping
27monkey-patch (`ViTLayer.forward` returns a plain tensor now, not a tuple), and
28the `hf_model.encoder_layer = hf_model.encoder.layer` aliasing hack are all
29gone — the component mapping below points straight at the real attributes.
30"""
32from typing import Any, Dict
34from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion
35from transformer_lens.conversion_utils.param_processing_conversion import (
36 ParamProcessingConversion,
37)
38from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
39from transformer_lens.model_bridge.generalized_components import (
40 AttentionBridge,
41 BlockBridge,
42 LinearBridge,
43 MLPBridge,
44 NormalizationBridge,
45)
46from transformer_lens.model_bridge.generalized_components.vision_classifier_head import (
47 VisionClassifierHeadBridge,
48)
49from transformer_lens.model_bridge.generalized_components.vision_embeddings import (
50 VisionEmbeddingsBridge,
51)
54class ViTArchitectureAdapter(ArchitectureAdapter):
55 """Architecture adapter for ViT and (non-distilled-head) DeiT vision models."""
57 supports_generation: bool = False
59 # Vision models have no tokenizer, so of the text phases only Phase 1 (HF
60 # parity on pixel input) applies — Phases 2/3 need a HookedTransformer
61 # counterpart and Phase 4 needs text generation. Phase 9 (vision hook/cache
62 # tests) is gated by is_visual_model, not this list; _full_and_core_phases()
63 # routes "vision" architectures to {1, 9}.
64 applicable_phases: list[int] = [1]
66 def __init__(self, cfg: Any) -> None:
67 super().__init__(cfg)
69 self.cfg.is_visual_model = True
70 self.cfg.normalization_type = "LN"
71 self.cfg.positional_embedding_type = "standard"
72 self.cfg.final_rms = False
73 self.cfg.gated_mlp = False
74 self.cfg.attn_only = False
75 self.supports_fold_ln = True
77 n_heads = self.cfg.n_heads
79 self.weight_processing_conversions = {
80 "blocks.{i}.attn.q.weight": ParamProcessingConversion(
81 tensor_conversion=RearrangeTensorConversion(
82 "(h d_head) d_model -> h d_model d_head", h=n_heads
83 ),
84 ),
85 "blocks.{i}.attn.k.weight": ParamProcessingConversion(
86 tensor_conversion=RearrangeTensorConversion(
87 "(h d_head) d_model -> h d_model d_head", h=n_heads
88 ),
89 ),
90 "blocks.{i}.attn.v.weight": ParamProcessingConversion(
91 tensor_conversion=RearrangeTensorConversion(
92 "(h d_head) d_model -> h d_model d_head", h=n_heads
93 ),
94 ),
95 "blocks.{i}.attn.q.bias": ParamProcessingConversion(
96 tensor_conversion=RearrangeTensorConversion("(h d_head) -> h d_head", h=n_heads),
97 ),
98 "blocks.{i}.attn.k.bias": ParamProcessingConversion(
99 tensor_conversion=RearrangeTensorConversion("(h d_head) -> h d_head", h=n_heads),
100 ),
101 "blocks.{i}.attn.v.bias": ParamProcessingConversion(
102 tensor_conversion=RearrangeTensorConversion("(h d_head) -> h d_head", h=n_heads),
103 ),
104 "blocks.{i}.attn.o.weight": ParamProcessingConversion(
105 tensor_conversion=RearrangeTensorConversion(
106 "d_model (h d_head) -> h d_head d_model", h=n_heads
107 ),
108 ),
109 }
111 self.component_mapping = self._build_component_mapping(prefix="", with_classifier=False)
113 def _build_component_mapping(self, prefix: str, with_classifier: bool) -> Dict[str, Any]:
114 p = prefix
115 mapping: Dict[str, Any] = {
116 "embed": VisionEmbeddingsBridge(name=f"{p}embeddings"),
117 "blocks": BlockBridge(
118 # HF's ViTModel/DeiTModel no longer wrap blocks in a `.encoder`
119 # module — they sit directly on `<prefix>.layers`.
120 name=f"{p}layers",
121 hook_alias_overrides={
122 "hook_mlp_out": "mlp.out.hook_out",
123 "hook_mlp_in": "mlp.in.hook_in",
124 },
125 submodules={
126 "ln1": NormalizationBridge(
127 name="layernorm_before",
128 config=self.cfg,
129 use_native_layernorm_autograd=True,
130 ),
131 "ln2": NormalizationBridge(
132 name="layernorm_after",
133 config=self.cfg,
134 use_native_layernorm_autograd=True,
135 ),
136 "attn": AttentionBridge(
137 name="attention",
138 config=self.cfg,
139 submodules={
140 # Submodule paths here are resolved relative to the
141 # already-resolved "attn" bridge module itself
142 # (block.attention). ViTAttention/DeiTAttention now
143 # owns q_proj/k_proj/v_proj/o_proj directly
144 # (flattened, no nested self-attention + separate
145 # output.dense module), so no "attention." prefix.
146 "q": LinearBridge(name="q_proj"),
147 "k": LinearBridge(name="k_proj"),
148 "v": LinearBridge(name="v_proj"),
149 "o": LinearBridge(name="o_proj"),
150 },
151 ),
152 "mlp": MLPBridge(
153 name="mlp",
154 config=self.cfg,
155 submodules={
156 # ViTLayer.mlp is a real ViTMLP module now
157 # (fc1/fc2) — no more intermediate/output split,
158 # so no wrapper shim is needed.
159 "in": LinearBridge(name="fc1"),
160 "out": LinearBridge(name="fc2"),
161 },
162 ),
163 },
164 ),
165 "ln_final": NormalizationBridge(
166 name=f"{p}layernorm",
167 config=self.cfg,
168 use_native_layernorm_autograd=True,
169 ),
170 }
171 if with_classifier:
172 mapping["unembed"] = VisionClassifierHeadBridge(name="classifier")
173 return mapping
175 def prepare_model(self, hf_model: Any) -> None:
176 """Detect ViTForImageClassification vs DeiTForImageClassification vs a bare
177 *Model, and add/omit the classifier head + prefix accordingly.
179 No structural patching of the HF model is needed any more: current
180 transformers ViT/DeiT blocks already expose a flat `.mlp` (fc1/fc2) and
181 return plain tensors from `forward`, and the blocks live directly on
182 `<prefix>.layers` rather than behind a now-removed `.encoder` wrapper.
183 This method only has to figure out the right prefix/classifier and
184 build the component mapping to point at those real attributes.
185 """
186 if hasattr(hf_model, "cls_classifier") and hasattr(hf_model, "distillation_classifier"):
187 raise NotImplementedError(
188 "DeiTForImageClassificationWithTeacher (dual cls_classifier + "
189 "distillation_classifier head, averaged) isn't supported by this "
190 "adapter yet — see vision_classifier_head.py's docstring for why, "
191 "and what to check before adding it."
192 )
194 # Check for the wrapper module created by HF's ForImageClassification classes.
195 # Bare ViTModel / DeiTModel have components at the root, so prefix must be "".
196 if hasattr(hf_model, "deit") and not hasattr(hf_model, "vit"):
197 prefix = "deit."
198 elif hasattr(hf_model, "vit"):
199 prefix = "vit."
200 else:
201 prefix = ""
203 with_classifier = hasattr(hf_model, "classifier")
204 self.component_mapping = self._build_component_mapping(
205 prefix=prefix, with_classifier=with_classifier
206 )