Coverage for transformer_lens/model_bridge/supported_architectures/gidd.py: 73%
45 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Gidd architecture adapter.
3Dimitri von Rütte's GIDD (``GiddForDiffusionLM``, remote code): the only
4open uniform-noise (non-masked) diffusion LM at scale, with self-correction
5sampling. The decoder is bidirectional (config.is_causal=False) with
6softcap attention variants, optional per-head QK norms, ScaledLinear
7projections (weight-scaled at forward), per-layer scaled residual adds
8(resid_scale/num_layers), an ungated up/down MLP, and rotary positions
9held as a model-level buffer rather than a module.
11Everything nonstandard lives inside delegated modules: attention delegates
12wholesale (softcap + bidirectional), ScaledLinear wraps as plain hookable
13Linears, and generation is the model's own diffusion sampler, reached via
14``bridge.diffusion_generate`` — no autoregressive generation, no folding
15into scaled projections.
16"""
18from typing import Any
20from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
21from transformer_lens.model_bridge.generalized_components import (
22 AttentionBridge,
23 BlockBridge,
24 EmbeddingBridge,
25 LinearBridge,
26 RMSNormalizationBridge,
27 UnembeddingBridge,
28)
29from transformer_lens.model_bridge.supported_architectures._remote_code_compat import (
30 disable_tied_weights_lookup,
31 force_import_remote_class,
32 patch_init_weights_skip_loaded,
33)
36def restore_frequencies(hf_model: Any) -> bool:
37 """Recompute GIDD's non-persistent ``frequencies`` rotary table; under v5's
38 meta-device load it materializes as uninitialized memory that silently corrupts
39 every forward (applied to both bridge and HF reference so they agree)."""
40 import sys
42 inner = getattr(hf_model, "model", None)
43 old = getattr(inner, "frequencies", None)
44 if inner is None or old is None:
45 return False
46 module = sys.modules.get(type(inner).__module__)
47 compute = getattr(module, "compute_basic_frequencies", None)
48 if compute is None: 48 ↛ 49line 48 didn't jump to line 49 because the condition on line 48 was never true
49 return False
50 config = hf_model.config
51 freqs = compute(
52 base=config.rope_theta,
53 rotary_dim=config.hidden_size // config.num_attention_heads,
54 max_position_embeddings=config.max_position_embeddings,
55 )
56 inner.frequencies = freqs.to(device=old.device, dtype=old.dtype)
57 return True
60class GiddArchitectureAdapter(ArchitectureAdapter):
61 """Architecture adapter for GiddForDiffusionLM models."""
63 applicable_phases: list[int] = [1, 2, 3, 4]
64 supports_generation: bool = False
65 # Bidirectional masked-denoising objective; shifted causal CE is undefined.
66 supports_causal_loss: bool = False
67 # Block-wise denoising with self-correction, shipped on the model class.
68 native_sampler: str = "generate"
69 # ScaledLinear applies a runtime weight scale; folding norms into those
70 # projections (or centering through scaled residual adds) is unsound.
71 supports_fold_ln = False
72 # Delegated attention reads the rotary buffer inside HF; nothing to wire.
73 _testing_eager = None
74 _testing_wire_rotary = False
76 def native_sampler_kwargs(self, max_new_tokens: int, prompt_len: int) -> dict:
77 """Gidd's max_length counts generated tokens: its windows start at
78 prompt_length and span max_length, so adding the prompt over-generates."""
79 return {
80 "max_length": max_new_tokens,
81 "block_length": min(128, max_new_tokens),
82 "steps": max_new_tokens,
83 }
85 def __init__(self, cfg: Any) -> None:
86 """Initialize the Gidd architecture adapter."""
87 super().__init__(cfg)
89 self._set_rms_rotary_defaults(gated=False)
91 self.weight_processing_conversions = {}
93 self.component_mapping = {
94 "embed": EmbeddingBridge(name="model.embed_tokens"),
95 "blocks": BlockBridge(
96 name="model.layers",
97 config=self.cfg,
98 submodules={
99 "ln1": RMSNormalizationBridge(name="attn_layernorm", config=self.cfg),
100 "ln2": RMSNormalizationBridge(name="mlp_layernorm", config=self.cfg),
101 # Bidirectional softcap attention: delegate; QK norms only
102 # exist when use_qk_norm is set.
103 "attn": AttentionBridge(
104 name="self_attn",
105 config=self.cfg,
106 submodules={
107 "q": LinearBridge(name="q_proj"),
108 "k": LinearBridge(name="k_proj"),
109 "v": LinearBridge(name="v_proj"),
110 "o": LinearBridge(name="o_proj"),
111 "q_norm": RMSNormalizationBridge(
112 name="q_norm", config=self.cfg, optional=True
113 ),
114 "k_norm": RMSNormalizationBridge(
115 name="k_norm", config=self.cfg, optional=True
116 ),
117 },
118 maintain_native_attention=True,
119 ),
120 "mlp": self._ungated_mlp(),
121 },
122 ),
123 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg),
124 "unembed": UnembeddingBridge(name="lm_head"),
125 }
127 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
128 """Patch the remote class before from_pretrained runs.
130 Like BD3LM, the remote code's attribute handling raises on v5's
131 all_tied_weights_keys lookup (the checkpoint is untied anyway).
132 """
133 try:
134 model_class = force_import_remote_class(model_name, "modeling_gidd.GiddForDiffusionLM")
135 if model_class is not None:
136 disable_tied_weights_lookup(model_class)
137 # v5 walks _init_weights over every module post-materialization;
138 # the remote _init_weights assumes module.weight exists (crashes
139 # on containers) and would re-randomize loaded tensors anyway.
140 # The remote pretrained base sits one MRO step above the LM class.
141 patch_init_weights_skip_loaded(model_class.__mro__[1])
142 except Exception:
143 pass
144 super().prepare_loading(model_name, model_kwargs)
146 def prepare_model(self, hf_model: Any) -> None:
147 """Restore the rotary table lost to meta-device loading."""
148 super().prepare_model(hf_model)
149 restore_frequencies(hf_model)