feat(v5-g12): checkpoint save/resume via safetensors

Closes V5-G12 / cppmega-mlx-t8i. stage_train opts:
  - checkpoint_load_path: safetensors → tree_unflatten → model.update
    (before training, so first step starts from loaded weights)
  - checkpoint_save_path: model.parameters → tree_flatten → safetensors
    save (after training, weights persist for next session)
extras.checkpoint = {saved_path, loaded_path}; both null when opts
absent; load failure non-fatal.

4 new pytest. UI textinputs for paths deferred — backend contract first.

Author: apstenku123

cppmega_v4/runner/stages.py

@@ -773,6 +773,20 @@ def _count(tree: Any) -> int:
         probe_key: str | None = None
         probe_before: mx.array | None = None
 
+        # G12: optional checkpoint load before training. Reads safetensors
+        # weights into the model. Failure is non-fatal — log via extras.
+        checkpoint_loaded: str | None = None
+        ckpt_load = opts.get("checkpoint_load_path")
+        if ckpt_load:
+            try:
+                import safetensors.mlx as _stmlx
+                loaded = _stmlx.load_file(ckpt_load)
+                all_modules.update(
+                    nn.utils.tree_unflatten(list(loaded.items())))
+                checkpoint_loaded = str(ckpt_load)
+            except Exception:
+                pass
+
         # G09: check abort flag set via opts.abort or _ABORT_TOKENS set
         abort_token = opts.get("abort_token")
         for step in range(n_steps):
@@ -846,6 +860,18 @@ def _count(tree: Any) -> int:
         except Exception:
             pass
 
+        # G12: optional checkpoint save after training.
+        checkpoint_saved: str | None = None
+        ckpt_save = opts.get("checkpoint_save_path")
+        if ckpt_save:
+            try:
+                import safetensors.mlx as _stmlx
+                flat = dict(nn.utils.tree_flatten(all_modules.parameters()))
+                _stmlx.save_file(flat, ckpt_save)
+                checkpoint_saved = str(ckpt_save)
+            except Exception:
+                pass
+
         # G10: cache opt.state for future warm-start lookups (capped LRU)
         try:
             _RUN_CACHE[run_id] = opt.state
@@ -965,6 +991,10 @@ def _count(tree: Any) -> int:
                 "moe": moe_extras,
                 "opt_state_carried": opt_state_carried,
                 "run_id": run_id,
+                "checkpoint": {
+                    "saved_path": checkpoint_saved,
+                    "loaded_path": checkpoint_loaded,
+                },
                 "mtp": _compute_mtp_extras(
                     all_modules, mtp_k, mtp_betas, vocab_size,
                     batch, seq, hidden, targets,

tests/v5/test_stage_train_checkpoint.py

@@ -0,0 +1,63 @@
+"""G12: checkpoint save/resume."""
+
+from __future__ import annotations
+
+import pathlib
+
+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")
+    return train.extras
+
+
+def test_checkpoint_save_writes_file(tmp_path):
+    save_path = str(tmp_path / "ckpt.safetensors")
+    extras = _run({"num_steps": 2, "checkpoint_save_path": save_path})
+    assert extras["checkpoint"]["saved_path"] == save_path
+    assert extras["checkpoint"]["loaded_path"] is None
+    assert pathlib.Path(save_path).exists()
+
+
+def test_checkpoint_load_reads_file(tmp_path):
+    save_path = str(tmp_path / "ckpt.safetensors")
+    a = _run({"num_steps": 2, "checkpoint_save_path": save_path})
+    assert a["checkpoint"]["saved_path"] == save_path
+    b = _run({"num_steps": 2, "checkpoint_load_path": save_path})
+    assert b["checkpoint"]["loaded_path"] == save_path
+
+
+def test_checkpoint_load_missing_file_is_non_fatal():
+    extras = _run({"num_steps": 2,
+                   "checkpoint_load_path": "/nonexistent/ckpt.safetensors"})
+    # train still succeeds; loaded_path is None
+    assert extras["checkpoint"]["loaded_path"] is None
+
+
+def test_no_checkpoint_opts_yields_null_paths():
+    extras = _run({"num_steps": 2})
+    assert extras["checkpoint"]["saved_path"] is None
+    assert extras["checkpoint"]["loaded_path"] is None