Coverage for transformer_lens/model_bridge/supported_architectures/llada.py: 82%
184 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""LLaDA architecture support for one masked-token transformer forward pass.
3LLaDA is a masked discrete-diffusion model. Its transformer uses bidirectional
4self-attention and returns logits for every input position; the iterative
5denoising sampler lives outside the model. This adapter intentionally supports
6only the transformer forward pass. TransformerBridge's autoregressive
7generation APIs are disabled for this architecture.
9The remote ``LLaDALlamaBlock`` owns its Q/K/V, output, and gated-MLP projections
10directly instead of grouping them in attention and MLP modules. The adapter
11preserves the reviewed remote block forward and replaces only its
12``attention(...)`` method with a hook-aware reconstruction. This keeps the
13native pre-norm residual order, dropout, RoPE implementation, and MLP math while
14exposing the standard attention scores and pattern hooks.
15"""
17from __future__ import annotations
19import math
20from typing import Any, Dict, Optional, cast
22import torch
24from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
25from transformer_lens.model_bridge.generalized_components import (
26 AttentionBridge,
27 BlockBridge,
28 EmbeddingBridge,
29 GatedMLPBridge,
30 LinearBridge,
31 RMSNormalizationBridge,
32 UnembeddingBridge,
33)
34from transformer_lens.model_bridge.generalized_components.base import (
35 GeneralizedComponent,
36)
39class _LLaDAAttentionBridge(AttentionBridge):
40 """Reconstruct LLaDA's block-local, bidirectional attention method."""
42 supports_split_qkv_fork: bool = False
43 supports_attn_result: bool = True
45 def set_original_component(self, original_component: torch.nn.Module) -> None:
46 """Patch the owning block without registering it as a child module."""
47 object.__setattr__(self, "_original_block", original_component)
48 object.__setattr__(original_component, "attention", self)
50 @property
51 def original_component(self) -> Optional[torch.nn.Module]:
52 """Return the owning LLaDA block."""
53 return self.__dict__.get("_original_block")
55 def forward(
56 self,
57 q: torch.Tensor,
58 k: torch.Tensor,
59 v: torch.Tensor,
60 attention_bias: Optional[torch.Tensor] = None,
61 layer_past: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
62 use_cache: bool = False,
63 ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:
64 """Run LLaDA attention with observable pre/post-softmax tensors."""
65 block = self.original_component
66 if block is None: 66 ↛ 67line 66 didn't jump to line 67 because the condition on line 66 was never true
67 raise RuntimeError("LLaDA attention is not attached to an owning block")
68 config = self.config
69 if config is None: 69 ↛ 70line 69 didn't jump to line 70 because the condition on line 69 was never true
70 raise RuntimeError("LLaDA attention is not attached to a bridge config")
72 batch_size, query_len, q_width = q.shape
73 n_heads = int(config.n_heads)
74 n_kv_heads = int(getattr(config, "n_key_value_heads", None) or n_heads)
75 head_dim = q_width // n_heads
76 input_dtype = k.dtype
78 q_norm = getattr(block, "q_norm", None)
79 k_norm = getattr(block, "k_norm", None)
80 if q_norm is not None and k_norm is not None: 80 ↛ 81line 80 didn't jump to line 81 because the condition on line 80 was never true
81 q = q_norm(q).to(dtype=input_dtype)
82 k = k_norm(k).to(dtype=input_dtype)
84 q = q.view(batch_size, query_len, n_heads, head_dim).transpose(1, 2)
85 k = k.view(batch_size, query_len, n_kv_heads, head_dim).transpose(1, 2)
86 v = v.view(batch_size, query_len, n_kv_heads, head_dim).transpose(1, 2)
88 if layer_past is not None: 88 ↛ 89line 88 didn't jump to line 89 because the condition on line 88 was never true
89 past_key, past_value = layer_past
90 k = torch.cat((past_key, k), dim=-2)
91 v = torch.cat((past_value, v), dim=-2)
92 present = (k, v) if use_cache else None
94 if bool(getattr(block.config, "rope", False)): 94 ↛ 96line 94 didn't jump to line 96 because the condition on line 94 was always true
95 q, k = cast(Any, block.rotary_emb)(q, k)
96 if hasattr(self, "hook_rot_q"): 96 ↛ 98line 96 didn't jump to line 98 because the condition on line 96 was always true
97 q = self.hook_rot_q(q)
98 if hasattr(self, "hook_rot_k"): 98 ↛ 101line 98 didn't jump to line 101 because the condition on line 98 was always true
99 k = self.hook_rot_k(k)
101 key_len = k.shape[-2]
102 if attention_bias is not None:
103 attention_bias = cast(Any, block._cast_attn_bias)(
104 attention_bias[:, :, key_len - query_len : key_len, :key_len],
105 input_dtype,
106 )
108 if n_heads != n_kv_heads: 108 ↛ 117line 108 didn't jump to line 117 because the condition on line 108 was always true
109 if n_heads % n_kv_heads != 0: 109 ↛ 110line 109 didn't jump to line 110 because the condition on line 109 was never true
110 raise ValueError(
111 f"n_heads ({n_heads}) must be divisible by n_key_value_heads " f"({n_kv_heads})"
112 )
113 groups = n_heads // n_kv_heads
114 k = k.repeat_interleave(groups, dim=1, output_size=n_heads)
115 v = v.repeat_interleave(groups, dim=1, output_size=n_heads)
117 attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(head_dim)
118 attn_scores = self._apply_reconstruct_attention_mask(
119 attn_scores=attn_scores,
120 attention_mask=attention_bias,
121 seq_len=key_len,
122 q_seq_len=query_len,
123 )
124 attn_scores = self.hook_attn_scores(attn_scores)
126 pattern = torch.nn.functional.softmax(attn_scores, dim=-1, dtype=torch.float32).to(q.dtype)
127 pattern = self._scrub_compatibility_pattern_nans(pattern)
128 dropout = float(getattr(block.config, "attention_dropout", 0.0))
129 if block.training and dropout > 0.0: 129 ↛ 130line 129 didn't jump to line 130 because the condition on line 129 was never true
130 pattern = torch.nn.functional.dropout(pattern, p=dropout, training=True)
131 pattern = self.hook_pattern(pattern)
133 z = torch.matmul(pattern, v).transpose(1, 2).contiguous()
134 flat_z = z.view(batch_size, query_len, n_heads * head_dim)
135 if bool(getattr(self.config, "use_attn_result", False)):
136 flat_z = self.o.hook_in(flat_z)
137 z = flat_z.view(batch_size, query_len, n_heads, head_dim)
138 output = self._compute_per_head_result(z, n_heads, head_dim)
139 else:
140 output = self.o(flat_z)
141 output = self.hook_out(output)
142 return output, present
145class _LLaDABlockBridge(BlockBridge):
146 """Preserve native block math while routing container-level hooks."""
148 def __init__(
149 self,
150 name: str,
151 config: Optional[Any] = None,
152 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
153 ) -> None:
154 super().__init__(
155 name,
156 config=config,
157 submodules=submodules,
158 hook_alias_overrides={
159 "hook_attn_in": "attn.hook_in",
160 "hook_q_input": "attn.q.hook_in",
161 "hook_k_input": "attn.k.hook_in",
162 "hook_v_input": "attn.v.hook_in",
163 },
164 )
165 self._llada_container_hooks_wired = False
166 self._llada_container_hook_handles: list[torch.utils.hooks.RemovableHandle] = []
168 def _route_attn_input(
169 self, _module: torch.nn.Module, _args: tuple[Any, ...], output: torch.Tensor
170 ) -> torch.Tensor:
171 return self.attn.hook_in(output)
173 def _route_mlp_input(
174 self, _module: torch.nn.Module, _args: tuple[Any, ...], output: torch.Tensor
175 ) -> torch.Tensor:
176 return self.mlp.hook_in(output)
178 def _route_mlp_output(
179 self, _module: torch.nn.Module, _args: tuple[Any, ...], output: torch.Tensor
180 ) -> torch.Tensor:
181 if bool(getattr(self.mlp, "_executing_container", False)): 181 ↛ 182line 181 didn't jump to line 182 because the condition on line 181 was never true
182 return output
183 return self.mlp.hook_out(output)
185 def _route_pre_mlp_norm(
186 self, _module: torch.nn.Module, args: tuple[Any, ...]
187 ) -> Optional[tuple[Any, ...]]:
188 if not self._read_use_hook_mlp_in(): 188 ↛ 190line 188 didn't jump to line 190 because the condition on line 188 was always true
189 return None
190 if args and isinstance(args[0], torch.Tensor):
191 return (self.hook_mlp_in(args[0]),) + args[1:]
192 return None
194 def _maybe_wire_capture_hooks(self) -> None:
195 """Use deepcopy-safe bound hooks for LLaDA's pre-MLP residual hook."""
196 if self._capture_hooks_wired:
197 return
198 if self.ln2.original_component is not None: 198 ↛ 202line 198 didn't jump to line 202 because the condition on line 198 was always true
199 self._capture_hook_handles.append(
200 self.ln2.register_forward_pre_hook(self._route_pre_mlp_norm)
201 )
202 self._capture_hooks_wired = True
204 def _wire_llada_container_hooks(self) -> None:
205 if self._llada_container_hooks_wired:
206 return
208 self._llada_container_hook_handles.extend(
209 (
210 self.ln1.register_forward_hook(self._route_attn_input),
211 self.ln2.register_forward_hook(self._route_mlp_input),
212 self.mlp.out.register_forward_hook(self._route_mlp_output),
213 )
214 )
215 self._llada_container_hooks_wired = True
217 def forward(self, *args: Any, **kwargs: Any) -> Any:
218 """Wire LLaDA's symbolic containers, then delegate to the native block."""
219 self._wire_llada_container_hooks()
220 return super().forward(*args, **kwargs)
223class _LLaDAGatedMLPBridge(GatedMLPBridge):
224 """Executable view over LLaDA's block-local gated MLP projections."""
226 # The block calls this view's forward, which fires hook_in/hook_out itself;
227 # the setup-time mirror would double-apply interventions.
228 mirror_placeholder_hooks = False
230 def set_original_component(self, original_component: torch.nn.Module) -> None:
231 """The executable view needs the block's children, not ownership of the block."""
232 del original_component
234 def forward(self, hidden_states: torch.Tensor, **kwargs: Any) -> torch.Tensor:
235 """Run SiLU-gated MLP math using the live wrapped projections."""
236 del kwargs
237 hidden_states = self.hook_in(hidden_states)
238 self._executing_container = True
239 try:
240 up_projection = getattr(self, "in")
241 gate = self.gate(hidden_states)
242 up = up_projection(hidden_states)
243 activated = self.act(gate)
244 gated = activated * up
245 output = self.out(gated)
246 finally:
247 self._executing_container = False
248 return self.hook_out(output)
251class LLaDAArchitectureAdapter(ArchitectureAdapter):
252 """Adapter for the dense ``LLaDAModelLM`` architecture.
254 Support is deliberately limited to the released dense LLaDA block contract:
255 Llama-style blocks, RMSNorm, separate bias-free projections, RoPE,
256 bidirectional attention, an untied LM head, and no KV cache. The external
257 iterative denoising/remasking loop is not a TransformerBridge generation API.
258 Loading the Hugging Face checkpoint requires the caller to opt in with
259 ``trust_remote_code=True``.
260 """
262 applicable_phases: list[int] = []
263 supports_generation: bool = False
264 supports_hf_output_attentions: bool = False
265 supports_causal_loss: bool = False
267 def __init__(self, cfg: Any) -> None:
268 super().__init__(cfg)
269 self._validate_config()
271 self.cfg.d_vocab_out = self.cfg.d_vocab
272 self.cfg.normalization_type = "RMS"
273 self.cfg.uses_rms_norm = True
274 self.cfg.positional_embedding_type = "rotary"
275 self.cfg.rotary_adjacent_pairs = False
276 self.cfg.attention_dir = "bidirectional"
277 self.cfg.final_rms = True
278 self.cfg.gated_mlp = True
279 self.cfg.attn_only = False
280 self.cfg.default_prepend_bos = False
281 self.cfg.default_padding_side = "right"
283 self.weight_processing_conversions = {
284 **self._qkvo_weight_conversions(),
285 }
287 self.component_mapping = {
288 "embed": EmbeddingBridge(name="model.transformer.wte", config=self.cfg),
289 "blocks": _LLaDABlockBridge(
290 name="model.transformer.blocks",
291 config=self.cfg,
292 submodules={
293 "ln1": RMSNormalizationBridge(name="attn_norm", config=self.cfg),
294 "attn": _LLaDAAttentionBridge(
295 name=None,
296 config=self.cfg,
297 is_causal=False,
298 submodules={
299 "q": LinearBridge(name="q_proj"),
300 "k": LinearBridge(name="k_proj"),
301 "v": LinearBridge(name="v_proj"),
302 "o": LinearBridge(name="attn_out"),
303 },
304 ),
305 "ln2": RMSNormalizationBridge(name="ff_norm", config=self.cfg),
306 "mlp": _LLaDAGatedMLPBridge(
307 name=None,
308 config=self.cfg,
309 submodules={
310 "gate": LinearBridge(name="ff_proj"),
311 "in": LinearBridge(name="up_proj"),
312 "act": GeneralizedComponent(name="act"),
313 "out": LinearBridge(name="ff_out"),
314 },
315 ),
316 },
317 ),
318 "ln_final": RMSNormalizationBridge(name="model.transformer.ln_f", config=self.cfg),
319 "unembed": UnembeddingBridge(name="model.transformer.ff_out", config=self.cfg),
320 }
322 @staticmethod
323 def _enum_value(value: Any) -> Any:
324 return getattr(value, "value", value)
326 def _require_value(self, name: str, expected: Any) -> None:
327 if not hasattr(self.cfg, name):
328 raise ValueError(f"LLaDAModelLM config is missing required field '{name}'")
329 actual = self._enum_value(getattr(self.cfg, name))
330 if actual != expected:
331 raise ValueError(f"LLaDAModelLM requires {name}={expected!r}; got {actual!r}")
333 def _validate_config(self) -> None:
334 self._require_value("block_type", "llama")
335 self._require_value("block_group_size", 1)
336 self._require_value("rope", True)
337 self._require_value("rope_full_precision", True)
338 self._require_value("alibi", False)
339 self._require_value("attention_layer_norm", False)
340 self._require_value("include_bias", False)
341 self._require_value("include_qkv_bias", False)
342 self._require_value("scale_logits", False)
343 self._require_value("input_emb_norm", False)
344 self._require_value("layer_norm_type", "rms")
345 if self.cfg.act_fn != "silu":
346 raise ValueError(
347 f"LLaDAModelLM requires activation_type='silu'; got {self.cfg.act_fn!r}"
348 )
349 if bool(self.cfg.tie_word_embeddings):
350 raise ValueError(
351 "LLaDAModelLM tied embeddings are not supported by the initial dense adapter"
352 )
353 embedding_size = int(getattr(self.cfg, "embedding_size", self.cfg.d_vocab))
354 if embedding_size != self.cfg.d_vocab:
355 raise ValueError(
356 "LLaDAModelLM requires embedding_size == vocab_size in the initial "
357 f"dense adapter; got {embedding_size} and {self.cfg.d_vocab}"
358 )
360 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
361 """Disable remote branches incompatible with single-pass hook support."""
362 del model_name
363 config = model_kwargs.get("config")
364 if config is not None: 364 ↛ exitline 364 didn't return from function 'prepare_loading' because the condition on line 364 was always true
365 config.output_attentions = False
366 config.use_cache = False
368 def prepare_model(self, hf_model: Any) -> None:
369 """Keep the wrapper and underlying model on the no-cache path."""
370 for config in (
371 getattr(hf_model, "config", None),
372 getattr(getattr(hf_model, "model", None), "config", None),
373 ):
374 if config is not None: 374 ↛ 370line 374 didn't jump to line 370 because the condition on line 374 was always true
375 config.output_attentions = False
376 config.use_cache = False