""" config.py -- SpikeWhale: combined config from SpikeTransformer (My Project) + NanoWhale (DeepSeek-V4). Features carried from My Project (not in NanoWhale): - DERF attention: erf(alpha*score+bias)*gamma replaces softmax - XSA (Exclusive Self-Attention): orthogonality correction removes self-echo from attn output - Engram N-gram module: hash-table N-gram lookup with DERF gate injected into embeddings - Three-tier optimizer: embed/table params trained at lower LR Features carried from NanoWhale (not in My Project): - MLA (Multi-Head Latent Attention): low-rank Q projection + direct K,V (MQA) - Partial RoPE: rotary embeddings on only qk_rope_head_dim dims of Q and K - Low-rank grouped output projection (o_lora_rank) - Hyper-Connections: hc_mult residual streams with learned routing between layers - Shared expert in MoE (always-active expert alongside routed experts) - sqrtsoftplus expert scoring (vs softmax in My Project) - Hash-based routing for first num_hash_layers layers - norm_topk_prob + routed_scaling_factor - Multi-Token Prediction (MTP): extra heads predict k steps ahead - torch.compile, FineWeb-Edu streaming, Trackio, YAML configs in train.py """ from transformers import PretrainedConfig class SpikeWhaleConfig(PretrainedConfig): model_type = "spike_whale" def __init__( self, # Standard vocab_size: int = 16512, # SpikeTokenizer: 16384 base + 128 padded special slots hidden_size: int = 1024, num_hidden_layers: int = 16, max_position_embeddings: int = 4096, rms_norm_eps: float = 1e-6, initializer_range: float = 0.02, # Byrne-500M: embeddings UNTIED. At 15M tying was mandatory (embed was 42% # of the budget); at 500M it is ~7%, so the input and output embeddings get # their own weights -- more capacity, standard for models this size. This # is the "MoE instead of tied embeddings" toggle: `use_moe` stays True and # `tie_word_embeddings` flips to False. Flip it back to True to save ~17M. tie_word_embeddings: bool = False, hidden_dropout: float = 0.0, bos_token_id: int = 0, eos_token_id: int = 1, # MLA Attention (NanoWhale) num_attention_heads: int = 16, # 16 x head_dim(64) = 1024 = hidden_size num_key_value_heads: int = 2, # 1 = MQA; >1 = GQA q_lora_rank: int = 512, # low-rank Q: hidden -> q_lora_rank -> num_heads*head_dim head_dim: int = 64, # total per-head dim = nope_head_dim + qk_rope_head_dim qk_rope_head_dim: int = 32, # RoPE applied only to these dims o_lora_rank: int = 341, # low-rank output: num_heads*head_dim -> o_lora_rank -> hidden attention_dropout: float = 0.0, rope_theta: float = 10000.0, # DERF + XSA (My Project). DERF off: Elo takes precedence over it anyway. use_derf: bool = False, use_xsa: bool = True, # Elo / Bradley-Terry tournament attention (ported from wheelerv2). # Keys carry a PERSISTENT rating that accumulates across queries and # biases future logits. Gated by a zero-init per-head gain, so it is an # exact no-op at init. Takes precedence over DERF when enabled. use_elo: bool = False, elo_k_init: float = 0.5, # Elo K-factor (rating step per "match") # Fractal RoPE (ported from wheelerv2): place RoPE inverse frequencies on # a Cantor spectrum instead of the geometric grid. gamma=0 == standard RoPE. use_fractal_rope: bool = True, fractal_rope_gamma: float = 1.0, # MoE (combined) -- Byrne-700M: 8 routed + 1 shared, top-2. # This is the capacity mechanism at 700M: total params live in the 8 # experts/layer, active stays ~2.3x smaller since only top-2 (+shared) # fire per token. use_moe must stay True for the 700M budget to hold. use_moe: bool = True, moe_intermediate_size: int = 1408, # /32 aligned; 700M n_routed_experts: int = 8, n_shared_experts: int = 1, # NanoWhale: always-active shared expert num_experts_per_tok: int = 2, norm_topk_prob: bool = True, # NanoWhale: normalize top-k routing weights scoring_func: str = "sqrtsoftplus", # NanoWhale: sqrt(softplus(x)) vs softmax routed_scaling_factor: float = 1.0, # NanoWhale: scale routed expert weights num_hash_layers: int = 2, # NanoWhale: first N layers use hash routing moe_aux_loss_coef: float = 0.01, moe_layers: list = None, # Hyper-Connections (NanoWhale) use_hyper_connections: bool = True, hc_mult: int = 2, # number of parallel residual streams hc_sinkhorn_iters: int = 20, hc_eps: float = 1e-6, # Multi-Token Prediction (NanoWhale) num_nextn_predict_layers: int = 1, # extra MTP heads (0 = disabled) # Engram N-gram module (My Project) # Engram scaled back up now that hidden=1024 gives it room (it was shrunk # to fit the 15M budget). table 4096 x 4 heads. use_engram: bool = True, engram_compress_dim: int = 64, engram_num_heads: int = 4, engram_table_size: int = 4096, engram_max_ngram: int = 3, # DERF gate init bias. MUST NOT go below about -2.5: the gate is # gamma * (erf(alpha*logits + bias) + 1) / 2 # and erf saturates far faster than the sigmoid this constant was # evidently picked for. At bias=-4, erf(-4) rounds to exactly -1.0 in # float32, so the gate is a HARD ZERO -- the Engram output, gamma's # gradient, and the gradient into the n-gram tables are all exactly 0 # forever. (sigmoid(-4)=0.018 would have been a sane "mostly off"; # erf gives 0.0.) -1.0 => gate 0.079, off-ish but alive. engram_gate_init_bias: float = -1.0, # Engram hash. False (default) = sign-bit LSH, scale-invariant. # True = the pre-v3 `h.abs().long() % table_size` hash, which collapsed # every token into bucket 0 because `h` was O(0.1) and .long() truncates. # Only set True to load a pre-v3 checkpoint bit-identically. engram_legacy_hash: bool = False, # HRM-inspired iterative refinement (EXPERIMENTAL; off by default). # Adds one small block that refines the final hidden state over N inner # steps before the output norm. This is the "iterative refinement" part # that the ARC-Prize ablation found carried most of HRM's benefit -- NOT # the full two-timescale H/L hierarchy. Honestly labeled HRM-inspired. use_hrm_refine: bool = True, hrm_refine_steps: int = 3, # inner refinement iterations hrm_refine_dim: int = 256, # bottleneck width of the refine MLP # Per-step gate init. MUST STAY NON-ZERO. The refine update is # h += tanh(gate[t]) * up(silu(down(...))) # with up zero-init, so gate and up sit on one multiplicative path: if # BOTH start at zero (as they did) then d/d(up) is proportional to # tanh(gate)=0, d/d(gate) is proportional to up(...)=0, and down/norm sit # behind up.weight=0 -- every tensor in the block gets exactly zero # gradient forever and HRM refinement silently never trains. Only ONE # endpoint of the path may be zero; `up` keeps it (that is what makes # the block an exact no-op at step 0), so this must not be 0. # 1.0 -> tanh 0.76, derivative 0.42: open but far from saturation. hrm_gate_init: float = 1.0, # --- LDT-adapted (arXiv:2605.08605); see docs/LDT-notes.md --- # Deep supervision: LDT supervises ALL 16 of its internal iterations # equally (eq. 1) and reads only the final one at inference. The paper's # own Limitations section blames its failed ARC port partly on lacking # "deep supervision that pushes a useful gradient through every # intermediate state". Applied here to hrm_refine's inner steps, whose # outputs are currently unsupervised. Requires use_hrm_refine=True. # Cost: one lm_head projection per refinement step during training. hrm_deep_supervision: bool = True, # Abstention head: the LM-side analogue of LDT's CLS conflict sigmoid. # A scalar per position trained to predict "the model is wrong here". # This is a calibrated abstention signal, NOT the paper's soundness # guarantee -- there is no verifier over text to make it sound. use_abstain_head: bool = True, abstain_loss_weight: float = 0.1, # paper's lambda_cls abstain_pos_weight: float = 8.0, # paper's w+/w- asymmetry abstain_threshold: float = 0.6, # inference threshold, raised vs 0.5 # --- Looped transformer (Byrne-700M-Looped) --- # Apply the whole layer stack `loop_count` times, reusing the SAME # weights each pass (Universal-Transformer / LDT style recurrence). # PARAM-MATCHED: looping changes compute and effective depth, not the # parameter count (+3*hidden for pass embeds). loop_count=1 is the dense # model. Default 3 here == the looped variant that beat the dense baseline # at 15M (see ../Byrne-15M-Looped/FINDINGS.md). loop_count: int = 3, # Add a learned per-pass embedding to x at the start of each pass, so the # shared weights can tell which iteration they are on. Zero-init => exact # no-op at start; without it every pass sees an identical input signature. loop_pass_embed: bool = True, # Which layers the loop covers. # "full" -- every layer, `loop_count` times (Byrne-700M-Looped # original; effective depth = L * loop_count). # "middle_split" -- Nanbeige-4.2 `enable_double_loop_split`: keep a # fixed prefix and suffix of unlooped layers and # repeat only the MIDDLE block. Rationale: the first # and last layers do input decoding / output shaping, # which is not the work that benefits from being # iterated; the middle is where the "reasoning" is. # Cheaper per unit of effective depth than "full". loop_mode: str = "full", # Width of the looped middle block; None => num_hidden_layers // 2. # Must divide num_hidden_layers. Ignored unless loop_mode="middle_split". loop_middle_layers: int = None, # --- Depth Attention (ported from Nanbeige-4.2) --- # Cross-LAYER attention: at layer L, mix the current V with the V of # earlier checkpointed layers, weighting by softmax over the depth axis # of q.k computed per (batch, kv-head, position). Gives a token direct # access to the value it built at shallower depths -- a learned, # content-addressed skip connection along the depth axis rather than the # sequence axis. Adds ~1 parameter per layer (the gate); the mix itself # is parameter-free. # Gated by a zero-init per-layer gain => exact no-op at init. use_depth_attention: bool = True, depth_attention_stride: int = 4, # checkpoint every Nth layer as a source # --- Engram mid-network fusion (ported from Nanbeige NgramLayerFusion) --- # Engram currently injects n-gram evidence ONCE, at the embedding. Nanbeige # re-injects its n-gram features again mid-stack, gated by agreement between # the n-gram key and the current hidden state. Cheap, and the mid-stack # residual is where surface-form evidence is most useful again. # List of layer indices; None/[] disables. Requires use_engram=True. engram_fusion_layers: list = None, engram_fusion_dim: int = 256, # bottleneck width of the fusion gate # --- v2 additions --- use_qk_norm: bool = True, # per-head RMSNorm on Q,K before RoPE zloss_coef: float = 1e-4, # log^2(Z) penalty on lm_head logits (0=off) mtp_loss_weight: float = 0.3, # down-weight for MTP CE loss use_value_embed: bool = False, # per-layer value-embedding residual (zero-init) # --- Memory Caching (arXiv 2602.24281) --- # A linear-attention memory branch PARALLEL to MLA that gives growing # cross-segment memory: the sequence is split into segments of length # `mc_segment_len`; each segment's linear-attention memory is cached, and # a token reads out a Gated Residual Memory over its own (causal) segment # plus all earlier cached segments, gated by softmax(). # Added into the attention sub-layer through a zero-init gate => exact # no-op at init, so it cannot destabilize the baseline. Lightweight heads # (mc_num_heads x mc_head_dim) keep the added parameter count small. use_memory_cache: bool = False, mc_segment_len: int = 256, mc_num_heads: int = 4, mc_head_dim: int = 32, mc_gate_dim: int = 64, **kwargs, ): super().__init__( bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs, ) self.vocab_size = vocab_size self.hidden_size = hidden_size self.num_hidden_layers = num_hidden_layers self.max_position_embeddings = max_position_embeddings self.rms_norm_eps = rms_norm_eps self.initializer_range = initializer_range self.hidden_dropout = hidden_dropout # Memory Caching (arXiv 2602.24281) self.use_memory_cache = use_memory_cache self.mc_segment_len = mc_segment_len self.mc_num_heads = mc_num_heads self.mc_head_dim = mc_head_dim self.mc_gate_dim = mc_gate_dim self.num_attention_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.q_lora_rank = q_lora_rank self.head_dim = head_dim self.qk_rope_head_dim = qk_rope_head_dim self.nope_head_dim = head_dim - qk_rope_head_dim self.o_lora_rank = o_lora_rank self.attention_dropout = attention_dropout self.rope_theta = rope_theta self.use_derf = use_derf self.use_xsa = use_xsa self.use_elo = use_elo self.elo_k_init = elo_k_init self.use_fractal_rope = use_fractal_rope self.fractal_rope_gamma = fractal_rope_gamma self.use_moe = use_moe self.moe_intermediate_size = moe_intermediate_size self.n_routed_experts = n_routed_experts self.n_shared_experts = n_shared_experts self.num_experts_per_tok = num_experts_per_tok self.norm_topk_prob = norm_topk_prob self.scoring_func = scoring_func self.routed_scaling_factor = routed_scaling_factor self.num_hash_layers = num_hash_layers self.moe_aux_loss_coef = moe_aux_loss_coef self.moe_layers = moe_layers if moe_layers is not None else list(range(num_hidden_layers)) self.use_hyper_connections = use_hyper_connections self.hc_mult = hc_mult self.hc_sinkhorn_iters = hc_sinkhorn_iters self.hc_eps = hc_eps self.num_nextn_predict_layers = num_nextn_predict_layers self.use_engram = use_engram self.engram_compress_dim = engram_compress_dim self.engram_num_heads = engram_num_heads self.engram_table_size = engram_table_size self.engram_max_ngram = engram_max_ngram self.engram_gate_init_bias = engram_gate_init_bias self.engram_legacy_hash = engram_legacy_hash self.engram_fusion_layers = engram_fusion_layers if engram_fusion_layers is not None else [] self.engram_fusion_dim = engram_fusion_dim self.use_depth_attention = use_depth_attention self.depth_attention_stride = depth_attention_stride self.loop_mode = loop_mode self.loop_middle_layers = loop_middle_layers self.use_hrm_refine = use_hrm_refine self.hrm_refine_steps = hrm_refine_steps self.hrm_refine_dim = hrm_refine_dim if use_hrm_refine and hrm_gate_init == 0.0: raise ValueError( "hrm_gate_init=0 with the zero-init `up` projection makes every " "parameter in HRMRefinementBlock permanently un-trainable " "(gradient exactly 0 on every step). Use a non-zero value; see " "the note in config.py." ) self.hrm_gate_init = hrm_gate_init self.hrm_deep_supervision = hrm_deep_supervision self.use_abstain_head = use_abstain_head self.abstain_loss_weight = abstain_loss_weight self.abstain_pos_weight = abstain_pos_weight self.abstain_threshold = abstain_threshold self.loop_count = loop_count self.loop_pass_embed = loop_pass_embed self.use_qk_norm = use_qk_norm self.zloss_coef = zloss_coef self.mtp_loss_weight = mtp_loss_weight self.use_value_embed = use_value_embed