soft_pmap alternative

[Joshuaalbert]

I have come across a few problems with soft_pmap (#5070) and pmap (#5065) which have hampered my ability to distributed moderate chunks of work over devices (virtual or not). To that end, I made this small function that utilised a single pmap call and scan. The single call to pmap avoids #5065, and the scan effectively avoids #5070, namely so that only as much memory as needed is requested on the device. Here is the function,

from timeit import default_timer
from jax import local_device_count, tree_map, numpy as jnp, pmap
from jax.lax import scan


def chunked_pmap(f, *args, chunksize=None):
    """
    Calls pmap on chunks of moderate work to be distributed over devices.
    Automatically handle non-dividing chunksizes, by adding filler elements.

    Args:
        f: jittable, callable
        *args: ndarray arguments to map down first dimension
        chunksize: optional chunk size, should be <= local_device_count(). None is local_device_count.

    Returns: pytree mapped result.
    """
    if chunksize is None:
        chunksize = local_device_count()
    if chunksize > local_device_count():
        raise ValueError("chunksize should be <= {}".format(local_device_count()))
    N = args[0].shape[0]
    remainder = N % chunksize
    if (remainder != 0) and (N > chunksize):
        extra = chunksize - remainder
        args = tree_map(lambda arg: jnp.concatenate([arg, arg[:extra]], axis=0), args)
        N = args[0].shape[0]
    args = tree_map(lambda arg: jnp.reshape(arg, (chunksize, N//chunksize)+arg.shape[1:]), args)
    print("Running on {}".format(chunksize))
    t0 = default_timer()
    def pmap_body(*args):
        def body(state, args):
            return state, f(*args)
        _, result = scan(body, (), args, unroll=1)
        return result
    result = pmap(pmap_body)(*args)
    result = tree_map(lambda arg: jnp.reshape(arg, (-1,)+arg.shape[2:]), result)
    if remainder != 0:
        result = tree_map(lambda x: x[:-extra], result)
    dt = default_timer() - t0
    print("Time to run: {}, rate: {} / s".format(dt, N/dt))
    return result

[Joshuaalbert]

Sadly, this also still suffers from pmap hanging at some point with a few devices spinning at 100% indefinitely.

[Joshuaalbert]

It seems to be non-deterministic, because sometimes it gets stuck and other times not.

[matpalm]

i wonder if the newly introduced sharded_jit will help?