Fit the TrainerRank layout cost model (coefficient version 2) with a calibration campaign

[bradhilton]

Summary

Replaces the TrainerRank layout scorer's constants with a fitted cost model (coefficient version 2) and lands the calibration tooling that produced it.

The landing (#826) shipped the research thread's layout score. The research thread has since confirmed its constants were hand-set, not fitted: 1 µs per token per layer as a structural scale, 96 + 32·cp per segment and 64 + 32·cp per shared edge hand-shaped, 768 µs per GDN layer as the smallest quantum preserving four measured winners, 256 µs as a launch-floor proxy — applied to the total layer count instead of the GDN layer count (Qwen3.5-4B: 24 of 32), and blind to TP. This PR fits the score from measurements.

What changed in production

• _planner_cost.py: the score is Σ coefficient[term] · term(features, facts) over ten interpretable integer term functions (TERM_FUNCTIONS) of four O(segments) layout features (layout_features: packed tokens, segment count, dependency levels, a segment-length histogram) and the planner facts (cp, tp, layers, GDN layers). The ten terms are the ones that carried weight: per-rank token work with CP-exchange and TP-collective terms, a GDN per-token surcharge, attention KV exchange across CP ranks, tiny-per-rank segments per layer, dependency levels crossing CP or TP ranks, and GDN level hand-offs with their TP interaction; candidates that fitted to zero were dropped (a future campaign can reintroduce them). Coefficients are integer milli-microseconds per feature unit (COEFFICIENTS_MILLI_US); everything stays integer and bit-identical across ranks. COEFFICIENT_VERSION = 2 invalidates cached recipes.
• Planner facts: the layout cache key and select_prefix_tree_layout carry (cp, tp, layers, gdn_layers, uses_gdn, coefficient_version); the GDN layer count is read from the model.
• Calibrated domain (review, narrowed in re-review): the fitted table applies only inside CalibrationProfile, which is exactly the measured envelope and bound to the certificate by test — compute capability 9.0 and an H200-class memory system (device memory ≥ 120 GiB; the 80 GB H100 shares the capability and is excluded), bf16, hidden size 2,560 (hidden size is not a score feature, so neighbouring widths are not admitted), dense models. TrainerRank keeps the version-1 score (kept verbatim as the fallback) outside it, logging once. Capability-based, never model-name-based; CPU-only planning (unit tests) uses the fitted table. Extending the domain means running the calibration cells on the new device or width and regenerating the certificate.
• Tie-breaks unchanged: packed tokens, segment count, maximum depth.

Calibration protocol (new tooling)

• dev/trainer_rank_landing_acceptance.py --phase cost-calibrate: for one cell (workload × model × layers × topology), forces every mandatory candidate layout through the public API with the test-only anchor hook, times forward + backward through an active LoRA slot (compile-free, max-rank), and also times the production selector's own choice; records features, topology and model facts to JSONL. Sky recipes for 4× and 2× H200 cells, a local 1-GPU runner.
• dev/trainer_rank_cost_fit.py: paired within-cell timing deltas (cells weighted equally, pairs by scale), exact non-negative least squares over the production term functions, then a deterministic regret-minimizing coordinate refinement; whole-cell holdouts; the review's noise-qualified gates; --integerize emits the production table; --selector-check runs the shipped table through the real selector on every measured cell.
• Cells: hierarchical GRPO g8 / g16 / g4×4 shapes on Qwen3.5-4B (GDN) and Qwen3-4B (attention) at 2 layers and full height, three heterogeneous controls, real Ellavox groups (corpus mounted, never committed), at TP1/TP2 × CP1/CP2/CP4 on H200 bf16.
• Failure propagation (review): the runners carry torchrun's exit status through pipefail, record failed cells and exit nonzero; the fitter's --require-complete N refuses evidence where any mandatory candidate has fewer than N usable rows, --exclude-cells makes every exclusion explicit and recorded in the certificate, and (re-review) --manifest validates exact cell identities against dev/trainer_rank_cost_calibration_manifest.json — a missing whole cell, an unexpected cell, or duplicate cells with differing execution fingerprints (source, workload hash, device, dtype, hidden size) fail the fit.
• Certificate (review): dev/trainer_rank_cost_calibration_certificate.json (compact: one line per cell, ~80 lines) binds the shipped table to its data — per-cell candidate features, median timings, counts, spreads and fingerprints (no tokens, no per-sample rows), the exact fit arguments, the integer table and its hash, the headline metrics. tests/unit/test_planner_cost_certificate.py asserts the shipped table is the certified table and that the certified metrics hold on the recorded aggregates; --from-certificate re-fits from it (opt-in full reproduction in the test via ART_COST_CERTIFICATE_REFIT=1).

What the data showed

• The cost of an additional shared prefix level is a GDN effect that grows when its state hand-offs cross CP or TP ranks: on the GDN model, sharing a 1,023-token system across two prompt groups wins at TP1/CP1 but loses at CP4 and TP2; the attention model pays almost nothing for the level and benefits from it even at CP4.
• Per-rank token work scales with tp × cp; GDN layers cost more per token than attention layers; rows in segments that are short per rank (threshold × cp) run inefficient kernels.
• The sealed "full sharing 876 ms vs automatic 1,133 ms" gap on the win cell was the research run's online calibration wandering between five layouts, not a ranking error of the frozen score. The frozen score's real misranking there: prompt-level sharing (791 ms) vs full sharing (872 ms), an ~80 ms level cost priced at under 1 ms.

Gates and evidence

Held-out gates (per the review): pairwise ordering ≥ 90% on pairs separated by more than 3%, median regret ≤ 2%, p95 ≤ 5%, none above 10%, clear winners selected within 5%.

Final table: fitted on 45 cells and evaluated on all 56 (3,844 within-cell pairs); the 11 odd Ellavox groups are the pre-registered holdout, never used in the fit. The checked-in certificate refit reproduces the shipped integer table from exactly this procedure.

Selector	Cells	Pairwise (>3%)	Median regret	p95	Max
Version 1 (hand-set)	56	78.6%	0.0%	18.6%	67.5% (Ellavox g6 at CP2)
Version 2 final table (fit on 45), evaluated on all 56	56	98.1%	0.0%	2.9%	4.2%
Version 2, pre-registered holdout (odd Ellavox groups)	11	98.2%	0.0%	4.2%	4.2%
Ablation: every TP2 cell withheld	18	98.3%	0.0%	4.2%	4.2%
Ablation: every CP2 cell withheld	22	96.9%	0.0%	2.0%	4.2%
Ablation: attention model withheld	19	97.9%	0.0%	4.2%	4.2%
Ablation: every heterogeneous cell withheld	25	97.2%	0.0%	4.2%	9.5% (one CP4 heterogeneous cell)
Ablation: every CP4 cell withheld	24	94.8%	0.0%	13.3%	25.0% (does not extrapolate; CP4 cells stay in the fit)
Campaign-1 table (38 cells) on the 18 later cells, through the real selector	18	—	0.0%	—	4.2%
Production selection timed in the campaign-2 CP2/TP2 cells (prospective)	9	—	−0.2%	—	0.4%

The final table run through the real selector (beam search included) on every measured cell: median regret 0%, max 4.2%, never a layout outside the measured family. The four cells above 2% are all decisions between two deep-sharing layouts a few percent apart (a 2-layer attention CP4 cell, an Ellavox group at CP2, a heterogeneous shape at CP4, an Ellavox group at CP4).

Landing gates re-derived and passing: the sealed win-cell shape still selects deep sharing (prompt-level sharing at CP4, where it measures fastest; full sharing at CP1), the heterogeneous control and the tiny sealed families still decline sharing, selection is deterministic. Two width-search unit tests that relied on the old score sharing 40–1,000-token toy prefixes now use prefixes where sharing genuinely pays.

Also in this PR

• The compile watch used by the TP2 gate in Admit tensor-parallel runtimes in TrainerRank (TP>1 support) #832 attached to a guessed logger name and observed nothing; it now takes the telemetry module's logger. The TP2 cell's "compile-free" claim was therefore unverified in Admit tensor-parallel runtimes in TrainerRank (TP>1 support) #832 (its fixed warm-ups made it true in practice); the calibration harness verifies compile-free rows properly.

Known limitations

• Fitted on H200 bf16 at hidden size 2,560; the calibrated profile admits exactly that envelope (H200-class SM 9.0 by capability and device memory, bf16, hidden 2,560, dense) and everything else uses the version-1 score.
• Two Ellavox CP4 cells (groups 1 and 4) are excluded explicitly from the certificate: their CP4 executions hang deterministically in NCCL all-to-alls in the context-parallel group (reproduced on a second fresh cluster at the identical collective sequence numbers). This is a pre-existing CP4 execution bug, filed as TrainerRank CP4: deterministic NCCL all-to-all hang in the context-parallel group on two real Ellavox groups #840; notably the version-1 score selects the hanging layout for group 4 at CP4, while the fitted table happens to select its neighbour.
• Withholding every CP4 cell does not extrapolate (a CP4 heterogeneous cell is then misranked by 25%); CP4 cells stay in the fit. Unseen topologies (TP4, TP×CP) are untested.
• Per-subforward execution overhead and head chunking are still not planner decisions.

🤖 Generated with Claude Code

[bradhilton]

Final table installed (fitted on all 56 cells, 3,844 within-cell pairs) and the PR description now carries the complete evidence. Summary for review:

• Model: integer term table (COEFFICIENTS_MILLI_US × TERM_FUNCTIONS) over O(segments) layout features and the planner facts (cp, tp, layers, GDN layers). Non-zero terms: per-rank token work (plus a GDN-layer surcharge), TP collective terms on tokens, tiny-segment cost per layer, dependency-level costs that scale with (cp−1) and (tp−1), GDN level cost and its TP interaction, and a small attention-area term. Segment launch, fan-out and shared-token terms carried no weight and are kept at zero for table/term parity.
• Fit: paired within-cell deltas, exact active-set NNLS, deterministic regret-minimizing refinement, pre-registered holdout (odd Ellavox groups) plus ablations by topology, shape and model. Reweighting rounds toward mis-ordered pairs were tried and rejected (they distort coefficients); least squares alone misranked close pairs, hence the regret refinement.
• Results: 98.1% pairwise on separated pairs, p95 regret 2.9%, max 4.2%, no clear misses; holdout passes; TP2/CP2/attention ablations pass; the all-heterogeneous ablation leaves one CP4 heterogeneous cell at 9.5%, and the all-CP4 ablation does not extrapolate (CP4 cells stay in the fit). Prospective: the campaign-1 table, through the real selector on the 18 later cells, was within 4.2% everywhere, and the production selection timed in the later CP2/TP2 cells had median regret −0.2%, max 0.4%. The hand-set score on the same cells: 78.6%, max regret 67%.
• Landing gates unchanged in meaning and passing; two width-search unit tests re-tuned to prefixes where sharing genuinely pays.
• Process notes: the sealed research rows are ~9× slower than this harness on the same cell (research-runtime overhead) and were excluded from the fit; the TP2 gate's compile watch in Admit tensor-parallel runtimes in TrainerRank (TP>1 support) #832 was vacuous (guessed logger name) and is fixed here.

[bradhilton]

Codex review — request changes

The scoring implementation is much stronger, but I found two blockers and one evidence gap.

[P1] Gate the fitted model to its validated domain

ScoringFacts contains only TP, CP, total layers, and GDN layers, so the H200/bf16 4B coefficients are applied unchanged to every GPU, dtype, hidden width, and model family. In particular, the relative token-vs-boundary economics remain independent of hidden/FFN width, although that ratio changes substantially between 0.6B, 4B, and 35B/MoE models. This can silently make the default selector worse on supported but unmeasured configurations. Either add the relevant capability/model-shape facts and validate their scaling, or enable V2 only for a calibrated capability profile and retain the previous automatic scorer outside it. This should remain capability-based, not model-name-based.

[P1] Calibration failures are reported as success

The trailing || true defeats set -o pipefail, masking torchrun exceptions, timeouts, and grep/pipeline failures. The same pattern occurs in the CP4 and local runners. The campaign then continues and exits successfully. Since the fitter silently omits candidates with fewer than two usable samples, a failed cell or candidate can disappear while the reduced dataset still passes its gates. Preserve the torchrun exit status in all three runners, then validate expected cells, mandatory candidates, and repetition counts before fitting.

[P2] Check in a reproducible calibration certificate

The coefficient table comment says the table can be regenerated with the fitter, but the PR contains no evidence aggregate, fitted report, command manifest, or hash binding these exact constants to the claimed 56-cell result. The synthetic fitter tests cannot verify the production table or reported regret. Please retain a compact artifact containing per-cell/per-candidate features, medians, counts and fingerprints, the exact fit arguments, and the expected integer-table/report hash; raw tokens and full samples are unnecessary.

Minor cleanup: _impl.py still says the cost model has no TP terms, which is now stale.

Verification: the 30 targeted planner-feature, fitter, and weird-shape tests pass; GitHub CI is also green. I did not find a runtime correctness defect in the new feature algebra itself.

[bradhilton]

Thanks — all three points addressed in 3f9ee53, plus one finding the second point surfaced.

[P1] Calibrated domain. Added CalibrationProfile (capability-based: compute capability 9.0, bf16, hidden size 2,048–3,072, non-MoE) and coefficient_version_for(...). TrainerRank evaluates it once at construction from the model's device capability, parameter dtype, hidden size and the model-support spec's is_moe; outside the profile it keeps the version-1 score, which is retained verbatim as the fallback (_prefix_tree_layout_score_v1) and logs once. The version is part of the planner facts, so it flows into the layout cache key, select_prefix_tree_layout, and the telemetry signature. CPU-only planning (no CUDA device, i.e. the unit tests) uses the fitted table. tests/unit/test_planner_cost_profile.py covers the profile edges (Ampere, fp16, 0.6B- and 8B-class widths, MoE) and that the fallback reproduces the landing formula bit for bit. I took the "gate to a profile" option rather than adding width facts to the terms: hidden width was constant across the campaign, so any width scaling would have been assumed, not fitted.

[P1] Masked failures. The three runners now carry torchrun's status through pipefail (the log filter is wrapped so only its status is ignored), record failed cells, and exit nonzero. The fitter gained --require-complete N (every mandatory candidate of every cell must have N usable rows, else it refuses) and --exclude-cells (exclusions are explicit and recorded in the certificate's fit arguments). Running the completeness check on the campaign evidence immediately found what you predicted: two Ellavox CP4 cells (groups 1 and 4) had zero usable rows for several candidates and had been dropped silently. Their torchrun jobs hung in NCCL all-to-alls in the context-parallel group (ranks stopped at different sequence numbers) — in both logs right after backward kernels were being loaded from the shared on-disk compile cache; one hang on the first run of uniform_depth_2, the other on the fourth run of a layout that had already succeeded three times. That pattern reads like a compile-cache race more than a layout-specific desync, but I'm re-measuring exactly those two cells on a fresh cluster now; if it reproduces at the same layout it's deterministic and I'll escalate it separately as a CP4 execution bug (it predates this PR either way, but the new table can steer selections into deeper layouts on such data). The certificate excludes them explicitly until then.

