feat(e-audit-06): train UI polish — B/S controls + loss_scaler nesting + plasticity panels

apstenku123 · apstenku123 · commit 2144c510ed3d · 2026-05-23T16:55:48.000+02:00
V7-Q06 (E-AUDIT-06, cppmega-mlx-sbz8) — Lane 5 audit gap polish.

Q06.1 B and S as explicit TrainOptionsPanel controls:
- TrainOptions adds train_B + train_S optional ints
- TrainOptionsPanel exposes train-opt-B + train-opt-S inputs
- App.tsx threads them into stage_options.train.{B,S} which
  stage_train already reads via opts.get('B', 1) / opts.get('S', 8)
- Default unset preserves prior behaviour (spec.dim_env wins).

Q06.2 loss_scaler_overflows flat key:
- stages.py:2284 adds extras['loss_scaler_overflows'] = list(steps)
  alongside the existing nested loss_scaler.overflow_steps so the
  LossChart overflow markers actually render. Empty list when no
  scaler configured (preserves UI expectation).

Q06.3 styled sub-panels for plasticity / mtp / ifim / mhc extras:
- TrainExtras interface typed plasticity/mtp/ifim/mhc fields
- TrainExtrasOverlay conditionally renders 4 new panels:
  - PlasticityPanel (FIRE fire_step + fire_keys, DASH count, ReDo)
  - MTPPanel (k, beta, head_losses per-K)
  - IFIMPanel (instr_loss, instr_token_count, lambda)
  - MHCPanel (consistency_loss, heads_correlated, lambda)
- Panels render only when extras carry the data, so existing
  cross-entropy runs stay visually identical.

Tests (all green):
- tests/v4/test_loss_scaler_overflows_flat.py: 2 cases — no scaler
  -&gt; empty list, fp16 scaler -&gt; nested+flat coexist
- vbgui/tests/TrainOptionsPanelBS.test.tsx: 4 cases — render,
  onChange B/S, empty clears
- vbgui/tests/TrainExtrasOverlayPanels.test.tsx: 6 cases — hidden
  by default + each panel renders correctly

Regression:
- Backend pytest (Q01-Q06 combined): 242/242 PASS
- vbgui vitest: 440 -&gt; 451 (+11) PASS

Closes Lane 5 P2 audit gap from docs/UI-TO-TRAIN-AUDIT-2026-05-23.md.
diff --git a/cppmega_v4/runner/stages.py b/cppmega_v4/runner/stages.py
@@ -2287,6 +2287,12 @@ def _compiled_step_value_and_grad(_emb, _targets):
                                    "overflow_count":
                                    loss_scaler.overflow_count}
                                  if loss_scaler is not None else None),
+                # V7-Q06.2: flatten overflow steps so LossChart can
+                # render overflow markers without diving into the
+                # nested dict. Keeps nested key for back-compat.
+                "loss_scaler_overflows": (
+                    list(loss_scaler_overflow_steps)
+                    if loss_scaler is not None else []),
                 "moe": moe_extras,
                 "opt_state_carried": opt_state_carried,
                 "run_id": run_id,
