GPU summing up large array results in different errors

[pawelc]

Hi,

Thank you for amazing library. I have been user since first versions of autograd. Recently I have noticed that when summing up large arrays of numbers using jax and GPU device can give different errors when executing the sum over the same arrray multiple times. I observed the same when executed the code on my system and on the colab. When using CPU device the error is much smaller and is always the same. Also what is strange the error when using jax and CPU device is the same as the error when doing the same operation using pycuda on GPU.

The notebook: https://www.google.com/url?q=https://colab.research.google.com/drive/124TRFMMSybeWKWNu9mmR30sBZx-Px_jC%23scrollTo%3DbdHo1QZpF80c&sa=D&source=hangouts&ust=1620146534289000&usg=AFQjCNHoVBEGAwDyR-ebg33wTfnxgV64OQ

• The result when running on colab with CPU:

True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on cpu:0: 50002916.0, error: 0.4821309447288513
True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on cpu:0: 50002916.0, error: 0.4821309447288513
True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on cpu:0: 50002916.0, error: 0.4821309447288513
True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on cpu:0: 50002916.0, error: 0.4821309447288513
True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on cpu:0: 50002916.0, error: 0.4821309447288513
True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on cpu:0: 50002916.0, error: 0.4821309447288513
True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on cpu:0: 50002916.0, error: 0.4821309447288513
True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on cpu:0: 50002916.0, error: 0.4821309447288513
True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on cpu:0: 50002916.0, error: 0.4821309447288513
True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on cpu:0: 50002916.0, error: 0.4821309447288513

• The result when running on colab with GPU:

True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on gpu:0: 50002932.0, error: 16.48213094472885
True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on gpu:0: 50002932.0, error: 16.48213094472885
True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on gpu:0: 50002936.0, error: 20.48213094472885
True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on gpu:0: 50002936.0, error: 20.48213094472885
True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on gpu:0: 50002928.0, error: 12.482130944728851
True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on gpu:0: 50002936.0, error: 20.48213094472885
True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on gpu:0: 50002932.0, error: 16.48213094472885
True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on gpu:0: 50002928.0, error: 12.482130944728851
True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on gpu:0: 50002928.0, error: 12.482130944728851
True sum: 50002915.5178690553
CPU 32bit precision: 50002928.0, error: 12.482130944728851
jax 32bit precision on gpu:0: 50002928.0, error: 12.482130944728851

The source code:

import numpy as np
from jax import numpy as jnp
import jax

vector_length = int(1e8)
np.random.seed(101)
numpy_double_precision_vector = np.random.rand(vector_length)
numpy_single_precision_vector = numpy_double_precision_vector.astype(np.float32)
jax_single_precision_vector = jax.device_put(jnp.array(numpy_double_precision_vector), jax.devices()[0])

for i in range(10):
  numpy_double_precision_sum = numpy_double_precision_vector.sum()
  numpy_single_precision_sum = numpy_single_precision_vector.sum()
  jax_single_precision_sum = jax_single_precision_vector.sum().block_until_ready()

  numpy_single_precision_error = np.abs(numpy_double_precision_sum - np.float64(numpy_single_precision_sum))
  jax_single_precision_sum_error = np.abs(numpy_double_precision_sum - np.float64(jax_single_precision_sum))

  print(f"True sum: {numpy_double_precision_sum:.10f}")
  print(f"CPU 32bit precision: {numpy_single_precision_sum}, error: {numpy_single_precision_error}")  
  print(f"jax 32bit precision on {jax_single_precision_vector.device_buffer.device()}: {jax_single_precision_sum}, error: {jax_single_precision_sum_error}")

[hawkinsp]

Can you try with XLA_FLAGS=--xla_gpu_deterministic_reductions?

By default GPU uses a non-deterministic reduction algorithm, but there's a deterministic version, not yet enabled by default, that is deterministic.

[pawelc]

Thank you Peter, now it works perfectly.
So it is expected that non-deterministic reduction algorithm tends to give less precise results in such computations (for example compared to pycuda or JAX CPU)?
... next time I should look also at you not yet merged PRs like #4824 :).

Thank you,
Pawel
p

[hawkinsp]

We (well, the XLA folks) are actively working on enabling deterministic reductions by default. Hopefully that will happen soon. It currently regresses performance in some benchmarks, which is why it isn't enabled by default.

I am unsure why the nondeterministic reduction is less accurate, but hopefully it will soon be a moot point when we enable deterministic reductions...

[pawelc]

Thank you for very quick answer.

[cheshire]

It is now enabled by default.