Coverage for transformer_lens/model_bridge/supported_architectures/phimoe.py: 88%

1"""PhiMoE architecture adapter.""" 

2 

3from typing import Any 

4 

5from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

6from transformer_lens.model_bridge.generalized_components import ( 

7 AttentionBridge, 

8 BlockBridge, 

9 EmbeddingBridge, 

10 LinearBridge, 

11 MoEBridge, 

12 MoERouterBridge, 

13 NormalizationBridge, 

14 UnembeddingBridge, 

15) 

16 

17 

18class PhiMoEArchitectureAdapter(ArchitectureAdapter): 

19 """Architecture adapter for Microsoft PhiMoE models. 

20 

21 PhiMoE is a Phi-style decoder with LayerNorm, split Q/K/V attention, and a 

22 sparse MoE block. This adapter targets the native Transformers implementation 

23 (``trust_remote_code=False``); the archived remote implementation is not 

24 compatible with modern Transformers generation/cache semantics. 

25 """ 

26 

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

28 """Initialize the PhiMoE architecture adapter.""" 

29 super().__init__(cfg) 

30 

31 self.cfg.normalization_type = "LN" 

32 self.cfg.positional_embedding_type = "rotary" 

33 self.cfg.final_rms = False 

34 self.cfg.gated_mlp = True 

35 self.cfg.attn_only = False 

36 self.cfg.uses_rms_norm = False 

37 self.cfg.attn_implementation = "eager" 

38 self.cfg.default_prepend_bos = False 

39 

40 if hasattr(cfg, "num_experts"): 40 ↛ 42line 40 didn't jump to line 42 because the condition on line 40 was always true

41 self.cfg.num_experts = cfg.num_experts 

42 if hasattr(cfg, "experts_per_token"): 42 ↛ 44line 42 didn't jump to line 44 because the condition on line 42 was always true

43 self.cfg.experts_per_token = cfg.experts_per_token 

44 if hasattr(cfg, "router_jitter_noise"): 

45 setattr(self.cfg, "router_jitter_noise", cfg.router_jitter_noise) 

46 if hasattr(cfg, "input_jitter_noise"): 

47 setattr(self.cfg, "input_jitter_noise", cfg.input_jitter_noise) 

48 if hasattr(cfg, "attention_bias"): 

49 setattr(self.cfg, "attention_bias", cfg.attention_bias) 

50 if hasattr(cfg, "lm_head_bias"): 

51 setattr(self.cfg, "lm_head_bias", cfg.lm_head_bias) 

52 if hasattr(cfg, "eos_token_id") and cfg.eos_token_id is not None: 

53 # PhiMoE chat templates terminate assistant turns with <|end|>, while 

54 # the tokenizer's primary EOS is <|endoftext|>. Stop on either by 

55 # default so generate() does not continue into a new assistant turn. 

56 setattr(self.cfg, "eos_token_id", [cfg.eos_token_id, 32007]) 

57 

58 rope_parameters = getattr(cfg, "rope_parameters", None) or {} 

59 rope_theta = rope_parameters.get("rope_theta") or getattr(cfg, "rope_theta", None) 

60 if rope_theta is not None: 

61 self.cfg.rotary_base = rope_theta 

62 

63 self.weight_processing_conversions = { 

64 **self._qkvo_weight_conversions( 

65 include_biases=bool(getattr(self.cfg, "attention_bias", False)) 

66 ), 

67 } 

68 

69 self.component_mapping = { 

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

71 "blocks": BlockBridge( 

72 name="model.layers", 

73 submodules={ 

74 "ln1": NormalizationBridge(name="input_layernorm", config=self.cfg), 

75 "ln2": NormalizationBridge(name="post_attention_layernorm", config=self.cfg), 

76 # Keep PhiMoE attention delegated to HF so native RoPE, GQA, 

77 # and cache behavior stay aligned with Transformers. 

78 "attn": AttentionBridge( 

79 name="self_attn", 

80 config=self.cfg, 

81 submodules={ 

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

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

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

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

86 }, 

87 maintain_native_attention=True, 

88 requires_attention_mask=True, 

89 ), 

90 # Native Transformers names the sparse MoE block "mlp" and 

91 # its router "router"; the archived remote code used other names. 

92 "mlp": MoEBridge( 

93 name="mlp", 

94 config=self.cfg, 

95 submodules={ 

96 "gate": MoERouterBridge(name="router"), 

97 }, 

98 ), 

99 }, 

100 ), 

101 "ln_final": NormalizationBridge(name="model.norm", config=self.cfg), 

102 "unembed": UnembeddingBridge(name="lm_head", config=self.cfg), 

103 } 

104 

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

106 """Disable remote code; base hook forces eager attention.""" 

107 # The archived remote PhiMoE code is incompatible with current 

108 # Transformers cache/generation semantics; always use the native class. 

109 model_kwargs["trust_remote_code"] = False 

110 super().prepare_loading(model_name, model_kwargs) 

111 

112 def prepare_model(self, hf_model: Any) -> None: 

113 """Also force eager on the inner model module (PhiMoE re-derives it there).""" 

114 super().prepare_model(hf_model) 

115 if hasattr(hf_model, "model") and hasattr(hf_model.model, "_attn_implementation"): 

116 hf_model.model._attn_implementation = "eager"