Coverage for transformer_lens/model_bridge/supported_architectures/bd3lm.py: 76%
62 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +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)
19from transformer_lens.model_bridge.supported_architectures._remote_code_compat import (
20 disable_tied_weights_lookup,
21 force_import_remote_class,
22)
25class BD3LMArchitectureAdapter(ArchitectureAdapter):
26 """Architecture adapter for BD3LM (Block Diffusion LM, ICLR 2025).
28 BD3LM uses adaLN conditioning on diffusion timesteps, a custom Rotary
29 embedding, joint QKV projections, and non-causal block-diffusion masking.
30 Because adaLN modulation varies per-timestep, it cannot be folded into
31 weights — the adapter uses ``DelegatedAttentionBlockBridge`` to delegate
32 each ``DDiTBlock.forward()`` wholesale to the original HF module.
33 Hooks fire at block boundaries and on mapped submodules.
34 """
36 # Phases 1–3 cover component mapping, weight conversion, and forward-pass
37 # parity. Phase 4 (autoregressive generation) is excluded because BD3LM
38 # uses iterative diffusion sampling.
39 applicable_phases: list[int] = [1, 2, 3]
41 # BD3LM uses diffusion sampling, not autoregressive generation.
42 supports_generation: bool = False
44 def __init__(self, cfg: Any) -> None:
45 super().__init__(cfg)
47 # ── Config attributes ──────────────────────────────────────────
48 self.cfg.normalization_type = "LN"
49 self.cfg.uses_rms_norm = False
50 self.cfg.positional_embedding_type = "none" # custom Rotary, not HF
51 self.cfg.gated_mlp = False # standard GELU MLP, not gated
52 self.cfg.attn_only = False
53 self.cfg.final_rms = False # final norm is custom LayerNorm
54 self.cfg.tokenizer_name = "gpt2"
56 # BD3LM-specific config fields. These live on the HF config and are
57 # forwarded via _HF_PASSTHROUGH_ATTRS; we also store them explicitly
58 # so unit tests can assert them without loading a real model.
59 block_size = getattr(self.cfg, "block_size", 4)
60 setattr(self.cfg, "block_size", block_size)
62 cond_dim = getattr(self.cfg, "cond_dim", 128)
63 setattr(self.cfg, "cond_dim", cond_dim)
65 adaln = getattr(self.cfg, "adaln", True)
66 setattr(self.cfg, "adaln", adaln)
68 cross_attn = getattr(self.cfg, "cross_attn", True)
69 setattr(self.cfg, "cross_attn", cross_attn)
71 # Compute d_mlp from the MLP ratio (hardcoded to 4 in DDiTBlock).
72 mlp_ratio = 4
73 d_mlp = mlp_ratio * self.cfg.d_model
74 setattr(self.cfg, "d_mlp", d_mlp)
76 # ── Weight processing conversions ──────────────────────────────
77 # Wrap-don't-reimplement: no weight rearrangement needed since we
78 # delegate forward to the original modules.
79 self.weight_processing_conversions = {}
81 # ── Component mapping ──────────────────────────────────────────
82 # DelegatedAttentionBlockBridge delegates forward() wholesale to the
83 # original DDiTBlock (wrap-don't-reimplement), but exposes standard
84 # attn/mlp-shaped hook aliases (hook_resid_mid, hook_mlp_in) instead
85 # of the SSM-specific hook_mixer_in alias that SSMBlockBridge would
86 # have incorrectly implied for this attn+mlp architecture. Submodule
87 # bridges map to the HF module paths relative to each block.
88 #
89 # Module hierarchy (HF paths relative to backbone.blocks[i]):
90 # norm1 → pre-attention LayerNorm
91 # attn_qkv → joint Q/K/V projection (no bias)
92 # attn_out → output projection (no bias)
93 # adaLN_modulation → conditioning linear (cond_dim → 6*dim)
94 # norm2 → pre-MLP LayerNorm
95 # mlp.0 → MLP in-projection (bias)
96 # mlp.2 → MLP out-projection (bias)
97 self.component_mapping = {
98 "embed": EmbeddingBridge(name="backbone.vocab_embed"),
99 "blocks": DelegatedAttentionBlockBridge(
100 name="backbone.blocks",
101 submodules={
102 "ln1": NormalizationBridge(name="norm1", config=self.cfg),
103 "attn": SymbolicBridge(
104 submodules={
105 "qkv": LinearBridge(name="attn_qkv"),
106 "o": LinearBridge(name="attn_out"),
107 },
108 ),
109 "adaln_modulation": LinearBridge(name="adaLN_modulation"),
110 "ln2": NormalizationBridge(name="norm2", config=self.cfg),
111 "mlp": MLPBridge(
112 name="mlp",
113 config=self.cfg,
114 submodules={
115 "in": LinearBridge(name="0"),
116 "out": LinearBridge(name="2"),
117 },
118 ),
119 },
120 # hook_attn_out captures the raw projection, not the gate_msa-scaled
121 # value added to residual stream, as gating is fused inside a
122 # torch.jit.script function with no hookable module boundary.
123 hook_alias_overrides={
124 "hook_attn_out": "attn.o.hook_out",
125 },
126 ),
127 "sigma_map": MLPBridge(
128 name="backbone.sigma_map.mlp",
129 config=self.cfg,
130 submodules={
131 "in": LinearBridge(name="0"),
132 "out": LinearBridge(name="2"),
133 },
134 ),
135 "ln_final": NormalizationBridge(
136 name="backbone.output_layer.norm_final", config=self.cfg
137 ),
138 "unembed": UnembeddingBridge(name="backbone.output_layer.linear"),
139 }
141 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
142 """Patch BD3LM dynamic class before from_pretrained runs.
144 Modeling code has a custom __getattr__ that fails to delegate back
145 to PreTrainedModel, raising AttributeError on all_tied_weights_keys.
146 """
147 # Best-effort: never block loading if the dynamic module is missing/modified.
148 model_class = force_import_remote_class(model_name, "modeling_bd3lm.BD3LM")
149 if model_class is not None:
150 disable_tied_weights_lookup(model_class)
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 {}