[P2] Certificate. dev/trainer_rank_cost_calibration_certificate.json (307 KB): per-cell candidate features, median timings, counts, spreads and source/workload fingerprints — no tokens, no per-sample rows — the exact fit arguments including exclusions, the integer table and its SHA-256, and the headline metrics (56 cells, 3,844 pairs, 98.1% pairwise, p95 2.9%, max 4.2%). tests/unit/test_planner_cost_certificate.py asserts the shipped COEFFICIENTS_MILLI_US is the certified table (hash) and that the certified metrics are exactly what the table produces on the recorded aggregates; ART_COST_CERTIFICATE_REFIT=1 additionally re-runs the full fit from the aggregates and requires the identical integer table. --from-certificate re-fits from the artifact without the raw evidence.

Minor: the stale "no TP terms" comment in _impl.py is updated. I'll follow up here with the re-measurement result (and a regenerated certificate if the two cells come back clean).

[bradhilton]

Re-measurement result: both Ellavox CP4 cells (groups 1 and 4) hung again on a fresh cluster at the identical NCCL collective sequence numbers, so this is a deterministic CP4 execution bug, not a compile-cache race. Group 4 hangs on the first execution of its uniform_depth_2 layout (nested sharing with short unequal tails); group 1 hangs on the first no_sharing run after other layouts have executed. Filed as #840 with the repro command and details. Relevant for this PR: the version-1 score selects the hanging layout for group 4 at CP4 (production would hang on that data today), while the fitted table selects its neighbour minimum_effective_span_495 — incidental, not a fix. The two cells stay explicitly excluded from the certificate (recorded in its fit arguments); the other 56 cells are complete. The new runner correctly reported both as CELL FAILED and exited nonzero.

