DynamicMind-MoE: 30.2M total / 8.9M active sparse MoE, upcycled from DynamicMind-Mini

README.md

@@ -0,0 +1,183 @@
+---
+license: apache-2.0
+datasets:
+- HuggingFaceFW/fineweb-edu
+- HuggingFaceTB/finemath
+- HuggingFaceTB/smollm-corpus
+language:
+- en
+pipeline_tag: text-generation
+library_name: transformers
+base_model: DedeProGames/DynamicMind-Mini
+tags:
+- causal-lm
+- language-model
+- base-model
+- mixture-of-experts
+- sparse-moe
+- dynamicmind
+- digit-tokenizer
+- pytorch
+- custom-code
+- trust-remote-code
+---
+
+![Banner](https://cdn-uploads.huggingface.co/production/uploads/685ea8ff7b4139b6845ce395/0QdBKT5iKdrEy3iwLosAH.png)
+
+# DynamicMind-MoE
+
+DynamicMind-MoE is a sparse mixture-of-experts decoder-only model trained on [FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu), [SmolLM-Corpus](https://huggingface.co/datasets/HuggingFaceTB/smollm-corpus) and [FineMath](https://huggingface.co/datasets/HuggingFaceTB/finemath)
+
+The model has **30.2M total parameters but activates only 8.9M per token**, and was initialised by sparse-upcycling [DynamicMind-Mini](https://huggingface.co/DedeProGames/DynamicMind-Mini). It keeps the same custom 8k-token byte-level BPE tokenizer with digit-aware tokenization, so inference cost per token is unchanged from the dense model while total capacity is 3.4x larger.
+
+## Model Details
+
+| Field | Value |
+|---|---:|
+| Total parameters | 30,150,912 |
+| Active parameters per token | 8,917,248 |
+| Architecture | Sparse MoE Llama-style decoder |
+| Layers | 9 |
+| Hidden size | 256 |
+| Routed experts per layer | 14 |
+| Shared experts per layer | 1 |
+| Experts activated per token | 2 routed + 1 shared |
+| Expert intermediate size | 256 |
+| Attention heads | 8 |
+| KV heads | 2 |
+| Vocabulary size | 8,192 |
+| Context length | 1,024 |
+| Embeddings | Tied input/output embeddings |
+| Weight format | safetensors |
+
+## Mixture of Experts
+
+Each block's dense MLP is replaced by one always-on shared expert plus 14 fine-grained routed experts, of which the router selects the top 2 per token. Shared + top-2 at intermediate size 256 reproduces the dense model's exact active parameter count.
+
+Load balancing uses an auxiliary-loss-free bias (DeepSeek-V3 style): a per-expert bias steers *selection* toward idle experts while the combining weights come from the unbiased softmax, so balancing costs no gradient interference.
+
+Routing is **per token, per layer** — a single sequence touches many different experts, and the 9 routers are independent of each other.
+
+## Tokenizer
+
+DynamicMind-MoE uses the same digit-aware 8k tokenizer as [DynamicMind-Mini](https://huggingface.co/DedeProGames/DynamicMind-Mini).
+
+Digits are kept as separate tokens so numbers do not collapse into large number tokens during tokenization.
+
+Digit IDs:
+
+| Token | ID |
+|---|---:|
+| `1` | 9 |
+| `2` | 10 |
+| `3` | 11 |
+| `4` | 12 |
+| `5` | 13 |
+| `6` | 14 |
+| `7` | 15 |
+| `8` | 16 |
+| `9` | 17 |
+| `0` | 18 |
+
+## Training
+
+| Field | Value |
+|---|---:|
+| Initialisation | Sparse upcycling from DynamicMind-Mini |
+| Tokens seen | 10,008,133,632 |
+| Optimizer steps | 8,484 |
+| Sequence length | 1,024 |
+| Tokens per optimizer step | 1,179,648 |
+| Peak learning rate | 1e-4 |
+| Min learning rate | 1e-5 |
+| LR schedule | Cosine decay after 200 warmup steps |
+| Optimizer | AdamW (betas 0.9, 0.95) |
+| Weight decay | 0.1 |
+| Gradient clipping | 1.0 |
+| Seed | 1337 |
+| Hardware | 1x RTX 3060 12GB, 39.5h |
+
+Data mixture: FineWeb-Edu `sample-10BT` 50.0%, Cosmopedia-v2 27.8%, FineMath-4plus 22.2%.
+
+Embeddings, attention and norms transferred 1:1 from the dense model. The dense MLP (intermediate size 768) splits exactly into three 256-wide slices — because `down_proj` sums over the intermediate axis — and each expert was seeded from one slice plus small noise to break router symmetry.
+
+## Benchmarks
+
+Self-reported results from the official [BananaMind Base Bench 1.1](https://huggingface.co/datasets/BananaMind/BananaMind-Base-Bench-1.1) script, all measured with the same runner, dtype (bfloat16) and GPU.
+
+| Model | Total params | Active params | Overall Elo |
+|---|---:|---:|---:|
+| BananaMind-2-Medium | 55.9M | 55.9M | 1,037 |
+| GPT-2 | 124M | 124M | 990 |
+| BananaMind-2-Nano | 12.1M | 12.1M | 915 |
+| **DynamicMind-MoE** | **30.2M** | **8.9M** | **912** |
+| Pythia-14m-deduped | 14M | 14M | 909 |
+| BananaMind-2-MoE | 26.1M | — | 903 |
+| DynamicMind-Mini | 8.9M | 8.9M | 868 |
+
+**Detailed DynamicMind-MoE result**
+
+| Category | Passed | Elo |
+|---|---:|---:|
+| **Overall** | **137 / 350** | **912** |
+| Language completion | 38 / 50 | 1,115 |
+| Logical reasoning | 17 / 50 | 970 |
+| World knowledge | 25 / 50 | 951 |
+| Commonsense | 23 / 50 | 923 |
+| Context tracking | 15 / 50 | 843 |
+| Quantitative | 11 / 50 | 822 |
+| Code completion | 8 / 50 | 793 |
+
+Against the dense DynamicMind-Mini the MoE gains +44 Elo overall at identical inference cost, improving in six of seven categories. Quantitative is the exception (822 vs 837), despite FineMath being 22% of the corpus.
+
+Scores are self-evaluated and may vary with the benchmark revision, Transformers version, dtype, hardware, and generation settings. Cross-tokenizer comparisons (e.g. against GPT-2's 50k vocabulary) carry a residual tokenizer effect that mean-per-token log-probability does not fully remove.
+
+## Usage
+
+This model uses custom architecture code, so load it with `trust_remote_code=True`.
+
+Install dependencies:
+
+```bash
+pip install -U transformers safetensors torch
+```
+
+Run inference:
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+model_id = "DedeProGames/DynamicMind-MoE"
+
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16,
+).cuda().eval()
+
+prompt = "The meaning of life is "
+input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
+
+with torch.no_grad():
+    output = model.generate(
+        input_ids=input_ids,
+        max_new_tokens=64,
+        do_sample=False,
+        repetition_penalty=1.1,
+        pad_token_id=tokenizer.eos_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+    )
+
+print(tokenizer.decode(output[0], skip_special_tokens=True))
+```
+
+## Limitations
+
+This is a base model, not instruction-tuned — it continues text rather than following instructions. At 8.9M active parameters it reproduces register and structure well (encyclopedic text reads encyclopedic, code keeps valid indentation) but is frequently wrong on facts and arithmetic. Keep a finite generation limit and do not use it for high-stakes decisions.
+
+## License
+
+Apache 2.0

config.json

@@ -0,0 +1,51 @@
+{
+  "transformers_version": "5.5.3",
+  "architectures": [
+    "DynamicMindMoEForCausalLM"
+  ],
+  "output_hidden_states": false,
+  "return_dict": true,
+  "dtype": "float32",
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "problem_type": null,
+  "_name_or_path": "",
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "pad_token_id": 1,
+  "tie_word_embeddings": true,
+  "vocab_size": 8192,
+  "hidden_size": 256,
+  "intermediate_size": 768,
+  "moe_intermediate_size": 256,
+  "num_hidden_layers": 9,
+  "num_attention_heads": 8,
+  "num_key_value_heads": 2,
+  "num_routed_experts": 14,
+  "num_shared_experts": 1,
+  "num_experts_per_token": 2,
+  "first_k_dense_layers": 0,
+  "norm_topk_prob": true,
+  "router_aux_loss_coef": 0.01,
+  "router_z_loss_coef": 0.001,
+  "router_bias_update_rate": 0.001,
+  "use_aux_loss_free_balancing": true,
+  "max_position_embeddings": 1024,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 10000.0,
+  "attention_dropout": 0.0,
+  "model_type": "dynamicmind_moe",
+  "output_attentions": false,
+  "auto_map": {
+    "AutoConfig": "configuration_dynamicmind_moe.DynamicMindMoEConfig",
+    "AutoModelForCausalLM": "modeling_dynamicmind_moe.DynamicMindMoEForCausalLM"
+  }
+}

configuration_dynamicmind_moe.py

@@ -0,0 +1,74 @@
+from transformers.configuration_utils import PretrainedConfig
+
+
+class DynamicMindMoEConfig(PretrainedConfig):
+    """DynamicMind-MoE: sparse mixture-of-experts variant of DynamicMind-Mini.
+
+    The dense MLP (intermediate 768) is replaced by one always-on shared expert
+    plus `num_routed_experts` fine-grained experts (intermediate 256), of which
+    `num_experts_per_token` are selected. Shared + top-2 reproduces the dense
+    layer's exact active parameter count, so inference cost per token is
+    unchanged while total capacity grows ~3.4x.
+    """
+
+    model_type = "dynamicmind_moe"
+
+    def __init__(
+        self,
+        vocab_size=8192,
+        hidden_size=256,
+        intermediate_size=768,        # kept for dense layers / upcycling source
+        moe_intermediate_size=256,    # per-expert width (768 / 3)
+        num_hidden_layers=9,
+        num_attention_heads=8,
+        num_key_value_heads=2,
+        num_routed_experts=14,
+        num_shared_experts=1,
+        num_experts_per_token=2,
+        first_k_dense_layers=0,       # keep the first K blocks dense if desired
+        norm_topk_prob=True,
+        router_aux_loss_coef=0.01,
+        router_z_loss_coef=0.001,
+        router_bias_update_rate=0.001,  # aux-loss-free balancing (DeepSeek-V3)
+        use_aux_loss_free_balancing=True,
+        max_position_embeddings=1024,
+        rms_norm_eps=1e-5,
+        rope_theta=10000.0,
+        attention_dropout=0.0,
+        tie_word_embeddings=True,
+        bos_token_id=0,
+        eos_token_id=0,
+        pad_token_id=1,
+        **kwargs,
+    ):
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.moe_intermediate_size = moe_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.num_routed_experts = num_routed_experts
+        self.num_shared_experts = num_shared_experts
+        self.num_experts_per_token = num_experts_per_token
+        self.first_k_dense_layers = first_k_dense_layers
+        self.norm_topk_prob = norm_topk_prob
+
+        self.router_aux_loss_coef = router_aux_loss_coef
+        self.router_z_loss_coef = router_z_loss_coef
+        self.router_bias_update_rate = router_bias_update_rate
+        self.use_aux_loss_free_balancing = use_aux_loss_free_balancing
+
+        self.max_position_embeddings = max_position_embeddings
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = rope_theta
+        self.attention_dropout = attention_dropout

model.safetensors

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

modeling_dynamicmind_moe.py

@@ -0,0 +1,339 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.modeling_utils import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast
+from transformers.cache_utils import Cache, DynamicCache
+
+from .configuration_dynamicmind_moe import DynamicMindMoEConfig
+
+
+class DynamicMindRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, x):
+        dtype = x.dtype
+        x = x.float()
+        x = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        return (self.weight * x).to(dtype)
+
+
+class DynamicMindRotaryEmbedding(nn.Module):
+    """RoPE with a cached inv_freq.
+
+    The dense model rebuilt inv_freq on every forward of every layer; caching it
+    removes 9 redundant allocations per step.
+
+    inv_freq is a constant derived from config, so persistent=False looks
+    correct — but from_pretrained materialises tensors straight from the
+    checkpoint onto meta-device modules, never running __init__'s value nor
+    _load_from_state_dict for it. A non-persistent buffer therefore survives
+    loading as uninitialised `torch.empty` garbage, silently scrambling RoPE:
+    measured 833 vs 908 Elo on identical weights. Persisting the 16 floats is
+    the only variant that loads correctly through every path.
+    """
+
+    def __init__(self, head_dim, rope_theta, max_position_embeddings):
+        super().__init__()
+        self.head_dim = head_dim
+        self.rope_theta = rope_theta
+        self.register_buffer("inv_freq", self._compute(), persistent=True)
+        self.max_seq_len_cached = 0
+
+    def _compute(self):
+        return 1.0 / (self.rope_theta ** (
+            torch.arange(0, self.head_dim, 2).float() / self.head_dim))
+
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+        with torch.no_grad():
+            self.inv_freq.copy_(self._compute().to(self.inv_freq.device))
+
+    def forward(self, x, position_ids):
+        freqs = position_ids[:, :, None].float() * self.inv_freq[None, None, :]
+        return freqs.cos().to(x.dtype), freqs.sin().to(x.dtype)
+
+
+def apply_rope(q, k, cos, sin):
+    cos = cos[:, None, :, :]
+    sin = sin[:, None, :, :]
+
+    def rotate(x):
+        even, odd = x[..., 0::2], x[..., 1::2]
+        return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
+
+    return rotate(q), rotate(k)
+
+
+class DynamicMindAttention(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.attention_dropout = config.attention_dropout
+
+        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+
+    def forward(self, x, cos, sin, attention_mask=None, past_key_values=None, cache_position=None):
+        bsz, q_len, _ = x.shape
+
+        q = self.q_proj(x).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        q, k = apply_rope(q, k, cos, sin)
+
+        if past_key_values is not None:
+            k, v = past_key_values.update(k, v, self.layer_idx, {"cache_position": cache_position})
+
+        if self.num_kv_heads != self.num_heads:
+            repeats = self.num_heads // self.num_kv_heads
+            k = k.repeat_interleave(repeats, dim=1)
+            v = v.repeat_interleave(repeats, dim=1)
+
+        is_causal = attention_mask is None and q_len > 1
+        y = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=attention_mask,
+            dropout_p=self.attention_dropout if self.training else 0.0,
+            is_causal=is_causal,
+        )
+        y = y.transpose(1, 2).contiguous().view(bsz, q_len, -1)
+        return self.o_proj(y)
+
+
+class DynamicMindMLP(nn.Module):
+    def __init__(self, hidden_size, intermediate_size):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class DynamicMindMoE(nn.Module):
+    """Shared expert + top-k routed fine-grained experts.
+
+    The shared expert runs on every token and absorbs knowledge common to all
+    inputs, so the routed experts are free to specialise instead of each
+    re-learning the same basics.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.num_routed = config.num_routed_experts
+        self.top_k = config.num_experts_per_token
+        self.norm_topk_prob = config.norm_topk_prob
+        self.aux_free = config.use_aux_loss_free_balancing
+
+        self.experts = nn.ModuleList([
+            DynamicMindMLP(config.hidden_size, config.moe_intermediate_size)
+            for _ in range(self.num_routed)
+        ])
+        self.shared_experts = nn.ModuleList([
+            DynamicMindMLP(config.hidden_size, config.moe_intermediate_size)
+            for _ in range(config.num_shared_experts)
+        ])
+
+        self.router = nn.Linear(config.hidden_size, self.num_routed, bias=False)
+
+        # Aux-loss-free balancing: a per-expert bias nudged toward even load.
+        # It steers selection only — never the combining weights — so it costs
+        # no gradient interference, unlike an auxiliary loss.
+        self.register_buffer("expert_bias", torch.zeros(self.num_routed), persistent=True)
+        self.bias_update_rate = config.router_bias_update_rate
+
+    def forward(self, x):
+        bsz, seq_len, hidden = x.shape
+        flat = x.view(-1, hidden)
+        n_tokens = flat.size(0)
+
+        logits = self.router(flat)                       # [T, E]
+        probs = F.softmax(logits, dim=-1, dtype=torch.float)
+
+        scores = probs + self.expert_bias if self.aux_free else probs
+        _, topk_idx = torch.topk(scores, self.top_k, dim=-1)
+        topk_w = probs.gather(-1, topk_idx)              # weights from unbiased probs
+        if self.norm_topk_prob:
+            topk_w = topk_w / topk_w.sum(dim=-1, keepdim=True).clamp_min(1e-9)
+        topk_w = topk_w.to(x.dtype)
+
+        out = torch.zeros_like(flat)
+        for expert in self.shared_experts:
+            out = out + expert(flat)
+
+        # one-hot over experts -> per-expert token lists
+        mask = torch.zeros(n_tokens, self.num_routed, dtype=torch.bool, device=x.device)
+        mask.scatter_(1, topk_idx, True)
+        load = mask.sum(0)
+
+        for e in range(self.num_routed):
+            idx = mask[:, e].nonzero(as_tuple=True)[0]
+            if idx.numel() == 0:
+                continue
+            slot = (topk_idx[idx] == e).float().argmax(dim=-1)
+            w = topk_w[idx].gather(-1, slot[:, None])
+            out.index_add_(0, idx, self.experts[e](flat[idx]) * w)
+
+        if self.training and self.aux_free:
+            with torch.no_grad():
+                target = n_tokens * self.top_k / self.num_routed
+                self.expert_bias += self.bias_update_rate * (target - load.float()).sign()
+
+        # Reported for logging even when aux-free balancing is on.
+        frac_tokens = load.float() / (n_tokens * self.top_k)
+        frac_probs = probs.mean(dim=0)
+        aux_loss = self.num_routed * (frac_tokens * frac_probs).sum()
+        z_loss = torch.logsumexp(logits.float(), dim=-1).pow(2).mean()
+
+        return out.view(bsz, seq_len, hidden), aux_loss, z_loss, load
+
+
+class DynamicMindBlock(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.input_layernorm = DynamicMindRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.self_attn = DynamicMindAttention(config, layer_idx)
+        self.post_attention_layernorm = DynamicMindRMSNorm(config.hidden_size, config.rms_norm_eps)
+
+        self.is_moe = layer_idx >= config.first_k_dense_layers
+        if self.is_moe:
+            self.mlp = DynamicMindMoE(config)
+        else:
+            self.mlp = DynamicMindMLP(config.hidden_size, config.intermediate_size)
+
+    def forward(self, x, cos, sin, attention_mask=None, past_key_values=None, cache_position=None):
+        x = x + self.self_attn(self.input_layernorm(x), cos, sin,
+                               attention_mask, past_key_values, cache_position)
+        h = self.post_attention_layernorm(x)
+        if self.is_moe:
+            delta, aux, z, load = self.mlp(h)
+            return x + delta, aux, z, load
+        return x + self.mlp(h), None, None, None
+
+
+class DynamicMindMoEPreTrainedModel(PreTrainedModel):
+    config_class = DynamicMindMoEConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["DynamicMindBlock"]
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+
+class DynamicMindMoEForCausalLM(DynamicMindMoEPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "embed_tokens.weight"}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList([
+            DynamicMindBlock(config, i) for i in range(config.num_hidden_layers)
+        ])
+        self.norm = DynamicMindRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.rotary = DynamicMindRotaryEmbedding(
+            config.hidden_size // config.num_attention_heads,
+            config.rope_theta,
+            config.max_position_embeddings,
+        )
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.embed_tokens.weight
+        self.post_init()
+
+    def tie_weights(self, *args, **kwargs):
+        if getattr(self.config, "tie_word_embeddings", True):
+            self.lm_head.weight = self.embed_tokens.weight
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
+                past_key_values=None, labels=None, use_cache=True, **kwargs):
+        # cache_position is read from kwargs rather than declared: transformers
+        # warns about remote-code models whose signature expects it, and plans
+        # to stop passing it. It is derived below whenever it is absent.
+        cache_position = kwargs.get("cache_position")
+        x = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+        past_len = past_key_values.get_seq_length() if isinstance(past_key_values, Cache) else 0
+
+        if cache_position is None:
+            cache_position = torch.arange(past_len, past_len + x.size(1), device=x.device)
+        if position_ids is None:
+            position_ids = cache_position[None, :]
+
+        cos, sin = self.rotary(x, position_ids)
+
+        causal_mask = None
+        if x.size(1) > 1:
+            total = past_len + x.size(1)
+            causal = torch.tril(torch.ones(x.size(1), total, dtype=torch.bool, device=x.device),
+                                diagonal=past_len)
+            causal_mask = torch.zeros(x.size(1), total, dtype=x.dtype, device=x.device)
+            causal_mask.masked_fill_(~causal, torch.finfo(x.dtype).min)
+            causal_mask = causal_mask[None, None, :, :]
+
+        aux_total = x.new_zeros(())
+        z_total = x.new_zeros(())
+        loads = []
+        for layer in self.layers:
+            x, aux, z, load = layer(x, cos, sin, causal_mask, past_key_values, cache_position)
+            if aux is not None:
+                aux_total = aux_total + aux
+                z_total = z_total + z
+                loads.append(load)
+
+        logits = self.lm_head(self.norm(x))
+
+        loss = None
+        if labels is not None:
+            shift_labels = torch.cat(
+                [labels[:, 1:], labels.new_full((labels.size(0), 1), -100)], dim=1
+            )
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), shift_labels.view(-1))
+            n_moe = max(len(loads), 1)
+            loss = loss + self.config.router_aux_loss_coef * aux_total / n_moe
+            loss = loss + self.config.router_z_loss_coef * z_total / n_moe
+
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=past_key_values if use_cache else None,
+            aux_loss=aux_total / max(len(loads), 1) if loads else None,
+        )
+
+    def expert_load(self):
+        """Per-layer expert token counts from the last forward, for monitoring."""
+        return [m.expert_bias for m in self.modules() if isinstance(m, DynamicMindMoE)]
+
+    def state_dict(self, *args, **kwargs):
+        sd = super().state_dict(*args, **kwargs)
+        if getattr(self.config, "tie_word_embeddings", True):
+            for k in list(sd.keys()):
+                if k == "lm_head.weight" or k.endswith(".lm_head.weight"):
+                    del sd[k]
+        return sd

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|pad|>",
+  "unk_token": "<|unk|>",
+  "model_max_length": 1024,
+  "tokenizer_class": "TokenizersBackend"
+}