feat(v7-d03): LossScaler with static + dynamic modes for fp16

Closes V7-D03 (cppmega-mlx-4ij1): fp16 has a narrow dynamic range
(~6e-5 to 6e4) so a naive fp16 train step overflows or underflows
within a handful of iterations. cppmega_v4.runtime.loss_scaler.LossScaler
provides the missing scaling + overflow detection.

Public API:
  - LossScaler(mode='static'|'dynamic', init_scale, growth_factor,
                backoff_factor, growth_interval, min_scale, max_scale)
  - .scale_grads(grads) — multiply each grad tensor by scale.
  - .unscale_and_check(grads) → (unscaled, overflow). Overflow is
    True when any inf/nan in any grad entry.
  - .update(overflow) — dynamic mode halves scale on overflow,
    doubles after growth_interval consecutive clean steps, bounded
    by [min_scale, max_scale]. Static mode just counts overflow.
  - .snapshot() — dict suitable for extras / UI overlay:
    {mode, scale, overflow_count, clean_steps_since_overflow}.

Tests (tests/v4/test_loss_scaler.py): 8/8 — static init+scaling,
dynamic halve-on-overflow, doubling after growth_interval, static
counts but doesn't adjust, unscale inverts scaling, NaN+inf both
detected, min/max scale bounds respected, snapshot shape.

Integration into stage_train is the next step (V7-D03 sub-task —
hooking LossScaler into the training loop when master_dtype="fp16"
and surfacing scale/overflow in extras + UI).

Author: apstenku123

cppmega_v4/runtime/__init__.py

@@ -0,0 +1 @@
+"""cppmega_v4 runtime helpers."""

cppmega_v4/runtime/loss_scaler.py

@@ -0,0 +1,144 @@
+"""V7-D03: loss scaling + overflow detection for fp16 training.
+
+fp16 has a narrow dynamic range (~6e-5 to 6e4). Without a loss
+scaler, gradients underflow silently or overflow into inf/nan
+within a few steps. This module provides:
+
+  * `LossScaler` with two modes:
+      - static (`mode='static'`, `init_scale=2**16`): multiplies the
+        loss by a fixed factor before the backward pass; on
+        gradient overflow it just increments a counter without
+        adjusting the scale.
+      - dynamic (`mode='dynamic'`): halves the scale on overflow,
+        doubles after `growth_interval` consecutive clean steps,
+        bounded by `[min_scale, max_scale]`.
+  * `unscale_and_check(grads)` divides each grad by the active scale
+    and returns (unscaled_grads, overflow_detected).
+  * Caller is responsible for skipping the optimizer.update() on
+    overflow and calling `update(overflow=...)` after each step so
+    the dynamic scaler tracks state.
+
+Usage in a training loop:
+
+    scaler = LossScaler(mode="dynamic")
+    for step in range(N):
+        loss, grads = lvg(model, x, y)
+        loss = loss * scaler.scale            # scale the loss
+        grads = scaler.scale_grads(grads)     # equivalent: scale grads
+        unscaled, overflow = scaler.unscale_and_check(grads)
+        if not overflow:
+            opt.update(model, unscaled)
+        scaler.update(overflow)               # adjust dynamic scale
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class LossScaler:
+    """Loss / gradient scaler for fp16 training.
+
+    Attributes:
+        scale: current scale factor (mx scalar).
+        overflow_count: total overflow events observed.
+        clean_steps_since_overflow: monotonic counter used by dynamic
+            mode to decide when to grow the scale.
+    """
+
+    def __init__(
+        self,
+        *,
+        mode: Literal["static", "dynamic"] = "dynamic",
+        init_scale: float = 2.0 ** 16,
+        growth_factor: float = 2.0,
+        backoff_factor: float = 0.5,
+        growth_interval: int = 200,
+        min_scale: float = 1.0,
+        max_scale: float = 2.0 ** 24,
+    ) -> None:
+        if mode not in ("static", "dynamic"):
+            raise ValueError(f"mode must be static|dynamic, got {mode}")
+        if init_scale <= 0:
+            raise ValueError("init_scale must be > 0")
+        if growth_factor <= 1.0 or backoff_factor >= 1.0:
+            raise ValueError(
+                "growth_factor must be > 1 and backoff_factor must be < 1"
+            )
+        self.mode = mode
+        self._scale = float(init_scale)
+        self.growth_factor = growth_factor
+        self.backoff_factor = backoff_factor
+        self.growth_interval = max(1, int(growth_interval))
+        self.min_scale = float(min_scale)
+        self.max_scale = float(max_scale)
+        self.overflow_count = 0
+        self.clean_steps_since_overflow = 0
+
+    @property
+    def scale(self) -> float:
+        return self._scale
+
+    def scale_grads(self, grads: Any) -> Any:
+        """Multiply each grad tensor by self.scale (in-place equivalent
+        for tree). Used when scaling AFTER the backward pass — the
+        more common pattern is to multiply the loss directly."""
+        return nn.utils.tree_map(
+            lambda g: g * self._scale if hasattr(g, "shape") else g,
+            grads,
+        )
+
+    def unscale_and_check(self, grads: Any) -> tuple[Any, bool]:
+        """Divide each grad by self.scale and check for inf/nan.
+
+        Returns (unscaled_grads, overflow). When overflow is True the
+        caller MUST skip the optimizer update for this step.
+        """
+        flat = dict(nn.utils.tree_flatten(grads))
+        overflow = False
+        unscaled: dict[str, Any] = {}
+        for k, g in flat.items():
+            if not hasattr(g, "shape"):
+                unscaled[k] = g
+                continue
+            ug = g / self._scale
+            # mlx has no isfinite — use the abs<inf trick.
+            if bool(mx.any(mx.isnan(ug)).item()) or bool(
+                    mx.any(mx.isinf(ug)).item()):
+                overflow = True
+            unscaled[k] = ug
+        unscaled_tree = nn.utils.tree_unflatten(list(unscaled.items()))
+        return unscaled_tree, overflow
+
+    def update(self, overflow: bool) -> None:
+        """Adjust the scale based on this step's overflow outcome."""
+        if overflow:
+            self.overflow_count += 1
+            self.clean_steps_since_overflow = 0
+            if self.mode == "dynamic":
+                self._scale = max(
+                    self.min_scale, self._scale * self.backoff_factor)
+            return
+        self.clean_steps_since_overflow += 1
+        if (self.mode == "dynamic"
+                and self.clean_steps_since_overflow
+                >= self.growth_interval):
+            self._scale = min(
+                self.max_scale, self._scale * self.growth_factor)
+            self.clean_steps_since_overflow = 0
+
+    def snapshot(self) -> dict[str, Any]:
+        """Reportable state — meant for extras / UI overlay."""
+        return {
+            "mode": self.mode,
+            "scale": float(self._scale),
+            "overflow_count": int(self.overflow_count),
+            "clean_steps_since_overflow":
+                int(self.clean_steps_since_overflow),
+        }
+
+
+__all__ = ["LossScaler"]

