Memmap multi-shard data pipeline + GPU prefetch + LeakyReLU² + Legal TTT + Parallel Muon

val_bpb: 1.1147 (3-seed mean, std 0.0005) | ~15.23 MB (mean int6+lzma + code) | 8×H100 80GB HBM3

Results (8×H100 80GB HBM3, PyTorch 2.10.0+cu126)

Pre-TTT bpb is final_int6_sliding_window_exact (GPTQ-lite int6 dequantized weights, sliding window, stride 64). Post-TTT bpb is legal_ttt_exact after score-first chunk adaptation. Artifact is Total submission size int6+lzma (compressed weights + UTF-8 train_gpt.py bytes).

Key changes

1. Training data pipeline

The old loader walked shards in fixed order and took contiguous spans per rank. The new DistributedTokenLoader:

• Memory-maps each .bin shard (numpy.memmap) with a small header-token cache to avoid repeated header reads.
• Samples global training windows across shards: per batch, it draws multiple shards with probability weighted by remaining usable blocks (exponent tightens as training progresses), uses a coprime stride over valid 2048-token blocks inside each shard, and interleaves windows for batch diversity.
• Merges nearby reads on the same shard into one slab copy to reduce mmap churn.
• Overlaps I/O and compute: a daemon thread builds CPU (x, y) batches into a queue while the main thread trains; CUDA streams and events prefetch the next batch to GPU while the current step runs.

Together this replaces the old strictly sequential per-rank stream with a stratified, multi-shard mixture and async H2D overlap. Model forward, optimizer (Parallel Muon + AdamW), EMA/SWA, late QAT, and export paths are otherwise in the same family as the previous script.

2. Defaults and TTT

• TTT_ENABLED defaults to 1 in train_gpt.py (it was 0 in train_gpt_old.py), so legal score-first TTT runs at the end unless disabled.
• TTT_FREEZE_BLOCKS defaults to 2 for the logged runs (freeze_blocks=2 in ttt_sliding:start in the training logs). The prior README’s command used TTT_FREEZE_BLOCKS=0 (all blocks); this submission’s numbers are with two frozen bottom blocks during TTT unless you override the env var.

3. LeakyReLU² (unchanged from prior submission)

Same MLP nonlinearity as before:

x = F.leaky_relu(F.linear(x, up_w.to(x.dtype)), negative_slope=0.5)
return F.linear(x.square(), down_w.to(x.dtype))

LeakyReLU(0.5) keeps gradients on negative pre-activations; squaring preserves the non-negative “relu²-style” bias.