class

GradScaler

GradScaler(init_scale: float = 2.0 ** 16, growth_factor: float = 2.0, backoff_factor: float = 0.5, growth_interval: int = 2000, enabled: bool = True)

source edit

Dynamic loss-scaling helper for mixed-precision training.

Mixed-precision training runs the forward pass in fp16 to halve memory bandwidth and exploit fp16-fast hardware paths, but fp16's narrow dynamic range causes small gradients to underflow to zero — the network stops learning. GradScaler works around this by multiplying the loss by a large constant $s$ before backpropagation:

\tilde{L} = s \cdot L, \qquad \frac{\partial \tilde{L}}{\partial \theta} = s \cdot \frac{\partial L}{\partial \theta}.

The scaled gradients sit comfortably inside fp16's representable range; before the optimizer step they are unscaled by $1/s$ in fp32 so the update is mathematically equivalent to ordinary training.

The scale itself is adapted dynamically. After every step the unscaled gradients are checked for inf / NaN:

Overflow detected — the step is skipped and $s$ is multiplied by backoff_factor (typically 0.5).
No overflow for growth_interval consecutive steps — $s$ is multiplied by growth_factor (typically 2.0).

This produces a sawtooth schedule that tracks the largest scale the current gradient distribution can tolerate.

Parameters

init_scalefloat= 2**16

Initial loss scaling factor applied by scale.

growth_factorfloat= 2.0

Multiplier applied to the scale after growth_interval consecutive non-overflowing steps. Must be > 1.0.

backoff_factorfloat= 0.5

Multiplier applied when an inf / NaN gradient is detected. Must be in (0, 1).

growth_intervalint= 2000

Number of overflow-free steps required before the scale grows.

enabledbool= True

When False the scaler degenerates into a transparent pass-through — scale returns its input unchanged, step calls the optimizer directly, and update is a no-op.

Notes

The canonical training-loop pattern is scale-loss, then step, then update:

scale multiplies the loss by $s$ before backward() so the gradients land safely inside fp16 range.
step unscales the gradients, checks for inf / NaN, and either runs optimizer.step() or skips the update.
update adjusts $s$ according to the growth / backoff schedule for the next iteration.

To work on the real gradients between backward() and the step — gradient clipping is the usual reason — call unscale_ first. The scaler remembers, per optimizer, that its gradients are already unscaled, so the following step does not divide them a second time. Each optimizer moves through ready → unscaled → stepped once per iteration and update returns every optimizer to ready; calling unscale_ twice, unscale_ after step, or step twice before update raises RuntimeError.

With several optimizers on one scaler, an overflow in one optimizer's gradients skips only that optimizer's step; update backs the scale off if any optimizer overflowed.

The parameters themselves stay float32: lucid.amp.autocast moves the forward pass into half precision, not the weights, so the gradients arrive in float32 too. bfloat16 and float64 gradients are unscaled as well, and unscaling never changes a gradient's dtype. A float16 gradient — from a float16 parameter — makes unscale_ and step raise ValueError.

Examples

>>> import lucid
>>> import lucid.nn as nn
>>> import lucid.optim as optim
>>> from lucid.amp import GradScaler, autocast
>>> model = nn.Linear(4, 1).to("metal")
>>> optimizer = optim.SGD(model.parameters(), lr=0.01)
>>> loss_fn = nn.MSELoss()
>>> dataloader = [
...     (lucid.randn(8, 4, device="metal"), lucid.randn(8, 1, device="metal"))
...     for _ in range(3)
... ]
>>> scaler = GradScaler()
>>> for x, y in dataloader:
...     optimizer.zero_grad()
...     with autocast():
...         out = model(x)
...     loss = loss_fn(out.float(), y)
...     scaler.scale(loss).backward()
...     scaler.step(optimizer)
...     scaler.update()
Clipping the unscaled gradients — step sees that
unscale_ already ran and leaves them alone:
>>> from lucid.nn.utils import clip_grad_norm_
>>> x, y = dataloader[0]
>>> optimizer.zero_grad()
>>> scaler.scale(loss_fn(model(x), y)).backward()
>>> scaler.unscale_(optimizer)
>>> norm = clip_grad_norm_(model.parameters(), max_norm=1.0)
>>> scaler.step(optimizer)
>>> scaler.update()

Used by 2

Constructors

dunder

init

→None

__init__(init_scale: float = 2.0 ** 16, growth_factor: float = 2.0, backoff_factor: float = 0.5, growth_interval: int = 2000, enabled: bool = True)

source edit

Initialize the scaler state.

Parameters

init_scalefloat= 2**16

Initial loss scaling factor applied by scale.

growth_factorfloat= 2.0

Multiplier applied to the scale after growth_interval consecutive non-overflowing steps.

backoff_factorfloat= 0.5

Multiplier applied when an inf/NaN gradient is detected.

growth_intervalint= 2000

Number of overflow-free steps required before the scale grows.

enabledbool= True

When False the scaler is a transparent pass-through.

Instance methods

get_scale

→float

get_scale()

source edit

Return the current scale factor.

load_state_dict

→None

load_state_dict(state_dict: dict[str, float])

source edit

Load state from a dict.

