From e6dfc02994ab058a5ce5e48e4e4df29a78450c21 Mon Sep 17 00:00:00 2001
From: Johannes Rausch <jrausch@nvidia.com>
Date: Thu, 6 Aug 2026 15:06:24 +0200
Subject: [PATCH 1/6] Fail closed on Puzzletron sanity failures

Treat sort and slicing-equivalence mismatches as blocking correctness errors while keeping ranking-quality findings advisory. Propagate the same contract through parent sweeps and progress reports.

Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
---
 .../diagnostics/campaign_findings.py          |   5 +-
 .../diagnostics/campaign_progress_report.py   |  56 +++++++--
 .../puzzletron/diagnostics/sanity_verdict.py  |  53 ++++++--
 .../puzzletron/diagnostics/width_sanity.py    |   1 +
 .../torch/puzzletron/stages/diagnostics.py    |  52 ++++++--
 .../puzzletron/test_campaign_findings.py      |   5 +-
 .../test_campaign_progress_report.py          |  39 ++++--
 .../test_diagnostic_scoring_config.py         |  26 ----
 .../test_hidden_width_diagnostic.py           |  21 ++++
 .../torch/puzzletron/test_sanity_verdict.py   | 114 ++++++++++++++++--
 .../test_width_sanity_aggregation.py          |   5 +-
 .../test_width_slice_equivalence.py           |  10 +-
 12 files changed, 310 insertions(+), 77 deletions(-)

diff --git a/modelopt/torch/puzzletron/diagnostics/campaign_findings.py b/modelopt/torch/puzzletron/diagnostics/campaign_findings.py
index 94aaca69e8e..9fbd4465c86 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:
@@ -111,6 +111,7 @@ def equivalence_findings(
                         "delta": delta,
                         "tolerance": allowed,
                     },
+                    severity="error",
                 )
             )
     return findings
diff --git a/modelopt/torch/puzzletron/diagnostics/campaign_progress_report.py b/modelopt/torch/puzzletron/diagnostics/campaign_progress_report.py
index 1cbb598b52e..898fc094a1e 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}"
             "</tr>"
         )
-    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(
+            "<p class='gate {}'>Width ranking: {}</p>".format(
+                "passed" if width_passed else "warning",
+                "passed" if width_passed else "warning (advisory)",
+            )
+        )
+    if summary.get("slicing_present"):
+        slicing_passed = summary.get("slicing_passed") is True
+        gates.append(
+            "<p class='gate {}'>Dynamic/physical equivalence: {}</p>".format(
+                "passed" if slicing_passed else "failed",
+                "passed" if slicing_passed else "failed (blocking correctness)",
+            )
+        )
+    gate_summary = f"<div class='sanity-gates'>{''.join(gates)}</div>" if gates else ""
     if not rows:
