feat(v7-h33): loss_surface.run dispatcher route + LossSurfaceModal UI

apstenku123 · apstenku123 · commit eed403218b96 · 2026-05-23T03:51:24.000+02:00
Closes the orphan-backend half of item 15 (loss_surface): the
loss_surface_method.py file existed but no dispatcher route, so the
UI got -32601 method_not_found. UI also had no consumer.

- Register loss_surface.run in dispatcher._ROUTES.
- New LossSurfaceModal: N×M heatmap with viridis-style coloring,
  best cell outlined yellow, Apply-best button commits lr/wd
  multipliers back into spec.optim.groups via dispatch.
- App.tsx wires a floating 'Loss surface…' launcher button.

bd: cppmega-mlx-sahc
2/2 pytest (dispatcher route + grid matrix) + 4/4 vitest.
diff --git a/cppmega_v4/jsonrpc/dispatcher.py b/cppmega_v4/jsonrpc/dispatcher.py
@@ -62,6 +62,10 @@
     SideChannelApplyParams,
     apply_side_channels,
 )
+from cppmega_v4.jsonrpc.loss_surface_method import (
+    LossSurfaceParams,
+    loss_surface_run,
+)
 from cppmega_v4.jsonrpc.schema import (
     BuildPresetSpecsParams,
     CatalogExplainParams,
@@ -182,6 +186,10 @@
         SideChannelApplyParams,
         lambda p, c: apply_side_channels(p, cache=c),
     ),
+    "loss_surface.run": (
+        LossSurfaceParams,
+        lambda p, c: loss_surface_run(p, cache=c),
+    ),
 }
 
 
diff --git a/tests/v4/test_loss_surface_dispatcher.py b/tests/v4/test_loss_surface_dispatcher.py
@@ -0,0 +1,66 @@
+"""V7-H33: loss_surface.run RPC dispatcher route registration."""
+
+from __future__ import annotations
+
+from cppmega_v4.jsonrpc.dispatcher import dispatch
+
+
+_SPEC = {
+    "graph": {
+        "nodes": [
+            {"id": "attn", "kind": "attention",
+             "params": {"num_heads": 4, "head_dim": 64}},
+            {"id": "mlp", "kind": "mlp", "params": {}},
+        ],
+        "edges": [{"src": "attn", "dst": "mlp"}],
+    },
+    "dim_env": {"B": 1, "S": 8, "H": 128,
+                "nh": 2, "nkv": 1, "head_dim": 64},
+    "loss": {"kind": "cross_entropy", "head_outputs": ["mlp"]},
+    "optim": {"kind": "adamw",
+              "groups": [{"matcher": "all", "lr": 1e-3,
+                          "weight_decay": 0.01,
+                          "betas": [0.9, 0.95]}]},
+}
+
+
+def test_v7_h33_loss_surface_run_method_is_routed():
+    """Before this fix loss_surface_method.py existed but the route
+    was orphan — UI got -32601 method_not_found."""
+    response = dispatch({
+        "jsonrpc": "2.0", "id": "T1", "method": "loss_surface.run",
+        "params": {
+            "spec": _SPEC,
+            "lr_deltas": [1.0],
+            "wd_deltas": [1.0],
+            "k_steps": 2,
+        },
+    })
+    assert response.error is None, response.error
+    assert response.result is not None
+    assert "rows" in response.result
+    assert len(response.result["rows"]) == 1
+    assert response.result["lr_deltas"] == [1.0]
+    assert response.result["wd_deltas"] == [1.0]
+    assert response.result["best_lr_mult"] == 1.0
+    assert response.result["best_wd_mult"] == 1.0
+    assert response.result["best_loss"] is not None
+
+
+def test_v7_h33_loss_surface_run_grid_returns_full_matrix():
+    response = dispatch({
+        "jsonrpc": "2.0", "id": "T2", "method": "loss_surface.run",
+        "params": {
+            "spec": _SPEC,
+            "lr_deltas": [0.5, 1.0],
+            "wd_deltas": [0.5, 1.0, 2.0],
+            "k_steps": 2,
+        },
+    })
+    assert response.error is None
+    rows = response.result["rows"]
+    assert len(rows) == 2
+    assert all(len(r) == 3 for r in rows)
+    for row in rows:
+        for cell in row:
+            assert cell["status"] in {"ok"} or cell["status"].startswith("fail")
diff --git a/vbgui/src/App.tsx b/vbgui/src/App.tsx
@@ -21,6 +21,7 @@ import { TopBar, type RunMode } from "@/components/TopBar";
 import { BottomStrip } from "@/components/BottomStrip";
 import { AppTabs, type AppTab } from "@/components/AppTabs";
 import { RunResultModal, type RunReport } from "@/components/RunResultModal";
+import { LossSurfaceModal } from "@/components/LossSurfaceModal";
 import { TokenizerPlayground } from "@/components/TokenizerPlayground";
 import { DataInspector } from "@/components/DataInspector";
 import { BrickContextPanel } from "@/components/BrickContextPanel";
