determinism.py

"""
CAN THE FOLD BE MADE BIT-REPRODUCIBLE, AND WHAT DOES IT COST?

Two runs of the same experiment, same seed, same configuration, produced
models that differed — by about 0.002 in final accuracy, diverging from
epoch 10 onward. The cause is the scatter-add: folding a gradient sums many
weights into one stored value using hardware atomics, atomics complete in
whatever order the scheduler picks, and floating-point addition is not
associative. The same numbers added in a different order give a different
answer.

That is a 3e-6 relative perturbation amplified by chaotic dynamics into
0.002 of accuracy. It is too small to be useful as noise — the minibatch
gradient noise already present is thirty thousand times larger — and too
large for a published digest to survive, so it sits in the gap where it is
neither an asset nor harmless.

It is also avoidable, and the fold is the reason: THE PARTITION NEVER
CHANGES DURING TRAINING. The accumulation order can therefore be fixed once
and reused for every step of every epoch. This script builds that fixed
order and answers three questions:

  1. Is the current scatter actually non-deterministic on this hardware?
  2. Is the fixed-order scatter bit-identical across repeats?
  3. What does it cost per call?

and then trains a small folded model twice each way, on random data, to see
whether determinism survives a few hundred optimiser steps.

No dataset, no download. Random inputs are enough, because the question is
bit-exactness and not accuracy.
"""

import numpy as np
import time

try:
    import cupy as _cp
    _GPU = _cp.cuda.runtime.getDeviceCount() > 0
except Exception:
    _GPU = False
xp = _cp if _GPU else np
DT = np.float32


def to_dev(a, dtype=DT):
    a = np.asarray(a, dtype=dtype)
    return xp.asarray(a) if _GPU else a


def to_host(a):
    return _cp.asnumpy(a) if _GPU and isinstance(a, _cp.ndarray) else np.asarray(a)


def scatter_atomic(g, idx, K):
    """What every script in this programme uses. On a GPU this is an atomic
    accumulation, so the summation order is whatever the hardware picks."""
    out = xp.zeros(K, DT)
    if _GPU:
        import cupyx
        cupyx.scatter_add(out, idx, g.reshape(-1))
    else:
        np.add.at(out, idx, g.reshape(-1))
    return out


class FixedScatter:
    """The same sum, in an order fixed once and reused forever.

    Sort the index array, then lay the gradients into a rectangular buffer
    of shape (K, longest group) by ASSIGNMENT rather than accumulation --
    every destination is written by exactly one source, so there is no race
    at all -- and reduce along the second axis. A tree reduction over a
    fixed layout gives the same answer every time.

    Building this costs one sort. The partition never changes, so it is
    paid once for the whole run rather than once per step."""

    def __init__(self, idx, K):
        h = to_host(idx).astype(np.int64)
        order = np.argsort(h, kind="stable")
        counts = np.bincount(h, minlength=K)
        starts = np.cumsum(counts) - counts
        pos = np.arange(len(h)) - np.repeat(starts, counts)
        self.K, self.width = K, int(counts.max())
        self.order = to_dev(order, np.int64) if _GPU else order
        slot = h[order]*self.width + pos
        self.slot = to_dev(slot, np.int64) if _GPU else slot
        self.buf = xp.zeros(K*self.width, DT)
        self.bytes = K*self.width*4

    def __call__(self, g):
        self.buf[:] = 0
        self.buf[self.slot] = g.reshape(-1)[self.order]
        return self.buf.reshape(self.K, self.width).sum(1)


def bits_equal(a, b):
    return np.array_equal(to_host(a).view(np.uint8), to_host(b).view(np.uint8))


CFG = dict(D=768, hid=4096, K=433, reps=5, steps=200, batch=128, lr=1e-3)


