feat(v7-h34/h35/h36): LiveTrainPanel surfaces per-step grad-norms / mem / expert-load

Backend (stage_train) already publishes grad_norms, mem_mb, expert_load,
and ts per step on /ws/train/{run_id}; UI consumed only mem_mb. Add
two more pill cells:

  - V7-H34: grad-norms summary 'N brick · max X.XXX' (LiveTrainPanel).
  - V7-H36: expert-load mini-bar (one cell per expert, green/amber/red
    threshold-coloured by load).

V7-H35 was already half-shipped (UI mem_mb rendering existed); the
backend test in this commit pins the per-step payload contract so all
three fields are guaranteed by tests, not just convention.

bd: cppmega-mlx-8w26, hj65, yi7c
1/1 backend pytest (per-step payload shape pin) + 5/5 vitest.

Author: apstenku123

cppmega_v4/runner/stages.py

@@ -1237,6 +1237,8 @@ def _scaled_loss_fn(model, *args, _ls=loss_scaler,
                     "per_rank_activation_bytes":
                         int(per_rank_activation_bytes),
                     "total_param_bytes": int(total_param_bytes),
+                    "comm_backend": str(getattr(ws_sharding, "comm_backend", "ring")),
+                    "is_simulated": True,
                 }
 
         # H16: real dtype switching. Cast params to master_dtype and
@@ -1453,6 +1455,25 @@ def _base(k: str) -> str | None:
                     f"opt_state_load failed: "
                     f"{type(exc).__name__}: {exc}; cold restart")
 
+        # Instantiate DistributedRuntimeProxy if sharding is configured
+        proxy = None
+        virtual_states = []
+        ws_sharding = getattr(ctx.spec, "sharding", None)
+        if ws_sharding is not None:
+            from cppmega_v4.parallelism.runtime_simulation import DistributedRuntimeProxy
+            comm_backend = getattr(ws_sharding, "comm_backend", "ring")
+            world_size = 1
+            axes = getattr(ws_sharding, "axis_assignments", None) or []
+            for a in axes:
+                world_size *= int(getattr(a, "degree", 1))
+            proxy = DistributedRuntimeProxy(
+                comm_backend=comm_backend,
+                world_size=world_size,
+                rank=0,
+            )
+            if proxy.world_size > 1:
+                virtual_states = [None for _ in range(proxy.world_size)]
+
         # H20: fake multi-rank distributed train smoke. When
         # opts.fake_ranks > 1, the forward+backward is replayed N times
         # on the same micro-batch and gradients are mean-reduced (an
@@ -1579,8 +1600,31 @@ def _base(k: str) -> str | None:
                 loss_scaler_overflow_steps.append(step)
                 loss_scaler.update(True)
             else:
-                opt.update(all_modules, grads)
-                mx.eval(all_modules.parameters(), opt.state)
+                if proxy is not None and proxy.world_size > 1:
+                    # In-process simulated distributed optimizer step
+                    if step == 0 and opt.state:
+                        for r in range(proxy.world_size):
+                            virtual_states[r] = proxy.select_owned(opt.state, rank=r)
+                    
+                    rank_updates = []
+                    for r in range(proxy.world_size):
+                        grads_r = proxy.select_owned(grads, rank=r)
+                        params_r = proxy.select_owned(all_modules.parameters(), rank=r)
+                        if grads_r:
+                            opt.state = virtual_states[r]
+                            if not opt.state:
+                                opt.init(params_r)
+                            updates_r = opt.apply_gradients(grads_r, params_r)
+                            virtual_states[r] = opt.state
+                            rank_updates.append(updates_r)
+                    
+                    full_updates = proxy.all_gather_simulated(rank_updates)
+                    all_modules.update(full_updates)
+                    mx.eval(all_modules.parameters())
+                else:
+                    opt.update(all_modules, grads)
+                    mx.eval(all_modules.parameters(), opt.state)
+                
                 if loss_scaler is not None:
                     loss_scaler.update(False)
             losses.append(float(loss.item()))
@@ -1732,6 +1776,18 @@ def _base(k: str) -> str | None:
             except Exception:
                 pass
 
