vmap scales approximately linear in wallclock

[joeryjoery]

Hi, I'm using vmap to batch experiments (i.e., batch over jax.random.key seeds) but when I scale-up the batch-dimension it seems to slow-down approximately linearly on the GPU.

Is this expected? I would've expected this to only slightly slow down the full computation but not to such a large degree.

Example

The way that I've designed my code is as follows;

1. Generate batch of random keys: jax.vmap(jax.random.key)(jnp.asarray(seeds))
2. Define a Learner and DataGeneration object each with a reset and step function. These functions, internally implement batching over random keys acceleration with jax.jit.
3. Data between Learner and DataGeneration is passed in a Python loop so that I can do logging between steps.

So abstractly, this looks something like,

# Do imports

learner = Learner(...)
gen = DataGen(...)

seeds = list(range(5))
keys = jax.vmap(jax.random.key)(jnp.asarray(seeds))
state = learner.reset(keys)  # Internally batches over keys

while True:
    data, metrics_gen = gen.step(state)  # Internally jits and batches over state
    state, metrics_learner = learner.step(state, data)

    # Log metrics from state
    ...

Now when I scale up seeds to e.g., list(range(20)) I notice that the wallclock time per loop-iteration on my RTX 4080 scales almost linearly (bit below). But I somehow expected better performance...

Any thoughts?

[joeryjoery]

So, after some profiling I found where the slow-down arises, it is unsurprisingly isolated to Learner.step.

When doing some offline testing with doing timeit on jax.vmap(jax.vmap(... functions I indeed found that the compute-time scales linearly with the input size. But there is a constant improvement in comparison to using jax.lax.scan and an ever bigger improvement over jax.lax.map.

I think I'm misunderstanding the potential performance benefits with vmap because intuitively it does scale somewhat linearly with the array sizes, but not as aggresively as a loop based method. However, when the array sizes get very very large, it seems that all performance benefits seem to dissapear and using a loop sometimes becomes better. Why is that?

[jakevdp]

Hi, thanks for the question! All told, the results you're seeing aren't that surprising: after all, the number of operations in your program increases linearly with the size of the batch dimension, so absent any sort of implicit device parallelism, you'd expect the wall clock time to increase linearly as well.

If you're running on an accelerator like GPU or TPU, there are situations in which vmapped operations will scale sub-linearly with batch size, particularly when a single operation is not large enough to take advantage of the intrinsic parallelization in the chip's architecture. For example, a vmapped vector product becomes a matrix product, which will generally be faster than performing N separate vector products. But if a single batch is a relatively involved operation, you're probably out of this regime.

Regarding large arrays: when you run a vmap operation, it will attempt to compute everything at once, which can lead to large intermediate array sizes. This memory footprint of vmap can make it less performant than a serial loop in cases where the intermediate arrays have a large size.

Does that help?