GDN-Hybrid: Fix warmdown/SWA/QAT timing — 1.0208 BPB

[OE-GOD]

Summary

PR #1545's GDN-Hybrid has three training systems (warmdown, SWA, Late QAT) that are configured but never activate due to a timing mismatch:

• ITERATIONS=9999 + WARMDOWN_ITERS=3000 → warmdown starts at step 6999
• Training only reaches ~2100 steps in 590s wallclock
• Result: zero warmdown, zero SWA, zero QAT

The fix (2 lines)

# Before (PR #1545):
iterations = int(os.environ.get("ITERATIONS", 9999))
warmdown_iters = int(os.environ.get("WARMDOWN_ITERS", 3000))

# After:
iterations = int(os.environ.get("ITERATIONS", 2200))
warmdown_iters = int(os.environ.get("WARMDOWN_ITERS", 400))

This activates all three systems:

• Warmdown starts at step 1800 (cosine decay lr 1.0 → 0.0)
• SWA triggers at step 2100 (lr_mul < 0.2), collects 3 checkpoints
• Late QAT triggers at step 2099, reducing quantization degradation

Also fixes a SWA device mismatch bug on line 1037 where swa_avg tensors on CPU were mixed with avg_state tensors on CUDA.

Results

8×H100, seed=42, SP1024, runpod/parameter-golf:latest, 590s wallclock:

	Without fix (PR #1545 defaults)	With fix
Training steps	2119	2200
Warmdown	None (never triggers)	Steps 1800-2200
SWA checkpoints	0	3
Late QAT	Never	Step 2099
EMA BPB	1.0109	1.0064
Quantized BPB	1.0243	1.0208
Artifact	14.91 MB	15.03 MB

Improvement: -0.0035 BPB from fixing the training pipeline alone.

The insight

The model architecture is already good. The waste is in the pipeline around it — three systems designed to improve final model quality were silently disabled by a config mismatch. Same pattern as PR #1510 (ANS compression): optimize the boring parts.

How to reproduce

ARCH_MODE=D VOCAB_SIZE=1024 \
DATA_PATH=./data/datasets/fineweb10B_sp1024 \
TOKENIZER_PATH=./data/tokenizers/fineweb_1024_bpe.model \
VAL_LOSS_EVERY=9999 SEED=42 \
ITERATIONS=2200 WARMDOWN_ITERS=400 \
torchrun --standalone --nproc_per_node=8 gdn/gdn_train_gpt.py

Test plan

• Verified warmdown triggers at step 1800 (lr_mul starts decreasing)
• Verified SWA starts at step 2100 (3 checkpoints collected)
• Verified Late QAT at step 2099
• Quantized BPB: 1.0208 (artifact 15.03 MB, under 16 MB)
• Compared with unfixed baseline on same hardware/seed

Based on PR #1545 GDN-Hybrid architecture.

[OE-GOD]

Closing — BPB numbers computed with incorrect byte counting. Will resubmit with corrected evaluation.