Parallelisation Strategies for Multiple Program

[YigitElma]

Hello,

First of all, thank you for the great tool!

I am trying to parallelize our optimization code which uses Jax for automatic differentiation and jit. My primary aim was to use multiple device to prevent the memory problem we have during certain calculations (usually the gradient of multi-objective cost functions). I think I kind of achieved that. Unfortunately, by doing so, I lost the ability to jit the main function that calculates the jacobian (because the computations occur in multiple devices). A simplified version of my current implementation is like this,

# we want to JIT the methods of a class on specific devices
# for convenience, we define a decorator that does this for us, this will use
# the device attribute of the class to determine the device to JIT on
def jit_with_device(method):
    @functools.wraps(method)
    def wrapper(self, *args, **kwargs):
        # Compile the method with jax.jit for the specific device
        wrapped = jax.jit(method, device=self._device)
        return wrapped(self, *args, **kwargs)

    return wrapper

def pconcat(arrays):
    """Concatenate arrays from multiple devices"""
    # we will use either CPU or GPU[0] for the matrix decompositions, so the
    # array of float64 should fit into single device
    device = jax.devices("gpu")[0]
    out = jnp.concatenate([jax.device_put(x, device=device) for x in arrays])
    return out


# This will be registered as proper pytree
class Optimizable:
    def __init__(self, N, coefs):
        self.N = N
        self.coefs = coefs
        # actual Optimizable class has more attributes, but we don't need them here


# This will be registered as proper pytree
# in the actual code, we have many Objective classes like this with different output shapes
class Objective:
    def __init__(self, opt, grid, target, device_id=0):
        self.opt = opt
        self.grid = grid
        self.target = target
        self._device_id = device_id
        self._device = jax.devices("gpu")[device_id]

    def build(self):
        # construct the constant arrays that will be used in the
        # optimization problem
        # For example,
        # the transform matrix A such that A @ coefs gives the
        # values of the function at the grid points
        self.A = jnp.vstack([jnp.cos(i * self.grid) for i in range(self.opt.N)]).T

    @jit_with_device
    def compute_error(self, coefs, A=None):
        if A is None:
            A = self.A
        vals = A @ coefs
        return vals - self.target

    @jit_with_device
    def jac_error(self, coefs, A=None):
        if A is None:
            A = self.A
        return jax.jacfwd(self.compute_error)(coefs, A)


# This will be registered as proper pytree
class ObjectiveFunctionParallel:
    def __init__(self, objectives):
        self.objectives = objectives
        self.num_device = len(objectives)

    def build(self):
        # construct the constant arrays
        ...

    def compute_error(self, coefs=None, A=None):
        # compute the error for each objective and concatenate them
        fs = [
            obj.compute_error(jax.device_put(coefi, device=obj._device), Ai)
            for obj, coefi, Ai in zip(self.objectives, coefs, A)
        ]
        return pconcat(fs)

    def jac_error(self, coefs=None, A=None):
        # compute the jacobian for each objective and concatenate them
        fs = [
            obj.jac_error(jax.device_put(coefi, device=obj._device), Ai)
            for obj, coefi, Ai in zip(self.objectives, coefs, A)
        ]
        return pconcat(fs)

I skipped the parts that I declare these classes as proper pytrees, but I use register_pytree_node() to be able to jit the methods.

We are using this kind of classes as the inputs of our optimizer, which packages everything using ObjectiveFunctionParallel. I distribute all the data on the objectives to specified devices like this,

N = 40
num_nodes = 15
coefs = np.zeros(N)
coefs[2] = 3
eq = Optimizable(N, coefs)
grid1 = jnp.linspace(-jnp.pi, 0, num_nodes, endpoint=False)
grid2 = jnp.linspace(0, jnp.pi, num_nodes, endpoint=False)
target1 = grid1**2
target2 = grid2**2

obj1 = Objective(eq, grid1, target1, device_id=0)
obj2 = Objective(eq, grid2, target2, device_id=1)
obj1.build()
obj2.build()

# we will put different objectives to different devices
obj1 = jax.device_put(obj1, jax.devices("gpu")[0])
obj2 = jax.device_put(obj2, jax.devices("gpu")[1])
# if we don't assign the eq again, there will be no connection
# between obj.opt.coefs. Since they are supposed to be the same optimizable,
# they need to have same pointers (jax.device_put creates a copy which has
# different memory location)
obj1.opt = eq
obj2.opt = eq

objp_fun = ObjectiveFunctionParallel([obj1, obj2])
objp_fun.build()

Re-assigning the optimizable of each objective with the original one that lives on the default device is necessary for the rest of our optimizer.

With these design considerations in mind, I would like to parallelize this specific portion,

    def jac_error(self, coefs=None, A=None):
        # compute the jacobian for each objective and concatenate them
        fs = [
            obj.jac_error(jax.device_put(coefi, device=obj._device), Ai)
            for obj, coefi, Ai in zip(self.objectives, coefs, A)
        ]
        return pconcat(fs)

Note that in general case, self.objectives are completely different classes with differenr functions and outlut shapes, so I cannot simply use pmap. Since the data is already transferred to the corresponding device, all I need is a 'spark' that will ignite obj.jac_error on each GPU. MPI
(with mpi4jax) seems to be a good option, previously I used openMP in C++ #pragma omp parallel for, I didn't see jax compatible way of doing this though. If there is a native jax way to do this, I would prefer it to reduce the complexity of the installation instruction of our package. I know jax parallelism is primarily for SPMD but I wanted to ask anyway.

Sorry, I couldn't make a shorter explanation of the problem!
I described what I am trying to do also here mpi4jax/mpi4jax#275.

Best regards,

[YigitElma]

In mpi4jax/mpi4jax#275, it has been pointed out to me that I can use mpi4py by casting the arrays into numpy outside the jit.