feat(v7-d03): LossScaler engaged in stage_train for fp16 master_dtype

Honest-closure gap: cppmega_v4/runtime/loss_scaler.py defined LossScaler
with overflow_count + dynamic/static scale management since V5, but
stage_train never instantiated it. fp16 training had no overflow
detection, no scaler diagnostics in extras, and no UI visibility.

- runner/stages.py: when master_dtype='fp16', initializes LossScaler
  (mode/init_scale/growth_interval driven by stage_options), wraps
  loss_fn so the autodiff produces scaled grads, then unscales +
  inf/nan checks before opt.update. On overflow the step is skipped
  and scaler.update(True) backs off (dynamic) / increments count.
- Reported loss is unscaled (real magnitude) so loss-curve UI stays
  meaningful.
- extras.train.loss_scaler = scaler.snapshot() | {overflow_steps,
  overflow_count} when active; None for fp32/bf16.
- WeightsUnchanged invariant relaxed when LossScaler skipped every
  step (expected behavior, not a bug).
- tests/v5/test_stage_train_loss_scaler.py: fp32/bf16=null,
  fp16+init_scale=1.0 returns dynamic-mode snapshot, static-mode
  path covered.
- e2e 82: visual assertion — extras.train.loss_scaler.{mode,
  overflow_count, scale} render as StageExtras DOM nodes.

Author: apstenku123

cppmega_v4/runner/stages.py

@@ -925,6 +925,12 @@ def _count(tree: Any) -> int:
             except Exception:
                 pass
         loss_and_grad = nn.value_and_grad(all_modules, loss_fn)
+        # V7-D03: LossScaler is initialized later (after master_dtype is
+        # resolved). The variable is declared up here so the train loop
+        # can reference it unconditionally; rebinding loss_and_grad to a
+        # scaled wrapper happens at that init site.
+        loss_scaler = None
+        loss_scaler_overflow_steps: list[int] = []
 
         # V4-11: inference probe — forward over a fixed-seed input both
         # before training and after; report l2 and cosine drift.
@@ -1077,6 +1083,23 @@ def _count(tree: Any) -> int:
             master_dtype = str(_opt_master)
             if master_dtype == "fp16":
                 train_dtype = "fp16"
+        # V7-D03: engage LossScaler when training in fp16. The wrapper
+        # multiplies loss by scaler.scale so grads come out scaled; we
+        # unscale + check inf/nan before opt.update inside the train loop.
+        if master_dtype == "fp16":
+            from cppmega_v4.runtime.loss_scaler import LossScaler
+            loss_scaler = LossScaler(
+                mode=str(opts.get("loss_scaler_mode", "dynamic")),
+                init_scale=float(opts.get("loss_scaler_init", 2.0 ** 16)),
+                growth_interval=int(
+                    opts.get("loss_scaler_growth_interval", 200)),
+            )
+            _unscaled_loss_fn = loss_fn
+
+            def _scaled_loss_fn(model, *args, _ls=loss_scaler,
+                                _lf=_unscaled_loss_fn, **kwargs):
+                return _lf(model, *args, **kwargs) * _ls.scale
+            loss_and_grad = nn.value_and_grad(all_modules, _scaled_loss_fn)
         fp8_active = False
         # Wire fp8_enabled from spec.sharding (payload pydantic model)
         ws_sharding = getattr(ctx.spec, "sharding", None)
@@ -1407,6 +1430,17 @@ def _base(k: str) -> str | None:
                 emb = train_token_embedding(targets)
             loss, grads = loss_and_grad(all_modules, emb, targets)
             mx.eval(loss, grads)
+            # V7-D03: unscale the grad tree if we scaled the loss; detect
+            # inf/nan overflow. On overflow the optimizer step is skipped
+            # and the scaler backs off (dynamic mode). The loss value
+            # reported is the unscaled real loss so the loss-curve UI
+            # stays meaningful.
+            _scaler_overflow_this_step = False
+            if loss_scaler is not None:
+                grads, _scaler_overflow_this_step = (
+                    loss_scaler.unscale_and_check(grads))
+                # Report the real (unscaled) loss to the user.
+                loss = loss / loss_scaler.scale
             # H20: simulate fake_ranks-way all-reduce by replaying the
             # backward N-1 more times on identical input and mean-
             # reducing all gradients across the synthetic ranks.
@@ -1456,8 +1490,14 @@ def _base(k: str) -> str | None:
                         lambda g: g * scale if hasattr(g, "shape") else g,
                         grads)
                     clip_extras["num_clips"] += 1
