Package: jaxline
JAXline - Experiment framework for JAXWhat is JAXlineJAXline is a distributed JAX training and evaluation framework.
It is designed to be forked, covering only the most general aspects of
experiment boilerplate. This ensures that it can serve as an effective starting
point for a wide variety of use cases.Many users will only need to fork theexperiment.pyfile and rely on JAXline for everything else. Other users with more custom
requirements will want to (and are encouraged to) fork other components of
JAXline too, depending on their particular use case.ContentsQuickstartCheckpointingLoggingLaunchingDistribution strategyRandom number handlingDebuggingContributingQuickstartInstallationJAXline is written in pure Python, but depends on C++ code via JAX and
TensorFlow (the latter is used for writing summaries).Because JAX / TensorFlow installation is different depending on your CUDA
version, JAXline does not list JAX or TensorFlow as a dependencies inrequirements.txt.First, follow the instructions to installJAXandTensorFlowrespectively with the relevant accelerator support.Then, install JAXline using pip:$pipinstallgit+https://github.com/deepmind/jaxlineBuilding your own experimentCreate anexperiment.pyfile and inside it define anExperimentclass
that inherits fromexperiment.AbstractExperiment.Implement the methods required byAbstractExperimentin your ownExperimentclass (i.e. theabstractmethods). Optionally override the default implementations ofAbstractExperiment's other methods.Define aconfig, either inexperiment.pyor elsewhere, defining any
settings that you do not wish to inherit frombase_config.
At the very least this will includeconfig.experiment_kwargsto define the
config required by yourExperiment. Make sure thisconfigobject is
included in theflagsaccessible toexperiment.py.Add the following lines to the bottom of yourexperiment.pyto ensure that
yourExperimentobject is correctly passed through toplatform.py:if __name__ == '__main__':
  flags.mark_flag_as_required('config')
  platform.main(Experiment, sys.argv[1:])Run yourexperiment.py.CheckpointingSo far this version of JAXline only supports in-memory checkpointing, as handled
