Remat inside remat

[LeoXinhaoLee]

Hi, I'm trying to scan through a sequence by groups of chunks, and use remat to avoid saving intermediate variables produced during the processing of each group of chunks.

Specifically, in the below code, compute_chunk processes a chunk of tokens, and compute_group processes a group of chunks. I'm wondering:

1. To avoid saving all Y, do we need the remat decorator on compute_chunk? Or the remat on compute_group already has effect on compute_chunk? If it's needed, how should we set prevent_cse for compute_chunk?

2. The remat on compute_chunk will save the W of each chunk since it's the output of compute_chunk. However, the remat on compute_group only saves the final W of a group of chunks. How does JAX handle this conflict?

from functools import partial
import jax
import jax.numpy as jnp

rng = jax.random.PRNGKey(0)
T = 2048
Chunk_Size = 16
Num_Chunk = T // Chunk_Size
Group_Size = 4
Hidden_Size = 64

X = jax.random.normal(rng, (T, Hidden_Size), dtype=jnp.float32)
W = jax.random.normal(rng, (Hidden_Size, Hidden_Size), dtype=jnp.float32)

@partial(jax.checkpoint, prevent_cse=False)  # TODO: Do we need remat here? How to set prevent_cse here?
def compute_chunk(W, X):
    """
    Under Remat, don't save intermediate variable `Y`
    But `W_new` and `Z` will be saved  # TODO: `compute_group` doesn't want to save `W_new` that is *not the last one in a group*!
    """
    Y = X @ W  # intermediate variable
    Z = Y * X
    W_new = W - Z.transpose(1,0) @ X
    return W_new, Z

@partial(jax.checkpoint, prevent_cse=False)
def compute_group(W, input_group):
    """
    Only want to save the *final* `W` and `output_group` to memory
    Don't save intermediate `W`
    """
    output_group = []
    for i in range(Group_Size):
        W, output_chunk = compute_chunk(W, input_group[i])  # W here is intermediate variable, don't save except for the last one
        output_group.append(output_chunk)
    output_group = jnp.stack(output_group)
    return W, output_group

W_new, Z = jax.lax.scan(compute_group, W, X.reshape(Num_Chunk // Group_Size, Group_Size, Chunk_Size, Hidden_Size))
Z = Z.reshape(-1, Hidden_Size)
print(Z.shape)

Thank you very much for your time and help! @jakevdp