[RFC] Scan & Gradient checkpointing in Flax

[patrickvonplaten]

Feature request

We should add scan and remat (gradient checkpointing) to the most important Flax/JAX models (BERT, GPT2, OPT, T5, BART, Wav2Vec2).

Motivation

Scan allows for much faster compilation and memory savings and remat is the equivalent of gradient_checkpointing in PyTorch.

@sanchit-gandhi already uses both features in the Flax Seq2Seq Speech project - see: https://github.com/sanchit-gandhi/seq2seq-speech so it'd be quite trivial to get them working.

Implementation details:

Given that both scan and remat are not related to the model architecture, they should IMO not be in the model's config (We've done this mistake in PyTorch and don't want to repeat it here).

I would advocate for the following API:

model = FlaxBertForMaskedLM.from_pretrained("bert-base-cased")
model.scan()  # or model.scan_enable()
model.unscan()  # or model.scan_disable()

and

model = FlaxBertForMaskedLM.from_pretrained("bert-base-cased")
model.gradient_checkpoint_enable()
model.gradient_checkpoint_disable()

As can be seen here: https://github.com/sanchit-gandhi/seq2seq-speech/blob/b28d0c25c8fad0f9ffa6707f91f7aba320d44a4b/models/modeling_flax_wav2vec2.py#L504

We'll need to re-initialize the flax.linen.module inside the model. However this should be fine since it just means that we do

self.module = self.module_class(config=config, dtype=dtype, use_scan=True, **kwargs)
 self. _is_scan_enabled = True

similar to this line:

transformers/src/transformers/models/wav2vec2/modeling_flax_wav2vec2.py

Line 868 in 71e6027

module = self.module_class(config=config, dtype=dtype, **kwargs)

We can see along the PR how much logic can reside in modeling_flax_utils.py and how much would go into the specific models, e.g. modeling_flax_wav2vec2.py.

The same API / logic could be used for the gradient_checkpointing.

Your contribution

Happy to give this implementation a shot with @sanchit-gandhi and @patil-suraj .

Also would love to hear feedback from @borisdayma @marcvanzee about the API

[borisdayma]

I'm not sure you would need both versions within a same script (scan and unscanned, or with and without checkpointing which affects only training anyway).

Then maybe you could just add it directly as an arg to model.from_pretrained(..., scan=False, gradient_checkpointing=False)

You would just have to use some naming conventions on your params to see if you need to scan/unscan when loading a checkpoint.

[sanchit-gandhi]

Suppose you have a training script, it would be useful to be able to use scan and remat during training for faster compile times and larger batch sizes, and then switch to unscan and no remat during eval for faster inference?

[borisdayma]

I'm not sure it would be worth it:

• Most of the time evaluation is relatively fast
• You would have to reformat your parameters each time between eval and train, potentially leading to memory fragmentation

[KMFODA]

Hey @patrickvonplaten, I'm keen to get gradient checkpointing working in JAX for long-t5. If this is not on the cards to be added soon happy to work on a PR for it if that works with you all?

[sanchit-gandhi]

Hey @KMFODA! There's a PR that is close to being merged: #17843 I'll let you know once it's complete, and you can copy the logic across to Flax T5 in a new PR if that sounds good to you!