11L Int6 QAT + Per-Dim SmearGate + SWA: 1.1480 BPB (3-seed mean)

[baudrillardsgh0st]

Summary

• Mean val_bpb: 1.1480 (3 seeds, self-verified)
• Best single seed: 1.1453 (seed 1337)
• Artifact: 15.33 MiB (int6-in-int8 + zstd-22)
• 11 layers, 512 dim, 8 heads, 4 KV heads, MLP 3x, 26.5M params

Statistical Significance

Improvement over SOTA (1.1748): 0.0268 BPB (required: ≥0.005)

Seed	val_loss	val_bpb	Steps	ms/step	SWA ckpts
1337	1.9339	1.1453	8052	74.49	30
7	1.9400	1.1490	8040	74.34	29
42	1.9413	1.1498	7772	77.17	27

Mean: 1.1480 | Std: 0.0024 | t = 16.04 (df=2) | p < 0.001

One-sided t-test: H₀ that improvement < 0.005 nats rejected at p < 0.001.

Key Techniques

1. Per-Dimension SmearGate: Learned sigmoid(Parameter(dim)) gate blending current and previous token embeddings. More expressive than scalar gating — each embedding dimension gets its own blend ratio. Zero-initialized, 512 params.

2. Stochastic Weight Averaging (SWA): Every 50 steps over the last 50% of training (~29-30 checkpoints averaged). Produces smoother weight distributions that are more robust to int6 quantization. Pre-quant BPB 1.1666 → post-quant 1.1453 (only 0.021 degradation).

3. Int6 QAT with STE: Fake int6 quantization during forward pass. Per-row symmetric, 6-bit clipping. Int6 values stored in int8 containers — zstd-22 compresses ~35%.

4. High Muon weight decay (0.038): Keeps weights small for better int6 quantization fidelity and generalization.

5. FP16 tied embedding passthrough: Embedding/unembedding kept in fp16 to avoid compounding quantization error.

Training Config

Parameter	Value
Layers	11
Matrix LR	0.02
Scalar LR	0.02
Tied Embed LR	0.03
Muon Momentum	0.99 (warmup 0.92→0.99 over 1500 steps)
Muon Weight Decay	0.038
Warmdown Steps	3000
QAT Bits	6
SWA Every	50 steps
Batch tokens	524,288
Seq len	2048

Run command

NCCL_NVLS_ENABLE=0 \
VOCAB_SIZE=1024 NUM_HEADS=8 NUM_KV_HEADS=4 MODEL_DIM=512 \
NUM_LAYERS=11 MLP_MULT=3 \
QAT=1 QUANT_BITS=6 FP16_EMBED=1 LATE_K_LAYERS=0 \
EVAL_STRIDE=64 EVAL_BATCH_SEQS=32 \
MATRIX_LR=0.02 SCALAR_LR=0.02 TIED_EMBED_LR=0.03 \
MUON_MOMENTUM=0.99 MUON_WEIGHT_DECAY=0.038 \
SMEAR_GATE=1 BIGRAM_HASH=0 \
TRAIN_SEQ_LEN=2048 TRAIN_BATCH_TOKENS=524288 \
WARMDOWN_STEPS=3000 MAX_WALLCLOCK_SECONDS=600 \
SWA_ENABLED=1 SWA_START_FRAC=0.5 SWA_EVERY=50 \
torchrun --standalone --nproc_per_node=8 train_gpt.py