auto parallelize

[pharringtonp19]

The following bit of code generates a matrix and a vector and then checks whether the vector lies in the convex hull formed from the matrix.

I would like to make use of a cluster of gpus when determining whether the vector lies in the convex hull.

When I jit the check_convex_hull function, I can see from jax.debug.visualize_array_sharding, that the matrix is not sharded across the gpus. However, it is sharded if I don't jit the function. I am surprised by this.

When I apply this function to real data, this results in the following type of gpu utilization:

def test_sample(key, convex_hull, n, d):
    k1, k2 = jax.random.split(key)
    X = jax.random.normal(k1, shape=(n, d))
    w = jax.random.normal(k2, shape=(n,))
    if convex_hull:
        w = jax.nn.softmax(w)
    p = w.T @ X
    return X, p, w

#@jax.jit
def check_convex_hull(X, p):
    X = jax.device_put(X, sharding)
    jax.debug.visualize_array_sharding(X)
    b = jnp.hstack((jnp.ones(shape=(1,)),p))
    Q = jnp.zeros(shape=(X.shape[0], X.shape[0]))
    c = jnp.zeros(shape=(X.shape[0],))
    A = jnp.vstack((jnp.ones(shape=(X.shape[0],)), X.T))
    G = -jnp.eye(X.shape[0])
    h = jnp.zeros(shape=(X.shape[0],))
    qp = OSQP()
    sol = qp.run(
                init_params=None,
                params_obj=(Q, c), 
                params_eq=(A, b), 
                params_ineq=(G, h)).params.primal
    jax.debug.visualize_array_sharding(sol)
    return jnp.allclose(A @sol, b, atol=1e-3)

  ### --- Testing --- ###
  X, p, w = test_sample(jax.random.PRNGKey(0), True, 1000, 10)
  print(X.shape, p.shape, w.shape)
  status = check_convex_hull(X, p)
  print(f"Status: {status}")

[jakevdp]

Hi - the issue is that device_put is basically a no-op within JIT, and so when you wrap your function in jit you are no longer sharding your data (the sharding must be specified at the JIT boundary, and at the JIT boundary your data is unsharded). This seems a bit surprising from the user perspective – perhaps we should make device_put with sharding specifications an error when run under JIT. @yashk2810, what do you think?

[pharringtonp19]

@jakevdp Thank you for the explanation -- that makes sense!

[yashk2810]

You need to use jax.lax.with_sharding_constraint if you want to constraint shardings of intermediates inside jit.

For inputs/outputs, you can use in_shardings or out_shardings or shard the inputs before you pass them to jitted function (I would recommend doing this).

we should make device_put with sharding specifications an error when run under JIT

Yeah this is in the cards. it's just a semantics breaking change so we need to do it carefully. There have been lots of discussions to merge device_put and with_sharding_constraint but we had decided to keep them separate. But we will revist that soon with the memories work.

[pharringtonp19]

@yashk2810 This is great. I'll shard the input before passing it to the jitted function. Thanks for the explanation.