SubSixteen v2: Int6 QAT + MLP 3x + SWA + Sliding Window

val_bpb: 1.1708 | Artifact: 14,603,588 bytes (under 16MB)

Results

Architecture

• 9-layer GPT, 512-dim, GQA (8 heads, 4 KV heads), U-Net skip connections
• MLP 3x expansion (hidden=1536), relu² activation
• 1024 vocab SentencePiece BPE, tied embeddings
• Overtone SVD init, phase-transition resid_mix, NTK-aware RoPE
• Logit soft-capping (tanh, cap=30)

Key Techniques (v2)

1. STE fake-int6 QAT: CastedLinear weights fake-quantized to [-31,31] during forward via Straight-Through Estimator. Model learns distributions that survive 6-bit post-training quantization.
2. MLP 3x expansion: Hidden dim 1536 (up from 1024). Enabled by int6 saving ~4MB artifact space. 9 layers (down from 10) to fit under 16MB.
3. SWA (Stochastic Weight Averaging): 16 checkpoints collected every 200 steps during warmdown, averaged for export.
4. zstd-22 compression: Better compression ratio than zlib-9 for the quantized artifact.
5. Sliding window eval: stride=64, seq_len=4096 for near-full-context scoring.

Optimizer

• Muon (momentum=0.99, lr=0.02) for matrix params with Newton-Schulz orthogonalization
• AdamW for embeddings (lr=0.03) and scalars (lr=0.02)
• Warmdown: 3000 iters, wallclock-aware schedule
• Momentum warmup: 0.92 → 0.99 over 1500 steps

Run command

pip install zstandard
torchrun --standalone --nproc_per_node=8 train_gpt.py

Trained on 8xH100 SXM, 600s wallclock cap.