Upload SHIVIK-M3 FP32 (2.43B params, 28 layers, 200K vocab)

config.json

@@ -0,0 +1,21 @@
+{
+  "architectures": [
+    "ShivikM3Model"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_shivikM3.ShivikM3Config",
+    "AutoModelForCausalLM": "modeling_shivikM3.ShivikM3Model"
+  },
+  "dtype": "float32",
+  "hidden_size": 2048,
+  "intermediate_size": 7168,
+  "kv_head_split_layer": 14,
+  "model_type": "shivik-m3",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 28,
+  "num_kv_heads": 8,
+  "num_kv_heads_high": 32,
+  "rms_norm_eps": 1e-05,
+  "transformers_version": "4.57.3",
+  "vocab_size": 200018
+}

configuration_shivikM3.py

@@ -0,0 +1,30 @@
+from transformers import PretrainedConfig
+
+class ShivikM3Config(PretrainedConfig):
+    model_type = "shivik-m3"
+
+    def __init__(
+        self,
+        vocab_size=200018,
+        hidden_size=2048,
+        num_hidden_layers=28,
+        num_attention_heads=32,
+        intermediate_size=7168,
+        kv_head_split_layer=14,
+        num_kv_heads=8,
+        num_kv_heads_high=32,
+        tie_word_embeddings=True,
+        rms_norm_eps=1e-5,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.kv_head_split_layer = kv_head_split_layer
+        self.num_kv_heads = num_kv_heads
+        self.num_kv_heads_high = num_kv_heads_high
+        self.tie_word_embeddings = tie_word_embeddings
+        self.rms_norm_eps = rms_norm_eps

generation_config.json

@@ -0,0 +1 @@
+{"bos_token_id": 0, "eos_token_id": 0, "pad_token_id": 0}

model-00001-of-00002.safetensors

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

model-00002-of-00002.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,263 @@
+{
+  "metadata": {
+    "total_parameters": 2434197504,
+    "total_size": 9736790016
+  },
+  "weight_map": {
+    "embed_tokens.weight": "model-00001-of-00002.safetensors",
+    "layers.0.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.0.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.0.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.0.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.0.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.0.mlp.fc.weight": "model-00001-of-00002.safetensors",
+    "layers.0.mlp.gate.weight": "model-00001-of-00002.safetensors",
+    "layers.0.mlp.proj.weight": "model-00001-of-00002.safetensors",
+    "layers.0.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.1.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.1.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.1.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.1.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.1.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.1.mlp.fc.weight": "model-00001-of-00002.safetensors",
+    "layers.1.mlp.gate.weight": "model-00001-of-00002.safetensors",
+    "layers.1.mlp.proj.weight": "model-00001-of-00002.safetensors",
+    "layers.1.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.10.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.10.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.10.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.10.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.10.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.10.mlp.fc.weight": "model-00001-of-00002.safetensors",
+    "layers.10.mlp.gate.weight": "model-00001-of-00002.safetensors",
+    "layers.10.mlp.proj.weight": "model-00001-of-00002.safetensors",
+    "layers.10.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.11.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.11.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.11.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.11.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.11.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.11.mlp.fc.weight": "model-00001-of-00002.safetensors",
+    "layers.11.mlp.gate.weight": "model-00001-of-00002.safetensors",
+    "layers.11.mlp.proj.weight": "model-00001-of-00002.safetensors",
+    "layers.11.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.12.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.12.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.12.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.12.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.12.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.12.mlp.fc.weight": "model-00001-of-00002.safetensors",
+    "layers.12.mlp.gate.weight": "model-00001-of-00002.safetensors",
+    "layers.12.mlp.proj.weight": "model-00001-of-00002.safetensors",
+    "layers.12.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.13.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.13.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.13.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.13.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.13.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.13.mlp.fc.weight": "model-00001-of-00002.safetensors",
+    "layers.13.mlp.gate.weight": "model-00001-of-00002.safetensors",
+    "layers.13.mlp.proj.weight": "model-00001-of-00002.safetensors",
+    "layers.13.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.14.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.14.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.14.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.14.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.14.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.14.mlp.fc.weight": "model-00001-of-00002.safetensors",
+    "layers.14.mlp.gate.weight": "model-00001-of-00002.safetensors",
+    "layers.14.mlp.proj.weight": "model-00001-of-00002.safetensors",
+    "layers.14.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.15.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.15.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.15.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.15.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.15.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.15.mlp.fc.weight": "model-00002-of-00002.safetensors",
+    "layers.15.mlp.gate.weight": "model-00002-of-00002.safetensors",
+    "layers.15.mlp.proj.weight": "model-00002-of-00002.safetensors",
+    "layers.15.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.16.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.16.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.16.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.16.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.16.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.16.mlp.fc.weight": "model-00002-of-00002.safetensors",
+    "layers.16.mlp.gate.weight": "model-00002-of-00002.safetensors",
+    "layers.16.mlp.proj.weight": "model-00002-of-00002.safetensors",
+    "layers.16.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.17.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.17.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.17.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.17.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.17.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.17.mlp.fc.weight": "model-00002-of-00002.safetensors",
+    "layers.17.mlp.gate.weight": "model-00002-of-00002.safetensors",
+    "layers.17.mlp.proj.weight": "model-00002-of-00002.safetensors",
+    "layers.17.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.18.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.18.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.18.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.18.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.18.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.18.mlp.fc.weight": "model-00002-of-00002.safetensors",
+    "layers.18.mlp.gate.weight": "model-00002-of-00002.safetensors",
+    "layers.18.mlp.proj.weight": "model-00002-of-00002.safetensors",
+    "layers.18.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.19.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.19.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.19.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.19.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.19.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.19.mlp.fc.weight": "model-00002-of-00002.safetensors",
+    "layers.19.mlp.gate.weight": "model-00002-of-00002.safetensors",
+    "layers.19.mlp.proj.weight": "model-00002-of-00002.safetensors",
+    "layers.19.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.2.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.2.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.2.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.2.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.2.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.2.mlp.fc.weight": "model-00001-of-00002.safetensors",
+    "layers.2.mlp.gate.weight": "model-00001-of-00002.safetensors",
+    "layers.2.mlp.proj.weight": "model-00001-of-00002.safetensors",
+    "layers.2.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.20.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.20.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.20.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.20.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.20.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.20.mlp.fc.weight": "model-00002-of-00002.safetensors",
+    "layers.20.mlp.gate.weight": "model-00002-of-00002.safetensors",
+    "layers.20.mlp.proj.weight": "model-00002-of-00002.safetensors",
+    "layers.20.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.21.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.21.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.21.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.21.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.21.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.21.mlp.fc.weight": "model-00002-of-00002.safetensors",
+    "layers.21.mlp.gate.weight": "model-00002-of-00002.safetensors",
+    "layers.21.mlp.proj.weight": "model-00002-of-00002.safetensors",
+    "layers.21.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.22.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.22.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.22.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.22.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.22.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.22.mlp.fc.weight": "model-00002-of-00002.safetensors",
+    "layers.22.mlp.gate.weight": "model-00002-of-00002.safetensors",
+    "layers.22.mlp.proj.weight": "model-00002-of-00002.safetensors",
+    "layers.22.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.23.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.23.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.23.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.23.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.23.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.23.mlp.fc.weight": "model-00002-of-00002.safetensors",
+    "layers.23.mlp.gate.weight": "model-00002-of-00002.safetensors",
+    "layers.23.mlp.proj.weight": "model-00002-of-00002.safetensors",
+    "layers.23.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.24.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.24.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.24.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.24.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.24.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.24.mlp.fc.weight": "model-00002-of-00002.safetensors",
+    "layers.24.mlp.gate.weight": "model-00002-of-00002.safetensors",
+    "layers.24.mlp.proj.weight": "model-00002-of-00002.safetensors",
+    "layers.24.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.25.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.25.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.25.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.25.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.25.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.25.mlp.fc.weight": "model-00002-of-00002.safetensors",
+    "layers.25.mlp.gate.weight": "model-00002-of-00002.safetensors",
+    "layers.25.mlp.proj.weight": "model-00002-of-00002.safetensors",
+    "layers.25.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.26.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.26.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.26.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.26.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.26.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.26.mlp.fc.weight": "model-00002-of-00002.safetensors",
+    "layers.26.mlp.gate.weight": "model-00002-of-00002.safetensors",
+    "layers.26.mlp.proj.weight": "model-00002-of-00002.safetensors",
+    "layers.26.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.27.attn.k_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.27.attn.o_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.27.attn.q_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.27.attn.v_proj.weight": "model-00002-of-00002.safetensors",
+    "layers.27.attn_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.27.mlp.fc.weight": "model-00002-of-00002.safetensors",
+    "layers.27.mlp.gate.weight": "model-00002-of-00002.safetensors",
+    "layers.27.mlp.proj.weight": "model-00002-of-00002.safetensors",
+    "layers.27.mlp_norm.weight": "model-00002-of-00002.safetensors",
+    "layers.3.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.3.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.3.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.3.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.3.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.3.mlp.fc.weight": "model-00001-of-00002.safetensors",
+    "layers.3.mlp.gate.weight": "model-00001-of-00002.safetensors",
+    "layers.3.mlp.proj.weight": "model-00001-of-00002.safetensors",
+    "layers.3.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.4.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.4.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.4.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.4.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.4.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.4.mlp.fc.weight": "model-00001-of-00002.safetensors",
+    "layers.4.mlp.gate.weight": "model-00001-of-00002.safetensors",
+    "layers.4.mlp.proj.weight": "model-00001-of-00002.safetensors",
+    "layers.4.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.5.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.5.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.5.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.5.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.5.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.5.mlp.fc.weight": "model-00001-of-00002.safetensors",
+    "layers.5.mlp.gate.weight": "model-00001-of-00002.safetensors",
+    "layers.5.mlp.proj.weight": "model-00001-of-00002.safetensors",
+    "layers.5.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.6.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.6.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.6.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.6.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.6.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.6.mlp.fc.weight": "model-00001-of-00002.safetensors",
+    "layers.6.mlp.gate.weight": "model-00001-of-00002.safetensors",
+    "layers.6.mlp.proj.weight": "model-00001-of-00002.safetensors",
+    "layers.6.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.7.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.7.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.7.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.7.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.7.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.7.mlp.fc.weight": "model-00001-of-00002.safetensors",
+    "layers.7.mlp.gate.weight": "model-00001-of-00002.safetensors",
+    "layers.7.mlp.proj.weight": "model-00001-of-00002.safetensors",
+    "layers.7.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.8.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.8.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.8.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.8.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.8.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.8.mlp.fc.weight": "model-00001-of-00002.safetensors",
+    "layers.8.mlp.gate.weight": "model-00001-of-00002.safetensors",
+    "layers.8.mlp.proj.weight": "model-00001-of-00002.safetensors",
+    "layers.8.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.9.attn.k_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.9.attn.o_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.9.attn.q_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.9.attn.v_proj.weight": "model-00001-of-00002.safetensors",
+    "layers.9.attn_norm.weight": "model-00001-of-00002.safetensors",
+    "layers.9.mlp.fc.weight": "model-00001-of-00002.safetensors",
+    "layers.9.mlp.gate.weight": "model-00001-of-00002.safetensors",
+    "layers.9.mlp.proj.weight": "model-00001-of-00002.safetensors",
+    "layers.9.mlp_norm.weight": "model-00001-of-00002.safetensors",
+    "lm_head.weight": "model-00002-of-00002.safetensors",
+    "norm.weight": "model-00002-of-00002.safetensors"
+  }
+}

modeling_shivikM3.py

@@ -0,0 +1,135 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from configuration_shivikM3 import ShivikM3Config
+
+class RMSNorm(nn.Module):
+    def __init__(self, d, eps=1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(d))
+        self.eps = eps
+
+    def forward(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+def apply_rope(q, k, cos, sin):
+    q = (q * cos) + (rotate_half(q) * sin)
+    k = (k * cos) + (rotate_half(k) * sin)
+    return q, k
+
+class ShivikM3Attention(nn.Module):
+    def __init__(self, c, idx):
+        super().__init__()
+        self.num_q = c.num_attention_heads
+        self.head_dim = c.hidden_size // self.num_q
+        self.num_kv = c.num_kv_heads if idx < c.kv_head_split_layer else c.num_kv_heads_high
+        self.q_proj = nn.Linear(c.hidden_size, self.num_q * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(c.hidden_size, self.num_kv * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(c.hidden_size, self.num_kv * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(c.hidden_size, c.hidden_size, bias=False)
+
+    def forward(self, x, cos, sin, mask=None):
+        B, T, C = x.size()
+        q = self.q_proj(x).view(B, T, self.num_q, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.num_kv, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.num_kv, self.head_dim).transpose(1, 2)
+
+        if self.num_kv != self.num_q:
+            k = k.repeat_interleave(self.num_q // self.num_kv, dim=1)
+            v = v.repeat_interleave(self.num_q // self.num_kv, dim=1)
+
+        q, k = apply_rope(q, k, cos, sin)
+
+        attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        if mask is not None:
+            attn = attn + mask
+        attn = F.softmax(attn, dim=-1)
+        out = attn @ v
+
+        return self.o_proj(out.transpose(1, 2).reshape(B, T, C))
+
+class ShivikM3MLP(nn.Module):
+    def __init__(self, c):
+        super().__init__()
+        self.gate = nn.Linear(c.hidden_size, c.intermediate_size, bias=False)
+        self.fc = nn.Linear(c.hidden_size, c.intermediate_size, bias=False)
+        self.proj = nn.Linear(c.intermediate_size, c.hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.proj(F.silu(self.gate(x)) * self.fc(x))
+
+class ShivikM3Block(nn.Module):
+    def __init__(self, c, idx):
+        super().__init__()
+        self.attn_norm = RMSNorm(c.hidden_size, c.rms_norm_eps)
+        self.mlp_norm = RMSNorm(c.hidden_size, c.rms_norm_eps)
+        self.attn = ShivikM3Attention(c, idx)
+        self.mlp = ShivikM3MLP(c)
+
+    def forward(self, x, cos, sin, mask=None):
+        h = x + self.attn(self.attn_norm(x), cos, sin, mask)
+        h = h + self.mlp(self.mlp_norm(h))
+        return h
+
+class ShivikM3Model(PreTrainedModel):
+    config_class = ShivikM3Config
+    supports_gradient_checkpointing = True
+
+    def __init__(self, c):
+        super().__init__(c)
+        self.config = c
+        self.gradient_checkpointing = False
+        self.embed_tokens = nn.Embedding(c.vocab_size, c.hidden_size)
+        self.layers = nn.ModuleList([ShivikM3Block(c, i) for i in range(c.num_hidden_layers)])
+        self.norm = RMSNorm(c.hidden_size, c.rms_norm_eps)
+        self.lm_head = nn.Linear(c.hidden_size, c.vocab_size, bias=False)
+        self.post_init()
+
+    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
+        B, T = input_ids.size()
+        x = self.embed_tokens(input_ids)
+
+        # Create RoPE - match model dtype
+        head_dim = self.config.hidden_size // self.config.num_attention_heads
+        device = x.device
+        dtype = x.dtype
+
+        pos = torch.arange(T, device=device, dtype=dtype)
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, head_dim, 2, device=device, dtype=dtype) / head_dim))
+        freqs = torch.outer(pos, inv_freq)
+        freqs = torch.cat((freqs, freqs), dim=-1)
+        cos = freqs.cos().view(1, 1, T, head_dim)
+        sin = freqs.sin().view(1, 1, T, head_dim)
+
+        # Create causal mask - match model dtype
+        mask = torch.triu(
+            torch.full((T, T), float('-inf'), device=device, dtype=dtype),
+            diagonal=1
+        ).view(1, 1, T, T)
+
+        if attention_mask is not None:
+            mask = mask + (1.0 - attention_mask[:, None, None, :]) * torch.finfo(dtype).min
+
+        for block in self.layers:
+            if self.gradient_checkpointing and self.training:
+                x = torch.utils.checkpoint.checkpoint(block, x, cos, sin, mask, use_reentrant=False)
+            else:
+                x = block(x, cos, sin, mask)
+
+        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, self.config.vocab_size), shift_labels.view(-1))
+
+        return CausalLMOutputWithPast(loss=loss, logits=logits)

special_tokens_map.json

@@ -0,0 +1,13 @@
+{
+  "unk_token": "<unk>",
+  "pad_token": "<pad>",
+  "bos_token": "<think>",
+  "eos_token": "</think>",
+  "additional_special_tokens": [
+    "<answer>",
+    "</answer>",
+    "<context>",
+    "</context>",
+    "<end>"
+  ]
+}

tokenizer_config.json

@@ -0,0 +1,18 @@
+{
+  "model_type": "wordlevel",
+  "unk_token": "<unk>",
+  "pad_token": "<pad>",
+  "bos_token": "<think>",
+  "eos_token": "</think>",
+  "special_tokens": [
+    "<unk>",
+    "<pad>",
+    "<think>",
+    "</think>",
+    "<answer>",
+    "</answer>",
+    "<context>",
+    "</context>",
+    "<end>"
+  ]
+}