jmp_mixed_precision_policy_with_loss_scaling_jax.py

1import jax
2import jax.numpy as jnp
3import jmp
4
5# Define a mixed-precision policy.
6# In this example: compute in float16, store/calculate parameters in float32.
7policy = jmp.get_policy("float32=float16")
8
9# Sample data and parameters
10x = jnp.ones((4, 4), dtype=jnp.float32)
11w = jnp.ones((4, 4), dtype=jnp.float32)
12
13# Use the policy to cast inputs to the compute dtype
14x_compute, w_compute = policy.cast_to_compute((x, w))
15
16# Perform computation in the compute dtype (float16)
17y_compute = jnp.matmul(x_compute, w_compute)
18
19# Cast the result back to the output dtype (float32)
20y = policy.cast_to_output(y_compute)
21
22# Example of Loss Scaling for stability
23loss_scale = jmp.StaticLossScale(2**15)
24loss = jnp.array(1.0, dtype=jnp.float32)
25
26# Scale the loss before computing gradients
27scaled_loss = loss_scale.scale(loss)
28
29# After computing gradients, unscale them
30grads = jnp.ones_like(w)  # Dummy gradient
31unscaled_grads = loss_scale.unscale(grads)
32
33print(f"Compute dtype: {y_compute.dtype}")
34print(f"Output dtype: {y.dtype}")