"Write your model in JAX they said. It'll be deterministic they said."

[samuela]

What is necessary in order to get deterministic results on a GPU? I was under the impression that we get determinism "for free" by virtue of using jax and jax.random's explicit PRNGs, but I'm now learning that this is not the case. For example,

• Non fully reproducible results on GPU flax#33
• add gpu determinism note #4824
• GPU determinism flag #565
• SAME Code & Two Runs ==> Varied Results #4823

... and my own code. Each of these issues/PRs has a variety of comments from different people with varying claims of what flags are sufficient to receive deterministic results, but they vary in exactly what those flags are, whether or not the flags are deprecated, and some of the comments are relatively old. There's also been discussion of a jax.config.deterministic() mode although this does not seem to have materialized yet.

What is the official stance on how to achieve deterministic results with JAX on a GPU? Is there documentation for this? When should I file a determinism bug report?

[mattjj]

Hey! It seems like a good time to answer this question! Or at least, better late than never...

I want to distinguish two things:

1. whether jax.random generates deterministic random bits as a pure function of the random seed given, and
2. other sources of GPU nondeterminism that have nothing to do with jax.random or pseudorandom number generation.

I think the question here is really about the latter. GPUs have several sources of nondeterminism, as mentioned in the (unmerged) #4824. For example, XLA:GPU might autotune cudnn conv kernels or Triton kernels, deciding which to employ based on which it measures to be fastest on your card at the moment. As another example, any multithreaded reduction operation implemented using atomics will be nondeterministic from execution to execution just by virtue of how GPU hardware works. These don't have to do with pseudorandom numbers at all, and would potentially affect any code run on GPU.

I'm 95% sure that we get full determinism, at some performance penalty, by setting XLA_FLAGS='--xla_gpu_deterministic_ops=true --xla_gpu_autotune_level=0' TF_DETERMINISTIC_OPS=1. I'm 90% sure that we can do it with just the XLA_FLAGS, ie that we don't need TF_DETERMINISTIC_OPS=1, unless you're using an AMD GPU.

I'm working on finding out and then finally merging a version of #4824.

EDIT The above is sufficient for determinism, but not necessary. All you need for Nvidia GPUs is XLA_FLAGS=--xla_gpu_deterministic_ops=true. For AMD GPUs you currently need XLA_FLAGS=--xla_gpu_deterministic_ops=true TF_DETERMINISTIC_OPS=1, though we're working on fixing that to only need the XLA flag.

[brett-daley]

This worked for me! But, I had to add both flags for an Nvidia Tesla V100 (SXM2 16GB). I'm using relatively old versions (jax 0.4.9 and jaxlib 0.4.7+cuda11.cudnn82) which might be why.