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

1"""Arcee 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 BlockBridge, 

8 EmbeddingBridge, 

9 LinearBridge, 

10 MLPBridge, 

11 PositionEmbeddingsAttentionBridge, 

12 RMSNormalizationBridge, 

13 RotaryEmbeddingBridge, 

14 UnembeddingBridge, 

15) 

16 

17 

18class ArceeArchitectureAdapter(ArchitectureAdapter): 

19 """Architecture adapter for Arcee models (ArceeForCausalLM / AFM-4.5B). 

20 

21 Arcee is a Llama-style dense decoder: pre-norm RMSNorm, rotary position 

22 embeddings (RoPE), grouped query attention (GQA), and no biases on any 

23 projection. The single distinguishing feature is the MLP: an *ungated* 

24 feed-forward block (``up_proj -> ReLU^2 -> down_proj``) using the squared-ReLU 

25 activation (HF ``hidden_act = "relu2"``) instead of the gated SiLU/GeLU used by 

26 Llama. The post-activation neurons are exposed via the MLP bridge's 

27 ``hook_post`` (``mlp.out.hook_in``), which is useful for inspecting the sparse 

28 activation structure ReLU^2 produces. 

29 

30 Structurally identical to Llama except for the ungated ReLU^2 MLP; unlike 

31 Apertus it uses standard ``input_layernorm`` / ``post_attention_layernorm`` 

32 names and has no Q/K normalization. 

33 

34 Optional Parameters (may not exist in state_dict): 

35 ------------------------------------------------- 

36 Arcee models do NOT have biases on attention or MLP projections 

37 (``attention_bias = false``, ``mlp_bias = false``): 

38 

39 - blocks.{i}.attn.b_Q / b_K / b_V / b_O - No bias on attention projections 

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

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

42 - blocks.{i}.ln1.b / ln2.b / ln_final.b - RMSNorm has no bias 

43 

44 Weight processing handles these missing biases gracefully via 

45 ProcessWeights._safe_get_tensor(). 

46 """ 

47 

48 _testing_eager = None 

49 

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

51 """Initialize the Arcee architecture adapter.""" 

52 super().__init__(cfg) 

53 

54 # Set config variables for weight processing 

55 self.cfg.normalization_type = "RMS" 

56 self.cfg.positional_embedding_type = "rotary" 

57 self.cfg.final_rms = True 

58 self.cfg.gated_mlp = False # ungated ReLU^2 MLP (up_proj -> act -> down_proj) 

59 self.cfg.attn_only = False 

60 self.cfg.uses_rms_norm = True 

61 

62 # Use eager attention so output_attentions works for hook_attn_scores / 

63 # hook_pattern; SDPA does not support output_attentions. 

64 self.cfg.attn_implementation = "eager" 

65 

66 self.weight_processing_conversions = { 

67 **self._qkvo_weight_conversions(), 

68 } 

69 

70 self.component_mapping = { 

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

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

73 "blocks": BlockBridge( 

74 name="model.layers", 

75 submodules={ 

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

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

78 "attn": PositionEmbeddingsAttentionBridge( 

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 requires_attention_mask=True, 

88 requires_position_embeddings=True, 

89 ), 

90 "mlp": MLPBridge( 

91 name="mlp", 

92 submodules={ 

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

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

95 }, 

96 ), 

97 }, 

98 ), 

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

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

101 }