Coverage for transformer_lens/model_bridge/supported_architectures/bd3lm.py: 77%
64 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-08-11 18:50 +0000
1"""BD3LM (Block Diffusion Language Model) architecture adapter."""
3from typing import Any
5import torch
7from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
8from transformer_lens.model_bridge.generalized_components import (
9 EmbeddingBridge,
10 LinearBridge,
11 MLPBridge,
12 NormalizationBridge,
13 SymbolicBridge,
14 UnembeddingBridge,
15)
16from transformer_lens.model_bridge.generalized_components.block import (
17 DelegatedAttentionBlockBridge,
18)
21class BD3LMArchitectureAdapter(ArchitectureAdapter):
22 """Architecture adapter for BD3LM (Block Diffusion LM, ICLR 2025).
24 BD3LM uses adaLN conditioning on diffusion timesteps, a custom Rotary
25 embedding, joint QKV projections, and non-causal block-diffusion masking.
26 Because adaLN modulation varies per-timestep, it cannot be folded into
27 weights — the adapter uses ``DelegatedAttentionBlockBridge`` to delegate
28 each ``DDiTBlock.forward()`` wholesale to the original HF module.
29 Hooks fire at block boundaries and on mapped submodules.
30 """
32 # Phases 1–3 cover component mapping, weight conversion, and forward-pass
33 # parity. Phase 4 (autoregressive generation) is excluded because BD3LM
34 # uses iterative diffusion sampling.
35 applicable_phases: list[int] = [1, 2, 3]
37 # BD3LM uses diffusion sampling, not autoregressive generation.
38 supports_generation: bool = False
40 def __init__(self, cfg: Any) -> None:
41 super().__init__(cfg)
43 # ── Config attributes ──────────────────────────────────────────
44 self.cfg.normalization_type = "LN"
45 self.cfg.uses_rms_norm = False
46 self.cfg.positional_embedding_type = "none" # custom Rotary, not HF
47 self.cfg.gated_mlp = False # standard GELU MLP, not gated
48 self.cfg.attn_only = False
49 self.cfg.final_rms = False # final norm is custom LayerNorm
50 self.cfg.tokenizer_name = "gpt2"
52 # BD3LM-specific config fields. These live on the HF config and are
53 # forwarded via _HF_PASSTHROUGH_ATTRS; we also store them explicitly
54 # so unit tests can assert them without loading a real model.
55 block_size = getattr(self.cfg, "block_size", 4)
56 setattr(self.cfg, "block_size", block_size)
58 cond_dim = getattr(self.cfg, "cond_dim", 128)
59 setattr(self.cfg, "cond_dim", cond_dim)
61 adaln = getattr(self.cfg, "adaln", True)
62 setattr(self.cfg, "adaln", adaln)
64 cross_attn = getattr(self.cfg, "cross_attn", True)
65 setattr(self.cfg, "cross_attn", cross_attn)
67 # Compute d_mlp from the MLP ratio (hardcoded to 4 in DDiTBlock).
68 mlp_ratio = 4
69 d_mlp = mlp_ratio * self.cfg.d_model
70 setattr(self.cfg, "d_mlp", d_mlp)
72 # ── Weight processing conversions ──────────────────────────────
73 # Wrap-don't-reimplement: no weight rearrangement needed since we
74 # delegate forward to the original modules.
75 self.weight_processing_conversions = {}
77 # ── Component mapping ──────────────────────────────────────────
78 # DelegatedAttentionBlockBridge delegates forward() wholesale to the
79 # original DDiTBlock (wrap-don't-reimplement), but exposes standard
80 # attn/mlp-shaped hook aliases (hook_resid_mid, hook_mlp_in) instead
81 # of the SSM-specific hook_mixer_in alias that SSMBlockBridge would
82 # have incorrectly implied for this attn+mlp architecture. Submodule
83 # bridges map to the HF module paths relative to each block.
84 #
85 # Module hierarchy (HF paths relative to backbone.blocks[i]):
86 # norm1 → pre-attention LayerNorm
87 # attn_qkv → joint Q/K/V projection (no bias)
88 # attn_out → output projection (no bias)
89 # adaLN_modulation → conditioning linear (cond_dim → 6*dim)
90 # norm2 → pre-MLP LayerNorm
91 # mlp.0 → MLP in-projection (bias)
92 # mlp.2 → MLP out-projection (bias)
93 self.component_mapping = {
94 "embed": EmbeddingBridge(name="backbone.vocab_embed"),
95 "blocks": DelegatedAttentionBlockBridge(
96 name="backbone.blocks",
97 submodules={
98 "ln1": NormalizationBridge(name="norm1", config=self.cfg),
99 "attn": SymbolicBridge(
100 submodules={
101 "qkv": LinearBridge(name="attn_qkv"),
102 "o": LinearBridge(name="attn_out"),
103 },
104 ),
105 "adaln_modulation": LinearBridge(name="adaLN_modulation"),
106 "ln2": NormalizationBridge(name="norm2", config=self.cfg),
107 "mlp": MLPBridge(
108 name="mlp",
109 config=self.cfg,
110 submodules={
111 "in": LinearBridge(name="0"),
112 "out": LinearBridge(name="2"),
113 },
114 ),
115 },
116 # hook_attn_out captures the raw projection, not the gate_msa-scaled
117 # value added to residual stream, as gating is fused inside a
118 # torch.jit.script function with no hookable module boundary.
119 hook_alias_overrides={
120 "hook_attn_out": "attn.o.hook_out",
121 },
122 ),
123 "sigma_map": MLPBridge(
124 name="backbone.sigma_map.mlp",
125 config=self.cfg,
126 submodules={
127 "in": LinearBridge(name="0"),
128 "out": LinearBridge(name="2"),
129 },
130 ),
131 "ln_final": NormalizationBridge(
132 name="backbone.output_layer.norm_final", config=self.cfg
133 ),
134 "unembed": UnembeddingBridge(name="backbone.output_layer.linear"),
135 }
137 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
138 """Patch BD3LM dynamic class before from_pretrained runs.
140 Modeling code has a custom __getattr__ that fails to delegate back
141 to PreTrainedModel, raising AttributeError on all_tied_weights_keys.
142 """
143 try:
144 from transformers.dynamic_module_utils import get_class_from_dynamic_module
146 model_class = get_class_from_dynamic_module("modeling_bd3lm.BD3LM", model_name)
147 setattr(model_class, "all_tied_weights_keys", {})
148 except Exception:
149 # Best-effort patch: should never block loading if dynamic module is missing/modified.
150 pass
152 def prepare_model(self, hf_model: Any) -> None:
153 """Patch BD3LM quirks that prevent standard bridge construction.
155 Three issues must be fixed before the bridge can wrap the model:
157 1. ``vocab_embed`` is an ``nn.Parameter``, not ``nn.Embedding``, so it
158 lacks a ``.weight`` attribute that ``EmbeddingBridge`` expects.
159 2. The ``flex`` attention backend crashes on CPU; fall back to ``sdpa``
160 and regenerate ``block_diff_mask`` for the new backend.
161 3. The HF ``forward()`` does not accept ``output_attentions`` and other
162 kwargs the bridge unconditionally injects; patch at runtime because
163 no other hook point allows filtering them before HF's forward call.
164 """
165 # Patch vocab_embed to have a weight attribute for EmbeddingBridge
166 if hasattr(hf_model, "backbone") and hasattr(hf_model.backbone, "vocab_embed"): 166 ↛ 172line 166 didn't jump to line 172 because the condition on line 166 was always true
167 embed_mod = hf_model.backbone.vocab_embed
168 if not hasattr(embed_mod, "weight") and hasattr(embed_mod, "embedding"):
169 embed_mod.weight = embed_mod.embedding
171 # Fix attention backend for CPU and ensure mask is on a real device
172 if hasattr(hf_model, "backbone"): 172 ↛ 190line 172 didn't jump to line 190 because the condition on line 172 was always true
173 if not torch.cuda.is_available(): 173 ↛ 179line 173 didn't jump to line 179 because the condition on line 173 was always true
174 backend = "sdpa"
175 setattr(self.cfg, "attn_backend", "sdpa")
176 for b in hf_model.backbone.blocks:
177 b.attn_backend = "sdpa"
178 else:
179 first_block = hf_model.backbone.blocks[0]
180 backend = getattr(first_block, "attn_backend", "sdpa")
182 if hasattr(hf_model.backbone, "gen_mask"):
183 hf_model.backbone.gen_mask(
184 getattr(self.cfg, "model_length", getattr(self.cfg, "n_ctx", 2048)),
185 getattr(hf_model.backbone, "block_size", 4),
186 attn_backend=backend,
187 )
189 # Patch the model's forward method to filter out unsupported kwargs (like output_attentions)
190 original_forward = hf_model.forward
191 import inspect
193 def patched_forward(*args, **kwargs):
194 sig = inspect.signature(original_forward)
195 valid_params = set(sig.parameters.keys())
197 # Check if VAR_KEYWORD (like **kwargs) is accepted
198 accepts_var_keyword = any(
199 p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()
200 )
202 if not accepts_var_keyword:
203 kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
204 return original_forward(*args, **kwargs)
206 hf_model.forward = patched_forward
208 def convert_weights(self) -> dict[str, torch.Tensor]:
209 """Return empty dict — delegation means no weight rearrangement."""
210 return {}