Add SFT weights, tokenizer, chat template and training metadata

chat_template.jinja

@@ -0,0 +1,10 @@
+{{ bos_token }}{% for message in messages %}{% if message['role'] == 'assistant' %}<|assistant|>
+{{ message['content'] }}{{ eos_token }}
+{% elif message['role'] == 'system' %}<|system|>
+{{ message['content'] }}
+{% elif message['role'] == 'user' %}<|user|>
+{{ message['content'] }}
+{% else %}<|{{ message['role'] }}|>
+{{ message['content'] }}
+{% endif %}{% endfor %}{% if add_generation_prompt %}<|assistant|>
+{% endif %}

config.json

@@ -0,0 +1,26 @@
+{
+  "architectures": [
+    "DynamicMindForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_dynamicmind.DynamicMindConfig",
+    "AutoModelForCausalLM": "modeling_dynamicmind.DynamicMindForCausalLM"
+  },
+  "bos_token_id": 0,
+  "dtype": "float32",
+  "eos_token_id": 0,
+  "hidden_size": 256,
+  "intermediate_size": 768,
+  "max_position_embeddings": 1024,
+  "model_type": "dynamicmind",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 9,
+  "num_key_value_heads": 2,
+  "pad_token_id": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.5.3",
+  "vocab_size": 8192
+}

configuration_dynamicmind.py

@@ -0,0 +1,42 @@
+from transformers.configuration_utils import PretrainedConfig
+
+
+class DynamicMindConfig(PretrainedConfig):
+    model_type = "dynamicmind"
+
+    def __init__(
+        self,
+        vocab_size=8192,
+        hidden_size=256,
+        intermediate_size=768,
+        num_hidden_layers=9,
+        num_attention_heads=8,
+        num_key_value_heads=2,
+        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.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.max_position_embeddings = max_position_embeddings
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = rope_theta
+        self.attention_dropout = attention_dropout

generation_config.json

@@ -0,0 +1,9 @@
+{
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "pad_token_id": 1,
+  "do_sample": false,
+  "repetition_penalty": 1.1,
+  "max_new_tokens": 192,
+  "transformers_version": "5.7.0"
+}

model.safetensors

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

modeling_dynamicmind.py

@@ -0,0 +1,204 @@
+import math
+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 CausalLMOutput
+from .configuration_dynamicmind import DynamicMindConfig
+
+
+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()
+        var = x.pow(2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.eps)
+        return (self.weight * x).to(dtype)
+
+
+def apply_rope(q, k, rope_theta):
+    # q: [B, H, T, D], k: [B, KVH, T, D]
+    device = q.device
+    dtype = q.dtype
+    seq_len = q.size(-2)
+    head_dim = q.size(-1)
+
+    inv_freq = 1.0 / (
+        rope_theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)
+    )
+    t = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(t, inv_freq)
+
+    cos = freqs.cos()[None, None, :, :].to(dtype)
+    sin = freqs.sin()[None, None, :, :].to(dtype)
+
+    def rotate(x):
+        x_even = x[..., 0::2]
+        x_odd = x[..., 1::2]
+        x_rot_even = x_even * cos - x_odd * sin
+        x_rot_odd = x_even * sin + x_odd * cos
+        return torch.stack((x_rot_even, x_rot_odd), dim=-1).flatten(-2)
+
+    return rotate(q), rotate(k)
+
+
+class DynamicMindAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        self.hidden_size = config.hidden_size
+        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.rope_theta = config.rope_theta
+        self.attention_dropout = config.attention_dropout
+
+        assert self.hidden_size % self.num_heads == 0
+        assert self.num_heads % self.num_kv_heads == 0
+
+        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):
+        bsz, seq_len, _ = x.shape
+
+        q = self.q_proj(x).view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        q, k = apply_rope(q, k, self.rope_theta)
+
+        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)
+
+        y = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=None,
+            dropout_p=self.attention_dropout if self.training else 0.0,
+            is_causal=True,
+        )
+
+        y = y.transpose(1, 2).contiguous().view(bsz, seq_len, self.hidden_size)
+        return self.o_proj(y)
+
+
+class DynamicMindMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class DynamicMindBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.input_layernorm = DynamicMindRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.self_attn = DynamicMindAttention(config)
+        self.post_attention_layernorm = DynamicMindRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.mlp = DynamicMindMLP(config)
+
+    def forward(self, x):
+        x = x + self.self_attn(self.input_layernorm(x))
+        x = x + self.mlp(self.post_attention_layernorm(x))
+        return x
+
+
+class DynamicMindPreTrainedModel(PreTrainedModel):
+    config_class = DynamicMindConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["DynamicMindBlock"]
+
+    def _init_weights(self, module):
+        std = 0.02
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+            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=std)
+
+
+class DynamicMindForCausalLM(DynamicMindPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "embed_tokens.weight"}
+    _keys_to_ignore_on_load_missing = [r"lm_head.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) for _ 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)
+
+        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 get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+
+    def forward(self, input_ids=None, labels=None, **kwargs):
+        x = self.embed_tokens(input_ids)
+
+        for layer in self.layers:
+            x = layer(x)
+
+        x = self.norm(x)
+        logits = self.lm_head(x)
+
+        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, shift_logits.size(-1)),
+                shift_labels.view(-1),
+            )
+
+        return CausalLMOutput(loss=loss, logits=logits)
+
+    def state_dict(self, *args, **kwargs):
+        sd = super().state_dict(*args, **kwargs)
+        # lm_head.weight is tied to embed_tokens.weight. Safetensors cannot store
+        # duplicate shared tensors unless one key is removed.
+        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
+
+    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
+        # Accepted so stock generate() validation passes; causality comes from
+        # is_causal=True and generation is unpadded batch-size 1.
+        return {"input_ids": input_ids, "attention_mask": attention_mask}

sft_metadata.json

@@ -0,0 +1,27 @@
+{
+  "base_model": "DedeProGames/DynamicMind-Mini",
+  "dataset": "HuggingFaceTB/smol-smoltalk",
+  "dataset_config": "default",
+  "dataset_split": "train",
+  "training_type": "full_parameter_supervised_finetuning",
+  "loss_mask": "assistant_content_and_eos_only",
+  "global_step": 3440,
+  "tokens_seen": 507248640,
+  "supervised_tokens_seen": 361226420,
+  "sequence_length": 1024,
+  "effective_batch_sequences": 144,
+  "tokens_per_step": 147456,
+  "peak_learning_rate": 0.0002,
+  "warmup_steps": 100,
+  "lr_schedule": "constant_after_warmup",
+  "optimizer": "AdamW",
+  "betas": [
+    0.9,
+    0.95
+  ],
+  "weight_decay": 0.1,
+  "grad_clip": 1.0,
+  "seed": 1337,
+  "final_loss": 1.917157530784607,
+  "recipe_source": "BananaMind/BananaMind-2-Nano-Chat (adapted to 1024-token context)"
+}

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"
+}