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