pmap inside jit created by lax.scan raises jit-of-pmap warning (single-process, multi-device parallelism)

[carlosgmartin]

I have a program of the following form:

import os
os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=4'
import jax
import mpi4py
import mpi4jax
import sys

def mpi_map(f):
    def g(*args, **kwargs):
        rank = mpi4py.MPI.COMM_WORLD.Get_rank()
        result = f(
            *jax.tree_map(lambda x: x[rank], args),
            **jax.tree_map(lambda x: x[rank], kwargs),
        )
        results, _ = mpi4jax.allgather(result)
        return results
    return g

def openai_es(f, std, samples, map_operator, x, key):
    """https://arxiv.org/abs/1703.03864"""
    def g(z):
        u = z * std
        return f(x + u) - f(x - u)
    zs = jax.random.normal(key, (samples, x.size))
    ys = map_operator(g)(zs)
    return ys @ zs / (2 * len(ys) * std)

def f(x):
    return -x @ x

def main():
    # for mpi_map, launch program as mpiexec -n 4 python3 es.py
    map_operator = {
        'vmap': jax.vmap,
        'pmap': jax.pmap,
        'mpi_map': mpi_map,
    }[sys.argv[1]]
    std = 1e-2
    samples = jax.device_count()
    stepsize = 1e-3
    steps = 10**6

    def update(x, key):
        g = openai_es(f, std, samples, map_operator, x, key)
        return x + g * stepsize, x

    x = jax.numpy.ones(10)
    key = jax.random.PRNGKey(0)
    keys = jax.random.split(key, steps)
    x_final, xs = jax.lax.scan(update, x, keys)
    jax.block_until_ready(xs)
    

if __name__ == '__main__':
    main()

Outputs:

$ time python3 es.py vmap        
python3 es.py vmap  3.12s user 0.29s system 110% cpu 3.101 total

$ time mpiexec -n 4 python3 es.py mpi_map
mpiexec -n 4 python3 es.py mpi_map  14.82s user 1.63s system 406% cpu 4.041 total

$ time python3 es.py pmap                
/usr/local/lib/python3.11/site-packages/jax/_src/dispatch.py:285: UserWarning: The jitted function scan includes a pmap. Using jit-of-pmap can lead to inefficient data movement, as the outer jit does not preserve sharded data representations and instead collects input and output arrays onto a single device. Consider removing the outer jit unless you know what you're doing. See https://github.com/google/jax/issues/2926.
  warnings.warn(
python3 es.py pmap  61.50s user 24.82s system 206% cpu 41.821 total

As you can see, the pmap version takes a much longer time. It also raises the warning discussed in #2926. My question is this: What is the recommended way to write a program of this form (in particular, one that uses pmap inside a scan loop) in such a way that it employs single-process, multi-device parallelism effectively? MPI message-passing seems to incur some overhead in the mpi_map version; I'd like to avoid that overhead, if possible.

[BugMaker-Boyan]

I meet the problem too, can anyone help?