Star-ReLU MLP (learnable scale+bias) + run_swiglu.sh for openai#462 architecture + cosine TTT

Author: mrdavtan

records/track_10min_16mb/2026-03-21_11L_XSA_EMA_TTT/train_gpt.py

@@ -880,17 +880,20 @@ def forward(self, x: Tensor, lora: AttentionLoRA | None = None, v_embed: Tensor
 
 
 class MLP(nn.Module):
-    # relu^2 MLP from the original modded-nanogpt setup
+    # Star-ReLU: relu(x)^2 with learnable per-channel scale and bias (PR #462)
     def __init__(self, dim: int, mlp_mult: int, mlp_hidden: int = 0):
         super().__init__()
         hidden = mlp_hidden if mlp_hidden > 0 else mlp_mult * dim
         self.fc = CastedLinear(dim, hidden, bias=False)
         self.proj = CastedLinear(hidden, dim, bias=False)
         self.proj._zero_init = True
+        self.scale = nn.Parameter(torch.ones(hidden, dtype=torch.float32))
+        self.bias = nn.Parameter(torch.zeros(hidden, dtype=torch.float32))
 
     def forward(self, x: Tensor) -> Tensor:
-        x = torch.relu(self.fc(x))
-        return self.proj(x.square())
+        activated = torch.relu(self.fc(x)).square()
+        activated = activated * self.scale.to(dtype=activated.dtype) + self.bias.to(dtype=activated.dtype)
+        return self.proj(activated)
 
 
 class Block(nn.Module):

run_swiglu.sh

@@ -0,0 +1,57 @@
+#!/bin/bash
+# RUN_SWIGLU: Star-ReLU architecture (#462) + our cosine per-layer TTT
+#
+# Architecture from PR #462 (JoeProAI):
+#   Star-ReLU MLP with learnable scale+bias
+#   8 KV heads (full MHA, not GQA)
+#   MLP hidden=1792 (wider)
+#   BigramHash 8192 buckets
+#   XSA4 enabled, EMA decay=0.9985, warmdown=6000
+#
+# TTT schedule (ours):
+#   Cosine lr decay, per-layer lr, 50 epochs
+#
+# Target: sub-1.05
+
+set -e
+cd /workspace/parameter-golf
+git fetch origin && git checkout swiglu-cosine-ttt && git reset --hard origin/swiglu-cosine-ttt
+
+# Architecture (from #462)
+export TRAIN_SEQ_LEN=2048 EVAL_SEQ_LEN=2048 UNET_SKIPS=1
+export ROPE_DIMS=16 LN_SCALE=1 ROPE_BASE=10000
+export EVAL_STRIDE=64 DOC_ISOLATED_EVAL=0
+export NUM_KV_HEADS=8
+export MLP_HIDDEN=1792
+export BIGRAM_HASH_BUCKETS=8192
+export XSA_LAST_N=4
+export EMA_DECAY=0.9985
+export WARMDOWN_ITERS=6000
+export QAT=0
+export LATE_K_FP16=0 FP16_EMBED_EXPORT=0
+
+# Cosine per-layer TTT (ours)
+export TTT_OPTIMIZER=adamw
+export TTT_LR=0.0005
+export TTT_EPOCHS=50
+export TTT_COSINE=1
+export TTT_PERLAYER=1
+export TTT_FREEZE_BLOCKS=0
+export TTT_BATCH_SEQS=64
+export TTT_MAX_STEPS=9999
+
+# Seed from argument or default 1337
+export SEED=${1:-1337}
+
+unset MLP_HIDDEN_OLD QUANT_BITS RUN_ID TIER2_MODE MLP_MULT \
+  BACKOUT LAYER_DROP HEAD_DROP EVAL_TEMPERATURE \
+  MLP_QUANT_BITS USE_FA3 TRAIN_BATCH_TOKENS SWA PRUNE_PCT \
+  REPTILE_TTT VE_ENABLED TTT_TWO_PHASE
+
+echo "=== SWIGLU + COSINE TTT ==="
+echo "SEED=$SEED KV=$NUM_KV_HEADS MLP=$MLP_HIDDEN BIGRAM=$BIGRAM_HASH_BUCKETS XSA=$XSA_LAST_N"
+echo "TTT: AdamW ${TTT_EPOCHS}ep cosine perlayer EMA=$EMA_DECAY WD=$WARMDOWN_ITERS"
+echo "==========================="
+
+torchrun --standalone --nproc_per_node=8 \
+  records/track_10min_16mb/2026-03-21_11L_XSA_EMA_TTT/train_gpt.py