"""Configuration for the SHRAM transformer. All architectural parameters that vary across model scales or are meaningful research variables are expressed here. Architectural constants (no bias in linear layers, SwiGLU activation with SiLU gate) are implemented in the relevant modules and documented at the point of use — they are not config parameters because they do not vary and changing them produces a different architecture, not a different scale. RoPE configuration is owned entirely by this config. Each attention path reads its parameters directly and constructs its own RotaryEmbedding instance explicitly — no HuggingFace rope infrastructure is used. See Unit 5.A design decisions in plan.md. """ from transformers import PretrainedConfig class ShramConfig(PretrainedConfig): """Configuration class for the SHRAM decoder-only transformer. SHRAM (Sparse Hybrid Token Routed Attention Mixture) replaces every standard attention layer with a hybrid layer H(x) = h_l(x) + h_s(x), where h_l is a local sliding-window causal attention path and h_s is the MoSRAH sparse routed path. All other components follow the Llama 3 baseline. This config is the single source of truth for every architectural dimension of the model. Nothing in the architecture may use a literal number that belongs here. Two independent RoPE configurations exist — one per attention path: - h_l always uses standard RoPE with ``local_rope_theta``. - BEA always uses YaRN with ``mosrah_rope_theta``, ``training_sequence_length``, ``inference_sequence_length``, ``alpha``, and ``beta``. When ``inference_sequence_length == training_sequence_length`` the YaRN scale factor ``s = 1`` and YaRN reduces exactly to standard RoPE — this is the default state and the correct setting for experiments that do not require context extension. Registered with HuggingFace AutoClass via ``auto_map``. Instantiate from the Hub:: config = AutoConfig.from_pretrained( "your-namespace/advanced-transformers-lib", trust_remote_code=True, num_hidden_layers=12, ) model = AutoModelForCausalLM.from_config(config) Args: vocab_size: Vocabulary size. Controls the embedding table and output logits dimension. Must match the tokenizer. embedding_width: Model width ``d``. The dimension of the residual stream. mlp_width: FFN hidden dimension. num_decoder_layers: Number of transformer blocks stacked in sequence. num_sliding_window_heads: Number of heads in the local sliding-window path h_l. num_mosrah_heads: Total MoSRAH expert heads available ``L``. num_selected_heads: MoSRAH heads each token selects ``K``. head_dim: Per-head dimension, shared by both attention paths. Must be even (RoPE rotates dimensions in pairs). Paper uses 16. window_size: Sliding window size for h_l. Paper uses 128. rope_mode: RoPE position encoding mode for BEA. ``"main_sequence"`` supplies original sequence positions; ``"semantic_sequence"`` supplies local slot indices. Both are required; experimentally correct mode is undetermined (paper §4). Default ``"main_sequence"``. rms_norm_eps: Epsilon for RMSNorm layers. local_rope_theta: RoPE base frequency ``b`` for the local attention path h_l. Paper uses b=10000. mosrah_rope_theta: RoPE base frequency ``b`` for the BEA path. Paper uses b=10000. training_sequence_length: Context length ``C_train`` the model was or will be trained at. Used to compute the YaRN scale factor for BEA. inference_sequence_length: Context length ``C_target`` the model must support at inference. Optional; defaults to ``training_sequence_length`` so that ``scale=1`` and YaRN reduces to standard RoPE unless explicitly extended. alpha: YaRN ramp lower boundary α (paper §A.2). Frequency dimensions with ``r(d) < alpha`` are fully interpolated by scale s. Paper value: 1.0. beta: YaRN ramp upper boundary β (paper §A.2). Frequency dimensions with ``r(d) > beta`` are left unscaled. Paper value: 32.0. attention_dropout: Dropout probability on attention weights. Default 0.0. use_cache: Whether to return past_key_values for KV caching. output_hidden_states: Whether to return hidden states after each layer. tie_word_embeddings: Whether input embedding and LM head share weights. """ model_type = "shram" auto_map = { "AutoConfig": "configuration.ShramConfig", "AutoModelForCausalLM": "huggingface.ShramForCausalLM", } def __init__( self, vocab_size: int = 50277, embedding_width: int = 512, mlp_width: int = 1366, num_decoder_layers: int = 12, num_sliding_window_heads: int = 16, num_mosrah_heads: int = 16, num_selected_heads: int = 16, head_dim: int = 16, window_size: int = 128, rope_mode: str = "main_sequence", rms_norm_eps: float = 1e-5, local_rope_theta: float = 10000.0, mosrah_rope_theta: float = 10000.0, training_sequence_length: int = 1024, inference_sequence_length: int | None = None, alpha: float = 1.0, beta: float = 32.0, attention_dropout: float = 0.0, use_cache: bool = True, output_hidden_states: bool = False, tie_word_embeddings: bool = False, **kwargs, ): if head_dim % 2 != 0: raise ValueError( f"head_dim must be even (RoPE rotates dimensions in pairs). " f"Got head_dim={head_dim}." ) if rope_mode not in {"main_sequence", "semantic_sequence"}: raise ValueError( f"rope_mode must be 'main_sequence' or 'semantic_sequence', " f"got '{rope_mode}'." ) if training_sequence_length <= 0: raise ValueError( f"training_sequence_length must be positive, " f"got {training_sequence_length}." ) if inference_sequence_length is None: inference_sequence_length = training_sequence_length if inference_sequence_length <= 0: raise ValueError( f"inference_sequence_length must be positive, " f"got {inference_sequence_length}." ) self.vocab_size = vocab_size self.hidden_size = embedding_width self.intermediate_size = mlp_width self.num_hidden_layers = num_decoder_layers self.num_sliding_window_heads = num_sliding_window_heads self.num_mosrah_heads = num_mosrah_heads self.num_selected_heads = num_selected_heads self.head_dim = head_dim self.window_size = window_size self.rope_mode = rope_mode self.rms_norm_eps = rms_norm_eps self.local_rope_theta = local_rope_theta self.mosrah_rope_theta = mosrah_rope_theta self.training_sequence_length = training_sequence_length self.inference_sequence_length = inference_sequence_length self.alpha = alpha self.beta = beta self.attention_dropout = attention_dropout self.use_cache = use_cache super().__init__( tie_word_embeddings=tie_word_embeddings, output_hidden_states=output_hidden_states, **kwargs, ) # Promote auto_map to an instance attribute so PretrainedConfig.to_dict() # serialises it into config.json. self.auto_map = type(self).auto_map @property def scale(self) -> float: """YaRN context extension scale factor s = inference_sequence_length / training_sequence_length. When scale == 1.0, YaRN reduces exactly to standard RoPE — all frequency adjustments cancel and A_rope = 1. This is the default state. """ return self.inference_sequence_length / self.training_sequence_length