diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md
index d5149969d13..3f3074205ae 100644
--- a/examples/puzzletron/README.md
+++ b/examples/puzzletron/README.md
@@ -468,7 +468,7 @@ The authoritative dependencies and enabled-stage rules live in
| `width_importance` | Collect activation-based rankings for every enabled width axis. |
| `sort` | Reorder the teacher so nested prefixes implement ranked width choices. |
| `sort_sanity` | Check that sorting preserves teacher outputs. |
-| `width_sanity` | Compare ranked, random, and reverse slices on representative layers. |
+| `width_sanity` | Compare ranked, original-order, and reverse slices on representative layers. |
| `slicing_sanity` | Verify dynamic slicing against physical materialization. |
| `bypass_sanity` | Overfit small local-distillation cases before production bypass. |
| `bypass` | Train nested replacement blocks across the configured search space. |
@@ -481,6 +481,27 @@ The authoritative dependencies and enabled-stage rules live in
| `global_distillation` | Distill the selected architecture at the configured production scale. |
| `post_distillation_evaluation` | Evaluate the final distilled checkpoint. |
+## Interpret width sanity results
+
+Puzzletron separates implementation correctness from ranking quality:
+
+- Sort and slicing equivalence failures are correctness errors and always fail
+ their stages.
+- A width-ranking miss means the activation-sorted candidate underperformed an
+ original or reverse control. It is a quality warning by default.
+
+To also fail a sanity stage on ranking-quality warnings, enable strict warning
+handling:
+
+```yaml
+sanity:
+ fail_on_warnings: true
+```
+
+See [Sorting, width ranking, and slicing sanity](docs/sanity_validation.md) for
+the slicing mental model, measured metrics, comparison controls, tolerances,
+worked example, and qualification guidance.
+
Independent DAG branches may run concurrently when they have disjoint writers.
Long-running stages should resume their durable checkpoints or immutable shards
rather than restarting completed work.
diff --git a/examples/puzzletron/docs/sanity_validation.md b/examples/puzzletron/docs/sanity_validation.md
new file mode 100644
index 00000000000..40f30151ace
--- /dev/null
+++ b/examples/puzzletron/docs/sanity_validation.md
@@ -0,0 +1,124 @@
+# Sorting, width ranking, and slicing sanity
+
+Puzzletron uses separate sanity checks to answer two different questions:
+
+1. Did sorting and slicing preserve the intended model operation?
+2. Did the importance ranking produce a better reduced candidate than its
+ controls?
+
+The first question is about implementation correctness. The second is about
+ranking quality. Keeping them separate helps diagnose a failed campaign and
+lets an acceptance plan decide how strict ranking quality must be.
+
+## Sorting and slicing mental model
+
+Width pruning makes selected model dimensions smaller, such as hidden channels,
+attention heads, expert widths, or state-space dimensions. Puzzletron first
+scores the relative importance of the units in each supported dimension. It
+then permutes every coupled tensor into the same importance order so that a
+prefix of the sorted dimension represents a reduced-width candidate. At full
+width, sorting is only a permutation and must not change model behavior.
+
+A **dynamic slice** runs a reduced candidate from the resident full-size sorted
+teacher. Runtime hooks or views select only the requested prefix without first
+writing a smaller checkpoint. A **physical slice**, also called physical
+materialization, actually removes the excluded tensor rows, columns, heads, or
+groups and updates the model config. The physical model is the exportable
+ground truth for the dynamic path.
+
+For example, reducing hidden width from 4096 to 3072 does not mean truncating
+one weight matrix. Puzzletron must select the same 3072 channel identities
+across every coupled embedding, residual, normalization, and projection
+dimension. Sorting first moves the chosen identities into a common prefix;
+slicing then applies that prefix consistently throughout the model.
+
+## Compared views
+
+The width sanity stages compare four views of one target width:
+
+| View | Meaning |
+|---|---|
+| Activation-sorted | Keep the prefix that the importance scores rank highest. This is the candidate whose ranking quality is being tested. |
+| Original-order | Keep the same-sized prefix in the teacher's original order. Some artifacts retain the legacy method key `random`, but this control is deterministic and is not a seeded-random permutation. |
+| Reverse-sorted | Keep the prefix ranked least important. This is a negative control for the importance ordering. |
+| Physical | Rewrite the config and tensors as a smaller standalone model using the same target geometry as the activation-sorted candidate. |
+
+## Measured values
+
+Each comparison uses the same axis, layer, target width, model revision, and
+input samples. Depending on the diagnostic configuration, the measured values
+can include:
+
+- replacement loss or language-model loss;
+- hidden-state cosine distance, mean squared error, or mean absolute error;
+- output-distribution KL divergence; and
+- token top-k accuracy or consistency.
+
+Loss, distance, error, and divergence metrics are lower-is-better. Accuracy and
+consistency metrics are higher-is-better. The artifact records the metric,
+direction, compared methods, observed degradation or difference, and allowed
+tolerance so a finding can be traced to measured evidence.
+
+## Width-ranking misses
+
+For the same axis, layer, target width, samples, and metric, width sanity asks
+whether the activation-sorted candidate is at least as good as the original
+and reverse controls within the configured comparison tolerance. A ranking
+miss occurs when the activation-sorted value is worse than a control by more
+than that tolerance.
+
+For example, if lower loss is better and the activation-sorted loss is `2.10`
+while the original control is `2.00` with tolerance `0.01`, the degradation is
+`0.10` and the check reports a ranking miss. This means the importance ordering
+did not demonstrate the expected quality benefit for that case. It does not by
+itself prove that sorting or slicing produced an invalid model.
+
+## Equivalence failures
+
+Equivalence checks compare two routes that are intended to represent the same
+model operation:
+
+- Sort sanity compares the full-width teacher with full-width sorted and
+ reverse-sorted teachers. These permutations must preserve behavior within
+ descriptor-owned tolerances.
+- Slicing sanity compares a dynamic slice with a physically materialized model
+ for the same axis, layers, target value, batch, and model revision. It checks
+ loss, output shape, and output tensor differences using explicit absolute and
+ relative tolerances.
+
+An equivalence failure means those supposedly identical routes disagree beyond
+their declared tolerance. Downstream measurements from the dynamic candidate
+can no longer be assumed to describe the exportable physical checkpoint, so
+the corresponding sanity stage fails unconditionally.
+
+## Stage completion and qualification
+
+Puzzletron currently records correctness failures and ranking-quality findings
+through the same sanity-stage result, but they have different default effects:
+
+| Result | Current stage behavior |
+|---|---|
+| Sort or slicing equivalence failure | Always fails the stage and is reported as a correctness error. |
+| Width-ranking miss | Records a quality warning. By default the stage may complete successfully with `passed: false`. |
+| Width-ranking miss with strict warning policy | Fails the stage when `sanity.fail_on_warnings` is enabled. |
+
+Enable strict warning handling with:
+
+```yaml
+sanity:
+ fail_on_warnings: true
+```
+
+This setting changes stage-completion policy; it does not reclassify a ranking
+miss as an implementation-correctness error.
+
+Campaign or release qualification is a separate policy decision. The current
+code does not emit a distinct `qualification_blocked` state. Before treating a
+campaign as accepted, define which ranking controls must pass, the metrics and
+directions, sample counts, tolerances, covered axes and target widths, and the
+aggregation rule. A stage-level quality warning may therefore remain
+non-correctness evidence while still blocking scientific or release
+acceptance.
+
+For the broader pipeline design, see the
+[semantic validation gates](v2_architecture.md#semantic-validation-gates).
diff --git a/examples/puzzletron/docs/v2_architecture.md b/examples/puzzletron/docs/v2_architecture.md
index cd285dcfcee..1b9edd39ff8 100644
--- a/examples/puzzletron/docs/v2_architecture.md
+++ b/examples/puzzletron/docs/v2_architecture.md
@@ -10,14 +10,14 @@ Puzzletron v2 turns pruning from a sequence of model-specific scripts into a
distributed, validated, end-to-end campaign. The design is driven by three
goals:
-1. **Scalability** — run large and long-context models with stage-specific
+1. **Scalability**, run large and long-context models with stage-specific
tensor, context, pipeline, expert, data, and sequence parallelism; avoid
repeatedly loading or materializing checkpoints; and distribute independent
work through persistent or sharded workers.
-2. **Semantic correctness** — describe every model and pruning axis explicitly,
+2. **Semantic correctness**, describe every model and pruning axis explicitly,
collect all compatible importance statistics together, and gate the campaign
with sorting, ranking, slicing, bypass, and distillation checks.
-3. **End-to-end automation** — generate a campaign from a setup wizard, execute
+3. **End-to-end automation**, generate a campaign from a setup wizard, execute
its dependency graph, search for architectures, run downstream processing,
and continuously assemble a durable HTML report.
@@ -138,8 +138,8 @@ flowchart TB
- **Granularity is explicit.** Bypass, replacement scoring, vLLM statistics,
and depth decisions can operate at block or subblock granularity.
- **Correctness and quality are separate.** Sort and slice equivalence are
- correctness gates. Ranking against reverse and unsorted/random controls is a
- quality gate whose warning remains visible.
+ correctness gates. Ranking against original-order and reverse controls is a
+ quality check whose warning remains visible.
- **Artifacts are APIs.** Stages communicate through versioned, hashed,
transactionally published artifacts rather than in-memory coupling. This
enables resume, parallel execution, and report regeneration.
@@ -284,7 +284,7 @@ flowchart LR
width["Multi-axis width importance hooks"]
sort["Sort teacher once"]
sortcheck["Sort equivalence"]
- widthcheck["Ranking quality sorted vs reverse vs unsorted"]
+ widthcheck["Ranking quality sorted vs original vs reverse"]
slicecheck["Dynamic vs physical slicing equivalence"]
bypasscheck["Bypass overfit checks"]
bypass["Nested bypass"]
@@ -358,13 +358,33 @@ durable manual decision gate.
|---|---|---|
| Capability validation | Does the model support every requested axis, backend, and parallel mode? | Invalid campaign configuration fails before expensive work |
| Sort sanity | Does full-width sorting or reverse sorting preserve teacher behavior? | Difference beyond dtype-aware tolerance is a correctness failure |
-| Width sanity | Does the proposed ranking outperform reverse and unsorted/random controls at reduced width? | Poor ranking is surfaced as a quality warning |
+| Width sanity | Does the proposed ranking outperform original-order and reverse controls at reduced width? | Poor ranking is a quality finding; warning policy determines whether it also fails the stage |
| Slicing sanity | Does dynamic slicing agree with physical materialization? | Disagreement is a correctness failure |
| Bypass sanity | Can a fixed small candidate and a sampled nested search overfit one batch? | Validates boundaries, gradients, and sampling mechanics |
| Depth evaluation | Are removal scores recomputed after every selected removal? | Produces a conditional trajectory rather than independent linear scores |
| Global KD sanity | Can the student overfit with the configured CE/KLD/MTP loss path? | Validates forward/backward and loss semantics before a long run |
| Artifact completion | Are all expected identities, shards, candidates, and outputs present? | Partial work remains resumable progress, not a completed stage |
+These gates answer different questions. Width ranking compares the quality of
+different reduced candidates at the same target geometry. Sort and slicing
+equivalence compare routes that are supposed to represent the same model
+operation. A poor ranking can show that the importance heuristic is not useful
+for a case even when every candidate is structurally valid. An equivalence
+failure shows that a permutation or runtime slice does not reproduce its
+reference implementation, so later measurements cannot be trusted as evidence
+for the physical checkpoint.
+
+The implementation has two stage-completion policies. Correctness failures
+always fail their stage. Ranking-quality findings remain warnings by default
+and fail the stage when `sanity.fail_on_warnings` is enabled. Scientific,
+customer, or release qualification is a third layer and is not currently a
+separate Puzzletron verdict. A qualification plan must declare its required
+controls, metrics, sample count, tolerances, axis and target coverage, and
+aggregation rule. It may reject a campaign for a ranking-quality warning
+without reclassifying that warning as a correctness error. The
+[sanity validation guide](sanity_validation.md) provides definitions, the
+slicing mental model, and a worked example.
+
## Current implementation versus design direction
### Implemented in the current v2 code
diff --git a/modelopt/torch/puzzletron/diagnostics/campaign_findings.py b/modelopt/torch/puzzletron/diagnostics/campaign_findings.py
index 94aaca69e8e..03cdaf49b48 100644
--- a/modelopt/torch/puzzletron/diagnostics/campaign_findings.py
+++ b/modelopt/torch/puzzletron/diagnostics/campaign_findings.py
@@ -34,12 +34,12 @@ class MetricSpec:
@dataclass(frozen=True)
class Finding:
- """One non-blocking warning derived from report evidence."""
+ """One evidence-derived advisory warning or correctness error."""
stage: str
message: str
evidence: Mapping[str, Any]
- severity: Literal["warning"] = "warning"
+ severity: Literal["warning", "error"] = "warning"
def _allowed(left: float, right: float, spec: MetricSpec) -> float:
@@ -72,7 +72,7 @@ def equivalence_findings(
group_keys: Sequence[str],
method_key: str = "method",
) -> list[Finding]:
- """Warn when paired methods differ beyond a metric's equivalence tolerance."""
+ """Report a correctness error when equivalent methods disagree beyond tolerance."""
findings = []
for group, values in _groups(rows, group_keys).items():
@@ -111,6 +111,7 @@ def equivalence_findings(
"delta": delta,
"tolerance": allowed,
},
+ severity="error",
)
)
return findings
@@ -126,7 +127,12 @@ def ranking_findings(
group_keys: Sequence[str],
method_key: str = "method",
) -> list[Finding]:
- """Warn when a preferred ranking is worse than a comparison beyond tolerance."""
+ """Report a quality warning when a preferred ranking is worse than a control.
+
+ The warning evaluates the ranking heuristic, not whether the compared model
+ transformations are equivalent. A caller may still promote it through a
+ stricter qualification policy.
+ """
findings = []
for group, values in _groups(rows, group_keys).items():
diff --git a/modelopt/torch/puzzletron/diagnostics/campaign_progress_report.py b/modelopt/torch/puzzletron/diagnostics/campaign_progress_report.py
index 1cbb598b52e..c00e8e47212 100644
--- a/modelopt/torch/puzzletron/diagnostics/campaign_progress_report.py
+++ b/modelopt/torch/puzzletron/diagnostics/campaign_progress_report.py
@@ -39,6 +39,7 @@
publish_report_transaction,
stable_digest,
)
+from .sanity_verdict import is_correctness_sanity_stage
from .width_sanity import descriptor_realization_findings
_STAGES = tuple(spec.stage_id for spec in STAGE_SPECS)
@@ -173,7 +174,7 @@ def _stage_artifact_present(root: Path, spec: StageSpec) -> bool:
def _pipeline_state(root: Path, spec: StageSpec, config: dict[str, Any]) -> str:
- """Return completed, disabled, or pending from artifacts and configuration."""
+ """Return the report state from the manifest, artifacts, and configuration."""
post_mip_stages = {
"zero_shot_evaluation",
@@ -184,6 +185,15 @@ def _pipeline_state(root: Path, spec: StageSpec, config: dict[str, Any]) -> str:
}
if (config.get("post_mip") or {}).get("flows") and spec.stage_id in post_mip_stages:
return "disabled"
+ if is_correctness_sanity_stage(spec.stage_id):
+ manifest = _manifest(root, spec.stage_id)
+ summary = _load_optional(root / "artifacts" / spec.stage_id / "summary.json")
+ if (
+ manifest.get("status") == "failed"
+ or summary.get("passed") is False
+ or summary.get("verdict") == "failed"
+ ):
+ return "failed"
if _stage_artifact_present(root, spec):
return "completed"
section = config.get(spec.stage_id)
@@ -284,8 +294,12 @@ def _sort_table(summary: dict[str, Any]) -> str:
f"{reverse_cell}"
""
)
- gate = "passed" if summary.get("passed") is True else "warning"
- gate_label = "passed" if summary.get("passed") is True else "warning"
+ gate = (
+ "passed"
+ if summary.get("passed") is True
+ else "failed (blocking correctness)"
+ )
+ gate_label = "passed" if summary.get("passed") is True else "failed"
findings = list(summary.get("findings") or ())
finding_notes = ""
gate_attributes = ""
@@ -402,6 +416,7 @@ def _activation_diagnostic_summary(root: Path) -> dict[str, Any]:
if identity not in seen:
rows.append(row)
seen.add(identity)
+ width_findings = list(width.get("findings") or ())
slicing_findings = list(slicing.get("findings") or ())
if not slicing_findings:
slicing_findings = [
@@ -413,7 +428,13 @@ def _activation_diagnostic_summary(root: Path) -> dict[str, Any]:
return {
"rows": rows,
"axes": sorted({str(row.get("axis")) for row in rows if row.get("axis")}),
- "width_findings": list(width.get("findings") or ()),
+ "width_present": bool(width),
+ "width_passed": width.get("passed", not width_findings) is True if width else None,
+ "slicing_present": bool(slicing),
+ "slicing_passed": (
+ slicing.get("passed", not slicing_findings) is True if slicing else None
+ ),
+ "width_findings": width_findings,
"slicing_findings": slicing_findings,
"sort_findings": list((_load_optional(root / "artifacts" / "sort_sanity" / "summary.json")).get("findings") or ()),
}
@@ -461,8 +482,26 @@ def _activation_diagnostic_section(summary: dict[str, Any]) -> str:
summary.get("slicing_findings") or ()
) + list(summary.get("sort_findings") or ())
rows = _activation_diagnostic_rows(summary)
+ gates = []
+ if summary.get("width_present"):
+ width_passed = summary.get("width_passed") is True
+ gates.append(
+ "
Width ranking: {}
".format(
+ "passed" if width_passed else "warning",
+ "passed" if width_passed else "quality warning",
+ )
+ )
+ if summary.get("slicing_present"):
+ slicing_passed = summary.get("slicing_passed") is True
+ gates.append(
+ "
" if gates else ""
if not rows:
- return (
+ return gate_summary + (
"
The sanity artifact contains no plottable numeric metrics.
"
if findings
else "
No width or slicing sanity artifact.
"
@@ -486,7 +525,9 @@ def _activation_diagnostic_section(summary: dict[str, Any]) -> str:
for metric in metrics
)
if not axes or not metrics:
- return "
The sanity artifact contains no plottable numeric metrics.
"
+ return gate_summary + (
+ "
The sanity artifact contains no plottable numeric metrics.
Every sliced model is compared with the full, unsliced original teacher. "
"The original baseline is the teacher channel order sliced to the same target. "
- "Physical is a materialized slice of the sorted checkpoint and is the slicing ground truth.
"
+ "Physical is a materialized slice of the sorted checkpoint and is the slicing ground truth. "
+ "A width-ranking quality warning means the activation-sorted candidate was worse than "
+ "an original or reverse control beyond tolerance; it does not mean the dynamic and "
+ "physical implementations disagree. Campaign qualification may still require the "
+ "ranking warning to pass."
""
f''
""
@@ -3262,6 +3307,7 @@ def _stage_dag(
"orient='auto'>"
f"{''.join(edges)}{''.join(nodes)}"
"