[bradhilton]

Codex re-review — two residual changes requested

The original three findings are substantially addressed: the version-1 fallback is faithful, the runners now propagate failures, the certificate is compact and useful, and the full certificate refit reproduces the shipped integer table locally. I found two remaining evidence/scope gaps and one documentation correction.

[P1] The calibrated profile still admits unmeasured hardware and model widths

CalibrationProfile says extrapolation is avoided, but the certificate contains only NVIDIA H200 at hidden size exactly 2,560, while the runtime admits every SM90 device and hidden sizes 2,048–3,072. SM90 includes H100, whose memory system differs from H200, and 2,048 admits materially smaller models; neither expansion is backed by a calibration row or a fitted scaling term. Because hidden size is not a score feature, every width in that range gets identical token-vs-boundary economics. Please either narrow the profile to the actually measured device/width or add guardrail measurements supporting the wider envelope. This can remain capability-based and does not require model-name checks. The certificate/profile test should bind the admitted envelope to the certified evidence or an explicit validated extension.

[P2] Completeness still cannot detect a missing whole cell or mixed cell identities

validate_completeness derives the expected set only from calibration_cell rows that already exist. A process that fails before emitting that row, an omitted evidence file, or an unlaunched recipe is therefore invisible and --require-complete can still pass. The certificate test compounds this by accepting any self-recorded count >= 50 rather than the intended cell identities. Also, _cell_key omits source, workload hash, device and dtype, so repeated evidence with the same display key but incompatible fingerprints is pooled while the certificate records only the last fingerprint. Add a small expected-cell manifest and validate exact identities, and reject or explicitly authorize duplicate keys whose execution fingerprints differ. The current certificate should assert the exact 56 retained cells plus the two explicit exclusions.

