Any way to submit multi-node CPU jobs

[srijaniiserprinceton]

It would be AMAZING if there is any way to submit jobs across nodes in JAX. We want to use a total of 80 CPU cores but our cluster has 40 cores per node. So we want to use 2 nodes. On trying this out we find that ALL the cores on only one of the nodes is being used while NONE in the other is being used.

We did come across pjit but it seems it is only for GPU or TPU?

[PhilipVinc]

https://mpi4jax.readthedocs.io (note: I'm the author)

Disclaimer: this works great on university clusters that are already setup with MPI.
If the cluster is not setup for MPI this might be somewhat complicated.

[srijaniiserprinceton]

Oh wow! This looks like just the thing I might be looking for. Thanks a lot for providing something like this for the community, much needed! We do have a cluster setup perfectly for MPI.

I have no experience in MPI parallelization (although I have used codes which allowed trivial pooling using mpi4py to cast across nodes (see MPI section here). So, please feel free to correct me if I am saying something wrong below. The function that I jit looks like the following (below is a sketch of the actual problem we have):

# M_all = <a list of sparse matrices M of shape (N x N) each>

def get_sparse_matrix(c, M):
    # M_summed = <sparse matrix>
    for i in range(len(c)):
        M_summed +=  c[i] * M
        
    # returning a sparse matrix
    return M_summed

def model():
    for mat_ind in range(len(M_all)):
        # c = some  column vector of shape (N,)
        M = M_all[mat_ind]  # extracting a sparse matrix
        M_summed = get_sparse_matrix(c, M)
        
        # solving the eigenvalue problem for the densified matrix
        return jnp.eigh(M_summed.todense())

So, as you see the operations are fairly simple. A loop where a weighted sum is happening and the resultant matrix is getting acted upon by an eigenvalue decomposition. So, ideally I would like the for loop in model to be distributed across cores as much as possible. This is because the eigenvalue problem takes the most time. get_sparse_matrix is very cheap.

Since I have never coded internal MPI step myself, do you think implementing the above in mpi4jax would be fairly simple?