Router precision: protect the selection margin as well as softmax

For finite logits, subtracting their maximum before exponentiation avoids positive exponential overflow in the mathematical normalization. It does not undo errors already introduced while computing or storing the logits, and tiny probabilities can still underflow. Top-k ordering can be computed from logits without exponentiating because exact softmax preserves their order. In finite arithmetic, however, rounding can collapse nearby values into ties. Define nonfinite handling, tie policy, and where precision promotion occurs rather than declaring the router stable because it subtracts a maximum.

Switch's selective-precision design uses higher precision locally in routing while keeping other work and communicated tensors at lower precision under its stated implementation. Megatron-Core 0.15.0 exposes router dtype choices with their own scope. Promoting a previously rounded low-precision logit cannot recover its missing bits; performing the relevant projection and normalization in a wider path is a different intervention. Compare these choices with matched model parameters and inputs, and inspect the generated execution path rather than assuming a configuration name proves every router operation uses that dtype.

Suppose every computed logit differs from a reference by at most epsilon. If the gap between the kth and next logit exceeds twice epsilon, no selected logit can fall below an unselected one under that bound. This is a sufficient ordering condition, not an estimate of the actual error or a quality guarantee. Near the boundary, compare selected identities explicitly. A probability difference that looks negligible can change capacity survivors or route a representation to an expert with very different parameters.

Router z-loss, introduced in ST-MoE, penalizes squared log-sum-exp under a stated coefficient and batch reduction. It controls a logit-scale-related quantity rather than directly enforcing uniform hard assignments. Adding a common constant leaves softmax probabilities unchanged but changes log-sum-exp and thus this penalty. It is distinct from entropy regularization and the hard-count/soft-mass auxiliary objective. If several regularizers are combined, log each scaled component and inspect their gradients; improved numerical behavior does not establish that the expert workload became balanced.

Test well-separated logits, exact ties, near ties, large common offsets, and the declared nonfinite cases. Compare reference logits, selected expert identities, selected weights, overflow masks, auxiliary terms, and router-gradient norms. A count histogram can miss swapped token identities. Use fixed routing in one ablation to isolate expert numerical error, then restore live routing to measure selection sensitivity. For a stability intervention, compare convergence at matched data order and token budget as well as immediate throughput and memory cost.

The network lab only evaluates a stable binary sigmoid. Reset, set x=1 and w=1, predict p≈0.622459, then Play or Step. Separately Reset, set x=1 and w=1.05, predict p≈0.634136, then Play or Step. Both control values are reachable. The instrument has no dtype conversion, top-k boundary, logit projection, or z-loss. It shows a continuous probability response; the hand example establishes how a separate finite representation can collapse ordering before such normalization occurs.