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

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 # Semi-autoregressive block remasking, shipped on the model class. 

59 native_sampler: str = "generate" 

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

61 supports_fold_ln = False 

62 

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

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

65 block_length = min(32, max_new_tokens) 

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

67 return { 

68 "gen_length": blocks * block_length, 

69 "block_length": block_length, 

70 "steps": max_new_tokens, 

71 } 

72 

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

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

75 super().__init__(cfg) 

76 

77 self._set_rms_rotary_defaults() 

78 

79 self.weight_processing_conversions = {} 

80 

81 self.component_mapping = { 

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

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

84 "blocks": BlockBridge( 

85 name="model.layers", 

86 config=self.cfg, 

87 submodules={ 

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

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

90 # Bidirectional diffusion attention with fused QKV; the 

91 # bridge reimplementation assumes causal masking, so 

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

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

94 "attn": _LLaDA2FusedAttentionBridge( 

95 name="attention", 

96 config=self.cfg, 

97 submodules={ 

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

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

100 "q_norm": RMSNormalizationBridge( 

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

102 ), 

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

104 }, 

105 maintain_native_attention=True, 

106 ), 

107 # Dense on the first first_k_dense_replace layers, routed 

108 # MoE elsewhere — gate and shared_experts are optional so 

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

110 "mlp": MoEBridge( 

111 name="mlp", 

112 config=self.cfg, 

113 sparse_required=("gate",), 

114 submodules={ 

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

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

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

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

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

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

121 }, 

122 ), 

123 }, 

124 ), 

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

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

127 } 

128 

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

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

131 _register_default_rope_init() 

132 super().prepare_loading(model_name, model_kwargs) 

133 

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

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

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

137 if model is None or getattr(model, "_llada2_mask_guard", False): 137 ↛ 138line 137 didn't jump to line 138 because the condition on line 137 was never true

138 return 

139 

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

141 import torch 

142 

143 mask = kwargs.get("attention_mask") 

144 if mask is None: 

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

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

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

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

149 batch, seq = ids.shape 

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

151 return args, kwargs 

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

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

154 raise NotImplementedError( 

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

156 "batched padded inputs are unsupported — pass " 

157 "equal-length sequences." 

158 ) 

159 batch, seq = mask.shape 

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

161 return args, kwargs 

162 

163 model.register_forward_pre_hook(_mask_guard, with_kwargs=True) 

164 model._llada2_mask_guard = True 

165 

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

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