test(v7-b11..18): distributed-semantics parity suite (8 tests)

Closes the 8-item B-block 'simulation without real multi-process'
audit:

B11 test_fake_ranks_bit_identical.py — fake_ranks=4 mean-reduced
    grads match single-rank grads at atol=1e-6 (the H20 replay path
    must be bit-equivalent on identical input).
B12 test_tp_proxy_loss_parity.py — column_split_forward and
    row_split_forward parity vs unsharded matmul for tp∈{1,2,4,8}
    plus a 2-brick chain (ColumnParallel→RowParallel composition).
B13 test_pp_proxy_loss_parity.py — pipeline_forward matches
    sequential composition for num_microbatches∈{1,2,4,8}; the new
    pipeline_forward_real with 1f1b schedule matches the legacy
    pipeline_forward on world_size==1.
B14 test_collective_proxy_sum_scatter.py — reduce_scatter(arr,4)
    mean-reduces 4 N-element chunks to a single chunk equal to
    arr.mean(axis=0) on the reshaped (4,N) view; all_gather+
    reduce_scatter round-trips to shard for world∈{1,2,4,8};
    misaligned shape raises ValueError.
B15 bench_allreduce_distributed.py --simulated — forces single-
    process path, writes reports/allreduce_simulated_<date>.csv
    with a warm_ms column; test_pf0g_csv_shape.py asserts the 10
    required columns + warm_ms > 0.
B16 test_distributed_world1_fallback.py — distributed.init(
    force_single=True) returns world=1 / rank=0 / real=False;
    stage_train survives compile_mode='regional' + fake_ranks=1 +
    m3_ultra_solo with extras.distributed.real=False.
B17 test_parallelism_gotcha_matrix.py — pytest.parametrize 3
    topology-agnostic gotchas (fsdp2_whole_compile,
    megatron_tp_whole_compile, fp8_grad_duplication) × 4 topology
    factories (h100_8x, gb10_quarter, m3_ultra_solo, tpu_v6e_8) =
    12 cells. Each cell asserts the gotcha fires on the matched
    spec; topologies without the required mesh axis (e.g.
    gb10_quarter has no tp axis) skip gracefully via
    _topology_supports_kinds(). Plus a clean-spec test that the
    three headline gotchas stay silent, plus a size-floor of 15
    on the GOTCHAS table.
B18 test_runtime_simulation_allsum_accuracy.py — world_size=1 is
    identity; simulated mode is zero-copy; manual mean reduce
    over 4 distinct fp32 arrays equals (a+b+c+d)/4 bit-exact
    regardless of summation order.

Bench update: scripts/bench_allreduce_distributed.py grew the
--simulated CLI flag (force single-process), the warm_ms column,
and the allreduce_simulated_<date>.csv filename pattern.

26 passed, 5 gracefully skipped (topology without the requested
mesh axis); no failures.

Author: apstenku123

scripts/bench_allreduce_distributed.py

@@ -46,24 +46,47 @@ def main() -> int:
                          default=[0.25, 1, 4, 16, 64])
     parser.add_argument("--n-iter", type=int, default=50)
     parser.add_argument("--out-dir", default="reports")
+    # V7-B15: --simulated emits the proxy bench result under a
+    # canonical filename pattern (reports/allreduce_simulated_<date>.csv)
+    # so the existing pf0g bd ticket has a deterministic artifact.
+    parser.add_argument("--simulated", action="store_true",
+                         help="Force the single-process simulation path "
+                              "and write reports/allreduce_simulated_"
+                              "<date>.csv with warm_ms column.")
     args = parser.parse_args()
 
-    info = _d.init()
+    if args.simulated:
+        # V7-B15: force the simulation path regardless of mx.distributed
+        # availability so the artifact filename is deterministic.
+        _d.reset_for_test()
+        info = _d.init(force_single=True)
+    else:
+        info = _d.init()
     out_dir = pathlib.Path(args.out_dir)
     out_dir.mkdir(parents=True, exist_ok=True)
 
