Skip to content

Stochastic Optimizers

About

The d9d.optim.stochastic module provides optimizers for training with bf16 parameters and no fp32 master copy. They round the updated parameters to bf16 with stochastic rounding. The update and the rounding run in one fused Triton kernel. The kernels live in d9d.kernel.stochastic and can also be called directly.

Stochastic Rounding

A standard cast such as tensor.to(torch.bfloat16) rounds to the nearest representable value. In bf16 training this can stall learning. If an update is smaller than half the gap between two bf16 values, rounding removes it completely.

Stochastic rounding rounds up or down at random. For example, if a value \(x\) is \(30\%\) of the way from representable value \(A\) to representable value \(B\), it rounds to \(B\) with \(30\%\) probability and to \(A\) with \(70\%\) probability. The expected result equals the exact value, \(E[Round(x)] = x\). Small updates therefore add up over many steps instead of being lost.

For details, see:

Limits

  • The kernels need a GPU that runs Triton.
  • StochasticAdamW parameters must be bf16. Gradients can be bf16 or fp32. Optimizer states can be fp32 or bf16 (state_dtype).
  • StochasticAdamW does not support closures or sparse gradients.

Usage

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
import torch

from d9d.optim.stochastic import StochasticAdamW

model: torch.nn.Module = ...
model = model.to(torch.bfloat16)

optimizer = StochasticAdamW(
    model.parameters(),
    lr=1e-4,
    weight_decay=0.01,
    state_dtype=torch.bfloat16,
)

StochasticAdamW saves the state of its random number generator in state_dict(). A resumed run gets the same rounding noise.

Benchmarks

The plots show throughput in GB/s against the number of elements. The baselines are a PyTorch eager and a torch.compile implementation of the same operation.

copy_fp32_to_bf16_stochastic_ (H100 80GB, fp32 source, bf16 target)

adamw_stochastic_bf16_ (H100 80GB, bf16 parameters, gradients and states)

API Reference

d9d.optim.stochastic

Optimizers that update bf16 parameters with stochastic rounding.

StochasticAdamW

Bases: Optimizer

AdamW optimizer for bf16 parameters with stochastic rounding.

Parameters must be bf16. Gradients can be bf16 or fp32. Parameters can be DTensors.

The optimizer owns its random number generator and saves the generator state in state_dict. A resumed run therefore gets the same rounding noise.

__init__(params, lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.01, generator=None, state_dtype=torch.float32)

Constructs the StochasticAdamW object.

Parameters:

Name Type Description Default
params ParamsT

Parameters to optimize, or dicts that define parameter groups.

 required
lr float

Learning rate.

 required
betas tuple[float, float]

Decay rates of the first and second moments.

 (0.9, 0.999)
eps float

Term added to the denominator for numerical stability.

 1e-08
weight_decay float

Decoupled weight decay coefficient.

 0.01
generator Generator | None

Random number generator for stochastic rounding. If None, a new CPU generator is seeded from the default PyTorch generator.

 None
state_dtype dtype

Dtype of the optimizer states. Can be torch.float32 or torch.bfloat16.

 float32

Raises:

Type Description
ValueError

If lr, eps, betas or weight_decay is out of range.
load_state_dict(state_dict)

Loads the optimizer state, including the generator state if it is present.

state_dict()

Returns the optimizer state, including the generator state.

d9d.kernel.stochastic

Triton kernels that round fp32 values to bf16 stochastically.

adamw_stochastic_bf16_(params, grads, exp_avg, exp_avg_sq, lr, beta1, beta2, eps, weight_decay, step, generator=None)

Performs one in-place AdamW step on bf16 parameters with stochastic rounding.

The update is computed in fp32, and the new parameters are rounded to bf16 stochastically. Gradients and optimizer states can be fp32 or bf16. States stored in bf16 are also rounded stochastically.

Parameters:

Name Type Description Default
params Tensor

Parameters to update in place. Must be bf16 and contiguous.

 required
grads Tensor

Gradients, in bf16 or fp32.

 required
exp_avg Tensor

First moment (moving average of the gradients), updated in place. Must be contiguous.

 required
exp_avg_sq Tensor

Second moment (moving average of the squared gradients), updated in place. Must be contiguous and have the same dtype as exp_avg.

 required
lr float

Learning rate.

 required
beta1 float

Decay rate of the first moment.

 required
beta2 float

Decay rate of the second moment.

 required
eps float

Term added to the denominator for numerical stability.

 required
weight_decay float

Decoupled weight decay coefficient.

 required
step int

Number of the current step, used for bias correction. Must be at least 1.

 required
generator Generator | None

Random number generator that seeds the stochastic rounding. If None, the default PyTorch generator is used.

 None

Raises:

Type Description
ValueError

If grads, exp_avg or exp_avg_sq differs in shape from params, if params is not bf16, if params, exp_avg or exp_avg_sq is not contiguous, or if exp_avg and exp_avg_sq have different dtypes.

copy_fp32_to_bf16_stochastic_(target, source, generator=None)

Copies an fp32 tensor into a bf16 tensor in place with stochastic rounding.

Stochastic rounding rounds each value up or down at random. The probability of rounding up equals the fraction that the cast drops, so the expected value is preserved: E[round(x)] = x. Small updates then accumulate in bf16 instead of being lost.

Parameters:

Name Type Description Default
target Tensor

Output tensor. Must be bf16 and contiguous.

 required
source Tensor

Input tensor. Must be fp32 and have the same shape as target.

 required
generator Generator | None

Random number generator that seeds the rounding noise. If None, the default PyTorch generator is used.

 None

Returns:

Type Description
Tensor

The target tensor.

Raises:

Type Description
ValueError

If target is not contiguous, if the shapes differ, or if source is not fp32 or target is not bf16.