Skip to content

← Flax NNX step 57 of 100

Hard Primitives

NNX Mixed-Precision Step

Jump to the task

Why this matters

Mixed precision uses different numerical formats for different parts of a computation. A common choice keeps parameters and updates in float32 while using bfloat16 for selected forward operations. A bfloat16 array uses half the storage bytes of a float32 array of the same shape. Whole-model memory and speed depend on the operations, hardware and additional state retained.

Bfloat16 has roughly float32's exponent range but fewer significant bits. Small values such as 1e-9 are representable. A different issue is adding a small update to a much larger parameter: limited precision can round away that update. Keeping float32 parameters helps preserve such changes.

Float16 has a narrower exponent range and more often needs loss scaling. Bfloat16 training commonly runs without it. This exercise includes an explicit scale to teach the scale-and-unscale procedure, not because every bfloat16 computation needs it.

Scaling the loss

For a fixed positive scale S, exact arithmetic gives:

grad(S * loss) / S = grad(loss)

Finite-precision rounding and overflow can break exact equality. Scaling cannot recover information already lost during the forward pass. A scale that is too large can itself overflow gradients. This exercise uses a fixed scale and has no dynamic overflow detection or skipped-update logic.

A small precision example

import jax.numpy as jnp

small = jnp.asarray(1e-9, dtype=jnp.bfloat16)
assert float(small) > 0
one = jnp.asarray(1.0, dtype=jnp.bfloat16)
step = jnp.asarray(1e-4, dtype=jnp.bfloat16)
assert float(one + step) == 1.0
assert float(jnp.float32(1.0) + jnp.float32(1e-4)) > 1.0

The small number exists; the problem is resolving an increment near one. Range and precision are different properties.

Problem

Implement mixed_precision_step(seed, x, y, lr, loss_scale):

  1. Initialize nnx.Linear(x.shape[-1], y.shape[-1], rngs=nnx.Rngs(int(seed))) with default initializers. Create an nnx.Optimizer with optax.sgd(lr) and wrt=nnx.Param. Retain the float32 parameters on this model.
  2. Inside the loss function, split the model with gdef, state = nnx.split(model, nnx.Param). Cast the array leaves to bfloat16 and merge that state into a temporary model.
  3. Cast x and y to bfloat16, run the temporary model, and compute jnp.mean((prediction - y_bf16) ** 2).
  4. Cast that loss to float32 and multiply by loss_scale. Differentiate with nnx.value_and_grad with respect to the original model.
  5. Cast the gradient leaves to float32, divide by loss_scale, and call optimizer.update(model, grads) exactly once.
  6. Using the updated parameters, repeat the temporary bfloat16 forward and MSE calculation without scaling. Cast to float32 and reshape to (1,).

Inputs x and y are float32 arrays of shapes (N, D_in) and (N, D_out); all dimensions are positive. The seed is nonnegative and integer-valued, lr is finite and nonnegative, and loss_scale is finite and positive. Return a float32 array containing the unscaled loss after one step.

The temporary state transformation does not mutate the original parameters. That is a property of building new arrays with tree_map here, not a general claim that NNX state is frozen or that every state manipulation is harmless. This simple Linear model has only parameter state; a more complex model may require preserving other collections when splitting and merging.

Common pitfalls

  • Differentiate through the bfloat16 casts while retaining the original float32 parameters. Returned parameter gradients have the original parameters' dtype; a float32 loss alone does not determine their dtype.
  • Unscale before the optimizer update. Otherwise, the scale also changes the effective update size, subject to rounding.
  • Casting y affects the subtraction and loss, not the preceding matrix multiplication. Inputs, parameters and layer dtype settings determine the forward computation's dtypes.
  • This exercise deliberately computes an MSE with bfloat16 inputs and a bfloat16 result. Real training often computes sensitive losses in float32. Casting the already rounded result afterward does not restore precision.
  • An operation's output dtype does not specify every internal accumulator's precision. Reduction and matrix-multiplication implementations vary.
  • Report the new, unscaled loss. A fixed SGD step is not guaranteed to lower it for every input or learning rate.
  • Keep the result as a JAX array. Python float conversions obstruct tracing. Setup arguments converted with Python int/float must be static under JIT. Compiled fusion can also change bfloat16 rounding, so bit-for-bit eager/JIT agreement is not a general guarantee.