by ourInMemoryCheckpointerIt allows you to save in memory multiple separate checkpoint series in your
train and eval jobs (see below).The user is expected to override theCHECKPOINT_ATTRSandNON_BROADCAST_CHECKPOINT_ATTRSdicts (at least one of these) in order to map checkpointable attributes of their
ownExperimentclass to names they wish them to be stored under in the
checkpoint.CHECKPOINT_ATTRSspecifies jaxDeviceArraysfor which jaxline should only
take the first slice (corresponding to device 0) for checkpointing.NON_BROADCAST_CHECKPOINT_ATTRSspecifies any other picklable object that
jaxline should checkpoint whole.You can specify the frequency with which to save checkpoints, as well as whether
to checkpoint based on step or seconds, by setting thesave_checkpoint_intervalandinterval_typeconfig flagshere.config.max_checkpoints_to_keepcan be used to specify the maximum number of
checkpoints to keep. By default this is set to 5.By settingconfig.best_model_eval_metric, you can specify which value in thescalarsdictionary returned by yourevaluatefunction to use as a 'fitness score'. JAXline will then save a separate series
of checkpoints corresponding to steps at which the fitness score is better than
previously seen. Depending on whether you are maximizing or minimizing the eval
metric, setconfig.best_model_eval_metric_higher_is_betterto True or False.LoggingSo far this version of JAXline only supports logging to Tensorboard via ourTensorBoardLoggerThe user is expected to return a dictionary of scalars from theirstepandevaluatemethods, andTensorBoardLogger.write_scalarswill periodically write these scalars toTensorBoard.All logging will happen asynchronously to the main thread so as not to interrupt
the training loop.You can specify the frequency with which to log, as well as whether to log by
step or by seconds, by setting thelog_train_data_intervalandinterval_typeconfig flagshere.
Ifconfig.log_all_train_datais set toTrue(Falseby default) JAXline
will cache the scalars from intermediate steps and log them all at once at the
end of the period.JAXline passes theTensorBoardLoggerinstance through to thestepandevaluatemethods to allow the user to perform additional logging inside theirExperimentclass if they so wish. A particular use case for this is if you
want to write images, which can be achieved viaExperimentWriter.write_images.LaunchingSo far this version of JAXline does not support launching remotely.Distribution strategyJAX makes it super simple to distribute your jobs across multiple hosts and
cores. As such, JAXline leaves it up to the user to implement distributed
training and evaluation.Essentially, by decorating a function withjax.pmapyou tell JAX to slice the inputs along the first dimension and then run the
function in parallel on each input slice, across all available local devices (or
a subset thereof). In other words,jax.pmapinvokes the single-program multiple-data (SPMD) paradigm. Then by usingjax.laxcollective
communications operations from within your pmapped function, you can tell JAX to
communicate results between all deviceson all hosts. For example, you may
want to usejax.lax.psumto sum up the gradients across all devices on all hosts, and return the result
to each device (an all-reduce).JAX will then automatically detect which devices are available on each host
allowingjax.pmapandjax.laxto work their
magic.One very important thing to bear in mind is that each time you calljax.pmap,
a separate TPU program will be compiled for the computation it wraps. Therefore
you do not want to be doing this regularly! In particular, for a standard ML
experiment you will want to calljax.pmaponce to wrap your parameter update function,
and then you call this wrapped function on each step, rather than callingjax.pmapon each step, which will kill your performance! This is a very common mistake
for new JAX starters. Luckily it has quite an extreme downside so should be
easily noticeable. In JAXline we actually calljax.pmaponce more innext_device_stateto wrap our function to update device state between steps, so end up with 2 TPU
programs rather than just 1 (but this adds negligible overhead).Random number handlingRandom numbers in JAX might seem a bit unfamiliar to users coming from ordinarynumpyandTensorflow. In these languages we have global stateful PRNGs.
Every time you call a random op it updates the PRNGs global state. However,
stateful PRNGs in JAX would be incompatible with JAX's functional design
semantics, leading to problems with reproducibility and parallelizability. JAX
introduces stateless PRNGs to avoid these issues. The downside of this is that
the user needs to thread random state through their program, splitting a new
PRNG off from the old one every time they want to draw a new random number. This
can be quite onerous, especially in a distributed setting, where you may have
independent PRNGs on each device.In JAXline we take care of this for you. On each step, innext_device_state,
we split a new PRNG from the old one, and optionally specialize it to the host
and/or device based on therandom_mode_trainconfigvalue you specify. We then pass this new PRNG through to yourstepfunction to use on that particular step. At evaluation time, we pass a fresh
PRNG to yourevaluatemethod, initialized according to therandom_mode_evalconfigvalue
you specify. This PRNG will be the same on each call toevaluate(as normally you want your evaluation to be deterministic). If you want
different random behaviour on each call, a simple solution would be to fold in
theglobal_stepi.e.jax.random.fold_in(rng, global_step)at the top of yourevaluatemethod.Of course you are free to completely ignore the PRNGs we pass through to yourstepandevaluatemethods and handle random numbers in your own way, should you have different
requirements.DebuggingPost mortem debuggingBy setting the flag--jaxline_post_mortem(definedhere) on the command-line,
tasks will pause on exceptions (exceptSystemExitandKeyboardInterrupt) and
enter post-mortem debugging using pdb. Paused tasks will hang until you attach
a debugger.Disabling pmap and jitBy setting the flag--jaxline_disable_pmap_jiton the command-line, all pmaps
and jits will be disabled, making it easier to inspect and trace code in a
debugger.Citing JaxlinePlease usethis reference.ContributingThank you for your interest in JAXline. The primary goal of open-sourcing
JAXline was to allow us to open-source our research more easily. Unfortunately,
we are not currently able to accept pull requests from external contributors,
though we hope to do so in future. Please feel free to open GitHub issues.