feat(v4/buildspec): MBSpec Stage C — Rewriter protocol + MTPRewriter

apstenku123 · apstenku123 · commit da6f8623357c · 2026-05-20T16:36:59.000+02:00
Adds the first concrete graph rewriter (ModelBuildSpec.md §3.4, §6 C).

cppmega_v4/buildspec/rewriters.py:
  - MTPRewriter(k, beta=None, share_backbone=True, add_head_param_group=True)
    Multi-Token-Prediction rewrite. K=1 is no-op fast path; K≥2:
      * locates the head brick (explicit params['is_head']=True wins;
        falls back to last mlp/attention node in graph order)
      * renames it to &lt;head&gt;_0 and clones K-1 more (&lt;head&gt;_1 .. &lt;head&gt;_{k-1})
        with weight modules cleared (build-time re-init)
      * rewires every edge into the old head onto every new head_i so
        each clone receives the same backbone activations
      * rewrites loss CE → mtp_weighted with the original head name as
        prefix, honouring optional custom beta tuple
      * optionally prepends a regex:.*&lt;head&gt;_\d+$ ParamGroup to the
        optimizer, inheriting betas (AdamW) / ns_steps (Muon) from the
        first existing group
    Protocol-conformant: name + required_preconditions={"single_head"} +
    provided_postconditions={"mtp_k_heads"} (empty for K=1 no-op).
  - HeadDetectionError when no head candidate exists.
  - LossRewriteError when input loss isn't CE (auto-conversion unsafe).

cppmega_v4/buildspec/model_build_spec.py:
  - Rewriter Protocol marked @runtime_checkable so isinstance(r, Rewriter)
    works for tests / GUI lookups.

Tests (tests/v4/test_mbspec_stage_c.py, 25 cases):
  - protocol conformance (isinstance check); name/precondition/post
    advertised correctly; K=1 advertises no post; frozen dataclass
  - construction rejection: k&lt;1; mismatched beta length
  - K=1 fast-path returns input
  - K=2 / K=3 graph rewrite: head node count, edge wiring (producer →
    every clone), backbone untouched
  - loss rewrite: CE → mtp_weighted with k/beta/head_outputs; custom
    head-name prefix; non-CE input raises LossRewriteError
  - optim rewrite: head-only group prepended with regex matcher;
    inheritance of betas (AdamW) / ns_steps (Muon) / neither (SGD);
    opt-out via add_head_param_group=False
  - head detection: explicit is_head marker beats kind heuristic; no
    candidate raises HeadDetectionError
  - composition: apply_rewrites adds mtp_k_heads state token;
    verify_build_spec clean on Qwen3-Next with MTPRewriter
  - system: K=2 doubles head count only; share_backbone=False keeps
    backbone intact (Stage C scope — backbone replication is future work)

Full v4 regression: 858 passed / 5 skipped / 0 failed.
diff --git a/cppmega_v4/buildspec/__init__.py b/cppmega_v4/buildspec/__init__.py
@@ -35,6 +35,11 @@
     RewriteOrderError,
     Rewriter,
 )
