det-train / README.md
phanerozoic's picture
Card: standardized form with hero
6afe4b2 verified
|
Raw
History Blame
5.42 kB
metadata
library_name: kernels
license: apache-2.0

det-train

Bitwise-deterministic training arithmetic at any parallelism, loadable through kernels. The contract: outputs are fixed functions of the inputs, independent of tiling, K-sharding, world size, and repetition, so a loss curve is torch.equal-identical whether the same global computation runs on one GPU or sixty-four. The reference baseline is an fp64 matmul, against which results sit at the fp32 rounding floor.

Ordinary GPU training is not reproducible: floating-point addition is not associative, so changing the tiling, the shard boundaries, or the world size changes the bits, and two runs of the same experiment diverge. That makes bisection, audit, and silent-data-corruption detection guesswork. This kernel removes the effect at the arithmetic level: every GEMM accumulates exactly, so any decomposition of the computation produces identical bytes, and a training run becomes a reproducible object.

The loss-difference channel: float32's gap between an unsharded and a sharded run jitters near 1e-8 while det-train's gap is a flat zero

The same 160-step training run computed unsharded and 4-way K-sharded. In float32 the per-step loss difference wanders around 10^-8 and the final weights differ by 3e-7 relative; under det-train the difference is exactly zero at every step and the weights are identical bytes.

Usage

import torch
from kernels import get_kernel

dt = get_kernel("phanerozoic/det-train", version=1, trust_remote_code=True)

layer = dt.DetLinear(4096, 4096, device="cuda")
y = layer(x); y.square().mean().backward()      # deterministic fwd and bwd

# tensor parallelism over K: shards compose bit-identically
a1 = dt.det_gemm_partial(x[:, :2048], w[:, :2048])
a2 = dt.det_gemm_partial(x[:, 2048:], w[:, 2048:])
y_sharded = dt.det_finalize(dt.det_combine([a1, a2]))
assert torch.equal(y_sharded, dt.det_gemm(x, w))

# data/tensor parallel: sum long-accumulator digits across ranks, exactly
acc = dt.det_all_reduce(dt.det_gemm_partial(x_shard, w_shard))
y = dt.det_finalize(acc)                        # identical at any world size

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark. K-shard boundaries are arbitrary; any partition composes exactly.

API

Symbol Purpose
det_gemm(x [M,K], w [N,K]) deterministic x @ w^T, f32 out
det_gemm_partial(x, w) long-accumulator digits [M,N,NB] int64 for composition
det_combine([...]) / det_finalize(acc) exact digit-array sum / compose to f32
det_all_reduce(acc) world-size-invariant reduction (one int64 all-reduce)
det_sum(x, dim) permutation-invariant reduction
DetLinear linear module, deterministic forward and backward

Method

Each GEMM output owns a fixed array of signed int64 digits in radix 2^32 spanning the full fp32 product exponent range (a Kulisch long accumulator). Every product is formed exactly in fp64 (bf16 and f32 significands make every product exact) and its base-2^32 limbs are added into the digits. Accumulation is exact integer addition into fixed bins, associative by construction, so any fold tree over any K-decomposition is bit-identical. Sharded partials compose by adding digit arrays; distributed reduction is a single int64 all-reduce of digit arrays. There is no data-dependent window and no floor loss, so the result is not merely run-to-run repeatable but invariant to how the computation is split. Results carry the exact fp64-product sum rounded once to fp32.

Measured

The accumulation runs on fp64 scalar cores, not tensor cores; the property is bought with throughput:

M x K x N det_gemm cuBLAS f32
256 x 1024 x 256 1.6 ms 2.3 ms
1024 x 4096 x 1024 90.3 ms 0.28 ms

Small shapes hide inside launch overhead; large shapes pay the scalar-core price. The intended regimes are debugging runs, bisection across cluster topologies, evaluation training, silent-data-corruption detection, and audits.

Correctness

  • Sharding and grouping invariance: det_gemm over full K equals the det_combine/det_finalize composition of 2-, 4-, 8-, and 16-way K-shards, and any fold grouping of the partials, bitwise, for bf16 and f32 (L4 sm_89 and RTX PRO 6000 sm_120).
  • Real multi-GPU: a K-sharded GEMM reduced across 4 GPUs via one NCCL int64 all-reduce of the long-accumulator digits is torch.equal to the single-process result, and every rank at world size 1 and 4 emits the identical result digest.
  • Accuracy: results carry the exact fp64-product sum rounded once to fp32; maximum relative error against an fp64 matmul reference is at the fp32 rounding floor.
  • Determinism: repeated runs, and a training run whose forward is computed by sharded composition, are bitwise identical.

Requirements and limits

  • NVIDIA GPU with compute capability 8.0+.
  • K <= 2^19 per GEMM call; bf16 and f32 inputs.
  • Throughput is the price of the invariance; route production training elsewhere and use this where the property is the point.

References

Kulisch, "Computer Arithmetic and Validity" (long-accumulator exact summation); reproducible training under tensor and sequence parallelism.

License

Apache-2.0.