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

1"""Switch Transformers (``SwitchTransformersForConditionalGeneration``) adapter: 

2T5 encoder-decoder with top-1-routed sparse-MoE feed-forwards. Blocks delegate 

3wholesale (v5's plain tensor protocol), FF is a delegated MoEBridge with an optional router.""" 

4 

5from typing import Any 

6 

7from transformer_lens.model_bridge.generalized_components import ( 

8 AttentionBridge, 

9 BlockBridge, 

10 LinearBridge, 

11 MoEBridge, 

12 RMSNormalizationBridge, 

13) 

14from transformer_lens.model_bridge.generalized_components.base import ( 

15 GeneralizedComponent, 

16) 

17from transformer_lens.model_bridge.supported_architectures.t5 import ( 

18 T5ArchitectureAdapter, 

19) 

20 

21 

22class _SwitchBlockBridge(BlockBridge): 

23 """v5 Switch blocks take and return bare tensors; the stack's minimal 

24 layer call would otherwise trip the tuple-normalizing heuristic.""" 

25 

26 @staticmethod 

27 def _is_standalone_hidden_state_call(args: tuple, kwargs: dict) -> bool: 

28 return False 

29 

30 

31class SwitchTransformersArchitectureAdapter(T5ArchitectureAdapter): 

32 """Architecture adapter for SwitchTransformersForConditionalGeneration models.""" 

33 

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

35 """Initialize the Switch Transformers architecture adapter.""" 

36 super().__init__(cfg) 

37 

38 def attention(name: str, *, cross: bool = False) -> AttentionBridge: 

39 return AttentionBridge( 

40 name=name, 

41 config=self.cfg, 

42 submodules={ 

43 "q": LinearBridge(name="q"), 

44 "k": LinearBridge(name="k"), 

45 "v": LinearBridge(name="v"), 

46 "o": LinearBridge(name="o"), 

47 }, 

48 is_cross_attention=cross, 

49 maintain_native_attention=True, 

50 ) 

51 

52 self.components["encoder_blocks"] = _SwitchBlockBridge( 

53 name="encoder.block", 

54 config=self.cfg, 

55 submodules={ 

56 "ln1": RMSNormalizationBridge(name="layer.0.layer_norm", config=self.cfg), 

57 "attn": attention("layer.0.SelfAttention"), 

58 "ln2": RMSNormalizationBridge(name="layer.1.layer_norm", config=self.cfg), 

59 "mlp": self._build_ff_bridge("layer.1"), 

60 }, 

61 ) 

62 self.components["decoder_blocks"] = _SwitchBlockBridge( 

63 name="decoder.block", 

64 config=self.cfg, 

65 hook_alias_overrides={ 

66 "hook_attn_in": "self_attn.hook_attn_in", 

67 "hook_attn_out": "self_attn.hook_out", 

68 "hook_q_input": "self_attn.hook_q_input", 

69 "hook_k_input": "self_attn.hook_k_input", 

70 "hook_v_input": "self_attn.hook_v_input", 

71 }, 

72 submodules={ 

73 "ln1": RMSNormalizationBridge(name="layer.0.layer_norm", config=self.cfg), 

74 "self_attn": attention("layer.0.SelfAttention"), 

75 "ln2": RMSNormalizationBridge(name="layer.1.layer_norm", config=self.cfg), 

76 "cross_attn": attention("layer.1.EncDecAttention", cross=True), 

77 "ln3": RMSNormalizationBridge(name="layer.2.layer_norm", config=self.cfg), 

78 "mlp": self._build_ff_bridge("layer.2"), 

79 }, 

80 ) 

81 

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

83 """The google/switch-base-* repos ship pytorch_model.bin only; skip 

84 v5's Hub-side safetensors auto-conversion (it needs a conversion PR).""" 

85 model_kwargs.setdefault("use_safetensors", False) 

86 super().prepare_loading(model_name, model_kwargs) 

87 

88 def _build_ff_bridge(self, layer_prefix: str) -> MoEBridge: 

89 """Dense or sparse per layer parity; router only exists on sparse.""" 

90 return MoEBridge( 

91 name=f"{layer_prefix}.mlp", 

92 config=self.cfg, 

93 submodules={ 

94 "gate": GeneralizedComponent(name="router", optional=True), 

95 # Dense-layer projections (absent on sparse layers); in/out per 

96 # T5-family convention (nonstandard names dodge the component 

97 # prober's down-projection skip and get crash-probed). 

98 "in": LinearBridge(name="wi", optional=True), 

99 "out": LinearBridge(name="wo", optional=True), 

100 }, 

101 )