Upload 12 files

.gitattributes

@@ -0,0 +1,2 @@
+AI-Transfer/data/spiderportal_combined.pkl filter=lfs diff=lfs merge=lfs -text
+SpiderPortal-Base-v5/model.safetensors filter=lfs diff=lfs merge=lfs -text

AI-Transfer/README.md

@@ -0,0 +1,81 @@
+# SpiderPortal v5 — GPU Training Transfer Package
+
+## Contents
+- `scripts/train_single_gpu.py` — Optimized single-GPU training script
+- `scripts/train_ddp.py` — Original DDP script (for reference)
+- `data/spiderportal_combined.pkl` — Training dataset (491K samples, 537MB)
+- `notebooks/spiderportal_gpu.ipynb` — Original Kaggle notebook (for reference)
+
+## Quick Start
+
+### 1. Setup
+```bash
+pip install torch transformers pandas
+```
+
+### 2. Run Training
+```bash
+python scripts/train_single_gpu.py
+```
+
+## Configuration (edit in `train_single_gpu.py`)
+
+| Parameter | Default | RTX 6000 (96GB) |
+|-----------|---------|-----------------|
+| `BATCH_SIZE` | 64 | 64-128 |
+| `MAX_LEN` | 256 | 256-512 |
+| `EPOCHS` | 3 | 3 |
+| `N_LOOPS` | 2 | 2 |
+| `BASE_LR` | 2e-5 | 2e-5 |
+| `WARMUP_STEPS` | 1000 | 1000 |
+
+## Expected Performance (RTX PRO 6000, 96GB)
+
+| Metric | Value |
+|--------|-------|
+| Model size | 659M params |
+| Active params (MoE) | ~59M |
+| VRAM usage | ~15-25GB |
+| Batch size | 64 |
+| Steps per epoch | ~7,672 |
+| Time per epoch | 40-100 min |
+| **Total (3 epochs)** | **2-5 hours** |
+
+## Checkpoints
+
+Saved to `checkpoints/` directory:
+
+| Checkpoint type | What's saved | Size | Purpose |
+|----------------|--------------|------|---------|
+| Every 500 steps | Model weights only | ~1.3GB | Testing, transfer, inference |
+| End of epoch | Model + optimizer | ~6.6GB | Resume training |
+| Best loss | Model + optimizer | ~6.6GB | Resume training |
+
+Step checkpoints are auto-deleted at the start of each new epoch to free disk space.
+
+### Loading a weights-only checkpoint (inference)
+```python
+state_dict = torch.load("checkpoints/spiderportal-v5-ep1-step500.pt", map_location="cpu")
+model.load_state_dict(state_dict)
+```
+
+### Resuming training from an epoch checkpoint
+```python
+ckpt = torch.load("checkpoints/spiderportal-v5-ep1.pt", map_location="cpu")
+model.load_state_dict(ckpt["model_state_dict"])
+optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+start_epoch = ckpt["epoch"]
+```
+
+### Peak disk usage during training: ~20GB
+
+## Model Architecture
+
+- **Type**: MoE-RDT (Mixture of Experts + Recurrent Depth Transformer)
+- **Total params**: 659M
+- **Active params**: ~59M (1 routed expert + 1 shared expert per token)
+- **Experts**: 64 routed + 1 shared
+- **Layers**: 2 prelude → 8 recurrent (MoE) → 2 coda
+- **Attention**: GQA (8 heads, 2 KV heads)
+- **Context**: 131K tokens (YaRN RoPE scaling)
+- **Loop**: ACT halting + LTI injection + LoRA adapters

AI-Transfer/data/spiderportal_combined.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:706bb0870f785241df36c3dcc1925deb8a47ce6003dfe506b6c0751624c34e63
+size 563067610

AI-Transfer/notebooks/spiderportal_gpu.ipynb