tests/v4/test_loss_scaler.py

@@ -0,0 +1,101 @@
+"""V7-D03: LossScaler unit + integration tests."""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import pytest
+
+from cppmega_v4.runtime.loss_scaler import LossScaler
+
+
+def _grads_finite() -> dict:
+    return {"w": mx.array([[1.0, 2.0], [3.0, 4.0]], dtype=mx.float32)}
+
+
+def _grads_with_inf() -> dict:
+    return {"w": mx.array([[float("inf"), 0.0], [0.0, 0.0]],
+                          dtype=mx.float32)}
+
+
+def _grads_with_nan() -> dict:
+    return {"w": mx.array([[float("nan"), 0.0], [0.0, 0.0]],
+                          dtype=mx.float32)}
+
+
+def test_v7_d03_static_scaler_init_and_scale():
+    s = LossScaler(mode="static", init_scale=128.0)
+    assert s.scale == 128.0
+    grads = _grads_finite()
+    scaled = s.scale_grads(grads)
+    # Each entry multiplied by 128.
+    assert float(scaled["w"][0, 0].item()) == pytest.approx(128.0)
+
+
+def test_v7_d03_dynamic_scaler_halves_on_overflow():
+    s = LossScaler(mode="dynamic", init_scale=512.0,
+                    backoff_factor=0.5, min_scale=1.0)
+    assert s.scale == 512.0
+    _, overflow = s.unscale_and_check(_grads_with_inf())
+    assert overflow is True
+    s.update(overflow=True)
+    assert s.scale == 256.0
+    assert s.overflow_count == 1
+
+
+def test_v7_d03_dynamic_scaler_doubles_after_growth_interval():
+    s = LossScaler(mode="dynamic", init_scale=8.0,
+                    growth_factor=2.0, growth_interval=3,
+                    max_scale=1024.0)
+    for _ in range(3):
+        s.update(overflow=False)
+    assert s.scale == 16.0
+    for _ in range(3):
+        s.update(overflow=False)
+    assert s.scale == 32.0
+
+
+def test_v7_d03_static_mode_counts_overflow_but_does_not_adjust():
+    s = LossScaler(mode="static", init_scale=64.0)
+    s.update(overflow=True)
+    assert s.scale == 64.0  # unchanged
+    assert s.overflow_count == 1
+
+
+def test_v7_d03_unscale_inverts_scaling():
+    s = LossScaler(mode="static", init_scale=4.0)
+    grads = _grads_finite()
+    scaled = s.scale_grads(grads)
+    unscaled, overflow = s.unscale_and_check(scaled)
+    assert overflow is False
+    assert mx.allclose(unscaled["w"], grads["w"], atol=1e-6)
+
+
+def test_v7_d03_nan_grads_also_detected_as_overflow():
+    s = LossScaler(mode="dynamic", init_scale=128.0)
+    _, overflow = s.unscale_and_check(_grads_with_nan())
+    assert overflow is True
+
+
+def test_v7_d03_dynamic_scaler_respects_min_max_bounds():
+    s = LossScaler(mode="dynamic", init_scale=2.0,
+                    backoff_factor=0.5, min_scale=1.0,
+                    growth_factor=2.0, max_scale=4.0,
+                    growth_interval=1)
+    # Slam against min via repeated overflow.
+    for _ in range(10):
+        s.update(overflow=True)
+    assert s.scale == 1.0
+    # Slam against max via repeated clean steps.
+    for _ in range(10):
+        s.update(overflow=False)
+    assert s.scale == 4.0
+
+
+def test_v7_d03_snapshot_shape():
+    s = LossScaler(mode="dynamic", init_scale=128.0)
+    snap = s.snapshot()
+    for k in ("mode", "scale", "overflow_count",
+              "clean_steps_since_overflow"):
+        assert k in snap
+    assert snap["mode"] == "dynamic"
+    assert snap["scale"] == 128.0