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

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 # Block-wise denoising with self-correction, shipped on the model class. 

61 native_sampler: str = "generate" 

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

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

64 supports_fold_ln = False 

65 

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

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

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

69 return { 

70 "max_length": max_new_tokens, 

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

72 "steps": max_new_tokens, 

73 } 

74 

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

76 """Initialize the Gidd architecture adapter.""" 

77 super().__init__(cfg) 

78 

79 self.cfg.normalization_type = "RMS" 

80 self.cfg.uses_rms_norm = True 

81 self.cfg.positional_embedding_type = "rotary" 

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

83 self.cfg.attn_only = False 

84 self.cfg.final_rms = True 

85 

86 self.weight_processing_conversions = {} 

87 

88 self.component_mapping = { 

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

90 "blocks": BlockBridge( 

91 name="model.layers", 

92 config=self.cfg, 

93 submodules={ 

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

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

96 # Bidirectional softcap attention: delegate; QK norms only 

97 # exist when use_qk_norm is set. 

98 "attn": AttentionBridge( 

99 name="self_attn", 

100 config=self.cfg, 

101 submodules={ 

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

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

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

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

106 "q_norm": RMSNormalizationBridge( 

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

108 ), 

109 "k_norm": RMSNormalizationBridge( 

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

111 ), 

112 }, 

113 maintain_native_attention=True, 

114 ), 

115 "mlp": self._ungated_mlp(), 

116 }, 

117 ), 

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

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

120 } 

121 

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

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

124 

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

126 all_tied_weights_keys lookup (the checkpoint is untied anyway). 

127 """ 

128 try: 

129 from transformers.dynamic_module_utils import get_class_from_dynamic_module 

130 

131 model_class = get_class_from_dynamic_module( 

132 "modeling_gidd.GiddForDiffusionLM", model_name 

133 ) 

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

135 # v5 walks _init_weights over every module post-materialization; 

136 # the remote _init_weights assumes module.weight exists (crashes 

137 # on containers) and would re-randomize loaded tensors anyway. 

138 # Skip modules whose params are already real (internlm2 pattern). 

139 pretrained_cls = model_class.__mro__[1] 

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

141 original_init_weights = getattr(pretrained_cls, "_init_weights") 

142 

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

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

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

146 return 

147 _original(self, mod) 

148 

149 setattr(pretrained_cls, "_init_weights", safe_init_weights) 

150 setattr(pretrained_cls, "_tl_patched", True) 

151 except Exception: 

152 pass 

153 super().prepare_loading(model_name, model_kwargs) 

154 

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

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

157 super().prepare_model(hf_model) 

158 restore_frequencies(hf_model) 

159 

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

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