Trouble pmapping a vmap-defined matrix multiplication

[oconnoob]

Description

Hi there,

I've defined a matrix multiplication operation using vmap()

vmap_mv = vmap(lambda v1, v2: jnp.vdot(v1, v2), (0, None), 0)

vmap_mm = vmap(vmap_mv, (None, 1), 1)

I've tried pmapping this operation so that when multiplying e.g. two 8x8 matrices, the matrix-vector multiplication gets parallelized across 8 cloud TPUs, with each TPU handling the mv multiplication between the first matrix and one column of the 2nd matrix, but I'm getting an error.

The multiplication is between two identical upper-triangular 8x8 matrices:

Testing

My vmap_mm function yields the same output as the native jnp.matmul function, and (as expected) the native jnp.matmul operation pmaps just fine

Error

However, when I try to pmap my vmap_mm function, I get errors.

pmap(vmap_mm)(aaa, aaa) yields the error ValueError: vmap was requested to map its argument along axis 1, which implies that its rank should be at least 2, but is only 1 (its shape is (8,))

while pmap(vmap_mm, in_axes=(0, None))(aaa, aaa) yields the error TypeError: Incompatible shapes for dot: got (1,) and (8,).

If anyone could clarify why the error is occurring and how to remedy it, I'd very much appreciate it!

Extra

Side note: if we had 64 TPUs available could we parallelize it so, not only is each MV operation parallelized, but each VV operation within each MV operation is parallelized? I.e. so we have 64 concurrent vector-vector dot products, each between one row of the first matrix and one column of the second matrix (the result of which representing one value in the final matrix from the overarching MM multiplication). If so, how would we change the code to accomplish this?

[jakevdp]

pmap has a similar API as vmap, so I suspect what you want is this:

pmap_mm = pmap(vmap_mv, in_axes=(None, 1), out_axes=1)
pmap_mm(aaa, aaa)

That is, don't use pmap in addition to vmap, use pmap in place of vmap.

[oconnoob]

I see the issue now - so since I want to parallelize the vmap_mv operation over the columns of the matrix B in A*B, I'm pmaping vmap_mv instead of vmap_mm. I would pmap vmap_mm if I were to multiply several B_i's by the same A, to parallelize A*B_i for i in {0, 1, ... n}.

I appreciate the help! And, to clarify, if I wanted to parallelize across 64 TPUs each of the 64 total vector-vector dot products that's involved in the matrix multiplication of two 8x8 matrices, I would completely omit vmap and either (i) pmap(pmap(jnp.vdot())) or (ii) pmap(jnp.vdot()) over multiple axes (? - not sure which one if either).

Thanks again!

[jakevdp]

I don't think nested pmaps are currently supported; see e.g. #1753 actually I think it can be done using named axes; see e.g. https://jax.readthedocs.io/en/latest/notebooks/xmap_tutorial.html

[oconnoob]

Also, in general, can we know (or know how to determine) which method would be faster? That is,

1. pmap(vmap_mv()) will distribute the vmap_mv loops involved with computing the solution across the 8 TPUs, while
2. vmap_mm will vectorize these the vmap_mv loops and compute on one TPU.

Are these therefore equivalent (if you have enough memory)? Or is either

1. vmap_mm faster because you don't have to deal with the overhead of distributing and then recombining the results (not sure if they're actually "recombined" - I assume XLA doesn't do this unless necessary?), or is
2. pmap(vmap_mv()) faster because you break the computation down into 8 pieces (but not sure if this matters for XLA or accelerator computation in general - not very familiar with how the computation actually happens under the hood)

Maybe it depends on the sizes of the matrices involved as well?

Thanks again!

EDIT: Just saw your previous response - I will experiment with named axes! Thanks!

[jakevdp]

I don't think there's any generic answer to that question: it depends on what computation you're doing, the size of your arrays, where they live (if they are pre-sharded to the devices, for example), etc.

[oconnoob]

Got it, thanks again for the help!