pmap slower than jit(vmap)

[krzysztofrusek]

Hi,

I have two fairly complex and independent computations that I want to run on two GPUs with pmap
Surprisingly the pmap-ed version is much slower.
I know that doubling the performance is almost impossible but I expected a better performance.

Below is the complete example that mimics my workload (it requires haiku and chex).
The model is different but the logic is more or less the same.
Any suggestions on what I can do to speed up the whole computation?

import contextlib
import time
from typing import NamedTuple

import chex
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from absl import app, logging

@contextlib.contextmanager
def timer(name: str):
    begin = time.time_ns()
    try:
        yield begin
    finally:
        logging.info(f'Timer {name}[ms] {(time.time_ns() - begin) / int(1e6)}')


class Input(NamedTuple):
    x: chex.Array
    a: chex.Array
    b: chex.Array


@hk.without_apply_rng
@hk.transform
def net(x: Input) -> chex.Array:
    return hk.nets.MLP([128, 128,128, 1])((x.x - x.a) * x.b)


apply_for_many = jax.vmap(net.apply, in_axes=(None, Input(None, 0, 0)))


def main(_):
    rng = jax.random.PRNGKey(42)

    x = np.random.normal(size=(16, 8))
    a = np.random.normal(size=(64))
    b = np.random.normal(size=(64))

    params = net.init(rng, Input(x, a[0], b[0]))

    lr = jnp.logspace(-5, -2, jax.device_count())

    def cost(x: chex.Array, a: chex.Array, b: chex.Array, params: hk.Params) -> chex.Array:
        yhat = apply_for_many(params, Input(x, a, b))
        return jnp.mean(yhat)

    def update(params: hk.Params, x: Input, lr: chex.Array) -> chex.Array:
        df = jax.grad(cost)(*x, params)
        proposal = x.x - lr * df
        proposal = jnp.where(proposal < 0., 0., proposal)
        value = cost(proposal, a, b, params)
        return value, proposal

    jv_update = jax.jit(jax.vmap(update, in_axes=(None, None, 0)))
    p_update = jax.pmap(update, in_axes=(None, None, 0))

    @jax.jit
    def select(values: chex.Array, proposals: chex.Array) -> chex.Array:
        min_idx = jnp.argmin(values)
        return proposals[min_idx, ...]

    # jit
    v, p = jv_update(params, Input(x, a, b), lr)
    v, p = p_update(params, Input(x, a, b), lr)
    new_x = select(v, p).block_until_ready()

    with timer('pmap()'):
        v, p = p_update(params, Input(x + 0.1, a, b), lr)
        new_x = select(v, p).block_until_ready()

    with timer('jit(vmap())'):
        v, p = jv_update(params, Input(x + 0.1, a, b), lr)
        new_x = select(v, p).block_until_ready()


    return 0


if __name__ == '__main__':
    app.run(main)

And example results:

I1209 20:33:09.194915 47073222112192 pmap_vmap_benchmark.py:18] Timer pmap()[ms] 2.554325
I1209 20:33:09.195681 47073222112192 pmap_vmap_benchmark.py:18] Timer jit(vmap())[ms] 0.430292