scale

→Tensor | list[Tensor]

scale(outputs: Tensor | list[Tensor])

source edit

Multiply outputs by the current scale factor.

Parameters

outputsTensor | list[Tensor]

A Tensor or list of Tensors to scale.

Returns

Tensor | list[Tensor]

Scaled Tensor(s) — same structure as input.

state_dict

→dict[str, float]

state_dict()

source edit

Return serializable state dict.

step

→Tensor or None

step(optimizer: Optimizer, args: object = (), kwargs: object = {})

source edit

Unscale the gradients if needed, then step unless they overflowed.

If unscale_ already ran for optimizer in this iteration its gradients are used as they are; otherwise they are unscaled here. The optimizer's step() is skipped when its gradients held an inf or NaN.

Parameters

optimizerOptimizer

The optimizer to step.

*argsobject= ()

Forwarded to optimizer.step().

**kwargsobject= {}

Forwarded to optimizer.step(). closure is rejected while the scaler is enabled: a closure re-runs backward() and its gradients would never be unscaled.

Returns

Tensor or None

What optimizer.step() returned, or None when the step was skipped.

Raises

RuntimeError

If step already ran for optimizer since the last update, or closure is passed while enabled.

ValueError

If the gradients still need unscaling and one of them is float16 — see unscale_.

NotImplementedError

If the gradients still need unscaling and one of them is complex.

update

→None

update(new_scale: float | None = None)

source edit

Update the scale factor and start the next iteration.

If a scale is provided, it is set directly. Otherwise, the scale is reduced by backoff_factor if any optimizer's gradients overflowed since the last update, or grown by growth_factor after growth_interval overflow-free updates. Either way every optimizer returns to the state where unscale_ and step may be called again.

Parameters

new_scalefloat= None

Explicit new scale value.

In-place ops

unscale_

→None

unscale_(optimizer: Optimizer)

source edit

Divide the optimizer's gradients by the current scale, in place.

Call it between backward() and step when the real gradients are needed — to clip them, or to inspect their norm. The scaler records that this optimizer's gradients are unscaled, so the step that follows does not divide them again.

Parameters

optimizerOptimizer

The optimizer whose parameters' gradients are unscaled.

Raises

RuntimeError

If unscale_ already ran for optimizer since the last update, or step already ran for it.

ValueError

If any of optimizer's gradients is float16. Nothing is unscaled and the optimizer's stage does not change.

NotImplementedError

If any of optimizer's gradients is complex, again before anything is unscaled.

Notes

Every gradient is divided, including those of a parameter whose gradient holds an inf / NaN (those entries stay non-finite); the overflow is recorded for this optimizer and its step is skipped.

Every gradient keeps its dtype. A bfloat16 gradient is cast to float32, divided and finiteness-checked there, then rounded back to bfloat16 once. float32 and float64 gradients are unscaled in their own dtype.

Float16 gradients are refused, as the reference framework refuses them. They come from float16 parameters, and loss scaling assumes float32 master weights: lucid.amp.autocast runs the forward in half precision while the parameters and their gradients stay float32. Dividing a float16 gradient by the default scale 2**16 puts most of its entries below float16's smallest normal value 6.1e-5. Metal flushes those subnormals to zero. Keep the parameters in float32 and let lucid.amp.autocast choose the half-precision ops.

>>> import lucid >>> import lucid.nn as nn >>> import lucid.optim as optim >>> from lucid.amp import GradScaler, autocast >>> model = nn.Linear(4, 1).to("metal") >>> optimizer = optim.SGD(model.parameters(), lr=0.01) >>> loss_fn = nn.MSELoss() >>> dataloader = [ ... (lucid.randn(8, 4, device="metal"), lucid.randn(8, 1, device="metal")) ... for _ in range(3) ... ] >>> scaler = GradScaler() >>> for x, y in dataloader: ... optimizer.zero_grad() ... with autocast(): ... out = model(x) ... loss = loss_fn(out.float(), y) ... scaler.scale(loss).backward() ... scaler.step(optimizer) ... scaler.update() Clipping the unscaled gradients — step sees that unscale_ already ran and leaves them alone: >>> from lucid.nn.utils import clip_grad_norm_ >>> x, y = dataloader[0] >>> optimizer.zero_grad() >>> scaler.scale(loss_fn(model(x), y)).backward() >>> scaler.unscale_(optimizer) >>> norm = clip_grad_norm_(model.parameters(), max_norm=1.0) >>> scaler.step(optimizer) >>> scaler.update()

unscale_

→None

unscale_(optimizer: Optimizer)

source edit

Divide the optimizer's gradients by the current scale, in place.

Parameters

optimizerOptimizer

The optimizer whose parameters' gradients are unscaled.

Raises

RuntimeError

If unscale_ already ran for optimizer since the last update, or step already ran for it.

ValueError

If any of optimizer's gradients is float16. Nothing is unscaled and the optimizer's stage does not change.

NotImplementedError

If any of optimizer's gradients is complex, again before anything is unscaled.

Notes

Every gradient is divided, including those of a parameter whose gradient holds an inf / NaN (those entries stay non-finite); the overflow is recorded for this optimizer and its step is skipped.