+    import time as _time
+    t_warm = _time.perf_counter()
     rows = [_bench(m, args.n_iter) for m in args.buf_mb]
+    warm_ms = (_time.perf_counter() - t_warm) * 1000.0
+    for r in rows:
+        # V7-B15: add warm_ms column so the pf0g CSV-shape test has a
+        # field to assert on.
+        r["warm_ms"] = round(warm_ms, 4)
 
     # Only rank 0 writes the report so a multi-process launch
     # doesn't fight on the same file.
     if info.rank != 0:
         return 0
 
     date = _dt.datetime.utcnow().strftime("%Y%m%d_%H%M%S")
-    csv_path = out_dir / f"bench_allreduce_distributed_{date}.csv"
+    if args.simulated:
+        csv_path = out_dir / f"allreduce_simulated_{date}.csv"
+    else:
+        csv_path = out_dir / f"bench_allreduce_distributed_{date}.csv"
     keys = ["buf_mb", "bytes", "world_size", "rank", "backend", "real",
             "all_reduce_ms_per_iter", "all_gather_ms_per_iter",
-            "reduce_scatter_ms_per_iter"]
+            "reduce_scatter_ms_per_iter", "warm_ms"]
     with csv_path.open("w", newline="") as f:
         w = csv.DictWriter(f, fieldnames=keys)
         w.writeheader()

tests/v4/test_collective_proxy_sum_scatter.py

@@ -0,0 +1,51 @@
+"""V7-B14: collective_proxy.reduce_scatter bit-exact vs naive mean."""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+from cppmega_v4.runtime.collective_proxy import (
+    all_gather, reduce_scatter,
+)
+
+
+def test_reduce_scatter_world_size_4_mean_reduces_chunks():
+    """Build a deliberately-replicated arr of shape (4*N,) where each
+    chunk of N elements represents a per-rank value; reduce_scatter
+    must mean-reduce them into a single chunk identical to arr.mean(0)
+    when reshaped to (4, N)."""
+    N = 8
+    world = 4
+    chunks = [
+        mx.full((N,), float(r + 1), dtype=mx.float32)
+        for r in range(world)
+    ]
+    arr = mx.concatenate(chunks, axis=0)
+    out = reduce_scatter(arr, world)
+    # Mean across the 4 chunks = (1+2+3+4)/4 = 2.5 broadcast over N.
+    expected = mx.full((N,), 2.5, dtype=mx.float32)
+    assert mx.array_equal(out, expected).item()
+
+
+def test_reduce_scatter_world_size_1_is_identity():
+    arr = mx.arange(8.0)
+    out = reduce_scatter(arr, 1)
+    assert mx.array_equal(out, arr).item()
+
+
+def test_all_gather_round_trips_with_reduce_scatter():
+    """all_gather(shard, W) followed by reduce_scatter(full, W) returns
+    the original shard — V7-B contract."""
+    for world in (1, 2, 4, 8):
+        shard = mx.arange(4.0)
+        full = all_gather(shard, world)
+        back = reduce_scatter(full, world)
+        assert mx.array_equal(back, shard).item(), (
+            f"world={world} all_gather→reduce_scatter broke identity")
+
+
+def test_reduce_scatter_rejects_misaligned_shape():
+    import pytest
+    arr = mx.arange(7.0)   # 7 % 4 != 0
+    with pytest.raises(ValueError):
+        reduce_scatter(arr, 4)

tests/v4/test_distributed_world1_fallback.py

