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

1"""SmolLM3 architecture adapter. 

2 

3SmolLM3 (the HuggingFaceTB SmolLM3 family, base and instruct) is a Llama-family 

4decoder. It pairs pre-norm RMSNorm blocks with grouped-query attention (GQA), a 

5SwiGLU gated MLP, rotary position embeddings (RoPE), tied input and output 

6embeddings, and no biases on any projection. The one feature that sets it apart 

7from a plain Llama or Qwen2 decoder is NoPE (No Positional Encoding): RoPE is 

8skipped on a periodic subset of layers. That behaviour is the only piece of this 

9adapter that is not a near-verbatim clone of qwen2.py, and it is handled by the 

10small _SmolLM3AttentionBridge subclass below. 

11""" 

12 

13from typing import Any 

14 

15import torch 

16 

17from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

18from transformer_lens.model_bridge.generalized_components import ( 

19 BlockBridge, 

20 EmbeddingBridge, 

21 LinearBridge, 

22 PositionEmbeddingsAttentionBridge, 

23 RMSNormalizationBridge, 

24 RotaryEmbeddingBridge, 

25 UnembeddingBridge, 

26) 

27 

28 

29class _SmolLM3AttentionBridge(PositionEmbeddingsAttentionBridge): 

30 """Attention bridge that honours SmolLM3's per-layer NoPE setting. 

31 

32 SmolLM3 disables rotary position embeddings on a periodic subset of layers 

33 (every no_rope_layer_interval-th layer, default every 4th, controlled by 

34 config.no_rope_layers). The wrapped HF SmolLM3Attention module records this 

35 choice as an integer flag use_rope: 1 means apply RoPE, 0 means this is a 

36 NoPE layer. HF honours the flag inside its own forward by only calling 

37 apply_rotary_pos_emb when use_rope is truthy. 

38 

39 The base PositionEmbeddingsAttentionBridge reimplements attention so that all 

40 hook points fire at the right stage, and it applies RoPE whenever a 

41 position_embeddings tuple is passed. It never consults use_rope. On a NoPE 

42 layer that would rotate Q and K while native HF does not, diverging from the 

43 reference and failing logit-equivalence checks on roughly a quarter of the 

44 layers. 

45 

46 To match HF exactly we suppress position_embeddings on NoPE layers before 

47 delegating to the base forward. The base forward only rotates when 

48 position_embeddings is not None, so passing None skips the rotation while 

49 every non-rotary hook (hook_q, hook_k, hook_v, hook_attn_scores, 

50 hook_pattern, hook_z) still fires identically. RoPE layers (use_rope == 1) 

51 are left untouched and behave exactly like the qwen2.py attention bridge. 

52 """ 

53 

54 # Nulls position_embeddings on NoPE layers by design. 

55 rope_optional = True 

56 

57 def forward(self, *args: Any, **kwargs: Any) -> Any: 

58 """Drop position_embeddings on NoPE layers, then run the base forward.""" 

59 hf_attn = self.original_component 

60 # use_rope is 1 on RoPE layers and 0 on NoPE layers. Default to RoPE-on 

61 # when the attribute is somehow absent so standard layers never break. 

62 if hf_attn is not None and not getattr(hf_attn, "use_rope", 1): 

63 kwargs["position_embeddings"] = None 

64 # SmolLM3DecoderLayer (inherited from LlamaDecoderLayer) passes 

65 # position_embeddings as a keyword, so the line above is what fires 

66 # in practice. The positional branch below is defensive: if a caller 

67 # ever passes (hidden_states, position_embeddings, ...) positionally, 

68 # the second slot holds the (cos, sin) tuple, not a tensor, so we 

69 # null it out there too. 

70 if len(args) >= 2 and not isinstance(args[1], torch.Tensor): 

71 args = (args[0], None) + args[2:] 

72 return super().forward(*args, **kwargs) 

73 

74 

75class SmolLM3ArchitectureAdapter(ArchitectureAdapter): 

76 """Architecture adapter for SmolLM3 models. 

77 

78 SmolLM3 is a pre-norm decoder with RMSNorm, grouped-query attention (GQA), 

79 a SwiGLU gated MLP, rotary position embeddings (RoPE), tied input and output 

80 embeddings, and no biases on any projection. The block shape matches Llama 

81 and Qwen2 exactly, so the component mapping and weight conversions mirror 

82 qwen2.py. 

83 

84 NoPE (No Positional Encoding): SmolLM3 disables RoPE on every 

85 no_rope_layer_interval-th layer (default every 4th) via config.no_rope_layers. 

86 That per-layer toggle lives inside HF's SmolLM3Attention.forward, but the 

87 bridge reimplements attention and would otherwise rotate Q and K on those 

88 layers. The _SmolLM3AttentionBridge subclass handles it by suppressing 

89 position embeddings on NoPE layers, so the reimplemented attention matches HF. 

90 

91 No Q/K normalization: unlike Qwen3, SmolLM3 has no per-head Q or K RMSNorm, 

92 so the attention block uses the plain q/k/v/o submodules. 

93 

94 Optional Parameters (may not exist in state_dict): 

95 ------------------------------------------------- 

96 SmolLM3 models do NOT have biases on any linear layers: 

97 

98 - blocks.{i}.attn.b_Q - No bias on query projection 

99 - blocks.{i}.attn.b_K - No bias on key projection 

100 - blocks.{i}.attn.b_V - No bias on value projection 

101 - blocks.{i}.attn.b_O - No bias on output projection 

102 - blocks.{i}.mlp.b_in - No bias on MLP input (up_proj) 

103 - blocks.{i}.mlp.b_gate - No bias on MLP gate projection 

104 - blocks.{i}.mlp.b_out - No bias on MLP output (down_proj) 

105 - blocks.{i}.ln1.b - RMSNorm has no bias 

106 - blocks.{i}.ln2.b - RMSNorm has no bias 

107 - ln_final.b - RMSNorm has no bias 

108 

109 Weight processing must handle these missing biases gracefully using 

110 ProcessWeights._safe_get_tensor() or by checking for None values. 

111 """ 

112 

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

114 """Initialize the SmolLM3 architecture adapter.""" 

115 super().__init__(cfg) 

116 

117 self._set_rms_rotary_defaults() 

118 

119 self.cfg.default_prepend_bos = False 

120 # The bridge reimplements attention and reads output_attentions, so the 

121 # HF model must run in eager mode for the scores and pattern hooks to 

122 # match the reference. Set it on cfg so weight processing and 

123 # setup_component_testing agree without relying on boot()'s default. 

124 self.cfg.attn_implementation = "eager" 

125 

126 # Standard separate q_proj/k_proj/v_proj/o_proj layout, GQA-aware. No 

127 # biases anywhere (attention_bias=False, mlp_bias=False), so no bias 

128 # conversions are needed. 

129 self.weight_processing_conversions = { 

130 **self._qkvo_weight_conversions(), 

131 } 

132 

133 self.component_mapping = { 

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

135 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg), 

136 "blocks": BlockBridge( 

137 name="model.layers", 

138 config=self.cfg, 

139 submodules={ 

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

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

142 "attn": _SmolLM3AttentionBridge( 

143 name="self_attn", 

144 config=self.cfg, 

145 submodules={ 

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

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

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

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

150 }, 

151 requires_attention_mask=True, 

152 requires_position_embeddings=True, 

153 ), 

154 "mlp": self._gated_mlp(), 

155 }, 

156 ), 

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

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

159 }