def main(**over):
    CFG.update(over)
    D, hid, K = CFG["D"], CFG["hid"], CFG["K"]
    print("=" * 78)
    print("CAN THE FOLD BE MADE BIT-REPRODUCIBLE, AND WHAT DOES IT COST?")
    print("=" * 78)
    print(f"  backend: {'cupy (GPU)' if _GPU else 'numpy (CPU)'}")
    for k, v in CFG.items():
        print(f"  {k:7s} = {v}")
    if not _GPU:
        print()
        print("  NOTE: on CPU numpy's add.at is already sequential and")
        print("  therefore deterministic, so part 1 will show no difference.")
        print("  The non-determinism this tests is a GPU atomics effect and")
        print("  the timings here do not predict GPU timings either.")
    print("=" * 78, flush=True)

    rng = np.random.default_rng(0)
    idx_h = rng.integers(0, K, D*hid)
    idx = to_dev(idx_h, np.int64) if _GPU else idx_h
    g = to_dev(rng.normal(size=D*hid))

    print("\n  1. IS THE CURRENT SCATTER NON-DETERMINISTIC HERE?\n")
    a0 = scatter_atomic(g, idx, K)
    diffs = []
    for _ in range(CFG["reps"]):
        a = scatter_atomic(g, idx, K)
        diffs.append(0 if bits_equal(a0, a)
                     else float(np.abs(to_host(a)-to_host(a0)).max()))
    nd = sum(1 for d in diffs if d)
    print(f"     {nd} of {CFG['reps']} repeats differed from the first")
    if nd:
        print(f"     largest difference {max(diffs):.3e}")
        print(f"     — the same additions, a different order, a different sum")
    else:
        print(f"     all identical (expected on CPU; see the note above)")

    print("\n  2. IS THE FIXED ORDER BIT-IDENTICAL?\n")
    t0 = time.time()
    fs = FixedScatter(idx, K)
    prep = time.time()-t0
    b0 = fs(g)
    same = all(bits_equal(b0, fs(g)) for _ in range(CFG["reps"]))
    print(f"     preparation {prep*1000:.0f} ms, ONCE for the whole run")
    print(f"     buffer {fs.K} x {fs.width} = {fs.bytes/1e6:.1f} MB")
    print(f"     bit-identical across {CFG['reps']} repeats: {same}")
    print(f"     agrees with the atomic result to "
          f"{float(np.abs(to_host(b0)-to_host(a0)).max()):.3e}")

    print("\n  3. WHAT DOES IT COST?\n")
    def timeit(f, n=20):
        f()
        if _GPU:
            _cp.cuda.Stream.null.synchronize()
        t = time.time()
        for _ in range(n):
            f()
        if _GPU:
            _cp.cuda.Stream.null.synchronize()
        return (time.time()-t)/n*1000
    ta = timeit(lambda: scatter_atomic(g, idx, K))
    tf = timeit(lambda: fs(g))
    # the matmul the scatter sits beside, for scale
    x = to_dev(rng.normal(size=(CFG["batch"], D)))
    W = to_dev(rng.normal(size=(D, hid)))
    tm = timeit(lambda: x @ W)
    print(f"     {'atomic scatter':>22s} {ta:8.3f} ms")
    print(f"     {'fixed-order scatter':>22s} {tf:8.3f} ms   "
          f"{tf/ta:.2f}x")
    print(f"     {'one forward matmul':>22s} {tm:8.3f} ms   (for scale)")
    over_step = (tf-ta)/max(tm, 1e-9)
    print(f"\n     the extra cost is {tf-ta:+.3f} ms a step, which is "
          f"{over_step:+.2f}")
    print(f"     of a forward matmul")

    print("\n  4. DOES A TRAINING RUN BECOME REPRODUCIBLE?\n")
    def run(fixed, seed=0):
        r = np.random.default_rng(seed)
        v = to_dev(r.normal(0, 0.05, K))
        m = xp.zeros_like(v); vv = xp.zeros_like(v)
        X = to_dev(r.normal(size=(CFG["batch"], D)))
        s = FixedScatter(idx, K) if fixed else None
        for t in range(1, CFG["steps"]+1):
            W = v[idx].reshape(D, hid)
            h = xp.maximum(X @ W, 0)
            d = (h - 1.0)/CFG["batch"]
            gW = X.T @ (d*(h > 0))
            gr = s(gW) if fixed else scatter_atomic(gW, idx, K)
            m = 0.9*m + 0.1*gr
            vv = 0.999*vv + 0.001*gr*gr
            v = v - CFG["lr"]*(m/(1-0.9**t))/(xp.sqrt(vv/(1-0.999**t))+1e-8)
        return v
    for lab, fixed in (("atomic", False), ("fixed order", True)):
        r1, r2 = run(fixed), run(fixed)
        eq = bits_equal(r1, r2)
        dd = float(np.abs(to_host(r1)-to_host(r2)).max())
        print(f"     {lab:>12s}: two runs bit-identical after "
              f"{CFG['steps']} steps: {eq}"
              + ("" if eq else f"   (max difference {dd:.3e})"))

    print("\n" + "=" * 78)
    print("  READOUT")
    print("=" * 78)
    if not _GPU:
        print("  CPU ONLY, so this run cannot answer the question it was")
        print("  written for. Rerun where cupy sees a device.")
    elif nd and same:
        print(f"  DETERMINISM IS AVAILABLE, at {tf/ta:.2f}x the scatter's own")
        print(f"  cost and {abs(over_step)*100:.0f}% of a forward matmul a step.")
        print(f"  The current scatter is genuinely non-deterministic here and")
        print(f"  the fixed order is exact. Whether it is worth paying for")
        print(f"  depends on whether anything downstream needs a base to have")
        print(f"  a stable digest.")
    elif not nd:
        print("  THE CURRENT SCATTER IS ALREADY DETERMINISTIC on this")
        print("  hardware, so there is nothing to fix here — though the 0.002")
        print("  drift measured across sessions says something else varies,")
        print("  and it would be worth finding what.")
    else:
        print("  THE FIXED ORDER IS NOT REPRODUCING EXACTLY, which should not")
        print("  happen — the reduction is over a fixed layout. Read the")
        print("  numbers above before trusting either path.")


if __name__ == "__main__":
    main()