Equivalence tests: compare the update state, not only the loss

Start with a tiny deterministic model and a known global example or token multiset. Use a separable objective so that splitting inputs does not change model semantics; rank-local batch normalization, randomness, and stateful layers need separate intended contracts. Compute a high-precision unsharded reference and preserve masks, reduction mode, eligible counts, parameter identities, and optimizer hyperparameter groups. Reconstruct sharded tensors in canonical logical order before comparison. Equal norms cannot detect a permutation that attaches correct numerical moments to the wrong coordinates.

Compare at several boundaries: local loss-sum accounting, normalized aggregate gradients, post-clipping gradients, updated parameters, and optimizer state after the step. Include zero gradients, unequal token counts, uneven parameter shapes, and a final partial accumulation window. A reference that repeats the implementation’s mistaken denominator will agree for the wrong reason, so derive expected values independently for at least a few scalar fixtures. Record which differences are algorithmic changes and which are supposed to preserve the same update.

The first optimizer step can hide a lost momentum or adaptive-moment checkpoint because the initial buffers are often zero. Test multiple steps and restore a fixture with known nonzero state. Check the step counter, schedule position, weight-decay mode, and parameter-group settings as well as tensor values. Advancing an optimizer or scheduler once per microbatch instead of once per accumulation window can change every later step even if the initial gradient matches. A state-sharding test must verify each owner’s identity and state, not just that the final model contains finite numbers.

Mixed precision needs explicit expected ordering: accumulate under a consistent scale, synchronize according to the wrapper’s contract, unscale at the proper boundary, and clip the intended global gradient. Overflow and skipped-step behavior must leave ranks compatible. Floating-point reductions are not associative, so a different legal collective or bucket order may introduce small drift without an objective error. Choose tolerances based on dtype and the operation, then investigate structured discrepancies such as an exact world-size factor, a swapped parameter block, or a step counter that differs by the accumulation count.

Change the number of workers while preserving the intended global batch and objective, or repartition the same canonical gradient among owners, then compare the next state. Such transformations should preserve the stated mathematics only when their assumptions remain true. A local-batch-dependent model or changed sample duplication invalidates that expectation. Also test a save-and-resume boundary against uninterrupted execution, including optimizer and input progress state. Keep the result scoped to the evaluated implementation and configurations, and retain a failing fixture rather than weakening tolerance until every difference disappears.

For a simple update-boundary check, Reset the optimization lab, explicitly set step size 0.1 and initial position two, predict two iterates, then Step twice. Compare a fresh Reset with both step size 0.2 and initial position two set before one Step. Results are 1.8 then 1.62 versus 1.6. The fixed quadratic is a small independent arithmetic oracle, not a momentum, Adam, DDP, or sharding implementation. It demonstrates why distinct update schedules should not be labeled equivalent merely because their total nominal step sizes match.