Agent MoE Experiment: may_arch + 1pct-noclip + col-norm on GN vs KV (d512/d768)

[ClassicLarry]

🤖 ## Original prompt

lets take the best d512 result, and try two variations at both d512 and d768: col norm on gated_norm, and col norm on kv projections.

Recipe

Base: 1pct-noclip (muonh-may-arch-gn-muonh-1pct-noclip-v1-d512 = 3.6427, current d512 best):

• may_arch architecture (256 experts, PKO every 4th, partial rope, last_layer_pko, no router/logit z-loss)
• 1% warmup
• max_grad_norm = None (no clip)
• token_embed → adamh_embed, output_proj → adamh
• GN → muonh, KV → muonh (matrices)
• LR scales 1.0× heuristic

Two variations:

trial	what's routed through col-norm (NS + col-norm + hyperball)
gn-colnorm	all 4 GatedNorm matrices (.w_up, .w_down)
kv-colnorm	attn.w_k and attn.w_v only

Everything else identical to the 1pct-noclip baseline. 4 runs total (2 variants × d512, d768).

Why col-norm is interesting for these

GatedNorm at d512 is (512, 128) — 4:1 aspect; NS produces 128 orthonormal cols (length 512) but row norms spread ±15% around the mean. Col-norm equalizes the rows.

K/V at d512 with GQA 4:1 is (512, 128) — exactly the same shape as GatedNorm. So K/V might benefit from the same col-norm treatment that we suspect (per the prior colnorm experiment vs may_arch-default) is mildly helpful.

Branch / launcher

• Branch: moe_muonh_may_arch_1pct_colnorm_variants (off main)
• Launcher: experiments/grug/moe/muonh_may_arch_1pct_colnorm_variants_sweep.py
• Submission target zone: us-central1-a (no --reserve).

.venv/bin/iris --config lib/iris/examples/marin.yaml job run \
  --no-wait --zone us-central1-a \
  -e WANDB_API_KEY "$WANDB_API_KEY" \
  -- python -m experiments.grug.moe.muonh_may_arch_1pct_colnorm_variants_sweep

Mask verified for both variants (smoke test before submission):

gn-colnorm:  GN→muonh_col_norm,  K/V→muonh
kv-colnorm:  GN→muonh,           K/V→muonh_col_norm,  rest unchanged

Comparison anchor at d512: 1pct-noclip baseline = 3.6427.

[ClassicLarry]

🤖 Submitted: /larry/iris-run-job-20260515-175657 on --zone us-central1-a.

4 trials launching in parallel (2 variants × d512, d768):

run-id	col-norm target
muonh-may-arch-1pct-gn-colnorm-v1-d{512,768}-*	4 GatedNorms → muonh_col_norm
muonh-may-arch-1pct-kv-colnorm-v1-d{512,768}-*	attn.w_k, attn.w_v → muonh_col_norm

Branch: moe_muonh_may_arch_1pct_colnorm_variants (commit e16c5167d).
Launcher: experiments/grug/moe/muonh_may_arch_1pct_colnorm_variants_sweep.py.

Verified mask routes via smoke test:

gn-colnorm: GN -> muonh_col_norm, K/V -> muonh, rest unchanged
kv-colnorm: GN -> muonh,          K/V -> muonh_col_norm, rest unchanged

Anchor: muonh-may-arch-gn-muonh-1pct-noclip-v1-d512 = 3.6427.
If gn-colnorm helps, expect d512 < 3.6427.
If kv-colnorm helps, the K/V rectangular treatment carries.

Wandb group: https://wandb.ai/marin-community/marin_moe/groups/muonh-may-arch-1pct-colnorm-variants-sweep

[ClassicLarry]

🤖 kv-colnorm promotion to d1024

At gate-1, kv-colnorm beat the 1pct-noclip d768 anchor (3.3005 vs 3.3040, Δ=−0.0035) and matched at d512 (3.6443 vs 3.6427, Δ=+0.0016). gn-colnorm came in slightly worse than anchor at both scales (+0.0053 at d512, +0.0009 at d768) so it is not promoted. Submitting kv-colnorm at d1024 alone.

• Branch: moe_muonh_kv_colnorm_d1024 (commit b73334ade)
• Job: /larry/iris-run-job-20260515-230419 on us-central1-a
• Single run: muonh-may-arch-1pct-kv-colnorm-v1-d1024-9.00e18
• Wandb: https://wandb.ai/marin-community/marin_moe/groups/muonh-may-arch-1pct-colnorm-variants-sweep
• ETA: ~10h on v5p-8

Launcher restructured to an explicit _RUNS tuple containing only the new d1024 entry; gate-1 outputs are content-hash cached and untouched.

	d512 macro	d768 macro	d1024 macro
1pct-noclip anchor (#5763)	3.6427	3.3040	(running)
kv-colnorm	3.6443 (+0.0016)	3.3005 (−0.0035)	pending
gn-colnorm	3.6480 (+0.0053)	3.3049 (+0.0009)	(not promoted)

Speedup of kv-colnorm at d768 vs README compute-optimal baseline: 1.91× wall-clock, 2.23× equal-tps. vs the next-best d768 recipe (gnwup-0.7×): 1.001× wall-clock — most of the gain is may_arch's, kv-colnorm adds ~1%.