Add DynamicMind-MoE-Instruct: 2-epoch SFT on smol-smoltalk, custom MoE code, EOS-weighted loss

README.md

@@ -0,0 +1,115 @@
+---
+license: apache-2.0
+language:
+- en
+library_name: transformers
+pipeline_tag: text-generation
+base_model: DedeProGames/DynamicMind-MoE
+datasets:
+- HuggingFaceTB/smol-smoltalk
+tags:
+- dynamicmind
+- mixture-of-experts
+- sparse-moe
+- instruct
+- sft
+- conversational
+- custom-code
+- trust-remote-code
+---
+
+# DynamicMind-MoE-Instruct
+
+An instruction-tuned version of [DynamicMind-MoE](https://huggingface.co/DedeProGames/DynamicMind-MoE), a sparse mixture-of-experts model with 30.2M total parameters and 8.9M active per token. It was fully fine-tuned for two epochs on [smol-smoltalk](https://huggingface.co/datasets/HuggingFaceTB/smol-smoltalk) with the loss computed only on assistant tokens (end-of-turn tokens weighted 5×, so the model learns to stop).
+
+At this size it is a small research model. It follows the chat format and answers simple requests, but factual knowledge, reasoning and arithmetic are weak. See the limitations below.
+
+## Quickstart
+
+This model uses custom architecture code, so load it with `trust_remote_code=True`. Use the recommended sampling settings (stored in `generation_config.json`). With near-greedy decoding this model can fall into repetition loops.
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+repo = "DedeProGames/DynamicMind-MoE-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(repo)
+model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True, torch_dtype=torch.float32).eval()
+
+messages = [{"role": "user", "content": "What is photosynthesis?"}]
+input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=False)
+out = model.generate(input_ids, max_new_tokens=256)  # temperature 0.7, top-p 0.9, repetition penalty 1.1 from generation_config.json
+print(tokenizer.decode(out[0, input_ids.shape[1]:], skip_special_tokens=True))
+```
+
+Keep prompt + answer within the 1,024-token context.
+
+## Chat format
+
+This vocabulary has no chat tokens, so the format is plain text. `apply_chat_template` renders it:
+
+```text
+<|endoftext|>### User:
+{user message}
+### Assistant:
+{assistant reply}
+<|endoftext|>### User:
+...
+```
+
+`<|endoftext|>` (id 0) is the BOS, the end-of-turn token and the EOS.
+
+## Architecture
+
+Same as DynamicMind-MoE: 9 layers, hidden size 256, 8 attention heads with 2 KV heads, one always-on shared expert plus 14 routed experts per layer with the top 2 selected per token, auxiliary-loss-free balancing, tied embeddings, 8,192-token digit-aware vocabulary, 1,024-token context. 30,150,912 total parameters, 8,917,248 active per token. Weights in BF16 (`model.safetensors`, 60 MB).
+
+## Training
+
+| | |
+|---|---|
+| Base | DedeProGames/DynamicMind-MoE |
+| Data | smol-smoltalk train split, conversations of up to 1,024 tokens (221,811 with this tokenizer; 68.9M supervised tokens) |
+| Epochs | 2 (13,864 steps, batch 32 = 16 × 2 gradient accumulation) |
+| Optimizer | AdamW (β = 0.9, 0.95), weight decay 0.1 on 2D matrices, grad clip 1.0 |
+| Schedule | lr 1e-4 (the base model's pretraining peak), 300 warmup steps, cosine to 1e-5 |
+| Loss | assistant tokens; `<|endoftext|>` weighted 5×; router balance loss (coefficient 0.01) added |
+| Precision | bf16 autocast, fp32 master weights |
+| Hardware | 1× NVIDIA RTX 3060 12 GB, 239 minutes |
+
+Held-out loss on assistant tokens (smol-smoltalk test split): 1.422 at the end of training.
+
+## Evaluation
+
+Benchmarks follow the protocol of the [Kiyo-230M-Instruct](https://huggingface.co/DedeProGames/Kiyo-230M-Instruct) card: FP16, the model's own chat template, no system prompt, new-token limits IFEval 1,024, GPQA 384, LiveBench 512. The context is 1,024 tokens, so the generation limit is capped at 1,024 minus the prompt length. The main setting uses temperature 0.1, top-k 50 and repetition penalty 1.05, as on the card (GPQA is sampled at T = 0.6, as on the card). The second setting is the recommended one shipped in `generation_config.json`.
+
+| Benchmark | Metric | DynamicMind-MoE-Instruct, protocol | DynamicMind-MoE-Instruct, recommended sampling |
+|---|---|---:|---:|
+| GPQA Diamond | pass@1 (198 × 5 samples, T = 0.6) | 12.32 | not run |
+| IFEval | prompt-level strict | 18.85 | 14.97 |
+| LiveBench instruction following | strict, 200 questions | 18.66 | not run |
+| CaseReportBench | token F1 | not run (not on the Hub) | not run |
+| EOS rate (IFEval) | share of answers that stop | 0.88 | 0.99 |
+
+For reference, the Kiyo card reports (not re-run here): Kiyo-230M-Instruct GPQA 21.11, IFEval 17.38, LiveBench 18.77; SmolLM2-135M-Instruct GPQA 16.57, IFEval 20.89, LiveBench 24.23; BananaMind-2-Pro GPQA 15.05, IFEval 20.15, LiveBench 19.75; Supra2-100M-Instruct GPQA 5.56, IFEval 9.06, LiveBench 11.28.
+
+Notes:
+
+- With the protocol setting, about 12% of IFEval answers never stop and loop. Under the recommended sampling this drops to about 1%, at a cost of about 4 IFEval points.
+- GPQA: only 46% of answers contain an identifiable option letter, and 40 prompts were truncated to fit the 1,024-token context. Among answered questions, accuracy is 26.6%, close to the 25% chance level.
+- The GPQA prompt was written for this run, not taken from the Kiyo card, so the number is not strictly comparable.
+- Single run, single seed. Differences of a point or two are within noise.
+
+Per-sample outputs and summaries are in [`eval/`](eval/).
+
+## Limitations
+
+- It is small (8.9M active parameters per token). Expect factual errors, invented details, arithmetic mistakes, weak multi-step reasoning, and repetitive text in long answers.
+- Under near-greedy decoding it can loop. Keep the recommended sampling settings.
+- Context is 1,024 tokens.
+- English only.
+- Requires `trust_remote_code=True`, since the architecture is custom.
+- No safety tuning and no safety evaluation. Do not use it for medical, legal, financial or other high-stakes decisions.
+
+## License
+
+Apache-2.0, as for the base model. The SFT data (smol-smoltalk) is also Apache-2.0.

config.json

@@ -0,0 +1,36 @@
+{
+  "architectures": [
+    "DynamicMindMoEForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_dynamicmind_moe.DynamicMindMoEConfig",
+    "AutoModelForCausalLM": "modeling_dynamicmind_moe.DynamicMindMoEForCausalLM"
+  },
+  "bos_token_id": 0,
+  "dtype": "bfloat16",
+  "eos_token_id": 0,
+  "first_k_dense_layers": 0,
+  "hidden_size": 256,
+  "intermediate_size": 768,
+  "max_position_embeddings": 1024,
+  "model_type": "dynamicmind_moe",
+  "moe_intermediate_size": 256,
+  "norm_topk_prob": true,
+  "num_attention_heads": 8,
+  "num_experts_per_token": 2,
+  "num_hidden_layers": 9,
+  "num_key_value_heads": 2,
+  "num_routed_experts": 14,
+  "num_shared_experts": 1,
+  "pad_token_id": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 10000.0,
+  "router_aux_loss_coef": 0.01,
+  "router_bias_update_rate": 0.001,
+  "router_z_loss_coef": 0.001,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.53.3",
+  "use_aux_loss_free_balancing": true,
+  "vocab_size": 8192
+}

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

eval/protocol/summary.json

@@ -0,0 +1,25 @@
+{
+  "ifeval": {
+    "metric": "prompt-level strict",
+    "score": 18.853974121996302,
+    "n": 541,
+    "eos_rate": 0.8835489833641405,
+    "minutes": 29.08
+  },
+  "gpqa": {
+    "metric": "pass@1 (990 attempts, no-letter counts wrong)",
+    "score": 12.323232323232324,
+    "n_attempts": 990,
+    "answered_rate": 46.26262626262626,
+    "accuracy_among_answered": 26.637554585152838,
+    "prompts_truncated": 40,
+    "minutes": 25.69
+  },
+  "livebench": {
+    "metric": "per-question mean(all followed, fraction followed), strict",
+    "score": 18.658333333333335,
+    "n": 200,
+    "eos_rate": 0.835,
+    "minutes": 9.9
+  }
+}

eval/recommended_sampling/summary.json

@@ -0,0 +1,9 @@
+{
+  "ifeval": {
+    "metric": "prompt-level strict",
+    "score": 14.972273567467653,
+    "n": 541,
+    "eos_rate": 0.988909426987061,
+    "minutes": 21.12
+  }
+}

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "bos_token_id": 0,
+  "pad_token_id": 1,
+  "eos_token_id": 0,
+  "do_sample": true,
+  "temperature": 0.7,
+  "top_p": 0.9,
+  "repetition_penalty": 1.1,
+  "max_new_tokens": 512
+}

model.safetensors

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

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 @@
+{
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|pad|>",
+  "unk_token": "<|unk|>",
+  "model_max_length": 1024,
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "chat_template": "{{- '<|endoftext|>' -}}{%- for message in messages -%}{%- if message['role'] == 'system' -%}{{ '### System:\n' + message['content'].strip() + '\n' }}{%- elif message['role'] == 'user' -%}{{ '### User:\n' + message['content'].strip() + '\n' }}{%- elif message['role'] == 'assistant' -%}{{ '### Assistant:\n' + message['content'].strip() + '\n<|endoftext|>' }}{%- endif -%}{%- endfor -%}{%- if add_generation_prompt -%}{{ '### Assistant:\n' }}{%- endif -%}"
+}