Upload 11 files

README.md

@@ -1,3 +1,29 @@
 ---
 license: apache-2.0
+language:
+  - fa
+  - en
+  - multilingual
+tags:
+  - aethermind
+  - decoder-only
+  - rope
+  - gqa
+  - swiglu
+  - moe
+  - safetensors
+library_name: transformers
 ---
+
+# AetherMind-Tiny (pipeline-validation prototype)
+
+A real, trainable 5.83M-parameter decoder-only model produced by the AetherMind
+training pipeline (`training/train_proto.py`, 900 steps on CPU). Loads via
+`AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True)` with
+safetensors weights, config.json, tokenizer JSONs and a ChatML tool-aware
+chat template.
+
+Real metrics: train loss 0.49 · val loss 9.05 (memorization regime —
+prototype scale) · task checks 4/4 multilingual FAQ, 3/3 chat format.
+This is NOT a production chat model; it validates the full pipeline that
+scales to the 1B–70B presets.

chat_template.jinja

@@ -0,0 +1 @@
+{%- if tools %}{{- '<|im_start|>system\nYou are AetherMind, a helpful, honest and safe multilingual assistant.\nYou have access to the following tools (JSON schema):\n' + tools | tojson + '\nWhen you decide to call a tool, output one JSON object between <|tool_start|> and <|tool_end|>. After receiving tool results, answer the user.<|im_end|>\n' }}{%- endif %}{%- for message in messages %}{%- if message.role == 'assistant' and message.tool_calls %}{{- '<|im_start|>assistant\n' + (message.content or '') }}{%- for tc in message.tool_calls %}{{- '<|tool_start|>' + {'name': tc['function']['name'], 'arguments': tc['function']['arguments']} | tojson + '<|tool_end|>' }}{%- endfor %}{{- '<|im_end|>\n' }}{%- else %}{{- '<|im_start|>' + message['role'] + '\n' + (message['content'] or '') + '<|im_end|>\n' }}{%- endif %}{%- endfor %}{%- if add_generation_prompt %}{{- '<|im_start|>assistant\n' }}{%- endif %}

chat_template.json

@@ -0,0 +1,3 @@
+{
+  "chat_template": "{%- if tools %}{{- '<|im_start|>system\\nYou are AetherMind, a helpful, honest and safe multilingual assistant.\\nYou have access to the following tools (JSON schema):\\n' + tools | tojson + '\\nWhen you decide to call a tool, output one JSON object between <|tool_start|> and <|tool_end|>. After receiving tool results, answer the user.<|im_end|>\\n' }}{%- endif %}{%- for message in messages %}{%- if message.role == 'assistant' and message.tool_calls %}{{- '<|im_start|>assistant\\n' + (message.content or '') }}{%- for tc in message.tool_calls %}{{- '<|tool_start|>' + {'name': tc['function']['name'], 'arguments': tc['function']['arguments']} | tojson + '<|tool_end|>' }}{%- endfor %}{{- '<|im_end|>\\n' }}{%- else %}{{- '<|im_start|>' + message['role'] + '\\n' + (message['content'] or '') + '<|im_end|>\\n' }}{%- endif %}{%- endfor %}{%- if add_generation_prompt %}{{- '<|im_start|>assistant\\n' }}{%- endif %}"
+}

config.json

