Fix the CP4 calibration hang: decide warm-ups from world-wide compile telemetry (#840)

[bradhilton]

Closes #840.

Root cause: a harness desynchronization, not a context-parallel runtime bug

The "deterministic CP4 hang" on Ellavox groups 1 and 4 came from dev/trainer_rank_landing_acceptance.py --phase cost-calibrate: the warm-up loop stopped a candidate's warm-ups when the local rank's forward was compile-free. Compile telemetry is per process, and one CP rank (rank 3, whose local shard shapes differ) recompiled on its second depth_one warm-up while ranks 0–2 did not. Ranks 0–2 moved on to the next candidate; rank 3 ran a third depth_one. The context-parallel all-to-alls then paired two different layouts and deadlocked.

Evidence (instrumented run of the unmodified harness, every torch.distributed collective logged per rank):

Cell	Ranks 0–2 at the hang	Rank 3 at the hang	First divergent collective
Ellavox g4	forward 6 = first uniform_depth_2	forward 6 = third depth_one	#778, GDN state exchange: rank 3 expects 7,216 rows from rank 2, rank 2 sends 0
Ellavox g1	forward 8 = measured no_sharing	forward 8 = automatic (one warm-up behind)	#1390, GDN state exchange, split mismatch

This also explains why a standalone repro of the same layouts on a fresh cluster never hung (identical per-rank sequences), and why the same cells at CP1/CP2 were fine. The static plan checks (attention _build_runtime_plan, GDN rank plans, plan-cache keys) were all consistent across ranks; no runtime code changes.

Fix

• Every recorded compile_statuses is now the gathered union over all ranks (_world_compile_statuses, one all_gather_object per run), and warm-up completion is a pure function of it (_warmup_complete). The tp2-public compile-free gate uses the same world-wide view.
• CPU regression test: tests/unit/test_trainer_rank_calibration_harness.py.
• Validation on 4×H200 (k8s): both cells run to completion with the fixed harness (CELL_SET=ellavox ELLAVOX_GROUPS="1 4", 8 measured rounds). The evidence rows show the mechanism directly: the warm-up attempts that used to desynchronize now record world statuses ["none", "recompile"] and all ranks take one more warm-up together.

Certificate: the two cells folded in, table unchanged

The certificate now covers all 58 manifest cells (3,849 within-cell pairs) with no exclusions. Group 1 is an odd Ellavox group (pre-registered holdout) and group 4 is held out by argument (--holdout "cp4|g4"), so the 45 training cells are the same as before and the fitter reproduces the shipped table bit for bit (same hash; the opt-in ART_COST_CERTIFICATE_REFIT=1 test passes). Metrics on all 58: 98.1% pairwise, p95 regret 2.9%, max 4.2%, no clear misses; shipped-table regret on the two new cells is 0% (g1) and 2.8% (g4, picks minimum_effective_span_495 at 911 ms over uniform_depth_2 at 887 ms).

For the record: a plain refit on all 58 cells lands on an equivalent-quality but different table (gdn_token_per_rank → 0, token_per_rank +25%, identical metrics); the coordinate search has equivalent optima there, so swapping production coefficients for no measurable gain was rejected.

Tooling

• dev/trainer_rank_collective_trace.py: runs a calibration cell with every collective traced per rank (operation, caller, splits, layout label, forward index) plus a stall watchdog that dumps all thread stacks and exits, so a deadlock surfaces in minutes instead of the 30-minute NCCL timeout.
• dev/trainer_rank_collective_diff.py: prints the layout each rank ran per forward and the first collective where ranks disagree.
• dev/trainer_rank_cost_calibration_cp4.sky.yaml: CELL_SET=ellavox runs only the groups in ELLAVOX_GROUPS.

The Ellavox corpus stays git-ignored and mounted; only timings and features were recorded.

Checks

uv run prek run --all-files green; tests/acceptance/trainer_rank_planner + tests/unit/test_trainer_rank*.py + tests/unit/test_planner_cost*.py: 582 passed, 17 skipped (including the opt-in full refit). All SkyPilot clusters torn down.

🤖 Generated with Claude Code

[bradhilton]

Codex review — no blocking findings

The root-cause analysis and fix are coherent. The divergent behavior was in the calibration harness: a rank-local compile-status decision could advance CP peers to different forced layouts. Gathering compile statuses before making the next-forward decision gives every rank the same bounded warm-up sequence; the new regression test covers the decisive rank-3-recompile case. The extra collective is confined to the dev/acceptance harness and does not affect TrainerRank's production path.

The evidence update is also internally consistent: both formerly excluded CP4 cells are complete, the explicit cp4|g4 holdout plus the existing odd-Ellavox rule preserves the original 45 training cells, all 58 cells are evaluated, and the opt-in refit reproduces the unchanged production coefficient table. The revised manifest/recipe tests and certificate metrics agree.

Verification on 4a08d15fa: 13 focused tests pass, including ART_COST_CERTIFICATE_REFIT=1; diff --check and compilation of the three dev scripts pass. GitHub quality and TrainerRank GPU validation are green. I found no production runtime change or unresolved correctness issue in this PR.