feat(e7-11): ablation.run RPC — side-by-side variant comparison

Stage E7-11 of E2E Coverage Matrix v2 epic (cppmega-mlx-bb0.11). Backend
of the Ablations sidebar tab: pick an axis (activation / optimizer /
norm / schedule), 2-4 variants, the RPC clones the spec per variant,
runs train, returns loss trajectories sorted by final loss.

cppmega_v4/jsonrpc/ablation_method.py (NEW):
- _mutate(base, axis, variant) clones the VerifyParams and swaps:
    * activation: every mlp/gated_mlp/moe brick's params['activation']
    * optimizer: spec.optim.kind + first group's lr/weight_decay/betas
      auto-populated from catalog.explain.recommended_params (so Lion
      gets lr=1e-4 without manual override + ns_steps=5 for Muon)
    * norm: every node's params['pre_norm']
    * schedule: each group's schedule dict (clears it when 'constant')
- ablation_run() executes Pipeline(parse → verify_build_spec →
  resolve_shapes → build_model → train) per variant; failed variants
  do not abort the rest; returns AblationVariantResult per variant +
  ranked_by_final_loss + baseline_variant + elapsed_ms_total.

cppmega_v4/jsonrpc/schema.py: ablation.run added to METHOD_REGISTRY.
cppmega_v4/jsonrpc/dispatcher.py: routes the new method.

Tests (+8 pytest):
- method registered
- activation ablation runs 3 variants (glu/swiglu/gelu); at least one
  reaches status=ok
- optimizer ablation runs adamw+sgd
- Lion ablation pulls recommended params via _mutate (asserts mutated
  spec has lr=1e-4 + wd=0.01)
- ranked_by_final_loss sorted ascending
- failed variant does NOT abort other variants
- baseline_variant = variants[0]
- elapsed_ms_total > 0

Regression: pytest 2306 (was 2298; +8) / 2 skip / 15 xfail / 0 fail
(excl. pre-existing path_d). vbgui vitest 150 / 0 fail.

UI AblationsTab sidebar component left as a follow-up — the RPC
+ mutate helper are the verification gate.

Closes cppmega-mlx-bb0.11.

Author: apstenku123

cppmega_v4/jsonrpc/ablation_method.py