Documentation correction

The PR repeatedly says the final table was “fitted on all 56 cells,” but the checked-in refit uses 45 training cells and preserves 11 odd-Ellavox cells as holdout; I ran ART_COST_CERTIFICATE_REFIT=1 and that 45-cell fit reproduces the table. That is arguably stronger evidence, but the text should say “fit on 45 and evaluated on all 56,” unless the intended final procedure is actually to refit on all 56 after the holdout decision.

Verification on head 8826c179d: 37 targeted tests pass, including the opt-in full certificate refit; GitHub CI is green. I found no new runtime correctness defect in the scorer/fallback plumbing.

[bradhilton]

Re-review items addressed in 624d46c:

• [P1] Calibrated profile narrowed to the measured envelope. CalibrationProfile now admits exactly what the certificate contains: compute capability 9.0 and an H200-class memory system (device memory ≥ 120 GiB, so the 80 GB H100 that shares the capability is excluded), bf16, hidden size exactly 2,560, dense models. Neighbouring widths (2,048, 3,072) are no longer admitted, since hidden size is not a score feature. TrainerRank passes the device's total memory alongside the capability. The certificate test binds the profile's device names, dtypes and hidden sizes to the certified evidence (measured_envelope is recorded in the certificate), so widening the envelope without new calibration cells fails the test. Still capability-based; names appear only as documentation.
• [P2] Exact cell identities. dev/trainer_rank_cost_calibration_manifest.json lists the 58 cells the recipes launch plus the two explicit exclusions with reasons. The fitter's --manifest validation requires every non-excluded cell to be present, rejects unexpected cells, rejects an exclusion the manifest does not list, and rejects a cell key recorded with more than one execution fingerprint (source, workload hash, device, dtype, hidden size). The certificate records the manifest and the resolved exclusions; its test asserts the exact 56 retained identities and the two exclusions (replacing the >= 50 count) and that every retained cell carries its fingerprints. A negative unit test covers the missing/unexpected/mixed cases.
• Wording corrected: fitted on 45 cells (the 11 odd Ellavox groups are the pre-registered holdout, never used in the fit) and evaluated on all 56; the checked-in refit reproduces the table from that procedure. PR description, design brief, README and the module docstring say so.
• The harness now records device capability and total memory in evidence for future campaigns.

