The virtual GPU and the byte counter

"""A small MLP with an explicit flat parameter vector.

Keeping every parameter in one flat float32 array is what makes sharding and
byte accounting simple: a shard is a slice, not a tree walk.
"""
from __future__ import annotations
import numpy as np


class MLP:
    def __init__(self, dims=(64, 128, 128, 64), seed=0):
        self.dims = dims
        rng = np.random.default_rng(seed)
        self.shapes, self.slices, n = [], [], 0
        for a, b in zip(dims[:-1], dims[1:]):
            for shape in ((a, b), (b,)):
                size = int(np.prod(shape))
                self.shapes.append(shape)
                self.slices.append(slice(n, n + size))
                n += size
        self.n_params = n
        theta = np.zeros(n, dtype=np.float32)
        for shape, sl in zip(self.shapes, self.slices):
            if len(shape) == 2:
                theta[sl] = (rng.standard_normal(sl.stop - sl.start)
                             * np.sqrt(2.0 / shape[0])).astype(np.float32)
        self.theta0 = theta

    def unpack(self, theta):
        return [theta[sl].reshape(s) for s, sl in zip(self.shapes, self.slices)]

    def loss_and_grad(self, theta, x, y):
        """Forward and backward for MSE on a plain ReLU MLP."""
        ps = self.unpack(theta)
        acts, h = [x], x
        pre = []
        for i in range(0, len(ps), 2):
            W, b = ps[i], ps[i + 1]
            z = h @ W + b
            pre.append(z)
            h = np.maximum(z, 0) if i + 2 < len(ps) else z
            acts.append(h)
        diff = h - y
        loss = float(np.mean(diff ** 2))
        g = (2.0 / (x.shape[0] * h.shape[1])) * diff
        grads = [None] * len(ps)
        for i in range(len(ps) - 2, -1, -2):
            W = ps[i]
            a_in = acts[i // 2]
            grads[i] = (a_in.T @ g).astype(np.float32)
            grads[i + 1] = g.sum(axis=0).astype(np.float32)
            if i > 0:
                g = (g @ W.T) * (pre[i // 2 - 1] > 0)
        flat = np.concatenate([gr.ravel() for gr in grads]).astype(np.float32)
        return loss, flat

th = m.theta0.copy(); eps = 1e-3
for idx in (0, 500, 5000, 12000):
    p = th.copy(); p[idx] += eps; lp, _ = m.loss_and_grad(p, x, y)
    p = th.copy(); p[idx] -= eps; lm, _ = m.loss_and_grad(p, x, y)
    print(idx, g[idx], (lp - lm) / (2 * eps))

# idx      0 analytic  0.001464 numeric  0.001431 diff 3.3e-05
# idx    500 analytic -0.005282 numeric -0.005245 diff 3.7e-05
# idx   5000 analytic -0.011015 numeric -0.011086 diff 7.1e-05
# idx  12000 analytic  0.019557 numeric  0.019670 diff 1.1e-04

"""Collectives over a list of per-rank tensors, with every byte counted."""
from __future__ import annotations
from collections import defaultdict
import numpy as np


class Comm:
    """Simulates a communicator over `world` ranks.

    Every collective takes a list of per-rank arrays and returns a list of
    per-rank arrays. Bytes are attributed to the sending rank, using the ring
    cost 2(N-1)/N * bytes for all-reduce and (N-1)/N * bytes for the two halves,
    which is what a real NCCL ring moves.
    """

    def __init__(self, world: int):
        self.world = world
        self.sent = defaultdict(float)     # rank -> bytes sent
        self.by_op = defaultdict(float)    # op name -> bytes sent, all ranks

    # ---- accounting -------------------------------------------------
    def _charge(self, op: str, per_rank_bytes: float):
        for r in range(self.world):
            self.sent[r] += per_rank_bytes
        self.by_op[op] += per_rank_bytes * self.world

    def reset(self):
        self.sent.clear()
        self.by_op.clear()

    def bytes_per_rank(self) -> float:
        return max(self.sent.values()) if self.sent else 0.0

    # ---- collectives ------------------------------------------------
    def all_reduce_mean(self, shards):
        """Every rank in, every rank out with the mean. Ring cost 2(N-1)/N."""
        n = self.world
        total = sum(shards) / n
        self._charge("all_reduce", 2 * (n - 1) / n * shards[0].nbytes)
        return [total.copy() for _ in range(n)]

    def reduce_scatter_mean(self, shards):
        """Every rank in, rank i out with the mean of slice i. Cost (N-1)/N."""
        n = self.world
        total = sum(shards) / n
        parts = np.array_split(total, n)
        self._charge("reduce_scatter", (n - 1) / n * shards[0].nbytes)
        return [p.copy() for p in parts]

    def all_gather(self, parts):
        """Rank i in with slice i, every rank out with the whole. Cost (N-1)/N."""
        n = self.world
        full = np.concatenate(parts)
        self._charge("all_gather", (n - 1) / n * full.nbytes)
        return [full.copy() for _ in range(n)]

import numpy as np
from zerosim.comm import Comm

def test_reduce_scatter_then_all_gather_equals_all_reduce():
    N = 4
    shards = [np.full(8, i + 1, dtype=np.float32) for i in range(N)]

    a = Comm(N); ar = a.all_reduce_mean(shards)[0]
    b = Comm(N); rs = b.reduce_scatter_mean(shards); ag = b.all_gather(rs)[0]

    assert np.array_equal(ar, ag)                       # same result
    assert b.bytes_per_rank() == a.bytes_per_rank()     # same cost