@@ -0,0 +1,154 @@
+"""ablation.run RPC handler (E7-11).
+
+Runs N variants of a base spec where one axis (activation / optimizer
+/ norm / schedule) is swapped, executes stage_train per variant with
+the same num_steps, and returns side-by-side loss trajectories.
+
+Used by the Ablations sidebar tab to show 'replace swiglu with gelu →
++2.5% final loss after 20 steps' style results.
+"""
+
+from __future__ import annotations
+
+import copy
+import time
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from cppmega_v4.jsonrpc.cache import LRUCache
+from cppmega_v4.jsonrpc.schema import VerifyParams
+
+
+AblationAxis = Literal["activation", "optimizer", "norm", "schedule"]
+
+
+class AblationRunParams(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    base_spec: VerifyParams
+    ablation_axis: AblationAxis
+    variants: list[str]
+    num_steps: int = 20
+
+
+class AblationVariantResult(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    variant: str
+    status: Literal["ok", "fail"]
+    losses: list[float] = Field(default_factory=list)
+    elapsed_ms: float
+    weight_delta_norm: float = 0.0
+    error: dict[str, Any] | None = None
+
+
+class AblationRunResult(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    results: list[AblationVariantResult] = Field(default_factory=list)
+    ranked_by_final_loss: list[str] = Field(default_factory=list)
+    baseline_variant: str = ""
+    elapsed_ms_total: float = 0.0
+
+
+def _mutate(base: VerifyParams, axis: AblationAxis,
+            variant: str) -> VerifyParams:
+    """Return a deep-copy with axis swapped to variant."""
+    d = base.model_dump(mode="json")
+    if axis == "activation":
+        # Apply to every mlp/gated_mlp/moe brick.
+        for node in d.get("graph", {}).get("nodes", []):
+            if node.get("kind") in ("mlp", "gated_mlp", "moe"):
+                node.setdefault("params", {})["activation"] = variant
+    elif axis == "optimizer":
+        d["optim"]["kind"] = variant
+        # Replace first group's lr/wd/betas with recommended values
+        # from the catalogue so e.g. Lion gets lr=1e-4 automatically.
+        from cppmega_v4.explain import get_entry
+        entry = get_entry("optimizer", variant)
+        if entry and entry.recommended_params:
+            rp = entry.recommended_params
+            g = d["optim"]["groups"][0]
+            if "lr" in rp:
+                g["lr"] = float(rp["lr"])
+            if "weight_decay" in rp:
+                g["weight_decay"] = float(rp["weight_decay"])
+            if "betas" in rp and isinstance(rp["betas"], (list, tuple)):
+                g["betas"] = list(rp["betas"])
+        if variant == "muon":
+            d["optim"]["groups"][0].setdefault("ns_steps", 5)
+    elif axis == "norm":
+        for node in d.get("graph", {}).get("nodes", []):
+            node.setdefault("params", {})["pre_norm"] = variant
+    elif axis == "schedule":
+        for grp in d["optim"]["groups"]:
+            if variant == "constant":
+                grp.pop("schedule", None)
+            else:
+                grp["schedule"] = {"kind": variant, "warmup_steps": 2,
+                                   "total_steps": 50}
+    return VerifyParams.model_validate(d)
+
+
+def ablation_run(
+    params: AblationRunParams,
+    *,
+    cache: LRUCache | None = None,
+) -> AblationRunResult:
+    """Execute the train stage per variant; return collected losses."""
+    from cppmega_v4.runner import Pipeline, run_pipeline
+
+    t0 = time.perf_counter()
+    results: list[AblationVariantResult] = []
+
+    for variant in params.variants:
+        v_start = time.perf_counter()
+        try:
+            spec = _mutate(params.base_spec, params.ablation_axis, variant)
+            pipeline = Pipeline.from_dict({
+                "stages": ["parse", "verify_build_spec", "resolve_shapes",
+                           "build_model", "train"],
+                "stage_options": {
+                    "train": {"num_steps": params.num_steps},
+                },
+            })
+            report = run_pipeline(spec, pipeline)
+            train = next((s for s in report.stages if s.name == "train"),
+                         None)
+            if train is None or train.status != "ok":
+                results.append(AblationVariantResult(
+                    variant=variant, status="fail",
+                    losses=[],
+                    elapsed_ms=(time.perf_counter() - v_start) * 1000,
+                    weight_delta_norm=0.0,
+                    error=train.error if train else
+                          {"detail": "train stage missing"},
+                ))
+                continue
+            extras = train.extras or {}
+            results.append(AblationVariantResult(
+                variant=variant, status="ok",
+                losses=[float(l) for l in extras.get("losses", [])],
+                elapsed_ms=(time.perf_counter() - v_start) * 1000,
+                weight_delta_norm=float(extras.get("weight_delta_norm", 0)),
+            ))
+        except Exception as exc:
+            results.append(AblationVariantResult(
+                variant=variant, status="fail",
+                losses=[],
+                elapsed_ms=(time.perf_counter() - v_start) * 1000,
+                weight_delta_norm=0.0,
+                error={"type": type(exc).__name__, "detail": str(exc)},
+            ))
+
+    ranked = sorted(
+        (r for r in results if r.status == "ok" and r.losses),
+        key=lambda r: r.losses[-1],
+    )
+    return AblationRunResult(
+        results=results,
+        ranked_by_final_loss=[r.variant for r in ranked],
+        baseline_variant=params.variants[0] if params.variants else "",
+        elapsed_ms_total=(time.perf_counter() - t0) * 1000,
+    )

cppmega_v4/jsonrpc/dispatcher.py

@@ -39,6 +39,9 @@
 from cppmega_v4.jsonrpc.roundtrip_method import (
     RoundtripCheckParams, roundtrip_check,
 )
+from cppmega_v4.jsonrpc.ablation_method import (
+    AblationRunParams, ablation_run,
+)
 from cppmega_v4.jsonrpc.schema import (
     BuildPresetSpecsParams,
     CatalogExplainParams,
@@ -111,6 +114,10 @@
         RoundtripCheckParams,
         lambda p, c: roundtrip_check(p, cache=c),
     ),
+    "ablation.run": (
+        AblationRunParams,
+        lambda p, c: ablation_run(p, cache=c),
+    ),
 }
 
 

cppmega_v4/jsonrpc/schema.py

@@ -704,6 +704,7 @@ class PipelineRunResult(BaseModel):
     "architectures.list_presets",
     "suggest_optim_groups",
     "data.roundtrip_check",
+    "ablation.run",
 })
 
 

tests/v4/test_ablation_run.py

