Coverage for transformer_lens/model_bridge/supported_architectures/gidd.py: 65%

1"""Gidd architecture adapter. 

2 

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. 

10 

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""" 

17 

18from typing import Any 

19 

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) 

29 

30 

31def restore_frequencies(hf_model: Any) -> bool: 

32 """Recompute GIDD's non-persistent ``frequencies`` rotary table; under v5's 

33 meta-device load it materializes as uninitialized memory that silently corrupts 

34 every forward (applied to both bridge and HF reference so they agree).""" 

35 import sys 

36 

37 inner = getattr(hf_model, "model", None) 

38 old = getattr(inner, "frequencies", None) 

39 if inner is None or old is None: 

40 return False 

41 module = sys.modules.get(type(inner).__module__) 

42 compute = getattr(module, "compute_basic_frequencies", None) 

43 if compute is None: 43 ↛ 44line 43 didn't jump to line 44 because the condition on line 43 was never true

44 return False 

45 config = hf_model.config 

46 freqs = compute( 

47 base=config.rope_theta, 

48 rotary_dim=config.hidden_size // config.num_attention_heads, 

49 max_position_embeddings=config.max_position_embeddings, 

50 ) 

51 inner.frequencies = freqs.to(device=old.device, dtype=old.dtype) 

52 return True 

53 

54 

55class GiddArchitectureAdapter(ArchitectureAdapter): 

56 """Architecture adapter for GiddForDiffusionLM models.""" 

57 

58 applicable_phases: list[int] = [1, 2, 3, 4] 

59 supports_generation: bool = False 

60 # Bidirectional masked-denoising objective; shifted causal CE is undefined. 

61 supports_causal_loss: bool = False 

62 # Block-wise denoising with self-correction, shipped on the model class. 

63 native_sampler: str = "generate" 

64 # ScaledLinear applies a runtime weight scale; folding norms into those 

65 # projections (or centering through scaled residual adds) is unsound. 

66 supports_fold_ln = False 

67 

68 def native_sampler_kwargs(self, max_new_tokens: int, prompt_len: int) -> dict: 

69 """Gidd's max_length counts generated tokens: its windows start at 

70 prompt_length and span max_length, so adding the prompt over-generates.""" 

71 return { 

72 "max_length": max_new_tokens, 

73 "block_length": min(128, max_new_tokens), 

74 "steps": max_new_tokens, 

75 } 

76 

77 def __init__(self, cfg: Any) -> None: 

78 """Initialize the Gidd architecture adapter.""" 

79 super().__init__(cfg) 

80 

81 self.cfg.normalization_type = "RMS" 

82 self.cfg.uses_rms_norm = True 

83 self.cfg.positional_embedding_type = "rotary" 

84 self.cfg.gated_mlp = False # ungated up/down MLP 

85 self.cfg.attn_only = False 

86 self.cfg.final_rms = True 

87 

88 self.weight_processing_conversions = {} 

89 

90 self.component_mapping = { 

91 "embed": EmbeddingBridge(name="model.embed_tokens"), 

92 "blocks": BlockBridge( 

93 name="model.layers", 

94 config=self.cfg, 

95 submodules={ 

96 "ln1": RMSNormalizationBridge(name="attn_layernorm", config=self.cfg), 

97 "ln2": RMSNormalizationBridge(name="mlp_layernorm", config=self.cfg), 

98 # Bidirectional softcap attention: delegate; QK norms only 

99 # exist when use_qk_norm is set. 

100 "attn": AttentionBridge( 

101 name="self_attn", 

102 config=self.cfg, 

103 submodules={ 

104 "q": LinearBridge(name="q_proj"), 

105 "k": LinearBridge(name="k_proj"), 

106 "v": LinearBridge(name="v_proj"), 

107 "o": LinearBridge(name="o_proj"), 

108 "q_norm": RMSNormalizationBridge( 

109 name="q_norm", config=self.cfg, optional=True 

110 ), 

111 "k_norm": RMSNormalizationBridge( 

112 name="k_norm", config=self.cfg, optional=True 

113 ), 

114 }, 

115 maintain_native_attention=True, 

116 ), 

117 "mlp": self._ungated_mlp(), 

118 }, 

119 ), 

120 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), 

121 "unembed": UnembeddingBridge(name="lm_head"), 

122 } 

123 

124 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None: 

125 """Patch the remote class before from_pretrained runs. 

126 

127 Like BD3LM, the remote code's attribute handling raises on v5's 

128 all_tied_weights_keys lookup (the checkpoint is untied anyway). 

129 """ 

130 try: 

131 from transformers.dynamic_module_utils import get_class_from_dynamic_module 

132 

133 model_class = get_class_from_dynamic_module( 

134 "modeling_gidd.GiddForDiffusionLM", model_name 

135 ) 

136 setattr(model_class, "all_tied_weights_keys", {}) 

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 # Skip modules whose params are already real (internlm2 pattern). 

141 pretrained_cls = model_class.__mro__[1] 

142 if not getattr(pretrained_cls, "_tl_patched", False): 

143 original_init_weights = getattr(pretrained_cls, "_init_weights") 

144 

145 def safe_init_weights(self, mod, _original=original_init_weights): 

146 first_param = next(mod.parameters(), None) 

147 if first_param is not None and first_param.device.type != "meta": 

148 return 

149 _original(self, mod) 

150 

151 setattr(pretrained_cls, "_init_weights", safe_init_weights) 

152 setattr(pretrained_cls, "_tl_patched", True) 

153 except Exception: 

154 pass 

155 super().prepare_loading(model_name, model_kwargs) 

156 

157 def prepare_model(self, hf_model: Any) -> None: 

158 """Restore the rotary table lost to meta-device loading.""" 

159 super().prepare_model(hf_model) 

160 restore_frequencies(hf_model) 

161 

162 def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None: 

163 """Delegated attention reads the rotary buffer inside HF; nothing to wire."""