Coverage for transformer_lens/model_bridge/supported_architectures/llada.py: 83%
183 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +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 def set_original_component(self, original_component: torch.nn.Module) -> None:
227 """The executable view needs the block's children, not ownership of the block."""
228 del original_component
230 def forward(self, hidden_states: torch.Tensor, **kwargs: Any) -> torch.Tensor:
231 """Run SiLU-gated MLP math using the live wrapped projections."""
232 del kwargs
233 hidden_states = self.hook_in(hidden_states)
234 self._executing_container = True
235 try:
236 up_projection = getattr(self, "in")
237 gate = self.gate(hidden_states)
238 up = up_projection(hidden_states)
239 activated = self.act(gate)
240 gated = activated * up
241 output = self.out(gated)
242 finally:
243 self._executing_container = False
244 return self.hook_out(output)
247class LLaDAArchitectureAdapter(ArchitectureAdapter):
248 """Adapter for the dense ``LLaDAModelLM`` architecture.
250 Support is deliberately limited to the released dense LLaDA block contract:
251 Llama-style blocks, RMSNorm, separate bias-free projections, RoPE,
252 bidirectional attention, an untied LM head, and no KV cache. The external
253 iterative denoising/remasking loop is not a TransformerBridge generation API.
254 Loading the Hugging Face checkpoint requires the caller to opt in with
255 ``trust_remote_code=True``.
256 """
258 applicable_phases: list[int] = []
259 supports_generation: bool = False
260 supports_hf_output_attentions: bool = False
261 supports_causal_loss: bool = False
263 def __init__(self, cfg: Any) -> None:
264 super().__init__(cfg)
265 self._validate_config()
267 self.cfg.d_vocab_out = self.cfg.d_vocab
268 self.cfg.normalization_type = "RMS"
269 self.cfg.uses_rms_norm = True
270 self.cfg.positional_embedding_type = "rotary"
271 self.cfg.rotary_adjacent_pairs = False
272 self.cfg.attention_dir = "bidirectional"
273 self.cfg.final_rms = True
274 self.cfg.gated_mlp = True
275 self.cfg.attn_only = False
276 self.cfg.default_prepend_bos = False
277 self.cfg.default_padding_side = "right"
279 self.weight_processing_conversions = {
280 **self._qkvo_weight_conversions(),
281 }
283 self.component_mapping = {
284 "embed": EmbeddingBridge(name="model.transformer.wte", config=self.cfg),
285 "blocks": _LLaDABlockBridge(
286 name="model.transformer.blocks",
287 config=self.cfg,
288 submodules={
289 "ln1": RMSNormalizationBridge(name="attn_norm", config=self.cfg),
290 "attn": _LLaDAAttentionBridge(
291 name=None,
292 config=self.cfg,
293 is_causal=False,
294 submodules={
295 "q": LinearBridge(name="q_proj"),
296 "k": LinearBridge(name="k_proj"),
297 "v": LinearBridge(name="v_proj"),
298 "o": LinearBridge(name="attn_out"),
299 },
300 ),
301 "ln2": RMSNormalizationBridge(name="ff_norm", config=self.cfg),
302 "mlp": _LLaDAGatedMLPBridge(
303 name=None,
304 config=self.cfg,
305 submodules={
306 "gate": LinearBridge(name="ff_proj"),
307 "in": LinearBridge(name="up_proj"),
308 "act": GeneralizedComponent(name="act"),
309 "out": LinearBridge(name="ff_out"),
310 },
311 ),
312 },
313 ),
314 "ln_final": RMSNormalizationBridge(name="model.transformer.ln_f", config=self.cfg),
315 "unembed": UnembeddingBridge(name="model.transformer.ff_out", config=self.cfg),
316 }
318 @staticmethod
319 def _enum_value(value: Any) -> Any:
320 return getattr(value, "value", value)
322 def _require_value(self, name: str, expected: Any) -> None:
323 if not hasattr(self.cfg, name):
324 raise ValueError(f"LLaDAModelLM config is missing required field '{name}'")
325 actual = self._enum_value(getattr(self.cfg, name))
326 if actual != expected:
327 raise ValueError(f"LLaDAModelLM requires {name}={expected!r}; got {actual!r}")
329 def _validate_config(self) -> None:
330 self._require_value("block_type", "llama")
331 self._require_value("block_group_size", 1)
332 self._require_value("rope", True)
333 self._require_value("rope_full_precision", True)
334 self._require_value("alibi", False)
335 self._require_value("attention_layer_norm", False)
336 self._require_value("include_bias", False)
337 self._require_value("include_qkv_bias", False)
338 self._require_value("scale_logits", False)
339 self._require_value("input_emb_norm", False)
340 self._require_value("layer_norm_type", "rms")
341 if self.cfg.act_fn != "silu":
342 raise ValueError(
343 f"LLaDAModelLM requires activation_type='silu'; got {self.cfg.act_fn!r}"
344 )
345 if bool(self.cfg.tie_word_embeddings):
346 raise ValueError(
347 "LLaDAModelLM tied embeddings are not supported by the initial dense adapter"
348 )
349 embedding_size = int(getattr(self.cfg, "embedding_size", self.cfg.d_vocab))
350 if embedding_size != self.cfg.d_vocab:
351 raise ValueError(
352 "LLaDAModelLM requires embedding_size == vocab_size in the initial "
353 f"dense adapter; got {embedding_size} and {self.cfg.d_vocab}"
354 )
356 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
357 """Disable remote branches incompatible with single-pass hook support."""
358 del model_name
359 config = model_kwargs.get("config")
360 if config is not None: 360 ↛ exitline 360 didn't return from function 'prepare_loading' because the condition on line 360 was always true
361 config.output_attentions = False
362 config.use_cache = False
364 def prepare_model(self, hf_model: Any) -> None:
365 """Keep the wrapper and underlying model on the no-cache path."""
366 for config in (
367 getattr(hf_model, "config", None),
368 getattr(getattr(hf_model, "model", None), "config", None),
369 ):
370 if config is not None: 370 ↛ 366line 370 didn't jump to line 366 because the condition on line 370 was always true
371 config.output_attentions = False
372 config.use_cache = False