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

1"""Gemma 3n text-only architecture adapter. 

2 

3Bridges the text path of the full tri-modal ``Gemma3nForConditionalGeneration`` 

4(``model.language_model`` + ``lm_head``); the vision/audio towers stay referenced but 

5unbridged (see the vision+audio follow-up). The decoder layers run on a stacked AltUp 

64-stream residual, so blocks use ``AltUpBlockBridge`` rather than ``BlockBridge``. All 

7math is deferred to HF; submodules are decomposed only for hooks (parity-safe delegation). 

8""" 

9 

10from typing import Any 

11 

12from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

13from transformer_lens.model_bridge.generalized_components import ( 

14 AltUpBlockBridge, 

15 AttentionBridge, 

16 EmbeddingBridge, 

17 GatedMLPBridge, 

18 LinearBridge, 

19 RotaryEmbeddingBridge, 

20 UnembeddingBridge, 

21) 

22from transformer_lens.model_bridge.generalized_components.base import ( 

23 GeneralizedComponent, 

24) 

25 

26 

27class _SparsityPreservingGatedMLPBridge(GatedMLPBridge): 

28 """GatedMLPBridge that always delegates, even under processed weights. 

29 

30 Gemma3nTextMLP applies `_gaussian_topk` activation sparsity (0.95 on the 

31 first layers of E2B/E4B) before the activation; the functional 

32 processed-weights branch computes plain act(gate)*up and inflates those 

33 layers' output ~10x. Delegation stays numerically correct in compat mode 

34 because set_processed_weights has already written the processed tensors 

35 into the wrapped Linears — and the gate/in/out hooks still fire from 

36 inside HF's forward, which is all the bridge swap wanted. 

37 """ 

38 

39 def set_processed_weights(self, weights, verbose: bool = False) -> None: 

40 super().set_processed_weights(weights, verbose=verbose) 

41 self._use_processed_weights = False 

42 

43 

44class Gemma3nArchitectureAdapter(ArchitectureAdapter): 

45 """Text-only adapter for Gemma 3n (`Gemma3nForConditionalGeneration`).""" 

46 

47 _testing_lm_attr = "model.language_model" 

48 _testing_wire_rotary = False 

49 

50 # The full model includes a timm-based vision tower (TimmWrapperModel), so timm is needed 

51 # even for text-only use (the towers stay referenced). 

52 required_libraries: list[str] = ["timm"] 

53 required_libraries_group: str = "multimodal" 

54 

55 # Phase 3 (processed/compatibility mode) folds LN into a single residual stream, which 

56 # AltUp's 4-stream residual can't represent. Phases 1 (HF parity), 2 (hooks), and 4 (text 

57 # quality) do apply and pass. 

58 applicable_phases: list[int] = [1, 2, 4] 

59 

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

61 super().__init__(cfg) 

62 

63 self.cfg.is_multimodal = False 

64 self.cfg.gated_mlp = True 

65 self.cfg.uses_rms_norm = True 

66 self.cfg.normalization_type = "RMS" 

67 self.cfg.rmsnorm_uses_offset = True # Gemma RMSNorm uses (1.0 + weight) 

68 self.cfg.positional_embedding_type = "rotary" 

69 self.cfg.attn_implementation = "eager" 

70 # AltUp + per-layer-embedding residual topology isn't fold-safe. 

71 self.supports_fold_ln = False 

72 self.weight_processing_conversions: dict = {} 

73 

74 self.component_mapping = { 

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

76 "rotary_emb": RotaryEmbeddingBridge(name="model.language_model.rotary_emb"), 

77 "blocks": AltUpBlockBridge( 

78 name="model.language_model.layers", 

79 config=self.cfg, 

80 submodules={ 

81 "input_layernorm": GeneralizedComponent(name="input_layernorm"), 

82 "post_attention_layernorm": GeneralizedComponent( 

83 name="post_attention_layernorm" 

84 ), 

85 "pre_feedforward_layernorm": GeneralizedComponent( 

86 name="pre_feedforward_layernorm" 

87 ), 

88 "post_feedforward_layernorm": GeneralizedComponent( 

89 name="post_feedforward_layernorm" 

90 ), 

91 "post_per_layer_input_norm": GeneralizedComponent( 

92 name="post_per_layer_input_norm" 

93 ), 

94 "altup": GeneralizedComponent(name="altup"), 

95 "laurel": GeneralizedComponent(name="laurel"), 

96 "per_layer_input_gate": GeneralizedComponent(name="per_layer_input_gate"), 

97 "per_layer_projection": GeneralizedComponent(name="per_layer_projection"), 

98 # AttentionBridge for the per-head hook surface a bare 

99 # component lacks. Delegated semantics: pattern AND scores 

100 # both fire HF's post-softmax weights (no pre-softmax tensor 

101 # exists, unlike gemma1/2/3); writes to them are inert. 

102 "self_attn": AttentionBridge( 

103 name="self_attn", 

104 config=self.cfg, 

105 maintain_native_attention=True, 

106 submodules={ 

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

108 # The last num_kv_shared_layers layers reuse earlier KV and 

109 # drop their own k/v projections and norms. 

110 "k": LinearBridge(name="k_proj", optional=True), 

111 "v": LinearBridge(name="v_proj", optional=True), 

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

113 "q_norm": GeneralizedComponent(name="q_norm"), 

114 "k_norm": GeneralizedComponent(name="k_norm", optional=True), 

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

116 }, 

117 ), 

118 "mlp": _SparsityPreservingGatedMLPBridge( 

119 name="mlp", 

120 config=self.cfg, 

121 submodules={ 

122 "gate": LinearBridge(name="gate_proj"), 

123 "in": LinearBridge(name="up_proj"), 

124 "out": LinearBridge(name="down_proj"), 

125 }, 

126 ), 

127 }, 

128 ), 

129 "ln_final": GeneralizedComponent(name="model.language_model.norm"), 

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

131 }