Upload folder using huggingface_hub

Dockerfile

@@ -36,10 +36,10 @@ COPY --chown=user . $HOME/app
 RUN cd neural_engine && cargo build --release
 
 # Download all model weights from our model repository directly into the root folder
-RUN curl -L -o ./model.onnx "https://huggingface.co/dpv007/Neurex-Weights/resolve/main/model.onnx?v=4" || true
-RUN curl -L -o ./model.onnx.data "https://huggingface.co/dpv007/Neurex-Weights/resolve/main/model.onnx.data?v=4" || true
-RUN curl -L -o ./model_int8.onnx "https://huggingface.co/dpv007/Neurex-Weights/resolve/main/model_int8.onnx?v=4" || true
-RUN curl -L -o ./model_int8.onnx.data "https://huggingface.co/dpv007/Neurex-Weights/resolve/main/model_int8.onnx.data?v=4" || true
+RUN curl -L -o ./model.onnx "https://huggingface.co/dpv007/Neurex-Weights/resolve/main/model.onnx?v=5" || true
+RUN curl -L -o ./model.onnx.data "https://huggingface.co/dpv007/Neurex-Weights/resolve/main/model.onnx.data?v=5" || true
+RUN curl -L -o ./model_int8.onnx "https://huggingface.co/dpv007/Neurex-Weights/resolve/main/model_int8.onnx?v=5" || true
+RUN curl -L -o ./model_int8.onnx.data "https://huggingface.co/dpv007/Neurex-Weights/resolve/main/model_int8.onnx.data?v=5" || true
 
 # Prepare the Neurex_Engine folder
 RUN mkdir -p Neurex_Engine && \

config.yaml

@@ -0,0 +1,28 @@
+data:
+  url: "https://database.lichess.org/standard/lichess_db_standard_rated_2024-04.pgn.zst"
+  chunk_file: "dataset/stream_chunk.pgn"
+  games_per_chunk: 5000
+  max_length: 120
+  val_split: 0.1
+
+model:
+  d_model: 512
+  nhead: 8
+  num_layers: 6
+
+training:
+  epochs: 10
+  chunk_size: 50000
+  epochs_per_chunk: 2
+  batch_size: 32
+  max_lr: 0.0005
+  min_lr: 0.00001
+  warmup_steps: 50000      # 10% of max_steps
+  max_steps: 500000
+  val_every_n_steps: 100   # Checkpoint frequently so you never lose progress
+  patience: 50
+  scheduler_t0: 100000     # T0 aligned with max_steps
+  grad_accum_steps: 2
+  log_file: "dataset/fast_training_log.txt"
+  weights_path: "weights/chess_fast_best.pth"
+  latest_weights_path: "weights/chess_fast_latest.pth"

model.py

