Stability fixes

Misc stability fixes

Grad_norms_2

A mixture of a scalar grad norm parameter, and a mode that can be engaged by config that shuts it off entirely.

Gradient Norms Fix

Load balancing is finally fixed. This time it was an issue with me not assigning separate residual gates to separate residual streams, causing gradient explosion later on in training.

Mechanical Load balancing solution

For this to not stay load balanced math must fail.

Causal Loss

A causal load balancing loss has been introduced capable of penalizing the trajectory of decisions which lead up to a violation state occuring. This should close the loophole the model used to have that let it get in a state with no valid alternative choices that let it beat the last load balancing formula.

Router loss rebuild

We are adapting an entirely new form of loss to ensure the model stays load balanced and is load balancible indefinitely.

Router rebuild

Router has been entirely rebuilt to be impossible to reward hack. This will balance the load, or go down in flames trying.

Load balancing Fix 2

This time, scaled down logits too.

Telemetry update

Hopefully this goes through

Load balancing Fxi

The deepseek style load balancing was too slow and did not correct itself as the imbalance worsens. This update implements loss choices that correct this issue. The gshard, ce, and bce options are all available, and the ce and bce have familiar logprob properties which make them crank up the pain as load balancing worsens. Like Deepseek, we still perform load balancing on decoupled logits.