adamw-sr

Fused AdamW with stochastic rounding, loadable through kernels: the decoupled-weight-decay Adam step fused into one kernel per parameter, with unbiased bf16 write-back and deterministic, shard-invariant randomness. The reference baseline is torch.optim.AdamW, tracked to < 1e-4 over 100 fp32 steps.

Training in bf16 without an fp32 master copy fails under ordinary rounding: any update smaller than a bf16 ULP rounds away, so small updates, which is most of them late in training, simply vanish and the weights stall. This kernel rounds stochastically, up or down with probability set by the truncated bits, so the write-back is unbiased and sub-ULP updates accumulate statistically. bf16 weights then train with no master copy, removing 4 bytes per parameter from the optimizer footprint.

6,000 optimizer steps of 1e-4, forty times below one bf16 ULP at this weight scale. Round-to-nearest drops every update and the mean weight never leaves 1.0000; stochastic rounding accumulates them and converges toward the 1.5 target, with repeated runs bitwise identical.

Usage

import torch
from kernels import get_kernel

asr = get_kernel("phanerozoic/adamw-sr", version=1, trust_remote_code=True)

model = MyModel().cuda().to(torch.bfloat16)      # bf16 weights, no fp32 master
opt = asr.AdamWSR(model.parameters(), lr=1e-3, weight_decay=0.01, seed=0)

for batch in loader:
    opt.zero_grad()
    loss(model(batch)).backward()
    opt.step()                                    # stochastic rounding auto-on for bf16

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark. Drop-in for torch.optim.AdamW. Optimizer states are float32; parameters may be f32 or bf16; gradients f32 or bf16.

API

Method

The rounding randomness is counter-based Philox keyed by (seed, step) and indexed by each element's global parameter position, so the step is bitwise reproducible run to run and identical whether a parameter is updated whole or split across shards and ranks. This composes with det-train and det-attn: matmul, attention, and the optimizer step are each reproducible under any parallel decomposition.

Correctness

• fp32 parameters track torch.optim.AdamW over 100 steps to < 1e-4.
• Stochastic rounding is unbiased: repeated bf16 write-back of a sub-ULP value averages to the value; round-to-nearest drops it.
• A bf16-only layer with no fp32 master converges with SR and out-trains round-to-nearest.
• Same (seed, step) gives a bitwise-identical step; a parameter updated in shards with per-shard base_index equals the whole update bitwise.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+.
• Optimizer states are float32 regardless of parameter dtype; the saving is the removed fp32 master copy, not the states.

References

Gupta et al., "Deep Learning with Limited Numerical Precision" (stochastic rounding, 2015); decoupled weight decay (Loshchilov and Hutter, 2019); counter-based Philox RNG.

License

Apache-2.0.