-        return (
+        return gate_summary + (
             "<p class='empty'>The sanity artifact contains no plottable numeric metrics.</p>"
             if findings
             else "<p class='empty'>No width or slicing sanity artifact.</p>"
@@ -539,7 +578,7 @@ def _activation_diagnostic_section(summary: dict[str, Any]) -> str:
             + "".join(cells)
             + "</tr>"
         )
-    return (
+    return gate_summary + (
         "<p class='note'>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.</p>"
@@ -3262,6 +3301,7 @@ def _stage_dag(
         "orient='auto'><path d='M0,0 L7,3.5 L0,7 Z'/></marker></defs>"
         f"{''.join(edges)}{''.join(nodes)}</svg></div>"
         "<div class='dag-legend'><span class='completed'>Completed</span>"
+        "<span class='failed'>Failed</span>"
         "<span class='pending'>Pending</span>"
         "<span class='disabled'>Disabled</span></div>"
         "<div class='dag-legend dag-type-legend'><span class='required-node'>Required</span>"
@@ -3934,8 +3974,8 @@ def build_section(
 .hoverlayer .hovertext path,.hoverlayer .hovertext rect,.hoverlayer .axistext path{{fill-opacity:.32!important}} .cardinality-formula{{padding:14px 16px;border:1px solid #4f8cff66;border-radius:10px;background:#0a1421;color:#cfe0ff;font:600 15px/1.5 ui-monospace,SFMono-Regular,Menlo,monospace}}
 pre{{max-height:480px;overflow:auto;background:#07101b;border:1px solid #1d2a3d;border-radius:12px;padding:16px;color:#cfe0ff;white-space:pre-wrap}}
 .summary-grid{{display:grid;grid-template-columns:repeat(auto-fit,minmax(230px,1fr));gap:10px}} .summary-grid article{{display:flex;min-height:86px;flex-direction:column;justify-content:center;gap:7px;padding:14px 16px;border:1px solid var(--line);border-radius:12px;background:#0a1421}} .summary-grid span{{color:var(--muted);font-size:12px;text-transform:uppercase;letter-spacing:.06em}} .summary-grid strong{{color:#cfe0ff;font-size:16px}}
-.dag-scroll{{overflow-x:auto;padding:6px 0 12px}} .stage-dag{{display:block;min-width:100%;font-family:Inter,ui-sans-serif,system-ui,sans-serif}} .dag-edge{{fill:none;stroke:#4f8cff88;stroke-width:2;marker-end:url(#dag-arrow)}} .dag-edge.muted{{stroke:#55647a55;stroke-dasharray:5 6}} .dag-node rect{{fill:#0b1421;stroke:#34445e;stroke-width:1.5;transition:fill .15s,stroke .15s}} .dag-node.optional rect{{stroke-dasharray:6 4}} .dag-node circle{{fill:#55647a}} .dag-node .dag-label{{fill:var(--ink);font-size:11px;font-weight:650}} .dag-node .dag-status{{fill:var(--muted);font-size:9px;letter-spacing:.08em}} .dag-node.completed circle{{fill:var(--green)}} .dag-node.completed rect{{stroke:#35d07f99}} .dag-node.pending circle{{fill:var(--amber)}} .dag-node.pending rect{{stroke:#ffbd4577}} .dag-node.disabled{{opacity:.38}} .stage-dag a:hover .dag-node rect{{fill:#101d30;stroke:var(--blue)}} .dag-legend{{display:flex;gap:16px;flex-wrap:wrap;color:var(--muted);font-size:12px;margin-top:5px}} .dag-legend span::before{{content:'';display:inline-block;width:8px;height:8px;border-radius:50%;margin-right:6px;background:#55647a}} .dag-legend .completed::before{{background:var(--green)}} .dag-legend .pending::before{{background:var(--amber)}} .dag-legend .disabled::before{{background:#55647a;opacity:.45}} .dag-type-legend span::before{{width:13px;height:9px;border-radius:2px;background:transparent;border:1.5px solid #93a4bd}} .dag-type-legend .optional-node::before{{border-style:dashed}}
-@keyframes pulse{{50%{{transform:scale(1.7);box-shadow:0 0 18px #ffbd45aa}}}} .table-wrap{{overflow:auto}} table{{width:100%;border-collapse:collapse}} th,td{{padding:10px 12px;border-bottom:1px solid var(--line);text-align:right}} th:first-child{{text-align:left}} thead th{{color:#a9bee0;background:#0a1421;position:sticky;top:0}} tr.selected-candidate{{background:#4f8cff18}} tr.selected-candidate td{{box-shadow:inset 0 1px #4f8cff44,inset 0 -1px #4f8cff44}} .candidate-swatch{{display:inline-block;width:10px;height:10px;margin-right:8px;border-radius:50%}} td.warning-cell{{background:#ffbd4524;color:#ffe3a3;box-shadow:inset 0 0 0 1px #ffbd4570}} .warning-value{{cursor:help;outline-offset:3px}} #warning-tooltip{{position:fixed;z-index:10000;max-width:min(420px,calc(100vw - 24px));padding:9px 11px;border:1px solid #ffbd4588;border-radius:8px;background:#241b0d;color:#ffe3a3;font-size:12px;line-height:1.4;box-shadow:0 10px 30px #0009;pointer-events:none;white-space:pre-wrap}} .gate{{display:inline-block;border-radius:999px;padding:5px 10px;background:#1b2839}} .gate.passed{{color:var(--green)}} .empty,.note{{color:var(--muted)}} select{{margin:0 18px 18px 8px;padding:8px 12px;border:1px solid var(--line);border-radius:8px;background:#0a1421;color:var(--ink)}} .selector-label{{color:var(--muted)}} code{{color:#bcd2ff}} .probe-summaries{{display:grid;grid-template-columns:repeat(auto-fit,minmax(260px,1fr));gap:12px;margin:16px 0}} .probe-summary{{border:1px solid var(--line);border-radius:12px;padding:14px;background:#0a1421}} .probe-summary.passed{{border-color:#35d07f66}} .probe-summary.failed{{border-color:#ff657766}} .probe-summary h3{{margin:0 0 10px}} .probe-summary dl{{display:grid;grid-template-columns:1fr auto;gap:5px 14px;margin:0}} .probe-summary dt{{color:var(--muted)}} .probe-summary dd{{margin:0;text-align:right}} .probe-plots{{display:grid;grid-template-columns:repeat(auto-fit,minmax(430px,1fr));gap:14px}} .probe-plot-panel{{min-width:0;border:1px solid var(--line);border-radius:12px;padding:12px;background:#091321}} .probe-plot-panel h3{{margin:0 0 4px}} .plotly-chart{{width:100%;height:390px}} .depth-chart{{height:460px}}
+.dag-scroll{{overflow-x:auto;padding:6px 0 12px}} .stage-dag{{display:block;min-width:100%;font-family:Inter,ui-sans-serif,system-ui,sans-serif}} .dag-edge{{fill:none;stroke:#4f8cff88;stroke-width:2;marker-end:url(#dag-arrow)}} .dag-edge.muted{{stroke:#55647a55;stroke-dasharray:5 6}} .dag-node rect{{fill:#0b1421;stroke:#34445e;stroke-width:1.5;transition:fill .15s,stroke .15s}} .dag-node.optional rect{{stroke-dasharray:6 4}} .dag-node circle{{fill:#55647a}} .dag-node .dag-label{{fill:var(--ink);font-size:11px;font-weight:650}} .dag-node .dag-status{{fill:var(--muted);font-size:9px;letter-spacing:.08em}} .dag-node.completed circle{{fill:var(--green)}} .dag-node.completed rect{{stroke:#35d07f99}} .dag-node.failed circle{{fill:var(--red)}} .dag-node.failed rect{{stroke:#ff657799}} .dag-node.pending circle{{fill:var(--amber)}} .dag-node.pending rect{{stroke:#ffbd4577}} .dag-node.disabled{{opacity:.38}} .stage-dag a:hover .dag-node rect{{fill:#101d30;stroke:var(--blue)}} .dag-legend{{display:flex;gap:16px;flex-wrap:wrap;color:var(--muted);font-size:12px;margin-top:5px}} .dag-legend span::before{{content:'';display:inline-block;width:8px;height:8px;border-radius:50%;margin-right:6px;background:#55647a}} .dag-legend .completed::before{{background:var(--green)}} .dag-legend .failed::before{{background:var(--red)}} .dag-legend .pending::before{{background:var(--amber)}} .dag-legend .disabled::before{{background:#55647a;opacity:.45}} .dag-type-legend span::before{{width:13px;height:9px;border-radius:2px;background:transparent;border:1.5px solid #93a4bd}} .dag-type-legend .optional-node::before{{border-style:dashed}}
+@keyframes pulse{{50%{{transform:scale(1.7);box-shadow:0 0 18px #ffbd45aa}}}} .table-wrap{{overflow:auto}} table{{width:100%;border-collapse:collapse}} th,td{{padding:10px 12px;border-bottom:1px solid var(--line);text-align:right}} th:first-child{{text-align:left}} thead th{{color:#a9bee0;background:#0a1421;position:sticky;top:0}} tr.selected-candidate{{background:#4f8cff18}} tr.selected-candidate td{{box-shadow:inset 0 1px #4f8cff44,inset 0 -1px #4f8cff44}} .candidate-swatch{{display:inline-block;width:10px;height:10px;margin-right:8px;border-radius:50%}} td.warning-cell{{background:#ffbd4524;color:#ffe3a3;box-shadow:inset 0 0 0 1px #ffbd4570}} .warning-value{{cursor:help;outline-offset:3px}} #warning-tooltip{{position:fixed;z-index:10000;max-width:min(420px,calc(100vw - 24px));padding:9px 11px;border:1px solid #ffbd4588;border-radius:8px;background:#241b0d;color:#ffe3a3;font-size:12px;line-height:1.4;box-shadow:0 10px 30px #0009;pointer-events:none;white-space:pre-wrap}} .sanity-gates{{display:flex;flex-wrap:wrap;gap:8px;margin:0 0 14px}} .gate{{display:inline-block;border-radius:999px;padding:5px 10px;background:#1b2839}} .gate.passed{{color:var(--green)}} .gate.warning{{color:var(--amber)}} .gate.failed{{color:var(--red)}} .empty,.note{{color:var(--muted)}} select{{margin:0 18px 18px 8px;padding:8px 12px;border:1px solid var(--line);border-radius:8px;background:#0a1421;color:var(--ink)}} .selector-label{{color:var(--muted)}} code{{color:#bcd2ff}} .probe-summaries{{display:grid;grid-template-columns:repeat(auto-fit,minmax(260px,1fr));gap:12px;margin:16px 0}} .probe-summary{{border:1px solid var(--line);border-radius:12px;padding:14px;background:#0a1421}} .probe-summary.passed{{border-color:#35d07f66}} .probe-summary.failed{{border-color:#ff657766}} .probe-summary h3{{margin:0 0 10px}} .probe-summary dl{{display:grid;grid-template-columns:1fr auto;gap:5px 14px;margin:0}} .probe-summary dt{{color:var(--muted)}} .probe-summary dd{{margin:0;text-align:right}} .probe-plots{{display:grid;grid-template-columns:repeat(auto-fit,minmax(430px,1fr));gap:14px}} .probe-plot-panel{{min-width:0;border:1px solid var(--line);border-radius:12px;padding:12px;background:#091321}} .probe-plot-panel h3{{margin:0 0 4px}} .plotly-chart{{width:100%;height:390px}} .depth-chart{{height:460px}}
 .vllm-overview-cards{{display:grid;grid-template-columns:repeat(auto-fit,minmax(170px,1fr));gap:10px;margin:16px 0}} .vllm-overview-cards article{{display:flex;flex-direction:column;padding:14px;border:1px solid var(--line);border-radius:10px;background:#0a1421}} .vllm-overview-cards strong{{font-size:24px;color:#cfe0ff}} .vllm-overview-cards span{{color:var(--muted)}} .vllm-panel{{margin:14px 0}} .vllm-controls{{display:grid;grid-template-columns:repeat(auto-fit,minmax(230px,1fr));gap:10px;margin:14px 0}} .vllm-controls>span{{display:contents}} .vllm-controls label{{display:flex;flex-direction:column;gap:5px;color:var(--muted)}} .vllm-controls select{{margin:0}}
 .vllm-connect-toggle{{display:flex;align-items:center;gap:7px;margin:8px 0 12px;color:var(--muted)}} .vllm-connect-toggle input{{accent-color:#4f8cff}}
 .replacement-layer-grid{{display:grid;grid-template-columns:repeat(auto-fit,minmax(105px,1fr));gap:6px;margin:12px 0;padding:10px;border:1px solid var(--line);border-radius:10px;background:#091321}} .layer-toggle,.replacement-connect-toggle,.replacement-all-toggle{{display:flex;align-items:center;gap:6px;color:var(--muted);font-size:13px}} .layer-toggle input,.replacement-connect-toggle input,.replacement-all-toggle input{{accent-color:#4f8cff}} .replacement-connect-toggle{{margin:8px 0}} .replacement-layer-toolbar{{display:flex;align-items:center;justify-content:space-between;gap:18px;flex-wrap:wrap;margin:10px 0}} .replacement-all-toggle{{font-weight:600;color:var(--ink)}} .replacement-layer-color-key{{display:flex;align-items:center;gap:8px;color:var(--muted);font-size:12px}} .replacement-layer-color-key i{{display:block;width:min(280px,32vw);height:8px;border-radius:999px;background:linear-gradient(90deg,#ff6577,#4f8cff)}}
diff --git a/modelopt/torch/puzzletron/diagnostics/sanity_verdict.py b/modelopt/torch/puzzletron/diagnostics/sanity_verdict.py
index 2960a202b6c..c1d08273497 100644
--- a/modelopt/torch/puzzletron/diagnostics/sanity_verdict.py
+++ b/modelopt/torch/puzzletron/diagnostics/sanity_verdict.py
@@ -6,21 +6,35 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Any
+from typing import Any, Literal
 
 from ..manifest import StageManifest
 from ..stage_runner import StageResult
 from ..stages.common import complete_stage
 
-__all__ = ["SanityVerdict", "complete_sanity_stage", "finding_from_message"]
+__all__ = [
+    "SanityVerdict",
+    "complete_sanity_stage",
+    "finding_from_message",
+    "is_correctness_sanity_stage",
+]
+
+_CORRECTNESS_STAGES = frozenset({"sort_sanity", "slicing_sanity"})
+
+
+def is_correctness_sanity_stage(stage: str) -> bool:
+    """Return whether a failed stage verdict is a blocking correctness failure."""
+
+    return stage in _CORRECTNESS_STAGES
 
 
 @dataclass
 class SanityVerdict:
-    """Quality outcome for one sanity stage execution."""
+    """Observed outcome for one sanity stage execution."""
 
     passed: bool
     findings: list[dict[str, Any]] = field(default_factory=list)
+    blocking: bool = False
     message: str | None = None
 
 
@@ -29,12 +43,13 @@ def finding_from_message(
     stage: str,
     message: str,
     evidence: dict[str, Any] | None = None,
+    severity: Literal["warning", "error"] = "warning",
 ) -> dict[str, Any]:
     return {
         "stage": stage,
         "message": message,
         "evidence": evidence or {},
-        "severity": "warning",
+        "severity": severity,
     }
 
 
@@ -46,14 +61,36 @@ def complete_sanity_stage(
     verdict: SanityVerdict,
     message: str | None = None,
 ) -> StageResult:
-    """Complete a sanity stage according to the global warning policy."""
+    """Complete a sanity stage without downgrading correctness failures."""
 
     merged = dict(outputs or {})
+    correctness_failure = not verdict.passed and (
+        verdict.blocking or is_correctness_sanity_stage(manifest.stage)
+    )
+    findings = []
+    for finding in verdict.findings:
+        normalized = dict(finding)
+        finding_stage = str(normalized.get("stage", manifest.stage))
+        if correctness_failure and (
+            is_correctness_sanity_stage(manifest.stage)
+            or is_correctness_sanity_stage(finding_stage)
+        ):
+            normalized["severity"] = "error"
+        else:
+            normalized.setdefault("severity", "warning")
+        findings.append(normalized)
     merged["passed"] = verdict.passed
-    merged["findings"] = list(verdict.findings)
-    merged["verdict"] = "passed" if verdict.passed else "warning"
+    merged["findings"] = findings
+    merged["blocking"] = correctness_failure
+    merged["verdict"] = (
+        "passed" if verdict.passed else "failed" if correctness_failure else "warning"
+    )
     fail_on_warnings = bool((config.get("sanity") or {}).get("fail_on_warnings", False))
-    status = "failed" if fail_on_warnings and not verdict.passed else "success"
+    status = (
+        "failed"
+        if correctness_failure or (fail_on_warnings and not verdict.passed)
+        else "success"
+    )
     return complete_stage(
         config,
         manifest,
diff --git a/modelopt/torch/puzzletron/diagnostics/width_sanity.py b/modelopt/torch/puzzletron/diagnostics/width_sanity.py
index 556ec856c26..e1d4a54597b 100644
--- a/modelopt/torch/puzzletron/diagnostics/width_sanity.py
+++ b/modelopt/torch/puzzletron/diagnostics/width_sanity.py
@@ -101,6 +101,7 @@ def descriptor_realization_findings(
                         "right_method": "physical",
                         "delta": delta,
                     },
+                    severity="error",
                 )
             )
     return findings
diff --git a/modelopt/torch/puzzletron/stages/diagnostics.py b/modelopt/torch/puzzletron/stages/diagnostics.py
index 2bdf12389fb..4f4fb8d2637 100644
--- a/modelopt/torch/puzzletron/stages/diagnostics.py
+++ b/modelopt/torch/puzzletron/stages/diagnostics.py
@@ -1607,6 +1607,27 @@ def _merge_reused_sort_equivalence(
     return merged
 
 
+def _parent_sweep_sanity_verdict(
+    width_summary: dict[str, Any], sort_summary: dict[str, Any]
+):
+    """Combine advisory width quality with blocking reused-sort correctness."""
+
+    from ..diagnostics.sanity_verdict import SanityVerdict
+
+    width_findings = list(width_summary.get("findings") or ())
+    sort_findings = [
+        {**finding, "stage": "sort_sanity", "severity": "error"}
+        for finding in sort_summary.get("findings") or ()
+    ]
+    width_passed = bool(width_summary.get("passed", not width_findings))
+    sort_passed = sort_summary.get("passed") is True
+    return SanityVerdict(
+        passed=width_passed and sort_passed,
+        findings=[*width_findings, *sort_findings],
+        blocking=not sort_passed,
+    )
+
+
 def _extract_rows(method: str, output_dir: Path) -> list[dict[str, Any]]:
     rows = []
     for result_path in sorted(output_dir.glob("solution_*.json")):
@@ -2046,7 +2067,13 @@ def _publish_parent_sweep_sanity(
         ("slicing_sanity", slicing_summary),
     ):
         payload["passed"] = not payload.get("findings")
-        payload["verdict"] = "passed" if payload["passed"] else "warning"
+        payload["verdict"] = (
+            "passed"
+            if payload["passed"]
+            else "failed"
+            if stage == "slicing_sanity"
+            else "warning"
+        )
         payload["provenance"] = provenance
         output = puzzle_dir / "artifacts" / stage / "summary.json"
         output.parent.mkdir(parents=True, exist_ok=True)
@@ -2657,7 +2684,10 @@ def _activation_diagnostic_parent_sweep(
             activation_equivalence = (
                 (sweep_manifest.get("parents") or {}).get("activation") or {}
             ).get("equivalence") or {}
-            equivalence_findings = list(activation_equivalence.get("findings") or ())
+            equivalence_findings = [
+                {**finding, "stage": "sort_sanity", "severity": "error"}
+                for finding in activation_equivalence.get("findings") or ()
+            ]
             sort_passed = activation_equivalence.get("passed") is True
             sort_equivalence_dir = puzzle_dir / "artifacts" / "sort_sanity"
             sort_equivalence_dir.mkdir(parents=True, exist_ok=True)
@@ -2672,6 +2702,8 @@ def _activation_diagnostic_parent_sweep(
                 "sorted_teacher_dir": str(sorted_dir),
                 "equivalence": activation_equivalence,
                 "findings": equivalence_findings,
+                "blocking": not sort_passed,
+                "verdict": "passed" if sort_passed else "failed",
                 "parent_sweep_manifest": str(load_manifest_path),
             }
             sort_summary_path.write_text(
@@ -2706,8 +2738,9 @@ def _activation_diagnostic_parent_sweep(
 
     width_summary_path = puzzle_dir / "artifacts" / "width_sanity" / "summary.json"
     width_verdict = json.loads(width_summary_path.read_text(encoding="utf-8"))
-    width_findings = list(width_verdict.get("findings") or ())
-    from ..diagnostics.sanity_verdict import SanityVerdict, complete_sanity_stage
+    sort_summary_path = puzzle_dir / "artifacts" / "sort_sanity" / "summary.json"
+    sort_verdict = json.loads(sort_summary_path.read_text(encoding="utf-8"))
+    from ..diagnostics.sanity_verdict import complete_sanity_stage
 
     return complete_sanity_stage(
         config,
@@ -2733,10 +2766,7 @@ def _activation_diagnostic_parent_sweep(
                 puzzle_dir / "artifacts" / "slicing_sanity" / "summary.json"
             ),
         },
-        verdict=SanityVerdict(
-            passed=bool(width_verdict.get("passed", not width_findings)),
-            findings=width_findings,
-        ),
+        verdict=_parent_sweep_sanity_verdict(width_verdict, sort_verdict),
     )
 
 
@@ -3310,7 +3340,7 @@ def sort_equivalence_stage(config: dict[str, Any], manifest: StageManifest):
             "max_abs_reverse_delta": reverse_tolerance,
             "passed": passed,
             "findings": [],
-            "verdict": "passed" if passed else "warning",
+            "verdict": "passed" if passed else "failed",
             "teacher_result": str(scoring_output_dir / "teacher.json"),
             "sorted_result": str(sorted_result_path),
             "reverse_sorted_dir": str(reverse_dir) if include_reverse else None,
@@ -3365,6 +3395,7 @@ def sort_equivalence_stage(config: dict[str, Any], manifest: StageManifest):
                             f"tolerance={tolerance:.6g}"
                         ),
                         evidence={"metric": metric, "delta": delta, "tolerance": tolerance},
+                        severity="error",
                     )
                 )
             if include_reverse and not reverse_passed:
@@ -3381,6 +3412,7 @@ def sort_equivalence_stage(config: dict[str, Any], manifest: StageManifest):
                             "reverse_delta": reverse_delta,
                             "tolerance": reverse_tolerance,
                         },
+                        severity="error",
                     )
                 )
             summary["findings"] = findings
@@ -3560,7 +3592,7 @@ def width_slice_equivalence_stage(config: dict[str, Any], manifest: StageManifes
             "stage": "slicing_sanity",
             "message": f"width-slice equivalence failed for case {case.get('case_id')}",
             "evidence": {"case": case},
-            "severity": "warning",
+            "severity": "error",
         }
         for case in summary.get("cases", ())
         if not case.get("passed", True)
diff --git a/tests/unit/torch/puzzletron/test_campaign_findings.py b/tests/unit/torch/puzzletron/test_campaign_findings.py
index f86ad14d4a2..2da32e84ab6 100644
--- a/tests/unit/torch/puzzletron/test_campaign_findings.py
+++ b/tests/unit/torch/puzzletron/test_campaign_findings.py
@@ -12,7 +12,7 @@
 )
 
 
-def test_equivalence_warning_is_derived_from_values_not_axis_name():
+def test_equivalence_error_is_derived_from_values_not_axis_name():
     rows = [
         {"axis": "arbitrary_axis", "method": "dynamic", "lm_loss": 4.9582},
         {"axis": "arbitrary_axis", "method": "physical", "lm_loss": 4.9697},
@@ -29,6 +29,7 @@ def test_equivalence_warning_is_derived_from_values_not_axis_name():
 
     assert findings[0].evidence["delta"] == pytest.approx(0.0115)
     assert findings[0].evidence["group"] == {"axis": "arbitrary_axis"}
+    assert findings[0].severity == "error"
     assert "arbitrary_axis" not in findings[0].message
 
 
@@ -52,6 +53,8 @@ def test_ranking_findings_obey_metric_direction(direction, preferred, comparison
     )
 
     assert bool(findings) is warns
+    if findings:
+        assert findings[0].severity == "warning"
 
 
 def test_loss_trend_warns_when_ending_window_does_not_improve():
diff --git a/tests/unit/torch/puzzletron/test_campaign_progress_report.py b/tests/unit/torch/puzzletron/test_campaign_progress_report.py
index d4837a67dad..a1d5188ef79 100644
--- a/tests/unit/torch/puzzletron/test_campaign_progress_report.py
+++ b/tests/unit/torch/puzzletron/test_campaign_progress_report.py
@@ -186,11 +186,15 @@ def test_progress_report_renders_canonical_sort_sanity_metrics(tmp_path: Path):
     assert 'data-stage="sort_sanity" data-status="completed"' in document
 
 
-def test_sort_sanity_failure_renders_warning_without_failed_dag_node(tmp_path: Path):
+def test_sort_sanity_failure_renders_blocking_failure_and_failed_dag_node(tmp_path: Path):
     message = "sorted teacher loss drift exceeded tolerance"
     _write(
         tmp_path / "manifests/sort_sanity.json",
-        {"config": {"sort_sanity": {"enabled": True}}},
+        {
+            "status": "failed",
+            "config": {"sort_sanity": {"enabled": True}},
+            "outputs": {"passed": False, "verdict": "failed", "blocking": True},
+        },
     )
     _write(
         tmp_path / "artifacts/sort_sanity/summary.json",
@@ -201,7 +205,7 @@ def test_sort_sanity_failure_renders_warning_without_failed_dag_node(tmp_path: P
             "findings": [
                 {
                     "stage": "sort_sanity",
-                    "severity": "warning",
+                    "severity": "error",
                     "message": message,
                     "evidence": {"metric": "lm_loss"},
                 }
@@ -212,10 +216,10 @@ def test_sort_sanity_failure_renders_warning_without_failed_dag_node(tmp_path: P
     result = generate_campaign_progress_report(tmp_path)
     document = Path(result["html"]).read_text(encoding="utf-8")
 
-    assert "Equivalence gate: warning" in document
+    assert "Equivalence gate: failed (blocking correctness)" in document
     assert "warning-value" in document
     assert f"data-warning='{message}'" in document
-    assert 'data-stage="sort_sanity" data-status="completed"' in document
+    assert 'data-stage="sort_sanity" data-status="failed"' in document
 
 
 def test_progress_report_uses_pending_instead_of_transient_running_state(tmp_path: Path):
@@ -257,18 +261,36 @@ def test_width_and_slicing_findings_render_on_affected_cells(tmp_path: Path):
     _write(
         tmp_path / "artifacts/width_sanity/summary.json",
         {
+            "passed": False,
             "rows": [
                 {**common, "method": "sorted", "raw_replacement_loss": 1.2},
                 {**common, "method": "original", "raw_replacement_loss": 1.1},
                 {**common, "method": "reverse", "raw_replacement_loss": 1.3},
             ],
-            "findings": [],
+            "findings": [
+                {
+                    "stage": "width_sanity",
+                    "message": "sorted ranking is worse than original.",
+                    "severity": "warning",
+                    "evidence": {
+                        "group": {
+                            "axis": "arbitrary_axis",
+                            "layer_idx": 3,
+                            "target_value": 8,
+                        },
+                        "metric": "raw_replacement_loss",
+                        "preferred_method": "sorted",
+                        "comparison_method": "original",
+                    },
+                }
+            ],
         },
     )
     message = "sorted and physical differ for raw_replacement_loss."
     _write(
         tmp_path / "artifacts/slicing_sanity/summary.json",
         {
+            "passed": False,
             "rows": [
                 {**common, "method": "sorted", "raw_replacement_loss": 1.2},
                 {**common, "method": "physical", "raw_replacement_loss": 1.0},
@@ -277,7 +299,7 @@ def test_width_and_slicing_findings_render_on_affected_cells(tmp_path: Path):
                 {
                     "stage": "slicing_sanity",
                     "message": message,
-                    "severity": "warning",
+                    "severity": "error",
                     "evidence": {
                         "group": {
                             "axis": "arbitrary_axis",
@@ -296,6 +318,9 @@ def test_width_and_slicing_findings_render_on_affected_cells(tmp_path: Path):
     result = generate_campaign_progress_report(tmp_path)
     document = Path(result["html"]).read_text(encoding="utf-8")
 
+    assert "Width ranking: warning (advisory)" in document
+    assert "Dynamic/physical equivalence: failed (blocking correctness)" in document
+    assert 'data-stage="slicing_sanity" data-status="failed"' in document
     assert "class='warning-cell'" in document
     assert "class='warning-value'" in document
     assert "tabindex='0'" in document
diff --git a/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py b/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py
index 696a1425fe9..4a7379db886 100644
--- a/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py
+++ b/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py
@@ -13,12 +13,6 @@
     _validate_worker_topology,
     _worker_config,
 )
-from modelopt.torch.puzzletron.diagnostics.sanity_verdict import (
-    SanityVerdict,
-    complete_sanity_stage,
-    finding_from_message,
-)
-from modelopt.torch.puzzletron.manifest import StageManifest
 from modelopt.torch.puzzletron.stages import diagnostics
 from modelopt.torch.puzzletron.stages.diagnostics import _scoring_cfg_for_method
 
@@ -100,26 +94,6 @@ def test_sort_equivalence_keeps_production_and_reverse_control_tolerances_separa
     )["passed"]
 
 
-def test_failed_sort_verdict_completes_manifest_with_warning(tmp_path: Path):
-    config = {"experiment": {"dir": str(tmp_path)}}
-    manifest = StageManifest(stage="sort_sanity", config=config)
-    finding = finding_from_message(
-        stage="sort_sanity",
-        message="sorted teacher drift exceeded tolerance",
-    )
-
-    result = complete_sanity_stage(
-        config,
-        manifest,
-        verdict=SanityVerdict(passed=False, findings=[finding]),
-    )
-
-    assert result.status == "success"
-    assert manifest.status == "success"
-    assert manifest.outputs["passed"] is False
-    assert manifest.outputs["verdict"] == "warning"
-
-
 def test_axis_worker_preserves_requested_layers_and_targets_per_axis(tmp_path: Path):
     parallel = {
         "tp": 1,
diff --git a/tests/unit/torch/puzzletron/test_hidden_width_diagnostic.py b/tests/unit/torch/puzzletron/test_hidden_width_diagnostic.py
index ecb939a7a1f..d17e1ae4796 100644
--- a/tests/unit/torch/puzzletron/test_hidden_width_diagnostic.py
+++ b/tests/unit/torch/puzzletron/test_hidden_width_diagnostic.py
@@ -7,6 +7,7 @@
     _hidden_width_result_metrics,
     _merge_reused_sort_equivalence,
     _near_teacher_axis_targets,
+    _parent_sweep_sanity_verdict,
     _ratio_aligned_hidden_widths,
     _select_diagnostic_hidden_width,
     _select_layers,
@@ -219,3 +220,23 @@ def test_reused_parent_sweep_preserves_existing_sort_diagnosis_metrics():
     assert merged["sorted_teacher"] == existing["sorted_teacher"]
     assert merged["reverse_sorted"] == existing["reverse_sorted"]
     assert merged["reused_parent_sweep"] is True
+
+
+def test_parent_sweep_sort_miss_is_blocking_but_width_miss_remains_advisory():
+    verdict = _parent_sweep_sanity_verdict(
+        {
+            "passed": False,
+            "findings": [{"stage": "width_sanity", "message": "ranking regressed"}],
+        },
+        {
+            "passed": False,
+            "findings": [{"stage": "width_sanity", "message": "teacher drifted"}],
+        },
+    )
+
+    assert verdict.passed is False
+    assert verdict.blocking is True
+    assert verdict.findings == [
+        {"stage": "width_sanity", "message": "ranking regressed"},
+        {"stage": "sort_sanity", "message": "teacher drifted", "severity": "error"},
+    ]
diff --git a/tests/unit/torch/puzzletron/test_sanity_verdict.py b/tests/unit/torch/puzzletron/test_sanity_verdict.py
index 1900b9915ca..851b8ff2d81 100644
--- a/tests/unit/torch/puzzletron/test_sanity_verdict.py
+++ b/tests/unit/torch/puzzletron/test_sanity_verdict.py
@@ -13,21 +13,21 @@
 from modelopt.torch.puzzletron.manifest import StageManifest
 
 
-def test_complete_sanity_stage_allows_warnings_by_default(tmp_path: Path):
+def test_complete_sanity_stage_allows_advisory_warnings_by_default(tmp_path: Path):
     config = {"experiment": {"dir": str(tmp_path)}}
-    manifest = StageManifest(stage="sort_sanity", config=config)
+    manifest = StageManifest(stage="width_sanity", config=config)
     (tmp_path / "manifests").mkdir(parents=True)
 
     result = complete_sanity_stage(
         config,
         manifest,
-        outputs={"summary_path": "artifacts/sort_sanity/summary.json"},
+        outputs={"summary_path": "artifacts/width_sanity/summary.json"},
         verdict=SanityVerdict(
             passed=False,
             findings=[
                 finding_from_message(
-                    stage="sort_sanity",
-                    message="sorted teacher drift too large",
+                    stage="width_sanity",
+                    message="activation ranking is worse than reverse",
                 )
             ],
         ),
@@ -37,15 +37,16 @@ def test_complete_sanity_stage_allows_warnings_by_default(tmp_path: Path):
     assert manifest.status == "success"
     assert manifest.outputs["passed"] is False
     assert manifest.outputs["verdict"] == "warning"
+    assert manifest.outputs["blocking"] is False
     assert manifest.outputs["findings"]
 
 
-def test_complete_sanity_stage_fails_when_warnings_are_strict(tmp_path: Path):
+def test_complete_sanity_stage_fails_advisory_warnings_when_strict(tmp_path: Path):
     config = {
         "experiment": {"dir": str(tmp_path)},
         "sanity": {"fail_on_warnings": True},
     }
-    manifest = StageManifest(stage="bypass_sanity", config=config)
+    manifest = StageManifest(stage="width_sanity", config=config)
     (tmp_path / "manifests").mkdir(parents=True)
 
     result = complete_sanity_stage(
@@ -55,8 +56,8 @@ def test_complete_sanity_stage_fails_when_warnings_are_strict(tmp_path: Path):
             passed=False,
             findings=[
                 finding_from_message(
-                    stage="bypass_sanity",
-                    message="overfit probe did not improve",
+                    stage="width_sanity",
+                    message="activation ranking is worse than reverse",
                 )
             ],
         ),
@@ -66,9 +67,94 @@ def test_complete_sanity_stage_fails_when_warnings_are_strict(tmp_path: Path):
     assert manifest.status == "failed"
     assert manifest.outputs["passed"] is False
     assert manifest.outputs["verdict"] == "warning"
+    assert manifest.outputs["blocking"] is False
     assert manifest.outputs["findings"]
 
 
+def test_complete_sanity_stage_fails_sort_correctness_by_default(tmp_path: Path):
+    config = {"experiment": {"dir": str(tmp_path)}}
+    manifest = StageManifest(stage="sort_sanity", config=config)
+    (tmp_path / "manifests").mkdir(parents=True)
+
+    result = complete_sanity_stage(
+        config,
+        manifest,
+        verdict=SanityVerdict(
+            passed=False,
+            findings=[
+                finding_from_message(
+                    stage="sort_sanity",
+                    message="sorted teacher drift too large",
+                )
+            ],
+        ),
+    )
+
+    assert result.status == "failed"
+    assert manifest.status == "failed"
+    assert manifest.outputs["verdict"] == "failed"
+    assert manifest.outputs["blocking"] is True
+    assert manifest.outputs["findings"][0]["severity"] == "error"
+
+
+def test_warning_policy_cannot_downgrade_slicing_correctness_failure(tmp_path: Path):
+    config = {
+        "experiment": {"dir": str(tmp_path)},
+        "sanity": {"fail_on_warnings": False},
+    }
+    manifest = StageManifest(stage="slicing_sanity", config=config)
+    (tmp_path / "manifests").mkdir(parents=True)
+
+    result = complete_sanity_stage(
+        config,
+        manifest,
+        verdict=SanityVerdict(
+            passed=False,
+            findings=[
+                finding_from_message(
+                    stage="slicing_sanity",
+                    message="dynamic and physical slices disagree",
+                )
+            ],
+        ),
+    )
+
+    assert result.status == "failed"
+    assert manifest.status == "failed"
+    assert manifest.outputs["verdict"] == "failed"
+    assert manifest.outputs["blocking"] is True
+
+
+def test_folded_correctness_failure_preserves_advisory_finding_severity(tmp_path: Path):
+    config = {"experiment": {"dir": str(tmp_path)}}
+    manifest = StageManifest(stage="width_sanity", config=config)
+    (tmp_path / "manifests").mkdir(parents=True)
+
+    result = complete_sanity_stage(
+        config,
+        manifest,
+        verdict=SanityVerdict(
+            passed=False,
+            blocking=True,
+            findings=[
+                finding_from_message(
+                    stage="width_sanity", message="ranking quality regressed"
+                ),
+                finding_from_message(
+                    stage="sort_sanity", message="sorted teacher drifted"
+                ),
+            ],
+        ),
+    )
+
+    assert result.status == "failed"
+    assert manifest.outputs["blocking"] is True
+    assert [finding["severity"] for finding in manifest.outputs["findings"]] == [
+        "warning",
+        "error",
+    ]
+
+
 def test_complete_sanity_stage_strict_policy_does_not_fail_passed_verdict(tmp_path: Path):
     config = {
         "experiment": {"dir": str(tmp_path)},
@@ -92,3 +178,13 @@ def test_finding_from_message_shape():
     assert finding["stage"] == "width_sanity"
     assert finding["severity"] == "warning"
     assert finding["evidence"]["x"] == 1
+
+
+def test_finding_from_message_accepts_correctness_severity():
+    finding = finding_from_message(
+        stage="sort_sanity",
+        message="example",
+        severity="error",
+    )
+
+    assert finding["severity"] == "error"
diff --git a/tests/unit/torch/puzzletron/test_width_sanity_aggregation.py b/tests/unit/torch/puzzletron/test_width_sanity_aggregation.py
index 482c8b4d7e7..2de72a70df9 100644
--- a/tests/unit/torch/puzzletron/test_width_sanity_aggregation.py
+++ b/tests/unit/torch/puzzletron/test_width_sanity_aggregation.py
@@ -216,7 +216,7 @@ def test_parent_sweep_publication_accepts_per_metric_physical_tolerances(tmp_pat
     }
 
 
-def test_parent_sweep_physical_miss_is_published_as_warning(tmp_path):
+def test_parent_sweep_physical_miss_is_published_as_correctness_failure(tmp_path):
     parent_summary = {
         "rows": [
             {
@@ -244,5 +244,6 @@ def test_parent_sweep_physical_miss_is_published_as_warning(tmp_path):
 
     summary = json.loads(slicing_path.read_text())
     assert summary["passed"] is False
+    assert summary["verdict"] == "failed"
     assert summary["findings"]
-    assert all(finding["severity"] == "warning" for finding in summary["findings"])
+    assert all(finding["severity"] == "error" for finding in summary["findings"])
diff --git a/tests/unit/torch/puzzletron/test_width_slice_equivalence.py b/tests/unit/torch/puzzletron/test_width_slice_equivalence.py
index 900291e3d51..796a4676ff2 100644
--- a/tests/unit/torch/puzzletron/test_width_slice_equivalence.py
+++ b/tests/unit/torch/puzzletron/test_width_slice_equivalence.py
@@ -944,7 +944,7 @@ def multimodal_loader(args, *, checkpoint_dir, data_layout):
     assert selected["data_layout"] == "fixed"
 
 
-def test_distributed_slicing_verdict_warns_without_failing_stage(
+def test_distributed_slicing_verdict_fails_closed_and_normalizes_error(
     tmp_path: Path, monkeypatch: pytest.MonkeyPatch
 ):
     summary_path = tmp_path / "artifacts" / "slicing_sanity" / "summary.json"
@@ -980,7 +980,9 @@ def test_distributed_slicing_verdict_warns_without_failing_stage(
         manifest,
     )
 
-    assert result.status == "success"
-    assert manifest.status == "success"
+    assert result.status == "failed"
+    assert manifest.status == "failed"
     assert manifest.outputs["passed"] is False
-    assert manifest.outputs["findings"] == [finding]
+    assert manifest.outputs["verdict"] == "failed"
+    assert manifest.outputs["blocking"] is True
+    assert manifest.outputs["findings"] == [{**finding, "severity": "error"}]

From dab1cbd343b15d2e3607e4fa6fb72cd5e7c68500 Mon Sep 17 00:00:00 2001
From: Johannes Rausch <jrausch@nvidia.com>
Date: Fri, 7 Aug 2026 16:25:45 +0200
Subject: [PATCH 2/6] Clarify Puzzletron sanity qualification

Explain how width-ranking quality checks differ from sort and slicing equivalence failures, including what is measured and how strict qualification can promote warnings.

Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
---
 examples/puzzletron/README.md                 | 108 ++++++++++++++++++
 examples/puzzletron/docs/v2_architecture.md   |  30 ++++-
 .../diagnostics/campaign_findings.py          |   9 +-
 .../diagnostics/campaign_progress_report.py   |   8 +-
 .../puzzletron/diagnostics/sanity_verdict.py  |   7 +-
 .../test_campaign_progress_report.py          |   4 +-
 6 files changed, 155 insertions(+), 11 deletions(-)

diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md
index d5149969d13..2d299eac9d4 100644
--- a/examples/puzzletron/README.md
+++ b/examples/puzzletron/README.md
@@ -481,6 +481,114 @@ 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. |
 
+## Understand sorting, width ranking, and slicing
+
+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.
+
+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 or unsorted | Keep the same-sized prefix in the teacher's original order. Some older paths call this the random control, although the original prefix is deterministic. |
+| 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. |
+
+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.
+
+### What does width sanity measure?
+
+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.
+
+### What is a width-ranking miss?
+
+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. Loss-like
+metrics are lower-is-better; accuracy-like metrics are higher-is-better. 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.
+
+### What is an equivalence failure?
+
+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.
+
+### How stage completion and qualification differ
+
+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
+```
+
+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. See
+[`docs/v2_architecture.md`](docs/v2_architecture.md#semantic-validation-gates)
+for the broader design context.
+
 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/v2_architecture.md b/examples/puzzletron/docs/v2_architecture.md
index cd285dcfcee..c6bce52ed87 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.
 
@@ -139,7 +139,7 @@ flowchart TB
   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.
+  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.
@@ -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 reverse and unsorted/random 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
+[Puzzletron README](../README.md#understand-sorting-width-ranking-and-slicing)
+provides definitions and a worked example for new users.
+
 ## 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 9fbd4465c86..03cdaf49b48 100644
--- a/modelopt/torch/puzzletron/diagnostics/campaign_findings.py
+++ b/modelopt/torch/puzzletron/diagnostics/campaign_findings.py
@@ -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():
@@ -127,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 898fc094a1e..80bf1caf7f1 100644
--- a/modelopt/torch/puzzletron/diagnostics/campaign_progress_report.py
+++ b/modelopt/torch/puzzletron/diagnostics/campaign_progress_report.py
@@ -488,7 +488,7 @@ def _activation_diagnostic_section(summary: dict[str, Any]) -> str:
         gates.append(
             "<p class='gate {}'>Width ranking: {}</p>".format(
                 "passed" if width_passed else "warning",
-                "passed" if width_passed else "warning (advisory)",
+                "passed" if width_passed else "quality warning",
             )
         )
     if summary.get("slicing_present"):
@@ -581,7 +581,11 @@ def _activation_diagnostic_section(summary: dict[str, Any]) -> str:
     return gate_summary + (
         "<p class='note'>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.</p>"
+        "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.</p>"
         "<label class='selector-label' for='activation-axis-select'>Swept axis</label>"
         f'<select id="activation-axis-select">{axis_options}</select>'
         "<label class='selector-label' for='activation-metric-select'>Metric</label>"
diff --git a/modelopt/torch/puzzletron/diagnostics/sanity_verdict.py b/modelopt/torch/puzzletron/diagnostics/sanity_verdict.py
index c1d08273497..410a4686481 100644
--- a/modelopt/torch/puzzletron/diagnostics/sanity_verdict.py
+++ b/modelopt/torch/puzzletron/diagnostics/sanity_verdict.py
@@ -61,7 +61,12 @@ def complete_sanity_stage(
     verdict: SanityVerdict,
     message: str | None = None,
 ) -> StageResult:
-    """Complete a sanity stage without downgrading correctness failures."""
+    """Complete a sanity stage using separate correctness and warning policies.
+
+    Correctness failures always fail the stage. Other failed verdicts remain
+    quality warnings unless ``sanity.fail_on_warnings`` promotes them to a
+    failed stage; campaign qualification remains a caller-owned policy.
+    """
 
     merged = dict(outputs or {})
     correctness_failure = not verdict.passed and (
diff --git a/tests/unit/torch/puzzletron/test_campaign_progress_report.py b/tests/unit/torch/puzzletron/test_campaign_progress_report.py
index a1d5188ef79..556df6311bc 100644
--- a/tests/unit/torch/puzzletron/test_campaign_progress_report.py
+++ b/tests/unit/torch/puzzletron/test_campaign_progress_report.py
@@ -318,8 +318,10 @@ def test_width_and_slicing_findings_render_on_affected_cells(tmp_path: Path):
     result = generate_campaign_progress_report(tmp_path)
     document = Path(result["html"]).read_text(encoding="utf-8")
 
-    assert "Width ranking: warning (advisory)" in document
+    assert "Width ranking: quality warning" in document
     assert "Dynamic/physical equivalence: failed (blocking correctness)" in document
+    assert "it does not mean the dynamic and physical implementations disagree" in document
+    assert "Campaign qualification may still require the ranking warning to pass" in document
     assert 'data-stage="slicing_sanity" data-status="failed"' in document
     assert "class='warning-cell'" in document
     assert "class='warning-value'" in document

From 7ca3cd27563446e477fe5bd5fc8a53900cf1e55b Mon Sep 17 00:00:00 2001
From: Johannes Rausch <jrausch@nvidia.com>
Date: Fri, 7 Aug 2026 17:12:14 +0200
Subject: [PATCH 3/6] Move Puzzletron sanity details into guide

Keep the main README focused on the first-run workflow while retaining the detailed ranking, slicing, and qualification reference in a linked topic guide.

Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
---
 examples/puzzletron/README.md                 | 109 ++-------------
 examples/puzzletron/docs/sanity_validation.md | 124 ++++++++++++++++++
 examples/puzzletron/docs/v2_architecture.md   |   4 +-
 3 files changed, 137 insertions(+), 100 deletions(-)
 create mode 100644 examples/puzzletron/docs/sanity_validation.md

diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md
index 2d299eac9d4..502ae64d3fe 100644
--- a/examples/puzzletron/README.md
+++ b/examples/puzzletron/README.md
@@ -481,113 +481,26 @@ 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. |
 
-## Understand sorting, width ranking, and slicing
+## Interpret width sanity results
 
-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.
+Puzzletron separates implementation correctness from ranking quality:
 
-The width sanity stages compare four views of one target width:
+- 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.
 
-| View | Meaning |
-|---|---|
-| Activation-sorted | Keep the prefix that the importance scores rank highest. This is the candidate whose ranking quality is being tested. |
-| Original or unsorted | Keep the same-sized prefix in the teacher's original order. Some older paths call this the random control, although the original prefix is deterministic. |
-| 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. |
-
-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.
-
-### What does width sanity measure?
-
-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.
-
-### What is a width-ranking miss?
-
-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. Loss-like
-metrics are lower-is-better; accuracy-like metrics are higher-is-better. 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.
-
-### What is an equivalence failure?
-
-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.
-
-### How stage completion and qualification differ
-
-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:
+To also fail a sanity stage on ranking-quality warnings, enable strict warning
+handling:
 
 ```yaml
 sanity:
   fail_on_warnings: true
 ```
 
-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. See
-[`docs/v2_architecture.md`](docs/v2_architecture.md#semantic-validation-gates)
-for the broader design context.
+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
diff --git a/examples/puzzletron/docs/sanity_validation.md b/examples/puzzletron/docs/sanity_validation.md
new file mode 100644
index 00000000000..48900a7fcc5
--- /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 or unsorted | Keep the same-sized prefix in the teacher's original order. Some older paths call this the random control, although the original prefix is deterministic. |
+| 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 c6bce52ed87..3ed5073658a 100644
--- a/examples/puzzletron/docs/v2_architecture.md
+++ b/examples/puzzletron/docs/v2_architecture.md
@@ -382,8 +382,8 @@ 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
-[Puzzletron README](../README.md#understand-sorting-width-ranking-and-slicing)
-provides definitions and a worked example for new users.
+[sanity validation guide](sanity_validation.md) provides definitions, the
+slicing mental model, and a worked example.
 
 ## Current implementation versus design direction
 

From 5ebbf951828db40ec96daec4748ef128704354fa Mon Sep 17 00:00:00 2001
From: Johannes Rausch <jrausch@nvidia.com>
Date: Sat, 8 Aug 2026 11:06:22 +0200
Subject: [PATCH 4/6] Fix Puzzletron sanity review findings

Make every rank consume the persisted sort verdict, keep empty-state gates visible, and distinguish the deterministic original-order control from seeded random permutations.

Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
---
 examples/puzzletron/README.md                 |  2 +-
 examples/puzzletron/docs/sanity_validation.md |  2 +-
 examples/puzzletron/docs/v2_architecture.md   |  6 +-
 .../diagnostics/campaign_progress_report.py   |  4 +-
 .../torch/puzzletron/stages/diagnostics.py    | 76 ++++++++-----------
 .../test_campaign_progress_report.py          | 16 ++++
 .../test_diagnostic_scoring_config.py         | 40 ++++++++++
 7 files changed, 97 insertions(+), 49 deletions(-)

diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md
index 502ae64d3fe..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. |
diff --git a/examples/puzzletron/docs/sanity_validation.md b/examples/puzzletron/docs/sanity_validation.md
index 48900a7fcc5..40f30151ace 100644
--- a/examples/puzzletron/docs/sanity_validation.md
+++ b/examples/puzzletron/docs/sanity_validation.md
@@ -39,7 +39,7 @@ 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 or unsorted | Keep the same-sized prefix in the teacher's original order. Some older paths call this the random control, although the original prefix is deterministic. |
+| 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. |
 
diff --git a/examples/puzzletron/docs/v2_architecture.md b/examples/puzzletron/docs/v2_architecture.md
index 3ed5073658a..1b9edd39ff8 100644
--- a/examples/puzzletron/docs/v2_architecture.md
+++ b/examples/puzzletron/docs/v2_architecture.md
@@ -138,7 +138,7 @@ 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
+  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
@@ -284,7 +284,7 @@ flowchart LR
     width["Multi-axis width<br/>importance hooks"]
     sort["Sort teacher once"]
     sortcheck["Sort equivalence"]
-    widthcheck["Ranking quality<br/>sorted vs reverse vs unsorted"]
+    widthcheck["Ranking quality<br/>sorted vs original vs reverse"]
     slicecheck["Dynamic vs physical<br/>slicing equivalence"]
     bypasscheck["Bypass overfit checks"]
     bypass["Nested bypass"]
@@ -358,7 +358,7 @@ 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 a quality finding; warning policy determines whether it also fails the stage |
+| 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 |
diff --git a/modelopt/torch/puzzletron/diagnostics/campaign_progress_report.py b/modelopt/torch/puzzletron/diagnostics/campaign_progress_report.py
index 80bf1caf7f1..c00e8e47212 100644
--- a/modelopt/torch/puzzletron/diagnostics/campaign_progress_report.py
+++ b/modelopt/torch/puzzletron/diagnostics/campaign_progress_report.py
@@ -525,7 +525,9 @@ def _activation_diagnostic_section(summary: dict[str, Any]) -> str:
         for metric in metrics
     )
     if not axes or not metrics:
-        return "<p class='empty'>The sanity artifact contains no plottable numeric metrics.</p>"
+        return gate_summary + (
+            "<p class='empty'>The sanity artifact contains no plottable numeric metrics.</p>"
+        )
     first_axis = axes[0]
     first_metric = metrics[0]
     cases: dict[tuple[Any, Any, Any], dict[str, dict[str, Any]]] = {}
diff --git a/modelopt/torch/puzzletron/stages/diagnostics.py b/modelopt/torch/puzzletron/stages/diagnostics.py
index 4f4fb8d2637..c0202fe7c9d 100644
--- a/modelopt/torch/puzzletron/stages/diagnostics.py
+++ b/modelopt/torch/puzzletron/stages/diagnostics.py
@@ -40,6 +40,7 @@
     maybe_cast_block_configs,
 )
 from ..diagnostics.campaign_findings import MetricSpec
+from ..diagnostics.sanity_verdict import SanityVerdict, complete_sanity_stage, finding_from_message
 from ..diagnostics.width_sanity import aggregate_parent_sweep_sanity
 from ..diagnostics.width_slice_equivalence import (
     evaluate_width_slice_equivalence,
@@ -1612,8 +1613,6 @@ def _parent_sweep_sanity_verdict(
 ):
     """Combine advisory width quality with blocking reused-sort correctness."""
 
-    from ..diagnostics.sanity_verdict import SanityVerdict
-
     width_findings = list(width_summary.get("findings") or ())
     sort_findings = [
         {**finding, "stage": "sort_sanity", "severity": "error"}
@@ -2740,7 +2739,6 @@ def _activation_diagnostic_parent_sweep(
     width_verdict = json.loads(width_summary_path.read_text(encoding="utf-8"))
     sort_summary_path = puzzle_dir / "artifacts" / "sort_sanity" / "summary.json"
     sort_verdict = json.loads(sort_summary_path.read_text(encoding="utf-8"))
-    from ..diagnostics.sanity_verdict import complete_sanity_stage
 
     return complete_sanity_stage(
         config,
@@ -3151,6 +3149,34 @@ def _sort_equivalence_tolerances(
     return tolerance, reverse_tolerance
 
 
+def _complete_sort_equivalence_stage(
+    config: dict[str, Any],
+    manifest: StageManifest,
+    *,
+    summary_path: Path,
+    table_path: Path,
+    fallback_metric: str,
+):
+    """Complete one rank from the verdict persisted by the master rank."""
+
+    summary = json.loads(summary_path.read_text(encoding="utf-8"))
+    return complete_sanity_stage(
+        config,
+        manifest,
+        outputs={
+            "summary_path": str(summary_path),
+            "table_path": str(table_path),
+            "metric": summary.get("metric", fallback_metric),
+            "delta": summary.get("delta"),
+            "reverse_delta": summary.get("reverse_delta"),
+        },
+        verdict=SanityVerdict(
+            passed=summary.get("passed") is True,
+            findings=list(summary.get("findings") or ()),
+        ),
+    )
+
+
 def sort_equivalence_stage(config: dict[str, Any], manifest: StageManifest):
     """Evaluate teacher and sorted teacher with the chunked AutoModel scorer."""
 
@@ -3270,8 +3296,6 @@ def sort_equivalence_stage(config: dict[str, Any], manifest: StageManifest):
 
     summary_path = artifacts_dir / "summary.json"
     md_path = artifacts_dir / "table.md"
-    delta_value = None
-    reverse_delta_value = None
     if dist.is_master():
         teacher_raw = json.loads((scoring_output_dir / "teacher.json").read_text())
         sorted_result_path = scoring_output_dir / "sliced_teacher.json"
@@ -3288,7 +3312,6 @@ def sort_equivalence_stage(config: dict[str, Any], manifest: StageManifest):
                 f"Could not find metric {metric!r}; teacher={teacher_raw} sorted={sorted_raw}"
             )
         delta = float(sorted_value) - float(teacher_value)
-        delta_value = delta
         reverse_value = _metric_avg(reverse_raw, metric) if reverse_raw is not None else None
         if include_reverse and reverse_value is None:
             raise RuntimeError(
@@ -3297,7 +3320,6 @@ def sort_equivalence_stage(config: dict[str, Any], manifest: StageManifest):
         reverse_delta = (
             float(reverse_value) - float(teacher_value) if reverse_value is not None else None
         )
-        reverse_delta_value = reverse_delta
         decision = _sort_equivalence_decision(
             delta=delta,
             reverse_delta=reverse_delta,
@@ -3380,12 +3402,6 @@ def sort_equivalence_stage(config: dict[str, Any], manifest: StageManifest):
         mprint(md_path.read_text(encoding="utf-8"))
         findings = []
         if not passed:
-            from ..diagnostics.sanity_verdict import (
-                SanityVerdict,
-                complete_sanity_stage,
-                finding_from_message,
-            )
-
             if not sorted_passed:
                 findings.append(
                     finding_from_message(
@@ -3419,33 +3435,13 @@ def sort_equivalence_stage(config: dict[str, Any], manifest: StageManifest):
             summary_path.write_text(
                 json.dumps(canonicalize(summary), indent=2, sort_keys=True) + "\n"
             )
-            dist.barrier()
-            return complete_sanity_stage(
-                config,
-                manifest,
-                outputs={
-                    "summary_path": str(summary_path),
-                    "table_path": str(md_path),
-                    "metric": metric,
-                    "delta": delta_value,
-                    "reverse_delta": reverse_delta_value,
-                },
-                verdict=SanityVerdict(passed=False, findings=findings),
-            )
     dist.barrier()
-    from ..diagnostics.sanity_verdict import SanityVerdict, complete_sanity_stage
-
-    return complete_sanity_stage(
+    return _complete_sort_equivalence_stage(
         config,
         manifest,
-        outputs={
-            "summary_path": str(summary_path),
-            "table_path": str(md_path),
-            "metric": metric,
-            "delta": delta_value,
-            "reverse_delta": reverse_delta_value,
-        },
-        verdict=SanityVerdict(passed=True),
+        summary_path=summary_path,
+        table_path=md_path,
+        fallback_metric=metric,
     )
 
 
@@ -3469,8 +3465,6 @@ def width_slice_equivalence_stage(config: dict[str, Any], manifest: StageManifes
             raise RuntimeError(f"distributed slicing summary has no evidence: {summary_path}")
         findings = list(summary.get("findings") or ())
         passed = bool(summary.get("passed", not findings))
-        from ..diagnostics.sanity_verdict import SanityVerdict, complete_sanity_stage
-
         return complete_sanity_stage(
             config,
             manifest,
@@ -3597,8 +3591,6 @@ def width_slice_equivalence_stage(config: dict[str, Any], manifest: StageManifes
         for case in summary.get("cases", ())
         if not case.get("passed", True)
     ]
-    from ..diagnostics.sanity_verdict import SanityVerdict, complete_sanity_stage
-
     return complete_sanity_stage(
         config,
         manifest,
@@ -4164,8 +4156,6 @@ def bypass_diagnostic_stage(config: dict[str, Any], manifest: StageManifest):
                 )
         dist.barrier()
 
-    from ..diagnostics.sanity_verdict import SanityVerdict, complete_sanity_stage
-
     findings = list((summary or {}).get("findings") or ())
     return complete_sanity_stage(
         config,
diff --git a/tests/unit/torch/puzzletron/test_campaign_progress_report.py b/tests/unit/torch/puzzletron/test_campaign_progress_report.py
index 556df6311bc..ab0f4c963d9 100644
--- a/tests/unit/torch/puzzletron/test_campaign_progress_report.py
+++ b/tests/unit/torch/puzzletron/test_campaign_progress_report.py
@@ -672,6 +672,22 @@ def test_progress_report_renders_axis_selectable_activation_diagnostic_tables(tm
     assert "unadjusted hidden-state mean squared error" in document
 
 
+def test_activation_diagnostic_empty_metrics_keep_gate_outcomes_visible():
+    section = report_module._activation_diagnostic_section(
+        {
+            "rows": [{"axis": "ffn_intermediate", "method": "sorted"}],
+            "width_present": True,
+            "width_passed": False,
+            "slicing_present": True,
+            "slicing_passed": False,
+        }
+    )
+
+    assert "Width ranking: quality warning" in section
+    assert "Dynamic/physical equivalence: failed (blocking correctness)" in section
+    assert "no plottable numeric metrics" in section
+
+
 def test_progress_report_recovers_sort_table_from_compact_reuse_summary(tmp_path: Path):
     _write(
         tmp_path / "manifests/sort_sanity.json",
diff --git a/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py b/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py
index 4a7379db886..52992b111c6 100644
--- a/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py
+++ b/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py
@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import json
 from pathlib import Path
 from types import SimpleNamespace
 
@@ -13,6 +14,7 @@
     _validate_worker_topology,
     _worker_config,
 )
+from modelopt.torch.puzzletron.manifest import StageManifest
 from modelopt.torch.puzzletron.stages import diagnostics
 from modelopt.torch.puzzletron.stages.diagnostics import _scoring_cfg_for_method
 
@@ -94,6 +96,44 @@ def test_sort_equivalence_keeps_production_and_reverse_control_tolerances_separa
     )["passed"]
 
 
+def test_sort_equivalence_uses_master_failure_verdict_on_every_rank(tmp_path: Path):
+    config = {"experiment": {"dir": str(tmp_path)}}
+    manifest = StageManifest(stage="sort_sanity", config=config)
+    summary_path = tmp_path / "artifacts" / "sort_sanity" / "summary.json"
+    summary_path.parent.mkdir(parents=True)
+    (tmp_path / "manifests").mkdir()
+    finding = {
+        "stage": "sort_sanity",
+        "message": "sorted teacher drift too large",
+        "severity": "error",
+    }
+    summary_path.write_text(
+        json.dumps(
+            {
+                "passed": False,
+                "metric": "lm_loss",
+                "delta": 0.25,
+                "reverse_delta": 0.0,
+                "findings": [finding],
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    result = diagnostics._complete_sort_equivalence_stage(
+        config,
+        manifest,
+        summary_path=summary_path,
+        table_path=summary_path.with_name("table.md"),
+        fallback_metric="unused",
+    )
+
+    assert result.status == "failed"
+    assert manifest.outputs["verdict"] == "failed"
+    assert manifest.outputs["delta"] == 0.25
+    assert manifest.outputs["findings"] == [finding]
+
+
 def test_axis_worker_preserves_requested_layers_and_targets_per_axis(tmp_path: Path):
     parallel = {
         "tp": 1,

From 39ed8441585817f814dc7593a707ca649529285e Mon Sep 17 00:00:00 2001
From: Johannes Rausch <jrausch@nvidia.com>
Date: Sat, 8 Aug 2026 11:17:53 +0200
Subject: [PATCH 5/6] Keep sort sanity finalization distributed

Finalize summary evidence and stage manifests before process-group teardown so one rank writes the result and every rank consumes the same completed verdict.

Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
---
 .../torch/puzzletron/stages/diagnostics.py    | 357 +++++++++++-------
 .../test_diagnostic_scoring_config.py         |  81 +++-
 2 files changed, 289 insertions(+), 149 deletions(-)

diff --git a/modelopt/torch/puzzletron/stages/diagnostics.py b/modelopt/torch/puzzletron/stages/diagnostics.py
index c0202fe7c9d..ec463cdca1a 100644
--- a/modelopt/torch/puzzletron/stages/diagnostics.py
+++ b/modelopt/torch/puzzletron/stages/diagnostics.py
@@ -3177,6 +3177,203 @@ def _complete_sort_equivalence_stage(
     )
 
 
+def _write_sort_equivalence_summary(
+    *,
+    teacher_dir: Path,
+    sorted_dir: Path,
+    reverse_dir: Path,
+    scoring_output_dir: Path,
+    reverse_output_dir: Path | None,
+    summary_path: Path,
+    table_path: Path,
+    metric: str,
+    include_reverse: bool,
+    tolerance: float,
+    reverse_tolerance: float,
+) -> None:
+    """Persist the canonical sort-equivalence evidence from rank zero."""
+
+    teacher_raw = json.loads((scoring_output_dir / "teacher.json").read_text())
+    sorted_result_path = scoring_output_dir / "sliced_teacher.json"
+    sorted_raw = json.loads(sorted_result_path.read_text())
+    reverse_raw = (
+        json.loads((reverse_output_dir / "sliced_teacher.json").read_text())
+        if reverse_output_dir is not None
+        else None
+    )
+    teacher_value = _metric_avg(teacher_raw, metric)
+    sorted_value = _metric_avg(sorted_raw, metric)
+    if teacher_value is None or sorted_value is None:
+        raise RuntimeError(
+            f"Could not find metric {metric!r}; teacher={teacher_raw} sorted={sorted_raw}"
+        )
+    delta = float(sorted_value) - float(teacher_value)
+    reverse_value = _metric_avg(reverse_raw, metric) if reverse_raw is not None else None
+    if include_reverse and reverse_value is None:
+        raise RuntimeError(
+            f"Could not find metric {metric!r} for reverse-sorted checkpoint: {reverse_raw}"
+        )
+    reverse_delta = (
+        float(reverse_value) - float(teacher_value) if reverse_value is not None else None
+    )
+    decision = _sort_equivalence_decision(
+        delta=delta,
+        reverse_delta=reverse_delta,
+        tolerance=tolerance,
+        reverse_tolerance=reverse_tolerance,
+    )
+    sorted_passed = decision["sorted_passed"]
+    reverse_passed = decision["reverse_passed"]
+    passed = decision["passed"]
+    findings = []
+    if not sorted_passed:
+        findings.append(
+            finding_from_message(
+                stage="sort_sanity",
+                message=(
+                    f"sorted teacher {metric} drift too large: delta={delta:.6g} "
+                    f"tolerance={tolerance:.6g}"
+                ),
+                evidence={"metric": metric, "delta": delta, "tolerance": tolerance},
+                severity="error",
+            )
+        )
+    if include_reverse and not reverse_passed:
+        findings.append(
+            finding_from_message(
+                stage="sort_sanity",
+                message=(
+                    f"reverse-sorted teacher {metric} drift too large: "
+                    f"reverse_delta={reverse_delta:.6g} "
+                    f"tolerance={reverse_tolerance:.6g}"
+                ),
+                evidence={
+                    "metric": metric,
+                    "reverse_delta": reverse_delta,
+                    "tolerance": reverse_tolerance,
+                },
+                severity="error",
+            )
+        )
+    summary = {
+        "metric": metric,
+        "teacher_dir": str(teacher_dir),
+        "sorted_teacher_dir": str(sorted_dir),
+        "teacher": {
+            key: _metric_avg(teacher_raw, key)
+            for key in _PRIMARY_METRICS
+            if _metric_avg(teacher_raw, key) is not None
+        },
+        "sorted_teacher": {
+            key: _metric_avg(sorted_raw, key)
+            for key in _PRIMARY_METRICS
+            if _metric_avg(sorted_raw, key) is not None
+        },
+        "reverse_sorted": (
+            {
+                key: _metric_avg(reverse_raw, key)
+                for key in _PRIMARY_METRICS
+                if _metric_avg(reverse_raw, key) is not None
+            }
+            if reverse_raw is not None
+            else None
+        ),
+        "delta": delta,
+        "abs_delta": abs(delta),
+        "sorted_passed": sorted_passed,
+        "reverse_delta": reverse_delta,
+        "reverse_abs_delta": abs(reverse_delta) if reverse_delta is not None else None,
+        "reverse_passed": reverse_passed if include_reverse else None,
+        "max_abs_delta": tolerance,
+        "max_abs_reverse_delta": reverse_tolerance,
+        "passed": passed,
+        "findings": findings,
+        "verdict": "passed" if passed else "failed",
+        "teacher_result": str(scoring_output_dir / "teacher.json"),
+        "sorted_result": str(sorted_result_path),
+        "reverse_sorted_dir": str(reverse_dir) if include_reverse else None,
+        "reverse_result": (
+            str(reverse_output_dir / "sliced_teacher.json")
+            if reverse_output_dir is not None
+            else None
+        ),
+    }
+    summary_path.write_text(json.dumps(canonicalize(summary), indent=2, sort_keys=True) + "\n")
+    table_path.write_text(
+        "\n".join(
+            [
+                f"# Sorted Teacher Equivalence ({metric})",
+                "",
+                "| checkpoint | value | delta_vs_teacher |",
+                "| --- | --- | --- |",
+                f"| teacher | {teacher_value:.8g} | 0 |",
+                f"| sorted_teacher | {sorted_value:.8g} | {delta:.8g} |",
+                *(
+                    [
+                        f"| reverse_sorted_teacher | {reverse_value:.8g} | "
+                        f"{reverse_delta:.8g} |"
+                    ]
+                    if reverse_value is not None and reverse_delta is not None
+                    else []
+                ),
+                "",
+                f"pass: {'yes' if passed else 'no'} "
+                f"(max_abs_delta={tolerance:.3g}, "
+                f"max_abs_reverse_delta={reverse_tolerance:.3g})",
+            ]
+        )
+        + "\n",
+        encoding="utf-8",
+    )
+    mprint(table_path.read_text(encoding="utf-8"))
+
+
+def _finalize_sort_equivalence_stage(
+    config: dict[str, Any],
+    manifest: StageManifest,
+    *,
+    teacher_dir: Path,
+    sorted_dir: Path,
+    reverse_dir: Path,
+    scoring_output_dir: Path,
+    reverse_output_dir: Path | None,
+    summary_path: Path,
+    table_path: Path,
+    metric: str,
+    include_reverse: bool,
+    tolerance: float,
+    reverse_tolerance: float,
+):
+    """Publish and consume one verdict before distributed teardown."""
+
+    if not dist.is_initialized():
+        raise RuntimeError("sort-equivalence finalization requires an active process group")
+    if dist.is_master():
+        _write_sort_equivalence_summary(
+            teacher_dir=teacher_dir,
+            sorted_dir=sorted_dir,
+            reverse_dir=reverse_dir,
+            scoring_output_dir=scoring_output_dir,
+            reverse_output_dir=reverse_output_dir,
+            summary_path=summary_path,
+            table_path=table_path,
+            metric=metric,
+            include_reverse=include_reverse,
+            tolerance=tolerance,
+            reverse_tolerance=reverse_tolerance,
+        )
+    dist.barrier()
+    result = _complete_sort_equivalence_stage(
+        config,
+        manifest,
+        summary_path=summary_path,
+        table_path=table_path,
+        fallback_metric=metric,
+    )
+    dist.barrier()
+    return result
+
+
 def sort_equivalence_stage(config: dict[str, Any], manifest: StageManifest):
     """Evaluate teacher and sorted teacher with the chunked AutoModel scorer."""
 
@@ -3294,155 +3491,23 @@ def sort_equivalence_stage(config: dict[str, Any], manifest: StageManifest):
             launch_score_solutions_automodel(reverse_scoring_cfg)
             dist.barrier()
 
-    summary_path = artifacts_dir / "summary.json"
-    md_path = artifacts_dir / "table.md"
-    if dist.is_master():
-        teacher_raw = json.loads((scoring_output_dir / "teacher.json").read_text())
-        sorted_result_path = scoring_output_dir / "sliced_teacher.json"
-        sorted_raw = json.loads(sorted_result_path.read_text())
-        reverse_raw = (
-            json.loads((reverse_output_dir / "sliced_teacher.json").read_text())
-            if reverse_output_dir is not None
-            else None
-        )
-        teacher_value = _metric_avg(teacher_raw, metric)
-        sorted_value = _metric_avg(sorted_raw, metric)
-        if teacher_value is None or sorted_value is None:
-            raise RuntimeError(
-                f"Could not find metric {metric!r}; teacher={teacher_raw} sorted={sorted_raw}"
-            )
-        delta = float(sorted_value) - float(teacher_value)
-        reverse_value = _metric_avg(reverse_raw, metric) if reverse_raw is not None else None
-        if include_reverse and reverse_value is None:
-            raise RuntimeError(
-                f"Could not find metric {metric!r} for reverse-sorted checkpoint: {reverse_raw}"
-            )
-        reverse_delta = (
-            float(reverse_value) - float(teacher_value) if reverse_value is not None else None
-        )
-        decision = _sort_equivalence_decision(
-            delta=delta,
-            reverse_delta=reverse_delta,
+        summary_path = artifacts_dir / "summary.json"
+        table_path = artifacts_dir / "table.md"
+        return _finalize_sort_equivalence_stage(
+            config,
+            manifest,
+            teacher_dir=teacher_dir,
+            sorted_dir=sorted_dir,
+            reverse_dir=reverse_dir,
+            scoring_output_dir=scoring_output_dir,
+            reverse_output_dir=reverse_output_dir,
+            summary_path=summary_path,
+            table_path=table_path,
+            metric=metric,
+            include_reverse=include_reverse,
             tolerance=tolerance,
             reverse_tolerance=reverse_tolerance,
         )
-        sorted_passed = decision["sorted_passed"]
-        reverse_passed = decision["reverse_passed"]
-        passed = decision["passed"]
-        summary = {
-            "metric": metric,
-            "teacher_dir": str(teacher_dir),
-            "sorted_teacher_dir": str(sorted_dir),
-            "teacher": {
-                key: _metric_avg(teacher_raw, key)
-                for key in _PRIMARY_METRICS
-                if _metric_avg(teacher_raw, key) is not None
-            },
-            "sorted_teacher": {
-                key: _metric_avg(sorted_raw, key)
-                for key in _PRIMARY_METRICS
-                if _metric_avg(sorted_raw, key) is not None
-            },
-            "reverse_sorted": (
-                {
-                    key: _metric_avg(reverse_raw, key)
-                    for key in _PRIMARY_METRICS
-                    if _metric_avg(reverse_raw, key) is not None
-                }
-                if reverse_raw is not None
-                else None
-            ),
-            "delta": delta,
-            "abs_delta": abs(delta),
-            "sorted_passed": sorted_passed,
-            "reverse_delta": reverse_delta,
-            "reverse_abs_delta": abs(reverse_delta) if reverse_delta is not None else None,
-            "reverse_passed": reverse_passed if include_reverse else None,
-            "max_abs_delta": tolerance,
-            "max_abs_reverse_delta": reverse_tolerance,
-            "passed": passed,
-            "findings": [],
-            "verdict": "passed" if passed else "failed",
-            "teacher_result": str(scoring_output_dir / "teacher.json"),
-            "sorted_result": str(sorted_result_path),
-            "reverse_sorted_dir": str(reverse_dir) if include_reverse else None,
-            "reverse_result": (
-                str(reverse_output_dir / "sliced_teacher.json")
-                if reverse_output_dir is not None
-                else None
-            ),
-        }
-        summary_path.write_text(json.dumps(canonicalize(summary), indent=2, sort_keys=True) + "\n")
-        md_path.write_text(
-            "\n".join(
-                [
-                    f"# Sorted Teacher Equivalence ({metric})",
-                    "",
-                    "| checkpoint | value | delta_vs_teacher |",
-                    "| --- | --- | --- |",
-                    f"| teacher | {teacher_value:.8g} | 0 |",
-                    f"| sorted_teacher | {sorted_value:.8g} | {delta:.8g} |",
-                    *(
-                        [
-                            f"| reverse_sorted_teacher | {reverse_value:.8g} | "
-                            f"{reverse_delta:.8g} |"
-                        ]
-                        if reverse_value is not None and reverse_delta is not None
-                        else []
-                    ),
-                    "",
-                    f"pass: {'yes' if passed else 'no'} "
-                    f"(max_abs_delta={tolerance:.3g}, "
-                    f"max_abs_reverse_delta={reverse_tolerance:.3g})",
-                ]
-            )
-            + "\n",
-            encoding="utf-8",
-        )
-        mprint(md_path.read_text(encoding="utf-8"))
-        findings = []
-        if not passed:
-            if not sorted_passed:
-                findings.append(
-                    finding_from_message(
-                        stage="sort_sanity",
-                        message=(
-                            f"sorted teacher {metric} drift too large: delta={delta:.6g} "
-                            f"tolerance={tolerance:.6g}"
-                        ),
-                        evidence={"metric": metric, "delta": delta, "tolerance": tolerance},
-                        severity="error",
-                    )
-                )
-            if include_reverse and not reverse_passed:
-                findings.append(
-                    finding_from_message(
-                        stage="sort_sanity",
-                        message=(
-                            f"reverse-sorted teacher {metric} drift too large: "
-                            f"reverse_delta={reverse_delta:.6g} "
-                            f"tolerance={reverse_tolerance:.6g}"
-                        ),
-                        evidence={
-                            "metric": metric,
-                            "reverse_delta": reverse_delta,
-                            "tolerance": reverse_tolerance,
-                        },
-                        severity="error",
-                    )
-                )
-            summary["findings"] = findings
-            summary_path.write_text(
-                json.dumps(canonicalize(summary), indent=2, sort_keys=True) + "\n"
-            )
-    dist.barrier()
-    return _complete_sort_equivalence_stage(
-        config,
-        manifest,
-        summary_path=summary_path,
-        table_path=md_path,
-        fallback_metric=metric,
-    )
 
 
 def width_slice_equivalence_stage(config: dict[str, Any], manifest: StageManifest):
diff --git a/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py b/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py
index 52992b111c6..4daf7f6854e 100644
--- a/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py
+++ b/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py
@@ -5,6 +5,7 @@
 from pathlib import Path
 from types import SimpleNamespace
 
+import pytest
 from omegaconf import OmegaConf
 
 import examples.puzzletron.run_axis_diagnostic_worker as axis_worker
@@ -96,7 +97,40 @@ def test_sort_equivalence_keeps_production_and_reverse_control_tolerances_separa
     )["passed"]
 
 
-def test_sort_equivalence_uses_master_failure_verdict_on_every_rank(tmp_path: Path):
+def test_sort_equivalence_summary_records_blocking_drift(tmp_path: Path):
+    scoring_output_dir = tmp_path / "scoring"
+    scoring_output_dir.mkdir()
+    (scoring_output_dir / "teacher.json").write_text(
+        json.dumps({"lm_loss": {"avg": 1.0}}), encoding="utf-8"
+    )
+    (scoring_output_dir / "sliced_teacher.json").write_text(
+        json.dumps({"lm_loss": {"avg": 1.25}}), encoding="utf-8"
+    )
+    summary_path = tmp_path / "summary.json"
+
+    diagnostics._write_sort_equivalence_summary(
+        teacher_dir=tmp_path / "teacher",
+        sorted_dir=tmp_path / "sorted",
+        reverse_dir=tmp_path / "reverse",
+        scoring_output_dir=scoring_output_dir,
+        reverse_output_dir=None,
+        summary_path=summary_path,
+        table_path=tmp_path / "table.md",
+        metric="lm_loss",
+        include_reverse=False,
+        tolerance=0.01,
+        reverse_tolerance=0.01,
+    )
+
+    summary = json.loads(summary_path.read_text(encoding="utf-8"))
+    assert summary["passed"] is False
+    assert summary["delta"] == 0.25
+    assert summary["findings"][0]["stage"] == "sort_sanity"
+
+
+def test_sort_equivalence_uses_master_failure_verdict_on_every_rank(
+    monkeypatch, tmp_path: Path
+):
     config = {"experiment": {"dir": str(tmp_path)}}
     manifest = StageManifest(stage="sort_sanity", config=config)
     summary_path = tmp_path / "artifacts" / "sort_sanity" / "summary.json"
@@ -119,21 +153,62 @@ def test_sort_equivalence_uses_master_failure_verdict_on_every_rank(tmp_path: Pa
         ),
         encoding="utf-8",
     )
+    barriers = []
+    monkeypatch.setattr(diagnostics.dist, "is_initialized", lambda: True)
+    monkeypatch.setattr(diagnostics.dist, "is_master", lambda: False)
+    monkeypatch.setattr(diagnostics.dist, "barrier", lambda: barriers.append("barrier"))
+    monkeypatch.setattr(
+        diagnostics,
+        "_write_sort_equivalence_summary",
+        lambda **_kwargs: (_ for _ in ()).throw(AssertionError("non-master wrote summary")),
+    )
 
-    result = diagnostics._complete_sort_equivalence_stage(
+    result = diagnostics._finalize_sort_equivalence_stage(
         config,
         manifest,
+        teacher_dir=tmp_path / "teacher",
+        sorted_dir=tmp_path / "sorted",
+        reverse_dir=tmp_path / "reverse",
+        scoring_output_dir=tmp_path / "scoring",
+        reverse_output_dir=None,
         summary_path=summary_path,
         table_path=summary_path.with_name("table.md"),
-        fallback_metric="unused",
+        metric="lm_loss",
+        include_reverse=False,
+        tolerance=0.01,
+        reverse_tolerance=0.01,
     )
 
     assert result.status == "failed"
+    assert barriers == ["barrier", "barrier"]
     assert manifest.outputs["verdict"] == "failed"
     assert manifest.outputs["delta"] == 0.25
     assert manifest.outputs["findings"] == [finding]
 
 
+def test_sort_equivalence_rejects_finalization_after_distributed_cleanup(
+    monkeypatch, tmp_path: Path
+):
+    monkeypatch.setattr(diagnostics.dist, "is_initialized", lambda: False)
+
+    with pytest.raises(RuntimeError, match="active process group"):
+        diagnostics._finalize_sort_equivalence_stage(
+            {"experiment": {"dir": str(tmp_path)}},
+            StageManifest(stage="sort_sanity"),
+            teacher_dir=tmp_path / "teacher",
+            sorted_dir=tmp_path / "sorted",
+            reverse_dir=tmp_path / "reverse",
+            scoring_output_dir=tmp_path / "scoring",
+            reverse_output_dir=None,
+            summary_path=tmp_path / "summary.json",
+            table_path=tmp_path / "table.md",
+            metric="lm_loss",
+            include_reverse=False,
+            tolerance=0.01,
+            reverse_tolerance=0.01,
+        )
+
+
 def test_axis_worker_preserves_requested_layers_and_targets_per_axis(tmp_path: Path):
     parallel = {
         "tp": 1,

From cc70a2a2446b32aaeb1bdd7d5626c4c4b058c480 Mon Sep 17 00:00:00 2001
From: Johannes Rausch <jrausch@nvidia.com>
Date: Sat, 8 Aug 2026 11:58:18 +0200
Subject: [PATCH 6/6] Propagate sort summary write failures

Broadcast rank-zero write errors before consuming the summary so all ranks exit consistently instead of hanging at a barrier.

Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
---
 .../torch/puzzletron/stages/diagnostics.py    | 36 ++++++++-----
 .../test_diagnostic_scoring_config.py         | 52 +++++++++++++++++++
 2 files changed, 75 insertions(+), 13 deletions(-)

diff --git a/modelopt/torch/puzzletron/stages/diagnostics.py b/modelopt/torch/puzzletron/stages/diagnostics.py
index ec463cdca1a..60e5841f087 100644
--- a/modelopt/torch/puzzletron/stages/diagnostics.py
+++ b/modelopt/torch/puzzletron/stages/diagnostics.py
@@ -3348,21 +3348,31 @@ def _finalize_sort_equivalence_stage(
 
     if not dist.is_initialized():
         raise RuntimeError("sort-equivalence finalization requires an active process group")
+    write_error = None
     if dist.is_master():
-        _write_sort_equivalence_summary(
-            teacher_dir=teacher_dir,
-            sorted_dir=sorted_dir,
-            reverse_dir=reverse_dir,
-            scoring_output_dir=scoring_output_dir,
-            reverse_output_dir=reverse_output_dir,
-            summary_path=summary_path,
-            table_path=table_path,
-            metric=metric,
-            include_reverse=include_reverse,
-            tolerance=tolerance,
-            reverse_tolerance=reverse_tolerance,
-        )
+        try:
+            _write_sort_equivalence_summary(
+                teacher_dir=teacher_dir,
+                sorted_dir=sorted_dir,
+                reverse_dir=reverse_dir,
+                scoring_output_dir=scoring_output_dir,
+                reverse_output_dir=reverse_output_dir,
+                summary_path=summary_path,
+                table_path=table_path,
+                metric=metric,
+                include_reverse=include_reverse,
+                tolerance=tolerance,
+                reverse_tolerance=reverse_tolerance,
+            )
+        except Exception as exc:
+            write_error = {"type": type(exc).__name__, "message": str(exc)}
     dist.barrier()
+    write_error = dist.broadcast(write_error, src=0)
+    if write_error is not None:
+        raise RuntimeError(
+            "sort-equivalence summary write failed on the master rank: "
+            f"{write_error['type']}: {write_error['message']}"
+        )
     result = _complete_sort_equivalence_stage(
         config,
         manifest,
diff --git a/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py b/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py
index 4daf7f6854e..34c2f689086 100644
--- a/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py
+++ b/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py
@@ -126,6 +126,8 @@ def test_sort_equivalence_summary_records_blocking_drift(tmp_path: Path):
     assert summary["passed"] is False
     assert summary["delta"] == 0.25
     assert summary["findings"][0]["stage"] == "sort_sanity"
+    assert summary["findings"][0]["severity"] == "error"
+    assert summary["verdict"] == "failed"
 
 
 def test_sort_equivalence_uses_master_failure_verdict_on_every_rank(
@@ -157,6 +159,7 @@ def test_sort_equivalence_uses_master_failure_verdict_on_every_rank(
     monkeypatch.setattr(diagnostics.dist, "is_initialized", lambda: True)
     monkeypatch.setattr(diagnostics.dist, "is_master", lambda: False)
     monkeypatch.setattr(diagnostics.dist, "barrier", lambda: barriers.append("barrier"))
+    monkeypatch.setattr(diagnostics.dist, "broadcast", lambda value, src=0: value)
     monkeypatch.setattr(
         diagnostics,
         "_write_sort_equivalence_summary",
@@ -186,6 +189,55 @@ def test_sort_equivalence_uses_master_failure_verdict_on_every_rank(
     assert manifest.outputs["findings"] == [finding]
 
 
+@pytest.mark.parametrize("is_master", [True, False])
+def test_sort_equivalence_propagates_master_write_failure(
+    monkeypatch, tmp_path: Path, is_master: bool
+):
+    barriers = []
+    write_calls = []
+    monkeypatch.setattr(diagnostics.dist, "is_initialized", lambda: True)
+    monkeypatch.setattr(diagnostics.dist, "is_master", lambda: is_master)
+    monkeypatch.setattr(diagnostics.dist, "barrier", lambda: barriers.append("barrier"))
+    monkeypatch.setattr(
+        diagnostics.dist,
+        "broadcast",
+        lambda value, src=0: value or {"type": "ValueError", "message": "missing metric"},
+    )
+
+    def fail_write(**_kwargs):
+        write_calls.append("write")
+        raise ValueError("missing metric")
+
+    monkeypatch.setattr(
+        diagnostics,
+        "_write_sort_equivalence_summary",
+        fail_write,
+    )
+
+    with pytest.raises(
+        RuntimeError,
+        match="summary write failed on the master rank: ValueError: missing metric",
+    ):
+        diagnostics._finalize_sort_equivalence_stage(
+            {"experiment": {"dir": str(tmp_path)}},
+            StageManifest(stage="sort_sanity"),
+            teacher_dir=tmp_path / "teacher",
+            sorted_dir=tmp_path / "sorted",
+            reverse_dir=tmp_path / "reverse",
+            scoring_output_dir=tmp_path / "scoring",
+            reverse_output_dir=None,
+            summary_path=tmp_path / "summary.json",
+            table_path=tmp_path / "table.md",
+            metric="lm_loss",
+            include_reverse=False,
+            tolerance=0.01,
+            reverse_tolerance=0.01,
+        )
+
+    assert barriers == ["barrier"]
+    assert write_calls == (["write"] if is_master else [])
+
+
 def test_sort_equivalence_rejects_finalization_after_distributed_cleanup(
     monkeypatch, tmp_path: Path
 ):