Upload 5 files

Wikitext mlm model training codes. All baselines and experimental architectures.

FNet_Hybrid_Wikitext_Training.ipynb

@@ -0,0 +1,600 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "A100"
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "code",
+      "source": [
+        "!pip install -q x-transformers\n",
+        "!pip install -q flash-attn --no-build-isolation"
+      ],
+      "metadata": {
+        "id": "6q9RTvlf5IiS"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import torch\n",
+        "import torch.nn as nn\n",
+        "import torch.nn.functional as F\n",
+        "import torch.optim as optim\n",
+        "import math\n",
+        "import os\n",
+        "import sys\n",
+        "import subprocess\n",
+        "import hashlib\n",
+        "import gc\n",
+        "import platform\n",
+        "from datetime import datetime\n",
+        "from tqdm.auto import tqdm\n",
+        "from torch.utils.data import DataLoader\n",
+        "from torch.utils.tensorboard import SummaryWriter\n",
+        "from transformers import RobertaTokenizerFast, get_cosine_schedule_with_warmup, DataCollatorForLanguageModeling\n",
+        "from datasets import load_dataset\n",
+        "from x_transformers import Encoder\n",
+        "\n",
+        "# ==========================================\n",
+        "# 1. CONFIGURATION\n",
+        "# ==========================================\n",
+        "# YOUR REPO ID (Created in previous step)\n",
+        "HF_ID = \"prism-lab/wikitext-103-prism-32k-seq4k\"\n",
+        "\n",
+        "# Hyperparameters\n",
+        "VOCAB_SIZE = 32768\n",
+        "SEQ_LEN = 4096\n",
+        "BATCH_SIZE = 8\n",
+        "EPOCHS = 40\n",
+        "LR = 1e-3\n",
+        "D_MODEL = 512\n",
+        "RESUME_PATH = None\n",
+        "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+        "torch.set_float32_matmul_precision(\"high\")\n",
+        "\n",
+        "# ==========================================\n",
+        "# 2. DATA PIPELINE (The \"Pro\" Way)\n",
+        "# ==========================================\n",
+        "def prepare_data_from_hub():\n",
+        "    print(f\"⬇️ Pulling Pre-Tokenized Data from {HF_ID}...\")\n",
+        "\n",
+        "    # 1. Load Tokenizer (Instant)\n",
+        "    # This pulls the exact tokenizer you uploaded\n",
+        "    tokenizer = RobertaTokenizerFast.from_pretrained(HF_ID)\n",
+        "\n",
+        "    # 2. Load Dataset (Instant)\n",
+        "    # This pulls the already chunked/tokenized data\n",
+        "    dataset = load_dataset(HF_ID)\n",
+        "\n",
+        "    print(f\"✅ Loaded {len(dataset['train'])} training chunks.\")\n",
+        "\n",
+        "    # 3. Collator\n",
+        "    data_collator = DataCollatorForLanguageModeling(\n",
+        "        tokenizer=tokenizer,\n",
+        "        mlm=True,\n",
+        "        mlm_probability=0.15\n",
+        "    )\n",
+        "\n",
+        "    return dataset, data_collator\n",
+        "\n",
+        "\n",
+        "class FNetBlock(nn.Module):\n",
+        "    def __init__(self, d_model, d_ff, dropout):\n",
+        "        super().__init__()\n",
+        "        self.norm_mix = nn.LayerNorm(d_model) # LayerNorm is safer for FNet than RMSNorm\n",
+        "        self.norm_ff = nn.LayerNorm(d_model)\n",
+        "\n",
+        "        self.mix_dropout = nn.Dropout(dropout)\n",
+        "\n",
+        "        self.ff = nn.Sequential(\n",
+        "            nn.Linear(d_model, d_ff),\n",
+        "            nn.GELU(),\n",
+        "            nn.Dropout(dropout),\n",
+        "            nn.Linear(d_ff, d_model),\n",
+        "            nn.Dropout(dropout)\n",
+        "        )\n",
+        "\n",
+        "    def forward(self, x):\n",
+        "        # 1. Fourier Mixing Branch\n",
+        "        residual = x\n",
+        "        x = self.norm_mix(x)\n",
+        "\n",
+        "        # --- THE FIX ---\n",
+        "        with torch.cuda.amp.autocast(enabled=False):\n",
+        "            x = x.float()\n",
+        "            # norm='ortho' makes the FFT energy-preserving.\n",
+        "            # Output magnitude will match input magnitude (~1).\n",
+        "            x = torch.fft.fftn(x, dim=(-2, -1), norm='ortho').real\n",
+        "            x = x.to(dtype=residual.dtype)\n",
+        "        # ---------------\n",
+        "\n",
+        "        # Now 'x' and 'residual' have roughly same magnitude.\n",
+        "        # The skip connection works again.\n",
+        "        x = self.mix_dropout(x)\n",
+        "        x = x + residual\n",
+        "\n",
+        "        # 2. Feed Forward Branch\n",
+        "        residual = x\n",
+        "        x = self.norm_ff(x)\n",
+        "        x = self.ff(x)\n",
+        "        return x + residual\n",
+        "\n",
+        "\n",
+        "class FNetEncoder(nn.Module):\n",
+        "    def __init__(self, depth, d_model, d_ff, dropout):\n",
+        "        super().__init__()\n",
+        "        self.layers = nn.ModuleList([\n",
+        "            FNetBlock(d_model, d_ff, dropout) for _ in range(depth)\n",
+        "        ])\n",
+        "        # [FIX] Use LayerNorm here to match the blocks\n",
+        "        self.norm_out = nn.LayerNorm(d_model)\n",
+        "\n",
+        "    def forward(self, x):\n",
+        "        for layer in self.layers:\n",
+        "            x = layer(x)\n",
+        "        return self.norm_out(x)\n",
+        "\n",
+        "\n",
+        "class HybridFNetMLM(nn.Module):\n",
+        "    def __init__(self, vocab_size, d_model, seq_len, d_ff, dropout):\n",
+        "        super().__init__()\n",
+        "\n",
+        "        # 1. Standard Embeddings + Absolute Positions\n",
+        "        # (FNet NEEDS these because FFT is position-blind)\n",
+        "        self.token_emb = nn.Embedding(vocab_size, d_model)\n",
+        "        self.pos_emb = nn.Parameter(torch.randn(1, seq_len, d_model) * 0.02)\n",
+        "        self.dropout = nn.Dropout(dropout)\n",
+        "\n",
+        "        self.fnet_encoder = FNetEncoder(\n",
+        "            depth=6,\n",
+        "            d_model=d_model,\n",
+        "            d_ff=d_ff,\n",
+        "            dropout=dropout\n",
+        "        )\n",
+        "\n",
+        "        # 3. The Attention Cap (1 Layer) -> YOUR CONFIGURATION\n",
+        "        self.transformer_cap = Encoder(\n",
+        "            dim=d_model,\n",
+        "            depth=1,                # Just 1 layer\n",
+        "            heads=8,\n",
+        "            rotary_pos_emb=True,    # RoPE (Hybrid Positioning: Absolute for FNet, Rotary for Attn)\n",
+        "            attn_flash=True,\n",
+        "            attn_dropout=dropout,\n",
+        "            ff_dropout=dropout\n",
+        "            # Removed 'dim_head' (fixes your error)\n",
+        "            # Removed 'use_rmsnorm' (matches your snippet)\n",
+        "            # Removed 'ff_glu' (matches your snippet)\n",
+        "        )\n",
+        "\n",
+        "        # 4. MLM Head\n",
+        "        self.final_norm = nn.LayerNorm(d_model)\n",
+        "        self.to_logits = nn.Linear(d_model, vocab_size)\n",
+        "\n",
+        "        # Weight Tying\n",
+        "        self.to_logits.weight = self.token_emb.weight\n",
+        "\n",
+        "    def forward(self, input_ids):\n",
+        "        # A. Embedding\n",
+        "        x = self.token_emb(input_ids)\n",
+        "        b, n, d = x.shape\n",
+        "\n",
+        "        # Add Absolute Positions (Crucial for FNet layers)\n",
+        "        x = x + self.pos_emb[:, :n, :]\n",
+        "        x = self.dropout(x)\n",
+        "\n",
+        "        x = self.fnet_encoder(x)\n",
+        "\n",
+        "        # C. Attention Refinement (1 Layer)\n",
+        "        # Note: This layer will internally apply RoPE to Q/K\n",
+        "        x = self.transformer_cap(x)\n",
+        "\n",
+        "        # D. Output\n",
+        "        x = self.final_norm(x)\n",
+        "        return self.to_logits(x)\n",
+        "\n",
+        "# ==========================================\n",
+        "# INSTANTIATE MODEL\n",
+        "# ==========================================\n",
+        "print(\"🏗️  Constructing Hybrid FNet (6-Spectral + 1-Attention)...\")\n",
+        "\n",
+        "def count_active_parameters(model):\n",
+        "    print(f\"\\n{'='*60}\")\n",
+        "    print(f\"🧩 DETAILED PARAMETER BREAKDOWN\")\n",
+        "    print(f\"{'='*60}\")\n",
+        "\n",
+        "    # 1. Identify Parameter Groups\n",
+        "    # ----------------------------\n",
+        "    embedding_ids = set()\n",
+        "    active_ids = set()\n",
+        "    unique_params = set()\n",
+        "\n",
+        "    # --- MEMORY (Embeddings & Encodings) ---\n",
+        "    embedding_params = 0\n",
+        "    for p in model.token_emb.parameters():\n",
+        "        embedding_params += p.numel()\n",
+        "        embedding_ids.add(id(p))\n",
+        "        unique_params.add(id(p))\n",
+        "\n",
+        "    pos_params = model.pos_emb.numel()\n",
+        "    embedding_ids.add(id(model.pos_emb))\n",
+        "    unique_params.add(id(model.pos_emb))\n",
+        "\n",
+        "    total_memory = embedding_params + pos_params\n",
+        "\n",
+        "    # --- LOGIC (Active Processing) ---\n",
+        "    active_count = 0\n",
+        "    for name, param in model.named_parameters():\n",
+        "        if id(param) in embedding_ids:\n",
+        "            continue\n",
+        "        if id(param) not in active_ids:\n",
+        "            active_count += param.numel()\n",
+        "            active_ids.add(id(param))\n",
+        "            unique_params.add(id(param))\n",
+        "\n",
+        "    # 2. Calculate Totals\n",
+        "    # -------------------\n",
+        "    total_physical_params = total_memory + active_count\n",
+        "\n",
+        "    # 3. Print Report (FIXED SYNTAX)\n",
+        "    # -------------------\n",
+        "    print(f\"{'Component':<25} | {'Count':<15} | {'% of Model':<10}\")\n",
+        "    print(f\"{'-'*60}\")\n",
+        "\n",
+        "    print(f\"{'Token Embeddings':<25} | {embedding_params:<15,} | {embedding_params/total_physical_params:.1%}\")\n",
+        "    print(f\"{'Positional Encodings':<25} | {pos_params:<15,} | {pos_params/total_physical_params:.1%}\")\n",
+        "    print(f\"{'[MEMORY TOTAL]':<25} | {total_memory:<15,} | {total_memory/total_physical_params:.1%}\")\n",
+        "    print(f\"{'-'*60}\")\n",
+        "\n",
+        "    fnet_params = sum(p.numel() for p in model.fnet_encoder.parameters())\n",
+        "    cap_params = sum(p.numel() for p in model.transformer_cap.parameters())\n",
+        "    misc_params = active_count - fnet_params - cap_params\n",
+        "\n",
+        "    print(f\"{'FNet Encoder (6 Layers)':<25} | {fnet_params:<15,} | {fnet_params/total_physical_params:.1%}\")\n",
+        "    print(f\"{'Transformer Cap (1 Layer)':<25}| {cap_params:<15,} | {cap_params/total_physical_params:.1%}\")\n",
+        "    print(f\"{'Norms & Biases':<25} | {misc_params:<15,} | {misc_params/total_physical_params:.1%}\")\n",
+        "    print(f\"{'[ACTIVE LOGIC TOTAL]':<25} | {active_count:<15,} | {active_count/total_physical_params:.1%}\")\n",
+        "\n",
+        "    print(f\"{'='*60}\")\n",
+        "    print(f\"📢 FINAL ACTIVE PARAMETERS: {active_count / 1_000_000:.2f} M\")\n",
+        "    print(f\"{'='*60}\\n\")\n",
+        "\n",
+        "    return active_count\n",
+        "\n",
+        "\n"
+      ],
+      "metadata": {
+        "id": "V7DOwmmUjyin"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# ==========================================\n",
+        "# 3. INITIALIZATION FUNCTION (FNet Specific)\n",
+        "# ==========================================\n",
+        "def init_fnet_weights(model):\n",
+        "    print(\"✨ Applying BERT-Style Initialization (N(0, 0.02))...\")\n",
+        "\n",
+        "    for name, module in model.named_modules():\n",
+        "        # A. Linear Layers (Projections, FFNs)\n",
+        "        if isinstance(module, nn.Linear):\n",
+        "            module.weight.data.normal_(mean=0.0, std=0.02)\n",
+        "            if module.bias is not None:\n",
+        "                module.bias.data.zero_()\n",
+        "\n",
+        "        # B. Embeddings (Tokens)\n",
+        "        elif isinstance(module, nn.Embedding):\n",
+        "            module.weight.data.normal_(mean=0.0, std=0.02)\n",
+        "            if module.padding_idx is not None:\n",
+        "                module.weight.data[module.padding_idx].zero_()\n",
+        "\n",
+        "        # C. LayerNorms (Stability)\n",
+        "        elif isinstance(module, nn.LayerNorm):\n",
+        "            if module.bias is not None:  # <--- FIX IS HERE\n",
+        "                module.bias.data.zero_()\n",
+        "            if module.weight is not None:\n",
+        "                module.weight.data.fill_(1.0)\n",
+        "\n",
+        "    # D. Positional Embeddings (Manually handle the nn.Parameter)\n",
+        "    if hasattr(model, 'pos_emb') and model.pos_emb is not None:\n",
+        "        model.pos_emb.data.normal_(mean=0.0, std=0.02)\n",
+        "\n",
+        "    print(\"✅ Initialization Complete.\")\n",
+        "\n",
+        "\n",
+        "# ==========================================\n",
+        "# 4. LOGGING UTILITIES\n",
+        "# ==========================================\n",
+        "def generate_run_id():\n",
+        "    raw = datetime.now().strftime(\"%Y%m%d%H%M%S%f\")\n",
+        "    return hashlib.md5(raw.encode()).hexdigest()[:8]\n",
+        "\n",
+        "def log_full_environment(save_dir, run_id, config):\n",
+        "    log_path = os.path.join(save_dir, f\"env_metadata_{run_id}.txt\")\n",
+        "\n",
+        "    # 1. Gather System Info\n",
+        "    sys_info = {\n",
+        "        \"Python Version\": sys.version.split()[0],\n",
+        "        \"OS\": platform.platform(),\n",
+        "        \"PyTorch Version\": torch.__version__,\n",
+        "        \"CUDA Available\": torch.cuda.is_available(),\n",
+        "        \"CUDNN Version\": torch.backends.cudnn.version() if torch.cuda.is_available() else \"N/A\"\n",
+        "    }\n",
+        "\n",
+        "    # 2. Gather GPU Info\n",
+        "    gpu_info = []\n",
+        "    if torch.cuda.is_available():\n",
+        "        for i in range(torch.cuda.device_count()):\n",
+        "            props = torch.cuda.get_device_properties(i)\n",
+        "            gpu_info.append(f\"GPU {i}: {props.name} | VRAM: {props.total_memory / 1e9:.2f} GB\")\n",
+        "    else:\n",
+        "        gpu_info.append(\"No GPU Detected\")\n",
+        "\n",
+        "    # 3. Gather Pip Freeze\n",
+        "    try:\n",
+        "        pip_packages = subprocess.check_output([sys.executable, '-m', 'pip', 'freeze']).decode('utf-8')\n",
+        "    except Exception as e:\n",
+        "        pip_packages = f\"Could not retrieve pip packages: {e}\"\n",
+        "\n",
+        "    # 4. Write to File\n",
+        "    with open(log_path, \"w\") as f:\n",
+        "        f.write(f\"🧪 EXPERIMENT METADATA | Run ID: {run_id}\\n\")\n",
+        "        f.write(f\"{'='*60}\\n\\n\")\n",
+        "\n",
+        "        f.write(f\"--- [1] CONFIGURATION ---\\n\")\n",
+        "        for k, v in config.items():\n",
+        "            f.write(f\"{k}: {v}\\n\")\n",
+        "        f.write(\"\\n\")\n",
+        "\n",
+        "        f.write(f\"--- [2] SYSTEM HARDWARE ---\\n\")\n",
+        "        for k, v in sys_info.items():\n",
+        "            f.write(f\"{k}: {v}\\n\")\n",
+        "        for g in gpu_info:\n",
+        "            f.write(f\"{g}\\n\")\n",
+        "        f.write(\"\\n\")\n",
+        "\n",
+        "        f.write(f\"--- [3] INSTALLED PACKAGES (pip freeze) ---\\n\")\n",
+        "        f.write(pip_packages)\n",
+        "\n",
+        "    print(f\"📝 Full Environment Snapshot (GPU + Pip) saved to: {log_path}\")\n",
+        "\n",
+        "\n",
+        "def save_checkpoint(path, model, optimizer, scheduler, epoch, best_loss, config):\n",
+        "    torch.save({\n",
+        "        'epoch': epoch,\n",
+        "        'model_state_dict': model.state_dict(),\n",
+        "        'optimizer_state_dict': optimizer.state_dict(),\n",
+        "        'scheduler_state_dict': scheduler.state_dict(),\n",
+        "        'best_val_loss': best_loss,\n",
+        "        'config': config\n",
+        "    }, path)\n",
+        "\n",
+        "# ==========================================\n",
+        "# 5. TRAINING LOOP\n",
+        "# ==========================================\n",
+        "def run_wikitext_training(experiment_name=\"FNet_Encoder\"):\n",
+        "    from google.colab import drive\n",
+        "    if not os.path.exists('/content/drive'): drive.mount('/content/drive')\n",
+        "\n",
+        "    # --- SETUP DIRS ---\n",
+        "    if RESUME_PATH and os.path.exists(RESUME_PATH):\n",
+        "        print(f\"🔄 RESUMING FROM: {RESUME_PATH}\")\n",
+        "        checkpoint = torch.load(RESUME_PATH, map_location=DEVICE)\n",
+        "        SAVE_DIR = os.path.dirname(RESUME_PATH)\n",
+        "        run_id = checkpoint.get('config', {}).get('run_id', 'resumed')\n",
+        "    else:\n",
+        "        run_id = generate_run_id()\n",
+        "        timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
+        "        folder_name = f\"{experiment_name}_{timestamp}_{run_id}\"\n",
+        "        SAVE_DIR = os.path.join(\"/content/drive/My Drive/PRISM_Experiments\", folder_name)\n",
+        "        os.makedirs(SAVE_DIR, exist_ok=True)\n",
+        "        print(f\"💾 Checkpoints: {SAVE_DIR}\")\n",
+        "\n",
+        "    writer = SummaryWriter(log_dir=SAVE_DIR)\n",
+        "    GRAD_ACCUM = 4\n",
+        "\n",
+        "    # Load Data\n",
+        "    lm_datasets, data_collator = prepare_data_from_hub()\n",
+        "\n",
+        "    # Create Loaders\n",
+        "    train_loader = DataLoader(\n",
+        "        lm_datasets[\"train\"], batch_size=BATCH_SIZE, shuffle=True,\n",
+        "        collate_fn=data_collator, num_workers=2, pin_memory=True,\n",
+        "        prefetch_factor=2, persistent_workers=True\n",
+        "    )\n",
+        "    valid_loader = DataLoader(\n",
+        "        lm_datasets[\"validation\"], batch_size=BATCH_SIZE,\n",
+        "        collate_fn=data_collator, num_workers=2, pin_memory=True\n",
+        "    )\n",
+        "    test_loader = DataLoader(\n",
+        "        lm_datasets[\"test\"], batch_size=BATCH_SIZE,\n",
+        "        collate_fn=data_collator, num_workers=2, pin_memory=True\n",
+        "    )\n",
+        "\n",
+        "    print(\"\\n⚡ INITIALIZING HYBRID FNET MODEL...\")\n",
+        "\n",
+        "    # INSTANTIATE\n",
+        "    model = HybridFNetMLM(\n",
+        "        vocab_size=VOCAB_SIZE,\n",
+        "        d_model=D_MODEL,\n",
+        "        seq_len=SEQ_LEN,\n",
+        "        d_ff=D_MODEL * 4,\n",
+        "        dropout=0.1\n",
+        "    ).to(DEVICE)\n",
+        "\n",
+        "    # OPTIMIZER\n",
+        "    optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=0.01)\n",
+        "\n",
+        "    total_steps = (len(train_loader) // GRAD_ACCUM) * EPOCHS\n",
+        "    scheduler = get_cosine_schedule_with_warmup(\n",
+        "        optimizer, num_warmup_steps=int(0.05 * total_steps), num_training_steps=total_steps\n",
+        "    )\n",
+        "    criterion = nn.CrossEntropyLoss()\n",
+        "\n",
+        "    start_epoch = 0\n",
+        "    best_val_loss = float('inf')\n",
+        "\n",
+        "    # RESUME OR INIT\n",
+        "    if RESUME_PATH and os.path.exists(RESUME_PATH):\n",
+        "        model.load_state_dict(checkpoint['model_state_dict'])\n",
+        "        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
+        "        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])\n",
+        "        start_epoch = checkpoint['epoch'] + 1\n",
+        "        best_val_loss = checkpoint['best_val_loss']\n",
+        "        del checkpoint\n",
+        "        torch.cuda.empty_cache()\n",
+        "    else:\n",
+        "        # [UPDATED] CALL THE FNET INITIALIZATION\n",
+        "        init_fnet_weights(model)\n",
+        "\n",
+        "    # METRICS\n",
+        "    try:\n",
+        "        active_params = count_active_parameters(model) # Uses function defined in prev step\n",
+        "    except:\n",
+        "        print(\"⚠️ Parameter counter not found, skipping detailed breakdown.\")\n",
+        "\n",
+        "    total_params = sum(p.numel() for p in model.parameters())\n",
+        "    print(f\"✅ Model Ready. Total Raw Params: {total_params/1e6:.2f}M\")\n",
+        "\n",
+        "    log_full_environment(SAVE_DIR, run_id, {\n",
+        "        \"model\": \"HybridFNetMLM\",\n",
+        "        \"d_model\": D_MODEL,\n",
+        "        \"depth\": \"6+1\",\n",
+        "        \"vocab\": VOCAB_SIZE,\n",
+        "        \"batch\": BATCH_SIZE,\n",
+        "        \"lr\": LR,\n",
+        "        \"active_params\": f\"{active_params/1e6:.2f}M\"\n",
+        "    })\n",
+        "\n",
+        "    print(f\"\\n🚀 STARTING (Ep {start_epoch+1} to {EPOCHS})\")\n",
+        "    global_step = (len(train_loader) // GRAD_ACCUM) * start_epoch\n",
+        "\n",
+        "    for epoch in range(start_epoch, EPOCHS):\n",
+        "        model.train()\n",
+        "        pbar = tqdm(train_loader, desc=f\"Ep {epoch+1}/{EPOCHS}\")\n",
+        "\n",
+        "        for step, batch in enumerate(pbar):\n",
+        "            x, y = batch['input_ids'].to(DEVICE), batch['labels'].to(DEVICE)\n",
+        "\n",
+        "            # FNet Forward Pass\n",
+        "            logits = model(x)\n",
+        "\n",
+        "            # Loss Calculation\n",
+        "            loss = criterion(logits.view(-1, VOCAB_SIZE), y.view(-1)) / GRAD_ACCUM\n",
+        "            loss.backward()\n",
+        "\n",
+        "            if (step + 1) % GRAD_ACCUM == 0:\n",
+        "                # 1. Calc Norm\n",
+        "                grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
+        "\n",
+        "                # 2. Step\n",
+        "                optimizer.step()\n",
+        "                scheduler.step()\n",
+        "                optimizer.zero_grad()\n",
+        "                global_step += 1\n",
+        "\n",
+        "                # 3. LOGGING\n",
+        "                actual_loss = loss.item() * GRAD_ACCUM\n",
+        "                writer.add_scalar('Train/Loss', actual_loss, global_step)\n",
+        "                writer.add_scalar('Train/GradNorm', grad_norm.item(), global_step)\n",
+        "                writer.add_scalar('Train/LR', scheduler.get_last_lr()[0], global_step)\n",
+        "\n",
+        "                # 4. Progress Bar\n",
+        "                pbar.set_postfix({\n",
+        "                    'loss': f\"{actual_loss:.4f}\",\n",
+        "                    'gnorm': f\"{grad_norm.item():.2f}\"\n",
+        "                })\n",
+        "\n",
+        "        # VALIDATION\n",
+        "        model.eval()\n",
+        "        val_loss = 0\n",
+        "        with torch.no_grad():\n",
+        "            for batch in valid_loader:\n",
+        "                x, y = batch['input_ids'].to(DEVICE), batch['labels'].to(DEVICE)\n",
+        "                val_loss += criterion(model(x).view(-1, VOCAB_SIZE), y.view(-1)).item()\n",
+        "\n",
+        "        avg_val_loss = val_loss / len(valid_loader)\n",
+        "        ppl = math.exp(avg_val_loss) if avg_val_loss < 100 else float('inf')\n",
+        "\n",
+        "        print(f\"✨ Epoch {epoch+1} | Val Loss: {avg_val_loss:.4f} | PPL: {ppl:.2f}\")\n",
+        "        writer.add_scalar('Val/PPL', ppl, epoch+1)\n",
+        "        writer.add_scalar('Val/Loss', avg_val_loss, epoch+1)\n",
+        "\n",
+        "        config_dump = {\"epoch\": epoch, \"run_id\": run_id}\n",
+        "        save_checkpoint(os.path.join(SAVE_DIR, \"last.pt\"), model, optimizer, scheduler, epoch, best_val_loss, config_dump)\n",
+        "\n",
+        "        if avg_val_loss < best_val_loss:\n",
+        "            best_val_loss = avg_val_loss\n",
+        "            torch.save(model.state_dict(), os.path.join(SAVE_DIR, \"best.pt\"))\n",
+        "            print(\"   🏆 New Best Model Saved!\")\n",
+        "\n",
+        "    # FINAL TEST\n",
+        "    best_path = os.path.join(SAVE_DIR, \"best.pt\")\n",
+        "    if os.path.exists(best_path):\n",
+        "        model.load_state_dict(torch.load(best_path))\n",
+        "        model.eval()\n",
+        "        test_loss = 0\n",
+        "        with torch.no_grad():\n",
+        "            for batch in tqdm(test_loader, desc=\"Testing\"):\n",
+        "                x, y = batch['input_ids'].to(DEVICE), batch['labels'].to(DEVICE)\n",
+        "                test_loss += criterion(model(x).view(-1, VOCAB_SIZE), y.view(-1)).item()\n",
+        "        print(f\"🏆 FINAL TEST PPL: {math.exp(test_loss/len(test_loader)):.2f}\")\n",
+        "\n",
+        "    writer.close()\n",
+        "    return model"
+      ],
+      "metadata": {
+        "id": "-TNEv89gkS1k"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "if __name__ == \"__main__\":\n",
+        "\n",
+        "\n",
+        "    # 1. Run the Training Routine\n",
+        "    # This handles Model Creation -> Analysis -> Training -> Saving\n",
+        "    trained_prism = run_wikitext_training()\n",
+        "\n",
+        "    # 2. Cleanup & Shutdown\n"
+      ],
+      "metadata": {
+        "id": "KaiJU0tPkVp-"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "from google.colab import runtime\n",
+        "runtime.unassign()"
+      ],
+      "metadata": {
+        "id": "bxFTYWHVqcSI"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}

HSSM_Wikitext_Training.ipynb

@@ -0,0 +1,951 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "A100"
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "code",
+      "source": [
+        "!pip install -q x-transformers\n",
+        "!pip install -q flash-attn --no-build-isolation"
+      ],
+      "metadata": {
+        "id": "6q9RTvlf5IiS"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import torch\n",
+        "import torch.nn as nn\n",
+        "import torch.nn.functional as F\n",
+        "import torch.optim as optim\n",
+        "import math\n",
+        "import os\n",
+        "import sys\n",
+        "import subprocess\n",
+        "import hashlib\n",
+        "import gc\n",
+        "from datetime import datetime\n",
+        "from tqdm.auto import tqdm\n",
+        "from torch.utils.data import DataLoader\n",
+        "from torch.utils.tensorboard import SummaryWriter\n",
+        "from transformers import RobertaTokenizerFast, get_cosine_schedule_with_warmup, DataCollatorForLanguageModeling\n",
+        "from datasets import load_dataset\n",
+        "from x_transformers import Encoder\n",
+        "\n",
+        "# ==========================================\n",
+        "# 1. CONFIGURATION\n",
+        "# ==========================================\n",
+        "# YOUR REPO ID (Created in previous step)\n",
+        "HF_ID = \"prism-lab/wikitext-103-prism-32k-seq4k\"\n",
+        "\n",
+        "# Hyperparameters\n",
+        "VOCAB_SIZE = 32768\n",
+        "SEQ_LEN = 4096\n",
+        "BATCH_SIZE = 8\n",
+        "EPOCHS = 40\n",
+        "LR = 1e-3\n",
+        "D_MODEL = 512\n",
+        "D_BRANCH = 256\n",
+        "DEPTH = 9\n",
+        "RESUME_PATH = None #\"/content/drive/MyDrive/PRISM_Experiments/PILLARS_SplitStream_8Layer_20260116_025321_8438ce62/last.pt\"\n",
+        "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+        "torch.set_float32_matmul_precision(\"high\")\n",
+        "\n",
+        "# ==========================================\n",
+        "# 2. DATA PIPELINE (The \"Pro\" Way)\n",
+        "# ==========================================\n",
+        "def prepare_data_from_hub():\n",
+        "    print(f\"⬇️ Pulling Pre-Tokenized Data from {HF_ID}...\")\n",
+        "\n",
+        "    # 1. Load Tokenizer (Instant)\n",
+        "    # This pulls the exact tokenizer you uploaded\n",
+        "    tokenizer = RobertaTokenizerFast.from_pretrained(HF_ID)\n",
+        "\n",
+        "    # 2. Load Dataset (Instant)\n",
+        "    # This pulls the already chunked/tokenized data\n",
+        "    dataset = load_dataset(HF_ID)\n",
+        "\n",
+        "    print(f\"✅ Loaded {len(dataset['train'])} training chunks.\")\n",
+        "\n",
+        "    # 3. Collator\n",
+        "    data_collator = DataCollatorForLanguageModeling(\n",
+        "        tokenizer=tokenizer,\n",
+        "        mlm=True,\n",
+        "        mlm_probability=0.15\n",
+        "    )\n",
+        "\n",
+        "    return dataset, data_collator\n",
+        "# ==========================================\n",
+        "# 3. PRISM ARCHITECTURE (Complex-Valued)\n",
+        "# ==========================================\n",
+        "\n",
+        "class ComplexDropout(nn.Module):\n",
+        "    def __init__(self, p=0.5):\n",
+        "        super().__init__()\n",
+        "        self.p = p\n",
+        "    def forward(self, z):\n",
+        "        if not self.training or self.p == 0.0: return z\n",
+        "        mask = torch.ones_like(z.real)\n",
+        "        mask = F.dropout(mask, self.p, self.training, inplace=False)\n",
+        "        return z * mask\n",
+        "\n",
+        "class RobustPhaseNorm(nn.Module):\n",
+        "    def __init__(self, d_model, eps=1e-5):\n",
+        "        super().__init__()\n",
+        "        self.scale = nn.Parameter(torch.ones(d_model))\n",
+        "        self.eps = eps\n",
+        "    def forward(self, x):\n",
+        "        mag = torch.abs(x)\n",
+        "        rms = torch.sqrt(torch.mean(mag**2, dim=-1, keepdim=True) + self.eps)\n",
+        "        return (x / rms) * self.scale\n",
+        "\n",
+        "class ModReLU(nn.Module):\n",
+        "    def __init__(self, features):\n",
+        "        super().__init__()\n",
+        "        self.b = nn.Parameter(torch.zeros(features))\n",
+        "    def forward(self, z):\n",
+        "        mag = torch.abs(z)\n",
+        "        new_mag = F.relu(mag + self.b)\n",
+        "        phase = z / (mag + 1e-6)\n",
+        "        return new_mag * phase\n",
+        "\n",
+        "class ComplexToRealBridge(nn.Module):\n",
+        "    def __init__(self, d_model):\n",
+        "        super().__init__()\n",
+        "        self.proj = nn.Linear(d_model * 2, d_model)\n",
+        "        self.norm = nn.LayerNorm(d_model)\n",
+        "    def forward(self, x_complex):\n",
+        "        cat = torch.cat([x_complex.real, x_complex.imag], dim=-1)\n",
+        "        return self.norm(self.proj(cat))\n",
+        "\n",
+        "# ==========================================\n",
+        "# 4. DYNAMIC RoSE (Mamba-3 Engine)\n",
+        "# ==========================================\n",
+        "class DynamicRoSE(nn.Module):\n",
+        "    def __init__(self, num_embeddings, embedding_dim, max_period=10000.0):\n",
+        "        super().__init__()\n",
+        "        self.embedding_dim = embedding_dim\n",
+        "\n",
+        "        # 1. Master Real Embedding (The \"Particle\")\n",
+        "        self.raw_embedding = nn.Embedding(num_embeddings, embedding_dim)\n",
+        "\n",
+        "        # 2. Complex Adapter (The \"Wave\" Magnitude/Initial Phase)\n",
+        "        self.adapter = nn.Linear(embedding_dim, embedding_dim * 2)\n",
+        "\n",
+        "        # 3. Static Frequencies (Positional)\n",
+        "        freqs = torch.exp(torch.arange(0, embedding_dim, dtype=torch.float32) * -(math.log(max_period) / embedding_dim))\n",
+        "        self.register_buffer('freqs', freqs)\n",
+        "\n",
+        "        self.rotation_predictor = nn.Linear(embedding_dim, embedding_dim * 2)\n",
+        "\n",
+        "    def forward(self, input_ids):\n",
+        "        # A. Raw Particle\n",
+        "        real_base = self.raw_embedding(input_ids)\n",
+        "        B, L, D = real_base.shape\n",
+        "\n",
+        "        # B. Complex Wave Content\n",
+        "        complex_params = self.adapter(real_base)\n",
+        "        z_t = torch.complex(complex_params[..., :D], complex_params[..., D:])\n",
+        "\n",
+        "        rot_raw = self.rotation_predictor(real_base)\n",
+        "        rot_x, rot_y = rot_raw.chunk(2, dim=-1)\n",
+        "\n",
+        "        rot_mag = torch.sqrt(rot_x**2 + rot_y**2 + 1e-6)\n",
+        "        dynamic_rot = torch.complex(rot_x / rot_mag, rot_y / rot_mag)\n",
+        "\n",
+        "        # D. Static Positional Rotation\n",
+        "        pos = torch.arange(L, device=input_ids.device).float()\n",
+        "        static_angles = torch.outer(pos, self.freqs) # [L, D]\n",
+        "        static_rot = torch.polar(torch.ones_like(static_angles), static_angles) # [L, D]\n",
+        "\n",
+        "        z_final = z_t * static_rot.unsqueeze(0) * dynamic_rot\n",
+        "\n",
+        "        return z_final, real_base\n",
+        "\n",
+        "# ==========================================\n",
+        "# 5. HYENA FILTER\n",
+        "# ==========================================\n",
+        "class HyenaNeuralFilter(nn.Module):\n",
+        "    def __init__(self, d_model, max_len=1024, hidden_dim=64):\n",
+        "        super().__init__()\n",
+        "        self.d_model = d_model\n",
+        "        freqs = torch.exp(torch.arange(0, hidden_dim, 2, dtype=torch.float32) * -(math.log(10000.0) / hidden_dim))\n",
+        "        self.register_buffer(\"freqs\", freqs)\n",
+        "        self.mlp = nn.Sequential(\n",
+        "            nn.Linear(hidden_dim, hidden_dim), nn.SiLU(),\n",
+        "            nn.Linear(hidden_dim, hidden_dim), nn.SiLU(),\n",
+        "            nn.Linear(hidden_dim, d_model * 2)\n",
+        "        )\n",
+        "    def forward(self, L, device):\n",
+        "        t = torch.linspace(0, 1, steps=L, device=device).unsqueeze(-1)\n",
+        "        emb = torch.cat([torch.sin(t * self.freqs), torch.cos(t * self.freqs)], dim=-1)\n",
+        "        out = self.mlp(emb).view(L, self.d_model, 2)\n",
+        "        return torch.complex(out[..., 0], out[..., 1])\n",
+        "\n",
+        "# ==========================================\n",
+        "# 6. GATED HARMONIC CONVOLUTION (Lean)\n",
+        "# ==========================================\n",
+        "class GatedHarmonicConvolution(nn.Module):\n",
+        "    def __init__(self, d_model, max_len=1024, dropout=0.1):\n",
+        "        super().__init__()\n",
+        "        self.d_model = d_model\n",
+        "        self.filter_len = max_len\n",
+        "        self.neural_filter = HyenaNeuralFilter(d_model, max_len=max_len)\n",
+        "        self.gate_proj = nn.Linear(d_model * 2, d_model * 2)\n",
+        "        self.mix_real = nn.Linear(d_model, d_model)\n",
+        "        self.mix_imag = nn.Linear(d_model, d_model)\n",
+        "        self.out_real = nn.Linear(d_model, d_model)\n",
+        "        self.out_imag = nn.Linear(d_model, d_model)\n",
+        "        self.activation = ModReLU(d_model)\n",
+        "        self.norm = RobustPhaseNorm(d_model)\n",
+        "        self.dropout = ComplexDropout(dropout)\n",
+        "\n",
+        "    def forward(self, x, src_mask=None):\n",
+        "        residual = x\n",
+        "        x_norm = self.norm(x)\n",
+        "        if src_mask is not None:\n",
+        "             x_norm = x_norm.masked_fill(src_mask.unsqueeze(-1), 0.0)\n",
+        "\n",
+        "        # 1. Global Beam (FFT + Hyena)\n",
+        "        B, L, D = x_norm.shape\n",
+        "        eff_L = min(L, self.filter_len)\n",
+        "        x_freq = torch.fft.fft(x_norm, n=eff_L, dim=1, norm='ortho')\n",
+        "        h = self.neural_filter(eff_L, x.device).unsqueeze(0)\n",
+        "        x_filtered = x_freq * h\n",
+        "        x_time = torch.fft.ifft(x_filtered, n=eff_L, dim=1, norm='ortho')\n",
+        "        if L > eff_L: x_time = F.pad(x_time, (0,0,0,L-eff_L))\n",
+        "        else: x_time = x_time[:, :L, :]\n",
+        "\n",
+        "        # 2. Gating\n",
+        "        gates = torch.sigmoid(self.gate_proj(torch.cat([x_norm.real, x_norm.imag], dim=-1)))\n",
+        "        g_r, g_i = gates.chunk(2, dim=-1)\n",
+        "        x_gated = torch.complex(x_time.real * g_r, x_time.imag * g_i)\n",
+        "\n",
+        "        # 3. Mixing & Out\n",
+        "        mr, mi = self.mix_real, self.mix_imag\n",
+        "        x_mixed = torch.complex(mr(x_gated.real) - mi(x_gated.imag), mr(x_gated.imag) + mi(x_gated.real))\n",
+        "        x_act = self.activation(x_mixed)\n",
+        "        or_, oi = self.out_real, self.out_imag\n",
+        "        out = torch.complex(or_(x_act.real) - oi(x_act.imag), or_(x_act.imag) + oi(x_act.real))\n",
+        "        return self.dropout(out) + residual\n",
+        "\n",
+        "# ==========================================\n",
+        "# 7. MODEL WRAPPERS\n",
+        "# ==========================================\n",
+        "class PRISMEncoder(nn.Module):\n",
+        "    def __init__(self, num_layers, d_model, max_len, dropout=0.1):\n",
+        "        super().__init__()\n",
+        "        self.layers = nn.ModuleList([\n",
+        "            GatedHarmonicConvolution(d_model, max_len, dropout)\n",
+        "            for _ in range(num_layers)\n",
+        "        ])\n",
+        "        self.final_norm = RobustPhaseNorm(d_model)\n",
+        "    def forward(self, x, src_mask=None):\n",
+        "        for layer in self.layers:\n",
+        "            if self.training: x = torch.utils.checkpoint.checkpoint(layer, x, src_mask, use_reentrant=False)\n",
+        "            else: x = layer(x, src_mask)\n",
+        "        return self.final_norm(x)\n",
+        "\n",
+        "class PRISM_WikiText_Model(nn.Module):\n",
+        "    def __init__(self, vocab_size, d_model, max_len, prism_depth=5, trans_depth=1, dropout=0.1):\n",
+        "        super().__init__()\n",
+        "        self.d_model = d_model\n",
+        "\n",
+        "        # 1. PRISM Core (The Optical/Passive Part)\n",
+        "        self.rose = DynamicRoSE(vocab_size, d_model)\n",
+        "        self.prism_encoder = PRISMEncoder(prism_depth, d_model, max_len=max_len, dropout=dropout)\n",
+        "        self.bridge = ComplexToRealBridge(d_model)\n",
+        "        self.periscope_proj = nn.Sequential(nn.Linear(d_model * 2, d_model), nn.LayerNorm(d_model), nn.GELU())\n",
+        "\n",
+        "        # 2. Refiner (The Digital/Active Part)\n",
+        "        # 🔄 SWAPPED: Replaced Standard Transformer with RoPE-Enabled Encoder\n",
+        "        if trans_depth > 0:\n",
+        "            self.refiner = Encoder(\n",
+        "                dim=d_model,\n",
+        "                depth=trans_depth,\n",
+        "                heads=8,\n",
+        "                rotary_pos_emb=True,\n",
+        "                attn_flash=True,\n",
+        "                attn_dropout=dropout,\n",
+        "                ff_dropout=dropout,\n",
+        "\n",
+        "            )\n",
+        "        else:\n",
+        "            self.refiner = None\n",
+        "\n",
+        "        # 3. Output\n",
+        "        self.lm_head = nn.Linear(d_model, vocab_size)\n",
+        "        self.lm_head.weight = self.rose.raw_embedding.weight\n",
+        "\n",
+        "    def forward(self, input_ids):\n",
+        "        # A. Wave Physics\n",
+        "        wave_src, particle_src = self.rose(input_ids)\n",
+        "        wave_out = self.prism_encoder(wave_src)\n",
+        "        wave_real = self.bridge(wave_out)\n",
+        "\n",
+        "        # B. Interface\n",
+        "        mixed_memory = self.periscope_proj(torch.cat([wave_real, particle_src], dim=-1))\n",
+        "\n",
+        "        # C. Digital Refinement (Now with RoPE)\n",
+        "        if self.refiner:\n",
+        "            out = self.refiner(mixed_memory)\n",
+        "        else:\n",
+        "            out = mixed_memory\n",
+        "\n",
+        "        return self.lm_head(out)\n",
+        "\n",
+        "class FNetBlock(nn.Module):\n",
+        "    def __init__(self, d_model, d_ff, dropout):\n",
+        "        super().__init__()\n",
+        "        self.norm_mix = nn.LayerNorm(d_model) # LayerNorm is safer for FNet than RMSNorm\n",
+        "        self.norm_ff = nn.LayerNorm(d_model)\n",
+        "\n",
+        "        self.mix_dropout = nn.Dropout(dropout)\n",
+        "\n",
+        "        self.ff = nn.Sequential(\n",
+        "            nn.Linear(d_model, d_ff),\n",
+        "            nn.GELU(),\n",
+        "            nn.Dropout(dropout),\n",
+        "            nn.Linear(d_ff, d_model),\n",
+        "            nn.Dropout(dropout)\n",
+        "        )\n",
+        "\n",
+        "    def forward(self, x):\n",
+        "        # 1. Fourier Mixing Branch\n",
+        "        residual = x\n",
+        "        x = self.norm_mix(x)\n",
+        "\n",
+        "        # --- THE FIX ---\n",
+        "        with torch.cuda.amp.autocast(enabled=False):\n",
+        "            x = x.float()\n",
+        "            # norm='ortho' makes the FFT energy-preserving.\n",
+        "            # Output magnitude will match input magnitude (~1).\n",
+        "            x = torch.fft.fftn(x, dim=(-2, -1), norm='ortho').real\n",
+        "            x = x.to(dtype=residual.dtype)\n",
+        "        # ---------------\n",
+        "\n",
+        "        # Now 'x' and 'residual' have roughly same magnitude.\n",
+        "        # The skip connection works again.\n",
+        "        x = self.mix_dropout(x)\n",
+        "        x = x + residual\n",
+        "\n",
+        "        # 2. Feed Forward Branch\n",
+        "        residual = x\n",
+        "        x = self.norm_ff(x)\n",
+        "        x = self.ff(x)\n",
+        "        return x + residual\n",
+        "\n",
+        "\n",
+        "class FNetEncoder(nn.Module):\n",
+        "    def __init__(self, depth, d_model, d_ff, dropout):\n",
+        "        super().__init__()\n",
+        "        self.layers = nn.ModuleList([\n",
+        "            FNetBlock(d_model, d_ff, dropout) for _ in range(depth)\n",
+        "        ])\n",
+        "        # [FIX] Use LayerNorm here to match the blocks\n",
+        "        self.norm_out = nn.LayerNorm(d_model)\n",
+        "\n",
+        "    def forward(self, x):\n",
+        "        for layer in self.layers:\n",
+        "            x = layer(x)\n",
+        "        return self.norm_out(x)\n",
+        "\n",
+        "class Pillars_DualStream(nn.Module):\n",
+        "    def __init__(self, vocab_size, d_model=512, d_branch=384, seq_len=4096, depth=4):\n",
+        "        super().__init__()\n",
+        "        self.d_branch = d_branch\n",
+        "        self.d_refiner = d_model\n",
+        "\n",
+        "        # --- A. Rate Stream (FNet) ---\n",
+        "        self.fnet_emb = nn.Embedding(vocab_size, d_branch)\n",
+        "        self.fnet_pos = nn.Embedding(seq_len, d_branch)\n",
+        "        self.stream_rate = FNetEncoder(depth=depth, d_model=d_branch, d_ff=d_branch*4, dropout=0.1)\n",
+        "\n",
+        "        # --- B. Phase Stream (PRISM) ---\n",
+        "        self.stream_phase_emb = DynamicRoSE(vocab_size, d_branch)\n",
+        "        self.stream_phase = PRISMEncoder(num_layers=depth, d_model=d_branch, max_len=seq_len, dropout=0.1)\n",
+        "        self.phase_bridge = ComplexToRealBridge(d_branch)\n",
+        "\n",
+        "        # --- C. Fusion (The Funnel) ---\n",
+        "        self.fusion_proj = nn.Linear(d_branch * 2, d_model)\n",
+        "        self.fusion_norm = nn.LayerNorm(d_model)\n",
+        "\n",
+        "        # --- D. Refiner ---\n",
+        "        self.refiner = Encoder(\n",
+        "            dim=d_model, depth=1, heads=8, attn_flash=True,\n",
+        "            rotary_pos_emb=True, attn_dropout=0.1, ff_dropout=0.1\n",
+        "        )\n",
+        "        self.lm_head = nn.Linear(d_model, vocab_size)\n",
+        "\n",
+        "    def forward(self, x):\n",
+        "        # 1. Rate Path\n",
+        "        f_emb = self.fnet_emb(x) + self.fnet_pos(torch.arange(x.shape[1], device=x.device))\n",
+        "        rate_out = self.stream_rate(f_emb)\n",
+        "\n",
+        "        # 2. Phase Path\n",
+        "        p_src, _ = self.stream_phase_emb(x)\n",
+        "        phase_out = self.phase_bridge(self.stream_phase(p_src))\n",
+        "\n",
+        "        # 3. Fusion\n",
+        "        fused = self.fusion_norm(self.fusion_proj(torch.cat([rate_out, phase_out], dim=-1)))\n",
+        "\n",
+        "        # 4. Refine & Output\n",
+        "        return self.lm_head(self.refiner(fused))\n",
+        "\n",
+        "\n",
+        "class Pillars_Compact(nn.Module):\n",
+        "    def __init__(self, vocab_size, d_model=512, d_branch=384, seq_len=4096, depth=4):\n",
+        "        super().__init__()\n",
+        "        self.d_model = d_model\n",
+        "        self.d_branch = d_branch\n",
+        "\n",
+        "        # 1. SHARED ROOT\n",
+        "        self.rose = DynamicRoSE(vocab_size, d_model)\n",
+        "\n",
+        "        # 2. DOWNSAMPLE (512 -> 384)\n",
+        "        self.particle_down = nn.Linear(d_model, d_branch)\n",
+        "        self.wave_down = nn.Linear(d_model * 2, d_branch * 2)\n",
+        "\n",
+        "        # 3. RATE STREAM (FNet, Depth 4)\n",
+        "        self.fnet_pos = nn.Embedding(seq_len, d_branch)\n",
+        "        self.stream_rate = FNetEncoder(depth=depth, d_model=d_branch, d_ff=d_branch*4, dropout=0.1)\n",
+        "\n",
+        "        # 4. PHASE STREAM (PRISM, Depth 4)\n",
+        "        self.stream_phase = PRISMEncoder(num_layers=depth, d_model=d_branch, max_len=seq_len, dropout=0.1)\n",
+        "        self.phase_bridge = ComplexToRealBridge(d_branch)\n",
+        "\n",
+        "        # 5. FUSION (Clean Projection)\n",
+        "        # Input: 384 (Rate) + 384 (Phase) = 768\n",
+        "        # Output: 512 (Refiner Dim)\n",
+        "        self.fusion_proj = nn.Linear(d_branch * 2, d_model)\n",
+        "        self.fusion_norm = nn.LayerNorm(d_model)\n",
+        "\n",
+        "        # 6. REFINER (The Brain)\n",
+        "        self.refiner = Encoder(\n",
+        "            dim=d_model, depth=1, heads=8, attn_flash=True,\n",
+        "            rotary_pos_emb=True, attn_dropout=0.1, ff_dropout=0.1\n",
+        "        )\n",
+        "\n",
+        "        # 7. TIED HEAD\n",
+        "        self.head_bias = nn.Parameter(torch.zeros(vocab_size))\n",
+        "\n",
+        "    def forward(self, input_ids):\n",
+        "        # A. Shared Root\n",
+        "        wave_src, particle_src = self.rose(input_ids)\n",
+        "\n",
+        "        # B. Downsample\n",
+        "        p_small = self.particle_down(particle_src)\n",
+        "        w_flat = torch.cat([wave_src.real, wave_src.imag], dim=-1)\n",
+        "        w_small_flat = self.wave_down(w_flat)\n",
+        "        w_small = torch.complex(w_small_flat[..., :self.d_branch], w_small_flat[..., self.d_branch:])\n",
+        "\n",
+        "        # C. Branches\n",
+        "        pos_emb = self.fnet_pos(torch.arange(input_ids.shape[1], device=input_ids.device))\n",
+        "        rate_out = self.stream_rate(p_small + pos_emb)\n",
+        "        phase_out = self.phase_bridge(self.stream_phase(w_small))\n",
+        "\n",
+        "        # D. Fusion (Concat -> Project)\n",
+        "        # We rely on the Transformer Refiner to attend to the right parts.\n",
+        "        stacked = torch.cat([rate_out, phase_out], dim=-1)\n",
+        "        context = self.fusion_norm(self.fusion_proj(stacked))\n",
+        "\n",
+        "        # E. Refiner\n",
+        "        refined = self.refiner(context)\n",
+        "\n",
+        "        # F. Output\n",
+        "        logits = F.linear(refined, self.rose.raw_embedding.weight, self.head_bias)\n",
+        "\n",
+        "        return logits\n",
+        "\n",
+        "import torch\n",
+        "import torch.nn as nn\n",
+        "from prettytable import PrettyTable # Optional, but makes tables nice.\n",
+        "# If you don't have prettytable, the code below uses standard f-strings.\n",
+        "\n",
+        "import torch\n",
+        "import torch.nn as nn\n",
+        "\n",
+        "import torch\n",
+        "import torch.nn as nn\n",
+        "\n",
+        "def deep_analyze_pillars(model):\n",
+        "    def get_p(obj):\n",
+        "        \"\"\"Safely returns parameter count for Modules OR raw Parameters.\"\"\"\n",
+        "        if isinstance(obj, nn.Parameter):\n",
+        "            return obj.numel()\n",
+        "        return sum(p.numel() for p in obj.parameters() if p.requires_grad)\n",
+        "\n",
+        "    def format_num(n):\n",
+        "        if n > 1e6: return f\"{n/1e6:.2f}M\"\n",
+        "        if n > 1e3: return f\"{n/1e3:.2f}K\"\n",
+        "        return str(n)\n",
+        "\n",
+        "    print(\"\\n\" + \"=\"*80)\n",
+        "    print(f\"🏗️  PILLARS (COMPACT) - DEEP LAYER ANALYSIS\")\n",
+        "    print(\"=\"*80)\n",
+        "    print(f\"{'MODULE / LAYER':<40} | {'PARAMS':<15} | {'TYPE'}\")\n",
+        "    print(\"-\" * 80)\n",
+        "\n",
+        "    total_params = get_p(model)\n",
+        "\n",
+        "    # -----------------------------------------------\n",
+        "    # 1. STATIC MEMORY (Embeddings)\n",
+        "    # -----------------------------------------------\n",
+        "    vocab_emb = get_p(model.rose.raw_embedding)\n",
+        "    fnet_pos  = get_p(model.fnet_pos)\n",
+        "\n",
+        "    print(f\"{'Shared Vocab Embedding':<40} | {format_num(vocab_emb):<15} | 💾 STORAGE\")\n",
+        "    print(f\"{'FNet Positional Embedding':<40} | {format_num(fnet_pos):<15} | 💾 STORAGE\")\n",
+        "\n",
+        "    # -----------------------------------------------\n",
+        "    # 2. INPUT LOGIC (RoSE & Downsampling)\n",
+        "    # -----------------------------------------------\n",
+        "    rose_total = get_p(model.rose)\n",
+        "    rose_logic = rose_total - vocab_emb # Subtract the embedding matrix we already counted\n",
+        "\n",
+        "    print(\"-\" * 80)\n",
+        "    print(f\"{'Dynamic RoSE (Adapters)':<40} | {format_num(rose_logic):<15} | 🌊 PHASE INIT\")\n",
+        "    print(f\"{'Particle Downsample (512->384)':<40} | {format_num(get_p(model.particle_down)):<15} | 📉 PROJ\")\n",
+        "    print(f\"{'Wave Downsample (1024->768)':<40} | {format_num(get_p(model.wave_down)):<15} | 📉 PROJ\")\n",
+        "\n",
+        "    # -----------------------------------------------\n",
+        "    # 3. STREAM A: RATE (FNet)\n",
+        "    # -----------------------------------------------\n",
+        "    print(\"-\" * 80)\n",
+        "    print(f\"TRACK A: RATE STREAM (FNet) - Depth {len(model.stream_rate.layers)}\")\n",
+        "\n",
+        "    fnet_encoder_total = 0\n",
+        "    for i, layer in enumerate(model.stream_rate.layers):\n",
+        "        p = get_p(layer)\n",
+        "        fnet_encoder_total += p\n",
+        "        print(f\"  ├─ FNet Block {i:<24} | {format_num(p):<15} | ⚡ RATE\")\n",
+        "\n",
+        "    fnet_norm = get_p(model.stream_rate.norm_out)\n",
+        "    fnet_encoder_total += fnet_norm\n",
+        "    print(f\"  └─ Final Norm {i:<24} | {format_num(fnet_norm):<15} | ⚡ RATE\")\n",
+        "\n",
+        "    # -----------------------------------------------\n",
+        "    # 4. STREAM B: PHASE (PRISM)\n",
+        "    # -----------------------------------------------\n",
+        "    print(\"-\" * 80)\n",
+        "    print(f\"TRACK B: PHASE STREAM (PRISM) - Depth {len(model.stream_phase.layers)}\")\n",
+        "\n",
+        "    prism_encoder_total = 0\n",
+        "    for i, layer in enumerate(model.stream_phase.layers):\n",
+        "        p = get_p(layer)\n",
+        "        prism_encoder_total += p\n",
+        "        print(f\"  ├─ PRISM Block {i:<23} | {format_num(p):<15} | 🌊 PHASE\")\n",
+        "\n",
+        "    prism_norm = get_p(model.stream_phase.final_norm)\n",
+        "    prism_encoder_total += prism_norm\n",
+        "    print(f\"  └─ Final Norm {i:<24} | {format_num(prism_norm):<15} | 🌊 PHASE\")\n",
+        "\n",
+        "    bridge_p = get_p(model.phase_bridge)\n",
+        "    print(f\"{'Phase Bridge (Complex->Real)':<40} | {format_num(bridge_p):<15} | 🌉 BRIDGE\")\n",
+        "\n",
+        "    # -----------------------------------------------\n",
+        "    # 5. THE BRAIN (Fusion & Refiner)\n",
+        "    # -----------------------------------------------\n",
+        "    print(\"-\" * 80)\n",
+        "    fusion_p = get_p(model.fusion_proj) + get_p(model.fusion_norm)\n",
+        "    print(f\"{'Fusion (Concat -> Proj -> Norm)':<40} | {format_num(fusion_p):<15} | 🧠 FUSION\")\n",
+        "\n",
+        "    refiner_p = get_p(model.refiner)\n",
+        "    print(f\"{'Transformer Refiner (1 Layer)':<40} | {format_num(refiner_p):<15} | 🧠 ATTENTION\")\n",
+        "\n",
+        "    # [FIX] Handle nn.Parameter directly\n",
+        "    head_bias_p = get_p(model.head_bias)\n",
+        "    print(f\"{'Output Head Bias':<40} | {format_num(head_bias_p):<15} | 🎯 OUTPUT\")\n",
+        "\n",
+        "    # -----------------------------------------------\n",
+        "    # 6. SUMMARY\n",
+        "    # -----------------------------------------------\n",
+        "    print(\"=\"*80)\n",
+        "\n",
+        "    storage = vocab_emb + fnet_pos + head_bias_p\n",
+        "    active = total_params - storage\n",
+        "\n",
+        "    print(f\"TOTAL PARAMETERS:      {total_params/1e6:.2f} M\")\n",
+        "    print(f\"   ├─ 💾 Storage:      {storage/1e6:.2f} M  (Embeddings)\")\n",
+        "    print(f\"   └─ 🧠 Compute:      {active/1e6:.2f} M  (Logic/Weights)\")\n",
+        "    print(\"-\" * 80)\n",
+        "    print(f\"STREAM BREAKDOWN:\")\n",
+        "    print(f\"   ├─ ⚡ Rate Stream:   {fnet_encoder_total/1e6:.2f} M\")\n",
+        "    print(f\"   └─ 🌊 Phase Stream:  {prism_encoder_total/1e6:.2f} M\")\n",
+        "    print(\"=\"*80 + \"\\n\")\n",
+        "\n",
+        "    return total_params\n",
+        "\n",
+        "model = Pillars_Compact(\n",
+        "    vocab_size=VOCAB_SIZE,\n",
+        "    d_model=D_MODEL,\n",
+        "    d_branch=D_BRANCH,\n",
+        "    seq_len=SEQ_LEN,\n",
+        "    depth=DEPTH\n",
+        ").to(DEVICE)\n",
+        "deep_analyze_pillars(model)"
+      ],
+      "metadata": {
+        "id": "V7DOwmmUjyin"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "\n",
+        "# Run the parameter analysis to confirm strict adherence to budget\n",
+        "def analyze_pillars_compact(model):\n",
+        "    print(\"\\n\" + \"=\"*70)\n",
+        "    print(\"🏛️  PILLARS COMPACT: ARCHITECTURAL COST ANALYSIS\")\n",
+        "    print(\"=\"*70)\n",
+        "\n",
+        "    stats = {\n",
+        "        \"Shared Memory (Storage)\": 0,\n",
+        "        \"Rate Stream (FNet)\": 0,\n",
+        "        \"Phase Stream (PRISM)\": 0,\n",
+        "        \"Fusion & Refiner\": 0,\n",
+        "        \"Tied Head Bias\": 0\n",
+        "    }\n",
+        "\n",
+        "    total_params = 0\n",
+        "\n",
+        "    for name, param in model.named_parameters():\n",
+        "        if not param.requires_grad: continue\n",
+        "        n = param.numel()\n",
+        "        total_params += n\n",
+        "\n",
+        "        if \"rose.raw_embedding\" in name:\n",
+        "            stats[\"Shared Memory (Storage)\"] += n\n",
+        "        elif \"rose.adapter\" in name or \"rose.rotation\" in name or \"stream_phase\" in name or \"phase_bridge\" in name:\n",
+        "            stats[\"Phase Stream (PRISM)\"] += n\n",
+        "        elif \"fnet_pos\" in name or \"stream_rate\" in name:\n",
+        "            stats[\"Rate Stream (FNet)\"] += n\n",
+        "        elif \"gate\" in name or \"mix\" in name or \"refiner\" in name or \"down\" in name or \"proj\" in name or \"norm\" in name:\n",
+        "            stats[\"Fusion & Refiner\"] += n\n",
+        "        elif \"head_bias\" in name:\n",
+        "            stats[\"Tied Head Bias\"] += n\n",
+        "        else:\n",
+        "            print(f\"⚠️ Uncategorized: {name} ({n})\")\n",
+        "\n",
+        "    print(f\"{'COMPONENT':<30} | {'PARAMS':<12} | {'% TOTAL':<8}\")\n",
+        "    print(\"-\" * 60)\n",
+        "\n",
+        "    for category, count in stats.items():\n",
+        "        if count > 0:\n",
+        "            pct = (count / total_params) * 100\n",
+        "            print(f\"{category:<30} | {count:12,} | {pct:6.1f}%\")\n",
+        "\n",
+        "    print(\"-\" * 60)\n",
+        "    print(f\"{'TOTAL PARAMETERS':<30} | {total_params:12,} | 100.0%\")\n",
+        "    print(\"=\" * 70)\n",
+        "\n",
+        "\n",
+        "    active_params = total_params - stats[\"Shared Memory (Storage)\"] - stats[\"Tied Head Bias\"]\n",
+        "    print(f\"   1. Total Model Size:      {total_params/1e6:.1f}M\")\n",
+        "    print(f\"   2. Baseline Target:       ~32.5M\")\n",
+        "    print(f\"   3. Active Reasoning Params: {active_params/1e6:.1f}M (The actual brain)\")\n",
+        "    print(\"=\"*70 + \"\\n\")\n",
+        "\n",
+        "\n"
+      ],
+      "metadata": {
+        "id": "ke4fYT8UX5zH"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# ==========================================\n",
+        "# 4. LOGGING UTILITIES\n",
+        "# ==========================================\n",
+        "def generate_run_id():\n",
+        "    raw = datetime.now().strftime(\"%Y%m%d%H%M%S%f\")\n",
+        "    return hashlib.md5(raw.encode()).hexdigest()[:8]\n",
+        "\n",
+        "def log_environment(save_dir, run_id, config):\n",
+        "    log_path = os.path.join(save_dir, f\"env_metadata_{run_id}.txt\")\n",
+        "    with open(log_path, \"w\") as f:\n",
+        "        f.write(f\"PRISM EXPERIMENT METADATA | Run ID: {run_id}\\n{'='*50}\\n\")\n",
+        "        for k, v in config.items(): f.write(f\"{k}: {v}\\n\")\n",
+        "    print(f\"📝 Environment Snapshot saved to: {log_path}\")\n",
+        "\n",
+        "def log_metrics(save_dir, run_id, epoch, train_loss, val_loss, ppl):\n",
+        "    log_path = os.path.join(save_dir, f\"metrics_log_{run_id}.csv\")\n",
+        "    if not os.path.exists(log_path):\n",
+        "        with open(log_path, \"w\") as f: f.write(\"Timestamp,Epoch,Train_Loss,Val_Loss,Perplexity\\n\")\n",
+        "    with open(log_path, \"a\") as f:\n",
+        "        ts = datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n",
+        "        f.write(f\"{ts},{epoch},{train_loss:.6f},{val_loss:.6f},{ppl:.6f}\\n\")\n",
+        "\n",
+        "\n",
+        "def save_checkpoint(path, model, optimizer, scheduler, epoch, best_loss, config):\n",
+        "    torch.save({\n",
+        "        'epoch': epoch,\n",
+        "        'model_state_dict': model.state_dict(),\n",
+        "        'optimizer_state_dict': optimizer.state_dict(),\n",
+        "        'scheduler_state_dict': scheduler.state_dict(),\n",
+        "        'best_val_loss': best_loss,\n",
+        "        'config': config\n",
+        "    }, path)\n",
+        "\n",
+        "def init_pillars_weights(model):\n",
+        "    print(\"✨ APPLYING PILLARS INITIALIZATION PROTOCOL...\")\n",
+        "\n",
+        "    # 1. SHARED ROOT (RoSE) - MATCHING YOUR ORIGINAL LOGIC\n",
+        "    # Standard embedding init\n",
+        "    nn.init.normal_(model.rose.raw_embedding.weight, std=model.d_model ** -0.5)\n",
+        "\n",
+        "    # Adapter: Orthogonal ensures clean entry to complex plane\n",
+        "    nn.init.orthogonal_(model.rose.adapter.weight)\n",
+        "\n",
+        "    # --- THE ROSE IDENTITY TRICK (From your original code) ---\n",
+        "    # Start with almost zero rotation influence from content\n",
+        "    nn.init.normal_(model.rose.rotation_predictor.weight, std=0.01)\n",
+        "    with torch.no_grad():\n",
+        "        # Force initial vector to (1, 0) -> Angle 0, Mag 1\n",
+        "        # This allows the model to start with \"Safe\" static physics\n",
+        "        model.rose.rotation_predictor.bias[:model.d_model].fill_(1.0)\n",
+        "        model.rose.rotation_predictor.bias[model.d_model:].fill_(0.0)\n",
+        "    # -------------------------------------------------------\n",
+        "\n",
+        "    # 2. DOWNSAMPLERS (The Split)\n",
+        "    # Scale gain by 1.414 (sqrt 2) to preserve energy when halving dimensions\n",
+        "    nn.init.orthogonal_(model.particle_down.weight, gain=1.414)\n",
+        "    nn.init.orthogonal_(model.wave_down.weight, gain=1.414)\n",
+        "\n",
+        "    # 3. FNET BRANCH (Rate Stream)\n",
+        "    # Kaiming Normal (Good for GELU)\n",
+        "    for name, m in model.stream_rate.named_modules():\n",
+        "        if isinstance(m, nn.Linear):\n",
+        "            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n",
+        "            if m.bias is not None: nn.init.zeros_(m.bias)\n",
+        "\n",
+        "    # 4. PRISM BRANCH (Phase Stream)\n",
+        "    # Xavier Uniform (Good for Complex/Linear)\n",
+        "    for name, m in model.stream_phase.named_modules():\n",
+        "        if isinstance(m, nn.Linear):\n",
+        "            nn.init.xavier_uniform_(m.weight, gain=1.0)\n",
+        "            if m.bias is not None: nn.init.zeros_(m.bias)\n",
+        "        # Initialize ModReLU bias slightly positive to avoid dead neurons\n",
+        "        if isinstance(m, ModReLU):\n",
+        "            nn.init.constant_(m.b, 0.01)\n",
+        "\n",
+        "    # 5. FUSION & REFINER\n",
+        "    # Start neutral\n",
+        "    nn.init.xavier_uniform_(model.fusion_proj.weight, gain=1.0)\n",
+        "\n",
+        "    for p in model.refiner.parameters():\n",
+        "        if p.dim() > 1:\n",
+        "            nn.init.xavier_uniform_(p)\n",
+        "\n",
+        "    # 6. TIED HEAD BIAS\n",
+        "    nn.init.zeros_(model.head_bias)\n",
+        "\n",
+        "    print(\"✅ INITIALIZATION COMPLETE.\")\n",
+        "\n",
+        "def run_wikitext_training(experiment_name=\"PILLARS_SplitStream_9Layer\"):\n",
+        "    from google.colab import drive\n",
+        "    if not os.path.exists('/content/drive'): drive.mount('/content/drive')\n",
+        "\n",
+        "    # --- SETUP DIRS ---\n",
+        "    if RESUME_PATH and os.path.exists(RESUME_PATH):\n",
+        "        print(f\"🔄 RESUMING FROM: {RESUME_PATH}\")\n",
+        "        checkpoint = torch.load(RESUME_PATH, map_location=DEVICE)\n",
+        "        SAVE_DIR = os.path.dirname(RESUME_PATH)\n",
+        "        run_id = checkpoint.get('config', {}).get('run_id', 'resumed')\n",
+        "    else:\n",
+        "        run_id = hashlib.md5(datetime.now().strftime(\"%Y%m%d%H%M%S%f\").encode()).hexdigest()[:8]\n",
+        "        timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
+        "        folder_name = f\"{experiment_name}_{timestamp}_{run_id}\"\n",
+        "        SAVE_DIR = os.path.join(\"/content/drive/My Drive/PRISM_Experiments\", folder_name)\n",
+        "        os.makedirs(SAVE_DIR, exist_ok=True)\n",
+        "        print(f\"💾 Checkpoints: {SAVE_DIR}\")\n",
+        "\n",
+        "    writer = SummaryWriter(log_dir=SAVE_DIR)\n",
+        "    GRAD_ACCUM = 4\n",
+        "\n",
+        "    lm_datasets, data_collator = prepare_data_from_hub()\n",
+        "\n",
+        "    train_loader = DataLoader(\n",
+        "        lm_datasets[\"train\"], batch_size=BATCH_SIZE, shuffle=True,\n",
+        "        collate_fn=data_collator, num_workers=2, pin_memory=True,\n",
+        "        prefetch_factor=2, persistent_workers=True\n",
+        "    )\n",
+        "    valid_loader = DataLoader(\n",
+        "        lm_datasets[\"validation\"], batch_size=BATCH_SIZE,\n",
+        "        collate_fn=data_collator, num_workers=2, pin_memory=True\n",
+        "    )\n",
+        "    test_loader = DataLoader(\n",
+        "        lm_datasets[\"test\"], batch_size=BATCH_SIZE,\n",
+        "        collate_fn=data_collator, num_workers=2, pin_memory=True\n",
+        "    )\n",
+        "\n",
+        "    print(\"\\n⚡ INITIALIZING PILLARS MODEL...\")\n",
+        "\n",
+        "    # INSTANTIATE THE NEW MODEL\n",
+        "    model = Pillars_Compact(\n",
+        "        vocab_size=VOCAB_SIZE,\n",
+        "        d_model=D_MODEL,\n",
+        "        d_branch=D_BRANCH,\n",
+        "        seq_len=SEQ_LEN,\n",
+        "        depth=DEPTH\n",
+        "    ).to(DEVICE)\n",
+        "\n",
+        "    optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=0.01) # Added decay for stabilization\n",
+        "    total_steps = (len(train_loader) // GRAD_ACCUM) * EPOCHS\n",
+        "    scheduler = get_cosine_schedule_with_warmup(\n",
+        "        optimizer, num_warmup_steps=int(0.05 * total_steps), num_training_steps=total_steps\n",
+        "    )\n",
+        "    criterion = nn.CrossEntropyLoss()\n",
+        "\n",
+        "    start_epoch = 0\n",
+        "    best_val_loss = float('inf')\n",
+        "\n",
+        "    if RESUME_PATH and os.path.exists(RESUME_PATH):\n",
+        "        model.load_state_dict(checkpoint['model_state_dict'])\n",
+        "        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
+        "        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])\n",
+        "        start_epoch = checkpoint['epoch'] + 1\n",
+        "        best_val_loss = checkpoint['best_val_loss']\n",
+        "        del checkpoint\n",
+        "        torch.cuda.empty_cache()\n",
+        "    else:\n",
+        "        # APPLY THE NEW INIT LOGIC\n",
+        "        init_pillars_weights(model)\n",
+        "    print(model)\n",
+        "    analyze_pillars_compact(model)\n",
+        "    print(f\"\\n🚀 STARTING (Ep {start_epoch+1} to {EPOCHS})\")\n",
+        "    global_step = (len(train_loader) // GRAD_ACCUM) * start_epoch\n",
+        "\n",
+        "    for epoch in range(start_epoch, EPOCHS):\n",
+        "        model.train()\n",
+        "        pbar = tqdm(train_loader, desc=f\"Ep {epoch+1}/{EPOCHS}\")\n",
+        "\n",
+        "        for step, batch in enumerate(pbar):\n",
+        "            x, y = batch['input_ids'].to(DEVICE), batch['labels'].to(DEVICE)\n",
+        "\n",
+        "            loss = criterion(model(x).view(-1, VOCAB_SIZE), y.view(-1)) / GRAD_ACCUM\n",
+        "            loss.backward()\n",
+        "\n",
+        "            if (step + 1) % GRAD_ACCUM == 0:\n",
+        "                # 1. Calc Norm\n",
+        "                grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
+        "\n",
+        "                # 2. Step\n",
+        "                optimizer.step()\n",
+        "                scheduler.step()\n",
+        "                optimizer.zero_grad()\n",
+        "                global_step += 1\n",
+        "\n",
+        "                # 3. LOGGING\n",
+        "                actual_loss = loss.item() * GRAD_ACCUM\n",
+        "\n",
+        "                # [FIX] Log Grad Norm to TensorBoard now\n",
+        "                writer.add_scalar('Train/Loss', actual_loss, global_step)\n",
+        "                writer.add_scalar('Train/GradNorm', grad_norm.item(), global_step)\n",
+        "\n",
+        "                # 4. Progress Bar\n",
+        "                pbar.set_postfix({\n",
+        "                    'loss': f\"{actual_loss:.4f}\",\n",
+        "                    'gnorm': f\"{grad_norm.item():.2f}\"\n",
+        "                })\n",
+        "\n",
+        "        # VALIDATION\n",
+        "        model.eval()\n",
+        "        val_loss = 0\n",
+        "        with torch.no_grad():\n",
+        "            for batch in valid_loader:\n",
+        "                x, y = batch['input_ids'].to(DEVICE), batch['labels'].to(DEVICE)\n",
+        "                val_loss += criterion(model(x).view(-1, VOCAB_SIZE), y.view(-1)).item()\n",
+        "\n",
+        "        avg_val_loss = val_loss / len(valid_loader)\n",
+        "        ppl = math.exp(avg_val_loss) if avg_val_loss < 100 else float('inf')\n",
+        "\n",
+        "        print(f\"✨ Epoch {epoch+1} | Val Loss: {avg_val_loss:.4f} | PPL: {ppl:.2f}\")\n",
+        "        writer.add_scalar('Val/PPL', ppl, epoch+1)\n",
+        "\n",
+        "        config_dump = {\"epoch\": epoch, \"run_id\": run_id}\n",
+        "        save_checkpoint(os.path.join(SAVE_DIR, \"last.pt\"), model, optimizer, scheduler, epoch, best_val_loss, config_dump)\n",
+        "\n",
+        "        if avg_val_loss < best_val_loss:\n",
+        "            best_val_loss = avg_val_loss\n",
+        "            torch.save(model.state_dict(), os.path.join(SAVE_DIR, \"best.pt\"))\n",
+        "            print(\"   🏆 New Best Model Saved!\")\n",
+        "\n",
+        "    best_path = os.path.join(SAVE_DIR, \"best.pt\")\n",
+        "    if os.path.exists(best_path):\n",
+        "        model.load_state_dict(torch.load(best_path))\n",
+        "        model.eval()\n",
+        "        test_loss = 0\n",
+        "        with torch.no_grad():\n",
+        "            for batch in tqdm(test_loader, desc=\"Testing\"):\n",
+        "                x, y = batch['input_ids'].to(DEVICE), batch['labels'].to(DEVICE)\n",
+        "                test_loss += criterion(model(x).view(-1, VOCAB_SIZE), y.view(-1)).item()\n",
+        "        print(f\"🏆 FINAL PPL: {math.exp(test_loss/len(test_loader)):.2f}\")\n",
+        "\n",
+        "    writer.close()\n",
+        "    return model"
+      ],
+      "metadata": {
+        "id": "-TNEv89gkS1k"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "if __name__ == \"__main__\":\n",
+        "\n",
+        "    print(\"🔥 IGNITING PILLARS TRAINING PIPELINE...\")\n",
+        "\n",
+        "    # 1. Run the Training Routine\n",
+        "    # This handles Model Creation -> Analysis -> Training -> Saving\n",
+        "    trained_prism = run_wikitext_training()\n",
+        "\n",
+        "    # 2. Cleanup & Shutdown\n",
+        "    print(\"✅ Experiment Complete. Shutting down runtime...\")\n",
+        "    from google.colab import runtime\n",
+        "    runtime.unassign()"
+      ],
+      "metadata": {
+        "id": "KaiJU0tPkVp-"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "from google.colab import runtime\n",
+        "runtime.unassign()"
+      ],
+      "metadata": {
+        "id": "bxFTYWHVqcSI"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}

PRISM_wikitext_103_last.ipynb

@@ -0,0 +1,589 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "A100"
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "code",
+      "source": [
+        "!pip install -q x-transformers"
+      ],
+      "metadata": {
+        "id": "6q9RTvlf5IiS"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import torch\n",
+        "import torch.nn as nn\n",
+        "import torch.nn.functional as F\n",
+        "import torch.optim as optim\n",
+        "import math\n",
+        "import os\n",
+        "import sys\n",
+        "import subprocess\n",
+        "import hashlib\n",
+        "import gc\n",
+        "from datetime import datetime\n",
+        "from tqdm.auto import tqdm\n",
+        "from torch.utils.data import DataLoader\n",
+        "from torch.utils.tensorboard import SummaryWriter\n",
+        "from transformers import RobertaTokenizerFast, get_cosine_schedule_with_warmup, DataCollatorForLanguageModeling\n",
+        "from datasets import load_dataset\n",
+        "from x_transformers import Encoder\n",
+        "\n",
+        "# ==========================================\n",
+        "# 1. CONFIGURATION\n",
+        "# ==========================================\n",
+        "# YOUR REPO ID (Created in previous step)\n",
+        "HF_ID = \"prism-lab/wikitext-103-prism-32k-seq4k\"\n",
+        "\n",
+        "# Hyperparameters\n",
+        "VOCAB_SIZE = 32768\n",
+        "SEQ_LEN = 4096\n",
+        "BATCH_SIZE = 8\n",
+        "EPOCHS = 40\n",
+        "LR = 1e-3\n",
+        "D_MODEL = 512\n",
+        "DEPTH = 6\n",
+        "DROPOUT = 0.1\n",
+        "RESUME_PATH = None\n",
+        "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+        "torch.set_float32_matmul_precision(\"high\")\n",
+        "\n",
+        "# ==========================================\n",
+        "# 2. DATA PIPELINE (The \"Pro\" Way)\n",
+        "# ==========================================\n",
+        "def prepare_data_from_hub():\n",
+        "    print(f\"⬇️ Pulling Pre-Tokenized Data from {HF_ID}...\")\n",
+        "\n",
+        "    # 1. Load Tokenizer (Instant)\n",
+        "    # This pulls the exact tokenizer you uploaded\n",
+        "    tokenizer = RobertaTokenizerFast.from_pretrained(HF_ID)\n",
+        "\n",
+        "    # 2. Load Dataset (Instant)\n",
+        "    # This pulls the already chunked/tokenized data\n",
+        "    dataset = load_dataset(HF_ID)\n",
+        "\n",
+        "    print(f\"✅ Loaded {len(dataset['train'])} training chunks.\")\n",
+        "\n",
+        "    # 3. Collator\n",
+        "    data_collator = DataCollatorForLanguageModeling(\n",
+        "        tokenizer=tokenizer,\n",
+        "        mlm=True,\n",
+        "        mlm_probability=0.15\n",
+        "    )\n",
+        "\n",
+        "    return dataset, data_collator\n",
+        "# ==========================================\n",
+        "# 3. PRISM ARCHITECTURE (Complex-Valued)\n",
+        "# ==========================================\n",
+        "\n",
+        "class ComplexDropout(nn.Module):\n",
+        "    def __init__(self, p=0.5):\n",
+        "        super().__init__()\n",
+        "        self.p = p\n",
+        "    def forward(self, z):\n",
+        "        if not self.training or self.p == 0.0: return z\n",
+        "        mask = torch.ones_like(z.real)\n",
+        "        mask = F.dropout(mask, self.p, self.training, inplace=False)\n",
+        "        return z * mask\n",
+        "\n",
+        "class RobustPhaseNorm(nn.Module):\n",
+        "    def __init__(self, d_model, eps=1e-5):\n",
+        "        super().__init__()\n",
+        "        self.scale = nn.Parameter(torch.ones(d_model))\n",
+        "        self.eps = eps\n",
+        "    def forward(self, x):\n",
+        "        mag = torch.abs(x)\n",
+        "        rms = torch.sqrt(torch.mean(mag**2, dim=-1, keepdim=True) + self.eps)\n",
+        "        return (x / rms) * self.scale\n",
+        "\n",
+        "class ModReLU(nn.Module):\n",
+        "    def __init__(self, features):\n",
+        "        super().__init__()\n",
+        "        self.b = nn.Parameter(torch.zeros(features))\n",
+        "    def forward(self, z):\n",
+        "        mag = torch.abs(z)\n",
+        "        new_mag = F.relu(mag + self.b)\n",
+        "        phase = z / (mag + 1e-6)\n",
+        "        return new_mag * phase\n",
+        "\n",
+        "class ComplexToRealBridge(nn.Module):\n",
+        "    def __init__(self, d_model):\n",
+        "        super().__init__()\n",
+        "        self.proj = nn.Linear(d_model * 2, d_model)\n",
+        "        self.norm = nn.LayerNorm(d_model)\n",
+        "    def forward(self, x_complex):\n",
+        "        cat = torch.cat([x_complex.real, x_complex.imag], dim=-1)\n",
+        "        return self.norm(self.proj(cat))\n",
+        "\n",
+        "# ==========================================\n",
+        "# 4. DYNAMIC RoSE (Mamba-3 Engine)\n",
+        "# ==========================================\n",
+        "class DynamicRoSE(nn.Module):\n",
+        "    def __init__(self, num_embeddings, embedding_dim, max_period=10000.0):\n",
+        "        super().__init__()\n",
+        "        self.embedding_dim = embedding_dim\n",
+        "\n",
+        "        # 1. Master Real Embedding (The \"Particle\")\n",
+        "        self.raw_embedding = nn.Embedding(num_embeddings, embedding_dim)\n",
+        "\n",
+        "        # 2. Complex Adapter (The \"Wave\" Magnitude/Initial Phase)\n",
+        "        self.adapter = nn.Linear(embedding_dim, embedding_dim * 2)\n",
+        "\n",
+        "        # 3. Static Frequencies (Positional)\n",
+        "        freqs = torch.exp(torch.arange(0, embedding_dim, dtype=torch.float32) * -(math.log(max_period) / embedding_dim))\n",
+        "        self.register_buffer('freqs', freqs)\n",
+        "\n",
+        "        self.rotation_predictor = nn.Linear(embedding_dim, embedding_dim * 2)\n",
+        "\n",
+        "    def forward(self, input_ids):\n",
+        "        # A. Raw Particle\n",
+        "        real_base = self.raw_embedding(input_ids)\n",
+        "        B, L, D = real_base.shape\n",
+        "\n",
+        "        # B. Complex Wave Content\n",
+        "        complex_params = self.adapter(real_base)\n",
+        "        z_t = torch.complex(complex_params[..., :D], complex_params[..., D:])\n",
+        "\n",
+        "        rot_raw = self.rotation_predictor(real_base)\n",
+        "        rot_x, rot_y = rot_raw.chunk(2, dim=-1)\n",
+        "\n",
+        "        rot_mag = torch.sqrt(rot_x**2 + rot_y**2 + 1e-6)\n",
+        "        dynamic_rot = torch.complex(rot_x / rot_mag, rot_y / rot_mag)\n",
+        "\n",
+        "        # D. Static Positional Rotation\n",
+        "        pos = torch.arange(L, device=input_ids.device).float()\n",
+        "        static_angles = torch.outer(pos, self.freqs) # [L, D]\n",
+        "        static_rot = torch.polar(torch.ones_like(static_angles), static_angles) # [L, D]\n",
+        "\n",
+        "        z_final = z_t * static_rot.unsqueeze(0) * dynamic_rot\n",
+        "\n",
+        "        return z_final, real_base\n",
+        "\n",
+        "# ==========================================\n",
+        "# 5. HYENA FILTER\n",
+        "# ==========================================\n",
+        "class HyenaNeuralFilter(nn.Module):\n",
+        "    def __init__(self, d_model, max_len=1024, hidden_dim=64):\n",
+        "        super().__init__()\n",
+        "        self.d_model = d_model\n",
+        "        freqs = torch.exp(torch.arange(0, hidden_dim, 2, dtype=torch.float32) * -(math.log(10000.0) / hidden_dim))\n",
+        "        self.register_buffer(\"freqs\", freqs)\n",
+        "        self.mlp = nn.Sequential(\n",
+        "            nn.Linear(hidden_dim, hidden_dim), nn.SiLU(),\n",
+        "            nn.Linear(hidden_dim, hidden_dim), nn.SiLU(),\n",
+        "            nn.Linear(hidden_dim, d_model * 2)\n",
+        "        )\n",
+        "    def forward(self, L, device):\n",
+        "        t = torch.linspace(0, 1, steps=L, device=device).unsqueeze(-1)\n",
+        "        emb = torch.cat([torch.sin(t * self.freqs), torch.cos(t * self.freqs)], dim=-1)\n",
+        "        out = self.mlp(emb).view(L, self.d_model, 2)\n",
+        "        return torch.complex(out[..., 0], out[..., 1])\n",
+        "\n",
+        "# ==========================================\n",
+        "# 6. GATED HARMONIC CONVOLUTION (Lean)\n",
+        "# ==========================================\n",
+        "class GatedHarmonicConvolution(nn.Module):\n",
+        "    def __init__(self, d_model, max_len=1024, dropout=0.1):\n",
+        "        super().__init__()\n",
+        "        self.d_model = d_model\n",
+        "        self.filter_len = max_len\n",
+        "        self.neural_filter = HyenaNeuralFilter(d_model, max_len=max_len)\n",
+        "        self.gate_proj = nn.Linear(d_model * 2, d_model * 2)\n",
+        "        self.mix_real = nn.Linear(d_model, d_model)\n",
+        "        self.mix_imag = nn.Linear(d_model, d_model)\n",
+        "        self.out_real = nn.Linear(d_model, d_model)\n",
+        "        self.out_imag = nn.Linear(d_model, d_model)\n",
+        "        self.activation = ModReLU(d_model)\n",
+        "        self.norm = RobustPhaseNorm(d_model)\n",
+        "        self.dropout = ComplexDropout(dropout)\n",
+        "\n",
+        "    def forward(self, x, src_mask=None):\n",
+        "        residual = x\n",
+        "        x_norm = self.norm(x)\n",
+        "        if src_mask is not None:\n",
+        "             x_norm = x_norm.masked_fill(src_mask.unsqueeze(-1), 0.0)\n",
+        "\n",
+        "        # 1. Global Beam (FFT + Hyena)\n",
+        "        B, L, D = x_norm.shape\n",
+        "        eff_L = min(L, self.filter_len)\n",
+        "        x_freq = torch.fft.fft(x_norm, n=eff_L, dim=1, norm='ortho')\n",
+        "        h = self.neural_filter(eff_L, x.device).unsqueeze(0)\n",
+        "        x_filtered = x_freq * h\n",
+        "        x_time = torch.fft.ifft(x_filtered, n=eff_L, dim=1, norm='ortho')\n",
+        "        if L > eff_L: x_time = F.pad(x_time, (0,0,0,L-eff_L))\n",
+        "        else: x_time = x_time[:, :L, :]\n",
+        "\n",
+        "        # 2. Gating\n",
+        "        gates = torch.sigmoid(self.gate_proj(torch.cat([x_norm.real, x_norm.imag], dim=-1)))\n",
+        "        g_r, g_i = gates.chunk(2, dim=-1)\n",
+        "        x_gated = torch.complex(x_time.real * g_r, x_time.imag * g_i)\n",
+        "\n",
+        "        # 3. Mixing & Out\n",
+        "        mr, mi = self.mix_real, self.mix_imag\n",
+        "        x_mixed = torch.complex(mr(x_gated.real) - mi(x_gated.imag), mr(x_gated.imag) + mi(x_gated.real))\n",
+        "        x_act = self.activation(x_mixed)\n",
+        "        or_, oi = self.out_real, self.out_imag\n",
+        "        out = torch.complex(or_(x_act.real) - oi(x_act.imag), or_(x_act.imag) + oi(x_act.real))\n",
+        "        return self.dropout(out) + residual\n",
+        "\n",
+        "# ==========================================\n",
+        "# 7. MODEL WRAPPERS\n",
+        "# ==========================================\n",
+        "class PRISMEncoder(nn.Module):\n",
+        "    def __init__(self, num_layers, d_model, max_len, dropout=0.1):\n",
+        "        super().__init__()\n",
+        "        self.layers = nn.ModuleList([\n",
+        "            GatedHarmonicConvolution(d_model, max_len, dropout)\n",
+        "            for _ in range(num_layers)\n",
+        "        ])\n",
+        "        self.final_norm = RobustPhaseNorm(d_model)\n",
+        "    def forward(self, x, src_mask=None):\n",
+        "        for layer in self.layers:\n",
+        "            if self.training: x = torch.utils.checkpoint.checkpoint(layer, x, src_mask, use_reentrant=False)\n",
+        "            else: x = layer(x, src_mask)\n",
+        "        return self.final_norm(x)\n",
+        "\n",
+        "class PRISM_WikiText_Model(nn.Module):\n",
+        "    def __init__(self, vocab_size, d_model, max_len, prism_depth=5, trans_depth=1, dropout=0.1):\n",
+        "        super().__init__()\n",
+        "        self.d_model = d_model\n",
+        "\n",
+        "        # 1. PRISM Core (The Optical/Passive Part)\n",
+        "        self.rose = DynamicRoSE(vocab_size, d_model)\n",
+        "        self.prism_encoder = PRISMEncoder(prism_depth, d_model, max_len=max_len, dropout=dropout)\n",
+        "        self.bridge = ComplexToRealBridge(d_model)\n",
+        "        self.periscope_proj = nn.Sequential(nn.Linear(d_model * 2, d_model), nn.LayerNorm(d_model), nn.GELU())\n",
+        "\n",
+        "        # 2. Refiner (The Digital/Active Part)\n",
+        "        # 🔄 SWAPPED: Replaced Standard Transformer with RoPE-Enabled Encoder\n",
+        "        if trans_depth > 0:\n",
+        "            self.refiner = Encoder(\n",
+        "                dim=d_model,\n",
+        "                depth=trans_depth,\n",
+        "                heads=8,\n",
+        "                rotary_pos_emb=True,\n",
+        "                attn_flash=True,\n",
+        "                attn_dropout=dropout,\n",
+        "                ff_dropout=dropout,\n",
+        "\n",
+        "            )\n",
+        "        else:\n",
+        "            self.refiner = None\n",
+        "\n",
+        "        # 3. Output\n",
+        "        self.lm_head = nn.Linear(d_model, vocab_size)\n",
+        "        self.lm_head.weight = self.rose.raw_embedding.weight\n",
+        "\n",
+        "    def forward(self, input_ids):\n",
+        "        # A. Wave Physics\n",
+        "        wave_src, particle_src = self.rose(input_ids)\n",
+        "        wave_out = self.prism_encoder(wave_src)\n",
+        "        wave_real = self.bridge(wave_out)\n",
+        "\n",
+        "        # B. Interface\n",
+        "        mixed_memory = self.periscope_proj(torch.cat([wave_real, particle_src], dim=-1))\n",
+        "\n",
+        "        # C. Digital Refinement (Now with RoPE)\n",
+        "        if self.refiner:\n",
+        "            out = self.refiner(mixed_memory)\n",
+        "        else:\n",
+        "            out = mixed_memory\n",
+        "\n",
+        "        return self.lm_head(out)"
+      ],
+      "metadata": {
+        "id": "V7DOwmmUjyin"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# ==========================================\n",
+        "# 4. LOGGING UTILITIES\n",
+        "# ==========================================\n",
+        "def generate_run_id():\n",
+        "    raw = datetime.now().strftime(\"%Y%m%d%H%M%S%f\")\n",
+        "    return hashlib.md5(raw.encode()).hexdigest()[:8]\n",
+        "\n",
+        "def log_environment(save_dir, run_id, config):\n",
+        "    log_path = os.path.join(save_dir, f\"env_metadata_{run_id}.txt\")\n",
+        "    with open(log_path, \"w\") as f:\n",
+        "        f.write(f\"PRISM EXPERIMENT METADATA | Run ID: {run_id}\\n{'='*50}\\n\")\n",
+        "        for k, v in config.items(): f.write(f\"{k}: {v}\\n\")\n",
+        "    print(f\"📝 Environment Snapshot saved to: {log_path}\")\n",
+        "\n",
+        "def log_metrics(save_dir, run_id, epoch, train_loss, val_loss, ppl):\n",
+        "    log_path = os.path.join(save_dir, f\"metrics_log_{run_id}.csv\")\n",
+        "    if not os.path.exists(log_path):\n",
+        "        with open(log_path, \"w\") as f: f.write(\"Timestamp,Epoch,Train_Loss,Val_Loss,Perplexity\\n\")\n",
+        "    with open(log_path, \"a\") as f:\n",
+        "        ts = datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n",
+        "        f.write(f\"{ts},{epoch},{train_loss:.6f},{val_loss:.6f},{ppl:.6f}\\n\")\n",
+        "\n",
+        "\n",
+        "def save_checkpoint(path, model, optimizer, scheduler, epoch, best_loss, config):\n",
+        "    torch.save({\n",
+        "        'epoch': epoch,\n",
+        "        'model_state_dict': model.state_dict(),\n",
+        "        'optimizer_state_dict': optimizer.state_dict(),\n",
+        "        'scheduler_state_dict': scheduler.state_dict(),\n",
+        "        'best_val_loss': best_loss,\n",
+        "        'config': config\n",
+        "    }, path)\n",
+        "\n",
+        "\n",
+        "def run_wikitext_training(experiment_name=\"PRISM2_WT103_40epochs\"):\n",
+        "    from google.colab import drive\n",
+        "    if not os.path.exists('/content/drive'): drive.mount('/content/drive')\n",
+        "\n",
+        "    # --- SETUP DIRS ---\n",
+        "    if RESUME_PATH and os.path.exists(RESUME_PATH):\n",
+        "        print(f\"🔄 RESUMING FROM: {RESUME_PATH}\")\n",
+        "        checkpoint = torch.load(RESUME_PATH, map_location=DEVICE)\n",
+        "        SAVE_DIR = os.path.dirname(RESUME_PATH)\n",
+        "        run_id = checkpoint.get('config', {}).get('run_id', 'resumed')\n",
+        "    else:\n",
+        "        run_id = hashlib.md5(datetime.now().strftime(\"%Y%m%d%H%M%S%f\").encode()).hexdigest()[:8]\n",
+        "        timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
+        "        folder_name = f\"{experiment_name}_{timestamp}_{run_id}\"\n",
+        "        SAVE_DIR = os.path.join(\"/content/drive/My Drive/PRISM_Experiments\", folder_name)\n",
+        "        os.makedirs(SAVE_DIR, exist_ok=True)\n",
+        "        print(f\"💾 Checkpoints: {SAVE_DIR}\")\n",
+        "\n",
+        "    writer = SummaryWriter(log_dir=SAVE_DIR)\n",
+        "    GRAD_ACCUM = 4\n",
+        "\n",
+        "    lm_datasets, data_collator = prepare_data_from_hub()\n",
+        "\n",
+        "    # WORKERS=2 (Safe for Colab)\n",
+        "    train_loader = DataLoader(\n",
+        "        lm_datasets[\"train\"], batch_size=BATCH_SIZE, shuffle=True,\n",
+        "        collate_fn=data_collator, num_workers=2, pin_memory=True,\n",
+        "        prefetch_factor=2, persistent_workers=True\n",
+        "    )\n",
+        "    valid_loader = DataLoader(\n",
+        "        lm_datasets[\"validation\"], batch_size=BATCH_SIZE,\n",
+        "        collate_fn=data_collator, num_workers=2, pin_memory=True\n",
+        "    )\n",
+        "    test_loader = DataLoader(\n",
+        "        lm_datasets[\"test\"], batch_size=BATCH_SIZE,\n",
+        "        collate_fn=data_collator, num_workers=2, pin_memory=True\n",
+        "    )\n",
+        "\n",
+        "    print(\"\\n⚡ INITIALIZING MODEL...\")\n",
+        "    model = PRISM_WikiText_Model(\n",
+        "        vocab_size=VOCAB_SIZE, d_model=D_MODEL, max_len=SEQ_LEN,\n",
+        "        prism_depth=DEPTH-1, trans_depth=1, dropout=DROPOUT\n",
+        "    ).to(DEVICE)\n",
+        "\n",
+        "    optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=0.0)\n",
+        "    total_steps = (len(train_loader) // GRAD_ACCUM) * EPOCHS\n",
+        "    scheduler = get_cosine_schedule_with_warmup(\n",
+        "        optimizer, num_warmup_steps=int(0.1 * total_steps), num_training_steps=total_steps\n",
+        "    )\n",
+        "    criterion = nn.CrossEntropyLoss()\n",
+        "\n",
+        "    start_epoch = 0\n",
+        "    best_val_loss = float('inf')\n",
+        "\n",
+        "    if RESUME_PATH and os.path.exists(RESUME_PATH):\n",
+        "        model.load_state_dict(checkpoint['model_state_dict'])\n",
+        "        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
+        "        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])\n",
+        "        start_epoch = checkpoint['epoch'] + 1\n",
+        "        best_val_loss = checkpoint['best_val_loss']\n",
+        "        del checkpoint\n",
+        "        torch.cuda.empty_cache()\n",
+        "    else:\n",
+        "        def init_weights_PRISM(m):\n",
+        "            if isinstance(m, nn.Linear):\n",
+        "                nn.init.xavier_uniform_(m.weight)\n",
+        "                if m.bias is not None: nn.init.zeros_(m.bias)\n",
+        "            elif isinstance(m, nn.Embedding):\n",
+        "                nn.init.normal_(m.weight, std=D_MODEL**-0.5)\n",
+        "        model.apply(init_weights_PRISM)\n",
+        "        nn.init.normal_(model.rose.rotation_predictor.weight, std=0.01)\n",
+        "        with torch.no_grad():\n",
+        "            model.rose.rotation_predictor.bias[:D_MODEL].fill_(1.0)\n",
+        "            model.rose.rotation_predictor.bias[D_MODEL:].fill_(0.0)\n",
+        "\n",
+        "    print(f\"\\n🚀 STARTING (Ep {start_epoch+1} to {EPOCHS})\")\n",
+        "    global_step = (len(train_loader) // GRAD_ACCUM) * start_epoch\n",
+        "\n",
+        "    for epoch in range(start_epoch, EPOCHS):\n",
+        "        model.train()\n",
+        "        pbar = tqdm(train_loader, desc=f\"Ep {epoch+1}/{EPOCHS}\")\n",
+        "\n",
+        "        for step, batch in enumerate(pbar):\n",
+        "            x, y = batch['input_ids'].to(DEVICE), batch['labels'].to(DEVICE)\n",
+        "\n",
+        "            # Forward\n",
+        "            loss = criterion(model(x).view(-1, VOCAB_SIZE), y.view(-1)) / GRAD_ACCUM\n",
+        "            loss.backward()\n",
+        "\n",
+        "            # Step (Every 4 batches)\n",
+        "            if (step + 1) % GRAD_ACCUM == 0:\n",
+        "                # 1. Calc Norm\n",
+        "                grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
+        "\n",
+        "                # 2. Step\n",
+        "                optimizer.step()\n",
+        "                scheduler.step()\n",
+        "                optimizer.zero_grad()\n",
+        "                global_step += 1\n",
+        "\n",
+        "                # 3. UPDATE BAR WITH GNORM\n",
+        "                actual_loss = loss.item() * GRAD_ACCUM\n",
+        "                writer.add_scalar('Train/Loss', actual_loss, global_step)\n",
+        "\n",
+        "                # <--- FIXED LINE HERE:\n",
+        "                pbar.set_postfix({\n",
+        "                    'loss': f\"{actual_loss:.4f}\",\n",
+        "                    'gnorm': f\"{grad_norm.item():.2f}\"\n",
+        "                })\n",
+        "\n",
+        "        # VALIDATION\n",
+        "        model.eval()\n",
+        "        val_loss = 0\n",
+        "        with torch.no_grad():\n",
+        "            for batch in valid_loader:\n",
+        "                x, y = batch['input_ids'].to(DEVICE), batch['labels'].to(DEVICE)\n",
+        "                val_loss += criterion(model(x).view(-1, VOCAB_SIZE), y.view(-1)).item()\n",
+        "\n",
+        "        avg_val_loss = val_loss / len(valid_loader)\n",
+        "        ppl = math.exp(avg_val_loss) if avg_val_loss < 100 else float('inf')\n",
+        "\n",
+        "        print(f\"✨ Epoch {epoch+1} | Val Loss: {avg_val_loss:.4f} | PPL: {ppl:.2f}\")\n",
+        "        writer.add_scalar('Val/PPL', ppl, epoch+1)\n",
+        "\n",
+        "        config_dump = {\"epoch\": epoch, \"run_id\": run_id}\n",
+        "        save_checkpoint(os.path.join(SAVE_DIR, \"last.pt\"), model, optimizer, scheduler, epoch, best_val_loss, config_dump)\n",
+        "\n",
+        "        if avg_val_loss < best_val_loss:\n",
+        "            best_val_loss = avg_val_loss\n",
+        "            torch.save(model.state_dict(), os.path.join(SAVE_DIR, \"best.pt\"))\n",
+        "            print(\"   🏆 New Best Model Saved!\")\n",
+        "\n",
+        "    # TEST\n",
+        "    best_path = os.path.join(SAVE_DIR, \"best.pt\")\n",
+        "    if os.path.exists(best_path):\n",
+        "        model.load_state_dict(torch.load(best_path))\n",
+        "        model.eval()\n",
+        "        test_loss = 0\n",
+        "        with torch.no_grad():\n",
+        "            for batch in tqdm(test_loader, desc=\"Testing\"):\n",
+        "                x, y = batch['input_ids'].to(DEVICE), batch['labels'].to(DEVICE)\n",
+        "                test_loss += criterion(model(x).view(-1, VOCAB_SIZE), y.view(-1)).item()\n",
+        "        print(f\"🏆 FINAL PPL: {math.exp(test_loss/len(test_loader)):.2f}\")\n",
+        "\n",
+        "    writer.close()\n",
+        "    return model"
+      ],
+      "metadata": {
+        "id": "-TNEv89gkS1k"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "def analyze_prism_params(model):\n",
+        "    print(\"=\"*80)\n",
+        "    print(f\"📊 LEAN PRISM-2 PARAMETER ANALYSIS\")\n",
+        "    print(\"=\"*80)\n",
+        "    total_params = sum(p.numel() for p in model.parameters())\n",
+        "    # Embeddings (Shared)\n",
+        "    vocab_params = model.rose.raw_embedding.weight.numel()\n",
+        "    # Wave Engine\n",
+        "    enc_params = sum(p.numel() for p in model.prism_encoder.parameters())\n",
+        "    # Transformer Refiner\n",
+        "    ref_params = sum(p.numel() for p in model.refiner.parameters()) if model.refiner else 0\n",
+        "    # Other\n",
+        "    other_params = total_params - vocab_params - enc_params - ref_params\n",
+        "\n",
+        "    print(f\"{'Shared Embeddings (Particle)':<35} | {vocab_params:<15,} | {vocab_params/total_params:.1%} | Tied\")\n",
+        "    print(f\"{'PRISM Optical Engine':<35} | {enc_params:<15,} | {enc_params/total_params:.1%} | 5 Layers\")\n",
+        "    print(f\"{'Digital Refiner':<35} | {ref_params:<15,} | {ref_params/total_params:.1%} | 1 Layer\")\n",
+        "    print(\"=\"*80)\n",
+        "    print(f\"{'TOTAL PARAMETERS':<35} | {total_params:<15,} | 100.0%\")\n",
+        "    print(\"=\"*80)\n",
+        "\n",
+        "\n",
+        "if __name__ == \"__main__\":\n",
+        "\n",
+        "    print(\"🏗️ INSTANTIATING MODEL FOR INSPECTION...\")\n",
+        "    # Initialize a temporary model just for counting\n",
+        "    dummy_model = PRISM_WikiText_Model(\n",
+        "        vocab_size=VOCAB_SIZE,\n",
+        "        d_model=D_MODEL,\n",
+        "        max_len=SEQ_LEN,\n",
+        "        prism_depth=DEPTH-1,\n",
+        "        trans_depth=1,\n",
+        "        dropout=DROPOUT\n",
+        "    )\n",
+        "\n",
+        "    # 2. Run Analysis\n",
+        "    analyze_prism_params(dummy_model)\n",
+        "\n",
+        "    # 3. Clean up to free RAM for actual training\n",
+        "    del dummy_model\n",
+        "    gc.collect()\n",
+        "    torch.cuda.empty_cache()\n",
+        "\n",
+        "    # 4. Ask for confirmation (Optional, or just proceed)\n",
+        "    print(\"\\n✅ Analysis Complete. Starting Training Routine in 5 seconds...\")\n",
+        "    import time\n",
+        "    time.sleep(5)\n",
+        "\n",
+        "    # 5. Start Training\n",
+        "    trained_prism = run_wikitext_training()\n",
+        "\n",
+        "    # 6. Final Analysis (Post-training check)\n",
+        "    analyze_prism_params(trained_prism)\n",
+        "\n",
+        "    # 7. Kill Runtime (Colab specific)\n",
+        "    from google.colab import runtime\n",
+        "    runtime.unassign()"
+      ],
+      "metadata": {
+        "id": "KaiJU0tPkVp-"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "from google.colab import runtime\n",
+        "runtime.unassign()"
+      ],
+      "metadata": {
+        "id": "bxFTYWHVqcSI"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}

WPT_Wikitext_103_Training.ipynb

@@ -0,0 +1,1061 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "A100"
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "code",
+      "source": [
+        "!pip install -q x-transformers\n",
+        "!pip install -q flash-attn --no-build-isolation"
+      ],
+      "metadata": {
+        "id": "6q9RTvlf5IiS"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import torch\n",
+        "import torch.nn as nn\n",
+        "import torch.nn.functional as F\n",
+        "import torch.optim as optim\n",
+        "import math\n",
+        "import os\n",
+        "import sys\n",
+        "import subprocess\n",
+        "import hashlib\n",
+        "import gc\n",
+        "from datetime import datetime\n",
+        "from tqdm.auto import tqdm\n",
+        "from torch.utils.data import DataLoader\n",
+        "from torch.utils.tensorboard import SummaryWriter\n",
+        "from transformers import RobertaTokenizerFast, get_cosine_schedule_with_warmup, DataCollatorForLanguageModeling\n",
+        "from datasets import load_dataset\n",
+        "from x_transformers import Encoder\n",
+        "\n",
+        "# ==========================================\n",
+        "# 1. CONFIGURATION\n",
+        "# ==========================================\n",
+        "# YOUR REPO ID (Created in previous step)\n",
+        "HF_ID = \"prism-lab/wikitext-103-prism-32k-seq4k\"\n",
+        "\n",
+        "# Hyperparameters\n",
+        "VOCAB_SIZE = 32768\n",
+        "SEQ_LEN = 4096\n",
+        "BATCH_SIZE = 8\n",
+        "EPOCHS = 40\n",
+        "LR = 1e-3\n",
+        "D_MODEL = 512\n",
+        "D_BRANCH = 256\n",
+        "DEPTH = 6\n",
+        "RESUME_PATH = None #\"/content/drive/MyDrive/PRISM_Experiments/PILLARS_SplitStream_8Layer_20260116_025321_8438ce62/last.pt\"\n",
+        "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+        "torch.set_float32_matmul_precision(\"high\")\n",
+        "\n",
+        "# ==========================================\n",
+        "# 2. DATA PIPELINE (The \"Pro\" Way)\n",
+        "# ==========================================\n",
+        "def prepare_data_from_hub():\n",
+        "    print(f\"⬇️ Pulling Pre-Tokenized Data from {HF_ID}...\")\n",
+        "\n",
+        "    # 1. Load Tokenizer (Instant)\n",
+        "    # This pulls the exact tokenizer you uploaded\n",
+        "    tokenizer = RobertaTokenizerFast.from_pretrained(HF_ID)\n",
+        "\n",
+        "    # 2. Load Dataset (Instant)\n",
+        "    # This pulls the already chunked/tokenized data\n",
+        "    dataset = load_dataset(HF_ID)\n",
+        "\n",
+        "    print(f\"✅ Loaded {len(dataset['train'])} training chunks.\")\n",
+        "\n",
+        "    # 3. Collator\n",
+        "    data_collator = DataCollatorForLanguageModeling(\n",
+        "        tokenizer=tokenizer,\n",
+        "        mlm=True,\n",
+        "        mlm_probability=0.15\n",
+        "    )\n",
+        "\n",
+        "    return dataset, data_collator\n",
+        "# ==========================================\n",
+        "# 3. PRISM ARCHITECTURE (Complex-Valued)\n",
+        "# ==========================================\n",
+        "\n",
+        "class ComplexDropout(nn.Module):\n",
+        "    def __init__(self, p=0.5):\n",
+        "        super().__init__()\n",
+        "        self.p = p\n",
+        "    def forward(self, z):\n",
+        "        if not self.training or self.p == 0.0: return z\n",
+        "        mask = torch.ones_like(z.real)\n",
+        "        mask = F.dropout(mask, self.p, self.training, inplace=False)\n",
+        "        return z * mask\n",
+        "\n",
+        "class RobustPhaseNorm(nn.Module):\n",
+        "    def __init__(self, d_model, eps=1e-5):\n",
+        "        super().__init__()\n",
+        "        self.scale = nn.Parameter(torch.ones(d_model))\n",
+        "        self.eps = eps\n",
+        "    def forward(self, x):\n",
+        "        mag = torch.abs(x)\n",
+        "        rms = torch.sqrt(torch.mean(mag**2, dim=-1, keepdim=True) + self.eps)\n",
+        "        return (x / rms) * self.scale\n",
+        "\n",
+        "class ModReLU(nn.Module):\n",
+        "    def __init__(self, features):\n",
+        "        super().__init__()\n",
+        "        self.b = nn.Parameter(torch.zeros(features))\n",
+        "\n",
+        "    def forward(self, z):\n",
+        "        # 1. FORCE FLOAT32 FOR GEOMETRY\n",
+        "        # We must calculate magnitude in high precision to prevent\n",
+        "        # square-law overflow (Re^2 + Im^2) from killing the gradients.\n",
+        "        z_32 = z.to(torch.complex64)\n",
+        "\n",
+        "        # 2. Calculate Magnitude (Safe)\n",
+        "        mag = torch.abs(z_32)\n",
+        "\n",
+        "        # 3. Activation Logic (Still FP32)\n",
+        "        new_mag = F.relu(mag + self.b.float())\n",
+        "\n",
+        "        # 4. Reconstruct Phase (Safe Division)\n",
+        "        # (z / mag) is the unit vector (phase)\n",
+        "        phase = z_32 / (mag + 1e-6)\n",
+        "\n",
+        "        # 5. Result\n",
+        "        out = new_mag * phase\n",
+        "\n",
+        "        # 6. Cast back to network dtype (BF16/FP16)\n",
+        "        return out.to(z.dtype)\n",
+        "\n",
+        "class ComplexToRealBridge(nn.Module):\n",
+        "    def __init__(self, d_model):\n",
+        "        super().__init__()\n",
+        "        self.proj = nn.Linear(d_model * 2, d_model)\n",
+        "        self.norm = nn.LayerNorm(d_model)\n",
+        "    def forward(self, x_complex):\n",
+        "        cat = torch.cat([x_complex.real, x_complex.imag], dim=-1)\n",
+        "        return self.norm(self.proj(cat))\n",
+        "\n",
+        "# ==========================================\n",
+        "# 4. DYNAMIC RoSE (Mamba-3 Engine)\n",
+        "# ==========================================\n",
+        "class DynamicRoSE(nn.Module):\n",
+        "    def __init__(self, num_embeddings, embedding_dim, max_period=10000.0):\n",
+        "        super().__init__()\n",
+        "        self.embedding_dim = embedding_dim\n",
+        "\n",
+        "        # 1. Master Real Embedding (The \"Particle\")\n",
+        "        self.raw_embedding = nn.Embedding(num_embeddings, embedding_dim)\n",
+        "\n",
+        "        # 2. Complex Adapter (The \"Wave\" Magnitude/Initial Phase)\n",
+        "        self.adapter = nn.Linear(embedding_dim, embedding_dim * 2)\n",
+        "\n",
+        "        # 3. Static Frequencies (Positional)\n",
+        "        freqs = torch.exp(torch.arange(0, embedding_dim, dtype=torch.float32) * -(math.log(max_period) / embedding_dim))\n",
+        "        self.register_buffer('freqs', freqs)\n",
+        "\n",
+        "        self.rotation_predictor = nn.Linear(embedding_dim, embedding_dim * 2)\n",
+        "\n",
+        "    def forward(self, input_ids):\n",
+        "        # A. Raw Particle\n",
+        "        real_base = self.raw_embedding(input_ids)\n",
+        "        B, L, D = real_base.shape\n",
+        "\n",
+        "        # B. Complex Wave Content\n",
+        "        complex_params = self.adapter(real_base)\n",
+        "        z_t = torch.complex(complex_params[..., :D], complex_params[..., D:])\n",
+        "\n",
+        "        rot_raw = self.rotation_predictor(real_base)\n",
+        "        rot_x, rot_y = rot_raw.chunk(2, dim=-1)\n",
+        "\n",
+        "        rot_mag = torch.sqrt(rot_x**2 + rot_y**2 + 1e-6)\n",
+        "        dynamic_rot = torch.complex(rot_x / rot_mag, rot_y / rot_mag)\n",
+        "\n",
+        "        # D. Static Positional Rotation\n",
+        "        pos = torch.arange(L, device=input_ids.device).float()\n",
+        "        static_angles = torch.outer(pos, self.freqs) # [L, D]\n",
+        "        static_rot = torch.polar(torch.ones_like(static_angles), static_angles) # [L, D]\n",
+        "\n",
+        "        z_final = z_t * static_rot.unsqueeze(0) * dynamic_rot\n",
+        "\n",
+        "        return z_final, real_base\n",
+        "\n",
+        "# ==========================================\n",
+        "# 5. HYENA FILTER\n",
+        "# ==========================================\n",
+        "class HyenaNeuralFilter(nn.Module):\n",
+        "    def __init__(self, d_model, max_len=1024, hidden_dim=64):\n",
+        "        super().__init__()\n",
+        "        self.d_model = d_model\n",
+        "        freqs = torch.exp(torch.arange(0, hidden_dim, 2, dtype=torch.float32) * -(math.log(10000.0) / hidden_dim))\n",
+        "        self.register_buffer(\"freqs\", freqs)\n",
+        "        self.mlp = nn.Sequential(\n",
+        "            nn.Linear(hidden_dim, hidden_dim), nn.SiLU(),\n",
+        "            nn.Linear(hidden_dim, hidden_dim), nn.SiLU(),\n",
+        "            nn.Linear(hidden_dim, d_model * 2)\n",
+        "        )\n",
+        "    def forward(self, L, device):\n",
+        "        t = torch.linspace(0, 1, steps=L, device=device).unsqueeze(-1)\n",
+        "        emb = torch.cat([torch.sin(t * self.freqs), torch.cos(t * self.freqs)], dim=-1)\n",
+        "        out = self.mlp(emb).view(L, self.d_model, 2)\n",
+        "        return torch.complex(out[..., 0], out[..., 1])\n",
+        "\n",
+        "# ==========================================\n",
+        "# 6. GATED HARMONIC CONVOLUTION (Lean)\n",
+        "# ==========================================\n",
+        "# @title 🛠️ Fixed PRISM Layer (Precision-Gated)\n",
+        "\n",
+        "# @title 🛠️ Fixed PRISM Layer (Type-Safe)\n",
+        "\n",
+        "class GatedHarmonicConvolution(nn.Module):\n",
+        "    def __init__(self, d_model, max_len=1024, dropout=0.1):\n",
+        "        super().__init__()\n",
+        "        self.d_model = d_model\n",
+        "        self.filter_len = max_len\n",
+        "        self.neural_filter = HyenaNeuralFilter(d_model, max_len=max_len)\n",
+        "        self.gate_proj = nn.Linear(d_model * 2, d_model * 2)\n",
+        "\n",
+        "        self.mix_real = nn.Linear(d_model, d_model)\n",
+        "        self.mix_imag = nn.Linear(d_model, d_model)\n",
+        "        self.out_real = nn.Linear(d_model, d_model)\n",
+        "        self.out_imag = nn.Linear(d_model, d_model)\n",
+        "\n",
+        "        self.activation = ModReLU(d_model)\n",
+        "        self.norm = RobustPhaseNorm(d_model)\n",
+        "        self.dropout = ComplexDropout(dropout)\n",
+        "\n",
+        "    def forward(self, x, src_mask=None):\n",
+        "        residual = x\n",
+        "        x_norm = self.norm(x)\n",
+        "        if src_mask is not None:\n",
+        "             x_norm = x_norm.masked_fill(src_mask.unsqueeze(-1), 0.0)\n",
+        "\n",
+        "        # 🛑 PRECISION GATE 🛑\n",
+        "        # Force operations to Float32 Complex to preserve Phase Physics\n",
+        "        with torch.amp.autocast('cuda', enabled=False):\n",
+        "\n",
+        "            # --- THE FIX IS HERE ---\n",
+        "            # Old: x_32 = x_norm.float()  <-- This stripped the imaginary part\n",
+        "            # New: Explicit cast to Complex64\n",
+        "            x_32 = x_norm.to(torch.complex64)\n",
+        "            # -----------------------\n",
+        "\n",
+        "            B, L, D = x_32.shape\n",
+        "            eff_L = min(L, self.filter_len)\n",
+        "\n",
+        "            # 1. FFT (Now safe because x_32 is definitely complex)\n",
+        "            x_freq = torch.fft.fft(x_32, n=eff_L, dim=1, norm='ortho')\n",
+        "\n",
+        "            # 2. Filter (Ensure filter is also complex64)\n",
+        "            h = self.neural_filter(eff_L, x.device).unsqueeze(0).to(torch.complex64)\n",
+        "            x_filtered = x_freq * h\n",
+        "\n",
+        "            # 3. IFFT\n",
+        "            x_time = torch.fft.ifft(x_filtered, n=eff_L, dim=1, norm='ortho')\n",
+        "\n",
+        "            if L > eff_L: x_time = F.pad(x_time, (0,0,0,L-eff_L))\n",
+        "            else: x_time = x_time[:, :L, :]\n",
+        "\n",
+        "            # 4. Gating (Sigmoid logic)\n",
+        "            # Safe concatenation because x_32 is complex64\n",
+        "            x_cat = torch.cat([x_32.real, x_32.imag], dim=-1)\n",
+        "\n",
+        "            # Cast weights to Float32 for the calculation\n",
+        "            gate_w = self.gate_proj.weight.to(torch.float32)\n",
+        "            gate_b = self.gate_proj.bias.to(torch.float32)\n",
+        "\n",
+        "            gate_out = F.linear(x_cat, gate_w, gate_b)\n",
+        "            gates = torch.sigmoid(gate_out)\n",
+        "\n",
+        "            g_r, g_i = gates.chunk(2, dim=-1)\n",
+        "            x_gated_32 = torch.complex(x_time.real * g_r, x_time.imag * g_i)\n",
+        "\n",
+        "            # 🏁 EXIT GATE: Cast back to original dtype (likely BFloat16 from autocast)\n",
+        "            # We cast real/imag separately to be safe\n",
+        "            target_dtype = x.dtype\n",
+        "            # If x was complex, target is complex. If x was real, we have an issue.\n",
+        "            # Assuming x comes from autocast, it might be complex16.\n",
+        "\n",
+        "            x_gated = x_gated_32.to(target_dtype)\n",
+        "\n",
+        "        # 5. Mixing (Back in mixed precision)\n",
+        "        mr, mi = self.mix_real, self.mix_imag\n",
+        "        x_mixed = torch.complex(mr(x_gated.real) - mi(x_gated.imag), mr(x_gated.imag) + mi(x_gated.real))\n",
+        "\n",
+        "        x_act = self.activation(x_mixed)\n",
+        "\n",
+        "        or_, oi = self.out_real, self.out_imag\n",
+        "        out = torch.complex(or_(x_act.real) - oi(x_act.imag), or_(x_act.imag) + oi(x_act.real))\n",
+        "\n",
+        "        return self.dropout(out) + residual\n",
+        "# ==========================================\n",
+        "# 7. MODEL WRAPPERS\n",
+        "# ==========================================\n",
+        "class PRISMEncoder(nn.Module):\n",
+        "    def __init__(self, num_layers, d_model, max_len, dropout=0.1):\n",
+        "        super().__init__()\n",
+        "        self.layers = nn.ModuleList([\n",
+        "            GatedHarmonicConvolution(d_model, max_len, dropout)\n",
+        "            for _ in range(num_layers)\n",
+        "        ])\n",
+        "        self.final_norm = RobustPhaseNorm(d_model)\n",
+        "    def forward(self, x, src_mask=None):\n",
+        "        for layer in self.layers:\n",
+        "            if self.training: x = torch.utils.checkpoint.checkpoint(layer, x, src_mask, use_reentrant=False)\n",
+        "            else: x = layer(x, src_mask)\n",
+        "        return self.final_norm(x)\n",
+        "\n",
+        "class PRISM_WikiText_Model(nn.Module):\n",
+        "    def __init__(self, vocab_size, d_model, max_len, prism_depth=5, trans_depth=1, dropout=0.1):\n",
+        "        super().__init__()\n",
+        "        self.d_model = d_model\n",
+        "\n",
+        "        # 1. PRISM Core (The Optical/Passive Part)\n",
+        "        self.rose = DynamicRoSE(vocab_size, d_model)\n",
+        "        self.prism_encoder = PRISMEncoder(prism_depth, d_model, max_len=max_len, dropout=dropout)\n",
+        "        self.bridge = ComplexToRealBridge(d_model)\n",
+        "        self.periscope_proj = nn.Sequential(nn.Linear(d_model * 2, d_model), nn.LayerNorm(d_model), nn.GELU())\n",
+        "\n",
+        "        # 2. Refiner (The Digital/Active Part)\n",
+        "        # 🔄 SWAPPED: Replaced Standard Transformer with RoPE-Enabled Encoder\n",
+        "        if trans_depth > 0:\n",
+        "            self.refiner = Encoder(\n",
+        "                dim=d_model,\n",
+        "                depth=trans_depth,\n",
+        "                heads=8,\n",
+        "                rotary_pos_emb=True,\n",
+        "                attn_flash=True,\n",
+        "                attn_dropout=dropout,\n",
+        "                ff_dropout=dropout,\n",
+        "\n",
+        "            )\n",
+        "        else:\n",
+        "            self.refiner = None\n",
+        "\n",
+        "        # 3. Output\n",
+        "        self.lm_head = nn.Linear(d_model, vocab_size)\n",
+        "        self.lm_head.weight = self.rose.raw_embedding.weight\n",
+        "\n",
+        "    def forward(self, input_ids):\n",
+        "        # A. Wave Physics\n",
+        "        wave_src, particle_src = self.rose(input_ids)\n",
+        "        wave_out = self.prism_encoder(wave_src)\n",
+        "        wave_real = self.bridge(wave_out)\n",
+        "\n",
+        "        # B. Interface\n",
+        "        mixed_memory = self.periscope_proj(torch.cat([wave_real, particle_src], dim=-1))\n",
+        "\n",
+        "        # C. Digital Refinement (Now with RoPE)\n",
+        "        if self.refiner:\n",
+        "            out = self.refiner(mixed_memory)\n",
+        "        else:\n",
+        "            out = mixed_memory\n",
+        "\n",
+        "        return self.lm_head(out)\n",
+        "\n",
+        "# ==========================================\n",
+        "# 1. SENSORY STREAM (Transformer + RoPE)\n",
+        "# ==========================================\n",
+        "class SensoryStream(nn.Module):\n",
+        "    def __init__(self, depth, d_model, dropout=0.1):\n",
+        "        super().__init__()\n",
+        "        self.encoder = Encoder(\n",
+        "            dim=d_model,\n",
+        "            depth=depth,\n",
+        "            heads=4,                # 256 dim / 64 head_dim = 4 heads\n",
+        "            attn_flash=True,        # Flash Attention\n",
+        "            rotary_pos_emb=True,    # <--- CRITICAL: RoPE Enabled\n",
+        "            attn_dropout=dropout,\n",
+        "            ff_dropout=dropout,\n",
+        "            use_rmsnorm=True,       # RMSNorm (Llama style)\n",
+        "            ff_glu=True             # SwiGLU (Llama style)\n",
+        "        )\n",
+        "\n",
+        "    def forward(self, x):\n",
+        "        return self.encoder(x)\n",
+        "\n",
+        "# ==========================================\n",
+        "# 2. PILLARS-DAT (Dual Attention with RoPE)\n",
+        "# ==========================================\n",
+        "class Pillars_DAT(nn.Module):\n",
+        "    def __init__(self, vocab_size, d_model=512, d_branch=256, seq_len=4096, depth=4):\n",
+        "        super().__init__()\n",
+        "        self.d_model = d_model\n",
+        "        self.d_branch = d_branch\n",
+        "\n",
+        "        # --- A. SHARED ROOT ---\n",
+        "        self.rose = DynamicRoSE(vocab_size, d_model)\n",
+        "\n",
+        "        # --- B. DOWNSAMPLE ---\n",
+        "        self.particle_down = nn.Linear(d_model, d_branch)\n",
+        "        self.wave_down = nn.Linear(d_model * 2, d_branch * 2)\n",
+        "\n",
+        "        # --- C. STREAM 1: SENSORY (Object Attributes) ---\n",
+        "        # REPLACED: FNet -> Transformer with RoPE\n",
+        "        # NOTE: No self.sensory_pos anymore! RoPE handles it.\n",
+        "        self.stream_sensory = SensoryStream(depth=depth, d_model=d_branch, dropout=0.1)\n",
+        "\n",
+        "        # --- D. STREAM 2: RELATIONAL (Structure / Phase) ---\n",
+        "        # PRISM handles positions internally via RoSE frequencies\n",
+        "        self.stream_relational = PRISMEncoder(num_layers=depth, d_model=d_branch, max_len=seq_len, dropout=0.1)\n",
+        "        self.relational_bridge = ComplexToRealBridge(d_branch)\n",
+        "\n",
+        "        # --- E. FUSION ---\n",
+        "        self.fusion_proj = nn.Linear(d_branch * 2, d_model)\n",
+        "        self.fusion_norm = nn.LayerNorm(d_model)\n",
+        "\n",
+        "        # --- F. REFINER ---\n",
+        "        self.refiner = Encoder(\n",
+        "            dim=d_model, depth=1, heads=8, attn_flash=True,\n",
+        "            rotary_pos_emb=True, attn_dropout=0.1, ff_dropout=0.1\n",
+        "        )\n",
+        "\n",
+        "        # --- G. OUTPUT ---\n",
+        "        self.head_bias = nn.Parameter(torch.zeros(vocab_size))\n",
+        "\n",
+        "    def forward(self, input_ids):\n",
+        "        # 1. Root Physics\n",
+        "        wave_src, particle_src = self.rose(input_ids)\n",
+        "\n",
+        "        # 2. Downsample\n",
+        "        p_small = self.particle_down(particle_src)\n",
+        "\n",
+        "        # Prepare complex wave input\n",
+        "        w_flat = torch.cat([wave_src.real, wave_src.imag], dim=-1)\n",
+        "        w_small_flat = self.wave_down(w_flat)\n",
+        "        w_small = torch.complex(w_small_flat[..., :self.d_branch], w_small_flat[..., self.d_branch:])\n",
+        "\n",
+        "        # 3. Parallel Processing\n",
+        "\n",
+        "        # --- Stream A: Sensory (Transformer + RoPE) ---\n",
+        "        # Pass pure features. RoPE adds position info inside the attention layer.\n",
+        "        sensory_out = self.stream_sensory(p_small)\n",
+        "\n",
+        "        # --- Stream B: Relational (PRISM) ---\n",
+        "        relational_out_complex = self.stream_relational(w_small)\n",
+        "        relational_out = self.relational_bridge(relational_out_complex)\n",
+        "\n",
+        "        # 4. Integration\n",
+        "        stacked = torch.cat([sensory_out, relational_out], dim=-1)\n",
+        "        context = self.fusion_norm(self.fusion_proj(stacked))\n",
+        "\n",
+        "        # 5. Refinement\n",
+        "        refined = self.refiner(context)\n",
+        "\n",
+        "        # 6. Output\n",
+        "        logits = F.linear(refined, self.rose.raw_embedding.weight, self.head_bias)\n",
+        "\n",
+        "        return logits\n",
+        "\n",
+        "import torch\n",
+        "import torch.nn as nn\n",
+        "from prettytable import PrettyTable # Optional, but makes tables nice.\n",
+        "# If you don't have prettytable, the code below uses standard f-strings.\n",
+        "\n",
+        "import torch\n",
+        "import torch.nn as nn\n",
+        "\n",
+        "import torch\n",
+        "import torch.nn as nn\n",
+        "\n",
+        "def deep_analyze_pillars(model):\n",
+        "    def get_p(obj):\n",
+        "        \"\"\"Safely returns parameter count for Modules OR raw Parameters.\"\"\"\n",
+        "        if isinstance(obj, nn.Parameter):\n",
+        "            return obj.numel()\n",
+        "        return sum(p.numel() for p in obj.parameters() if p.requires_grad)\n",
+        "\n",
+        "    def format_num(n):\n",
+        "        if n > 1e6: return f\"{n/1e6:.2f}M\"\n",
+        "        if n > 1e3: return f\"{n/1e3:.2f}K\"\n",
+        "        return str(n)\n",
+        "\n",
+        "    print(\"\\n\" + \"=\"*80)\n",
+        "    print(f\"🏗️  PILLARS (COMPACT) - DEEP LAYER ANALYSIS\")\n",
+        "    print(\"=\"*80)\n",
+        "    print(f\"{'MODULE / LAYER':<40} | {'PARAMS':<15} | {'TYPE'}\")\n",
+        "    print(\"-\" * 80)\n",
+        "\n",
+        "    total_params = get_p(model)\n",
+        "\n",
+        "    # -----------------------------------------------\n",
+        "    # 1. STATIC MEMORY (Embeddings)\n",
+        "    # -----------------------------------------------\n",
+        "    vocab_emb = get_p(model.rose.raw_embedding)\n",
+        "    fnet_pos  = get_p(model.fnet_pos)\n",
+        "\n",
+        "    print(f\"{'Shared Vocab Embedding':<40} | {format_num(vocab_emb):<15} | 💾 STORAGE\")\n",
+        "    print(f\"{'FNet Positional Embedding':<40} | {format_num(fnet_pos):<15} | 💾 STORAGE\")\n",
+        "\n",
+        "    # -----------------------------------------------\n",
+        "    # 2. INPUT LOGIC (RoSE & Downsampling)\n",
+        "    # -----------------------------------------------\n",
+        "    rose_total = get_p(model.rose)\n",
+        "    rose_logic = rose_total - vocab_emb # Subtract the embedding matrix we already counted\n",
+        "\n",
+        "    print(\"-\" * 80)\n",
+        "    print(f\"{'Dynamic RoSE (Adapters)':<40} | {format_num(rose_logic):<15} | 🌊 PHASE INIT\")\n",
+        "    print(f\"{'Particle Downsample (512->384)':<40} | {format_num(get_p(model.particle_down)):<15} | 📉 PROJ\")\n",
+        "    print(f\"{'Wave Downsample (1024->768)':<40} | {format_num(get_p(model.wave_down)):<15} | 📉 PROJ\")\n",
+        "\n",
+        "    # -----------------------------------------------\n",
+        "    # 3. STREAM A: RATE (FNet)\n",
+        "    # -----------------------------------------------\n",
+        "    print(\"-\" * 80)\n",
+        "    print(f\"TRACK A: RATE STREAM (FNet) - Depth {len(model.stream_rate.layers)}\")\n",
+        "\n",
+        "    fnet_encoder_total = 0\n",
+        "    for i, layer in enumerate(model.stream_rate.layers):\n",
+        "        p = get_p(layer)\n",
+        "        fnet_encoder_total += p\n",
+        "        print(f\"  ├─ FNet Block {i:<24} | {format_num(p):<15} | ⚡ RATE\")\n",
+        "\n",
+        "    fnet_norm = get_p(model.stream_rate.norm_out)\n",
+        "    fnet_encoder_total += fnet_norm\n",
+        "    print(f\"  └─ Final Norm {i:<24} | {format_num(fnet_norm):<15} | ⚡ RATE\")\n",
+        "\n",
+        "    # -----------------------------------------------\n",
+        "    # 4. STREAM B: PHASE (PRISM)\n",
+        "    # -----------------------------------------------\n",
+        "    print(\"-\" * 80)\n",
+        "    print(f\"TRACK B: PHASE STREAM (PRISM) - Depth {len(model.stream_phase.layers)}\")\n",
+        "\n",
+        "    prism_encoder_total = 0\n",
+        "    for i, layer in enumerate(model.stream_phase.layers):\n",
+        "        p = get_p(layer)\n",
+        "        prism_encoder_total += p\n",
+        "        print(f\"  ├─ PRISM Block {i:<23} | {format_num(p):<15} | 🌊 PHASE\")\n",
+        "\n",
+        "    prism_norm = get_p(model.stream_phase.final_norm)\n",
+        "    prism_encoder_total += prism_norm\n",
+        "    print(f\"  └─ Final Norm {i:<24} | {format_num(prism_norm):<15} | 🌊 PHASE\")\n",
+        "\n",
+        "    bridge_p = get_p(model.phase_bridge)\n",
+        "    print(f\"{'Phase Bridge (Complex->Real)':<40} | {format_num(bridge_p):<15} | 🌉 BRIDGE\")\n",
+        "\n",
+        "    # -----------------------------------------------\n",
+        "    # 5. THE BRAIN (Fusion & Refiner)\n",
+        "    # -----------------------------------------------\n",
+        "    print(\"-\" * 80)\n",
+        "    fusion_p = get_p(model.fusion_proj) + get_p(model.fusion_norm)\n",
+        "    print(f\"{'Fusion (Concat -> Proj -> Norm)':<40} | {format_num(fusion_p):<15} | 🧠 FUSION\")\n",
+        "\n",
+        "    refiner_p = get_p(model.refiner)\n",
+        "    print(f\"{'Transformer Refiner (1 Layer)':<40} | {format_num(refiner_p):<15} | 🧠 ATTENTION\")\n",
+        "\n",
+        "    # [FIX] Handle nn.Parameter directly\n",
+        "    head_bias_p = get_p(model.head_bias)\n",
+        "    print(f\"{'Output Head Bias':<40} | {format_num(head_bias_p):<15} | 🎯 OUTPUT\")\n",
+        "\n",
+        "    # -----------------------------------------------\n",
+        "    # 6. SUMMARY\n",
+        "    # -----------------------------------------------\n",
+        "    print(\"=\"*80)\n",
+        "\n",
+        "    storage = vocab_emb + fnet_pos + head_bias_p\n",
+        "    active = total_params - storage\n",
+        "\n",
+        "    print(f\"TOTAL PARAMETERS:      {total_params/1e6:.2f} M\")\n",
+        "    print(f\"   ├─ 💾 Storage:      {storage/1e6:.2f} M  (Embeddings)\")\n",
+        "    print(f\"   └─ 🧠 Compute:      {active/1e6:.2f} M  (Logic/Weights)\")\n",
+        "    print(\"-\" * 80)\n",
+        "    print(f\"STREAM BREAKDOWN:\")\n",
+        "    print(f\"   ├─ ⚡ Rate Stream:   {fnet_encoder_total/1e6:.2f} M\")\n",
+        "    print(f\"   └─ 🌊 Phase Stream:  {prism_encoder_total/1e6:.2f} M\")\n",
+        "    print(\"=\"*80 + \"\\n\")\n",
+        "\n",
+        "    return total_params\n"
+      ],
+      "metadata": {
+        "id": "V7DOwmmUjyin"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "\n",
+        "# Run the parameter analysis to confirm strict adherence to budget\n",
+        "def deep_analyze_pillars_dat(model):\n",
+        "    def get_p(obj):\n",
+        "        if isinstance(obj, nn.Parameter): return obj.numel()\n",
+        "        return sum(p.numel() for p in obj.parameters() if p.requires_grad)\n",
+        "\n",
+        "    def format_num(n):\n",
+        "        if n > 1e6: return f\"{n/1e6:.2f}M\"\n",
+        "        if n > 1e3: return f\"{n/1e3:.2f}K\"\n",
+        "        return str(n)\n",
+        "\n",
+        "    print(\"\\n\" + \"=\"*80)\n",
+        "    print(f\"🏛️  PILLARS-DAT (Hybrid Transformer-PRISM) - ANALYSIS\")\n",
+        "    print(\"=\"*80)\n",
+        "    print(f\"{'MODULE / LAYER':<40} | {'PARAMS':<12} | {'TYPE'}\")\n",
+        "    print(\"-\" * 80)\n",
+        "\n",
+        "    total_params = get_p(model)\n",
+        "\n",
+        "    # --- 1. MEMORY ---\n",
+        "    vocab_emb = get_p(model.rose.raw_embedding)\n",
+        "    print(f\"{'Shared Vocab Embedding':<40} | {format_num(vocab_emb):<12} | 💾 STORAGE\")\n",
+        "\n",
+        "    # --- 2. INPUT PHYSICS ---\n",
+        "    rose_logic = get_p(model.rose) - vocab_emb\n",
+        "    print(f\"{'Dynamic RoSE (Adapters)':<40} | {format_num(rose_logic):<12} | 🌊 PHYSICS\")\n",
+        "\n",
+        "    down_p = get_p(model.particle_down) + get_p(model.wave_down)\n",
+        "    print(f\"{'Stream Splitters (Downsample)':<40} | {format_num(down_p):<12} | 📉 PROJ\")\n",
+        "\n",
+        "    # --- 3. STREAM A: SENSORY (TRANSFORMER) ---\n",
+        "    print(\"-\" * 80)\n",
+        "    print(f\"STREAM A: SENSORY (Identity/Magnitude)\")\n",
+        "    sensory_p = get_p(model.stream_sensory)\n",
+        "    # Attempt to count depth if accessible, else generic\n",
+        "    try:\n",
+        "        depth_s = len(model.stream_sensory.encoder.layers)\n",
+        "        print(f\"  ├─ Transformer Encoder (Depth {depth_s})     | {format_num(sensory_p):<12} | ⚡ ATTENTION\")\n",
+        "    except:\n",
+        "        print(f\"  ├─ Transformer Encoder (Fused)           | {format_num(sensory_p):<12} | ⚡ ATTENTION\")\n",
+        "\n",
+        "    # --- 4. STREAM B: RELATIONAL (PRISM) ---\n",
+        "    print(\"-\" * 80)\n",
+        "    print(f\"STREAM B: RELATIONAL (Structure/Phase)\")\n",
+        "    relational_core = get_p(model.stream_relational)\n",
+        "    relational_bridge = get_p(model.relational_bridge)\n",
+        "\n",
+        "    try:\n",
+        "        depth_r = len(model.stream_relational.layers)\n",
+        "        print(f\"  ├─ PRISM Encoder (Depth {depth_r})           | {format_num(relational_core):<12} | 🌊 SPECTRAL\")\n",
+        "    except:\n",
+        "        print(f\"  ├─ PRISM Encoder (Fused)                 | {format_num(relational_core):<12} | 🌊 SPECTRAL\")\n",
+        "\n",
+        "    print(f\"  └─ Bridge (Complex->Real)                | {format_num(relational_bridge):<12} | 🌉 PROJ\")\n",
+        "\n",
+        "    # --- 5. FUSION & OUTPUT ---\n",
+        "    print(\"-\" * 80)\n",
+        "    fusion_p = get_p(model.fusion_proj) + get_p(model.fusion_norm)\n",
+        "    print(f\"{'Fusion (Concat -> Proj)':<40} | {format_num(fusion_p):<12} | 🧠 MIX\")\n",
+        "\n",
+        "    refiner_p = get_p(model.refiner)\n",
+        "    print(f\"{'Refiner (1-Layer Transformer)':<40} | {format_num(refiner_p):<12} | 🧠 REASONING\")\n",
+        "\n",
+        "    bias_p = get_p(model.head_bias)\n",
+        "    print(f\"{'Output Head Bias':<40} | {format_num(bias_p):<12} | 🎯 OUT\")\n",
+        "\n",
+        "    # --- SUMMARY ---\n",
+        "    print(\"=\"*80)\n",
+        "    storage = vocab_emb + bias_p\n",
+        "    active = total_params - storage\n",
+        "\n",
+        "    print(f\"TOTAL PARAMETERS:       {total_params/1e6:.2f} M\")\n",
+        "    print(f\"   ├─ 💾 Storage:       {storage/1e6:.2f} M  (Embeddings)\")\n",
+        "    print(f\"   └─ 🧠 Compute:       {active/1e6:.2f} M  (Active Weights)\")\n",
+        "    print(\"-\" * 80)\n",
+        "    print(f\"RATIO CHECK:\")\n",
+        "    print(f\"   ⚡ Sensory (Transf): {sensory_p/1e6:.2f} M\")\n",
+        "    print(f\"   🌊 Relation (PRISM): {(relational_core + relational_bridge)/1e6:.2f} M\")\n",
+        "    print(\"=\"*80 + \"\\n\")\n"
+      ],
+      "metadata": {
+        "id": "ke4fYT8UX5zH"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# ==========================================\n",
+        "# 4. LOGGING & ANALYSIS UTILITIES\n",
+        "# ==========================================\n",
+        "def deep_analyze_pillars_dat(model):\n",
+        "    def get_p(obj):\n",
+        "        if isinstance(obj, nn.Parameter): return obj.numel()\n",
+        "        return sum(p.numel() for p in obj.parameters() if p.requires_grad)\n",
+        "\n",
+        "    def format_num(n):\n",
+        "        if n > 1e6: return f\"{n/1e6:.2f}M\"\n",
+        "        if n > 1e3: return f\"{n/1e3:.2f}K\"\n",
+        "        return str(n)\n",
+        "\n",
+        "    print(\"\\n\" + \"=\"*80)\n",
+        "    print(f\"🏛️  PILLARS-DAT (Hybrid Transformer-PRISM) - ANALYSIS\")\n",
+        "    print(\"=\"*80)\n",
+        "    print(f\"{'MODULE / LAYER':<40} | {'PARAMS':<12} | {'TYPE'}\")\n",
+        "    print(\"-\" * 80)\n",
+        "\n",
+        "    total_params = get_p(model)\n",
+        "\n",
+        "    # --- 1. MEMORY ---\n",
+        "    vocab_emb = get_p(model.rose.raw_embedding)\n",
+        "    print(f\"{'Shared Vocab Embedding':<40} | {format_num(vocab_emb):<12} | 💾 STORAGE\")\n",
+        "\n",
+        "    # --- 2. INPUT PHYSICS ---\n",
+        "    rose_logic = get_p(model.rose) - vocab_emb\n",
+        "    print(f\"{'Dynamic RoSE (Adapters)':<40} | {format_num(rose_logic):<12} | 🌊 PHYSICS\")\n",
+        "\n",
+        "    down_p = get_p(model.particle_down) + get_p(model.wave_down)\n",
+        "    print(f\"{'Stream Splitters (Downsample)':<40} | {format_num(down_p):<12} | 📉 PROJ\")\n",
+        "\n",
+        "    # --- 3. STREAM A: SENSORY (TRANSFORMER) ---\n",
+        "    print(\"-\" * 80)\n",
+        "    print(f\"STREAM A: SENSORY (Identity/Magnitude)\")\n",
+        "    sensory_p = get_p(model.stream_sensory)\n",
+        "    try:\n",
+        "        depth_s = len(model.stream_sensory.encoder.layers)\n",
+        "        print(f\"  ├─ Transformer Encoder (Depth {depth_s})     | {format_num(sensory_p):<12} | ⚡ ATTENTION\")\n",
+        "    except:\n",
+        "        print(f\"  ├─ Transformer Encoder (Fused)           | {format_num(sensory_p):<12} | ⚡ ATTENTION\")\n",
+        "\n",
+        "    # --- 4. STREAM B: RELATIONAL (PRISM) ---\n",
+        "    print(\"-\" * 80)\n",
+        "    print(f\"STREAM B: RELATIONAL (Structure/Phase)\")\n",
+        "    relational_core = get_p(model.stream_relational)\n",
+        "    relational_bridge = get_p(model.relational_bridge)\n",
+        "\n",
+        "    try:\n",
+        "        depth_r = len(model.stream_relational.layers)\n",
+        "        print(f\"  ├─ PRISM Encoder (Depth {depth_r})           | {format_num(relational_core):<12} | 🌊 SPECTRAL\")\n",
+        "    except:\n",
+        "        print(f\"  ├─ PRISM Encoder (Fused)                 | {format_num(relational_core):<12} | 🌊 SPECTRAL\")\n",
+        "\n",
+        "    print(f\"  └─ Bridge (Complex->Real)                | {format_num(relational_bridge):<12} | 🌉 PROJ\")\n",
+        "\n",
+        "    # --- 5. FUSION & OUTPUT ---\n",
+        "    print(\"-\" * 80)\n",
+        "    fusion_p = get_p(model.fusion_proj) + get_p(model.fusion_norm)\n",
+        "    print(f\"{'Fusion (Concat -> Proj)':<40} | {format_num(fusion_p):<12} | 🧠 MIX\")\n",
+        "\n",
+        "    refiner_p = get_p(model.refiner)\n",
+        "    print(f\"{'Refiner (1-Layer Transformer)':<40} | {format_num(refiner_p):<12} | 🧠 REASONING\")\n",
+        "\n",
+        "    bias_p = get_p(model.head_bias)\n",
+        "    print(f\"{'Output Head Bias':<40} | {format_num(bias_p):<12} | 🎯 OUT\")\n",
+        "\n",
+        "    # --- SUMMARY ---\n",
+        "    print(\"=\"*80)\n",
+        "    storage = vocab_emb + bias_p\n",
+        "    active = total_params - storage\n",
+        "\n",
+        "    print(f\"TOTAL PARAMETERS:       {total_params/1e6:.2f} M\")\n",
+        "    print(f\"   ├─ 💾 Storage:       {storage/1e6:.2f} M  (Embeddings)\")\n",
+        "    print(f\"   └─ 🧠 Compute:       {active/1e6:.2f} M  (Active Weights)\")\n",
+        "    print(\"-\" * 80)\n",
+        "    print(f\"RATIO CHECK:\")\n",
+        "    print(f\"   ⚡ Sensory (Transf): {sensory_p/1e6:.2f} M\")\n",
+        "    print(f\"   🌊 Relation (PRISM): {(relational_core + relational_bridge)/1e6:.2f} M\")\n",
+        "    print(\"=\"*80 + \"\\n\")\n",
+        "\n",
+        "def generate_run_id():\n",
+        "    raw = datetime.now().strftime(\"%Y%m%d%H%M%S%f\")\n",
+        "    return hashlib.md5(raw.encode()).hexdigest()[:8]\n",
+        "\n",
+        "def log_environment(save_dir, run_id, config):\n",
+        "    log_path = os.path.join(save_dir, f\"env_metadata_{run_id}.txt\")\n",
+        "    with open(log_path, \"w\") as f:\n",
+        "        f.write(f\"PRISM EXPERIMENT METADATA | Run ID: {run_id}\\n{'='*50}\\n\")\n",
+        "        for k, v in config.items(): f.write(f\"{k}: {v}\\n\")\n",
+        "    print(f\"📝 Environment Snapshot saved to: {log_path}\")\n",
+        "\n",
+        "def log_metrics(save_dir, run_id, epoch, train_loss, val_loss, ppl):\n",
+        "    log_path = os.path.join(save_dir, f\"metrics_log_{run_id}.csv\")\n",
+        "    if not os.path.exists(log_path):\n",
+        "        with open(log_path, \"w\") as f: f.write(\"Timestamp,Epoch,Train_Loss,Val_Loss,Perplexity\\n\")\n",
+        "    with open(log_path, \"a\") as f:\n",
+        "        ts = datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n",
+        "        f.write(f\"{ts},{epoch},{train_loss:.6f},{val_loss:.6f},{ppl:.6f}\\n\")\n",
+        "\n",
+        "def save_checkpoint(path, model, optimizer, scheduler, scaler, epoch, best_loss, config):\n",
+        "    torch.save({\n",
+        "        'epoch': epoch,\n",
+        "        'model_state_dict': model.state_dict(),\n",
+        "        'optimizer_state_dict': optimizer.state_dict(),\n",
+        "        'scheduler_state_dict': scheduler.state_dict(),\n",
+        "        'scaler_state_dict': scaler.state_dict(), # <--- IMPORTANT FOR AMP\n",
+        "        'best_val_loss': best_loss,\n",
+        "        'config': config\n",
+        "    }, path)\n",
+        "\n",
+        "# ==========================================\n",
+        "# 5. A100 TRAINING LOOP (WITH LOGGING)\n",
+        "# ==========================================\n",
+        "# ==========================================\n",
+        "# 4. LOGGING & ANALYSIS UTILITIES\n",
+        "# ==========================================\n",
+        "def deep_analyze_pillars_dat(model):\n",
+        "    def get_p(obj):\n",
+        "        if isinstance(obj, nn.Parameter): return obj.numel()\n",
+        "        return sum(p.numel() for p in obj.parameters() if p.requires_grad)\n",
+        "\n",
+        "    def format_num(n):\n",
+        "        if n > 1e6: return f\"{n/1e6:.2f}M\"\n",
+        "        if n > 1e3: return f\"{n/1e3:.2f}K\"\n",
+        "        return str(n)\n",
+        "\n",
+        "    print(\"\\n\" + \"=\"*80)\n",
+        "    print(f\"🏛️  PILLARS-DAT (Hybrid Transformer-PRISM) - ANALYSIS\")\n",
+        "    print(\"=\"*80)\n",
+        "    print(f\"{'MODULE / LAYER':<40} | {'PARAMS':<12} | {'TYPE'}\")\n",
+        "    print(\"-\" * 80)\n",
+        "\n",
+        "    total_params = get_p(model)\n",
+        "\n",
+        "    # --- 1. MEMORY ---\n",
+        "    vocab_emb = get_p(model.rose.raw_embedding)\n",
+        "    print(f\"{'Shared Vocab Embedding':<40} | {format_num(vocab_emb):<12} | 💾 STORAGE\")\n",
+        "\n",
+        "    # --- 2. INPUT PHYSICS ---\n",
+        "    rose_logic = get_p(model.rose) - vocab_emb\n",
+        "    print(f\"{'Dynamic RoSE (Adapters)':<40} | {format_num(rose_logic):<12} | 🌊 PHYSICS\")\n",
+        "\n",
+        "    down_p = get_p(model.particle_down) + get_p(model.wave_down)\n",
+        "    print(f\"{'Stream Splitters (Downsample)':<40} | {format_num(down_p):<12} | 📉 PROJ\")\n",
+        "\n",
+        "    # --- 3. STREAM A: SENSORY (TRANSFORMER) ---\n",
+        "    print(\"-\" * 80)\n",
+        "    print(f\"STREAM A: SENSORY (Identity/Magnitude)\")\n",
+        "    sensory_p = get_p(model.stream_sensory)\n",
+        "    try:\n",
+        "        depth_s = len(model.stream_sensory.encoder.layers)\n",
+        "        print(f\"  ├─ Transformer Encoder (Depth {depth_s})     | {format_num(sensory_p):<12} | ⚡ ATTENTION\")\n",
+        "    except:\n",
+        "        print(f\"  ├─ Transformer Encoder (Fused)           | {format_num(sensory_p):<12} | ⚡ ATTENTION\")\n",
+        "\n",
+        "    # --- 4. STREAM B: RELATIONAL (PRISM) ---\n",
+        "    print(\"-\" * 80)\n",
+        "    print(f\"STREAM B: RELATIONAL (Structure/Phase)\")\n",
+        "    relational_core = get_p(model.stream_relational)\n",
+        "    relational_bridge = get_p(model.relational_bridge)\n",
+        "\n",
+        "    try:\n",
+        "        depth_r = len(model.stream_relational.layers)\n",
+        "        print(f\"  ├─ PRISM Encoder (Depth {depth_r})           | {format_num(relational_core):<12} | 🌊 SPECTRAL\")\n",
+        "    except:\n",
+        "        print(f\"  ├─ PRISM Encoder (Fused)                 | {format_num(relational_core):<12} | 🌊 SPECTRAL\")\n",
+        "\n",
+        "    print(f\"  └─ Bridge (Complex->Real)                | {format_num(relational_bridge):<12} | 🌉 PROJ\")\n",
+        "\n",
+        "    # --- 5. FUSION & OUTPUT ---\n",
+        "    print(\"-\" * 80)\n",
+        "    fusion_p = get_p(model.fusion_proj) + get_p(model.fusion_norm)\n",
+        "    print(f\"{'Fusion (Concat -> Proj)':<40} | {format_num(fusion_p):<12} | 🧠 MIX\")\n",
+        "\n",
+        "    refiner_p = get_p(model.refiner)\n",
+        "    print(f\"{'Refiner (1-Layer Transformer)':<40} | {format_num(refiner_p):<12} | 🧠 REASONING\")\n",
+        "\n",
+        "    bias_p = get_p(model.head_bias)\n",
+        "    print(f\"{'Output Head Bias':<40} | {format_num(bias_p):<12} | 🎯 OUT\")\n",
+        "\n",
+        "    # --- SUMMARY ---\n",
+        "    print(\"=\"*80)\n",
+        "    storage = vocab_emb + bias_p\n",
+        "    active = total_params - storage\n",
+        "\n",
+        "    print(f\"TOTAL PARAMETERS:       {total_params/1e6:.2f} M\")\n",
+        "    print(f\"   ├─ 💾 Storage:       {storage/1e6:.2f} M  (Embeddings)\")\n",
+        "    print(f\"   └─ 🧠 Compute:       {active/1e6:.2f} M  (Active Weights)\")\n",
+        "    print(\"-\" * 80)\n",
+        "    print(f\"RATIO CHECK:\")\n",
+        "    print(f\"   ⚡ Sensory (Transf): {sensory_p/1e6:.2f} M\")\n",
+        "    print(f\"   🌊 Relation (PRISM): {(relational_core + relational_bridge)/1e6:.2f} M\")\n",
+        "    print(\"=\"*80 + \"\\n\")\n",
+        "\n",
+        "def init_pillars_dat_weights(model):\n",
+        "    print(\"✨ APPLYING PILLARS-DAT INITIALIZATION PROTOCOL...\")\n",
+        "    # 1. SHARED ROOT (RoSE)\n",
+        "    nn.init.normal_(model.rose.raw_embedding.weight, std=model.d_model ** -0.5)\n",
+        "    nn.init.orthogonal_(model.rose.adapter.weight)\n",
+        "\n",
+        "    # --- ROSE IDENTITY TRICK ---\n",
+        "    nn.init.normal_(model.rose.rotation_predictor.weight, std=0.01)\n",
+        "    with torch.no_grad():\n",
+        "        model.rose.rotation_predictor.bias[:model.d_model].fill_(1.0) # Real=1\n",
+        "        model.rose.rotation_predictor.bias[model.d_model:].fill_(0.0) # Imag=0\n",
+        "\n",
+        "    # 2. DOWNSAMPLERS\n",
+        "    nn.init.orthogonal_(model.particle_down.weight, gain=1.414)\n",
+        "    nn.init.orthogonal_(model.wave_down.weight, gain=1.414)\n",
+        "\n",
+        "    # 3. SENSORY STREAM (Transformer + RoPE)\n",
+        "    print(\"   ├─ Initializing Sensory Stream (Transformer)...\")\n",
+        "    for name, p in model.stream_sensory.named_parameters():\n",
+        "        if p.dim() > 1:\n",
+        "            nn.init.xavier_uniform_(p)\n",
+        "        elif \"norm\" in name.lower() and p.dim() == 1:\n",
+        "            if \"weight\" in name: nn.init.ones_(p)\n",
+        "            if \"bias\" in name:   nn.init.zeros_(p)\n",
+        "\n",
+        "    # 4. RELATIONAL STREAM (PRISM)\n",
+        "    print(\"   ├─ Initializing Relational Stream (PRISM)...\")\n",
+        "    for name, m in model.stream_relational.named_modules():\n",
+        "        if isinstance(m, nn.Linear):\n",
+        "            nn.init.xavier_uniform_(m.weight, gain=1.0)\n",
+        "            if m.bias is not None: nn.init.zeros_(m.bias)\n",
+        "        if isinstance(m, ModReLU):\n",
+        "            nn.init.constant_(m.b, 0.01)\n",
+        "\n",
+        "    # 5. FUSION & REFINER\n",
+        "    nn.init.xavier_uniform_(model.fusion_proj.weight, gain=1.0)\n",
+        "    for p in model.refiner.parameters():\n",
+        "        if p.dim() > 1: nn.init.xavier_uniform_(p)\n",
+        "\n",
+        "    # 6. TIED HEAD BIAS\n",
+        "    nn.init.zeros_(model.head_bias)\n",
+        "    print(\"✅ DAT INITIALIZATION COMPLETE.\")\n",
+        "\n",
+        "# ==========================================\n",
+        "# 5. A100 TRAINING LOOP (WITH LOGGING)\n",
+        "# ==========================================\n",
+        "def run_a100_training(experiment_name=\"PILLARS_DAT_A100_Final\"):\n",
+        "    from torch.cuda.amp import autocast, GradScaler\n",
+        "    from torch.utils.tensorboard import SummaryWriter\n",
+        "\n",
+        "    # --- 1. SETUP DRIVE & LOGGING ---\n",
+        "    from google.colab import drive\n",
+        "    if not os.path.exists('/content/drive'): drive.mount('/content/drive')\n",
+        "\n",
+        "    run_id = generate_run_id()\n",
+        "    timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
+        "    SAVE_DIR = os.path.join(\"/content/drive/My Drive/PRISM_Experiments\", f\"{experiment_name}_{timestamp}_{run_id}\")\n",
+        "    os.makedirs(SAVE_DIR, exist_ok=True)\n",
+        "\n",
+        "    writer = SummaryWriter(log_dir=SAVE_DIR)\n",
+        "\n",
+        "    # Config for Logs\n",
+        "    config_dump = {\n",
+        "        \"run_id\": run_id, \"batch_size\": 6, \"accum\": 8, \"d_model\": D_MODEL, \"depth\": DEPTH, \"seq_len\": SEQ_LEN\n",
+        "    }\n",
+        "    log_environment(SAVE_DIR, run_id, config_dump)\n",
+        "\n",
+        "    # --- 2. MODEL & DATA ---\n",
+        "    SAFE_BATCH_SIZE = BATCH_SIZE\n",
+        "    GRAD_ACCUM = 4\n",
+        "    print(f\"\\n⚡ A100 DETECTED. CONFIGURING FLASH ATTENTION PIPELINE...\")\n",
+        "\n",
+        "    lm_datasets, data_collator = prepare_data_from_hub()\n",
+        "    train_loader = DataLoader(lm_datasets[\"train\"], batch_size=SAFE_BATCH_SIZE, shuffle=True, collate_fn=data_collator, num_workers=4, pin_memory=True)\n",
+        "    valid_loader = DataLoader(lm_datasets[\"validation\"], batch_size=SAFE_BATCH_SIZE, collate_fn=data_collator, num_workers=2)\n",
+        "\n",
+        "    model = Pillars_DAT(vocab_size=VOCAB_SIZE, d_model=D_MODEL, d_branch=D_BRANCH, seq_len=SEQ_LEN, depth=DEPTH).to(DEVICE)\n",
+        "    init_pillars_dat_weights(model)\n",
+        "    print(model)\n",
+        "    deep_analyze_pillars_dat(model) # <--- Parameter Analysis\n",
+        "\n",
+        "    optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=0.01)\n",
+        "    total_steps = (len(train_loader) // GRAD_ACCUM) * EPOCHS\n",
+        "    warmup_steps = int(total_steps * 0.1)\n",
+        "    scheduler = get_cosine_schedule_with_warmup(optimizer, warmup_steps, total_steps)\n",
+        "    criterion = nn.CrossEntropyLoss()\n",
+        "    scaler = GradScaler() # For AMP\n",
+        "\n",
+        "    print(f\"\\n🚀 IGNITING FUSION DRIVE... Saving to: {SAVE_DIR}\")\n",
+        "\n",
+        "    global_step = 0\n",
+        "    best_val_loss = float('inf')\n",
+        "\n",
+        "    for epoch in range(EPOCHS):\n",
+        "        model.train()\n",
+        "        pbar = tqdm(train_loader, desc=f\"Ep {epoch+1}\")\n",
+        "\n",
+        "        for step, batch in enumerate(pbar):\n",
+        "            x, y = batch['input_ids'].to(DEVICE), batch['labels'].to(DEVICE)\n",
+        "\n",
+        "            # ⚡ AMP CONTEXT\n",
+        "            with autocast(dtype=torch.float16):\n",
+        "                logits = model(x).view(-1, VOCAB_SIZE)\n",
+        "                loss = criterion(logits, y.view(-1)) / GRAD_ACCUM\n",
+        "\n",
+        "            scaler.scale(loss).backward()\n",
+        "\n",
+        "            if (step + 1) % GRAD_ACCUM == 0:\n",
+        "                scaler.unscale_(optimizer)\n",
+        "                # 🛑 CALC GRAD NORM HERE FOR PBAR 🛑\n",
+        "                grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
+        "\n",
+        "                scaler.step(optimizer)\n",
+        "                scaler.update()\n",
+        "                scheduler.step()\n",
+        "                optimizer.zero_grad()\n",
+        "                global_step += 1\n",
+        "\n",
+        "                # 📝 STEP LOGGING\n",
+        "                actual_loss = loss.item() * GRAD_ACCUM\n",
+        "                writer.add_scalar('Train/Loss', actual_loss, global_step)\n",
+        "                writer.add_scalar('Train/GradNorm', grad_norm.item(), global_step)\n",
+        "                writer.add_scalar('Train/LR', scheduler.get_last_lr()[0], global_step)\n",
+        "\n",
+        "                # ✨ UPDATE PBAR WITH GNORM ✨\n",
+        "                pbar.set_postfix({\n",
+        "                    'loss': f\"{actual_loss:.4f}\",\n",
+        "                    'gnorm': f\"{grad_norm.item():.2f}\"\n",
+        "                })\n",
+        "\n",
+        "        # --- VALIDATION ---\n",
+        "        model.eval()\n",
+        "        val_loss = 0\n",
+        "        with torch.no_grad(), autocast():\n",
+        "            for batch in valid_loader:\n",
+        "                x, y = batch['input_ids'].to(DEVICE), batch['labels'].to(DEVICE)\n",
+        "                val_loss += criterion(model(x).view(-1, VOCAB_SIZE), y.view(-1)).item()\n",
+        "\n",
+        "        avg_val_loss = val_loss / len(valid_loader)\n",
+        "        # Prevent overflow if loss is exploding\n",
+        "        ppl = math.exp(avg_val_loss) if avg_val_loss < 20 else float('inf')\n",
+        "\n",
+        "        print(f\"✨ Ep {epoch+1} | Val Loss: {avg_val_loss:.4f} | PPL: {ppl:.2f}\")\n",
+        "\n",
+        "        # 📝 EPOCH LOGGING\n",
+        "        writer.add_scalar('Val/Loss', avg_val_loss, epoch+1)\n",
+        "        writer.add_scalar('Val/PPL', ppl, epoch+1)\n",
+        "        log_metrics(SAVE_DIR, run_id, epoch+1, 0.0, avg_val_loss, ppl)\n",
+        "\n",
+        "        # 💾 SAVE CHECKPOINTS (Includes Scaler/Optim/Sched)\n",
+        "        save_checkpoint(os.path.join(SAVE_DIR, \"last.pt\"), model, optimizer, scheduler, scaler, epoch, best_val_loss, config_dump)\n",
+        "\n",
+        "        if avg_val_loss < best_val_loss:\n",
+        "            best_val_loss = avg_val_loss\n",
+        "            print(f\"   🏆 New Best Model! Saving best.pt...\")\n",
+        "            save_checkpoint(os.path.join(SAVE_DIR, \"best.pt\"), model, optimizer, scheduler, scaler, epoch, best_val_loss, config_dump)\n",
+        "\n",
+        "    writer.close()\n",
+        "    return model\n",
+        "\n",
+        "if __name__ == \"__main__\":\n",
+        "    run_a100_training()"
+      ],
+      "metadata": {
+        "id": "-TNEv89gkS1k"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "from google.colab import runtime\n",
+        "runtime.unassign()"
+      ],
+      "metadata": {
+        "id": "bxFTYWHVqcSI"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}

WT103_Transformer_Baseline.ipynb

@@ -0,0 +1,357 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "A100"
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "code",
+      "source": [
+        "!pip install -q x-transformers\n",
+        "!pip install -q flash-attn --no-build-isolation\n",
+        "\n",
+        "import torch\n",
+        "import torch.nn as nn\n",
+        "import torch.optim as optim\n",
+        "import math\n",
+        "import os\n",
+        "import hashlib\n",
+        "from datetime import datetime\n",
+        "from tqdm.auto import tqdm\n",
+        "from torch.utils.data import DataLoader\n",
+        "from torch.utils.tensorboard import SummaryWriter\n",
+        "from transformers import RobertaTokenizerFast, get_cosine_schedule_with_warmup, DataCollatorForLanguageModeling\n",
+        "from datasets import load_dataset\n",
+        "from x_transformers import TransformerWrapper, Encoder\n",
+        "\n",
+        "# ==========================================\n",
+        "# 1. CONFIGURATION\n",
+        "# ==========================================\n",
+        "HF_ID = \"prism-lab/wikitext-103-prism-32k-seq4k\"\n",
+        "EXPERIMENT_NAME = \"BASELINE_Pure_XTransformers\"\n",
+        "\n",
+        "VOCAB_SIZE = 32768\n",
+        "SEQ_LEN = 4096\n",
+        "BATCH_SIZE = 8\n",
+        "GRAD_ACCUM = 4\n",
+        "EPOCHS = 40\n",
+        "LR = 1e-3\n",
+        "D_MODEL = 512\n",
+        "DEPTH = 5\n",
+        "HEADS = 8\n",
+        "DROPOUT = 0.1\n",
+        "WEIGHT_DECAY = 0.01\n",
+        "\n",
+        "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+        "torch.set_float32_matmul_precision(\"high\")\n",
+        "def print_detailed_param_count(model):\n",
+        "    \"\"\"\n",
+        "    Detailed parameter breakdown for x_transformers models.\n",
+        "    \"\"\"\n",
+        "    total_params = 0\n",
+        "    categories = {\n",
+        "        \"Embeddings\": 0,\n",
+        "        \"Attention\": 0,\n",
+        "        \"FeedForward\": 0,\n",
+        "        \"Norms\": 0,\n",
+        "        \"Head/Other\": 0\n",
+        "    }\n",
+        "\n",
+        "    # Track distinct tensors to handle tied weights correctly\n",
+        "    seen_pointers = set()\n",
+        "\n",
+        "    for name, param in model.named_parameters():\n",
+        "        if param.requires_grad:\n",
+        "            # Check for shared weights (tied embeddings)\n",
+        "            if param.data_ptr() in seen_pointers:\n",
+        "                continue\n",
+        "            seen_pointers.add(param.data_ptr())\n",
+        "\n",
+        "            num_params = param.numel()\n",
+        "            total_params += num_params\n",
+        "\n",
+        "            # Logic for x_transformers naming conventions\n",
+        "            if \"token_emb\" in name or \"pos_emb\" in name:\n",
+        "                categories[\"Embeddings\"] += num_params\n",
+        "            elif \"attn\" in name or \"to_q\" in name or \"to_k\" in name or \"to_v\" in name:\n",
+        "                categories[\"Attention\"] += num_params\n",
+        "            elif \"ff\" in name:\n",
+        "                categories[\"FeedForward\"] += num_params\n",
+        "            elif \"norm\" in name:\n",
+        "                categories[\"Norms\"] += num_params\n",
+        "            else:\n",
+        "                categories[\"Head/Other\"] += num_params\n",
+        "\n",
+        "    print(f\"\\n{'='*60}\")\n",
+        "    print(f\"{'COMPONENT':<25} | {'PARAMS':<12} | {'%':<6}\")\n",
+        "    print(f\"{'-'*60}\")\n",
+        "\n",
+        "    for cat, count in categories.items():\n",
+        "        if count > 0:\n",
+        "            percentage = (count / total_params) * 100\n",
+        "            print(f\"{cat:<25} | {count:12,} | {percentage:5.1f}%\")\n",
+        "\n",
+        "    print(f\"{'='*60}\")\n",
+        "    print(f\"{'TOTAL':<25} | {total_params:12,} | 100.0%\")\n",
+        "    print(f\"{'='*60}\\n\")\n",
+        "\n",
+        "    return total_params\n",
+        "\n",
+        "# ==========================================\n",
+        "# 2. TRAINING ROUTINE\n",
+        "# ==========================================\n",
+        "def run_baseline_training():\n",
+        "    from google.colab import drive\n",
+        "    if not os.path.exists('/content/drive'): drive.mount('/content/drive')\n",
+        "\n",
+        "    # Setup Logging\n",
+        "    timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
+        "    run_id = hashlib.md5(timestamp.encode()).hexdigest()[:8]\n",
+        "    SAVE_DIR = os.path.join(\"/content/drive/My Drive/PRISM_Experiments\", f\"{EXPERIMENT_NAME}_{timestamp}_{run_id}\")\n",
+        "    os.makedirs(SAVE_DIR, exist_ok=True)\n",
+        "    writer = SummaryWriter(log_dir=SAVE_DIR)\n",
+        "\n",
+        "    # Data\n",
+        "    print(\"⬇️ Loading Data...\")\n",
+        "    tokenizer = RobertaTokenizerFast.from_pretrained(HF_ID)\n",
+        "    dataset = load_dataset(HF_ID)\n",
+        "    pad_id = tokenizer.pad_token_id # We need this for the mask\n",
+        "\n",
+        "    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=True, mlm_probability=0.15)\n",
+        "\n",
+        "    train_loader = DataLoader(dataset[\"train\"], batch_size=BATCH_SIZE, shuffle=True, collate_fn=data_collator, num_workers=2, pin_memory=True, persistent_workers=True)\n",
+        "    valid_loader = DataLoader(dataset[\"validation\"], batch_size=BATCH_SIZE, collate_fn=data_collator, num_workers=2, pin_memory=True)\n",
+        "    test_loader = DataLoader(dataset[\"test\"], batch_size=BATCH_SIZE, collate_fn=data_collator, num_workers=2, pin_memory=True)\n",
+        "\n",
+        "    # ======================================================\n",
+        "    # 3. MODEL: DIRECTLY USING THE LIBRARY\n",
+        "    # ======================================================\n",
+        "    print(\"\\n⚡ INITIALIZING x_transformers LIBRARY MODEL...\")\n",
+        "\n",
+        "    model = TransformerWrapper(\n",
+        "        num_tokens=VOCAB_SIZE,\n",
+        "        max_seq_len=SEQ_LEN,\n",
+        "        use_abs_pos_emb=False,    # FALSE because we use RoPE\n",
+        "        tie_embedding=True,       # Match PRISM (if PRISM tied embeddings)\n",
+        "        attn_layers=Encoder(\n",
+        "            dim=D_MODEL,\n",
+        "            depth=DEPTH,\n",
+        "            heads=HEADS,\n",
+        "            layer_dropout=DROPOUT,\n",
+        "            attn_dropout=DROPOUT,\n",
+        "            ff_dropout=DROPOUT,\n",
+        "            rotary_pos_emb=True,  # Match PRISM Refiner\n",
+        "            attn_flash=True,      # Match PRISM Refiner\n",
+        "            use_scalenorm=False   # <--- CHANGE THIS TO FALSE (Or remove it)\n",
+        "        )\n",
+        "    ).to(DEVICE)\n",
+        "    print(model)\n",
+        "    print_detailed_param_count(model)\n",
+        "    # Simple Parameter Count\n",
+        "    print(f\"📊 Parameters: {sum(p.numel() for p in model.parameters()):,}\")\n",
+        "\n",
+        "    # ======================================================\n",
+        "\n",
+        "    optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n",
+        "    total_steps = (len(train_loader) // GRAD_ACCUM) * EPOCHS\n",
+        "    scheduler = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=int(0.1*total_steps), num_training_steps=total_steps)\n",
+        "    criterion = nn.CrossEntropyLoss()\n",
+        "    scaler = torch.cuda.amp.GradScaler()\n",
+        "\n",
+        "    print(f\"\\n🚀 STARTING TRAINING (Ep 1 to {EPOCHS})\")\n",
+        "\n",
+        "    best_val_loss = float('inf')\n",
+        "    global_step = 0\n",
+        "\n",
+        "    for epoch in range(EPOCHS):\n",
+        "        model.train()\n",
+        "        pbar = tqdm(train_loader, desc=f\"Ep {epoch+1}/{EPOCHS}\", dynamic_ncols=True)\n",
+        "\n",
+        "        for step, batch in enumerate(pbar):\n",
+        "            x, y = batch['input_ids'].to(DEVICE), batch['labels'].to(DEVICE)\n",
+        "\n",
+        "            # --- CRITICAL CHANGE: PASS MASK MANUALLY ---\n",
+        "            # x_transformers requires 'mask' to know what is padding\n",
+        "            mask = (x != pad_id)\n",
+        "\n",
+        "            with torch.cuda.amp.autocast():\n",
+        "                # We pass the mask directly here\n",
+        "                outputs = model(x, mask=mask)\n",
+        "                loss = criterion(outputs.view(-1, VOCAB_SIZE), y.view(-1)) / GRAD_ACCUM\n",
+        "\n",
+        "            scaler.scale(loss).backward()\n",
+        "\n",
+        "            if (step + 1) % GRAD_ACCUM == 0:\n",
+        "                scaler.unscale_(optimizer)\n",
+        "                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
+        "                scaler.step(optimizer)\n",
+        "                scaler.update()\n",
+        "                scheduler.step()\n",
+        "                optimizer.zero_grad()\n",
+        "                global_step += 1\n",
+        "\n",
+        "                loss_val = loss.item() * GRAD_ACCUM\n",
+        "                pbar.set_postfix({'loss': f\"{loss_val:.4f}\", 'lr': f\"{scheduler.get_last_lr()[0]:.2e}\"})\n",
+        "                writer.add_scalar('Train/Loss', loss_val, global_step)\n",
+        "\n",
+        "        # Validation\n",
+        "        model.eval()\n",
+        "        val_loss = 0\n",
+        "        total_batches = 0\n",
+        "        with torch.no_grad():\n",
+        "            for batch in valid_loader:\n",
+        "                x, y = batch['input_ids'].to(DEVICE), batch['labels'].to(DEVICE)\n",
+        "                mask = (x != pad_id)\n",
+        "\n",
+        "                with torch.cuda.amp.autocast():\n",
+        "                    outputs = model(x, mask=mask)\n",
+        "                    val_loss += criterion(outputs.view(-1, VOCAB_SIZE), y.view(-1)).item()\n",
+        "                total_batches += 1\n",
+        "\n",
+        "        avg_val_loss = val_loss / total_batches\n",
+        "        ppl = math.exp(avg_val_loss) if avg_val_loss < 100 else float('inf')\n",
+        "        print(f\"✨ Epoch {epoch+1} | Val Loss: {avg_val_loss:.4f} | PPL: {ppl:.2f}\")\n",
+        "        writer.add_scalar('Val/PPL', ppl, epoch+1)\n",
+        "\n",
+        "        # Save Best\n",
+        "        if avg_val_loss < best_val_loss:\n",
+        "            best_val_loss = avg_val_loss\n",
+        "            torch.save(model.state_dict(), os.path.join(SAVE_DIR, \"best.pt\"))\n",
+        "            print(\"   🏆 New Best Model Saved!\")\n",
+        "\n",
+        "        # Save Last State (Optimizer, Scheduler, Scaler)\n",
+        "        state = {\n",
+        "            'epoch': epoch,\n",
+        "            'model': model.state_dict(),\n",
+        "            'optimizer': optimizer.state_dict(),\n",
+        "            'scheduler': scheduler.state_dict(),\n",
+        "            'scaler': scaler.state_dict(),\n",
+        "            'best_loss': best_val_loss\n",
+        "        }\n",
+        "        torch.save(state, os.path.join(SAVE_DIR, \"last.pt\"))\n",
+        "\n",
+        "    # Test\n",
+        "    print(f\"\\n🧪 Testing Best Model...\")\n",
+        "    if os.path.exists(os.path.join(SAVE_DIR, \"best.pt\")):\n",
+        "        model.load_state_dict(torch.load(os.path.join(SAVE_DIR, \"best.pt\")))\n",
+        "\n",
+        "    model.eval()\n",
+        "    test_loss = 0\n",
+        "    total_batches = 0\n",
+        "    with torch.no_grad():\n",
+        "        for batch in tqdm(test_loader, desc=\"Testing\"):\n",
+        "            x, y = batch['input_ids'].to(DEVICE), batch['labels'].to(DEVICE)\n",
+        "            mask = (x != pad_id)\n",
+        "            with torch.cuda.amp.autocast():\n",
+        "                outputs = model(x, mask=mask)\n",
+        "                test_loss += criterion(outputs.view(-1, VOCAB_SIZE), y.view(-1)).item()\n",
+        "            total_batches += 1\n",
+        "\n",
+        "    final_ppl = math.exp(test_loss / total_batches)\n",
+        "    print(f\"🏆 FINAL BASELINE PPL: {final_ppl:.2f}\")\n",
+        "\n",
+        "    writer.close()\n",
+        "    return model\n",
+        "\n",
+        "if __name__ == \"__main__\":\n",
+        "    run_baseline_training()"
+      ],
+      "metadata": {
+        "id": "FnQHarBgVBnU"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "!pip freeze"
+      ],
+      "metadata": {
+        "id": "3YRhdCC-O2tW"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import sys\n",
+        "import torch\n",
+        "import platform\n",
+        "import os\n",
+        "\n",
+        "# Define output file\n",
+        "log_file = \"environment_snapshot.txt\"\n",
+        "\n",
+        "print(f\"📝 Generating {log_file}...\")\n",
+        "\n",
+        "with open(log_file, \"w\") as f:\n",
+        "    # 1. PYTHON & SYSTEM INFO\n",
+        "    f.write(\"=\"*40 + \"\\n\")\n",
+        "    f.write(\"SYSTEM INFORMATION\\n\")\n",
+        "    f.write(\"=\"*40 + \"\\n\")\n",
+        "    f.write(f\"Python Version: {sys.version}\\n\")\n",
+        "    f.write(f\"Platform:       {platform.platform()}\\n\")\n",
+        "    f.write(f\"Architecture:   {platform.machine()}\\n\")\n",
+        "    f.write(f\"Processor:      {platform.processor()}\\n\")\n",
+        "\n",
+        "    # 2. GPU & CUDA INFO\n",
+        "    f.write(\"\\n\" + \"=\"*40 + \"\\n\")\n",
+        "    f.write(\"GPU / CUDA INFORMATION\\n\")\n",
+        "    f.write(\"=\"*40 + \"\\n\")\n",
+        "    f.write(f\"PyTorch Version: {torch.__version__}\\n\")\n",
+        "    f.write(f\"CUDA Available:  {torch.cuda.is_available()}\\n\")\n",
+        "\n",
+        "    if torch.cuda.is_available():\n",
+        "        f.write(f\"CUDA Version:    {torch.version.cuda}\\n\")\n",
+        "        f.write(f\"CUDNN Version:   {torch.backends.cudnn.version()}\\n\")\n",
+        "        f.write(f\"Device Name:     {torch.cuda.get_device_name(0)}\\n\")\n",
+        "        f.write(f\"Device Count:    {torch.cuda.device_count()}\\n\")\n",
+        "    else:\n",
+        "        f.write(\"NO GPU DETECTED\\n\")\n",
+        "\n",
+        "    # 3. COLAB SPECIFICS (If applicable)\n",
+        "    f.write(\"\\n\" + \"=\"*40 + \"\\n\")\n",
+        "    f.write(\"ENV VARIABLES (FILTERED)\\n\")\n",
+        "    f.write(\"=\"*40 + \"\\n\")\n",
+        "    # Capturing useful Colab/Jupyter env vars without exposing secrets\n",
+        "    keys_to_log = ['COLAB_GPU', 'CUDA_VERSION', 'TBE_RUNTIME_ADDR']\n",
+        "    for k, v in os.environ.items():\n",
+        "        if any(x in k for x in ['COLAB', 'CUDA', 'LD_LIBRARY_PATH']):\n",
+        "            f.write(f\"{k}: {v}\\n\")\n",
+        "\n",
+        "    # 4. INSTALLED PACKAGES (Pip Freeze)\n",
+        "    f.write(\"\\n\" + \"=\"*40 + \"\\n\")\n",
+        "    f.write(\"PIP FREEZE (FULL LIBRARY LIST)\\n\")\n",
+        "    f.write(\"=\"*40 + \"\\n\")\n",
+        "\n",
+        "# Append pip freeze output directly to the file\n",
+        "os.system(f\"pip freeze >> {log_file}\")\n",
+        "\n",
+        "print(f\"✅ Log saved to {log_file}\")\n",
+        "print(\"   (Check the Files tab on the left to download it)\")\n",
+        "\n",
+        "# Optional: Print head of file to verify\n",
+        "!head -n 20 environment_snapshot.txt"
+      ],
+      "metadata": {
+        "id": "63cDp0cJL2c2"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}