@@ -138,6 +139,8 @@ export function App(): JSX.Element {
     setRunErrorRaw({ code, message, data });
   }, []);
   const [trainInFlight, setTrainInFlight] = useState(false);
+  // V7-H33: loss-surface explorer modal open flag.
+  const [lossSurfaceOpen, setLossSurfaceOpen] = useState<boolean>(false);
   // V7-I03: synchronous lock. React's setTrainInFlight schedules an
   // async commit, so two button clicks within the same microtask
   // both read trainInFlight=false and both call rpc.call. The ref
@@ -1323,6 +1326,34 @@ export function App(): JSX.Element {
         <RunResultModal report={runReport} error={runError}
                         onClose={() => { setRunReport(null);
                                          setRunError(null); }} />
+        {/* V7-H33: loss-surface explorer launcher + modal. */}
+        <button data-testid="loss-surface-open"
+                onClick={() => setLossSurfaceOpen(true)}
+                style={{ position: "fixed", bottom: 56, right: 18,
+                         padding: "6px 10px", fontSize: 12,
+                         background: "#eef2ff",
+                         border: "1px solid #c7d2fe", borderRadius: 6,
+                         fontFamily: "system-ui, sans-serif" }}>
+          Loss surface…
+        </button>
+        <LossSurfaceModal
+          rpc={rpc}
+          spec={buildVerifyParams(
+            nodes, edges, spec, availableSideChannels)}
+          open={lossSurfaceOpen}
+          onClose={() => setLossSurfaceOpen(false)}
+          onApplyBest={(lrMult, wdMult) => {
+            // V7-H33: scale every optim group's lr+wd by the chosen
+            // multipliers and commit through dispatch so verify+history
+            // see the change.
+            const nextGroups = spec.optim.groups.map((g) => ({
+              ...g,
+              lr: g.lr * lrMult,
+              weight_decay: (g.weight_decay ?? 0) * wdMult,
+            }));
+            dispatch({ type: "optim.set",
+                       optim: { ...spec.optim, groups: nextGroups } });
+          }} />
       </div>
     </ReactFlowProvider>
   );
