Coverage for transformer_lens/model_bridge/supported_architectures/jetmoe.py: 100%

1"""JetMoE architecture adapter. 

2 

3MIT-IBM's JetMoE (``JetMoeForCausalLM``, native in transformers): the only 

4open at-scale Mixture-of-Attention-heads model — attention Q and output 

5projections are per-expert parallel 3D tensors behind a top-k router 

6(``experts``: JetMoeMoA), with a shared fused KV projection, alongside a 

7conventional parallel-experts MoE MLP. Both routers are hookable; the 

8mixers delegate to HF (per-expert 3D projections have no uniform 

9reconstruction, so no fold target exists either). 

10""" 

11 

12from typing import Any 

13 

14from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

15from transformer_lens.model_bridge.generalized_components import ( 

16 AttentionBridge, 

17 BlockBridge, 

18 EmbeddingBridge, 

19 LinearBridge, 

20 MoEBridge, 

21 MoERouterBridge, 

22 RMSNormalizationBridge, 

23 RotaryEmbeddingBridge, 

24 UnembeddingBridge, 

25) 

26from transformer_lens.model_bridge.generalized_components.base import ( 

27 GeneralizedComponent, 

28) 

29 

30 

31class _JetMoeAttentionBridge(AttentionBridge): 

32 """Mixture-of-Attention: no separate q/k/v/o Linears to alias — Q and O 

33 live inside the per-expert MoA; only the shared fused KV is a Linear.""" 

34 

35 hook_aliases = { 

36 "hook_kv": "kv.hook_out", 

37 } 

38 

39 

40class JetMoeArchitectureAdapter(ArchitectureAdapter): 

41 """Architecture adapter for JetMoeForCausalLM models.""" 

42 

43 # Per-expert 3D Q/O projections: nothing to fold a norm into. 

44 supports_fold_ln = False 

45 # TopKGating's forward sorts/scatters expert assignments and crashes on 

46 # the harness's isolated probes; routers stay hookable at runtime. 

47 component_test_skip_suffixes = ("mlp.gate", "attn.experts.router") 

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

49 _testing_eager = None 

50 _testing_wire_rotary = False 

51 

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

53 """Initialize the JetMoE architecture adapter.""" 

54 super().__init__(cfg) 

55 

56 self._set_rms_rotary_defaults() 

57 

58 self.weight_processing_conversions = {} 

59 

60 self.component_mapping = { 

61 "embed": EmbeddingBridge(name="model.embed_tokens"), 

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

63 "blocks": BlockBridge( 

64 name="model.layers", 

65 config=self.cfg, 

66 submodules={ 

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

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

69 # MoA: delegate; the attention router and shared KV are 

70 # hookable, per-expert Q/O stay inside the delegated MoA. 

71 "attn": _JetMoeAttentionBridge( 

72 name="self_attention", 

73 config=self.cfg, 

74 submodules={ 

75 "kv": LinearBridge(name="kv_proj"), 

76 "experts": GeneralizedComponent( 

77 name="experts", 

78 submodules={ 

79 # JetMoeTopKGating puts logits last in its 5-tuple. 

80 "router": MoERouterBridge( 

81 name="router", 

82 logits_index=-1, 

83 weights_index=None, 

84 indices_index=None, 

85 ), 

86 }, 

87 ), 

88 }, 

89 maintain_native_attention=True, 

90 ), 

91 "mlp": MoEBridge( 

92 name="mlp", 

93 config=self.cfg, 

94 submodules={ 

95 "gate": MoERouterBridge( 

96 name="router", 

97 logits_index=-1, 

98 weights_index=None, 

99 indices_index=None, 

100 ), 

101 }, 

102 ), 

103 }, 

104 ), 

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

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

107 }