Test MoE Arch at 1e21 and 1e22 Flop Scales

[ClassicLarry]

Summary

This issue is the tracking thread for pushing Marin's current MoE recipe from the earlier smaller-scale work in #2167 up to roughly 1e21 and 1e22 non-embedding FLOPs, with the immediate question being whether routing and loss stay stable at those scales. As of March 29, 2026, the 1e22 run is live on v4-512 and the thread records that the team is using quantile balancing (QB) instead of older auxiliary-loss-based load balancing; the current expectation is not that the recipe is fully tuned, but that this run will show whether the larger recipe is stable and how it compares against prior dense Delphi and MoE baselines. The latest thread context also notes a provisional predicted paloma/macro_loss of 2.3887 for the 1e22 run, while batch-size scheduling remains static for now and is called out as future tuning work rather than part of this experiment's core claim.

Helpful links

• Earlier MoE scaling recipe issue
• 1e22 run
• Comparison report for dense Delphi and prior MoE runs
• Comment with predicted loss and QB explanation
• Launch file for the 1e22 run and isoflop sweeps

Description

Test if the current scaling recipe in #2167 scales to 1e21 and 1e22 Flop scales without router or general instabilities.

Hypothesis or Goal

Will the loss curve or routing destabilize? Hypothesizing that similar to 2048 width runs at smaller scales, the initial routing on layer zero will look choppy during lr warmup, then stabilize.

Links

1e21 Run: 14B total params, 2B active, 75B tokens.
https://wandb.ai/marin-community/dial_moe/runs/moe-d2304-1e21?nw=nwuserlarrydial

1e22 Run: 35B total params, 5B active, 326B tokens.
https://wandb.ai/marin-community/dial_moe/runs/moe-v7-1e22-d3200?nw=nwuserlarrydial

Results

Pending

[ClassicLarry]

🤖 1e22 Plan

Parameter	Value
Hidden dim	3,200
Layers	32
Heads	25
Experts	64 (K=4, all MoE layers)
Active params	5.1B
Total params	34.6B
Tokens	326B
Batch size	1,024
Steps	~77,724
EP	4
TPU	v4-512
Capacity factor	1.0

Optimizer (CompletedAdamH heuristic):

Parameter	Value
AdamH LR	0.00408
Adam LR	0.00087
Beta1	0.9062
Beta2	0.9685
Epsilon	2.7e-15
Decay	0.2 (20% of steps)
Warmup	0.1 (10% of steps)
LR schedule	linear
Max grad norm	1.0
Z-loss weight	1e-4

Architecture features: QB routing, sigmoid combine weights, GatedNorm (rank 128), XSA, QK norm, QK mult=1.3, sliding window=4096, no dense layers, no load-balancing loss (router z-loss only).

Non-embedding FLOPs: ~1.0e22 (excludes lm_head, which is 7.4% of total at this width).

[ClassicLarry]

Setting Capacity Factor to 1 is based on experiment here.

Some sanity checks against https://arxiv.org/html/2507.17702.
They recommend:

• D = 5.2232*C^0.4905 = 323B tokens
• M = 0.1915*C^0.5095 = 31B total parameters.
• LR = 1.1576*C^-0.1529 = 0.0003 learning rate.
• Batch Size = 0.0694*C^0.3644 = 16M.
  -beta1=0.9, beta2=0.95.

Our batch size is 4x smaller (4k seq*1024), which we compensate for with higher beta2. Data and model size roughly match, which comes out to 64 chinchilla with respect to active params.

Main difference is in the learning rate- our embeds have almost 3x the lr for AdamW, and our AdamH learning rate looks much higher, despite 4x smaller batch size. Can't compare directly since we are on AdamH instead of AdamW.

Will likely move to linear decay schedule and add additional lr tuning in future runs. Main objective of this run is to see if we can get a stable loss curve and routing with clear gains over prior 1e22 run. Recipe looked stable at 1e20 scale in experiment linked above.

Running on v4-512, ballparking 10 days to complete. May need to update batch size or EP if model doesn't fit.

[ClassicLarry]

Run started: https://wandb.ai/marin-community/dial_moe/runs/moe-v7-1e22-d3200?nw=nwuserlarrydial

488k tokens/s, 23MFU -> estimating 7.7 days.

[ClassicLarry]

🤖 Launch file for the 1e22 run and isoflop sweeps: experiments/grug/moe_mar26/launch.py

[ClassicLarry]

