feat(v01): rng_key roundtrip → strict 1e-5 H19 loss continuation

V6 H19 left the strict 1e-5 bit-identity claim unreachable because
stage_train's per-step rng_key (mx.random.split each step) was NOT
saved in the checkpoint, so a resumed Train consumed a fresh
key(0) data stream — different synthetic targets at step N+1 vs
the contiguous run's step N+1.

Fix: opt_state save bundle now also writes the active rng_key under
the "_rng_key" entry; load pops it back before tree_unflatten and
seeds the resumed loop with the restored key. New
extras.checkpoint.rng_key_loaded boolean lets tests pin which path
was taken.

Tests (tests/v4/test_stage_train_rng_roundtrip.py): 3/3 —
  * Contiguous 5-step losses[4] == (4-save + load + 1-step) losses[0]
    within 1e-5 once weights + opt.state + rng_key all restored.
  * rng_key_loaded=False without opt_state_load_path.
  * rng_key_loaded=True after opt_state_load with the side-car file.
- pytest stage_train + checkpoint regression: 57/57 passing.

Author: apstenku123

cppmega_v4/runner/stages.py

@@ -682,7 +682,11 @@ def loss_fn(model: nn.Module, emb: mx.array, tgt: mx.array) -> mx.array:
         # when absent (preserves E-4 matrix behaviour). data_source +
         # token_count surface in extras so deep e2e can prove the real-
         # data path actually executed instead of silently degrading.
+        # V01: per-step random data key. Round-tripped through the
+        # opt.state safetensors side-car (key "_rng_key") so a resumed
+        # Train picks up the exact same data stream as a contiguous run.
         rng_key = mx.random.key(0)
+        rng_key_loaded = False
         data_source = "synthetic"
         token_count = 0
         tokenizer_used: str | None = None
@@ -1118,7 +1122,14 @@ def _count(tree: Any) -> int:
             try:
                 import safetensors.mlx as _stmlx
                 loaded_st = _stmlx.load_file(opt_state_load)
+                # V01: pull rng_key out of the opt-state bundle before
+                # tree_unflatten so it doesn't pollute opt.state. mlx
+                # rng keys are mx.uint32 arrays of shape (2,).
+                rng_buf = loaded_st.pop("_rng_key", None)
                 opt.state = nn.utils.tree_unflatten(list(loaded_st.items()))
+                if rng_buf is not None:
+                    rng_key = rng_buf
+                    rng_key_loaded = True
                 opt_state_loaded_path = str(opt_state_load)
             except FileNotFoundError as exc:
                 opt_state_warning = (
@@ -1260,17 +1271,20 @@ def _count(tree: Any) -> int:
                 pass
         # H19: opt.state save → separate file so a follow-up Train can
         # resume exactly where this one left off (Adam moments + step).
+        # V01: also persist rng_key under the "_rng_key" entry so the
+        # resumed run consumes the same synthetic data stream → enables
+        # strict 1e-5 loss continuation in the matched-fragment test.
         opt_state_save = opts.get("opt_state_save_path")
         if opt_state_save:
             try:
                 import safetensors.mlx as _stmlx
                 opt_flat = dict(nn.utils.tree_flatten(opt.state))
-                # safetensors only accepts mx.array values; strip
-                # scalars / non-array entries from opt.state.
                 opt_arrays = {
                     k: v for k, v in opt_flat.items()
                     if hasattr(v, "shape")
                 }
+                if hasattr(rng_key, "shape"):
+                    opt_arrays["_rng_key"] = rng_key
                 _stmlx.save_file(opt_arrays, opt_state_save)
                 opt_state_saved_path = str(opt_state_save)
             except Exception:
@@ -1414,6 +1428,7 @@ def _count(tree: Any) -> int:
                     "opt_state_saved_path": opt_state_saved_path,
                     "opt_state_loaded_path": opt_state_loaded_path,
                     "opt_state_warning": opt_state_warning,
+                    "rng_key_loaded": rng_key_loaded,
                 },
                 "mtp": _compute_mtp_extras(
                     all_modules, mtp_k, mtp_betas, vocab_size,

tests/v4/test_stage_train_rng_roundtrip.py

@@ -0,0 +1,81 @@
+"""V01: rng_key roundtrip enables strict 1e-5 loss continuation.
+
+Without rng_key save+load, the per-step random target stream restarts
+from key(0) on a resumed Train, so resumed.losses[0] does NOT equal
+the contiguous run's losses[N]. V01 persists the active key in the
+opt-state safetensors side-car under "_rng_key" and restores it on
+load. This test pins the strict 1e-5 bound.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+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": {"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 _train(spec, **opts) -> dict:
+    rep = run_pipeline(spec, 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", f"train failed: {tr.error}"
+    return tr.extras
+
+
+def test_v01_rng_key_roundtrip_strict_loss_parity(tmp_path):
+    """Contiguous 5-step Run A losses[4] == (4-save + load + 1-step)
+    Run B losses[0] within 1e-5 once weights + opt.state + rng_key
+    are all restored."""
+    a = _train(_spec(), num_steps=5)
+    save_ckpt = str(tmp_path / "ck.safetensors")
+    save_opt = str(tmp_path / "opt.safetensors")
+    _train(_spec(), num_steps=4,
+           checkpoint_save_path=save_ckpt,
+           opt_state_save_path=save_opt)
+    resumed = _train(_spec(), num_steps=1,
+                     checkpoint_load_path=save_ckpt,
+                     opt_state_load_path=save_opt)
+    assert resumed["checkpoint"]["rng_key_loaded"] is True
+
+    diff = abs(a["losses"][4] - resumed["losses"][0])
+    assert diff < 1e-5, (
+        f"V01 rng_key roundtrip insufficient: contig l4={a['losses'][4]}, "
+        f"resumed l0={resumed['losses'][0]}, diff={diff}"
+    )
+
+
+def test_v01_rng_key_loaded_flag_false_without_opt_state():
+    """Without opt_state_load_path, rng_key_loaded stays False."""
+    extras = _train(_spec(), num_steps=2)
+    assert extras["checkpoint"]["rng_key_loaded"] is False
+
+
+def test_v01_rng_key_loaded_flag_when_present(tmp_path):
+    """opt_state_load_path with a key-bearing file flips the flag."""
+    save_opt = str(tmp_path / "opt.safetensors")
+    _train(_spec(), num_steps=2, opt_state_save_path=save_opt)
+    extras = _train(_spec(), num_steps=1, opt_state_load_path=save_opt)
+    assert extras["checkpoint"]["rng_key_loaded"] is True