pengxiang/looped-moe-1.7b-tensorboard

Looped MoE Transformer (1.7B) training logs

TensorBoard scalar logs for a loop-count sweep of a looped MoE transformer. The loop count H is the only variable between runs.

Model

• 15 layers, hidden 1280, 20 heads / 10 KV heads
• MoE: 30 routed + 2 shared experts, top-8 (shared included), expert width 768
• 1.69B total parameters, about 0.63B active per loop
• The 15-layer block is applied H times (H = 1, 3, 6, 9, 12)

Training stages

Global batch 1024 sequences x 1024 tokens (1,048,576 tokens per step). AdamW (0.9, 0.95), weight decay 0.1, gradient clip 1.0.

Stage	Folder	Steps	Tokens	LR	Data
1	20b/H*	20,000	20.97B	WSD, peak 6e-5, 10% warmup, 10% linear decay to 6e-6	dolma3 + FineWeb-Edu, about 50:50
2	cont60b/H*	34,000 (aligned stop point)	35.65B	re-warm over 570 steps to 6e-5, then constant	dolma3 only, unseen data
3	anneal2b/H*	2,000	2.10B	linear 6e-5 to 6e-6	dolma3 only, continues stage 2's data order

Stage 2 starts from the stage 1 final weights and optimizer state. Stage 3 starts from the stage 2 step-34000 checkpoint. Stages 1 + 2 are 56.6B tokens, and stage 3 adds 2.1B.

eval/loss in every stage uses the same held-out split (80M tokens held out from the stage 1 mixture). Each evaluation reads the first ~1M tokens of that split. The loss is token-mean cross-entropy in nats.

Final eval/loss

H	stage 1 (step 20000)	stage 2 (step 34000)	stage 3 (step 2000)
1	4.7957 (diverged, see notes)	-	-
3	2.5483	2.2169	2.1903
6	2.3186	2.0992	2.0685
9	2.2638	2.0880	2.0500
12	2.2493	2.1044	2.0586

Notes

• H=1 diverged after step 11k in stage 1 and was not continued.
• In stage 2, H=3 and H=9 ran past step 34000 (to 35.7k and 40.9k) before they were stopped. Compare runs at step 34000 or earlier.
• Stage 2 changes the data distribution (pure dolma3), so the loss curves have a kink at the stage boundary.
• module_diag/* tags are activation norm and variance of the attention and MoE branches.