transformer_lens.model_bridge.supported_architectures.arcee module

Arcee architecture adapter.

class transformer_lens.model_bridge.supported_architectures.arcee.ArceeArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

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

Arcee is a Llama-style dense decoder: pre-norm RMSNorm, rotary position embeddings (RoPE), grouped query attention (GQA), and no biases on any projection. The single distinguishing feature is the MLP: an ungated feed-forward block (up_proj -> ReLU^2 -> down_proj) using the squared-ReLU activation (HF hidden_act = "relu2") instead of the gated SiLU/GeLU used by Llama. The post-activation neurons are exposed via the MLP bridge’s hook_post (mlp.out.hook_in), which is useful for inspecting the sparse activation structure ReLU^2 produces.

Structurally identical to Llama except for the ungated ReLU^2 MLP; unlike Apertus it uses standard input_layernorm / post_attention_layernorm names and has no Q/K normalization.

Optional Parameters (may not exist in state_dict):

Arcee models do NOT have biases on attention or MLP projections (attention_bias = false, mlp_bias = false):

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

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

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

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

Weight processing handles these missing biases gracefully via ProcessWeights._safe_get_tensor().

__init__(cfg: Any) → None

Initialize the Arcee architecture adapter.