SIGN IN SIGN UP

test(pp): add deepseek_v4 PP and EP parity tests [2/4] (#3613)

* test(pp): add deepseek_v4 PP and EP parity tests [2/4]

DeepSeek-V4 owns get_pipeline_stage_metas, so it carries hyper-connection
streams across a pipeline stage boundary rather than just hidden states, and it
is the main expert-parallel recipe. Neither contract had a GitHub-CI test.

Add two parity tests that train a shrunk DeepSeek-V4 proxy twice with the same
seed and data order and fail when the loss or gradient-norm trajectories
disagree: PP2 against a single-rank baseline, and EP2 against a 2-rank
data-parallel reference so both legs are FSDP-wrapped identically.

The proxy generates its own synthetic token sequences, so this stages no
tokenizer or dataset in TEST_DATA_DIR or HF_CACHE.

Raise the gemma4 PP2 test's global_batch_size to 4 so it accumulates two
micro-batches, and assert positively that stage shapes were precomputed --
the pre-existing negative grep cannot catch the precompute being skipped
outright, because then the fallback log line is not emitted either.

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

* test(pp): match the recipe's expert backend in the deepseek proxy

`experts: torch_mm` is what the shipped recipe uses and it runs on the CI
runners, so the proxy no longer substitutes `torch` for it.

`attn` stays sdpa because tilelang JIT-compiles its kernels on every run, and
`dispatcher` stays torch because hybridep is not installed in the CI container.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
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
3d647a29e51cf624f252935ee53d609411ef2232
Parent: b3375b8
Committed by GitHub <noreply@github.com> on 8/24/2026, 5:58:48 PM