Report to compare to dense Delphi runs and prior MoE runs: https://wandb.ai/marin-community/marin/reports/Large-Scale-MoE-and-Dense-Runs--VmlldzoxNjM1MzQ3NA

[hammer]

Do we have a predicted loss for this run that we are tracking to see how well we are calling our shots with MoEs? Something like what Will posted for Delphi...

[hammer]

Also I see "no load-balancing loss". I assume, possibly incorrectly, that refers to the Auxiliary-Loss-Free Load Balancing Strategy from DeepSeek? Are we not adding this loss just because this MoE is small and we don't expect router instability issues at this scale? I noticed in Arcee Trinity Large that they introduced some smoothing on the expert bias updates (they call it SMEBU) to fix instability issues and assume we'll need to do something similar at larger scales. Sorry if it's a dumb question or has been discussed before, I searched Discord and GitHub for "Trinity" and "SMEBU" and "Arcee" and couldn't find any previous discussion on it...

[ClassicLarry]

Do we have a predicted loss for this run

Here would be our best estimate: 2.3887 paloma/macro_loss.

However, this run is not preceded by weeks of isoflops on the same architecture, so this prediction is made from much fewer data points with lower confidence. (We improved the optimizer last week) One of the goals for early April is to get better at MoE performance prediction. We decided it would be worthwhile to kick off this run now instead of waiting 2 weeks, as the compute is readily available and we can learn a lot from how it performs.

Also I see "no load-balancing loss".

3 terms to define here:

1. Load-balancing-loss. This was the standard approach before Deepseek, and what we initially tested. Its an auxiliary loss term, where total loss = prediction_loss + coef * auxillary_loss. Every parameter in the model gets pushed to minimize the auxillary_loss, with strength proportional to the coefficient. Finding the optimal coefficient for each scale can be tricky.
2. Auxillary Loss Free Load Balancing. This was introduced by Deepseek, and what we tested next. Instead of forcing every parameter to focus on good balancing, it simply gives a bonus to the routing scores of the under-loaded experts. Arcee found at large scales this was insufficient for their model, so they added SMEBU. Nemotron 3 Super used the Deepseek approach, with a minor load balancing loss in the background.
3. Quantile Balancing (QB). This is like SMEBU on steroids, and what we are now using, thanks to some earlier work for Pranshu and optimization from Will. Original Deepseek approach makes the expert bonus on each step a fixed number, like 0.01. SMEBU lets this number shrink when its oscillating. QB directly solves for the bonus update that would optimize routing, based on all routing scores from the last step. Details. So far, this seems to be working extremely well.

[hammer]

Thanks Larry! I had followed the transition from 1 to 2, pulling expert balancing out of the loss function. I had not followed the transition from 2 to 3. I am not surprised to learn you all already have an improvement on SMEBU, I was just having a hard time tracking it down, and now I know to search on "Quantile Balancing"/"QB".

I wonder if it would be possible to use agents to generate a paper-level description of the architecture, optimizer, initialization, and other parameters of the recipe we use for each of our larger runs? It seems like there is enough structure to the recipe that it may be feasible.

[hammer]

Okay sorry to be annoying but I'm also curious to know if we're using any of the latest batch size scheduling ideas from Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Laws. My smooth brain assessment from reading the code is that we are using static batch sizes and I'm curious again if that's just because it's a smaller training run where we don't expect instabilities to be an issue or if I am once again just missing stuff that you all know way better than me...

[ClassicLarry]

Okay sorry to be annoying but I'm also curious to know if we're using any of the latest batch size scheduling ideas from Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Laws. My smooth brain assessment from reading the code is that we are using static batch sizes and I'm curious again if that's just because it's a smaller training run where we don't expect instabilities to be an issue or if I am once again just missing stuff that you all know way better than me...

We haven't done much work yet on batch size optimization. It is static, and we could likely get better MFU with a bigger batch size, and better performance from scaling up the batch size over the run. We will definitely get to tuning this in April, but it has been lower priority than some other items. I haven't read that paper and will check it out later, along with some other papers in the space. At this stage I bet Will, David, and others know much more about batch sizing than me.

I wonder if it would be possible to use agents to generate a paper-level description of the architecture, optimizer, initialization, and other parameters of the recipe we use for each of our larger runs? It seems like there is enough structure to the recipe that it may be feasible.

This sounds like a good idea and is very feasible. We can have something that does this automatically, ill have this generated next week for this run. This also may be a good point to do a more general blog of MoE progress to date.