pmap(lambda: matmul) takes forever to compile with 2M input shape

[ivandustin]

Hi,

matmul inside pmap takes forever to compile. I tried this in Google Colab using TPU.

So I set up the environment to use TPU.

import jax.tools.colab_tpu
jax.tools.colab_tpu.setup_tpu()

import jax
jax.devices()

[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),
 TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1),
 TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0),
 TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1),
 TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0),
 TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1),
 TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0),
 TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]

I created a test input.

key          = jax.random.PRNGKey(0)
key_a, key_b = jax.random.split(key)
feature_size = 2_048_000
batch_size   = 128
a            = jax.random.uniform(key_a, shape=(batch_size, feature_size))
b            = jax.random.uniform(key_b, shape=(feature_size, batch_size))

First, I tested without pmap and it compiles quickly.

%timeit jax.numpy.matmul(a,b)

100 loops, best of 5: 8.75 ms per loop

However, when I tried to run matmul to all cores in parallel by using pmap, it takes forever to compile.

jax.pmap(lambda _: jax.numpy.matmul(a, b))(jax.numpy.arange(jax.local_device_count()))