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

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`` restores it in the remote 

16modeling module with the v4 semantics (plain inverse-frequency rope, 

17attention factor 1.0). 

18""" 

19 

20from typing import Any 

21 

22import torch 

23 

24from transformer_lens.model_bridge.generalized_components import ( 

25 AttentionBridge, 

26 LinearBridge, 

27) 

28from transformer_lens.model_bridge.supported_architectures._remote_code_compat import ( 

29 force_import_remote_class, 

30 restore_default_rope_init, 

31) 

32from transformer_lens.model_bridge.supported_architectures.qwen2 import ( 

33 Qwen2ArchitectureAdapter, 

34) 

35 

36 

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

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

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

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

41 return 

42 

43 original_forward = attn_cls.forward 

44 

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

46 mask = kwargs.get("attention_mask") 

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

48 kwargs["attention_mask"] = None 

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

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

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

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

53 

54 setattr(attn_cls, "forward", forward) 

55 setattr(attn_cls, "_tl_mask_sentinel_patched", True) 

56 

57 

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

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

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

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

62 return 

63 

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

65 generation_config = cls() 

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

67 value = getattr(model_config, key, None) 

68 if value is not None: 

69 setattr(generation_config, key, value) 

70 generation_config._from_model_config = True 

71 return generation_config 

72 

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

74 setattr(gen_cfg_cls, "_tl_from_model_config_patched", True) 

75 

76 

77class DreamArchitectureAdapter(Qwen2ArchitectureAdapter): 

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

79 

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

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

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

83 supports_generation: bool = False 

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

85 supports_causal_loss: bool = False 

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

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

88 # so bridge hooks fire during sampling. 

89 native_sampler: str = "diffusion_generate" 

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

91 _testing_eager = None 

92 _testing_wire_rotary = False 

93 

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

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

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

97 

98 def _build_attention_bridge(self): 

99 """Bidirectional diffusion attention; the bridge reimplementation 

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

101 return AttentionBridge( 

102 name="self_attn", 

103 config=self.cfg, 

104 submodules={ 

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

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

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

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

109 }, 

110 maintain_native_attention=True, 

111 ) 

112 

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

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

115 restore_default_rope_init( 

116 model_name, 

117 "modeling_dream.DreamModel", 

118 "DreamRotaryEmbedding", 

119 revision=model_kwargs.get("code_revision") or model_kwargs.get("revision"), 

120 ) 

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

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

123 # kwargs-tolerant no-op. 

124 gen_cfg_cls = force_import_remote_class( 

125 model_name, 

126 "generation_utils.DreamGenerationConfig", 

127 revision=model_kwargs.get("code_revision") or model_kwargs.get("revision"), 

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( 

133 model_name, 

134 "modeling_dream.DreamAttention", 

135 revision=model_kwargs.get("code_revision") or model_kwargs.get("revision"), 

136 ) 

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

138 _patch_eager_attention_mask(attn_cls) 

139 super().prepare_loading(model_name, model_kwargs)