@@ -0,0 +1,116 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_seq_len=512):
+        super().__init__()
+        self.dim = dim
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        t = torch.arange(max_seq_len, dtype=torch.float)
+        freqs = torch.outer(t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos(), persistent=False)
+        self.register_buffer("sin_cached", emb.sin(), persistent=False)
+
+    def _rotate_half(self, x):
+        return torch.cat((-x[..., self.dim // 2:], x[..., :self.dim // 2]), dim=-1)
+
+    def forward(self, x, seq_len, offset=0):
+        cos = self.cos_cached[offset:offset+seq_len, :].unsqueeze(0).unsqueeze(0)
+        sin = self.sin_cached[offset:offset+seq_len, :].unsqueeze(0).unsqueeze(0)
+        return (x * cos) + (self._rotate_half(x) * sin)
+
+class SwiGLU(nn.Module):
+    def __init__(self, d_model, d_ff):
+        super().__init__()
+        self.w1 = nn.Linear(d_model, d_ff, bias=False)
+        self.w2 = nn.Linear(d_model, d_ff, bias=False)
+        self.w3 = nn.Linear(d_ff, d_model, bias=False)
+
+    def forward(self, x):
+        return self.w3(F.silu(self.w1(x)) * self.w2(x))
+
+class OptimizedAttention(nn.Module):
+    def __init__(self, d_model, nhead, rope, num_kv_heads=2):
+        super().__init__()
+        self.nhead, self.head_dim = nhead, d_model // nhead
+        self.num_kv_heads = num_kv_heads
+        self.num_queries_per_kv = nhead // num_kv_heads
+        self.rope = rope
+        
+        self.q_proj = nn.Linear(d_model, d_model, bias=False)
+        self.k_proj = nn.Linear(d_model, num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(d_model, num_kv_heads * self.head_dim, bias=False)
+        self.out_proj = nn.Linear(d_model, d_model, bias=False)
+
+    def forward(self, x, past_key_value=None):
+        B, S, C = x.shape
+        q = self.q_proj(x).view(B, S, self.nhead, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        
+        offset = 0 if past_key_value is None else past_key_value[0].size(2)
+        q = self.rope(q, S, offset)
+        k = self.rope(k, S, offset)
+        
+        if past_key_value is not None:
+            past_k, past_v = past_key_value
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+            
+        present_key_value = (k, v)
+        
+        if self.num_queries_per_kv > 1:
+            k = torch.repeat_interleave(k, repeats=self.num_queries_per_kv, dim=1)
+            v = torch.repeat_interleave(v, repeats=self.num_queries_per_kv, dim=1)
+            
+        attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        out = self.out_proj(attn_out.transpose(1, 2).contiguous().view(B, S, C))
+        return out, present_key_value
+
+class EfficientDecoderLayer(nn.Module):
+    def __init__(self, d_model, nhead, d_ff, rope, num_kv_heads=2):
+        super().__init__()
+        self.attn = OptimizedAttention(d_model, nhead, rope, num_kv_heads)
+        self.ffn = SwiGLU(d_model, d_ff)
+        self.ln1, self.ln2 = nn.LayerNorm(d_model), nn.LayerNorm(d_model)
+
+    def forward(self, x, past_key_value=None):
+        attn_out, present_key_value = self.attn(self.ln1(x), past_key_value)
+        x = x + attn_out
+        return x + self.ffn(self.ln2(x)), present_key_value
+
+class ChessTransformer(nn.Module):
+    def __init__(self, vocab_size, d_model=512, nhead=8, num_layers=6, max_length=120, num_kv_heads=2):
+        super().__init__()
+        self.token_embedding = nn.Embedding(vocab_size, d_model)
+        d_ff = int(2 * (d_model * 4) / 3)
+        self.rope = RotaryEmbedding(dim=d_model // nhead, max_seq_len=max_length)
+        self.layers = nn.ModuleList([EfficientDecoderLayer(d_model, nhead, d_ff, self.rope, num_kv_heads) for _ in range(num_layers)])
+        self.ln_final = nn.LayerNorm(d_model)
+        self.fc_out = nn.Linear(d_model, vocab_size, bias=False)
+        self.value_head = nn.Sequential(
+            nn.Linear(d_model, d_model // 2),
+            nn.SiLU(),
+            nn.Linear(d_model // 2, 1),
+            nn.Tanh()
+        )
+        
+    def forward(self, src, past_key_values=None, use_cache=False):
+        x = self.token_embedding(src)
+        present_key_values = []
+        for i, layer in enumerate(self.layers):
+            past_kv = past_key_values[i] if past_key_values is not None else None
+            x, present_kv = layer(x, past_kv)
+            if use_cache:
+                present_key_values.append(present_kv)
+            
+        hidden = self.ln_final(x)
+        policy_logits = self.fc_out(hidden)
+        value = self.value_head(hidden).squeeze(-1)
+        
+        if use_cache:
+            return policy_logits, value, tuple(present_key_values)
+        return policy_logits, value

neural_engine/src/main.rs

@@ -290,18 +290,18 @@ fn evaluate_onnx(
     let cache = if let Some(c) = root_cache { c } else { &empty_cache };
     let past_seq_len = cache.seq_len;
     
-    let pk_0 = ort::value::Tensor::from_array((vec![1, 8, past_seq_len, 64], cache.layers[0].0.clone())).unwrap();
-    let pv_0 = ort::value::Tensor::from_array((vec![1, 8, past_seq_len, 64], cache.layers[0].1.clone())).unwrap();
-    let pk_1 = ort::value::Tensor::from_array((vec![1, 8, past_seq_len, 64], cache.layers[1].0.clone())).unwrap();
-    let pv_1 = ort::value::Tensor::from_array((vec![1, 8, past_seq_len, 64], cache.layers[1].1.clone())).unwrap();
-    let pk_2 = ort::value::Tensor::from_array((vec![1, 8, past_seq_len, 64], cache.layers[2].0.clone())).unwrap();
-    let pv_2 = ort::value::Tensor::from_array((vec![1, 8, past_seq_len, 64], cache.layers[2].1.clone())).unwrap();
-    let pk_3 = ort::value::Tensor::from_array((vec![1, 8, past_seq_len, 64], cache.layers[3].0.clone())).unwrap();
-    let pv_3 = ort::value::Tensor::from_array((vec![1, 8, past_seq_len, 64], cache.layers[3].1.clone())).unwrap();
-    let pk_4 = ort::value::Tensor::from_array((vec![1, 8, past_seq_len, 64], cache.layers[4].0.clone())).unwrap();
-    let pv_4 = ort::value::Tensor::from_array((vec![1, 8, past_seq_len, 64], cache.layers[4].1.clone())).unwrap();
-    let pk_5 = ort::value::Tensor::from_array((vec![1, 8, past_seq_len, 64], cache.layers[5].0.clone())).unwrap();
-    let pv_5 = ort::value::Tensor::from_array((vec![1, 8, past_seq_len, 64], cache.layers[5].1.clone())).unwrap();
+    let pk_0 = ort::value::Tensor::from_array((vec![1, 2, past_seq_len, 64], cache.layers[0].0.clone())).unwrap();
+    let pv_0 = ort::value::Tensor::from_array((vec![1, 2, past_seq_len, 64], cache.layers[0].1.clone())).unwrap();
+    let pk_1 = ort::value::Tensor::from_array((vec![1, 2, past_seq_len, 64], cache.layers[1].0.clone())).unwrap();
+    let pv_1 = ort::value::Tensor::from_array((vec![1, 2, past_seq_len, 64], cache.layers[1].1.clone())).unwrap();
+    let pk_2 = ort::value::Tensor::from_array((vec![1, 2, past_seq_len, 64], cache.layers[2].0.clone())).unwrap();
+    let pv_2 = ort::value::Tensor::from_array((vec![1, 2, past_seq_len, 64], cache.layers[2].1.clone())).unwrap();
+    let pk_3 = ort::value::Tensor::from_array((vec![1, 2, past_seq_len, 64], cache.layers[3].0.clone())).unwrap();
+    let pv_3 = ort::value::Tensor::from_array((vec![1, 2, past_seq_len, 64], cache.layers[3].1.clone())).unwrap();
+    let pk_4 = ort::value::Tensor::from_array((vec![1, 2, past_seq_len, 64], cache.layers[4].0.clone())).unwrap();
+    let pv_4 = ort::value::Tensor::from_array((vec![1, 2, past_seq_len, 64], cache.layers[4].1.clone())).unwrap();
+    let pk_5 = ort::value::Tensor::from_array((vec![1, 2, past_seq_len, 64], cache.layers[5].0.clone())).unwrap();
+    let pv_5 = ort::value::Tensor::from_array((vec![1, 2, past_seq_len, 64], cache.layers[5].1.clone())).unwrap();
 
     let outputs = session.run(ort::inputs![
         "input_ids" => input_value,
@@ -356,18 +356,18 @@ fn compute_root_kv_cache(
     let seq_len = seq.len();
     let input_value = ort::value::Tensor::from_array((vec![1, seq_len], seq)).unwrap();
     
-    let pk_0 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 8, 0, 64))).unwrap();
-    let pv_0 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 8, 0, 64))).unwrap();
-    let pk_1 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 8, 0, 64))).unwrap();
-    let pv_1 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 8, 0, 64))).unwrap();
-    let pk_2 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 8, 0, 64))).unwrap();
-    let pv_2 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 8, 0, 64))).unwrap();
-    let pk_3 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 8, 0, 64))).unwrap();
-    let pv_3 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 8, 0, 64))).unwrap();
-    let pk_4 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 8, 0, 64))).unwrap();
-    let pv_4 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 8, 0, 64))).unwrap();
-    let pk_5 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 8, 0, 64))).unwrap();
-    let pv_5 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 8, 0, 64))).unwrap();
+    let pk_0 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 2, 0, 64))).unwrap();
+    let pv_0 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 2, 0, 64))).unwrap();
+    let pk_1 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 2, 0, 64))).unwrap();
+    let pv_1 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 2, 0, 64))).unwrap();
+    let pk_2 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 2, 0, 64))).unwrap();
+    let pv_2 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 2, 0, 64))).unwrap();
+    let pk_3 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 2, 0, 64))).unwrap();
+    let pv_3 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 2, 0, 64))).unwrap();
+    let pk_4 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 2, 0, 64))).unwrap();
+    let pv_4 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 2, 0, 64))).unwrap();
+    let pk_5 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 2, 0, 64))).unwrap();
+    let pv_5 = ort::value::Tensor::from_array(ndarray::Array4::<f32>::zeros((1, 2, 0, 64))).unwrap();
 
     let outputs = session.run(ort::inputs![
         "input_ids" => input_value,
@@ -694,12 +694,14 @@ fn main() {
                     // Time Advantage Scaling
                     if time_left > 0 && opp_time > 0 {
                         let ratio = time_left as f64 / opp_time as f64;
-                        if ratio > 1.5 {
-                            // We have >50% more time than opponent, spend 30% more time to press the advantage
-                            movetime_ms = (movetime_ms as f64 * 1.3) as u64;
-                        } else if ratio < 0.6 {
-                            // We are behind on time, play faster to catch up
-                            movetime_ms = (movetime_ms as f64 * 0.85) as u64;
+                        if ratio > 2.0 {
+                            movetime_ms = (movetime_ms as f64 * 1.5) as u64;
+                        } else if ratio > 1.2 {
+                            movetime_ms = (movetime_ms as f64 * 1.25) as u64;
+                        } else if ratio < 0.8 {
+                            movetime_ms = (movetime_ms as f64 * 0.9) as u64;
+                        } else if ratio < 0.5 {
+                            movetime_ms = (movetime_ms as f64 * 0.75) as u64;
                         }
                     }
                     

rl_train.py

@@ -284,7 +284,13 @@ def learner_worker():
             log_prob = F.log_softmax(p_logits, dim=-1)
             action_log_probs = log_prob[torch.arange(batch_size), a_batch]
             
-            policy_loss = -(action_log_probs * adv_batch).mean()
+            # CRITICAL FIX: Only train on positive advantages (Advantage-Weighted Behavioral Cloning).
+            # If we allow negative advantages, the optimizer pushes log_prob to negative infinity,
+            # causing the policy_loss to explode into massive negative numbers (e.g. -59.7) and
+            # completely destroying the Neural Network's weights (including the Value Head).
+            positive_adv = torch.clamp(adv_batch, min=0.0)
+            policy_loss = -(action_log_probs * positive_adv).mean()
+            
             value_loss = F.mse_loss(v_pred, sf_val_batch)
             
             loss = policy_loss + 0.5 * value_loss