Remat (aka checkpoint) with tied variables from custom variable collections

[giganttheo]

Hi.
I am working on a custom model using JAX/Flax, and I use a custom variable collection to be able to pass layer-specific configurations to the model, depending on the input. To be more precise, I pass some Attention Mask that can be defined layer-by-layer and is input-specific.

To save memory usage in cases where I want this info to be shared between layers, I tied the variables using nn.map_variables and a tie method inspired by this comment. It works as expected in most cases, but when I activate gradient checkpointing with the nn.partitioning.remat function, wherever I train the model, there is a warning saying:

('graph', 'encoder', 'block', '0', 'layer', '0', 'SelfAttention') [note: scope of the source of the tied variable] or ('graph', 'encoder', 'block', '1', 'layer', '0', 'SelfAttention') [note: scope of the target of the tied variable] is not a valid variable

I checked the code in Flax and I didn't find the origin of this warning (it doesn't throw an error). And I am not sure how this affect the variable values during the gradient checkpointing, and if it does affect, if there is a way to fix this.

Do you have any ideas to help me please?

[cgarciae]

Hey! Sorry for the late reply. It seems you managed to solve the issue?

[giganttheo]

Hi, thank you for the answer.
I found a way to avoid this warning by using other means to broadcast the collection, for instance by using nn.scan.

For this issue, the Warning is thrown from my own code in the tie_in function that is called by nn.map_variables, I missed that.

From my understanding, with remat, the scope is changed in such a way that it is not possible anymore for the model to access a variable in another scope. In that case it is not possible for a layer to access the parameters, or other custom collections, that are defined for the scope of another layer. Which was possible without remat with the tie_in + nn.map_variables method. Are there prefered ways to broadcast a collection with remat/checkpoint, or to access variables outside of scope, in a similar way as without remat?

The issue is not a blocking one at all but if it is possible to avoid I would like to know how.