looped-moe-329m-ablation-tensorboard

Looped MoE Transformer (329M) ablation logs

TensorBoard scalar logs and per-loop activation diagnostics for seven ablations of a 329M looped MoE transformer. Each run changes exactly one switch of the base recipe.

Base recipe

• 10 layers, hidden 512, 8 heads / 4 KV heads
• MoE: 30 routed + 2 shared experts, top-8 (shared included), expert width 512
• The 10-layer block is applied H times
• lr 2e-4, WSD schedule: 10% warmup, 10% linear decay to 0.1x
• Global batch 1024 sequences x 1024 tokens, micro-batch 16, AdamW (0.9, 0.95), weight decay 0.1, gradient clip 1.0, bf16
• Data: FineWeb-Edu (10B-token sample), SmolLM tokenizer
• Loop residual scale gamma = lambda / (H sqrt(L)) and embedding re-injection g = lambda / (t sqrt(L)), with lambda = 0.5 frozen
• RMSNorm (affine) on the MoE branch output
• Segmented backprop: every 3 loops the model computes the loss, backpropagates, takes an optimizer step, and detaches the state (bp_segment_len=3)

Runs

bp_segment_len=0 runs all H loops as one graph and backpropagates once per step. It also removes the intermediate-loop losses and the extra optimizer steps of the segmented scheme. With H=3 the segmented scheme has a single segment, so nobpk_H3 is mathematically the same as the base recipe at H=3.

H=9 and H=12 without segmentation ran out of memory at micro-batch 16 on 96 GB cards, so they use micro-batch 8 (global batch unchanged). MoE capacity is computed per micro-batch, so this is a second difference from the base recipe for these two runs.

The 5000-step runs keep the 10% / 10% WSD ratios, so each of them finishes its own decay at step 5000.

Files

• events.out.tfevents.*: scalars. train/grad_norm is the global gradient norm before clipping. module_diag/*/act_norm_mean and act_var_mean are activation norm and variance of the attention and MoE branches. module_diag/*/grad_norm_mean is always 0: per-module gradient norms were not collected under FSDP2, so ignore it.
• module_diag.csv: every 50 steps, per layer and per loop (cHxx), the branch activation norm, variance, effective rank, top singular value and mean cosine. The empty per-module gradient-norm columns were removed.
• In nobpk_H9_mb8 and nobpk_H12_mb8, the training script restarted from a snapshot after the non-finite gradient (from step 1000 and step 8000), so those steps appear twice. The first pass is the result. The H=9 restart hit a second non-finite gradient at step 1852.

Notes

• No baseline runs are included. The base recipe was trained on a different cluster. Compare against it after measuring the platform offset with nobpk_H3, which is equivalent to the base recipe.
• Evaluation uses a held-out part of the same FineWeb-Edu pack. The loss is token-mean cross-entropy in nats.