# coding=utf-8 """Swarm-MoE model configuration — HuggingFace compatible. This defines a real, trainable Mixture-of-Experts causal language model that keeps the identity of the project (sparse "swarm" expert routing, GQA, RoPE, RMSNorm, SwiGLU, optional BitNet b1.58 ternary weights) while being fully loadable through ``transformers`` (``AutoConfig`` / ``AutoModelForCausalLM``). """ from transformers.configuration_utils import PretrainedConfig class SwarmMoEConfig(PretrainedConfig): """Configuration for :class:`SwarmMoEForCausalLM`. Args mirror standard decoder-only LLM configs, plus Swarm-MoE specifics: a pool of ``num_experts`` per layer of which only ``num_experts_per_tok`` are activated per token (the "swarm" sparsity), and an optional BitNet b1.58 ternary-weight path enabled with ``quantization="bitnet"``. """ model_type = "swarm_moe" keys_to_ignore_at_inference = ["past_key_values"] def __init__( self, vocab_size: int = 32000, hidden_size: int = 640, intermediate_size: int = 1280, num_hidden_layers: int = 8, num_attention_heads: int = 10, num_key_value_heads: int = 2, head_dim: int | None = None, # ---- Swarm Mixture-of-Experts (fully exposed for HF) ---- num_experts: int = 8, num_experts_per_tok: int = 2, num_shared_experts: int = 1, norm_topk_prob: bool = True, router_aux_loss_coef: float = 1e-2, router_jitter_noise: float = 0.0, output_router_logits: bool = False, # ---- deeper routing + stability (added identity-at-init; retains old weights) ---- deeper_router: bool = True, router_z_loss_coef: float = 1e-3, use_v_norm: bool = True, # ---- runtime cognition/plasticity hooks (zero-param, disabled by default) ---- test_time_compute_steps: int = 0, test_time_compute_layer: int = -1, memory_injection_strength: float = 0.0, router_bias_strength: float = 0.0, plasticity_strength: float = 0.0, # ---- cross-layer weight sharing (ALBERT-style); None => no sharing ---- num_unique_layers: int | None = None, # ---- optional runtime vision sidecar (lazy, outside forward/generate) ---- vision_sidecar_enabled: bool = False, vision_sidecar_path: str | None = None, vision_sidecar_max_blocks: int = 24, vision_snapshot_policy: str = "retain_latest_only", vision_source_checkpoint: str | None = None, vision_token_map: dict | None = None, # ---- positions / attention (RoPE fully exposed for HF) ---- max_position_embeddings: int = 2048, rope_theta: float = 10000.0, rope_scaling: dict | None = None, attention_dropout: float = 0.0, attention_bias: bool = False, sliding_window: int | None = None, # ---- norm / init ---- rms_norm_eps: float = 1e-6, initializer_range: float = 0.02, hidden_act: str = "silu", # ---- quantization ("none" => bf16 dense, "bitnet" => 1.58-bit STE) ---- quantization: str = "none", # ---- misc / generation ---- tie_word_embeddings: bool = True, use_cache: bool = True, pad_token_id: int = 0, bos_token_id: int = 1, eos_token_id: int = 2, **kwargs, ): self.vocab_size = vocab_size self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads self.num_experts = num_experts self.num_experts_per_tok = num_experts_per_tok self.num_shared_experts = num_shared_experts self.norm_topk_prob = norm_topk_prob self.router_aux_loss_coef = router_aux_loss_coef self.router_jitter_noise = router_jitter_noise self.output_router_logits = output_router_logits self.deeper_router = deeper_router self.router_z_loss_coef = router_z_loss_coef self.use_v_norm = use_v_norm self.test_time_compute_steps = test_time_compute_steps self.test_time_compute_layer = test_time_compute_layer self.memory_injection_strength = memory_injection_strength self.router_bias_strength = router_bias_strength self.plasticity_strength = plasticity_strength self.num_unique_layers = num_unique_layers self.vision_sidecar_enabled = bool(vision_sidecar_enabled) self.vision_sidecar_path = vision_sidecar_path self.vision_sidecar_max_blocks = int(vision_sidecar_max_blocks) self.vision_snapshot_policy = vision_snapshot_policy self.vision_source_checkpoint = vision_source_checkpoint self.vision_token_map = vision_token_map or {} self.max_position_embeddings = max_position_embeddings self.rope_theta = rope_theta self.rope_scaling = rope_scaling self.attention_dropout = attention_dropout self.attention_bias = attention_bias self.sliding_window = sliding_window self.rms_norm_eps = rms_norm_eps self.initializer_range = initializer_range self.hidden_act = hidden_act if quantization not in ("none", "bitnet"): raise ValueError(f"quantization must be 'none' or 'bitnet', got {quantization!r}") self.quantization = quantization self.use_cache = use_cache super().__init__( pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs, ) @property def num_active_params_note(self) -> str: if self.num_experts <= 0: return "shared-expert only; no routed expert pool" frac = self.num_experts_per_tok / max(self.num_experts, 1) return f"~{frac:.0%} of expert FFN params active per token (top-{self.num_experts_per_tok} of {self.num_experts})" __all__ = ["SwarmMoEConfig"]