transformer_lens.model_bridge.supported_architectures.stablelm module

StableLM architecture adapter.

class transformer_lens.model_bridge.supported_architectures.stablelm.StableLmArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for StableLM models.

StableLM uses a Llama-like architecture with separate Q/K/V projections and gated MLP, but differs in using standard LayerNorm (not RMSNorm) and partial rotary embeddings (25% of head dimensions by default).

Supports optional features: - Grouped Query Attention (num_key_value_heads != num_attention_heads) - QKV bias (use_qkv_bias=True on some models like stable-code-3b) - Parallel residual connections (use_parallel_residual=True) - Per-head QK LayerNorm (qk_layernorm=True)

Optional Parameters (may not exist in state_dict):

  • blocks.{i}.attn.b_Q - Only present when use_qkv_bias=True

  • blocks.{i}.attn.b_K - Only present when use_qkv_bias=True

  • blocks.{i}.attn.b_V - Only present when use_qkv_bias=True

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

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

  • blocks.{i}.mlp.b_gate - No bias on MLP gate_proj

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

__init__(cfg: Any) → None

Initialize the StableLM architecture adapter.