Coverage for transformer_lens/model_bridge/supported_architectures/llada2_moe.py: 93%

1"""LLaDA 2.0 MoE architecture adapter. 

2 

3Ant Group's LLaDA 2.x (``LLaDA2MoeModelLM``, remote code): masked 

4block-diffusion language models on a DeepSeek-V3-style MoE decoder — 

5fused ``query_key_value`` attention with full-width query/key layernorms, 

6per-expert routed MLPs behind a bias-corrected router plus shared 

7experts, and dense MLPs on the first ``first_k_dense_replace`` layers. 

8 

9Attention is bidirectional (``is_causal = False``) and generation is 

10block-diffusion sampling via the model's own ``generate``, reached through 

11``bridge.diffusion_generate``; attention delegates to HF and the bridge's 

12autoregressive generation stays off. The fused QKV ships no 

13HookedTransformer-format weight conversions, so LN folding is disabled. 

14 

15The remote forward validates attention_mask strictly: it must be the 4D 

16block-diffusion form (batch, 1, seq, seq) — full-ones for full 

17bidirectional visibility — not a 2D padding mask. A forward pre-hook 

18drops all-ones 2D masks (informationless) and rejects padded ones with a 

19clear error instead of the remote validator's opaque failure. 

20""" 

21 

22from typing import Any 

23 

24from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

25from transformer_lens.model_bridge.generalized_components import ( 

26 AttentionBridge, 

27 BlockBridge, 

28 EmbeddingBridge, 

29 LinearBridge, 

30 MoEBridge, 

31 RMSNormalizationBridge, 

32 RotaryEmbeddingBridge, 

33 UnembeddingBridge, 

34) 

35from transformer_lens.model_bridge.generalized_components.base import ( 

36 GeneralizedComponent, 

37) 

38from transformer_lens.model_bridge.supported_architectures.dream import ( 

39 _register_default_rope_init, 

40) 

41 

42 

43class _LLaDA2FusedAttentionBridge(AttentionBridge): 

44 """Fused query_key_value projection: no separate q/k/v submodules to 

45 alias — expose the fused output instead.""" 

46 

47 hook_aliases = { 

48 "hook_qkv": "qkv.hook_out", 

49 "hook_z": "o.hook_in", 

50 } 

51 

52 

53class LLaDA2MoeArchitectureAdapter(ArchitectureAdapter): 

54 """Architecture adapter for LLaDA2MoeModelLM models.""" 

55 

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

57 supports_generation: bool = False 

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

59 supports_causal_loss: bool = False 

60 # Semi-autoregressive block remasking, shipped on the model class. 

61 native_sampler: str = "generate" 

62 # Fused query_key_value with no per-projection conversions to fold into. 

63 supports_fold_ln = False 

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

65 _testing_eager = None 

66 _testing_wire_rotary = False 

67 

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

69 """gen_length must cover whole blocks; block_length caps at the budget.""" 

70 block_length = min(32, max_new_tokens) 

71 blocks = max(1, -(-max_new_tokens // block_length)) 

72 return { 

73 "gen_length": blocks * block_length, 

74 "block_length": block_length, 

75 "steps": max_new_tokens, 

76 } 

77 

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

79 """Initialize the LLaDA 2.0 MoE architecture adapter.""" 

80 super().__init__(cfg) 

81 

82 self._set_rms_rotary_defaults() 

83 

84 self.weight_processing_conversions = {} 

85 

86 self.component_mapping = { 

87 "embed": EmbeddingBridge(name="model.word_embeddings"), 

88 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"), 

89 "blocks": BlockBridge( 

90 name="model.layers", 

91 config=self.cfg, 

92 submodules={ 

93 "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), 

94 "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg), 

95 # Bidirectional diffusion attention with fused QKV; the 

96 # bridge reimplementation assumes causal masking, so 

97 # delegate to HF. The fused projection and dense output 

98 # are hookable; q/k layernorms are full-width. 

99 "attn": _LLaDA2FusedAttentionBridge( 

100 name="attention", 

101 config=self.cfg, 

102 submodules={ 

103 "qkv": LinearBridge(name="query_key_value"), 

104 "o": LinearBridge(name="dense"), 

105 "q_norm": RMSNormalizationBridge( 

106 name="query_layernorm", config=self.cfg 

107 ), 

108 "k_norm": RMSNormalizationBridge(name="key_layernorm", config=self.cfg), 

109 }, 

110 maintain_native_attention=True, 

111 ), 

112 # Dense on the first first_k_dense_replace layers, routed 

113 # MoE elsewhere — gate and shared_experts are optional so 

114 # setup skips them on dense layers (deepseek_v3 pattern). 

115 "mlp": MoEBridge( 

116 name="mlp", 

117 config=self.cfg, 

118 sparse_required=("gate",), 

119 submodules={ 

120 "gate": GeneralizedComponent(name="gate", optional=True), 

121 "shared_experts": self._gated_mlp(name="shared_experts", optional=True), 

122 # Dense-layer projections (absent on MoE layers). 

123 "dense_gate": LinearBridge(name="gate_proj", optional=True), 

124 "dense_in": LinearBridge(name="up_proj", optional=True), 

125 "dense_out": LinearBridge(name="down_proj", optional=True), 

126 }, 

127 ), 

128 }, 

129 ), 

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

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

132 } 

133 

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

135 """Restore the v4 'default' rope init the remote code looks up (Dream shim).""" 

136 _register_default_rope_init() 

137 super().prepare_loading(model_name, model_kwargs) 

138 

139 def setup_hook_compatibility(self, bridge: Any) -> None: 

140 """Guard the remote forward against auto-passed 2D padding masks.""" 

141 model = getattr(bridge, "original_model", None) 

142 if model is None or getattr(model, "_llada2_mask_guard", False): 142 ↛ anywhereline 142 didn't jump anywhere: it always raised an exception.

143 return 

144 

145 def _mask_guard(module: Any, args: Any, kwargs: Any) -> Any: 

146 import torch 

147 

148 mask = kwargs.get("attention_mask") 

149 if mask is None: 

150 # Remote forward calls attention_mask.size() unconditionally; 

151 # synthesize the full-visibility 4D mask from the input shape. 

152 ids = kwargs.get("input_ids", args[0] if args else None) 

153 if isinstance(ids, torch.Tensor) and ids.ndim == 2: 153 ↛ 156line 153 didn't jump to line 156 because the condition on line 153 was always true

154 batch, seq = ids.shape 

155 kwargs["attention_mask"] = torch.ones(batch, 1, seq, seq, device=ids.device) 

156 return args, kwargs 

157 if isinstance(mask, torch.Tensor) and mask.ndim == 2: 

158 if not bool(mask.all()): 

159 raise NotImplementedError( 

160 "LLaDA2's remote forward rejects 2D padding masks; " 

161 "batched padded inputs are unsupported — pass " 

162 "equal-length sequences." 

163 ) 

164 batch, seq = mask.shape 

165 kwargs["attention_mask"] = torch.ones(batch, 1, seq, seq, device=mask.device) 

166 return args, kwargs 

167 

168 model.register_forward_pre_hook(_mask_guard, with_kwargs=True) 

169 model._llada2_mask_guard = True