feat(v4/parallelism): PSpec Stage B — distributed memory accounting

Adds cppmega_v4.parallelism.distributed_memory (ParallelismSpec.md §3.3,
§6 B). Per-rank byte breakdown that covers the lessons-learned pitfalls
from ../cppmega (Megatron EP=4/8 production) and ../nanochat
(/memory_estimator.py 1265 lines).

cppmega_v4/parallelism/distributed_memory.py:

  PerRankMemory frozen dataclass — full byte breakdown for ONE rank:
    weights / grads / optimizer_state / master_weights /
    activations / fsdp_allgather_peak / kv_cache /
    moe_routing_buffers / collective_workspace / framework_overhead /
    total. .fits_on_device(device, headroom) gate.

  DistributedMemoryReport — mesh-wide rollup with per_rank tuple,
    explicit duplication_bytes (FP8+bf16 grad overhead), explicit
    master_weights_overhead_bytes, explicit
    kernel_boundary_materialisation_bytes (Megatron RowParallel),
    worst-rank picker, bottleneck_diagnostics surfacing the three
    pain points loudly. .fits_on_topology(headroom) + .summary() for GUI.

  estimate_distributed_memory(build_spec, sharding, *, training):
    1) Single-device baseline from cppmega_v4.spec (VBSpec)
    2) Apply shard factors:
       - TP divides per-layer weights by tp_degree
       - EP divides MoE expert weights by ep_degree
       - FSDP/ZeRO steady-state divides optimizer + grads (FSDP/ZeRO-2)
       - FSDP all-gather peak adds ~10% (one block briefly unsharded)
    3) Per-rank duplication accounting:
       - fp8_enabled → fp8 weights + bf16 grads (no fp8 backward;
         see nanochat/fp8_training.py:1-32)
       - master_weights_fp32 → +fp32 copy alongside bf16/fp8
       - Megatron RowParallel (TP>1) → +full (B,S,H) at AllReduce
         boundary on backward
       - SP → small-replicated-param AllReduce workspace
    4) Framework overhead: CUDA 2 GB / XLA 6.4 GB / MLX 1 GB
       depending on first-device kind.

  Optimizer-state width via OPTIM_BYTES_PER_PARAM:
    AdamW=8 (m+v fp32), Muon=2 (momentum bf16), Hybrid=5, SGD=0.

  Activation checkpointing scales activations: off → 1.5×, full → 0.5×,
  selective → unchanged baseline.

Re-exports from cppmega_v4.parallelism: DistributedMemoryReport,
PerRankMemory, estimate_distributed_memory.

Tests (tests/v4/test_pspec_stage_b.py, 27 cases):
  - PerRankMemory.fits_on_device validates headroom
  - single_device yields num_devices per_rank entries; no duplication
    /master/kernel-boundary terms
  - FSDP2 reduces optimizer_state by dp_degree (specifically ≥4×)
  - FSDP2 adds all_gather_peak term; no-FSDP → 0
  - TP reduces per-rank weights; TP>1 adds kernel_boundary; TP=1 → 0
  - EP reduces MoE routing buffers; no MoE bricks → 0 routing buffers
  - fp8_enabled → duplication_bytes > 0 + diagnostic surfaced;
    bf16-only → 0
  - master_weights_fp32 → overhead > 0 + diagnostic; default → 0
  - Muon optimizer state < AdamW; SGD = 0
  - Inference drops grads/optim; can include kv_cache
  - activation_checkpointing="off" > "full" activations
  - Framework overhead: TPU > CUDA > MLX
  - fits_on_topology validates headroom; summary dict shape
  - worst_rank_idx in range and matches max-total
  - system: every preset under {gb10_quarter, h100_8x, h200_8x}
    produces a finite report

Full v4 regression: 999 passed / 5 skipped / 0 failed.

Author: apstenku123

cppmega_v4/parallelism/__init__.py

@@ -13,6 +13,11 @@
 
 from __future__ import annotations
 
+from cppmega_v4.parallelism.distributed_memory import (
+    DistributedMemoryReport,
+    PerRankMemory,
+    estimate_distributed_memory,
+)
 from cppmega_v4.parallelism.sharding_spec import (
     AxisAssignment,
     ParallelismKind,
@@ -42,11 +47,14 @@
     "DeviceKind",
     "DeviceSpec",
     "DeviceTopology",
+    "DistributedMemoryReport",
     "ParallelismKind",
+    "PerRankMemory",
     "ShardingSpec",
     "TOPOLOGY_BUILTINS",
     "a100_8x",
     "b100_8x",
+    "estimate_distributed_memory",
     "fsdp2_only",
     "fsdp2_plus_tp",
     "gb10_quarter",