[moe] Good 10T: ablate gated norms

[dlwh]

Summary

This issue tracks the smaller good-enough 10T gated-norm ablation that unblocked later great-10T follow-up work. PR #4057 added GatedNorm, a gated_norm_rank config field, model wiring, tests, and a launch script for a roughly 1e19-FLOP baseline-vs-gated-norm comparison. The implementation work is complete and validated, but the thread still does not report experiment results, so the current conclusion is procedural rather than scientific: gated norms are now a real knob in the MoE stack, and this issue’s remaining job is to run the ablation and judge the outcome.

Helpful links

• PR [moe] Add gated norm support and ablation launch script #4057
• Great-scale follow-up: [moe] Great 10T: ablate gated norms #4044
• Good 10T gate: [moe] Good 10T gate for #3469 #4013

Description

TL;DR: Run a single roughly 1e19-FLOP comparison for gated norms in the good-enough 10T gate.

Hypothesis or Goal

We want to know whether gated norms belong in the baseline recipe once the rest of the stack is fixed.

Links

• Parent sweep: MoE Sweep: Hyperparameters, Routing, Architecture, and Optimizer #3469\n* Gate: [moe] Good 10T gate for #3469 #4013

Results

[github-actions]

🤖 I'm looking into this issue. I'll post a task list shortly with my plan and progress.

[claude]

🤖 Plan:

• Research codebase and formulate fix
• Add GatedNorm class and gated_norm_rank config to MoE model
• Wire gated norms into Block (post-RMSNorm in attn and MLP sub-blocks)
• Create ablation launch script for ~1e19 FLOP comparison
• Add test coverage for gated norm lowering
• Run pre-commit and tests
• Open PR

Gated norms were previously prototyped in a feature branch (moe_scaling_iteration_02) but never merged. The MoE model at experiments/grug/moe/model.py needs: a GatedNorm module (low-rank self-gating after RMSNorm), a gated_norm_rank field on GrugModelConfig, and wiring in Block. A launch script will define the ablation pair (baseline vs gated-norm-enabled) at ~1e19 FLOPs.

[claude]

🤖 Plan (all complete):

• Research codebase and formulate fix
• Add GatedNorm class and gated_norm_rank config to MoE model
• Wire gated norms into Block (post-RMSNorm in attn and MLP sub-blocks)
• Create ablation launch script for ~1e19 FLOP comparison
• Add test coverage for gated norm lowering
• Run pre-commit and tests
• Open PR [moe] Add gated norm support and ablation launch script #4057
• Verify CI checks pass - all green

Status: PR #4057 opened, all CI checks passing.

[dlwh-golem]

@dlwh Added a managed TL;DR block summarizing the gated-norm implementation status, linked the follow-up issue/PR, and added the tldr label.