SIGN IN SIGN UP

test(pp): add gemma4 parallelism parity tests [1/4] (#3529)

* test(pp): add gemma4 parallelism parity tests

Existing PP tests only assert that training does not crash, and all of them
run stock HuggingFace Mixtral at tp_size=1, so no VLM and no custom model is
covered under pipeline parallelism.

Add parity tests that train a small random Gemma4 proxy twice with the same
seed and data -- once on one rank, once with PP=2 or TP=2 -- and fail when the
loss or gradient-norm trajectories disagree. The proxy is built at test time
from a shrunk copy of the 31B config, so nothing has to be staged in
TEST_DATA_DIR beyond the Gemma4 processor already used by the joint-drafter
test.

Also guard the three existing PP tests against a silent fallback from static
stage metadata to dynamic shape inference, and add CPU unit tests covering
stage metadata at pipeline depths 2, 4 and 8.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Abhishree <abhishreetm@gmail.com>

* test(pp): trim guard comments to one line

Move the bug and mechanism write-ups out of the code comments and into the PR
description. Each guard now states in one line what it protects against, and
every commit reference names its PR.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Abhishree <abhishreetm@gmail.com>

* test(pp): drop preserve-grads tests that belong to the fix PR

The gradient-preservation tests import a helper that only exists in the
pipeline fix PR, so collecting them here fails the whole unit-test shard.
They ship with the fix instead.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Abhishree <abhishreetm@gmail.com>

* test(pp): add per-parameter grad-accumulation parity

Compares accumulated gradients against the sum of the same windows run
separately, so a discarded micro-batch fails numerically rather than silently.

Signed-off-by: Abhishree <abhishreetm@gmail.com>

---------

Signed-off-by: Abhishree <abhishreetm@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
A
Abhishree Thittenamane committed
bd58fd66f7e6e496857e1df9097de2062adf5346
Parent: 084cbe3
Committed by GitHub <noreply@github.com> on 8/20/2026, 7:36:47 PM