@@ -0,0 +1,466 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# SpiderPortal MoE-RDT v5 \u2014 Multi-GPU Training (DDP)\n",
+    "\n",
+    "Optimized for 2\u00d7 T4 GPUs (32GB total VRAM).\n",
+    "- **Total params**: ~659M (64 experts)\n",
+    "- **Active params**: ~59M per token\n",
+    "- **Training**: bf16, DDP, gradient accumulation\n",
+    "- **Expected time**: ~1-1.5 hours for 3 epochs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install -q transformers pandas safetensors\n",
+    "\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "import math, os, json, gc, random, time\n",
+    "from pathlib import Path\n",
+    "from dataclasses import dataclass\n",
+    "from typing import Optional, Tuple, Dict, List\n",
+    "from torch.nn import CrossEntropyLoss\n",
+    "\n",
+    "print(f\"PyTorch: {torch.__version__}\")\n",
+    "print(f\"CUDA available: {torch.cuda.is_available()}\")\n",
+    "print(f\"GPU count: {torch.cuda.device_count()}\")\n",
+    "for i in range(torch.cuda.device_count()):\n",
+    "    p = torch.cuda.get_device_properties(i)\n",
+    "    print(f\"  GPU {i}: {torch.cuda.get_device_name(i)} ({p.total_memory / 1e9:.1f}GB)\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "@dataclass\n",
+    "class SpiderPortalConfig:\n",
+    "    vocab_size: int = 50278\n",
+    "    hidden_size: int = 384\n",
+    "    num_hidden_layers: int = 8\n",
+    "    num_attention_heads: int = 8\n",
+    "    num_key_value_heads: int = 2\n",
+    "    intermediate_size: int = 1024\n",
+    "    hidden_act: str = \"silu\"\n",
+    "    num_experts: int = 64\n",
+    "    num_experts_per_tok: int = 1\n",
+    "    num_shared_experts: int = 1\n",
+    "    router_aux_loss_coef: float = 0.05\n",
+    "    max_loop_iters: int = 4\n",
+    "    act_threshold: float = 0.5\n",
+    "    max_position_embeddings: int = 131072\n",
+    "    rope_theta: float = 10000000.0\n",
+    "    rope_scaling: dict = None\n",
+    "    sliding_window: int = 4096\n",
+    "    attention_dropout: float = 0.0\n",
+    "    rms_norm_eps: float = 1e-6\n",
+    "    initializer_range: float = 0.02\n",
+    "    use_cache: bool = True\n",
+    "    tie_word_embeddings: bool = True\n",
+    "    prelude_layers: int = 2\n",
+    "    coda_layers: int = 2\n",
+    "    lora_rank: int = 32\n",
+    "    loop_embed_dim: int = 48\n",
+    "    vision_hidden_size: int = 384\n",
+    "    audio_hidden_size: int = 512\n",
+    "    vision_num_frames: int = 60\n",
+    "    vision_tokens_per_frame: int = 256\n",
+    "    vision_temporal_tokens: int = 64\n",
+    "    vision_temporal_layers: int = 2\n",
+    "    model_type: str = \"spiderportal\"\n",
+    "    torch_dtype: str = \"bfloat16\"\n",
+    "\n",
+    "def loop_index_embedding(h, loop_t, loop_dim, theta=10000.0):\n",
+    "    freqs = 1.0 / (theta ** (torch.arange(0, loop_dim, 2, device=h.device, dtype=h.dtype) / loop_dim))\n",
+    "    angles = loop_t * freqs\n",
+    "    emb = torch.cat([angles.sin(), angles.cos()], dim=-1)[:loop_dim]\n",
+    "    emb_full = torch.zeros(h.shape[-1], device=h.device, dtype=h.dtype)\n",
+    "    emb_full[:loop_dim] = emb\n",
+    "    return h + emb_full.unsqueeze(0).unsqueeze(0)\n",
+    "\n",
+    "class SpiderPortalRMSNorm(nn.Module):\n",
+    "    def __init__(self, hidden_size, eps=1e-6):\n",
+    "        super().__init__()\n",
+    "        self.weight = nn.Parameter(torch.ones(hidden_size))\n",
+    "        self.variance_epsilon = eps\n",
+    "    def forward(self, hidden_states):\n",
+    "        input_dtype = hidden_states.dtype\n",
+    "        hidden_states = hidden_states.to(torch.float32)\n",
+    "        variance = hidden_states.pow(2).mean(-1, keepdim=True)\n",
+    "        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)\n",
+    "        return self.weight.to(input_dtype) * hidden_states.to(input_dtype)\n",
+    "\n",
+    "def compute_yarn_inv_freq(head_dim, rope_theta, factor, orig_max, beta_fast=32.0, beta_slow=1.0):\n",
+    "    dim = head_dim\n",
+    "    orig_inv_freq = 1.0 / (rope_theta ** (torch.arange(0, dim, 2).float() / dim))\n",
+    "    pos_freqs = torch.arange(0, dim, 2).float() / dim\n",
+    "    beta = (pos_freqs * math.log(rope_theta) / math.log(orig_max))\n",
+    "    scale = torch.where(beta < beta_slow, torch.ones_like(beta), torch.where(beta > beta_fast, torch.ones_like(beta) / factor, 1.0 - (beta - beta_slow) / (beta_fast - beta_slow) * (1.0 - 1.0 / factor)))\n",
+    "    return orig_inv_freq * scale\n",
+    "\n",
+    "class SpiderPortalGQA(nn.Module):\n",
+    "    def __init__(self, config):\n",
+    "        super().__init__()\n",
+    "        self.config = config\n",
+    "        self.hidden_size = config.hidden_size\n",
+    "        self.num_heads = config.num_attention_heads\n",
+    "        self.num_kv_heads = config.num_key_value_heads\n",
+    "        self.head_dim = config.hidden_size // config.num_attention_heads\n",
+    "        self.num_key_value_groups = self.num_heads // self.num_kv_heads\n",
+    "        self.q_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)\n",
+    "        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)\n",
+    "        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)\n",
+    "        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)\n",
+    "        self.attention_dropout = config.attention_dropout\n",
+    "        rope_scaling = getattr(config, 'rope_scaling', None)\n",
+    "        if rope_scaling and rope_scaling.get(\"type\") == \"yarn\":\n",
+    "            factor = rope_scaling.get(\"factor\", 1.0)\n",
+    "            orig_max_pos = rope_scaling.get(\"original_max_position_embeddings\", config.max_position_embeddings)\n",
+    "            inv_freq = compute_yarn_inv_freq(self.head_dim, config.rope_theta, factor, orig_max_pos)\n",
+    "        else:\n",
+    "            inv_freq = 1.0 / (config.rope_theta ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim))\n",
+    "        self.register_buffer(\"inv_freq\", inv_freq, persistent=False)\n",
+    "    def _rotate_half(self, x):\n",
+    "        x1 = x[..., :x.shape[-1] // 2]\n",
+    "        x2 = x[..., x.shape[-1] // 2:]\n",
+    "        return torch.cat((-x2, x1), dim=-1)\n",
+    "    def _apply_rotary(self, x, cos, sin):\n",
+    "        return (x * cos) + (self._rotate_half(x) * sin)\n",
+    "    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, use_cache=False):\n",
+    "        bsz, q_len, _ = hidden_states.size()\n",
+    "        query_states = self.q_proj(hidden_states)\n",
+    "        key_states = self.k_proj(hidden_states)\n",
+    "        value_states = self.v_proj(hidden_states)\n",
+    "        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
+    "        key_states = key_states.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)\n",
+    "        value_states = value_states.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)\n",
+    "        if position_ids is None:\n",
+    "            position_ids = torch.arange(q_len, device=hidden_states.device).unsqueeze(0).expand(bsz, -1)\n",
+    "        max_pos = position_ids.max().item() + 1\n",
+    "        seq_len = max(max_pos, q_len)\n",
+    "        t = torch.arange(seq_len, device=hidden_states.device, dtype=self.inv_freq.dtype)\n",
+    "        freqs = torch.outer(t, self.inv_freq)\n",
+    "        emb = torch.cat((freqs, freqs), dim=-1)\n",
+    "        cos, sin = emb.cos(), emb.sin()\n",
+    "        cos = cos[position_ids].unsqueeze(1)\n",
+    "        sin = sin[position_ids].unsqueeze(1)\n",
+    "        query_states = self._apply_rotary(query_states, cos, sin)\n",
+    "        key_states = self._apply_rotary(key_states, cos, sin)\n",
+    "        if past_key_value is not None:\n",
+    "            key_states = torch.cat([past_key_value[0], key_states], dim=2)\n",
+    "            value_states = torch.cat([past_key_value[1], value_states], dim=2)\n",
+    "        past_kv = (key_states, value_states) if use_cache else None\n",
+    "        key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1)\n",
+    "        value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1)\n",
+    "        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)\n",
+    "        if attention_mask is not None:\n",
+    "            attn_weights = attn_weights + attention_mask\n",
+    "        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(value_states.dtype)\n",
+    "        attn_weights = F.dropout(attn_weights, p=self.attention_dropout, training=self.training)\n",
+    "        attn_output = torch.matmul(attn_weights, value_states)\n",
+    "        attn_output = attn_output.transpose(1, 2).contiguous()\n",
+    "        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)\n",
+    "        return self.o_proj(attn_output), past_kv\n",
+    "\n",
+    "class SpiderPortalExpert(nn.Module):\n",
+    "    def __init__(self, config, intermediate_size=None):\n",
+    "        super().__init__()\n",
+    "        inter_size = intermediate_size or config.intermediate_size\n",
+    "        self.gate_proj = nn.Linear(config.hidden_size, inter_size, bias=False)\n",
+    "        self.up_proj = nn.Linear(config.hidden_size, inter_size, bias=False)\n",
+    "        self.down_proj = nn.Linear(inter_size, config.hidden_size, bias=False)\n",
+    "        self.act_fn = nn.SiLU()\n",
+    "    def forward(self, hidden_states):\n",
+    "        return self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))\n",
+    "\n",
+    "class SpiderPortalRouter(nn.Module):\n",
+    "    def __init__(self, config):\n",
+    "        super().__init__()\n",
+    "        self.num_experts = config.num_experts\n",
+    "        self.num_experts_per_tok = config.num_experts_per_tok\n",
+    "        self.weight = nn.Parameter(torch.randn(config.hidden_size, config.num_experts) * config.initializer_range)\n",
+    "        self.register_buffer(\"router_bias\", torch.zeros(config.num_experts))\n",
+    "    def forward(self, hidden_states):\n",
+    "        router_logits = hidden_states.view(-1, hidden_states.size(-1)) @ self.weight\n",
+    "        routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float32)\n",
+    "        biased_logits = router_logits + self.router_bias\n",
+    "        biased_weights = F.softmax(biased_logits, dim=-1, dtype=torch.float32)\n",
+    "        top_weights, top_indices = torch.topk(biased_weights, self.num_experts_per_tok, dim=-1)\n",
+    "        top_weights = top_weights / top_weights.sum(dim=-1, keepdim=True)\n",
+    "        top_weights = top_weights.to(hidden_states.dtype)\n",
+    "        mean_probs = routing_weights.mean(dim=0)\n",
+    "        aux_loss = self.num_experts * (mean_probs * mean_probs).sum()\n",
+    "        return top_weights, top_indices, aux_loss\n",
+    "\n",
+    "class SpiderPortalMoE(nn.Module):\n",
+    "    def __init__(self, config):\n",
+    "        super().__init__()\n",
+    "        self.config = config\n",
+    "        self.num_experts = config.num_experts\n",
+    "        self.num_experts_per_tok = config.num_experts_per_tok\n",
+    "        self.experts = nn.ModuleList([SpiderPortalExpert(config) for _ in range(config.num_experts)])\n",
+    "        self.shared_expert = SpiderPortalExpert(config)\n",
+    "        self.router = SpiderPortalRouter(config)\n",
+    "    def forward(self, hidden_states):\n",
+    "        batch_size, seq_len, hidden_dim = hidden_states.shape\n",
+    "        top_weights, top_indices, aux_loss = self.router(hidden_states)\n",
+    "        flat_hidden = hidden_states.view(-1, hidden_dim)\n",
+    "        final_output = torch.zeros_like(flat_hidden)\n",
+    "        for expert_idx in range(self.num_experts_per_tok):\n",
+    "            expert_ids = top_indices[:, expert_idx]\n",
+    "            expert_weights = top_weights[:, expert_idx:expert_idx+1]\n",
+    "            for e in range(self.num_experts):\n",
+    "                mask = expert_ids == e\n",
+    "                if mask.any():\n",
+    "                    expert_output = self.experts[e](flat_hidden[mask])\n",
+    "                    final_output[mask] += expert_output * expert_weights[mask]\n",
+    "        shared_output = self.shared_expert(flat_hidden)\n",
+    "        final_output = final_output + shared_output\n",
+    "        return final_output.view(batch_size, seq_len, hidden_dim), aux_loss\n",
+    "\n",
+    "class SpiderPortalDenseLayer(nn.Module):\n",
+    "    def __init__(self, config):\n",
+    "        super().__init__()\n",
+    "        self.self_attn = SpiderPortalGQA(config)\n",
+    "        dense_intermediate = config.hidden_size * 4 // 3\n",
+    "        self.ffn = SpiderPortalExpert(config, intermediate_size=dense_intermediate)\n",
+    "        self.input_layernorm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)\n",
+    "        self.post_attention_layernorm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)\n",
+    "    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, use_cache=False):\n",
+    "        attn_input = self.input_layernorm(hidden_states)\n",
+    "        attn_output, past_kv = self.self_attn(attn_input, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, use_cache=use_cache)\n",
+    "        hidden_states = hidden_states + attn_output\n",
+    "        ffn_input = self.post_attention_layernorm(hidden_states)\n",
+    "        ffn_output = self.ffn(ffn_input)\n",
+    "        hidden_states = hidden_states + ffn_output\n",
+    "        return hidden_states, past_kv\n",
+    "\n",
+    "class SpiderPortalMoELayer(nn.Module):\n",
+    "    def __init__(self, config, layer_idx):\n",
+    "        super().__init__()\n",
+    "        self.layer_idx = layer_idx\n",
+    "        self.self_attn = SpiderPortalGQA(config)\n",
+    "        self.moe = SpiderPortalMoE(config)\n",
+    "        self.input_layernorm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)\n",
+    "        self.post_attention_layernorm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)\n",
+    "    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, use_cache=False):\n",
+    "        attn_input = self.input_layernorm(hidden_states)\n",
+    "        attn_output, past_kv = self.self_attn(attn_input, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, use_cache=use_cache)\n",
+    "        hidden_states = hidden_states + attn_output\n",
+    "        moe_input = self.post_attention_layernorm(hidden_states)\n",
+    "        moe_output, aux_loss = self.moe(moe_input)\n",
+    "        hidden_states = hidden_states + moe_output\n",
+    "        return hidden_states, aux_loss, past_kv\n",
+    "\n",
+    "class LTIInjection(nn.Module):\n",
+    "    def __init__(self, config):\n",
+    "        super().__init__()\n",
+    "        self.hidden_size = config.hidden_size\n",
+    "        self.log_A = nn.Parameter(torch.full((config.hidden_size,), -2.0))\n",
+    "        self.delta_t = nn.Parameter(torch.tensor(1.0))\n",
+    "        self.B = nn.Linear(config.hidden_size, config.hidden_size, bias=False)\n",
+    "        with torch.no_grad():\n",
+    "            self.B.weight.data.normal_(mean=0.0, std=0.01)\n",
+    "    def get_A(self):\n",
+    "        return -torch.exp(self.log_A)\n",
+    "    def forward(self, h_t, e):\n",
+    "        A = self.get_A()\n",
+    "        return A * h_t + self.B(e)\n",
+    "\n",
+    "class ACTHalting(nn.Module):\n",
+    "    def __init__(self, config):\n",
+    "        super().__init__()\n",
+    "        self.halt_predictor = nn.Linear(config.hidden_size, 1)\n",
+    "        self.threshold = config.act_threshold\n",
+    "    def forward(self, hidden_states):\n",
+    "        return torch.sigmoid(self.halt_predictor(hidden_states))\n",
+    "\n",
+    "class LoRAAdapter(nn.Module):\n",
+    "    def __init__(self, config):\n",
+    "        super().__init__()\n",
+    "        rank = config.lora_rank\n",
+    "        self.down = nn.Linear(config.hidden_size, rank, bias=False)\n",
+    "        self.B = nn.Parameter(torch.randn(rank, config.hidden_size) * 0.02)\n",
+    "        self.scale = nn.Embedding(config.max_loop_iters, rank)\n",
+    "        with torch.no_grad():\n",
+    "            self.scale.weight.data.zero_()\n",
+    "            self.down.weight.data.normal_(mean=0.0, std=0.001)\n",
+    "    def forward(self, x, loop_t):\n",
+    "        max_t = self.scale.num_embeddings - 1\n",
+    "        t_idx = min(loop_t, max_t)\n",
+    "        s = self.scale(torch.tensor(t_idx, device=x.device))\n",
+    "        down = self.down(x) * s\n",
+    "        return down @ self.B\n",
+    "\n",
+    "class SpiderPortalMoEModel(nn.Module):\n",
+    "    def __init__(self, config):\n",
+    "        super().__init__()\n",
+    "        self.config = config\n",
+    "        self.prelude_layers = nn.ModuleList([SpiderPortalDenseLayer(config) for _ in range(config.prelude_layers)])\n",
+    "        self.recurrent_layers = nn.ModuleList([SpiderPortalMoELayer(config, i) for i in range(config.num_hidden_layers)])\n",
+    "        self.coda_layers = nn.ModuleList([SpiderPortalDenseLayer(config) for _ in range(config.coda_layers)])\n",
+    "        self.norm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)\n",
+    "        self.injection = LTIInjection(config)\n",
+    "        self.act_halting = ACTHalting(config)\n",
+    "        self.lora_adapter = LoRAAdapter(config)\n",
+    "        self.loop_embed_dim = config.loop_embed_dim\n",
+    "    def forward(self, hidden_states, input_embedding=None, attention_mask=None, position_ids=None, past_key_values=None, use_cache=False, n_loops=None):\n",
+    "        n_loops = n_loops or self.config.max_loop_iters\n",
+    "        input_embedding = input_embedding if input_embedding is not None else hidden_states\n",
+    "        total_aux_loss = 0.0\n",
+    "        for layer in self.prelude_layers:\n",
+    "            hidden_states, _ = layer(hidden_states, attention_mask=attention_mask, position_ids=position_ids)\n",
+    "        e = hidden_states.clone()\n",
+    "        B, T_seq, D = hidden_states.shape\n",
+    "        halted = torch.zeros(B, T_seq, device=hidden_states.device, dtype=torch.bool)\n",
+    "        cumulative_p = torch.zeros(B, T_seq, device=hidden_states.device, dtype=hidden_states.dtype)\n",
+    "        h_out = torch.zeros_like(hidden_states)\n",
+    "        past_key_values = past_key_values if past_key_values is not None else [None] * len(self.recurrent_layers)\n",
+    "        for t in range(n_loops):\n",
+    "            h_loop = loop_index_embedding(hidden_states, t, self.loop_embed_dim)\n",
+    "            if t > 0:\n",
+    "                injection = self.injection(hidden_states, input_embedding)\n",
+    "                hidden_states = hidden_states + injection\n",
+    "            new_past_key_values = []\n",
+    "            for i, layer in enumerate(self.recurrent_layers):\n",
+    "                hidden_states, aux_loss, past_kv = layer(hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_values[i] if t == 0 else None, use_cache=use_cache)\n",
+    "                total_aux_loss = total_aux_loss + aux_loss\n",
+    "                new_past_key_values.append(past_kv)\n",
+    "            lora_delta = self.lora_adapter(hidden_states, t)\n",
+    "            hidden_states = hidden_states + lora_delta\n",
+    "            halt_prob = self.act_halting(hidden_states).squeeze(-1)\n",
+    "            still_running = ~halted\n",
+    "            remainder = (1.0 - cumulative_p).clamp(min=0)\n",
+    "            weight = torch.where(cumulative_p + halt_prob >= self.config.act_threshold, remainder, halt_prob)\n",
+    "            weight = weight * still_running.to(hidden_states.dtype)\n",
+    "            h_out = h_out + weight.unsqueeze(-1) * hidden_states\n",
+    "            cumulative_p = cumulative_p + halt_prob * still_running.to(hidden_states.dtype)\n",
+    "            halted = halted | (cumulative_p >= self.config.act_threshold)\n",
+    "            if halted.all() and not self.training:\n",
+    "                break\n",
+    "        never_halted = (~halted).to(hidden_states.dtype).unsqueeze(-1)\n",
+    "        hidden_states = h_out + never_halted * hidden_states\n",
+    "        for layer in self.coda_layers:\n",
+    "            hidden_states, _ = layer(hidden_states, attention_mask=attention_mask, position_ids=position_ids)\n",
+    "        hidden_states = self.norm(hidden_states)\n",
+    "        return hidden_states, total_aux_loss, new_past_key_values\n",
+    "\n",
+    "class SpiderPortalForConditionalGeneration(nn.Module):\n",
+    "    def __init__(self, config):\n",
+    "        super().__init__()\n",
+    "        self.config = config\n",
+    "        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)\n",
+    "        self.model = SpiderPortalMoEModel(config)\n",
+    "        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)\n",
+    "        if config.tie_word_embeddings:\n",
+    "            self.lm_head.weight = self.embed_tokens.weight\n",
+    "        self.apply(self._init_weights)\n",
+    "    def _init_weights(self, module):\n",
+    "        if isinstance(module, nn.Linear):\n",
+    "            if hasattr(self, 'model') and module is self.model.injection.B:\n",
+    "                return\n",
+    "            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)\n",
+    "            if module.bias is not None:\n",
+    "                module.bias.data.zero_()\n",
+    "        elif isinstance(module, nn.Embedding):\n",
+    "            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)\n",
+    "    def forward(self, input_ids, attention_mask=None, position_ids=None, labels=None, n_loops=None, use_cache=False):\n",
+    "        hidden_states = self.embed_tokens(input_ids)\n",
+    "        model_dtype = next(self.model.parameters()).dtype\n",
+    "        hidden_states = hidden_states.to(model_dtype)\n",
+    "        input_embedding = hidden_states.clone()\n",
+    "        if attention_mask is None:\n",
+    "            attention_mask = torch.ones_like(input_ids, dtype=torch.bool)\n",
+    "        causal_mask = torch.full((attention_mask.size(0), 1, attention_mask.size(1), attention_mask.size(1)), 0.0, dtype=hidden_states.dtype, device=hidden_states.device)\n",
+    "        causal_mask = causal_mask.masked_fill(~attention_mask.unsqueeze(1).unsqueeze(2), torch.finfo(hidden_states.dtype).min)\n",
+    "        causal_mask = causal_mask.triu(1)\n",
+    "        hidden_states, aux_loss, past_kv = self.model(hidden_states, input_embedding=input_embedding, attention_mask=causal_mask, position_ids=position_ids, use_cache=use_cache, n_loops=n_loops)\n",
+    "        logits = self.lm_head(hidden_states)\n",
+    "        loss = None\n",
+    "        if labels is not None:\n",
+    "            shift_logits = logits[..., :-1, :].contiguous()\n",
+    "            shift_labels = labels[..., 1:].contiguous()\n",
+    "            loss_fct = CrossEntropyLoss()\n",
+    "            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))\n",
+    "            loss = loss + self.config.router_aux_loss_coef * aux_loss\n",
+    "        return {\"loss\": loss, \"logits\": logits, \"aux_loss\": aux_loss, \"past_key_values\": past_kv}\n",
+    "    def get_num_params(self):\n",
+    "        total = sum(p.numel() for p in self.parameters())\n",
+    "        return {\"total\": total, \"trainable\": total}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Multi-GPU DDP Training\n",
+    "# DDP requires a standalone script (mp.spawn doesn't work in notebooks)\n",
+    "# The script is at scripts/train_ddp.py\n",
+    "!python scripts/train_ddp.py"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Test generation (after training)\n",
+    "device = torch.device(\"cuda:0\")\n",
+    "config = SpiderPortalConfig(\n",
+    "    prelude_layers=2, coda_layers=2, lora_rank=32,\n",
+    "    rope_theta=10000000.0, tie_word_embeddings=True,\n",
+    ")\n",
+    "model = SpiderPortalForConditionalGeneration(config)\n",
+    "\n",
+    "checkpoint = torch.load(\"/kaggle/working/spiderportal-v5-ep1/model.pt\", map_location=\"cpu\")\n",
+    "model.load_state_dict(checkpoint)\n",
+    "model = model.to(torch.bfloat16).to(device)\n",
+    "model.eval()\n",
+    "\n",
+    "from transformers import AutoTokenizer\n",
+    "tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n",
+    "tokenizer.pad_token = tokenizer.eos_token\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    prompt = \"Question: Instruction: What is the capital of France?\\nAnswer:\"\n",
+    "    input_ids = tokenizer(prompt, return_tensors=\"pt\")[\"input_ids\"].to(device)\n",
+    "    generated = input_ids.clone()\n",
+    "    for _ in range(32):\n",
+    "        outputs = model(generated, n_loops=1)\n",
+    "        next_token = torch.argmax(outputs[\"logits\"][:, -1, :], dim=-1, keepdim=True)\n",
+    "        generated = torch.cat([generated, next_token], dim=1)\n",
+    "        if next_token.item() == tokenizer.eos_token_id:\n",
+    "            break\n",
+    "    print(f\"Generated: {tokenizer.decode(generated[0])}\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  },
+  "accelerator": "GPU"
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

