Update strict-small architecture files (sliding window [64, 16, 8, 4], ln3, ln_post_moe, no res3)

README.md

@@ -0,0 +1,84 @@
+---
+license: cc-by-nc-4.0
+language:
+  - en
+tags:
+  - babylm
+  - babylm-2026
+  - mixture-of-experts
+  - msit
+  - xpertgpt
+  - custom_code
+  - safetensors
+library_name: transformers
+pipeline_tag: text-generation
+---
+
+# XpertGPT (Sliding Window 16, 16, 4, 4)
+
+XpertGPT is a sparse Mixture of Experts (MoE) language model designed for data-efficient pretraining under the **BabyLM 2026 challenge (Strict-Small 10M track)**. It leverages **Parallelized Multi-Scale Information Transmission (MSIT)** and **Expert Choice Routing** to maximize representational capacity within restricted token budgets.
+
+This version implements corrected LayerNorms, redundant residual removal, and **expanded sliding window sizes** across parallel expert channels.
+
+---
+
+## 1. Expert Sliding Window Layout
+
+The four parallel MoE experts are configured with distinct sliding window attention constraints to capture varying context ranges (multi-scale sequence features):
+* **Expert 1**: Window size `16` tokens
+* **Expert 2**: Window size `16` tokens
+* **Expert 3**: Window size `4` tokens
+* **Expert 4**: Window size `4` tokens
+
+---
+
+## 2. Architectural Layout & Changes
+
+This model implements:
+1. **Removal of Redundant Residual (`res3`)**:
+   * Removed redundant residual connection around the global dense block. Gated input is now simply $X_2 = X_1$.
+2. **Introduction of Post-Block LayerNorm (`ln3`)**:
+   * LayerNorm `ln3` is added after the Residual 2 Feed-Forward Network addition inside every `MSITBranchBlock` (global block and all expert blocks).
+   * Formulation: $X_{\text{out}} = \text{LayerNorm}(X^{(2)})$.
+3. **Introduction of Post-MoE LayerNorm (`ln_post_moe`)**:
+   * LayerNorm `ln_post_moe` is added after the Residual 4 MoE aggregation.
+   * Formulation: $X_{\text{out}} = \text{LayerNorm}(X_2 + X_{3, \text{full}})$.
+
+---
+
+## 3. Checkpoint Branch Layout
+
+Checkpoints are saved as separate git branches (revisions) on this repository:
+* **1M to 10M words**: Saved every 1M words (`chck_1M` through `chck_10M`).
+* **10M to 100M words**: Saved every 10M words (`chck_10M` through `chck_100M`).
+* **Final Model**: Saved under the `main` branch.
+
+---
+
+## 4. How to Load and Use Checkpoints (Bypass Retraining)
+
+### A. Loading the Final Model (`main` branch)
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+    "SRJ5035/correct_small_sw_16_16_4_4_norm_residuls_xpert_strcit_small",
+    revision="main",
+    trust_remote_code=True
+).eval()
+
+tokenizer = AutoTokenizer.from_pretrained(
+    "SRJ5035/correct_small_sw_16_16_4_4_norm_residuls_xpert_strcit_small",
+    revision="main"
+)
+```
+
+### B. Loading an Intermediate Milestone (e.g. `chck_5M`)
+```python
+model_5m = AutoModelForCausalLM.from_pretrained(
+    "SRJ5035/correct_small_sw_16_16_4_4_norm_residuls_xpert_strcit_small",
+    revision="chck_5M",
+    trust_remote_code=True
+).eval()
+```

configuration_xpertgpt.py

@@ -0,0 +1,36 @@
+from transformers import PretrainedConfig
+
+class XpertGPTConfig(PretrainedConfig):
+    model_type = "xpertgpt"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 256,
+        d_thin: int = 384,
+        num_layers: int = 6,
+        num_blocks: int = 4,
+        capacity_factor: float = 2.0,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+        self.num_hidden_layers = num_layers
+
+        super().__init__(**kwargs)

generate_plots_and_report.py

@@ -0,0 +1,207 @@
+import os
+import json
+import math
+import torch
+import argparse
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate training reports, learning curves, and router plots.")
+    parser.add_argument("--model-dir", type=str, default="./checkpoints/xpertgpt_fresh", help="Path to checkpoints directory")
+    args = parser.parse_args()
+
+    print("=== 1. PARAMETER COUNT & MODEL METRICS REPORT ===")
+    model_path = os.path.join(args.model_dir, "main")
+    weight_file = os.path.join(model_path, "pytorch_model.bin")
+    
+    if os.path.exists(weight_file):
+        try:
+            state_dict = torch.load(weight_file, map_location="cpu")
+            total_params = 0
+            non_emb_params = 0
+            for name, param in state_dict.items():
+                size = param.numel()
+                total_params += size
+                # Exclude embedding and head tie
+                if "wte.weight" not in name and "lm_head.weight" not in name:
+                    non_emb_params += size
+            print(f"Model weights loaded successfully from: {weight_file}")
+            print(f"- Total Parameters (with tied embeddings): {total_params:,}")
+            print(f"- Non-Embedding Parameters: {non_emb_params:,}")
+        except Exception as e:
+            print(f"Warning: Could not read parameter counts from weight file: {e}")
+            # Fallback to config math
+            fallback_param_math()
+    else:
+        print(f"Info: Checked {weight_file} but weights not found. Using fallback config computation:")
+        fallback_param_math()
+
+    # Load validation metrics
+    val_file = "validation_metrics.json"
+    if not os.path.exists(val_file):
+        val_file = os.path.join(args.model_dir, "validation_metrics.json")
+        
+    if os.path.exists(val_file):
+        with open(val_file, "r") as f:
+            val_logs = json.load(f)
+        if len(val_logs) > 0:
+            final_val = val_logs[-1]
+            print(f"- Final Validation Step: {final_val['step']}")
+            print(f"- Final Validation Loss: {final_val['val_loss']:.4f}")
+            print(f"- Final Validation Perplexity: {final_val['val_perplexity']:.2f}")
+    else:
+        print(f"Warning: {val_file} not found! Cannot report validation perplexity.")
+
+    print("\n=== 2. GENERATING LEARNING CURVES ===")
+    train_file = "training_metrics.json"
+    if not os.path.exists(train_file):
+        train_file = os.path.join(args.model_dir, "training_metrics.json")
+
+    try:
+        import matplotlib
+        matplotlib.use('Agg')
+        import matplotlib.pyplot as plt
+        
+        has_plots = False
+        
+        # Setup plot grid
+        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
+        
+        if os.path.exists(train_file):
+            with open(train_file, "r") as f:
+                train_logs = json.load(f)
+            steps = [log["step"] for log in train_logs]
+            losses = [log["train_loss"] for log in train_logs]
+            ax1.plot(steps, losses, color="#ef4444", linewidth=1.5, label="Train Loss")
+            ax1.set_xlabel("Global Step")
+            ax1.set_ylabel("Cross Entropy Loss")
+            ax1.set_title("Training Loss Curve")
+            ax1.grid(True, linestyle="--", alpha=0.6)
+            ax1.legend()
+            has_plots = True
+        else:
+            ax1.text(0.5, 0.5, "training_metrics.json not found", ha="center")
+            
+        if os.path.exists(val_file):
+            with open(val_file, "r") as f:
+                val_logs = json.load(f)
+            steps = [log["step"] for log in val_logs]
+            ppl = [log["val_perplexity"] for log in val_logs]
+            ax2.plot(steps, ppl, marker="o", color="#3b82f6", linewidth=2, label="Val Perplexity")
+            ax2.set_xlabel("Global Step")
+            ax2.set_ylabel("Perplexity")
+            ax2.set_title("Validation Perplexity Curve")
+            ax2.grid(True, linestyle="--", alpha=0.6)
+            ax2.legend()
+            has_plots = True
+        else:
+            ax2.text(0.5, 0.5, "validation_metrics.json not found", ha="center")
+
+        if has_plots:
+            plt.tight_layout()
+            plt.savefig("learning_curves.png", dpi=150)
+            plt.close()
+            print("Successfully saved learning curve diagrams to 'learning_curves.png'.")
+    except Exception as e:
+        print(f"Warning: Could not plot learning curves: {e}")
+
+    print("\n=== 3. ROUTER UTILIZATION ANALYSIS ===")
+    router_details_file = "router_details.json"
+    if not os.path.exists(router_details_file):
+        router_details_file = os.path.join(args.model_dir, "router_details.json")
+        
+    if os.path.exists(router_details_file):
+        try:
+            with open(router_details_file, "r") as f:
+                records = json.load(f)
+                
+            print(f"Loaded router mapping records containing {len(records)} logged steps.")
+            
+            # Reconstruct selection count mapping
+            # records format: [{"step": s, "layer": l, "expert_assignments": [[tokens for exp1], [exp2], ...]}]
+            # Find the total number of unique token indices in each step
+            total_expert_counts = [0, 0, 0, 0, 0] # indices representing 0, 1, 2, 3, 4 experts
+            
+            for record in records:
+                assignments = record["expert_assignments"]
+                # Number of experts is 4
+                num_blocks = len(assignments)
+                # Find maximum token index present in assignments to know the token range
+                max_token_idx = 0
+                for exp_list in assignments:
+                    if len(exp_list) > 0:
+                        max_token_idx = max(max_token_idx, max(exp_list))
+                
+                n_tokens = max_token_idx + 1
+                token_counts = [0] * n_tokens
+                for exp_list in assignments:
+                    for token_idx in exp_list:
+                        token_counts[token_idx] += 1
+                        
+                for count in token_counts:
+                    if count < 5:
+                        total_expert_counts[count] += 1
+                        
+            total_tokens_sum = sum(total_expert_counts)
+            if total_tokens_sum > 0:
+                fractions = [total_expert_counts[i] / total_tokens_sum for i in range(5)]
+            else:
+                fractions = [0.0] * 5
+                
+            fraction_no_expert = fractions[0]
+            fraction_multiple_experts = sum(fractions[2:])
+            
+            report_text = f"""=== ROUTER UTILIZATION REPORT ===
+Total tokens tracked: {total_tokens_sum:,}
+
+Number of experts selected per token:
+- 0 experts: {total_expert_counts[0]:,} tokens ({fractions[0]:.2%})
+- 1 expert:  {total_expert_counts[1]:,} tokens ({fractions[1]:.2%})
+- 2 experts: {total_expert_counts[2]:,} tokens ({fractions[2]:.2%})
+- 3 experts: {total_expert_counts[3]:,} tokens ({fractions[3]:.2%})
+- 4 experts: {total_expert_counts[4]:,} tokens ({fractions[4]:.2%})
+
+Summary metrics:
+- Fraction of tokens receiving NO expert: {fraction_no_expert:.2%}
+- Fraction of tokens receiving MULTIPLE experts (2, 3, or 4): {fraction_multiple_experts:.2%}
+"""
+            with open("router_report.txt", "w") as f:
+                f.write(report_text)
+            print(report_text)
+            
+            # Plot utilization chart
+            try:
+                import matplotlib
+                matplotlib.use('Agg')
+                import matplotlib.pyplot as plt
+                
+                categories = ['0', '1', '2', '3', '4']
+                plt.figure(figsize=(8, 5))
+                plt.bar(categories, [f * 100 for f in fractions], color='#10b981', edgecolor='#059669', width=0.6)
+                plt.xlabel('Number of Experts Selected per Token')
+                plt.ylabel('Percentage of Tokens (%)')
+                plt.title('Router Utilization / Expert Choices per Token')
+                plt.grid(axis='y', linestyle='--', alpha=0.6)
+                
+                for i, f in enumerate(fractions):
+                    plt.text(i, f * 100 + 1, f"{f:.2%}", ha='center', fontweight='bold')
+                    
+                plt.ylim(0, max([f * 100 for f in fractions]) + 10)
+                plt.savefig('router_utilization.png', dpi=150)
+                plt.close()
+                print("Successfully saved router utilization plot to 'router_utilization.png'.")
+            except Exception as plot_err:
+                print(f"Warning: Could not plot router utilization: {plot_err}")
+                
+        except Exception as e:
+            print(f"Warning: Could not analyze router assignments file: {e}")
+    else:
+        print(f"Warning: {router_details_file} not found! Cannot analyze router utilization.")
+
+def fallback_param_math():
+    # Fallback mathematical projection computation based on model config
+    print("Fallback configuration architecture math:")
+    print("- Config: 6 layers, d_model=256, d_thin=384, 4 experts, capacity_factor=2.0")
+    print("- Non-Embedding parameter count is approximately ~48.4M parameters.")
+
+if __name__ == "__main__":
+    main()

modeling_xpertgpt.py

@@ -0,0 +1,460 @@
+import os
+import math
+import random
+import inspect
+from typing import Optional, Tuple, Dict, Any
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PretrainedConfig, PreTrainedModel, GenerationMixin, AutoConfig, AutoModel, AutoModelForCausalLM
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+# Global router tracker for validation & learning curve analysis
+ROUTER_TRACKER = {
+    "enabled": False,
+    "expert_counts": []  # count of experts selected per token [0, 1, 2, 3, 4]
+}
+
+# ─────────────────────────────────────────────────────────────
+# Configuration Classes
+# ─────────────────────────────────────────────────────────────
+
+@dataclass
+class XpertGPTModelConfig:
+    vocab_size: int = 16384
+    block_size: int = 512
+    d_model: int = 256
+    d_thin: int = 384
+    num_layers: int = 6
+    num_blocks: int = 4
+    capacity_factor: float = 2.0
+    dropout: float = 0.1
+
+class XpertGPTConfig(PretrainedConfig):
+    model_type = "xpertgpt"
+    auto_map = {
+        "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
+        "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
+        "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        block_size: int = 512,
+        d_model: int = 256,
+        d_thin: int = 384,
+        num_layers: int = 6,
+        num_blocks: int = 4,
+        capacity_factor: float = 2.0,
+        dropout: float = 0.1,
+        **kwargs
+    ):
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("bos_token_id", 2)  # [CLS]
+        kwargs.setdefault("eos_token_id", 3)  # [SEP]
+        kwargs.setdefault("pad_token_id", 1)  # [PAD]
+
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_layers = num_layers
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+        self.dropout = dropout
+        
+        # Attribute parity for classification heads
+        self.hidden_size = d_model
+        self.num_hidden_layers = num_layers
+
+        super().__init__(**kwargs)
+
+# ─────────────────────────────────────────────────────────────
+# ROPE HELPERS
+# ─────────────────────────────────────────────────────────────
+
+def _precompute_rope_freqs(head_dim: int, seq_len: int, device: torch.device, theta: float = 10000.0):
+    assert head_dim % 2 == 0, "head_dim must be divisible by 2 for RoPE"
+    inv_freq = 1.0 / (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)
+    emb = torch.cat((freqs, freqs), dim=-1)
+    return emb.cos(), emb.sin()
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    L = x.size(2)
+    cos = cos[:L, :].unsqueeze(0).unsqueeze(1) 
+    sin = sin[:L, :].unsqueeze(0).unsqueeze(1) 
+    
+    half_dim = x.size(-1) // 2
+    x1 = x[..., :half_dim]
+    x2 = x[..., half_dim:]
+    rotated_x = torch.cat((-x2, x1), dim=-1)
+    
+    return (x * cos) + (rotated_x * sin)
+
+# ─────────────────────────────────────────────────────────────
+# 1.  SLIDING WINDOW ATTENTION
+# ─────────────────────────────────────────────────────────────
+
+class SlidingWindowAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size=None):
+        super().__init__()
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+        self.num_heads  = num_heads
+        self.window_size = window_size
+        self.head_dim   = dim // num_heads
+
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+        self.o_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        B, L, D = x.size()
+
+        q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            past_len = past_k.size(2)
+            q_cos, q_sin = _precompute_rope_freqs(self.head_dim, past_len + L, x.device)
+            q = _apply_rope(q, q_cos[past_len:, :], q_sin[past_len:, :])
+            k = _apply_rope(k, q_cos[past_len:, :], q_sin[past_len:, :])
+            
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+        else:
+            cos, sin = _precompute_rope_freqs(self.head_dim, L, x.device)
+            q = _apply_rope(q, cos, sin)
+            k = _apply_rope(k, cos, sin)
+
+        L_kv = k.size(2)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+
+        if bidirectional:
+            if self.window_size is not None:
+                past_len = L_kv - L
+                pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+                pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+                dist  = torch.abs(pos_i - pos_j)
+                win_mask = dist < self.window_size
+                scores = scores.masked_fill(
+                    ~win_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+                )
+        else:
+            past_len = L_kv - L
+            pos_i = (past_len + torch.arange(L, device=x.device)).unsqueeze(1)
+            pos_j = torch.arange(L_kv, device=x.device).unsqueeze(0)
+            dist  = pos_i - pos_j
+
+            causal_mask = dist >= 0
+            if self.window_size is not None:
+                causal_mask = causal_mask & (dist < self.window_size)
+
+            scores = scores.masked_fill(
+                ~causal_mask.unsqueeze(0).unsqueeze(0), float('-inf')
+            )
+
+        attn = torch.softmax(scores, dim=-1)
+        out = torch.matmul(attn, v)                                
+        out = out.transpose(1, 2).contiguous().view(B, L, D)
+        out = self.o_proj(out)
+
+        if use_cache:
+            if self.window_size is not None:
+                present_kv = (
+                    k[:, :, -self.window_size:, :],
+                    v[:, :, -self.window_size:, :]
+                )
+            else:
+                present_kv = (k, v)
+        else:
+            present_kv = None
+
+        return out, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 2.  MSIT BRANCH BLOCK (Pre-Norm)
+# ─────────────────────────────────────────────────────────────
+
+class MSITBranchBlock(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window_size):
+        super().__init__()
+        self.ln1  = nn.LayerNorm(dim)
+        self.attn = SlidingWindowAttention(dim, num_heads, window_size)
+        self.ln2  = nn.LayerNorm(dim)
+        self.ffn  = nn.Sequential(
+            nn.Linear(dim, dim * 4, bias=False),
+            nn.GELU(),
+            nn.Linear(dim * 4, dim, bias=False),
+        )
+        self.ln3  = nn.LayerNorm(dim)
+
+    def forward(self, x: torch.Tensor,
+                past_kv=None,
+                use_cache: bool = False,
+                bidirectional: bool = False):
+        attn_out, present_kv = self.attn(
+            self.ln1(x), past_kv, use_cache, bidirectional
+        )
+        x = x + attn_out
+        x = x + self.ffn(self.ln2(x))
+        x = self.ln3(x)
+        return x, present_kv
+
+# ─────────────────────────────────────────────────────────────
+# 3.  MoEP-MSIT ARCHITECTURE BLOCK  (Expert Choice routing)
+# ─────────────────────────────────────────────────────────────
+
+class MoEPMSITBlock(nn.Module):
+    def __init__(self, d_model: int = 512, d_thin: int = 192, num_blocks: int = 14,
+                 capacity_factor: float = 2.0):
+        super().__init__()
+        self.d_model = d_model
+        self.d_thin = d_thin
+        self.num_blocks = num_blocks
+        self.capacity_factor = capacity_factor
+
+        # 1. Global Block (Dense, d_model)
+        num_heads_global = max(1, d_model // 64)
+        self.global_block = MSITBranchBlock(d_model, num_heads_global, window_size=None)
+
+        # 2. Router
+        self.router_ln = nn.LayerNorm(d_model)
+        self.w_router = nn.Linear(d_model, num_blocks, bias=False)
+
+        # 3. Shrink Projection
+        self.w_down = nn.Linear(d_model, d_thin, bias=False)
+
+        # 4. Thin Parallel Blocks (d_thin)
+        self.windows = [64, 16, 8, 4] + [None] * (num_blocks - 4)
+        self.heads = [max(1, d_thin // 64)] * num_blocks  
+        
+        self.thin_blocks = nn.ModuleList([
+            MSITBranchBlock(d_thin, self.heads[i], self.windows[i])
+            for i in range(num_blocks)
+        ])
+
+        # 6. Grow Projection
+        self.w_up = nn.Linear(d_thin, d_model, bias=False)
+        self.last_topk_indices = None
+        self.ln_post_moe = nn.LayerNorm(d_model)
+
+    def forward(self, x_0: torch.Tensor, past_kvs=None, use_cache: bool = False, bidirectional: bool = False):
+        B, T, D = x_0.size()
+        n_tokens = B * T
+
+        # Step 1: Global Block
+        pkv_g = past_kvs[0] if past_kvs else None
+        x_1, nkv_g = self.global_block(x_0, pkv_g, use_cache, bidirectional)
+
+        # Step 2: Gated input stream
+        x_2 = x_1
+
+        # Router scores
+        r_logits = self.w_router(self.router_ln(x_2)).view(n_tokens, self.num_blocks)
+        r_probs = F.softmax(r_logits, dim=-1)
+
+        # Per-expert capacity: k = (n * c) / e
+        k_capacity = max(1, int(round(n_tokens * self.capacity_factor / self.num_blocks)))
+        k_capacity = min(k_capacity, n_tokens)
+
+        # Expert Choice routing: topk over the token axis for each expert
+        expert_token_scores = r_probs.transpose(0, 1) # (num_blocks, n_tokens)
+        topk_scores, topk_token_idx = torch.topk(expert_token_scores, k_capacity, dim=-1)
+        self.last_topk_indices = topk_token_idx
+
+        # Track routing decisions
+        if ROUTER_TRACKER["enabled"]:
+            token_expert_counts = torch.zeros(n_tokens, dtype=torch.long, device=x_2.device)
+            for idx_exp in range(self.num_blocks):
+                token_expert_counts.scatter_add_(0, topk_token_idx[idx_exp], torch.ones_like(topk_token_idx[idx_exp]))
+            counts = torch.bincount(token_expert_counts, minlength=5).cpu().tolist()
+            ROUTER_TRACKER["expert_counts"].append(counts)
+
+        # Load balancing is guaranteed by construction in Expert Choice
+        layer_aux_loss = x_2.new_zeros(())
+
+        # Shrink projection
+        x_2_thin = self.w_down(x_2) # (B, T, d_thin)
+        x_2_thin_flat = x_2_thin.view(n_tokens, self.d_thin)
+
+        # Expert computations
+        new_kvs = [nkv_g]
+        expert_outputs_flat = torch.zeros(n_tokens, self.d_model, device=x_2.device, dtype=x_2.dtype)
+
+        for i, block in enumerate(self.thin_blocks):
+            sel_idx = topk_token_idx[i]
+            bucket_in = x_2_thin_flat[sel_idx].unsqueeze(0) # (1, k_capacity, d_thin)
+
+            pkv_i = past_kvs[i + 1] if past_kvs else None
+            bucket_out, nkv_i = block(bucket_in, pkv_i, use_cache, bidirectional)
+            if use_cache:
+                new_kvs.append(nkv_i)
+
+            bucket_out = bucket_out.squeeze(0) # (k_capacity, d_thin)
+            bucket_out_full = self.w_up(bucket_out) # (k_capacity, d_model)
+
+            gate = topk_scores[i].unsqueeze(-1) # (k_capacity, 1)
+            expert_outputs_flat.index_add_(0, sel_idx, bucket_out_full * gate)
+
+        x_3_full = expert_outputs_flat.view(B, T, self.d_model)
+        out = x_2 + x_3_full
+        out = self.ln_post_moe(out)
+
+        present_kvs = tuple(new_kvs) if use_cache else None
+        return out, present_kvs, layer_aux_loss
+
+# ─────────────────────────────────────────────────────────────
+# 4.  RAW XpertGPT MODEL
+# ─────────────────────────────────────────────────────────────
+
+class XpertGPTModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.wte      = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.drop_emb = nn.Dropout(cfg.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(
+                d_model=cfg.d_model, 
+                d_thin=cfg.d_thin, 
+                num_blocks=cfg.num_blocks, 
+                capacity_factor=cfg.capacity_factor
+            )
+            for _ in range(cfg.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(cfg.d_model)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        self.wte.weight = self.lm_head.weight
+
+    def forward(self, input_ids: torch.Tensor, targets: torch.Tensor = None, bidirectional: bool = False):
+        x = self.drop_emb(self.wte(input_ids))
+        total_aux_loss = 0.0
+        
+        for block in self.blocks:
+            x, _, layer_aux = block(x, past_kvs=None, use_cache=False, bidirectional=bidirectional)
+            total_aux_loss += layer_aux
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if targets is not None:
+            ce_loss = F.cross_entropy(logits.view(-1, self.cfg.vocab_size), targets.view(-1), ignore_index=-100)
+            avg_aux_loss = total_aux_loss / self.cfg.num_layers
+            loss = ce_loss + (0.01 * avg_aux_loss)
+
+        return logits, loss
+
+# ─────────────────────────────────────────────────────────────
+# 5.  HUGGING FACE MODEL WRAPPERS
+# ─────────────────────────────────────────────────────────────
+
+class XpertGPTModelWrapper(PreTrainedModel):
+    config_class = XpertGPTConfig
+    base_model_prefix = "transformer"
+
+    def __init__(self, config: XpertGPTConfig):
+        super().__init__(config)
+        self.wte      = nn.Embedding(config.vocab_size, config.d_model)
+        self.drop_emb = nn.Dropout(config.dropout)
+        self.blocks   = nn.ModuleList([
+            MoEPMSITBlock(config.d_model, config.d_thin, config.num_blocks, config.capacity_factor)
+            for _ in range(config.num_layers)
+        ])
+        self.ln_f    = nn.LayerNorm(config.d_model)
+        self.post_init()
+
+    def forward(self, input_ids, **kwargs):
+        x = self.drop_emb(self.wte(input_ids))
+        for block in self.blocks:
+            x, _, _ = block(x, past_kvs=None, use_cache=False, bidirectional=False)
+        x = self.ln_f(x)
+        return BaseModelOutputWithPast(last_hidden_state=x)
+
+
+class XpertGPTForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = XpertGPTConfig
+    base_model_prefix = "transformer"
+    _no_split_modules = ["MoEPMSITBlock"]
+    _tied_weights_keys = {"transformer.lm_head.weight": "transformer.wte.weight"}
+
+    def __init__(self, config: XpertGPTConfig):
+        super().__init__(config)
+        self.transformer = XpertGPTModelWrapper(config)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.post_init()
+        
+        # State-dict pre-hook for backwards compatibility with checkpoint key naming
+        def _prefix_cleaner(state_dict, prefix, local_metadata, Moore, missing_keys, unexpected_keys, error_msgs):
+            keys = list(state_dict.keys())
+            for k in keys:
+                if k.startswith("transformer."):
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                elif f"{prefix}transformer." in k:
+                    state_dict[k.replace("transformer.", "", 1)] = state_dict.pop(k)
+                    
+        self._register_load_state_dict_pre_hook(_prefix_cleaner)
+
+    def tie_weights(self, **kwargs):
+        if hasattr(self, "transformer") and hasattr(self.transformer, "wte") and hasattr(self.transformer, "lm_head"):
+            self.transformer.wte.weight = self.lm_head.weight
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(self,
+                input_ids: Optional[torch.LongTensor] = None,
+                attention_mask: Optional[torch.FloatTensor] = None,
+                labels: Optional[torch.LongTensor] = None,
+                **kwargs) -> CausalLMOutputWithPast:
+        outputs = self.transformer(input_ids)
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+        
+        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),
+                ignore_index=-100
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+# Register with auto-mapping
+AutoConfig.register("xpertgpt", XpertGPTConfig)
+AutoModel.register(XpertGPTConfig, XpertGPTModelWrapper)
+AutoModelForCausalLM.register(XpertGPTConfig, XpertGPTForCausalLM)

train.py

@@ -0,0 +1,1539 @@
+import os
+import math
+import time
+import json
+import random
+import inspect
+import shutil
+import subprocess
+import argparse
+
+# ─────────────────────────────────────────────────────────────
+# 1. CORE PIPELINE FUNCTION AND HELPERS
+# ─────────────────────────────────────────────────────────────
+
+def is_zero_shot_cache_valid(cache_dir):
+    if not os.path.exists(cache_dir):
+        return False
+    pred_count = 0
+    for root, dirs, files in os.walk(cache_dir):
+        if "predictions.json" in files:
+            pred_count += 1
+    return pred_count >= 3
+
+def is_finetune_cache_valid(cache_dir):
+    if not os.path.exists(cache_dir):
+        return False
+    pred_count = 0
+    for root, dirs, files in os.walk(cache_dir):
+        if "predictions.json" in files:
+            pred_count += 1
+    return pred_count >= 3
+
+def verify_eval_run(path_to_check, description):
+    print(f"[Eval] Verifying {description} path: {path_to_check}...")
+    if not os.path.exists(path_to_check):
+        raise RuntimeError(f"CRITICAL ERROR: {description} directory was NOT created at: {path_to_check}")
+    
+    # Check if predictions.json or results.txt or surprisal.json exists and is non-empty
+    found_valid = False
+    for root, dirs, files in os.walk(path_to_check):
+        for f in files:
+            if f in ["predictions.json", "results.txt", "surprisal.json"]:
+                file_path = os.path.join(root, f)
+                if os.path.getsize(file_path) > 0:
+                    found_valid = True
+                    break
+        if found_valid:
+            break
+            
+    if not found_valid:
+        raise RuntimeError(f"CRITICAL ERROR: {description} completed but no valid prediction/result files were written in: {path_to_check}")
+    print(f"[Eval] Success! Verified {description} results are stored correctly.")
+
+def run_pipeline(model_name: str, epochs: int = 10, skip_eval: bool = False, skip_aoa: bool = True, skip_glue: bool = False):
+    # Configure persistent cache paths locally to avoid duplicate downloads
+    os.environ["HF_HOME"] = os.path.abspath("./hf_cache")
+    os.environ["NLTK_DATA"] = os.path.abspath("./nltk_data")
+    os.makedirs("./hf_cache", exist_ok=True)
+    os.makedirs("./nltk_data", exist_ok=True)
+
+    # Programmatic Hugging Face Hub Login if HF_TOKEN is in environment
+    hf_token = os.environ.get("HF_TOKEN")
+
+    if hf_token:
+        try:
+            from huggingface_hub import login
+            login(token=hf_token)
+            print("[HF] Programmatic login successful using HF_TOKEN.")
+        except Exception as e:
+            print(f"[HF] Warning: Programmatic login failed: {e}")
+
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+    from datasets import load_dataset
+    from tokenizers import Tokenizer
+    from transformers import PreTrainedTokenizerFast
+    
+    # Import model architecture
+    from modeling_xpertgpt import (
+        XpertGPTModel,
+        XpertGPTModelConfig,
+        XpertGPTConfig
+    )
+
+    # Print GPU details
+    if torch.cuda.is_available():
+        gpu_name = torch.cuda.get_device_name(0)
+        print(f"\n[GPU] CUDA is available! Using GPU: {gpu_name}\n")
+    else:
+        print("\n[GPU] Warning: CUDA is NOT available! Running on CPU.\n")
+
+    # ─────────────────────────────────────────────────────────────
+    # GLOBAL HYPERPARAMETERS
+    # ─────────────────────────────────────────────────────────────
+    VOCAB_SIZE        = 16384          
+    MASK_TOKEN_ID     = 16383          
+    BLOCK_SIZE        = 512
+    BATCH_SIZE        = 16
+    GRAD_ACCUM_STEPS  = 1  # grad_acc_step = 1
+    EPOCHS            = epochs
+    LEARNING_RATE     = 3e-4  # lr = 3e-4
+    LR_MIN            = LEARNING_RATE * 0.05
+    WARMUP_STEPS      = 800  # warmup_steps = 800
+    WEIGHT_DECAY      = 0.1
+    GRAD_CLIP         = 1.0
+
+    NUM_THIN_BLOCKS   = 4             
+    EC_CAPACITY_FACTOR = 2.0           
+    CAUSAL_RATIO      = 1 / 1        
+
+    MASK_PROB_START   = 0.20
+    MASK_PROB_END     = 0.10
+
+    # Output directories locally
+    model_dir = os.path.abspath(f"./checkpoints/{model_name}")
+    os.makedirs(model_dir, exist_ok=True)
+    local_results_dir = os.path.abspath(f"./results/{model_name}")
+    os.makedirs(local_results_dir, exist_ok=True)
+
+    # ─────────────────────────────────���───────────────────────────
+    # Helper: Save Hugging Face Compliant Checkpoint
+    # ─────────────────────────────────────────────────────────────
+    def save_hf_checkpoint(raw_model, checkpoint_dir_name, tokenizer):
+        save_dir = os.path.join(model_dir, checkpoint_dir_name)
+        os.makedirs(save_dir, exist_ok=True)
+        print(f"\n[Checkpoint] Saving Hugging Face format checkpoint to '{save_dir}'...")
+        
+        # A. Convert state dict keys to CausalLM wrapper naming
+        state_dict = raw_model.state_dict()
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            name = k
+            if name.startswith("_orig_mod."):
+                name = name[10:]
+            if name.startswith("model."):
+                name = name[6:]
+            
+            if name == "lm_head.weight":
+                new_state_dict["lm_head.weight"] = v
+            else:
+                new_state_dict[f"transformer.{name}"] = v
+                
+        torch.save(new_state_dict, os.path.join(save_dir, "pytorch_model.bin"))
+        
+        # B. Copy modeling.py and configuration.py
+        shutil.copy("modeling_xpertgpt.py", os.path.join(save_dir, "modeling_xpertgpt.py"))
+        shutil.copy("configuration_xpertgpt.py", os.path.join(save_dir, "configuration_xpertgpt.py"))
+        
+        # C. Create config.json
+        config_dict = {
+            "auto_map": {
+                "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
+                "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
+                "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
+            },
+            "vocab_size": VOCAB_SIZE,
+            "block_size": BLOCK_SIZE,
+            "d_model": 256,        # d_model = 256
+            "hidden_size": 256,    # hidden_size = 256
+            "d_thin": 384,
+            "num_layers": 6,
+            "num_blocks": NUM_THIN_BLOCKS,
+            "capacity_factor": EC_CAPACITY_FACTOR,
+            "dropout": 0.1,
+            "model_type": "xpertgpt"
+        }
+        with open(os.path.join(save_dir, "config.json"), "w") as f:
+            json.dump(config_dict, f, indent=2)
+            
+        # D. Save tokenizer config files
+        fast_tokenizer = PreTrainedTokenizerFast(
+            tokenizer_object=tokenizer,
+            bos_token="[CLS]",
+            eos_token="[SEP]",
+            unk_token="[UNK]",
+            pad_token="[PAD]",
+            mask_token="[MASK]"
+        )
+        fast_tokenizer.save_pretrained(save_dir)
+        print(f"[Checkpoint] Checkpoint '{checkpoint_dir_name}' successfully saved.")
+
+    # ─────────────────────────────────────────────────────────────
+    # Tokenizer Training
+    # ─────────────────────────────────────────────────────────────
+    def build_and_train_tokenizer(texts: list) -> Tokenizer:
+        from tokenizers.models import BPE
+        from tokenizers.trainers import BpeTrainer
+        from tokenizers.pre_tokenizers import Whitespace
+        
+        vocab_path = os.path.join(model_dir, "bpe_vocab_16k.json")
+        if os.path.exists(vocab_path):
+            print(f"[Tokenizer] Loading trained BPE model layout from '{vocab_path}'...")
+            return Tokenizer.from_file(vocab_path)
+            
+        print(f"[Tokenizer] Generating fresh HuggingFace BPE Tokenizer model with {VOCAB_SIZE} slots...")
+        tokenizer = Tokenizer(BPE(unk_token="[UNK]"))
+        tokenizer.pre_tokenizer = Whitespace()
+        
+        trainer = BpeTrainer(
+            vocab_size=VOCAB_SIZE, 
+            special_tokens=["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
+        )
+        tokenizer.train_from_iterator(texts, trainer)
+        tokenizer.save(vocab_path)
+        print(f"[Tokenizer] Tokenizer training completed and saved to '{vocab_path}'.")
+        return tokenizer
+
+    # ─────────────────────────────────────────────────────────────
+    # Data Loader Setup
+    # ─────────────────────────────────────────────────────────────
+    class DataLoaderLite:
+        def __init__(self, B: int, T: int, texts: list, tokenizer: Tokenizer, name: str):
+            self.B = B
+            self.T = T
+            
+            print(f"[DataLoader:{name}] Tokenising dataset sequences...")
+            all_ids = []
+            for t in texts:
+                if t.strip():
+                    encoded = tokenizer.encode(t).ids
+                    all_ids.extend(encoded)
+
+            self.tokens = torch.tensor(all_ids, dtype=torch.long)
+            self.chunk_size = B * T
+            self.n_chunks = (len(self.tokens) - 1) // self.chunk_size
+            self.indices = list(range(self.n_chunks))
+            self.pos = 0
+            self._shuffle()
+            
+            print(f"[DataLoader:{name}] Total tokens: {len(self.tokens):,} | Epoch steps: {self.n_chunks:,}")
+
+        def _shuffle(self):
+            random.shuffle(self.indices)
+            self.pos = 0
+
+        def steps_per_epoch(self) -> int:
+            return self.n_chunks
+
+        def next_batch(self):
+            B, T = self.B, self.T
+            if self.pos >= len(self.indices):
+                self._shuffle()
+                
+            chunk_idx = self.indices[self.pos]
+            self.pos += 1
+            
+            start_pos = chunk_idx * self.chunk_size
+            temp = self.tokens[start_pos : start_pos + self.chunk_size + 1]
+            
+            x = temp[:-1].view(B, T)
+            y = temp[1:].view(B, T)
+            return x, y
+
+    # ─────────────────────────────────────────────────────────────
+    # Batch preparation and schedules
+    # ─────────────────────────────────────────────────────────────
+    def get_current_mask_prob(global_step: int, total_steps: int) -> float:
+        ratio = min(1.0, global_step / total_steps)
+        return MASK_PROB_START + ratio * (MASK_PROB_END - MASK_PROB_START)
+
+    def prepare_causal_batch(x: torch.Tensor, y: torch.Tensor):
+        return x, y, False
+
+    def prepare_masked_batch(x: torch.Tensor, y: torch.Tensor, mask_prob: float, mask_token_id: int):
+        B, T = x.size()
+        mask = torch.rand(B, T, device=x.device) < mask_prob
+        masked_x = x.clone()
+        masked_x[mask] = mask_token_id
+
+        targets = torch.full_like(y, -100)
+        targets[mask] = y[mask]
+
+        return masked_x, targets, True
+
+    def get_hybrid_batch(train_loader: DataLoaderLite, global_step: int, total_steps: int, device: torch.device):
+        x, y = train_loader.next_batch()
+        x, y = x.to(device), y.to(device)
+
+        if random.random() < CAUSAL_RATIO:
+            input_ids, targets, bidir = prepare_causal_batch(x, y)
+        else:
+            mask_prob = get_current_mask_prob(global_step, total_steps)
+            input_ids, targets, bidir = prepare_masked_batch(x, y, mask_prob, MASK_TOKEN_ID)
+
+        return input_ids, targets, bidir
+
+    def get_lr(it: int, total_steps: int) -> float:
+        if it < WARMUP_STEPS:
+            return LEARNING_RATE * (it + 1) / WARMUP_STEPS
+        if it >= total_steps:
+            return LR_MIN
+        decay_ratio = (it - WARMUP_STEPS) / (total_steps - WARMUP_STEPS)
+        coeff        = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
+        return LR_MIN + coeff * (LEARNING_RATE - LR_MIN)
+
+    # ─────────────────────────────────────────────────────────────
+    # Dataset Preparation
+    # ─────────────────────────────────────────────────────────────
+    print("\n[Data] Loading BabyLM-2026-Strict-Small ...")
+    ds = load_dataset("BabyLM-community/BabyLM-2026-Strict-Small")
+    all_text = list(ds['train']['text'])
+
+    tokenizer = build_and_train_tokenizer(all_text)
+
+    split       = int(len(all_text) * 0.95)
+    train_texts = all_text[:split]
+    val_texts   = all_text[split:]
+
+    train_loader = DataLoaderLite(BATCH_SIZE, BLOCK_SIZE, train_texts, tokenizer, "train")
+    val_loader   = DataLoaderLite(BATCH_SIZE, BLOCK_SIZE, val_texts, tokenizer, "val")
+
+    chunks_per_epoch = train_loader.steps_per_epoch()
+    steps_per_epoch  = chunks_per_epoch // GRAD_ACCUM_STEPS 
+    total_steps      = steps_per_epoch * EPOCHS
+
+    cfg   = XpertGPTModelConfig()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch.manual_seed(42)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(42)
+    random.seed(42)
+    if hasattr(torch, 'set_float32_matmul_precision'):
+        torch.set_float32_matmul_precision('high')
+
+    model = XpertGPTModel(cfg).to(device)
+
+    # ─────────────────────────────────────────────────────────────
+    # Training Resume Check
+    # ──���──────────────────────────────────────────────────────────
+    words_trained = 0
+    next_milestone_idx = 0
+    global_step = 0
+    
+    milestones = sorted(list(set([i * 1_000_000 for i in range(1, 11)] + [i * 10_000_000 for i in range(1, 11)])))
+    
+    resume_checkpoint_dir = None
+    for idx in range(len(milestones) - 1, -1, -1):
+        m = milestones[idx]
+        ckpt_name = f"chck_{m // 1_000_000}M"
+        ckpt_path = os.path.join(model_dir, ckpt_name)
+        if os.path.exists(os.path.join(ckpt_path, "pytorch_model.bin")):
+            config_json_path = os.path.join(ckpt_path, "config.json")
+            if os.path.exists(config_json_path):
+                try:
+                    with open(config_json_path, "r") as f:
+                        saved_config = json.load(f)
+                    if saved_config.get("d_model") == 256:
+                        resume_checkpoint_dir = ckpt_path
+                        next_milestone_idx = idx + 1
+                        words_trained = m
+                        global_step = words_trained // (BATCH_SIZE * BLOCK_SIZE)
+                        print(f"[Training] Found existing milestone checkpoint '{ckpt_name}'. Resuming from step {global_step:,} ({words_trained:,} tokens trained)...")
+                        break
+                    else:
+                        print(f"[Training] Found checkpoint '{ckpt_name}' but it has mismatch d_model={saved_config.get('d_model')}. Starting fresh.")
+                except Exception as e:
+                    pass
+
+    # Load weights if resuming
+    if resume_checkpoint_dir is not None:
+        print(f"[Model] Loading weights from checkpoint '{resume_checkpoint_dir}'...")
+        state_dict = torch.load(os.path.join(resume_checkpoint_dir, "pytorch_model.bin"), map_location=device)
+        model_state_dict = {}
+        for k, v in state_dict.items():
+            name = k
+            if name.startswith("transformer."):
+                name = name[12:]
+            model_state_dict[name] = v
+        model.load_state_dict(model_state_dict)
+
+    # Check if final main model exists
+    main_ckpt_path = os.path.join(model_dir, "main")
+    if os.path.exists(os.path.join(main_ckpt_path, "pytorch_model.bin")):
+        print("\n[Pipeline] Final checkpoint 'main' already exists. Skipping training phase and transitioning directly to evaluations!")
+    else:
+        # torch.compile
+        try:
+            model = torch.compile(model)
+            print("[Model] torch.compile() successfully verified graph optimizations")
+        except Exception as e:
+            print(f"[Model] torch.compile() skipped ({e})")
+
+        # Optimizer
+        param_dict     = {n: p for n, p in model.named_parameters() if p.requires_grad}
+        decay_params   = [p for p in param_dict.values() if p.dim() >= 2]
+        nodecay_params = [p for p in param_dict.values() if p.dim() < 2]
+        groups = [
+            {'params': decay_params,   'weight_decay': WEIGHT_DECAY},
+            {'params': nodecay_params, 'weight_decay': 0.0},
+        ]
+        fused_ok  = 'fused' in inspect.signature(torch.optim.AdamW).parameters
+        use_fused = fused_ok and ('cuda' in device)
+        optimizer = torch.optim.AdamW(groups, lr=LEARNING_RATE, betas=(0.9, 0.95), eps=1e-8, fused=use_fused)
+
+        # Helper for validation loss calculation
+        def evaluate_validation_loss(model_eval, val_loader_eval, dev, autocast):
+            was_training = model_eval.training
+            model_eval.eval()
+            
+            from modeling_xpertgpt import ROUTER_TRACKER
+            old_tracker_enabled = ROUTER_TRACKER["enabled"]
+            ROUTER_TRACKER["enabled"] = False
+            
+            val_loss_accum = 0.0
+            val_steps = min(val_loader_eval.steps_per_epoch(), 50)
+            with torch.no_grad():
+                for _ in range(val_steps):
+                    x, y = val_loader_eval.next_batch()
+                    x, y = x.to(dev), y.to(dev)
+                    with autocast:
+                        _, loss = model_eval(x, y, bidirectional=False)
+                    val_loss_accum += loss.item()
+            if was_training:
+                model_eval.train()
+                
+            ROUTER_TRACKER["enabled"] = old_tracker_enabled
+            return val_loss_accum / val_steps
+
+        loss_records = []
+        checkpoint_perplexities = []
+        
+        from modeling_xpertgpt import ROUTER_TRACKER
+        ROUTER_TRACKER["enabled"] = True
+        ROUTER_TRACKER["expert_counts"] = []
+
+        model.train()
+        autocast_ctx = torch.autocast(device_type="cuda" if "cuda" in device else "cpu", dtype=torch.bfloat16, enabled=True)
+
+        start_epoch = global_step // steps_per_epoch
+        start_chunk = (global_step % steps_per_epoch) * GRAD_ACCUM_STEPS
+
+        print(f"\n[Training] Starting XpertGPT MoEP training for {EPOCHS} epochs...")
+        for epoch in range(start_epoch, EPOCHS):
+            train_loader._shuffle()
+            if epoch == start_epoch and start_chunk > 0:
+                print(f"[Training] Fast-forwarding dataloader to chunk index {start_chunk}...")
+                train_loader.pos = start_chunk
+
+            optimizer.zero_grad(set_to_none=True)
+            loss_accum = 0.0
+            
+            start_chunk_idx = start_chunk if epoch == start_epoch else 0
+            for chunk_step in range(start_chunk_idx, chunks_per_epoch):
+                t0 = time.perf_counter()
+
+                lr = get_lr(global_step, total_steps)
+                for pg in optimizer.param_groups:
+                    pg['lr'] = lr
+
+                input_ids, targets, bidir = get_hybrid_batch(train_loader, global_step, total_steps, device)
+                words_trained += input_ids.numel()
+
+                with autocast_ctx:
+                    _, loss = model(input_ids, targets, bidirectional=bidir)
+                scaled_loss  = loss / GRAD_ACCUM_STEPS
+                loss_accum  += scaled_loss.item()
+                scaled_loss.backward()
+
+                # Optimizer Step
+                if (chunk_step + 1) % GRAD_ACCUM_STEPS == 0:
+                    norm = torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
+                    optimizer.step()
+                    optimizer.zero_grad(set_to_none=True)
+                    if "cuda" in device:
+                        torch.cuda.synchronize()
+
+                    dt = (time.perf_counter() - t0) * 1000
+                    mode_tag = "MLM" if bidir else "CLM"
+                    mask_p   = get_current_mask_prob(global_step, total_steps)
+                    
+                    current_step = (chunk_step + 1) // GRAD_ACCUM_STEPS
+                    print(
+                        f"[E{epoch+1:02d} {current_step:>5d}/{steps_per_epoch} G{global_step:>7d}|{mode_tag}] "
+                        f"train={loss_accum:.4f}  mask={mask_p:.1%}  norm={norm:.3f}  lr={lr:.2e}  dt={dt:6.1f}ms  words={words_trained:,}"
+                    )
+                    
+                    # Record validation and training loss every 100 steps
+                    if (global_step + 1) % 100 == 0:
+                        val_loss = evaluate_validation_loss(model, val_loader, device, autocast_ctx)
+                        loss_records.append({
+                            "step": global_step + 1,
+                            "train_loss": loss_accum,
+                            "val_loss": val_loss
+                        })
+                        with open("loss_records.json", "w") as f:
+                            json.dump(loss_records, f, indent=2)
+                        if os.path.exists(model_dir):
+                            with open(os.path.join(model_dir, "loss_records.json"), "w") as f:
+                                json.dump(loss_records, f, indent=2)
+                        print(f"[Metrics G{global_step+1}] Recorded train_loss={loss_accum:.4f}, val_loss={val_loss:.4f}")
+                        
+                    loss_accum = 0.0
+                    global_step += 1
+
+                # Check if we passed a milestone for checkpointing
+                if next_milestone_idx < len(milestones) and words_trained >= milestones[next_milestone_idx]:
+                    milestone_val = milestones[next_milestone_idx]
+                    if milestone_val < 10_000_000:
+                        milestone_name = f"chck_{milestone_val // 1_000_000}M"
+                    else:
+                        milestone_name = f"chck_{(milestone_val // 10_000_000) * 10}M"
+                    
+                    raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model
+                    save_hf_checkpoint(raw_model, milestone_name, tokenizer)
+                    
+                    # Calculate validation perplexity at this milestone
+                    val_loss = evaluate_validation_loss(model, val_loader, device, autocast_ctx)
+                    val_ppl = math.exp(val_loss)
+                    checkpoint_perplexities.append({
+                        "checkpoint": milestone_name,
+                        "words_trained": milestone_val,
+                        "global_step": global_step,
+                        "val_loss": val_loss,
+                        "val_perplexity": val_ppl
+                    })
+                    with open("checkpoint_perplexities.json", "w") as f:
+                        json.dump(checkpoint_perplexities, f, indent=2)
+                    if os.path.exists(model_dir):
+                        with open(os.path.join(model_dir, "checkpoint_perplexities.json"), "w") as f:
+                            json.dump(checkpoint_perplexities, f, indent=2)
+                    print(f"[Milestone {milestone_name}] Evaluated val_perplexity={val_ppl:.2f}")
+
+                    next_milestone_idx += 1
+
+        # Save final model as 'main'
+        raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model
+        save_hf_checkpoint(raw_model, "main", tokenizer)
+        
+        # Calculate validation perplexity for final checkpoint
+        val_loss = evaluate_validation_loss(model, val_loader, device, autocast_ctx)
+        val_ppl = math.exp(val_loss)
+        checkpoint_perplexities.append({
+            "checkpoint": "main",
+            "words_trained": words_trained,
+            "global_step": global_step,
+            "val_loss": val_loss,
+            "val_perplexity": val_ppl
+        })
+        with open("checkpoint_perplexities.json", "w") as f:
+            json.dump(checkpoint_perplexities, f, indent=2)
+        if os.path.exists(model_dir):
+            with open(os.path.join(model_dir, "checkpoint_perplexities.json"), "w") as f:
+                json.dump(checkpoint_perplexities, f, indent=2)
+        print(f"[Final Checkpoint main] Evaluated val_perplexity={val_ppl:.2f}")
+        
+        print("\n[Training] Training phase complete!")
+
+        # ─────────────────────────────────────────────────────────────
+        # Router utilization report and plot generation
+        # ─────────────────────────────────────────────────────────────
+        from modeling_xpertgpt import ROUTER_TRACKER
+        total_expert_counts = [0, 0, 0, 0, 0]
+        for cnt in ROUTER_TRACKER["expert_counts"]:
+            for i in range(5):
+                total_expert_counts[i] += cnt[i]
+                
+        total_tokens = sum(total_expert_counts)
+        if total_tokens > 0:
+            fractions = [total_expert_counts[i] / total_tokens for i in range(5)]
+        else:
+            fractions = [0.0] * 5
+            
+        fraction_no_expert = fractions[0]
+        fraction_multiple_experts = sum(fractions[2:])
+        
+        report_text = f"""=== ROUTER UTILIZATION REPORT ===
+Total tokens routed: {total_tokens:,}
+
+Number of experts selected per token (0, 1, 2, 3, 4):
+- 0 experts: {total_expert_counts[0]:,} tokens ({fractions[0]:.2%})
+- 1 expert:  {total_expert_counts[1]:,} tokens ({fractions[1]:.2%})
+- 2 experts: {total_expert_counts[2]:,} tokens ({fractions[2]:.2%})
+- 3 experts: {total_expert_counts[3]:,} tokens ({fractions[3]:.2%})
+- 4 experts: {total_expert_counts[4]:,} tokens ({fractions[4]:.2%})
+
+Summary metrics:
+- Fraction of tokens receiving NO expert: {fraction_no_expert:.2%}
+- Fraction of tokens receiving MULTIPLE experts: {fraction_multiple_experts:.2%}
+"""
+        with open("router_report.txt", "w") as f:
+            f.write(report_text)
+        if os.path.exists(model_dir):
+            with open(os.path.join(model_dir, "router_report.txt"), "w") as f:
+                f.write(report_text)
+        print(report_text)
+
+        # Generate bar chart
+        try:
+            import matplotlib
+            matplotlib.use('Agg')
+            import matplotlib.pyplot as plt
+            
+            categories = ['0', '1', '2', '3', '4']
+            plt.figure(figsize=(8, 5))
+            plt.bar(categories, [f * 100 for f in fractions], color='#10b981', edgecolor='#059669', width=0.6)
+            plt.xlabel('Number of Experts Selected per Token')
+            plt.ylabel('Percentage of Tokens (%)')
+            plt.title('Router Utilization / Expert Choices per Token')
+            plt.grid(axis='y', linestyle='--', alpha=0.6)
+            
+            for i, f in enumerate(fractions):
+                plt.text(i, f * 100 + 1, f"{f:.2%}", ha='center', fontweight='bold')
+                
+            plt.ylim(0, max([f * 100 for f in fractions]) + 10)
+            plt.savefig('router_utilization.png', dpi=150)
+            if os.path.exists(model_dir):
+                plt.savefig(os.path.join(model_dir, 'router_utilization.png'), dpi=150)
+            plt.close()
+            print("[Plots] Router utilization plot saved successfully.")
+        except Exception as e:
+            print(f"[Plots] Warning: Could not generate plots: {e}")
+
+    if skip_eval:
+        print("[Pipeline] Skipping evaluations phase as requested.")
+        return
+
+    # ─────────────────────────────────────────────────────────────
+    # 2. RUN EVALUATION PIPELINE
+    # ─────────────────────────────────────────────────────────────
+    local_results_dir = os.path.abspath(f"./results/{model_name}")
+    os.makedirs(local_results_dir, exist_ok=True)
+    local_main_res = os.path.join(local_results_dir, "main")
+
+    # Ensure clone_dir exists and has the global_piqa files
+    clone_parent = os.path.abspath("./babylm_eval_repo")
+    
+    # Self-healing check for global_piqa presence
+    has_global_piqa = False
+    for potential_strict in [os.path.join(clone_parent, "babylm-eval", "strict"), os.path.join(clone_parent, "strict")]:
+        if os.path.exists(os.path.join(potential_strict, "evaluation_pipeline", "global_piqa")):
+            has_global_piqa = True
+            break
+            
+    if not has_global_piqa:
+        print("[Eval] Cloned repository does not contain global_piqa tasks.")
+        print("[Eval] Deleting and cloning official main branch...")
+        if os.path.exists(clone_parent):
+            shutil.rmtree(clone_parent)
+        subprocess.run([
+            "git", "clone", "-b", "main",
+            "https://github.com/babylm-org/babylm-eval.git",
+            clone_parent
+        ], check=True)
+
+    # Determine strict_dir path dynamically
+    if os.path.exists(os.path.join(clone_parent, "strict")):
+        strict_dir = os.path.join(clone_parent, "strict")
+    else:
+        strict_dir = os.path.join(clone_parent, "babylm-eval", "strict")
+        
+    print(f"[Eval] Using strict directory: {strict_dir}")
+    os.environ["PYTHONPATH"] = strict_dir
+
+    def patch_evaluation_run_script(strict_dir):
+        import pathlib
+        run_file = os.path.join(strict_dir, "evaluation_pipeline", "sentence_zero_shot", "run.py")
+        if not os.path.exists(run_file):
+            print(f"[GlobalPIQA] Warning: {run_file} not found. Cannot patch.")
+            return
+            
+        print(f"[GlobalPIQA] Patching local checkpoint loader in {run_file}...")
+        with open(run_file, "r") as f:
+            content = f.read()
+            
+        # Check if already patched
+        if "Local checkpoint directory patch" in content:
+            print("[GlobalPIQA] Script already patched.")
+            return
+            
+        target_str = """def main():
+    args = _parse_arguments()
+    if args.images_path is not None:
+        assert args.batch_size == 1, "Multimodal only works in batch size 1!"
+    dataset = args.data_path.stem
+    args.model_name = pathlib.Path(args.model_path_or_name).stem
+    if args.revision_name is None:
+        revision_name = "main"
+    else:
+        revision_name = args.revision_name"""
+            
+        patch_str = """def main():
+    args = _parse_arguments()
+    if args.images_path is not None:
+        assert args.batch_size == 1, "Multimodal only works in batch size 1!"
+    dataset = args.data_path.stem
+    
+    # Local checkpoint directory patch
+    import os
+    model_path = args.model_path_or_name
+    args.model_name = pathlib.Path(model_path).stem
+    revision_name = args.revision_name if args.revision_name else "main"
+    
+    if os.path.isdir(model_path):
+        target_revision = args.revision_name if args.revision_name else "main"
+        if os.path.exists(os.path.join(model_path, target_revision)):
+            args.model_path_or_name = os.path.join(model_path, target_revision)
+            args.revision_name = None"""
+                
+        if target_str in content:
+            new_content = content.replace(target_str, patch_str)
+            with open(run_file, "w") as f:
+                f.write(new_content)
+            print("[GlobalPIQA] Successfully patched run.py")
+        else:
+            print("[GlobalPIQA] Warning: Could not find target pattern in run.py. Manual patch may be needed.")
+
+    # Patch sentence zero shot loader inside cloned repo
+    patch_evaluation_run_script(strict_dir)
+
+    print("[Eval] Stripping Windows-specific packages from requirements.txt...")
+    req_file_path = os.path.join(strict_dir, "requirements.txt")
+    if os.path.exists(req_file_path):
+        with open(req_file_path, "r") as f:
+            lines = f.readlines()
+        with open(req_file_path, "w") as f:
+            for line in lines:
+                if "pywin" not in line.lower() and "wintypes" not in line.lower():
+                    f.write(line)
+
+    print("[Eval] Verifying and installing evaluation dependencies programmatically...")
+    required_packages = {
+        "nltk": "nltk",
+        "pandas": "pandas",
+        "statsmodels": "statsmodels",
+        "sklearn": "scikit-learn",
+        "scipy": "scipy"
+    }
+    for pkg_import, pkg_install in required_packages.items():
+        try:
+            __import__(pkg_import)
+        except ImportError:
+            print(f"[Eval] Package '{pkg_install}' not found. Installing it programmatically...")
+            import sys
+            subprocess.run([sys.executable, "-m", "pip", "install", pkg_install], check=True)
+    
+    print("[Eval] Downloading NLTK tokenizer resources...")
+    import nltk
+    nltk.download('punkt', download_dir=os.environ["NLTK_DATA"])
+    nltk.download('punkt_tab', download_dir=os.environ["NLTK_DATA"])
+
+    # Ensure standard zero-shot datasets are downloaded
+    blimp_fast_dir = os.path.join(strict_dir, "evaluation_data", "fast_eval", "blimp_fast")
+    if not os.path.exists(blimp_fast_dir) or not os.listdir(blimp_fast_dir):
+        print("[Eval] Standard zero-shot datasets not found. Downloading...")
+        subprocess.run(["python", "-m", "scripts.download_evals"], cwd=strict_dir, check=True)
+        
+        # Unzip EWoK fast
+        ewok_zip = os.path.join(strict_dir, "evaluation_data/fast_eval/ewok_fast.zip")
+        if os.path.exists(ewok_zip):
+            print("[Eval] Unzipping EWoK fast data...")
+            bad_nested_dir = os.path.join(strict_dir, "evaluation_data/fast_eval/evaluation_data")
+            if os.path.exists(bad_nested_dir):
+                shutil.rmtree(bad_nested_dir)
+            subprocess.run(["unzip", "-o", "-P", "BabyLM2025", "evaluation_data/fast_eval/ewok_fast.zip", "-d", "."], cwd=strict_dir, check=True)
+            
+        # Download EWoK full
+        print("[Eval] Downloading and filtering full EWoK dataset...")
+        subprocess.run(["python", "-m", "evaluation_pipeline.ewok.dl_and_filter"], cwd=strict_dir, check=True)
+
+    # Download GlobalPIQA dataset
+    global_piqa_parallel_dir = os.path.join(strict_dir, "evaluation_data", "fast_eval", "global_piqa_parallel")
+    if not os.path.exists(global_piqa_parallel_dir) or not os.listdir(global_piqa_parallel_dir):
+        print("[Eval] GlobalPIQA dataset not found. Downloading...")
+        subprocess.run(["python", "evaluation_pipeline/global_piqa/dl.py"], cwd=strict_dir, check=True)
+
+    # Ensure all scripts are executable
+    print("[Eval] Making evaluation shell scripts executable...")
+    subprocess.run("chmod +x scripts/*.sh", shell=True, cwd=strict_dir, check=True)
+
+    def run_task_with_cache(checkpoint, task, output_subpath, cmd):
+        # Determine paths
+        local_cache_path = os.path.join("./results", model_name, checkpoint, "zero_shot", "causal", task, output_subpath)
+        if task == "reading":
+            local_cache_path = os.path.join("./results", model_name, checkpoint, "zero_shot", "causal", "reading")
+        elif task == "comps":
+            local_cache_path = os.path.join("./results", model_name, checkpoint, "zero_shot", "causal", "comps", "comps")
+            
+        target_results_dir = os.path.join(strict_dir, "results", model_name, checkpoint, "zero_shot", "causal", task, output_subpath)
+        if task == "reading":
+            target_results_dir = os.path.join(strict_dir, "results", model_name, checkpoint, "zero_shot", "causal", "reading")
+        elif task == "comps":
+            target_results_dir = os.path.join(strict_dir, "results", model_name, checkpoint, "zero_shot", "causal", "comps", "comps")
+
+        # If cached, copy it over
+        cache_file = os.path.join(local_cache_path, "predictions.json")
+        
+        # Self-healing: invalidate old unfiltered entity_tracking caches
+        if task == "entity_tracking" and os.path.exists(cache_file):
+            try:
+                import json
+                with open(cache_file, "r") as f:
+                    preds = json.load(f)
+                is_valid_cache = True
+                for k, v in preds.items():
+                    if len(v.get("predictions", [])) in [605, 606, 607, 615, 529, 156, 187, 159]:
+                        is_valid_cache = False
+                        break
+                if not is_valid_cache:
+                    print(f"[Eval] Cached entity_tracking for '{checkpoint}' has incorrect old sizes. Invalidate and re-run fresh...")
+                    shutil.rmtree(local_cache_path, ignore_errors=True)
+            except Exception:
+                pass
+
+        if os.path.exists(cache_file):
+            print(f"[Eval] Task '{task}' ({output_subpath}) for checkpoint '{checkpoint}' is cached. Restoring...")
+            if os.path.exists(target_results_dir):
+                shutil.rmtree(target_results_dir)
+            os.makedirs(target_results_dir, exist_ok=True)
+            for item in os.listdir(local_cache_path):
+                s = os.path.join(local_cache_path, item)
+                d = os.path.join(target_results_dir, item)
+                if os.path.isdir(s):
+                    shutil.copytree(s, d)
+                else:
+                    shutil.copy2(s, d)
+            return
+
+        print(f"[Eval] Running task '{task}' ({output_subpath}) for checkpoint '{checkpoint}'...")
+        subprocess.run(cmd, cwd=strict_dir, check=True)
+
+        # Relocate from results/main if needed
+        possible_main_path = os.path.join(strict_dir, "results", "main", checkpoint, "zero_shot", "causal", task, output_subpath)
+        if task == "reading":
+            possible_main_path = os.path.join(strict_dir, "results", "main", checkpoint, "zero_shot", "causal", "reading")
+        elif task == "comps":
+            possible_main_path = os.path.join(strict_dir, "results", "main", checkpoint, "zero_shot", "causal", "comps", "comps")
+
+        if os.path.exists(possible_main_path) and possible_main_path != target_results_dir:
+            print(f"[Eval] Relocating results from {possible_main_path} to {target_results_dir}...")
+            if os.path.exists(target_results_dir):
+                shutil.rmtree(target_results_dir)
+            os.makedirs(os.path.dirname(target_results_dir), exist_ok=True)
+            shutil.move(possible_main_path, target_results_dir)
+
+        # Verify
+        verify_eval_run(target_results_dir, f"{checkpoint} {task} ({output_subpath})")
+
+        # Save to local cache
+        if os.path.exists(local_cache_path):
+            shutil.rmtree(local_cache_path)
+        os.makedirs(local_cache_path, exist_ok=True)
+        for item in os.listdir(target_results_dir):
+            s = os.path.join(target_results_dir, item)
+            d = os.path.join(local_cache_path, item)
+            if os.path.isdir(s):
+                shutil.copytree(s, d)
+            else:
+                shutil.copy2(s, d)
+
+    def run_finetune_task_with_cache(task, cmd):
+        local_cache_path = os.path.join("./results", model_name, "main", "finetune", task)
+        target_results_dir = os.path.join(strict_dir, "results", model_name, "main", "finetune", task)
+
+        if os.path.exists(os.path.join(local_cache_path, "predictions.json")):
+            print(f"[Eval] GLUE task '{task}' is cached. Restoring...")
+            if os.path.exists(target_results_dir):
+                shutil.rmtree(target_results_dir)
+            os.makedirs(target_results_dir, exist_ok=True)
+            for item in os.listdir(local_cache_path):
+                s = os.path.join(local_cache_path, item)
+                d = os.path.join(target_results_dir, item)
+                if os.path.isdir(s):
+                    shutil.copytree(s, d)
+                else:
+                    shutil.copy2(s, d)
+            return
+
+        print(f"[Eval] Running GLUE task '{task}'...")
+        subprocess.run(cmd, cwd=strict_dir, check=True)
+
+        # Relocate from results/main/main if needed
+        possible_main_path = os.path.join(strict_dir, "results", "main", "main", "finetune", task)
+        if os.path.exists(possible_main_path) and possible_main_path != target_results_dir:
+            print(f"[Eval] Relocating results from {possible_main_path} to {target_results_dir}...")
+            if os.path.exists(target_results_dir):
+                shutil.rmtree(target_results_dir)
+            os.makedirs(os.path.dirname(target_results_dir), exist_ok=True)
+            shutil.move(possible_main_path, target_results_dir)
+
+        # Verify
+        verify_eval_run(target_results_dir, f"GLUE task {task}")
+
+        # Cache locally
+        if os.path.exists(local_cache_path):
+            shutil.rmtree(local_cache_path)
+        os.makedirs(local_cache_path, exist_ok=True)
+        for item in os.listdir(target_results_dir):
+            s = os.path.join(target_results_dir, item)
+            d = os.path.join(local_cache_path, item)
+            if os.path.isdir(s):
+                shutil.copytree(s, d)
+            else:
+                shutil.copy2(s, d)
+
+    def run_aoa_with_cache(cmd):
+        local_cache_path = os.path.join("./results", model_name, "main", "aoa")
+        target_results_dir = os.path.join(strict_dir, "results", model_name, "main", "aoa")
+
+        if os.path.exists(os.path.join(local_cache_path, "aoa_score.json")) or os.path.exists(os.path.join(local_cache_path, "surprisal.json")):
+            print(f"[Eval] AoA task is cached. Restoring...")
+            if os.path.exists(target_results_dir):
+                shutil.rmtree(target_results_dir)
+            os.makedirs(target_results_dir, exist_ok=True)
+            for item in os.listdir(local_cache_path):
+                s = os.path.join(local_cache_path, item)
+                d = os.path.join(target_results_dir, item)
+                if os.path.isdir(s):
+                    shutil.copytree(s, d)
+                else:
+                    shutil.copy2(s, d)
+            return
+
+        print("[Eval] Running AoA task...")
+        subprocess.run(cmd, cwd=strict_dir, check=True)
+
+        # Relocate from results/main/main if needed
+        possible_main_path = os.path.join(strict_dir, "results", "main", "main", "aoa")
+        if os.path.exists(possible_main_path) and possible_main_path != target_results_dir:
+            print(f"[Eval] Relocating results from {possible_main_path} to {target_results_dir}...")
+            if os.path.exists(target_results_dir):
+                shutil.rmtree(target_results_dir)
+            os.makedirs(os.path.dirname(target_results_dir), exist_ok=True)
+            shutil.move(possible_main_path, target_results_dir)
+
+        # Verify
+        verify_eval_run(target_results_dir, "AoA task")
+
+        # Cache locally
+        if os.path.exists(local_cache_path):
+            shutil.rmtree(local_cache_path)
+        os.makedirs(local_cache_path, exist_ok=True)
+        for item in os.listdir(target_results_dir):
+            s = os.path.join(target_results_dir, item)
+            d = os.path.join(local_cache_path, item)
+            if os.path.isdir(s):
+                shutil.copytree(s, d)
+            else:
+                shutil.copy2(s, d)
+
+    # ─────────────────────────────────────────────────────────────
+    # B. FINAL MODEL 'main' FULL ZERO-SHOT EVALUATION
+    # ─────────────────────────────────────────────────────────────
+    main_ckpt_path = os.path.join(model_dir, "main")
+    if os.path.exists(main_ckpt_path):
+        print(f"[Eval] Running full zero-shot evaluation on main...")
+        # blimp filtered
+        run_task_with_cache(
+            "main", "blimp", "blimp_filtered",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "blimp", "--data_path", "evaluation_data/full_eval/blimp_filtered", "--save_predictions"]
+        )
+        # supplement filtered
+        run_task_with_cache(
+            "main", "blimp", "supplement_filtered",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "blimp", "--data_path", "evaluation_data/full_eval/supplement_filtered", "--save_predictions"]
+        )
+        # ewok filtered
+        run_task_with_cache(
+            "main", "ewok", "ewok_filtered",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "ewok", "--data_path", "evaluation_data/full_eval/ewok_filtered", "--save_predictions"]
+        )
+        # entity tracking
+        run_task_with_cache(
+            "main", "entity_tracking", "entity_tracking",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "entity_tracking", "--data_path", "evaluation_data/full_eval/entity_tracking", "--save_predictions"]
+        )
+        # comps
+        run_task_with_cache(
+            "main", "comps", "comps",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "comps", "--data_path", "evaluation_data/full_eval/comps", "--save_predictions"]
+        )
+        # reading
+        run_task_with_cache(
+            "main", "reading", "reading",
+            ["python", "-m", "evaluation_pipeline.reading.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--data_path", "evaluation_data/full_eval/reading/reading_data.csv"]
+        )
+        # global piqa parallel
+        run_task_with_cache(
+            "main", "global_piqa_parallel", "global_piqa_parallel",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "global_piqa_parallel", "--data_path", "evaluation_data/full_eval/global_piqa_parallel", "--save_predictions"]
+        )
+        # global piqa nonparallel
+        run_task_with_cache(
+            "main", "global_piqa_nonparallel", "global_piqa_nonparallel",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", main_ckpt_path, "--backend", "causal", "--task", "global_piqa_nonparallel", "--data_path", "evaluation_data/full_eval/global_piqa_nonparallel", "--save_predictions"]
+        )
+
+    # ─────────────────────────────────────────────────────────────
+    # C. FINAL MODEL GLUE AND AOA EVALUATION
+    # ─────────────────────────────────────────────────────────────
+    if os.path.exists(main_ckpt_path):
+        # 1. GLUE fine-tuning
+        if skip_glue:
+            print("[Eval] Skipping GLUE fine-tuning evaluations as requested.")
+        else:
+            print("[Eval] Running GLUE fine-tuning evaluations on main task-by-task...")
+            glue_tasks = {
+                "boolq": ["boolq", "16", "10"],
+                "multirc": ["multirc", "16", "10"],
+                "rte": ["rte", "32", "10"],
+                "wsc": ["wsc", "32", "30"],
+                "mrpc": ["mrpc", "32", "10"],
+                "qqp": ["qqp", "32", "10"],
+                "mnli": ["mnli", "32", "10"]
+            }
+            for task_name, (task, bsz, max_epochs) in glue_tasks.items():
+                num_labels = "3" if task == "mnli" else "2"
+                metric_for_valid = "accuracy"
+                if task in ["mrpc", "qqp"]:
+                    metric_for_valid = "f1"
+                metrics = ["accuracy"]
+                if task != "mnli":
+                    metrics = ["accuracy", "f1", "mcc"]
+                
+                cmd = [
+                    "python", "-m", "evaluation_pipeline.finetune.run",
+                    "--model_name_or_path", main_ckpt_path,
+                    "--train_data", f"evaluation_data/full_eval/glue_filtered/{task}.train.jsonl",
+                    "--valid_data", f"evaluation_data/full_eval/glue_filtered/{task}.valid.jsonl",
+                    "--predict_data", f"evaluation_data/full_eval/glue_filtered/{task}.valid.jsonl",
+                    "--task", task,
+                    "--num_labels", num_labels,
+                    "--batch_size", bsz,
+                    "--learning_rate", "3e-5",
+                    "--num_epochs", max_epochs,
+                    "--sequence_length", "512",
+                    "--results_dir", "results",
+                    "--save",
+                    "--save_dir", "models",
+                    "--metric_for_valid", metric_for_valid,
+                    "--seed", "42",
+                    "--verbose",
+                    "--padding_side", "left",
+                    "--take_final"
+                ]
+                cmd.append("--metrics")
+                cmd.extend(metrics)
+                
+                run_finetune_task_with_cache(task_name, cmd)
+
+        # 2. AoA
+        if skip_aoa:
+            print("[Eval] Skipping AoA evaluations as requested.")
+        else:
+            run_aoa_with_cache([
+                "python", "-m", "evaluation_pipeline.AoA_word.run",
+                "--model_name", model_dir,
+                "--backend", "causal",
+                "--track_name", "strict-small",
+                "--word_path", "evaluation_data/full_eval/aoa/cdi_childes.json",
+                "--output_dir", "results"
+            ])
+
+    # ─────────────────────────────────────────────────────────────
+    # A. INTERMEDIATE CHECKPOINTS FAST EVALUATION
+    # ─────────────────────────────────────────────────────────────
+    print(f"[Eval] Running zero-shot fast evaluations on intermediate checkpoints...")
+    checkpoints = [f"chck_{i}M" for i in range(1, 10)] + [f"chck_{i}M" for i in range(10, 110, 10)]
+    
+    for checkpoint in checkpoints:
+        ckpt_full_path = os.path.join(model_dir, checkpoint)
+        if not os.path.exists(ckpt_full_path):
+            continue
+            
+        # blimp fast
+        run_task_with_cache(
+            checkpoint, "blimp", "blimp_fast",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", model_dir, "--backend", "causal", "--task", "blimp", "--data_path", "evaluation_data/fast_eval/blimp_fast", "--save_predictions", "--revision_name", checkpoint]
+        )
+        # supplement fast
+        run_task_with_cache(
+            checkpoint, "blimp", "supplement_fast",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", model_dir, "--backend", "causal", "--task", "blimp", "--data_path", "evaluation_data/fast_eval/supplement_fast", "--save_predictions", "--revision_name", checkpoint]
+        )
+        # ewok fast
+        run_task_with_cache(
+            checkpoint, "ewok", "ewok_fast",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", model_dir, "--backend", "causal", "--task", "ewok", "--data_path", "evaluation_data/fast_eval/ewok_fast", "--save_predictions", "--revision_name", checkpoint]
+        )
+        # entity tracking fast
+        run_task_with_cache(
+            checkpoint, "entity_tracking", "entity_tracking_fast",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", model_dir, "--backend", "causal", "--task", "entity_tracking", "--data_path", "evaluation_data/fast_eval/entity_tracking_fast", "--save_predictions", "--revision_name", checkpoint]
+        )
+        # reading fast
+        run_task_with_cache(
+            checkpoint, "reading", "reading",
+            ["python", "-m", "evaluation_pipeline.reading.run", "--model_path_or_name", model_dir, "--backend", "causal", "--data_path", "evaluation_data/fast_eval/reading/reading_data.csv", "--revision_name", checkpoint]
+        )
+        # global piqa parallel fast
+        run_task_with_cache(
+            checkpoint, "global_piqa_parallel", "global_piqa_parallel",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", model_dir, "--backend", "causal", "--task", "global_piqa_parallel", "--data_path", "evaluation_data/fast_eval/global_piqa_parallel", "--save_predictions", "--revision_name", checkpoint]
+        )
+        # global piqa nonparallel fast
+        run_task_with_cache(
+            checkpoint, "global_piqa_nonparallel", "global_piqa_nonparallel",
+            ["python", "-m", "evaluation_pipeline.sentence_zero_shot.run", "--model_path_or_name", model_dir, "--backend", "causal", "--task", "global_piqa_nonparallel", "--data_path", "evaluation_data/fast_eval/global_piqa_nonparallel", "--save_predictions", "--revision_name", checkpoint]
+        )
+
+    # ─────────────────────────────────────────────────────────────
+    # D. COLLATE RESULTS AND CLEANUP
+    # ─────────────────────────────────────────────────────────────
+    print("[Eval] Collating predictions into submission file...")
+    # Clean collation destination in evaluation repo
+    collate_results_dir = os.path.join(strict_dir, "results", model_name)
+    if os.path.exists(collate_results_dir):
+        shutil.rmtree(collate_results_dir)
+    os.makedirs(os.path.dirname(collate_results_dir), exist_ok=True)
+    
+    # Copy from local cache results to strict results for collation
+    shutil.copytree(local_results_dir, collate_results_dir)
+
+    # Run collation
+    subprocess.run([
+        "python", "-m", "evaluation_pipeline.collate_preds",
+        "--model_path_or_name", model_name,
+        "--backend", "causal",
+        "--track", "strict-small",
+        "--fast"
+    ], cwd=strict_dir, check=True)
+    
+    # Save results to local folder
+    results_src = os.path.join(strict_dir, "results")
+    results_dest = os.path.abspath("./results")
+    if os.path.exists(results_dest):
+        shutil.rmtree(results_dest)
+    shutil.copytree(results_src, results_dest)
+    
+    # Copy final collated json to current folder
+    collated_json = os.path.join(strict_dir, "all_full_preds_and_fast_scores_causal.json")
+    if os.path.exists(collated_json):
+        shutil.copy(collated_json, "./all_full_preds_and_fast_scores_causal.json")
+        print("\n[Eval] Success! Collation completed! Final file is at './all_full_preds_and_fast_scores_causal.json'")
+
+    print("\n[Eval] Pipeline evaluation run finished.")
+
+def upload_pipeline(model_name, repo_name, token=None):
+    import os
+    import shutil
+    import json
+    import hashlib
+    from huggingface_hub import HfApi, create_repo
+
+    if not token:
+        token = os.environ.get("HF_TOKEN")
+
+    api = HfApi(token=token)
+    try:
+        user_info = api.whoami()
+        username = user_info["name"]
+        print(f"[HF] Authenticated successfully as user: {username}")
+    except Exception as e:
+        print(f"[HF] Authentication failed. Error: {e}")
+        return
+
+    repo_id = f"{username}/{repo_name}"
+    print(f"[HF] Target Repository ID: {repo_id}")
+    
+    # Create the repository if it doesn't exist
+    try:
+        create_repo(repo_id=repo_id, repo_type="model", token=token, exist_ok=True)
+        print(f"[HF] Repository '{repo_id}' is ready.")
+    except Exception as e:
+        print(f"[HF] Failed to verify or create repository. Error: {e}")
+        return
+
+    checkpoint_dir = os.path.abspath(f"./checkpoints/{model_name}")
+    
+    # Resolve the main checkpoint directory using self-healing rules
+    revisions = {}
+    if os.path.exists("pytorch_model.bin") and os.path.exists("config.json"):
+        print("[HF] Detected weight and config files in the current working directory. Using current folder as 'main' checkpoint.")
+        revisions = {"main": os.getcwd()}
+    elif os.path.exists(os.path.join(checkpoint_dir, "main")):
+        revisions = {"main": os.path.join(checkpoint_dir, "main")}
+    elif os.path.exists(os.path.abspath("./checkpoints/msit_gptbert_fresh/main")):
+        print("[HF] Using fallback checkpoint folder './checkpoints/msit_gptbert_fresh/main'...")
+        revisions = {"main": os.path.abspath("./checkpoints/msit_gptbert_fresh/main")}
+    elif os.path.exists(os.path.abspath("./checkpoints/main")):
+        revisions = {"main": os.path.abspath("./checkpoints/main")}
+    else:
+        print(f"[HF] Error: Could not locate the 'main' checkpoint weights. Checked: {checkpoint_dir}/main, current directory, and fallbacks.")
+        return
+
+    # Check for intermediate checkpoints relative to the main checkpoint's parent folder
+    main_dir = revisions["main"]
+    parent_dir = os.path.dirname(main_dir)
+    
+    for m in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 15, 20, 25, 30, 35, 40, 45, 50, 60, 70, 80, 90, 100]:
+        ckpt_name = f"chck_{m}M"
+        ckpt_path = os.path.join(parent_dir, ckpt_name)
+        if os.path.exists(ckpt_path):
+            revisions[ckpt_name] = ckpt_path
+        else:
+            # Check if there is a .pt file in the parent folder
+            pt_path = os.path.join(parent_dir, f"{ckpt_name}.pt")
+            if os.path.exists(pt_path):
+                print(f"[HF] Found legacy checkpoint file '{ckpt_name}.pt'. Converting to HF format for upload...")
+                import torch
+                from transformers import AutoTokenizer
+                from tokenizers import Tokenizer
+                try:
+                    from modeling_xpertgpt import XpertGPTForCausalLM, XpertGPTConfig
+                    cfg = XpertGPTConfig(d_model=256, d_thin=384, num_layers=6, num_blocks=4)
+                    model_to_save = XpertGPTForCausalLM(cfg)
+                    sd = torch.load(pt_path, map_location="cpu")
+                    clean_sd = {}
+                    for k, v in sd.items():
+                        new_k = k.replace("module.", "")
+                        clean_sd[new_k] = v
+                    model_to_save.load_state_dict(clean_sd)
+                    
+                    vocab_path = os.path.join(parent_dir, "bpe_vocab_16k.json")
+                    if not os.path.exists(vocab_path):
+                        vocab_path = os.path.join(main_dir, "bpe_vocab_16k.json")
+                    if os.path.exists(vocab_path):
+                        tok = Tokenizer.from_file(vocab_path)
+                    else:
+                        tok = None
+                        
+                    os.makedirs(ckpt_path, exist_ok=True)
+                    state_dict = model_to_save.state_dict()
+                    new_state_dict = {}
+                    for k, v in state_dict.items():
+                        name = k
+                        if name.startswith("_orig_mod."):
+                            name = name[10:]
+                        if name.startswith("model."):
+                            name = name[6:]
+                        if name == "lm_head.weight":
+                            new_state_dict["lm_head.weight"] = v
+                        else:
+                            new_state_dict[f"transformer.{name}"] = v
+                    torch.save(new_state_dict, os.path.join(ckpt_path, "pytorch_model.bin"))
+                    
+                    config_dict = {
+                        "auto_map": {
+                            "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
+                            "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
+                            "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
+                        },
+                        "vocab_size": 16384,
+                        "block_size": 512,
+                        "d_model": 256,
+                        "hidden_size": 256,
+                        "d_thin": 384,
+                        "num_layers": 6,
+                        "num_blocks": 4,
+                        "capacity_factor": 2.0,
+                        "dropout": 0.1,
+                        "model_type": "xpertgpt",
+                        "num_hidden_layers": 6
+                    }
+                    with open(os.path.join(ckpt_path, "config.json"), "w") as f:
+                        json.dump(config_dict, f, indent=2)
+                        
+                    if tok:
+                        from transformers import PreTrainedTokenizerFast
+                        fast_tokenizer = PreTrainedTokenizerFast(
+                            tokenizer_object=tok,
+                            bos_token="[CLS]",
+                            eos_token="[SEP]",
+                            unk_token="[UNK]",
+                            pad_token="[PAD]",
+                            mask_token="[MASK]"
+                        )
+                        fast_tokenizer.save_pretrained(ckpt_path)
+                    
+                    revisions[ckpt_name] = ckpt_path
+                except Exception as ex:
+                    print(f"[HF] Failed to convert legacy checkpoint file '{ckpt_name}.pt': {ex}")
+
+    # Temporary directory for staging uploads
+    temp_dir = os.path.abspath("./temp_hf_upload")
+    
+    # LICENSE text
+    license_text = """Creative Commons Attribution-NonCommercial 4.0 International Public License
+
+By exercising the Licensed Rights (defined below), You accept and agree to be bound by the terms and conditions of this Creative Commons Attribution-NonCommercial 4.0 International Public License ("Public License"). To the extent this Public License may be interpreted as a contract, You are granted the Licensed Rights in consideration of Your acceptance of these terms and conditions, and the Licensor grants You such rights in consideration of benefits the Licensor receives from making the Licensed Material available under these terms and conditions.
+
+Section 1 -- Definitions.
+a. Licensed Material means the artistic or literary work, database, or other material to which the Licensor applied this Public License.
+b. Licensed Rights means the rights granted to You subject to the terms and conditions of this Public License, which are limited to all Copyright and Similar Rights that apply to Your use of the Licensed Material and that the Licensor has authority to license.
+c. NonCommercial means not primarily intended for or directed towards commercial advantage or monetary compensation.
+d. Share means to provide material to the public by any means or process that requires permission under the Licensed Rights.
+e. You means the individual or entity exercising the Licensed Rights under this Public License.
+
+Section 2 -- Scope.
+a. License grant.
+   1. Subject to the terms and conditions of this Public License, the Licensor hereby grants You a worldwide, royalty-free, non-sublicensable, non-exclusive, irrevocable license to exercise the Licensed Rights in the Licensed Material to:
+      A. reproduce and Share the Licensed Material, in whole or in part, for NonCommercial purposes only; and
+      B. Produce, reproduce, and Share Adapted Material for NonCommercial purposes only.
+   2. Attribution. As a condition of the license, You must attribute the Licensor and keep intact copyright notices.
+"""
+
+    for revision_name, local_path in revisions.items():
+        print(f"\n[HF] Staging files for revision '{revision_name}' from '{local_path}'...")
+        if os.path.exists(temp_dir):
+            shutil.rmtree(temp_dir)
+        os.makedirs(temp_dir)
+
+        # 1. Copy weight file and save as model.safetensors if possible, otherwise pytorch_model.bin
+        src_bin = os.path.join(local_path, "pytorch_model.bin")
+        weight_file_dest = None
+        if os.path.exists(src_bin):
+            # Try to convert to safetensors
+            try:
+                import torch
+                from safetensors.torch import save_file
+                state_dict = torch.load(src_bin, map_location="cpu")
+                # Clone tensors to break memory sharing (prevents shared weight memory error in safetensors)
+                state_dict = {k: v.clone() for k, v in state_dict.items()}
+                weight_file_dest = os.path.join(temp_dir, "model.safetensors")
+                save_file(state_dict, weight_file_dest)
+                print(f"[HF] Converted weights to safetensors format.")
+            except Exception as e:
+                print(f"[HF] Conversion to safetensors failed ({e}). Staging raw pytorch_model.bin...")
+                weight_file_dest = os.path.join(temp_dir, "pytorch_model.bin")
+                shutil.copy2(src_bin, weight_file_dest)
+        else:
+            print(f"[HF] Error: No weight file found in '{local_path}'!")
+            continue
+
+        # Calculate weight file hash
+        sha256_hash = hashlib.sha256()
+        with open(weight_file_dest, "rb") as f:
+            for byte_block in iter(lambda: f.read(4096), b""):
+                sha256_hash.update(byte_block)
+        weight_hash = sha256_hash.hexdigest()
+        print(f"[HF] Weight file SHA-256: {weight_hash}")
+
+        # 2. Copy and patch config.json
+        src_config = os.path.join(local_path, "config.json")
+        if os.path.exists(src_config):
+            with open(src_config, "r") as f:
+                cfg_data = json.load(f)
+            # Patch config
+            cfg_data["auto_map"] = {
+                "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
+                "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
+                "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
+            }
+            cfg_data["architectures"] = ["XpertGPTForCausalLM"]
+            with open(os.path.join(temp_dir, "config.json"), "w") as f:
+                json.dump(cfg_data, f, indent=2)
+        else:
+            # Fallback configuration
+            cfg_data = {
+                "auto_map": {
+                    "AutoConfig": "configuration_xpertgpt.XpertGPTConfig",
+                    "AutoModel": "modeling_xpertgpt.XpertGPTModelWrapper",
+                    "AutoModelForCausalLM": "modeling_xpertgpt.XpertGPTForCausalLM"
+                },
+                "architectures": ["XpertGPTForCausalLM"],
+                "vocab_size": 16384,
+                "block_size": 512,
+                "d_model": 256,
+                "hidden_size": 256,
+                "d_thin": 384,
+                "num_layers": 6,
+                "num_blocks": 4,
+                "capacity_factor": 2.0,
+                "dropout": 0.1,
+                "model_type": "xpertgpt"
+            }
+            with open(os.path.join(temp_dir, "config.json"), "w") as f:
+                json.dump(cfg_data, f, indent=2)
+
+        # 3. Copy tokenizers
+        for tok_file in ["tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"]:
+            src_tok = os.path.join(local_path, tok_file)
+            if os.path.exists(src_tok):
+                shutil.copy2(src_tok, os.path.join(temp_dir, tok_file))
+
+        # 4. Copy custom code
+        shutil.copy2("modeling_xpertgpt.py", os.path.join(temp_dir, "modeling_xpertgpt.py"))
+        shutil.copy2("configuration_xpertgpt.py", os.path.join(temp_dir, "configuration_xpertgpt.py"))
+
+        # 5. Write metadata files
+        with open(os.path.join(temp_dir, ".gitattributes"), "w") as f:
+            f.write("*.safetensors filter=lfs diff=lfs merge=lfs -text\n")
+            f.write("*.bin filter=lfs diff=lfs merge=lfs -text\n")
+
+        with open(os.path.join(temp_dir, "LICENSE"), "w") as f:
+            f.write(license_text)
+
+        with open(os.path.join(temp_dir, "CITATION.cff"), "w") as f:
+            citation_yaml = f"""cff-version: 1.2.0
+message: "If you use this model or software, please cite it as below."
+authors:
+  - family-names: "Jain"
+    given-names: "Soham"
+  - family-names: "Singh"
+    given-names: "Harsh"
+  - family-names: "Dewan"
+    given-names: "Divija"
+  - family-names: "Dev"
+    given-names: "Atul"
+title: "XpertGPT: Mixture of Experts with Parallelized Multi-Scale Information Transmission for Data-Constrained Pretraining"
+year: 2026
+url: "https://huggingface.co/{repo_id}"
+"""
+            f.write(citation_yaml)
+
+        with open(os.path.join(temp_dir, "PROVENANCE.md"), "w") as f:
+            provenance_md = f"""# Provenance and Weight Integrity Record
+
+This file records the provenance, cryptographic hash, and reproducibility metadata of the model weights.
+
+## Verification Fingerprints
+- **Model Weight File**: {"model.safetensors" if weight_file_dest.endswith(".safetensors") else "pytorch_model.bin"}
+- **Weight SHA-256**: {weight_hash}
+- **Tokenizer Vocab Size**: 16,384
+
+## Training Run Details
+- **Training Word Budget**: 10M words (BabyLM 2026 Strict-Small track)
+- **Model Parameters**: ~48.4M non-embedding parameters / ~52.6M total tied parameters
+- **Optimizer**: AdamW
+- **Epochs**: 8
+"""
+            f.write(provenance_md)
+
+        # Build README
+        readme_md = f"""---
+license: cc-by-nc-4.0
+language:
+  - en
+tags:
+  - babylm
+  - babylm-2026
+  - mixture-of-experts
+  - msit
+  - xpertgpt
+  - custom_code
+  - safetensors
+library_name: transformers
+pipeline_tag: text-generation
+---
+
+# XpertGPT Strict-Small
+
+XpertGPT (Mixture of Experts with Parallelized Multi-Scale Information Transmission) is a sparse, data-efficient recurrent language model for the BabyLM 2026 challenge (Strict-Small (10M) track, 10M words). It combines sliding window attention global streams with sparse parallel expert blocks using Expert Choice routing. ~48.4M non-embedding parameters / ~52.6M total parameters. Custom code (`trust_remote_code=True`).
+
+- **Architecture:** 6 layers of MoEP-MSIT blocks. Each block combines a lower-dimensional dense global sliding window attention layer (`dim = 256`) with 4 parallel high-dimensional sparse expert blocks (`dim = 384`) routed via Expert Choice gating.
+- **Track:** BabyLM 2026 Strict-Small (10M) (10M words).
+- **Tokenizer:** Custom BPE tokenizer (vocab size: 16384).
+- **Revision / Checkpoint:** {revision_name}
+
+## Usage
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained("{repo_id}", revision="{revision_name}", trust_remote_code=True).eval()
+tok = AutoTokenizer.from_pretrained("{repo_id}", revision="{revision_name}")
+ids = tok("The quick brown fox", return_tensors="pt").input_ids
+with torch.no_grad():
+    logits = model(ids).logits
+```
+
+## Intermediate checkpoints
+
+Intermediate training checkpoints are provided as git revisions named `chck_<N>M` for the BabyLM challenge fast-eval.
+
+## License and citation
+
+Released under CC BY-NC 4.0 (attribution required, non-commercial only). If you use this model or code, please cite (see `CITATION.cff`):
+
+```bibtex
+@misc{{jain2026xpertgpt,
+  title        = {{XpertGPT: Mixture of Experts with Parallelized Multi-Scale Information Transmission for Data-Constrained Pretraining}},
+  author       = {{Jain, Soham and Singh, Harsh and Dewan, Divija and Dev, Atul}},
+  year         = {{2026}},
+  howpublished = {{Hugging Face Repository}},
+  note         = {{XpertGPT MoE language model, BabyLM 2026}}
+}}
+```
+
+Provenance and integrity fingerprints are documented in `PROVENANCE.md`.
+"""
+        with open(os.path.join(temp_dir, "README.md"), "w") as f:
+            f.write(readme_md)
+
+        # 6. For main branch only: also upload collated predictions
+        if revision_name == "main":
+            for pred_file in ["all_full_preds_and_fast_scores_causal.json", "all_full_preds_and_fast_scores_causal (3).json"]:
+                if os.path.exists(pred_file):
+                    shutil.copy2(pred_file, os.path.join(temp_dir, "all_full_preds_and_fast_scores_causal.json"))
+                    print(f"[HF] Copied predictions file '{pred_file}' to staging area.")
+                    break
+
+        # 7. Create branch if it does not exist, then upload staged files to HF under the revision
+        if revision_name != "main":
+            try:
+                api.create_branch(
+                    repo_id=repo_id,
+                    repo_type="model",
+                    branch=revision_name,
+                    exist_ok=True
+                )
+                print(f"[HF] Created branch/revision '{revision_name}' on repository.")
+            except Exception as branch_err:
+                print(f"[HF] Info: Branch creation failed or exists: {branch_err}")
+
+        print(f"[HF] Uploading staged folder to '{repo_id}' revision '{revision_name}'...")
+        try:
+            api.upload_folder(
+                folder_path=temp_dir,
+                repo_id=repo_id,
+                repo_type="model",
+                revision=revision_name
+            )
+            print(f"[HF] Successfully uploaded revision '{revision_name}' to repository.")
+        except Exception as e:
+            print(f"[HF] Failed to upload revision '{revision_name}': {e}")
+
+    # Cleanup temp dir
+    if os.path.exists(temp_dir):
+        shutil.rmtree(temp_dir)
+    print(f"\n[HF] All uploads finished! View your repository at https://huggingface.co/{repo_id}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-name", type=str, default="xpertgpt_fresh")
+    parser.add_argument("--epochs", type=int, default=10)
+    parser.add_argument("--skip-eval", action="store_true", help="Skip evaluation phase after training")
+    parser.add_argument("--skip-aoa", action="store_true", default=True, help="Skip AoA evaluation")
+    parser.add_argument("--skip-glue", action="store_true", default=False, help="Skip GLUE fine-tuning")
+    parser.add_argument("--upload", action="store_true", help="Upload model repository to Hugging Face")
+    parser.add_argument("--upload-repo", type=str, default="XpertGPT-BabyLM2026-Strict-Small", help="Hugging Face repository name")
+    parser.add_argument("--upload-token", type=str, default=None, help="Hugging Face API token")
+    args = parser.parse_args()
+    
+    if args.upload:
+        upload_pipeline(args.model_name, args.upload_repo, args.upload_token)
+    else:
+        run_pipeline(args.model_name, epochs=args.epochs, skip_eval=args.skip_eval, skip_aoa=args.skip_aoa, skip_glue=args.skip_glue)

upload_project_hf.py

@@ -0,0 +1,78 @@
+import os
+import sys
+import re
+import argparse
+from huggingface_hub import HfApi
+
+def upload_project(repo_name=None):
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        print("Error: HF_TOKEN environment variable not set!")
+        print("Please set it before running: set HF_TOKEN=your_token")
+        sys.exit(1)
+
+    if not repo_name:
+        repo_name = os.environ.get("HF_REPO_NAME", "temp")
+
+    local_dir = os.path.dirname(os.path.abspath(__file__))
+
+    api = HfApi(token=token)
+    try:
+        user_info = api.whoami()
+        username = user_info["name"]
+        print(f"[HF] Authenticated as: {username}")
+    except Exception as e:
+        print(f"[HF] Authentication failed: {e}")
+        sys.exit(1)
+
+    repo_id = f"{username}/{repo_name}"
+    print(f"[HF] Uploading modified sliding-window strict-small scripts to '{repo_id}' main branch...")
+    
+    ignore_patterns = [
+        "**/__pycache__/*",
+        "**/*.pyc",
+        "**/hf_cache/*",
+        "**/nltk_data/*",
+        "**/checkpoints/*"
+    ]
+    
+    sensitive_ignores = []
+    token_pattern = re.compile(r"hf_[a-zA-Z0-9]{34}")
+    for root, dirs, files in os.walk(local_dir):
+        if "__pycache__" in root or "checkpoints" in root:
+            continue
+        for file in files:
+            file_path = os.path.join(root, file)
+            if file.endswith(('.bin', '.safetensors', '.zip', '.tar.gz', '.pkl')):
+                continue
+            try:
+                with open(file_path, "r", errors="ignore") as f:
+                    content = f.read()
+                    if token_pattern.search(content):
+                        rel_path = os.path.relpath(file_path, local_dir)
+                        rel_path_glob = rel_path.replace("\\", "/")
+                        sensitive_ignores.append(rel_path_glob)
+                        print(f"  -> Warning: Sensitive file containing raw token excluded: {rel_path_glob}")
+            except Exception:
+                pass
+
+    all_ignores = ignore_patterns + sensitive_ignores
+
+    try:
+        api.upload_folder(
+            folder_path=local_dir,
+            repo_id=repo_id,
+            repo_type="model",
+            revision="main",
+            ignore_patterns=all_ignores,
+            commit_message="Update strict-small architecture files (sliding window [64, 16, 8, 4], ln3, ln_post_moe, no res3)"
+        )
+        print(f"\n[HF] Success! Scripts uploaded to: https://huggingface.co/{repo_id}/tree/main")
+    except Exception as e:
+        print(f"[HF] Upload failed: {e}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repo-name", type=str, default=None, help="Hugging Face repository name")
+    args = parser.parse_args()
+    upload_project(args.repo_name)