feat(v7-c04): streaming safetensors checkpoint loader

Closes V7-C04 (cppmega-mlx-9t1m): cppmega_v4.runtime.checkpoint_streaming
iterates per-tensor via safetensors.safe_open so a multi-GB resume
doesn't double host memory.

API:
  - streaming_load(path) → Iterator[(key, mx.array)]
  - streaming_load_all(path, *, progress_cb) → dict, with every-5%
    callback fired as (loaded, total).
  - streaming_load_sharded(paths) → cross-shard iterator (V7-C02
    composition).

Tests (tests/v4/test_checkpoint_streaming.py): 3/3 —
  * Streaming load bit-equal to safetensors.mlx.load_file.
  * Progress callback fires with total count.
  * Sharded iterator merges two files' keys + values.

CheckpointStage._load_weights swap to streaming_load is a separate
follow-up that ALSO needs to retire the existing bulk call without
breaking V01 rng-key roundtrip; this commit lands the building
block + contract.

Author: apstenku123

cppmega_v4/runtime/checkpoint_streaming.py

@@ -0,0 +1,65 @@
+"""V7-C04: streaming safetensors loader for large checkpoints.
+
+safetensors.mlx.load_file mmaps/reads the full blob into one host
+buffer. For multi-GB models this doubles peak memory at resume.
+This module iterates per-tensor via safe_open, loads, hands to the
+caller, and drops the host reference before the next tensor.
+
+Usage:
+    for key, tensor in streaming_load(path):
+        live_param_tree[key] = tensor
+        # implicit drop of `tensor` reference at next iter step.
+"""
+
+from __future__ import annotations
+
+import pathlib
+from typing import Callable, Iterator, Iterable
+
+import mlx.core as mx
+
+
+def streaming_load(path: str | pathlib.Path) -> Iterator[
+        tuple[str, mx.array]]:
+    """Yield (key, mx.array) per tensor in the safetensors file."""
+    from safetensors import safe_open
+    with safe_open(str(path), framework="mlx") as f:
+        for k in f.keys():
+            yield k, f.get_tensor(k)
+
+
+def streaming_load_all(path: str | pathlib.Path,
+                        *,
+                        progress_cb: Callable[[int, int], None] | None
+                        = None) -> dict[str, mx.array]:
+    """Eagerly collect all tensors. Same as load_file but goes
+    through the per-tensor iterator; useful when callers want a
+    progress callback (loaded, total) every 5%.
+    """
+    from safetensors import safe_open
+    out: dict[str, mx.array] = {}
+    with safe_open(str(path), framework="mlx") as f:
+        keys = list(f.keys())
+        total = len(keys)
+        last_pct = -1
+        for i, k in enumerate(keys):
+            out[k] = f.get_tensor(k)
+            if progress_cb:
+                pct = int((i + 1) / max(1, total) * 100)
+                if pct // 5 != last_pct // 5:
+                    progress_cb(i + 1, total)
+                    last_pct = pct
+    return out
+
+
+def streaming_load_sharded(paths: Iterable[str | pathlib.Path]
+                            ) -> Iterator[tuple[str, mx.array]]:
+    """V7-C04 + V7-C02 composition: walk multiple shard files,
+    yielding (key, tensor) across all of them."""
+    for p in paths:
+        yield from streaming_load(p)
+
+
+__all__ = [
+    "streaming_load", "streaming_load_all", "streaming_load_sharded",
+]

tests/v4/test_checkpoint_streaming.py

@@ -0,0 +1,59 @@
+"""V7-C04: streaming safetensors load equality + sharded composition."""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import pytest
+import safetensors.mlx as st
+
+from cppmega_v4.runtime.checkpoint_streaming import (
+    streaming_load, streaming_load_all, streaming_load_sharded,
+)
+
+
+def _write(tmp_path, tensors):
+    p = str(tmp_path / "weights.safetensors")
+    st.save_file(tensors, p)
+    return p
+
+
+def test_v7_c04_streaming_load_matches_bulk_load(tmp_path):
+    tensors = {
+        "layer0.w": mx.random.normal(shape=(8, 16),
+                                       key=mx.random.key(0)),
+        "layer1.w": mx.random.normal(shape=(16, 8),
+                                       key=mx.random.key(1)),
+    }
+    path = _write(tmp_path, tensors)
+    bulk = st.load_file(path)
+    stream: dict[str, mx.array] = {}
+    for k, t in streaming_load(path):
+        stream[k] = t
+    assert set(bulk.keys()) == set(stream.keys())
+    for k in bulk:
+        assert mx.allclose(bulk[k], stream[k], atol=0.0)
+
+
+def test_v7_c04_streaming_load_all_progress_callback(tmp_path):
+    tensors = {f"t{i}": mx.zeros((2, 2)) for i in range(20)}
+    path = _write(tmp_path, tensors)
+    seen: list[tuple[int, int]] = []
+    out = streaming_load_all(
+        path, progress_cb=lambda done, tot: seen.append((done, tot)))
+    assert len(out) == 20
+    # At least one progress callback fired and totals match.
+    assert seen
+    assert all(t == 20 for _, t in seen)
+
+
+def test_v7_c04_streaming_load_sharded(tmp_path):
+    t1 = {"a.w": mx.array([[1.0, 2.0]])}
+    t2 = {"b.w": mx.array([[3.0, 4.0]])}
+    p1 = str(tmp_path / "s0.safetensors")
+    p2 = str(tmp_path / "s1.safetensors")
+    st.save_file(t1, p1)
+    st.save_file(t2, p2)
+    merged = dict(streaming_load_sharded([p1, p2]))
+    assert set(merged.keys()) == {"a.w", "b.w"}
+    assert mx.allclose(merged["a.w"], t1["a.w"], atol=0.0)
+    assert mx.allclose(merged["b.w"], t2["b.w"], atol=0.0)