Coverage for transformer_lens/model_bridge/supported_architectures/llada.py: 83%
182 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-08-11 18:50 +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 dropout = float(getattr(block.config, "attention_dropout", 0.0))
128 if block.training and dropout > 0.0: 128 ↛ 129line 128 didn't jump to line 129 because the condition on line 128 was never true
129 pattern = torch.nn.functional.dropout(pattern, p=dropout, training=True)
130 pattern = self.hook_pattern(pattern)
132 z = torch.matmul(pattern, v).transpose(1, 2).contiguous()
133 flat_z = z.view(batch_size, query_len, n_heads * head_dim)
134 if bool(getattr(self.config, "use_attn_result", False)):
135 flat_z = self.o.hook_in(flat_z)
136 z = flat_z.view(batch_size, query_len, n_heads, head_dim)
137 output = self._compute_per_head_result(z, n_heads, head_dim)
138 else:
139 output = self.o(flat_z)
140 output = self.hook_out(output)
141 return output, present
144class _LLaDABlockBridge(BlockBridge):
145 """Preserve native block math while routing container-level hooks."""
147 def __init__(
148 self,
149 name: str,
150 config: Optional[Any] = None,
151 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
152 ) -> None:
153 super().__init__(
154 name,
155 config=config,
156 submodules=submodules,
157 hook_alias_overrides={
158 "hook_attn_in": "attn.hook_in",
159 "hook_q_input": "attn.q.hook_in",
160 "hook_k_input": "attn.k.hook_in",
161 "hook_v_input": "attn.v.hook_in",
162 },
163 )
164 self._llada_container_hooks_wired = False
165 self._llada_container_hook_handles: list[torch.utils.hooks.RemovableHandle] = []
167 def _route_attn_input(
168 self, _module: torch.nn.Module, _args: tuple[Any, ...], output: torch.Tensor
169 ) -> torch.Tensor:
170 return self.attn.hook_in(output)
172 def _route_mlp_input(
173 self, _module: torch.nn.Module, _args: tuple[Any, ...], output: torch.Tensor
174 ) -> torch.Tensor:
175 return self.mlp.hook_in(output)
177 def _route_mlp_output(
178 self, _module: torch.nn.Module, _args: tuple[Any, ...], output: torch.Tensor
179 ) -> torch.Tensor:
180 if bool(getattr(self.mlp, "_executing_container", False)): 180 ↛ 181line 180 didn't jump to line 181 because the condition on line 180 was never true
181 return output
182 return self.mlp.hook_out(output)
184 def _route_pre_mlp_norm(
185 self, _module: torch.nn.Module, args: tuple[Any, ...]
186 ) -> Optional[tuple[Any, ...]]:
187 if not self._read_use_hook_mlp_in(): 187 ↛ 189line 187 didn't jump to line 189 because the condition on line 187 was always true
188 return None
189 if args and isinstance(args[0], torch.Tensor):
190 return (self.hook_mlp_in(args[0]),) + args[1:]
191 return None
193 def _maybe_wire_pre_ln_capture(self) -> None:
194 """Use deepcopy-safe bound hooks for LLaDA's pre-MLP residual hook."""
195 if self._pre_ln_capture_wired:
196 return
197 if self.ln2.original_component is not None: 197 ↛ 201line 197 didn't jump to line 201 because the condition on line 197 was always true
198 self._pre_ln_capture_handles.append(
199 self.ln2.register_forward_pre_hook(self._route_pre_mlp_norm)
200 )
201 self._pre_ln_capture_wired = True
203 def _wire_llada_container_hooks(self) -> None:
204 if self._llada_container_hooks_wired:
205 return
207 self._llada_container_hook_handles.extend(
208 (
209 self.ln1.register_forward_hook(self._route_attn_input),
210 self.ln2.register_forward_hook(self._route_mlp_input),
211 self.mlp.out.register_forward_hook(self._route_mlp_output),
212 )
213 )
214 self._llada_container_hooks_wired = True
216 def forward(self, *args: Any, **kwargs: Any) -> Any:
217 """Wire LLaDA's symbolic containers, then delegate to the native block."""
218 self._wire_llada_container_hooks()
219 return super().forward(*args, **kwargs)
222class _LLaDAGatedMLPBridge(GatedMLPBridge):
223 """Executable view over LLaDA's block-local gated MLP projections."""
225 def set_original_component(self, original_component: torch.nn.Module) -> None:
226 """The executable view needs the block's children, not ownership of the block."""
227 del original_component
229 def forward(self, hidden_states: torch.Tensor, **kwargs: Any) -> torch.Tensor:
230 """Run SiLU-gated MLP math using the live wrapped projections."""
231 del kwargs
232 hidden_states = self.hook_in(hidden_states)
233 self._executing_container = True
234 try:
235 up_projection = getattr(self, "in")
236 gate = self.gate(hidden_states)
237 up = up_projection(hidden_states)
238 activated = self.act(gate)
239 gated = activated * up
240 output = self.out(gated)
241 finally:
242 self._executing_container = False
243 return self.hook_out(output)
246class LLaDAArchitectureAdapter(ArchitectureAdapter):
247 """Adapter for the dense ``LLaDAModelLM`` architecture.
249 Support is deliberately limited to the released dense LLaDA block contract:
250 Llama-style blocks, RMSNorm, separate bias-free projections, RoPE,
251 bidirectional attention, an untied LM head, and no KV cache. The external
252 iterative denoising/remasking loop is not a TransformerBridge generation API.
253 Loading the Hugging Face checkpoint requires the caller to opt in with
254 ``trust_remote_code=True``.
255 """
257 applicable_phases: list[int] = []
258 supports_generation: bool = False
259 supports_hf_output_attentions: bool = False
260 supports_causal_loss: bool = False
262 def __init__(self, cfg: Any) -> None:
263 super().__init__(cfg)
264 self._validate_config()
266 self.cfg.d_vocab_out = self.cfg.d_vocab
267 self.cfg.normalization_type = "RMS"
268 self.cfg.uses_rms_norm = True
269 self.cfg.positional_embedding_type = "rotary"
270 self.cfg.rotary_adjacent_pairs = False
271 self.cfg.attention_dir = "bidirectional"
272 self.cfg.final_rms = True
273 self.cfg.gated_mlp = True
274 self.cfg.attn_only = False
275 self.cfg.default_prepend_bos = False
276 self.cfg.default_padding_side = "right"
278 self.weight_processing_conversions = {
279 **self._qkvo_weight_conversions(),
280 }
282 self.component_mapping = {
283 "embed": EmbeddingBridge(name="model.transformer.wte", config=self.cfg),
284 "blocks": _LLaDABlockBridge(
285 name="model.transformer.blocks",
286 config=self.cfg,
287 submodules={
288 "ln1": RMSNormalizationBridge(name="attn_norm", config=self.cfg),
289 "attn": _LLaDAAttentionBridge(
290 name=None,
291 config=self.cfg,
292 is_causal=False,
293 submodules={
294 "q": LinearBridge(name="q_proj"),
295 "k": LinearBridge(name="k_proj"),
296 "v": LinearBridge(name="v_proj"),
297 "o": LinearBridge(name="attn_out"),
298 },
299 ),
300 "ln2": RMSNormalizationBridge(name="ff_norm", config=self.cfg),
301 "mlp": _LLaDAGatedMLPBridge(
302 name=None,
303 config=self.cfg,
304 submodules={
305 "gate": LinearBridge(name="ff_proj"),
306 "in": LinearBridge(name="up_proj"),
307 "act": GeneralizedComponent(name="act"),
308 "out": LinearBridge(name="ff_out"),
309 },
310 ),
311 },
312 ),
313 "ln_final": RMSNormalizationBridge(name="model.transformer.ln_f", config=self.cfg),
314 "unembed": UnembeddingBridge(name="model.transformer.ff_out", config=self.cfg),
315 }
317 @staticmethod
318 def _enum_value(value: Any) -> Any:
319 return getattr(value, "value", value)
321 def _require_value(self, name: str, expected: Any) -> None:
322 if not hasattr(self.cfg, name):
323 raise ValueError(f"LLaDAModelLM config is missing required field '{name}'")
324 actual = self._enum_value(getattr(self.cfg, name))
325 if actual != expected:
326 raise ValueError(f"LLaDAModelLM requires {name}={expected!r}; got {actual!r}")
328 def _validate_config(self) -> None:
329 self._require_value("block_type", "llama")
330 self._require_value("block_group_size", 1)
331 self._require_value("rope", True)
332 self._require_value("rope_full_precision", True)
333 self._require_value("alibi", False)
334 self._require_value("attention_layer_norm", False)
335 self._require_value("include_bias", False)
336 self._require_value("include_qkv_bias", False)
337 self._require_value("scale_logits", False)
338 self._require_value("input_emb_norm", False)
339 self._require_value("layer_norm_type", "rms")
340 if self.cfg.act_fn != "silu":
341 raise ValueError(
342 f"LLaDAModelLM requires activation_type='silu'; got {self.cfg.act_fn!r}"
343 )
344 if bool(self.cfg.tie_word_embeddings):
345 raise ValueError(
346 "LLaDAModelLM tied embeddings are not supported by the initial dense adapter"
347 )
348 embedding_size = int(getattr(self.cfg, "embedding_size", self.cfg.d_vocab))
349 if embedding_size != self.cfg.d_vocab:
350 raise ValueError(
351 "LLaDAModelLM requires embedding_size == vocab_size in the initial "
352 f"dense adapter; got {embedding_size} and {self.cfg.d_vocab}"
353 )
355 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
356 """Disable remote branches incompatible with single-pass hook support."""
357 del model_name
358 config = model_kwargs.get("config")
359 if config is not None: 359 ↛ exitline 359 didn't return from function 'prepare_loading' because the condition on line 359 was always true
360 config.output_attentions = False
361 config.use_cache = False
363 def prepare_model(self, hf_model: Any) -> None:
364 """Keep the wrapper and underlying model on the no-cache path."""
365 for config in (
366 getattr(hf_model, "config", None),
367 getattr(getattr(hf_model, "model", None), "config", None),
368 ):
369 if config is not None: 369 ↛ 365line 369 didn't jump to line 365 because the condition on line 369 was always true
370 config.output_attentions = False
371 config.use_cache = False