Memory Consumption of linen.Dense

[SimonPre]

Dear Community,
i am currently working on the replication of a GNN model for Knowledge Graph completion using Jax and Flax.
The model involves the computation of new node representations in an iterative fashion using linen.Dense layers. This works fine during the forward pass, but during the backwards pass the memory consumption of my model is substantial. It seems like a copy of all the inputs to my dense layer is saved to memory for each iteration. Is there a way to reduce this memory consumption during the backwards pass?
Best regards and thanks in advance,
Simon
`

class GraphModel(nn.Module):

    dims_during_message_passing: list[int]

    @nn.compact
    def __call__(self, input_):
        ...
        for dim in self.dims_during_message_passing:
           ...
           # node_reps is a large array > (1280000*64)
           node_reps = nn.Dense(dim)(node_reps)
           ...

`

[marcvanzee]

Hi Simon, are you wrapping your backward pass in jax.jit?

[SimonPre]

Hi @marcvanzee, yes I wrapped the whole update step for the batches in jit, including the jax.value_and_grad call as well as the state.apply_gradients call.

[cgarciae]

Hey @SimonPre, consider using https://flax.readthedocs.io/en/latest/api_reference/_autosummary/flax.linen.remat.html over your Dense Modules to reduce memory consumption during the backward pass. @8bitmp3 @marcvanzee maybe we should create a guide for this?

[SimonPre]

@cgarciae Thanks for your recommendation, i will check that out.