feat(v7-d04): hybrid master-fp32 + bf16-fwd helpers

Closes V7-D04 (cppmega-mlx-m9i4): pure helpers for the
best-practice mixed-precision pattern:

  - cast_for_forward(param, fwd_dtype): bf16 view; master untouched.
  - cast_grads_to_master(grads, master_dtype=fp32): tree-cast.
  - hybrid_step(master_params, fwd_dtype, loss_and_grad,
                 apply_gradients) → updated master_params:
      1. cast master→bf16 for forward
      2. compute (loss, grads_bf16)
      3. cast grads back to fp32
      4. apply_gradients(master, fp32 grads) → fp32 master

Tests (tests/v4/test_hybrid_precision.py): 4/4 — fwd cast leaves
master fp32, grads cast preserves fp32, end-to-end step keeps
master fp32 and applies fp32 grads.

hybrid_lm.set_dtype('hybrid') integration is V7-D04 follow-up.

Author: apstenku123

cppmega_v4/runtime/hybrid_precision.py

@@ -0,0 +1,55 @@
+"""V7-D04: hybrid master-fp32-grad + bf16-forward helpers.
+
+Best-practice mixed precision: params live in fp32 (master copy);
+forward materialises a bf16 view (cast on the fly); gradients are
+cast back to fp32 for optimiser apply.
+
+  cast_for_forward(param, fwd_dtype) → fwd_dtype view
+  cast_grads_to_master(grads, master_dtype=fp32) → fp32 grads
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+def cast_for_forward(param: mx.array, fwd_dtype: mx.Dtype) -> mx.array:
+    """Return a non-master view of `param` in fwd_dtype. The original
+    fp32 master is preserved by the caller."""
+    return param.astype(fwd_dtype)
+
+
+def cast_grads_to_master(grads, master_dtype: mx.Dtype = mx.float32):
+    """Walk a grad tree and cast every leaf to master_dtype."""
+    return nn.utils.tree_map(
+        lambda g: g.astype(master_dtype) if hasattr(g, "shape") else g,
+        grads,
+    )
+
+
+def hybrid_step(*, master_params: dict, fwd_dtype: mx.Dtype,
+                 loss_and_grad,
+                 apply_gradients) -> dict:
+    """One end-to-end hybrid mixed-precision step.
+
+    Args:
+        master_params: dict of name → fp32 mx.array master params
+            (mutated in place by apply_gradients).
+        fwd_dtype: e.g. mx.bfloat16.
+        loss_and_grad: callable (fwd_params_dict) → (loss, grads).
+        apply_gradients: callable (master_params, fp32_grads) → updated
+            master_params dict.
+
+    Returns the updated master params dict.
+    """
+    fwd_params = {k: cast_for_forward(v, fwd_dtype)
+                  for k, v in master_params.items()}
+    _, grads = loss_and_grad(fwd_params)
+    grads_fp32 = cast_grads_to_master(grads, master_dtype=mx.float32)
+    return apply_gradients(master_params, grads_fp32)
+
+
+__all__ = [
+    "cast_for_forward", "cast_grads_to_master", "hybrid_step",
+]

tests/v4/test_hybrid_precision.py

@@ -0,0 +1,70 @@
+"""V7-D04: hybrid master-fp32 + bf16-fwd helpers."""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import pytest
+
+from cppmega_v4.runtime.hybrid_precision import (
+    cast_for_forward, cast_grads_to_master, hybrid_step,
+)
+
+
+def test_v7_d04_cast_for_forward_returns_bf16_view():
+    p = mx.random.normal(shape=(8, 8), key=mx.random.key(0))
+    assert p.dtype == mx.float32
+    v = cast_for_forward(p, mx.bfloat16)
+    assert v.dtype == mx.bfloat16
+    # Master unchanged.
+    assert p.dtype == mx.float32
+
+
+def test_v7_d04_cast_grads_to_master_preserves_fp32():
+    grads = {"a": mx.zeros((4,), dtype=mx.bfloat16),
+             "b": mx.zeros((4,), dtype=mx.float16)}
+    out = cast_grads_to_master(grads, master_dtype=mx.float32)
+    assert out["a"].dtype == mx.float32
+    assert out["b"].dtype == mx.float32
+
+
+def test_v7_d04_hybrid_step_updates_master_in_fp32():
+    master = {"w": mx.array([1.0, 2.0, 3.0], dtype=mx.float32)}
+
+    def loss_and_grad(fwd):
+        # Always-positive grad of 0.1 per element.
+        return mx.array(0.0), {"w": mx.array([0.1, 0.1, 0.1],
+                                              dtype=mx.bfloat16)}
+
+    def apply_gradients(m, g):
+        return {k: m[k] - g[k] for k in m}
+
+    new_master = hybrid_step(
+        master_params=master, fwd_dtype=mx.bfloat16,
+        loss_and_grad=loss_and_grad,
+        apply_gradients=apply_gradients,
+    )
+    # Updated master stays fp32.
+    assert new_master["w"].dtype == mx.float32
+    # Each entry decreased by 0.1.
+    assert mx.allclose(new_master["w"],
+                        mx.array([0.9, 1.9, 2.9], dtype=mx.float32),
+                        atol=1e-3)
+
+
+def test_v7_d04_hybrid_step_grads_cast_to_master_before_apply():
+    """apply_gradients must receive fp32 grads, not bf16."""
+    seen_dtype: dict[str, mx.Dtype] = {}
+
+    master = {"w": mx.zeros((2,), dtype=mx.float32)}
+
+    def loss_and_grad(fwd):
+        return mx.array(0.0), {"w": mx.ones((2,), dtype=mx.bfloat16)}
+
+    def apply_gradients(m, g):
+        seen_dtype["w"] = g["w"].dtype
+        return m
+
+    hybrid_step(master_params=master, fwd_dtype=mx.bfloat16,
+                 loss_and_grad=loss_and_grad,
+                 apply_gradients=apply_gradients)
+    assert seen_dtype["w"] == mx.float32