Optimize jax.vmap for memory

[vcharraut]

Hello,

I have been trying to vectorize functions with jax.vmap but I encounter a problem where it takes too much memory on a higher number of sampes making the system crash.

I try to compute the GAE function for an actor critic algorithm in RL, and I compared a manuel vectorization vs vmap.

Using this data:

step = 512
envs = 4

rewards = np.ones((step, envs))
values = np.ones((step, envs))
flags = np.zeros((step, envs))
next_state_value = np.ones((envs))
gamma = 0.99
gae = 0.95

Manual vectorization:

@functools.partial(jax.jit, static_argnums=(4, 5))
def compute_gae_jax(rewards, state_values, flags, next_state_value, gamma, gae):
    advantages = jnp.zeros_like(rewards)
    adv = jnp.zeros(rewards.shape[1])

    for i in reversed(range(rewards.shape[0])):
        terminal = 1.0 - flags[i]

        returns = rewards[i] + gamma * next_state_value * terminal
        delta = returns - state_values[i]

        adv = gamma * gae * adv * terminal + delta
        advantages = advantages.at[i].set(adv)

        next_state_value = state_values[i]

    return advantages

result_1 = compute_gae_jax(rewards, values, flags, next_state_value, gamma, gae)
%timeit -n 100 compute_gae_jax(rewards, values, flags, next_state_value, gamma, gae)
%memit compute_gae_jax(rewards, values, flags, next_state_value, gamma, gae)

Result:

166 µs ± 21.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
peak memory: 2374.65 MiB, increment: 0.45 MiB

Vmap vectorization

(same implementatiom as here https://github.com/google/flax/tree/main/examples/ppo)

@functools.partial(jax.jit, static_argnums=(4, 5))
@functools.partial(jax.vmap, in_axes=(1, 1, 1, 0, None, None), out_axes=1)
def compute_gae_jax_vmap(rewards, state_values, flags, last_value, gamma, gae):
    advantages = []
    adv = 0.0

    for i in reversed(range(len(rewards))):
        terminal = 1.0 - flags[i]

        returns = rewards[i] + gamma * last_value * terminal
        delta = returns - state_values[i]

        adv = gamma * gae * adv * terminal + delta
        advantages.append(adv)

        last_value = state_values[i]

    advantages = advantages[::-1]
    return jnp.array(advantages)

result_2 = compute_gae_jax_vmap(rewards, values, flags, next_state_value, gamma, gae)
%timeit -n 100 compute_gae_jax_vmap(rewards, values, flags, next_state_value, gamma, gae)
%memit compute_gae_jax_vmap(rewards, values, flags, next_state_value, gamma, gae)

Result:

115 µs ± 3.88 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
peak memory: 10743.67 MiB, increment: 0.00 MiB

Question

Is there anyway to optimize the vmap vectorization memory wise or is it just a trade-off ?
Because it is impossible to use for steps => 1024