The integer table is unchanged; the certificate was regenerated with the manifest (validation: 58 expected, 2 excluded, fingerprints consistent, every mandatory candidate ≥ 8 rows). Targeted tests: 380 passing plus the opt-in full refit.

[bradhilton]

Codex re-review — one residual evidence fix

The prior findings are otherwise resolved on 624d46cd5: the fitted scorer is now gated to the measured H200/bf16/hidden-2560/dense envelope, exact evidence identities and exclusions are asserted, the certificate refit still reproduces the shipped table, and the fit/holdout wording is accurate.

[P2] The checked-in recipes do not produce three cells required by the manifest

The manifest says it is the exact union of the checked-in recipes, but it contains three campaign: local2 TP1×CP1 cells (cal-hetero2 on Qwen3 and Qwen3.5, plus cal-hetero3 on Qwen3.5) that no checked-in runner launches. The local runner launches 15 cells when the Ellavox corpus is present; the two 2-GPU modes launch 22 and the two CP4 modes launch 18, for 55 total. The manifest requires 58. Consequently, a clean rerun of every documented recipe produces only 53 retained cells after the two exclusions, then correctly fails its new manifest gate instead of reproducing the certified 56.

Please add those three invocations to trainer_rank_cost_calibration_local.sh (or check in the missing local2 recipe) so the manifest and recipes are actually mutually reproducible. This is a small provenance fix; I found no remaining production scorer defect.

