Package: jaxamp
jaxamp: automatic mixed precision in JAXTable of ContentsInstallationUsageLicenseInstallationpip install jaxampUsageTL;DR: Likepytorch amp, but for JAX.Replaceloss_fn(model, minibatch)withjaxamp.amp(loss_fn)(model, minibatch)to run with withmixed precision. Usescaler_state = jaxamp.DynamicScalerState()andjaxamp.dynamic_scale_gradorjaxamp.dynamic_scale_value_and_gradto apply a dynamic loss scaler:defloss(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 detailsYour usual training loop might look like this:defloss_fn(model_state,minibatch):...returnloss,accuracydeftrain_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)returnmodel_state,opt_state,loss,accuracydeftrain_loop(model_state,opt_state,optimizer,dataloader):train_step_jit=jax.jit(train_step,static_argnums=3)forminibatchindataloader:model_state,opt_state,loss,accuracy=train_step_jit(model_state,opt_state,minibatch,optimizer)log_metrics(loss,accuracy)returnmodel_state,opt_stateNow, you can replace this with:deftrain_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)returnmodel_state,opt_state,dynamic_scaler_state,loss,accuracydeftrain_loop(model_state,opt_state,optimizer,dataloader):train_step_jit=jax.jit(train_step,static_argnums=3)dynamic_scaler_state=amp.DynamicScalerState()forminibatchindataloader:model_state,opt_state,dynamic_scaler_state,loss,accuracy=train_step_jit(model_state,opt_state,minibatch,optimizer)log_metrics(loss,accuracy)returnmodel_state,opt_stateIt should now be faster!More details onampTheampfunction transforms an arbitrary function into one in which some operations are performed in low precision. This precision can be controlled via thecompute_dtypekeyword-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 theamp_policykeyword-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:defprecision_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 argumentsto the op we are about to transform were going  to be.invars: the input arrays to this operation (note that these dtypes may not matchoriginal_dtypes because of previous casting we might have performed).bind_params: the "meta" parameters to the op (things like axis specifications).returnsnew_invars, new_bind_params: the transformed values for invars and bind_params.'''For example, the function used to cast tocompute_dtypeis:defuse_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"inbind_params:bind_params["preferred_element_type"]=compute_dtypereturninvars,bind_paramsampwill walk through all the ops in your function and look up each op in youramp_policydict. 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 inamp_policy. In this case, if the current operation
is executed inside a scope declared withjax.named_scope, we will apply the  specified transformation function. If two or more active scopes match policies inamp_policytheoutermostscope 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 AMPYou can disableampfor a specific function (or area of code) using the context/decoratorjaxamp.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 zMore details on dynamic loss scalersWe supply a loss scaling operation viaDynamicScalerStateand corresponding functionsdynamic_scale_gradanddynamic_scale_value_and_grad.DynamicScalerStatehas the following structure:classDynamicScalerState(NamedTuple):patience:jax.Array=jnp.array(2000)# number of non-inf/NaN iterations to wait before increasing the scaleradjust_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 valuecount:jax.Array=jnp.array(0)# number of non-inf/NaN iterations since the scaler was last increased.The gradient functions then have behavior like:defdynamic_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 differentiatehas_aux: same meaning as in jax.gradredo_on_nan: if the output is nan, we will decrease the scalerand recompute this many times. If the output remains nan, give upand return it.filter: if True, differentiate with equinox.filter_value_and_grad, otherwise use jax.value_and_gradReturns:grad_fn: a function that behaves like the output of jax.value_and_grad except:1. has an extra required keyword argument dynamic_scaler_state2. 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 tipsWhen usingoptax, you may want to wrap your optimizers inoptax.apply_if_finiteto automatically skip NaN gradients. Alternatively, you could use theredo_on_nanoption.Licensejaxampis distributed under the terms of theApache 2.0license.