feat(v7-c02): sharded checkpoint save + index.json + per-rank load

Closes V7-C02 (cppmega-mlx-z6dj) building block:
cppmega_v4.runtime.checkpoint_shard.save_sharded writes one
safetensors file per shard plus a model.index.json mapping
param_key → shard_file_basename. load_sharded reconstructs the
full state; load_sharded_for_rank skips shards not owned by the
calling rank (round-robin shard%world_size==rank).

Tests (tests/v4/test_checkpoint_shard.py): 4/4 —
  * 4-shard round trip equal to source state.
  * index.json carries metadata + weight_map for every key.
  * world_size=2 round-robin: rank 0 + rank 1 combined keys
    equal full state.
  * num_shards=0 rejected.

CheckpointStage._save_weights swap from single-file to save_sharded
is a separate follow-up that needs to coordinate with the
distributed-train path (V7-B01).

Author: apstenku123

cppmega_v4/runtime/checkpoint_shard.py

@@ -0,0 +1,110 @@
+"""V7-C02: sharded safetensors save with index.json manifest.
+
+  save_sharded(state, out_dir, prefix='model', num_shards=N)
+    Writes:
+      out_dir/{prefix}.shard-00001-of-NNNNN.safetensors  (N files)
+      out_dir/{prefix}.index.json
+        { "metadata": {"total_size": int, "num_shards": int},
+          "weight_map": {"param_key": "shard_file_basename"} }
+
+  load_sharded(index_path) → dict[str, mx.array]
+    Reads each shard listed in the index and merges into one dict.
+
+  load_sharded_for_rank(index_path, rank, world_size) → dict
+    Skips shards not owned by this rank (round-robin assignment).
+"""
+
+from __future__ import annotations
+
+import json
+import pathlib
+from typing import Any
+
+import mlx.core as mx
+import safetensors.mlx as st
+
+
+def _shard_name(prefix: str, idx: int, total: int) -> str:
+    return f"{prefix}.shard-{idx + 1:05d}-of-{total:05d}.safetensors"
+
+
+def save_sharded(state: dict[str, mx.array],
+                  out_dir: str | pathlib.Path,
+                  *, prefix: str = "model",
+                  num_shards: int = 1,
+                  metadata: dict[str, str] | None = None
+                  ) -> dict[str, Any]:
+    if num_shards <= 0:
+        raise ValueError("num_shards must be > 0")
+    out_dir = pathlib.Path(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    keys = sorted(state.keys())
+    shards: list[dict[str, mx.array]] = [
+        {} for _ in range(num_shards)]
+    weight_map: dict[str, str] = {}
+    total_size = 0
+    for i, k in enumerate(keys):
+        shard_idx = i % num_shards
+        shards[shard_idx][k] = state[k]
+        fname = _shard_name(prefix, shard_idx, num_shards)
+        weight_map[k] = fname
+        if hasattr(state[k], "shape"):
+            total_size += int(state[k].size) * int(state[k].dtype.size)
+    for i, shard in enumerate(shards):
+        fname = _shard_name(prefix, i, num_shards)
+        st.save_file(shard, str(out_dir / fname), metadata=metadata)
+    index = {
+        "metadata": {"total_size": int(total_size),
+                      "num_shards": int(num_shards)},
+        "weight_map": weight_map,
+    }
+    index_path = out_dir / f"{prefix}.index.json"
+    index_path.write_text(json.dumps(index, indent=2, sort_keys=True))
+    return index
+
+
+def load_sharded(index_path: str | pathlib.Path) -> dict[str, mx.array]:
+    index_path = pathlib.Path(index_path)
+    index = json.loads(index_path.read_text())
+    base = index_path.parent
+    out: dict[str, mx.array] = {}
+    seen_shards: set[str] = set()
+    for k, shard_name in index["weight_map"].items():
+        if shard_name in seen_shards:
+            continue
+        seen_shards.add(shard_name)
+        out.update(st.load_file(str(base / shard_name)))
+    return out
+
+
+def load_sharded_for_rank(index_path: str | pathlib.Path,
+                           *, rank: int, world_size: int
+                           ) -> dict[str, mx.array]:
+    """Round-robin shard ownership: rank r reads shards i where
+    i % world_size == r."""
+    if world_size <= 0 or not (0 <= rank < world_size):
+        raise ValueError("invalid rank/world_size")
+    index_path = pathlib.Path(index_path)
+    index = json.loads(index_path.read_text())
+    base = index_path.parent
+    num_shards = int(index["metadata"]["num_shards"])
+    owned = {i for i in range(num_shards) if i % world_size == rank}
+    out: dict[str, mx.array] = {}
+    for k, shard_name in index["weight_map"].items():
+        # Reverse-derive shard index from filename "*-NNNNN-of-MMMMM"
+        try:
+            idx = int(shard_name.split("-")[-3]) - 1
+        except (ValueError, IndexError):
+            continue
+        if idx in owned:
+            if shard_name not in {n for n in
+                                  (index["weight_map"][kk]
+                                   for kk in index["weight_map"])}:
+                continue
+            # Load each shard once.
+            if shard_name not in out.get("__loaded__", set()):
+                out.update(st.load_file(str(base / shard_name)))
+    return out
+
+
+__all__ = ["save_sharded", "load_sharded", "load_sharded_for_rank"]

tests/v4/test_checkpoint_shard.py

@@ -0,0 +1,69 @@
+"""V7-C02: sharded checkpoint save + index.json + load round-trip."""
+
+from __future__ import annotations
+
+import json
+
+import mlx.core as mx
+import pytest
+
+from cppmega_v4.runtime.checkpoint_shard import (
+    load_sharded, load_sharded_for_rank, save_sharded,
+)
+
+
+def _state(seed: int = 0) -> dict[str, mx.array]:
+    return {
+        "layer0.w": mx.random.normal(shape=(8, 16),
+                                       key=mx.random.key(seed)),
+        "layer0.b": mx.random.normal(shape=(16,),
+                                       key=mx.random.key(seed + 1)),
+        "layer1.w": mx.random.normal(shape=(16, 8),
+                                       key=mx.random.key(seed + 2)),
+        "layer1.b": mx.random.normal(shape=(8,),
+                                       key=mx.random.key(seed + 3)),
+        "head.w":   mx.random.normal(shape=(8, 4),
+                                       key=mx.random.key(seed + 4)),
+    }
+
+
+def test_v7_c02_round_trip_4_shards(tmp_path):
+    state = _state()
+    save_sharded(state, tmp_path, num_shards=4)
+    index = json.loads((tmp_path / "model.index.json").read_text())
+    assert index["metadata"]["num_shards"] == 4
+    assert set(index["weight_map"].keys()) == set(state.keys())
+
+    loaded = load_sharded(tmp_path / "model.index.json")
+    assert set(loaded.keys()) == set(state.keys())
+    for k in state:
+        assert mx.allclose(loaded[k], state[k], atol=0.0)
+
+
+def test_v7_c02_index_json_has_weight_map_and_total_size(tmp_path):
+    state = _state()
+    save_sharded(state, tmp_path, num_shards=2)
+    idx = json.loads((tmp_path / "model.index.json").read_text())
+    assert "metadata" in idx and "weight_map" in idx
+    assert idx["metadata"]["total_size"] > 0
+    # Every param key maps to a shard filename.
+    for k in state:
+        assert idx["weight_map"][k].endswith(".safetensors")
+
+
+def test_v7_c02_load_sharded_for_rank_round_robin(tmp_path):
+    state = _state()
+    save_sharded(state, tmp_path, num_shards=4)
+    # 4 shards, world_size=2 → rank 0 owns shards {0,2}; rank 1 owns {1,3}
+    r0 = load_sharded_for_rank(
+        tmp_path / "model.index.json", rank=0, world_size=2)
+    r1 = load_sharded_for_rank(
+        tmp_path / "model.index.json", rank=1, world_size=2)
+    # Combined coverage equals full state.
+    combined = {**r0, **r1}
+    assert set(combined.keys()) == set(state.keys())
+
+
+def test_v7_c02_num_shards_validation(tmp_path):
+    with pytest.raises(ValueError):
+        save_sharded({"k": mx.zeros((2,))}, tmp_path, num_shards=0)