Non-record: Muon-Aware QAT + LAWA + Adaptive LR Scheduling

Status: Non-record (pending 8×H100 verification)

Summary

7 independently toggleable training improvements targeting the Muon optimizer's quantization sensitivity. Each technique was selected based on analysis of Muon's Newton-like orthogonalized update dynamics and the int8/int6 quantization pipeline's failure modes.

Techniques

What's Novel

Muon-aware QAT design: Most QAT implementations use vanilla STE, which injects discrete rounding noise into gradients. With Muon's orthogonalization step, this directional noise gets amplified rather than averaged out (unlike Adam). Our implementation:

• Two modes: STE (standard) and Gaussian noise (smoother for Muon)
• Late activation at 75% of training to preserve early momentum subspace
• Automatic 50% LR reduction when QAT activates
• Only targets large matrices (>65K params) — small control tensors stay clean

LAWA + QAT synergy: QAT teaches quant-robustness during training; LAWA smooths remaining outlier weights post-convergence. Complementary mechanisms targeting the same goal from different angles.

Research Foundation

Backed by targeted analysis of:

• Muon optimizer dynamics under quantization noise
• Depth recurrence at 15–50M parameter scale (Universal Transformers, "From Growing to Looping")
• BitNet b1.58 viability (concluded: not viable at this scale, confirmed by PR Non-record: BitNet b1.58 + depth recurrence + NorMuon (1.7510 BPB, 3.78 MB) #126)
• modded-nanogpt speedrun techniques

Next Steps (in-progress)

• Int6 quantization + MLP 3× expansion
• Sliding window evaluation
• NorMuon optimizer
• Depth recurrence (middle-looped: 1 prelude + 3 shared × 3 loops + 1 coda)
• 8×H100 verification runs

Usage

# Full stack
QAT_ENABLED=1 QAT_MODE=noise LAWA_ENABLED=1 SEQ_LEN_WARMUP=1 COMPRESSION=zstd \
  torchrun --standalone --nproc_per_node=8 train_gpt.py

All features default to safe values and can be independently toggled via environment variables. See README in submission folder for full documentation.