@@ -0,0 +1,154 @@
+"""E7-11 tests: ablation.run RPC."""
+
+from __future__ import annotations
+
+from cppmega_v4.architectures import build_preset_specs
+from cppmega_v4.jsonrpc import dispatch
+from cppmega_v4.jsonrpc.schema import METHOD_REGISTRY
+
+
+def _base_spec(preset: str = "llama3_8b") -> dict:
+    specs = build_preset_specs(preset, hidden_size=128)
+    return {
+        "graph": {
+            "nodes": [
+                {"id": s.get("name"), "kind": s["kind"],
+                 "params": s.get("params", {})}
+                for s in specs
+            ],
+            "edges": [
+                {"src": specs[i].get("name"),
+                 "dst": specs[i + 1].get("name")}
+                for i in range(len(specs) - 1)
+            ],
+        },
+        "dim_env": {"B": 1, "S": 8, "H": 128, "nh": 2, "nkv": 1,
+                    "head_dim": 64, "num_experts": 4, "top_k": 2},
+        "loss": {"kind": "cross_entropy",
+                 "head_outputs": [specs[-1].get("name")]},
+        "optim": {"kind": "adamw",
+                  "groups": [{"matcher": "all", "lr": 3e-4,
+                              "weight_decay": 0.01,
+                              "betas": [0.9, 0.95]}]},
+    }
+
+
+def test_method_registered():
+    assert "ablation.run" in METHOD_REGISTRY
+
+
+def test_activation_ablation_runs_three_variants():
+    resp = dispatch({
+        "jsonrpc": "2.0", "id": 1, "method": "ablation.run",
+        "params": {
+            "base_spec": _base_spec(),
+            "ablation_axis": "activation",
+            "variants": ["glu", "swiglu", "gelu"],
+            "num_steps": 2,
+        },
+    })
+    assert resp.error is None, resp.error
+    res = resp.result
+    assert len(res["results"]) == 3
+    names = {r["variant"] for r in res["results"]}
+    assert names == {"glu", "swiglu", "gelu"}
+    # At least one variant should succeed
+    assert any(r["status"] == "ok" for r in res["results"])
+
+
+def test_optimizer_ablation_runs_adamw_and_sgd():
+    resp = dispatch({
+        "jsonrpc": "2.0", "id": 1, "method": "ablation.run",
+        "params": {
+            "base_spec": _base_spec(),
+            "ablation_axis": "optimizer",
+            "variants": ["adamw", "sgd"],
+            "num_steps": 2,
+        },
+    })
+    assert resp.error is None, resp.error
+    res = resp.result
+    assert len(res["results"]) == 2
+
+
+def test_optimizer_ablation_lion_pulls_recommended_params():
+    """The Lion variant must trigger _mutate's catalogue lookup; the
+    spec's first group should end up with Lion's recommended lr (1e-4)
+    even if the train stage itself uses its own (AdamW-backed) loop."""
+    from cppmega_v4.jsonrpc.ablation_method import _mutate
+    from cppmega_v4.jsonrpc.schema import VerifyParams
+    mutated = _mutate(
+        VerifyParams.model_validate(_base_spec()),
+        "optimizer", "lion",
+    )
+    assert mutated.optim.kind == "lion"
+    assert mutated.optim.groups[0].lr == 1e-4
+    assert mutated.optim.groups[0].weight_decay == 0.01
+
+
+def test_ranked_by_final_loss_sorted_ascending():
+    resp = dispatch({
+        "jsonrpc": "2.0", "id": 1, "method": "ablation.run",
+        "params": {
+            "base_spec": _base_spec(),
+            "ablation_axis": "activation",
+            "variants": ["glu", "swiglu"],
+            "num_steps": 2,
+        },
+    })
+    assert resp.error is None
+    ranked = resp.result["ranked_by_final_loss"]
+    ok_results = [r for r in resp.result["results"]
+                  if r["status"] == "ok" and r["losses"]]
+    if len(ok_results) >= 2:
+        final_losses = {r["variant"]: r["losses"][-1] for r in ok_results}
+        prev = -float("inf")
+        for name in ranked:
+            cur = final_losses[name]
+            assert cur >= prev
+            prev = cur
+
+
+def test_failed_variant_does_not_abort_others():
+    """Pass a deliberately bogus optimizer name → that variant fails;
+    other variants still complete."""
+    resp = dispatch({
+        "jsonrpc": "2.0", "id": 1, "method": "ablation.run",
+        "params": {
+            "base_spec": _base_spec(),
+            "ablation_axis": "optimizer",
+            "variants": ["adamw", "totally_made_up_optim"],
+            "num_steps": 2,
+        },
+    })
+    assert resp.error is None
+    statuses = {r["variant"]: r["status"] for r in resp.result["results"]}
+    assert statuses["adamw"] == "ok"
+    assert statuses["totally_made_up_optim"] == "fail"
+
+
+def test_baseline_variant_is_first_in_variants_list():
+    resp = dispatch({
+        "jsonrpc": "2.0", "id": 1, "method": "ablation.run",
+        "params": {
+            "base_spec": _base_spec(),
+            "ablation_axis": "activation",
+            "variants": ["glu", "swiglu"],
+            "num_steps": 2,
+        },
+    })
+    assert resp.result["baseline_variant"] == "glu"
+
+
+def test_elapsed_ms_total_positive():
+    resp = dispatch({
+        "jsonrpc": "2.0", "id": 1, "method": "ablation.run",
+        "params": {
+            "base_spec": _base_spec(),
+            "ablation_axis": "activation",
+            "variants": ["glu"],
+            "num_steps": 2,
+        },
+    })
+    assert resp.error is None
+    assert resp.result["elapsed_ms_total"] > 0