Add MoeAdamHHeuristic, drop dense layers, fix align_kv_heads sharding

[ClassicLarry]

Summary

• Port CompletedAdamHHeuristic from the moe_isoflop_apr_2026 branch into experiments/grug/moe/heuristic.py as MoeAdamHHeuristic. Adds compute_flops_per_token, compute_tokens_and_batch, and build_from_heuristic so launch.py can derive (model, optimizer, batch, steps) from a compute budget + hidden_dim.
• Remove the initial-dense-layer path from experiments/grug/moe/model.py (_NUM_DENSE_LAYERS, dense_intermediate_dim, Block.dense_mlp, dense_only init branch) to match the isoflop MoE architecture. Default EP capacity factor back to 1.0. Restore the aligned_v reshard the isoflop branch had.
• Rewire experiments/grug/moe/launch.py to use build_from_heuristic for the baseline ExecutorStep instead of hardcoding GrugModelConfig / GrugMoeAdamHConfig. Manual specification is still supported by passing configs directly.
• Add experiments/grug/moe/README.md covering architecture, the scaling heuristic, v16 isoflop best runs per compute budget, projections (Paloma macro L∞ pinned at 1.6), and promotion criteria.
• Fix lib/levanter/src/levanter/grug/attention.py::align_kv_heads: replace jnp.repeat with a reshape+broadcast pattern. Resolves the ValueError: Please pass sharding to jnp.repeat via out_sharding parameter crash under abstract-mesh training contexts.

Test plan

• Verified loss curves match the published isoflop-moe-v16-1e+18-d1024 run to within ±0.005 at matching steps on a v5p-8 iris run (4_10_test_moe).
• build_from_heuristic(budget=1e18, hidden_dim=1024) reproduces the isoflop d1024 hyperparameters exactly (lr=0.01, adam_lr=0.002308, beta2=0.999, bs=32, steps=5622).
• Pre-commit linter + type check (./infra/pre-commit.py) passing.
• Reviewer: confirm the align_kv_heads reshape is equivalent to the old jnp.repeat under GQA.

🤖 Generated with Claude Code

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: c333125126

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Propagate seq_len through heuristic model construction

build_from_heuristic accepts a seq_len override, but it builds the model with h.build_model_config(hidden_dim) and then computes FLOPs from that config, so FLOPs are still based on the hardcoded 4096 sequence length. If a caller passes a non-default seq_len, the returned (batch_size, num_steps) is computed with the override while model shape/FLOP estimation remains at 4096, causing the compute budget calculation to drift and mis-size experiments.

Useful? React with 👍 / 👎.