[levanter] Make Snowball training deterministic and bounded

[yonromai]

This draft carries the three Levanter changes required by MarinSkyRL PR #578:

• Checkpoint each Snowball scan layer during reverse mode. The target 67B backward estimate fell from 292.59 GiB to 59.49 GiB per H100.
• Replace the local MoE scatter-add output combine with a reverse-position gather and ordered top-k sum. The scatter path used collision-prone GPU atomics; the new path is deterministic and has dense-reference value and gradient regressions.
• Initialize accumulated MIN/MAX metrics with +inf/-inf rather than zero so PPO ratio extrema remain valid across microbatches.

Numerical evidence

Before the deterministic combine, repeated unchanged-policy scoring was nondeterministic. After it, repeated scoring became bitwise exact. Standalone score-versus-training drift remained deterministic at maximum/mean 0.0089493/0.00014377; backend, attention, accumulation, and traced-score A/Bs did not remove it. The paired MarinSkyRL path therefore runs independent old-policy scoring through the same compiled accumulated gradient/optimizer program in no-commit mode. Its declared gate requires exact zero score drift, exact unit ratios, and zero clipping.

Representative eight-H100 gate /romain/snowball-numerical-exact-mean-01a0a1be passed that rule over a 4,096-token, two-accumulation-step batch. The independent native PyTorch Grug oracle remained within 2.384e-7 for token log probabilities, 4.470e-8 for loss, 8.941e-8 for gradients, and 1.192e-7 for the first AdamW update.

Corrected target run

Paired job /romain/snowball-67b-exact-two-cycle-r4-01a0a1be used this PR at 779cd403521e02d1615d8c49290cd8de63efdde0, MarinSkyRL 34056d565ae2e7a7b79ef9c28b4604aaa8613e53, the pinned 67B-A2B model, 32 learner H100s over four hosts, and a separate eight-H100 TP1/DP8/EP8 vLLM host. All five tasks exited 0 without failures, retries, or preemptions.

Both synchronous GRPO/AdamW updates passed exact zero score-difference, exact unit-ratio, and zero-clipping checks. They consumed 445,242 and 717,953 valid tokens; synchronized update bodies took 26.848 and 26.348 seconds; parameter-probe L2 deltas were 1.1653e-4 and 1.0050e-4. Policy versions 1 and 2 each verified 19,968 expert projection slices across all eight inference workers. A single final native checkpoint committed at global_step_2.

Closure job /romain/snowball-67b-final-policy-generation-01a0a1be loaded that target-size checkpoint in 142.12 seconds, republished version 2 in 449.776 seconds with all 19,968 slices verified, skipped training, and generated a correct retained RLVR-math response. This completes generation after both newly installed policies.

Focused validation:

• uv run --frozen pytest -q lib/levanter/tests/test_grad_accum.py -k "extrema_identities or accumulate_gradients_sharded": 7 passed.
• Dense-reference deterministic MoE value/gradient regression passed in the Marin PR matrix.
• GitHub Actions run 34907576545 and the rest of the PR matrix are green at 779cd403.

The two-update run is feasibility evidence. Learning quality, sustained throughput, and matched backend benchmarks remain out of scope.