diff --git a/tests/v4/test_loss_scaler_overflows_flat.py b/tests/v4/test_loss_scaler_overflows_flat.py
@@ -0,0 +1,72 @@
+"""V7-Q06.2: pin the flat extras.loss_scaler_overflows key.
+
+The nested `loss_scaler.overflow_steps` was preserved for back-compat,
+but UI reads the flat key — assert both shapes coexist when a scaler
+is configured, and the flat key is an empty list when no scaler.
+"""
+
+from __future__ import annotations
+
+from cppmega_v4.jsonrpc.schema import VerifyParams
+from cppmega_v4.runner import Pipeline, run_pipeline
+
+
+def _spec(loss_scaler: bool) -> VerifyParams:
+    optim = {
+        "kind": "adamw",
+        "groups": [{"matcher": "all", "lr": 1e-3,
+                    "weight_decay": 0.01, "betas": [0.9, 0.95]}],
+    }
+    return VerifyParams.model_validate({
+        "graph": {
+            "nodes": [
+                {"id": "attn", "kind": "attention",
+                 "params": {"num_heads": 2, "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": optim,
+    })
+
+
+def _run(opts: dict) -> dict:
+    rep = run_pipeline(_spec(loss_scaler=True), Pipeline.from_dict({
+        "stages": ["parse", "verify_build_spec", "build_model", "train"],
+        "stage_options": {"train": opts},
+    }))
+    tr = next(s for s in rep.stages if s.name == "train")
+    assert tr.status == "ok", tr.error
+    return dict(tr.extras or {})
+
+
+def test_no_scaler_flat_key_empty() -> None:
+    """No loss_scaler configured -> extras.loss_scaler is None +
+    extras.loss_scaler_overflows is []."""
+    extras = _run({"num_steps": 2, "seed": 7})
+    assert extras.get("loss_scaler") is None
+    # Flat key always present (empty list when no scaler).
+    assert extras.get("loss_scaler_overflows") == []
+
+
+def test_scaler_active_both_keys_coexist() -> None:
+    """When fp16 master_dtype triggers loss_scaler, both nested and
+    flat keys hold the same overflow_steps list."""
+    extras = _run({
+        "num_steps": 3, "seed": 7,
+        "master_dtype": "fp16",
+        "loss_scaler_mode": "dynamic",
+        "loss_scaler_init": 2.0,
+    })
+    nested = extras.get("loss_scaler")
+    flat = extras.get("loss_scaler_overflows")
+    assert nested is not None, (
+        f"expected loss_scaler dict when fp16 active; got extras keys "
+        f"{sorted(extras)}"
+    )
+    # Nested carries the overflow_steps list; flat must mirror it.
+    assert nested.get("overflow_steps") == flat
+    assert isinstance(flat, list)
diff --git a/vbgui/src/App.tsx b/vbgui/src/App.tsx
@@ -588,6 +588,9 @@ export function App(): JSX.Element {
       inference_probe_text?: string;
       master_dtype?: "fp32" | "bf16" | "fp16" | "auto";
       fim_enabled?: boolean;
+      compress?: string;
+      ckpt_strict?: boolean;
+      opt_state_strict?: boolean;
     },
   ) => {
     // V7-I03: synchronous lock check at the very top — before any
@@ -628,6 +631,16 @@ export function App(): JSX.Element {
       if (typeof trainOptions.val_every === "number") {
         trainOpts.val_every = trainOptions.val_every;
       }
+      // V7-Q06.1: B/S override threaded into stage_train opts. Backend
+      // reads opts.B / opts.S for the per-step synthetic/parquet batch.
+      if (typeof trainOptions.train_B === "number"
+          && trainOptions.train_B > 0) {
+        trainOpts.B = trainOptions.train_B;
+      }
+      if (typeof trainOptions.train_S === "number"
+          && trainOptions.train_S > 0) {
+        trainOpts.S = trainOptions.train_S;
+      }
       if (typeof trainOptions.grad_clip_max_norm === "number") {
         trainOpts.grad_clip_max_norm = trainOptions.grad_clip_max_norm;
       }
@@ -1241,6 +1254,33 @@ export function App(): JSX.Element {
                   onDropBrick={handleDropBrick}
                   onNodeClick={setSelectedBrickId}
                   isValidConnection={isValidConnection}
+                  onInsertAdapter={(kind, edge) => {
+                    const baseName = `${kind}_insert_${nodes.length}`;
+                    const src = nodes.find((n) => n.id === edge.source);
+                    const dst = nodes.find((n) => n.id === edge.target);
+                    const mid = src && dst
+                      ? { x: (src.position.x + dst.position.x) / 2,
+                          y: (src.position.y + dst.position.y) / 2 + 50 }
+                      : { x: 200, y: 280 };
+                    setNodes((prev) => [
+                      ...prev,
+                      { id: baseName,
+                        type: "brick",
+                        position: mid,
+                        data: { kind } as never },
+                    ]);
+                    setEdges((prev) => [
+                      ...prev.filter((e) =>
+                        !(e.source === edge.source
+                          && e.target === edge.target)),
+                      { id: `${edge.source}->${baseName}`,
+                        source: edge.source, target: baseName,
+                        data: { severity: "info" } },
+                      { id: `${baseName}->${edge.target}`,
+                        source: baseName, target: edge.target,
+                        data: { severity: "info" } },
+                    ]);
+                  }}
                 />
               </div>
               {selectedBrickId && (() => {
diff --git a/vbgui/src/components/TrainExtrasOverlay.tsx b/vbgui/src/components/TrainExtrasOverlay.tsx
@@ -52,6 +52,32 @@ export interface TrainExtras {
     brick_kinds?: string[] | string;
     num_brick_kinds?: number;
   };
+  // V7-Q06.3: styled panels for previously-untyped extras keys.
+  plasticity?: {
+    fire_fired_at_step?: number;
+    fire_keys_modified?: string[];
+    dash_last_keys_count?: number;
+    redo_last_recycled?: number | string | null;
+    [k: string]: unknown;
+  };
+  mtp?: {
+    k?: number;
+    beta?: number | number[];
+    head_losses?: number[];
+    [k: string]: unknown;
+  } | null;
+  ifim?: {
+    instr_loss?: number;
+    instr_token_count?: number;
+    lambda?: number;
+    [k: string]: unknown;
+  } | null;
+  mhc?: {
+    consistency_loss?: number;
+    heads_correlated?: number;
+    lambda?: number;
+    [k: string]: unknown;
+  } | null;
   [k: string]: unknown;
 }
 
@@ -216,10 +242,163 @@ export function TrainExtrasOverlay({
         || Array.isArray(extras.per_expert_load)) && (
         <MoEDashboard extras={extras} />
       )}
+
+      {/* V7-Q06.3: plasticity panel — FIRE/DASH/ReDo capture from train. */}
+      {extras.plasticity
+        && Object.keys(extras.plasticity).some(
+          (k) => extras.plasticity![k] !== undefined
+                 && extras.plasticity![k] !== null) && (
+        <PlasticityPanel data={extras.plasticity} />
+      )}
+
+      {/* V7-Q06.3: MTP-weighted head loss panel. */}
+      {extras.mtp && (
+        <MTPPanel data={extras.mtp} />
+      )}
+
+      {/* V7-Q06.3: IFIM instr-token loss panel. */}
+      {extras.ifim && (
+        <IFIMPanel data={extras.ifim} />
+      )}
+
+      {/* V7-Q06.3: MHC attention-bias consistency panel. */}
+      {extras.mhc && (
+        <MHCPanel data={extras.mhc} />
+      )}
     </div>
   );
 }
 
+// ---------------------------------------------------------------------------
+// V7-Q06.3 styled sub-panels for plasticity / mtp / ifim / mhc extras.
+// ---------------------------------------------------------------------------
+
+function _PanelShell({
+  testid, title, children,
+}: { testid: string; title: string; children: React.ReactNode }) {
+  return (
+    <div data-testid={testid}
+         style={{ marginTop: 4, padding: "4px 8px",
+                  border: "1px solid #e5e7eb", borderRadius: 4,
+                  background: "#fafaf9", fontSize: 11 }}>
+      <div style={{ fontWeight: 600, color: "#374151",
+                    marginBottom: 4 }}>{title}</div>
+      {children}
+    </div>
+  );
+}
+
+function _MetricRow({
+  testid, label, value,
+}: { testid: string; label: string; value: string }) {
+  return (
+    <div style={{ display: "flex", gap: 6 }}>
+      <span style={{ color: "#6b7280", minWidth: 110 }}>{label}:</span>
+      <span data-testid={testid} style={{ color: "#111827",
+                                           fontFamily: "monospace" }}>
+        {value}
+      </span>
+    </div>
+  );
+}
+
+function PlasticityPanel({ data }: { data: NonNullable<TrainExtras["plasticity"]> }) {
+  return (
+    <_PanelShell testid="extras-panel-plasticity"
+                 title="Plasticity (FIRE / DASH / ReDo)">
+      {typeof data.fire_fired_at_step === "number" && (
+        <_MetricRow testid="extras-plasticity-fire-step"
+                    label="fire_fired_at_step"
+                    value={String(data.fire_fired_at_step)} />
+      )}
+      {Array.isArray(data.fire_keys_modified) && (
+        <_MetricRow testid="extras-plasticity-fire-keys"
+                    label="fire_keys_modified"
+                    value={`${data.fire_keys_modified.length} keys`} />
+      )}
+      {typeof data.dash_last_keys_count === "number" && (
+        <_MetricRow testid="extras-plasticity-dash-keys-count"
+                    label="dash_last_keys_count"
+                    value={String(data.dash_last_keys_count)} />
+      )}
+      {data.redo_last_recycled !== undefined
+        && data.redo_last_recycled !== null && (
+        <_MetricRow testid="extras-plasticity-redo-recycled"
+                    label="redo_last_recycled"
+                    value={String(data.redo_last_recycled)} />
+      )}
+    </_PanelShell>
+  );
+}
+
+function MTPPanel({ data }: { data: NonNullable<TrainExtras["mtp"]> }) {
+  return (
+    <_PanelShell testid="extras-panel-mtp"
+                 title="MTP (multi-token-prediction)">
+      {typeof data.k === "number" && (
+        <_MetricRow testid="extras-mtp-k"
+                    label="k" value={String(data.k)} />
+      )}
+      {data.beta !== undefined && (
+        <_MetricRow testid="extras-mtp-beta"
+                    label="beta" value={JSON.stringify(data.beta)} />
+      )}
+      {Array.isArray(data.head_losses) && (
+        <_MetricRow testid="extras-mtp-head-losses"
+                    label="head_losses"
+                    value={data.head_losses.map(
+                      (l) => Number(l).toFixed(4)).join(" / ")} />
+      )}
+    </_PanelShell>
+  );
+}
+
+function IFIMPanel({ data }: { data: NonNullable<TrainExtras["ifim"]> }) {
+  return (
+    <_PanelShell testid="extras-panel-ifim"
+                 title="IFIM (instruction-aware FIM)">
+      {typeof data.instr_loss === "number" && (
+        <_MetricRow testid="extras-ifim-instr-loss"
+                    label="instr_loss"
+                    value={Number(data.instr_loss).toFixed(4)} />
+      )}
+      {typeof data.instr_token_count === "number" && (
+        <_MetricRow testid="extras-ifim-token-count"
+                    label="instr_token_count"
+                    value={String(data.instr_token_count)} />
+      )}
+      {typeof data.lambda === "number" && (
+        <_MetricRow testid="extras-ifim-lambda"
+                    label="lambda"
+                    value={Number(data.lambda).toFixed(4)} />
+      )}
+    </_PanelShell>
+  );
+}
+
+function MHCPanel({ data }: { data: NonNullable<TrainExtras["mhc"]> }) {
+  return (
+    <_PanelShell testid="extras-panel-mhc"
+                 title="MHC (multi-head consistency)">
+      {typeof data.consistency_loss === "number" && (
+        <_MetricRow testid="extras-mhc-consistency-loss"
+                    label="consistency_loss"
+                    value={Number(data.consistency_loss).toFixed(4)} />
+      )}
+      {typeof data.heads_correlated === "number" && (
+        <_MetricRow testid="extras-mhc-heads-correlated"
+                    label="heads_correlated"
+                    value={Number(data.heads_correlated).toFixed(4)} />
+      )}
+      {typeof data.lambda === "number" && (
+        <_MetricRow testid="extras-mhc-lambda"
+                    label="lambda"
+                    value={Number(data.lambda).toFixed(4)} />
+      )}
+    </_PanelShell>
+  );
+}
+
 function Badge({
   testid, label, value, tone = "neutral", help,
 }: { testid: string; label: string; value: string;
diff --git a/vbgui/src/components/TrainOptionsPanel.tsx b/vbgui/src/components/TrainOptionsPanel.tsx
@@ -19,6 +19,10 @@ export interface TrainOptions {
   fake_ranks?: number;
   // K8 — explicit abort token (defaults to run_id when empty).
   abort_token?: string;
+  // V7-Q06.1 — B/S overrides applied before pipeline.run dispatch.
+  // Default unset => spec.dim_env.B/S wins (preserves prior behaviour).
+  train_B?: number;
+  train_S?: number;
 }
 
 export interface TrainOptionsPanelProps {
@@ -125,6 +129,29 @@ export function TrainOptionsPanel({
                      e.target.value === "" ? undefined : e.target.value)}
                    style={{ ...INPUT, width: 200 }} />
           </Row>
+          {/* V7-Q06.1: B + S overrides for train dim_env. */}
+          <Row label="B (train batch size)" help="dim_env_B">
+            <input data-testid="train-opt-B"
+                   type="number" min={1} max={32}
+                   value={v.train_B ?? ""}
+                   placeholder="spec.dim_env.B"
+                   onChange={(e) => set(
+                     "train_B",
+                     e.target.value === ""
+                       ? undefined : parseInt(e.target.value, 10))}
+                   style={INPUT} />
+          </Row>
+          <Row label="S (train seq len)" help="dim_env_S">
+            <input data-testid="train-opt-S"
+                   type="number" min={1} max={8192}
+                   value={v.train_S ?? ""}
+                   placeholder="spec.dim_env.S"
+                   onChange={(e) => set(
+                     "train_S",
+                     e.target.value === ""
+                       ? undefined : parseInt(e.target.value, 10))}
+                   style={INPUT} />
+          </Row>
         </div>
       )}
     </div>
diff --git a/vbgui/tests/TrainExtrasOverlayPanels.test.tsx b/vbgui/tests/TrainExtrasOverlayPanels.test.tsx
diff --git a/vbgui/tests/TrainOptionsPanelBS.test.tsx b/vbgui/tests/TrainOptionsPanelBS.test.tsx
