Equivalence tests: reconstruct values, gradients, and the objective
Validate partitioned computation against a dense reference with deliberate shard boundaries, uneven tokens, masks, and numerical tolerances.
On this page
The question for today
What you’ll understand
- Compare forward output, input gradients, parameter gradients, and one optimizer step.
- Normalize a global valid-token mean correctly across uneven shards.
- Distinguish acceptable floating-point differences from a missing reduction or wrong layout.
A partitioned model should compute the intended logical model, not merely avoid a communication error. A strong test reconstructs the global outputs and gradients and compares them with a small dense reference under the same objective. It also verifies ownership and loss scaling. Communication changes floating-point reduction order, so exact bit identity is often too strict, while a loose tolerance can hide a missing sum. Design examples that make semantic errors large and distinguishable.
Start with a deterministic small reference
Use fixed distinguishable weights, nonzero biases, and inputs that cross each shard boundary. Disable dropout initially, use a supported even partition, and compare a simple scalar loss. Reconstruct column shards by concatenation and row-partial outputs by summation according to the declared layout. Then compare input gradients, every parameter shard in its global position, shared normalization gradients, and tied embedding contributions. A correct forward output does not prove the backward mappings are correct.
After basic equivalence, compare one optimizer update using identical optimizer state and loss scaling. An accidental extra averaging factor may barely affect a forward-only test but move weights differently. Add activation recomputation, mixed precision, and the intended randomness policy incrementally. A same seed on each rank is not necessarily the same global dropout mask. Record versions, dtype, group definitions, parameter initialization, and the exact layout graph so a discrepancy can be attributed to a changed assumption rather than guessed away.
Unequal valid-token counts change the loss denominator
For a global mean token loss, sum losses over valid tokens and divide by the total valid count. Equally averaging local means is wrong when ranks hold different numbers of valid tokens, including padding and ignored labels. The framework may average or sum gradients in different groups, so account for those reductions once when scaling the loss. TP participants cooperating on the same examples must not be treated as additional independent examples. Check the gradient normalization algebra as well as the printed scalar loss.
Packed sequences and context-sharded masks deserve separate tests. A query must see exactly the same globally permitted keys as the dense model, including causal boundaries and example boundaries. Vocabulary padding must not create extra valid classes. For unsupported uneven DTensor operations, a padded execution path needs a mask and offset mapping that restores the intended logical input. Test all-masked chunks using the implementation’s explicit neutral handling rather than relying on undefined softmax divisions.
Use adversarial values before performance traces
Include cancellation in row-partial sums, a bias that exposes duplication, a vocabulary maximum and target on different ranks, and shared KV gradients under GQA. Exercise both a sequence-local norm and context attention requiring remote keys. Assert global shapes, local offsets, and collective groups, and keep test timeouts that report rank progress when a collective hangs. A performance trace is useful only after the partitioned values and training step match the specified model within defensible tolerance.
The optimization lab gives a small numerical sensitivity illustration. Reset, set η=0.5 and x₀=2, predict x₁=1, then Step. Reset and set η=2.5 and x₀=2, predict x₁=−3, then Step; further steps leave the plotting range as the unstable magnitude grows. The model is one exact quadratic descent recurrence, not distributed training or an equivalence tester. It shows why a changed update rule can accumulate meaningful differences, while actual partitioned tests must compare the model’s real gradients and optimizer state.
Work it through
Averaging two local means changes an uneven objective
Rank 0 owns two valid token losses[1,3]. Rank 1 owns one valid loss[8] plus ignored padding. Compute the intended global mean and the incorrect equal average of local means.
Rank 0’s sum is 4 over two valid tokens; rank 1’s is 8 over one. The global sum is 12 and valid count 3, so the intended mean is 4.
Local means are 2 and 8. Their equal average is 5, giving the single-token shard the same weight as the two-token shard and changing the objective.
Weight the local means by their valid counts: (2×2+1×8)/(2+1)=4. For backward, use equivalent scaling that matches the actual gradient sum/average operations; printing the right scalar alone does not prove parameter gradients use that denominator.
The correct denominator follows valid examples, not rank count. Padding and cooperating TP ranks must not silently alter it.
The living notebook
Explore the model
Interactive lab · optimization
A changed update rule can accumulate
Reset, set η=0.5 and x₀=2, predict x₁=1, then Step. Reset and set η=2.5 and x₀=2, predict x₁=−3, then Step. This quadratic recurrence is not tensor-parallel training or a gradient-equivalence checker; the unstable case reports out-of-range values as it grows.
- Position x0
- 2
- Objective x²/2
- 2
- Iteration behavior
- Stable step
Gradient descent on f(x) = x²/2: xₙ = (1 − η)ⁿx₀. Playback advances two discrete iterations per second. Values outside |x| ≤ 4 are explicitly reported. Parameter changes pause playback and recompute the model at the current elapsed time. Reset restores both parameters and returns time to zero. Playback starts paused. Adjust either parameter, then choose Play or Step.
Look forCompare the intended update and objective before interpreting later model differences as harmless numerical noise.
Pause & reason
Check your understanding
Your turn
Try it yourself
Rank 0 has three valid losses[2,4,6] and rank 1 one valid loss[8]. Compute the global mean and equal average of local means. Which objective should an equivalence test use?
Take a moment to make a prediction or work it out on paper. Open a hint when you need one.
Give me a hint
Combine sums and valid counts, then compare with rank-weighted averaging.
Compare with the solution
The global sum is 20 over four tokens, so mean 5. Local means are 4 and 8; their equal average is 6. Use the intended valid-token mean 5 and verify gradients and the optimizer step follow that same objective, accounting exactly once for the framework’s gradient reductions.
Keep the idea
Take it with you
- Forward shape agreement does not establish backward or optimizer equivalence.
- Global valid-token normalization must survive padding and uneven ownership.
- Use targeted semantic counterexamples and justified numerical tolerances before benchmarking.
Read further
Original sources and deeper explanations to follow at your own pace.
Optional progress, saved in this browser when available. Revisit any lesson at any time.
Lesson not marked complete.