Coverage for transformer_lens/model_bridge/supported_architectures/bloom.py: 100%
40 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Bloom architecture adapter."""
3from typing import Any
5import torch
7from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion
8from transformer_lens.conversion_utils.param_processing_conversion import (
9 ParamProcessingConversion,
10)
11from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
12from transformer_lens.model_bridge.generalized_components import (
13 BloomAttentionBridge,
14 BloomBlockBridge,
15 BloomMLPBridge,
16 EmbeddingBridge,
17 LinearBridge,
18 NormalizationBridge,
19 UnembeddingBridge,
20)
23class BloomArchitectureAdapter(ArchitectureAdapter):
24 """Architecture adapter for Bloom models."""
26 def __init__(self, cfg: Any) -> None:
27 """Initialize the Bloom architecture adapter."""
28 super().__init__(cfg)
30 # Set config variables for weight processing
31 self.cfg.normalization_type = "LN"
32 self.cfg.positional_embedding_type = "alibi"
33 self.cfg.final_rms = False
34 self.cfg.gated_mlp = False
35 self.cfg.attn_only = False
36 # HF BloomAttention is always full MHA and ignores num_key_value_heads;
37 # a stray config field would otherwise send weight processing down GQA paths.
38 self.cfg.n_key_value_heads = None
40 self.cfg.default_prepend_bos = False
41 # After split_qkv_matrix, Q/K/V are individual [n_heads*d_head, d_model] weights.
42 # Convert to TL format [n_heads, d_model, d_head].
43 self.weight_processing_conversions = {
44 "blocks.{i}.attn.q": ParamProcessingConversion(
45 tensor_conversion=RearrangeTensorConversion(
46 "(n h) m -> n m h",
47 n=self.cfg.n_heads,
48 ),
49 ),
50 "blocks.{i}.attn.k": ParamProcessingConversion(
51 tensor_conversion=RearrangeTensorConversion(
52 "(n h) m -> n m h",
53 n=self.cfg.n_heads,
54 ),
55 ),
56 "blocks.{i}.attn.v": ParamProcessingConversion(
57 tensor_conversion=RearrangeTensorConversion(
58 "(n h) m -> n m h",
59 n=self.cfg.n_heads,
60 ),
61 ),
62 "blocks.{i}.attn.o": ParamProcessingConversion(
63 tensor_conversion=RearrangeTensorConversion("m (n h) -> n h m", n=self.cfg.n_heads),
64 ),
65 }
67 self.component_mapping = {
68 "embed": EmbeddingBridge(name="transformer.word_embeddings"),
69 "embed_ln": NormalizationBridge(
70 name="transformer.word_embeddings_layernorm", config=self.cfg
71 ),
72 "blocks": BloomBlockBridge(
73 name="transformer.h",
74 config=self.cfg,
75 hook_alias_overrides={
76 "hook_attn_out": "attn.o.hook_out",
77 "hook_mlp_out": "mlp.out.hook_out",
78 },
79 submodules={
80 "ln1": NormalizationBridge(name="input_layernorm", config=self.cfg),
81 "ln2": NormalizationBridge(name="post_attention_layernorm", config=self.cfg),
82 "attn": BloomAttentionBridge(
83 name="self_attention",
84 config=self.cfg,
85 split_qkv_matrix=self.split_qkv_matrix,
86 submodules={
87 "qkv": LinearBridge(name="query_key_value"),
88 "o": LinearBridge(name="dense"),
89 },
90 ),
91 "mlp": BloomMLPBridge(
92 name="mlp",
93 submodules={
94 "in": LinearBridge(name="dense_h_to_4h"),
95 "out": LinearBridge(name="dense_4h_to_h"),
96 },
97 ),
98 },
99 ),
100 "ln_final": NormalizationBridge(name="transformer.ln_f", config=self.cfg),
101 "unembed": UnembeddingBridge(name="lm_head"),
102 }
104 def split_qkv_matrix(
105 self, original_attention_component: Any
106 ) -> tuple[torch.nn.Linear, torch.nn.Linear, torch.nn.Linear]:
107 """Split the QKV matrix into separate linear transformations.
108 Args:
109 attention_component: The original attention layer component
110 Returns:
111 Tuple of nn.Linear modules for Q, K, and V transformations
112 """
114 # Keep mypy happy
115 assert original_attention_component is not None
116 assert original_attention_component.query_key_value is not None
118 qkv_weights = original_attention_component.query_key_value.weight
120 # Keep mypy happy
121 assert isinstance(qkv_weights, torch.Tensor)
123 # Bloom QKV weights are interleaved: [Q0,K0,V0, Q1,K1,V1, ...]
124 # i.e. layout is (n_heads, 3, d_head), not (3, n_heads*d_head).
125 # Reshape to [d_model, n_heads, 3, d_head] to correctly deinterleave.
126 W_split = qkv_weights.T.reshape(self.cfg.d_model, self.cfg.n_heads, 3, self.cfg.d_head)
128 # W_Q/K/V shape: [d_model, n_heads, d_head]
129 W_Q, W_K, W_V = W_split[..., 0, :], W_split[..., 1, :], W_split[..., 2, :]
131 qkv_bias = original_attention_component.query_key_value.bias
133 # Keep mypy happy
134 assert isinstance(qkv_bias, torch.Tensor)
136 # Same interleaved layout for bias: reshape to [n_heads, 3, d_head]
137 qkv_bias = qkv_bias.reshape(self.cfg.n_heads, 3, self.cfg.d_head)
139 # b_Q/K/V shape: [n_heads, d_head]
140 b_Q, b_K, b_V = qkv_bias[:, 0, :], qkv_bias[:, 1, :], qkv_bias[:, 2, :]
142 # Create nn.Linear modules
143 # W_Q shape is [d_model, n_heads, d_head] -> flatten to [d_model, n_heads*d_head]
144 # nn.Linear expects weight shape [out_features, in_features] = [n_heads*d_head, d_model]
145 d_out = self.cfg.n_heads * self.cfg.d_head
147 W_Q_transformation = torch.nn.Linear(self.cfg.d_model, d_out, bias=True)
148 W_Q_transformation.weight = torch.nn.Parameter(W_Q.reshape(self.cfg.d_model, d_out).T)
149 W_Q_transformation.bias = torch.nn.Parameter(b_Q.reshape(d_out))
151 W_K_transformation = torch.nn.Linear(self.cfg.d_model, d_out, bias=True)
152 W_K_transformation.weight = torch.nn.Parameter(W_K.reshape(self.cfg.d_model, d_out).T)
153 W_K_transformation.bias = torch.nn.Parameter(b_K.reshape(d_out))
155 W_V_transformation = torch.nn.Linear(self.cfg.d_model, d_out, bias=True)
156 W_V_transformation.weight = torch.nn.Parameter(W_V.reshape(self.cfg.d_model, d_out).T)
157 W_V_transformation.bias = torch.nn.Parameter(b_V.reshape(d_out))
159 return W_Q_transformation, W_K_transformation, W_V_transformation