AI-Transfer/scripts/train_ddp.py

@@ -0,0 +1,601 @@
+#!/usr/bin/env python3
+"""SpiderPortal v5 — Multi-GPU DDP Training.
+
+Run from Kaggle notebook cell:
+    !python train_ddp.py
+
+Or directly:
+    python -m torch.distributed.run --nproc_per_node=2 train_ddp.py
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+import math
+import os
+import json
+import gc
+import random
+import time
+import subprocess
+from pathlib import Path
+from dataclasses import dataclass
+from typing import Optional, Tuple, Dict, List
+from torch.nn import CrossEntropyLoss
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+@dataclass
+class SpiderPortalConfig:
+    vocab_size: int = 50278
+    hidden_size: int = 384
+    num_hidden_layers: int = 8
+    num_attention_heads: int = 8
+    num_key_value_heads: int = 2
+    intermediate_size: int = 1024
+    hidden_act: str = "silu"
+    num_experts: int = 64
+    num_experts_per_tok: int = 1
+    num_shared_experts: int = 1
+    router_aux_loss_coef: float = 0.05
+    max_loop_iters: int = 4
+    act_threshold: float = 0.5
+    max_position_embeddings: int = 131072
+    rope_theta: float = 10000000.0
+    rope_scaling: dict = None
+    sliding_window: int = 4096
+    attention_dropout: float = 0.0
+    rms_norm_eps: float = 1e-6
+    initializer_range: float = 0.02
+    use_cache: bool = True
+    tie_word_embeddings: bool = True
+    prelude_layers: int = 2
+    coda_layers: int = 2
+    lora_rank: int = 32
+    loop_embed_dim: int = 48
+    vision_hidden_size: int = 384
+    audio_hidden_size: int = 512
+    vision_num_frames: int = 60
+    vision_tokens_per_frame: int = 256
+    vision_temporal_tokens: int = 64
+    vision_temporal_layers: int = 2
+    model_type: str = "spiderportal"
+    torch_dtype: str = "bfloat16"
+
+def loop_index_embedding(h, loop_t, loop_dim, theta=10000.0):
+    freqs = 1.0 / (theta ** (torch.arange(0, loop_dim, 2, device=h.device, dtype=h.dtype) / loop_dim))
+    angles = loop_t * freqs
+    emb = torch.cat([angles.sin(), angles.cos()], dim=-1)[:loop_dim]
+    emb_full = torch.zeros(h.shape[-1], device=h.device, dtype=h.dtype)
+    emb_full[:loop_dim] = emb
+    return h + emb_full.unsqueeze(0).unsqueeze(0)
+
+class SpiderPortalRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight.to(input_dtype) * hidden_states.to(input_dtype)
+
+def compute_yarn_inv_freq(head_dim, rope_theta, factor, orig_max, beta_fast=32.0, beta_slow=1.0):
+    dim = head_dim
+    orig_inv_freq = 1.0 / (rope_theta ** (torch.arange(0, dim, 2).float() / dim))
+    pos_freqs = torch.arange(0, dim, 2).float() / dim
+    beta = (pos_freqs * math.log(rope_theta) / math.log(orig_max))
+    scale = torch.where(beta < beta_slow, torch.ones_like(beta), torch.where(beta > beta_fast, torch.ones_like(beta) / factor, 1.0 - (beta - beta_slow) / (beta_fast - beta_slow) * (1.0 - 1.0 / factor)))
+    return orig_inv_freq * scale
+
+class SpiderPortalGQA(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.num_key_value_groups = self.num_heads // self.num_kv_heads
+        self.q_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.attention_dropout = config.attention_dropout
+        rope_scaling = getattr(config, 'rope_scaling', None)
+        if rope_scaling and rope_scaling.get("type") == "yarn":
+            factor = rope_scaling.get("factor", 1.0)
+            orig_max_pos = rope_scaling.get("original_max_position_embeddings", config.max_position_embeddings)
+            inv_freq = compute_yarn_inv_freq(self.head_dim, config.rope_theta, factor, orig_max_pos)
+        else:
+            inv_freq = 1.0 / (config.rope_theta ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+    def _rotate_half(self, x):
+        x1 = x[..., :x.shape[-1] // 2]
+        x2 = x[..., x.shape[-1] // 2:]
+        return torch.cat((-x2, x1), dim=-1)
+    def _apply_rotary(self, x, cos, sin):
+        return (x * cos) + (self._rotate_half(x) * sin)
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, use_cache=False):
+        bsz, q_len, _ = hidden_states.size()
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        if position_ids is None:
+            position_ids = torch.arange(q_len, device=hidden_states.device).unsqueeze(0).expand(bsz, -1)
+        max_pos = position_ids.max().item() + 1
+        seq_len = max(max_pos, q_len)
+        t = torch.arange(seq_len, device=hidden_states.device, dtype=self.inv_freq.dtype)
+        freqs = torch.outer(t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        cos, sin = emb.cos(), emb.sin()
+        cos = cos[position_ids].unsqueeze(1)
+        sin = sin[position_ids].unsqueeze(1)
+        query_states = self._apply_rotary(query_states, cos, sin)
+        key_states = self._apply_rotary(key_states, cos, sin)
+        if past_key_value is not None:
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+        past_kv = (key_states, value_states) if use_cache else None
+        key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1)
+        value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1)
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+        if attention_mask is not None:
+            attn_weights = attn_weights + attention_mask
+        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(value_states.dtype)
+        attn_weights = F.dropout(attn_weights, p=self.attention_dropout, training=self.training)
+        attn_output = torch.matmul(attn_weights, value_states)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+        return self.o_proj(attn_output), past_kv
+
+class SpiderPortalExpert(nn.Module):
+    def __init__(self, config, intermediate_size=None):
+        super().__init__()
+        inter_size = intermediate_size or config.intermediate_size
+        self.gate_proj = nn.Linear(config.hidden_size, inter_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, inter_size, bias=False)
+        self.down_proj = nn.Linear(inter_size, config.hidden_size, bias=False)
+        self.act_fn = nn.SiLU()
+    def forward(self, hidden_states):
+        return self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))
+
+class SpiderPortalRouter(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+        self.weight = nn.Parameter(torch.randn(config.hidden_size, config.num_experts) * config.initializer_range)
+        self.register_buffer("router_bias", torch.zeros(config.num_experts))
+    def forward(self, hidden_states):
+        router_logits = hidden_states.view(-1, hidden_states.size(-1)) @ self.weight
+        routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float32)
+        biased_logits = router_logits + self.router_bias
+        biased_weights = F.softmax(biased_logits, dim=-1, dtype=torch.float32)
+        top_weights, top_indices = torch.topk(biased_weights, self.num_experts_per_tok, dim=-1)
+        top_weights = top_weights / top_weights.sum(dim=-1, keepdim=True)
+        top_weights = top_weights.to(hidden_states.dtype)
+        mean_probs = routing_weights.mean(dim=0)
+        aux_loss = self.num_experts * (mean_probs * mean_probs).sum()
+        return top_weights, top_indices, aux_loss
+
+class SpiderPortalMoE(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.num_experts = config.num_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+        self.experts = nn.ModuleList([SpiderPortalExpert(config) for _ in range(config.num_experts)])
+        self.shared_expert = SpiderPortalExpert(config)
+        self.router = SpiderPortalRouter(config)
+    def forward(self, hidden_states):
+        batch_size, seq_len, hidden_dim = hidden_states.shape
+        top_weights, top_indices, aux_loss = self.router(hidden_states)
+        flat_hidden = hidden_states.view(-1, hidden_dim)
+        final_output = torch.zeros_like(flat_hidden)
+        for expert_idx in range(self.num_experts_per_tok):
+            expert_ids = top_indices[:, expert_idx]
+            expert_weights = top_weights[:, expert_idx:expert_idx+1]
+            for e in range(self.num_experts):
+                mask = expert_ids == e
+                if mask.any():
+                    expert_output = self.experts[e](flat_hidden[mask])
+                    final_output[mask] += expert_output * expert_weights[mask]
+        shared_output = self.shared_expert(flat_hidden)
+        final_output = final_output + shared_output
+        return final_output.view(batch_size, seq_len, hidden_dim), aux_loss
+
+class SpiderPortalDenseLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.self_attn = SpiderPortalGQA(config)
+        dense_intermediate = config.hidden_size * 4 // 3
+        self.ffn = SpiderPortalExpert(config, intermediate_size=dense_intermediate)
+        self.input_layernorm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, use_cache=False):
+        attn_input = self.input_layernorm(hidden_states)
+        attn_output, past_kv = self.self_attn(attn_input, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, use_cache=use_cache)
+        hidden_states = hidden_states + attn_output
+        ffn_input = self.post_attention_layernorm(hidden_states)
+        ffn_output = self.ffn(ffn_input)
+        hidden_states = hidden_states + ffn_output
+        return hidden_states, past_kv
+
+class SpiderPortalMoELayer(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.self_attn = SpiderPortalGQA(config)
+        self.moe = SpiderPortalMoE(config)
+        self.input_layernorm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, use_cache=False):
+        attn_input = self.input_layernorm(hidden_states)
+        attn_output, past_kv = self.self_attn(attn_input, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, use_cache=use_cache)
+        hidden_states = hidden_states + attn_output
+        moe_input = self.post_attention_layernorm(hidden_states)
+        moe_output, aux_loss = self.moe(moe_input)
+        hidden_states = hidden_states + moe_output
+        return hidden_states, aux_loss, past_kv
+
+class LTIInjection(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.log_A = nn.Parameter(torch.full((config.hidden_size,), -2.0))
+        self.delta_t = nn.Parameter(torch.tensor(1.0))
+        self.B = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        with torch.no_grad():
+            self.B.weight.data.normal_(mean=0.0, std=0.01)
+    def get_A(self):
+        return -torch.exp(self.log_A)
+    def forward(self, h_t, e):
+        A = self.get_A()
+        return A * h_t + self.B(e)
+
+class ACTHalting(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.halt_predictor = nn.Linear(config.hidden_size, 1)
+        self.threshold = config.act_threshold
+    def forward(self, hidden_states):
+        return torch.sigmoid(self.halt_predictor(hidden_states))
+
+class LoRAAdapter(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        rank = config.lora_rank
+        self.down = nn.Linear(config.hidden_size, rank, bias=False)
+        self.B = nn.Parameter(torch.randn(rank, config.hidden_size) * 0.02)
+        self.scale = nn.Embedding(config.max_loop_iters, rank)
+        with torch.no_grad():
+            self.scale.weight.data.zero_()
+            self.down.weight.data.normal_(mean=0.0, std=0.001)
+    def forward(self, x, loop_t):
+        max_t = self.scale.num_embeddings - 1
+        t_idx = min(loop_t, max_t)
+        s = self.scale(torch.tensor(t_idx, device=x.device))
+        down = self.down(x) * s
+        return down @ self.B
+
+class SpiderPortalMoEModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.prelude_layers = nn.ModuleList([SpiderPortalDenseLayer(config) for _ in range(config.prelude_layers)])
+        self.recurrent_layers = nn.ModuleList([SpiderPortalMoELayer(config, i) for i in range(config.num_hidden_layers)])
+        self.coda_layers = nn.ModuleList([SpiderPortalDenseLayer(config) for _ in range(config.coda_layers)])
+        self.norm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.injection = LTIInjection(config)
+        self.act_halting = ACTHalting(config)
+        self.lora_adapter = LoRAAdapter(config)
+        self.loop_embed_dim = config.loop_embed_dim
+    def forward(self, hidden_states, input_embedding=None, attention_mask=None, position_ids=None, past_key_values=None, use_cache=False, n_loops=None):
+        n_loops = n_loops or self.config.max_loop_iters
+        input_embedding = input_embedding if input_embedding is not None else hidden_states
+        total_aux_loss = 0.0
+        for layer in self.prelude_layers:
+            hidden_states, _ = layer(hidden_states, attention_mask=attention_mask, position_ids=position_ids)
+        e = hidden_states.clone()
+        B, T_seq, D = hidden_states.shape
+        halted = torch.zeros(B, T_seq, device=hidden_states.device, dtype=torch.bool)
+        cumulative_p = torch.zeros(B, T_seq, device=hidden_states.device, dtype=hidden_states.dtype)
+        h_out = torch.zeros_like(hidden_states)
+        past_key_values = past_key_values if past_key_values is not None else [None] * len(self.recurrent_layers)
+        for t in range(n_loops):
+            h_loop = loop_index_embedding(hidden_states, t, self.loop_embed_dim)
+            if t > 0:
+                injection = self.injection(hidden_states, input_embedding)
+                hidden_states = hidden_states + injection
+            new_past_key_values = []
+            for i, layer in enumerate(self.recurrent_layers):
+                hidden_states, aux_loss, past_kv = layer(hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_values[i] if t == 0 else None, use_cache=use_cache)
+                total_aux_loss = total_aux_loss + aux_loss
+                new_past_key_values.append(past_kv)
+            lora_delta = self.lora_adapter(hidden_states, t)
+            hidden_states = hidden_states + lora_delta
+            halt_prob = self.act_halting(hidden_states).squeeze(-1)
+            still_running = ~halted
+            remainder = (1.0 - cumulative_p).clamp(min=0)
+            weight = torch.where(cumulative_p + halt_prob >= self.config.act_threshold, remainder, halt_prob)
+            weight = weight * still_running.to(hidden_states.dtype)
+            h_out = h_out + weight.unsqueeze(-1) * hidden_states
+            cumulative_p = cumulative_p + halt_prob * still_running.to(hidden_states.dtype)
+            halted = halted | (cumulative_p >= self.config.act_threshold)
+            if halted.all() and not self.training:
+                break
+        never_halted = (~halted).to(hidden_states.dtype).unsqueeze(-1)
+        hidden_states = h_out + never_halted * hidden_states
+        for layer in self.coda_layers:
+            hidden_states, _ = layer(hidden_states, attention_mask=attention_mask, position_ids=position_ids)
+        hidden_states = self.norm(hidden_states)
+        return hidden_states, total_aux_loss, new_past_key_values
+
+class SpiderPortalForConditionalGeneration(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.model = SpiderPortalMoEModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.embed_tokens.weight
+        self.apply(self._init_weights)
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            if hasattr(self, 'model') and module is self.model.injection.B:
+                return
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+    def forward(self, input_ids, attention_mask=None, position_ids=None, labels=None, n_loops=None, use_cache=False):
+        hidden_states = self.embed_tokens(input_ids)
+        model_dtype = next(self.model.parameters()).dtype
+        hidden_states = hidden_states.to(model_dtype)
+        input_embedding = hidden_states.clone()
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
+        causal_mask = torch.full((attention_mask.size(0), 1, attention_mask.size(1), attention_mask.size(1)), 0.0, dtype=hidden_states.dtype, device=hidden_states.device)
+        causal_mask = causal_mask.masked_fill(~attention_mask.unsqueeze(1).unsqueeze(2), torch.finfo(hidden_states.dtype).min)
+        causal_mask = causal_mask.triu(1)
+        hidden_states, aux_loss, past_kv = self.model(hidden_states, input_embedding=input_embedding, attention_mask=causal_mask, position_ids=position_ids, use_cache=use_cache, n_loops=n_loops)
+        logits = self.lm_head(hidden_states)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+            loss = loss + self.config.router_aux_loss_coef * aux_loss
+        return {"loss": loss, "logits": logits, "aux_loss": aux_loss, "past_key_values": past_kv}
+    def get_num_params(self):
+        total = sum(p.numel() for p in self.parameters())
+        return {"total": total, "trainable": total}
+
+def train_ddp(local_rank, world_size):
+    dist.init_process_group("nccl", rank=local_rank, world_size=world_size)
+    torch.cuda.set_device(local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+    is_master = local_rank == 0
+
+    if is_master:
+        print(f"Training on {world_size} GPUs (DDP)")
+        for i in range(world_size):
+            p = torch.cuda.get_device_properties(i)
+            print(f"  GPU {i}: {torch.cuda.get_device_name(i)} ({p.total_memory / 1e9:.1f}GB)")
+
+    config = SpiderPortalConfig(
+        hidden_size=384, num_hidden_layers=8, num_attention_heads=8,
+        num_key_value_heads=2, intermediate_size=1024,
+        num_experts=64, num_experts_per_tok=1, num_shared_experts=1,
+        router_aux_loss_coef=0.05,         max_loop_iters=2,
+        prelude_layers=2, coda_layers=2, lora_rank=32,
+        rope_theta=10000000.0,
+        rope_scaling={"type": "yarn", "factor": 4.0, "original_max_position_embeddings": 32768},
+        max_position_embeddings=131072, sliding_window=4096,
+        tie_word_embeddings=True,
+    )
+
+    model = SpiderPortalForConditionalGeneration(config)
+    model = model.to(torch.bfloat16).to(device)
+    model = DDP(model, device_ids=[local_rank], find_unused_parameters=True)
+
+    if is_master:
+        params = model.module.get_num_params()
+        print(f"Model: {params['total']/1e6:.1f}M params")
+        print(f"Experts: {config.num_experts} routed + {config.num_shared_experts} shared")
+
+    BASE_LR = 2e-5
+    WARMUP_STEPS = 1000
+    optimizer = torch.optim.AdamW(model.parameters(), lr=BASE_LR, weight_decay=0.01)
+
+    import pandas as pd
+    data_dir = Path("/kaggle/input/datasets/cliworks/spiderp-custom-1")
+    all_records = []
+    if data_dir.exists():
+        pkl_file = data_dir / "spiderportal_combined.pkl"
+        if pkl_file.exists():
+            df = pd.read_pickle(pkl_file)
+            all_records = df.to_dict("records")
+        else:
+            for f in data_dir.glob("*.parquet"):
+                try:
+                    df_pq = pd.read_parquet(f)
+                    if all(c in df_pq.columns for c in ["instruction", "input", "output"]):
+                        all_records.extend(df_pq[["instruction", "input", "output"]].to_dict("records"))
+                except:
+                    pass
+
+    if not all_records:
+        if is_master:
+            print("No data found, creating synthetic data...")
+        all_records = [{"instruction": f"Question {i}: What is {i} + {i}?", "input": "", "output": f"The answer is {i+i}."} for i in range(10000)]
+
+    if is_master:
+        print(f"Loaded {len(all_records)} samples")
+
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    tokenizer.pad_token = tokenizer.eos_token
+
+    BATCH_SIZE = 8
+    GRAD_ACCUM = 2
+    MAX_LEN = 256
+    EPOCHS = 3
+    N_LOOPS = 1
+
+    if is_master:
+        print(f"Batch size: {BATCH_SIZE}, Grad accum: {GRAD_ACCUM}")
+        print(f"Effective batch: {BATCH_SIZE * GRAD_ACCUM * world_size}")
+        print(f"LR: {BASE_LR} with {WARMUP_STEPS}-step warmup")
+
+    def build_prompt(sample):
+        instruction = str(sample.get("instruction", "")).strip()
+        inp = str(sample.get("input", "")).strip()
+        output = str(sample.get("output", "")).strip()
+        if inp:
+            return f"Question: Instruction: {instruction}\nInput: {inp}\nAnswer: {output}\n"
+        return f"Question: Instruction: {instruction}\nAnswer: {output}\n"
+
+    def format_sample(sample, tokenizer, max_len):
+        text = build_prompt(sample) + tokenizer.eos_token
+        enc = tokenizer(text, truncation=True, max_length=max_len, padding="max_length", return_tensors="pt")
+        input_ids = enc["input_ids"].squeeze(0)
+        labels = input_ids.clone()
+        prefix_ids = tokenizer("Question:", add_special_tokens=False)["input_ids"]
+        mask_len = min(len(prefix_ids), labels.shape[0])
+        labels[:mask_len] = -100
+        return {"input_ids": input_ids, "labels": labels}
+
+    shard_size = len(all_records) // world_size
+    start_idx = local_rank * shard_size
+    end_idx = start_idx + shard_size if local_rank < world_size - 1 else len(all_records)
+    local_samples = all_records[start_idx:end_idx]
+
+    if is_master:
+        print(f"Data sharding: {len(all_records)} total -> {len(local_samples)} per GPU")
+
+    global_step = 0
+    best_loss = float('inf')
+    start_time = time.time()
+
+    for epoch in range(1, EPOCHS + 1):
+        random.shuffle(local_samples)
+        total_loss = 0
+        num_batches = 0
+        optimizer.zero_grad()
+
+        for i in range(0, len(local_samples), BATCH_SIZE):
+            batch_samples = local_samples[i:i+BATCH_SIZE]
+            if len(batch_samples) < BATCH_SIZE:
+                continue
+
+            if global_step < WARMUP_STEPS:
+                lr = BASE_LR * (global_step + 1) / WARMUP_STEPS
+                for param_group in optimizer.param_groups:
+                    param_group['lr'] = lr
+
+            batch = [format_sample(s, tokenizer, MAX_LEN) for s in batch_samples]
+            input_ids = torch.stack([b["input_ids"] for b in batch]).to(device)
+            labels = torch.stack([b["labels"] for b in batch]).to(device)
+
+            if global_step == 0 and is_master:
+                print("  [First forward pass - CUDA graph building...]")
+
+            outputs = model(input_ids=input_ids, labels=labels, n_loops=N_LOOPS)
+            loss = outputs["loss"] / GRAD_ACCUM
+            loss.backward()
+
+            if (i // BATCH_SIZE + 1) % GRAD_ACCUM == 0:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+                optimizer.step()
+                optimizer.zero_grad()
+                global_step += 1
+
+            total_loss += loss.item() * GRAD_ACCUM
+            num_batches += 1
+
+            if is_master and ((i // BATCH_SIZE) == 0 or (global_step < 20) or (global_step % 50 == 0)):
+                avg_loss = total_loss / max(num_batches, 1)
+                elapsed = time.time() - start_time
+                steps_per_hour = (global_step + 1) / elapsed * 3600 if elapsed > 0 else 0
+                current_lr = optimizer.param_groups[0]['lr']
+                print(f"Epoch {epoch}/{EPOCHS} | Step {global_step} | avg_loss={avg_loss:.4f} | LR={current_lr:.2e} | {steps_per_hour:.1f} steps/hr")
+
+            if is_master and global_step > 0 and global_step % 500 == 0:
+                output_dir = Path(f"/kaggle/working/spiderportal-v5-ep{epoch}-step{global_step}")
+                output_dir.mkdir(parents=True, exist_ok=True)
+                state_dict = {k: v.cpu() for k, v in model.module.state_dict().items()}
+                torch.save(state_dict, output_dir / "model.pt")
+                print(f"Saved checkpoint at step {global_step}")
+
+        if is_master:
+            avg_loss = total_loss / max(num_batches, 1)
+            print(f"Epoch {epoch}/{EPOCHS} | avg_loss={avg_loss:.4f} | Time: {(time.time()-start_time)/60:.1f}min")
+            output_dir = Path(f"/kaggle/working/spiderportal-v5-ep{epoch}")
+            output_dir.mkdir(parents=True, exist_ok=True)
+            state_dict = {k: v.cpu() for k, v in model.module.state_dict().items()}
+            torch.save(state_dict, output_dir / "model.pt")
+            if avg_loss < best_loss:
+                best_loss = avg_loss
+
+    if is_master:
+        print(f"Training complete! Best loss: {best_loss:.4f}")
+        print(f"Total time: {(time.time() - start_time)/60:.1f} minutes")
+
+    dist.destroy_process_group()
+
+if __name__ == "__main__":
+    import sys
+    print(f"Script started, CUDA devices: {torch.cuda.device_count()}")
+    sys.stdout.flush()
+    
+    num_gpus = torch.cuda.device_count()
+    
+    if num_gpus <= 1:
+        dist.init_process_group = lambda *a, **k: None
+        dist.destroy_process_group = lambda: None
+        train_ddp(0, 1)
+    else:
+        print(f"Launching DDP with {num_gpus} GPUs...")
+        sys.stdout.flush()
+        import os
+        import signal
+        
+        # Set NCCL environment before forking
+        os.environ["MASTER_ADDR"] = "127.0.0.1"
+        os.environ["MASTER_PORT"] = "29500"
+        os.environ["WORLD_SIZE"] = str(num_gpus)
+        os.environ["LOCAL_WORLD_SIZE"] = str(num_gpus)
+        
+        procs = []
+        for rank in range(num_gpus):
+            print(f"Forking rank {rank}...")
+            sys.stdout.flush()
+            pid = os.fork()
+            if pid == 0:
+                # Child process
+                os.environ["RANK"] = str(rank)
+                os.environ["LOCAL_RANK"] = str(rank)
+                print(f"Child {rank} starting training...")
+                sys.stdout.flush()
+                train_ddp(rank, num_gpus)
+                os._exit(0)
+            else:
+                procs.append(pid)
+                print(f"Parent: forked child {rank} with PID {pid}")
+                sys.stdout.flush()
+        
+        # Parent process: wait for children
+        print("Parent waiting for children...")
+        sys.stdout.flush()
+        for pid in procs:
+            os.waitpid(pid, 0)

AI-Transfer/scripts/train_single_gpu.py

@@ -0,0 +1,587 @@
+#!/usr/bin/env python3
+"""SpiderPortal v5 — Single-GPU Optimized Training.
+
+For RTX PRO 6000 (96GB) or similar large-VRAM GPU.
+No DDP, maximal batch size, torch.compile, pre-tokenized data.
+
+Usage:
+    python train_single_gpu.py
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import os
+import json
+import gc
+import random
+import time
+from pathlib import Path
+from dataclasses import dataclass
+from typing import Optional, Tuple, Dict, List
+from torch.nn import CrossEntropyLoss
+
+@dataclass
+class SpiderPortalConfig:
+    vocab_size: int = 50278
+    hidden_size: int = 384
+    num_hidden_layers: int = 8
+    num_attention_heads: int = 8
+    num_key_value_heads: int = 2
+    intermediate_size: int = 1024
+    hidden_act: str = "silu"
+    num_experts: int = 64
+    num_experts_per_tok: int = 1
+    num_shared_experts: int = 1
+    router_aux_loss_coef: float = 0.05
+    max_loop_iters: int = 4
+    act_threshold: float = 0.5
+    max_position_embeddings: int = 131072
+    rope_theta: float = 10000000.0
+    rope_scaling: dict = None
+    sliding_window: int = 4096
+    attention_dropout: float = 0.0
+    rms_norm_eps: float = 1e-6
+    initializer_range: float = 0.02
+    use_cache: bool = True
+    tie_word_embeddings: bool = True
+    prelude_layers: int = 2
+    coda_layers: int = 2
+    lora_rank: int = 32
+    loop_embed_dim: int = 48
+    vision_hidden_size: int = 384
+    audio_hidden_size: int = 512
+    vision_num_frames: int = 60
+    vision_tokens_per_frame: int = 256
+    vision_temporal_tokens: int = 64
+    vision_temporal_layers: int = 2
+    model_type: str = "spiderportal"
+    torch_dtype: str = "bfloat16"
+
+def loop_index_embedding(h, loop_t, loop_dim, theta=10000.0):
+    freqs = 1.0 / (theta ** (torch.arange(0, loop_dim, 2, device=h.device, dtype=h.dtype) / loop_dim))
+    angles = loop_t * freqs
+    emb = torch.cat([angles.sin(), angles.cos()], dim=-1)[:loop_dim]
+    emb_full = torch.zeros(h.shape[-1], device=h.device, dtype=h.dtype)
+    emb_full[:loop_dim] = emb
+    return h + emb_full.unsqueeze(0).unsqueeze(0)
+
+class SpiderPortalRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight.to(input_dtype) * hidden_states.to(input_dtype)
+
+def compute_yarn_inv_freq(head_dim, rope_theta, factor, orig_max, beta_fast=32.0, beta_slow=1.0):
+    dim = head_dim
+    orig_inv_freq = 1.0 / (rope_theta ** (torch.arange(0, dim, 2).float() / dim))
+    pos_freqs = torch.arange(0, dim, 2).float() / dim
+    beta = (pos_freqs * math.log(rope_theta) / math.log(orig_max))
+    scale = torch.where(beta < beta_slow, torch.ones_like(beta), torch.where(beta > beta_fast, torch.ones_like(beta) / factor, 1.0 - (beta - beta_slow) / (beta_fast - beta_slow) * (1.0 - 1.0 / factor)))
+    return orig_inv_freq * scale
+
+class SpiderPortalGQA(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.num_key_value_groups = self.num_heads // self.num_kv_heads
+        self.q_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        self.attention_dropout = config.attention_dropout
+        rope_scaling = getattr(config, 'rope_scaling', None)
+        if rope_scaling and rope_scaling.get("type") == "yarn":
+            factor = rope_scaling.get("factor", 1.0)
+            orig_max_pos = rope_scaling.get("original_max_position_embeddings", config.max_position_embeddings)
+            inv_freq = compute_yarn_inv_freq(self.head_dim, config.rope_theta, factor, orig_max_pos)
+        else:
+            inv_freq = 1.0 / (config.rope_theta ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+    def _rotate_half(self, x):
+        x1 = x[..., :x.shape[-1] // 2]
+        x2 = x[..., x.shape[-1] // 2:]
+        return torch.cat((-x2, x1), dim=-1)
+    def _apply_rotary(self, x, cos, sin):
+        return (x * cos) + (self._rotate_half(x) * sin)
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, use_cache=False):
+        bsz, q_len, _ = hidden_states.size()
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        if position_ids is None:
+            position_ids = torch.arange(q_len, device=hidden_states.device).unsqueeze(0).expand(bsz, -1)
+        max_pos = position_ids.max().item() + 1
+        seq_len = max(max_pos, q_len)
+        t = torch.arange(seq_len, device=hidden_states.device, dtype=self.inv_freq.dtype)
+        freqs = torch.outer(t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        cos, sin = emb.cos(), emb.sin()
+        cos = cos[position_ids].unsqueeze(1)
+        sin = sin[position_ids].unsqueeze(1)
+        query_states = self._apply_rotary(query_states, cos, sin)
+        key_states = self._apply_rotary(key_states, cos, sin)
+        if past_key_value is not None:
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+        past_kv = (key_states, value_states) if use_cache else None
+        key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1)
+        value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1)
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+        if attention_mask is not None:
+            attn_weights = attn_weights + attention_mask
+        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(value_states.dtype)
+        attn_weights = F.dropout(attn_weights, p=self.attention_dropout, training=self.training)
+        attn_output = torch.matmul(attn_weights, value_states)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+        return self.o_proj(attn_output), past_kv
+
+class SpiderPortalExpert(nn.Module):
+    def __init__(self, config, intermediate_size=None):
+        super().__init__()
+        inter_size = intermediate_size or config.intermediate_size
+        self.gate_proj = nn.Linear(config.hidden_size, inter_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, inter_size, bias=False)
+        self.down_proj = nn.Linear(inter_size, config.hidden_size, bias=False)
+        self.act_fn = nn.SiLU()
+    def forward(self, hidden_states):
+        return self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))
+
+class SpiderPortalRouter(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+        self.weight = nn.Parameter(torch.randn(config.hidden_size, config.num_experts) * config.initializer_range)
+        self.register_buffer("router_bias", torch.zeros(config.num_experts))
+    def forward(self, hidden_states):
+        router_logits = hidden_states.view(-1, hidden_states.size(-1)) @ self.weight
+        routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float32)
+        biased_logits = router_logits + self.router_bias
+        biased_weights = F.softmax(biased_logits, dim=-1, dtype=torch.float32)
+        top_weights, top_indices = torch.topk(biased_weights, self.num_experts_per_tok, dim=-1)
+        top_weights = top_weights / top_weights.sum(dim=-1, keepdim=True)
+        top_weights = top_weights.to(hidden_states.dtype)
+        mean_probs = routing_weights.mean(dim=0)
+        aux_loss = self.num_experts * (mean_probs * mean_probs).sum()
+        return top_weights, top_indices, aux_loss
+
+class SpiderPortalMoE(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.num_experts = config.num_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+        self.experts = nn.ModuleList([SpiderPortalExpert(config) for _ in range(config.num_experts)])
+        self.shared_expert = SpiderPortalExpert(config)
+        self.router = SpiderPortalRouter(config)
+    def forward(self, hidden_states):
+        batch_size, seq_len, hidden_dim = hidden_states.shape
+        top_weights, top_indices, aux_loss = self.router(hidden_states)
+        flat_hidden = hidden_states.view(-1, hidden_dim)
+        final_output = torch.zeros_like(flat_hidden)
+        for expert_idx in range(self.num_experts_per_tok):
+            expert_ids = top_indices[:, expert_idx]
+            expert_weights = top_weights[:, expert_idx:expert_idx+1]
+            for e in range(self.num_experts):
+                mask = expert_ids == e
+                if mask.any():
+                    expert_output = self.experts[e](flat_hidden[mask])
+                    final_output[mask] += expert_output * expert_weights[mask]
+        shared_output = self.shared_expert(flat_hidden)
+        final_output = final_output + shared_output
+        return final_output.view(batch_size, seq_len, hidden_dim), aux_loss
+
+class SpiderPortalDenseLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.self_attn = SpiderPortalGQA(config)
+        dense_intermediate = config.hidden_size * 4 // 3
+        self.ffn = SpiderPortalExpert(config, intermediate_size=dense_intermediate)
+        self.input_layernorm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, use_cache=False):
+        attn_input = self.input_layernorm(hidden_states)
+        attn_output, past_kv = self.self_attn(attn_input, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, use_cache=use_cache)
+        hidden_states = hidden_states + attn_output
+        ffn_input = self.post_attention_layernorm(hidden_states)
+        ffn_output = self.ffn(ffn_input)
+        hidden_states = hidden_states + ffn_output
+        return hidden_states, past_kv
+
+class SpiderPortalMoELayer(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.self_attn = SpiderPortalGQA(config)
+        self.moe = SpiderPortalMoE(config)
+        self.input_layernorm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, use_cache=False):
+        attn_input = self.input_layernorm(hidden_states)
+        attn_output, past_kv = self.self_attn(attn_input, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, use_cache=use_cache)
+        hidden_states = hidden_states + attn_output
+        moe_input = self.post_attention_layernorm(hidden_states)
+        moe_output, aux_loss = self.moe(moe_input)
+        hidden_states = hidden_states + moe_output
+        return hidden_states, aux_loss, past_kv
+
+class LTIInjection(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.log_A = nn.Parameter(torch.full((config.hidden_size,), -2.0))
+        self.delta_t = nn.Parameter(torch.tensor(1.0))
+        self.B = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+        with torch.no_grad():
+            self.B.weight.data.normal_(mean=0.0, std=0.01)
+    def get_A(self):
+        return -torch.exp(self.log_A)
+    def forward(self, h_t, e):
+        A = self.get_A()
+        return A * h_t + self.B(e)
+
+class ACTHalting(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.halt_predictor = nn.Linear(config.hidden_size, 1)
+        self.threshold = config.act_threshold
+    def forward(self, hidden_states):
+        return torch.sigmoid(self.halt_predictor(hidden_states))
+
+class LoRAAdapter(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        rank = config.lora_rank
+        self.down = nn.Linear(config.hidden_size, rank, bias=False)
+        self.B = nn.Parameter(torch.randn(rank, config.hidden_size) * 0.02)
+        self.scale = nn.Embedding(config.max_loop_iters, rank)
+        with torch.no_grad():
+            self.scale.weight.data.zero_()
+            self.down.weight.data.normal_(mean=0.0, std=0.001)
+    def forward(self, x, loop_t):
+        max_t = self.scale.num_embeddings - 1
+        t_idx = min(loop_t, max_t)
+        s = self.scale(torch.tensor(t_idx, device=x.device))
+        down = self.down(x) * s
+        return down @ self.B
+
+class SpiderPortalMoEModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.prelude_layers = nn.ModuleList([SpiderPortalDenseLayer(config) for _ in range(config.prelude_layers)])
+        self.recurrent_layers = nn.ModuleList([SpiderPortalMoELayer(config, i) for i in range(config.num_hidden_layers)])
+        self.coda_layers = nn.ModuleList([SpiderPortalDenseLayer(config) for _ in range(config.coda_layers)])
+        self.norm = SpiderPortalRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.injection = LTIInjection(config)
+        self.act_halting = ACTHalting(config)
+        self.lora_adapter = LoRAAdapter(config)
+        self.loop_embed_dim = config.loop_embed_dim
+    def forward(self, hidden_states, input_embedding=None, attention_mask=None, position_ids=None, past_key_values=None, use_cache=False, n_loops=None):
+        n_loops = n_loops or self.config.max_loop_iters
+        input_embedding = input_embedding if input_embedding is not None else hidden_states
+        total_aux_loss = 0.0
+        for layer in self.prelude_layers:
+            hidden_states, _ = layer(hidden_states, attention_mask=attention_mask, position_ids=position_ids)
+        e = hidden_states.clone()
+        B, T_seq, D = hidden_states.shape
+        halted = torch.zeros(B, T_seq, device=hidden_states.device, dtype=torch.bool)
+        cumulative_p = torch.zeros(B, T_seq, device=hidden_states.device, dtype=hidden_states.dtype)
+        h_out = torch.zeros_like(hidden_states)
+        past_key_values = past_key_values if past_key_values is not None else [None] * len(self.recurrent_layers)
+        for t in range(n_loops):
+            h_loop = loop_index_embedding(hidden_states, t, self.loop_embed_dim)
+            if t > 0:
+                injection = self.injection(hidden_states, input_embedding)
+                hidden_states = hidden_states + injection
+            new_past_key_values = []
+            for i, layer in enumerate(self.recurrent_layers):
+                hidden_states, aux_loss, past_kv = layer(hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_values[i] if t == 0 else None, use_cache=use_cache)
+                total_aux_loss = total_aux_loss + aux_loss
+                new_past_key_values.append(past_kv)
+            lora_delta = self.lora_adapter(hidden_states, t)
+            hidden_states = hidden_states + lora_delta
+            halt_prob = self.act_halting(hidden_states).squeeze(-1)
+            still_running = ~halted
+            remainder = (1.0 - cumulative_p).clamp(min=0)
+            weight = torch.where(cumulative_p + halt_prob >= self.config.act_threshold, remainder, halt_prob)
+            weight = weight * still_running.to(hidden_states.dtype)
+            h_out = h_out + weight.unsqueeze(-1) * hidden_states
+            cumulative_p = cumulative_p + halt_prob * still_running.to(hidden_states.dtype)
+            halted = halted | (cumulative_p >= self.config.act_threshold)
+            if halted.all() and not self.training:
+                break
+        never_halted = (~halted).to(hidden_states.dtype).unsqueeze(-1)
+        hidden_states = h_out + never_halted * hidden_states
+        for layer in self.coda_layers:
+            hidden_states, _ = layer(hidden_states, attention_mask=attention_mask, position_ids=position_ids)
+        hidden_states = self.norm(hidden_states)
+        return hidden_states, total_aux_loss, new_past_key_values
+
+class SpiderPortalForConditionalGeneration(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.model = SpiderPortalMoEModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.embed_tokens.weight
+        self.apply(self._init_weights)
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            if hasattr(self, 'model') and module is self.model.injection.B:
+                return
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+    def forward(self, input_ids, attention_mask=None, position_ids=None, labels=None, n_loops=None, use_cache=False):
+        hidden_states = self.embed_tokens(input_ids)
+        model_dtype = next(self.model.parameters()).dtype
+        hidden_states = hidden_states.to(model_dtype)
+        input_embedding = hidden_states.clone()
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
+        causal_mask = torch.full((attention_mask.size(0), 1, attention_mask.size(1), attention_mask.size(1)), 0.0, dtype=hidden_states.dtype, device=hidden_states.device)
+        causal_mask = causal_mask.masked_fill(~attention_mask.unsqueeze(1).unsqueeze(2), torch.finfo(hidden_states.dtype).min)
+        causal_mask = causal_mask.triu(1)
+        hidden_states, aux_loss, past_kv = self.model(hidden_states, input_embedding=input_embedding, attention_mask=causal_mask, position_ids=position_ids, use_cache=use_cache, n_loops=n_loops)
+        logits = self.lm_head(hidden_states)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+            loss = loss + self.config.router_aux_loss_coef * aux_loss
+        return {"loss": loss, "logits": logits, "aux_loss": aux_loss, "past_key_values": past_kv}
+    def get_num_params(self):
+        total = sum(p.numel() for p in self.parameters())
+        return {"total": total, "trainable": total}
+
+def train_single_gpu():
+    device = torch.device("cuda")
+    gpu_name = torch.cuda.get_device_name(0)
+    gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
+    print(f"GPU: {gpu_name} ({gpu_mem:.1f}GB)")
+
+    config = SpiderPortalConfig(
+        hidden_size=384, num_hidden_layers=8, num_attention_heads=8,
+        num_key_value_heads=2, intermediate_size=1024,
+        num_experts=64, num_experts_per_tok=1, num_shared_experts=1,
+        router_aux_loss_coef=0.05, max_loop_iters=2,
+        prelude_layers=2, coda_layers=2, lora_rank=32,
+        rope_theta=10000000.0,
+        rope_scaling={"type": "yarn", "factor": 4.0, "original_max_position_embeddings": 32768},
+        max_position_embeddings=131072, sliding_window=4096,
+        tie_word_embeddings=True,
+    )
+
+    print("Building model...")
+    model = SpiderPortalForConditionalGeneration(config)
+    model = model.to(torch.bfloat16).to(device)
+
+    params = model.get_num_params()
+    print(f"Model: {params['total']/1e6:.1f}M params")
+    print(f"Experts: {config.num_experts} routed + {config.num_shared_experts} shared")
+
+    try:
+        model = torch.compile(model, mode="reduce-overhead")
+        print("torch.compile: enabled")
+    except Exception:
+        print("torch.compile: not available, using eager mode")
+
+    BASE_LR = 2e-5
+    WARMUP_STEPS = 1000
+    optimizer = torch.optim.AdamW(model.parameters(), lr=BASE_LR, weight_decay=0.01)
+
+    import pandas as pd
+    data_dir = Path(__file__).parent / "data"
+    all_records = []
+    pkl_file = data_dir / "spiderportal_combined.pkl"
+    if pkl_file.exists():
+        print(f"Loading dataset from {pkl_file}...")
+        df = pd.read_pickle(pkl_file)
+        all_records = df.to_dict("records")
+    else:
+        print(f"No dataset found at {pkl_file}, creating synthetic data...")
+        all_records = [{"instruction": f"Question {i}: What is {i} + {i}?", "input": "", "output": f"The answer is {i+i}."} for i in range(10000)]
+
+    print(f"Loaded {len(all_records):,} samples")
+
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    tokenizer.pad_token = tokenizer.eos_token
+
+    BATCH_SIZE = 128
+    MAX_LEN = 256
+    EPOCHS = 3
+    N_LOOPS = 2
+
+    print(f"Batch size: {BATCH_SIZE} (no grad accum)")
+    print(f"Effective batch: {BATCH_SIZE}")
+    print(f"LR: {BASE_LR} with {WARMUP_STEPS}-step warmup")
+    print(f"Max seq len: {MAX_LEN}, N_LOOPS: {N_LOOPS}")
+
+    def build_prompt(sample):
+        instruction = str(sample.get("instruction", "")).strip()
+        inp = str(sample.get("input", "")).strip()
+        output = str(sample.get("output", "")).strip()
+        if inp:
+            return f"Question: Instruction: {instruction}\nInput: {inp}\nAnswer: {output}\n"
+        return f"Question: Instruction: {instruction}\nAnswer: {output}\n"
+
+    print("Pre-tokenizing dataset...")
+    prefix_ids = tokenizer("Question:", add_special_tokens=False)["input_ids"]
+    mask_len = len(prefix_ids)
+
+    pre_tokenized = []
+    for i, sample in enumerate(all_records):
+        instruction = str(sample.get("instruction", "")).strip()
+        inp = str(sample.get("input", "")).strip()
+        output = str(sample.get("output", "")).strip()
+        if inp:
+            text = f"Question: Instruction: {instruction}\nInput: {inp}\nAnswer: {output}\n" + tokenizer.eos_token
+        else:
+            text = f"Question: Instruction: {instruction}\nAnswer: {output}\n" + tokenizer.eos_token
+        enc = tokenizer(text, truncation=True, max_length=MAX_LEN, padding="max_length")
+        input_ids = enc["input_ids"]
+        labels = input_ids[:]
+        for j in range(min(mask_len, len(labels))):
+            labels[j] = -100
+        pre_tokenized.append((input_ids, labels))
+        if (i + 1) % 50000 == 0:
+            print(f"  Tokenized {i+1:,}/{len(all_records):,}")
+
+    print(f"Pre-tokenization complete: {len(pre_tokenized):,} samples")
+    del all_records
+    gc.collect()
+
+    global_step = 0
+    best_loss = float('inf')
+    start_time = time.time()
+    checkpoint_dir = Path("checkpoints")
+    checkpoint_dir.mkdir(exist_ok=True)
+    step_ckpt_files = []
+
+    for epoch in range(1, EPOCHS + 1):
+        if epoch > 1:
+            for f in step_ckpt_files:
+                if f.exists():
+                    f.unlink()
+                    print(f"  Deleted old step checkpoint: {f.name}")
+            step_ckpt_files.clear()
+            gc.collect()
+
+        indices = list(range(len(pre_tokenized)))
+        random.shuffle(indices)
+        total_loss = 0
+        num_batches = 0
+        optimizer.zero_grad()
+
+        for batch_start in range(0, len(indices), BATCH_SIZE):
+            batch_indices = indices[batch_start:batch_start + BATCH_SIZE]
+            if len(batch_indices) < BATCH_SIZE:
+                continue
+
+            if global_step < WARMUP_STEPS:
+                lr = BASE_LR * (global_step + 1) / WARMUP_STEPS
+                for param_group in optimizer.param_groups:
+                    param_group['lr'] = lr
+
+            batch_input_ids = []
+            batch_labels = []
+            for idx in batch_indices:
+                input_ids, labels = pre_tokenized[idx]
+                batch_input_ids.append(input_ids)
+                batch_labels.append(labels)
+
+            input_ids = torch.tensor(batch_input_ids, dtype=torch.long, device=device)
+            labels = torch.tensor(batch_labels, dtype=torch.long, device=device)
+
+            if global_step == 0:
+                print("  [First forward pass - compiling...]")
+
+            outputs = model(input_ids=input_ids, labels=labels, n_loops=N_LOOPS)
+            loss = outputs["loss"]
+            loss.backward()
+
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+            optimizer.step()
+            optimizer.zero_grad()
+            global_step += 1
+
+            total_loss += loss.item()
+            num_batches += 1
+
+            if (batch_start // BATCH_SIZE) == 0 or global_step < 20 or global_step % 100 == 0:
+                avg_loss = total_loss / max(num_batches, 1)
+                elapsed = time.time() - start_time
+                steps_per_hour = (global_step + 1) / elapsed * 3600 if elapsed > 0 else 0
+                current_lr = optimizer.param_groups[0]['lr']
+                samples_per_sec = (global_step * BATCH_SIZE) / elapsed if elapsed > 0 else 0
+                print(f"Epoch {epoch}/{EPOCHS} | Step {global_step} | loss={avg_loss:.4f} | LR={current_lr:.2e} | {steps_per_hour:.0f} steps/hr | {samples_per_sec:.0f} samples/sec")
+
+            if global_step > 0 and global_step % 500 == 0:
+                ckpt_path = checkpoint_dir / f"spiderportal-v5-ep{epoch}-step{global_step}.pt"
+                state_dict = {k: v.cpu() for k, v in model.state_dict().items()}
+                torch.save(state_dict, ckpt_path)
+                step_ckpt_files.append(ckpt_path)
+                size_mb = ckpt_path.stat().st_size / (1024 * 1024)
+                print(f"Saved weights-only checkpoint: {ckpt_path.name} ({size_mb:.0f}MB)")
+
+        avg_loss = total_loss / max(num_batches, 1)
+        epoch_time = (time.time() - start_time) / 60
+        print(f"Epoch {epoch}/{EPOCHS} complete | avg_loss={avg_loss:.4f} | Time: {epoch_time:.1f}min")
+
+        ckpt_path = checkpoint_dir / f"spiderportal-v5-ep{epoch}.pt"
+        torch.save({
+            "step": global_step,
+            "epoch": epoch,
+            "model_state_dict": {k: v.cpu() for k, v in model.state_dict().items()},
+            "optimizer_state_dict": optimizer.state_dict(),
+            "config": config.__dict__,
+        }, ckpt_path)
+        size_mb = ckpt_path.stat().st_size / (1024 * 1024)
+        print(f"Saved epoch checkpoint: {ckpt_path.name} ({size_mb:.0f}MB)")
+
+        if avg_loss < best_loss:
+            best_loss = avg_loss
+            best_path = checkpoint_dir / "spiderportal-v5-best.pt"
+            torch.save({
+                "step": global_step,
+                "epoch": epoch,
+                "model_state_dict": {k: v.cpu() for k, v in model.state_dict().items()},
+                "optimizer_state_dict": optimizer.state_dict(),
+                "config": config.__dict__,
+            }, best_path)
+            size_mb = best_path.stat().st_size / (1024 * 1024)
+            print(f"Saved best checkpoint: {best_path.name} ({size_mb:.0f}MB)")
+
+    total_time = (time.time() - start_time) / 3600
+    print(f"\nTraining complete!")
+    print(f"Best loss: {best_loss:.4f}")
+    print(f"Total time: {total_time:.2f} hours")
+    print(f"Total steps: {global_step}")
+    print(f"Checkpoints saved to: {checkpoint_dir}")
+
+if __name__ == "__main__":
+    train_single_gpu()

SpiderPortal-Base-v5/config.json

@@ -0,0 +1,45 @@
+{
+  "architectures": [
+    "SpiderPortalForConditionalGeneration"
+  ],
+  "model_type": "spiderportal",
+  "vocab_size": 50278,
+  "hidden_size": 384,
+  "num_hidden_layers": 8,
+  "num_attention_heads": 8,
+  "num_key_value_heads": 2,
+  "intermediate_size": 1024,
+  "hidden_act": "silu",
+  "max_position_embeddings": 131072,
+  "rope_theta": 10000000.0,
+  "rope_scaling": {
+    "type": "yarn",
+    "factor": 4.0,
+    "original_max_position_embeddings": 32768
+  },
+  "sliding_window": 4096,
+  "rms_norm_eps": 1e-06,
+  "initializer_range": 0.02,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "num_experts": 64,
+  "num_experts_per_tok": 1,
+  "num_shared_experts": 1,
+  "router_aux_loss_coef": 0.05,
+  "max_loop_iters": 4,
+  "act_threshold": 0.5,
+  "prelude_layers": 2,
+  "coda_layers": 2,
+  "lora_rank": 32,
+  "enable_talker": true,
+  "vision_encoder_type": "eupe_vit_s",
+  "vision_encoder_path": "facebook/EUPE-ViT-S",
+  "vision_hidden_size": 384,
+  "audio_encoder_type": "moonshine_base",
+  "audio_encoder_path": "UsefulSensors/moonshine-base",
+  "audio_hidden_size": 512,
+  "vision_num_frames": 60,
+  "vision_tokens_per_frame": 256,
+  "vision_temporal_tokens": 64,
+  "vision_temporal_layers": 2
+}

SpiderPortal-Base-v5/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:256d05574c2c2ffe38f5c3b5eb7badfc614e35ac7898452a4e43e922cc41d19c
+size 1288358204

SpiderPortal-Base-v5/special_tokens_map.json

@@ -0,0 +1,4 @@
+{
+  "pad_token": "<|endoftext|>",
+  "eos_token": "<|im_end|>"
+}

SpiderPortal-Base-v5/tokenizer_config.json

@@ -0,0 +1 @@
+{"model_max_length": 1024}