GitLocker: The Coding Marketplace

Description:

jaxamp 0.0.4

jaxamp: automatic mixed precision in JAX

Table of Contents

Installation
Usage
License

Installation
pip install jaxamp

Usage
TL;DR: Like pytorch amp, but for JAX.
Replace loss_fn(model, minibatch) with jaxamp.amp(loss_fn)(model, minibatch) to run with with mixed precision. Use scaler_state = jaxamp.DynamicScalerState() and jaxamp.dynamic_scale_grad or jaxamp.dynamic_scale_value_and_grad to apply a dynamic loss scaler:
def loss(model, minibatch):
...

scaler_state= jaxamp.DynamicLossScaler()
amp_loss = jaxamp.amp(loss)
grad_fn = jaxamp.dynamic_scale_grad(amp_loss)
scaler_state, grad = grad_fn(model, minibatch, dynamic_scaler_state=scaler_state)

More details
Your usual training loop might look like this:
def loss_fn(model_state, minibatch):
...
return loss, accuracy

def train_step(model_state, opt_state, minibatch, optimizer):

value_and_grad_fn = jax.value_and_grad(loss_fn, has_aux=True)

(loss, accuracy), grads = value_and_grad_fn(model_state, minibatch)

updates, opt_state = optimizer.update(grads, opt_state, model_state)
model_state = optax.apply_updates(model_state, updates)
return model_state, opt_state, loss, accuracy

def train_loop(model_state, opt_state, optimizer, dataloader):
train_step_jit = jax.jit(train_step, static_argnums=3)

for minibatch in dataloader:
model_state, opt_state, loss, accuracy = train_step_jit(model_state, opt_state, minibatch, optimizer)
log_metrics(loss, accuracy)
return model_state, opt_state

Now, you can replace this with:
def train_step(
model_state,
opt_state,
minibatch,
dynamic_scaler_state,
optimizer):
amp_loss_fn = jaxamp.amp(loss_fn)

value_and_grad_fn = jaxamp.dynamic_scale_value_and_grad(amp_loss_fn, has_aux=True)

dynamic_scaler_state, ((loss, accuracy), grads) = value_and_grad_fn(
model_state,
minibatch,
dynamic_scaler_state=dynamic_scaler_state)

updates, opt_state = optimizer.update(grads, opt_state, model_state)
model_state = optax.apply_updates(model_state, updates)
return model_state, opt_state, dynamic_scaler_state, loss, accuracy

def train_loop(model_state, opt_state, optimizer, dataloader):
train_step_jit = jax.jit(train_step, static_argnums=3)
dynamic_scaler_state = amp.DynamicScalerState()
for minibatch in dataloader:
model_state, opt_state, dynamic_scaler_state, loss, accuracy = train_step_jit(
model_state,
opt_state,
minibatch,
optimizer)
log_metrics(loss, accuracy)
return model_state, opt_state

It should now be faster!
More details on amp
The amp function transforms an arbitrary function into one in which some operations are performed in low precision. This precision can be controlled via the compute_dtype keyword-only argument:
amp_loss_fn = amp(loss_fn, compute_dtype=jnp.float16). You can also control which operations are performed in low precision (and how) via the amp_policy keyword-only argument. This argument should
take a dictionary whose keys must be either strings or jax primitives (e.g. jax.lax.add_p). The values are functions that will be called to cast arrays into relevant dtypes. These functions should have signature:
def precision_fn(
compute_dtype: Type,
original_dtypes: Sequence[Type],
*invars: Sequence[Array],
*bind_params: Dict[str, Any]) -> Sequence[Array], Dict[str, Any]:
'''
Args:
compute_dtype: this is the compute_dtype provided to `amp`.
original_dtypes: these are the dtypes that original user code expected the arguments
to the op we are about to transform were going to be.
invars: the input arrays to this operation (note that these dtypes may not match
original_dtypes because of previous casting we might have performed).
bind_params: the "meta" parameters to the op (things like axis specifications).
returns
new_invars, new_bind_params: the transformed values for invars and bind_params.
'''

For example, the function used to cast to compute_dtype is:
def use_compute_precision(
compute_dtype: Type,
original_dtypes: Sequence[Type],
*invars: Sequence[Any],
**bind_params: Dict[str, Any]
) -> (Sequence[Any], Dict[str, Any]):
invars = cast_tree(compute_dtype, invars)
bind_params = cast_tree(compute_dtype, bind_params)
bind_params = dict(bind_params)
if "preferred_element_type" in bind_params:
bind_params["preferred_element_type"] = compute_dtype
return invars, bind_params

amp will walk through all the ops in your function and look up each op in your amp_policy dict. If the op is present, it will apply the specified function
Otherwise it will cast the inputs to their original values and apply the op unchanged. You can also provide string keys in amp_policy. In this case, if the current operation
is executed inside a scope declared with jax.named_scope, we will apply the specified transformation function. If two or more active scopes match policies in amp_policy the outermost scope is used. There are two special scopes "amp_step" and "amp_default".
By default these both stop any automatic mixed precision from happening inside them.
Selectively Disabling AMP
You can disable amp for a specific function (or area of code) using the context/decorator jaxamp.amp_stop:
@jaxamp.amp_stop
def high_precision_matmul(W, x):
return jnp.dot(W, x)

def high_precision_with_context(W, x):
# will be low precision
y = jnp.dot(W, x)

with jaxamp.amp_stop():
# in fp32 precision
z = jnp.dot(W, y)
return z

More details on dynamic loss scalers
We supply a loss scaling operation via DynamicScalerState and corresponding functions dynamic_scale_grad and dynamic_scale_value_and_grad.
DynamicScalerState has the following structure:
class DynamicScalerState(NamedTuple):
patience: jax.Array = jnp.array(2000) # number of non-inf/NaN iterations to wait before increasing the scaler
adjust_factor: jax.Array = jnp.array(2.0) # When increasing or decreasing the scaler, multiply or divide by this factor.
scaler: jax.Array = jnp.array(2**15, dtype=jnp.float32) # current scaler value
count: jax.Array = jnp.array(0) # number of non-inf/NaN iterations since the scaler was last increased.

The gradient functions then have behavior like:
def dynamic_scale_value_and_grad(
fun: Callable,
*,
has_aux: bool = False,
redo_on_nan: bool = 0,
filter=True,
**kwargs
):
'''
apply dynamic scalar to the value_and_grad function.

Args:
fun: function to differentiate
has_aux: same meaning as in jax.grad
redo_on_nan: if the output is nan, we will decrease the scaler
and recompute this many times. If the output remains nan, give up
and return it.
filter: if True, differentiate with equinox.filter_value_and_grad, otherwise use jax.value_and_grad

Returns:
grad_fn: a function that behaves like the output of jax.value_and_grad except:
1. has an extra required keyword argument dynamic_scaler_state
2. the return value is now a tuple (next_dynamic_scaler_state, (value, grads))
of (next_dynamic_scaler_state, ((value, aux), grads)) if has_aux=True
'''

More usage tips
When using optax, you may want to wrap your optimizers in optax.apply_if_finite to automatically skip NaN gradients. Alternatively, you could use the redo_on_nan option.
License
jaxamp is distributed under the terms of the Apache 2.0 license.