| --- |
| 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 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 |
|
|
| ```python |
| 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. |
|
|
