Device of JIT Compiled Functions

[YigitElma]

Hi,

I am trying to understand where the JIT-compiled operations will run. I am familiar with operation follows the data concept but I am not sure about certain cases. For example, let's say I have 2 arrays,

# some large array on CPU
A = jnp.zeros((50000000, 50000000), device=jax.devices("cpu")[0])
# smaller array on default device (GPU id0)
b = jnp.zeros(5)

Now, if I have a JIT-compiled function which takes A and b as arguments, where will the computations happen? I assume it is the default device GPU id0. Or is there a heuristic check like 'A doesn't fit in GPU memory, so the computation will be done on CPU'? Also, is the behavior the same for jax.numpy functions (that are not explicitly JIT-compiled by the user but internally compiled by JAX)?

I am specifically forming a large array on the CPU due to memory constraints. The answer will be helpful to prevent any possible attempts to transfer that large array to GPU, and get OOM.

Thanks!

[jakevdp]

In this case, since you explicitly specified the device of A and did not explicitly specify the device of b, array A is committed to the device and b is uncommitted. So if you do a simple operation like A + b (assuming shapes are compatible) then the computation will happen on the device were A resides. For example:

>>> import jax
>>> A = jax.numpy.ones((4, 4), device=jax.devices('cpu')[0])
>>> print(A.device, A.committed)
TFRT_CPU_0 True

>>> b = jax.numpy.arange(4.0)
>>> print(b.device, b.committed)
cuda:0 False

>>> c = A + b
>>> print(c.device, c.committed)
TFRT_CPU_0 True

Now, if I have a JIT-compiled function which takes A and b as arguments, where will the computations happen?

For a general jit-compiled function, that question is not possible to answer without more information, because the function definition may include sharding logic that moves the data to one device or the other (for example, using one or more of the mechanisms discussed at https://docs.jax.dev/en/latest/sharded-computation.html).

[jakevdp]

Note that if both arrays are committed to device and you attempt an operation on them without specifying sharding, you'll get an error. For example:

>>> A = jax.numpy.ones((4, 4), device=jax.devices('cpu')[0])
>>> b = jax.numpy.arange(4.0, device=jax.devices('gpu')[0])
>>> A + b
ValueError: Received incompatible devices for jitted computation. Got argument x of add with shape 
            float32[4,4] and device ids [0] on platform CPU and argument y of add with shape float32[4] 
            and device ids [0] on platform GPU

The fix in this case is to explicitly specify which data should be transferred to which device, as at the documentation link above.

[YigitElma]

@jakevdp Thank you very much for the quick reply! Currently, I don't have sharding, and the only committed array is the large one, so I guess things should work out fine.

[jakevdp]

"sharding" is a more general way of saying "device placement", so what I wrote is still relevant to you, since you're placing arrays on different devices.