feat(v4/buildspec): MBSpec Stage B — ModelBuildSpec composer + verify_build_spec

Adds the composer and coherence checker layer (ModelBuildSpec.md §3.3,
§6 B).

cppmega_v4/buildspec/model_build_spec.py:
  - Rewriter Protocol: name + required_preconditions + provided_postconditions
    + __call__. Forward-ref so model_build_spec.py stays decoupled from
    the concrete rewriters that land in Stage C.
  - ModelBuildSpec frozen dataclass: graph + loss + optim + rewrites +
    dim_env + state_tokens (defaults to {"single_head"}).
    .replace() returns a new spec with field overrides (immutable).
    .apply_rewrites() walks the chain in declared order, checks each
    rewriter's required_preconditions against the running state_tokens
    set, runs it, and merges provided_postconditions into the result.
  - RewriteOrderError for chain-order violations at apply time.

cppmega_v4/buildspec/diagnostics.py:
  - BuildDiagnosticSeverity / BuildDiagnostic / BuildDiagnostics
  - verify_build_spec(spec, *, check_shapes=True) runs four checks:
      1. Rewrite-chain ordering simulation (precondition satisfaction
         + duplicate-name warning)
      2. Loss head_outputs vs (simulated) post-rewrite brick names —
         when rewrites are present, the loss's declared head names are
         treated as future-valid (defers final check to apply time)
      3. Optimizer matcher coverage (warns on a group whose matcher
         selects zero brick names — dead group = silent param skip)
      4. Shape coherence via cppmega_v4.spec.resolve_shapes (no-op
         when dim_env is empty)
  - All severities mirror cppmega_v4.spec.DiagnosticSeverity so the GUI
    can use one legend.

Tests (tests/v4/test_mbspec_stage_b.py, 22 cases):
  - ModelBuildSpec rejection: non-BrickGraph / non-LossSpec / non-OptimSpec
    / rewriter missing required attribute
  - default state_token / replace() immutability
  - apply_rewrites: empty chain identity; postcondition added; missing
    precondition raises RewriteOrderError; chained pre→post satisfaction
  - verify happy paths: matching head brick = clean; summary dict shape;
    Qwen-preset without head brick surfaces the right loss error
  - verify warning surfaces: dead optim matcher (regex matches nothing);
    "all" matcher never flagged; rewrite duplicate name warning
  - verify error surfaces: rewrite-chain bad order; resolver errors
    surfaced as component=shape; check_shapes=False skips
  - system: MTP-style head naming (logits_0 / logits_1) accepted when
    rewriters are present in the chain

Full v4 regression: 833 passed / 5 skipped / 0 failed.

Author: apstenku123

cppmega_v4/buildspec/__init__.py

@@ -2,15 +2,24 @@
 
 See ``ModelBuildSpec.md`` (repo root) for the design.
 
-Stage A surface (this commit):
-  - loss_spec.LossKind / LossSpec / cross_entropy_loss / mtp_weighted_loss
+Stage A+B surface:
+  - loss_spec: LossKind / LossSpec / cross_entropy_loss / mtp_weighted_loss
     / ifim_shaped_loss / mhc_attn_bias_loss / custom_loss / LOSS_BUILTINS
-  - optim_spec.OptimKind / OptimSpec / ParamGroup / adamw / muon /
+  - optim_spec: OptimKind / OptimSpec / ParamGroup / adamw / muon /
     muon_adamw_hybrid / sgd / OPTIM_BUILTINS
