Skip to content

← Flax step 63 of 100

Hard Primitives

Mixed-Precision Training 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. Define a compact Linen module whose call returns nn.Dense(1)(x). Keep the default layer name Dense_0 and default initializers. Initialize with jax.random.PRNGKey(int(seed)) and retain the float32 params.
  2. Create a TrainState with those params, model.apply and optax.sgd(lr).
  3. Inside the loss function, cast a copy of the parameter leaves, x and y to bfloat16. Apply the model and flatten its (N, 1) prediction to (N,).
  4. Compute jnp.mean((prediction - y_bf16) ** 2), cast the loss to float32, and multiply by loss_scale. Differentiate with respect to the original float32 parameters using jax.grad.
  5. Cast gradient leaves to float32, divide by loss_scale, and perform one state.apply_gradients(grads=grads) update.
  6. Recompute the bfloat16 forward and MSE with the updated parameters, without scaling. Cast that loss to float32 and reshape to (1,).

Inputs x and y are float32 arrays of shapes (N, D) and (N,), with positive N and D. 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.

Use a wrapper module containing Dense, as specified. Initializing a standalone Dense instead changes the parameter naming and seeded initialization.

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.