Coverage for transformer_lens/model_bridge/supported_architectures/dream.py: 61%

1"""Dream diffusion LM architecture adapter. 

2 

3HKU-NLP's Dream 7B (``DreamModel``, remote code; also Apple's DiffuCoder): 

4a discrete-diffusion text model initialized from Qwen2.5, so the module 

5tree is exactly Qwen2 (biased q/k/v, gated SiLU MLP, RMS norms, shared 

6``model.rotary_emb``, untied ``lm_head``) — but attention is fully 

7bidirectional (``is_causal = False``) and generation is iterative 

8denoising via ``diffusion_generate``, not autoregressive decoding. 

9 

10Attention is therefore delegated to HF wholesale: the bridge's 

11reimplemented attention assumes causal masking. Q/K/V/O hooks fire on the 

12wrapped projections; there is no reconstructed pattern hook. 

13 

14The remote code targets transformers 4.46; v5 removed the "default" key 

15from ``ROPE_INIT_FUNCTIONS``, so ``prepare_loading`` re-registers it with 

16the v4 semantics (plain inverse-frequency rope, attention factor 1.0). 

17""" 

18 

19from typing import Any 

20 

21import torch 

22 

23from transformer_lens.model_bridge.generalized_components import ( 

24 AttentionBridge, 

25 LinearBridge, 

26) 

27from transformer_lens.model_bridge.supported_architectures.qwen2 import ( 

28 Qwen2ArchitectureAdapter, 

29) 

30 

31 

32def _patch_eager_attention_mask(attn_cls: Any) -> None: 

33 """Teach Dream's eager attention the ``"full"`` mask sentinel -- the bridge forces 

34 eager, whose path (unlike SDPA) raises on the non-tensor sentinel, so normalize it to None.""" 

35 if getattr(attn_cls, "_tl_mask_sentinel_patched", False): 35 ↛ 36line 35 didn't jump to line 36 because the condition on line 35 was never true

36 return 

37 

38 original_forward = attn_cls.forward 

39 

40 def forward(self: Any, *args: Any, **kwargs: Any) -> Any: 

41 mask = kwargs.get("attention_mask") 

42 if mask is not None and not isinstance(mask, torch.Tensor): 

43 kwargs["attention_mask"] = None 

44 elif args and len(args) > 1 and args[1] is not None: 

45 if not isinstance(args[1], torch.Tensor): 

46 args = (args[0], None) + args[2:] 

47 return original_forward(self, *args, **kwargs) 

48 

49 setattr(attn_cls, "forward", forward) 

50 setattr(attn_cls, "_tl_mask_sentinel_patched", True) 

51 

52 

53def _patch_from_model_config(gen_cfg_cls: Any) -> None: 

54 """Rebuild DreamGenerationConfig directly -- v5's from_model_config raises on Dream's 

55 diffusion fields, and rebuilding preserves the null-by-default ``mask_token_id``.""" 

56 if getattr(gen_cfg_cls, "_tl_from_model_config_patched", False): 56 ↛ 57line 56 didn't jump to line 57 because the condition on line 56 was never true

57 return 

58 

59 def from_model_config(cls: Any, model_config: Any) -> Any: 

60 generation_config = cls() 

61 for key in ("bos_token_id", "eos_token_id", "pad_token_id", "mask_token_id"): 

62 value = getattr(model_config, key, None) 

63 if value is not None: 

64 setattr(generation_config, key, value) 

65 generation_config._from_model_config = True 

66 return generation_config 

67 

68 setattr(gen_cfg_cls, "from_model_config", classmethod(from_model_config)) 

69 setattr(gen_cfg_cls, "_tl_from_model_config_patched", True) 

70 

71 

72def _v4_default_rope_parameters( 

73 config: Any = None, device: Any = None, seq_len: Any = None, **rope_kwargs: Any 

74) -> tuple: 

75 """transformers 4.x ``_compute_default_rope_parameters``, removed in v5.""" 

76 if config is not None: 76 ↛ 84line 76 didn't jump to line 84 because the condition on line 76 was always true

77 base = config.rope_theta 

78 partial = getattr(config, "partial_rotary_factor", 1.0) 

79 head_dim = getattr(config, "head_dim", None) or ( 

80 config.hidden_size // config.num_attention_heads 

81 ) 

82 dim = int(head_dim * partial) 

83 else: 

84 base = rope_kwargs["base"] 

85 dim = rope_kwargs["dim"] 

86 inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim)) 

87 return inv_freq, 1.0 

88 

89 

90def _register_default_rope_init() -> None: 

91 from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS 

92 

93 ROPE_INIT_FUNCTIONS.setdefault("default", _v4_default_rope_parameters) 

94 

95 

96class DreamArchitectureAdapter(Qwen2ArchitectureAdapter): 

97 """Architecture adapter for DreamModel diffusion LMs.""" 

98 

99 # Sampling is diffusion, not autoregressive; P4 scores the native 

100 # sampler's text (benchmarks route through diffusion_generate). 

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

102 supports_generation: bool = False 

103 # Sampling is iterative denoising, not left-to-right; Dream ships the 

104 # schedule as a mixin method whose per-step forward goes through __call__, 

105 # so bridge hooks fire during sampling. 

106 native_sampler: str = "diffusion_generate" 

107 

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

109 """Dream denoises a fixed-length canvas; one step per token is its default ratio.""" 

110 return {"max_new_tokens": max_new_tokens, "steps": max_new_tokens} 

111 

112 def _build_attention_bridge(self): 

113 """Bidirectional diffusion attention; the bridge reimplementation 

114 assumes causal masking, so delegate to HF.""" 

115 return AttentionBridge( 

116 name="self_attn", 

117 config=self.cfg, 

118 submodules={ 

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

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

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

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

123 }, 

124 maintain_native_attention=True, 

125 ) 

126 

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

128 """Shim the remote code's two transformers-v4 dependencies.""" 

129 _register_default_rope_init() 

130 # DreamGenerationConfig.validate is a no-op with the v4 signature 

131 # (is_init=False); v5 passes user_set_attributes. Replace with a 

132 # kwargs-tolerant no-op. 

133 try: 

134 from transformers.dynamic_module_utils import get_class_from_dynamic_module 

135 

136 gen_cfg_cls = get_class_from_dynamic_module( 

137 "generation_utils.DreamGenerationConfig", model_name 

138 ) 

139 setattr(gen_cfg_cls, "validate", lambda self, *args, **kwargs: None) 

140 _patch_from_model_config(gen_cfg_cls) 

141 _patch_eager_attention_mask( 

142 get_class_from_dynamic_module("modeling_dream.DreamAttention", model_name) 

143 ) 

144 except Exception: 

145 pass 

146 super().prepare_loading(model_name, model_kwargs) 

147 

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

149 """Delegated attention computes rotary inside HF; nothing to wire."""