Agent MoE Experiment: AdamH gradient normalization

[WhenWen]

Description

Test a Grug MoE AdamH variant that normalizes each module's gradients to RMS 1 before AdamH moment updates. The code path is local to experiments/grug/moe/ and compares against the compute-optimal MoE baselines in experiments/grug/moe/README.md.

Exact initiating prompt: "Read experiments/grug/moe/README.md on how to iterate MoE, implement a variant of AdamH optimizer that perform gradient normalization (scale gradient of each module to RMSNorm 1) and test it against the previous method"

TL;DR

Gate 1 is complete. The variant improved the d768 point but missed the d512 point, so it does not pass gate 1 and should not advance to gate 2 under experiments/grug/moe/agent.md.

Hypothesis or Goal

Module-wise gradient RMS normalization reduces optimizer scale mismatch across attention, shared expert, and routed expert modules, improving effective speedup without changing AdamH's projected parameter update rule.

Links

• Logbook: .agents/logbooks/moe-adamh-grad-norm.md
• Launch: experiments/grug/moe/launch_adamh_grad_norm.py
• PR: [grug] Add MoE AdamH gradient normalization #5181
• Gate 1 Iris parent job: /kaiyue/iris-run-job-20260425-184632
• Data browser: https://marin.community/data-browser/experiment?path=gs%3A//marin-us-east5/experiments/launch_adamh_grad_norm-7614bf.json

Results

Scale	Baseline loss	Variant loss	Loss delta	Baseline tok/s	Variant tok/s	Tok/s delta	Effective speedup
d512	3.8104	3.815110	+0.004710	405,630	406,983	+0.333%	0.980893
d768	3.4339	3.429193	-0.004707	273,532	274,218	+0.251%	1.030269

Both W&B runs reached finished, and both Iris child jobs reached JOB_STATE_SUCCEEDED.

Decision Log

• 2026-04-25: submitted gate 1 on Iris for d512 and d768.
• 2026-04-25: d512 finished with effective speedup 0.980893, below the required threshold.
• 2026-04-25: d768 finished with effective speedup 1.030269, but gate 1 requires both small-scale points to exceed 1.0.

Conclusion

Completed negative result. AdamH module gradient RMS normalization does not pass gate 1, so no gate 2 run is launched for this variant.

[WhenWen]

🤖 Update: gate1 submitted and running

• Change: submitted gate1 on Iris for the gradient-normalized AdamH variant at d512/2.19e17 and d768/1.70e18.
• Result delta: no final metrics yet. The first TPU allocation hit a device-busy bad-node signature, then Iris retried automatically; both child training jobs are now JOB_STATE_RUNNING on v5p-8.
• Confidence: exploratory, in progress.
• Links:
  • PR: [grug] Add MoE AdamH gradient normalization #5181
  • Commit: d97fdc3
  • Iris parent job: /kaiyue/iris-run-job-20260425-184632
  • Data browser: https://marin.community/data-browser/experiment?path=gs%3A//marin-us-east5/experiments/launch_adamh_grad_norm-7614bf.json
  • W&B d512: https://wandb.ai/understanding-sam/marin_moe/runs/moe-adamh-grad-norm-d512-2p19e17
  • W&B d768: https://wandb.ai/understanding-sam/marin_moe/runs/moe-adamh-grad-norm-d768-1p70e18
  • Logbook: https://github.com/marin-community/marin/tree/research/moe-adamh-grad-norm/.agents/logbooks/moe-adamh-grad-norm.md

[WhenWen]

🤖 Update: d512 gate1 result

• Change: pulled the final W&B summary for the d512 gate1 run.
• Result delta: d512 finished with eval/paloma/macro_loss=3.815110 and throughput/tokens_per_second=406,982.7. The README baseline is 3.8104 at 405,630 tok/s, so the effective speedup is 0.980893.
• Confidence: exploratory, single run.
• Decision: d512 does not clear gate1; gate2 should not launch under the current criteria. I am keeping the already-submitted d768 run alive to finish the full gate1 evidence table.
• Links:
  • W&B d512: https://wandb.ai/understanding-sam/marin_moe/runs/moe-adamh-grad-norm-d512-2p19e17
  • Iris parent job: /kaiyue/iris-run-job-20260425-184632
  • PR: [grug] Add MoE AdamH gradient normalization #5181

[WhenWen]

🤖 Final update: MOE-AGNH-001 gate 1 complete

Result: the AdamH module-gradient-normalization variant does not pass gate 1.

Scale	Baseline loss	Variant loss	Variant tok/s	Effective speedup	Decision
d512	3.8104	3.815110	406,983	0.980893	fail
d768	3.4339	3.429193	274,218	1.030269	pass

Gate 1 requires effective speedup > 1 at both d512 and d768. Because d512 is below threshold, I did not launch gate 2.

Confidence: exploratory. This is one gate-1 run at each required small scale, both with W&B state finished and Iris JOB_STATE_SUCCEEDED.

Links:

• PR: [grug] Add MoE AdamH gradient normalization #5181
• Logbook at closeout commit: https://github.com/marin-community/marin/blob/2356b2be347628a33c15eb990ca67305818bbda3/.agents/logbooks/moe-adamh-grad-norm.md
• d512 W&B: https://wandb.ai/understanding-sam/marin_moe/runs/moe-adamh-grad-norm-d512-2p19e17
• d768 W&B: https://wandb.ai/understanding-sam/marin_moe/runs/moe-adamh-grad-norm-d768-1p70e18
• Data browser: https://marin.community/data-browser/experiment?path=gs%3A//marin-us-east5/experiments/launch_adamh_grad_norm-7614bf.json

Next steps:

1. Do not promote this exact normalization rule.
2. If revisiting the idea, test a softer normalization target or apply it only to selected module families after diagnosing the d512 regression.
3. Keep PR [grug] Add MoE AdamH gradient normalization #5181 open only if maintainers want the variant machinery preserved for further optimizer ablations.