Update tokenizer, configuration and training reports

balanced_model.py

@@ -0,0 +1,82 @@
+"""Native Qwen3-style GQA model with explicit value residuals and tied embeddings.
+
+State names match the reviewed JugnuVR checkpoint. Explicit v0 avoids a shared
+mutable context dictionary in the training forward. No KV cache is implemented.
+"""
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+class RMSNorm(nn.Module):
+    def __init__(self, width, eps=1e-6):
+        super().__init__(); self.weight=nn.Parameter(torch.ones(width)); self.eps=eps
+    def forward(self,x):
+        dtype=x.dtype; x=x.float(); x=x*torch.rsqrt(x.square().mean(-1,keepdim=True)+self.eps)
+        return self.weight*x.to(dtype)
+
+class ValueProjection(nn.Linear):
+    def __init__(self,width,out,first):
+        super().__init__(width,out,bias=False)
+        if not first:self.vr_lambda=nn.Parameter(torch.zeros(1))
+
+class Attention(nn.Module):
+    def __init__(self,c,index):
+        super().__init__();h=c['hidden_size'];self.heads=c['num_attention_heads'];self.kv=c['num_key_value_heads'];self.d=c.get('head_dim',h//self.heads);self.first=index==0
+        self.q_proj=nn.Linear(h,self.heads*self.d,bias=False);self.k_proj=nn.Linear(h,self.kv*self.d,bias=False)
+        self.v_proj=ValueProjection(h,self.kv*self.d,self.first);self.o_proj=nn.Linear(self.heads*self.d,h,bias=False)
+        self.q_norm=RMSNorm(self.d,c.get('rms_norm_eps',1e-6));self.k_norm=RMSNorm(self.d,c.get('rms_norm_eps',1e-6))
+    def forward(self,x,cos,sin,v0):
+        b,t,_=x.shape
+        q=self.q_norm(self.q_proj(x).view(b,t,self.heads,self.d)).transpose(1,2)
+        k=self.k_norm(self.k_proj(x).view(b,t,self.kv,self.d)).transpose(1,2)
+        v=self.v_proj(x)
+        if self.first:v0=v
+        else:v=v+self.v_proj.vr_lambda*v0
+        v=v.view(b,t,self.kv,self.d).transpose(1,2)
+        def rotate(z):return torch.cat((-z[...,self.d//2:],z[...,:self.d//2]),dim=-1)
+        q=q*cos+rotate(q)*sin;k=k*cos+rotate(k)*sin
+        k=k.repeat_interleave(self.heads//self.kv,dim=1);v=v.repeat_interleave(self.heads//self.kv,dim=1)
+        out=F.scaled_dot_product_attention(q,k,v,is_causal=True)
+        return self.o_proj(out.transpose(1,2).contiguous().view(b,t,-1)),v0
+
+class MLP(nn.Module):
+    def __init__(self,c):
+        super().__init__();h=c['hidden_size'];f=c['intermediate_size']
+        self.gate_proj=nn.Linear(h,f,bias=False);self.up_proj=nn.Linear(h,f,bias=False);self.down_proj=nn.Linear(f,h,bias=False)
+    def forward(self,x):return self.down_proj(F.silu(self.gate_proj(x))*self.up_proj(x))
+
+class Block(nn.Module):
+    def __init__(self,c,i):
+        super().__init__();self.self_attn=Attention(c,i);self.mlp=MLP(c)
+        self.input_layernorm=RMSNorm(c['hidden_size'],c.get('rms_norm_eps',1e-6));self.post_attention_layernorm=RMSNorm(c['hidden_size'],c.get('rms_norm_eps',1e-6))
+    def forward(self,x,cos,sin,v0):
+        a,v0=self.self_attn(self.input_layernorm(x),cos,sin,v0);x=x+a
+        return x+self.mlp(self.post_attention_layernorm(x)),v0
+
+class Backbone(nn.Module):
+    def __init__(self,c):
+        super().__init__();self.embed_tokens=nn.Embedding(c['vocab_size'],c['hidden_size']);self.layers=nn.ModuleList([Block(c,i) for i in range(c['num_hidden_layers'])]);self.norm=RMSNorm(c['hidden_size'],c.get('rms_norm_eps',1e-6))
+
+class BalancedLM(nn.Module):
+    def __init__(self,c):
+        super().__init__();self.config=dict(c)
+        assert c['num_attention_heads']%c['num_key_value_heads']==0
+        assert c['hidden_size']==c['num_attention_heads']*c.get('head_dim',64)
+        self.model=Backbone(c);self.lm_head=nn.Linear(c['hidden_size'],c['vocab_size'],bias=False);self.lm_head.weight=self.model.embed_tokens.weight
+        d=c.get('head_dim',64);theta=c.get('rope_parameters',{}).get('rope_theta',10000.)
+        inv=1./(theta**(torch.arange(0,d,2,dtype=torch.float32)/d))
+        phases=torch.outer(torch.arange(c['max_position_embeddings'],dtype=torch.float32),inv);phases=torch.cat((phases,phases),dim=-1)
+        self.register_buffer('cos',phases.cos()[None,None],persistent=False);self.register_buffer('sin',phases.sin()[None,None],persistent=False)
+    def forward(self,ids,targets=None):
+        x=self.model.embed_tokens(ids);v0=None;t=ids.shape[1]
+        for layer in self.model.layers:x,v0=layer(x,self.cos[:,:,:t],self.sin[:,:,:t],v0)
+        logits=self.lm_head(self.model.norm(x))
+        loss=None if targets is None else F.cross_entropy(logits.float().flatten(0,1),targets.flatten())
+        return logits,loss
+    @torch.no_grad()
+    def initialize(self,seed):
+        torch.manual_seed(seed)
+        for name,p in self.named_parameters():
+            if p.ndim==2:p.normal_(mean=0,std=0.02)
+            elif name.endswith('vr_lambda'):p.zero_()
+            else:p.fill_(1)

config.json

@@ -0,0 +1,65 @@
+{
+  "architectures": [
+    "JugnuVRForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_jugnu_vr.JugnuVRConfig",
+    "AutoModelForCausalLM": "modeling_jugnu_vr.JugnuVRForCausalLM"
+  },
+  "bos_token_id": 0,
+  "dtype": "float32",
+  "eos_token_id": 0,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 576,
+  "initializer_range": 0.02,
+  "intermediate_size": 2496,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 2048,
+  "max_window_layers": 28,
+  "model_type": "jugnu_vr",
+  "num_attention_heads": 9,
+  "num_hidden_layers": 26,
+  "num_key_value_heads": 3,
+  "pad_token_id": null,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.16.1",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 24576
+}

load_model.py

@@ -0,0 +1,11 @@
+import json
+from pathlib import Path
+import torch
+from safetensors.torch import load_file
+from tokenizers import Tokenizer
+from balanced_model import BalancedLM
+
+def load_model(directory,device='cpu'):
+ path=Path(directory);m=BalancedLM(json.loads((path/'config.json').read_text()))
+ state=load_file(str(path/'model.safetensors'));state['lm_head.weight']=state['model.embed_tokens.weight']
+ m.load_state_dict(state,strict=True);return m.eval().to(device),Tokenizer.from_file(str(path/'tokenizer.json'))

reports/benchmark/result.json

@@ -0,0 +1,7 @@
+{
+  "tokens_per_second": 999441.1017687789,
+  "timed_updates": 20,
+  "seconds": 10.491623749956489,
+  "parameters": 149333081,
+  "note": "Synthetic systems benchmark, not training or quality evidence"
+}

reports/deadline.json

@@ -0,0 +1 @@
+{"ready_by_utc":"2026-10-01T13:00:00+00:00","ready_by_warsaw":"2026-10-01 15:00","free_reservation_confirmed_by_user":true,"reserve_seconds_for_evaluation_and_upload":4200}

reports/sources/finemath/manifest.json

@@ -0,0 +1,20 @@
+{
+  "repo": "HuggingFaceTB/finemath",
+  "revision": "e92b25a616738fe95dc186b64dfb19f9c8525594",
+  "license": "odc-by",
+  "used_files": [
+    {
+      "path": "finemath-4plus/train-00045-of-00064.parquet",
+      "sha256": "2c3e25eccc80c71af05c1c27a260fb25d5743169dbace5e3ecdd18e6b96c277e"
+    }
+  ],
+  "tokens": 100117355,
+  "validation_tokens": 503280,
+  "docs": 66751,
+  "overlap_rejected": 23,
+  "sha256": "f03351e6b661140e6fb0db3afb8d83e72a668b76a4d8715608d24b0ec85a26b4",
+  "validation_sha256": "180832b085910a6a2b31ae42e42118edd380b7007c23268d60f6d69f89e01688",
+  "tokenizer_sha256": "d1ffe47f6fd8cc19ba7c3e5d2537699d6a99d09b584d3319944f4e40fe6c44ea",
+  "split": "document hash 1% held out",
+  "decontamination": "normalized 13-word exact overlap and exact short sentences; not semantic guarantee"
+}

reports/status.json

@@ -0,0 +1 @@
+{"stage":"preparing_corpus_and_benchmarking","initialization":"from_scratch","parameters":149333081,"vocab_size":24576,"ready_by_warsaw":"2026-10-01 15:00","leaderboard_result":"not_measured"}

reports/tokenizer-selection.json

@@ -0,0 +1,152 @@
+{
+  "selected": {
+    "vocab": 24576,
+    "ffn": 2496,
+    "parameters": 149333081,
+    "weighted_tokens_per_byte": 0.21982151683772885,
+    "measurements": {
+      "fineweb_edu": {
+        "tokens": 108258,
+        "bytes": 500330,
+        "bytes_per_token": 4.621644589776275
+      },
+      "dclm_edu": {
+        "tokens": 113785,
+        "bytes": 502955,
+        "bytes_per_token": 4.42022234916729
+      },
+      "cosmopedia": {
+        "tokens": 96667,
+        "bytes": 500651,
+        "bytes_per_token": 5.179130416791666
+      },
+      "finemath": {
+        "tokens": 149111,
+        "bytes": 509813,
+        "bytes_per_token": 3.4190167056756375
+      },
+      "edge_cases": {
+        "tokens": 67,
+        "bytes": 130,
+        "bytes_per_token": 1.9402985074626866
+      }
+    },
+    "fit_seconds": 6.842154794954695
+  },
+  "candidates": [
+    {
+      "vocab": 16384,
+      "ffn": 2592,
+      "parameters": 148927577,
+      "weighted_tokens_per_byte": 0.2296576476963851,
+      "measurements": {
+        "fineweb_edu": {
+          "tokens": 113155,
+          "bytes": 500330,
+          "bytes_per_token": 4.421634041801069
+        },
+        "dclm_edu": {
+          "tokens": 118553,
+          "bytes": 502955,
+          "bytes_per_token": 4.242448525132219
+        },
+        "cosmopedia": {
+          "tokens": 101369,
+          "bytes": 500651,
+          "bytes_per_token": 4.938896506821612
+        },
+        "finemath": {
+          "tokens": 155627,
+          "bytes": 509813,
+          "bytes_per_token": 3.275864727842855
+        },
+        "edge_cases": {
+          "tokens": 72,
+          "bytes": 130,
+          "bytes_per_token": 1.8055555555555556
+        }
+      },
+      "fit_seconds": 9.067338421940804
+    },
+    {
+      "vocab": 24576,
+      "ffn": 2496,
+      "parameters": 149333081,
+      "weighted_tokens_per_byte": 0.21982151683772885,
+      "measurements": {
+        "fineweb_edu": {
+          "tokens": 108258,
+          "bytes": 500330,
+          "bytes_per_token": 4.621644589776275
+        },
+        "dclm_edu": {
+          "tokens": 113785,
+          "bytes": 502955,
+          "bytes_per_token": 4.42022234916729
+        },
+        "cosmopedia": {
+          "tokens": 96667,
+          "bytes": 500651,
+          "bytes_per_token": 5.179130416791666
+        },
+        "finemath": {
+          "tokens": 149111,
+          "bytes": 509813,
+          "bytes_per_token": 3.4190167056756375
+        },
+        "edge_cases": {
+          "tokens": 67,
+          "bytes": 130,
+          "bytes_per_token": 1.9402985074626866
+        }
+      },
+      "fit_seconds": 6.842154794954695
+    },
+    {
+      "vocab": 32768,
+      "ffn": 2384,
+      "parameters": 149019737,
+      "weighted_tokens_per_byte": 0.2144127319251325,
+      "measurements": {
+        "fineweb_edu": {
+          "tokens": 105600,
+          "bytes": 500330,
+          "bytes_per_token": 4.737973484848485
+        },
+        "dclm_edu": {
+          "tokens": 111026,
+          "bytes": 502955,
+          "bytes_per_token": 4.530065029812836
+        },
+        "cosmopedia": {
+          "tokens": 94072,
+          "bytes": 500651,
+          "bytes_per_token": 5.321998044051365
+        },
+        "finemath": {
+          "tokens": 145642,
+          "bytes": 509813,
+          "bytes_per_token": 3.50045316598234
+        },
+        "edge_cases": {
+          "tokens": 63,
+          "bytes": 130,
+          "bytes_per_token": 2.0634920634920637
+        }
+      },
+      "fit_seconds": 7.331924270954914
+    }
+  ],
+  "criterion": "Smallest vocabulary within 3% of best weighted held-out token compression; provisional pending speed check.",
+  "tokenizer_training_sha256": "c15bad040cf3eaf9e0053539bff3039601f9a4db7e89a0acd6db0e03e4cced49",
+  "tokenizer_sha256": "d1ffe47f6fd8cc19ba7c3e5d2537699d6a99d09b584d3319944f4e40fe6c44ea",
+  "normalization": "none; byte fallback alphabet; preserve Unicode and whitespace; digit groups max3",
+  "parameters_cap": 149500000,
+  "data_mix": {
+    "fineweb_edu": 0.65,
+    "dclm_edu": 0.2,
+    "cosmopedia": 0.1,
+    "finemath": 0.05
+  },
+  "status": "roundtrip_passed"
+}

reports/train.toml

@@ -0,0 +1,11 @@
+seed = 1337
+block = 2048
+micro_batch = 16
+accumulation = 2
+adam_lr = 0.0015
+muon_lr = 0.02
+weight_decay = 0.1
+clip = 1.0
+checkpoint_every = 500
+log_every = 20
+compile = true