@@ -0,0 +1,42 @@
+{
+  "architectures": [
+    "AetherMindForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 2,
+  "dtype": "float32",
+  "eos_token_id": 3,
+  "head_dim": 32,
+  "hidden_act": "silu",
+  "hidden_size": 256,
+  "initializer_range": 0.02,
+  "intermediate_size": 768,
+  "max_position_embeddings": 1024,
+  "model_type": "aethermind",
+  "moe_aux_loss_coeff": 0.01,
+  "moe_intermediate_size": 768,
+  "moe_layers_freq": 1,
+  "num_attention_heads": 8,
+  "num_experts": 8,
+  "num_experts_per_tok": 2,
+  "num_hidden_layers": 6,
+  "num_key_value_heads": 4,
+  "pad_token_id": 0,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.17.0",
+  "use_cache": true,
+  "use_moe": false,
+  "vocab_size": 4326,
+  "auto_map": {
+    "AutoConfig": "configuration_aethermind.AetherMindConfig",
+    "AutoModel": "modeling_aethermind.AetherMindModel",
+    "AutoModelForCausalLM": "modeling_aethermind.AetherMindForCausalLM"
+  }
+}

configuration_aethermind.py

@@ -0,0 +1,167 @@
+# AetherMind — configuration_aethermind.py
+# Copyright 2026 AetherMind Project. Apache License 2.0.
+"""AetherMind model configuration.
+
+AetherMind is a fully open, trainable family of decoder-only Transformer
+language models with a modern architecture:
+
+    * RMSNorm pre-normalization
+    * Rotary Position Embeddings (RoPE, with optional linear scaling)
+    * SwiGLU feed-forward blocks
+    * Grouped-Query Attention (GQA)
+    * Scaled-dot-product / Flash attention (via torch sdpa on CUDA)
+    * Dynamic KV cache for fast generation
+    * Optional Mixture-of-Experts (MoE) feed-forward for very large configs
+    * Optional gradient checkpointing (enable via ``model.gradient_checkpointing_enable()``)
+
+The same code path scales from the CPU-trainable ``tiny`` prototype up to
+70B-class configs; every hyper-parameter lives here.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+from transformers import PretrainedConfig
+
+
+# ----------------------------------------------------------------------
+# Named size presets (scaling path: tiny -> 70B)
+# ----------------------------------------------------------------------
+SIZE_PRESETS: Dict[str, Dict[str, Any]] = {
+    # CPU-trainable prototype used to validate the full pipeline.
+    "tiny": dict(
+        hidden_size=256, intermediate_size=768, num_hidden_layers=6,
+        num_attention_heads=8, num_key_value_heads=4, head_dim=32,
+        max_position_embeddings=1024, tie_word_embeddings=True,
+    ),
+    "1B": dict(
+        hidden_size=2048, intermediate_size=5632, num_hidden_layers=24,
+        num_attention_heads=16, num_key_value_heads=4, head_dim=128,
+        max_position_embeddings=32768, vocab_size=65536,
+    ),
+    "3B": dict(
+        hidden_size=3200, intermediate_size=8704, num_hidden_layers=32,
+        num_attention_heads=25, num_key_value_heads=5, head_dim=128,
+        max_position_embeddings=32768, vocab_size=65536,
+    ),
+    "7B": dict(
+        hidden_size=4096, intermediate_size=11008, num_hidden_layers=32,
+        num_attention_heads=32, num_key_value_heads=8, head_dim=128,
+        max_position_embeddings=32768, vocab_size=65536,
+    ),
+    "14B": dict(
+        hidden_size=5120, intermediate_size=13824, num_hidden_layers=40,
+        num_attention_heads=40, num_key_value_heads=8, head_dim=128,
+        max_position_embeddings=32768, vocab_size=65536,
+    ),
+    # From here up, sparse MoE feed-forward keeps FLOPs per token bounded.
+    "32B": dict(
+        hidden_size=5120, intermediate_size=13824, num_hidden_layers=48,
+        num_attention_heads=40, num_key_value_heads=8, head_dim=128,
+        max_position_embeddings=32768, vocab_size=65536,
+        use_moe=True, num_experts=64, num_experts_per_tok=6,
+        moe_intermediate_size=1408, moe_layers_freq=1,
+    ),
+    "70B": dict(
+        hidden_size=8192, intermediate_size=28672, num_hidden_layers=80,
+        num_attention_heads=64, num_key_value_heads=8, head_dim=128,
+        max_position_embeddings=32768, vocab_size=65536,
+        use_moe=True, num_experts=128, num_experts_per_tok=8,
+        moe_intermediate_size=2048, moe_layers_freq=1,
+        rope_scaling={"type": "linear", "factor": 4.0},  # 32K -> 128K effective
+    ),
+}
+
+
+class AetherMindConfig(PretrainedConfig):
+    r"""Configuration class for :class:`~model.modeling_aethermind.AetherMindModel`."""
+
+    model_type = "aethermind"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size: int = 32768,
+        hidden_size: int = 512,
+        intermediate_size: int = 1376,
+        num_hidden_layers: int = 8,
+        num_attention_heads: int = 8,
+        num_key_value_heads: int = 4,
+        head_dim: Optional[int] = None,
+        hidden_act: str = "silu",
+        max_position_embeddings: int = 4096,
+        initializer_range: float = 0.02,
+        rms_norm_eps: float = 1e-6,
+        use_cache: bool = True,
+        pad_token_id: Optional[int] = None,
+        bos_token_id: Optional[int] = None,
+        eos_token_id: Optional[int] = None,
+        tie_word_embeddings: bool = False,
+        rope_theta: float = 10000.0,
+        rope_scaling: Optional[Dict[str, Any]] = None,
+        attention_bias: bool = False,
+        attention_dropout: float = 0.0,
+        # ---- Mixture of Experts (for large-scale configs) ----
+        use_moe: bool = False,
+        num_experts: int = 8,
+        num_experts_per_tok: int = 2,
+        moe_intermediate_size: Optional[int] = None,
+        moe_layers_freq: int = 1,
+        moe_aux_loss_coeff: float = 0.01,
+        **kwargs,
+    ) -> None:
+        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
+        if num_key_value_heads is None:
+            num_key_value_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.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        if rope_scaling is not None and rope_scaling.get("type", rope_scaling.get("rope_type")) not in (None, "linear"):
+            raise ValueError("AetherMind currently supports rope_scaling type='linear' (or None).")
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+
+        self.use_moe = use_moe
+        self.num_experts = num_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.moe_intermediate_size = moe_intermediate_size or intermediate_size
+        self.moe_layers_freq = moe_layers_freq
+        self.moe_aux_loss_coeff = moe_aux_loss_coeff
+
+        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,
+        )
+
+    @classmethod
+    def from_preset(
+        cls,
+        name: str = "tiny",
+        vocab_size: Optional[int] = None,
+        **overrides: Any,
+    ) -> "AetherMindConfig":
+        """Build a config from a named preset (tiny, 1B, 3B, 7B, 14B, 32B, 70B)."""
+        if name not in SIZE_PRESETS:
+            raise KeyError(f"Unknown preset '{name}'. Available: {sorted(SIZE_PRESETS)}")
+        params: Dict[str, Any] = dict(SIZE_PRESETS[name])
+        if vocab_size is not None:
+            params["vocab_size"] = vocab_size
+        params.update(overrides)
+        return cls(**params)
+
+
+__all__ = ["AetherMindConfig", "SIZE_PRESETS"]

generation_config.json

@@ -0,0 +1,8 @@
+{
+  "bos_token_id": 2,
+  "do_sample": false,
+  "eos_token_id": 3,
+  "max_new_tokens": 64,
+  "pad_token_id": 0,
+  "transformers_version": "5.17.0"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c08e260e255960a1a1ea2ee8b37c46c106168b73a600f2ba73b1b0038ed8a03b
+size 23323560

modeling_aethermind.py

@@ -0,0 +1,684 @@
+# AetherMind — modeling_aethermind.py
+# Copyright 2026 AetherMind Project. Apache License 2.0.
+"""PyTorch AetherMind model: a modern decoder-only Transformer.
+
+Components implemented from scratch (no third-party LLM code):
+
+    * ``AetherMindRMSNorm``          — RMSNorm with fp32 statistics
+    * ``AetherMindRotaryEmbedding``  — RoPE with optional linear scaling
+    * ``AetherMindAttention``        — Grouped-Query Attention + KV cache
+    * ``AetherMindMLP``              — SwiGLU feed-forward
+    * ``AetherMindSparseMoE``        — top-k Mixture-of-Experts + aux loss
+    * ``AetherMindDecoderLayer``     — pre-norm Transformer block
+    * ``AetherMindModel``            — embedding + blocks + final norm
+    * ``AetherMindForCausalLM``      — LM head + loss + generation
+
+Attention runs through ``torch.nn.functional.scaled_dot_product_attention``:
+on CUDA this dispatches to FlashAttention kernels when no explicit mask is
+needed (no padding), and to the memory-efficient backend otherwise.
+"""
+
+from __future__ import annotations
+
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import GenerationMixin
+from transformers.cache_utils import DynamicCache
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+from .configuration_aethermind import AetherMindConfig
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "AetherMindConfig"
+
+
+# ======================================================================
+# Building blocks
+# ======================================================================
+class AetherMindRMSNorm(nn.Module):
+    """Root-mean-square layer normalization (as in the Llama/RWKV family)."""
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+class AetherMindRotaryEmbedding(nn.Module):
+    """Rotary position embeddings with optional linear interpolation scaling."""
+
+    def __init__(
+        self,
+        head_dim: int,
+        max_position_embeddings: int = 4096,
+        base: float = 10000.0,
+        rope_scaling: Optional[dict] = None,
+    ) -> None:
+        super().__init__()
+        self.head_dim = head_dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        factor = 1.0
+        if rope_scaling is not None:
+            rtype = rope_scaling.get("rope_type", rope_scaling.get("type", "linear"))
+            if rtype == "linear":
+                factor = float(rope_scaling.get("factor", 1.0))
+        self.scaling_factor = factor
+        # inv_freq is kept as a plain attribute (NOT a registered buffer):
+        # it is deterministic from (head_dim, base, scaling_factor), so saving it
+        # is redundant; keeping it out of the state dict avoids meta-device
+        # materialization bugs and accidental dtype casts (e.g. model.to(bf16)).
+        self._inv_freq: Optional[torch.Tensor] = None
+
+    def _get_inv_freq(self, device: torch.device) -> torch.Tensor:
+        if (
+            self._inv_freq is None
+            or self._inv_freq.device != device
+            or self._inv_freq.dtype != torch.float32
+        ):
+            inv_freq = 1.0 / (
+                self.base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32, device=device) / self.head_dim)
+            )
+            self._inv_freq = inv_freq / self.scaling_factor
+        return self._inv_freq
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor, position_ids: torch.Tensor):
+        """Return cos/sin of shape ``(batch, seq_len, head_dim)`` in float32."""
+        inv_freq = self._get_inv_freq(x.device)[None, None, :]        # (1, 1, hd/2)
+        pos = position_ids[:, :, None].to(torch.float32)              # (b, s, 1)
+        freqs = pos * inv_freq                                        # (b, s, hd/2)
+        emb = torch.cat((freqs, freqs), dim=-1)                       # (b, s, hd)
+        if position_ids.max() >= self.max_position_embeddings:
+            warnings.warn(
+                "Sequence length exceeds max_position_embeddings; RoPE positions "
+                "beyond the trained window degrade quality. Consider context extension."
+            )
+        return emb.cos().to(torch.float32), emb.sin().to(torch.float32)
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim: int = 1):
+    cos = cos.unsqueeze(unsqueeze_dim)  # (b, 1, s, d)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """Expand KV heads to match the number of query heads (GQA)."""
+    if n_rep == 1:
+        return hidden_states
+    batch, num_kv_heads, slen, head_dim = hidden_states.shape
+    hidden_states = hidden_states[:, :, None, :, :].expand(
+        batch, num_kv_heads, n_rep, slen, head_dim
+    )
+    return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim)
+
+
+# ======================================================================
+# Attention
+# ======================================================================
+class AetherMindAttention(nn.Module):
+    """Grouped-Query Attention with RoPE and a dynamic KV cache."""
+
+    def __init__(self, config: AetherMindConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = config.head_dim or config.hidden_size // config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        if self.num_heads % self.num_key_value_heads != 0:
+            raise ValueError(
+                f"num_attention_heads ({self.num_heads}) must be divisible by "
+                f"num_key_value_heads ({self.num_key_value_heads})"
+            )
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.scaling = self.head_dim ** -0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+
+        op_size = self.num_heads * self.head_dim
+        self.q_proj = nn.Linear(self.hidden_size, op_size, bias=config.attention_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
+        self.o_proj = nn.Linear(op_size, self.hidden_size, bias=config.attention_bias)
+
+        self.rotary_emb = AetherMindRotaryEmbedding(
+            self.head_dim,
+            config.max_position_embeddings,
+            config.rope_theta,
+            config.rope_scaling,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_values: Optional[DynamicCache] = None,
+        use_cache: bool = False,
+        cache_position: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        bsz, q_len, _ = hidden_states.shape
+
+        query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        cos, sin = self.rotary_emb(value_states, position_ids)
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_values is not None:
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx, cache_kwargs
+            )
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        if output_attentions:
+            # Explicit path (returns attention probabilities, CPU-friendly)
+            attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scaling
+            if attention_mask is not None:
+                attn_weights = attn_weights + attention_mask
+            attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+            attn_output = torch.matmul(attn_weights, value_states)
+            attn_weights_out = attn_weights
+        else:
+            full_prefill = attention_mask is None and q_len == key_states.shape[-2] and q_len > 1
+            attn_output = F.scaled_dot_product_attention(
+                query_states,
+                key_states,
+                value_states,
+                attn_mask=None if full_prefill else attention_mask,
+                dropout_p=self.attention_dropout if self.training else 0.0,
+                is_causal=full_prefill,
+                scale=self.scaling,
+            )
+            attn_weights_out = None
+
+        attn_output = attn_output.transpose(1, 2).contiguous().reshape(bsz, q_len, -1)
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights_out
+
+
+# ======================================================================
+# Feed-forward blocks
+# ======================================================================
+ACT2FN = {"silu": F.silu, "gelu": F.gelu, "relu": F.relu}
+
+
+class AetherMindMLP(nn.Module):
+    """SwiGLU feed-forward: down( act(gate(x)) * up(x) )."""
+
+    def __init__(self, config: AetherMindConfig, intermediate_size: Optional[int] = None) -> None:
+        super().__init__()
+        inter = intermediate_size or config.intermediate_size
+        self.gate_proj = nn.Linear(config.hidden_size, inter, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, inter, bias=False)
+        self.down_proj = nn.Linear(inter, config.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class AetherMindSparseMoE(nn.Module):
+    """Top-k Mixture-of-Experts feed-forward with a load-balancing aux loss.
+
+    The router picks ``num_experts_per_tok`` experts per token; outputs are
+    the probability-weighted sum of the selected experts. A switch-transformer
+    style auxiliary loss (importance x load) encourages balanced routing and
+    is accumulated in ``AetherMindModel.moe_aux_loss`` during training.
+    """
+
+    def __init__(self, config: AetherMindConfig) -> None:
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.top_k = config.num_experts_per_tok
+        self.gate = nn.Linear(config.hidden_size, self.num_experts, bias=False)
+        self.experts = nn.ModuleList(
+            [AetherMindMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(self.num_experts)]
+        )
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        bsz, seq_len, hidden = x.shape
+        flat = x.reshape(-1, hidden)
+        router_logits = self.gate(flat)                                   # (N, E)
+        routing_probs = F.softmax(router_logits, dim=-1, dtype=torch.float32)
+        topk_weights, topk_ids = torch.topk(routing_probs, self.top_k, dim=-1)
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+
+        out = torch.zeros_like(flat)
+        for e_idx, expert in enumerate(self.experts):
+            mask = topk_ids == e_idx                                      # (N, k)
+            token_mask = mask.any(dim=-1)
+            if token_mask.any():
+                weights = (topk_weights * mask).sum(dim=-1)[token_mask].to(flat.dtype)
+                out[token_mask] += expert(flat[token_mask]) * weights.unsqueeze(-1)
+
+        # Auxiliary load-balancing loss (importance x load), switch-transformer style
+        importance = routing_probs.mean(dim=0)                            # (E,)
+        load = topk_ids.new_zeros(self.num_experts, dtype=torch.float32)
+        for k in range(self.top_k):
+            load += torch.bincount(topk_ids[:, k], minlength=self.num_experts).to(torch.float32)
+        load = load / (flat.shape[0] * self.top_k)
+        aux_loss = self.num_experts * torch.sum(importance * load)
+
+        return out.reshape(bsz, seq_len, hidden), aux_loss
+
+
+# ======================================================================
+# Decoder layer
+# ======================================================================
+class AetherMindDecoderLayer(nn.Module):
+    """Pre-norm block: x + Attn(RMSNorm(x)); x + FFN(RMSNorm(x))."""
+
+    def __init__(self, config: AetherMindConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.self_attn = AetherMindAttention(config, layer_idx=layer_idx)
+        use_moe = config.use_moe and (layer_idx % max(1, config.moe_layers_freq) == 0)
+        self.mlp = AetherMindSparseMoE(config) if use_moe else AetherMindMLP(config)
+        self.input_layernorm = AetherMindRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = AetherMindRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_values: Optional[DynamicCache] = None,
+        use_cache: bool = False,
+        cache_position: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        **kwargs,
+    ):
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+
+        attn_out, attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            output_attentions=output_attentions,
+        )
+        hidden_states = residual + attn_out
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        aux_loss = None
+        if isinstance(self.mlp, AetherMindSparseMoE):
+            hidden_states, aux_loss = self.mlp(hidden_states)
+        else:
+            hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states, attn_weights, aux_loss)
+        return outputs
+
+
+# ======================================================================
+# Base model
+# ======================================================================
+class AetherMindPreTrainedModel(PreTrainedModel):
+    config_class = AetherMindConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["AetherMindDecoderLayer"]
+    _supports_flash_attn_2 = False
+    _supports_sdpa = True
+    _supports_cache_class = True
+
+    def _init_weights(self, module: nn.Module) -> None:
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+class AetherMindModel(AetherMindPreTrainedModel):
+    """Transformer backbone without the LM head."""
+
+    def __init__(self, config: AetherMindConfig) -> None:
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [AetherMindDecoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)]
+        )
+        self.norm = AetherMindRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.gradient_checkpointing = False
+        self.moe_aux_loss: Optional[torch.Tensor] = None
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.embed_tokens = value
+
+    # ------------------------------------------------------------------
+    def _create_causal_mask(
+        self,
+        bsz: int,
+        q_len: int,
+        kv_len: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        attention_mask: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        """Additive (min-val) mask combining causality and padding."""
+        min_val = torch.finfo(dtype).min
+        q_pos = torch.arange(kv_len - q_len, kv_len, device=device).unsqueeze(-1)  # (q, 1)
+        k_pos = torch.arange(0, kv_len, device=device).unsqueeze(0)                # (1, kv)
+        causal = (k_pos > q_pos).unsqueeze(0).unsqueeze(0)                         # (1, 1, q, kv)
+        mask = torch.zeros((bsz, 1, q_len, kv_len), dtype=dtype, device=device)
+        mask = mask.masked_fill(causal, min_val)
+        if attention_mask is not None:
+            pad = (attention_mask[:, None, None, :kv_len] == 0)
+            mask = mask.masked_fill(pad, min_val)
+        return mask
+
+    # ------------------------------------------------------------------
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[DynamicCache, List[Tuple[torch.Tensor, torch.Tensor]]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else False
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else False
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else True
+        if self.gradient_checkpointing and self.training:
+            use_cache = False  # the cache is mutated in-place and breaks recomputation
+
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Pass exactly one of input_ids or inputs_embeds.")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        bsz, seq_len, _ = inputs_embeds.shape
+        device = inputs_embeds.device
+
+        # ---- cache handling (accepts DynamicCache or legacy tuple list) ----
+        if past_key_values is not None and not hasattr(past_key_values, "update"):
+            legacy = past_key_values
+            past_key_values = DynamicCache()
+            for layer_idx, (k, v) in enumerate(legacy):
+                past_key_values.update(k.to(device), v.to(device), layer_idx)
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
+        if cache_position is None:
+            cache_position = torch.arange(past_seen, past_seen + seq_len, device=device)
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        # Explicit mask is required for padding, or for chunked prefill where the
+        # fast is_causal=True kernel cannot express the offset (kv_len > q_len > 1).
+        kv_len = past_seen + seq_len
+        need_explicit_mask = attention_mask is not None or (seq_len > 1 and kv_len > seq_len)
+        causal_mask = (
+            self._create_causal_mask(bsz, seq_len, kv_len, device, inputs_embeds.dtype, attention_mask)
+            if need_explicit_mask
+            else None
+        )
+
+        hidden_states = inputs_embeds
+        all_hidden_states: Tuple = ()
+        all_self_attns: Tuple = ()
+        aux_total: Optional[torch.Tensor] = None
+
+        for decoder_layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    causal_mask,
+                    position_ids,
+                    past_key_values if use_cache else None,
+                    use_cache,
+                    cache_position,
+                    output_attentions,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=causal_mask,
+                    position_ids=position_ids,
+                    past_key_values=past_key_values if use_cache else None,
+                    use_cache=use_cache,
+                    cache_position=cache_position,
+                    output_attentions=output_attentions,
+                )
+
+            hidden_states = layer_outputs[0]
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+            if layer_outputs[2] is not None:
+                aux_total = layer_outputs[2] if aux_total is None else aux_total + layer_outputs[2]
+
+        hidden_states = self.norm(hidden_states)
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+        self.moe_aux_loss = aux_total
+
+        if not return_dict:
+            return tuple(
+                v for v in (hidden_states, past_key_values, all_hidden_states, all_self_attns) if v is not None
+            )
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+            hidden_states=all_hidden_states if output_hidden_states else None,
+            attentions=all_self_attns if output_attentions else None,
+        )
+
+
+# ======================================================================
+# Causal LM
+# ======================================================================
+class AetherMindForCausalLM(AetherMindPreTrainedModel, GenerationMixin):
+    """AetherMind with a tied/untied LM head for causal language modeling."""
+
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+    def __init__(self, config: AetherMindConfig) -> None:
+        super().__init__(config)
+        self.model = AetherMindModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.lm_head
+
+    def set_output_embeddings(self, value: nn.Linear) -> None:
+        self.lm_head = value
+
+    # ------------------------------------------------------------------
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[DynamicCache, List[Tuple[torch.Tensor, torch.Tensor]]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else False
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else False
+        return_dict = return_dict if return_dict is not None else True
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            cache_position=cache_position,
+        )
+        hidden_states = outputs.last_hidden_state
+
+        if isinstance(logits_to_keep, int):
+            slice_indices = slice(None, -logits_to_keep if logits_to_keep > 0 else None)
+        else:
+            slice_indices = logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :]).float()
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+            if self.config.use_moe and self.model.moe_aux_loss is not None:
+                loss = loss + self.config.moe_aux_loss_coeff * self.model.moe_aux_loss
+
+        if not return_dict:
+            output = (logits, outputs.past_key_values)
+            if output_hidden_states:
+                output += (outputs.hidden_states,)
+            if output_attentions:
+                output += (outputs.attentions,)
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    # ------------------------------------------------------------------
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[DynamicCache] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        use_cache: bool = True,
+        **kwargs,
+    ):
+        past_length = 0
+        if past_key_values is not None and hasattr(past_key_values, "get_seq_length"):
+            past_length = past_key_values.get_seq_length()
+
+        model_inputs = {}
+        if inputs_embeds is not None and past_length == 0:
+            model_inputs["inputs_embeds"] = inputs_embeds
+        elif cache_position is not None:
+            model_inputs["input_ids"] = input_ids[:, cache_position].contiguous()
+        else:
+            model_inputs["input_ids"] = input_ids[:, past_length:].contiguous()
+            cache_position = torch.arange(
+                past_length, past_length + model_inputs["input_ids"].shape[1], device=input_ids.device
+            )
+
+        input_length = model_inputs["inputs_embeds"].shape[1] if "inputs_embeds" in model_inputs \
+            else model_inputs["input_ids"].shape[1]
+
+        position_ids = None
+        if attention_mask is not None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            position_ids = position_ids[:, -input_length:]
+
+        model_inputs.update(
+            {
+                "past_key_values": past_key_values,
+                "use_cache": use_cache,
+                "attention_mask": attention_mask,
+                "position_ids": position_ids,
+                "cache_position": cache_position,
+            }
+        )
+        return model_inputs
+
+
+__all__ = [
+    "AetherMindConfig",
+    "AetherMindModel",
+    "AetherMindForCausalLM",
+    "AetherMindPreTrainedModel",
+    "AetherMindRMSNorm",
+    "AetherMindRotaryEmbedding",
+    "AetherMindAttention",
+    "AetherMindMLP",
+    "AetherMindSparseMoE",
+    "AetherMindDecoderLayer",
+]

special_tokens_map.json

@@ -0,0 +1,12 @@
+{
+  "bos_token": "<|im_start|>",
+  "eos_token": "<|im_end|>",
+  "pad_token": "<|pad|>",
+  "unk_token": "<|unk|>",
+  "additional_special_tokens": [
+    "<|tool_start|>",
+    "<|tool_end|>",
+    "<|reasoning_start|>",
+    "<|reasoning_end|>"
+  ]
+}

tokenizer_config.json

@@ -0,0 +1,11 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|im_start|>",
+  "eos_token": "<|im_end|>",
+  "is_local": true,
+  "local_files_only": false,
+  "model_max_length": 1024,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<|unk|>"
+}