+  - model_build_spec: ModelBuildSpec / Rewriter / RewriteOrderError
+  - diagnostics: BuildDiagnostic / BuildDiagnosticSeverity /
+    BuildDiagnostics / verify_build_spec
 """
 
 from __future__ import annotations
 
+from cppmega_v4.buildspec.diagnostics import (
+    BuildDiagnostic,
+    BuildDiagnosticSeverity,
+    BuildDiagnostics,
+    verify_build_spec,
+)
 from cppmega_v4.buildspec.loss_spec import (
     LOSS_BUILTINS,
     LossKind,
@@ -21,6 +30,11 @@
     mhc_attn_bias_loss,
     mtp_weighted_loss,
 )
+from cppmega_v4.buildspec.model_build_spec import (
+    ModelBuildSpec,
+    RewriteOrderError,
+    Rewriter,
+)
 from cppmega_v4.buildspec.optim_spec import (
     OPTIM_BUILTINS,
     OptimKind,
@@ -33,13 +47,19 @@
 )
 
 __all__ = [
+    "BuildDiagnostic",
+    "BuildDiagnosticSeverity",
+    "BuildDiagnostics",
     "LOSS_BUILTINS",
     "LossKind",
     "LossSpec",
+    "ModelBuildSpec",
     "OPTIM_BUILTINS",
     "OptimKind",
     "OptimSpec",
     "ParamGroup",
+    "RewriteOrderError",
+    "Rewriter",
     "adamw",
     "cross_entropy_loss",
     "custom_loss",
@@ -49,4 +69,5 @@
     "muon",
     "muon_adamw_hybrid",
     "sgd",
+    "verify_build_spec",
 ]

cppmega_v4/buildspec/diagnostics.py

@@ -0,0 +1,325 @@
+"""Coherence checker for :class:`ModelBuildSpec`.
+
+Walks the spec and reports any incoherence between its three components
+plus the rewrite chain:
+
+  - LossSpec.head_outputs ⊆ post-rewrite graph.names
+    (the loss reads from brick outputs that must actually exist after
+    all rewrites apply)
+  - OptimSpec.groups[*].matcher matches at least one parameter on the
+    post-rewrite graph (no dead matchers)
+  - Rewrite chain preconditions are satisfied at each step (dry-run via
+    state-token simulation, without actually invoking the rewriters)
+  - Shape contracts on the post-rewrite graph still resolve cleanly
+    under ``spec.dim_env`` (delegates to cppmega_v4.spec.resolve_shapes)
+
+Severities mirror cppmega_v4.spec.DiagnosticSeverity (ERROR / WARNING /
+INFO) so the GUI can render them with the same legend.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Iterable
+
+from cppmega_v4.buildspec.model_build_spec import (
+    ModelBuildSpec,
+    RewriteOrderError,
+)
+from cppmega_v4.buildspec.optim_spec import ParamGroup
+
+
+class BuildDiagnosticSeverity(str, Enum):
+    INFO = "info"
+    WARNING = "warning"
+    ERROR = "error"
+
+
+@dataclass(frozen=True)
+class BuildDiagnostic:
+    """One issue surfaced by :func:`verify_build_spec`."""
+
+    severity: BuildDiagnosticSeverity
+    component: str   # "loss" | "optim" | "graph" | "rewrites" | "shape"
+    message: str
+    suggested_fix: str | None = None
+
+
+@dataclass(frozen=True)
+class BuildDiagnostics:
+    """Bundle returned by :func:`verify_build_spec`."""
+
+    diagnostics: tuple[BuildDiagnostic, ...] = field(default_factory=tuple)
+
+    @property
+    def errors(self) -> tuple[BuildDiagnostic, ...]:
+        return tuple(
+            d for d in self.diagnostics
+            if d.severity is BuildDiagnosticSeverity.ERROR
+        )
+
+    @property
+    def warnings(self) -> tuple[BuildDiagnostic, ...]:
+        return tuple(
+            d for d in self.diagnostics
+            if d.severity is BuildDiagnosticSeverity.WARNING
+        )
+
+    @property
+    def has_errors(self) -> bool:
+        return any(
+            d.severity is BuildDiagnosticSeverity.ERROR
+            for d in self.diagnostics
+        )
+
+    def summary(self) -> dict[str, int]:
+        return {
+            "errors": len(self.errors),
+            "warnings": len(self.warnings),
+            "total": len(self.diagnostics),
+        }
+
+
+# ---------------------------------------------------------------------------
+# Per-component checkers
+# ---------------------------------------------------------------------------
+
+
+def _check_rewrite_chain(spec: ModelBuildSpec) -> list[BuildDiagnostic]:
+    """Simulate rewrite preconditions without invoking the rewriters.
+
+    Anti-cycle: a token added by an earlier rewriter and then removed
+    isn't a thing (postconditions are additive). We catch ordering
+    errors by simulating the state-token set as we walk."""
+    diags: list[BuildDiagnostic] = []
+    if not spec.rewrites:
+        return diags
+    state = set(spec.state_tokens)
+    names_seen: set[str] = set()
+    for r in spec.rewrites:
+        if r.name in names_seen:
+            diags.append(
+                BuildDiagnostic(
+                    severity=BuildDiagnosticSeverity.WARNING,
+                    component="rewrites",
+                    message=(
+                        f"rewriter {r.name!r} appears more than once in the "
+                        "chain; later instances may be no-ops or contradict "
+                        "earlier ones"
+                    ),
+                    suggested_fix="dedupe the rewrite chain",
+                )
+            )
+        names_seen.add(r.name)
+        missing = r.required_preconditions - state
+        if missing:
+            diags.append(
+                BuildDiagnostic(
+                    severity=BuildDiagnosticSeverity.ERROR,
+                    component="rewrites",
+                    message=(
+                        f"rewriter {r.name!r} requires preconditions "
+                        f"{sorted(missing)} but the chain reaches it with "
+                        f"state {sorted(state)}"
+                    ),
+                    suggested_fix=(
+                        "reorder the rewrite chain so a rewriter providing "
+                        f"{sorted(missing)} runs first"
+                    ),
+                )
+            )
+        state |= r.provided_postconditions
+    return diags
+
+
+def _matcher_matches_any(
+    matcher: str, param_names: Iterable[str],
+) -> bool:
+    """Return True if ``matcher`` selects at least one of ``param_names``.
+
+    Mirrors the runtime matcher logic; kept here so the verifier doesn't
+    need to materialise an nn.Module. Built-in matchers map to substring
+    rules; ``regex:<pattern>`` compiles a Python regex."""
+    if matcher == "all":
+        return True
+    if matcher.startswith("regex:"):
+        import re
+        pattern = re.compile(matcher.removeprefix("regex:"))
+        return any(pattern.search(n) for n in param_names)
+    needle = {
+        "moe_experts": "expert",
+        "embeddings":  "embed",
+        "attention":   "attn",
+        "mlp":         "mlp",
+        "head":        "head",
+    }.get(matcher)
+    if needle is None:
+        return False
+    return any(needle in n for n in param_names)
+
+
+def _check_optim_matchers(spec: ModelBuildSpec) -> list[BuildDiagnostic]:
+    """Warn on optimizer matchers that don't select any brick name.
+
+    The brick names are a coarse proxy for parameter names, but it's
+    enough to catch the common typo cases. A matcher selecting zero
+    params silently drops weight updates — surface that loudly."""
+    diags: list[BuildDiagnostic] = []
+    brick_names = [n.name for n in spec.graph.nodes]
+    if not brick_names:
+        return diags
+    for g in spec.optim.groups:
+        if g.matcher == "all":
+            continue
+        if not _matcher_matches_any(g.matcher, brick_names):
+            diags.append(
+                BuildDiagnostic(
+                    severity=BuildDiagnosticSeverity.WARNING,
+                    component="optim",
+                    message=(
+                        f"optimizer matcher {g.matcher!r} (lr={g.lr}) "
+                        "matches no brick names in the graph; the group "
+                        "will receive no parameters"
+                    ),
+                    suggested_fix=(
+                        "rename the matcher to 'all' or check the brick "
+                        "names match the matcher's pattern"
+                    ),
+                )
+            )
+    return diags
+
+
+def _check_loss_head_outputs(
+    spec: ModelBuildSpec, post_rewrite_names: set[str],
+) -> list[BuildDiagnostic]:
+    diags: list[BuildDiagnostic] = []
+    for head in spec.loss.head_outputs:
+        # The MTP head naming convention is "<prefix>_<i>"; before
+        # rewrites are applied, the bare "logits" name is what bricks
+        # produce. We accept either:
+        #   - exact match in post-rewrite brick names
+        #   - prefix match (head_output starts with a brick name)
+        if head in post_rewrite_names:
+            continue
+        if any(head.startswith(n + "_") or n.startswith(head + "_")
+               for n in post_rewrite_names):
+            continue
+        diags.append(
+            BuildDiagnostic(
+                severity=BuildDiagnosticSeverity.ERROR,
+                component="loss",
+                message=(
+                    f"LossSpec.head_outputs entry {head!r} does not match "
+                    "any brick name in the post-rewrite graph "
+                    f"(have: {sorted(post_rewrite_names)[:5]}{'...' if len(post_rewrite_names) > 5 else ''})"
+                ),
+                suggested_fix=(
+                    "add a head brick with this name OR apply MTPRewriter "
+                    "before verifying"
+                ),
+            )
+        )
+    return diags
+
+
+def _check_shape_coherence(spec: ModelBuildSpec) -> list[BuildDiagnostic]:
+    """Defer to cppmega_v4.spec.resolve_shapes (lenient). Surface any
+    ERROR diagnostics as 'shape' build-diagnostics."""
+    if not spec.dim_env:
+        return []
+    diags: list[BuildDiagnostic] = []
+    try:
+        from cppmega_v4.spec import resolve_shapes
+    except ImportError:
+        return []  # spec layer not present — skip silently
+    try:
+        resolved = resolve_shapes(
+            spec.graph, spec.dim_env, strict=False,
+            available_side_channels=frozenset({"doc_ids", "token_ids"}),
+        )
+    except Exception as exc:
+        return [
+            BuildDiagnostic(
+                severity=BuildDiagnosticSeverity.ERROR,
+                component="shape",
+                message=f"shape resolution raised: {exc}",
+                suggested_fix="fix the graph or supply missing dim_env entries",
+            )
+        ]
+    for d in resolved.errors:
+        diags.append(
+            BuildDiagnostic(
+                severity=BuildDiagnosticSeverity.ERROR,
+                component="shape",
+                message=d.message,
+                suggested_fix=d.suggested_fix,
+            )
+        )
+    for d in resolved.warnings:
+        diags.append(
+            BuildDiagnostic(
+                severity=BuildDiagnosticSeverity.WARNING,
+                component="shape",
+                message=d.message,
+                suggested_fix=d.suggested_fix,
+            )
+        )
+    return diags
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def verify_build_spec(
+    spec: ModelBuildSpec,
+    *,
+    check_shapes: bool = True,
+) -> BuildDiagnostics:
+    """Run all coherence checks; return a :class:`BuildDiagnostics`.
+
+    Args:
+      spec: the spec to verify.
+      check_shapes: when True (default), runs shape coherence via
+        cppmega_v4.spec.resolve_shapes (no-op if dim_env is empty).
+    """
+    diags: list[BuildDiagnostic] = []
+
+    # 1. Rewrite-chain ordering
+    diags.extend(_check_rewrite_chain(spec))
+
+    # 2. Loss head_outputs vs (simulated) post-rewrite graph names.
+    # We simulate the post-rewrite names by walking the rewrite chain
+    # WITHOUT actually invoking the rewriters — instead, we add the
+    # head-output names declared by the loss spec to the brick-name
+    # universe. This is a lenient check: the actual post-rewrite names
+    # land at apply_rewrites time. We only ERROR when the head_output
+    # references something neither the graph nor the rewrites can
+    # plausibly produce.
+    post_rewrite_names = {n.name for n in spec.graph.nodes}
+    if spec.rewrites:
+        # The MTPRewriter (and similar) declare provided_postconditions
+        # that hint at the head naming convention; we accept all
+        # head_outputs of the loss spec as future-valid.
+        post_rewrite_names |= set(spec.loss.head_outputs)
+    diags.extend(_check_loss_head_outputs(spec, post_rewrite_names))
+
+    # 3. Optimizer matcher coverage
+    diags.extend(_check_optim_matchers(spec))
+
+    # 4. Shape coherence
+    if check_shapes:
+        diags.extend(_check_shape_coherence(spec))
+
+    return BuildDiagnostics(diagnostics=tuple(diags))
+
+
+__all__ = [
+    "BuildDiagnostic",
+    "BuildDiagnosticSeverity",
+    "BuildDiagnostics",
+    "verify_build_spec",
+]