diff --git a/vbgui/src/components/LossSurfaceModal.tsx b/vbgui/src/components/LossSurfaceModal.tsx
@@ -0,0 +1,215 @@
+// V7-H33: LossSurfaceModal — N×M heatmap of (lr_mult × wd_mult)
+// triggered by RunResultModal "Explore loss surface" button. Calls
+// loss_surface.run RPC, renders the best cell + Apply-best button that
+// commits the chosen lr/wd multipliers back to the optim spec.
+
+import { useState } from "react";
+import type { RpcClient } from "@/lib/rpc";
+import { HelpIcon } from "@/components/HelpIcon";
+
+export interface LossSurfaceCell {
+  lr_mult: number;
+  wd_mult: number;
+  status: string;
+  final_loss?: number | null;
+  throughput_tok_s?: number | null;
+  mem_mb?: number | null;
+  elapsed_ms?: number;
+}
+
+export interface LossSurfaceResult {
+  rows: LossSurfaceCell[][];
+  lr_deltas: number[];
+  wd_deltas: number[];
+  best_lr_mult: number | null;
+  best_wd_mult: number | null;
+  best_loss: number | null;
+}
+
+export interface LossSurfaceModalProps {
+  rpc: RpcClient | null;
+  spec: unknown;
+  open: boolean;
+  onClose: () => void;
+  onApplyBest: (lrMult: number, wdMult: number) => void;
+}
+
+const DEFAULT_LR_DELTAS = [0.5, 1.0, 2.0];
+const DEFAULT_WD_DELTAS = [0.5, 1.0, 2.0];
+
+function cellColor(loss: number | null | undefined,
+                    min: number | null, max: number | null): string {
+  if (loss == null || min == null || max == null || max <= min) {
+    return "#f3f4f6";
+  }
+  const t = (loss - min) / (max - min);
+  const r = Math.round(40 + 215 * t);
+  const g = Math.round(180 - 120 * t);
+  return `rgb(${r}, ${g}, 60)`;
+}
+
+export function LossSurfaceModal({
+  rpc, spec, open, onClose, onApplyBest,
+}: LossSurfaceModalProps): JSX.Element | null {
+  const [running, setRunning] = useState(false);
+  const [result, setResult] = useState<LossSurfaceResult | null>(null);
+  const [error, setError] = useState<string | null>(null);
+  const [kSteps, setKSteps] = useState<number>(2);
+
+  if (!open) return null;
+
+  async function run() {
+    setRunning(true);
+    setError(null);
+    setResult(null);
+    if (!rpc) {
+      setError("no backend connection");
+      setRunning(false);
+      return;
+    }
+    try {
+      const r = await rpc.call<LossSurfaceResult>("loss_surface.run", {
+        spec,
+        lr_deltas: DEFAULT_LR_DELTAS,
+        wd_deltas: DEFAULT_WD_DELTAS,
+        k_steps: kSteps,
+      });
+      setResult(r);
+    } catch (e) {
+      setError(e instanceof Error ? e.message : String(e));
+    } finally {
+      setRunning(false);
+    }
+  }
+
+  const allLosses = result?.rows.flat()
+    .map((c) => c.final_loss)
+    .filter((v): v is number => v != null) ?? [];
+  const minLoss = allLosses.length ? Math.min(...allLosses) : null;
+  const maxLoss = allLosses.length ? Math.max(...allLosses) : null;
+
+  return (
+    <div data-testid="loss-surface-modal"
+         role="dialog" aria-modal="true"
+         style={{ position: "fixed", inset: 0, zIndex: 1000,
+                  background: "rgba(0,0,0,0.45)", display: "flex",
+                  alignItems: "center", justifyContent: "center",
+                  padding: 24, fontFamily: "system-ui, sans-serif" }}>
+      <div style={{ background: "white", borderRadius: 8, padding: 18,
+                    minWidth: 460, maxWidth: 720, maxHeight: "85vh",
+                    overflow: "auto" }}>
+        <div style={{ display: "flex", justifyContent: "space-between",
+                      alignItems: "center", marginBottom: 8 }}>
+          <h3 data-testid="loss-surface-title"
+              style={{ margin: 0, fontSize: 16 }}>
+            Loss surface · lr_mult × wd_mult
+            <HelpIcon topic="loss_surface_explorer" />
+          </h3>
+          <button data-testid="loss-surface-close"
+                  onClick={onClose}>×</button>
+        </div>
+
+        <div style={{ fontSize: 12, marginBottom: 8 }}>
+          <label style={{ marginRight: 8 }}>k_steps:</label>
+          <input data-testid="loss-surface-k-steps"
+                 type="number" min={1} max={64} value={kSteps}
+                 onChange={(e) =>
+                   setKSteps(Math.max(1, Math.min(64,
+                     Number(e.target.value) || 2)))}
+                 style={{ width: 60 }} />
+          <button data-testid="loss-surface-run"
+                  onClick={() => void run()}
+                  disabled={running}
+                  style={{ marginLeft: 10 }}>
+            {running ? "Running…" : "Run sweep"}
+          </button>
+        </div>
+
+        {error && (
+          <div data-testid="loss-surface-error"
+               style={{ color: "#b91c1c", fontSize: 12 }}>
+            {error}
+          </div>
+        )}
+
+        {result && (
+          <div data-testid="loss-surface-result">
+            <table style={{ borderCollapse: "collapse",
+                            margin: "8px 0" }}>
+              <thead>
+                <tr>
+                  <th style={{ padding: "4px 8px", fontSize: 11 }}>
+                    lr↓ / wd→
+                  </th>
+                  {result.wd_deltas.map((wd) => (
+                    <th key={wd}
+                        data-testid={`loss-surface-col-${wd}`}
+                        style={{ padding: "4px 8px", fontSize: 11 }}>
+                      ×{wd}
+                    </th>
+                  ))}
+                </tr>
+              </thead>
+              <tbody>
+                {result.rows.map((row, ri) => (
+                  <tr key={ri}>
+                    <th data-testid={`loss-surface-row-${result.lr_deltas[ri]}`}
+                        style={{ padding: "4px 8px", fontSize: 11 }}>
+                      ×{result.lr_deltas[ri]}
+                    </th>
+                    {row.map((cell, ci) => {
+                      const isBest = result.best_lr_mult === cell.lr_mult
+                                   && result.best_wd_mult === cell.wd_mult;
+                      return (
+                        <td key={ci}
+                            data-testid={`loss-surface-cell-${ri}-${ci}`}
+                            style={{
+                              padding: "10px 14px",
+                              background: cellColor(cell.final_loss,
+                                                   minLoss, maxLoss),
+                              color: "white", fontSize: 11,
+                              border: isBest
+                                ? "3px solid #facc15" : "1px solid #e5e7eb",
+                              textAlign: "center",
+                              minWidth: 56,
+                            }}>
+                          {cell.status === "ok" && cell.final_loss != null
+                            ? cell.final_loss.toFixed(3)
+                            : cell.status}
+                        </td>
+                      );
+                    })}
+                  </tr>
+                ))}
+              </tbody>
+            </table>
+
+            <div data-testid="loss-surface-best"
+                 style={{ fontSize: 12, marginTop: 4 }}>
+              {result.best_loss != null ? (
+                <>
+                  best: lr×{result.best_lr_mult}, wd×{result.best_wd_mult}
+                  {" → loss "}{result.best_loss.toFixed(4)}
+                </>
+              ) : (
+                "no successful cells"
+              )}
+            </div>
+
+            {result.best_lr_mult != null && result.best_wd_mult != null && (
+              <button data-testid="loss-surface-apply-best"
+                      onClick={() => {
+                        onApplyBest(result.best_lr_mult!,
+                                    result.best_wd_mult!);
+                        onClose();
+                      }}
+                      style={{ marginTop: 8 }}>
+                Apply best to optim
+              </button>
+            )}
+          </div>
+        )}
+      </div>
+    </div>
+  );
+}
diff --git a/vbgui/tests/LossSurfaceModal.test.tsx b/vbgui/tests/LossSurfaceModal.test.tsx
