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

[athitten]

What

File	What it does
L2_Parallelism_DeepSeekV4_PP2_Parity.sh	Trains a tiny DeepSeek-V4 twice — once on 1 GPU, once split over 2 GPUs with pipeline parallelism — and fails if the loss or gradient-norm curves drift apart.
L2_Parallelism_DeepSeekV4_EP2_Parity.sh	Same for expert parallelism: 2 GPUs both times, once with all 8 experts on every rank, once with 4 per rank.
deepseek_v4_proxy.yaml	A shrunk deepseek_v4_flash_hellaswag.yaml (6 layers, ~20M params) keeping every field the parallelism code branches on.

Also raises the gemma4 PP2 test's global_batch_size 2 → 4 so it accumulates two micro-batches.

Why

#2983 broke every pipeline-parallel recipe while GitHub CI stayed green, because no test ran a real recipe topology twice and compared the numbers. #3529 covered gemma4. DeepSeek-V4 is the next one that matters: it owns get_pipeline_stage_metas, so it carries hyper-connection streams across the stage boundary rather than just hidden states, and it is the main expert-parallel recipe.

Loss on its own is a weak detector; gradient norm is the sensitive one. Both tests assert on it, and the comparison script fails if a log has no grad_norm.

No new CI assets

The proxy generates its own fixed-length synthetic token sequences (llm.mock.build_unpacked_dataset, std_len: 0) — no tokenizer, no dataset to stage. Same pattern as baichuan_2_7b_mock_fp8.yaml.

Three things the proxy has to copy from the shipped recipe

Each was found by running the test, not by reading it:

1. patch_inner_model / patch_causal_lm_model: false — as in every shipped deepseek recipe. At their true default, patch_hf_model_for_pp replaces DeepSeek-V4's own PP-aware forward with the generic CausalLM one, which skips the [B,S,dim] → [B,S,hc_mult,dim] expansion and dies in DeepseekV4HyperConnection.
2. 6 layers, not 4 — with 4 over pp_size=2 a hash-routing layer lands on stage 1, gets input_ids=None, and takes its score-based-topk fallback, routing differently from the baseline. The shipped 43-layer/pp4 recipe keeps all 3 hash layers on stage 0.
3. No fake_balanced_gate — model.py:181 only installs the hash gate when it is off, so it disabled the routing num_hash_layers: 3 is there to exercise.

Measured

Test	Bounds	Worst observed
PP2	loss 0.10, grad-norm rtol 0.20	0.064, 0.19%
EP2	loss 0.02, grad-norm rtol 0.05	0.0006, 0.064%

PP's loss bound is 0.10 because at a loss of ~10.9 in bf16 the representable resolution is already ~0.04. EP is 50x tighter because both its legs are 2-rank and FSDP-wrapped identically — which is also the control proving the PP gap is the wrapped/unwrapped asymmetry, not pipeline parallelism.

CI cost: L2_Parallelism goes ~4m42s → ~9min against a 40-minute timeout, and finishes ~8 minutes before the critical-path job. No change to total pipeline wall-time.

Backends

experts: torch_mm matches the shipped recipe. Two differ:

• attn is sdpa, not tilelang — tilelang JIT-compiles its kernels on every run (~1min for this suite), and the parity numbers are identical either way.
• dispatcher is torch, not hybridep — hybridep is not installed in the CI container, and at world_size == 1 the deepep/hybridep path falls back to GroupedExperts (layers.py:771), so the PP test's single-rank baseline would run a different expert implementation than its pp2 leg.

Neither changes what EP shards: apply_ep → ExpertParallel distributes the expert weights on the expert axis in every case.

Also

Both PP scripts now assert positively that "Precomputed pipeline stage shapes" appears. The pre-existing negative grep cannot catch _precompute_stage_shapes being skipped outright, since the fallback line disappears with it — the exact function #2983 broke.

Known gap: the EP test catches EP changing the numbers, but not EP silently never being applied — that would make both legs identical and pass. PP closes the equivalent hole via the log line; apply_ep has none to grep. Documented in the script header.

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[athitten]

/ok to test f6f9ceb