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

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._remote_code_compat import ( 

28 compute_default_rope_inv_freq, 

29 force_import_remote_class, 

30) 

31from transformer_lens.model_bridge.supported_architectures.qwen2 import ( 

32 Qwen2ArchitectureAdapter, 

33) 

34 

35 

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

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

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

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

40 return 

41 

42 original_forward = attn_cls.forward 

43 

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

45 mask = kwargs.get("attention_mask") 

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

47 kwargs["attention_mask"] = None 

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

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

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

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

52 

53 setattr(attn_cls, "forward", forward) 

54 setattr(attn_cls, "_tl_mask_sentinel_patched", True) 

55 

56 

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

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

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

60 if getattr(gen_cfg_cls, "_tl_from_model_config_patched", False): 60 ↛ anywhereline 60 didn't jump anywhere: it always raised an exception.

61 return 

62 

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

64 generation_config = cls() 

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

66 value = getattr(model_config, key, None) 

67 if value is not None: 

68 setattr(generation_config, key, value) 

69 generation_config._from_model_config = True 

70 return generation_config 

71 

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

73 setattr(gen_cfg_cls, "_tl_from_model_config_patched", True) 

74 

75 

76def _register_default_rope_init() -> None: 

77 """Restore the global ``ROPE_INIT_FUNCTIONS["default"]`` entry v5 removed; 

78 Dream's remote code (and llada2_moe's) looks it up by that key.""" 

79 from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS 

80 

81 ROPE_INIT_FUNCTIONS.setdefault("default", compute_default_rope_inv_freq) 

82 

83 

84class DreamArchitectureAdapter(Qwen2ArchitectureAdapter): 

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

86 

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

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

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

90 supports_generation: bool = False 

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

92 supports_causal_loss: bool = False 

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

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

95 # so bridge hooks fire during sampling. 

96 native_sampler: str = "diffusion_generate" 

97 # Delegated attention computes rotary inside HF; nothing to wire. 

98 _testing_eager = None 

99 _testing_wire_rotary = False 

100 

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

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

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

104 

105 def _build_attention_bridge(self): 

106 """Bidirectional diffusion attention; the bridge reimplementation 

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

108 return AttentionBridge( 

109 name="self_attn", 

110 config=self.cfg, 

111 submodules={ 

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

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

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

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

116 }, 

117 maintain_native_attention=True, 

118 ) 

119 

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

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

122 _register_default_rope_init() 

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

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

125 # kwargs-tolerant no-op. 

126 gen_cfg_cls = force_import_remote_class( 

127 model_name, "generation_utils.DreamGenerationConfig" 

128 ) 

129 if gen_cfg_cls is not None: 129 ↛ 132line 129 didn't jump to line 132 because the condition on line 129 was always true

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

131 _patch_from_model_config(gen_cfg_cls) 

132 attn_cls = force_import_remote_class(model_name, "modeling_dream.DreamAttention") 

133 if attn_cls is not None: 133 ↛ 135line 133 didn't jump to line 135 because the condition on line 133 was always true

134 _patch_eager_attention_mask(attn_cls) 

135 super().prepare_loading(model_name, model_kwargs)