+from cppmega_v4.buildspec.rewriters import (
+    HeadDetectionError,
+    LossRewriteError,
+    MTPRewriter,
+)
 from cppmega_v4.buildspec.optim_spec import (
     OPTIM_BUILTINS,
     OptimKind,
@@ -50,7 +55,10 @@
     "BuildDiagnostic",
     "BuildDiagnosticSeverity",
     "BuildDiagnostics",
+    "HeadDetectionError",
     "LOSS_BUILTINS",
+    "LossRewriteError",
+    "MTPRewriter",
     "LossKind",
     "LossSpec",
     "ModelBuildSpec",
diff --git a/cppmega_v4/buildspec/model_build_spec.py b/cppmega_v4/buildspec/model_build_spec.py
@@ -15,7 +15,7 @@
 
 from collections.abc import Mapping
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Protocol
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
 
 from cppmega_v4.buildspec.loss_spec import LossSpec
 from cppmega_v4.buildspec.optim_spec import OptimSpec
@@ -25,6 +25,7 @@
     pass
 
 
+@runtime_checkable
 class Rewriter(Protocol):
     """A pure :class:`ModelBuildSpec` transformation.
 
diff --git a/cppmega_v4/buildspec/rewriters.py b/cppmega_v4/buildspec/rewriters.py
@@ -0,0 +1,276 @@
+"""Concrete graph rewriters.
+
+Stage C scope (this commit): :class:`MTPRewriter`.
+
+A rewriter is a callable ``(ModelBuildSpec) -> ModelBuildSpec`` that
+matches the :class:`cppmega_v4.buildspec.Rewriter` Protocol. It:
+
+  - reads the input graph + loss + optim
+  - materialises a new graph (adding/removing nodes / edges)
+  - rewrites the loss spec to match the new outputs (e.g. CE → MTP-weighted)
+  - optionally adds a new optimizer parameter group (e.g. head-only lr)
+
+Rewriters are PURE — they never mutate the input spec; they return a
+fresh one via :meth:`ModelBuildSpec.replace`.
+
+MTPRewriter
+-----------
+
+Multi-Token Prediction rewrite. Given a spec with a single head brick
+(``state_token "single_head"``), materialises K-1 additional head copies
+and rewrites the loss to :func:`mtp_weighted_loss`. The new state token
+``"mtp_k_heads"`` advertises the post-condition for later rewriters.
+
+Naming convention: the original head brick is renamed
+``<head>_0`` and copies become ``<head>_1 ... <head>_{k-1}``. Loss
+head_outputs match this scheme (``logits_0 / logits_1 / ...``).
+
+Param-group integration: when ``add_head_param_group=True`` (default
+True), the optimizer gains a new high-priority ``regex:.*_head_\\d+$``
+group with the same lr as the base — callers can override the head-only
+lr after the fact via :meth:`OptimSpec.replace`.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Final
+
+from cppmega_v4.buildspec.loss_spec import (
+    LossKind,
+    LossSpec,
+    mtp_weighted_loss,
+)
+from cppmega_v4.buildspec.model_build_spec import (
+    ModelBuildSpec,
+    Rewriter,
+)
+from cppmega_v4.buildspec.optim_spec import (
+    OptimKind,
+    OptimSpec,
+    ParamGroup,
+)
+from cppmega_v4.fusion.brick_graph import BrickGraph, BrickNode
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+_HEAD_KINDS: Final[frozenset[str]] = frozenset({
+    # bricks that produce the "logits" head — Stage C scope. Add more
+    # as we wire heads explicitly. For now we recognise mlp-as-head and
+    # accept an explicit ``is_head=True`` attribute on a BrickNode.
+    "mlp",
+    "attention",
+})
+
+
+def _find_head_node(graph: BrickGraph) -> BrickNode:
+    """Locate the head brick.
+
+    Resolution order:
+      1. Last node in the graph whose ``params`` contains ``"is_head": True``.
+      2. Last node in the graph whose kind is in :data:`_HEAD_KINDS`.
+      3. Else: raises :class:`HeadDetectionError`.
+    """
+    explicit = [
+        n for n in graph.nodes if n.params.get("is_head") is True
+    ]
+    if explicit:
+        return explicit[-1]
+    typed = [n for n in graph.nodes if n.kind in _HEAD_KINDS]
+    if not typed:
+        raise HeadDetectionError(
+            "no head brick found — annotate a brick with "
+            "`params={'is_head': True}` or end the graph with one of "
+            f"{sorted(_HEAD_KINDS)}"
+        )
+    return typed[-1]
+
+
+class HeadDetectionError(RuntimeError):
+    """Raised by :class:`MTPRewriter` when no head brick can be located."""
+
+
+# ---------------------------------------------------------------------------
+# MTPRewriter
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class MTPRewriter:
+    """K-head Multi-Token-Prediction rewriter.
+
+    Fields:
+      k: number of prediction heads (≥ 1). K=1 is a no-op (returns input
+        spec unchanged besides the postcondition token).
+      beta: optional per-head weights for the mtp_weighted loss. When
+        None, falls back to the default geometric decay used by
+        :func:`mtp_weighted_loss`.
+      share_backbone: when True (default), the backbone bricks are kept
+        as-is and only the head is replicated K times. When False, the
+        original head node is renamed ``_0`` and new ``_1.._{k-1}``
+        clones are added (still pointing into the backbone end-node).
+      add_head_param_group: when True (default), the optimizer spec gets
+        a new head-only param group with the same lr as the first
+        existing group. Useful when the GUI wants to surface a separate
+        slider for head lr.
+    """
+
+    k: int = 2
+    beta: tuple[float, ...] | None = None
+    share_backbone: bool = True
+    add_head_param_group: bool = True
+
+    # Rewriter protocol
+    name: str = field(init=False)
+    required_preconditions: frozenset[str] = field(init=False)
+    provided_postconditions: frozenset[str] = field(init=False)
+
+    def __post_init__(self) -> None:
+        if self.k < 1:
+            raise ValueError(f"MTPRewriter.k must be ≥ 1, got {self.k}")
+        if self.beta is not None and len(self.beta) != self.k:
+            raise ValueError(
+                f"MTPRewriter.beta length ({len(self.beta)}) must equal "
+                f"k ({self.k})"
+            )
+        # Init the Rewriter-protocol fields.
+        object.__setattr__(self, "name", f"MTPRewriter(k={self.k})")
+        object.__setattr__(
+            self, "required_preconditions", frozenset({"single_head"}),
+        )
+        object.__setattr__(
+            self, "provided_postconditions",
+            frozenset() if self.k == 1 else frozenset({"mtp_k_heads"}),
+        )
+
+    def __call__(self, spec: ModelBuildSpec) -> ModelBuildSpec:
+        if self.k == 1:
+            # No-op fast path. Don't touch the loss either — caller asked
+            # for K=1 explicitly.
+            return spec
+
+        if spec.loss.kind is not LossKind.CROSS_ENTROPY:
+            raise LossRewriteError(
+                f"MTPRewriter requires LossKind.CROSS_ENTROPY input "
+                f"(got {spec.loss.kind!r}) — automatic rewrite to "
+                "mtp_weighted only safe from a pure CE base"
+            )
+
+        head = _find_head_node(spec.graph)
+        original_loss_name = spec.loss.head_outputs[0]
+
+        # ---- rewrite graph: rename head -> head_0; add head_1..k-1 ----
+        new_nodes: list[BrickNode] = []
+        existing_names = {n.name for n in spec.graph.nodes}
+        existing_names.discard(head.name)
+
+        renamed_head = BrickNode(
+            kind=head.kind,
+            name=f"{head.name}_0",
+            params=dict(head.params),
+            module=head.module,
+        )
+        head_names: list[str] = [renamed_head.name]
+        for n in spec.graph.nodes:
+            if n.name == head.name:
+                new_nodes.append(renamed_head)
+            else:
+                new_nodes.append(n)
+
+        for i in range(1, self.k):
+            clone_name = f"{head.name}_{i}"
+            # Ensure unique
+            if clone_name in existing_names:
+                clone_name = f"{head.name}_mtp_{i}"
+            existing_names.add(clone_name)
+            new_nodes.append(
+                BrickNode(
+                    kind=head.kind,
+                    name=clone_name,
+                    params=dict(head.params),
+                    module=None,  # weights re-init at build time
+                )
+            )
+            head_names.append(clone_name)
+
+        # ---- rewrite edges: every edge into old head_name -> head_0 ----
+        producer_of_head = [
+            p for p, c in spec.graph.edges if c == head.name
+        ]
+        new_edges: list[tuple[str, str]] = []
+        for p, c in spec.graph.edges:
+            if c == head.name:
+                new_edges.append((p, renamed_head.name))
+                # Also wire the producer into every new head clone.
+                for clone in head_names[1:]:
+                    new_edges.append((p, clone))
+            elif p == head.name:
+                # Anything that consumed the old head now consumes head_0
+                # (the "real" prediction); MTP heads are leaves that
+                # only feed the loss, not subsequent bricks.
+                new_edges.append((renamed_head.name, c))
+            else:
+                new_edges.append((p, c))
+
+        # If the head had no producers (size-1 graph), wire the clones
+        # directly into the head_0 slot — they'll consume the same input
+        # the head_0 brick gets at runtime.
+        if not producer_of_head and self.k > 1:
+            for clone in head_names[1:]:
+                new_edges.append((renamed_head.name, clone))
+
+        new_graph = BrickGraph(nodes=tuple(new_nodes), edges=tuple(new_edges))
+
+        # ---- rewrite loss: CE -> mtp_weighted -------------------------
+        new_loss = mtp_weighted_loss(
+            k=self.k,
+            beta=self.beta,
+            head_output_prefix=original_loss_name,
+        )
+
+        # ---- rewrite optim: optionally add head-only group ------------
+        new_optim = spec.optim
+        if self.add_head_param_group:
+            base_lr = spec.optim.groups[0].lr
+            head_matcher = f"regex:.*{original_loss_name}_\\d+$"
+            head_group = ParamGroup(
+                matcher=head_matcher,
+                lr=base_lr,
+                weight_decay=spec.optim.groups[0].weight_decay,
+                betas=(
+                    spec.optim.groups[0].betas
+                    if spec.optim.kind in {OptimKind.ADAMW, OptimKind.MUON_ADAMW_HYBRID}
+                    else None
+                ),
+                ns_steps=(
+                    spec.optim.groups[0].ns_steps
+                    if spec.optim.kind is OptimKind.MUON
+                    else None
+                ),
+            )
+            # Prepend so the head matcher wins over the catch-all "all".
+            new_optim = OptimSpec(
+                kind=spec.optim.kind,
+                groups=(head_group, *spec.optim.groups),
+                gradient_clip_norm=spec.optim.gradient_clip_norm,
+                mixed_precision=spec.optim.mixed_precision,
+            )
+
+        return spec.replace(
+            graph=new_graph, loss=new_loss, optim=new_optim,
+        )
+
+
+class LossRewriteError(RuntimeError):
+    """Raised by MTPRewriter when the input loss can't be safely converted."""
+
+
+__all__ = [
+    "HeadDetectionError",
+    "LossRewriteError",
+    "MTPRewriter",
+]
diff --git a/tests/v4/test_mbspec_stage_c.py b/tests/v4/test_mbspec_stage_c.py
