Chunked batching with xmap and SerialLoop

[Edenhofer]

For a reconstruction in astrophysics, I need to perform n_los line-of-sight (LOS) integrals within my model. The integrals boil down to weighted sums along one axis of the input. The weights are computed using non-linear functions of some constants, but the LOS integrals themselves are linear in the input. In other words, the LOS integrals can be though of as a very large implicit matrix. Unfortunately, due to memory constraints, it is no possible to batch all LOS integrals in parallel. Instead, I would like to batch through the integrals in chunks.

I have two main constraints for computing the LOS integrals, (1) they need to be fast and (2) they need to be memory efficient. Both constraints seem pretty standard but require some extra though especially when differentiating through the model. Naively, the backwards pass would e.g. store all intermediate LOS weights instead of rematerializing them on the fly. Furthermore, because different LOS integrals access the same elements in an array, the backwards pass has a lot of atomic adds.

My current approach to fulfill both constraints involves manually batching the n_los LOS integrals and defining a custom JVP and VJP (via jax.custom_jvp and jax.custom_derivatives.linear_call). There are multiple downsides to this approach, including that the model gets un-batchable after using linear_call and second-derivatives become inaccessible. (Hopefully all of this will get better with #9073 and #9129.)

To get rid of all of my problems, simplify the code and use more experimental JAX features (:D), I decided to give xmap and SerialLoop a try. Unfortunately, I can't get them to work for my problem. I would be glad if someone could help me get SerialLoop into axis_resources and guide me through how I should tell JAX about my desired rematerialization strategy.

from functools import partial

from jax import numpy as jnp
from jax import random
from jax.experimental.maps import SerialLoop, xmap


def one_los(
    inp,
    idx,
    dist,
    par_std,
    coo_bounds,
):
    sf = jnp.exp(-(dist - coo_bounds) / (dist * par_std * coo_bounds))
    wgt = -jnp.diff(sf)
    return jnp.dot(wgt, inp[idx])


# let's not worry about the random numbers and use one key for everything
key = random.PRNGKey(42)

sky_shape = (786432, 516)  # (plane-of-sky, radius)
n_los = 132400
r_minmax = (40, 1700)
jk_idx_map = random.randint(key, minval=0, maxval=sky_shape[0], shape=(n_los, ))
coo_bounds = jnp.geomspace(*r_minmax, num=sky_shape[1] + 1, endpoint=True)
data_r = random.normal(key, shape=(n_los, ))
data_r = jnp.exp(jnp.log(r_minmax[1]) / 2 + data_r).clip(*r_minmax)
par_std = 0.1 * data_r
synth_sky = jnp.exp(random.normal(key, shape=sky_shape))

ol = partial(one_los, coo_bounds=coo_bounds)
ol = xmap(
    ol,
    in_axes=[[...], ['star', ...], ['star', ...], ['star', ...]],
    out_axes=['star'],
    axis_resources={'star': SerialLoop(20)}  # works without this line
)
_ = ol(synth_sky, jk_idx_map, data_r, par_std)

[PhilipVinc]

Hi @Edenhofer, we encountered a similar problem in a package we developed and ended up implementing a vmap_chunked function that behaves like vmap but operates on chunks of a certain size to bound the memory cost.
In here you can see a few other functions that similarly operates on chunks of a certain size, like vjp_chunked.

They are all implemented on top of jax.lax.scan so it should be differentiable unless I'm mistaken.

You can try to download netket and quickly see if netket.jax.vmap_chunked and similar functions solve your problem. If they do, feel free to steal those functions..

[Edenhofer]

Thank you for you reply.

Unfortunately, vmap_chunked restrictions on in_axes are a bit too restrictive for my use-case. Furthermore, I suspect I would run into the same re-materializing issues that I worked around with custom_jvp. I would much prefer to use native JAX features like e.g. xmap whenever possible instead of writing custom code. If I do the latter, I would basically end up where I am currently at.