@@ -0,0 +1,67 @@
+"""V7-B16: mx.distributed world=1 fallback + stage_train survives."""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import pytest
+
+from cppmega_v4.jsonrpc.schema import VerifyParams
+from cppmega_v4.runner import Pipeline, run_pipeline
+from cppmega_v4.runtime import distributed as _d
+
+
+@pytest.fixture(autouse=True)
+def _clean():
+    _d.reset_for_test()
+    yield
+    _d.reset_for_test()
+
+
+def test_mx_distributed_init_strict_false_returns_world_1():
+    """mx.distributed.init(strict=False) returns a singleton size 1
+    group when no launcher is active. Our wrapper must reflect that."""
+    info = _d.init(force_single=True)
+    assert info.world_size == 1
+    assert info.rank == 0
+    assert info.real is False
+    # is_distributed() must be False in fallback mode.
+    assert _d.is_distributed() is False
+
+
+def test_stage_train_survives_compile_regional_plus_fake_ranks_1():
+    """compile_mode='regional' + fake_ranks=1 + world=1 fallback must
+    not raise inside stage_train. This is the canonical 'works on my
+    laptop' path."""
+    spec_payload = {
+        "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]}]},
+        "sharding": {
+            "topology": {"factory": "m3_ultra_solo", "kwargs": {}},
+            "axis_assignments": [],
+            "compile_mode": "regional",
+        },
+    }
+    spec = VerifyParams.model_validate(spec_payload)
+    report = run_pipeline(spec, Pipeline.from_dict({
+        "stages": ["parse", "verify_build_spec", "build_model", "train"],
+        "stage_options": {"train": {"num_steps": 2, "fake_ranks": 1}},
+    }))
+    train = next(s for s in report.stages if s.name == "train")
+    assert train.status == "ok", f"stage_train failed: {train.error}"
+    # distributed.real must be False on the single-process path.
+    dist = (train.extras or {}).get("distributed", {})
+    assert dist.get("real") is False
+    assert dist.get("world_size") == 1

tests/v4/test_fake_ranks_bit_identical.py

@@ -0,0 +1,54 @@
+"""V7-B11: fake_ranks=4 mean-reduced grads are bit-equivalent to single-rank.
+
+The H20 replay path runs the same forward+backward N times on identical
+inputs and mean-reduces the gradient tree. Since (g + g + g + g) / 4 == g
+when every replay sees the same batch, fake_ranks=4 must produce grads
+indistinguishable from fake_ranks=1 within fp32 rounding.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+def _build_tiny():
+    m = nn.Sequential(
+        nn.Linear(8, 16),
+        nn.Linear(16, 8),
+    )
+    return m
+
+
+def _grad_tree(model: nn.Module, x: mx.array) -> dict:
+    def loss_fn(m: nn.Module, _x: mx.array) -> mx.array:
+        return mx.mean(m(_x) ** 2)
+    lvg = nn.value_and_grad(model, loss_fn)
+    _, grads = lvg(model, x)
+    mx.eval(grads)
+    return dict(nn.utils.tree_flatten(grads))
+
+
+def test_fake_ranks_4_mean_matches_single_rank():
+    mx.random.seed(0xBEEF)
+    m1 = _build_tiny()
+    mx.eval(m1.parameters())
+    x = mx.random.normal(shape=(2, 8), key=mx.random.key(0))
+
+    g1 = _grad_tree(m1, x)
+
+    # fake_ranks=4: replay the same backward 4× on identical input then
+    # mean-reduce — that's exactly the H20 simulation in stage_train.
+    accum: dict = {}
+    fake_ranks = 4
+    for _ in range(fake_ranks):
+        gi = _grad_tree(m1, x)
+        for k, v in gi.items():
+            accum[k] = v if k not in accum else (accum[k] + v)
+    g_mean = {k: v / float(fake_ranks) for k, v in accum.items()}
+    mx.eval(g_mean)
+
+    for k in g1:
+        assert mx.allclose(g_mean[k], g1[k], atol=1e-6).item(), (
+            f"grad {k!r} drifted: max-diff="
+            f"{float(mx.max(mx.abs(g_mean[k] - g1[k])).item())}")