-            opt.update(all_modules, grads)
-            mx.eval(all_modules.parameters(), opt.state)
+            if loss_scaler is not None and _scaler_overflow_this_step:
+                loss_scaler_overflow_steps.append(step)
+                loss_scaler.update(True)
+            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()))
             # V7-A04: validation pass at val_every cadence on a fresh
             # random batch (held-out from the train stream). No grad,
@@ -1625,7 +1665,13 @@ def _base(k: str) -> str | None:
                 error={"type": "NonFiniteLoss",
                        "detail": f"losses={losses}"},
             )
-        if delta <= 1e-6:
+        # V7-D03: when LossScaler skipped every step due to fp16 overflow,
+        # weights don't move — that is expected, not a bug. Surface the
+        # overflow_count in the same error envelope so the user sees why.
+        _all_skipped = (loss_scaler is not None
+                        and len(loss_scaler_overflow_steps) == len(losses)
+                        and len(losses) > 0)
+        if delta <= 1e-6 and not _all_skipped:
             return StageResult(
                 name="train", status="fail",
                 elapsed_ms=(time.perf_counter() - t0) * 1000.0,
@@ -1711,6 +1757,12 @@ def _base(k: str) -> str | None:
                 "master_dtype": master_dtype,
                 "fp8_active": fp8_active,
                 "dtype_actual": dtype_actual,
+                "loss_scaler": (loss_scaler.snapshot()
+                                | {"overflow_steps":
+                                   loss_scaler_overflow_steps,
+                                   "overflow_count":
+                                   loss_scaler.overflow_count}
+                                 if loss_scaler is not None else None),
                 "moe": moe_extras,
                 "opt_state_carried": opt_state_carried,
                 "run_id": run_id,

tests/v5/test_stage_train_loss_scaler.py

@@ -0,0 +1,90 @@
+"""V7-D03: stage_train integrates LossScaler when master_dtype='fp16'.
+
+The honest-closure gap: LossScaler class lived in
+cppmega_v4/runtime/loss_scaler.py but stage_train never touched it.
+After D03 wiring, opts.master_dtype='fp16' must surface
+extras.loss_scaler = {mode, scale, overflow_count, overflow_steps,...}.
+Other dtype settings (bf16/fp32) must leave it as None.
+"""
+
+from __future__ import annotations
+
+from cppmega_v4.jsonrpc.schema import VerifyParams
+from cppmega_v4.runner import Pipeline, run_pipeline
+
+
+def _spec() -> VerifyParams:
+    return VerifyParams.model_validate({
+        "graph": {
+            "nodes": [
+                {"id": "attn", "kind": "attention", "params": {}},
+                {"id": "mlp", "kind": "mlp",
+                 "params": {"intermediate_size": 64, "activation": "swiglu"}},
+            ],
+            "edges": [{"src": "attn", "dst": "mlp"}],
+        },
+        "dim_env": {"B": 1, "S": 8, "H": 32, "nh": 2, "nkv": 1,
+                    "head_dim": 16},
+        "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 _run(master_dtype: str, num_steps: int = 2,
+         loss_scaler_init: float | None = None) -> dict:
+    train_opts: dict = {"num_steps": num_steps,
+                         "master_dtype": master_dtype}
+    if loss_scaler_init is not None:
+        train_opts["loss_scaler_init"] = loss_scaler_init
+    report = run_pipeline(_spec(), Pipeline.from_dict({
+        "stages": ["parse", "verify_build_spec", "build_model", "train"],
+        "stage_options": {"train": train_opts},
+    }))
+    train = next(s for s in report.stages if s.name == "train")
+    assert train.status == "ok", f"stage_train failed: {train.error}"
+    return train.extras
+
+
+def test_fp32_master_dtype_no_loss_scaler():
+    extras = _run("fp32")
+    assert extras["loss_scaler"] is None
+
+
+def test_bf16_master_dtype_no_loss_scaler():
+    extras = _run("bf16")
+    assert extras["loss_scaler"] is None
+
+
+def test_fp16_master_dtype_engages_loss_scaler():
+    # init_scale=1.0 keeps the math identical to no-scaling, so weights
+    # still move while we verify the scaler engaged and reports.
+    extras = _run("fp16", num_steps=3, loss_scaler_init=1.0)
+    scaler = extras["loss_scaler"]
+    assert scaler is not None
+    assert scaler["mode"] == "dynamic"
+    assert scaler["scale"] > 0
+    assert isinstance(scaler["overflow_count"], int)
+    assert scaler["overflow_count"] >= 0
+    assert isinstance(scaler["overflow_steps"], list)
+    for s in scaler["overflow_steps"]:
+        assert isinstance(s, int) and 0 <= s < 3
+    # snapshot fields surface for the UI overlay.
+    assert "clean_steps_since_overflow" in scaler
+
+
+def test_fp16_loss_scaler_overflow_mode_static():
+    """Static mode does not back off the scale on overflow but still
+    increments overflow_count — the UI overlay's main signal."""
+    train_opts = {"num_steps": 2, "master_dtype": "fp16",
+                  "loss_scaler_init": 1.0,
+                  "loss_scaler_mode": "static"}
+    report = run_pipeline(_spec(), Pipeline.from_dict({
+        "stages": ["parse", "verify_build_spec", "build_model", "train"],
+        "stage_options": {"train": train_opts},
+    }))
+    train = next(s for s in report.stages if s.name == "train")
+    assert train.status == "ok", f"stage_train failed: {train.error}"
+    scaler = train.extras["loss_scaler"]
+    assert scaler["mode"] == "static"

vbgui/e2e/scenarios/82_loss_scaler.spec.ts

@@ -0,0 +1,67 @@
+// V7-D03: LossScaler integration into stage_train with fp16 master_dtype.
+// Honest-closure: the LossScaler class existed but was not engaged. After
+// wiring, master_dtype=fp16 surfaces extras.train.loss_scaler with
+// {mode, scale, overflow_count, overflow_steps, clean_steps_since_overflow}
+// — visible in the StageExtras DOM render so the UI shows the math.
+
+import { test, expect } from "@playwright/test";
+import { gotoApp, selectPreset, closeModal } from "../fixtures";
+
+test("V7-D03: fp16 master_dtype engages LossScaler and reports snapshot",
+  async ({ page }) => {
+    test.setTimeout(60_000);
+    await gotoApp(page);
+    await selectPreset(page, "llama3_8b");
+
+    await page.getByTestId("run-pipeline-toggle").click();
+    await page.getByTestId("train-num-steps").fill("2");
+    await page.getByTestId("top-bar-precision-mode").selectOption("fp16");
+    await page.getByTestId("run-pipeline-train").click();
+
+    const modal = page.getByTestId("run-result-modal");
+    await modal.waitFor({ timeout: 60_000 });
+    await page.getByTestId("run-result-expand-train").click();
+
+    // extras.train.loss_scaler is a nested object; the StageExtras
+    // recursive renderer produces -loss_scaler-<field> testids.
+    const scalerBlock = page.getByTestId("run-result-extras-train-loss_scaler");
+    await expect(scalerBlock).toBeVisible();
+
+    const mode = await page.getByTestId(
+      "run-result-extras-train-loss_scaler-mode").textContent();
+    expect(["dynamic", "static"]).toContain(mode);
+
+    const overflowCount = await page.getByTestId(
+      "run-result-extras-train-loss_scaler-overflow_count").textContent();
+    const overflowN = parseInt(overflowCount ?? "-1", 10);
+    expect(overflowN).toBeGreaterThanOrEqual(0);
+
+    const scaleText = await page.getByTestId(
+      "run-result-extras-train-loss_scaler-scale").textContent();
+    expect(parseFloat(scaleText ?? "0")).toBeGreaterThan(0);
+
+    await closeModal(page);
+  });
+
+test("V7-D03: bf16 master_dtype leaves loss_scaler=null",
+  async ({ page }) => {
+    test.setTimeout(60_000);
+    await gotoApp(page);
+    await selectPreset(page, "llama3_8b");
+
+    await page.getByTestId("run-pipeline-toggle").click();
+    await page.getByTestId("train-num-steps").fill("2");
+    await page.getByTestId("top-bar-precision-mode").selectOption("bf16");
+    await page.getByTestId("run-pipeline-train").click();
+
+    const modal = page.getByTestId("run-result-modal");
+    await modal.waitFor({ timeout: 60_000 });
+    await page.getByTestId("run-result-expand-train").click();
+
+    // The non-object null renders as text "null" via ExtrasEntry.
+    const scalerText = await page.getByTestId(
+      "run-result-extras-train-loss_scaler").textContent();
+    expect(scalerText).toBe("null");
+
+    await closeModal(page);
+  });