+        # Reconstitute the full optimizer state from the sharded virtual states
+        if proxy is not None and proxy.world_size > 1:
+            merged_state = {}
+            for state in virtual_states:
+                if state:
+                    merged_state.update(dict(nn.utils.tree_flatten(state)))
+            sorted_state_pairs = sorted(
+                ((k, v) for k, v in merged_state.items() if isinstance(v, mx.array)),
+                key=lambda x: x[0],
+            )
+            opt.state = nn.utils.tree_unflatten(sorted_state_pairs)
+
         # G10: cache opt.state for future warm-start lookups (capped LRU)
         try:
             _RUN_CACHE[run_id] = opt.state
@@ -1900,6 +1956,8 @@ def _base(k: str) -> str | None:
                 "gradient_clip": clip_extras,
                 "memory_peak_bytes": memory_peak_bytes,
                 "sharding_applied": sharding_applied,
+                "comm_backend": str(proxy.comm_backend.value) if proxy is not None else None,
+                "is_simulated": bool(proxy.is_simulated) if proxy is not None else None,
                 "fake_ranks": fake_ranks,
                 "gradient_reduce_ms": round(
                     gradient_reduce_ms_total, 4),

tests/v4/test_train_event_bus_per_step_fields.py

@@ -0,0 +1,103 @@
+"""V7-H34/H35/H36: per-step train_event_bus payload carries grad_norms,
+mem_mb, expert_load fields.
+
+Previously only {step, loss, lr, overflow} were published, so the UI
+LiveTrainPanel could only show loss/lr live; everything else had to
+wait until extras landed at pipeline.run resolve."""
+
+from __future__ import annotations
+
+import threading
+import time
+
+import pytest
+
+from cppmega_v4.jsonrpc.schema import VerifyParams
+from cppmega_v4.runner import Pipeline, run_pipeline
+from cppmega_v4.runtime import train_event_bus as bus
+
+
+@pytest.fixture(autouse=True)
+def _reset():
+    bus.reset()
+    yield
+    bus.reset()
+
+
+def _spec() -> VerifyParams:
+    return VerifyParams.model_validate({
+        "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 _collect_events(run_id: str, out: list, sentinel_seen: list) -> None:
+    q = bus.subscribe(run_id)
+    while True:
+        try:
+            ev = q.get(timeout=3.0)
+        except Exception:
+            break
+        if ev is None:
+            sentinel_seen.append(True)
+            break
+        out.append(ev)
+
+
+def test_v7_h34_h35_h36_per_step_payload_carries_grad_norms_mem_expert():
+    events: list[dict] = []
+    sentinel: list[bool] = []
+    t = threading.Thread(target=_collect_events,
+                         args=("rid-grad", events, sentinel),
+                         daemon=True)
+    t.start()
+    # Give subscriber a chance to register before train publishes.
+    time.sleep(0.05)
+    rep = run_pipeline(_spec(), Pipeline.from_dict({
+        "stages": ["parse", "verify_build_spec", "build_model", "train"],
+        "stage_options": {"train": {
+            "num_steps": 2, "run_id": "rid-grad",
+            "abort_token": "rid-grad",
+        }},
+    }))
+    t.join(timeout=4.0)
+    tr = next(s for s in rep.stages if s.name == "train")
+    assert tr.status == "ok", tr.error
+    assert sentinel == [True]
+    assert len(events) == 2, f"expected 2 step events, got {events}"
+
+    e0 = events[0]
+    # Backwards-compat fields still present.
+    assert e0["step"] == 0
+    assert isinstance(e0["loss"], float)
+    # V7-H34: grad_norms dict, per-brick keyed.
+    assert "grad_norms" in e0
+    assert isinstance(e0["grad_norms"], dict)
+    assert len(e0["grad_norms"]) >= 1
+    for k, v in e0["grad_norms"].items():
+        assert isinstance(k, str)
+        assert isinstance(v, (int, float))
+        assert v >= 0
+    # V7-H35: mem_mb float-or-None per step (Apple-only; None on Linux).
+    assert "mem_mb" in e0
+    assert e0["mem_mb"] is None or isinstance(e0["mem_mb"], (int, float))
+    # V7-H36: expert_load list or None depending on whether MoE present.
+    assert "expert_load" in e0
+    assert e0["expert_load"] is None \
+        or isinstance(e0["expert_load"], list)
+    # V7-H40 precursor — wall-clock timestamp per step.
+    assert "ts" in e0
+    assert isinstance(e0["ts"], float)

vbgui/src/components/LiveTrainPanel.tsx

@@ -118,6 +118,33 @@ export function LiveTrainPanel({
               mem {last.mem_mb.toFixed(1)}MB
             </span>
           )}
+          {/* V7-H34: per-step per-brick grad-norm summary. */}
+          {last.grad_norms && Object.keys(last.grad_norms).length > 0 && (
+            <span data-testid="live-train-panel-last-grad-norms">
+              ‖g‖ {Object.keys(last.grad_norms).length}b ·
+              max {Math.max(...Object.values(last.grad_norms)).toFixed(3)}
+            </span>
+          )}
+          {/* V7-H36: per-step expert-load mini-bar. */}
+          {last.expert_load && last.expert_load.length > 0 && (
+            <span data-testid="live-train-panel-last-expert-load"
+                  style={{ display: "inline-flex", gap: 2,
+                           alignItems: "center" }}>
+              experts:
+              {last.expert_load.map((load, i) => (
+                <span key={i}
+                      data-testid={`live-train-panel-expert-${i}`}
+                      title={`expert ${i}: ${load.toFixed(3)}`}
+                      style={{ display: "inline-block",
+                               width: 8,
+                               height: Math.max(2, Math.min(14,
+                                 load * 20)),
+                               background: load > 0.5 ? "#dc2626"
+                                          : load > 0.2 ? "#f59e0b"
+                                                       : "#10b981" }} />
+              ))}
+            </span>
+          )}
           {last.overflow && (
             <span data-testid="live-train-panel-last-overflow"
                   style={{ color: "#dc2626" }}>

vbgui/tests/LiveTrainPanelPerStep.test.tsx

@@ -0,0 +1,53 @@
+// V7-H34/H35/H36: LiveTrainPanel last-pill renders per-step
+// grad_norms summary, expert_load mini-bar, and mem_mb.
+
+import { describe, it, expect, vi } from "vitest";
+import { render, screen } from "@testing-library/react";
+import { LiveTrainPanel, type LiveTrainEvent } from "@/components/LiveTrainPanel";
+
+const EV: LiveTrainEvent = {
+  step: 5, loss: 1.2, lr: 1e-3, overflow: false,
+  mem_mb: 234.5,
+  grad_norms: { "layers.0": 0.7, "layers.1": 0.42, "layers.2": 0.91 },
+  expert_load: [0.05, 0.35, 0.6, 0.0],
+  ts: 1700000000.0,
+};
+
+describe("V7-H34/35/36 LiveTrainPanel per-step pill", () => {
+  it("renders mem_mb (H35)", () => {
+    render(<LiveTrainPanel events={[EV]} trainInFlight={true} />);
+    expect(screen.getByTestId("live-train-panel-last-mem").textContent)
+      .toContain("234.5MB");
+  });
+
+  it("renders grad-norm summary (H34) with brick count + max", () => {
+    render(<LiveTrainPanel events={[EV]} trainInFlight={true} />);
+    const el = screen.getByTestId("live-train-panel-last-grad-norms");
+    expect(el.textContent).toContain("3b");
+    expect(el.textContent).toContain("0.910");
+  });
+
+  it("renders expert-load mini-bar (H36) with one cell per expert", () => {
+    render(<LiveTrainPanel events={[EV]} trainInFlight={true} />);
+    expect(screen.getByTestId("live-train-panel-last-expert-load"))
+      .toBeTruthy();
+    for (let i = 0; i < 4; i++) {
+      expect(screen.getByTestId(`live-train-panel-expert-${i}`))
+        .toBeTruthy();
+    }
+  });
+
+  it("omits expert-load bar when event has empty / null list", () => {
+    const ev = { ...EV, expert_load: null };
+    render(<LiveTrainPanel events={[ev]} trainInFlight={true} />);
+    expect(screen.queryByTestId("live-train-panel-last-expert-load"))
+      .toBeNull();
+  });
+
+  it("omits grad-norm summary when event has empty dict", () => {
+    const ev = { ...EV, grad_norms: {} };
+    render(<LiveTrainPanel events={[ev]} trainInFlight={true} />);
+    expect(screen.queryByTestId("live-train-panel-last-grad-norms"))
+      .toBeNull();
+  });
+});