Verification on 624d46cd5: 40 focused tests pass, including ART_COST_CERTIFICATE_REFIT=1. Quality CI is green; the 2×H200 check is still running at review time.

[bradhilton]

Two more commits:

[P2] Recipes ↔ manifest (090e19b). The local runner now launches the three heterogeneous-variant TP1×CP1 cells (hetero2 on both models, hetero3) that had been run ad hoc; the manifest's campaign labels name checked-in recipes; and tests/unit/test_planner_cost_manifest_recipes.py parses the three runners (both CELL_SET modes, Ellavox loops expanded, layers=0 resolved per model) and asserts they launch exactly the manifest's 58 cells. A clean rerun of every documented recipe now reproduces the certified set minus the two explicit exclusions.

Right-sizing pass (7d76cfc), at the author's request. The PR was ~21k lines, 83% of it the pretty-printed certificate. Now 2,974 insertions:

• The certificate and manifest are written one line per cell (81 and 68 lines instead of 17,368 and 601), identical content.
• The production module keeps only the ten term functions that carried weight and the four features they read (packed tokens, segment count, dependency levels, segment-length histogram). The table was refitted on exactly those terms; the regret refinement lands on an equivalent optimum, so the coefficients changed while every metric is identical (98.1% pairwise, p95 regret 2.9%, max 4.2%, holdout pass, selector check median 0% / max 4.2%, no unmeasured selections) and every coefficient is non-zero. Certificate regenerated; the opt-in full refit reproduces it. Candidate terms that fitted to zero (segment launches, fan-out, shared tokens, attention area, small-M token surcharges) are gone from the module and can be reintroduced by a future campaign.
• The fitter drops the rejected reweighting rounds, the one-off feature refresh, and the version-1 comparison report.

The profile envelope, manifest validation, exclusions and wording from the previous round are unchanged. 557 planner-related tests pass; prek green.