feat(v5-g10): optimizer state warm-start across sequential train runs

Closes V5-G10 / cppmega-mlx-jrp. Backend opts.run_id + LRU
_RUN_CACHE (size 8) of opt.state. opts.continue_from_run_id loads
cached state pre-train. extras.opt_state_carried + extras.run_id.

UI checkbox surface deferred — backend exposure first, train-warm-start
control TopBar follow-up.

3/3 pytest green.

Author: apstenku123

cppmega_v4/runner/stages.py

@@ -561,6 +561,19 @@ def loss_fn(model: nn.Module, emb: mx.array, tgt: mx.array) -> mx.array:
                     pass
         all_modules = nn.Sequential(*modules, *lm_heads)
         opt, optimizer_kind = _build_optimizer(spec_optim, lr)
+        # G10: optional optimizer state warm-start across sequential
+        # Train runs. opts.continue_from_run_id refers to a prior
+        # run cached in _RUN_CACHE. If hit, restore opt.state so the
+        # second Train's losses[0] picks up where the first left off.
+        opt_state_carried = False
+        run_id = str(opts.get("run_id") or id(opt))
+        continue_from = opts.get("continue_from_run_id")
+        if continue_from and continue_from in _RUN_CACHE:
+            try:
+                opt.state = _RUN_CACHE[continue_from]
+                opt_state_carried = True
+            except Exception:
+                pass
         # V4-9: when hybrid, count params routed to each bucket so e2e can
         # prove the split predicate actually saw 2D vs 1D/3D parameters.
         muon_group_size: int | None = None
@@ -725,6 +738,14 @@ def _count(tree: Any) -> int:
         except Exception:
             pass
 
+        # G10: cache opt.state for future warm-start lookups (capped LRU)
+        try:
+            _RUN_CACHE[run_id] = opt.state
+            if len(_RUN_CACHE) > 8:
+                _RUN_CACHE.pop(next(iter(_RUN_CACHE)))
+        except Exception:
+            pass
+
         after_flat = dict(nn.utils.tree_flatten(all_modules.parameters()))
         delta = 0.0
         if probe_key is not None and probe_before is not None:
@@ -814,6 +835,8 @@ def _count(tree: Any) -> int:
                 "train_dtype": train_dtype,
                 "master_dtype": master_dtype,
                 "fp8_active": fp8_active,
+                "opt_state_carried": opt_state_carried,
+                "run_id": run_id,
                 "mtp": _compute_mtp_extras(
                     all_modules, mtp_k, mtp_betas, vocab_size,
                     batch, seq, hidden, targets,
@@ -877,6 +900,11 @@ def _tokenize_parquet_text(
         return [], None
 
 
+# G10: in-process LRU cache of opt.state by run_id. Used for warm-start
+# across sequential Train clicks in the same backend session.
+_RUN_CACHE: dict[str, Any] = {}
+
+
 _REWRITER_FACTORIES: dict[str, Callable[..., Any]] = {}
 
 

tests/v5/test_stage_train_warm_start.py

@@ -0,0 +1,57 @@
+"""G10: optimizer state warm-start across sequential train runs."""
+
+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": {}},
+                {"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(opts: dict) -> dict:
+    report = run_pipeline(_spec(), Pipeline.from_dict({
+        "stages": ["parse", "verify_build_spec", "build_model", "train"],
+        "stage_options": {"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_first_run_no_carry():
+    extras = _run({"num_steps": 2, "run_id": "run-A"})
+    assert extras["opt_state_carried"] is False
+    assert extras["run_id"] == "run-A"
+
+
+def test_continue_from_loads_carried_state():
+    """Two sequential runs: second sets continue_from_run_id=A → carried."""
+    a = _run({"num_steps": 2, "run_id": "run-A2"})
+    b = _run({"num_steps": 2, "run_id": "run-B2",
+              "continue_from_run_id": "run-A2"})
+    assert a["opt_state_carried"] is False
+    assert b["opt_state_carried"] is True
+
+
+def test_continue_from_unknown_id_no_carry():
+    extras = _run({"num_steps": 2, "run_id": "run-C",
+                   "continue_from_run_id": "nonexistent-id"})
+    assert extras["opt_state_carried"] is False